import numpy as np
import scipy
import skimage
import skimage.feature
import skimage.segmentation
from tqdm.auto import tqdm

# ----------------------- Helper functions ---------------------------------


def maskpk(pk, mask):
    """remove points in nx2 array which are located outside mask
    wp@tl20190709
    """
    return pk[[mask[tuple(p)] > 0 for p in np.round(pk).astype(int)], :]


def disk(s):
    """make a disk shaped structural element to be used with
    morphological functions
    wp@tl20190709
    """
    d = np.zeros((s, s))
    c = (s - 1) / 2.0
    x, y = np.meshgrid(range(s), range(s))
    d2 = (x - c) ** 2 + (y - c) ** 2
    d[d2 < s**2 / 4.0] = 1
    return d


def fill_nan(im):
    """Assigns the value of the nearest finite valued pixel to infinite valued pixels
    wp@tl20190910
    """
    im = im.copy()
    b = np.where(~np.isfinite(im))
    if len(b[0]):
        a = np.where(np.isfinite(im))
        v = scipy.interpolate.griddata(a, im[a[0], a[1]], b, "nearest")
        for i, j, w in zip(b[0], b[1], v):
            im[i, j] = w
    return im


def collectr(cf):
    """Makes a 1d corrfun out of a 2d corrfun
    The result of the maximum of cf at each distance from the center
    wp@tl20191102
    """
    x, y = np.meshgrid(range(cf.shape[0]), range(cf.shape[1]))
    c = [(i - 1) / 2 for i in cf.shape]
    r = np.sqrt((x - c[0]) ** 2 + (y - c[1]) ** 2).flatten()
    idx = np.argsort(r)
    r = np.round(r[idx]).astype("int")
    cf = np.round(cf.flatten()[idx]).astype("int")
    rr = np.unique(r)
    cf = [np.max(cf[r == i]) for i in rr]
    return rr, cf


def corrfft(im, jm):
    """
    usage: d, cfunc = corrfft(images)
    input:
       im, jm: images to be correlated
    output:
       d:      offset (x,y) in px
       cfunc:  correlation function
    """

    im = im.astype(float)
    jm = jm.astype(float)

    im -= np.nanmean(im)
    im /= np.nanstd(im)
    jm -= np.nanmean(jm)
    jm /= np.nanstd(jm)

    im[np.isnan(im)] = 0
    jm[np.isnan(jm)] = 0

    n_y = np.shape(im)[0]
    n_x = np.shape(im)[1]

    cfunc = np.real(
        np.fft.fftshift(np.fft.ifft2(np.fft.fft2(im) * np.conj(np.fft.fft2(jm))))
    )
    y, x = np.unravel_index(np.nanargmax(cfunc), cfunc.shape)

    d = [x - np.floor(n_x / 2), y - np.floor(n_y / 2)]

    # peak at x=nX-1 means xoffset=-1
    if d[0] > n_x / 2:
        d[0] -= n_x
    if d[1] > n_y / 2:
        d[1] -= n_y
    return d, cfunc


def otsu_local(im, res=16):
    # Executes otsu on blocks of an image, then interpolates to give a local threshold
    # im: array with the image
    # res: image will be cut in res x res blocks
    # Image size should be an integer multiple of res!
    #
    # wp@tl20191121
    if res == 1:  # global otsu
        d = im.flatten()
        d = d[np.isfinite(d)]
        if len(np.unique(d)) > 1:
            return skimage.filters.threshold_otsu(d)
        else:
            return d.flatten()[0]
    th = np.zeros((res, res))
    s = im.shape[0] // res
    for i in range(res):
        for j in range(res):
            d = im[i * s : (i + 1) * s, j * s : (j + 1) * s].flatten()
            d = d[np.isfinite(d)]
            if len(d) == 0:
                th[i, j] = np.nan
            elif len(np.unique(d)) > 1:
                th[i, j] = skimage.filters.threshold_otsu(d)
            else:
                th[i, j] = d.flatten()[0]
    th = skimage.transform.resize(th, im.shape, anti_aliasing=False, mode="edge")
    return scipy.ndimage.gaussian_filter(th, im.shape[0] / 8, mode="reflect")


def fill_holes(im):
    """Fill holes (value==0) in each label individually in a faster way
    wp@tl20200120
    """
    im = im.copy()
    z = skimage.morphology.dilation(im) * (im == 0)
    holes_lbl = skimage.measure.label(im == 0)
    lbl_edges = np.unique(
        np.hstack(
            (holes_lbl[:, 0], holes_lbl[:, -1], holes_lbl[0, :], holes_lbl[-1, :])
        )
    )
    lbl_h = list(np.unique(holes_lbl))
    for i in lbl_h:
        if i == 0 or i in lbl_edges:
            continue
        lbl_edge = list(np.unique((holes_lbl == i) * z))
        if 0 in lbl_edge:
            lbl_edge.remove(0)
        if len(lbl_edge) == 1:
            im += lbl_edge[0] * (holes_lbl == i)
    return im


def make_mask(im, threshold, thres, smooth=2.5, min_feature_size=5, dilate=5):
    lb = fill_nan(im)
    if smooth:
        la = scipy.ndimage.gaussian_filter(lb, smooth, mode="nearest")
    else:
        la = lb.copy()
    la[~np.isfinite(im)] = np.nan
    if threshold is None or isinstance(threshold, str) and threshold.lower() == "none":
        th = otsu_local(la, thres)
    else:
        th = float(threshold)
    mask = la > th
    if min_feature_size > 0:
        mask = skimage.morphology.binary_opening(mask, disk(min_feature_size))
    if dilate == 0:
        return mask, la
    elif dilate > 0:
        return skimage.morphology.binary_dilation(mask, disk(dilate)), la
    else:
        return skimage.morphology.binary_erosion(mask, disk(-dilate)), la


# ----------------------- Finally the function itself ---------------------------------


def findcells(
    im,
    imnuc=None,
    cellcolormask=None,
    ccdist=None,
    threshold=None,
    thresholdnuc=None,
    thres=1,
    thresnuc=1,
    smooth=2.5,
    smoothnuc=2.5,
    minfeatsize=5,
    minfeatsizenuc=5,
    dilate=5,
    dilatenuc=5,
    removeborders=True,
    **_kwargs,
):
    """segment cells and nuclei from an image (nxm array)
    wp@tl20190710

    im: 2d numpy array with an image

    optional:
    imnuc: 2d numpy array with an image containing the nuclei, for example with DAPI
    cellcolormask: use a cellmask from another frame to assign the same colors
                   to the cells found by findcells
    ccdist:        (approximate) minimum distance between the centers of cells, good values: yeast: 25,
       mamalian: 150

    optional parameters to debug/finetune findcells:
    threshold:     value used to threshold the image, default: use Otsu
    thres:         divide the image in thres x thres blocks and use a different threshold in each of them,
                   the image size should be an integer multiple of thres
    smooth:        smooth the image before using it
    minfeatsize:   remove features smaller than minfeatsize
    dilate:        >0: make every cell dilate px larger, <0: make them smaller
    thresholdnuc, thresnuc, smoothnuc, minfeatsizenuc, dilatenuc: parameters applying to imnuc
    removeborders: remove any cells (and their nuclei) touching any borders (case with nan's unhandled atm),
        default: True
    """
    im = np.asarray(im).astype("float")
    if imnuc is not None:
        imnuc = np.asarray(imnuc).astype("float")

    if ccdist is None:  # try to determine good ccdist if not given
        if imnuc is None:
            cf = corrfft(im.copy(), im.copy())[1]
        else:
            cf = corrfft(imnuc.copy(), imnuc.copy())[1]
        r, cf = collectr(cf)
        ccdist = r[scipy.signal.find_peaks(-np.array(cf))[0][0]]
        ccdist = np.round(np.clip(ccdist, 10, np.sqrt(np.prod(im.shape)) / 5)).astype(
            "int"
        )

    if imnuc is None:  # only one channel
        mask, la = make_mask(im, threshold, thres, smooth, minfeatsize, dilate)
        pk = skimage.feature.peak_local_max(
            fill_nan(la), min_distance=ccdist // 2, exclude_border=False
        )
        pk = maskpk(pk, mask)
        pk = np.array(sorted([q.tolist() for q in pk])[::-1])
        markers = np.zeros(im.shape, int)
        for i in range(pk.shape[0]):
            if cellcolormask is None:
                markers[pk[i, 0], pk[i, 1]] = i + 1
            else:
                markers[pk[i, 0], pk[i, 1]] = cellcolormask[pk[i, 0], pk[i, 1]]
        cells = skimage.segmentation.watershed(-la, markers=markers, mask=mask)
        cells *= mask  # in Python2, watershed sometimes colors outside the mask
        lbl = list(np.unique(cells))
        if 0 in lbl:
            lbl.remove(0)
        cells = fill_holes(cells)
        nuclei = np.zeros(im.shape)
        for i in tqdm(
            lbl, disable=len(lbl) < 25, leave=False
        ):  # threshold each cell to find its nucleus
            cellmask = cells == i
            cm_la = cellmask * la
            cell = cm_la.flatten()
            cell = cell[cell > 0]
            cell = cell[cell > np.percentile(cell, 25)]
            pxval = np.unique(cell)
            if len(pxval) == 0:
                th = 0
            elif len(pxval) == 1:
                th = cell[0]
            else:
                th = skimage.filters.threshold_otsu(cell)
            nucleus = cm_la > th
            nuclei += float(i) * nucleus
    else:  # extra channel with nuclei
        nuclei, _ = findcells(
            imnuc,
            cellcolormask=cellcolormask,
            ccdist=ccdist,
            threshold=thresholdnuc,
            thres=thresnuc,
            smooth=smoothnuc,
            minfeatsize=minfeatsizenuc,
            dilate=dilatenuc,
            removeborders=False,
        )
        mask, la = make_mask(im, threshold, thres, smooth, minfeatsize, dilate)
        cells = skimage.segmentation.watershed(-la, markers=nuclei, mask=mask)
        cells *= mask  # in Python2, watershed sometimes colors outside the mask
        cells = np.dstack((cells, nuclei)).max(
            2
        )  # make sure each nucleus has a cell body
        cells = fill_holes(cells)

    if removeborders:
        for lbl in np.unique(
            np.hstack((cells[:, 0], cells[:, -1], cells[0, :], cells[-1, :]))
        ):
            cells[cells == lbl] = 0
            nuclei[nuclei == lbl] = 0
    return cells, nuclei
