"""f19_rk_vs_vdw -- Z vs P_r for propane at T_r = 1.2 (single-phase,
supercritical) and T_r = 0.9 (subcritical vapour branch only, up to the
saturation pressure) computed with vdW, RK and PR (SPEC.md section 8
constants). PR is labelled as the modern reference.

For T_r = 0.9 the vapour-branch upper bound is the model's own isofugacity
Psat (figB_common.Psat_solve); one value is checked against an independent
Maxwell equal-area integration (figB_common.maxwell_equal_area_check).
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from figB_common import *

sp = "propane"
d = SPECIES[sp]
Pc = d["Pc"]

models = ["vdw", "rk", "pr"]
mcolor = {"vdw": NAVY, "rk": SKYBLUE, "pr": AMBER}
# PR set over two lines so the label stays inside the axes instead of
# running off the right spine
mname = {"vdw": "vdW (1873)", "rk": "RK (1949)",
         "pr": "PR (1976),\nmodern reference"}

fig, ax = newfig(7.2, 4.8)

# ---- Tr = 1.2: supercritical, single vapour-like root over the whole range
Tr_hi = 1.2
T_hi = Tr_hi * d["Tc"]
Pr_hi = np.linspace(0.02, 3.0, 250)
for m in models:
    Z = []
    for Pr in Pr_hi:
        r = Zroots(m, sp, T_hi, Pr * Pc)
        Z.append(r[-1] if len(r) else np.nan)
    ax.plot(Pr_hi, Z, color=mcolor[m], lw=2.0, linestyle="-")

# ---- Tr = 0.9: subcritical vapour branch, up to each model's own Psat
Tr_lo = 0.9
T_lo = Tr_lo * d["Tc"]
psat_check = {}
for m in models:
    Psat = Psat_solve(m, sp, T_lo)
    psat_check[m] = Psat
    Pr_lo = np.linspace(0.02, 0.985 * Psat / Pc, 200)
    Z = []
    for Pr in Pr_lo:
        r = Zroots(m, sp, T_lo, Pr * Pc)
        Z.append(r[-1] if len(r) else np.nan)
    ax.plot(Pr_lo, Z, color=mcolor[m], lw=2.0, linestyle="--")

# numerical check: PR Maxwell equal-area at Tr = 0.9
maxwell_equal_area_check("pr", sp, T_lo, psat_check["pr"])
maxwell_equal_area_check("vdw", sp, T_lo, psat_check["vdw"])
for m in models:
    print(f"[info] {m.upper()} propane Psat(Tr=0.9) = {psat_check[m]:.3f} bar "
          f"(Pr_sat = {psat_check[m]/Pc:.4f})")

ax.axhline(1.0, color=GRAPHITE, lw=1.0, linestyle=":", zorder=1)
ax.text(0.06, 1.015, r"$T_r=1.2$ (single phase, solid)", fontsize=10.5,
        color=GRAPHITE, ha="left", va="bottom")

# --- right-edge direct labels for the Tr=1.2 curves, staggered so they
# do not overlap (the three curves converge near Pr=3)
x_end = 3.0
ends12 = []
for m in models:
    r = Zroots(m, sp, T_hi, x_end * Pc)
    ends12.append((m, r[-1]))
ends12.sort(key=lambda t: -t[1])
min_gap = 0.075                     # ~17 pt pitch at 10.5 pt text
ytexts = []
for i, (m, z) in enumerate(ends12):
    extra = 0.055 * mname[m].count("\n")   # room for the two-line PR label
    yt = z if i == 0 else min(z, ytexts[-1] - min_gap - extra)
    ytexts.append(yt)
for (m, z), yt in zip(ends12, ytexts):
    col = mcolor[m]
    ax.plot([x_end, x_end + 0.10], [z, yt], color=col, lw=0.8, alpha=0.7)
    ax.text(x_end + 0.14, yt, mname[m], color=col, fontsize=10.5, va="center",
            ha="left", linespacing=1.3)

# --- labels for the Tr=0.9 vapour branches, staggered near each Psat_r end
ends09 = []
for m in models:
    Pr_end = 0.985 * psat_check[m] / Pc
    r = Zroots(m, sp, T_lo, Pr_end * Pc)
    ends09.append((m, Pr_end, r[-1]))
ends09.sort(key=lambda t: -t[2])
# the Tr = 1.2 curves pass straight through the region where these branches
# terminate, so the labels are dropped into the empty band below them and
# joined to the true end point with a leader
X_LAB, Y_TOP, Y_PITCH = 1.05, 0.45, 0.075
for i, (m, x, z) in enumerate(ends09):
    col = mcolor[m]
    yt = Y_TOP - i * Y_PITCH
    ax.plot([x, X_LAB - 0.03], [z, yt], color=col, lw=0.9, alpha=0.75)
    ax.text(X_LAB, yt, m.upper(), color=col, fontsize=10.5, va="center", ha="left")

ax.text(0.06, 0.10, r"$T_r=0.9$ vapour branch (dashed), each"
        "\nending at that model's own $P_r^{sat}$", fontsize=10.5,
        color=GRAPHITE, ha="left", va="bottom", linespacing=1.3)

ax.set_xlabel(r"Reduced pressure, $P_r$")
ax.set_ylabel(r"Compressibility factor, $Z$")
ax.set_xlim(0, 4.35)
ax.set_xticks(np.arange(0.0, 3.01, 0.5))   # ticks pinned to the data range
ax.set_ylim(0.0, 1.06)
fig.subplots_adjust(right=0.85)

save(fig, "f19_rk_vs_vdw")
