from loguru import logger
from pymop.problem import Problem
import numpy as np
from search import micro_encoding
from projectcode.training import train_and_evaluate, pickled_train_and_evaluate
from compress_pickle import dump, load
import os
import peewee
from projectcode.server.database import (
    Task,
    pickle_to_str,
    get_list_of_not_completed_tasks,
    get_task,
    pickle_load_from_str,
    get_list_of_inprogress_tasks,
    get_list_of_completed_tasks,
    delete_tasks_of_generation,
)
from pathlib import Path
import time
import main_config


class NAS(Problem):
    # first define the NAS problem (inherit from pymop)
    def __init__(
        self,
        search_space="micro",
        n_var=20,
        n_obj=1,
        n_constr=0,
        lb=None,
        ub=None,
        init_channels=24,
        layers=8,
        epochs=25,
        save_dir=None,
        args=None,
        debug=False,
        mode="single",
    ):
        super().__init__(n_var=n_var, n_obj=n_obj, n_constr=n_constr, type_var=np.int)
        self.xl = lb
        self.xu = ub
        self._search_space = search_space
        self._init_channels = init_channels
        self._layers = layers
        self._epochs = epochs
        self._save_dir = save_dir
        self._n_evaluated = 0  # keep track of how many architectures are sampled
        self.args = args
        self.debug = debug
        self.mode = mode
        self.generation = None

        assert self.args is not None

    def _evaluate(self, x, out, individuals, algorithm, *args, **kwargs):

        self.generation = algorithm.n_gen

        if algorithm.pop is None:
            logger.info("Initializing population")
            first_gen = True
        else:
            logger.info("Evaluating New Members")
            first_gen = False

        objs = np.full((x.shape[0], self.n_obj), np.nan)

        pkl_paths = []

        logger.info("Maybe Creating Directory for Saving Training Pickles")
        Path(os.path.join(self.args.save, "train_spec")).mkdir(
            parents=True, exist_ok=True
        )

        for i in range(x.shape[0]):

            # unique id assigned to the individual
            arch_id = self._n_evaluated + 1

            individuals[i].id = arch_id
            individual = individuals[i]
            print("\n")
            logger.info("Network id = {}".format(arch_id))

            # call back-propagation training
            if self._search_space == "micro":
                genome = micro_encoding.convert(x[i, :])
            else:
                raise KeyError("unknown search space")

            if self.debug:
                performance = {
                    "valid_acc": 100,
                    "params": 100,
                    "flops": 100,
                }
            else:
                logger.info("pickling training specification")
                save = "arch_{}".format(arch_id)

                spec = {
                    "genome": genome,
                    "individual": individual,
                    "save": save,
                    "args": self.args,
                    "first_gen": first_gen,
                }

                pkl_path = os.path.join(
                    self.args.save, "train_spec", f"individual_{arch_id:05d}.pkl"
                )
                pkl_paths.append(pkl_path)

                dump(spec, pkl_path)

                self._n_evaluated += 1

                # performance = pickled_train_and_evaluate(pkl_path)

                # performance = train_and_evaluate(
                #     genome=genome,
                #     individual=individual,
                #     save=save,
                #     args=self.args,
                #     first_gen=first_gen,
                # )

        results = self.get_evaluation(pkl_paths)

        logger.info("Evaluation Complete, Assigning Results to Population")

        for i, performance in enumerate(results):

            # all objectives assume to be MINIMIZED !!!!!
            objs[i, 0] = 100 - performance["valid_acc"]
            objs[i, 1] = performance["flops"]

        out["F"] = objs
        # self.generation += 1

    def get_evaluation(self, pkl_paths):

        results = []

        if self.mode == "single":

            logger.info("Evaluating in single mode")

            for i, pkl_path in enumerate(pkl_paths):
                performance = pickled_train_and_evaluate(pkl_path)
                results.append(performance)

        elif self.mode == "server":

            logger.info("Evaluating in Distributed Server Mode")

            removed = len(delete_tasks_of_generation(self.generation, self.args.code))
            logger.info(
                "Removed %i Tasks from DB of generation %i" % (removed, self.generation)
            )

            logger.info("Creating Tasks in DB")

            for pkl_path in pkl_paths:

                spec = load(pkl_path)
                network_id = spec["individual"].id

                task = Task()
                task.pkl_path = pkl_path
                task.generation = self.generation
                task.pkl_data = pickle_to_str(pkl_path)
                task.network_id = network_id
                task.run_code = self.args.code
                task.client_stream = self.args.stream
                task.save()

                # TODO: add integrity check to ensure that every task created can actually inherit genes if lammarckian and not first gen

            logger.info("Waiting for Task Completion")
            run_completed = False
            previous_completed = 0
            for i in range(360 * 6):

                s = "Checking Completion for Generation %i" % self.generation
                logger.debug(s)

                completed = len(
                    get_list_of_completed_tasks(
                        self.generation, run_code=self.args.code
                    )
                )
                inprogress = len(get_list_of_inprogress_tasks(run_code=self.args.code))
                total = len(pkl_paths)

                s = "Generation %i tasks inprogress/completed/total : %i/%i/%i" % (
                    self.generation,
                    inprogress,
                    completed,
                    total,
                )
                if previous_completed != completed:
                    logger.info(s)
                else:
                    logger.debug(s)

                previous_completed = completed

                incompletes = len(
                    get_list_of_not_completed_tasks(run_code=self.args.code)
                )

                if incompletes == 0:
                    logger.info("all tasks are complete")


                    run_completed = True

                    l = list(Task.collection.filter(run_code=self.args.code, generation=self.generation).fetch())
                    l = sorted(l, key=lambda x: x.network_id)

                    results = []

                    for task in l:
                        try:
                            result = pickle_load_from_str(task.result_pkl_data)
                            results.append(result)
                        except Exception as e:
                            logger.exception()
                            logger.warning("Task with network id of %i has no pickle data" %task.network_id)
                            run_completed = False


                    if run_completed:
                        break
                else:
                    time.sleep(main_config.server_task_progress_polling_interval)


            assert run_completed
            if not run_completed:
                raise KeyError("Tasks still incomplete after timingout")

            logger.info("Waiting a bit before reading results")
            time.sleep(30)

            logger.info("Reading Results of Completed Tasks")
            results = results






        else:
            raise NameError("unrecognized mode")

        return results


def do_every_generations(algorithm):
    # this function will be call every generation
    # it has access to the whole algorithm class
    gen = algorithm.n_gen
    pop_var = algorithm.pop.get("X")
    pop_obj = algorithm.pop.get("F")

    # report generation info to files
    logger.info("generation = {}".format(gen))
    logger.info(
        "population error: best = {}, mean = {}, "
        "median = {}, worst = {}".format(
            np.min(pop_obj[:, 0]),
            np.mean(pop_obj[:, 0]),
            np.median(pop_obj[:, 0]),
            np.max(pop_obj[:, 0]),
        )
    )
    logger.info(
        "population complexity: best = {}, mean = {}, "
        "median = {}, worst = {}".format(
            np.min(pop_obj[:, 1]),
            np.mean(pop_obj[:, 1]),
            np.median(pop_obj[:, 1]),
            np.max(pop_obj[:, 1]),
        )
    )
