import numpy as np
import scipy
from sklearn.base import BaseEstimator

# Building a class Fourier for better use of Fourier Analysis.

class Fourier:
    """
    Apply the Discrete Fourier Transform (DFT) on the signal using the Fast Fourier
    Transform (FFT) from the scipy package.

    Example:
      fourier = Fourier(signal, sampling_rate=2000.0)
    """

    def __init__(self, signal, sampling_rate):
        """
        Initialize the Fourier class.

        Args:
            signal (np.ndarray): The samples of the signal
            sampling_rate (float): The sampling per second of the signal

        Additional parameters,which are required to generate Fourier calculations, are
        calculated and defined to be initialized here too:
            time_step (float): 1.0/sampling_rate
            time_axis (np.ndarray): Generate the time axis from the duration and
                                  the time_step of the signal. The time axis is
                                  for better representation of the signal.
            duration (float): The duration of the signal in seconds.
            frequencies (numpy.ndarray): The frequency axis to generate the spectrum.
            fourier (numpy.ndarray): The DFT using rfft from the scipy package.
        """
        self.signal = signal
        self.sampling_rate = sampling_rate
        self.time_step = 1.0 / self.sampling_rate
        self.duration = len(self.signal) / self.sampling_rate
        self.time_axis = np.arange(0, self.duration, self.time_step)
        self.frequencies = scipy.fft.rfftfreq(len(self.signal), d=self.time_step)
        self.fourier = scipy.fft.rfft(self.signal)

    # Generate the actual amplitudes of the spectrum
    def amplitude(self):
        """
        Method of Fourier

        Returns:
            numpy.ndarray of the actual amplitudes of the sinusoids.
        """
        return 2 * np.abs(self.fourier) / len(self.signal)

    # Generate the phase information from the output of rfft
    def phase(self, degree=False):
        """
        Method of Fourier

        Args:
            degree: To choose the type of phase representation (Radian, Degree).
                    By default, it's in radian.

        Returns:
            numpy.ndarray of the phase information of the Fourier output.
        """
        return np.angle(self.fourier, deg=degree)



# define the transformer
class FourierPreprocessor(BaseEstimator):
    def __init__(self, sampling_rate=12000, len_sig=36000):
        print("Initialising transformer...")
        self.sampling_rate = sampling_rate
        self.len_sig = len_sig

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        transformed_X = []
        for sig in X:
            try:
                if len(sig) != self.len_sig:
                    sig = scipy.signal.resample(sig, self.len_sig)
                # Convert signal to frequency domain
                fourier = np.array(
                    Fourier(signal=sig, sampling_rate=self.sampling_rate).amplitude()
                )
            except Exception as e:
                print(e)
                fourier = np.zeros(shape=(18001,))
            transformed_X.append(fourier)
        return np.array(transformed_X)