classify_benchmark.py

#!/usr/bin/env python3
"""
Classification accuracy benchmark for CL-bench tasks.

Runs CL-bench tasks through the pipeline's classifiers, compares against
ground-truth metadata.context_category, and outputs accuracy reports.

Modes:
  --mode heuristic  Heuristic-only (default, no LLM calls, fast)
  --mode full       Full classification pipeline (heuristic + LLM + guards)
  --mode ensemble   Ensemble voter system (Phase 3 — all voters + confidence)

Usage:
    python3 classify_benchmark.py --input CL-bench.jsonl
    python3 classify_benchmark.py --input CL-bench.jsonl --mode full
    python3 classify_benchmark.py --input CL-bench.jsonl --mode ensemble
"""

import argparse
import asyncio
import json
import sys
import os
from pathlib import Path
from collections import Counter, defaultdict

# Add project root so serve2 imports work
_project_root = str(Path(__file__).resolve().parent.parent)
if _project_root not in sys.path:
    sys.path.insert(0, _project_root)
os.chdir(_project_root)

# Import classifier functions from serve2.py
from reasoner_2.serve2 import _classify_context_category_heuristic


# CL-bench category name → pipeline category mapping
CATEGORY_MAP = {
    "Domain Knowledge Reasoning": "domain_knowledge",
    "Rule System Application": "rule_system",
    "Procedural Task Execution": "procedural",
    "Empirical Discovery & Simulation": "empirical",
}


def load_tasks(input_path: str) -> list[dict]:
    """Load tasks from JSONL file."""
    tasks = []
    with open(input_path, "r") as f:
        for line in f:
            line = line.strip()
            if line:
                tasks.append(json.loads(line))
    return tasks


def extract_system_and_query(messages: list[dict]) -> tuple[str, str]:
    """Extract system prompt and last user query from messages."""
    system_prompt = ""
    query = ""
    last_user_idx = -1
    for i, msg in enumerate(messages):
        if msg.get("role") == "system":
            system_prompt = msg.get("content", "")
        if msg.get("role") == "user":
            last_user_idx = i
    if last_user_idx >= 0:
        content = messages[last_user_idx].get("content", "")
        if isinstance(content, list):
            content = " ".join(p.get("text", "") for p in content if isinstance(p, dict))
        query = content
    return system_prompt, query


def classify_all_heuristic(tasks: list[dict]) -> list[dict]:
    """Run heuristic classifier on all tasks."""
    results = []
    for task in tasks:
        meta = task.get("metadata", {})
        task_id = meta.get("task_id", "unknown")
        gt_category = meta.get("context_category", "unknown")
        sub_category = meta.get("sub_category", "unknown")
        gt_pipeline = CATEGORY_MAP.get(gt_category, "unknown")

        messages = task.get("messages", [])
        system_prompt, query = extract_system_and_query(messages)
        heuristic = _classify_context_category_heuristic(system_prompt, query)

        results.append({
            "task_id": task_id,
            "gt_category": gt_category,
            "gt_pipeline": gt_pipeline,
            "sub_category": sub_category,
            "heuristic": heuristic,
            "heuristic_match": heuristic == gt_pipeline,
            "predicted": heuristic,
            "predicted_match": heuristic == gt_pipeline,
            "messages_count": len(messages),
            "is_multi_turn": any(m.get("role") == "assistant" for m in messages),
            "context_chars": sum(len(m.get("content", "")) for m in messages if isinstance(m.get("content"), str)),
            "rubric_count": len(task.get("rubrics", [])),
        })
    return results


async def classify_all_full(tasks: list[dict]) -> list[dict]:
    """Run full classification pipeline (heuristic + LLM + guards)."""
    from reasoner_2.serve2 import _classify_context_category

    # Need a processor instance for LLM dispatch
    try:
        from reasoner_2.serve2 import CLBenchProcessor
        processor = CLBenchProcessor()
    except Exception as e:
        print(f"WARNING: Could not create processor: {e}")
        print("Falling back to heuristic-only mode")
        return classify_all_heuristic(tasks)

    results = []
    for i, task in enumerate(tasks):
        meta = task.get("metadata", {})
        task_id = meta.get("task_id", "unknown")
        gt_category = meta.get("context_category", "unknown")
        sub_category = meta.get("sub_category", "unknown")
        gt_pipeline = CATEGORY_MAP.get(gt_category, "unknown")

        messages = task.get("messages", [])
        system_prompt, query = extract_system_and_query(messages)
        heuristic = _classify_context_category_heuristic(system_prompt, query)

        # Extract first user content for context_sample
        first_user_content = ""
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "")
                if isinstance(content, list):
                    content = " ".join(p.get("text", "") for p in content if isinstance(p, dict))
                first_user_content = (content or "")[:5000]
                break

        try:
            full_result = await _classify_context_category(
                processor, system_prompt, query,
                context_sample=first_user_content
            )
        except Exception as e:
            print(f"  LLM classification failed for task {i}: {e}")
            full_result = heuristic  # fallback

        results.append({
            "task_id": task_id,
            "gt_category": gt_category,
            "gt_pipeline": gt_pipeline,
            "sub_category": sub_category,
            "heuristic": heuristic,
            "heuristic_match": heuristic == gt_pipeline,
            "full": full_result,
            "full_match": full_result == gt_pipeline,
            "predicted": full_result,
            "predicted_match": full_result == gt_pipeline,
            "messages_count": len(messages),
            "is_multi_turn": any(m.get("role") == "assistant" for m in messages),
            "context_chars": sum(len(m.get("content", "")) for m in messages if isinstance(m.get("content"), str)),
            "rubric_count": len(task.get("rubrics", [])),
        })

        if (i + 1) % 50 == 0:
            print(f"  Classified {i + 1}/{len(tasks)}...")

    return results


async def classify_all_ensemble(tasks: list[dict]) -> list[dict]:
    """Run ensemble voter classification (Phase 3)."""
    try:
        from reasoner_2.serve2 import _classify_ensemble
    except ImportError:
        print("WARNING: _classify_ensemble not available — falling back to full mode")
        return await classify_all_full(tasks)

    from reasoner_2.serve2 import CLBenchProcessor
    processor = CLBenchProcessor()

    results = []
    for i, task in enumerate(tasks):
        meta = task.get("metadata", {})
        task_id = meta.get("task_id", "unknown")
        gt_category = meta.get("context_category", "unknown")
        sub_category = meta.get("sub_category", "unknown")
        gt_pipeline = CATEGORY_MAP.get(gt_category, "unknown")

        messages = task.get("messages", [])
        system_prompt, query = extract_system_and_query(messages)

        first_user_content = ""
        for msg in messages:
            if msg.get("role") == "user":
                content = msg.get("content", "")
                if isinstance(content, list):
                    content = " ".join(p.get("text", "") for p in content if isinstance(p, dict))
                first_user_content = (content or "")[:5000]
                break

        try:
            ensemble_result = await _classify_ensemble(
                processor, system_prompt, query,
                context_sample=first_user_content
            )
            predicted = ensemble_result["category"]
            confidence = ensemble_result["confidence"]
        except Exception as e:
            print(f"  Ensemble classification failed for task {i}: {e}")
            predicted = _classify_context_category_heuristic(system_prompt, query)
            confidence = 0.0

        heuristic = _classify_context_category_heuristic(system_prompt, query)

        results.append({
            "task_id": task_id,
            "gt_category": gt_category,
            "gt_pipeline": gt_pipeline,
            "sub_category": sub_category,
            "heuristic": heuristic,
            "heuristic_match": heuristic == gt_pipeline,
            "predicted": predicted,
            "predicted_match": predicted == gt_pipeline,
            "confidence": confidence,
            "messages_count": len(messages),
            "is_multi_turn": any(m.get("role") == "assistant" for m in messages),
            "context_chars": sum(len(m.get("content", "")) for m in messages if isinstance(m.get("content"), str)),
            "rubric_count": len(task.get("rubrics", [])),
        })

        if (i + 1) % 50 == 0:
            print(f"  Classified {i + 1}/{len(tasks)}...")

    return results


def print_confusion_matrix(results: list[dict], label: str = "Predicted"):
    """Print confusion matrix: rows=ground truth, cols=predicted."""
    categories = ["rule_system", "empirical", "domain_knowledge", "procedural"]
    matrix = defaultdict(Counter)
    for r in results:
        matrix[r["gt_pipeline"]][r["predicted"]] += 1

    print(f"\n=== CONFUSION MATRIX ({label}) ===")
    print(f"{'Ground Truth':<25}", end="")
    for c in categories:
        print(f"{c:>18}", end="")
    print(f"{'Total':>10}{'Accuracy':>10}")
    print("-" * 100)

    total_correct = 0
    total_tasks = 0
    for gt in categories:
        row = matrix.get(gt, Counter())
        row_total = sum(row.values())
        correct = row.get(gt, 0)
        total_correct += correct
        total_tasks += row_total
        print(f"{gt:<25}", end="")
        for c in categories:
            count = row.get(c, 0)
            print(f"{count:>18}", end="")
        acc = correct / row_total if row_total > 0 else 0
        print(f"{row_total:>10}{acc:>9.1%}")

    overall_acc = total_correct / total_tasks if total_tasks > 0 else 0
    print("-" * 100)
    print(f"{'Overall':<25}", end="")
    print(f"{'':>72}{total_tasks:>10}{overall_acc:>9.1%}")


def print_subcategory_accuracy(results: list[dict], pred_key: str = "predicted"):
    """Print per-subcategory classification accuracy."""
    subcats = defaultdict(lambda: {"total": 0, "correct": 0, "predictions": Counter()})
    for r in results:
        key = f"{r['gt_category']} / {r['sub_category']}"
        subcats[key]["total"] += 1
        if r["predicted_match"]:
            subcats[key]["correct"] += 1
        subcats[key]["predictions"][r[pred_key]] += 1

    print(f"\n=== PER-SUBCATEGORY ACCURACY ({pred_key}) ===")
    print(f"{'Subcategory':<55}{'Total':>8}{'Correct':>10}{'Accuracy':>10}{'Predictions':>30}")
    print("-" * 115)

    for key in sorted(subcats.keys()):
        s = subcats[key]
        acc = s["correct"] / s["total"] if s["total"] > 0 else 0
        preds = ", ".join(f"{k}:{v}" for k, v in s["predictions"].most_common())
        print(f"{key:<55}{s['total']:>8}{s['correct']:>10}{acc:>9.1%}  {preds}")


def print_dataset_stats(results: list[dict]):
    """Print overall dataset statistics."""
    print("\n=== DATASET STATISTICS ===")
    print(f"Total tasks: {len(results)}")

    # Per-category stats
    cats = defaultdict(list)
    for r in results:
        cats[r["gt_category"]].append(r)

    print(f"\n{'Category':<40}{'Tasks':>8}{'Multi-Turn':>12}{'Avg Context':>14}{'Avg Rubrics':>14}")
    print("-" * 90)
    for cat in sorted(cats.keys()):
        tasks = cats[cat]
        n = len(tasks)
        multi = sum(1 for t in tasks if t["is_multi_turn"])
        avg_ctx = sum(t["context_chars"] for t in tasks) / n
        avg_rub = sum(t["rubric_count"] for t in tasks) / n
        print(f"{cat:<40}{n:>8}{multi/n:>11.0%}{avg_ctx:>13,.0f}{avg_rub:>14.1f}")


def print_confidence_bands(results: list[dict]):
    """Print accuracy by confidence band (ensemble mode only)."""
    if not any("confidence" in r for r in results):
        return

    bands = [
        ("low (< 0.4)", lambda c: c < 0.4),
        ("medium (0.4-0.6)", lambda c: 0.4 <= c < 0.6),
        ("high (0.6-0.8)", lambda c: 0.6 <= c < 0.8),
        ("very high (>= 0.8)", lambda c: c >= 0.8),
    ]

    print("\n=== ACCURACY BY CONFIDENCE BAND ===")
    print(f"{'Band':<25}{'Tasks':>8}{'Correct':>10}{'Accuracy':>10}{'Fallback%':>12}")
    print("-" * 65)

    for label, pred_fn in bands:
        band_results = [r for r in results if "confidence" in r and pred_fn(r["confidence"])]
        if not band_results:
            print(f"{label:<25}{'0':>8}{'—':>10}{'—':>10}{'—':>12}")
            continue
        correct = sum(1 for r in band_results if r["predicted_match"])
        fallback = sum(1 for r in band_results if r["predicted"] == "procedural")
        n = len(band_results)
        print(f"{label:<25}{n:>8}{correct:>10}{correct/n:>9.1%}{fallback/n:>11.1%}")


def print_sample_coverage(sample_path: str, label: str):
    """Print subcategory coverage of a sample file."""
    try:
        tasks = load_tasks(sample_path)
    except FileNotFoundError:
        print(f"\n{label}: FILE NOT FOUND ({sample_path})")
        return

    cats = Counter()
    for task in tasks:
        meta = task.get("metadata", {})
        key = f"{meta.get('context_category', '?')} / {meta.get('sub_category', '?')}"
        cats[key] += 1

    print(f"\n=== {label} ({len(tasks)} tasks) ===")
    for k, v in sorted(cats.items()):
        print(f"  {k}: {v}")


def save_results(results: list[dict], output_path: str):
    """Save per-task results to JSONL."""
    with open(output_path, "w") as f:
        for r in results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    print(f"\nPer-task results saved to: {output_path}")


def main():
    parser = argparse.ArgumentParser(description="CL-bench classification accuracy benchmark")
    parser.add_argument("--input", default="CL-bench.jsonl", help="Input JSONL file")
    parser.add_argument("--output", default="outputs/classification_results.jsonl", help="Output JSONL file")
    parser.add_argument("--mode", choices=["heuristic", "full", "ensemble"], default="heuristic",
                        help="Classification mode: heuristic (fast, no LLM), full (heuristic+LLM+guards), ensemble (all voters)")
    # Keep --llm for backward compatibility
    parser.add_argument("--llm", action="store_true", help="(Deprecated) Use --mode full instead")
    args = parser.parse_args()

    # Backward compat: --llm maps to --mode full
    mode = args.mode
    if args.llm and mode == "heuristic":
        mode = "full"

    print(f"Loading tasks from {args.input}...")
    tasks = load_tasks(args.input)
    print(f"Loaded {len(tasks)} tasks")

    # Run classifier
    print(f"\nRunning {mode} classifier...")
    if mode == "heuristic":
        results = classify_all_heuristic(tasks)
        label = "Heuristic"
    elif mode == "full":
        results = asyncio.run(classify_all_full(tasks))
        label = "Full (Heuristic+LLM+Guards)"
    elif mode == "ensemble":
        results = asyncio.run(classify_all_ensemble(tasks))
        label = "Ensemble"

    # Print results
    print_dataset_stats(results)
    print_confusion_matrix(results, label=label)
    print_subcategory_accuracy(results)
    if mode == "ensemble":
        print_confidence_bands(results)

    # Sample file coverage
    base_dir = Path(args.input).parent
    print_sample_coverage(str(base_dir / "CL-bench-10.jsonl"), "CL-bench-10.jsonl")
    print_sample_coverage(str(base_dir / "CL-bench-20curated.jsonl"), "CL-bench-20curated.jsonl")

    # Save results
    os.makedirs(Path(args.output).parent, exist_ok=True)
    save_results(results, args.output)

    # Summary
    correct = sum(1 for r in results if r["predicted_match"])
    print(f"\n=== SUMMARY ===")
    print(f"Mode: {mode}")
    print(f"Accuracy: {correct}/{len(results)} ({correct/len(results):.1%})")
    for cat in ["rule_system", "empirical", "domain_knowledge", "procedural"]:
        cat_total = sum(1 for r in results if r["gt_pipeline"] == cat)
        cat_correct = sum(1 for r in results if r["gt_pipeline"] == cat and r["predicted_match"])
        print(f"  {cat}: {cat_correct}/{cat_total}")


if __name__ == "__main__":
    main()