#!/usr/bin/env python3
"""Apply the frozen BabyAI B0 clean-selector rule."""

import argparse
import json
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_RESULTS = ROOT / "results" / "babyai_shared" / "b0"
DEFAULT_OUT = ROOT / "results" / "babyai_shared" / "b0_selector.json"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--results", type=Path, default=DEFAULT_RESULTS)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    args = parser.parse_args()
    rows = []
    for depth in (2, 4):
        for learning_rate in (0.01, 0.03):
            records = {}
            for condition in ("bp", "clean_kp"):
                path = args.results / (
                    f"d{depth}_lr{learning_rate}_{condition}.json")
                with open(path, encoding="utf-8") as handle:
                    records[condition] = json.load(handle)
            bp = records["bp"]
            kp = records["clean_kp"]
            bp_success = float(bp["rollout"]["success"])
            kp_success = float(kp["rollout"]["success"])
            bp_lesion = float(bp["mission_lesion_rollout"]["success"])
            checks = {
                "both_finite": bool(bp["finite"] and kp["finite"]),
                "bp_success_at_least_0p8": bp_success >= 0.8,
                "clean_kp_success_at_least_0p8": kp_success >= 0.8,
                "bp_mission_lesion_drop_at_least_0p2": (
                    bp_success - bp_lesion >= 0.2),
            }
            rows.append({
                "hidden_layers": depth,
                "learning_rate": learning_rate,
                "bp_rollout_success": bp_success,
                "clean_kp_rollout_success": kp_success,
                "bp_mission_lesion_success": bp_lesion,
                "bp_mission_lesion_drop": bp_success - bp_lesion,
                "bp_action_accuracy": float(bp["validation"]["accuracy"]),
                "clean_kp_action_accuracy": float(
                    kp["validation"]["accuracy"]),
                "eligible": all(checks.values()),
                "checks": checks,
                "source_files": [str(
                    (args.results / f"d{depth}_lr{learning_rate}_{condition}.json")
                    .relative_to(ROOT)) for condition in ("bp", "clean_kp")],
            })
    eligible = [row for row in rows if row["eligible"]]
    selected = max(eligible, key=lambda row: (
        row["clean_kp_rollout_success"], row["clean_kp_action_accuracy"],
        -row["hidden_layers"], -row["learning_rate"])) if eligible else None
    report = {
        "stage": "babyai_shared_b0_selector",
        "gate": "pass" if selected is not None else "fail",
        "selection_rule": (
            "highest clean-KP rollout success, then action accuracy, then "
            "fewer layers, then smaller learning rate, among eligible rows"),
        "candidates": rows,
        "selected": ({
            "hidden_layers": selected["hidden_layers"],
            "width": 256,
            "learning_rate": selected["learning_rate"],
            "context_gain": 1.0,
            "b1_epochs": 40,
            "b1_model_and_shuffle_seeds": [4101, 4102, 4103],
        } if selected is not None else None),
        "raw_or_sdil_results_read": False,
        "test_split_generated_or_read": False,
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as handle:
        json.dump(report, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()