# analysis_crossre/scripts/phase3_validate.py
"""Phase 3: closed-loop validation using cross-Re SINDy control law.

Usage::

    conda run -n pycuda_3_10 python phase3_validate.py \\
        --device 2 --out-dir output/analysis_crossre/sindy_val

    conda run -n pycuda_3_10 python phase3_validate.py \\
        --validate-re 35,70,150 --device 2

    conda run -n pycuda_3_10 python phase3_validate.py \\
        --baseline-only --validate-re 35 --device 2
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import deque

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 utils import (
    nu_from_re,
    load_legacy_configs,
    build_karman_cloak_env,
    add_pinball,
    build_observation,
    scale_action,
    action_to_physical,
    compute_dimensionless,
    compute_v3_symbols,
    save_vorticity_png,
    vorticity_from_ddf,
    compute_similarity,
)
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_CASES_VALIDATION,
    RE_LABEL_MAP,
)

DATA_TYPE = np.float32


def load_cross_re_coef(sindy_results_path: str, threshold: float) -> dict:
    """Load v3 cross-Re coefficients.

    Returns dict ``{cylinder_name: {"coef": np.ndarray, "feat_names": list, "has_bias": bool}}``
    """
    with open(sindy_results_path) as f:
        data = json.load(f)

    cross = data["cross_re"]
    coefs = {}
    for ch_entry in cross["channels"]:
        name = ch_entry["cylinder"]
        feat_names = ch_entry["feature_names"]
        coef_full = np.array(ch_entry["best_coef"], dtype=np.float64)
        has_bias = ch_entry["has_bias"]

        scale = np.max(np.abs(coef_full))
        if scale > 0 and threshold > 0:
            mask = np.abs(coef_full) / scale >= threshold
        else:
            mask = np.ones_like(coef_full, dtype=bool)
        coef = coef_full * mask

        nz = int(np.sum(mask))
        print(f"  {name}: total={len(coef_full)} nz={nz} threshold={threshold} "
              f"R2={ch_entry['best']['r2']:.4f}")

        coefs[name] = {"coef": coef, "feat_names": feat_names,
                       "has_bias": has_bias, "nz": nz, "r2": ch_entry["best"]["r2"]}
    return coefs


def predict_omega_v3(
    obs_slice: np.ndarray,
    actions_prev: np.ndarray,
    actions_prev2: np.ndarray,
    coefs: dict,
    mu: float,
    u0: float = 0.01,
) -> np.ndarray:
    """Predict physical omega using v3 dimensionless features.

    Front: no bias term.
    Bottom/Top: with bias term.
    All 3 independently (no exchange symmetry constraint).

    Parameters
    ----------
    obs_slice : (12,)  raw [sensor(6), force(6)] in lattice units
    actions_prev : (3,)  omega(t-1)
    actions_prev2 : (3,)  omega(t-2)
    """
    sensors = obs_slice[0:6].astype(np.float64).reshape(1, 6)
    forces = obs_slice[6:12].astype(np.float64).reshape(1, 6)
    a_prev = actions_prev.astype(np.float64).reshape(1, 3)
    a_prev2 = actions_prev2.astype(np.float64).reshape(1, 3)

    # Dimensionless
    dim = compute_dimensionless(sensors, forces, u0=u0, d=20.0)

    # Build v3 features
    Theta_f, Theta_top, names = compute_v3_symbols(
        dim, a_prev, a_prev2, mu=mu, include_mu=(mu > 0))

    # Predict
    omega = np.zeros(3, dtype=np.float64)
    omega[0] = float(Theta_f[0] @ coefs["front"]["coef"])  # front (no bias)
    omega[1] = float(Theta_top[0] @ coefs["bottom"]["coef"])  # bottom
    omega[2] = float(Theta_top[0] @ coefs["top"]["coef"])  # top

    return omega


def run_sindy_controlled(
    re_code: int,
    coefs: dict,
    device_id: int,
    output_root: str,
    *,
    n_steps: int = 150,
) -> dict:
    """Run closed-loop validation with SINDy control law."""
    os.makedirs(output_root, exist_ok=True)

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

    # Build environment (same as Phase 1)
    cuda_cfg, field_cfg = load_legacy_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))

    # Phase 1: dist + sensors + target
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    target_states, _ = build_karman_cloak_env(
        ff, u0=U0, l0=20.0, sample_interval=SAMPLE_INTERVAL,
        fifo_len=FIFO_LEN, data_type=DATA_TYPE,
    )

    # Phase 2: pinball + norm
    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,
    )
    np.savez(os.path.join(output_root, "target.npz"), target_states=target_states)

    # --- Uncontrolled rollout ---
    print("  uncontrolled rollout ...")
    ff.restore_ddf()
    ff.apply_ddf()
    sens_unc, forc_unc = [], []
    for _ in range(n_steps):
        ff.run(SAMPLE_INTERVAL, np.zeros(7, dtype=DATA_TYPE))
        obs_slice = ff.obs.copy()[2:14]
        sens_unc.append(obs_slice[0:6])
        forc_unc.append(obs_slice[6:12])
    np.savez(os.path.join(output_root, "uncontrolled.npz"),
             sensors=np.array(sens_unc, dtype=np.float32),
             forces=np.array(forc_unc, dtype=np.float32))

    # Uncontrolled vorticity
    omega_unc = vorticity_from_ddf(ff, u0=U0)
    save_vorticity_png(os.path.join(output_root, "vorticity_uncontrolled.png"),
                       omega_unc, title=f"{label} uncontrolled")

    # --- SINDy controlled rollout ---
    print(f"  SINDy controlled rollout ({n_steps} steps) ...")
    ff.restore_ddf()
    ff.apply_ddf()

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

    sens_sc, forc_sc, omega_sc = [], [], []
    omega_bias = action_to_physical(
        np.zeros((1, 3), dtype=np.float32),
        scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0,
    ).flatten()
    actions_prev = omega_bias.copy()
    actions_prev2 = omega_bias.copy()

    for step in range(n_steps):
        obs_slice = fifo[-1] if len(fifo) > 0 else np.zeros(12, dtype=np.float32)

        omega_pred = predict_omega_v3(obs_slice, actions_prev, actions_prev2, coefs, mu, u0=U0)
        omega_sc.append(omega_pred.copy())

        # Convert action to legacy array and apply
        norm_action = (omega_pred / U0 - np.array(ACTION_BIAS, dtype=np.float64)) / ACTION_SCALE
        norm_action = np.clip(norm_action, -1.0, 1.0).astype(np.float32)
        action_arr = scale_action(
            norm_action,
            scale=ACTION_SCALE, bias=ACTION_BIAS, u0=U0, n_total_bodies=7,
        )
        ff.run(SAMPLE_INTERVAL, action_arr)

        obs_slice_new = ff.obs.copy()[2:14]
        fifo.append(obs_slice_new)
        sens_sc.append(obs_slice_new[0:6])
        forc_sc.append(obs_slice_new[6:12])
        actions_prev = omega_pred

    sens_sc_arr = np.array(sens_sc, dtype=np.float32)
    forc_sc_arr = np.array(forc_sc, dtype=np.float32)
    omega_sc_arr = np.array(omega_sc, dtype=np.float32)
    np.savez(os.path.join(output_root, "sindy_controlled.npz"),
             sensors=sens_sc_arr, forces=forc_sc_arr,
             omegas=omega_sc_arr)

    # Vorticity
    omega_vort = vorticity_from_ddf(ff, u0=U0)
    save_vorticity_png(os.path.join(output_root, "vorticity_sindy_controlled.png"),
                       omega_vort, title=f"{label} SINDy-controlled")

    # Similarity
    sim = compute_similarity(target_states, sens_sc_arr, CONV_LEN)
    print(f"  SINDy similarity: {sim:.4f}")

    del ff
    result = {"re_code": re_code, "mu": mu,
              "sindy_similarity": sim,
              "n_steps": n_steps}
    with open(os.path.join(output_root, "result.json"), "w") as f:
        json.dump(result, f, indent=2)
    return result


def main():
    ap = argparse.ArgumentParser(description="Phase 3: cross-Re SINDy validation")
    ap.add_argument("--sindy-results", type=str,
                    default=os.path.join(OUTPUT_DIR, "sindy", "sindy_results_v3.json"),
                    help="Path to Phase 2 SINDy results JSON (v3 dimensionless)")
    ap.add_argument("--validate-re", type=str, default="35,70,150",
                    help="Comma-separated validation Re codes")
    ap.add_argument("--device", type=int, default=0, help="GPU device ID")
    ap.add_argument("--steps", type=int, default=150,
                    help="Number of inference steps")
    ap.add_argument("--threshold", type=float, default=0.002,
                    help="SINDy sparsity threshold (default: 0.002)")
    ap.add_argument("--out-dir", type=str,
                    default=os.path.join(OUTPUT_DIR, "sindy_val"),
                    help="Output root for validation results")
    args = ap.parse_args()

    validate_re = [int(r) for r in args.validate_re.split(",")]
    os.makedirs(args.out_dir, exist_ok=True)

    # Load cross-Re coefficients
    print(f"\nLoading cross-Re coefficients from {args.sindy_results}")
    coefs = load_cross_re_coef(args.sindy_results, args.threshold)
    for name in ["front", "bottom", "top"]:
        print(f"  {name}: nz={coefs[name]['nz']}, R2={coefs[name]['r2']:.4f}, "
              f"threshold={args.threshold}")

    t_start = time.time()

    # Run for each validation Re
    for re_code in validate_re:
        out_dir = os.path.join(args.out_dir, f"re{re_code}")
        result = run_sindy_controlled(
            re_code, coefs, args.device, out_dir, n_steps=args.steps,
        )
        print(f"  Done: Re{re_code}  -> {out_dir}")

    elapsed = time.time() - t_start
    print(f"\nTotal time: {elapsed:.1f}s")
    return 0


if __name__ == "__main__":
    sys.exit(main())