# Copyright 2026 FrontiersMind. Licensed under the Apache License, Version 2.0.
"""Lumma-Fev decision model for transformers (loaded with trust_remote_code).

One document (the state) and a set of typed questions in, one probability distribution per question out. The model
reads; it never generates text.

    from transformers import AutoModel
    model = AutoModel.from_pretrained("FrontiersMind/<model>", trust_remote_code=True)
    model.decide("I was charged twice.", {"billing": {"type": "noul", "instructions": "Is this about billing?"}})

Every question is scored on its own causal row, the state followed by that question's tokens:

    <state> state text <q> instructions <opt> option 1 </opt> <opt> option 2 </opt> ... <decide>

so questions never see each other. The pointer head compares the hidden state at <decide> with the hidden state at
each </opt> and a softmax over the options gives the distribution. All rows of a request run in one forward pass
(several when a very long request would not fit one pass).
"""
import math
import re

import torch
import torch.nn as nn
from transformers import AutoTokenizer, PreTrainedModel

from .configuration_fev import FevConfig
from .fev_backbone import BackboneConfig, BackboneModel

MAX_OPTIONS = 255
MAX_ROW = 8192          # tokens in one row (state + one question), lowered to the backbone's context window
QUESTION_RESERVE = 640  # when the window lowers the row limit, the state gives way so each question keeps this budget
PASS_TOKENS = 16384     # tokens per forward pass; the rows of a longer request run in several passes
_CONTROL = re.compile(r"<\|([A-Za-z0-9_]+)\|>")


def render(value, indent=0):
    """Flatten a JSON state, instruction or option (str, number, object, array) into the text the model reads."""
    pad = "  " * indent
    if value is None: return ""
    if isinstance(value, (str, int, float, bool)): return str(value)
    if isinstance(value, list): return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in value)
    return "\n".join(f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}"
                     for k, x in value.items())


def option_text(name, description):
    return name if description is None or description == "" else f"{name}: {render(description)}"


def as_dict(question):
    """A question as a plain dict: dicts pass through, pydantic question objects (TypeSafe SDK types) are dumped."""
    if hasattr(question, "model_dump"): question = question.model_dump(mode="json")
    if not isinstance(question, dict) or "type" not in question:
        raise ValueError("every question must be a dict with a 'type' (noul, choice or score)")
    return question


def pack(questions):
    """TypeSafe-shaped questions -> per question (instruction text, option texts, answer metadata)."""
    if not questions: raise ValueError("questions must be a non-empty mapping of names to questions")
    packed = []
    for qid, q in questions.items():
        q = as_dict(q)
        kind, criteria = q["type"], q.get("criteria")
        if kind == "noul":
            criteria = criteria or {}
            options, meta = [option_text("no", criteria.get("false")), option_text("yes", criteria.get("true"))], {}
        elif kind == "choice":
            if not isinstance(criteria, dict) or not 1 <= len(criteria) <= MAX_OPTIONS:
                raise ValueError(f"choice {qid!r}: criteria must map 1 to {MAX_OPTIONS} option names to descriptions")
            options, meta = [option_text(k, v) for k, v in criteria.items()], {"keys": list(criteria)}
        elif kind == "score":
            if not isinstance(criteria, list) or not 1 <= len(criteria) <= MAX_OPTIONS:
                raise ValueError(f"score {qid!r}: criteria must list 1 to {MAX_OPTIONS} ordered levels")
            options = [render(x) for x in criteria]
            meta = {"legend": {str(i): text for i, text in enumerate(options)}}
        else:
            raise ValueError(f"question {qid!r}: unknown type {kind!r} (noul, choice or score)")
        packed.append((render(q.get("instructions")), options, {"id": qid, "type": kind, **meta}))
    return packed


def choice_confidence(p):
    k = len(p)
    return 1.0 if k == 1 else (max(p) - 1 / k) / (1 - 1 / k)


def score_confidence(p):
    n, mode = len(p), max(range(len(p)), key=lambda i: p[i])
    return 1.0 if n == 1 else 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / (n - 1)


def answer(p, meta):
    """One question's probabilities -> its TypeSafe-shaped answer (4 decimals keep 255 rounded values summing to ~1)."""
    r = lambda x: round(float(x), 4)
    if meta["type"] == "noul":
        return {"type": "noul", "noul": r(p[1])}
    if meta["type"] == "choice":
        best = max(range(len(p)), key=lambda i: p[i])
        return {"type": "choice", "choice": meta["keys"][best], "confidence": r(choice_confidence(p)),
                "probabilities": {k: r(v) for k, v in zip(meta["keys"], p)}}
    return {"type": "score", "score": r(sum(i * pi for i, pi in enumerate(p))), "confidence": r(score_confidence(p)),
            "legend": meta["legend"], "probabilities": {str(i): r(v) for i, v in enumerate(p)}}


class PointerHead(nn.Module):
    def __init__(self, hidden_size, head_dim):
        super().__init__()
        self.q, self.k = nn.Linear(hidden_size, head_dim), nn.Linear(hidden_size, head_dim)
        self.scale = 1 / math.sqrt(head_dim)

    def forward(self, h_decide, h_options):   # [d], [K, d] -> logits [K]
        dtype = self.q.weight.dtype
        return ((self.k(h_options.to(dtype)) @ self.q(h_decide.to(dtype))) * self.scale).float()


class FevForDecision(PreTrainedModel):
    """Decision model: backbone (no language-model head) + pointer head. There is no generate()."""

    config_class = FevConfig
    base_model_prefix = "lm"
    _keep_in_fp32_modules = ["head"]
    _keep_in_fp32_modules_strict = ["head"]   # the plain list upcasts under fp16 only; the head must stay fp32 under bf16 too
    _no_split_modules = []

    def __init__(self, config, tokenizer=None):
        super().__init__(config)
        backbone = {k: v for k, v in config.backbone_config.items() if k not in ("auto_map", "architectures", "_name_or_path")}
        self.lm = BackboneModel(BackboneConfig(**backbone))
        self.head = PointerHead(config.hidden_size, config.head_dim)
        self.tokenizer = tokenizer
        self.post_init()

    def _init_weights(self, module):
        if isinstance(module, PointerHead):
            for linear in (module.q, module.k):
                nn.init.normal_(linear.weight, std=0.02); nn.init.zeros_(linear.bias)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded   # output_loading_info=True returns (model, info)
        hub = {k: kwargs[k] for k in ("revision", "cache_dir", "token", "local_files_only", "subfolder") if k in kwargs}
        model.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True, **hub)
        model.eval()
        return loaded

    # --- encoding ------------------------------------------------------------------------------------------------

    def limits(self):
        """(state tokens, row tokens) a request may use: MAX_ROW, lowered to the backbone's context window."""
        window = int(self.lm.config.max_position_embeddings)
        if MAX_ROW <= window: return MAX_ROW, MAX_ROW
        return min(MAX_ROW, window - QUESTION_RESERVE), window

    def delimiter_ids(self, tok):
        ids = [tok.convert_tokens_to_ids(t) for t in self.config.delimiters]
        if any(i is None or i == tok.unk_token_id for i in ids):
            raise ValueError(f"the tokenizer lacks the delimiter tokens {self.config.delimiters}")
        return ids

    @staticmethod
    def text_ids(tok, text):
        """Token ids of caller text. `<|name|>` becomes `<¦name¦>` first, so no input can forge a delimiter or control
        token; the tokenizers backend is called directly because the tokenizer's __call__ adds a leading marker."""
        return tok.backend_tokenizer.encode(_CONTROL.sub(r"<¦\1¦>", text), add_special_tokens=False).ids

    def encode(self, tok, state, packed):
        """-> (state ids, one row per question {"ids", "decide", "options"}). A long state is truncated to the state
        limit; a question that cannot fit its row raises ValueError."""
        s_id, q_id, o_id, c_id, d_id = self.delimiter_ids(tok)
        max_state, max_row = self.limits()
        S = [s_id] + self.text_ids(tok, render(state))[: max_state - 1]
        rows = []
        for instruction, options, _ in packed:
            ids, ends = [q_id] + self.text_ids(tok, instruction), []
            for option in options:
                ids += [o_id] + self.text_ids(tok, option) + [c_id]; ends.append(len(ids) - 1)
            ids.append(d_id)
            if len(S) + len(ids) > max_row:
                raise ValueError(f"a question needs {len(ids)} tokens after a {len(S)}-token state; one row holds {max_row}")
            rows.append({"ids": ids, "decide": len(ids) - 1, "options": ends})
        return S, rows

    # --- scoring -------------------------------------------------------------------------------------------------

    @torch.no_grad()
    def probabilities(self, tok, S, rows):
        """One probability tensor per row: every row is the state followed by its question, scored causally."""
        device, pad = self.head.q.weight.device, self.config.pad_token_id or 0
        full = [S + r["ids"] for r in rows]
        per_pass = max(1, PASS_TOKENS // max(len(x) for x in full))
        out = []
        for start in range(0, len(full), per_pass):
            part = full[start:start + per_pass]
            L = max(len(x) for x in part)
            ids = torch.full((len(part), L), pad, dtype=torch.long, device=device)
            mask = torch.zeros((len(part), L), dtype=torch.long, device=device)
            for i, x in enumerate(part):
                ids[i, : len(x)] = torch.tensor(x, device=device); mask[i, : len(x)] = 1
            pos = torch.arange(L, device=device).expand(len(part), L)
            h = self.lm(input_ids=ids, attention_mask=mask, position_ids=pos, use_cache=False).last_hidden_state
            for i, r in enumerate(rows[start:start + per_pass]):
                z = self.head(h[i, len(S) + r["decide"]], h[i, [len(S) + o for o in r["options"]]])
                out.append(torch.softmax(z / self.config.temperature, -1).cpu())
        return out

    def decide(self, state, questions, tokenizer=None, return_usage=False):
        """TypeSafe-shaped request -> {question name: answer}. With return_usage, also {"input_tokens", "output_tokens"}.

        state: text, a JSON object or an array. questions: {name: {"type": "noul" | "choice" | "score",
        "instructions": ..., "criteria": ...}} (dicts or TypeSafe SDK question objects)."""
        tok = tokenizer or self.tokenizer
        if tok is None: raise ValueError("pass tokenizer=..., or load the model with from_pretrained")
        packed = pack(questions)
        S, rows = self.encode(tok, state, packed)
        answers = {meta["id"]: answer(p.tolist(), meta) for p, (_, _, meta) in zip(self.probabilities(tok, S, rows), packed)}
        if not return_usage: return answers
        import json
        usage = {"input_tokens": len(S) + sum(len(r["ids"]) for r in rows),
                 "output_tokens": len(tok.backend_tokenizer.encode(json.dumps(answers), add_special_tokens=False).ids)}
        return answers, usage

    def forward(self, input_ids, attention_mask=None, position_ids=None, **kwargs):
        """The backbone's hidden states (the pointer readout needs the row layout; use decide())."""
        return self.lm(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False, **kwargs)