# %%
"""Compare PNIPAAm/butyl acetate/n-butanol gel swelling traces with
DOS data for fitted models."""

from __future__ import annotations
import re

import json
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.parameters import (
    add_mixing_association,
    get_subset_of_parameters,
)
from equi_gel.plot_binary import (
    binary_axes,
    get_style_colors,
)
from equi_gel.save_load_json import (
    BASE_FOLDER_EXP,
    load_dos_record,
)
from equi_gel.cost_func_binodal import get_minimal_distance_swelling

from gel_tracer_ternary import trace_binodal_gel

# -----------------------------
# Config
# -----------------------------
T_TRACE_HIGH = 350.0
T_TRACE_LOW = 298.0


NETWORK_TERMS = ["Flory"]
V_MAX_LOW_LST = [0.001, 0.01, 0.1]
MARKERS = ["o", "s", "^", "v", "D", "1", "2", "x", "+", "*"]
LINESTYLES = ["-", ":", "-"]


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

# -----------------------------
# Core computations
# -----------------------------
base_sets = [
    (
        "Jovein Water",
        "optimize_parameter_gel_other_solvents/pnipaam_best_params_BA_B_"
        + "comp13_kij_intercept.json",
        r"Model Opt. $k_\mathrm{P-B}$",
    ),
    (
        "Jovein Water",
        "optimize_parameter_gel_other_solvents/pnipaam_best_params_BA_B_"
        + "comp13_kij_intercept_comp23_kij_intercept.json",
        r"Model Opt. $k_\mathrm{P-B}, k_\mathrm{B-BA}$",
    ),
]

initial_guess_y_vec = np.load(
    "optimize_parameter_gel_other_solvents/pnipaam_best_params_AA_BA_"
    + "last_y_lst.npy"
)
addings = {}

# -----------------------------
# Calculations
# -----------------------------
# -----------------------------
# Load experimental DOS data (Poschlad)
# -----------------------------
dos_records = []
for p in DOS_FILES:
    if Path(p).exists():
        dos_records.append(load_dos_record(p))
        current_parameter_set = json.loads(
            Path(base_sets[0][1]).read_text(encoding="utf-8")
        )
        if "w_3_mol_per_mol" in dos_records[-1]:
            binary_parameters = get_subset_of_parameters(
                current_parameter_set, [2, 3]
            )
            binary_parameters = add_mixing_association(binary_parameters)
            dos_records[-1]["w_3_g_per_g"] = [
                get_feos_npt_state(
                    T_TRACE_HIGH, 1.13, binary_parameters, xi_lst=[1.0 - x, x]
                ).massfracs()[1]
                for x in dos_records[-1]["w_3_mol_per_mol"]
            ]

            # invert the list
            dos_records[-1]["w_3_g_per_g"] = dos_records[-1]["w_3_g_per_g"][
                ::-1
            ]
            dos_records[-1]["w_3_mol_per_mol"] = dos_records[-1][
                "w_3_mol_per_mol"
            ][::-1]

        w_II_lst = [[0.0, 1.0 - w, w] for w in dos_records[-1]["w_3_g_per_g"]]

# -----------------------------
# Calculations: one binodal per DOS record (use its crosslinker)
# -----------------------------
binodal_dicts = []  # aligned with dos_records
labels_model = []

for base_set in base_sets:

    label_part1 = base_set[2]

    for rec in dos_records:
        y_cl = float(rec["crosslinker"])

        current_parameter_set = json.loads(
            Path(base_set[1]).read_text(encoding="utf-8")
        )

        binodal = trace_binodal_gel(
            current_parameter_set,
            t_high=T_TRACE_HIGH,
            t_low=T_TRACE_LOW,
            y_crosslinker=y_cl,
            verbosity=True,
            y_init=initial_guess_y_vec,
        )

        score_gel = get_minimal_distance_swelling(binodal, rec)

        binodal_dicts.append(binodal)
        labels_model.append(f"{label_part1} (y={y_cl:g})")

# -----------------------------
# %% Plot
# -----------------------------
colors = get_style_colors()
_, ax = binary_axes(
    xlabel=r"Mass fraction Butyl Acetate [-]",
    ylabel=r"DOS [g/g]",
)
ax.set_xlim([0.0, 1.0])

# model curves
for idx, (curve, lab) in enumerate(zip(binodal_dicts, labels_model)):
    w_poly = np.array([w[0] for w in curve["w_II_lst"]], dtype=float)
    w_aa = np.array([w[2] for w in curve["w_I_lst"]], dtype=float)

    mask = (
        np.isfinite(w_poly)
        & np.isfinite(w_aa)
        & (w_poly > 1e-12)
        & (w_aa >= 0.0)
        & (w_aa <= 1.0)
    )
    dos_model = 1.0 / w_poly[mask]

    ax.plot(
        w_aa[mask],
        dos_model,
        # marker="X",
        label=lab,
        color=colors[idx],
        # linestyle=LINESTYLES[int(idx / len(DOS_FILES))],
    )

# experimental points
for rec_idx, rec in enumerate(dos_records):
    dos_exp = np.asarray(rec["DOS_g_per_g"], dtype=float)

    w = rec.get("w_3_g_per_g")
    w = np.asarray(w, dtype=float)

    y_cl = float(rec["crosslinker"])
    src = rec.get("source", "").split(",")[0].strip()
    lab = f"Exp. {src} (y={y_cl:g})" if src else f"Exp (y={y_cl:g})"

    mask = np.isfinite(w) & np.isfinite(dos_exp)
    ax.scatter(
        w[mask],
        dos_exp[mask],
        marker=MARKERS[rec_idx % len(MARKERS)],
        label=lab,
        color="black",
        zorder=2,
    )

ax.legend(
    loc="upper left",
    bbox_to_anchor=(1.02, 1.0),
    ncols=1,
    borderaxespad=0.0,
)
ax.set_ylim(bottom=0.0)

figures_dir = Path("figures")
figures_dir.mkdir(parents=True, exist_ok=True)

outfile = figures_dir / f"{Path(__file__).stem}.pdf"
plt.savefig(outfile, dpi=600, bbox_inches="tight")
print(f"Saved: {outfile.resolve()}")

# %%
# %%
# -----------------------------
# Build LaTeX table (minimal)
# -----------------------------
rows = []

for i, entry in enumerate(base_sets):
    params = json.loads(Path(entry[1]).read_text(encoding="utf-8"))

    k_b_p = params["comp13_kij_intercept"]
    k_b_ba = params["comp23_kij_intercept"]
    score = params.get("score", "")

    tabel_lable = re.search(r"\$.*?\$", entry[2]).group(0)

    # first row: italic second entry
    if i == 0:
        k_b_ba_str = rf"\textit{{{k_b_ba:.3f}}}"
    else:
        k_b_ba_str = f"{k_b_ba:.3f}"

    row = f"{tabel_lable} & " f"{k_b_p:.3f} & " f"{k_b_ba_str} \\\\"
    rows.append(row)

LATEX_TABLE = "\n".join(
    [
        r"\begin{tabular}{lcc}",
        r"\hline",
        r"Optimized parameters & $k_\mathrm{B-P}$ & $k_\mathrm{B-BA}$\\",
        r"\hline",
        *rows,
        r"\hline",
        r"\end{tabular}",
    ]
)
tables_dir = Path("tables")
tables_dir.mkdir(parents=True, exist_ok=True)

outfile = tables_dir / f"{Path(__file__).stem}_table.tex"
outfile.write_text(LATEX_TABLE, encoding="utf-8")

# %%
