"""ModernBERT-architecture encoder (Ettin weights), written from scratch.

Attention is controlled by per-token vectors instead of a precomputed mask:
  position_ids: RoPE positions (branches restart after the state).
  segment_ids:  0 = state, k > 0 = question branch k, -1 = padding.
  option_ids:   0 = not an option, m > 0 = option m of its question.
A token attends to state tokens, its question's own tokens, and (for option tokens)
its own option. Local layers additionally restrict attention to
|pos_i - pos_j| <= window // 2, measured in position ids (not sequence index) so
packed and separate requests match.
Comparison mode lets each question and its options attend to the whole question branch.
"""

import json
from dataclasses import dataclass

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from torch import nn


@dataclass
class Config:
    vocab_size: int
    hidden_size: int
    intermediate_size: int
    num_hidden_layers: int
    num_attention_heads: int
    global_attn_every_n_layers: int = 3
    local_attention: int = 128
    global_rope_theta: float = 160000.0
    local_rope_theta: float = 10000.0
    norm_eps: float = 1e-5
    pad_token_id: int = 50283
    cls_token_id: int = 50281
    sep_token_id: int = 50282
    model_type: str = "modernbert"
    compare_options: bool = False
    date_facts: bool = False  # an input-format setting: facts.date_facts runs on every state (model.pack_request)

    @classmethod
    def from_hf(cls, d: dict) -> "Config":
        """Reads a Hugging Face config or a run's saved config, ignoring fields this class does not define."""
        return cls(**{k: d[k] for k in cls.__dataclass_fields__ if k in d})


def rope(x: torch.Tensor, pos: torch.Tensor, theta: float) -> torch.Tensor:
    """x: [B, H, L, D], pos: [B, L]. Half-split rotation (rotate_half convention)."""
    d = x.shape[-1]
    inv_freq = 1.0 / (theta ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d))
    freqs = pos[:, None, :, None].float() * inv_freq  # [B, 1, L, D/2]
    cos, sin = freqs.cos().to(x.dtype), freqs.sin().to(x.dtype)
    x1, x2 = x[..., : d // 2], x[..., d // 2 :]
    return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)


class Layer(nn.Module):
    def __init__(self, c: Config, idx: int):
        super().__init__()
        self.is_global = idx % c.global_attn_every_n_layers == 0
        self.theta = c.global_rope_theta if self.is_global else c.local_rope_theta
        self.heads = c.num_attention_heads
        self.attn_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False) if idx else nn.Identity()
        self.Wqkv = nn.Linear(c.hidden_size, 3 * c.hidden_size, bias=False)
        self.Wo = nn.Linear(c.hidden_size, c.hidden_size, bias=False)
        self.mlp_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False)
        self.Wi = nn.Linear(c.hidden_size, 2 * c.intermediate_size, bias=False)
        self.Wo_mlp = nn.Linear(c.intermediate_size, c.hidden_size, bias=False)

    def forward(self, h, pos, global_mask, local_mask, prefix=None):
        """prefix: cached (k, v) of earlier tokens, prepended to this call's keys and values."""
        B, L, _ = h.shape
        q, k, v = self.Wqkv(self.attn_norm(h)).view(B, L, 3, self.heads, -1).permute(2, 0, 3, 1, 4)
        q, k = rope(q, pos, self.theta), rope(k, pos, self.theta)
        kv = (k, v)
        if prefix is not None:
            k, v = torch.cat([prefix[0], k], 2), torch.cat([prefix[1], v], 2)
        mask = global_mask if self.is_global else local_mask
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        h = h + self.Wo(a.transpose(1, 2).reshape(B, L, -1))
        x, gate = self.Wi(self.mlp_norm(h)).chunk(2, dim=-1)
        return h + self.Wo_mlp(F.gelu(x) * gate), kv


class Encoder(nn.Module):
    def __init__(self, c: Config):
        super().__init__()
        self.config = c
        self.tok_embeddings = nn.Embedding(c.vocab_size, c.hidden_size)
        self.emb_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False)
        self.layers = nn.ModuleList(Layer(c, i) for i in range(c.num_hidden_layers))
        self.final_norm = nn.LayerNorm(c.hidden_size, eps=c.norm_eps, bias=False)

    @staticmethod
    def masks(pos, seg, opt, window, pos_k=None, seg_k=None, opt_k=None, *, compare_options=False):
        """Boolean [B, 1, Lq, Lk] masks, True = may attend. Global and local variants.

        Keys default to the queries; pass *_k to attend over a longer key sequence (cached state + queries).
        The masks are dense, so their memory grows with Lq * Lk.
        """
        pos_k, seg_k, opt_k = (x if x is not None else y for x, y in ((pos_k, pos), (seg_k, seg), (opt_k, opt)))
        si, sj, oi, oj = seg[:, :, None], seg_k[:, None, :], opt[:, :, None], opt_k[:, None, :]
        within = (sj == si) & ((si > 0) if compare_options else ((oj == 0) | (oj == oi)))
        allowed = (sj != -1) & ((sj == 0) | within)
        local = allowed & ((pos[:, :, None] - pos_k[:, None, :]).abs() <= window // 2)
        return allowed[:, None], local[:, None]

    def forward(self, input_ids, position_ids=None, segment_ids=None, option_ids=None, inputs_embeds=None):
        B, L = input_ids.shape
        if position_ids is None:
            position_ids = torch.arange(L, device=input_ids.device).expand(B, L)
        if segment_ids is None:
            segment_ids = torch.zeros_like(input_ids)
        if option_ids is None:
            option_ids = torch.zeros_like(input_ids)
        emb = self.tok_embeddings(input_ids) if inputs_embeds is None else inputs_embeds
        return self.run(emb, position_ids, segment_ids, option_ids)[0]

    def run(self, emb, pos, seg, opt, cache=None):
        """Encode embeddings, optionally on top of a cached prefix. Returns (hidden, cache for these tokens).

        A cache holds each layer's rotated keys and values plus the pos/seg/opt ids of its tokens. The state
        never attends to questions, so a state encoded once can serve any number of later questions.
        """
        if cache is None:
            g, l = self.masks(pos, seg, opt, self.config.local_attention, compare_options=self.config.compare_options)
        else:
            g, l = self.masks(pos, seg, opt, self.config.local_attention,
                              torch.cat([cache["pos"], pos], 1), torch.cat([cache["seg"], seg], 1), torch.cat([cache["opt"], opt], 1),
                              compare_options=self.config.compare_options)
        h, kvs = self.emb_norm(emb), []
        for i, layer in enumerate(self.layers):
            h, kv = layer(h, pos, g, l, None if cache is None else cache["kv"][i])
            kvs.append(kv)
        return self.final_norm(h), {"kv": kvs, "pos": pos, "seg": seg, "opt": opt}

    @classmethod
    def from_pretrained(cls, repo: str) -> "Encoder":
        with open(hf_hub_download(repo, "config.json")) as f:
            model = cls(Config.from_hf(json.load(f)))
        sd = torch.load(hf_hub_download(repo, "pytorch_model.bin"), map_location="cpu", weights_only=True)
        rename = {
            "model.embeddings.tok_embeddings.": "tok_embeddings.",
            "model.embeddings.norm.": "emb_norm.",
            "model.final_norm.": "final_norm.",
            ".attn.Wqkv.": ".Wqkv.",
            ".attn.Wo.": ".Wo.",
            ".mlp.Wi.": ".Wi.",
            ".mlp.Wo.": ".Wo_mlp.",
            "model.layers.": "layers.",
        }
        out = {}
        for k, v in sd.items():
            if not k.startswith("model."):  # MLM head, unused
                continue
            for a, b in rename.items():
                k = k.replace(a, b)
            out[k] = v
        model.load_state_dict(out, strict=True)
        return model.eval()