# OID_analysis/analysis/phase6_comparison.py
"""
Phase 6: POD vs OID vs PCD predictive comparison.

Linear regression with 70/30 time-series split.
Compares:
  - force ~ POD coords
  - force ~ OID coords
  - future signature ~ POD coords
  - future signature ~ OID coords
  - future signature ~ PCD coords

Usage:
    python3 src/OID_analysis/analysis/phase6_comparison.py
"""
from __future__ import annotations

import argparse
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)

from OID_analysis.configs import DATA_DIR  # noqa: E402
from OID_analysis.utils.analysis import standardize  # noqa: E402

try:
    from sklearn.linear_model import LinearRegression
    from sklearn.metrics import r2_score
    HAS_SKLEARN = True
except ImportError:
    HAS_SKLEARN = False
    # sklearn unavailable — comparison will be skipped


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


def r2_manual(y_true, y_pred):
    ss_res = np.sum((y_true - y_pred) ** 2)
    ss_tot = np.sum((y_true - np.mean(y_true, axis=0, keepdims=True)) ** 2)
    return 1.0 - ss_res / (ss_tot + 1e-30)


def run_comparison(scene_key: str):
    print(f"\n--- Phase 6: Comparison for {scene_key} ---")

    pod_dir = os.path.join(DATA_DIR, "derived", "pod", scene_key)
    obs_dir = os.path.join(DATA_DIR, "derived", "observables", scene_key)
    oid_dir = os.path.join(DATA_DIR, "derived", "oid")

    # Load POD coefs (r=10)
    pod_fp = os.path.join(pod_dir, "pod_coefs_r10.npy.npz")
    if not os.path.isfile(pod_fp):
        print(f"  SKIPPED: POD not found")
        return False

    pod_npz = np.load(pod_fp, allow_pickle=True)
    A_all = pod_npz["coefs"]  # (N, r)
    N = A_all.shape[0]

    # OID coordinates
    oid_sources = {}

    # Try loading force-OID
    force_oid_fp = os.path.join(oid_dir, "force", scene_key, "force_oid.npz")
    if os.path.isfile(force_oid_fp):
        d = np.load(force_oid_fp)
        oid_sources["force-oid"] = d["z"]

    # Try loading signature-OID (delayed)
    sig_oid_fp = os.path.join(oid_dir, "signature", scene_key, "sig_oid_delayed.npz")
    if os.path.isfile(sig_oid_fp):
        d = np.load(sig_oid_fp)
        oid_sources["sig-oid"] = d["z"]

    # Try loading PCD
    pcd_fp = os.path.join(oid_dir, "pcd", scene_key, "sig_pcd.npz")
    if os.path.isfile(pcd_fp):
        d = np.load(pcd_fp)
        oid_sources["sig-pcd"] = d["z_pcd"]

    # Target observables for regression
    targets = {}

    # Force observable
    force_fp = os.path.join(obs_dir, "force_total.npz")
    if os.path.isfile(force_fp):
        targets["force"] = np.load(force_fp)["standardized"]

    # Sensor error (delayed)
    sig_fp = os.path.join(obs_dir, "sensor_error_delayed.npz")
    if os.path.isfile(sig_fp):
        targets["future_sig"] = np.load(sig_fp)["standardized"]

    # For steady cloak: RMS
    rms_fp = os.path.join(obs_dir, "rms_uy.npz")
    if os.path.isfile(rms_fp):
        targets["suppression"] = np.load(rms_fp)["standardized"]

    if not HAS_SKLEARN:
        print("  sklearn not available, using manual R^2")
        # Still do manual comparison for a single test
        if "force" in targets and "force-oid" in oid_sources:
            A_std, _, _ = standardize(A_all[:N])
            Y = targets["force"][:N].squeeze()
            r2_pod = r2_manual(Y, A_std[:, :3] @ np.linalg.lstsq(A_std[:, :3], Y, rcond=None)[0])
            r2_oid = r2_manual(Y, oid_sources["force-oid"][:N, :3] @
                              np.linalg.lstsq(oid_sources["force-oid"][:N, :3], Y, rcond=None)[0])
            print(f"  force: POD R2={r2_pod:.4f}, OID R2={r2_oid:.4f}")
        print("  Install sklearn for full comparison")
        return True

    results = {}
    for target_name, Y in targets.items():
        Y = Y[:N]
        results[target_name] = {}

        for src_name, z in oid_sources.items():
            z = z[:N]

            # 70/30 time-series split
            split = int(N * 0.7)
            for m in [1, 2, 3, 5]:
                X = z[:, :m]
                X_train, X_test = X[:split], X[split:]
                Y_train, Y_test = Y[:split], Y[split:]

                reg = LinearRegression().fit(X_train, Y_train)
                Y_pred = reg.predict(X_test)
                r2 = r2_score(Y_test, Y_pred) if Y_test.ndim == 1 else \
                     r2_score(Y_test, Y_pred, multioutput="variance_weighted")
                results[target_name][f"{src_name}_m{m}"] = float(r2)

        # Also POD baseline
        A_std, _, _ = standardize(A_all[:N])
        for m in [1, 2, 3, 5]:
            X = A_std[:, :m]
            X_train, X_test = X[:split], X[split:]
            Y_train, Y_test = Y[:split], Y[split:]
            reg = LinearRegression().fit(X_train, Y_train)
            Y_pred = reg.predict(X_test)
            r2 = r2_score(Y_test, Y_pred) if Y_test.ndim == 1 else \
                 r2_score(Y_test, Y_pred, multioutput="variance_weighted")
            results[target_name][f"pod_m{m}"] = float(r2)

        print(f"  {target_name}:")
        for k, v in sorted(results[target_name].items()):
            print(f"    {k}: R2={v:.4f}")

    # Save
    out_dir = os.path.join(DATA_DIR, "derived", "comparison")
    os.makedirs(out_dir, exist_ok=True)
    with open(os.path.join(out_dir, f"{scene_key}.json"), "w") as f:
        json.dump(results, f, indent=2)

    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, default=None)
    args = ap.parse_args()

    targets = [args.scene] if args.scene and args.scene in SCENES else SCENES

    for sn in targets:
        run_comparison(sn)

    print("\nPhase 6 complete.")


if __name__ == "__main__":
    main()
