"""Decision 2.0 System One on Core ML: every question of a request in one packed call.

Supports the Qwen3 package (Kai, `L*_N*` functions) and the Qwen3.5 hybrid package (Eos, `S*_C*_N*` functions).

Needs only numpy, tokenizers and coremltools. Prompt rendering, option endpoints, Score
offsets and answer normalization follow the upstream runtime (vllm-sr Decision 2.0,
Apache-2.0: decision2/_vendor/dev2model/{decision_model,infer,score_bias}.py).

    from decision2_coreml import Decision2CoreML
    model = Decision2CoreML(".")  # this repo's folder
    model.system_one(state="...", questions={...})
"""

from __future__ import annotations

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

import coremltools as ct
import numpy as np
from tokenizers import Tokenizer

MAX_OPTIONS = 255
NEG = -1e4


def canonical(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False)


def _payload(value: Any) -> str:
    return value if isinstance(value, str) else canonical(value)


def question_to_row(state: Any, question: dict[str, Any]) -> dict[str, Any]:
    kind = question.get("type")
    if kind not in ("choice", "noul", "score"):
        raise ValueError("unsupported question type")
    if not question.get("instructions"):
        raise ValueError("missing question instructions")
    criteria = question.get("criteria")
    if kind == "score":
        if not isinstance(criteria, list) or not 2 <= len(criteria) <= 10:
            raise ValueError("score criteria must be an ordered list of 2..10 levels")
        options = [{"key": str(i), "description": d} for i, d in enumerate(criteria)]
    else:
        if kind == "noul":
            criteria = criteria or {}
            if not isinstance(criteria, dict) or set(criteria) - {"false", "true"}:
                raise ValueError("noul requires only false and true criteria")
            if len(criteria) < 2:
                criteria = {"false": criteria.get("false", "No"), "true": criteria.get("true", "Yes")}
        if not isinstance(criteria, dict) or not 2 <= len(criteria) <= MAX_OPTIONS:
            raise ValueError("choice criteria must be an object with 2..255 options")
        options = [{"key": k, "description": d} for k, d in criteria.items()]
    return {"state": state, "instructions": question["instructions"], "options": options, "task_type": kind}


def encode(row: dict[str, Any], tokenizer: Tokenizer) -> dict[str, Any]:
    prefix = (
        f"Context:\n{_payload(row['state'])}\n\n"
        f"Task type: {row['task_type']}\nQuestion:\n{_payload(row['instructions'])}\nOptions:"
    )
    ids = tokenizer.encode(prefix, add_special_tokens=False).ids
    endpoints = []
    for option in row["options"]:
        text = "\n<option>\n" + canonical({"key": option["key"], "description": option["description"]}) + "\n</option>"
        ids.extend(tokenizer.encode(text, add_special_tokens=False).ids)
        endpoints.append(len(ids) - 1)
    suffix = "\n\nSelect the single option best supported by the context and instructions.\nDecision:"
    ids.extend(tokenizer.encode(suffix, add_special_tokens=False).ids)
    return {"ids": ids, "candidate_positions": endpoints, "query_position": len(ids) - 1,
            "keys": [o["key"] for o in row["options"]]}


def _prefix(encoded: list[dict[str, Any]]) -> int:
    seqs = [e["ids"] for e in encoded]
    limit = min(min(e["candidate_positions"][0] for e in encoded), min(len(s) for s in seqs) - 1)
    for i in range(limit):
        if any(s[i] != seqs[0][i] for s in seqs):
            return i
    return limit


def packed_size(encoded: list[dict[str, Any]]) -> tuple[int, int]:
    P = _prefix(encoded)
    return P + sum(len(e["ids"]) - P for e in encoded), sum(len(e["keys"]) for e in encoded)


def pack(encoded: list[dict[str, Any]], L: int, N: int, pad_id: int) -> tuple[dict[str, np.ndarray], list[int]]:
    """Shared prefix once, then each question's suffix; a suffix sees the prefix and itself (causal)."""
    P = _prefix(encoded)
    ids, pos, seg = list(encoded[0]["ids"][:P]), list(range(P)), [-1] * P
    cand, qry, owner = [], [], []
    for j, e in enumerate(encoded):
        shift = len(ids) - P
        ids += e["ids"][P:]
        pos += range(P, len(e["ids"]))
        seg += [j] * (len(e["ids"]) - P)
        for c in e["candidate_positions"]:
            cand.append(c if c < P else c + shift)
            qry.append(e["query_position"] + shift)
            owner.append(j)
    T = len(ids)
    if T > L or len(cand) > N:
        raise ValueError(f"{T} tokens / {len(cand)} options exceed L{L}_N{N}")
    s = np.array(seg + [-2] * (L - T))
    i = np.arange(L)
    allow = (i[None, :] <= i[:, None]) & ((s[None, :] == s[:, None]) | (s[None, :] == -1)) & (s[None, :] != -2)
    allow[i, i] = True
    return {
        "input_ids": np.array([ids + [pad_id] * (L - T)], dtype=np.int32),
        "position_ids": np.array([pos + [0] * (L - T)], dtype=np.int32),
        "mask": np.where(allow, 0.0, NEG).astype(np.float16)[None, None],
        "cand_idx": np.array(cand + [0] * (N - len(cand)), dtype=np.int32),
        "query_idx": np.array(qry + [0] * (N - len(cand)), dtype=np.int32),
    }, owner


def pack_hybrid(encoded: list[dict[str, Any]], S: int, P: int, N: int, pad_id: int, rope: dict[str, Any],
                lags: int = 3, chunk: int = 64) -> tuple[dict[str, np.ndarray], list[int]]:
    """Qwen3.5 (Gated DeltaNet hybrid) graph inputs: the shared prefix right-padded to S, the questions' suffixes
    packed in P, plus the masks that make every suffix restart from the prefix's state (see README)."""
    pre = _prefix(encoded)
    total = sum(len(e["ids"]) - pre for e in encoded)
    options = sum(len(e["keys"]) for e in encoded)
    if not 1 <= pre <= S or total > P or options > N:
        raise ValueError(f"prefix {pre} / {total} question tokens / {options} options exceed S{S}_C{P}_N{N}")
    ids = np.full(S + P, pad_id, dtype=np.int32)
    pos = np.zeros(S + P, dtype=np.float64)
    ids[:pre] = encoded[0]["ids"][:pre]
    pos[:pre] = np.arange(pre)
    pos[pre:] = pre
    valid = np.zeros(S, dtype=np.float16)
    valid[:pre] = 1
    tail = np.zeros((lags, S), dtype=np.float16)
    for i in range(lags):
        if pre - lags + i >= 0:
            tail[i, pre - lags + i] = 1
    segment = np.eye(P, dtype=np.float16)
    keep = np.zeros((lags, P), dtype=np.float16)
    lag_tail = np.zeros((lags, P, lags), dtype=np.float16)
    cand, qry, owner = [], [], []
    start = 0
    for j, e in enumerate(encoded):
        n = len(e["ids"]) - pre
        ids[S + start : S + start + n] = e["ids"][pre:]
        pos[S + start : S + start + n] = np.arange(pre, pre + n)
        segment[start : start + n, start : start + n] = np.tril(np.ones((n, n)))
        for p in range(n):
            for s in range(1, lags + 1):
                if p >= s:
                    keep[s - 1, start + p] = 1
                else:
                    lag_tail[s - 1, start + p, lags + p - s] = 1
        for c in e["candidate_positions"]:
            cand.append(start + c - pre)
            qry.append(start + e["query_position"] - pre)
            owner.append(j)
        start += n
    dim = rope["rotary_dim"]
    inv_freq = 1.0 / (rope["rope_theta"] ** (np.arange(0, dim, 2, dtype=np.float64) / dim))
    freqs = pos[:, None] * inv_freq[None, :]  # text-only M-RoPE: all three axes share the position
    emb = np.concatenate([freqs, freqs], axis=-1)
    M = P // chunk
    starts = np.argmax(segment > 0, axis=1)
    pad = N - len(cand)
    return {
        "input_ids": ids[None], "cos": np.cos(emb).astype(np.float16), "sin": np.sin(emb).astype(np.float16),
        "valid": valid, "tail_onehot": tail, "segment": segment, "lag_keep": keep, "lag_tail": lag_tail,
        "seg_chunks": np.stack([segment[i * chunk : (i + 1) * chunk, i * chunk : (i + 1) * chunk] for i in range(M)]),
        "cont": (starts < (np.arange(P) // chunk) * chunk).astype(np.float16),
        "last_seg": np.stack([segment[(i + 1) * chunk - 1, i * chunk : (i + 1) * chunk] for i in range(M)]),
        "cand_idx": np.array(cand + [0] * pad, dtype=np.int32),
        "query_idx": np.array(qry + [0] * pad, dtype=np.int32),
    }, owner


def answer(kind: str, keys: list[str], logits: list[float], descriptions: list[Any]) -> dict[str, Any]:
    top = max(logits)
    exps = [math.exp(v - top) for v in logits]
    total = sum(exps)
    probs = [v / total for v in exps]
    pmap = dict(zip(keys, probs))
    if kind == "noul":
        return {"type": "noul", "noul": pmap["true"]}
    entropy = -sum(p * math.log(p) for p in probs if p > 0)
    confidence = max(0.0, min(1.0, 1.0 - entropy / math.log(len(keys))))
    if kind == "score":
        return {"type": "score", "score": sum(int(k) * pmap[k] for k in keys), "probabilities": pmap,
                "confidence": confidence,
                "legend": {k: d if isinstance(d, str) else canonical(d) for k, d in zip(keys, descriptions)}}
    return {"type": "choice", "choice": keys[probs.index(max(probs))], "probabilities": pmap, "confidence": confidence}


def _dims(name: str) -> dict[str, int]:
    return {part[0]: int(part[1:]) for part in name.split("_")}


class Decision2CoreML:
    def __init__(self, root: str | Path = ".", package: str | None = None,
                 compute_units: ct.ComputeUnit = ct.ComputeUnit.CPU_AND_GPU):
        root = Path(root)
        self.config = json.loads((root / "coreml_config.json").read_text())
        self.tokenizer = Tokenizer.from_file(str(root / "tokenizer.json"))
        self.pad_id = self.config["pad_token_id"]
        self.hybrid = self.config.get("backbone", "qwen3") == "qwen3_5"
        bias = root / "score_bias.json"
        offsets = json.loads(bias.read_text())["offsets"] if bias.exists() else {}
        self.score_bias = {int(k): v for k, v in offsets.items()}
        path = str(root / (package or self.config["package"]))
        self.cost_ms = self.config["functions"]
        self.buckets = [(name, _dims(name), ct.models.MLModel(path, function_name=name, compute_units=compute_units))
                        for name in self.config["functions"]]  # smallest first

    def _fits(self, d: dict[str, int], encoded: list[dict[str, Any]]) -> bool:
        options = sum(len(e["keys"]) for e in encoded)
        if self.hybrid:
            pre = _prefix(encoded)
            return pre <= d["S"] and sum(len(e["ids"]) - pre for e in encoded) <= d["C"] and options <= d["N"]
        return packed_size(encoded)[0] <= d["L"] and options <= d["N"]

    def _inputs(self, d: dict[str, int], encoded: list[dict[str, Any]]):
        if self.hybrid:
            return pack_hybrid(encoded, d["S"], d["C"], d["N"], self.pad_id, self.config["rope"])
        return pack(encoded, d["L"], d["N"], self.pad_id)

    def _calls(self, encoded: list[dict[str, Any]]) -> list[tuple[Any, dict[str, int], list[int]]]:
        """Fewest estimated milliseconds: one call in the smallest bucket that fits, or greedy chunks."""
        best = None
        for name, d, m in self.buckets:
            groups, group = [], []
            for j in range(len(encoded)):
                if not self._fits(d, [encoded[i] for i in group + [j]]):
                    if not group or not self._fits(d, [encoded[j]]):
                        break  # this question alone does not fit this bucket
                    groups.append(group)
                    group = []
                group.append(j)
            else:
                groups.append(group)
                cost = len(groups) * self.cost_ms[name]
                if best is None or cost < best[0]:
                    best = (cost, [(m, d, g) for g in groups])
        if best is None:
            raise ValueError(f"a question exceeds the largest function {self.buckets[-1][0]}")
        return best[1]

    def logits(self, encoded: list[dict[str, Any]]) -> list[list[float]]:
        per: list[list[float]] = [[] for _ in encoded]
        for m, d, group in self._calls(encoded):
            x, owner = self._inputs(d, [encoded[i] for i in group])
            out = m.predict(x)["logits"]
            for k, o in enumerate(owner):
                per[group[o]].append(float(out[k]))
        return per
    def system_one(self, state: Any, questions: dict[str, dict[str, Any]]) -> dict[str, Any]:
        rows, encoded, answers = {}, {}, {}
        for qid, q in questions.items():
            try:
                rows[qid] = question_to_row(state, q)
                encoded[qid] = encode(rows[qid], self.tokenizer)
            except ValueError:
                answers[qid] = {"type": q.get("type") if isinstance(q, dict) else None, "error": "invalid_question"}
        qids = list(encoded)
        if qids:
            for qid, values in zip(qids, self.logits([encoded[q] for q in qids])):
                row = rows[qid]
                offsets = self.score_bias.get(len(values)) if row["task_type"] == "score" else None
                if offsets:
                    values = [v + b for v, b in zip(values, offsets)]
                answers[qid] = answer(row["task_type"], encoded[qid]["keys"], values,
                                      [o["description"] for o in row["options"]])
        return {"model": self.config["model_name"], "answers": {q: answers[q] for q in questions},
                "usage": {"input_tokens": sum(len(e["ids"]) for e in encoded.values()), "output_tokens": 0}}