"""Decider 2B: a decision model for Jev/SystemOne requests. One file, as released on Hugging Face.

Reads a state and a schema of typed questions (choice, score, noul) and returns a probability for
every allowed option of every question in one forward pass: no text generation, no parsing.

  from decider_model import load_release_model, systemone
  model = load_release_model(path, device="cuda")
  systemone(model, {"state": "...", "questions": {...}})

Sections: the Jev/SystemOne request and response format, the token layout of a model pass, and the
model with its loader. The training code in training/ uses the same format and layout.
"""
import json
import math
import os
import re


# ================================================================================================
# Jev / SystemOne format
# ================================================================================================

QUESTION_TYPES = ("choice", "score", "noul")
MAX_REQUEST_TOKENS = 65536            # Jev: whole request
MAX_STATE_PLUS_QUESTION_TOKENS = 32768  # Jev: state plus the longest single question
MAX_QUESTIONS = 128                 # as pplx-decider; Cloudflare allows 64
MAX_OPTIONS = 512
MAX_KEY_CHARS = 128
NOUL_OPTIONS = ("true", "false")
# Selective answering: when a choice offers "I don't know" and no answer reaches this probability,
# the model abstains. A wrong answer costs more than an abstention wherever hallucination is
# counted (AA-Omniscience-Decision: answering pays only when p > 1 / (2 - accuracy), about 0.55).
# 0 turns it off. Tuned with aa_eval/run_aa.py --evals omniscience_proxy.
ABSTAIN_THRESHOLD = float(os.environ.get("DECIDER_ABSTAIN_THRESHOLD", "0.55"))
ABSTAIN_TEXT = re.compile(r"^\s*(i\s*do\s*n'?\s*t\s*know|i\s*do\s*not\s*know)\s*[.!]?\s*$", re.I)


class SchemaError(ValueError):
    """The request or record does not follow the format; the message says where."""


def _text(value):
    return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False, sort_keys=True)


def state_text(state):
    """One canonical text for a state, used in training and serving alike."""
    if isinstance(state, str):
        return state if state.strip() else "(empty)"
    if isinstance(state, list) and all(isinstance(item, str) for item in state):
        return "\n\n".join(state) if state else "(empty)"
    return json.dumps(state, ensure_ascii=False, indent=2)


def is_empty_state(state):
    return (isinstance(state, str) and not state.strip()) or (isinstance(state, (dict, list)) and not state)


def request_state_text(state, questions):
    """The state text a request is served with. When the state is empty, the content is in the
    question instructions (as in some Jev benchmark requests), so those are shown as the state:
    the model learned to read content from the state."""
    if is_empty_state(state):
        texts = list(dict.fromkeys(_text(q.get("instructions")).strip() for q in questions.values()
                                   if q.get("instructions") is not None))
        texts = [t for t in texts if t]
        if texts:
            return "\n\n".join(texts)
    return state_text(state)


def level_text(level):
    if isinstance(level, dict):
        label = level.get("label")
        extra = {k: v for k, v in level.items() if k != "label"}
        return f"{label} ({_text(extra)})" if extra else str(label)
    return str(level)


def validate_request(body):
    """Raise SchemaError unless body is a valid SystemOne request."""
    if not isinstance(body, dict):
        raise SchemaError("request must be a JSON object")
    if "state" not in body:
        raise SchemaError("missing 'state'")
    state = body["state"]
    if not isinstance(state, (str, dict, list)):
        raise SchemaError("'state' must be a string, an object or a list of strings")
    if isinstance(state, list) and not all(isinstance(item, (str, dict)) for item in state):
        raise SchemaError("'state' list items must be strings or objects")
    questions = body.get("questions")
    if not isinstance(questions, dict) or not questions:
        raise SchemaError("'questions' must be a non-empty object")
    if len(questions) > MAX_QUESTIONS:
        raise SchemaError(f"at most {MAX_QUESTIONS} questions per request")
    for qid, question in questions.items():
        where = f"questions.{qid}"
        if not qid or len(qid) > MAX_KEY_CHARS:
            raise SchemaError(f"{where}: question ids must be 1-{MAX_KEY_CHARS} characters")
        if not isinstance(question, dict):
            raise SchemaError(f"{where}: must be an object")
        kind = question.get("type")
        if kind not in QUESTION_TYPES:
            raise SchemaError(f"{where}.type: must be one of {', '.join(QUESTION_TYPES)}")
        # Jev requests may carry structured instructions (an object, e.g. a candidate document); any
        # JSON value is read as its canonical text.
        instructions = question.get("instructions")
        criteria = question.get("criteria")
        if kind == "choice":
            if not isinstance(criteria, dict) or len(criteria) < 2:
                raise SchemaError(f"{where}.criteria: a choice needs an object with at least 2 answers")
            if len(criteria) > MAX_OPTIONS:
                raise SchemaError(f"{where}.criteria: at most {MAX_OPTIONS} answers")
            for key, description in criteria.items():
                if not key or len(key) > MAX_KEY_CHARS:
                    raise SchemaError(f"{where}.criteria: answer keys must be 1-{MAX_KEY_CHARS} characters")
                # A description may be structured too (a chord's pitch classes, a palette, a move in two
                # notations); it is read as its canonical text.
        elif kind == "score":
            if not isinstance(criteria, list) or len(criteria) < 2:
                raise SchemaError(f"{where}.criteria: a score needs a list of at least 2 levels, lowest first")
            if len(criteria) > MAX_OPTIONS:
                raise SchemaError(f"{where}.criteria: at most {MAX_OPTIONS} levels")
            for i, level in enumerate(criteria):
                if not (isinstance(level, str) and level.strip()) and not (
                        isinstance(level, dict) and isinstance(level.get("label"), str) and level["label"].strip()):
                    raise SchemaError(f"{where}.criteria[{i}]: a level is a non-empty string or an object with 'label'")
        else:
            if not (instructions and _text(instructions).strip()):
                raise SchemaError(f"{where}.instructions: a noul question needs the statement to judge")
            # Jev takes no criteria on noul; Clef allows optional true/false wording. Accept that shape
            # (it is not used) so Clef-style clients work, and reject anything else.
            if criteria is not None and not (isinstance(criteria, dict) and set(criteria) <= {"true", "false"}):
                raise SchemaError(f"{where}.criteria: noul questions take no criteria")


def render(qid, question):
    """What the model reads for one question: prompt text, option lines, and the answer keys."""
    kind = question["type"]
    instructions = question.get("instructions")
    text = _text(instructions).strip() if instructions is not None else ""
    prompt = text or str(qid)
    criteria = question.get("criteria")
    if kind == "choice":
        keys = list(criteria)
        labels = [k if d in (None, "", k) else f"{k}: {_text(d)}" for k, d in criteria.items()]
    elif kind == "score":
        keys = list(range(len(criteria)))
        labels = [f"level {i}: {level_text(level)}" for i, level in enumerate(criteria)]
        prompt = f"{prompt}\n(Rate on the ordered levels below, lowest first.)"
    else:
        keys = [True, False]
        labels = list(NOUL_OPTIONS)
        prompt = prompt if prompt.endswith("?") else f"Is this statement true? {prompt}"
    return {"id": qid, "kind": kind, "prompt": prompt, "labels": labels, "keys": keys}


def items(questions):
    return [render(qid, question) for qid, question in questions.items()]


def answer(item, probabilities, criteria=None):
    """Jev answer object for one question from its option probabilities (in render order)."""
    total = sum(probabilities) or 1.0
    p = [float(x) / total for x in probabilities]
    best = max(range(len(p)), key=p.__getitem__)
    if item["kind"] == "choice" and criteria:
        p, best = abstain(item["keys"], criteria, p, best)
    if item["kind"] == "choice":
        return {"type": "choice", "choice": item["keys"][best], "confidence": p[best],
                "probabilities": dict(zip(item["keys"], p))}
    if item["kind"] == "score":
        return {"type": "score", "score": sum(i * v for i, v in enumerate(p)), "confidence": p[best],
                "legend": {str(i): level_text(level) for i, level in enumerate(criteria)},
                "probabilities": {str(i): v for i, v in enumerate(p)}}
    return {"type": "noul", "noul": p[0]}


def abstain_index(keys, criteria):
    """Index of the option that means "I don't know", or None."""
    for i, key in enumerate(keys):
        description = criteria.get(key) if isinstance(criteria, dict) else None
        if ABSTAIN_TEXT.match(str(description or "")) or ABSTAIN_TEXT.match(str(key).replace("_", " ")):
            return i
    return None


def abstain(keys, criteria, p, best, threshold=None):
    """(probabilities, best) after selective answering: below the threshold, "I don't know" wins.

    The answers keep their relative probabilities; "I don't know" takes the mass needed to be
    the highest, so the response stays a valid distribution.
    """
    threshold = ABSTAIN_THRESHOLD if threshold is None else threshold
    idk = abstain_index(keys, criteria)
    if idk is None or threshold <= 0 or best == idk or p[best] >= threshold:
        return p, best
    others = 1.0 - p[idk]
    top = max(v for i, v in enumerate(p) if i != idk)
    if others <= 0:
        return p, idk
    scale = 0.9 / (top + others)        # top * scale < 1 - others * scale, so "I don't know" is highest
    q = [v * scale if i != idk else 0.0 for i, v in enumerate(p)]
    q[idk] = 1.0 - sum(q)
    return q, idk


def response(model, questions, probabilities, input_tokens):
    return {"model": model,
            "answers": {item["id"]: answer(item, probs, questions[item["id"]].get("criteria"))
                        for item, probs in zip(items(questions), probabilities)},
            "usage": {"input_tokens": int(input_tokens), "output_tokens": 0}}


# ---- training records -------------------------------------------------------------------------

RECORD_FIELDS = ("id", "task", "source", "license", "state", "questions", "expected")


def validate_record(record):
    """Raise SchemaError unless the record is a valid request with a right answer for every question."""
    for field in RECORD_FIELDS:
        if field not in record:
            raise SchemaError(f"record is missing '{field}'")
    validate_request(record)
    expected = record["expected"]
    if not isinstance(expected, dict) or set(expected) != set(record["questions"]):
        raise SchemaError("'expected' must have exactly one answer per question")
    for qid, question in record["questions"].items():
        value = expected[qid]
        kind = question["type"]
        if kind == "choice" and value not in question["criteria"]:
            raise SchemaError(f"expected.{qid}: '{value}' is not one of the criteria keys")
        if kind == "score" and not (isinstance(value, int) and not isinstance(value, bool)
                                    and 0 <= value < len(question["criteria"])):
            raise SchemaError(f"expected.{qid}: must be a level index 0..{len(question['criteria']) - 1}")
        if kind == "noul" and not isinstance(value, bool):
            raise SchemaError(f"expected.{qid}: must be true or false")
    teacher = record.get("teacher")
    if teacher is not None:
        for qid, probs in teacher.items():
            if qid not in record["questions"] or not isinstance(probs, dict):
                raise SchemaError(f"teacher.{qid}: unknown question or not an object")
            if not math.isclose(sum(probs.values()), 1.0, abs_tol=1e-2):
                raise SchemaError(f"teacher.{qid}: probabilities must sum to 1")


def gold_index(item, value):
    """Position of the expected answer among the rendered options."""
    return item["keys"].index(value)


def teacher_vector(item, probs):
    """Teacher probabilities in render order, or None."""
    if not probs:
        return None
    if item["kind"] == "noul":
        p = float(probs.get("true", probs.get(True, 0.0)))
        return [p, 1.0 - p]
    return [float(probs.get(str(key), probs.get(key, 0.0))) for key in item["keys"]]


# ================================================================================================
# Token layout of one model pass
# ================================================================================================

INTRO = "Read the state, then answer each question by choosing one of its options.\n\nState:\n"


class TooLong(ValueError):
    pass


def question_pieces(k, item):
    """Text pieces of question k; marks say which pieces end at a scored token."""
    pieces = [(f"\n\nQuestion {k + 1}: {item['prompt']}\nOptions:", None)]
    pieces += [(f"\n[{i + 1}] {label}", "option") for i, label in enumerate(item["labels"])]
    pieces.append((f"\nAnswer {k + 1}:", "answer"))
    return pieces


def encode_pieces(tokenizer, texts):
    return tokenizer(texts, add_special_tokens=False)["input_ids"] if texts else []


def question_tokens(tokenizer, items):
    """Token pieces per question, numbered as if each were first in its pass (k is fixed later)."""
    return [list(zip(encode_pieces(tokenizer, [t for t, _ in question_pieces(0, item)]),
                     [m for _, m in question_pieces(0, item)])) for item in items]


def plan(tokenizer, state, items, max_length):
    """Split questions into passes that fit max_length with the whole state in each.

    Returns (passes, state_tokens, question_sizes); each pass is a list of item indices.
    Raises TooLong with the index of the first question that cannot fit even alone.
    """
    intro = encode_pieces(tokenizer, [INTRO, state])
    base = len(intro[0]) + len(intro[1])
    sizes = [sum(len(ids) for ids, _ in pieces) + 4 for pieces in question_tokens(tokenizer, items)]
    passes, current, used = [], [], base
    for index, size in enumerate(sizes):
        if base + size > max_length:
            raise TooLong(index)
        if current and used + size > max_length:
            passes.append(current)
            current, used = [], base
        current.append(index)
        used += size
    if current:
        passes.append(current)
    return passes, len(intro[1]), sizes


def build(tokenizer, state, items):
    """Token ids plus the scored positions for one pass over these items."""
    texts, marks = [INTRO, state], [None, None]
    for k, item in enumerate(items):
        for text, mark in question_pieces(k, item):
            texts.append(text)
            marks.append(mark)
    ids, answers, options, current = [], [], [], []
    for piece, mark in zip(encode_pieces(tokenizer, texts), marks):
        ids += piece
        if mark == "option":
            current.append(len(ids) - 1)
        elif mark == "answer":
            answers.append(len(ids) - 1)
            options.append(current)
            current = []
    return {"ids": ids, "answers": answers, "options": options}


# ================================================================================================
# Model: merged Qwen3.5-2B backbone (bidirectional full attention) + option-pointer decision head
# ================================================================================================

import torch

HEAD_WEIGHTS = "decider_head.safetensors"
HEAD_CONFIG = "decider_head_config.json"
NEG = -1e4


class Encoder:
    """Batch tensors for a list of {"state": text, "items": items(questions)} records."""

    def __init__(self, tokenizer, pad_to_multiple=8):
        self.tokenizer = tokenizer
        self.pad_to_multiple = pad_to_multiple

    def __call__(self, batch):
        layouts = [build(self.tokenizer, x["state"], x["items"]) for x in batch]
        width = max(len(x["ids"]) for x in layouts)
        width += (-width) % self.pad_to_multiple
        most = max(len(o) for x in layouts for o in x["options"])
        input_ids = torch.full((len(batch), width), self.tokenizer.pad_token_id, dtype=torch.long)
        attention_mask = torch.zeros((len(batch), width), dtype=torch.long)
        answer_pos, option_pos, counts = [], [], []
        for b, layout in enumerate(layouts):
            input_ids[b, : len(layout["ids"])] = torch.tensor(layout["ids"])
            attention_mask[b, : len(layout["ids"])] = 1
            for answer, positions in zip(layout["answers"], layout["options"]):
                answer_pos.append(b * width + answer)
                option_pos.append([b * width + p for p in positions] + [0] * (most - len(positions)))
                counts.append(len(positions))
        return {"input_ids": input_ids, "attention_mask": attention_mask,
                "answer_pos": torch.tensor(answer_pos, dtype=torch.long),
                "option_pos": torch.tensor(option_pos, dtype=torch.long),
                "option_mask": torch.arange(most)[None] < torch.tensor(counts)[:, None]}


def encode_record(record, questions=None):
    """{"state", "items"} for one Jev/SystemOne request (optionally a subset of its question ids)."""
    chosen = record["questions"] if questions is None else {q: record["questions"][q] for q in questions}
    return {"state": request_state_text(record["state"], record["questions"]), "items": items(chosen)}


class DeciderModel(torch.nn.Module):
    """The released model: one forward pass scores every option of every question in a request."""

    def __init__(self, path, dtype=torch.bfloat16, quantization=None, device=None):
        super().__init__()
        from safetensors.torch import load_file
        from transformers import AutoModelForCausalLM, AutoTokenizer

        with open(os.path.join(path, HEAD_CONFIG)) as file:
            self.config = json.load(file)
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.encoder = Encoder(self.tokenizer)
        if quantization:
            # 8-bit or 4-bit backbone via bitsandbytes (CUDA only); the decision head stays in float32.
            from transformers import BitsAndBytesConfig
            if quantization == "8bit":
                bnb = BitsAndBytesConfig(load_in_8bit=True)
            elif quantization == "4bit":
                bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)
            else:
                raise ValueError("quantization must be None, '8bit' or '4bit'")
            lm = AutoModelForCausalLM.from_pretrained(path, dtype=dtype, quantization_config=bnb,
                                                      device_map={"": device or "cuda"})
        else:
            lm = AutoModelForCausalLM.from_pretrained(path, dtype=dtype)
        head = lm.get_output_embeddings()
        hidden = head.weight.shape[1]
        # The vocabulary head is unused: the backbone returns its final hidden states.
        for module in list(lm.modules()):
            for name, child in module.named_children():
                if child is head:
                    setattr(module, name, torch.nn.Identity())
        self.text_config = getattr(lm.config, "text_config", lm.config)
        self.bidirectional = bool(self.config.get("bidirectional", True))
        if self.bidirectional:
            for module in lm.modules():
                if hasattr(module, "is_causal") and hasattr(module, "q_proj"):
                    module.is_causal = False
        self.lm = lm
        self.head_dim = int(self.config["head_dim"])
        self.norm = torch.nn.RMSNorm(hidden, eps=1e-6)
        self.query = torch.nn.Linear(hidden, self.head_dim, bias=False)
        self.key = torch.nn.Linear(hidden, self.head_dim, bias=False)
        self.prior = torch.nn.Linear(hidden, 1)
        state = load_file(os.path.join(path, HEAD_WEIGHTS))
        for name in ("norm", "query", "key", "prior"):
            getattr(self, name).load_state_dict({k.split(".", 1)[1]: v for k, v in state.items() if k.startswith(name + ".")})
        self.temperature = float(self.config.get("temperature", 1.0))
        self.max_length = int(self.config.get("max_length", 8192))

    def encode(self, batch):
        return self.encoder(batch)

    def masks(self, embeds, attention_mask):
        """Bidirectional masks for the full-attention layers, padding-only for linear attention."""
        from transformers.masking_utils import create_bidirectional_mask, create_recurrent_attention_mask

        return {"full_attention": create_bidirectional_mask(config=self.text_config, inputs_embeds=embeds,
                                                            attention_mask=attention_mask),
                "linear_attention": create_recurrent_attention_mask(self.text_config, embeds, attention_mask)}

    def forward(self, input_ids, attention_mask, answer_pos, option_pos, option_mask, **_):
        """Option logits [questions, most options] (masked options are -1e4). Divide by
        self.temperature and softmax per question for calibrated probabilities."""
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=input_ids.is_cuda):
            if self.bidirectional:
                embeds = self.lm.get_input_embeddings()(input_ids)
                hidden = self.lm(inputs_embeds=embeds, attention_mask=self.masks(embeds, attention_mask)).logits
            else:
                hidden = self.lm(input_ids=input_ids, attention_mask=attention_mask).logits
        flat = hidden.reshape(-1, hidden.shape[-1]).float()
        query = self.query(self.norm(flat[answer_pos]))
        options = self.norm(flat[option_pos])
        logits = torch.einsum("qd,qmd->qm", query, self.key(options)) / math.sqrt(self.head_dim)
        logits = logits + self.prior(options).squeeze(-1)
        return logits.masked_fill(~option_mask, NEG)


def load_release_model(path, device="cuda", dtype=None, quantization=None):
    """Load the released model from a local folder (e.g. huggingface_hub.snapshot_download(...)).

    quantization: None (bf16 on GPU, float32 on CPU), "8bit" or "4bit" (bitsandbytes, CUDA only:
    about 2.5 GB or 1.8 GB of GPU memory instead of 4.5 GB).
    The tuned abstention threshold in decider_head_config.json is applied unless the
    DECIDER_ABSTAIN_THRESHOLD environment variable is set.
    """
    global ABSTAIN_THRESHOLD
    dtype = dtype or (torch.bfloat16 if str(device).startswith("cuda") else torch.float32)
    model = DeciderModel(path, dtype=dtype, quantization=quantization, device=device)
    model.head_device = device
    if "abstain_threshold" in model.config and "DECIDER_ABSTAIN_THRESHOLD" not in os.environ:
        ABSTAIN_THRESHOLD = float(model.config["abstain_threshold"])
    model.norm.float(); model.query.float(); model.key.float(); model.prior.float()
    if quantization:  # the backbone is already placed; move only the head
        for part in (model.norm, model.query, model.key, model.prior):
            part.to(device)
        return model.eval()
    return model.to(device).eval()


@torch.inference_mode()
def systemone(model, body, name="decider-2b"):
    """Answer one Jev/SystemOne request body; returns {"model", "answers", "usage"}.

    Questions that do not fit next to the full state in one pass are split across passes, each
    with the whole state; nothing is truncated.
    """
    validate_request(body)
    questions = body["questions"]
    state = request_state_text(body["state"], questions)
    rendered = items(questions)
    device = next(model.parameters()).device
    passes, state_tokens, sizes = plan(model.tokenizer, state, rendered, model.max_length)
    probabilities = [None] * len(rendered)
    for indices in passes:
        tensors = model.encode([{"state": state, "items": [rendered[i] for i in indices]}])
        logits = model(**{k: v.to(device) for k, v in tensors.items()}).float() / model.temperature
        for row, i in enumerate(indices):
            count = len(rendered[i]["labels"])
            probabilities[i] = logits[row, :count].softmax(-1).tolist()
    total = len(passes) * state_tokens + sum(sizes)
    return response(body.get("model") or name, questions, probabilities, total)