"""f06_vdw_isotherms -- P-V isotherms for CO2 from vdW at T_r = 0.85, 0.90,
0.95, 1.00, 1.10; sand-shaded two-phase dome from the Maxwell construction;
critical point marked.

Maxwell equal-area construction is solved numerically for each subcritical
isotherm: for given T < Tc, find P_sat such that the two lobes of
P_vdw(V) - P_sat integrate to zero area between the liquid and vapour
roots (standard Maxwell criterion), using SPEC.md vdW a, b formulas.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *
from scipy.optimize import brentq
from scipy.integrate import quad

sp = "CO2"
d = SPECIES[sp]
Tc, Pc = d["Tc"], d["Pc"]
a, b, _, _ = cubic_params("vdw", Tc, Pc, d["w"], Tc, Rb)  # a is T-independent for vdW
Vc = 3*b

def Pvdw(V, T):
    return Rb*T/(V - b) - a/V**2


def three_roots_V(T, P):
    """Solve cubic in V for given T,P; return sorted real roots > b."""
    # Rg T V^3 ... rewrite as: P V^3 - (Pb+RgT) V^2 + a V - a b = 0
    coeffs = [P, -(P*b + Rb*T), a, -a*b]
    r = np.roots(coeffs)
    r = np.sort(np.real(r[np.abs(np.imag(r)) < 1e-7*max(1,abs(r).max())]))
    return r[r > b*1.001]


def maxwell_Psat(T):
    """Find P such that equal-area (Maxwell) construction holds for vdW at T<Tc."""
    def area_diff(P):
        roots = three_roots_V(T, P)
        if len(roots) < 3:
            return np.nan
        Vl, Vm, Vv = roots
        # integral of (P_vdw(V) - P) dV from Vl to Vv should be zero
        val, _ = quad(lambda V: Pvdw(V, T) - P, Vl, Vv, limit=200)
        return val
    # bracket search between small and near-critical pressure
    Plo, Phi = 0.001*Pc, Pc*0.999
    # scan for sign change
    Ps = np.linspace(Plo, Phi, 400)
    vals = []
    for P in Ps:
        v = area_diff(P)
        vals.append(v)
    vals = np.array(vals)
    idx = np.where(np.isfinite(vals))[0]
    Ps, vals = Ps[idx], vals[idx]
    sign_changes = np.where(np.diff(np.sign(vals)) != 0)[0]
    if len(sign_changes) == 0:
        return None
    i = sign_changes[0]
    return brentq(area_diff, Ps[i], Ps[i+1], xtol=1e-6)


fig, ax = newfig(6.8, 4.6)

Trs = [0.85, 0.90, 0.95, 1.00, 1.10]
# Each isotherm is drawn out to its own V, staggered by 80 cm3/mol, so that the
# five curve ends form a diagonal staircase in the empty right-hand strip and
# each label can sit directly at the end of its own curve. (At a common V the
# five branches are only 2-4 bar apart, which is far too close to label.)
VEND = {0.85: 705.0, 0.90: 625.0, 0.95: 545.0, 1.00: 465.0, 1.10: 385.0}
VMAXC = max(VEND.values())
Vgrid = np.linspace(1.01*b, VMAXC, 2400)

dome_liq, dome_vap, dome_T = [], [], []
for Tr in Trs:
    T = Tr*Tc
    Vend = VEND[Tr]
    if Tr < 1.0:
        Psat = maxwell_Psat(T)
        roots = three_roots_V(T, Psat)
        Vl, Vm, Vv = roots
        dome_liq.append(Vl); dome_vap.append(Vv); dome_T.append(T)
        # plot: stable liquid branch, unstable+metastable suppressed here (full loop in f07)
        Vleft = Vgrid[(Vgrid > 1.01*b) & (Vgrid <= Vl)]
        Vright = Vgrid[(Vgrid >= Vv) & (Vgrid <= Vend)]
        Pleft = Pvdw(Vleft, T)
        Pright = Pvdw(Vright, T)
        Vloop = Vgrid[(Vgrid >= Vl) & (Vgrid <= Vv)]
        Ploop = Pvdw(Vloop, T)
        ax.plot(Vleft, Pleft, color=NAVY, lw=1.9)
        ax.plot(Vright, Pright, color=NAVY, lw=1.9)
        ax.plot(Vloop, Ploop, color=NAVY, lw=1.1, linestyle=":", alpha=0.55)
        ax.plot([Vl, Vv], [Psat, Psat], color=SLATE, lw=1.1, linestyle="--", alpha=0.6)
    else:
        mask = Vgrid <= Vend
        P = Pvdw(Vgrid[mask], T)
        ax.plot(Vgrid[mask], P, color=NAVY, lw=1.9)

# critical point
Pc_check = Pvdw(Vc, Tc)
ax.plot([Vc], [Pc], marker="o", ms=7, color=AMBER, zorder=6, markeredgecolor="white", markeredgewidth=1.0)
ax.annotate("critical point", xy=(Vc, Pc), xytext=(Vc*1.55, Pc*1.18),
            fontsize=10.5, color=AMBER, arrowprops=dict(arrowstyle="-", color=AMBER, lw=1.1))

# two-phase dome (sand shading) using Maxwell points across subcritical Trs
dome_T = np.array(dome_T); dome_liq = np.array(dome_liq); dome_vap = np.array(dome_vap)
order = np.argsort(dome_T)
Vdome = np.concatenate([dome_liq[order], dome_vap[order][::-1]])
Pdome = np.concatenate([Pc*np.array(Trs[:len(dome_T)])[order]*0 + 0, []])  # placeholder unused
Psat_arr = np.array([Pc]*0)
# build dome polygon using actual Psat at each subcritical T (recompute to keep Psat values)
Psats = []
for T in dome_T[order]:
    Psats.append(maxwell_Psat(T))
Psats = np.array(Psats)
Vpoly = np.concatenate([dome_liq[order], [Vc], dome_vap[order][::-1]])
Ppoly = np.concatenate([Psats, [Pc], Psats[::-1]])
ax.fill(Vpoly, Ppoly, color=SAND, alpha=0.55, zorder=1, edgecolor="none")

# direct label at the end of each isotherm: text starts just past the last
# point of its own curve and sits above it, so nothing else is drawn in that
# corner (every hotter isotherm has already stopped further to the left).
for Tr in Trs:
    Vend = VEND[Tr]
    Pend = Pvdw(Vend, Tr*Tc)
    ax.text(Vend + 9, Pend + 1.2, f"$T_r$={Tr:.2f}", fontsize=10.5, color=NAVY,
            va="bottom", ha="left")

ax.annotate("two-phase\nregion", xy=(350.0, 40.0), xytext=(425.0, 15.0),
            fontsize=10.5, color="#8a7328", ha="left", va="center", linespacing=1.3,
            arrowprops=dict(arrowstyle="-", color="#8a7328", lw=1.0, alpha=0.9))

ax.set_xlim(1.01*b, 800)
ax.set_ylim(0, 130)
ax.set_xlabel(r"Molar volume, $V$ / (cm$^3$ mol$^{-1}$)")
ax.set_ylabel("Pressure, $P$ / bar")

fig.subplots_adjust(right=0.97)
save(fig, "f06_vdw_isotherms")
