import sys
import time
import random
import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


nparameters = 2   # 2 or 4

method = "all"   # ga, sa, pso, remc, simplex

# 変数の範囲
if nparameters == 4:
    bounds = [[-10, 10], [-10, 10], [-10, 10], [-10, 10]]
else:
    bounds = [[-10, 10], [-10, 10]]
#bounds = [(-20, 20), (-20, 20)]
#bounds = [(-30, 30), (-30, 30)]
#初期値 (simplex)
if nparameters == 4:
    initial_guess = [5, 5, 5, 5]
else:
    initial_guess = [5, 5]
#initial_guess = [50, 50]

# 母集団の個体数
nparents = 50

nmaxiter = 100
tol = 1.0e-3

print_level = 0


class tkParams():
    def __init__(self, **kwargs):
        self.__dict__.update(**kwargs)


def update_vars():
    global method, nparents, tol, nmaxiter, bounds, print_level
    
    nargs = len(sys.argv)
    if nargs > 1: method = sys.argv[1]
    if nargs > 2: nparents = int(sys.argv[2])
    if nargs > 3: tol      = float(sys.argv[3])
    if nargs > 4: nmaxiter = int(sys.argv[4])
    if nargs > 5: 
        bounds[0][0] = float(sys.argv[5])
        bounds[1][0] = bounds[0][0]
    if nargs > 6:
        bounds[0][1] = float(sys.argv[6])
        bounds[1][1] = bounds[0][1]
    if nargs > 7: print_level = int(sys.argv[7])


def AckleyFunc(X, Y):
# Ackley function
    t1 = 20
    t2 = -20 * np.exp(-0.2 * np.sqrt(1.0 / 2 * (X**2 + Y**2 )))
    t3 = np.e
    t4 = -np.exp(1.0 / 2 * (np.cos(2 * np.pi * X)+np.cos(2 * np.pi * Y)))

    return t1 + t2 + t3 + t4


icall = 0
def minimize_func(xk):
    global icall

    icall += 1

    if nparameters == 4:
        return AckleyFunc(xk[0], xk[1]) + 2.0 * AckleyFunc(xk[2], xk[3])
    else:
        return AckleyFunc(xk[0], xk[1])

iter = 0
def callback(xk):
    global iter

    if print_level:
        print(f"iter={iter}: fmin=", minimize_func(xk), "  xk:", xk)

    iter += 1

def usage():
    print(f"\nUsage: python {sys.argv[0]} method=[ga|pso|remc|simplex] nparents tol nmaxiter\n")


def sampling(bounds = None, x0 = None):
    if bounds is None:
        return np.array([x0])
    else:
        return np.array([random.uniform(low, high) for low, high in bounds])

#######################
# Genetic algorithm
#######################
def mutate(individual, bounds, mutation_rate = 0.01):
    for i in range(len(individual)):
        if random.random() < mutation_rate:
            individual[i] = random.uniform(bounds[i][0], bounds[i][1])

def crossover(parent1, parent2):
    crossover_point = random.randint(1, len(parent1) - 1)
    child1 = np.concatenate((parent1[:crossover_point], parent2[crossover_point:]))
    child2 = np.concatenate((parent2[:crossover_point], parent1[crossover_point:]))
    return child1, child2

def select_best_from_random_group(population, scores, k = 3):
    selected = random.choices(population, k = k)

# selectedから、scoresに対応するindexを検索してscoresに入れる
#    selected_scores = [scores[i] for i, ind in enumerate(population) if any((ind == s).all() for s in selected)]
    selected_scores = []
    selected_params = []
    for i, ind in enumerate(population):
        for s in selected:
            if (ind == s).all():
                selected_scores.append(scores[i])
                selected_params.append(s)
                break
#    if not selected_scores:  # selected_scoresが空の場合の対処
#        return random.choice(population)

    idx = np.argmin(selected_scores)
    return selected_params[idx]

def genetic_algorithm(objective_function, 
                    x0 = None, pk = None, optid = None, dx = None, kmin = None, kmax = None, bounds = None,
                    population_size = 20, mutation_rate = 0.01, 
                    callback = None, tol = 1.0e-5, nmaxiter = 1000, ):
    if bounds is None:
        bounds = []
        for i in range(len(pk)):
            if optid is not None and optid[i] == 0: continue

            _dx = 1.0 * dx[i]
            bounds.append([pk[i] - _dx, pk[i] + _dx])

    if x0 is not None:
        population = [x0]
        n = population_size - 1
    else:
        replicas = []
        n = population_size
    for i in range(n):
        new_individual = [random.uniform(low, high) for low, high in bounds]
        population.append(new_individual)

    population = np.array(population)
#    print("population=", population)
    
    best_individual = population[0]
    best_score = objective_function(best_individual)

    for generation in range(nmaxiter):
        scores = [objective_function(ind) for ind in population]
        new_population = []
        for i in range(population_size // 2):
            parent1 = select_best_from_random_group(population, scores)
            parent2 = select_best_from_random_group(population, scores)
            child1, child2 = crossover(parent1, parent2)
            mutate(child1, bounds, mutation_rate)
            mutate(child2, bounds, mutation_rate)
            new_population.extend([child1, child2])

        population = new_population
        current_best = min(population, key=objective_function)
        current_best_score = objective_function(current_best)
        if current_best_score < best_score:
            best_individual, best_score = current_best, current_best_score

        if callback: callback(best_individual)

        if current_best_score < tol: 
            res = tkParams(x = best_individual, fun = best_score, nit = generation, success = True)
            return res

    res = tkParams(x = best_individual, fun = best_score, nit = generation, success = False)
    return res


##########################
# Particle sworm optimization
##########################
class Particle:
    def __init__(self, bounds = None, x0 = None):
        if bounds is None:
            self.position = np.array(x0, dtype = float)
            n = 1
        else:
            self.position = np.array([random.uniform(low, high) for low, high in bounds], dtype = float)
            n = len(bounds)

        self.velocity = np.array([0.0 for _ in range(n)], dtype = float)
        self.best_position = self.position.copy()
        self.best_score = float('inf')

    def update_velocity(self, global_best_position, inertia, cognitive, social):
        r1, r2 = np.random.rand(2)
        cognitive_velocity = cognitive * r1 * (self.best_position - self.position)
        social_velocity = social * r2 * (global_best_position - self.position)
        self.velocity = inertia * self.velocity + cognitive_velocity + social_velocity

    def update_position(self, bounds):
        self.position += self.velocity
        for i, (low, high) in enumerate(bounds):
            if self.position[i] < low:
                self.position[i] = low
            elif self.position[i] > high:
                self.position[i] = high

def particle_swarm_optimization(objective_function, 
                    x0 = None, pk = None, optid = None, dx = None, kmin = None, kmax = None, bounds = None,
                    nparticles = 20, inertia = 0.5, cognitive = 1.5, social = 1.5,
                    callback = None, tol = 1.0e-5, nmaxiter = 1000, ):
    if bounds is None:
        n1 = nparticles // 3 - 1
        n2 = nparticles // 3
        n3 = nparticles - n1 - n2 - 1
        bounds1 = []
        bounds2 = []
        bounds3 = []
        for i in range(len(pk)):
            if optid[i] == 0: continue

            _dx = 1.0 * dx[i]
            bounds1.append([pk[i] - _dx, pk[i] + _dx])
            _dx = 10.0 * dx[i]
            bounds2.append([pk[i] - _dx, pk[i] + _dx])
            bounds3.append([kmin[i], kmax[i]])
    else:
        n1 = nparticles
        n2 = 0
        n3 = 0
        bounds1 = bounds
        bounds2 = []
        bounds3 = []
    
    if x0 is not None:
        particles  = [Particle(bounds = None, x0 = x0)]
    else:
        particles  = [Particle(bounds1)]
    particles += [Particle(bounds1) for _ in range(n1)]
    particles += [Particle(bounds2) for _ in range(n2)]
    particles += [Particle(bounds3) for _ in range(n3)]

    global_best_position = particles[0].position.copy()
    global_best_score = float('inf')

    for iter in range(nmaxiter):
        for particle in particles:
            score = objective_function(particle.position)
            if score < particle.best_score:
                particle.best_score = score
                particle.best_position = particle.position.copy()
            if score < global_best_score:
                global_best_score = score
                global_best_position = particle.position.copy()
                if callback: callback(global_best_position)

        for particle in particles:
            particle.update_velocity(global_best_position, inertia, cognitive, social)
            particle.update_position(bounds1)

        if global_best_score < tol: 
            res = tkParams(x = global_best_position, fun = global_best_score, nit = iter, success = True)
            return res

    res = tkParams(x = global_best_position, fun = global_best_score, nit = iter, success = False)
    return res

##############################
# Replica Exchange MonteCarlo
##############################
def exchange_acceptance(delta, temp1, temp2):
    return np.exp(delta * (1/temp1 - 1/temp2))

def replica_exchange_monte_carlo(objective_function,
                    x0 = None, pk = None, optid = None, dx = None, kmin = None, kmax = None, bounds = None,
                    ntemperatures = 5, REMC_Tmin = 1, REMC_Tmax = None,
                    callback = None, tol = 1.0e-5, nmaxiter = 1000):
    if REMC_Tmax is None: REMC_Tmax = ntemperatures

    if bounds is None:
        bounds1 = []
        for i in range(len(pk)):
            if optid is not None and optid[i] == 0: continue

            _dx = 1.0 * dx[i]
            bounds1.append([pk[i] - _dx, pk[i] + _dx])
    else:
        bounds1 = bounds

    temperatures = np.linspace(REMC_Tmin, REMC_Tmax, ntemperatures)
    if x0 is not None:
        replicas = [x0]
        n = ntemperatures - 1
    else:
        replicas = []
        n = ntemperatures
    for i in range(n):
        new_individual = [random.uniform(low, high) for low, high in bounds1]
        replicas.append(new_individual)

    best_individual = replicas[0]
    best_score = objective_function(best_individual)

    for iteration in range(nmaxiter):
#        print(f"iteration={iteration}")
        for i in range(ntemperatures):
#            new_individual = sampling(bounds)
            new_individual = [random.uniform(low, high) for low, high in bounds1]
            delta = objective_function(new_individual) - objective_function(replicas[i])
            if delta < 0 or np.random.rand() < np.exp(-delta / temperatures[i]):
#                print("  update:", new_individual, delta)
                replicas[i] = new_individual

        for i in range(ntemperatures - 1):
            delta = objective_function(replicas[i+1]) - objective_function(replicas[i])
            if np.random.rand() < exchange_acceptance(delta, temperatures[i], temperatures[i+1]):
#                print("  exchange:", replicas[i], replicas[i+1])
                replicas[i], replicas[i+1] = replicas[i+1], replicas[i]

        current_best = min(replicas, key = objective_function)
        current_best_score = objective_function(current_best)
        if current_best_score < tol: 
            res = tkParams(x = best_individual, fun = best_score, nit = iteration, success = True)
            return res

        if current_best_score < best_score:
            best_individual, best_score = current_best, current_best_score

        if callback: callback(best_individual)

    res = tkParams(x = best_individual, fun = best_score, nit = iteration, success = False)
    return res


##############################
# Simulated annealing
##############################
h = 1.0e-5
def gradient(xk):
    grad = []
    for i in range(len(xk)):
        _xk = xk.copy()
        _xk[i] += h
        grad.append(minimize_func(_xk - minimize_func(xk)) / h)

    return grad
    
def simulated_annealing(minimize_func, 
                    x0 = None, pk = None, optid = None, dx = None, kmin = None, kmax = None, bounds = None,
                    T = None, cooling_rate = None, step_size = None,
                    method2 = 'cg', tol2 = 0.1, nmaxiter2 = 10,
                    callback = callback, tol = tol, nmaxiter = nmaxiter, print_level = 0):
    def cool(i):
        nextT = cooling_rate * i
        return max([nextT, 1.0e-10])

    if T is None:
        if x0 is not None:
            T = minimize_func(x0)
        else:
            T = 10.0

    if cooling_rate is None:
        cooling_rate = T / (nmaxiter / 2)

    curr_x = x0.copy()
    best_x = x0.copy()
    prev_x = x0.copy()
    curr_obj = minimize_func(x0)
    best_obj = curr_obj

    for i in range(nmaxiter):
        T = cool(i)

#        if 'rand' not in method2 and 'grad' not in method2:
        if 'rand' not in method2 and 'grad' not in method2 and curr_x != prev_x:
            _res = minimize(minimize_func, curr_x, method = method2, tol = tol2, options={'maxiter': nmaxiter2})
#            print(f"  method {method2}=", _res.x, _res.fun, _res.nit)
            new_obj = _res.fun
            new_x   = _res.x
            if new_obj < best_obj:
                best_x = new_x.copy()
                best_obj = new_obj

            prev_x = curr_x.copy()

        if best_obj < tol: break

        if method2 == 'grad':
# 勾配を使用して新しい候補解を生成
            grad_x = gradient(curr_x)
            new_x = []
            if dx is not None:
                for i in range(len(curr_x)):
                    new_x.append(curr_x[i] - step_size * grad_x[i])
            else:
                for i in range(len(curr_x)):
                    new_x.append(curr_x[i] - step_size * dx[i] * grad_x[i])
        elif method2 == 'gradrand':
# 勾配を使用して新しい候補解を生成
            grad_x = gradient(curr_x)
            new_x = []
            if dx is not None:
                for i in range(len(curr_x)):
                    new_x.append(curr_x[i] - step_size * grad_x[i] + np.random.uniform(-step_size * dx[i], step_size * dx[i]))
            else:
                for i in range(len(curr_x)):
                    new_x.append(curr_x[i] - step_size * grad_x[i] + np.random.uniform(-step_size, step_size))
        else: #if method2 == 'rand':
#ランダム
            new_x = []
            if dx is not None:
                for i in range(len(x0)):
                    new_x.append(curr_x[i] + np.random.uniform(-step_size * dx[i], step_size * dx[i]))
            else:
                for i in range(len(x0)):
                    new_x.append(curr_x[i] + np.random.uniform(-step_size, step_size))

        new_obj = minimize_func(new_x)

        if new_obj < best_obj:
            if print_level:
                print(f"Update next x:", new_x, new_y, new_obj)
            curr_x = new_x.copy()
            best_x = new_x.copy()
            curr_obj = new_obj
            best_obj = new_obj
        elif T > 0.0 and np.random.rand() < np.exp((best_obj - new_obj) / T):
            if print_level:
                print(f"Try next x:", new_x, new_obj)
            curr_x = new_x.copy()
            curr_obj = new_obj
        else:
            if print_level:
                print(f"Rejected x:", new_x, new_obj)
    else:
        res = tkParams(x = best_x, fun = best_obj, nit = i, success = False)
        return res

    if 'rand' not in method2 and 'grad' not in method2 and curr_x != prev_x:
        _res = minimize(minimize_func, best_x, method = method2, tol = tol, options={'maxiter': nmaxiter})
        res = tkParams(x = _res.x, fun = _res.fun, nit = i, success = True)
    else:
        res = tkParams(x = best_x, fun = best_obj, nit = i, success = True)
    
    return res


def optimize_direct(_method):
    global iter, icall
    
    _aa = _method.split(':')
    if len(_aa) == 1:
        method = _aa[0]
        method2 = ''
    else:
        method  = _aa[0]
        method2 = _aa[1]
    
    if method2 == 'nelder-mead':
        nmaxiter2 = 1000
        tol2 = 0.1
    else:
        nmaxiter2 = 10
        tol2 = 0.1
        
    iter = 0
    icall = 0

    print()
    print(f"method: {method}:{method2}")
    start_time = time.perf_counter()
    
    if method == 'ga':
        res = genetic_algorithm(minimize_func, 
                    x0 = initial_guess, pk = initial_guess, optid = None, dx = None, kmin = None, kmax = None, bounds = bounds,
                    population_size = nparents, mutation_rate = 0.01, 
                    callback = callback, tol = tol, nmaxiter = nmaxiter)
    elif method == 'simplex':
        if nparameters == 4:
            initial_simplex = np.array([
                initial_guess,
                (bounds[0][0], initial_guess[1], initial_guess[2], initial_guess[3]),
                (initial_guess[0], bounds[1][1], initial_guess[2], initial_guess[3]),
                (initial_guess[0], initial_guess[1], bounds[2][0], initial_guess[3]),
                (initial_guess[0], initial_guess[1], initial_guess[2], bounds[3][1]),
                ])
        else:
            initial_simplex = np.array([
                initial_guess,
                (bounds[0][0], initial_guess[1]),
                (initial_guess[0], bounds[1][1]),
                ])
        res = minimize(minimize_func, initial_guess, method = 'Nelder-Mead', tol = tol,
                        callback = callback, 
                        options = {'initial_simplex': initial_simplex, "maxiter": nmaxiter})
    elif method == 'pso':
        res = particle_swarm_optimization(minimize_func, 
                    x0 = initial_guess, pk = initial_guess, optid = None, dx = None, kmin = None, kmax = None, bounds = bounds,
                    nparticles = nparents, inertia = 0.5, cognitive = 1.5, social = 1.5,
                    callback = callback, tol = tol, nmaxiter = nmaxiter)
    elif method == 'remc':
        res = replica_exchange_monte_carlo(minimize_func, 
                    x0 = initial_guess, pk = initial_guess, optid = None, dx = None, kmin = None, kmax = None, bounds = bounds,
                    ntemperatures = nparents, REMC_Tmin = 1, REMC_Tmax = 10,
                    callback = callback, tol = tol, nmaxiter = nmaxiter)
    elif method == 'sa':
        res = simulated_annealing(minimize_func, 
                    x0 = initial_guess, pk = initial_guess, optid = None, dx = None, kmin = None, kmax = None, bounds = bounds,
                    T = None, cooling_rate = None, step_size = 1.0,
                    method2 = method2, nmaxiter2 = nmaxiter2, tol2 = tol2,
                    callback = callback, tol = tol, nmaxiter = nmaxiter)
    else:
        print(f"Invalide method [{method}]")
        exit()

    end_time = time.perf_counter()

    if res.success:
        print(f"  Converged at iter={res.nit}, icall={icall}")
    else:
        print(f"  Not converged in {nmaxiter} iterations, icall={icall}")
    print(f"    Elapsed time: {end_time - start_time} sec")
    print(f"    Best parameters: {res.x}")
    print(f"    Best score: {res.fun}")
    print()

def plot():
    x = np.linspace(-5, 5, 400)
    y = np.linspace(-5, 5, 400)
    x, y = np.meshgrid(x, y)    

    z = minimize_func([x, y])

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.plot_surface(x, y, z, cmap='viridis')

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.pause(0.01)


def main():
    update_vars()
    
    print()
    print(f"method: {method}")
    print(f"bounds:", bounds)
    print(f"initial guess:", initial_guess)
    print(f"nmaxiter: {nmaxiter}")

    if method == 'all':
        for m in ['ga', 'pso', 'remc', 'simplex', 'sa:rand', 'sa:grad', 'sa:gradrand', 'sa:cg', 'sa:bfgs']:
            optimize_direct(m)
    else:
       optimize_direct(method)


if __name__ == '__main__':
    if nparameters == 2: plot()

    main()
    
    usage()

    input("\nPress ENTER to terminate>>\n")
    
    