from __future__ import division

import socket, time, errno, ctypes, struct
import numpy as np
import math
import os
import matplotlib.pyplot as plt
import pickle
from sklearn import linear_model
import scipy


from math import log



def read_u32():
    global data, data_p
    temp=int.from_bytes( data[data_p+0:data_p+4], byteorder='little', signed=False )
    data_p += 4
    return temp
    
def read_32():
    i = read_u32()
    return ctypes.c_int32(i).value

def read_16():
    global data, data_p
    temp = 0
    temp = temp + ((data[data_p+1])<<8)
    temp = temp + ((data[data_p+0]))
    data_p += 2
    return ctypes.c_short(temp).value

def read_f():
    i = read_u32()    
    return struct.unpack('f',struct.pack('I',i))[0]

def read_char():
    global data, data_p
    temp=int.from_bytes( data[data_p+0:data_p+1], byteorder='little', signed=True )
    data_p += 1

    return temp

def read_string_nullify(l):
    global data, data_p
    temp = data[data_p:data_p+l-1].decode("latin-1") + ""
    data_p = data_p + l -1
    return temp

def read_string(l):
    global data, data_p
    temp = data[data_p:data_p+l]
    data_p = data_p + l
    return temp.decode("latin1")

def get_real_timescale(i):
    return [2.0e-9, 5.0e-9,  
	1.0e-8, 2.0e-8, 5.0e-8, # 10 ns
	1.0e-7, 2.0e-7, 5.0e-7, # 100 ns
	1.0e-6, 2.0e-6, 5.0e-6, # 1 us
	1.0e-5, 2.0e-5, 5.0e-5, # 10 us
	1.0e-4, 2.0e-4, 5.0e-4, # 100 us
	1.0e-3, 2.0e-3, 5.0e-3, # 1 ms
	1.0e-2, 2.0e-2, 5.0e-2, # 10 ms
	1.0e-1, 2.0e-1, 5.0e-1, # 100 ms
	1.0e+0, 2.0e+0, 5.0e+0, # 1 s
	1.0e+1, 2.0e+1, 5.0e+1, # 10 s
        1.0e+2 # 100 s
            ][i]
    

def get_real_voltscale(i):
    return [2.0e-2, 5.0e-2, 1.0e-1, 2.0e-1, 5.0e-1, 1.0e+0, 2.0e+0, 5.0e+0, 1.0e+1, 2.0e+1, 5.0e+1, 1.0e+2][i]

def get_real_attenuation(i):
    return [1.0e0, 1.0e1, 1.0e2, 1.0e3][i]


# a2-a1 is the voltage across the DUT
# a2 is the voltage across the reference resistor    

def zimp(a1,a2,p):
    temp=a1*a1+a2*a2-2*a1*a2*math.cos(p)
    temp=math.sqrt(temp)
    return(temp/a2)

def zimpa(a1,a2,p):
    t1=a1*math.sin(p)
    t2=a1*math.cos(p)-a2
    return(math.atan2(t1,t2))



def parse_channel():
    print ("Parse Channel")
    channel ={}
    channel['name'] = read_string_nullify(4)
    channel['unknownint'] = read_32()
    channel['datatype'] = read_32()
    channel['unknown4'] = read_string(4)
    channel['samples_count'] = read_u32()
    channel['samples_file'] = read_u32()
    channel['samples3'] = read_u32()
    channel['timediv'] = get_real_timescale(read_u32())
    channel['offsety'] = read_32()
    channel['voltsdiv'] = get_real_voltscale(read_u32());
    channel['attenuation'] = get_real_attenuation(read_u32())
    channel['time_mul'] = read_f()
    channel['frequency'] = read_f()
    channel['period'] = read_f()
    channel['volts_mul'] = read_f()
    channel['data'] = []
    for i in range (0,channel['samples_file']-1):
        if (channel['datatype'] == 2):
            channel['data'].append(read_16())
        else:
            channel['data'].append(read_char())
            
    return channel



def process_file(filename,freq):
    global data, data_p
    
 
    
    with open(filename, "rb") as text_file:
        full_data=text_file.read()
    
  
    
    data = full_data
    data_p = 0
    data_len = len(full_data)
    head = {}
    
    if(data[12:12+3] == 'SPB'.encode("ascii")):
        head['len'] = read_32()
        head['unknown1'] = read_32()
        head['type'] = read_32()
    
    head['model'] = read_string_nullify(7)
    print(head['model'])
    head['intsize'] = read_32()
    
    if(head['intsize'] != 0xFFFFFF):
        head['serial'] = read_string_nullify(30)
        head['triggerstatus'] = read_char()
        head['unknownstatus'] = read_char()
        head['unknownvalue1'] = read_u32()
        head['unknownvalue2'] = read_char()
        head['unknown3'] = read_string(8)
    
    head['channels_count'] = 0
    
    while(data_p < data_len-1):
        if(data[data_p:data_p+2]=='CH'.encode("ascii")):
            head['channels_count'] = head['channels_count'] + 1
            head['channel'+str(head['channels_count'])] = parse_channel()
            v_mul = head['channel'+str(head['channels_count'])]['volts_mul'] / 100.0
            v_ofs = head['channel'+str(head['channels_count'])]['offsety']
            print("vmul",v_mul)

            head['channel'+str(head['channels_count'])]['datav'] = [((x - v_ofs) * v_mul) for x in head['channel'+str(head['channels_count'])]['data']]

        data_p = data_p + 1
    

#going to do 100000 samples at 500MSPs 200us
#cycles=200us/(period)    
    cycles=0.0002*float(freq)

    samples=100000-1
    length=samples/cycles
    sx=np.empty(samples)
    cx=np.empty(samples)
    for i in range(0,samples):
        sx[i]=np.sin((i*2*np.pi)/length)
        cx[i]=np.cos((i*2*np.pi)/length)

    x = np.asarray(sx)
    x = np.vstack([np.vstack([x,cx])]).transpose()

    clf = linear_model.LinearRegression()


    data=head['channel1']['datav']
    clf.fit(x,data)

    a,b=clf.coef_

    amplitude1=math.sqrt(a*a+b*b)
    phase1=-math.atan2(b,a)
    print("amplitude1",amplitude1,"phase1",phase1)

    data=head['channel2']['datav']
    clf.fit(x,data)

    a,b=clf.coef_

    amplitude2=math.sqrt(a*a+b*b)
    phase2=-math.atan2(b,a)
    print("amplitude2",amplitude2,"phase2",phase2)


    phase=phase1-phase2
    if(phase>np.pi):
        phase=phase-2*np.pi

    return(phase,amplitude1,amplitude2,head)

    
def plotfreq(zhead):
#    plt.figure()
    fig,ax1=plt.subplots()
    ax2=ax1.twinx()
    ax1.plot(zhead['channel1']['datav'])
    ax2.plot(zhead['channel2']['datav'])
    plt.show()




def zi(f,r,l,c,cp):
    r=r*rguess
    l=l*lguess
    c=c*cguess
    cp=cp*cpguess
    
    w=2*np.pi*f
    temp=r+1j*(w*l-1.0/(w*c))
    temp=temp*(1.0-1j*w*cp*temp)/(1.0+w*w*cp*cp*temp*temp)
    return(np.abs(temp)/1000.0)

def za(f,r,l,c,cp):
    r=r*rguess
    l=l*lguess
    c=c*cguess
    cp=cp*cpguess
    
    w=2*np.pi*f
    temp=r+1j*(w*l-1.0/(w*c))
    temp=temp*(1.0-1j*w*cp*temp)/(1.0+w*w*cp*cp*temp*temp)
    return(np.angle(temp))


def zi2(f2,r,l,c,cp):
    [f1,f2]=np.hsplit(f2,2)
    return(np.concatenate((zi(f1,r,l,c,cp),za(f2,r,l,c,cp))))


rguess=340
cguess=300e-12
lguess=0.048
cpguess=2150e-12


results={}

#folder1="./Logs/20210123-175558"
#folder1="./Logs/20210123-181344"
#folder1="./Logs/20210123-184358"
#folder1="./Logs/20210123-215957"
#folder1="./Logs/20210123-223005"

graphicname=""

#graphicname,folder1="Transimp_HCTX_DP1.svg","./Logs/20210125-180342" #1K resistor yellow signal generator 2V div red resistor 1V div 10V China TX 5us depth 100K 500MS/s
#folder1="./Logs/20210125-145532" #1K resistor yellow signal generator 2V div red resistor 1V div 10V Murata RX 5us depth 100K 500MS/s
#folder1="./Logs/20210125-124546"#1K resistor yellow signal generator 2V div red resistor 1V div 10V Chinese RX 5us depth 100K 500MS/s

#folder1="./Logs/20210126-174538"#1K resistor yellow signal generator 2V div red resistor 1V div 10V Murata S 5us depth 100K 500MS/s
#graphicname,folder1="Transimp_HCRX_DP0.svg","./Logs/20210126-142632"#1K resistor yellow signal generator 2V div red resistor 1V div 10V China R 5us depth 100K 500MS/s
#graphicname,folder1="Transimp_JSN_DP0.svg","./Logs/20210126-124839"#1K resistor yellow signal generator 2V div red resistor 1V div 10V JSR 5us depth 100K 500MS/s

#graphicname,folder1="Transimp_Mouse_DP0.svg","./Logs/20210131-191548"#1K resistor yellow signal generator 2V div red resistor 1V div 10V Mouse 5us depth 100K 500MS/s
#graphicname,folder1="Transimp_Murata_R_DP0.svg","./Logs/20210131-173959"#1K resistor yellow signal generator 2V div red resistor 1V div 10V Murata R 5us depth 100K 500MS/s
#graphicname,folder1="Transimp_Murata_S_DP0.svg","./Logs/20210131-155546"#1K resistor yellow signal generator 2V div red resistor 1V div 10V Murata S 5us depth 100K 500MS/s


graphicname,folder1="Transimp_MurataMA40S4S_DP1.svg","./Logs/20210212-123204"#1K resistor yellow signal generator 2V div red resistor 1V div 10V 4040SR 5us depth 100K 500MS/s

#folder1="./Logs/20210215-124054"#1K resistor yellow signal generator 2V div red resistor 1V div 10V inductor 5us depth 100K 500MS/s
#folder1="./Logs/20210215-155108"#1K resistor yellow signal generator 2V div red resistor 1V div 10V inductor 5us depth 100K 500MS/s

#folder1="./Logs/20210215-221651"#1K resistor yellow signal generator 2V div red resistor 1V div 10V inductor||2000pF 5us depth 100K 500MS/s

#folder1="./Logs/20210216-140539"#1K resistor yellow signal generator 2V div red resistor 1V div 10V 2000pF 5us depth 100K 500MS/s



pck=folder1+".pck"

if os.path.isfile(pck):# and False:
    pickle_in = open(pck,"rb")
    results = pickle.load(pickle_in)
    
else:    
    files = os.listdir(folder1)
    for name in files:
        print(name)
        freq=name.rsplit('.', 1)[0]
        ext=name.rsplit('.', 1)[1]
        if(ext=="bin"):
            phase,amplitude1,amplitude2,head=process_file(folder1+"/"+name,freq)
            results[freq]=phase,amplitude1,amplitude2#,head
   #     break;


    pickle_out = open(folder1+".pck","wb")
    pickle.dump(results, pickle_out)
    pickle_out.close()




fig,ax1=plt.subplots(figsize=[8.0,4.0], tight_layout = {'pad': 0})
ax2=ax1.twinx()



a1=[]
a2=[]
p=[]
z=[]
freq=[]
for f in sorted(list(results)):
    p.append(zimpa(results[f][2],results[f][1],results[f][0]))
    a1.append(results[f][1])
    a2.append(results[f][2])
    z.append(zimp(results[f][2],results[f][1],results[f][0]))
    freq.append(float(f))

freq=np.asarray(freq)
z=np.asarray(z)

p=np.unwrap(p)
max_index = max( range( len(z) ), key = lambda index : z[ index ] )
min_index = min( range( len(z) ), key = lambda index : z[ index ] )

ax1.plot(freq,z,'g+',label='Z data')
ax2.plot(freq,p,'r+',label='phase data')
ax2.axhline(y=0, color='r', linestyle='dashed')


freq2=np.concatenate((freq,freq))
z2=np.concatenate((z,p))

#popt,pcov=scipy.optimize.curve_fit(zi2,freq2,z2)
popt,pcov=scipy.optimize.curve_fit(zi,freq,z)
#popt=1.0,1.0,1.0,1.0
print(popt)
print(rguess*popt[0],lguess*popt[1],cguess*popt[2],cpguess*popt[3])

z1=zi(freq,*popt)
ax1.plot(freq,z1,'g',label='model')

p=za(freq,*popt)
ax2.plot(freq,p,'r',label='model')





#ax1.axvline(freq[max_index],color='g', linestyle='dashed')
#ax1.axvline(freq[min_index],color='g', linestyle='dashed')

fig.legend(loc="upper right", bbox_to_anchor=(1,0.8), bbox_transform=ax2.transAxes)
plt.yticks(ax2.get_yticks(),[r"$" + format(r/np.pi, ".2g")+ r"\pi$" for r in ax2.get_yticks()])

ax1.set_xlabel("frequency in Hertz");
ax2.set_ylabel("phase in radians");
ax1.set_ylabel("impedance in kΩ");

mstring='max '+str(round(z[max_index],2))+' kΩ at '+str(freq[max_index])+' Hz '+str(round(p[max_index],2))+' radians'
mstring+="\n"
mstring+='min '+str(round(z[min_index],2))+' kΩ at '+str(freq[min_index])+' Hz '+str(round(p[min_index],2))+' radians'

#ax1.text(0.05, 0.6, mstring,
#        horizontalalignment='left',
#        verticalalignment='top',
#        transform=ax1.transAxes)


ax1.xaxis.set_major_locator(plt.MaxNLocator(10))
ax2.xaxis.set_major_locator(plt.MaxNLocator(10))
plt.savefig('/home/david/'+graphicname, bbox_inches='tight',dpi=200)
plt.show()