#!/usr/bin/env python3
"""Run deterministic open-weight LVLM baselines for the RAVEL revision.

The script writes hard-label LVLM results only. It does not fabricate
calibration metrics when normalized label-token probabilities are unavailable.
"""

from __future__ import annotations

import argparse
import csv
import json
import random
import re
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np
import torch
from PIL import Image
from sklearn.metrics import accuracy_score, f1_score
from transformers import AutoModelForImageTextToText, AutoProcessor

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


ROOT = Path("ravel_revision_results")
OUT = ROOT / "lvlm_results"
AGG = ROOT / "aggregate_results"

MODEL_ALIASES = {
    "lvlm_small": "Qwen/Qwen2.5-VL-3B-Instruct",
    "lvlm_large": "Qwen/Qwen2.5-VL-7B-Instruct",
}

MVSA_LABELS = {0: "Positive", 1: "Neutral", 2: "Negative"}
HFM_LABELS = {0: "Non-hateful", 1: "Hateful"}


@dataclass
class EvalSample:
    sample_id: str
    image_path: str
    text: str
    label: int


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run LVLM baselines for RAVEL revision.")
    parser.add_argument("--models", nargs="+", default=["lvlm_small", "lvlm_large"])
    parser.add_argument("--datasets", nargs="+", default=["mvsa_multiple", "hfm_deleak"])
    parser.add_argument("--settings", nargs="+", default=["zero_shot", "four_shot"])
    parser.add_argument("--split-seed", type=int, default=1)
    parser.add_argument("--max-new-tokens", type=int, default=16)
    parser.add_argument("--dtype", choices=["bf16", "fp16", "fp32"], default="bf16")
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--limit-test-samples", type=int, default=None)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--output-root", default=str(ROOT))
    parser.add_argument(
        "--hfm-deleak-manifest",
        default="ravel_revision_results/data_audit/hfm_split_manifest_deleaked.csv",
    )
    return parser.parse_args()


def write_csv(path: Path, rows: Iterable[Dict[str, Any]], columns: Sequence[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(columns), extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def load_samples(dataset: str, split_seed: int, manifest: str) -> Tuple[List[EvalSample], List[EvalSample]]:
    if dataset == "mvsa_multiple":
        from src.mvsa_multiple_pipeline import DEFAULT_MVSA_MULTIPLE_CONFIG, LABEL_NAME_TO_ID, MVSALoader

        cfg = dict(DEFAULT_MVSA_MULTIPLE_CONFIG)
        loader = MVSALoader(cfg["text_dir"], cfg["label_file"])
        loader.load(
            preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")),
            require_unanimous=bool(cfg.get("require_unanimous", False)),
            require_cross_agree=bool(cfg.get("require_cross_agree", False)),
            paper_exact_counts=True,
        )
        train, _val, test = loader.split(
            train_ratio=float(cfg["train_ratio"]),
            val_ratio=float(cfg["val_ratio"]),
            seed=split_seed,
            paper_811=True,
        )

        def convert(sample: Any) -> EvalSample:
            return EvalSample(
                sample_id=str(sample.image_id),
                image_path=str(sample.image_path),
                text=str(sample.text),
                label=int(LABEL_NAME_TO_ID[sample.combined_majority]),
            )

        return [convert(s) for s in train], [convert(s) for s in test]

    if dataset == "hfm_deleak":
        from src.hfm_pipeline import HFMLoader, DEFAULT_HFM_CONFIG

        cfg = dict(DEFAULT_HFM_CONFIG)
        loader = HFMLoader(cfg["text_dir"], cfg["image_root"])
        loader.load_from_manifest(manifest)

        def convert(sample: Any) -> EvalSample:
            return EvalSample(
                sample_id=str(sample.image_id),
                image_path=str(sample.image_path),
                text=str(sample.text),
                label=int(sample.label),
            )

        return [convert(s) for s in loader.get_split("train")], [convert(s) for s in loader.get_split("test")]

    raise ValueError(f"Unsupported LVLM dataset: {dataset}")


def label_names(dataset: str) -> Dict[int, str]:
    return MVSA_LABELS if dataset.startswith("mvsa") else HFM_LABELS


def task_prompt(dataset: str, text: str) -> str:
    if dataset.startswith("mvsa"):
        return (
            "You are given one image and one accompanying text.\n"
            "Classify the overall sentiment.\n\n"
            f"Text: {text}\n\n"
            "Return exactly one label:\n"
            "Negative\n"
            "Neutral\n"
            "Positive"
        )
    return (
        "You are given one meme image and its accompanying text.\n"
        "Determine whether the combined image-text content is hateful.\n\n"
        f"Text: {text}\n\n"
        "Return exactly one label:\n"
        "Non-hateful\n"
        "Hateful"
    )


def choose_few_shot(train: List[EvalSample], dataset: str, k: int, seed: int) -> List[EvalSample]:
    if k <= 0:
        return []
    rng = random.Random(seed)
    grouped: Dict[int, List[EvalSample]] = {}
    for sample in train:
        grouped.setdefault(sample.label, []).append(sample)
    for values in grouped.values():
        rng.shuffle(values)
    labels = sorted(grouped)
    selected: List[EvalSample] = []
    while len(selected) < k and labels:
        progressed = False
        for label in labels:
            bucket = grouped.get(label, [])
            if bucket and len(selected) < k:
                selected.append(bucket.pop(0))
                progressed = True
        if not progressed:
            break
    return selected[:k]


def build_messages(dataset: str, sample: EvalSample, demos: List[EvalSample]) -> Tuple[List[Dict[str, Any]], List[Image.Image]]:
    images: List[Image.Image] = []
    messages: List[Dict[str, Any]] = []
    labels = label_names(dataset)
    for demo in demos:
        images.append(Image.open(demo.image_path).convert("RGB"))
        messages.append(
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": task_prompt(dataset, demo.text)},
                ],
            }
        )
        messages.append({"role": "assistant", "content": labels[demo.label]})
    images.append(Image.open(sample.image_path).convert("RGB"))
    messages.append(
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": task_prompt(dataset, sample.text)},
            ],
        }
    )
    return messages, images


def parse_label(dataset: str, raw_output: str) -> Tuple[Optional[int], bool]:
    text = raw_output.strip().lower()
    text = re.sub(r"[^a-z\\-\\s]", " ", text)
    text = re.sub(r"\\s+", " ", text).strip()
    if dataset.startswith("mvsa"):
        if re.search(r"\\bnegative\\b", text):
            return 2, True
        if re.search(r"\\bneutral\\b", text):
            return 1, True
        if re.search(r"\\bpositive\\b", text):
            return 0, True
        return None, False
    if "non-hateful" in text or "non hateful" in text or "not hateful" in text:
        return 0, True
    if re.search(r"\\bhateful\\b", text):
        return 1, True
    return None, False


def tensor_device(model: torch.nn.Module) -> torch.device:
    try:
        return next(model.parameters()).device
    except StopIteration:
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def move_inputs(inputs: Dict[str, Any], device: torch.device) -> Dict[str, Any]:
    return {key: value.to(device) if torch.is_tensor(value) else value for key, value in inputs.items()}


def load_model(model_id: str, dtype: str, device_map: str) -> Tuple[Any, Any, int]:
    torch_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[dtype]
    processor = AutoProcessor.from_pretrained(model_id)
    model = AutoModelForImageTextToText.from_pretrained(
        model_id,
        torch_dtype=torch_dtype,
        device_map=device_map,
        low_cpu_mem_usage=True,
    )
    model.eval()
    total_params = int(sum(parameter.numel() for parameter in model.parameters()))
    return processor, model, total_params


@torch.no_grad()
def run_suite(
    processor: Any,
    model: Any,
    total_params: int,
    model_alias: str,
    model_id: str,
    dataset: str,
    setting: str,
    args: argparse.Namespace,
) -> Dict[str, Any]:
    train, test = load_samples(dataset, args.split_seed, args.hfm_deleak_manifest)
    if args.limit_test_samples and args.limit_test_samples > 0:
        test = test[: args.limit_test_samples]
    demos = choose_few_shot(train, dataset, 4 if setting == "four_shot" else 0, args.split_seed)

    suite_dir = Path(args.output_root) / "lvlm_results" / model_alias / dataset / setting
    pred_path = suite_dir / "predictions.csv"
    if pred_path.exists() and not args.overwrite:
        return aggregate_suite(model_alias, model_id, dataset, setting, suite_dir, total_params, args.dtype)

    suite_dir.mkdir(parents=True, exist_ok=True)
    prompt_record = {
        "dataset": dataset,
        "setting": setting,
        "base_prompt": task_prompt(dataset, "{sample_text}"),
        "decoding": {"temperature": 0, "do_sample": False, "max_new_tokens": args.max_new_tokens},
        "probability_source": "hard_label_only",
    }
    (suite_dir / "prompt.txt").write_text(json.dumps(prompt_record, indent=2), encoding="utf-8")
    (suite_dir / "few_shot_examples.json").write_text(
        json.dumps([asdict(sample) for sample in demos], indent=2),
        encoding="utf-8",
    )

    device = tensor_device(model)
    rows: List[Dict[str, Any]] = []
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()
    for sample in test:
        try:
            messages, images = build_messages(dataset, sample, demos)
            text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            inputs = processor(text=[text], images=images, return_tensors="pt")
            prompt_tokens = int(inputs["input_ids"].shape[-1]) if "input_ids" in inputs else 0
            inputs = move_inputs(inputs, device)
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            start = time.perf_counter()
            generated = model.generate(
                **inputs,
                do_sample=False,
                max_new_tokens=args.max_new_tokens,
            )
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            latency_ms = (time.perf_counter() - start) * 1000.0
            input_len = int(inputs["input_ids"].shape[-1])
            generated_ids = generated[:, input_len:]
            raw_output = processor.batch_decode(
                generated_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            )[0].strip()
            generated_tokens = int(generated_ids.shape[-1])
            parsed, ok = parse_label(dataset, raw_output)
        except Exception as exc:
            latency_ms = float("nan")
            prompt_tokens = 0
            generated_tokens = 0
            raw_output = f"ERROR: {type(exc).__name__}: {exc}"
            parsed, ok = None, False
        rows.append(
            {
                "sample_id": sample.sample_id,
                "true_label": sample.label,
                "raw_output": raw_output,
                "parsed_label": "" if parsed is None else int(parsed),
                "parse_success": bool(ok),
                "latency_ms": latency_ms,
                "prompt_tokens": prompt_tokens,
                "generated_tokens": generated_tokens,
            }
        )

    write_csv(
        pred_path,
        rows,
        [
            "sample_id",
            "true_label",
            "raw_output",
            "parsed_label",
            "parse_success",
            "latency_ms",
            "prompt_tokens",
            "generated_tokens",
        ],
    )
    return aggregate_suite(model_alias, model_id, dataset, setting, suite_dir, total_params, args.dtype)


def aggregate_suite(
    model_alias: str,
    model_id: str,
    dataset: str,
    setting: str,
    suite_dir: Path,
    total_params: int,
    dtype: str,
) -> Dict[str, Any]:
    pred_path = suite_dir / "predictions.csv"
    rows: List[Dict[str, str]] = []
    with pred_path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.DictReader(f)
        rows = list(reader)
    true: List[int] = []
    pred: List[int] = []
    latencies: List[float] = []
    failures = 0
    for row in rows:
        ok = str(row.get("parse_success", "")).lower() == "true"
        if not ok or row.get("parsed_label", "") == "":
            failures += 1
            continue
        true.append(int(row["true_label"]))
        pred.append(int(row["parsed_label"]))
        try:
            latencies.append(float(row["latency_ms"]))
        except Exception:
            pass
    if true:
        accuracy = float(accuracy_score(true, pred))
        macro_f1 = float(f1_score(true, pred, average="macro", zero_division=0))
        weighted_f1 = float(f1_score(true, pred, average="weighted", zero_division=0))
    else:
        accuracy = macro_f1 = weighted_f1 = float("nan")
    lat = np.asarray(latencies, dtype=float)
    peak_vram = torch.cuda.max_memory_reserved() / (1024**2) if torch.cuda.is_available() else float("nan")
    efficiency = {
        "total_parameters": total_params,
        "precision": dtype,
        "peak_vram_mb": float(peak_vram),
        "mean_latency_ms": float(np.nanmean(lat)) if lat.size else float("nan"),
        "p50_latency_ms": float(np.nanpercentile(lat, 50)) if lat.size else float("nan"),
        "p95_latency_ms": float(np.nanpercentile(lat, 95)) if lat.size else float("nan"),
    }
    (suite_dir / "efficiency.json").write_text(json.dumps(efficiency, indent=2), encoding="utf-8")
    return {
        "model": model_alias,
        "model_id": model_id,
        "dataset": dataset,
        "setting": setting,
        "total_params": total_params,
        "trainable_params": 0,
        "accuracy": accuracy,
        "macro_f1": macro_f1,
        "weighted_f1": weighted_f1,
        "f1": weighted_f1 if dataset.startswith("mvsa") else macro_f1,
        "auroc": "N/A",
        "raw_ece": "N/A",
        "nll": "N/A",
        "brier": "N/A",
        "latency_per_sample_ms": efficiency["mean_latency_ms"],
        "peak_vram_mb": efficiency["peak_vram_mb"],
        "parse_failure_rate": failures / max(1, len(rows)),
        "probability_source": "hard_label_only",
        "status": "complete" if rows else "missing",
    }


def write_aggregate(rows: List[Dict[str, Any]]) -> None:
    AGG.mkdir(parents=True, exist_ok=True)
    write_csv(
        AGG / "lvlm_results.csv",
        rows,
        [
            "model",
            "model_id",
            "dataset",
            "setting",
            "total_params",
            "trainable_params",
            "accuracy",
            "macro_f1",
            "weighted_f1",
            "f1",
            "auroc",
            "raw_ece",
            "nll",
            "brier",
            "latency_per_sample_ms",
            "peak_vram_mb",
            "parse_failure_rate",
            "probability_source",
            "status",
        ],
    )


def main() -> None:
    args = parse_args()
    global ROOT, OUT, AGG
    ROOT = Path(args.output_root)
    OUT = ROOT / "lvlm_results"
    AGG = ROOT / "aggregate_results"
    if args.dry_run:
        for model_alias in args.models:
            model_id = MODEL_ALIASES.get(model_alias, model_alias)
            for dataset in args.datasets:
                train, test = load_samples(dataset, args.split_seed, args.hfm_deleak_manifest)
                if args.limit_test_samples and args.limit_test_samples > 0:
                    test = test[: args.limit_test_samples]
                for setting in args.settings:
                    demos = choose_few_shot(train, dataset, 4 if setting == "four_shot" else 0, args.split_seed)
                    print(
                        f"DRY-RUN model={model_alias} ({model_id}) dataset={dataset} "
                        f"setting={setting} train={len(train)} test={len(test)} demos={[d.sample_id for d in demos]}",
                        flush=True,
                    )
        return
    all_rows: List[Dict[str, Any]] = []
    for model_alias in args.models:
        model_id = MODEL_ALIASES.get(model_alias, model_alias)
        processor, model, total_params = load_model(model_id, args.dtype, args.device_map)
        for dataset in args.datasets:
            for setting in args.settings:
                row = run_suite(processor, model, total_params, model_alias, model_id, dataset, setting, args)
                all_rows.append(row)
                write_aggregate(all_rows)
        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    write_aggregate(all_rows)
    print(f"Wrote {AGG / 'lvlm_results.csv'}")


if __name__ == "__main__":
    main()