"""Transformer encoder/decoder layers for sign language segmentation.

Adapted from segmenterv2/transformer/ with VLT-compatible imports.
"""

from __future__ import annotations

import math

import torch
from torch import Tensor, nn

# ---------------------------------------------------------------------------
# Helpers (inlined from segmenterv2/helpers.py)
# ---------------------------------------------------------------------------


class ConfigurationError(Exception):
    """Custom exception for misspecifications of configuration."""


def freeze_params(module: nn.Module) -> None:
    """Freeze all parameters of a module."""
    for _, p in module.named_parameters():
        p.requires_grad = False


# ---------------------------------------------------------------------------
# Activation builder
# ---------------------------------------------------------------------------


def build_activation(activation: str = "relu"):
    """Return the activation function class for the given name."""
    if activation == "relu":
        return nn.ReLU
    elif activation == "gelu":
        return nn.GELU
    elif activation == "tanh":
        return torch.tanh
    elif activation == "swish":
        return nn.SiLU
    else:
        raise ConfigurationError(
            "Invalid activation function. Valid options: " "'relu', 'gelu', 'tanh', 'swish'."
        )


# ---------------------------------------------------------------------------
# Multi-Headed Attention
# ---------------------------------------------------------------------------


class MultiHeadedAttention(nn.Module):
    """Multi-Head Attention module from 'Attention is All You Need'."""

    def __init__(self, num_heads: int, size: int, dropout: float = 0.1) -> None:
        super().__init__()
        assert size % num_heads == 0

        self.head_size = head_size = size // num_heads
        self.model_size = size
        self.num_heads = num_heads

        self.k_layer = nn.Linear(size, num_heads * head_size)
        self.v_layer = nn.Linear(size, num_heads * head_size)
        self.q_layer = nn.Linear(size, num_heads * head_size)

        self.output_layer = nn.Linear(size, size)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        k: Tensor,
        v: Tensor,
        q: Tensor,
        mask: Tensor | None = None,
        return_weights: bool | None = None,
    ):
        batch_size = k.size(0)
        key_len = k.size(1)
        query_len = q.size(1)

        k = self.k_layer(k)
        v = self.v_layer(v)
        q = self.q_layer(q)

        k = k.view(batch_size, -1, self.num_heads, self.head_size).transpose(1, 2)
        v = v.view(batch_size, -1, self.num_heads, self.head_size).transpose(1, 2)
        q = q.view(batch_size, -1, self.num_heads, self.head_size).transpose(1, 2)

        q = q / math.sqrt(self.head_size)
        scores = torch.matmul(q, k.transpose(2, 3))

        if mask is not None:
            scores = scores.masked_fill(~mask.unsqueeze(1), float("-inf"))

        attention_weights = self.softmax(scores)
        attention_probs = self.dropout(attention_weights)

        context = torch.matmul(attention_probs, v)
        context = (
            context.transpose(1, 2)
            .contiguous()
            .view(batch_size, -1, self.num_heads * self.head_size)
        )

        output = self.output_layer(context)

        if return_weights:
            attention_output_weights = attention_weights.view(
                batch_size, self.num_heads, query_len, key_len
            )
            return output, attention_output_weights.sum(dim=1) / self.num_heads
        return output, None


# ---------------------------------------------------------------------------
# Position-wise Feed-Forward
# ---------------------------------------------------------------------------


class PositionwiseFeedForward(nn.Module):
    """Position-wise Feed-forward layer."""

    def __init__(
        self,
        input_size: int,
        ff_size: int,
        dropout: float = 0.1,
        alpha: float = 1.0,
        layer_norm: str = "post",
        activation: str = "relu",
    ) -> None:
        super().__init__()
        activation_fnc = build_activation(activation=activation)

        self.layer_norm = nn.LayerNorm(input_size, eps=1e-6)
        self.pwff_layer = nn.Sequential(
            nn.Linear(input_size, ff_size),
            activation_fnc(),
            nn.Dropout(dropout),
            nn.Linear(ff_size, input_size),
            nn.Dropout(dropout),
        )

        self.alpha = alpha
        self._layer_norm_position = layer_norm
        assert self._layer_norm_position in {"pre", "post"}

    def forward(self, x: Tensor) -> Tensor:
        residual = x
        if self._layer_norm_position == "pre":
            x = self.layer_norm(x)

        x = self.pwff_layer(x) + self.alpha * residual

        if self._layer_norm_position == "post":
            x = self.layer_norm(x)
        return x


# ---------------------------------------------------------------------------
# Positional Encoding
# ---------------------------------------------------------------------------


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

    def __init__(self, size: int = 0, max_len: int = 5000) -> None:
        if size % 2 != 0:
            raise ValueError(
                f"Cannot use sin/cos positional encoding with odd dim (got dim={size})"
            )
        pe = torch.zeros(max_len, size)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, size, 2, dtype=torch.float) * -(math.log(10000.0) / size)
        )
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        pe = pe.unsqueeze(0)
        super().__init__()
        self.register_buffer("pe", pe)
        self.dim = size

    def forward(self, emb: Tensor) -> Tensor:
        return emb + self.pe[:, : emb.size(1)]


# ---------------------------------------------------------------------------
# Transformer Encoder Layer
# ---------------------------------------------------------------------------


class TransformerEncoderLayer(nn.Module):
    """One Transformer encoder layer: self-attention + feed-forward."""

    def __init__(
        self,
        size: int = 0,
        ff_size: int = 0,
        num_heads: int = 0,
        dropout: float = 0.1,
        alpha: float = 1.0,
        layer_norm: str = "post",
        activation: str = "relu",
    ) -> None:
        super().__init__()

        self.layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.src_src_att = MultiHeadedAttention(num_heads, size, dropout=dropout)

        self.feed_forward = PositionwiseFeedForward(
            size,
            ff_size=ff_size,
            dropout=dropout,
            alpha=alpha,
            layer_norm=layer_norm,
            activation=activation,
        )

        self.dropout = nn.Dropout(dropout)
        self.size = size
        self.alpha = alpha
        self._layer_norm_position = layer_norm
        assert self._layer_norm_position in {"pre", "post"}

    def forward(self, x: Tensor, mask: Tensor) -> Tensor:
        residual = x
        if self._layer_norm_position == "pre":
            x = self.layer_norm(x)

        x, _ = self.src_src_att(x, x, x, mask)
        x = self.dropout(x) + self.alpha * residual

        if self._layer_norm_position == "post":
            x = self.layer_norm(x)

        out = self.feed_forward(x)
        return out


# ---------------------------------------------------------------------------
# Transformer Encoder
# ---------------------------------------------------------------------------


class TransformerEncoder(nn.Module):
    """Full Transformer Encoder (stacked encoder layers)."""

    def __init__(
        self,
        hidden_size: int = 512,
        ff_size: int = 2048,
        num_layers: int = 8,
        num_heads: int = 4,
        dropout: float = 0.1,
        emb_dropout: float = 0.1,
        freeze: bool = False,
        **kwargs,
    ):
        super().__init__()
        self._output_size = hidden_size

        self.layers = nn.ModuleList(
            [
                TransformerEncoderLayer(
                    size=hidden_size,
                    ff_size=ff_size,
                    num_heads=num_heads,
                    dropout=dropout,
                    alpha=kwargs.get("alpha", 1.0),
                    layer_norm=kwargs.get("layer_norm", "pre"),
                    activation=kwargs.get("activation", "relu"),
                )
                for _ in range(num_layers)
            ]
        )

        self.pe = PositionalEncoding(hidden_size)
        self.emb_dropout = nn.Dropout(p=emb_dropout)

        self.layer_norm = (
            nn.LayerNorm(hidden_size, eps=1e-6)
            if kwargs.get("layer_norm", "post") == "pre"
            else None
        )

        if freeze:
            freeze_params(self)

    @property
    def output_size(self):
        return self._output_size

    def forward(self, src_embed: Tensor, src_length: Tensor, mask: Tensor = None) -> Tensor:
        x = self.pe(src_embed)
        x = self.emb_dropout(x)

        for layer in self.layers:
            x = layer(x, mask)

        if self.layer_norm is not None:
            x = self.layer_norm(x)
        return x


# ---------------------------------------------------------------------------
# Transformer Decoder Layer
# ---------------------------------------------------------------------------


class TransformerDecoderLayer(nn.Module):
    """Transformer decoder layer: self-attn + cross-attn + feed-forward."""

    def __init__(
        self,
        size: int = 0,
        ff_size: int = 0,
        num_heads: int = 0,
        dropout: float = 0.1,
        alpha: float = 1.0,
        layer_norm: str = "post",
        activation: str = "relu",
        **kwargs,
    ) -> None:
        super().__init__()
        self.size = size

        self.trg_trg_att = MultiHeadedAttention(num_heads, size, dropout=dropout)
        self.src_trg_att = MultiHeadedAttention(num_heads, size, dropout=dropout)

        self.feed_forward = PositionwiseFeedForward(
            size,
            ff_size=ff_size,
            dropout=dropout,
            alpha=alpha,
            layer_norm=layer_norm,
            activation=activation,
        )

        self.x_layer_norm = nn.LayerNorm(size, eps=1e-6)
        self.dec_layer_norm = nn.LayerNorm(size, eps=1e-6)

        self.dropout = nn.Dropout(dropout)
        self.alpha = alpha
        self._layer_norm_position = layer_norm
        assert self._layer_norm_position in {"pre", "post"}

    def forward(
        self,
        x: Tensor,
        memory: Tensor,
        src_mask: Tensor,
        trg_mask: Tensor,
        return_attention: bool = False,
        **kwargs,
    ) -> Tensor:
        # 1. target-target self-attention
        residual = x
        if self._layer_norm_position == "pre":
            x = self.x_layer_norm(x)

        h1, _ = self.trg_trg_att(x, x, x, mask=trg_mask)
        h1 = self.dropout(h1) + self.alpha * residual

        if self._layer_norm_position == "post":
            h1 = self.x_layer_norm(h1)

        # 2. source-target cross-attention
        h1_residual = h1
        if self._layer_norm_position == "pre":
            h1 = self.dec_layer_norm(h1)

        h2, att = self.src_trg_att(
            memory, memory, h1, mask=src_mask, return_weights=return_attention
        )
        h2 = self.dropout(h2) + self.alpha * h1_residual

        if self._layer_norm_position == "post":
            h2 = self.dec_layer_norm(h2)

        # 3. feed-forward
        out = self.feed_forward(h2)

        if return_attention:
            return out, att
        return out, None
