"""Embedding modules used by the transformer architecture."""

from __future__ import annotations

from typing import cast

import torch
import torch.nn as nn
from torch import Tensor

from ..utils import sinusoidal_positional_encoding

__all__ = ["InputEmbedding", "PositionalEmbedding", "LearnedPositionalEmbedding"]


class InputEmbedding(nn.Module):
    """Token and positional embedding with optional dropout."""

    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        max_seq_len: int,
        pad_id: int,
        dropout_rate: float = 0.0,
    ) -> None:
        super().__init__()

        if not isinstance(vocab_size, int):
            raise TypeError(f"vocab_size must be int, got {type(vocab_size)}")
        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be int, got {type(d_model)}")
        if not isinstance(max_seq_len, int):
            raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}")
        if not isinstance(pad_id, int):
            raise TypeError(f"pad_id must be int, got {type(pad_id)}")
        if not isinstance(dropout_rate, float):
            raise TypeError(f"dropout_rate must be a float, got {type(dropout_rate)}")

        if vocab_size <= 0:
            raise ValueError("vocab_size must be > 0")
        if d_model <= 0:
            raise ValueError("d_model must be > 0")
        if max_seq_len <= 0:
            raise ValueError("max_seq_len must be > 0")
        if not (0 <= pad_id < vocab_size):
            raise ValueError(f"pad_id must be in [0, {vocab_size - 1}]")
        if not 0.0 <= dropout_rate < 1.0:
            raise ValueError("dropout_rate must be in [0, 1)")

        self.vocab_size = vocab_size
        self.d_model = d_model
        self.max_seq_len = max_seq_len
        self.pad_id = pad_id

        self.token_embed = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
        self.pos_embed = PositionalEmbedding(max_seq_len, d_model)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, x: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"x must be Tensor, got {type(x)}")
        if x.dtype != torch.long:
            raise TypeError(f"x must be torch.long, got {x.dtype}")
        if x.dim() != 2:
            raise ValueError(f"x must be 2D (B, S), got shape {tuple(x.shape)}")
        if x.size(1) > self.max_seq_len:
            raise ValueError(f"Sequence length {x.size(1)} exceeds max_seq_len {self.max_seq_len}")

        tokens = self.token_embed(x) * (self.d_model**0.5)
        tokens = self.pos_embed(tokens)
        return self.dropout(tokens)


class PositionalEmbedding(nn.Module):
    """Fixed sinusoidal positional encodings."""

    def __init__(self, max_seq_len: int, d_model: int) -> None:
        super().__init__()

        if not isinstance(max_seq_len, int):
            raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}")
        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be int, got {type(d_model)}")
        if max_seq_len < 0:
            raise ValueError("max_seq_len must be >= 0")
        if d_model <= 0:
            raise ValueError("d_model must be > 0")

        self.max_seq_len = max_seq_len
        self.d_model = d_model

        pe = sinusoidal_positional_encoding(max_seq_len, d_model)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"x must be Tensor, got {type(x)}")
        if x.dim() != 3:
            raise ValueError(f"x must be 3D (B, S, D), got shape {tuple(x.shape)}")

        _, seq_len, dim = x.shape
        if dim != self.d_model:
            raise ValueError(f"d_model mismatch: got {dim}, expected {self.d_model}")
        if seq_len > self.max_seq_len:
            raise ValueError(f"seq_len {seq_len} exceeds max_seq_len {self.max_seq_len}")

        pe_buffer = cast(Tensor, self.pe)
        return x + pe_buffer[:, :seq_len].to(dtype=x.dtype, device=x.device)


class LearnedPositionalEmbedding(nn.Module):
    """Learned positional embeddings compatible with :class:`InputEmbedding`."""

    def __init__(self, max_seq_len: int, d_model: int) -> None:
        super().__init__()

        if not isinstance(max_seq_len, int):
            raise TypeError(f"max_seq_len must be int, got {type(max_seq_len)}")
        if not isinstance(d_model, int):
            raise TypeError(f"d_model must be int, got {type(d_model)}")
        if max_seq_len < 0:
            raise ValueError("max_seq_len must be >= 0")
        if d_model <= 0:
            raise ValueError("d_model must be > 0")

        self.max_seq_len = max_seq_len
        self.d_model = d_model
        self.pos_table = nn.Embedding(max_seq_len, d_model)
        nn.init.normal_(self.pos_table.weight, mean=0.0, std=0.02)

    def forward(self, x: Tensor) -> Tensor:
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"x must be Tensor, got {type(x)}")
        if x.dim() != 3:
            raise ValueError(f"x must be 3D (B, S, D), got {tuple(x.shape)}")

        batch, seq_len, dim = x.shape
        if dim != self.d_model:
            raise ValueError(f"d_model mismatch: got {dim}, expected {self.d_model}")
        if seq_len > self.max_seq_len:
            raise ValueError(f"seq_len {seq_len} exceeds max_seq_len {self.max_seq_len}")

        positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch, seq_len)
        pos_emb = self.pos_table(positions)
        return x + pos_emb