#*******************************************************************#
#                                                                   #
# USER DEFINED PARAMETERS                                           #
#                                                                   #
#*******************************************************************#
#------------------------                                           #
# HDF5 FILES (INPUT)                                                #
#------------------------                                           #
# File with transmission spectrum for pristine wires                #
inputfile_pristine = 'sinw_pristine_trans.hdf5'                     #
# File with transmission spectra for phosporous doped wires         #
inputfile_P_doping = 'sinw_P_dopant_trans_1234.hdf5'                #
#------------------------                                           #
# OUTPUT FILEA                                                      #
#------------------------                                           #
# File for saving calculated data for later analysis                #
savefile = 'sinw.pckl'                                              #
# PNG or PDF file with plot                                         #
pngfile = 'T_and_mfp.png'                                                 #
#------------------------                                           #
# DOPING                                                            #
#------------------------                                           #
# Doping density  (cm**3)                                           #
doping_density = 1e19                                               #
# Number of Si atoms in unit cell                                   #
nua0 = 48                                                           #
# Length of unit cell in z-direction                                #
Lz = 5.431*Angstrom                                                 #
# Volume of each atom in a bulk Si crystal                          #
atom_volume = 0.5*40.039*Angstrom**3                                #
#                                                                   #
#*******************************************************************#
#                                                                   #
# NO CHANGES NEEDED BELOW                                           #
#                                                                   #
#*******************************************************************#

#====================================================================
# Load transmissions
#====================================================================
# Load pristine wire transmission
trans_elec = nlread(inputfile_pristine, TransmissionSpectrum)[-1]
T0 = trans_elec.evaluate()
E0 = trans_elec.energies().inUnitsOf(eV)
# Load dopant transmissions
trans_dopant = nlread(inputfile_P_doping, TransmissionSpectrum)
E1 = trans_dopant[0].energies().inUnitsOf(eV)
T_doped = []
for transmission in trans_dopant:
    T_doped.append(transmission.evaluate())
T1 = T_doped[0]

#====================================================================
# Align conduction band minima and slice out transmissions
#====================================================================
# Conduction band edge for pristine wire
Ec0 = E0[numpy.where(T0<1e-9)[0][-1]]
# Conduction band edge for wires with dopants
Ec1 = E1[numpy.where(T1<1e-9)[0][-1]-1]
# Shift E1 to fit E0
E1 -= (Ec1 - Ec0)
# Extract transmission on a common energy grid
Emin = max(min(E0),min(E1),)
Emax = min(max(E0),max(E1))
dE = E0[1]-E0[0]
# Energy range for transmissions
E = numpy.arange(Emin,Emax,dE)
# Slice out part of transmission covered by the common energy grid
T0 = T0[numpy.argmin(abs(E0-E[0])):numpy.argmin(abs(E0-E[-1]))]
for i,transmission in enumerate(T_doped):
    T_doped[i] = transmission[numpy.argmin(abs(E1-E[0])):numpy.argmin(abs(E1-E[-1]))]
E = E[:-1]

#====================================================================
# Calculate average transmission and scattering resistance
#====================================================================
# Average transmission
Tave = numpy.mean(T_doped, axis=0)
# Contact resistance
Rc = 1/(T0 + 1e-19)
# Average defect resistances
Rave = 1/(Tave + 1e-20)
# Scattering resistances
Rs = Rave - Rc

#====================================================================
# Save data for later use
#====================================================================
import pickle
f = open(savefile, 'w')
pickle.dump((E,Ec0,T0,T_doped,Tave,Rc,Rave,Rs), f)
f.close()

#====================================================================
# Calculate the scattering mean free path for a given dopant density
#====================================================================
# Volume of one Si atom in units of cm^3
atom_volume_cm = atom_volume.inUnitsOf(Units.cm**3)
# Average distance between two dopant atoms
dopant_distance = 1.0/(doping_density * atom_volume_cm * nua0) * Lz
print('')
print('Doping density:  %2.2e cm^-3'    % doping_density)
print('Dopant-distance: %4.2f Angstrom' % (dopant_distance.inUnitsOf(Ang)))
print('')
# Compute the mean free path
mfp = Rc/(Rs) * dopant_distance

#====================================================================
# Plotting results
#====================================================================
import matplotlib.pyplot as plt
fig = plt.figure(figsize=(6,7))
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212, sharex=ax1)
# Plot the transmissions
ax1.plot(E,T0,'k-',label='Pristine')
for i in range(0,len(T_doped)):
    ax1.plot(E,T_doped[i],label='Dopant %i'%(i+1))
ax1.plot(E,Tave,'g-',label='Average',lw=3)
ax1.set_xlim((1.0,1.8))
ax1.set_ylabel('Transmission')
ax1.legend(loc='upper left',ncol=2,frameon=False)
# Plot the mean free path
ax2.plot(E,mfp.inUnitsOf(nm),'b-',lw=2,label='MFP')
ax2.set_xlabel('Energy (eV)')
ax2.set_ylabel('Mean free path (nm)')
ax2.legend(loc='upper left',frameon=False)
# Finalize figure
plt.tight_layout()
plt.savefig(pngfile)
plt.show()