import numpy as np


class GeneticAlgorithm():

    def __init__(self, chromosomeShape,
                 errorFunction,
                 elitism=1,
                 populationSize=25,
                 mutationProbability=.1,
                 mutationScale=.5,
                 numIterations=10000,
                 errorTreshold=1e-6,
                 ):

        self.populationSize = populationSize
        self.p = mutationProbability
        self.numIter = numIterations
        self.e = errorTreshold
        self.f = errorFunction
        self.keep = elitism
        self.k = mutationScale
        self.i = -1

        self.population = []
        for _ in range(populationSize):
            chromosome = np.random.randn(chromosomeShape)

            fitness = self.calculateFitness(chromosome)
            self.population.append((chromosome, fitness))

        self.population = sorted(self.population, key=lambda t: -t[-1])

    def run(self, NN):

        done = False
        while not done:
            done, iteration, best = self.step()
            if iteration % 100 == 0:
                print(iteration, 1/best[-1])
            NN.set_weights(best[0])
        NN.set_weights(best[0])

    def step(self):

        self.i += 1

        bestUnits = self.bestN(self.keep)
        isFinished = False
        newPopulation = []
        for bu in bestUnits:
            newPopulation.append(bu)

        while not len(newPopulation) == len(self.population):
            parents = self.selectParents()
            child = self.crossover(parents[0], parents[1])
            child = self.mutate(child)
            #child = self.mutate1(child)
            newPopulation.append((child, self.calculateFitness(child)))

        self.population = newPopulation

        self.population = sorted(self.population, key=lambda t: -t[-1])

        bestUnit = self.population[0]

        fitness = bestUnit[1]

        if self.i == self.numIter or fitness <= self.e:
            isFinished = True

        return (isFinished, self.i, bestUnit)

    def calculateFitness(self, chromosome):

        chromosomeError = self.f(chromosome)

        return 1 / chromosomeError

    def bestN(self, n):

        return self.population[:n]

    def best(self):

        return self.population[0]

    def selectParents(self):

        sumFit = 0
        prob = 0
        parents = []
        probabilities = []
        for (unit, fit) in self.population:
            sumFit += fit
        for (unit, fit) in self.population:
            probabilities.append((unit, prob, prob + fit / sumFit))
            prob += fit / sumFit

        random1 = np.random.uniform()
        random2 = np.random.uniform()
        for probab in probabilities:
            if (random1 > probab[1] and random1 < probab[2]):
                parents.append(probab[0])
            if (random2 > probab[1] and random2 < probab[2]):
                parents.append(probab[0])
            if len(parents) == 2:
                break
        return parents

    def crossover(self, p1, p2):

        #child = (p1 + p2) / 2

        n = np.random.uniform(0,1)

        child = n*p1 + (1-n)*p2

        return child

    def mutate(self, chromosome):

        rand = np.random.uniform()

        if rand < 0.33:
            for i in range(0, chromosome.shape[0]):
                prob = np.random.uniform(0, 1)
                if prob < self.p:
                    gauss = np.random.normal(0, self.k)
                    chromosome[i] += gauss
            return chromosome

        elif 0.33 < rand < 0.66:
            for i in range(0, chromosome.shape[0]):
                prob = np.random.uniform(0, 1)
                if prob < self.p:
                    gauss = np.random.normal(0, self.k)
                    chromosome[i] += gauss
            return chromosome

        elif 0.66 < rand < 1:
            for i in range(0, chromosome.shape[0]):
                prob = np.random.uniform(0, 1)
                if prob < self.p:
                    gauss = np.random.normal(0, self.k)
                    chromosome[i] = gauss
            return chromosome

    def mutate1(self, chromosome):
        for i in range(0, chromosome.shape[0]):
            rand = np.random.uniform(0, 1)
            if rand < self.p:
                rand = np.random.normal(0, self.k)
                chromosome[i] += rand
        return chromosome