#######################################################################
# Copyright (C) kgmills@ualberta.ca, June 2019                        #
# Keith G. Mills, Dept. Electrical and Computer Engineering, Ualberta #
# Permission to modify so long as you keep this declaration           #
#######################################################################

import numpy as np
from math import inf

class CETuner:

    def __init__(self, agent):

        # Population size, alpha (parameter for matrices) and rho (used for determining best candidates)
        self.rho = 0.01
        self.alpha = 0.6

        # TODO make this like a config with insertable popSize
        self.popSize = max(1000, int(1 / self.rho))

        self.nElite = int(self.popSize * self.rho)

        self.agent = agent

        # Commented out, but these are the relevant weights
        # self.agent.network.fc_action.parameters()

        self.bestParams = [-inf, -inf, -inf]
        self.bestGen = 0

        # The 2 is there for the biases
        self.nFeatures = agent.network.actor_body.feature_dim * \
                         agent.network.action_dim + agent.network.action_dim

        # Weight matrix - One more column is added for the scores
        self.weightsAndScores = np.zeros((self.popSize, self.nFeatures + 1),)

        # Means and standard deviations of each weight (as well as the rewards)
        if agent.finalWeightHistory is None:
            self.Mus = np.zeros(self.nFeatures + 3,)
            self.Sigmas = np.ones(self.nFeatures + 3,) * 50

            self.vanillaWeights = agent.getFinalActorWeights()

        else:
            self.Mus = np.append(np.mean(agent.finalWeightHistory, axis = 0), [0, 0, 0])
            self.Sigmas = np.append(np.std(agent.finalWeightHistory, axis = 0), [0, 0, 0])

        self.muList = None
        self.sigList = None

    # A single generation
    def executeGeneration(self, gen):

        # Generate all the weights and dummy scores at once.
        self.weightsAndScores = np.random.normal(self.Mus, self.Sigmas,
                                                 (self.popSize, self.nFeatures + 3))

        # Run for each individual
        for i in range(self.popSize):

            if self.agent.finalWeightHistory is None:
                self.agent.setAgentFinalLayerWeights(
                    np.add(self.vanillaWeights, self.weightsAndScores[i, :-3])
                )
            else:
                self.agent.setAgentFinalLayerWeights(self.weightsAndScores[i, :-3])

            # Run for the number of episodes specified
            for j in range(self.agent.config.tunerEpisodes):
                self.agent.stepNoUpdate()
                self.agent.switch_task()
            rewardDict = self.agent.eval_episodes(log = False, tune = True)
            self.weightsAndScores[i, -3] = rewardDict['episodic_return_test']
            self.weightsAndScores[i, -2] = rewardDict['test_std']
            self.weightsAndScores[i, -1] = rewardDict['episodic_return_test'] - rewardDict['test_std']

        # https://stackoverflow.com/questions/2828059/sorting-arrays-in-numpy-by-column
        # Sort the matrix of weights and scores by the scores, in descending order.
        self.weightsAndScores = self.weightsAndScores[self.weightsAndScores[:, -1].argsort()[::-1]]

        # Select the top performing weight vectors
        eliteWeights = self.weightsAndScores[:self.nElite, :]

        # TODO fix this for early generations
        # Compute Mu and Sigma for each weight as well as the reward
        self.Mus = np.add((1 - self.alpha) * self.Mus,
                          self.alpha * np.mean(eliteWeights, axis = 0))

        self.Sigmas = np.add((1 - self.alpha) * self.Sigmas,
                             self.alpha * np.std(eliteWeights, axis = 0))

        # Check if the best is better than the global best
        if eliteWeights[0, -3] > self.bestParams[-3]:
            self.bestParams = eliteWeights[0, :]
            self.bestGen = gen

        self.recordMuSigma()

    def recordMuSigma(self):

        if self.muList is None:
            self.muList = np.asmatrix(self.weightsAndScores[:, -3]).T
            self.sigList = np.asmatrix(self.weightsAndScores[:, -2]).T

        else:
            self.muList = np.append(self.muList, np.asmatrix(
                self.weightsAndScores[:, -3]).T, axis = 1)
            self.sigList = np.append(self.sigList, np.asmatrix(
                self.weightsAndScores[:, -2]).T, axis = 1)

    def saveMats(self):

        print(self.agent.config.tag)
        fileName = "mu" + str(self.agent.config.tag) + ".csv"
        np.savetxt(fileName, self.muList, delimiter = ',')

        fileName = "sig" + str(self.agent.config.tag) + ".csv"
        np.savetxt(fileName, self.sigList, delimiter = ',')
