# OID_analysis/li22b/phase_b.py
"""Phase B: OID on Li22b DB + Cross-mapping + Joint-input OID.

B.1: Correction-field OID: Delta_q_ctl(b) = q_ctl(b) - q_blk,
     Force-OID with Y = force per command.
B.2: Li22b POD modes (phase A.2) vs OID modes cross-mapping via subspace overlap.
B.3: Joint-input OID: Y = [forces, sensors] concatenated.

Usage:
  PYTHONPATH="src:$PYTHONPATH" python3 src/OID_analysis/li22b/phase_b.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, compute_force_oid, standardize, reconstruct_oid_modes

DATA_BASE = os.path.join(os.path.dirname(__file__), "..", "data_li22b")
DERIVED = os.path.join(DATA_BASE, "derived")

ROI_X0, ROI_X1 = 800, 1400
ROI_Y0, ROI_Y1 = 200, 400
NY_ROI = ROI_Y1 - ROI_Y0; NX_ROI = ROI_X1 - ROI_X0


def load_qblk():
    fp = os.path.join(DATA_BASE, "q_blk", "fields.npz")
    d = np.load(fp)
    return d["ux"][:, ROI_Y0:ROI_Y1, ROI_X0:ROI_X1], d["uy"][:, ROI_Y0:ROI_Y1, ROI_X0:ROI_X1]


def load_cmd(cmd_dir):
    d = np.load(os.path.join(cmd_dir, "fields.npz"))
    sens = np.load(os.path.join(cmd_dir, "sensors.npz"))["sensors"]
    force = np.load(os.path.join(cmd_dir, "forces.npz"))["forces"]
    with open(os.path.join(cmd_dir, "config.json")) as f:
        b = json.load(f)["b"]
    return (d["ux"][:, ROI_Y0:ROI_Y1, ROI_X0:ROI_X1],
            d["uy"][:, ROI_Y0:ROI_Y1, ROI_X0:ROI_X1], sens, force, b)


def snapshot_matrix(ux, uy):
    N = ux.shape[0]; DOF = NY_ROI * NX_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():
    out_dir = os.path.join(DERIVED, "oid_li22b")
    os.makedirs(out_dir, exist_ok=True)

    ux_blk, uy_blk = load_qblk()
    print(f"q_blk: {ux_blk.shape}")

    cmd_dirs = sorted(glob.glob(os.path.join(DATA_BASE, "[0-9][0-9][0-9]")))
    valid_dirs = [d for d in cmd_dirs if os.path.isfile(os.path.join(d, "forces.npz"))]
    print(f"Commands: {len(valid_dirs)}")

    # --- B.1: Correction-field OID per command ---
    all_delta = []; all_force = []; all_sens = []; all_b = []
    for d in valid_dirs:
        ux_ctl, uy_ctl, sens, force, b = load_cmd(d)
        N = min(ux_ctl.shape[0], ux_blk.shape[0])
        dux = ux_ctl[:N] - ux_blk[:N]; duy = uy_ctl[:N] - uy_blk[:N]
        all_delta.append((dux, duy))
        # Aggregate force: mean over time
        all_force.append(np.mean(force[:N], axis=0))
        all_sens.append(np.mean(sens[:N], axis=0))
        all_b.append(b)

    # Concatenate all Delta q_ctl
    ux_all = np.concatenate([d[0] for d in all_delta], axis=0)
    uy_all = np.concatenate([d[1] for d in all_delta], axis=0)
    Q = snapshot_matrix(ux_all, uy_all)
    print(f"Delta-q_ctl snapshot: {Q.shape}")

    # POD on Delta-q_ctl
    pod_li22b = compute_pod(Q, rank=min(20, Q.shape[0]-1))
    print(f"Delta-q POD: cum5={pod_li22b['cum_energy'][4]:.4f}")

    # Force-OID: Y = force per command, expanded to all snapshots
    snaps_per_cmd = ux_all.shape[0] // len(valid_dirs)
    Y_force = np.repeat(np.array(all_force), snaps_per_cmd, axis=0)[:ux_all.shape[0]]
    Yf_std, _, _ = standardize(Y_force.astype(np.float64))
    A_std, _, _ = standardize(pod_li22b["coefs"].astype(np.float64))
    oid_force = compute_force_oid(A_std, Yf_std)
    print(f"Force-OID on Li22b DB: S[0]={oid_force['S'][0]:.4f}")

    # --- B.2: Cross-mapping Li22b POD ↔ OID modes ---
    # Load Li22b full POD modes (from phase A.2)
    li22b_pod_fp = os.path.join(DERIVED, "pod", "pod_modes.npz")
    if os.path.isfile(li22b_pod_fp):
        li22b_modes = np.load(li22b_pod_fp)["modes"]
        r_min = min(li22b_modes.shape[1], pod_li22b["modes"].shape[1], 10)
        # Subspace overlap: O(Phi_Li22b, Phi_OID)
        overlap_matrix = np.zeros((r_min, r_min))
        for i in range(r_min):
            for j in range(r_min):
                overlap_matrix[i,j] = abs(np.dot(li22b_modes[:,i], pod_li22b["modes"][:,j]))
        print(f"\nLi22b POD x OID subspace overlap (first {r_min} modes):")
        print("    " + "  ".join([f"OID{j}" for j in range(r_min)]))
        for i in range(r_min):
            row = "  ".join([f"{overlap_matrix[i,j]:.3f}" for j in range(r_min)])
            print(f"Li{i:>2d} {row}")
        np.savez(os.path.join(out_dir, "crossmap.npz"), overlap=overlap_matrix)
        # Top matches
        for i in range(min(3, r_min)):
            top_j = int(np.argmax(overlap_matrix[i]))
            print(f"  Li22b mode {i} best matches OID mode {top_j} (overlap={overlap_matrix[i,top_j]:.3f})")

    # --- B.3: Joint-input OID ---
    Y_sens = np.repeat(np.array(all_sens), snaps_per_cmd, axis=0)[:ux_all.shape[0]]
    Ys_std, _, _ = standardize(Y_sens.astype(np.float64))
    Y_joint = np.hstack([Yf_std, Ys_std])
    oid_joint = compute_force_oid(A_std, Y_joint)
    print(f"\nJoint-OID: S[0]={oid_joint['S'][0]:.4f}, cum2={oid_joint['cum_corr'][1]:.4f}")

    # Compare: does joint combine both channels?
    print("\n=== Comparison ===")
    print(f"Force-OID S[0]: {oid_force['S'][0]:.4f}")
    print(f"Joint-OID S[0]: {oid_joint['S'][0]:.4f}")

    # Save
    np.savez(os.path.join(out_dir, "oid_force.npz"),
             U=oid_force["U"], S=oid_force["S"], z=oid_force["z"])
    np.savez(os.path.join(out_dir, "oid_joint.npz"),
             U=oid_joint["U"], S=oid_joint["S"], z=oid_joint["z"])
    np.savez(os.path.join(out_dir, "pod_delta_q.npz"),
             modes=pod_li22b["modes"], mean=pod_li22b["mean"],
             energy=pod_li22b["energy"])
    print(f"\nSaved to {out_dir}")


if __name__ == "__main__":
    main()
