"""f16_psat_reduced -- log10(P_r^sat) vs 1/T_r from the PR EOS for methane,
propane, n-butane, water; near-linear lines with different slopes, all
passing near (1, 0).

P_sat(T) is found for each species/T by solving the PR EOS for equal
fugacity of the liquid and vapour roots (standard isofugacity condition),
using SPEC.md PR constants.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *
from scipy.optimize import brentq

def ln_phi(model, sp, T, P, Z, Rg=Rb):
    """Fugacity coefficient (ln) for a cubic EOS root Z, PR form."""
    d = SPECIES[sp]
    a, b, d1, d2 = cubic_params(model, d["Tc"], d["Pc"], d["w"], T, Rg)
    A = a*P/(Rg*T)**2
    B = b*P/(Rg*T)
    # PR: ln phi = Z-1 - ln(Z-B) - A/(B*(d1-d2)) * ln((Z+d1*B)/(Z+d2*B))
    term = A/(B*(d1 - d2)) * np.log((Z + d1*B)/(Z + d2*B))
    return (Z - 1) - np.log(Z - B) - term

def Psat_PR(sp, T):
    d = SPECIES[sp]
    def f(P):
        r = Zroots("pr", sp, T, P)
        if len(r) < 2:
            return np.nan
        Zl, Zv = r[0], r[-1]
        return ln_phi("pr", sp, T, P, Zv) - ln_phi("pr", sp, T, P, Zl)
    # bracket search between small P and Pc (log-spaced so very small Psat
    # at low Tr, e.g. for water/propane far below Tc, is still bracketed)
    Ps = np.geomspace(1e-7*d["Pc"], 0.999*d["Pc"], 400)
    vals = np.array([f(P) for P in Ps])
    idx = np.isfinite(vals)
    Ps_f, vals_f = Ps[idx], vals[idx]
    sc = np.where(np.diff(np.sign(vals_f)) != 0)[0]
    if len(sc) == 0:
        return np.nan
    i = sc[0]
    return brentq(f, Ps_f[i], Ps_f[i+1], xtol=1e-8)

species_list = ["methane", "propane", "n-butane", "water"]
colors = {"methane": NAVY, "propane": TEAL, "n-butane": SKYBLUE, "water": RUST}

fig, ax = newfig(6.8, 4.6)

ends = []
for sp in species_list:
    d = SPECIES[sp]
    Trs = np.linspace(0.5, 0.99, 60)
    invTr, logPr = [], []
    for Tr in Trs:
        T = Tr*d["Tc"]
        Ps = Psat_PR(sp, T)
        if np.isnan(Ps):
            continue
        invTr.append(1.0/Tr)
        logPr.append(np.log10(Ps/d["Pc"]))
    invTr = np.array(invTr); logPr = np.array(logPr)
    ax.plot(invTr, logPr, color=colors[sp], lw=2.1)
    ends.append((sp, invTr[0], logPr[0]))

# direct labels, staggered vertically at the cold (right) end of each line
ends_sorted = sorted(ends, key=lambda t: -t[2])
min_gap = 0.30
ytexts = []
for i, (sp, x, y) in enumerate(ends_sorted):
    yt = y if i == 0 else min(y, ytexts[-1] - min_gap)
    ytexts.append(yt)
for (sp, x, y), yt in zip(ends_sorted, ytexts):
    col = colors[sp]
    if abs(yt - y) > 1e-6:
        ax.plot([x, x + 0.03], [y, yt], color=col, lw=0.8, alpha=0.7)
    ax.text(x + 0.045, yt, sp, color=col, fontsize=10.5, ha="left", va="center", zorder=6,
            bbox=dict(facecolor="white", edgecolor="none", alpha=0.9, pad=1.5))

ax.plot([1.0], [0.0], marker="o", ms=6, color=GRAPHITE, zorder=5, markeredgecolor="white", markeredgewidth=0.8)
ax.annotate("all curves pass near\n$(1/T_r,\\log_{10}P_r^{sat})=(1,0)$\n(definition of $T_c,P_c$)",
            xy=(1.0, 0.0), xytext=(1.12, 0.55), fontsize=10.5, color=GRAPHITE, linespacing=1.25,
            arrowprops=dict(arrowstyle="-", color=GRAPHITE, lw=1.0))

ax.set_xlabel(r"$1/T_r$")
ax.set_ylabel(r"$\log_{10}\!\left(P_r^{sat}\right)$")
ax.set_xlim(1.0, 2.45)
ax.set_ylim(-4.6, 0.6)

save(fig, "f16_psat_reduced")
