import pytest

torch = pytest.importorskip("torch")

from heapr.grouped_model import (
    ExpandedGroupedLagunaContext,
    RepackedExpandedGroupedLagunaContext,
    default_contiguous_group_indices,
    full_group_mask,
)


def test_default_contiguous_group_indices():
    indices = default_contiguous_group_indices(16, group_width=4)

    assert indices.tolist() == [
        [0, 1, 2, 3],
        [4, 5, 6, 7],
        [8, 9, 10, 11],
        [12, 13, 14, 15],
    ]


def test_full_group_mask_shape():
    mask = full_group_mask(2, 3, 4)

    assert mask.shape == (2, 3, 4)
    assert mask.all()


class TinyRouter(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden_dim = 2
        self.top_k = 1
        self.weight = torch.nn.Parameter(torch.zeros(2, 2))
        self.e_score_correction_bias = torch.nn.Parameter(torch.zeros(2), requires_grad=False)

    def forward(self, hidden_states):
        logits = hidden_states @ self.weight.T
        scores = torch.sigmoid(logits)
        selected = scores.topk(self.top_k, dim=-1).indices
        weights = scores.gather(1, selected)
        weights = weights / weights.sum(dim=-1, keepdim=True)
        return logits, weights.to(hidden_states.dtype), selected


class TinyExperts(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.gate_up_proj = torch.nn.Parameter(torch.zeros(2, 4, 2))
        self.down_proj = torch.nn.Parameter(torch.zeros(2, 2, 2))


class TinySparseMlp(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.num_experts_per_tok = 1
        self.gate = TinyRouter()
        self.experts = TinyExperts()
        self.routed_scaling_factor = 1.0

    def forward(self, hidden_states):
        return hidden_states


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()
        self.config = type("Config", (), {"moe_routed_scaling_factor": 1.0})()


def test_expanded_grouped_context_patches_and_restores():
    model = TinyModel()
    original_forward = model.model.layers[0].mlp.forward
    mask = full_group_mask(1, 2, 2)
    indices = default_contiguous_group_indices(2, group_width=1)

    with ExpandedGroupedLagunaContext(model, keep_group_masks=mask, group_indices=indices):
        assert model.model.layers[0].mlp.forward != original_forward

    assert model.model.layers[0].mlp.forward == original_forward


def test_repacked_expanded_context_rewrites_and_restores_shapes():
    model = TinyModel()
    mlp = model.model.layers[0].mlp
    mask = full_group_mask(1, 2, 2)
    indices = default_contiguous_group_indices(2, group_width=1)

    with RepackedExpandedGroupedLagunaContext(model, keep_group_masks=mask, group_indices=indices):
        assert mlp.experts.gate_up_proj.shape == (4, 2, 2)
        assert mlp.experts.down_proj.shape == (4, 2, 1)
        assert mlp.gate.weight.shape == (4, 2)
        assert mlp.gate.top_k == 2

    assert mlp.experts.gate_up_proj.shape == (2, 4, 2)
    assert mlp.experts.down_proj.shape == (2, 2, 2)
    assert mlp.gate.weight.shape == (2, 2)
    assert mlp.gate.top_k == 1


def test_repacked_layer_scaled_context_reduces_child_top_k():
    model = TinyModel()
    mlp = model.model.layers[0].mlp
    mask = full_group_mask(1, 2, 2)
    mask[0, :, 1] = False
    indices = default_contiguous_group_indices(2, group_width=1)

    with RepackedExpandedGroupedLagunaContext(
        model,
        keep_group_masks=mask,
        group_indices=indices,
        child_budget_mode="layer-scaled",
    ):
        assert mlp.gate.top_k == 1


def test_repacked_parent_weighted_context_zeros_children_past_budget():
    model = TinyModel()
    mlp = model.model.layers[0].mlp
    mlp.gate.e_score_correction_bias.data[:] = torch.tensor([1.0, 0.0])
    mask = full_group_mask(1, 2, 2)
    mask[0, 0, 1] = False
    indices = default_contiguous_group_indices(2, group_width=1)

    with RepackedExpandedGroupedLagunaContext(
        model,
        keep_group_masks=mask,
        group_indices=indices,
        child_budget_mode="parent-weighted",
    ):
        _, routing_weights, selected_children = mlp.gate(torch.zeros(1, 2))

    assert selected_children.tolist() == [[0, 2]]
    assert routing_weights.tolist() == [[1.0, 0.0]]