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

from configuration_axiom import axiomConfig
from config import ModelConfig
from model import LLM


class axiomForCausalLM(PreTrainedModel):
    config_class = axiomConfig
    base_model_prefix = "model"
    _no_split_modules = ["TransformerBlock"]

    def __init__(self, config: axiomConfig):
        super().__init__(config)
        core_cfg = ModelConfig(
            vocab_size=config.vocab_size,
            dim=config.dim,
            n_layers=config.n_layers,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads,
            ffn_dim_multiplier=config.ffn_dim_multiplier,
            max_seq_len=config.max_seq_len,
            rope_theta=config.rope_theta,
            norm_eps=config.norm_eps,
            dropout=config.dropout,
        )
        self.model = LLM(core_cfg)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed

    def set_input_embeddings(self, value):
        self.model.embed = value
        self.model.lm_head.weight = value.weight

    def get_output_embeddings(self):
        return self.model.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.model.lm_head = new_embeddings

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        return {"input_ids": input_ids, "past_key_values": past_key_values}

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        if past_key_values is None:
            return past_key_values
        out = []
        for layer in past_key_values:
            if layer is None:
                out.append(layer)
                continue
            k, v = layer
            out.append((k.index_select(0, beam_idx), v.index_select(0, beam_idx)))
        return out

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        past_key_values=None,
        use_cache=None,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("input_ids must be provided")
        if use_cache is None:
            use_cache = True
        logits, loss, new_cache = self.model(
            input_ids,
            targets=labels,
            cache=past_key_values,
            use_grad_ckpt=False,
            return_cache=bool(use_cache),
        )
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=new_cache if use_cache else None,
        )