import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from .configuration_approxdumb import ApproxDumbConfig

class ApproxDumbForCausalLM(PreTrainedModel):
    config_class = ApproxDumbConfig

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

        # Token & Position Embeddings
        self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_emb = nn.Embedding(config.max_len, config.d_model)

        # Causal Self-Attention
        self.q = nn.Linear(config.d_model, config.d_model)
        self.k = nn.Linear(config.d_model, config.d_model)
        self.v = nn.Linear(config.d_model, config.d_model)
        self.out_proj = nn.Linear(config.d_model, config.d_model)

        # Feed Forward Network
        self.ffn = nn.Sequential(
            nn.Linear(config.d_model, 12),
            nn.ReLU(),
            nn.Linear(12, config.d_model)
        )

        # LM Head
        self.lm_head = nn.Linear(config.d_model, config.vocab_size)

        self.post_init()

    def forward(
        self,
        input_ids=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        **kwargs
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        B, T = input_ids.shape

        pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
        h = self.token_emb(input_ids) + self.pos_emb(pos)

        # Causal Attention
        q, k, v = self.q(h), self.k(h), self.v(h)
        weights = (q @ k.transpose(-2, -1)) * (self.d_model ** -0.5)

        mask = torch.tril(torch.ones(T, T, device=input_ids.device))
        weights = weights.masked_fill(mask == 0, float('-inf'))
        weights = torch.softmax(weights, dim=-1)

        attn_out = weights @ v
        h = h + self.out_proj(attn_out)
        h = h + self.ffn(h)

        logits = self.lm_head(h)

        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))

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

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

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