"""d3 text requests on Core ML (Python host), for the checkpoints whose text graph takes embeddings (d3-mini).

Upstream's own request preparation runs on a weightless (meta-device) d3, so prompts and token ids are exactly
upstream's; Core ML runs the text graph per question (host embedding gather, text positions on all three M-RoPE axes);
FP32 readout on the host.

    m = D3CoreMLText("<d3 checkpoint>", "<coreml dir>", text_buckets=[256, 640])
    m.system_one(state=..., questions={...})"""
import sys
import time
from pathlib import Path

import coremltools as ct
import numpy as np
import torch

HERE = Path(__file__).parent


class D3CoreMLText:
    def __init__(self, root, coreml_dir, text_buckets=None, compute_units=ct.ComputeUnit.CPU_AND_GPU):
        root, cdir = Path(root), Path(coreml_dir)
        sys.path.insert(0, str(root))
        sys.path.insert(0, str(HERE.parent))
        import d3_runtime as up
        from d3_graph import load_text_config
        from qwen35_export import rope_cos_sin
        from safetensors.torch import load_file

        self.up, self.rope_cos_sin = up, rope_cos_sin
        self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none")  # no weights: prompts only
        self.temperature = self.prep.temperature
        self.tcfg = load_text_config(str(root))
        self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy()
        self.embed = np.memmap(cdir / "embeddings.f16", dtype=np.float16, mode="r").reshape(-1, self.tcfg.hidden_size)
        self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units)
                     for p in sorted(cdir.glob("d3_E*.mlpackage")) if not text_buckets or int(p.stem[4:]) in text_buckets}
        self.pad = self.prep.pad_id
        self.model_name = root.name + "-coreml"
        self.timing = []

    def warm(self):
        for L, model in self.text.items():
            for _ in range(2):
                self._final(list(range(10)), L)

    def _final(self, ids, L):
        n = len(ids)
        x = np.zeros((1, L, self.tcfg.hidden_size), dtype=np.float16)
        x[0, :n] = self.embed[np.asarray(ids)]
        x[0, n:] = self.embed[self.pad]
        pos = torch.arange(L).clamp(max=n - 1)[None].expand(3, -1)
        cos, sin = self.rope_cos_sin(self.tcfg, pos)
        valid = np.zeros(L, dtype=np.float16); valid[:n] = 1
        last = np.zeros((1, L), dtype=np.float16); last[0, n - 1] = 1
        return self.text[L].predict({"hidden": x, "cos": cos.numpy().astype(np.float16),
                                     "sin": sin.numpy().astype(np.float16), "key_valid": valid,
                                     "last_onehot": last})["final"].reshape(-1)

    def system_one(self, *, state, questions):
        prepared = self.prep.prepare(state, questions)
        answers, used, self.timing = {}, 0, []
        for key in prepared.keys:
            if key in prepared.errors:
                answers[key] = prepared.errors[key]
                continue
            ids, question = prepared.sequences[key], prepared.questions[key]
            L = next((b for b in sorted(self.text) if b >= len(ids)), None)
            if L is None:
                answers[key] = {"type": question.kind, "error": "max_length_exceeded",
                                "message": f"{len(ids)} tokens exceed the largest Core ML bucket {max(self.text)}"}
                continue
            t = time.perf_counter()
            h = self._final(ids, L)
            self.timing.append((len(ids), L, round(time.perf_counter() - t, 3)))
            logits = self.readout[:len(question.keys)].astype(np.float64) @ h.astype(np.float64) / self.temperature
            p = np.exp(logits - logits.max())
            answers[key] = self.up.product_answer(question, (p / p.sum()).tolist())
            used += len(ids)
        return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}