#!/usr/bin/env python3
"""Measure Piko-9b's memory footprint precisely.

Separates weight residency from activation and KV-cache growth, and finds the
longest context that fits on the present GPU by bisection.  Every configuration
that fails is recorded rather than dropped.

    python benchmarks/profile_memory.py --model <path> --quantization 4bit \
        --output benchmarks/results/memory_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 reset() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()
        torch.cuda.synchronize()


def allocated_gb() -> float:
    return round(torch.cuda.memory_allocated() / 1024**3, 3)


def peak_gb() -> float:
    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, dict[str, Any]]:
    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)
        )

    reset()
    rss_before = host_rss_gb()
    began = time.perf_counter()
    model = AutoModelForMultimodalLM.from_pretrained(model_path, **kwargs)
    model.eval()
    torch.cuda.synchronize()
    seconds = time.perf_counter() - began
    processor = AutoProcessor.from_pretrained(model_path)

    vision_parameters = sum(p.numel() for p in model.model.visual.parameters())
    total_parameters = sum(p.numel() for p in model.parameters())

    stats = {
        "cold_load_seconds": round(seconds, 2),
        "weights_vram_gb": allocated_gb(),
        "peak_during_load_gb": peak_gb(),
        "host_rss_before_gb": rss_before,
        "host_rss_after_gb": host_rss_gb(),
        "reported_parameters": total_parameters,
        "reported_vision_parameters": vision_parameters,
        "note": "quantized parameters report packed element counts, not logical parameters",
    }
    return model, processor, stats


def measure_context(model: Any, processor: Any, tokens: int) -> dict[str, Any]:
    tokenizer = processor.tokenizer
    text = FILLER * max(1, tokens // 12)
    while len(tokenizer(text)["input_ids"]) < tokens:
        text += FILLER * 32
    ids = tokenizer(text, return_tensors="pt")["input_ids"][:, :tokens].to(model.device)

    reset()
    baseline = allocated_gb()
    began = time.perf_counter()
    with torch.inference_mode():
        model(input_ids=ids, use_cache=True)
    torch.cuda.synchronize()
    seconds = time.perf_counter() - began

    return {
        "context_tokens": int(ids.shape[1]),
        "prefill_seconds": round(seconds, 3),
        "baseline_vram_gb": baseline,
        "peak_vram_gb": peak_gb(),
        "activation_and_cache_gb": round(peak_gb() - baseline, 3),
    }


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, 32768])
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA required: CPU offload corrupts this architecture.")

    import transformers

    model, processor, load_stats = 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),
            "vram_total_gb": round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2),
            "dtype": args.dtype,
            "quantization": args.quantization,
        },
        "load": load_stats,
        "contexts": [],
        "failures": [],
    }

    for tokens in sorted(args.context_lengths):
        print(f"measuring context {tokens}", flush=True)
        try:
            measurement = measure_context(model, processor, tokens)
            report["contexts"].append(measurement)
            print(
                f"    peak {measurement['peak_vram_gb']} GB "
                f"(+{measurement['activation_and_cache_gb']} GB), "
                f"prefill {measurement['prefill_seconds']}s",
                flush=True,
            )
        except torch.cuda.OutOfMemoryError:
            report["failures"].append({"context_tokens": tokens, "error": "CUDA out of memory"})
            print(f"    OOM at {tokens} tokens — recorded, stopping context sweep", flush=True)
            reset()
            break
        except Exception as exc:  # noqa: BLE001
            report["failures"].append(
                {"context_tokens": tokens, "error": f"{type(exc).__name__}: {exc}"}
            )
            print(f"    FAILED at {tokens}: {exc}", flush=True)
            reset()
            break

    if report["contexts"]:
        report["max_context_measured"] = max(c["context_tokens"] for c in report["contexts"])

    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 __name__ == "__main__":
    main()