import random

MAX_GEN = 100
POP_SIZE = 50
IND_LEN = 25
MUT_PROB = 0.2
XOVER_PROB = 0.8
MUT_FLIP_PROB = 1/IND_LEN

# s = [random.randint(0,100) for _ in range(IND_LEN)]

def fitness(x):
    return sum(x)

# def fitness(x):
#     a = sum(si*xi for si,xi in zip(s, x))
#     if a <= 956:
#         return a
#     return 0

def select(pop, fits):
    return random.choices(pop, fits, k=POP_SIZE)

def cross(p1, p2):
    p = random.randrange(1,IND_LEN)
    return p1[:p] + p2[p:], p2[:p] + p1[p:]

def crossover(pop):
    o = []
    for p1, p2 in zip(pop[::2], pop[1::2]):
        o1, o2 = p1, p2
        if random.random() < XOVER_PROB:
            o1, o2 = cross(p1, p2)
        o.append(o1[:])
        o.append(o2[:])
    return o

def mutate(ind):
    return [1-v if random.random() < MUT_FLIP_PROB else v for v in ind]

def mutation(pop):
    return [mutate(ind) if random.random() < MUT_PROB else ind[:] for ind in pop]

def random_individual():
    return [0 if random.random() < 0.5 else 1 for _ in range(IND_LEN)]

def random_initial_population():
    return [random_individual() for _ in range(POP_SIZE)]

def evolutionary_algorithm(pop, elitism=False):
    log = []
    for _ in range(MAX_GEN):
        fits = [fitness(ind) for ind in pop]
        log.append(max(fits)) 
        mating_pool = select(pop, fits)
        o = crossover(mating_pool)
        offspring = mutation(o)
        if elitism:
            offspring[0] = max(pop, key=fitness)
        pop = offspring[:]
    return pop, log

i1 = random_individual()
i2 = random_individual()
print('Par1:', i1)
print('Par1:', i2)

o = cross(i1, i2)
print('Cro0:', o[0])
print('Cro1:', o[1])

o2 = mutate(o[0])
print('Mut:', o2)

logs = []
for _ in range(10):
    pop = random_initial_population()
    pop, log = evolutionary_algorithm(pop)
    logs.append(log)

logs2 = []
for _ in range(10):
    pop = random_initial_population()
    pop, log = evolutionary_algorithm(pop, elitism=True)
    logs2.append(log)

import matplotlib.pyplot as plt
import numpy as np

logs = np.array(logs)

plt.plot(logs.mean(axis=0))
plt.fill_between(list(range(MAX_GEN)), np.percentile(logs, axis=0, q=25), 
                                       np.percentile(logs, axis=0, q=75), alpha=0.5)
logs = logs2
logs = np.array(logs)
plt.plot(logs.mean(axis=0))
plt.fill_between(list(range(MAX_GEN)), np.percentile(logs, axis=0, q=25), 
                                       np.percentile(logs, axis=0, q=75), alpha=0.5)
plt.show()