"""CaptionBERT with the final key-at-every-position rotary adaptation.

This file and configuration_captionbert.py are deliberately self-contained for
AutoModel remote-code loading. No training package is needed for inference.
"""
import copy
from dataclasses import dataclass
import math
from typing import Optional

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
from transformers import PreTrainedModel
from transformers.utils import ModelOutput
from .configuration_captionbert import CaptionBertConfig


@dataclass
class CaptionBertOutput(ModelOutput):
    last_hidden_state: Optional[torch.Tensor] = None
    pooler_output: Optional[torch.Tensor] = None
    hidden_states: Optional[tuple] = None
    attentions: Optional[tuple] = None
    pointer_loss: Optional[torch.Tensor] = None


def hierarchical_frequencies():
    """v0.1.103's 8x3 form verbatim in value: 2 relative, 2 content, 12 anchor pairs."""
    base = 10000.0 * 4.0 ** (4.0 / 2.0)
    w = [base ** (-j / 2) for j in range(2)] + [0.0, 0.0]
    kinds = [0, 0, 1, 1]
    for level in range(3):
        w += [2 * math.pi * k / 9 for k in range(1, 5)]
        kinds += [2 + level] * 4
    return torch.tensor(w, dtype=torch.float64), torch.tensor(kinds, dtype=torch.long)


def rotate_half(t, angle):
    pairs = t.shape[-1] // 2
    a, b = t[..., :pairs].float(), t[..., pairs:].float()
    c, s = angle.float().cos(), angle.float().sin()
    return torch.cat((a * c - b * s, a * s + b * c), dim=-1).to(t.dtype)


def hold_pairs(key, on, rho):
    pairs = key.shape[-1] // 2
    a, b = key[..., :pairs], key[..., pairs:]
    # Clamp before sqrt: the same forward floor, with finite derivatives at
    # zero/padded pairs (sqrt(0) followed by clamp can yield 0 * inf in backward).
    length = (a.float().square() + b.float().square()).clamp_min(1e-12).sqrt()
    factor = torch.where(on[..., None], rho / length, torch.ones_like(length)).to(key.dtype)
    return torch.cat((a * factor, b * factor), dim=-1)


def word_features(ids, valid, continuation, special, pairs=16, seed=1662197431):
    """Stateless whole-WordPiece-key binding; all pieces and offsets contribute.

    The source's fixed byte/offset phase table becomes a fixed 31-bit counter
    hash over WordPiece IDs and reverse offsets. Integer accumulation is exact
    and independent of scatter order. It has no key-length truncation. This is
    a declared tokenizer adaptation, not the original byte table's draw.
    """
    batch, length = ids.shape
    position = torch.arange(length, device=ids.device)[None].expand(batch, -1)
    active = valid & ~special[ids]
    starts = valid & ~continuation[ids]
    group = starts.long().cumsum(1)
    first = torch.where(starts, position, torch.zeros_like(position)).cummax(1).values
    offset = position - first
    counts = torch.zeros(batch, length + 1, device=ids.device, dtype=torch.long)
    counts.scatter_add_(1, group, active.long())
    reverse = (counts.gather(1, group) - 1 - offset).clamp_min(0)
    pair = torch.arange(1, pairs + 1, device=ids.device, dtype=torch.long)
    mask31 = 2147483647
    u = (ids[..., None].long() * 1103515245 + reverse[..., None] * 1234567891
         + pair * 1013904223 + int(seed)) & mask31
    u = ((u ^ (u >> 16)) * 73244475) & mask31
    u = ((u ^ (u >> 16)) * 73244475) & mask31
    u = (u ^ (u >> 16)) & mask31
    sums = torch.zeros(batch, length + 1, pairs, device=ids.device, dtype=torch.long)
    sums.scatter_add_(1, group[..., None].expand(-1, -1, pairs), u * active[..., None])
    keys = sums.gather(1, group[..., None].expand(-1, -1, pairs)) & mask31
    phase = keys.float() * (2 * math.pi / 2147483648.0)
    next_start = F.pad(starts[:, 1:], (0, 1), value=True)
    next_active = F.pad(active[:, 1:], (0, 1), value=False)
    ends = active & (next_start | ~next_active)

    # First occurrence's ending is the record; later occurrences are queries.
    # Two independent 31-bit components identify candidate groups. Training
    # targets are computed from exact token tuples separately, never this hash.
    identity = keys[..., 0] | (keys[..., 1] << 31)
    identity = torch.where(active, identity, torch.full_like(identity, -1))
    order = identity.argsort(dim=1, stable=True)
    sorted_id = identity.gather(1, order)
    changed = F.pad(sorted_id[:, 1:] != sorted_id[:, :-1], (1, 0), value=True)
    sorted_group = changed.long().cumsum(1)
    candidate = torch.where(ends, position, torch.full_like(position, length)).gather(1, order)
    minimum = torch.full((batch, length + 1), length, device=ids.device, dtype=torch.long)
    minimum.scatter_reduce_(1, sorted_group, candidate, reduce="amin", include_self=True)
    sorted_record = minimum.gather(1, sorted_group)
    record_at = torch.empty_like(sorted_record).scatter_(1, order, sorted_record)
    records = ends & (record_at == position)
    return phase, active, ends, records


class CaptionBertLayer(nn.TransformerEncoderLayer):
    def __init__(self, config):
        super().__init__(config.hidden_size, config.num_attention_heads,
                         config.intermediate_size, config.hidden_dropout_prob,
                         activation="gelu", batch_first=True, norm_first=True)
        self.lookup_head = config.lookup_head
        self.rho = config.key_rho

    def forward(self, x, valid, aq, ak, phase, binding_on, records,
                pointer_targets=None, bind=False, output_attentions=False):
        z = self.norm1(x)
        batch, length, width = z.shape
        heads = self.self_attn.num_heads
        qkv = F.linear(z, self.self_attn.in_proj_weight, self.self_attn.in_proj_bias)
        q, k, v = [a.reshape(batch, length, heads, width // heads).transpose(1, 2)
                   for a in qkv.chunk(3, -1)]
        qr, kr = rotate_half(q, aq), rotate_half(k, ak)
        if bind:
            h = self.lookup_head
            qa = torch.where(binding_on[..., None], phase, aq[:, 0])
            ka = torch.where(binding_on[..., None], phase, ak[:, 0])
            kh = hold_pairs(k[:, h], records, self.rho) if self.rho else k[:, h]
            qh, kh = rotate_half(q[:, h], qa), rotate_half(kh, ka)
            qr = torch.cat((qr[:, :h], qh[:, None], qr[:, h + 1:]), 1)
            kr = torch.cat((kr[:, :h], kh[:, None], kr[:, h + 1:]), 1)
        pointer_loss = x.new_zeros((), dtype=torch.float32)
        if pointer_targets is not None:
            logits = qr[:, self.lookup_head].float() @ kr[:, self.lookup_head].float().transpose(1, 2)
            logits = logits / math.sqrt(width // heads)
            logits = logits.masked_fill(~valid[:, None, :], torch.finfo(logits.dtype).min)
            selected = pointer_targets >= 0
            nll = -logits.log_softmax(-1).gather(-1, pointer_targets.clamp_min(0)[..., None]).squeeze(-1)
            pointer_loss = (nll * selected).sum() / selected.sum().clamp_min(1)
        y = F.scaled_dot_product_attention(qr, kr, v, attn_mask=valid[:, None, None, :],
            dropout_p=self.self_attn.dropout if self.training else 0.0)
        x = x + self.dropout1(self.self_attn.out_proj(y.transpose(1, 2).reshape(batch, length, width)))
        x = x + self._ff_block(self.norm2(x))
        if output_attentions:
            logits = qr.float() @ kr.float().transpose(-1, -2) / math.sqrt(width // heads)
            attention = logits.masked_fill(~valid[:, None, None, :], torch.finfo(logits.dtype).min).softmax(-1)
            return x, pointer_loss, attention
        return x, pointer_loss


class CaptionBertModel(PreTrainedModel):
    config_class = CaptionBertConfig
    base_model_prefix = "captionbert"
    main_input_name = "input_ids"
    supports_gradient_checkpointing = True
    _supports_sdpa = True

    def __init__(self, config):
        super().__init__(config)
        self.gradient_checkpointing = config.gradient_checkpointing
        self.token_emb = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.emb_norm = nn.LayerNorm(config.hidden_size)
        self.emb_drop = nn.Dropout(config.hidden_dropout_prob)
        template = CaptionBertLayer(config)
        self.layers = nn.ModuleList([copy.deepcopy(template) for _ in range(config.num_hidden_layers)])
        self.output_proj = nn.Sequential(nn.Linear(config.hidden_size, config.hidden_size), nn.GELU(),
                                        nn.LayerNorm(config.hidden_size), nn.Linear(config.hidden_size, config.output_dim))
        self.structured = config.rotary_mode in ("hierarchical", "key_end", "key_all")
        self.value_address = nn.Parameter(torch.zeros(3, 8, config.hidden_size)) if self.structured and config.use_value_addresses else None
        continuation = torch.zeros(config.vocab_size, dtype=torch.bool)
        continuation[config.continuation_ids] = True
        special = torch.zeros_like(continuation)
        special[[i for i in config.special_ids if i < config.vocab_size]] = True
        self.register_buffer("wordpiece_continuation", continuation)
        self.register_buffer("special_token", special)
        if self.structured:
            frequencies, kinds = hierarchical_frequencies()
        else:
            pairs = config.hidden_size // config.num_attention_heads // 2
            dim = pairs * 2
            base = 10000.0 * (4.0 ** (dim / (dim - 2)) if config.rotary_mode == "ntk" else 1.0)
            frequencies = base ** (-torch.arange(pairs, dtype=torch.float64) / pairs)
            kinds = torch.zeros(pairs, dtype=torch.long)
        # Save the small immutable rotary tables alongside the weights. Current
        # Transformers materializes non-persistent buffers with empty_like when
        # loading on meta; constructor-only initialization would then be lost.
        self.register_buffer("frequencies", frequencies.float())
        self.register_buffer("kinds", kinds)
        self.post_init()

    def _init_weights(self, module):
        # Keep PyTorch's native CaptionBERT initialization; value-address tables
        # are born zero. No HF reinitialization changes the paired controls.
        pass

    def get_input_embeddings(self):
        return self.token_emb

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

    def addresses(self, ids, valid):
        punctuation = torch.zeros_like(valid)
        for token in self.config.punctuation_ids:
            punctuation = punctuation | (ids == token)
        punctuation = punctuation & valid
        sentence = punctuation.long().cumsum(1) - punctuation.long()
        digits = torch.stack(((sentence // 64) % 8, (sentence // 8) % 8, sentence % 8), -1)
        return torch.where((valid & (sentence < 512))[..., None], digits, torch.full_like(digits, 8))

    def angles(self, ids, valid, query_addresses=None, key_addresses=None):
        batch, length = ids.shape
        position = torch.arange(length, device=ids.device, dtype=torch.float32)
        base = position[:, None] * torch.where(self.kinds == 0, self.frequencies, 0)[None]
        aq = base[None].expand(batch, -1, -1).clone()
        ak = aq.clone()
        addresses = self.addresses(ids, valid) if self.structured else None
        if self.structured:
            dq = addresses if query_addresses is None else query_addresses
            dk = addresses if key_addresses is None else key_addresses
            if dq.shape != addresses.shape or dk.shape != addresses.shape:
                raise ValueError("Address streams must have shape (batch, length, 3)")
            for level in range(3):
                selected = self.kinds == level + 2
                aq[..., selected] = (dq[..., level:level + 1].float() - 8) * self.frequencies[selected]
                ak[..., selected] = (dk[..., level:level + 1].float() - 8) * self.frequencies[selected]
            addresses = dk
        return aq[:, None], ak[:, None], addresses

    def forward(self, input_ids, attention_mask=None, token_type_ids=None,
                query_addresses=None, key_addresses=None, pointer_targets=None,
                output_hidden_states=None, output_attentions=None, return_dict=None, **kwargs):
        if input_ids.ndim != 2 or not 1 <= input_ids.shape[1] <= self.config.max_position_embeddings:
            raise ValueError("input_ids must be (batch, length), with length in the configured range")
        valid = input_ids.ne(self.config.pad_token_id) if attention_mask is None else attention_mask.bool()
        if valid.shape != input_ids.shape:
            raise ValueError("attention_mask must match input_ids")
        output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
        output_attentions = self.config.output_attentions if output_attentions is None else output_attentions
        return_dict = self.config.use_return_dict if return_dict is None else return_dict
        aq, ak, addresses = self.angles(input_ids, valid, query_addresses, key_addresses)
        binding = self.config.rotary_mode in ("key_end", "key_all")
        if binding:
            phase, active, ends, records = word_features(input_ids, valid, self.wordpiece_continuation,
                                                       self.special_token, seed=self.config.phase_seed)
            on = ends if self.config.rotary_mode == "key_end" else active
        else:
            phase = aq[:, 0]
            on = records = torch.zeros_like(valid)
        x = self.token_emb(input_ids)
        if self.value_address is not None:
            for level in range(3):
                d = addresses[..., level].long()
                d = torch.where((d >= 0) & (d < 8), d, torch.full_like(d, 8))
                table = torch.cat((self.value_address[level], x.new_zeros(1, x.shape[-1])), 0)
                x = x + F.embedding(d, table, padding_idx=8)
        x = self.emb_drop(self.emb_norm(x))
        states = (x,) if output_hidden_states else None
        attentions = () if output_attentions else None
        pointer_loss = x.new_zeros((), dtype=torch.float32)
        for index, layer in enumerate(self.layers):
            target = pointer_targets if index == self.config.lookup_layer else None
            if target is not None and self.config.pointer_positions == "end":
                if not binding:
                    raise ValueError("End-position supervision requires the word-binding features")
                target = torch.where(ends, target, torch.full_like(target, -1))
            args = (x, valid, aq, ak, phase, on, records)
            opts = dict(pointer_targets=target, bind=binding and index == self.config.lookup_layer,
                        output_attentions=bool(output_attentions))
            if self.training and self.gradient_checkpointing and not output_attentions:
                result = checkpoint(layer, *args, use_reentrant=False, **opts)
            else:
                result = layer(*args, **opts)
            x, term = result[:2]
            pointer_loss = pointer_loss + term
            if output_hidden_states:
                states += (x,)
            if output_attentions:
                attentions += (result[2],)
        mask = valid[..., None].float()
        pooled = (x.float() * mask).sum(1) / mask.sum(1).clamp_min(1)
        embedding = F.normalize(self.output_proj(pooled).float(), dim=-1)
        embedding = torch.where(valid.any(1)[:, None], embedding, torch.zeros_like(embedding))
        if not return_dict:
            return tuple(v for v in (x, embedding, states, attentions, pointer_loss) if v is not None)
        return CaptionBertOutput(last_hidden_state=x, pooler_output=embedding,
                                 hidden_states=states, attentions=attentions, pointer_loss=pointer_loss)