"""f20_alpha_functions -- a(T)/a_c vs T_r for vdW (constant = 1), RK
(T_r^-0.5), SRK and PR at omega = 0.15 (solid), plus SRK and PR at
omega = 0.35 (dashed). All formulas from SPEC.md section 8:
  vdW:  a(T)/a_c = 1
  RK:   a(T)/a_c = T_r^-1/2   (since a(T) = 0.42748 R^2 Tc^2.5/Pc / sqrt(T)
        and a_c = a(Tc), a(T)/a_c = sqrt(Tc/T) = Tr^-1/2)
  SRK:  alpha = [1 + m(1-sqrt(Tr))]^2,  m = 0.480 + 1.574*w - 0.176*w^2
  PR:   alpha = [1 + k(1-sqrt(Tr))]^2,  k = 0.37464 + 1.54226*w - 0.26992*w^2
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *

Tr = np.linspace(0.4, 2.5, 300)

def alpha_srk(Tr, w):
    m = 0.480 + 1.574*w - 0.176*w**2
    return (1 + m*(1 - np.sqrt(Tr)))**2

def alpha_pr(Tr, w):
    k = 0.37464 + 1.54226*w - 0.26992*w**2
    return (1 + k*(1 - np.sqrt(Tr)))**2

LBL = 12.5      # annotation size: at 10.5 the mathtext superscripts were
                # unreadable on a projected slide
fig, ax = newfig(7.2, 4.8)

ax.plot(Tr, np.ones_like(Tr), color=NAVY, lw=2.0, label="vdW")
ax.plot(Tr, Tr**-0.5, color=SKYBLUE, lw=2.0, label="RK")
ax.plot(Tr, alpha_srk(Tr, 0.15), color=TEAL, lw=2.0, linestyle="-")
ax.plot(Tr, alpha_pr(Tr, 0.15), color=AMBER, lw=2.0, linestyle="-")
ax.plot(Tr, alpha_srk(Tr, 0.35), color=TEAL, lw=1.8, linestyle="--")
ax.plot(Tr, alpha_pr(Tr, 0.35), color=AMBER, lw=1.8, linestyle="--")

ax.axvline(1.0, color=GRAPHITE, lw=0.9, linestyle=":", zorder=1)
ax.plot([1.0], [1.0], marker="o", ms=6, color=GRAPHITE, zorder=5,
         markeredgecolor="white", markeredgewidth=0.8)
ax.text(1.06, 2.58, r"$T_r=1$: all $\alpha=a/a_c=1$ by construction", fontsize=LBL,
        color=GRAPHITE, ha="left", va="top")

# direct labels: staggered leader labels at the right edge (x=2.5) since
# several curves converge there, plus two labels placed in open space
ax.annotate("vdW: $a/a_c=1$ (no $T$-dependence)", xy=(2.3, 1.0), xytext=(2.3, 1.40),
            fontsize=LBL, color=NAVY, ha="center", va="bottom", linespacing=1.3,
            arrowprops=dict(arrowstyle="-", color=NAVY, lw=0.9, alpha=0.85))

x_end = 2.5
end_curves = [
    ("SRK, $\\omega=0.15$", TEAL, alpha_srk(x_end, 0.15)),
    ("PR, $\\omega=0.15$", AMBER, alpha_pr(x_end, 0.15)),
    ("RK: $T_r^{-1/2}$", SKYBLUE, x_end**-0.5),
    ("SRK, $\\omega=0.35$ (dashed)", TEAL, alpha_srk(x_end, 0.35)),
    ("PR, $\\omega=0.35$ (dashed)", AMBER, alpha_pr(x_end, 0.35)),
]
end_curves.sort(key=lambda t: -t[2])
# the five curve ends span only ~0.46 in alpha, far too little room for five
# stacked labels, so they are fanned out on a fixed generous pitch in the
# empty right-hand gutter and each joined to its own curve end by a leader
# (the fan preserves the curve order, so the leaders never cross)
Y_TOP, Y_PITCH, X_LAB = 1.10, 0.24, x_end + 0.14
for i, (lbl, col, y) in enumerate(end_curves):
    yt = Y_TOP - i * Y_PITCH
    ax.plot([x_end, X_LAB - 0.04], [y, yt], color=col, lw=0.7, alpha=0.6)
    ax.text(X_LAB, yt, lbl, color=col, fontsize=LBL, va="center", ha="left")

ax.set_xlabel(r"Reduced temperature, $T_r$")
ax.set_ylabel(r"$\alpha(T_r,\omega) = a(T)/a_c$")
ax.set_xlim(0.4, 4.12)
ax.set_xticks([0.5, 1.0, 1.5, 2.0, 2.5])   # ticks pinned to the data range
ax.set_ylim(0.0, 2.7)
fig.subplots_adjust(right=0.85)

save(fig, "f20_alpha_functions")
