#    This file is part of NEORL.

#    Copyright (c) 2021 Exelon Corporation and MIT Nuclear Science and Engineering
#    NEORL is free software: you can redistribute it and/or modify
#    it under the terms of the MIT LICENSE

#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#    SOFTWARE.

# -*- coding: utf-8 -*-
#"""
#Created on Thu Dec  3 14:42:29 2020
#
#@author: Majdi
#"""


from scipy import dot, eye, randn, asarray, array, trace, log, exp, sqrt 
from scipy import mean, sum, argsort, arange
from scipy.stats import multivariate_normal, norm
from scipy.linalg import det, expm
import joblib
import random
import numpy as np
import copy
from collections import defaultdict
from neorl.utils.seeding import set_neorl_seed

import multiprocessing
import multiprocessing.pool
class NoDaemonProcess(multiprocessing.Process):
    # make 'daemon' attribute always return False
    def _get_daemon(self):
        return False
    def _set_daemon(self, value):
        pass
    daemon = property(_get_daemon, _set_daemon)

# We sub-class multiprocessing.pool.Pool instead of multiprocessing.Pool
# because the latter is only a wrapper function, not a proper class.
class MyPool(multiprocessing.pool.Pool):
    Process = NoDaemonProcess

class XNESmod(object):
    """
    Exponential Natural Evolution Strategies
    
    :param bounds: (dict) input parameter type and lower/upper bounds in dictionary form. Example: {'x1': ['int', 1, 4], 'x2': ['float', 0.1, 0.8], 'x3': ['float', 2.2, 6.2]}
    :param fit: (function) the fitness function 
    :param npop: (int) total number of individuals in the population (default: if None, it will make an approximation, see **Notes** below)
    :param A: (np.array): initial guess of the covariance matrix A (default: identity matrix, see **Notes** below)
    :param eta_mu: (float) learning rate for updating the center of the search distribution ``mu`` (see **Notes** below)
    :param eta_sigma: (float) learning rate for updating the step size ``sigma`` (default: if None, it will make an approximation, see **Notes** below)
    :param eta_Bmat: (float) learning rate for updating the normalized transformation matrix ``B``  (default: if None, it will make an approximation, see **Notes** below)
    :param adapt_sampling: (bool): activate the adaption sampling option
    :param ncores: (int) number of parallel processors
    :param seed: (int) random seed for sampling
    """
    def __init__(self, bounds, fit, A=None, npop=None,
                 eta_mu=1.0, eta_sigma=None, eta_Bmat=None, 
                 adapt_sampling=False, ncores=1, seed=None):
        
        set_neorl_seed(seed)
            
        patience=100
        self.fitness_hom=-np.inf
        
        self.f = fit
        self.eta_mu = eta_mu
        self.use_adasam = adapt_sampling
        self.ncores = ncores
        self.bounds=bounds

        dim = len(bounds)
        A = np.eye(dim) if A is None else A
        sigma = abs(det(A))**(1.0/dim)
        bmat = A*(1.0/sigma)
        self.dim = dim
        self.sigma = sigma
        self.bmat = bmat

        # default population size and learning rates
        npop = int(4 + 3*log(dim)) if npop is None else npop
        eta_sigma = 3*(3+log(dim))*(1.0/(5*dim*sqrt(dim))) if eta_sigma is None else eta_sigma
        eta_Bmat = 3*(3+log(dim))*(1.0/(5*dim*sqrt(dim))) if eta_Bmat is None else eta_Bmat
        self.npop = npop
        self.eta_sigma = eta_sigma
        self.eta_bmat = eta_Bmat
        
        use_fshape=True
        # compute utilities if using fitness shaping
        if use_fshape:
            a = log(1+0.5*npop)
            utilities = array([max(0, a-log(k)) for k in range(1,npop+1)])
            utilities /= sum(utilities)
            utilities -= 1.0/npop           # broadcast
            utilities = utilities[::-1]  # ascending order
        else:
            utilities = None
        self.use_fshape = use_fshape
        self.utilities = utilities

        # stuff for adasam
        self.eta_sigma_init = eta_sigma
        self.sigma_old = None

        # logging
        self.fitness_best = -np.inf
        self.mu_best = None
        self.done = False
        self.counter = 0
        self.patience = patience
        self.history = {'eta_sigma':[], 'sigma':[], 'fitness':[]}
        

        # do not use options below when one individual in population is used
        if npop == 1:
            self.use_fshape = False
            self.use_adasam = False

    def init_sample(self, bounds):
    
        indv=[]
        for key in bounds:
            if bounds[key][0] == 'int':
                indv.append(random.randint(bounds[key][1], bounds[key][2]))
            elif bounds[key][0] == 'float':
                indv.append(random.uniform(bounds[key][1], bounds[key][2]))
            elif bounds[key][0] == 'grid':
                indv.append(random.sample(bounds[key][1],1)[0])
            else:
                raise Exception ('unknown data type is given, either int, float, or grid are allowed for parameter bounds')   
        return indv
    
    def ensure_bounds(self, vec, bounds):
    
        vec_new = []
        # cycle through each variable in vector 
        for i, (key, val) in enumerate(bounds.items()):
    
            # variable exceedes the minimum boundary
            if vec[i] < bounds[key][1]:
                vec_new.append(bounds[key][1])
    
            # variable exceedes the maximum boundary
            if vec[i] > bounds[key][2]:
                vec_new.append(bounds[key][2])
    
            # the variable is fine
            if bounds[key][1] <= vec[i] <= bounds[key][2]:
                vec_new.append(vec[i])
            
        return vec_new

    def evolute(self, ngen, x0=None, verbose=True):
        """
        This function evolutes the XNES algorithm for number of generations.
        
        :param ngen: (int) number of generations to evolute
        :param x0: (list) initial guess for the search (must be of same size as ``len(bounds)``)
        :param verbose: (bool) print statistics to screen
        
        :return: (dict) dictionary containing major XNES search results
        """
        f = self.f
        self.verbose=verbose
        if x0 is not None:
            assert len(x0) == self.npop, 'the length of x0 ({}) MUST equal to npop ({})'.format(len(x0), self.npop)
            self.mu=np.array(x0)
        else:
            self.mu=self.init_sample(self.bounds)
        
        mu, sigma, bmat = self.mu, self.sigma, self.bmat
        eta_mu, eta_sigma, eta_bmat = self.eta_mu, self.eta_sigma, self.eta_bmat
        npop = self.npop
        dim = self.dim
        sigma_old = self.sigma_old

        eyemat = eye(dim)

        for i in range(ngen):
            s_try = randn(npop, dim)
            z_try = mu + sigma * dot(s_try, bmat)     # broadcast
            
            for k in range (len(z_try)):
                z_try[k] = self.ensure_bounds(vec=z_try[k], bounds=self.bounds)

            p=MyPool(self.ncores)
            f_try = p.map(f, z_try)
            p.close(); p.join()   
                
            #f_try = parallel(joblib.delayed(f)(z) for z in z_try)
            f_try = asarray(f_try)
            
            # save if best
            fitness = mean(f_try)

            isort = argsort(f_try)                
            f_try = f_try[isort]
            s_try = s_try[isort]
            z_try = z_try[isort]
            
            for m in range (len(f_try)):
                if f_try[m] > self.fitness_best:
                    self.fitness_best=f_try[m]
                    self.x_best=copy.deepcopy(z_try[m])
                    
            if fitness - 1e-8 > self.fitness_best:
                self.mu_best = mu.copy()
                self.counter = 0
            else: 
                self.counter += 1
            
            u_try = self.utilities if self.use_fshape else f_try

            if self.use_adasam and sigma_old is not None:  # sigma_old must be available
                eta_sigma = self.adasam(eta_sigma, mu, sigma, bmat, sigma_old, z_try)

            dj_delta = dot(u_try, s_try)
            dj_mmat = dot(s_try.T, s_try*u_try.reshape(npop,1)) - sum(u_try)*eyemat
            dj_sigma = trace(dj_mmat)*(1.0/dim)
            dj_bmat = dj_mmat - dj_sigma*eyemat

            sigma_old = sigma

            # update
            mu += eta_mu * sigma * dot(bmat, dj_delta)
            sigma *= exp(0.5 * eta_sigma * dj_sigma)
            bmat = dot(bmat, expm(0.5 * eta_bmat * dj_bmat))

            # logging
            self.history['fitness'].append(self.fitness_best)
            self.history['sigma'].append(sigma)
            self.history['eta_sigma'].append(eta_sigma)

            # Print data
            if self.verbose and i % self.npop:
                print('^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^')
                print('XNES step {}/{}, NPOP={}, ETA_MU={}, ETA_SIGMA={}, ETA_BMAT={}, Ncores={}'.format((i+1)*self.npop, ngen*self.npop, self.npop, np.round(self.eta_mu,2), np.round(self.eta_sigma,2), np.round(self.eta_bmat,2), self.ncores))
                print('^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^')
                print('Best Swarm Fitness:', np.round(self.fitness_best,6))
                print('Best Swarm Position:', np.round(self.x_best,6))
                print('MU:', np.round(mu,3))
                print('Sigma:', np.round(sigma,3))
                print('BMAT:', np.round(bmat,3))
                print('^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^')
            
        # keep last results
        self.mu, self.sigma, self.bmat = mu, sigma, bmat
        self.eta_sigma = eta_sigma
        self.sigma_old = sigma_old
        
        if self.verbose:
            print('------------------------ NES Summary --------------------------')
            print('Best fitness (y) found:', self.fitness_best)
            print('Best individual (x) found:', self.x_best)
            print('--------------------------------------------------------------')  

        #-------------------------------------
        #return population ranked for PESA2
        #-------------------------------------
        pesa_pop=defaultdict(list)
        for i in range(len(z_try)):
            pesa_pop[i].append(list(z_try[i]))
            pesa_pop[i].append(f_try[i])
                
        return self.x_best, self.fitness_best, pesa_pop

    def adasam(self, eta_sigma, mu, sigma, bmat, sigma_old, z_try):
        #Adaptation sampling
        eta_sigma_init = self.eta_sigma_init
        dim = self.dim
        c = .1
        rho = 0.5 - 1./(3*(dim+1))  # empirical

        bbmat = dot(bmat.T, bmat)
        cov = sigma**2 * bbmat
        sigma_ = sigma * sqrt(sigma*(1./sigma_old))  # increase by 1.5
        cov_ = sigma_**2 * bbmat
        
        try:
            p0 = multivariate_normal.logpdf(z_try, mean=mu[0], cov=cov)
            p1 = multivariate_normal.logpdf(z_try, mean=mu[0], cov=cov_)
        except:
            raise Exception('--error: LinAlgError: singular matrix, turn off adapt_sampling=False to escape this error')            
        w = exp(p1-p0)

        # Mann-Whitney. It is assumed z_try was in ascending order.
        n = self.npop
        n_ = sum(w)
        u_ = sum(w * (arange(n)+0.5))

        u_mu = n*n_*0.5
        u_sigma = sqrt(n*n_*(n+n_+1)/12.)
        cum = norm.cdf(u_, loc=u_mu, scale=u_sigma)

        if cum < rho:
            return (1-c)*eta_sigma + c*eta_sigma_init
        else:
            return min(1, (1+c)*eta_sigma)