import pytest
import torch

import kernels

dt = kernels.get_kernel("phanerozoic/det-train", version=1, trust_remote_code=True)

requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")


@requires_cuda
@pytest.mark.kernels_ci
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_sharding_invariance_bitwise(dtype):
    """The core contract: any K-sharding on canonical block boundaries
    composes to the bitwise-identical result."""
    torch.manual_seed(0)
    M, N, K = 64, 96, 8192
    x = (torch.randn(M, K, device="cuda") *
         torch.exp2(torch.randint(-20, 21, (M, K), device="cuda").float())).to(dtype)
    w = (torch.randn(N, K, device="cuda") *
         torch.exp2(torch.randint(-20, 21, (N, K), device="cuda").float())).to(dtype)
    full = dt.det_gemm(x, w)
    for shards in (2, 4, 8):
        bs = K // shards
        assert bs % 1024 == 0
        states = [dt.det_gemm_partial(x[:, i*bs:(i+1)*bs].contiguous(),
                                      w[:, i*bs:(i+1)*bs].contiguous())
                  for i in range(shards)]
        y = dt.det_finalize(*dt.det_combine(states))
        assert torch.equal(y, full), f"{shards}-way shard differs"
    # grouping invariance: ((a,b),(c,d)) vs (((a,b),c),d)
    st = [dt.det_gemm_partial(x[:, i*2048:(i+1)*2048].contiguous(),
                              w[:, i*2048:(i+1)*2048].contiguous()) for i in range(4)]
    left = dt.det_combine([dt.det_combine(st[:2]), dt.det_combine(st[2:])])
    seq = dt.det_combine(st)
    assert torch.equal(left[0], seq[0]) and torch.equal(left[1], seq[1])


@requires_cuda
@pytest.mark.kernels_ci
def test_repeat_bitwise_and_accuracy():
    torch.manual_seed(1)
    x = torch.randn(128, 4096, device="cuda")
    w = torch.randn(256, 4096, device="cuda")
    y1, y2 = dt.det_gemm(x, w), dt.det_gemm(x, w)
    assert torch.equal(y1, y2)
    ref = (x.double() @ w.double().t())
    rel = ((y1.double() - ref).abs().max() / ref.abs().max()).item()
    assert rel < 2 ** -30, rel   # window accuracy far exceeds fp32 accumulation


@requires_cuda
@pytest.mark.kernels_ci
def test_training_bitwise_repeatable_and_shard_equal():
    """Two full training runs are bitwise identical; a run whose forward is
    computed via 4-way K-sharded composition equals the unsharded run."""
    def train(sharded: bool):
        torch.manual_seed(7)
        l1 = dt.DetLinear(2048, 256, device="cuda")
        opt = torch.optim.SGD(l1.parameters(), lr=1e-3)
        x = torch.randn(32, 2048, device="cuda")
        t = torch.randn(32, 256, device="cuda")
        losses = []
        for _ in range(50):
            if sharded:
                st = [dt.det_gemm_partial(x[:, i*1024:(i+1)*1024].contiguous(),
                                          l1.weight[:, i*1024:(i+1)*1024].contiguous())
                      for i in range(2)]
                y = dt.det_finalize(*dt.det_combine(st))
                # backward through the unsharded op on the same values
                y2 = _DetProxy.apply(x, l1.weight, y)
            else:
                y2 = l1(x)
            loss = (y2 - t).square().mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            losses.append(loss.detach().clone())
        return torch.stack(losses)

    class _DetProxy(torch.autograd.Function):
        @staticmethod
        def forward(ctx, x, W, y):
            ctx.save_for_backward(x, W)
            return y
        @staticmethod
        def backward(ctx, dY):
            x, W = ctx.saved_tensors
            dX = dt.det_gemm(dY.contiguous(), W.t().contiguous())
            dW = dt.det_gemm(dY.t().contiguous(), x.t().contiguous())
            return dX, dW, None
    globals()['_DetProxy'] = _DetProxy

    a = train(False)
    b = train(False)
    assert torch.equal(a, b), "repeat runs must be bitwise identical"
    c = train(True)
    assert torch.equal(a, c), "sharded-forward run must equal unsharded bitwise"


@requires_cuda
@pytest.mark.kernels_ci
def test_det_sum_matches_permutation():
    torch.manual_seed(3)
    x = torch.randn(64, 4096, device="cuda") * torch.exp2(
        torch.randint(-15, 16, (64, 4096), device="cuda").float())
    a = dt.det_sum(x, dim=-1)
    idx = torch.randperm(4096, device="cuda")
    b = dt.det_sum(x[:, idx], dim=-1)
    assert torch.equal(a, b), "det_sum must be permutation-invariant"