#!/usr/bin/env python3
"""
Multi-Head Attention Implementation
====================================
Implements scaled dot-product multi-head attention mechanism
following the Transformer architecture (Vaswani et al., 2017).

This module provides a configurable multi-head attention layer
suitable for financial time series analysis.

Key Features:
1. Configurable number of heads, d_model, d_k, d_v
2. Returns both output and attention weights
3. Supports optional masking
4. Uses einsum operations for clarity

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

import math
from typing import Tuple, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):
    """
    Multi-Head Attention mechanism with scaled dot-product attention.

    Implements the attention mechanism from "Attention Is All You Need":
        Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) * V

    Then applies multi-head split and linear projections:
        MultiHead(Q, K, V) = Concat(head_1, ..., head_h) * W^O
        where head_i = Attention(Q*W_i^Q, K*W_i^K, V*W_i^V)

    Args:
        d_model: Total dimension of the model (default: 256)
        n_heads: Number of attention heads (default: 8)
        d_k: Dimension of keys/queries per head (default: None, computed as d_model // n_heads)
        d_v: Dimension of values per head (default: None, same as d_k)
        dropout: Dropout probability on attention weights (default: 0.1)
        bias: Whether to use bias in linear projections (default: True)

    Example:
        >>> mha = MultiHeadAttention(d_model=256, n_heads=8)
        >>> x = torch.randn(32, 100, 256)  # (batch, seq_len, d_model)
        >>> output, attn_weights = mha(x, x, x)
        >>> print(output.shape)  # (32, 100, 256)
        >>> print(attn_weights.shape)  # (32, 8, 100, 100)
    """

    def __init__(
        self,
        d_model: int = 256,
        n_heads: int = 8,
        d_k: Optional[int] = None,
        d_v: Optional[int] = None,
        dropout: float = 0.1,
        bias: bool = True
    ) -> None:
        super().__init__()

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_k if d_k is not None else d_model // n_heads
        self.d_v = d_v if d_v is not None else self.d_k

        # Validate dimensions
        if self.d_k * n_heads != d_model:
            raise ValueError(
                f"d_k ({self.d_k}) * n_heads ({n_heads}) must equal d_model ({d_model}) "
                f"when using standard multi-head attention. "
                f"Got {self.d_k * n_heads} != {d_model}"
            )

        # Linear projections for Q, K, V
        self.w_q = nn.Linear(d_model, n_heads * self.d_k, bias=bias)
        self.w_k = nn.Linear(d_model, n_heads * self.d_k, bias=bias)
        self.w_v = nn.Linear(d_model, n_heads * self.d_v, bias=bias)

        # Output projection
        self.w_o = nn.Linear(n_heads * self.d_v, d_model, bias=bias)

        # Dropout on attention weights
        self.dropout = nn.Dropout(dropout)

        # Scaling factor for dot-product attention
        self.scale = math.sqrt(self.d_k)

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize weights using Xavier uniform initialization."""
        nn.init.xavier_uniform_(self.w_q.weight)
        nn.init.xavier_uniform_(self.w_k.weight)
        nn.init.xavier_uniform_(self.w_v.weight)
        nn.init.xavier_uniform_(self.w_o.weight)

        if self.w_q.bias is not None:
            nn.init.zeros_(self.w_q.bias)
            nn.init.zeros_(self.w_k.bias)
            nn.init.zeros_(self.w_v.bias)
            nn.init.zeros_(self.w_o.bias)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        return_attention: bool = True
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Forward pass of multi-head attention.

        Args:
            query: Query tensor of shape (batch, seq_len_q, d_model)
            key: Key tensor of shape (batch, seq_len_k, d_model)
            value: Value tensor of shape (batch, seq_len_v, d_model)
                   Note: seq_len_k must equal seq_len_v
            mask: Optional attention mask of shape (batch, 1, 1, seq_len_k)
                  or (batch, 1, seq_len_q, seq_len_k) or (batch, n_heads, seq_len_q, seq_len_k)
                  Positions with True/1 will be masked (set to -inf before softmax)
            return_attention: Whether to return attention weights (default: True)

        Returns:
            output: Attended output of shape (batch, seq_len_q, d_model)
            attention_weights: Attention weights of shape (batch, n_heads, seq_len_q, seq_len_k)
                              Returns None if return_attention is False
        """
        batch_size, seq_len_q, _ = query.shape
        seq_len_k = key.shape[1]
        seq_len_v = value.shape[1]

        # Validate key and value sequence lengths match
        if seq_len_k != seq_len_v:
            raise ValueError(
                f"Key sequence length ({seq_len_k}) must match value sequence length ({seq_len_v})"
            )

        # Project to Q, K, V
        # Shape: (batch, seq_len, n_heads * d_k) -> (batch, seq_len, n_heads, d_k)
        q = self.w_q(query).view(batch_size, seq_len_q, self.n_heads, self.d_k)
        k = self.w_k(key).view(batch_size, seq_len_k, self.n_heads, self.d_k)
        v = self.w_v(value).view(batch_size, seq_len_v, self.n_heads, self.d_v)

        # Transpose for attention: (batch, n_heads, seq_len, d_k/d_v)
        q = q.transpose(1, 2)  # (batch, n_heads, seq_len_q, d_k)
        k = k.transpose(1, 2)  # (batch, n_heads, seq_len_k, d_k)
        v = v.transpose(1, 2)  # (batch, n_heads, seq_len_v, d_v)

        # Compute attention scores: QK^T / sqrt(d_k)
        # (batch, n_heads, seq_len_q, d_k) @ (batch, n_heads, d_k, seq_len_k)
        # -> (batch, n_heads, seq_len_q, seq_len_k)
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale

        # Apply mask if provided
        if mask is not None:
            # Convert boolean mask to float mask if needed
            if mask.dtype == torch.bool:
                scores = scores.masked_fill(mask, float('-inf'))
            else:
                # Assume mask contains 0s for keep and 1s for mask
                scores = scores.masked_fill(mask > 0.5, float('-inf'))

        # Apply softmax to get attention weights
        attention_weights = F.softmax(scores, dim=-1)

        # Apply dropout to attention weights
        attention_weights = self.dropout(attention_weights)

        # Apply attention to values
        # (batch, n_heads, seq_len_q, seq_len_k) @ (batch, n_heads, seq_len_v, d_v)
        # -> (batch, n_heads, seq_len_q, d_v)
        attended = torch.matmul(attention_weights, v)

        # Transpose and reshape: (batch, n_heads, seq_len_q, d_v) -> (batch, seq_len_q, n_heads * d_v)
        attended = attended.transpose(1, 2).contiguous().view(batch_size, seq_len_q, -1)

        # Final linear projection
        output = self.w_o(attended)

        if return_attention:
            return output, attention_weights
        else:
            return output, None


class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention (single head).

    This is the core attention mechanism without multi-head split.
    Useful for understanding or debugging attention patterns.

    Computes:
        Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) * V

    Args:
        d_k: Dimension of keys/queries (for scaling)
        dropout: Dropout probability (default: 0.1)
    """

    def __init__(self, d_k: int, dropout: float = 0.1) -> None:
        super().__init__()
        self.scale = math.sqrt(d_k)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Compute scaled dot-product attention.

        Args:
            q: Query tensor (batch, ..., seq_len_q, d_k)
            k: Key tensor (batch, ..., seq_len_k, d_k)
            v: Value tensor (batch, ..., seq_len_v, d_v)
            mask: Optional mask (batch, ..., seq_len_q, seq_len_k)

        Returns:
            output: Attended values (batch, ..., seq_len_q, d_v)
            attention_weights: Attention weights (batch, ..., seq_len_q, seq_len_k)
        """
        # Compute attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale

        # Apply mask if provided
        if mask is not None:
            if mask.dtype == torch.bool:
                scores = scores.masked_fill(mask, float('-inf'))
            else:
                scores = scores.masked_fill(mask > 0.5, float('-inf'))

        # Softmax and dropout
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # Apply attention to values
        output = torch.matmul(attention_weights, v)

        return output, attention_weights


def create_causal_mask(seq_len: int, device: torch.device = None) -> torch.Tensor:
    """
    Create a causal (lower triangular) attention mask.

    This mask prevents attention to future positions, suitable for
    autoregressive models or sequence prediction tasks.

    Args:
        seq_len: Length of the sequence
        device: Device to create the mask on

    Returns:
        mask: Boolean mask of shape (1, 1, seq_len, seq_len)
              True values indicate positions to mask (future positions)
    """
    # Create upper triangular matrix (True = mask)
    mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1)
    # Add batch and head dimensions
    return mask.unsqueeze(0).unsqueeze(0)


def create_padding_mask(
    lengths: torch.Tensor,
    max_len: int,
    device: torch.device = None
) -> torch.Tensor:
    """
    Create a padding mask from sequence lengths.

    Args:
        lengths: Tensor of shape (batch,) containing sequence lengths
        max_len: Maximum sequence length
        device: Device to create the mask on

    Returns:
        mask: Boolean mask of shape (batch, 1, 1, max_len)
              True values indicate padding positions to mask
    """
    batch_size = lengths.shape[0]
    # Create position indices
    positions = torch.arange(max_len, device=device).unsqueeze(0).expand(batch_size, -1)
    # Create mask where positions >= lengths
    mask = positions >= lengths.unsqueeze(1)
    # Add head and query dimensions
    return mask.unsqueeze(1).unsqueeze(2)
