#!/usr/bin/env python3
"""Build the hands-DTW-vs-keypoint-weight table from the sweep's eval JSONs.

One row per stage-1 weight setting. Reported in shoulder widths on the same 300
test clips (`--seed 0`), so rows are paired and comparable to the reference table
in `RESULTS_VSL.md`.

Two columns exist for a reason:
  ceiling   stage-1 VQ round-trip -- this is what the weighting DIRECTLY optimizes
  model     stage-2 free-running generation -- this is what we actually care about
A weighting that moves the ceiling but not the model tells us the bottleneck is
stage 2, which is the hypothesis under test.

Every model number is a mean over repeated SAMPLING seeds with the clip subset
held fixed, because categorial sampling carries ~+-0.01 run-to-run noise on hands
and single-seed margins below that are unresolvable.
"""
import json
import os
import re

import numpy as np

FLOOR_HANDS = 0.6950   # random real train clip of the wrong sentence (trivis_floors.json)

# tag: (label, vq_dir, gpt_dir)
ROWS = [
    ("uni", "1.0 / 1.0 / 1.0",            "vq_vsl_fl_uni", "gpt_vsl_fl_uni"),
    ("wh1", "1.0 / 0.5 / 1.0",            "vq_vsl_fl_wh1", "gpt_vsl_fl_wh1"),
    ("wh3", "1.0 / 0.5 / 3.0  (current)", "vq_vsl_front_lab", "gpt_vsl_front_lab_v2"),
    ("wh6", "1.0 / 0.5 / 6.0",            "vq_vsl_fl_wh6", "gpt_vsl_fl_wh6"),
    ("ft",  "1.0 / 0.5 / 3.0 +tip 9.0",   "vq_vsl_fl_ft",  "gpt_vsl_fl_ft"),
]
OUT = "output_vsl"


def stage1_info(vq_dir):
    """Raw (unweighted) val hands MPJPE + codes used, from the run log's best line."""
    log = os.path.join(OUT, vq_dir, "run.log")
    if not os.path.exists(log):
        return None, None
    best_hands, codes = None, None
    with open(log, errors="replace") as f:
        for line in f:
            m = re.search(r"hands ([0-9.]+) \| codes used (\d+)", line)
            if m:
                h, c = float(m.group(1)), int(m.group(2))
                if best_hands is None or h < best_hands:
                    best_hands, codes = h, c
    return best_hands, codes


def collect(gpt_dir, which="best"):
    """hands (shoulder, gt_anchor) per condition, across sampling seeds.

    which='best' reads the net_best evals, which='last' the matched 30k net_last
    evals. Read 'last' for the weighting comparison: net_best is selected on the
    flat/noisy 80-clip metric and landed on 2000-26000 iters depending on config.
    """
    pat = r"wsweep_s(\d+)\.json" if which == "best" else r"wsweep_last_s(\d+)\.json"
    d = os.path.join(OUT, gpt_dir)
    if not os.path.isdir(d):
        return {}
    acc = {}
    for fn in sorted(os.listdir(d)):
        m = re.fullmatch(pat, fn)
        if not m:
            continue
        with open(os.path.join(d, fn)) as f:
            r = json.load(f).get("results", {})
        for cond, v in r.items():
            h = v.get("gt_anchor", {}).get("hands", {}).get("mean")
            if h is not None:
                acc.setdefault(cond, []).append(h)
    return acc


def fmt(vals):
    if not vals:
        return "--"
    a = np.asarray(vals, float)
    if len(a) == 1:
        return f"{a[0]:.4f}"
    return f"{a.mean():.4f} ±{a.std(ddof=1):.4f}"


def best_iter(gpt_dir):
    log = os.path.join(OUT, gpt_dir, "run.log")
    if not os.path.exists(log):
        return None
    for line in reversed(open(log, errors="replace").readlines()):
        m = re.search(r"best gen hands MPJPE [0-9.]+ @ iter (\d+)", line)
        if m:
            return int(m.group(1))
    return None


def table(which, note):
    print(f"\n### stage-2 net_{which} -- {note}")
    print(f"| {'tag':4} | {'body/face/hand':26} | {'s1 hands':8} | {'ceiling':8} "
          f"| {'model (oracle)':16} | {'shuffled':16} | {'shuf pen':8} "
          f"| {'vs floor':8} | {'s2 iter':7} | {'n':1} |")
    print("|" + "|".join("-" * w for w in
                         [6, 28, 10, 10, 18, 18, 10, 10, 9, 3]) + "|")
    rows = {}
    for tag, label, vq_dir, gpt_dir in ROWS:
        s1h, _ = stage1_info(vq_dir)
        acc = collect(gpt_dir, which)
        orc, shf, cei = (acc.get("oracle-gloss", []), acc.get("shuffled-half", []),
                         acc.get("ceiling", []))
        pen = (f"{100 * (np.mean(shf) - np.mean(orc)) / np.mean(orc):+.1f}%"
               if orc and shf else "--")
        vsf = (f"{100 * (FLOOR_HANDS - np.mean(orc)) / FLOOR_HANDS:+.1f}%"
               if orc else "--")
        it = 30000 if which == "last" else best_iter(gpt_dir)
        print(f"| {tag:4} | {label:26} | "
              f"{(f'{s1h:.5f}' if s1h is not None else '--'):8} | "
              f"{(f'{np.mean(cei):.4f}' if cei else '--'):8} | "
              f"{fmt(orc):16} | {fmt(shf):16} | {pen:8} | {vsf:8} | "
              f"{(str(it) if it else '--'):7} | {len(orc):1} |")
        rows[tag] = {"label": label, "stage1_val_hands": s1h, "evals": acc,
                     "stage2_iter": it}
    return rows


def main():
    print()
    print("hands DTW-MJE, shoulder widths, gt_anchor, 300 test clips (--seed 0, paired)")
    print(f"floor (random wrong real clip) = {FLOOR_HANDS:.4f}")
    print("ceiling is deterministic (no sampling) -> no +-; model/shuffled are mean +-sd "
          "over sampling seeds")

    last = table("last", "MATCHED 30k iters -- READ THIS ONE for the weighting effect")
    best = table("best", "net_best: selection iter varies 2k-26k, CONFOUNDED")

    print()
    print("s1 hands = stage-1 val MPJPE, frame-normalized, UNWEIGHTED -> comparable across")
    print("           settings even though val_recon (weighted) is not. Stage 1 is converged")
    print("           flat, so its own best/last choice moves nothing.")
    json_out = os.path.join(OUT, "hand_weight_sweep.json")
    with open(json_out, "w") as f:
        json.dump({"floor_hands": FLOOR_HANDS, "net_last": last, "net_best": best}, f,
                  indent=1)
    print(f"\nwrote {json_out}")


if __name__ == "__main__":
    main()