import numpy as np
import matplotlib.pyplot as plt
from exo_skryer.vert_Tp import isothermal
from exo_skryer.vert_cloud import const_profile
from exo_skryer.data_constants import bar, kb, amu

# Create pressure grid
nlev = 100
p_lev = np.logspace(np.log10(100.0), np.log10(1e-6), nlev) * bar
p_lay = (p_lev[1:] - p_lev[:-1]) / np.log(p_lev[1:] / p_lev[:-1])

# Background state (only needed to satisfy kernel signature)
T_lev, T_lay = isothermal(p_lev, {"T_iso": 1200.0})
mu_lay = np.full_like(T_lay, 2.33)  # amu
nd_lay = p_lay / (kb * T_lay)       # cm^-3
rho_lay = nd_lay * mu_lay * amu     # g cm^-3

q_c_lay = const_profile(
    p_lay, T_lay, mu_lay, rho_lay, nd_lay,
    params={"log_10_q_c": -6.0},
)
q_c_lay = np.asarray(q_c_lay)

# Plot
fig, ax = plt.subplots(figsize=(9, 4.5))
ax.loglog(q_c_lay, p_lay / bar, c="orchid")
ax.set_xlim(1e-12, 1e-5)
ax.set_xlabel(r"$q_c$ [g g$^{-1}$]", fontsize=14)
ax.set_ylabel("pressure [bar]", fontsize=14)
ax.set_title("Constant Cloud Profile", fontsize=12)
ax.invert_yaxis()
plt.tight_layout()