#!/usr/bin/env python3
"""Batch-1 inference latency benchmark for revised RAVEL checkpoints."""

from __future__ import annotations

import argparse
import json
import os
import platform
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from PIL import Image
from transformers import DebertaV2Tokenizer
from transformers.utils import logging as hf_logging

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

from run_revised_experiments import (  # noqa: E402
    METHODS,
    apply_method_trainability,
    freeze_non_lora,
    load_dataset,
    logits_for_target,
    set_seed,
)

hf_logging.set_verbosity_error()

DATASETS = ["mvsa_multiple", "hfm_deleak"]
MODELS = ["lora_concat", "legacy_global", "token_coattn", "token_aux", "param_mlp", "full_revised"]
CLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
CLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run E10 inference benchmark.")
    parser.add_argument("--output-root", default="ravel_revision_results")
    parser.add_argument("--datasets", nargs="+", default=DATASETS, choices=DATASETS)
    parser.add_argument("--models", nargs="+", default=MODELS, choices=MODELS)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--precision", choices=["fp32", "fp16", "bf16"], default="fp16")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--warmup-iterations", type=int, default=50)
    parser.add_argument("--measurement-iterations", type=int, default=500)
    parser.add_argument("--max-length", type=int, default=None)
    parser.add_argument("--hfm-deleak-manifest", default="ravel_revision_results/data_audit/hfm_split_manifest_deleaked.csv")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def image_to_tensor(image_path: str) -> torch.Tensor:
    try:
        image = Image.open(image_path).convert("RGB").resize((224, 224))
    except Exception:
        image = Image.new("RGB", (224, 224), (0, 0, 0))
    arr = np.asarray(image, dtype=np.float32) / 255.0
    arr = (arr - CLIP_MEAN) / CLIP_STD
    return torch.from_numpy(np.transpose(arr, (2, 0, 1))).float()


def sample_label(dataset: str, sample: Any) -> int:
    if dataset.startswith("mvsa"):
        return {"positive": 0, "neutral": 1, "negative": 2}[str(sample.combined_majority).lower()]
    return int(sample.label)


def sample_text(sample: Any) -> str:
    return str(getattr(sample, "text", ""))


def make_batch(sample: Any, dataset: str, tokenizer: DebertaV2Tokenizer, max_length: int, device: torch.device) -> Dict[str, torch.Tensor]:
    text_inputs = tokenizer(
        [sample_text(sample)],
        max_length=max_length,
        padding=True,
        truncation=True,
        return_tensors="pt",
    )
    batch = {
        "pixel_values": image_to_tensor(sample.image_path).unsqueeze(0),
        "input_ids": text_inputs["input_ids"],
        "attention_mask": text_inputs["attention_mask"],
        "labels": torch.tensor([sample_label(dataset, sample)], dtype=torch.long),
    }
    return {key: value.to(device, non_blocking=True) for key, value in batch.items()}


def cuda_sync(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def autocast_context(device: torch.device, precision: str):
    if device.type != "cuda" or precision == "fp32":
        return torch.autocast(device_type="cpu", enabled=False)
    dtype = torch.float16 if precision == "fp16" else torch.bfloat16
    return torch.autocast(device_type="cuda", dtype=dtype)


def run_forward(model: nn.Module, method_key: str, batch: Dict[str, torch.Tensor], precision: str, device: torch.device) -> None:
    criterion = nn.CrossEntropyLoss()
    labels = batch["labels"].long()
    with torch.no_grad(), autocast_context(device, precision):
        logits_for_target(model, batch, METHODS[method_key].train_target, criterion, labels)


def load_model_and_sample(args: argparse.Namespace, dataset: str, method_key: str) -> Tuple[nn.Module, Any, int, torch.device, Dict[str, Any]]:
    root = Path(args.output_root)
    run_dir = root / "runs" / dataset / method_key / f"seed_{args.seed}"
    ckpt_path = run_dir / "checkpoint.pt"
    if not ckpt_path.exists():
        raise FileNotFoundError(ckpt_path)
    start = time.perf_counter()
    ckpt = torch.load(ckpt_path, map_location="cpu")
    ckpt_cfg = ckpt.get("cfg", {}) if isinstance(ckpt, dict) else {}
    max_length = int(args.max_length or ckpt_cfg.get("max_length") or 96)
    method = METHODS[method_key]
    set_seed(args.seed)
    device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
    (
        model_cls,
        cfg,
        _train_loader,
        _val_loader,
        _test_loader,
        _train_samples,
        _val_samples,
        test_samples,
        _num_classes,
        _label_names,
    ) = load_dataset(
        dataset_key=dataset,
        seed=args.seed,
        batch_size=1,
        max_length=max_length,
        num_workers=0,
        method=method,
        hfm_deleak_manifest=args.hfm_deleak_manifest,
        limits=(None, None, None),
    )
    cfg.update(ckpt_cfg)
    cfg.update(
        {
            "architecture": method.architecture,
            "enable_clip_lora": method.enable_lora,
            "enable_text_lora": method.enable_lora,
            "seed": args.seed,
            "batch_size": 1,
            "max_length": max_length,
        }
    )
    model = model_cls(cfg).to(device)
    if hasattr(model, "vision_lora"):
        freeze_non_lora(model.vision_lora)
    if hasattr(model, "text"):
        freeze_non_lora(model.text)
    apply_method_trainability(model, method)
    model.load_state_dict(ckpt.get("model_state", {}), strict=False)
    model.eval()
    load_time = time.perf_counter() - start
    stats = model.parameter_stats() if hasattr(model, "parameter_stats") else {
        "total": sum(p.numel() for p in model.parameters()),
        "trainable": sum(p.numel() for p in model.parameters() if p.requires_grad),
    }
    meta = {
        "checkpoint_size_mb": ckpt_path.stat().st_size / (1024 * 1024),
        "model_load_time_seconds": load_time,
        "total_parameters": int(stats.get("total", 0)),
        "trainable_parameters": int(stats.get("trainable", 0)),
        "max_length": max_length,
    }
    return model, test_samples[0], max_length, device, meta


def latency_stats(values_ms: Sequence[float]) -> Dict[str, float]:
    arr = np.array(values_ms, dtype=float)
    return {
        "mean": float(arr.mean()),
        "std": float(arr.std(ddof=1)) if arr.size > 1 else 0.0,
        "p50": float(np.percentile(arr, 50)),
        "p95": float(np.percentile(arr, 95)),
    }


def benchmark_one(args: argparse.Namespace, dataset: str, method_key: str) -> Dict[str, Any]:
    model, sample, max_length, device, meta = load_model_and_sample(args, dataset, method_key)
    tokenizer = DebertaV2Tokenizer.from_pretrained("microsoft/deberta-v3-base")
    batch = make_batch(sample, dataset, tokenizer, max_length, device)
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats(device)
    for _ in range(args.warmup_iterations):
        run_forward(model, method_key, batch, args.precision, device)
    cuda_sync(device)

    model_only: List[float] = []
    for _ in range(args.measurement_iterations):
        cuda_sync(device)
        start = time.perf_counter()
        run_forward(model, method_key, batch, args.precision, device)
        cuda_sync(device)
        model_only.append((time.perf_counter() - start) * 1000.0)

    end_to_end: List[float] = []
    for _ in range(args.measurement_iterations):
        cuda_sync(device)
        start = time.perf_counter()
        fresh_batch = make_batch(sample, dataset, tokenizer, max_length, device)
        run_forward(model, method_key, fresh_batch, args.precision, device)
        cuda_sync(device)
        end_to_end.append((time.perf_counter() - start) * 1000.0)

    model_stats = latency_stats(model_only)
    e2e_stats = latency_stats(end_to_end)
    peak_vram = torch.cuda.max_memory_allocated(device) / (1024 * 1024) if device.type == "cuda" else float("nan")
    del model
    if device.type == "cuda":
        torch.cuda.empty_cache()
    return {
        "model": method_key,
        "dataset": dataset,
        "device": torch.cuda.get_device_name(device) if device.type == "cuda" else platform.processor(),
        "precision": args.precision.upper(),
        "batch_size": args.batch_size,
        **meta,
        "model_only_latency_mean_ms": model_stats["mean"],
        "model_only_latency_std_ms": model_stats["std"],
        "end_to_end_latency_mean_ms": e2e_stats["mean"],
        "end_to_end_latency_std_ms": e2e_stats["std"],
        "latency_mean_ms": e2e_stats["mean"],
        "latency_std_ms": e2e_stats["std"],
        "latency_p50_ms": e2e_stats["p50"],
        "latency_p95_ms": e2e_stats["p95"],
        "throughput_samples_per_second": 1000.0 / max(e2e_stats["mean"], 1e-9),
        "peak_inference_vram_mb": peak_vram,
        "warmup_iterations": args.warmup_iterations,
        "measurement_iterations": args.measurement_iterations,
        "preprocessing_included": "end_to_end includes PIL image decode, CLIP normalization, and tokenization for one sample; model_only excludes preprocessing",
    }


def main() -> None:
    args = parse_args()
    out_path = Path(args.output_root) / "aggregate_results" / "efficiency_results.csv"
    out_path.parent.mkdir(parents=True, exist_ok=True)
    existing = pd.read_csv(out_path) if out_path.exists() and out_path.stat().st_size else pd.DataFrame()
    rows: List[Dict[str, Any]] = [] if args.overwrite or existing.empty else existing.to_dict("records")
    done = {(str(row.get("model")), str(row.get("dataset")), str(row.get("precision"))) for row in rows if row.get("model_only_latency_mean_ms") == row.get("model_only_latency_mean_ms")}
    for dataset in args.datasets:
        for model in args.models:
            key = (model, dataset, args.precision.upper())
            if key in done and not args.overwrite:
                print(f"SKIP E10 {dataset} {model} {args.precision}", flush=True)
                continue
            print(f"RUN E10 dataset={dataset} model={model} precision={args.precision}", flush=True)
            row = benchmark_one(args, dataset, model)
            rows = [r for r in rows if not (str(r.get("model")) == model and str(r.get("dataset")) == dataset and str(r.get("precision")) == args.precision.upper())]
            rows.append(row)
            pd.DataFrame(rows).to_csv(out_path, index=False)
            print(
                f"DONE E10 {dataset} {model} p50={row['latency_p50_ms']:.2f}ms "
                f"p95={row['latency_p95_ms']:.2f}ms vram={row['peak_inference_vram_mb']:.1f}MB",
                flush=True,
            )
    pd.DataFrame(rows).to_csv(out_path, index=False)


if __name__ == "__main__":
    main()