#!/usr/bin/env python3
"""
SFT sanity checks — run BEFORE any long training run.

Checks
------
1. Manifest counts and category distribution
2. Dataset: sample counts, label ranges, no train/val video overlap
3. Model: HazardHead init (output ≈ 0.27), TTAHead init (output ≈ 5.0)
4. Single-step forward pass: shapes, loss components
5. Single-step backward pass: gradients flow through all heads

Usage
-----
python -m training.SFT.sanity_check \
    --manifest_dir PROJECT_ROOT/data/sft_manifests \
    --model_name   PROJECT_ROOT/models/Qwen2.5-VL-3B-Instruct \
    --pretrained_lora PROJECT_ROOT/checkpoints/pretrain_v2/stage_b/best_model
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path
from typing import List

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

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

PASS = "✅"
FAIL = "❌"
WARN = "⚠️"


def check(cond: bool, msg: str, fatal: bool = True):
    if cond:
        logger.info(f"{PASS} {msg}")
    else:
        if fatal:
            logger.error(f"{FAIL} {msg}")
            sys.exit(1)
        else:
            logger.warning(f"{WARN} {msg}")


# ── 1. Manifest checks ────────────────────────────────────────────────────────

def check_manifests(manifest_dir: Path):
    logger.info("\n=== 1. Manifest Checks ===")
    manifest_dir = Path(manifest_dir)

    expected_train = ["nexar_train", "dada_pos_train", "dada_noneego_train", "dada_neg_train"]
    expected_val   = ["nexar_val",   "dada_pos_val",   "dada_noneego_val"]

    for name in expected_train + expected_val:
        p = manifest_dir / f"{name}.json"
        check(p.exists(), f"Manifest exists: {p.name}")
        if not p.exists():
            continue
        with open(p) as f:
            m = json.load(f)
        nv = m.get("num_videos", 0)
        cc = m.get("category_counts", {})
        logger.info(f"    {name}: {nv} videos  cats={cc}")
        check(nv > 0, f"{name} has > 0 videos")

    # Check no video_id overlap between train and val
    train_ids: set = set()
    val_ids:   set = set()
    for name in expected_train:
        p = manifest_dir / f"{name}.json"
        if not p.exists():
            continue
        with open(p) as f:
            m = json.load(f)
        for v in m.get("videos", []):
            train_ids.add(v["video_id"])
    for name in expected_val:
        p = manifest_dir / f"{name}.json"
        if not p.exists():
            continue
        with open(p) as f:
            m = json.load(f)
        for v in m.get("videos", []):
            val_ids.add(v["video_id"])

    overlap = train_ids & val_ids
    check(len(overlap) == 0,
          f"Zero train/val video_id overlap (found {len(overlap)})" if overlap else "Zero train/val video_id overlap",
          fatal=True)
    logger.info(f"    train videos: {len(train_ids)}  val videos: {len(val_ids)}")


# ── 2. Dataset checks ─────────────────────────────────────────────────────────

def check_datasets(manifest_dir: Path, n_samples: int = 200):
    logger.info("\n=== 2. Dataset Checks ===")
    from .dataset import SFTDataset, sft_collate_fn, MAX_TTA, MIN_TTA

    manifest_dir = Path(manifest_dir)
    train_manifests = [
        manifest_dir / "nexar_train.json",
        manifest_dir / "dada_pos_train.json",
        manifest_dir / "dada_noneego_train.json",
        manifest_dir / "dada_neg_train.json",
    ]
    val_manifests = [
        manifest_dir / "nexar_val.json",
        manifest_dir / "dada_pos_val.json",
        manifest_dir / "dada_noneego_val.json",
    ]
    train_manifests = [m for m in train_manifests if m.exists()]
    val_manifests   = [m for m in val_manifests   if m.exists()]

    logger.info("  Loading train dataset (debug mode)...")
    train_ds = SFTDataset(manifests=train_manifests, split="train", debug=True, debug_samples=n_samples)
    logger.info("  Loading val dataset (debug mode)...")
    val_ds   = SFTDataset(manifests=val_manifests,   split="val",   debug=True, debug_samples=n_samples // 2)

    check(len(train_ds) > 0, f"Train dataset non-empty ({len(train_ds)} samples)")
    check(len(val_ds)   > 0, f"Val dataset non-empty ({len(val_ds)} samples)")

    # Inspect a few samples
    loader = DataLoader(train_ds, batch_size=4, shuffle=False, collate_fn=sft_collate_fn, num_workers=0)
    batch = next(iter(loader))

    tta_labels  = batch["tta_labels"]
    haz_labels  = batch["hazard_labels"]
    haz_weights = batch["hazard_weights"]
    is_ego      = batch["is_ego_positive"]
    is_cen      = batch["is_censored"]

    check(tta_labels.shape[0] == 4,        f"Batch size correct: {tta_labels.shape[0]}")
    check((haz_labels >= 0).all() and (haz_labels <= 1).all(),
          f"hazard_labels in [0,1]: min={haz_labels.min():.2f} max={haz_labels.max():.2f}")
    check((haz_weights > 0).all() and (haz_weights <= 1.01).all(),
          f"hazard_weights in (0,1]: min={haz_weights.min():.2f} max={haz_weights.max():.2f}")
    # TTA labels for censored can be > 10 (will be clamped in loss)
    check((tta_labels >= 0).all(),
          f"tta_labels >= 0: min={tta_labels.min():.2f}")

    # ego_positive must have hazard_label=1
    if is_ego.any():
        ego_haz = haz_labels[is_ego]
        check((ego_haz == 1).all(),
              f"ego_positive samples have hazard_label=1 (all {int(is_ego.sum())} checked)")

    # non-ego and safe-neg must always be censored (no TTA supervision)
    is_ne    = batch.get("is_non_ego", torch.zeros(tta_labels.shape[0], dtype=torch.bool))
    non_pos  = ~is_ego
    if non_pos.any():
        check(is_cen[non_pos].all(),
              f"All non-ego/safe-neg samples are censored ({int(non_pos.sum())} non-pos, {int(is_cen[non_pos].sum())} censored)")

    logger.info(
        f"    Sample breakdown — ego_pos={int(is_ego.sum())}  "
        f"censored={int(is_cen.sum())}  "
        f"is_non_ego={int(batch.get('is_non_ego', torch.zeros(4)).sum())}"
    )

    # Check images shape
    imgs = batch["images"]
    check(len(imgs) == 4, f"images list length == batch_size ({len(imgs)})")
    check(len(imgs[0]) > 0, f"Each sample has ≥1 frame ({len(imgs[0])} frames)")
    logger.info(f"    Image size: {imgs[0][0].size}")

    return train_ds, val_ds


# ── 3. Model checks ───────────────────────────────────────────────────────────

def check_model(model_name: str, pretrained_lora: str | None):
    logger.info("\n=== 3. Model Checks ===")
    from .trainer import SFTModel

    logger.info("  Loading model (this may take a while)...")
    model = SFTModel(
        model_name=model_name,
        pretrained_lora_path=pretrained_lora,
        belief_strategy="mean_pool",
        use_lora=True,
        use_bf16=True,
        device="auto",
    )

    # HazardHead init: sigmoid(-1) ≈ 0.269
    dummy = torch.zeros(1, model.hidden_dim, device=model.device, dtype=model.dtype)
    with torch.no_grad():
        h_logit = model.hazard_head(dummy).float().item()
        h_prob  = torch.sigmoid(torch.tensor(h_logit)).item()
        t_mean, _ = model.tta_head(dummy)
        t_mean = t_mean.float().item()

    check(0.20 <= h_prob <= 0.35,
          f"HazardHead init: hazard_logit={h_logit:.3f}  hazard_prob={h_prob:.3f} (expected ≈0.27, range [0.20,0.35])")
    check(3.0 <= t_mean <= 8.0,
          f"TTAHead init: tta_mean={t_mean:.3f} (expected ≈5.0, range [3,8])")

    logger.info(f"    Hidden dim: {model.hidden_dim}")
    logger.info(f"    Device: {model.device}")
    logger.info(f"    Dtype: {model.dtype}")

    # Trainable params
    lora_params = [(n, p) for n, p in model.vlm.named_parameters() if p.requires_grad and "lora_" in n.lower()]
    head_params = list(model.hazard_head.parameters()) + list(model.tta_head.parameters()) + list(model.belief_aggregator.parameters())
    check(len(lora_params) > 0, f"LoRA has trainable params ({len(lora_params)} tensors)")
    check(len(head_params) > 0, f"Head params exist ({len(head_params)} tensors)")

    return model


# ── 4 & 5. Forward/backward checks ───────────────────────────────────────────

def check_forward_backward(model, train_ds, val_ds):
    logger.info("\n=== 4. Forward Pass Check ===")
    from .dataset import sft_collate_fn
    from .trainer import compute_sft_loss, SFTTrainer

    loader = DataLoader(train_ds, batch_size=2, shuffle=False, collate_fn=sft_collate_fn, num_workers=0)
    batch  = next(iter(loader))

    proc = model.processor
    apply_chat = (
        proc.apply_chat_template
        if hasattr(proc, "apply_chat_template")
        else proc.tokenizer.apply_chat_template
    )
    SYSTEM = "You are a driving safety AI analyzing dashcam footage for collision risk."
    images = batch["images"]
    texts = []
    for i in range(len(batch["video_ids"])):
        frames  = images[i]
        content = [{"type": "image"} for _ in range(len(frames))]
        content.append({"type": "text", "text": "Estimate time to collision. Output a single number."})
        msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": content}]
        texts.append(apply_chat(msgs, tokenize=False, add_generation_prompt=False))

    inputs = proc(text=texts, images=images, return_tensors="pt", padding=True, truncation=True)

    dev = model.device
    t = {k: batch[k].to(dev) for k in ["tta_labels", "hazard_labels", "hazard_weights", "is_ego_positive", "is_censored"]}

    model.train()
    from torch.amp import autocast
    with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
        out = model(inputs)

    check("hazard_logit" in out, "Output has hazard_logit")
    check("hazard_prob"  in out, "Output has hazard_prob")
    check("tta_mean"     in out, "Output has tta_mean")
    check("tta_logvar"   in out, "Output has tta_logvar")
    check(out["hazard_logit"].shape == (2,), f"hazard_logit shape == (2,): {tuple(out['hazard_logit'].shape)}")
    check(out["tta_mean"].shape     == (2,), f"tta_mean shape == (2,): {tuple(out['tta_mean'].shape)}")

    hp = out["hazard_prob"].float()
    check((hp >= 0).all() and (hp <= 1).all(),
          f"hazard_prob in [0,1]: [{hp.min().item():.3f}, {hp.max().item():.3f}]")
    check((out["tta_mean"] > 0).all(),
          f"tta_mean > 0: [{out['tta_mean'].min().item():.3f}, {out['tta_mean'].max().item():.3f}]")

    with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
        loss, metrics = compute_sft_loss(
            hazard_logit=out["hazard_logit"],
            tta_mean=out["tta_mean"],
            tta_logvar=out["tta_logvar"],
            hazard_label=t["hazard_labels"],
            hazard_weight=t["hazard_weights"],
            is_ego_positive=t["is_ego_positive"],
            is_censored=t["is_censored"],
            tta_label=t["tta_labels"],
            nll_weight=0.5,
        )

    check(torch.isfinite(loss), f"Loss is finite: {loss.item():.4f}")
    check(loss.item() > 0,     f"Loss > 0: {loss.item():.4f}")
    check("loss_hazard" in metrics, "metrics has loss_hazard")
    logger.info(f"    loss={loss.item():.4f}  loss_hazard={metrics['loss_hazard']:.4f}  "
                f"loss_tta_obs={metrics['loss_tta_obs']:.4f}  hazard_acc={metrics['hazard_acc']:.3f}")

    logger.info("\n=== 5. Backward Pass Check ===")
    loss.backward()

    # Check hazard_head gets gradient
    hh_grad = model.hazard_head.fc.weight.grad
    check(hh_grad is not None,               "hazard_head.fc.weight has gradient")
    check(hh_grad is not None and hh_grad.abs().sum() > 0, "hazard_head gradient is non-zero")

    # Check tta_head gets gradient
    tta_last = model.tta_head.net[-1]
    th_grad  = tta_last.weight.grad
    if t["is_ego_positive"].any():
        check(th_grad is not None,               "tta_head last layer has gradient")
        check(th_grad is not None and th_grad.abs().sum() > 0, "tta_head gradient is non-zero")
    else:
        logger.info(f"    {WARN} No ego_positive in batch; tta_head gradient may be zero (OK)")

    # Check LoRA params have gradients (may be zero at init if hazard_head.weight=zeros)
    lora_total    = [(n, p) for n, p in model.vlm.named_parameters() if p.requires_grad and "lora_" in n.lower()]
    lora_with_grad = [(n, p) for n, p in lora_total if p.grad is not None and p.grad.abs().sum() > 0]
    if len(lora_with_grad) > 0:
        check(True, f"LoRA params have non-zero gradient ({len(lora_with_grad)} tensors)")
    else:
        # At init, HazardHead.fc.weight==0 → d(loss)/d(belief)=0 → LoRA grad=0.
        # This is expected; gradient will flow once weights are non-zero.
        logger.info(
            f"    {WARN} LoRA gradient is zero at init (HazardHead.fc.weight=0 → no VLM grad path from hazard loss alone). "
            f"Expected: LoRA grads appear after first optimizer step. "
            f"LoRA params with requires_grad: {len(lora_total)}"
        )
        check(len(lora_total) > 0, f"LoRA params exist ({len(lora_total)} tensors with requires_grad)")


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

def main():
    parser = argparse.ArgumentParser("SFT sanity check")
    parser.add_argument("--manifest_dir",   type=str, default="PROJECT_ROOT/data/sft_manifests")
    parser.add_argument("--model_name",     type=str, default="PROJECT_ROOT/models/Qwen2.5-VL-3B-Instruct")
    parser.add_argument("--pretrained_lora", type=str, default="PROJECT_ROOT/checkpoints/pretrain_v2/stage_b/best_model")
    parser.add_argument("--n_samples",      type=int, default=200, help="Debug samples for dataset check")
    parser.add_argument("--skip_model",     action="store_true", help="Skip model load (much faster)")
    args = parser.parse_args()

    logger.info("=" * 60)
    logger.info("LKAlert SFT Sanity Check")
    logger.info("=" * 60)

    check_manifests(Path(args.manifest_dir))
    train_ds, val_ds = check_datasets(Path(args.manifest_dir), n_samples=args.n_samples)

    if args.skip_model:
        logger.info("\n⚠️  Skipping model checks (--skip_model)")
    else:
        model = check_model(args.model_name, args.pretrained_lora)
        check_forward_backward(model, train_ds, val_ds)

    logger.info("\n" + "=" * 60)
    logger.info("✅ All sanity checks passed!")
    logger.info("=" * 60)


if __name__ == "__main__":
    main()