eval/eval_e2e_gemma.py

"""
End-to-End Evaluation Benchmark: Text -> Gemma 4 31B (Low Thinking) -> DSL v6.0 -> Solve -> Render.
Tests 5 national exam geometry problems across diverse categories.
"""

import asyncio
import json
import os
import sys
import time
from typing import Dict, Any, List

# Ensure backend root is in sys.path
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
from geometry_render.renderer import RendererAgent

BENCHMARK_CASES = [
    {
        "id": "geo_2016_essay_qessay_V",
        "name": "Oxyz Projection & Plane (THPT QG 2016)",
        "category": "oxyz",
        "text": (
            "Trong không gian với hệ tọa độ Oxyz cho ba điểm A(3; 2; -2), B(1; 0; 1) và C(2; -1; 3). "
            "Viết phương trình mặt phẳng đi qua A và vuông góc với đường thẳng BC. "
            "Tìm tọa độ hình chiếu vuông góc của A trên đường thẳng BC."
        ),
        "ground_truth_answer": "(P): x - y + 2z + 3 = 0; H(0; 1; -1)",
        "expected_points": ["A", "B", "C", "H"],
    },
    {
        "id": "geo_2017_123_q16",
        "name": "Regular Square Pyramid (THPT QG 2017)",
        "category": "solid_geometry (regular pyramid)",
        "text": (
            "Cho khối chóp tứ giác đều S.ABCD có cạnh đáy bằng a, cạnh bên gấp hai lần cạnh đáy (SA = 2a). "
            "Tính thể tích V của khối chóp đã cho."
        ),
        "ground_truth_answer": "V = sqrt(14)/6 * a^3",
        "expected_points": ["S", "A", "B", "C", "D"],
    },
    {
        "id": "geo_2016_essay_qessay_VII",
        "name": "Oblique Triangular Prism (THPT QG 2016)",
        "category": "solid_geometry (oblique prism)",
        "text": (
            "Cho lăng trụ ABC.A'B'C' có đáy ABC là tam giác vuông cân tại B, AC = 2a. "
            "Hình chiếu vuông góc của A' trên mặt phẳng (ABC) là trung điểm H của cạnh AC, "
            "đường thẳng A'B tạo với mặt phẳng (ABC) một góc 45°. "
            "Tính thể tích khối lăng trụ ABC.A'B'C'."
        ),
        "ground_truth_answer": "V = a^3 (A'H = a, S_ABC = a^2)",
        "expected_points": ["A", "B", "C", "A1", "B1", "C1", "H"],
    },
    {
        "id": "geo_2017_123_q18",
        "name": "Cylinder Round Body (THPT QG 2017)",
        "category": "round_bodies (cylinder)",
        "text": (
            "Tính thể tích V của khối trụ có bán kính đáy r = 4 và chiều cao h = 4*sqrt(2)."
        ),
        "ground_truth_answer": "V = 64*sqrt(2)*pi",
        "expected_points": ["O", "O1"],
    },
    {
        "id": "geo_2018_101_q25",
        "name": "Pyramid SA Perpendicular to Base (THPT QG 2018)",
        "category": "solid_geometry (pyramid)",
        "text": (
            "Cho hình chóp S.ABC có đáy là tam giác vuông đỉnh B, AB = 1, BC = 2, "
            "SA vuông góc với mặt phẳng đáy (ABC) và SA = 2. "
            "Tính khoảng cách từ A đến mặt phẳng (SBC)."
        ),
        "ground_truth_answer": "d(A, (SBC)) = 2*sqrt(5)/5 * a",
        "expected_points": ["S", "A", "B", "C"],
    },
]


async def run_single_case(case: Dict[str, Any], parser_agent: GeometryParserAgent) -> Dict[str, Any]:
    cid = case["id"]
    name = case["name"]
    text = case["text"]
    gt_ans = case.get("ground_truth_answer", "")
    print(f"\n========================================================")
    print(f"[*] RUNNING CASE: [{cid}] {name}")
    print(f"    Category: {case['category']}")
    print(f"    Đáp án chuẩn: {gt_ans}")
    print(f"    Text: {text}")

    report: Dict[str, Any] = {
        "id": cid,
        "name": name,
        "category": case["category"],
        "input_text": text,
        "ground_truth_answer": gt_ans,
        "stages": {},
        "status": "FAILED",
    }

    # Stage 1: Text -> LLM (Gemma 4 31B low thinking) -> DSL
    t0 = time.time()
    try:
        parsed_agent_out = await parser_agent.process(text)
        t_llm = (time.time() - t0) * 1000
        dsl_text = parsed_agent_out.get("geometry_dsl", "")
        shape_type = parsed_agent_out.get("type", "")
        print(f"  [1] LLM Success ({t_llm:.1f}ms): type={shape_type}")
        print(f"      Generated DSL:\n{dsl_text}")
        report["stages"]["llm_parsing"] = {
            "status": "OK",
            "latency_ms": round(t_llm, 1),
            "type": shape_type,
            "dsl": dsl_text,
            "entities": parsed_agent_out.get("entities", []),
        }
    except Exception as e:
        t_llm = (time.time() - t0) * 1000
        print(f"  [1] LLM FAILED ({t_llm:.1f}ms): {e}")
        report["stages"]["llm_parsing"] = {"status": "ERROR", "error": str(e), "latency_ms": round(t_llm, 1)}
        return report

    # Stage 2: DSL -> AST (DSLParser v6.0)
    t0 = time.time()
    try:
        dsl_parser = DSLParser()
        points, constraints, is_3d = dsl_parser.parse(dsl_text)
        t_dsl = (time.time() - t0) * 1000
        print(f"  [2] DSLParser Success ({t_dsl:.1f}ms): {len(points)} points, {len(constraints)} constraints, is_3d={is_3d}")
        report["stages"]["dsl_compiler"] = {
            "status": "OK",
            "latency_ms": round(t_dsl, 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_dsl = (time.time() - t0) * 1000
        print(f"  [2] DSLParser FAILED ({t_dsl:.1f}ms): {e}")
        report["stages"]["dsl_compiler"] = {"status": "ERROR", "error": str(e), "latency_ms": round(t_dsl, 1)}
        return report

    # Stage 3: AST -> Engine Solving (GeometryEngine)
    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("GeometryEngine returned empty coordinates or failed solve.")
        coords = solved_ir.get("coordinates", {})
        print(f"  [3] GeometryEngine Success ({t_solve:.1f}ms): {len(coords)} coordinates solved")
        for pid, pos in list(coords.items())[:6]:
            print(f"        {pid}: {[round(v, 4) for v in pos]}")
        report["stages"]["geometry_engine"] = {
            "status": "OK",
            "latency_ms": round(t_solve, 1),
            "coordinates": {k: [round(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
        print(f"  [3] GeometryEngine FAILED ({t_solve:.1f}ms): {e}")
        report["stages"]["geometry_engine"] = {"status": "ERROR", "error": str(e), "latency_ms": round(t_solve, 1)}
        return report

    # Stage 4: Validation (GeometryValidator)
    validator = GeometryValidator()
    val_res = validator.validate(solved_ir, constraints, is_3d=is_3d)
    diag = val_res.to_dict()
    print(f"  [4] Validator: is_valid={diag['is_valid']}, checked={diag['checked_count']}, errors={len(diag['errors'])}")
    report["stages"]["validation"] = {
        "status": "OK" if diag["is_valid"] else "WARNING",
        "diagnostics": diag,
    }

    # Stage 5: Solving IR -> Renderer (Manim Script & Video)
    t0 = time.time()
    try:
        renderer = RendererAgent()
        manim_script = renderer.generate_manim_script(solved_ir)
        job_id = f"e2e_gemma_{cid}"
        video_path = renderer.run_manim(manim_script, job_id)
        t_render = (time.time() - t0) * 1000
        print(f"  [5] Renderer Success ({t_render:.1f}ms): video={video_path}")
        report["stages"]["renderer"] = {
            "status": "OK" if video_path else "SCRIPT_ONLY",
            "latency_ms": round(t_render, 1),
            "video_path": video_path,
            "script_length_lines": len(manim_script.splitlines()),
        }
    except Exception as e:
        t_render = (time.time() - t0) * 1000
        print(f"  [5] Renderer FAILED ({t_render:.1f}ms): {e}")
        report["stages"]["renderer"] = {"status": "ERROR", "error": str(e), "latency_ms": round(t_render, 1)}

    report["status"] = "SUCCESS"
    report["total_latency_ms"] = round(
        report["stages"]["llm_parsing"]["latency_ms"]
        + report["stages"]["dsl_compiler"]["latency_ms"]
        + report["stages"]["geometry_engine"]["latency_ms"]
        + report["stages"]["renderer"].get("latency_ms", 0),
        1
    )
    return report


async def main():
    print("================================================================")
    print("STARTING E2E BENCHMARK EVALUATION (GEMMA 4 31B LOW THINKING)")
    print("================================================================")

    parser_agent = GeometryParserAgent()
    results = []

    for case in BENCHMARK_CASES:
        res = await run_single_case(case, parser_agent)
        results.append(res)

    out_path = os.path.join(os.path.dirname(__file__), "benchmark_e2e_gemma_results.json")
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)

    print(f"\n================================================================")
    print(f"COMPLETED ALL 5 CASES. Results written to: {out_path}")
    print(f"================================================================")


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