from QuantumATK import *
import pylab
import sys

filename = "au_vacuumgap.hdf5"

# Read the data
potentials = nlread(filename, ElectrostaticDifferencePotential)
densities = nlread(filename, ElectronDifferenceDensity)
configuration_list = nlread(filename, DeviceConfiguration, read_state=False)

# Make a fingerprint/bias lookup table
bias_dict = {}
for configuration in configuration_list:
    fingerprint = configuration.calculator()._fingerPrint()
    voltages = configuration.calculator().electrodeVoltages()
    bias_dict[fingerprint] = abs(voltages[1].inUnitsOf(Volt)-voltages[0].inUnitsOf(Volt))
    
# Calculate the volume element
dX, dY, dZ = densities[0].volumeElement().inUnitsOf(Ang)
volume = numpy.abs(numpy.dot(dX,numpy.cross(dY,dZ) ))*Ang**3

# Get the bias used for each Mulliken object
voltage_list_e = [bias_dict[i._fingerPrint()] for i in potentials]
voltage_list_d = [bias_dict[i._fingerPrint()] for i in densities]

# Sort all lists by bias!
sort_order_e = numpy.argsort(voltage_list_e)
sort_order_d = numpy.argsort(voltage_list_d)
biases = numpy.take(voltage_list_d,sort_order_d,axis=0)
potentials = numpy.take(potentials,sort_order_e,axis=0)
densities = numpy.take(densities,sort_order_d,axis=0)

# Calculate the electrostatic energy
energies = []
for i in range(len(potentials)):
    induced_potential = potentials[i]-potentials[0]
    induced_density = densities[i]-densities[0]
    energy = -0.5*(induced_density[:,:,:]*induced_potential[:,:,:]*elementary_charge).sum()*volume
    energies += [energy.inUnitsOf(eV)]

# Fit a parabola
energies = numpy.array(energies)
pol = numpy.polyfit(biases, energies, 2)
c = 2.0*pol[0]*elementary_charge/Volt

# Save the computed arrays for future use
numpy.savez('%s_energy_data.npz' % filename.replace(".hdf5",""), v=biases, e=energies)

# Parallel plate, c = epsilon A/d, calculate d
vector_a = configuration_list[0].centralRegion().bravaisLattice().primitiveVectors()[0]
vector_b = configuration_list[0].centralRegion().bravaisLattice().primitiveVectors()[1]

A = numpy.linalg.norm(numpy.cross(vector_a, vector_b))*Angstrom**2
d = vacuum_permitivity*A/c

# --------- Plotting the fit ---------
pylab.figure()

# Regression line
v = numpy.arange(0,1,0.01)
line = pol[0]*v**2 + pol[1]*v + pol[2]

pylab.plot(v,line,'r-',biases, energies,'o')
pylab.xlabel('Bias (Volt)')
pylab.ylabel('Energy (eV)')

pylab.savefig("%s_electrostatic_energy.png" % filename.replace(".hdf5",""), dpi=120)
pylab.show()

# --------- Report the results ---------
print('Capacitance    C = %g Farad'    % c.inUnitsOf(Coulomb/Volt))
print('Parallel plate d = %g Angstrom' % d.inUnitsOf(Ang))