#!/usr/bin/env python3
"""Render one TriVis sentence across every system, straight from the scored dumps.

Reads the same .npz dumps `eval_trivis_metrics.py` scores, so the picture is
guaranteed to show exactly what the numbers describe -- no second generation pass,
no risk of a render that disagrees with the table.

All dumps are 50 joints in TriVis frame coordinates: body 0:8 (OpenPose 0-7),
left hand 8:29, right hand 29:50. Drawn in pixels so proportions are correct.
"""
import argparse
import glob
import os

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

# OpenPose 0-7: 0 nose, 1 neck, 2 Rsho, 3 Relb, 4 Rwri, 5 Lsho, 6 Lelb, 7 Lwri
BODY_EDGES = [(1, 0), (1, 2), (2, 3), (3, 4), (1, 5), (5, 6), (6, 7)]
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)]
LH, RH = 8, 29
WRIST_L, WRIST_R = 7, 4          # body wrists, to attach each hand root


def draw(ax, xy, color, hand_color, W, H, lo, hi, lw=1.6):
    ax.clear()
    px, py = xy[:, 0] * W, xy[:, 1] * H
    for a, b in BODY_EDGES:
        ax.plot([px[a], px[b]], [py[a], py[b]], color=color, lw=lw, solid_capstyle='round')
    for off, wr in ((LH, WRIST_L), (RH, WRIST_R)):
        ax.plot([px[wr], px[off]], [py[wr], py[off]], color=color, lw=lw * 0.8)
        for a, b in HAND_EDGES:
            ax.plot([px[off + a], px[off + b]], [py[off + a], py[off + b]],
                    color=hand_color, lw=lw * 1.1, solid_capstyle='round')
    ax.set_xlim(lo[0], hi[0]); ax.set_ylim(hi[1], lo[1])
    ax.set_aspect('equal'); ax.axis('off')


def load(p):
    d = np.load(p, allow_pickle=True)
    return {str(n): q for n, q in zip(d['names'], d['poses'])}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--panels', nargs='+', required=True,
                    help='label=path/to.npz, in display order')
    ap.add_argument('--clip', default=None, help='clip name (default: first shared)')
    ap.add_argument('--frames', type=int, default=6)
    ap.add_argument('--fps', type=int, default=30)
    ap.add_argument('--frame-w', type=float, default=1176.0)
    ap.add_argument('--frame-h', type=float, default=1288.0)
    ap.add_argument('--gloss', default=None, help='caption text')
    ap.add_argument('--out', default='qual_trivis/strip')
    ap.add_argument('--mp4', action='store_true')
    ap.add_argument('--gif', action='store_true',
                    help='also write a gif (ffmpeg palettegen: far smaller and cleaner '
                         'than a direct matplotlib gif)')
    ap.add_argument('--gif-fps', type=int, default=15)
    ap.add_argument('--gif-width', type=int, default=900)
    args = ap.parse_args()

    panels = []
    for spec in args.panels:
        lab, path = spec.split('=', 1)
        panels.append((lab, load(path)))
    shared = set(panels[0][1])
    for _, m in panels[1:]:
        shared &= set(m)
    if not shared:
        raise SystemExit('no clip shared by all panels')
    clip = args.clip or sorted(shared)[0]
    if clip not in shared:
        raise SystemExit(f'{clip} not in all panels')
    print(f'clip {clip}  ({len(shared)} shared)')

    seqs = [(lab, np.asarray(m[clip], np.float32)) for lab, m in panels]
    W, H = args.frame_w, args.frame_h
    ref = seqs[0][1]
    lo = (ref.reshape(-1, 2).min(0) * np.array([W, H])) - 60
    hi = (ref.reshape(-1, 2).max(0) * np.array([W, H])) + 60

    cols = [('#111111', '#c0392b'), ('#1f6f3f', '#27ae60'), ('#1a4f8a', '#2980b9'),
            ('#8e44ad', '#c39bd3'), ('#b9770e', '#e59866'), ('#117a65', '#45b39d')]
    os.makedirs(os.path.dirname(args.out) or '.', exist_ok=True)
    title = f"{args.gloss or clip}"

    # ---- keyframe grid: one row per system ----
    nf = args.frames
    fig, axes = plt.subplots(len(seqs), nf, figsize=(1.55 * nf, 1.9 * len(seqs)))
    if len(seqs) == 1:
        axes = axes[None, :]
    fig.suptitle(title, fontsize=9)
    for r, (lab, seq) in enumerate(seqs):
        ts = np.linspace(0, len(seq) - 1, nf).astype(int)
        c, hc = cols[r % len(cols)]
        for ci, t in enumerate(ts):
            draw(axes[r, ci], seq[t], c, hc, W, H, lo, hi, lw=1.2)
            if ci == 0:
                axes[r, ci].set_ylabel(lab, fontsize=7)
            axes[r, ci].set_title(f'{t+1}', fontsize=6)
    plt.tight_layout(rect=[0, 0, 1, 0.94])
    png = args.out + '.png'
    fig.savefig(png, dpi=140)
    plt.close(fig)
    print(f'wrote {png}')

    if args.mp4 or args.gif:
        nT = max(len(s) for _, s in seqs)
        fig, axes = plt.subplots(1, len(seqs), figsize=(2.1 * len(seqs), 3.4))
        fig.suptitle(title, fontsize=9)

        def frame(t):
            for k, (lab, seq) in enumerate(seqs):
                tt = min(t, len(seq) - 1)
                c, hc = cols[k % len(cols)]
                draw(axes[k], seq[tt], c, hc, W, H, lo, hi)
                axes[k].set_title(f'{lab}\n{tt+1}/{len(seq)}', fontsize=7)
            return []

        anim = FuncAnimation(fig, frame, frames=nT, interval=1000 / args.fps)
        mp4 = args.out + '.mp4'
        anim.save(mp4, writer=FFMpegWriter(fps=args.fps, bitrate=2400))
        plt.close(fig)
        print(f'wrote {mp4}')

        if args.gif:
            import subprocess
            gif = args.out + '.gif'
            pal = args.out + '.pal.png'
            vf = f'fps={args.gif_fps},scale={args.gif_width}:-1:flags=lanczos'
            subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', mp4,
                            '-vf', vf + ',palettegen=stats_mode=diff', pal], check=True)
            subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', mp4, '-i', pal,
                            '-lavfi', vf + '[x];[x][1:v]paletteuse=dither=bayer:bayer_scale=3',
                            gif], check=True)
            os.remove(pal)
            print(f'wrote {gif} ({os.path.getsize(gif)/1e6:.2f} MB)')


if __name__ == '__main__':
    main()