# OID_analysis/analysis/phase4b_signature_pcd.py
"""
Phase 4b: Signature-PCD (whitened cross-correlation).
Uses delay-stacked signature observable p_sig(t) with whitening.

Only for periodic scenes: karman_re100, illusion_0.75L, 1.0L, 1.5L.

Usage:
    python3 src/OID_analysis/analysis/phase4b_signature_pcd.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_pcd, standardize, reconstruct_oid_modes,
)


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


def run_signature_pcd(scene_key: str):
    print(f"\n--- Phase 4b: Signature-PCD for {scene_key} ---")

    # Load POD coefs (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")
        return
    pod_npz = np.load(pod_fp, allow_pickle=True)
    coefs = pod_npz["coefs"]

    # Load delay-stacked observable
    obs_dir = os.path.join(DATA_DIR, "derived", "observables", scene_key)
    p_sig_fp = os.path.join(obs_dir, "p_sig_stack.npz")
    if not os.path.isfile(p_sig_fp):
        print(f"  WARNING: p_sig_stack not found at {p_sig_fp}")
        return

    p_data = np.load(p_sig_fp)
    P = p_data["standardized"]

    N = min(coefs.shape[0], P.shape[0])
    A = coefs[:N]
    P_cut = P[:N]
    A_std, _, _ = standardize(A)

    # PCD
    pcd = compute_pcd(A_std, P_cut, tikhonov_eps=1e-6)

    # Reconstruct modes
    modes_fp = os.path.join(pod_dir, "pod_modes_r10.npz")
    pcd_modes = None
    if os.path.isfile(modes_fp):
        md = np.load(modes_fp)
        pod_modes = md["modes"]
        pcd_modes = reconstruct_oid_modes(pod_modes[:, :A.shape[1]], pcd["W"])

    # Save
    out_dir = os.path.join(DATA_DIR, "derived", "oid", "pcd", scene_key)
    os.makedirs(out_dir, exist_ok=True)
    np.savez(os.path.join(out_dir, "sig_pcd.npz"),
             z_pcd=pcd["z_pcd"], S=pcd["S"], W=pcd["W"], cum_corr=pcd["cum_corr"])
    if pcd_modes is not None:
        np.savez_compressed(os.path.join(out_dir, "sig_pcd_modes.npz"),
                            modes=pcd_modes[:, :6])

    print(f"  PCD sig vals: {pcd['S'][:6]}")
    print(f"  Cum corr (3): {pcd['cum_corr'][:3]}")


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

    print("Phase 4b complete.")


if __name__ == "__main__":
    main()
