#!/usr/bin/env python3
"""
Evaluate baseline query-translation prediction files against gold UPPAAL models.

Each prediction file contains exactly 100 plain-text formulas (one per line),
ordered S01-Q1 … S01-Q5, S02-Q1 … S20-Q5.

For each formula the script:
  1. Runs verifyta against the matching gold model (datasets/gold_models/S{id:02d}/model.xml).
  2. Records the T/F verdict (or '?' on compile failure).
  3. Compares against the expected answer from datasets/query_translation_eval.csv.

Outputs:
  artifacts/baselines/<name>/trans_query_eval.csv   — per-query detail
  artifacts/results/rq2_baseline_trans_query_summary.txt — per-spec + overall table

Usage
-----
  python scripts/eval_baselines_trans_query.py
  python scripts/eval_baselines_trans_query.py --baselines gpt claude
"""
from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from frame.pipeline.model_checking_pipeline import ModelCheckingPipeline
from frame.pipeline.nl_uppaal_query import parse_verifyta_text_verdict

EVAL_CSV   = ROOT / "datasets" / "query_translation_eval.csv"
GOLD_DIR   = ROOT / "datasets" / "gold_models"
BASELINES  = ROOT / "artifacts" / "baselines"
RESULTS    = ROOT / "artifacts" / "results"

FIELDNAMES = ["spec_id", "nl_query", "uppaal_query", "status", "verdict", "expected", "correct"]

ALL_BASELINES = ["gpt", "claude", "grok"]


def load_eval_csv() -> list[dict]:
    with EVAL_CSV.open(encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


def load_predictions(path: Path) -> list[str]:
    """Read up to 100 non-empty lines from a prediction file."""
    lines = []
    for ln in path.read_text(encoding="utf-8", errors="replace").splitlines():
        ln = ln.strip()
        if ln:
            lines.append(ln)
        if len(lines) == 100:
            break
    return lines


def verdict_char(v: str) -> str:
    return {"satisfied": "T", "not_satisfied": "F"}.get(v, "?")


def evaluate(name: str, preds: list[str], eval_rows: list[dict]) -> list[dict]:
    checkers: dict[int, ModelCheckingPipeline] = {}
    results: list[dict] = []

    for idx, row in enumerate(eval_rows):
        sid = int(row["spec_id"])
        nl  = row["nl_query"].strip()
        exp = row["answer"].strip().upper()

        formula = preds[idx].strip() if idx < len(preds) else ""

        if formula:
            if sid not in checkers:
                xml = GOLD_DIR / f"S{sid:02d}" / "model.xml"
                checkers[sid] = ModelCheckingPipeline(str(xml))
            res, _, errs = checkers[sid].verify(formula)
            v = parse_verifyta_text_verdict(res, errors=errs)
            verdict = verdict_char(v)
        else:
            verdict = "?"

        correct = "Y" if verdict == exp else "N"
        results.append({
            "spec_id":      sid,
            "nl_query":     nl,
            "uppaal_query": formula,
            "status":       "ok" if verdict in ("T", "F") else "compile_fail",
            "verdict":      verdict,
            "expected":     exp,
            "correct":      correct,
        })

    return results


def write_csv(out_path: Path, rows: list[dict]) -> None:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=FIELDNAMES)
        w.writeheader()
        w.writerows(rows)


def summarise(name: str, rows: list[dict]) -> dict:
    total       = len(rows)
    correct     = sum(1 for r in rows if r["correct"] == "Y")
    compile_ok  = sum(1 for r in rows if r["verdict"] in ("T", "F"))
    return {
        "name": name,
        "total": total,
        "correct": correct,
        "compile_ok": compile_ok,
        "qar": 100 * correct // total if total else 0,
        "qcr": 100 * compile_ok // total if total else 0,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--baselines", nargs="+", default=ALL_BASELINES,
                    help="Which baselines to evaluate (default: gpt claude grok)")
    ap.add_argument("--pred-file", default="trans_query.txt",
                    help="Prediction filename inside each baseline folder (default: trans_query.txt)")
    args = ap.parse_args()

    eval_rows = load_eval_csv()
    if len(eval_rows) != 100:
        print(f"WARNING: eval CSV has {len(eval_rows)} rows (expected 100)", file=sys.stderr)

    RESULTS.mkdir(parents=True, exist_ok=True)
    summaries: list[dict] = []

    for name in args.baselines:
        pred_path = BASELINES / name / args.pred_file
        if not pred_path.is_file():
            print(f"[{name}] prediction file not found: {pred_path} — skipping")
            continue

        preds = load_predictions(pred_path)
        print(f"[{name}] loaded {len(preds)} predictions from {pred_path.name}")

        rows = evaluate(name, preds, eval_rows)
        out_csv = BASELINES / name / "trans_query_eval.csv"
        write_csv(out_csv, rows)

        s = summarise(name, rows)
        summaries.append(s)

        # Per-spec breakdown
        print(f"\n  {'Spec':<6} {'Score':>5}   Q1 Q2 Q3 Q4 Q5")
        print(f"  {'-'*38}")
        for sid in range(1, 21):
            spec_rows = [r for r in rows if r["spec_id"] == sid]
            marks = [("+" if r["correct"] == "Y" else "x") for r in spec_rows]
            sc = marks.count("+")
            print(f"  S{sid:02d}    {sc}/5     " + "  ".join(marks))
        print(f"  {'-'*38}")
        print(f"  Total  {s['correct']}/{s['total']}  QAR={s['qar']}%  QCR={s['qcr']}%")
        print(f"  Saved  -> {out_csv}\n")

    # Overall comparison table
    print("\n" + "=" * 55)
    print(f"{'Baseline':<10} {'Correct':>8} {'QAR':>6} {'QCR':>6}")
    print("-" * 55)
    for s in summaries:
        print(f"{s['name']:<10} {s['correct']:>4}/{s['total']:<4}  {s['qar']:>4}%  {s['qcr']:>4}%")

    # Also include ours if available
    ours_csv = ROOT / "artifacts" / "ours" / "trans_query.csv"
    if ours_csv.is_file():
        ours_rows = list(csv.DictReader(ours_csv.open(encoding="utf-8", newline="")))
        if "correct" in (ours_rows[0] if ours_rows else {}):
            oc = sum(1 for r in ours_rows if r["correct"] == "Y")
            ot = len(ours_rows)
            oqcr = sum(1 for r in ours_rows if r.get("verdict") in ("T", "F"))
            print(f"{'ours':<10} {oc:>4}/{ot:<4}  {100*oc//ot if ot else 0:>4}%  {100*oqcr//ot if ot else 0:>4}%")

    print("=" * 55)

    # Write summary text file
    summary_path = RESULTS / "rq2_baseline_trans_query_summary.txt"
    with summary_path.open("w", encoding="utf-8") as f:
        f.write(f"{'Baseline':<10} {'Correct':>8} {'QAR':>6} {'QCR':>6}\n")
        f.write("-" * 35 + "\n")
        for s in summaries:
            f.write(f"{s['name']:<10} {s['correct']:>4}/{s['total']:<4}  {s['qar']:>4}%  {s['qcr']:>4}%\n")
        if ours_csv.is_file() and ours_rows and "correct" in ours_rows[0]:
            f.write(f"{'ours':<10} {oc:>4}/{ot:<4}  {100*oc//ot if ot else 0:>4}%  {100*oqcr//ot if ot else 0:>4}%\n")
    print(f"\nSummary saved -> {summary_path}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())