"""Plot a binary free-energy curve and illustrate the alternating-
tangents construction."""

from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq
from mpl_toolkits.axes_grid1.inset_locator import inset_axes, mark_inset

from equi_gel.plot_ternary import STYLE_PATH

# Model parameters
A = 2.9
B = 5.0

# Numerical domain (avoid log singularities at 0 and 1)
EPS = 1e-6
x_grid = np.linspace(EPS, 1.0 - EPS, 1000)


def g_x(x_val):
    """Free energy of mixing."""
    return (
        A * x_val * np.log(x_val)
        + (1 - x_val) * np.log(1 - x_val)
        + B * x_val * (1 - x_val)
    )


def dg_dx(x_val):
    """First derivative of g(x)."""
    return A * np.log(x_val) + A - 2 * B * x_val + B - np.log(1 - x_val) - 1


def ddg_ddx(x_val):
    """Second derivative of g(x)."""
    return A / x_val - 2 * B + 1 / (1 - x_val)


def tangent_condition(x_k, x_ref):
    """Secant slope between x_k and x_ref equals dg/dx at x_k."""
    return (g_x(x_ref) - g_x(x_k)) / (x_ref - x_k) - dg_dx(x_k)


def line_tangent(x_val, _x_a, _x_b):
    """Line through (x_a, g(x_a)) with slope from (x_a, x_b)."""
    return g_x(_x_a) + (g_x(_x_b) - g_x(_x_a)) / (_x_b - _x_a) * (x_val - _x_a)


def find_tangent_root(x_ref, a, b):
    """Root of tangent_condition(x, x_ref) in [a, b]."""
    return brentq(
        lambda x_val: tangent_condition(x_val, x_ref),
        a=a,
        b=b,
        xtol=1e-10,
        rtol=1e-10,
        maxiter=100,
    )


# Spinodal points (roots of ddg/ddx)
DISCRIMINANT = A**2 - 2 * A * (2 * B + 1) + (1 - 2 * B) ** 2
x_sp_1 = (-np.sqrt(DISCRIMINANT) + A + 2 * B - 1) / (4 * B)
x_sp_2 = (+np.sqrt(DISCRIMINANT) + A + 2 * B - 1) / (4 * B)

# Alternating tangent points
x_1 = find_tangent_root(x_sp_1, a=x_sp_2, b=1.0 - EPS)
x_2 = find_tangent_root(x_1, a=EPS, b=x_sp_1)
x_3 = find_tangent_root(x_2, a=x_1, b=1.0 - EPS)


def plot_tangents(axis, tangents, *, color):
    """Plot a sequence of tangent lines on an axis."""
    for _x_a, _x_b in tangents:
        axis.plot(
            x_grid,
            line_tangent(x_grid, _x_a, _x_b),
            linewidth=0.5,
            linestyle="--",
            color=color,
        )


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

fig, ax = plt.subplots(1)
ax.plot(x_grid, g_x(x_grid), label=r"$g^\mathrm{mix}(x_1)$", zorder=50)

# Key points
points_main = np.array([x_sp_1, x_sp_2, x_1, x_2, x_3])
ax.scatter(points_main, g_x(points_main), marker="x", color="black")

plot_tangents(ax, [(x_sp_1, x_1), (x_2, x_1), (x_2, x_3)], color=colors[1])

ax.annotate(
    r"$x^\mathrm{sp}_1$",
    (x_sp_1, g_x(x_sp_1)),
    xytext=(-1, 2),
    textcoords="offset points",
    ha="right",
    va="bottom",
)
ax.annotate(
    r"$x^\mathrm{sp}_2$",
    (x_sp_2, g_x(x_sp_2)),
    xytext=(-1, 3),
    textcoords="offset points",
    ha="right",
    va="bottom",
)
ax.annotate(
    r"$x^{(2)}$",
    (x_2, g_x(x_2)),
    xytext=(2, 2),
    textcoords="offset points",
    ha="right",
    va="bottom",
)

ax.set_xlim(0, 1)
ax.set_ylim(-0.4, 0.05)
ax.set_xlabel(r"Composition $x$")
ax.set_ylabel(r"$g(x)$")
ax.tick_params(labelleft=False)
ax.legend()

# Inset (zoom near x_3)
axins = inset_axes(
    ax,
    width="35%",
    height="35%",
    loc="lower right",
    borderpad=1,
)
axins.plot(x_grid, g_x(x_grid), zorder=50)

axins.scatter([x_sp_1], [g_x(x_sp_1)], marker="^", color="black")
points_inset = np.array([x_sp_2, x_1, x_2, x_3])
axins.scatter(points_inset, g_x(points_inset), marker="x", color="black")

plot_tangents(axins, [(x_sp_1, x_1), (x_2, x_1), (x_2, x_3)], color=colors[1])

axins.annotate(
    r"$x^{(1)}$",
    (x_1, g_x(x_1)),
    xytext=(2, 2),
    textcoords="offset points",
    ha="right",
    va="bottom",
)
axins.annotate(
    r"$x^{(3)}$",
    (x_3, g_x(x_3)),
    xytext=(9, 2),
    textcoords="offset points",
    ha="right",
    va="bottom",
)

INSET_SIZE_X = 0.02
INSET_SIZE_Y = 0.01
axins.set_xlim(x_3 - INSET_SIZE_X, x_3 + INSET_SIZE_X)
axins.set_ylim(g_x(x_3) - INSET_SIZE_Y, g_x(x_3) + INSET_SIZE_Y)
axins.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)

mark_inset(ax, axins, loc1=2, loc2=1, fc="none", ec="black", lw=0.6)

fig.tight_layout()

output_dir = Path("figures")
output_dir.mkdir(exist_ok=True)
output_path = output_dir / f"{Path(__file__).stem}.pdf"
fig.savefig(output_path, format="pdf", bbox_inches="tight")
plt.show()

# --------------------
# TOC plot
# --------------------
fig_toc, ax_toc = plt.subplots(figsize=(2.0, 2.0))

ax_toc.plot(
    x_grid,
    g_x(x_grid),
    linewidth=2.0,
    color=colors[0],
    zorder=50,
)

ax_toc.scatter(
    points_main,
    g_x(points_main),
    marker="x",
    s=28,
    linewidths=2.0,
    color="black",
    zorder=100,
)

for x_a, x_b in [(x_sp_1, x_1), (x_2, x_3)]:
    ax_toc.plot(
        x_grid,
        line_tangent(x_grid, x_a, x_b),
        linewidth=1.0,
        linestyle="--",
        color=colors[1],
    )

ax_toc.set_xlim(0, 1)
ax_toc.set_ylim(-0.4, 0.05)

ax_toc.set_xticks([])
ax_toc.set_yticks([])
ax_toc.set_xlabel("")
ax_toc.set_ylabel("")
ax_toc.tick_params(left=False, bottom=False)

for spine in ax_toc.spines.values():
    spine.set_linewidth(1.6)

fig_toc.tight_layout(pad=0.05)

toc_output_path = output_dir / f"{Path(__file__).stem}_toc.pdf"
fig_toc.savefig(
    toc_output_path, format="pdf", bbox_inches="tight", pad_inches=0.01
)
plt.show()
