#!/usr/bin/env python3
"""Generation diversity: spread of generated poses ACROSS sampling seeds, per clip.

FGD only *implies* mode collapse (its source paper calls it a diversity measure). This
measures it directly: for each clip, take the N sampled generations, resample to a common
length, and take the per-frame/per-joint sd across seeds. Falling = the conditional
distribution is narrowing, i.e. decoding is becoming deterministic. Reported alongside
mean length and length sd so a pure length collapse can be told apart from a real one.
"""
import argparse
import glob
import os
import re

import numpy as np


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


def resamp(x, T):
    idx = np.linspace(0, len(x) - 1, T)
    lo, hi = np.floor(idx).astype(int), np.ceil(idx).astype(int)
    w = (idx - lo)[:, None, None]
    return x[lo] * (1 - w) + x[hi] * w


def diversity(paths):
    ds = [load(p) for p in paths]
    if len(ds) < 2:
        return None
    common = set(ds[0])
    for d in ds[1:]:
        common &= set(d)
    sds, lsd, mlen = [], [], []
    for n in sorted(common):
        seqs = [np.asarray(d[n], np.float32) for d in ds]
        L = [len(s) for s in seqs]
        T = int(np.median(L))
        if T < 2:
            continue
        lsd.append(float(np.std(L, ddof=1)))
        mlen.append(float(np.mean(L)))
        st = np.stack([resamp(s, T) for s in seqs])
        sds.append(float(st.std(axis=0, ddof=1).mean()))
    return {"n_clips": len(sds), "cross_seed_sd": float(np.mean(sds)),
            "len_sd": float(np.mean(lsd)), "mean_len": float(np.mean(mlen))}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dump-dir", required=True)
    ap.add_argument("--pattern", default="it{it}_s*.npz",
                    help="{it} is substituted with each --iters value")
    ap.add_argument("--iters", nargs="+", type=int, required=True)
    args = ap.parse_args()
    print(f"{'iter':>8} {'cross-seed sd':>14} {'len sd':>9} {'mean len':>9} {'clips':>6}")
    for it in args.iters:
        ps = sorted(glob.glob(os.path.join(args.dump_dir, args.pattern.format(it=it))))
        r = diversity(ps)
        if r:
            print(f"{it:>8} {r['cross_seed_sd']:>14.5f} {r['len_sd']:>9.2f} "
                  f"{r['mean_len']:>9.1f} {r['n_clips']:>6}")


if __name__ == "__main__":
    main()