"""f34_Z_landscape -- filled contour map of Z(T_r, P_r) from PR at
omega = 0 (synthetic Tc=Pc=1 species, matching figA_04/figA_11 convention),
with the Z = 1 contour highlighted in amber and the saturation line drawn.

Root selection (avoids an unphysical discontinuity away from the true
first-order transition): for T_r < 1, at each P_r the saturation pressure
Pr_sat(Tr) is located by the isofugacity condition; for Pr > Pr_sat the
fluid is vapour (largest real root), for Pr < Pr_sat it is liquid
(smallest real root). For T_r >= 1 (supercritical, single phase) the
largest real root is used at every Pr. This makes Z continuous everywhere
except across the true coexistence line, where the physical jump in Z is
expected (first-order phase transition).
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from figB_common import *

SPECIES["_simple"] = dict(Tc=1.0, Pc=1.0, w=0.0)

Trs = np.linspace(0.55, 3.0, 340)
Prs = np.linspace(0.02, 7.0, 340)
TT, PP = np.meshgrid(Trs, Prs)
ZZ = np.full_like(TT, np.nan)

# saturation curve Pr_sat(Tr) for Tr < 1 (synthetic species, Rg=1)
Tr_sat = np.linspace(0.55, 0.9995, 300)
Pr_sat = np.array([Psat_solve("pr", "_simple", Tr, Rg=1.0) for Tr in Tr_sat])  # Pc=1 -> Pr=P

for i in range(TT.shape[0]):
    for j in range(TT.shape[1]):
        Tr, Pr = TT[i, j], PP[i, j]
        r = Zroots("pr", "_simple", Tr, Pr, Rg=1.0)
        if len(r) == 0:
            continue
        if Tr < 1.0:
            ps = np.interp(Tr, Tr_sat, Pr_sat)
            if np.isnan(ps):
                ZZ[i, j] = r[-1]
            elif Pr >= ps:
                ZZ[i, j] = r[-1]     # vapour branch
            else:
                ZZ[i, j] = r[0]      # liquid branch
        else:
            ZZ[i, j] = r[-1]         # supercritical, single physical root

# numerical check: at Tr=0.8 confirm the isofugacity Psat used for the
# root-selection boundary satisfies the Maxwell equal-area condition
Tchk = 0.8
Pchk = Psat_solve("pr", "_simple", Tchk, Rg=1.0)
maxwell_equal_area_check("pr", "_simple", Tchk, Pchk, Rg=1.0)

fig, ax = newfig(7.2, 5.0)
ax.grid(False)

levels = np.linspace(0.0, 2.2, 23)
cf = ax.contourf(TT, PP, ZZ, levels=levels, cmap="RdBu_r", extend="max")
cbar = fig.colorbar(cf, ax=ax, pad=0.02)
cbar.set_label(r"$Z$", fontsize=11, color=INK)
cbar.ax.tick_params(labelsize=9.5, color=GRAPHITE)

# Z = 1 contour, highlighted amber
ax.contour(TT, PP, ZZ, levels=[1.0], colors=[AMBER], linewidths=2.4)

# saturation line
ax.plot(Tr_sat, Pr_sat, color=GRAPHITE, lw=2.2, zorder=5)
ax.plot([1.0], [1.0], marker="o", ms=7, color=GRAPHITE, zorder=6,
         markeredgecolor="white", markeredgewidth=1.0)
ax.annotate("critical point\n" r"$(T_r,P_r)=(1,1)$", xy=(1.0, 1.0), xytext=(1.55, 1.55),
            fontsize=10.5, color=GRAPHITE, linespacing=1.3,
            arrowprops=dict(arrowstyle="-", color=GRAPHITE, lw=1.0))
# moved off the dark end of the colour fill into the pale region right of the
# critical point, joined to the coexistence curve by a leader
ax.annotate("saturation line\n(liquid-vapour coexistence)",
            xy=(0.85, np.interp(0.85, Tr_sat, Pr_sat)),
            xytext=(1.22, 0.40), fontsize=10.5, color=GRAPHITE, linespacing=1.3,
            ha="left", va="bottom",
            arrowprops=dict(arrowstyle="-", color=GRAPHITE, lw=1.0))

# tip of the Z=1 contour on the compressed-liquid (top-left) branch, read off
# the plotted grid so the leader lands exactly on the drawn contour; the label
# itself is pulled clear of both the contour and the left spine
row = int(np.argmin(np.abs(Prs - 6.45)))
zr = ZZ[row, :]
j = int(np.where(np.isfinite(zr) & (zr < 1.0))[0][0])
Tr_tip = np.interp(1.0, [zr[j], zr[j - 1]], [Trs[j], Trs[j - 1]])
ax.annotate("compressed-liquid\nbranch, $Z>1$", xy=(Tr_tip, Prs[row]),
            xytext=(0.795, 6.90), fontsize=10.5, color=AMBER,
            ha="left", va="top", linespacing=1.25,
            arrowprops=dict(arrowstyle="-", color=AMBER, lw=1.0))
from scipy.optimize import brentq as _brentq
Tr_pt = 2.3
Pr_on_contour = _brentq(lambda Pr: Zroots("pr", "_simple", Tr_pt, Pr, Rg=1.0)[-1] - 1.0, 0.05, 6.9)
# label moved to the empty pale band on the low-T side of the contour (it used
# to sit astride the contour and ran into the colourbar)
ax.annotate("$Z=1$ contour\n(supercritical branch)", xy=(Tr_pt, Pr_on_contour),
            xytext=(1.45, 5.80), fontsize=10.5, color=AMBER, linespacing=1.3,
            ha="left", va="center",
            arrowprops=dict(arrowstyle="-", color=AMBER, lw=1.0))

ax.set_xlabel(r"Reduced temperature, $T_r$")
ax.set_ylabel(r"Reduced pressure, $P_r$")
ax.set_xlim(0.55, 3.0)
ax.set_ylim(0.02, 7.0)

save(fig, "f34_Z_landscape")
