#!/usr/bin/env python3
"""
Pretrain V2 Evaluation Script
==============================
Compares three model states side-by-side:
  (1) base  : Qwen2.5-VL-3B-Instruct, no LoRA
  (2) stage_a: base + Stage-A LoRA (BDD100K + DAD, risk vocabulary)
  (3) stage_b: base + Stage-B LoRA (DADA + NEXAR + DAD, TTA estimation)

Stage-A metrics (on stage_a_val.json):
  - risk_format_rate : fraction of outputs containing "Risk: X/5"
  - risk_accuracy    : predicted risk level == ground-truth risk level
  - anti_bias_rate   : fraction of safe inputs correctly predicted as Risk 1-2/5
                       (tests for the old "always safe" anti-crash bias)

Stage-B metrics (on stage_b_val.json):
  - tta_format_rate  : fraction of outputs containing "TTA: X.Xs"
  - tta_mae          : mean absolute error of extracted TTA vs ground-truth TTA
  - risk_accuracy    : predicted risk level == ground-truth risk level
  - tta_mae_by_risk  : TTA MAE broken down by ground-truth risk level

Usage:
    cd PROJECT_ROOT
    conda activate lkalert

    # Evaluate all three models on both stages (default, takes ~30 min)
    python training/pretrain_v2/evaluate.py

    # Evaluate only on Stage-A (faster)
    python training/pretrain_v2/evaluate.py --stage a

    # Evaluate only on Stage-B
    python training/pretrain_v2/evaluate.py --stage b

    # Limit samples for a quick sanity check
    python training/pretrain_v2/evaluate.py --n_samples 50

Output:
    eval_results/pretrain_v2/eval_report.md   (human-readable table)
    eval_results/pretrain_v2/eval_raw.json    (raw per-sample predictions)
"""

import argparse
import json
import re
import sys
import random
from collections import defaultdict
from pathlib import Path
from datetime import datetime

import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForVision2Seq
from peft import PeftModel
from tqdm import tqdm

# ── paths (mirrors config.py) ─────────────────────────────────────────────────
MODEL_PATH      = "PROJECT_ROOT/models/Qwen2.5-VL-3B-Instruct"
STAGE_A_CKPT    = "PROJECT_ROOT/checkpoints/pretrain_v2/stage_a/best_model"
STAGE_B_CKPT    = "PROJECT_ROOT/checkpoints/pretrain_v2/stage_b/best_model"
STAGE_A_VAL_JSON = "PROJECT_ROOT/data/pretrain_v2/stage_a_val.json"
STAGE_B_VAL_JSON = "PROJECT_ROOT/data/pretrain_v2/stage_b_val.json"
OUTPUT_DIR      = "PROJECT_ROOT/eval_results/pretrain_v2"

MAX_NEW_TOKENS  = 80
SEED            = 42


# ── regex helpers ─────────────────────────────────────────────────────────────

def extract_risk(text: str):
    """Extract integer risk level from 'Risk: X/5' pattern. Returns None if absent."""
    m = re.search(r"[Rr]isk[:\s]+(\d)/5", text)
    return int(m.group(1)) if m else None


def extract_tta(text: str):
    """Extract float TTA from 'TTA: X.Xs' or 'TTA: Xs' pattern. Returns None if absent."""
    m = re.search(r"TTA[:\s]+([\d.]+)\s*s", text, re.IGNORECASE)
    return float(m.group(1)) if m else None


def extract_gt_risk(label: str):
    """Extract ground-truth risk from label string."""
    return extract_risk(label)


def extract_gt_tta(label: str):
    """Extract ground-truth TTA from label string."""
    return extract_tta(label)


# ── model loader ──────────────────────────────────────────────────────────────

def load_model(model_name: str, device: torch.device):
    """
    Load model + processor for one of: 'base', 'stage_a', 'stage_b'.
    Returns (model, processor, processor_seq).
    """
    print(f"\n{'='*55}")
    print(f"Loading model: {model_name}")

    processor = AutoProcessor.from_pretrained(
        MODEL_PATH, trust_remote_code=True,
        min_pixels=4 * 28 * 28,
        max_pixels=768 * 28 * 28,
    )
    processor_seq = AutoProcessor.from_pretrained(
        MODEL_PATH, trust_remote_code=True,
        min_pixels=4 * 28 * 28,
        max_pixels=128 * 28 * 28,
    )
    for proc in (processor, processor_seq):
        if proc.tokenizer.pad_token is None:
            proc.tokenizer.pad_token    = proc.tokenizer.eos_token
            proc.tokenizer.pad_token_id = proc.tokenizer.eos_token_id

    base = AutoModelForVision2Seq.from_pretrained(
        MODEL_PATH,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    )
    base.config.use_cache = True

    if model_name == "base":
        model = base
    elif model_name == "stage_a":
        model = PeftModel.from_pretrained(base, STAGE_A_CKPT, is_trainable=False)
        model = model.merge_and_unload()
    elif model_name == "stage_b":
        # Stage-B LoRA was trained on top of Stage-A LoRA weights
        model = PeftModel.from_pretrained(base, STAGE_B_CKPT, is_trainable=False)
        model = model.merge_and_unload()
    else:
        raise ValueError(f"Unknown model_name: {model_name}")

    model.eval()
    model.to(device)
    print(f"  {model_name} ready on {device}")
    return model, processor, processor_seq


# ── single inference ──────────────────────────────────────────────────────────

@torch.no_grad()
def run_inference(model, processor, processor_seq, frames, prompt, device):
    """
    Run one forward pass with greedy decoding.
    frames: List[PIL.Image]
    Returns generated text (excluding prompt).
    """
    content = [{"type": "image", "image": f} for f in frames]
    content.append({"type": "text", "text": prompt})
    messages = [{"role": "user", "content": content}]

    prompt_text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    is_seq = len(frames) > 1
    proc = processor_seq if is_seq else processor

    enc = proc(
        text=[prompt_text],
        images=[frames],
        return_tensors="pt",
        padding=True,
    )
    enc = {k: v.to(device) if torch.is_tensor(v) else v for k, v in enc.items()}

    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        out_ids = model.generate(
            **enc,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=proc.tokenizer.pad_token_id,
        )

    # Decode only the generated part
    input_len = enc["input_ids"].shape[1]
    gen_ids = out_ids[0, input_len:]
    return proc.tokenizer.decode(gen_ids, skip_special_tokens=True).strip()


# ── dataset loading ────────────────────────────────────────────────────────────

def load_val_samples(json_path: str, n_samples: int, seed: int = SEED):
    """Load and subsample val set. Returns list of dicts."""
    data = json.loads(Path(json_path).read_text(encoding="utf-8"))
    rng = random.Random(seed)
    rng.shuffle(data)
    samples = data[:n_samples]
    print(f"  Loaded {len(samples)} / {len(data)} samples from {Path(json_path).name}")
    return samples


def load_frames(sample: dict):
    """Load PIL images from a sample dict (stage A or B format)."""
    if "frame_paths" in sample:
        frames = []
        for fp in sample["frame_paths"]:
            try:
                frames.append(Image.open(fp).convert("RGB"))
            except Exception:
                pass
        return frames or [Image.new("RGB", (224, 224), (128, 128, 128))]
    else:
        try:
            return [Image.open(sample["image_path"]).convert("RGB")]
        except Exception:
            return [Image.new("RGB", (224, 224), (128, 128, 128))]


# ── metric computation ─────────────────────────────────────────────────────────

def compute_stage_a_metrics(results):
    """
    results: list of {gt_label, prediction, task}
    Returns metric dict.
    """
    total = len(results)
    risk_fmt = sum(1 for r in results if extract_risk(r["prediction"]) is not None)
    risk_correct = 0
    anti_bias_total = 0
    anti_bias_correct = 0

    for r in results:
        gt_risk  = extract_gt_risk(r["gt_label"])
        pr_risk  = extract_risk(r["prediction"])
        if gt_risk is not None and pr_risk is not None:
            if gt_risk == pr_risk:
                risk_correct += 1
        # Anti-bias test: BDD100K/DAD negative samples should be Risk 1-2
        if gt_risk is not None and gt_risk <= 2:
            anti_bias_total += 1
            if pr_risk is not None and pr_risk <= 2:
                anti_bias_correct += 1

    return {
        "n": total,
        "risk_format_rate":  risk_fmt / total if total else 0,
        "risk_accuracy":     risk_correct / total if total else 0,
        "anti_bias_rate":    anti_bias_correct / anti_bias_total if anti_bias_total else 0,
        "anti_bias_total":   anti_bias_total,
    }


def compute_stage_b_metrics(results):
    """
    results: list of {gt_label, prediction, task}
    Returns metric dict.
    """
    total = len(results)
    tta_fmt = sum(1 for r in results if extract_tta(r["prediction"]) is not None)
    risk_correct = 0

    tta_errors = []
    tta_errors_by_risk = defaultdict(list)

    for r in results:
        gt_tta   = extract_gt_tta(r["gt_label"])
        pr_tta   = extract_tta(r["prediction"])
        gt_risk  = extract_gt_risk(r["gt_label"])
        pr_risk  = extract_risk(r["prediction"])

        if gt_risk is not None and pr_risk is not None and gt_risk == pr_risk:
            risk_correct += 1

        if gt_tta is not None and pr_tta is not None:
            err = abs(gt_tta - pr_tta)
            tta_errors.append(err)
            if gt_risk is not None:
                tta_errors_by_risk[gt_risk].append(err)

    tta_mae = sum(tta_errors) / len(tta_errors) if tta_errors else float("nan")
    tta_mae_by_risk = {
        f"Risk{k}": round(sum(v) / len(v), 3)
        for k, v in sorted(tta_errors_by_risk.items())
    }

    return {
        "n": total,
        "tta_format_rate":  tta_fmt / total if total else 0,
        "tta_mae":          round(tta_mae, 3) if tta_errors else "n/a (no TTA parsed)",
        "tta_mae_n":        len(tta_errors),
        "risk_accuracy":    risk_correct / total if total else 0,
        "tta_mae_by_risk":  tta_mae_by_risk,
    }


# ── evaluation loop ────────────────────────────────────────────────────────────

def evaluate_model_on_stage(
    model_name: str,
    model, processor, processor_seq,
    samples: list,
    device: torch.device,
    stage: str,
):
    """Run inference on all samples. Returns list of result dicts."""
    results = []
    for s in tqdm(samples, desc=f"  [{model_name}] Stage-{stage.upper()}"):
        frames = load_frames(s)
        try:
            pred = run_inference(model, processor, processor_seq, frames, s["prompt"], device)
        except Exception as e:
            pred = f"[ERROR: {e}]"
        results.append({
            "model":      model_name,
            "task":       s.get("task", ""),
            "gt_label":   s["label"],
            "prediction": pred,
            "prompt":     s["prompt"],
        })
    return results


# ── report generation ──────────────────────────────────────────────────────────

def format_pct(v):
    if isinstance(v, float):
        return f"{v*100:.1f}%"
    return str(v)


def build_report(stage_a_metrics: dict, stage_b_metrics: dict):
    lines = []
    lines.append("# Pretrain V2 Evaluation Report")
    lines.append(f"> Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    lines.append("")

    # ── Stage-A table ────────────────────────────────────────────────────────
    lines.append("## Stage-A: Risk Vocabulary (BDD100K + DAD)")
    lines.append("")
    lines.append("| Metric | base (no LoRA) | stage_a | stage_b |")
    lines.append("|--------|---------------|---------|---------|")

    def row(name, key, fmt=format_pct):
        vals = [
            fmt(stage_a_metrics.get(m, {}).get(key, "—"))
            for m in ("base", "stage_a", "stage_b")
        ]
        lines.append(f"| {name} | {' | '.join(vals)} |")

    row("Risk format rate",  "risk_format_rate")
    row("Risk accuracy",     "risk_accuracy")
    row("Anti-bias rate ↑",  "anti_bias_rate",
        fmt=lambda v: format_pct(v) if isinstance(v, float) else str(v))
    lines.append("")
    lines.append("> **Anti-bias rate**: fraction of safe inputs (Risk ≤ 2/5) correctly predicted as Risk 1-2/5.")
    lines.append("> Old pretrain failure mode: base model predicts Risk 1 for everything → anti-bias=100% but risk_accuracy=low.")
    lines.append("")

    # ── Stage-B table ────────────────────────────────────────────────────────
    lines.append("## Stage-B: TTA Estimation (DADA + NEXAR + DAD)")
    lines.append("")
    lines.append("| Metric | base (no LoRA) | stage_a | stage_b |")
    lines.append("|--------|---------------|---------|---------|")

    def row_b(name, key, fmt=format_pct):
        vals = []
        for m in ("base", "stage_a", "stage_b"):
            v = stage_b_metrics.get(m, {}).get(key, "—")
            vals.append(fmt(v) if isinstance(v, (int, float)) else str(v))
        lines.append(f"| {name} | {' | '.join(vals)} |")

    row_b("TTA format rate ↑", "tta_format_rate")
    row_b("TTA MAE (s) ↓",    "tta_mae",    fmt=lambda v: f"{v:.3f}s" if isinstance(v, float) else str(v))
    row_b("Risk accuracy ↑",  "risk_accuracy")
    lines.append("")

    # TTA MAE by risk breakdown for stage_b
    if "stage_b" in stage_b_metrics:
        br = stage_b_metrics["stage_b"].get("tta_mae_by_risk", {})
        if br:
            lines.append("### Stage-B TTA MAE by Risk Level (stage_b model)")
            lines.append("")
            lines.append("| Risk Level | TTA MAE (s) |")
            lines.append("|------------|-------------|")
            for k, v in br.items():
                lines.append(f"| {k} | {v:.3f}s |")
            lines.append("")

    # ── Interpretation ────────────────────────────────────────────────────────
    lines.append("## Interpretation")
    lines.append("")
    lines.append("| Signal | What it proves |")
    lines.append("|--------|----------------|")
    lines.append("| stage_a risk_format_rate ↑ vs base | Model has acquired `Risk: X/5` vocabulary |")
    lines.append("| stage_a anti_bias_rate reasonable | No collapse to always-safe prediction |")
    lines.append("| stage_b tta_format_rate ↑↑↑ vs base | TTA regression format learned from pretrain |")
    lines.append("| stage_b tta_mae < 2.0s | TTA estimation is meaningful, not random |")
    lines.append("| SFT starts from stage_b → faster convergence than from base | Main benefit for CVPR paper |")

    return "\n".join(lines)


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

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--stage",    choices=["a", "b", "both"], default="both",
                        help="Which stage to evaluate (default: both)")
    parser.add_argument("--models",   nargs="+", default=["base", "stage_a", "stage_b"],
                        choices=["base", "stage_a", "stage_b"],
                        help="Which models to evaluate")
    parser.add_argument("--n_samples", type=int, default=200,
                        help="Samples per stage per model (default: 200)")
    args = parser.parse_args()

    rng = random.Random(SEED)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    out_dir = Path(OUTPUT_DIR)
    out_dir.mkdir(parents=True, exist_ok=True)

    do_a = args.stage in ("a", "both")
    do_b = args.stage in ("b", "both")

    # Pre-load and shuffle val sets once so all models see identical samples
    a_samples = load_val_samples(STAGE_A_VAL_JSON, args.n_samples) if do_a else []
    b_samples = load_val_samples(STAGE_B_VAL_JSON, args.n_samples) if do_b else []

    stage_a_metrics = {}
    stage_b_metrics = {}
    all_results     = []

    for model_name in args.models:
        # Check checkpoint availability
        if model_name == "stage_a" and not Path(STAGE_A_CKPT).exists():
            print(f"⚠  Stage-A checkpoint not found at {STAGE_A_CKPT}, skipping.")
            continue
        if model_name == "stage_b" and not Path(STAGE_B_CKPT).exists():
            print(f"⚠  Stage-B checkpoint not found at {STAGE_B_CKPT}, skipping.")
            continue

        model, processor, processor_seq = load_model(model_name, device)

        if do_a:
            res_a = evaluate_model_on_stage(
                model_name, model, processor, processor_seq, a_samples, device, "a"
            )
            stage_a_metrics[model_name] = compute_stage_a_metrics(res_a)
            all_results.extend(res_a)
            print(f"  Stage-A metrics ({model_name}): {stage_a_metrics[model_name]}")

        if do_b:
            res_b = evaluate_model_on_stage(
                model_name, model, processor, processor_seq, b_samples, device, "b"
            )
            stage_b_metrics[model_name] = compute_stage_b_metrics(res_b)
            all_results.extend(res_b)
            print(f"  Stage-B metrics ({model_name}): {stage_b_metrics[model_name]}")

        # Free GPU memory before loading next model
        del model
        torch.cuda.empty_cache()

    # ── Save outputs ─────────────────────────────────────────────────────────
    raw_path = out_dir / "eval_raw.json"
    raw_path.write_text(
        json.dumps({
            "stage_a_metrics": stage_a_metrics,
            "stage_b_metrics": stage_b_metrics,
            "samples":         all_results,
        }, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"\n✓ Raw results saved → {raw_path}")

    report = build_report(stage_a_metrics, stage_b_metrics)
    report_path = out_dir / "eval_report.md"
    report_path.write_text(report, encoding="utf-8")
    print(f"✓ Report saved → {report_path}")

    # Print report to stdout
    print("\n" + "="*65)
    print(report)
    print("="*65)


if __name__ == "__main__":
    main()