import pprint
import random

DIM = 25

MAX_GEN = 100
POP_SIZE = 50
MUT_FLIP_PROB = 1/DIM
MUT_PROB = 0.8
CROSS_PROB = 0.8

def random_individual():
    return [random.randint(0,1) for _ in range(DIM)]

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

def cross(p1, p2):
    point = random.randrange(DIM)
    o1 = p1[:point] + p2[point:]
    o2 = p2[:point] + p1[point:]
    return o1, o2

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

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

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

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

def fitness(ind):
    return sum(ind)

def evolutionary_algorithm():
    pop = random_population()
    log = []
    for g in range(MAX_GEN):
        fits = [fitness(ind) for ind in pop]
        print(min(fits), sum(fits)/POP_SIZE, max(fits))
        log.append((min(fits), sum(fits)/POP_SIZE, max(fits)))
        mating_pool = selection(pop, fits, k=POP_SIZE)
        off = crossover(mating_pool)
        off = mutation(off)
#        off[0] = max(pop, key=fitness)[:]
        pop = off

    return pop, log

logs = []
for _ in range(50):
    pop, log = evolutionary_algorithm()
    logs.append([l[2] for l in log])

logs2 = []
for _ in range(50):
    pop, log = evolutionary_algorithm()
    logs2.append([l[2] for l in log])
import numpy as np
logs = np.array(logs)
logs2 = np.array(logs2)

import matplotlib.pyplot as plt
plt.plot(np.mean(logs, axis=0))
plt.fill_between(range(MAX_GEN), 
    np.percentile(logs, 25, axis=0),
    np.percentile(logs, 75, axis=0),
    alpha=0.2)
plt.plot(np.mean(logs2, axis=0))
plt.fill_between(range(MAX_GEN), 
    np.percentile(logs2, 25, axis=0),
    np.percentile(logs2, 75, axis=0),
    alpha=0.2)
plt.show()