"""Trace a PNIPAAm/water/ethanol ternary binodal and plot the
envelope."""

from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

from equi_gel.plot_ternary import (
    close_envelope,
    plot_ternary,
    scatter_ternary,
    ternary_axes,
    ternary_coords,
)

from equi_gel import parameters
from equi_gel.newton_raphson_tracing import (
    add_component_to_y_vec,
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
)

COMPONENTS = ["pnipaam", "water", "ethanol"]
SOURCE = "arndt"
T_HIGH = 350.0
T_LOW = 298.15
PRESSURE = 1.13


params_pw = parameters.get_parameter_set_from_tables(COMPONENTS[:-1], SOURCE)
params_pwe = parameters.get_parameter_set_from_tables(COMPONENTS, SOURCE)


# Binary miscibility gap -> seed ternary

y_init_binary = build_initial_guess_binary(T_HIGH, PRESSURE, params_pw)

y_init_ternary = add_component_to_y_vec(y_init_binary, 1e-5, 1e-5)

y_high_t, _, _, _, _, _ = trace_parameter(
    trace_index=2,
    y_init=y_init_ternary,
    trace_start=1e-5,
    trace_target=0.8,
    pressure=PRESSURE,
    params=params_pwe,
    step_size_init=0.001,
    newton_tol=1e-7,
    fixed_index=[-1],
    critical_point_tol=1e-3,
)

sol_high_t = interpret_y_vector(y_high_t, PRESSURE, params_pwe)

# Pick a fixed concentration for temperature tracing
FIXED_CONCENTRATION = 0.2
idx = np.argmin(
    np.abs(np.asarray(y_high_t)[:, 2] - np.log(FIXED_CONCENTRATION))
)

y_seed = y_high_t[idx]

# Trace temperature down

y_t_trace, _, _, _, _, _ = trace_parameter(
    trace_index=-1,
    y_init=y_seed,
    trace_start=T_HIGH,
    trace_target=T_LOW,
    pressure=PRESSURE,
    params=params_pwe,
    step_size_init=0.01,
    newton_tol=1e-6,
    fixed_index=[2],
    max_step_fraction=10,
)

sol_t_trace = interpret_y_vector(y_t_trace, PRESSURE, params_pwe)

# Trace low-temperature branches

y_low_left, _, _, _, _, _ = trace_parameter(
    trace_index=2,
    y_init=y_t_trace[-1],
    trace_start=np.exp(y_t_trace[-1][2]),
    trace_target=1e-2,
    pressure=PRESSURE,
    params=params_pwe,
    step_size_init=0.0001,
    newton_tol=1e-8,
    fixed_index=[-1],
    critical_point_tol=5e-5,
    stuck_progress_tol=6e-4,
    max_step_fraction=1e-4,
)

y_low_right, _, _, _, _, _ = trace_parameter(
    trace_index=2,
    y_init=y_t_trace[-1],
    trace_start=np.exp(y_t_trace[-1][2]),
    trace_target=1.0,
    pressure=PRESSURE,
    params=params_pwe,
    step_size_init=0.001,
    newton_tol=1e-6,
    fixed_index=[-1],
    critical_point_tol=5e-5,
    stuck_progress_tol=2e-4,
)

y_low = (y_low_left[::-1])[:-1] + y_low_right
sol_t_low = interpret_y_vector(y_low, PRESSURE, params_pwe)

# Plot

fig, ax = ternary_axes(labels=("A", "B", "C"))

# Set labels (Update these if you want actual component names instead of A/B)
ax.set_tlabel("pNIPAM")
ax.set_llabel("Water")
ax.set_rlabel("Ethanol")
colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]

high_t_ii = ternary_coords(sol_high_t["w_II_lst"])
high_t_i = ternary_coords(sol_high_t["w_I_lst"])

trace_t_ii = ternary_coords(sol_t_trace["w_II_lst"])
trace_t_i = ternary_coords(sol_t_trace["w_I_lst"])

low_t_ii = ternary_coords(sol_t_low["w_II_lst"])
low_t_i = ternary_coords(sol_t_low["w_I_lst"])

envelope_high = close_envelope(
    high_t_i, high_t_ii, reverse="right", close=False
)
envelope_low = close_envelope(low_t_i, low_t_ii, reverse="left", close=True)

plot_ternary(
    ax,
    envelope_high,
    label=r"$T_1 > T_2$",
    linewidth=1.4,
    zorder=10,
    c="red",
)

plot_ternary(
    ax,
    envelope_low,
    label=r"$T_2$",
    linewidth=1.4,
    zorder=12,
    c=colors[0],
)

scatter_ternary(
    ax,
    trace_t_ii,
    idx=None,
    linewidth=1.4,
    label="Tracing T",
    marker="x",
    c="red",
    zorder=15,
)

scatter_ternary(
    ax,
    trace_t_i,
    idx=None,
    linewidth=1.4,
    marker="x",
    c="red",
    zorder=15,
)

plot_ternary(
    ax,
    np.vstack([trace_t_i[0], trace_t_ii[0]]),
    linewidth=0.7,
    linestyle="--",
    c="red",
    zorder=1,
)

ax.legend(loc="upper left", bbox_to_anchor=(0.75, 1.0), frameon=False)

out_dir = Path.cwd() / "figures"
out_dir.mkdir(exist_ok=True)
plt.savefig(
    out_dir / f"{Path(__file__).stem}.pdf", dpi=600, bbox_inches="tight"
)


# --------------------
# TOC plot
# --------------------
fig_toc, ax_toc = ternary_axes(labels=("", "", ""))

# Make it smaller
fig_toc.set_size_inches(2.2, 2.2)

# No ternary labels
ax_toc.set_tlabel("")
ax_toc.set_llabel("")
ax_toc.set_rlabel("")

# No regular matplotlib ticks
ax_toc.set_xticks([])
ax_toc.set_yticks([])
ax_toc.tick_params(
    left=False, bottom=False, labelleft=False, labelbottom=False
)


# Thicker lines for the small TOC figure
plot_ternary(
    ax_toc,
    envelope_high,
    linewidth=2.6,
    zorder=10,
    c="red",
)

plot_ternary(
    ax_toc,
    envelope_low,
    linewidth=2.6,
    zorder=12,
    c=colors[0],
)

scatter_ternary(
    ax_toc,
    trace_t_ii,
    idx=None,
    linewidth=2.2,
    marker="x",
    c="red",
    zorder=15,
)

scatter_ternary(
    ax_toc,
    trace_t_i,
    idx=None,
    linewidth=2.2,
    marker="x",
    c="red",
    zorder=15,
)

plot_ternary(
    ax_toc,
    np.vstack([trace_t_i[0], trace_t_ii[0]]),
    linewidth=1.6,
    linestyle="--",
    c="red",
    zorder=1,
)

fig_toc.tight_layout(pad=0.02)

plt.savefig(
    out_dir / f"{Path(__file__).stem}_toc.pdf",
    dpi=600,
    bbox_inches="tight",
    pad_inches=0.01,
)
