#!/usr/bin/env python3
"""Audit the frozen oral-B-v2 cold-start recovery development validation."""
import argparse
import glob
import hashlib
import json
import os

from experiments.bci_v2_run import CONDITIONS
from experiments.bci_v2_recovery_run import (
    CHALLENGE_EPISODES,
    FIXED_CONFIG,
    MODEL_SEEDS,
    TARGETS,
    TASK_SEEDS,
    build_recovery_config,
    source_paths,
)


def require(condition, message):
    if not condition:
        raise ValueError(message)


def sha256(path):
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def checks(row):
    conditions = row["conditions"]
    signatures = row["signatures"]
    intact = conditions["intact"]
    return {
        "final_success_at_least_0p70":
            intact["final_success"] >= 0.70,
        "learning_gain_at_least_0p10":
            intact["learning_gain"] >= 0.10,
        "fixed_role_gap_at_least_0p20": (
            intact["final_success"]
            - conditions["fixed_role"]["final_success"]
        ) >= 0.20,
        "oracle_deficit_at_most_0p10": (
            conditions["oracle_role"]["final_success"]
            - intact["final_success"]
        ) <= 0.10,
        "plasticity_lesion_retains_at_most_half_gain":
            conditions["plasticity_lesion"]["learning_gain"]
            <= 0.5 * intact["learning_gain"],
        "role_cosine_at_least_0p80":
            intact["role_cosine_after_training"] >= 0.80,
        "challenge_fraction_between_0p15_and_0p85":
            0.15 <= signatures["challenge_success_fraction"] <= 0.85,
        "terminal_outcome_accuracy_at_least_0p75":
            signatures[
                "terminal_residual_outcome_balanced_acc"
            ] >= 0.75,
        "terminal_role_separation_at_least_0p20":
            signatures[
                "terminal_role_aligned_outcome_separation"
            ] >= 0.20,
        "outcome_lesion_separation_drop_at_least_0p20":
            signatures[
                "terminal_outcome_separation_drop_under_acute_lesion"
            ] >= 0.20,
        "critic_expectedness_at_least_0p05":
            signatures[
                "mean_critic_expectedness_contribution"
            ] >= 0.05,
        "critic_contribution_value_corr_at_least_0p95":
            signatures[
                "critic_contribution_value_prediction_corr"
            ] >= 0.95,
        "residual_soma_corr_at_most_0p10":
            signatures["mean_abs_residual_soma_corr"] <= 0.10,
        "raw_residual_corr_gap_at_least_0p20":
            signatures[
                "raw_minus_residual_abs_soma_corr"
            ] >= 0.20,
        "surrounding_event_accuracy_at_least_0p52":
            signatures[
                "surrounding_event_decoder_balanced_acc"
            ] >= 0.52,
        "decoder_distance_corr_at_least_0p02":
            signatures[
                "decoder_distance_residual_corr"
            ] >= 0.02,
        "sign_inversion_at_least_0p01":
            signatures[
                "causal_role_sign_inversion_index"
            ] >= 0.01,
        "velocity_advantage_at_least_0p05":
            signatures[
                "velocity_minus_error_abs_cv_corr"
            ] >= 0.05,
    }


def validate(row, path, task_seed, digests):
    require(row.get("schema_version") == 3, f"{path}: schema")
    require(
        row.get("args", {}).get("task_seed") == task_seed
        and row["args"].get("model_seed") in MODEL_SEEDS,
        f"{path}: seeds",
    )
    protocol = row.get("protocol", {})
    require(
        protocol.get("name")
        == "oral_b_v2_cold_start_recovery_development_v1"
        and protocol.get("selection_split") == "development_validation"
        and protocol.get("hyperparameter_selection") is False
        and protocol.get("confirmation_seeds_touched") is False
        and protocol.get("fixed_config") == FIXED_CONFIG,
        f"{path}: protocol",
    )
    require(
        protocol.get("protocol_sha256") == digests["protocol"]
        and protocol.get("d4_gate_sha256") == digests["d4_gate"]
        and protocol.get("old_r2_gate_sha256")
        == digests["old_r2_gate"]
        and protocol.get("failed_v2_gate_sha256")
        == digests["failed_v2_gate"],
        f"{path}: parent digest",
    )
    provenance = row.get("provenance", {})
    require(
        provenance.get("git_tracked_dirty") is False
        and provenance.get("tracked_inputs")
        and all(provenance["tracked_inputs"].values())
        and set(provenance["tracked_inputs"]) == set(digests)
        and provenance.get("input_sha256") == digests,
        f"{path}: provenance",
    )
    require(
        row.get("config") == vars(build_recovery_config()),
        f"{path}: config",
    )
    require(
        row.get("split") == "development"
        and row.get("finite") is True,
        f"{path}: finite/split",
    )
    require(
        set(row.get("conditions", {})) == set(CONDITIONS),
        f"{path}: conditions",
    )
    for name in CONDITIONS:
        warmup = row["warmup"][name]
        require(
            warmup["batches"] == 100
            and warmup["examples"] == 6400
            and warmup["instruction_present"] is False
            and warmup["role_cursor_scalar_observations"] == 12800
            and warmup["predictor_max_abs_error"] <= 1e-5,
            f"{path}: warmup {name}",
        )
        cost = row["conditions"][name]["cost"]
        require(
            len(row["conditions"][name]["daily_success"]) == 14
            and cost["maximum_state_episode_steps"] == 25088
            and cost["cursor_scalar_observations"]
            == 2 * cost["role_probe_examples"]
            and cost["task_loss_queries"] == 0
            and cost["reverse_mode_calls"] == 0,
            f"{path}: condition {name}",
        )
    challenge = row["assays"]["challenge"]
    require(
        challenge["targets"] == list(TARGETS)
        and challenge["episodes_per_target"] == CHALLENGE_EPISODES
        and challenge["selection_over_targets"] is False
        and challenge["maximum_steps_per_episode"] == 28,
        f"{path}: target ladder",
    )
    require(
        row["signatures"]["challenge_episodes"]
        == len(TARGETS) * CHALLENGE_EPISODES,
        f"{path}: challenge count",
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--results", default="results/bci_v2_recovery_dev"
    )
    parser.add_argument(
        "--out", default="results/bci_v2_recovery_dev_gate.json"
    )
    args = parser.parse_args()
    digests = {
        name: sha256(path)
        for name, path in source_paths().items()
    }
    expected = {
        f"bci_v2_recovery_t{task_seed}_m0.json"
        for task_seed in TASK_SEEDS
    }
    observed = {
        os.path.basename(path)
        for path in glob.glob(os.path.join(args.results, "*.json"))
    }
    require(
        observed == expected,
        (
            f"recovery grid drift: missing={sorted(expected-observed)}, "
            f"extra={sorted(observed-expected)}"
        ),
    )
    records = {}
    commits = set()
    source_sha256 = {}
    for task_seed in TASK_SEEDS:
        path = os.path.join(
            args.results,
            f"bci_v2_recovery_t{task_seed}_m0.json",
        )
        with open(path) as handle:
            row = json.load(handle)
        validate(row, path, task_seed, digests)
        records[task_seed] = row
        commits.add(row["provenance"]["git_commit"])
        source_sha256[path] = sha256(path)
    require(len(commits) == 1, "recovery source revision drift")
    checks_by_seed = {
        str(seed): checks(records[seed]) for seed in TASK_SEEDS
    }
    passed = all(
        value for seed_checks in checks_by_seed.values()
        for value in seed_checks.values()
    )
    output = {
        "protocol":
            "oral_b_v2_cold_start_recovery_development_v1",
        "status": "passed" if passed else "failed",
        "complete_grid": True,
        "fixed_config": FIXED_CONFIG,
        "checks_by_task_seed": checks_by_seed,
        "intact_final_by_task_seed": [
            records[seed]["conditions"]["intact"]["final_success"]
            for seed in TASK_SEEDS
        ],
        "challenge_success_fraction_by_task_seed": [
            records[seed]["signatures"]["challenge_success_fraction"]
            for seed in TASK_SEEDS
        ],
        "terminal_outcome_accuracy_by_task_seed": [
            records[seed]["signatures"][
                "terminal_residual_outcome_balanced_acc"
            ]
            for seed in TASK_SEEDS
        ],
        "confirmation_seeds_touched": False,
        "recovery_confirmation_opened": passed,
        "source_commit": next(iter(commits)),
        "source_sha256": source_sha256,
        "input_sha256": digests,
        "review_score_before": 7,
        "review_score_after": 7,
        "score_change_rule": (
            "development validation never changes the formal score"
        ),
    }
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    if os.path.exists(args.out):
        with open(args.out) as handle:
            existing = json.load(handle)
        require(
            existing == output,
            "existing recovery development gate differs",
        )
    else:
        with open(args.out, "w") as handle:
            json.dump(output, handle, indent=2, sort_keys=True)
            handle.write("\n")
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()