# OID_analysis/li22b/phase_a3_lse.py
"""Phase A.3: LSE — [sensors, b] → POD coefficients.

Li22b Eq 2.7-2.8: T_ij solves a_i = sum_j T_ij * q_j
where q = [sensors (18 channels), b (3)] = 21 inputs.

Usage:
  PYTHONPATH="src:$PYTHONPATH" python3 src/OID_analysis/li22b/phase_a3_lse.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)

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


def load_all():
    cmd_dirs = sorted(glob.glob(os.path.join(DATA_BASE, "[0-9][0-9][0-9]")))
    all_sens, all_forces, all_b_snap, all_configs = [], [], [], []
    for d in cmd_dirs:
        sfp = os.path.join(d, "sensors.npz")
        ffp = os.path.join(d, "forces.npz")
        if not os.path.isfile(sfp): continue
        sens = np.load(sfp)["sensors"]   # (200, 18)
        force = np.load(ffp)["forces"]   # (200, 6)
        all_sens.append(sens); all_forces.append(force)
        with open(os.path.join(d, "config.json")) as f:
            cfg = json.load(f)
        all_configs.append(cfg)
        # Repeat b for each time step
        b_rep = np.tile(cfg["b"], (sens.shape[0], 1))  # (200, 3)
        all_b_snap.append(b_rep)
    S = np.concatenate(all_sens, axis=0)
    B = np.concatenate(all_b_snap, axis=0)
    N = S.shape[0]
    print(f"Loaded {N} snapshots from {len(cmd_dirs)} commands")
    return S, B, all_configs, len(cmd_dirs)


def main():
    S, B, configs, n_cmds = load_all()
    N = S.shape[0]

    # Load POD coefs (handles .npy or .npy.npz naming)
    coef_base = os.path.join(DERIVED, "pod", "pod_coefs.npy")
    for ext in [".npy", ".npy.npz", ".npz"]:
        fp = coef_base + ext
        if os.path.isfile(fp):
            pod = np.load(fp)
            break
    else:
        raise FileNotFoundError(f"No pod_coefs found")
    A = pod["coefs"]  # (N, r)
    r = A.shape[1]
    print(f"POD coefs: {A.shape}, rank={r}")

    # Build input matrix Q = [S, B] (21-dim)
    # Standardize sensors per-channel
    S_std = np.zeros_like(S)
    S_mean, S_std_scale = [], []
    for j in range(18):
        m = np.mean(S[:, j]); s = np.std(S[:, j]) or 1.0
        S_std[:, j] = (S[:, j] - m) / s
        S_mean.append(m); S_std_scale.append(s)
    Q = np.hstack([S_std, B])  # (N, 21)

    # Train/test split by commands (80/20 per Li22b)
    rng = np.random.RandomState(42)
    cmd_indices = np.arange(n_cmds)
    rng.shuffle(cmd_indices)
    n_train = int(0.8 * n_cmds)
    # Map back to snapshot indices
    snaps_per_cmd = N // n_cmds
    train_mask = np.zeros(N, dtype=bool)
    for ci in cmd_indices[:n_train]:
        train_mask[ci*snaps_per_cmd:(ci+1)*snaps_per_cmd] = True
    test_mask = ~train_mask

    # Solve LSE: T = (Q^T Q)^{-1} Q^T A
    Q_train = Q[train_mask]; A_train = A[train_mask]
    QTQ = Q_train.T @ Q_train
    T = np.linalg.solve(QTQ + 1e-6 * np.eye(21), Q_train.T @ A_train)  # (21, r)
    print(f"T matrix shape: {T.shape}")

    # Predict test set
    Q_test = Q[test_mask]; A_test = A[test_mask]
    A_pred = Q_test @ T

    # Error metrics (Li22b Eq 3.3, 3.5)
    # Per-command error ε_a(b)
    eps_a = []
    for ci in cmd_indices[n_train:]:
        i0 = ci * snaps_per_cmd; i1 = (ci + 1) * snaps_per_cmd
        mask = np.zeros(N, dtype=bool); mask[i0:i1] = True
        mask_test = mask[test_mask]
        if np.sum(mask_test) == 0: continue
        a_true = A_test[mask_test]; a_pred = A_pred[mask_test]
        num = np.mean(np.sum((a_true - a_pred)**2, axis=1))
        den = np.mean(np.sum(a_true**2, axis=1))
        eps_a.append(float(np.sqrt(num / den)) if den > 1e-30 else 0.0)
    E = float(np.sqrt(np.mean(np.array(eps_a)**2)))

    print(f"\n=== LSE Results ===")
    print(f"Train commands: {n_train}, Test: {n_cmds - n_train}")
    print(f"Mean ε_a: {np.mean(eps_a):.4f} ± {np.std(eps_a):.4f}")
    print(f"Overall E: {E:.4f} (Li22b ~0.15-0.25 for periodic)")
    print(f"Best ε_a: {np.min(eps_a):.4f}, Worst: {np.max(eps_a):.4f}")

    # Per-mode estimation error
    per_mode_mse = np.mean((A_test - A_pred)**2, axis=0)
    per_mode_energy = np.mean(A_test**2, axis=0)
    per_mode_err = np.sqrt(per_mode_mse / (per_mode_energy + 1e-30))
    print(f"\nTop 10 mode errors: {per_mode_err[:10].round(4).tolist()}")

    # Save
    out_dir = os.path.join(DERIVED, "lse")
    np.savez(os.path.join(out_dir, "lse_T.npz"), T=T,
             S_mean=S_mean, S_std=S_std_scale)
    json.dump({"n_train": n_train, "n_test": n_cmds - n_train,
               "E": E, "mean_eps_a": float(np.mean(eps_a)),
               "std_eps_a": float(np.std(eps_a)),
               "eps_a_all": eps_a,
               "per_mode_err": per_mode_err[:20].tolist()},
              open(os.path.join(out_dir, "lse_results.json"), "w"), indent=2)
    print(f"\nSaved to {out_dir}")


if __name__ == "__main__":
    main()