import matplotlib.pyplot as plt
import numpy as np
from solcore.absorption_calculator import calculate_absorption_profile
from solcore import material, si
from solcore.structure import Structure, Layer

GaAs = material('GaAs')(T=300)
AlGaAs = material('AlGaAs')(T=300, Al=0.3)

my_structure = Structure([
    Layer(si(30, 'nm'), material=AlGaAs),
    Layer(si(3000, 'nm'), material=GaAs),
    Layer(si(300, 'um'), material=GaAs),
])

wl = np.linspace(400, 1000, 200)
out = calculate_absorption_profile(my_structure, wl, steps_size=1, z_limit=3000)

plt.figure(1)
ax = plt.contourf(out['position'], wl, out['absorption'], 200)
plt.xlabel('Position (nm)')
plt.ylabel('Wavelength (nm)')
cbar = plt.colorbar()
cbar.set_label('Absorption (1/nm)')

A = np.zeros_like(wl)

for i, absorption in enumerate(out['absorption'][:]):
    A[i] = np.trapz(absorption, out['position'])

plt.figure(2)
plt.plot(wl, A)
plt.xlabel('Wavelength (nm)')
plt.ylabel('Absorption')

plt.show()