#!/usr/bin/env python3
"""Vortex Cloak environment (V5 — no_bias, 2000x600, transfer from Karman).

Two-phase initialization:
  Phase 1: sensors only -> warmup -> add_vortex(x=10) -> record target(150 steps)
  Phase 2: sensors + pinball -> warmup -> add_vortex(x=15) -> FIFO -> snapshot

reset() restores to vortex+pinball state. MAX_STEPS=150 (time-bounded).
After step 150, done=True. No disturbance cylinder.

Observation (12-dim, physical norm, NO clip):
  [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 * 12 + [0,0,0]) * U0 / R

Reward: Gaussian + EMA + normalized DTW, same as Karman cloak.
"""
from __future__ import annotations

import json
import 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 CelerisLab.lbm.initializers import add_vortex

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

PINBALL_FRONT_X = 1000.0
PINBALL_REAR_X = 1026.0
SENSOR_X = 1200.0

VORTEX_X_TARGET = 200.0   # x=10*L0 for target phase
VORTEX_X_TRAIN = 300.0    # x=15*L0 for training (closer to pinball)
VORTEX_RADIUS = 40.0      # 2*L0
VORTEX_STRENGTH_LAMB = 0.5 * U0
VORTEX_STRENGTH_TAYLOR = 0.03 * U0

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

# ---------------------------------------------------------------------------
# DTW utilities (same as env_karman.py)
# ---------------------------------------------------------------------------
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
    t_slice = target[:conv_len]
    s_slice = state[-conv_len:]
    sim = 0.0
    for i in range(6):
        sim += calc_dtw_sim(t_slice[:, i], s_slice[:, i], 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


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

    def __init__(self, device_id=0, seed=42, calibration=None,
                 config_path=None, vortex_type="lamb"):
        super().__init__()
        self.device_id = device_id
        self.seed = seed
        np.random.seed(seed)
        self._vortex_type = vortex_type

        if calibration is None:
            raise ValueError("calibration dict is required")
        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._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 = None
        self.target_states = 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.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 (sensors + vortex only, no pinball) ----
        print("  [vortex] Phase 1: Target recording...", flush=True)
        sim_t = Simulation(lbm_config_path=self._config_path, device_id=self.device_id)
        sim_t._assert_object_count_contract = lambda *a, **kw: None
        s0 = sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0)
        s1 = sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0)
        s2 = sim_t.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0)
        sensor_ids_t = [s0, s1, s2]
        sim_t.initialize()
        sim_t.run(warmup, zero_obs=True)
        print(f"  [vortex] Target warmup done ({warmup} steps)")

        # Inject vortex and record
        add_vortex(sim_t.field, (VORTEX_X_TARGET, CENTER_Y),
                   VORTEX_RADIUS, VORTEX_STRENGTH_TAYLOR if self._vortex_type == "taylor" else VORTEX_STRENGTH_LAMB,
                   self._vortex_type)

        target = np.zeros((MAX_STEPS, 6), dtype=np.float32)
        for i in range(MAX_STEPS):
            sim_t.run(self._si, zero_obs=True)
            target[i] = [
                sim_t.read_sensor(s0, normalize=True)[0] * SENSOR_CC,
                sim_t.read_sensor(s0, normalize=True)[1] * SENSOR_CC,
                sim_t.read_sensor(s1, normalize=True)[0] * SENSOR_CC,
                sim_t.read_sensor(s1, normalize=True)[1] * SENSOR_CC,
                sim_t.read_sensor(s2, normalize=True)[0] * SENSOR_CC,
                sim_t.read_sensor(s2, normalize=True)[1] * SENSOR_CC,
            ]
        sim_t.close()
        self.target_states = target

        # ---- Phase 2: Training sim (sensors + pinball + vortex) ----
        print("  [vortex] 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"  [vortex] Pinball warmup done ({warmup} steps)")

        # Inject vortex at training position (closer to pinball)
        self._gpu_block(lambda: add_vortex(self.sim.field, (VORTEX_X_TRAIN, CENTER_Y),
                                            VORTEX_RADIUS,
                                            VORTEX_STRENGTH_TAYLOR if self._vortex_type == "taylor" else VORTEX_STRENGTH_LAMB,
                                            self._vortex_type))
        print(f"  [vortex] Vortex ({self._vortex_type}) injected at x={VORTEX_X_TRAIN}")

        # Zero-action FIFO (no rotation, vortex evolves with pinball)
        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]  # 6 sensor channels (legacy-equiv)
            sl = sl * SENSOR_CC
            fifo_save.append(sl.copy())
        self.save_states = np.array(fifo_save, dtype=np.float32)

        self._gpu_block(lambda: self.sim.snapshot())
        print(f"  [vortex] 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 _compute_reward(self, obs_slice):
        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:
            sim_val = compute_similarity(self.target_states,
                                          np.array(list(self.fifo_states)),
                                          CONV_LEN, 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):
        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, :])
        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)
        return obs, {}

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

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

        self.current_step += 1
        done = self.current_step >= MAX_STEPS
        return obs, reward, done, 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)
    parser.add_argument("--vortex-type", type=str, default="lamb",
                        choices=["lamb", "taylor"])
    args = parser.parse_args()

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

    print(f"=== VortexCloakEnv V5 Quick Test ({args.vortex_type}) ===")
    env = VortexCloakEnv(device_id=args.device_id, calibration=cal,
                          config_path=args.config, vortex_type=args.vortex_type)

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

    rewards = []
    for step in range(MAX_STEPS):
        obs, reward, done, *_ = env.step(np.zeros(3, dtype=np.float32))
        rewards.append(reward)
        if done:
            break
    print(f"  Zero-action avg reward: {np.mean(rewards):.4f}")
    print(f"  Steps: {len(rewards)}")

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