"""
======================
Apodization functions
======================
"""


import numpy as np
from scipy.signal.windows import general_hamming


class GaussWindow:
    r"""
    Gaussian function

    .. math::
        y = \exp\Biggl[-\Biggl(\frac{\pi x GB}{4\sqrt{\ln2}}\Biggr)^2\Biggr]

    """
    name = 'Gaussian'
    equation = r'exp(-(\pi GB x/[2*sqrt(ln2)])^{2})'
    params = ['GB/Hz']

    @staticmethod
    def apod(x, gb):
        r"""
        Gaussian function

        .. math::
            y = \exp\Biggl[-\Biggl(\frac{\pi x gb}{4\sqrt{\ln2}}\Biggr)^2\Biggr]

        Args:
            x (array-like): Input values
            gb (float): FWHM (in the frequency domain)
        """
        return np.exp(-(x*gb*np.pi/4./np.sqrt(np.log(2)))**2)


class ExpWindow:
    r"""
    Exponential function

    .. math::
        y = \exp(-\pi x LB)
    """
    name = 'Exponential'
    equation = r'exp(-\pi LB x)'
    params = ['LB/Hz']

    @staticmethod
    def apod(x, lb):
        r"""
        Exponential function

        .. math::
            y = \exp(-\pi x lb)

        Args:
            x (array-like): Input values
            lb (float): FWHM (in the frequency domain)
        """
        return np.exp(-np.pi*x*lb)


class GaussExpWindow:
    name = 'Exp/Gauss'
    equation = r'exp(\pi LB x) exp(-(\pi GB x/[2*sqrt(ln2)])^{2}))'
    params = ['LB/Hz', 'GB/Hz']

    @staticmethod
    def apod(x, gb, lb):
        return ExpWindow.apod(x, gb) * GaussWindow.apod(x, lb)


class Traficante:
    name = 'Traficante'
    equation = r'exp(-\pi LB x) / [exp(-\pi LB x)^{2} + exp(-\pi LB (x_{max}-x))^{2}]'
    params = ['LB/Hz']

    @staticmethod
    def apod(x, lb):
        e = ExpWindow.apod(x, lb)
        E = ExpWindow.apod(np.max(x)-x, lb)
        return e / (e**2 + E**2)


# class GaussExpShiftWindow:
#     name = 'Exp/Gauss with GB shift'
#     equation = r'exp(-\pi LB x) exp(-(\pi GB x(1-S)/[2*sqrt(ln2)])^{2})'
#     params = ['GB/Hz', 'LB/Hz', 'GB shift (0-1)']
#
#     @staticmethod
#     def apod(x, gb, lb, shift):
#         shift *= max(x)
#         return GaussWindow.apod(x-shift, gb) * ExpWindow.apod(x, lb)


class Trapezoid:
    name = 'Trapezoid'
    equation = 'x/(p_{1}x_{max}) - 1 - (x_{max}-x)/(x_{max}(1-p_{2}))'
    params = ['Left ramp (0-1)', 'Right ramp (0-1)']

    @staticmethod
    def apod(x, lramp, rramp):
        size = len(x)
        lramp = min(lramp, 1)
        rramp = min(rramp, 1)
        window = np.r_[np.linspace(0, 1, int(lramp*size)),
                       np.ones(int((rramp-lramp)*size)),
                       np.linspace(1, 0, int((1-rramp)*size))]
        if len(window) < size:
            window = np.r_[window, np.zeros(size-len(window))]
        elif len(window) > size:
            window = window[:size]

        return window


class SineBell:
    name = 'Sine-bell'
    equation = r'sin(\phi + (\pi-\phi) x/x_{max})^{n}'
    params = [r'\phi/deg.', 'n']

    @staticmethod
    def apod(x, phi, pow):
        phi_rad = np.pi*phi/180.
        return np.sin((np.pi - phi_rad)*x/np.max(x) + phi_rad)**pow


class Box:
    name = 'Box'
    equation = 'y[p*x_{max}:x_{max}] = 0'
    params = ['rel. position']

    @staticmethod
    def apod(x, step):
        one = np.ones_like(x)
        one[int(step*len(x)):] = 0.0
        return one


class Hamming:
    name = 'Hamming'
    equation = r'0.54 - 0.46cos(2\pi x/x_{max})'
    params = []

    @staticmethod
    def apod(x):
        return general_hamming(len(x), 0.54)


class Hann:
    name = 'Hann'
    equation = r'0.5 - 0.5cos(2\pi x/x_{max})'
    params = []

    @staticmethod
    def apod(x):
        return general_hamming(len(x), 0.5)




