from __future__ import annotations

from dataclasses import dataclass

import torch
from torch import nn
from torch.nn import functional as F
from transformers import GenerationMixin, PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import CausalLMOutputWithPast


class TinyGPTConfig(PretrainedConfig):
    model_type = "tiny_gpt"

    def __init__(
        self,
        vocab_size: int = 2048,
        context_length: int = 256,
        n_layers: int = 4,
        n_heads: int = 4,
        d_model: int = 128,
        d_ff: int = 512,
        dropout: float = 0.1,
        tie_embeddings: bool = True,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        pad_token_id: int = 3,
        unk_token_id: int = 0,
        use_cache: bool = False,
        **kwargs,
    ) -> None:
        is_decoder = kwargs.pop("is_decoder", True)
        tie_word_embeddings = kwargs.pop("tie_word_embeddings", tie_embeddings)
        use_cache = kwargs.pop("use_cache", use_cache)
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            unk_token_id=unk_token_id,
            use_cache=use_cache,
            tie_word_embeddings=tie_word_embeddings,
            is_decoder=is_decoder,
            **kwargs,
        )
        self.vocab_size = vocab_size
        self.context_length = context_length
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.d_model = d_model
        self.d_ff = d_ff
        self.dropout = dropout
        self.tie_embeddings = tie_embeddings

        # Standard Transformer aliases used by generation helpers.
        self.hidden_size = d_model
        self.intermediate_size = d_ff
        self.num_attention_heads = n_heads
        self.num_hidden_layers = n_layers
        self.max_position_embeddings = context_length
        self.head_dim = d_model // n_heads

        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")


class TokenEmbedding(nn.Module):
    def __init__(self, config: TinyGPTConfig) -> None:
        super().__init__()
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)

    @property
    def weight(self) -> torch.Tensor:
        return self.embedding.weight

    def forward(self, idx: torch.Tensor) -> torch.Tensor:
        return self.embedding(idx)


class PositionEmbedding(nn.Module):
    def __init__(self, config: TinyGPTConfig) -> None:
        super().__init__()
        self.embedding = nn.Embedding(config.context_length, config.d_model)

    def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
        positions = torch.arange(seq_len, device=device).unsqueeze(0)
        return self.embedding(positions)


class CausalSelfAttention(nn.Module):
    def __init__(self, config: TinyGPTConfig) -> None:
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads

        self.q_proj = nn.Linear(config.d_model, config.d_model)
        self.k_proj = nn.Linear(config.d_model, config.d_model)
        self.v_proj = nn.Linear(config.d_model, config.d_model)
        self.out_proj = nn.Linear(config.d_model, config.d_model)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(config.context_length, config.context_length))
        self.register_buffer(
            "causal_mask",
            mask.view(1, 1, config.context_length, config.context_length),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, seq_len, d_model = x.shape

        query = self._split_heads(self.q_proj(x), batch_size, seq_len)
        key = self._split_heads(self.k_proj(x), batch_size, seq_len)
        value = self._split_heads(self.v_proj(x), batch_size, seq_len)

        scores = query @ key.transpose(-2, -1)
        scores = scores / (self.head_dim**0.5)
        scores = scores.masked_fill(
            self.causal_mask[:, :, :seq_len, :seq_len] == 0,
            float("-inf"),
        )

        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.attn_dropout(attention_weights)

        out = attention_weights @ value
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        out = self.out_proj(out)
        return self.resid_dropout(out)

    def _split_heads(self, x: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor:
        x = x.view(batch_size, seq_len, self.n_heads, self.head_dim)
        return x.transpose(1, 2)


class FeedForward(nn.Module):
    def __init__(self, config: TinyGPTConfig) -> None:
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff),
            nn.GELU(),
            nn.Linear(config.d_ff, config.d_model),
            nn.Dropout(config.dropout),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class TransformerBlock(nn.Module):
    def __init__(self, config: TinyGPTConfig) -> None:
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.d_model)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.d_model)
        self.mlp = FeedForward(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x


class TinyGPTForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TinyGPTConfig
    base_model_prefix = "tiny_gpt"
    main_input_name = "input_ids"
    _tied_weights_keys = {"lm_head.weight": "token_embedding.embedding.weight"}

    def __init__(self, config: TinyGPTConfig) -> None:
        super().__init__(config)

        self.token_embedding = TokenEmbedding(config)
        self.position_embedding = PositionEmbedding(config)
        self.dropout = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
        self.final_ln = nn.LayerNorm(config.d_model)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        self.post_init()
        self.tie_weights()
        if getattr(self, "generation_config", None) is not None:
            self.generation_config.use_cache = False
            self.generation_config.cache_implementation = None

    def get_input_embeddings(self) -> nn.Module:
        return self.token_embedding.embedding

    def set_input_embeddings(self, value: nn.Module) -> None:
        self.token_embedding.embedding = value

    def get_output_embeddings(self) -> nn.Module:
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
        self.lm_head = new_embeddings

    def tie_weights(self, *args, **kwargs) -> None:
        del args, kwargs
        if self.config.tie_embeddings:
            self.lm_head.weight = self.token_embedding.weight

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        past_key_values=None,
        use_cache: bool | None = None,
        return_dict: bool = True,
        **kwargs,
    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
        del attention_mask, past_key_values, use_cache, kwargs

        batch_size, seq_len = input_ids.shape
        if seq_len > self.config.context_length:
            raise ValueError(
                f"Sequence length {seq_len} exceeds context length {self.config.context_length}"
            )
        if labels is not None and labels.shape != input_ids.shape:
            raise ValueError(f"labels shape {labels.shape} must match input_ids shape {input_ids.shape}")

        token_embeddings = self.token_embedding(input_ids)
        position_embeddings = self.position_embedding(seq_len, input_ids.device)
        x = token_embeddings + position_embeddings
        x = self.dropout(x)

        for block in self.blocks:
            x = block(x)

        x = self.final_ln(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            if seq_len < 2:
                raise ValueError("Need at least 2 tokens to compute causal LM loss")
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.reshape(-1, self.config.vocab_size),
                shift_labels.reshape(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits,)
            if loss is not None:
                output = (loss,) + output
            return output

        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        past_key_values=None,
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> dict[str, torch.Tensor | None]:
        del past_key_values
        if input_ids.shape[1] > self.config.context_length:
            input_ids = input_ids[:, -self.config.context_length :]
            if attention_mask is not None:
                attention_mask = attention_mask[:, -self.config.context_length :]
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
            **kwargs,
        }


TinyGPTConfig.register_for_auto_class()
TinyGPTForCausalLM.register_for_auto_class("AutoModelForCausalLM")