"""f15_corresponding_states -- Z vs P_r at T_r = 1.10 for methane, nitrogen,
ethane (simple, near-collapse) and water, ammonia (clear departure); all
computed with PR (using each species' own omega) so the comparison is
like-for-like; annotate that departure grows with omega.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *

Tr = 1.10
species_list = ["methane", "nitrogen", "ethane", "ammonia", "water"]
colors = {"methane": NAVY, "nitrogen": SKYBLUE, "ethane": TEAL, "ammonia": "#c9812f", "water": RUST}

Pr = np.linspace(0.02, 6.0, 300)

fig, ax = newfig(7.0, 4.6)

ends = []
for sp in species_list:
    d = SPECIES[sp]
    T = Tr*d["Tc"]
    Z = []
    for pr in Pr:
        P = pr*d["Pc"]
        r = Zroots("pr", sp, T, P)
        Z.append(r[-1] if len(r) else np.nan)
    Z = np.array(Z)
    ax.plot(Pr, Z, color=colors[sp], lw=2.1)
    ends.append((sp, Z[-1]))

ax.plot([0.0, 6.0], [1.0, 1.0], color=GRAPHITE, lw=1.1, linestyle=":", zorder=2)
ax.text(0.1, 1.012, "$Z=1$", color=GRAPHITE, fontsize=10.5, va="bottom")

# the leader ends on the visible spread between the water and methane
# isotherms, which is what the sentence is pointing at.
ax.annotate("departure from the simple-fluid\ncurve grows with $\\omega$", xy=(2.15, 0.435),
            xytext=(1.30, 0.29), fontsize=10.5, color=INK, ha="left", va="center",
            linespacing=1.3, arrowprops=dict(arrowstyle="-", color=GRAPHITE, lw=1.0))

ax.set_xlabel(r"Reduced pressure, $P_r$")
ax.set_ylabel(r"Compressibility factor, $Z$")
# P_r > 6 carries no data and is used as a gutter for the direct labels
ax.set_xlim(0, 8.8)
ax.set_ylim(0.2, 1.15)
ax.set_xticks([0, 1, 2, 3, 4, 5, 6])

# direct labels staggered at the right edge, with a minimum pitch set from
# the real axes height so that adjacent lines cannot touch
LABEL_FS = 10.5
PITCH_PT = 15.0
fig.tight_layout()
fig.canvas.draw()
ax_h_pt = ax.get_window_extent().height * 72.0 / fig.dpi
y0, y1 = ax.get_ylim()
min_gap = PITCH_PT * (y1 - y0) / ax_h_pt

ends_sorted = sorted(ends, key=lambda t: -t[1])
ytexts = []
for i, (sp, Z) in enumerate(ends_sorted):
    yt = Z if i == 0 else min(Z, ytexts[-1] - min_gap)
    ytexts.append(yt)
for (sp, Z), yt in zip(ends_sorted, ytexts):
    col = colors[sp]
    w = SPECIES[sp]["w"]
    ax.plot([6.02, 6.32], [Z, yt], color=col, lw=0.9, alpha=0.85, zorder=3)
    ax.text(6.40, yt, f"{sp} ($\\omega$={w:.3f})", color=col, fontsize=LABEL_FS,
            va="center", ha="left", zorder=6,
            bbox=dict(facecolor="white", edgecolor="none", alpha=0.9, pad=1.5))

save(fig, "f15_corresponding_states")
