scripts/analysis/102_train_on_dingwall_variantA.py

"""variant A — semi-supervised panda on dingwall using data s3 marker panels with score+margin gate."""
from __future__ import annotations
from pathlib import Path
import warnings, json, sys, time
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import anndata as ad
import scanpy as sc
import scipy.sparse as sp
from scipy.stats import fisher_exact
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

import os as _os
from pathlib import Path as _Path
PANDA_ROOT = _Path(_os.environ.get("PANDA_ROOT", str(_Path(__file__).resolve().parents[2])))
sys.path.insert(0, str(PANDA_ROOT))
from panda.model import (
    PANDAEncoder, supcon_loss, vicreg_loss, hsic_biased, subcenter_angular_infonce
)

ROOT = Path(str(PANDA_ROOT))
RAW_H5 = ROOT / "data/raw/GSE220977_combined.h5ad"
DERM_MARKERS = ROOT / "data/external_labels/dingwall_supp/biorxiv_media-3.xlsx"
OUT_DIR = ROOT / "discovery/pan_skin/marker"
CK_DIR = ROOT / "checkpoints/pan_skin_dingwall_variantA"

# Dingwall GSM -> genotype (see 101_primary_eden_derm_scoring)
CKO_GSMS = {"GSM6833482", "GSM6833483"}  # CORRECTED: 480/481 are rttaControl (WT), not cKO
WT_GSMS  = {"GSM6833478", "GSM6833479", "GSM6833480", "GSM6833481"}  # CORRECTED: 4 Cre-neg controls per GEO metadata

TOP_N = 30                     # markers per Derm panel for scoring
SCORE_MIN = 0.10               # min score to accept a pseudo-label
MARGIN_MIN = 0.05              # min gap best - runner-up
N_HVG = 2000                   # matches paper
N_PCA = 40                     # matches paper
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Training config (mirrors 20_train_panda.py)
GUARANTEED_PER_CLASS = 6
NATURAL_SLOTS = 96
STAGE_EPOCHS = [15, 25, 40, 40]
BALANCE_MIX = 0.5


# ---------- data prep ----------

def load_derm_panels(top_n: int = TOP_N) -> dict[int, list[str]]:
    df = pd.read_excel(DERM_MARKERS)
    df = df.sort_values(["cluster", "avg_log2FC"], ascending=[True, False])
    return {int(c): df[df["cluster"] == c].head(top_n)["gene"].tolist()
            for c in sorted(df["cluster"].unique())}


def load_dingwall_dermal() -> ad.AnnData:
    a = ad.read_h5ad(RAW_H5)
    sample = a.obs["sample"].astype(str)
    a.obs["genotype"] = np.where(sample.isin(list(CKO_GSMS)), "En1-cKO",
                        np.where(sample.isin(list(WT_GSMS)), "WT", "other"))
    a = a[a.obs["genotype"].isin(["WT", "En1-cKO"])].copy()
    return a


def preprocess_paper_style(a: ad.AnnData) -> ad.AnnData:
    """lognormalize + hvg(2000) + pca(40) + harmony per-sample, matches dingwall STAR methods."""
    sc.pp.filter_genes(a, min_cells=10)
    sc.pp.normalize_total(a, target_sum=1e4)
    sc.pp.log1p(a)
    sc.pp.highly_variable_genes(a, n_top_genes=N_HVG, flavor="seurat", batch_key="sample")
    sc.pp.scale(a, max_value=10, zero_center=False)
    sc.tl.pca(a, n_comps=N_PCA, use_highly_variable=True, zero_center=False)
    try:
        import harmonypy as hm  # noqa
        sc.external.pp.harmony_integrate(a, key="sample", basis="X_pca",
                                         adjusted_basis="X_pca_harmony", max_iter_harmony=20)
        a.obsm["X_train"] = a.obsm["X_pca_harmony"]
    except Exception as exc:
        print(f"[preprocess] harmony skipped ({exc}); using raw PCA", flush=True)
        a.obsm["X_train"] = a.obsm["X_pca"]
    return a


# ---------- pseudo-labelling ----------

def score_and_gate(a: ad.AnnData, panels: dict[int, list[str]],

                   score_min: float = SCORE_MIN,

                   margin_min: float = MARGIN_MIN) -> ad.AnnData:
    """score cells on 12 derm panels; accept label if best>score_min and margin>margin_min."""
    for cl, genes in panels.items():
        present = [g for g in genes if g in a.var_names]
        if not present:
            a.obs[f"derm{cl}_score"] = 0.0
        else:
            sc.tl.score_genes(a, gene_list=present, score_name=f"derm{cl}_score",
                              random_state=0, use_raw=False)
    cols = [f"derm{cl}_score" for cl in sorted(panels)]
    S = a.obs[cols].values
    top1_ix = S.argmax(axis=1)
    top1 = S[np.arange(len(S)), top1_ix]
    S_copy = S.copy(); S_copy[np.arange(len(S)), top1_ix] = -np.inf
    top2 = S_copy.max(axis=1)
    margin = top1 - top2
    accept = (top1 > score_min) & (margin > margin_min)

    ids = np.array([int(cols[i].replace("derm", "").replace("_score", "")) for i in top1_ix])
    a.obs["derm_pseudo"] = ids
    a.obs["derm_pseudo_top1"] = top1
    a.obs["derm_pseudo_margin"] = margin
    a.obs["derm_pseudo_accept"] = accept
    return a


# ---------- PANDA training (mirrors 20_train_panda.py) ----------

class CorpusDataset(Dataset):
    def __init__(self, X, y, d, aux):
        self.X = X.astype(np.float32); self.y = y.astype(np.int64)
        self.d = d.astype(np.int64); self.aux = aux.astype(np.float32)
    def __len__(self): return self.X.shape[0]
    def __getitem__(self, i):
        return (torch.from_numpy(self.X[i]), torch.tensor(self.y[i]),
                torch.tensor(self.d[i]), torch.from_numpy(self.aux[i]))


class HybridSampler:
    def __init__(self, y, n_batches=100, seed=0):
        self.y = np.asarray(y); self.n_batches = n_batches
        self.rng = np.random.default_rng(seed)
        self.classes = np.unique(self.y)
        self.by_cls = {int(c): np.where(self.y == c)[0] for c in self.classes}
        counts = np.bincount(self.y, minlength=int(self.classes.max()) + 1).astype(float)
        self.natural_p = counts / counts.sum()
    def __iter__(self):
        for _ in range(self.n_batches):
            batch = []
            for c in self.classes:
                idx = self.by_cls[int(c)]
                take = min(GUARANTEED_PER_CLASS, len(idx))
                if take > 0:
                    batch.extend(self.rng.choice(idx, size=take, replace=(len(idx) < take)).tolist())
            for _ in range(NATURAL_SLOTS):
                c = self.rng.choice(len(self.natural_p), p=self.natural_p)
                idx = self.by_cls.get(int(c), self.by_cls[int(self.classes[0])])
                batch.append(int(self.rng.choice(idx)))
            yield batch
    def __len__(self): return self.n_batches


def train_panda(X_tr, y_tr, d_tr, aux_tr, n_classes, n_datasets, ck_out: Path):
    ck_out.mkdir(parents=True, exist_ok=True)
    counts = np.bincount(y_tr, minlength=n_classes)
    inv_sqrt = 1.0 / np.sqrt(counts + 1); inv_sqrt = inv_sqrt / inv_sqrt.mean()
    class_w = BALANCE_MIX * inv_sqrt + (1 - BALANCE_MIX) * np.ones_like(inv_sqrt)
    class_w = torch.tensor(class_w, dtype=torch.float32, device=DEVICE)

    ds = CorpusDataset(X_tr, y_tr, d_tr, aux_tr)
    loader = DataLoader(ds, batch_sampler=HybridSampler(y_tr, n_batches=100), num_workers=0)

    model = PANDAEncoder(variant="pca", n_pca=X_tr.shape[1], n_classes=n_classes,
                         n_datasets=n_datasets).to(DEVICE)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)

    for stage, n_ep in enumerate(STAGE_EPOCHS):
        print(f"[panda-A] stage {stage} ({n_ep} epochs)", flush=True)
        for e in range(n_ep):
            t0 = time.time(); losses = []
            for X_b, y_b, d_b, aux_b in loader:
                X_b = X_b.to(DEVICE); y_b = y_b.to(DEVICE); d_b = d_b.to(DEVICE); aux_b = aux_b.to(DEVICE)
                lam = 1.0 if stage >= 2 else 0.0
                out = model(X_b, aux_b, lam_dann=lam)
                L_supcon = supcon_loss(out["z"], y_b)
                L_vic = vicreg_loss(out["z"])
                L_ce = F.cross_entropy(out["logits"], y_b, weight=class_w, label_smoothing=0.05)
                total = L_supcon + 1.0 * L_vic + 0.4 * L_ce
                if stage >= 1:
                    proto_ref = model.prototypes.detach().clone()
                    total = total + 0.6 * subcenter_angular_infonce(out["z"], y_b, proto_ref)
                if stage >= 2:
                    total = total + F.cross_entropy(out["dom"], d_b)
                    total = total + 0.3 * F.mse_loss(out["depth"].squeeze(1), aux_b[:, 1])
                    total = total + 0.05 * hsic_biased(out["repr"], aux_b[:, 1:2])
                opt.zero_grad(); total.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
                opt.step()
                if stage >= 1:
                    model.update_prototypes(out["z"].detach(), y_b)
                losses.append(float(total.item()))
            if (e + 1) % 5 == 0:
                print(f"  ep {e+1}/{n_ep} loss={np.mean(losses):.3f} dt={time.time()-t0:.1f}s", flush=True)
        torch.save({"model": model.state_dict()}, ck_out / f"panda_stage{stage}.pt")

    torch.save({"model": model.state_dict(),
                "prototypes": model.prototypes.detach().cpu().numpy()},
               ck_out / "panda_final.pt")
    return model


@torch.no_grad()
def infer(model, X, aux):
    model.eval()
    Xt = torch.from_numpy(X.astype(np.float32)).to(DEVICE)
    at = torch.from_numpy(aux.astype(np.float32)).to(DEVICE)
    B = 4096; preds = []; confs = []
    for i in range(0, len(Xt), B):
        out = model(Xt[i:i+B], at[i:i+B])
        p = F.softmax(out["logits"], dim=1)
        preds.append(p.argmax(dim=1).cpu().numpy())
        confs.append(p.max(dim=1).values.cpu().numpy())
    return np.concatenate(preds), np.concatenate(confs)


# ---------- reporting ----------

def report_depletion(labels: np.ndarray, genotype: np.ndarray, n_classes: int) -> pd.DataFrame:
    n_wt = int((genotype == "WT").sum()); n_cko = int((genotype == "En1-cKO").sum())
    base = n_cko / max(n_wt + n_cko, 1)
    rows = []
    for c in range(n_classes):
        m = labels == c
        w = int(((genotype == "WT") & m).sum()); k = int(((genotype == "En1-cKO") & m).sum())
        if w + k == 0: continue
        try:
            odds, p = fisher_exact([[w, n_wt - w], [k, n_cko - k]], alternative="two-sided")
        except ValueError:
            odds, p = 1.0, 1.0
        rows.append({"derm_id": c, "n": w + k, "n_WT": w, "n_cKO": k,
                     "cko_frac": k / (w + k), "baseline": base,
                     "odds_ratio": float(odds), "fisher_p": float(p)})
    return pd.DataFrame(rows).sort_values("cko_frac")


def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True); CK_DIR.mkdir(parents=True, exist_ok=True)

    print("[A] load panels + dermal Dingwall", flush=True)
    panels = load_derm_panels()
    a = load_dingwall_dermal()
    # reuse existing panda-v3 fibroblast calls if present, else all cells
    pred_csv = ROOT / "discovery/pan_skin/marker/dingwall_predictions.csv"
    if pred_csv.exists():
        pred = pd.read_csv(pred_csv)
        pred_map = dict(zip(pred["cell_id"].astype(str), pred["pred_label"]))
        a.obs["v3_label"] = [pred_map.get(c, "unknown") for c in a.obs_names.astype(str)]
        a = a[np.isin(a.obs["v3_label"], ["fibroblast-papillary", "fibroblast-reticular"])].copy()
        print(f"[A] restricted to PANDA-v3 fibroblasts: n={a.n_obs}", flush=True)

    print("[A] paper-style preprocess", flush=True)
    a = preprocess_paper_style(a)

    print("[A] score + gate pseudo-labels", flush=True)
    a = score_and_gate(a, panels)
    n_acc = int(a.obs["derm_pseudo_accept"].sum())
    print(f"[A] pseudo-label acceptance: {n_acc}/{a.n_obs} ({100*n_acc/a.n_obs:.1f}%)", flush=True)

    # train/heldout split (gate = train; rest = infer)
    train_mask = a.obs["derm_pseudo_accept"].values.astype(bool)
    X_all = np.asarray(a.obsm["X_train"])
    y_all = a.obs["derm_pseudo"].astype(int).values
    sample_ix = {s: i for i, s in enumerate(sorted(a.obs["sample"].astype(str).unique()))}
    d_all = np.array([sample_ix[s] for s in a.obs["sample"].astype(str)])
    aux_all = np.stack([np.zeros(a.n_obs, dtype=np.float32),
                        np.log10(np.asarray(a.X.sum(axis=1)).ravel() + 1)], axis=1)
    aux_all[:, 1] = (aux_all[:, 1] - aux_all[:, 1].mean()) / (aux_all[:, 1].std() + 1e-6)

    classes = sorted(np.unique(y_all[train_mask]).tolist())
    if len(classes) < 2:
        print("[A] not enough classes accepted; abort", flush=True); return
    cls_ix = {c: i for i, c in enumerate(classes)}
    y_all_ix = np.array([cls_ix.get(int(c), -1) for c in y_all])
    y_tr = y_all_ix[train_mask]
    X_tr = X_all[train_mask]; d_tr = d_all[train_mask]; aux_tr = aux_all[train_mask]

    print(f"[A] train n={train_mask.sum()} on {len(classes)} classes: {classes}", flush=True)
    model = train_panda(X_tr, y_tr, d_tr, aux_tr, n_classes=len(classes),
                        n_datasets=len(sample_ix), ck_out=CK_DIR)

    # inference on held-out
    infer_mask = ~train_mask
    preds_ix, confs = infer(model, X_all[infer_mask], aux_all[infer_mask])
    preds_derm = np.array([classes[p] for p in preds_ix])

    # combine: use pseudo-label on train, prediction on inference
    final = np.where(train_mask, y_all,
                     np.concatenate([y_all[train_mask].astype(int) * 0 - 1,  # placeholder
                                     preds_derm.astype(int)])[:a.n_obs] if False else 0)
    # simpler: assemble directly
    final = y_all.astype(int).copy()
    final[infer_mask] = preds_derm.astype(int)

    df = pd.DataFrame({
        "cell_id": a.obs_names.astype(str).values,
        "sample": a.obs["sample"].astype(str).values,
        "genotype": a.obs["genotype"].astype(str).values,
        "pseudo_derm": y_all,
        "pseudo_accept": train_mask,
        "final_derm": final,
    })
    df.to_csv(OUT_DIR / "102_variantA_predictions.csv", index=False)

    dep = report_depletion(final, a.obs["genotype"].values, n_classes=12)
    dep.to_csv(OUT_DIR / "102_variantA_depletion.csv", index=False)
    summary = {
        "variant": "A_semi_supervised_S3_scoring",
        "score_min": SCORE_MIN, "margin_min": MARGIN_MIN, "top_n": TOP_N,
        "n_total": int(a.n_obs), "n_train_pseudo": int(train_mask.sum()),
        "classes_trained": classes,
        "derm10": dep[dep["derm_id"] == 10].to_dict("records"),
        "derm2": dep[dep["derm_id"] == 2].to_dict("records"),
        "derm9": dep[dep["derm_id"] == 9].to_dict("records"),
        "all": dep.to_dict("records"),
    }
    (OUT_DIR / "102_variantA_summary.json").write_text(json.dumps(summary, indent=2, default=str))
    print(f"[A] done -> {OUT_DIR}/102_variantA_*", flush=True)


if __name__ == "__main__":
    main()