"""Zone-wise CCD: force-CCD and signature-CCD per spatial zone.

Processes each of three spatial zones separately for illusion 0.75L and 1.0L
(unified geometry: pinball center ~613px, sensors at 800px):
  - near_body:  x[580:720]  (envelope around pinball)
  - body_wake:  x[720:850]  (body-connected near wake)
  - sensor_zone: x[780:850] (around sensor plane at x=40*L0=800)

For each zone: mask the snapshot matrix to keep only grid points in the zone,
build a target-only POD basis, project correction fields, and compute
force-CCD (SigmaFy) and signature-CCD (tau=0, tau=tau_corr) at r=6, Q_delay=6.

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

CCD_Q = 6
R = 6


# ---------------------------------------------------------------------------
# Zone masks
# ---------------------------------------------------------------------------

def _define_zones() -> dict:
    """Define three-zone masks for unified geometry (pinball center ~613px, sensors at 800px).

    Updated 2026-06-28 to match unified geometry after all scenes were re-collected
    with pinball at front x=30*L0=600, rear x=31.3*L0=626 (center ~613px), sensors at 40*L0=800.
    Same zones as diagnose_corrections.py:define_zones().
    """
    zones = {}

    # near_body: envelope around pinball (center ~613px)
    mask = np.zeros((NY, NX), dtype=bool)
    mask[:, 580:720] = True
    zones["near_body"] = mask

    # body_wake: near wake downstream of pinball
    mask = np.zeros((NY, NX), dtype=bool)
    mask[:, 720:850] = True
    zones["body_wake"] = mask

    # sensor_zone: around sensors at x=800 (40*L0)
    mask = np.zeros((NY, NX), dtype=bool)
    mask[:, 780:850] = True
    zones["sensor_zone"] = mask

    return zones


# ---------------------------------------------------------------------------
# Masking helper
# ---------------------------------------------------------------------------

def mask_field_matrix(Q_full: np.ndarray, ny: int, nx: int,
                      mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Zero out all grid points outside the mask.

    Parameters
    ----------
    Q_full : (2*nx*ny, N) ndarray
        Snapshot matrix — first half = ux, second half = uy.
    mask : (ny, nx) ndarray
        Boolean mask, True = keep.

    Returns
    -------
    Q_masked : (2*n_sum, N) ndarray
        Masked snapshot matrix.
    ux_idx : (n_sum,) ndarray
        Indices into the original ux ravel for kept points.
    """
    mask_flat = mask.ravel()  # (ny*nx,)
    ux_idx = np.where(mask_flat)[0]
    uy_idx = ux_idx + nx * ny
    all_idx = np.concatenate([ux_idx, uy_idx])
    return Q_full[all_idx, :], ux_idx


# ---------------------------------------------------------------------------
# tau_corr heuristic (from run_15L_correction.py)
# ---------------------------------------------------------------------------

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."""
    a1 = a_ctl[0, :]
    n = len(a1)
    a1_z = (a1 - a1.mean()) / (a1.std() + 1e-12)
    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)


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

def main():
    print("=" * 60, flush=True)
    print("Zone-wise CCD: force + signature per spatial zone", flush=True)
    print("=" * 60, flush=True)

    out_dir = os.path.join(DATA_DIR, "ccd")
    os.makedirs(out_dir, exist_ok=True)
    all_results = {}

    scene_types = ["illusion_0.75L", "illusion_1.0L"]
    zones = _define_zones()

    for scene_type in scene_types:
        print(f"\n{'=' * 60}", flush=True)
        print(f"Processing: {scene_type}", flush=True)
        print(f"{'=' * 60}", flush=True)

        # ---- 1. Load correction fields ----
        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)
            continue
        if dq_tar is None:
            print("  dq_tar is None — cannot proceed.", flush=True)
            continue

        N = dq_ctl["ux"].shape[0]
        print(f"  N={N}, diam={diam}, load_time={t_elapsed:.1f}s", flush=True)

        # ---- 2. Build full snapshot matrices ----
        Q_ctl = dict_to_field_matrix(dq_ctl)  # (2*NX*NY, N)
        Q_tar = dict_to_field_matrix(dq_tar)  # (2*NX*NY, N_tar)
        print(f"  Q_ctl: {Q_ctl.shape}, Q_tar: {Q_tar.shape}", flush=True)

        # ---- 3. Extract global observables (unmasked) ----
        # Force observable (SigmaFy)
        frc = dq_ctl.get("forces")
        if frc is None:
            print("  No force data — skipping.", flush=True)
            continue
        y_force = make_force_obs(frc, scene_type, mode="fy")[:, :N]  # (1, N)

        # Sensor error e_s = sensors_ctl - sensors_tar
        sensors_ctl = dq_ctl.get("sensors")
        sensors_tar = dq_tar.get("sensors")
        if sensors_ctl is None or sensors_tar is None:
            print("  Sensor data incomplete — skipping.", flush=True)
            continue
        n_min = min(sensors_ctl.shape[0], sensors_tar.shape[0], N)
        e_s = (sensors_ctl[:n_min] - sensors_tar[:n_min]).T  # (6, N)

        print(f"  y_force: {y_force.shape}, e_s: {e_s.shape}", flush=True)

        # ---- 4. Process each zone ----
        for zname, zmask in zones.items():
            n_masked = int(zmask.sum())
            print(f"\n  --- Zone: {zname} (N_grid={n_masked}) ---", flush=True)

            # Mask snapshot matrices
            Q_ctl_m, _ = mask_field_matrix(Q_ctl, NY, NX, zmask)
            Q_tar_m, _ = mask_field_matrix(Q_tar, NY, NX, zmask)
            n_field = Q_ctl_m.shape[0]
            print(f"    Masked field dim: {n_field}", flush=True)

            # Build target-only POD basis from masked dq_tar
            mf_tar, modes_tar, sv_tar, _ = compute_pod(Q_tar_m)
            en_tar = cumulative_energy(sv_tar)
            e95 = int(np.searchsorted(en_tar, 0.95) + 1) if len(en_tar) > 0 else 0
            print(f"    Target POD: E95={e95}, "
                  f"N_modes={len(sv_tar)}", flush=True)

            # Project masked dq_ctl into masked target basis
            proj_mean = mf_tar[:, None]
            a_ctl_all = modes_tar.T @ (Q_ctl_m - proj_mean).astype(np.float64)
            a_r = a_ctl_all[:R, :]  # (R, N)

            # ---- Force-CCD (SigmaFy) ----
            W_f, sig_f, _, _, _, _ = compute_reduced_ccd(
                a_r, y_force, 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
            frc_key = f"{scene_type}_{zname}_force_fy_r{R}"
            all_results[frc_key] = {
                "scene": scene_type,
                "zone": zname,
                "r": R,
                "N_masked_grid": n_masked,
                "m80": m80_f,
                "N_modes": int(sig_f.size),
                "sigma_top3": [float(sig_f[i])
                               for i in range(min(3, len(sig_f)))],
            }
            s1_f = sig_f[0]
            s2_f = sig_f[1] if len(sig_f) > 1 else float('nan')
            s3_f = sig_f[2] if len(sig_f) > 2 else float('nan')
            print(f"      force_fy r={R}: m80={m80_f}, "
                  f"s1={s1_f:.4f}, s2={s2_f:.4f}, s3={s3_f:.4f}",
                  flush=True)

            # ---- Compute tau_corr for this zone ----
            tau_corr = compute_tau_corr(a_ctl_all, e_s, max_lag=12)
            tau_candidates = [("tau0", 0), ("tau_corr", tau_corr)]
            print(f"      tau_corr = {tau_corr}", flush=True)

            # ---- Signature-CCD ----
            for tau_label, tau in tau_candidates:
                # Shift sensor error forward by tau
                if tau >= 0:
                    y_sig = e_s[:, tau: tau + N]
                    a_r_aligned = a_r[:, :N - tau] if tau > 0 else a_r
                else:
                    y_sig = e_s[:, :N + 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"        tau={tau}: too few samples "
                          f"({y_sig_aligned.shape[1]}), skipping", flush=True)
                    continue

                W_s, sig_s, _, _, _, _ = compute_reduced_ccd(
                    a_r_aligned, y_sig_aligned, Q_delay=CCD_Q)

                en_s = cumulative_energy(sig_s)
                m80_s = (int(np.searchsorted(en_s, 0.80) + 1)
                         if len(en_s) > 0 else 0)
                sig_key = f"{scene_type}_{zname}_sig_{tau_label}_r{R}"
                all_results[sig_key] = {
                    "scene": scene_type,
                    "zone": zname,
                    "r": R,
                    "tau": tau,
                    "N_masked_grid": n_masked,
                    "m80": m80_s,
                    "N_modes": int(sig_s.size),
                    "sigma_top3": [float(sig_s[i])
                                   for i in range(min(3, len(sig_s)))],
                }
                s1_s = sig_s[0]
                s2_s = sig_s[1] if len(sig_s) > 1 else float('nan')
                s3_s = sig_s[2] if len(sig_s) > 2 else float('nan')
                print(f"      sig_{tau_label}: m80={m80_s}, "
                      f"s1={s1_s:.4f}, s2={s2_s:.4f}, s3={s3_s:.4f}",
                      flush=True)

                # ---- Overlap O(force, sig) ----
                w_f0 = W_f[:, 0] / (np.linalg.norm(W_f[:, 0]) + 1e-12)
                w_s0 = W_s[:, 0] / (np.linalg.norm(W_s[:, 0]) + 1e-12)
                overlap = float(abs(w_f0 @ w_s0))

                ov_key = f"{scene_type}_{zname}_O_force_vs_sig_{tau_label}_r{R}"
                all_results[ov_key] = {
                    "scene": scene_type,
                    "zone": zname,
                    "r": R,
                    "tau": tau,
                    "overlap": overlap,
                }
                print(f"      O(force, sig)_{tau_label}: {overlap:.4f}",
                      flush=True)

    # ---- 5. Save results ----
    ccd_path = os.path.join(out_dir, "zone_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)

    # ---- 6. Print summary table ----
    print("\n" + "=" * 80, flush=True)
    print("SUMMARY: Zone CCD Results", flush=True)
    print("=" * 80, flush=True)

    for scene_type in scene_types:
        print(f"\n{'=' * 70}", flush=True)
        print(f"  {scene_type}", flush=True)
        print(f"{'=' * 70}", flush=True)
        header = (
            f"  {'Zone':<15s} | {'N_masked':>9s} | "
            f"{'force_fy':>20s} | {'sig_tau0':>20s} | {'sig_tau_corr':>22s} | "
            f"{'O_0':>6s} | {'O_corr':>6s}"
        )
        sep = "  " + "-" * (15 + 9 + 20 + 20 + 22 + 6 + 6 + 12)
        print(header, flush=True)
        print(sep, flush=True)

        for zname in zones:
            n_pts = all_results.get(
                f"{scene_type}_{zname}_force_fy_r{R}", {}
            ).get("N_masked_grid", 0)

            fd = all_results.get(f"{scene_type}_{zname}_force_fy_r{R}", {})
            sd0 = all_results.get(f"{scene_type}_{zname}_sig_tau0_r{R}", {})
            sdc = all_results.get(f"{scene_type}_{zname}_sig_tau_corr_r{R}", {})
            od0 = all_results.get(
                f"{scene_type}_{zname}_O_force_vs_sig_tau0_r{R}", {})
            odc = all_results.get(
                f"{scene_type}_{zname}_O_force_vs_sig_tau_corr_r{R}", {})

            def fmt_ccd(d):
                m = d.get("m80", "-")
                s1 = d.get("sigma_top3", ["-"])[0]
                if isinstance(s1, float):
                    return f"m80={m} s1={s1:.4f}"
                return f"m80={m} s1={s1}"

            def fmt_ov(d):
                v = d.get("overlap", "-")
                if isinstance(v, float):
                    return f"{v:.4f}"
                return f"{v}"

            print(
                f"  {zname:<15s} | {n_pts:>9d} | "
                f"{fmt_ccd(fd):>20s} | {fmt_ccd(sd0):>20s} | "
                f"{fmt_ccd(sdc):>22s} | {fmt_ov(od0):>6s} | "
                f"{fmt_ov(odc):>6s}",
                flush=True,
            )

    print(f"\n{'=' * 80}", flush=True)
    print("Done. Zone CCD analysis complete.", flush=True)
    print(f"{'=' * 80}", flush=True)
    return 0


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