_run_benchmark_full.py

"""Tier-1 benchmarks on ALL cells (no subsampling).

RNA pseudobulk is computed from the sparse matrix one donor at a time,
so the full 174k × 9660 dense matrix is never materialised.
Runs three feature sets: protein-only, RNA-only, RNA+protein.
"""
import gc, re, sys
from pathlib import Path

import anndata as ad
import muon as mu
import numpy as np
import pandas as pd
from scipy.sparse import issparse

sys.path.insert(0, str(Path(__file__).parent))

from data_utils import build_reactome_network
from binn.data import get_map
from binn.train import (
    pseudobulk_lr_benchmark,
    pseudobulk_xgb_benchmark,
    per_lineage_lr_benchmark,
    cell_composition_benchmark,
)

RNA_MOD     = "rna"
PROT_MOD    = "prot"
N_LEVELS    = 3
LINEAGE_COL = "manual.coarse.idents"
LINEAGE_MAP = {"nk": ["NK"], "t": ["CD4 T", "CD8 T"], "b": ["B", "PB"]}

def normalize_protein_name(name):
    return re.sub(r"-\d+$", "", str(name)).upper()


def sparse_pseudobulk(X_sparse, donor_labels, unique_donors):
    """Per-donor mean from sparse matrix, one donor at a time. Returns (n_donors, n_genes) float32."""
    donors_arr = np.array(donor_labels)
    out = np.zeros((len(unique_donors), X_sparse.shape[1]), dtype=np.float32)
    for i, donor in enumerate(unique_donors):
        mask = donors_arr == donor
        out[i] = np.asarray(X_sparse[mask].mean(axis=0)).flatten()
    return out


def make_prepseudобulked_adatas(pb_dict, donor_obs):
    """
    Wrap already-pseudobulked DataFrames as AnnData with one row per donor.
    The benchmark functions will groupby Donor and get the same matrix back.
    """
    adatas = {}
    for mod, df in pb_dict.items():
        obs = donor_obs.loc[df.index].copy()
        adata = ad.AnnData(X=df.values.astype(np.float32), obs=obs)
        adata.var_names = df.columns.tolist()
        adatas[mod] = adata
    return adatas


# ── 1. Load ──────────────────────────────────────────────────────────────────
print("[1/4] Loading MuData...")
mdata = mu.read("citeseq_logcounts.h5mu")
common_cells = sorted(set(mdata[RNA_MOD].obs_names) & set(mdata[PROT_MOD].obs_names))
obs_full = mdata[RNA_MOD].obs.loc[common_cells].copy()
print(f"  {len(common_cells)} cells, {obs_full['Donor'].nunique()} donors")

# ── 2. Reactome gene list ─────────────────────────────────────────────────────
print("[2/4] Building Reactome gene list...")
reactome_net  = build_reactome_network("train_data/reactome")
map_raw       = get_map(reactome_net, n_levels=N_LEVELS).copy()
map_raw["layer0"] = map_raw["layer0"].astype(str).str.upper()
reactome_genes = set(map_raw["layer0"])
del reactome_net
print(f"  {len(reactome_genes)} Reactome layer-0 genes")

# ── 3. Build pseudobulk matrices per lineage ─────────────────────────────────
print("[3/4] Computing pseudobulk per lineage (sparse RNA + dense protein)...")

# Protein: 174k × 209 dense is fine (~145 MB)
prot_raw       = mdata[PROT_MOD][common_cells]
norm_prot_names = [normalize_protein_name(c) for c in prot_raw.var_names]
prot_X          = prot_raw.X.toarray() if issparse(prot_raw.X) else np.asarray(prot_raw.X)
prot_df_cells   = pd.DataFrame(prot_X.astype(np.float32), index=common_cells, columns=norm_prot_names)
if prot_df_cells.columns.duplicated().any():
    prot_df_cells = prot_df_cells.T.groupby(level=0).mean().T
del prot_X; gc.collect()

# RNA: keep sparse, filter to Reactome genes only
rna_raw        = mdata[RNA_MOD][common_cells]
rna_names_upper = [str(c).upper() for c in rna_raw.var_names]
rna_keep_idx   = [i for i, n in enumerate(rna_names_upper) if n in reactome_genes]
rna_keep_names = [rna_names_upper[i] for i in rna_keep_idx]
rna_sparse     = rna_raw.X[:, rna_keep_idx]   # still sparse, 174k × 9660
print(f"  RNA sparse shape: {rna_sparse.shape}  "
      f"(nnz={rna_sparse.nnz:,}, density={rna_sparse.nnz/rna_sparse.shape[0]/rna_sparse.shape[1]:.2%})")

prot_cols = sorted(prot_df_cells.columns)
prot_df_cells = prot_df_cells[prot_cols]

# Pseudobulk per lineage
pb_rna  = {}   # mod -> donors × rna_genes  DataFrame
pb_prot = {}   # mod -> donors × prot_genes DataFrame
donor_obs_per_lineage = {}   # mod -> donors × obs_cols DataFrame

for name, labels in LINEAGE_MAP.items():
    cell_mask   = obs_full[LINEAGE_COL].astype(str).isin(labels)
    cell_ids    = obs_full.index[cell_mask]
    cell_pos    = [common_cells.index(c) for c in cell_ids]   # positions in common_cells
    lineage_obs = obs_full.loc[cell_ids]
    donors      = lineage_obs["Donor"].values
    unique_d    = sorted(set(donors))

    # donor-level obs (label + donor id)
    donor_obs_per_lineage[name] = (
        lineage_obs[["Donor", "Status"]]
        .drop_duplicates("Donor")
        .set_index("Donor")
        .loc[unique_d]
        .reset_index()
        .set_index("Donor")
    )

    # RNA pseudobulk (sparse, one donor at a time)
    rna_lin = rna_sparse[cell_pos]
    pb_rna_arr = sparse_pseudobulk(rna_lin, donors, unique_d)
    pb_rna[name] = pd.DataFrame(pb_rna_arr, index=unique_d, columns=rna_keep_names)

    # Protein pseudobulk (already dense slice)
    prot_lin = prot_df_cells.loc[cell_ids]
    pb_prot[name] = prot_lin.groupby(lineage_obs["Donor"].values).mean()[prot_cols]

    counts = pd.Series(donors).value_counts()
    print(f"  {name}: {len(cell_ids)} cells, {len(unique_d)} donors | "
          f"cells/donor: min={counts.min()}  median={counts.median():.0f}  max={counts.max()}")

del rna_sparse, prot_df_cells; gc.collect()

# ── Merge RNA + protein per lineage (handle overlapping gene names) ──────────
pb_combined = {}
for name in LINEAGE_MAP:
    rna_d  = pb_rna[name]
    prot_d = pb_prot[name]
    overlap = set(rna_d.columns) & set(prot_d.columns)
    # average overlapping names, keep rest separate
    if overlap:
        overlap_mean = (rna_d[list(overlap)].values + prot_d[list(overlap)].values) / 2
        overlap_df   = pd.DataFrame(overlap_mean, index=rna_d.index, columns=sorted(overlap))
        rna_only   = rna_d[[c for c in rna_d.columns  if c not in overlap]]
        prot_only  = prot_d[[c for c in prot_d.columns if c not in overlap]]
        pb_combined[name] = pd.concat([rna_only, prot_only, overlap_df], axis=1)
    else:
        pb_combined[name] = pd.concat([rna_d, prot_d], axis=1)

# ── Build a shared donor_obs covering all lineages ───────────────────────────
all_donor_obs = pd.concat(donor_obs_per_lineage.values()).reset_index()
all_donor_obs = all_donor_obs.drop_duplicates("Donor").set_index("Donor")

def wrap_adatas(pb_dict):
    """Wrap pre-computed pseudobulk DataFrames as AnnData (one row = one donor)."""
    adatas = {}
    for mod, df in pb_dict.items():
        obs = all_donor_obs.loc[df.index].copy()
        obs["Donor"] = obs.index  # benchmark needs Donor as a column
        adata = ad.AnnData(X=df.values.astype(np.float32), obs=obs)
        adata.var_names = [str(c) for c in df.columns]
        adatas[mod] = adata
    return adatas

adatas_rna      = wrap_adatas(pb_rna)
adatas_prot     = wrap_adatas(pb_prot)
adatas_combined = wrap_adatas(pb_combined)

# ── obs-only AnnData for cell composition ────────────────────────────────────
obs_adatas_obs = {}
for name, labels in LINEAGE_MAP.items():
    mask = obs_full[LINEAGE_COL].astype(str).isin(labels)
    obs_adatas_obs[name] = ad.AnnData(obs=obs_full[mask].copy())

# ── 4. Run benchmarks ─────────────────────────────────────────────────────────
print("\n[4/4] Running benchmarks...\n")

results = {}

for tag, adatas in [
    ("RNA only (Reactome-filtered)",   adatas_rna),
    ("Protein only",                   adatas_prot),
    ("RNA + Protein",                  adatas_combined),
]:
    print(f"=== {tag} ===")
    n_feat = sum(a.n_vars for a in adatas.values())

    print(f"  -- LR (ElasticNet) --")
    lr = pseudobulk_lr_benchmark(
        obs_adatas=adatas, label_col="Status", donor_col="Donor", n_splits=5, random_state=42,
    )

    print(f"  -- XGBoost --")
    xgb = pseudobulk_xgb_benchmark(
        obs_adatas=adatas, label_col="Status", donor_col="Donor", n_splits=5, random_state=42,
    )

    results[tag] = {"lr": lr, "xgb": xgb}

print("  === Per-lineage LR ===")
for tag, adatas in [
    ("RNA only",    adatas_rna),
    ("Protein only", adatas_prot),
]:
    print(f"  -- {tag} --")
    pl = per_lineage_lr_benchmark(
        obs_adatas=adatas, label_col="Status", donor_col="Donor", n_splits=5, random_state=42,
    )
    results[f"per_lineage_{tag}"] = pl

print("\n  === Cell composition ===")
comp = cell_composition_benchmark(
    obs_adatas=obs_adatas_obs, label_col="Status", donor_col="Donor",
    celltype_col="manual.fine.idents", n_splits=5, random_state=42,
)
results["composition"] = comp

# ── Summary ───────────────────────────────────────────────────────────────────
W = 62
print("\n" + "="*W)
print("TIER-1 BENCHMARK SUMMARY  (ALL CELLS)")
print("="*W)
print(f"  {'Model':<38} {'Acc':>10}  {'AUC':>10}")
print(f"  {'-'*38} {'-'*10}  {'-'*10}")

for feat_tag in ["RNA only (Reactome-filtered)", "Protein only", "RNA + Protein"]:
    r_lr  = results[feat_tag]["lr"]
    r_xgb = results[feat_tag]["xgb"]
    short = feat_tag.replace(" (Reactome-filtered)", "")
    print(f"  PB LR  {short:<31} {r_lr['mean_acc']:.3f}±{r_lr['std_acc']:.3f}  {r_lr['mean_auc']:.3f}±{r_lr['std_auc']:.3f}")
    print(f"  PB XGB {short:<31} {r_xgb['mean_acc']:.3f}±{r_xgb['std_acc']:.3f}  {r_xgb['mean_auc']:.3f}±{r_xgb['std_auc']:.3f}")

print(f"  {'-'*38} {'-'*10}  {'-'*10}")
for feat_tag in ["RNA only", "Protein only"]:
    pl = results[f"per_lineage_{feat_tag}"]
    for mod in ["nk", "t", "b"]:
        r = pl[mod]
        print(f"  Per-lineage LR {feat_tag} {mod:<14} {r['mean_acc']:.3f}±{r['std_acc']:.3f}  {r['mean_auc']:.3f}±{r['std_auc']:.3f}")

print(f"  {'-'*38} {'-'*10}  {'-'*10}")
r = results["composition"]
print(f"  Cell composition only (9 features)     {r['mean_acc']:.3f}±{r['std_acc']:.3f}  {r['mean_auc']:.3f}±{r['std_auc']:.3f}")
print("="*W)