"""vela3_coreml.py: vllm-sr/vela3-fast on Core ML (macOS 15+), no PyTorch weights needed.

The encoder runs as a fixed-length fp16 Core ML function (the smallest bucket that fits); the small heads run in fp32 on the host.
Input rendering, span decoding and thresholds are the release's own code (upstream/, Apache-2.0, unmodified).

    from vela3_coreml import Vela3CoreML
    m = Vela3CoreML(".")                                   # this folder; compute_units="CPU_AND_GPU" by default
    m.decide("My card was charged twice.", {"intent": {"type": "choice", "instructions": "What does the customer want?",
                                                       "criteria": {"refund": "Money back", "cancel": "Cancel the order"}},
                                            "urgent": {"type": "noul", "instructions": "The customer needs a fast response."}})
    m.pii("Call Dana Whitfield at 312-555-0110.")         # -> [{"label", "start", "end", "text", "score"}]
    m.unsupported(context="...", answer="...")              # words of the answer the context does not support

Scope: decisions (choice / noul) on the decision backbone, and the two label sets of the span_sets adapter (17 PII types,
unsupported claims). Inputs are limited to 1,024 tokens per encoder pass: decision inputs longer than that raise Unsupported;
span inputs longer than that are windowed at 1,024 tokens (the release windows at 8,192), so long-document spans can differ."""
import json, sys, time
from pathlib import Path

import numpy as np
import torch
from torch import nn

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "upstream"))
import vela3_schema as VS  # noqa: E402
from vela3_encoder import DecisionEncoder, Vela3Encoder, Unsupported  # noqa: E402

PII = None  # filled from vela3_config.json


class _Tok:
    """The slice of the HF tokenizer API the release's DecisionEncoder uses (ids identical to AutoTokenizer on this tokenizer.json)."""

    def __init__(self, path, cfg):
        from tokenizers import Tokenizer
        self.t = Tokenizer.from_file(str(path / "tokenizer.json")); self.t.no_padding(); self.t.no_truncation()
        self.bos_token_id, self.eos_token_id, self.pad_token_id = cfg["bos_token_id"], cfg["eos_token_id"], cfg["pad_token_id"]

    def __call__(self, s, add_special_tokens=False):
        return {"input_ids": self.t.encode(s, add_special_tokens=add_special_tokens).ids}


class _Encoder:
    """One multifunction Core ML package: L<n> functions, input_ids / attention_mask int32 [1, n] -> hidden fp16 [1, n, 1024]."""

    def __init__(self, package, lengths, compute_units):
        import coremltools as ct
        cu = getattr(ct.ComputeUnit, compute_units)
        self.lengths = tuple(lengths)
        self.fn = {L: ct.models.MLModel(str(package), function_name=f"L{L}", compute_units=cu) for L in self.lengths}
        self.last_ms = []

    def __call__(self, ids):
        n = len(ids)
        L = next((L for L in self.lengths if n <= L), None)
        if L is None:
            raise Unsupported(f"input of {n} tokens is longer than the largest Core ML function (L{self.lengths[-1]})")
        x = np.zeros((1, L), np.int32); x[0, :n] = ids; m = np.zeros((1, L), np.int32); m[0, :n] = 1
        t = time.perf_counter(); h = self.fn[L].predict({"input_ids": x, "attention_mask": m})["hidden"]
        self.last_ms.append((L, 1e3 * (time.perf_counter() - t)))
        return torch.from_numpy(h[0, :n].astype(np.float32))


class _SpanRun:
    """Drop-in for vela3_schema._GatedBackend: __call__(x, m, qi, oi, ui, ei) -> (choice logits, span logits)."""

    def __init__(self, net, enc):
        self.net, self.enc = net, enc

    def __call__(self, x, m, qi, oi, ui, ei):
        B, W = x.shape
        H = torch.zeros(B, W, 1024)
        for b in range(B):
            n = int(m[b].sum()); H[b, :n] = self.enc(x[b, :n].tolist())
        with torch.no_grad():
            lg, sp = self.net.heads(H, *(torch.from_numpy(a) for a in (qi, oi, ui, ei)))
        return lg.float().numpy(), sp.float().numpy()


class Vela3CoreML:
    def __init__(self, path=HERE, compute_units="CPU_AND_GPU"):
        global PII
        path = Path(path)
        cfg = json.loads((path / "config.json").read_text()); self.cfg = json.loads((path / "vela3_config.json").read_text())
        PII = self.cfg["span"]["router_label_sets"]["pii"]
        from safetensors.torch import load_file
        # decisions: the release DecisionEncoder without its torch encoder (heads + rendering only), encoder on Core ML
        d = DecisionEncoder.__new__(DecisionEncoder); nn.Module.__init__(d)
        h, p = cfg["hidden_size"], cfg["decision"]["proj_dim"]
        d.norm = nn.LayerNorm(h); d.w_o, d.w_q, d.w_p = nn.Linear(h, p), nn.Linear(h, p), nn.Linear(h, p)
        d.log_tau = nn.Parameter(torch.zeros(())); d.cls_mlp = nn.Sequential(nn.Linear(h, 2 * h), nn.ReLU(), nn.Dropout(0.1), nn.Linear(2 * h, 1))
        d.load_state_dict(load_file(str(path / "heads_decision.safetensors")), strict=True); d.eval()
        d.tok = _Tok(path, cfg); d.mk = {t: int(i) for t, i in cfg["decision"]["marker_ids"].items()}
        d.device, d.max_len, d.budget = "cpu", 1024, 65536
        self.dec_enc = _Encoder(path / "Vela3Decision.mlpackage", (128, 256, 512, 1024), compute_units)
        d.forward_recs = self._decision_forward
        self.decision = d
        # spans: the release SchemaNet heads of the span_sets adapter + its SchemaEngine, encoder on Core ML
        net = VS.SchemaNet(nn.Identity(), h, p, span=True)
        net.load_state_dict(load_file(str(path / "heads_span_sets.safetensors")), strict=True); net.eval()
        self.span_enc = _Encoder(path / "Vela3SpanSets.mlpackage", (512, 1024), compute_units)
        eng = VS.SchemaEngine(path, net, "cpu", max_len=1024)
        eng.run = _SpanRun(net, self.span_enc)
        v = Vela3Encoder.__new__(Vela3Encoder)
        v.cfg, v._VS, v.span, v.span_ev, v.span_set, v.task, v.entity_route = self.cfg, VS, eng, eng, eng, {}, "ner"
        self.v = v

    @torch.no_grad()
    def _decision_forward(self, recs):
        out = []
        for r in recs:
            H = self.dec_enc(r["ids"])[None]
            z = self.decision._logits(H, [r])[0].double().numpy(); e = np.exp(z - z.max()); out.append(e / e.sum())
        return out

    def decide(self, state, questions):
        """questions: {key: {type: choice | noul, instructions, criteria}} -> {key: answer} (the release's predict())."""
        return self.decision.predict(state, questions)[0]

    def pii(self, text, threshold=None):
        q = {"type": "span", "criteria": PII}
        if threshold is not None:
            q["threshold"] = threshold
        return self.v.span_answer(text, q)[0]["spans"]

    def unsupported(self, context, answer, threshold=None):
        q = {"type": "span", "criteria": ["unsupported"]}
        if threshold is not None:
            q["threshold"] = threshold
        return self.v.span_answer({"context": context, "answer": answer}, q)[0]["spans"]


if __name__ == "__main__":
    m = Vela3CoreML()
    print(m.decide("My card was charged twice and I need this fixed today.",
                   {"intent": {"type": "choice", "instructions": "What does the customer want?",
                               "criteria": {"refund": "Money back for the duplicate charge", "cancel": "Cancel the order"}},
                    "urgent": {"type": "noul", "instructions": "The customer needs a fast response."}}))
    print(m.pii("Hi, I'm Dana Whitfield (dana@example.com, 312-555-0110), 4410 Elm Ridge Drive, Dayton, OH 45429."))
    print(m.unsupported(context="Adults may take up to 4 g of paracetamol in 24 hours.",
                        answer="Adults may take up to 6 g of paracetamol in 24 hours."))