revision/scripts/benchmark_inference.py

"""Benchmark warm-checkpoint inference for one frozen outer run."""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any, Sequence

import joblib
import numpy as np
import pandas as pd
import torch
from sklearn.preprocessing import LabelEncoder
from torch_geometric.loader import DataLoader as PyGDataLoader


PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.data import MolecularFeatureExtractor
from src.models import GATModel, GraphConvModel, HybridModel
from src.neural_models import FingerprintNN

from revision.scripts.reanalysis_core import annotate_structures
from revision.scripts.reanalysis_pipeline import _extract_features


MODEL_CLASSES = {"gat": GATModel, "gcn": GraphConvModel}


def summarize_durations(durations: Sequence[float], *, n_records: int) -> dict[str, float | int]:
    values = np.asarray(durations, dtype=float)
    if len(values) == 0 or n_records <= 0:
        raise ValueError("Durations and a positive record count are required.")
    mean = float(values.mean())
    sd = float(values.std(ddof=1)) if len(values) > 1 else 0.0
    return {
        "repeats": int(len(values)),
        "seconds_mean": mean,
        "seconds_sd": sd,
        "milliseconds_per_record_mean": mean / n_records * 1000.0,
    }


def _load_hybrid_model(
    checkpoint_path: Path,
    model_name: str,
    descriptor_dim: int,
    device: torch.device,
) -> HybridModel:
    payload: dict[str, Any] = torch.load(
        checkpoint_path, map_location="cpu", weights_only=True
    )
    config = payload["config"]
    model = HybridModel(
        graph_model_class=MODEL_CLASSES[model_name],
        descriptor_dim=descriptor_dim,
        graph_model_kwargs=config.get("graph_model_kwargs"),
        graph_feature_dim=config.get("graph_feature_dim"),
        descriptor_hidden_dims=config.get("descriptor_hidden_dims"),
        final_hidden_dims=config.get("final_hidden_dims"),
        dropout=config.get("dropout", 0.2),
        use_batch_norm=config.get("use_batch_norm", True),
        output_dim=config.get("output_dim", 1),
    )
    model.load_state_dict(payload["model_state"])
    model.target_mean = float(payload["target_mean"])
    model.target_std = float(payload["target_std"])
    model.eval()
    return model.to(device)


def _load_fpnn_model(checkpoint_path: Path, device: torch.device) -> FingerprintNN:
    payload: dict[str, Any] = torch.load(
        checkpoint_path, map_location="cpu", weights_only=True
    )
    if "config" in payload:
        model_config = payload["config"]
    else:
        state = payload["model_state"]
        laboratory_shape = state["lab_embedding.weight"].shape
        linear_keys = sorted(
            (
                key
                for key, value in state.items()
                if key.startswith("network.")
                and key.endswith(".weight")
                and value.ndim == 2
            ),
            key=lambda key: int(key.split(".")[1]),
        )
        linear_shapes = [state[key].shape for key in linear_keys]
        model_config = {
            "input_dim": int(linear_shapes[0][1] - laboratory_shape[1]),
            "hidden_dims": [int(shape[0]) for shape in linear_shapes[:-1]],
            "dropout": 0.15,
            "use_batch_norm": any("running_mean" in key for key in state),
            "num_labs": int(laboratory_shape[0]),
            "lab_embed_dim": int(laboratory_shape[1]),
            "input_dropout": 0.1,
        }
    model = FingerprintNN(**model_config)
    model.load_state_dict(payload["model_state"])
    model.target_mean = float(payload["target_mean"])
    model.target_std = float(payload["target_std"])
    model.eval()
    return model.to(device)


def _predict_hybrid(
    model: HybridModel,
    graphs: Sequence[Any],
    labs: np.ndarray,
    descriptors: np.ndarray,
    device: torch.device,
) -> np.ndarray:
    prepared = []
    for index, graph in enumerate(graphs):
        graph_copy = graph.clone()
        graph_copy.lab_feature = torch.tensor([int(labs[index])], dtype=torch.long)
        graph_copy.descriptors = torch.tensor(
            descriptors[index], dtype=torch.float32
        ).reshape(1, -1)
        prepared.append(graph_copy)
    loader = PyGDataLoader(prepared, batch_size=64, shuffle=False)
    predictions: list[float] = []
    with torch.no_grad():
        for batch in loader:
            batch = batch.to(device)
            lab_tensor = (
                batch.lab_feature.squeeze(-1)
                if batch.lab_feature.dim() > 1
                else batch.lab_feature
            )
            descriptor_tensor = batch.descriptors.reshape(batch.num_graphs, -1)
            values = model(
                batch.x,
                batch.edge_index,
                batch.batch,
                lab_tensor,
                descriptor_tensor,
                getattr(batch, "edge_attr", None),
            )
            predictions.extend(values.detach().cpu().numpy().reshape(-1).tolist())
    values = np.asarray(predictions, dtype=np.float32)
    return values * model.target_std + model.target_mean


def _predict_fpnn(
    model: FingerprintNN,
    fingerprints: np.ndarray,
    labs: np.ndarray,
    device: torch.device,
) -> np.ndarray:
    features = torch.tensor(fingerprints, dtype=torch.float32)
    laboratory = torch.tensor(labs, dtype=torch.long)
    loader = torch.utils.data.DataLoader(
        torch.utils.data.TensorDataset(features, laboratory),
        batch_size=256,
        shuffle=False,
    )
    predictions: list[float] = []
    with torch.no_grad():
        for batch_features, batch_labs in loader:
            values = model(
                batch_features.to(device), batch_labs.to(device)
            ).detach().cpu().numpy().reshape(-1)
            predictions.extend(values.tolist())
    values = np.asarray(predictions, dtype=np.float32)
    return values * model.target_std + model.target_mean


def _synchronize(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Measure loaded six-fold neural-stack inference cost."
    )
    parser.add_argument("--data", required=True)
    parser.add_argument("--artifacts-root", required=True)
    parser.add_argument("--strategy", default="canonical_grouped")
    parser.add_argument("--seed", type=int, default=123456)
    parser.add_argument("--expected-folds", type=int, default=6)
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--output", required=True)
    arguments = parser.parse_args()

    if arguments.expected_folds <= 0 or arguments.repeats <= 0:
        raise ValueError("expected-folds and repeats must be positive.")
    device = torch.device(
        "cuda"
        if arguments.device == "auto" and torch.cuda.is_available()
        else "cpu"
        if arguments.device == "auto"
        else arguments.device
    )
    artifacts_root = Path(arguments.artifacts_root).resolve()
    run_dir = artifacts_root / arguments.strategy / f"seed_{arguments.seed}"
    neural_dir = run_dir / "neural_stack"
    output_path = Path(arguments.output).resolve()
    if output_path.exists():
        raise FileExistsError(f"Refusing to overwrite inference benchmark: {output_path}")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    config = json.loads(
        (artifacts_root / "FROZEN_CONFIG.json").read_text(encoding="utf-8")
    )
    annotated = annotate_structures(pd.read_csv(arguments.data))
    indices = np.load(run_dir / "split_indices.npz", allow_pickle=True)
    test_indices = np.asarray(indices["test_indices"], dtype=int)
    test_frame = annotated.iloc[test_indices].reset_index(drop=True)

    feature_started = time.perf_counter()
    descriptors, fingerprints = _extract_features(
        test_frame,
        config["descriptor_features"],
        config["fingerprint"],
    )
    extractor = MolecularFeatureExtractor()
    graphs = [extractor.smiles_to_graph(value) for value in test_frame["SMILES"].astype(str)]
    if any(graph is None for graph in graphs):
        raise ValueError("Graph generation failed during inference benchmark.")
    feature_seconds = time.perf_counter() - feature_started

    encoder_payload = json.loads(
        (neural_dir / "fold_preprocessing" / "lab_encoder.json").read_text(
            encoding="utf-8"
        )
    )
    lab_encoder = LabelEncoder()
    lab_encoder.classes_ = np.asarray(encoder_payload["classes_in_index_order"], dtype=object)
    laboratory_indices = lab_encoder.transform(test_frame["Lab"].astype(str)).astype(np.int64)

    hybrid_models: dict[str, list[tuple[HybridModel, np.ndarray]]] = {
        "gat": [],
        "gcn": [],
    }
    fpnn_models: list[FingerprintNN] = []
    for fold in range(arguments.expected_folds):
        preprocessing = joblib.load(
            neural_dir / "fold_preprocessing" / f"fold_{fold}.joblib"
        )
        scaled_descriptors = preprocessing["scaler"].transform(descriptors).astype(
            np.float32
        )
        for model_name in ("gat", "gcn"):
            checkpoint = (
                neural_dir
                / "checkpoints"
                / model_name
                / f"{model_name}_fold_{fold}.pt"
            )
            if not checkpoint.is_file():
                raise FileNotFoundError(f"Incomplete checkpoint matrix: {checkpoint}")
            hybrid_models[model_name].append(
                (
                    _load_hybrid_model(
                        checkpoint,
                        model_name,
                        descriptor_dim=descriptors.shape[1],
                        device=device,
                    ),
                    scaled_descriptors,
                )
            )
        fpnn_checkpoint = (
            neural_dir / "checkpoints" / "fpnn" / f"fp_nn_fold_{fold}.pt"
        )
        if not fpnn_checkpoint.is_file():
            raise FileNotFoundError(f"Incomplete checkpoint matrix: {fpnn_checkpoint}")
        fpnn_models.append(_load_fpnn_model(fpnn_checkpoint, device))

    stack_models = joblib.load(neural_dir / "stack_models.joblib")
    primary_stack = stack_models["stack_all_plus_descriptors"]

    component_durations = {"gat": [], "gcn": [], "fpnn": [], "stack": []}
    total_durations: list[float] = []
    for repeat in range(arguments.repeats + 1):
        _synchronize(device)
        total_started = time.perf_counter()
        averaged: dict[str, np.ndarray] = {}
        for model_name in ("gat", "gcn"):
            _synchronize(device)
            component_started = time.perf_counter()
            fold_predictions = [
                _predict_hybrid(
                    model,
                    graphs,
                    laboratory_indices,
                    scaled_descriptors,
                    device,
                )
                for model, scaled_descriptors in hybrid_models[model_name]
            ]
            _synchronize(device)
            elapsed = time.perf_counter() - component_started
            averaged[model_name] = np.mean(np.vstack(fold_predictions), axis=0)
            if repeat > 0:
                component_durations[model_name].append(elapsed)

        _synchronize(device)
        fpnn_started = time.perf_counter()
        fpnn_predictions = [
            _predict_fpnn(model, fingerprints, laboratory_indices, device)
            for model in fpnn_models
        ]
        _synchronize(device)
        fpnn_elapsed = time.perf_counter() - fpnn_started
        averaged["fpnn"] = np.mean(np.vstack(fpnn_predictions), axis=0)

        stack_features = np.column_stack(
            [averaged["gat"], averaged["gcn"], averaged["fpnn"], descriptors]
        )
        stack_started = time.perf_counter()
        _ = primary_stack.predict(stack_features)
        stack_elapsed = time.perf_counter() - stack_started
        total_elapsed = time.perf_counter() - total_started
        if repeat > 0:
            component_durations["fpnn"].append(fpnn_elapsed)
            component_durations["stack"].append(stack_elapsed)
            total_durations.append(total_elapsed)

    payload = {
        "strategy": arguments.strategy,
        "seed": arguments.seed,
        "device": str(device),
        "gpu_name": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
        "n_records": int(len(test_frame)),
        "n_fold_models_per_base": arguments.expected_folds,
        "scope": "warm loaded checkpoints; includes batching and device transfer; excludes checkpoint loading and RDKit feature generation",
        "feature_generation": {
            "seconds": float(feature_seconds),
            "milliseconds_per_record": float(feature_seconds / len(test_frame) * 1000.0),
        },
        "full_bundle": summarize_durations(total_durations, n_records=len(test_frame)),
        "components": {
            name: summarize_durations(values, n_records=len(test_frame))
            for name, values in component_durations.items()
        },
    }
    output_path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    print(f"Wrote inference benchmark to: {output_path}")
    return 0


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