import random

MAX_GEN = 100
POP_SIZE = 50
IND_LEN = 25
MUT_PROB = 0.3
CROSS_PROB = 0.8
MUT_FLIP_PROB = 1/IND_LEN

def fitness(ind):
    return sum(ind)

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

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

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

def cross(p1, p2):
    cut = random.randrange(1, len(p1))
    o1 = p1[:cut] + p2[cut:]
    o2 = p2[:cut] + p1[cut:]
    return o1, o2

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

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

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

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

i1 = random_individual()
i2 = random_individual()

o1, o2 = cross(i1, i2)

o1m = mutate(o1)

print(f'{i1=}')
print(f'{i2=}')
print(f'{o1=}')
print(f'{o2=}')
print(f'{o1m=}')

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

logs2 = []
for _ in range(10):
    pop = random_initial_population()
    end_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)
logs2 = np.array(logs2)
plt.plot(logs2.mean(axis=0))
plt.fill_between(list(range(MAX_GEN)), np.percentile(logs2, axis=0, q=25), 
                                       np.percentile(logs2, axis=0, q=75), alpha=0.5)
plt.show()