import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate
from scipy.signal import find_peaks, resample, detrend
import csv

"""
NOTES

When you compute the Fourier transform of a signal, you obtain a set of complex-valued coefficients corresponding to different frequency components present in the signal. These coefficients represent the amplitude and phase of sinusoidal components at specific frequencies.
    
The frequency range covered by the Fourier transform output is divided into discrete frequency bins, each representing a specific frequency component. The width of these bins depends on the sampling rate and the length of the input signal. In a typical implementation, the frequency bins are evenly spaced.

In practical terms, the power spectrum bins correspond to the frequency components at which the power spectral density (PSD) or magnitude squared of the Fourier coefficients are evaluated. These bins are used to represent the distribution of signal power across different frequency components, providing insights into the frequency content of the signal.
    
The bin width, also called the resolution bandwidth, is simply sampling rate / Total number of samples = (1/dt)/N

N can be re-written in terms of t, dt. Putting this in to the expression for Sxx = 2 * |fft|^2/(Resolution Bandwidth * Noise Power Bandwidth) We get the expression used here. Note that the fft computed above is scaled by N which results eventually in the factor dt^2/T.

The noise power bandwidth is typically 1 Hz if no windowing/tapering function is used.


Compute the broadband noise level in Vrms2/Hz by summing all the power spectrum bins, excluding any peaks and the DC component, and dividing the sum by the equivalent noise bandwidth of the window

The equivalent noise bandwidth (ENBW) of a window in the context of a power spectrum refers to a measure of the effective bandwidth of the window function applied to the signal before taking the Fourier transform.

When you compute the power spectrum of a signal using a windowing function (e.g., Hamming window, Hann window, etc.), the window modifies the original signal by tapering its edges. This tapering reduces the spectral leakage and improves frequency resolution but also introduces a smoothing effect, which can affect the estimation of the signal's power at different frequencies.

The equivalent noise bandwidth provides a way to quantify the effective bandwidth of the window function in terms of its impact on noise power. It represents the width of a rectangular filter that would have the same noise power as the windowed signal.

In practical terms, when calculating the power spectrum of a signal using a window, the ENBW is used to adjust the power spectrum to account for the smoothing effect of the window. Dividing the sum of the power spectrum bins by the ENBW yields an estimate of the noise power per unit frequency bandwidth.

ENBW is often used in the context of noise measurements or signal processing applications where accurate estimation of noise power is important. It helps ensure that the power spectrum accurately reflects the true power distribution of the signal, accounting for the effects of windowing.


The noise floor refers to the minimum level of signal that can be reliably distinguished from the background noise. It represents the lowest amplitude of a signal that can be detected or measured with reasonable accuracy.

The noise floor is often defined as the RMS (Root Mean Square) value of the background noise in a given frequency band.

To compute the noise floor value from the power spectral density (PSD) values, you typically need to analyze the portion of the PSD that corresponds to the background noise.

"""

def compute_autocorrelation(data):
    
    print("Calculating autocorrelation...")
    yunbiased = data-np.mean(data)
    ynorm = np.sum(yunbiased**2)
    acor = np.correlate(yunbiased, yunbiased, "same")/ynorm
    return acor

def pre_process_data(data, re_sample = False, new_sampling_rate = 10000):
    # Resample the time series
    # Define the new sampling rate and calculate the new time values
    if re_sample:
        n_points = int((len(data[:, 0]) * (data[:, 0][1] - data[:, 0][0])) * new_sampling_rate)
        t        = np.linspace(data[:, 0][0], data[:, 0][-1], n_points)
        x_array  = resample(data[:, 1], n_points)
        x        = detrend(x_array - x_array.mean())
    else:
        t        = data[:, 0]
        x_array  = data[:, 1]
        x        = detrend(x_array - x_array.mean())

    return np.column_stack((t, x))

def smoothen_data(data, window_size):
    # Smooth data by doing a moving average
    return np.convolve(data, np.ones(window_size, dtype=int)/window_size, mode='valid')

def compute_psd(data, re_sample = False, new_sampling_rate = 10000, window_size = 21):
    """
    A power spectral density (PSD) takes the amplitude of the FFT, multiplies it by its complex conjugate and normalizes it to the frequency bin width.

    """

    processed_data = pre_process_data(data, re_sample, new_sampling_rate)
    t, x           = processed_data[:, 0], processed_data[:, 1]

    N  = t.shape[0]                                         # Total number of data points
    
    # Calculate fft
    print("Calculating power spectrum...")
    fft_ts = np.fft.fft(x)                                 # Compute Fourier transform of x                       
    Sxx = (fft_ts * fft_ts.conj()) / N**2                  # Compute spectrum
    Sxx = 2 * Sxx[:int(len(x) / 2)]                        # Ignore negative frequencies, we have accounted for this by the scaling factor of 2

    return processed_data, smoothen_data(Sxx.real, window_size)

def compute_RIN(dc_voltages, delta_f, Sxx_smooth):
    
    # Compute the average power
    average_P = np.mean(np.square(dc_voltages))

    # Calculate the RIN
    Skk                = Sxx_smooth / (average_P * delta_f)
    RIN_Sxx_smooth     = 10 * np.log10(Skk)
    
    return Skk, RIN_Sxx_smooth

def find_noise_peaks(psd, faxis, freq_range, threshold):
    """
    Compute the peak power in the specified frequency range.

    Parameters:
        psd_values: array-like
            Power spectral density values.
        faxis: array-like
            Frequencies corresponding to the PSD values.
        freq_range: tuple
            Tuple containing the start and end frequencies of the range of interest.
        threshold: scalar
            Threshold for peak heights

    Returns:
        float: Peak power in the specified frequency range.
    """
    start_freq, end_freq = freq_range
    idx_start            = np.argmax(faxis >= start_freq)
    idx_end              = np.argmax(faxis >= end_freq)
    sliced_psd           = psd[idx_start:idx_end]
    sliced_faxis         = faxis[idx_start:idx_end]

    peak_indices, _  = find_peaks(sliced_psd, height=threshold)
    peak_powers      = 10 * np.log10(sliced_psd[peak_indices])
    peak_frequencies = np.around(sliced_faxis[peak_indices], 2)

    return peak_powers, peak_frequencies

def compute_noise_level(psd, resolution_bandwidth, exclude_peaks=False, faxis=None, freq_range=None, threshold=None):
    """
    Compute the noise level from a power spectral density (PSD).

    Parameters:
        psd: array-like
            One-sided power spectral density.
        resolution_bandwidth: float
            Bin width
    Returns:
        float: Noise level (Vrms^2).
    """

    noise_level = None
    # Exclude peaks from the sum
    if exclude_peaks and threshold is not None:
        
        threshold                = 10**(threshold/10)

        if freq_range is None:
            peak_indices, _      = find_peaks(psd, height=threshold)
            noise_level          = resolution_bandwidth * np.sum([psd[i] for i in range(len(psd)) if i not in peak_indices])
        else:
            start_freq, end_freq = freq_range
            idx_start            = np.argmax(faxis >= start_freq)
            idx_end              = np.argmax(faxis >= end_freq)
            sliced_psd           = psd[idx_start:idx_end]
            peak_indices, _      = find_peaks(sliced_psd, height=threshold)
            noise_level          = resolution_bandwidth * np.sum([sliced_psd[i] for i in range(len(sliced_psd)) if i not in peak_indices])
    else:
        
        if freq_range is None:
            noise_level          = resolution_bandwidth * np.sum([psd[i] for i in range(len(psd))])
        else:
            start_freq, end_freq = freq_range
            idx_start            = np.argmax(faxis >= start_freq)
            idx_end              = np.argmax(faxis >= end_freq)
            sliced_psd           = psd[idx_start:idx_end]
            noise_level          = resolution_bandwidth * np.sum([sliced_psd[i] for i in range(len(sliced_psd))])    
    
    return noise_level
    
def extract_data(filepath):
    
    # Open the CSV file
    with open(filepath, newline='') as csvfile:
        # Skip the first line (header)
        next(csvfile)
        
        # Read the CSV file using csv.reader
        reader = csv.reader(csvfile)

        # Read the headers from the second line
        next_header = next(reader)
        string_number = next_header[-1]
        
        try:
            time_step = int(string_number)
        except ValueError:
            try:
                time_step = float(string_number)
            except ValueError:
                print("The string does not represent a valid number.")
        
        # Initialize lists to store the first and second values
        first_column = []
        second_column = []
        
        # Iterate over each row in the CSV file
        for row in reader:
            # Extract the first and second values from the row and convert to float
            first_value = float(row[0])
            second_value = float(row[1])
            
            # Append the values to their respective lists
            first_column.append(first_value)
            second_column.append(second_value)

        # Convert the lists into numpy arrays
        time_array    = np.arange(0, len(first_column)*time_step, time_step)
        voltage_array = np.array(second_column)

        # Stack the arrays horizontally to form a single 2D array
        data_array = np.column_stack((time_array, voltage_array))

        return data_array

def closest(list, num):
    idx = int(min(range(len(list)), key=lambda i: abs(list[i]-num)))
    return idx

def plot_psd(data, Sxx, data_str, peak_find = False, peak_find_threshold = -100, window_size = 21):

    if isinstance(data, list) and isinstance(Sxx, list) and isinstance(data_str, list):
        plt.figure(figsize=(12, 8))

        for idx in range(len(data)):

            dat     = data[idx]
            psd     = Sxx[idx]

            time    = dat[:, 0]
            dt      = time[1] - time[0]                                    # Sampling interval / time step
            N       = time.shape[0]                                        # Total number of data points
            f_s     = 1 / dt                                               # Sampling frequency
            fNQ     = f_s / 2                                              # Determine Nyquist frequency
            faxis   = smoothen_data(np.linspace(0,fNQ,N//2), window_size)  # Construct frequency axis
            
            """ Noise peak finder"""
            if peak_find:
                freq_range                    = (50, max(faxis))
                threshold                     = 10**(peak_find_threshold/10)
                peak_powers, peak_frequencies = find_noise_peaks(psd, faxis, freq_range, threshold)

            """ Noise levels in units of V/sqrt(Hz)"""
            idx_ini       = closest(faxis, 100)
            idx_fin       = closest(faxis, 1000)
            noiseat1hunHz = np.sqrt(psd[idx_ini])
            noiseat1kHz   = np.sqrt(psd[idx_fin])
            cum_noise     = np.sqrt(integrate.cumtrapz(psd[idx_ini:idx_fin], initial=0))
            print('Noise level @ 100 Hz = {}'.format(noiseat1hunHz))
            print('Noise level @ 1  kHz = {}'.format(noiseat1kHz))
            print('Integrated Noise from 100 Hz to 1 kHz = {}'.format(cum_noise[-1]))

            # Plot Power Spectrum in dB
            plt.semilogx(faxis, 10 * np.log10(psd), linewidth=2, label = data_str[idx])
            
            if peak_find:
                plt.plot(peak_frequencies, peak_powers, 'o', markerfacecolor='none', markeredgecolor='r', markersize=10) # Plot power against frequency as hollow circles
                for freq, power in zip(peak_frequencies, peak_powers):
                    plt.text(freq, power, str(freq)+' Hz', verticalalignment='bottom', horizontalalignment='right') # Add text next to each circle indicating the frequency

        plt.grid(True, which="both", linestyle='-', linewidth=0.5, color='gray')  # Thin lines for non-decade grid
        plt.grid(True, which="both", linestyle=':', linewidth=1, color='gray', axis='x')  # Thick lines for decade grid
        
        # Calculate the x-axis values for multiples of 10
        x_multiples_of_10 = [10**i for i in range(int(np.log10(min(faxis[faxis > 0]))), int(np.log10(max(faxis[faxis > 0]))) + 1)]
        # Add thick lines for multiples of 10
        for val in x_multiples_of_10:
            plt.axvline(x=val, color='black', linestyle='-', linewidth=2)  # Thick lines for multiples of 10
        
        plt.xlim([min(faxis), max(faxis)])
        # plt.ylim([-150, -80])
        plt.legend(loc = 3, fontsize=12)
        plt.xlabel('Frequency [Hz]', fontsize=14)
        plt.ylabel('PSD [dB/Hz]', fontsize=14)

        # Adjust layout
        plt.tight_layout()
    else:   
        time    = data[:, 0]
        dt      = time[1] - time[0]                                    # Sampling interval / time step
        N       = time.shape[0]                                        # Total number of data points
        f_s     = 1 / dt                                               # Sampling frequency
        fNQ     = f_s / 2                                              # Determine Nyquist frequency
        faxis   = smoothen_data(np.linspace(0,fNQ,N//2), window_size)  # Construct frequency axis
        
        """ Noise peak finder"""
        if peak_find:
            freq_range                    = (50, max(faxis))
            threshold                     = 10**(peak_find_threshold/10)
            peak_powers, peak_frequencies = find_noise_peaks(Sxx, faxis, freq_range, threshold)
        
        """ Noise levels in units of V/sqrt(Hz)"""
        idx_ini       = closest(faxis, 100)
        idx_fin       = closest(faxis, 1000)
        noiseat1hunHz = np.sqrt(Sxx[idx_ini])
        noiseat1kHz   = np.sqrt(Sxx[idx_fin])
        cum_noise     = np.sqrt(integrate.cumtrapz(Sxx[idx_ini:idx_fin], initial=0))
        print('Noise level @ 100 Hz = {}'.format(noiseat1hunHz))
        print('Noise level @ 1  kHz = {}'.format(noiseat1kHz))
        print('Integrated Noise from 100 Hz to 1 kHz = {}'.format(cum_noise[-1]))
        
        plt.figure(figsize=(12, 8))
            
        # Plot Power Spectrum in dB
        plt.semilogx(faxis, 10 * np.log10(Sxx), color='green', linewidth=2, label = data_str)
        
        if peak_find:
            plt.plot(peak_frequencies, peak_powers, 'o', markerfacecolor='none', markeredgecolor='r', markersize=10) # Plot power against frequency as hollow circles
            for freq, power in zip(peak_frequencies, peak_powers):
                plt.text(freq, power, str(freq)+' Hz', verticalalignment='bottom', horizontalalignment='right') # Add text next to each circle indicating the frequency

        plt.grid(True, which="both", linestyle='-', linewidth=0.5, color='gray')  # Thin lines for non-decade grid
        plt.grid(True, which="both", linestyle=':', linewidth=1, color='gray', axis='x')  # Thick lines for decade grid
        
        # Calculate the x-axis values for multiples of 10
        x_multiples_of_10 = [10**i for i in range(int(np.log10(min(faxis[faxis > 0]))), int(np.log10(max(faxis[faxis > 0]))) + 1)]
        # Add thick lines for multiples of 10
        for val in x_multiples_of_10:
            plt.axvline(x=val, color='black', linestyle='-', linewidth=2)  # Thick lines for multiples of 10
        
        plt.xlim([min(faxis), max(faxis)])
        # plt.ylim([-50, 0])
        plt.legend(loc = 3, fontsize=12)
        plt.xlabel('Frequency [Hz]', fontsize=14)
        plt.ylabel('PSD [dB/Hz]', fontsize=14)

        # Adjust layout
        plt.tight_layout()

        # Show plot
        plt.show()  

def plot_rin(data, rin, data_str, peak_find = False, peak_find_threshold = -100, window_size = 21):

    if isinstance(data, list) and isinstance(rin, list) and isinstance(data_str, list):
        plt.figure(figsize=(12, 8))

        for idx in range(len(data)):

            dat     = data[idx]
            rn      = rin[idx]
            
            time    = dat[:, 0]
            dt      = time[1] - time[0]                                    # Sampling interval / time step
            N       = time.shape[0]                                        # Total number of data points
            f_s     = 1 / dt                                               # Sampling frequency
            fNQ     = f_s / 2                                              # Determine Nyquist frequency
            faxis   = smoothen_data(np.linspace(0,fNQ,N//2), window_size)  # Construct frequency axis
            
            """ Noise peak finder"""
            if peak_find:
                freq_range                    = (50, max(faxis))
                threshold                     = 10**(peak_find_threshold/10)
                peak_powers, peak_frequencies = find_noise_peaks(rn, faxis, freq_range, threshold)

            # Plot RIN
            plt.semilogx(faxis, rn, linewidth=2, label = data_str[idx])
            
            if peak_find:
                plt.plot(peak_frequencies, peak_powers, 'o', markerfacecolor='none', markeredgecolor='r', markersize=10) # Plot power against frequency as hollow circles
                for freq, power in zip(peak_frequencies, peak_powers):
                    plt.text(freq, power, str(freq)+' Hz', verticalalignment='bottom', horizontalalignment='right') # Add text next to each circle indicating the frequency

        plt.grid(True, which="both", linestyle='-', linewidth=0.5, color='gray')  # Thin lines for non-decade grid
        plt.grid(True, which="both", linestyle=':', linewidth=1, color='gray', axis='x')  # Thick lines for decade grid
        
        # Calculate the x-axis values for multiples of 10
        x_multiples_of_10 = [10**i for i in range(int(np.log10(min(faxis[faxis > 0]))), int(np.log10(max(faxis[faxis > 0]))) + 1)]
        # Add thick lines for multiples of 10
        for val in x_multiples_of_10:
            plt.axvline(x=val, color='black', linestyle='-', linewidth=2)  # Thick lines for multiples of 10
        
        plt.xlim([min(faxis), max(faxis)])
        # plt.ylim([-150, -80])
        plt.legend(loc = 3, fontsize=12)
        plt.xlabel('Frequency [Hz]', fontsize=14)
        plt.ylabel('RIN [dB/Hz]', fontsize=14)

        # Adjust layout
        plt.tight_layout()
    else:   
        time    = data[:, 0]

        dt      = time[1] - time[0]                                    # Sampling interval / time step
        N       = time.shape[0]                                        # Total number of data points
        f_s     = 1 / dt                                               # Sampling frequency
        fNQ     = f_s / 2                                              # Determine Nyquist frequency
        faxis   = smoothen_data(np.linspace(0,fNQ,N//2), window_size)  # Construct frequency axis
        
        """ Noise peak finder"""
        if peak_find:
            freq_range                    = (50, max(faxis))
            threshold                     = 10**(peak_find_threshold/10)
            peak_powers, peak_frequencies = find_noise_peaks(rin, faxis, freq_range, threshold)
        
        plt.figure(figsize=(12, 8))
            
        # Plot Power Spectrum in dB
        plt.semilogx(faxis, rin, color='green', linewidth=2, label = data_str)
        
        if peak_find:
            plt.plot(peak_frequencies, peak_powers, 'o', markerfacecolor='none', markeredgecolor='r', markersize=10) # Plot power against frequency as hollow circles
            for freq, power in zip(peak_frequencies, peak_powers):
                plt.text(freq, power, str(freq)+' Hz', verticalalignment='bottom', horizontalalignment='right') # Add text next to each circle indicating the frequency

        plt.grid(True, which="both", linestyle='-', linewidth=0.5, color='gray')  # Thin lines for non-decade grid
        plt.grid(True, which="both", linestyle=':', linewidth=1, color='gray', axis='x')  # Thick lines for decade grid
        
        # Calculate the x-axis values for multiples of 10
        x_multiples_of_10 = [10**i for i in range(int(np.log10(min(faxis[faxis > 0]))), int(np.log10(max(faxis[faxis > 0]))) + 1)]
        # Add thick lines for multiples of 10
        for val in x_multiples_of_10:
            plt.axvline(x=val, color='black', linestyle='-', linewidth=2)  # Thick lines for multiples of 10
        
        plt.xlim([min(faxis), max(faxis)])
        # plt.ylim([-50, 0])
        plt.legend(loc = 3, fontsize=12)
        plt.xlabel('Frequency [Hz]', fontsize=14)
        plt.ylabel('RIN [dB/Hz]', fontsize=14)
        
        # Adjust layout
        plt.tight_layout()

        # Show plot
        plt.show()

def plot_e_folding_time(data, Skk, fstart, fend, data_str, window_size = 21):

    
    time, voltages          = data[:, 0], data[:, 1]

    dt = time[1] - time[0]           # Define the sampling interval
    N  = voltages.shape[0]           # Define the total number of data points
    T  = N * dt                      # Define the total duration of the data
    df = 1 / T.max()

    fNQ = 1 / dt / 2                                             # Determine Nyquist frequency
    faxis = smoothen_data(np.linspace(0,fNQ,N//2), window_size)  # Construct frequency axis

    idx_ini        = closest(faxis, 2*fstart)
    idx_fin        = closest(faxis, 2*fend)
    Skk            = Skk[idx_ini:idx_fin]
    freqs          = np.linspace(fstart, fend, len(Skk))
    e_folding_time = [1/(np.pi**2 * freqs[idx]**2 * Skk[idx]) for idx in range(len(freqs))]

    plt.figure(figsize=(12, 8))
                
    # Plot Power Spectrum in dB
    plt.semilogx(freqs, e_folding_time, color='red', linewidth=2, alpha=0.75, label = data_str)

    plt.grid(True, which="both", linestyle='-', linewidth=0.5, color='gray')  # Thin lines for non-decade grid
    plt.grid(True, which="both", linestyle=':', linewidth=1, color='gray', axis='x')  # Thick lines for decade grid

    # Calculate the x-axis values for multiples of 10
    x_multiples_of_10 = [10**i for i in range(int(np.log10(min(freqs[freqs > 0]))), int(np.log10(max(freqs[freqs > 0]))) + 1)]
    # Add thick lines for multiples of 10
    for val in x_multiples_of_10:
        plt.axvline(x=val, color='black', linestyle='-', linewidth=2)  # Thick lines for multiples of 10

    # plt.xlim([min(freqs), max(freqs)])
    # plt.ylim([-50, 0])
    plt.legend(loc = 3, fontsize=12)
    plt.xlabel('Frequency [Hz]', fontsize=14)
    plt.ylabel('$T_e$ [s]', fontsize=14)
    plt.title('Lower Bound= %.5f s' % (min(e_folding_time)), fontsize=14)
    # Adjust layout
    plt.tight_layout()

    # Show plot
    plt.show()