processing.py

"""
Core image-processing pipeline for cochlear neurofilament tracing.

Handles:
  * Loading Zeiss CZI z-stacks and generic TIFF stacks (with voxel sizes).
  * Channel identification (Neurofilament vs Myo7a).
  * 3D tracing of the neurofilament network into a single continuous skeleton.
  * Myo7a-guided splitting of the field into an IHC region and an OHC region.
  * Per-region quantification: number of fibers, diameter, length,
    branch points and area covered within the field of view.

The module has no Gradio dependency so it can be unit-tested on its own.
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass, field
from typing import Optional

import numpy as np
from scipy import ndimage as ndi
from skimage.filters import gaussian, threshold_otsu
from skimage.morphology import remove_small_objects, skeletonize
from skan import Skeleton, summarize

# --------------------------------------------------------------------------- #
#  Data containers
# --------------------------------------------------------------------------- #

FREQ_CHOICES = ["8kHz", "16kHz", "22kHz", "32kHz", "64kHz", "Other / unknown"]


@dataclass
class LoadedImage:
    """A loaded multi-channel z-stack."""

    data: np.ndarray               # (C, Z, Y, X) float32
    channels: list                 # list of dicts: {name, dye, color}
    voxel: tuple                   # (dz, dy, dx) in microns
    source_name: str = ""

    @property
    def n_channels(self) -> int:
        return self.data.shape[0]


@dataclass
class RegionMetrics:
    """Quantification for one region (whole field, IHC region or OHC region)."""

    region: str = ""
    n_fibers: int = 0
    total_length_um: float = 0.0
    mean_diameter_um: float = 0.0
    median_diameter_um: float = 0.0
    n_branch_points: int = 0
    area_covered_um2: float = 0.0
    fov_area_um2: float = 0.0
    pct_area_covered: float = 0.0

    def as_row(self) -> dict:
        return {
            "Region": self.region,
            "Number of fibers": self.n_fibers,
            "Total length (um)": round(self.total_length_um, 2),
            "Mean diameter (um)": round(self.mean_diameter_um, 3),
            "Median diameter (um)": round(self.median_diameter_um, 3),
            "Branch points": self.n_branch_points,
            "Area covered (um^2)": round(self.area_covered_um2, 2),
            "FOV area (um^2)": round(self.fov_area_um2, 2),
            "Area covered (% of FOV)": round(self.pct_area_covered, 2),
        }


@dataclass
class TraceResult:
    """Everything produced by tracing one image."""

    mask: np.ndarray               # 3D bool
    skeleton: np.ndarray           # 3D bool
    distance_um: np.ndarray        # 3D float, EDT in microns
    voxel: tuple
    metrics: dict = field(default_factory=dict)   # region name -> RegionMetrics


# --------------------------------------------------------------------------- #
#  Loading
# --------------------------------------------------------------------------- #

def detect_frequency(filename: str) -> str:
    """Guess the frequency-region label from a filename (e.g. '16kHz')."""
    m = re.search(r"(\d+)\s*k\s*hz", filename, re.IGNORECASE)
    if m:
        label = f"{int(m.group(1))}kHz"
        if label in FREQ_CHOICES:
            return label
    return "Other / unknown"


def _czi_channel_meta(czi) -> list:
    """Extract per-channel name/dye/color from CZI metadata (best effort)."""
    meta = czi.meta
    seen = {}
    order = []
    for ch in meta.iter("Channel"):
        name = ch.get("Name")
        if not name:
            continue
        dye = ch.findtext("DyeName") or ch.findtext("Fluor")
        color = ch.findtext("Color")
        ex = ch.findtext("ExcitationWavelength")
        if name not in seen:
            seen[name] = {"name": name, "dye": None, "color": None, "ex": None}
            order.append(name)
        rec = seen[name]
        rec["dye"] = rec["dye"] or dye
        rec["color"] = rec["color"] or color
        rec["ex"] = rec["ex"] or ex
    return [seen[n] for n in order]


def _czi_voxel(czi) -> tuple:
    """Return (dz, dy, dx) in microns from CZI scaling metadata."""
    scale = {}
    for d in czi.meta.iter("Distance"):
        idv = d.get("Id")
        val = d.findtext("Value")
        if idv in ("X", "Y", "Z") and val:
            scale[idv] = float(val) * 1e6  # metres -> microns
    dx = scale.get("X", 0.0895)
    dy = scale.get("Y", dx)
    dz = scale.get("Z", 0.35)
    return (dz, dy, dx)


def load_czi(path: str) -> LoadedImage:
    from aicspylibczi import CziFile

    czi = CziFile(path)
    img, shp = czi.read_image()
    dims = [d for d, _ in shp]
    arr = np.asarray(img)
    # collapse everything except C, Z, Y, X
    # find axis indices
    idx = {d: i for i, d in enumerate(dims)}
    # move to C,Z,Y,X ordering, squeezing singletons (B,V,T,...)
    keep = ["C", "Z", "Y", "X"]
    order = [idx[k] for k in keep if k in idx]
    other = [i for i in range(arr.ndim) if i not in order]
    arr = np.transpose(arr, other + order)
    arr = arr.reshape((-1,) + arr.shape[len(other):]) if other else arr
    # after reshape leading axis is product of others -> take first
    if other:
        arr = arr[0]
    # arr now (C,Z,Y,X) or missing Z
    if "Z" not in idx:
        arr = arr[:, None]
    arr = arr.astype(np.float32)
    channels = _czi_channel_meta(czi)
    if len(channels) != arr.shape[0]:
        channels = [{"name": f"Channel {i}", "dye": None, "color": None}
                    for i in range(arr.shape[0])]
    return LoadedImage(arr, channels, _czi_voxel(czi), os.path.basename(path))


def load_tiff(path: str, dz=0.35, dy=0.0895, dx=0.0895) -> LoadedImage:
    import tifffile

    arr = tifffile.imread(path)
    arr = np.squeeze(arr)
    # Heuristic to reach (C, Z, Y, X). The two largest axes are Y, X.
    if arr.ndim == 2:                       # (Y, X) single channel, single plane
        arr = arr[None, None]
    elif arr.ndim == 3:
        # Could be (Z,Y,X) single channel or (C,Y,X). Assume small first axis = C
        if arr.shape[0] <= 5:
            arr = arr[:, None]              # (C, 1, Y, X)
        else:
            arr = arr[None]                 # (1, Z, Y, X)
    elif arr.ndim == 4:
        # find the two largest axes -> Y, X; of the remaining two the smaller = C
        yx = sorted(range(4), key=lambda a: arr.shape[a])[-2:]
        rest = [a for a in range(4) if a not in yx]
        c_axis = min(rest, key=lambda a: arr.shape[a])
        z_axis = [a for a in rest if a != c_axis][0]
        arr = np.transpose(arr, (c_axis, z_axis, *sorted(yx)))
    else:
        raise ValueError(f"Unsupported TIFF with {arr.ndim} dimensions")
    arr = arr.astype(np.float32)
    channels = [{"name": f"Channel {i}", "dye": None, "color": None}
                for i in range(arr.shape[0])]
    return LoadedImage(arr, channels, (dz, dy, dx), os.path.basename(path))


def load_image(path: str, dz=0.35, dy=0.0895, dx=0.0895) -> LoadedImage:
    ext = os.path.splitext(path)[1].lower()
    if ext == ".czi":
        return load_czi(path)
    if ext in (".tif", ".tiff"):
        return load_tiff(path, dz, dy, dx)
    raise ValueError(f"Unsupported file type: {ext}")


# --------------------------------------------------------------------------- #
#  Channel identification
# --------------------------------------------------------------------------- #

# Wavelength-based hints: neurofilament here is Alexa-555 (red/green range),
# Myo7a is Alexa-405 (blue). Transmitted-light PMT channels have no dye.
def guess_channels(img: LoadedImage) -> tuple:
    """Best-effort (neurofilament_index, myo7a_index) from metadata + content."""
    nf_idx, myo_idx = None, None
    blue_like, red_like, plain = [], [], []
    for i, ch in enumerate(img.channels):
        dye = (ch.get("dye") or "").lower()
        color = (ch.get("color") or "").upper()
        ex = ch.get("ex")
        ex = float(ex) if ex else None
        if "405" in dye or (ex and ex < 430) or color == "#0000FF":
            blue_like.append(i)
        elif dye and dye not in ("", "none"):
            red_like.append(i)
        else:
            plain.append(i)     # e.g. transmitted-light PMT
    if blue_like:
        myo_idx = blue_like[0]
    if red_like:
        nf_idx = red_like[0]

    # Fall back to image content when metadata is missing (e.g. plain TIFF).
    # Fluorescence channels have a mostly-dark background; transmitted-light
    # (brightfield) channels fill the frame, so we skip those.  We cannot
    # reliably tell fibers from blobs automatically, so we default NF to the
    # first fluorescent channel and Myo7a to the last — the user confirms
    # via the channel previews in the UI.
    if nf_idx is None or myo_idx is None:
        fluo = [i for i in range(img.n_channels)
                if _dark_fraction(img.data[i]) >= 0.2]
        if not fluo:
            fluo = list(range(img.n_channels))
        if nf_idx is None:
            nf_idx = fluo[0]
        if myo_idx is None or myo_idx == nf_idx:
            myo_idx = fluo[-1] if fluo[-1] != nf_idx else fluo[0]
    return nf_idx, myo_idx


def _dark_fraction(vol: np.ndarray) -> float:
    """Fraction of the (normalised) MIP that is near-black background."""
    mip = vol.max(0).astype(np.float32)
    mip = (mip - mip.min()) / (np.ptp(mip) + 1e-6)
    return float((mip < 0.15).mean())


# --------------------------------------------------------------------------- #
#  Neurofilament tracing
# --------------------------------------------------------------------------- #

def _threshold_volume(vol: np.ndarray, sensitivity: float) -> np.ndarray:
    """Smooth + Otsu threshold. `sensitivity` (0.5-1.5) scales the threshold;
    higher sensitivity -> lower threshold -> more fibers captured."""
    lo, hi = np.percentile(vol, 1), np.percentile(vol, 99.8)
    norm = np.clip((vol - lo) / (hi - lo + 1e-6), 0, 1)
    sm = gaussian(norm, sigma=(0.6, 1.0, 1.0), preserve_range=True)
    fg = sm[sm > sm.mean() * 0.3]
    base = threshold_otsu(fg) if fg.size else sm.mean()
    thr = base * (2.0 - sensitivity)        # sensitivity 1.0 -> base
    return sm > thr


def trace_neurites(nf_vol: np.ndarray, voxel: tuple,
                   sensitivity: float = 1.0,
                   min_object_vox: int = 64) -> TraceResult:
    """Segment and skeletonise the neurofilament network in 3D."""
    dz, dy, dx = voxel
    mask = _threshold_volume(nf_vol, sensitivity)
    mask = remove_small_objects(mask, min_object_vox)
    mask = ndi.binary_closing(mask, iterations=1)
    if mask.sum() == 0:
        z = np.zeros_like(mask)
        return TraceResult(mask, z, np.zeros(mask.shape, np.float32), voxel)
    skel = skeletonize(mask)
    dist = ndi.distance_transform_edt(mask, sampling=(dz, dy, dx)).astype(np.float32)
    return TraceResult(mask, skel, dist, voxel)


# --------------------------------------------------------------------------- #
#  Region definition (IHC vs OHC) from Myo7a
# --------------------------------------------------------------------------- #

def myo_band_profile(myo_vol: np.ndarray) -> np.ndarray:
    """Row-wise (Y) intensity profile of the Myo7a hair-cell band."""
    mip = myo_vol.max(0).astype(np.float32)
    sm = gaussian(mip, 3, preserve_range=True)
    return sm.sum(1)


def suggest_boundary(myo_vol: np.ndarray) -> float:
    """Suggest an IHC/OHC boundary as a fraction (0-1) along the Y axis.

    Places the boundary at the centre of the Myo7a hair-cell band, which sits
    between the (single) IHC row and the (three) OHC rows in a well-oriented
    organ-of-Corti image.  Users can refine this manually.
    """
    prof = myo_band_profile(myo_vol)
    if prof.sum() == 0:
        return 0.5
    ys = np.arange(prof.size)
    centroid = float((ys * prof).sum() / prof.sum())
    return centroid / prof.size


def make_region_masks(shape_yx: tuple, boundary_frac: float,
                      ihc_side: str = "low", axis: str = "Y") -> tuple:
    """Return (ihc_roi, ohc_roi) boolean 2D masks split by a straight line.

    axis="Y" splits along rows (radial axis, the usual case); axis="X" splits
    along columns.  ihc_side selects which side of the boundary is IHC.
    """
    ny, nx = shape_yx
    low = np.zeros((ny, nx), bool)
    if axis.upper() == "Y":
        b = int(round(np.clip(boundary_frac, 0, 1) * ny))
        low[:b] = True
    else:
        b = int(round(np.clip(boundary_frac, 0, 1) * nx))
        low[:, :b] = True
    ihc = low if ihc_side == "low" else ~low
    return ihc, ~ihc


# --------------------------------------------------------------------------- #
#  Quantification
# --------------------------------------------------------------------------- #

def _branch_point_count(skel: np.ndarray) -> int:
    if skel.sum() == 0:
        return 0
    k = np.ones((3, 3, 3), int) if skel.ndim == 3 else np.ones((3, 3), int)
    nb = ndi.convolve(skel.astype(np.uint8), k, mode="constant") - skel
    return int((skel & (nb > 2)).sum())


def compute_metrics(trace: TraceResult, region_name: str,
                    roi_yx: Optional[np.ndarray] = None,
                    min_fiber_um: float = 5.0) -> RegionMetrics:
    """Quantify the skeleton, optionally restricted to a 2D ROI (broadcast in Z)."""
    dz, dy, dx = trace.voxel
    skel = trace.skeleton
    mask = trace.mask
    if roi_yx is not None:
        roi3d = np.broadcast_to(roi_yx, skel.shape)
        skel = skel & roi3d
        mask = mask & roi3d

    m = RegionMetrics(region=region_name)
    m.fov_area_um2 = float(roi_yx.sum() if roi_yx is not None
                           else mask.shape[1] * mask.shape[2]) * dx * dy

    foot = mask.max(0)
    m.area_covered_um2 = float(foot.sum()) * dx * dy
    m.pct_area_covered = (100.0 * m.area_covered_um2 / m.fov_area_um2
                          if m.fov_area_um2 else 0.0)

    if skel.sum() < 2:
        return m

    m.n_branch_points = _branch_point_count(skel)

    diam = 2.0 * trace.distance_um[skel]
    if diam.size:
        m.mean_diameter_um = float(diam.mean())
        m.median_diameter_um = float(np.median(diam))

    # Length and fiber count via skan (per connected skeleton component).
    try:
        S = Skeleton(skel, spacing=(dz, dy, dx))
        df = summarize(S, separator="_")
        comp_len = df.groupby("skeleton_id")["branch_distance"].sum()
        kept = comp_len[comp_len >= min_fiber_um]
        m.n_fibers = int(kept.size)
        m.total_length_um = float(kept.sum())
    except Exception:
        # Fallback: label components, approximate length by voxel count.
        lbl, n = ndi.label(skel, structure=np.ones((3, 3, 3)))
        m.n_fibers = int(n)
        m.total_length_um = float(skel.sum()) * np.mean([dz, dy, dx])
    return m


# --------------------------------------------------------------------------- #
#  Rendering
# --------------------------------------------------------------------------- #

def skeleton_image(skel: np.ndarray, dilate: int = 1) -> np.ndarray:
    """White skeleton on black background (2D uint8), as a MIP over Z."""
    flat = skel.max(0) if skel.ndim == 3 else skel
    if dilate:
        flat = ndi.binary_dilation(flat, iterations=dilate)
    return (flat * 255).astype(np.uint8)


def region_overlay(skel: np.ndarray, ihc_roi: np.ndarray, ohc_roi: np.ndarray,
                   boundary_frac: float, axis: str = "Y",
                   dilate: int = 1) -> np.ndarray:
    """Colour-coded RGB preview: IHC fibers cyan, OHC fibers magenta,
    with the boundary line drawn in yellow."""
    flat = skel.max(0) if skel.ndim == 3 else skel
    if dilate:
        flat = ndi.binary_dilation(flat, iterations=dilate)
    ny, nx = flat.shape
    rgb = np.zeros((ny, nx, 3), np.uint8)
    ihc_pix = flat & ihc_roi
    ohc_pix = flat & ohc_roi
    rgb[ihc_pix] = (0, 220, 255)      # cyan  = IHC
    rgb[ohc_pix] = (255, 60, 200)     # magenta = OHC
    if axis.upper() == "Y":
        b = int(round(np.clip(boundary_frac, 0, 1) * ny))
        b = min(max(b, 0), ny - 1)
        rgb[b, :] = (255, 255, 0)
    else:
        b = int(round(np.clip(boundary_frac, 0, 1) * nx))
        b = min(max(b, 0), nx - 1)
        rgb[:, b] = (255, 255, 0)
    return rgb


def channel_preview(vol: np.ndarray) -> np.ndarray:
    """Contrast-stretched MIP of a channel for display (uint8)."""
    mip = vol.max(0).astype(np.float32)
    lo, hi = np.percentile(mip, 1), np.percentile(mip, 99.5)
    return (np.clip((mip - lo) / (hi - lo + 1e-6), 0, 1) * 255).astype(np.uint8)