"""
The optimizer module provides wrappers which connect optimization packages to Evaluators.
Currently we support `platypus <https://github.com/Project-Platypus/Platypus>`_
and `rbf_opt <https://github.com/coin-or/rbfopt>`_

The supported Platypus algorithms are:

- GeneticAlgorithm
- EvolutionaryStrategy
- NSGAII
- EpsMOEA
- GDE3
- SPEA2
- MOEAD
- NSGAIII
- ParticleSwarm
- OMOPSO
- SMPSO
- CMAES
- IBEA
- PAES
- PESA2
- EpsNSGAII

These algorithms have the same configuration options as their counterparts in platypus,
with the parts provided by besos evaluators filled automatically.
"""

# Python Core Libraries
import os
import sys
from collections import OrderedDict
from functools import partial
from inspect import signature, Parameter, Signature
from typing import Type, List

# External Libraries
import dask_utils
import numpy as np
import pandas as pd
import platypus
import rbfopt

# BESOS Imports
from config import optimizer as conf
from evaluator import AbstractEvaluator
from problem import Problem


# TODO: evaluations or end criterion?
def platypus_alg(
    evaluator: AbstractEvaluator,
    algorithm: Type[platypus.Algorithm],
    evaluations: int = conf["evaluations"],
    *args,
    **kwargs,
) -> pd.DataFrame:
    """Uses a platypus algorithm to optimise over an evaluator
        If the evaluator is configured to run in parrallel (via 'multi' flag),
        the platypus algorithm will also run in parrallel.

    :param evaluator: An evaluation function to optimise over.
    :param algorithm: The platypus algorithm to use.
    :param evaluations: The algorithm will be stopped once it uses more than this many evaluations.
    :param args: arguments to pass to `algorithm`
    :param kwargs: keyword arguments to pass to `algorithm`.
    :return: the non-dominated solutions found by the algorithm.
    """

    if evaluator.multi:
        parrallel_platypus_evaluator = platypus.MapEvaluator(dask_utils.dask_map)
        kwargs.update({"evaluator": parrallel_platypus_evaluator})

    problem = evaluator.to_platypus()
    evaluator.problem.pre_optimisation()

    alg: platypus.AbstractGeneticAlgorithm = algorithm(problem, *args, **kwargs)
    alg.run(evaluations)
    # TODO: Have a smarter default for parts (instead of 'all')
    return solutions_to_df(platypus.unique(alg.result), evaluator.problem)


_ordered_argtypes = [
    (Parameter.POSITIONAL_OR_KEYWORD, False),
    (Parameter.POSITIONAL_OR_KEYWORD, True),
    (Parameter.VAR_POSITIONAL, False),
    (Parameter.KEYWORD_ONLY, False),
    (Parameter.KEYWORD_ONLY, True),
    (Parameter.VAR_KEYWORD, False),
]


def get_operator(problem: platypus.Problem, mutation=False):
    """Creates a valid operator for the given platypus Problem.
    This is provided to help when optimizing over problems with mixed types.

    :param problem: the platypus problem this operator should apply to
    :param mutation: if True create a mutation operator, if False create a variation operator
    :return:
    """
    operators = []
    if mutation:
        defaults = platypus.config.PlatypusConfig.default_mutator
        class_ = platypus.CompoundMutation
    else:
        defaults = platypus.config.PlatypusConfig.default_variator
        class_ = platypus.CompoundOperator
    for t in problem.types:
        if t in defaults:
            operators.append(defaults[t])
        else:
            for super_t, operator in defaults.items():
                if isinstance(t, super_t):
                    operators.append(operator)
    if mutation:
        return class_(*operators)
    return class_(*operators)


def _alg_t(algorithm):
    """A function that wraps a platypus algorithm in a common workflow.
    This function propagates the algorithm's function signature out to produce a more descriptive
    function signature."""

    # This is perhaps an overly fancy way of doing it, but it avoids having to
    # do the same process manually for all platypus algorithms, and will
    # update the signatures automatically if they change in platypus.

    defaults = {
        "variator": get_operator,
        "mutator": partial(get_operator, mutation=True),
    }

    def template(evaluator, evaluations=conf["evaluations"], *args, **kwargs):
        for kwarg, default in defaults.items():
            if kwarg not in kwargs or kwargs[kwarg] == "automatic":
                kwargs[kwarg] = default(evaluator.to_platypus())
        return platypus_alg(evaluator, algorithm, evaluations, *args, **kwargs)

    temp_name = (
        # (function, arguments_to_omit)
        (template, ["args", "kwargs"]),
        (algorithm.__init__, ["self", "problem"]),
    )

    argdict = {kind: [] for kind in _ordered_argtypes}
    for i, (func, omit) in enumerate(temp_name):
        sig = signature(func)
        p: Parameter
        for name, p in sig.parameters.items():
            if name in omit:
                continue
            if i > 0:
                # convert non-default keyword arguments to keyword only arguments
                if (
                    p.kind is Parameter.POSITIONAL_OR_KEYWORD
                    and p.default is Parameter.empty
                ):
                    p = p.replace(kind=Parameter.KEYWORD_ONLY)
                # use the defaults specified above
                if name in defaults:
                    p = p.replace(default="automatic")

            argdict[(p.kind, p.default is not Parameter.empty)].append(p)
    arglist = []
    for kind, default in _ordered_argtypes:
        arglist.extend(argdict[kind, default])

    template.__signature__ = Signature(
        parameters=arglist, return_annotation=signature(template).return_annotation
    )
    template.__name__ = algorithm.__name__
    return template


# The following functions are convenience wrappers for platypus' implementations
# of the algorithms with the same name
# they return the non-dominated solutions found by the algorithm

GeneticAlgorithm = _alg_t(platypus.GeneticAlgorithm)
EvolutionaryStrategy = _alg_t(platypus.EvolutionaryStrategy)
NSGAII = _alg_t(platypus.NSGAII)
EpsMOEA = _alg_t(platypus.EpsMOEA)
GDE3 = _alg_t(platypus.GDE3)
SPEA2 = _alg_t(platypus.SPEA2)
MOEAD = _alg_t(platypus.MOEAD)
NSGAIII = _alg_t(platypus.NSGAIII)
ParticleSwarm = _alg_t(platypus.ParticleSwarm)
OMOPSO = _alg_t(platypus.OMOPSO)
SMPSO = _alg_t(platypus.SMPSO)
CMAES = _alg_t(platypus.CMAES)
IBEA = _alg_t(platypus.IBEA)
PAES = _alg_t(platypus.PAES)
PESA2 = _alg_t(platypus.PESA2)
EpsNSGAII = _alg_t(platypus.EpsNSGAII)


def values_to_solution(
    values: list, problem: platypus.Problem, evaluate=True
) -> platypus.Solution:
    """Produces an platypus solution with the given inputs

    :param values: a list of input values to convert to a platypus solution
    :param problem: the platypus problem this solution applies to
    :param evaluate: whether to evaluate the solution before returning (solution will be unevaluated otherwise)
    :return: the corresponding platypus.Solution object
    """
    solution = platypus.Solution(problem)
    if len(values) != problem.nvars:
        raise ValueError(
            f"Length of values is {len(values)} but expected length of {problem.nvars}"
        )
    solution.variables = list(values)
    if evaluate:
        solution.evaluate()
    return solution


# TODO: Unify `parts` handling between this and Problem to remove duplicate code if possible


def solution_to_values(solution: platypus.Solution, parts="all") -> list:
    """Converts a platypus solution to a list containing the same values

    :param solution: a platypus solution to convert
    :param parts: which parts of the solution should be kept
    :return: a list of the requested values from the solution
    """
    part_options = OrderedDict(
        inputs="variables",
        outputs="objectives",
        constraints="constraints",
        violation="violation",
    )
    if parts == "all":
        parts = part_options.keys()
    assert set(parts) <= set(
        part_options
    ), f'parts must be a subset of {set(part_options.keys())} or "all"'
    result = []
    for part in parts:
        if part == "violation":
            result.append(solution.constraint_violation)
        elif part == "inputs":
            encoded = getattr(solution, part_options[part])
            decoded = [
                plat_type.decode(val)
                for val, plat_type in zip(encoded, solution.problem.types)
            ]
            result.extend(decoded)
        else:
            result.extend(getattr(solution, part_options[part]))
    return result


def solutions_to_df(
    solutions: List[platypus.Solution], problem, parts="all", flag_optimal=True
) -> pd.DataFrame:
    """Converts a list of platypus solutions to a DataFrame, with one row corresponding to each solution

    :param solutions: list of solutions to convert
    :param problem: the column names for DataFrame
    :param parts: which parts of the solutions should be kept
    :param flag_optimal: whether to include a boolean column denoting whether each solution is pareto-optimal
    :return: a DataFrame
    """

    def to_col_vals(solution_list):
        return list(
            zip(*(solution_to_values(solution, parts) for solution in solution_list))
        )

    solutions = platypus.unique(solutions)
    non_dominated = platypus.nondominated(solutions)
    columns = problem.names(parts)
    values, non_dom_vals = to_col_vals(solutions), to_col_vals(non_dominated)
    assert len(columns) == len(
        values
    ), f"{len(values)} values does not match {len(columns)} columns"
    # TODO: Intuit the dataframe column types based on the types of the parameters of the problem
    # or use the to_df method of the problem object
    solution_df = pd.DataFrame(
        {column: data for column, data in zip(columns, values)}
    )  # , dtype=float
    if flag_optimal:
        non_dom_df = pd.DataFrame(
            {column: data for column, data in zip(columns, non_dom_vals)}
        )  # , dtype=float
        df = pd.merge(solution_df, non_dom_df, how="outer", indicator="pareto-optimal")
        df["pareto-optimal"] = df["pareto-optimal"] == "both"
        if problem.depend_params or problem.add_outputs:
            df = problem.overwrite_df(df)
            problem.post_optimazation()
        return df
    if problem.depend_params or problem.add_outputs:
        solution_df = problem.overwrite_df(solution_df)
        problem.post_optimazation()
    return solution_df


def df_solution_to_solutions(
    df: pd.DataFrame,
    platypus_problem: platypus.Problem,
    besos_problem: Problem,
) -> List[platypus.Solution]:
    """Converts a solution DataFrame to a list of platypus solutions, with each row converted to one solution

    :param df: DataFrame to convert
    :param platypus_problem: platypus problem that the solutions apply to
    :param besos_problem: besos problem that the solutions apply to
    :return: a list of platypus solutions
    """
    input_df = df.copy(deep=True)
    labels = df.columns.values.tolist()
    for input in besos_problem.inputs:
        labels.remove(input.name)
    for label in labels:
        input_df.drop(columns=label, inplace=True)
    return list(input_df.apply(values_to_solution, problem=platypus_problem, axis=1))


def df_to_solutions(
    df: pd.DataFrame, problem: platypus.Problem
) -> List[platypus.Solution]:
    """Converts a DataFrame to a list of platypus solutions, with each row converted to one solution

    :param df: DataFrame to convert
    :param problem: platypus problem that the solutions apply to
    :return: a list of platypus solutions
    """
    return list(df.apply(values_to_solution, problem=problem, axis=1))


# rbf opt

# TODO: Add support for the other termination criterion
def rbf_opt(
    evaluator, evaluations, hide_output: bool = True, bonmin_path=None, rand_seed=None
):
    """
    This is wrapper to the rbfopt function.

    :param evaluator: An evaluation function to optimise over.
    :param evaluations: The algorithm will be stopped once it uses more than this many evaluations.
    :param hide_output: Whether to suppress output from the rbf_opt algorithm
    :param bonmin_path: The path to Bonmin, a specific feature useful for the cluster.
    :param rand_seed: Give a seed number to make sure the random has numbers.
    :return: a list of solutions

    """

    # default setting from rbfopt or do something custom (specifically for the cluster)
    # This parameter points rbfopt to the solver it was written to use.
    if bonmin_path == None:
        bonmin_path_set = "bonmin"
    else:
        bonmin_path_set = bonmin_path
    # This parameter seeds the random number generator of numpy.
    # To do a proper assessment this needs to be randomly set by the user.
    # else every run will just do the same search.
    if rand_seed == None:
        rand_seed_set = 937627691
    else:
        rand_seed_set = rand_seed

    problem = evaluator.problem
    inputs = problem.inputs
    dimension = problem.num_inputs
    num_outputs = problem.num_outputs

    assert (
        num_outputs == 1
    ), f"This is a single objective optimizer, but you have {num_outputs} objectives"
    assert problem.num_constraints == 0, "This optimizer cannot handle constraints"
    # may throw an attribute error if invalid parameters are used (i.e categorical)
    param_attrs = [
        (p.value_descriptor.min, p.value_descriptor.max, p.value_descriptor.rbf_type)
        for p in inputs
    ]
    var_lower, var_upper, var_type = (np.array(x) for x in zip(*param_attrs))
    # with the current classes var_type is just an array of 'R's
    # this would change if we add a descriptor for integers
    print(var_lower, var_upper, var_type)

    def wrapper(values):
        return evaluator(values)[0]

    def run_alg():
        bb = rbfopt.RbfoptUserBlackBox(
            dimension=dimension,
            var_lower=var_lower,
            var_upper=var_upper,
            var_type=var_type,
            obj_funct=wrapper,
        )
        settings = rbfopt.RbfoptSettings(
            max_evaluations=evaluations,
            minlp_solver_path=bonmin_path_set,
            rand_seed=rand_seed_set,
        )
        alg = rbfopt.RbfoptAlgorithm(settings, bb)
        return alg.optimize()

    if hide_output:
        with open(os.devnull, "w") as null:
            old, sys.stdout = sys.stdout, null
            try:
                val, x, itercount, evalcount, fast_evalcount = run_alg()
            finally:
                sys.stdout = old
    else:
        val, x, itercount, evalcount, fast_evalcount = run_alg()

    return problem.to_df(np.array([np.concatenate([x, [val]])]))
