#!/usr/bin/env python3
"""Phase 0: Baseline measurement for 2000x600 Karman Cloak.

Collects Stage0 (zero rotation) and Stage1 (bias [0,-4,4]*U0) data to inform
reward design. Records per-step Cd/Cl/Sim and obs norm health.

Outputs to output/stage_baseline_2000x600/:
  - stage_data.pkl : raw obs + forces + sim per stage
  - norm_health.json : obs distribution stats
  - summary.txt : human-readable comparison table

Usage:
    conda run -n pycuda_3_10 python -u phase0_baseline_measure.py --device-id 0
"""
from __future__ import annotations

import argparse
import json
import os
import pickle
import sys
import time
from collections import deque
from pathlib import Path

import numpy as np
import pycuda.driver as cuda

cuda.init()

_REPO = Path(__file__).resolve().parents[3]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

from CelerisLab import Simulation

# ---------------------------------------------------------------------------
# Hard-coded CelerisLab paths for cache cleaning
# ---------------------------------------------------------------------------
_CELERIS = _REPO / "CelerisLab"
_CONFIG_H = _CELERIS / "src/CelerisLab/lbm/kernels/config/config_objects.h"
_PTX = _CELERIS / "src/CelerisLab/lbm/kernels/kernel.ptx"


def _clean_cache():
    for p in [_CONFIG_H, _PTX]:
        if p.exists():
            p.unlink()


# ---------------------------------------------------------------------------
# Physics / geometry constants (match env_karman_2000x600.py)
# ---------------------------------------------------------------------------
L0 = 20.0
D_CYL = L0
U0 = 0.01
RADIUS = L0 / 2.0

NX = 2000
NY = 600
CENTER_Y = float(NY - 1) / 2.0

DIST_X = 600.0
PINBALL_FRONT_X = 1000.0
PINBALL_REAR_X = 1026.0
SENSOR_X = 1200.0

SI = 800
FIFO_LEN = 150
CONV_LEN = 30
ACTION_SCALE = 8.0
ACTION_BIAS = np.array([0.0, -4.0, 4.0], dtype=np.float32)

# Current env norm constants (to assess health)
FORCE_SCALE = np.float32(0.005)
SENS_SCALE = np.float32(U0)

WARMUP_STEPS = int(4.0 * NX / U0)
CFG_PATH = str(_REPO / "configs" / "config_lbm_karman_2000x600.json")
SENSOR_CC = 78.0

N_MEASURE = 100  # steps to measure per stage (after FIFO filled)

OUT_DIR = Path(__file__).resolve().parent / "output" / "stage_baseline_2000x600"


# ---------------------------------------------------------------------------
# DTW utilities (match env_karman_2000x600.py exactly)
# ---------------------------------------------------------------------------
def calc_lag(target: np.ndarray, state: np.ndarray) -> int:
    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: np.ndarray, state: np.ndarray) -> float:
    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
    return float(1.0 - dtw[n, m] / float(n))


def compute_similarity(target_states: np.ndarray, fifo_states: np.ndarray,
                        conv_len: int = CONV_LEN) -> float:
    target = np.asarray(target_states, dtype=np.float64)
    state = np.asarray(fifo_states, dtype=np.float64)
    id_sens = 3  # s1_uy (center sensor y-velocity) — matches 2000x600 env
    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)
    return float(sim / 6.0)


# ---------------------------------------------------------------------------
# Context guard
# ---------------------------------------------------------------------------
class CtxGuard:
    def __init__(self, sim):
        self.sim = sim

    def __enter__(self):
        if self.sim is not None:
            self.sim.ctx._ctx.push()

    def __exit__(self, *exc):
        if self.sim is not None:
            self.sim.ctx._ctx.pop()
        return False


def gpu_block(sim, fn):
    with CtxGuard(sim):
        fn()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def action_to_omega(action_norm: np.ndarray) -> np.ndarray:
    """[-1,1] action -> lattice omega (new kernel sign convention)."""
    sv = (np.asarray(action_norm, dtype=np.float32) * ACTION_SCALE + ACTION_BIAS) * U0
    return -sv / RADIUS


def read_obs(sim, dist_id, sensor_ids, pinball_ids) -> np.ndarray:
    """14-dim: [dist_fx,fy, 6x raw_sensor, 6x raw_force]."""
    obs = list(sim.read_force(dist_id, normalize=True))
    for sid in sensor_ids:
        s = sim.read_sensor(sid, normalize=True)
        obs.extend([float(s[0]), float(s[1])])
    for pid in pinball_ids:
        obs.extend(sim.read_force(pid, normalize=True))
    return np.array(obs, dtype=np.float32)


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


def cd_cl_from_obs(obs_slice, force_scale=FORCE_SCALE):
    """obs_slice = [6 sensor (legacy-equiv), 6 force (raw)].
    Returns (cd, cl) normalized by force_scale.
    """
    forces = obs_slice[6:12] / force_scale
    cd = (forces[0] + forces[2] + forces[4]) / 3.0
    cl = (forces[1] + forces[3] + forces[5]) / 3.0
    return float(cd), float(cl)


# ---------------------------------------------------------------------------
# Stage runner
# ---------------------------------------------------------------------------
def run_stage(sim, dist_id, sensor_ids, pinball_ids, target_states,
              omega_vec, stage_name, n_measure=N_MEASURE):
    """Run one stage: fill FIFO with given omega, then measure n_measure steps.

    Returns dict with per-step obs_slice, cd, cl, sim, and FIFO.
    """
    print(f"  [{stage_name}] Filling FIFO ({FIFO_LEN} x SI)...", end=" ", flush=True)
    t0 = time.perf_counter()
    fifo = deque(maxlen=FIFO_LEN)
    with CtxGuard(sim):
        set_omega(sim, pinball_ids, omega_vec)
        for _ in range(FIFO_LEN):
            sim.run(SI, zero_obs=True)
            obs = read_obs(sim, dist_id, sensor_ids, pinball_ids)
            sl = obs[2:14].copy()
            sl[0:6] *= SENSOR_CC  # legacy-equiv for DTW
            fifo.append(sl)
    print(f"done ({time.perf_counter()-t0:.0f}s).")

    print(f"  [{stage_name}] Measuring ({n_measure} steps)...", end=" ", flush=True)
    t0 = time.perf_counter()
    obs_slices = []
    cds, cls, sims = [], [], []
    with CtxGuard(sim):
        set_omega(sim, pinball_ids, omega_vec)
        for _ in range(n_measure):
            sim.run(SI, zero_obs=True)
            obs = read_obs(sim, dist_id, sensor_ids, pinball_ids)
            sl = obs[2:14].copy()
            sl[0:6] *= SENSOR_CC
            fifo.append(sl)
            obs_slices.append(sl.copy())
            cd, cl = cd_cl_from_obs(sl)
            cds.append(cd)
            cls.append(cl)
            # Sim requires >= CONV_LEN*2 samples in fifo
            sim_val = compute_similarity(target_states, np.array(list(fifo)))
            sims.append(sim_val)
    print(f"done ({time.perf_counter()-t0:.0f}s).")

    return {
        "obs_slices": np.array(obs_slices, dtype=np.float32),  # (N, 12)
        "cds": np.array(cds, dtype=np.float32),
        "cls": np.array(cls, dtype=np.float32),
        "sims": np.array(sims, dtype=np.float32),
        "fifo": np.array(list(fifo), dtype=np.float32),  # (FIFO_LEN, 12)
    }


# ---------------------------------------------------------------------------
# Norm health
# ---------------------------------------------------------------------------
def norm_health(obs_slices, force_scale=FORCE_SCALE, sens_scale=SENS_SCALE):
    """Assess obs normalization health.

    obs_slices: (N, 12) where [0:6]=sensor (legacy-equiv), [6:12]=force (raw).
    Returns dict with per-dim stats after normalization.
    """
    # Normalize the same way env does (but env uses raw sensor, not legacy-equiv;
    # for norm health we check the RAW values that env actually normalizes)
    # Note: env _normalize_obs uses obs_slice[0:6] (raw sensor) / SENS_SCALE
    # and obs_slice[6:12] (raw force) / FORCE_SCALE.
    # But our obs_slices here have sensor already *CC. For norm health we need RAW.
    # We stored sensor*CC for DTW; for norm health we divide back.
    raw = obs_slices.copy()
    raw[:, 0:6] = raw[:, 0:6] / SENSOR_CC  # back to raw (area-normalized) sensor

    forces_n = raw[:, 6:12] / force_scale
    sens_n = raw[:, 0:6] / sens_scale
    normed = np.hstack([forces_n, sens_n])  # (N, 12) [forces(6), sens(6)] — env order

    names = ["f_front_fx", "f_front_fy", "f_top_fx", "f_top_fy", "f_bot_fx", "f_bot_fy",
             "s0_ux", "s0_uy", "s1_ux", "s1_uy", "s2_ux", "s2_uy"]
    stats = {}
    clip_count = 0
    total = normed.size
    for i, name in enumerate(names):
        col = normed[:, i]
        clipped = np.sum(np.abs(col) > 1.0)
        clip_count += clipped
        stats[name] = {
            "mean": float(np.mean(col)),
            "std": float(np.std(col)),
            "min": float(np.min(col)),
            "max": float(np.max(col)),
            "clip_ratio": float(clipped / len(col)),
        }
    stats["_overall"] = {
        "clip_ratio_total": float(clip_count / total),
        "force_clip": float(np.sum(np.abs(forces_n) > 1.0) / forces_n.size),
        "sens_clip": float(np.sum(np.abs(sens_n) > 1.0) / sens_n.size),
    }
    return stats, normed


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
    parser = argparse.ArgumentParser(description="Phase 0 Baseline Measure")
    parser.add_argument("--device-id", type=int, default=0)
    args = parser.parse_args()

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    log_path = OUT_DIR / "phase0.log"

    def log(msg, **kwargs):
        line = f"[{time.strftime('%H:%M:%S')}] {msg}"
        print(line, **kwargs)
        with open(log_path, "a") as f:
            f.write(line + "\n")
            f.flush()

    log(f"=== Phase 0 Baseline Measure (2000x600, device={args.device_id}) ===")
    log(f"  Output: {OUT_DIR}")

    # ---- Step 1: Record target (dist_cyl + 3 sensors, no pinball) ----
    log("  Step 1: Recording target signal...")
    _clean_cache()
    sim = Simulation(lbm_config_path=CFG_PATH, device_id=args.device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    dist_id = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * 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)
    sensor_ids = [s0, s1, s2]
    sim.initialize()

    log(f"  Warmup target sim ({WARMUP_STEPS} steps)...", end=" ", flush=True)
    t0 = time.perf_counter()
    with CtxGuard(sim):
        sim.run(WARMUP_STEPS, zero_obs=True)
    log(f"done ({time.perf_counter()-t0:.0f}s).")

    target = np.zeros((FIFO_LEN, 6), dtype=np.float32)
    with CtxGuard(sim):
        for i in range(FIFO_LEN):
            sim.run(SI, zero_obs=True)
            target[i] = [
                sim.read_sensor(s0, normalize=True)[0] * SENSOR_CC,
                sim.read_sensor(s0, normalize=True)[1] * SENSOR_CC,
                sim.read_sensor(s1, normalize=True)[0] * SENSOR_CC,
                sim.read_sensor(s1, normalize=True)[1] * SENSOR_CC,
                sim.read_sensor(s2, normalize=True)[0] * SENSOR_CC,
                sim.read_sensor(s2, normalize=True)[1] * SENSOR_CC,
            ]
    sim.close()
    log(f"  Target recorded. s1_uy range: [{target[:,3].min():.4f}, {target[:,3].max():.4f}], std={target[:,3].std():.4f}")

    # ---- Step 2: Create training sim with ALL objects ----
    log("  Step 2: Creating training sim (all 7 objects)...")
    _clean_cache()
    sim = Simulation(lbm_config_path=CFG_PATH, device_id=args.device_id)
    sim._assert_object_count_contract = lambda *a, **kw: None
    dist_id = sim.add_body("circle", center=(DIST_X, CENTER_Y, 0.0), radius=1.0 * L0)
    sensor_ids = [
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y + 40.0, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y, 0.0), radius=5.0),
        sim.add_body("sensor", center=(SENSOR_X, CENTER_Y - 40.0, 0.0), radius=5.0),
    ]
    sim.add_body("circle", center=(PINBALL_FRONT_X, CENTER_Y, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y + 15.0, 0.0), radius=RADIUS)
    sim.add_body("circle", center=(PINBALL_REAR_X, CENTER_Y - 15.0, 0.0), radius=RADIUS)
    sim.initialize()
    pinball_ids = [4, 5, 6]

    log(f"  Warmup pinball sim ({WARMUP_STEPS} steps)...", end=" ", flush=True)
    t0 = time.perf_counter()
    with CtxGuard(sim):
        sim.run(WARMUP_STEPS, zero_obs=True)
    log(f"done ({time.perf_counter()-t0:.0f}s).")

    # Snapshot at zero-action state (for restoring between stages)
    with CtxGuard(sim):
        sim.snapshot()
    log("  Snapshot saved (zero-action pinball state).")

    # ---- Step 3: Stage 0 — zero rotation ----
    log("  Step 3: Stage 0 (zero rotation)...")
    with CtxGuard(sim):
        sim.restore()
    zero_omega = np.zeros(3, dtype=np.float32)
    stage0 = run_stage(sim, dist_id, sensor_ids, pinball_ids, target,
                       zero_omega, "Stage0")

    # ---- Step 4: Stage 1 — bias [0,-4,4]*U0 ----
    log("  Step 4: Stage 1 (bias [0,-4,4]*U0)...")
    with CtxGuard(sim):
        sim.restore()
    # bias action_norm = [0,0,0] -> omega = -(0*8 + [0,-4,4])*U0 / R
    bias_omega = action_to_omega(np.zeros(3, dtype=np.float32))
    log(f"  bias_omega = {bias_omega}")
    stage1 = run_stage(sim, dist_id, sensor_ids, pinball_ids, target,
                       bias_omega, "Stage1")

    sim.close()

    # ---- Step 5: Analysis ----
    log("  Step 5: Analysis...")

    # 5a. Cd/Cl/Sim summary per stage
    def stage_summary(stage, name):
        cd_mean, cd_std = float(np.mean(stage["cds"])), float(np.std(stage["cds"]))
        cl_mean, cl_std = float(np.mean(stage["cls"])), float(np.std(stage["cls"]))
        sim_mean, sim_std = float(np.mean(stage["sims"])), float(np.std(stage["sims"]))
        # Current reward formula
        r_cd = float(np.mean(np.exp(-np.abs(stage["cds"] * 20))))
        r_cl = float(np.mean(np.exp(-np.abs(stage["cls"] * 80))))
        r_sim = float(np.mean(np.exp(-10 * np.abs(stage["sims"] - 1))))
        reward = float(np.mean(np.minimum(0.3 * np.exp(-np.abs(stage["cds"] * 20)) +
                                          0.4 * np.exp(-np.abs(stage["cls"] * 80)) +
                                          0.3 * np.exp(-10 * np.abs(stage["sims"] - 1)), 1.0)))
        return {
            "cd_mean": cd_mean, "cd_std": cd_std,
            "cl_mean": cl_mean, "cl_std": cl_std,
            "sim_mean": sim_mean, "sim_std": sim_std,
            "r_cd": r_cd, "r_cl": r_cl, "r_sim": r_sim, "reward": reward,
        }

    s0_sum = stage_summary(stage0, "Stage0")
    s1_sum = stage_summary(stage1, "Stage1")

    # 5b. Force range (for "combined range" normalization design)
    def force_range(stage):
        f = stage["obs_slices"][:, 6:12]  # raw forces
        return {
            "max_abs_per_dim": [float(np.max(np.abs(f[:, i]))) for i in range(6)],
            "max_abs_combined": float(np.max(np.abs(f))),
            "mean_abs_combined": float(np.mean(np.abs(f))),
            "std_combined": float(np.std(f)),
        }

    def sensor_range(stage):
        s = stage["obs_slices"][:, 0:6]  # legacy-equiv sensors
        return {
            "max_abs_per_dim": [float(np.max(np.abs(s[:, i]))) for i in range(6)],
            "max_abs_combined": float(np.max(np.abs(s))),
            "mean_abs_combined": float(np.mean(np.abs(s))),
            "std_combined": float(np.std(s)),
        }

    fr0, fr1 = force_range(stage0), force_range(stage1)
    sr0, sr1 = sensor_range(stage0), sensor_range(stage1)

    # 5c. Norm health (using current FORCE_SCALE=0.005, SENS_SCALE=U0)
    nh0, normed0 = norm_health(stage0["obs_slices"])
    nh1, normed1 = norm_health(stage1["obs_slices"])

    # ---- Print summary ----
    lines = []
    lines.append("=" * 100)
    lines.append("PHASE 0 BASELINE MEASURE SUMMARY (2000x600)")
    lines.append("=" * 100)

    lines.append("\n--- Target signal stats ---")
    names_s = ["s0_ux", "s0_uy", "s1_ux", "s1_uy", "s2_ux", "s2_uy"]
    for i, n in enumerate(names_s):
        lines.append(f"  {n:>8}: mean={target[:,i].mean():.4f}, std={target[:,i].std():.4f}, "
                     f"range=[{target[:,i].min():.4f}, {target[:,i].max():.4f}]")

    lines.append("\n--- Stage comparison: Cd / Cl / Sim / Reward ---")
    lines.append(f"  {'Stage':>8} {'Cd_mean':>10} {'Cd_std':>10} {'Cl_mean':>10} {'Cl_std':>10} "
                 f"{'Sim_mean':>10} {'Sim_std':>10} {'r_cd':>8} {'r_cl':>8} {'r_sim':>8} {'reward':>8}")
    for name, s in [("Stage0", s0_sum), ("Stage1", s1_sum)]:
        lines.append(f"  {name:>8} {s['cd_mean']:>10.4f} {s['cd_std']:>10.4f} {s['cl_mean']:>10.4f} "
                     f"{s['cl_std']:>10.4f} {s['sim_mean']:>10.4f} {s['sim_std']:>10.4f} "
                     f"{s['r_cd']:>8.4f} {s['r_cl']:>8.4f} {s['r_sim']:>8.4f} {s['reward']:>8.4f}")

    lines.append("\n--- Force range (raw, for combined-range norm design) ---")
    fnames = ["front_fx", "front_fy", "top_fx", "top_fy", "bot_fx", "bot_fy"]
    for name, fr in [("Stage0", fr0), ("Stage1", fr1)]:
        lines.append(f"  {name}: max_abs_combined={fr['max_abs_combined']:.6f}, "
                     f"mean_abs={fr['mean_abs_combined']:.6f}, std={fr['std_combined']:.6f}")
        for i, fn in enumerate(fnames):
            lines.append(f"    {fn:>10}: max_abs={fr['max_abs_per_dim'][i]:.6f}")

    lines.append("\n--- Sensor range (legacy-equiv, for combined-range norm design) ---")
    for name, sr in [("Stage0", sr0), ("Stage1", sr1)]:
        lines.append(f"  {name}: max_abs_combined={sr['max_abs_combined']:.4f}, "
                     f"mean_abs={sr['mean_abs_combined']:.4f}, std={sr['std_combined']:.4f}")
        for i, sn in enumerate(names_s):
            lines.append(f"    {sn:>8}: max_abs={sr['max_abs_per_dim'][i]:.4f}")

    lines.append("\n--- Obs norm health (current: FORCE_SCALE=0.005, SENS_SCALE=U0=0.01) ---")
    nh_names = ["f_front_fx", "f_front_fy", "f_top_fx", "f_top_fy", "f_bot_fx", "f_bot_fy",
                "s0_ux", "s0_uy", "s1_ux", "s1_uy", "s2_ux", "s2_uy"]
    for name, nh in [("Stage0", nh0), ("Stage1", nh1)]:
        lines.append(f"\n  {name} (env order: forces[6], sensors[6]):")
        lines.append(f"    {'dim':>14} {'mean':>10} {'std':>10} {'min':>10} {'max':>10} {'clip%':>8}")
        for i, n in enumerate(nh_names):
            d = nh[n]
            lines.append(f"    {n:>14} {d['mean']:>10.4f} {d['std']:>10.4f} {d['min']:>10.4f} "
                         f"{d['max']:>10.4f} {d['clip_ratio']*100:>7.1f}%")
        ov = nh["_overall"]
        lines.append(f"    {'OVERALL':>14} clip_ratio_total={ov['clip_ratio_total']*100:.1f}%, "
                     f"force_clip={ov['force_clip']*100:.1f}%, sens_clip={ov['sens_clip']*100:.1f}%")

    lines.append("\n--- Gradient assessment ---")
    lines.append(f"  Cd:  Stage0={s0_sum['cd_mean']:.4f} -> Stage1={s1_sum['cd_mean']:.4f} "
                 f"(Δ={s1_sum['cd_mean']-s0_sum['cd_mean']:+.4f})")
    lines.append(f"  Cl:  Stage0={s0_sum['cl_mean']:.4f} -> Stage1={s1_sum['cl_mean']:.4f} "
                 f"(Δ={s1_sum['cl_mean']-s0_sum['cl_mean']:+.4f})")
    lines.append(f"  Sim: Stage0={s0_sum['sim_mean']:.4f} -> Stage1={s1_sum['sim_mean']:.4f} "
                 f"(Δ={s1_sum['sim_mean']-s0_sum['sim_mean']:+.4f})")
    lines.append(f"  Reward: Stage0={s0_sum['reward']:.4f} -> Stage1={s1_sum['reward']:.4f} "
                 f"(Δ={s1_sum['reward']-s0_sum['reward']:+.4f})")

    summary = "\n".join(lines)
    print("\n" + summary)
    with open(OUT_DIR / "summary.txt", "w") as f:
        f.write(summary + "\n")

    # Save raw data
    with open(OUT_DIR / "stage_data.pkl", "wb") as f:
        pickle.dump({
            "target": target,
            "stage0": stage0, "stage1": stage1,
            "s0_summary": s0_sum, "s1_summary": s1_sum,
            "force_range": {"stage0": fr0, "stage1": fr1},
            "sensor_range": {"stage0": sr0, "stage1": sr1},
        }, f)

    with open(OUT_DIR / "norm_health.json", "w") as f:
        json.dump({"stage0": nh0, "stage1": nh1}, f, indent=2)

    log(f"\nPhase 0 complete. Files in {OUT_DIR}:")
    log(f"  summary.txt, stage_data.pkl, norm_health.json, phase0.log")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())