import pylab

T=300

directory = './'

fig_name = "IV_density_plot.png"
filename_device = directory + '0K-zero-bias.py/0K-zero-bias.py.hdf5'

device_configuration = nlread(filename_device,DeviceConfiguration)[-1]
left = device_configuration.electrodes()[0]
v=left.bravaisLattice().conventionalVectors().inUnitsOf(Ang)
area =  v[0][0]*v[1][1]
area_cm = area * 1e-16

symbols = ["sk-", "or--"]

m = 0
for label in ['Noninteracting','STD']:
    if label is 'Noninteracting':
        transmission_spectrum_list = nlread(directory + '0K-iv/biasscan_TS2.hdf5',TransmissionSpectrum)
    elif label is 'STD':
        transmission_spectrum_list = nlread(directory + '300K-iv/biasscan_TS_std.hdf5',TransmissionSpectrum)

    n = len(transmission_spectrum_list)

    # Make list of relevant gate voltages.
    voltage_list = numpy.zeros(n)
    # Make list to hold the current calculations.
    current_list = numpy.zeros(n)

    # Loop through the voltages.
    for i in range(n):
        voltage_list[i] = transmission_spectrum_list[i].bias().inUnitsOf(Volt)
        current_list[i] = transmission_spectrum_list[i].current(electrode_temperatures=[T,T]*Kelvin).inUnitsOf(Ampere)

    voltage_list = numpy.append(voltage_list,0.0)
    current_list = numpy.append(current_list,1e-24*area_cm)
    nn = numpy.argsort(voltage_list)
    pylab.semilogy(-voltage_list[nn], abs(current_list[nn]/area_cm), symbols[m], ms=15,linewidth=3 , label=label)
    m += 1
    print(label)
    for i in range(len(voltage_list)):
        for j in range(len(voltage_list)):
            if i == j:
                continue
            if abs(voltage_list[nn][i]) == abs(voltage_list[nn][j]):
                print('Vbias i ', abs(voltage_list[nn][i]), ' Vbias j :', abs(voltage_list[nn][j]))
                print(abs(current_list[nn][i]/area_cm)/abs(current_list[nn][j]/area_cm))

# Note that we need to change the sign of the bias to reproduce the figure in the paper. This is because the final implementation introduced a phase change compared to the prototype used in the paper, meaning the atomic positions have been mirrored in our configuration, compared to the paper.

#################################
# Read LOE
#################################
Load_Path = directory + 'loe/'
str_file = 'loe_bias_'

transmission_spectrum_list = nlread(directory + '0K-iv/biasscan_TS2.hdf5',TransmissionSpectrum)

I_loe = []
Vbias = []
kpt_factor = 1.0/(21**2)
for i, transmission_spectrum in enumerate(transmission_spectrum_list):
    V_L = transmission_spectrum.electrodeVoltages()[0].inUnitsOf(Volt)
    V_R = transmission_spectrum.electrodeVoltages()[1].inUnitsOf(Volt)
    bias = V_L - V_R
    if abs(bias)>0:
        Vbias.append(bias)
        # Load LOE
        filename = Load_Path+'loe_bias_%4.4f_kpts_1x1_qpts_1x1.hdf5' %(bias)   
        I0=transmission_spectrum.current(spin=Spin.Up,electrode_temperatures=[T,T]*Kelvin).inUnitsOf(Ampere)
        it = nlread(filename ,InelasticTransmissionSpectrum)[0]
        # Calculate the inelastic contribution to the current
        I_loe.append(I0+it.inelasticCurrent(bias=bias*Volt,temperature=T*Kelvin).inUnitsOf(Ampere)[0]*kpt_factor)

Vbias.append(0.0)
I_loe.append(1e-24*area_cm)
Vbias = numpy.array(Vbias)
I_loe = numpy.array(I_loe)
nn = numpy.argsort(Vbias)
pylab.semilogy(-Vbias[nn],abs(I_loe[nn]/area_cm),'<b-.',ms=12,linewidth=3,label='LOE')
print('LOE')
for i in range(len(Vbias)):
    for j in range(len(Vbias)):
        if i == j:
            continue
        if abs(Vbias[nn][i]) == abs(Vbias[nn][j]):
            print('Vbias i ', abs(Vbias[nn][i]), ' Vbias j :', abs(Vbias[nn][j]))
            print(abs(I_loe[nn][i]/area_cm)/abs(I_loe[nn][j]/area_cm))

#################################
pylab.xlabel('Bias (V)')
pylab.ylabel('|Current density| ($A/cm^2$)')  
pylab.xlim([-0.8,0.7])
pylab.ylim([1e-6,1e8])
pylab.grid(True)
pylab.legend(loc=0)
pylab.savefig(fig_name)
pylab.show()