"""Signature-line CCD on dq_ctl: which correction structures determine future sensor mismatch.

Force/action line CCD on dq_ctl is complete (Phase 1-2). Now we need the
SIGNATURE LINE — answering which correction structures most determine future
sensor error (rather than instantaneous force).

Key idea: the observable for signature-CCD is the FUTURE sensor error
e(t+tau) = sensors_ctl(t+tau) - sensors_tar(t+tau).

Usage:
    conda run -n pycuda_3_10 python correction_analysis/run_signature_line.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, SCENES
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_LIST = [6, 8, 10]
CCD_Q = 6
SCENE_TYPES = ["illusion_0.75L", "illusion_1.0L"]
TAU_GEOM = 3


# ---------------------------------------------------------------------------
# Signature observable
# ---------------------------------------------------------------------------

def make_signature_obs(sensors_ctl: np.ndarray, sensors_tar: np.ndarray,
                       step_indices: list, tau: int = 0) -> np.ndarray:
    """Construct signature observable: future sensor error e(t+tau).

    Parameters
    ----------
    sensors_ctl : (N_total_raw, 6) — full raw sensor telemetry from controlled data.
    sensors_tar : (N_total_raw, 6) — full raw sensor telemetry from target data.
    step_indices : list of int — absolute frame indices (from dq_ctl step_indices).
    tau : int — future shift in simulation steps.

    Returns
    -------
    e : (6, N_valid) — sensor error at shifted indices (6 channels).
    """
    si = np.asarray(step_indices, dtype=int)
    max_idx = min(len(sensors_ctl), len(sensors_tar)) - 1
    shifted = np.clip(si + tau, 0, max_idx)
    e = sensors_ctl[shifted] - sensors_tar[shifted]
    return e.T  # (6, N_valid)


# ---------------------------------------------------------------------------
# Tau computation
# ---------------------------------------------------------------------------

def compute_tau_corr(sensors_ctl: np.ndarray, sensors_tar: np.ndarray,
                     step_indices: list, max_lag: int = 50) -> int:
    """Compute optimal tau via cross-correlation of target/illusion sensor[:,3].

    Cross-correlates the target cylinder sensor[:,3] with the illusion
    sensor[:,3] at the snapshot-aligned times. Returns the absolute lag
    (in steps) that maximises cross-correlation.

    Parameters
    ----------
    sensors_ctl : (N_total_raw, 6) — full raw sensors from controlled data.
    sensors_tar : (N_total_raw, 6) — full raw sensors from target data.
    step_indices : list of int — absolute frame indices.
    max_lag : int — maximum lag to consider (in steps).

    Returns
    -------
    tau : int — optimal lag in steps (non-negative).
    """
    si = np.asarray(step_indices, dtype=int)
    s_ctl = sensors_ctl[si, 3]
    s_tar = sensors_tar[si, 3]

    n = len(s_ctl)
    ctl = s_ctl - np.mean(s_ctl)
    tar = s_tar - np.mean(s_tar)

    xcorr = np.correlate(tar, ctl, mode='same')
    mid = n // 2
    lags = np.arange(-mid, mid + 1)
    if n % 2 == 0:
        lags = lags[:-1]

    valid = np.abs(lags) <= max_lag
    if not np.any(valid):
        return 0

    best_idx = np.argmax(xcorr[valid])
    tau = lags[valid][best_idx]
    return int(abs(tau))


# ---------------------------------------------------------------------------
# Raw sensor loader (full telemetry, before step-index subsampling)
# ---------------------------------------------------------------------------

def _load_raw_sensors(scene_name: str) -> np.ndarray:
    """Load full raw sensor telemetry (before step-index subsampling).

    Returns
    -------
    sensors : (N_total_raw, 6) ndarray — the full sensor time series.
    """
    cfg = SCENES.get(scene_name)
    if cfg is None:
        raise KeyError(f"Unknown scene: {scene_name}")

    scene_id = cfg["scene_id"]
    sd = os.path.join(DATA_DIR, scene_id, scene_name)

    tele_path = None
    for p in [os.path.join(sd, "controlled.npz"), os.path.join(sd, "sensors.npz")]:
        if os.path.isfile(p):
            tele_path = p
            break
    if tele_path is None:
        raise FileNotFoundError(f"No telemetry (*.npz) found in {sd}")

    td = np.load(tele_path)
    sensors = td["sensors"]
    td.close()
    return sensors


# ---------------------------------------------------------------------------
# Modal overlap
# ---------------------------------------------------------------------------

def modal_overlap(W_a: np.ndarray, W_b: np.ndarray, n_modes: int = 5) -> list:
    """Pairwise modal overlap between two CCD direction matrices.

    Returns list of {mode, O} dicts.
    """
    n = min(W_a.shape[1], W_b.shape[1], n_modes)
    results = []
    for k in range(n):
        u_a = W_a[:, k] / (np.linalg.norm(W_a[:, k]) + 1e-12)
        u_b = W_b[:, k] / (np.linalg.norm(W_b[:, k]) + 1e-12)
        ov = float(abs(u_a @ u_b))
        results.append({"mode": k + 1, "O": ov})
    return results


# ---------------------------------------------------------------------------
# LOCO validation helpers
# ---------------------------------------------------------------------------

def r2_score(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """Coefficient of determination."""
    ss_r = np.sum((y_true - y_pred) ** 2)
    ss_t = np.sum((y_true - np.mean(y_true)) ** 2)
    return float(1.0 - ss_r / (ss_t + 1e-12))


def reconstruct_from_ccd(W, sigma, R, a_test, y_train, CCD_Q, m_obs):
    """Reconstruct observable from CCD modes. Returns dict with 'mode1' and 'm80'."""
    am = np.mean(a_test, axis=1, keepdims=True)
    as_ = np.std(a_test, axis=1, keepdims=True) + 1e-12
    a_test_z = (a_test - am) / as_
    z_test = W.T @ a_test_z

    ym = np.mean(y_train, axis=1, keepdims=True)
    ys = np.std(y_train, axis=1, keepdims=True) + 1e-12
    half = CCD_Q // 2

    en = cumulative_energy(sigma)
    m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 1

    results = {}

    # Mode-1
    if R.shape[1] >= 1:
        pz_1 = R[:, :1] * sigma[:1] @ z_test[:1, :]
        yp_1 = pz_1[half * m_obs:(half + 1) * m_obs, :] * ys + ym
        results["mode1"] = yp_1
    else:
        results["mode1"] = np.zeros_like(y_train[:, :a_test.shape[1]])

    # M80
    n_rm = min(m80, R.shape[1])
    if n_rm >= 1:
        pz_m = R[:, :n_rm] * sigma[:n_rm] @ z_test[:n_rm, :]
        yp_m = pz_m[half * m_obs:(half + 1) * m_obs, :] * ys + ym
        results["m80"] = yp_m
    else:
        results["m80"] = np.zeros_like(y_train[:, :a_test.shape[1]])

    return results


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    print("=" * 60, flush=True)
    print("Signature-line CCD on 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 = {}            # CCD direction matrices keyed by label
    pod_basis_cache = {}   # (scene_type, r) -> (mean_field, modes_r)
    raw_sensors_cache = {}  # scene_name -> raw sensors

    # ---- 1. Load correction fields ----
    print("\n--- Step 1: 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"sensors={'✓' if corr['q_ctl'] is not None else '✗'}, "
                      f"forces={'✓' if dq['forces'] is not None else '✗'}, "
                      f"{time.time()-t0:.1f}s", flush=True)
        except Exception as e:
            print(f"  {st}: FAILED — {e}", flush=True)

    # ---- 2. Load raw sensor data for all needed scenes ----
    print("\n--- Step 2: Loading raw sensor telemetry ---", flush=True)
    for st in SCENE_TYPES:
        corr = cache.get(st)
        if corr is None:
            continue
        diam = corr.get("diam")
        tar_name = f"target_cylinder_{diam}L"
        for name in [st, tar_name]:
            if name not in raw_sensors_cache:
                try:
                    raw_sensors_cache[name] = _load_raw_sensors(name)
                    print(f"  {name}: raw sensors {raw_sensors_cache[name].shape}", flush=True)
                except Exception as e:
                    print(f"  {name}: FAILED — {e}", flush=True)

    # ---- 3. For each scene, pre-compute POD basis ----
    print("\n--- Step 3: Building target-only POD basis ---", flush=True)
    for st in SCENE_TYPES:
        corr = cache.get(st)
        if corr is None:
            continue
        dq_tar = corr["dq_tar"]
        if dq_tar is None:
            print(f"  {st}: dq_tar is None, skipping", flush=True)
            continue

        Q_tar = dict_to_field_matrix(dq_tar)
        mf_tar, modes_tar, sv_tar, _ = compute_pod(Q_tar)
        e95 = e95_index(cumulative_energy(sv_tar))
        print(f"  {st}: target-only POD E95={e95}", flush=True)

        for r in R_LIST:
            pod_basis_cache[(st, r)] = (mf_tar, modes_tar[:, :r])

    # ---- 4. Compute tau_corr for each scene ----
    print("\n--- Step 4: Computing tau values ---", flush=True)
    tau_config = {}  # scene_type -> list of tau
    for st in SCENE_TYPES:
        corr = cache.get(st)
        if corr is None:
            continue
        raw_ctl = raw_sensors_cache.get(st)
        diam = corr.get("diam")
        tar_name = f"target_cylinder_{diam}L"
        raw_tar = raw_sensors_cache.get(tar_name)
        step_idx = corr["q_ctl"].get("step_indices", [])
        if step_idx is None or len(step_idx) == 0:
            step_idx = list(range(corr["q_ctl"]["ux"].shape[0]))

        if raw_ctl is not None and raw_tar is not None:
            tau_corr = compute_tau_corr(raw_ctl, raw_tar, step_idx)
        else:
            tau_corr = TAU_GEOM

        taus = sorted(set([0, TAU_GEOM, tau_corr]))
        tau_config[st] = taus
        print(f"  {st}: tau_corr={tau_corr}, taus={taus}", flush=True)

    # ---- 5. Signature-CCD and Force-CCD for each (scene, r, tau) ----
    print("\n--- Step 5: Running CCD ---", flush=True)
    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 or dq_tar is None:
            continue

        diam = corr.get("diam")
        taus = tau_config.get(st, [0, TAU_GEOM])
        step_idx = corr["q_ctl"].get("step_indices", [])
        if step_idx is None or len(step_idx) == 0:
            step_idx = list(range(dq_ctl["ux"].shape[0]))

        raw_ctl = raw_sensors_cache.get(st)
        diam = corr.get("diam")
        tar_name = f"target_cylinder_{diam}L"
        raw_tar = raw_sensors_cache.get(tar_name)

        Q_ctl = dict_to_field_matrix(dq_ctl)
        N = Q_ctl.shape[1]
        print(f"\n  --- {st} (diam={diam}) ---", flush=True)

        for r in R_LIST:
            mf_r, modes_r = pod_basis_cache[(st, r)]
            a_r = modes_r.T @ (Q_ctl - mf_r[:, None]).astype(np.float64)
            Nv = a_r.shape[1]
            print(f"  r={r}: N={Nv}", flush=True)

            # -- Signature-CCD --
            if raw_ctl is not None and raw_tar is not None and step_idx is not None:
                for tau in taus:
                    e_sig = make_signature_obs(raw_ctl, raw_tar, step_idx, tau=tau)
                    # Trim to match a_r length
                    Ne = e_sig.shape[1]
                    a_r_use = a_r[:, :Ne] if Ne < Nv else a_r
                    e_use = e_sig[:, :Nv] if Nv < Ne else e_sig
                    N_use = min(Nv, Ne)

                    W, sig, Rmat, z, N_orig, N_valid = compute_reduced_ccd(
                        a_r_use[:, :N_use], e_use[:, :N_use], 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}_sig_tau{tau}_r{r}"
                    W_dict[key] = W
                    all_results[key] = {
                        "scene": st, "diam": diam, "obs": f"sig_tau{tau}", "r": r,
                        "tau": tau, "m80": m80, "N": sig.size, "N_valid": N_valid,
                        "N_original": N_orig,
                        "sigma_top3": [float(sig[i]) for i in range(min(3, len(sig)))],
                    }
                    print(f"    {key}: m80={m80}  "
                          f"s1={float(sig[0]):.4f}  N_valid={N_valid}", flush=True)

            # -- Force-CCD reference (SigmaFy, tau=0) --
            frc = dq_ctl.get("forces")
            if frc is not None:
                y_f = make_force_obs(frc, st, mode="fy")[:, :Nv]
                W_f, sig_f, _, _, N_orig_f, N_valid_f = compute_reduced_ccd(
                    a_r, y_f, Q_delay=CCD_Q
                )
                en_f = cumulative_energy(sig_f)
                m80_f = int(np.searchsorted(en_f, 0.80) + 1) if len(en_f) > 0 else 0
                key_f = f"{st}_force_fy_r{r}"
                W_dict[key_f] = W_f
                all_results[key_f] = {
                    "scene": st, "diam": diam, "obs": "force_fy", "r": r,
                    "tau": 0, "m80": m80_f, "N": sig_f.size, "N_valid": N_valid_f,
                    "N_original": N_orig_f,
                    "sigma_top3": [float(sig_f[i]) for i in range(min(3, len(sig_f)))],
                }
                print(f"    {key_f}: m80={m80_f}  "
                      f"s1={float(sig_f[0]):.4f}  N_valid={N_valid_f}", flush=True)

    # ---- 6. Force vs Signature modal overlap comparison (r=6) ----
    print("\n\n--- Step 6: Force vs Signature modal overlap (r=6) ---", flush=True)
    for st in SCENE_TYPES:
        corr = cache.get(st)
        if corr is None:
            continue
        diam = corr.get("diam")
        taus = tau_config.get(st, [0, TAU_GEOM])

        force_key = f"{st}_force_fy_r{6}"
        if force_key not in W_dict:
            print(f"  {st}: no force key, skipping overlap", flush=True)
            continue

        W_force = W_dict[force_key]
        print(f"\n  {st} (diam={diam}):", flush=True)
        for tau in taus:
            sig_key = f"{st}_sig_tau{tau}_r{6}"
            if sig_key not in W_dict:
                continue
            W_sig = W_dict[sig_key]
            ovs = modal_overlap(W_force, W_sig, n_modes=5)
            for ov in ovs:
                key = f"{st}_O_force_vs_sig_tau{tau}_r6_mode{ov['mode']}"
                all_results[key] = {
                    "scene": st, "diam": diam, "r": 6,
                    "tau_sig": tau, "mode": ov["mode"],
                    "overlap": ov["O"],
                }
            ov_str = ", ".join([f"mode{ov['mode']}={ov['O']:.4f}" for ov in ovs])
            print(f"    O(force, sig_tau{tau}) r=6: {ov_str}", flush=True)

    # ---- 7. LOCO validation (signature observable) ----
    N_PTS = 24
    N_CYCLES = 4

    print("\n\n--- Step 7: LOCO validation (signature observable, r=6) ---", flush=True)
    loco_results = {}
    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 or dq_tar is None:
            continue

        diam = corr.get("diam")
        taus = tau_config.get(st, [0, TAU_GEOM])
        step_idx = corr["q_ctl"].get("step_indices", [])
        if step_idx is None or len(step_idx) == 0:
            step_idx = list(range(dq_ctl["ux"].shape[0]))

        tar_name = f"target_cylinder_{diam}L"
        raw_ctl = raw_sensors_cache.get(st)
        raw_tar = raw_sensors_cache.get(tar_name)
        if raw_ctl is None or raw_tar is None:
            print(f"  {st}: raw sensors missing, skipping LOCO", flush=True)
            continue

        Q_ctl = dict_to_field_matrix(dq_ctl)
        N_total = Q_ctl.shape[1]

        for tau in taus:
            # Build the full signature observable
            e_full = make_signature_obs(raw_ctl, raw_tar, step_idx, tau=tau)

            r = 6
            fold_r2_m1, fold_r2_m80 = [], []

            for fold in range(N_CYCLES):
                test_cyc = fold
                train_cyc = [c for c in range(N_CYCLES) if c != test_cyc]
                train_idx = sorted([c * N_PTS + p for c in train_cyc for p in range(N_PTS)])
                test_idx = sorted([c * N_PTS + p for c in [test_cyc] for p in range(N_PTS)])

                # Trim to valid range
                train_idx = [i for i in train_idx if i < N_total]
                test_idx = [i for i in test_idx if i < N_total]
                if len(train_idx) < N_PTS or len(test_idx) < N_PTS // 2:
                    continue

                # Build LOCO POD basis from target-only data
                Q_tar = dict_to_field_matrix(dq_tar)
                Q_ref = Q_tar[:, train_idx]
                mf = np.mean(Q_ref, axis=1)
                U, _, _ = np.linalg.svd(Q_ref - mf[:, None], full_matrices=False)
                modes_r = U[:, :r]

                a_train = modes_r.T @ (Q_ctl[:, train_idx] - mf[:, None])
                a_test = modes_r.T @ (Q_ctl[:, test_idx] - mf[:, None])

                y_train = e_full[:, train_idx]
                y_test = e_full[:, test_idx]

                # Handle length mismatches
                na = a_train.shape[1]
                ny = y_train.shape[1]
                n_min = min(na, ny)
                a_train = a_train[:, :n_min]
                y_train = y_train[:, :n_min]

                try:
                    W, sigma, Rmat, _, _, _ = compute_reduced_ccd(a_train, y_train, Q_delay=CCD_Q)
                except Exception as exc:
                    print(f"    LOCO fold {fold}: CCD failed — {exc}", flush=True)
                    continue

                recon = reconstruct_from_ccd(W, sigma, Rmat, a_test, y_train, CCD_Q, m_obs=6)

                na_test = a_test.shape[1]
                ny_test = y_test.shape[1]
                n_test = min(na_test, ny_test)
                ch_m1 = [r2_score(y_test[c, :n_test], recon["mode1"][c, :n_test])
                         for c in range(min(y_test.shape[0], recon["mode1"].shape[0]))]
                ch_m80 = [r2_score(y_test[c, :n_test], recon["m80"][c, :n_test])
                          for c in range(min(y_test.shape[0], recon["m80"].shape[0]))]
                fold_r2_m1.append(float(np.mean(ch_m1)))
                fold_r2_m80.append(float(np.mean(ch_m80)))

            if fold_r2_m1:
                key = f"{st}_LOCO_sig_tau{tau}_r{r}"
                loco_results[key] = {
                    "scene": st, "diam": diam, "tau": tau, "r": r,
                    "mode1": {
                        "mean": float(np.mean(fold_r2_m1)),
                        "std": float(np.std(fold_r2_m1)),
                    },
                    "m80": {
                        "mean": float(np.mean(fold_r2_m80)),
                        "std": float(np.std(fold_r2_m80)),
                    },
                }
                print(f"  {key}: R2_m1={loco_results[key]['mode1']['mean']:.4f}+-"
                      f"{loco_results[key]['mode1']['std']:.4f}  "
                      f"R2_m80={loco_results[key]['m80']['mean']:.4f}+-"
                      f"{loco_results[key]['m80']['std']:.4f}", flush=True)
            else:
                print(f"  {st} tau={tau}: LOCO skipped (no valid folds)", flush=True)

    all_results["_loco"] = loco_results

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

    # ---- 9. Summary ----
    print("\n" + "=" * 60, flush=True)
    print("SUMMARY", flush=True)
    print("=" * 60, flush=True)

    for st in SCENE_TYPES:
        corr = cache.get(st)
        if corr is None:
            continue
        diam = corr.get("diam")
        taus = tau_config.get(st, [0, TAU_GEOM])
        print(f"\n  {st} (diam={diam}):", flush=True)

        # Signature R2_m80 from LOCO (r=6)
        for tau in taus:
            lk = f"{st}_LOCO_sig_tau{tau}_r{6}"
            if lk in loco_results:
                r2_m80 = loco_results[lk]["m80"]["mean"]
                flag = "✓" if r2_m80 >= 0.4 else "✗"
                print(f"    LOCO sig_tau{tau} R2_m80={r2_m80:.4f} {flag}", flush=True)

        # Overlaps
        for tau in taus:
            ok = f"{st}_O_force_vs_sig_tau{tau}_r6_mode1"
            if ok in all_results:
                ov = all_results[ok]["overlap"]
                cat = "shared" if ov > 0.8 else ("partial" if ov > 0.5 else "separated")
                print(f"    O(force, sig_tau{tau}) mode1={ov:.4f} ({cat})", flush=True)

    return 0


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