#!/usr/bin/env python3
"""Dump gpt_mvsl generations in NSLP-G's dump format, so both models can be scored
with one FGD/MAEJ/DTW protocol.

Two alignment details, both required for the comparison to be fair:

  * joints  -- gpt_mvsl emits the 124-kpt `upper` layout; NSLP-G emits 50 joints
               (8 body + 42 hands, no face). We slice to the same 50 here.
               KEEP_50 = upper[0:8] + upper[82:124], matching
               0.NSLP-G/.../modules/data/mvsl.py.
  * clips   -- NSLP-G selects with np.random.default_rng(0); eval_vsl.py uses
               np.random.RandomState(0). Those give DIFFERENT subsets, so nothing is
               assumed: every clip is dumped with its name and the scorer intersects
               on names.

Output matches their `score(..., dump_dir)`: poses (object array of [T,50,2]),
labels, names.
"""
import argparse
import os

import numpy as np
import torch
from tqdm import tqdm

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 = list(range(8)) + list(range(82, 124))     # upper-layout indices


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/MVSL')
    ap.add_argument('--vq', default='output_vsl/vq_mvsl/net_best.pth')
    ap.add_argument('--gpt', default='output_vsl/gpt_mvsl/net_best.pth')
    ap.add_argument('--split', default='test')
    ap.add_argument('--n', type=int, default=0, help='0 = all clips in the split')
    ap.add_argument('--dump-dir', default='dumps_fgd_t2m')
    ap.add_argument('--sampling', default='categorial', choices=['categorial', 'greedy'])
    ap.add_argument('--device', default='cuda')
    args = ap.parse_args()

    device = torch.device(args.device)
    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
    ids = list(range(len(st.index)))
    if args.n:
        ids = ids[:args.n]
    unit = 2 ** targs.down_t

    gen_p, ceil_p, gt_p, labels, names = [], [], [], [], []
    with torch.no_grad():
        for i in tqdm(ids, desc='dump gpt_mvsl'):
            c = st.index[i]
            motion, _ = st.get(i)
            gt = (motion * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50]

            mt = torch.from_numpy(motion).unsqueeze(0).to(device)
            T = (len(motion) // unit) * unit
            rec = net.decode_batch(net.encode(mt[:, :T]))[0].cpu().numpy()
            rec = (rec * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50]

            feat = text_enc([c[g.text_field]])
            idx = gpt.sample(feat, if_categorial=(args.sampling == 'categorial'))
            if idx is None or idx.numel() == 0:
                continue
            idx = idx.clamp(max=g.nb_code - 1)
            gen = net.decode_batch(idx)[0].cpu().numpy()
            gen = (gen * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50]

            gen_p.append(gen.astype(np.float32))
            ceil_p.append(rec.astype(np.float32))
            gt_p.append(gt.astype(np.float32))
            labels.append(int(c.get('label', -1)))
            names.append(c['name'])

    os.makedirs(args.dump_dir, exist_ok=True)
    for tag, arr in (('t2mgpt_gen', gen_p), ('t2mgpt_ceiling', ceil_p), ('gt', gt_p)):
        np.savez(os.path.join(args.dump_dir, f'{tag}.npz'),
                 poses=np.array(arr, dtype=object),
                 labels=np.array(labels), names=np.array(names))
        print(f'  {tag}: {len(arr)} clips -> {args.dump_dir}/{tag}.npz')


if __name__ == '__main__':
    main()