import pytest
import numpy as np
import torch
import time
import tracemalloc

from vicentin.kernels import Kernel, SumKernel, ProductKernel
from vicentin.kernels.laplacian_kernel import LaplacianKernel


# --- Dummy Kernel for Base Testing ---
class LinearKernel(Kernel):
    def _compute(self, x, y=None):
        if y is None:
            y = x
        if isinstance(x, torch.Tensor):
            return x @ y.T
        return x @ y.T


# --- Fixtures ---
@pytest.fixture
def np_data():
    np.random.seed(42)
    return np.random.randn(10, 5)


@pytest.fixture
def torch_data(np_data):
    return torch.tensor(np_data, dtype=torch.float32)


@pytest.fixture
def v_np():
    np.random.seed(42)
    return np.random.randn(10, 2)


@pytest.fixture
def v_torch(v_np):
    return torch.tensor(v_np, dtype=torch.float32)


# --- Tests ---


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_backend_dispatch(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data
    k = LinearKernel(data)
    assert k.backend is not None
    assert type(k.backend).__name__.lower().endswith(backend)
    assert k.N == 10


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_dense_vs_chunked_eval(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data

    k_dense = LinearKernel(data, chunk_size=None, alpha=1e-5)
    k_chunked = LinearKernel(data, chunk_size=3, alpha=1e-5)

    out_dense = k_dense(data)
    out_chunked = k_chunked(data)

    if backend == "numpy":
        np.testing.assert_allclose(out_dense, out_chunked, rtol=1e-5)
    else:
        torch.testing.assert_close(out_dense, out_chunked)


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_cache_and_getitem(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data
    k = LinearKernel(data, cache_size=2, chunk_size=3)

    # Trigger cache for specific rows
    row_0 = k[0]
    row_1 = k[1]

    assert 0 in k.backend._cache
    assert 1 in k.backend._cache

    # Cache eviction (size is 2)
    _ = k[2]
    assert len(k.backend._cache) == 2
    assert 0 not in k.backend._cache or 1 not in k.backend._cache

    # Test tuple slicing
    block = k[0:2, 1:3]
    assert block.shape == (2, 2)


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_matmul_and_apply(np_data, torch_data, v_np, v_torch, backend):
    data = np_data if backend == "numpy" else torch_data
    v = v_np if backend == "numpy" else v_torch

    k_dense = LinearKernel(data, chunk_size=None)
    k_chunked = LinearKernel(data, chunk_size=4)

    res_dense = k_dense @ v
    res_chunked = k_chunked @ v
    res_apply = k_chunked.apply(data, v)

    if backend == "numpy":
        np.testing.assert_allclose(res_dense, res_chunked, rtol=1e-5)
        np.testing.assert_allclose(res_dense, res_apply, rtol=1e-5)
    else:
        torch.testing.assert_close(res_dense, res_chunked)
        torch.testing.assert_close(res_dense, res_apply)


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_trace(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data

    k_dense = LinearKernel(data, chunk_size=None, alpha=0.1)
    k_chunked = LinearKernel(data, chunk_size=3, alpha=0.1)

    tr_dense = k_dense.trace()
    tr_chunked = k_chunked.trace()

    assert pytest.approx(tr_dense, rel=1e-4) == tr_chunked


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_kernel_algebra(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data

    k1 = LinearKernel(data)
    k2 = LinearKernel(data)

    k_sum = k1 + k2
    k_prod = k1 * k2
    k_scaled = 2.5 * k1
    k_pow = k1**2

    assert isinstance(k_sum, SumKernel)
    assert isinstance(k_prod, ProductKernel)

    x_eval = data[:2]

    if backend == "numpy":
        np.testing.assert_allclose(
            k_sum(x_eval), k1(x_eval) + k2(x_eval), rtol=1e-5
        )
        np.testing.assert_allclose(
            k_prod(x_eval), k1(x_eval) * k2(x_eval), rtol=1e-5
        )
        np.testing.assert_allclose(
            k_scaled(x_eval), 2.5 * k1(x_eval), rtol=1e-5
        )
        np.testing.assert_allclose(k_pow(x_eval), k1(x_eval) ** 2, rtol=1e-5)
    else:
        torch.testing.assert_close(k_sum(x_eval), k1(x_eval) + k2(x_eval))
        torch.testing.assert_close(k_pow(x_eval), k1(x_eval) ** 2)


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_centering(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data
    k = LinearKernel(data)
    k_c = k.center()

    K_centered = k_c(data)

    # The sum of rows/cols of a centered kernel matrix should be near 0
    if backend == "numpy":
        assert np.allclose(K_centered.sum(axis=0), 0, atol=1e-4)
    else:
        assert torch.allclose(K_centered.sum(dim=0), torch.zeros(10), atol=1e-4)


@pytest.mark.parametrize("backend", ["numpy", "torch"])
def test_laplacian_kernel(np_data, torch_data, backend):
    data = np_data if backend == "numpy" else torch_data
    k = LaplacianKernel(data)

    assert k.gamma is not None  # Auto-gamma should trigger

    K_mat = k(data)

    # Distance to self is 0, exp(0) is 1. Diagonals should be 1 + alpha
    diag = np.diag(K_mat) if backend == "numpy" else torch.diagonal(K_mat)
    expected_diag = 1.0 + k.alpha

    if backend == "numpy":
        np.testing.assert_allclose(diag, expected_diag, rtol=1e-5)
    else:
        torch.testing.assert_close(diag, torch.full_like(diag, expected_diag))


def test_caching_performance():
    # Large feature dimension to make compute artificially expensive
    data = np.random.randn(500, 5000)

    k_cached = LinearKernel(data, cache_size=10)
    k_uncached = LinearKernel(data, cache_size=0)

    # Warm up cache
    _ = k_cached[0]

    start_cached = time.perf_counter()
    for _ in range(50):
        _ = k_cached[0]
    time_cached = time.perf_counter() - start_cached

    start_uncached = time.perf_counter()
    for _ in range(50):
        _ = k_uncached[0]
    time_uncached = time.perf_counter() - start_uncached

    assert time_cached < time_uncached, "Caching did not improve access time."


def test_chunking_memory():
    # Large N to force significant memory allocation (N=2000 -> 4M elements -> ~32MB for float64)
    data = np.random.randn(2000, 10)
    v = np.random.randn(2000, 1)

    # 1. Measure Dense Memory (Matmul triggers full N x N instantiation if not chunked)
    tracemalloc.start()
    k_dense = LinearKernel(data, chunk_size=None)
    _ = k_dense @ v
    _, peak_dense = tracemalloc.get_traced_memory()
    tracemalloc.stop()

    # 2. Measure Chunked Memory
    tracemalloc.start()
    k_chunked = LinearKernel(data, chunk_size=100, cache_size=0)
    _ = k_chunked @ v
    _, peak_chunked = tracemalloc.get_traced_memory()
    tracemalloc.stop()

    # The chunked peak memory should be strictly less than the dense peak memory
    assert (
        peak_chunked < peak_dense
    ), "Chunking did not reduce peak memory usage."


def test_stress_cache_and_chunks():
    # Stress test: large N, small chunks, small cache, heavy trace and matmul ops
    data = np.random.randn(3000, 50)
    v = np.random.randn(3000, 5)

    k_stress = LinearKernel(data, chunk_size=200, cache_size=5)

    try:
        # Operations that should route through chunks and avoid OOM
        trace_val = k_stress.trace()
        res_matmul = k_stress @ v

        # Thrash the cache
        for i in range(20):
            _ = k_stress[i]

        success = True
    except MemoryError:
        success = False

    assert success, "Stress test failed due to memory constraints."
    assert res_matmul.shape == (3000, 5)
