"""Fit PNIPAAm/acetic acid/butyl acetate gel parameters to DOS data
and save tracing seeds."""

# pylint: disable=duplicate-code

from __future__ import annotations
from pathlib import Path
import json
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.interface_feos import get_feos_npt_state

from equi_gel.parameters import (
    get_parameter_set_from_tables,
    get_subset_of_parameters,
    add_mixing_association,
)

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


# DOS files
DOS_FILES = [
    BASE_FOLDER_EXP
    / ("EXP_DOS_pnipaam_aceticacid" "_butylacetate_eckert_0.01.json"),
]

keys_to_optimize_lst = [
    # [
    #     "comp13_kij_intercept",
    # ],
    [
        "NETWORK_FUNCTIONALITY",
    ],
    # [
    #     "comp13_kij_intercept",
    #     "NETWORK_FUNCTIONALITY",
    # ],
    # [
    #     "comp13_kij_intercept",
    #     "NETWORK_FUNCTIONALITY",
    #     "comp23_kij_intercept",
    # ],
]

for keys_to_optimize in keys_to_optimize_lst:

    max_change_of_parameters = [2.0] * len(keys_to_optimize)

    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.7,
        "comp23_kij_slope": 0.00,
        "comp13_kij_intercept": 0.00,
        "comp13_kij_slope": 0.00,
        "comp12_kij_intercept": -0.06744727108093351,
        "comp12_kij_slope": 0.00,
    }

    current_parameter_set = (
        get_parameter_set_from_tables(
            ["pnipaam", "acetic_acid", "butyl_acetate"],
            ["arndt", "bahrabadi", "bahrabadi"],
            binary_sources={},
        )
        | optimized_params_bahrabadi
        | add_params
    )

    initial_guess_y_vec = np.load(
        "optimize_parameter_gel_other_solvents/"
        + "pnipaam_best_params_W_AA_last_y_lst.npy"
    )

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

            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_HIGH,
                        PRESSURE,
                        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"]
            ]

    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:  # pylint: disable=cell-var-from-loop
                    # pylint: disable=cell-var-from-loop
                    _binodal_dict = trace_binodal_gel(
                        current_args,
                        t_high=T_HIGH,
                        t_low=T_LOW,
                        y_crosslinker=dos_record["crosslinker"],
                        verbosity=True,
                        y_init=initial_guess_y_vec,
                    )
                    # pylint: enable=cell-var-from-loop

                    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")

    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

        keys_tag = "_".join(keys_to_optimize)

        outpath = Path(f"pnipaam_best_params_AA_BA_{keys_tag}.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=dos_records[0]["crosslinker"],
            verbosity=True,
            y_init=initial_guess_y_vec,
        )
        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(
            f"pnipaam_best_params_AA_BA_last_y_lst_{keys_tag}.npy",
            np.array(y_lst),
        )
