#!/usr/bin/env python3
"""The TriVis SENTENCE-level T2M-GPT, scored in shoulder widths against real floors.

`eval_vsl.py` measures this model in TriVis frame-widths and against its tokenizer
ceiling only. Two consequences worth fixing before calling its 0.1537 hands "good":

  * frame-widths divide x by W=1176 and y by H=1288, so the unit is anisotropic and not
    comparable to any Multi-VSL/NSLP-G number. Here both prediction and reference are
    converted to isotropic shoulder widths with the reference clip's own median anchor.
  * there was no floor. `check_trivis_floors.py` supplies them; this script adds the
    model rows in the same unit on the same 300 clips so one table can be read.

Anchor choice matters and both are reported, because the word-level pipelines are
neck-centred by construction and therefore get global placement for free:

  gt-anchor    prediction is placed with the REFERENCE clip's neck/shoulder-width, so the
               model is charged for getting global position and body scale wrong.
  own-anchor   prediction is normalized by ITS OWN median neck/shoulder width, which
               removes global placement from the score -- the apples-to-apples row
               against the composed word models.
"""
import argparse
import csv
import json
import os

import numpy as np
import torch
from tqdm import tqdm

import eval_vsl
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

NECK, RSHO, LSHO = 1, 2, 5
# Same rule as `prepare_multivsl_data.normalize_clip`: below this the shoulder-width scale
# reference is untrustworthy. One `right`-view clip of the 3-view pack has both shoulders
# on the same pixel (sw = 0), which would silently divide by zero.
MIN_SHOULDER_PX = 8.0


def anchor_of(xy, valid, W, H):
    px = xy * np.array([W, H], np.float32)
    ok = valid[:, NECK] & valid[:, RSHO] & valid[:, LSHO]
    if ok.sum() < 3:
        ok = np.ones(len(xy), bool)
    neck = np.median(px[ok, NECK, :], axis=0)
    sw = float(np.median(np.linalg.norm(px[ok, RSHO, :] - px[ok, LSHO, :], axis=-1)))
    return neck, sw


def to_shoulder(xy, neck, sw, W, H):
    return (xy * np.array([W, H], np.float32) - neck[None, None, :].astype(np.float32)) / sw


class Agg:
    def __init__(self, groups):
        self.vals = {g: [] for g in groups}

    def add(self, d):
        for g, v in d.items():
            self.vals[g].append(v)

    def summary(self):
        out = {}
        for g, v in self.vals.items():
            a = np.asarray(v, np.float64)
            out[g] = {"mean": float(a.mean()), "n": len(a),
                      "sem": float(a.std(ddof=1) / np.sqrt(len(a))) if len(a) > 1 else 0.0}
        return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/VSL')
    ap.add_argument('--vq', default='output_vsl/vq_vsl_front_lab/net_best.pth')
    ap.add_argument('--gpt', default='output_vsl/gpt_vsl_front_lab_v2/net_best.pth')
    ap.add_argument('--gloss-json', default='output_vsl/text2gloss/pred_full_trivis.json')
    ap.add_argument('--split', default='test')
    ap.add_argument('--n', type=int, default=300)
    ap.add_argument('--seed', type=int, default=0)
    # --seed fixes the 300-clip subset AND the shuffle donors, so every table stays
    # paired; it must not change between runs. --sample-seed reseeds only categorial
    # generation, which is what makes repeat runs differ (~+-0.01 on hands DTW).
    # Defaults to --seed, i.e. exactly the old single-seed behaviour.
    ap.add_argument('--sample-seed', type=int, default=None)
    ap.add_argument('--sampling', default='categorial', choices=['greedy', 'categorial'])
    ap.add_argument('--conditions', default='oracle-gloss,pred-gloss,shuffled-rotate,'
                                           'shuffled-half,shuffled-random')
    ap.add_argument('--no-ceiling', action='store_true')
    ap.add_argument('--by-view', action='store_true',
                    help='break every row down by camera view (3-view packs)')
    ap.add_argument('--view-csv', default='../Full_TriVis/split_lab_3view.csv',
                    help='authoritative clip -> view map; file names are NOT reliable '
                         "(Group4's front clips are named _center_)")
    ap.add_argument('--frame-w', type=float, default=1176.0)
    ap.add_argument('--frame-h', type=float, default=1288.0)
    ap.add_argument('--device', default='cuda')
    ap.add_argument('--out-json', default='output_vsl/gpt_vsl_front_lab_v2/on_trivis_sw.json')
    args = ap.parse_args()

    W, H = args.frame_w, args.frame_h
    device = torch.device(args.device)
    torch.manual_seed(args.seed if args.sample_seed is None else args.sample_seed)

    net, targs, ck1 = 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
    groups = st.layout.metric_groups()
    eval_vsl.GROUPS = groups

    with open(args.gloss_json, encoding='utf-8') as f:
        glosses = json.load(f)

    rs = np.random.RandomState(args.seed)
    items = sorted(rs.choice(len(st.index), size=min(args.n, len(st.index)),
                             replace=False).tolist())

    # view comes from the CSV, never from the file name: Group4's front-view clips are
    # named `_center_`, so a name-based guess mislabels 5,852 front clips (8% of the pack)
    view_of = {}
    if args.by_view and os.path.exists(args.view_csv):
        with open(args.view_csv, newline='', encoding='utf-8') as f:
            for r in csv.DictReader(f):
                view_of[os.path.splitext(os.path.basename(r['npz_path']))[0]] = r['view']
        print(f'view map: {len(view_of)} clips from {args.view_csv}')

    refs, keep, dropped = [], [], 0
    for i in items:
        motion, mask = st.get(i)
        vd = mask[:, ::2].astype(bool)
        gt = (motion * st.std + st.mean).reshape(-1, NK, 2)
        neck, sw = anchor_of(gt, vd, W, H)
        if not np.isfinite(sw) or sw < MIN_SHOULDER_PX:
            dropped += 1
            continue
        refs.append((gt, vd, neck, sw, to_shoulder(gt, neck, sw, W, H), motion))
        keep.append(i)
    items = keep
    if dropped:
        print(f'[drop] {dropped} clip(s) with median shoulder width < {MIN_SHOULDER_PX} px')
    print(f'{len(items)} clips of {args.split}, {NK} keypoints, text_field={g.text_field}')

    # conditioning texts. The model was trained on the `gloss` field, which is the
    # space-joined Sign_sentence; the predicted-gloss row substitutes the BARTpho output.
    oracle = [st.index[i][g.text_field] for i in items]
    pred = [glosses[st.index[i]['name']]['pred_gloss'] for i in items]
    # Three ways to mismatch the text, because the choice changes the answer. Rotating a
    # name-sorted clip list pairs neighbours -- same signer, same session, 25% same
    # category (token Jaccard 0.031) -- so it is the LEAKIEST control. eval_vsl.py's
    # i+N/2 jump and a uniform random draw over the whole split both land at 0.021-0.022.
    # On a 3-view pack the CSV lists front/left/right of the SAME sentence consecutively,
    # so a donor drawn by index can carry the identical gloss. Every rule below is
    # therefore filtered: if the donor text equals the reference's, redraw.
    perm = list(range(1, len(items))) + [0]
    N = len(st.index)
    rs3 = np.random.RandomState(args.seed + 7)

    def clean(donors, ref):
        out, fixed = [], 0
        for d, t in zip(donors, ref):
            if d == t:
                for _ in range(50):
                    d = st.index[int(rs3.randint(N))][g.text_field]
                    if d != t:
                        break
                fixed += 1
            out.append(d)
        if fixed:
            print(f'  redrew {fixed} donors that carried the reference gloss')
        return out

    sources = {
        'oracle-gloss': oracle,
        'pred-gloss': pred,
        'shuffled-rotate': clean([oracle[p] for p in perm], oracle),
        'shuffled-half': clean([st.index[(i + N // 2) % N][g.text_field] for i in items],
                               oracle),
        'shuffled-random': clean([st.index[j][g.text_field]
                                  for j in rs3.randint(0, N, len(items))], oracle),
    }
    keep = args.conditions.split(',')
    sources = {k: v for k, v in sources.items() if k in keep}

    results = {}

    def score(tag, preds):
        # The SAME predictions scored in both units, so any difference in the
        # real-vs-shuffled gap is attributable to the unit alone -- not to a different
        # run, seed, shuffle rule or code path.
        a_gt, a_own, a_frame = Agg(groups), Agg(groups), Agg(groups)
        per_view = {}
        ratios, n = [], 0
        for k, ((gt, vd, neck, sw, gt_sw, _), p) in enumerate(zip(refs, preds)):
            if p is None:
                continue
            d_gt = eval_vsl.dtw_mje(to_shoulder(p, neck, sw, W, H).astype(np.float64),
                                    gt_sw.astype(np.float64), vd)
            a_gt.add(d_gt)
            pn, psw = anchor_of(p, np.ones((len(p), NK), bool), W, H)
            a_own.add(eval_vsl.dtw_mje(to_shoulder(p, pn, psw, W, H).astype(np.float64),
                                       gt_sw.astype(np.float64), vd))
            a_frame.add(eval_vsl.dtw_mje(p.astype(np.float64), gt.astype(np.float64), vd))
            if args.by_view:
                nm = st.index[items[k]]['name']
                v = view_of.get(nm, 'unknown')
                per_view.setdefault(v, Agg(groups)).add(d_gt)
            ratios.append(len(p) / len(gt))
            n += 1
        sg, so, sf = a_gt.summary(), a_own.summary(), a_frame.summary()
        print(f"{tag:<20} gt-anchor: " + " ".join(f"{q} {sg[q]['mean']:.4f}"
                                                 for q in ('all', 'body', 'hands'))
              + " | own-anchor: " + " ".join(f"{q} {so[q]['mean']:.4f}"
                                            for q in ('all', 'body', 'hands'))
              + " | frame: " + " ".join(f"{q} {sf[q]['mean']:.4f}"
                                       for q in ('all', 'body', 'hands'))
              + f"  len_ratio {np.mean(ratios):.3f} (n={n})", flush=True)
        results[tag] = {'gt_anchor': sg, 'own_anchor': so, 'frame': sf,
                        'len_ratio': float(np.mean(ratios)), 'n': n}
        if args.by_view:
            results[tag]['by_view'] = {v: a.summary() for v, a in sorted(per_view.items())}
            for v, a in sorted(per_view.items()):
                sv = a.summary()
                print(f"    {v:<18} n={sv['hands']['n']:<5}" + " ".join(
                    f"{q} {sv[q]['mean']:.4f}" for q in ('all', 'body', 'hands')), flush=True)

    unit = 2 ** targs.down_t
    empties = 0
    with torch.no_grad():
        # tokenizer ceiling
        if not args.no_ceiling:
            preds = []
            for (gt, vd, neck, sw, gt_sw, motion) in tqdm(refs, desc='ceiling', leave=False):
                T = max(unit, (len(motion) // unit) * unit)
                x = torch.from_numpy(motion[:T]).unsqueeze(0).to(device)
                rec = net.decode_batch(net.encode(x))[0].cpu().numpy() * st.std + st.mean
                preds.append(rec.reshape(-1, NK, 2))
            score('ceiling', preds)

        for tag, texts in sources.items():
            preds = []
            for t in tqdm(texts, desc=tag, leave=False):
                idx = gpt.sample(text_enc([t]), if_categorial=(args.sampling == 'categorial'))
                if idx is None or idx.numel() == 0:
                    empties += 1
                    preds.append(None)
                    continue
                p = net.decode_batch(idx.clamp(max=g.nb_code - 1))[0].cpu().numpy()
                preds.append((p * st.std + st.mean).reshape(-1, NK, 2))
            score(tag, preds)

    print(f'\nunit: shoulder widths, isotropic, per-clip.  {empties} empty samples')
    with open(args.out_json, 'w') as f:
        json.dump({'args': vars(args), 'clips': [st.index[i]['name'] for i in items],
                   'n_empty': empties, 'results': results}, f, indent=2)
    print(f'wrote {args.out_json}')


if __name__ == '__main__':
    main()