"""
Module that implements the CMA evolutionary strategy.
"""

import numpy as np

from evolutionary_strategy import Candidate, EvolutionaryStrategy
from utils import clamp, en0i


class CMAEvolutionaryStrategy(EvolutionaryStrategy):
    """
    The module implementing the CMA evolutionary strategy.
    References: https://en.wikipedia.org/wiki/CMA-ES, https://arxiv.org/pdf/1604.00772.pdf,
                https://www.youtube.com/watch?v=Hk9rG1_CMKA&t=1119s
    """

    def __init__(self, left_limit, right_limit, car, starting_velocity=0.0, weight_group_size=20, smoothing_length=9,
                 smoothing_order=1, population_size=50, iterations=5, standard_deviation=0.3):
        """
        Method to initialize the CMA evolutionary strategy.

        Args:
            left_limit(RacingLine): Left side track limit.
            right_limit(RacingLine): Right side track limit.
            car(Car): Car to evaluate the racing line for.
            starting_velocity(float): Starting velocity of the car in m/s. Defaults to 0.0 m/s.
            weight_group_size(int): Number of weights to group together. Defaults to 20.
            smoothing_length(int): Length of weights to smoothen at the ends of each weight group. Defaults to 9.
            smoothing_order(int): Order of interpolation to use for smoothing. Defaults to 1 (Linear).
            population_size(int): Size of the population to be used. Defaults to 50.
            iterations(int): Number of iterations to run the algorithm for. Defaults to 5.
            standard_deviation(float): Standard deviation to use when constructing the path vector. Defaults to 0.3.

        """
        super(CMAEvolutionaryStrategy, self).__init__(
            left_limit=left_limit, right_limit=right_limit, car=car, starting_velocity=starting_velocity,
            weight_group_size=weight_group_size, smoothing_length=smoothing_length, smoothing_order=smoothing_order,
            population_size=population_size, iterations=iterations, standard_deviation=standard_deviation)

        # Population variables.
        self.population_matrix = np.array([[]])
        self.means = np.array([[]])
        self.covariance_matrix = np.array([[]])

        # Path vectors
        self.p_c = np.array([[]])
        self.p_sigma = np.array([[]])

        # Number of parents/candidates to select.
        self.mu = int(self.population_size / 2)

        # Weight of each selected element.
        self.weight = 1.0 / self.mu

        # Number of relevant individuals to update covariance matrix.
        self.mu_eff = ((self.mu * self.weight) ** 2) / (self.mu * (self.weight ** 2))

        # Parameter for weighing between rank-one and rank-mu update.
        self.mu_cov = self.mu_eff

        # Learning rate for cumulation of the rank-one update of the covariance matrix.
        self.c_cov = (1.0 / self.mu_cov) * 2.0 / ((self.candidate_length + 1.4) ** 2) + (1.0 - 1.0 / self.mu_cov) * \
            ((2.0 * self.mu_eff - 1.0) / ((self.candidate_length + 2.0) ** 2 + 2.0 * self.mu_eff))

        # Parameters for backward time horizon.
        self.c_c = 1.0 / self.candidate_length
        self.c_sigma = self.c_c

        # Damping parameter for step-size update.
        self.d_sigma = 1.0

        # Stores the Eigen values and vectors of the covariance matrix.
        self.eigen_values = np.array([[]])
        self.eigen_vectors = np.array([[]])

    def generate_population(self):
        """
        Method to generate the population for the solution set.

        Returns:
            population(list): List of Candidates in the population.

        """
        super(CMAEvolutionaryStrategy, self).generate_population()
        self.population_matrix = np.array([candidate.weights for candidate in self.population])
        self.means = np.reshape(self.population_matrix.sum(axis=0), (-1, 1)) / float(self.population_size)
        self.covariance_matrix = self.generate_empirical_covariance_matrix()
        self.p_c = np.reshape(np.array([[0.0] * self.candidate_length]), (-1, 1))
        self.p_sigma = np.reshape(np.array([[0.0] * self.candidate_length]), (-1, 1))
        self.sort_population_by_fitness()
        return self.population

    def generate_offspring(self):
        """
        Method to generate the offspring for the population.

        Returns:
            offspring(list): List of offspring Candidates that were generated by mutating the parents.

        """
        offspring = []

        for r_idx in range(self.population_size):
            random_vector = np.random.normal(0.0, 1.0, self.candidate_length)
            random_vector = np.reshape(random_vector, (-1, 1))

            new_weights = self.means + self.standard_deviation * (self.eigen_vectors.dot(
                self.eigen_values.dot(random_vector)))
            new_weights = new_weights.transpose()[0]
            new_weights = [clamp(w, 0.0, 1.0) for w in new_weights]

            new_candidate = Candidate(list(new_weights))
            offspring.append(new_candidate)
            self.population_matrix[r_idx] = new_weights

        self.population = offspring
        return offspring

    def run(self):
        """
        Method to run the strategy.

        Returns:
            best_candidate(Candidate): Candidate with the best fitness (least lap time) in the final population.

        """
        self.generate_population()

        for generation in range(self.iterations):
            average_fitness = self.get_average_fitness()
            print("Generation {}".format(generation))
            print("Average fitness: {}".format(average_fitness))
            print("Standard deviation: {}".format(self.standard_deviation))

            # Getting the Eigen decomposition of the covariance matrix.
            self.eigen_values, self.eigen_vectors = np.linalg.eigh(self.covariance_matrix)
            self.eigen_values = np.sqrt(np.diag(self.eigen_values))

            # Calculating the new population.
            self.generate_offspring()
            self.sort_population_by_fitness()

            # Calculating the new means for the mu best individuals.
            new_means = np.array([[0.0] * self.candidate_length]).transpose()
            for idx in range(self.mu):
                weights = np.reshape(self.population_matrix[idx], (-1, 1))
                new_means = new_means + (self.weight * weights)

            # Constructing the evolution path.
            self.p_c = (1 - self.c_c) * self.p_c + np.sqrt(self.c_c * (2.0 - self.c_c) * self.mu_eff) * \
                (new_means - self.means) / self.standard_deviation

            # Rank update of the covariance matrix.
            mu_update = np.array([[0.0] * self.candidate_length] * self.candidate_length)
            for idx in range(self.mu):
                rhs = (np.reshape(self.population_matrix[idx], (-1, 1)) - self.means) / self.standard_deviation
                mu_update = mu_update + (self.weight * (rhs.dot(rhs.transpose())))

            self.covariance_matrix = (1 - self.c_cov) * self.covariance_matrix + self.c_cov / self.mu_cov * \
                (self.p_c.dot(self.p_c.transpose())) + self.c_cov * (1 - 1 / self.mu_cov) * mu_update

            # Updating the step size.
            decomposed_cov = self.eigen_vectors.dot(np.linalg.inv(self.eigen_values).dot(
                self.eigen_vectors.transpose()))
            self.p_sigma = (1 - self.c_sigma) * self.p_sigma + np.sqrt(self.c_sigma * (2 - self.c_sigma) * \
                self.mu_eff) * decomposed_cov * (new_means - self.means) / self.standard_deviation
            p_sigma_sum_squares = np.sum(np.square(self.p_sigma))
            self.standard_deviation = self.standard_deviation * np.exp(self.c_sigma / self.d_sigma * \
                (np.sqrt(p_sigma_sum_squares) / en0i(self.candidate_length) - 1))
            self.means = new_means

            print("")

        print("Completed running the CMA-ES algorithm.")
        best_candidate = self.population[0]
        for idx in range(1, self.population_size):
            # Select the candidate with the least lap time.
            if (self.population[idx].fitness < best_candidate.fitness):
                best_candidate = self.population[idx]

        print("Best fitness after running the algorithm: {}".format(best_candidate.fitness))
        return best_candidate

    def get_average_fitness(self):
        """
        Method to get the average fitness of the population.

        Returns:
            average_fitness(float): Average fitness of the population in seconds.

        """
        total_fitness = 0.0

        for candidate in self.population:
            total_fitness += self.calculate_fitness(candidate)

        return total_fitness / self.population_size

    def sort_population_by_fitness(self):
        """
        Method to sort the population in the increasing order of fitness (increasing lap times).
        """
        for i in range(self.population_size - 1):
            min_elem = i
            for j in range(i + 1, self.population_size):
                if (self.calculate_fitness(self.population[j]) < self.calculate_fitness(self.population[min_elem])):
                    min_elem = j

            if min_elem != i:
                self.population[min_elem], self.population[i] = self.population[i], self.population[min_elem]
                self.population_matrix[min_elem], self.population_matrix[i] = self.population_matrix[i], \
                    self.population_matrix[min_elem]

    def generate_empirical_covariance_matrix(self):
        """
        Method to generate the empirical covariance matrix for a population.

        Returns:
            covariance_matrix(np.array): Covariance matrix for the population.

        """
        cols = len(self.population_matrix[0])
        covariance_matrix = np.array([[0.0] * cols] * cols)
        column_sums = self.population_matrix.sum(axis=0)
        column_sums = np.reshape(column_sums, (-1, 1)).transpose()

        for row in self.population_matrix:
            mat_row = np.reshape(row, (-1, 1)).transpose()
            bracket = mat_row - (1.0 / self.population_size) * column_sums
            bracket = bracket.transpose()
            covariance_matrix = covariance_matrix + (bracket.dot(bracket.transpose()))

        self.covariance_matrix = 1.0 / (self.population_size - 1) * covariance_matrix
        return self.covariance_matrix
