import numpy
import numpy as np


class GeneticAlgorithm:
    """Genetic algorithm.
    Implementation of the real-valued Genetic algorithm. The mutations are
    normally distributed perturbations, the selection mechanism is a tournament
    selection, and the crossover oepration is the standard linear combination
    taken at a randomly generated cutting point.
    The total number of evaluations are popsize x ngen
    :param function: Function that can be used to evaluate the entire
        population. It needs to take an input of size pop_size x dim and
        return a numpy.array of size pop_size x 1
    :type function: Object
    :param dim: Number of dimensions
    :type dim: int
    :param lb: Lower variable bounds, of length dim
    :type lb: numpy.array
    :param ub: Lower variable bounds, of length dim
    :type ub: numpy.array
    :param int_var: List of indices with the integer valued variables
        (e.g., [0, 1, 5])
    :type int_var: list
    :param pop_size: Population size
    :type pop_size: int
    :param num_gen: Number of generations
    :type num_gen: int
    :param start: Method for generating the initial population
    :type start: string
    :ivar nvariables: Number of variables (dimensions)
    :ivar nindividuals: population size
    :ivar lower_boundary: lower bounds for the optimization problem
    :ivar upper_boundary: upper bounds for the optimization problem
    :ivar integer_variables: List of variables that are integer valued
    :ivar start: Method for generating the initial population
    :ivar sigma: Perturbation radius. Each pertubation is N(0, sigma)
    :ivar p_mutation: Mutation probability (1/dim)
    :ivar tournament_size: Size of the tournament (5)
    :ivar p_cross: Cross-over probability (0.9)
    :ivar ngenerations: Number of generations
    :ivar function: Object that can be used to evaluate the objective function
    """

    def __init__(self, function, dim, lb, ub, int_var=None, pop_size=20, num_gen=300, start="Random"):

        self.nvariables = dim  # column
        self.nindividuals = pop_size + (pop_size % 2)  # Make sure this is even row
        self.lower_boundary = np.array(lb)
        self.upper_boundary = np.array(ub)
        self.integer_variables = []
        if int_var is not None:
            self.integer_variables = np.array(int_var)
        self.start = start
        self.sigma = 0.2
        self.p_mutation = 1.0 / dim
        self.tournament_size = 5
        self.p_cross = 0.9
        self.ngenerations = num_gen
        self.function = function

    def optimize(self):
        """Method used to run the Genetic algorithm
        :return: Returns the best individual and its function value
        :rtype: numpy.array, float
        """
        #  Initialize population
        if isinstance(self.start, np.ndarray):
            if self.start.shape[0] != self.nindividuals or self.start.shape[1] != self.nvariables:
                raise ValueError("Initial population has incorrect size")
            if any(np.min(self.start, axis=0) < self.lower_boundary) or any(
                    np.max(self.start, axis=0) > self.upper_boundary
            ):
                raise ValueError("Initial population is outside the domain", self.lower_boundary, self.upper_boundary,
                                 self.start)
            population = self.start
        elif self.start == "SLHD":
            from pySOT.experimental_design import SymmetricLatinHypercube

            exp_des = SymmetricLatinHypercube(self.nvariables, self.nindividuals)
            population = self.lower_boundary + exp_des.generate_points() * (self.upper_boundary - self.lower_boundary)
        elif self.start == "LHD":
            from pySOT.experimental_design import LatinHypercube

            exp_des = LatinHypercube(self.nvariables, self.nindividuals)
            population = self.lower_boundary + exp_des.generate_points() * (self.upper_boundary - self.lower_boundary)
        elif self.start == "Random":
            population = self.lower_boundary + np.random.rand(self.nindividuals, self.nvariables) * (
                    self.upper_boundary - self.lower_boundary
            )
        else:
            raise ValueError("Unknown argument for initial population")

        new_population = []
        #  Round positions
        if len(self.integer_variables) > 0:  # 对特定列进行操作
            new_population = np.copy(population)
            population[:, self.integer_variables] = np.round(population[:, self.integer_variables])
            for i in self.integer_variables:
                ind = np.where(population[:, i] < self.lower_boundary[i])
                population[ind, i] += 1
                ind = np.where(population[:, i] > self.upper_boundary[i])
                population[ind, i] -= 1

        #  Evaluate all individuals
        # function_values = self.function(population) we cannot compute in this way to ensure x is one-dim in policy
        n_row, n_dim = population.shape
        function_values = []
        for r in range(n_row):
            function_values.append(self.function(population[r, :]))
        function_values = np.array(function_values)

        if len(function_values.shape) == 2:
            function_values = np.squeeze(np.asarray(function_values))

        # Save the best individual
        ind = np.argmin(function_values)
        best_individual = np.copy(population[ind, :])  # 找到最优个体
        best_value = function_values[ind]

        if len(self.integer_variables) > 0:
            population = new_population

        # Main loop
        for _ in range(self.ngenerations):
            print('------------------------------')
            print("当前为第{}代".format(_))
            print("最优个体为:{}".format(best_individual))
            print("最优值为:{}".format(best_value))
            print("------------------------------")
            # Do tournament selection to select the parents
            competitors = np.random.randint(0, self.nindividuals, (self.nindividuals, self.tournament_size))
            ind = np.argmin(function_values[competitors], axis=1)
            winner_indices = np.zeros(self.nindividuals, dtype=int)
            for i in range(self.tournament_size):  # This loop is short
                winner_indices[np.where(ind == i)] = competitors[np.where(ind == i), i]

            parent1 = population[winner_indices[0: self.nindividuals // 2], :]
            parent2 = population[winner_indices[self.nindividuals // 2: self.nindividuals], :]

            # Averaging Crossover
            cross = np.where(np.random.rand(self.nindividuals // 2) < self.p_cross)[0]
            nn = len(cross)  # Number of crossovers
            alpha = np.random.rand(nn, 1)

            # Create the new chromosomes
            parent1_new = np.multiply(alpha, parent1[cross, :]) + np.multiply(1 - alpha, parent2[cross, :])
            parent2_new = np.multiply(alpha, parent2[cross, :]) + np.multiply(1 - alpha, parent1[cross, :])
            parent1[cross, :] = parent1_new
            parent2[cross, :] = parent2_new
            population = np.concatenate((parent1, parent2))

            # Apply mutation
            scale_factors = self.sigma * (self.upper_boundary - self.lower_boundary)  # Scale
            perturbation = np.random.randn(self.nindividuals, self.nvariables)  # Generate perturbations
            perturbation = np.multiply(perturbation, scale_factors)  # Scale accordingly
            perturbation = np.multiply(
                perturbation, (np.random.rand(self.nindividuals, self.nvariables) < self.p_mutation)
            )

            perturbation = round_vars(perturbation, self.integer_variables, self.lower_boundary, self.upper_boundary)
            population = round_vars(population, self.integer_variables, self.lower_boundary, self.upper_boundary)

            population += perturbation  # Add perturbation
            population = np.maximum(np.reshape(self.lower_boundary, (1, self.nvariables)), population)
            population = np.minimum(np.reshape(self.upper_boundary, (1, self.nvariables)), population)

            # Round chromosomes
            new_population = []
            if len(self.integer_variables) > 0:
                new_population = np.copy(population)
                population = round_vars(population, self.integer_variables, self.lower_boundary, self.upper_boundary)

            # Keep the best individual
            population[0, :] = best_individual

            #  Evaluate all individuals
            # function_values = self.function(population) we cannot compute in this way to ensure x is one-dim in policy
            n_row, n_dim = population.shape
            function_values = []
            for r in range(n_row):
                function_values.append(self.function(population[r, :]))
            function_values = np.array(function_values)

            if len(function_values.shape) == 2:
                function_values = np.squeeze(np.asarray(function_values))

            # Save the best individual
            ind = np.argmin(function_values)
            best_individual = np.copy(population[ind, :])
            best_value = function_values[ind]

            # Use the positions that are not rounded
            if len(self.integer_variables) > 0:
                population = new_population

        # return best_individual, best_value


def round_vars(x: np.ndarray, int_var, lb, ub):
    """Round integer variables to closest integer in the domain.
    :param x: Set of points, of size npts x dim
    :type x: numpy.array
    :param int_var: Set of indices of integer variables
    :type int_var: numpy.array
    :param lb: Lower bounds, of size 1 x dim
    :type lb: numpy.array
    :param ub: Upper bounds, of size 1 x dim
    :type ub: numpy.array
    :return: The set of points with the integer variables
        rounded to the closest integer in the domain
    :rtype: numpy.array
    """
    # Make sure we don't violate the bound constraints
    for i in int_var:
        ind = np.where(x[:, i] < lb[i])
        x[ind, i] = lb[i]
        ind = np.where(x[:, i] > ub[i])
        x[ind, i] = ub[i]

    if len(int_var) > 0:
        # Round the original ranged integer variables
        x[:, int_var] = np.round(x[:, int_var])
        x = x.astype(numpy.int32, copy=True)
    else:
        x = x.astype(numpy.float64, copy=True)
    return x