import math
from typing import Tuple, Union

import mlx.core as mx
import mlx.nn as nn


class PositionalEncoding(nn.Module):
    """
    Sinusoidal positional encoding.

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

    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
        """
        Args:
            d_model: Embedding dimension
            dropout_rate: Dropout rate
            max_len: Maximum input length
        """
        super().__init__()
        self.d_model = d_model
        self.xscale = math.sqrt(self.d_model)
        self.dropout_rate = dropout_rate
        self.max_len = max_len

        # Precompute positional encodings
        self.pe = self._create_pe(max_len, d_model)

    def _create_pe(self, max_len: int, d_model: int) -> mx.array:
        """Create positional encoding matrix."""
        pe = mx.zeros((max_len, d_model))
        position = mx.arange(0, max_len, dtype=mx.float32)[:, None]
        div_term = mx.exp(
            mx.arange(0, d_model, 2, dtype=mx.float32) * -(math.log(10000.0) / d_model)
        )

        pe_sin = mx.sin(position * div_term)
        pe_cos = mx.cos(position * div_term)

        # Interleave sin and cos
        pe = mx.zeros((max_len, d_model))
        pe = pe.at[:, 0::2].add(pe_sin)
        pe = pe.at[:, 1::2].add(pe_cos)

        return mx.expand_dims(pe, 0)  # (1, max_len, d_model)

    def __call__(self, x: mx.array, offset: int = 0) -> Tuple[mx.array, mx.array]:
        """
        Add positional encoding.

        Args:
            x: Input tensor (B, T, D)
            offset: Position offset

        Returns:
            x: Encoded input (B, T, D)
            pos_emb: Positional embeddings (1, T, D)
        """
        pos_emb = self.position_encoding(offset, x.shape[1])
        x = x * self.xscale + pos_emb

        if self.training and self.dropout_rate > 0:
            x = nn.Dropout(self.dropout_rate)(x)
            pos_emb = nn.Dropout(self.dropout_rate)(pos_emb)

        return x, pos_emb

    def position_encoding(self, offset: int, size: int) -> mx.array:
        """Get positional encoding for a range."""
        assert offset + size <= self.max_len
        return self.pe[:, offset : offset + size, :]


class RelPositionalEncoding(PositionalEncoding):
    """
    Relative positional encoding module.

    See Appendix B in https://arxiv.org/abs/1901.02860

    Unlike PositionalEncoding, this does NOT add pos_emb to input.
    The pos_emb is returned separately for use in relative attention.
    """

    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
        super().__init__(d_model, dropout_rate, max_len)

    def __call__(self, x: mx.array, offset: int = 0) -> Tuple[mx.array, mx.array]:
        """
        Compute positional encoding.

        Args:
            x: Input tensor (B, T, D)
            offset: Position offset

        Returns:
            x: Scaled input (B, T, D) - NOT added to pos_emb
            pos_emb: Positional embeddings (1, T, D)
        """
        x = x * self.xscale
        pos_emb = self.position_encoding(offset, x.shape[1])

        if self.training and self.dropout_rate > 0:
            x = nn.Dropout(self.dropout_rate)(x)
            pos_emb = nn.Dropout(self.dropout_rate)(pos_emb)

        return x, pos_emb


class EspnetRelPositionalEncoding(nn.Module):
    """
    Relative positional encoding module (ESPnet implementation).

    This version computes both positive and negative position encodings
    for bidirectional relative attention.

    See https://github.com/espnet/espnet/pull/2816
    and Appendix B in https://arxiv.org/abs/1901.02860
    """

    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
        super().__init__()
        self.d_model = d_model
        self.xscale = math.sqrt(self.d_model)
        self.dropout_rate = dropout_rate
        self.max_len = max_len

        # Initialize PE - will be extended as needed
        self.pe = None
        self._extend_pe(max_len)

    def _extend_pe(self, size: int):
        """Create or extend positional encodings."""
        if self.pe is not None and self.pe.shape[1] >= size * 2 - 1:
            return

        # Create positive and negative position encodings
        position = mx.arange(0, size, dtype=mx.float32)[:, None]
        div_term = mx.exp(
            mx.arange(0, self.d_model, 2, dtype=mx.float32)
            * -(math.log(10000.0) / self.d_model)
        )

        # Positive positions
        pe_positive_sin = mx.sin(position * div_term)
        pe_positive_cos = mx.cos(position * div_term)
        pe_positive = mx.zeros((size, self.d_model))
        pe_positive = pe_positive.at[:, 0::2].add(pe_positive_sin)
        pe_positive = pe_positive.at[:, 1::2].add(pe_positive_cos)

        # Negative positions
        pe_negative_sin = mx.sin(-1 * position * div_term)
        pe_negative_cos = mx.cos(-1 * position * div_term)
        pe_negative = mx.zeros((size, self.d_model))
        pe_negative = pe_negative.at[:, 0::2].add(pe_negative_sin)
        pe_negative = pe_negative.at[:, 1::2].add(pe_negative_cos)

        # Reverse positive and concatenate
        # pe_positive: [0, 1, 2, ...] -> reversed: [..., 2, 1, 0]
        # Use slicing to reverse since MLX doesn't have flip
        pe_positive = mx.expand_dims(pe_positive[::-1], 0)
        # pe_negative: skip first (which is 0) -> [1, 2, 3, ...]
        pe_negative = mx.expand_dims(pe_negative[1:], 0)

        # Concatenate: [..., 2, 1, 0, 1, 2, ...]
        self.pe = mx.concatenate([pe_positive, pe_negative], axis=1)

    def __call__(self, x: mx.array, offset: int = 0) -> Tuple[mx.array, mx.array]:
        """
        Add positional encoding.

        Args:
            x: Input tensor (B, T, D)
            offset: Position offset (not used in this implementation)

        Returns:
            x: Scaled input (B, T, D)
            pos_emb: Positional embeddings (1, 2*T-1, D)
        """
        self._extend_pe(x.shape[1])
        x = x * self.xscale
        pos_emb = self.position_encoding(x.shape[1], offset)

        if self.training and self.dropout_rate > 0:
            x = nn.Dropout(self.dropout_rate)(x)
            pos_emb = nn.Dropout(self.dropout_rate)(pos_emb)

        return x, pos_emb

    def position_encoding(self, size: int, offset: int = 0) -> mx.array:
        """
        Get positional encoding for relative attention.

        Args:
            size: Required size
            offset: Position offset (not used)

        Returns:
            pos_emb: Positional embeddings (1, 2*size-1, d_model)
        """
        # Extract from center: positions from -(size-1) to +(size-1)
        center = self.pe.shape[1] // 2
        start = center - size + 1
        end = center + size
        return self.pe[:, start:end, :]


class NoPositionalEncoding(nn.Module):
    """No positional encoding - returns zeros."""

    def __init__(self, d_model: int, dropout_rate: float):
        super().__init__()
        self.d_model = d_model
        self.dropout_rate = dropout_rate

    def __call__(self, x: mx.array, offset: int = 0) -> Tuple[mx.array, mx.array]:
        """Return input unchanged with zero positional embeddings."""
        pos_emb = mx.zeros((1, x.shape[1], self.d_model))

        if self.training and self.dropout_rate > 0:
            x = nn.Dropout(self.dropout_rate)(x)

        return x, pos_emb

    def position_encoding(self, offset: int, size: int) -> mx.array:
        """Return zero positional encoding."""
        return mx.zeros((1, size, self.d_model))
