import matplotlib.pyplot as plt
import numpy as np
from kkcalc2.models import asp_db_complex
import XEFI

# Construct atomic scattering polynomial (ASP) models from
# database atomic scattering factors for a few materials
refractive_air = asp_db_complex("(N78O20Ar1)0.01", density=0.001225, name="Air")
refractive_P3HT = asp_db_complex("C10H14S", density=1.33, name="P3HT")
refractive_PS = asp_db_complex("C8H8", density=1.05, name="PS")
refractive_Si = asp_db_complex("Si", density=2.329, name="Si")

# Evaluate the real and imaginary parts of the refractive index for a range of energies
energies = np.linspace(200, 10000, 1000)  # Energy range from 200 eV to 10 keV
n_air = refractive_air.eval_refractive_index(energies)
n_P3HT = refractive_P3HT.eval_refractive_index(energies)
n_PS = refractive_PS.eval_refractive_index(energies)
n_Si = refractive_Si.eval_refractive_index(energies)

# Wavelength / Beam Energy
wav = 1.54  # Å
beam_energy = XEFI.utils.wav2en(wav) # in eV
print(f"{beam_energy:0.2f} eV")

angles = np.linspace(0.1, 0.4, 3000)  # Angles of Incidence in degrees
z = [0, -800, -1340]  # Z-coordinates for the multilayer interface

# Refractive indexes
refractive_indicies: list[kk.models.asp_complex] = [
    refractive_air,
    refractive_PS,
    refractive_P3HT,
    refractive_Si,
]

result = XEFI.XEF_Basic(
    energies=beam_energy,
    angles=angles,
    z=z,
    refractive_indices=refractive_indicies,
    method=XEFI.XEF_method.DEV,
)

fig, ax = result.generate_graphic_XEFI_map()
ax.set_title(
    rf"X-ray Electric Field Intensity at $\lambda$={wav} Å, {beam_energy:0.2f} eV",
    pad=20,
)
plt.show()