# Only works with positive parameters, crossover isnt defined for the possibility of negative parameters

import numpy as np
import ionbench
import copy
from functools import lru_cache


def run(bm, x0=[], nGens=10, popSize=2500, tournementSize=5, debug=False):
    """
    Runs the genetic algorithm from Cairns et al 2017.

    Parameters
    ----------
    bm : Benchmarker
        A benchmarker to evaluate the performance of the optimisation algorithm.
    x0 : list, optional
        Initial parameter guess. Population is generated by randomly perturbing this initial guess +-50%, clamped to bounds if necessary. If x0=[] (the default), then the population will be sampled using bm.sample().
    nGens : int, optional
        The number of generations to run the optimisation algorithm for. The default is 10.
    popSize : int, optional
        The size of the population in each generation. The default is 2500.
    tournementSize : int, optional
        The size of the tournement to run during tournement selection. The default is 5.
    debug : bool, optional
        If True, debug information will be printed, reporting that status of the optimisation each generation. The default is False.

    Returns
    -------
    xbest : list
        The best parameters identified.

    """
    if not bm._bounded:
        raise RuntimeError('Cairns et al 2017 GA optimiser requires bounds.')

    boundsRange = bm.input_parameter_space(bm.ub) - bm.input_parameter_space(bm.lb)

    class Individual():
        def __init__(self):
            if len(x0) == 0:
                self.x = bm.sample()
            else:
                self.x = bm.input_parameter_space(bm.original_parameter_space(x0) * np.random.uniform(low=0.5, high=1.5, size=bm.n_parameters()))
            self.x = bm.clamp(self.x)
            self.cost = None

        def find_cost(self):
            self.cost = cost_func(tuple(self.x))

    @lru_cache(maxsize=None)
    def cost_func(x):
        return bm.cost(x)

    pop = [None] * popSize
    for i in range(popSize):
        pop[i] = Individual()
        pop[i].find_cost()

    for gen in range(nGens):
        minCost = np.inf
        for i in range(popSize):
            if pop[i].cost < minCost:
                minCost = pop[i].cost
                elite = copy.deepcopy(pop[i])
        if debug:
            print("------------")
            print(f'Gen {gen}, Best cost: {minCost}')

        # Tournement selection
        parents = []
        # Generate 2*popSize number of parents
        for j in range(popSize * 2):
            # Tournement indices
            tour = np.random.choice(np.arange(popSize), tournementSize, replace=False)
            # Find best individual in the tournement
            bestInd = pop[tour[0]]
            for i in tour[1:]:
                if pop[i].cost < bestInd.cost:
                    bestInd = pop[i]
            # Save best individual as parent
            parents.append(copy.deepcopy(bestInd))

        # Crossover
        newPop = []
        for j in range(popSize):
            # Get parents 1 and 2
            par1 = parents[2 * j]
            par2 = parents[2 * j + 1]
            # Construct child
            child = Individual()
            for i in range(bm.n_parameters()):
                # Get parameter i from parents 1 and 2
                num1 = par1.x[i]
                num2 = par2.x[i]
                # Want to convert to integer and maintain 6 decimal places of accuracy
                num1 = int(np.round(num1 * 1e6))
                num2 = int(np.round(num2 * 1e6))
                # Then map to binary string
                # Find length of binary strings
                binStrLen = int(np.log2(max(num1, num2))) + 1
                bin1 = format(num1, '0' + str(binStrLen) + 'b')
                bin2 = format(num2, '0' + str(binStrLen) + 'b')
                # Random position between 0 and binStrLen (inclusive)
                randPos = np.random.randint(0, binStrLen + 1)
                # Crossover in binary
                binNew = bin1[:randPos] + bin2[randPos:]
                # Convert back to decimal float
                numNew = int(binNew, 2) / 1e6
                child.x[i] = numNew
            newPop.append(child)

        # Mutation
        for j in range(popSize):
            for i in range(bm.n_parameters()):
                # Change of mutation is 0.1%
                if np.random.rand() < 0.001:
                    perturb = np.random.normal(0.05 * boundsRange[i], 0.05 * boundsRange[i] / 3)
                    if np.random.rand() < 0.5:
                        newPop[j].x[i] += perturb
                    else:
                        newPop[j].x[i] -= perturb
        pop = newPop
        if debug:
            print(f'Finishing gen {gen}')
        # Find costs
        for i in range(popSize):
            pop[i].find_cost()
        # Elitism
        maxCost = -np.inf
        for i in range(popSize):
            if pop[i].cost > maxCost:
                maxCost = pop[i].cost
                maxIndex = i
        pop[maxIndex] = copy.deepcopy(elite)

    minCost = np.inf
    for i in range(popSize):
        if pop[i].cost < minCost:
            minCost = pop[i].cost
            elite = pop[i]
    bm.evaluate(elite.x)
    return elite.x


def get_modification(modNum=1):
    """
    modNum = 1 -> Cairns2017

    Returns
    -------
    mod : modification
        Modification corresponding to inputted modNum. Default is modNum = 1, so Cairns2017.

    """
    mod = ionbench.modification.Cairns2017()
    return mod


if __name__ == '__main__':
    bm = ionbench.problems.staircase.HH_Benchmarker()
    mod = get_modification()
    mod.apply(bm)
    run(bm, popSize=50, debug=True, **mod.kwargs)
