"""One process per Core ML chunk (workaround: several long-context chunks in one process crash the Neural Engine runtime
on the second request on this Mac, while each chunk alone, or all on the GPU, runs fine). The parent runs the host side
(embeddings, masks, pointer head) and passes tensors along the chain through pipes.
  WorkerPipeline(export_dir, chunk_dir, L, units).probs(enc)"""
import glob
import os
import multiprocessing as mp
import re
import shutil
import sys

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


BUCKETS = tuple(int(x) for x in os.environ.get("D1A_COREML_BUCKETS", "512").split(","))   # each length is its own ANE program; fewer keeps the ANE within its memory


def compiled_chunks(chunk_dir):
    """-> [(first layer, end layer, .mlmodelc, its .mlpackage)]. A released bundle's coreml/L512_<a>-<b>.mlpackage are
    compiled once, each as a one-function package (L512), into coreml/.compiled/; a build tree's compiled/*.mlmodelc
    (multifunction, one function per length) are used as they are."""
    import coremltools as ct
    out = []
    for f in sorted(glob.glob(f"{chunk_dir}/*.mlmodelc")):
        a, b = map(int, re.search(r"(\d+)-(\d+)\.mlmodelc$", f).groups())
        out.append((a, b, f, f.replace("compiled/", "buckets/").replace(".mlmodelc", ".mlpackage")))
    if out:
        return out
    os.makedirs(f"{chunk_dir}/.compiled", exist_ok=True)
    for src in sorted(glob.glob(f"{chunk_dir}/L*_*.mlpackage")):
        L, a, b = map(int, re.search(r"L(\d+)_(\d+)-(\d+)\.mlpackage$", src).groups())
        dst = f"{chunk_dir}/.compiled/{a:02d}-{b:02d}.mlmodelc"
        if not os.path.exists(dst):
            d = ct.utils.MultiFunctionDescriptor(); d.add_function(src, src_function_name="main", target_function_name=f"L{L}")
            d.default_function_name = f"L{L}"; tmp = f"{chunk_dir}/.compiled/{a:02d}-{b:02d}.mlpackage"
            ct.utils.save_multifunction(d, tmp); shutil.move(ct.models.utils.compile_model(tmp), dst); shutil.rmtree(tmp)
        out.append((a, b, dst, src))
    return out


def _worker(path, spec_path, units, conn):
    import coremltools as ct
    cu = getattr(ct.ComputeUnit, units)
    # a precompiled .mlmodelc (compiled once by coremltools.models.utils.compile_model): loading each length function
    # from it never recompiles, where MLModel(.mlpackage) compiles the whole package again for every function
    models = {L: ct.models.CompiledMLModel(path, compute_units=cu, function_name=f"L{L}") for L in BUCKETS}
    spec = ct.utils.load_spec(spec_path)
    fns = {f.name: f for f in spec.description.functions}
    fn = fns.get(f"L{BUCKETS[0]}") or spec.description   # a one-function package keeps its inputs in the description
    conn.send(([i.name for i in fn.input], [o.name for o in fn.output]))
    while True:
        msg = conn.recv()
        if msg is None:
            break
        L, feed = msg
        import time as _t
        t0 = _t.perf_counter()
        out = models[L].predict(feed)
        t1 = _t.perf_counter()
        res = {k: np.array(v, dtype=np.float16 if v.dtype == np.float32 else v.dtype, copy=True) for k, v in out.items()}
        res["_predict_ms"] = np.array((t1 - t0) * 1000)
        conn.send(res)


class WorkerPipeline:
    def __init__(self, export_dir, chunk_dir, L=None, units="CPU_AND_NE"):
        from d1a_coreml import Export, HostEmbed
        self.ex = Export(export_dir)
        self.host = HostEmbed(self.ex)
        self.hq, self.hk = self.ex.head()
        self.T, self.L = self.ex.d1a["temperature"], L
        ctx = mp.get_context("spawn")
        self.chunks = []
        for a, b, f, spec_path in compiled_chunks(chunk_dir):
            parent, child = ctx.Pipe()
            p = ctx.Process(target=_worker, args=(f, spec_path, units, child), daemon=True)
            p.start()
            self.chunks.append((a, b, p, parent))
            self.io = getattr(self, "io", []) + [parent.recv()]   # one at a time: concurrent first loads race in Core ML's compiler

    def close(self):
        for _, _, p, conn in self.chunks:
            conn.send(None); p.join(5)

    def hidden(self, enc):
        from d1a_coreml import masks_for
        self.last_predict_ms = []
        n = len(enc["ids"])
        L = next(b for b in BUCKETS if b >= n)   # the smallest compiled length that fits
        ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
        emb, ple = self.host(ids)
        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(128, 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": emb.numpy().astype(np.float16), "ple_tok": ple.numpy().astype(np.float16), "picks": pk}
        for (a, b, p, conn), (ins, outs) in zip(self.chunks, self.io):
            conn.send((L, {k: (common[k] if k in common else state[k]) for k in ins}))
            out = conn.recv()
            self.last_predict_ms = getattr(self, "last_predict_ms", []) + [float(out.pop("_predict_ms"))]
            if "hidden" in out:
                return torch.tensor(out["hidden"][: len(picks)].astype(np.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