"""
AGoT Experiment Runner — Evaluate AGoT on benchmark datasets.

Compares AGoT against IO, CoT, CoT-SC, and ToT baselines.
Saves results to data/results/.
"""

import json
import os
import sys
import time
from pathlib import Path

# Add project root to path
_project_root = Path(__file__).resolve().parent.parent.parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

from tqdm import tqdm
from src.agot.agot import AGoTEngine
from src.baseline.benchmark import BaselineRunner, extract_answer, is_correct


def evaluate_agot(
    dataset: list,
    task_name: str = "gsm8k",
    model: str = "qwen2.5:32b",
    max_depth: int = 2,
    max_layers: int = 2,
    max_samples: int = -1,
    results_dir: str = "data/results",
) -> dict:
    """Evaluate AGoT on a dataset."""
    if max_samples > 0:
        dataset = dataset[:max_samples]

    results_dir = Path(results_dir)
    results_dir.mkdir(parents=True, exist_ok=True)

    engine = AGoTEngine({
        "model": model,
        "max_depth": max_depth,
        "max_layers": max_layers,
        "temperature": 0.3,
    }, task_name=task_name)

    results = []
    correct = 0

    print(f"\n{'='*60}")
    print(f"AGoT Evaluation — {model} on {task_name}")
    print(f"  depth={max_depth}, layers={max_layers}, samples={len(dataset)}")
    print(f"{'='*60}")

    for item in tqdm(dataset, desc="AGoT"):
        question = item["question"]
        expected = item["answer"]

        start = time.time()
        result = engine.solve(question, task_name)
        elapsed = time.time() - start

        predicted = result["answer"]
        c = is_correct(predicted, expected)
        if c:
            correct += 1

        results.append({
            "question": question[:100],
            "expected": expected,
            "predicted": predicted,
            "correct": c,
            "time_s": round(elapsed, 2),
            "stats": result["stats"],
        })

        if len(results) <= 2 or not c:
            status = "✓" if c else "✗"
            print(f"  [{status}] Pred: {predicted} | Exp: {expected} ({elapsed:.1f}s)")

    accuracy = correct / len(dataset) if dataset else 0
    summary = {
        "method": "AGoT",
        "model": model,
        "task": task_name,
        "config": {"max_depth": max_depth, "max_layers": max_layers},
        "samples": len(dataset),
        "correct": correct,
        "accuracy": round(accuracy, 4),
        "avg_time_s": round(sum(r["time_s"] for r in results) / len(results), 2) if results else 0,
    }

    # Save
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    filename = f"{task_name}_AGoT_d{max_depth}_l{max_layers}_{model.replace(':', '_')}_{timestamp}.json"
    filepath = results_dir / filename
    with open(filepath, "w") as f:
        json.dump({"summary": summary, "details": results}, f, indent=2, ensure_ascii=False)

    print(f"\nResults saved: {filepath}")
    print(f"  Accuracy: {summary['accuracy']:.2%} ({summary['correct']}/{summary['samples']})")
    print(f"  Avg time: {summary['avg_time_s']}s")

    return summary


# ── CLI ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Run AGoT evaluation")
    parser.add_argument("--model", default="qwen2.5:32b")
    parser.add_argument("--task", choices=["gsm8k", "math", "aime", "c_gsm8k", "cmath", "c_aime"], default="gsm8k")
    parser.add_argument("--samples", type=int, default=3, help="Number of test samples")
    parser.add_argument("--max-depth", type=int, default=2)
    parser.add_argument("--max-layers", type=int, default=2)
    parser.add_argument("--compare", action="store_true", help="Also run baselines for comparison")
    args = parser.parse_args()

    from src.samples import get_dataset
    dataset = get_dataset(args.task, args.samples)

    # Run AGoT
    agot_result = evaluate_agot(
        dataset, args.task, args.model,
        max_depth=args.max_depth, max_layers=args.max_layers,
    )

    if args.compare:
        print("\n--- Running Baselines for Comparison ---")
        runner = BaselineRunner(model=args.model)
        runner.run_all(dataset, task_name=args.task, methods=["IO", "CoT"])

    print("\nDone!")
