#!/usr/bin/env python3
"""Compare FP32 FLAN-T5 Small with a quantized checkpoint on fixed prompts."""

from __future__ import annotations

import argparse
import difflib
import gc
import json
import platform
import statistics
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import torch
import transformers
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer


DEFAULT_PROMPTS = [
    "translate English to German: How old are you?",
    "Answer this question: What is the capital of France?",
    "Classify the sentiment as positive or negative: I loved the thoughtful story and acting.",
    "summarize: The James Webb Space Telescope observes the universe in infrared light, allowing it to see through dust and study very distant galaxies.",
    "Premise: All roses are flowers. Some flowers fade quickly. Question: Does it follow that some roses fade quickly? Explain briefly.",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-model", default="google/flan-t5-small")
    parser.add_argument("--base-revision", default="main")
    parser.add_argument("--quantized-model", default="ShinpacheShimura/t5-smaller")
    parser.add_argument("--quantized-revision", default="main")
    parser.add_argument("--quantized-subfolder", default=None)
    parser.add_argument(
        "--prompts-file",
        type=Path,
        help="Optional UTF-8 text file with one non-empty prompt per line.",
    )
    parser.add_argument("--output-dir", type=Path, default=Path("benchmark-results"))
    parser.add_argument("--warmup-runs", type=int, default=2)
    parser.add_argument("--runs", type=int, default=10)
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--num-beams", type=int, default=1)
    parser.add_argument("--seed", type=int, default=42)
    return parser.parse_args()


def load_prompts(path: Path | None) -> list[str]:
    if path is None:
        return DEFAULT_PROMPTS
    prompts = [line.strip() for line in path.read_text(encoding="utf-8").splitlines()]
    prompts = [prompt for prompt in prompts if prompt]
    if not prompts:
        raise SystemExit("The prompts file contains no non-empty prompts.")
    return prompts


def synchronize() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def reset_peak_memory() -> None:
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()


def input_device(model: torch.nn.Module) -> torch.device:
    try:
        return model.device
    except AttributeError:
        return next(model.parameters()).device


def generate_once(
    model: torch.nn.Module,
    tokenizer: Any,
    prompt: str,
    max_new_tokens: int,
    num_beams: int,
) -> tuple[str, int, float]:
    encoded = tokenizer(prompt, return_tensors="pt").to(input_device(model))
    synchronize()
    started = time.perf_counter()
    with torch.inference_mode():
        generated = model.generate(
            **encoded,
            max_new_tokens=max_new_tokens,
            num_beams=num_beams,
            do_sample=False,
        )
    synchronize()
    elapsed = time.perf_counter() - started
    text = tokenizer.decode(generated[0], skip_special_tokens=True)
    generated_tokens = int(generated.shape[-1])
    return text, generated_tokens, elapsed


def benchmark_model(
    label: str,
    model_id: str,
    revision: str,
    prompts: list[str],
    warmup_runs: int,
    runs: int,
    max_new_tokens: int,
    num_beams: int,
    quantized: bool,
    subfolder: str | None = None,
) -> dict[str, Any]:
    tokenizer_kwargs: dict[str, Any] = {"revision": revision}
    model_kwargs: dict[str, Any] = {"revision": revision}
    if subfolder:
        tokenizer_kwargs["subfolder"] = subfolder
        model_kwargs["subfolder"] = subfolder

    tokenizer = AutoTokenizer.from_pretrained(model_id, **tokenizer_kwargs)
    if quantized:
        model_kwargs["device_map"] = "auto"
    else:
        model_kwargs["torch_dtype"] = torch.float32
        model_kwargs["device_map"] = "auto"

    reset_peak_memory()
    load_started = time.perf_counter()
    model = AutoModelForSeq2SeqLM.from_pretrained(model_id, **model_kwargs)
    model.eval()
    synchronize()
    load_seconds = time.perf_counter() - load_started

    for index in range(warmup_runs):
        generate_once(
            model,
            tokenizer,
            prompts[index % len(prompts)],
            max_new_tokens,
            num_beams,
        )

    timings: list[float] = []
    total_output_tokens = 0
    outputs: list[dict[str, str]] = []
    for prompt in prompts:
        output, generated_tokens, _ = generate_once(
            model, tokenizer, prompt, max_new_tokens, num_beams
        )
        outputs.append({"prompt": prompt, "output": output})

    for index in range(runs):
        _, generated_tokens, elapsed = generate_once(
            model,
            tokenizer,
            prompts[index % len(prompts)],
            max_new_tokens,
            num_beams,
        )
        timings.append(elapsed)
        total_output_tokens += generated_tokens

    footprint = int(model.get_memory_footprint())
    peak_cuda = int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None
    result = {
        "label": label,
        "model_id": model_id,
        "revision": revision,
        "subfolder": subfolder,
        "load_seconds": load_seconds,
        "model_memory_footprint_bytes": footprint,
        "peak_cuda_allocated_bytes": peak_cuda,
        "latency_seconds": {
            "mean": statistics.mean(timings),
            "median": statistics.median(timings),
            "min": min(timings),
            "max": max(timings),
        },
        "examples_per_second": runs / sum(timings),
        "output_tokens_per_second": total_output_tokens / sum(timings),
        "outputs": outputs,
    }

    del model, tokenizer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return result


def output_agreement(base: dict[str, Any], quantized: dict[str, Any]) -> dict[str, Any]:
    pairs = []
    exact_count = 0
    similarities: list[float] = []
    for base_item, quantized_item in zip(base["outputs"], quantized["outputs"], strict=True):
        base_text = base_item["output"].strip()
        quantized_text = quantized_item["output"].strip()
        exact = base_text == quantized_text
        similarity = difflib.SequenceMatcher(None, base_text, quantized_text).ratio()
        exact_count += int(exact)
        similarities.append(similarity)
        pairs.append(
            {
                "prompt": base_item["prompt"],
                "fp32_output": base_text,
                "quantized_output": quantized_text,
                "exact_match": exact,
                "text_similarity": similarity,
            }
        )
    return {
        "exact_match_rate": exact_count / len(pairs),
        "mean_text_similarity": statistics.mean(similarities),
        "note": "These are regression diagnostics, not ground-truth quality metrics.",
        "pairs": pairs,
    }


def mib(value: int | None) -> str:
    return "N/A" if value is None else f"{value / (1024**2):.2f} MiB"


def render_markdown(results: dict[str, Any]) -> str:
    base = results["models"]["fp32"]
    quantized = results["models"]["quantized"]
    agreement = results["agreement"]
    lines = [
        "# Benchmark results",
        "",
        f"Generated: `{results['created_at_utc']}`",
        "",
        "| Measurement | FP32 base | NF4 checkpoint |",
        "|---|---:|---:|",
        f"| Model memory footprint | {mib(base['model_memory_footprint_bytes'])} | {mib(quantized['model_memory_footprint_bytes'])} |",
        f"| Peak CUDA allocation | {mib(base['peak_cuda_allocated_bytes'])} | {mib(quantized['peak_cuda_allocated_bytes'])} |",
        f"| Load time | {base['load_seconds']:.4f} s | {quantized['load_seconds']:.4f} s |",
        f"| Median generation latency | {base['latency_seconds']['median']:.4f} s | {quantized['latency_seconds']['median']:.4f} s |",
        f"| Examples/second | {base['examples_per_second']:.3f} | {quantized['examples_per_second']:.3f} |",
        f"| Output tokens/second | {base['output_tokens_per_second']:.3f} | {quantized['output_tokens_per_second']:.3f} |",
        "",
        f"Exact output agreement: **{agreement['exact_match_rate']:.1%}**  ",
        f"Mean text similarity: **{agreement['mean_text_similarity']:.3f}**",
        "",
        "> Agreement and text similarity compare outputs with FP32. They do not measure correctness against labels.",
        "",
        "## Environment",
        "",
        "```json",
        json.dumps(results["environment"], indent=2),
        "```",
        "",
        "## Outputs",
        "",
    ]
    for index, pair in enumerate(agreement["pairs"], start=1):
        lines.extend(
            [
                f"### Prompt {index}",
                "",
                f"**Input:** {pair['prompt']}",
                "",
                f"**FP32:** {pair['fp32_output']}",
                "",
                f"**NF4:** {pair['quantized_output']}",
                "",
            ]
        )
    return "\n".join(lines)


def main() -> None:
    args = parse_args()
    if args.runs < 1 or args.warmup_runs < 0:
        raise SystemExit("--runs must be at least 1 and --warmup-runs cannot be negative.")

    torch.manual_seed(args.seed)
    prompts = load_prompts(args.prompts_file)
    args.output_dir.mkdir(parents=True, exist_ok=True)

    environment: dict[str, Any] = {
        "python": platform.python_version(),
        "platform": platform.platform(),
        "torch": torch.__version__,
        "transformers": transformers.__version__,
        "cuda_available": torch.cuda.is_available(),
        "torch_cuda_version": torch.version.cuda,
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "seed": args.seed,
        "warmup_runs": args.warmup_runs,
        "timed_runs": args.runs,
        "max_new_tokens": args.max_new_tokens,
        "num_beams": args.num_beams,
        "prompt_count": len(prompts),
    }

    print("Benchmarking FP32 base model...")
    base = benchmark_model(
        "fp32",
        args.base_model,
        args.base_revision,
        prompts,
        args.warmup_runs,
        args.runs,
        args.max_new_tokens,
        args.num_beams,
        quantized=False,
    )
    print("Benchmarking NF4 checkpoint...")
    quantized = benchmark_model(
        "nf4",
        args.quantized_model,
        args.quantized_revision,
        prompts,
        args.warmup_runs,
        args.runs,
        args.max_new_tokens,
        args.num_beams,
        quantized=True,
        subfolder=args.quantized_subfolder,
    )

    results = {
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "environment": environment,
        "models": {"fp32": base, "quantized": quantized},
        "agreement": output_agreement(base, quantized),
    }
    json_path = args.output_dir / "results.json"
    markdown_path = args.output_dir / "results.md"
    json_path.write_text(json.dumps(results, indent=2) + "\n", encoding="utf-8")
    markdown_path.write_text(render_markdown(results), encoding="utf-8")
    print(f"Wrote {json_path} and {markdown_path}")


if __name__ == "__main__":
    main()