from QuantumATK import *
import sys
import pylab

from NL.ComputerScienceUtilities.Functions import realVectorToNumpy

# ------------ Preparation: Read in the Mulliken data and sort it by bias ---------

filename = "au_vacuumgap.hdf5"

# Read the device configs and Mulliken populations from the NC file
configuration_list = nlread(filename, DeviceConfiguration, read_state=False)
mulliken_list = nlread(filename, MullikenPopulation)

# 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))
    
# Get the atomic occupations from the zero-bias config
calculator = configuration_list[0].calculator()
builder = calculator._builder(configuration_list[0])
initial_population = builder.neutralAtomicOccupations()

# Calculate the net Mulliken charge on each atom
mulliken_charges = [mulliken.atoms(spin=Spin.Sum)-initial_population for mulliken in mulliken_list]

# Get the bias used for each Mulliken object
voltage_list = [bias_dict[i._fingerPrint()] for i in mulliken_list]

# Sort by bias!
sort_order = numpy.argsort(voltage_list)
voltages = numpy.take(voltage_list,sort_order,axis=0)
net_charges = numpy.take(mulliken_charges,sort_order,axis=0)

# Save the computed arrays for future use
numpy.savez('%s_mulliken_data.npz' % filename.replace(".hdf5",""), v=voltages, m=net_charges)

# ------------ Part 1: Plotting the net Mulliken charges ---------

f, axarr = pylab.subplots(2, sharex=True)
axarr[0].set_title('Net Mulliken Charge')
axarr[1].set_title('Accumulated Mulliken Charge')
axarr[1].set_xlabel('Atom Index')
for c in net_charges:
    axarr[0].plot(numpy.arange(len(c)),c-net_charges[0])
    axarr[1].plot(numpy.arange(len(c)),numpy.cumsum(c-net_charges[0]))

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

pylab.clf()

# ------------ Part 2: Fitting a line to the induced charge vs bias ---------

# Compute the total induced charge on one surface, defining the first 7 atoms as belonging to the left surface
q = [-numpy.cumsum(c-net_charges[0])[10] for c in net_charges]
# Fit a straight line to C vs V
M = numpy.array([voltages, numpy.ones(len(q))])
w = numpy.linalg.lstsq(M.T,q)[0]

# Find the capacitance
c = w[0]*elementary_charge/Volt

# Regression line
line = w[0]*voltages+w[1] 

# Parallel plate, c = epsilon A/d, calculate d
# A = cross section area of device in XY plane
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 line vs the data points
pylab.xlabel('Bias (Volt)')
pylab.ylabel('Charge (e)')
pylab.title('Induced charge vs. bias voltage')
pylab.plot(voltages, line, 'r-', voltages, q, 'o')

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

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