# Copyright 2018 Olivier Scholder <o.scholder@gmail.com>

import numpy as np
import scipy.optimize as opt

"""
Provides some useful mathematical functions which are not present in numpy.
"""


def prod(x):
    p = 1
    for k in x:
        p *= k
    return p


def perm(x):
    """
    permutation with repetition of n factors
    x = [n1,n2,n3,...nk]
    n = n1+n2+...+nk
    return n!/(n1!*n2*...*nk!)
    """
    n = sum(x)
    num = list(range(2, n + 1))
    denom = list(range(2, x[0] + 1))
    for k in x[1:]:
        denom += list(range(2, k + 1))
    for x in num:
        if x in denom:
            num.remove(x)
            denom.remove(x)
    return prod(num) / prod(denom)


def closest_arg(array, value):
    return np.argmin(np.abs(array - value))


def strictly_positify(x):
    """
    Make the result strictly positive by setting the minimum value to
    the lower allowed float value
    """
    return np.fmax(x, np.finfo(x.dtype).eps)


def positify(x):
    """
    Set to zero all negative values
    """
    return np.fmax(0, x)


def clip01(x):
    """
    clip data x between 0 and 1
    """
    return np.fmax(np.fmin(x, 1), 0)


def fact(x):
    """
    Return the factors of an integer as a list.
    Warning: This function is not a factorial!
    """
    if x < 0 or type(x) is not int:
        raise ValueError("input must be a positive integer")
    if x < 2:
        return x
    f = []
    i = 2
    while True:
        while x % i == 0:
            f.append(i)
            x /= i
        i += 1
        if x == 1:
            return f


def moving_average(x, N):
    assert len(x) > N
    c = np.cumsum(x)
    return (c[N:] - c[:-N]) / N


def butter_lowpass(cutOff, fs, order=5):
    import scipy.signal

    nyq = 0.5 * fs
    normalCutoff = cutOff / nyq
    b, a = scipy.signal.butter(order, normalCutoff, btype="low", analog=True)
    return b, a


def butter_lowpass_filter(data, cutOff, fs, order=4):
    import scipy.signal

    b, a = butter_lowpass(cutOff, fs, order=order)
    y = scipy.signal.lfilter(b, a, data)
    return y


def Gauss(x, x0, s, amp=None, **kargs):
    if "A" in kargs:
        from warnings import warn

        warn(
            "Parameter A is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        amp = kargs.pop("A")
    elif "Amp" in kargs:
        from warnings import warn

        warn(
            "Parameter Amp is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        amp = kargs.pop("Amp")
    R = np.exp(-((x - x0) ** 2) / (2 * s**2))
    if amp is None:
        R /= s * np.sqrt(2 * np.pi)
    else:
        R *= amp
    R[s == 0] = (x[s == 0] == x0) * 1.0
    return R


def Lorentz(x, x0, gamma, amp=None, **kargs):
    if "A" in kargs:
        from warnings import warn

        warn(
            "Parameter A is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        kargs["A"]
    elif "Amp" in kargs:
        from warnings import warn

        warn(
            "Parameter Amp is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        amp = kargs.pop("Amp")
    R = 1 / ((x - x0) ** 2 + (0.5 * gamma) ** 2)
    if amp is None:
        return 0.5 * gamma * R / np.pi
    return amp * R * (0.5 * gamma) ** 2


def CDF(x, mu, sig, amp=1, lg=0, **kargs):
    if "Amp" in kargs:
        from warnings import warn

        warn(
            "Parameter Amp is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        amp = kargs.pop("Amp")
    from scipy.special import erf

    g = sig * np.sqrt(2 * np.log(2))
    return amp * lg * (0.5 + np.arctan2(x - mu, g) / np.pi) + (1 - lg) * amp * 0.5 * (
        1 + erf((x - mu) / (sig * np.sqrt(2)))
    )


def LG(x, x0, sig=None, amp=None, lg=0.5, asym=1, FWHM=None, **kargs):
    if "Amp" in kargs:
        from warnings import warn

        warn(
            "Parameter Amp is deprecated. Please use amp in order to set the amplitude!",
            stacklevel=2,
        )
        amp = kargs.pop("Amp")

    assert sig is not None or FWHM is not None

    if FWHM is None:
        FWHM = 2 * np.sqrt(2 * np.log(2)) * sig
    if sig is None:
        sig = FWHM / (2 * np.sqrt(2 * np.log(2)))
    Y = (1 - lg) * Gauss(x, x0, sig, amp=amp) + lg * Lorentz(x, x0, FWHM, amp=amp)
    if asym != 1:
        Yr = (1 - lg) * Gauss(x, x0, sig * asym, amp=amp) + lg * Lorentz(
            x, x0, FWHM * asym, amp=amp
        )
        Y[x > x0] = Yr[x > x0]
    return Y


def logistic(x, lower=0, upper=1, growth=1, x0=0, nu=1, C=1):
    return lower + (upper - lower) / (C + np.exp(-growth * (x - x0))) ** (1 / nu)


def fitCDF1line(A):
    line = np.zeros(A.shape[1])
    for x in range(A.shape[1]):
        popt, pcov = opt.curve_fit(
            CDF,
            np.arange(A.shape[0]),
            A[:, x],
            (A.shape[0] / 2, 1, np.max(A[:, x])),
            bounds=(0, (A.shape[1], np.inf, np.inf)),
        )
        line[x] = popt[0]
    return line


def FT(x, ufunc=np.real, real=False):
    """
    shortcut for 1D/2D Fourier Transform (real and centered)
    """
    assert isinstance(x, np.ndarray)
    if len(x.shape) == 1:
        F = np.fft.rfft(x) if real else np.fft.fft(x)
    elif len(x.shape) == 2:
        F = np.fft.rfft2(x) if real else np.fft.fft2(x)
    else:
        raise TypeError("The array should be 1D or 2D")
    return ufunc(np.fft.fftshift(F))


def binning(data, N=2, axis=0, ufunc=np.sum):
    w = int(np.floor(data.shape[axis] / N))
    r = np.copy(data)
    size = list(data.shape)
    size[axis] = w
    size = size[: axis + 1] + [N] + size[axis + 1 :]
    r.resize(size)
    return ufunc(r, axis=axis + 1)


def stat_info(data):
    import matplotlib.pyplot as plt

    D = np.ravel(data)
    U = np.unique(D)
    if len(U) > 1:
        sep = np.min(U[1:] - U[:-1])
        N = min(100, int(np.ceil((np.max(D) - np.min(D)) / sep)))
    else:
        N = 1

    mean = np.mean(D)
    std = np.std(D)

    fig, ax = plt.subplots(2, 1, figsize=(21, 4))
    ax[0].boxplot(D, 0, "ro", 0)
    ax[1].hist(D, N, density=True)
    ax[1].axvline(mean, color="r", label="mean")
    ax[1].axvline(mean + std, color="r", linestyle="--", label="1$\\sigma$")
    ax[1].axvline(mean - std, color="r", linestyle="--", label="1$\\sigma$")
    if mean - 2 * std >= U[0]:
        ax[1].axvline(mean - 2 * std, color="r", linestyle=":", label="2$\\sigma$")
    if mean + 2 * std <= U[-1]:
        ax[1].axvline(mean + 2 * std, color="r", linestyle=":", label="2$\\sigma$")
    ax[1].legend()
    print("Stats")
    print("\tAverage:", mean)
    print("\tStandard-deviation:", std)
    print("\tMinimum:", np.min(D))
    print("\tQ1:", np.percentile(D, 25))
    print("\tMedian:", np.percentile(D, 50))
    print("\tQ3:", np.percentile(D, 75))
    print("\tMaximum:", np.max(D))


def ellipse(a, b, phi):
    """
    return the x,y coordinates of an ellipse with major axis=a and minor axis=b at angle phi(radians)
    """
    r = np.zeros(phi.shape)  # if a & b is zero, result is zero
    m = np.logical_and(a != 0, b != 0)
    r[m] = a * b / np.sqrt((b * np.cos(phi[m])) ** 2 + (a * np.sin(phi[m])) ** 2)
    return r


def asymm_ellipse(left, right, upper, lower, phi):
    phi = np.divmod(phi + 2 * np.pi, 2 * np.pi)[1]  # Be sure phi ∈ [0,2π]
    b = np.where(phi <= np.pi, upper, lower)
    a = np.where(np.logical_or(phi <= np.pi / 2, phi >= 3 * np.pi / 2), right, left)
    m = np.logical_and(a != 0, b != 0)
    r = np.zeros(phi.shape)  # if a & b is zero, result is zero
    r[m] = (
        a[m]
        * b[m]
        / np.sqrt((b[m] * np.cos(phi[m])) ** 2 + (a[m] * np.sin(phi[m])) ** 2)
    )
    return r


def LG2D(
    XY,
    amplitude=1,
    angle=0,
    sig_x=10,
    sig_y=10,
    x0=None,
    y0=None,
    LG_x=0,
    LG_y=0,
    assym_x=1,
    assym_y=1,
    bg=0,
):
    """
    Return a 2D Lorentz-Gauss.
    XY: (X,Y) tuple
    A: amplitude
    a: angle
    sx: sigma for x-axis
    sy: sigma fdor y-axis
    x0,y0 : center coordinates of the peak
    lgx, lgy: Lorentz-Gauss proportion (for x,y axis)
    assym_x, assym_y: The assymetry in sig_x/sig_y for the left/right or upper/lower part of the curve
    """
    if x0 is None:
        x0 = XY[0][0, XY[0].shape[1] // 2]
    if y0 is None:
        y0 = XY[1][XY[1].shape[0] // 2, 0]
    X1 = (XY[0] - x0) * np.cos(angle) - (XY[1] - y0) * np.sin(angle)
    Y1 = (XY[0] - x0) * np.sin(angle) + (XY[1] - y0) * np.cos(angle)

    R1 = np.sqrt(X1**2 + Y1**2)
    angle = np.arctan2(Y1, X1)
    sig = asymm_ellipse(sig_x, sig_x * assym_x, sig_y, sig_y * assym_y, angle)
    gamma = np.sqrt(2 * np.log(2)) * sig
    LG = ellipse(LG_x + 1, LG_y + 1, angle) - 1

    Gxy = Gauss(R1, 0, sig, 1)
    Lxy = 1 / ((R1 / gamma) ** 2 + 1)

    f = (1 - LG) * Gxy + LG * Lxy
    out = bg + amplitude * f
    return out


def LG2Da(
    XY,
    amplitude=1,
    angle=0,
    sigN=10,
    sigS=None,
    sigE=10,
    sigW=None,
    x0=None,
    y0=None,
    LGN=0,
    LGS=None,
    LGE=0,
    LGW=None,
    bg=0,
):
    if x0 is None:
        x0 = XY[0][0, XY[0].shape[1] // 2]
    if y0 is None:
        y0 = XY[1][XY[1].shape[0] // 2, 0]
    if sigS is None:
        sigS = sigN
    if sigW is None:
        sigW = sigE
    if LGS is None:
        LGS = LGN
    if LGW is None:
        LGW = LGE

    X1 = (XY[0] - x0) * np.cos(angle) - (XY[1] - y0) * np.sin(angle)
    Y1 = (XY[0] - x0) * np.sin(angle) + (XY[1] - y0) * np.cos(angle)

    R1 = np.sqrt(X1**2 + Y1**2)
    angle = np.arctan2(Y1, X1)
    sig = asymm_ellipse(sigW, sigE, sigN, sigS, angle)
    gamma = np.sqrt(2 * np.log(2)) * sig  # HFHM
    LG = asymm_ellipse(LGW + 1, LGE + 1, LGN + 1, LGS + 1, angle) - 1

    Gxy = Gauss(R1, 0, sig, 1)
    Lxy = 1 / ((R1 / gamma) ** 2 + 1)

    f = (1 - LG) * Gxy + LG * Lxy
    out = bg + amplitude * f
    return out


def MaxwellBoltzmann(E, T):
    from . import constants as const

    return (
        2
        * const.qe
        * np.sqrt(E / np.pi)
        * np.exp(-E / (const.kb * T))
        / (const.kb * T) ** 1.5
    )


def Voigt(x, x0, sig, gamma, A=1):
    import scipy

    L = Lorentz(x, x0, gamma, A=1)
    G = Gauss(x, x0, sig)
    out = scipy.signal.convolve(L, G, "same")
    out /= np.max(out)
    return A * out
