#/bin/python
import random

# tasks = [4, 5, 8, 2, 10, 7]
# tasks = [random.randint(1,20) for _ in range(100)]
tasks = [19, 12, 12, 3, 5, 10, 2, 17, 18, 8, 20, 6, 3, 10, 6, 1, 13, 16, 17, 1, 11, 6, 12, 2, 11, 17, 8, 12, 7, 15, 1, 5, 17, 19, 5, 16, 15, 20, 7, 4, 12, 6, 17, 6, 8, 10, 8, 13, 15, 2, 6, 16, 11, 16, 16, 16, 1, 2, 17, 9, 4, 3, 7, 4, 1, 14, 16, 1, 1, 15, 20, 2, 13, 15, 15, 11, 5, 18, 1, 14, 19, 19, 19, 19, 5, 5, 12, 5, 2, 2, 2, 9, 8, 12, 6, 20, 18, 12, 12, 19]
sol1=[1, 0, 1, 1, 0, 1]
sol2=[0, 0, 1, 0, 1, 1]

def generate_initial_population(problem_data, k, popsize):
    return [[random.randint(0, k-1) for _ in range(len(problem_data))] for _ in range(popsize)]

def fitness(solution, problem_data, k):
    return max([sum(problem_data[i] for i in range(len(solution)) if solution[i] == j) for j in range(k)])

def crossover(solutionA, solutionB):
    return [solutionA[i] if i % 2 == 0 else solutionB[i] for i in range(len(solutionA))]

def mutation(solution, k):
    i = random.randint(0,len(solution)-1)
    solution[i] = random.choice(list({i for i in range(k)} - {solution[i]}))
    return solution

def genetic(problem_data, k, popsize, mutation_chance, stop, reproduce):
    population = generate_initial_population(problem_data, k, popsize)
    population.sort(key = lambda x : fitness(x, problem_data, k), reverse=True)
    loop = 0
    while loop<stop:
        best = fitness(population[-1], problem_data, k)
        new = []
        for _ in range(reproduce):
            weightsP = [(i+1)**2 for i in range(popsize)]
            parents = random.choices(population, weights=weightsP, k=2)
            x = random.random()
            offspring = crossover(*parents)
            new.append(mutation(offspring, k) if mutation_chance>x else offspring)
        population = sorted(population + new, key=lambda x: fitness(x, problem_data, k), reverse=True)[-popsize:]

        new_best = fitness(population[-1], problem_data, k)
        loop = loop + 1 if best<=new_best else 0
        best = new_best

    return(population[-1])

ans = genetic(tasks, 4, 10, 0.7, 10, 5)
print(ans)
print(fitness(ans, tasks, 4))