#!/usr/bin/env python3
import argparse
import csv
import json
from pathlib import Path


CASES = {
    "baseline_0_1": {"group_size": 0, "num_in_group": 1},
    "group_8_num_1": {"group_size": 8, "num_in_group": 1},
    "group_8_num_2": {"group_size": 8, "num_in_group": 2},
}
DATASETS = ("gsm8k", "gpqa")


def load_result(results_dir, case, dataset):
    path = results_dir / case / f"{dataset}-result.json"
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise RuntimeError(f"unable to read accuracy result: {path}") from error
    if not isinstance(payload, dict):
        raise RuntimeError(f"accuracy result must be an object: {path}")
    return payload


def validate_result(payload, case, dataset, expected_settings):
    if payload.get("dataset") != dataset:
        raise RuntimeError(f"dataset metadata mismatch for {case}/{dataset}")
    if payload.get("offload_settings") != expected_settings:
        raise RuntimeError(f"offload settings mismatch for {case}/{dataset}")

    total = payload.get("total")
    correct = payload.get("correct")
    accuracy = payload.get("accuracy")
    if (
        isinstance(total, bool)
        or not isinstance(total, int)
        or total <= 0
        or isinstance(correct, bool)
        or not isinstance(correct, int)
        or correct < 0
        or correct > total
        or not isinstance(accuracy, (int, float))
    ):
        raise RuntimeError(f"invalid score metadata for {case}/{dataset}")
    if abs(float(accuracy) - correct / total) > 1e-12:
        raise RuntimeError(f"accuracy does not match correct/total for {case}/{dataset}")
    if not payload.get("dataset_sha256"):
        raise RuntimeError(f"missing dataset hash for {case}/{dataset}")


def verify_matrix(results_dir, max_accuracy_drop):
    if max_accuracy_drop < 0:
        raise ValueError("max_accuracy_drop must be non-negative")
    results_dir = Path(results_dir)
    rows = []

    for dataset in DATASETS:
        baseline = load_result(results_dir, "baseline_0_1", dataset)
        validate_result(
            baseline, "baseline_0_1", dataset, CASES["baseline_0_1"]
        )
        baseline_accuracy = float(baseline["accuracy"])
        baseline_hash = baseline["dataset_sha256"]
        baseline_total = baseline["total"]

        for case, expected_settings in CASES.items():
            result = baseline if case == "baseline_0_1" else load_result(
                results_dir, case, dataset
            )
            validate_result(result, case, dataset, expected_settings)
            if result["dataset_sha256"] != baseline_hash:
                raise RuntimeError(f"dataset hash mismatch for {case}/{dataset}")
            if result["total"] != baseline_total:
                raise RuntimeError(f"sample count mismatch for {case}/{dataset}")

            accuracy = float(result["accuracy"])
            change = accuracy - baseline_accuracy
            passed = change >= -max_accuracy_drop - 1e-12
            rows.append(
                {
                    "case": case,
                    "dataset": dataset,
                    "correct": result["correct"],
                    "total": result["total"],
                    "accuracy": accuracy,
                    "baseline_accuracy": baseline_accuracy,
                    "accuracy_change": change,
                    "passed": passed,
                    "dataset_sha256": result["dataset_sha256"],
                }
            )
            if not passed:
                raise RuntimeError(
                    f"accuracy drop for {case}/{dataset}: "
                    f"{change:.12f} < -{max_accuracy_drop:.12f}"
                )

    return {
        "baseline_case": "baseline_0_1",
        "max_accuracy_drop": max_accuracy_drop,
        "all_passed": all(row["passed"] for row in rows),
        "rows": rows,
    }


def write_summaries(summary, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    json_path = output_dir / "accuracy-summary.json"
    csv_path = output_dir / "accuracy-summary.csv"
    markdown_path = output_dir / "accuracy-summary.md"

    json_path.write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    fieldnames = list(summary["rows"][0])
    with csv_path.open("w", encoding="utf-8", newline="") as file:
        writer = csv.DictWriter(file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(summary["rows"])

    lines = [
        "# Accuracy Matrix",
        "",
        f"Maximum allowed accuracy drop: `{summary['max_accuracy_drop']}`",
        "",
        "| Case | Dataset | Correct | Total | Accuracy | Change | Passed |",
        "| --- | --- | ---: | ---: | ---: | ---: | --- |",
    ]
    for row in summary["rows"]:
        lines.append(
            f"| {row['case']} | {row['dataset']} | {row['correct']} | "
            f"{row['total']} | {row['accuracy']:.6f} | "
            f"{row['accuracy_change']:+.6f} | {row['passed']} |"
        )
    markdown_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return json_path, csv_path, markdown_path


def parse_args(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results-dir", type=Path, required=True)
    parser.add_argument("--max-accuracy-drop", type=float, default=0.0)
    args = parser.parse_args(argv)
    if args.max_accuracy_drop < 0:
        parser.error("--max-accuracy-drop must be non-negative")
    return args


def main(argv=None):
    args = parse_args(argv)
    summary = verify_matrix(args.results_dir, args.max_accuracy_drop)
    paths = write_summaries(summary, args.results_dir)
    print(json.dumps({"all_passed": True, "artifacts": [str(path) for path in paths]}))


if __name__ == "__main__":
    main()
