revision/scripts/generate_revision_assets.py

"""Generate reviewer-requested, source-data-backed revision figures."""

from __future__ import annotations

import argparse
import hashlib
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit.Chem import Crippen, Descriptors, rdMolDescriptors


STRATEGIES = ("canonical_grouped", "scaffold_aware")
SEEDS = (123456, 123457, 123458)
STRATEGY_LABELS = {
    "canonical_grouped": "Identity-grouped",
    "scaffold_aware": "Scaffold-aware",
}
COLORS = {
    "canonical_grouped": "#3B6FB6",
    "scaffold_aware": "#D97936",
}


def _set_style() -> None:
    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 8.5,
            "axes.labelsize": 9,
            "axes.titlesize": 9,
            "legend.fontsize": 7.5,
            "xtick.labelsize": 7.5,
            "ytick.labelsize": 7.5,
            "axes.spines.top": False,
            "axes.spines.right": False,
            "savefig.dpi": 300,
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )


def _panel_label(axis: plt.Axes, label: str) -> None:
    axis.text(
        -0.13,
        1.06,
        label,
        transform=axis.transAxes,
        fontsize=11,
        fontweight="bold",
        va="top",
    )


def _save_figure(figure: plt.Figure, output_dir: Path, name: str) -> None:
    figure.savefig(output_dir / f"{name}.pdf", bbox_inches="tight")
    figure.savefig(output_dir / f"{name}.png", bbox_inches="tight", dpi=300)
    plt.close(figure)


def _distribution_table(values: np.ndarray, property_name: str, bins: int = 8) -> pd.DataFrame:
    counts, edges = np.histogram(np.asarray(values, dtype=float), bins=bins)
    percentages = counts.astype(float) / counts.sum() * 100.0
    return pd.DataFrame(
        {
            "property": property_name,
            "bin_lower": edges[:-1],
            "bin_upper": edges[1:],
            "bin_center": (edges[:-1] + edges[1:]) / 2.0,
            "count": counts,
            "percentage": percentages,
        }
    )


def _molecular_properties(data: pd.DataFrame) -> pd.DataFrame:
    rows = []
    for row_index, smiles in enumerate(data["SMILES"].astype(str)):
        molecule = Chem.MolFromSmiles(smiles)
        if molecule is None:
            raise ValueError(f"Invalid SMILES at input row {row_index}: {smiles}")
        rows.append(
            {
                "MW": Descriptors.MolWt(molecule),
                "LogP": Crippen.MolLogP(molecule),
                "TPSA": rdMolDescriptors.CalcTPSA(molecule),
                "RT": float(data.iloc[row_index]["RT"]),
            }
        )
    return pd.DataFrame(rows)


def _toc_graphic(output_dir: Path) -> None:
    """Draw a result-neutral TOC workflow synchronized with the revised scope."""
    figure, axis = plt.subplots(figsize=(3.25, 1.75))
    axis.set_xlim(0, 10)
    axis.set_ylim(0, 4)
    axis.axis("off")

    def box(
        x: float,
        y: float,
        width: float,
        height: float,
        text: str,
        color: str,
        fontsize: float = 6.2,
    ) -> None:
        patch = FancyBboxPatch(
            (x, y),
            width,
            height,
            boxstyle="round,pad=0.08,rounding_size=0.12",
            linewidth=1.0,
            edgecolor=color,
            facecolor=color + "20",
        )
        axis.add_patch(patch)
        axis.text(
            x + width / 2,
            y + height / 2,
            text,
            ha="center",
            va="center",
            fontsize=fontsize,
            fontweight="bold",
            color="#233447",
        )

    def arrow(start: tuple[float, float], end: tuple[float, float]) -> None:
        axis.add_patch(
            FancyArrowPatch(
                start,
                end,
                arrowstyle="-|>",
                mutation_scale=12,
                linewidth=1.2,
                color="#536878",
            )
        )

    box(0.2, 2.35, 1.55, 0.75, "Molecular\nstructure", "#3B6FB6", 5.8)
    box(0.2, 0.85, 1.55, 0.75, "Laboratory\nlabel", "#7A6AA6", 5.8)
    box(2.35, 2.7, 1.35, 0.65, "GAT", "#3B6FB6")
    box(2.35, 1.75, 1.35, 0.65, "GCN", "#3B6FB6")
    box(2.35, 0.8, 1.35, 0.65, "FPNN", "#3B6FB6")
    box(4.55, 1.55, 2.15, 1.05, "OOF base\npredictions", "#D97936", 5.1)
    box(7.25, 1.55, 1.45, 1.05, "ExtraTrees\nstack", "#D97936", 5.4)
    box(9.15, 1.55, 0.65, 1.05, "RT", "#4A8C67")

    for y in (3.02, 2.07, 1.12):
        arrow((1.75, 2.72), (2.30, y))
        arrow((1.75, 1.22), (2.30, y))
        arrow((3.70, y), (4.50, 2.07))
    arrow((6.70, 2.07), (7.20, 2.07))
    arrow((8.70, 2.07), (9.10, 2.07))
    axis.text(
        5.0,
        0.20,
        "Identity-grouped and scaffold-aware evaluation\nwithin 23 represented laboratories",
        ha="center",
        va="center",
        fontsize=6.0,
        color="#36454F",
    )
    _save_figure(figure, output_dir, "toc_graphic")


def _figure_performance(summary: pd.DataFrame, output_dir: Path, source_dir: Path) -> None:
    model_order = ("gat", "gcn", "fpnn", "stack_all_plus_descriptors")
    labels = ("GAT", "GCN", "FPNN", "Full stack")
    source = summary.loc[summary["model"].isin(model_order)].copy()
    source["model"] = pd.Categorical(source["model"], model_order, ordered=True)
    source = source.sort_values(["strategy", "model"])
    source.to_csv(source_dir / "figure1_performance.csv", index=False)

    figure, axes = plt.subplots(1, 3, figsize=(7.2, 2.45))
    x = np.arange(len(model_order), dtype=float)
    width = 0.36
    for strategy_index, strategy in enumerate(STRATEGIES):
        subset = source.loc[source["strategy"] == strategy].set_index("model").reindex(model_order)
        offset = (strategy_index - 0.5) * width
        for axis, metric, ylabel in zip(
            axes,
            ("r2", "mae", "rmse"),
            ("$R^2$", "MAE (min)", "RMSE (min)"),
        ):
            axis.bar(
                x + offset,
                subset[f"{metric}_mean"],
                width,
                yerr=subset[f"{metric}_sd"],
                color=COLORS[strategy],
                alpha=0.9,
                capsize=2,
                label=STRATEGY_LABELS[strategy],
            )
            axis.set_ylabel(ylabel)
            axis.set_xticks(x, labels, rotation=25, ha="right")
            axis.grid(axis="y", color="#D9D9D9", linewidth=0.6, alpha=0.8)
    handles, legend_labels = axes[0].get_legend_handles_labels()
    figure.legend(
        handles,
        legend_labels,
        frameon=False,
        loc="lower center",
        bbox_to_anchor=(0.5, 0.005),
        ncol=2,
    )
    for axis, label in zip(axes, "ABC"):
        _panel_label(axis, label)
    figure.tight_layout(rect=(0.0, 0.12, 1.0, 1.0), w_pad=1.1)
    _save_figure(figure, output_dir, "figure1_performance")


def _figure_property_distributions(data: pd.DataFrame, output_dir: Path, source_dir: Path) -> None:
    properties = _molecular_properties(data)
    distribution = pd.concat(
        [_distribution_table(properties[column].to_numpy(), column) for column in properties],
        ignore_index=True,
    )
    distribution.to_csv(source_dir / "figure2_property_distributions.csv", index=False)

    labels = {
        "MW": "Molecular weight (Da)",
        "LogP": "RDKit LogP",
        "TPSA": "TPSA ($\AA^2$)",
        "RT": "Retention time (min)",
    }
    figure, axes = plt.subplots(2, 2, figsize=(7.2, 5.0))
    for axis, column, panel in zip(axes.flat, properties.columns, "ABCD"):
        subset = distribution.loc[distribution["property"] == column]
        widths = subset["bin_upper"].to_numpy() - subset["bin_lower"].to_numpy()
        axis.bar(
            subset["bin_center"],
            subset["percentage"],
            width=widths * 0.92,
            color="#4F81BD",
            edgecolor="white",
            linewidth=0.5,
        )
        axis.set_xlabel(labels[column])
        axis.set_ylabel("Observations (%)")
        axis.grid(axis="y", color="#D9D9D9", linewidth=0.6, alpha=0.8)
        _panel_label(axis, panel)
    figure.tight_layout(h_pad=1.4, w_pad=1.3)
    _save_figure(figure, output_dir, "figure2_property_distributions")


def _figure_ablations(
    neural_summary: pd.DataFrame,
    classical_summary: pd.DataFrame,
    output_dir: Path,
    source_dir: Path,
) -> None:
    neural_models = (
        "arithmetic_mean",
        "stack_all_base_only",
        "stack_all_plus_descriptors",
        "stack_gat_gcn",
        "stack_gat_fpnn",
        "stack_gcn_fpnn",
    )
    classical_models = (
        "lab_median",
        "descriptor_only_ridge",
        "fingerprint_no_lab_et",
        "fingerprint_no_lab_plus_lab_affine",
        "fingerprint_only_et",
        "fingerprint_plus_descriptors_et",
    )
    neural = neural_summary.loc[neural_summary["model"].isin(neural_models)].copy()
    neural["family"] = "Neural-stack ablation"
    classical = classical_summary.loc[classical_summary["model"].isin(classical_models)].copy()
    classical["family"] = "Fixed classical baseline"
    source = pd.concat([neural, classical], ignore_index=True)
    source.to_csv(source_dir / "figure3_ablations.csv", index=False)

    display = {
        "arithmetic_mean": "Arithmetic mean",
        "stack_all_base_only": "3-base stack",
        "stack_all_plus_descriptors": "3-base + descriptors",
        "stack_gat_gcn": "GAT + GCN",
        "stack_gat_fpnn": "GAT + FPNN",
        "stack_gcn_fpnn": "GCN + FPNN",
        "lab_median": "Lab median",
        "descriptor_only_ridge": "Descriptors + lab ridge",
        "fingerprint_no_lab_et": "Morgan ET, no lab",
        "fingerprint_no_lab_plus_lab_affine": "Morgan ET + lab affine",
        "fingerprint_only_et": "Morgan + one-hot lab ET",
        "fingerprint_plus_descriptors_et": "Morgan + descriptors + lab ET",
    }
    figure, axes = plt.subplots(1, 2, figsize=(7.2, 4.9))
    for axis, (family, model_order), panel in zip(
        axes,
        (("Neural-stack ablation", neural_models), ("Fixed classical baseline", classical_models)),
        "AB",
    ):
        y = np.arange(len(model_order), dtype=float)
        height = 0.36
        for strategy_index, strategy in enumerate(STRATEGIES):
            subset = source.loc[
                (source["family"] == family) & (source["strategy"] == strategy)
            ].set_index("model").reindex(model_order)
            offset = (strategy_index - 0.5) * height
            axis.barh(
                y + offset,
                subset["mae_mean"],
                height,
                xerr=subset["mae_sd"],
                color=COLORS[strategy],
                capsize=2,
                label=STRATEGY_LABELS[strategy],
            )
        axis.set_yticks(y, [display[name] for name in model_order])
        axis.invert_yaxis()
        axis.set_xlabel("MAE (min; mean $\pm$ SD)")
        axis.grid(axis="x", color="#D9D9D9", linewidth=0.6, alpha=0.8)
        axis.set_title(family)
        _panel_label(axis, panel)
    handles, legend_labels = axes[0].get_legend_handles_labels()
    figure.legend(
        handles,
        legend_labels,
        frameon=False,
        loc="lower center",
        bbox_to_anchor=(0.5, 0.005),
        ncol=2,
    )
    figure.tight_layout(rect=(0.0, 0.08, 1.0, 1.0), w_pad=1.8)
    _save_figure(figure, output_dir, "figure3_ablations")


def _collect_predictions(artifacts_root: Path) -> pd.DataFrame:
    frames = []
    for strategy in STRATEGIES:
        for seed in SEEDS:
            path = artifacts_root / strategy / f"seed_{seed}" / "neural_stack" / "test_predictions.csv"
            if not path.is_file():
                raise FileNotFoundError(f"Incomplete prediction matrix: {path}")
            frame = pd.read_csv(path)
            frame.insert(0, "seed", seed)
            frame.insert(0, "strategy", strategy)
            frames.append(frame)
    return pd.concat(frames, ignore_index=True)


def _figure_diagnostics(
    predictions: pd.DataFrame,
    per_lab: pd.DataFrame,
    output_dir: Path,
    source_dir: Path,
) -> None:
    predictions = predictions.copy()
    predictions["absolute_error"] = np.abs(
        predictions["prediction_stack_all_plus_descriptors"] - predictions["RT"]
    )
    predictions["similarity_bin"] = pd.cut(
        predictions["maximum_development_tanimoto"],
        bins=[0.0, 0.2, 0.4, 0.6, 0.8, 1.000001],
        right=False,
        include_lowest=True,
    ).astype(str)
    predictions["rt_quartile"] = predictions.groupby(["strategy", "seed"])["RT"].transform(
        lambda values: pd.qcut(values, q=4, labels=("Q1", "Q2", "Q3", "Q4"), duplicates="drop")
    )
    similarity_bins = (
        predictions.groupby(["strategy", "similarity_bin"], observed=False)
        .agg(n=("absolute_error", "size"), mae=("absolute_error", "mean"))
        .reset_index()
    )
    rt_bins = (
        predictions.groupby(["strategy", "rt_quartile"], observed=False)
        .agg(n=("absolute_error", "size"), mae=("absolute_error", "mean"))
        .reset_index()
    )
    lab_primary = per_lab.loc[per_lab["model"] == "stack_all_plus_descriptors"].copy()
    lab_summary = (
        lab_primary.groupby(["strategy", "Lab"])
        .agg(n=("n", "sum"), mae=("mae", "mean"), mae_sd=("mae", "std"))
        .reset_index()
    )
    predictions.to_csv(source_dir / "figure4_predictions.csv", index=False)
    similarity_bins.to_csv(source_dir / "figure4_similarity_bins.csv", index=False)
    rt_bins.to_csv(source_dir / "figure4_rt_quartiles.csv", index=False)
    lab_summary.to_csv(source_dir / "figure4_per_lab.csv", index=False)

    figure, axes = plt.subplots(2, 2, figsize=(7.2, 5.6))
    canonical = predictions.loc[predictions["strategy"] == "canonical_grouped"]
    axes[0, 0].scatter(
        canonical["maximum_development_tanimoto"],
        canonical["absolute_error"],
        s=10,
        alpha=0.28,
        color=COLORS["canonical_grouped"],
        edgecolors="none",
    )
    axes[0, 0].set_xlabel("Maximum development-set Tanimoto")
    axes[0, 0].set_ylabel("Absolute error (min)")

    similarity_order = list(similarity_bins["similarity_bin"].drop_duplicates())
    x = np.arange(len(similarity_order), dtype=float)
    width = 0.36
    for strategy_index, strategy in enumerate(STRATEGIES):
        subset = similarity_bins.loc[similarity_bins["strategy"] == strategy].set_index(
            "similarity_bin"
        ).reindex(similarity_order)
        axes[0, 1].bar(
            x + (strategy_index - 0.5) * width,
            subset["mae"],
            width,
            color=COLORS[strategy],
            label=STRATEGY_LABELS[strategy],
        )
    axes[0, 1].set_xticks(x, similarity_order, rotation=25, ha="right")
    axes[0, 1].set_xlabel("Maximum Tanimoto bin")
    axes[0, 1].set_ylabel("MAE (min)")
    axes[0, 1].legend(frameon=False)

    canonical_labs = lab_summary.loc[lab_summary["strategy"] == "canonical_grouped"].sort_values(
        "mae", ascending=True
    )
    axes[1, 0].barh(
        np.arange(len(canonical_labs)),
        canonical_labs["mae"],
        color=COLORS["canonical_grouped"],
    )
    axes[1, 0].set_yticks(np.arange(len(canonical_labs)), canonical_labs["Lab"])
    axes[1, 0].set_xlabel("Identity-grouped MAE (min)")
    axes[1, 0].tick_params(axis="y", labelsize=5.8)

    quartile_order = ("Q1", "Q2", "Q3", "Q4")
    x = np.arange(4, dtype=float)
    for strategy_index, strategy in enumerate(STRATEGIES):
        subset = rt_bins.loc[rt_bins["strategy"] == strategy].set_index("rt_quartile").reindex(
            quartile_order
        )
        axes[1, 1].bar(
            x + (strategy_index - 0.5) * width,
            subset["mae"],
            width,
            color=COLORS[strategy],
            label=STRATEGY_LABELS[strategy],
        )
    axes[1, 1].set_xticks(x, quartile_order)
    axes[1, 1].set_xlabel("Within-run experimental RT quartile")
    axes[1, 1].set_ylabel("MAE (min)")
    for axis, panel in zip(axes.flat, "ABCD"):
        axis.grid(axis="both", color="#D9D9D9", linewidth=0.5, alpha=0.7)
        _panel_label(axis, panel)
    figure.tight_layout(h_pad=1.4, w_pad=1.2)
    _save_figure(figure, output_dir, "figure4_diagnostics")


def _write_manifest(output_dir: Path) -> None:
    rows = []
    for path in sorted(output_dir.rglob("*")):
        if path.is_file() and path.name != "FIGURE_MANIFEST.csv":
            rows.append(
                {
                    "relative_path": path.relative_to(output_dir).as_posix(),
                    "bytes": path.stat().st_size,
                    "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
                }
            )
    pd.DataFrame(rows).to_csv(output_dir / "FIGURE_MANIFEST.csv", index=False)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", required=True)
    parser.add_argument("--artifacts-root", required=True)
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()

    data_path = Path(args.data).resolve()
    artifacts_root = Path(args.artifacts_root).resolve()
    output_dir = Path(args.output_dir).resolve()
    if output_dir.exists() and any(output_dir.iterdir()):
        raise FileExistsError(f"Refusing nonempty figure output directory: {output_dir}")
    output_dir.mkdir(parents=True, exist_ok=True)
    source_dir = output_dir / "source_data"
    source_dir.mkdir()
    _set_style()

    data = pd.read_csv(data_path)
    summary_dir = artifacts_root / "summary"
    neural_summary = pd.read_csv(summary_dir / "neural_metrics_aggregate.csv")
    classical_summary = pd.read_csv(summary_dir / "classical_metrics_aggregate.csv")
    per_lab = pd.read_csv(summary_dir / "neural_per_lab_by_run.csv")
    predictions = _collect_predictions(artifacts_root)

    _figure_performance(neural_summary, output_dir, source_dir)
    _figure_property_distributions(data, output_dir, source_dir)
    _figure_ablations(neural_summary, classical_summary, output_dir, source_dir)
    _figure_diagnostics(predictions, per_lab, output_dir, source_dir)
    _toc_graphic(output_dir)
    _write_manifest(output_dir)
    print(f"Wrote revision figures and source data to: {output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())