"""Run a chunked D1A Core ML trunk end to end: host embeddings -> chunk 0 .. chunk N (h, per-layer inputs and the shared
keys/values passed along) -> hidden states at the picks -> pointer head on the host.
Usage as a library: Pipeline(export_dir, chunk_dir, L, units).probs(enc) -> [per-question probabilities]."""
import glob
import os
import re
import sys
import time

import coremltools as ct
import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))   # d1a_coreml.py sits next to this file
from d1a_coreml import Export, HostEmbed, masks_for  # noqa: E402

P_PICKS = 128


class Pipeline:
    def __init__(self, export_dir, chunk_dir, L, units="CPU_AND_NE"):
        self.ex = Export(export_dir)
        self.host = HostEmbed(self.ex)
        self.hq, self.hk = self.ex.head()
        self.T = self.ex.d1a["temperature"]
        self.L = L
        cu = getattr(ct.ComputeUnit, units)
        self.chunks = []
        if chunk_dir.endswith(".mlpackage"):   # one multifunction package: functions c00_07, c07_15, ...
            names = ["c00_07", "c07_15", "c15_22", "c22_29", "c29_35"]
            for name in names:
                a, b = map(int, name[1:].split("_"))
                m = ct.models.MLModel(chunk_dir, compute_units=cu, function_name=name)
                fn = next(f for f in m.get_spec().description.functions if f.name == name)
                self.chunks.append((a, b, m, [i.name for i in fn.input], [o.name for o in fn.output]))
            return
        files = sorted(glob.glob(f"{chunk_dir}/L{L}_*.mlpackage"))
        for f in files:
            a, b = map(int, re.search(r"_(\d+)-(\d+)\.mlpackage$", f).groups())
            m = ct.models.MLModel(f, compute_units=cu)
            spec = m.get_spec()
            self.chunks.append((a, b, m, [i.name for i in spec.description.input], [o.name for o in spec.description.output]))

    def hidden(self, enc):
        L, n = self.L, len(enc["ids"])
        # right-pad with the pad token's own embeddings (as d1a's torch backend pads ids), never zero rows: an all-zero row
        # makes every RMSNorm compute 0 / sqrt(eps), and eps (1e-6) is subnormal in fp16, which the Neural Engine flushes
        ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
        emb, ple = self.host(ids)
        e = emb.numpy().astype(np.float16)
        p = ple.numpy().astype(np.float16)
        mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
        picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
        pk = np.zeros(P_PICKS, np.int32); pk[: len(picks)] = picks
        common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
                  "mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
        state = {"emb": e, "ple_tok": p, "picks": pk}
        for a, b, m, ins, outs in self.chunks:
            feed = {k: np.ascontiguousarray(common[k] if k in common else state[k]) for k in ins}
            feed = {k: (v.astype(np.float16) if v.dtype == np.float32 else v) for k, v in feed.items()}   # chunk outputs come back fp32
            out = {k: np.array(v, copy=True) for k, v in m.predict(feed).items()}   # Core ML may reuse its output buffers on the next call
            if "hidden" in out:
                return torch.tensor(out["hidden"][: len(picks)], dtype=torch.float32)
            state["h_in"] = out["h"]
            if "ple_out" in out:
                state["ple"] = out["ple_out"]
            for k, v in out.items():
                if k.endswith("_out") and k != "ple_out":
                    state[k[: -len("_out")]] = v
        raise RuntimeError("no chunk returned hidden")

    def probs(self, enc):
        h = self.hidden(enc)
        Q = len(enc["decide_idx"])
        out, j = [], Q
        with torch.no_grad():
            for qi, oo in enumerate(enc["opt_idx"]):
                z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
                out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
        return out