#!/usr/bin/env python3
"""Apply the frozen PickupLoc history selector."""

import argparse
import json
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_RESULTS = ROOT / "results" / "babyai_shared" / "pickup_p1"
DEFAULT_OUT = ROOT / "results" / "babyai_shared" / "pickup_p1_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()
    candidates = []
    for history_steps in (4, 8):
        records = {}
        for condition in ("bp", "clean_kp"):
            path = args.results / f"h{history_steps}_{condition}.json"
            with open(path, encoding="utf-8") as handle:
                records[condition] = json.load(handle)
        bp, kp = records["bp"], records["clean_kp"]
        bp_success = 100.0 * bp["rollout"]["success"]
        kp_success = 100.0 * kp["rollout"]["success"]
        lesion = 100.0 * bp["mission_lesion_rollout"]["success"]
        checks = {
            "both_finite": bool(bp["finite"] and kp["finite"]),
            "bp_success_at_least_70": bp_success >= 70.0,
            "clean_kp_success_at_least_60": kp_success >= 60.0,
            "bp_mission_lesion_drop_at_least_20": bp_success - lesion >= 20.0,
        }
        candidates.append({
            "history_steps": history_steps,
            "bp_rollout_success_percent": bp_success,
            "clean_kp_rollout_success_percent": kp_success,
            "bp_mission_lesion_success_percent": lesion,
            "bp_mission_lesion_drop_points": bp_success - lesion,
            "bp_action_accuracy_percent": 100.0 * bp["validation"]["accuracy"],
            "clean_kp_action_accuracy_percent": (
                100.0 * kp["validation"]["accuracy"]),
            "eligible": all(checks.values()),
            "checks": checks,
        })
    eligible = [row for row in candidates if row["eligible"]]
    selected = max(eligible, key=lambda row: (
        row["clean_kp_rollout_success_percent"],
        row["clean_kp_action_accuracy_percent"],
        -row["history_steps"])) if eligible else None
    report = {
        "stage": "babyai_pickup_p1_history_selector",
        "gate": "pass" if selected is not None else "fail",
        "candidates": candidates,
        "selected": ({
            "history_steps": selected["history_steps"],
            "hidden_layers": 4,
            "width": 256,
            "learning_rate": 0.03,
            "context_gain": 1.0,
            "p2_epochs": 40,
            "p2_model_and_shuffle_seeds": [4101, 4102, 4103],
        } if selected is not None else None),
        "raw_or_sdil_results_run_or_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()