eval/eval_all_benchmark_core.py

"""
Comprehensive Benchmark Evaluation for MathSolver AI Core (Excluding Manim)
Evaluates all 165 questions from the National High School Exam Geometry Benchmark (2016-2026).
Model: Google Gemma 4 26B A4B (gemini/gemma-4-26b-a4b-it) in Low Thinking mode.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import sys
import time
from pathlib import Path
from typing import Dict, Any, List, Optional
import numpy as np

# Suppress debug logs
logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from agents.geometry_parser_agent import GeometryParserAgent
from solver.dsl_parser import DSLParser
from solver.engine import GeometryEngine
from solver.validator import GeometryValidator


DATASET_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "../../research/geometry_dataset/all_geometry_questions.jsonl")
)
RESULTS_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "full_benchmark_gemma26b_results.json")
)
METRICS_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "full_benchmark_gemma26b_metrics.json")
)

SEMAPHORE_LIMIT = 3  # Concurrency limit across 10 API keys


async def evaluate_single_question(
    q: Dict[str, Any],
    parser_agent: GeometryParserAgent,
    sem: asyncio.Semaphore,
    progress_info: Dict[str, Any],
) -> Dict[str, Any]:
    qid = q["id"]
    year = q["year"]
    topic = q.get("topic", "unknown")
    subtopic = q.get("subtopic", "unknown")
    difficulty = q.get("difficulty", "medium")
    raw_text = q.get("raw_text", "").strip()
    gt_answer = q.get("answer", "")

    result: Dict[str, Any] = {
        "id": qid,
        "year": year,
        "topic": topic,
        "subtopic": subtopic,
        "difficulty": difficulty,
        "raw_text": raw_text,
        "ground_truth_answer": gt_answer,
        "stages": {},
        "overall_status": "FAILED",
        "failure_reason": None,
        "total_latency_ms": 0.0,
    }

    async with sem:
        # Stage 1: LLM Parsing -> Geometry DSL (Gemma 4 26B A4B Low Thinking)
        t0 = time.time()
        try:
            parsed = await parser_agent.process(raw_text)
            t_llm = (time.time() - t0) * 1000
            dsl_text = parsed.get("geometry_dsl", "")
            shape_type = parsed.get("type", "")

            if not dsl_text:
                raise ValueError("Empty DSL generated by parser")

            result["stages"]["llm_parsing"] = {
                "status": "OK",
                "latency_ms": round(t_llm, 1),
                "type": shape_type,
                "dsl": dsl_text,
                "entities": parsed.get("entities", []),
            }
        except Exception as e:
            t_llm = (time.time() - t0) * 1000
            result["stages"]["llm_parsing"] = {
                "status": "ERROR",
                "latency_ms": round(t_llm, 1),
                "error": str(e),
            }
            result["overall_status"] = "PARSE_FAILED"
            result["failure_reason"] = f"LLM parsing error: {e}"
            result["total_latency_ms"] = round(t_llm, 1)
            _log_progress(progress_info, qid, year, result)
            return result

        # Stage 2: DSL Parser / Compiler -> AST & Primitive Constraints
        t0 = time.time()
        try:
            dsl_parser = DSLParser()
            points, constraints, is_3d = dsl_parser.parse(dsl_text)
            t_compile = (time.time() - t0) * 1000

            result["stages"]["dsl_compiler"] = {
                "status": "OK",
                "latency_ms": round(t_compile, 1),
                "num_points": len(points),
                "num_constraints": len(constraints),
                "is_3d": is_3d,
                "point_ids": [p.id for p in points],
            }
        except Exception as e:
            t_compile = (time.time() - t0) * 1000
            result["stages"]["dsl_compiler"] = {
                "status": "ERROR",
                "latency_ms": round(t_compile, 1),
                "error": str(e),
            }
            result["overall_status"] = "COMPILE_FAILED"
            result["failure_reason"] = f"DSL syntax/compile error: {e}"
            result["total_latency_ms"] = round(t_llm + t_compile, 1)
            _log_progress(progress_info, qid, year, result)
            return result

        # Stage 3: GeometryEngine -> Numerical Coordinate Solving
        t0 = time.time()
        try:
            engine = GeometryEngine()
            solved_ir = engine.solve(points, constraints, is_3d=is_3d)
            t_solve = (time.time() - t0) * 1000

            if not solved_ir or not solved_ir.get("coordinates"):
                raise ValueError("Engine returned empty coordinates")

            coords = solved_ir.get("coordinates", {})
            result["stages"]["geometry_engine"] = {
                "status": "OK",
                "latency_ms": round(t_solve, 1),
                "num_coords": len(coords),
                "coordinates": {k: [round(float(v), 4) for v in val] for k, val in coords.items()},
                "num_faces": len(solved_ir.get("faces", [])),
                "num_segments": len(solved_ir.get("segments", [])),
            }
        except Exception as e:
            t_solve = (time.time() - t0) * 1000
            result["stages"]["geometry_engine"] = {
                "status": "ERROR",
                "latency_ms": round(t_solve, 1),
                "error": str(e),
            }
            result["overall_status"] = "SOLVE_FAILED"
            result["failure_reason"] = f"GeometryEngine solve error: {e}"
            result["total_latency_ms"] = round(t_llm + t_compile + t_solve, 1)
            _log_progress(progress_info, qid, year, result)
            return result

        # Stage 4: GeometryValidator -> Invariant & Semantic Verification
        t0 = time.time()
        try:
            validator = GeometryValidator()
            val_res = validator.validate(solved_ir, constraints, is_3d=is_3d)
            diag = val_res.to_dict()
            t_val = (time.time() - t0) * 1000

            is_valid = diag.get("is_valid", False)
            result["stages"]["validation"] = {
                "status": "OK" if is_valid else "WARNING",
                "latency_ms": round(t_val, 1),
                "is_valid": is_valid,
                "checked_count": diag.get("checked_count", 0),
                "error_count": len(diag.get("errors", [])),
                "errors": diag.get("errors", []),
                "warnings": diag.get("warnings", []),
            }
            if is_valid:
                result["overall_status"] = "SUCCESS"
            else:
                result["overall_status"] = "SOLVED_WITH_WARNINGS"
                result["failure_reason"] = f"Validator detected invariant violations: {diag.get('errors')[:2]}"
        except Exception as e:
            t_val = (time.time() - t0) * 1000
            result["stages"]["validation"] = {
                "status": "ERROR",
                "latency_ms": round(t_val, 1),
                "error": str(e),
            }
            result["overall_status"] = "VALIDATION_ERROR"
            result["failure_reason"] = f"Validator error: {e}"

        result["total_latency_ms"] = round(
            t_llm + t_compile + t_solve + result["stages"]["validation"].get("latency_ms", 0), 1
        )
        _log_progress(progress_info, qid, year, result)
        return result


def _log_progress(p: Dict[str, Any], qid: str, year: int, res: Dict[str, Any]):
    p["completed"] += 1
    st = res["overall_status"]
    t_llm = res["stages"].get("llm_parsing", {}).get("latency_ms", 0)
    t_eng = res["stages"].get("geometry_engine", {}).get("latency_ms", 0)
    t_tot = res["total_latency_ms"]

    icon = "✅" if st == "SUCCESS" else ("⚠️" if st == "SOLVED_WITH_WARNINGS" else "❌")
    print(
        f"[{p['completed']:3d}/{p['total']}] {icon} {year} | {qid[:24]:<24} | "
        f"Status: {st:<20} | LLM: {t_llm:6.1f}ms | Eng: {t_eng:5.1f}ms | Total: {t_tot/1000:4.1f}s"
    )
    if res.get("failure_reason"):
        print(f"      ↳ Reason: {res['failure_reason'][:120]}")


def calculate_metrics(results: List[Dict[str, Any]]) -> Dict[str, Any]:
    n = len(results)
    if n == 0:
        return {}

    # Pipeline stage counts
    parse_ok = sum(1 for r in results if r["stages"].get("llm_parsing", {}).get("status") == "OK")
    compile_ok = sum(1 for r in results if r["stages"].get("dsl_compiler", {}).get("status") == "OK")
    solve_ok = sum(1 for r in results if r["stages"].get("geometry_engine", {}).get("status") == "OK")
    val_ok = sum(1 for r in results if r["stages"].get("validation", {}).get("is_valid") is True)
    full_success = sum(1 for r in results if r["overall_status"] == "SUCCESS")
    solved_any = sum(1 for r in results if r["overall_status"] in ("SUCCESS", "SOLVED_WITH_WARNINGS"))

    # Latencies
    llm_lats = [r["stages"]["llm_parsing"]["latency_ms"] for r in results if "llm_parsing" in r["stages"] and "latency_ms" in r["stages"]["llm_parsing"]]
    comp_lats = [r["stages"]["dsl_compiler"]["latency_ms"] for r in results if "dsl_compiler" in r["stages"] and "latency_ms" in r["stages"]["dsl_compiler"]]
    eng_lats = [r["stages"]["geometry_engine"]["latency_ms"] for r in results if "geometry_engine" in r["stages"] and "latency_ms" in r["stages"]["geometry_engine"]]
    tot_lats = [r["total_latency_ms"] for r in results]

    def lat_stats(arr):
        if not arr:
            return {}
        a = np.array(arr)
        return {
            "mean_ms": round(float(np.mean(a)), 1),
            "median_ms": round(float(np.median(a)), 1),
            "p90_ms": round(float(np.percentile(a, 90)), 1),
            "p95_ms": round(float(np.percentile(a, 95)), 1),
            "min_ms": round(float(np.min(a)), 1),
            "max_ms": round(float(np.max(a)), 1),
        }

    # Breakdown helper
    def group_stats(key: str):
        groups: Dict[str, List[Dict[str, Any]]] = {}
        for r in results:
            val = str(r.get(key, "unknown"))
            groups.setdefault(val, []).append(r)
        
        stat = {}
        for g_name, g_items in sorted(groups.items()):
            g_n = len(g_items)
            g_parse = sum(1 for r in g_items if r["stages"].get("llm_parsing", {}).get("status") == "OK")
            g_solve = sum(1 for r in g_items if r["stages"].get("geometry_engine", {}).get("status") == "OK")
            g_success = sum(1 for r in g_items if r["overall_status"] == "SUCCESS")
            stat[g_name] = {
                "total": g_n,
                "parse_rate_pct": round(g_parse / g_n * 100, 1),
                "solve_rate_pct": round(g_solve / g_n * 100, 1),
                "success_rate_pct": round(g_success / g_n * 100, 1),
            }
        return stat

    # Top Failure reasons
    failure_counts: Dict[str, int] = {}
    for r in results:
        reason = r.get("failure_reason")
        if reason:
            # categorize
            cat = "Other"
            if "DSL syntax" in reason or "compile error" in reason:
                cat = "DSL Compilation Syntax Error"
            elif "Empty DSL" in reason or "LLM parsing" in reason:
                cat = "LLM Generation Failure"
            elif "Engine returned empty" in reason or "solve error" in reason:
                cat = "Engine Constraint Solver Error"
            elif "Validator detected" in reason:
                cat = "Geometric Invariant Warning"
            failure_counts[cat] = failure_counts.get(cat, 0) + 1

    return {
        "total_questions": n,
        "summary": {
            "llm_parse_success_count": parse_ok,
            "llm_parse_success_rate_pct": round(parse_ok / n * 100, 2),
            "dsl_compiler_success_count": compile_ok,
            "dsl_compiler_success_rate_pct": round(compile_ok / n * 100, 2),
            "geometry_engine_solve_count": solve_ok,
            "geometry_engine_solve_rate_pct": round(solve_ok / n * 100, 2),
            "geometry_validator_pass_count": val_ok,
            "geometry_validator_pass_rate_pct": round(val_ok / n * 100, 2),
            "full_pipeline_success_count": full_success,
            "full_pipeline_success_rate_pct": round(full_success / n * 100, 2),
            "solved_any_rate_pct": round(solved_any / n * 100, 2),
        },
        "latency_stats": {
            "llm_parsing": lat_stats(llm_lats),
            "dsl_compiler": lat_stats(comp_lats),
            "geometry_engine": lat_stats(eng_lats),
            "total_core": lat_stats(tot_lats),
        },
        "by_topic": group_stats("topic"),
        "by_difficulty": group_stats("difficulty"),
        "by_year": group_stats("year"),
        "failure_categories": failure_counts,
    }


async def main():
    print(f"=== STARTING FULL BENCHMARK EVALUATION (N=165) ===")
    print(f"Model: Google Gemma 4 26B A4B (gemini/gemma-4-26b-a4b-it) Low Thinking")
    print(f"Pipeline: Natural Text -> GeometryParserAgent -> DSLParser -> GeometryEngine -> GeometryValidator")
    print(f"Excluding: Manim Video Rendering (AI Core only)\n")

    with open(DATASET_PATH, "r", encoding="utf-8") as f:
        questions = [json.loads(line) for line in f]

    print(f"Loaded {len(questions)} questions from {DATASET_PATH}")
    assert len(questions) == 165

    # Check for existing partial results to enable resume
    existing_results: Dict[str, Dict[str, Any]] = {}
    if os.path.exists(RESULTS_PATH):
        try:
            with open(RESULTS_PATH, "r", encoding="utf-8") as f:
                loaded = json.load(f)
                if isinstance(loaded, list):
                    for item in loaded:
                        if "id" in item:
                            existing_results[item["id"]] = item
            print(f"[*] Found {len(existing_results)} existing results in {RESULTS_PATH}. Resuming...")
        except Exception as e:
            print(f"[*] Could not load existing results: {e}. Starting fresh.")

    results_dict: Dict[str, Dict[str, Any]] = dict(existing_results)
    # Filter remaining questions (including questions that failed earlier due to rate limit)
    remaining_questions = [
        q for q in questions
        if q["id"] not in results_dict or results_dict[q["id"]].get("overall_status") == "PARSE_FAILED"
    ]
    # Keep already successful/solved results
    valid_existing = {k: v for k, v in results_dict.items() if v.get("overall_status") != "PARSE_FAILED"}
    results_dict = dict(valid_existing)

    print(f"[*] Already solved/evaluated: {len(results_dict)} questions.")
    print(f"[*] Remaining questions to evaluate: {len(remaining_questions)} / {len(questions)}")

    parser_agent = GeometryParserAgent()
    sem = asyncio.Semaphore(2)  # Controlled concurrency of 2
    progress = {"completed": len(results_dict), "total": len(questions)}
    results_lock = asyncio.Lock()

    async def wrapped_eval(q):
        try:
            res = await evaluate_single_question(q, parser_agent, sem, progress)
            async with results_lock:
                results_dict[q["id"]] = res
                # Save checkpoint on every completion
                with open(RESULTS_PATH, "w", encoding="utf-8") as f_out:
                    ordered = [results_dict.get(orig["id"]) for orig in questions if orig["id"] in results_dict]
                    json.dump(ordered, f_out, ensure_ascii=False, indent=2)
        except Exception as err:
            logger.error(f"Error in wrapped_eval for {q.get('id')}: {err}", exc_info=True)
            print(f"Error in wrapped_eval for {q.get('id')}: {err}")

    start_time = time.time()
    if remaining_questions:
        running_tasks = []
        for q in remaining_questions:
            t = asyncio.create_task(wrapped_eval(q))
            running_tasks.append(t)
            await asyncio.sleep(2.5)  # 2.5s stagger ensures 25s per-key interval across 10 API keys
        await asyncio.gather(*running_tasks)

    total_time = time.time() - start_time
    print(f"\n========================================================")
    print(f"EVALUATION COMPLETE IN {total_time:.1f}s ({total_time/60:.2f} min)")

    # Final ordered list
    results = [results_dict[orig["id"]] for orig in questions if orig["id"] in results_dict]

    # Save final results
    with open(RESULTS_PATH, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"Saved detailed results to {RESULTS_PATH}")

    # Calculate metrics
    metrics = calculate_metrics(results)
    with open(METRICS_PATH, "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)
    print(f"Saved metrics summary to {METRICS_PATH}")

    # Pretty print summary
    s = metrics["summary"]
    print("\n================== BENCHMARK METRICS SUMMARY ==================")
    print(f"Total Questions Evaluated:         {metrics['total_questions']}")
    print(f"1. LLM DSL Parse Success Rate:     {s['llm_parse_success_count']}/{metrics['total_questions']} ({s['llm_parse_success_rate_pct']}%)")
    print(f"2. DSL Syntax Compile Rate:        {s['dsl_compiler_success_count']}/{metrics['total_questions']} ({s['dsl_compiler_success_rate_pct']}%)")
    print(f"3. Geometry Engine Solve Rate:     {s['geometry_engine_solve_count']}/{metrics['total_questions']} ({s['geometry_engine_solve_rate_pct']}%)")
    print(f"4. Validator Strict Invariant Rate:{s['geometry_validator_pass_count']}/{metrics['total_questions']} ({s['geometry_validator_pass_rate_pct']}%)")
    print(f"5. End-to-End Strict Core Success: {s['full_pipeline_success_count']}/{metrics['total_questions']} ({s['full_pipeline_success_rate_pct']}%)")
    print(f"6. End-to-End Solved Coordinates:  {s['geometry_engine_solve_count']}/{metrics['total_questions']} ({s['solved_any_rate_pct']}%)")

    print("\n--- Latency Breakdown (Mean / Median / P95) ---")
    llm_lat = metrics["latency_stats"]["llm_parsing"]
    eng_lat = metrics["latency_stats"]["geometry_engine"]
    print(f"LLM Parsing (Gemma 4 26B): Mean={llm_lat.get('mean_ms')}ms | Median={llm_lat.get('median_ms')}ms | P95={llm_lat.get('p95_ms')}ms")
    print(f"GeometryEngine:             Mean={eng_lat.get('mean_ms')}ms | Median={eng_lat.get('median_ms')}ms | P95={eng_lat.get('p95_ms')}ms")

    print("\n--- By Topic ---")
    for top, stat in metrics["by_topic"].items():
        print(f"  * {top:<20}: Solve Rate={stat['solve_rate_pct']:5.1f}% | Full Strict Success={stat['success_rate_pct']:5.1f}% (N={stat['total']})")

    print("\n--- By Difficulty ---")
    for diff, stat in metrics["by_difficulty"].items():
        print(f"  * {diff:<10}: Solve Rate={stat['solve_rate_pct']:5.1f}% | Full Strict Success={stat['success_rate_pct']:5.1f}% (N={stat['total']})")

    print("\n--- Top Failure Categories ---")
    for cat, count in metrics["failure_categories"].items():
        print(f"  * {cat}: {count}")


if __name__ == "__main__":
    asyncio.run(main())