from __future__ import annotations

import argparse
import json
from pathlib import Path

from policyfaultbench.generalization_eval import write_generalization_report


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    report = write_generalization_report(args.root, args.output_dir)
    summary = {
        "policy_count": report["policy_count"],
        "corpus_count": report["corpus_count"],
        "offline_policy_scenario_evaluations": report[
            "offline_policy_scenario_evaluations"
        ],
        "reports": [
            {
                key: value
                for key, value in item.items()
                if key
                in {
                    "policy",
                    "corpus",
                    "reference_cases",
                    "operator_count",
                    "core_killed",
                    "expanded_killed",
                    "core_mutation_score",
                    "expanded_mutation_score",
                }
            }
            for item in report["reports"]
        ],
    }
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()
