"""Correctness and performance comparison between Python and C++ hull caches."""

from __future__ import annotations

import time

import mlx.core as mx
import numpy as np
import pytest

from mlx_transformer_vm.attention.hull_cache import HullKVCache, has_hull_extension
from mlx_transformer_vm.attention.hull_python import PythonHullKVCache

N_LAYERS = 2
N_HEADS = 18
HEAD_DIM = 2
D_MODEL = N_HEADS * HEAD_DIM

skip_no_ext = pytest.mark.skipif(
    not has_hull_extension(),
    reason="C++ hull extension not available",
)


@pytest.fixture
def seed():
    np.random.seed(42)


def _random_vectors(n):
    """Return n tuples of (keys, queries, values) as mx.array."""
    vecs = []
    for _ in range(n):
        k = mx.array(np.random.randn(D_MODEL).astype(np.float32))
        q = mx.array(np.random.randn(D_MODEL).astype(np.float32))
        v = mx.array(np.random.randn(D_MODEL).astype(np.float32))
        vecs.append((k, q, v))
    return vecs


def _setup_tiebreak(cache, layer, latest_heads):
    for h in latest_heads:
        cache.set_tiebreak(layer, h, True)


@skip_no_ext
def test_hull_python_cpp_correctness(seed):
    """PythonHullKVCache and C++ HullKVCache produce identical outputs."""
    py_cache = PythonHullKVCache(N_LAYERS, N_HEADS)
    cpp_cache = HullKVCache(N_LAYERS, N_HEADS)

    latest_heads = [0, 3, 7, 12, 15]
    for layer in range(N_LAYERS):
        _setup_tiebreak(py_cache, layer, latest_heads)
        _setup_tiebreak(cpp_cache, layer, latest_heads)

    vecs = _random_vectors(200)
    for step, (k, q, v) in enumerate(vecs):
        for layer in range(N_LAYERS):
            out_py = np.array(py_cache.layer_step(layer, k, q, v))
            out_cpp = np.array(cpp_cache.layer_step(layer, k, q, v))
            assert np.allclose(out_py, out_cpp, atol=1e-6), (
                f"mismatch at step {step} layer {layer}: "
                f"max diff {np.max(np.abs(out_py - out_cpp)):.2e}"
            )


def _bench_cache(cache, vecs, n_layers):
    """Run insert+query sequence, return elapsed seconds."""
    start = time.perf_counter()
    for k, q, v in vecs:
        for layer in range(n_layers):
            cache.layer_step(layer, k, q, v)
    return time.perf_counter() - start


@skip_no_ext
def test_hull_python_cpp_performance(seed, capsys):
    """Benchmark Python vs C++ hull across varying sequence lengths."""
    seq_lengths = [100, 500, 1000, 5000]
    results = []

    for n in seq_lengths:
        np.random.seed(42)
        vecs = _random_vectors(n)

        py_cache = PythonHullKVCache(1, N_HEADS)
        _setup_tiebreak(py_cache, 0, [0, 7, 15])
        t_py = _bench_cache(py_cache, vecs, 1)

        np.random.seed(42)
        vecs = _random_vectors(n)

        cpp_cache = HullKVCache(1, N_HEADS)
        _setup_tiebreak(cpp_cache, 0, [0, 7, 15])
        t_cpp = _bench_cache(cpp_cache, vecs, 1)

        ratio = t_py / t_cpp if t_cpp > 0 else float("inf")
        results.append((n, t_py, t_cpp, ratio))

    with capsys.disabled():
        print("\n")
        print(f"{'Seq Len':>8}  {'Python (s)':>11}  {'C++ (s)':>11}  {'Ratio':>7}")
        print(f"{'-------':>8}  {'-----------':>11}  {'-----------':>11}  {'-----':>7}")
        for n, t_py, t_cpp, ratio in results:
            print(f"{n:>8}  {t_py:>11.4f}  {t_cpp:>11.4f}  {ratio:>7.1f}x")
