#!/usr/bin/env python3
"""Measure Piko-9b's practical inference characteristics.

Reports cold-load time, time to first token, prompt-processing throughput,
decode throughput, and image-preprocessing time, across context lengths and
batch sizes.  Only configurations that actually run on the present hardware are
recorded; nothing is extrapolated.

    python benchmarks/profile_inference.py --model <path> --quantization 4bit \
        --context-lengths 512 2048 8192 --batch-sizes 1 2 4 \
        --output benchmarks/results/inference_4bit.json
"""

from __future__ import annotations

import argparse
import gc
import json
import platform
import time
from pathlib import Path
from typing import Any

import torch

FILLER = "Routine operations log entry: all monitored systems reported nominal status. "


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


def reset_memory() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()


def peak_vram_gb() -> float | None:
    if not torch.cuda.is_available():
        return None
    return round(torch.cuda.max_memory_allocated() / 1024**3, 3)


def host_rss_gb() -> float | None:
    try:
        import psutil
    except ImportError:
        return None
    return round(psutil.Process().memory_info().rss / 1024**3, 3)


def load(model_path: str, quantization: str, dtype: str) -> tuple[Any, Any, float]:
    from transformers import AutoModelForMultimodalLM, AutoProcessor

    torch_dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16}[dtype]
    kwargs: dict[str, Any] = {"dtype": torch_dtype, "device_map": {"": 0}}
    if quantization in ("4bit", "8bit"):
        from transformers import BitsAndBytesConfig

        kwargs["quantization_config"] = (
            BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch_dtype,
                bnb_4bit_use_double_quant=True,
            )
            if quantization == "4bit"
            else BitsAndBytesConfig(load_in_8bit=True)
        )

    began = time.perf_counter()
    model = AutoModelForMultimodalLM.from_pretrained(model_path, **kwargs)
    model.eval()
    synchronize()
    load_seconds = time.perf_counter() - began

    processor = AutoProcessor.from_pretrained(model_path)
    return model, processor, load_seconds


def build_prompt(processor: Any, target_tokens: int) -> str:
    """Grow filler text until the tokenised prompt is close to the target."""
    tokenizer = processor.tokenizer
    repeats = max(1, target_tokens // 12)
    text = FILLER * repeats
    while len(tokenizer(text)["input_ids"]) < target_tokens:
        text += FILLER * max(1, repeats // 8)
    return text + "\n\nSummarise the log in one sentence."


def time_generation(
    model: Any,
    processor: Any,
    prompts: list[str],
    max_new_tokens: int,
) -> dict[str, Any]:
    texts = [
        processor.apply_chat_template(
            [{"role": "user", "content": [{"type": "text", "text": prompt}]}],
            add_generation_prompt=True,
            tokenize=False,
        )
        for prompt in prompts
    ]
    inputs = processor(text=texts, return_tensors="pt", padding=True).to(model.device)
    prompt_tokens = int(inputs["input_ids"].shape[1])
    batch = len(prompts)

    # Warm up once so kernel autotuning and lazy allocation do not land in the timings.
    with torch.inference_mode():
        model.generate(**inputs, max_new_tokens=1, do_sample=False)
    synchronize()

    reset_memory()

    # Time-to-first-token: a full generate capped at one new token. Measuring a bare
    # forward pass instead would exclude generate()'s own setup, and subtracting it
    # from the full run can go negative and produce absurd decode rates.
    synchronize()
    began = time.perf_counter()
    with torch.inference_mode():
        model.generate(**inputs, max_new_tokens=1, do_sample=False)
    synchronize()
    ttft_seconds = time.perf_counter() - began

    synchronize()
    began = time.perf_counter()
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
    synchronize()
    total_seconds = time.perf_counter() - began

    new_tokens = int(output.shape[1]) - prompt_tokens
    decode_tokens = new_tokens - 1  # the first token is attributed to TTFT
    decode_seconds = total_seconds - ttft_seconds

    # If decode time is not positive the run was too short to measure; report null
    # rather than a number produced by dividing by a floor.
    if decode_tokens > 0 and decode_seconds > 1e-3:
        decode_rate: float | None = round(decode_tokens * batch / decode_seconds, 2)
    else:
        decode_rate = None

    return {
        "batch_size": batch,
        "prompt_tokens": prompt_tokens,
        "new_tokens_per_sequence": new_tokens,
        "time_to_first_token_s": round(ttft_seconds, 3),
        "prefill_tokens_per_s": round(prompt_tokens * batch / max(ttft_seconds, 1e-9), 1),
        "decode_tokens_per_s": decode_rate,
        "decode_seconds": round(decode_seconds, 3),
        "total_seconds": round(total_seconds, 2),
        "peak_vram_gb": peak_vram_gb(),
        "measurement_note": None if decode_rate else "decode window too short to measure",
    }


def time_image_preprocessing(processor: Any, image: Path, repeats: int = 5) -> dict[str, Any]:
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "url": str(image)},
                {"type": "text", "text": "Describe this image."},
            ],
        }
    ]
    timings = []
    for _ in range(repeats):
        began = time.perf_counter()
        inputs = processor.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        )
        timings.append(time.perf_counter() - began)
    timings.sort()
    return {
        "image": image.name,
        "repeats": repeats,
        "median_seconds": round(timings[len(timings) // 2], 4),
        "min_seconds": round(timings[0], 4),
        "max_seconds": round(timings[-1], 4),
        "prompt_tokens_with_image": int(inputs["input_ids"].shape[1]),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--label", default=None)
    parser.add_argument("--quantization", default="4bit", choices=["none", "4bit", "8bit"])
    parser.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16"])
    parser.add_argument("--context-lengths", type=int, nargs="+", default=[512, 2048, 8192])
    parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 2, 4])
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--image", type=Path, default=None)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    import transformers

    reset_memory()
    rss_before = host_rss_gb()
    model, processor, load_seconds = load(args.model, args.quantization, args.dtype)

    report: dict[str, Any] = {
        "model": args.model,
        "label": args.label or Path(args.model).name,
        "environment": {
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
            "python": platform.python_version(),
            "platform": platform.platform(),
            "torch": torch.__version__,
            "transformers": transformers.__version__,
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
            "vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2)
            if torch.cuda.is_available()
            else None,
            "dtype": args.dtype,
            "quantization": args.quantization,
            "decoding": "greedy (do_sample=False)",
            "linear_attention_kernels": "not installed (pure PyTorch fallback)",
        },
        "cold_load": {
            "seconds": round(load_seconds, 2),
            "resident_vram_gb": peak_vram_gb(),
            "host_rss_before_gb": rss_before,
            "host_rss_after_gb": host_rss_gb(),
        },
        "text": [],
        "image": [],
        "failures": [],
    }

    for context in args.context_lengths:
        prompt = build_prompt(processor, context)
        for batch in args.batch_sizes:
            label = f"context={context} batch={batch}"
            print(f"measuring {label}", flush=True)
            try:
                measurement = time_generation(
                    model, processor, [prompt] * batch, args.max_new_tokens
                )
                measurement["requested_context"] = context
                report["text"].append(measurement)
                rate = measurement["decode_tokens_per_s"]
                print(
                    f"    ttft {measurement['time_to_first_token_s']}s, "
                    f"decode {rate if rate is not None else 'not measurable'} tok/s, "
                    f"peak {measurement['peak_vram_gb']} GB",
                    flush=True,
                )
            except torch.cuda.OutOfMemoryError:
                report["failures"].append({"case": label, "error": "CUDA out of memory"})
                print(f"    OOM at {label} — recorded, continuing", flush=True)
                reset_memory()
            except Exception as exc:  # noqa: BLE001
                report["failures"].append({"case": label, "error": f"{type(exc).__name__}: {exc}"})
                print(f"    FAILED {label}: {exc}", flush=True)
                reset_memory()

    if args.image and args.image.is_file():
        print("measuring image preprocessing", flush=True)
        try:
            report["image"].append(time_image_preprocessing(processor, args.image))
        except Exception as exc:  # noqa: BLE001
            report["failures"].append({"case": "image_preprocessing", "error": str(exc)})

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(f"\nwrote {args.output}")
    if report["failures"]:
        print(f"{len(report['failures'])} configuration(s) failed and were recorded.")


if __name__ == "__main__":
    main()