#!/usr/bin/env python3
"""Phase 2 (Illusion) + Phase 3 (Vortex) reproduction on new CelerisLab.

Uses same proven approach as Karman: legacy norm + old-equiv sensor conversion.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
_SRC = os.path.join(_REPO, "src")
for p in [_REPO, _SRC]:
    if p not in sys.path:
        sys.path.insert(0, p)

import pycuda.driver as cuda
cuda.init()

from CelerisLab import Simulation
from CelerisLab.common.render import compute_vorticity, render_vorticity_field
from CelerisLab.lbm.initializers import add_vortex
from drl_pinball.reproduce.configs.model_inventory import ModelInventory
from drl_pinball.reproduce.core.action_wrapper import (
    ActionSmoother, norm_action_to_omega, surface_vel_to_omega, U0 as _U0, RADIUS,
)
from drl_pinball.reproduce.core.obs_normalizer import normalize_observation

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
L0 = 20.0
U0 = float(_U0)
NX = 1280
NY = 512
CENTER_Y = float(NY - 1) / 2.0
FIFO_LEN = 150
WARMUP = int(4.0 * NX / U0)
CFG_PATH = "configs/config_lbm_pinball_legacy_compat.json"
OUT_BASE = os.path.join(os.path.dirname(__file__), "output")


def get_cc(sim, sid):
    nx, ny = sim.lbm_cfg.nx, sim.lbm_cfg.ny
    cells_arr, _ = sim.bodies.get(sid).get_sensor_list(nx, ny)
    return float(len(cells_arr))


def read_obs_legacy(sim, sensor_ids, pinball_ids, cc):
    """Return [s0_ux,uy, s1_ux,uy, s2_ux,uy, front_fx,fy, top_fx,fy, bottom_fx,fy].
    No dist_cylinder — this is for 6-object envs (illusion, vortex, steady).
    """
    obs = []
    for sid in sensor_ids:
        s = sim.read_sensor(sid, normalize=True)
        obs.extend(s * cc)
    for pid in pinball_ids:
        obs.extend(sim.read_force(pid, normalize=True))
    return np.array(obs, dtype=np.float32)


def load_norm(path):
    with open(path) as f:
        d = json.load(f)
    return {
        "force_norm_fact": np.float32(d["force_norm_fact"]),
        "sens_deviation": np.array(d["sens_deviation"], dtype=np.float32),
        "sens_norm_fact": np.array(d["sens_norm_fact"], dtype=np.float32),
    }


def save_vorticity(sim, out_path, cylinders):
    macro = sim.get_macroscopic()
    vort = compute_vorticity(macro["ux"], macro["uy"])
    render_vorticity_field(vort, nx=NX, ny=NY, out_path=out_path, cylinders=cylinders)


# =========================================================================
# Phase 2: Illusion
# =========================================================================
def run_illusion(device_id, target_label, model_name, sample_interval, target_diameter,
                 ref_dir, out_dir):
    """Run Illusion inference using legacy norm + legacy target harmonics."""
    print(f"\n{'='*70}")
    print(f"Illusion {target_label} (model={model_name}, SI={sample_interval})")
    print(f"{'='*70}")
    os.makedirs(out_dir, exist_ok=True)

    norm = load_norm(os.path.join(ref_dir, "norm.json"))
    with open(os.path.join(ref_dir, "target_harmonics.json")) as f:
        target_harmonics = json.load(f)
    legacy_target = np.load(os.path.join(ref_dir, "target.npz"))
    legacy_target_states = legacy_target["target_states"]

    def gen_target_states_at(t, harmonics):
        D = len(harmonics)
        result = np.zeros(D, dtype=np.float32)
        for d, h in enumerate(harmonics):
            val = np.float32(h["dc"])
            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
        return result

    # Pinball env (no disturbance cylinder)
    sim = Simulation(lbm_config_path=CFG_PATH, device_id=device_id)
    sensor_ids = [
        sim.add_body("sensor", center=(30.0 * L0, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(30.0 * L0, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(30.0 * L0, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.add_body("circle", center=(19.0 * L0, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(20.3 * L0, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(20.3 * L0, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])
    fid, tid, bid = 3, 4, 5

    # Bias FIFO: init surface_vel = [0, -1, 1] * U0 (matching legacy_env_imit.py)
    init_bias_surf = np.array([0.0, -1.0, 1.0], dtype=np.float32) * U0
    init_bias_omega = surface_vel_to_omega(init_bias_surf)

    ema_bias = ActionSmoother(weight=0.1)
    ema_bias.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema_bias(init_bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(sample_interval, zero_obs=True)
    sim.snapshot()

    # DRL inference — EMA starts from converged init-bias state
    sim.restore()
    ema = ActionSmoother(weight=0.1)
    ema.reset(init_bias_omega.copy())

    model = ModelInventory().load(model_name, device="cpu")
    print(f"  Model loaded on CPU")

    obs_init = read_obs_legacy(sim, sensor_ids, [fid, tid, bid], cc)
    obs_norm = normalize_observation(obs_init, norm)
    t0 = gen_target_states_at(0, target_harmonics)
    target_cd = np.float32(t0[0] / norm["force_norm_fact"])
    target_cl = np.float32(t0[1] / norm["force_norm_fact"])
    obs_norm = np.clip(np.hstack([obs_norm, [target_cd, target_cl]]), -1, 1).astype(np.float32)

    num_steps = 200
    sig_s = np.zeros((num_steps, 6), dtype=np.float32)
    sig_f = np.zeros((num_steps, 6), dtype=np.float32)
    sig_a = np.zeros((num_steps, 3), dtype=np.float32)

    for step in range(num_steps):
        action, _ = model.predict(obs_norm, deterministic=True)
        action = np.asarray(action, dtype=np.float32).flatten()
        target_omega = norm_action_to_omega(action, scale=8.0, bias=(0.0, -2.0, 2.0))
        smoothed = ema(target_omega)

        sim.set_body(fid, omega=smoothed[0]); sim.set_body(tid, omega=smoothed[1]); sim.set_body(bid, omega=smoothed[2])
        sim.run(sample_interval, zero_obs=True)

        obs = read_obs_legacy(sim, sensor_ids, [fid, tid, bid], cc)
        sig_s[step] = obs[0:6]
        sig_f[step] = obs[6:12]
        sig_a[step] = action

        obs_norm_base = normalize_observation(obs, norm)
        t_h = gen_target_states_at(step, target_harmonics)
        target_cd = np.float32(t_h[0] / norm["force_norm_fact"])
        target_cl = np.float32(t_h[1] / norm["force_norm_fact"])
        obs_norm = np.clip(np.hstack([obs_norm_base, [target_cd, target_cl]]), -1, 1).astype(np.float32)

    save_vorticity(sim, os.path.join(out_dir, f"vorticity.png"),
                   cylinders=[((19.0*L0, CENTER_Y), RADIUS),
                              ((20.3*L0, CENTER_Y+15.0), RADIUS),
                              ((20.3*L0, CENTER_Y-15.0), RADIUS)])
    np.savez_compressed(os.path.join(out_dir, "signals.npz"),
                        sensors=sig_s, forces=sig_f, actions=sig_a)
    sim.close()

    # Compare with reference
    ref = np.load(os.path.join(ref_dir, "controlled.npz"))
    print(f"\n  Actions vs ref:")
    for i, name in enumerate(["aF","aB","aT"]):
        c = np.corrcoef(ref["actions"][:num_steps,i], sig_a[:,i])[0,1]
        print(f"    {name}: ref_mean={ref['actions'][:num_steps,i].mean():+.4f} "
              f"our_mean={sig_a[:,i].mean():+.4f} corr={c:+.4f}")

    # DTW similarity
    n_c = 36
    def dtw_sim(t, s):
        n = len(t)
        D = np.full((n+1, n+1), np.inf)
        D[0,0] = 0
        for i in range(1, n+1):
            for j in range(1, n+1):
                D[i,j] = abs(t[i-1]-s[j-1]) + min(D[i-1,j], D[i,j-1], D[i-1,j-1])
        return 1 - D[n,n] / n

    t_seq = legacy_target_states[n_c:2*n_c, 1+2]
    s_seq = sig_s[-n_c:, 1]
    if np.std(t_seq) > 1e-10 and np.std(s_seq) > 1e-10:
        corr = np.correlate(t_seq - t_seq.mean(), s_seq - s_seq.mean(), mode="full")
        lag = np.argmax(corr) - (len(t_seq) - 1)
    else:
        lag = 0

    sim_val = 0.0
    for i in range(6):
        t_rolled = np.roll(legacy_target_states[:, i+2], -lag)[n_c:2*n_c]
        sim_val += dtw_sim(t_rolled, sig_s[-n_c:, i]) / 6.0
    print(f"  DTW similarity (legacy target vs our sensors): {sim_val:.4f}")
    print(f"  Reference DTW similarity: {0.9754}")


# =========================================================================
# Phase 3: Vortex (lamb / taylor)
# =========================================================================
def run_vortex(device_id, vortex_type, model_name, vortex_strength, ref_dir, out_dir):
    """Run Vortex cloak inference using legacy norm."""
    print(f"\n{'='*70}")
    print(f"Vortex {vortex_type} (model={model_name}, strength={vortex_strength})")
    print(f"{'='*70}")
    os.makedirs(out_dir, exist_ok=True)

    norm = load_norm(os.path.join(ref_dir, "norm.json"))
    legacy_target = np.load(os.path.join(ref_dir, "target.npz"))["target_states"]

    MAX_STEPS = 150
    CONV_LEN = 30
    SI = 800

    # Stage 1: sensor-only env, record target with vortex
    sim = Simulation(lbm_config_path=CFG_PATH, device_id=device_id)
    sensor_ids = [
        sim.add_body("sensor", center=(40.0 * L0, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(40.0 * L0, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(40.0 * L0, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])

    sim.snapshot()

    add_vortex(sim.field, center=(10.0 * L0, CENTER_Y), radius=2.0 * L0,
               strength=vortex_strength, vortex_type=vortex_type)

    target_states = np.empty((0, 6), dtype=np.float32)
    for _ in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        obs = read_obs_legacy(sim, sensor_ids, [], cc)
        target_states = np.vstack((target_states, obs))
    np.savez_compressed(os.path.join(out_dir, "target.npz"), target_states=target_states)

    # Stage 2: add pinball + vortex
    sim.restore()
    n0 = sim.bodies.count
    sim.add_body("circle", center=(30.0 * L0, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(31.3 * L0, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(31.3 * L0, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.sync_bodies()
    fid, tid, bid = list(range(n0, n0 + 3))

    # Bias warmup: surface_vel = [0, -4, 4] * U0
    bias_surf = np.array([0.0, -4.0, 4.0], dtype=np.float32) * U0
    bias_omega = surface_vel_to_omega(bias_surf)

    ema_init = ActionSmoother(weight=0.1)
    ema_init.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema_init(bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(SI, zero_obs=True)

    # Add vortex at pinball-phase position
    add_vortex(sim.field, center=(15.0 * L0, CENTER_Y), radius=2.0 * L0,
               strength=vortex_strength, vortex_type=vortex_type)
    sim.snapshot()
    print(f"  Post-vortex DDF saved")

    # Norm FIFO (zero action, with pinball+vortex)
    fifo = []
    for _ in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        obs = read_obs_legacy(sim, sensor_ids, [fid, tid, bid], cc)
        fifo.append(obs)
    fifo_arr = np.array(fifo, dtype=np.float32)
    new_fn = 6.0 * np.max(np.abs(fifo_arr[:, 6:12]))
    print(f"  New CFD force_norm_fact: {new_fn:.6f} (legacy: {norm['force_norm_fact']:.6f})")

    # Bias FIFO after vortex
    sim.restore()
    ema_bias = ActionSmoother(weight=0.1)
    ema_bias.reset(np.zeros(3, dtype=np.float32))
    for _ in range(FIFO_LEN):
        s = ema_bias(bias_omega)
        sim.set_body(fid, omega=s[0]); sim.set_body(tid, omega=s[1]); sim.set_body(bid, omega=s[2])
        sim.run(SI, zero_obs=True)
    sim.snapshot()

    # DRL inference
    sim.restore()
    ema = ActionSmoother(weight=0.1)
    ema.reset(bias_omega.copy())

    model = ModelInventory().load(model_name, device="cpu")
    print(f"  Model loaded on CPU")

    obs_init = read_obs_legacy(sim, sensor_ids, [fid, tid, bid], cc)
    obs_norm = normalize_observation(obs_init, norm)

    sig_s = np.zeros((MAX_STEPS, 6), dtype=np.float32)
    sig_f = np.zeros((MAX_STEPS, 6), dtype=np.float32)
    sig_a = np.zeros((MAX_STEPS, 3), dtype=np.float32)

    for step in range(MAX_STEPS):
        action, _ = model.predict(obs_norm, deterministic=True)
        action = np.asarray(action, dtype=np.float32).flatten()
        target_omega = norm_action_to_omega(action, scale=4.0, bias=(0.0, -4.0, 4.0))
        smoothed = ema(target_omega)

        sim.set_body(fid, omega=smoothed[0]); sim.set_body(tid, omega=smoothed[1]); sim.set_body(bid, omega=smoothed[2])
        sim.run(SI, zero_obs=True)

        obs = read_obs_legacy(sim, sensor_ids, [fid, tid, bid], cc)
        sig_s[step] = obs[0:6]
        sig_f[step] = obs[6:12]
        sig_a[step] = action
        obs_norm = normalize_observation(obs, norm)

    save_vorticity(sim, os.path.join(out_dir, f"vorticity.png"),
                   cylinders=[((30.0*L0, CENTER_Y), RADIUS),
                              ((31.3*L0, CENTER_Y+15.0), RADIUS),
                              ((31.3*L0, CENTER_Y-15.0), RADIUS)])
    np.savez_compressed(os.path.join(out_dir, "signals.npz"),
                        sensors=sig_s, forces=sig_f, actions=sig_a)
    sim.close()

    # Compare with reference
    ref = np.load(os.path.join(ref_dir, "controlled.npz"))
    print(f"\n  Actions vs ref:")
    for i, name in enumerate(["aF","aB","aT"]):
        c = np.corrcoef(ref["actions"][:MAX_STEPS,i], sig_a[:,i])[0,1]
        print(f"    {name}: ref_mean={ref['actions'][:MAX_STEPS,i].mean():+.4f} "
              f"our_mean={sig_a[:,i].mean():+.4f} corr={c:+.4f}")

    # DTW similarity
    def dtw_sim(t, s):
        n = len(t)
        D = np.full((n+1, n+1), np.inf); D[0,0] = 0
        for i in range(1, n+1):
            for j in range(1, n+1):
                D[i,j] = abs(t[i-1]-s[j-1]) + min(D[i-1,j], D[i,j-1], D[i-1,j-1])
        return 1 - D[n,n] / n

    sim_val = 0.0
    for i in range(6):
        t_rolled = np.roll(target_states[-CONV_LEN:, i], -MAX_STEPS)
        sim_val += dtw_sim(t_rolled, sig_s[-CONV_LEN:, i]) / 6.0
    print(f"  DTW similarity (our target vs our sensors): {sim_val:.4f}")


# =========================================================================
# Main
# =========================================================================
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", type=int, default=0)
    parser.add_argument("--case", type=str, required=True,
                        choices=["illusion_075L", "illusion_1L", "illusion_15L",
                                 "vortex_lamb", "vortex_taylor"])
    args = parser.parse_args()

    _SRC2 = os.path.join(os.path.dirname(__file__), "..", "..", "..", "src")

    if "illusion" in args.case:
        scenes = {
            "illusion_075L": ("d1a3o14_250525_imit_075L_2U_400S", 400, 0.75 * L0,
                             os.path.join(_SRC2, "SR_analysis", "data", "illusion", "illusion_0.75L")),
            "illusion_1L":   ("d1a3o14_250525_imit_1L_2U_600S", 600, 1.0 * L0,
                             os.path.join(_SRC2, "SR_analysis", "data", "illusion", "illusion_1L")),
            "illusion_15L":  ("d1a3o14_250525_imit_15L_2U", 800, 1.5 * L0,
                             os.path.join(_SRC2, "SR_analysis", "data", "illusion", "illusion_1.5L")),
        }
        model_name, si, diam, ref_dir = scenes[args.case]
        out_dir = os.path.join(OUT_BASE, args.case)
        run_illusion(args.device, args.case, model_name, si, diam, ref_dir, out_dir)

    elif "vortex" in args.case:
        vtype = "lamb" if "lamb" in args.case else "taylor"
        scenes = {
            "vortex_lamb": ("vortex_lamb", 0.5 * U0,
                           os.path.join(_SRC2, "SR_analysis", "data", "vortex", "vortex_lamb")),
            "vortex_taylor": ("vortex_taylor", 0.03 * U0,
                             os.path.join(_SRC2, "SR_analysis", "data", "vortex", "vortex_taylor")),
        }
        model_name, strength, ref_dir = scenes[args.case]
        out_dir = os.path.join(OUT_BASE, args.case)
        run_vortex(args.device, vtype, model_name, strength, ref_dir, out_dir)

    print("\nDone.")
