"""f25_saft_terms -- four-panel schematic of the SAFT construction:
(1) hard-sphere fluid, (2) chain formation by bonding segments,
(3) dispersion attraction between segments, (4) association through
short-range directional sites A and B. Amber highlights the physics added
at each step, left to right reading order. Pure schematic (matplotlib
primitives only, no data).

Layout note: the canvas is deliberately kept near 16:9 and fairly small
(9.4 x 5.3 in) so that the annotation point sizes below are large relative
to the figure width -- the figure is projected about 6 in wide, so what
matters for legibility is the ratio of font size to canvas width, not the
absolute point size. The stage-to-stage arrows run along a clear band
(y = 0.42) below the sphere diagrams so they never overlap the segments.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *
from matplotlib.patches import Ellipse, FancyArrowPatch, FancyBboxPatch

W, H = 9.4, 5.3
YMAX = 1.12
fig, ax = blank(W, H)
ax.set_ylim(0, YMAX)
ASPECT = (W / H) * YMAX  # true-circle correction: x-span=1 over W in, y-span=YMAX over H in

FS_TITLE = 16.0     # figure title
FS_PANEL = 13.2     # panel headings
FS_CAP = 12.8       # panel captions
FS_EQ = 12.0        # running-formula boxes

panel_w = 1.0 / 4
Y_DIAG = 0.655      # vertical centre of the sphere diagrams
Y_ARROW = 0.42      # clear band for the stage-to-stage arrows
Y_CAP = 0.335       # top of the panel captions
Y_EQ = 0.075        # centre of the running-formula boxes


def circ(cx, cy, r, **kw):
    e = Ellipse((cx, cy), width=2*r, height=2*r*ASPECT, **kw)
    ax.add_patch(e)
    return e

# panel dividers (stopped above the formula row so the widest formula box
# never runs into a divider line)
for k in range(1, 4):
    ax.plot([k*panel_w, k*panel_w], [0.15, 0.90], color=MIST, lw=1.0, zorder=1)

titles = ["1. Hard-sphere\nreference fluid", "2. Chain formation\n(bonded segments)",
          "3. Dispersion\n(segment attraction)", "4. Association\n(sites A, B)"]
subtitle = [r"$a_{hs}$", r"$a_{hs}+a_{chain}$", r"$a_{hs}+a_{chain}+a_{disp}$",
            r"$a_{hs}+a_{chain}+a_{disp}+a_{assoc}$"]

for k in range(4):
    cx = (k + 0.5) * panel_w
    ax.text(cx, 0.995, titles[k], ha="center", va="top", fontsize=FS_PANEL, color=NAVY,
            weight="bold", linespacing=1.25)
    ax.text(cx, Y_EQ, subtitle[k], ha="center", va="center", fontsize=FS_EQ, color=INK,
            bbox=dict(boxstyle="round,pad=0.28", fc="#F4F6F7", ec=GRAPHITE, lw=1.1))

r_seg = 0.022

# ---------------- Panel 1: hard-sphere fluid (disordered hard spheres) ----
cx1 = 0.129
np.random.seed(7)
pts1 = [(cx1 + dx, Y_DIAG + dy) for dx, dy in
        [(-0.0589, 0.0328), (0.0157, 0.0546), (0.0668, -0.0109), (-0.0236, -0.0601),
         (0.0550, -0.0765), (-0.0746, -0.0546), (0.0, 0.0)]]
for (x, y) in pts1:
    circ(x, y, r_seg, facecolor="#DCE2E5", edgecolor=GRAPHITE, lw=1.4, zorder=4)
ax.text(cx1, Y_CAP, "hard, impenetrable\nsegments,\nrandom packing", ha="center", va="top",
        fontsize=FS_CAP, color=SLATE, linespacing=1.3)

# ---------------- Panel 2: chain formation ---------------------------------
cx2 = 1.5 * panel_w
chain_y = Y_DIAG
n_chain = 5
d_bond = 2.4 * r_seg           # slightly more than tangent so the amber bonds show
xs = np.linspace(cx2 - 2*d_bond, cx2 + 2*d_bond, n_chain)
for i in range(n_chain - 1):
    ax.plot([xs[i], xs[i+1]], [chain_y, chain_y], color=AMBER, lw=2.6, zorder=3)
for x in xs:
    circ(x, chain_y, r_seg, facecolor="#DCE2E5", edgecolor=GRAPHITE, lw=1.4, zorder=4)
ax.text(cx2, Y_CAP, "m tangent segments\nbonded into a chain\n(amber bonds)", ha="center",
        va="top", fontsize=FS_CAP, color=SLATE, linespacing=1.3)

# ---------------- Panel 3: dispersion attraction ---------------------------
cx3 = 2.5 * panel_w
chain_y2a, chain_y2b = Y_DIAG + 0.075, Y_DIAG - 0.095
xsA = np.linspace(cx3 - 0.022 - 1.5*d_bond, cx3 - 0.022 + 1.5*d_bond, 4)
xsB = np.linspace(cx3 + 0.022 - 1.5*d_bond, cx3 + 0.022 + 1.5*d_bond, 4)
for i in range(3):
    ax.plot([xsA[i], xsA[i+1]], [chain_y2a, chain_y2a], color=GRAPHITE, lw=2.2, zorder=3)
    ax.plot([xsB[i], xsB[i+1]], [chain_y2b, chain_y2b], color=GRAPHITE, lw=2.2, zorder=3)
for x in xsA:
    circ(x, chain_y2a, r_seg, facecolor="#DCE2E5", edgecolor=GRAPHITE, lw=1.3, zorder=4)
for x in xsB:
    circ(x, chain_y2b, r_seg, facecolor="#DCE2E5", edgecolor=GRAPHITE, lw=1.3, zorder=4)
# dispersion attraction arrows between the two chains (amber)
for x0, x1 in [(xsA[1], xsB[1]), (xsA[2], xsB[2])]:
    arr = FancyArrowPatch((x0, chain_y2a - r_seg*ASPECT*1.15), (x1, chain_y2b + r_seg*ASPECT*1.15),
                           arrowstyle="<|-|>", mutation_scale=10, linewidth=1.4, color=AMBER, zorder=3)
    ax.add_patch(arr)
ax.text(cx3, Y_CAP, "attraction between\nsegments on different\nchains (amber)", ha="center",
        va="top", fontsize=FS_CAP, color=SLATE, linespacing=1.3)

# ---------------- Panel 4: association --------------------------------
cx4 = 3.5 * panel_w
m1x, m1y = cx4 - 0.058, Y_DIAG + 0.045
m2x, m2y = cx4 + 0.058, Y_DIAG - 0.055
circ(m1x, m1y, r_seg*1.3, facecolor="#DCE2E5", edgecolor=GRAPHITE, lw=1.4, zorder=4)
circ(m2x, m2y, r_seg*1.3, facecolor="#DCE2E5", edgecolor=GRAPHITE, lw=1.4, zorder=4)
# site markers A (on molecule 1) and B (on molecule 2), pointing toward each other
siteA = (m1x + r_seg*1.3, m1y - 0.012)
siteB = (m2x - r_seg*1.3, m2y + 0.012)
ax.plot([siteA[0]], [siteA[1]], marker="o", ms=8, color=AMBER, zorder=6,
         markeredgecolor=GRAPHITE, markeredgewidth=0.9)
ax.plot([siteB[0]], [siteB[1]], marker="o", ms=8, color=RUST, zorder=6,
         markeredgecolor=GRAPHITE, markeredgewidth=0.9)
ax.text(siteA[0] + 0.008, siteA[1] + 0.055, "A", ha="center", va="bottom", fontsize=FS_CAP,
        color=AMBER, weight="bold")
ax.text(siteB[0] - 0.008, siteB[1] - 0.055, "B", ha="center", va="top", fontsize=FS_CAP,
        color=RUST, weight="bold")
ax.plot([siteA[0], siteB[0]], [siteA[1], siteB[1]], color=AMBER, lw=1.6, linestyle=(0, (2, 1.5)), zorder=3)
ax.text(cx4, Y_CAP, "short-range,\ndirectional A-B\nbonding site pairs", ha="center", va="top",
        fontsize=FS_CAP, color=SLATE, linespacing=1.3)

# ---------------- connecting arrows between panels --------------------
# placed in the clear horizontal band below the diagrams, so no arrow tail or
# head lands inside a segment circle
for k in range(3):
    x0 = (k + 1) * panel_w - 0.055
    x1 = (k + 1) * panel_w + 0.055
    arr = FancyArrowPatch((x0, Y_ARROW), (x1, Y_ARROW), arrowstyle="-|>", mutation_scale=17,
                           linewidth=1.8, color=GRAPHITE, zorder=5)
    ax.add_patch(arr)

ax.text(0.5, YMAX - 0.015, "Building the SAFT residual Helmholtz energy, term by term",
        ha="center", va="top", fontsize=FS_TITLE, color=NAVY, weight="bold")

save(fig, "f25_saft_terms")
