#!/usr/bin/env python3
"""Run auditable FlashKDA correctness and latency probes on one CUDA GPU.

The upstream torch reference is imported from a pinned FlashKDA checkout. The
script deliberately uses valid synthetic A_log[H] tensors: it does not resolve
the public Kimi K3 checkpoint's A_log[128] versus H=96 inconsistency.
"""

from __future__ import annotations

import argparse
import hashlib
import importlib.metadata
import json
import math
import platform
import statistics
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable

import torch
import torch.nn.functional as F

import flash_kda
import flash_kda_C


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--flashkda-source", type=Path, required=True)
    parser.add_argument("--wheel", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--iters", type=int, default=100)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--seed", type=int, default=20260729)
    parser.add_argument("--captured-at", default=None)
    return parser.parse_args()


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def percentile(sorted_values: list[float], fraction: float) -> float:
    if not sorted_values:
        return float("nan")
    index = fraction * (len(sorted_values) - 1)
    lower = math.floor(index)
    upper = math.ceil(index)
    if lower == upper:
        return sorted_values[lower]
    weight = index - lower
    return sorted_values[lower] * (1 - weight) + sorted_values[upper] * weight


def nvidia_smi() -> dict[str, Any]:
    fields = [
        "name",
        "driver_version",
        "memory.total",
        "power.limit",
        "clocks.max.sm",
        "clocks.max.memory",
    ]
    output = subprocess.check_output(
        [
            "nvidia-smi",
            f"--query-gpu={','.join(fields)}",
            "--format=csv,noheader,nounits",
        ],
        text=True,
    ).strip()
    values = [value.strip() for value in output.split(",")]
    return dict(zip(fields, values, strict=True))


def make_inputs(
    sequence_lengths: list[int],
    heads: int,
    seed: int,
) -> dict[str, torch.Tensor | float | None]:
    batch = 1
    dimension = 128
    total_tokens = sum(sequence_lengths)
    sequences = len(sequence_lengths)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    q = F.normalize(
        torch.randn(
            (batch, total_tokens, heads, dimension),
            dtype=torch.float32,
            device="cuda",
            generator=generator,
        ),
        p=2,
        dim=-1,
    ).to(torch.bfloat16)
    k = F.normalize(
        torch.randn(
            (batch, total_tokens, heads, dimension),
            dtype=torch.float32,
            device="cuda",
            generator=generator,
        ),
        p=2,
        dim=-1,
    ).to(torch.bfloat16)
    v = torch.randn(
        (batch, total_tokens, heads, dimension),
        dtype=torch.bfloat16,
        device="cuda",
        generator=generator,
    )
    g = torch.randn(
        (batch, total_tokens, heads, dimension),
        dtype=torch.bfloat16,
        device="cuda",
        generator=generator,
    )
    beta = torch.randn(
        (batch, total_tokens, heads),
        dtype=torch.bfloat16,
        device="cuda",
        generator=generator,
    )
    a_log = torch.rand(
        heads, dtype=torch.float32, device="cuda", generator=generator
    )
    dt_bias = torch.rand(
        (heads, dimension),
        dtype=torch.float32,
        device="cuda",
        generator=generator,
    )
    initial_state = torch.randn(
        (sequences, heads, dimension, dimension),
        dtype=torch.bfloat16,
        device="cuda",
        generator=generator,
    )
    cu_seqlens = None
    if len(sequence_lengths) > 1:
        offsets = [0]
        for length in sequence_lengths:
            offsets.append(offsets[-1] + length)
        cu_seqlens = torch.tensor(offsets, dtype=torch.long, device="cuda")

    return {
        "q": q,
        "k": k,
        "v": v,
        "g": g,
        "beta": beta,
        "A_log": a_log,
        "dt_bias": dt_bias,
        "initial_state": initial_state,
        "cu_seqlens": cu_seqlens,
        "scale": 1 / math.sqrt(dimension),
    }


def run_correctness_case(
    torch_ref: Callable[..., None],
    name: str,
    sequence_lengths: list[int],
    heads: int,
    seed: int,
) -> dict[str, Any]:
    inputs = make_inputs(sequence_lengths, heads, seed)
    q = inputs["q"]
    assert isinstance(q, torch.Tensor)
    initial_state = inputs["initial_state"]
    assert isinstance(initial_state, torch.Tensor)

    out_kernel = torch.zeros_like(q)
    out_reference = torch.zeros_like(q)
    state_kernel = torch.zeros_like(initial_state)
    state_reference = torch.zeros_like(initial_state)
    common = {
        "A_log": inputs["A_log"],
        "dt_bias": inputs["dt_bias"],
        "lower_bound": -5.0,
        "initial_state": initial_state.clone(),
        "cu_seqlens": inputs["cu_seqlens"],
    }

    start = time.perf_counter()
    flash_kda.fwd(
        inputs["q"],
        inputs["k"],
        inputs["v"],
        inputs["g"],
        inputs["beta"],
        inputs["scale"],
        out_kernel,
        final_state=state_kernel,
        **common,
    )
    torch.cuda.synchronize()
    kernel_ms = (time.perf_counter() - start) * 1000

    start = time.perf_counter()
    torch_ref(
        inputs["q"],
        inputs["k"],
        inputs["v"],
        inputs["g"],
        inputs["beta"],
        inputs["scale"],
        out_reference,
        final_state=state_reference,
        **common,
    )
    torch.cuda.synchronize()
    reference_ms = (time.perf_counter() - start) * 1000

    output_diff = (out_kernel.float() - out_reference.float()).abs()
    state_diff = (state_kernel.float() - state_reference.float()).abs()
    result = {
        "name": name,
        "sequence_lengths": sequence_lengths,
        "heads": heads,
        "dimension": 128,
        "kernel_ms_first_measured": kernel_ms,
        "reference_ms": reference_ms,
        "output_exact": bool(torch.equal(out_kernel, out_reference)),
        "state_exact": bool(torch.equal(state_kernel, state_reference)),
        "output_max_abs_diff": output_diff.max().item(),
        "output_mean_abs_diff": output_diff.mean().item(),
        "state_max_abs_diff": state_diff.max().item(),
        "state_mean_abs_diff": state_diff.mean().item(),
    }
    if not result["output_exact"] or not result["state_exact"]:
        raise AssertionError(f"exact correctness failed: {result}")
    return result


def bench_events(
    function: Callable[[], None],
    warmup: int,
    iters: int,
    repeats: int,
) -> list[float]:
    for _ in range(max(warmup, 1)):
        function()
    torch.cuda.synchronize()

    elapsed: list[float] = []
    for _ in range(repeats):
        starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
        ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
        torch.cuda.synchronize()
        for index in range(iters):
            starts[index].record()
            function()
            ends[index].record()
        torch.cuda.synchronize()
        elapsed.extend(
            start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)
        )
    return elapsed


def summarize_latency(
    values: list[float],
    total_tokens: int,
) -> dict[str, Any]:
    ordered = sorted(float(value) for value in values)
    mean = statistics.fmean(ordered)
    return {
        "samples": len(ordered),
        "mean_ms": mean,
        "min_ms": ordered[0],
        "p50_ms": percentile(ordered, 0.50),
        "p95_ms": percentile(ordered, 0.95),
        "max_ms": ordered[-1],
        "sequence_tokens_per_second": total_tokens / (mean / 1000),
    }


def run_benchmark_case(
    name: str,
    sequence_lengths: list[int],
    heads: int,
    seed: int,
    warmup: int,
    iters: int,
    repeats: int,
) -> dict[str, Any]:
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    inputs = make_inputs(sequence_lengths, heads, seed)
    q = inputs["q"]
    initial_state = inputs["initial_state"]
    assert isinstance(q, torch.Tensor)
    assert isinstance(initial_state, torch.Tensor)
    out = torch.zeros_like(q)
    state_bf16 = torch.zeros_like(initial_state)
    initial_fp32 = initial_state.float()
    state_fp32 = torch.zeros_like(initial_fp32)
    common = {
        "A_log": inputs["A_log"],
        "dt_bias": inputs["dt_bias"],
        "lower_bound": -5.0,
        "cu_seqlens": inputs["cu_seqlens"],
    }

    def invoke(initial: torch.Tensor | None, final: torch.Tensor | None) -> None:
        flash_kda.fwd(
            inputs["q"],
            inputs["k"],
            inputs["v"],
            inputs["g"],
            inputs["beta"],
            inputs["scale"],
            out,
            initial_state=initial,
            final_state=final,
            **common,
        )

    variants = {
        "bf16_state": lambda: invoke(initial_state, state_bf16),
        "no_state": lambda: invoke(None, None),
        "fp32_state": lambda: invoke(initial_fp32, state_fp32),
    }
    timings = {
        variant: summarize_latency(
            bench_events(function, warmup, iters, repeats),
            sum(sequence_lengths),
        )
        for variant, function in variants.items()
    }
    torch.cuda.synchronize()
    return {
        "name": name,
        "sequence_lengths": sequence_lengths,
        "total_tokens": sum(sequence_lengths),
        "heads": heads,
        "dimension": 128,
        "warmup": warmup,
        "iters": iters,
        "repeats": repeats,
        "timings": timings,
        "output_abs_mean_after_last_run": out.float().abs().mean().item(),
        "peak_allocated_mib": torch.cuda.max_memory_allocated() / 2**20,
        "peak_reserved_mib": torch.cuda.max_memory_reserved() / 2**20,
    }


def main() -> None:
    args = parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU is required")
    if torch.cuda.get_device_capability(0) < (9, 0):
        raise RuntimeError("FlashKDA requires SM90 or newer")
    tests_dir = args.flashkda_source / "tests"
    sys.path.insert(0, str(tests_dir))
    from torch_ref import torch_ref

    revision = subprocess.check_output(
        ["git", "-C", str(args.flashkda_source), "rev-parse", "HEAD"],
        text=True,
    ).strip()
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    correctness_specs = [
        ("one_chunk", [16], 1),
        ("partial_tail", [17], 1),
        ("two_chunks", [32], 1),
        ("multi_chunk_multi_head", [65], 2),
        ("k3_head_count", [17], 96),
        ("varlen_partial_chunks", [17, 31], 2),
    ]
    correctness = [
        run_correctness_case(torch_ref, name, lengths, heads, args.seed + index)
        for index, (name, lengths, heads) in enumerate(correctness_specs)
    ]

    benchmark_specs = [
        ("teaching_scale", [512], 8),
        ("intermediate", [2048], 32),
        ("k3_fixed_shape", [8192], 96),
        ("k3_varlen_shape", [1300, 547, 2048, 963, 271, 3063], 96),
    ]
    benchmarks = [
        run_benchmark_case(
            name,
            lengths,
            heads,
            args.seed + 100 + index,
            args.warmup,
            args.iters,
            args.repeats,
        )
        for index, (name, lengths, heads) in enumerate(benchmark_specs)
    ]

    captured_at = args.captured_at or datetime.now(timezone.utc).isoformat()
    result = {
        "schema_version": 1,
        "captured_at": captured_at,
        "evidence_identity": "X / local execution on deterministic synthetic tensors",
        "boundary": {
            "k3_checkpoint_loaded": False,
            "real_token_hidden_states": False,
            "a_log_shape_conflict_resolved": False,
            "benchmark_comparison": "local FlashKDA timings only; author H20/GB200 tables remain separate",
        },
        "provenance": {
            "flashkda_revision": revision,
            "flashkda_package": importlib.metadata.version("flash-kda"),
            "wheel_filename": args.wheel.name,
            "wheel_sha256": sha256(args.wheel),
            "runner": str(Path(__file__).relative_to(Path.cwd())),
        },
        "environment": {
            "python": platform.python_version(),
            "platform": platform.platform(),
            "libc": list(platform.libc_ver()),
            "torch": torch.__version__,
            "torch_cuda": torch.version.cuda,
            "gpu": torch.cuda.get_device_name(0),
            "capability": list(torch.cuda.get_device_capability(0)),
            "nvidia_smi": nvidia_smi(),
            "flash_kda_module": flash_kda.__file__,
            "flash_kda_extension": flash_kda_C.__file__,
        },
        "correctness": {
            "all_exact": all(
                case["output_exact"] and case["state_exact"]
                for case in correctness
            ),
            "cases": correctness,
        },
        "benchmarks": benchmarks,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()