# CelerisLab/config.py
"""
Unified configuration system.

Two JSON files:
  config_lbm.json  – grid, physics, method, cuda
  config_body.json – object list

LBMConfig / BodyConfig are plain dataclasses; they translate to
macro dicts consumed by `cuda/compiler_v2.py` for header generation.
"""

import json
import os
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

# ---------------------------------------------------------------------------
# String ↔ integer mappings (used in JSON and macro generation)
# ---------------------------------------------------------------------------
LATTICE_MODEL_MAP = {
    "D2Q9":  {"dim": 2, "nq": 9, "id": 0},
    "D3Q19": {"dim": 3, "nq": 19, "id": 1},
}
COLLISION_MAP = {"SRT": 0, "TRT": 1, "MRT": 2}
STREAMING_MAP = {"double_buffer": 0, "esopull": 1}
PRECISION_MAP = {"FP32": 0, "FP16S": 1, "FP16C": 2}
INLET_MAP = {"uniform": 0, "parabolic": 1}
INLET_SCHEME_MAP = {
    "zou_he_local": 0,
    "channel_stabilized": 1,
    "equilibrium": 2,
    "regularized": 3,
}
OUTLET_MAP = {"neq_extrap": 0, "zero_gradient": 1, "blended": 2}
Y_WALL_BC_MAP = {"bounce_back": 0, "free_slip": 1}
DTYPE_MAP = {"FP32": "float", "FP64": "double"}


# ---------------------------------------------------------------------------
# LBM config
# ---------------------------------------------------------------------------
# ⚠ THIS IS NOT THE CONFIGURATION LOCATION.
# All parameters come from configs/config_lbm.json (the single source of truth).
# Use load_lbm_config() to construct; do not instantiate LBMConfig() with
# bare defaults.  Tests may pass explicit kwargs for override only.
# See configs/CONFIG.md for parameter documentation.
# ---------------------------------------------------------------------------
@dataclass
class LBMConfig:
    # Lattice model (canonical: determines dim and nq)
    lattice_model: str = "D2Q9"

    # Grid
    nx: int = 512
    ny: int = 256
    nz: int = 1

    # Physics
    data_type: str = "FP32"
    viscosity: float = 0.002
    velocity: float = 0.03
    rho: float = 1.0

    # Method – collision / streaming / storage
    collision: str = "SRT"
    streaming: str = "double_buffer"
    store_precision: str = "FP32"
    ddf_shifting: bool = False
    les_enabled: bool = False
    les_cs: float = 0.16
    les_closed_form: bool = True
    trt_magic_param: float = 0.1875
    inlet_profile: str = "parabolic"
    inlet_scheme: str = "zou_he_local"
    outlet_mode: str = "neq_extrap"
    outlet_blend_alpha: float = 0.7
    outlet_backflow_clamp: bool = True
    inlet_trt_neq_damp: float = 0.5
    inlet_regularized_neq_damp: float = 0.5
    outlet_srt_neq_damp: float = 0.5
    y_wall_bc: str = "bounce_back"
    omega_min: float = 0.01
    omega_max: float = 1.99

    # CUDA
    threads_per_block: int = 256
    compute_capability: str = "auto"

    # Derived (computed on validate — do not set manually)
    dim: int = 0
    nq: int = 0
    omega: float = 0.0

    @property
    def is_d2q9(self) -> bool:
        return self.lattice_model == "D2Q9"

    @property
    def is_d3q19(self) -> bool:
        return self.lattice_model == "D3Q19"

    def validate(self):
        assert self.lattice_model in LATTICE_MODEL_MAP, (
            f"Unknown lattice_model: {self.lattice_model}. "
            f"Supported: {list(LATTICE_MODEL_MAP.keys())}")
        lm = LATTICE_MODEL_MAP[self.lattice_model]
        self.dim = lm["dim"]
        self.nq = lm["nq"]
        self.lattice_model_id = lm["id"]
        assert self.collision in COLLISION_MAP, f"Unknown collision: {self.collision}"
        assert self.streaming in STREAMING_MAP, f"Unknown streaming: {self.streaming}"
        assert self.y_wall_bc in Y_WALL_BC_MAP, f"Unknown y_wall_bc: {self.y_wall_bc}"
        assert self.data_type in DTYPE_MAP, f"Unknown data_type: {self.data_type}"
        assert self.store_precision in PRECISION_MAP, (
            f"Unknown store_precision: {self.store_precision}")
        assert self.inlet_scheme in INLET_SCHEME_MAP, (
            f"Unknown inlet_scheme: {self.inlet_scheme}")
        if self.store_precision == "FP16C":
            raise ValueError(
                "store_precision='FP16C' is not supported in the current runtime path. "
                "Supported store_precision values: FP32, FP16S."
            )
        if self.data_type != "FP32":
            raise ValueError(
                "Only data_type='FP32' is supported: CUDA kernels and LBMField "
                "paths use float32; LBtype is not wired through device code yet."
            )
        assert self.nx > 0 and self.ny > 0 and self.nz > 0
        if self.dim == 2:
            assert self.nz == 1, "nz must be 1 for 2D"
        if not (0.0 <= self.inlet_trt_neq_damp <= 1.0):
            raise ValueError("method.inlet.trt_neq_damp must lie in [0, 1].")
        if not (0.0 <= self.inlet_regularized_neq_damp <= 1.0):
            raise ValueError("method.inlet.regularized_neq_damp must lie in [0, 1].")
        if not (0.0 <= self.outlet_srt_neq_damp <= 1.0):
            raise ValueError("method.outlet.srt_neq_damp must lie in [0, 1].")
        if self.dim == 3 and self.y_wall_bc == "free_slip":
            raise ValueError("y_wall_bc='free_slip' is currently supported for D2Q9 only.")
        if self.omega_max >= 2.0:
            raise ValueError("method.omega_guard.max must be < 2.0.")
        if self.omega_min >= self.omega_max:
            raise ValueError(
                f"Invalid omega guard window: min={self.omega_min}, max={self.omega_max}."
            )
        self.omega = 1.0 / (3.0 * self.viscosity + 0.5)
        if not (self.omega_min <= self.omega <= self.omega_max):
            raise ValueError(
                f"Computed omega={self.omega:.6f} outside guard "
                f"[{self.omega_min:.6f}, {self.omega_max:.6f}]."
            )

    def to_macros(self) -> Dict[str, Any]:
        """Return flat dict of macro_name → value for header generation."""
        self.validate()
        return {
            "NT": self.threads_per_block,
            "MULT_GPU": 0,
            "NX": self.nx,
            "NY": self.ny,
            "NZ": self.nz,
            "LATTICE_MODEL": self.lattice_model_id,
            "LBtype": DTYPE_MAP[self.data_type],
            "VIS": f"{self.viscosity:.10f}",
            "RHO": f"{self.rho:.1f}",
            "U0": f"{self.velocity}",
            "COLLISION_MODEL": COLLISION_MAP[self.collision],
            "STREAMING_MODEL": STREAMING_MAP[self.streaming],
            "STORE_PRECISION": PRECISION_MAP[self.store_precision],
            "USE_DDF_SHIFTING": int(self.ddf_shifting),
            "USE_LES": int(self.les_enabled),
            "LES_CS": f"{self.les_cs:.6f}f",
            "LES_CLOSED_FORM": int(self.les_closed_form),
            "INLET_PROFILE": INLET_MAP[self.inlet_profile],
            "INLET_SCHEME": INLET_SCHEME_MAP[self.inlet_scheme],
            "OUTLET_MODE": OUTLET_MAP[self.outlet_mode],
            "OUTLET_BLEND_ALPHA": f"{self.outlet_blend_alpha:.3f}f",
            "OUTLET_BACKFLOW_CLAMP": int(self.outlet_backflow_clamp),
            "Y_WALL_BC": Y_WALL_BC_MAP[self.y_wall_bc],
            "OMEGA_COLLISION_MIN": f"{self.omega_min:.2f}f",
            "OMEGA_COLLISION_MAX": f"{self.omega_max:.3f}f",
            "TRT_MAGIC_PARAM": f"{self.trt_magic_param:.6f}f",
            "INLET_TRT_NEQ_DAMP": f"{self.inlet_trt_neq_damp:.4f}f",
            "INLET_REGULARIZED_NEQ_DAMP": f"{self.inlet_regularized_neq_damp:.4f}f",
            "OUTLET_SRT_NEQ_DAMP": f"{self.outlet_srt_neq_damp:.4f}f",
        }


# ---------------------------------------------------------------------------
# Body config
# ---------------------------------------------------------------------------
@dataclass
class BodyConfig:
    objects: List[Dict[str, Any]] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Loaders
# ---------------------------------------------------------------------------
def _find_config(filename: str, explicit_path: Optional[str] = None) -> str:
    """Search for config file in standard locations."""
    candidates = []
    if explicit_path:
        candidates.append(explicit_path)
    env_dir = os.environ.get("CELERISLAB_CONFIG_DIR")
    if env_dir:
        candidates.append(os.path.join(env_dir, filename))
    candidates.append(os.path.join(os.getcwd(), "configs", filename))
    pkg_root = os.path.dirname(os.path.abspath(__file__))
    candidates.append(os.path.join(pkg_root, "configs", filename))
    for p in candidates:
        if os.path.isfile(p):
            return os.path.abspath(p)
    raise FileNotFoundError(
        f"Config '{filename}' not found.  Searched:\n"
        + "\n".join(f"  - {p}" for p in candidates)
    )


def load_lbm_config(path: Optional[str] = None) -> LBMConfig:
    fpath = _find_config("config_lbm.json", path)
    with open(fpath) as f:
        d = json.load(f)
    g, p, m, c = d["grid"], d["physics"], d["method"], d["cuda"]
    cfg = LBMConfig(
        lattice_model=g["lattice_model"],
        nx=g["nx"], ny=g["ny"], nz=g["nz"],
        data_type=p["data_type"], viscosity=p["viscosity"],
        velocity=p["velocity"], rho=p["rho"],
        collision=m["collision"], streaming=m["streaming"],
        store_precision=m["store_precision"],
        ddf_shifting=m["ddf_shifting"],
        les_enabled=m["les"]["enabled"], les_cs=m["les"]["cs"],
        les_closed_form=m["les"].get("closed_form", True),
        trt_magic_param=m["trt"]["magic_param"],
        inlet_profile=m["inlet"]["profile"],
        inlet_scheme=m["inlet"].get("scheme", "zou_he_local"),
        outlet_mode=m["outlet"]["mode"],
        outlet_blend_alpha=m["outlet"]["blend_alpha"],
        outlet_backflow_clamp=m["outlet"]["backflow_clamp"],
        inlet_trt_neq_damp=m["inlet"].get("trt_neq_damp", 0.5),
        inlet_regularized_neq_damp=m["inlet"].get("regularized_neq_damp", 0.5),
        outlet_srt_neq_damp=m["outlet"].get("srt_neq_damp", 0.5),
        y_wall_bc=m.get("y_wall_bc", "bounce_back"),
        omega_min=m["omega_guard"]["min"],
        omega_max=m["omega_guard"]["max"],
        threads_per_block=c["threads_per_block"],
        compute_capability=c["compute_capability"],
    )
    cfg.validate()
    return cfg


def load_body_config(path: Optional[str] = None) -> BodyConfig:
    fpath = _find_config("config_body.json", path)
    with open(fpath) as f:
        d = json.load(f)
    return BodyConfig(objects=d.get("objects", []))