"""Re-align tick_label across all per_tick PTs to a single canonical scheme.

Problem: different scorers used different labeling rules and different manifest
snapshots, so the same (video_id, tick_idx) row can have different
`tick_label` and `tta_raw` across PT files. This makes the comparison unfair
(each method evaluated against its OWN ground truth).

Fix: pick ONE canonical (video_id, tick_idx) → (tick_label, tta_raw) mapping
from a reference PT (vlalert_x_c1_seed5.pt, the winner, which uses the
sft_x_v3 belief cache labels), then overwrite the corresponding fields in
every other PT in eval_results/benchmark_v1_val/per_tick/.

Backs up originals to per_tick_orig/ before rewriting.

Run: python tools/relabel_per_tick_canonical.py
"""
from __future__ import annotations
import shutil
from collections import Counter
from pathlib import Path

import torch

ROOT = Path("PROJECT_ROOT")
PT_DIR  = ROOT / "eval_results/benchmark_v1_val/per_tick"
BACKUP  = ROOT / "eval_results/benchmark_v1_val/per_tick_orig"
REF_PT  = PT_DIR / "vlalert_x_c1_seed5.pt"


def main():
    print(f"[ref] {REF_PT.name}")
    ref = torch.load(REF_PT, weights_only=False, map_location="cpu")
    canonical = {}   # (vid, tick_idx) → (label, tta_raw)
    for i, (vid, ti, lab, tta) in enumerate(zip(
            ref["ids"], ref["tick_idx"].tolist(),
            ref["tick_label"].tolist(), ref["tta_raw"].tolist())):
        canonical[(vid, int(ti))] = (int(lab), float(tta))

    # Drop the dummy ('', 0) bucket that collects DoTA frame-folder failures
    canonical.pop(("", 0), None)
    print(f"[ref] {len(canonical):,} canonical (vid, tick_idx) entries")
    print(f"[ref] label dist: {Counter(l for l, _ in canonical.values())}")

    BACKUP.mkdir(parents=True, exist_ok=True)

    for pt in sorted(PT_DIR.glob("*.pt")):
        if pt == REF_PT:
            continue   # skip the reference
        # Backup once
        bk = BACKUP / pt.name
        if not bk.exists():
            shutil.copy2(pt, bk)
        d = torch.load(pt, weights_only=False, map_location="cpu")
        ids  = list(d["ids"])
        tidx = d["tick_idx"].tolist()
        new_labels = torch.zeros(len(ids), dtype=torch.long)
        new_tta    = torch.zeros(len(ids), dtype=torch.float)
        n_match = n_miss = 0
        for i, (vid, ti) in enumerate(zip(ids, tidx)):
            key = (vid, int(ti))
            if key in canonical:
                lab, tta = canonical[key]
                new_labels[i] = lab
                new_tta[i] = tta
                n_match += 1
            else:
                # No canonical entry → mark INVALID (-1) so aggregators skip.
                # This applies to (a) DoTA frame-folder failures, (b) any tick
                # in the manifest that the belief cache couldn't materialize.
                new_labels[i] = -1
                new_tta[i] = float("nan")
                n_miss += 1
        old_dist = Counter(d["tick_label"].tolist())
        d["tick_label"] = new_labels
        d["tta_raw"] = new_tta
        torch.save(d, pt)
        new_dist = Counter(new_labels.tolist())
        change = "no-change" if old_dist == new_dist else "RELABELED"
        print(f"  {pt.name:35s} n_match={n_match:5d} n_miss={n_miss:3d}  "
              f"old {dict(old_dist)} → new {dict(new_dist)}  {change}")


if __name__ == "__main__":
    main()