"""D1A (Gemma 4 text) -> Core ML pilot (jonpol01/d1a#224): rebuild a D1A MLX export as a torch module, split into
  - host side: the two embedding tables (token embedding, per-layer embedding: 2.75B of E2B's 4.6B parameters) read
    row by row from the export's safetensors (dequantized per request, never held whole), and
  - Core: everything that computes (per-layer projection, 35 decoder layers, final norm, pointer head), which is what
    gets converted to Core ML.
Inputs follow D1A's packed encoding (d1a.core.encoding.encode): one record, block-causal mask, positions restarting per
question after the state, readout at each question's <decide> and </opt> tokens."""
import json
import math
from pathlib import Path

import numpy as np
import torch
from safetensors import safe_open
from torch import nn
from transformers import AutoConfig
from transformers.models.gemma4 import modeling_gemma4 as G

PREFIX = "language_model.model."


def dequant(packed, scales, biases, bits, group=64):
    """MLX affine dequantization: packed uint32 [out, in*bits/32] (lowest bits first) -> float32 [out, in]."""
    p = packed.view(torch.int32).to(torch.int64) & 0xFFFFFFFF
    per = 32 // bits
    shifts = torch.arange(per, dtype=torch.int64) * bits
    q = ((p[..., None] >> shifts) & ((1 << bits) - 1)).reshape(p.shape[0], -1).to(torch.float32)
    s, b = scales.to(torch.float32), biases.to(torch.float32)
    return (q.view(q.shape[0], -1, group) * s[..., None] + b[..., None]).reshape(q.shape[0], -1)


class _Shards:
    """model.safetensors or model-0000N-of-0000M.safetensors (via model.safetensors.index.json), read like one safe_open."""

    def __init__(self, path):
        idx = Path(path) / "model.safetensors.index.json"
        if idx.exists():
            wm = json.loads(idx.read_text())["weight_map"]
            files = {f: safe_open(str(Path(path) / f), "pt") for f in sorted(set(wm.values()))}
            self.where = {k: files[f] for k, f in wm.items()}
        else:   # a released Core ML bundle ships only the embedding tables, as embeddings.safetensors
            name = "model.safetensors" if (Path(path) / "model.safetensors").exists() else "embeddings.safetensors"
            f = safe_open(str(Path(path) / name), "pt")
            self.where = {k: f for k in f.keys()}

    def keys(self):
        return self.where.keys()

    def get_tensor(self, k):
        return self.where[k].get_tensor(k)

    def get_slice(self, k):
        return self.where[k].get_slice(k)


class Export:
    """A D1A MLX export directory (config.json, d1a_config.json, model.safetensors, head.safetensors)."""

    def __init__(self, path):
        self.path = Path(path)
        self.cfg = AutoConfig.from_pretrained(self.path)
        self.tc = self.cfg.text_config
        self.tc._attn_implementation = "eager"
        self.d1a = json.loads((self.path / "d1a_config.json").read_text())
        q = self.d1a["quantization"]
        self.bits = {"default": (q["bits"], q["group_size"]), "embed_tokens_per_layer": (q[PREFIX.replace("model.", "") + "model.embed_tokens_per_layer"]["bits"] if False else q["language_model.model.embed_tokens_per_layer"]["bits"], q["language_model.model.embed_tokens_per_layer"]["group_size"])}
        self.f = _Shards(self.path)

    def bits_of(self, name):
        return self.bits["embed_tokens_per_layer"] if name.endswith("embed_tokens_per_layer") else self.bits["default"]

    def tensor(self, name):
        """A float32 tensor for `name` (without PREFIX), dequantized if it is stored quantized."""
        full = PREFIX + name
        keys = set(self.f.keys())
        if full + ".weight" in keys and full + ".scales" in keys:
            bits, group = self.bits_of(name)
            return dequant(self.f.get_tensor(full + ".weight"), self.f.get_tensor(full + ".scales"), self.f.get_tensor(full + ".biases"), bits, group)
        return self.f.get_tensor(full).to(torch.float32)

    def text_model(self, dtype=torch.float32):
        """The Gemma 4 text trunk with every non-embedding weight loaded (embedding tables stay empty: host side)."""
        with torch.device("meta"):
            m = G.Gemma4TextModel(self.tc)
        m = m.to_empty(device="cpu")
        state = {}
        keys = set(self.f.keys())
        for name, p in m.named_parameters():
            if name.startswith("embed_tokens"):
                continue
            base = name[: -len(".weight")] if name.endswith(".weight") else name
            if PREFIX + base + ".scales" in keys:
                state[name] = self.tensor(base)
            else:
                state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
            assert state[name].shape == p.shape, (name, state[name].shape, p.shape)
        for name, b in m.named_buffers():
            if PREFIX + name in keys:
                state[name] = self.f.get_tensor(PREFIX + name).to(torch.float32)
        missing, unexpected = m.load_state_dict(state, strict=False)
        missing = [k for k in missing if not k.startswith("embed_tokens")]
        assert not missing and not unexpected, (missing, unexpected)
        # buffers transformers computes itself (rotary inv_freq, embed_scale): re-run their init on CPU
        m.rotary_emb = G.Gemma4TextRotaryEmbedding(self.tc)
        return m.to(dtype).eval()

    def head(self):
        h = safe_open(str(self.path / "head.safetensors"), "pt")
        dp, d = h.get_tensor("q.weight").shape                     # pointer width, hidden size (E2B 1536, E4B 2560)
        q, k = nn.Linear(d, dp), nn.Linear(d, dp)
        q.weight.data, q.bias.data = h.get_tensor("q.weight").float(), h.get_tensor("q.bias").float()
        k.weight.data, k.bias.data = h.get_tensor("k.weight").float(), h.get_tensor("k.bias").float()
        return q, k


class HostEmbed:
    """Row lookups of the two embedding tables, dequantized per request: emb = E[ids] * sqrt(hidden) and
    ple = E_pl[ids] * sqrt(ple_dim), reshaped [L, layers, ple_dim] (transformers' get_per_layer_inputs)."""

    def __init__(self, ex):
        self.ex, self.tc = ex, ex.tc
        self.f = ex.f

    def _rows(self, name, ids):
        full = PREFIX + name
        bits, group = self.ex.bits_of(name)
        w, s, b = (self.f.get_slice(full + suf) for suf in (".weight", ".scales", ".biases"))
        uniq, inv = np.unique(np.asarray(ids), return_inverse=True)
        rows = [dequant(w[int(i): int(i) + 1], s[int(i): int(i) + 1], b[int(i): int(i) + 1], bits, group) for i in uniq]
        return torch.cat(rows)[torch.as_tensor(inv)]

    def __call__(self, ids):
        emb = self._rows("embed_tokens", ids) * (self.tc.hidden_size ** 0.5)
        ple = self._rows("embed_tokens_per_layer", ids) * (self.tc.hidden_size_per_layer_input ** 0.5)
        return emb[None], ple.view(1, len(ids), self.tc.num_hidden_layers, self.tc.hidden_size_per_layer_input)


def masks_for(enc, window, L=None, dtype=torch.float32):
    """D1A's packed masks (d1a.backends.torch.branch_masks) for one record, padded to L: additive [1, 1, L, L]."""
    seg, pos = enc["seg"], enc["pos"]
    n = len(seg)
    L = L or n
    s = torch.full((L,), -1)
    s[:n] = torch.tensor(seg)
    p = torch.zeros(L, dtype=torch.long)
    p[:n] = torch.tensor(pos)
    causal = torch.tril(torch.ones(L, L, dtype=torch.bool))
    same = (s[None, :] == s[:, None]) | (s[None, :] == 0)
    allow = (causal & same & (s != -1)[None, :]) | torch.eye(L, dtype=torch.bool)
    neg = torch.finfo(dtype).min
    full = torch.zeros(L, L, dtype=dtype).masked_fill(~allow, neg)
    far = (p[:, None] - p[None, :] >= window) & ~torch.eye(L, dtype=torch.bool)
    slide = full.masked_fill(far, neg)
    return full[None, None], slide[None, None], p[None]


class Core(nn.Module):
    """Everything that computes: per-layer projection, decoder layers, final norm, pointer head.
    forward(emb [1,L,d], ple_tok [1,L,layers,256], pos [1,L] int, mask_full, mask_slide [1,1,L,L],
            decide [Q] int, opts [Q,K] int, opt_valid [Q,K] float (1 used, 0 padding)) -> logits [Q,K]
    (raw, temperature not applied; padded options get a large negative)."""

    def __init__(self, text, head_q, head_k, silu_tanh=True):
        super().__init__()
        self.text, self.hq, self.hk = text, head_q, head_k
        self.types = text.config.layer_types
        self.scale = 1 / math.sqrt(head_q.out_features)

    def forward(self, emb, ple_tok, pos, mask_full, mask_slide, decide, opts, opt_valid):
        t = self.text
        ple = t.project_per_layer_inputs(emb, ple_tok)
        rope = {lt: t.rotary_emb(emb, pos, lt) for lt in t.unique_layer_types}
        masks = {"full_attention": mask_full, "sliding_attention": mask_slide}
        from collections import UserDict
        kv = UserDict()
        h = emb
        for i, layer in enumerate(t.layers):
            h = layer(h, ple[:, :, i, :], shared_kv_states=kv, position_embeddings=rope[self.types[i]],
                      attention_mask=masks[self.types[i]], position_ids=pos)
        h = t.norm(h)[0]                                    # [L, d]
        q = self.hq(h[decide])                              # [Q, dp]
        k = self.hk(h[opts.reshape(-1)]).view(*opts.shape, -1)   # [Q, K, dp]
        z = (k * q[:, None, :]).sum(-1) * self.scale
        return z + (opt_valid - 1.0) * 1e4