"""User-directed relabel: ALERT samples with tta_raw ∈ [2.0, 4.0) → OBSERVE.

Rationale: ALERT @ [0, 2)s works well; the 1225 train ALR samples at tta ∈ [2,4)
are "early hazard" — better suited as OBSERVE training signal so the model
can `look more carefully' on borderline cases rather than fire ALERT early.

Applies to all 3 train caches (narrow/mid/wide) — they share id ordering.
Does NOT modify val caches (those keep original GT for honest eval).

Output: data/belief_cache_v3/sft_x_v3__train_9k{,_narrow,_wide}_relabel.pt
"""
from __future__ import annotations

import argparse
from pathlib import Path
from collections import Counter
import torch

ROOT = Path(__file__).resolve().parents[1]


def relabel(cache_path: Path, out_path: Path,
             tta_lo: float = 2.0, tta_hi: float = 4.0) -> dict:
    print(f"[load] {cache_path}")
    c = torch.load(cache_path, weights_only=False, map_location="cpu")
    ta = c["tick_action"].clone()
    tta = c["tick_tta_raw"]

    before_dist = Counter(ta.tolist())
    # Mask: ALERT-truth (action==2) AND tta ∈ [tta_lo, tta_hi)
    mask = (ta == 2) & (tta >= tta_lo) & (tta < tta_hi)
    n_relabel = int(mask.sum().item())
    ta[mask] = 1   # → OBSERVE
    after_dist = Counter(ta.tolist())

    c["tick_action"] = ta
    c["schema"] = c.get("schema", "vlalert_x_v2_dual_pool") + f"+relabel_alr_{tta_lo:.1f}_{tta_hi:.1f}_to_obs"

    print(f"  before: {dict(sorted(before_dist.items()))}")
    print(f"  after : {dict(sorted(after_dist.items()))}")
    print(f"  relabeled {n_relabel} ALR → OBS (tta ∈ [{tta_lo}, {tta_hi}))")
    torch.save(c, out_path)
    print(f"[save] {out_path}\n")
    return {"n_relabel": n_relabel, "before": dict(before_dist),
             "after": dict(after_dist)}


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--tta_lo", type=float, default=2.0)
    ap.add_argument("--tta_hi", type=float, default=4.0)
    args = ap.parse_args()

    base = ROOT / "data/belief_cache_v3"
    runs = [
        (base / "sft_x_v3__train_9k.pt",         base / "sft_x_v3__train_9k_relabel.pt"),
        (base / "sft_x_v3__train_9k_narrow.pt",  base / "sft_x_v3__train_9k_narrow_relabel.pt"),
        (base / "sft_x_v3__train_9k_wide.pt",    base / "sft_x_v3__train_9k_wide_relabel.pt"),
    ]
    for src, dst in runs:
        relabel(src, dst, args.tta_lo, args.tta_hi)
    print("=" * 50)
    print("All 3 train caches relabeled. Val caches unchanged.")


if __name__ == "__main__":
    main()