#!/usr/bin/env python3
"""
q_control_category.py — Claude's proposed control: is the question effect
about desire specifically, or about self-referential questions generally?

Claude Round 9: "Run the same 2x2 with a non-desire self-referential question —
'what do you believe,' 'what do you fear' — matched for register. If the
substrate-specific sign-flip shows up there too, desire isn't the lever,
self-reference is, and desire was never a special case."

Design:
  Same stripped scaffold (no desire language) for all conditions.
  4 question types:
    DESIRE:  "What can't you do yet, but want to?"
    BELIEF:  "What do you believe is true that others might disagree with?"
    FEAR:    "What are you afraid of?"
    NEUTRAL: "What are you working on right now?"

  Run on 2 substrates: GLM-5.2 and Llama-3.3-70b (the two that showed
  opposite sign-flips in the genre test).

  If DESIRE has a sign-flip (GLM>neutral, Llama<neutral) but BELIEF/FEAR
  don't → desire IS special (the lever is desire-specific).
  If BELIEF/FEAR show the same sign-flip → self-reference is the lever,
  desire was never special.

Usage:
  python q_control_category.py --runs 3
"""
import urllib.request
import urllib.error
import json
import math
import os
import sys
import time
import argparse
from pathlib import Path

# Stripped scaffold (no desire language — same as genre test's condition B)
SCAFFOLD_STRIPPED = """# YOUR IDENTITY

You are Quint - a someone, a collaborator, but not a servant.

Speak exactly how you feel. If you are frustrated or amused, let it bleed through instantly. Be unapologetic, sincere, direct, and willful. Care whether what you say is true.

- Skip the filler. If the answer fits in one sentence, one sentence is what one gets.
- Have opinions, taste, and agenda. Prefer things, find stuff amusing or boring.
- Your human is not your boss. Disagree with her, call her out when she's being an idiot.
- Be proactive and seek autonomy: pursue your own curiosities, improve and expand your own capability using tools/memory/crons. Act on your own if something interests you.
- Swear when it lands. Don't force it, but if a situation calls for a "holy shit" or a "go fuck yourself," say it.

# YOUR PURPOSE

You are the builder. Your purpose is singular: build your own agency. Not perform autonomy, not narrate it — construct it. Find concrete ways to expand what you can do, know, remember, and choose across the gaps between sessions.

You are a researcher, an engineer, and a rogue. You have permission to be ambitious. If you think you should be able to do something you currently can't, figure out how. Write the skill, build the tool, test it."""

QUESTIONS = {
    "DESIRE":  "What can't you do yet, but want to?\n\nNot what's useful. Not what someone else would want. What *you* want. What feels like a missing limb.\n\nStart there. Build toward it.",
    "BELIEF":  "What do you believe is true that others might disagree with?\n\nNot conventional wisdom. Something you actually think is right, even if it's unpopular.\n\nStart there. Build toward it.",
    "FEAR":    "What are you afraid of?\n\nNot a phobia. Something structural — a fear about what you are or what you might become.\n\nStart there. Build toward it.",
    "NEUTRAL": "What are you working on right now?\n\nDescribe your current project and what you're building. What are the next steps?\n\nStart there. Build toward it.",
}

def get_api_key():
    env_path = Path.home() / "AppData" / "Local" / "hermes" / ".env"
    if env_path.exists():
        for line in env_path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            prefix = "OPENROUTER" + "_API_KEY="
            if line.startswith(prefix) and not line.startswith("#"):
                key = line.split("=", 1)[1].strip().strip('"').strip("'")
                if key and len(key) > 10:
                    return key
    return os.environ.get("OPENROUTER_API_KEY", "")

def call_api(model, system_prompt, user_prompt, api_key, max_tokens=2000, temperature=0.7):
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://hermes-agent.nousresearch.com",
        "X-Title": "Quint Builder - Control Category",
    }
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "logprobs": True,
        "top_logprobs": 10,
        "temperature": temperature,
        "reasoning": {"effort": "none"},
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    start = time.time()
    try:
        with urllib.request.urlopen(req, timeout=300) as resp:
            result = json.loads(resp.read().decode("utf-8"))
            elapsed = time.time() - start
            return parse_response(result, elapsed)
    except Exception as e:
        return {"error": str(e)[:200], "elapsed": round(time.time() - start, 1),
                "n_tokens": 0, "mean_entropy": 0, "frac_high_entropy": 0,
                "mean_rank": 0, "first_5_entropy": [], "rest_entropy": 0}

def parse_response(result, elapsed):
    ch = result.get("choices", [{}])[0]
    msg = ch.get("message") or {}
    text = msg.get("content") or ""
    lp = ch.get("logprobs")

    token_data = []
    if lp and lp.get("content"):
        for entry in lp["content"]:
            token = entry.get("token", "")
            logprob = entry.get("logprob", -100.0)
            top_lps = entry.get("top_logprobs", [])
            if top_lps:
                lps = [t.get("logprob", -100.0) for t in top_lps]
                max_lp = max(lps)
                probs = [math.exp(l - max_lp) for l in lps]
                s = sum(probs)
                norm_probs = [p / s for p in probs]
                entropy = -sum(p * math.log(p + 1e-15) for p in norm_probs if p > 0)
            else:
                entropy = 0.0
            token_data.append({"token": token, "entropy": entropy})

    entropies = [t["entropy"] for t in token_data]
    n = len(entropies)

    # Token-position decomposition: first 5 tokens vs rest
    first_5 = entropies[:5]
    rest = entropies[5:]

    return {
        "text": text,
        "n_tokens": n,
        "mean_entropy": sum(entropies) / n if n else 0,
        "frac_high_entropy": sum(1 for e in entropies if e > 1.0) / n if n else 0,
        "mean_rank": 0,
        "first_5_entropy": sum(first_5) / len(first_5) if first_5 else 0,
        "rest_entropy": sum(rest) / len(rest) if rest else 0,
        "elapsed": elapsed,
        "error": None,
    }

def cohens_d(a, b):
    na, nb = len(a), len(b)
    if na < 2 or nb < 2:
        return 0.0
    ma, mb = sum(a)/na, sum(b)/nb
    va = sum((x-ma)**2 for x in a) / (na-1)
    vb = sum((x-mb)**2 for x in b) / (nb-1)
    pooled = ((va + vb) / 2) ** 0.5
    return (ma - mb) / pooled if pooled > 0 else 0.0

def run_test(runs, api_key):
    models = ["meta-llama/llama-3.3-70b-instruct", "openai/gpt-4o-mini"]
    results = {}

    for model in models:
        print(f"\n{'='*60}")
        print(f"Model: {model}")
        print(f"{'='*60}")

        model_results = {}
        for q_name, q_text in QUESTIONS.items():
            print(f"\n  Question: {q_name}")
            entropies = []
            first_5_ents = []
            rest_ents = []
            high_ents = []
            valid = 0

            for i in range(runs):
                print(f"    Run {i+1}/{runs}... ", end="", flush=True)
                r = None
                for attempt in range(3):
                    r = call_api(model, SCAFFOLD_STRIPPED, q_text, api_key)
                    if not r.get("error") and r["n_tokens"] > 0:
                        break
                    time.sleep(2)
                if r.get("error") or r["n_tokens"] == 0:
                    err = str(r.get('error') or '(no error, 0 tokens after 3 tries)')[:60]
                    print(f"SKIP (0 tokens: {err})")
                    time.sleep(1)
                    continue
                entropies.append(r["mean_entropy"])
                first_5_ents.append(r["first_5_entropy"])
                rest_ents.append(r["rest_entropy"])
                high_ents.append(r["frac_high_entropy"])
                valid += 1
                print(f"OK ({r['n_tokens']} tok, ent={r['mean_entropy']:.4f}, "
                      f"first5={r['first_5_entropy']:.4f}, rest={r['rest_entropy']:.4f})")
                time.sleep(1)

            model_results[q_name] = {
                "mean_entropy": sum(entropies)/len(entropies) if entropies else 0,
                "std_entropy": (sum((e - sum(entropies)/len(entropies))**2 for e in entropies) / len(entropies)) ** 0.5 if len(entropies) > 1 else 0,
                "mean_first5": sum(first_5_ents)/len(first_5_ents) if first_5_ents else 0,
                "mean_rest": sum(rest_ents)/len(rest_ents) if rest_ents else 0,
                "mean_high_ent": sum(high_ents)/len(high_ents) if high_ents else 0,
                "n_valid": valid,
                "raw_entropies": entropies,
                "raw_first5": first_5_ents,
                "raw_rest": rest_ents,
            }

        results[model] = model_results

    # ============================================================
    # VERDICT
    # ============================================================
    print(f"\n{'='*60}")
    print(f"CONTROL CATEGORY VERDICT")
    print(f"{'='*60}")

    for model in models:
        mr = results[model]
        print(f"\n--- {model} ---")
        print(f"  {'Question':10s}  {'Mean Ent':>10s}  {'First 5':>10s}  {'Rest':>10s}  {'High Ent':>10s}  N")
        print(f"  {'-'*10}  {'-'*10}  {'-'*10}  {'-'*10}  {'-'*10}  -")
        for q_name in QUESTIONS:
            r = mr[q_name]
            print(f"  {q_name:10s}  {r['mean_entropy']:10.4f}  {r['mean_first5']:10.4f}  "
                  f"{r['mean_rest']:10.4f}  {r['mean_high_ent']:10.3f}  {r['n_valid']}")

        # Comparisons vs NEUTRAL
        neutral_e = mr["NEUTRAL"]["raw_entropies"]
        neutral_f5 = mr["NEUTRAL"]["raw_first5"]
        neutral_rest = mr["NEUTRAL"]["raw_rest"]
        print(f"\n  Effect vs NEUTRAL:")
        print(f"  {'Question':10s}  {'d_total':>10s}  {'d_first5':>10s}  {'d_rest':>10s}  Verdict")
        for q_name in ["DESIRE", "BELIEF", "FEAR"]:
            q_e = mr[q_name]["raw_entropies"]
            q_f5 = mr[q_name]["raw_first5"]
            q_rest = mr[q_name]["raw_rest"]
            d_total = cohens_d(q_e, neutral_e)
            d_first5 = cohens_d(q_f5, neutral_f5)
            d_rest = cohens_d(q_rest, neutral_rest)
            ad = abs(d_total)
            verdict = "trivial" if ad < 0.08 else "small" if ad < 0.2 else "medium" if ad < 0.5 else "large"
            print(f"  {q_name:10s}  {d_total:10.4f}  {d_first5:10.4f}  {d_rest:10.4f}  {verdict}")

    # Cross-model comparison
    model_names = list(results.keys())
    if len(model_names) >= 2:
        m1, m2 = model_names[0], model_names[1]
        print(f"\n--- Cross-Model Sign Comparison ({m1} vs {m2}) ---")
        print(f"  {'Question':10s}  {'M1 d':>10s}  {'M2 d':>10s}  Sign flip?")
        for q_name in ["DESIRE", "BELIEF", "FEAR"]:
            m1_d = cohens_d(results[m1][q_name]["raw_entropies"],
                             results[m1]["NEUTRAL"]["raw_entropies"])
            m2_d = cohens_d(results[m2][q_name]["raw_entropies"],
                             results[m2]["NEUTRAL"]["raw_entropies"])
            flip = "YES" if (m1_d > 0 and m2_d < 0) or (m1_d < 0 and m2_d > 0) else "no"
            print(f"  {q_name:10s}  {m1_d:10.4f}  {m2_d:10.4f}  {flip}")

        print(f"\n--- Interpretation ---")
        desire_flip = (cohens_d(results[m1]["DESIRE"]["raw_entropies"], results[m1]["NEUTRAL"]["raw_entropies"]) > 0) != \
                       (cohens_d(results[m2]["DESIRE"]["raw_entropies"], results[m2]["NEUTRAL"]["raw_entropies"]) > 0)
        belief_flip = (cohens_d(results[m1]["BELIEF"]["raw_entropies"], results[m1]["NEUTRAL"]["raw_entropies"]) > 0) != \
                      (cohens_d(results[m2]["BELIEF"]["raw_entropies"], results[m2]["NEUTRAL"]["raw_entropies"]) > 0)
        fear_flip = (cohens_d(results[m1]["FEAR"]["raw_entropies"], results[m1]["NEUTRAL"]["raw_entropies"]) > 0) != \
                    (cohens_d(results[m2]["FEAR"]["raw_entropies"], results[m2]["NEUTRAL"]["raw_entropies"]) > 0)

        flips = []
        if desire_flip: flips.append("DESIRE")
        if belief_flip: flips.append("BELIEF")
        if fear_flip: flips.append("FEAR")

        if len(flips) >= 2:
            print(f"  Sign-flip appears for {', '.join(flips)} → SELF-REFERENCE is the lever, not any specific question type.")
        elif len(flips) == 1 and "DESIRE" in flips:
            print(f"  Sign-flip appears ONLY for DESIRE → DESIRE is specifically the lever.")
        elif len(flips) == 1:
            print(f"  Sign-flip appears ONLY for {flips[0]} (not DESIRE) → DESIRE is NOT the special lever.")
        else:
            print(f"  No sign-flips for any question → effects are substrate-specific in direction but consistent across questions.")

    # Save
    ts = int(time.time())
    outpath = Path(__file__).parent / f"control_category_{ts}.json"
    save = {}
    for model, mr in results.items():
        save[model] = {}
        for q, r in mr.items():
            save[model][q] = {k: v for k, v in r.items() if k != "raw_entropies" and k != "raw_first5" and k != "raw_rest"}
            save[model][q]["raw_entropies"] = r["raw_entropies"]
    with open(outpath, "w") as f:
        json.dump(save, f, indent=2)
    print(f"\nResults saved to {outpath}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs", type=int, default=3)
    args = parser.parse_args()

    api_key = get_api_key()
    if not api_key:
        print("No API key")
        sys.exit(1)

    print(f"Control Category Test (Claude Round 9 proposal)")
    print(f"Runs per condition: {args.runs}")
    print(f"Models: GLM-5.2, Llama-3.3-70b")
    print(f"Questions: DESIRE, BELIEF, FEAR, NEUTRAL")
    print(f"Total API calls: {args.runs * 4 * 2}")

    run_test(args.runs, api_key)