"""f18_omega_vs_shape -- horizontal bar chart of the Pitzer acentric factor
omega for twelve species, grouped and coloured by molecular family
(spherical, chain, polar, associating).

Values for methane, nitrogen, ethane, propane, n-butane, CO2, ammonia,
water, benzene are the SPEC.md section-8 table values (fixed everywhere in
this package). Argon, n-octane and methanol are not in that table; their
omega values are taken from the standard literature compilation (Poling,
Prausnitz & O'Connell, "The Properties of Gases and Liquids", 5th ed.,
2001, Appendix A, SPEC.md reference list): argon 0.000, n-octane 0.395,
methanol 0.556. These three are used only for this comparison chart, never
elsewhere in the package.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *

data = [
    ("argon",    0.000, "spherical"),
    ("methane",  SPECIES["methane"]["w"], "spherical"),
    ("nitrogen", SPECIES["nitrogen"]["w"], "spherical"),
    ("ethane",   SPECIES["ethane"]["w"], "chain"),
    ("propane",  SPECIES["propane"]["w"], "chain"),
    ("n-butane", SPECIES["n-butane"]["w"], "chain"),
    ("n-octane", 0.395, "chain"),
    ("benzene",  SPECIES["benzene"]["w"], "polar"),
    ("CO2",      SPECIES["CO2"]["w"], "polar"),
    ("ammonia",  SPECIES["ammonia"]["w"], "associating"),
    ("water",    SPECIES["water"]["w"], "associating"),
    ("methanol", 0.556, "associating"),
]

fam_color = {"spherical": NAVY, "chain": TEAL, "polar": SKYBLUE, "associating": RUST}
fam_label = {"spherical": "spherical (near-zero ω)", "chain": "chain / elongated",
             "polar": "polar, non-associating", "associating": "hydrogen-bonding / associating"}

# group and order by family, then by increasing omega within a family
order = ["spherical", "chain", "polar", "associating"]
rows = []
for fam in order:
    items = sorted([d for d in data if d[2] == fam], key=lambda t: t[1])
    rows.extend(items)

names = [r[0] for r in rows]
omegas = [r[1] for r in rows]
cols = [fam_color[r[2]] for r in rows]

fig, ax = newfig(7.0, 6.0)
ax.grid(axis="x", zorder=0)
ax.grid(axis="y", visible=False)

y = np.arange(len(rows))[::-1]
ax.barh(y, omegas, height=0.62, color=cols, zorder=3, edgecolor="white", linewidth=0.6)

for yi, om in zip(y, omegas):
    ax.text(om + 0.012, yi, f"{om:.3f}", va="center", ha="left", fontsize=10.5, color=INK)

ax.set_yticks(y)
ax.set_yticklabels(names, fontsize=10.5)
ax.set_xlabel(r"Pitzer acentric factor, $\omega = -1-\log_{10}P_r^{sat}(T_r=0.7)$")
ax.set_xlim(-0.03, 0.66)
ax.set_ylim(-0.8, len(rows) + 0.9)

# family separators + labels
fam_bounds = {}
idx = 0
for fam in order:
    n = sum(1 for r in rows if r[2] == fam)
    fam_bounds[fam] = (idx, idx + n - 1)
    idx += n
for fam in order:
    lo, hi = fam_bounds[fam]
    ylo, yhi = y[hi], y[lo]
    if lo > 0:
        ax.axhline(y[lo] + 0.5, color=MIST, lw=0.9, zorder=1)
    ax.text(0.615, (ylo + yhi) / 2, fam_label[fam], color=fam_color[fam],
            fontsize=10.5, ha="left", va="center", rotation=0, weight="bold",
            linespacing=1.2)

ax.text(0.0, len(rows) + 0.75, "spherical, non-polar molecules cluster near $\\omega\\approx 0$;\n"
        "chain length and polarity/association push $\\omega$ up",
        ha="left", va="top", fontsize=10.5, color=SLATE, linespacing=1.35)

save(fig, "f18_omega_vs_shape")
