# CelerisLab/simulation.py
"""
Top-level orchestrator — assembles LBM field, stepper, body manager,
and CUDA context into a single coherent simulation.

Usage::

    sim = Simulation("configs/config_lbm.json")
    sim.add_cylinder((100, 50), radius=10)
    sim.initialize()
    sim.run(1000)
    macro = sim.get_macroscopic()   # {"rho": ..., "ux": ..., "uy": ...}
"""

from typing import Dict, Optional, Tuple, Any

import os

import numpy as np
import pycuda.driver as cuda

from .config import LBMConfig, BodyConfig, load_lbm_config, load_body_config
from .cuda.context import CudaContext
from .cuda import compiler_v2 as compiler
from .lbm.field import LBMField
from .lbm.stepper import LBMStepper
from .body.objects import Cylinder, Sensor, SimObject
from .body.manager import ObjectManager


class Simulation:
    """High-level simulation handle.

    LBM field/stepper and body manager live at the same level;
    this class orchestrates them.
    """

    def __init__(self,
                 lbm_config_path: Optional[str] = None,
                 body_config_path: Optional[str] = None,
                 device_id: int = 0):
        # Load configs
        self.lbm_cfg = load_lbm_config(lbm_config_path)
        self.body_cfg = load_body_config(body_config_path)

        # CUDA context
        self.ctx = CudaContext(device_id)
        arch = self._resolve_compile_arch()

        # Compile kernel
        compiler.generate_config(self.lbm_cfg, n_objects=0)
        self._ptx_path = compiler.compile_kernel(arch=arch)
        self._module = compiler.load_module(self._ptx_path)

        # LBM field & stepper
        self.field = LBMField(self.lbm_cfg, self._module)
        self.stepper = LBMStepper(self.field, self._module, self.lbm_cfg)

        # Body manager
        self.bodies = ObjectManager(
            self.lbm_cfg.nx, self.lbm_cfg.ny, self.lbm_cfg.nz,
            self.lbm_cfg.nq, self.lbm_cfg,
        )

        self._initialized = False
        self._load_objects_from_body_config()
        if self.bodies.count > 0:
            self.recompile()
        self._assert_object_count_contract(expected_count=self.bodies.count)

    # -- Object management ---------------------------------------------------
    def add_cylinder(self, center: Tuple[float, ...],
                     radius: float) -> int:
        obj = Cylinder(obj_id=-1, center=center, radius=radius)
        return self.bodies.add(obj)

    def add_sensor(self, center: Tuple[float, ...],
                   radius: float) -> int:
        obj = Sensor(obj_id=-1, center=center, radius=radius)
        return self.bodies.add(obj)

    def add_object(self, obj: SimObject) -> int:
        return self.bodies.add(obj)

    def _load_objects_from_body_config(self) -> None:
        """Instantiate simple objects declared in ``config_body.json``."""
        for spec in self.body_cfg.objects:
            obj_type = str(spec.get("type", "")).strip().lower()
            center = tuple(spec.get("center", ()))
            radius = float(spec.get("radius", 0.0))
            omega = spec.get("omega", None)

            if obj_type == "":
                raise ValueError("Body config object is missing required key 'type'.")
            if len(center) not in (2, 3):
                raise ValueError(
                    f"Body config object '{obj_type}' requires center with 2 or 3 entries."
                )

            if obj_type == "cylinder":
                body_id = self.add_cylinder(center=center, radius=radius)
                if omega is not None:
                    self.bodies.get(body_id).state.omega = float(omega)
            elif obj_type == "sensor":
                self.add_sensor(center=center, radius=radius)
            else:
                raise ValueError(
                    f"Unsupported body config object type '{obj_type}'. "
                    "Supported types: cylinder, sensor."
                )

    # -- Compilation ---------------------------------------------------------
    def recompile(self):
        """Re-generate config headers and recompile kernel.

        Call after changing compile-time parameters (collision model, etc.).
        """
        if self._initialized:
            raise RuntimeError(
                "recompile() must not be called after initialize(); "
                "rebuild Simulation instead.")
        arch = self._resolve_compile_arch()
        compiler.generate_config(self.lbm_cfg, n_objects=self.bodies.count)
        self._ptx_path = compiler.compile_kernel(arch=arch)
        self._module = compiler.load_module(self._ptx_path)
        # Reconnect field and stepper to new module
        self.field.module = self._module
        self.field._upload_params()
        _prev_step_count = self.stepper._step_count
        self.stepper = LBMStepper(
            self.field, self._module, self.lbm_cfg,
        )
        self.stepper._step_count = _prev_step_count
        self._assert_object_count_contract(expected_count=self.bodies.count)

    # -- Initialization ------------------------------------------------------
    def initialize(self):
        """Initialize flow field and sync objects to GPU."""
        # Recompile if objects were added after construction.
        if self.bodies.count > 0 and self._read_generated_n_objects() != self.bodies.count:
            self.recompile()

        self.field.flag = self.field.build_channel_flags()
        if self.bodies.count > 0:
            self.field.flag = self.bodies.build_flags(self.field.flag)
        self.field.upload_flags()

        self.stepper.initialize()
        if self.bodies.count > 0:
            self.bodies.sync_to_gpu(self.field, rebuild_flags=False)
        else:
            self.bodies.ensure_packed_buffers()
        self._assert_runtime_contracts()
        if self.lbm_cfg.streaming == "esopull" and self.field.n_curved > 0:
            raise RuntimeError(
                "Unsupported configuration: streaming='esopull' with curved links. "
                "Use streaming='double_buffer' or remove curved-link objects."
            )
        self._initialized = True

    # -- Stepping ------------------------------------------------------------
    def run(self, steps: int, checkpoint_interval: int = 0):
        """Advance simulation by *steps* time steps.

        Args:
            checkpoint_interval: If >0, save checkpoint every N steps.
        """
        if not self._initialized:
            raise RuntimeError("Call initialize() first")
        self._assert_runtime_contracts()
        if checkpoint_interval > 0:
            done = 0
            while done < steps:
                batch = min(checkpoint_interval, steps - done)
                self.stepper.step(
                    batch,
                    action_gpu=self.bodies.action_gpu,
                    obs_gpu=self.bodies.obs_gpu,
                )
                done += batch
                if done < steps or done % checkpoint_interval == 0:
                    self.save_checkpoint()
        else:
            self.stepper.step(
                steps,
                action_gpu=self.bodies.action_gpu,
                obs_gpu=self.bodies.obs_gpu,
            )

    def step(self, n: int = 1):
        """Advance *n* steps (convenience for interactive use)."""
        self.run(n)

    # -- Data access ---------------------------------------------------------
    def get_macroscopic(self) -> Dict[str, np.ndarray]:
        """Download DDF and return rho, ux, uy [, uz]."""
        return self.field.get_macroscopic()

    def get_ddf(self) -> np.ndarray:
        self.field.download_ddf()
        return self.field.ddf.copy()

    def get_flags(self) -> np.ndarray:
        return self.field.flag.copy()

    # -- Runtime parameter updates -------------------------------------------
    def update_runtime_params(self, **kwargs):
        """Update __constant__ d_params without recompiling.

        Accepted: omega, fx, fy, fz, n_objects.
        """
        self.field.update_params(**kwargs)

    # -- Snapshots -----------------------------------------------------------
    def snapshot(self):
        self.field.snapshot()

    def restore(self):
        self.field.restore()

    # -- Checkpoint / Restore (HDF5) ----------------------------------------
    def save_checkpoint(self, path: str = None):
        """Save full simulation state to an HDF5 file."""
        from .common.checkpoint import save_checkpoint
        return save_checkpoint(self.field, self.stepper, self.lbm_cfg,
                               self.bodies, path)

    def load_checkpoint(self, path: str):
        """Restore simulation state from an HDF5 checkpoint."""
        from .common.checkpoint import load_checkpoint
        load_checkpoint(path, self.field, self.stepper, self.lbm_cfg,
                        self.bodies)
        self._initialized = True

    # -- Cleanup -------------------------------------------------------------
    def close(self):
        self.ctx.close()

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def __del__(self):
        try:
            self.close()
        except Exception:
            pass

    def _resolve_compile_arch(self) -> str:
        """Return nvcc arch, respecting explicit compute_capability override."""
        if self.lbm_cfg.compute_capability == "auto":
            return self.ctx.sm_arch
        return f"sm_{''.join(self.lbm_cfg.compute_capability.split('.'))}"

    def _read_generated_n_objects(self) -> int:
        """Read N_OBJS from generated config_objects.h for safety checks."""
        cfg_path = compiler.kernel_path("config/config_objects.h")
        with open(cfg_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line.startswith("#define N_OBJS"):
                    return int(line.split()[-1])
        raise RuntimeError(f"N_OBJS not found in generated header: {cfg_path}")

    def _assert_object_count_contract(self, expected_count: int | None = None) -> None:
        """Ensure generated compile-time N_OBJS matches runtime object count."""
        runtime_count = self.bodies.count if expected_count is None else int(expected_count)
        compiled_count = self._read_generated_n_objects()
        if compiled_count != runtime_count:
            raise RuntimeError(
                "Object count contract mismatch: generated N_OBJS="
                f"{compiled_count}, runtime count={runtime_count}. Recompile required."
            )

    def _assert_obs_layout_contract(self) -> None:
        """Ensure host obs layout matches runtime body count and generated contract."""
        lay = compiler.obs_layout(self.lbm_cfg.dim, self.bodies.count)
        mgr = self.bodies
        if mgr.obs_n_slots != lay.n_slots:
            raise RuntimeError(
                f"obs layout mismatch: n_slots={mgr.obs_n_slots}, expected={lay.n_slots}."
            )
        if (
            mgr.slot_stride_floats != lay.slot_stride_floats
            or mgr.torque0_floats != lay.torque0_floats
            or mgr.sensor0_floats != lay.sensor0_floats
            or mgr.obs_total_floats != lay.total_floats
        ):
            raise RuntimeError(
                "obs layout mismatch between runtime buffers and expected layout "
                f"for dim={self.lbm_cfg.dim}, n_objects={self.bodies.count}."
            )

    def _assert_runtime_contracts(self) -> None:
        """Validate object-count and packed-observation contracts before stepping."""
        self._assert_object_count_contract()
        self._assert_obs_layout_contract()