"""
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)
    fibers_per_hc: float = -1.0    # -1 = not applicable (hair-cell count unknown)
    length_per_hc_um: float = -1.0
    n_radial: int = -1             # fibers running IHC->OHC (radial); -1 = n/a
    n_offaxis: int = -1            # fibers running along the rows (misdirected)
    pct_radial: float = -1.0       # % of fibers that are radially directed

    def as_row(self) -> dict:
        def opt(value, ndigits=2):
            """Blank when the metric was not applicable (negative sentinel)."""
            return round(value, ndigits) if value is not None and value >= 0 else ""

        row = {
            "Region": self.region,
            "Number of fibers": self.n_fibers,
            "Hair cells (Myo7a)": (self.n_hair_cells
                                   if self.n_hair_cells >= 0 else ""),
            "Fibers / hair cell": opt(self.fibers_per_hc, 3),
            "Total length (um)": round(self.total_length_um, 2),
            "Length / hair cell (um)": opt(self.length_per_hc_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,
            "Radial fibers (IHC->OHC)": (self.n_radial
                                         if self.n_radial >= 0 else ""),
            "Off-axis fibers": (self.n_offaxis if self.n_offaxis >= 0 else ""),
            "Radial fibers (%)": opt(self.pct_radial, 1),
            "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 _reorder_tiff_by_axes(arr: np.ndarray, axes: str):
    """Reorder a TIFF array to (C, Z, Y, X) using tifffile's axes string.

    Treats colour samples ('S', e.g. RGB) and channels ('C') as channels, keeps
    'Z', and squeezes any other axes (T, etc.).  Returns None if it cannot map.
    """
    axes = axes.upper()
    if "Y" not in axes or "X" not in axes or len(axes) != arr.ndim:
        return None
    # Rank: channel-like first, then Z, then Y, X; unknown axes go last.
    prio = {"C": 0, "S": 0, "Z": 1, "Y": 2, "X": 3}
    order = sorted(range(len(axes)), key=lambda i: (prio.get(axes[i], 4), i))
    arr = np.transpose(arr, order)
    new_axes = "".join(axes[i] for i in order)
    # Collapse all leading channel-like axes (C and/or S) into one channel axis.
    n_chan = sum(1 for c in new_axes if c in ("C", "S"))
    if n_chan >= 1:
        arr = arr.reshape((int(np.prod(arr.shape[:n_chan])),) + arr.shape[n_chan:])
        new_axes = "C" + new_axes[n_chan:]
    else:
        arr = arr[None]
        new_axes = "C" + new_axes
    if "Z" not in new_axes:                 # single-plane image -> add Z=1
        arr = arr[:, None]
        new_axes = new_axes[0] + "Z" + new_axes[1:]
    while arr.ndim > 4:                      # drop any leftover trailing axes
        arr = arr[..., 0]
    return arr


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

    with tifffile.TiffFile(path) as tf:
        arr = tf.asarray()
        try:
            axes = tf.series[0].axes
        except Exception:
            axes = None

    arr = np.asarray(arr)
    reordered = _reorder_tiff_by_axes(arr, axes) if axes else None

    if reordered is not None:
        arr = reordered
    else:
        # Fallback heuristic when axis metadata is missing/unusable.
        arr = np.squeeze(arr)
        if arr.ndim == 2:                          # (Y, X)
            arr = arr[None, None]
        elif arr.ndim == 3:
            # trailing RGB/RGBA samples -> channels; else small first axis = C
            if arr.shape[-1] in (3, 4) and arr.shape[-1] < min(arr.shape[:2]):
                arr = np.moveaxis(arr, -1, 0)[:, None]   # (C, 1, Y, X)
            elif arr.shape[0] <= 5:
                arr = arr[:, None]                       # (C, 1, Y, X)
            else:
                arr = arr[None]                          # (1, Z, Y, X)
        elif arr.ndim == 4:
            if arr.shape[-1] in (3, 4) and arr.shape[-1] < min(arr.shape[1:3]):
                arr = np.moveaxis(arr, -1, 0)            # (C, Z, Y, X)
            else:
                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 = np.ascontiguousarray(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}")


# --------------------------------------------------------------------------- #
#  IMARIS ground-truth ingestion
# --------------------------------------------------------------------------- #

def read_imaris_filament_length(source: str) -> Optional[float]:
    """Return IMARIS 'Filament Length (sum)' in microns for an image, or None.

    IMARIS exports this exact statistic to
    ``<name>_Statistics/<name>_Filament_Length_(sum).csv`` beside the image.
    ``source`` may be that CSV, a ``_Statistics`` folder, or the image path
    (``.czi``/``.tif``) whose sibling ``_Statistics`` folder is searched.  When a
    filament CSV lists several filament objects the per-object sums are added.

    This is the ONLY way to reproduce IMARIS's exact number — it reads IMARIS's
    own output rather than re-estimating it with a different algorithm.
    """
    import glob
    import re

    csvs = []
    if source and os.path.isfile(source) and source.lower().endswith(".csv"):
        csvs = [source]
    else:
        cand = None
        if source and source.lower().endswith((".czi", ".tif", ".tiff")):
            stem = os.path.splitext(os.path.basename(source))[0]
            cand = os.path.join(os.path.dirname(source), f"{stem}_Statistics")
        elif source and os.path.isdir(source):
            cand = source
        if cand and os.path.isdir(cand):
            csvs = sorted(glob.glob(
                os.path.join(cand, "*_Filament_Length_(sum).csv")))
    if not csvs:
        return None

    total, found = 0.0, False
    for csv in csvs[:1]:                       # one filament CSV per image
        try:
            with open(csv, errors="ignore") as fh:
                for line in fh:
                    m = re.match(r"\s*([-+]?[0-9]*\.?[0-9]+)\s*,\s*µm", line)
                    if m:
                        total += float(m.group(1))
                        found = True
        except OSError:
            return None
    return round(total, 3) if found else None


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

# Common fluorophores in these cochlear stacks.  Each entry maps a token that
# may appear in the CZI dye/fluor name to a friendly name and a colour family.
# Emission colour drives the family; the family drives auto channel assignment.
#   blue   -> reference marker (Myo7a is Alexa-405 here)
#   orange -> neurofilament (Alexa-555 here) — preferred trace channel
#   green / red -> additional markers, selectable in the UI.
KNOWN_DYES = [
    ("atto 390", "ATTO 390", "blue"),
    ("405", "Alexa Fluor 405", "blue"),
    ("dapi", "DAPI", "blue"),
    ("488", "Alexa Fluor 488", "green"),
    ("atto 488", "ATTO 488", "green"),
    ("fitc", "FITC", "green"),
    ("gfp", "GFP", "green"),
    ("514", "Alexa Fluor 514", "green"),
    ("532", "Alexa Fluor 532", "orange"),
    ("546", "Alexa Fluor 546", "orange"),
    ("555", "Alexa Fluor 555", "orange"),
    ("568", "Alexa Fluor 568", "orange"),
    ("cy3", "Cy3", "orange"),
    ("594", "Alexa Fluor 594", "red"),
    ("633", "Alexa Fluor 633", "red"),
    ("647", "Alexa Fluor 647", "red"),
    ("atto 647", "ATTO 647", "red"),
    ("cy5", "Cy5", "red"),
    ("680", "Alexa Fluor 680", "red"),
]


def _channel_color_family(ch: dict) -> str:
    """Classify a channel as blue / green / orange / red / none (transmitted).

    Uses the dye/fluor name first, then the excitation wavelength, then the
    metadata colour swatch.  This lets the app recognise many more dyes
    (405/488/514/532/555/568/594/633/647/ATTO/Cy…) than just 405 vs "any dye".
    """
    dye = (ch.get("dye") or "").lower()
    color = (ch.get("color") or "").upper().lstrip("#")
    ex = ch.get("ex")
    try:
        ex = float(ex) if ex else None
    except (TypeError, ValueError):
        ex = None

    for token, _name, family in KNOWN_DYES:
        if token in dye:
            return family
    if ex is not None:
        if ex < 430:
            return "blue"
        if ex < 505:
            return "green"
        if ex < 565:
            return "orange"
        return "red"
    # metadata colour swatch (Zeiss stores ARGB or RGB hex)
    if color:
        hexrgb = color[-6:] if len(color) >= 6 else color
        try:
            r = int(hexrgb[0:2], 16); g = int(hexrgb[2:4], 16); b = int(hexrgb[4:6], 16)
            if b > r and b > g:
                return "blue"
            if g > r and g > b:
                return "green"
            if r > g and r > b and g > b // 2:
                return "orange"
            if r > g and r > b:
                return "red"
        except (ValueError, IndexError):
            pass
    if dye and dye not in ("", "none"):
        return "other"
    return "none"     # e.g. transmitted-light PMT


def channel_dye_label(ch: dict) -> str:
    """Friendly dye label for the channel dropdown (recognises known dyes)."""
    dye = (ch.get("dye") or "").strip()
    if dye:
        low = dye.lower()
        for token, name, _family in KNOWN_DYES:
            if token in low:
                return name
        return dye
    fam = _channel_color_family(ch)
    return "no dye / transmitted" if fam == "none" else f"{fam} marker"


def guess_channels(img: LoadedImage) -> tuple:
    """Best-effort (neurofilament_index, myo7a_index) from metadata + content."""
    nf_idx, myo_idx = None, None
    families = [_channel_color_family(ch) for ch in img.channels]
    blue_like = [i for i, f in enumerate(families) if f == "blue"]
    orange_like = [i for i, f in enumerate(families) if f == "orange"]
    red_like = [i for i, f in enumerate(families) if f == "red"]
    green_like = [i for i, f in enumerate(families) if f == "green"]

    # Myo7a reference = the bluest (405-like) fluorescent channel.
    if blue_like:
        myo_idx = blue_like[0]
    # Neurofilament = Alexa-555 (orange) here; then fall back to red/green.
    for group in (orange_like, red_like, green_like):
        if group:
            nf_idx = next((i for i in group if i != myo_idx), group[0])
            break

    # 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 _remove_small_components(skel: np.ndarray, voxel: tuple,
                             min_len_um: float) -> np.ndarray:
    """Drop connected skeleton components shorter than ``min_len_um``."""
    if min_len_um <= 0 or skel.sum() == 0:
        return skel
    dz, dy, dx = voxel
    step = float(np.mean([dy, dx]))              # ~ length of one skeleton step
    min_vox = max(2, int(round(min_len_um / step)))
    struct = np.ones((3,) * skel.ndim, int)
    lbl, n = ndi.label(skel, structure=struct)
    if n == 0:
        return skel
    counts = np.bincount(lbl.ravel())
    keep = np.zeros(counts.size, bool)
    keep[1:] = counts[1:] >= min_vox
    return keep[lbl]


def prune_skeleton(skel: np.ndarray, voxel: tuple,
                   spur_um: float = 3.0, min_component_um: Optional[float] = None,
                   iterations: int = 4) -> np.ndarray:
    """Remove short terminal spurs and small isolated fragments.

    Terminal twigs (junction→endpoint branches) shorter than ``spur_um`` are
    trimmed while their junction pixel is preserved, so the backbone stays
    connected.  Isolated components shorter than ``min_component_um`` are dropped
    entirely.  Iterated a few times because trimming one spur can expose another.
    """
    if spur_um <= 0 and (min_component_um or 0) <= 0:
        return skel
    if min_component_um is None:
        min_component_um = spur_um
    cur = skel.copy()
    for _ in range(iterations):
        if cur.sum() == 0 or spur_um <= 0:
            break
        try:
            S = Skeleton(cur, spacing=voxel)
            df = summarize(S, separator="_")
        except Exception:
            break
        k = np.ones((3,) * cur.ndim, int)
        nb = ndi.convolve(cur.astype(np.uint8), k, mode="constant") - cur
        junction = cur & (nb > 2)               # protect branch points
        drop = df.index[
            ((df["branch_type"] == 1) & (df["branch_distance"] < spur_um)) |
            (df["branch_type"].isin([0, 3]) &
             (df["branch_distance"] < min_component_um))]
        if len(drop) == 0:
            break
        removal = np.zeros_like(cur)
        for i in drop:
            coords = S.path_coordinates(i).astype(int)
            removal[tuple(coords.T)] = True
        removal &= ~junction
        new = cur & ~removal
        if int(new.sum()) == int(cur.sum()):
            break
        cur = new
    return _remove_small_components(cur, voxel, min_component_um)


def trace_neurites(nf_vol: np.ndarray, voxel: tuple,
                   sensitivity: float = 1.0,
                   min_object_vox: int = 64,
                   prune_um: float = 3.0) -> TraceResult:
    """Segment and skeletonise the neurofilament network in 3D.

    ``prune_um`` removes terminal spurs and isolated fragments shorter than this
    many microns from the skeleton, giving cleaner fiber lines (0 = no pruning).
    """
    dz, dy, dx = voxel
    mask = _threshold_volume(nf_vol, sensitivity)
    mask = remove_small_objects(mask, min_object_vox)
    # Close within each z-plane. A full 3D element would erode a single-plane
    # (Z=1) stack to nothing, since its z-neighbours fall outside the array.
    mask = ndi.binary_closing(mask, structure=np.ones((1, 3, 3), bool))
    if mask.sum() == 0:
        z = np.zeros_like(mask)
        return TraceResult(mask, z, np.zeros(mask.shape, np.float32), voxel)
    skel = skeletonize(mask)
    if prune_um and prune_um > 0:
        skel = prune_skeleton(skel, voxel, spur_um=float(prune_um))
    dist = ndi.distance_transform_edt(mask, sampling=(dz, dy, dx)).astype(np.float32)
    return TraceResult(mask, skel, dist, voxel)


def trace_from_probability(prob: np.ndarray, voxel: tuple,
                           threshold: float = 0.5,
                           min_object_px: int = 32,
                           prune_um: float = 3.0) -> TraceResult:
    """Trace a 2D fiber-probability map (e.g. a MedCLIP map) into a skeleton.

    ``prob`` is a 2D float map in [0, 1] at the MIP resolution; ``threshold``
    binarises it into a fiber mask which is cleaned, skeletonised and pruned.
    The result is wrapped as a 3D (Z=1) ``TraceResult`` so the standard
    ``compute_metrics`` / ``skeleton_image`` helpers apply unchanged and the
    total length is measured with the same (dy, dx) spacing-aware code path.
    """
    dz, dy, dx = voxel
    prob = np.asarray(prob, dtype=np.float32)
    mask = prob >= float(threshold)
    if mask.any():
        mask = remove_small_objects(mask, int(min_object_px))
        mask = ndi.binary_closing(mask, structure=np.ones((3, 3), bool))
    if not mask.any():
        z = np.zeros((1,) + prob.shape, bool)
        return TraceResult(z, z.copy(), np.zeros(z.shape, np.float32), voxel)
    # Wrap the 2D skeleton as a single-plane (Z=1) volume so the 3D-voxel
    # pruning / metrics helpers apply unchanged (dz never affects in-plane steps).
    skel = skeletonize(mask)[None]
    if prune_um and prune_um > 0:
        skel = prune_skeleton(skel, voxel, spur_um=float(prune_um))
    dist = ndi.distance_transform_edt(mask, sampling=(dy, dx)).astype(np.float32)
    return TraceResult(mask[None], skel, dist[None], 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


RADIAL_AXIS_INDEX = {"Z": 0, "Y": 1, "X": 2}


def boundary_line_coords(shape_yx: tuple, boundary_frac: float, axis: str = "Y",
                         angle_deg: float = 0.0, curvature: float = 0.0) -> tuple:
    """Pixel coordinates of the (possibly tilted / curved) boundary line.

    Returns (ys, xs) integer arrays to index a 2D image, so the same line is
    drawn on every overlay and used to build the region masks — they can never
    disagree.  ``angle_deg`` tilts the line; ``curvature`` bows it (a parabola,
    positive = bow toward higher coordinates in the middle).  For axis="Z" there
    is no in-plane line, so empty arrays are returned.
    """
    ny, nx = shape_yx
    if axis.upper() == "Y":
        xs = np.arange(nx)
        cx, hx = nx / 2.0, max(nx / 2.0, 1.0)
        ys = (boundary_frac * ny + np.tan(np.deg2rad(angle_deg)) * (xs - cx)
              + curvature * ((xs - cx) / hx) ** 2 * ny)
        ys = np.clip(np.round(ys), 0, ny - 1).astype(int)
        return ys, xs
    if axis.upper() == "X":
        ys = np.arange(ny)
        cy, hy = ny / 2.0, max(ny / 2.0, 1.0)
        xs = (boundary_frac * nx + np.tan(np.deg2rad(angle_deg)) * (ys - cy)
              + curvature * ((ys - cy) / hy) ** 2 * nx)
        xs = np.clip(np.round(xs), 0, nx - 1).astype(int)
        return ys, xs
    return np.array([], int), np.array([], int)   # axis == "Z"


def make_region_masks(shape: tuple, boundary_frac: float,
                      ihc_side: str = "low", axis: str = "Y",
                      angle_deg: float = 0.0, curvature: float = 0.0) -> tuple:
    """Return (ihc_roi, ohc_roi) boolean masks split by the boundary.

    ``shape`` may be (ny, nx) or (nz, ny, nx).  axis="Y" splits along rows
    (radial axis, the usual case), axis="X" splits along columns; both accept a
    tilt (``angle_deg``) and a bow (``curvature``) and return 2D masks that are
    broadcast over Z during quantification.  axis="Z" splits by depth and
    returns 3D masks.  ``ihc_side`` selects which side ("low" = smaller
    coordinate) is IHC.
    """
    axis = axis.upper()
    if axis == "Z":
        if len(shape) != 3:
            raise ValueError("Z split needs a 3D (nz, ny, nx) shape")
        nz, ny, nx = shape
        if nz < 2:
            raise ValueError("Z (depth) split needs a multi-plane stack "
                             "(nz >= 2); this image has a single plane.")
        b = int(round(np.clip(boundary_frac, 0, 1) * nz))
        low = np.zeros((nz, ny, nx), bool)
        low[:b] = True
        ihc = low if ihc_side == "low" else ~low
        return ihc, ~ihc

    ny, nx = shape[-2], shape[-1]
    yy, xx = np.mgrid[0:ny, 0:nx]
    ang = np.tan(np.deg2rad(angle_deg))
    if axis == "Y":
        cx, hx = nx / 2.0, max(nx / 2.0, 1.0)
        yline = (boundary_frac * ny + ang * (xx - cx)
                 + curvature * ((xx - cx) / hx) ** 2 * ny)
        low = yy < yline
    else:                                        # axis == "X"
        cy, hy = ny / 2.0, max(ny / 2.0, 1.0)
        xline = (boundary_frac * nx + ang * (yy - cy)
                 + curvature * ((yy - cy) / hy) ** 2 * nx)
        low = xx < xline
    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 custom_cellpose_model_path() -> Optional[str]:
    """Path to a fine-tuned Cellpose hair-cell model, if one is provided.

    Looked up in order:
      1. the ``NEURON_TRACER_CP_MODEL`` environment variable, then
      2. a bundled ``models/hair_cell_cpsam`` file next to this module.
    Returns None to use the stock pretrained Cellpose model.  This is how a
    model trained with ``training/train_hair_cells.py`` gets picked up by the
    app: drop it at ``models/hair_cell_cpsam`` (or point the env var at it).
    """
    p = os.environ.get("NEURON_TRACER_CP_MODEL")
    if p and os.path.exists(p):
        return p
    here = os.path.dirname(os.path.abspath(__file__))
    bundled = os.path.join(here, "models", "hair_cell_cpsam")
    return bundled if os.path.exists(bundled) else None


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


def prep_detection_image(myo_vol: np.ndarray, voxel: tuple) -> tuple:
    """The exact Myo7a image the detector feeds to Cellpose, plus its scale.

    Returns ``(img_uint8, scale)`` where ``img_uint8`` is the contrast-
    normalised max-projection rescaled so hair cells sit near Cellpose's
    working size, and ``scale`` maps its coordinates back to full resolution
    (full = detected / scale).  Exposed so ``training/prepare_data.py`` can
    export annotation images at the *same* scale the model sees at inference —
    keeping training and inference consistent.
    """
    dx = voxel[2]
    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)
    return (small * 255).astype(np.uint8), scale


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.  If a fine-tuned model is provided (see
    ``custom_cellpose_model_path``) it is used automatically and the engine is
    reported as ``cellpose-custom``.  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.
    """
    mip_shape = myo_vol.max(0).shape
    small, scale = prep_detection_image(myo_vol, voxel)
    small = small.astype(np.float32) / 255.0

    engine = "watershed"
    masks = None
    if prefer_cellpose and CELLPOSE_AVAILABLE:
        try:
            masks = _get_cellpose_model().eval(small, diameter=_CP_TARGET_PX)[0]
            engine = ("cellpose-custom" if custom_cellpose_model_path()
                      else "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}


# Anatomy-based thresholds for accepting an IHC/OHC split (mouse organ of Corti).
_MIN_TWO_ROW_EXTENT_UM = 18.0   # IHC row + tunnel + OHC rows span well past this;
                                # a single hair-cell row spans < ~15 µm radially.
_MIN_TUNNEL_UM = 7.0            # the IHC/OHC gap (tunnel of Corti) is ~10-25 µm;
                                # inter-cell spacing within a row is smaller.


def auto_regions_from_cells(centroids: np.ndarray, shape_yx: tuple,
                            um_per_px: Optional[float] = None) -> 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.

    Crucially, a split is only proposed when the cells actually span more than
    one row (radial extent) *and* the gap is a real cell-free tunnel.  When only
    a single hair-cell row is present, ``single_row`` is returned True and no
    split is made — the field should be reported whole-field, not forced into a
    meaningless IHC/OHC division.
    """
    ny, nx = shape_yx
    result = {"boundary_frac": 0.5, "ihc_side": "low", "confidence": "low",
              "n_ihc_cells": 0, "n_ohc_cells": 0, "single_row": False,
              "reason": "not enough hair cells for an automatic split"}
    if len(centroids):
        result["boundary_frac"] = float(np.clip(centroids[:, 0].mean() / ny, 0, 1))
    if len(centroids) < 6:
        result["single_row"] = True
        result["reason"] = (f"only {len(centroids)} hair cell(s) detected — "
                            "cannot separate IHC vs OHC; report whole-field")
        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)

    # Convert radial distances to microns (or fall back to image-fraction units).
    if um_per_px and um_per_px > 0:
        to_um = um_per_px
        extent = (rs[-1] - rs[0]) * to_um
        min_extent, min_gap = _MIN_TWO_ROW_EXTENT_UM, _MIN_TUNNEL_UM
        unit = "µm"
    else:                                               # unknown scale
        to_um = 1.0
        extent = (rs[-1] - rs[0]) / max(ny, 1)          # fraction of image
        min_extent, min_gap = 0.12, 0.04
        unit = "frac"

    # Candidate splits must leave >=2 cells on each side (ignore stray outliers).
    valid = [(i, gaps[i]) for i in range(1, len(gaps) - 1)]
    biggest_gap = (max(valid, key=lambda t: t[1])[1] * to_um) if valid else 0.0

    # Single-row / no-tunnel guard: refuse to split an unresolved band.
    if extent < min_extent or not valid or biggest_gap < min_gap:
        result["single_row"] = True
        result["reason"] = (
            f"hair cells span only {extent:.1f} {unit} radially with no clear "
            f"tunnel gap — looks like a single row; IHC vs OHC cannot be "
            f"separated from this field. Report whole-field or set the boundary "
            f"by hand if you know the anatomy.")
        return result

    gi = max(valid, key=lambda t: t[1])[0]
    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))

    ratio = n_ohc / max(n_ihc, 1)
    strong_gap = (biggest_gap >= 2 * min_gap)
    if strong_gap and extent >= 1.6 * min_extent and 1.8 <= ratio <= 5.0:
        conf = "high"
    else:
        conf = "medium"

    return {"boundary_frac": bfrac, "ihc_side": ihc_side, "confidence": conf,
            "n_ihc_cells": n_ihc, "n_ohc_cells": n_ohc, "single_row": False,
            "reason": f"tunnel gap {biggest_gap:.1f} {unit}, radial extent "
                      f"{extent:.1f} {unit}, 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 _draw_boundary_rgb(rgb: np.ndarray, boundary_frac: Optional[float],
                       axis: str = "Y", angle_deg: float = 0.0,
                       curvature: float = 0.0, color=(255, 255, 0)) -> np.ndarray:
    """Draw the (tilted / curved) boundary line onto an RGB image in place.

    For axis="Z" there is no in-plane line, so nothing is drawn.
    """
    if boundary_frac is None:
        return rgb
    ny, nx = rgb.shape[:2]
    ys, xs = boundary_line_coords((ny, nx), boundary_frac, axis, angle_deg, curvature)
    if ys.size:
        rgb[ys, xs] = color
    return rgb


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,
                      angle_deg: float = 0.0, curvature: float = 0.0) -> np.ndarray:
    """Myo7a MIP with detected hair cells marked and the boundary drawn.

    Cells are coloured by the *actual* region mask (so tilt/curvature are
    honoured): cyan = IHC region, magenta = OHC region, neutral green when no
    split is defined.
    """
    rgb = np.stack([myo_mip_u8] * 3, axis=-1).copy()
    ny, nx = myo_mip_u8.shape
    ihc_mask = None
    if (boundary_frac is not None and ihc_side is not None
            and axis.upper() in ("Y", "X")):
        ihc_mask, _ = make_region_masks((ny, nx), boundary_frac, ihc_side,
                                        axis, angle_deg, curvature)
    for (y, x) in centroids.astype(int):
        yc, xc = min(max(y, 0), ny - 1), min(max(x, 0), nx - 1)
        col = (0, 255, 0)
        if ihc_mask is not None:
            col = (0, 220, 255) if ihc_mask[yc, xc] else (255, 60, 200)
        ys, xs = slice(max(0, y - 3), y + 4), slice(max(0, x - 3), x + 4)
        rgb[ys, xs] = col
    _draw_boundary_rgb(rgb, boundary_frac, axis, angle_deg, curvature)
    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 _fiber_stats(skel: np.ndarray, distance_um: np.ndarray, voxel: tuple,
                 min_fiber_um: float, min_diam_um: float, max_diam_um: float,
                 radial_axis: Optional[int]) -> tuple:
    """Per-fiber (connected-component) filtering and radial-direction counting.

    A fiber is kept only if its length >= ``min_fiber_um`` **and** its mean
    diameter is within [``min_diam_um``, ``max_diam_um``] (``max_diam_um`` <= 0
    means no upper limit).  When ``radial_axis`` is given, each kept fiber is
    classified as radial (its principal axis aligns with the IHC->OHC axis) or
    off-axis (running along the rows — a "misdirected" fiber).

    Returns (kept_skel_bool, n_fibers, total_length_um, n_radial, n_offaxis).
    ``n_radial``/``n_offaxis`` are -1 when direction was not computed.
    """
    struct = np.ones((3,) * skel.ndim, int)
    lbl, n = ndi.label(skel, structure=struct)
    if n == 0:
        return np.zeros_like(skel), 0, 0.0, -1, -1

    # Length per component: skan gives spacing-aware branch lengths; map each
    # branch to its connected component via a pixel lying on the branch.
    comp_len = {}
    try:
        S = Skeleton(skel, spacing=tuple(voxel))
        df = summarize(S, separator="_")
        for pi in df.index:
            coords = np.round(S.path_coordinates(pi)).astype(int)
            for a in range(coords.shape[1]):
                coords[:, a] = np.clip(coords[:, a], 0, skel.shape[a] - 1)
            c = 0
            for p in (coords[len(coords) // 2], coords[0], coords[-1]):
                c = int(lbl[tuple(p)])
                if c:
                    break
            if c:
                comp_len[c] = comp_len.get(c, 0.0) + float(df.loc[pi, "branch_distance"])
    except Exception:
        step = float(np.mean(voxel))
        for c in range(1, n + 1):
            comp_len[c] = float((lbl == c).sum()) * step

    spacing = np.asarray(voxel, float)
    slices = ndi.find_objects(lbl)
    kept = np.zeros(n + 1, bool)
    n_radial = n_offaxis = 0
    do_dir = radial_axis is not None
    for c in range(1, n + 1):
        sl = slices[c - 1]
        if sl is None:
            continue
        # A component with no measured length is a degenerate fragment (e.g. an
        # isolated voxel skan emits no branch for) — never count it as a fiber,
        # even when min_fiber_um is 0.  This also guarantees kept[c] => c in
        # comp_len, so the total_length sum below can index comp_len safely.
        if c not in comp_len or comp_len[c] < min_fiber_um:
            continue
        sub = lbl[sl] == c
        dsub = 2.0 * distance_um[sl][sub]
        mean_d = float(dsub.mean()) if dsub.size else 0.0
        if mean_d < min_diam_um or (max_diam_um > 0 and mean_d > max_diam_um):
            continue
        kept[c] = True
        if do_dir:
            offset = np.array([s.start for s in sl], float)
            pts = (np.argwhere(sub).astype(float) + offset) * spacing
            if len(pts) >= 2:
                pc = pts - pts.mean(0)
                _, _, vt = np.linalg.svd(pc, full_matrices=False)
                principal = vt[0]
                cos_r = abs(principal[radial_axis]) / (np.linalg.norm(principal) + 1e-9)
                n_radial += int(cos_r >= 0.5)
                n_offaxis += int(cos_r < 0.5)
            else:
                n_offaxis += 1

    kept_skel = kept[lbl]
    n_fibers = int(kept[1:].sum())
    total_length = float(sum(comp_len.get(c, 0.0)
                             for c in range(1, n + 1) if kept[c]))
    if not do_dir:
        n_radial = n_offaxis = -1
    return kept_skel, n_fibers, total_length, n_radial, n_offaxis


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,
                    min_diameter_um: float = 0.0,
                    max_diameter_um: float = 0.0,
                    manual_hair_cells: Optional[int] = None,
                    radial_axis_name: str = "Y") -> RegionMetrics:
    """Quantify the skeleton, optionally restricted to an ROI (2D or 3D).

    Fibers are filtered by length (``min_fiber_um``) and, when set, by mean
    diameter (``min_diameter_um`` / ``max_diameter_um``).  ``manual_hair_cells``
    overrides the detected hair-cell count for normalisation; otherwise
    ``hair_cell_centroids`` (if given) are counted inside the region.
    ``radial_axis_name`` ("Y"/"X"/"Z") is the IHC->OHC axis used to classify
    fiber direction.
    """
    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) if roi_yx.ndim == 2 else roi_yx
        skel = skel & roi3d
        mask = mask & roi3d

    m = RegionMetrics(region=region_name)

    # 2D footprint of the ROI drives area / FOV (handles 2D and 3D ROIs).
    if roi_yx is not None:
        roi_foot = roi_yx if roi_yx.ndim == 2 else roi_yx.max(0)
        m.fov_area_um2 = float(roi_foot.sum()) * dx * dy
    else:
        roi_foot = None
        m.fov_area_um2 = float(mask.shape[1] * mask.shape[2]) * dx * dy

    # Hair-cell count: manual override wins, else count detected centroids.
    if manual_hair_cells is not None and manual_hair_cells >= 0:
        m.n_hair_cells = int(manual_hair_cells)
    elif hair_cell_centroids is not None:
        if roi_foot is not None:
            m.n_hair_cells = count_cells_in_roi(hair_cell_centroids, roi_foot)
        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:
        name = (radial_axis_name or "Y").upper()
        if skel.ndim == 3:
            radial_axis = {"Z": 0, "Y": 1, "X": 2}.get(name, 1)
        else:
            radial_axis = {"Y": 0, "X": 1}.get(name, 0)
        kept_skel, n_fibers, total_len, n_radial, n_offaxis = _fiber_stats(
            skel, trace.distance_um, (dz, dy, dx), float(min_fiber_um),
            float(min_diameter_um), float(max_diameter_um), radial_axis)
        m.n_fibers = n_fibers
        m.total_length_um = total_len
        m.n_branch_points = _branch_point_count(kept_skel)
        diam = 2.0 * trace.distance_um[kept_skel]
        if diam.size:
            m.mean_diameter_um = float(diam.mean())
            m.median_diameter_um = float(np.median(diam))
        if n_radial >= 0:
            m.n_radial = n_radial
            m.n_offaxis = n_offaxis
            tot = n_radial + n_offaxis
            m.pct_radial = (100.0 * n_radial / tot) if tot else 0.0

    # Normalisation by hair-cell count (when known and positive).
    if m.n_hair_cells > 0:
        m.fibers_per_hc = m.n_fibers / m.n_hair_cells
        m.length_per_hc_um = m.total_length_um / m.n_hair_cells

    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",
                   angle_deg: float = 0.0, curvature: float = 0.0,
                   dilate: int = 1) -> np.ndarray:
    """Colour-coded RGB preview: IHC fibers cyan, OHC fibers magenta, with the
    boundary line in yellow.  Works with 2D (broadcast) or 3D (Z-split) ROIs."""
    if skel.ndim == 3:
        ihc3 = (np.broadcast_to(ihc_roi, skel.shape) if ihc_roi.ndim == 2
                else ihc_roi)
        ohc3 = (np.broadcast_to(ohc_roi, skel.shape) if ohc_roi.ndim == 2
                else ohc_roi)
        ihc_pix = (skel & ihc3).max(0)
        ohc_pix = (skel & ohc3).max(0)
        ny, nx = skel.shape[1], skel.shape[2]
    else:
        ihc_pix = skel & ihc_roi
        ohc_pix = skel & ohc_roi
        ny, nx = skel.shape
    if dilate:
        ihc_pix = ndi.binary_dilation(ihc_pix, iterations=dilate)
        ohc_pix = ndi.binary_dilation(ohc_pix, iterations=dilate)
    rgb = np.zeros((ny, nx, 3), np.uint8)
    rgb[ihc_pix] = (0, 220, 255)      # cyan  = IHC
    rgb[ohc_pix] = (255, 60, 200)     # magenta = OHC
    _draw_boundary_rgb(rgb, boundary_frac, axis, angle_deg, curvature)
    return rgb


def overlay_on_original(base_gray_u8: np.ndarray, skel: np.ndarray,
                        color=(255, 70, 70), dilate: int = 1) -> np.ndarray:
    """Traced skeleton drawn in colour over the original (grayscale) MIP.

    Lets the trace be checked against the raw signal — the reviewers' request
    to view the tracing overlaid on the original image.
    """
    flat = skel.max(0) if skel.ndim == 3 else skel
    if dilate:
        flat = ndi.binary_dilation(flat, iterations=dilate)
    rgb = np.stack([base_gray_u8] * 3, axis=-1).copy()
    rgb[flat] = color
    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)