import pylab
import numpy.polynomial.polynomial as poly

filename_list=["from-niflheim/tutorial-calcs/gatescan_TS.hdf5", "from-niflheim/tutorial-calcs/drain-underlap/gatescan_TS_du.hdf5"]
line_type=["*b-","or-",]
label=["p-i-n junction","10 nm underlap",""]

# Figure filename
fig_name = "gate_bias.png"

# Shit, in V, to apply to the gate bias to match the work function of a gold metal gate
shift=0.21

# Gate voltage list
pos_labels=numpy.arange(0.0,1.0,0.1)
neg_labels = numpy.flipud(-0.1*numpy.arange(1,6,1))
gate_voltage_list = numpy.append(neg_labels, pos_labels)

for index,filename_TS in enumerate(filename_list):

    # Read the transmission spectra
    transmission_spectrum_list = nlread(filename_TS,TransmissionSpectrum)
    n = len(transmission_spectrum_list)

    # Make list to hold the current calculations
    current_list = numpy.zeros(n)

    # Calculate the current and convert the units in A/m, considering the width of the device
    # in the periodic direction of (B) of 4.28387 Ang
    for i in range(n):
        current_list[i] = transmission_spectrum_list[i].current().inUnitsOf(Ampere)/    \
                        ((4.28387*Ang).inUnitsOf(Meter))
    [pos_current_list,neg_current_list] = numpy.split(current_list,[10])
    neg_current_list = numpy.flipud(neg_current_list[1:])
    current_list = numpy.append(neg_current_list,pos_current_list)
    a = poly.polyfit(gate_voltage_list[9:12]-shift, numpy.log10(abs(current_list[9:12])), deg=1)
    print(a)
    print(1.0/a)
    p = poly.Polynomial(a)

    # Plot in a semi-logarithmic  scale
    pylab.semilogy(gate_voltage_list-shift, abs(current_list), line_type[index], label=label[index])
    if index is 0:
        pylab.semilogy(gate_voltage_list[9:12]-shift, 10**(numpy.array(p(gate_voltage_list[9:12]-shift))), linewidth=4,color='k', linestyle='--')
    else:
        pylab.semilogy(gate_voltage_list[9:12]-shift, 10**(numpy.array(p(gate_voltage_list[9:12]-shift))), linewidth=4,color='k', linestyle='--', label='Best linear fit')

# Set plot properties
pylab.xlabel('V(GS) (V)')
pylab.ylabel('|Current| (A/m)')

pylab.xlim([-0.8,0.8])
pylab.grid(True)
pylab.legend(loc='lower left')
pylab.title("V(SD) = -0.5 Volt")

pylab.savefig(fig_name)
pylab.show()

'''
# Using 0 and 60 as the endpoints is based on a visual inspection of the plot.
first_image = 0
last_image = 60 

# Doing the linear fit and constructing the polynomial from the coefficients.


# Compute the R^2 value
yhat = p(t)                        
ybar = numpy.sum(msd_data)/len(msd_data)          
ssres = numpy.sum((msd_data - yhat)**2)   
sstot = numpy.sum((msd_data - ybar)**2) 
r2 = 1 - ssres / sstot
print('R^2: ', r2)


# Calculating the diffusion coefficient in Ang**2/ps.
diffusion_coefficient = a[-1]/6.0
print('Diffusion coeffcient for liquid copper at T=2000K is: ',diffusion_coefficient, ' Ang**2/ps')


# Plot the data using pylab:

fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_xlabel('t (ps)')
ax.set_ylabel('MSD(t) (Ang**2)')
ax.text(0.03, 0.95, 'Fit parameters: \n$MSD=%.2ft + %.2f$ \n$R^2=%.2f$' %(a[1],a[0],r2), verticalalignment='top', horizontalalignment='left', transform=ax.transAxes,fontsize=14, color='k', bbox={'facecolor':'white','alpha':0.5, 'pad':10})
ax.plot(t, msd_data, linewidth=2,label='MSD of liquid copper')
ax.plot(t[first_image:last_image], p(t[first_image:last_image]), linewidth=4,color='k', linestyle='--', label='Best linear fit')
ax.legend()
plt.ylim(ymin = 0)
plt.show()
'''