from ErwinJr2 import SaveLoad
from ErwinJr2.QCPlotter import plotPotential, plotWF
from ErwinJr2.QCLayers import QCLayers, auto_gain
import matplotlib.pyplot as plt
import numpy as np

qcl = QCLayers(
   substrate='InP',
   EField=51.0,
   layerWidths=[
      34.0, 14.0, 33.0, 13.0, 32.0, 15.0, 31.0, 19.0, 29.0, 23.0, 27.0, 25.0,
      27.0, 44.0, 18.0, 9.0, 57.0, 11.0, 54.0, 12.0, 45.0, 25.0],
   layerMtrls=[
      0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
   layerDopings=[
      0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0,
      0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
   layerARs=[
      False, False, False, False, False, False, False, False, False, False,
      False, False, False, True, True, True, True, True, True, True, True,
      False],
   mtrlIFRDelta=[1.2, 1.2],
   mtrlIFRLambda=[90.0, 90.0]
)
# IFR must be manually enabled, otherwise the IFR parameters will be ignored
qcl.includeIFR = True
# or load from file
# with open("../../ErwinJr2/example/std8um.json", 'r') as f:
#     qcl = SaveLoad.qclLoad(f)

# To use customized IFR settings
# qcl.customIFR = True
# array to define individual IFR for interfaces.. the length should be the
# same as the number of layers, which means IFR parameters between layer[n]
# and layer[n+1]
# qcl.ifrDelta = [...]
# qcl.ifrLambda = [...]

qcl.populate_x()
qcl.solve_whole()
qcl.period_recognize()
qcl.full_population()
plotPotential(qcl)
plotWF(qcl)
plt.xlabel('Position (Å)')
plt.ylabel('Energy (eV)')
plt.show()