"""Trace PNIPAAm/water/ethanol solution and gel binodals against DOS
swelling data."""

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

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,
)
from equi_gel.plot_ternary import (
    close_envelope,
    plot_ternary,
    ternary_coords,
)
from equi_gel.save_load_json import (
    load_dos_record,
)
from equi_gel.plot_binary import STYLE_PATH

from gel_tracer_ternary import trace_binodal_gel

# -----------------------------
# Configuration
# -----------------------------
# General Conditions
PRESSURE = 1.13
T_HIGH = 350.0
T_LOW = 298.15

# Components & Parameters
COMPONENTS = ["pnipaam", "water", "ethanol"]
SOURCE = "arndt"
MOLAR_WEIGHT = -1

# Gel Network Configuration
NETWORK_FUNCTIONALITY = 2.69
V_MAX_FAC = 1.0
Y_CROSSLINKER = 0.01
C_FAC = 1

ADD_PARAMS = {
    "number_components": 3,
    "gel": True,
    "elasticity": "Flory",
    "ref_volume_frac": 1,
    "NETWORK_FUNCTIONALITY": NETWORK_FUNCTIONALITY,
    "V_MAX_FAC": V_MAX_FAC,
    "y_crosslinker": Y_CROSSLINKER,
    "C_FAC": C_FAC,
}

# Experimental Data Files
DOS_FILES = []

# Plotting Styles
LINESTYLES = ["-", "--", ":", "-."]
MARKERS = ["o", "s", "^", "D"]


# -----------------------------
# Calculations: Polymer LLE
# -----------------------------
print("Starting Polymer LLE Calculations...")

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

# 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_fixed_conc = np.argmin(
    np.abs(np.asarray(y_high_t)[:, 2] - np.log(FIXED_CONCENTRATION))
)
y_seed = y_high_t[idx_fixed_conc]

# 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)


# -----------------------------
# Calculations: Gel LLE
# -----------------------------
print("Starting Gel LLE Calculations...")

params_pw |= ADD_PARAMS
params_pwe |= ADD_PARAMS

gel_lle_results = trace_binodal_gel(
    params_pwe,
    t_high=T_HIGH,
    t_low=T_LOW,
    y_crosslinker=Y_CROSSLINKER,
    verbosity=True,
)


# -----------------------------
# Load Experimental DOS Data
# -----------------------------
dos_records = []
for p in DOS_FILES:
    if Path(p).exists():
        dos_records.append(load_dos_record(p))
    else:
        print(f"Warning: Could not find DOS file {p}")


# -----------------------------
# %% Plotting (Side-by-Side)
# -----------------------------
print("Generating plots...")


plt.style.use(STYLE_PATH)
fig = plt.figure(figsize=(6.7, 2.2))

ternary_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
ternary_colors[1] = "red"

# ==========================================
# --- Left Subplot: DOS Plot (Cartesian) ---
# ==========================================
# Add standard subplot in a 1x2 grid at position 1
ax_dos = fig.add_subplot(1, 2, 2)
# ax_dos.set_title("Degree of Swelling")
ax_dos.set_xlabel(r"Mass fraction Ethanol [-]")
ax_dos.set_ylabel(r"DOS [g/g]")
ax_dos.set_xlim([0.0, 1.0])

binodal_dicts = [gel_lle_results, sol_t_low]
labels_model = [f"Gel System (y={Y_CROSSLINKER})", "Polymer Solution"]

# Plot model curves
for idx, (curve, lab) in enumerate(zip(binodal_dicts, labels_model)):
    w_poly = np.array([w[0] for w in curve["w_II_lst"]], dtype=float)
    w_aa = np.array([w[2] for w in curve["w_I_lst"]], dtype=float)

    mask = (
        np.isfinite(w_poly)
        & np.isfinite(w_aa)
        & (w_poly > 1e-12)
        & (w_aa >= 0.0)
        & (w_aa <= 1.0)
    )
    dos_model = 1.0 / w_poly[mask]

    ax_dos.plot(
        w_aa[mask],
        dos_model,
        label=lab,
        color=ternary_colors[idx % len(ternary_colors)],
        linestyle=LINESTYLES[idx % len(LINESTYLES)],
    )

# Plot experimental points
for rec_idx, rec in enumerate(dos_records):
    dos_exp = np.atleast_1d(np.asarray(rec["DOS_g_per_g"], dtype=float))

    w = rec.get("w_3_g_per_g")
    if w is None:
        w = rec.get("w_3_mol_per_mol")
    w = np.atleast_1d(np.asarray(w, dtype=float))

    if w.size == 1 and dos_exp.size > 1:
        w = np.full_like(dos_exp, w[0])

    y_cl = float(rec["crosslinker"])
    src = rec.get("source", "").split(",")[0].strip()
    lab = f"Exp. {src} (y={y_cl:g})" if src else f"Exp (y={y_cl:g})"

    mask = np.isfinite(w) & np.isfinite(dos_exp)

    ax_dos.scatter(
        w[mask],
        dos_exp[mask],
        marker=MARKERS[rec_idx % len(MARKERS)],
        label=lab,
        color=ternary_colors[rec_idx % len(ternary_colors)],
        zorder=2,
    )

ax_dos.legend(loc="upper left", borderaxespad=0.0)

# ==========================================
# --- Right Subplot: Ternary Plot (mpltern) ---
# ==========================================
# Add ternary subplot in a 1x2 grid at position 2
ax_ternary = fig.add_subplot(1, 2, 1, projection="ternary")
# ax_ternary.set_title("Ternary Phase Diagram", y=1.1)

# Set labels (Update these if you want actual component names instead of A/B)
ax_ternary.set_tlabel("pNIPAM")
ax_ternary.set_llabel("Water")
ax_ternary.set_rlabel("Ethanol")

for axis in (ax_ternary.taxis, ax_ternary.laxis, ax_ternary.raxis):
    axis.set_label_position("corner")
    axis.set_label_rotation_mode("horizontal")

ax_ternary.tick_params(labelrotation="horizontal")

ax_ternary.annotate(
    "Mass fractions [-]",
    xy=(0.5, -0.1),
    xycoords="axes fraction",
    ha="center",
    va="top",
)

# 1. Generate Coordinates for Low Temp Polymer LLE
low_t_ii = ternary_coords(sol_t_low["w_II_lst"])
low_t_i = ternary_coords(sol_t_low["w_I_lst"])
envelope_low = close_envelope(low_t_i, low_t_ii, reverse="left", close=True)

# 2. Generate Coordinates for Gel LLE
gel_t_ii = ternary_coords(gel_lle_results["w_II_lst"])
gel_t_i = ternary_coords(gel_lle_results["w_I_lst"])
# Setting close=False for the gel in case the binodal doesn't close on the axes
envelope_gel = close_envelope(gel_t_i, gel_t_ii, reverse="left", close=False)

# Plotting the ternary data
plot_ternary(
    ax_ternary,
    envelope_low,
    label=r"Polymer LLE ($T_{low}$)",
    linewidth=1.4,
    zorder=12,
    c=ternary_colors[1],
    linestyle="--",  # Dashed line to distinguish from the Polymer LLE
)

plot_ternary(
    ax_ternary,
    envelope_gel,
    label=f"Gel Model (y={Y_CROSSLINKER})",
    linewidth=1.4,
    zorder=13,
    c=ternary_colors[0],  # Matches the Gel color from the DOS plot
    linestyle="-",  # Dashed line to distinguish from the Polymer LLE
)

# ax_ternary.legend(loc="upper left", bbox_to_anchor=(1.1, 1.0), frameon=False)

# Finalize Layout & Save
plt.subplots_adjust(wspace=0.3)

figures_dir = Path("figures")
figures_dir.mkdir(parents=True, exist_ok=True)

try:
    outfile = figures_dir / Path(__file__).with_suffix(".pdf").name
except NameError:
    outfile = figures_dir / "binodal_gel_plot.pdf"

plt.savefig(outfile, dpi=600, bbox_inches="tight")
print(f"Saved: {outfile.resolve()}")
plt.show()

# %%
