#!/usr/bin/env python3
"""Run one untouched oral-B-v2 cold-start recovery confirmation cell."""
import argparse
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from experiments.bci_v2_recovery_run import (
    FIXED_CONFIG,
    PROTOCOL_PATH,
    finite_tree,
    provenance,
    require_parent_gates,
    run_cell,
    sha256,
    source_paths,
)


ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
R1_GATE_PATH = os.path.join(
    ROOT, "results", "bci_v2_recovery_dev_gate.json"
)
TASK_SEEDS = tuple(range(30, 36))
MODEL_SEEDS = tuple(range(5))


def require_r1_gate(r1):
    digests = {
        name: sha256(path)
        for name, path in source_paths().items()
    }
    if not (
        r1.get("protocol")
        == "oral_b_v2_cold_start_recovery_development_v1"
        and r1.get("status") == "passed"
        and r1.get("complete_grid") is True
        and r1.get("fixed_config") == FIXED_CONFIG
        and r1.get("confirmation_seeds_touched") is False
        and r1.get("recovery_confirmation_opened") is True
        and r1.get("review_score_after") == 7
        and r1.get("input_sha256") == digests
    ):
        raise ValueError(
            "recovery R2 requires the complete eligible R1 gate"
        )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--task-seed", type=int, choices=TASK_SEEDS, required=True
    )
    parser.add_argument(
        "--model-seed", type=int, choices=MODEL_SEEDS, required=True
    )
    parser.add_argument(
        "--r1-gate",
        default="results/bci_v2_recovery_dev_gate.json",
    )
    parser.add_argument(
        "--outdir", default="results/bci_v2_recovery_confirmation"
    )
    args = parser.parse_args()
    require_parent_gates()
    with open(args.r1_gate) as handle:
        r1 = json.load(handle)
    require_r1_gate(r1)
    source = provenance({
        "development_runner": os.path.join(
            ROOT, "experiments", "bci_v2_recovery_run.py"
        ),
        "runner": os.path.abspath(__file__),
        "r1_gate": os.path.abspath(args.r1_gate),
    })
    if (
        source["git_tracked_dirty"]
        or not all(source["tracked_inputs"].values())
    ):
        raise RuntimeError(
            "recovery R2 requires clean, tracked, frozen inputs"
        )
    cell = run_cell(
        args.task_seed,
        args.model_seed,
        split="untouched_confirmation",
        performance_seed_offset=520_000,
        challenge_seed_offset=530_000,
    )
    result = {
        "schema_version": 3,
        "protocol": {
            "name":
                "oral_b_v2_cold_start_recovery_confirmation_v1",
            "split": "untouched_confirmation",
            "training_task_seed": args.task_seed,
            "model_seed": args.model_seed,
            "fixed_config": FIXED_CONFIG,
            "no_further_selection": True,
            "confirmation_grid_size": (
                len(TASK_SEEDS) * len(MODEL_SEEDS)
            ),
            "protocol_sha256": sha256(PROTOCOL_PATH),
            "r1_gate_sha256": sha256(args.r1_gate),
        },
        "args": vars(args),
        "provenance": source,
        **cell,
    }
    result["finite"] = finite_tree(result)
    if not result["finite"]:
        raise RuntimeError("non-finite recovery confirmation record")
    os.makedirs(args.outdir, exist_ok=True)
    path = os.path.join(
        args.outdir,
        (
            f"bci_v2_recovery_confirm_t{args.task_seed}"
            f"_m{args.model_seed}.json"
        ),
    )
    if os.path.exists(path):
        raise FileExistsError(f"refusing to overwrite {path}")
    with open(path, "w") as handle:
        json.dump(result, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps({
        "path": path,
        "intact_final": result["conditions"]["intact"]["final_success"],
        "challenge_success_fraction": result["signatures"][
            "challenge_success_fraction"
        ],
        "terminal_outcome_accuracy": result["signatures"][
            "terminal_residual_outcome_balanced_acc"
        ],
        "finite": result["finite"],
    }, indent=2))


if __name__ == "__main__":
    main()