"""f14_virial_range -- Z vs reduced density for the two-term and three-term
virial truncations against a "full" reference (PR isotherm, labelled as
such), showing where truncation breaks down.

Virial coefficients B, C are obtained from the PR EOS itself by series
expansion of Z(rho) at rho->0 (so the truncations are directly comparable
to the "full" PR curve on the same footing): Z = 1 + B*rho + C*rho^2 + ...
B and C are computed by finite-difference differentiation of Z(rho) at
rho=0 using the analytic PR pressure expression, for CO2 at T_r = 1.05.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *

sp = "CO2"
d = SPECIES[sp]
Tc, Pc = d["Tc"], d["Pc"]
Tr = 1.05
T = Tr*Tc

a, b, d1, d2 = cubic_params("pr", Tc, Pc, d["w"], T, Rb)

def Z_of_rho(rho):
    # rho = 1/V (mol/cm3); P = RgT/(V-b) - a/((V+d1 b)(V+d2 b)); Z = PV/(RgT)
    V = 1.0/rho
    P = Rb*T/(V - b) - a/((V + d1*b)*(V + d2*b))
    return P*V/(Rb*T)

# series coefficients via numerical differentiation at rho -> 0
h = 1e-6
Z0 = 1.0  # Z(rho->0) = 1 by construction of any EOS reducing to ideal gas
Zp = Z_of_rho(h)
Zm = Z_of_rho(2*h)
# Z(rho) = 1 + B rho + C rho^2 + ...
# use three points (0 implied, h, 2h) to fit B, C
# Z(h) - 1 = B h + C h^2 ; Z(2h) - 1 = B (2h) + C (2h)^2
A_mat = np.array([[h, h**2], [2*h, 4*h**2]])
rhs = np.array([Zp - 1.0, Zm - 1.0])
B, C = np.linalg.solve(A_mat, rhs)
print(f"f14: virial coefficients from PR series expansion at T_r={Tr}: "
      f"B = {B:.4f} cm3/mol, C = {C:.4f} (cm3/mol)^2")

rho_c = 1.0/(3*b) if False else None  # not used; b is cubic co-volume, not vdW b here directly
# reduced density axis: rho* = rho * b_PR (dimensionless, order-1 near
# the liquid-like density limit b_PR = 1/V_min)
rho = np.linspace(1e-6, 0.9/b, 400)   # up to rho*b = 0.9 (approaching V->b)
rho_star = rho*b

Z_full = np.array([Z_of_rho(r) for r in rho])
Z_2term = 1.0 + B*rho
Z_3term = 1.0 + B*rho + C*rho**2

fig, ax = newfig(6.8, 4.6)

ax.plot(rho_star, Z_full, color=NAVY, lw=2.2, zorder=4)
ax.plot(rho_star, Z_2term, color=AMBER, lw=2.0, linestyle="--", zorder=3)
ax.plot(rho_star, Z_3term, color=TEAL, lw=2.0, linestyle="-.", zorder=3)

# annotation anchors are placed on the visible part of each curve (inside the
# axes limits) and the text blocks sit in the empty band below the full-PR
# minimum and in the clear strip along the top.
ax.annotate("full PR isotherm\n(reference)", xy=(0.47, Z_of_rho(0.47/b)), xytext=(0.55, 0.14),
            color=NAVY, fontsize=10.5, ha="left", va="center", linespacing=1.25,
            arrowprops=dict(arrowstyle="-", color=NAVY, lw=1.0))
ax.annotate("two-term virial\n$Z=1+B\\rho$", xy=(0.215, 1.0+B*0.215/b), xytext=(0.28, 0.10),
            fontsize=10.5, color="#8a6b16", ha="left", va="center", linespacing=1.25,
            arrowprops=dict(arrowstyle="-", color=AMBER, lw=1.0))
ax.annotate("three-term virial\n$Z=1+B\\rho+C\\rho^2$", xy=(0.435, 1.0+B*0.435/b+C*(0.435/b)**2),
            xytext=(0.13, 1.30), fontsize=10.5, color=TEAL, ha="left", va="center",
            linespacing=1.25, arrowprops=dict(arrowstyle="-", color=TEAL, lw=1.0))

# mark where the two-term truncation deviates from the full curve by >5%
dev2 = np.abs(Z_2term - Z_full)/np.abs(Z_full)
idx_break = np.argmax(dev2 > 0.05) if np.any(dev2 > 0.05) else None
if idx_break:
    ax.axvline(rho_star[idx_break], color=RUST, lw=1.3, linestyle=":", zorder=2)
    ax.annotate("truncation breaks down\n(>5% error, two-term)", xy=(rho_star[idx_break], 1.35),
                xytext=(rho_star[idx_break]+0.045, 1.90), fontsize=10.5, color=RUST,
                ha="left", va="top", linespacing=1.25,
                arrowprops=dict(arrowstyle="-", color=RUST, lw=1.0))

ax.set_xlabel(r"Reduced density, $\rho^* = \rho\, b_{PR}$ (dimensionless)")
ax.set_ylabel("Compressibility factor, $Z$")
ax.set_xlim(0, 0.9)
ax.set_ylim(-0.1, 1.95)

save(fig, "f14_virial_range")
