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

from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

from equi_gel.parameters import (
    get_parameter_set_from_tables,
)
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,
)

from gel_tracer_binary import trace_binodal_binary_gel

# -----------------------------
# Config
# -----------------------------
PRESSURE = 1.13

T_TRACE_HIGH = 350.0
T_TRACE_LOW = 270.0


NETWORK_TERMS = ["Flory", "Phantom", "Miao"]
NETWORK_TERM_NAMES = ["Affine", "Phantom", "FE"]
V_MAX_LOW_LST = [0.001, 0.01, 0.1]

C_FAC = 1
NETWORK_FUNCTIONALITY = 2.69
V_MAX_FAC = 1.0
Y_CROSSLINKER = 0.01


ADD_PARAMS = {
    "number_components": 2,
    "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,
}

MOLAR_WEIGHT = -1

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

# -----------------------------
# Build parameter sets
# -----------------------------
opt_arndt_water = {}

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

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

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


# -----------------------------
# Calculations
# -----------------------------
binodal_dicts = []
third_curve_dicts = []

for set_name, base_params in base_sets:
    for net_idx, network_theory in enumerate(NETWORK_TERMS):

        params = base_params | ADD_PARAMS
        params["elasticity"] = network_theory

        current_dict, _ = trace_binodal_binary_gel(
            params, t_high=T_TRACE_HIGH, t_low=T_TRACE_LOW
        )

        binodal_dicts.append(current_dict)


# -----------------------------
# Load experimental DOS data (Poschlad)
# -----------------------------
dos_records = []
for p in DOS_FILES:
    if Path(p).exists():
        dos_records.append(load_dos_record(p))


# %%Plot
# -----------------------------
# Two subfigures side by side
# Left: mass fraction vs T
# Right: DOS-space (1/w) vs T
# Legend only in left subplot, using labels from first plot
# -----------------------------
color_lst = get_style_colors()


plt.style.use(STYLE_PATH)
fig, (ax1, ax2) = plt.subplots(
    1,
    2,
    figsize=(6.7, 2.2),
    # constrained_layout=True,
)

# match your usual axis styling
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 mass fraction to DOS space."""
    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, (set_name, _) in enumerate(base_sets):
    for net_idx, net_term in enumerate(NETWORK_TERMS):
        curve = binodal_dicts[set_idx * len(NETWORK_TERMS) + net_idx]

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

        if net_term == "Miao":
            color = color_lst[2]
            LINESTYLE = ["-", "--", "-."][net_idx - 2]
            LABEL = rf"{NETWORK_TERM_NAMES[net_idx]}"
        else:
            color = color_lst[net_idx]
            LINESTYLE = "-"
            LABEL = rf"{NETWORK_TERM_NAMES[net_idx]}"

        # -----------------------------
        # LEFT SUBPLOT: mass fraction vs T
        # -----------------------------
        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 SUBPLOT: DOS-space
        # -----------------------------
        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,
            )


# -----------------------------
# Experimental points
# -----------------------------
for rec in dos_records:
    T = np.asarray(rec["temperature_K"])
    w = np.asarray(rec["w_polymer"])
    cl = rec["crosslinker"]
    lab = f"Poschlad (y={cl})"

    # left subplot
    ax1.scatter(w, T, zorder=1, color="black")

    # right subplot
    mask = np.isfinite(w) & (w > 0) & np.isfinite(T)
    inv_w = 1.0 / w[mask]
    Tm = T[mask]

    ax2.scatter(inv_w, Tm, label=lab, zorder=1, color="black")

# legend only in left subplot, inside plot area
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)

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