"""
*********************
Correlation functions
*********************

Decay functions (KWW, ML,...)
"""

import numpy as np
from scipy.special import gammaincc

from ..math.mittagleffler import mlf


class Exponential:
    """
    Stretched exponential function
    """
    type = 'Correlation function'
    name = 'Stretched Exponential'
    equation = r'C*exp(-(x/\tau)^{\beta})'
    params = ['C', r'\tau', r'\beta']

    @staticmethod
    def func(x, c, x0, beta):
        r"""
        Stretched exponential function

        .. math::
            x = c\exp(-x/x_0)

        Args:
            x (array-like): Input values
            c (float): Amplitude
            x0 (float):
            beta (float): Stretching parameter
        """
        return c*np.exp(-(x/x0)**beta)


class MittagLefflerCC:
    type = 'Correlation function'
    name = 'Cole-Cole'
    equation = r'C * E_{\alpha}(-(x/\tau)^{\alpha})'
    params = ['C', r'\tau', r'\alpha']
    bounds = [(None, None), (0, None), (0, 1)]

    @staticmethod
    def func(x, c, tau, alpha):
        return c * mlf(-(x/tau)**alpha, alpha)


class GammaCD:
    type = 'Correlation function'
    name = 'Cole-Davidson'
    equation = r'C * \Gamma(\gamma, x/\tau) / \Gamma(\gamma)'
    params = ['C', r'\tau', r'\gamma']
    bounds = [(None, None), (0, None), (0, 1)]

    @staticmethod
    def func(x, c, tau, gamma):
        return c * gammaincc(gamma, x/tau)


class MLHavNeg:
    type = 'Correlation function'
    name = 'Havriliak-Negami'
    equation = r'C * E_{\alpha,\alpha\gamma+1}^{\gamma}(-(x/\tau)^{\alpha})'
    params = ['C', r'\tau', r'\gamma']
    bounds = [(None, None), (0, None), (0, 1)]

    @staticmethod
    def func(x, c, tau, alpha, gamma):
        return c * (1 - (x/tau)**(alpha*gamma) * mlf(-(x/tau)**alpha, alpha, alpha*gamma+1, gamma))
