"""JEV-9B inference core, ported 1:1 from the model authors' documented `transformers` +
`peft` path (https://huggingface.co/autotrust/JEV-9B -> "Other ways to run it").

The Blocks-of-Experts architecture:
  * backbone  = frozen Qwen3.5-9B text tower (bit-identical to the release)
  * System 1  = LoRA r=16 (40.1M, shipped unmerged in `adapter/`) + a 24-slot fp32
               decision head (`head.safetensors`) initialised from `lm_head` verbalizer rows
  * System 2  = the original `lm_head`; used with the adapter off (pristine Qwen3.5-9B)

System 1 answers a typed question (`noul` yes/no, `choice` over 2-16 options, `score` on a
0-5 scale) in ONE prefill pass and returns a calibrated probability distribution. The
decision prompt follows the authors' "bare-v1" template; the head is applied to the last
token's post-norm hidden state, divided by the per-kind calibration temperature.
"""

import json
import os

import torch
from huggingface_hub import hf_hub_download
from peft import PeftModel
from safetensors.torch import load_file
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "autotrust/JEV-9B"
LETTERS = "ABCDEFGHIJKLMNOP"
KINDS = ("noul", "choice", "score")
MAX_STATE_TOKENS = 1024  # the authors truncate long inputs to 1,024 tokens (head 60% / tail 40%)
DEVICE = os.environ.get("JEV_DEVICE", "cuda")
# Optional local-dir override (used only for offline verification of the wiring).
_LOCAL = os.environ.get("JEV_MODEL_PATH")


def _file(name: str) -> str:
    """Resolve a model-repo file to a local path, falling back to the Hub."""
    if _LOCAL:
        p = os.path.join(_LOCAL, name)
        if os.path.exists(p):
            return p
    return hf_hub_download(REPO, name)


_SRC = _LOCAL or REPO

# ---------------------------------------------------------------------------
# Model loading. On ZeroGPU the final `_peft.to(DEVICE)` is intercepted at
# startup: tensors are packed to disk and streamed into VRAM on the first
# @spaces.GPU entry, so the load cost is never paid inside a request.
#
# Everything is built on CPU first, then moved once:
#   * peft's default would be `torch_device=infer_device()` -> "cuda", making
#     safetensors read the adapter straight onto the GPU. That bypasses the
#     packer and crashes the GPU-less main process ("No CUDA GPUs are
#     available"). `torch_device="cpu"` keeps the read on the CPU.
#   * one `.to(DEVICE)` at the end moves every tensor (backbone, LoRA, head)
#     through the standard machinery.
# ---------------------------------------------------------------------------
_tok = AutoTokenizer.from_pretrained(_SRC)
_base = AutoModelForCausalLM.from_pretrained(_SRC, dtype=torch.bfloat16)

# Keep the adapter UNMERGED so System 2 (generation) runs with the adapter
# disabled -- `with _peft.disable_adapter():` -- i.e. pristine Qwen3.5-9B.
_peft = PeftModel.from_pretrained(_base, _SRC, subfolder="adapter", torch_device="cpu")

# System 1 head: fp32 linear, 4096 -> 24 slots.
_HEAD = torch.nn.Linear(4096, 24, dtype=torch.float32)
_head_sd = load_file(_file("head.safetensors"))
with torch.no_grad():
    _HEAD.weight.copy_(_head_sd["proj.weight"])
    _HEAD.bias.copy_(_head_sd["proj.bias"])
_HEAD.eval()

_JUDGE_CFG = json.load(open(_file("judge_config.json")))
_CALIB = json.load(open(_file("calibration.json")))["per_kind"]
SLOT_RANGES = {k: tuple(v) for k, v in _JUDGE_CFG["slots"]["ranges"].items()}
TEMPERATURES = {k: float(_CALIB[k]) for k in KINDS}

# Module handles: PeftModel -> LoraModel -> Qwen3_5ForCausalLM -> Qwen3_5TextModel.
_CAUSAL = _peft.base_model.model          # the Qwen3_5ForCausalLM with LoRA modules inside
_BACKBONE = _CAUSAL.model                 # the text tower (LoRA active by default)

# Single move of the whole tree (backbone + LoRA + head) to the compute device.
_peft.to(DEVICE)
_HEAD.to(DEVICE)

# ---------------------------------------------------------------------------
# Prompt template -- exactly the authors' `bare-v1` decision format.
# ---------------------------------------------------------------------------


def build_decision_prompt(kind: str, state: str, question: str, options) -> tuple:
    """Return (prompt_text, option_labels) in the authors' `bare-v1` format."""
    if kind == "noul":
        options = ["false", "true"]
    elif kind == "score":
        options = [str(i) for i in range(6)]
    options = list(options)
    if kind == "choice":
        lines = [f"{LETTERS[i]}) {o}" for i, o in enumerate(options)]
    else:
        lines = options
    prompt = (
        f"[kind] {kind}\n[state] {state}\n[question] {question}\n[options]\n"
        + "\n".join(lines)
        + "\n[decision]:"
    )
    return prompt, options


def truncate_state(state: str) -> tuple:
    """Head-60% / tail-40% truncation of the state to `MAX_STATE_TOKENS` tokens.

    Returns (state_used, was_truncated, n_tokens).
    """
    ids = _tok(state, add_special_tokens=False)["input_ids"]
    if len(ids) <= MAX_STATE_TOKENS:
        return state, False, len(ids)
    keep_head = int(MAX_STATE_TOKENS * 0.6)
    keep_tail = MAX_STATE_TOKENS - keep_head
    return _tok.decode(ids[:keep_head] + ids[-keep_tail:], skip_special_tokens=True), True, len(ids)


def decide(kind: str, state: str, question: str, options=None) -> dict:
    """Run ONE System 1 typed decision; return the calibrated probability distribution.

    This mirrors the model card's reference `decide()` byte for byte: one prefill pass
    through the backbone with the LoRA active, the 24-slot head applied to the last
    token's hidden state, per-kind temperature, softmax over the active slots.
    """
    if kind not in KINDS:
        raise ValueError(f"unknown kind {kind!r}")
    if kind == "choice":
        if not isinstance(options, (list, tuple)) or not 2 <= len(options) <= 16:
            raise ValueError("`choice` needs between 2 and 16 options")
        if len(set(options)) != len(options):
            raise ValueError("`choice` options must be unique")
    prompt, labels = build_decision_prompt(kind, state, question, options)
    device = _BACKBONE.embed_tokens.weight.device   # resolve inside the GPU call
    ids = _tok(prompt, return_tensors="pt", add_special_tokens=False).to(device)
    with torch.no_grad():
        h = _BACKBONE(**ids).last_hidden_state[0, -1].float()
        z = _HEAD(h.to(_HEAD.weight.device))       # [24] fp32 logits
    lo = SLOT_RANGES[kind][0]
    z = z[lo : lo + len(labels)] / TEMPERATURES[kind]
    p = torch.softmax(z, dim=0)
    return {label: float(v) for label, v in zip(labels, p.tolist())}


def n_state_tokens(state: str) -> int:
    return len(_tok(state, add_special_tokens=False)["input_ids"])


# ---------------------------------------------------------------------------
# System 2 -- ordinary generation from the pristine base model (adapter off).
# ---------------------------------------------------------------------------


def build_escalation_prompt(state: str, question: str, options) -> str:
    """The model card's confidence-gated escalation prompt (System 1 -> System 2)."""
    opts = [str(o) for o in options]
    if opts == ["false", "true"]:
        tail = "Answer `true` or `false` on the last line."
    else:
        tail = "Options: " + "; ".join(opts) + "\nThink it through, then give exactly one option on the last line."
    return f"{state}\n\nQuestion: {question}\n{tail}"


def generate(state: str, question: str, options, max_new_tokens: int = 512,
             enable_thinking: bool = True) -> tuple:
    """System 2: generate a reasoned answer with the adapter OFF -- pristine Qwen3.5-9B.

    Returns (text, completed): `completed` is False when generation hit the token
    budget mid-sentence (the model is verbose in thinking mode), which the UI needs
    to avoid presenting a cut-off trace as a finished answer.
    """
    prompt = build_escalation_prompt(state, question, options)
    msgs = [{"role": "user", "content": prompt}]
    enc = _tok.apply_chat_template(
        msgs, add_generation_prompt=True, enable_thinking=enable_thinking,
        return_tensors="pt", return_dict=True,
    ).to(_BACKBONE.embed_tokens.weight.device)
    n_prompt = enc["input_ids"].shape[1]
    with torch.no_grad(), _peft.disable_adapter():
        out = _CAUSAL.generate(**enc, max_new_tokens=int(max_new_tokens), do_sample=False)
    text = _tok.decode(out[0, n_prompt:], skip_special_tokens=True)
    completed = out.shape[1] - n_prompt < int(max_new_tokens)  # stopped before the cap
    return text.strip(), completed