# OID_analysis/li22b/phase_a2_pod.py
"""Phase A.2: Snapshot POD on Li22b open-loop database.

Replicates Li22b 4.1:
  - Load all commands, compute ensemble mean over b,t
  - POD via method-of-snapshots on ROI (x=[800:1400], y=[200:400])
  - Truncate to 99% energy
  - Visualize mode shapes + physical interpretation

Usage (after data collected):
  PYTHONPATH="src:$PYTHONPATH" python3 src/OID_analysis/li22b/phase_a2_pod.py
"""
import os, sys, json, glob
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.utils.analysis import compute_pod

DATA_BASE = os.path.join(os.path.dirname(__file__), "..", "data_li22b")
DERIVED = os.path.join(DATA_BASE, "derived", "pod")
os.makedirs(DERIVED, exist_ok=True)

NX, NY = 2000, 600
# ROI: pinball wake region
ROI_X0, ROI_X1 = 800, 1400
ROI_Y0, ROI_Y1 = 200, 400
NX_ROI = ROI_X1 - ROI_X0; NY_ROI = ROI_Y1 - ROI_Y0


def load_all_fields():
    """Load all ux, uy snapshots from all commands, plus b vectors."""
    cmd_dirs = sorted(glob.glob(os.path.join(DATA_BASE, "[0-9][0-9][0-9]")))
    all_ux, all_uy, all_b = [], [], []
    for d in cmd_dirs:
        fp = os.path.join(d, "fields.npz")
        if not os.path.isfile(fp): continue
        try:
            data = np.load(fp)
        except Exception as e:
            print(f"  SKIP {os.path.basename(d)}: {e}")
            continue
        ux = data["ux"][:, ROI_Y0:ROI_Y1, ROI_X0:ROI_X1]
        uy = data["uy"][:, ROI_Y0:ROI_Y1, ROI_X0:ROI_X1]
        all_ux.append(ux); all_uy.append(uy)
        with open(os.path.join(d, "config.json")) as f:
            cfg = json.load(f)
        all_b.append(cfg["b"])
        print(f"  Loaded {os.path.basename(d)}: {ux.shape[0]} snaps, b={cfg['b']}")
    return (np.concatenate(all_ux, axis=0), np.concatenate(all_uy, axis=0),
            np.array(all_b), len(cmd_dirs))


def build_snapshot_matrix(ux, uy):
    N = ux.shape[0]
    DOF = NX_ROI * NY_ROI * 2
    Q = np.zeros((N, DOF), dtype=np.float64)
    for t in range(N):
        Q[t] = np.concatenate([ux[t].ravel(), uy[t].ravel()])
    return Q


def main():
    ux_all, uy_all, b_all, n_cmds = load_all_fields()
    N = ux_all.shape[0]
    print(f"\nTotal: {n_cmds} commands, {N} snapshots, ROI={NX_ROI}x{NY_ROI}")

    # POD
    Q = build_snapshot_matrix(ux_all, uy_all)
    print(f"Snapshot matrix: {Q.shape} ({Q.nbytes/1e9:.2f} GB)")
    pod = compute_pod(Q, rank=min(100, N-1))
    S = pod["S"]; energy = pod["energy"]; cum = pod["cum_energy"]
    print(f"POD: S[0]={S[0]:.2e}, cum10={cum[9]:.4f}, cum2={cum[1]:.4f}")

    # Truncation to 99%
    r99 = int(np.searchsorted(cum, 0.99)) + 1
    print(f"99% energy: {r99} modes (Li22b: 78 modes)")

    # Save
    np.savez_compressed(os.path.join(DERIVED, "pod_modes.npz"),
                        modes=pod["modes"], mean=pod["mean"])
    np.savez(os.path.join(DERIVED, "pod_coefs.npy"),
             coefs=pod["coefs"], S=S, energy=energy, cum_energy=cum)
    json.dump({"n_commands": n_cmds, "n_snapshots": N,
               "dof": Q.shape[1], "roi": [ROI_X0,ROI_X1,ROI_Y0,ROI_Y1],
               "r99": r99, "e10": float(cum[9]), "e2": float(cum[1]),
               "e1": float(energy[0]), "top10": energy[:10].tolist()},
              open(os.path.join(DERIVED, "energy.json"), "w"), indent=2)

    # Comparison table
    print(f"\n=== Energy Comparison ===")
    print(f"Li22b: e1+2=44.9%, e1-10=78.9%, r99=78")
    print(f"Ours:  e1+2={cum[1]*100:.1f}%, e1-10={cum[9]*100:.1f}%, r99={r99}")
    print(f"\nResults saved to {DERIVED}")


if __name__ == "__main__":
    main()
