"""Dual Brain Architecture: Encoder-Decoder Hybrid for Reasoning + Generation.

Combines:
- Encoder (left brain): Bidirectional understanding, analysis, reasoning
- Decoder (right brain): Autoregressive generation, creativity

Each can operate independently or together for complex tasks.
"""

import torch
import torch.nn as nn
from typing import Optional, Literal, Tuple, Union, List
from composennent.models.nlp_base import BaseLanguageModel
from composennent.modules.attention import MultiHeadAttention, causal_mask


class CrossAttentionLayer(nn.Module):
    """Cross-attention layer for decoder to attend to encoder outputs.
    
    Args:
        dim: Model dimension.
        num_heads: Number of attention heads.
        dropout: Dropout probability.
    """
    def __init__(self, dim: int, num_heads: int, dropout: float = 0.1):
        super().__init__()
        self.cross_attn = MultiHeadAttention(
            embed_dim=dim,
            num_heads=num_heads,
            dropout=dropout,
        )
        self.norm = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(
        self,
        x: torch.Tensor,
        encoder_output: torch.Tensor,
        encoder_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Apply cross-attention.
        
        Args:
            x: Decoder hidden states (batch, dec_seq, dim).
            encoder_output: Encoder output (batch, enc_seq, dim).
            encoder_mask: Mask for encoder padding.
            
        Returns:
            Updated decoder states.
        """
        normed = self.norm(x)
        attn_out, _, _ = self.cross_attn(
            normed,
            key=encoder_output,
            value=encoder_output,
            key_padding_mask=encoder_mask,
        )
        return x + self.dropout(attn_out)


class EncoderBlock(nn.Module):
    """Bidirectional encoder block (BERT-style)."""
    
    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: int = 4,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.self_attn = MultiHeadAttention(dim, num_heads, dropout)
        self.norm1 = nn.LayerNorm(dim)
        
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim * mlp_ratio, dim),
        )
        self.norm2 = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # Self-attention (bidirectional - no causal mask)
        normed = self.norm1(x)
        attn_out, _, _ = self.self_attn(normed, key_padding_mask=mask)
        x = x + self.dropout(attn_out)
        
        # MLP
        normed = self.norm2(x)
        x = x + self.dropout(self.mlp(normed))
        
        return x


class DecoderBlock(nn.Module):
    """Causal decoder block with optional cross-attention."""
    
    def __init__(
        self,
        dim: int,
        num_heads: int,
        mlp_ratio: int = 4,
        dropout: float = 0.1,
        use_cross_attention: bool = True,
    ):
        super().__init__()
        self.self_attn = MultiHeadAttention(dim, num_heads, dropout)
        self.norm1 = nn.LayerNorm(dim)
        
        self.use_cross_attention = use_cross_attention
        if use_cross_attention:
            self.cross_attn = CrossAttentionLayer(dim, num_heads, dropout)
        
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim * mlp_ratio),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim * mlp_ratio, dim),
        )
        self.norm2 = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(
        self,
        x: torch.Tensor,
        causal_mask: Optional[torch.Tensor] = None,
        encoder_output: Optional[torch.Tensor] = None,
        encoder_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        """Forward pass with optional KV-cache support.
        
        Args:
            x: Input tensor (batch, seq, dim).
            causal_mask: Causal attention mask.
            encoder_output: Encoder output for cross-attention.
            encoder_mask: Encoder padding mask.
            past_key_value: Cached (key, value) from previous forward.
            use_cache: Whether to return updated cache.
            
        Returns:
            Tuple of (output, present_key_value).
        """
        # Causal self-attention with KV-cache
        normed = self.norm1(x)
        attn_out, _, present_key_value = self.self_attn(
            normed,
            attn_mask=causal_mask,
            past_key_value=past_key_value,
            use_cache=use_cache,
        )
        x = x + self.dropout(attn_out)
        
        # Cross-attention to encoder (if available)
        if self.use_cross_attention and encoder_output is not None:
            x = self.cross_attn(x, encoder_output, encoder_mask)
        
        # MLP
        normed = self.norm2(x)
        x = x + self.dropout(self.mlp(normed))
        
        return x, present_key_value


class DualBrainModel(BaseLanguageModel):
    """Encoder-Decoder Hybrid Architecture.
    
    Inspired by the human brain's hemispheric specialization:
    - Left brain (Encoder): Analysis, logic, structured reasoning
    - Right brain (Decoder): Creativity, generation, intuition
    
    Can operate in three modes:
    - "encoder_only": Just encoding (for classification, embedding)
    - "decoder_only": Just decoding (standard GPT-style generation)
    - "full": Both encoder and decoder with cross-attention (reasoning + generation)
    
    Args:
        vocab_size: Vocabulary size.
        latent_dim: Model dimension.
        encoder_layers: Number of encoder layers.
        decoder_layers: Number of decoder layers.
        num_heads: Number of attention heads.
        max_seq_len: Maximum sequence length.
        mlp_ratio: MLP expansion ratio.
        dropout: Dropout probability.
        tie_weights: Tie embedding weights between encoder/decoder.
        
    Example:
        >>> model = DualBrainModel(
        ...     vocab_size=32000,
        ...     latent_dim=768,
        ...     encoder_layers=6,
        ...     decoder_layers=6,
        ...     num_heads=12,
        ... )
        >>> 
        >>> # Full reasoning mode
        >>> output = model(input_ids, mode="full")
        >>> 
        >>> # Encoder only (for embeddings/classification)
        >>> embeddings = model(input_ids, mode="encoder_only")
        >>> 
        >>> # Decoder only (standard generation)
        >>> logits = model(input_ids, mode="decoder_only")
    """
    
    def __init__(
        self,
        vocab_size: int,
        latent_dim: int,
        encoder_layers: int = 6,
        decoder_layers: int = 6,
        num_heads: int = 8,
        max_seq_len: int = 512,
        mlp_ratio: int = 4,
        dropout: float = 0.1,
        tie_weights: bool = True,
    ):
        super().__init__()
        
        self.vocab_size = vocab_size
        self.latent_dim = latent_dim
        self.encoder_layers = encoder_layers
        self.decoder_layers = decoder_layers
        self.num_heads = num_heads
        self.max_seq_len = max_seq_len
        self.mlp_ratio = mlp_ratio
        self.dropout_rate = dropout
        
        # Shared embeddings
        self.token_embedding = nn.Embedding(vocab_size, latent_dim)
        self.position_embedding = nn.Embedding(max_seq_len, latent_dim)
        self.dropout = nn.Dropout(dropout)
        
        # Encoder (Left Brain)
        self.encoder = nn.ModuleList([
            EncoderBlock(latent_dim, num_heads, mlp_ratio, dropout)
            for _ in range(encoder_layers)
        ])
        self.encoder_norm = nn.LayerNorm(latent_dim)
        
        # Decoder (Right Brain)
        self.decoder = nn.ModuleList([
            DecoderBlock(latent_dim, num_heads, mlp_ratio, dropout, use_cross_attention=True)
            for _ in range(decoder_layers)
        ])
        self.decoder_norm = nn.LayerNorm(latent_dim)
        
        # Output projection
        self.lm_head = nn.Linear(latent_dim, vocab_size, bias=False)
        
        # Weight tying
        if tie_weights:
            self.lm_head.weight = self.token_embedding.weight
        
        self.init_weights()
    
    def get_config(self) -> dict:
        return {
            "vocab_size": self.vocab_size,
            "latent_dim": self.latent_dim,
            "encoder_layers": self.encoder_layers,
            "decoder_layers": self.decoder_layers,
            "num_heads": self.num_heads,
            "max_seq_len": self.max_seq_len,
            "mlp_ratio": self.mlp_ratio,
            "drop_out": self.dropout_rate,
        }
    
    def embed(self, input_ids: torch.Tensor) -> torch.Tensor:
        """Get embeddings for input tokens."""
        seq_len = input_ids.size(1)
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        return self.dropout(x)
    
    def encode(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Run encoder (left brain) only.
        
        Args:
            input_ids: Token IDs.
            attention_mask: Padding mask (1=attend, 0=ignore).
            
        Returns:
            Encoder hidden states.
        """
        x = self.embed(input_ids)
        
        # Convert mask format
        key_padding_mask = None
        if attention_mask is not None:
            key_padding_mask = attention_mask == 0
        
        for layer in self.encoder:
            x = layer(x, mask=key_padding_mask)
        
        return self.encoder_norm(x)
    
    def decode(
        self,
        input_ids: torch.Tensor,
        encoder_output: Optional[torch.Tensor] = None,
        encoder_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[List[Tuple[torch.Tensor, torch.Tensor]]]]:
        """Run decoder (right brain) with optional KV-cache.
        
        Args:
            input_ids: Token IDs for decoder.
            encoder_output: Optional encoder hidden states for cross-attention.
            encoder_mask: Mask for encoder outputs.
            past_key_values: List of (key, value) tuples for each layer.
            use_cache: Whether to return updated KV-cache.
            
        Returns:
            Tuple of (hidden_states, present_key_values).
        """
        x = self.embed(input_ids)
        seq_len = input_ids.size(1)
        
        # Calculate total sequence length for causal mask
        past_len = 0
        if past_key_values is not None and past_key_values[0] is not None:
            past_len = past_key_values[0][0].size(2)
        total_len = past_len + seq_len
        
        # Create causal mask for the new tokens only
        # New tokens can attend to all past tokens + themselves
        mask = causal_mask(total_len, input_ids.device)
        # Only need the last seq_len rows (for the new tokens)
        mask = mask[-seq_len:, :] if seq_len < total_len else mask
        
        present_key_values = [] if use_cache else None
        
        for i, layer in enumerate(self.decoder):
            layer_past = None
            if past_key_values is not None:
                layer_past = past_key_values[i] if i < len(past_key_values) else None
            
            x, present_kv = layer(
                x,
                causal_mask=mask,
                encoder_output=encoder_output,
                encoder_mask=encoder_mask,
                past_key_value=layer_past,
                use_cache=use_cache,
            )
            
            if use_cache:
                present_key_values.append(present_kv)
        
        return self.decoder_norm(x), present_key_values
    
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        mode: Literal["encoder_only", "decoder_only", "full"] = "full",
        **kwargs,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """Forward pass with mode selection.
        
        Args:
            input_ids: Token IDs.
            attention_mask: Attention mask.
            mode: Operation mode:
                - "encoder_only": Returns encoder hidden states
                - "decoder_only": Returns logits (standard generation)
                - "full": Encoder + decoder with cross-attention
                
        Returns:
            Depends on mode:
            - encoder_only: Hidden states (batch, seq, dim)
            - decoder_only: Logits (batch, seq, vocab)
            - full: Logits (batch, seq, vocab)
        """
        if mode == "encoder_only":
            return self.encode(input_ids, attention_mask)
        
        elif mode == "decoder_only":
            hidden, _ = self.decode(input_ids)
            return self.lm_head(hidden)
        
        else:  # full mode
            # Encoder processes input
            encoder_output = self.encode(input_ids, attention_mask)
            
            # Decoder generates with cross-attention to encoder
            decoder_hidden, _ = self.decode(
                input_ids,
                encoder_output=encoder_output,
                encoder_mask=attention_mask == 0 if attention_mask is not None else None,
            )
            
            return self.lm_head(decoder_hidden)
    
    def generate(
        self,
        prompt: str,
        tokenizer,
        max_length: int = 100,
        mode: Literal["decoder_only", "full"] = "full",
        temperature: float = 1.0,
        top_k: int = 50,
        top_p: float = 0.9,
        use_cache: bool = True,
        **kwargs,
    ) -> str:
        """Generate text with KV-cache for efficient autoregressive generation.
        
        Uses KV-caching to achieve O(N) generation complexity instead of O(N²).
        On first step, processes full prompt and caches K/V states.
        On subsequent steps, only processes the new token using cached K/V.
        
        Args:
            prompt: Input text.
            tokenizer: Tokenizer instance.
            max_length: Maximum generation length.
            mode: "decoder_only" or "full" (with reasoning).
            temperature: Sampling temperature.
            top_k: Top-k filtering.
            top_p: Nucleus sampling threshold.
            use_cache: Whether to use KV-caching (recommended for speed).
            
        Returns:
            Generated text.
        """
        self.eval()
        device = next(self.parameters()).device
        
        # Tokenize
        input_ids = tokenizer.encode(prompt)
        input_ids = torch.tensor([input_ids], device=device)
        
        # Get encoder output for full mode
        encoder_output = None
        if mode == "full":
            with torch.no_grad():
                encoder_output = self.encode(input_ids)
        
        generated = input_ids.tolist()[0]
        past_key_values = None
        
        with torch.no_grad():
            for step in range(max_length):
                # First step: process full prompt
                # Subsequent steps: only process last token (with cache)
                if step == 0 or not use_cache:
                    curr_input = torch.tensor([generated], device=device)
                else:
                    # Only process the new token
                    curr_input = torch.tensor([[generated[-1]]], device=device)
                
                if mode == "full" and encoder_output is not None:
                    hidden, past_key_values = self.decode(
                        curr_input,
                        encoder_output=encoder_output,
                        past_key_values=past_key_values if use_cache else None,
                        use_cache=use_cache,
                    )
                else:
                    hidden, past_key_values = self.decode(
                        curr_input,
                        past_key_values=past_key_values if use_cache else None,
                        use_cache=use_cache,
                    )
                
                logits = self.lm_head(hidden)[:, -1, :] / temperature
                
                # Top-k filtering
                if top_k > 0:
                    indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
                    logits[indices_to_remove] = float('-inf')
                
                # Top-p filtering
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                    sorted_indices_to_remove[..., 0] = 0
                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                    logits[indices_to_remove] = float('-inf')
                
                probs = torch.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1).item()
                
                generated.append(next_token)
                
                if hasattr(tokenizer, 'eos_id') and next_token == tokenizer.eos_id:
                    break
        
        return tokenizer.decode(generated)
