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


def _cellpose_available() -> bool:
    try:
        import cellpose  # noqa: F401
        return True
    except Exception:
        return False


CELLPOSE_AVAILABLE = _cellpose_available()
_CP_MODEL = None       # lazily-created, cached Cellpose model

# --------------------------------------------------------------------------- #
#  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
    n_hair_cells: int = -1         # -1 = not measured (no Myo7a detection run)

    def as_row(self) -> dict:
        row = {
            "Region": self.region,
            "Number of fibers": self.n_fibers,
            "Hair cells (Myo7a)": (self.n_hair_cells
                                   if self.n_hair_cells >= 0 else ""),
            "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),
        }
        return row


@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


# --------------------------------------------------------------------------- #
#  Hair-cell detection (Cellpose assist, with a classical fallback)
# --------------------------------------------------------------------------- #

# A mouse cochlear hair cell body is roughly this wide; used to pick the
# working scale so cells land near Cellpose's preferred pixel size.
HAIR_CELL_DIAMETER_UM = 8.0
_CP_TARGET_PX = 30


def _norm(mip: np.ndarray) -> np.ndarray:
    lo, hi = np.percentile(mip, 1), np.percentile(mip, 99.5)
    return np.clip((mip - lo) / (hi - lo + 1e-6), 0, 1).astype(np.float32)


def _get_cellpose_model():
    global _CP_MODEL
    if _CP_MODEL is None:
        from cellpose import models
        _CP_MODEL = models.CellposeModel(gpu=False)
    return _CP_MODEL


def _watershed_cells(img: np.ndarray, cell_px: float) -> np.ndarray:
    """Classical blob segmentation fallback (no deep-learning dependency)."""
    from skimage.feature import peak_local_max
    from skimage.segmentation import watershed
    sm = gaussian(img, max(cell_px / 6.0, 1.0), preserve_range=True)
    try:
        mask = sm > threshold_otsu(sm)
    except Exception:
        return np.zeros(img.shape, int)
    mask = ndi.binary_opening(mask, iterations=1)
    dist = ndi.distance_transform_edt(mask)
    coords = peak_local_max(dist, min_distance=max(int(cell_px * 0.5), 3),
                            labels=mask)
    if len(coords) == 0:
        return np.zeros(img.shape, int)
    markers = np.zeros(img.shape, int)
    markers[tuple(coords.T)] = np.arange(1, len(coords) + 1)
    return watershed(-dist, markers, mask=mask)


def detect_hair_cells(myo_vol: np.ndarray, voxel: tuple,
                      prefer_cellpose: bool = True) -> dict:
    """Detect Myo7a hair-cell bodies on the max-projection.

    Uses Cellpose when available (better on touching cells), otherwise a
    watershed fallback.  Returns full-resolution centroids plus the count and
    which engine ran.  The image is rescaled so cells are ~30 px, which is
    where Cellpose performs best.
    """
    dz, dy, dx = voxel
    mip = _norm(myo_vol.max(0).astype(np.float32))
    cell_px_full = HAIR_CELL_DIAMETER_UM / dx           # e.g. ~89 px
    scale = float(np.clip(_CP_TARGET_PX / cell_px_full, 0.2, 1.0))

    from skimage.transform import rescale
    small = rescale(mip, scale, anti_aliasing=True, preserve_range=True
                    ).astype(np.float32) if scale < 0.999 else mip
    small = _norm(small)

    engine = "watershed"
    masks = None
    if prefer_cellpose and CELLPOSE_AVAILABLE:
        try:
            masks = _get_cellpose_model().eval(small, diameter=_CP_TARGET_PX)[0]
            engine = "cellpose"
        except Exception:
            masks = None
    if masks is None:
        masks = _watershed_cells(small, _CP_TARGET_PX)

    n = int(masks.max())
    if n:
        cent_small = np.array(ndi.center_of_mass(
            np.ones_like(masks), masks, range(1, n + 1)))
        centroids = cent_small / scale                  # back to full-res Y,X
    else:
        centroids = np.zeros((0, 2))
    return {"centroids": centroids, "count": n, "engine": engine,
            "scale": scale, "mip_shape": mip.shape}


def auto_regions_from_cells(centroids: np.ndarray, shape_yx: tuple) -> dict:
    """Suggest an IHC/OHC boundary from hair-cell centroids.

    IHCs form a single row and OHCs form three rows separated from the IHCs by
    the tunnel of Corti.  We project cells onto the radial axis (perpendicular
    to the hair-cell band), find the widest cell-free gap that leaves at least
    two cells on each side, and call the sparser/tighter side IHC.  A
    confidence label reflects how clearly the gap and the ~1:3 cell ratio
    appear — real fields are often ambiguous, so this is a suggestion.
    """
    ny, nx = shape_yx
    result = {"boundary_frac": 0.5, "ihc_side": "low", "confidence": "low",
              "n_ihc_cells": 0, "n_ohc_cells": 0,
              "reason": "not enough hair cells for an automatic split"}
    if len(centroids) < 6:
        # fall back to band centroid
        result["boundary_frac"] = float(np.clip(
            centroids[:, 0].mean() / ny, 0, 1)) if len(centroids) else 0.5
        return result

    c = centroids - centroids.mean(0)
    _, _, vt = np.linalg.svd(c, full_matrices=False)
    radial = vt[1]
    if radial[0] < 0:
        radial = -radial                                # point toward +Y
    r = c @ radial
    order = np.argsort(r)
    rs = r[order]
    gaps = np.diff(rs)
    # only accept splits leaving >=2 cells on each side (ignore stray outliers)
    valid = [(i, gaps[i]) for i in range(1, len(gaps) - 1)]
    if not valid:
        result["boundary_frac"] = float(np.clip(
            centroids[:, 0].mean() / ny, 0, 1))
        return result
    gi = max(valid, key=lambda t: t[1])[0]
    biggest_gap = gaps[gi]
    split_r = (rs[gi] + rs[gi + 1]) / 2.0

    left = r <= split_r
    n_left, n_right = int(left.sum()), int((~left).sum())
    s_left = float(r[left].std()) if n_left > 1 else 0.0
    s_right = float(r[~left].std()) if n_right > 1 else 0.0
    # IHC = single row: fewer cells and tighter spread
    score = (1 if n_right > n_left else -1) + (1 if s_right > s_left else -1)
    ihc_is_left = score >= 0
    ihc_side = "low" if ihc_is_left else "high"
    n_ihc = n_left if ihc_is_left else n_right
    n_ohc = n_right if ihc_is_left else n_left

    ymid = centroids[:, 0].mean() + split_r * radial[0]
    bfrac = float(np.clip(ymid / ny, 0, 1))

    med_gap = float(np.median(gaps[gaps > 0])) if (gaps > 0).any() else 1.0
    gap_ratio = biggest_gap / (med_gap + 1e-6)
    ratio = n_ohc / max(n_ihc, 1)
    if gap_ratio >= 3.0 and 1.8 <= ratio <= 5.0:
        conf = "high"
    elif gap_ratio >= 2.0:
        conf = "medium"
    else:
        conf = "low"

    return {"boundary_frac": bfrac, "ihc_side": ihc_side, "confidence": conf,
            "n_ihc_cells": n_ihc, "n_ohc_cells": n_ohc,
            "reason": f"tunnel gap {gap_ratio:.1f}x median spacing, "
                      f"IHC:OHC cell ratio 1:{ratio:.1f}"}


def count_cells_in_roi(centroids: np.ndarray, roi_yx: np.ndarray) -> int:
    """Count hair-cell centroids falling inside a 2D ROI mask."""
    if len(centroids) == 0:
        return 0
    ny, nx = roi_yx.shape
    yy = np.clip(centroids[:, 0].round().astype(int), 0, ny - 1)
    xx = np.clip(centroids[:, 1].round().astype(int), 0, nx - 1)
    return int(roi_yx[yy, xx].sum())


def hair_cell_overlay(myo_mip_u8: np.ndarray, centroids: np.ndarray,
                      boundary_frac: Optional[float] = None,
                      axis: str = "Y", ihc_side: Optional[str] = None) -> np.ndarray:
    """Myo7a MIP with detected hair cells marked and the boundary drawn."""
    rgb = np.stack([myo_mip_u8] * 3, axis=-1).copy()
    ny, nx = myo_mip_u8.shape
    b = None
    if boundary_frac is not None:
        if axis.upper() == "Y":
            b = min(max(int(round(boundary_frac * ny)), 0), ny - 1)
        else:
            b = min(max(int(round(boundary_frac * nx)), 0), nx - 1)
    for (y, x) in centroids.astype(int):
        # colour by region if we know the split, else neutral green
        col = (0, 255, 0)
        if b is not None and ihc_side is not None:
            on_low = (y <= b) if axis.upper() == "Y" else (x <= b)
            is_ihc = (on_low and ihc_side == "low") or \
                     (not on_low and ihc_side == "high")
            col = (0, 220, 255) if is_ihc else (255, 60, 200)
        ys, xs = slice(max(0, y - 3), y + 4), slice(max(0, x - 3), x + 4)
        rgb[ys, xs] = col
    if b is not None:
        if axis.upper() == "Y":
            rgb[b, :] = (255, 255, 0)
        else:
            rgb[:, b] = (255, 255, 0)
    return rgb


# --------------------------------------------------------------------------- #
#  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,
                    hair_cell_centroids: Optional[np.ndarray] = None
                    ) -> RegionMetrics:
    """Quantify the skeleton, optionally restricted to a 2D ROI (broadcast in Z).

    If ``hair_cell_centroids`` is given, the number of Myo7a hair cells within
    the region is also reported.
    """
    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

    if hair_cell_centroids is not None:
        if roi_yx is not None:
            m.n_hair_cells = count_cells_in_roi(hair_cell_centroids, roi_yx)
        else:
            m.n_hair_cells = int(len(hair_cell_centroids))

    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)