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

from .configuration_tinyindianlm import TinyIndianLMConfig
from .model import ModelConfig, TinyIndianLM


class TinyIndianLMForCausalLM(PreTrainedModel):
    config_class = TinyIndianLMConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False

    def __init__(self, config: TinyIndianLMConfig):
        super().__init__(config)
        internal_cfg = ModelConfig(
            vocab_size=config.vocab_size,
            d_model=config.d_model,
            n_layers=config.n_layers,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads,
            max_seq_len=config.max_seq_len,
            d_ffn=config.d_ffn,
            dropout=config.dropout,
            rope_theta=config.rope_theta,
        )
        self.model = TinyIndianLM(internal_cfg)
        self.post_init()

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.embedding

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.model.embedding = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.model.output

    def set_output_embeddings(self, new_embeddings: nn.Linear) -> None:
        self.model.output = new_embeddings

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        **kwargs,
    ):
        del attention_mask, kwargs

        # Safety guard for exported repos where tokenizer vocab and
        # model embedding size diverge. Any out-of-range token id is
        # mapped to unk to avoid embedding lookup crashes.
        if input_ids is not None:
            vocab_size = self.model.embedding.num_embeddings
            if (input_ids >= vocab_size).any():
                unk_id = getattr(self.config, "unk_token_id", 1)
                input_ids = input_ids.clone()
                input_ids[input_ids >= vocab_size] = unk_id

        logits, loss = self.model(
            input_ids,
            targets=labels,
            pad_id=self.config.pad_token_id,
        )
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )

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

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        del beam_idx
        return past_key_values