import numpy as np

from ..lib.utils import ArrayLike

__all__ = ['Gaussian', 'Lorentzian', 'PseudoVoigt']


class Gaussian:
    name = 'Gaussian'
    equation = 'A*((4ln2/\pi)^{0.5}*(1/w) exp(-4ln2 [(x-\mu)/w]^{2})+A_{0}'
    params = ['A', r'\mu', 'w', 'A_{0}']
    type = 'Spectrum'
    bounds = [(0, None), (None, None), (0, None), (None, None)]

    @staticmethod
    def func(x: ArrayLike, a: float, mu: float, sigma: float, off: float) -> ArrayLike:
        r"""

        Calculates Gaussian lineshape

        .. math::
            f(x) = \frac{A}{\sqrt{2\pi} \sigma} \exp\left[-\frac{(x-\mu)^2}{2\sigma^2} \right] + A_0

        Args:
            x (float): Frequncy axis
            a (float): Amplitude
            mu (float): Maximum position
            sigma (float): FWHM
            off (float): Baseline

        """
        return a / (np.sqrt(4*np.log(2)/np.pi) * sigma) * np.exp(-4*np.log(2) * ((x-mu) / sigma)**2) + off


class Lorentzian:
    type = 'Spectrum'
    name = 'Lorentzian'
    equation = 'A (2/\pi)w/[4*(x-\mu)^{2} + w^{2}] + A_{0}'
    params = ['A', '\mu', 'w', 'A_{0}']
    ext_params = None
    bounds = [(0, None), (None, None), (0, None), (None, None)]

    @staticmethod
    def func(x: ArrayLike, a: float, mu: float, sigma: float, off: float) -> ArrayLike:
        r"""
        Calculate Lorentzian lineshape

        .. math::
            f(x) = \frac{A\sigma}{(x-\mu)^2 + \sigma^2} + A_0

        Args:
            x (float): Frequency axis
            a (float): Amplitude
            mu (float): Maximum position
            sigma (float): FWHM
            off (float): baseline

        """
        return (a/np.pi) * 2*sigma / (4*(x-mu)**2 + sigma**2) + off


class PseudoVoigt:
    type = 'Spectrum'
    name = 'Pseudo Voigt'
    equation = 'A [R*2/\pi*w/[4*(x-\mu)^{2} + w^{2}] + ' \
               '(1-R)*sqrt(4*ln(2)/pi)/w*exp(-4*ln(2)[(x-\mu)/w]^{2})] + A_{0}'
    params = ['A', 'R', '\mu', 'w', 'A_{0}']
    ext_params = None
    bounds = [(0, None), (0, 1), (None, None), (0, None)]

    @staticmethod
    def func(x: ArrayLike, a: float, r_lor: float, mu: float, sigma: float, off: float) -> ArrayLike:
        r"""
        Calculate Pseudo-Voigt lineshape, superposition of Lorentzian and Gaussian with shared center and FWHM.

        .. math::
                f(x) = A \left[ R \frac{2}{\pi} \frac{\sigma}{4(x-\mu)^2 + \sigma^2} +
                        (1-R)\sqrt{\frac{4\ln2}{\pi}} \frac{1}{\sigma} \exp\left(-4\ln2 \frac{(x-\mu)^2}{\sigma^2} \right) \right] + A_0

        Args:
            x (array-like): Frequency
            a (float): Amplitude
            r_lor (float): Relative Lorentzian contribution
            mu (float): Maximum position
            sigma (float): FWHM
            off (float): Baseline

        """
        return a * (r_lor * (2/np.pi) * (sigma/(4*(x-mu)**2 + sigma**2)) +
                    (1-r_lor) * np.sqrt(4*np.log(2)/np.pi)/sigma * np.exp(-4*np.log(2)*((x-mu)/sigma)**2)) + off
