from __future__ import annotations

import math

import torch
from torch import nn


def sinusoidal_time_embedding(t: torch.Tensor, dimension: int) -> torch.Tensor:
    half = dimension // 2
    frequencies = torch.exp(
        -math.log(10_000)
        * torch.arange(half, device=t.device, dtype=torch.float32)
        / max(half - 1, 1)
    )
    angles = t.float().unsqueeze(1) * frequencies.unsqueeze(0)
    embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)
    if dimension % 2:
        embedding = torch.nn.functional.pad(embedding, (0, 1))
    return embedding


class MultiMDMTransformer(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        clean_vocab_size: int,
        seq_len: int,
        num_masks: int,
        d_model: int = 128,
        nhead: int = 4,
        num_layers: int = 2,
        dim_feedforward: int = 256,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.clean_vocab_size = clean_vocab_size
        self.seq_len = seq_len
        self.num_masks = num_masks
        self.d_model = d_model
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(seq_len, d_model)
        self.time_mlp = nn.Sequential(
            nn.Linear(d_model, d_model * 2),
            nn.SiLU(),
            nn.Linear(d_model * 2, d_model),
        )
        layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers)
        self.norm = nn.LayerNorm(d_model)
        self.output_head = nn.Linear(d_model, clean_vocab_size)
        self.mask_class_head = nn.Linear(d_model, num_masks)

    def forward(
        self,
        input_ids: torch.Tensor,
        t: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        batch, length = input_ids.shape
        if length > self.seq_len:
            raise ValueError(f"input length {length} exceeds seq_len {self.seq_len}")
        positions = torch.arange(length, device=input_ids.device)
        hidden = self.token_embedding(input_ids)
        hidden = hidden + self.position_embedding(positions).unsqueeze(0)
        time_hidden = self.time_mlp(
            sinusoidal_time_embedding(t, self.d_model)
        ).unsqueeze(1)
        hidden = hidden + time_hidden
        padding_mask = None if attention_mask is None else ~attention_mask.bool()
        hidden = self.transformer(hidden, src_key_padding_mask=padding_mask)
        hidden = self.norm(hidden)
        return self.output_head(hidden), self.mask_class_head(hidden)