import numpy
from . import autophase


class DamarisFFT:
    """
    Class for Fourier transforming data.
    Provides several helper and apodization functions

    """

    def clip( self, start=None, stop=None ):
        """
        Method for clipping data, returns only the data between start and stop
        start and stop can be either time or frequency.
        The unit is automatically determined (Hz or s).

        :param float start: beginning of clipping in s
        :param float stop: end of clipping in s

        """
        # check if start/stop order is properly
        if start > stop:
            start, stop = stop, start

        # do nothing if one uses clip as a "placeholder"
        if start is None and stop is None:
            return self

        if start is None:
            start = self.x[ 0 ]
        if stop is None:
            stop = self.x[ -1 ]

        # check if data is fft which changes the start/stop units
        if self.xlabel == "Frequency / Hz":
            start = self.x.size * (0.5 + start / self.sampling_rate)
            stop = self.x.size * (0.5 + stop / self.sampling_rate)
        else:
            # get the corresponding indices
            start *= self.sampling_rate
            stop *= self.sampling_rate
        # check if boundaries make sense, raise exception otherwise
        if numpy.abs( int( start ) - int( stop ) ) <= 0:
            raise ValueError( "start stop too close: There are no samples in the given boundaries!" )

        # clip the data for each channel
        for ch in range( len( self.y ) ):
            self.y[ ch ] = self.y[ ch ][ int( start ):int( stop ) ]
        self.x = self.x[ int( start ):int( stop ) ]
        return self

    def baseline( self, last_part=0.1 ):
        """
        Correct the baseline of your data by subtracting the mean of the 
        last_part fraction of your data.

        :param float last_part: last section of your timesignal used to calculate baseline

        last_part defaults to 0.1, i.e. last 10% of your data
        """
        # TODO baseline correction for spectra after:
        # Heuer, A; Haeberlen, U.: J. Mag. Res.(1989) 85, Is 1, 79-94
        n = int( self.x.size * last_part )
        for ch in range( len( self.y ) ):
            self.y[ ch ] = self.y[ ch ] - self.y[ ch ][ -n: ].mean( )
        return self

    def exp_window( self, line_broadening=10 ):
        """
        Exponential window function

        :param float line_broadening: default 10, line broadening factor in Hz

        .. math::

            \\exp\\left(-\\pi\\cdot \\textsf{line_broadening} \\cdot t\\right)

        """
        apod = numpy.exp( -self.x * numpy.pi * line_broadening )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def gauss_window( self, line_broadening=10 ):
        """
        Gaussian window function

        :param float line_broadening: default 10, line broadening factor in Hz

        .. math:: \\exp\\left(- (\\textsf{line_broadening} \\cdot t)^2\\right)

        """
        apod = numpy.exp( -(self.x * line_broadening) ** 2 )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def dexp_window( self, line_broadening=-10, gaussian_multiplicator=0.3 ):
        apod = numpy.exp( -(self.x * line_broadening - gaussian_multiplicator * self.x.max( )) ** 2 )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def traf_window( self, line_broadening=10 ):
        apod = (numpy.exp( -self.x * line_broadening )) ** 2 / ( (numpy.exp( -self.x * line_broadening )) ** 3
                                                                 + (
        numpy.exp( -self.x.max( ) * line_broadening )) ** 3  )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def hanning_window( self ):
        """
        Symmetric centered window (hanning)
        """
        apod = numpy.hanning( self.x.size )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def hamming_window( self ):
        """
        Symmetric centered window (hamming)
        """

        apod = numpy.hamming( self.x.size )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def blackman_window( self ):
        """
        Symmetric centered window (blackmann)
        """

        apod = numpy.blackman( self.x.size )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def bartlett_window( self ):
        """
        Symmetric centered window (bartlett)
        """
        apod = numpy.bartlett( self.x.size )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def kaiser_window( self, beta=4, use_scipy=None ):
        """
        Symmetric centered window (kaiser)
        """
        apod = numpy.kaiser( self.x.size, beta )
        for i in range( 2 ):
            self.y[ i ] = self.y[ i ] * apod
        return self

    def autophase( self ):
        """
        Automatically phases the data to maximize real part.

        Works nice with a SNR above 20 dB, i.e.
        10 V signal to 0.1 V noise amplitude.
        """
        autophase.get_phase( self )
        return self

    def fft( self, samples=None ):
        """
        Calculate the Fourier transform of the data inplace.
        For zero filling set **samples** to a value higher than your data length,
        smaller values will truncate your data.

        :param int samples: default=None, if given, number of samples returned

        """
        fft_of_signal = numpy.fft.fft( self.y[ 0 ] + 1j * self.y[ 1 ], n=samples )
        fft_of_signal = numpy.fft.fftshift( fft_of_signal )
        dwell = 1.0 / self.sampling_rate
        n = fft_of_signal.size
        fft_frequencies = numpy.fft.fftfreq( n, dwell )
        self.x = numpy.fft.fftshift( fft_frequencies )
        self.y[ 0 ] = fft_of_signal.real
        self.y[ 1 ] = fft_of_signal.imag
        self.set_xlabel( "Frequency / Hz" )
        return self

    def magnitude( self ):
        """
        Return absolute signal, i.e.:

        .. math::

            y[0] &= \\sqrt{y[0]^2 + y[1]^2} \\\\
            y[1] &= 0


        """
        # this should calculate the absolute value, and set the imag channel to zero
        self.y[ 0 ] = numpy.sqrt( self.y[ 0 ] ** 2 + self.y[ 1 ] ** 2 )
        self.y[ 1 ] *= 0  # self.y[0].copy()
        return self

    def ppm(self, f_ref):
        """
        Return result scaled to PPM compared to f_ref

        :param f_ref: larmor frequency in MHz
        :return:
        """
        self.x /= f_ref
        self.set_xlabel( "PPM" )
        return self