"""
CyMeta Linear Layer.

Drop-in replacement for nn.Linear that uses CyMeta compression internally.
"""

import torch
import torch.nn as nn
from typing import Optional
from cymeta.dictionary import MetaDictionary
from cymeta.index_map import IndexMap
from cymeta.gating import GatingNetwork
from cymeta.reconstruct import reconstruct_weights


class CyMetaLinear(nn.Module):
    """
    Linear layer with CyMeta compression.
    
    This is a drop-in replacement for nn.Linear that internally uses
    CyMeta compression. The weights are reconstructed on-the-fly during
    forward passes using dictionary atoms, index maps, and gating networks.
    
    Args:
        in_features: Input feature dimension
        out_features: Output feature dimension
        dict_size: Number of dictionary atoms (default: 32)
        shift_bits: Number of bits for shift encoding (determines max_shift)
        gating_hidden: Hidden dimension for gating network (default: 16)
        bias: Whether to include bias term (default: True)
        original_weight: Optional original weight matrix to compress (for initialization)
        use_cached_weights: Whether to cache reconstructed weights (default: False)
    
    Attributes:
        dictionary: MetaDictionary for this layer
        index_map: IndexMap for this layer
        gating_network: GatingNetwork for this layer
        bias: Optional bias parameter
    """
    
    def __init__(
        self,
        in_features: int,
        out_features: int,
        dict_size: int = 32,
        shift_bits: int = 8,
        gating_hidden: int = 16,
        bias: bool = True,
        original_weight: Optional[torch.Tensor] = None,
        use_cached_weights: bool = False,
    ):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.dict_size = dict_size
        self.shift_bits = shift_bits
        self.gating_hidden = gating_hidden
        self.use_cached_weights = use_cached_weights
        
        # Calculate max shift from shift_bits
        max_shift = (1 << (shift_bits - 1)) - 1  # e.g., 8 bits -> max_shift = 127
        
        # Initialize dictionary
        self.dictionary = MetaDictionary(
            num_atoms=dict_size,
            atom_dim=in_features,
        )
        
        # Initialize index map
        if original_weight is not None:
            # Initialize from original weight (would need compression algorithm)
            # For now, create empty and let training handle it
            self.index_map = IndexMap.create_empty(
                out_dim=out_features,
                in_dim=in_features,
                num_atoms=dict_size,
                max_shift=max_shift,
            )
        else:
            self.index_map = IndexMap.create_empty(
                out_dim=out_features,
                in_dim=in_features,
                num_atoms=dict_size,
                max_shift=max_shift,
            )
        
        # Initialize gating network
        self.gating_network = GatingNetwork(
            input_dim=in_features,  # Using input dimension for token embedding
            num_atoms=dict_size,
            hidden_dim=gating_hidden,
        )
        
        # Initialize bias
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features))
        else:
            self.register_parameter("bias", None)
        
        # Cache for reconstructed weights (if enabled)
        self._cached_weights: Optional[torch.Tensor] = None
        self._cache_token_embedding: Optional[torch.Tensor] = None
    
    def forward(
        self,
        input: torch.Tensor,
        token_embedding: Optional[torch.Tensor] = None,
        head: Optional[int] = None,
    ) -> torch.Tensor:
        """
        Forward pass with on-the-fly weight reconstruction.
        
        Args:
            input: Input tensor of shape (..., in_features)
            token_embedding: Optional token embedding for context-dependent gating.
                           If None, uses input as token embedding.
            head: Optional attention head index
        
        Returns:
            Output tensor of shape (..., out_features)
        """
        # Use input as token embedding if not provided
        if token_embedding is None:
            token_embedding = input
        
        # Reconstruct weights
        if self.use_cached_weights and self._cached_weights is not None:
            # Check if cache is still valid (simple check - can be improved)
            if torch.equal(token_embedding, self._cache_token_embedding):
                weight = self._cached_weights
            else:
                weight = self._reconstruct_weights(token_embedding, head)
                self._cached_weights = weight
                self._cache_token_embedding = token_embedding.clone()
        else:
            weight = self._reconstruct_weights(token_embedding, head)
            if self.use_cached_weights:
                self._cached_weights = weight
                self._cache_token_embedding = token_embedding.clone()
        
        # Apply linear transformation
        output = torch.nn.functional.linear(input, weight, self.bias)
        
        return output
    
    def _reconstruct_weights(
        self,
        token_embedding: torch.Tensor,
        head: Optional[int] = None,
    ) -> torch.Tensor:
        """
        Internal method to reconstruct weights.
        
        Args:
            token_embedding: Token embedding for gating
            head: Optional head index
        
        Returns:
            Reconstructed weight matrix
        """
        return reconstruct_weights(
            dictionary=self.dictionary,
            index_map=self.index_map,
            gating_network=self.gating_network,
            token_embedding=token_embedding,
            head=head,
        )
    
    def get_compression_ratio(self) -> float:
        """
        Calculate compression ratio compared to standard nn.Linear.
        
        Returns:
            Compression ratio (original_size / compressed_size)
        """
        # Original size: weight matrix + bias (if present)
        original_size = (
            self.out_features * self.in_features * 4 +  # float32 weights
            (self.out_features * 4 if self.bias is not None else 0)  # float32 bias
        )
        
        # Compressed size: dictionary + index maps + gating network
        dict_size = self.dict_size * self.in_features * 4  # float32 atoms
        index_map_size = (
            self.out_features * self.in_features * 4 +  # int32 atom indices
            self.out_features * self.in_features * 4    # int32 shift indices
        )
        gating_size = sum(p.numel() * 4 for p in self.gating_network.parameters())  # float32
        
        compressed_size = dict_size + index_map_size + gating_size
        
        return original_size / compressed_size if compressed_size > 0 else float("inf")
    
    def clear_cache(self) -> None:
        """Clear cached weights."""
        self._cached_weights = None
        self._cache_token_embedding = None
    
    def __repr__(self) -> str:
        return (
            f"CyMetaLinear(in_features={self.in_features}, "
            f"out_features={self.out_features}, "
            f"dict_size={self.dict_size}, "
            f"compression={self.get_compression_ratio():.2f}x)"
        )

