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

  For each DoTA clip with valid anomaly_start / anomaly_end:
    - [anomaly_start, anomaly_end]  → ALERT
    - In the 3-second pre-anomaly window [anomaly_start - 30, anomaly_start - 1]
      (30 frames @ 10fps), use BADAS to find t_observe:
        - t_observe = first frame index where BADAS p_alert > threshold
        - If no frame crosses threshold, no OBSERVE labels (SILENT→ALERT direct)
    - [t_observe, anomaly_start - 1] → OBSERVE
    - All other frames → SILENT
  For DoTA clips without anomaly (negatives) → all frames SILENT.

Threshold is derived from the per-clip BADAS @ anomaly_start distribution
(eval_results/badas_dota_anomaly_start.json). Default: 25th percentile of
that distribution. Override via --threshold or --threshold_strategy.

USAGE
  # First, score BADAS at anomaly_start (one-time):
  python tools/badas_dota_anomaly_start.py

  # Then score BADAS on every pre-anomaly anchor frame (this script):
  python tools/relabel_dota_corpus.py --threshold_strategy mean
  # or:
  python tools/relabel_dota_corpus.py --threshold 0.05

Reads: data/cot_corpus_v3/v4_sft_{train,val,test}_full_relabeled.jsonl
        eval_results/badas_dota_anomaly_start.json
Writes: data/cot_corpus_v3/v4_sft_{train,val,test}_full_relabeled2.jsonl
         eval_results/badas_dota_pre_anomaly_scores.json   (per-clip pre-window BADAS)
"""
from __future__ import annotations
import argparse
import json
import logging
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from tqdm import tqdm
from torch.utils.data import DataLoader, Dataset

ROOT = Path("PROJECT_ROOT")
BADAS_REPO = Path("~/.cache/huggingface/hub/models--nexar-ai--badas-open/"
                   "snapshots/8fda93711e79d72401b0a4efc151b56455885cd2")
sys.path.insert(0, str(BADAS_REPO / "src"))
import train.video_training  # noqa: F401
from models.vjepa import VJEPAModel

DOTA_FRAMES   = ROOT / "DoTA/frames"
META_TRAIN    = ROOT / "DoTA/metadata_train.json"
META_VAL      = ROOT / "DoTA/metadata_val.json"
COT_DIR       = ROOT / "data/cot_corpus_v3"
ANOMALY_JSON  = ROOT / "eval_results/badas_dota_anomaly_start.json"
PREWIN_JSON   = ROOT / "eval_results/badas_dota_pre_anomaly_scores.json"

DOTA_FPS = 10.0
PREWIN_SECONDS = 2.0
PREWIN_FRAMES = int(PREWIN_SECONDS * DOTA_FPS)  # 20 (20 frames @ 10fps = 2s)
FRAME_COUNT = 16
IMG_SIZE = 224
MODEL_NAME = "facebook/vjepa2-vitl-fpc16-256-ssv2"
CKPT_PATH = str(BADAS_REPO / "weights" / "badas_open.pth")
TEMPERATURE = 2.0

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


# ─────────────────────────── frame loading + BADAS ───────────────────────────

def load_pil_frames_causal(video_name: str, anchor_frame: int,
                            frame_count: int = FRAME_COUNT) -> list[Image.Image]:
    folder = DOTA_FRAMES / video_name / "images"
    if not folder.is_dir():
        return []
    avail = sorted(int(p.stem) for p in folder.glob("*.jpg"))
    if not avail: return []
    avail_np = np.array(avail)
    wanted = list(range(anchor_frame - frame_count + 1, anchor_frame + 1))
    out = []
    for w in wanted:
        if w < avail[0]: w = avail[0]
        k = int(avail_np[np.abs(avail_np - w).argmin()])
        for width in (6, 5, 4, 3):
            cand = folder / f"{k:0{width}d}.jpg"
            if cand.exists():
                out.append(Image.open(cand).convert("RGB"))
                break
    return out


class AnchorDS(Dataset):
    """Each item is (video, anchor_frame) for the pre-anomaly window."""
    def __init__(self, items: list[tuple[str, int]], processor):
        self.items = items
        self.processor = processor

    def __len__(self): return len(self.items)

    def __getitem__(self, i):
        vname, anchor = self.items[i]
        frames = load_pil_frames_causal(vname, anchor)
        if len(frames) < FRAME_COUNT:
            if frames:
                frames = [frames[0]] * (FRAME_COUNT - len(frames)) + frames
            else:
                frames = [Image.new("RGB", (IMG_SIZE, IMG_SIZE))] * FRAME_COUNT
        proc = self.processor(videos=[frames], return_tensors="pt")
        if "pixel_values_videos" in proc:
            video = proc["pixel_values_videos"].squeeze(0)
        elif "pixel_values" in proc:
            video = proc["pixel_values"].squeeze(0)
        else:
            video = list(proc.values())[0].squeeze(0)
        return {"video": video, "video_name": vname, "anchor": int(anchor)}


def coll(batch):
    return {
        "videos":     torch.stack([b["video"] for b in batch]),
        "video_name": [b["video_name"] for b in batch],
        "anchor":     [b["anchor"] for b in batch],
    }


@torch.no_grad()
def forward(model, videos, device):
    """bf16 autocast for 2× speedup; softmax computed in fp32 for numerical safety."""
    videos = videos.to(device, non_blocking=True)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out = model(videos)
    logits = out.float() / TEMPERATURE
    probs = torch.softmax(logits, dim=1)[:, 1]
    return probs.cpu().numpy()


# ─────────────────────────── threshold derivation ───────────────────────────

def derive_threshold(strategy: str, override: float | None = None) -> float:
    if override is not None and override > 0:
        logger.info(f"[threshold] using override = {override:.4f}")
        return float(override)
    if not ANOMALY_JSON.exists():
        raise FileNotFoundError(f"{ANOMALY_JSON} not found — run tools/badas_dota_anomaly_start.py first")
    d = json.loads(ANOMALY_JSON.read_text())
    scores = [r["p_alert_at_anomaly_start"] for r in d["per_clip"].values()]
    arr = np.asarray(scores, dtype=np.float64)
    options = {
        "mean":   float(arr.mean()),
        "median": float(np.median(arr)),
        "p25":    float(np.percentile(arr, 25)),
        "p10":    float(np.percentile(arr, 10)),
    }
    logger.info(f"[threshold] distribution at anomaly_start (N={arr.size}):")
    for k, v in options.items():
        logger.info(f"  {k:8s} = {v:.4f}")
    return options[strategy]


# ─────────────────────────── label rewriter ───────────────────────────

def rewrite_dota_labels(actions_pf: list[str], tta_pf: list[float],
                          tick_action: str, tick_tta: float,
                          anomaly_start: int, anomaly_end: int,
                          t_observe: int | None,
                          frame_indices: list[int]) -> tuple[list[str], str]:
    """For each of the 8 frame indices in this tick, assign:
       [anomaly_start, anomaly_end] → ALERT
       [t_observe, anomaly_start - 1] → OBSERVE (if t_observe is not None)
       else → SILENT
    """
    new_actions = []
    for f in frame_indices:
        if anomaly_start is not None and anomaly_end is not None and \
           anomaly_start <= f <= anomaly_end:
            new_actions.append("ALERT")
        elif (t_observe is not None and anomaly_start is not None
              and t_observe <= f < anomaly_start):
            new_actions.append("OBSERVE")
        else:
            new_actions.append("SILENT")
    # Tick label = last frame of the 8-frame window (per existing convention)
    new_tick = new_actions[-1]
    return new_actions, new_tick


# ─────────────────────────── main ───────────────────────────

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--threshold_strategy", choices=["mean", "median", "p25", "p10"],
                    default="p25",
                    help="how to derive the OBSERVE threshold from the per-clip "
                         "BADAS @ anomaly_start distribution")
    ap.add_argument("--threshold", type=float, default=0.0,
                    help="override threshold (>0)")
    ap.add_argument("--batch_size",  type=int, default=8)
    ap.add_argument("--num_workers", type=int, default=2)
    ap.add_argument("--skip_badas", action="store_true",
                    help="reuse existing pre-window BADAS scores (no GPU run)")
    args = ap.parse_args()

    threshold = derive_threshold(args.threshold_strategy, args.threshold or None)
    logger.info(f"[threshold] FINAL = {threshold:.4f}  (strategy={args.threshold_strategy})")

    # ── Build per-clip list of (video, pre-window anchors) ──
    meta = {}
    for p in (META_TRAIN, META_VAL):
        meta.update(json.loads(p.read_text()))
    items = []
    skipped = 0
    for vid, m in meta.items():
        a_start = m.get("anomaly_start"); a_end = m.get("anomaly_end")
        if a_start is None or a_start <= 0:
            skipped += 1; continue
        if not (DOTA_FRAMES / vid / "images").is_dir():
            skipped += 1; continue
        win_lo = max(0, a_start - PREWIN_FRAMES)
        win_hi = a_start - 1
        items.append({"video_name": vid, "anomaly_start": int(a_start),
                       "anomaly_end": int(a_end) if a_end else None,
                       "pre_anchors": list(range(win_lo, win_hi + 1))})
    logger.info(f"DoTA clips with anomaly_start: {len(items)} (skipped {skipped})")

    # ── Pre-window BADAS scoring (one anchor per pre-window frame) ──
    if not args.skip_badas:
        logger.info(f"Loading V-JEPA2 …")
        vjepa = VJEPAModel(model_name=MODEL_NAME, checkpoint_path=CKPT_PATH,
                            frame_count=FRAME_COUNT, img_size=IMG_SIZE,
                            window_stride=1, target_fps=8.0,
                            use_sliding_window=False)
        vjepa.load()
        device = vjepa.device

        flat = [(it["video_name"], a) for it in items for a in it["pre_anchors"]]
        logger.info(f"  total anchors to score: {len(flat)}")

        # ── Resume support: skip anchors already in checkpoint ──
        per_anchor: dict[tuple[str, int], float] = {}
        ckpt_path = PREWIN_JSON.parent / "_pre_anomaly_anchors_ckpt.json"
        if ckpt_path.exists():
            ck = json.loads(ckpt_path.read_text())
            # JSON keys are strings "vname|anchor"
            for k, v in ck.items():
                vname, anchor = k.rsplit("|", 1)
                per_anchor[(vname, int(anchor))] = float(v)
            logger.info(f"  [resume] loaded {len(per_anchor)} anchors from {ckpt_path}")
            flat = [t for t in flat if t not in per_anchor]
            logger.info(f"  {len(flat)} anchors remaining")

        ds = AnchorDS(flat, processor=vjepa.processor)
        loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False,
                              num_workers=args.num_workers, collate_fn=coll,
                              pin_memory=True,
                              persistent_workers=(args.num_workers > 0),
                              prefetch_factor=4 if args.num_workers > 0 else None)

        def _save_ckpt():
            tmp = {f"{k[0]}|{k[1]}": v for k, v in per_anchor.items()}
            ckpt_path.parent.mkdir(parents=True, exist_ok=True)
            tmp_path = ckpt_path.with_suffix(".json.tmp")
            tmp_path.write_text(json.dumps(tmp))
            tmp_path.replace(ckpt_path)

        SAVE_EVERY = 5000     # incremental save cadence (anchors)
        pbar = tqdm(total=len(flat), desc="badas", ncols=110,
                     unit="anc", smoothing=0.05, dynamic_ncols=False)
        n_done = 0
        for batch in loader:
            probs = forward(vjepa.model, batch["videos"], device)
            for vn, an, p in zip(batch["video_name"], batch["anchor"], probs):
                per_anchor[(vn, int(an))] = float(p)
            n_done += len(probs)
            pbar.update(len(probs))
            if n_done % 200 == 0:
                pbar.set_postfix(gpu_GB=f"{torch.cuda.memory_allocated()/1e9:.1f}")
            if n_done % SAVE_EVERY == 0:
                _save_ckpt()
        pbar.close()

        _save_ckpt()
        logger.info(f"[ckpt] final save → {ckpt_path}")

        # Save per-window BADAS scores per clip
        scores_by_clip: dict[str, dict] = {}
        for it in items:
            vname = it["video_name"]
            per_frame = {int(a): per_anchor.get((vname, int(a)), float("nan"))
                          for a in it["pre_anchors"]}
            scores_by_clip[vname] = {
                "anomaly_start": it["anomaly_start"],
                "pre_anchors":   it["pre_anchors"],
                "scores":        per_frame,
            }
        PREWIN_JSON.parent.mkdir(parents=True, exist_ok=True)
        PREWIN_JSON.write_text(json.dumps(scores_by_clip, indent=2))
        logger.info(f"[save] {PREWIN_JSON}  ({len(scores_by_clip)} clips)")
    else:
        if not PREWIN_JSON.exists():
            raise FileNotFoundError(f"--skip_badas set but {PREWIN_JSON} doesn't exist")
        scores_by_clip = json.loads(PREWIN_JSON.read_text())
        logger.info(f"[skip_badas] loaded {len(scores_by_clip)} clips from {PREWIN_JSON}")

    # ── Determine t_observe per clip ──
    t_observe_by_clip: dict[str, int | None] = {}
    n_with_obs = n_without = 0
    for vname, info in scores_by_clip.items():
        anchors = info["pre_anchors"]
        scs = info["scores"]
        # Sort anchors ascending and find the FIRST one that crosses threshold
        first_cross = None
        for a in sorted(anchors):
            v = scs.get(str(a), scs.get(a))  # handle JSON int-as-str keys
            if v is None or not np.isfinite(v): continue
            if v > threshold:
                first_cross = int(a); break
        t_observe_by_clip[vname] = first_cross
        if first_cross is None: n_without += 1
        else: n_with_obs += 1
    logger.info(f"[t_observe] {n_with_obs} clips have OBSERVE window, "
                 f"{n_without} go SILENT→ALERT direct (no crossing)")

    # ── Rewrite corpus jsonl ──
    for split_tag in ["v4_sft_train_full", "v4_sft_val_full", "v4_sft_test_full"]:
        in_path = COT_DIR / f"{split_tag}_relabeled.jsonl"
        out_path = COT_DIR / f"{split_tag}_relabeled2.jsonl"
        if not in_path.exists():
            logger.warning(f"[skip] {in_path} not found")
            continue
        n_total = n_dota = n_changed = 0
        before = Counter(); after = 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", "")
                if src != "dota":
                    fout.write(json.dumps(rec) + "\n"); continue
                n_dota += 1
                # DoTA video id in corpus has "dota_" prefix; metadata keys don't
                vid_raw = rec.get("video_id") or rec.get("clip_id") or ""
                vid = vid_raw.replace("dota_", "", 1) if vid_raw.startswith("dota_") else vid_raw
                m = meta.get(vid, {})
                a_start = m.get("anomaly_start"); a_end = m.get("anomaly_end")
                t_obs = t_observe_by_clip.get(vid)

                frame_idx = rec.get("frame_indices", [])
                if len(frame_idx) != 8 or a_start is None or a_start <= 0:
                    # No anomaly window or malformed → keep all SILENT
                    new_acts = ["SILENT"] * 8
                    new_tick = "SILENT"
                else:
                    new_acts, new_tick = rewrite_dota_labels(
                        rec.get("actions_per_frame", []),
                        rec.get("tta_per_frame", []),
                        rec.get("tick_action", ""),
                        rec.get("tick_tta_raw", -1.0),
                        a_start, a_end, t_obs, frame_idx)
                before[rec.get("tick_action", "?")] += 1
                rec["actions_per_frame"] = new_acts
                rec["tick_action"] = new_tick
                after[new_tick] += 1
                if rec.get("tick_action") != before:
                    n_changed += 1
                fout.write(json.dumps(rec) + "\n")
        logger.info(f"[{split_tag}] N={n_total}  DoTA={n_dota}  saved → {out_path}")
        logger.info(f"  before tick_action: {dict(before)}")
        logger.info(f"  after  tick_action: {dict(after)}")


if __name__ == "__main__":
    main()