import numpy as np;

from .covariance_target import CovarianceTarget;

class AMaLGaMCovariance(CovarianceTarget):
    _API=2.
    def __init__(self,
        theta_SDR = 1.,
        eta_DEC = 0.9,
        alpha_Sigma = [-1.1,1.2,1.6],
        NIS_MAX = 25,
        tau = 0.35,
        epsilon = 1e-30,
        condition_number_epsilon = 1e6):

        self.epsilon = epsilon;
        self.theta_SDR = theta_SDR;
        self.eta_DEC = eta_DEC;
        self.eta_INC = 1./eta_DEC;
        self.NIS_MAX = NIS_MAX;
        self.alpha_Sigma = alpha_Sigma;
        self.tau = tau;
        self.condition_number_epsilon = condition_number_epsilon;

    def _reset(self, initial_mean, initial_covariance):
        self.mean = initial_mean;
        self.old_mean = initial_mean;

        self.covariance = initial_covariance;

        self.d = len(initial_mean);

        self.Sigma = initial_covariance;
        self.c_multiplier = 1.;
        self.NIS = 0;
        self.t = 0;

        self.best_f = -float('inf');


    def _update_mean(self, mean):
        self.old_mean = self.mean;
        self.mean = mean;

    def _calculate(self, population):
        self.update_matrix(population);
        self.update_multiplier(population);
        self.t += 1;
        self.best_f = max(self.best_f, np.max(population.fitness));

        new_covariance = self.Sigma*self.c_multiplier;

        u,s,_ = np.linalg.svd(new_covariance);
        s_max  = np.max(s)
        s_max  = np.clip(s_max, self.epsilon*self.condition_number_epsilon, 1e3);
        s = np.clip(s, s_max/self.condition_number_epsilon, s_max);
        new_covariance = u*s@u.T

        self.covariance = new_covariance

        return self.covariance;

    def update_matrix(self, population):
        F = population.fitness;
        sel_idx = F.argsort()[-np.ceil(self.tau*len(population)).astype(int):][::-1]

        alpha = self.alpha_Sigma;
        eta_Sigma = 1.-np.exp(alpha[0]*len(sel_idx)**alpha[1]/self.d**alpha[2]);

        current_update = np.zeros((self.d,self.d));
        selection = population.population[sel_idx];
        for individual in selection:
            delta = individual-self.old_mean;
            current_update += np.outer(delta,delta)
        current_update /= (selection.shape[0]);

        self.Sigma *= (1-eta_Sigma);
        self.Sigma += eta_Sigma*current_update;

        # We need to ensure the condition number is OK to avoid singular matrix.
        u,s,_ = np.linalg.svd(self.Sigma);
        s_max  = np.max(s)
        s_max  = np.clip(s_max, self.epsilon*self.condition_number_epsilon, None);
        s = np.clip(s, s_max/self.condition_number_epsilon, s_max);
        self.Sigma = u*s@u.T

    def update_multiplier(self, population):
        if np.any(population.fitness>self.best_f):
            self.NIS = 0;
            self.c_multiplier = max(1., self.c_multiplier);
            self.SDR(population);

        else:
            if self.c_multiplier <= 1:
                self.NIS += 1;
            if self.c_multiplier > 1 or self.NIS >= self.NIS_MAX:
                self.c_multiplier *= self.eta_DEC;
            if self.c_multiplier < 1 and self.NIS < self.NIS_MAX:
                self.c_multiplier = 1;

    def SDR(self, population):
        x_avg = np.mean(population.population[population.fitness>self.best_f], axis=0);
        delta = np.abs(self.mean-x_avg);
        variances = np.abs(np.diag(self.covariance));
        if np.any(delta/np.sqrt(variances)>self.theta_SDR):
            self.c_multiplier *= self.eta_INC;






    def _calculate_deterministic(self,population):
        return self._calculate(population);
    def _terminating(self, population):
        pass;

    @classmethod
    def _get_kwargs(cls, config, key = ""):
        cls._config_required(
            'theta_SDR',
            'eta_DEC',
            'alpha_Sigma',
            'NIS_MAX',
            'tau',
            'epsilon',
            'condition_number_epsilon'
        )
        cls._config_defaults(
            theta_SDR = 1.,
            eta_DEC = 0.9,
            alpha_Sigma = [-1.1,1.2,1.6],
            NIS_MAX = 25,
            tau = 0.35,
            epsilon = 1e-30,
            condition_number_epsilon = 1e6
        )

        return super()._get_kwargs(config, key = key);

[((53, 16, 53, 45), 'numpy.linalg.svd', 'np.linalg.svd', ({(53, 30, 53, 44): 'new_covariance'}, {}), '(new_covariance)', True, 'import numpy as np\n'), ((54, 17, 54, 26), 'numpy.max', 'np.max', ({(54, 24, 54, 25): 's'}, {}), '(s)', True, 'import numpy as np\n'), ((55, 17, 55, 80), 'numpy.clip', 'np.clip', ({(55, 25, 55, 30): 's_max', (55, 32, 55, 74): 'self.epsilon * self.condition_number_epsilon', (55, 76, 55, 79): '1000.0'}, {}), '(s_max, self.epsilon * self.condition_number_epsilon, 1000.0)', True, 'import numpy as np\n'), ((56, 12, 56, 66), 'numpy.clip', 'np.clip', ({(56, 20, 56, 21): 's', (56, 23, 56, 58): 's_max / self.condition_number_epsilon', (56, 60, 56, 65): 's_max'}, {}), '(s, s_max / self.condition_number_epsilon, s_max)', True, 'import numpy as np\n'), ((70, 25, 70, 50), 'numpy.zeros', 'np.zeros', ({(70, 34, 70, 49): '(self.d, self.d)'}, {}), '((self.d, self.d))', True, 'import numpy as np\n'), ((81, 16, 81, 41), 'numpy.linalg.svd', 'np.linalg.svd', ({(81, 30, 81, 40): 'self.Sigma'}, {}), '(self.Sigma)', True, 'import numpy as np\n'), ((82, 17, 82, 26), 'numpy.max', 'np.max', ({(82, 24, 82, 25): 's'}, {}), '(s)', True, 'import numpy as np\n'), ((83, 17, 83, 81), 'numpy.clip', 'np.clip', ({(83, 25, 83, 30): 's_max', (83, 32, 83, 74): 'self.epsilon * self.condition_number_epsilon', (83, 76, 83, 80): 'None'}, {}), '(s_max, self.epsilon * self.condition_number_epsilon, None)', True, 'import numpy as np\n'), ((84, 12, 84, 66), 'numpy.clip', 'np.clip', ({(84, 20, 84, 21): 's', (84, 23, 84, 58): 's_max / self.condition_number_epsilon', (84, 60, 84, 65): 's_max'}, {}), '(s, s_max / self.condition_number_epsilon, s_max)', True, 'import numpy as np\n'), ((88, 11, 88, 49), 'numpy.any', 'np.any', ({(88, 18, 88, 48): '(population.fitness > self.best_f)'}, {}), '(population.fitness > self.best_f)', True, 'import numpy as np\n'), ((102, 16, 102, 86), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((103, 16, 103, 39), 'numpy.abs', 'np.abs', ({(103, 23, 103, 38): 'self.mean - x_avg'}, {}), '(self.mean - x_avg)', True, 'import numpy as np\n'), ((49, 39, 49, 65), 'numpy.max', 'np.max', ({(49, 46, 49, 64): 'population.fitness'}, {}), '(population.fitness)', True, 'import numpy as np\n'), ((74, 30, 74, 51), 'numpy.outer', 'np.outer', ({(74, 39, 74, 44): 'delta', (74, 45, 74, 50): 'delta'}, {}), '(delta, delta)', True, 'import numpy as np\n'), ((104, 27, 104, 51), 'numpy.diag', 'np.diag', ({(104, 35, 104, 50): 'self.covariance'}, {}), '(self.covariance)', True, 'import numpy as np\n'), ((105, 24, 105, 42), 'numpy.sqrt', 'np.sqrt', ({(105, 32, 105, 41): 'variances'}, {}), '(variances)', True, 'import numpy as np\n')]