# %%
"""Plot PNIPAAm-water binodals and mass-density predictions for
several molar masses."""

from __future__ import annotations

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.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_LST = [5e3, 50e3, 500e3]

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_[AP]*.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.1709406623294618,
    "comp12_kij_slope": 0.013543916811717197,
    "comp12_epsilon_k_ab": 2283.9489103460455,
}
opt_arndt_water = {
    "comp12_kij_intercept": -0.16411647669662233,
    "comp12_kij_slope": 0.001275948739486485,
    # "comp12_epsilon_k_ab": 2283.9489103460455,
}

params_binary = {
    "comp12_kij_intercept": -0.169591782339968,
    "comp12_kij_slope": 0.0011862254418824325,
    "comp1_epsilon_k_ab": 180.0896988893679,
    "comp1_kappa_ab": 0.03979848465351869,
}

params_binary_scaled_for_grob = {
    "comp12_kij_intercept": -0.16278582725153543,
    "comp12_kij_slope": 0.0014890861866361637,
    "comp1_epsilon_k_ab": 434.52906009559996,
    "comp1_kappa_ab": 0.06051154634995117,
}

params_binary_scaled_for_esper = {
    "comp12_kij_intercept": 0.003786830327696722,
    "comp12_kij_slope": 0.001125978172894745,
    "comp1_epsilon_k_ab": 57.7486017638068,
    "comp1_kappa_ab": 0.012611862305695054,
}


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 = [
    (
        "Arndt Lit.",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "arndt"],
            binary_sources={
                (1, 2): "arndt",
            },
        ),
    ),
    (
        "Arndt Opt.",
        get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", "arndt"],
            binary_sources={
                (1, 2): "arndt",
            },
        )
        | opt_arndt_water,
    ),
]


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


def estimate_lcst_from_curve(_curve):
    """Estimate the LCST as lowest point in curve"""
    return int(np.rint(np.nanmin(_curve["temperature_K"])))


table_path = Path.cwd() / "tables" / f"{Path(__file__).stem}_params_table.txt"
table_path.parent.mkdir(exist_ok=True)

rows = []
NO_MWS = len(MOLAR_WEIGHT_LST)

for set_idx, (set_name, raw_params) in enumerate(base_sets):
    kij = raw_params.get("comp12_kij_intercept", np.nan)
    kij_slope = raw_params.get("comp12_kij_slope", np.nan)

    # LCST depends on molar mass in your plot.
    # For a 2-row summary table, use the lowest binodal temperature
    # among all curves belonging to the same parameter set.
    curves_this_set = binodal_dicts[set_idx * NO_MWS : (set_idx + 1) * NO_MWS]
    lcst_vals = [estimate_lcst_from_curve(curve) for curve in curves_this_set]
    lcst = int(np.rint(np.nanmin(lcst_vals)))

    label = set_name
    if "Opt." in set_name:
        label += r" ($k_{ij}(T)$)"

    rows.append(f"{label} & {kij:.3f} & {kij_slope:.3f} & {lcst:d} \\\\")

TABLE_TEXT = "\n".join(
    [
        r"\begin{tabular}{lccc}",
        r"\hline",
        r"- & $k_{ij}$ & Slope $k_{ij}$ & LCST \\",
        r"\hline",
        *rows,
        r"\hline",
        r"\end{tabular}",
        "",
    ]
)

table_path.write_text(TABLE_TEXT, encoding="utf-8")
print(f"Saved: {table_path}")

# 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()
"""Plot binodal curves and experimental points."""
_, 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[mw_idx],
            # marker=MARKERS[mw_idx % len(MARKERS)],
            linestyle=LINESTYLES[set_idx % len(LINESTYLES)],
            zorder=5,
        )
        plot_binary(
            ax,
            binodal_right,
            color=color_lst[mw_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"
)
