import torch

from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_halos import HaloSConfig

# Tu implementación real
from halo import HaloConfig
from halo import HaloSModel


class HaloSPreTrainedModel(PreTrainedModel):
    config_class = HaloSConfig
    base_model_prefix = "halo"

    def _init_weights(self, module):
        pass


class HaloSHFModel(HaloSPreTrainedModel):

    def __init__(self, config):
        super().__init__(config)

        halo_cfg = HaloConfig(
            vocab_size=config.vocab_size,
            hidden_size=config.hidden_size,
            num_layers=config.num_layers,
            num_heads=config.num_heads,
            num_kv_heads=config.num_kv_heads,
            num_globals=config.num_globals,
            local_window=config.local_window,
            dilated_offsets=config.dilated_offsets,
            num_random=config.num_random,
            dropout=config.dropout,
            max_seq_len=config.max_seq_len,
        )

        self.halo = HaloSModel(halo_cfg)

        print("ANTES:", self.halo.token_emb.weight.abs().mean())
        self.post_init()
        print("DESPUES:", self.halo.token_emb.weight.abs().mean())
        
    def forward(
        self,
        input_ids=None,
        labels=None,
        **kwargs
    ):
        result = self.halo(input_ids)
        
        logits = result[0]
        
        return CausalLMOutput(
            logits=logits,
            loss=None
        )


class HaloSForCausalLM(HaloSHFModel):
    pass