#!/usr/bin/env python3
"""Render the audited oral-B-v2 confirmation figure from all 30 records."""
import argparse
import hashlib
import json
import math
import os
import statistics

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np


TASK_SEEDS = tuple(range(30, 36))
MODEL_SEEDS = tuple(range(5))
PROTOCOL = "oral_b_v2_calibrated_recovery_confirmation_v1"
BLUE = "#0072B2"
LIGHT_BLUE = "#56B4E9"
GREEN = "#009E73"
ORANGE = "#E69F00"
RED = "#D55E00"
GRAY = "#5A5A5A"
LIGHT_GRAY = "#B7B7B7"
PDF_METADATA = {
    "Creator": "SDIL audited oral-B-v2 figure pipeline",
    "Producer": "Matplotlib",
    "CreationDate": None,
    "ModDate": None,
}


def load_json(path):
    with open(path) as handle:
        return json.load(handle)


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 require(condition, message):
    if not condition:
        raise RuntimeError(message)


def finite_leaves(value):
    if isinstance(value, bool) or value is None:
        return True
    if isinstance(value, (int, float)):
        return math.isfinite(float(value))
    if isinstance(value, dict):
        return all(finite_leaves(child) for child in value.values())
    if isinstance(value, (list, tuple)):
        return all(finite_leaves(child) for child in value)
    return True


def all_true(value):
    if isinstance(value, dict):
        return all(all_true(child) for child in value.values())
    return value is True


def load_panel(input_dir, gate_path):
    gate = load_json(gate_path)
    require(gate.get("protocol") == PROTOCOL, "oral-B-v2 protocol drift")
    require(gate.get("status") == "passed", "oral-B-v2 gate is not passed")
    require(gate.get("review_score_after") == 8, "oral-B-v2 score drift")
    require(all_true(gate.get("checks", {})), "oral-B-v2 has false checks")
    records = {}
    paths = []
    for task_seed in TASK_SEEDS:
        for model_seed in MODEL_SEEDS:
            name = (
                f"bci_v2_calibrated_confirm_t{task_seed}"
                f"_m{model_seed}.json"
            )
            path = os.path.join(input_dir, name)
            require(os.path.isfile(path), f"missing oral-B-v2 record: {path}")
            row = load_json(path)
            paths.append(path)
            require(row.get("schema_version") == 4, f"{path}: schema")
            require(
                row.get("protocol", {}).get("name") == PROTOCOL
                and row["protocol"].get("split")
                == "untouched_confirmation"
                and row.get("split") == "untouched_confirmation",
                f"{path}: split/protocol",
            )
            require(
                row.get("args", {}).get("task_seed") == task_seed
                and row["args"].get("model_seed") == model_seed,
                f"{path}: seed",
            )
            require(
                row.get("provenance", {}).get("git_tracked_dirty") is False
                and row["provenance"].get("git_commit")
                == gate.get("source_commit"),
                f"{path}: provenance",
            )
            require(
                row.get("finite") is True and finite_leaves(row),
                f"{path}: finite",
            )
            records[(task_seed, model_seed)] = row
    observed = sorted(
        name for name in os.listdir(input_dir) if name.endswith(".json")
    )
    expected = sorted(
        (
            f"bci_v2_calibrated_confirm_t{task_seed}"
            f"_m{model_seed}.json"
        )
        for task_seed in TASK_SEEDS
        for model_seed in MODEL_SEEDS
    )
    require(observed == expected, "oral-B-v2 record-set drift")
    return gate, records, paths


def task_cluster_array(records, getter):
    return np.asarray([
        np.mean([
            getter(records[(task_seed, model_seed)])
            for model_seed in MODEL_SEEDS
        ], axis=0)
        for task_seed in TASK_SEEDS
    ], dtype=float)


def style_axis(axis):
    axis.spines["top"].set_visible(False)
    axis.spines["right"].set_visible(False)
    axis.tick_params(width=1.0, length=4)
    axis.grid(axis="y", color="#E7E7E7", linewidth=0.8, zorder=0)


def panel_label(axis, label):
    axis.text(
        -0.14,
        1.08,
        label,
        transform=axis.transAxes,
        fontsize=15,
        fontweight="bold",
        va="top",
    )


def plot_learning(axis, records):
    conditions = (
        ("intact", "SDIL + TD critic", BLUE),
        ("critic_training_lesion", "critic lesion", ORANGE),
        ("outcome_training_lesion", "outcome lesion", GREEN),
        ("fixed_role", "fixed role", LIGHT_GRAY),
    )
    days = np.arange(1, 15)
    for condition, label, color in conditions:
        clustered = task_cluster_array(
            records,
            lambda row: np.asarray(
                row["conditions"][condition]["daily_success"]
            ) * 100,
        )
        mean = clustered.mean(0)
        interval = 1.96 * clustered.std(0, ddof=1) / np.sqrt(len(TASK_SEEDS))
        axis.plot(days, mean, color=color, linewidth=2.0, label=label)
        axis.fill_between(
            days, mean - interval, mean + interval,
            color=color, alpha=0.12, linewidth=0,
        )
    axis.set_xlim(1, 14)
    axis.set_ylim(-3, 104)
    axis.set_xticks([1, 4, 7, 10, 14])
    axis.set_xlabel("Training day")
    axis.set_ylabel("Successful episodes (%)")
    axis.set_title(
        "Local actor–critic learns",
        loc="left",
        fontsize=11.5,
        fontweight="bold",
        pad=10,
    )
    axis.legend(frameon=False, fontsize=8, loc="upper left")
    style_axis(axis)
    panel_label(axis, "a")


def plot_residualization(axis, records):
    raw = task_cluster_array(
        records,
        lambda row: row["signatures"]["mean_abs_raw_soma_corr"],
    )
    residual = task_cluster_array(
        records,
        lambda row: row["signatures"]["mean_abs_residual_soma_corr"],
    )
    for index in range(len(TASK_SEEDS)):
        axis.plot(
            [0, 1],
            [raw[index], residual[index]],
            color=LIGHT_GRAY,
            linewidth=1.3,
            zorder=1,
        )
    axis.scatter(
        np.zeros_like(raw), raw, s=32, facecolor="white",
        edgecolor=GRAY, zorder=2,
    )
    axis.scatter(
        np.ones_like(residual), residual, s=34, facecolor=BLUE,
        edgecolor="white", zorder=3,
    )
    axis.scatter(
        [0, 1], [raw.mean(), residual.mean()],
        marker="D", s=65, c=[GRAY, BLUE], edgecolor="#111111",
        linewidth=0.7, zorder=4,
    )
    axis.axhline(0.10, color=RED, linestyle="--", linewidth=1.0)
    axis.text(
        0.98, 0.115, "frozen residual ceiling",
        ha="right", va="bottom", color=RED, fontsize=8,
    )
    axis.set_xlim(-0.35, 1.35)
    axis.set_ylim(0, 1.06)
    axis.set_xticks([0, 1], ["Raw apical", "Innovation"])
    axis.set_ylabel("Mean |soma correlation|")
    axis.set_title(
        "Prediction removes ordinary traffic",
        loc="left",
        fontsize=11.5,
        fontweight="bold",
        pad=10,
    )
    axis.text(
        0.05,
        0.48,
        f"raw $-$ residual = {raw.mean() - residual.mean():.3f}\n"
        f"surrounding decoder = "
        f"{100 * task_cluster_array(records, lambda row: row['signatures']['surrounding_event_decoder_balanced_acc']).mean():.2f}%",
        transform=axis.transAxes,
        fontsize=8.5,
        bbox={
            "boxstyle": "round,pad=0.35",
            "facecolor": "white",
            "edgecolor": "#D6D6D6",
        },
    )
    style_axis(axis)
    panel_label(axis, "b")


def plot_psychometric(axis, records):
    quantiles = np.asarray([0.20, 0.35, 0.50, 0.65, 0.80])

    def success_curve(row):
        targets = row["assays"]["challenge"]["targets"]
        values = row["signatures"]["target_success_fraction"]
        return np.asarray([values[str(target)] for target in targets]) * 100

    clustered = task_cluster_array(records, success_curve)
    mean = clustered.mean(0)
    interval = 1.96 * clustered.std(0, ddof=1) / np.sqrt(len(TASK_SEEDS))
    expected = (1.0 - quantiles) * 100
    axis.plot(
        quantiles, expected, color=LIGHT_GRAY, linestyle="--",
        linewidth=1.5, label="calibration expectation",
    )
    axis.plot(
        quantiles, mean, color=GREEN, marker="o", linewidth=2.0,
        label="independent assay",
    )
    axis.fill_between(
        quantiles, mean - interval, mean + interval,
        color=GREEN, alpha=0.15, linewidth=0,
    )
    axis.axhline(50, color="#DDDDDD", linewidth=0.9)
    axis.set_xlim(0.16, 0.84)
    axis.set_ylim(10, 90)
    axis.set_xticks(quantiles)
    axis.set_xlabel("Calibration quantile used as target")
    axis.set_ylabel("Rewarded trials (%)")
    axis.set_title(
        "Label-free calibration generalizes",
        loc="left",
        fontsize=11.5,
        fontweight="bold",
        pad=10,
    )
    axis.legend(frameon=False, fontsize=8, loc="upper right")
    axis.text(
        0.04,
        0.05,
        f"pooled success = "
        f"{100 * task_cluster_array(records, lambda row: row['signatures']['challenge_success_fraction']).mean():.2f}%",
        transform=axis.transAxes,
        fontsize=8.6,
    )
    style_axis(axis)
    panel_label(axis, "c")


def plot_outcome(axis, records, gate):
    keys = (
        (
            "terminal_previous_soma_outcome_balanced_acc",
            "Pre-outcome\nsoma",
            LIGHT_GRAY,
        ),
        (
            "acute_outcome_lesion_outcome_balanced_acc",
            "Innovation\noutcome lesion",
            ORANGE,
        ),
        (
            "terminal_residual_outcome_balanced_acc",
            "Innovation\nintact",
            BLUE,
        ),
    )
    values = [
        task_cluster_array(
            records, lambda row, key=key: row["signatures"][key]
        ) * 100
        for key, _, _ in keys
    ]
    for index, ((_, _, color), clustered) in enumerate(zip(keys, values)):
        axis.bar(
            index, clustered.mean(), color=color, width=0.66,
            edgecolor="white", zorder=2,
        )
        jitter = np.linspace(-0.13, 0.13, len(clustered))
        axis.scatter(
            index + jitter, clustered, s=22, facecolor="white",
            edgecolor="#333333", linewidth=0.7, zorder=3,
        )
    axis.axhline(50, color=RED, linestyle="--", linewidth=1.0)
    axis.set_ylim(45, 103)
    axis.set_xticks(range(len(keys)), [item[1] for item in keys])
    axis.set_ylabel("Outcome balanced accuracy (%)")
    axis.set_title(
        "Residual carries outcome surprise",
        loc="left",
        fontsize=11.5,
        fontweight="bold",
        pad=10,
    )
    drop = gate["metrics"]["outcome_lesion_drop"]
    critic = gate["metrics"]["critic_expectedness"]
    axis.text(
        0.04,
        0.95,
        "role separation lesion\n"
        f"$\\Delta$ = {drop['mean']:.3f} "
        f"(LCB {drop['one_sided_95pct_lower']:.3f})\n"
        "critic expectedness\n"
        f"{critic['mean']:.3f} "
        f"(LCB {critic['one_sided_95pct_lower']:.3f})",
        transform=axis.transAxes,
        va="top",
        fontsize=8.3,
        bbox={
            "boxstyle": "round,pad=0.35",
            "facecolor": "white",
            "edgecolor": "#D6D6D6",
            "alpha": 0.96,
        },
    )
    style_axis(axis)
    panel_label(axis, "d")


def write_caption(path):
    caption = """# Oral-B-v2 confirmation figure caption

**Figure 5 | Role-vectorized somato-dendritic innovation carries
temporal-difference outcome surprise.** All summaries use the complete
untouched six-task by five-model confirmation; uncertainty is computed after
averaging models within each task seed. **a,** Mean episode success over
training. The learned causal-role vector is necessary, the forward-plasticity
lesion (omitted for visual overlap with zero) does not learn, and critic or
terminal-outcome training lesions are reported rather than hidden. Bands are
95% normal intervals over six task clusters. **b,** Ordinary apical traffic is
almost perfectly soma-correlated before subtraction; the innovation is below
the frozen correlation ceiling in every task cluster. **c,** Five targets are
chosen from fixed quantiles of 512 outcome-free cursor-max calibration trials.
The displayed rewarded fractions come from independent assay trajectories;
no evaluation label selects or reweights a target. **d,** Terminal innovation
decodes rewarded versus timeout outcomes, while the acute terminal-outcome
lesion removes role-aligned separation. The critic contribution tracks the
stored value prediction and attenuates expected outcomes. Exact role maps,
outcome labels, and gradients are diagnostic-only and never used by the
learner's updates.
"""
    with open(path, "w") as handle:
        handle.write(caption)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input-dir",
        default="results/bci_v2_calibrated_confirmation",
    )
    parser.add_argument(
        "--gate",
        default="results/bci_v2_calibrated_confirmation_gate.json",
    )
    parser.add_argument("--outdir", default="results/figs")
    args = parser.parse_args()
    gate, records, paths = load_panel(args.input_dir, args.gate)
    os.makedirs(args.outdir, exist_ok=True)
    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 9.5,
        "axes.linewidth": 1.0,
        "xtick.major.width": 1.0,
        "ytick.major.width": 1.0,
        "pdf.fonttype": 42,
        "ps.fonttype": 42,
    })
    figure, axes = plt.subplots(2, 2, figsize=(10.5, 7.8))
    plot_learning(axes[0, 0], records)
    plot_residualization(axes[0, 1], records)
    plot_psychometric(axes[1, 0], records)
    plot_outcome(axes[1, 1], records, gate)
    figure.subplots_adjust(
        left=0.09, right=0.98, bottom=0.10, top=0.93,
        hspace=0.42, wspace=0.30,
    )
    pdf_path = os.path.join(args.outdir, "figure5_bci_v2.pdf")
    png_path = os.path.join(args.outdir, "figure5_bci_v2.png")
    caption_path = os.path.join(
        args.outdir, "figure5_bci_v2_caption.md"
    )
    manifest_path = os.path.join(
        args.outdir, "figure5_bci_v2_manifest.json"
    )
    figure.savefig(pdf_path, metadata=PDF_METADATA)
    figure.savefig(png_path, dpi=320)
    plt.close(figure)
    write_caption(caption_path)

    def mean_metric(key):
        return float(task_cluster_array(
            records, lambda row: row["signatures"][key]
        ).mean())

    terminal_previous_soma = mean_metric(
        "terminal_previous_soma_outcome_balanced_acc"
    )
    terminal_outcome_lesion = mean_metric(
        "acute_outcome_lesion_outcome_balanced_acc"
    )
    manifest = {
        "strict": True,
        "protocol": PROTOCOL,
        "gate_status": gate["status"],
        "task_seeds": list(TASK_SEEDS),
        "model_seeds": list(MODEL_SEEDS),
        "record_count": len(records),
        "statistics": {
            "intact_final_mean": gate["metrics"]["intact_final"]["mean"],
            "learning_gain_mean": gate["metrics"]["intact_gain"]["mean"],
            "fixed_role_gap_mean": gate["metrics"]["fixed_role_gap"]["mean"],
            "role_cosine_mean": gate["metrics"]["role_cosine"]["mean"],
            "residual_soma_corr_mean":
                gate["metrics"]["residual_soma_corr"]["mean"],
            "raw_residual_corr_gap_mean":
                gate["metrics"]["raw_residual_corr_gap"]["mean"],
            "surrounding_accuracy_mean":
                gate["metrics"]["surrounding_accuracy"]["mean"],
            "surrounding_accuracy_lower":
                gate["metrics"]["surrounding_accuracy"][
                    "one_sided_95pct_lower"
                ],
            "decoder_corr_mean": gate["metrics"]["decoder_corr"]["mean"],
            "sign_inversion_mean":
                gate["metrics"]["sign_inversion"]["mean"],
            "positive_sign_count": gate["positive_sign_count"],
            "velocity_advantage_mean":
                gate["metrics"]["velocity_advantage"]["mean"],
            "challenge_fraction_mean":
                gate["metrics"]["challenge_fraction"]["mean"],
            "terminal_accuracy_mean":
                gate["metrics"]["terminal_accuracy"]["mean"],
            "terminal_accuracy_lower":
                gate["metrics"]["terminal_accuracy"][
                    "one_sided_95pct_lower"
                ],
            "terminal_previous_soma_accuracy_mean":
                terminal_previous_soma,
            "terminal_outcome_lesion_accuracy_mean":
                terminal_outcome_lesion,
            "terminal_separation_mean":
                gate["metrics"]["terminal_separation"]["mean"],
            "outcome_lesion_drop_mean":
                gate["metrics"]["outcome_lesion_drop"]["mean"],
            "outcome_lesion_drop_lower":
                gate["metrics"]["outcome_lesion_drop"][
                    "one_sided_95pct_lower"
                ],
            "critic_expectedness_mean":
                gate["metrics"]["critic_expectedness"]["mean"],
            "critic_expectedness_lower":
                gate["metrics"]["critic_expectedness"][
                    "one_sided_95pct_lower"
                ],
        },
        "source_sha256": {
            os.path.relpath(path): sha256(path)
            for path in sorted(paths + [args.gate])
        },
    }
    with open(manifest_path, "w") as handle:
        json.dump(manifest, handle, indent=2, sort_keys=True)
        handle.write("\n")
    print(json.dumps({
        "manifest": manifest_path,
        "png": png_path,
        "status": "passed",
    }, indent=2))


if __name__ == "__main__":
    main()