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

# Pressure grid (levels → layers)
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])

# Simple isothermal background state
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 = no_cloud(p_lay, T_lay, mu_lay, rho_lay, nd_lay, params={})
q_c_lay = np.asarray(q_c_lay)

fig, ax = plt.subplots(figsize=(9, 4.5))
# Show zero regions on log-x by flooring at a tiny value.
q_floor = 1e-12
q_plot = np.maximum(q_c_lay, q_floor)
ax.loglog(q_plot, p_lay / bar, c="k")
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(r"No Cloud (floored at $10^{-12}$ for log-x)", fontsize=12)
ax.invert_yaxis()
plt.tight_layout()