PDGrapher/scripts/eval_drug_retrieval.py

"""Drug retrieval evaluation for DrugRank-Flow.

Loads a trained Phase-1 checkpoint, encodes all 189 SciPlex3 drugs into a gallery,
then retrieves the correct drug from unseen test conditions and reports:

  A-class retrieval  : Hit@1/5/10, MRR, NDCG@10, Median Rank
  B-class scPerturBench : PCC-delta, Energy Distance, Common DEGs@50
  MOA-AUC            : ROC-AUC for ranking same-MOA drugs first
  Bootstrap 95% CI   : 1000 resample iterations on retrieval metrics

Usage
-----
    python scripts/eval_drug_retrieval.py \\
        --checkpoint outputs/drug_rank/phase1_best.pt \\
        --config configs/drug_rank_phase1.yaml \\
        --split drug_disjoint \\
        --output results/drug_retrieval/ \\
        [--moa-stratified] [--device cuda] [--batch-size 16]
"""

from __future__ import annotations

import argparse
import csv
import json
import logging
import os
import sys
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import yaml

# ── Project path ─────────────────────────────────────────────────────────────
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from gidflow.models.population_encoder import PopulationEncoder
from gidflow.models.gap_encoder import GapEncoder
from gidflow.models.drug_encoder import DrugEncoder
from gidflow.models.drug_gene_bridge import DrugGeneBridge

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger(__name__)

# ── Constants ─────────────────────────────────────────────────────────────────
CELL_LINE_MAP: Dict[str, int] = {"A549": 1, "K562": 2, "MCF7": 3}
ANNOTATION_DIR = Path("/data/boom/ICLR/data/annotation")
SPLITS_DIR = Path("/data/boom/ICLR/data/splits")

# ── Argument parsing ──────────────────────────────────────────────────────────

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Drug retrieval evaluation for DrugRank-Flow")
    p.add_argument("--checkpoint", default="outputs/drug_rank/phase1_best.pt")
    p.add_argument("--config", default="configs/drug_rank_phase1.yaml")
    p.add_argument("--split", default="drug_disjoint",
                   help="Split name (file in data/splits/<split>.json)")
    p.add_argument("--output", default="results/drug_retrieval/",
                   help="Directory for output JSON and CSV")
    p.add_argument("--device", default="auto",
                   help="cuda | cpu | auto")
    p.add_argument("--batch-size", type=int, default=16,
                   help="Batch size for query encoding")
    p.add_argument("--gallery-batch-size", type=int, default=32,
                   help="Batch size when building drug gallery")
    p.add_argument("--n-bootstrap", type=int, default=1000,
                   help="Bootstrap iterations for CI estimation")
    p.add_argument("--max-cells", type=int, default=200,
                   help="Max cells per condition for e-distance (tractability)")
    p.add_argument("--moa-stratified", action="store_true",
                   help="Compute per-MOA-class breakdown of Hit@1 / MRR")
    p.add_argument("--no-pcc", action="store_true",
                   help="Skip PCC-delta and e-distance (faster eval)")
    return p.parse_args()


# ── Model loading ─────────────────────────────────────────────────────────────

def build_models(cfg: dict, device: torch.device):
    """Instantiate all four sub-models from config."""
    m = cfg["model"]
    num_proteins = len(json.load(open(ANNOTATION_DIR / "protein_target_vocab.json")))

    source_enc = PopulationEncoder(
        num_genes=m["num_genes"],
        hidden_dim=m["encoder_hidden"],
        output_dim=m["encoder_output"],
    ).to(device)

    target_enc = PopulationEncoder(
        num_genes=m["num_genes"],
        hidden_dim=m["encoder_hidden"],
        output_dim=m["encoder_output"],
    ).to(device)

    gap_enc = GapEncoder(
        input_dim=m["encoder_output"],
        hidden_dim=m["gap_hidden"],
        output_dim=m["gap_output"],
        proj_dim=m["gap_proj_dim"],
        num_cell_lines=m["num_cell_lines"],
        num_genes=m["num_genes"],  # enable reconstruct_expression
    ).to(device)

    drug_enc = DrugEncoder(
        encoding=m.get("drug_encoder", "morgan"),
        emb_dim=m["drug_emb_dim"],
        freeze=True,
    ).to(device)

    bridge = DrugGeneBridge(
        num_proteins=num_proteins,
        drug_emb_dim=m["drug_emb_dim"],
        hidden_dim=m["bridge_hidden_dim"],
        proj_dim=m["bridge_proj_dim"],
        protein_emb_dim=m["bridge_protein_emb_dim"],
    ).to(device)

    return source_enc, target_enc, gap_enc, drug_enc, bridge


def load_checkpoint(ckpt_path: str, models: tuple, device: torch.device) -> None:
    """Load state_dicts from checkpoint into (source_enc, target_enc, gap_enc, drug_enc, bridge)."""
    source_enc, target_enc, gap_enc, drug_enc, bridge = models
    log.info("Loading checkpoint: %s", ckpt_path)
    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)

    source_enc.load_state_dict(ckpt["source_enc"])
    target_enc.load_state_dict(ckpt["target_enc"])
    gap_enc.load_state_dict(ckpt["gap_enc"])
    drug_enc.load_state_dict(ckpt["drug_enc"])
    bridge.load_state_dict(ckpt["bridge"])

    for m in models:
        m.eval()
    log.info("All models loaded and set to eval().")


# ── Drug gallery ──────────────────────────────────────────────────────────────

def load_drug_smiles(annotation_dir: Path) -> Dict[str, str]:
    """Build drug_name → SMILES mapping from drug_annotation_master.csv."""
    smiles_map: Dict[str, str] = {}
    csv_path = annotation_dir / "drug_annotation_master.csv"
    with open(csv_path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            name = row.get("drug_name", "").strip()
            smi = row.get("smiles", "").strip()
            if name and smi:
                smiles_map[name] = smi
    log.info("Loaded SMILES for %d drugs", len(smiles_map))
    return smiles_map


@torch.no_grad()
def build_drug_gallery(
    drug_order: List[str],
    smiles_map: Dict[str, str],
    drug_enc: DrugEncoder,
    bridge: DrugGeneBridge,
    batch_size: int = 32,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """Compute drug_proj for every drug in drug_order.

    Returns
    -------
    gallery : [189, 128]  (proj_dim)
    """
    n_drugs = len(drug_order)
    proj_dim = bridge.proj_dim
    gallery = torch.zeros(n_drugs, proj_dim, device=device)

    missing = 0
    for start in range(0, n_drugs, batch_size):
        batch_names = drug_order[start : start + batch_size]
        batch_smiles = []
        valid_mask = []
        for name in batch_names:
            smi = smiles_map.get(name, "")
            batch_smiles.append(smi if smi else "C")  # placeholder for missing
            valid_mask.append(bool(smi))

        try:
            drug_emb = drug_enc(batch_smiles)          # [B, emb_dim]
            out = bridge(drug_emb)
            projs = out["drug_proj"]                   # [B, 128]
        except Exception as e:
            log.warning("Gallery batch %d failed: %s — using zeros", start, e)
            continue

        for i, (proj, valid) in enumerate(zip(projs, valid_mask)):
            idx = start + i
            if valid:
                gallery[idx] = proj
            else:
                gallery[idx] = torch.zeros(proj_dim, device=device)
                missing += 1

    log.info(
        "Gallery built: %d drugs, %d missing SMILES (zero embeddings)",
        n_drugs, missing,
    )
    return gallery


# ── Data helpers ──────────────────────────────────────────────────────────────

def load_sciplex3_raw(h5ad_path: str, num_genes: int = 2000):
    """Load SciPlex3 h5ad and return (X_hvg [n_cells, G], obs DataFrame).

    Applies the same normalization as Sciplex3Dataset:
    library-size normalize → log1p → top-variance HVG selection.
    """
    log.info("Loading SciPlex3 from %s ...", h5ad_path)
    try:
        import anndata as ad
        import scipy.sparse as sp
    except ImportError:
        raise ImportError("pip install anndata scipy")

    adata = ad.read_h5ad(h5ad_path)
    log.info("  raw shape: %s", adata.shape)

    obs = adata.obs.copy()

    # filter out dose=0 for drug-treated (keeps vehicle too for source lookup)
    # We keep ALL cells here; we split into vehicle/drug in query building
    X_raw = adata.X.toarray() if sp.issparse(adata.X) else np.array(adata.X)

    # Library-size normalize
    lib_sizes = X_raw.sum(axis=1, keepdims=True).clip(min=1)
    X_norm = np.log1p(X_raw / lib_sizes * 1e4).astype(np.float32)
    del X_raw

    # HVG selection by variance
    var = X_norm.var(axis=0)
    hvg_idx = np.argsort(var)[::-1][:num_genes]
    X_hvg = X_norm[:, hvg_idx].astype(np.float32)

    log.info("  Cells: %d, HVGs: %d", X_hvg.shape[0], X_hvg.shape[1])
    return X_hvg, obs


def parse_pair_id(pair_id: str) -> Tuple[str, str, float]:
    """Parse pair_id like 'vorinostat_A549_10.0' into (drug, cell_line, dose).

    Strategy: known cell lines are used as anchors to split the string.
    """
    for cl in ["A549", "K562", "MCF7"]:
        marker = f"_{cl}_"
        pos = pair_id.find(marker)
        if pos != -1:
            drug_name = pair_id[:pos]
            rest = pair_id[pos + len(marker):]
            try:
                dose = float(rest)
            except ValueError:
                dose = float("nan")
            return drug_name, cl, dose
    # Fallback: last two underscore-separated tokens are cell_line and dose
    parts = pair_id.rsplit("_", 2)
    if len(parts) == 3:
        return parts[0], parts[1], float(parts[2])
    return pair_id, "unknown", float("nan")


def build_queries_from_dataset(
    cfg: dict,
    test_pair_ids: List[str],
    drug_order: List[str],
    seed: int = 42,
) -> List[Dict]:
    """Build eval queries by REUSING the exact training Sciplex3Dataset.

    This guarantees the HVG gene basis, library-size normalization, log1p
    ordering, and gene ordering are byte-for-byte identical to what the model
    saw during training. The previous free-standing loader selected HVGs on
    the log-normalized matrix (training selects on the pre-log matrix), which
    silently fed the encoder a DIFFERENT 2000-gene basis and corrupted all
    retrieval metrics. It also densified the full 799k x 111k matrix (330 GiB).

    Parameters
    ----------
    cfg           : loaded YAML config (uses data.* and model.num_genes)
    test_pair_ids : list of "<drug>_<cell_line>_<dose>" ids from the split file
    drug_order    : ordered drug names (index = gallery row / true_drug_idx)
    seed          : RNG seed for cell sampling

    Returns
    -------
    list of query dicts with keys:
        pair_id, drug_name, cell_line, dose, true_drug_idx,
        source_cells [Ns, G], target_cells [Nt, G]
    """
    from gidflow.data.sciplex_dataset import Sciplex3Dataset

    data_cfg = cfg["data"]
    model_cfg = cfg["model"]
    annotation_dir = data_cfg["annotation_dir"]
    smiles_csv = os.path.join(annotation_dir, "drug_annotation_master.csv")
    if not os.path.exists(smiles_csv):
        smiles_csv = ""

    max_source_cells = int(data_cfg.get("max_source_cells", 64))
    max_target_cells = int(data_cfg.get("max_target_cells", 64))

    log.info("Instantiating Sciplex3Dataset (identical preprocessing to training) ...")
    dataset = Sciplex3Dataset(
        h5ad_path=data_cfg["sciplex3_h5ad"],
        n_hvg=model_cfg["num_genes"],
        max_source_cells=max_source_cells,
        max_target_cells=max_target_cells,
        seed=seed,
        drug_emb_dim=model_cfg["drug_emb_dim"],
        preprocessed_path=None,          # force raw h5ad (all 3 cell lines)
        drug_smiles_csv=smiles_csv,
    )

    X = dataset._X  # [n_cells, G] float32, training gene basis
    drug_to_idx = {name: i for i, name in enumerate(drug_order)}

    # Map "<drug>_<cell_line>_<dose>" -> condition index (matches training script)
    pair_id_to_idx: Dict[str, int] = {}
    for i, cond in enumerate(dataset._conditions):
        pid = f"{cond['drug_name']}_{cond.get('cell_line', '')}_{cond.get('dose', '')}"
        pair_id_to_idx[pid] = i

    rng = np.random.default_rng(seed)
    queries: List[Dict] = []
    skipped_unmatched = 0
    skipped_nodrug = 0

    for pair_id in test_pair_ids:
        idx = pair_id_to_idx.get(pair_id)
        if idx is None:
            skipped_unmatched += 1
            continue
        cond = dataset._conditions[idx]
        # Strip trailing/leading whitespace: 11 SciPlex3 drug names carry a
        # trailing space in the h5ad ('Busulfan ', 'Mesna ', ...) while
        # drug_order.json stores them stripped. Without this the true index
        # lookup returns None and these 11 drugs are silently dropped.
        drug_name = cond["drug_name"].strip()
        if drug_name not in drug_to_idx:
            log.warning("Drug not in drug_order: %s (pair_id=%s)", drug_name, pair_id)
            skipped_nodrug += 1
            continue

        veh_rows = np.asarray(cond["vehicle_cell_idx"])
        drug_rows = np.asarray(cond["drug_cell_idx"])
        if len(veh_rows) == 0 or len(drug_rows) == 0:
            skipped_unmatched += 1
            continue

        ns = min(max_source_cells, len(veh_rows))
        nt = min(max_target_cells, len(drug_rows))
        chosen_src = rng.choice(veh_rows, size=ns, replace=False)
        chosen_tgt = rng.choice(drug_rows, size=nt, replace=False)

        queries.append({
            "pair_id": pair_id,
            "drug_name": drug_name,
            "cell_line": cond["cell_line"],
            "dose": float(cond["dose"]),
            "true_drug_idx": drug_to_idx[drug_name],
            "source_cells": np.asarray(X[chosen_src], dtype=np.float32),
            "target_cells": np.asarray(X[chosen_tgt], dtype=np.float32),
        })

    log.info(
        "Queries built from dataset: %d matched, %d unmatched pair_ids, %d drug-missing (of %d test)",
        len(queries), skipped_unmatched, skipped_nodrug, len(test_pair_ids),
    )
    return queries


def build_queries(
    test_pair_ids: List[str],
    drug_order: List[str],
    X_hvg: np.ndarray,
    obs,
    max_source_cells: int = 64,
    max_target_cells: int = 64,
    seed: int = 42,
) -> List[Dict]:
    """[DEPRECATED — kept for reference] Match each pair_id to cells in the dataset.

    Returns list of dicts with keys:
        drug_name, cell_line, dose, true_drug_idx,
        source_cells [Ns, G], target_cells [Nt, G]
    Missing/unmatched queries are skipped.
    """
    rng = np.random.default_rng(seed)
    drug_to_idx = {name: i for i, name in enumerate(drug_order)}

    # Pre-index by (perturbation, cell_line, dose_value) for fast lookup
    obs = obs.copy()
    obs["_row"] = np.arange(len(obs))

    # vehicle rows per cell_line
    vehicle_mask = (
        obs["perturbation"].str.lower().str.contains("vehicle", na=False)
        | (obs["dose_value"] == 0)
    )
    vehicle_by_cl = {}
    for cl in ["A549", "K562", "MCF7"]:
        rows = obs["_row"][vehicle_mask & (obs["cell_line"] == cl)].values
        if len(rows) > 0:
            vehicle_by_cl[cl] = rows

    queries = []
    skipped = 0

    for pair_id in test_pair_ids:
        drug_name, cell_line, dose = parse_pair_id(pair_id)

        if drug_name not in drug_to_idx:
            log.warning("Drug not in drug_order: %s (pair_id=%s)", drug_name, pair_id)
            skipped += 1
            continue

        true_drug_idx = drug_to_idx[drug_name]

        # Locate drug-treated cells
        drug_rows = obs["_row"][
            (~vehicle_mask)
            & (obs["perturbation"] == drug_name)
            & (obs["cell_line"] == cell_line)
            & (np.abs(obs["dose_value"] - dose) < 1e-3)
        ].values

        if len(drug_rows) == 0:
            log.debug("No drug cells for %s (dose=%.1f, cl=%s) — skipping", drug_name, dose, cell_line)
            skipped += 1
            continue

        # Locate vehicle (source) cells for same cell line
        src_rows = vehicle_by_cl.get(cell_line, np.array([], dtype=int))
        if len(src_rows) == 0:
            # Fallback: any vehicle
            src_rows = obs["_row"][vehicle_mask].values

        if len(src_rows) == 0:
            log.warning("No vehicle cells found for cell_line=%s", cell_line)
            skipped += 1
            continue

        # Sample cells
        ns = min(max_source_cells, len(src_rows))
        nt = min(max_target_cells, len(drug_rows))
        chosen_src = rng.choice(src_rows, size=ns, replace=False)
        chosen_tgt = rng.choice(drug_rows, size=nt, replace=False)

        queries.append({
            "pair_id": pair_id,
            "drug_name": drug_name,
            "cell_line": cell_line,
            "dose": dose,
            "true_drug_idx": true_drug_idx,
            "source_cells": X_hvg[chosen_src],  # [Ns, G]
            "target_cells": X_hvg[chosen_tgt],  # [Nt, G]
        })

    log.info(
        "Queries built: %d matched, %d skipped out of %d total",
        len(queries), skipped, len(test_pair_ids),
    )
    return queries


# ── Encoding ──────────────────────────────────────────────────────────────────

@torch.no_grad()
def encode_queries(
    queries: List[Dict],
    source_enc: PopulationEncoder,
    target_enc: PopulationEncoder,
    gap_enc: GapEncoder,
    batch_size: int = 16,
    device: torch.device = torch.device("cpu"),
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Encode all test queries.

    Returns
    -------
    gap_embs    : [N, 128]
    true_indices: [N] int64
    query_meta  : list of N dicts (drug_name, cell_line, dose, ...)
    """
    all_gap_embs: List[torch.Tensor] = []
    all_src_means: List[np.ndarray] = []
    all_tgt_means: List[np.ndarray] = []
    all_src_cells: List[np.ndarray] = []
    all_tgt_cells: List[np.ndarray] = []
    all_true_indices: List[int] = []
    query_meta: List[Dict] = []

    def _pad_cells(cell_list: List[np.ndarray]) -> Tuple[torch.Tensor, torch.Tensor]:
        """Pad cell arrays to same N and return (cells [B,N,G], mask [B,N])."""
        max_n = max(c.shape[0] for c in cell_list)
        G = cell_list[0].shape[1]
        cells_t = torch.zeros(len(cell_list), max_n, G)
        mask_t = torch.zeros(len(cell_list), max_n, dtype=torch.bool)
        for i, c in enumerate(cell_list):
            n = c.shape[0]
            cells_t[i, :n, :] = torch.from_numpy(c)
            mask_t[i, :n] = True
        return cells_t.to(device), mask_t.to(device)

    for start in range(0, len(queries), batch_size):
        batch = queries[start : start + batch_size]

        src_list = [q["source_cells"] for q in batch]
        tgt_list = [q["target_cells"] for q in batch]

        src_t, src_mask = _pad_cells(src_list)   # [B, Ns, G]
        tgt_t, tgt_mask = _pad_cells(tgt_list)   # [B, Nt, G]

        # IMPORTANT: training (phase1 & phase2) NEVER passes cell_line_ids, so the
        # cell_line embedding rows for A549/K562/MCF7 were never trained (random
        # init). Passing them here injected untrained noise and HALVED retrieval
        # (Hit@10 0.30 -> 0.15). Match training: do not condition on cell line.
        z_src = source_enc(src_t, src_mask)       # [B, H]
        z_tgt = target_enc(tgt_t, tgt_mask)       # [B, H]
        gap_out = gap_enc(z_src, z_tgt)
        gap_emb = gap_out["gap_emb"]              # [B, 128]

        all_gap_embs.append(gap_emb.cpu())

        # Store per-cell arrays for B-class metrics
        for q, src_arr, tgt_arr in zip(batch, src_list, tgt_list):
            all_src_means.append(src_arr.mean(axis=0))
            all_tgt_means.append(tgt_arr.mean(axis=0))
            all_src_cells.append(src_arr)
            all_tgt_cells.append(tgt_arr)
            all_true_indices.append(q["true_drug_idx"])
            query_meta.append({k: v for k, v in q.items()
                                if k not in ("source_cells", "target_cells")})

    gap_embs = torch.cat(all_gap_embs, dim=0)       # [N, 128]
    true_indices = torch.tensor(all_true_indices, dtype=torch.long)

    return gap_embs, true_indices, query_meta, all_src_means, all_tgt_means, all_src_cells, all_tgt_cells


# ── Ranking ───────────────────────────────────────────────────────────────────

def rank_drugs(
    gap_embs: torch.Tensor,
    gallery: torch.Tensor,
    true_indices: torch.Tensor,
) -> np.ndarray:
    """Compute rank of the true drug for each query.

    Returns
    -------
    ranks : [N] int (1-indexed)
    all_scores : [N, 189]
    """
    scores = gap_embs @ gallery.T.to(gap_embs.device)   # [N, 189]
    ranked_indices = torch.argsort(scores, dim=-1, descending=True)  # [N, 189]

    ranks = []
    for i in range(len(true_indices)):
        true_idx = true_indices[i].item()
        # position of true_idx in ranked_indices[i]
        pos = (ranked_indices[i] == true_idx).nonzero(as_tuple=True)[0]
        if len(pos) == 0:
            rank = len(gallery)  # worst case
        else:
            rank = pos[0].item() + 1  # 1-indexed
        ranks.append(rank)

    return np.array(ranks, dtype=int), scores.cpu().numpy()


# ── Metrics ───────────────────────────────────────────────────────────────────

def hit_at_k(ranks: np.ndarray, k: int) -> float:
    return float((ranks <= k).mean())


def mrr(ranks: np.ndarray) -> float:
    return float((1.0 / ranks).mean())


def ndcg_at_10(ranks: np.ndarray) -> float:
    """NDCG@10 assuming a single relevant item per query."""
    # ideal DCG = 1 / log2(2) = 1.0  (best possible rank = 1)
    ideal_dcg = 1.0 / np.log2(2)
    dcgs = np.where(ranks <= 10, 1.0 / np.log2(ranks + 1), 0.0)
    return float((dcgs / ideal_dcg).mean())


def pcc_delta(
    gap_emb: torch.Tensor,
    gap_enc: GapEncoder,
    src_means: List[np.ndarray],
    tgt_means: List[np.ndarray],
    device: torch.device,
    batch_size: int = 64,
) -> float:
    """Mean per-sample Pearson correlation between predicted and true delta expression."""
    from scipy.stats import pearsonr

    pccs = []
    gap_enc.eval()
    with torch.no_grad():
        for start in range(0, len(src_means), batch_size):
            emb_batch = gap_emb[start : start + batch_size].to(device)
            pred_deltas = gap_enc.reconstruct_expression(emb_batch).cpu().numpy()
            for i, (src_m, tgt_m, pred_d) in enumerate(
                zip(src_means[start:start+batch_size],
                    tgt_means[start:start+batch_size],
                    pred_deltas)
            ):
                true_d = tgt_m - src_m
                if true_d.std() < 1e-8 or pred_d.std() < 1e-8:
                    continue
                r, _ = pearsonr(pred_d, true_d)
                pccs.append(r)

    return float(np.mean(pccs)) if pccs else float("nan")


def energy_distance(X: np.ndarray, Y: np.ndarray) -> float:
    """Energy distance D_E(X, Y) between two sets of vectors.

    D_E(X,Y) = 2/(n*m)*sum_ij||Xi-Yj|| - 1/n^2*sum_ij||Xi-Xj|| - 1/m^2*sum_ij||Yi-Yj||
    """
    n, m = len(X), len(Y)
    if n == 0 or m == 0:
        return float("nan")

    def mean_pairwise_dist(A: np.ndarray, B: np.ndarray) -> float:
        # Efficient vectorized pairwise L2 using broadcasting
        # For large arrays this can be memory-heavy; chunks help
        chunk = 50
        total = 0.0
        count = 0
        for i in range(0, len(A), chunk):
            Ai = A[i : i + chunk]
            diff = Ai[:, None, :] - B[None, :, :]  # [ci, len(B), G]
            total += np.sqrt((diff ** 2).sum(axis=-1)).sum()
            count += Ai.shape[0] * B.shape[0]
        return total / count if count > 0 else 0.0

    cross = mean_pairwise_dist(X, Y)
    self_x = mean_pairwise_dist(X, X)
    self_y = mean_pairwise_dist(Y, Y)
    return float(2 * cross - self_x - self_y)


def compute_e_distance(
    gap_emb: torch.Tensor,
    gap_enc: GapEncoder,
    src_cells: List[np.ndarray],
    tgt_cells: List[np.ndarray],
    device: torch.device,
    max_cells: int = 200,
    batch_size: int = 64,
) -> float:
    """Mean energy distance across queries."""
    edists = []
    gap_enc.eval()
    with torch.no_grad():
        for start in range(0, len(src_cells), batch_size):
            emb_batch = gap_emb[start : start + batch_size].to(device)
            pred_deltas = gap_enc.reconstruct_expression(emb_batch).cpu().numpy()

            for i, (src_arr, tgt_arr, pred_d) in enumerate(
                zip(src_cells[start:start+batch_size],
                    tgt_cells[start:start+batch_size],
                    pred_deltas)
            ):
                src_arr = src_arr[:max_cells]
                tgt_arr = tgt_arr[:max_cells]
                # Predicted cells: source cells shifted by predicted delta
                pred_cells = src_arr + pred_d[np.newaxis, :]  # broadcast
                ed = energy_distance(pred_cells, tgt_arr)
                edists.append(ed)

    return float(np.mean(edists)) if edists else float("nan")


def common_degs_at_50(
    gap_emb: torch.Tensor,
    gap_enc: GapEncoder,
    src_means: List[np.ndarray],
    tgt_means: List[np.ndarray],
    device: torch.device,
    batch_size: int = 64,
) -> float:
    """Fraction of top-50 predicted DEGs that overlap with true top-50 DEGs."""
    overlaps = []
    gap_enc.eval()
    with torch.no_grad():
        for start in range(0, len(src_means), batch_size):
            emb_batch = gap_emb[start : start + batch_size].to(device)
            pred_deltas = gap_enc.reconstruct_expression(emb_batch).cpu().numpy()

            for i, (src_m, tgt_m, pred_d) in enumerate(
                zip(src_means[start:start+batch_size],
                    tgt_means[start:start+batch_size],
                    pred_deltas)
            ):
                true_d = tgt_m - src_m
                k = min(50, len(true_d))
                pred_top = set(np.argsort(np.abs(pred_d))[::-1][:k])
                true_top = set(np.argsort(np.abs(true_d))[::-1][:k])
                overlap = len(pred_top & true_top) / k
                overlaps.append(overlap)

    return float(np.mean(overlaps)) if overlaps else float("nan")


def compute_moa_auc(
    scores_all: np.ndarray,
    true_indices: np.ndarray,
    moa_mask: np.ndarray,
) -> float:
    """Mean per-query ROC-AUC for ranking same-MOA drugs first.

    Queries where the drug has no same-MOA peers (row sum <= 1) are skipped.
    """
    from sklearn.metrics import roc_auc_score

    aucs = []
    for i, true_idx in enumerate(true_indices):
        moa_labels = moa_mask[true_idx].copy()
        # Exclude the drug itself from the labels
        moa_labels[true_idx] = 0
        if moa_labels.sum() == 0:
            continue
        try:
            auc = roc_auc_score(moa_labels, scores_all[i])
            aucs.append(auc)
        except Exception:
            pass

    return float(np.mean(aucs)) if aucs else float("nan")


# ── Bootstrap CI ──────────────────────────────────────────────────────────────

def bootstrap_ci(
    ranks: np.ndarray,
    n_iter: int = 1000,
    seed: int = 42,
) -> Dict[str, Dict[str, float]]:
    """Bootstrap 95% CI for retrieval metrics."""
    rng = np.random.default_rng(seed)
    N = len(ranks)

    h1_boot, h5_boot, h10_boot, mrr_boot, ndcg_boot = [], [], [], [], []

    for _ in range(n_iter):
        idx = rng.integers(0, N, size=N)
        r = ranks[idx]
        h1_boot.append(hit_at_k(r, 1))
        h5_boot.append(hit_at_k(r, 5))
        h10_boot.append(hit_at_k(r, 10))
        mrr_boot.append(mrr(r))
        ndcg_boot.append(ndcg_at_10(r))

    def ci(arr: List[float]) -> Dict[str, float]:
        a = np.array(arr)
        return {
            "mean": float(a.mean()),
            "ci_lo": float(np.percentile(a, 2.5)),
            "ci_hi": float(np.percentile(a, 97.5)),
        }

    return {
        "hit@1": ci(h1_boot),
        "hit@5": ci(h5_boot),
        "hit@10": ci(h10_boot),
        "mrr": ci(mrr_boot),
        "ndcg@10": ci(ndcg_boot),
    }


# ── MOA-stratified breakdown ──────────────────────────────────────────────────

def moa_stratified_breakdown(
    query_meta: List[Dict],
    ranks: np.ndarray,
    drug_order: List[str],
    annotation_dir: Path,
) -> Dict[str, Dict]:
    """Per-MOA breakdown of Hit@1 and MRR."""
    csv_path = annotation_dir / "drug_annotation_master.csv"
    drug_to_moa: Dict[str, str] = {}
    with open(csv_path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            drug_to_moa[row["drug_name"].strip()] = row.get("moa_class", "Unknown").strip()

    moa_groups: Dict[str, List[int]] = {}
    for i, meta in enumerate(query_meta):
        moa = drug_to_moa.get(meta["drug_name"], "Unknown")
        moa_groups.setdefault(moa, []).append(i)

    breakdown: Dict[str, Dict] = {}
    for moa, idxs in sorted(moa_groups.items()):
        r = ranks[np.array(idxs)]
        breakdown[moa] = {
            "n_queries": len(r),
            "hit@1": round(hit_at_k(r, 1), 4),
            "mrr": round(mrr(r), 4),
            "median_rank": float(np.median(r)),
        }

    return breakdown


# ── Main ──────────────────────────────────────────────────────────────────────

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

    # ── Device ───────────────────────────────────────────────────────────────
    if args.device == "auto":
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    log.info("Using device: %s", device)

    # ── Config ────────────────────────────────────────────────────────────────
    cfg_path = Path(args.config)
    if not cfg_path.is_absolute():
        cfg_path = Path(__file__).resolve().parents[1] / cfg_path
    with open(cfg_path) as f:
        cfg = yaml.safe_load(f)

    # ── Build & load models ───────────────────────────────────────────────────
    models = build_models(cfg, device)
    source_enc, target_enc, gap_enc, drug_enc, bridge = models

    ckpt_path = Path(args.checkpoint)
    if not ckpt_path.is_absolute():
        ckpt_path = Path(__file__).resolve().parents[1] / ckpt_path
    load_checkpoint(str(ckpt_path), models, device)

    # ── Drug order & SMILES ───────────────────────────────────────────────────
    drug_order_raw: List[str] = json.load(open(ANNOTATION_DIR / "drug_order.json"))
    moa_mask_raw: np.ndarray = np.load(ANNOTATION_DIR / "moa_mask.npy")  # [189, 189]

    # Drop non-drug placeholder rows (e.g. the 189th 'control'/vehicle row, which
    # has no SMILES -> methane 'C' and is a spurious retrieval competitor that
    # inflates the rank denominator to 189). Keep only real drugs and slice the
    # MOA mask to match, so the gallery denominator is the true 188 drugs.
    _NON_DRUG = {"control", "vehicle", "Vehicle", "DMSO", ""}
    keep_idx = [i for i, d in enumerate(drug_order_raw) if d not in _NON_DRUG]
    drug_order: List[str] = [drug_order_raw[i] for i in keep_idx]
    moa_mask: np.ndarray = moa_mask_raw[np.ix_(keep_idx, keep_idx)]
    n_dropped = len(drug_order_raw) - len(drug_order)
    log.info("Gallery drugs: %d (dropped %d non-drug rows: %s)",
             len(drug_order), n_dropped,
             [drug_order_raw[i] for i in range(len(drug_order_raw)) if i not in set(keep_idx)])

    smiles_map = load_drug_smiles(ANNOTATION_DIR)

    # ── Gallery ───────────────────────────────────────────────────────────────
    gallery = build_drug_gallery(
        drug_order=drug_order,
        smiles_map=smiles_map,
        drug_enc=drug_enc,
        bridge=bridge,
        batch_size=args.gallery_batch_size,
        device=device,
    )  # [188, 128]
    n_valid_smiles = sum(1 for d in drug_order if smiles_map.get(d))
    log.info("Gallery built with %d/%d drugs having SMILES", n_valid_smiles, len(drug_order))

    # ── Load split ────────────────────────────────────────────────────────────
    split_path = SPLITS_DIR / f"{args.split}.json"
    with open(split_path) as f:
        split_data = json.load(f)
    test_pair_ids: List[str] = split_data["test"]
    log.info("Test set size: %d pairs (split=%s)", len(test_pair_ids), args.split)

    # ── Build queries (reuse training Sciplex3Dataset for identical basis) ─────
    queries = build_queries_from_dataset(
        cfg=cfg,
        test_pair_ids=test_pair_ids,
        drug_order=drug_order,
        seed=cfg.get("seed", 42),
    )

    if len(queries) == 0:
        log.error("No queries matched — check split and dataset alignment.")
        return

    # ── Encode queries ────────────────────────────────────────────────────────
    log.info("Encoding %d test queries ...", len(queries))
    gap_embs, true_indices, query_meta, src_means, tgt_means, src_cells, tgt_cells = encode_queries(
        queries=queries,
        source_enc=source_enc,
        target_enc=target_enc,
        gap_enc=gap_enc,
        batch_size=args.batch_size,
        device=device,
    )

    # ── Rank drugs ────────────────────────────────────────────────────────────
    log.info("Ranking drugs ...")
    ranks, scores_all = rank_drugs(gap_embs, gallery.cpu(), true_indices)
    log.info(
        "Median rank: %.1f  |  Hit@1: %.3f  |  MRR: %.4f",
        float(np.median(ranks)),
        hit_at_k(ranks, 1),
        mrr(ranks),
    )

    # ── A-class retrieval metrics ─────────────────────────────────────────────
    retrieval_metrics: Dict[str, float] = {
        "hit@1": round(hit_at_k(ranks, 1), 4),
        "hit@5": round(hit_at_k(ranks, 5), 4),
        "hit@10": round(hit_at_k(ranks, 10), 4),
        "mrr": round(mrr(ranks), 4),
        "ndcg@10": round(ndcg_at_10(ranks), 4),
        "median_rank": round(float(np.median(ranks)), 2),
        "mean_rank": round(float(ranks.mean()), 2),
        "n_queries": int(len(ranks)),
    }

    # ── Bootstrap CI ──────────────────────────────────────────────────────────
    log.info("Computing bootstrap CI (%d iterations) ...", args.n_bootstrap)
    bootstrap = bootstrap_ci(ranks, n_iter=args.n_bootstrap, seed=cfg.get("seed", 42))

    # ── B-class scPerturBench metrics ─────────────────────────────────────────
    perturbench_metrics: Dict[str, float] = {}
    if not args.no_pcc:
        log.info("Computing PCC-delta ...")
        perturbench_metrics["pcc_delta"] = round(
            pcc_delta(gap_embs, gap_enc, src_means, tgt_means, device=device), 4
        )
        log.info("Computing common DEGs@50 ...")
        perturbench_metrics["common_degs_50"] = round(
            common_degs_at_50(gap_embs, gap_enc, src_means, tgt_means, device=device), 4
        )
        log.info("Computing energy distance ...")
        perturbench_metrics["e_distance"] = round(
            compute_e_distance(
                gap_embs, gap_enc, src_cells, tgt_cells,
                device=device, max_cells=args.max_cells
            ), 4
        )
    else:
        log.info("Skipping PCC / e-distance (--no-pcc)")
        perturbench_metrics = {"pcc_delta": None, "common_degs_50": None, "e_distance": None}

    # ── MOA-AUC ──────────────────────────────────────────────────────────────
    log.info("Computing MOA-AUC ...")
    moa_auc = compute_moa_auc(scores_all, true_indices.numpy(), moa_mask)
    perturbench_metrics["moa_auc"] = round(moa_auc, 4) if not np.isnan(moa_auc) else None

    # ── MOA-stratified breakdown ──────────────────────────────────────────────
    moa_breakdown: Optional[Dict] = None
    if args.moa_stratified:
        log.info("Computing MOA-stratified breakdown ...")
        moa_breakdown = moa_stratified_breakdown(query_meta, ranks, drug_order, ANNOTATION_DIR)

    # ── Output ────────────────────────────────────────────────────────────────
    out_dir = Path(args.output)
    if not out_dir.is_absolute():
        out_dir = Path(__file__).resolve().parents[1] / out_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    # JSON metrics
    results = {
        "split": args.split,
        "checkpoint": str(ckpt_path),
        "n_queries": int(len(ranks)),
        "retrieval": retrieval_metrics,
        "bootstrap_ci": bootstrap,
        "perturbench": perturbench_metrics,
    }
    if moa_breakdown is not None:
        results["moa_stratified"] = moa_breakdown

    metrics_path = out_dir / f"{args.split}_metrics.json"
    with open(metrics_path, "w") as f:
        json.dump(results, f, indent=2)
    log.info("Metrics saved: %s", metrics_path)

    # CSV per-query rankings
    rankings_path = out_dir / f"{args.split}_rankings.csv"
    with open(rankings_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["pair_id", "drug_name", "cell_line", "dose", "true_drug_idx", "rank"])
        for meta, rank in zip(query_meta, ranks):
            writer.writerow([
                meta["pair_id"],
                meta["drug_name"],
                meta["cell_line"],
                meta["dose"],
                meta["true_drug_idx"],
                int(rank),
            ])
    log.info("Rankings saved: %s", rankings_path)

    # ── Summary print ─────────────────────────────────────────────────────────
    print("\n" + "=" * 60)
    print(f"  DrugRank-Flow Evaluation  |  split={args.split}")
    print("=" * 60)
    print(f"  Queries : {len(ranks)}")
    print(f"  Hit@1   : {retrieval_metrics['hit@1']:.4f}  "
          f"[{bootstrap['hit@1']['ci_lo']:.4f}, {bootstrap['hit@1']['ci_hi']:.4f}]")
    print(f"  Hit@5   : {retrieval_metrics['hit@5']:.4f}  "
          f"[{bootstrap['hit@5']['ci_lo']:.4f}, {bootstrap['hit@5']['ci_hi']:.4f}]")
    print(f"  Hit@10  : {retrieval_metrics['hit@10']:.4f}  "
          f"[{bootstrap['hit@10']['ci_lo']:.4f}, {bootstrap['hit@10']['ci_hi']:.4f}]")
    print(f"  MRR     : {retrieval_metrics['mrr']:.4f}  "
          f"[{bootstrap['mrr']['ci_lo']:.4f}, {bootstrap['mrr']['ci_hi']:.4f}]")
    print(f"  NDCG@10 : {retrieval_metrics['ndcg@10']:.4f}  "
          f"[{bootstrap['ndcg@10']['ci_lo']:.4f}, {bootstrap['ndcg@10']['ci_hi']:.4f}]")
    print(f"  Median R: {retrieval_metrics['median_rank']}")
    if not args.no_pcc:
        print(f"  PCC-delta   : {perturbench_metrics.get('pcc_delta')}")
        print(f"  E-distance  : {perturbench_metrics.get('e_distance')}")
        print(f"  DEGs@50     : {perturbench_metrics.get('common_degs_50')}")
    print(f"  MOA-AUC     : {perturbench_metrics.get('moa_auc')}")
    print("=" * 60)
    print(f"  Saved: {metrics_path}")
    print(f"  Saved: {rankings_path}")


if __name__ == "__main__":
    main()