"""
Unit and integration tests for torch-supercluster.
Run with: pytest tests/
"""

import torch
import torch.nn as nn
import pytest

from supercluster import (
    SuperCluster,
    center_diversity_loss,
    effective_prototype_count,
    MLPBlock,
)


# ── Fixtures ──────────────────────────────────────────────────────────────────

@pytest.fixture
def binary_model():
    return SuperCluster(input_dim=10, embedding_dim=32, num_clusters=4,
                        num_attn_heads=4, d_ff=64)


@pytest.fixture
def multiclass_model():
    return SuperCluster(input_dim=10, embedding_dim=32, num_clusters=4,
                        num_classes=3, num_attn_heads=4, d_ff=64)


@pytest.fixture
def x():
    torch.manual_seed(0)
    return torch.randn(16, 10)


# ── Shape tests ────────────────────────────────────────────────────────────────

class TestOutputShapes:
    def test_binary_logits_shape(self, binary_model, x):
        logits, cw = binary_model(x)
        assert logits.shape == (16, 1)

    def test_binary_cluster_weights_shape(self, binary_model, x):
        _, cw = binary_model(x)
        assert cw.shape == (16, 4)

    def test_cluster_weights_sum_to_one(self, binary_model, x):
        _, cw = binary_model(x)
        assert torch.allclose(cw.sum(dim=-1), torch.ones(16), atol=1e-5)

    def test_multiclass_logits_shape(self, multiclass_model, x):
        logits, _ = multiclass_model(x)
        assert logits.shape == (16, 3)

    def test_prototype_scores_binary_shape(self, binary_model):
        assert binary_model.prototype_scores.shape == (4,)

    def test_prototype_scores_multiclass_shape(self, multiclass_model):
        assert multiclass_model.prototype_scores.shape == (4, 3)


# ── Loss and regulariser tests ─────────────────────────────────────────────────

class TestLosses:
    def test_center_diversity_loss_scalar(self, binary_model):
        loss = center_diversity_loss(binary_model.centers)
        assert loss.ndim == 0

    def test_center_diversity_loss_in_range(self, binary_model):
        loss = center_diversity_loss(binary_model.centers)
        assert -1.0 <= loss.item() <= 1.0

    def test_effective_prototype_count_range(self, binary_model, x):
        _, cw = binary_model(x)
        k_eff = effective_prototype_count(cw)
        assert 1.0 <= k_eff <= 4.0

    def test_effective_prototype_count_uniform(self):
        """Uniform routing → K_eff = K."""
        K = 6
        cw = torch.ones(32, K) / K
        k_eff = effective_prototype_count(cw)
        assert abs(k_eff - K) < 0.01

    def test_effective_prototype_count_collapsed(self):
        """All mass on one prototype → K_eff ≈ 1."""
        K = 6
        cw = torch.zeros(32, K)
        cw[:, 0] = 1.0
        k_eff = effective_prototype_count(cw)
        assert abs(k_eff - 1.0) < 0.01


# ── Gradient flow tests ───────────────────────────────────────────────────────

class TestGradients:
    def test_binary_backward(self, binary_model, x):
        y = torch.randint(0, 2, (16, 1)).float()
        logits, cw = binary_model(x)
        loss = nn.BCEWithLogitsLoss()(logits, y)
        loss += 0.1 * center_diversity_loss(binary_model.centers)
        loss.backward()
        assert binary_model.prototype_scores.grad is not None
        assert binary_model.centers.grad is not None

    def test_multiclass_backward(self, multiclass_model, x):
        y = torch.randint(0, 3, (16,))
        logits, cw = multiclass_model(x)
        loss = nn.CrossEntropyLoss()(logits, y)
        loss.backward()
        assert multiclass_model.prototype_scores.grad is not None

    def test_regression_backward(self, binary_model, x):
        y = torch.randn(16, 1)
        preds, _ = binary_model(x)
        loss = nn.MSELoss()(preds, y)
        loss.backward()
        assert binary_model.prototype_scores.grad is not None


# ── active_prototypes tests ───────────────────────────────────────────────────

class TestActivePrototypes:
    def test_returns_list(self, binary_model, x):
        _, cw = binary_model(x)
        active = binary_model.active_prototypes(cw)
        assert isinstance(active, list)

    def test_all_zero_weights_returns_empty(self, binary_model):
        cw = torch.zeros(8, 4)
        cw[:, 0] = 1.0   # only prototype 0
        active = binary_model.active_prototypes(cw, threshold=0.01)
        assert active == [0]


# ── MLPBlock tests ────────────────────────────────────────────────────────────

class TestMLPBlock:
    def test_output_shape(self):
        block = MLPBlock(10, 32, num_layers=3, hidden_size=64)
        x = torch.randn(8, 10)
        assert block(x).shape == (8, 32)

    def test_skip_identity_when_same_dim(self):
        block = MLPBlock(32, 32, num_layers=2, hidden_size=32)
        assert isinstance(block.skip, nn.Identity)


# ── Backward-compat alias ─────────────────────────────────────────────────────

def test_backward_compat_alias():
    from supercluster import SupervisedClusteringModelV2
    m = SupervisedClusteringModelV2(input_dim=5, embedding_dim=16,
                                    num_clusters=3, num_attn_heads=4, d_ff=32)
    assert isinstance(m, SuperCluster)
