"""Tests for encrypted neural network modules."""

import pytest
import torch
import torch.nn as nn

from cukks.nn import (
    EncryptedModule,
    EncryptedLinear,
    EncryptedConv2d,
    EncryptedSequential,
    EncryptedFlatten,
    EncryptedSquare,
    EncryptedReLU,
    EncryptedGELU,
    EncryptedAvgPool2d,
    EncryptedMaxPool2d,
    EncryptedApproxAttention,
    EncryptedLayerNorm,
    EncryptedDropout,
)
from cukks.converter import ModelConverter


class TestEncryptedLinear:
    """Test EncryptedLinear layer."""
    
    def test_from_torch(self):
        """Test creation from PyTorch Linear."""
        torch_linear = nn.Linear(10, 5, bias=True)
        torch_linear.eval()
        
        enc_linear = EncryptedLinear.from_torch(torch_linear)
        
        assert enc_linear.in_features == 10
        assert enc_linear.out_features == 5
        assert enc_linear.bias is not None
    
    def test_from_torch_no_bias(self):
        """Test creation from PyTorch Linear without bias."""
        torch_linear = nn.Linear(10, 5, bias=False)
        torch_linear.eval()
        
        enc_linear = EncryptedLinear.from_torch(torch_linear)
        
        assert enc_linear.bias is None
    
    def test_weight_shape(self):
        """Test that weights have correct shape."""
        weight = torch.randn(5, 10)
        enc_linear = EncryptedLinear(10, 5, weight)
        
        assert enc_linear.weight.shape == (5, 10)
    
    def test_mult_depth(self):
        """Linear layer should have mult_depth 1."""
        enc_linear = EncryptedLinear(10, 5, torch.randn(5, 10))
        assert enc_linear.mult_depth() == 1
    
    def test_repr(self):
        """Test string representation."""
        enc_linear = EncryptedLinear(10, 5, torch.randn(5, 10))
        repr_str = repr(enc_linear)
        
        assert "EncryptedLinear" in repr_str
        assert "in_features=10" in repr_str
        assert "out_features=5" in repr_str


class TestEncryptedConv2d:
    """Test EncryptedConv2d layer."""
    
    def test_from_torch(self):
        """Test creation from PyTorch Conv2d."""
        torch_conv = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        torch_conv.eval()
        
        enc_conv = EncryptedConv2d.from_torch(torch_conv)
        
        assert enc_conv.in_channels == 3
        assert enc_conv.out_channels == 16
        assert enc_conv.kernel_size == (3, 3)
        assert enc_conv.padding == (1, 1)
    
    def test_weight_matrix_shape(self):
        """Test that weight is reshaped for matmul."""
        torch_conv = nn.Conv2d(3, 16, kernel_size=3)
        enc_conv = EncryptedConv2d.from_torch(torch_conv)
        
        # Should be (out_channels, in_channels * kH * kW)
        assert enc_conv.weight_matrix.shape == (16, 3 * 3 * 3)
    
    def test_output_size_calculation(self):
        """Test output size calculation."""
        enc_conv = EncryptedConv2d(
            in_channels=3,
            out_channels=16,
            kernel_size=(3, 3),
            weight=torch.randn(16, 3, 3, 3),
            stride=(2, 2),
            padding=(1, 1),
        )
        
        out_h, out_w = enc_conv.get_output_size(28, 28)
        
        # (28 + 2*1 - 3) // 2 + 1 = 14
        assert out_h == 14
        assert out_w == 14
    
    def test_unfold_input(self):
        """Test input unfolding for im2col."""
        x = torch.randn(3, 8, 8)  # C, H, W
        
        patches = EncryptedConv2d.unfold_input(
            x,
            kernel_size=(3, 3),
            stride=(1, 1),
            padding=(0, 0),
        )
        
        # Output: (6*6, 3*3*3) = (36, 27)
        assert patches.shape == (36, 27)
    
    def test_conv2d_auto_unfold(self, mock_enc_context):
        """Test that 4D input raises RuntimeError in pure HE mode."""
        torch_conv = nn.Conv2d(3, 8, kernel_size=3, padding=1)
        torch_conv.eval()
        enc_conv = EncryptedConv2d.from_torch(torch_conv)
        
        x = torch.randn(1, 3, 8, 8)
        
        enc_tensor = mock_enc_context.encrypt(x)
        with pytest.raises(RuntimeError, match="encrypt_cnn_input"):
            enc_conv(enc_tensor)
    
    def test_conv2d_chained(self, mock_enc_context):
        """Test that chained conv on 4D input raises RuntimeError in pure HE mode."""
        conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
        conv1.eval()
        
        enc_conv1 = EncryptedConv2d.from_torch(conv1)
        
        x = torch.randn(1, 3, 8, 8)
        enc_tensor = mock_enc_context.encrypt(x)
        
        with pytest.raises(RuntimeError, match="encrypt_cnn_input"):
            enc_conv1(enc_tensor)


class TestEncryptedSequential:
    """Test EncryptedSequential container."""
    
    def test_construction(self):
        """Test creating a sequential model."""
        model = EncryptedSequential(
            EncryptedLinear(10, 5, torch.randn(5, 10)),
            EncryptedSquare(),
            EncryptedLinear(5, 2, torch.randn(2, 5)),
        )
        
        assert len(model) == 3
    
    def test_indexing(self):
        """Test accessing layers by index."""
        linear = EncryptedLinear(10, 5, torch.randn(5, 10))
        square = EncryptedSquare()
        
        model = EncryptedSequential(linear, square)
        
        assert model[0] is linear
        assert model[1] is square
    
    def test_iteration(self):
        """Test iterating over layers."""
        layers = [
            EncryptedLinear(10, 5, torch.randn(5, 10)),
            EncryptedSquare(),
        ]
        model = EncryptedSequential(*layers)
        
        for i, layer in enumerate(model):
            assert layer is layers[i]
    
    def test_mult_depth(self):
        """Test total multiplicative depth."""
        model = EncryptedSequential(
            EncryptedLinear(10, 5, torch.randn(5, 10)),  # depth 1
            EncryptedSquare(),  # depth 1
            EncryptedLinear(5, 2, torch.randn(2, 5)),  # depth 1
        )
        
        assert model.mult_depth() == 3


class TestActivations:
    """Test activation functions."""
    
    def test_square_mult_depth(self):
        """Square activation has mult_depth 1."""
        assert EncryptedSquare().mult_depth() == 1
    
    def test_relu_degree(self):
        """Test ReLU with different degrees."""
        relu4 = EncryptedReLU(degree=4)
        relu7 = EncryptedReLU(degree=7)
        
        assert relu4.degree == 4
        assert relu7.degree == 7
        assert len(relu4.coeffs) == 5  # degree + 1
    
    def test_relu_mult_depth(self):
        """Higher degree should have higher depth."""
        relu2 = EncryptedReLU(degree=2)
        relu8 = EncryptedReLU(degree=8)
        
        assert relu2.mult_depth() <= relu8.mult_depth()
    
    def test_gelu_creation(self):
        """Test GELU activation."""
        gelu = EncryptedGELU(degree=4)
        assert gelu.degree == 4


class TestEncryptedFlatten:
    """Test EncryptedFlatten layer."""
    
    def test_mult_depth(self):
        """Flatten with absorbed permutation is free; without costs 1 level."""
        flatten = EncryptedFlatten()
        assert flatten.mult_depth() == 1
        
        absorbed = EncryptedFlatten._with_absorbed_permutation()
        assert absorbed.mult_depth() == 0
    
    def test_repr(self):
        """Test string representation."""
        flatten = EncryptedFlatten(start_dim=1, end_dim=-1)
        repr_str = repr(flatten)
        
        assert "EncryptedFlatten" in repr_str


class TestMaxPool2d:

    def test_maxpool2d_creation(self):
        pool = EncryptedMaxPool2d(kernel_size=2, stride=2, padding=0, degree=4)
        
        assert pool.kernel_size == (2, 2)
        assert pool.stride == (2, 2)
        assert pool.padding == (0, 0)
        assert pool.degree == 4
        assert pool.pool_area == 4

    def test_maxpool2d_from_torch(self):
        torch_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        enc_pool = EncryptedMaxPool2d.from_torch(torch_pool)
        
        assert enc_pool.kernel_size == (3, 3)
        assert enc_pool.stride == (2, 2)
        assert enc_pool.padding == (1, 1)

    def test_maxpool2d_output_size(self):
        pool = EncryptedMaxPool2d(kernel_size=2, stride=2)
        
        out_h, out_w = pool.get_output_size(8, 8)
        assert out_h == 4
        assert out_w == 4
        
        out_h2, out_w2 = pool.get_output_size(7, 7)
        assert out_h2 == 3
        assert out_w2 == 3

    def test_maxpool2d_mult_depth(self):
        pool2x2 = EncryptedMaxPool2d(kernel_size=2, degree=4)
        pool3x3 = EncryptedMaxPool2d(kernel_size=3, degree=4)
        
        assert pool2x2.mult_depth() >= 1
        assert pool3x3.mult_depth() >= pool2x2.mult_depth()

    def test_maxpool2d_accuracy(self):
        import numpy as np
        
        pool = EncryptedMaxPool2d(kernel_size=2, stride=2, degree=4)
        
        x_vals = np.linspace(-1, 1, 100)
        errors = []
        for a in x_vals[::10]:
            for b in x_vals[::10]:
                exact_max = max(a, b)
                diff = a - b
                eps = 0.01
                approx_abs = np.sqrt(diff**2 + eps)
                approx_max = (a + b + approx_abs) / 2
                abs_error = abs(approx_max - exact_max)
                errors.append(abs_error)
        
        avg_error = np.mean(errors)
        assert avg_error < 0.10, f"Average absolute error {avg_error:.4f} exceeds 0.10"

    def test_maxpool2d_repr(self):
        pool = EncryptedMaxPool2d(kernel_size=2, stride=2, padding=1, degree=6)
        repr_str = repr(pool)
        
        assert "EncryptedMaxPool2d" in repr_str
        assert "kernel_size=(2, 2)" in repr_str
        assert "degree=6" in repr_str


class TestAttention:
    """Test EncryptedApproxAttention module."""

    def test_attention_creation(self):
        """Test creating attention module with valid parameters."""
        attn = EncryptedApproxAttention(embed_dim=64, num_heads=4, softmax_degree=4)
        
        assert attn.embed_dim == 64
        assert attn.num_heads == 4
        assert attn.head_dim == 16
        assert attn.softmax_degree == 4
        assert attn.scale == 1.0 / (16 ** 0.5)

    def test_attention_creation_invalid_heads(self):
        """Test that invalid num_heads raises ValueError."""
        with pytest.raises(ValueError, match="must be divisible"):
            EncryptedApproxAttention(embed_dim=64, num_heads=5)

    def test_attention_from_torch(self):
        """Test conversion from PyTorch MultiheadAttention."""
        torch_attn = nn.MultiheadAttention(embed_dim=32, num_heads=4, batch_first=True)
        torch_attn.eval()
        
        enc_attn = EncryptedApproxAttention.from_torch(torch_attn, softmax_degree=4)
        
        assert enc_attn.embed_dim == 32
        assert enc_attn.num_heads == 4
        assert enc_attn.q_weight is not None
        assert enc_attn.k_weight is not None
        assert enc_attn.v_weight is not None
        assert enc_attn.out_weight is not None

    def test_attention_mult_depth(self):
        """Test multiplicative depth estimation."""
        attn_no_proj = EncryptedApproxAttention(embed_dim=64, num_heads=4, softmax_degree=4)
        depth_no_proj = attn_no_proj.mult_depth()
        
        assert depth_no_proj == 8  # 1 (Q@K) + 6 (Power-Softmax: sq+recip+wt) + 1 (attn@V)
        
        torch_attn = nn.MultiheadAttention(embed_dim=32, num_heads=4)
        enc_attn_with_proj = EncryptedApproxAttention.from_torch(torch_attn, softmax_degree=4)
        depth_with_proj = enc_attn_with_proj.mult_depth()
        
        assert depth_with_proj == 10  # 1 (Q,K,V proj parallel) + 1 + 6 + 1 + 1 (out proj)

    def test_attention_forward(self, mock_enc_context):
        """Test forward pass with encrypted tensors."""
        embed_dim = 8
        
        attn = EncryptedApproxAttention(embed_dim=embed_dim, num_heads=2, softmax_degree=4)
        
        query = torch.randn(1, embed_dim) * 0.5
        key = torch.randn(1, embed_dim) * 0.5
        value = torch.randn(1, embed_dim) * 0.5
        
        enc_query = mock_enc_context.encrypt(query)
        enc_key = mock_enc_context.encrypt(key)
        enc_value = mock_enc_context.encrypt(value)
        
        output = attn.forward_attention(enc_query, enc_key, enc_value)
        
        assert output is not None
        output_plain = mock_enc_context.decrypt(output)
        assert torch.isfinite(output_plain).all()

    def test_attention_approximate_accuracy(self, mock_enc_context):
        """Test that approximate attention produces reasonable outputs.
        
        Note: This is an approximation, so we check that outputs are
        in a reasonable range rather than exact equality.
        """
        import numpy as np
        
        embed_dim = 8
        
        attn = EncryptedApproxAttention(embed_dim=embed_dim, num_heads=2, softmax_degree=4)
        
        query = torch.randn(1, embed_dim) * 0.5
        key = torch.randn(1, embed_dim) * 0.5
        value = torch.randn(1, embed_dim) * 0.5
        
        enc_query = mock_enc_context.encrypt(query)
        enc_key = mock_enc_context.encrypt(key)
        enc_value = mock_enc_context.encrypt(value)
        
        output = attn.forward_attention(enc_query, enc_key, enc_value)
        output_plain = mock_enc_context.decrypt(output)
        
        assert torch.isfinite(output_plain).all(), "Output contains non-finite values"
        assert output_plain.abs().max() < 100, "Output values too large"

    def test_attention_repr(self):
        """Test string representation."""
        attn = EncryptedApproxAttention(embed_dim=64, num_heads=8, softmax_degree=6)
        repr_str = repr(attn)
        
        assert "EncryptedApproxAttention" in repr_str
        assert "embed_dim=64" in repr_str
        assert "num_heads=8" in repr_str
        assert "softmax_degree=6" in repr_str

    def test_taylor_exp_coeffs(self):
        """Test Taylor expansion coefficients for exp(x)."""
        from cukks.nn.attention import _taylor_exp_coeffs
        
        coeffs = _taylor_exp_coeffs(4)
        
        expected = [1.0, 1.0, 0.5, 1/6, 1/24]  # 1, 1, 1/2!, 1/3!, 1/4!
        
        assert len(coeffs) == 5
        for i, (c, e) in enumerate(zip(coeffs, expected)):
            assert abs(c - e) < 1e-10, f"Coefficient {i} mismatch: {c} vs {e}"


class TestAttentionMultiToken:
    """Tests for seq_len > 1 attention using Power-Softmax."""
    
    @pytest.mark.parametrize("seq_len", [2, 4, 8])
    def test_attention_multi_token_forward(self, mock_enc_context, seq_len):
        """Test forward pass with multiple tokens."""
        embed_dim = 16
        attn = EncryptedApproxAttention(embed_dim=embed_dim, num_heads=2)
        
        q = [mock_enc_context.encrypt(torch.randn(embed_dim) * 0.5) for _ in range(seq_len)]
        k = [mock_enc_context.encrypt(torch.randn(embed_dim) * 0.5) for _ in range(seq_len)]
        v = [mock_enc_context.encrypt(torch.randn(embed_dim) * 0.5) for _ in range(seq_len)]
        
        out = attn.forward_attention(q, k, v)
        
        assert isinstance(out, list)
        assert len(out) == seq_len
        for o in out:
            dec = mock_enc_context.decrypt(o)
            assert torch.isfinite(dec).all()
    
    def test_attention_seq_len_9_raises(self, mock_enc_context):
        """Test that seq_len > 8 raises NotImplementedError."""
        embed_dim = 8
        attn = EncryptedApproxAttention(embed_dim=embed_dim, num_heads=2)
        
        q = [mock_enc_context.encrypt(torch.randn(embed_dim)) for _ in range(9)]
        k = [mock_enc_context.encrypt(torch.randn(embed_dim)) for _ in range(9)]
        v = [mock_enc_context.encrypt(torch.randn(embed_dim)) for _ in range(9)]
        
        with pytest.raises(NotImplementedError, match="Maximum is 8"):
            attn.forward_attention(q, k, v)


class TestLayerNorm:
    """Test EncryptedLayerNorm module."""

    def test_layernorm_creation(self):
        """Test creating LayerNorm with normalized_shape."""
        ln = EncryptedLayerNorm(normalized_shape=64)
        
        assert ln.normalized_shape == [64]
        assert ln.eps == 1e-5
        assert ln.weight is not None
        assert ln.bias is not None

    def test_layernorm_from_torch(self):
        """Test conversion from PyTorch LayerNorm."""
        torch_ln = nn.LayerNorm(64)
        torch_ln.eval()
        
        enc_ln = EncryptedLayerNorm.from_torch(torch_ln)
        
        assert enc_ln.normalized_shape == [64]
        assert enc_ln.eps == 1e-5
        assert enc_ln.weight is not None
        assert enc_ln.bias is not None

    def test_layernorm_from_torch_no_affine(self):
        """Test conversion from PyTorch LayerNorm without affine."""
        torch_ln = nn.LayerNorm(64, elementwise_affine=False)
        torch_ln.eval()
        
        enc_ln = EncryptedLayerNorm.from_torch(torch_ln)
        
        assert enc_ln.normalized_shape == [64]

    def test_layernorm_mult_depth(self):
        """LayerNorm should have positive mult_depth for pure HE polynomial approximation."""
        ln = EncryptedLayerNorm(normalized_shape=64)
        assert ln.mult_depth() > 0

    def test_layernorm_repr(self):
        """Test string representation."""
        ln = EncryptedLayerNorm(normalized_shape=64)
        repr_str = repr(ln)
        
        assert "EncryptedLayerNorm" in repr_str
        assert "normalized_shape" in repr_str

    def test_layernorm_forward(self, mock_enc_context):
        """Test forward pass with encrypted tensors.
        
        Note: This test uses a custom mock context that handles float64
        since LayerNorm uses float64 internally. Tolerance is relaxed
        to 0.15 due to Chebyshev polynomial approximation of inv_sqrt.
        """
        from cukks.tensor import EncryptedTensor
        
        class Float64MockContext:
            def __init__(self, base_ctx):
                self._base = base_ctx
            
            def encrypt(self, tensor):
                cipher = self._base._ctx.encrypt(tensor.to(torch.float64))
                return EncryptedTensor(cipher, tuple(tensor.shape), self)
            
            def decrypt(self, enc_tensor, shape=None):
                target_shape = shape if shape else enc_tensor.shape
                result = self._base._ctx.decrypt(enc_tensor._cipher, shape=target_shape)
                return result.to(torch.float64)
        
        float64_ctx = Float64MockContext(mock_enc_context)
        
        ln = EncryptedLayerNorm(normalized_shape=8)
        
        x = torch.randn(1, 8)
        enc_x = float64_ctx.encrypt(x)
        
        enc_output = ln(enc_x)
        
        assert enc_output is not None
        assert enc_output.shape == (1, 8)
        
        output = float64_ctx.decrypt(enc_output)
        assert torch.isfinite(output).all()
        
        expected = torch.nn.functional.layer_norm(
            x.to(torch.float64), [8], ln.weight, ln.bias, ln.eps
        )
        # Polynomial inv_sqrt introduces systematic scaling error (~20-25%).
        # Use cosine similarity to verify normalization direction is correct.
        cos_sim = torch.nn.functional.cosine_similarity(
            output.flatten().unsqueeze(0),
            expected.flatten().unsqueeze(0),
        ).item()
        assert cos_sim > 0.99, f"Cosine similarity {cos_sim:.4f} too low"


class TestDropout:
    """Test EncryptedDropout module."""

    def test_dropout_creation(self):
        """Test creating Dropout with p=0.3."""
        dropout = EncryptedDropout(p=0.3)
        assert dropout.p == 0.3

    def test_dropout_default_p(self):
        """Test Dropout with default p."""
        dropout = EncryptedDropout()
        assert dropout.p == 0.5

    def test_dropout_from_torch(self):
        """Test conversion from PyTorch Dropout."""
        torch_dropout = nn.Dropout(p=0.2)
        torch_dropout.eval()
        
        enc_dropout = EncryptedDropout.from_torch(torch_dropout)
        assert enc_dropout.p == 0.2

    def test_dropout_mult_depth(self):
        """Dropout should have mult_depth 0 (passthrough)."""
        dropout = EncryptedDropout(p=0.3)
        assert dropout.mult_depth() == 0

    def test_dropout_repr(self):
        """Test string representation."""
        dropout = EncryptedDropout(p=0.3)
        repr_str = repr(dropout)
        
        assert "p=" in repr_str

    def test_dropout_forward_passthrough(self, mock_enc_context):
        """Test that dropout passes through input unchanged during inference."""
        dropout = EncryptedDropout(p=0.5)
        
        x = torch.randn(1, 8)
        enc_x = mock_enc_context.encrypt(x)
        
        enc_output = dropout(enc_x)
        
        output = mock_enc_context.decrypt(enc_output)
        input_decrypted = mock_enc_context.decrypt(enc_x)
        
        assert torch.allclose(output, input_decrypted, atol=1e-6)


class TestConvGroupedDilated:
    """Test EncryptedConv2d with grouped and dilated convolutions."""

    def test_conv2d_from_torch_grouped(self):
        """Test conversion from PyTorch Conv2d with groups."""
        torch_conv = nn.Conv2d(8, 16, kernel_size=3, groups=2)
        torch_conv.eval()
        
        enc_conv = EncryptedConv2d.from_torch(torch_conv)
        
        assert enc_conv.groups == 2
        assert enc_conv.in_channels == 8
        assert enc_conv.out_channels == 16

    def test_conv2d_from_torch_dilated(self):
        """Test conversion from PyTorch Conv2d with dilation."""
        torch_conv = nn.Conv2d(3, 8, kernel_size=3, dilation=2)
        torch_conv.eval()
        
        enc_conv = EncryptedConv2d.from_torch(torch_conv)
        
        assert enc_conv.dilation == (2, 2)
        assert enc_conv.in_channels == 3
        assert enc_conv.out_channels == 8

    def test_conv2d_output_size_dilated(self):
        """Test output size calculation with dilation."""
        enc_conv = EncryptedConv2d(
            in_channels=3,
            out_channels=8,
            kernel_size=(3, 3),
            weight=torch.randn(8, 3, 3, 3),
            stride=(1, 1),
            padding=(0, 0),
            dilation=(2, 2),
        )
        
        out_h, out_w = enc_conv.get_output_size(8, 8)
        
        # Effective kernel = 3 + (3-1)*2 = 5
        # Output = (8 - 5) // 1 + 1 = 4
        assert out_h == 4
        assert out_w == 4

    def test_conv2d_output_size_no_dilation(self):
        """Test output size calculation without dilation."""
        enc_conv = EncryptedConv2d(
            in_channels=3,
            out_channels=8,
            kernel_size=(3, 3),
            weight=torch.randn(8, 3, 3, 3),
            stride=(1, 1),
            padding=(0, 0),
            dilation=(1, 1),
        )
        
        out_h, out_w = enc_conv.get_output_size(8, 8)
        
        # Output = (8 - 3) // 1 + 1 = 6
        assert out_h == 6
        assert out_w == 6


class TestMaxPool2dForward:
    """Test EncryptedMaxPool2d forward pass."""

    def test_maxpool2d_forward_4d(self, mock_enc_context):
        """Test that 4D input raises RuntimeError in pure HE mode."""
        pool = EncryptedMaxPool2d(kernel_size=2, stride=2)
        
        x = torch.randn(1, 1, 4, 4)
        enc_x = mock_enc_context.encrypt(x)
        
        with pytest.raises(RuntimeError, match="encrypt_cnn_input"):
            pool(enc_x)


class TestAttentionSelfForward:
    """Test EncryptedApproxAttention self-attention forward."""

    def test_attention_self_forward(self, mock_enc_context):
        """Test self-attention forward(x) delegates to forward_attention(x, x, x)."""
        attn = EncryptedApproxAttention(embed_dim=8, num_heads=2, softmax_degree=4)
        
        x = torch.randn(1, 8) * 0.5
        enc_x = mock_enc_context.encrypt(x)
        
        output = attn(enc_x)
        
        assert output is not None
        output_plain = mock_enc_context.decrypt(output)
        assert torch.isfinite(output_plain).all()


class TestConverterNewModules:
    """Test ModelConverter with new modules."""

    def test_converter_layernorm(self):
        """Test converter handles LayerNorm."""
        model = nn.Sequential(
            nn.Linear(8, 8),
            nn.LayerNorm(8),
        )
        model.eval()
        
        converter = ModelConverter()
        enc_model = converter.convert(model)
        
        assert isinstance(enc_model[1], EncryptedLayerNorm)

    def test_converter_dropout(self):
        """Test converter handles Dropout."""
        model = nn.Sequential(
            nn.Linear(8, 8),
            nn.Dropout(0.3),
        )
        model.eval()
        
        converter = ModelConverter()
        enc_model = converter.convert(model)
        
        assert isinstance(enc_model[1], EncryptedDropout)

    def test_converter_dropout2d(self):
        """Test converter handles Dropout2d."""
        model = nn.Sequential(
            nn.Linear(8, 8),
            nn.Dropout2d(0.3),
        )
        model.eval()
        
        converter = ModelConverter()
        enc_model = converter.convert(model)
        
        assert isinstance(enc_model[1], EncryptedDropout)

    def test_converter_attention(self):
        """Test converter handles MultiheadAttention."""
        torch_attn = nn.MultiheadAttention(embed_dim=8, num_heads=2, batch_first=True)
        torch_attn.eval()
        
        converter = ModelConverter()
        enc_attn = converter._convert_attention(torch_attn)
        
        assert isinstance(enc_attn, EncryptedApproxAttention)

    def test_converter_grouped_conv(self):
        """Test converter handles grouped Conv2d."""
        model = nn.Sequential(
            nn.Conv2d(8, 16, kernel_size=3, groups=2),
        )
        model.eval()
        
        converter = ModelConverter()
        enc_model = converter.convert(model)
        
        assert isinstance(enc_model[0], EncryptedConv2d)
        assert enc_model[0].groups == 2

