# OID_analysis/analysis/robustness_analysis.py
"""
Comprehensive robustness analysis for force-OID vs signature-OID separation.
Covers:
1. POD rank sensitivity (r=6,8,10,12,16) -- all 5 scenes
2. Time-window sensitivity (split data into 2 halves)  
3. Karman tau_c sensitivity sweep
4. Zone OID: force-OID / sig-OID in near-body / near-wake / downstream zones

Usage:
    cd /home/frank14f/DynamisLab && PYTHONPATH="src:$PYTHONPATH" conda run -n sr_env python3 src/OID_analysis/analysis/robustness_analysis.py
"""
from __future__ import annotations

import json
import os
import sys

import numpy as np

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

from OID_analysis.configs import DATA_DIR  # noqa: E402
from OID_analysis.utils.analysis import (  # noqa: E402
    compute_pod, compute_force_oid, compute_pcd, standardize,
)

SCENES = ["steady_cloak", "karman_re100",
          "illusion_0.75L", "illusion_1.0L", "illusion_1.5L"]


def cos_sim(a, b):
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))


def load_pod_coefs(scene: str, rank: int = 10):
    """Load correction-field POD coefficients for a scene at given rank."""
    pod_dir = os.path.join(DATA_DIR, "derived", "pod", scene)
    fp = os.path.join(pod_dir, f"pod_coefs_r{rank}.npy.npz")
    if not os.path.isfile(fp):
        # Try loading delta fields and re-running POD
        delta_fp = os.path.join(pod_dir, "delta_q_ctl.npz")
        if not os.path.isfile(delta_fp):
            return None
        d = np.load(delta_fp)
        ux, uy = d["ux"], d["uy"]
        # Flatten to snapshot matrix
        N, ny, nx = ux.shape
        Q = np.zeros((N, ny * nx * 2), dtype=np.float64)
        for t in range(N):
            Q[t] = np.concatenate([ux[t].ravel(), uy[t].ravel()])
        pod = compute_pod(Q, rank=rank)
        return pod["coefs"]
    else:
        d = np.load(fp, allow_pickle=True)
        return d["coefs"]


def load_observable(scene: str, obs_name: str):
    """Load standardized observable."""
    obs_dir = os.path.join(DATA_DIR, "derived", "observables", scene)
    fp = os.path.join(obs_dir, f"{obs_name}.npz")
    if not os.path.isfile(fp):
        return None
    d = np.load(fp)
    return d["standardized"]


def compute_overlap(coefs_a, obs_a, coefs_b, obs_b):
    """Compute cosine similarity between force-OID mode 1 and sig-OID mode 1."""
    N = min(coefs_a.shape[0], obs_a.shape[0], coefs_b.shape[0], obs_b.shape[0])
    A1 = standardize(coefs_a[:N])[0]
    A2 = standardize(coefs_b[:N])[0]
    Y1 = obs_a[:N]
    Y2 = obs_b[:N]

    # Force-OID
    C1 = (1.0 / N) * A1.T @ Y1
    U1, _, _ = np.linalg.svd(C1, full_matrices=False)

    # Sig-OID
    C2 = (1.0 / N) * A2.T @ Y2
    U2, _, _ = np.linalg.svd(C2, full_matrices=False)

    return cos_sim(U1[:, 0], U2[:, 0])


def compute_r2_sweep(A, Y, max_m=5):
    """Compute R2 for each m=1..max_m using 70/30 split."""
    from sklearn.linear_model import LinearRegression
    from sklearn.metrics import r2_score

    N = min(A.shape[0], Y.shape[0])
    A = A[:N]
    Y = Y[:N]
    split = int(N * 0.7)

    results = {}
    for m in range(1, max_m + 1):
        X = standardize(A[:, :m])[0]
        X_train, X_test = X[:split], X[split:]
        Y_train, Y_test = Y[:split], Y[split:]
        reg = LinearRegression().fit(X_train, Y_train)
        r2 = r2_score(Y_test, reg.predict(X_test))
        results[m] = float(max(r2, -1.0))  # clamp
    return results


def main():
    results = {
        "rank_sensitivity": {},
        "window_sensitivity": {},
        "tauc_sensitivity": {},
        "zone_oid": {},
        "overlap_table": {},
    }

    # ===========================
    # 1. POD RANK SENSITIVITY
    # ===========================
    print("=" * 60)
    print("1. POD RANK SENSITIVITY")
    print("=" * 60)
    ranks = [6, 8, 10, 12, 16]
    overlap_by_rank = {}
    for scene in SCENES:
        print(f"\n  Scene: {scene}")
        # Force observable
        y_force = load_observable(scene, "force_total")
        y_sig = load_observable(scene, "sensor_error_delayed")
        if scene == "steady_cloak":
            y_sig = load_observable(scene, "rms_uy")

        if y_force is None:
            print("    SKIP: no force observable")
            continue
        if y_sig is None:
            print("    SKIP: no sig observable")
            continue

        scene_rank_data = {}
        for r in ranks:
            coefs = load_pod_coefs(scene, rank=r)
            if coefs is None:
                print(f"    r={r}: no POD data")
                continue

            overlap = compute_overlap(coefs, y_force, coefs, y_sig)
            scene_rank_data[r] = overlap
            print(f"    r={r}: overlap={overlap:.4f}")

        overlap_by_rank[scene] = scene_rank_data

    results["rank_sensitivity"] = overlap_by_rank

    # ===========================
    # 2. TIME-WINDOW SENSITIVITY
    # ===========================
    print("\n" + "=" * 60)
    print("2. TIME-WINDOW SENSITIVITY")
    print("=" * 60)
    for scene in SCENES:
        print(f"\n  Scene: {scene}")
        # Load delta-q_ctl fields
        pod_dir = os.path.join(DATA_DIR, "derived", "pod", scene)
        delta_fp = os.path.join(pod_dir, "delta_q_ctl.npz")
        if not os.path.isfile(delta_fp):
            print("    SKIP: no delta fields")
            continue
        d = np.load(delta_fp)
        ux, uy = d["ux"], d["uy"]
        N = ux.shape[0]

        y_force = load_observable(scene, "force_total")
        y_sig = load_observable(scene, "sensor_error_delayed")
        if scene == "steady_cloak":
            y_sig = load_observable(scene, "rms_uy")
        if y_force is None or y_sig is None:
            continue

        # Split into 2 windows
        half = N // 2
        window_results = {}
        for wname, start, end in [("first_half", 0, half), ("second_half", half, 2 * half)]:
            # POD on this window
            Q = np.zeros((end - start, ux.shape[1] * ux.shape[2] * 2), dtype=np.float64)
            for t in range(start, end):
                Q[t - start] = np.concatenate([ux[t].ravel(), uy[t].ravel()])
            pod = compute_pod(Q, rank=10)

            A = pod["coefs"]
            Y_f = y_force[start:end]
            Y_s = y_sig[start:end]
            overlap = compute_overlap(A, Y_f, A, Y_s)
            window_results[wname] = overlap
            print(f"    {wname}: overlap={overlap:.4f}")

        # Full window reference
        if "first_half" in window_results and "second_half" in window_results:
            print(f"    delta: {abs(window_results['first_half'] - window_results['second_half']):.4f}")

        results["window_sensitivity"][scene] = window_results

    # ===========================
    # 3. KARMAN tau_c SENSITIVITY
    # ===========================
    print("\n" + "=" * 60)
    print("3. KARMAN tau_c SENSITIVITY")
    print("=" * 60)

    # Load Karman data
    pod_dir = os.path.join(DATA_DIR, "derived", "pod", "karman_re100")
    delta_fp = os.path.join(pod_dir, "delta_q_ctl.npz")
    obs_dir = os.path.join(DATA_DIR, "derived", "observables", "karman_re100")

    if os.path.isfile(delta_fp) and os.path.isdir(obs_dir):
        d = np.load(delta_fp)
        ux, uy = d["ux"], d["uy"]
        N = ux.shape[0]

        # Load target for re-computing delayed error
        ctl_fp = os.path.join(
            os.path.join(DATA_DIR, "karman_cloak", "karman_re100"), "controlled.npz")
        if os.path.isfile(ctl_fp):
            ctl = np.load(ctl_fp)
            sensors = ctl["sensors"]
            target_dd = os.path.join(DATA_DIR, "karman_cloak", "karman_re100")
            target_fp = os.path.join(target_dd, "target.npz")
            if os.path.isfile(target_fp):
                target_states = np.load(target_fp)["target_states"]
            else:
                target_states = sensors * 0.0

            # Force observable
            forces = ctl["forces"]
            Fx_total = np.sum(forces[:, 0::2], axis=1, keepdims=True)
            Fy_total = np.sum(forces[:, 1::2], axis=1, keepdims=True)
            y_force_std, _, _ = standardize(np.hstack([Fx_total, Fy_total]))

            tlen = target_states.shape[0]
            tau_candidates = [0, 10, 15, 20, 25, 30, 35, 40, 50, 60]
            NT = 30  # approx shedding period in steps

            # POD
            Q = np.zeros((N, ux.shape[1] * ux.shape[2] * 2), dtype=np.float64)
            for t in range(N):
                Q[t] = np.concatenate([ux[t].ravel(), uy[t].ravel()])
            pod = compute_pod(Q, rank=10)
            A_std, _, _ = standardize(pod["coefs"])

            tauc_results = {}
            for tc in tau_candidates:
                e_s_delayed = np.zeros_like(sensors)
                if N > tc:
                    # Create a reference that matches N length
                    ref_full = np.zeros((N, 6), dtype=np.float32)
                    for i in range(6):
                        if target_states.shape[1] > i:
                            t_col = target_states[:, i]
                            if len(t_col) < N:
                                # tile the target to match N
                                repeats = (N + len(t_col) - 1) // len(t_col)
                                ref_full[:, i] = np.tile(t_col, repeats)[:N]
                            else:
                                ref_full[:, i] = t_col[:N]
                        else:
                            ref_full[:, i] = sensors[:, i]
                        e_s_delayed[tc:, i] = sensors[tc:, i] - ref_full[tc:, i]
                        e_s_delayed[:tc, i] = 0.0

                y_sig_std, _, _ = standardize(e_s_delayed)

                # Compute overlap
                N_use = min(A_std.shape[0], y_force_std.shape[0], y_sig_std.shape[0])
                A_u = A_std[:N_use]
                Y_f = y_force_std[:N_use]
                Y_s = y_sig_std[:N_use]

                C_f = (1.0 / N_use) * A_u.T @ Y_f
                C_s = (1.0 / N_use) * A_u.T @ Y_s
                U_f, _, _ = np.linalg.svd(C_f, full_matrices=False)
                U_s, _, _ = np.linalg.svd(C_s, full_matrices=False)
                overlap = cos_sim(U_f[:, 0], U_s[:, 0])

                # R2 for force prediction
                r2_force = compute_r2_sweep(A_u, Y_f, max_m=3)
                r2_sig = compute_r2_sweep(A_u, Y_s, max_m=3)

                tauc_results[tc] = {
                    "overlap": overlap,
                    "r2_force_m2": r2_force.get(2, None),
                    "r2_sig_m2": r2_sig.get(2, None),
                }
                print(f"    tau_c={tc:3d}: overlap={overlap:.4f}, "
                      f"force_R2={r2_force.get(2, None):.4f}, sig_R2={r2_sig.get(2, None):.4f}")

            results["tauc_sensitivity"]["karman_re100"] = tauc_results

    # ===========================
    # 4. ZONE OID
    # ===========================
    print("\n" + "=" * 60)
    print("4. ZONE OID")
    print("=" * 60)

    # Zone masks (in lattice coords for 1280x512)
    ny, nx = 512, 1280
    zones = {
        "near-body": (200, 310, 580, 660),
        "near-wake": (180, 330, 660, 800),
        "downstream": (180, 330, 790, 810),
    }

    for scene in SCENES:
        print(f"\n  Scene: {scene}")
        pod_dir = os.path.join(DATA_DIR, "derived", "pod", scene)
        delta_fp = os.path.join(pod_dir, "delta_q_ctl.npz")
        if not os.path.isfile(delta_fp):
            print("    SKIP: no delta fields")
            continue
        d = np.load(delta_fp)
        ux, uy = d["ux"], d["uy"]
        N = ux.shape[0]
        ny_a, nx_a = ux.shape[1], ux.shape[2]
        print(f"    field shape: ({ny_a}, {nx_a})")

        y_force = load_observable(scene, "force_total")
        y_sig = load_observable(scene, "sensor_error_delayed")
        if scene == "steady_cloak":
            y_sig = load_observable(scene, "rms_uy")
        if y_force is None or y_sig is None:
            print("    SKIP: missing observable")
            continue

        try:
            for zname, (y0, y1, x0, x1) in zones.items():
                y0 = max(0, y0); y1 = min(ny_a, y1); x0 = max(0, x0); x1 = min(nx_a, x1)
                if y1 <= y0 or x1 <= x0: continue
                
                ux_z = ux[:, y0:y1, x0:x1]
                uy_z = uy[:, y0:y1, x0:x1]
                DOF = ux_z.shape[1] * ux_z.shape[2] * 2
                if DOF == 0: continue
                
                Q = np.zeros((N, DOF), dtype=np.float64)
                for t in range(N):
                    Q[t] = np.concatenate([ux_z[t].ravel(), uy_z[t].ravel()])
                pod = compute_pod(Q, rank=10)
                A = pod["coefs"]
                min_n = min(A.shape[0], y_force.shape[0], y_sig.shape[0])
                A_s = A[:min_n]; Y_f = y_force[:min_n]; Y_s = y_sig[:min_n]
                C_f = (1.0 / min_n) * A_s.T @ Y_f
                C_s = (1.0 / min_n) * A_s.T @ Y_s
                U_f, S_f, _ = np.linalg.svd(C_f, full_matrices=False)
                U_s, S_s, _ = np.linalg.svd(C_s, full_matrices=False)
                overlap = cos_sim(U_f[:, 0], U_s[:, 0])
                scene_zone_results[zname] = {"overlap": overlap,
                    "force_S0": float(S_f[0]) if len(S_f) > 0 else None,
                    "sig_S0": float(S_s[0]) if len(S_s) > 0 else None}
                print(f"    {zname:15s}: overlap={overlap:.4f}, force_S0={S_f[0]:.4f}")
        except Exception as e:
            print(f"    Zone analysis error: {e}")
            for t in range(N):
                Q[t] = np.concatenate([ux_z[t].ravel(), uy_z[t].ravel()])

            pod = compute_pod(Q, rank=10)
            A = pod["coefs"]

            # Limit observable length
            min_n = min(A.shape[0], y_force.shape[0], y_sig.shape[0])
            A_s = A[:min_n]
            Y_f = y_force[:min_n]
            Y_s = y_sig[:min_n]

            C_f = (1.0 / min_n) * A_s.T @ Y_f
            C_s = (1.0 / min_n) * A_s.T @ Y_s
            U_f, S_f, _ = np.linalg.svd(C_f, full_matrices=False)
            U_s, S_s, _ = np.linalg.svd(C_s, full_matrices=False)

            overlap = cos_sim(U_f[:, 0], U_s[:, 0])
            scene_zone_results[zname] = {
                "overlap": overlap,
                "force_S0": float(S_f[0]) if len(S_f) > 0 else None,
                "sig_S0": float(S_s[0]) if len(S_s) > 0 else None,
            }
            print(f"    {zname:15s}: overlap={overlap:.4f}, "
                  f"force_S0={S_f[0]:.4f}, sig_S0={S_s[0]:.4f}")

        results["zone_oid"][scene] = scene_zone_results

    # ===========================
    # FINAL OVERLAP TABLE (r=10)
    # ===========================
    print("\n" + "=" * 60)
    print("FINAL OVERLAP TABLE (r=10, full data)")
    print("=" * 60)

    print(f"\n{'Scene':<20s} {'Overlap':>8s} {'Force_R2_m2':>12s} {'Sig_R2_m2':>12s}")
    print("-" * 52)
    overlap_table = {}
    for scene in SCENES:
        y_force = load_observable(scene, "force_total")
        y_sig = load_observable(scene, "sensor_error_delayed")
        if scene == "steady_cloak":
            y_sig = load_observable(scene, "rms_uy")
        if y_force is None or y_sig is None:
            continue

        coefs = load_pod_coefs(scene, rank=10)
        if coefs is None:
            continue

        N = min(coefs.shape[0], y_force.shape[0], y_sig.shape[0])
        A_s, _, _ = standardize(coefs[:N])
        Y_f = y_force[:N]
        Y_s = y_sig[:N]

        C_f = (1.0 / N) * A_s.T @ Y_f
        C_s = (1.0 / N) * A_s.T @ Y_s
        U_f, _, _ = np.linalg.svd(C_f, full_matrices=False)
        U_s, _, _ = np.linalg.svd(C_s, full_matrices=False)
        overlap = cos_sim(U_f[:, 0], U_s[:, 0])

        r2_f = compute_r2_sweep(A_s, Y_f, max_m=3)
        r2_s = compute_r2_sweep(A_s, Y_s, max_m=3)

        overlap_table[scene] = {
            "overlap": overlap,
            "force_R2_m2": r2_f.get(2, None),
            "force_R2_m3": r2_f.get(3, None),
            "sig_R2_m2": r2_s.get(2, None),
            "sig_R2_m3": r2_s.get(3, None),
        }
        print(f"{scene:<20s} {overlap:>8.4f} {r2_f.get(2, 0):>12.4f} {r2_s.get(2, 0):>12.4f}")

    results["overlap_table"] = overlap_table

    # ===========================
    # SAVE
    # ===========================
    out_dir = os.path.join(DATA_DIR, "derived", "robustness")
    os.makedirs(out_dir, exist_ok=True)

    # Convert numpy arrays to floats
    def clean(obj):
        if isinstance(obj, np.floating):
            return float(obj)
        if isinstance(obj, np.integer):
            return int(obj)
        if isinstance(obj, dict):
            return {k: clean(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple)):
            return [clean(v) for v in obj]
        return obj

    results_clean = clean(results)
    with open(os.path.join(out_dir, "robustness_results.json"), "w") as f:
        json.dump(results_clean, f, indent=2)

    print(f"\nAll results saved to {out_dir}/robustness_results.json")


if __name__ == "__main__":
    main()
