"""Transformers inference adapter using the unchanged native model equations."""
from dataclasses import fields
import torch
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_purrence import PurrenceConfig
from .purrence_core import ModelConfig, RecurrentLM


class PurrenceForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = PurrenceConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    _supports_cache_class = False

    def __init__(self, config, dtype=None):
        super().__init__(config)
        native = RecurrentLM(ModelConfig(**{
            f.name: getattr(config, f.name) for f in fields(ModelConfig)
        }))
        self.embedding = native.embedding
        self.block = native.block
        self.queries = native.queries

    def get_input_embeddings(self):
        return self.embedding

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

    def get_output_embeddings(self):
        return None  # The native forward uses the embedding weight directly.

    def read(self, sources, query):
        return RecurrentLM.read(self, sources, query)

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

    def forward(self, input_ids=None, attention_mask=None, return_dict=None,
                use_cache=False, past_key_values=None, labels=None, **kwargs):
        if labels is not None or past_key_values is not None:
            raise ValueError("This inference adapter does not implement training or a KV cache")
        if input_ids is None or input_ids.ndim != 2 or input_ids.numel() == 0:
            raise ValueError("Expected a nonempty batch of token IDs")
        if attention_mask is None:
            mask = torch.ones_like(input_ids, dtype=torch.bool)
        else:
            if attention_mask.shape != input_ids.shape:
                raise ValueError("Attention mask shape differs from input IDs")
            mask = attention_mask.bool()
        lengths = mask.sum(-1)
        if (lengths == 0).any() or (lengths > self.config.max_seq_len).any():
            raise ValueError("Each sequence must contain 1 to 8192 unmasked tokens")
        packed = input_ids[mask]
        starts = torch.cat((lengths.new_zeros(1), lengths.cumsum(0)[:-1]))
        document_starts = packed.eq(self.config.eos_token_id)
        document_starts[starts] = True
        offsets = document_starts.nonzero().flatten()
        cu = torch.cat((offsets, offsets.new_tensor([packed.numel()]))).to(torch.int32)
        positions = torch.arange(packed.numel(), device=packed.device) - torch.repeat_interleave(offsets, cu.diff().long())
        with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
            values = RecurrentLM.forward(self, packed, cu, positions)
        logits = values.new_zeros((*input_ids.shape, self.config.vocab_size))
        logits[mask] = values
        if return_dict is False or (return_dict is None and not self.config.use_return_dict):
            return (logits,)
        return CausalLMOutput(logits=logits)