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

# 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)

fig,ax = plt.subplots(2,1, figsize=(10,8), sharex=True)
ax[0].plot(energies, n_air.real, label="Air")
ax[0].plot(energies, n_P3HT.real, label="P3HT")
ax[0].plot(energies, n_PS.real, label="PS")
ax[0].plot(energies, n_Si.real, label="Si")
ax[0].set_ylabel("Real part of refractive index")
ax[0].legend()
ax[1].plot(energies, n_air.imag, label="Air")
ax[1].plot(energies, n_P3HT.imag, label="P3HT")
ax[1].plot(energies, n_PS.imag, label="PS")
ax[1].plot(energies, n_Si.imag, label="Si")
ax[1].set_xlabel("Energy (eV)")
ax[1].set_ylabel("Imaginary part of refractive index")
ax[1].legend()
ax[1].set_yscale("log")
ax[1].set_xscale("log")
ax[0].patch.set_alpha(0)
ax[1].patch.set_alpha(0)
ax[0].patch.set_facecolor('none')
ax[1].patch.set_facecolor('none')
fig.tight_layout()
plt.show()