"""Tests for fused kernels integration.

This module tests the fused kernel integration implemented for issue #39.
"""

from __future__ import annotations

import pytest
import torch
import torch.nn as nn

from sagellm_core.kernel.fusion_detector import (
    FusionDetector,
    FusionStats,
    detect_and_apply_fusions,
)
from sagellm_core.model.layers.fused_layers import (
    FusedSiLUMul,
    FusedAddRMSNorm,
    FusedQKVProjection,
)


# Test fixtures
class SimpleBackend:
    """Simple backend for testing."""

    def __init__(self, has_fused_kernels: bool = True):
        self._has_fused_kernels = has_fused_kernels

    def capability(self):
        """Return capability."""

        class BackendCapability:
            device_name = "cpu"

        return BackendCapability()


class SimpleModel(nn.Module):
    """Simple model for testing fusion detection."""

    def __init__(self):
        super().__init__()
        self.fused_silu_mul = FusedSiLUMul()
        self.fused_add_rmsnorm = FusedAddRMSNorm(hidden_size=512)
        self.fused_qkv = FusedQKVProjection(
            hidden_size=512,
            num_heads=8,
            num_kv_heads=8,
            head_dim=64,
        )
        self.other_layer = nn.Linear(512, 512)


# Tests for FusionStats
def test_fusion_stats_defaults():
    """Test FusionStats default values."""
    stats = FusionStats()
    assert stats.fused_silu_mul_count == 0
    assert stats.fused_add_rmsnorm_count == 0
    assert stats.fused_qkv_projection_count == 0
    assert stats.total_fusion_opportunities == 0
    assert stats.fusion_enabled is False


def test_fusion_stats_get_fusion_ratio():
    """Test fusion ratio calculation."""
    stats = FusionStats()
    assert stats.get_fusion_ratio() == 0.0

    stats.total_fusion_opportunities = 10
    stats.fused_silu_mul_count = 3
    stats.fused_add_rmsnorm_count = 2
    stats.fused_qkv_projection_count = 1
    assert stats.get_fusion_ratio() == 0.6


def test_fusion_stats_to_dict():
    """Test stats conversion to dict."""
    stats = FusionStats(
        fused_silu_mul_count=2,
        fused_add_rmsnorm_count=1,
        fused_qkv_projection_count=1,
        total_fusion_opportunities=4,
        fusion_enabled=True,
    )
    stats_dict = stats.to_dict()

    assert stats_dict["fused_silu_mul_count"] == 2
    assert stats_dict["fused_add_rmsnorm_count"] == 1
    assert stats_dict["fused_qkv_projection_count"] == 1
    assert stats_dict["total_fusion_opportunities"] == 4
    assert stats_dict["fusion_ratio"] == 1.0
    assert stats_dict["fusion_enabled"] is True


# Tests for FusionDetector
def test_fusion_detector_creation():
    """Test FusionDetector creation."""
    backend = SimpleBackend()
    detector = FusionDetector(backend=backend, enable_fusion=True)

    assert detector._backend == backend
    assert detector._enable_fusion is True
    assert detector._stats.fusion_enabled is True


def test_fusion_detector_disabled():
    """Test FusionDetector when fusion is disabled."""
    backend = SimpleBackend()
    detector = FusionDetector(backend=backend, enable_fusion=False)
    model = SimpleModel()

    stats = detector.apply_fusions(model)

    assert stats.fused_silu_mul_count == 0
    assert stats.fused_add_rmsnorm_count == 0
    assert stats.fused_qkv_projection_count == 0
    assert stats.total_fusion_opportunities == 0


def test_fusion_detector_no_backend():
    """Test FusionDetector without backend."""
    detector = FusionDetector(backend=None, enable_fusion=True)
    model = SimpleModel()

    stats = detector.apply_fusions(model)

    # Should still detect opportunities but not inject kernels
    assert stats.fused_silu_mul_count == 0


def test_fusion_detector_apply_fusions():
    """Test fusion detector applies fusions correctly."""
    backend = SimpleBackend()
    detector = FusionDetector(backend=backend, enable_fusion=True)
    model = SimpleModel()

    stats = detector.apply_fusions(model)

    # Should detect all three fused layers
    assert stats.fused_silu_mul_count == 1
    assert stats.fused_add_rmsnorm_count == 1
    assert stats.fused_qkv_projection_count == 1
    assert stats.total_fusion_opportunities == 3
    assert stats.get_fusion_ratio() == 1.0


def test_fusion_detector_reset_stats():
    """Test fusion detector stats reset."""
    backend = SimpleBackend()
    detector = FusionDetector(backend=backend, enable_fusion=True)
    model = SimpleModel()

    # Apply fusions and collect stats
    stats = detector.apply_fusions(model)
    assert stats.total_fusion_opportunities > 0

    # Reset stats
    detector.reset_stats()
    new_stats = detector.get_stats()
    assert new_stats.total_fusion_opportunities == 0
    assert new_stats.fused_silu_mul_count == 0


# Tests for fused layers
def test_fused_silu_mul_forward():
    """Test FusedSiLUMul forward pass."""
    layer = FusedSiLUMul()
    gate = torch.randn(2, 4, 512)
    up = torch.randn(2, 4, 512)

    output = layer(gate, up)

    # Check output shape
    assert output.shape == gate.shape

    # Verify computation (without kernel injection, uses PyTorch fallback)
    expected = torch.nn.functional.silu(gate) * up
    assert torch.allclose(output, expected, atol=1e-6)


def test_fused_add_rmsnorm_forward():
    """Test FusedAddRMSNorm forward pass."""
    layer = FusedAddRMSNorm(hidden_size=512)
    x = torch.randn(2, 4, 512)
    residual = torch.randn(2, 4, 512)

    output, new_residual = layer(x, residual)

    # Check output shapes
    assert output.shape == x.shape
    assert new_residual.shape == residual.shape

    # Verify residual update
    expected_residual = x + residual
    assert torch.allclose(new_residual, expected_residual, atol=1e-6)


def test_fused_qkv_projection_forward():
    """Test FusedQKVProjection forward pass."""
    layer = FusedQKVProjection(
        hidden_size=512,
        num_heads=8,
        num_kv_heads=8,
        head_dim=64,
    )
    hidden_states = torch.randn(2, 4, 512)

    q, k, v = layer(hidden_states)

    # Check output shapes
    batch_size, seq_len, _ = hidden_states.shape
    assert q.shape == (batch_size, seq_len, 8, 64)
    assert k.shape == (batch_size, seq_len, 8, 64)
    assert v.shape == (batch_size, seq_len, 8, 64)


def test_fused_qkv_projection_with_bias():
    """Test FusedQKVProjection with bias."""
    layer = FusedQKVProjection(
        hidden_size=512,
        num_heads=8,
        num_kv_heads=8,
        head_dim=64,
        bias=True,
    )

    assert layer.qkv_bias is not None
    assert layer.qkv_bias.shape == (8 * 64 + 2 * 8 * 64,)


# Integration tests
def test_detect_and_apply_fusions():
    """Test convenience function for fusion detection and application."""
    backend = SimpleBackend()
    model = SimpleModel()

    stats = detect_and_apply_fusions(model, backend, enable_fusion=True)

    assert stats.fused_silu_mul_count == 1
    assert stats.fused_add_rmsnorm_count == 1
    assert stats.fused_qkv_projection_count == 1
    assert stats.get_fusion_ratio() == 1.0


def test_detect_and_apply_fusions_disabled():
    """Test convenience function with fusion disabled."""
    backend = SimpleBackend()
    model = SimpleModel()

    stats = detect_and_apply_fusions(model, backend, enable_fusion=False)

    assert stats.total_fusion_opportunities == 0
    assert stats.get_fusion_ratio() == 0.0


def test_nested_model_fusion():
    """Test fusion detection in nested models."""

    class NestedModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.layer1 = nn.ModuleList(
                [
                    FusedSiLUMul(),
                    FusedSiLUMul(),
                ]
            )
            self.layer2 = nn.Sequential(
                FusedAddRMSNorm(512),
                nn.Linear(512, 512),
            )

    backend = SimpleBackend()
    model = NestedModel()

    stats = detect_and_apply_fusions(model, backend, enable_fusion=True)

    # Should detect all nested fused layers
    assert stats.fused_silu_mul_count == 2
    assert stats.fused_add_rmsnorm_count == 1


# Performance and edge cases
def test_fused_silu_mul_extra_repr():
    """Test FusedSiLUMul extra_repr."""
    layer = FusedSiLUMul()
    repr_str = layer.extra_repr()
    assert "unfused" in repr_str

    # After setting kernel, should show "fused"
    backend = SimpleBackend()
    layer.set_backend(backend)
    layer.set_kernel(lambda g, u: torch.nn.functional.silu(g) * u)
    repr_str = layer.extra_repr()
    assert "fused" in repr_str


def test_fused_add_rmsnorm_extra_repr():
    """Test FusedAddRMSNorm extra_repr."""
    layer = FusedAddRMSNorm(hidden_size=512, eps=1e-6)
    repr_str = layer.extra_repr()
    assert "512" in repr_str
    assert "1e-06" in repr_str or "1e-6" in repr_str


def test_fused_qkv_projection_extra_repr():
    """Test FusedQKVProjection extra_repr."""
    layer = FusedQKVProjection(
        hidden_size=512,
        num_heads=8,
        num_kv_heads=8,
        head_dim=64,
    )
    repr_str = layer.extra_repr()
    assert "512" in repr_str
    assert "8" in repr_str
    assert "64" in repr_str


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
