# %%
"""Compare PNIPAAm/water/acetic acid gel swelling traces with DOS data
for fitted models."""

from __future__ import annotations

import json
from pathlib import Path
import re
import matplotlib.pyplot as plt
import numpy as np


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 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_water_aceticacid_pnipaam_orlov_0.01_xigel0.08.json",
    BASE_FOLDER_EXP
    / "EXP_DOS_water_aceticacid_pnipaam_orlov_0.015_xigel0.08.json",
    BASE_FOLDER_EXP
    / "EXP_DOS_water_aceticacid_pnipaam_orlov_0.02_xigel0.08.json",
    # BASE_FOLDER_EXP
    # / "EXP_DOS_water_aceticacid_pnipaam_orlov_0.015_xigel0.10.json",
    # BASE_FOLDER_EXP
    # / "EXP_DOS_pnipaam_aceticacid_butylacetate_eckert_0.0138.json",
    # BASE_FOLDER_EXP
    # / "EXP_DOS_pnipaam_butanol_butylacetate_eckert_0.0138.json",
]

# -----------------------------
# Core computations
# -----------------------------
base_sets = [
    (
        "Jovein Water",
        "optimize_parameter_gel_other_solvents/pnipaam_best_params_W_AA.json",
        r"Model Opt. $k_\mathrm{AA-P}, k_\mathrm{AA-W}$",
    ),
    (
        "Jovein Water",
        "optimize_parameter_gel_other_solvents/pnipaam_best_params"
        + "_just13.json",
        r"Model Opt. $k_\mathrm{AA-P}$",
    ),
]


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

# -----------------------------
# 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=None,
        )

        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 Acetic Acid [-]",
    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 % len(DOS_FILES)],
        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")
    if w is None:
        w = rec.get("w_3_mol_per_mol")
    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=colors[rec_idx % len(DOS_FILES)],
        zorder=2,
    )

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

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")
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_aa_p = params["comp13_kij_intercept"]
    k_aa_w = params["comp23_kij_intercept"]
    score = params.get("score", "")

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

    # first row: italic second entry
    if i == 1:
        k_aa_w_str = rf"\textit{{{k_aa_w:.3f}}}"
    else:
        k_aa_w_str = f"{k_aa_w:.3f}"

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

LATEX_TABLE = "\n".join(
    [
        r"\begin{tabular}{lcc}",
        r"\hline",
        r"Optimized parameters & $k_\mathrm{AA-P}$ & $k_\mathrm{AA-W}$\\",
        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")


# %%
