#!/usr/bin/env python3
"""Dump ground-truth poses in the shared 50-joint format, for FGD/MAEJ scoring.

`eval_fgd_maej.py` needs a `--gt-npz` in the same format as the model dumps. It exists
for Multi-VSL (`dumps_fgd_t2m/gt.npz`) but not for the 3-view TriVis pack, and FGD is
biased by sample size, so the GT must cover exactly the clips the conditions cover --
hence `--like`, which copies the clip list from an existing dump.

KEEP_50 = upper[0:8] + upper[82:124], matching 0.NSLP-G/.../modules/data/mvsl.py.
"""
import argparse
import os

import numpy as np

from dataset import dataset_vsl

KEEP_50 = list(range(8)) + list(range(82, 124))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/VSL3_upper')
    ap.add_argument('--split', default='test')
    ap.add_argument('--like', default=None,
                    help='npz whose `names` define the clip set (keeps FGD comparable)')
    ap.add_argument('--out', required=True)
    args = ap.parse_args()

    st = dataset_vsl.VSLStore(args.data_dir, args.split)
    NK = st.layout.n_kpts
    assert NK == 124, f'expects the 124-kpt upper layout, got {NK}'

    want = None
    if args.like:
        z = np.load(args.like, allow_pickle=True)
        want = [str(n) for n in z['names']]
        print(f'{len(want)} clip names taken from {args.like}')

    by_name = {c['name']: k for k, c in enumerate(st.index)}
    names = want if want is not None else [c['name'] for c in st.index]

    poses, kept = [], []
    for nm in names:
        if nm not in by_name:
            continue
        motion, _ = st.get(by_name[nm])
        p = (motion * st.std + st.mean).reshape(-1, NK, 2)[:, KEEP_50]
        poses.append(p.astype(np.float32))
        kept.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(kept))
    print(f'wrote {args.out}: {len(kept)} clips, 50 joints')


if __name__ == '__main__':
    main()