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

Design: Two-phase initialization to AVOID runtime sync_bodies().
  Phase 1: Temporary Simulation(target cyl + sensors) -> record 150-step signal
           + FFT harmonics -> close
  Phase 2: Training Simulation(3 sensors + 3 pinball upfront) -> warmup ->
           zero-action FIFO -> snapshot

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

Geometry: pinball@19/20.3L0, sensor@30L0, target cylinder@20L0 (1L diameter).
  Equivalent pinball-to-sensor distance matched to illusion target.

Observation (14-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)
  [12:14] = target_cd, target_cl       (from FFT harmonics reconstruction)

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

Reward (V3: Gaussian + EMA smoothing + normalized DTW):
  r_cd  = EMA(exp(-(cd - target_cd)^2 * K_CD), EMA_FAST)
  r_cl  = EMA(exp(-(cl - target_cl)^2 * K_CL), EMA_FAST)
  r_sim = piecewise_map(sim, SIM_BP, SIM_VAL)
  reward = W_CD*r_cd + W_CL*r_cl + W_SIM*r_sim - floor_penalty

K_CD/K_CL are LOWER than Karman (12/25 vs 50/100) because 3-pinball total force
is compared to 1-target-cylinder force, giving wider normalized error range.
"""
from __future__ import annotations

import json, 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
from normalization import dtw_sensor_factor, policy_sensor_scale

# ---------------------------------------------------------------------------
L0 = 20.0; U0 = 0.01; RADIUS = L0 / 2.0
NX = 2000; NY = 600
CENTER_Y = float(NY - 1) / 2.0

# Illusion geometry: pinball and sensors shifted left vs Karman
PINBALL_FRONT_X = 380.0    # 19*L0
PINBALL_REAR_X = 406.0     # 20.3*L0
SENSOR_X = 600.0           # 30*L0
TARGET_CYL_X = 400.0       # 20*L0 (target cylinder position during recording)
TARGET_CYL_RADIUS = 1.0 * L0  # 1L target diameter

FIFO_LEN = 150; CONV_LEN = 30; MAX_STEPS = 500
EMA_FAST = 0.2
S_DIM = 14; A_DIM = 3
ACTION_SCALE = 12.0
ACTION_BIAS = np.array([0.0, 0.0, 0.0], dtype=np.float32)

# ---------------------------------------------------------------------------
# DTW utilities
# ---------------------------------------------------------------------------
def calc_lag(target, state):
    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, state, norm_scale=1.0):
    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, fifo_arr, conv_len, norm_scale):
    target = np.asarray(target, dtype=np.float64)
    state = np.asarray(fifo_arr, dtype=np.float64)
    if len(state) < conv_len:
        return 0.0
    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=0.1):
        self.weight = weight; self._state = None
    def __call__(self, target):
        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=None):
        self._state = np.asarray(value, dtype=np.float32).copy() if value is not None else None


# ---------------------------------------------------------------------------
# Harmonics analysis
# ---------------------------------------------------------------------------
def analyze_harmonics(states, n_harmonics=5):
    N, D = states.shape
    result = []
    for d in range(D):
        y = states[:, d]
        fft_coef = np.fft.rfft(y)
        freqs = np.fft.rfftfreq(N, d=1)
        amps = 2 * np.abs(fft_coef) / N
        phases = np.angle(fft_coef)
        idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
        harmonics = {
            'dc': float(np.real(fft_coef[0]) / N),
            'amps': amps[idx].tolist(),
            'freqs': freqs[idx].tolist(),
            'phases': phases[idx].tolist(),
        }
        result.append(harmonics)
    return result


def gen_target_states_at(t, harmonics):
    t = np.asarray(t)
    D = len(harmonics)
    result = np.zeros((t.size, D), dtype=np.float32)
    for d, h in enumerate(harmonics):
        val = np.full(t.shape, h['dc'], dtype=np.float32)
        for amp, freq, phase in zip(h['amps'], h['freqs'], h['phases']):
            val += amp * np.cos(2 * np.pi * freq * t + phase)
        result[:, d] = val
    if result.shape[0] == 1:
        return result[0]
    return result


# ---------------------------------------------------------------------------
def record_illusion_target(config_path, device_id, si, target_diam=1.0, sensor_factor=1.0):
    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=(TARGET_CYL_X, CENTER_Y, 0.0),
                 radius=target_diam * 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_states = np.zeros((FIFO_LEN, 8), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(si, zero_obs=True)
        target_states[i] = [
            sim.read_sensor(s0, normalize=True)[0] * sensor_factor,
            sim.read_sensor(s0, normalize=True)[1] * sensor_factor,
            sim.read_sensor(s1, normalize=True)[0] * sensor_factor,
            sim.read_sensor(s1, normalize=True)[1] * sensor_factor,
            sim.read_sensor(s2, normalize=True)[0] * sensor_factor,
            sim.read_sensor(s2, normalize=True)[1] * sensor_factor,
            sim.read_force(0, normalize=True)[0],
            sim.read_force(0, normalize=True)[1],
        ]
    sim.close()
    harmonics = analyze_harmonics(target_states, n_harmonics=5)
    return target_states, harmonics


# ---------------------------------------------------------------------------
class IllusionCloakEnv(gym.Env):
    metadata = {"render_modes": ["human"]}

    def __init__(self, device_id=0, seed=42, calibration=None,
                 config_path=None, target_states=None, target_harmonics=None,
                 target_diam=1.0, policy_sensor_units=None):
        super().__init__()
        self.device_id = device_id
        self.seed = seed
        np.random.seed(seed)
        self._target_diam = target_diam

        if calibration is None:
            raise ValueError("calibration dict is required")
        self._cal = calibration.copy()
        # Allow calibration.json to override target_diam
        if "target_diam" in self._cal:
            self._target_diam = float(self._cal["target_diam"])
        self._si = int(self._cal["SI"])
        self._force_scale = np.float32(self._cal["FORCE_SCALE"])
        self._sens_scale = np.float32(policy_sensor_scale(self._cal, policy_sensor_units))
        self._dtw_sensor_factor = np.float32(dtw_sensor_factor(self._cal))
        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._config_path = config_path or self._cal.get("config_path")
        if self._config_path is None:
            raise ValueError("config_path is required")
        self._target_states = target_states
        self._target_harmonics = target_harmonics

        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 = None
        self.target_sensor = None  # (FIFO_LEN, 6) for DTW
        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.sensor_ids = []
        self.pinball_ids = []
        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)

        # ---- Phase 1: Target ----
        if self._target_states is None or self._target_harmonics is None:
            print("  [illusion] Phase 1: Recording target + harmonics...", flush=True)
            self._target_states, self._target_harmonics = record_illusion_target(
                self._config_path, self.device_id, self._si, self._target_diam,
                sensor_factor=float(self._dtw_sensor_factor))
            print("  [illusion] Target recorded.", flush=True)

        self.target_sensor = self._target_states[:, 0:6].copy()

        # ---- Phase 2: Training sim ----
        print("  [illusion] Phase 2: Training sim...", flush=True)
        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.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 = [3, 4, 5]

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

        zero_omega = self._action_to_omega(np.zeros(3, dtype=np.float32))
        self.smoother.reset(zero_omega.copy())
        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[0:6] * self._dtw_sensor_factor
            fifo_save.append(sl.copy())
        self.save_states = np.array(fifo_save, dtype=np.float32)

        self._gpu_block(lambda: self.sim.snapshot())
        self._snapshot_fifo = self.save_states.copy()
        self._snapshot_omega = zero_omega.copy()
        self._snapshot_ema = (0.0, 0.0)
        self._snapshot_step = 0
        snapshot_raw = self._read_obs()
        snapshot_target = gen_target_states_at(
            self._snapshot_step, self._target_harmonics)
        self._snapshot_obs = self._make_obs(
            snapshot_raw,
            float(snapshot_target[6]) / self._force_scale,
            float(snapshot_target[7]) / self._force_scale,
        ).copy()
        print(f"  [illusion] Init done ({time.perf_counter()-t0:.0f}s)")

    def _read_obs(self):
        obs = []
        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):
        sv = (np.asarray(action_norm, dtype=np.float32) * ACTION_SCALE + ACTION_BIAS) * U0
        return -sv / RADIUS

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

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

    def _make_obs(self, obs_slice, target_cd, target_cl):
        base = self._normalize_obs(obs_slice)
        return np.hstack([base, np.float32(target_cd), np.float32(target_cl)])

    def _compute_reward(self, obs_slice):
        forces_raw = obs_slice[6:12]
        cd_raw = forces_raw[0] + forces_raw[2] + forces_raw[4]
        cl_raw = forces_raw[1] + forces_raw[3] + forces_raw[5]
        cd_norm = cd_raw / self._force_scale
        cl_norm = cl_raw / self._force_scale

        tgt = gen_target_states_at(self.current_step, self._target_harmonics)
        target_cd = float(tgt[6]) / self._force_scale
        target_cl = float(tgt[7]) / self._force_scale

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

        r_cd_raw = float(np.exp(-(cd_norm - target_cd)**2 * self._k_cd))
        r_cl_raw = float(np.exp(-(cl_norm - target_cl)**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),
                "target_cd": float(target_cd), "target_cl": float(target_cl),
                "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):
        super().reset(seed=seed)
        self._gpu_block(lambda: self.sim.restore())
        self.smoother.reset(self._snapshot_omega)
        self.fifo_states.clear()
        for row in self._snapshot_fifo:
            self.fifo_states.append(row.copy())
        self.current_step = self._snapshot_step
        self._ema_r_cd, self._ema_r_cl = self._snapshot_ema
        return self._snapshot_obs.copy(), {}

    def step(self, action):
        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))

        self.current_step += 1
        obs_raw = self._read_obs()
        self.fifo_states.append(obs_raw[0:6] * self._dtw_sensor_factor)
        reward, info = self._compute_reward(obs_raw)

        tgt = gen_target_states_at(self.current_step, self._target_harmonics)
        tgt_cd = float(tgt[6]) / self._force_scale
        tgt_cl = float(tgt[7]) / self._force_scale
        obs = self._make_obs(obs_raw, tgt_cd, tgt_cl)

        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)
    parser.add_argument("--config", type=str, required=True)
    args = parser.parse_args()

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

    print("=== IllusionCloakEnv V5 Quick Test ===")
    env = IllusionCloakEnv(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 ===")
