import numpy as np
import matplotlib.pyplot as plt

from shone.opacity import Opacity

wavelength = np.geomspace(0.5, 5, 500)
pressure = np.geomspace(1e-6, 1)  # [bar]
temperature = 700 * (pressure / 0.1) ** 0.05  # [K]

opacity_samples = []
molecules = ['H2O', 'CO2']
for molecule in molecules:
    opacity = Opacity.load_demo_species(molecule)
    interp_opacity = opacity.get_interpolator()
    opacity_samples.append(
        interp_opacity(wavelength, temperature, pressure)
    )

ax = plt.gca()
for molecule, op in zip(molecules, opacity_samples):
    ax.semilogy(wavelength, op[30], label=molecule.replace('2', '$_2$'))
plt.legend()
ax.set(
    xlabel='Wavelength [µm]',
    ylabel='Opacity [cm$^2$ g$^{-1}$]',
    ylim=(1e-6, 1e5)
)