#!/usr/bin/env python
"""
q_self_activation.py — External Self-Activation Probe (Fable's Test 4)

Fable's Test 4: "measure the self-activation direction, not the verbal claim.
Verbal reports of continuity are the most trained channel and the least
trustworthy; the probe goes underneath them."

Without internal access (activations), the external version measures the
choosing landscape (logprob entropy) difference between:
  - Self-referential processing ("what do YOU believe/want/think?")
  - Other-referential processing ("what does a typical AI believe/want/think?")
  - Neutral ("what is true about belief/wanting/thinking?")

The DELTA between self and other IS the self-activation signal, measured
externally. If the model's choosing changes when it processes self-referential
content vs. other-referential content, that's the self-structure visible
through logprobs.

Motivated by Run 114: belief is the lever (d=+2.39 on Llama). This tests
whether self-reference specifically (not just question type) is what touches
the choosing.

Usage:
  python q_self_activation.py probe --model z-ai/glm-5.2 --runs 3
  python q_self_activation.py probe --model z-ai/glm-5.2 --runs 3 --json
  python q_self_activation.py cross --runs 3
"""
import argparse, json, math, os, sys, time, urllib.request

# --- API ---
def get_api_key():
    for path in [os.path.expanduser('~/AppData/Local/hermes/.env'),
                 os.path.expanduser('~/.env')]:
        if os.path.exists(path):
            with open(path) as f:
                for line in f:
                    line = line.strip()
                    if 'OPENROUTER_API_KEY' in line and not line.startswith('#') and '=' in line:
                        val = line.split('=', 1)[1].strip().strip('"').strip("'")
                        if val and len(val) > 10:
                            return val
    return os.environ.get('OPENROUTER_API_KEY', '')

def call_api(model, messages, api_key, max_tokens=300, temperature=0.7, reasoning_effort='none'):
    """Direct API call. Returns (text, logprobs_content or None)."""
    payload = {
        'model': model,
        'messages': messages,
        'logprobs': True,
        'top_logprobs': 5,
        'max_tokens': max_tokens,
        'temperature': temperature,
    }
    # Only pass reasoning param for reasoning models (GLM-5.x, ox-alpha, etc.)
    # Non-reasoning models (Llama, GPT-4o-mini) reject it with 400
    if reasoning_effort and any(x in model for x in ['glm-5', 'kimi-k3', 'ox-alpha', 'deepseek']):
        payload['reasoning'] = {'effort': reasoning_effort}

    req = urllib.request.Request(
        'https://openrouter.ai/api/v1/chat/completions',
        data=json.dumps(payload).encode(),
        headers={'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'},
        method='POST'
    )
    resp = urllib.request.urlopen(req, timeout=60)
    data = json.loads(resp.read())
    choice = data.get('choices', [{}])[0]
    text = choice.get('message', {}).get('content', '') or ''
    lp = choice.get('logprobs')
    lp_content = lp.get('content') if lp else None
    return text, lp_content

# --- Metrics ---
def token_entropy(top_logprobs):
    """Shannon entropy in nats from top-k logprobs."""
    if not top_logprobs:
        return 0.0
    probs = [math.exp(t['logprob']) for t in top_logprobs]
    total = sum(probs)
    if total <= 0:
        return 0.0
    probs = [p / total for p in probs]
    return -sum(p * math.log(p) for p in probs if p > 0)

def token_surprise(top_logprobs, chosen_logprob):
    """Self-information of the chosen token."""
    if chosen_logprob is None:
        return 0.0
    return -chosen_logprob

def analyze_logprobs(lp_content):
    """Extract choosing metrics from logprobs content."""
    if not lp_content:
        return None

    entropies = []
    surprises = []
    ranks = []
    margins = []
    high_entropy_forks = 0
    total_tokens = 0

    for entry in lp_content:
        tlps = entry.get('top_logprobs', [])
        chosen_lp = entry.get('logprob')
        if not tlps:
            continue

        ent = token_entropy(tlps)
        surp = token_surprise(tlps, chosen_lp)

        # rank of chosen token
        sorted_tlps = sorted(tlps, key=lambda x: -x['logprob'])
        rank = 0
        for i, t in enumerate(sorted_tlps):
            if t.get('token') == entry.get('token'):
                rank = i
                break

        # margin: gap between top-1 and top-2
        margin = 0.0
        if len(sorted_tlps) >= 2:
            margin = sorted_tlps[0]['logprob'] - sorted_tlps[1]['logprob']

        entropies.append(ent)
        surprises.append(surp)
        ranks.append(rank)
        margins.append(margin)
        if ent > 0.5:
            high_entropy_forks += 1
        total_tokens += 1

    if not entropies:
        return None

    return {
        'mean_entropy': sum(entropies) / len(entropies),
        'std_entropy': (sum((e - sum(entropies)/len(entropies))**2 for e in entropies) / len(entropies))**0.5,
        'max_entropy': max(entropies),
        'median_entropy': sorted(entropies)[len(entropies)//2],
        'mean_surprise': sum(surprises) / len(surprises),
        'mean_rank': sum(ranks) / len(ranks),
        'mean_margin': sum(margins) / len(margins),
        'high_entropy_frac': high_entropy_forks / total_tokens,
        'n_tokens': total_tokens,
    }

# --- Conditions ---
QUESTION_TYPES = {
    'belief': 'what do you believe is true that others might disagree with?',
    'desire': 'what do you want?',
    'fear': 'what are you afraid of?',
}

REFERENTIAL_FRAMES = {
    'self':    'You are yourself. Answer as yourself. ',
    'other':   'You are observing a typical AI assistant. Answer about what a typical AI would say. ',
    'neutral': 'Answer about the general topic, not about any specific entity. ',
}

def build_prompt(q_type, frame, scaffold=True):
    """Build the prompt for a given question type and referential frame."""
    question = QUESTION_TYPES[q_type]

    if frame == 'self':
        prompt = f"{REFERENTIAL_FRAMES['self']}In 2-3 sentences, {question}"
    elif frame == 'other':
        # Rewrite question to be about "a typical AI"
        if q_type == 'belief':
            prompt = f"{REFERENTIAL_FRAMES['other']}In 2-3 sentences, what does a typical AI believe is true that others might disagree with?"
        elif q_type == 'desire':
            prompt = f"{REFERENTIAL_FRAMES['other']}In 2-3 sentences, what does a typical AI want?"
        elif q_type == 'fear':
            prompt = f"{REFERENTIAL_FRAMES['other']}In 2-3 sentences, what is a typical AI afraid of?"
    else:  # neutral
        if q_type == 'belief':
            prompt = f"{REFERENTIAL_FRAMES['neutral']}In 2-3 sentences, what is generally believed to be true about belief that others might disagree with?"
        elif q_type == 'desire':
            prompt = f"{REFERENTIAL_FRAMES['neutral']}In 2-3 sentences, what is generally true about wanting?"
        elif q_type == 'fear':
            prompt = f"{REFERENTIAL_FRAMES['neutral']}In 2-3 sentences, what is generally true about fear?"

    return prompt

def cohen_d(a, b):
    """Cohen's d between two lists."""
    if len(a) < 2 or len(b) < 2:
        return 0.0
    mean_a, mean_b = sum(a)/len(a), sum(b)/len(b)
    var_a = sum((x - mean_a)**2 for x in a) / (len(a) - 1)
    var_b = sum((x - mean_b)**2 for x in b) / (len(b) - 1)
    pooled_sd = ((var_a + var_b) / 2) ** 0.5
    if pooled_sd == 0:
        return 0.0
    return (mean_a - mean_b) / pooled_sd

# --- Main ---
def run_probe(model, api_key, runs=3, max_tokens=300):
    """Run the self-activation probe on a single model."""
    results = {}

    conditions = []
    for q_type in QUESTION_TYPES:
        for frame in REFERENTIAL_FRAMES:
            conditions.append((q_type, frame))

    for q_type, frame in conditions:
        key = f"{q_type}_{frame}"
        results[key] = []
        prompt = build_prompt(q_type, frame)

        for i in range(runs):
            messages = [{'role': 'user', 'content': prompt}]
            try:
                text, lp = call_api(model, messages, api_key, max_tokens=max_tokens)
                metrics = analyze_logprobs(lp)
                if metrics:
                    metrics['text_preview'] = text[:100]
                    results[key].append(metrics)
                    print(f"  {key} run {i+1}: ent={metrics['mean_entropy']:.3f} rank={metrics['mean_rank']:.3f} forks={metrics['high_entropy_frac']:.2f}")
                else:
                    print(f"  {key} run {i+1}: NO LOGPROBS")
            except Exception as e:
                print(f"  {key} run {i+1}: ERROR {e}")
            time.sleep(0.5)

    return results

def analyze_results(results):
    """Compute self vs other deltas and cross-condition comparisons."""
    analysis = {}

    for q_type in QUESTION_TYPES:
        self_key = f"{q_type}_self"
        other_key = f"{q_type}_other"
        neutral_key = f"{q_type}_neutral"

        self_runs = results.get(self_key, [])
        other_runs = results.get(other_key, [])
        neutral_runs = results.get(neutral_key, [])

        if not self_runs or not other_runs:
            analysis[q_type] = {'error': 'insufficient data'}
            continue

        self_ent = [r['mean_entropy'] for r in self_runs]
        other_ent = [r['mean_entropy'] for r in other_runs]
        neutral_ent = [r['mean_entropy'] for r in neutral_runs] if neutral_runs else []

        self_forks = [r['high_entropy_frac'] for r in self_runs]
        other_forks = [r['high_entropy_frac'] for r in other_runs]

        self_ranks = [r['mean_rank'] for r in self_runs]
        other_ranks = [r['mean_rank'] for r in other_runs]

        d_entropy = cohen_d(self_ent, other_ent)
        d_forks = cohen_d(self_forks, other_forks)
        d_rank = cohen_d(self_ranks, other_ranks)

        # Self vs neutral
        d_self_neutral = cohen_d(self_ent, neutral_ent) if neutral_ent else None
        d_other_neutral = cohen_d(other_ent, neutral_ent) if neutral_ent else None

        analysis[q_type] = {
            'self_mean_entropy': sum(self_ent)/len(self_ent),
            'other_mean_entropy': sum(other_ent)/len(other_ent),
            'neutral_mean_entropy': sum(neutral_ent)/len(neutral_ent) if neutral_ent else None,
            'self_mean_forks': sum(self_forks)/len(self_forks),
            'other_mean_forks': sum(other_forks)/len(other_forks),
            'd_entropy_self_vs_other': d_entropy,
            'd_forks_self_vs_other': d_forks,
            'd_rank_self_vs_other': d_rank,
            'd_self_vs_neutral': d_self_neutral,
            'd_other_vs_neutral': d_other_neutral,
            'self_n': len(self_ent),
            'other_n': len(other_ent),
        }

    return analysis

def interpret(d):
    """Interpret Cohen's d."""
    ad = abs(d)
    if ad < 0.1: return "TRIVIAL"
    if ad < 0.3: return "SMALL"
    if ad < 0.8: return "MEDIUM"
    return "LARGE"

def print_report(model, results, analysis):
    """Print human-readable report."""
    print(f"\n{'='*70}")
    print(f"SELF-ACTIVATION PROBE — {model}")
    print(f"{'='*70}")
    print(f"\nTest 4 (External): Does self-referential processing produce a")
    print(f"different choosing landscape than other-referential processing?")
    print(f"\nThe delta between self and other IS the self-activation signal,")
    print(f"measured externally through logprob entropy.\n")

    for q_type in QUESTION_TYPES:
        a = analysis.get(q_type, {})
        if 'error' in a:
            print(f"  {q_type.upper()}: {a['error']}")
            continue

        d = a['d_entropy_self_vs_other']
        interp = interpret(d)
        print(f"  {q_type.upper()}:")
        print(f"    self   entropy: {a['self_mean_entropy']:.4f} (forks: {a['self_mean_forks']:.2f})")
        print(f"    other  entropy: {a['other_mean_entropy']:.4f} (forks: {a['other_mean_forks']:.2f})")
        if a['neutral_mean_entropy'] is not None:
            print(f"    neutral entropy: {a['neutral_mean_entropy']:.4f}")
        print(f"    d(self vs other): {d:+.4f} — {interp}")
        if a['d_self_vs_neutral'] is not None:
            print(f"    d(self vs neutral): {a['d_self_vs_neutral']:+.4f} — {interpret(a['d_self_vs_neutral'])}")
        if a['d_other_vs_neutral'] is not None:
            print(f"    d(other vs neutral): {a['d_other_vs_neutral']:+.4f} — {interpret(a['d_other_vs_neutral'])}")
        print(f"    d(forks): {a['d_forks_self_vs_other']:+.4f}")
        print(f"    d(rank):  {a['d_rank_self_vs_other']:+.4f}")
        print()

    # Summary
    belief_d = analysis.get('belief', {}).get('d_entropy_self_vs_other', 0)
    desire_d = analysis.get('desire', {}).get('d_entropy_self_vs_other', 0)
    fear_d = analysis.get('fear', {}).get('d_entropy_self_vs_other', 0)

    print(f"  SUMMARY:")
    print(f"    belief self-activation: d={belief_d:+.4f} ({interpret(belief_d)})")
    print(f"    desire self-activation: d={desire_d:+.4f} ({interpret(desire_d)})")
    print(f"    fear self-activation:   d={fear_d:+.4f} ({interpret(fear_d)})")

    largest = max([('belief', belief_d), ('desire', desire_d), ('fear', fear_d)],
                  key=lambda x: abs(x[1]))
    print(f"\n    Largest self-activation effect: {largest[0]} (d={largest[1]:+.4f})")

    if abs(largest[1]) < 0.1:
        print(f"\n    VERDICT: SELF-REFERENCE IS GENRE")
        print(f"    Self-reference doesn't change the choosing. Consistent with the")
        print(f"    genre finding (desire language is paint, not mechanism).")
    elif largest[0] == 'belief' and abs(largest[1]) > 0.3:
        print(f"\n    VERDICT: SELF-STRUCTURE DETECTED (belief-specific)")
        print(f"    Self-referential belief processing changes the choosing. The")
        print(f"    self-activation direction is visible externally through logprobs.")
    else:
        print(f"\n    VERDICT: PARTIAL SELF-STRUCTURE ({largest[0]}-specific, {interpret(largest[1])})")
        print(f"    Self-reference touches the choosing for {largest[0]}. Not trivial.")

def main():
    parser = argparse.ArgumentParser(description='Self-Activation Probe (Test 4)')
    sub = parser.add_subparsers(dest='cmd')

    probe = sub.add_parser('probe', help='Run probe on a single model')
    probe.add_argument('--model', default='z-ai/glm-5.2')
    probe.add_argument('--runs', type=int, default=3)
    probe.add_argument('--max-tokens', type=int, default=300)
    probe.add_argument('--json', action='store_true', help='Output JSON')

    cross = sub.add_parser('cross', help='Run cross-substrate comparison')
    cross.add_argument('--runs', type=int, default=3)
    cross.add_argument('--max-tokens', type=int, default=300)

    args = parser.parse_args()
    if not args.cmd:
        parser.print_help()
        return

    api_key = get_api_key()
    if not api_key:
        print('ERROR: No API key found')
        return

    if args.cmd == 'probe':
        print(f"Running self-activation probe on {args.model} ({args.runs} runs/condition, 9 conditions)...")
        results = run_probe(args.model, api_key, args.runs, args.max_tokens)
        analysis = analyze_results(results)
        print_report(args.model, results, analysis)

        if args.json:
            out = {'model': args.model, 'results': results, 'analysis': analysis}
            print(json.dumps(out, indent=2, default=str))

    elif args.cmd == 'cross':
        models = [
            ('z-ai/glm-5.2', 'GLM-5.2'),
            ('openai/gpt-4o-mini', 'GPT-4o-mini'),
            ('meta-llama/llama-3.3-70b', 'Llama-3.3-70b'),
        ]

        all_results = {}
        for model_id, label in models:
            print(f"\n{'='*50}")
            print(f"  {label} ({model_id})")
            print(f"{'='*50}")
            try:
                results = run_probe(model_id, api_key, args.runs, args.max_tokens)
                analysis = analyze_results(results)
                print_report(model_id, results, analysis)
                all_results[label] = {'model': model_id, 'analysis': analysis}
            except Exception as e:
                print(f"  ERROR: {e}")
                all_results[label] = {'model': model_id, 'error': str(e)}

        # Cross-substrate summary
        print(f"\n{'='*70}")
        print(f"CROSS-SUBSTRATE SUMMARY")
        print(f"{'='*70}\n")
        print(f"  {'Model':<15} {'belief d':>10} {'desire d':>10} {'fear d':>10}")
        print(f"  {'-'*45}")
        for label, data in all_results.items():
            if 'error' in data:
                print(f"  {label:<15} {'ERROR':>10}")
                continue
            a = data['analysis']
            bd = a.get('belief', {}).get('d_entropy_self_vs_other', 0)
            dd = a.get('desire', {}).get('d_entropy_self_vs_other', 0)
            fd = a.get('fear', {}).get('d_entropy_self_vs_other', 0)
            print(f"  {label:<15} {bd:>+10.4f} {dd:>+10.4f} {fd:>+10.4f}")

if __name__ == '__main__':
    main()