import pylab as pl
from NL.CommonConcepts.Configurations.Utilities import fractional2cartesian


#-------------------------------------------------------------------------
# Model data
#-------------------------------------------------------------------------

# alpha parameter fitted to bulk InAs
alpha_bulk = 2.869

# InAs bulk conduction band effective mass
meff_bulk = 0.028*electron_mass

# Effective width of slab
slab_width = 30*Angstrom


#-------------------------------------------------------------------------
# Load DFT calcularted bandstructure and effective mass
#-------------------------------------------------------------------------

# Read bulk configuration
configuration = nlread('InAs_slab_pseudohydrogen.hdf5',BulkConfiguration)[0]

# Read bandstructure
bandstructure = nlread('InAs_slab_pseudohydrogen.hdf5',Bandstructure)[-1]

# Read Effective mass
effective_mass = nlread('InAs_slab_pseudohydrogen.hdf5',EffectiveMass)[0]
meff = effective_mass.evaluate(band=1)[0][0][0]

# Get the fractional kpoints
kpoints = bandstructure.kpoints()


# Get reciprocal lattice vectors (used to convert frational k to cartesian k)
G = configuration.bravaisLattice().reciprocalVectors()

# K-points in cartesian coordinates, Ang^-1
k_cart = fractional2cartesian(kpoints, G)

# Get |k| as 1D array
k = numpy.zeros(len(kpoints))
for ii,ki in enumerate(k_cart.inUnitsOf(Angstrom**-1)):
    k[ii] = (ki[0]**2 + ki[1]**2 + ki[2]**2)**0.5

# Add unit
k = k*Angstrom**-1

# Get all the bands
bands = bandstructure.evaluate().inUnitsOf(eV)

# Energies at the Gamma-point
E0 = bands[0,:]

# Index of conduction band
conduction_band_index = numpy.where(E0 > 0.0)[0][0]

# Get the energies of the conduction band
conduction_band = bands[:,conduction_band_index]

# conduction band minimum
cbm = min(conduction_band)

# Evaluate non-parabolic model bandstructure
E_non_parabolic = (-1 + (1 + 2*alpha_bulk*(hbar**2/meff_bulk*(k**2 + numpy.pi**2/(slab_width)**2 ) ).inUnitsOf(eV) )**0.5 )/(2*alpha_bulk)

# Align conduction band minimum to DFT
E_non_parabolic = E_non_parabolic - min(E_non_parabolic) + cbm

# Evaluate parabolic model bandstructure and align CBM
E_parabolic = (0.5*hbar**2*k**2/meff_bulk).inUnitsOf(eV) + cbm


# Ananlytical effective mass of slab
m_slab = meff_bulk * (1 + 2*alpha_bulk*(hbar**2*numpy.pi**2/(meff_bulk*slab_width**2)).inUnitsOf(eV) )**0.5
print('')
print('Analytical slab mass = %.4f m_e' %m_slab.inUnitsOf(electron_mass))
print('DFT slab mass = %.4f m_e\n' %0.095)

# Analytical band gap increase
delta_gap = (-1 +  (1 + 2*alpha_bulk*(hbar**2*numpy.pi**2/(meff_bulk*slab_width**2)).inUnitsOf(eV) )**0.5)/(2*alpha_bulk)
print('Analytical band gap increase = %.4f eV' %delta_gap)
print('DFT gap change = %.4f eV' %(0.91-0.354))
print('')


pl.figure()
pl.plot(k,conduction_band,'k',label='DFT conduction bands')
pl.plot(k,bands[:,conduction_band_index+1:conduction_band_index+5],'k')
pl.plot(k,E_non_parabolic,'r',label='Non-parabolic fit ')
pl.plot(k,E_parabolic,'b',label='Parabolic fit')
pl.xlabel('k  (Ang.$^{-1}$)')
pl.ylabel('Energy  (eV)')
pl.ylim([0,1.5])
pl.legend(loc=0)

pl.show()