"""
Integration tests for Calabi compression library
"""

import unittest
import torch
import torch.nn as nn
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

try:
    from transformers import AutoConfig
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False

from calabi import (
    OptimizedCalabiYauLinear,
    OptimizedCYModelUtils,
    replace_transformer_linear_layers,
    compress_bert_model,
    compress_gpt_model
)


class TestIntegration(unittest.TestCase):
    
    def test_end_to_end_compression(self):
        """Test end-to-end compression workflow."""
        # Create a model
        model = nn.Sequential(
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )
        
        original_params = sum(p.numel() for p in model.parameters())
        
        # Compress the model
        compressed_count = OptimizedCYModelUtils.replace_linear_layers(
            model,
            compression_ratio=0.5,
            min_features=128
        )
        
        new_params = sum(p.numel() for p in model.parameters())
        
        # Test inference still works
        x = torch.randn(32, 256)
        with torch.no_grad():
            output = model(x)
        
        # Assertions
        self.assertGreater(compressed_count, 0)
        self.assertLess(new_params, original_params)
        self.assertEqual(output.shape, (32, 10))
    
    def test_gradient_flow_after_compression(self):
        """Test that gradients flow properly after compression."""
        model = nn.Sequential(
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 64)
        )
        
        # Compress
        OptimizedCYModelUtils.replace_linear_layers(
            model,
            compression_ratio=0.5,
            min_features=64
        )
        
        # Test gradient flow
        x = torch.randn(16, 128, requires_grad=True)
        output = model(x)
        loss = output.mean()
        loss.backward()
        
        # Check that gradients exist and are finite
        self.assertIsNotNone(x.grad)
        self.assertTrue(torch.all(torch.isfinite(x.grad)))
        
        for param in model.parameters():
            if param.requires_grad:
                self.assertIsNotNone(param.grad)
                if param.grad is not None:
                    self.assertTrue(torch.all(torch.isfinite(param.grad)))
    
    def test_multiple_compression_methods_consistency(self):
        """Test that different compression methods produce consistent results."""
        # Create identical models
        model1 = nn.Linear(128, 64)
        model2 = nn.Linear(128, 64)
        
        # Compress using direct layer method
        compressed_layer1 = OptimizedCalabiYauLinear(model1, compression_ratio=0.5)
        
        # Compress using utility method
        seq_model = nn.Sequential(model2)
        OptimizedCYModelUtils.replace_linear_layers(seq_model, compression_ratio=0.5, min_features=1)
        compressed_layer2 = seq_model[0]  # Get the compressed layer
        
        # Both should have the same rank
        self.assertEqual(compressed_layer1.rank, compressed_layer2.rank)
    
    def test_transformer_compression_integration(self):
        """Test transformer-specific compression integration."""
        # Create a model that mimics transformer structure
        transformer_like = nn.Module()
        transformer_like.layers = nn.ModuleList()
        
        for i in range(2):
            layer = nn.Module()
            layer.attention = nn.Linear(128, 128)
            layer.feed_forward = nn.Sequential(
                nn.Linear(128, 512),
                nn.ReLU(),
                nn.Linear(512, 128)
            )
            transformer_like.layers.append(layer)
        
        original_params = sum(p.numel() for p in transformer_like.parameters())
        
        # Use transformer-specific compression
        compressed_count = replace_transformer_linear_layers(
            transformer_like,
            compression_ratio=0.5,
            verbose=False
        )
        
        new_params = sum(p.numel() for p in transformer_like.parameters())
        
        self.assertGreater(compressed_count, 0)
        self.assertLess(new_params, original_params)
    
    def test_error_handling_integration(self):
        """Test error handling in integrated workflow."""
        # Test with very small layers
        model = nn.Sequential(
            nn.Linear(2, 4),  # Very small layer
            nn.Linear(4, 2)
        )
        
        # This should work even with small layers
        try:
            compressed_count = OptimizedCYModelUtils.replace_linear_layers(
                model,
                compression_ratio=0.5,
                min_features=1  # Allow compression of small layers
            )
            self.assertGreaterEqual(compressed_count, 0)
        except Exception as e:
            self.fail(f"Compression failed on small layers: {e}")
    
    def test_numerical_stability(self):
        """Test numerical stability of compressed models."""
        model = nn.Sequential(
            nn.Linear(64, 128),
            nn.ReLU(),
            nn.Linear(128, 64)
        )
        
        # Compress
        OptimizedCYModelUtils.replace_linear_layers(
            model,
            compression_ratio=0.5,
            min_features=32
        )
        
        # Run multiple forward passes to check for stability
        for i in range(5):
            x = torch.randn(16, 64)
            output = model(x)
            
            # Check for NaN or Inf
            self.assertTrue(torch.all(torch.isfinite(output)))
    
    def test_memory_efficiency(self):
        """Test that compressed models can be more memory efficient with appropriate settings."""
        # Use more reasonable layer sizes where compression is beneficial
        large_model = nn.Sequential(*[nn.Linear(256, 512) for _ in range(3)])
        
        original_params = sum(p.numel() for p in large_model.parameters())
        
        # Compress with moderate compression
        OptimizedCYModelUtils.replace_linear_layers(
            large_model,
            compression_ratio=0.5,  # Moderate compression
            min_features=64
        )
        
        compressed_params = sum(p.numel() for p in large_model.parameters())
        
        # For appropriate layer sizes, compression should reduce parameters
        # Allow for some overhead but expect net reduction
        self.assertLess(compressed_params, original_params * 1.2)  # Should be close to original or less


@unittest.skipUnless(TRANSFORMERS_AVAILABLE, "Transformers library not available")
class TestTransformerIntegration(unittest.TestCase):
    """Integration tests requiring transformers library."""
    
    def test_simulated_bert_compression(self):
        """Test BERT compression with simulated model."""
        # Create a model that looks like BERT
        bert_sim = nn.Module()
        bert_sim.encoder = nn.Module()
        bert_sim.encoder.layer = nn.ModuleList()
        
        for _ in range(2):  # 2 transformer layers like BERT
            layer = nn.Module()
            # Simulate BERT components
            layer.attention = nn.Module()
            layer.attention.self = nn.Linear(128, 128)
            layer.attention.output = nn.Linear(128, 128)
            layer.intermediate = nn.Linear(128, 512)  # FFN intermediate
            layer.output = nn.Linear(512, 128)       # FFN output
            bert_sim.encoder.layer.append(layer)
        
        original_params = sum(p.numel() for p in bert_sim.parameters())
        
        # Apply BERT-specific compression
        compressed_count = compress_bert_model(
            bert_sim,
            compression_ratio=0.5,
            verbose=False
        )
        
        new_params = sum(p.numel() for p in bert_sim.parameters())
        
        self.assertGreater(compressed_count, 0)
        self.assertLess(new_params, original_params)
    
    def test_simulated_gpt_compression(self):
        """Test GPT compression with simulated model."""
        # Create a model that looks like GPT
        gpt_sim = nn.Module()
        gpt_sim.transformer = nn.Module()
        gpt_sim.transformer.h = nn.ModuleList()
        
        for _ in range(2):  # 2 transformer layers like GPT
            layer = nn.Module()
            # Simulate GPT components
            layer.attn = nn.Linear(128, 128)  # Attention
            layer.mlp = nn.Sequential(        # MLP
                nn.Linear(128, 512),
                nn.Linear(512, 128)
            )
            gpt_sim.transformer.h.append(layer)
        
        original_params = sum(p.numel() for p in gpt_sim.parameters())
        
        # Apply GPT-specific compression
        compressed_count = compress_gpt_model(
            gpt_sim,
            compression_ratio=0.5,
            verbose=False
        )
        
        new_params = sum(p.numel() for p in gpt_sim.parameters())
        
        self.assertGreater(compressed_count, 0)
        self.assertLess(new_params, original_params)


if __name__ == '__main__':
    unittest.main()