# drl_pinball/cfd/pinball_env.py
"""
PinballEnv — wraps CelerisLab.Simulation for DRL inference.

This class provides the same telemetry interface as LegacyCelerisLab.FlowField.run(),
but using the new Simulation API. The key difference is that the new API returns
N-step cumulative values, while the old API returned per-step averages.

Usage::

    from pinball_env import PinballEnv

    env = PinballEnv(lbm_config, body_config, device_id=0)
    env.set_cylinders({front_id: 0.0, bottom_id: -0.04, top_id: 0.04})
    result = env.run_and_read(800)
    # result['forces'][body_id] = [fx_per_step, fy_per_step]
    # result['sensors'][body_id] = [ux_per_step, uy_per_step]
    
    env.snapshot()
    env.restore()
    
    env.save_field_tecplot("output.dat")
    env.export_vorticity_png("vorticity.png")
"""

from __future__ import annotations

import json
import os
import sys
import tempfile
from typing import Dict, List, Optional, Tuple, Any

import numpy as np
import pycuda.driver as cuda

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)
_DEFAULT_LBM = os.path.join(_REPO, "configs", "config_lbm_pinball.json")
_DEFAULT_BODY = os.path.join(_REPO, "configs", "config_body.json")

# LBM constants
_CS2 = 1.0 / 3.0  # lattice speed of sound squared


class PinballEnv:
    """High-level wrapper around CelerisLab.Simulation for pinball DRL tasks.
    
    Responsibilities:
    - Create and manage a Simulation instance
    - Provide run_and_read() that matches old API semantics (per-step averages)
    - Manage body ids for sensors and cylinders
    - Support snapshot/restore for checkpointing
    - Export macroscopic fields
    
    Body ID convention (all envs follow this order):
        sensors[0], sensors[1], sensors[2], [disturbance_cylinder], 
        front_cylinder, bottom_cylinder, top_cylinder
    
    Some scenes (illusion, vortex) omit the disturbance cylinder.
    """
    
    def __init__(
        self,
        lbm_config_path: Optional[str] = None,
        body_config_path: Optional[str] = None,
        device_id: int = 0,
        *,
        viscosity: Optional[float] = None,
        velocity: Optional[float] = None,
    ):
        # Build config with optional physics override
        if lbm_config_path is None:
            lbm_config_path = _DEFAULT_LBM
        if body_config_path is None:
            body_config_path = _DEFAULT_BODY
        
        if viscosity is not None or velocity is not None:
            # Create a temp config with overridden physics
            with open(lbm_config_path) as f:
                cfg = json.load(f)
            if viscosity is not None:
                cfg["physics"]["viscosity"] = float(viscosity)
            if velocity is not None:
                cfg["physics"]["velocity"] = float(velocity)
            tmpd = tempfile.mkdtemp(prefix="pinball_env_cfg_")
            tmp_cfg_path = os.path.join(tmpd, "config_lbm.json")
            with open(tmp_cfg_path, "w") as f:
                json.dump(cfg, f, indent=2)
            lbm_config_path = tmp_cfg_path
            
        from CelerisLab import Simulation
        
        self.sim = Simulation(
            lbm_config_path=lbm_config_path,
            body_config_path=body_config_path,
            device_id=device_id,
        )
        
        self._velocity = float(velocity) if velocity is not None else 0.01
        self._device_id = device_id
        self._stream = cuda.Stream()
        
        # Body tracking
        self._body_ids: Dict[str, List[int]] = {
            "sensors": [],
            "cylinders": [],
            "disturbance": [],
        }
        self._body_id_to_name: Dict[int, str] = {}
        
    # -------------------------------------------------------------------
    # Geometry construction
    # -------------------------------------------------------------------
    
    def add_cylinder(self, center: Tuple[float, float], radius: float) -> int:
        """Add a cylinder body. Returns body_id."""
        from CelerisLab import Simulation
        body_id = self.sim.add_body("circle", center=center, radius=radius)
        self._body_ids["cylinders"].append(body_id)
        self._body_id_to_name[body_id] = f"cylinder_{len(self._body_ids['cylinders'])}"
        return body_id
    
    def add_sensor(self, center: Tuple[float, float], radius: float) -> int:
        """Add a sensor body. Returns body_id."""
        body_id = self.sim.add_body("sensor", center=center, radius=radius)
        self._body_ids["sensors"].append(body_id)
        self._body_id_to_name[body_id] = f"sensor_{len(self._body_ids['sensors'])}"
        return body_id
    
    def add_disturbance_cylinder(self, center: Tuple[float, float], radius: float) -> int:
        """Add a disturbance cylinder (upstream). Returns body_id."""
        body_id = self.sim.add_body("circle", center=center, radius=radius)
        self._body_ids["disturbance"].append(body_id)
        self._body_id_to_name[body_id] = "disturbance"
        return body_id
    
    def reinitialize(self):
        """Recompile and reinitialize after adding bodies."""
        self.sim.initialize()
    
    # -------------------------------------------------------------------
    # Runtime control
    # -------------------------------------------------------------------
    
    def set_cylinder_omega(self, body_id: int, omega: float):
        """Set cylinder rotation speed in lattice units."""
        self.sim.set_body(body_id, omega=float(omega))
    
    def set_cylinders(self, omegas: Dict[int, float]):
        """Set multiple cylinder omegas at once. {body_id: omega}."""
        for bid, omega in omegas.items():
            self.sim.set_body(bid, omega=float(omega))
    
    def run_and_read(
        self,
        steps: int,
        omegas: Optional[Dict[int, float]] = None,
        *,
        read_fields: bool = False,
    ) -> Dict[str, Any]:
        """Run N LBM steps and read telemetry.
        
        Parameters
        ----------
        steps : int
            Number of LBM steps to run.
        omegas : dict, optional
            Cylinder omegas to set before running. {body_id: omega}
        read_fields : bool
            If True, also return macroscopic field.
        
        Returns
        -------
        dict with:
            forces : dict {body_id: [fx, fy]} per-step average
            sensors : dict {body_id: [ux, uy]} per-step average
            fields : dict (only if read_fields=True)
        """
        # Set omegas if provided
        if omegas is not None:
            self.set_cylinders(omegas)
        
        # Zero GPU telemetry
        self.sim.bodies.zero_force_segment_async(self._stream)
        if self.sim.field.n_sensor > 0:
            self.sim.bodies.zero_sensor_segment_async(self._stream)
        
        # Run steps
        self.sim.stepper.step(
            int(steps),
            action_gpu=self.sim.bodies.action_gpu,
            obs_gpu=self.sim.bodies.obs_gpu,
            stream=self._stream,
        )
        
        # Download telemetry
        self.sim.bodies.download_obs_full_async(self._stream)
        self._stream.synchronize()
        
        # Read forces (per-step average)
        forces = {}
        for bid in self._body_ids["cylinders"]:
            f = self.sim.read_force(bid)
            forces[bid] = (np.array(f, dtype=np.float32) / float(steps)).tolist()
        
        for bid in self._body_ids["disturbance"]:
            f = self.sim.read_force(bid)
            forces[bid] = (np.array(f, dtype=np.float32) / float(steps)).tolist()
        
        # Read sensors (per-step average, raw sum / steps, NO cell count division)
        sensors = {}
        for bid in self._body_ids["sensors"]:
            s = self.sim.read_sensor(bid, normalize=False)
            sensors[bid] = (np.array(s, dtype=np.float32) / float(steps)).tolist()
        
        result = {
            "forces": forces,
            "sensors": sensors,
            "n_steps": int(steps),
        }
        
        if read_fields:
            macro = self.sim.get_macroscopic()
            result["fields"] = {
                "ux": np.asarray(macro["ux"], dtype=np.float32),
                "uy": np.asarray(macro["uy"], dtype=np.float32),
                "rho": np.asarray(macro["rho"], dtype=np.float32),
            }
        
        return result
    
    def get_sensor_array(self, sensors: Dict[int, list]) -> np.ndarray:
        """Convert sensor dict to flat array in body_id order: [s0_ux, s0_uy, s1_ux, ...]."""
        arr = []
        for bid in sorted(sensors.keys()):
            arr.extend(sensors[bid])
        return np.array(arr, dtype=np.float32)
    
    def get_force_array(self, forces: Dict[int, list]) -> np.ndarray:
        """Convert force dict to flat array in cylinder body_id order."""
        arr = []
        for bid in self._body_ids["cylinders"]:
            if bid in forces:
                arr.extend(forces[bid])
        for bid in self._body_ids["disturbance"]:
            if bid in forces:
                arr.extend(forces[bid])
        return np.array(arr, dtype=np.float32)
    
    def get_force_array_legacy_order(self, forces: Dict[int, list]) -> np.ndarray:
        """Return forces in legacy obs order: dist_cyl first, then front, bottom, top.
        
        This matches the old API's flat obs array layout:
        [dist_fx, dist_fy, front_fx, front_fy, bottom_fx, bottom_fy, top_fx, top_fy]
        """
        arr = []
        # Disturbance cylinders first
        for bid in self._body_ids["disturbance"]:
            if bid in forces:
                arr.extend(forces[bid])
        # Then regular cylinders
        for bid in self._body_ids["cylinders"]:
            if bid in forces:
                arr.extend(forces[bid])
        return np.array(arr, dtype=np.float32)
    
    # -------------------------------------------------------------------
    # Checkpoint / Snapshot
    # -------------------------------------------------------------------
    
    def snapshot(self):
        """Save in-memory snapshot of current DDF state."""
        self.sim.snapshot()
    
    def restore(self):
        """Restore from in-memory snapshot."""
        self.sim.restore()
    
    def save_checkpoint(self, path: str):
        """Save HDF5 checkpoint to disk."""
        self.sim.save_checkpoint(path)
    
    def load_checkpoint(self, path: str):
        """Load HDF5 checkpoint from disk."""
        self.sim.load_checkpoint(path)
    
    # -------------------------------------------------------------------
    # Field export
    # -------------------------------------------------------------------
    
    def get_macroscopic(self) -> Dict[str, np.ndarray]:
        """Download macroscopic field. Returns {rho, ux, uy}."""
        return self.sim.get_macroscopic()
    
    def save_field_tecplot(self, filename: str):
        """Save current flow field in Tecplot format.
        
        Matches the format of old save_field() in legacy envs.
        """
        macro = self.get_macroscopic()
        ux = np.asarray(macro["ux"], dtype=np.float32)
        uy = np.asarray(macro["uy"], dtype=np.float32)
        
        nx, ny = ux.shape
        u0 = self._velocity
        
        with open(filename, "w") as f:
            f.write('Title= "LBM 2D"\r\n')
            f.write('VARIABLES= "X","Y","flag","U","V",\r\n')
            f.write(f"ZONE T= \"BOX\",I= {nx},J= {ny},F=POINT\r\n")
            for j in range(ny):
                for i in range(nx):
                    u_val = ux[i, j] / u0
                    v_val = uy[i, j] / u0
                    f.write(f"{i},{j},0,{u_val},{v_val}\r\n")
    
    def export_vorticity_png(self, path: str, title: str = ""):
        """Export vorticity field as PNG."""
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            return
        
        macro = self.get_macroscopic()
        ux = np.asarray(macro["ux"], dtype=np.float64)
        uy = np.asarray(macro["uy"], dtype=np.float64)
        
        omega = np.gradient(uy, axis=1) - np.gradient(ux, axis=0)
        
        abs_o = np.abs(omega[np.isfinite(omega)])
        vmax = float(np.percentile(abs_o, 99.5)) if abs_o.size > 0 else 1.0
        if vmax <= 0:
            vmax = 1.0
        
        ny, nx = omega.shape
        fig, ax = plt.subplots(figsize=(min(18, max(8, nx / 60)), min(10, max(3, ny / 40))))
        im = ax.imshow(omega, origin="lower", aspect="equal", cmap="RdBu_r",
                       vmin=-vmax, vmax=vmax, extent=(0, nx - 1, 0, ny - 1))
        ax.set_xlabel("x (lattice)")
        ax.set_ylabel("y (lattice)")
        if title:
            ax.set_title(title)
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label=r"$\omega_z$")
        fig.tight_layout()
        fig.savefig(path, dpi=150, bbox_inches="tight")
        plt.close(fig)
    
    # -------------------------------------------------------------------
    # Cleanup
    # -------------------------------------------------------------------
    
    def close(self):
        """Release GPU resources."""
        self.sim.close()
    
    def __del__(self):
        try:
            self.close()
        except Exception:
            pass


def _omega_from_nu(nu: float) -> float:
    """Convert kinematic viscosity to relaxation parameter omega."""
    cs2 = 1.0 / 3.0
    return 1.0 / (3.0 * nu / 1.0 + 0.5)