"""f03_ideal_vs_real_Z -- Z vs P (0-400 bar) at T = 300 K for methane, CO2,
ethane and nitrogen, computed with Peng-Robinson (SPEC.md PR constants).
Navy dashed Z = 1 line labelled "ideal gas".
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *

fig, ax = newfig()

T = 300.0
P = np.linspace(2, 400, 300)
species_colors = {
    "methane": AMBER,
    "CO2": RUST,
    "ethane": TEAL,
    "nitrogen": SKYBLUE,
}

for sp, col in species_colors.items():
    Z = np.array([Zroots("pr", sp, T, p)[-1] for p in P])
    ax.plot(P, Z, color=col, lw=2.0)

# Direct labels in a clear strip to the right of the data (the x axis is
# extended past the last data point to open that strip). Methane and ethane
# end only 0.03 apart in Z, so their labels are pulled apart vertically and
# tied back to the curve ends with short leader lines.
label_y = {"methane": 1.062, "CO2": 0.6765, "ethane": 0.898, "nitrogen": 1.179}
for sp, col in species_colors.items():
    Zend = Zroots("pr", sp, T, 400)[-1]
    yl = label_y[sp]
    ax.annotate("", xy=(401, Zend), xytext=(424, yl),
                arrowprops=dict(arrowstyle="-", color=col, lw=1.1,
                                shrinkA=0, shrinkB=0), zorder=4)
    ax.text(429, yl, sp, color=col, fontsize=10.5, ha="left", va="center",
            weight="bold", zorder=5)

ax.axhline(1.0, color=NAVY, lw=1.8, linestyle="--", zorder=3)
ax.text(15, 1.014, "ideal gas, $Z=1$", color=NAVY, fontsize=10.5, va="bottom")

ax.set_xlim(0, 505)
ax.set_xticks(np.arange(0, 401, 50))
ax.set_xlabel("Pressure, $P$ / bar")
ax.set_ylabel("Compressibility factor, $Z = PV/RT$")
ax.set_title(None)

save(fig, "f03_ideal_vs_real_Z")
