"""Shared style + helpers for the EOS lecture package.

Usage in every figure script:

    import sys; sys.path.insert(0, "/home/claude/eos_lecture/code")
    from eosstyle import *          # noqa
    fig, ax = newfig()
    ...
    save(fig, "vdw_isotherms")
"""
import os
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

FIGDIR = "/home/claude/eos_lecture/figs"
os.makedirs(FIGDIR, exist_ok=True)

# ---------------------------------------------------------------- palette
INK      = "#1C242B"
GRAPHITE = "#333F4A"
SLATE    = "#66727C"
MIST     = "#B9C1C6"
PAPER    = "#FFFFFF"
NAVY     = "#1F4E79"
SKYBLUE  = "#5A91BE"
AMBER    = "#E29A2D"
TEAL     = "#0F6E6B"
RUST     = "#BE654C"
SAND     = "#DFC98F"

CYCLE = [NAVY, AMBER, TEAL, SKYBLUE, RUST, SLATE]

plt.rcParams.update({
    "figure.facecolor": PAPER,
    "savefig.facecolor": PAPER,
    "axes.facecolor": PAPER,
    "axes.edgecolor": GRAPHITE,
    "axes.linewidth": 0.9,
    "axes.labelcolor": INK,
    "axes.labelsize": 13,
    "axes.titlesize": 14,
    "axes.spines.top": False,
    "axes.spines.right": False,
    "axes.prop_cycle": matplotlib.cycler(color=CYCLE),
    "xtick.color": GRAPHITE,
    "ytick.color": GRAPHITE,
    "xtick.labelsize": 12,
    "ytick.labelsize": 12,
    "text.color": INK,
    "grid.color": MIST,
    "grid.linewidth": 0.6,
    "grid.alpha": 0.55,
    "lines.linewidth": 2.0,
    "lines.markersize": 5.5,
    "legend.frameon": False,
    "legend.fontsize": 11.5,
    "font.family": "sans-serif",
    "font.sans-serif": ["DejaVu Sans"],
    "mathtext.fontset": "dejavusans",
    "figure.dpi": 110,
})


def newfig(w=6.4, h=4.2, **kw):
    fig, ax = plt.subplots(figsize=(w, h), **kw)
    ax.grid(True, zorder=0)
    ax.set_axisbelow(True)
    return fig, ax


def blank(w=8.0, h=4.5):
    """Canvas for schematics / concept maps: no axes, unit coordinates."""
    fig, ax = plt.subplots(figsize=(w, h))
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.axis("off")
    return fig, ax


def save(fig, name, tight=True):
    if tight:
        fig.tight_layout()
    fig.savefig(f"{FIGDIR}/{name}.pdf", bbox_inches="tight", pad_inches=0.03)
    fig.savefig(f"{FIGDIR}/{name}.png", dpi=600, bbox_inches="tight", pad_inches=0.03)
    plt.close(fig)
    print("wrote", name)


# ------------------------------------------------------------- cubic EOS
R = 8.314            # J/mol/K
Rb = 83.14           # cm3 bar /mol/K

SPECIES = {
    "methane":  dict(Tc=190.6, Pc=45.99, w=0.012),
    "ethane":   dict(Tc=305.3, Pc=48.72, w=0.100),
    "propane":  dict(Tc=369.8, Pc=42.48, w=0.152),
    "n-butane": dict(Tc=425.1, Pc=37.96, w=0.200),
    "CO2":      dict(Tc=304.2, Pc=73.83, w=0.224),
    "nitrogen": dict(Tc=126.2, Pc=34.00, w=0.038),
    "water":    dict(Tc=647.1, Pc=220.55, w=0.345),
    "ammonia":  dict(Tc=405.7, Pc=112.80, w=0.253),
    "benzene":  dict(Tc=562.2, Pc=48.98, w=0.210),
}


def cubic_params(model, Tc, Pc, w, T, Rg=Rb):
    """Return (aT, b, delta1, delta2) with P = RT/(V-b) - aT/((V+d1 b)(V+d2 b))."""
    Tr = T / Tc
    if model == "vdw":
        a = 27 * Rg**2 * Tc**2 / (64 * Pc); b = Rg * Tc / (8 * Pc)
        return a, b, 0.0, 0.0
    if model == "rk":
        a = 0.42748 * Rg**2 * Tc**2.5 / Pc / np.sqrt(T)
        b = 0.08664 * Rg * Tc / Pc
        return a, b, 1.0, 0.0
    if model == "srk":
        ac = 0.42748 * Rg**2 * Tc**2 / Pc
        m = 0.480 + 1.574*w - 0.176*w**2
        al = (1 + m*(1 - np.sqrt(Tr)))**2
        b = 0.08664 * Rg * Tc / Pc
        return ac*al, b, 1.0, 0.0
    if model == "pr":
        ac = 0.45724 * Rg**2 * Tc**2 / Pc
        k = 0.37464 + 1.54226*w - 0.26992*w**2
        al = (1 + k*(1 - np.sqrt(Tr)))**2
        b = 0.07780 * Rg * Tc / Pc
        return ac*al, b, 1 + np.sqrt(2), 1 - np.sqrt(2)
    raise ValueError(model)


def P_of_V(model, sp, T, V, Rg=Rb):
    d = SPECIES[sp]
    a, b, d1, d2 = cubic_params(model, d["Tc"], d["Pc"], d["w"], T, Rg)
    if model == "vdw":
        return Rg*T/(V - b) - a/V**2
    return Rg*T/(V - b) - a/((V + d1*b)*(V + d2*b))


def Zroots(model, sp, T, P, Rg=Rb):
    """Cubic in Z; returns sorted real roots."""
    d = SPECIES[sp]
    a, b, d1, d2 = cubic_params(model, d["Tc"], d["Pc"], d["w"], T, Rg)
    A = a*P/(Rg*T)**2
    B = b*P/(Rg*T)
    u, wq = d1 + d2, d1*d2
    c = [1.0,
         (u - 1)*B - 1,
         A + (wq*B - u - u*B)*B,
         -(A*B + wq*B**2*(1 + B))]
    r = np.roots(c)
    r = np.sort(np.real(r[np.abs(np.imag(r)) < 1e-9]))
    return r[r > B*0.999]
