from typing import Callable
import pytest
import torch

pytest.importorskip("torch.cuda")
from .test_setup import DTYPES, DTYPE_TO_TOLS, PARAMS, SEED
from flex_sae import (
    triton_hierarchical_sae_loss,
    hierarchical_sae_loss,
    triton_topk_sae_loss,
    topk_sae_loss,
)


@pytest.fixture(autouse=True)
def _set_cuda_default_device():
    torch.set_default_device("cuda")


def run_funcs(B, K, F, D, dtype, *, kernel_foo: Callable, ref_foo: Callable):
    if dtype is torch.bfloat16 and not torch.cuda.is_bf16_supported():
        pytest.skip("BF16 not supported on this GPU")

    torch.manual_seed(SEED)

    indices = torch.randint(0, F, (B, K), dtype=torch.long, device="cuda")

    vals = torch.randn(B, K, dtype=dtype, device="cuda").abs().requires_grad_()
    decoder = torch.randn(F, D, dtype=dtype, device="cuda", requires_grad=True)
    bias = torch.randn(D, dtype=dtype, device="cuda", requires_grad=True)
    target = torch.randn(B, D, dtype=dtype, device="cuda")

    sv_ref = vals.clone().detach().requires_grad_()
    dec_ref = decoder.clone().detach().requires_grad_()
    bias_ref = bias.clone().detach().requires_grad_()

    loss_f = kernel_foo(indices, decoder, vals, bias, target)
    loss_r = ref_foo(indices, dec_ref, sv_ref, bias_ref, target)

    torch.testing.assert_close(loss_f, loss_r, **DTYPE_TO_TOLS[dtype])

    grad_out = torch.randn((), device="cuda", dtype=torch.float32)
    loss_f.backward(grad_out)
    loss_r.backward(grad_out.clone())

    torch.testing.assert_close(vals.grad, sv_ref.grad, **DTYPE_TO_TOLS[dtype])
    torch.testing.assert_close(decoder.grad, dec_ref.grad, **DTYPE_TO_TOLS[dtype])
    torch.testing.assert_close(bias.grad, bias_ref.grad, **DTYPE_TO_TOLS[dtype])

    assert indices.grad is None


@pytest.mark.parametrize("B, K, F, D", PARAMS)
@pytest.mark.parametrize("dtype", DTYPES)
def test_triton_hierarchical_sae_loss_and_grads(B, K, F, D, dtype):
    run_funcs(B, K, F, D, dtype, kernel_foo=triton_hierarchical_sae_loss, ref_foo=hierarchical_sae_loss)
    torch.cuda.empty_cache()


@pytest.mark.parametrize("B, K, F, D", PARAMS)
@pytest.mark.parametrize("dtype", DTYPES)
def test_topk_sae_loss_and_grads(B, K, F, D, dtype):
    run_funcs(
        B, K, F, D, dtype, kernel_foo=triton_topk_sae_loss, ref_foo=topk_sae_loss
    )
    torch.cuda.empty_cache()