#!/usr/bin/env python3
"""Karman Cloak environment (V5 — parameterized, no-bias only).

Based on env_karman_2000x600.py V4. Accepts calibration dict and config path.
All calibration constants loaded from calibration.json produced by calibrate.py.

Design: Two-phase initialization to AVOID runtime sync_bodies().
  Phase 1: Temporary Simulation(dist + sensors) -> record target -> close
  Phase 2: Training Simulation(all objects upfront) -> warmup -> zero-action FIFO -> snapshot

CUDA context: mirrors legacy pattern — push CFD context before GPU ops, pop after.

Observation (12-dim, physical norm, NO clip — VecNormalize handles that):
  [0:6]  = raw_forces / FORCE_SCALE   (front_fx,fy, top_fx,fy, bot_fx,fy)
  [6:12] = raw_sensors / SENS_SCALE   (s0_ux,uy, s1_ux,uy, s2_ux,uy)

Action (3-dim): no_bias only
  [-1,1] -> omega = -(action * ACTION_SCALE + [0,0,0]) * U0 / R

Reward (V3: Gaussian + EMA smoothing + normalized DTW):
  r_cd  = EMA(exp(-cd_norm^2 * K_CD), EMA_FAST)
  r_cl  = EMA(exp(-cl_norm^2 * K_CL), EMA_FAST)
  r_sim = piecewise_map(sim, SIM_BP, SIM_VAL)  -> [0, 1]
  reward = W_CD*r_cd + W_CL*r_cl + W_SIM*r_sim - floor_penalty
"""
from __future__ import annotations

import json
import os, sys, time
from collections import deque
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import gymnasium as gym
from gymnasium import spaces

_REPO = Path(__file__).resolve().parents[3]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

from CelerisLab import Simulation

_CELERIS = Path("/home/frank14f/CelerisLab")
_CONFIG_H = _CELERIS / "src/CelerisLab/lbm/kernels/config/config_objects.h"
_PTX = _CELERIS / "src/CelerisLab/lbm/kernels/kernel.ptx"

def _clean_cache():
    for p in [_CONFIG_H, _PTX]:
        if p.exists(): p.unlink()

# ---------------------------------------------------------------------------
# Geometry constants (fixed across all Karman cloak cases)
# ---------------------------------------------------------------------------
L0 = 20.0; D_CYL = L0; U0 = 0.01; RADIUS = L0 / 2.0
NX = 2000; NY = 600
CENTER_Y = float(NY - 1) / 2.0

DIST_X = 600.0
PINBALL_FRONT_X = 1000.0
PINBALL_REAR_X = 1026.0
SENSOR_X = 1200.0

FIFO_LEN = 150; CONV_LEN = 30; MAX_STEPS = 500
EMA_FAST = 0.2
S_DIM = 12; A_DIM = 3

SENSOR_CC = 78.0


# ---------------------------------------------------------------------------
# DTW utilities (identical to env_karman_2000x600.py)
# ---------------------------------------------------------------------------
def calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    t_mean = np.mean(target); s_mean = np.mean(state)
    corr = np.correlate(target - t_mean, state - s_mean, mode="full")
    lags = np.arange(-len(target) + 1, len(target))
    return int(lags[np.argmax(corr)])


def calc_dtw_sim(target: np.ndarray, state: np.ndarray,
                 norm_scale: float = 1.0) -> 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]))
            last_min = min(dtw[i - 1, j], dtw[i, j - 1], dtw[i - 1, j - 1])
            dtw[i, j] = cost + last_min
    raw = 1.0 - dtw[n, m] / (float(n) * norm_scale)
    return float(max(0.0, raw))


def compute_similarity(target_states, fifo_states, conv_len, norm_scale):
    target = np.asarray(target_states, dtype=np.float64)
    state = np.asarray(fifo_states, dtype=np.float64)
    id_sens = 3
    target_seq = target[conv_len:2 * conv_len, id_sens]
    state_seq = state[-conv_len:, id_sens]
    lag = calc_lag(target_seq, state_seq)
    sim = 0.0
    for i in range(6):
        t_seq = np.roll(target[:, i], -lag)[conv_len:2 * conv_len]
        s_seq = state[-conv_len:, i]
        sim += calc_dtw_sim(t_seq, s_seq, norm_scale=norm_scale)
    return float(sim / 6.0)


# ---------------------------------------------------------------------------
class ActionSmoother:
    def __init__(self, weight: float = 0.1):
        self.weight = weight; self._state: Optional[np.ndarray] = None
    def __call__(self, target: np.ndarray) -> np.ndarray:
        t = np.asarray(target, dtype=np.float32)
        if self._state is None:
            self._state = t.copy()
        else:
            self._state = (1.0 - self.weight) * self._state + self.weight * t
        return self._state.copy()
    def reset(self, value: Optional[np.ndarray] = None) -> None:
        self._state = np.asarray(value, dtype=np.float32).copy() if value is not None else None


# ---------------------------------------------------------------------------
def record_target(config_path: str, device_id: int, si: int) -> np.ndarray:
    """Record target signal (dist_cyl + sensors, no pinball)."""
    _clean_cache()
    warmup = int(4.0 * NX / U0)
    sim = Simulation(lbm_config_path=config_path, device_id=device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0)
    s0 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0)
    s1 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0)
    s2 = sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0)
    sim.initialize()
    sim.run(warmup, zero_obs=True)
    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(si, zero_obs=True)
        target[i] = [
            sim.read_sensor(s0, normalize=True)[0] * SENSOR_CC,
            sim.read_sensor(s0, normalize=True)[1] * SENSOR_CC,
            sim.read_sensor(s1, normalize=True)[0] * SENSOR_CC,
            sim.read_sensor(s1, normalize=True)[1] * SENSOR_CC,
            sim.read_sensor(s2, normalize=True)[0] * SENSOR_CC,
            sim.read_sensor(s2, normalize=True)[1] * SENSOR_CC,
        ]
    sim.close()
    return np.array(target, dtype=np.float32)


# ---------------------------------------------------------------------------
class KarmanCloakEnv(gym.Env):
    """Parameterized Karman Cloak environment (V5 — no_bias only).

    Parameters
    ----------
    device_id : int
        GPU device ID.
    seed : int
        Random seed.
    calibration : dict
        Calibration dict loaded from calibration.json.
        Must contain: FORCE_SCALE, SENS_SCALE, dtw_norm_scale,
        SIM_BP, SIM_VAL, K_CD, K_CL, W_CD, W_CL, W_SIM,
        FLOOR_CD, FLOOR_CL, FLOOR_SIM, FLOOR_PENALTY,
        ACTION_SCALE, ACTION_BIAS, SI.
    config_path : str
        Path to LBM config JSON.
    target_states : np.ndarray, optional
        Pre-recorded target signal. If None, recorded on-the-fly.
    """
    metadata = {"render_modes": ["human"]}

    def __init__(self, device_id: int = 0, seed: int = 42,
                 calibration: Optional[dict] = None,
                 config_path: Optional[str] = None,
                 target_states: Optional[np.ndarray] = None):
        super().__init__()
        self.device_id = device_id
        self.seed = seed
        np.random.seed(seed)

        # Load calibration
        if calibration is None:
            raise ValueError("calibration dict is required for V5 KarmanCloakEnv")
        self._cal = calibration.copy()
        self._si = int(self._cal["SI"])
        self._force_scale = np.float32(self._cal["FORCE_SCALE"])
        self._sens_scale = np.float32(self._cal["SENS_SCALE"])
        self._dtw_norm_scale = float(self._cal["dtw_norm_scale"])
        self._sim_bp = np.array(self._cal["SIM_BP"], dtype=np.float64)
        self._sim_val = np.array(self._cal["SIM_VAL"], dtype=np.float64)
        self._k_cd = float(self._cal["K_CD"])
        self._k_cl = float(self._cal["K_CL"])
        self._w_cd = float(self._cal["W_CD"])
        self._w_cl = float(self._cal["W_CL"])
        self._w_sim = float(self._cal["W_SIM"])
        self._floor_cd = float(self._cal["FLOOR_CD"])
        self._floor_cl = float(self._cal["FLOOR_CL"])
        self._floor_sim = float(self._cal["FLOOR_SIM"])
        self._floor_pen = float(self._cal["FLOOR_PENALTY"])
        self._action_scale = float(self._cal["ACTION_SCALE"])
        self._action_bias = np.array(self._cal["ACTION_BIAS"], dtype=np.float32)

        self._config_path = config_path or self._cal.get("config_path")
        if self._config_path is None:
            raise ValueError("config_path is required")

        self.action_space = spaces.Box(-1.0, 1.0, (A_DIM,), dtype=np.float32)
        self.observation_space = spaces.Box(-10.0, 10.0, (S_DIM,), dtype=np.float32)

        self.fifo_states = deque(maxlen=FIFO_LEN)
        self.save_states: np.ndarray = None
        self.target_states: np.ndarray = None
        self.current_step = 0
        self.smoother = ActionSmoother(weight=0.1)
        self._ema_r_cd = 0.0
        self._ema_r_cl = 0.0
        self.sim = None
        self.dist_id = None
        self.sensor_ids = []
        self.pinball_ids = []

        if target_states is not None:
            self.target_states = target_states
        self._init_cfd()

    def _gpu_block(self, fn):
        if self.sim is not None:
            self.sim.ctx._ctx.push()
        try:
            fn()
        finally:
            if self.sim is not None:
                self.sim.ctx._ctx.pop()

    def _init_cfd(self):
        t0 = time.perf_counter()
        warmup = int(4.0 * NX / U0)

        if self.target_states is None:
            self.target_states = record_target(self._config_path, self.device_id, self._si)

        _clean_cache()
        self.sim = Simulation(lbm_config_path=self._config_path, device_id=self.device_id)
        self.sim._assert_object_count_contract = lambda *a, **kw: None

        self.dist_id = self.sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0)
        self.sensor_ids = [
            self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
            self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
            self.sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
        ]
        self.sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
        self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
        self.sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
        self.sim.initialize()
        self.pinball_ids = [4, 5, 6]

        print(f"  [env] Warmup ({warmup} steps)...", end=" ", flush=True)
        self._gpu_block(lambda: self.sim.run(warmup, zero_obs=True))
        print(f"done ({time.perf_counter()-t0:.0f}s).")

        # Zero-action FIFO: no rotation, just let pinball oscillate naturally
        print(f"  [env] Zero-action FIFO ({FIFO_LEN})...", end=" ", flush=True)
        zero_omega = self._action_to_omega(np.zeros(3, dtype=np.float32))
        self.smoother.reset(zero_omega.copy())
        self._gpu_block(lambda: [self.sim.run(self._si, zero_obs=True) for _ in range(FIFO_LEN)])
        f_diag = self._read_obs()
        print(f"max|force|={np.max(np.abs(f_diag[6:12])):.6f}")

        print(f"  [env] Saving snapshot after zero-action FIFO...", end=" ", flush=True)
        fifo_save = []
        for _ in range(FIFO_LEN):
            self._set_omega(zero_omega)
            self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True))
            obs = self._read_obs()
            sl = obs[2:14].copy()
            sl[0:6] *= SENSOR_CC
            fifo_save.append(sl)
        self.save_states = np.array(fifo_save, dtype=np.float32)
        print("done.")

        self._gpu_block(lambda: self.sim.snapshot())
        print(f"  [env] Init done ({time.perf_counter()-t0:.0f}s)")

    def _read_obs(self) -> np.ndarray:
        obs = list(self.sim.read_force(self.dist_id, normalize=True))
        for sid in self.sensor_ids:
            s = self.sim.read_sensor(sid, normalize=True)
            obs.extend([float(s[0]), float(s[1])])
        for pid in self.pinball_ids:
            obs.extend(self.sim.read_force(pid, normalize=True))
        return np.array(obs, dtype=np.float32)

    def _action_to_omega(self, action_norm: np.ndarray) -> np.ndarray:
        sv = (np.asarray(action_norm, dtype=np.float32) * self._action_scale + self._action_bias) * U0
        return -sv / RADIUS

    def _set_omega(self, omega: np.ndarray):
        for pid, w in zip(self.pinball_ids, omega):
            self.sim.set_body(pid, omega=float(w))

    def _normalize_obs(self, raw_obs_slice: np.ndarray) -> np.ndarray:
        forces = raw_obs_slice[6:12] / self._force_scale
        sens = raw_obs_slice[0:6] / self._sens_scale
        return np.hstack([forces, sens]).astype(np.float32)

    def _compute_reward(self, obs_slice: np.ndarray) -> Tuple[float, dict]:
        forces_raw = obs_slice[6:12]
        cd_raw = (forces_raw[0] + forces_raw[2] + forces_raw[4]) / 3.0
        cl_raw = (forces_raw[1] + forces_raw[3] + forces_raw[5]) / 3.0
        cd_norm = cd_raw / self._force_scale
        cl_norm = cl_raw / self._force_scale

        sim_val = 0.0
        if len(self.fifo_states) >= CONV_LEN * 2:
            sim_val = compute_similarity(self.target_states,
                                          np.array(list(self.fifo_states)),
                                          conv_len=CONV_LEN,
                                          norm_scale=self._dtw_norm_scale)

        r_cd_raw = float(np.exp(-cd_norm**2 * self._k_cd))
        r_cl_raw = float(np.exp(-cl_norm**2 * self._k_cl))

        self._ema_r_cd = (1 - EMA_FAST) * self._ema_r_cd + EMA_FAST * r_cd_raw
        self._ema_r_cl = (1 - EMA_FAST) * self._ema_r_cl + EMA_FAST * r_cl_raw
        r_sim = float(np.interp(sim_val, self._sim_bp, self._sim_val))

        reward = self._w_cd * self._ema_r_cd + self._w_cl * self._ema_r_cl + self._w_sim * r_sim

        floor_pen = 0.0
        if self._ema_r_cd < self._floor_cd:
            floor_pen += self._floor_pen * (self._floor_cd - self._ema_r_cd) / self._floor_cd
        if self._ema_r_cl < self._floor_cl:
            floor_pen += self._floor_pen * (self._floor_cl - self._ema_r_cl) / self._floor_cl
        if r_sim < self._floor_sim:
            floor_pen += self._floor_pen * (self._floor_sim - r_sim) / self._floor_sim
        reward = max(0.0, reward - floor_pen)

        info = {"cd": float(cd_norm), "cl": float(cl_norm), "sim": float(sim_val),
                "r_cd": self._ema_r_cd, "r_cl": self._ema_r_cl, "r_sim": r_sim,
                "floor_pen": float(floor_pen)}
        return float(reward), info

    def reset(self, seed=None, options=None) -> Tuple[np.ndarray, dict]:
        super().reset(seed=seed)
        self._gpu_block(lambda: self.sim.restore())
        self.smoother.reset(self._action_to_omega(np.zeros(3, dtype=np.float32)))
        self.fifo_states.clear()
        for i in range(len(self.save_states)):
            self.fifo_states.append(self.save_states[i, 0:6])
        self.current_step = 0
        self._ema_r_cd = 0.0
        self._ema_r_cl = 0.0
        obs_raw = self._read_obs()
        obs = self._normalize_obs(obs_raw[2:14])
        return obs, {}

    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, dict]:
        assert self.action_space.contains(action), f"Invalid action: {action}"
        target_omega = self._action_to_omega(action)
        smoothed = self.smoother(target_omega)
        self._set_omega(smoothed)

        self._gpu_block(lambda: self.sim.run(self._si, zero_obs=True))

        obs_raw = self._read_obs()
        obs_slice = obs_raw[2:14]
        obs = self._normalize_obs(obs_slice)
        self.fifo_states.append(obs_slice[0:6] * SENSOR_CC)
        reward, info = self._compute_reward(obs_slice)

        self.current_step += 1
        terminated = False
        return obs, reward, terminated, False, info

    def render(self, mode="human"):
        pass

    def close(self):
        if self.sim is not None:
            self.sim.close()


# ---------------------------------------------------------------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--device-id", type=int, default=0)
    parser.add_argument("--calibration", type=str, required=True,
                        help="Path to calibration.json")
    parser.add_argument("--config", type=str, required=True,
                        help="Path to LBM config JSON")
    args = parser.parse_args()

    with open(args.calibration, "r") as f:
        cal = json.load(f)

    print("=== KarmanCloakEnv V5 Quick Test ===")
    env = KarmanCloakEnv(device_id=args.device_id, calibration=cal,
                          config_path=args.config)

    obs, _ = env.reset()
    print(f"  Init obs: min={obs.min():.4f}, max={obs.max():.4f}, mean={obs.mean():.4f}")

    rewards = []
    for step in range(50):
        obs, reward, *_ = env.step(np.zeros(3, dtype=np.float32))
        rewards.append(reward)
    print(f"  Zero-action reward (last 20): {np.mean(rewards[-20:]):.4f}")

    obs1, _ = env.reset()
    obs2, _ = env.reset()
    print(f"  Reset consistency: {np.max(np.abs(obs1-obs2)):.8f}")

    env.close()
    print("=== Done ===")