from itertools import product
from yoshix.yoshiegg import YoshiEgg, YoshiEggKeyException


class YoshiExperiment(object):
    """
    `YoshiExperiment` is the base class for every user created experiment.

    The class provide the basic interface and infrastructure to register data,
    run single experiments, generate input data and so on.

    The class should never be initialized by its own. This should be always
    extended by a child class.
    """

    def __init__(self):
        self.name = self.__class__.__name__
        self._generators = {}  # Dictionary to map a parameter to a generator.
        self.__egg = None
        self._egg_is_ready = False  # Egg is ready only after the experiment.
        self.__empty_row = None  # Store the initialization values for the egg rows.
        self.__run_counter = 0  # Store the iteration number.
        self.__fixed_parameters = {}  # Store the fixed parameters.
        self.__generators_iterator = None # Store the combined product-iterator for every generators.
        self.__parameter_transformer = {} # Store a transformation function for the egg representation of the parameter.
        self.__private_generators = [] # List of hidden generators. Hidden generators are not put in the egg.

    def setup(self):
        """
        This function is called before the experiment is started. This can be
        used to initialize variables, generators and every other detail.
        :return: None
        """
        pass

    def single_run(self, params):
        """
        This represent the atomic experiment run.
        :param params List of parameters for the experiment run. This is automatically
        generated by the parent method.
        :return: None
        """
        pass

    def _run_experiment(self):
        """
        The wrapping experiment loop. This function invokes single_run for
        every combination of **variable parameters** provided by the generators.
        :return: None
        """
        self.__run_counter = 0
        while True:  # TODO: Is there a better way to iterate until the Iterator is empty?
            try:
                # Generating Parameters dictionary
                params = self.__fixed_parameters.copy()
                params.update(self.__generate())

                self.__egg.add_row(self.__empty_row)
                self.__run_counter += 1

                # Add the variable parameters to the egg.
                for g in self._generators.keys():
                    if g not in self.__private_generators:
                        self.__egg[g] = self.__apply_transformer(g, params[g])

                self.single_run(params)
            except StopIteration:
                break
        self._egg_is_ready = True

    def after_run(self):
        """
        This method is invoked after the experiment is completed.

        Can be used to package the result Egg, clean up the disk, export to CSV and more.
        :return:
        """
        pass

    @property
    def partial_egg(self):
        """
        :return: Return an external reference to the experiment Egg to be used **during** the experiment.
        """
        if self.__egg is None:
            raise EggNotReady("Try to access an egg that is None")
        else:
            return self.__egg

    @property
    def egg(self):
        """
        :return: Return an external reference to the experiment **AFTER** the experiment.
        """
        if self._egg_is_ready:
            return self.__egg
        if self.__egg is None:
            raise EggNotReady("Try to access an egg that is None")
        elif not self._egg_is_ready:
            raise EggNotReady("The egg is there but the experiment is not completed yet!\n\
            Maybe you are looking for partial_egg?")
        else:
            raise Exception("Something is really wrong there!")

    @property
    def run_counter(self):
        """
        :return: Return the number of the current iteration.
        """
        return self.__run_counter

    def setup_egg(self, data_headers, row_initialization=None):
        """
        This method is used to initialize the experiment egg.
        :param data_headers: The tuple of the experiment parameters and desired computed outputs.
        :param row_initialization: A vector representing an empty row. Default is a vector of zeros.
        :return:
        """
        self.__egg = YoshiEgg(data_headers)
        # If row_init is None we assume all zeroes.
        if row_initialization is None:
            row_initialization = tuple((0 for _ in data_headers))
        if len(row_initialization) == len(data_headers):
            self.__empty_row = row_initialization
        else:
            raise YoshiEggKeyException("Initialization vector does not match the header.")

    def assign_generators(self, key, generator, private=False):
        """
        Link a generator with a particular parameter of the algorithm.
        :param key: The parameter key identifier.
        :param generator: The desired generator.
        :return:
        """
        if key not in self.__egg and not private:
            raise YoshiEggKeyException("It is not possible to attach a generator to an unknown key!")
        self._generators[key] = generator
        gen_list = [v for _, v in self._generators.items()]
        self.__generators_iterator = product(*gen_list)
        if private:
            self.__private_generators.append(key)

    def assign_transformer(self, key, transformer):
        if key not in self.__egg:
            raise YoshiEggKeyException("It is not possible to attach a transformator to an unknown key!")
        self.__parameter_transformer[key] = transformer

    def assign_fixed_parameter(self, key, value):
        """
        Link a parameter with a fixed value.
        :param key: The parameter key identifier.
        :param value: The desired value.
        :return:
        """
        if key not in self.__egg:
            raise YoshiEggKeyException("It is not possible to attach a value to an unknown key!")
        self.__fixed_parameters[key] = value

    def __generate(self):
        """
        This is used in order to generate a new set of variable parameters (using the generators list)
        :return: A dictionary with the current variable parameters values.
        """
        if self.__generators_iterator is not None:
            current_iteration = next(self.__generators_iterator)
            return {k: v for k, v in zip(self._generators.keys(), current_iteration)}

    def __apply_transformer(self, key, value):
        return self.__parameter_transformer[key](value) \
            if key in self.__parameter_transformer else value

    def run(self):
        self.setup()
        self._run_experiment()
        self.after_run()


class EggNotReady(Exception):
    pass


