import torch
import torch.nn as nn
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput

from models.config import TinyGPTConfig
from models.transformer_block import TransformerBlock


class TinyGPT(PreTrainedModel, GenerationMixin):
    config_class = TinyGPTConfig

    def __init__(self, config):
        super().__init__(config)
        self.config = config

        self.token_embedding = nn.Embedding(
            config.vocab_size,
            config.embed_dim,
        )

        self.position_embedding = nn.Embedding(
            config.max_seq_len,
            config.embed_dim,
        )

        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(
                config.embed_dim
            )
            for _ in range(config.num_layers)
        ])

        self.ln_f = nn.LayerNorm(
            config.embed_dim
        )

        self.lm_head = nn.Linear(
            config.embed_dim,
            config.vocab_size,
        )

        self.post_init()

    def forward(self, input_ids, **kwargs):
        batch_size, seq_len = input_ids.shape

        positions = torch.arange(
            seq_len,
            device=input_ids.device
        )

        token_emb = self.token_embedding(
            input_ids
        )

        pos_emb = self.position_embedding(
            positions
        )

        x = token_emb + pos_emb

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

        x = self.ln_f(x)

        logits = self.lm_head(x)

        return CausalLMOutput(logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

    def _init_weights(self, module):
        std = 0.02

        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)

        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


TinyGPT.register_for_auto_class("AutoModelForCausalLM")