"""HF AutoModel wrapper for mini-beatrix-1 (AlephLM byte-trigram craft).

Loads with:
    AutoModelForCausalLM.from_pretrained("AbstractPhil/mini-beatrix-1",
                                         trust_remote_code=True)
The model reads raw UTF-8 bytes: input_ids are byte values 0..255.
"""
from __future__ import annotations

import torch
from transformers import PretrainedConfig, PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

from .presets import AlephLMConfig
from .alephlm import AlephLM


class MiniBeatrixConfig(PretrainedConfig):
    model_type = "mini-beatrix"
    attribute_map = {"num_hidden_layers": "n_layers",
                     "hidden_size": "d_model",
                     "num_attention_heads": "n_heads",
                     "max_position_embeddings": "context"}

    def __init__(self, name="mini-beatrix-1", d_model=768, n_layers=16,
                 n_heads=12, context=2048, vocab_size=256,
                 tokenizer="byte-trigram", hub_layers=(4, 9, 14),
                 hub_K=512, hub_D=32, tau=0.1, bank_experts=3, bank_ff=None,
                 head_K=512, head_D=32, gate_init=-3.0,
                 tie_embeddings=False, hub_chunk=128, **kwargs):
        self.name = name
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.context = context
        self.vocab_size = vocab_size
        self.tokenizer = tokenizer
        self.hub_layers = list(hub_layers)
        self.hub_K = hub_K
        self.hub_D = hub_D
        self.tau = tau
        self.bank_experts = bank_experts
        self.bank_ff = bank_ff
        self.head_K = head_K
        self.head_D = head_D
        self.gate_init = gate_init
        self.tie_embeddings = tie_embeddings
        self.hub_chunk = hub_chunk
        super().__init__(**kwargs)

    def to_aleph(self) -> AlephLMConfig:
        return AlephLMConfig(
            name=self.name, d_model=self.d_model, n_layers=self.n_layers,
            n_heads=self.n_heads, context=self.context,
            vocab_size=self.vocab_size, tokenizer=self.tokenizer,
            hub_layers=tuple(self.hub_layers), hub_K=self.hub_K,
            hub_D=self.hub_D, tau=self.tau, bank_experts=self.bank_experts,
            bank_ff=self.bank_ff, head_K=self.head_K, head_D=self.head_D,
            gate_init=self.gate_init, tie_embeddings=self.tie_embeddings,
            hub_chunk=self.hub_chunk)


class MiniBeatrixForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = MiniBeatrixConfig
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False
    _tied_weights_keys = {}

    def __init__(self, config: MiniBeatrixConfig):
        super().__init__(config)
        self.model = AlephLM(config.to_aleph())
        self.post_init()

    def _init_weights(self, module):
        pass  # AlephLM initializes itself; loader must not re-init

    def forward(self, input_ids=None, labels=None, attention_mask=None,
                past_key_values=None, use_cache=None, **kwargs):
        # attention_mask is safely ignored: causal model, right-padding
        # with label masking is the trained convention. No KV cache in
        # this wrapper; generate() recomputes the prefix each step.
        if input_ids.shape[1] > self.config.context:
            input_ids = input_ids[:, -self.config.context:]
            if labels is not None:
                labels = labels[:, -self.config.context:]
        out = self.model(input_ids, labels=labels)
        logits = out[0]
        loss = out[1] if labels is not None else None
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids[:, -self.config.context:]}