#!/usr/bin/env python3
"""Dump generated + GT poses for the hand-weight sweep in the shared 50-joint space,
so `eval_fgd_maej.py` can score FGD/MAEJ on exactly the DTW table's 300 clips.

Joint mapping. NSLP-G's KEEP_50 is expressed in the 124-kpt UPPER layout
(`range(8) + range(82,124)`), but the sweep models are trained on the 128-kpt FULL
layout. The upper layout only drops body 9/10/12/13, so the same physical joints are
`range(8) + range(86,128)` in full indices -- body 0-7, then both 21-joint hands.
Verified equal: upper body 0-7 == full body 0-7 (all drops are >= 9), and the hand
blocks shift by exactly the 4 dropped body joints (82->86, 103->107).

MIND THE AE. `eval_fgd_maej.py` applies KEEP_50 directly to its `--data-dir` store, so
it must be pointed at `dataset/VSL_upper` (124 kpt), NOT `dataset/VSL`. On a 128-kpt
store those indices silently select face-tail + the wrong hand joints and the FGD comes
out plausible but wrong. Same clips either way, so the AE is unaffected.

Clip subset is `RandomState(--seed).choice(...)` sorted -- byte-identical to
`eval_trivis_sent_sw.py`, so FGD and DTW describe the same 300 clips.
"""
import argparse
import os

import numpy as np
import torch

import models.t2m_trans as trans
from dataset import dataset_vsl
from models.text_encoder_vi import ViTextEncoder
from train_t2m_trans_vsl import build_vqvae

KEEP_50_FULL = list(range(8)) + list(range(86, 128))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/VSL')
    ap.add_argument('--vq', required=True)
    ap.add_argument('--gpt', required=True)
    ap.add_argument('--split', default='test')
    ap.add_argument('--n', type=int, default=300)
    ap.add_argument('--seed', type=int, default=0, help='clip-subset seed; keep at 0')
    ap.add_argument('--sample-seed', type=int, default=0, help='generation seed only')
    ap.add_argument('--out', required=True)
    ap.add_argument('--gt-out', default=None, help='also write the GT dump here')
    ap.add_argument('--device', default='cuda')
    args = ap.parse_args()

    device = torch.device(args.device)
    torch.manual_seed(args.sample_seed)

    net, _, _ = build_vqvae(args.vq, device)
    tck = torch.load(args.gpt, map_location='cpu')
    g = argparse.Namespace(**tck['args'])
    text_enc = ViTextEncoder(g.text_model, device=args.device)
    gpt = trans.Text2Motion_Transformer(
        num_vq=g.nb_code, embed_dim=g.embed_dim_gpt, clip_dim=text_enc.dim,
        block_size=g.max_tokens + 1, num_layers=g.num_layers, n_head=g.n_head_gpt,
        drop_out_rate=g.drop_out_rate, fc_rate=g.ff_rate)
    gpt.load_state_dict(tck['trans'], strict=True)
    gpt.eval().to(device)

    st = dataset_vsl.VSLStore(args.data_dir, args.split)
    NK = st.layout.n_kpts
    if NK != 128:
        raise SystemExit(f'expected the 128-kpt full layout, got {NK} -- '
                         f'KEEP_50_FULL would select the wrong joints')
    rs = np.random.RandomState(args.seed)
    items = sorted(rs.choice(len(st.index), size=min(args.n, len(st.index)),
                             replace=False).tolist())

    poses, names, texts = [], [], []
    gt_poses, gt_names = [], []
    with torch.no_grad():
        for k in items:
            c = st.index[k]
            nm, txt = c['name'], c[g.text_field]
            idx = gpt.sample(text_enc([txt]), if_categorial=True)
            if idx is None or idx.numel() == 0:
                print(f'  [skip] {nm}: empty sample')
                continue
            p = net.decode_batch(idx.clamp(max=g.nb_code - 1))[0].cpu().numpy()
            p = (p * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50_FULL]
            poses.append(p.astype(np.float32)); names.append(nm); texts.append(txt)
            if args.gt_out:
                m, _ = st.get(k)
                gq = (m * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50_FULL]
                gt_poses.append(gq.astype(np.float32)); gt_names.append(nm)

    os.makedirs(os.path.dirname(args.out) or '.', exist_ok=True)
    np.savez(args.out, poses=np.array(poses, dtype=object),
             names=np.array(names), texts=np.array(texts))
    print(f'wrote {args.out}: {len(poses)} clips, 50 joints')
    if args.gt_out:
        np.savez(args.gt_out, poses=np.array(gt_poses, dtype=object),
                 names=np.array(gt_names), texts=np.array(gt_names))
        print(f'wrote {args.gt_out}: {len(gt_poses)} GT clips')


if __name__ == '__main__':
    main()