"""Tests for Linear/Norm layer operator replacement (Issue #38).

This module tests:
1. CustomLinear layer implementation and kernel injection
2. CustomRMSNorm layer implementation and kernel injection
3. Layer replacement logic
4. Numerical consistency with standard PyTorch implementations
"""

from __future__ import annotations

import pytest
import torch
import torch.nn as nn

from sagellm_core.kernel_injection import KernelInjector, OperatorRegistry
from sagellm_core.model.layers import (
    CustomLinear,
    CustomRMSNorm,
    replace_layers_in_model,
    replace_linear_with_custom,
)
from sagellm_protocol import OperatorConstraint, OperatorType


class DummyBackend:
    """Dummy backend for testing."""

    def __init__(self, device_type: str = "cpu"):
        self.device_type = device_type

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

        class Capability:
            def __init__(self, device_type):
                self.device_type = device_type

        return Capability(self.device_type)


# ============================================================================
# Tests for CustomLinear
# ============================================================================


def test_custom_linear_creation():
    """Test CustomLinear layer creation."""
    layer = CustomLinear(in_features=768, out_features=3072, bias=True)

    assert layer.in_features == 768
    assert layer.out_features == 3072
    assert layer.weight.shape == (3072, 768)
    assert layer.bias is not None
    assert layer.bias.shape == (3072,)


def test_custom_linear_forward_standard():
    """Test CustomLinear forward pass without kernel injection."""
    layer = CustomLinear(in_features=768, out_features=3072)

    input_tensor = torch.randn(2, 10, 768)
    output = layer(input_tensor)

    assert output.shape == (2, 10, 3072)


def test_custom_linear_forward_with_kernel():
    """Test CustomLinear forward pass with kernel injection."""

    def dummy_linear_kernel(input, weight, bias):
        """Dummy kernel that adds 1.0 to standard linear output."""
        return torch.nn.functional.linear(input, weight, bias) + 1.0

    layer = CustomLinear(in_features=768, out_features=3072)
    layer.set_kernel(dummy_linear_kernel)

    input_tensor = torch.randn(2, 10, 768)
    output = layer(input_tensor)

    # Compute expected output
    expected = torch.nn.functional.linear(input_tensor, layer.weight, layer.bias) + 1.0

    assert output.shape == (2, 10, 3072)
    assert torch.allclose(output, expected)


def test_custom_linear_numerical_consistency():
    """Test CustomLinear produces same results as nn.Linear."""
    standard_layer = nn.Linear(768, 3072)
    custom_layer = CustomLinear(768, 3072)

    # Copy weights
    with torch.no_grad():
        custom_layer.weight.copy_(standard_layer.weight)
        custom_layer.bias.copy_(standard_layer.bias)

    input_tensor = torch.randn(2, 10, 768)
    standard_output = standard_layer(input_tensor)
    custom_output = custom_layer(input_tensor)

    assert torch.allclose(standard_output, custom_output, rtol=1e-5, atol=1e-7)


# ============================================================================
# Tests for CustomRMSNorm
# ============================================================================


def test_custom_rmsnorm_creation():
    """Test CustomRMSNorm layer creation."""
    layer = CustomRMSNorm(hidden_size=768, eps=1e-6)

    assert layer.hidden_size == 768
    assert layer.eps == 1e-6
    assert layer.weight.shape == (768,)


def test_custom_rmsnorm_forward_standard():
    """Test CustomRMSNorm forward pass without kernel injection."""
    layer = CustomRMSNorm(hidden_size=768, eps=1e-6)

    input_tensor = torch.randn(2, 10, 768)
    output = layer(input_tensor)

    assert output.shape == (2, 10, 768)


def test_custom_rmsnorm_forward_with_kernel():
    """Test CustomRMSNorm forward pass with kernel injection."""

    def dummy_rmsnorm_kernel(input, weight, eps):
        """Dummy kernel that scales output by 2.0."""
        variance = input.pow(2).mean(dim=-1, keepdim=True)
        input_normalized = input * torch.rsqrt(variance + eps)
        return input_normalized * weight * 2.0

    layer = CustomRMSNorm(hidden_size=768, eps=1e-6)
    layer.set_kernel(dummy_rmsnorm_kernel)

    input_tensor = torch.randn(2, 10, 768)
    output = layer(input_tensor)

    # Compute expected output
    variance = input_tensor.pow(2).mean(dim=-1, keepdim=True)
    input_normalized = input_tensor * torch.rsqrt(variance + 1e-6)
    expected = input_normalized * layer.weight * 2.0

    assert output.shape == (2, 10, 768)
    assert torch.allclose(output, expected, rtol=1e-5)


def test_custom_rmsnorm_numerical_consistency():
    """Test CustomRMSNorm produces correct RMSNorm results."""
    layer = CustomRMSNorm(hidden_size=768, eps=1e-6)

    input_tensor = torch.randn(2, 10, 768)
    output = layer(input_tensor)

    # Compute expected output manually
    variance = input_tensor.pow(2).mean(dim=-1, keepdim=True)
    input_normalized = input_tensor * torch.rsqrt(variance + 1e-6)
    expected = input_normalized * layer.weight

    assert torch.allclose(output, expected, rtol=1e-5, atol=1e-7)


# ============================================================================
# Tests for Layer Replacement
# ============================================================================


def test_replace_linear_with_custom():
    """Test replacing nn.Linear with CustomLinear."""

    class SimpleModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(768, 3072)

    model = SimpleModel()
    original_weight = model.fc.weight.clone()
    original_bias = model.fc.bias.clone()

    replace_linear_with_custom(model, "fc", model.fc)

    # Check replacement
    assert isinstance(model.fc, CustomLinear)
    assert torch.allclose(model.fc.weight, original_weight)
    assert torch.allclose(model.fc.bias, original_bias)


def test_replace_layers_in_model():
    """Test replacing multiple layers in a model."""

    class MultiLayerModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc1 = nn.Linear(768, 3072)
            self.fc2 = nn.Linear(3072, 768)

    model = MultiLayerModel()

    stats = replace_layers_in_model(
        model,
        backend=None,
        injector=None,
        replace_linear=True,
        replace_norm=False,
    )

    assert stats["linear_replaced"] == 2
    assert isinstance(model.fc1, CustomLinear)
    assert isinstance(model.fc2, CustomLinear)


def test_replace_layers_with_kernel_injection():
    """Test replacing layers and injecting kernels."""

    def dummy_linear_kernel(input, weight, bias):
        return torch.nn.functional.linear(input, weight, bias)

    class SimpleModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(768, 3072)

    model = SimpleModel()
    backend = DummyBackend("cpu")

    # Setup registry and injector
    registry = OperatorRegistry()
    constraint = OperatorConstraint(device_type="cpu")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=dummy_linear_kernel,
        constraint=constraint,
        priority=10,
    )

    injector = KernelInjector(registry, backend)

    stats = replace_layers_in_model(
        model,
        backend=backend,
        injector=injector,
        replace_linear=True,
        replace_norm=False,
    )

    assert stats["linear_replaced"] == 1
    assert stats["linear_injected"] == 1
    assert isinstance(model.fc, CustomLinear)
    assert model.fc._kernel is not None


# ============================================================================
# Integration Tests
# ============================================================================


def test_end_to_end_linear_replacement():
    """Test end-to-end linear layer replacement with numerical consistency."""

    class TestModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(128, 256)

        def forward(self, x):
            return self.fc(x)

    # Create original model and compute output
    original_model = TestModel()
    input_tensor = torch.randn(4, 128)
    original_output = original_model(input_tensor)

    # Replace layers
    replace_layers_in_model(
        original_model,
        backend=None,
        injector=None,
        replace_linear=True,
    )

    # Compute output with replaced layers
    replaced_output = original_model(input_tensor)

    # Should be identical
    assert torch.allclose(original_output, replaced_output, rtol=1e-5, atol=1e-7)


def test_end_to_end_with_backend_kernel():
    """Test end-to-end with backend kernel injection."""
    from sagellm_backend import get_provider
    from sagellm_backend.kernels import LinearKernel

    class TestModel(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(128, 256)

        def forward(self, x):
            return self.fc(x)

    # Create model
    model = TestModel()
    input_tensor = torch.randn(4, 128)
    original_output = model(input_tensor)

    # Get backend and setup kernel registry
    backend = get_provider("cpu")
    registry = OperatorRegistry()

    # Register backend kernel
    linear_kernel = LinearKernel()
    constraint = OperatorConstraint(device_type="cpu")
    registry.register_kernel(
        operator_type=OperatorType.LINEAR,
        kernel_impl=linear_kernel.forward,
        constraint=constraint,
        priority=10,
    )

    injector = KernelInjector(registry, backend)

    # Replace layers and inject kernels
    stats = replace_layers_in_model(
        model,
        backend=backend,
        injector=injector,
        replace_linear=True,
    )

    assert stats["linear_replaced"] == 1
    assert stats["linear_injected"] == 1

    # Compute output with backend kernel
    kernel_output = model(input_tensor)

    # Should be very close (minor numerical differences possible)
    assert torch.allclose(original_output, kernel_output, rtol=1e-4, atol=1e-6)


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