evaluate_random_search_nested_holdout.py

#!/usr/bin/env python3
"""
Evaluate random-search trials with a nested holdout protocol.

For every trial, the script:
1. Reads the prediction files saved per epoch.
2. Splits the original held-out split into:
   - 1/3 validation subset for epoch selection
   - 2/3 test subset for reporting the hyperparameter combination
3. Selects the best epoch using validation AUPRC only.
4. Reports test metrics for that chosen epoch.

This keeps epoch selection separate from the final reported test score.
"""

from __future__ import annotations

import argparse
import json
import logging
import pickle
import re
from collections import defaultdict
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import numpy as np
import pandas as pd
from sklearn.metrics import (
    average_precision_score,
    balanced_accuracy_score,
    matthews_corrcoef,
    roc_auc_score,
)
from sklearn.model_selection import train_test_split


logger = logging.getLogger("nested_random_search_eval")
if not logger.handlers:
    logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate random-search trials with a 1/3 validation and 2/3 test split of the held-out predictions."
    )
    parser.add_argument(
        "--search-root",
        type=Path,
        action="append",
        required=True,
        help="Root directory to scan recursively for prediction TSV files. Repeatable.",
    )
    parser.add_argument(
        "--search-tag",
        type=str,
        required=True,
        help="Substring used to identify only the trials from one random-search batch.",
    )
    parser.add_argument(
        "--split-file",
        type=Path,
        required=True,
        help="Held-out split TSV (typically __ft_val.tsv) used during training.",
    )
    parser.add_argument(
        "--dataset-path",
        type=Path,
        help="Optional dataset pickle used to recover labels by Gene_ID.",
    )
    parser.add_argument(
        "--manifest",
        type=Path,
        help="Optional TSV manifest generated by the launcher script. Joined into the summary using run_tag.",
    )
    parser.add_argument(
        "--output-dir",
        type=Path,
        required=True,
        help="Directory for the evaluation outputs.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=2023,
        help="Random seed used for the nested validation/test split.",
    )
    parser.add_argument(
        "--validation-fraction",
        type=float,
        default=1.0 / 3.0,
        help="Fraction of the held-out split reserved for epoch selection (default: 1/3).",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=150,
        help="Top-K precision summary to report alongside AUPRC (default: 150).",
    )
    return parser.parse_args()


def load_label_map(dataset_path: Optional[Path]) -> Dict[str, int]:
    if dataset_path is None:
        return {}

    with open(dataset_path, "rb") as fh:
        dataset_df = pickle.load(fh)
    dataset_df = dataset_df.reset_index(drop=True)

    if "Gene_ID" not in dataset_df.columns or "positive" not in dataset_df.columns:
        raise KeyError(f"{dataset_path} does not contain the required Gene_ID/positive columns.")

    dedup = dataset_df[["Gene_ID", "positive"]].drop_duplicates(subset=["Gene_ID"], keep="first")
    if dedup["Gene_ID"].duplicated().any():
        raise ValueError(f"Gene_ID values are not unique in {dataset_path}; cannot build a stable label map.")

    label_map = {str(gene_id): int(label) for gene_id, label in zip(dedup["Gene_ID"], dedup["positive"])}
    logger.info("Loaded %d labels from %s", len(label_map), dataset_path)
    return label_map


def load_manifest(manifest_path: Optional[Path]) -> Optional[pd.DataFrame]:
    if manifest_path is None:
        return None
    manifest_df = pd.read_csv(manifest_path, sep="\t")
    if "run_tag" not in manifest_df.columns:
        raise KeyError(f"Manifest {manifest_path} must contain a 'run_tag' column.")
    return manifest_df


def safe_auprc(labels: np.ndarray, scores: np.ndarray) -> float:
    if len(labels) == 0 or len(np.unique(labels)) < 2:
        return float("nan")
    return float(average_precision_score(labels, scores))


def safe_auroc(labels: np.ndarray, scores: np.ndarray) -> float:
    if len(labels) == 0 or len(np.unique(labels)) < 2:
        return float("nan")
    return float(roc_auc_score(labels, scores))


def safe_threshold_metrics(labels: np.ndarray, scores: np.ndarray) -> Tuple[float, float]:
    if len(labels) == 0 or len(np.unique(labels)) < 2:
        return float("nan"), float("nan")
    preds = (scores >= 0.5).astype(int)
    return (
        float(matthews_corrcoef(labels, preds)),
        float(balanced_accuracy_score(labels, preds)),
    )


def precision_at_k(df: pd.DataFrame, k: int) -> float:
    if df.empty:
        return float("nan")
    top_k = min(k, len(df))
    return float(df.sort_values("score", ascending=False).head(top_k)["label"].mean())


def discover_prediction_files(search_roots: Iterable[Path], search_tag: str) -> List[Path]:
    discovered: List[Path] = []
    for root in search_roots:
        if not root.exists():
            logger.warning("Search root does not exist and will be skipped: %s", root)
            continue
        for path in root.rglob("*predictions_epoch_*.tsv"):
            if search_tag in str(path):
                discovered.append(path)

    discovered = sorted(set(discovered))
    logger.info("Discovered %d prediction files containing tag '%s'.", len(discovered), search_tag)
    return discovered


def parse_trial_and_epoch(path: Path) -> Optional[Tuple[str, int]]:
    stem = path.stem
    match = re.match(r"(.+)_predictions_epoch_(\d+)$", stem)
    if match:
        return match.group(1), int(match.group(2))

    match = re.match(r"predictions_epoch_(\d+)$", stem)
    if match:
        return path.parent.name, int(match.group(1))

    return None


def build_holdout_frame(
    split_file: Path,
    label_map: Dict[str, int],
    fallback_predictions: Optional[pd.DataFrame] = None,
) -> pd.DataFrame:
    holdout_df = pd.read_csv(split_file, sep="\t")
    holdout_df["Gene_ID"] = holdout_df["Gene_ID"].astype(str)

    if label_map:
        holdout_df["label"] = holdout_df["Gene_ID"].map(label_map)
    elif fallback_predictions is not None:
        pred_labels = (
            fallback_predictions[["Gene_ID", "label"]]
            .drop_duplicates(subset=["Gene_ID"], keep="first")
            .copy()
        )
        pred_labels["Gene_ID"] = pred_labels["Gene_ID"].astype(str)
        holdout_df = holdout_df.merge(pred_labels, on="Gene_ID", how="left")
    else:
        raise ValueError("No labels available. Provide --dataset-path or ensure predictions can be used as fallback.")

    missing = holdout_df["label"].isna().sum()
    if missing:
        raise ValueError(f"{missing} held-out genes have no label mapping. Check the dataset and split inputs.")

    holdout_df["label"] = holdout_df["label"].astype(int)
    return holdout_df


def split_nested_holdout(holdout_df: pd.DataFrame, validation_fraction: float, seed: int) -> Tuple[pd.DataFrame, pd.DataFrame]:
    if not 0.0 < validation_fraction < 1.0:
        raise ValueError("--validation-fraction must be strictly between 0 and 1.")

    stratify = holdout_df["label"] if holdout_df["label"].nunique() > 1 else None
    selection_df, test_df = train_test_split(
        holdout_df,
        test_size=1.0 - validation_fraction,
        random_state=seed,
        stratify=stratify,
    )

    selection_df = selection_df.sort_values("Gene_ID").reset_index(drop=True)
    test_df = test_df.sort_values("Gene_ID").reset_index(drop=True)
    return selection_df, test_df


def match_manifest_row(manifest_df: Optional[pd.DataFrame], trial_key: str, example_path: Path) -> Dict[str, object]:
    if manifest_df is None or manifest_df.empty:
        return {}

    path_text = str(example_path)
    candidates = manifest_df[
        manifest_df["run_tag"].astype(str).apply(lambda tag: bool(tag) and (tag in trial_key or tag in path_text))
    ]
    if candidates.empty:
        return {}

    if len(candidates) > 1:
        candidates = candidates.assign(_tag_len=candidates["run_tag"].astype(str).str.len()).sort_values("_tag_len", ascending=False)
    row = candidates.iloc[0].drop(labels=["_tag_len"], errors="ignore")
    return row.to_dict()


def load_prediction_table(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path, sep="\t")
    if "positive" in df.columns and "label" not in df.columns:
        df = df.rename(columns={"positive": "label"})

    required = {"Gene_ID", "score", "label"}
    missing = required.difference(df.columns)
    if missing:
        raise KeyError(f"{path} is missing columns: {sorted(missing)}")

    df = df[["Gene_ID", "score", "label"]].copy()
    df["Gene_ID"] = df["Gene_ID"].astype(str)
    df["label"] = df["label"].astype(int)
    df["score"] = df["score"].astype(float)
    return df


def evaluate_subset(df: pd.DataFrame, gene_ids: Iterable[str], top_k: int) -> Dict[str, float]:
    subset = df[df["Gene_ID"].isin(set(gene_ids))].copy()
    labels = subset["label"].to_numpy()
    scores = subset["score"].to_numpy()
    mcc, bacc = safe_threshold_metrics(labels, scores)

    return {
        "n_samples": int(len(subset)),
        "auprc": safe_auprc(labels, scores),
        "auroc": safe_auroc(labels, scores),
        "mcc": mcc,
        "bacc": bacc,
        f"precision_at_{top_k}": precision_at_k(subset, top_k),
    }


def make_json_safe(record: Dict[str, object]) -> Dict[str, object]:
    safe: Dict[str, object] = {}
    for key, value in record.items():
        if isinstance(value, (np.floating, float)):
            safe[key] = None if np.isnan(value) else float(value)
        elif isinstance(value, (np.integer, int)):
            safe[key] = int(value)
        else:
            safe[key] = value
    return safe


def main() -> None:
    args = parse_args()

    prediction_files = discover_prediction_files(args.search_root, args.search_tag)
    if not prediction_files:
        raise FileNotFoundError(
            f"No prediction files were found under {args.search_root} containing '{args.search_tag}'."
        )

    manifest_df = load_manifest(args.manifest)
    label_map = load_label_map(args.dataset_path)

    fallback_df = load_prediction_table(prediction_files[0]) if not label_map else None
    holdout_df = build_holdout_frame(args.split_file, label_map, fallback_predictions=fallback_df)
    selection_df, test_df = split_nested_holdout(holdout_df, args.validation_fraction, args.seed)

    output_dir = args.output_dir
    output_dir.mkdir(parents=True, exist_ok=True)
    selection_path = output_dir / "nested_validation_split.tsv"
    test_path = output_dir / "nested_test_split.tsv"
    selection_df.to_csv(selection_path, sep="\t", index=False)
    test_df.to_csv(test_path, sep="\t", index=False)
    logger.info("Saved nested validation split to %s", selection_path)
    logger.info("Saved nested test split to %s", test_path)

    grouped_paths: Dict[str, Dict[int, Path]] = defaultdict(dict)
    for path in prediction_files:
        parsed = parse_trial_and_epoch(path)
        if parsed is None:
            logger.warning("Skipping unrecognised prediction file naming scheme: %s", path)
            continue
        trial_key, epoch = parsed
        grouped_paths[trial_key][epoch] = path

    if not grouped_paths:
        raise RuntimeError("No usable prediction files were discovered after parsing trial/epoch names.")

    selection_gene_ids = selection_df["Gene_ID"].astype(str).tolist()
    test_gene_ids = test_df["Gene_ID"].astype(str).tolist()

    epoch_records: List[Dict[str, object]] = []
    leaderboard_records: List[Dict[str, object]] = []

    for trial_key, epoch_map in sorted(grouped_paths.items()):
        example_path = next(iter(epoch_map.values()))
        manifest_info = match_manifest_row(manifest_df, trial_key, example_path)

        per_epoch_rows: List[Dict[str, object]] = []
        for epoch, path in sorted(epoch_map.items()):
            pred_df = load_prediction_table(path)
            selection_metrics = evaluate_subset(pred_df, selection_gene_ids, args.top_k)
            test_metrics = evaluate_subset(pred_df, test_gene_ids, args.top_k)

            row: Dict[str, object] = {
                "trial": trial_key,
                "epoch": int(epoch),
                "prediction_file": str(path),
                "validation_auprc": selection_metrics["auprc"],
                "validation_n": selection_metrics["n_samples"],
                "test_auprc": test_metrics["auprc"],
                "test_auroc": test_metrics["auroc"],
                "test_mcc": test_metrics["mcc"],
                "test_bacc": test_metrics["bacc"],
                f"test_precision_at_{args.top_k}": test_metrics[f"precision_at_{args.top_k}"],
                "test_n": test_metrics["n_samples"],
            }
            row.update(manifest_info)
            per_epoch_rows.append(row)
            epoch_records.append(dict(row))

        if not per_epoch_rows:
            continue

        epoch_df = pd.DataFrame(per_epoch_rows)
        valid_epoch_df = epoch_df.dropna(subset=["validation_auprc"]).sort_values(
            ["validation_auprc", "epoch"], ascending=[False, True]
        )
        if valid_epoch_df.empty:
            logger.warning("Skipping trial without valid validation metrics: %s", trial_key)
            continue

        best_row = valid_epoch_df.iloc[0].to_dict()
        best_row["selection_metric"] = "validation_auprc"
        leaderboard_records.append(best_row)

    if not leaderboard_records:
        raise RuntimeError("No trials produced a valid nested-holdout evaluation.")

    epoch_df = pd.DataFrame(epoch_records).sort_values(["trial", "epoch"]).reset_index(drop=True)
    leaderboard_df = (
        pd.DataFrame(leaderboard_records)
        .sort_values(["validation_auprc", "test_auprc", "epoch"], ascending=[False, False, True])
        .reset_index(drop=True)
    )

    epoch_path = output_dir / "random_search_per_epoch.tsv"
    leaderboard_path = output_dir / "random_search_leaderboard.tsv"
    best_json_path = output_dir / "best_trial.json"
    epoch_df.to_csv(epoch_path, sep="\t", index=False)
    leaderboard_df.to_csv(leaderboard_path, sep="\t", index=False)

    with open(best_json_path, "w") as fh:
        json.dump(make_json_safe(leaderboard_df.iloc[0].to_dict()), fh, indent=2)

    logger.info("Saved per-epoch results to %s", epoch_path)
    logger.info("Saved leaderboard to %s", leaderboard_path)
    logger.info("Saved best-trial summary to %s", best_json_path)
    logger.info(
        "Best trial: %s (epoch %s, validation AUPRC %.4f, test AUPRC %.4f)",
        leaderboard_df.iloc[0]["trial"],
        leaderboard_df.iloc[0]["epoch"],
        leaderboard_df.iloc[0]["validation_auprc"],
        leaderboard_df.iloc[0]["test_auprc"],
    )


if __name__ == "__main__":
    main()