src/probes.py

"""
Novel foundation-model probes for quality <-> pathology disentanglement.

These operationalise the core research question: *does a retinal foundation
model conflate image degradation with disease pathology?*

1. Degradation Sensitivity Probe (DSP)
   Apply controlled, monotonic degradations (defocus, illumination gradient,
   contrast loss) at increasing severity. Track how FLAIR's confidence in the
   originally-predicted disease moves as image quality collapses. A *true*
   finding is degradation-robust; a quality-driven artefact decays or the
   'ungradable' mass rises. We summarise this with an Entanglement Index.

2. Occlusion Spatial Disentanglement
   Occlusion saliency localises the pixels driving FLAIR's disease call. We
   measure how much of that evidence sits on regions the QC pipeline flags as
   degraded (a spatial confound score).

All probes degrade gracefully when FLAIR is unavailable.
"""

from __future__ import annotations
import numpy as np
import cv2


# ----------------------------------------------------------- controlled degradations
def _defocus(rgb, sev):
    if sev <= 0:
        return rgb
    return cv2.GaussianBlur(rgb, (0, 0), 0.6 + 3.2 * sev)


def _illumination(rgb, sev, fov):
    if sev <= 0:
        return rgb
    h, w = rgb.shape[:2]
    yy, xx = np.mgrid[0:h, 0:w].astype(np.float32)
    cx, cy, r = fov["cx"], fov["cy"], max(fov["radius"], 1)
    d = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2) / r
    shade = 1.0 - sev * np.clip(d - 0.15, 0, 1) * 0.9        # vignetting shadow
    return np.clip(rgb.astype(np.float32) * shade[..., None], 0, 255).astype(np.uint8)


def _contrast(rgb, sev):
    if sev <= 0:
        return rgb
    m = rgb.reshape(-1, 3).mean(0)
    k = 1.0 - 0.75 * sev
    return np.clip((rgb.astype(np.float32) - m) * k + m, 0, 255).astype(np.uint8)


DEGRADATIONS = {"defocus": _defocus, "illumination": _illumination,
                "contrast": _contrast}


# --------------------------------------------------------------- degradation probe
def degradation_sensitivity(engine, rgb, fov, qc_composite_fn, diseases=None,
                            levels=6, kinds=("defocus", "illumination", "contrast")):
    """Sweep degradations; track FLAIR disease confidence vs QC quality.

    qc_composite_fn(rgb_uint8) -> float 0-100 (the geometric QC composite).
    Returns a dict of per-kind curves plus an aggregate entanglement index.
    """
    if not engine.load():
        return None
    diseases = diseases or engine_default_diseases(engine)

    # baseline top disease on the clean image
    base = engine.zero_shot_disease(rgb, diseases)
    top = next((d["label"] for d in base if d["label"] != "normal"), base[0]["label"])
    sev_grid = np.linspace(0, 1, levels)

    curves = {}
    all_disease, all_qc = [], []
    for kind in kinds:
        fn = DEGRADATIONS[kind]
        dz_prob, ungr, qc = [], [], []
        for s in sev_grid:
            im = fn(rgb, s) if kind != "illumination" else fn(rgb, s, fov)
            ranked = engine.zero_shot_disease(im, diseases)
            pmap = {d["label"]: d["prob"] for d in ranked}
            dz_prob.append(float(pmap.get(top, 0.0)))
            ungr.append(float(engine.ungradable_prob(im)))
            qc.append(float(qc_composite_fn(im)))
        curves[kind] = dict(severity=sev_grid.tolist(), disease_prob=dz_prob,
                            ungradable=ungr, qc=qc)
        all_disease += dz_prob
        all_qc += qc

    # Entanglement index: how strongly disease confidence co-moves with quality.
    a = np.array(all_disease); q = np.array(all_qc)
    if a.std() < 1e-6 or q.std() < 1e-6:
        ent = 0.0
    else:
        ent = float(np.clip(np.corrcoef(a, q)[0, 1], -1, 1))
    # positive corr => disease confidence falls with quality => entangled
    entanglement = max(ent, 0.0)
    robustness = 1.0 - entanglement

    return dict(top_disease=top, curves=curves,
                entanglement_index=entanglement, robustness=robustness,
                verdict=("Disease read is quality-entangled - interpret with caution"
                         if entanglement > 0.4 else
                         "Disease read is largely quality-robust"))


def engine_default_diseases(engine):
    from .flair_wrapper import DEFAULT_DISEASES
    return list(DEFAULT_DISEASES)


# --------------------------------------------------------- occlusion disentanglement
def occlusion_disentanglement(engine, rgb, fov, degradation_map, grid=7,
                              diseases=None):
    """Occlusion saliency for the top disease, and its spatial overlap with the
    QC degradation map. Returns saliency map, overlay, and a confound score."""
    if not engine.load():
        return None
    diseases = diseases or engine_default_diseases(engine)
    ranked = engine.zero_shot_disease(rgb, diseases)
    top = next((d["label"] for d in ranked if d["label"] != "normal"),
               ranked[0]["label"])
    base_p = {d["label"]: d["prob"] for d in ranked}[top]

    h, w = rgb.shape[:2]
    gh, gw = h // grid, w // grid
    sal = np.zeros((grid, grid), np.float32)
    mean_rgb = rgb.reshape(-1, 3).mean(0)
    for i in range(grid):
        for j in range(grid):
            occ = rgb.copy()
            occ[i*gh:(i+1)*gh, j*gw:(j+1)*gw] = mean_rgb
            p = {d["label"]: d["prob"]
                 for d in engine.zero_shot_disease(occ, diseases)}.get(top, 0.0)
            sal[i, j] = max(base_p - p, 0.0)                 # drop => important

    sal = cv2.resize(sal, (w, h), interpolation=cv2.INTER_CUBIC)
    sal = np.clip(sal, 0, None)
    if sal.max() > 0:
        sal /= sal.max()

    # spatial overlap with degradation
    dmap = degradation_map.astype(np.float32)
    if dmap.max() > 0:
        dmap /= dmap.max()
    mask = None
    from .qc_metrics import inner_disc_mask
    mask = inner_disc_mask(rgb.shape, fov, frac=0.95)
    s, d = sal[mask], dmap[mask]
    confound = float((s * d).sum() / (s.sum() + 1e-6))       # frac of evidence on degraded px

    return dict(top_disease=top, saliency=sal, confound=confound,
                note=("Disease evidence overlaps degraded regions - the call may "
                      "be a quality artefact." if confound > 0.35 else
                      "Disease evidence sits on clean retina."))