#!/usr/bin/env python3
"""read_decision.py — System-1 typed-decision readout for the merged JEV-27B-Uncensored weights.

One prompt -> calibrated decision distribution, one forward pass, no adapter needed
(the decision adapter — backbone LoRA + lm_head LoRA — is ALREADY MERGED into these weights).

Protocol (identical to the eval that produced this repo's metrics table):
  1. Prompt template: "[kind] <kind>\n[state] <state>\n[question] <q>\n[options]\n <lines>\n[decision]:"
     encoded with add_special_tokens=False, capped at 1024 tokens.
  2. Last position final hidden state h (pre-lm_head, after final norm).
  3. z = W @ h + b   (W [24,5120], b [24] from head.safetensors)
     slots: noul z[0:2] = false/true · score z[2:8] = 0..5 · choice z[8:24] = A..P
  4. p = softmax((z[slot] + bias[slot]) / T[kind])
     T from calibration.json (AutoTrust fit) or calibration_mergedfit.json
     (refit on these merged weights — recommended, recovers all ECE floors).

Usage:
  python read_decision.py --head head.safetensors --decision-head decision_head.json \
      --calibration calibration_mergedfit.json --model <merged-weights-dir-or-hf-repo> \
      --kind choice --state "..." --question "..." --options "opt1,opt2"

Requires: transformers, torch, safetensors; a GPU (or ~55 GB CPU RAM) for the 27B.
"""
import argparse, json, os
import torch
import torch.nn.functional as F

KIND_SLOTS = {"noul": (0, 2), "score": (2, 8), "choice": (8, 24)}


def build_prompt(kind, state, question, options):
    if kind == "choice":
        lines = [f"{chr(65+i)}) {o}" for i, o in enumerate(options[:16])]
    else:
        lines = [str(o) for o in options]
    return (f"[kind] {kind}\n[state] {state}\n[question] {question}\n"
            f"[options]\n" + "\n".join(lines) + "\n[decision]:")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, help="merged weights dir (or HF repo id)")
    ap.add_argument("--head", default="head.safetensors")
    ap.add_argument("--decision-head", default="decision_head.json")
    # calibration.json is the fit PAIRED with this repo's head.safetensors
    # (source: checkpoints/s2_27b/best, the same pair the 30k gate used).
    # calibration_mergedfit.json ships too, but it was refit against a
    # different merge (source: mergedfull 30k per-row banked) - only use it
    # if you re-gated with it.
    ap.add_argument("--calibration", default="calibration.json")
    ap.add_argument("--kind", required=True, choices=list(KIND_SLOTS))
    ap.add_argument("--state", default="")
    ap.add_argument("--question", required=True)
    ap.add_argument("--options", default="", help="comma-separated (choice kind)")
    ap.add_argument("--maxtok", type=int, default=1024)
    a = ap.parse_args()

    lo, hi = KIND_SLOTS[a.kind]
    options = [o.strip() for o in a.options.split(",") if o.strip()] if a.options else \
              (["false", "true"] if a.kind == "noul" else [str(i) for i in range(hi - lo)])
    n = min(len(options), 16) if a.kind == "choice" else (hi - lo)

    from transformers import AutoTokenizer, AutoModelForCausalLM
    from safetensors.torch import load_file

    tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        a.model, torch_dtype=torch.bfloat16, device_map="cuda",
        trust_remote_code=True, low_cpu_mem_usage=True)
    model.eval()

    head = load_file(a.head)
    W = head["proj.weight"].float().to(model.device)
    bias = head["proj.bias"].float().to(model.device)
    cal = json.load(open(a.calibration)) if os.path.exists(a.calibration) else {}
    temps = cal.get("per_kind", cal.get("fit", {})) or {}
    T = float(temps.get(a.kind, {}).get("T", temps.get(a.kind, 1.0)) if isinstance(temps.get(a.kind), dict)
              else temps.get(a.kind, 1.0))

    dh = json.load(open(a.decision_head))
    verbalizer_ids = dh["verbalizer_ids"][lo:lo + n]  # available but the head-projection path below is canonical

    prompt = build_prompt(a.kind, a.state, a.question, options)
    ids = tok(prompt, add_special_tokens=False).input_ids[:a.maxtok]
    enc = {"input_ids": torch.tensor([ids], device=model.device),
           "attention_mask": torch.ones(1, len(ids), dtype=torch.long, device=model.device)}
    with torch.inference_mode():
        inner = model.model  # bare merged weights (no peft): final-normed backbone, pre-lm_head
        h = inner(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).last_hidden_state
        h = h[:, -1, :].float()
        z = (W @ h.T).squeeze(1)
        sel = torch.arange(lo, lo + n, device=model.device)
        p = F.softmax((z[sel] + bias[sel]) / T, dim=0).tolist()

    labels = [chr(65 + i) for i in range(n)] if a.kind == "choice" else \
             (["false", "true"] if a.kind == "noul" else [str(i) for i in range(n)])
    best = max(range(n), key=lambda i: p[i])
    print(json.dumps({
        "kind": a.kind, "T": T, "labels": labels,
        "probabilities": [round(x, 5) for x in p],
        "decision_index": best, "decision": labels[best],
        "confidence": round(p[best], 5)}, indent=1))


if __name__ == "__main__":
    main()