"""Phase G.0a — Build 8-way hazard category labels for the v3 cache.

Heuristic mapping from (source, category) → hazard index, using the taxonomy
from `lkalert/models/adaptive_window.py:49-58`:
    0 = HAZARD_PEDESTRIAN
    1 = HAZARD_VRURIDER
    2 = HAZARD_VEHICLE_CROSS
    3 = HAZARD_VEHICLE_ONCOMING
    4 = HAZARD_VEHICLE_LEAD
    5 = HAZARD_WEATHER
    6 = HAZARD_INFRASTRUCTURE
    7 = HAZARD_NONE

This is an auxiliary-loss label set — it doesn't need to be ground truth.
The AdaptiveWindow uses hazard logits to bias window choice; even a noisy
3-way effective mapping (non_ego → cross, ego_positive → lead, safe → none)
gives the model a meaningful inductive bias for window selection.

Output: data/policy_labels/hazard_categories_{train_9k,multisrc_val}.json
"""
from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path

import torch

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


# (source, category) → hazard index
# Fallback HAZARD_VEHICLE_LEAD (4) for ambiguous accident cases
def map_to_hazard(source: str, category: str) -> int:
    src = (source or "").lower()
    cat = (category or "").lower()

    # Negative / safe → NONE
    if cat == "safe_neg" or cat.endswith("silent"):
        return 7

    # Non-ego cross-traffic
    if "non_ego" in cat or "cross" in cat:
        return 2  # VEHICLE_CROSS

    # Ego-involved accidents
    if "ego" in cat or cat in ("ego_alert", "ego_observe"):
        if src in ("dota",):
            return 4  # default DoTA ego = lead vehicle
        if src in ("dada",):
            return 3  # DADA ego often oncoming
        if src in ("nexar",):
            return 4  # Nexar ego mostly rear-end / lead
        return 4

    # ego_positive (Nexar / DADA) → lead vehicle
    if "positive" in cat:
        return 4

    # Source-only fallbacks
    if src == "dota":
        return 4  # most DoTA cases are ego-related vehicle
    if src == "dada":
        return 3
    if src == "nexar":
        return 4
    if src == "dad":
        return 4
    return 4   # generic fallback


def build_for_cache(cache_path: Path, out_path: Path):
    cache = torch.load(cache_path, weights_only=False, map_location="cpu")
    ids = cache["ids"]
    sources = cache["source"]
    cats = cache["category"]
    n = len(ids)
    print(f"[load] {cache_path}: N={n}")

    hazard_idx = []
    for i in range(n):
        h = map_to_hazard(sources[i], cats[i])
        hazard_idx.append(h)

    dist = Counter(hazard_idx)
    print(f"  hazard dist: {dict(sorted(dist.items()))}")
    src_dist = Counter(sources)
    cat_dist = Counter(cats)
    print(f"  source dist: {dict(src_dist.most_common(8))}")
    print(f"  category dist: {dict(cat_dist.most_common(8))}")

    out = {
        "schema": "v3_hazard_labels_v1",
        "cache_path": str(cache_path),
        "n_samples": n,
        "taxonomy": {
            0: "PEDESTRIAN", 1: "VRURIDER", 2: "VEHICLE_CROSS",
            3: "VEHICLE_ONCOMING", 4: "VEHICLE_LEAD", 5: "WEATHER",
            6: "INFRASTRUCTURE", 7: "NONE",
        },
        "rule_source": "heuristic (source × category) — auxiliary supervision",
        "labels": hazard_idx,    # parallel to cache["ids"]
        "ids": ids,
        "dist": dict(dist),
    }
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(out, indent=None))
    print(f"[save] {out_path}")


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--train_cache", type=Path,
                    default=ROOT / "data/belief_cache_v3/sft_x_v3__train_9k.pt")
    ap.add_argument("--val_cache", type=Path,
                    default=ROOT / "data/belief_cache_v3/sft_x_v3__multisrc_val.pt")
    ap.add_argument("--out_dir", type=Path,
                    default=ROOT / "data/policy_labels")
    args = ap.parse_args()

    build_for_cache(
        args.train_cache, args.out_dir / "hazard_categories_train_9k.json")
    build_for_cache(
        args.val_cache, args.out_dir / "hazard_categories_multisrc_val.json")


if __name__ == "__main__":
    main()