"""Tests for HuggingFace conversion utilities."""

import torch
import torch.nn as nn

from cymeta.exports.converter import (
    ConversionConfig,
    convert_model_linear_layers,
)
from cymeta.models.linear import CyMetaLinear
from cymeta.utils.cpu import ensure_cpu_model, get_compression_summary


class TinyModel(nn.Module):
    """Simple feed-forward network for conversion tests."""

    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(16, 32)
        self.block = nn.Sequential(
            nn.Linear(32, 32),
            nn.ReLU(),
            nn.Linear(32, 8),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.linear1(x)
        x = self.block(x)
        return x


def test_convert_model_linear_layers():
    """Validate linear layers get replaced with CyMetaLinear and remain functional."""
    model = TinyModel()
    config = ConversionConfig(dict_size=8, compress_iterations=5, verbose=False)

    converted_model, converted_count = convert_model_linear_layers(model, config=config)

    assert converted_count == 3
    assert isinstance(converted_model.linear1, CyMetaLinear)
    assert isinstance(converted_model.block[0], CyMetaLinear)
    assert isinstance(converted_model.block[2], CyMetaLinear)

    x = torch.randn(2, 16)
    y = converted_model(x)
    assert y.shape == (2, 8)


def test_cpu_optimized_config():
    """Test that cpu_optimized=True applies correct settings."""
    config = ConversionConfig(cpu_optimized=True)
    
    # Check that CPU-optimized settings are applied
    assert config.dict_size == 16
    assert config.gating_hidden == 8
    assert config.compress_iterations == 150
    assert config.device == torch.device("cpu")


def test_ensure_cpu_model():
    """Test ensure_cpu_model utility."""
    model = TinyModel()
    config = ConversionConfig(dict_size=8, compress_iterations=5, verbose=False)
    converted_model, _ = convert_model_linear_layers(model, config=config)
    
    # Move to CPU and ensure eval mode
    cpu_model = ensure_cpu_model(converted_model)
    
    assert next(cpu_model.parameters()).device.type == "cpu"
    assert not cpu_model.training  # Should be in eval mode
    assert not any(p.requires_grad for p in cpu_model.parameters())


def test_get_compression_summary():
    """Test compression summary utility."""
    model = TinyModel()
    config = ConversionConfig(dict_size=8, compress_iterations=5, verbose=False)
    converted_model, _ = convert_model_linear_layers(model, config=config)
    
    summary = get_compression_summary(converted_model)
    
    assert summary["total_layers"] == 3
    assert summary["compressed_layers"] == 3
    assert summary["compression_coverage"] == 1.0
    assert summary["average_compression_ratio"] > 0


