import numpy as np

from l2l.optimizees.optimizee import Optimizee


class FunctionGeneratorOptimizee(Optimizee):
    """
    Implements a simple function optimizee. Functions are generated using the FunctionGenerator.
    NOTE: Make sure the optimizee_fitness_weights is set to (-1,) to minimize the value of the function

    :param traj: The trajectory used to conduct the optimization.
    :param fg_instance: Instance of the FunctionGenerator class.
    :param seed: The random seed used for generation of optimizee individuals. It uses a copy of
        the fg_instance and overrides the random generator using one seeded by `seed`. Note that this
        random generator is also the one used by the :class:`.FunctionGeneratorOptimizee` itself.
        NOTE that this seed is converted to an np.uint32.
    """

    def __init__(self, traj, fg_instance, seed):
        super().__init__(traj)

        seed = np.uint32(seed)
        self.random_state = np.random.RandomState(seed=seed)

        self.fg_instance = fg_instance
        self.dims = self.fg_instance.dims
        self.cost_fn = self.fg_instance.cost_function
        self.bound = self.fg_instance.bound
        self.random_state = np.random.RandomState(seed=seed)

    def create_individual(self):
        """
        Creates a random value of parameter within given bounds
        """
        # Define the first solution candidate randomly
        return {'coords': self.random_state.rand(self.dims) * (self.bound[1] - self.bound[0]) + self.bound[0]}

    def bounding_func(self, individual):
        """
        Bounds the individual within the required bounds via coordinate clipping
        """
        return {'coords': np.clip(individual['coords'], a_min=self.bound[0], a_max=self.bound[1])}

    def simulate(self, traj):
        """
        Returns the value of the function chosen during initialization

        :param ~l2l.utils.trajectory.Trajectory traj: Trajectory
        :return: a single element :obj:`tuple` containing the value of the chosen function
        """

        individual = np.array(traj.individual.coords)
        return self.cost_fn(individual, random_state=self.random_state)
        
        
    def simulate_(self, traj):
        """
        Returns the fitness value and acts as a wrapper function for the multi-optimizer
        :param ~l2l.utils.trajectory.Trajectory traj: Trajectory
        :return: fitness values as a list for multi-optimzer, or single element
        """
        if(traj.inner_params > 1):
            fitness = []
            individuals  = np.array(traj.individual.coords)
            for i in range(traj.inner_params):
                traj.individual.coords = individuals[i]
                fitness.append(self.simulate(traj))
            return fitness
        else:
            return (self.simulate(traj),)
