from AddOns.TransportCoefficients.Utilities import evaluateMoments
from NL.CommonConcepts import PhysicalQuantity as Units

# Set input parameters
e_min = 0.0
e_max = 2.0
n_steps = 400
temperature = 300.00*Kelvin
energies = numpy.linspace(e_min, e_max, n_steps) * eV

# Read phonon transmission spectrum from file cnt_phonons.nc
# and get the phonon conductance at the specified temperature
phonon_transmission_spectrum = nlread('cnt_phonons.nc', PhononTransmissionSpectrum)[0]
conductance_phonons = phonon_transmission_spectrum.thermalConductance(phonon_temperature=temperature)

# Read phonon transmission spectrum from file cntc14_phonons.nc
# and get the phonon conductance at the specified temperature
phonon_transmission_spectrum_c14 = nlread('cntc14_phonons.nc', PhononTransmissionSpectrum)[0]
conductance_phonons_c14 = phonon_transmission_spectrum_c14.thermalConductance(phonon_temperature=temperature)

# Read electron transmission spectrum from file cnt_electrons.nc
electron_transmission_spectrum = nlread("cnt_electrons.nc",TransmissionSpectrum)[0]

# define lists for ZT data
zzt = []
zzt_c14 = []

# Calculate the ZTs at each energy
for energy in energies:
    # Calculate moments of the electron transmission spectrum
    k0, k1, k2 = evaluateMoments(electron_transmission_spectrum, temperature, energy)
    # Calculate transport properties.
    conductance       = k0*Units.e**2
    peltier           = k1/(k0*Units.e)
    seebeck           = k1/(k0*Units.e*temperature)
    thermal_electrons = (k2*k0-k1*k1)/(temperature*k0)
    total_thermal     = thermal_electrons + conductance_phonons
    total_thermal_c14 = thermal_electrons + conductance_phonons_c14
    zt                = conductance*seebeck**2*temperature/total_thermal
    zt_c14            = conductance*seebeck**2*temperature/total_thermal_c14
    zzt.append(zt)
    zzt_c14.append(zt_c14)

# Plot the ZT as function of energy
zzt = numpy.asarray(zzt)*1e3
zzt_c14 = numpy.asarray(zzt_c14)*1e3
import pylab
pylab.plot(energies,zzt,'-',label='CNT',linewidth=1.0,color='b')
pylab.plot(energies,zzt_c14,'--',label='14C-CNT',linewidth=1.0,color='r')
pylab.xlabel('Energy (eV)')
pylab.ylabel('ZT x 1000')
pylab.legend(loc='upper center')
pylab.show()