# OID_analysis/analysis/phase3_force_oid.py
"""
Phase 3: Force-OID for all scenes.
Cross-covariance SVD between correction-field POD coefficients and force observable.

Usage:
    python3 src/OID_analysis/analysis/phase3_force_oid.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 (  # noqa: E402
    compute_force_oid,
    standardize, reconstruct_oid_modes,
)


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


def run_force_oid(scene_key: str):
    print(f"\n--- Phase 3: Force-OID for {scene_key} ---")

    # Load POD coefficients (r=10)
    pod_dir = os.path.join(DATA_DIR, "derived", "pod", scene_key)
    pod_fp = os.path.join(pod_dir, "pod_coefs_r10.npy.npz")
    if not os.path.isfile(pod_fp):
        print(f"  WARNING: POD not found at {pod_fp}")
        return
    pod_data = np.load(pod_fp, allow_pickle=True)
    # pod_data is an npz file; need to extract coefs
    coefs = pod_data["coefs"]  # (N, r)

    # Load force observable
    obs_dir = os.path.join(DATA_DIR, "derived", "observables", scene_key)
    force_fp = os.path.join(obs_dir, "force_total.npz")
    if not os.path.isfile(force_fp):
        print(f"  WARNING: force observable not found at {force_fp}")
        return

    force_data = np.load(force_fp)
    Y_force = force_data["standardized"]  # (N, m)

    # Ensure same length
    N = min(coefs.shape[0], Y_force.shape[0])
    A = coefs[:N]
    Y = Y_force[:N]

    # Standardize POD coefs
    A_std, A_mean, A_std_val = standardize(A)

    # Compute force-OID
    oid = compute_force_oid(A_std, Y)

    # Load POD modes for reconstruction
    modes_fp = os.path.join(pod_dir, "pod_modes_r10.npz")
    if os.path.isfile(modes_fp):
        modes_data = np.load(modes_fp)
        pod_modes = modes_data["modes"]
        oid_modes = reconstruct_oid_modes(pod_modes[:, :A.shape[1]], oid["U"])
    else:
        oid_modes = None

    # Save
    out_dir = os.path.join(DATA_DIR, "derived", "oid", "force", scene_key)
    os.makedirs(out_dir, exist_ok=True)

    np.savez(os.path.join(out_dir, "force_oid.npz"),
             z=oid["z"], S=oid["S"], U=oid["U"],
             cum_corr=oid["cum_corr"])
    if oid_modes is not None:
        np.savez_compressed(os.path.join(out_dir, "force_oid_modes.npz"),
                            modes=oid_modes[:, :6])

    # Print summary
    print(f"  Singular values: {oid['S'][:6]}")
    print(f"  Cum corr (3): {oid['cum_corr'][:3]}")
    print(f"  Top 3 z std: {np.std(oid['z'][:, :3], axis=0)}")
    print(f"  Saved to {out_dir}")


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_force_oid(sn)

    print("Phase 3 complete.")


if __name__ == "__main__":
    main()
