"""Rewrite DAD per-frame action labels in cot_corpus_v3 manifests per user rule:

  DAD positives (event at t=3s of 4s @ 25fps clip):
      → all 8 frames of every tick → ALERT
      → tick_action = ALERT
  DAD negatives (no event):
      → all 8 frames → SILENT
      → tick_action = SILENT

No OBSERVE state for DAD.

Reads: data/cot_corpus_v3/v4_sft_{train,val,test}_full.jsonl
Writes: data/cot_corpus_v3/v4_sft_{train,val,test}_full_relabeled.jsonl
"""
from __future__ import annotations
import json
import logging
from collections import Counter
from pathlib import Path

ROOT = Path("PROJECT_ROOT")
COT_DIR = ROOT / "data/cot_corpus_v3"
SPLITS = ["v4_sft_train_full", "v4_sft_val_full", "v4_sft_test_full"]

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("dad_relabel")


def is_dad_positive(rec: dict) -> bool:
    """A DAD record is positive iff its tta_raw indicates a known accident.
    DAD positives have tta_raw > 0 in the manifest (they're aligned so the
    last frame is near t=3s, the hardcoded event time)."""
    tta = rec.get("tick_tta_raw", -1.0)
    return rec.get("source") == "dad" and tta is not None and tta >= 0


def relabel_dad(rec: dict) -> tuple[dict, str]:
    """Return (new_record, change_kind) where change_kind ∈ {kept, alert, silent}."""
    if rec.get("source") != "dad":
        return rec, "kept"

    new = dict(rec)
    if is_dad_positive(rec):
        # All 8 frames → ALERT
        new["actions_per_frame"] = ["ALERT"] * 8
        new["tick_action"] = "ALERT"
        change = "alert"
    else:
        # Safe / negative → all SILENT
        new["actions_per_frame"] = ["SILENT"] * 8
        new["tick_action"] = "SILENT"
        change = "silent"
    return new, change


def process_split(split_tag: str) -> dict:
    in_path = COT_DIR / f"{split_tag}.jsonl"
    out_path = COT_DIR / f"{split_tag}_relabeled.jsonl"
    if not in_path.exists():
        logger.warning(f"[skip] {in_path} not found")
        return {}

    n_total = n_dad = n_alert = n_silent = n_other = 0
    before_tick = Counter()
    after_tick = Counter()
    by_src = Counter()
    with in_path.open() as fin, out_path.open("w") as fout:
        for ln in fin:
            ln = ln.strip()
            if not ln: continue
            rec = json.loads(ln)
            n_total += 1
            src = rec.get("source", "?")
            by_src[src] += 1
            before_tick[(src, rec.get("tick_action", "?"))] += 1

            new, kind = relabel_dad(rec)
            if src == "dad":
                n_dad += 1
                if kind == "alert": n_alert += 1
                elif kind == "silent": n_silent += 1
                else: n_other += 1
            after_tick[(new.get("source", "?"), new.get("tick_action", "?"))] += 1
            fout.write(json.dumps(new) + "\n")

    logger.info(f"[{split_tag}] N={n_total}  DAD records={n_dad}  "
                 f"→ ALERT={n_alert}  → SILENT={n_silent}  unchanged={n_other}")
    logger.info(f"[{split_tag}] saved → {out_path}")
    return {
        "split": split_tag,
        "in_path": str(in_path),
        "out_path": str(out_path),
        "n_total": n_total,
        "n_dad": n_dad,
        "n_dad_positive_to_alert": n_alert,
        "n_dad_negative_to_silent": n_silent,
        "by_source_before": {f"{k[0]}/{k[1]}": v for k, v in sorted(before_tick.items())
                              if k[0] == "dad"},
        "by_source_after":  {f"{k[0]}/{k[1]}": v for k, v in sorted(after_tick.items())
                              if k[0] == "dad"},
    }


def main():
    out_summary = []
    for tag in SPLITS:
        out_summary.append(process_split(tag))
    summary_path = COT_DIR / "_relabel_dad_summary.json"
    summary_path.write_text(json.dumps(out_summary, indent=2))
    logger.info(f"[summary] saved → {summary_path}")

    # Verification log
    print("\n=== DAD RELABEL SUMMARY ===")
    for s in out_summary:
        print(f"\n{s['split']}:  {s['n_dad']} DAD records → "
              f"{s['n_dad_positive_to_alert']} ALERT, {s['n_dad_negative_to_silent']} SILENT")
        print("  before:")
        for k, v in s["by_source_before"].items():
            print(f"    {k}: {v}")
        print("  after:")
        for k, v in s["by_source_after"].items():
            print(f"    {k}: {v}")


if __name__ == "__main__":
    main()