"""
Unit tests for Hugging Face integration module
"""

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 BertModel, GPT2Model
    TRANSFORMERS_AVAILABLE = True
except ImportError:
    TRANSFORMERS_AVAILABLE = False

from calabi import (
    compress_hf_model,
    create_model_converter,
    save_compressed_model,
    load_compressed_model,
    replace_transformer_linear_layers
)

# Import private function directly from module
from calabi.hf_integration import _detect_model_type


class TestHFIntegration(unittest.TestCase):
    
    def test_detect_model_type_with_class_names(self):
        """Test model type detection based on class names."""
        # Create models with class names that should be detected
        class BertModel(nn.Module):
            def __init__(self):
                super().__init__()
        
        class GPT2Model(nn.Module):
            def __init__(self):
                super().__init__()
        
        bert_model = BertModel()
        gpt_model = GPT2Model()
        
        self.assertEqual(_detect_model_type(bert_model), 'bert')
        self.assertIn(_detect_model_type(gpt_model), ['gpt', 'decoder_only', 'generic'])
    
    def test_detect_model_type_generic(self):
        """Test generic model type detection."""
        # Test with a model that doesn't match known types
        generic_model = nn.Module()
        generic_model.__class__.__name__ = "MyCustomModel"
        
        result = _detect_model_type(generic_model)
        self.assertEqual(result, 'generic')
    
    def test_detect_model_type_variants(self):
        """Test detection of various model type variants."""
        # Test different model names that should map to 'bert'
        bert_models = [
            type('BertModel', (), {'__class__': type('Class', (), {})})(),
            type('RobertaModel', (), {'__class__': type('Class', (), {})})(),
            type('AlbertModel', (), {'__class__': type('Class', (), {})})()
        ]
        # Set class names manually since we can't easily modify __class__.__name__
        bert_models[0].__class__.__name__ = "BertModel"
        bert_models[1].__class__.__name__ = "RobertaModel"
        bert_models[2].__class__.__name__ = "AlbertModel"
        
        for model in bert_models:
            result = _detect_model_type(model)
            self.assertIn(result, ['bert', 'generic'])  # Some might be classified as generic
        
        # Test GPT models
        gpt_models = [
            type('GPT2Model', (), {'__class__': type('Class', (), {})})(),
            type('GPTNeoModel', (), {'__class__': type('Class', (), {})})()
        ]
        gpt_models[0].__class__.__name__ = "GPT2Model"
        gpt_models[1].__class__.__name__ = "GPTNeoModel"
        
        for model in gpt_models:
            result = _detect_model_type(model)
            self.assertIn(result, ['gpt', 'decoder_only', 'generic'])
    
    def test_create_simple_model_converter(self):
        """Test model converter with a simple model."""
        # Create a simple model to simulate transformer structure
        model = nn.Module()
        model.encoder = nn.Sequential(
            nn.Linear(256, 512),
            nn.Linear(512, 256)
        )
        model.decoder = nn.Linear(256, 100)
        
        config = {
            'method': 'calabi',
            'default_ratio': 0.5,
            'arch_specific': {
                'bert': {'attention': 0.4, 'mlp': 0.6}
            },
            'selective_layers': ['encoder', 'decoder'],
            'min_features': 128
        }
        
        # This test should work without transformers
        original_params = sum(p.numel() for p in model.parameters())
        
        # Simulate the conversion process by manually compressing
        from calabi import replace_transformer_linear_layers
        compressed_count = replace_transformer_linear_layers(
            model, 
            compression_ratio=0.5, 
            verbose=False
        )
        
        new_params = sum(p.numel() for p in model.parameters())
        
        # Should have compressed some layers
        self.assertGreaterEqual(compressed_count, 0)  # May be 0 if no layers match criteria
        # The test should not fail if no compression occurs, just check that it doesn't crash
    
    def test_compression_configs(self):
        """Test different compression configuration formats."""
        config1 = {
            'method': 'calabi',
            'default_ratio': 0.6,
            'arch_specific': {
                'bert': {'attention': 0.4, 'mlp': 0.7}
            }
        }
        
        config2 = {
            'default_ratio': 0.5,
            'patterns': {
                'attention': 0.3,
                'mlp': 0.6
            }
        }
        
        # Test that configs have required keys
        self.assertIn('default_ratio', config1)
        self.assertIn('default_ratio', config2)
        
        # Test that arch_specific is properly structured
        if 'arch_specific' in config1:
            self.assertIsInstance(config1['arch_specific'], dict)
            if 'bert' in config1['arch_specific']:
                self.assertIsInstance(config1['arch_specific']['bert'], dict)
    
    def test_save_load_simulation(self):
        """Simulate save/load functionality."""
        import tempfile
        import os
        import json
        
        # Create a simple model
        model = nn.Sequential(
            nn.Linear(128, 256),
            nn.Linear(256, 64)
        )
        
        # Compress it
        from calabi import replace_transformer_linear_layers
        replace_transformer_linear_layers(model, compression_ratio=0.5, verbose=False)
        
        # Simulate saving
        with tempfile.TemporaryDirectory() as temp_dir:
            save_path = os.path.join(temp_dir, "test_model")
            os.makedirs(save_path, exist_ok=True)
            
            # Save model state dict (simulating what save_compressed_model would do)
            model_state_dict = model.state_dict()
            torch.save(model_state_dict, os.path.join(save_path, "pytorch_model.bin"))
            
            # Save compression metadata (simulated)
            metadata = {
                'compression_method': 'calabi_yau_manifold',
                'compression_library': 'calabi',
                'original_model': 'test_model',
                'compressed_parameters': sum(p.numel() for p in model.parameters()),
            }
            
            with open(os.path.join(save_path, "calabi_compression_metadata.json"), 'w') as f:
                json.dump(metadata, f)
            
            # Verify files were created
            self.assertTrue(os.path.exists(os.path.join(save_path, "pytorch_model.bin")))
            self.assertTrue(os.path.exists(os.path.join(save_path, "calabi_compression_metadata.json")))


class TestHFIntegrationFunctionality(unittest.TestCase):
    """Tests for Hugging Face integration functionality that work without transformers."""
    
    def test_compression_metadata_extraction(self):
        """Test extraction of compression ratios from a model."""
        from calabi.hf_integration import _get_compression_ratios
        
        # Create a simple model with Calabi layers
        model = nn.Sequential(
            nn.Linear(128, 256),
            nn.Linear(256, 64)
        )
        
        # Manually compress layers to simulate
        from calabi import OptimizedCalabiYauLinear
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear) and module.in_features >= 64:
                compressed_layer = OptimizedCalabiYauLinear(module, compression_ratio=0.5)
                # Add attribute to simulate compression ratio
                compressed_layer.compression_ratio = 0.5
                name_parts = name.split('.')
                parent = model
                for part in name_parts[:-1]:
                    parent = getattr(parent, part)
                setattr(parent, name_parts[-1], compressed_layer)
        
        ratios = _get_compression_ratios(model)
        self.assertIsInstance(ratios, dict)
        # Should have entries for compressed layers
        self.assertGreater(len(ratios), 0)
    
    def test_detect_model_type_with_mock_models(self):
        """Test model type detection with mock model classes."""
        # Create mock models that mimic different transformer architectures
        bert_like = type('BertModel', (), {})()
        bert_like.__class__.__name__ = "BertModel"
        
        gpt_like = type('GPT2Model', (), {})()
        gpt_like.__class__.__name__ = "GPT2Model"
        
        roberta_like = type('RobertaModel', (), {})()
        roberta_like.__class__.__name__ = "RobertaModel"
        
        # Test detection
        from calabi.hf_integration import _detect_model_type
        
        self.assertEqual(_detect_model_type(bert_like), 'bert')
        self.assertEqual(_detect_model_type(gpt_like), 'gpt')
        self.assertEqual(_detect_model_type(roberta_like), 'bert')
    
    def test_save_compressed_model_simulation(self):
        """Test the save functionality with simulation."""
        import tempfile
        import os
        import json
        
        # Create a mock model that mimics a HuggingFace model
        class MockHFModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.model_layer = nn.Linear(128, 256)
                self.classifier = nn.Linear(256, 64)
            
            def save_pretrained(self, save_directory, **kwargs):
                # Mock save_pretrained method
                state_dict = self.state_dict()
                torch.save(state_dict, os.path.join(save_directory, "pytorch_model.bin"))
        
        model = MockHFModel()
        
        from calabi import replace_transformer_linear_layers
        replace_transformer_linear_layers(model, compression_ratio=0.5, verbose=False)
        
        # Test save_compressed_model function
        from calabi import save_compressed_model
        
        with tempfile.TemporaryDirectory() as temp_dir:
            save_path = os.path.join(temp_dir, "test_compressed_model")
            
            # This should not raise an exception
            save_compressed_model(model, save_path, "test_model")
            
            # Check that metadata file was created
            metadata_path = os.path.join(save_path, "calabi_compression_metadata.json")
            self.assertTrue(os.path.exists(metadata_path))
            
            # Check metadata content
            with open(metadata_path, 'r') as f:
                metadata = json.load(f)
            
            self.assertIn('compression_method', metadata)
            self.assertIn('compressed_parameters', metadata)
            self.assertEqual(metadata['compression_method'], 'calabi_yau_manifold')


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