import json
import numpy as np
import hyperopt
import random

from ...suggestion.models import Study
from ...suggestion.models import Trial
from ...suggestion.algorithm.abstract_algorithm import AbstractSuggestionAlgorithm


class BaseHyperoptAlgorithm(AbstractSuggestionAlgorithm):
    """
    The implementation is based on https://github.com/tobegit3hub/advisor
    """

    def __init__(self, algorithm_name="tpe"):

        if algorithm_name == "tpe":
            self.hyperopt_algorithm = hyperopt.tpe.suggest
        elif algorithm_name == "random_search":
            self.hyperopt_algorithm = hyperopt.rand.suggest
        elif algorithm_name == "anneal":
            self.hyperopt_algorithm = hyperopt.anneal.suggest

    def get_new_suggestions(self, study, trials=[], number=1):
        """
        Get the new suggested trials with TPE algorithm.
        """
        # Construct search space, example: {"x": hyperopt.hp.uniform('x', -10, 10), "x2": hyperopt.hp.uniform('x2', -10, 10)}
        hyperopt_search_space = {}

        # study = Study.objects.get(name=study_name)
        study_configuration_json = json.loads(study.study_configuration)
        params = study_configuration_json["params"]

        for param in params:
            param_name = param["parameterName"]

            if param["type"] == "INTEGER":
                # TODO: Support int type of search space)
                pass

            elif param["type"] == "DOUBLE":
                hyperopt_search_space[param_name] = hyperopt.hp.uniform(
                    param_name, param["minValue"], param["maxValue"]
                )

            elif param["type"] == "DISCRETE" or param["type"] == "CATEGORICAL":
                feasible_point_list = [
                    value.strip() for value in param["feasiblePoints"].split(",")
                ]
                hyperopt_search_space[param_name] = hyperopt.hp.choice(
                    param_name, feasible_point_list
                )

        # New hyperopt variables
        hyperopt_rstate = np.random.RandomState(random.randint(1, 2 ** 31 - 1))
        hyperopt_domain = hyperopt.Domain(
            None, hyperopt_search_space, pass_expr_memo_ctrl=None
        )

        hyperopt_trial_specs = []
        hyperopt_trial_results = []
        # Example: # Example: [{'tid': 0, 'idxs': {'l1_normalization': [0], 'learning_rate': [0], 'hidden2': [0], 'optimizer': [0]}, 'cmd': ('domain_attachment', 'FMinIter_Domain'), 'vals': {'l1_normalization': [0.1], 'learning_rate': [0.1], 'hidden2': [1], 'optimizer': [1]}, 'workdir': None}]
        hyperopt_trial_miscs = []
        hyperopt_trial_new_ids = []

        # Update hyperopt for trained trials with completed advisor trials
        completed_hyperopt_trials = hyperopt.Trials()

        # completed_advisor_trials = Trial.objects.filter(
        #    study_name=study_name, status="Completed")
        completed_advisor_trials = [i for i in trials if i.status == "Completed"]

        for index, advisor_trial in enumerate(completed_advisor_trials):
            # Example: {"learning_rate": 0.01, "optimizer": "ftrl"}
            parameter_values_json = json.loads(advisor_trial.parameter_values)

            # Example: {'l1_normalization': [0], 'learning_rate': [0], 'hidden2': [0], 'optimizer': [0]}
            hyperopt_trial_miscs_idxs = {}
            # Example: {'l1_normalization': [0.1], 'learning_rate': [0.1], 'hidden2': [1], 'optimizer': [1]}
            hyperopt_trial_miscs_vals = {}
            new_id = index
            hyperopt_trial_new_ids.append(new_id)
            hyperopt_trial_misc = dict(
                tid=new_id, cmd=hyperopt_domain.cmd, workdir=hyperopt_domain.workdir
            )

            for param in params:

                if param["type"] == "INTEGER":
                    pass

                elif param["type"] == "DOUBLE":
                    parameter_value = parameter_values_json[param["parameterName"]]
                    hyperopt_trial_miscs_idxs[param["parameterName"]] = [index]
                    hyperopt_trial_miscs_vals[param["parameterName"]] = [
                        parameter_value
                    ]

                elif param["type"] == "DISCRETE":
                    feasible_points_string = param["feasiblePoints"]
                    feasible_points = [
                        float(value.strip())
                        for value in feasible_points_string.split(",")
                    ]
                    parameter_value = parameter_values_json[param["parameterName"]]
                    index_of_value_in_list = feasible_points.index(parameter_value)
                    hyperopt_trial_miscs_idxs[param["parameterName"]] = [index]
                    hyperopt_trial_miscs_vals[param["parameterName"]] = [
                        index_of_value_in_list
                    ]

                elif param["type"] == "CATEGORICAL":
                    feasible_points_string = param["feasiblePoints"]
                    feasible_points = [
                        value.strip() for value in feasible_points_string.split(",")
                    ]
                    # Example: "ftrl"
                    parameter_value = parameter_values_json[param["parameterName"]]
                    index_of_value_in_list = feasible_points.index(parameter_value)
                    hyperopt_trial_miscs_idxs[param["parameterName"]] = [index]
                    hyperopt_trial_miscs_vals[param["parameterName"]] = [
                        index_of_value_in_list
                    ]

            hyperopt_trial_specs.append(None)

            hyperopt_trial_misc["idxs"] = hyperopt_trial_miscs_idxs
            hyperopt_trial_misc["vals"] = hyperopt_trial_miscs_vals
            hyperopt_trial_miscs.append(hyperopt_trial_misc)

            # TODO: Use negative objective value for loss or not

            loss_for_hyperopt = advisor_trial.objective_value
            if study_configuration_json["goal"] == "MAXIMIZE":
                # Now hyperopt only supports fmin and we need to reverse objective value for maximization
                loss_for_hyperopt = -1 * advisor_trial.objective_value

            hyperopt_trial_result = {
                "loss": loss_for_hyperopt,
                "status": hyperopt.STATUS_OK,
            }
            hyperopt_trial_results.append(hyperopt_trial_result)

        if len(completed_advisor_trials) > 0:
            # Example: {'refresh_time': datetime.datetime(2018, 9, 18, 12, 6, 41, 922000), 'book_time': datetime.datetime(2018, 9, 18, 12, 6, 41, 922000), 'misc': {'tid': 0, 'idxs': {'x2': [0], 'x': [0]}, 'cmd': ('domain_attachment', 'FMinIter_Domain'), 'vals': {'x2': [-8.137088361136204], 'x': [-4.849028446711832]}, 'workdir': None}, 'state': 2, 'tid': 0, 'exp_key': None, 'version': 0, 'result': {'status': 'ok', 'loss': 14.849028446711833}, 'owner': None, 'spec': None}
            hyperopt_trials = completed_hyperopt_trials.new_trial_docs(
                hyperopt_trial_new_ids,
                hyperopt_trial_specs,
                hyperopt_trial_results,
                hyperopt_trial_miscs,
            )
            for current_hyperopt_trials in hyperopt_trials:
                current_hyperopt_trials["state"] = hyperopt.JOB_STATE_DONE

            completed_hyperopt_trials.insert_trial_docs(hyperopt_trials)
            completed_hyperopt_trials.refresh()

        rval = hyperopt.FMinIter(
            self.hyperopt_algorithm,
            hyperopt_domain,
            completed_hyperopt_trials,
            max_evals=-1,
            rstate=hyperopt_rstate,
            verbose=0,
        )
        rval.catch_eval_exceptions = False

        new_ids = rval.trials.new_trial_ids(number)

        rval.trials.refresh()

        random_state = rval.rstate.randint(2 ** 31 - 1)
        new_trials = self.hyperopt_algorithm(
            new_ids, rval.domain, completed_hyperopt_trials, random_state
        )
        rval.trials.refresh()

        # Construct return advisor trials from new hyperopt trials
        return_trial_list = []

        for i in range(number):

            # Example: {u'hidden2': [2], u'learning_rate': [0.04633366105812467], u'l1_normalization': [0.16858448611765364], u'optimizer': [3]}
            vals = new_trials[i]["misc"]["vals"]

            new_advisor_trial = Trial.create(study.name, "TpeTrial")
            parameter_values_json = {}

            for param in params:

                if param["type"] == "INTEGER":
                    pass

                elif param["type"] == "DOUBLE":
                    suggest_value = vals[param["parameterName"]][0]
                    parameter_values_json[param["parameterName"]] = suggest_value

                elif param["type"] == "DISCRETE":
                    feasible_point_list = [
                        float(value.strip())
                        for value in param["feasiblePoints"].split(",")
                    ]
                    suggest_index = vals[param["parameterName"]][0]
                    suggest_value = feasible_point_list[suggest_index]

                elif param["type"] == "CATEGORICAL":
                    feasible_point_list = [
                        value.strip() for value in param["feasiblePoints"].split(",")
                    ]
                    suggest_index = vals[param["parameterName"]][0]
                    suggest_value = feasible_point_list[suggest_index]

                parameter_values_json[param["parameterName"]] = suggest_value

            new_advisor_trial.parameter_values = json.dumps(parameter_values_json)
            # new_advisor_trial.save()
            return_trial_list.append(new_advisor_trial)

        return return_trial_list
