"""Unit tests for the sparsity module (Task3.2).

Covers:
- Configs: UnstructuredSparsityConfig, NMSparsityConfig, BlockSparsityConfig
- Sparsifier: unstructured, N:M, block sparsity patterns
- StubSparsifier: deterministic CPU-only stub for CI
- SparseWeight / SparseModel properties
- compute_sparsity_ratio utility
- Fail-fast validation (target_sparsity out of range, n >= m)
"""

from __future__ import annotations

import numpy as np
import pytest

from sagellm_compression.sparsity import (
    BlockSparsityConfig,
    NMSparsityConfig,
    SparseModel,
    SparseWeight,
    Sparsifier,
    SparsityType,
    StubSparsifier,
    UnstructuredSparsityConfig,
)

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture()
def weight_2d() -> np.ndarray:
    """A deterministic 8x8 float32 weight matrix."""
    rng = np.random.default_rng(42)
    return rng.standard_normal((8, 8)).astype(np.float32)


@pytest.fixture()
def weight_1d() -> np.ndarray:
    """A deterministic 1-D weight vector of length 16."""
    rng = np.random.default_rng(0)
    return rng.standard_normal(16).astype(np.float32)


@pytest.fixture()
def model_weights(weight_2d: np.ndarray) -> dict[str, np.ndarray]:
    """Minimal model with two weight layers."""
    return {
        "layer0.weight": weight_2d.copy(),
        "layer1.weight": weight_2d.copy() * 2.0,
    }


# ---------------------------------------------------------------------------
# Config tests
# ---------------------------------------------------------------------------


def test_unstructured_config_defaults() -> None:
    cfg = UnstructuredSparsityConfig()
    assert cfg.sparsity_type == SparsityType.UNSTRUCTURED
    assert cfg.target_sparsity == 0.5
    assert cfg.sparsify_weights is True
    assert cfg.sparsify_activations is False


def test_nm_config_effective_sparsity() -> None:
    cfg = NMSparsityConfig(n=2, m=4)
    assert cfg.sparsity_type == SparsityType.NM
    assert cfg.effective_sparsity == pytest.approx(0.5)
    # target_sparsity is synced with effective_sparsity
    assert cfg.target_sparsity == pytest.approx(0.5)


def test_nm_config_invalid_n_ge_m() -> None:
    with pytest.raises(ValueError, match="n must be < m"):
        NMSparsityConfig(n=4, m=4)


def test_block_config_defaults() -> None:
    cfg = BlockSparsityConfig(target_sparsity=0.5)
    assert cfg.sparsity_type == SparsityType.BLOCK
    assert cfg.block_rows == 4
    assert cfg.block_cols == 4


def test_sparsity_config_target_out_of_range() -> None:
    """Sparsifier should reject target_sparsity >= 1."""
    cfg = UnstructuredSparsityConfig.model_construct(
        sparsity_type=SparsityType.UNSTRUCTURED,
        target_sparsity=1.0,
        sparsify_weights=True,
        sparsify_activations=False,
    )
    with pytest.raises(ValueError, match="target_sparsity"):
        Sparsifier(cfg)


# ---------------------------------------------------------------------------
# Sparsifier – unstructured
# ---------------------------------------------------------------------------


def test_sparsify_tensor_unstructured_sparsity_ratio(weight_2d: np.ndarray) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sp = Sparsifier(cfg)
    sw = sp.sparsify_tensor(weight_2d, layer_name="fc.weight")

    assert isinstance(sw, SparseWeight)
    # Accept small tolerance around 50% due to threshold rounding
    assert 0.45 <= sw.sparsity_ratio <= 0.60, f"Sparsity was {sw.sparsity_ratio:.2%}"


def test_sparsify_tensor_preserves_shape(weight_2d: np.ndarray) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.3)
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d)
    assert sw.data.shape == weight_2d.shape
    assert sw.mask.shape == weight_2d.shape


def test_sparsify_tensor_dtype_preserved(weight_2d: np.ndarray) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d)
    assert sw.data.dtype == weight_2d.dtype


def test_sparse_weight_to_dense(weight_2d: np.ndarray) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d)
    dense = sw.to_dense()
    assert dense.shape == weight_2d.shape
    np.testing.assert_array_equal(dense, sw.data)


def test_sparse_weight_repr(weight_2d: np.ndarray) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d, layer_name="test")
    r = repr(sw)
    assert "test" in r
    assert "sparsity=" in r


# ---------------------------------------------------------------------------
# Sparsifier – N:M
# ---------------------------------------------------------------------------


def test_sparsify_tensor_nm_pattern(weight_2d: np.ndarray) -> None:
    cfg = NMSparsityConfig(n=2, m=4)
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d)

    # Every M consecutive elements should have exactly N non-zero
    flat = sw.mask.ravel()
    for i in range(0, len(flat), 4):
        block = flat[i : i + 4]
        if len(block) == 4:
            assert int(np.sum(block)) == 2, (
                f"N:M block at offset {i} has {int(np.sum(block))} non-zeros"
            )


def test_sparsify_tensor_nm_sparsity_ratio(weight_2d: np.ndarray) -> None:
    cfg = NMSparsityConfig(n=1, m=4)  # 75% sparsity
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d)
    assert sw.sparsity_ratio == pytest.approx(0.75, abs=0.05)


# ---------------------------------------------------------------------------
# Sparsifier – block
# ---------------------------------------------------------------------------


def test_sparsify_tensor_block_sparsity_ratio(weight_2d: np.ndarray) -> None:
    cfg = BlockSparsityConfig(target_sparsity=0.5, block_rows=2, block_cols=2)
    sw = Sparsifier(cfg).sparsify_tensor(weight_2d)
    assert 0.40 <= sw.sparsity_ratio <= 0.65


def test_sparsify_tensor_block_1d_fallback(weight_1d: np.ndarray) -> None:
    """Block sparsifier should gracefully fall back to unstructured for 1-D tensors."""
    cfg = BlockSparsityConfig(target_sparsity=0.5, block_rows=4, block_cols=4)
    sw = Sparsifier(cfg).sparsify_tensor(weight_1d)
    assert sw.data.shape == weight_1d.shape


# ---------------------------------------------------------------------------
# Sparsifier – sparsify_model
# ---------------------------------------------------------------------------


def test_sparsify_model_all_layers(model_weights: dict[str, np.ndarray]) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sp = Sparsifier(cfg)
    sparse_model = sp.sparsify_model(model_weights)

    assert isinstance(sparse_model, SparseModel)
    assert set(sparse_model.weights.keys()) == set(model_weights.keys())


def test_sparse_model_overall_sparsity(model_weights: dict[str, np.ndarray]) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sparse_model = Sparsifier(cfg).sparsify_model(model_weights)
    assert 0.4 <= sparse_model.overall_sparsity <= 0.65


def test_sparse_model_repr(model_weights: dict[str, np.ndarray]) -> None:
    cfg = UnstructuredSparsityConfig(target_sparsity=0.5)
    sparse_model = Sparsifier(cfg).sparsify_model(model_weights)
    r = repr(sparse_model)
    assert "sparsity=" in r or "overall_sparsity=" in r or "SparseModel" in r


# ---------------------------------------------------------------------------
# compute_sparsity_ratio
# ---------------------------------------------------------------------------


def test_compute_sparsity_ratio_all_zeros() -> None:
    t = np.zeros((4, 4), dtype=np.float32)
    assert Sparsifier.compute_sparsity_ratio(t) == pytest.approx(1.0)


def test_compute_sparsity_ratio_no_zeros() -> None:
    t = np.ones((4, 4), dtype=np.float32)
    assert Sparsifier.compute_sparsity_ratio(t) == pytest.approx(0.0)


def test_compute_sparsity_ratio_half() -> None:
    t = np.array([0.0, 1.0, 0.0, 1.0], dtype=np.float32)
    assert Sparsifier.compute_sparsity_ratio(t) == pytest.approx(0.5)


# ---------------------------------------------------------------------------
# StubSparsifier
# ---------------------------------------------------------------------------


def test_stub_sparsifier_default_config() -> None:
    ms = StubSparsifier()
    assert ms.config.target_sparsity == pytest.approx(0.5)


def test_stub_sparsifier_custom_config(weight_2d: np.ndarray) -> None:
    cfg = NMSparsityConfig(n=1, m=4)
    ms = StubSparsifier(cfg)
    sw = ms.sparsify_tensor(weight_2d)
    assert sw.sparsity_ratio == pytest.approx(0.75, abs=0.05)


def test_stub_sparsifier_model_metadata(model_weights: dict[str, np.ndarray]) -> None:
    ms = StubSparsifier()
    sparse_model = ms.sparsify_model(model_weights)
    assert sparse_model.metadata.get("stub") is True


# ---------------------------------------------------------------------------
# Integration: public import from top-level package
# ---------------------------------------------------------------------------


def test_top_level_import() -> None:
    import sagellm_compression as sc

    assert hasattr(sc, "Sparsifier")
    assert hasattr(sc, "StubSparsifier")
    assert hasattr(sc, "UnstructuredSparsityConfig")
    assert hasattr(sc, "NMSparsityConfig")
    assert hasattr(sc, "BlockSparsityConfig")
    assert hasattr(sc, "SparseWeight")
    assert hasattr(sc, "SparseModel")
