"""Fit PNIPAAm/water/acetic acid gel parameters to DOS data and save
the best tracing seed."""

# pylint: disable=duplicate-code

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

from equi_gel.save_load_json import (
    BASE_FOLDER_EXP,
    load_dos_record,
)
from gel_tracer_ternary import trace_binodal_gel

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 = 298


# 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.02_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.015_xigel0.10.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_II_lst = [[0.0, 1.0 - w, w] for w in dos_records[-1]["w_3_g_per_g"]]


class GelObjectiveWrapper(ObjectiveWrapperBase):
    """Objective Wrapper"""

    def objective(self, current_args: dict, verbosity: bool = False) -> float:
        try:
            score_gel = 0.0

            for dos_record in dos_records:
                _binodal_dict = trace_binodal_gel(
                    current_args,
                    t_high=T_HIGH,
                    t_low=T_LOW,
                    y_crosslinker=dos_record["crosslinker"],
                    verbosity=True,
                )

                score_gel += get_minimal_distance_swelling(
                    _binodal_dict, dos_record
                )

            return float(score_gel)

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


keys_to_optimize = ["comp13_kij_intercept", "comp23_kij_intercept"]

max_change_of_parameters = [2.0, 2.0]  # comp12_kij_intercept


add_params = {
    "number_components": 3,
    "gel": True,
    "elasticity": "Flory",
    "NETWORK_FUNCTIONALITY": 2.6938877538943906,
    "Y_CROSSLINKER": 0.01,
    "V_MAX_FAC": 0.06,
    "ref_volume_frac": 1.0,
    "C_FAC": 1.0,
}

# -----------------------------
# Build parameter sets
# -----------------------------
optimized_params_bahrabadi = {
    "comp23_kij_intercept": -0.7073522163072805,
    "comp23_kij_slope": 0.00,
    "comp13_kij_intercept": 0.00,
    "comp13_kij_slope": 0.00,
}


current_parameter_set = (
    get_parameter_set_from_tables(
        ["pnipaam", "water", "acetic_acid"],
        ["arndt", "arndt", "bahrabadi"],
        binary_sources={(1, 2): "arndt"},
    )
    | optimized_params_bahrabadi
    | add_params
)

if __name__ == "__main__":

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

    objective_wrapper = GelObjectiveWrapper(
        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": "pnipaam_ter_",
        },
        population_kwargs={
            "popsize": 8,
            "sigma": 0.05,
            "sigma_random": 0.02,
        },
    )

    # Objective file 1:
    print(objective_wrapper.objective(current_parameter_set))

    best_params = current_parameter_set | {
        k: v for k, v in zip(keys_to_optimize, res["best_unscaled"][0])
    }
    print(objective_wrapper.objective(best_params))

    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("pnipaam_best_params_W_AA.json")
    outpath.write_text(
        json.dumps(best_params, indent=2, default=_jsonable), encoding="utf-8"
    )
    print(f"Saved: {outpath.resolve()}")

    # Create y_vec as initial guess for consecutive calculation
    binodal_dict = trace_binodal_gel(
        best_params,
        t_high=T_HIGH,
        t_low=T_LOW,
        y_crosslinker=0.015,
        verbosity=True,
    )
    ln_x_I_lst = np.log(binodal_dict["x_I_lst"][-1][[0, -1]])
    ln_x_II_lst = np.log(binodal_dict["x_II_lst"][-1][[0, -1]])
    ln_K_lst = ln_x_I_lst - ln_x_II_lst
    y_lst = np.concatenate(
        [
            ln_x_I_lst,
            ln_x_II_lst,
            ln_K_lst,
            [np.log(binodal_dict["temperature_K"][-1])],
        ]
    )
    np.save("pnipaam_best_params_W_AA_last_y_lst.npy", np.array(y_lst))
