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

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, [8, 12], btype='bandpass', 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)

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


length = 0.5
tt=np.linspace(0,length,int(length*fs)+1)[0:-1]
N=len(tt)
x_vec = fftfreq(N, 1.0/fs)[:N//2]
y1_vec = np.random.randn(N)
y2_vec = np.random.randn(N)
y3_vec = np.random.randn(N)
y4_vec = np.random.randn(N)

y1_fft = fft(y1_vec)
y2_fft = fft(y2_vec)
y3_fft = fft(y3_vec)


plt.style.use('ggplot')
plt.ion()
fig = plt.figure(figsize=(13,12))

ax1 = fig.add_subplot(211)
plt.ylabel('Amplitude (uV^2/Hz)')
plt.xlabel('Frequency (Hz)')
plt.title('Power Spectrum')

# create a variable for the line so we can later update it
line1, = ax1.plot(x_vec, 2.0/N * np.abs(y1_fft[0:N//2]),'-',alpha=0.8)
line2, = ax1.plot(x_vec, 2.0/N * np.abs(y2_fft[0:N//2]),'-',alpha=0.8)
line3, = ax1.plot(x_vec, 2.0/N * np.abs(y3_fft[0:N//2]),'-',alpha=0.8)
#update plot label/title

ax1.set_ylim(0,10)
ax1.set_xlim(0,30)

center_f=int(12*N/fs)

len2=5
t2=np.linspace(0,len2,int(len2*fs/20)+1)[0:-1]

#left_a=np.random.randn(len(t2))
#right_a=np.random.randn(len(t2))
left_a=  (25/N * np.abs(y1_fft[center_f-1]) + 50/N * np.abs(y1_fft[center_f]) + 25/N * np.abs(y1_fft[center_f+1]))/100
right_a=(25/N * np.abs(y2_fft[center_f-1]) + 50/N * np.abs(y2_fft[center_f]) + 25/N * np.abs(y2_fft[center_f+1]))/100

left_mov_avg=np.random.randn(2)
right_mov_avg=np.random.randn(2)

smoothing_factor=0.1
left_mov_avg[1]=smoothing_factor*left_a+(1-smoothing_factor)*left_mov_avg[0]
right_mov_avg[1]=smoothing_factor*right_a+(1-smoothing_factor)*right_mov_avg[0]


control_sig=np.random.randn(len(t2))
control_sig[-1] = left_mov_avg[1] - right_mov_avg[1]

ax2 = fig.add_subplot(212)
plt.ylabel('Time')
plt.xlabel('left-right mod')
ax2.set_xlim(-2,2)
#plt.title('')

# create a variable for the line so we can later update it
line4, = ax2.plot(control_sig, t2 ,'-',alpha=0.8)

plt.show()

left_mov_avg = np.append(left_mov_avg[1],0)
right_mov_avg = np.append(right_mov_avg[1],0)
control_sig = np.append(control_sig[1:],0)

a=np.zeros(20,dtype='int64')
b=np.zeros(20,dtype='int64')
c=np.zeros(20,dtype='int64')
d=np.zeros(20,dtype='int64')
while True:
    data, address = sock.recvfrom(4096)
    i=0
    while i<20:
        a[i] = int.from_bytes(data[i*8:i*8+2],byteorder='big',signed=True)
        b[i] = int.from_bytes(data[i*8+2:i*8+4],byteorder='big',signed=True)
        c[i] = int.from_bytes(data[i*8+4:i*8+6],byteorder='big',signed=True)
        d[i] = int.from_bytes(data[i*8+6:i*8+8],byteorder='big',signed=True)
        i+=1

    y1_vec[-20:] = a*256*scale*1000000
    y2_vec[-20:] = b*256*scale*1000000
    y3_vec[-20:] = c*256*scale*1000000
    #y4_vec[-20:] = d*256*scale*1000000

    #y1_filt = signal.sosfilt(sos,y1_vec)
    #y2_filt = signal.sosfilt(sos,y2_vec)
    #y3_filt = signal.sosfilt(sos,y3_vec)
    #y4_filt = signal.sosfilt(sos,y4_vec)
 

    y1_fft = fft(y1_vec)
    y2_fft = fft(y2_vec)
    y3_fft = fft(y3_vec)
    #y4_fft = signal.sosfilt(sos,y4_vec)

    


    left_a = (25/N * np.abs(y1_fft[center_f-1]) + 50/N * np.abs(y1_fft[center_f]) + 25/N * np.abs(y1_fft[center_f+1]))/100
    right_a = (25/N * np.abs(y2_fft[center_f-1]) + 50/N * np.abs(y2_fft[center_f]) + 25/N * np.abs(y2_fft[center_f+1]))/100

    left_mov_avg[1]=smoothing_factor*left_a+(1-smoothing_factor)*left_mov_avg[0]
    right_mov_avg[1]=smoothing_factor*right_a+(1-smoothing_factor)*right_mov_avg[0]
    control_sig[-1] = left_mov_avg[1] - right_mov_avg[1]


    line1.set_ydata(2.0/N * np.abs(y1_fft[0:N//2]))
    line2.set_ydata(2.0/N * np.abs(y2_fft[0:N//2]))
    line3.set_ydata(2.0/N * np.abs(y3_fft[0:N//2]))
    line4.set_xdata(control_sig)

    #ax1.set_ylim(0, 2/N * np.max(np.abs([y1_fft[5:N//2],y2_fft[5:N//2],y3_fft[5:N//2]])))
    
    y1_vec = np.append(y1_vec[20:],np.zeros(20))
    y2_vec = np.append(y2_vec[20:],np.zeros(20))
    y3_vec = np.append(y3_vec[20:],np.zeros(20))
    #y4_vec = np.append(y4_vec[20:],np.zeros(20))
    left_mov_avg = np.append(left_mov_avg[1],0)
    right_mov_avg = np.append(right_mov_avg[1],0)
    control_sig = np.append(control_sig[1:],0)



    plt.pause(.001)