# analysis_crossre/scripts/diagnose_equivariance.py
"""Phase A2-A3: diagnose PPO control-law equivariance under G operator.

Usage::

    conda run -n pycuda_3_10 python diagnose_equivariance.py --re 100 --device 0

    conda run -n pycuda_3_10 python diagnose_equivariance.py --re all --device 0

Output per Re: ``output/analysis_crossre/diagnostic/equivariance_re{re}.json``
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Dict, List, Tuple

import numpy as np

_PROJ = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _PROJ not in sys.path:
    sys.path.insert(0, _PROJ)
from LegacyCelerisLab import FlowField  # noqa: E402
from LegacyCelerisLab import utils as legacy_utils  # noqa: E402

from utils import (
    action_to_physical,
    compute_dimensionless,
    apply_G_x,
    apply_G_alpha,
    load_ppo_model,
    nu_from_re,
    load_legacy_configs,
    build_karman_cloak_env,
    add_pinball,
    build_observation,
    scale_action,
)
from cfg import (
    CONFIG_DIR,
    OUTPUT_DIR,
    MODEL_DIR,
    SAMPLE_INTERVAL,
    FIFO_LEN,
    CONV_LEN,
    S_DIM,
    A_DIM,
    ACTION_SCALE,
    ACTION_BIAS,
    U0,
    RE_CASES_TRAIN,
    RE_LABEL_MAP,
)

DATA_TYPE = np.float32


def diagnose_one_re(re_code: int, ppo_device: int, cfd_device: int, output_root: str) -> dict:
    """Run equivariance diagnosis for one Re case."""
    os.makedirs(output_root, exist_ok=True)

    nu = nu_from_re(re_code, u0=U0)
    mu = 2.0 / re_code
    label = RE_LABEL_MAP.get(re_code, f"Re{re_code}")
    print(f"\n{'='*60}")
    print(f"Diagnosing: {label}   nu={nu:.6f}   mu={mu:.6f}")
    print(f"{'='*60}")

    # Build full environment (dist + sensors + pinball)
    cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))
    ff = FlowField(field_cfg, cuda_cfg, device_id=cfd_device)

    # Stabilize and get to controlled state
    target_states, _ = build_karman_cloak_env(
        ff, u0=U0, l0=20.0, sample_interval=SAMPLE_INTERVAL,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
    )
    norm = add_pinball(
        ff, l0=20.0, u0=U0, sample_interval=SAMPLE_INTERVAL,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
        action_bias=ACTION_BIAS,
    )

    # Load PPO model
    model_path = None
    for rc, mn in RE_CASES_TRAIN:
        if rc == re_code:
            model_path = os.path.join(MODEL_DIR, "old", f"{mn}.zip")
            break
    if model_path is None or not os.path.isfile(model_path):
        return {"re_code": re_code, "error": f"No model for Re{re_code}"}

    model = load_ppo_model(model_path, device=f"cuda:{ppo_device}")
    model.set_random_seed(0)

    # Collect rollout data with PPO
    ff.restore_ddf()
    ff.apply_ddf()

    # Bias FIFO
    bias_action = scale_action(
        np.zeros(3, dtype=np.float32),
        scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0, n_total_bodies=7,
    )
    from collections import deque
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.context.push()
        try:
            ff.run(SAMPLE_INTERVAL, bias_action)
        finally:
            ff.context.pop()
        fifo.append(ff.obs.copy()[2:14])

    n_steps = 150
    obs_hist = np.zeros((n_steps, 12), dtype=np.float64)
    alpha_hist = np.zeros((n_steps, 3), dtype=np.float64)
    obs = np.zeros(S_DIM, dtype=np.float32)

    for step in range(n_steps):
        action, _ = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()

        # Convert to physical
        action_arr = scale_action(
            action, scale=ACTION_SCALE, bias=ACTION_BIAS,
            u0=U0, n_total_bodies=7,
        )
        ff.context.push()
        try:
            ff.run(SAMPLE_INTERVAL, action_arr)
        finally:
            ff.context.pop()

        obs_slice = ff.obs.copy()[2:14]
        fifo.append(obs_slice)
        alpha = action_to_physical(
            action.reshape(1, 3), scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
        ).flatten()

        obs_hist[step] = obs_slice
        alpha_hist[step] = alpha
        obs = build_observation(obs_slice, norm)

    del ff

    # ---- Equivariance diagnosis ----
    dim = compute_dimensionless(obs_hist[:, 0:6], obs_hist[:, 6:12], u0=U0, d=20.0)

    # Compute memory terms
    a_prev = np.zeros_like(alpha_hist)
    a_prev2 = np.zeros_like(alpha_hist)
    a_prev[1:] = alpha_hist[:-1]
    a_prev2[2:] = alpha_hist[:-2]

    # Diagnostic 1: front bias check (mean of alpha_F)
    mean_alpha_F = float(np.mean(alpha_hist[:, 0]))
    std_alpha_F = float(np.std(alpha_hist[:, 0]))
    front_bias_score = abs(mean_alpha_F) / (std_alpha_F + 1e-12)

    # Diagnostic 2: check front equivariance
    # For each point, compute PPO(Gx) by feeding G-transformed obs through model
    eq_front_errors = []
    eq_exchange_b_errors = []
    eq_exchange_t_errors = []
    eq_front_noise_floor = []

    for t in range(2, n_steps):
        # Get original obs and Gx
        Gx = apply_G_x(
            dim["u_hat_B"][t:t+1], dim["u_hat_C"][t:t+1], dim["u_hat_T"][t:t+1],
            dim["v_hat_B"][t:t+1], dim["v_hat_C"][t:t+1], dim["v_hat_T"][t:t+1],
            dim["Cd_F"][t:t+1], dim["Cd_T"][t:t+1], dim["Cd_B"][t:t+1],
            dim["Cl_F"][t:t+1], dim["Cl_T"][t:t+1], dim["Cl_B"][t:t+1],
            a_prev[t:t+1, 0], a_prev[t:t+1, 2], a_prev[t:t+1, 1],
            a_prev2[t:t+1, 0] - a_prev[t:t+1, 0],
            a_prev2[t:t+1, 2] - a_prev[t:t+1, 2],
            a_prev2[t:t+1, 1] - a_prev[t:t+1, 1],
        )

        # Build Gx observation for PPO: we need the normalized obs
        # The Gx in raw sensor/force space requires inverting the dimensionless transform
        # Actually easier: compute what PPO would predict for the G state
        # by transforming the raw obs and feeding it

        # Build raw obs corresponding to Gx
        raw_Gx = np.zeros(12, dtype=np.float64)
        # Sensors: reorder + sign flip
        # Original raw: [s0_ux, s0_uy, s1_ux, s1_uy, s2_ux, s2_uy] = top, center, bottom
        # G: bottom->top, center->center, top->bottom
        raw_Gx[0] = obs_hist[t, 4]   # s0_ux <- s2_ux (bottom -> top, streamwise no sign)
        raw_Gx[1] = -obs_hist[t, 5]  # s0_uy <- -s2_uy (bottom -> top, cross sign flip)
        raw_Gx[2] = obs_hist[t, 2]   # s1_ux maintains (center)
        raw_Gx[3] = -obs_hist[t, 3]  # s1_uy = -s1_uy (center cross sign flip)
        raw_Gx[4] = obs_hist[t, 0]   # s2_ux <- s0_ux (top -> bottom)
        raw_Gx[5] = -obs_hist[t, 1]  # s2_uy <- -s0_uy (top -> bottom, cross sign flip)
        # Forces: reorder + sign
        # ordering: [front_fx, front_fy, bottom_fx, bottom_fy, top_fx, top_fy]
        # G: front_fx -> front_fx (no sign), front_fy -> -front_fy
        #    bottom <-> top
        raw_Gx[6] = obs_hist[t, 6]    # front_fx unchanged
        raw_Gx[7] = -obs_hist[t, 7]   # front_fy sign flip
        raw_Gx[8] = obs_hist[t, 10]   # bottom_fx <- top_fx
        raw_Gx[9] = -obs_hist[t, 11]  # bottom_fy <- -top_fy
        raw_Gx[10] = obs_hist[t, 8]   # top_fx <- bottom_fx
        raw_Gx[11] = -obs_hist[t, 9]  # top_fy <- -bottom_fy

        # Build normalized PPO observation from Gx
        obs_Gx = build_observation(raw_Gx, norm)

        # Predict action for Gx
        action_Gx, _ = model.predict(obs_Gx, deterministic=True)
        action_Gx = action_Gx.astype(np.float32).flatten()
        alpha_Gx = action_to_physical(
            action_Gx.reshape(1, 3), scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
        ).flatten()

        # What equivariance says Gx should produce (with CORRECTED G)
        # G([aF, aT, aB]) = [-aF, -aB, -aT]
        alpha_Gx_expected = apply_G_alpha(alpha_hist[t])

        # Front error: PPO(Gx)[0] should == G(PPO(x))[0] = -aF(x)
        eq_front_errors.append(abs(float(alpha_Gx[0]) - float(alpha_Gx_expected[0])))

        # Rear error (CORRECTED): PPO(Gx)[1] should == G(PPO(x))[1] = -aT(x)
        # PPO(Gx)[2] should == G(PPO(x))[2] = -aB(x)
        # Previously this incorrectly checked alpha_B(x) == alpha_T(Gx)
        eq_exchange_b_errors.append(abs(float(alpha_Gx[1]) - float(alpha_Gx_expected[1])))
        eq_exchange_t_errors.append(abs(float(alpha_Gx[2]) - float(alpha_Gx_expected[2])))

        # Noise floor: difference between same-state replicate predictions
        # (we approximate by checking prediction consistency)
        action2, _ = model.predict(obs_Gx, deterministic=True)
        alpha_Gx2 = action_to_physical(
            action2.reshape(1, 3), scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
        ).flatten()
        eq_front_noise_floor.append(abs(float(alpha_Gx2[0]) - float(alpha_Gx[0])))

    eq_front_errors = np.array(eq_front_errors)
    eq_exchange_b = np.array(eq_exchange_b_errors)
    eq_exchange_t = np.array(eq_exchange_t_errors)
    eq_noise = np.array(eq_front_noise_floor)

    # Scale equivariance errors by action range for relative measure
    alpha_range = float(np.max(np.abs(alpha_hist[2:])))
    rel_front_err = float(np.mean(eq_front_errors) / (alpha_range + 1e-12))
    rel_exchange_b_err = float(np.mean(eq_exchange_b) / (alpha_range + 1e-12))
    rel_exchange_t_err = float(np.mean(eq_exchange_t) / (alpha_range + 1e-12))

    # Combined rear error (max of bottom and top)
    rel_exchange_err = max(rel_exchange_b_err, rel_exchange_t_err)

    # Diagnostic 3: cross-correlation between alpha_T and -alpha_B
    if len(alpha_hist) > 10:
        # After initial transient
        tail = n_steps // 2
        corr_TB = float(np.corrcoef(alpha_hist[tail:, 2], -alpha_hist[tail:, 1])[0, 1])
    else:
        corr_TB = float("nan")

    result = {
        "re_code": re_code,
        "mu": mu,
        "n_samples": n_steps,
        "alpha_range": alpha_range,
        "front_bias": {
            "mean_alpha_F": mean_alpha_F,
            "std_alpha_F": std_alpha_F,
            "bias_over_std": front_bias_score,
            "bias_significant": front_bias_score > 2.0,
        },
        "equivariance_front": {
            "mean_abs_error": float(np.mean(eq_front_errors)),
            "max_abs_error": float(np.max(eq_front_errors)),
            "relative_error": rel_front_err,
            "noise_floor": float(np.mean(eq_noise)),
            "signal_to_noise": float(np.mean(eq_front_errors) / (np.mean(eq_noise) + 1e-12)),
        },
        "equivariance_rear_bottom": {
            "mean_abs_error": float(np.mean(eq_exchange_b)),
            "max_abs_error": float(np.max(eq_exchange_b)),
            "relative_error": rel_exchange_b_err,
        },
        "equivariance_rear_top": {
            "mean_abs_error": float(np.mean(eq_exchange_t)),
            "max_abs_error": float(np.max(eq_exchange_t)),
            "relative_error": rel_exchange_t_err,
        },
        "top_bottom_correlation": {
            "corr_alphaT_vs_negAlphaB": corr_TB,
        },
        "equivariance_verdict": "PASS" if (rel_front_err < 0.20 and rel_exchange_err < 0.20) else "REVIEW",
    }

    print(f"  Front bias: mean_alpha_F={mean_alpha_F:.6f} |bias|/std={front_bias_score:.3f}")
    print(f"  Front equiv err: mean={np.mean(eq_front_errors):.6f} rel={rel_front_err:.3%}")
    print(f"  Rear-bot err:    mean={np.mean(eq_exchange_b):.6f} rel={rel_exchange_b_err:.3%}")
    print(f"  Rear-top err:    mean={np.mean(eq_exchange_t):.6f} rel={rel_exchange_t_err:.3%}")
    print(f"  T vs -B corr: {corr_TB:.4f}")
    print(f"  Verdict: {result['equivariance_verdict']}")

    with open(os.path.join(output_root, f"equivariance_re{re_code}.json"), "w") as f:
        json.dump(result, f, indent=2)
    print(f"  Saved to {output_root}/equivariance_re{re_code}.json")

    return result


def main():
    ap = argparse.ArgumentParser(description="Equivariance diagnosis for PPO cloak control")
    ap.add_argument("--re", type=str, default="all",
                    help='Re case: 50,100,200,400, or "all"')
    ap.add_argument("--device", type=int, default=0, help="GPU device for PPO model")
    ap.add_argument("--cfd-device", type=int, default=2, help="GPU device for CFD simulation")
    args = ap.parse_args()

    if args.re.lower() == "all":
        re_list = [rc for rc, _ in RE_CASES_TRAIN]
    else:
        re_list = [int(args.re)]

    # Store device args for use in diagnose_one_re
    device_id = args.device
    cfd_device = args.cfd_device

    diag_root = os.path.join(OUTPUT_DIR, "diagnostic")
    os.makedirs(diag_root, exist_ok=True)

    all_results = []
    for re_code in re_list:
        res = diagnose_one_re(re_code, device_id, cfd_device, diag_root)
        all_results.append(res)

    summary = {
        "summary": {
            "equivariance_verdicts": {r["re_code"]: r.get("equivariance_verdict", "ERROR")
                                      for r in all_results}
        },
        "details": all_results,
    }
    with open(os.path.join(diag_root, "equivariance_summary.json"), "w") as f:
        json.dump(summary, f, indent=2)
    print(f"\nSummary saved to {diag_root}/equivariance_summary.json")


if __name__ == "__main__":
    main()