import XEFI
import numpy as np

energy = 8050.92 # eV, corresponding to a wavelength of 1.54 Å
angles = np.linspace(0.1, 0.4, 3000) # Angles of Incidence in degrees
z = [0, -800, -1340]  # Z-coordinates for the multilayer interface
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 = [
    1.0        + 0j,         # Air/Vacuum
    0.99999637 + 4.96e-09j,  # Polystyrene (C8H8)
    0.99999536 + 3.31e-08j,  # Poly(3-hexylthiophene) (P3HT, C10H14S)
    0.99999243 + 1.72e-07j,  # Silicon (Si)
 ]

result = XEFI.XEF_Sliced(
    energies=energy,
    angles=angles,
    z=z,
    refractive_indices=refractive_indices,
    z_roughness=[40, 20, 5],
    slice_thickness=1.0,
    sigmas=4.0,
    layer_names=layer_names,
)

ax = result.graph_refractive_indexes()
ax[0].figure.tight_layout()
plt.show()

# Can provide z_vals argument to specify depth values for the map,
# but by default plots extra 10% of the total thickness into the semi-infinite layers.
fig, ax = result.generate_graphic_XEFI_map()
ax.set_title(
    rf"X-ray Electric Field Intensity at $\lambda$={XEFI.utils.en2wav(energy):0.2f} Å, {energy:0.2f} eV",
    pad=20,
)
plt.show()