import numpy as np
import matplotlib.pyplot as plt
from exo_skryer.vert_Tp import Modified_Milne
from exo_skryer.vert_chem import constant_vmr
from exo_skryer.data_constants import bar

nlev = 100
p_bot = np.log10(100.0)
p_top = np.log10(1e-4)
p_lev = np.logspace(p_bot, p_top, nlev) * bar

params_tp = {
    "T_int": 1200.0,
    "T_ratio": 0.333,
    "log_10_g": 4.5,
    "log_10_k_ir": -2.0,
    "log_10_p_t": 0.0,
    "beta": 0.55
}
T_lev, T_lay = Modified_Milne(p_lev, params_tp)
p_lay = (p_lev[1:] - p_lev[:-1]) / np.log(p_lev[1:] / p_lev[:-1])

species_order = ("H2O", "CO", "CO2", "CH4")
chem_kernel = constant_vmr(species_order)
params = {
    "log_10_f_H2O": -3.0,
    "log_10_f_CO": -4.0,
    "log_10_f_CO2": -8.0,
    "log_10_f_CH4": -6.0,
}
vmr_lay = chem_kernel(p_lay, T_lay, params, nlev - 1)

fig, ax = plt.subplots(figsize=(10, 5))
for key in ("H2O", "CO", "CO2", "CH4"):
    if key in vmr_lay:
        ax.semilogy(vmr_lay[key], p_lay / bar, label=key)
ax.set_xlabel("VMR", fontsize=16)
ax.set_ylabel("pressure [bar]", fontsize=16)
ax.set_title("Constant VMR Chemistry", fontsize=14)
ax.legend(fontsize=10)
ax.set_xscale('log')
ax.invert_yaxis()
plt.tight_layout()