"""
CyMeta Attention Layers.

Drop-in replacements for attention layers that use CyMeta compression.
"""

import torch
import torch.nn as nn
import math
from typing import Optional, Tuple
from cymeta.models.linear import CyMetaLinear


class CyMetaAttention(nn.Module):
    """
    Single-head attention with CyMeta-compressed projections.
    
    This replaces standard Q, K, V, and output projections with CyMeta
    compressed versions, enabling extreme compression while maintaining
    functional capacity.
    
    Args:
        embed_dim: Embedding dimension
        dict_size: Number of dictionary atoms (default: 32)
        shift_bits: Number of bits for shift encoding (default: 8)
        gating_hidden: Hidden dimension for gating network (default: 16)
        dropout: Dropout probability (default: 0.0)
        bias: Whether to include bias in projections (default: True)
    """
    
    def __init__(
        self,
        embed_dim: int,
        dict_size: int = 32,
        shift_bits: int = 8,
        gating_hidden: int = 16,
        dropout: float = 0.0,
        bias: bool = True,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.dict_size = dict_size
        self.shift_bits = shift_bits
        self.gating_hidden = gating_hidden
        
        # Q, K, V projections with CyMeta compression
        self.q_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        self.k_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        self.v_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        
        # Output projection with CyMeta compression
        self.out_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(embed_dim)
    
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        token_embedding: Optional[torch.Tensor] = None,
        attn_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass of attention.
        
        Args:
            query: Query tensor of shape (batch_size, seq_len, embed_dim)
            key: Key tensor of shape (batch_size, seq_len, embed_dim)
            value: Value tensor of shape (batch_size, seq_len, embed_dim)
            token_embedding: Optional token embedding for context-dependent gating
            attn_mask: Optional attention mask
        
        Returns:
            Tuple of (output, attention_weights)
        """
        batch_size, seq_len, _ = query.shape
        
        # Use query as token embedding if not provided
        if token_embedding is None:
            token_embedding = query
        
        # Project Q, K, V
        q = self.q_proj(query, token_embedding=token_embedding)
        k = self.k_proj(key, token_embedding=token_embedding)
        v = self.v_proj(value, token_embedding=token_embedding)
        
        # Compute attention scores
        # q: (batch_size, seq_len, embed_dim)
        # k: (batch_size, seq_len, embed_dim)
        # scores: (batch_size, seq_len, seq_len)
        scores = torch.bmm(q, k.transpose(1, 2)) * self.scale
        
        # Apply attention mask if provided
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask == 0, float("-inf"))
        
        # Softmax
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # Apply attention to values
        # attn_weights: (batch_size, seq_len, seq_len)
        # v: (batch_size, seq_len, embed_dim)
        # output: (batch_size, seq_len, embed_dim)
        output = torch.bmm(attn_weights, v)
        
        # Output projection
        output = self.out_proj(output, token_embedding=token_embedding)
        
        return output, attn_weights


class CyMetaMultiHeadAttention(nn.Module):
    """
    Multi-head attention with CyMeta-compressed projections.
    
    This implements multi-head attention where each head can have its own
    gating network for context-dependent weight reconstruction.
    
    Args:
        embed_dim: Embedding dimension
        num_heads: Number of attention heads
        dict_size: Number of dictionary atoms (default: 32)
        shift_bits: Number of bits for shift encoding (default: 8)
        gating_hidden: Hidden dimension for gating network (default: 16)
        dropout: Dropout probability (default: 0.0)
        bias: Whether to include bias in projections (default: True)
        use_head_specific_gating: Whether to use separate gating per head (default: True)
    """
    
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dict_size: int = 32,
        shift_bits: int = 8,
        gating_hidden: int = 16,
        dropout: float = 0.0,
        bias: bool = True,
        use_head_specific_gating: bool = True,
    ):
        super().__init__()
        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
        
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.dict_size = dict_size
        self.shift_bits = shift_bits
        self.gating_hidden = gating_hidden
        self.use_head_specific_gating = use_head_specific_gating
        
        # Q, K, V projections with CyMeta compression
        # Each projection outputs embed_dim, which is split into num_heads
        self.q_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        self.k_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        self.v_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        
        # Output projection
        self.out_proj = CyMetaLinear(
            embed_dim, embed_dim,
            dict_size=dict_size,
            shift_bits=shift_bits,
            gating_hidden=gating_hidden,
            bias=bias,
        )
        
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(self.head_dim)
    
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        token_embedding: Optional[torch.Tensor] = None,
        attn_mask: Optional[torch.Tensor] = None,
        key_padding_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass of multi-head attention.
        
        Args:
            query: Query tensor of shape (batch_size, seq_len, embed_dim)
            key: Key tensor of shape (batch_size, seq_len, embed_dim)
            value: Value tensor of shape (batch_size, seq_len, embed_dim)
            token_embedding: Optional token embedding for context-dependent gating
            attn_mask: Optional attention mask of shape (seq_len, seq_len)
            key_padding_mask: Optional key padding mask of shape (batch_size, seq_len)
        
        Returns:
            Tuple of (output, attention_weights)
        """
        batch_size, seq_len, _ = query.shape
        
        # Use query as token embedding if not provided
        if token_embedding is None:
            token_embedding = query
        
        # Project Q, K, V
        q = self.q_proj(query, token_embedding=token_embedding)
        k = self.k_proj(key, token_embedding=token_embedding)
        v = self.v_proj(value, token_embedding=token_embedding)
        
        # Reshape for multi-head attention
        # (batch_size, seq_len, embed_dim) -> (batch_size, seq_len, num_heads, head_dim)
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim)
        
        # Transpose for attention computation
        # (batch_size, num_heads, seq_len, head_dim)
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        # Compute attention scores
        # (batch_size, num_heads, seq_len, head_dim) @ (batch_size, num_heads, head_dim, seq_len)
        # -> (batch_size, num_heads, seq_len, seq_len)
        scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        
        # Apply attention masks
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask == 0, float("-inf"))
        
        if key_padding_mask is not None:
            # Expand key_padding_mask: (batch_size, seq_len) -> (batch_size, 1, 1, seq_len)
            scores = scores.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2) == 1,
                float("-inf")
            )
        
        # Softmax
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)
        
        # Apply attention to values
        # (batch_size, num_heads, seq_len, seq_len) @ (batch_size, num_heads, seq_len, head_dim)
        # -> (batch_size, num_heads, seq_len, head_dim)
        output = torch.matmul(attn_weights, v)
        
        # Concatenate heads
        # (batch_size, num_heads, seq_len, head_dim) -> (batch_size, seq_len, embed_dim)
        output = output.transpose(1, 2).contiguous()
        output = output.view(batch_size, seq_len, self.embed_dim)
        
        # Output projection
        output = self.out_proj(output, token_embedding=token_embedding)
        
        return output, attn_weights
    
    def get_compression_ratio(self) -> float:
        """
        Calculate average compression ratio across all projections.
        
        Returns:
            Average compression ratio
        """
        ratios = [
            self.q_proj.get_compression_ratio(),
            self.k_proj.get_compression_ratio(),
            self.v_proj.get_compression_ratio(),
            self.out_proj.get_compression_ratio(),
        ]
        return sum(ratios) / len(ratios)
    
    def __repr__(self) -> str:
        return (
            f"CyMetaMultiHeadAttention(embed_dim={self.embed_dim}, "
            f"num_heads={self.num_heads}, dict_size={self.dict_size}, "
            f"compression={self.get_compression_ratio():.2f}x)"
        )

