#/bin/python
import random

tasks = [4, 5, 8, 2, 10, 7]
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):
            parents = random.choices(population, cum_weights=[1]*popsize, 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)[-4:]

        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, 2, 4, 0.2, 10, 4)
print(ans)
print(fitness(ans, tasks, 2))