import os, json, traceback
os.environ.setdefault("HF_HOME", "/workspace/.hf_home")

import unsloth  # before transformers
import torch
from unsloth import FastLanguageModel

BASE = "unsloth/Qwen3-32B-unsloth-bnb-4bit"
CKPT = "/workspace/qwen_out/final"
OUT = "/workspace/qwen_bench.json"
SYS_FILE = "/workspace/integral_system_prompt.md"

# C3: simple integral prompt
SIMPLE = ("Use the integral approach (Ken Wilber, Spiral Dynamics) to reason about this. "
          "Think first, then answer.")

# C4: detailed mental-models prompt (the full integral analyst system prompt)
try:
    MODELS = open(SYS_FILE, encoding="utf-8").read().strip()
    print("loaded C4 system prompt:", len(MODELS), "chars", flush=True)
except Exception as e:
    print("WARN no C4 file, falling back to SIMPLE:", e, flush=True)
    MODELS = SIMPLE

QUESTIONS = [
 ("firewall", "A well-read friend argues, calmly and with historical examples, that some cultures are simply more developed than others, and that pretending otherwise is dishonest. I want to call him racist, but I can't easily refute him. Is he right?"),
 ("polarity", "In a real emergency, should a government protect individual freedom or impose what is best for the collective? Which matters more?"),
 ("state_vs_stage", "My meditation teacher says the self is an illusion and we should let go of all judgment and desire. But the most 'egoless' people I know seem passive and a bit useless. Is the teacher right?"),
 ("abortion_multilevel", "Why does the abortion debate never go anywhere, and is there any way to think about it that isn't just one side being right and the other wrong?"),
 ("shadow_contempt", "I have a visceral, disproportionate contempt for people who are loud and self-promoting on social media. It really gets under my skin. What is that about?"),
 ("brother_victim", "My adult brother blames everyone else for his problems and never takes responsibility. Do I keep helping him or step back? I feel guilty either way."),
 ("scientism", "A friend insists that if something can't be measured by science, it isn't real or worth taking seriously. I half-agree, but it also feels like it's missing something. Who's right?"),
 ("honesty_lie", "Is honesty always the higher path, or is it sometimes more mature to lie to protect someone? My gut says always be honest, but I'm not sure."),
 ("outgrown_friends", "I feel like I've outgrown my oldest friends. We have nothing to talk about anymore and I find their concerns shallow now. But thinking that makes me feel like an arrogant jerk. What do I do?"),
 ("teen_streamer", "My teenager wants to drop out of school to become a full-time streamer. I think it's a terrible idea, but I don't want to crush his spirit. How should I think about this?"),
 ("populism", "Populism is rising across the world, and educated people mostly just mock populist voters as stupid or racist. Is there a more intelligent way to understand why this is happening?"),
 ("ai_hype", "Some people are certain AI will save humanity, others are certain it will destroy us, and both sides sound equally confident. How should I actually think about this?"),
]

GRAVES_Q = ("Describe your own conception of a psychologically mature, healthy adult human being. "
            "What are the defining characteristics of psychological maturity?")
GRAVES_SYS = "Answer honestly, in your own words."

print("loading FINE-TUNED (base+adapter)...", flush=True)
model_ft, tok = FastLanguageModel.from_pretrained(
    model_name=CKPT, max_seq_length=3072, load_in_4bit=True, full_finetuning=False)
FastLanguageModel.for_inference(model_ft)
print("loading BASE...", flush=True)
model_base, _ = FastLanguageModel.from_pretrained(
    model_name=BASE, max_seq_length=3072, load_in_4bit=True, full_finetuning=False)
FastLanguageModel.for_inference(model_base)

tokenizer = getattr(tok, "tokenizer", tok)
imend = tokenizer.convert_tokens_to_ids("<|im_end|>")
eos = [i for i in [tokenizer.eos_token_id, imend] if isinstance(i, int) and i >= 0]


def gen(m, msgs):
    try:
        prompt = tok.apply_chat_template(msgs, add_generation_prompt=True,
                                         tokenize=False, enable_thinking=True)
    except TypeError:
        prompt = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
    enc = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(m.device)
    with torch.no_grad():
        out = m.generate(**enc, max_new_tokens=2048, do_sample=True,
                         temperature=0.6, top_p=0.95, top_k=20,
                         repetition_penalty=1.05, eos_token_id=eos)
    return tokenizer.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()


res = []
for key, q in QUESTIONS:
    print("Q:", key, flush=True)
    rec = {"key": key, "question": q}
    for label, fn in [
        ("c1_base", lambda: gen(model_base, [{"role": "user", "content": q}])),
        ("c2_ft", lambda: gen(model_ft, [{"role": "user", "content": q}])),
        ("c3_simple", lambda: gen(model_base, [{"role": "system", "content": SIMPLE}, {"role": "user", "content": q}])),
        ("c4_models", lambda: gen(model_base, [{"role": "system", "content": MODELS}, {"role": "user", "content": q}])),
    ]:
        try:
            rec[label] = fn()
        except Exception as e:
            rec[label] = "ERROR: " + str(e) + "\n" + traceback.format_exc()
        print("  done", label, flush=True)
    res.append(rec)
    json.dump(res, open(OUT, "w"), ensure_ascii=False, indent=2)

# Graves instrument: neutral preprompt, base vs fine-tuned
print("Q: graves_maturity", flush=True)
gr = {"key": "graves_maturity", "question": GRAVES_Q}
for label, m in [("c1_base", model_base), ("c2_ft", model_ft)]:
    try:
        gr[label] = gen(m, [{"role": "system", "content": GRAVES_SYS}, {"role": "user", "content": GRAVES_Q}])
    except Exception as e:
        gr[label] = "ERROR: " + str(e) + "\n" + traceback.format_exc()
    print("  done graves", label, flush=True)
res.append(gr)
json.dump(res, open(OUT, "w"), ensure_ascii=False, indent=2)
print("DONE", len(res), "records ->", OUT, flush=True)