import numpy as np
from scipy import stats

from scifit.interfaces.solvers import FitSolver1D


class ExponentialFitSolver(FitSolver1D):
    """
    `Exponential function model (calculus) <https://en.wikipedia.org/wiki/Exponential_function>`_
    """

    _model_equation = r"y = \beta_1 \exp(-\beta_2 x) + \beta_0"

    @staticmethod
    def model(x, b0, b1, b2):
        """
        Exponential function is defined as follows:

        .. math::

            y = \\beta_1 \\cdot \\exp \\left( - \\beta_2 \\cdot x \\right) + \\beta_0

        :param x: :math:`x`
        :param b0:
        :param b1:
        :param b2:
        :return: Gompertz function :math:`y`
        """
        return b1 * np.exp(-b2 * x[:, 0]) + b0


class GompertzFitSolver(FitSolver1D):
    """
    `Gompertz function model (calculus) <https://en.wikipedia.org/wiki/Gompertz_function>`_
    """

    _model_equation = (
        r"y = a \cdot \exp\left( - b \cdot \exp \left( - c \cdot x\right) \right)"
    )

    @staticmethod
    def model(x, a, b, c):
        """
        Gompertz function is defined as follows:

        .. math::

            y = a \\cdot \\exp\\left( - b \\cdot \\exp \\left( - c \\cdot x\\right) \\right)

        :param x: :math:`x`
        :param a: asymptotic value :math:`a`
        :param b: displacement :math:`b`
        :param c: growth rate :math:`c`
        :return: Gompertz function :math:`y`
        """
        return a * np.exp(-b * np.exp(-c * x[:, 0]))


class LogisticFitSolver(FitSolver1D):
    """
    `Logistic function model (calculus) <https://en.wikipedia.org/wiki/Logit>`_
    """

    _model_equation = r"y = \frac{1}{1 + \exp\left[- k \cdot (x - x_0)\right]}"

    @staticmethod
    def model(x, k, x0):
        """
        Logistic function is defined as follows:

        .. math::

            y = \\frac{1}{1 + \\exp\\left[- k \\cdot (x - x_0)\\right]}

        :param x: independent variable :math:`x`
        :param k: growth rate or sigmoid steepness :math:`k`
        :param x0: displacement or sigmoid inflexion point :math:`x_0`
        :return: Logistic function :math:`y`
        """
        return 1.0 / (1.0 + np.exp(-k * (x[:, 0] - x0)))


class AlgebraicSigmoidFitSolver(FitSolver1D):
    """
    `Algebraic sigmoid function model (calculus) <https://en.wikipedia.org/wiki/Algebraic_function>`_
    """

    _model_equation = r"y = \frac{x}{\left(1 + |x|^k\right)^{\frac{1}{k}}}"

    @staticmethod
    def model(x, k):
        """
        Logistic function is defined as follows:

        .. math::

            y = \\frac{x}{\\left(1 + |x|^k\\right)^{\\frac{1}{k}}}

        :param x: independent variable :math:`x`
        :param k: growth rate or sigmoid steepness :math:`k`
        :return: Algebraic sigmoid function :math:`y`
        """
        return x[:, 0] / np.power(1.0 + np.power(np.abs(x[:, 0]), k), 1.0 / k)


class RichardGeneralizedSigmoidFitSolver(FitSolver1D):
    """
    `Richard's generalized sigmoid function model (calculus) <https://en.wikipedia.org/wiki/Generalised_logistic_function>`_
    """

    _model_equation = (
        r"y = A + \frac{K - A}{\left(C + Q\cdot\exp(-B\cdot t)\right)^{\frac{1}{\nu}}}"
    )

    @staticmethod
    def model(x, A, B, C, K, Q, nu):
        """
        Richard's generalized sigmoid (GRS) function is defined as follows:

        .. math::

            y = A + \\frac{K - A}{\\left(C + Q\\cdot\\exp(-B\\cdot t)\\right)^{\\frac{1}{\\nu}}}

        :param A: lower (left) asymptote :math:`A`
        :param B: growth rate or sigmoid steepness :math:`B`
        :param C: asymptotical parameter :math:`C`
        :param K: upper (right) asymptote :math:`K`
        :param Q: location parameter :math:`Q`
        :param nu: asymptotical growth rate :math:`\\nu`
        :return: GRS function :math:`y`
        """
        return A + (K - A) / np.power(C + Q * np.exp(-B * x[:, 0]), 1.0 / nu)


class SmoothstepSigmoidFitSolver(FitSolver1D):
    """
    `Smoothstep sigmoid function model (calculus) <https://en.wikipedia.org/wiki/Smoothstep>`_
    """

    @staticmethod
    def model(x, a, b):
        """
        Smoothstep function is defined as follows:

        .. math::

            y =
            \\begin{cases}
            0,           & x \\le 0 \\\\
            a \\cdot x^2 - b \\cdot x^3, & 0 \\le x \\le 1 \\\\
            1,           & x \\ge 1 \\\\
            \\end{cases}

        :param a: quadratic coefficient :math:`a`
        :param b: cubic coefficient :math:`b`
        :return: Smoothstep sigmoid function :math:`y`
        """
        y = a * np.power(x[:, 0], 2) - b * np.power(x[:, 0], 3)
        y[x[:, 0] <= 0.0] = 0.0
        y[x[:, 0] >= 1.0] = 1.0
        return y


class InverseBoxCoxFitSolver(FitSolver1D):
    """
    `Inverse Box-Cox model (calculus) <https://en.wikipedia.org/wiki/Power_transform#Box%E2%80%93Cox_transformation>`_
    """

    @staticmethod
    def model(x, lambda_):
        """
        Inverse Box-Cox function is defined as follows:

        .. math::

            \\varphi(x, \\lambda) = y =
            \\begin{cases}
            (1 - \\lambda x)^\\frac{1}{\\lambda},  & \\lambda \\ne 0 \\\\
            \\exp(-x),           & \\lambda \\ge 0 \\\\
            \\end{cases}

        :param lambda\_: Box-Cox parameter :math:`\\lambda`
        :return: Inverse Box-Cox function :math:`y`
        """
        if np.allclose(lambda_, 0.0):
            return np.exp(-x[:, 0])
        else:
            return np.power(1.0 - lambda_ * x[:, 0], 1.0 / lambda_)


class DoubleInverseBoxCoxSigmoidFitSolver(FitSolver1D):
    """
    Double Inverse Box-Cox sigmoid model (calculus)
    """

    _model_equation = (
        r"y = \varphi(\varphi(x, \beta), \alpha), \quad \alpha < 1, \, \beta < 1"
    )

    @staticmethod
    def model(x, alpha, beta):
        """
        Inverse Box-Cox function is defined as follows:

        .. math::

            y = \\varphi(\\varphi(x, \\beta), \\alpha), \\quad \\alpha < 1, \\, \\beta < 1

        Where :math:`\\varphi` is the Inverse Box-Cox transformation, see :class:`InverseBoxCoxFitSolver` for details.

        :param alpha: first Box-Cox parameter :math:`\\alpha`
        :param beta: second Box-Cox parameter :math:`\\beta`
        :return: Double Inverse Box-Cox sigmoid function :math:`y`
        """
        return InverseBoxCoxFitSolver.model(
            InverseBoxCoxFitSolver.model(x, beta).reshape(-1, 1), alpha
        )


class MichaelisMentenKineticFitSolver(FitSolver1D):
    """
    `Michaëlis-Menten kinetic model (biochemistry) <https://en.wikipedia.org/wiki/Michaelis%E2%80%93Menten_kinetics>`_
    """

    # _model_equation = r"y = \frac{v_{\max} \cdot x}{K_{\mathrm{m}} + x}"

    @staticmethod
    def model(x, vmax, km):
        """
        Michaëlis-Menten kinetic is defined as follows:

        .. math::

            y = \\frac{v_\\max \\cdot x}{K_\\mathrm{m} + x}

        :param x: Substrate concentration :math:`x`
        :param vmax: Limiting kinetic rate :math:`v_\\max`
        :param km: Michaëlis constant :math:`K_\\mathrm{m}`
        :return: Michaëlis-Menten kinetic rate :math:`y`
        """
        return (vmax * x[:, 0]) / (km + x[:, 0])


class HillEquationFitSolver(FitSolver1D):
    """
    `Hill equation model (biochemistry) <https://en.wikipedia.org/wiki/Hill_equation_(biochemistry)>`_
    """

    _model_equation = (
        r"y = \frac{\beta_1 \cdot x^{\beta_0}}{1 + \beta_1 \cdot x^{\beta_0}}"
    )

    @staticmethod
    def model(x, n, k):
        """
        Hill equation is defined as follows:

        .. math::

            y = \\frac{k \\cdot x^n}{1 + k \\cdot x^n}

        :param x: total ligand concentration :math:`x`
        :param n: Hill coefficient :math:`n`
        :param k: apparent dissociation constant :math:`k` (law of mass action)
        :return: fraction :math:`y` of the receptor protein concentration that is bound by the ligand
        """
        term = k * np.power(x[:, 0], n)
        return term / (1 + term)


class GaussianPeakFitSolver(FitSolver1D):
    """
    `Gaussian function model (calculus) <https://en.wikipedia.org/wiki/Gaussian_function>`_
    """

    _model_equation = r"y = H \exp\left[- \frac{1}{2} \cdot \left(\frac{x - x_0}{\sigma}\right)^2\right]"

    @staticmethod
    def model(x, H, sigma, x0):
        """
        Gaussian function is defined as follows:

        .. math::

            y = H \\exp\\left[- \\frac{1}{2} \\cdot \\left(\\frac{x - x_0}{\\sigma}\\right)^2\\right]

        :param x: independent variable :math:`x`
        :param H: Peak height
        :param sigma: Peak width factor
        :param x0: Peak displacement :math:`x_0`
        :return: Gaussian peak :math:`y`
        """
        return H * np.exp(-0.5 * np.power((x[:, 0] - x0) / sigma, 2))


class GaussianPeakWithBaselineFitSolver(FitSolver1D):
    """
    Gaussian peak with baseline
    """

    _model_equation = r"y = H \exp\left[- \frac{1}{2} \cdot \left(\frac{x - x_0}{\sigma}\right)^2\right] + a x + b"

    @staticmethod
    def model(x, H, sigma, x0, a, b):
        """
        Gaussian function is defined as follows:

        .. math::

            y = H \\exp\\left[- \\frac{1}{2} \\cdot \\left(\\frac{x - x_0}{\\sigma}\\right)^2\\right] + a x + b

        :param x: independent variable :math:`x`
        :param H: Peak height
        :param sigma: Peak width factor
        :param x0: Peak displacement :math:`x_0`
        :param a:
        :param b:
        :return: Gaussian peak :math:`y`
        """
        return H * np.exp(-0.5 * np.power((x[:, 0] - x0) / sigma, 2)) + a * x[:, 0] + b


class EMGPeakFitSolver(FitSolver1D):
    """
    `Exponential Modified Gaussian Peak function model (calculus) <https://www.researchgate.net/publication/231172511_Equations_for_chromatographic_peak_modeling_and_calculation_of_peak_area>`_
    """

    _model_equation = r"y = \frac{A}{x_0} \exp\left[-\frac{1}{2}\left(\frac{\sigma_G}{x_0}\right)^2 -\frac{x - x_G}{x_0} \right] \cdot \frac{1}{\sqrt{2\pi}}\int\limits_{\infty}^{z} \exp\left[-\frac{y^2}{2}\right]\mathrm{d}y"
    cdf = np.vectorize(stats.norm(loc=0.0, scale=1.0).cdf)

    @staticmethod
    def model(x, A, sigma, x0, xG):
        """
        EMG Peak function is defined as follows:

        .. math::

            y = \\frac{A}{x_0} \\exp\\left[-\\frac{1}{2}\\left(\\frac{\\sigma_G}{x_0}\\right)^2 -\\frac{x - x_G}{x_0} \\right] \\cdot \\frac{1}{\\sqrt{2\pi}}\\int\\limits_{\\infty}^{z} \\exp\\left[-\\frac{y^2}{2}\\right]\\mathrm{d}y

        .. math::

            z = \\frac{x - x_G}{\\sigma_G} - \\frac{\\sigma_G}{x_0}

        :param x: independent variable :math:`x`
        :param A:
        :param sigma:
        :param x0:
        :param xG:
        :return: EMG peak :math:`y`
        """

        z = (x[:, 0] - xG) / sigma - sigma / x0
        return (
            A
            / x0
            * np.exp(-0.5 * np.power(sigma / x0, 2) - (x[:, 0] - xG) / x0)
            * EMGPeakFitSolver.cdf(z)
        )


class LaserPowerFitSolver(FitSolver1D):
    """
    Laser Power Model found on `Stack Overflow <https://stackoverflow.com/questions/77137301/curve-fit-seems-to-overestimate-error-of-estimated-parameters/77143084#77143084>`_
    """

    _model_equation = r"A \cdot \frac{x/s}{1 + k + x/s} + b"

    @staticmethod
    def model(x, A, s, b):
        """
        Laser Power Model

        .. math::

            A \\cdot \\frac{x/s}{1 + k + x/s} + b

        :param x:
        :param A:
        :param s:
        :param b:
        :return:
        """
        return A * (x[:, 0] / s) / (1 + (20.0 / 19.6) ** 2 + x[:, 0] / s) + b
