# OID_analysis/li22b/phase_a2_pod_incremental.py
"""Memory-safe incremental POD on full Li22b DB (50 commands).

Strategy: per-command POD (one at a time) -> stack modes -> meta-POD.
Peak RAM ~2.5 GB. Previous crash: (10000, 240000) Q = 19 GB.

Usage:
  PYTHONPATH="src:$PYTHONPATH" conda run -n sr_env python3 \
    src/OID_analysis/li22b/phase_a2_pod_incremental.py
"""
import os, sys, json, glob, time
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)

X0, X1, Y0, Y1 = 800, 1400, 200, 400
DOF = (X1 - X0) * (Y1 - Y0) * 2  # 240,000
K = 25   # modes per command
R = 80   # meta-POD truncation


def build_Q(ux, uy):
    N = ux.shape[0]; 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 project(ux, uy, modes, mean):
    Q = build_Q(ux, uy)
    return (Q - mean.reshape(1, -1)) @ modes


def main():
    t0 = time.time()
    dirs = sorted(glob.glob(os.path.join(DATA_BASE, "[0-9][0-9][0-9]")))
    n_cmds = len(dirs)
    if n_cmds == 0:
        print("No commands found!"); return
    print(f"{n_cmds} cmds, DOF={DOF}, K={K}, est peak: {DOF*K*n_cmds*8/1e9:.1f}GB + 0.4GB\n")

    # -- Per-command POD --
    modes_list, means_list, ns_list, b_list = [], [], [], []
    for i, d in enumerate(dirs):
        fp = os.path.join(d, "fields.npz")
        if not os.path.isfile(fp): continue
        data = np.load(fp)
        ux = data["ux"][:, Y0:Y1, X0:X1]; uy = data["uy"][:, Y0:Y1, X0:X1]
        N = ux.shape[0]
        Q = build_Q(ux, uy)
        pod = compute_pod(Q, rank=K)
        del Q
        modes_list.append(pod["modes"])
        means_list.append(pod["mean"])
        ns_list.append(N)
        with open(os.path.join(d, "config.json")) as f:
            b_list.append(json.load(f)["b"])
        t1 = time.time()
        eta = (t1 - t0) * (n_cmds - i - 1) / max(i + 1, 1) / 60
        print(f"  [{i:3d}] N={N}, cum5={pod['cum_energy'][4]:.4f}, {t1-t0:.0f}s, eta={eta:.0f}min")

    # -- Meta-POD: stack modes --
    M = np.column_stack(modes_list)  # (DOF, n_cmds*K)
    print(f"\nM_stack: {M.shape}, {M.nbytes/1e9:.1f}GB")
    M_centered = M - np.mean(M, axis=1, keepdims=True)
    U, S, Vt = np.linalg.svd(M_centered, full_matrices=False)
    S = S.astype(np.float64)
    gmodes = U[:, :R]
    gS = S[:R]
    energy = gS**2 / np.sum(S**2)
    cum = np.cumsum(energy)
    r99 = int(np.searchsorted(cum, 0.99)) + 1
    print(f"Meta-POD: S0={gS[0]:.1f}, cum5={cum[4]:.4f}, cum10={cum[9]:.4f}, r99={r99}")
    print(f"Li22b ref: E2=44.9%, E10=78.9%, r99=78")

    # -- Project snapshots --
    gmean = np.mean(np.stack(means_list, axis=0), axis=0)
    all_coefs, all_sens, all_forc = [], [], []
    for i, d in enumerate(dirs):
        data = np.load(os.path.join(d, "fields.npz"))
        ux = data["ux"][:, Y0:Y1, X0:X1]; uy = data["uy"][:, Y0:Y1, X0:X1]
        all_coefs.append(project(ux, uy, gmodes, gmean).astype(np.float64))
        sfp = os.path.join(d, "sensors.npz")
        if os.path.isfile(sfp):
            all_sens.append(np.load(sfp)["sensors"].astype(np.float64))
            all_forc.append(np.load(os.path.join(d, "forces.npz"))["forces"].astype(np.float64))
        if (i + 1) % 15 == 0:
            print(f"  projected {i+1}/{n_cmds}")

    coefs = np.concatenate(all_coefs, axis=0)
    sensors = np.concatenate(all_sens, axis=0) if all_sens else None
    forces = np.concatenate(all_forc, axis=0) if all_forc else None
    print(f"Coefs: {coefs.shape}, sensors: {sensors.shape if sensors is not None else 'N/A'}")

    # Save
    np.savez_compressed(os.path.join(DERIVED, "pod_modes.npz"), modes=gmodes, mean=gmean)
    np.savez(os.path.join(DERIVED, "pod_coefs.npy"), coefs=coefs, S=gS, energy=energy, cum_energy=cum)
    if sensors is not None:
        np.savez(os.path.join(DERIVED, "sensors_all.npy"), sensors=sensors)
        np.savez(os.path.join(DERIVED, "forces_all.npy"), forces=forces)
    json.dump({"n_cmds": n_cmds, "n_snaps": coefs.shape[0], "dof": DOF,
               "k": K, "r": R, "r99": r99,
               "e2": float(cum[1]), "e10": float(cum[9]),
               "top10": energy[:10].tolist()},
              open(os.path.join(DERIVED, "energy.json"), "w"), indent=2)
    print(f"Saved. Total: {(time.time()-t0)/60:.1f}min")


if __name__ == "__main__":
    main()
