"""Transformers wrapper for PakMosaic. Same weights as PakEncoderForMLM."""

from __future__ import annotations

from typing import Any

import torch
import torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput, MaskedLMOutput

try:
    from .configuration_pakmosaic import PakMosaicConfig
except ImportError:
    from configuration_pakmosaic import PakMosaicConfig


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 512, base: float = 10000.0):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.max_seq_len = max_seq_len

    def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
        positions = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(positions, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    cos = cos.unsqueeze(0).unsqueeze(0)
    sin = sin.unsqueeze(0).unsqueeze(0)
    return (x * cos) + (rotate_half(x) * sin)


def sliding_window_keep(seq_len: int, window: int, device: torch.device) -> torch.Tensor:
    idx = torch.arange(seq_len, device=device)
    return (idx[:, None] - idx[None, :]).abs() <= window


class EncoderAttention(nn.Module):
    def __init__(self, config: PakMosaicConfig, *, layer_index: int):
        super().__init__()
        self.config = config
        self.num_heads = config.num_attention_heads
        self.head_dim = config.head_dim
        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.use_bias)
        self.out = nn.Linear(config.hidden_size, config.hidden_size, bias=config.use_bias)
        self.dropout = config.attention_dropout
        self.rope = RotaryEmbedding(self.head_dim, config.max_position_embeddings, config.rope_theta)
        self.is_global = config.attention_kind == "full" or (
            config.attention_kind == "hybrid" and ((layer_index + 1) % max(1, config.global_every) == 0)
        )

    def _mask(self, attention_mask: torch.Tensor | None, seq_len: int) -> torch.Tensor | None:
        keep = None
        if attention_mask is not None:
            keep = attention_mask.bool()[:, None, None, :]
        if not self.is_global:
            window = sliding_window_keep(seq_len, self.config.local_window, self.qkv.weight.device)[None, None, :, :]
            keep = window if keep is None else keep & window
        return keep

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
        batch, seq_len, _ = hidden.shape
        qkv = self.qkv(hidden).view(batch, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        query, key, value = qkv[0], qkv[1], qkv[2]
        cos, sin = self.rope(seq_len, hidden.device, query.dtype)
        query = apply_rope(query, cos, sin)
        key = apply_rope(key, cos, sin)
        context = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=self._mask(attention_mask, seq_len),
            dropout_p=self.dropout if self.training else 0.0,
        )
        context = context.transpose(1, 2).contiguous().view(batch, seq_len, self.config.hidden_size)
        return self.out(context)


class EncoderFFN(nn.Module):
    def __init__(self, config: PakMosaicConfig):
        super().__init__()
        self.kind = config.ffn_activation
        bias = config.use_bias
        if self.kind in {"geglu", "swiglu"}:
            self.gate_up = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=bias)
            self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=bias)
        else:
            self.up = nn.Linear(config.hidden_size, config.intermediate_size, bias=bias)
            self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=bias)
            self.act = nn.GELU()
        self.dropout = nn.Dropout(config.hidden_dropout)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        if self.kind in {"geglu", "swiglu"}:
            gate, up = self.gate_up(hidden).chunk(2, dim=-1)
            hidden = (F.silu(gate) if self.kind == "swiglu" else F.gelu(gate)) * up
        else:
            hidden = self.act(self.up(hidden))
        return self.dropout(self.down(hidden))


class TokenEmbeddings(nn.Module):
    def __init__(self, config: PakMosaicConfig):
        super().__init__()
        self.word = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.norm(self.word(input_ids)))


class EncoderLayer(nn.Module):
    def __init__(self, config: PakMosaicConfig, layer_index: int):
        super().__init__()
        self.attn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.attn = EncoderAttention(config, layer_index=layer_index)
        self.ffn = EncoderFFN(config)

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
        hidden = hidden + self.attn(self.attn_norm(hidden), attention_mask=attention_mask)
        return hidden + self.ffn(self.ffn_norm(hidden))


class PakMosaicPreTrainedModel(PreTrainedModel):
    config_class = PakMosaicConfig
    base_model_prefix = "encoder"
    supports_gradient_checkpointing = True
    _no_split_modules = ["EncoderLayer"]

    def _init_weights(self, module: nn.Module) -> None:
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.weight.data.fill_(1.0)
            module.bias.data.zero_()


class PakMosaicModel(PakMosaicPreTrainedModel):
    def __init__(self, config: PakMosaicConfig):
        super().__init__(config)
        self.embeddings = TokenEmbeddings(config)
        self.layers = nn.ModuleList(EncoderLayer(config, i) for i in range(config.num_hidden_layers))
        self.final_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.embeddings.word

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.embeddings.word = value

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> BaseModelOutput:
        hidden = self.embeddings(input_ids)
        for layer in self.layers:
            if self.gradient_checkpointing and self.training:
                hidden = self._gradient_checkpointing_func(layer, hidden, attention_mask)
            else:
                hidden = layer(hidden, attention_mask=attention_mask)
        hidden = self.final_norm(hidden)
        return BaseModelOutput(last_hidden_state=hidden)


class MLMHead(nn.Module):
    def __init__(self, config: PakMosaicConfig, embedding: nn.Embedding):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=True)
        self.act = nn.GELU()
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.tie = config.tie_word_embeddings
        if config.tie_word_embeddings:
            self.weight = embedding.weight
            self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        else:
            self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        hidden = self.norm(self.act(self.dense(hidden)))
        if self.tie:
            return F.linear(hidden, self.weight, self.bias)
        return self.decoder(hidden)


class PakMosaicForMaskedLM(PakMosaicPreTrainedModel):
    _tied_weights_keys = {"mlm.weight": "encoder.embeddings.word.weight"}

    def __init__(self, config: PakMosaicConfig):
        super().__init__(config)
        self.encoder = PakMosaicModel(config)
        self.mlm = MLMHead(config, self.encoder.embeddings.word)
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.encoder.embeddings.word

    def forward(
        self,
        input_ids: torch.Tensor | None = None,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> MaskedLMOutput:
        hidden = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        logits = self.mlm(hidden)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100)
        return MaskedLMOutput(loss=loss, logits=logits, hidden_states=(hidden,))


PakMosaicModel.register_for_auto_class("AutoModel")
PakMosaicForMaskedLM.register_for_auto_class("AutoModelForMaskedLM")