#!/usr/bin/env python3
"""Benchmark numerically verified cached prefill/decode for dense and Q25."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import gc
import json
from pathlib import Path
import sys
import time

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
sys.path.insert(0, str(ROOT / "scripts"))

from run_strata_headquotient_v1_1_frontier import construct_export  # noqa: E402
from strata.experiments.compose_rf import load_dense_base  # noqa: E402
from strata.modeling.compose import HeadQuotientLM  # noqa: E402
from strata.training.lm_data import PackedLMDataset  # noqa: E402


DEFAULT_CONFIG = ROOT / "configs/experiments/strata_headquotient_v1_1.json"


def parse_ints(value: str) -> tuple[int, ...]:
    return tuple(int(item) for item in value.split(",") if item)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
    parser.add_argument("--device", default="cuda:0" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--prefix-lengths", type=parse_ints, default=(2048, 4096, 8160))
    parser.add_argument("--batch-sizes", type=parse_ints, default=(1, 4, 8, 16))
    parser.add_argument("--decode-steps", type=int, default=32)
    parser.add_argument("--warmup", type=int, default=1)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--smoke", action="store_true")
    return parser.parse_args()


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


def load_models(config: dict, device: torch.device):
    dense_base, model_config = load_dense_base(
        ROOT / config["model_config"], Path(config["dense_checkpoint"]), device,
    )
    dense = HeadQuotientLM(dense_base, model_config, gamma_max=0.05)
    for block in dense.base_model.blocks:
        block.strip_graph_adapter()
    dense.to(device).eval()

    output = Path(config["output_root"])
    plan = json.loads((output / "selection/Q25/plan.json").read_text(encoding="utf-8"))
    scoped = json.loads((output / "scoped_adapter/result.json").read_text(encoding="utf-8"))
    q25, _cfg, *_rest = construct_export(config, plan, scoped, device)
    q25.load_state_dict(
        torch.load(output / "frontier/Q25/exported_model.pt", map_location=device, weights_only=True),
        strict=True,
    )
    q25.eval()
    return dense, q25


@torch.inference_mode()
def verify_incremental(model, ids: torch.Tensor) -> dict[str, float | bool]:
    prefix = ids[:, :-1]
    final = ids[:, -1:]
    with sdpa_kernel(SDPBackend.MATH):
        full = model(ids, graph_enabled=False).logits[:, -1].float()
        cached = model.prefill_cache(prefix, maximum_length=ids.shape[1])
        incremental = model.decode_step(final, cached.cache).logits[:, -1].float()
    difference = (full - incremental).abs()
    probability_difference = (
        full.softmax(-1) - incremental.softmax(-1)
    ).abs()
    top_equal = bool(torch.equal(full.argmax(-1), incremental.argmax(-1)))
    return {
        "maximum_absolute_logit_difference": float(difference.max()),
        "mean_absolute_logit_difference": float(difference.mean()),
        "maximum_absolute_probability_difference": float(probability_difference.max()),
        "top_token_exact": top_equal,
        "numerically_equivalent": bool(
            top_equal
            and float(difference.max()) <= 0.15
            and float(probability_difference.max()) <= 5e-3
        ),
    }


@torch.inference_mode()
def profile_one(
    model,
    tokens: torch.Tensor,
    *,
    decode_steps: int,
    warmup: int,
    repeats: int,
    device: torch.device,
) -> dict[str, float | int]:
    prefix = tokens[:, :-decode_steps]
    continuation = tokens[:, -decode_steps:]
    maximum_length = tokens.shape[1]

    for _ in range(warmup):
        state = model.prefill_cache(prefix, maximum_length=maximum_length)
        for step in range(decode_steps):
            state = model.decode_step(continuation[:, step : step + 1], state.cache)
    synchronize(device)

    prefill_elapsed = 0.0
    decode_elapsed = 0.0
    first_token_elapsed = 0.0
    peak = 0
    persistent = 0
    for _ in range(repeats):
        gc.collect()
        if device.type == "cuda":
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats(device)
        started = time.perf_counter()
        state = model.prefill_cache(prefix, maximum_length=maximum_length)
        synchronize(device)
        prefill_elapsed += time.perf_counter() - started
        persistent = sum(
            tensor.numel() * tensor.element_size()
            for layer in state.cache
            for tensor in (
                layer.global_key, layer.global_value,
                layer.local_key, layer.local_value,
            )
        )
        started = time.perf_counter()
        state = model.decode_step(continuation[:, :1], state.cache)
        synchronize(device)
        first_token_elapsed += time.perf_counter() - started
        started = time.perf_counter()
        for step in range(1, decode_steps):
            state = model.decode_step(continuation[:, step : step + 1], state.cache)
        synchronize(device)
        decode_elapsed += time.perf_counter() - started
        if device.type == "cuda":
            peak = max(peak, int(torch.cuda.max_memory_allocated(device)))

    batch = tokens.shape[0]
    prefix_tokens = prefix.numel()
    decoded_tokens = batch * decode_steps
    total_decode_time = first_token_elapsed + decode_elapsed
    return {
        "batch_size": batch,
        "prefix_length": prefix.shape[1],
        "decode_steps": decode_steps,
        "prefill_seconds": prefill_elapsed / repeats,
        "prefill_tokens_per_second": repeats * prefix_tokens / prefill_elapsed,
        "first_cached_token_ms": 1000 * first_token_elapsed / repeats,
        "decode_seconds": total_decode_time / repeats,
        "decode_tokens_per_second": repeats * decoded_tokens / total_decode_time,
        "persistent_kv_bytes": persistent,
        "peak_allocated_bytes": peak,
    }


def main() -> None:
    args = parse_args()
    config = json.loads(args.config.read_text(encoding="utf-8"))
    device = torch.device(args.device)
    if device.type == "cuda":
        torch.cuda.set_device(device)
    if args.smoke:
        args.prefix_lengths = (64,)
        args.batch_sizes = (1,)
        args.decode_steps = 4
        args.repeats = 1
    maximum = max(args.prefix_lengths) + args.decode_steps
    if maximum > 8192:
        raise ValueError("prefix plus decode steps exceeds true-8k context")
    dataset = PackedLMDataset(config["heldout_corpus"], seq_len=maximum)
    source = dataset[min(700, len(dataset) - 1)].unsqueeze(0).to(device)
    dense, q25 = load_models(config, device)
    verification_ids = source[:, : min(maximum, 96)]
    verification = {
        "dense": verify_incremental(dense, verification_ids),
        "q25": verify_incremental(q25, verification_ids),
    }
    if not all(row["numerically_equivalent"] for row in verification.values()):
        raise RuntimeError(f"cached decoding differs from full decoding: {verification}")

    rows = []
    for prefix_length in args.prefix_lengths:
        length = prefix_length + args.decode_steps
        for batch_size in args.batch_sizes:
            tokens = source[:, :length].expand(batch_size, -1).contiguous()
            values = {}
            for name, model in (("dense", dense), ("q25", q25)):
                try:
                    values[name] = profile_one(
                        model, tokens, decode_steps=args.decode_steps,
                        warmup=args.warmup, repeats=args.repeats, device=device,
                    )
                except torch.OutOfMemoryError:
                    values[name] = {"out_of_memory": True}
                    if device.type == "cuda":
                        torch.cuda.empty_cache()
            row = {
                "prefix_length": prefix_length,
                "batch_size": batch_size,
                "dense": values["dense"],
                "q25": values["q25"],
            }
            if not any(value.get("out_of_memory", False) for value in values.values()):
                row["ratios"] = {
                    "prefill_throughput": (
                        values["q25"]["prefill_tokens_per_second"]
                        / values["dense"]["prefill_tokens_per_second"]
                    ),
                    "decode_throughput": (
                        values["q25"]["decode_tokens_per_second"]
                        / values["dense"]["decode_tokens_per_second"]
                    ),
                    "persistent_kv": (
                        values["q25"]["persistent_kv_bytes"]
                        / values["dense"]["persistent_kv_bytes"]
                    ),
                }
            rows.append(row)
            print(json.dumps(row, sort_keys=True), flush=True)

    payload = {
        "program": "STRATA-HEADQUOTIENT-Q25-CACHED-DECODE-PROFILE",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "device": str(device),
        "torch_version": torch.__version__,
        "graph_enabled": False,
        "measurement": "inference-only KV cache; BF16 numerical equivalence is checked against full-sequence logits; prefill and autoregressive decode are reported separately",
        "verification": verification,
        "rows": rows,
    }
    output = args.output or Path(config["output_root"]) / "cached_decode/profile.json"
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({"output": str(output), "rows": len(rows)}, sort_keys=True))


if __name__ == "__main__":
    main()