#!/usr/bin/env python3
"""
Positional Encoding Module
===========================
Implements learnable positional encoding for sequence models.

This module provides a time-agnostic positional encoding that learns
relative positions in the sequence rather than using fixed sinusoidal
patterns or real timestamp information.

Key Features:
1. Fully learnable position embeddings
2. NO dependency on real timestamps (time-agnostic)
3. Only encodes relative position in sequence
4. Configurable maximum sequence length
5. Optional dropout for regularization

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

import math
from typing import Optional

import torch
import torch.nn as nn


class LearnablePositionalEncoding(nn.Module):
    """
    Learnable Positional Encoding layer.

    Unlike sinusoidal positional encoding, this module learns position
    embeddings directly from data. This is more flexible for tasks where
    the optimal positional representation is unknown.

    IMPORTANT: This encoding is time-AGNOSTIC. It does NOT use any
    real timestamp information. It only encodes the relative position
    of each element within the input sequence (0, 1, 2, ..., seq_len-1).

    This is crucial for our price-focused attention architecture where
    we want the model to learn patterns based on price movements alone,
    without any temporal bias.

    Args:
        d_model: Dimension of the model embeddings (default: 256)
        max_seq_len: Maximum sequence length supported (default: 512)
        dropout: Dropout probability (default: 0.1)

    Example:
        >>> pos_enc = LearnablePositionalEncoding(d_model=256, max_seq_len=512)
        >>> x = torch.randn(32, 100, 256)  # (batch, seq_len, d_model)
        >>> x_encoded = pos_enc(x)
        >>> print(x_encoded.shape)  # (32, 100, 256)
    """

    def __init__(
        self,
        d_model: int = 256,
        max_seq_len: int = 512,
        dropout: float = 0.1
    ) -> None:
        super().__init__()

        self.d_model = d_model
        self.max_seq_len = max_seq_len

        # Learnable position embeddings
        # Shape: (max_seq_len, d_model)
        self.position_embeddings = nn.Parameter(
            torch.zeros(max_seq_len, d_model)
        )

        # Dropout for regularization
        self.dropout = nn.Dropout(p=dropout)

        # Initialize position embeddings
        self._init_embeddings()

    def _init_embeddings(self) -> None:
        """
        Initialize position embeddings using truncated normal distribution.

        Uses a standard deviation of 0.02, which is commonly used for
        transformer position embeddings (as in BERT).
        """
        nn.init.trunc_normal_(self.position_embeddings, std=0.02)

    def forward(
        self,
        x: torch.Tensor,
        offset: int = 0
    ) -> torch.Tensor:
        """
        Add positional encoding to input tensor.

        Args:
            x: Input tensor of shape (batch, seq_len, d_model)
            offset: Starting position offset (useful for incremental decoding)
                   Default is 0, meaning positions start from 0.

        Returns:
            Tensor with positional encoding added, same shape as input

        Raises:
            ValueError: If seq_len + offset exceeds max_seq_len
        """
        batch_size, seq_len, d_model = x.shape

        # Validate sequence length
        if seq_len + offset > self.max_seq_len:
            raise ValueError(
                f"Sequence length ({seq_len}) + offset ({offset}) = {seq_len + offset} "
                f"exceeds maximum sequence length ({self.max_seq_len})"
            )

        # Get position embeddings for the sequence
        # Shape: (seq_len, d_model)
        positions = self.position_embeddings[offset:offset + seq_len]

        # Add positional encoding to input
        # Broadcasting: (batch, seq_len, d_model) + (seq_len, d_model)
        x = x + positions.unsqueeze(0)

        # Apply dropout
        return self.dropout(x)

    def get_position_embedding(self, position: int) -> torch.Tensor:
        """
        Get the embedding for a specific position.

        Args:
            position: Position index (0 to max_seq_len - 1)

        Returns:
            Position embedding tensor of shape (d_model,)
        """
        if position < 0 or position >= self.max_seq_len:
            raise ValueError(
                f"Position {position} out of range [0, {self.max_seq_len})"
            )
        return self.position_embeddings[position]


class SinusoidalPositionalEncoding(nn.Module):
    """
    Sinusoidal Positional Encoding (fixed, not learnable).

    This is the original positional encoding from "Attention Is All You Need".
    Provided for comparison/reference but NOT recommended for price-focused
    attention since sinusoidal patterns may not be optimal for financial data.

    Formula:
        PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
        PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

    Args:
        d_model: Dimension of the model embeddings
        max_seq_len: Maximum sequence length
        dropout: Dropout probability
    """

    def __init__(
        self,
        d_model: int = 256,
        max_seq_len: int = 512,
        dropout: float = 0.1
    ) -> None:
        super().__init__()

        self.d_model = d_model
        self.max_seq_len = max_seq_len
        self.dropout = nn.Dropout(p=dropout)

        # Precompute sinusoidal positional encodings
        pe = self._compute_sinusoidal_encoding(max_seq_len, d_model)

        # Register as buffer (not a parameter, but saved with model)
        self.register_buffer('pe', pe)

    def _compute_sinusoidal_encoding(
        self,
        max_seq_len: int,
        d_model: int
    ) -> torch.Tensor:
        """Compute sinusoidal positional encoding matrix."""
        # Create position indices: (max_seq_len, 1)
        position = torch.arange(max_seq_len, dtype=torch.float).unsqueeze(1)

        # Create dimension indices for division term
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float) *
            (-math.log(10000.0) / d_model)
        )

        # Compute sinusoidal encoding
        pe = torch.zeros(max_seq_len, d_model)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        # Add batch dimension: (1, max_seq_len, d_model)
        return pe.unsqueeze(0)

    def forward(
        self,
        x: torch.Tensor,
        offset: int = 0
    ) -> torch.Tensor:
        """
        Add sinusoidal positional encoding to input tensor.

        Args:
            x: Input tensor of shape (batch, seq_len, d_model)
            offset: Starting position offset

        Returns:
            Tensor with positional encoding added
        """
        seq_len = x.shape[1]

        if seq_len + offset > self.max_seq_len:
            raise ValueError(
                f"Sequence length ({seq_len}) + offset ({offset}) exceeds "
                f"maximum sequence length ({self.max_seq_len})"
            )

        # Add positional encoding (buffer is already on correct device)
        x = x + self.pe[:, offset:offset + seq_len, :]

        return self.dropout(x)


class RelativePositionalEncoding(nn.Module):
    """
    Relative Positional Encoding for attention mechanisms.

    Instead of absolute positions, this encodes the relative distance
    between query and key positions. This is useful when the absolute
    position is less meaningful than the relative order.

    For price sequences, relative position (e.g., "3 candles ago")
    is often more meaningful than absolute position.

    Args:
        d_model: Dimension of position embeddings
        max_distance: Maximum relative distance to encode (default: 128)
                     Distances beyond this are clipped.
        n_heads: Number of attention heads (for head-specific biases)
    """

    def __init__(
        self,
        d_model: int = 256,
        max_distance: int = 128,
        n_heads: int = 8
    ) -> None:
        super().__init__()

        self.d_model = d_model
        self.max_distance = max_distance
        self.n_heads = n_heads

        # Total number of relative positions: [-max_distance, ..., 0, ..., max_distance]
        num_positions = 2 * max_distance + 1

        # Learnable relative position embeddings
        # Shape: (num_positions, d_model)
        self.relative_embeddings = nn.Parameter(
            torch.zeros(num_positions, d_model)
        )

        # Optional: per-head bias for relative positions
        # Shape: (n_heads, num_positions)
        self.head_bias = nn.Parameter(
            torch.zeros(n_heads, num_positions)
        )

        # Initialize
        nn.init.trunc_normal_(self.relative_embeddings, std=0.02)
        nn.init.zeros_(self.head_bias)

    def get_relative_positions(
        self,
        seq_len_q: int,
        seq_len_k: int,
        device: torch.device
    ) -> torch.Tensor:
        """
        Compute relative position matrix.

        Args:
            seq_len_q: Query sequence length
            seq_len_k: Key sequence length
            device: Device for the tensor

        Returns:
            Relative position indices of shape (seq_len_q, seq_len_k)
            Values are in range [0, 2*max_distance] after clipping and shifting
        """
        # Create position indices
        q_positions = torch.arange(seq_len_q, device=device).unsqueeze(1)
        k_positions = torch.arange(seq_len_k, device=device).unsqueeze(0)

        # Compute relative positions (q_pos - k_pos)
        relative_positions = q_positions - k_positions

        # Clip to [-max_distance, max_distance]
        relative_positions = torch.clamp(
            relative_positions,
            -self.max_distance,
            self.max_distance
        )

        # Shift to [0, 2*max_distance] for indexing
        relative_positions = relative_positions + self.max_distance

        return relative_positions

    def forward(
        self,
        seq_len_q: int,
        seq_len_k: int,
        device: torch.device
    ) -> torch.Tensor:
        """
        Get relative position embeddings for attention.

        Args:
            seq_len_q: Query sequence length
            seq_len_k: Key sequence length
            device: Device for the tensor

        Returns:
            Relative position embeddings of shape (seq_len_q, seq_len_k, d_model)
        """
        # Get relative position indices
        rel_pos = self.get_relative_positions(seq_len_q, seq_len_k, device)

        # Gather embeddings
        # rel_pos: (seq_len_q, seq_len_k)
        # relative_embeddings: (num_positions, d_model)
        embeddings = self.relative_embeddings[rel_pos]

        return embeddings

    def get_attention_bias(
        self,
        seq_len_q: int,
        seq_len_k: int,
        device: torch.device
    ) -> torch.Tensor:
        """
        Get per-head attention bias based on relative positions.

        This bias is added to attention scores before softmax.

        Args:
            seq_len_q: Query sequence length
            seq_len_k: Key sequence length
            device: Device for the tensor

        Returns:
            Attention bias of shape (n_heads, seq_len_q, seq_len_k)
        """
        # Get relative position indices
        rel_pos = self.get_relative_positions(seq_len_q, seq_len_k, device)

        # Gather per-head biases
        # head_bias: (n_heads, num_positions)
        # rel_pos: (seq_len_q, seq_len_k)
        bias = self.head_bias[:, rel_pos.view(-1)].view(
            self.n_heads, seq_len_q, seq_len_k
        )

        return bias
