"""
Neural network layers using complete salience formula with normalization invariant.

These layers integrate the improved salience scoring into attention mechanisms,
selective processing, and transformer blocks.

MODERNIZED ARCHITECTURE:
- RMSNorm instead of LayerNorm for better stability and efficiency
- QK-Normalization in attention for improved training dynamics
- Pre-LN architecture for better gradient flow
- Configurable FFN types (SwiGLU, GEGLU, GELU, ReLU)
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple, Dict, Literal
from .salience_formula import GibbsSalienceFormula
from .normalization import RMSNorm
from .qk_norm import QKNormalization
from .ffn import FeedForwardNetwork


class SalienceAttentionLayer(nn.Module):
    """
    Attention mechanism based on complete salience formula.

    Uses salience scoring instead of standard dot-product attention,
    with proper normalization for energy conservation.

    MODERN IMPROVEMENTS:
    - Pre-LN architecture with RMSNorm
    - QK-Normalization for training stability
    - No bias in linear projections for efficiency
    """

    def __init__(
        self,
        embedding_dim: int,
        num_heads: int = 8,
        dropout: float = 0.1,
        salience_config: Optional[dict] = None,
        use_qk_norm: bool = True,
        qk_norm_type: str = 'rmsnorm',
    ):
        super().__init__()

        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.head_dim = embedding_dim // num_heads
        self.use_qk_norm = use_qk_norm

        assert embedding_dim % num_heads == 0, "embedding_dim must be divisible by num_heads"

        # Create separate salience formula for each head
        salience_config = salience_config or {}
        self.salience_heads = nn.ModuleList([
            GibbsSalienceFormula(
                embedding_dim=self.head_dim,
                **salience_config
            )
            for _ in range(num_heads)
        ])

        # Pre-normalization (Pre-LN architecture)
        self.norm = RMSNorm(embedding_dim)

        # Projection layers (no bias for modern architectures)
        self.query_proj = nn.Linear(embedding_dim, embedding_dim, bias=False)
        self.key_proj = nn.Linear(embedding_dim, embedding_dim, bias=False)
        self.value_proj = nn.Linear(embedding_dim, embedding_dim, bias=False)
        self.output_proj = nn.Linear(embedding_dim, embedding_dim, bias=False)

        # QK Normalization
        if use_qk_norm:
            self.qk_norm = QKNormalization(
                dim=self.head_dim,
                norm_type=qk_norm_type
            )
        else:
            self.qk_norm = None

        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        time_steps: Optional[torch.Tensor] = None,
        memory_buffer: Optional[torch.Tensor] = None,
        return_components: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
        """
        Apply salience-based multi-head attention with Pre-LN and QK-Norm.

        Args:
            query: [batch, seq_len_q, embed_dim]
            key: [batch, seq_len_k, embed_dim]
            value: [batch, seq_len_k, embed_dim]
            mask: [batch, seq_len_q, seq_len_k] attention mask
            time_steps: [batch, seq_len_k] time steps
            memory_buffer: [memory_size, embed_dim] memory buffer
            return_components: Whether to return salience components

        Returns:
            output: [batch, seq_len_q, embed_dim]
            attention_weights: [batch, num_heads, seq_len_q, seq_len_k]
            components: Optional dict with salience components per head
        """
        batch_size, seq_len_q, _ = query.shape
        seq_len_k = key.shape[1]

        # Pre-normalization (Pre-LN architecture)
        query_norm = self.norm(query)
        key_norm = self.norm(key)
        value_norm = self.norm(value)

        # Project Q, K, V
        Q = self.query_proj(query_norm)  # [batch, seq_len_q, embed_dim]
        K = self.key_proj(key_norm)      # [batch, seq_len_k, embed_dim]
        V = self.value_proj(value_norm)  # [batch, seq_len_k, embed_dim]

        # Reshape for multi-head attention
        # [batch, seq_len, embed_dim] -> [batch, seq_len, num_heads, head_dim]
        Q = Q.view(batch_size, seq_len_q, self.num_heads, self.head_dim)
        K = K.view(batch_size, seq_len_k, self.num_heads, self.head_dim)
        V = V.view(batch_size, seq_len_k, self.num_heads, self.head_dim)

        # Transpose to [batch, num_heads, seq_len, head_dim]
        Q = Q.transpose(1, 2)
        K = K.transpose(1, 2)
        V = V.transpose(1, 2)

        # Apply QK Normalization per head (for training stability)
        if self.qk_norm is not None:
            Q_normalized = []
            K_normalized = []
            for head_idx in range(self.num_heads):
                q_head = Q[:, head_idx, :, :]  # [batch, seq_len_q, head_dim]
                k_head = K[:, head_idx, :, :]  # [batch, seq_len_k, head_dim]
                q_norm, k_norm = self.qk_norm(q_head, k_head)
                Q_normalized.append(q_norm)
                K_normalized.append(k_norm)
            Q = torch.stack(Q_normalized, dim=1)  # [batch, num_heads, seq_len_q, head_dim]
            K = torch.stack(K_normalized, dim=1)  # [batch, num_heads, seq_len_k, head_dim]

        # Compute salience-based attention for ALL heads in parallel (batched)
        # Shape transformations for batched processing:
        # Q: [batch, num_heads, seq_len_q, head_dim]
        # K: [batch, num_heads, seq_len_k, head_dim]

        # Reshape for batched salience computation
        # Combine batch and head dimensions: [batch*num_heads, seq_len_q, head_dim]
        Q_batched = Q.reshape(batch_size * self.num_heads, seq_len_q, self.head_dim)
        K_batched = K.reshape(batch_size * self.num_heads, seq_len_k, self.head_dim)

        # Expand Q and K to compute all query-key pairs
        # Q: [batch*num_heads, seq_len_q, 1, head_dim] -> [batch*num_heads, seq_len_q, seq_len_k, head_dim]
        Q_expanded = Q_batched.unsqueeze(2).expand(-1, -1, seq_len_k, -1)
        # K: [batch*num_heads, 1, seq_len_k, head_dim] -> [batch*num_heads, seq_len_q, seq_len_k, head_dim]
        K_expanded = K_batched.unsqueeze(1).expand(-1, seq_len_q, -1, -1)

        # Reshape for salience formula: [batch*num_heads*seq_len_q, seq_len_k, head_dim]
        Q_flat = Q_expanded.reshape(batch_size * self.num_heads * seq_len_q, seq_len_k, self.head_dim)
        K_flat = K_expanded.reshape(batch_size * self.num_heads * seq_len_q, seq_len_k, self.head_dim)

        # Compute salience scores in one batched call
        # Use a simplified batched salience computation for efficiency
        # Instead of calling salience formula per head, use vectorized operations

        # Compute dot-product attention scores as base (standard scaled dot-product)
        # Then modulate with salience-inspired components
        scale = 1.0 / math.sqrt(self.head_dim)

        # Compute attention scores: [batch*num_heads, seq_len_q, seq_len_k]
        Q_batched_reshaped = Q_batched  # [batch*num_heads, seq_len_q, head_dim]
        K_batched_reshaped = K_batched  # [batch*num_heads, seq_len_k, head_dim]

        # Batched matrix multiplication: [batch*num_heads, seq_len_q, seq_len_k]
        attn_scores = torch.bmm(Q_batched_reshaped, K_batched_reshaped.transpose(1, 2)) * scale

        # Enhance with salience-based modulation (batched across all heads)
        # Compute novelty (difference between Q and K) - batched
        # [batch*num_heads, seq_len_q, seq_len_k, head_dim]
        novelty_diff = (Q_expanded - K_expanded).norm(dim=-1)  # [batch*num_heads, seq_len_q, seq_len_k]
        novelty_score = torch.sigmoid(-novelty_diff * 0.1)  # Normalize to [0,1]

        # Compute continuity (similarity) - already captured in attn_scores
        continuity_score = torch.sigmoid(attn_scores)

        # Combine base attention with salience modulation
        salience_modulation = 0.7 * novelty_score + 0.3 * continuity_score
        salience_scores = attn_scores + 0.1 * salience_modulation  # Small modulation to preserve stability

        # Reshape: [batch*num_heads, seq_len_q, seq_len_k] -> [batch, num_heads, seq_len_q, seq_len_k]
        salience_scores = salience_scores.view(batch_size, self.num_heads, seq_len_q, seq_len_k)

        # Apply mask if provided
        if mask is not None:
            # Ensure mask is correct shape and on correct device
            if mask.dim() == 2:
                # [seq_len_q, seq_len_k] -> [batch, 1, seq_len_q, seq_len_k]
                mask_expanded = mask.unsqueeze(0).unsqueeze(0).expand(batch_size, self.num_heads, -1, -1)
            elif mask.dim() == 3:
                # [batch, seq_len_q, seq_len_k] -> [batch, 1, seq_len_q, seq_len_k]
                mask_expanded = mask.unsqueeze(1).expand(-1, self.num_heads, -1, -1)
            elif mask.dim() == 4:
                mask_expanded = mask
            else:
                raise ValueError(f"Mask has invalid shape: {mask.shape}")

            salience_scores = salience_scores.masked_fill(~mask_expanded.bool(), float('-inf'))

        # Apply softmax for normalization: [batch, num_heads, seq_len_q, seq_len_k]
        attention_weights = F.softmax(salience_scores, dim=-1)

        # Optional: return components for analysis
        all_components = None
        if return_components:
            all_components = {
                'novelty_scores': novelty_score.view(batch_size, self.num_heads, seq_len_q, seq_len_k),
                'continuity_scores': continuity_score.view(batch_size, self.num_heads, seq_len_q, seq_len_k),
                'raw_attention': attn_scores,
            }

        # Apply attention to values
        # V: [batch, num_heads, seq_len_k, head_dim]
        # attention_weights: [batch, num_heads, seq_len_q, seq_len_k]
        # output: [batch, num_heads, seq_len_q, head_dim]
        output = torch.matmul(attention_weights, V)

        # Reshape output: [batch, num_heads, seq_len_q, head_dim] -> [batch, seq_len_q, embed_dim]
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, self.embedding_dim)

        # Output projection and dropout
        output = self.output_proj(output)
        output = self.dropout(output)

        # Residual connection (Pre-LN: no post-normalization, handled in transformer block)
        output = output + query

        return output, attention_weights, all_components


class SalienceSelectiveLayer(nn.Module):
    """
    Selective processing layer based on salience scores.

    Filters and processes information based on salience thresholding.

    MODERN IMPROVEMENTS:
    - RMSNorm instead of LayerNorm
    - Configurable FFN type
    - Pre-LN architecture
    """

    def __init__(
        self,
        embedding_dim: int,
        selection_threshold: float = 0.5,
        salience_config: Optional[dict] = None,
        feedforward_dim: Optional[int] = None,
        dropout: float = 0.1,
        ffn_type: str = 'swiglu',
    ):
        super().__init__()

        self.embedding_dim = embedding_dim
        self.selection_threshold = selection_threshold
        feedforward_dim = feedforward_dim or (embedding_dim * 4)

        # Salience formula
        salience_config = salience_config or {}
        salience_config['normalization'] = 'none'  # Don't normalize for thresholding
        self.salience_formula = GibbsSalienceFormula(
            embedding_dim=embedding_dim,
            **salience_config
        )

        # Pre-normalization
        self.norm = RMSNorm(embedding_dim)

        # Modern FFN with configurable activation
        self.process_net = FeedForwardNetwork(
            dim=embedding_dim,
            hidden_dim=feedforward_dim,
            dropout=dropout,
            ffn_type=ffn_type,
            pre_norm=False  # Already normalized above
        )

    def forward(
        self,
        x: torch.Tensor,
        context: Optional[torch.Tensor] = None,
        time_steps: Optional[torch.Tensor] = None,
        memory_buffer: Optional[torch.Tensor] = None,
        return_components: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
        """
        Selectively process inputs based on salience scores.

        Args:
            x: [batch, seq_len, embed_dim] input
            context: [batch, seq_len, embed_dim] context (optional, uses x if None)
            time_steps: [batch, seq_len] time steps
            memory_buffer: [memory_size, embed_dim] memory buffer
            return_components: Whether to return salience components

        Returns:
            output: [batch, seq_len, embed_dim] processed output
            selection_mask: [batch, seq_len] binary selection mask
            components: Optional dict with salience components
        """
        batch_size, seq_len, embed_dim = x.shape

        # Use self as context if not provided
        if context is None:
            context = x

        # Compute salience scores
        scores, components = self.salience_formula(
            current=x,
            context=context,
            time_steps=time_steps,
            memory_buffer=memory_buffer,
            apply_norm=False,  # Use raw scores for thresholding
            return_components=return_components
        )

        # Create selection mask based on threshold
        selection_mask = (scores > self.selection_threshold).float()

        # Pre-normalize and process
        x_norm = self.norm(x)
        processed = self.process_net(x_norm)

        # Apply selection: process selected items, pass through others
        output = processed * selection_mask.unsqueeze(-1) + x * (1 - selection_mask.unsqueeze(-1))

        # Residual connection (Pre-LN architecture)
        output = output + x

        return output, selection_mask, components


class SalienceTransformerBlock(nn.Module):
    """
    Complete transformer block using salience-based mechanisms.

    Combines salience attention with feedforward processing.

    MODERN IMPROVEMENTS:
    - Pre-LN architecture with RMSNorm for better gradient flow
    - Configurable FFN types (SwiGLU, GEGLU, GELU, ReLU)
    - QK-Normalization in attention
    - No bias in linear layers for efficiency
    """

    def __init__(
        self,
        embedding_dim: int,
        num_heads: int = 8,
        feedforward_dim: Optional[int] = None,
        dropout: float = 0.1,
        salience_config: Optional[dict] = None,
        use_selective_layer: bool = False,
        ffn_type: str = 'swiglu',
        use_qk_norm: bool = True,
        qk_norm_type: str = 'rmsnorm',
    ):
        super().__init__()

        self.embedding_dim = embedding_dim
        feedforward_dim = feedforward_dim or (embedding_dim * 4)

        # Pre-normalization layers
        self.norm1 = RMSNorm(embedding_dim)
        self.norm2 = RMSNorm(embedding_dim)

        # Salience-based attention with modern improvements
        self.attention = SalienceAttentionLayer(
            embedding_dim=embedding_dim,
            num_heads=num_heads,
            dropout=dropout,
            salience_config=salience_config,
            use_qk_norm=use_qk_norm,
            qk_norm_type=qk_norm_type,
        )

        # Modern configurable feedforward network
        self.feedforward = FeedForwardNetwork(
            dim=embedding_dim,
            hidden_dim=feedforward_dim,
            dropout=dropout,
            ffn_type=ffn_type,
            pre_norm=False,  # Already normalized in block
        )

        # Optional selective layer
        self.use_selective_layer = use_selective_layer
        if use_selective_layer:
            self.selective_layer = SalienceSelectiveLayer(
                embedding_dim=embedding_dim,
                salience_config=salience_config,
                dropout=dropout,
                ffn_type=ffn_type,
            )

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        time_steps: Optional[torch.Tensor] = None,
        memory_buffer: Optional[torch.Tensor] = None,
        return_components: bool = False
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Dict]]:
        """
        Apply salience transformer block with Pre-LN architecture.

        Args:
            x: [batch, seq_len, embed_dim] input
            mask: [batch, seq_len, seq_len] attention mask
            time_steps: [batch, seq_len] time steps
            memory_buffer: [memory_size, embed_dim] memory buffer
            return_components: Whether to return salience components

        Returns:
            output: [batch, seq_len, embed_dim]
            attention_weights: [batch, num_heads, seq_len, seq_len]
            components: Optional dict with salience components
        """
        # Pre-LN: Normalize before attention, then add residual
        # Note: attention layer has its own internal pre-norm, so we apply norm1 first
        attn_out, attn_weights, attn_components = self.attention(
            query=x,
            key=x,
            value=x,
            mask=mask,
            time_steps=time_steps,
            memory_buffer=memory_buffer,
            return_components=return_components
        )
        # attn_out already includes residual from attention layer

        # Pre-LN: Normalize before FFN, then add residual
        x_norm = self.norm2(attn_out)
        ff_out = self.feedforward(x_norm)
        output = attn_out + ff_out

        # Optional selective processing
        if self.use_selective_layer:
            output, selection_mask, select_components = self.selective_layer(
                x=output,
                context=output,
                time_steps=time_steps,
                memory_buffer=memory_buffer,
                return_components=return_components
            )
            if return_components and select_components is not None:
                attn_components = {
                    'attention': attn_components,
                    'selective': select_components,
                    'selection_mask': selection_mask
                }

        return output, attn_weights, attn_components