"""Per-cell fluorescence quantification reproducing the analyst's ImageJ workflow.

Columns (same order as the Excel sheet 'GTTR Project Measurements'):
  Cell number, Total cell area, Nucleus area, Cytoplasm area, Total Cell Int Den, Nucleus Cell Int Den,
  Cytoplasm Int Den, Mean Cytoplasm Fluorescence (MCF), Background, MCF - background
Areas are in um^2 (pixel count / k^2, k = pixels per micron).  IntDen = sum of 8-bit red intensities
over the region / k^2 (this is ImageJ's IntDen = area * mean).  MCF = cytoplasm IntDen / cytoplasm area,
i.e. the mean 8-bit red intensity over the cytoplasm, independent of calibration.
"""
import numpy as np
from scipy import ndimage as ndi

COLUMNS = ['Cell number', 'Total cell area', 'Nucleus area', 'Cytoplasm area', 'Total Cell Int Den',
           'Nucleus Cell Int Den', 'Cytoplasm Int Den', 'Mean Cytoplasm Fluorescence (MCF)', 'Background',
           'MCF - background', 'Cell mean intensity (8-bit)', 'Nucleus mean intensity (8-bit)',
           'Cell mean intensity (ZEN scale)', 'Nucleus mean intensity (ZEN scale)']

ZEN_ESTIMATE_8BIT_EXPORT = {'mode': 'estimate', 'add': 1.5, 'mult': 256.0,
                            'desc': 'estimated as (8-bit mean + 1.5) x 256, assuming an 8-bit ZEN export of a 16-bit acquisition '
                                    '(the export maps the raw value v to floor(v / 256) - 1); checked on a ZEN-annotated image: '
                                    'nucleus means within -0.2% to +4.1% of ZEN'}

def zen_scale_value(mean8, zen):
    """ZEN-scale intensity mean from an 8-bit mean when the raw values are not available: zen = {'mode': 'estimate',
    'add': a, 'mult': m} gives (mean8 + a) * m; any other mode gives NaN."""
    if zen and zen.get('mode') == 'estimate' and mean8 == mean8:
        return (mean8 + float(zen['add'])) * float(zen['mult'])
    return float('nan')

def assign_nuclei(cells, nuclei):
    """Return a nucleus label map re-labelled with the id of the cell each nucleus belongs to.
    A nucleus goes to the cell that covers most of it (>= 50 % of its pixels); others are dropped.
    Nucleus pixels outside their cell are clipped to the cell."""
    out = np.zeros_like(cells)
    n_per_cell = {}
    for nl in np.unique(nuclei):
        if nl == 0: continue
        m = nuclei == nl
        labs, cnt = np.unique(cells[m], return_counts=True)
        order = np.argsort(cnt)[::-1]
        for o in order:
            if labs[o] == 0: continue
            if cnt[o] >= 0.5 * m.sum():
                out[m & (cells == labs[o])] = labs[o]
                n_per_cell[int(labs[o])] = n_per_cell.get(int(labs[o]), 0) + 1
            break
    return out, n_per_cell

DARKEST_PATCH_DESC = 'darkest empty 32x32 patch outside all detected cells'
FALLBACK_DESC = 'mean of all pixels outside the detected cells (no empty 32x32 patch in this image)'

def estimate_background(red, cells, erode_px=15, method='darkest_patch', return_method=False):
    """Background intensity from pixels outside every detected cell (eroded to avoid halos).
    With return_method=True returns (value, description of the estimator actually used)."""
    v, desc = _estimate_background(red, cells, erode_px, method)
    return (v, desc) if return_method else v

def _estimate_background(red, cells, erode_px, method):
    outside = cells == 0
    if erode_px:
        outside = ndi.binary_erosion(outside, iterations=erode_px)
    vals = red[outside]
    if vals.size < 100:
        vals = red[cells == 0]
    if vals.size == 0:
        return float('nan'), 'no pixels outside cells'
    if method == 'mode':
        h = np.bincount(vals.astype(np.int64).ravel(), minlength=256)
        return float(np.argmax(h)), 'mode of pixels outside cells'
    if method == 'p05':
        return float(np.percentile(vals, 5)), '5th percentile of pixels outside cells'
    if method == 'p10':
        return float(np.percentile(vals, 10)), '10th percentile of pixels outside cells'
    if method == 'median':
        return float(np.median(vals)), 'median of pixels outside cells'
    if method == 'mean':
        return float(vals.mean()), 'mean of pixels outside cells'
    if method == 'darkest_patch':
        # mean of the darkest 32x32 window that lies entirely outside every cell
        # (square erosion by 16 px keeps window centres >= 17 px, Chebyshev, from any cell pixel)
        sm = ndi.uniform_filter(red.astype(float), 32, mode='reflect')
        clear = ndi.binary_erosion(cells == 0, structure=np.ones((3, 3), bool), iterations=16)
        sm[~clear] = np.inf
        if not np.isfinite(sm).any(): return float(vals.mean()), FALLBACK_DESC
        return float(sm.min()), DARKEST_PATCH_DESC
    raise ValueError(f'unknown background method {method!r}')
    raise ValueError(method)

def measure(red, cells, nuclei, k_px_per_um, background=None, bg_method='darkest_patch', info=None, raw=None, zen=None):
    """red: (H,W) uint8 red channel; cells/nuclei: label maps; returns (rows, background, flags).
    If info (dict) is given, info['background_method'] records how the background was obtained.
    ZEN-scale means: exact means of `raw` (the acquisition values before the 8-bit mapping) when given, otherwise
    zen_scale_value(mean8, zen)."""
    red = red.astype(np.float64)
    nuc_of_cell, n_per_cell = assign_nuclei(cells, nuclei)
    if background is None:
        background, bg_desc = estimate_background(red, cells, method=bg_method, return_method=True)
    else:
        bg_desc = 'entered value'
    if info is not None:
        info['background_method'] = bg_desc
    if not np.isfinite(k_px_per_um) or k_px_per_um <= 0:
        raise ValueError('calibration (pixels per micron) must be a positive number')
    k2 = float(k_px_per_um) ** 2
    rows, flags = [], {}
    for i, lab in enumerate(np.unique(cells)):
        if lab == 0: continue
        cm = cells == lab
        nm = nuc_of_cell == lab
        cy = cm & ~nm
        ta, na, ca = cm.sum() / k2, nm.sum() / k2, cy.sum() / k2
        ti, ni, ci = red[cm].sum() / k2, red[nm].sum() / k2, red[cy].sum() / k2
        mcf = ci / ca if ca > 0 else float('nan')
        cell_mean = float(red[cm].mean()); nuc_mean = float(red[nm].mean()) if nm.any() else float('nan')
        if raw is not None:
            zc = float(raw[cm].mean()); zn = float(raw[nm].mean()) if nm.any() else float('nan')
        else:
            zc, zn = zen_scale_value(cell_mean, zen), zen_scale_value(nuc_mean, zen)
        rows.append({'Cell number': int(lab), 'Total cell area': ta, 'Nucleus area': na, 'Cytoplasm area': ca,
                     'Total Cell Int Den': ti, 'Nucleus Cell Int Den': ni, 'Cytoplasm Int Den': ci,
                     'Mean Cytoplasm Fluorescence (MCF)': mcf, 'Background': background,
                     'MCF - background': mcf - background,
                     'Cell mean intensity (8-bit)': cell_mean, 'Nucleus mean intensity (8-bit)': nuc_mean,
                     'Cell mean intensity (ZEN scale)': zc, 'Nucleus mean intensity (ZEN scale)': zn})
        nn = n_per_cell.get(int(lab), 0)
        if nn == 0: flags[int(lab)] = 'no nucleus found for this cell'
        elif nn > 1: flags[int(lab)] = f'{nn} nuclei merged'
        if ca <= 0: flags[int(lab)] = (flags.get(int(lab), '') + '; ' if int(lab) in flags else '') + 'no cytoplasm (nucleus covers the whole cell)'
    return rows, background, flags