# drl_pinball/scenes/karman_cloak/re100_scene.py
"""
Karman cloak re100 scene — inference orchestration for the Karman vortex street
cloaking scenario at code Reynolds number 100 (Re_D=50).

This scene exactly replicates the flow configuration of:
    env_karman_cloak_standard.py + model d1a3o12_re100

Geometry (in lattice units, L0=20):
    Disturbance cylinder:  center=(200, CENTER_Y),   radius=20
    3 sensors:             x=800, y=CENTER_Y + [-40, 0, 40],  radius=5
    Front pinball:         center=(600, CENTER_Y),           radius=10
    Bottom pinball:        center=(626, CENTER_Y-15),        radius=10
    Top pinball:           center=(626, CENTER_Y+15),        radius=10

Usage::

    from scenes.karman_cloak.re100_scene import KarmanRe100Scene

    scene = KarmanRe100Scene(device_id=0)
    
    # Record target
    scene.create_target_env()
    scene.record_target("output/target.npz")
    
    # Build full env with pinball
    scene.create_full_env()
    scene.collect_norm("output/norm.json")
    scene.build_checkpoints("output/")
    
    # Inference
    scene.load_steady("output/checkpoint_steady.h5")
    results = scene.run_controlled(model, n_steps=200)
    scene.export_fields("output/fields/", step_indices=[0, 50, 100, 200])
"""

from __future__ import annotations

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

import numpy as np

# Add src directory to sys.path for package imports
_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)

from drl_pinball.cfd.pinball_env import PinballEnv

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
U0 = 0.01
L0 = 20.0
SAMPLE_INTERVAL = 800
FIFO_LEN = 150
CONV_LEN = 30
S_DIM = 12
A_DIM = 3
ACTION_SCALE = 8.0
ACTION_BIAS = (0.0, -4.0, 4.0)  # front, bottom, top
DATA_TYPE = np.float32

# Pinball config for new API
NEW_LBM_CONFIG = os.path.join(_REPO, "configs", "config_lbm_pinball.json")


class KarmanRe100Scene:
    """Karman cloak re100 scene manager."""
    
    def __init__(self, device_id: int = 0, viscosity: float = 0.004):
        self.device_id = device_id
        self.viscosity = viscosity
        self.env: Optional[PinballEnv] = None
        
        # Body IDs (set during create_target_env / create_full_env)
        self.dist_cyl_id: Optional[int] = None
        self.sensor_ids: List[int] = []
        self.front_cyl_id: Optional[int] = None
        self.bottom_cyl_id: Optional[int] = None
        self.top_cyl_id: Optional[int] = None
        
        # Recorded data
        self.target_states: Optional[np.ndarray] = None
        self.norm_data: Optional[Dict] = None
        self.save_states: Optional[np.ndarray] = None
    
    def _center_y(self) -> float:
        """Return the center y of the domain (NY-1)/2 = 255.5."""
        return 255.5  # (512 - 1) / 2
    
    # -------------------------------------------------------------------
    # Phase 1: Target recording (disturbance cylinder + sensors only)
    # -------------------------------------------------------------------
    
    def create_target_env(self):
        """Create flow field with disturbance cylinder + 3 sensors (no pinball)."""
        self.env = PinballEnv(
            lbm_config_path=NEW_LBM_CONFIG,
            body_config_path=None,
            device_id=self.device_id,
            viscosity=self.viscosity,
            velocity=U0,
        )
        
        cy = self._center_y()
        
        # Disturbance cylinder (upstream)
        self.dist_cyl_id = self.env.add_disturbance_cylinder(
            center=(10.0 * L0, cy), radius=L0
        )
        
        # 3 sensors at x=40*L0
        for y_off in [2.0, 0.0, -2.0]:
            sid = self.env.add_sensor(
                center=(40.0 * L0, cy + y_off * L0), radius=L0 / 4.0
            )
            self.sensor_ids.append(sid)
        
        # Rebuild
        self.env.reinitialize()
    
    def record_target(self, out_dir: str) -> str:
        """Record target sensor signals (disturbance only, no pinball).
        
        Saves to {out_dir}/target.npz and returns the path.
        """
        if self.env is None:
            raise RuntimeError("Call create_target_env() first")
        
        os.makedirs(out_dir, exist_ok=True)
        out_path = os.path.join(out_dir, "target.npz")
        
        # Stabilize
        stabilize_steps = int(4 * 1280 / U0)
        self.env.run_and_read(stabilize_steps)
        
        # Record target
        target_list = []
        for _ in range(FIFO_LEN):
            result = self.env.run_and_read(SAMPLE_INTERVAL)
            sens_flat = self.env.get_sensor_array(result["sensors"])
            target_list.append(sens_flat)
        
        self.target_states = np.array(target_list, dtype=DATA_TYPE)
        np.savez(out_path, target_states=self.target_states)
        return out_path
    
    # -------------------------------------------------------------------
    # Phase 2: Full env (add pinball, compute norm, checkpoint)
    # -------------------------------------------------------------------
    
    def create_full_env(self):
        """Add pinball cylinders to existing target env (or create from scratch)."""
        if self.env is None:
            self.create_target_env()
        
        cy = self._center_y()
        
        # Front cylinder
        self.front_cyl_id = self.env.add_cylinder(
            center=(30.0 * L0, cy), radius=L0 / 2.0
        )
        # Bottom cylinder
        self.bottom_cyl_id = self.env.add_cylinder(
            center=(31.3 * L0, cy - 0.75 * L0), radius=L0 / 2.0
        )
        # Top cylinder
        self.top_cyl_id = self.env.add_cylinder(
            center=(31.3 * L0, cy + 0.75 * L0), radius=L0 / 2.0
        )
        
        self.env.reinitialize()
    
    def collect_norm(self, out_dir: str) -> Dict:
        """Compute normalisation factors from zero-action rollout.
        
        Saves to {out_dir}/norm.json.
        Returns the norm dict.
        """
        if self.env is None:
            raise RuntimeError("Call create_full_env() first")
        
        os.makedirs(out_dir, exist_ok=True)
        
        cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
        
        # Stabilize with pinball
        stabilize_steps = int(4 * 1280 / U0)
        self.env.run_and_read(stabilize_steps)
        
        # Snapshot the steady state
        self.env.snapshot()
        
        # Zero-action rollout for norm
        fifo = deque(maxlen=FIFO_LEN)
        for _ in range(FIFO_LEN):
            result = self.env.run_and_read(SAMPLE_INTERVAL, read_fields=False)
            
            # Forces in legacy order: dist, front, bottom, top
            force_arr = []
            force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
            for cid in cylinder_ids:
                force_arr.extend(result["forces"].get(cid, [0, 0]))
            
            sensors_flat = self.env.get_sensor_array(result["sensors"])
            obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
            fifo.append(obs_flat)
        
        # Compute norm from fifo
        temp_states = np.array(fifo, dtype=DATA_TYPE)
        force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, 6:12])))
        sens_deviation = np.mean(temp_states[:, 0:6], axis=0)
        sens_norm_fact = np.zeros(6, dtype=DATA_TYPE)
        for i in range(6):
            sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))
        
        self.norm_data = {
            "force_norm_fact": force_norm_fact,
            "sens_deviation": sens_deviation.tolist(),
            "sens_norm_fact": sens_norm_fact.tolist(),
            "action_bias": list(ACTION_BIAS),
            "action_scale": ACTION_SCALE,
        }
        
        with open(os.path.join(out_dir, "norm.json"), "w") as f:
            json.dump(self.norm_data, f, indent=2)
        
        # Bias-action rollout for save_states
        self.env.restore()
        bias_omegas = {
            self.front_cyl_id: float(ACTION_BIAS[0] * U0),
            self.bottom_cyl_id: float(ACTION_BIAS[1] * U0),
            self.top_cyl_id: float(ACTION_BIAS[2] * U0),
        }
        
        fifo.clear()
        for _ in range(FIFO_LEN):
            result = self.env.run_and_read(SAMPLE_INTERVAL, omegas=bias_omegas)
            
            force_arr = []
            force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
            for cid in cylinder_ids:
                force_arr.extend(result["forces"].get(cid, [0, 0]))
            
            sensors_flat = self.env.get_sensor_array(result["sensors"])
            obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
            fifo.append(obs_flat)
        
        self.save_states = np.array(fifo, dtype=DATA_TYPE)
        np.savez(os.path.join(out_dir, "save_states.npz"), save_states=self.save_states)
        
        # Save norm data as NPZ for easy loading
        np.savez(
            os.path.join(out_dir, "norm_data.npz"),
            force_norm_fact=np.array([force_norm_fact], dtype=np.float32),
            sens_deviation=np.array(sens_deviation, dtype=np.float32),
            sens_norm_fact=np.array(sens_norm_fact, dtype=np.float32),
        )
        
        # Restore steady state
        self.env.restore()
        
        return self.norm_data
    
    def build_checkpoints(self, out_dir: str):
        """Save steady-state and bias-state checkpoints."""
        os.makedirs(out_dir, exist_ok=True)
        
        # Steady state (already snapshot'd in collect_norm)
        self.env.snapshot()
        steady_path = os.path.join(out_dir, "checkpoint_steady.h5")
        self.env.save_checkpoint(steady_path)
        
        # Bias state: restore + run bias + save
        self.env.restore()
        cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
        bias_omegas = {
            self.front_cyl_id: float(ACTION_BIAS[0] * U0),
            self.bottom_cyl_id: float(ACTION_BIAS[1] * U0),
            self.top_cyl_id: float(ACTION_BIAS[2] * U0),
        }
        self.env.run_and_read(SAMPLE_INTERVAL * 10, omegas=bias_omegas)
        bias_path = os.path.join(out_dir, "checkpoint_bias.h5")
        self.env.save_checkpoint(bias_path)
        
        # Restore steady
        self.env.restore()
    
    # -------------------------------------------------------------------
    # Inference
    # -------------------------------------------------------------------
    
    def load_checkpoint(self, path: str):
        """Load from a saved checkpoint."""
        if self.env is None:
            raise RuntimeError("Env not created. Call create_full_env() first.")
        self.env.load_checkpoint(path)
    
    def run_uncontrolled(self, n_steps: int, out_dir: str) -> Dict:
        """Run uncontrolled inference (zero action)."""
        os.makedirs(out_dir, exist_ok=True)
        
        cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
        
        sens_list, forc_list = [], []
        
        for _ in range(n_steps):
            result = self.env.run_and_read(SAMPLE_INTERVAL)
            
            force_arr = []
            force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
            for cid in cylinder_ids:
                force_arr.extend(result["forces"].get(cid, [0, 0]))
            
            sensors_flat = self.env.get_sensor_array(result["sensors"])
            
            sens_list.append(sensors_flat)
            forc_list.append(np.array(force_arr, dtype=np.float32))
        
        out = {
            "sensors": np.array(sens_list, dtype=np.float32),
            "forces": np.array(forc_list, dtype=np.float32),
        }
        np.savez(os.path.join(out_dir, "uncontrolled.npz"), **out)
        
        # Final vorticity
        self.env.export_vorticity_png(
            os.path.join(out_dir, "vorticity_uncontrolled.png"),
            title="Karman re100 uncontrolled",
        )
        
        return out
    
    def run_controlled(
        self,
        model: Any,
        n_steps: int,
        out_dir: str,
        *,
        field_steps: Optional[List[int]] = None,
    ) -> Dict:
        """Run controlled inference with a PPO model.
        
        Parameters
        ----------
        model : PPO
            Trained PPO model (must have Sin activation).
        n_steps : int
            Number of inference steps.
        out_dir : str
            Output directory.
        field_steps : list of int, optional
            Step indices at which to save Tecplot field files.
        
        Returns
        -------
        dict with sensors, forces, obs, actions, rewards.
        """
        os.makedirs(out_dir, exist_ok=True)
        if field_steps is not None:
            os.makedirs(os.path.join(out_dir, "fields"), exist_ok=True)
        
        cylinder_ids = [self.front_cyl_id, self.bottom_cyl_id, self.top_cyl_id]
        cyl_map = {
            self.front_cyl_id: 0,
            self.bottom_cyl_id: 1,
            self.top_cyl_id: 2,
        }
        
        norm = self.norm_data
        if norm is None:
            raise RuntimeError("Call collect_norm() first")
        
        force_norm_fact = float(norm["force_norm_fact"])
        sens_deviation = np.array(norm["sens_deviation"], dtype=np.float32)
        sens_norm_fact = np.array(norm["sens_norm_fact"], dtype=np.float32)
        
        # Restore steady state
        self.env.restore()
        
        # Bias FIFO
        fifo = deque(maxlen=FIFO_LEN)
        bias_omegas = {
            self.front_cyl_id: float(ACTION_BIAS[0] * U0),
            self.bottom_cyl_id: float(ACTION_BIAS[1] * U0),
            self.top_cyl_id: float(ACTION_BIAS[2] * U0),
        }
        
        for _ in range(FIFO_LEN):
            result = self.env.run_and_read(SAMPLE_INTERVAL, omegas=bias_omegas)
            
            force_arr = []
            force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
            for cid in cylinder_ids:
                force_arr.extend(result["forces"].get(cid, [0, 0]))
            
            sensors_flat = self.env.get_sensor_array(result["sensors"])
            obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
            fifo.append(obs_flat)
        
        sens_list, forc_list, obs_list = [], [], []
        action_list, reward_list = [], []
        reward_cd_list, reward_cl_list, reward_sim_list = [], [], []
        
        obs = np.zeros(S_DIM, dtype=np.float32)
        
        for step in range(n_steps):
            # PPO action
            action, _states = model.predict(obs, deterministic=True)
            action = action.astype(np.float32).flatten()
            action_list.append(action.copy())
            
            # Convert to omegas
            omegas = {}
            for i, cid in enumerate(cylinder_ids):
                omega_val = (action[i] * ACTION_SCALE + ACTION_BIAS[i]) * U0
                omegas[cid] = float(omega_val)
            
            # Run CFD
            result = self.env.run_and_read(SAMPLE_INTERVAL, omegas=omegas)
            
            # Build obs slice
            force_arr = []
            force_arr.extend(result["forces"].get(self.dist_cyl_id, [0, 0]))
            for cid in cylinder_ids:
                force_arr.extend(result["forces"].get(cid, [0, 0]))
            
            sensors_flat = self.env.get_sensor_array(result["sensors"])
            obs_flat = np.concatenate([sensors_flat, np.array(force_arr, dtype=np.float32)])
            fifo.append(obs_flat)
            
            sens_list.append(sensors_flat)
            forc_list.append(np.array(force_arr, dtype=np.float32))
            
            # Build normalised observation
            forces_norm = np.array(force_arr[2:], dtype=np.float32) / force_norm_fact  # skip dist cy forces
            sens_norm = (sensors_flat - sens_deviation) / sens_norm_fact
            obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)
            obs_list.append(obs)
            
            # Compute reward
            states_arr = np.array(fifo, dtype=np.float32)
            if len(states_arr) >= CONV_LEN:
                forces = states_arr[-1, 6:12] / force_norm_fact
                cd = float((forces[0] + forces[2] + forces[4]) / 3.0)
                cl = float((forces[1] + forces[3] + forces[5]) / 3.0)
                
                sim = self._compute_similarity(states_arr)
                
                r_cd = float(np.exp(-abs(cd * 20.0)))
                r_cl = float(np.exp(-abs(cl * 80.0)))
                r_sim = float(np.exp(-10.0 * abs(sim - 1.0)))
                reward = float(min(0.3 * r_cd + 0.4 * r_cl + 0.3 * r_sim, 1.0))
            else:
                reward = 0.0
                r_cd = r_cl = r_sim = 0.0
            
            reward_list.append(reward)
            reward_cd_list.append(r_cd)
            reward_cl_list.append(r_cl)
            reward_sim_list.append(r_sim)
            
            # Field export
            if field_steps is not None and step in field_steps:
                fname = os.path.join(out_dir, "fields", f"field_{step:06d}.dat")
                self.env.save_field_tecplot(fname)
        
        out = {
            "sensors": np.array(sens_list, dtype=np.float32),
            "forces": np.array(forc_list, dtype=np.float32),
            "obs": np.array(obs_list, dtype=np.float32),
            "actions": np.array(action_list, dtype=np.float32),
            "rewards": np.array(reward_list, dtype=np.float32),
            "reward_cd": np.array(reward_cd_list, dtype=np.float32),
            "reward_cl": np.array(reward_cl_list, dtype=np.float32),
            "reward_sim": np.array(reward_sim_list, dtype=np.float32),
        }
        np.savez(os.path.join(out_dir, "controlled.npz"), **out)
        
        # Final vorticity
        self.env.export_vorticity_png(
            os.path.join(out_dir, "vorticity_controlled.png"),
            title="Karman re100 controlled (PPO)",
        )
        
        return out
    
    def _compute_similarity(self, states_arr: np.ndarray) -> float:
        """Compute lag-compensated DTW similarity (matches legacy env logic)."""
        if self.target_states is None:
            return 0.0
        
        target = self.target_states
        
        # Lag from middle sensor (index 1 = sensor1_uy in sensor[6] block)
        ref = target[CONV_LEN:2 * CONV_LEN, 1]
        cur = states_arr[-CONV_LEN:, 1]
        lag = self._calc_lag(ref, cur)
        
        sim_sum = 0.0
        for i in range(6):
            t_seq = np.roll(target[:, i], -lag)[CONV_LEN:2 * CONV_LEN]
            s_seq = states_arr[-CONV_LEN:, i]
            sim_sum += self._calc_dtw_sim(t_seq, s_seq) / 6.0
        
        return float(sim_sum)
    
    @staticmethod
    def _calc_lag(target: np.ndarray, state: np.ndarray) -> int:
        tm = np.mean(target)
        sm = np.mean(state)
        corr = np.correlate(target - tm, state - sm, mode="full")
        lags = np.arange(-len(target) + 1, len(target))
        return int(lags[np.argmax(corr)])
    
    @staticmethod
    def _calc_dtw_sim(target: np.ndarray, state: np.ndarray) -> float:
        n, m = len(target), len(state)
        dtw = np.full((n + 1, m + 1), np.inf)
        dtw[0, 0] = 0.0
        for i in range(1, n + 1):
            for j in range(1, m + 1):
                cost = abs(float(target[i - 1]) - float(state[j - 1]))
                dtw[i, j] = cost + min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
        return float(1.0 - dtw[n, m] / n)
    
    def close(self):
        if self.env is not None:
            self.env.close()
            self.env = None
