#!/usr/bin/env python3
"""Comprehensive run: Steady Cloak + Karman Cloak (new norm / legacy norm) + flow field output.

DEPRECATED (2026-07): Superseded by phase2_open_loop.py and phase3_reproduce.py
which use the legacy-compatible config (regularized inlet). This script uses
the old config_lbm_pinball.json with zou_he_local inlet, producing suboptimal results.

Keep for reference; use phase2/phase3 for new reproduce work.

Runs three cases:
  Case A: Steady Cloak (open-loop constant rotation, no DRL)
  Case B: Karman Cloak (new-CFD norm + DRL inference)
  Case C: Karman Cloak (legacy norm + DRL inference)

Each case saves: macroscopic.npz (rho/ux/uy), vorticity.png, sensors_forces.npz
"""
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 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 compute_norm, normalize_observation

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
L0 = 20.0
U0 = float(_U0)
NX = 1280
NY = 512
CENTER_Y = float(NY - 1) / 2.0
FIFO_LEN = 150
SI = 800
WARMUP = int(4.0 * NX / U0)
CFG_PATH = "configs/config_lbm_pinball_legacy_compat.json"
REF_DIR = os.path.join(_REPO, "src", "SR_analysis", "data", "karman", "karman_re100")
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, dist_id, pinball_ids, cc):
    obs = []
    if dist_id is not None:
        obs.extend(sim.read_force(dist_id, normalize=True))
    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 save_field(sim, out_dir, name):
    """Save macroscopic field and render vorticity."""
    macro = sim.get_macroscopic()
    np.savez_compressed(os.path.join(out_dir, f"macro_{name}.npz"),
                        rho=macro["rho"], ux=macro["ux"], uy=macro["uy"])
    vort = compute_vorticity(macro["ux"], macro["uy"])
    render_vorticity_field(
        vort, nx=NX, ny=NY,
        out_path=os.path.join(out_dir, f"vorticity_{name}.png"),
        cylinders=[
            ((10.0 * L0, CENTER_Y), 1.0 * L0),
            ((30.0 * L0, CENTER_Y), RADIUS),
            ((31.3 * L0, CENTER_Y + 15.0), RADIUS),
            ((31.3 * L0, CENTER_Y - 15.0), RADIUS),
        ],
    )
    print(f"  Saved {name}: macro + vorticity.png")


# =========================================================================
# Case A: Steady Cloak (open-loop constant rotation)
# =========================================================================
def run_steady_cloak(device_id, out_dir):
    """Steady cloak: pinball only, constant bias=[0,-5.1,5.1]*U0, no DRL."""
    print("\n" + "=" * 70)
    print("Case A: Steady Cloak (open-loop constant rotation)")
    print("=" * 70)
    os.makedirs(out_dir, exist_ok=True)

    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.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.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])

    save_field(sim, out_dir, "steady_uncontrolled")

    # Bias: surface_vel = [0, -5.1, 5.1] * U0 → omega
    bias_surf = np.array([0.0, -5.1, 5.1], dtype=np.float32) * U0
    bias_omega = surface_vel_to_omega(bias_surf)
    print(f"  Steady bias omega: {bias_omega}")
    sim.set_body(3, omega=bias_omega[0])  # front
    sim.set_body(4, omega=bias_omega[1])  # top
    sim.set_body(5, omega=bias_omega[2])  # bottom

    sim.run(WARMUP, zero_obs=True)

    sensors_f = []
    for _ in range(200):
        sim.run(SI, zero_obs=True)
        obs = read_obs_legacy(sim, sensor_ids, None, [3, 4, 5], cc)
        sensors_f.append(obs[0:12])
    sensors_f = np.array(sensors_f, dtype=np.float32)
    np.savez_compressed(os.path.join(out_dir, "steady_signals.npz"),
                        sensors=sensors_f[:, 0:6], forces=sensors_f[:, 6:12])

    save_field(sim, out_dir, "steady_controlled")
    print(f"  Forces: front_fy={sensors_f[:,7].mean():+.6f} "
          f"top_fy={sensors_f[:,9].mean():+.6f} bottom_fy={sensors_f[:,11].mean():+.6f}")
    sim.close()


# =========================================================================
# Karman shared build
# =========================================================================
def build_karman_env(device_id, out_dir):
    """Build Karman env, record target, add pinball, return (sim, ids, norm)."""
    sim = Simulation(lbm_config_path=CFG_PATH, device_id=device_id)

    dist_id = sim.add_body("circle", center=(10.0 * L0, CENTER_Y, 0.0), radius=1.0 * L0)
    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])

    target = 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, dist_id, [], cc)
        target = np.vstack((target, obs[2:8]))
    np.savez_compressed(os.path.join(out_dir, "target.npz"), target_states=target)

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

    sim.run(WARMUP, zero_obs=True)
    return sim, dist_id, sensor_ids, (fid, tid, bid), cc, target


def run_karman_drl(sim, dist_id, sensor_ids, pinball_ids, cc, target,
                   norm, model, out_dir, case_label, num_steps=200):
    """Run DRL inference with given norm."""
    fid, tid, bid = pinball_ids

    # Bias FIFO: surface_vel = [0, -4, 4] * U0  (legacy: bias_arr[-3:] = front, top, bottom)
    bias_surf = np.array([0.0, -4.0, 4.0], dtype=np.float32) * U0
    bias_omega = surface_vel_to_omega(bias_surf)

    ema = ActionSmoother(weight=0.1)
    ema.reset(np.zeros(3, dtype=np.float32))  # legacy starts from zero
    for _ in range(FIFO_LEN):
        s = ema(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.reset(bias_omega.copy())  # EMA starts from converged bias state

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

    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, -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, dist_id, [fid, tid, bid], cc)
        sl = obs[2:14]
        sig_s[step] = sl[0:6]
        sig_f[step] = sl[6:12]
        sig_a[step] = action
        obs_norm = normalize_observation(sl, norm)

    save_field(sim, out_dir, f"karman_{case_label}")
    np.savez_compressed(os.path.join(out_dir, f"signals_{case_label}.npz"),
                        sensors=sig_s, forces=sig_f, actions=sig_a)

    print(f"  [{case_label}] Actions: aF={sig_a[:,0].mean():+.4f} "
          f"aB={sig_a[:,1].mean():+.4f} aT={sig_a[:,2].mean():+.4f}")
    print(f"  [{case_label}] Forces:  front_fy={sig_f[:,1].mean():+.6f} "
          f"top_fy={sig_f[:,3].mean():+.6f} bottom_fy={sig_f[:,5].mean():+.6f}")


# =========================================================================
# Case B: Karman + new norm
# =========================================================================
def collect_new_norm(sim, sensor_ids, dist_id, pinball_ids, cc):
    """Collect norm values on new CFD from zero-action FIFO."""
    fid, tid, bid = pinball_ids
    fifo = []
    for _ in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        obs = read_obs_legacy(sim, sensor_ids, dist_id, [fid, tid, bid], cc)
        fifo.append(obs[2:14])
    f = np.array(fifo, dtype=np.float32)
    norm = compute_norm(f, force_slice=(6, 12), sens_slice=(0, 6))
    print(f"  New norm: fn={norm['force_norm_fact']:.6f}")
    return norm


def run_karman_new_norm(device_id, out_dir):
    print("\n" + "=" * 70)
    print("Case B: Karman Cloak (new-CFD norm + DRL)")
    print("=" * 70)
    os.makedirs(out_dir, exist_ok=True)

    model = ModelInventory().load("d1a3o12_re100", device="cpu")
    sim, dist_id, sensor_ids, pids, cc, target = build_karman_env(device_id, out_dir)
    norm = collect_new_norm(sim, sensor_ids, dist_id, pids, cc)
    np.savez(os.path.join(out_dir, "norm_new.npz"), **norm)
    run_karman_drl(sim, dist_id, sensor_ids, pids, cc, target,
                   norm, model, out_dir, "new_norm")
    sim.close()


# =========================================================================
# Case C: Karman + legacy norm
# =========================================================================
def run_karman_legacy_norm(device_id, out_dir):
    print("\n" + "=" * 70)
    print("Case C: Karman Cloak (legacy norm + DRL)")
    print("=" * 70)
    os.makedirs(out_dir, exist_ok=True)

    with open(os.path.join(REF_DIR, "norm.json")) as f:
        d = json.load(f)
    norm = {
        "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),
    }

    model = ModelInventory().load("d1a3o12_re100", device="cpu")
    sim, dist_id, sensor_ids, pids, cc, target = build_karman_env(device_id, out_dir)
    run_karman_drl(sim, dist_id, sensor_ids, pids, cc, target,
                   norm, model, out_dir, "legacy_norm")
    sim.close()


# =========================================================================
# Main
# =========================================================================
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--device", type=int, default=0)
    parser.add_argument("--cases", type=str, default="A,B,C",
                        help="Comma-separated: A=steady, B=new-norm, C=legacy-norm")
    args = parser.parse_args()

    cases = [c.strip().upper() for c in args.cases.split(",")]

    if "A" in cases:
        run_steady_cloak(args.device, os.path.join(OUT_BASE, "steady_cloak"))

    karman_dir = os.path.join(OUT_BASE, "karman_cloak")
    if "B" in cases:
        run_karman_new_norm(args.device, karman_dir)
    if "C" in cases:
        run_karman_legacy_norm(args.device, karman_dir)

    print("\nDone.")
