import aestimo.aestimo as solver
import aestimo.config as ac
ac.messagesoff = True # turn off logging in order to keep notebook from being flooded with messages.
import aestimo.database as adatabase
import numpy as np
import matplotlib.pyplot as plt
import copy
from pprint import pprint

meV2J = solver.meV2J # conversion factor
q = solver.q # electron charge

adatabase.materialproperty = {
    'GaAs':{
        'm_e':0.067,
        'm_e_alpha':0.0,
        'epsilonStatic':12.90,
        'Eg':0.0, #1.426, # set the energy scale origin to be at the GaAs condution band
        'Band_offset':0.67,
        },
    'AlAs':{
        'm_e':0.067, # normally Harrison would be using 0.15
        'm_e_alpha':0.0,
        'epsilonStatic':10.06,
        'Eg':2.673-1.426, #2.673, # set the energy scale origin to be at the GaAs condution band
        'Band_offset':0.67,
        },
    }

adatabase.alloyproperty = {
    'AlGaAs':{
        'Bowing_param':0.0,
        'Band_offset':0.67,
        'Material1':'AlAs',
        'Material2':'GaAs',
        'm_e_alpha':0.0,
        },
    }

s0 = {} # this will be our datastructure

# TEMPERATURE
s0['T'] = 300.0 #Kelvin

# COMPUTATIONAL SCHEME
# 0: Schrodinger
# 1: Schrodinger + nonparabolicity
# 2: Schrodinger-Poisson
# 3: Schrodinger-Poisson with nonparabolicity
# 4: Schrodinger-Exchange interaction
# 5: Schrodinger-Poisson + Exchange interaction
# 6: Schrodinger-Poisson + Exchange interaction with nonparabolicity
s0['comp_scheme'] = 0

# Non-parabolic effective mass function
# 0: no energy dependence
# 1: Nelson's effective 2-band model
# 2: k.p model from Vurgaftman's 2001 paper
s0['meff_method'] = 0 

# Non-parabolic Dispersion Calculations for Fermi-Dirac
s0['fermi_np_scheme'] = True

# QUANTUM
# Total subband number to be calculated for electrons
s0['subnumber_e'] = 10

# APPLIED ELECTRIC FIELD
s0['Fapp'] = 0.00 # (V/m)

# GRID
# For 1D, z-axis is choosen
gridfactor= 0.01 #nm
s0['dx'] = gridfactor*1e-9
maxgridpoints = 200000 #for controlling the size

# STRUCTURE
# a finite parabolic well between two barriers
a = 10.0 #nm #maximum width of parabolic well (at the barrier's energy level)
b = 10.0 #nm #width of the first barrier layer
b2 = 5.0 #nm #width of the first barrier layer
xmin = 0.0 #minimum Al alloy in structure
xmax = 10.0 #maximum Al alloy in structure

def alloy_profile(z):
    """function of alloy profile for finite parabolic well"""
    well = xmin + (z - (b+a/2.0))**2 / (a/2.0)**2 * (xmax - xmin)    
    return np.where(well<xmax,well,xmax) #cuts off parabolic well at barrier height

def bandstructure_profile(x):
    mat1 = adatabase.materialproperty['AlAs']
    mat2 = adatabase.materialproperty['GaAs']
    alloy = adatabase.alloyproperty['AlGaAs']
    return alloy['Band_offset']*(x*mat1['Eg'] + (1-x)* mat2['Eg']-alloy['Bowing_param']*x*(1-x))*solver.q # for electron. Joule
    #return alloy['Band_offset']*(x*mat1['Eg'] + (1-x)* mat2['Eg'])*aestimo.q # for electron. Joule

# Nb. we could have avoided using the database module at all since we are 
# defining our own Structure object and calculating its attribute arrays 
# ourselves but I wanted to show how easy it is to customise the database's values.

## Create structure arrays -----------------------------------------------------

# Calculate the required number of grid points
s0['x_max'] = sum([b,a,b2])*1e-9 #total thickness (m)
s0['n_max'] = n_max = int(s0['x_max']/s0['dx'])
# Check on n_max
if s0['n_max'] > maxgridpoints:
    solver.logger(" Grid number is exceeding the max number of %d", maxgridpoints)
    exit()
#
meff = adatabase.materialproperty['GaAs']['m_e']*solver.m_e
epsGaAs = adatabase.materialproperty['GaAs']['epsilonStatic']

z = np.arange(0.0,s0['x_max'],s0['dx'])*1e9 #nm

s0['cb_meff'] = np.ones(n_max)*meff	#conduction band effective mass
s0['cb_meff_alpha'] = np.zeros(n_max)   #non-parabolicity constant.
s0['eps'] = np.ones(n_max)*epsGaAs	#dielectric constant
s0['dop'] = np.ones(n_max)*0.0          #doping
s0['fi'] = bandstructure_profile(alloy_profile(z))   #Bandstructure potential

# Initialise structure class
model = solver.Structure(**s0)

#config.d_E = 1e-5*meV2J
#config.Estate_convergence_test = 3e-10*meV2J
result= solver.Poisson_Schrodinger(model)

#Plot QW representation
%matplotlib inline
ac.wavefunction_scalefactor = 5000
solver.QWplot(result)#,figno=None)

Energies = np.array(result.E_state)
Energies[1:] - Energies[:-1]