# This is free software for the public good of a permacomputer hosted at
# permacomputer.com, an always-on computer by the people, for the people.
"""ANDREA model for HuggingFace transformers.

A line-for-line port of our CUDA engine's forward pass. The reference is
`scripts/andrea_np_forward.py`, which is itself validated against the pure
Python `microgpt.py` to float32. Parity between this file and that one is
asserted in `microgpt/tests/test_hf_modeling.py` — if you change the math
here, that test is what catches it.

    x = wte[tok] + wpe[pos]
    x = rmsnorm(x)
    per layer:
        x = x + o_proj(attn(rmsnorm(x)))
        x = x + down_proj(relu(up_proj(rmsnorm(x))))
    logits = lm_head(rms_norm(x)) if config.final_norm else lm_head(x)

RMSNorm here is parameter-free: x * (mean(x^2) + eps) ** -0.5.

With config.qk_norm (andrea-v7+), q and k additionally get a parameter-free
per-head RMSNorm after projection, before the dot product.

Attention is full causal unless `config.sliding_window` is set, in which case a
query at position p sees only keys [max(0, p-w+1), p] — the same window the
CUDA engine's SWA_WINDOW applies during training. It must match the checkpoint.
"""
import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel

from .configuration_andrea import AndreaConfig


def rms_norm(x, eps):
    """Parameter-free RMSNorm. Computed in float32 so fp16 weights stay stable."""
    dtype = x.dtype
    x32 = x.float()
    scale = torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + eps)
    return (x32 * scale).to(dtype)


def cache_length(past_key_values, layer_idx=0):
    """Past sequence length, for either a transformers Cache or a legacy tuple."""
    if not past_key_values:
        return 0
    if hasattr(past_key_values, 'get_seq_length'):
        return int(past_key_values.get_seq_length(layer_idx) or 0)
    layer = past_key_values[layer_idx]
    return 0 if layer is None or layer[0] is None else layer[0].shape[2]


class AndreaAttention(nn.Module):
    def __init__(self, config, layer_idx=0):
        super().__init__()
        self.layer_idx = layer_idx
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.scaling = 1.0 / math.sqrt(self.head_dim)
        self.qk_norm = bool(getattr(config, 'qk_norm', False))
        self.eps = config.rms_norm_eps
        E = config.hidden_size
        self.q_proj = nn.Linear(E, E, bias=False)
        self.k_proj = nn.Linear(E, E, bias=False)
        self.v_proj = nn.Linear(E, E, bias=False)
        self.o_proj = nn.Linear(E, E, bias=False)

    def forward(self, x, attention_mask=None, past_key_values=None, use_cache=False):
        B, T, E = x.shape
        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        if self.qk_norm:
            # Gainless per-head RMSNorm over head_dim: attention logits are
            # bounded by construction, so attention cannot sharpen from q/k
            # magnitude growth. Trained-with means served-with.
            q = rms_norm(q, self.eps)
            k = rms_norm(k, self.eps)

        present = None
        if past_key_values is not None and hasattr(past_key_values, 'update'):
            k, v = past_key_values.update(k, v, self.layer_idx)
        else:
            if past_key_values:
                prev = past_key_values[self.layer_idx]
                if prev is not None and prev[0] is not None:
                    k = torch.cat([prev[0], k], dim=2)
                    v = torch.cat([prev[1], v], dim=2)
            if use_cache:
                present = (k, v)

        scores = torch.matmul(q, k.transpose(-1, -2)) * self.scaling
        if attention_mask is not None:
            scores = scores + attention_mask
        probs = F.softmax(scores.float(), dim=-1).to(q.dtype)
        ctx = torch.matmul(probs, v).transpose(1, 2).reshape(B, T, E)
        return self.o_proj(ctx), present


class AndreaMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(F.relu(self.up_proj(x)))


class AndreaDecoderLayer(nn.Module):
    def __init__(self, config, layer_idx=0):
        super().__init__()
        self.eps = config.rms_norm_eps
        self.self_attn = AndreaAttention(config, layer_idx)
        self.mlp = AndreaMLP(config)

    def forward(self, x, attention_mask=None, past_key_values=None, use_cache=False):
        attn_out, present = self.self_attn(
            rms_norm(x, self.eps), attention_mask, past_key_values, use_cache)
        x = x + attn_out
        x = x + self.mlp(rms_norm(x, self.eps))
        return x, present


class AndreaPreTrainedModel(PreTrainedModel):
    config_class = AndreaConfig
    base_model_prefix = 'model'
    supports_gradient_checkpointing = False
    _no_split_modules = ['AndreaDecoderLayer']

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=std)
            if isinstance(module, nn.Linear) and module.bias is not None:
                module.bias.data.zero_()


class AndreaModel(AndreaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.embed_positions = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.layers = nn.ModuleList(
            [AndreaDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
        self.eps = config.rms_norm_eps
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def _causal_mask(self, B, T, kv_len, attention_mask, dtype, device):
        """Additive mask of shape (B, 1, T, kv_len). Blocked positions get -inf."""
        neg = torch.finfo(dtype).min
        # Query i (absolute position kv_len - T + i) may attend to key j <= i.
        q_pos = torch.arange(kv_len - T, kv_len, device=device).unsqueeze(-1)
        k_pos = torch.arange(kv_len, device=device).unsqueeze(0)
        allowed = k_pos <= q_pos
        # Sliding window: also require j >= i - w + 1. A model trained under a
        # window and served without this attends context it never learned.
        window = getattr(self.config, 'sliding_window', None)
        if window:
            allowed = allowed & (k_pos > q_pos - window)
        # Built directly in `dtype`: python scalars would make a float32 tensor
        # first, and finfo(float64).min does not survive that round trip.
        zero_t = torch.zeros((), dtype=dtype, device=device)
        neg_t = torch.full((), neg, dtype=dtype, device=device)
        mask = torch.where(allowed, zero_t, neg_t)
        mask = mask.unsqueeze(0).unsqueeze(0).expand(B, 1, T, kv_len).clone()
        if attention_mask is not None:
            pad = (1.0 - attention_mask[:, None, None, :].to(dtype)) * neg
            mask = mask + pad
        return mask

    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
                past_key_values=None, inputs_embeds=None, use_cache=None,
                output_attentions=False, output_hidden_states=False,
                return_dict=True, **kwargs):
        use_cache = self.config.use_cache if use_cache is None else use_cache
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        B, T, _ = inputs_embeds.shape
        device = inputs_embeds.device

        past_len = cache_length(past_key_values)
        if position_ids is None:
            position_ids = torch.arange(past_len, past_len + T, device=device).unsqueeze(0)
        if position_ids.max() >= self.config.max_position_embeddings:
            raise ValueError(
                f'position {int(position_ids.max())} exceeds max_position_embeddings '
                f'{self.config.max_position_embeddings}: ANDREA uses learned absolute '
                f'positions and cannot extrapolate past its trained context window')

        x = inputs_embeds + self.embed_positions(position_ids)
        x = rms_norm(x, self.eps)  # input norm, before layer 0

        mask = self._causal_mask(B, T, past_len + T, attention_mask, x.dtype, device)

        uses_cache_object = past_key_values is not None and hasattr(past_key_values, 'update')
        presents = [] if (use_cache and not uses_cache_object) else None
        all_hidden = [] if output_hidden_states else None
        for layer in self.layers:
            if output_hidden_states:
                all_hidden.append(x)
            x, present = layer(x, mask, past_key_values, use_cache)
            if presents is not None:
                presents.append(present)
        if output_hidden_states:
            all_hidden.append(x)

        # No final norm — that is deliberate, see the module docstring.
        if uses_cache_object:
            next_cache = past_key_values if use_cache else None
        else:
            next_cache = tuple(presents) if presents else None

        if not return_dict:
            return tuple(v for v in (x, next_cache, all_hidden) if v is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=x,
            past_key_values=next_cache,
            hidden_states=tuple(all_hidden) if all_hidden else None,
        )


class AndreaForCausalLM(AndreaPreTrainedModel, GenerationMixin):
    _tied_weights_keys = []  # lm_head is genuinely untied from wte
    # Our cache is a plain tuple of (key, value) per layer rather than a Cache
    # object, so opt out of the machinery that would hand us one.
    _supports_cache_class = False
    _can_compile_fullgraph = False

    def __init__(self, config):
        super().__init__(config)
        self.model = AndreaModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new):
        self.lm_head = new

    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
                past_key_values=None, inputs_embeds=None, labels=None,
                use_cache=None, output_attentions=False, output_hidden_states=False,
                return_dict=True, **kwargs):
        out = self.model(
            input_ids=input_ids, attention_mask=attention_mask,
            position_ids=position_ids, past_key_values=past_key_values,
            inputs_embeds=inputs_embeds, use_cache=use_cache,
            output_hidden_states=output_hidden_states, return_dict=True)
        hidden = out.last_hidden_state
        if getattr(self.config, 'final_norm', False):
            hidden = rms_norm(hidden, self.config.rms_norm_eps)
        logits = self.lm_head(hidden).float()

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1))

        if not return_dict:
            return tuple(v for v in (loss, logits, out.past_key_values) if v is not None)
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=out.past_key_values,
            hidden_states=out.hidden_states,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
                                      attention_mask=None, **kwargs):
        past_len = cache_length(past_key_values)
        if past_len:
            input_ids = input_ids[:, past_len:]
        return {
            'input_ids': input_ids,
            'past_key_values': past_key_values,
            'attention_mask': attention_mask,
            'use_cache': kwargs.get('use_cache', True),
        }

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        if hasattr(past_key_values, 'reorder_cache'):
            past_key_values.reorder_cache(beam_idx)
            return past_key_values
        return tuple(
            tuple(state.index_select(0, beam_idx.to(state.device)) for state in layer)
            for layer in past_key_values)