#!/usr/bin/env bash
set -euo pipefail

if [[ $# -ne 6 ]]; then
  echo "usage: $0 <run-dir> <step> <gpu> <port> <arm> <output-root>" >&2
  exit 64
fi

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
runtime="${OPTR_RUNTIME_ROOT:?set OPTR_RUNTIME_ROOT to the runtime environment root}"
run_dir=$1
step=$2
gpu=$3
port=$4
arm=$5
output_root=$6

for policy_seed in 518957 765353 686765 516408; do
  "${script_dir}/evaluate_bridge_seed96_no_video.sh" \
    "${run_dir}" "${step}" "${gpu}" "${port}" "${policy_seed}" \
    "${arm}" "${output_root}"
done

"${runtime}/.venv/bin/python" - "${output_root}" "${step}" "${arm}" <<'PY'
import json
import os
import sys
from pathlib import Path

root = Path(sys.argv[1])
step = int(sys.argv[2])
arm = sys.argv[3]
seeds = [518957, 765353, 686765, 516408]
tasks = ("stack", "carrot", "spoon", "eggplant")
seed_rows = {}
task_rows = {task: {"success_count": 0, "rollout_count": 0} for task in tasks}
for seed in seeds:
    summary = json.loads(
        (root / f"policy_seed_{seed}" / f"step_{step}" / "four_task_summary.json")
        .read_text(encoding="utf-8")
    )
    seed_rows[str(seed)] = {
        "success_count": summary["total_success"],
        "rollout_count": summary["total_rollouts"],
        "strict_success_rate": summary["micro_success_rate"],
    }
    for task, row in summary["tasks"].items():
        task_rows[task]["success_count"] += row["success_count"]
        task_rows[task]["rollout_count"] += row["rollout_count"]
for row in task_rows.values():
    row["strict_success_rate"] = row["success_count"] / row["rollout_count"]
successes = sum(row["success_count"] for row in seed_rows.values())
rollouts = sum(row["rollout_count"] for row in seed_rows.values())
result = {
    "schema_version": 1,
    "status": "complete_no_video_four_seed_four_task384",
    "arm": arm,
    "checkpoint_step": step,
    "policy_seeds": seeds,
    "episodes_per_seed_task": 24,
    "tasks": task_rows,
    "seeds": seed_rows,
    "total_success": successes,
    "total_rollouts": rollouts,
    "strict_success_rate": successes / rollouts,
}
target = root / f"step_{step}_four_seed_four_task384.json"
temporary = target.with_suffix(f".json.tmp.{os.getpid()}")
temporary.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
os.replace(temporary, target)
print(json.dumps(result, sort_keys=True))
PY