#!/usr/bin/env python3
"""Karman Cloak environment for 2000x600 config (uniform, free-slip).

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

CUDA context: mirrors legacy pattern — push CFD context before GPU ops, pop after.
  This allows PPO (PyTorch) and CFD (PyCUDA) to coexist on the same GPU.

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):
  [-1,1] -> omega = -(action*8 + [0,-4,4]) * U0 / R

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

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 = _REPO / "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()

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

SI = 800
FIFO_LEN = 150; CONV_LEN = 30; MAX_STEPS = 500
ACTION_SCALE = 8.0
ACTION_BIAS = np.array([0.0, -4.0, 4.0], dtype=np.float32)

FORCE_SCALE = np.float32(0.0025)   # combined max|force| across Stage0+Bias (Phase 0)
SENS_SCALE = np.float32(0.8)       # combined max|sensor| legacy-equiv across stages

# Reward parameters V3 (Gaussian r_cd/r_cl + normalized DTW + raw r_sim)
K_CD = 50.0          # Gaussian: exp(-cd_norm^2 * K)
K_CL = 100.0         # Gaussian: exp(-cl_norm^2 * K)
EMA_FAST = 0.2       # EMA weight for r_cd/r_cl smoothing (r_sim NOT smoothed)
# Normalized DTW sim mapping (based on Phase 0: Stage0~0.36, Stage1~0.82)
SIM_BP = np.array([0.0, 0.36, 0.82, 0.90, 0.95, 1.0], dtype=np.float64)
SIM_VAL = np.array([0.0, 0.2, 0.5, 0.8, 0.9, 1.0], dtype=np.float64)
W_CD = 0.30
W_CL = 0.30
W_SIM = 0.40

# Floor penalty: when a component drops below its floor, extra penalty is applied
# This prevents DRL from sacrificing one component to optimize others
FLOOR_CD = 0.10    # r_cd below 0.10 gets penalized
FLOOR_CL = 0.10    # r_cl below 0.10 gets penalized
FLOOR_SIM = 0.10   # r_sim below 0.10 gets penalized
FLOOR_PENALTY = 0.05  # penalty per component below floor (subtracted from reward)

# NoBias action scale (wider range to cover Bias omega range)
ACTION_SCALE_NOBIAS = 12.0

WARMUP_STEPS = int(4.0 * NX / U0)
CFG_PATH = str(_REPO / "configs" / "config_lbm_karman_2000x600.json")
SENSOR_CC = 78.0

# ---------------------------------------------------------------------------
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:
    """DTW similarity with optional amplitude normalization.

    norm_scale: fixed reference amplitude (e.g. target's uy-avg std).
    sim = max(0, 1 - total_cost / (n * norm_scale))
    Without normalization (norm_scale=1.0), this is the raw formula.
    """
    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: np.ndarray, fifo_states: np.ndarray,
                        conv_len: int = CONV_LEN,
                        norm_scale: float = 1.0) -> float:
    """Compute DTW similarity across 6 sensor channels with fixed normalization."""
    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

# ---------------------------------------------------------------------------
S_DIM = 12; A_DIM = 3

def record_target(device_id: int) -> np.ndarray:
    _clean_cache()
    sim = Simulation(lbm_config_path=CFG_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_STEPS, 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):
    metadata = {"render_modes": ["human"]}

    def __init__(self, device_id: int = 0, seed: int = 42,
                 target_states: Optional[np.ndarray] = None,
                 action_bias: Optional[np.ndarray] = None,
                 action_scale: Optional[float] = None):
        super().__init__()
        self.device_id = device_id
        self.seed = seed
        np.random.seed(seed)

        # Configurable action bias and scale
        self.action_bias = action_bias if action_bias is not None else ACTION_BIAS.copy()
        self.action_scale = action_scale if action_scale is not None else ACTION_SCALE

        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)
        # EMA states for r_cd/r_cl smoothing only (r_sim used raw)
        self._ema_r_cd = 0.0
        self._ema_r_cl = 0.0
        # DTW normalization scale (computed from target uy-channel std)
        self._dtw_norm_scale = 1.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()

    # ---- Context guard (mirrors legacy context.push()/pop() pattern) ----
    def _ctx_guard(func):
        """Decorator: push CFD context before GPU ops, pop after."""
        def wrapper(self, *args, **kwargs):
            if self.sim is not None:
                self.sim.ctx._ctx.push()
            try:
                return func(self, *args, **kwargs)
            finally:
                if self.sim is not None:
                    self.sim.ctx._ctx.pop()
        return wrapper

    # ---- Init ----
    def _init_cfd(self):
        t0 = time.perf_counter()

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

        # Compute DTW normalization scale from target uy-channel std (fixed)
        target_std = np.std(self.target_states, axis=0)
        uy_std_avg = float(np.mean([target_std[1], target_std[3], target_std[5]]))
        self._dtw_norm_scale = max(uy_std_avg, 0.01)  # floor to avoid div-by-zero
        print(f"  [env] DTW norm scale (target uy-avg std): {self._dtw_norm_scale:.4f}")

        _clean_cache()
        self.sim = Simulation(lbm_config_path=CFG_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} steps)...", end=" ", flush=True)
        self._gpu_block(lambda: self.sim.run(WARMUP_STEPS, zero_obs=True))
        print(f"done ({time.perf_counter()-t0:.0f}s).")

        print(f"  [env] Zero-action FIFO ({FIFO_LEN})...", end=" ", flush=True)
        self._gpu_block(lambda: [self.sim.run(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}")

        bias_omega = self._action_to_omega(np.zeros(3, dtype=np.float32))
        self.smoother.reset(bias_omega.copy())
        print(f"  [env] Bias FIFO ({FIFO_LEN})...", end=" ", flush=True)
        fifo_save = []
        for _ in range(FIFO_LEN):
            self._set_omega(bias_omega)
            self._gpu_block(lambda: self.sim.run(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 _gpu_block(self, fn):
        """Execute fn with CFD context pushed, then pop."""
        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()

    # ---- Private ----
    def _read_obs(self) -> np.ndarray:
        """14-dim: [dist_fx,fy, 6×raw_sensor, 6×raw_force]."""
        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:
        """Physical normalization (no clip — VecNormalize handles that).

        obs_slice = [6 sensor (legacy-equiv after *CC), 6 force (raw)].
        Returns [forces(6)/FORCE_SCALE, sensors(6)/SENS_SCALE] in env order.
        """
        forces = raw_obs_slice[6:12] / FORCE_SCALE
        sens = raw_obs_slice[0:6] / SENS_SCALE
        return np.hstack([forces, sens]).astype(np.float32)

    def _compute_reward(self, obs_slice: np.ndarray) -> Tuple[float, dict]:
        """V3 reward: Gaussian r_cd/r_cl (EMA-smoothed) + raw normalized DTW r_sim.

        obs_slice = [6 sensor (legacy-equiv), 6 force (raw)].
        """
        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 / FORCE_SCALE
        cl_norm = cl_raw / FORCE_SCALE

        # Normalized DTW similarity (raw, no EMA — DTW is already smooth)
        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)

        # Gaussian reward for cd/cl (no zero-crossing spikes)
        r_cd_raw = float(np.exp(-cd_norm**2 * K_CD))
        r_cl_raw = float(np.exp(-cl_norm**2 * K_CL))

        # EMA smoothing only for cd/cl (r_sim used raw — DTW already smooth)
        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, SIM_BP, SIM_VAL))

        reward = W_CD * self._ema_r_cd + W_CL * self._ema_r_cl + W_SIM * r_sim

        # Floor penalty: discourage sacrificing any component below its floor
        floor_pen = 0.0
        if self._ema_r_cd < FLOOR_CD:
            floor_pen += FLOOR_PENALTY * (FLOOR_CD - self._ema_r_cd) / FLOOR_CD
        if self._ema_r_cl < FLOOR_CL:
            floor_pen += FLOOR_PENALTY * (FLOOR_CL - self._ema_r_cl) / FLOOR_CL
        if r_sim < FLOOR_SIM:
            floor_pen += FLOOR_PENALTY * (FLOOR_SIM - r_sim) / 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

    # ---- Gym interface ----
    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
        # Reset EMA states (r_cd/r_cl only)
        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(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
        # Never terminate — flow field stays continuous (legacy pattern).
        # Training loop controls episode length externally.
        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)
    args = parser.parse_args()

    print("=== KarmanCloakEnv 2000x600 Quick Test ===")
    env = KarmanCloakEnv(device_id=args.device_id)

    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"  Bias 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 ===")
