import torch
from torch.nn import functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_sol_pro import SolProHFConfig
from .modeling_sol_pro import SolProConfig, SolProForCausalLM


class SolProHFForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = SolProHFConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    _no_split_modules = ["Block"]

    def __init__(self, config):
        super().__init__(config)
        fields = SolProConfig.__dataclass_fields__
        core = SolProForCausalLM(SolProConfig(**{k: getattr(config, k) for k in fields}))
        # Keep original checkpoint names, including the single tied embedding.
        for name, module in core.named_children():
            self.add_module(name, module)
        for name, parameter in core.named_parameters(recurse=False):
            self.register_parameter(name, parameter)
        for name, buffer in core.named_buffers(recurse=False):
            self.register_buffer(name, buffer)
        self.post_init()

    def get_input_embeddings(self):
        return self.embedding

    def set_input_embeddings(self, value):
        self.embedding = value

    def get_output_embeddings(self):
        return self.embedding

    def set_output_embeddings(self, value):
        self.embedding = value

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}

    def forward(self, input_ids=None, attention_mask=None, labels=None,
                return_dict=None, use_cache=False, past_key_values=None, **kwargs):
        if input_ids is None:
            raise ValueError("input_ids are required")
        if past_key_values is not None:
            raise ValueError("KV caching is not implemented")
        if attention_mask is None or bool(attention_mask.bool().all()):
            logits = SolProForCausalLM.forward(self, input_ids)
        else:
            # Strip padding so n-gram memory and recurrent conditioning see the
            # same sequence as an unpadded standalone inference call.
            rows = []
            for ids, mask in zip(input_ids, attention_mask):
                keep = mask.bool()
                positions = keep.nonzero(as_tuple=True)[0]
                if positions.numel() == 0:
                    raise ValueError("An attention mask cannot be empty")
                values = SolProForCausalLM.forward(self, ids[keep][None])[0]
                row = values.new_zeros((ids.numel(), self.config.vocab_size))
                rows.append(row.index_copy(0, positions, values))
            logits = torch.stack(rows)
        loss = None
        if labels is not None:
            targets = labels[:, 1:].clone()
            if attention_mask is not None:
                targets.masked_fill_(~attention_mask[:, 1:].bool(), -100)
            loss = F.cross_entropy(logits[:, :-1].float().reshape(-1, logits.shape[-1]),
                                   targets.reshape(-1), ignore_index=-100)
        if return_dict is False:
            return ((loss,) if loss is not None else ()) + (logits,)
        return CausalLMOutput(loss=loss, logits=logits)