"""Numeric functions ported from p63_thresholding_test.py. CPython 3 only."""

import numpy as np

# skimage.measure.regionprops is preferred; fall back to scipy.ndimage if unavailable
try:
    from skimage.measure import regionprops
    _HAVE_SKIMAGE = True
except ImportError:
    _HAVE_SKIMAGE = False


def extract_red(img, channel):
    """Extract channel index `channel` from a multi-channel image array; return a 2-D float64 ndarray normalised to [0, 1].

    Handles both (C, H, W) and (H, W, C) layouts: if the first dimension is
    smaller than the last, it is treated as the channel axis (C, H, W);
    otherwise the last axis is used (H, W, C).
    """
    if img.ndim == 2:
        ch = img
    elif img.ndim == 3:
        # Channel axis is first when its size is smaller than the spatial dims
        if img.shape[0] < img.shape[2]:
            ch = img[channel]
        else:
            ch = img[:, :, channel]
    else:
        raise ValueError("img must be 2-D or 3-D, got shape {}".format(img.shape))
    ch = ch.astype(np.float64)
    return ch / 255.0 if ch.max() > 1.0 else ch


def get_cell_data(masks, red_channel):
    """Compute per-cell mean intensity in red_channel for each labelled region in masks.

    Returns a list of dicts with keys: label, pixels, mean, coords.
    Background (label == 0) is ignored.
    """
    cells = []
    if _HAVE_SKIMAGE:
        for rp in regionprops(masks):
            rr, cc = rp.coords[:, 0], rp.coords[:, 1]
            pixels = red_channel[rr, cc]
            cells.append({
                "label":  rp.label,
                "pixels": pixels,
                "mean":   float(np.mean(pixels)),
                "coords": (rr, cc),
            })
    else:
        # scipy.ndimage fallback — slower but avoids skimage dependency
        import scipy.ndimage as ndi
        labels = np.unique(masks)
        labels = labels[labels != 0]
        for lbl in labels:
            coords = np.where(masks == lbl)
            pixels = red_channel[coords]
            cells.append({
                "label":  int(lbl),
                "pixels": pixels,
                "mean":   float(np.mean(pixels)),
                "coords": coords,
            })
    return cells


def otsu_threshold(values, n_bins=254, cap=None):
    """Compute Otsu threshold over a 1-D array of intensities in [0, 1]; return the bin-centre float.

    Exact logic from p63_thresholding_test.py: background (v > 0) exclusion,
    n_bins=254 bins, sigma_b2[-1] forced to 0 to avoid last-bin artefact.

    Parameters
    ----------
    values : array-like of intensities
    n_bins : number of histogram bins (default 254 to match validated results)
    cap    : hard ceiling applied before histogram; activates T2/T14 attenuation variants
    """
    v = np.asarray(values, dtype=np.float64).ravel()
    v = v[v > 0]                            # exclude background / masked pixels
    if cap is not None:
        v = np.clip(v, 0.0, cap)
    if v.size == 0:
        return 0.0

    upper = cap if cap is not None else 1.0
    edges = np.linspace(0.0, upper, n_bins + 1)
    counts, _ = np.histogram(v, bins=edges)

    total = counts.sum()
    if total == 0:
        return 0.0

    p        = counts / total
    omega    = np.cumsum(p)                     # cumulative class probabilities
    mu       = np.cumsum(p * np.arange(len(p))) # cumulative weighted means
    mu_t     = mu[-1]                           # overall mean

    denom    = omega * (1.0 - omega)
    sigma_b2 = (mu_t * omega - mu) ** 2 / (denom + 1e-12)
    sigma_b2[-1] = 0.0                          # avoid artefact at last bin

    k           = int(np.argmax(sigma_b2))
    bin_centres = (edges[:-1] + edges[1:]) / 2.0
    return float(bin_centres[k])


def compute_thresholds(cells, attenuation_cap=0.7, n_bins=254):
    """Compute T1, T2, T7, T14 thresholds from per-cell data dicts; return dict of floats in [0, 1].

    T1  = pixel-level Otsu, full range
    T2  = pixel-level Otsu, capped at attenuation_cap
    T7  = cell-mean-level Otsu, full range
    T14 = cell-mean-level Otsu, capped at attenuation_cap  (default p63+ call)
    """
    all_pixels = np.concatenate([c["pixels"] for c in cells]) if cells else np.array([])
    cell_means = np.array([c["mean"] for c in cells])

    return {
        "T1":  otsu_threshold(all_pixels, n_bins=n_bins),
        "T2":  otsu_threshold(all_pixels, n_bins=n_bins, cap=attenuation_cap),
        "T7":  otsu_threshold(cell_means,  n_bins=n_bins),
        "T14": otsu_threshold(cell_means,  n_bins=n_bins, cap=attenuation_cap),
    }


def count_p63_positive(cells, threshold):
    """Return the count of cells whose mean red intensity >= threshold."""
    return sum(1 for c in cells if c["mean"] >= threshold)