#!/usr/bin/env python3
"""
Project Norn V18: Stratified Head-to-Head Ablation Test
Osakra Research — Academic Reference Implementation

Compares 3 Model Configurations on a Balanced 15-Pillar Stratified Evaluation Set:
1. Base Qwen (Raw Backbone without Cognitive Scaffolding)
2. Merged Adapter Checkpoint (norn-v17 fused weights)
3. Full Norn V18 Cognitive System (norn-v18)
"""

import os
import sys
import time
import json
import ast
import re
import urllib.request
from collections import defaultdict

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
if hasattr(sys.stderr, "reconfigure"):
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

OLLAMA_API_CHAT = "http://localhost:11434/api/chat"

MODELS_TO_TEST = [
    {"name": "Base Qwen 9B (Raw)", "ollama_id": "qwen-base-raw"},
    {"name": "Merged Adapter (V17 Fused)", "ollama_id": "norn-v17"},
    {"name": "Norn V18 (Full Cognitive)", "ollama_id": "norn-v18"}
]

script_dir = os.path.dirname(os.path.abspath(__file__))
suite_path = os.path.join(script_dir, "benchmark_1000_suite.json")
output_report_path = os.path.join(script_dir, "quick_ablation_results.json")

def extract_numbers(text):
    if not text:
        return []
    cleaned_text = re.sub(r'(?<=\d),(?=\d)', '', text)
    m_box = re.findall(r'\\boxed\{([0-9.,-]+)\}', text)
    if m_box:
        return [m.replace(',', '').strip() for m in m_box]
    m_hash = re.findall(r'####\s*([0-9.,-]+)', text)
    if m_hash:
        return [m.replace(',', '').strip() for m in m_hash]
    nums = re.findall(r'-?\d+(?:\.\d+)?', cleaned_text)
    return nums

def evaluate_smart(item, content, thinking):
    eval_type = item.get("eval_type", "exact_contains")
    expected = str(item.get("expected_answer", "")).strip().lower()
    content_clean = (content or "").strip()
    thinking_clean = (thinking or "").strip()
    
    if eval_type == "negative_constraint":
        forbidden = expected.lower()
        text_to_check = content_clean
        if not text_to_check or len(text_to_check) < 3:
            lines = [l.strip() for l in thinking_clean.splitlines() if l.strip()]
            for line in reversed(lines):
                if "," in line or " " in line:
                    text_to_check = line
                    break
        words = re.findall(r'\b[a-zA-Z]{3,}\b', text_to_check)
        meta_words = {
            'thinking', 'process', 'analyze', 'constraint', 'target', 'format',
            'word', 'words', 'none', 'letter', 'provide', 'only', 'common', 'english',
            'longer', 'contain', 'letters', 'input', 'output', 'option', 'plan'
        }
        valid_words = [w.lower() for w in words if w.lower() not in meta_words]
        has_forbidden = any(forbidden in w for w in valid_words)
        correct = (len(valid_words) >= 3) and (not has_forbidden)
        return correct, f"forbidden='{forbidden}', words={valid_words[:4]}, forbidden_present={has_forbidden}"
        
    elif eval_type == "numeric":
        nums_c = extract_numbers(content_clean)
        nums_t = extract_numbers(thinking_clean)
        if nums_c and (expected in nums_c or nums_c[-1] == expected):
            return True, f"Matched numeric in content: {nums_c[-1]}"
        if nums_t and (expected in nums_t or nums_t[-1] == expected):
            return True, f"Matched numeric in thinking: {nums_t[-1]}"
        if expected.endswith("000"):
            k_ver = expected[:-3] + "k"
            if k_ver in content_clean.lower() or k_ver in thinking_clean.lower():
                return True, f"Matched 'k' notation: {k_ver}"
        return False, f"Numeric mismatch: exp={expected}, content={nums_c}, think_tail={nums_t[-3:] if nums_t else []}"
        
    elif eval_type == "choice":
        m = re.search(r'\b(?:option|answer|choice)?\s*[:\(\[]?\s*([A-D])\b', content_clean, re.IGNORECASE)
        if m and m.group(1).upper() == expected.upper():
            return True, f"Matched choice in content: {m.group(1).upper()}"
        tail_think = thinking_clean[-500:]
        m_think_expl = re.search(r'\b(?:correct|final|best)\s+(?:option|choice|answer)\s+is\s+[:\(\[]?\s*([A-D])\b', tail_think, re.IGNORECASE)
        if m_think_expl and m_think_expl.group(1).upper() == expected.upper():
            return True, f"Matched choice in thinking deduction: {m_think_expl.group(1).upper()}"
        m_think = re.findall(r'\b([A-D])\b', tail_think)
        if m_think and m_think[-1].upper() == expected.upper():
            return True, f"Matched choice in thinking tail: {m_think[-1].upper()}"
        return False, f"Choice mismatch: exp={expected.upper()}, content='{content_clean[:60]}'"
        
    elif eval_type == "code_ast":
        block = content_clean if "```" in content_clean else thinking_clean
        if "```python" in block:
            block = block.split("```python")[1].split("```")[0].strip()
        elif "```" in block:
            block = block.split("```")[1].split("```")[0].strip()
        try:
            ast.parse(block)
            ast_ok = True
        except Exception:
            ast_ok = False
        has_def = "def " in block
        return (ast_ok and has_def), f"AST={ast_ok}, def={has_def}"
        
    else:  # exact_contains
        if expected in content_clean.lower():
            return True, f"Found '{expected}' in content"
        if expected in thinking_clean.lower():
            return True, f"Found '{expected}' in thinking"
        return False, f"'{expected}' not found in content or thinking"

def query_chat(model_id, prompt, num_predict=1024):
    payload = {
        "model": model_id,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "stream": False,
        "options": {
            "temperature": 0.2,
            "top_p": 0.95,
            "repeat_penalty": 1.15,
            "num_predict": num_predict
        }
    }
    t0 = time.time()
    for attempt in range(3):
        try:
            req = urllib.request.Request(
                OLLAMA_API_CHAT,
                data=json.dumps(payload).encode("utf-8"),
                headers={"Content-Type": "application/json"}
            )
            with urllib.request.urlopen(req, timeout=120) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                msg = data.get("message", {})
                content = msg.get("content", "")
                thinking = msg.get("thinking", "")
                done_reason = data.get("done_reason", "unknown")
                dt = time.time() - t0
                return content, thinking, done_reason, dt
        except Exception as e:
            time.sleep(1.0 * (attempt + 1))
    return "", "", "timeout_error", time.time() - t0

def main():
    print("=" * 80)
    print("  PROJECT NORN V18: QUICK HEAD-TO-HEAD ABLATION BENCHMARK")
    print("  Stratified Evaluation Across All 15 Reasoning Pillars")
    print("=" * 80)

    # 1. Load benchmark suite
    with open(suite_path, "r", encoding="utf-8") as f:
        full_suite = json.load(f)

    # 2. Select 1 pristine representative question per category (15 total, 1 per pillar)
    categories = defaultdict(list)
    for item in full_suite:
        cat = item.get("category", "General")
        categories[cat].append(item)

    stratified_sample = []
    for cat, items in categories.items():
        stratified_sample.append(items[0])

    print(f"\n[+] Selected {len(stratified_sample)} pristine stratified items (1 per pillar across all {len(categories)} pillars).")
    print("=" * 80)

    results_summary = {}

    for model_cfg in MODELS_TO_TEST:
        model_name = model_cfg["name"]
        model_id = model_cfg["ollama_id"]
        print(f"\n>>> Running Evaluation on: {model_name} (Ollama: {model_id})")

        correct_count = 0
        total_time = 0.0
        pillar_results = defaultdict(lambda: {"total": 0, "correct": 0})
        detailed_records = []

        for idx, item in enumerate(stratified_sample, 1):
            q_id = item["id"]
            cat = item["category"]
            prompt = item["prompt"]
            exp = item.get("expected_answer", "")
            n_pred = 1024
            if "Word Problems" in cat or "STEM" in cat:
                n_pred = 1280
            elif "Negative" in cat:
                n_pred = 768

            content, thinking, done_reason, dt = query_chat(model_id, prompt, num_predict=n_pred)
            total_time += dt
            correct, notes = evaluate_smart(item, content, thinking)

            if correct:
                correct_count += 1
                status = "[PASS]"
            else:
                status = "[FAIL]"

            pillar_results[cat]["total"] += 1
            if correct:
                pillar_results[cat]["correct"] += 1

            detailed_records.append({
                "id": q_id,
                "category": cat,
                "correct": correct,
                "latency_s": round(dt, 2),
                "notes": notes,
                "think_tokens": len(thinking.split()) if thinking else 0
            })

            print(f"  [{idx:02d}/{len(stratified_sample):02d}] {status} {cat[:28]:<28} ({dt:4.1f}s) | exp='{str(exp)[:15]}' | {notes[:40]}")

        acc = (correct_count / len(stratified_sample)) * 100.0
        avg_latency = total_time / len(stratified_sample) if stratified_sample else 0.0

        results_summary[model_name] = {
            "model_id": model_id,
            "correct": correct_count,
            "total": len(stratified_sample),
            "accuracy_pct": round(acc, 2),
            "avg_latency_s": round(avg_latency, 2),
            "total_time_s": round(total_time, 2),
            "pillar_breakdown": {
                c: f"{v['correct']}/{v['total']} ({(v['correct']/v['total'])*100:.0f}%)"
                for c, v in pillar_results.items()
            },
            "records": detailed_records
        }

        print(f"\n[RESULTS] {model_name}: {correct_count}/{len(stratified_sample)} ({acc:.1f}%) | Avg Latency: {avg_latency:.2f}s")
        print("-" * 80)

    # Save to disk
    with open(output_report_path, "w", encoding="utf-8") as f:
        json.dump(results_summary, f, indent=2)

    print("\n" + "=" * 80)
    print("  FINAL HEAD-TO-HEAD ABLATION SCOREBOARD (15-Pillar Stratified Suite)")
    print("=" * 80)
    print(f"{'Model Configuration':<32} | {'Score':<14} | {'Accuracy':<10} | {'Avg Latency':<12}")
    print("-" * 80)
    for name, data in results_summary.items():
        print(f"{name:<32} | {data['correct']}/{data['total']:<12} | {data['accuracy_pct']:>6.1f}%    | {data['avg_latency_s']:>6.2f}s/q")
    print("=" * 80)
    print(f"Detailed ablation log saved to: {output_report_path}\n")

if __name__ == "__main__":
    main()