"""1.5L correction-field CCD: force-CCD, action-CCD, signature-CCD on dq_ctl.

Extends the Phase 2 pipeline to the 1.5L "special mechanism" case.
Target-only POD basis, Q_delay=6, r=[6, 8, 10].

Usage:
    conda run -n pycuda_3_10 python correction_analysis/run_15L_correction.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, NX, NY, CENTER_Y
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_TYPE = "illusion_1.5L"

# ---------------------------------------------------------------------------
# Zone masks for illusion layout (sensors at x=30*L0=600)
# ---------------------------------------------------------------------------

def _define_zones() -> dict:
    """Define body_wake and sensor_zone masks for 1.5L illusion."""
    zones = {}
    # body_wake: immediate downstream of pinball, x=[500, 700)
    mask_bw = np.zeros((NY, NX), dtype=bool)
    mask_bw[:, 500:700] = True
    zones["body_wake"] = mask_bw
    # sensor_zone: around sensors at x=600, x=[580, 650)
    mask_sz = np.zeros((NY, NX), dtype=bool)
    mask_sz[:, 580:650] = True
    zones["sensor_zone"] = mask_sz
    return zones


def _zone_ke_ratio(dq: dict, zones: dict) -> dict:
    """Compute correction energy ratio body_wake / sensor_zone."""
    ux, uy = dq["ux"], dq["uy"]
    ke_field = 0.5 * np.mean(ux**2 + uy ** 2, axis=0)  # (NY, NX)
    body_ke = ke_field[zones["body_wake"]].sum()
    sensor_ke = ke_field[zones["sensor_zone"]].sum()
    ratio = body_ke / sensor_ke if sensor_ke > 0 else float("inf")
    return {
        "body_wake_KE": float(body_ke),
        "sensor_zone_KE": float(sensor_ke),
        "ratio_bw_over_sz": float(ratio),
    }


# ---------------------------------------------------------------------------
# Signature-CCD helpers
# ---------------------------------------------------------------------------

def compute_tau_corr(a_ctl: np.ndarray, e_s: np.ndarray,
                     max_lag: int = 12) -> int:
    """Find tau that maximises |cross-correlation| between a1 and sensor error.

    Computes average absolute cross-correlation across sensor channels,
    returns the lag (in snapshot steps) with the strongest correlation.
    """
    a1 = a_ctl[0, :]  # leading POD coefficient
    n = len(a1)
    # Normalise
    a1_z = (a1 - a1.mean()) / (a1.std() + 1e-12)
    # Average absolute correlation across sensor channels
    corr_avg = np.zeros(2 * max_lag + 1)
    for ch in range(e_s.shape[0]):
        ech = e_s[ch, :n]
        ech_z = (ech - ech.mean()) / (ech.std() + 1e-12)
        c = np.correlate(a1_z, ech_z, mode="full")
        c_mid = len(c) // 2
        seg = c[c_mid - max_lag:c_mid + max_lag + 1]
        corr_avg += np.abs(seg)
    corr_avg /= e_s.shape[0]
    best_lag = np.argmax(corr_avg) - max_lag
    return int(best_lag)


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


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

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

    # ---- 1. Load correction fields for 1.5L ----
    print(f"\n--- Loading correction: {SCENE_TYPE} ---", flush=True)
    t0 = time.time()
    corr = compute_correction(SCENE_TYPE)
    dq_ctl = corr["dq_ctl"]
    dq_tar = corr["dq_tar"]
    diam = corr.get("diam")
    t_elapsed = time.time() - t0

    if dq_ctl is None:
        print("  dq_ctl is None — cannot proceed.", flush=True)
        return 1

    print(f"  dq_ctl: {dq_ctl['ux'].shape[0]} frames, "
          f"forces={'✓' if dq_ctl['forces'] is not None else '✗'}, "
          f"actions={'✓' if dq_ctl['actions'] is not None else '✗'}, "
          f"sensors={'✓' if dq_ctl['sensors'] is not None else '✗'}, "
          f"{t_elapsed:.1f}s", flush=True)

    # ---- 2. Phase drift: zone energy ratio ----
    print(f"\n--- Phase drift: zone energy ratio (body_wake / sensor_zone) ---",
          flush=True)
    zones = _define_zones()
    ze = _zone_ke_ratio(dq_ctl, zones)
    print(f"  body_wake KE = {ze['body_wake_KE']:.4e}", flush=True)
    print(f"  sensor_zone KE = {ze['sensor_zone_KE']:.4e}", flush=True)
    print(f"  ratio (bw/sz) = {ze['ratio_bw_over_sz']:.4f}", flush=True)
    all_results["zone_energy_ratio"] = ze

    # ---- 3. POD: target-only basis ----
    print(f"\n--- POD: target-only basis ---", 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)

    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 "
              f"(E95={e95_index(cumulative_energy(sv_tar))})", flush=True)
    else:
        print("  dq_tar is None — cannot proceed.", flush=True)
        return 1

    # ---- 4. Sensor error for signature line ----
    sensors_ctl = dq_ctl.get("sensors")  # (N, 6) — illusion sensors
    sensors_tar = dq_tar.get("sensors")  # (N, 6) — target sensors
    if sensors_ctl is not None and sensors_tar is not None:
        # Both have 6 sensor channels: use all 6 dimensions
        n_min = min(sensors_ctl.shape[0], sensors_tar.shape[0], N)
        e_s_full = (sensors_ctl[:n_min] - sensors_tar[:n_min]).T  # (6, N)
        print(f"  Sensor error e_s: shape={e_s_full.shape}", flush=True)
    else:
        print("  Sensor data incomplete — signature-CCD skipped.", flush=True)
        e_s_full = None

    # ---- 5. Force-CCD, action-CCD, target-CCD ----
    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"),
                                  ("joint", "force_joint")]:
                y = make_force_obs(frc, SCENE_TYPE, 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"illusion_1.5L_dqctl_{flabel}_r{r}"
                W_dict[key] = W
                all_results[key] = {
                    "scene": SCENE_TYPE, "diam": diam, "obs": flabel, "r": r,
                    "m80": m80, "N": sig.size,
                    "sigma_top3": [
                        float(sig[i]) for i in range(min(3, len(sig)))
                    ],
                }
                if fmode == "fy":
                    print(f"    {key}: m80={m80} "
                          f"s1={float(sig[0]):.4f}", flush=True)

        # Action-CCD
        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"illusion_1.5L_dqctl_action_r{r}"
            W_dict[key] = W
            all_results[key] = {
                "scene": SCENE_TYPE, "diam": diam, "obs": "action", "r": r,
                "m80": m80, "N": sig.size,
                "sigma_top3": [
                    float(sig[i]) for i in range(min(3, len(sig)))
                ],
            }
            print(f"    {key}: m80={m80} s1={float(sig[0]):.4f}", flush=True)

        # Target force-CCD reference
        if dq_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(SCENE_TYPE) or f"{SCENE_TYPE}_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"illusion_1.5L_dqtar_force_fy_r{r}"
                W_dict[kt] = Wt
                all_results[kt] = {
                    "scene": SCENE_TYPE, "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)))
                    ],
                }
                # Overlap: dq_ctl vs dq_tar
                ck = f"illusion_1.5L_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"illusion_1.5L_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
                            )

    # ---- 6. Overlap at r=6 (comparison anchor) ----
    # Print explicit comparison with 0.75L (0.564) and 1.0L (0.913)
    key_r6 = "illusion_1.5L_O_dqctl_vs_dqtar_r6_mode1"
    ov_r6 = all_results.get(key_r6, {}).get("overlap")
    if ov_r6 is not None:
        verdict = "lower=special" if ov_r6 < 0.7 else "higher=normal"
        print(
            f"\n  1.5L O(dqctl, dqtar) = {ov_r6:.4f} "
            f"(0.75L: 0.564, 1.0L: 0.913 → {verdict})",
            flush=True,
        )

    # ---- 7. Signature-CCD ----
    print(f"\n--- Signature-CCD (future sensor error e_s(t+tau)) ---",
          flush=True)
    if e_s_full is not None:
        # tau candidates
        tau_geom = 3  # geometric advection delay (snapshot steps)
        tau_corr = compute_tau_corr(a_ctl, e_s_full, max_lag=12)
        tau_candidates = [("tau_0", 0), ("tau_geom", tau_geom),
                          ("tau_corr", tau_corr)]
        print(f"  tau_geom={tau_geom}, tau_corr={tau_corr}", flush=True)

        for tau_label, tau in tau_candidates:
            print(f"\n  --- tau={tau} ({tau_label}) ---", flush=True)
            for r in R_CANDIDATES:
                a_r = a_ctl[:r, :]
                Nv = a_r.shape[1]
                # Shift observable forward by tau
                if tau >= 0:
                    y_sig = e_s_full[:, tau: tau + Nv]
                    # Also shift POD coefficients to align: use a_r[:, :-tau]
                    a_r_aligned = a_r[:, :Nv - tau] if tau > 0 else a_r
                    y_sig_aligned = y_sig[:, :a_r_aligned.shape[1]]
                else:
                    # Negative tau: shift backward
                    y_sig = e_s_full[:, :Nv + tau]
                    a_r_aligned = a_r[:, -tau:]
                    y_sig_aligned = y_sig[:, :a_r_aligned.shape[1]]

                if y_sig_aligned.shape[1] < CCD_Q:
                    print(f"    r={r}: too few samples ({y_sig_aligned.shape[1]}), skipping",
                          flush=True)
                    continue

                W, sig, _, _, _, _ = compute_reduced_ccd(
                    a_r_aligned, y_sig_aligned, Q_delay=CCD_Q)
                en = cumulative_energy(sig)
                m80 = int(np.searchsorted(en, 0.80) + 1) if len(en) > 0 else 0
                key = f"illusion_1.5L_dqctl_signature_{tau_label}_r{r}"
                W_dict[key] = W
                all_results[key] = {
                    "scene": SCENE_TYPE, "diam": diam,
                    "obs": f"signature_{tau_label}", "r": r,
                    "tau": tau, "m80": m80, "N": sig.size,
                    "sigma_top3": [
                        float(sig[i]) for i in range(min(3, len(sig)))
                    ],
                }
                print(f"    {key}: m80={m80} "
                      f"s1={float(sig[0]):.4f}", flush=True)
    else:
        print("  Skipping: sensor error not available.", flush=True)

    # ---- 8. Save ----
    ccd_path = os.path.join(out_dir, "15L_correction_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)

    # ---- 9. Summary ----
    print("\n" + "=" * 60, flush=True)
    print("1.5L Correction-field CCD — Summary", flush=True)
    print("=" * 60, flush=True)
    print(f"  Zone KE ratio (body_wake/sensor_zone): {ze['ratio_bw_over_sz']:.4f}",
          flush=True)
    if ov_r6 is not None:
        print(f"  O(dqctl, dqtar) r=6 mode1: {ov_r6:.4f}", flush=True)

    for r in R_CANDIDATES:
        print(f"\n  r={r}:", flush=True)
        for obs in ["force_fy", "force_fx", "force_joint", "action"]:
            k = f"illusion_1.5L_dqctl_{obs}_r{r}"
            if k in all_results:
                d = all_results[k]
                print(f"    {obs:12s}: m80={d['m80']}, "
                      f"s1={d['sigma_top3'][0]:.4f}", flush=True)

    print(f"\nDone. Results saved.", flush=True)
    return 0


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