import os

import matplotlib.pyplot as plt
import numpy as np
import tifffile
from scipy.ndimage import *
from skimage import io as skio
from skimage.morphology import *
from skimage.transform import resize


def tifread(path):
    return tifffile.imread(path)


def tifwrite(I, path):
    tifffile.imsave(path, I)


def imshow(I, **kwargs):
    if not kwargs:
        plt.imshow(I, cmap="gray")
    else:
        plt.imshow(I, **kwargs)

    plt.axis("off")
    plt.show()


def imshowlist(L, **kwargs):
    n = len(L)
    for i in range(n):
        plt.subplot(1, n, i + 1)
        if not kwargs:
            plt.imshow(L[i], cmap="gray")
        else:
            plt.imshow(L[i], **kwargs)
        plt.axis("off")
    plt.show()


def imread(path):
    return skio.imread(path)


def imwrite(I, path):
    skio.imsave(path, I)


def im2double(I):
    if I.dtype == "uint16":
        return I.astype("float64") / 65535
    elif I.dtype == "uint8":
        return I.astype("float64") / 255
    elif I.dtype == "float32":
        return I.astype("float64")
    elif I.dtype == "float64":
        return I
    else:
        print("returned original image type: ", I.dtype)
        return I


def size(I):
    return list(I.shape)


def imresizeDouble(I, sizeOut):  # input and output are double
    return resize(I, (sizeOut[0], sizeOut[1]), mode="reflect")


def imresize3Double(I, sizeOut):  # input and output are double
    return resize(I, (sizeOut[0], sizeOut[1], sizeOut[2]), mode="reflect")


def imresizeUInt8(I, sizeOut):  # input and output are UInt8
    return np.uint8(
        resize(I.astype(float), (sizeOut[0], sizeOut[1]), mode="reflect", order=0)
    )


def imresize3UInt8(I, sizeOut):  # input and output are UInt8
    return np.uint8(
        resize(
            I.astype(float),
            (sizeOut[0], sizeOut[1], sizeOut[2]),
            mode="reflect",
            order=0,
        )
    )


def normalize(I):
    m = np.min(I)
    M = np.max(I)
    if M > m:
        return (I - m) / (M - m)
    else:
        return I


def snormalize(I):
    m = np.mean(I)
    s = np.std(I)
    if s > 0:
        return (I - m) / s
    else:
        return I


def cat(a, I, J):
    return np.concatenate((I, J), axis=a)


def imerode(I, r):
    return binary_erosion(I, disk(r))


def imdilate(I, r):
    return binary_dilation(I, disk(r))


def imerode3(I, r):
    return morphology.binary_erosion(I, ball(r))


def imdilate3(I, r):
    return morphology.binary_dilation(I, ball(r))


def sphericalStructuralElement(imShape, fRadius):
    if len(imShape) == 2:
        return disk(fRadius, dtype=float)
    if len(imShape) == 3:
        return ball(fRadius, dtype=float)


def medfilt(I, filterRadius):
    return median_filter(I, footprint=sphericalStructuralElement(I.shape, filterRadius))


def maxfilt(I, filterRadius):
    return maximum_filter(
        I, footprint=sphericalStructuralElement(I.shape, filterRadius)
    )


def minfilt(I, filterRadius):
    return minimum_filter(
        I, footprint=sphericalStructuralElement(I.shape, filterRadius)
    )


def ptlfilt(I, percentile, filterRadius):
    return percentile_filter(
        I, percentile, footprint=sphericalStructuralElement(I.shape, filterRadius)
    )


def imgaussfilt(I, sigma, **kwargs):
    return gaussian_filter(I, sigma, **kwargs)


def imlogfilt(I, sigma, **kwargs):
    return -gaussian_laplace(I, sigma, **kwargs)


def imgradmag(I, sigma):
    if len(I.shape) == 2:
        dx = imgaussfilt(I, sigma, order=[0, 1])
        dy = imgaussfilt(I, sigma, order=[1, 0])
        return np.sqrt(dx**2 + dy**2)
    if len(I.shape) == 3:
        dx = imgaussfilt(I, sigma, order=[0, 0, 1])
        dy = imgaussfilt(I, sigma, order=[0, 1, 0])
        dz = imgaussfilt(I, sigma, order=[1, 0, 0])
        return np.sqrt(dx**2 + dy**2 + dz**2)


def localstats(I, radius, justfeatnames=False):
    ptls = [10, 30, 50, 70, 90]
    featNames = []
    for i in range(len(ptls)):
        featNames.append("locPtl%d" % ptls[i])
    if justfeatnames == True:
        return featNames
    sI = size(I)
    nFeats = len(ptls)
    F = np.zeros((sI[0], sI[1], nFeats))
    for i in range(nFeats):
        F[:, :, i] = ptlfilt(I, ptls[i], radius)
    return F


def localstats3(I, radius, justfeatnames=False):
    ptls = [10, 30, 50, 70, 90]
    featNames = []
    for i in range(len(ptls)):
        featNames.append("locPtl%d" % ptls[i])
    if justfeatnames == True:
        return featNames
    sI = size(I)
    nFeats = len(ptls)
    F = np.zeros((sI[0], sI[1], sI[2], nFeats))
    for i in range(nFeats):
        F[:, :, :, i] = ptlfilt(I, ptls[i], radius)
    return F


def imderivatives(I, sigmas, justfeatnames=False):
    if type(sigmas) is not list:
        sigmas = [sigmas]
    derivPerSigmaFeatNames = [
        "d0",
        "dx",
        "dy",
        "dxx",
        "dxy",
        "dyy",
        "normGrad",
        "normHessDiag",
    ]
    if justfeatnames == True:
        featNames = []
        for i in range(len(sigmas)):
            for j in range(len(derivPerSigmaFeatNames)):
                featNames.append(
                    "derivSigma%d%s" % (sigmas[i], derivPerSigmaFeatNames[j])
                )
        return featNames
    nDerivativesPerSigma = len(derivPerSigmaFeatNames)
    nDerivatives = len(sigmas) * nDerivativesPerSigma
    sI = size(I)
    D = np.zeros((sI[0], sI[1], nDerivatives))
    for i in range(len(sigmas)):
        sigma = sigmas[i]
        dx = imgaussfilt(I, sigma, order=[0, 1])
        dy = imgaussfilt(I, sigma, order=[1, 0])
        dxx = imgaussfilt(I, sigma, order=[0, 2])
        dyy = imgaussfilt(I, sigma, order=[2, 0])
        D[:, :, nDerivativesPerSigma * i] = imgaussfilt(I, sigma)
        D[:, :, nDerivativesPerSigma * i + 1] = dx
        D[:, :, nDerivativesPerSigma * i + 2] = dy
        D[:, :, nDerivativesPerSigma * i + 3] = dxx
        D[:, :, nDerivativesPerSigma * i + 4] = imgaussfilt(I, sigma, order=[1, 1])
        D[:, :, nDerivativesPerSigma * i + 5] = dyy
        D[:, :, nDerivativesPerSigma * i + 6] = np.sqrt(dx**2 + dy**2)
        D[:, :, nDerivativesPerSigma * i + 7] = np.sqrt(dxx**2 + dyy**2)
    return D
    # derivatives are indexed by the last dimension, which is good for ML features but not for visualization,
    # in which case the expected dimensions are [plane,channel,y(row),x(col)]; to obtain that ordering, do
    # D = np.moveaxis(D,[0,3,1,2],[0,1,2,3])


def imderivatives3(I, sigmas, justfeatnames=False):
    if type(sigmas) is not list:
        sigmas = [sigmas]

    derivPerSigmaFeatNames = [
        "d0",
        "dx",
        "dy",
        "dz",
        "dxx",
        "dxy",
        "dxz",
        "dyy",
        "dyz",
        "dzz",
        "normGrad",
        "normHessDiag",
    ]

    # derivPerSigmaFeatNames = ['d0','normGrad','normHessDiag']

    if justfeatnames == True:
        featNames = []
        for i in range(len(sigmas)):
            for j in range(len(derivPerSigmaFeatNames)):
                featNames.append(
                    "derivSigma%d%s" % (sigmas[i], derivPerSigmaFeatNames[j])
                )
        return featNames
    nDerivativesPerSigma = len(derivPerSigmaFeatNames)
    nDerivatives = len(sigmas) * nDerivativesPerSigma
    sI = size(I)
    D = np.zeros((sI[0], sI[1], sI[2], nDerivatives))  # plane, channel, y, x
    for i in range(len(sigmas)):
        sigma = sigmas[i]
        dx = imgaussfilt(I, sigma, order=[0, 0, 1])  # z, y, x
        dy = imgaussfilt(I, sigma, order=[0, 1, 0])
        dz = imgaussfilt(I, sigma, order=[1, 0, 0])
        dxx = imgaussfilt(I, sigma, order=[0, 0, 2])
        dyy = imgaussfilt(I, sigma, order=[0, 2, 0])
        dzz = imgaussfilt(I, sigma, order=[2, 0, 0])

        D[:, :, :, nDerivativesPerSigma * i] = imgaussfilt(I, sigma)
        D[:, :, :, nDerivativesPerSigma * i + 1] = dx
        D[:, :, :, nDerivativesPerSigma * i + 2] = dy
        D[:, :, :, nDerivativesPerSigma * i + 3] = dz
        D[:, :, :, nDerivativesPerSigma * i + 4] = dxx
        D[:, :, :, nDerivativesPerSigma * i + 5] = imgaussfilt(
            I, sigma, order=[0, 1, 1]
        )
        D[:, :, :, nDerivativesPerSigma * i + 6] = imgaussfilt(
            I, sigma, order=[1, 0, 1]
        )
        D[:, :, :, nDerivativesPerSigma * i + 7] = dyy
        D[:, :, :, nDerivativesPerSigma * i + 8] = imgaussfilt(
            I, sigma, order=[1, 1, 0]
        )
        D[:, :, :, nDerivativesPerSigma * i + 9] = dzz
        D[:, :, :, nDerivativesPerSigma * i + 10] = np.sqrt(dx**2 + dy**2 + dz**2)
        D[:, :, :, nDerivativesPerSigma * i + 11] = np.sqrt(dxx**2 + dyy**2 + dzz**2)

        # D[:,:,:,nDerivativesPerSigma*i   ] = imgaussfilt(I,sigma)
        # D[:,:,:,nDerivativesPerSigma*i+1 ] = np.sqrt(dx**2+dy**2+dz**2)
        # D[:,:,:,nDerivativesPerSigma*i+2 ] = np.sqrt(dxx**2+dyy**2+dzz**2)
    return D
    # derivatives are indexed by the last dimension, which is good for ML features but not for visualization,
    # in which case the expected dimensions are [plane,y(row),x(col)]; to obtain that ordering, do
    # D = np.moveaxis(D,[2,0,1],[0,1,2])


def imfeatures(I=[], sigmaDeriv=1, sigmaLoG=1, locStatsRad=0, justfeatnames=False):
    if type(sigmaDeriv) is not list:
        sigmaDeriv = [sigmaDeriv]
    if type(sigmaLoG) is not list:
        sigmaLoG = [sigmaLoG]
    derivFeatNames = imderivatives([], sigmaDeriv, justfeatnames=True)
    nLoGFeats = len(sigmaLoG)
    locStatsFeatNames = []
    if locStatsRad > 1:
        locStatsFeatNames = localstats([], locStatsRad, justfeatnames=True)
    nLocStatsFeats = len(locStatsFeatNames)
    if justfeatnames == True:
        featNames = derivFeatNames
        for i in range(nLoGFeats):
            featNames.append("logSigma%d" % sigmaLoG[i])
        for i in range(nLocStatsFeats):
            featNames.append(locStatsFeatNames[i])
        return featNames
    nDerivFeats = len(derivFeatNames)
    nFeatures = nDerivFeats + nLoGFeats + nLocStatsFeats
    sI = size(I)
    F = np.zeros((sI[0], sI[1], nFeatures))
    F[:, :, :nDerivFeats] = imderivatives(I, sigmaDeriv)
    for i in range(nLoGFeats):
        F[:, :, nDerivFeats + i] = imlogfilt(I, sigmaLoG[i])
    if locStatsRad > 1:
        F[:, :, nDerivFeats + nLoGFeats :] = localstats(I, locStatsRad)
    return F


def imfeatures3(I=[], sigmaDeriv=2, sigmaLoG=2, locStatsRad=0, justfeatnames=False):
    if type(sigmaDeriv) is not list:
        sigmaDeriv = [sigmaDeriv]
    if type(sigmaLoG) is not list:
        sigmaLoG = [sigmaLoG]
    derivFeatNames = imderivatives3([], sigmaDeriv, justfeatnames=True)
    nLoGFeats = len(sigmaLoG)
    locStatsFeatNames = []
    if locStatsRad > 1:
        locStatsFeatNames = localstats3([], locStatsRad, justfeatnames=True)
    nLocStatsFeats = len(locStatsFeatNames)
    if justfeatnames == True:
        featNames = derivFeatNames
        for i in range(nLoGFeats):
            featNames.append("logSigma%d" % sigmaLoG[i])
        for i in range(nLocStatsFeats):
            featNames.append(locStatsFeatNames[i])
        return featNames
    nDerivFeats = len(derivFeatNames)
    nFeatures = nDerivFeats + nLoGFeats + nLocStatsFeats
    sI = size(I)
    F = np.zeros((sI[0], sI[1], sI[2], nFeatures))
    F[:, :, :, :nDerivFeats] = imderivatives3(I, sigmaDeriv)
    for i in range(nLoGFeats):
        F[:, :, :, nDerivFeats + i] = imlogfilt(I, sigmaLoG[i])
    if locStatsRad > 1:
        F[:, :, :, nDerivFeats + nLoGFeats :] = localstats3(I, locStatsRad)
    return F


def stack2list(S):
    L = []
    for i in range(size(S)[2]):
        L.append(S[:, :, i])
    return L


def thrsegment(I, wsBlr, wsThr):  # basic threshold segmentation
    G = imgaussfilt(I, sigma=(1 - wsBlr) + wsBlr * 5)  # min 1, max 5
    M = G > wsThr
    return M
