"""
Example: Compressing a Simple Transformer Layer with CyMeta

This example demonstrates how to:
1. Create a simple transformer layer
2. Compress it using CyMeta
3. Compare compression ratios
"""

import torch
import torch.nn as nn
from cymeta.models.linear import CyMetaLinear
from cymeta.models.attention import CyMetaMultiHeadAttention
from cymeta.models.ffn import CyMetaFFN
from cymeta.training.compressor import CyMetaCompressor


def create_simple_transformer_layer(embed_dim=256, num_heads=8, ffn_dim=1024):
    """Create a simple transformer layer."""
    class SimpleTransformerLayer(nn.Module):
        def __init__(self, embed_dim, num_heads, ffn_dim):
            super().__init__()
            self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
            self.ffn = nn.Sequential(
                nn.Linear(embed_dim, ffn_dim),
                nn.GELU(),
                nn.Linear(ffn_dim, embed_dim),
            )
            self.norm1 = nn.LayerNorm(embed_dim)
            self.norm2 = nn.LayerNorm(embed_dim)
        
        def forward(self, x):
            # Self-attention
            attn_out, _ = self.attention(x, x, x)
            x = self.norm1(x + attn_out)
            
            # FFN
            ffn_out = self.ffn(x)
            x = self.norm2(x + ffn_out)
            
            return x
    
    return SimpleTransformerLayer(embed_dim, num_heads, ffn_dim)


def create_cymeta_transformer_layer(embed_dim=256, num_heads=8, ffn_dim=1024, dict_size=32):
    """Create a CyMeta-compressed transformer layer."""
    class CyMetaTransformerLayer(nn.Module):
        def __init__(self, embed_dim, num_heads, ffn_dim, dict_size):
            super().__init__()
            self.attention = CyMetaMultiHeadAttention(
                embed_dim=embed_dim,
                num_heads=num_heads,
                dict_size=dict_size,
            )
            self.ffn = CyMetaFFN(
                embed_dim=embed_dim,
                ffn_dim=ffn_dim,
                dict_size=dict_size,
            )
            self.norm1 = nn.LayerNorm(embed_dim)
            self.norm2 = nn.LayerNorm(embed_dim)
        
        def forward(self, x):
            # Self-attention
            attn_out, _ = self.attention(x, x, x)
            x = self.norm1(x + attn_out)
            
            # FFN
            ffn_out = self.ffn(x)
            x = self.norm2(x + ffn_out)
            
            return x
    
    return CyMetaTransformerLayer(embed_dim, num_heads, ffn_dim, dict_size)


def count_parameters(model):
    """Count the number of parameters in a model."""
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def main():
    print("=" * 60)
    print("CyMeta Compression Example")
    print("=" * 60)
    
    # Create models
    embed_dim = 256
    num_heads = 8
    ffn_dim = 1024
    
    print("\n1. Creating standard transformer layer...")
    standard_model = create_simple_transformer_layer(embed_dim, num_heads, ffn_dim)
    standard_params = count_parameters(standard_model)
    print(f"   Standard model parameters: {standard_params:,}")
    
    print("\n2. Creating CyMeta-compressed transformer layer...")
    cymeta_model = create_cymeta_transformer_layer(embed_dim, num_heads, ffn_dim, dict_size=32)
    cymeta_params = count_parameters(cymeta_model)
    print(f"   CyMeta model parameters: {cymeta_params:,}")
    
    # Calculate compression ratio
    compression_ratio = standard_params / cymeta_params
    print(f"\n3. Compression ratio: {compression_ratio:.2f}x")
    
    # Test forward pass
    print("\n4. Testing forward pass...")
    batch_size = 2
    seq_len = 10
    x = torch.randn(batch_size, seq_len, embed_dim)
    
    with torch.no_grad():
        standard_out = standard_model(x)
        cymeta_out = cymeta_model(x)
    
    print(f"   Standard output shape: {standard_out.shape}")
    print(f"   CyMeta output shape: {cymeta_out.shape}")
    
    # Example: Compress a single linear layer
    print("\n5. Compressing a single Linear layer...")
    linear_layer = nn.Linear(128, 64)
    compressor = CyMetaCompressor(num_atoms=32, gating_hidden=16)
    
    print("   Compressing (this may take a moment)...")
    compressed = compressor.compress_layer(
        linear_layer,
        num_iterations=50,  # Reduced for example
        verbose=True,
    )
    
    print("   Compression complete!")
    print(f"   Original weight shape: {linear_layer.weight.shape}")
    print(f"   Dictionary atoms: {compressed['dictionary'].num_atoms}")
    print(f"   Index map shape: {compressed['index_map'].shape}")
    
    print("\n" + "=" * 60)
    print("Example complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()

