import random
from itertools import permutations

# genetic algorithm search of the one max optimization problem
from numpy.random import randint
from numpy.random import rand

# this is helper function for sum_distance function, it counts the taxi cab distance between 2 vectors [x,y]
def distance(x,y):
    temp1 = abs(x[0]-y[0])
    temp2 = abs(x[1]-y[1])
    vector_distance = temp1+temp2
    return vector_distance

# this is fitting function which tells how well specimen fits the environment
# this function counts the sum of distances between vectors for a specimen
# this was just for testing, it should be probably changed to A*
def sum_distance(speciment):
    sum = 0
    for i in range(0,len(speciment)-2):
        pom = distance(speciment[i],speciment[i+1])
        sum = sum + pom
    return sum


# tournament selection
# this function randomly puts speciments to groups, from each group one best speciment is taken for reproduction
def selection(pop, scores, k=3):
    # first random selection
    selection_ix = randint(len(pop))
    for ix in randint(0, len(pop), k- 1):
        # check if better (e.g. perform a tournament)
        if scores[ix] < scores[selection_ix]:
            selection_ix = ix
    return pop[selection_ix]


# crossover two parents to create two children
# this function creates speciments for new generation
def crossover(p1, p2, r_cross):
    p1=list(p1)
    p2=list(p2)
    # children are copies of parents by default
    c1, c2 = p1.copy(), p2.copy()
    # check for recombination
    if rand() < r_cross:
        # select crossover point that is not on the end of the string
        pt = randint(1, len(p1) - 2)
        # perform crossover
        temp1 = p1[:pt]
        for mine in p2:
            if mine not in p1[:pt]:
                temp1.append(mine)
        temp2 = p2[:pt]
        for mine in p1:
            if mine not in p2[:pt]:
                temp2.append(mine)
        c1 = temp1
        c2 = temp2
    return [c1, c2]


# mutation operator
# this function checks whether genes mutated
# if gene is mutated then it is swapped with randomly chosen gene from the same speciment
def mutation(speciment, r_mut):
    for i in range(len(speciment)-1):
        # check for a mutation
        if rand() < r_mut:
            # flip the bit
            temp = speciment[i]
            pom = random.randint(0,len(speciment)-1)
            speciment[i] = speciment[pom]
            speciment[pom] = temp



# genetic algorithm
def genetic_algorithm(objective, n_iter, n_pop, r_cross, r_mut):
    # this is hardcoded list of coordinates of all mines (for tests only) which represents one speciment in population
    # it is then permutated to get set number of species and create population
    speciment=[[1,1],[5,2],[1,3],[2,5],[2,1],[3,2],[3,4],[5,0]]
    pop = [random.choice(list(permutations(speciment,len(speciment)))) for _ in range(n_pop)]
    # permutation function returns tuples so I change them to lists
    for i in range(len(pop)):
        pop[i] = list(pop[i])

    # keep track of best solution
    best, best_eval = 0, objective(pop[0])

    # enumerate generations
    for gen in range(n_iter):
        # evaluate all candidates in the population
        scores = [objective(c) for c in pop]
        # check for new best solution
        for i in range(n_pop):
            if scores[i] < best_eval:
                best, best_eval = pop[i], scores[i]
                print(">%d, new best f(%s) = %.3f" % (gen, pop[i], scores[i]))
        # select parents
        selected = [selection(pop, scores) for _ in range(n_pop)]
        # create the next generation
        children = list()
        for i in range(0, n_pop, 2):
            # get selected parents in pairs
            p1, p2 = selected[i], selected[i + 1]
            # crossover and mutation
            for c in crossover(p1, p2, r_cross):
                # mutation
                mutation(c, r_mut)
                # store for next generation
                children.append(c)
        # replace population
        pop = children
    return [best, best_eval]


# define the total iterations
n_iter = 100
# bits
n_bits = 20
# define the population size
n_pop = 100
# crossover rate
r_cross = 0.9
# mutation rate
r_mut = 0.05
# perform the genetic algorithm search
best, score = genetic_algorithm(sum_distance, n_iter, n_pop, r_cross, r_mut)
print('Done!')
print('f(%s) = %f' % (best, score))