"""Preprocess MPM trajectories (L1/L2) into chart states for training:
per frame: Kabsch rigid pose (t, R) vs the torus rest cloud -> body-frame cloud -> CHART INVERSION
(warm-started Adam over the KQ=12 chart coords) -> xi(t). Stores per-trajectory npz with
(t, R, xi, inversion residual per frame). The MPM SUB particle indices index the SAME 381,414-particle
cloud the chart was trained on, so correspondence is exact.
Note (honest): fps=24 makes FD accelerations at impact frames coarse — the A' residual masks
high-jerk frames; short-window recon (loss C) carries contact. v2 data can re-export at fps=120.
Usage: prep_mpm_traj.py L1|L2 [GPU-visible via env]"""
import glob
import os
import sys
import time
import numpy as np
import torch

sys.path.insert(0, "/home/qiuyid/neurok/experiments/repro")
from neurok_faithful import NeurokFaithful, K as KTOK, DMODEL, make_functorch_safe, dq_apply_t

DEV = "cuda"; OUT = "/scr/qiuyid/dingqy/out"; DATA = os.environ.get("MPMDATA", "/home/qiuyid/neurok/data/mpm_torus")
MODE = sys.argv[1] if len(sys.argv) > 1 else "L1"
KQ = 12; NINV = 4096; INV_STEPS = 120; INV_STEPS_WARM = int(os.environ.get("WARM", "40"))
torch.manual_seed(0)

m = NeurokFaithful().to(DEV)                                                  # CHARTCK/CODESCK: enriched-chart override
m.load_state_dict(torch.load(os.environ.get("CHARTCK", f"{OUT}/neurok_torus.pt"), map_location=DEV)); m.eval()
make_functorch_safe(m.dec)
for p in m.parameters():
    p.requires_grad_(False)
CK = torch.load(os.environ.get("CODESCK", f"{OUT}/torus_codes.pt"), weights_only=False)
zf = CK["z"].to(DEV); rest = torch.tensor(CK["rest"], device=DEV)
z_rest = zf[0].reshape(-1)
Adev = (zf[1:].reshape(len(zf) - 1, -1) - z_rest[None])
_, S_, V_ = torch.linalg.svd(Adev, full_matrices=False)
Q = V_[:KQ].contiguous()
span = float((rest.max(0).values - rest.min(0).values).mean())


def chart_c(xi, P):
    y = dq_apply_t(m.dec((z_rest + xi @ Q).view(1, KTOK, DMODEL), P[None])[0], P)
    return y - y.mean(0, keepdim=True)


def kabsch(A, B):
    H = (A - A.mean(0)).T @ (B - B.mean(0)); U, S, Vt = np.linalg.svd(H)
    return (U @ np.diag([1, 1, np.sign(np.linalg.det(U @ Vt))]) @ Vt).astype(np.float32)


files = sorted(glob.glob(f"{DATA}/{MODE}_shard*_traj*.npz"))
print(f"{MODE}: {len(files)} trajectories to preprocess", flush=True)
gsel = np.random.RandomState(3).choice(30000, NINV, replace=False)

t0 = time.time()
for fi, fpath in enumerate(files):
    opath = fpath.replace(".npz", "_prep.npz")
    if os.path.exists(opath) or fpath.endswith("_prep.npz"):
        continue
    d = np.load(fpath)
    traj = d["traj"].astype(np.float32)                                     # (F, 30000, 3)
    sub_idx = d["sub_idx"]
    Psub = rest[torch.tensor(sub_idx[gsel], device=DEV)].contiguous()       # rest coords of inversion subsample
    rest_sub = rest[torch.tensor(sub_idx, device=DEV)].cpu().numpy()
    F = traj.shape[0]
    ts, Rs, xis, errs = [], [], [], []
    xi = torch.zeros(KQ, device=DEV, requires_grad=True)
    for f in range(F):
        Rk = kabsch(rest_sub, traj[f]).T                                    # body->world
        tk = traj[f].mean(0)
        yk = (traj[f] - tk) @ Rk                                            # world -> body
        ytgt = torch.tensor(yk[gsel], device=DEV)
        ytgt = ytgt - ytgt.mean(0, keepdim=True)
        opt = torch.optim.Adam([xi], lr=0.4 if f == 0 else 0.15)
        steps = INV_STEPS if f == 0 else INV_STEPS_WARM
        for it in range(steps):
            loss = ((chart_c(xi, Psub) - ytgt) ** 2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        ts.append(tk); Rs.append(Rk); xis.append(xi.detach().cpu().numpy()); errs.append(float(loss.detach().sqrt()) / span)
    np.savez_compressed(opath + ".tmp.npz", t=np.stack(ts), R=np.stack(Rs), xi=np.stack(xis),
                        inv_err=np.array(errs), E=d["E"], nu=d["nu"], vel=d["vel"], omega=d["omega"], fps=d["fps"])
    os.replace(opath + ".tmp.npz", opath)                                   # atomic (shared volume bursts ENOSPC)
    if fi % 10 == 0:
        el = time.time() - t0
        print(f"  {fi+1}/{len(files)}: inv_err mean {np.mean(errs)*100:.2f}% span, {el/60:.1f} min, "
              f"ETA {(el/max(fi+1,1))*(len(files)-fi-1)/60:.0f} min", flush=True)
print("PREP DONE", flush=True)