# 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_body("circle", center=(50, 50), radius=10)
    sim.initialize()
    sim.run(1000)
    macro = sim.get_macroscopic()   # {"rho": ..., "ux": ..., "uy": ...}
    force = sim.read_force(0)       # force on body 0
"""

from __future__ import annotations

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 SimObject
from .body.manager import ObjectManager
from .body.ddf_patch import patch_ddf_for_body_sync, upload_patched_ddf


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):
        # Internal stream for async DTOH coordination
        self._stream: Optional[cuda.Stream] = None

        # 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)

    # -- Stream access -------------------------------------------------------
    @property
    def stream(self) -> cuda.Stream:
        """Internal CUDA stream for async DTOH coordination."""
        if self._stream is None:
            self._stream = cuda.Stream()
        return self._stream

    # -- Object management ---------------------------------------------------
    def add_body(self, type: str = "circle", center=None,
                 radius: float = 0.0, **kwargs) -> int:
        """Add a simulation body. Returns body_id.

        Before ``initialize()`` the object is added to the formal registry
        immediately.  After ``initialize()`` the object is staged via
        ``stage_add()`` and will be committed at the next ``sync_bodies()``.

        Args:
            type: ``"circle"``, ``"sensor"``, or ``"force_region"``
                (future: ``"polygon"``, ``"mesh"``).
            center: (x, y) center coordinates (lattice units).
            radius: body radius (lattice units).
            **kwargs: passed to the geometry constructor.

        Returns:
            int body_id (formal id before initialize; -1 if staged).
        """
        type_lower = str(type).strip().lower()

        if type_lower == "circle":
            from .body.geometry.circle import CircleGeometry
            cx, cy = float(center[0]), float(center[1])
            geom = CircleGeometry(cx, cy, radius)
            obj = SimObject(obj_id=-1, geometry=geom,
                            center=center, radius=radius,
                            is_sensor=False)
        elif type_lower == "sensor":
            from .body.geometry.circle import CircleGeometry
            cx, cy = float(center[0]), float(center[1])
            geom = CircleGeometry(cx, cy, radius)
            obj = SimObject(obj_id=-1, geometry=geom,
                            center=center, radius=radius,
                            is_sensor=True)
        elif type_lower == "force_region":
            from .body.geometry.circle import CircleGeometry
            cx, cy = float(center[0]), float(center[1])
            geom = CircleGeometry(cx, cy, radius)
            obj = SimObject(obj_id=-1, geometry=geom,
                            center=center, radius=radius,
                            is_force_region=True)
        else:
            raise ValueError(
                f"Unknown body type '{type}'. "
                "Supported: 'circle', 'sensor', 'force_region'."
            )

        if self._initialized:
            self.bodies.stage_add(obj)
            return -1
        return self.bodies.add(obj)

    def remove_body(self, id: int) -> None:
        """Remove a body from the simulation.

        Before ``initialize()`` the object is removed from the formal registry
        immediately.  After ``initialize()`` the removal is staged and will be
        committed at the next ``sync_bodies()``.

        Args:
            id: body_id from ``add_body()``.
        """
        if self._initialized:
            self.bodies.stage_remove(id)
        else:
            self.bodies.remove(id)

    # -- Legacy convenience (keep for backward compat) -----------------------
    def add_cylinder(self, center: Tuple[float, ...],
                     radius: float) -> int:
        """Add a cylinder body (convenience wrapper around add_body)."""
        return self.add_body("circle", center=center, radius=radius)

    def add_sensor(self, center: Tuple[float, ...],
                   radius: float) -> int:
        """Add a sensor body (convenience wrapper around add_body)."""
        return self.add_body("sensor", center=center, radius=radius)

    def add_object(self, obj: SimObject) -> int:
        """Add a pre-configured SimObject directly."""
        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_body("circle", center=center, radius=radius)
                if omega is not None:
                    self.bodies.get(body_id).state.omega = float(omega)
            elif obj_type == "sensor":
                self.add_body("sensor", center=center, radius=radius)
            else:
                raise ValueError(
                    f"Unsupported body config object type '{obj_type}'. "
                    "Supported types: cylinder, sensor."
                )

    # -- Runtime control -----------------------------------------------------
    def set_body(self, id: int, omega: float = None,
                 vx: float = None, vy: float = None) -> None:
        """Set runtime body state (host action array only, no H2D).

        The GPU action buffer is updated at the next ``run()`` call.

        Args:
            id: body_id from add_body().
            omega: rotation speed (lattice units).  ``None`` = no change.
            vx: translation velocity x (future).  ``None`` = no change.
            vy: translation velocity y (future).  ``None`` = no change.
        """
        obj = self.bodies.get(id)
        changed = False
        if omega is not None:
            obj.state.omega = float(omega)
            changed = True
        if vx is not None:
            obj.state.vx = float(vx)
            changed = True
        if vy is not None:
            obj.state.vy = float(vy)
            changed = True
        if changed:
            self.bodies._refresh_action_from_objects()

    def set_force(self, id: int, fx: float = 0.0, fy: float = 0.0) -> None:
        """Set force density on a force_region object (host only, no H2D).

        The GPU action buffer is updated at the next ``run()`` call.

        Args:
            id: body_id from add_body(type='force_region', ...).
            fx: x-component of force density (lattice units).
            fy: y-component of force density (lattice units).
        """
        self.bodies.set_force_state(body_id=id, fx=float(fx), fy=float(fy))

    # -- Telemetry readback --------------------------------------------------
    def read_force(self, id: int) -> np.ndarray:
        """Return the force vector on body *id* from the pinned obs buffer."""
        if self.bodies.obs_pinned is None:
            raise RuntimeError("No obs buffer. Call run() first.")
        return self.bodies.read_force(id)

    def read_torque(self, id: int) -> np.ndarray:
        """Return the torque on body *id* from the pinned obs buffer."""
        if self.bodies.obs_pinned is None:
            raise RuntimeError("No obs buffer. Call run() first.")
        return self.bodies.read_torque(id)

    def read_sensor(self, id: int, *, normalize: bool = True) -> np.ndarray:
        """Return the sensor reading for body *id* from the pinned obs buffer.

        Args:
            normalize: If True, return area-averaged velocity. If False,
                return the raw sum accumulated by the GPU SensorKernel.
        """
        if self.bodies.obs_pinned is None:
            raise RuntimeError("No obs buffer. Call run() first.")
        return self.bodies.read_sensor(id, normalize=normalize)

    def read_body(self, id: int, *, stream: cuda.Stream | None = None):
        """Return unified telemetry for one body.

        Args:
            id: body_id from ``add_body()``.
            stream: Optional CUDA stream to synchronise before reading.
                If ``None``, uses the internal stream.

        Returns:
            BodyTelemetry with fields ``force``, ``torque``, ``sensor``.
            ``sensor`` is a dim-length array of zeros for non-sensor bodies.
        """
        if self.bodies.obs_pinned is None:
            raise RuntimeError("No obs buffer. Call run() first.")
        if stream is None:
            stream = self.stream
        if stream is not None:
            stream.synchronize()
        return self.bodies.read_body(id)

    def read_bodies(self, *, stream: cuda.Stream | None = None) -> np.ndarray:
        """Return all bodies' telemetry as a flat float32 array.

        Layout follows the packed obs buffer: ``[force_0..n-1, torque_0..n-1, sensor_0..n-1]``
        where each segment is ``max(n_objects, 1) * dim`` floats.
        Total length equals ``obs_layout(dim, n_objects).total_floats``.

        This is the most efficient readback for DRL control loops
        (single contiguous D2H copy, no Python per-body overhead).
        """
        if self.bodies.obs_pinned is None:
            raise RuntimeError("No obs buffer. Call run() first.")
        if stream is None:
            stream = self.stream
        if stream is not None:
            stream.synchronize()
        return self.bodies._obs_array()

    # -- Compilation ---------------------------------------------------------
    def _recompile(self, expected_n_objects: int | None = None) -> None:
        """Re-generate config headers and recompile the kernel.

        If *expected_n_objects* is ``None``, uses ``self.bodies.count``.
        Preserves ``step_count`` and GPU DDF/flag buffers across the
        module reload.  Does NOT preserve ``action_gpu`` / ``obs_gpu``
        (those are recreated by ``sync_to_gpu`` or ``run()``).

        This is safe to call before or after ``initialize()``.
        """
        if expected_n_objects is None:
            expected_n_objects = self.bodies.count
        arch = self._resolve_compile_arch()
        compiler.generate_config(self.lbm_cfg, n_objects=expected_n_objects)
        # Remove stale PTX to prevent PyCUDA module handle conflicts.
        import os as _os
        if _os.path.exists(self._ptx_path):
            _os.remove(self._ptx_path)
        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.invalidate_module_constant_cache()
        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=expected_n_objects)

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

        Call after changing compile-time parameters (collision model, etc.).
        Must not be called after ``initialize()`` -- use ``sync_bodies()``
        for runtime topology changes instead.
        """
        if self._initialized:
            raise RuntimeError(
                "recompile() must not be called after initialize(); "
                "use sync_bodies() for runtime body changes.")
        self._recompile()

    # -- 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()
        self._initialized = True

    # -- Runtime body topology sync -------------------------------------------
    def sync_bodies(self) -> None:
        """Apply pending body edits (add/remove) to a running simulation.

        This is the main entry point for runtime body topology changes.
        It downloads the current DDF, rebuilds topology, recompiles the
        kernel for the new object count, patches the DDF for geometry
        changes, and re-uploads everything to GPU.

        Currently only supports ``double_buffer`` streaming mode.
        ``esopull`` mode raises ``NotImplementedError``.

        Raises:
            RuntimeError: If called before ``initialize()``.
            NotImplementedError: If streaming mode is ``esopull``.
        """
        if not self._initialized:
            raise RuntimeError("Call initialize() before sync_bodies()")

        # 1. Synchronize GPU
        self.stream.synchronize()

        # 2. Check for pending edits
        if not self.bodies.has_pending_edit():
            return

        # 3. Check streaming mode
        if self.lbm_cfg.streaming == "esopull":
            raise NotImplementedError(
                "Runtime body sync is not yet supported for esopull "
                "streaming mode.  Use double_buffer instead.")

        # 4. Download current DDF to host (snapshot for patching)
        self.field.download_ddf(
            step_id=self.stepper.step_count, force=True)
        old_ddf = self.field.ddf.copy()
        old_flags = self.field.flag.copy()

        # 5. Build sync plan
        plan = self.bodies.build_sync_plan(self.field)

        # 6. Runtime recompile for new object count
        self._recompile(plan.next_count)

        # 7. Apply new topology data (flags, compact lists, params)
        self.bodies.apply_sync_plan(self.field, plan)

        # 8. DDF patch for geometry changes
        cl_fluid_idx = plan.curved_host[0]
        patch_ddf_for_body_sync(
            self.field,
            old_ddf, old_flags, plan.next_flags,
            plan.added_solid_mask, plan.released_fluid_mask,
            curved_fluid_indices=cl_fluid_idx,
        )
        upload_patched_ddf(
            self.field,
            plan.added_solid_mask, plan.released_fluid_mask,
            curved_fluid_indices=cl_fluid_idx,
        )

        # 9. Commit formal objects and clear pending
        self.bodies.commit_pending(plan.next_objects, plan.new_sensor_counts)

        # 10. Build action/obs telemetry for the new object set
        self.bodies._refresh_action_from_objects()
        self.bodies._allocate_packed_telemetry()
        assert self.bodies.obs_pinned is not None
        self.bodies.obs_pinned.fill(0)
        cuda.memcpy_htod(self.bodies.obs_gpu, self.bodies.obs_pinned)

        self._assert_runtime_contracts()

    # -- Stepping ------------------------------------------------------------
    def run(self, steps: int, *,
            stream: cuda.Stream | None = None,
            upload_act: bool = True,
            sync_obs: bool = True,
            checkpoint_interval: int = 0):
        """Advance simulation by *steps* time steps.

        Args:
            steps: Number of LBM steps.
            stream: CUDA stream for async operations.  ``None`` = internal stream.
            upload_act: If True, upload host action array to ``action_gpu``
                before the step group.
            sync_obs: If True, download ``obs_gpu`` to host pinned buffer
                after the step group.
            checkpoint_interval: If >0, save checkpoint every N steps.
        """
        if not self._initialized:
            raise RuntimeError("Call initialize() first")
        # Discard any uncommitted body edits before stepping.
        if self.bodies.has_pending_edit():
            self.bodies.clear_pending_edits()

        # Resolve stream
        if stream is None:
            stream = self.stream

        # Async upload action
        if upload_act and self.bodies.count > 0:
            self.bodies._upload_action_async(stream)

        # Zero obs force segment before step group
        self.bodies.zero_force_segment_async(stream)

        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,
                    stream=stream,
                )
                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,
                stream=stream,
            )

        # Async download obs
        if sync_obs:
            self.bodies.download_obs_full_async(stream)

        stream.synchronize()

    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(step_id=self.stepper.step_count)

    def get_ddf(self) -> np.ndarray:
        self.field.download_ddf(step_id=self.stepper.step_count)
        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 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(
                "obs layout mismatch: n_slots="
                f"{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 "
                f"layout for dim={self.lbm_cfg.dim}, "
                f"n_objects={self.bodies.count}."
            )

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