"""CCD analysis pipeline: POD + force/action CCD.

New data format (fields_aligned.npz + phase_plan.json).
Target-only POD basis. Per-force observable (primary=SigmaFy).
Short Q_delay=6 for force/action. 1.5L flagged as special_mechanism.

Usage:
    conda run -n pycuda_3_10 python ccd/run_ccd.py

Requires fields_aligned.npz and phase_plan.json in data/ directories.
"""
from __future__ import annotations

import json
import os
import sys
import time

import numpy as np

_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from CCD_analysis.configs import DATA_DIR, SCENES, NX, NY
from CCD_analysis.utils.resampling import (
    compute_pod, cumulative_energy, e95_index,
    compute_reduced_ccd,
    load_aligned_fields, make_force_obs,
    build_field_matrix, project_into_basis,
)

# -- Protocol constants ---------------------------------------------------
R_CANDIDATES = [6, 8, 10]
CCD_Q = 6  # short, near-synchronous window for force/action
DIAMETERS_MAIN = [0.75, 1.0]
DIAMETERS_ALL = [0.75, 1.0, 1.5]
CV_T_RELAXED = 0.12


# ---------------------------------------------------------------------------
# Preflight check (built-in)
# ---------------------------------------------------------------------------

def preflight(scene_name: str) -> dict:
    """Load and verify one scene's data. Returns meta or raises."""
    cfg = SCENES[scene_name]
    scene_id = cfg["scene_id"]
    data_dir = os.path.join(DATA_DIR, scene_id, scene_name)

    # Check fields_aligned.npz
    fa_path = os.path.join(data_dir, "fields_aligned.npz")
    if not os.path.isfile(fa_path):
        raise FileNotFoundError(f"{fa_path} not found")

    fd = np.load(fa_path)
    ux = fd["ux"]
    print(f"  {scene_name}: fields_aligned ux shape={ux.shape}  "
          f"(expect ({cfg.get('n_cycles', 4) * cfg.get('n_pts', 24)}, {NX}, {NY}))",
          flush=True)
    fd.close()

    # Check phase_plan.json
    plan_path = os.path.join(DATA_DIR, "resampled", scene_name, "phase_plan.json")
    if not os.path.isfile(plan_path):
        raise FileNotFoundError(f"{plan_path} not found")

    import json
    with open(plan_path) as f:
        plan = json.load(f)

    n_total = plan["n_cycles"] * plan["n_pts"]
    if n_total != ux.shape[0]:
        print(f"    WARNING: phase_plan has {n_total} snapshots but fields has {ux.shape[0]}",
              flush=True)

    gate = plan["gate"]
    cv_t = plan["CV_T"]
    print(f"    gate={gate}, CV_T={cv_t:.4f}, "
          f"N_raw={plan['N_raw_per_cycle']:.1f}, rho={plan['rho_interp']:.2f}",
          flush=True)

    if gate not in ("strict", "relaxed") and cv_t is not None and cv_t > CV_T_RELAXED:
        print(f"    WARNING: gate='{gate}' — does not pass relaxed gate (CV_T <= {CV_T_RELAXED})",
              flush=True)

    # Check telemetry
    tele_found = False
    for p in [os.path.join(data_dir, "controlled.npz"), os.path.join(data_dir, "sensors.npz")]:
        if os.path.isfile(p):
            td = np.load(p)
            if "forces" in td:
                print(f"    forces: {td['forces'].shape}", flush=True)
            if "actions" in td:
                print(f"    actions: {td['actions'].shape}", flush=True)
            td.close()
            tele_found = True
            break
    if not tele_found:
        raise FileNotFoundError(f"No telemetry found in {data_dir}")

    return {
        "gate": gate,
        "CV_T": cv_t,
        "n_snapshots": ux.shape[0],
        "N_raw_per_cycle": plan.get("N_raw_per_cycle"),
    }


def compute_modal_overlap(W_dict: dict, diam: float, r: int,
                          obs_label: str = "force") -> list:
    """Compute pairwise modal overlaps for a given diameter and r."""
    keys = [k for k in W_dict
            if f"{diam}L_" in k and f"_{obs_label}_r{r}" in k]
    overlaps = []
    for i, ka in enumerate(keys):
        for kb in keys[i + 1:]:
            Wa, Wb = W_dict[ka], W_dict[kb]
            n = min(Wa.shape[1], Wb.shape[1], 5)
            for k in range(n):
                ov = float(abs(
                    Wa[:, k] / (np.linalg.norm(Wa[:, k]) + 1e-12) @
                    Wb[:, k] / (np.linalg.norm(Wb[:, k]) + 1e-12)
                ))
                overlaps.append({
                    "case_a": ka.split(f"_{obs_label}_r{r}")[0],
                    "case_b": kb.split(f"_{obs_label}_r{r}")[0],
                    "mode": k + 1,
                    "O": ov,
                })
    return overlaps


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------

def main():
    print("=" * 60, flush=True)
    print("CCD Pipeline (Round 5 — fields_aligned, target-only basis)", flush=True)
    print("=" * 60, flush=True)

    out_dir = os.path.join(DATA_DIR, "ccd")
    os.makedirs(out_dir, exist_ok=True)
    all_results = {}
    W_dict = {}  # for modal overlap

    # -- Preflight --
    print("\n--- Preflight check ---", flush=True)
    all_scenes = ["pinball"]
    for diam in DIAMETERS_ALL:
        all_scenes.append(f"target_cylinder_{diam}L")
        all_scenes.append(f"illusion_{diam}L")

    preflight_ok = {}
    for sn in all_scenes:
        try:
            meta = preflight(sn)
            preflight_ok[sn] = meta
            print(f"    OK", flush=True)
        except (FileNotFoundError, AssertionError, KeyError) as e:
            print(f"    FAILED: {e}", flush=True)
            preflight_ok[sn] = None

    # -- Load all data --
    print("\n--- Loading data ---", flush=True)
    data_cache = {}
    for sn in all_scenes:
        if preflight_ok.get(sn) is None:
            continue
        t0 = time.time()
        try:
            d = load_aligned_fields(sn)
            data_cache[sn] = d
            print(f"  {sn}: loaded ({len(d['ux'])} snapshots, "
                  f"{time.time() - t0:.1f}s)", flush=True)
        except (FileNotFoundError, AssertionError, KeyError) as e:
            print(f"  {sn}: FAILED — {e}", flush=True)

    # -- Per-diameter CCD --
    print("\n--- CCD per diameter ---", flush=True)

    for diam in DIAMETERS_ALL:
        tgt_name = f"target_cylinder_{diam}L"
        ill_name = f"illusion_{diam}L"

        tgt_data = data_cache.get(tgt_name)
        ill_data = data_cache.get(ill_name)
        pin_data = data_cache.get("pinball")

        if tgt_data is None:
            print(f"\n  SKIP {diam}L: missing target data", flush=True)
            continue

        print(f"\n{'=' * 60}", flush=True)
        print(f"Diameter {diam}L", flush=True)
        print(f"{'=' * 60}", flush=True)

        is_special = (diam not in DIAMETERS_MAIN)
        if is_special:
            print(f"  Note: {diam}L flagged as special-mechanism case", flush=True)

        # -- Build target-only POD basis --
        Q_tgt = build_field_matrix(tgt_data["ux"], tgt_data["uy"])
        mean_f, modes, sv, coeffs = compute_pod(Q_tgt)
        energy = cumulative_energy(sv)
        e95 = e95_index(energy)
        print(f"  Target-only POD: E95={e95}", flush=True)
        for i in range(min(8, len(sv))):
            print(f"    mode {i + 1}: energy={energy[i]:.4f}", flush=True)

        # -- Project illusion and pinball into target basis --
        proj_cache = {tgt_name: coeffs}  # already in target basis

        if ill_data is not None:
            proj_cache[ill_name] = project_into_basis(
                ill_data["ux"], ill_data["uy"], modes, mean_f)

        if pin_data is not None:
            proj_cache["pinball"] = project_into_basis(
                pin_data["ux"], pin_data["uy"], modes, mean_f)

        # -- CCD for each r and each case --
        for r in R_CANDIDATES:
            print(f"\n  r={r}:", flush=True)
            modes_r = modes[:, :r]

            for name in [tgt_name, ill_name, "pinball"]:
                d = data_cache.get(name)
                if d is None:
                    continue
                if name not in proj_cache:
                    continue

                a_r = proj_cache[name][:r, :]
                N = a_r.shape[1]

                # --- Force-CCD (primary: SigmaFy) ---
                frc = d.get("forces")
                if frc is not None:
                    for f_mode, f_label in [("fy", "force_fy"),
                                             ("fx", "force_fx"),
                                             ("joint", "force_joint")]:
                        y_f = make_force_obs(frc, name, mode=f_mode)
                        y_f = y_f[:, :N]
                        W, sig, Rmat, z, No, Nv = compute_reduced_ccd(
                            a_r[:, :N], y_f, Q_delay=CCD_Q)
                        en = cumulative_energy(sig)
                        m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 0

                        key = f"{diam}L_{name}_{f_label}_r{r}"
                        W_dict[key] = W
                        all_results[key] = {
                            "diam": diam, "case": name,
                            "obs": f_label, "r": r,
                            "m80": m80, "N": Nv,
                            "sigma_top3": [float(sig[i])
                                           for i in range(min(3, len(sig)))],
                            "special_mechanism": is_special,
                        }
                        if f_mode == "fy":
                            print(f"    {key}: m80={m80}, "
                                  f"sigma1={float(sig[0]):.4f}", flush=True)

                # --- Action-CCD (illusion only) ---
                act = d.get("actions")
                if act is not None:
                    y_a = act.T  # (3, N)
                    W, sig, Rmat, z, No, Nv = compute_reduced_ccd(
                        a_r[:, :N], y_a[:, :N], Q_delay=CCD_Q)
                    en = cumulative_energy(sig)
                    m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 0

                    key = f"{diam}L_{name}_action_r{r}"
                    W_dict[key] = W
                    all_results[key] = {
                        "diam": diam, "case": name,
                        "obs": "action", "r": r,
                        "m80": m80, "N": Nv,
                        "sigma_top3": [float(sig[i])
                                       for i in range(min(3, len(sig)))],
                        "special_mechanism": is_special,
                    }
                    print(f"    {key}: m80={m80}, "
                          f"sigma1={float(sig[0]):.4f}", flush=True)

        # -- Modal overlaps (r=6, force_fy primary) --
        print(f"\n  Modal overlap (r=6, force_fy):", flush=True)
        ov_list = compute_modal_overlap(W_dict, diam, 6, "force_fy")
        for ov in ov_list:
            print(f"    O({ov['case_a']}, {ov['case_b']}) "
                  f"mode{ov['mode']} = {ov['O']:.4f}", flush=True)

        # -- Reconstruction quality (POD basis check) --
        # Project target fields back onto its own POD basis and check residual
        q_rec = modes[:, :r] @ coeffs[:r, :] + mean_f[:, None]
        res = Q_tgt.astype(np.float64) - q_rec
        r2 = 1.0 - np.sum(res ** 2) / np.sum(Q_tgt.astype(np.float64) ** 2)
        print(f"  Target self-reconstruction R2 (r={r}): {r2:.4f}", flush=True)

    # -- Cross-diameter comparison (0.75L illusion in 1.0L basis) --
    print("\n--- Cross-diameter: 0.75L -> 1.0L basis ---", flush=True)
    d10_cache = data_cache.get("target_cylinder_1.0L")
    d075_i = data_cache.get("illusion_0.75L")
    if d10_cache is not None and d075_i is not None:
        Q_10 = build_field_matrix(d10_cache["ux"], d10_cache["uy"])
        mf_10 = np.mean(Q_10, axis=1)
        U_10, _, _ = np.linalg.svd(Q_10 - mf_10[:, None], full_matrices=False)
        modes_10_6 = U_10[:, :6]

        # Project 0.75L illusion
        a_075 = project_into_basis(d075_i["ux"], d075_i["uy"],
                                   modes_10_6, mf_10)[:6, :]
        frc_075 = d075_i.get("forces")
        if frc_075 is not None:
            y_f = make_force_obs(frc_075, "illusion_0.75L", mode="fy")
            W_cross, _, _, _, _, _ = compute_reduced_ccd(a_075, y_f, Q_delay=CCD_Q)

            # Compare with 1.0L illusion in its own basis
            d10_i = data_cache.get("illusion_1.0L")
            if d10_i is not None:
                a_10 = project_into_basis(d10_i["ux"], d10_i["uy"],
                                          modes_10_6, mf_10)[:6, :]
                frc_10 = d10_i.get("forces")
                if frc_10 is not None:
                    y_f10 = make_force_obs(frc_10, "illusion_1.0L", mode="fy")
                    W_10, _, _, _, _, _ = compute_reduced_ccd(a_10, y_f10, Q_delay=CCD_Q)
                    n = min(W_cross.shape[1], W_10.shape[1], 5)
                    for k in range(n):
                        ov = float(abs(
                            W_cross[:, k] / (np.linalg.norm(W_cross[:, k]) + 1e-12) @
                            W_10[:, k] / (np.linalg.norm(W_10[:, k]) + 1e-12)
                        ))
                        print(f"  Cross-diam O(0.75L->1.0L) mode{k + 1} = {ov:.4f}",
                              flush=True)

    # -- Save --
    with open(os.path.join(out_dir, "ccd_results.json"), "w") as f:
        json.dump(all_results, f, indent=2)
    print(f"\nSaved to {out_dir}/ccd_results.json", flush=True)
    print(f"Total entries: {len(all_results)}", flush=True)
    return 0


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