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

from .configuration_bet import BETConfig
from .bet_model import BETConfig as CoreConfig, SparkBET, uniform_steps


class BETPreTrainedModel(PreTrainedModel):
    config_class=BETConfig
    base_model_prefix="core"
    supports_gradient_checkpointing=False
    _no_split_modules=["PlainBlock","LoopedBlock"]


class BETForCausalLM(BETPreTrainedModel,GenerationMixin):
    def __init__(self,config):
        super().__init__(config)
        core_cfg=CoreConfig(
            vocab_size=config.vocab_size,
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            prelude_layers=config.prelude_layers,
            body_blocks=config.body_blocks,
            coda_layers=config.coda_layers,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            head_dim=config.head_dim,
            lora_rank=config.lora_rank,
            hyper_lanes=config.hyper_lanes,
            max_seq_len=config.max_position_embeddings,
            max_loops=config.max_loops,
            rope_theta=config.rope_theta,
            rms_eps=config.rms_norm_eps,
            ddl_beta_init=config.ddl_beta_init,
            ddl_k_eps=config.ddl_k_eps,
            ddl_v_sigmoid_scale=config.ddl_v_sigmoid_scale,
        )
        self.core=SparkBET(core_cfg)

    def get_input_embeddings(self):return self.core.embed
    def set_input_embeddings(self,value):self.core.embed=value
    def get_output_embeddings(self):return None
    def set_output_embeddings(self,value):
        if value is not None:raise ValueError("SparkBET uses tied input/output embeddings")

    def _cycles(self,cycles=None):
        if cycles is None:
            cycles=int(os.environ.get("BET_EVAL_CYCLES",self.config.refinement_cycles))
        cycles=int(cycles)
        if not 1<=cycles<=self.config.max_loops:
            raise ValueError(f"refinement cycles must be in [1,{self.config.max_loops}]")
        return cycles

    def forward(
        self,input_ids=None,attention_mask=None,labels=None,cycles=None,
        past_key_values=None,use_cache=None,return_dict=True,**kwargs,
    ):
        if input_ids is None:raise ValueError("input_ids is required")
        if past_key_values is not None:raise ValueError("SparkBET does not implement a KV cache")
        logits=self.core(input_ids,uniform_steps(self._cycles(cycles)),attention_mask=attention_mask)
        loss=None
        if labels is not None:
            shift_logits=logits[:,:-1].contiguous().float();shift_labels=labels[:,1:].contiguous()
            loss=F.cross_entropy(shift_logits.view(-1,shift_logits.size(-1)),shift_labels.view(-1),ignore_index=-100)
        if not return_dict:return tuple(v for v in (loss,logits) if v is not None)
        return CausalLMOutput(loss=loss,logits=logits)

    def prepare_inputs_for_generation(self,input_ids,attention_mask=None,**kwargs):
        max_len=self.config.max_position_embeddings
        if input_ids.shape[1]>max_len:
            input_ids=input_ids[:,-max_len:]
            if attention_mask is not None:attention_mask=attention_mask[:,-max_len:]
        return {"input_ids":input_ids,"attention_mask":attention_mask,"use_cache":False}