"""System One decisions on Apple Silicon: one forward pass, typed answers.

A faithful MLX port of the transformers System One runner that ships with
LiquidAI/d1-3B (`prompt.py`, `runner.py`, `api.py`). The prompt rendering and
the logit readout are identical; the forward pass runs through mlx-vlm.

    python system_one.py --model DJLougen/d1-3B-MLX-bf16 \
        --state '{"text": "I was charged twice this month, please refund one of them."}' \
        --questions '{"refund": {"type": "noul", "instructions": "Is this a refund request?"}}'

The model sees the state and the question once, and the answer is read off the
logits at the answer slot. Nothing is decoded, so a schema violation is
impossible and tokens generated per decision is zero.
"""

from __future__ import annotations

import argparse
import json
import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping, Sequence

# --------------------------------------------------------------------------- #
# question types (Decision Index JSON schema)
# --------------------------------------------------------------------------- #


@dataclass(frozen=True)
class Choice:
    instructions: str
    criteria: Mapping[str, str | None]
    type: str = "choice"


@dataclass(frozen=True)
class Noul:
    instructions: str
    criteria: Mapping[str, Any] | None = None
    type: str = "noul"


@dataclass(frozen=True)
class Score:
    instructions: str
    criteria: Sequence[str]
    type: str = "score"


Question = Choice | Noul | Score


def as_question(q: Mapping | Question) -> Question:
    """`{"type": "noul" | "choice" | "score", "instructions", "criteria"}` as a class; a class passes through."""
    if not isinstance(q, Mapping):
        return q
    kind = q.get("type", "choice")
    if kind == "noul":
        return Noul(q["instructions"], q.get("criteria"))
    if kind == "score":
        return Score(q["instructions"], list(q["criteria"]))
    return Choice(q["instructions"], q["criteria"])


# --------------------------------------------------------------------------- #
# verbalizer: option -> single-token ids
# --------------------------------------------------------------------------- #


def option_codes(labels: Sequence[str]) -> list[str]:
    """Native letters when the labels already are letters, else A..Z, else 00.."""
    labs = [str(x).strip() for x in labels]
    if labs and all(len(k) == 1 and k.isalpha() for k in labs):
        return labs
    if len(labs) <= 26:
        return [chr(ord("A") + i) for i in range(len(labs))]
    return [f"{i:02d}" for i in range(len(labs))]


_FALLBACK_POOL = (
    [chr(c) for c in range(ord("A"), ord("Z") + 1)]
    + [f"{i:02d}" for i in range(100)]
    + [chr(c) for c in range(ord("a"), ord("z") + 1)]
    + [f"#{i}" for i in range(200)]
    + [chr(a) + chr(b) for a in range(ord("A"), ord("Z") + 1) for b in range(ord("A"), ord("Z") + 1)]
)


def aliases(tokenizer, labels: Sequence[str]) -> list[tuple[str, int]]:
    """Assign every label a distinct single-token code: [(code, token_id)]."""
    used: set[int] = set()
    out: list[tuple[str, int]] = []

    def take(raw: str) -> bool:
        enc = tokenizer.encode(raw, add_special_tokens=False)
        if len(enc) != 1 or enc[0] in used:
            return False
        out.append((raw, enc[0]))
        used.add(enc[0])
        return True

    for code in option_codes(labels):
        if take(code):
            continue
        if not any(take(raw) for raw in _FALLBACK_POOL):
            raise RuntimeError(f"no single-token alias left for {len(option_codes(labels))} options")
    return out


def _ids(tokenizer, texts: Sequence[str]) -> list[int]:
    """Single-token ids for the texts that encode to one token, deduped, in order."""
    out, seen = [], set()
    for t in texts:
        enc = tokenizer.encode(t, add_special_tokens=False)
        if len(enc) == 1 and enc[0] not in seen:
            out.append(enc[0])
            seen.add(enc[0])
    return out


YES_FORMS = ("yes", "Yes", "YES")
NO_FORMS = ("no", "No", "NO")


def readout_ids(tokenizer, q: Question) -> list[list[int]]:
    """Token ids to score, one group per option, max-pooled."""
    if isinstance(q, Noul):
        yes, no = _ids(tokenizer, YES_FORMS), _ids(tokenizer, NO_FORMS)
        if not yes or not no:
            raise RuntimeError("tokenizer has no single-token yes/no")
        return [yes, no]
    if isinstance(q, Score):
        groups = [_ids(tokenizer, [str(i)]) for i in range(len(q.criteria))]
        if any(not g for g in groups):
            raise RuntimeError(
                f"score with {len(q.criteria)} levels needs single-token digits; "
                "the primitive is defined for 2 to 10"
            )
        return groups
    groups = []
    for code, tid in aliases(tokenizer, list(q.criteria.keys())):
        extra = _ids(tokenizer, [f" {code}"])
        groups.append([tid] + [i for i in extra if i != tid])
    if not groups:
        raise RuntimeError("choice with no options")
    return groups


def readout(tokenizer, q: Question, logz, calibration=None) -> list[float]:
    """Option probabilities from the log-probabilities at the answer slot.

    `logz` is indexed by token id (a 1-D array over the vocabulary). Each
    option scores its best form; the answer is a softmax over these and
    nothing else.
    """
    scores = [max(float(logz[i]) for i in g) for g in readout_ids(tokenizer, q)]
    if calibration is not None:
        scores = calibration.apply(q, scores)
    m = max(scores)
    exps = [math.exp(s - m) for s in scores]
    return [e / sum(exps) for e in exps]


# --------------------------------------------------------------------------- #
# prompt rendering (identical to upstream prompt.py, state_style json_only)
# --------------------------------------------------------------------------- #

IM_START = "<|im_start|>"
IM_END = "<|im_end|>"


def state_block(state: Any) -> str:
    """The default `json_only` style: a string passes through, anything else as JSON."""
    if isinstance(state, str):
        return f"{state}\n\n"
    return json.dumps(state, ensure_ascii=False, indent=2) + "\n\n"


def _option_line(code: str, label: str, desc: str | None) -> str:
    return f"{code} {desc or label.replace('_', ' ')}"


def question_block(tokenizer, q: Question) -> str:
    if isinstance(q, Choice):
        labels = list(q.criteria.keys())
        codes = aliases(tokenizer, labels)
        lines = "\n".join(_option_line(codes[i][0], lab, q.criteria[lab]) for i, lab in enumerate(labels))
        return f"{q.instructions}\n\nOptions:\n{lines}\n\nReply with the option code only."
    if isinstance(q, Noul):
        extra = ""
        if q.criteria:
            extra = f"\nYes: {q.criteria.get('true')}\nNo: {q.criteria.get('false')}"
        return f"{q.instructions}{extra}\n\nReply with yes or no only."
    if isinstance(q, Score):
        legend = "\n".join(f"{i} {name}" for i, name in enumerate(q.criteria))
        return f"{q.instructions}\n\n{legend}\n\nReply with a single digit 0-{len(q.criteria) - 1} only."
    raise TypeError(f"unknown question type {type(q)}")


def prefix_text(state: Any, bos: str = "", images: str = "") -> str:
    """Everything before the question: pictures' markup at the head of the user turn, then the state."""
    body = "" if state is None else f"{state_block(state)}\nQUESTION:\n"
    return f"{bos}{IM_START}user\n{images}{body}"


def suffix_text(tokenizer, q: Question) -> str:
    """The question and the assistant header, up to the answer slot (no lead on LFM2-VL)."""
    return f"{question_block(tokenizer, q)}{IM_END}\n{IM_START}assistant\n"


def render(tokenizer, state: Any, q: Question, bos: str = "", images: str = "") -> str:
    return prefix_text(state, bos, images) + suffix_text(tokenizer, q)


def answer(q: Question, probs: Sequence[float]) -> dict:
    """A noul's P(yes); a choice's pick and its probabilities; a score's expected level."""
    if isinstance(q, Noul):
        return {"type": "noul", "noul": probs[0]}
    best = max(range(len(probs)), key=probs.__getitem__)
    if isinstance(q, Choice):
        names = list(q.criteria)
        return {"type": "choice", "choice": names[best], "confidence": probs[best],
                "probabilities": dict(zip(names, probs))}
    return {"type": "score", "score": sum(i * p for i, p in enumerate(probs)), "confidence": probs[best],
            "probabilities": {str(i): p for i, p in enumerate(probs)},
            "legend": {str(i): text for i, text in enumerate(q.criteria)}}


# --------------------------------------------------------------------------- #
# the MLX runner
# --------------------------------------------------------------------------- #

VISION_MAX_PIXELS = 1024 * 1024


def cap_pixels(image, max_pixels: int = VISION_MAX_PIXELS):
    """A picture downscaled to at most `max_pixels`, bicubic (as upstream does)."""
    image = image.convert("RGB") if hasattr(image, "convert") else image
    w, h = image.size
    if w * h <= max_pixels:
        return image
    from PIL import Image

    scale = math.sqrt(max_pixels / (w * h))
    return image.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.Resampling.BICUBIC)


class SystemOne:
    """State in, calibrated distribution out, one forward pass — on MLX."""

    def __init__(self, model_path: str, calibration=None):
        from mlx_vlm import load
        from mlx_vlm.utils import load_image  # noqa: F401  (re-exported for callers)

        self.model, self.processor = load(model_path)
        self.tokenizer = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor
        self.model_path = model_path
        self.calibration = calibration
        bos = getattr(self.tokenizer, "bos_token", None)
        self.bos = bos if isinstance(bos, str) else ""

    # ---------------------------------------------------------------- vision

    def _image_markup(self, n: int) -> str:
        """What the chat template writes for `n` images at the head of a user turn."""
        msgs = [{"role": "user", "content": [*([{"type": "image"}] * n), {"type": "text", "text": "\x00"}]}]
        text = self.processor.apply_chat_template(msgs, add_generation_prompt=False, tokenize=False)
        head = f"{IM_START}user\n"
        return text[text.index(head) + len(head):text.index("\x00")]

    # --------------------------------------------------------------- forward

    def _logz(self, prompt: str, images: Sequence = ()):
        """Log-softmax at the answer slot for one rendered prompt."""
        import mlx.core as mx
        from mlx_vlm.utils import prepare_inputs

        pics = [cap_pixels(im) for im in images]
        inputs = prepare_inputs(
            self.processor,
            images=list(pics) or None,
            prompts=[prompt],
            image_token_index=getattr(self.model.config, "image_token_index", None),
            add_special_tokens=False,  # the prompt carries its own BOS
        )
        input_ids = inputs["input_ids"]
        pixel_values = inputs.get("pixel_values", None)
        mask = inputs.get("attention_mask", None)
        data_kwargs = {k: v for k, v in inputs.items()
                       if k not in ("input_ids", "pixel_values", "attention_mask")}
        output = self.model(input_ids, pixel_values, mask, **data_kwargs)
        logits = output.logits if hasattr(output, "logits") else output
        row = logits[0, -1]
        return row - mx.logsumexp(row, axis=-1), int(input_ids.shape[1])

    # ------------------------------------------------------------------- api

    def probabilities(self, state: Any, questions: Sequence, images: Sequence | None = None) -> list[list[float]]:
        """Each question's distribution over its options (`yes`, `no` for a noul).

        One forward pass per question: the upstream runner packs several
        questions of one state into a shared trunk; this helper re-reads the
        state per question — same answers, more compute.
        """
        qs = [as_question(q) for q in questions]
        pics = [cap_pixels(im) for im in (images or ())]
        markup = self._image_markup(len(pics)) if pics else ""
        out = []
        for q in qs:
            logz, _ = self._logz(render(self.tokenizer, state, q, self.bos, markup), pics)
            out.append(readout(self.tokenizer, q, logz, self.calibration))
        return out

    def system_one(self, state: Any, questions: Mapping[str, Any], images: Sequence | None = None) -> dict:
        """Named questions over one state, and its pictures if any, in one forward pass each."""
        qs = {n: as_question(q) for n, q in questions.items()}
        probs = self.probabilities(state, list(qs.values()), images)
        return {"answers": {n: answer(q, p) for (n, q), p in zip(qs.items(), probs)},
                "usage": {"input_tokens": 0, "output_tokens": 0}}


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def main(argv=None):
    parser = argparse.ArgumentParser(description="System One decisions with MLX")
    parser.add_argument("--model", required=True, help="MLX model path or Hugging Face ID")
    parser.add_argument("--state", help="Text to evaluate")
    parser.add_argument("--state-file", type=Path, help="UTF-8 text file to evaluate")
    parser.add_argument("--questions", type=json.loads, help="Named questions as JSON")
    parser.add_argument("--questions-file", type=Path, help="JSON file of named questions")
    parser.add_argument("--image", action="append", default=[], help="Image path or URL (repeatable)")
    args = parser.parse_args(argv)
    try:
        state = args.state_file.read_text(encoding="utf-8") if args.state_file is not None else args.state
        questions = (json.loads(args.questions_file.read_text(encoding="utf-8"))
                     if args.questions_file is not None else args.questions)
        for name, q in questions.items():
            as_question(q)
    except (OSError, ValueError, KeyError, TypeError) as error:
        parser.error(str(error))
    engine = SystemOne(args.model)
    images = None
    if args.image:
        from mlx_vlm.utils import load_image

        images = [load_image(p) for p in args.image]
    result = engine.system_one(state, questions, images)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()