from __future__ import print_function
import time
import numpy as np
import pandas as pd 
import zhinst.utils
import matplotlib.pyplot as plt
from matplotlib import ticker, cm
import matplotlib as mpl
mpl.rcParams['text.color'] = '#009ee0'
mpl.rcParams['axes.labelcolor'] = '#009ee0'
mpl.rcParams['xtick.color'] = '#009ee0'
mpl.rcParams['ytick.color'] = '#009ee0'

device_id = 'dev4220' # replace with your device SN here
apilevel_example = 6
err_msg = "This example only supports instruments with IA option."
(daq, device, _) = zhinst.utils.create_api_session(device_id, apilevel_example,
                                                       required_options=['IA'],
                                                       required_err_msg=err_msg)
zhinst.utils.api_server_version_check(daq)
zhinst.utils.disable_everything(daq, device)

# We use the auto-range example before the next point


imp_index = 0
exp_settings = [['/%s/imps/%d/enable' % (device, imp_index), 1],
                    ['/%s/imps/%d/mode' % (device, imp_index), 0],
                    ['/%s/imps/%d/auto/output' % (device, imp_index), 1],
                    ['/%s/imps/%d/auto/bw' % (device, imp_index), 1],
                    ['/%s/imps/%d/bias/enable' % (device, imp_index), 1],
                    ['/%s/imps/%d/bias/value' % (device, imp_index), 0],
                    ['/%s/imps/%d/freq' % (device, imp_index), 1000],
                    ['/%s/imps/%d/auto/inputrange' % (device, imp_index), 1],
                    ['/%s/imps/%d/output/amplitude' % (device, imp_index), 0.1],
                    ['/%s/imps/%d/output/range'% (device, imp_index), 10],
                    ['/%s/imps/%d/model'% (device, imp_index), 0]]


# Subscribe to the impedance sample node path.
path = '/%s/imps/%d/sample' % (device, imp_index)
daq.set(exp_settings)
daq.sync()

amp_list = []
the_list = []
c_list=[]

'''for y in np.logspace(3, 6, 101): #freq
    daq.setDouble('/%s/imps/%d/freq' % (device, imp_index), y)
    for x in np.linspace(0.0, 2.0, 101): #bias/value, 2V is enough to turn on the LED
        daq.setDouble('/%s/imps/%d/bias/value' % (device, imp_index), x)'''  
# if sweep DC-offset, MFIA always auto-ranges, this is not good the instrument lifetime

for x in np.linspace(1.0, 1.7, 201): #bias/value, 2V is enough to turn on the LED
    daq.setDouble('/%s/imps/%d/bias/value' % (device, imp_index), x)
    for y in np.logspace(3, 6, 201): #freq
        daq.setDouble('/%s/imps/%d/freq' % (device, imp_index), y)


        daq.sync() #synchronize settings
            
        daq.subscribe(path)
        sleep_length = 1.0
        time.sleep(sleep_length) #this is to ensure the transient is not captured
        poll_length = 0.1  # [s]
        poll_timeout = 500  # [ms]
        poll_flags = 0
        poll_return_flat_dict = True
        data = daq.poll(poll_length, poll_timeout, poll_flags, poll_return_flat_dict) #in the 19.05 API, replace with flat=True
        daq.unsubscribe('*')

        impedanceSample = data[path]
        amp_ = np.abs(np.mean(impedanceSample['z']))
        the_ = np.angle(np.mean(impedanceSample['z']))*180/np.pi  #convert radian into deg, phase can also be extracted from data directly
        c_ = np.mean(impedanceSample['param1']) ##this presumes MFIA is using R||C model
        
        
        amp_list.append(amp_)
        the_list.append(the_)
        c_list.append(c_)

amp_array = np.array(amp_list).reshape(201,201) #convert into 2D matrix, and flip later x and y for easy plotting
the_array = np.array(the_list).reshape(201,201)
c_array = np.array(c_list).reshape(201,201)

real_array = amp_array*np.cos(the_array*np.pi/180) 
imag_array = amp_array*np.sin(the_array*np.pi/180)


'''This part can be used to show the original 2D amplitude and theta plots

plt.close('all')
X, Y = np.meshgrid(np.linspace(1.0, 1.7, 201), np.logspace(3, 6, 201))


fig0, ax0 = plt.subplots()
cs0 = ax0.contourf(X, Y, np.transpose(amp_array), locator=ticker.LogLocator(), cmap=cm.PuBu_r)
ax0.set_title('Impedance (Ohm)')
ax0.set_xlabel('DC bias (V)')
ax0.set_yscale('log')
ax0.set_ylabel('Frequency (Hz)')
ax0.autoscale()
cbar = fig0.colorbar(cs0)   
plt.draw()
plt.show()
    
fig1, ax1 = plt.subplots()
cs1 = ax1.contourf(X, Y, np.transpose(the_array), locator=ticker.LinearLocator(), cmap=cm.RuBu_r)
ax1.set_title('Phase (Deg)')
ax1.set_xlabel('DC bias (V)')
ax1.set_yscale('log')
ax1.set_ylabel('Frequency (Hz)')
ax1.autoscale()
cbar = fig1.colorbar(cs1)

'''

###here the data is excerpted only up to 1.56 V to avoid meaningless negative capacitance values at high DC bias

X, Y = np.meshgrid(np.linspace(1.0, 1.56, 160), np.logspace(3, 6, 201))


amp_array_sub=amp_array[0:160]
the_array_sub=the_array[0:160]
c_array_sub=c_array[0:160]

real_array_sub = real_array[::20]
imag_array_sub = imag_array[::20]


fig0, ax0 = plt.subplots()
cs0 = ax0.contourf(X, Y, np.transpose(amp_array_sub), locator=ticker.LogLocator(), cmap=cm.PuBu_r)
ax0.set_title('Impedance (Ohm)')
ax0.set_xlabel('DC bias (V)')
ax0.set_yscale('log')
ax0.set_ylabel('Frequency (Hz)')
ax0.autoscale()
cbar = fig0.colorbar(cs0)   
plt.draw()
plt.show()
    
fig1, ax1 = plt.subplots()
cs1 = ax1.contourf(X, Y, np.transpose(the_array_sub), locator=ticker.LinearLocator(), cmap=cm.PuBu_r)
ax1.set_title('Phase (Deg)')
ax1.set_xlabel('DC bias (V)')
ax1.set_yscale('log')
ax1.set_ylabel('Frequency (Hz)')
ax1.autoscale()
cbar = fig1.colorbar(cs1)   
plt.draw()
plt.show()

fig2, ax2 = plt.subplots()
cs2 = ax2.contourf(X, Y, np.transpose(c_array_sub), locator=ticker.LinearLocator(20), cmap=cm.PuBu_r)
ax2.set_title('Capacitance (F)')
ax2.set_xlabel('DC bias (V)')
ax2.set_yscale('log')
ax2.set_ylabel('Frequency (Hz)')
ax2.autoscale()
cbar = fig2.colorbar(cs2)   
plt.draw()
plt.show()

##Nyquist plot

fig3, ax3 = plt.subplots()
plt.plot(np.transpose(real_array_sub), -1*np.transpose(imag_array_sub)) #plot in 0.07V DC step
plt.ticklabel_format(style='sci', scilimits=(0,0))
ax3.set_title('    Nyquist plot at different DC bias at a step of 0.07 V')
ax3.set_xlabel('Real Z (Ohm)')
ax3.set_ylabel('-Imaginary Z (Ohm)')
plt.axis('equal')
ax3.autoscale()
ax3.set(xlim=(0, 6000000), ylim=(0,6000000)) #adjust to fit into the scale of the largest impedance range

plt.draw()
plt.show()
    
daq.setDouble('/%s/imps/%d/bias/value' % (device, imp_index), 0) # to turn off the LED
daq.setDouble('/%s/imps/%d/freq' % (device, imp_index), 1000)

pd.DataFrame(amp_array).to_csv('amp.csv') #save for post processing. These are enough to reconstruct other impedance parameters
pd.DataFrame(the_array).to_csv('the.csv')
pd.DataFrame(c_array).to_csv('c.csv')