import numpy as np
import pytest

torch = pytest.importorskip("torch")

from heapr.prune import (
    apply_group_mask_to_model,
    atomic_mask_from_scores,
    global_rank_atomic_scores,
    group_mask_from_scores,
)


def test_global_rank_atomic_scores_lowest_is_zero():
    scores = np.array([[[3.0, 1.0], [2.0, 4.0]]])
    ranks = global_rank_atomic_scores(scores)

    assert ranks.tolist() == [[[2, 0], [1, 3]]]


def test_atomic_mask_prunes_lowest_scores_with_min_keep():
    scores = np.array([[[1.0, 2.0, 3.0, 4.0]]])
    mask = atomic_mask_from_scores(scores, 0.5, min_keep_per_expert=1)

    assert mask.tolist() == [[[False, False, True, True]]]


def test_group_mask_keeps_at_least_one_group_per_expert():
    scores = np.array([[[1.0, 2.0, 3.0]]])
    mask = group_mask_from_scores(scores, 0.95, min_keep_per_expert=1)

    assert mask.sum() == 1
    assert mask.tolist() == [[[False, False, True]]]


class TinyExperts(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.gate_up_proj = torch.nn.Parameter(torch.ones(1, 8, 3))
        self.down_proj = torch.nn.Parameter(torch.ones(1, 3, 4))


class TinySparseMlp(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.experts = TinyExperts()


class TinyLayer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mlp = TinySparseMlp()


class TinyInner(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.ModuleList([TinyLayer()])


class TinyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.model = TinyInner()


def test_apply_group_mask_to_model_with_layer_specific_indices():
    tiny_laguna_model = TinyModel()
    keep = np.array([[[True, False]]])
    indices = np.array([[[[0, 2], [1, 3]]]])
    mlp = tiny_laguna_model.model.layers[0].mlp

    apply_group_mask_to_model(tiny_laguna_model, keep, group_width=2, group_indices=indices)

    gate_up = mlp.experts.gate_up_proj.detach()
    down = mlp.experts.down_proj.detach()
    assert gate_up[0, [1, 3], :].abs().sum().item() == 0
    assert gate_up[0, [5, 7], :].abs().sum().item() == 0
    assert down[0, :, [1, 3]].abs().sum().item() == 0