src/batch_effects.py

"""
Batch-effect analysis for retinal image cohorts.

When fundus images come from different cameras, sites or acquisition days, they
carry systematic technical variation ("batch effects") that can dominate the
biological signal - a well-known confounder in imaging cohorts and foundation-
model embeddings. This module provides:

  * feature extraction (FLAIR embeddings when available, else interpretable
    QC + colour descriptors),
  * batch-effect *detection* (a batch-classifier AUC + silhouette by batch),
  * batch-effect *visualisation* (PCA / optional UMAP scatter, per-feature
    distributions),
  * batch-effect *correction* via a self-contained ComBat (empirical-Bayes
    location/scale harmonisation) plus a simple z-standardisation option.
"""

from __future__ import annotations
import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_predict
from sklearn.metrics import roc_auc_score, silhouette_score


# ----------------------------------------------------------------- feature extraction
def descriptor_from_metrics(rgb, metrics, fov):
    """Interpretable per-image feature vector: the QC scores plus low-order colour
    / intensity statistics. Used as the embedding when FLAIR is not loaded."""
    import cv2
    from .qc_metrics import inner_disc_mask
    m = inner_disc_mask(rgb.shape, fov, 0.9)
    if m.sum() < 50:
        m = np.ones(rgb.shape[:2], bool)
    feats = [mm["score"] for mm in metrics]
    for c in range(3):
        ch = rgb[..., c][m].astype(np.float32)
        feats += [ch.mean() / 255, ch.std() / 255]
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
    feats += [hsv[..., 0][m].mean() / 180, hsv[..., 1][m].mean() / 255]
    feats += [fov["coverage"], fov["radius"] / max(rgb.shape[:2])]
    return np.array(feats, np.float32)


# ------------------------------------------------------------------------ detection
def detect_batch_effect(X, batches):
    """Quantify how strongly `batches` are separable in feature space.

    Returns dict with:
      auc        : cross-validated one-vs-rest batch-classifier AUC (0.5 == none)
      silhouette : silhouette of the batch labelling in PCA space (-1..1)
      severity   : qualitative label
    """
    X = np.asarray(X, np.float32)
    b = np.asarray(batches)
    uniq = np.unique(b)
    out = dict(auc=np.nan, silhouette=np.nan, severity="n/a")
    if len(uniq) < 2 or X.shape[0] < 6:
        out["severity"] = "insufficient data (need >=2 batches, >=6 images)"
        return out
    Xs = StandardScaler().fit_transform(X)

    # Batch classifier AUC (macro one-vs-rest)
    try:
        y = np.searchsorted(uniq, b)
        clf = LogisticRegression(max_iter=1000)
        proba = cross_val_predict(clf, Xs, y, cv=min(4, len(uniq) + 1),
                                  method="predict_proba")
        if proba.shape[1] == 2:
            auc = roc_auc_score(y, proba[:, 1])
        else:
            auc = roc_auc_score(y, proba, multi_class="ovr", average="macro")
        out["auc"] = float(auc)
    except Exception:
        pass

    try:
        emb = PCA(n_components=min(10, Xs.shape[1])).fit_transform(Xs)
        out["silhouette"] = float(silhouette_score(emb, b))
    except Exception:
        pass

    auc = out["auc"]
    if not np.isnan(auc):
        out["severity"] = ("negligible" if auc < 0.6 else
                           "mild" if auc < 0.75 else
                           "moderate" if auc < 0.9 else "severe")
    return out


def embed_2d(X, method="pca", seed=0):
    """2-D embedding for scatter plots. UMAP if available, else PCA."""
    X = StandardScaler().fit_transform(np.asarray(X, np.float32))
    if method == "umap":
        try:
            import umap
            n = X.shape[0]
            reducer = umap.UMAP(n_neighbors=min(15, max(2, n - 1)),
                                min_dist=0.1, random_state=seed)
            return reducer.fit_transform(X), "UMAP"
        except Exception:
            pass
    return PCA(n_components=2).fit_transform(X), "PCA"


# ------------------------------------------------------------------------ correction
def combat(X, batches, covariates=None):
    """Self-contained parametric ComBat harmonisation.

    Removes additive (location) and multiplicative (scale) batch effects with an
    empirical-Bayes shrinkage, optionally preserving biological covariates.
    Follows Johnson et al. 2007. X is (n_samples, n_features).
    """
    X = np.asarray(X, np.float64)
    b = np.asarray(batches)
    n, p = X.shape
    uniq = np.unique(b)
    if len(uniq) < 2:
        return X.astype(np.float32)

    # Design: intercept + optional covariates
    if covariates is not None and len(np.unique(covariates)) > 1:
        cov = np.asarray(covariates)
        cats = np.unique(cov)
        C = np.column_stack([(cov == c).astype(float) for c in cats])
        design = C  # already spans intercept
    else:
        design = np.ones((n, 1))

    # Standardise using pooled (grand) model
    B = np.linalg.lstsq(design, X, rcond=None)[0]
    grand = design @ B
    var = ((X - grand) ** 2).mean(axis=0) + 1e-8
    sd = np.sqrt(var)
    Z = (X - grand) / sd

    Xc = X.copy()
    for g in uniq:
        idx = np.where(b == g)[0]
        if len(idx) < 2:
            continue
        Zi = Z[idx]
        gamma_hat = Zi.mean(axis=0)                      # additive batch shift
        delta_hat = Zi.var(axis=0) + 1e-8                # multiplicative batch scale

        # Empirical-Bayes priors across features
        gbar, t2 = gamma_hat.mean(), gamma_hat.var() + 1e-8
        dbar, dvar = delta_hat.mean(), delta_hat.var() + 1e-8
        a = (2 * dvar + dbar ** 2) / dvar
        bpr = (dbar * dvar + dbar ** 3) / dvar

        # Iterative EB estimates (few fixed-point iterations)
        gamma_star, delta_star = gamma_hat.copy(), delta_hat.copy()
        for _ in range(30):
            gamma_star = (len(idx) * t2 * gamma_hat + delta_star * gbar) / \
                         (len(idx) * t2 + delta_star)
            resid = ((Zi - gamma_star) ** 2).sum(axis=0)
            delta_star = (0.5 * resid + bpr) / (0.5 * len(idx) + a - 1)
            delta_star = np.clip(delta_star, 1e-6, None)

        Zi_adj = (Zi - gamma_star) / np.sqrt(delta_star)
        Xc[idx] = Zi_adj * sd + grand[idx]

    return Xc.astype(np.float32)


def zstandardise_by_batch(X, batches):
    """Simpler baseline: per-batch z-standardisation to a common mean/var."""
    X = np.asarray(X, np.float64).copy()
    b = np.asarray(batches)
    gmean, gsd = X.mean(0), X.std(0) + 1e-8
    for g in np.unique(b):
        idx = b == g
        mu, sd = X[idx].mean(0), X[idx].std(0) + 1e-8
        X[idx] = (X[idx] - mu) / sd * gsd + gmean
    return X.astype(np.float32)