# %%
"""Compare PNIPAAm-water binodals and densities across several water
parameterizations."""

from __future__ import annotations

import json
import re
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
from equi_gel.interface_feos import get_feos_npt_state
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 (
    binary_axes,
    binary_coords,
    plot_binary,
    scatter_binary,
    get_style_colors,
)
from equi_gel.save_load_json import (
    BASE_FOLDER_EXP,
    load_exp_pnipaam_water_all,
    load_exp_pnipaam_density_kany,
)
from si_units import GRAM, METER

PRESSURE = 1.13
T_TRACE_HIGH = 350.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.07  # g / ccm SCHILD 1992

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

MOLAR_WEIGHT_LST = [5e3]

ADD_PARAMS = {
    "number_components": 2,
    "gel": False,
    "elasticity": "Flory",
    "C": 40,
    "ref_volume_frac": 1.0,
    "max_volume_poly_frac": 0.00003,
}

EXP_PNIPAAM_WATER_PATTERN = "EXP_SOL_PNIPAAM_water_Afroze_Mn36*.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_kij = {
    "comp12_kij_intercept": -0.16411647669662233,
    "comp12_kij_slope": 0.001275948739486485,
}


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)


def calculate_mass_density_curve(_params, temperature, w_vals):
    """Return mass density (g/cm^3) over weight fractions."""
    densities = []
    for _w1 in w_vals:
        wi_lst = [float(_w1), float(1.0 - _w1)]
        try:
            density = get_feos_npt_state(
                temperature=temperature,
                pressure=PRESSURE,
                params=_params,
                wi_lst=wi_lst,
            ).mass_density() / (GRAM / (METER * 1e-2) ** 3)
        except Exception:  # pylint: disable=broad-exception-caught
            density = np.nan
        densities.append(density)
    return np.array(densities, dtype=float)


# build the parameter sets

base_sets = [
    (
        r"Arndt Opt. ($k_{ij}(T)$)",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "arndt"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | opt_arndt_water_kij,
    ),
    (
        r"Arndt Opt.",
        json.loads(
            Path(
                "optimization_water_models/pnipaam_water_arndt_const_m.json"
            ).read_text(encoding="utf-8")
        ),
    ),
    (
        "Grob Opt.",
        json.loads(
            Path(
                "optimization_water_models/pnipaam_water_grob_const_m.json"
            ).read_text(encoding="utf-8")
        ),
    ),
    (
        "Jovein Opt.",
        json.loads(
            Path(
                "optimization_water_models/"
                + "pnipaam_water_bahrabadi_const_m.json"
            ).read_text(encoding="utf-8")
        ),
    ),
    (
        "Esper Opt.",
        json.loads(
            Path(
                "optimization_water_models/pnipaam_water_esper_const_m.json"
            ).read_text(encoding="utf-8")
        ),
    ),
]

# Calculations

binodal_dicts = []
mass_densities = []
for current_set in base_sets:

    for polymer_mw in MOLAR_WEIGHT_LST:

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

        current_params = current_params | ADD_PARAMS

        binodal_dicts.append(
            trace_binodal(current_params, T_TRACE_HIGH, T_TRACE_LOW)
        )

        mass_densities.append(
            calculate_mass_density_curve(
                current_params, MASS_DENSITY_TEMPERATURE, MASS_DENSITY_W_VALS
            )
        )


# load experimental data

exp_binodal_dicts = load_exp_pnipaam_water_all(EXP_PNIPAAM_WATER_PATTERN)
exp_mass_densities = load_exp_pnipaam_density_kany(EXP_PNIPAAM_DENS_KANY)

# %% Plot Binodal Curves
color_lst = get_style_colors()
color_lst.append("0.9")

_, ax = binary_axes(
    xlabel=r"Mass fraction pNIPAM [-]",
    ylabel=r"T [K]",
)
ax.set_xlim([0.0, 1.0])

for set_idx, current_set in enumerate(base_sets):

    for mw_idx, polymer_mw in enumerate(MOLAR_WEIGHT_LST):

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

        label = (
            current_set[0]
            + rf" $M_n={polymer_mw/1000.0:.0f}\, \mathrm{{kg/mol}}$"
        )

        binodal_left = binary_coords(
            curve["w_I_lst"],
            curve["temperature_K"],
            which_idx=0,
        )
        binodal_right = binary_coords(
            curve["w_II_lst"],
            curve["temperature_K"],
            which_idx=0,
        )
        plot_binary(
            ax,
            binodal_left,
            label=label,
            color=color_lst[set_idx],
            # marker=MARKERS[mw_idx % len(MARKERS)],
            # linestyle=LINESTYLES[set_idx % len(LINESTYLES)],
            zorder=5,
        )
        plot_binary(
            ax,
            binodal_right,
            color=color_lst[set_idx],
            # marker=MARKERS[mw_idx % len(MARKERS)],
            # linestyle=LINESTYLES[set_idx % len(LINESTYLES)],
            zorder=5,
        )

for idx, record in enumerate(exp_binodal_dicts):

    first_author = record["source"].split(",")[0]
    year = re.search(
        r"\b(19[5-9]\d|20[0-2]\d|2030)\b", record["source"]
    ).group(0)
    mw = record["Mn_g_mol"]

    label = (
        first_author
        + " "
        + year
        + rf", $M_n={mw/1000.0:.0f}\, \mathrm{{kg/mol}}$"
    )

    exp_xy = binary_coords(
        record["w_I_lst"],
        record["temperature_K"],
        which_idx=0,
    )
    marker = MARKERS[idx % len(MARKERS)]
    face_color = FACE_MARKERS[idx % len(FACE_MARKERS)]
    scatter_binary(
        ax,
        exp_xy,
        marker=marker,
        s=18,
        facecolors=face_color,
        edgecolors="black",
        zorder=3,
        label=label,
    )

ax.set_ylim([T_TRACE_LOW, T_TRACE_HIGH])
ax.legend(
    loc="upper left",
    bbox_to_anchor=(1.02, 1.0),
    ncols=1,
    borderaxespad=0.0,
)

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"
)
# %% Plot Mass Densities

_, ax = binary_axes(
    xlabel=r"Mass fraction pNIPAM [-]",
    ylabel=r"Mass density [g/cm$^3$]",
)
ax.set_xlim([0.0, 1.0])


for set_idx, current_set in enumerate(base_sets):

    mw_idx = 1  # medium molar weight

    label = (
        current_set[0] + rf" $M_n={polymer_mw/1000.0:.0f}\, \mathrm{{kg/mol}}$"
    )

    ax.plot(
        MASS_DENSITY_W_VALS,
        mass_densities[set_idx * len(MOLAR_WEIGHT_LST)],
        label=label,
        color=color_lst[set_idx],
    )


for idx, record in enumerate(exp_mass_densities):
    exp_temperature = float(record["Constant Value"][0][1])
    mn = float(record["Constant Value"][1][1])

    data_rows = record.get("Data Table", [])
    data = data_rows[1:]
    w1 = [float(row[0]) for row in data]
    rho = [float(row[1]) for row in data]

    label_parts = []
    label_parts.append(rf"$M_n={mn/1000.0:.0f}\,\mathrm{{kg/mol}}$")
    label_parts.append(rf"$T={exp_temperature:.2f}\,\mathrm{{K}}$")
    label = f"Kany 1998, {', '.join(label_parts)}"

    ax.scatter(
        w1,
        rho,
        marker=MARKERS[idx % len(MARKERS)],
        s=24,
        facecolors="none",
        edgecolors="black",
        label=label,
        zorder=3,
    )
ax.scatter(
    [0.995],  # visibility
    [AMORPH_POLYMER_DENSITY],
    marker=MARKERS[-1],
    s=24,
    facecolors="black",
    edgecolors="black",
    label="Schild 1992, amorph. solid",
    zorder=3,
)
ax.legend(
    loc="upper left",
    bbox_to_anchor=(1.02, 1.0),
    ncols=1,
    borderaxespad=0.0,
)

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


# %% Save Parameter Table


def fmt_or_dash(value, fmt):
    """Formatting helper"""
    if value is None:
        return "--"
    return format(float(value), fmt)


lit_params = get_parameter_set_from_tables(
    ["pnipaam", "water"],
    ["arndt", "arndt"],
    binary_sources={
        (1, 2): "arndt",
    },
)

table_sets = [("Arndt Lit.", lit_params), *base_sets]

table_rows = []
for label, params in table_sets:
    row = [
        label,
        fmt_or_dash(params.get("comp12_kij_intercept"), ".2f"),
        fmt_or_dash(params.get("comp12_kij_slope"), ".4f"),
        fmt_or_dash(params.get("comp1_epsilon_k_ab"), ".2f"),
        fmt_or_dash(params.get("comp1_kappa_ab"), ".3f"),
        fmt_or_dash(params.get("comp1_epsilon_k"), ".2f"),
        fmt_or_dash(params.get("comp1_m"), ".2f"),
        fmt_or_dash(params.get("comp1_sigma"), ".2f"),
    ]
    table_rows.append(" & ".join(row) + r" \\")

TABLE_TEXT = "\n".join(
    [
        r"\begin{tabular}{lrrrrrrr}",
        r"\hline",
        r" & $k_{ij}$ & Slope $k_{ij}$ & $\varepsilon^\mathrm{AB}_k$ &"
        + r" $\kappa^\mathrm{AB}$ & $\varepsilon_k$ & $m$ & $\sigma$\\",
        r"\hline",
        *table_rows,
        r"\hline",
        r"\end{tabular}",
        "",
    ]
)
table_path = Path.cwd() / "tables" / f"{Path(__file__).stem}_params_table.txt"
table_path.parent.mkdir(exist_ok=True)
table_path.write_text(TABLE_TEXT, encoding="utf-8")

print(f"Saved table to: {table_path}")
# %% TOC Graphic


# %% Plot TOC Binodal Graphic
# %% Plot TOC Binodal Graphic

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

_, ax = binary_axes(
    xlabel="",
    ylabel=r"$T$",
)
ax.set_ylabel(r"$T$", fontsize=14)
ax.set_xlim([0.0, 1.0])

for set_idx, current_set in enumerate(base_sets):

    set_label = current_set[0]

    if set_label not in toc_colors:
        continue

    for mw_idx, polymer_mw in enumerate(MOLAR_WEIGHT_LST):

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

        binodal_left = binary_coords(
            curve["w_I_lst"],
            curve["temperature_K"],
            which_idx=0,
        )
        binodal_right = binary_coords(
            curve["w_II_lst"],
            curve["temperature_K"],
            which_idx=0,
        )

        plot_binary(
            ax,
            binodal_left,
            color=toc_colors[set_label],
            linewidth=2.0,
            zorder=5,
        )
        plot_binary(
            ax,
            binodal_right,
            color=toc_colors[set_label],
            linewidth=2.0,
            zorder=5,
        )

for idx, record in enumerate(exp_binodal_dicts):

    exp_xy = binary_coords(
        record["w_I_lst"],
        record["temperature_K"],
        which_idx=0,
    )

    scatter_binary(
        ax,
        exp_xy,
        marker=MARKERS[idx % len(MARKERS)],
        s=28,
        facecolors="white",
        edgecolors="black",
        linewidths=1.2,
        zorder=2,
    )

ax.set_ylim([T_TRACE_LOW, T_TRACE_HIGH])

ax.set_xlabel(r"$w$", fontsize=14)
ax.set_xticks([])
ax.set_xticklabels([])

ax.set_ylabel(r"$T$")
ax.set_yticks([])
ax.set_yticklabels([])

if ax.legend_ is not None:
    ax.legend_.remove()

out_dir = Path.cwd() / "figures"
out_dir.mkdir(exist_ok=True)

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