"""f32_ml_eos -- schematic of a physics-informed ML EOS: inputs
(T, rho, molecular descriptors) -> neural Helmholtz function -> automatic
differentiation to P, mu, c_v, with hard constraints (ideal-gas limit,
Maxwell equal-area, positivity of c_v) shown as guard rails below.
Strict left-to-right reading order. Pure schematic.

Layout note: the canvas is deliberately small (9.0 x 5.6 in) and the point
sizes correspondingly large -- projected about 6 in wide, legibility follows
the ratio of font size to canvas width. The red guard-rail caption sits
below the guard-rail boxes so the three red constraint arrows, which run
upward from the box tops to the network box, no longer cross it.
"""
import sys
sys.path.insert(0, "/home/claude/eos_lecture/code")
from eosstyle import *
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch

W, H = 9.0, 5.6
fig, ax = blank(W, H)

FS_TITLE = 16.5
FS_MAIN = 12.0      # main-row box text
FS_GUARD = 11.5     # guard-rail box text
FS_CAP = 12.5       # red guard-rail caption


def box(cx, cy, w, h, text, fc=PAPER, ec=GRAPHITE, tc=INK, fs=FS_MAIN, lw=1.2, weight="normal"):
    b = FancyBboxPatch((cx - w/2, cy - h/2), w, h,
                        boxstyle="round,pad=0.012,rounding_size=0.02",
                        linewidth=lw, edgecolor=ec, facecolor=fc, zorder=3)
    ax.add_patch(b)
    ax.text(cx, cy, text, ha="center", va="center", fontsize=fs, color=tc,
             zorder=4, weight=weight, linespacing=1.3)
    return b


def arrow(p0, p1, color=GRAPHITE, lw=1.7):
    a = FancyArrowPatch(p0, p1, arrowstyle="-|>", mutation_scale=15,
                         linewidth=lw, color=color, zorder=2, shrinkA=2, shrinkB=2)
    ax.add_patch(a)

y_main = 0.72
h_main = 0.30

# main row: (centre x, width)
xs_main = [0.095, 0.365, 0.635, 0.895]
ws_main = [0.155, 0.220, 0.195, 0.175]

# 1. inputs
box(xs_main[0], y_main, ws_main[0], h_main, "Inputs\n\n$T,\\ \\rho$\nmolecular\ndescriptors",
    fc="#F4F6F7", ec=GRAPHITE, tc=INK, fs=FS_MAIN)

# 2. neural Helmholtz function
box(xs_main[1], y_main, ws_main[1], h_main, "Neural network\nHelmholtz function\n"
    r"$a^{res}_\theta(T,\rho)$", fc=NAVY, tc="white", fs=FS_MAIN, weight="bold")

# 3. automatic differentiation
box(xs_main[2], y_main, ws_main[2], h_main, "Automatic\ndifferentiation\n"
    r"$\partial a/\partial\rho,\ \partial a/\partial T$", fc=AMBER, tc="white",
    fs=FS_MAIN, weight="bold")

# 4. outputs
box(xs_main[3], y_main, ws_main[3], h_main, "Outputs\n\n$P,\\ \\mu,\\ c_v$",
    fc="#F4F6F7", ec=GRAPHITE, tc=INK, fs=FS_MAIN)

for k in range(3):
    x0 = xs_main[k] + ws_main[k]/2 + 0.010
    x1 = xs_main[k+1] - ws_main[k+1]/2 - 0.010
    arrow((x0, y_main), (x1, y_main))

# guard rails, below, feeding into the network box (training-time constraints)
y_guard = 0.235
h_guard = 0.200
w_guard = 0.215
guard_texts = [
    "Ideal-gas limit\n" r"$a^{res}\to 0$ as $\rho\to 0$",
    "Maxwell equal-area\n(consistent phase\nequilibrium)",
    r"Positivity of $c_v$" "\n(thermodynamic\nstability)",
]
gx = [0.15, 0.40, 0.65]
for x, txt in zip(gx, guard_texts):
    box(x, y_guard, w_guard, h_guard, txt, fc=SAND, ec=GRAPHITE, tc=INK, fs=FS_GUARD)
    arrow((x, y_guard + h_guard/2 + 0.012), (xs_main[1], y_main - h_main/2 - 0.014),
          color=RUST, lw=1.4)

# caption placed below the guard-rail row: the three red arrows all run upward
# from the box tops, so nothing crosses this line of text
ax.text(0.5, 0.055, "hard physical constraints (guard rails) shape training and evaluation",
        ha="center", va="center", fontsize=FS_CAP, color=RUST, weight="bold")

ax.text(0.5, 0.985, "A physics-informed machine-learning equation of state",
        ha="center", va="top", fontsize=FS_TITLE, color=NAVY, weight="bold")

save(fig, "f32_ml_eos")
