"""vela3-fast UltraBERT-380M encoder as a fixed-length, eager-attention graph (ids + mask -> final-norm hidden states).

Same contract as the CPU path of modeling_ultrabert.UltrabertModel.forward: local layers attend to keys with |i - j| <= 128, global
layers to all keys, padding masks keys only; RoPE theta per layer (10k local, 2.56M global). Masks use -1e4 instead of finfo.min
(fp16-safe); cos / sin and the two masks are precomputed for the bucket length."""
import sys
import torch
from torch import nn

MASK_VALUE = -1e4
sys.path.insert(0, "src")


def load_decision(path="src"):
    from vela3_encoder import DecisionEncoder
    return DecisionEncoder.load(path, device="cpu")


class Vela3EncoderExport(nn.Module):
    def __init__(self, enc, length):
        super().__init__()
        cfg = enc.config
        cfg._attn_implementation = "eager"
        self.length, self.layers, self.embed, self.norm = length, enc.layers, enc.embed_tokens, enc.norm
        self.windows, self.thetas = list(cfg.attention_windows), list(cfg.rope_theta)
        pos = torch.arange(length).unsqueeze(0)
        probe = torch.zeros(1, length, cfg.hidden_size)
        self.theta_key = {}
        for i, base in enumerate(sorted(set(self.thetas))):
            cos, sin = enc.rotary_emb[str(base).replace(".", "_")](probe, pos)
            self.register_buffer(f"cos{i}", cos.detach().clone()); self.register_buffer(f"sin{i}", sin.detach().clone())
            self.theta_key[base] = i
        idx = torch.arange(length)
        dist = (idx[:, None] - idx[None, :]).abs()
        self.window_key = {}
        for i, w in enumerate(sorted({w for w in self.windows if w is not None})):
            self.register_buffer(f"band{i}", torch.where(dist <= w, 0.0, MASK_VALUE).view(1, 1, length, length))
            self.window_key[w] = i

    def forward(self, input_ids, attention_mask):
        pad = (1.0 - attention_mask.float()).view(1, 1, 1, self.length) * MASK_VALUE
        bands = {w: getattr(self, f"band{i}") + pad for w, i in self.window_key.items()}
        x = self.embed(input_ids.long())
        for layer, w, t in zip(self.layers, self.windows, self.thetas):
            k = self.theta_key[t]
            x = layer(x, attention_mask=pad if w is None else bands[w],
                      position_embeddings=(getattr(self, f"cos{k}"), getattr(self, f"sin{k}")))
        return self.norm(x)


if __name__ == "__main__":  # parity of the export graph vs the release encoder (sdpa, fp32 CPU) on a real decision input
    import json
    import pyarrow.parquet as pq
    m = load_decision()
    m.encoder.config._attn_implementation = "eager"
    r = pq.read_table("data/typed/all/test-00000-of-00001.parquet").to_pylist()[0]
    state, qs = json.loads(r["state"]), json.loads(r["questions"])
    q = next(q for q in qs.values() if q["type"] == "choice")
    ids = m.encode(state, q)["ids"]; n = len(ids)
    with torch.no_grad():
        ref = m.encoder(input_ids=torch.tensor([ids]), attention_mask=torch.ones(1, n, dtype=torch.long)).last_hidden_state[0]
        for L in (n, 512):
            g = Vela3EncoderExport(m.encoder, L).eval()
            x = torch.zeros(1, L, dtype=torch.int32); x[0, :n] = torch.tensor(ids); mk = torch.zeros(1, L, dtype=torch.int32); mk[0, :n] = 1
            h = g(x, mk)[0, :n]
            print(f"n={n} L={L} max|dh|={(h - ref).abs().max():.2e} min cos={nn.functional.cosine_similarity(h, ref, dim=-1).min():.6f}")