#!/usr/bin/env python3
"""Summarize the frozen BabyAI GoToObj shared-feedback endpoint."""

import argparse
import json
from pathlib import Path
import statistics

import gymnasium as gym
import minigrid  # noqa: F401 - registers BabyAI environments


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_RESULTS = ROOT / "results" / "babyai_shared" / "b1"
DEFAULT_OUT = ROOT / "results" / "babyai_shared" / "b1_analysis.json"
SEEDS = (4101, 4102, 4103)
CONDITIONS = ("bp", "clean_kp", "raw_shared", "sdil")


def summarize(values):
    return {
        "values": values,
        "mean": statistics.mean(values),
        "sample_std": statistics.stdev(values),
    }


def object_counts(env_id, seeds):
    counts = []
    for seed in seeds:
        env = gym.make(env_id)
        try:
            env.reset(seed=seed)
            counts.append(sum(
                cell is not None and cell.type in {"ball", "box", "key"}
                for cell in env.unwrapped.grid.grid))
        finally:
            env.close()
    return counts


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()
    records = {}
    for seed in SEEDS:
        for condition in CONDITIONS:
            path = args.results / f"seed{seed}_{condition}.json"
            with open(path, encoding="utf-8") as handle:
                records[(seed, condition)] = json.load(handle)

    condition_summary = {}
    for condition in CONDITIONS:
        condition_records = [records[(seed, condition)] for seed in SEEDS]
        condition_summary[condition] = {
            "rollout_success_percent": summarize([
                100.0 * row["rollout"]["success"]
                for row in condition_records]),
            "expert_action_accuracy_percent": summarize([
                100.0 * row["validation"]["accuracy"]
                for row in condition_records]),
            "mission_lesion_rollout_success_percent": summarize([
                100.0 * row["mission_lesion_rollout"]["success"]
                for row in condition_records]),
            "training_wall_seconds": summarize([
                row["training"]["training_wall_seconds"]
                for row in condition_records]),
            "all_finite": all(row["finite"] for row in condition_records),
        }
    rollout_gaps = [100.0 * (
        records[(seed, "sdil")]["rollout"]["success"]
        - records[(seed, "raw_shared")]["rollout"]["success"])
                    for seed in SEEDS]
    action_gaps = [100.0 * (
        records[(seed, "sdil")]["validation"]["accuracy"]
        - records[(seed, "raw_shared")]["validation"]["accuracy"])
                   for seed in SEEDS]
    counts = object_counts("BabyAI-GoToObjS6-v1", range(256))
    task_required = min(counts) > 1
    checks = {
        "all_runs_finite": all(
            row["finite"] for row in records.values()),
        "bp_and_clean_kp_mean_rollout_at_least_80": (
            condition_summary["bp"]["rollout_success_percent"]["mean"] >= 80
            and condition_summary["clean_kp"]
            ["rollout_success_percent"]["mean"] >= 80),
        "sdil_rollout_above_raw_every_seed": all(
            value > 0 for value in rollout_gaps),
        "mission_structurally_required": task_required,
    }
    report = {
        "stage": "babyai_shared_b1_analysis",
        "gate": "pass" if all(checks.values()) else "fail",
        "checks": checks,
        "conditions": condition_summary,
        "paired_sdil_minus_raw": {
            "rollout_success_points": summarize(rollout_gaps),
            "expert_action_accuracy_points": summarize(action_gaps),
        },
        "task_structure_audit": {
            "env_id": "BabyAI-GoToObjS6-v1",
            "audited_episode_seeds": 256,
            "minimum_task_object_count": min(counts),
            "maximum_task_object_count": max(counts),
            "mission_structurally_required_to_identify_the_only_object": (
                task_required),
            "interpretation": (
                "The sole object is the target, so the mission does not select "
                "among alternatives. Mission lesion is an input ablation, not "
                "evidence that language is required by the task."),
        },
        "test_split_generated_or_read": False,
        "decision": (
            "Retain as an implementation diagnostic; do not use as the paper's "
            "task-required shared-feedback result. Move to PickupLoc."),
    }
    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()