#!/usr/bin/env python3
"""Generate T2M-GPT poses for named clips and save them in NSLP-G's 50-joint format.

Purpose: a visual comparison of the two families that is actually about quality rather
than about rendering. T2M-GPT emits 128 DWPose keypoints (including 68 face points),
NSLP-G emits 50 (8 body + 21 + 21 hand). Rendered with their own renderers the two look
different for reasons that have nothing to do with how good the signing is, so this maps
T2M-GPT's output onto exactly NSLP-G's 50 joints and writes the `poses`/`names` npz that
`0.NSLP-G/Word-level/NSLP-G/render_vsl.py` reads. Both families are then drawn by the
same code, in the same coordinate space, for the same clips.

Coordinate space: the VSL3 pack (T2M-GPT) and the VSL3_upper pack (NSLP-G) are both
per-axis frame-normalized [0,1] over the same clips, so no rescaling is needed -- only a
keypoint-index gather:
    original DWPose index = upper_layout.keep[j] for j in NSLP-G's KEEP_50
"""
import argparse
import json
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

# NSLP-G's modules/data/mvsl.py: 8 body + lhand + rhand inside the 124-kpt upper layout
KEEP_50 = list(range(8)) + list(range(82, 124))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/VSL3')
    ap.add_argument('--upper-dir', default='./dataset/VSL3_upper',
                    help='only its layout.json is read, for the index mapping')
    ap.add_argument('--vq', default='output_vsl/vq_vsl_3view/net_best.pth')
    ap.add_argument('--gpt', default='output_vsl/gpt_vsl_3view/net_best.pth')
    ap.add_argument('--split', default='test')
    ap.add_argument('--names', nargs='+', required=True, help='clip names to generate')
    ap.add_argument('--sampling', default='categorial', choices=['greedy', 'categorial'])
    ap.add_argument('--seed', type=int, default=0)
    ap.add_argument('--device', default='cuda')
    ap.add_argument('--out', required=True, help='output .npz')
    args = ap.parse_args()

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

    net, targs, _ = 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
    with open(os.path.join(args.upper_dir, 'layout.json')) as f:
        upper_keep = json.load(f)['keep']
    take = [upper_keep[j] for j in KEEP_50]          # -> original DWPose indices
    assert len(take) == 50

    by_name = {c['name']: k for k, c in enumerate(st.index)}
    poses, names, texts = [], [], []
    with torch.no_grad():
        for nm in args.names:
            if nm not in by_name:
                print(f'  [skip] {nm} not in {args.split}')
                continue
            c = st.index[by_name[nm]]
            txt = c[g.text_field]
            idx = gpt.sample(text_enc([txt]), if_categorial=(args.sampling == 'categorial'))
            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)[:, take]     # [T,50,2]
            poses.append(p.astype(np.float32))
            names.append(nm)
            texts.append(txt)
            print(f'  {nm}: T={len(p)}  gloss="{txt}"')

    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, frame-normalized coords')


if __name__ == '__main__':
    main()