# %%
"""Compare PNIPAAm-water gel binodals and DOS curves for several water
parameterizations."""

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,
    load_dos_record,
)

PRESSURE = 1.13
T_TRACE_HIGH = 330.0
T_TRACE_LOW = 290.0

MASS_DENSITY_TEMPERATURE = 298.15
MASS_DENSITY_W_VALS = np.linspace(1e-3, 1 - 1e-3, 50)
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

MOLAR_WEIGHT = 12e3

V_REF_LST = [1.0]

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"

DOS_FILES = [
    BASE_FOLDER_EXP / "EXP_DOS_water_pnipaam_poschlad_0.01.json",
]


dos_records = []
for p in DOS_FILES:
    if Path(p).exists():
        dos_records.append(load_dos_record(p))

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

params_binary_scaled_for_arndt = {
    "comp12_kij_intercept": -0.22192891216679728,
    "comp12_kij_slope": 0.002582638142251112,
    "comp1_epsilon_k_ab": 189.23270263964295,
    "comp1_kappa_ab": 0.19946876491055293,
    "comp1_epsilon_k": 262.27669033429044,
    "comp1_m": 56.36758016284652,
    "comp1_sigma": 5.521036667856562,
}

params_binary_scaled_for_grob = {
    "comp12_kij_intercept": -0.16991118210962414,
    "comp12_kij_slope": 0.002303734908111005,
    "comp1_epsilon_k_ab": 186.80206657240154,
    "comp1_kappa_ab": 0.1503662872679954,
    "comp1_epsilon_k": 227.0564422302229,
    "comp1_m": 112.80467398130423,
    "comp1_sigma": 4.402926225832167,
}

params_binary_scaled_for_esper = {
    "comp12_kij_intercept": 0.07022364567829675,
    "comp12_kij_slope": 0.0015676949621230424,
    "comp1_epsilon_k_ab": 50.95709992256691,
    "comp1_kappa_ab": 0.0876608468377187,
    "comp1_epsilon_k": 362.1240583130229,
    "comp1_m": 40.411743117621555,
    "comp1_sigma": 6.089678228942412,
}

params_binary_scaled_for_bahrabadi = {
    "comp12_kij_intercept": -0.21766778345962773,
    "comp12_kij_slope": 0.0016493415314367266,
    "comp1_epsilon_k_ab": 48.84370869369834,
    "comp1_kappa_ab": 0.11298537261191516,
    "comp1_epsilon_k": 322.2520089375094,
    "comp1_m": 31.48580617576581,
    "comp1_sigma": 6.71024725747544,
}

base_sets = [
    (
        "Arndt Opt.",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "arndt"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | params_binary_scaled_for_arndt,
    ),
    (
        "Grob Opt.",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "grob"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | params_binary_scaled_for_grob,
    ),
    (
        "Jovein Opt.",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "bahrabadi"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | params_binary_scaled_for_bahrabadi,
    ),
    (
        "Esper Opt.",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "esper"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | params_binary_scaled_for_esper,
    ),
]


# Calculations

binodal_dicts = []


for current_set in base_sets:

    for v_ref in V_REF_LST:

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

        current_params = current_params | ADD_PARAMS
        current_params["ref_volume_frac"] = v_ref

        current_dict = trace_binodal(current_params, T_TRACE_HIGH, T_TRACE_LOW)

        binodal_dicts.append(current_dict)


# %% Plot Binodal Curves - combined with DOS


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, 20])
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):

    for net_idx, net_term in enumerate(V_REF_LST):

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

        label = current_set[0]

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

        color = color_lst[set_idx]

        # left panel: mass fraction
        plot_binary(
            ax1,
            binodal_left_wT,
            color=color,
            zorder=5,
        )
        plot_binary(
            ax1,
            binodal_right_wT,
            color=color,
            zorder=5,
        )

        # right panel: DOS
        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,
                zorder=5,
            )
        if binodal_right_dos.size:
            plot_binary(
                ax2,
                binodal_right_dos,
                color=color,
                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()

# %% Plot TOC DOS Binodal Graphic with Legend

toc_colors = {
    r"Arndt Opt.": "#7030A0",
    r"Grob Opt.": "#3B7D23",
    r"Jovein Opt.": "#C04F15",
    r"Esper Opt.": "#215F9A",
}

toc_labels = {
    r"Arndt Opt.": r"Arndt et al.",
    r"Grob Opt.": r"Grob et al.",
    r"Jovein Opt.": r"Jovein et al.",
    r"Esper Opt.": r"Esper et al.",
}

plt.style.use(STYLE_PATH)

_, ax = plt.subplots(
    1,
    1,
    figsize=(3.35, 2.2),
)

ax.set_ylabel(r"$T$", fontsize=14)
ax.set_xlabel(r"$DOS$", fontsize=14)
ax.set_xlim([1, 20])
ax.set_ylim([T_TRACE_LOW, T_TRACE_HIGH])

for set_idx, current_set in enumerate(base_sets):

    set_label = current_set[0]

    if set_label not in toc_colors:
        continue

    for net_idx, net_term in enumerate(V_REF_LST):

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

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

        binodal_left_dos = to_dos_xy(binodal_left_wT)
        binodal_right_dos = to_dos_xy(binodal_right_wT)

        label = toc_labels[set_label] if net_idx == 0 else None

        if binodal_left_dos.size:
            plot_binary(
                ax,
                binodal_left_dos,
                label=label,
                color=toc_colors[set_label],
                linewidth=2.0,
                zorder=5,
            )

        if binodal_right_dos.size:
            plot_binary(
                ax,
                binodal_right_dos,
                color=toc_colors[set_label],
                linewidth=2.0,
                zorder=5,
            )

# experimental DOS points
for rec in dos_records:
    T = np.asarray(rec["temperature_K"])
    w = np.asarray(rec["w_polymer"])

    mask = np.isfinite(w) & (w > 0) & np.isfinite(T)

    ax.scatter(
        1.0 / w[mask],
        T[mask],
        color="black",
        s=14,
        zorder=1,
    )

ax.set_xticks([])
ax.set_xticklabels([])

ax.set_yticks([])
ax.set_yticklabels([])

ax.legend(
    loc="upper right",
    fontsize=12,
    handlelength=1.5,
    handletextpad=0.5,
    borderpad=0.4,
    labelcolor="linecolor",
)

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}_toc_dos.pdf",
    dpi=600,
    bbox_inches="tight",
)

plt.show()

# %%
