"""Minimal correction-field CCD: POD + force/action CCD on dq_ctl.

Simplified version — processes only illusion_0.75L and illusion_1.0L.
No LOCO validation (separate step). Outputs CCD results and overlaps.

Usage:
    conda run -n pycuda_3_10 python correction_analysis/decompose_corrections.py
"""
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
from CCD_analysis.utils.resampling import (
    compute_pod, cumulative_energy, e95_index,
    compute_reduced_ccd, make_force_obs,
)
from CCD_analysis.correction_analysis.compute_correction_fields import (
    compute_correction, dict_to_field_matrix,
)

R_CANDIDATES = [6, 8, 10]
CCD_Q = 6
SCENE_TYPES = ["illusion_0.75L", "illusion_1.0L", "illusion_1.5L", "steady_cloak"]
DIAMETERS_MAIN = [0.75, 1.0]
DIAMETER_SPECIAL = 1.5  # flagged as special_mechanism (high-freq modulation)


def compute_modal_overlap(W_dict, scene_label, r, obs_label="force_fy"):
    keys = [k for k in W_dict
            if scene_label 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


def _scene_to_target_name(scene_type):
    if "illusion" in scene_type:
        parts = scene_type.split("_")
        if len(parts) >= 2:
            return f"target_cylinder_{parts[1]}"
    return None


def main():
    print("=" * 60, flush=True)
    print("Correction-field CCD (Phase 3) — dq_ctl", 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 = {}

    # Load correction fields
    print("\n--- Loading correction fields ---", flush=True)
    cache = {}
    for st in SCENE_TYPES:
        t0 = time.time()
        try:
            corr = compute_correction(st)
            cache[st] = corr
            dq = corr["dq_ctl"]
            if dq is not None:
                print(f"  {st}: dq_ctl {dq['ux'].shape[0]} frames, "
                      f"forces={'✓' if dq['forces'] is not None else '✗'}, "
                      f"actions={'✓' if dq['actions'] is not None else '✗'}, "
                      f"{time.time()-t0:.1f}s", flush=True)
        except Exception as e:
            print(f"  {st}: FAILED — {e}", flush=True)

    # CCD on dq_ctl
    for st in SCENE_TYPES:
        corr = cache.get(st)
        if corr is None:
            continue
        dq_ctl = corr["dq_ctl"]
        dq_tar = corr["dq_tar"]
        if dq_ctl is None:
            continue

        diam = corr.get("diam")
        is_special = (diam is not None and diam >= DIAMETER_SPECIAL)
        flag = " [SPECIAL MECHANISM — high-freq modulation]" if is_special else ""
        print(f"\n--- {st} (diam={diam}){flag} ---", flush=True)

        Q_ctl = dict_to_field_matrix(dq_ctl)
        N = Q_ctl.shape[1]
        print(f"  dq_ctl: shape={Q_ctl.shape}", flush=True)

        # POD: target-only (dq_tar) or direct (dq_ctl)
        if dq_tar is not None:
            Q_tar = dict_to_field_matrix(dq_tar)
            print(f"  dq_tar: shape={Q_tar.shape}", flush=True)
            mf_tar, modes_tar, sv_tar, coeffs_tar = compute_pod(Q_tar)
            # Project dq_ctl into target basis
            dc = dq_ctl
            q_proj = np.column_stack([
                np.concatenate([dc["ux"][s].ravel(), dc["uy"][s].ravel()])
                for s in range(N)
            ])
            a_ctl = modes_tar.T @ (q_proj - mf_tar[:, None]).astype(np.float64)
            a_tar = coeffs_tar
            print(f"  POD: target-only basis (E95={e95_index(cumulative_energy(sv_tar))})", flush=True)
        else:
            mf, modes, sv, coeffs = compute_pod(Q_ctl)
            a_ctl = coeffs
            a_tar = None
            print(f"  POD: direct dq_ctl (E95={e95_index(cumulative_energy(sv))})", flush=True)

        for r in R_CANDIDATES:
            a_r = a_ctl[:r, :]
            Nv = a_r.shape[1]
            print(f"\n  r={r}: N={Nv}", flush=True)

            # Force-CCD
            frc = dq_ctl.get("forces")
            if frc is not None:
                for fmode, flabel in [("fy","force_fy"), ("fx","force_fx")]:
                    y = make_force_obs(frc, st, mode=fmode)[:, :Nv]
                    W, sig, _, _, _, _ = compute_reduced_ccd(a_r, y, Q_delay=CCD_Q)
                    en = cumulative_energy(sig)
                    m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 0
                    key = f"{st}_dqctl_{flabel}_r{r}"
                    W_dict[key] = W
                    all_results[key] = {
                        "scene": st, "diam": diam, "obs": flabel, "r": r,
                        "m80": m80, "N": sig.size,
                        "sigma_top3": [float(sig[i]) for i in range(min(3,len(sig)))],
                        "special_mechanism": is_special,
                    }
                    if fmode == "fy":
                        print(f"    {key}: m80={m80} s1={float(sig[0]):.4f}", flush=True)

            # Action-CCD (illusion only)
            act = dq_ctl.get("actions")
            if act is not None:
                y_a = act.T[:, :Nv]
                W, sig, _, _, _, _ = compute_reduced_ccd(a_r, y_a, Q_delay=CCD_Q)
                en = cumulative_energy(sig)
                m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 0
                key = f"{st}_dqctl_action_r{r}"
                W_dict[key] = W
                all_results[key] = {
                    "scene": st, "diam": diam, "obs": "action", "r": r,
                    "m80": m80, "N": sig.size,
                    "sigma_top3": [float(sig[i]) for i in range(min(3,len(sig)))],
                    "special_mechanism": is_special,
                }
                print(f"    {key}: m80={m80} s1={float(sig[0]):.4f}", flush=True)

            # Target force-CCD reference (if available)
            if dq_tar is not None and a_tar is not None:
                a_tr = a_tar[:r, :Nv]
                frc_t = dq_tar.get("forces")
                if frc_t is not None:
                    tname = _scene_to_target_name(st) or f"{st}_tar"
                    y_t = make_force_obs(frc_t[:Nv], tname, mode="fy")
                    Wt, sig_t, _, _, _, _ = compute_reduced_ccd(a_tr, y_t, Q_delay=CCD_Q)
                    kt = f"{st}_dqtar_force_fy_r{r}"
                    W_dict[kt] = Wt
                    all_results[kt] = {
                        "scene": st, "diam": diam, "obs": "force_fy_tar", "r": r,
                        "m80": int(np.searchsorted(cumulative_energy(sig_t), 0.80)+1) if len(sig_t) > 0 else 0,
                        "N": sig_t.size,
                        "sigma_top3": [float(sig_t[i]) for i in range(min(3,len(sig_t)))],
                        "special_mechanism": is_special,
                    }
                    # Overlap: dq_ctl vs dq_tar
                    ck = f"{st}_dqctl_force_fy_r{r}"
                    if ck in W_dict:
                        Wc = W_dict[ck]
                        n = min(Wc.shape[1], Wt.shape[1], 5)
                        for k in range(n):
                            ov = float(abs(
                                Wc[:, k] / (np.linalg.norm(Wc[:, k])+1e-12) @
                                Wt[:, k] / (np.linalg.norm(Wt[:, k])+1e-12)
                            ))
                            all_results[f"{st}_O_dqctl_vs_dqtar_r{r}_mode{k+1}"] = {
                                "overlap": ov, "mode": k+1, "r": r
                            }
                            if k == 0:
                                print(f"    O(dqctl, dqtar) mode1={ov:.4f}", flush=True)

        # Overlap dqctl_target vs dqctl_illusion at r=6
        print(f"  Modal overlaps r=6:", flush=True)
        ovs = compute_modal_overlap(W_dict, st, 6, "force_fy")
        for ov in ovs:
            print(f"    O({ov['case_a']}, {ov['case_b']}) mode{ov['mode']} = {ov['O']:.4f}", flush=True)

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


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