"""Decision Index engine over the Core ML window models:
`--engine dmodel_mac.engine:CoreMLDecisionEngine --option model_dir=... --option name=...`.

Per question: render windows (dmodel_mac.render), run each window through the Core ML bucket
that fits it, pool marker logits per option (log-mean-exp over occurrences), divide by the
calibrated temperature and softmax. Nothing is truncated and no option is dropped.
"""

from __future__ import annotations

import json
import math
import time
from pathlib import Path
from typing import Any

import numpy as np

from dmodel_mac.render import MAX_OPTIONS, Renderer, WindowConfig, option_list
from dmodel_mac.render import Unsupported as RenderUnsupported

try:
    from decision_index.engines import Engine, Unsupported
except ImportError:  # usable without the kit installed
    class Engine:  # type: ignore[no-redef]
        def __init__(self, **options: Any) -> None:
            self.options = options

        def warmup(self):
            pass

    class Unsupported(ValueError):  # type: ignore[no-redef]
        pass

UNITS = {"all": "ALL", "ane": "CPU_AND_NE", "gpu": "CPU_AND_GPU", "cpu": "CPU_ONLY"}


def to_answer(question: dict[str, Any], probs: np.ndarray) -> dict[str, Any]:
    keys, descriptions = option_list(question)
    values = [float(p) for p in probs]
    if len(values) != len(keys) or any(not math.isfinite(v) or v < 0 for v in values):
        raise RuntimeError("invalid distribution")
    total = sum(values)
    values = [v / total for v in values]
    if question["type"] == "noul":
        return {"type": "noul", "noul": values[1]}
    best = max(range(len(values)), key=values.__getitem__)
    dist = dict(zip(keys, values))
    if question["type"] == "choice":
        return {"type": "choice", "choice": keys[best], "probabilities": dist}
    return {"type": "score", "probabilities": dist, "legend": dict(zip(keys, descriptions)),
            "score": sum(i * p for i, p in enumerate(values)), "level": keys[best]}


class WindowRunner:
    """Loads one Core ML model per bucket and scores windows."""

    def __init__(self, model_dir: str, name: str, slots: int, buckets: tuple[int, ...], units: dict[int, str]):
        import coremltools as ct
        self.slots, self.models = slots, {}
        for b in buckets:
            unit = getattr(ct.ComputeUnit, UNITS[units[b]])
            self.models[b] = ct.models.MLModel(str(Path(model_dir) / f"{name}_L{b}_K{slots}.mlpackage"), compute_units=unit)
        self.buckets = buckets

    def window_logits(self, window) -> np.ndarray:
        """Logits of the window's markers, in marker order (chunks of `slots` markers per call)."""
        b = next(x for x in self.buckets if window.length <= x)
        ids = np.zeros((1, b), dtype=np.int32)
        ids[0, :] = 50283
        ids[0, :window.length] = window.ids
        mask = np.zeros((1, b), dtype=np.int32)
        mask[0, :window.length] = 1
        out = []
        for s in range(0, len(window.positions), self.slots):
            mm = np.zeros((1, self.slots, b), dtype=np.float16)
            chunk = window.positions[s:s + self.slots]
            mm[0, np.arange(len(chunk)), chunk] = 1
            res = self.models[b].predict({"input_ids": ids, "attention_mask": mask, "marker_map": mm})
            out.append(np.asarray(res["logits"], dtype=np.float32).reshape(-1)[:len(chunk)])
        return np.concatenate(out)


def pool_lme(windows, logits_per_window, n_opt: int) -> np.ndarray:
    occ: list[list[float]] = [[] for _ in range(n_opt)]
    for w, z in zip(windows, logits_per_window):
        for o, v in zip(w.options, z):
            occ[o].append(float(v))
    pooled = np.empty(n_opt, dtype=np.float64)
    for o, vals in enumerate(occ):
        a = np.asarray(vals, dtype=np.float64)
        m = a.max()
        pooled[o] = m + math.log(np.exp(a - m).mean())
    return pooled


class CoreMLDecisionEngine(Engine):
    name = "dmodel-mac"
    latency = ("in-process wall time of the whole request: rendering, tokenization, every Core ML window "
               "prediction, pooling and softmax; excludes model loading")

    def __init__(self, **options: Any) -> None:
        super().__init__(**options)
        cfg = json.loads(Path(options["config"]).read_text()) if "config" in options else {}
        cfg.update({k: v for k, v in options.items() if k != "config"})
        self.temperature = float(cfg.get("temperature", 1.0))
        self.slots = int(cfg.get("slots", 64))
        buckets = tuple(int(x) for x in str(cfg.get("buckets", "128,256,512")).split(","))
        default_units = {b: ("ane" if b <= 128 else "all") for b in buckets}
        for b in buckets:
            if f"units_{b}" in cfg:
                default_units[b] = cfg[f"units_{b}"]
        self.renderer = Renderer(cfg.get("tokenizer", "bases/modernbert-base/tokenizer.json"),
                                 WindowConfig(max_len=buckets[-1], buckets=buckets))
        self.runner = WindowRunner(cfg["model_dir"], cfg["name"], self.slots, buckets, default_units)
        self.provenance = {"model_dir": cfg["model_dir"], "name": cfg["name"], "temperature": self.temperature,
                           "buckets": list(buckets), "units": {str(k): v for k, v in default_units.items()},
                           "slots": self.slots, "window_config": self.renderer.cfg.__dict__}

    def logits(self, state: Any, question: dict[str, Any]) -> np.ndarray:
        keys, _ = option_list(question)
        if len(keys) > MAX_OPTIONS:
            raise Unsupported(f"{len(keys)} options exceeds the declared limit of {MAX_OPTIONS}")
        try:
            windows, n = self.renderer.windows(state, question)
        except RenderUnsupported as e:
            raise Unsupported(str(e)) from e
        return pool_lme(windows, [self.runner.window_logits(w) for w in windows], n)

    def __call__(self, state: Any, questions: dict[str, Any]):
        t0 = time.perf_counter()
        answers, windows = {}, 0
        for key, q in questions.items():
            z = self.logits(state, q) / self.temperature
            p = np.exp(z - z.max())
            p /= p.sum()
            answers[key] = to_answer(q, p)
        return {"model": self.name, "answers": answers}, {"wall_ms": (time.perf_counter() - t0) * 1000, "windows": windows}