import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import pandas as pd

def get_psd(data, fs, bin_width):
    f, psd = signal.welch(data, 
                          fs=fs, 
                          nperseg=fs/bin_width,
                          window='hann',
                          axis=0
                         )
 
    return f[1:], psd[1:] #drop the first value because it makes the plots look bad and is effectively 0


def setup_timeplot(figsize=(11, 6), plot_coeffs=False):
    global plot1, line1, plot2, line2, plot3, line3, plot4
    # time domain plot
    cols = 1
    rows = 4
    plt.figure(0, figsize=figsize)
    plt.clf()
    
    plot1 = plt.subplot2grid( (rows, cols), (0,0), 1)
    line1 = plot1.plot([0], label="mic")
    plt.legend()
    plot2 = plt.subplot2grid((rows,cols),(1,0), sharex=plot1)
    line2 = plot2.plot([0], label="acc")
    plt.legend()
    plot3 = plt.subplot2grid((rows, cols), (2,0), sharex=plot1)
    line3 = plot3.plot([0], label="out")
    plt.legend()
    
    if plot_coeffs:
        plot4 = plt.subplot2grid( (rows, cols), (3,0), sharex=plot1, rowspan=3)
        plt.legend(loc='upper right')
    
    plt.xlabel("sample")
    plt.tight_layout()
    plt.show()


def setup_freqplot(figsize=(11, 6)):
    global plot_freq, line_freq1, line_freq2, line_freq3
    cols = 1
    rows = 1
    # frequency domain plot
    plt.figure(1, figsize=figsize)
    plt.clf()
    plt.title("Power spectrum")
    plot_freq = plt.subplot2grid( (rows, cols), (0,0), 1)
    line_freq1 = plot_freq.loglog([0], label="mic")
    line_freq2 = plot_freq.loglog([0], label="acc")
    line_freq3 = plot_freq.loglog([0], label="out")
    plt.xlabel("frequency / Hz")
    plt.ylabel("dB²/Hz")
    plt.legend()
    plt.tight_layout()
    plt.show()


def plot_freqdomain(fs, ch1, ch2, output):

    x_min = []
    x_max = []
    y_min = []
    y_max = []

    f, psd = get_psd(ch1, fs, 10)
    x_min.append(min(f))
    x_max.append(max(f))
    y_min.append(min(psd))
    y_max.append(max(psd))
    line_freq1[0].set_data(f, psd)

    f, psd = get_psd(ch2, fs, 10)
    x_min.append(min(f))
    x_max.append(max(f))
    y_min.append(min(psd))
    y_max.append(max(psd))
    line_freq2[0].set_data(f, psd)

    f, psd = get_psd(output, fs, 10)
    x_min.append(min(f))
    x_max.append(max(f))
    y_min.append(min(psd))
    y_max.append(max(psd))
    line_freq3[0].set_data(f, psd)

    plot_freq.set_xlim(min(x_min), max(x_max))
    plot_freq.set_ylim(min(y_min), max(y_max))


def plot_timedomain(ch1, ch2, output, coeffs: pd.DataFrame =None):
    #Plot the result
    data= output
    line3[0].set_data(range(len(data)), data)
    plot3.set_xlim(0, len(data))
    plot3.set_ylim(np.min(data), np.max(data))

    # plot the chirp with noise
    data = ch1
    line1[0].set_data(range(len(data)), data)
    plot1.set_xlim(0, len(data))
    plot1.set_ylim(min(data), max(data))

    # Plot the noise
    data = ch2
    line2[0].set_data(range(len(data)), data)
    plot2.set_xlim(0, len(data))
    plot2.set_ylim(min(data), max(data))

    # Plot the coeffs
    if coeffs is not None:
        data = coeffs
        coeffs.plot(ax=plot4)
        #plot4.legend(bbox_to_anchor=(1.2, 1.05))
         
        #line4 = plot4.plot([0], label="coeffs")
        #line4[0].set_data(range(len(data)), data)
        
        #plot4.set_xlim(0, len(data))
        #plot4.set_ylim(min(data), max(data))