"""Vascular leakage quantification from FITC-dextran / IB4 dual-label images.

Designed for confocal images (orthogonal projections or single planes) of
cochlear / stria vascularis capillary networks labeled with:

  * FITC-dextran (green): intravascular tracer that extravasates into the
    surrounding tissue when the vascular barrier is compromised.
  * Isolectin B4 - 594 (red): endothelial / vessel wall marker.
  * Optional DAPI (blue): ignored.

Pipeline
--------
1.  Channel extraction. Merged RGB images and 2-3 plane stacks supported.
    Both channels are scaled by one common divisor so their intensities stay
    mutually comparable. Burned-in annotations (scale bars, text) are
    detected as co-saturated pixels and excluded.
2.  Tissue mask: the specimen footprint. Dark interstitium enclosed by
    tissue is kept (leak-free tissue must count in the denominator); wide
    camera-background regions connected to the frame border - the black
    corners of rotated acquisitions, empty space beside the specimen - are
    excluded.
3.  Vessel segmentation. IB4 staining quality varies strongly between
    specimens, so the vascular compartment is the union of two detectors:
      a. IB4 branch - background-flattened hysteresis threshold of the red
         channel.
      b. FITC-tube branch - multi-scale Sato vesselness on the green
         channel, gated by intensity. Perfused capillaries are bright
         *tubes* of FITC; extravasated dye is spatially diffuse and has no
         tubular geometry, so this branch recovers vessels that IB4 missed
         without absorbing genuine leak.
    A small dilation assigns the vessel wall to the intravascular
    compartment, and a guard-band annulus just beyond it is excluded from
    both compartments: it carries the point-spread halo of the bright
    lumen, which would otherwise read as leak even in leak-free controls.
4.  Metrics, defined so that each one moves with leakage alone:
      * leakage_ratio - background-corrected mean extravascular FITC over
        the background-corrected luminal FITC concentration. The luminal
        reference is a high percentile of FITC inside the vessel core, so
        it is estimated from vessel pixels only and cannot drift with the
        extravascular distribution.
      * leakage_area_fraction - fraction of extravascular tissue carrying
        tracer at or above LEAK_FRACTION_OF_LUMEN of luminal concentration.
        The criterion is anchored to the lumen rather than to a global
        Otsu split, which would jump discontinuously when the histogram
        changes mode.
    The FITC background floor is estimated from tracer-free pixels: camera
    background outside the specimen when available, otherwise the darkest
    tissue percentile, with a warning that whole-field diffuse leak cannot
    be separated from the offset in that case.

All steps are deterministic: the same image always yields the same numbers.
Every tunable scales with image resolution (baseline 1024 px, supported up
to 2048 px; larger inputs are flagged in the result warnings).
"""

from dataclasses import dataclass, field

import numpy as np
from scipy import ndimage as ndi
from skimage import filters, morphology


# Extravascular tracer at or above this fraction of the luminal
# concentration counts as leak-positive area.
LEAK_FRACTION_OF_LUMEN = 0.20

# Percentile of FITC within the vessel core taken as the luminal tracer
# concentration. High enough to represent filled lumen, robust to the
# unperfused vessel pixels and wall pixels also present in the core mask.
LUMEN_PERCENTILE = 90.0

MAX_SUPPORTED_SIDE = 2048


# ---------------------------------------------------------------------------
# Channel handling
# ---------------------------------------------------------------------------

def split_channels(img):
    """Return (fitc, ib4, annotation) as float32 / bool 2-D arrays.

    Accepts:
      * (H, W, 3) or (H, W, 4) RGB(A) merges -> green = FITC, red = IB4
      * (2|3, H, W) plane stacks             -> plane 0 = FITC, plane 1 = IB4
      * (H, W, 2) channel-last stacks        -> same convention

    Both channels are divided by one common scale factor, so absolute
    intensities remain comparable between the two channels and between
    images of the same bit depth.

    ``annotation`` marks burned-in white overlays (scale bars, text):
    compact clusters of pixels near saturation in every channel at once,
    which does not occur for genuine two-fluorophore signal.
    """
    a = np.asarray(img)
    if a.ndim == 2:
        raise ValueError(
            "Single-channel image: need both a FITC and an IB4 channel.")
    if a.ndim != 3:
        raise ValueError(f"Unsupported image shape {a.shape}.")

    divisor = _divisor(a)

    if a.shape[-1] in (3, 4) and a.shape[0] not in (2, 3):
        fitc, ib4, third = a[..., 1], a[..., 0], a[..., 2]
    elif a.shape[0] in (2, 3):
        fitc, ib4 = a[0], a[1]
        third = a[2] if a.shape[0] == 3 else None
    elif a.shape[-1] == 2:
        fitc, ib4 = a[..., 0], a[..., 1]
        third = None
    else:
        raise ValueError(f"Cannot infer channel layout for shape {a.shape}.")

    fitc = _scale_to_unit(fitc, divisor)
    ib4 = _scale_to_unit(ib4, divisor)

    # A grayscale image replicated across channels carries no FITC/IB4
    # separation; every downstream comparison would be degenerate.
    if np.allclose(fitc, ib4, atol=1e-4):
        raise ValueError(
            "The FITC and IB4 channels are identical: this looks like a "
            "grayscale image rather than a dual-label merge.")

    annotation = (fitc > 0.92) & (ib4 > 0.92)
    if third is not None:
        annotation &= _scale_to_unit(third, divisor) > 0.92
    annotation = _compact_only(annotation, fitc.shape)

    return fitc, ib4, annotation


def _divisor(a):
    """One scale factor for the whole image, from dtype and value range."""
    mx = float(a.max()) if a.size else 0.0
    if np.issubdtype(a.dtype, np.floating):
        if mx <= 1.0:
            return 1.0
        return 255.0 if mx <= 255 else (4095.0 if mx <= 4095 else 65535.0)
    if a.dtype == np.uint8 or a.dtype == np.int8:
        return 255.0
    if mx <= 255:
        return 255.0
    return 4095.0 if mx <= 4095 else 65535.0


def _scale_to_unit(ch, divisor):
    ch = np.asarray(ch, dtype=np.float32) / np.float32(divisor)
    return np.clip(ch, 0.0, 1.0)


def _compact_only(mask, shape):
    """Keep only locally dense parts of a mask, then pad them slightly."""
    if not mask.any():
        return mask
    s = _scale(shape)
    dens = ndi.uniform_filter(mask.astype(np.float32),
                              size=max(3, int(round(5 * s))))
    mask = mask & (dens > 0.6)
    if not mask.any():
        return mask
    return morphology.binary_dilation(
        mask, morphology.disk(max(2, int(round(2 * s)))))


def _scale(img_shape):
    """Resolution scale factor relative to the 1024 px design baseline."""
    return float(np.clip(max(img_shape) / 1024.0, 0.5, 2.0))


# ---------------------------------------------------------------------------
# Masks
# ---------------------------------------------------------------------------

def tissue_mask(fitc, ib4, annotation=None, min_signal=0.02):
    """Footprint of the imaged specimen.

    Signal-bearing regions are unified with a closing that bridges narrow
    dark inlets, then filled, so interstitium enclosed by tissue counts as
    tissue. Wide camera-background regions reachable from the frame border
    are then removed, which excludes the black wedges of rotated
    acquisitions and empty space beside the specimen.
    """
    s = _scale(fitc.shape)
    combined = np.maximum(fitc, ib4)
    signal = filters.gaussian(combined, sigma=4 * s) > min_signal
    signal = morphology.remove_small_objects(signal, min_size=int(1024 * s * s))

    if not signal.any():                  # uniformly dim image: use full frame
        mask = np.ones_like(signal, dtype=bool)
    else:
        closed = morphology.binary_closing(
            signal, morphology.disk(max(4, int(round(10 * s)))))
        mask = ndi.binary_fill_holes(closed | signal)
        mask &= ~_wide_border_background(signal, s)

    if annotation is not None:
        mask = mask & ~annotation
    return mask


def _wide_border_background(signal, s):
    """Signal-free regions reachable from the frame border and wide enough
    to be genuine empty background rather than a dark inlet in tissue."""
    bg = ~signal
    labels, n = ndi.label(bg)
    if n == 0:
        return np.zeros_like(signal)

    border_ids = set()
    for edge in (labels[0, :], labels[-1, :], labels[:, 0], labels[:, -1]):
        border_ids.update(np.unique(edge).tolist())
    border_ids.discard(0)
    if not border_ids:
        return np.zeros_like(signal)

    border_bg = np.isin(labels, list(border_ids))
    # "Wide" means the region survives an erosion of the expected inlet
    # width; only components retaining a core are treated as background.
    core = morphology.binary_erosion(
        border_bg, morphology.disk(max(4, int(round(8 * s)))))
    if not core.any():
        return np.zeros_like(signal)
    keep_ids = np.unique(labels[core])
    keep_ids = keep_ids[keep_ids != 0]
    wide = np.isin(labels, keep_ids)
    return ndi.binary_dilation(wide, iterations=2)


def _segment_ib4(ib4, tissue, s):
    """Intensity-based vessel detection on the IB4 (red) channel."""
    smooth = filters.gaussian(ib4, sigma=1.5 * s)
    background = filters.gaussian(ib4, sigma=50 * s)
    flat = np.clip(smooth - background, 0, None)

    inside = flat[tissue]
    if inside.size == 0 or float(inside.max()) <= 0:
        return np.zeros_like(tissue, dtype=bool)

    otsu = filters.threshold_otsu(inside)
    mask = filters.apply_hysteresis_threshold(flat, 0.5 * otsu, otsu)
    return mask & tissue


def _segment_fitc_tubes(fitc, tissue, s):
    """Geometry-based detection of FITC-perfused capillaries.

    Perfused vessels are bright tubular structures in the FITC channel;
    extravasated dye is diffuse. Multi-scale Sato vesselness responds to
    tubes across the capillary calibre range and is additionally gated by
    intensity so faint texture in leak clouds is not picked up.
    """
    smooth = filters.gaussian(fitc, sigma=1.0 * s)
    background = filters.gaussian(fitc, sigma=50 * s)
    flat = np.clip(smooth - background, 0, None)

    inside = flat[tissue]
    if inside.size == 0 or float(inside.max()) <= 0:
        return np.zeros_like(tissue, dtype=bool)

    sigmas = [c * s for c in (2, 4, 6, 9, 13)]
    vess = filters.sato(flat, sigmas=sigmas, black_ridges=False, mode="reflect")

    v_in = vess[tissue]
    if float(v_in.max()) <= 0:
        return np.zeros_like(tissue, dtype=bool)
    v_thr = filters.threshold_otsu(v_in)
    i_thr = 0.5 * filters.threshold_otsu(inside)

    return (vess > 0.5 * v_thr) & (flat > i_thr) & tissue


def segment_vessels(ib4, fitc, tissue, wall_dilation_px=None):
    """Segment the vascular compartment using both markers.

    Returns (vessel_core, vessel_dilated, guard).
    """
    s = _scale(ib4.shape)
    if wall_dilation_px is None:
        wall_dilation_px = max(2, int(round(3 * s)))
    guard_px = max(3, int(round(4 * s)))

    mask = _segment_ib4(ib4, tissue, s) | _segment_fitc_tubes(fitc, tissue, s)
    mask = morphology.binary_closing(mask, morphology.disk(max(2, int(round(2 * s)))))
    mask = morphology.remove_small_holes(mask, area_threshold=int(400 * s * s))
    mask = _fill_bright_holes(mask, fitc, tissue)
    mask = morphology.remove_small_objects(mask, min_size=int(96 * s * s))

    dilated = morphology.binary_dilation(mask, morphology.disk(wall_dilation_px))
    guard = morphology.binary_dilation(dilated, morphology.disk(guard_px)) & ~dilated
    return mask & tissue, dilated & tissue, guard & tissue


def _fill_bright_holes(mask, fitc, tissue):
    """Fill enclosed holes in the vessel mask whose FITC content is bright.

    Bright enclosed holes are unstained blood cells or uneven lumen texture
    surrounded by tracer; assigning them to the extravascular compartment
    would dilute it with intraluminal pixels. Dark enclosed holes are
    avascular tissue inside capillary loops and are left extravascular.
    """
    holes = ndi.binary_fill_holes(mask) & ~mask
    if not holes.any():
        return mask
    t_fitc = fitc[tissue]
    if t_fitc.size == 0 or float(t_fitc.max()) <= 0:
        return mask
    thr = filters.threshold_otsu(t_fitc)
    labels, n = ndi.label(holes)
    if n == 0:
        return mask
    means = np.asarray(ndi.mean(fitc, labels=labels, index=np.arange(1, n + 1)))
    bright = np.flatnonzero(means >= thr) + 1
    return mask | np.isin(labels, bright) if bright.size else mask


# ---------------------------------------------------------------------------
# Background floor
# ---------------------------------------------------------------------------

def estimate_background(fitc, tissue):
    """Estimate the FITC background floor (detector offset) from tracer-free
    pixels only.

    The floor is taken from camera background outside the specimen: those
    pixels cannot contain tracer, so the estimate is independent of how much
    dye has leaked. When the specimen fills the frame no tracer-free region
    exists and the floor is left at zero, with a note in the warnings.

    Estimating the floor from a low percentile of tissue instead - the
    obvious alternative - is not safe: diffuse extravasation raises the
    darkest tissue in lockstep with the leak, so the correction would
    subtract exactly the signal being measured and severe whole-field
    leakage would read as none.

    Returns (background, warnings).
    """
    outside = ~tissue
    if int(outside.sum()) >= max(256, int(0.005 * fitc.size)):
        return float(np.median(fitc[outside])), []
    return 0.0, [
        "The specimen fills the frame, so no tracer-free region is available "
        "to estimate the FITC background floor; intensities are not "
        "background-corrected and may include a detector or out-of-focus "
        "pedestal."]


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

@dataclass
class LeakageResult:
    leakage_ratio: float                 # corrected extravascular mean / luminal concentration
    leakage_area_fraction: float         # extravascular area at >= 20% of luminal concentration
    vessel_area_fraction: float          # vessel core px / tissue px
    perfused_vessel_area_fraction: float
    lumen_fitc_concentration: float      # corrected luminal reference (0-1)
    mean_fitc_extravascular: float       # raw mean FITC in extravascular tissue
    fitc_background: float               # estimated FITC background floor (0-1)
    leak_threshold: float                # raw FITC level counted as leak-positive
    tissue_px: int = 0
    vessel_px: int = 0
    perfused_vessel_px: int = 0
    extravascular_px: int = 0
    warnings: list = field(default_factory=list)


def quantify(fitc, ib4, tissue, vessel_core, vessel_dilated, guard=None,
             extra_warnings=None):
    """Compute leakage metrics from the channel images and masks."""
    warnings = list(extra_warnings or [])
    if guard is None:
        guard = np.zeros_like(tissue, dtype=bool)
    extravascular = tissue & ~vessel_dilated & ~guard
    intravascular = vessel_dilated & tissue

    n_in = int(intravascular.sum())
    n_out = int(extravascular.sum())
    n_tissue = int(tissue.sum())
    n_core = int(vessel_core.sum())

    background, bg_warn = estimate_background(fitc, tissue)
    warnings.extend(bg_warn)

    if n_in == 0 or n_core == 0:
        warnings.append("No vessels detected: leakage metrics are undefined.")
        return LeakageResult(float("nan"), float("nan"), 0.0, 0.0,
                             float("nan"), float("nan"), background,
                             float("nan"), n_tissue, 0, 0, n_out, warnings)

    # Luminal tracer concentration: a high percentile of FITC inside the
    # vessel core. Estimated from vessel pixels only, so it cannot drift
    # with the extravascular distribution being measured against it.
    lumen_raw = float(np.percentile(fitc[vessel_core], LUMEN_PERCENTILE))
    lumen = max(lumen_raw - background, 0.0)

    # Perfused lumen, for reporting: vessel pixels carrying tracer at or
    # above the same lumen-anchored criterion used for leak-positive area.
    leak_threshold = background + LEAK_FRACTION_OF_LUMEN * lumen
    perfused = intravascular & (fitc >= leak_threshold)
    n_perf = int(perfused.sum())

    if lumen <= 0.02:
        warnings.append(
            "Luminal FITC is barely above the background floor; the leakage "
            "metrics are unreliable for this image.")
    if n_out < max(256, n_tissue // 100):
        warnings.append(
            "Almost no extravascular tissue outside the vascular "
            "compartment; leakage metrics are not meaningful.")

    mean_out_raw = float(fitc[extravascular].mean()) if n_out else float("nan")
    if n_out:
        corr_out = float(np.clip(fitc[extravascular] - background, 0, None).mean())
        leak_px = int((fitc[extravascular] >= leak_threshold).sum())
        area_frac = leak_px / n_out
    else:
        corr_out, area_frac = float("nan"), float("nan")

    ratio = corr_out / lumen if lumen > 0 else float("nan")

    return LeakageResult(
        leakage_ratio=ratio,
        leakage_area_fraction=area_frac,
        vessel_area_fraction=n_core / max(n_tissue, 1),
        perfused_vessel_area_fraction=n_perf / max(n_tissue, 1),
        lumen_fitc_concentration=lumen,
        mean_fitc_extravascular=mean_out_raw,
        fitc_background=background,
        leak_threshold=leak_threshold,
        tissue_px=n_tissue,
        vessel_px=n_core,
        perfused_vessel_px=n_perf,
        extravascular_px=n_out,
        warnings=warnings,
    )


def analyze(img, wall_dilation_px=None):
    """Full pipeline on one image array. Returns (result, intermediates)."""
    fitc, ib4, annotation = split_channels(img)
    extra = []
    if max(fitc.shape) > MAX_SUPPORTED_SIDE:
        extra.append(
            f"Image is larger than {MAX_SUPPORTED_SIDE} px: internal length "
            "scales are capped, so vessels may be under-detected. Downscale "
            "the image before analysis.")
    tissue = tissue_mask(fitc, ib4, annotation)
    core, dilated, guard = segment_vessels(ib4, fitc, tissue, wall_dilation_px)
    result = quantify(fitc, ib4, tissue, core, dilated, guard, extra)
    inter = {
        "fitc": fitc, "ib4": ib4, "tissue": tissue,
        "vessel_core": core, "vessel_dilated": dilated, "guard": guard,
        "annotation": annotation,
    }
    return result, inter