import socket
import matplotlib.pyplot as plt
import numpy as np
from scipy import signal


MCLK=3276800
OSR=4096
PRESCALE=1

DRCLK=MCLK/(4*OSR*PRESCALE)

fs = DRCLK  # Sample frequency (Hz)
f0 = 60.0  # Frequency to be removed from signal (Hz)
Q = 30.0  # Quality factor
sos = signal.iirfilter(30, [55, 65], btype='bandstop', fs=fs, output='sos')
# Design notch filter
#b, a = signal.iirnotch(f0, Q, fs)

Vref=1.2
G=51*2
scale=Vref/(8388608*G*1.5)

plt.style.use('ggplot')

def live_plotter(x_vec,y1_data,line1,identifier='',pause_time=0.001):
    if line1==[]:
        # this is the call to matplotlib that allows dynamic plotting
        plt.ion()
        fig = plt.figure(figsize=(13,6))
        ax = fig.add_subplot(111)
        # create a variable for the line so we can later update it
        line1, = ax.plot(x_vec,y1_data,'-',alpha=0.8)        
        #update plot label/title
        plt.ylabel('Y Label')
        plt.title('Title: {}'.format(identifier))
        plt.show()
    
    # after the figure, axis, and line are created, we only need to update the y-data
    line1.set_ydata(y1_data)
    # adjust limits if new data goes beyond bounds
    if np.min(y1_data)<=line1.axes.get_ylim()[0] or np.max(y1_data)>=line1.axes.get_ylim()[1]:
        plt.ylim([np.min(y1_data)-np.std(y1_data),np.max(y1_data)+np.std(y1_data)])
    # this pauses the data so the figure/axis can catch up - the amount of pause can be altered above
    plt.pause(pause_time)
    
    # return line so we can update it again in the next iteration
    return line1


sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
server_address = ('', 10000)
sock.bind(server_address)
sock.setblocking(1)


length = 5
x_vec = np.linspace(0,length,int(length*fs)+1)[0:-1]
y_vec = np.random.randn(len(x_vec))
line1 = []

a=np.zeros(20,dtype='int64')
while True:
    data, address = sock.recvfrom(4096)
    i=0
    while i<20:
        a[i] = int.from_bytes(data[i*3:(i+1)*3],byteorder='big',signed=True)
        i+=1
    y_vec[-20:] = a*scale*1000000
    y_filt = signal.sosfiltfilt(sos,y_vec)
    line1 = live_plotter(x_vec,y_filt,line1)
    y_vec = np.append(y_vec[20:],np.zeros(20))