"""LLM as Jev-like classifier. Single-file Space."""
import json
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

LABELS = "0123456789abcdefghijklmnopqrstuvwxyz"
DEFAULT_MODEL = "Qwen/Qwen3-0.6B"
MODELS = [DEFAULT_MODEL, "Qwen/Qwen3-1.7B", "Qwen/Qwen3-4B",
          "Qwen/Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-1.5B-Instruct"]
_cache = {}


def get_model(mid):
    if mid not in _cache:
        tok = AutoTokenizer.from_pretrained(mid, trust_remote_code=True)
        tok.pad_token = tok.pad_token or tok.eos_token
        tok.padding_side = "left"
        m = AutoModelForCausalLM.from_pretrained(mid, trust_remote_code=True, dtype="auto", device_map="auto").eval()
        _cache[mid] = (m, tok)
    return _cache[mid]


def prompt(state, question, options):
    opts = "\n".join(f"<option_{i}>{o}</option_{i}>" for i, o in enumerate(options))
    return (f"For the following situation and question, select the desired option, by just returning the option number\n\n"
            f"<Situation>{state}</Situation>\n<Question>{question}</Question>\n\nOptions:\n\n{opts}\n\nAnswer: ")


def ids_for(tok, n):
    ids = []
    for c in LABELS[:n]:  # prompt ends with space, so no " c" variant needed
        e = tok.encode(c, add_special_tokens=False)
        assert len(e) == 1, f"label {c!r} is not single-token: {e}"
        ids.append(e[0])
    return ids


@torch.no_grad()
def predict_batch(prompts, ns, model, tok):
    enc = tok(prompts, return_tensors="pt", padding=True).to(model.device)
    last = model(**enc).logits[:, -1, :]  # [B, V]
    mn, b = max(ns), len(ns)
    idx = torch.zeros((b, mn), dtype=torch.long, device=last.device)
    for i, n in enumerate(ns):
        idx[i, :n] = torch.tensor(ids_for(tok, n), device=last.device)
    mask = torch.arange(mn, device=last.device)[None] < torch.tensor(ns, device=last.device)[:, None]
    scores = torch.where(mask, last.gather(1, idx), float("-inf"))
    return torch.softmax(scores, 1).tolist()


@spaces.GPU(duration=120)
def classify(mid, state, qjson):
    qs = json.loads(qjson)
    qs = [{"key": k, **v} for k, v in qs.items()] if isinstance(qs, dict) else qs
    prompts, ns, metas = [], [], []
    for q in qs:
        t = q.get("type", "choice")
        key = q.get("key", "q")
        if t == "noul":
            o = ["no", "yes"]
            meta = {"key": key, "type": t, "options": o}
            prompts.append(prompt(state, q.get("question", q.get("instructions", "")), o))
        else:
            o = q["options"]
            meta = {"key": key, "type": t, "options": o, "min": q.get("min", 0), "max": q.get("max", 100)}
            prompts.append(prompt(state, q.get("question", ""), o))
        ns.append(len(o))
        metas.append(meta)
    model, tok = get_model(mid)
    out = {}
    for pv, m in zip(predict_batch(prompts, ns, model, tok), metas):
        pv = [round(float(x), 6) for x in pv[:len(m["options"])]]
        probs = dict(zip(m["options"], pv))
        conf = max(pv)
        if m["type"] == "choice":
            out[m["key"]] = {"choice": m["options"][pv.index(conf)], "probabilities": probs, "confidence": conf}
        elif m["type"] == "noul":
            out[m["key"]] = {"noul": pv[1], "probabilities": probs, "confidence": conf}
        else:
            lo, hi = m["min"], m["max"]
            out[m["key"]] = {"score": round(sum(p * (lo + i * (hi - lo) / (len(pv) - 1)) for i, p in enumerate(pv)), 4),
                             "probabilities": probs, "confidence": conf}
    return out


EXAMPLE = json.dumps({
    "category": {"type": "choice", "question": "Which dept?", "options": ["billing", "technical", "sales"]},
    "urgency": {"type": "score", "question": "How urgent?", "options": ["low", "medium", "high"], "min": 0, "max": 100},
    "is_urgent": {"type": "noul", "question": "Urgent?"},
}, indent=2)

with gr.Blocks(title="QwenJev — Turn Context into Decisions") as demo:
    gr.Markdown("""
    # QwenJev — Turn Context into Decisions
    ### Make any LLM work as a decision model or zero-shot classifier

    Turn ordinary context into **structured decisions and probabilities**—without
    generating free-form text.

    1. Describe what is happening in **Context**.
    2. Define one or more decisions in **Questions**.
    3. Click **Evaluate decisions** to get typed, name-keyed results.
    """)
    with gr.Row():
        with gr.Column():
            mid = gr.Dropdown(
                MODELS, value=DEFAULT_MODEL, label="1. Model",
                info="The small Qwen3-0.6B model is the fastest default.",
                allow_custom_value=True)
            state = gr.Textbox(
                label="2. Context",
                info="Paste the facts or situation the model should judge.",
                lines=6, value="Stripe failing 3 days. Losing sales. ASAP.")
            qin = gr.Textbox(
                label="3. Questions (JSON)",
                info="Each named entry becomes one independent decision.",
                lines=18, value=EXAMPLE)
            with gr.Accordion("Question types", open=False):
                gr.Markdown("""
                - **choice** — select one item from `options`.
                - **score** — distribute probability over ordered `options` and
                  return their expected value between `min` and `max`.
                - **noul** — return a yes/no probability; no `options` required.

                The JSON key (for example, `urgency`) becomes the result key.
                """)
            btn = gr.Button("Evaluate decisions", variant="primary")
        with gr.Column():
            gr.Markdown("""
            ### Structured results
            Every answer is restricted to the requested choices.

            - **probabilities** show the model's relative preference among the
              allowed answers.
            - **confidence** is the highest option probability. It is a raw model
              score, not a calibrated guarantee of correctness.
            """)
            out = gr.JSON(label="Decisions and probabilities")
    btn.click(classify, [mid, state, qin], [out])

if __name__ == "__main__":
    demo.launch()