"""Fit esterification-system interaction parameters to binary and
ternary LLE data."""

# pylint: disable=duplicate-code

from __future__ import annotations

import numpy as np

from equi_gel.optimize_parameters import (
    ObjectiveWrapperBase,
    ParameterOptimization,
)
from equi_gel.cost_func_binodal import get_minimal_distance
from equi_gel.newton_raphson_tracing import (
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
    add_component_to_y_vec,
)
from equi_gel.parameters import (
    get_parameter_set_from_tables,
    add_mixing_association,
    get_subset_of_parameters,
)
from equi_gel.save_load_json import (
    load_record_groups,
    find_lle_paths_by_components,
)

COMPONENTS_ALL = ["water", "n_butanol", "butyl_acetate", "acetic_acid"]


PRESSURE = 1.13
T_INIT_BINARY_LLE = 270.0
T_TERNARY_LLE = 298.15
T_FINAL = 400.0


SYSTEMS = [
    # {
    #     "name": "water_n_butanol",
    #     "components": ["water", "n_butanol"],
    #     "subset_idx": [1, 2],
    #     "kind": "binary",
    #     "reverse_order": True,
    #     "step_size_init": 0.01,
    #     "critical_point_tol": 5e-2,
    #     "max_step_fraction": 1,
    #     "which_idx_binary": 1,
    # },
    # {
    #     "name": "water_butyl_acetate",
    #     "components": ["water", "butyl_acetate"],
    #     "subset_idx": [1, 3],
    #     "kind": "binary",
    #     "reverse_order": False,
    #     "step_size_init": 0.01,
    #     "critical_point_tol": 5e-2,
    #     "max_step_fraction": 1,
    #     "stuck_progress_tol": 0.0005,
    #     "which_idx_binary": 1,
    # },
    {
        "name": "water_n_butanol_acetic_acid",
        "components": ["water", "n_butanol", "acetic_acid"],
        "subset_idx": [1, 2, 4],
        "binary_subset_idx": [1, 2],
        "kind": "ternary",
        "reverse_order": True,
        "trace_index": 2,
        "step_size_init": 0.01,
        "critical_point_tol": 5e-2,
    },
    # {
    #     "name": "water_n_butanol_butyl_acetate",
    #     "components": ["water", "n_butanol", "butyl_acetate"],
    #     "subset_idx": [1, 2, 3],
    #     "binary_subset_idx": [1, 2],
    #     "kind": "ternary",
    #     "reverse_order": True,
    #     "trace_index": 5,
    #     "step_size_init": 0.01,
    #     "critical_point_tol": 5e-2,
    # },
    {
        "name": "water_butyl_acetate_acetic_acid",
        "components": ["water", "butyl_acetate", "acetic_acid"],
        "subset_idx": [1, 3, 4],
        "binary_subset_idx": [1, 3],
        "kind": "ternary",
        "reverse_order": False,
        "trace_index": 5,
        "step_size_init": 0.0001,
        "critical_point_tol": 5e-2,
        "stuck_progress_tol": 0.0005,
    },
]


class EsterificationObjectiveWrapper(ObjectiveWrapperBase):
    """Objective wrapper for esterification parameter optimization."""

    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
    # pylint: disable=broad-exception-caught
    def __init__(
        self,
        _base_parameters: dict,
        _keys_to_optimize: list[str],
        _exp_groups: dict,
    ) -> None:
        super().__init__(_base_parameters, _keys_to_optimize)
        self.exp_groups = _exp_groups

    def objective(self, current_args: dict, verbosity: bool = False) -> float:
        try:
            total_score = 0.0
            system_scores = []

            for system in SYSTEMS:
                exp_dicts, permutation_list = self.exp_groups[system["name"]]

                if system["kind"] == "binary":

                    permutation_list[0].append([i for i in range(2)])

                    params_binary = get_subset_of_parameters(
                        current_args, system["subset_idx"]
                    )
                    y_init = build_initial_guess_binary(
                        T_INIT_BINARY_LLE,
                        PRESSURE,
                        params_binary,
                        reverse_order=system["reverse_order"],
                    )
                    trace_kwargs = {}
                    if "stuck_progress_tol" in system:
                        trace_kwargs["stuck_progress_tol"] = system[
                            "stuck_progress_tol"
                        ]
                    y_lst, _, _, _, _, _ = trace_parameter(
                        trace_index=-1,
                        y_init=y_init,
                        trace_start=T_INIT_BINARY_LLE,
                        trace_target=T_FINAL,
                        pressure=PRESSURE,
                        params=params_binary,
                        step_size_init=system["step_size_init"],
                        newton_tol=1e-5,
                        critical_point_tol=system["critical_point_tol"],
                        max_step_fraction=system["max_step_fraction"],
                        **trace_kwargs,
                    )
                    binodal_dict = interpret_y_vector(
                        y_lst, PRESSURE, params_binary
                    )
                    score = get_minimal_distance(
                        binodal_dict,
                        exp_dicts,
                        permutation_list,
                        binodal_sampling=8,
                        which_idx_binary=system["which_idx_binary"],
                    )[0]

                else:
                    params_binary = get_subset_of_parameters(
                        current_args, system["binary_subset_idx"]
                    )
                    y_init_binary = build_initial_guess_binary(
                        T_TERNARY_LLE,
                        PRESSURE,
                        params_binary,
                        reverse_order=system["reverse_order"],
                    )
                    y_init_ternary = add_component_to_y_vec(
                        y_init_binary, 1e-5, 1e-5
                    )
                    params_ternary = get_subset_of_parameters(
                        current_args, system["subset_idx"]
                    )
                    trace_kwargs = {}
                    if "stuck_progress_tol" in system:
                        trace_kwargs["stuck_progress_tol"] = system[
                            "stuck_progress_tol"
                        ]
                    y_lst, _, _, _, _, _ = trace_parameter(
                        trace_index=system["trace_index"],
                        y_init=y_init_ternary,
                        trace_start=1e-5,
                        trace_target=1.0,
                        pressure=PRESSURE,
                        params=params_ternary,
                        step_size_init=system["step_size_init"],
                        newton_tol=1e-5,
                        fixed_index=[-1],
                        critical_point_tol=system["critical_point_tol"],
                        **trace_kwargs,
                    )
                    binodal_dict = interpret_y_vector(
                        y_lst, PRESSURE, params_ternary
                    )

                    permutation_list[0].append([i for i in range(3)])

                    score = get_minimal_distance(
                        binodal_dict,
                        exp_dicts,
                        permutation_list,
                        binodal_sampling=8,
                    )[0]

                system_scores.append(score)

                total_score += score

            if verbosity:
                print(f"Total score: {total_score:.6f}")

            return total_score

        except Exception as exc:
            print("Error:", exc)
            return float("inf")


SOURCES = ["bahrabadi", "grob", "esper"]

keys_to_optimize = [
    # "comp12_epsilon_k_ab",
    # "comp13_epsilon_k_ab",
    # "comp23_epsilon_k_ab",
    # "comp14_epsilon_k_ab",
    "comp14_kij_intercept",
    "comp24_kij_intercept",
    "comp34_kij_intercept",
]

max_change_of_parameters = [
    # 2000.0,
    # 2000.0,
    # 2000.0,
    # 2000.0,
    2.0,
    2.0,
    2.0,
]


keys_to_optimize_lst = [keys_to_optimize]
max_change_of_parameters_lst = [max_change_of_parameters]
label_key_set_lst = ["only_kij_with_AA"]


if __name__ == "__main__":

    exp_path_lst_lst = [
        find_lle_paths_by_components(system["components"], is_exp=True)
        for system in SYSTEMS
    ]

    exp_groups = {
        system["name"]: load_record_groups(
            path_exp_lst=exp_paths,
            component_list=system["components"],
        )[1:]
        for system, exp_paths in zip(SYSTEMS, exp_path_lst_lst)
    }

    for source_idx, source in enumerate(SOURCES):

        for set_idx, key_set in enumerate(label_key_set_lst):

            prefix_file = "optimizer_histories/" + source + "_" + key_set + "_"

            base_parameters = get_parameter_set_from_tables(
                COMPONENTS_ALL, source
            )
            base_parameters = add_mixing_association(base_parameters)

            current_keys_to_optimize = keys_to_optimize_lst[set_idx]
            current_max_change = max_change_of_parameters_lst[set_idx]

            initial_guess = np.array(
                [base_parameters[key] for key in current_keys_to_optimize]
            )

            objective_wrapper = EsterificationObjectiveWrapper(
                base_parameters,
                current_keys_to_optimize,
                exp_groups,
            )
            optimizer = ParameterOptimization(objective_wrapper)

            optimizer.run_optimization(
                initial_guess=initial_guess,
                max_change_of_parameters=current_max_change,
                verbosity=True,
                de_kwargs={
                    "prefix_file": prefix_file,
                },
                population_kwargs={
                    "popsize": 25,
                    "sigma": 0.05,
                    "sigma_random": 0.02,
                },
            )
