#!/usr/bin/env python

import os
import numpy as np
from mtuq.grid_search import grid_search
from mtuq.grid import Normal_Samples_MT
import cartopy
from mtuq.graphics import Hammer
import matplotlib.pyplot as plt
from mtuq.util.math import to_delta_gamma

def linear_transform(i, a, b):
    transformed = a + (b-a) * i / 10
    return transformed

def inverse_linear_transform(transformed, a, b):
    i = (10*(transformed-a))/(b-a)
    return i

class MT_CMA_ES(object):
    #
    def __init__(self, data_bw, greens_bw, misfit_bw, data_sw, greens_sw, misfit_sw, origin, n=6, sigma=2, lmbda=None, a=4, b=5, start=0, stop=None):

        #
        # Objective function definition
        #
        self.data_bw = data_bw
        self.greens_bw = greens_bw
        self.misfit_bw = misfit_bw
        self.data_sw = data_sw
        self.greens_sw = greens_sw
        self.misfit_sw = misfit_sw
        self.origin = origin

        self.n = n
        self.sigma = sigma
        if lmbda is not None:
            self.lmbda = lmbda
        else:
            self.lmbda = int(4 + np.floor(3*np.log(n)))

        self.a = a
        self.b = b
        self.start = start
        self.stop = stop

        self.counteval = 0
        # what is the length along each axis?

        self.xmean = np.random.randn(self.n,1)*0.1 + 5
        self.xmean = np.random.uniform(0,10, (6,1))


        xtart = self.xmean
        self.stopeval = 1e3*self.n**2*10

        # xmean[2] += 1e5

        self.mu = np.floor(self.lmbda/2)
        a = 1 # Original author uses 1/2 in tutorial and 1 in publication
        self.weights = np.array([np.log(self.mu+a) - np.log(np.arange(1, self.mu+1))]).T
        self.weights /= sum(self.weights)
        self.mueff = sum(self.weights)**2/sum(self.weights**2)

        # Step-size control
        self.cs = (self.mueff + 2) / (self.n + self.mueff + 5)
        self.damps = 1 + 2 * max(0, np.sqrt((self.mueff - 1)/(self.n + 1)) - 1) + self.cs

        # Covariance matrix adaptation
        self.cc = (4 + self.mueff / self.n)/(self.n + 4 + 2 * self.mueff / self.n)
        self.acov = 2
        self.c1 = self.acov / ((self.n + 1.3)**2 + self.mueff)
        self.cmu = min(1 - self.c1, self.acov * (self.mueff - 2 + 1 / self.mueff) / ((self.n + 2)**2 + self.acov*self.mueff/2))

        # INITIALIZATION
        self.ps = np.zeros_like(self.xmean)
        self.pc = np.zeros_like(self.xmean)
        self.B = np.eye(self.n,self.n)
        self.D = np.ones((self.n, 1))
        self.C = self.B @ np.diag(self.D[:,0]**2) @ self.B.T
        self.invsqrtC = self.B @ np.diag(self.D[:,0]**-1) @ self.B.T
        self.eigeneval = 0
        self.chin = self.n**0.5 * (1 - 1 / ( 4 * self.n) + 1 / (21 * self.n**2))
        self.mutants = np.zeros((self.n, self.lmbda))

        self.result_list = []
        self.mutant_list = []

    def rotate_kappa(self, i):
        deg = linear_transform(i, 0, 360)
        deg = deg%360
        i = inverse_linear_transform(deg, 0, 360)
        return(i)

    def rotate_sigma(self, i):
        deg = linear_transform(i, -180, 180) + 180
        deg = deg%360 - 180
        i = inverse_linear_transform(deg, -180, 180)
        return(i)

    def circular_mean(self, a, weights):
        a = linear_transform(a, 0, 360)-180
        mean = np.rad2deg(np.arctan2(np.sum(np.sin(np.deg2rad(a[range(int(self.mu))]))*self.weights.T), np.sum(np.cos(np.deg2rad(a[range(int(self.mu))]))*self.weights.T)))+180
        mean = inverse_linear_transform(mean, 0, 360)
        return mean

    def mean_diff(self, new, old):
        diff = new-old
        angular_diff = linear_transform(self.xmean[2], 0, 360)-linear_transform(self.xold[2], 0, 360)
        angular_diff = inverse_linear_transform((angular_diff+180)%360 - 180, 0, 360)
        diff[2] = angular_diff
        return diff

    def draw_muants(self):
        # Darw lambda mutants
        for i in range(self.lmbda):
            mutant = self.xmean + self.sigma * self.B @ (self.D * np.random.randn(self.n,1))
            while not (mutant[0] > 0 and mutant[0] < 10):
                mutant[0] = (self.xmean + self.sigma * self.B @ (self.D * np.random.randn(self.n,1)))[0]
            while not (mutant[1] > 0 and mutant[1] < 10):
                mutant[1] = (self.xmean + self.sigma * self.B @ (self.D * np.random.randn(self.n,1)))[1]
            while not (mutant[3] > 0 and mutant[3] < 10):
                mutant[3] = (self.xmean + self.sigma * self.B @ (self.D * np.random.randn(self.n,1)))[3]
            while not (mutant[4] > 0 and mutant[4] < 10):
                mutant[4] = (self.xmean + self.sigma * self.B @ (self.D * np.random.randn(self.n,1)))[4]
            while not (mutant[5] > 0 and mutant[5] < 10):
                mutant[5] = (self.xmean + self.sigma * self.B @ (self.D * np.random.randn(self.n,1)))[5]


            self.mutants[:,i] = mutant.T

        self.mutants[2] = self.rotate_kappa(self.mutants[2])
        self.counteval += self.lmbda

    def eval_norm(self):
        # Evaluate mutant fitness by calling grid_search and a custom mutant grid
        self.grid_bw = Normal_Samples_MT(self.mutants, self.a, self.b)
        results_bw = grid_search(self.data_bw, self.greens_bw, self.misfit_bw, self.origin, self.grid_bw)
        self.grid_sw = Normal_Samples_MT(self.mutants, self.a, self.b)
        results_sw = grid_search(self.data_sw, self.greens_sw, self.misfit_sw, self.origin, self.grid_sw)
        results = results_bw + results_sw
        self.result_list.append(results)
        self.mutant_list.append(self.mutants)
        self.arfitness = results.values

    def fitness_sort(self):
        # Sort by fitness
        self.mutants = self.mutants[:,np.argsort(self.arfitness.T)[0]]

    def update_mean(self):
        # Update the mean
        self.xold = self.xmean
        self.xmean = np.array([np.sum(self.mutants[:,range(int(self.mu))]*self.weights.T, axis=1)]).T
        self.xmean[2] = self.circular_mean(self.mutants[2], self.weights)

    def update_step_size(self):
        # Step size control
        self.ps = (1-self.cs)*self.ps + np.sqrt(self.cs*(2-self.cs)*self.mueff) * self.invsqrtC @ (self.mean_diff(self.xmean, self.xold) / self.sigma)

    def update_covariance(self):
        # Covariance matrix adaptation
        if np.linalg.norm(self.ps)/np.sqrt(1-(1-self.cs)**(2*self.counteval/self.lmbda))/self.chin < 1.4 + 2/(self.n+1):
            self.hsig = 1
        if np.linalg.norm(self.ps)/np.sqrt(1-(1-self.cs)**(2*self.counteval/self.lmbda))/self.chin >= 1.4 + 2/(self.n+1):
            self.hsig = 0
        self.pc = (1-self.cc) * self.pc + self.hsig * np.sqrt(self.cc*(2-self.cc)*self.mueff) * self.mean_diff(self.xmean, self.xold) / self.sigma

        artmp = (1/self.sigma) * self.mean_diff(self.mutants[:,0:int(self.mu)], self.xold)
        self.C = (1-self.c1-self.cmu) * self.C + self.c1 * (self.pc@self.pc.T + (1-self.hsig) * self.cc*(2-self.cc) * self.C) + self.cmu * artmp @ np.diag(self.weights.T[0]) @ artmp.T
        # Adapt step size
        self.sigma = self.sigma * np.exp((self.cs/self.damps)*(np.linalg.norm(self.ps)/self.chin - 1))

        if self.counteval - self.eigeneval > self.lmbda/(self.c1+self.cmu)/self.n/10:
            self.eigeneval = self.counteval
            self.C = np.triu(self.C) + np.triu(self.C,1).T
            self.D,self.B = np.linalg.eig(self.C)
            self.D = np.array([self.D]).T
            self.D = np.sqrt(self.D)
            self.invsqrtC = self.B @ np.diag(self.D[:,0]**-1) @ self.B.T

    def rslt(self):
        return Normal_Samples_MT(self.xmean, self.a, self.b)

    def plot_lune(self, ax):
        fig = plt.figure(figsize=(5,10))
        ax = plt.axes(projection=Hammer())
        coords = Normal_Samples_MT(np.hstack(self.mutant_list), self.a, self.b)
        reslist = self.result_list[0]
        for i in range(1, len(self.result_list)):
            reslist = reslist.append(self.result_list[i])
        y,x = to_delta_gamma(coords.coords[1], coords.coords[2])
        ax.scatter(x,y,s=5,c=reslist.values, transform = cartopy.crs.PlateCarree(), vmax=reslist.max()/1.5,alpha=0.5)
        coord = np.vstack((x,y)).T
        x = np.rad2deg(np.hstack((np.linspace(-1,1,100), np.ones(100), np.linspace(1,-1,100), np.ones(100)*-1, np.linspace(-1,1,100)[0]))*np.pi/6)
        ax.set_boundary(bound, transform=ccrs.PlateCarree())
        ax.gridlines()
        ax.grid()
        plt.show()
