import XEFI
import numpy as np
import kkcalc2 as kk

energies = np.linspace(2300, 2900, 1000)  # in eV
angles = 0.55 # Single angle of incidence in degrees, past the critical for all layers.
z = [0, -800, -1340]  # Z-coordinates for the multilayer interface
r = [20, 40, 10] # Roughness values for each interface in Å
layer_names = ["Air", "PS", "P3HT", "Si"]

# Calculated at 8050.92 eV using kkcalc2, but could be any complex
# refractive index or callable that returns a complex refractive index.
refractive_indices = [
    kk.models.asp_db_complex("(N78O20Ar1)0.01", density=0.001225, name="Air"),  # Air/Vacuum
    kk.models.asp_db_complex("C8H8", density=1.05, name="PS"),  # Polystyrene (C8H8)
    kk.models.asp_db_complex("C10H14S", density=1.33, name="P3HT"),  # Poly(3-hexylthiophene) (P3HT, C10H14S)
    kk.models.asp_db_complex("Si", density=2.329, name="Si"),  # Silicon (Si)
 ]

result: XEFI.BasicResult = XEFI.XEF_Basic(
    energies=energies,
    angles=angles,
    z=z,
    refractive_indices=refractive_indices,
    layer_names=layer_names,
    z_roughness=r,
)

fig, ax = result.generate_graphic_XEFI_map(
    grid_roughness=True,
    angles_in_deg=True,
)
ax.set_title(
    rf"X-ray Electric Field Intensity at $\theta$={angles}°",
    pad=20,
)
plt.show()