#################################################################################
# WaterTAP Copyright (c) 2020-2023, The Regents of the University of California,
# through Lawrence Berkeley National Laboratory, Oak Ridge National Laboratory,
# National Renewable Energy Laboratory, and National Energy Technology
# Laboratory (subject to receipt of any required approvals from the U.S. Dept.
# of Energy). All rights reserved.
#
# Please see the files COPYRIGHT.md and LICENSE.md for full copyright and license
# information, respectively. These files are also available online at the URL
# "https://github.com/watertap-org/watertap/"
#################################################################################
import numpy as np
from pyomo.common.config import ConfigValue

from watertap.tools.parameter_sweep.sampling_types import NormalSample
from watertap.tools.parameter_sweep.parameter_sweep import (
    _ParameterSweepBase,
    ParameterSweep,
    return_none,
)
from watertap.tools.parallel.single_process_parallel_manager import (
    SingleProcessParallelManager,
)


from pyomo.common.deprecation import deprecation_warning


class DifferentialParameterSweep(_ParameterSweepBase):
    CONFIG = _ParameterSweepBase.CONFIG()

    CONFIG.declare(
        "num_diff_samples",
        ConfigValue(
            default=1,
            domain=int,
            description="Number of differntial sweep samples",
        ),
    )

    CONFIG.declare(
        "guarantee_solves",
        ConfigValue(
            default=False,
            domain=bool,
            description="Guarantee a pre-specified number of solves.",
        ),
    )

    CONFIG.declare(
        "differential_sweep_specs",
        ConfigValue(
            default=dict(),
            domain=dict,
            description="Dictionary containing the specifications for the differential sweep",
            doc="""
            A specification dictionary that contains details for how to construct the parameter sweep dictionary for differential sweep.
            This is a nested dictionary where the first level denotes the variable names for which the differential sweep needs to be carried out.
            The second level denotes various options to be used for wach variable.
            The number of samples for each differential sweep is specified while initializing the DifferentialParameterSweep object wsing the keyword `num_diff_samples`
            e.g.
            
            {
                "fs.a": {
                    "diff_mode": "sum",
                    "diff_sample_type": NormalSample,
                    "std_dev": 0.01,
                    "pyomo_object": m.fs.input["a"],
                },
                "fs.b": {
                    "diff_mode": "product",
                    "diff_sample_type": UniformSample,
                    "relative_lb": 0.01,
                    "relative_ub": 0.01,
                    "pyomo_object": m.fs.input["b"],
                },
                "fs.c": {
                    "diff_mode": "sum",
                    "diff_sample_type": GeomSample,
                    "relative_lb": 0.01,
                    "relative_ub": 10.0,
                    "pyomo_object": m.fs.input["c"],
                },
            }

            """,
        ),
    )

    def __init__(
        self,
        **options,
    ):
        # Initialize the base Class
        super().__init__(**options)

        if self.config.guarantee_solves:
            raise NotImplementedError

    def _create_differential_sweep_params(self, local_values):
        differential_sweep_specs = self.config.differential_sweep_specs

        diff_sweep_param = {}
        for ctr, (param, specs) in enumerate(differential_sweep_specs.items()):
            nominal_val = local_values[self.diff_spec_index[ctr]]
            pyomo_object = specs["pyomo_object"]
            if specs["diff_sample_type"] == NormalSample:
                std_dev = specs["std_dev"]
                diff_sweep_param[param] = NormalSample(
                    pyomo_object, nominal_val, std_dev, self.config.num_diff_samples
                )
            else:
                relative_lb = specs["relative_lb"]
                relative_ub = specs["relative_ub"]
                if specs["diff_mode"] == "sum":
                    lb = nominal_val * (1 - relative_lb)
                    ub = nominal_val * (1 + relative_ub)
                elif specs["diff_mode"] == "product":
                    lb = nominal_val * relative_lb
                    ub = nominal_val * relative_ub
                elif specs["diff_mode"] == "percentile":
                    lower_nominal = specs["nominal_lb"]
                    upper_nominal = specs["nominal_ub"]
                    delta_nominal = abs(upper_nominal - lower_nominal)
                    lb = nominal_val + delta_nominal * relative_lb
                    ub = nominal_val + delta_nominal * relative_ub
                else:
                    raise NotImplementedError
                diff_sweep_param[param] = specs["diff_sample_type"](
                    pyomo_object, lb, ub, self.config.num_diff_samples
                )

        return diff_sweep_param

    def _check_differential_sweep_key_validity(self, sweep_params):
        diff_specs_keys = list(self.config.differential_sweep_specs.keys())
        sweep_param_keys = list(sweep_params.keys())

        if all(key in sweep_param_keys for key in diff_specs_keys):
            self.diff_spec_index = [
                sweep_param_keys.index(key) for key in diff_specs_keys
            ]
        else:
            raise ValueError(
                "differential_sweep_specs keys don't match with sweep_param keys"
            )

    def _define_differential_sweep_outputs(self, sweep_params):
        self.differential_outputs = self.outputs
        if self.outputs is not None:
            for key in sweep_params.keys():
                if key not in self.config.differential_sweep_specs.keys():
                    self.differential_outputs[key] = sweep_params[key].pyomo_object

    def _create_local_output_skeleton(self, model, sweep_params, outputs, num_samples):
        output_dict = super()._create_local_output_skeleton(
            model, sweep_params, outputs, num_samples
        )
        output_dict["nominal_idx"] = np.arange(
            num_samples, dtype=float
        )  # [*range(num_samples)]
        output_dict["differential_idx"] = np.array([np.nan] * num_samples)
        return output_dict

    def _append_differential_results(self, local_output_dict, diff_results_dict):
        for idx, diff_sol in diff_results_dict.items():
            for key, item in diff_sol.items():
                # Solve status
                if key == "solve_successful":
                    n_diff_samples = len(item)
                    local_output_dict["solve_successful"].extend(item)
                    local_output_dict["nominal_idx"] = np.concatenate(
                        (
                            local_output_dict["nominal_idx"],
                            np.array([np.nan] * n_diff_samples),
                        ),
                        axis=0,
                    )
                    local_output_dict["differential_idx"] = np.concatenate(
                        (
                            local_output_dict["differential_idx"],
                            np.array([idx] * n_diff_samples, dtype=float),
                        ),
                        axis=0,
                    )

                else:
                    for subkey, subitem in item.items():
                        local_output_dict[key][subkey]["value"] = np.concatenate(
                            (
                                local_output_dict[key][subkey]["value"],
                                subitem["value"],
                            )
                        )

                    # We also need to capture sweep_params variables that are not a part of differential_sweep_specs
                    if key == "sweep_params":
                        missing_sub_keys = set(item.keys()) ^ set(
                            local_output_dict[key].keys()
                        )
                        # This loop shouldn't run if the above set is empty
                        for subkey in missing_sub_keys:
                            # We are picking the unchanged sweep_params from the outputs. In the ideal world, they would be the same.
                            local_output_dict["sweep_params"][subkey][
                                "value"
                            ] = np.concatenate(
                                (
                                    local_output_dict["sweep_params"][subkey]["value"],
                                    diff_sol["outputs"][subkey]["value"],
                                )
                            )

    def _collect_local_inputs(self, local_results_dict):
        num_local_samples = len(local_results_dict["solve_successful"])
        local_inputs = np.zeros(
            (num_local_samples, len(local_results_dict["sweep_params"])),
            dtype=float,
        )

        for i, (key, item) in enumerate(local_results_dict["sweep_params"].items()):
            local_inputs[:, i] = item["value"]

        return local_inputs

    def _aggregate_input_arr(self, global_results_dict, num_global_samples):
        global_values = np.zeros(
            (num_global_samples, len(global_results_dict["sweep_params"])),
            dtype=float,
        )

        if self.parallel_manager.is_root_process():
            for i, (key, item) in enumerate(
                global_results_dict["sweep_params"].items()
            ):
                global_values[:, i] = item["value"]

        self.parallel_manager.sync_array_with_peers(global_values)

        return global_values

    def _aggregate_results(self, local_output_dict):
        # Create the global results dictionary
        global_results_dict = self._create_global_output(local_output_dict)

        # Broadcast the number of global samples to all ranks
        num_global_samples = len(global_results_dict["solve_successful"])
        num_global_samples = self.parallel_manager.sync_pyobject_with_peers(
            num_global_samples
        )

        global_results_arr = self._aggregate_results_arr(
            global_results_dict, num_global_samples
        )
        global_input_values = self._aggregate_input_arr(
            global_results_dict, num_global_samples
        )

        return (
            global_results_dict,
            global_results_arr,
            global_input_values,
            num_global_samples,
        )

    def _create_global_output(self, local_output_dict):  # , req_num_samples=None):
        global_output_dict = super()._create_global_output(local_output_dict)

        # We now need to get the mapping array. This only needs to happen on root
        local_num_cases_all = len(local_output_dict["solve_successful"])
        # AllGather the total size of the value array on each MPI rank
        sample_split_arr = self.parallel_manager.combine_data_with_peers(
            local_num_cases_all
        )
        num_total_samples = sum(sample_split_arr)

        # AllGather nominal values for creating the parallel offset
        nominal_sample_split_arr = self.parallel_manager.combine_data_with_peers(
            self.n_nominal_local
        )

        # We need to create a global index and offset items accordingly. This
        # needs to happen on all ranks/workers.
        my_rank = self.parallel_manager.get_rank()
        offset = 0
        if my_rank > 0:
            offset = sum(nominal_sample_split_arr[:my_rank])
        local_output_dict["nominal_idx"] = local_output_dict["nominal_idx"] + offset
        local_output_dict["differential_idx"] = (
            local_output_dict["differential_idx"] + offset
        )

        # Resize global index array
        if self.parallel_manager.is_root_process():
            global_output_dict["nominal_idx"] = np.zeros(num_total_samples, dtype=float)
            global_output_dict["differential_idx"] = np.zeros(
                num_total_samples, dtype=float
            )

        # Now we need to collect it on global_output_dict
        self.parallel_manager.gather_arrays_to_root(
            sendbuf=local_output_dict["nominal_idx"],
            recvbuf_spec=(
                global_output_dict["nominal_idx"],
                sample_split_arr,
            ),
        )
        self.parallel_manager.gather_arrays_to_root(
            sendbuf=local_output_dict["differential_idx"],
            recvbuf_spec=(
                global_output_dict["differential_idx"],
                sample_split_arr,
            ),
        )

        return global_output_dict

    def _run_differential_sweep(self, local_value):
        diff_sweep_param_dict = self._create_differential_sweep_params(local_value)

        # We want this instance of the parameter sweep to run in serial
        diff_ps = ParameterSweep(
            optimize_function=self.config.optimize_function,
            optimize_kwargs=self.config.optimize_kwargs,
            reinitialize_function=self.config.reinitialize_function,
            reinitialize_kwargs=self.config.reinitialize_kwargs,
            reinitialize_before_sweep=self.config.reinitialize_before_sweep,
            parallel_manager_class=SingleProcessParallelManager,
        )
        # pass model_manager from refernce sweep, to diff sweep
        # so we don't have to reijnit he model
        diff_ps.model_manager = self.model_manager

        _, differential_sweep_output_dict = diff_ps.parameter_sweep(
            diff_ps.model_manager.model,
            diff_sweep_param_dict,
            build_outputs=self.differential_outputs,
            num_samples=self.config.num_diff_samples,
            seed=self.seed,
        )

        return differential_sweep_output_dict

    def _run_sample(
        self,
        local_value_k,
        k,
        sweep_params,
        local_output_dict,
    ):
        run_successful = super()._run_sample(
            local_value_k,
            k,
            sweep_params,
            local_output_dict,
        )
        self.differential_sweep_output_dict[k] = self._run_differential_sweep(
            local_value_k
        )

        return run_successful

    def _do_param_sweep(self, sweep_params, outputs, local_values):
        self.differential_sweep_output_dict = {}

        local_output_dict = super()._do_param_sweep(sweep_params, outputs, local_values)

        # Now append the outputs of the differential solves
        self._append_differential_results(
            local_output_dict, self.differential_sweep_output_dict
        )

        return local_output_dict

    def parameter_sweep(
        self,
        build_model,
        build_sweep_params,
        build_outputs=None,
        build_outputs_kwargs=None,
        num_samples=None,
        seed=None,
        build_model_kwargs=None,
        build_sweep_params_kwargs=None,
    ):
        # Create a base sweep_params
        build_model_kwargs = (
            build_model_kwargs if build_model_kwargs is not None else dict()
        )
        build_outputs_kwargs = (
            build_outputs_kwargs if build_outputs_kwargs is not None else dict()
        )
        build_sweep_params_kwargs = (
            build_sweep_params_kwargs
            if build_sweep_params_kwargs is not None
            else dict()
        )

        if not callable(build_model):
            _model = build_model
            build_model = lambda: _model
            deprecation_warning(
                "Passing a model directly to the parameter_sweep function is deprecated \
                                and will not work with future implementations of parallelism.",
                version="0.10.0",
            )

        if not callable(build_sweep_params):
            _sweep_params = build_sweep_params
            build_sweep_params = lambda model: _sweep_params
            deprecation_warning(
                "Passing sweep params directly to the parameter_sweep function is deprecated \
                                and will not work with future implementations of parallelism.",
                version="0.10.0",
            )

        if build_outputs is None:
            build_outputs = return_none

        if not callable(build_outputs):
            _combined_outputs = build_outputs
            build_outputs = lambda model: _combined_outputs
            deprecation_warning(
                "Passing the output dict directly to the parameter_sweep function is deprecated \
                                and will not work with future implementations of parallelism.",
                version="0.10.0",
            )
        # This should be depreciated in future versions
        self.config.build_model = build_model
        self.config.build_sweep_params = build_sweep_params
        self.config.build_outputs = build_outputs
        self.config.build_outputs_kwargs = build_outputs_kwargs
        self.config.build_model_kwargs = build_model_kwargs
        self.config.build_sweep_params_kwargs = build_sweep_params_kwargs

        model = build_model(**build_model_kwargs)
        sweep_params = build_sweep_params(model, **build_sweep_params_kwargs)
        sweep_params, sampling_type = self._process_sweep_params(sweep_params)

        # Check if the keys in the differential sweep specs exist in sweep params
        self._check_differential_sweep_key_validity(sweep_params)

        # Define differential sweep outputs
        self.outputs = self.config.build_outputs(
            model, *self.config.build_outputs_kwargs
        )
        self._define_differential_sweep_outputs(sweep_params)

        # Set the seed before sampling
        self.seed = seed
        np.random.seed(self.seed)

        # Enumerate/Sample the parameter space
        global_values = self._build_combinations(
            sweep_params, sampling_type, num_samples
        )

        # divide the workload between processors
        local_values = self._divide_combinations(global_values)
        self.n_nominal_local = np.shape(local_values)[0]

        # Check if the outputs have the name attribute. If not, assign one.
        if self.outputs is not None:
            self.assign_variable_names(model, self.outputs)

        # Create a dictionary to store all the differential ps_objects
        self.diff_ps_dict = {}

        # Do the Loop
        if self.config.custom_do_param_sweep is None:
            local_results_dict = self._do_param_sweep(
                sweep_params,
                self.outputs,
                local_values,
            )
        else:
            local_results_dict = self.config.custom_do_param_sweep(
                self,
                sweep_params,
                self.outputs,
                local_values,
                **self.config.custom_do_param_sweep_kwargs,
            )

        # re-writing local_values
        local_values = self._collect_local_inputs(local_results_dict)

        # Aggregate results on Master
        (
            global_results_dict,
            global_results_arr,
            global_input_arr,
            num_global_samples,
        ) = self._aggregate_results(local_results_dict)

        # Save to file
        global_save_data = self.writer.save_results(
            sweep_params,
            local_values,
            global_input_arr,
            local_results_dict,
            global_results_dict,
            global_results_arr,
            self.parallel_manager.get_rank(),
        )

        return global_results_dict, global_save_data
