#!/usr/bin/env python3
"""Render qualitative gloss->pose examples for a VSL T2M-GPT model.

Three panels per example, which is the point of this script:

  GT                 the real skeleton
  CEILING            GT encoded to symbols and decoded back -- what stage 1 alone
                     can do, i.e. the best any generation could look
  GENERATED          from gloss text only

Comparing CEILING against GENERATED separates "the tokenizer lost detail" from
"the translator picked the wrong symbols".

Layout-aware: the skeleton topology is defined on the original DWPose-128 indices
and remapped through the dataset's layout.json, so a trimmed keypoint set (e.g.
`upper`, which drops knees/ankles and shifts every hand index) draws correctly.
Keypoints the model is never supervised on are suppressed in every panel --
otherwise the decoder's unconstrained outputs there render as a phantom limb.

Coordinate space is taken from the data: Full_TriVis is frame-normalized, while
Multi-VSL is per-clip shoulder-width normalized. The viewport is auto-fitted from
the clip, so both render correctly without per-dataset tuning.
"""
import argparse
import os

import numpy as np
import torch
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter, FuncAnimation

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

# Topology on ORIGINAL DWPose-128 indices (body 0:18, face 18:86, LH 86:107, RH 107:128)
BODY_EDGES = [(1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6, 7), (1, 8), (8, 9), (9, 10),
              (1, 11), (11, 12), (12, 13), (1, 0), (0, 14), (14, 16), (0, 15), (15, 17)]
HAND_EDGES = [(0, 1), (1, 2), (2, 3), (3, 4), (0, 5), (5, 6), (6, 7), (7, 8), (0, 9),
              (9, 10), (10, 11), (11, 12), (0, 13), (13, 14), (14, 15), (15, 16),
              (0, 17), (17, 18), (18, 19), (19, 20), (5, 9), (9, 13), (13, 17)]
# Knees/ankles are detected in only ~19% / ~0.1% of frames -- these clips crop below
# the hips, the masked loss never supervises them, so the decoder emits arbitrary
# values. Suppressed even when a layout keeps them.
UNSUPERVISED_ORIG = (9, 10, 12, 13)


class Topology:
    """Skeleton edges + group slices expressed in a layout's index space."""

    def __init__(self, layout):
        self.layout = layout
        o2n = layout.old2new
        drop = set(UNSUPERVISED_ORIG)
        self.body = [(o2n[a], o2n[b]) for a, b in BODY_EDGES
                     if a in o2n and b in o2n and a not in drop and b not in drop]
        g = layout.groups
        self.lh, self.rh = g['lhand'][0], g['rhand'][0]
        self.nlh = g['lhand'][1] - g['lhand'][0]
        self.face = g['face']
        self.hand = [(a, b) for a, b in HAND_EDGES if a < self.nlh and b < self.nlh]


def draw(ax, xy, topo, color, hand_color, scale, ctr, lw=1.6, valid=None):
    """xy: [n_kpts,2] in the dataset's own units. Drawn in display pixels."""
    ax.clear()
    ok = (lambda i: True) if valid is None else (lambda i: bool(valid[i]))
    px = (xy[:, 0] - ctr[0]) * scale
    py = (xy[:, 1] - ctr[1]) * scale

    def seg(edges, off, c, width):
        for a, b in edges:
            if not (ok(off + a) and ok(off + b)):
                continue
            ax.plot([px[off + a], px[off + b]], [py[off + a], py[off + b]],
                    color=c, lw=width, solid_capstyle='round')

    seg(topo.body, 0, color, lw)
    seg(topo.hand, topo.lh, hand_color, lw * 1.15)
    seg(topo.hand, topo.rh, hand_color, lw * 1.15)
    a, b = topo.face
    fx, fy = px[a:b], py[a:b]
    if valid is not None:
        m = valid[a:b].astype(bool)
        fx, fy = fx[m], fy[m]
    ax.scatter(fx, fy, s=1.2, color=color, alpha=0.45, linewidths=0)
    ax.set_aspect('equal')
    ax.axis('off')


def group_err(pred, gt, valid, groups, name='hands'):
    """Mean error over a keypoint group on the overlapping prefix, native units."""
    T = min(len(pred), len(gt))
    if T == 0:
        return float('nan')
    a, b = groups[name]
    d = np.linalg.norm(pred[:T, a:b] - gt[:T, a:b], axis=-1)
    v = valid[:T, a:b]
    return float((d * v).sum() / max(v.sum(), 1e-6))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--data-dir', default='./dataset/VSL')
    ap.add_argument('--resume-pth', default='output_vsl/vq_vsl_front_lab/net_best.pth')
    ap.add_argument('--resume-trans', default='output_vsl/gpt_vsl_front_lab_v2/net_best.pth')
    ap.add_argument('--split', default='test')
    ap.add_argument('--out-dir', default='qual_vsl')
    ap.add_argument('--n', type=int, default=6)
    ap.add_argument('--fps', type=int, default=30)
    ap.add_argument('--strip-frames', type=int, default=6)
    ap.add_argument('--display-px', type=float, default=0,
                    help='display pixels per data unit (0 = auto from the clip)')
    ap.add_argument('--unit-name', default='', help='label for the error unit (auto if blank)')
    ap.add_argument('--min-gloss', type=int, default=1)
    ap.add_argument('--max-gloss', type=int, default=99)
    ap.add_argument('--seed', type=int, default=1)
    ap.add_argument('--sampling', default='categorial', choices=['categorial', 'greedy'])
    ap.add_argument('--device', default='cuda')
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)
    device = torch.device(args.device)

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

    store = dataset_vsl.VSLStore(args.data_dir, args.split)
    layout = store.layout
    topo = Topology(layout)
    groups = layout.metric_groups()
    NK = layout.n_kpts
    unit = args.unit_name or ('shoulder-w' if 'MVSL' in args.data_dir else 'frame-w')
    print(f'stage-2 iter {tck.get("iter")} | {layout} | sampling={args.sampling} | unit={unit}')

    cand = [i for i, c in enumerate(store.index)
            if args.min_gloss <= len(str(c['gloss']).split()) <= args.max_gloss]
    rng = np.random.RandomState(args.seed)
    picks = rng.choice(cand, size=min(args.n, len(cand)), replace=False).tolist()

    unit_len = 2 ** targs.down_t
    for k, i in enumerate(picks):
        c = store.index[i]
        motion, mask = store.get(i)
        valid = mask[:, ::2]
        gt = (motion * store.std + store.mean).reshape(-1, NK, 2)

        with torch.no_grad():
            gt_t = torch.from_numpy(motion).unsqueeze(0).to(device)
            T = (len(motion) // unit_len) * unit_len
            rec = net.decode_batch(net.encode(gt_t[:, :T]))[0].cpu().numpy()
            rec = (rec * store.std + store.mean).reshape(-1, NK, 2)
            feat = text_enc([c[gargs.text_field]])
            idx = gpt.sample(feat, if_categorial=(args.sampling == 'categorial'))
            if idx is None or idx.numel() == 0:
                print(f'[{k}] empty generation, skipping')
                continue
            idx = idx.clamp(max=gargs.nb_code - 1)
            gen = net.decode_batch(idx)[0].cpu().numpy()
            gen = (gen * store.std + store.mean).reshape(-1, NK, 2)

        e_rec = group_err(rec, gt, valid, groups)
        e_gen = group_err(gen, gt, valid, groups)

        # auto-fit a shared viewport from the ground truth (works in either unit space)
        lo, hi = gt.reshape(-1, 2).min(0), gt.reshape(-1, 2).max(0)
        span = float(max(hi - lo)) or 1.0
        ctr = (lo + hi) / 2.0
        scale = args.display_px or (520.0 / span)
        half = span * scale * 0.62

        title = (f"{c['gloss']}   [{unit}]\n"
                 f"GT {len(gt)}f  |  ceiling {e_rec:.3f}  |  generated {e_gen:.3f}, "
                 f"{len(gen)}f ({len(gen)/len(gt):.2f}x)")
        print(f"[{k}] {c['name'][:52]}  gloss='{c['gloss']}'  "
              f"ceiling={e_rec:.3f} gen={e_gen:.3f} len={len(gen)}/{len(gt)}")

        panels = [('GROUND TRUTH', gt, valid, '#111111', '#c0392b'),
                  ('CEILING (tokenizer only)', rec, None, '#1f6f3f', '#27ae60'),
                  ('GENERATED (from gloss)', gen, None, '#1a4f8a', '#2980b9')]
        nT = max(len(p[1]) for p in panels)

        fig, axes = plt.subplots(1, 3, figsize=(11, 4.6))
        fig.suptitle(title, fontsize=9)

        def frame(t):
            for ax, (label, seq, vd, col, hcol) in zip(axes, panels):
                tt = min(t, len(seq) - 1)
                draw(ax, seq[tt], topo, col, hcol, scale, ctr,
                     valid=(vd[tt] if vd is not None else None))
                ax.set_xlim(-half, half); ax.set_ylim(half, -half)
                ax.set_title(f"{label}\nframe {tt+1}/{len(seq)}", fontsize=8)
            return []

        anim = FuncAnimation(fig, frame, frames=nT, interval=1000 / args.fps, blit=False)
        anim.save(os.path.join(args.out_dir, f'{k:02d}_{c["name"][:48]}.mp4'),
                  writer=FFMpegWriter(fps=args.fps, bitrate=2400))
        plt.close(fig)

        nf = args.strip_frames
        fig, ax2 = plt.subplots(3, nf, figsize=(1.7 * nf, 5.6))
        fig.suptitle(title, fontsize=9)
        for r, (label, seq, vd, col, hcol) in enumerate(panels):
            ts = np.linspace(0, len(seq) - 1, nf).astype(int)
            for ci, tt in enumerate(ts):
                a = ax2[r, ci]
                draw(a, seq[tt], topo, col, hcol, scale, ctr, lw=1.1,
                     valid=(vd[tt] if vd is not None else None))
                a.set_xlim(-half, half); a.set_ylim(half, -half)
                if ci == 0:
                    a.set_ylabel(label, fontsize=7)
                a.set_title(f'{tt+1}', fontsize=6)
        plt.tight_layout(rect=[0, 0, 1, 0.93])
        fig.savefig(os.path.join(args.out_dir, f'{k:02d}_{c["name"][:48]}.png'), dpi=130)
        plt.close(fig)

    print(f'\nwrote {args.out_dir}/  ({len(picks)} examples: mp4 + png each)')


if __name__ == '__main__':
    main()