"""Minimal working example: cribl-decision structured readout.

Loads the LoRA adapter, builds a prompt in the training format,
extracts hidden-state logits for candidate tokens, and returns
a probability distribution.  No text is generated.
"""

import json
import torch
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration
from peft import PeftModel

# ── 0.  Triton workaround (needed on some systems) ──────────────
# Uncomment the two lines below if you hit a Triton / Python.h
# compilation error on first inference:
#
# from torch._native import triton_utils
# triton_utils.deregister_op_overrides()

# ── 1.  Load model ──────────────────────────────────────────────
BASE_ID = "Qwen/Qwen3.5-4B-Base"
BASE_REV = "1001bb4d826a52d1f399e183466143f4da7b741b"
ADAPTER_ID = "cribl-ai/cribl-decision-1.0"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device.type == "cuda" else torch.float32

tokenizer = AutoTokenizer.from_pretrained(BASE_ID, revision=BASE_REV)

# IMPORTANT: use Qwen3_5ForConditionalGeneration, not AutoModelForCausalLM.
# AutoModelForCausalLM resolves to Qwen3_5ForCausalLM which silently
# ignores the LoRA adapter in transformers >= 5.19.
base = Qwen3_5ForConditionalGeneration.from_pretrained(
    BASE_ID, revision=BASE_REV, torch_dtype=dtype,
)
model = PeftModel.from_pretrained(base, ADAPTER_ID, is_trainable=False).to(device)
model.eval()

# Split into language_model (transformer) and lm_head (projection).
base_model = model.get_base_model()
language_model = base_model.model.language_model
lm_head = base_model.lm_head

# ── 2.  Label table ────────────────────────────────────────────
# A-Z for up to 26 options.  For >26, continues with two-letter
# single-token labels (AA, AB, ...) up to 255 total.
LETTERS = "ABCDEFGHIJ"


def letter_token_ids(tok, n):
    """Token IDs for the first n candidate letters."""
    return [tok.encode(LETTERS[i], add_special_tokens=False)[0] for i in range(n)]


# ── 3.  Build prompt (training format) ─────────────────────────
def build_prompt(state, question, options):
    """Build a prompt in the exact format the model was trained on.

    Format:
        Context:
        <JSON or plain-text state>

        Question: <text>
        Options:
        (A) <option 1>
        (B) <option 2>
        ...
        Answer: (
    """
    if isinstance(state, dict):
        state_str = json.dumps(state, ensure_ascii=False, separators=(", ", ": "))
    else:
        state_str = str(state)

    lines = [f"Context:\n{state_str}", "",
             f"Question: {question}", "Options:"]
    for i, opt in enumerate(options):
        lines.append(f"({LETTERS[i]}) {opt}")
    lines.append("Answer: (")
    return "\n".join(lines)


# ── 4.  Score candidates ───────────────────────────────────────
@torch.inference_mode()
def score(prompt, n_options):
    """Run one forward pass, return {option_index: probability}."""
    ids = tokenizer.encode(prompt, add_special_tokens=False)
    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
    mask = torch.ones_like(input_ids)

    hidden = language_model(input_ids=input_ids, attention_mask=mask).last_hidden_state
    h = hidden[0, -1]  # hidden state at the "(" token

    cand_ids = torch.tensor(letter_token_ids(tokenizer, n_options),
                            dtype=torch.long, device=device)
    weights = lm_head.weight[cand_ids]
    logits = (weights @ h).float()
    probs = torch.softmax(logits, dim=-1).cpu().tolist()
    return probs


# ── 5.  Demo ────────────────────────────────────────────────────
if __name__ == "__main__":
    # Example: route a support ticket
    state = {"message": "My card was charged twice for the same purchase."}
    question = "Which department should handle this?"
    options = ["billing", "technical support", "sales"]

    prompt = build_prompt(state, question, options)
    print("=== Prompt ===")
    print(prompt)
    print()

    probs = score(prompt, len(options))
    print("=== Probabilities ===")
    for i, (opt, p) in enumerate(zip(options, probs)):
        print(f"  ({LETTERS[i]}) {opt:20s}  {p:.4f}")
    print(f"\n  Pick: ({LETTERS[probs.index(max(probs))]}) {options[probs.index(max(probs))]}")