# %%
"""Compare PNIPAAm-water gel binodals while varying molar mass and the
network parameter C."""

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

from equi_gel.newton_raphson_tracing import (
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
)
from equi_gel.parameters import (
    add_mixing_association,
    get_parameter_set_from_tables,
    scale_component_molarweight,
)
from equi_gel.plot_binary import (
    STYLE_PATH,
    binary_coords,
    plot_binary,
    get_style_colors,
)
from equi_gel.save_load_json import (
    BASE_FOLDER_EXP,
)


PRESSURE = 1.13
T_TRACE_HIGH = 330.0
T_TRACE_LOW = 290.0

MASS_DENSITY_TEMPERATURE = 298.15

AMORPH_POLYMER_DENSITY = 1.11  # g / ccm

TRACE_STEP_SIZE_INIT = 1e-8
TRACE_NEWTON_TOL = 1e-6
TRACE_CRITICAL_TOL = 1e-5
TRACE_MAX_STEP_FRACTION = 0.1

JUMP_THRESH_VPT = 0.05

MOLAR_WEIGHT = 2e3

MOLAR_WEIGHT_LST = [1000, 6000, 60000]
C_LST = [5, 1 / 4]

ADD_PARAMS = {
    "number_components": 2,
    "gel": True,
    "elasticity": "Flory",
    "C": 1,
    "ref_volume_frac": 1.0,
    "max_volume_poly_frac": 0.003,
}

EXP_PNIPAAM_WATER_PATTERN = "EXP_SOL_PNIPAAM_water_*.json"
EXP_PNIPAAM_DENS_KANY = BASE_FOLDER_EXP / "EXP_POLY_DENS_Kany.json"

LINESTYLES = ["-", "--", "-.", ":"]
MARKERS = ["o", "s", "^", "v", "D", "1", "2", "x", "+", "*"]
FACE_MARKERS = [
    "none",
    "none",
    "none",
    "none",
    "none",
    "black",
    "black",
    "black",
    "black",
    "black",
]


opt_arndt_water = {
    "comp12_kij_intercept": -0.16411647669662233,
    "comp12_kij_slope": 0.001275948739486485,
}
opt_arndt_water = {}


def trace_binodal(params, t_high, t_low):
    """Trace binodal branches from high to low temperature."""
    y_init_binary = build_initial_guess_binary(t_high, PRESSURE, params)
    y_lst, _, _, _, _, _ = trace_parameter(
        trace_index=-1,
        y_init=y_init_binary,
        trace_start=t_high,
        trace_target=t_low,
        pressure=PRESSURE,
        params=params,
        step_size_init=TRACE_STEP_SIZE_INIT,
        newton_tol=TRACE_NEWTON_TOL,
        critical_point_tol=TRACE_CRITICAL_TOL,
        max_step_fraction=TRACE_MAX_STEP_FRACTION,
        verbosity=True,
    )

    return interpret_y_vector(y_lst, PRESSURE, params)


# build the parameter sets

base_sets = [
    (
        "Arndt 2012, optimized",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "arndt"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | opt_arndt_water,
    ),
]


# Calculations

binodal_dicts = []


for current_set in base_sets:

    for mw in MOLAR_WEIGHT_LST:

        current_params = scale_component_molarweight(
            current_set[1], 1, mw, in_place=False
        )
        current_params = add_mixing_association(current_params)

        current_params = current_params | ADD_PARAMS

        current_dict = trace_binodal(current_params, T_TRACE_HIGH, T_TRACE_LOW)

        binodal_dicts.append(current_dict)

    for C in C_LST:

        current_params = scale_component_molarweight(
            current_set[1], 1, MOLAR_WEIGHT_LST[1], in_place=False
        )
        current_params = add_mixing_association(current_params)

        current_params = current_params | ADD_PARAMS

        current_params["C"] = C

        print(C)

        current_dict = trace_binodal(current_params, T_TRACE_HIGH, T_TRACE_LOW)

        binodal_dicts.append(current_dict)


# %% Plot Binodal Curves - combined


color_lst = get_style_colors()
plt.style.use(STYLE_PATH)

fig, (ax1, ax2) = plt.subplots(
    1,
    2,
    figsize=(6.7, 2.2),
)

ax1.set_xlabel(r"Mass fraction pNIPAM [-]")
ax1.set_ylabel(r"T [K]")

ax2.set_xlabel(r"DOS: $1/w$ [-]")
ax2.set_ylabel(r"T [K]")

ax1.set_xlim([0.0, 1.0])
ax1.set_ylim([T_TRACE_LOW, T_TRACE_HIGH])

ax2.set_xlim([1, 50])
ax2.set_ylim([T_TRACE_LOW, T_TRACE_HIGH])


def to_dos_xy(xy):
    """Convert to DOS"""
    x = np.asarray(xy[:, 0])
    y = np.asarray(xy[:, 1])
    m = np.isfinite(x) & np.isfinite(y) & (x > 0)
    return np.column_stack([1.0 / x[m], y[m]])


for set_idx, current_set in enumerate(base_sets[:2]):

    for net_idx, net_term in enumerate(MOLAR_WEIGHT_LST):

        curve = binodal_dicts[set_idx * len(MOLAR_WEIGHT_LST) + net_idx]

        label = (
            rf"$M_n={net_term/1000.0:.0f}\,\mathrm{{kg/mol}}$"
            + rf"$\;C={curve['params']['C']}$"
        )

        color = color_lst[net_idx]
        linestyle = LINESTYLES[set_idx % len(LINESTYLES)]

        binodal_left_wT = binary_coords(
            curve["w_I_lst"],
            curve["temperature_K"],
            which_idx=0,
        )
        binodal_right_wT = binary_coords(
            curve["w_II_lst"],
            curve["temperature_K"],
            which_idx=0,
        )

        # left panel
        plot_binary(
            ax1,
            binodal_left_wT,
            color=color,
            linestyle=linestyle,
            zorder=5,
        )
        plot_binary(
            ax1,
            binodal_right_wT,
            color=color,
            linestyle=linestyle,
            zorder=5,
        )

        # right panel
        binodal_left_dos = to_dos_xy(binodal_left_wT)
        binodal_right_dos = to_dos_xy(binodal_right_wT)

        if binodal_left_dos.size:
            plot_binary(
                ax2,
                binodal_left_dos,
                label=label,
                color=color,
                linestyle=linestyle,
                zorder=5,
            )
        if binodal_right_dos.size:
            plot_binary(
                ax2,
                binodal_right_dos,
                color=color,
                linestyle=linestyle,
                zorder=5,
            )

    for net_idx, net_term in enumerate(C_LST):

        offset_idx = len(MOLAR_WEIGHT_LST)
        curve_idx = offset_idx + set_idx * len(C_LST) + net_idx
        curve = binodal_dicts[curve_idx]

        label = (
            rf"$M_n={MOLAR_WEIGHT_LST[1]/1000.0:.0f}\,\mathrm{{kg/mol}},$"
            + rf"$\;C={net_term}$"
        )

        color = color_lst[2 * net_idx]
        linestyle = "--"

        binodal_left_wT = binary_coords(
            curve["w_I_lst"],
            curve["temperature_K"],
            which_idx=0,
        )
        binodal_right_wT = binary_coords(
            curve["w_II_lst"],
            curve["temperature_K"],
            which_idx=0,
        )

        # left panel
        plot_binary(
            ax1,
            binodal_left_wT,
            label=label,
            color=color,
            linestyle=linestyle,
            zorder=5,
        )
        plot_binary(
            ax1,
            binodal_right_wT,
            color=color,
            linestyle=linestyle,
            zorder=5,
        )

        # right panel
        binodal_left_dos = to_dos_xy(binodal_left_wT)
        binodal_right_dos = to_dos_xy(binodal_right_wT)

        if binodal_left_dos.size:
            plot_binary(
                ax2,
                binodal_left_dos,
                label=label,
                color=color,
                linestyle=linestyle,
                zorder=5,
            )
        if binodal_right_dos.size:
            plot_binary(
                ax2,
                binodal_right_dos,
                color=color,
                linestyle=linestyle,
                zorder=5,
            )


ax2.legend(
    loc="upper right",
    frameon=True,
)

script_path = Path(__file__).resolve()
figures_dir = script_path.parent / "figures"
figures_dir.mkdir(parents=True, exist_ok=True)

plt.savefig(
    figures_dir / f"{script_path.stem}.pdf", dpi=600, bbox_inches="tight"
)
plt.show()
# %%
