import torch
from torch.nn import functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_kairo import KairoGPTConfig
from .kairo_model import KairoGPT, KairoGPTConfig as _InternalCfg


class KairoGPTForCausalLM(PreTrainedModel):
    config_class = KairoGPTConfig

    def __init__(self, config: KairoGPTConfig):
        super().__init__(config)
        internal_cfg = _InternalCfg(
            vocab_size=config.vocab_size,
            block_size=config.block_size,
            n_layer=config.n_layer,
            n_head=config.n_head,
            n_embd=config.n_embd,
            dropout=config.dropout,
        )
        self.transformer = KairoGPT(internal_cfg)

    def forward(self, input_ids, labels=None, **kwargs):
        logits, _ = self.transformer(input_ids)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
            )
        return CausalLMOutput(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids[:, -self.config.block_size:]}

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=400, temperature=0.8, top_k=40, **kwargs):
        return self.transformer.generate(
            input_ids, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k
        )