#!/usr/bin/env python3
# reproduce/phase2_open_loop.py
"""Phase 2: Open-loop target-signal validation.

Runs the target-recording phase on the new CelerisLab with the
legacy-compatible config (regularized inlet, NBB equivalent),
then compares against SR_analysis reference target signals.

This isolates CFD differences before DRL is involved.

Usage:
    conda run -n pycuda_3_10 python phase2_open_loop.py --device 0
    conda run -n pycuda_3_10 python phase2_open_loop.py --device 0 --scene karman
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
import pycuda.driver as cuda; cuda.init()

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

from CelerisLab import Simulation  # noqa: E402
from CelerisLab.lbm.initializers import add_vortex  # noqa: E402
from reproduce.core.open_loop_validator import compare_target_signals  # noqa: E402

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
CFG_PATH = "configs/config_lbm_pinball_legacy_compat.json"
L0 = 20.0
U0 = 0.01
NX = 1280
NY = 512
CENTER_Y = float(NY - 1) / 2.0
RADIUS = L0 / 2.0  # 10
FIFO_LEN = 150
SI = 800
WARMUP = int(4.0 * NX / U0)

REF_BASE = os.path.join(_SRC, "SR_analysis", "data")
OUT_BASE = os.path.join(os.path.dirname(__file__), "output", "phase2_validation")


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


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_sensors_legacy(sim, sensor_ids, cc):
    obs = []
    for sid in sensor_ids:
        s = sim.read_sensor(sid, normalize=True)
        obs.extend([float(s[0]) * cc, float(s[1]) * cc])
    return np.array(obs, dtype=np.float32)


# ---------------------------------------------------------------------------
# Karman target: dist-cyl + 3 sensors
# ---------------------------------------------------------------------------
def validate_karman(device_id: int, out_dir: str) -> dict:
    log("=== Karman target validation ===")
    ref_dir = os.path.join(REF_BASE, "karman", "karman_re100")
    os.makedirs(out_dir, exist_ok=True)

    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])
    log(f"  Sensor CC: {cc}")

    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        obs = read_sensors_legacy(sim, sensor_ids, cc)
        target[i] = obs

    np.savez_compressed(os.path.join(out_dir, "target.npz"), target_states=target)
    sim.close()

    result = compare_target_signals(ref_dir, target, label="karman")
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    return result


# ---------------------------------------------------------------------------
# Steady channel target: 3 sensors only
# ---------------------------------------------------------------------------
def validate_steady(device_id: int, out_dir: str) -> dict:
    log("=== Steady channel target validation ===")
    ref_dir = os.path.join(REF_BASE, "steady", "steady")
    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.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])
    log(f"  Sensor CC: {cc}")

    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        target[i] = read_sensors_legacy(sim, sensor_ids, cc)

    np.savez_compressed(os.path.join(out_dir, "target.npz"), target_states=target)
    sim.close()

    result = compare_target_signals(ref_dir, target, label="steady")
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    return result


# ---------------------------------------------------------------------------
# Illusion target: target cylinder + 3 sensors
# ---------------------------------------------------------------------------
def validate_illusion(device_id: int, out_dir: str, diam_L: float = 1.0) -> dict:
    # diam_L: 0.75 → "illusion_0.75L", 1.0 → "illusion_1L", 1.5 → "illusion_1.5L"
    if diam_L == int(diam_L):
        label_suffix = str(int(diam_L))
    else:
        label_suffix = str(diam_L).rstrip('0')
    label = f"illusion_{label_suffix}L"
    log(f"=== {label} target validation ===")
    ref_dir = os.path.join(REF_BASE, "illusion", label)
    os.makedirs(out_dir, exist_ok=True)

    sim = Simulation(lbm_config_path=CFG_PATH, device_id=device_id)
    sim.add_body("circle", center=(20.0 * L0, CENTER_Y, 0.0), radius=diam_L * L0)
    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.initialize()
    sim.run(WARMUP, zero_obs=True)
    cc = get_cc(sim, sensor_ids[0])
    log(f"  Sensor CC: {cc}")

    # Target: cyl_force(2) + sensors(6) = 8 channels
    target = np.zeros((FIFO_LEN, 8), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        f = list(sim.read_force(0, normalize=True))
        s = read_sensors_legacy(sim, sensor_ids, cc)
        target[i] = np.hstack([f, s])

    np.savez_compressed(os.path.join(out_dir, "target.npz"), target_states=target)
    sim.close()

    result = compare_target_signals(ref_dir, target, label=label, sensor_slice=slice(2, 8))
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    return result


# ---------------------------------------------------------------------------
# Vortex target: vortex + 3 sensors
# ---------------------------------------------------------------------------
def validate_vortex(device_id: int, out_dir: str, vortex_type: str = "lamb") -> dict:
    log(f"=== Vortex {vortex_type} target validation ===")
    ref_dir = os.path.join(REF_BASE, "vortex", f"vortex_{vortex_type}")
    os.makedirs(out_dir, exist_ok=True)

    strength = 0.5 * U0 if vortex_type == "lamb" else 0.03 * U0

    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)

    # Add vortex
    sim.field.download_ddf()
    add_vortex(sim.field, (10.0 * L0, CENTER_Y), 2.0 * L0, strength, vortex_type)

    cc = get_cc(sim, sensor_ids[0])
    log(f"  Sensor CC: {cc}")

    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    for i in range(FIFO_LEN):
        sim.run(SI, zero_obs=True)
        target[i] = read_sensors_legacy(sim, sensor_ids, cc)

    np.savez_compressed(os.path.join(out_dir, "target.npz"), target_states=target)
    sim.close()

    result = compare_target_signals(ref_dir, target, label=f"vortex_{vortex_type}")
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    return result


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    ap = argparse.ArgumentParser(description="Phase 2: Open-loop target validation")
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--scene", type=str, default="all",
                    help="Scene to validate: karman, steady, illusion_1L, vortex_lamb, vortex_taylor, all")
    args = ap.parse_args()

    results = {}
    scenes = {
        "karman": lambda: validate_karman(args.device, os.path.join(OUT_BASE, "karman")),
        "steady": lambda: validate_steady(args.device, os.path.join(OUT_BASE, "steady")),
        "illusion_1L": lambda: validate_illusion(args.device, os.path.join(OUT_BASE, "illusion_1L"), 1.0),
        "vortex_lamb": lambda: validate_vortex(args.device, os.path.join(OUT_BASE, "vortex_lamb"), "lamb"),
        "vortex_taylor": lambda: validate_vortex(args.device, os.path.join(OUT_BASE, "vortex_taylor"), "taylor"),
    }

    if args.scene == "all":
        for name, func in scenes.items():
            results[name] = func()
    else:
        for s in args.scene.split(","):
            s = s.strip()
            if s not in scenes:
                log(f"Unknown scene: {s}")
                continue
            results[s] = scenes[s]()

    # Summary
    log("\n=== Open-loop validation summary ===")
    all_pass = True
    for name, r in results.items():
        status = "PASS" if r["passed"] else "FAIL"
        log(f"  {name}: DTW={r['dtw_sim']:.4f}, corr={[f'{c:.3f}' for c in r['channel_corr']]} -> {status}")
        if not r["passed"]:
            all_pass = False

    if all_pass:
        log("\nALL SCENES PASSED open-loop validation. Proceed to Phase 3.")
    else:
        log("\nSOME SCENES FAILED. Review Phase 2 results before Phase 3.")
