"""Benchmark PNIPAAm-water gel models against DOS data and export
optimization score tables."""

# pylint: disable=duplicate-code

from __future__ import annotations
import csv
import json
from pathlib import Path
import numpy as np
from equi_gel.optimize_parameters import (
    ObjectiveWrapperBase,
    ParameterOptimization,
)
from equi_gel.cost_func_binodal import get_minimal_distance_swelling

from equi_gel.newton_raphson_tracing import (
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
)

from equi_gel.parameters import (
    get_parameter_set_from_tables,
    scale_component_molarweight,
)

from equi_gel.save_load_json import (
    BASE_FOLDER_EXP,
    load_dos_record,
)

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

PRESSURE = 1.13
T_HIGH = 350
T_LOW = 290


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

# -----------------------------
# Load experimental DOS data (Poschlad)
# -----------------------------
dos_records = []
for p in DOS_FILES:
    if Path(p).exists():
        dos_records.append(load_dos_record(p))
        w_I_lst = [[w, 1.0 - w] for w in dos_records[-1]["w_polymer"]]
        dos_records[-1]["w_I_lst"] = w_I_lst.copy()
        dos_records[-1]["type"] = "solubilities"

# -----------------------------
# Filter out DOS data with DOS < threshold
# -----------------------------
THRESHOLD = 3.0
for rec in dos_records:
    dos = rec.get("DOS_g_per_g")
    if not isinstance(dos, list) or not dos:
        continue

    mask = [d >= THRESHOLD for d in dos]
    n = len(dos)

    for k, v in list(rec.items()):
        if isinstance(v, list) and len(v) == n:
            rec[k] = [item for item, keep in zip(v, mask) if keep]


class GelObjectiveWrapper(ObjectiveWrapperBase):
    """Objective wrapper for PNIPAAm-water parameter optimization."""

    def __init__(self, _dos_records, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._dos_records = _dos_records

    def objective(self, current_args: dict, verbosity: bool = False) -> float:
        """Evaluate performance of PNIPAAm-water parameter set."""
        try:
            score_gel = 0
            for dos_record in self._dos_records:
                network_functionality = current_args["NETWORK_FUNCTIONALITY"]
                y_crosslinker = dos_record["crosslinker"]

                _c = (network_functionality - 2) / network_functionality
                m_desired = (
                    (2 / network_functionality)
                    * (1 - y_crosslinker)
                    / y_crosslinker
                )
                v_max = current_args["V_MAX_FAC"] * 0.94 * m_desired**2

                current_args["C"] = _c
                current_args["max_volume_poly_frac"] = 1.0 / v_max

                wrong_m = current_args["comp1_m"]
                wrong_mw = current_args["comp1_molarweight"]

                molar_weight = (m_desired / wrong_m) * wrong_mw

                current_args = scale_component_molarweight(
                    current_args, 1, molar_weight, in_place=False
                )

                y_init_binary = build_initial_guess_binary(
                    T_HIGH, PRESSURE, current_args
                )

                y_lst, *_ = trace_parameter(
                    trace_index=-1,
                    y_init=y_init_binary,
                    trace_start=T_HIGH,
                    trace_target=T_LOW,
                    pressure=PRESSURE,
                    params=current_args,
                    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,
                )
                binodal_dict = interpret_y_vector(
                    y_lst, PRESSURE, current_args
                )

                score_gel += get_minimal_distance_swelling(
                    binodal_dict,
                    dos_record,
                )

            return score_gel

        except Exception as exc:  # pylint: disable=broad-exception-caught
            print("Error:", exc)
            return float("inf")


original_parameter_set = get_parameter_set_from_tables(
    ["pnipaam", "water"],
    ["arndt", "arndt"],
    binary_sources={(1, 2): "arndt"},
) | {
    "number_components": 2,
    "gel": True,
    "elasticity": "Miao",
    "NETWORK_FUNCTIONALITY": 2.24,
    "Y_CROSSLINKER": 0.01,
    "V_MAX_FAC": 0.5,
    "ref_volume_frac": 1.0,
}
MAX_CHANGE_NET_FUNC = 2
MAX_CHANGE_V_MAX_FAC = 0.499

optimize_sets = [
    {
        "name": "Original Flory",
        "latex_section": "Original publication",
        "latex_row": "Affine",
        "current_params": original_parameter_set | {"elasticity": "Flory"},
        "keys_to_optimize": None,
        "max_change_of_parameters": None,
        "fitting_to": None,
    },
    {
        "name": "Original Phantom",
        "latex_section": "Original publication",
        "latex_row": "Phantom",
        "current_params": original_parameter_set | {"elasticity": "Phantom"},
        "keys_to_optimize": None,
        "max_change_of_parameters": None,
        "fitting_to": None,
    },
    {
        "name": "Original Miao",
        "latex_section": "Original publication",
        "latex_row": "FE",
        "current_params": original_parameter_set
        | {"elasticity": "Miao", "V_MAX_FAC": 1.0},
        "keys_to_optimize": None,
        "max_change_of_parameters": None,
        "fitting_to": None,
    },
    {
        "name": "One file Flory",
        "latex_section": "Optimized for y=0.01",
        "latex_row": "Affine",
        "current_params": original_parameter_set | {"elasticity": "Flory"},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY"],
        "max_change_of_parameters": [MAX_CHANGE_NET_FUNC],
        "fitting_to": [dos_records[0]],
    },
    {
        "name": "One file Phantom",
        "latex_section": "Optimized for y=0.01",
        "latex_row": "Phantom",
        "current_params": original_parameter_set | {"elasticity": "Phantom"},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY"],
        "max_change_of_parameters": [MAX_CHANGE_NET_FUNC],
        "fitting_to": [dos_records[0]],
    },
    {
        "name": "One file Miao",
        "latex_section": "Optimized for y=0.01",
        "latex_row": "FE",
        "current_params": original_parameter_set
        | {"elasticity": "Miao", "V_MAX_FAC": 1.0},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY"],
        "max_change_of_parameters": [MAX_CHANGE_NET_FUNC],
        "fitting_to": [dos_records[0]],
    },
    {
        "name": "One file Miao + VMAX",
        "latex_section": "Optimized for y=0.01",
        "latex_row": r"FE (adjusted $V_{\max}$)",
        "current_params": original_parameter_set | {"elasticity": "Miao"},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY", "V_MAX_FAC"],
        "max_change_of_parameters": [
            MAX_CHANGE_NET_FUNC,
            MAX_CHANGE_V_MAX_FAC,
        ],
        "fitting_to": [dos_records[0]],
    },
    {
        "name": "Two file Flory",
        "latex_section": "Optimized on both",
        "latex_row": "Affine",
        "current_params": original_parameter_set | {"elasticity": "Flory"},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY"],
        "max_change_of_parameters": [MAX_CHANGE_NET_FUNC],
        "fitting_to": dos_records,
    },
    {
        "name": "Two file Phantom",
        "latex_section": "Optimized on both",
        "latex_row": "Phantom",
        "current_params": original_parameter_set | {"elasticity": "Phantom"},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY"],
        "max_change_of_parameters": [MAX_CHANGE_NET_FUNC],
        "fitting_to": dos_records,
    },
    {
        "name": "Two file Miao",
        "latex_section": "Optimized on both",
        "latex_row": "FE",
        "current_params": original_parameter_set
        | {"elasticity": "Miao", "V_MAX_FAC": 1.0},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY"],
        "max_change_of_parameters": [MAX_CHANGE_NET_FUNC],
        "fitting_to": dos_records,
    },
    {
        "name": "Two file Miao + VMAX",
        "latex_section": "Optimized on both",
        "latex_row": r"FE (adjusted $V_{\max}$)",
        "current_params": original_parameter_set | {"elasticity": "Miao"},
        "keys_to_optimize": ["NETWORK_FUNCTIONALITY", "V_MAX_FAC"],
        "max_change_of_parameters": [
            MAX_CHANGE_NET_FUNC,
            MAX_CHANGE_V_MAX_FAC,
        ],
        "fitting_to": dos_records,
    },
]


if __name__ == "__main__":

    # Directory name based on this file's name without ".py"
    out_dir = Path(__file__).with_suffix("")
    out_dir.mkdir(parents=True, exist_ok=True)

    for optimize_set in optimize_sets:
        print("Running optimization set:", optimize_set["name"])
        current_parameter_set = optimize_set["current_params"]
        keys_to_optimize = optimize_set["keys_to_optimize"]
        max_change_of_parameters = optimize_set["max_change_of_parameters"]
        fitting_to = optimize_set["fitting_to"]

        if keys_to_optimize:
            initial_guess = np.array(
                [current_parameter_set[key] for key in keys_to_optimize]
            )

            objective_wrapper = GelObjectiveWrapper(
                fitting_to, current_parameter_set, keys_to_optimize
            )
            optimizer = ParameterOptimization(objective_wrapper)

            res = optimizer.run_optimization(
                initial_guess=initial_guess,
                max_change_of_parameters=max_change_of_parameters,
                verbosity=True,
                de_kwargs={
                    "popsize": 8,
                    "maxiter": 100,
                    "prefix_file": f"{__file__[:-3]}/{optimize_set['name']}_",
                },
                population_kwargs={
                    "popsize": 8,
                    "sigma": 0.05,
                    "sigma_random": 0.02,
                },
            )
            current_parameter_set |= {
                k: v for k, v in zip(keys_to_optimize, res["best_unscaled"][0])
            }

        # Objective file 1:
        objective_wrapper_1 = GelObjectiveWrapper(
            [dos_records[0]], current_parameter_set, keys_to_optimize
        )
        optimize_set["score_1"] = objective_wrapper_1.objective(
            current_parameter_set
        )

        # Objective file 2:
        objective_wrapper_2 = GelObjectiveWrapper(
            [dos_records[1]], current_parameter_set, keys_to_optimize
        )
        optimize_set["score_2"] = objective_wrapper_2.objective(
            current_parameter_set
        )

        def _jsonable(x):
            if isinstance(x, (np.floating, np.integer)):
                return x.item()
            if isinstance(x, np.ndarray):
                return x.tolist()
            return x

        outpath = Path(f"{__file__[:-3]}/{optimize_set['name']}_params.json")
        outpath.write_text(
            json.dumps(current_parameter_set, indent=2, default=_jsonable),
            encoding="utf-8",
        )
        print(f"Saved: {outpath.resolve()}")

    # Write summary table
    scores_path = out_dir / "00_scores.csv"
    fieldnames = [
        "Name",
        "elasticity",
        "NETWORK_FUNCTIONALITY",
        "ref_volume_frac",
        "V_MAX_FAC",
        "Score 1",
        "Score 2",
    ]

    with scores_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()

        for s in optimize_sets:
            params = s["current_params"]
            w.writerow(
                {
                    "Name": s.get("name", ""),
                    "elasticity": params.get("elasticity", ""),
                    "NETWORK_FUNCTIONALITY": params.get(
                        "NETWORK_FUNCTIONALITY", ""
                    ),
                    "ref_volume_frac": params.get("ref_volume_frac"),
                    "V_MAX_FAC": params.get("V_MAX_FAC"),
                    "Score 1": s.get("score_1", ""),
                    "Score 2": s.get("score_2", ""),
                }
            )

    def _get_phi(_params):
        return f"{_params.get('NETWORK_FUNCTIONALITY', ''):.2f}"

    def _get_phi_ref(_params):
        return f"{_params.get('ref_volume_frac'):.1f}"

    def _get_factor_max(_params):
        if _params.get("elasticity") == "Miao":
            return f"{_params.get('V_MAX_FAC')*100:.0f}$\\%$"
        else:
            return "-"

    def _score(opt_set, which):
        return f"{opt_set.get(which, ''):.2f}"

    def write_latex_table(_optimize_sets, out_path: Path):
        """Write the latex table"""
        # Order sections
        section_order = [
            "Original publication",
            "Optimized for y=0.01",
            "Optimized on both",
        ]

        section_titles = {
            "Original publication": r"\multicolumn{6}{l}"
            + r"{\textit{Original publication}} \\",
            "Optimized for y=0.01": r"\multicolumn{6}{l}"
            + r"{\textit{Optimized for $y=0.01$}} \\",
            "Optimized on both": r"\multicolumn{6}{l}"
            + r"{\textit{Optimized on both ($y=0.01$ and $y=0.015$)}} \\",
        }

        # Build rows grouped by section
        grouped = {k: [] for k in section_order}
        for _s in _optimize_sets:
            sec = _s.get("latex_section", "")
            if sec in grouped:
                grouped[sec].append(_s)
        print(grouped)

        def row_line(_s):
            _p = _s["current_params"]
            return (
                " & ".join(
                    [
                        _s.get("latex_row", _s.get("name", "")),
                        _get_phi(_p),
                        _get_phi_ref(_p),
                        _get_factor_max(_p),
                        _score(_s, "score_1"),
                        _score(_s, "score_2"),
                    ]
                )
                + r" \\"
            )

        lines = []
        lines.append(r"\begin{table}[h!]")
        lines.append(r"\centering")
        lines.append(
            r"\caption{Parameters and scores for different polymer network "
            + r"models. The original parameters show a higher score as the "
            + r"integration of the network resulted in a stiffer response. "
            + r"The score is defined as the average relative errors, where "
            + r"higher values indicate poorer performance.}"
        )
        lines.append(r"\label{tab:network_models}")
        lines.append(r"\begin{tabular}{lccccc}")
        lines.append(r"\hline")
        lines.append(
            r"Network model & $\Phi$ & $\phi_{\text{ref}}$ & Factor "
            + r"$V_\mathrm{max}$ & Score ($y=0.01$) & Score ($y=0.015$) \\"
        )
        lines.append(r"\hline")

        for sec in section_order:
            if not grouped[sec]:
                continue
            lines.append(section_titles[sec])
            for _s in grouped[sec]:
                lines.append(row_line(_s))
            lines.append(r"\hline")

        lines.append(r"\end{tabular}")
        lines.append(r"\end{table}")
        lines.append("")  # trailing newline

        out_path.write_text("\n".join(lines), encoding="utf-8")

    latex_path = out_dir / "00_scores_table.txt"
    write_latex_table(optimize_sets, latex_path)
