"""Plot representative network free energies and their derivatives."""

from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from equi_gel.plot_ternary import STYLE_PATH


# --------------------
# Parameters
# --------------------
V0 = 1.0
V_MAX = 10.0
V_MAX_LARGE = 100.0
N_P = 1.0
EPS = 1e-6
x_grid = np.linspace(EPS, 100.0, 5000)

output_dir = Path("figures")
output_dir.mkdir(exist_ok=True)


def f_affine(x_val):
    """Affine network free energy."""
    return 0.5 * N_P * (3 * (x_val / V0) ** (2 / 3) - 3 - np.log(x_val / V0))


def f_phantom(x_val):
    """Phantom network free energy."""
    return 0.5 * N_P * (3 * (x_val / V0) ** (2 / 3) - 3)


def f_finite_extensibility(x_val, vmax):
    """Finite extensibility free energy."""
    mask = x_val < vmax
    out = np.full_like(x_val, np.nan)
    out[mask] = (
        0.5
        * N_P
        * (
            3
            * ((x_val[mask] / V0) ** (2 / 3) - 1)
            * (1 - (x_val[mask] / vmax) ** (2 / 3)) ** -1
            - np.log(x_val[mask] / V0)
        )
    )
    return out


def d_affine(x_val):
    """Derivative of affine network free energy."""
    return N_P * (V0 ** (2 / 3) * x_val ** (-1 / 3) - 0.5 * x_val**-1)


def d_phantom(x_val):
    """Derivative of phantom network free energy."""
    return N_P * (V0 ** (2 / 3) * x_val ** (-1 / 3))


def d_finite_extensibility(x_val, vmax):
    """Derivative of finite extensibility free energy."""
    mask = x_val < vmax
    out = np.full_like(x_val, np.nan)
    out[mask] = (N_P / x_val[mask]) * (
        (x_val[mask] / V0) ** (2 / 3)
        * (1 - (x_val[mask] / vmax) ** (2 / 3)) ** -1
        + ((x_val[mask] / V0) ** (2 / 3) - 1)
        * (1 - (x_val[mask] / vmax) ** (2 / 3)) ** -2
        * (x_val[mask] / vmax) ** (2 / 3)
        / 2
        - 0.5
    )
    return out


def plot_family(axis, x_values, series, *, palette):
    """Plot a family of curves with consistent ordering/styling."""
    axis.plot(x_values, series[0][1], label=series[0][0], zorder=50)
    axis.plot(x_values, series[1][1], label=series[1][0], zorder=49)
    axis.plot(
        x_values, series[2][1], label=series[2][0], color=palette[2], zorder=48
    )
    axis.plot(
        x_values,
        series[3][1],
        linestyle="--",
        color=palette[2],
        label=series[3][0],
        zorder=47,
    )


# --------------------
# Plot free energies
# --------------------
plt.style.use(STYLE_PATH)
colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
colors[2] = "red"


fig, ax = plt.subplots(1)
ax.axvline(V0, linestyle="--", linewidth=0.75, color="black")
ax.axhline(0, linestyle="--", linewidth=0.75, color="black")

series_energy = [
    ("Affine Network", f_affine(x_grid)),
    ("Phantom Network", f_phantom(x_grid)),
    (
        rf"FE: $V_\mathrm{{max}}={V_MAX}\,V_\mathrm{{ref}}$",
        f_finite_extensibility(x_grid, V_MAX),
    ),
    (
        rf"FE: $V_\mathrm{{max}}={V_MAX_LARGE}\,V_\mathrm{{ref}}$",
        f_finite_extensibility(x_grid, V_MAX_LARGE),
    ),
]
plot_family(ax, x_grid, series_energy, palette=colors)

ax.set_xlim(0, 5)
ax.set_ylim(-2, 10)
ax.set_xlabel(r"Normalized Volume $V/V_\mathrm{ref}\,\,[-]$")
ax.set_ylabel(r"$A_{\mathrm{el}}/CRT\,\,[-]$")
ax.legend()
fig.tight_layout()
output_path = output_dir / f"{Path(__file__).stem}_potentials.pdf"
fig.savefig(output_path, format="pdf", bbox_inches="tight")


# --------------------
# Plot derivatives
# --------------------
fig, ax = plt.subplots(1)
ax.axvline(V0, linestyle="--", linewidth=0.75, color="black")
ax.axhline(0, linestyle="--", linewidth=0.75, color="black")

series_derivative = [
    ("Affine Network", d_affine(x_grid)),
    ("Phantom Network", d_phantom(x_grid)),
    (
        rf"FE: $V_\mathrm{{max}}={V_MAX}\,V_\mathrm{{ref}}$",
        d_finite_extensibility(x_grid, V_MAX),
    ),
    (
        rf"FE: $V_\mathrm{{max}}={V_MAX_LARGE}\,V_\mathrm{{ref}}$",
        d_finite_extensibility(x_grid, V_MAX_LARGE),
    ),
]
plot_family(ax, x_grid, series_derivative, palette=colors)

ax.set_xlim(0, 5)
ax.set_ylim(-5, 10)
ax.set_xlabel(r"Normalized Volume $V/V_\mathrm{ref}\,\,[-]$")
ax.set_ylabel(
    r"$\partial(A_{\mathrm{el}}/CRT)/\partial V\,\,[\mathrm{m^{-3}}]$"
)
ax.legend()
fig.tight_layout()
output_path = output_dir / f"{Path(__file__).stem}_derivatives.pdf"
fig.savefig(output_path, format="pdf", bbox_inches="tight")
