"""Original symmetric regular-solution plots; writes a PNG in the current directory."""

from pathlib import Path
import math

import matplotlib.pyplot as plt
import numpy as np


def mixing_free_energy(x, chi):
    x = np.asarray(x, dtype=float)
    log_x = np.zeros_like(x)
    log_other = np.zeros_like(x)
    np.log(x, out=log_x, where=x > 0)
    np.log(1 - x, out=log_other, where=x < 1)
    return x * log_x + (1 - x) * log_other + chi * x * (1 - x)


def low_binodal(chi):
    assert chi >= 2
    if chi == 2:
        return 0.5
    low = 1e-12
    high = (1 - math.sqrt(1 - 2 / chi)) / 2
    for _ in range(100):
        middle = (low + high) / 2
        derivative = math.log(middle / (1 - middle)) + chi * (1 - 2 * middle)
        if derivative > 0:
            high = middle
        else:
            low = middle
    return (low + high) / 2


def main():
    alpha = low_binodal(3)
    beta = 1 - alpha
    z = 0.30
    phase_fraction = (z - alpha) / (beta - alpha)
    np.testing.assert_allclose((1 - phase_fraction) * alpha + phase_fraction * beta, z)
    assert 0 < phase_fraction < 1
    for chi in (2.1, 3, 4):
        endpoint = low_binodal(chi)
        np.testing.assert_allclose(
            math.log(endpoint / (1 - endpoint)) + chi * (1 - 2 * endpoint),
            0, atol=1e-12,
        )
        np.testing.assert_allclose(
            mixing_free_energy(endpoint, chi), mixing_free_energy(1 - endpoint, chi),
            atol=1e-14,
        )
    x = np.linspace(0, 1, 1201)
    minimum = float(mixing_free_energy(alpha, 3))
    assert np.all(mixing_free_energy(x, 3) >= minimum - 1e-12)
    assert minimum < float(mixing_free_energy(z, 3))

    plt.rcParams.update({"font.size": 11, "axes.titlesize": 11})
    fig, (energy, diagram) = plt.subplots(1, 2, figsize=(9.6, 4.8), dpi=100,
                                         facecolor="white")
    fig.subplots_adjust(left=0.105, right=0.985, bottom=0.16, top=0.80, wspace=0.28)
    fig.suptitle("Symmetric regular solution: coexistence and local stability",
                 fontsize=15, y=0.97)
    fig.text(0.5, 0.885,
             r"$h(x)=x\ln x+(1-x)\ln(1-x)+\chi x(1-x)$; original model calculation",
             ha="center", fontsize=12)
    energy.plot(x, mixing_free_energy(x, 3), color="#0072b2", lw=2.2,
                label="homogeneous liquid")
    energy.plot([alpha, beta], [minimum, minimum], color="#d55e00", lw=2.5,
                label="two-phase minimum")
    energy.plot([alpha, beta], [minimum, minimum], "o", color="#d55e00")
    energy.plot(z, mixing_free_energy(z, 3), "o", color="#333333", ms=5)
    energy.annotate("overall z = 0.30", xy=(z, mixing_free_energy(z, 3)),
                    xytext=(0.04, 0.067), fontsize=10,
                    arrowprops={"arrowstyle": "->", "color": "#666666"})
    energy.text(alpha, minimum - 0.008, r"$x_\alpha$", ha="center", va="top")
    energy.text(beta, minimum - 0.008, r"$x_\beta$", ha="center", va="top")
    energy.set(xlabel="Mole fraction x", ylabel=r"Mixing free energy $h=g_{mix}/(RT)$",
               title=r"Supporting common tangent at $\chi=3$",
               xlim=(0, 1), ylim=(-0.08, 0.09))
    energy.legend(loc="upper right", fontsize=9)

    chi = np.linspace(2, 4, 301)
    binodal = np.array([low_binodal(c) for c in chi])
    spinodal = (1 - np.sqrt(1 - 2 / chi)) / 2
    diagram.plot(binodal, chi, color="#0072b2", lw=2.2, label="binodal: coexistence")
    diagram.plot(1 - binodal, chi, color="#0072b2", lw=2.2)
    diagram.plot(spinodal, chi, color="#cc79a7", lw=2, ls="--",
                 label="spinodal: local stability")
    diagram.plot(1 - spinodal, chi, color="#cc79a7", lw=2, ls="--")
    diagram.plot([alpha, beta], [3, 3], color="#777777", lw=1, ls=":")
    diagram.plot([alpha, beta], [3, 3], "o", color="#0072b2", ms=4)
    diagram.set(xlabel="Mole fraction x", ylabel=r"Interaction parameter $\chi$",
                title="Binodal and spinodal in the model",
                xlim=(0, 1), ylim=(1.95, 4.05))
    diagram.legend(loc="upper center", fontsize=9)
    for ax in (energy, diagram):
        ax.grid(alpha=0.20)
        ax.set_facecolor("white")
        ax.spines[["top", "right"]].set_visible(False)
    fig.savefig(Path.cwd() / "regular-solution-liquid-coexistence.png", dpi=100,
                transparent=False, facecolor="white")
    plt.close(fig)
    print(f"chi=3: x_alpha={alpha:.10f}, x_beta={beta:.10f}, "
          f"beta mole fraction at z=0.30: {phase_fraction:.10f}")


if __name__ == "__main__":
    main()