import pytest
import torch
from kernels.rank_estimator import sketch_rank, estimate_prune_counts
from tests.conftest import make_softmax_matrix


def test_single_matrix(device):
    P = make_softmax_matrix(1, 77, 196, device)
    svd = torch.linalg.matrix_rank(P[0]).item()
    skc = sketch_rank(P).item()
    assert abs(svd - skc) <= 2


def test_batched(device):
    B, T, V = 8, 77, 196
    P    = make_softmax_matrix(B, T, V, device)
    svd  = torch.stack([torch.linalg.matrix_rank(P[i]) for i in range(B)]).float()
    skc  = sketch_rank(P).float()
    assert (svd - skc).abs().max().item() <= 2


def test_known_low_rank(device):
    torch.manual_seed(0)
    U = torch.randn(50, 5, device=device)
    V = torch.randn(5, 100, device=device)
    A = (U @ V).unsqueeze(0)
    assert abs(sketch_rank(A).item() - 5) <= 2


def test_prune_counts_valid(device):
    P = make_softmax_matrix(4, 32, 196, device)
    counts = estimate_prune_counts(P, 196)
    assert counts.shape == (4,)
    assert (counts >= 0).all()
    assert (counts < 196).all()


def test_high_res(device):
    P = make_softmax_matrix(4, 128, 576, device)
    ranks = sketch_rank(P)
    assert ranks.shape == (4,)
    assert (ranks > 0).all()