"""MotionDNA v0 — modelo generativo de movimiento humano entrenado con la pose
extraida de music videos. Transformer causal que predice la pose del frame t+1
dada la secuencia 0..t. Genera movimiento nuevo por rollout autoregresivo.
Pensado para caber holgado en una L4 (23GB)."""
import os, glob, math, time, json
import numpy as np
import torch, torch.nn as nn
from torch.utils.data import Dataset, DataLoader

POSES = os.path.expanduser("~/motiondna/poses")
CKPT  = os.path.expanduser("~/motiondna/checkpoints")
SAMP  = os.path.expanduser("~/motiondna/samples")
os.makedirs(CKPT, exist_ok=True); os.makedirs(SAMP, exist_ok=True)
DEV = "cuda" if torch.cuda.is_available() else "cpu"
WIN = 64            # ventana temporal
DIM = 34           # 17 keypoints * 2
D_MODEL = 256; N_HEAD = 8; N_LAYER = 6; DROP = 0.1
BATCH = 64; EPOCHS = 60; LR = 3e-4

def normalize(seq):
    """seq [T,17,2] xy normalizado a frame -> centrado en cadera + escala torso."""
    s = seq.copy().astype(np.float32)
    hip = (s[:,11,:] + s[:,12,:]) / 2.0          # centro caderas
    s = s - hip[:,None,:]
    sh = (s[:,5,:] + s[:,6,:]) / 2.0             # centro hombros
    torso = np.linalg.norm(sh - 0.0, axis=1)     # dist hombros->cadera(0)
    scale = np.clip(np.median(torso), 1e-3, None)
    s = s / scale
    return s.reshape(s.shape[0], -1)             # [T,34]

class MotionDS(Dataset):
    def __init__(self):
        self.windows = []
        files = sorted(glob.glob(os.path.join(POSES, "*.npy")))
        for f in files:
            seq = np.load(f)
            if seq.shape[0] < WIN + 1: continue
            n = normalize(seq)
            for i in range(0, n.shape[0] - WIN - 1, WIN // 2):
                self.windows.append(n[i:i+WIN+1])
        if self.windows:
            allp = np.concatenate(self.windows, 0)
            self.mean = allp.mean(0, keepdims=True); self.std = allp.std(0, keepdims=True) + 1e-6
            self.windows = [(w - self.mean) / self.std for w in self.windows]
        print(f"Ventanas de entrenamiento: {len(self.windows)} | clips: {len(files)}", flush=True)
    def __len__(self): return len(self.windows)
    def __getitem__(self, i):
        w = torch.from_numpy(self.windows[i].astype(np.float32))
        return w[:-1], w[1:]   # input, target (shift 1)

class MotionDNA(nn.Module):
    def __init__(self):
        super().__init__()
        self.inp = nn.Linear(DIM, D_MODEL)
        self.pos = nn.Parameter(torch.randn(1, WIN, D_MODEL) * 0.02)
        layer = nn.TransformerEncoderLayer(D_MODEL, N_HEAD, D_MODEL*4, DROP, batch_first=True, activation="gelu")
        self.enc = nn.TransformerEncoder(layer, N_LAYER)
        self.out = nn.Linear(D_MODEL, DIM)
    def forward(self, x):
        T = x.size(1)
        h = self.inp(x) + self.pos[:, :T]
        mask = torch.triu(torch.ones(T, T, device=x.device), 1).bool()
        h = self.enc(h, mask=mask)
        return self.out(h)

def main():
    ds = MotionDS()
    if len(ds) == 0:
        print("ERROR: no hay ventanas. Corre extract_poses.py primero."); return
    dl = DataLoader(ds, batch_size=BATCH, shuffle=True, num_workers=4, drop_last=True)
    model = MotionDNA().to(DEV)
    nparam = sum(p.numel() for p in model.parameters())
    print(f"Device: {DEV} | Parametros: {nparam/1e6:.2f}M", flush=True)
    opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, EPOCHS)
    lossf = nn.SmoothL1Loss()
    best = 1e9; hist = []
    for ep in range(1, EPOCHS+1):
        model.train(); tot = 0; nb = 0
        for x, y in dl:
            x, y = x.to(DEV), y.to(DEV)
            pred = model(x)
            loss = lossf(pred, y)
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            tot += loss.item(); nb += 1
        sched.step()
        avg = tot / max(nb,1); hist.append(avg)
        mem = torch.cuda.max_memory_allocated()/1e9 if DEV=="cuda" else 0
        print(f"Epoch {ep:03d}/{EPOCHS}  loss={avg:.5f}  lr={sched.get_last_lr()[0]:.2e}  vram={mem:.2f}GB", flush=True)
        if avg < best:
            best = avg
            torch.save({"model": model.state_dict(), "mean": ds.mean, "std": ds.std,
                        "cfg": {"WIN":WIN,"DIM":DIM,"D_MODEL":D_MODEL,"N_HEAD":N_HEAD,"N_LAYER":N_LAYER},
                        "epoch": ep, "loss": best}, os.path.join(CKPT, "motiondna_v0_best.pt"))
    torch.save({"model": model.state_dict(), "mean": ds.mean, "std": ds.std,
                "cfg": {"WIN":WIN,"DIM":DIM,"D_MODEL":D_MODEL,"N_HEAD":N_HEAD,"N_LAYER":N_LAYER},
                "epoch": EPOCHS, "loss": hist[-1]}, os.path.join(CKPT, "motiondna_v0_final.pt"))
    json.dump(hist, open(os.path.join(CKPT, "loss_history.json"), "w"))
    # --- rollout de muestra: genera 120 frames de movimiento nuevo ---
    model.eval()
    with torch.no_grad():
        seed = ds[0][0][:16].unsqueeze(0).to(DEV)  # 16 frames semilla
        gen = seed.clone()
        for _ in range(120):
            nxt = model(gen[:, -WIN:])[:, -1:]
            gen = torch.cat([gen, nxt], 1)
        gen = gen[0].cpu().numpy() * ds.std + ds.mean   # des-normaliza
    np.save(os.path.join(SAMP, "motiondna_sample.npy"), gen.reshape(-1,17,2))
    print(f"LISTO. best_loss={best:.5f}  ckpt={CKPT}/motiondna_v0_best.pt  sample={SAMP}/motiondna_sample.npy", flush=True)

if __name__ == "__main__":
    main()