from __future__ import annotations

import torch
import torch.nn as nn
from transformers import PretrainedConfig
from transformers import PreTrainedModel

from src.transformer_block.layer_norm import LayerNorm
from src.transformer_block.t_block import GPTTransformerBlock


class GPTConfig(PretrainedConfig):
    """
    Configuration class for GPT-2 model
    """

    model_type = "gpt_fast_llm"

    def __init__(
        self,
        vocab_size: int = 200019,
        context_len: int = 256,
        embedding_dim: int = 768,
        n_heads: int = 12,
        n_layers: int = 12,
        drop_rate: float = 0.0,
        qkv_bias: bool = True,
        batch_size: int = 8,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.context_len = context_len
        self.embedding_dim = embedding_dim
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.drop_rate = drop_rate
        self.qkv_bias = qkv_bias
        self.batch_size = batch_size


class GPTModel(PreTrainedModel):
    """
    The base model for GPT-2 architecture
    Input:
        x : tensor of shape (batch_size, seq_len)
    Output:
        logits : tensor of shape (batch_size, seq_len, vocab_size)
    """

    config_class = GPTConfig

    def __init__(self, config: GPTConfig):
        super().__init__(config)
        self.token_embedding = nn.Embedding(
            num_embeddings=config.vocab_size,
            embedding_dim=config.embedding_dim,
        )
        self.pos_embedding = nn.Embedding(
            num_embeddings=config.context_len,
            embedding_dim=config.embedding_dim,
        )
        self.drop_embedding = nn.Dropout(config.drop_rate)
        self.transformer_blocks = nn.Sequential(
            *[GPTTransformerBlock(config=config) for _ in range(config.n_layers)],
        )
        self.final_norm = LayerNorm(config.embedding_dim)
        self.out_head = nn.Linear(
            config.embedding_dim,
            config.vocab_size,
            bias=False,
        )

    def forward(self, input_ids, attention_mask=None, labels=None):
        batch_size, seq_len = input_ids.shape
        token_emb = self.token_embedding(input_ids)
        pos_emb = self.pos_embedding(torch.arange(seq_len, device=input_ids.device))
        x = token_emb + pos_emb
        x = self.drop_embedding(x)
        x = self.transformer_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)

        if labels is not None:
            loss = torch.nn.functional.cross_entropy(
                logits.flatten(0, 1),
                labels.flatten(),
            )
            return {"logits": logits, "loss": loss}
        return {"logits": logits}