#pyright:basic
import random
import pprint

MAX_GEN = 100
POP_SIZE = 20
DIM = 25
MUT_FLIP_PROB = 1/DIM
MUT_PROB = 0.3
CROSS_PROB = 0.8

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

def create_random_population():
    return [create_random_individual() for _ in range(POP_SIZE)]

def fitness(ind):
    return sum(ind)

def select(pop, fits):
    return random.choices(pop, weights=fits, k=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 = p1[:], p2[:]
        if random.random() < CROSS_PROB:
            o1, o2 = cross(o1, o2)
        off.extend([o1, o2])
    return off

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 evolutionary_algorithm(elitism=False):
    pop = create_random_population()
    log = []
    for t 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:
            pop = offspring[1:] + [max(pop, key=fitness)]
        else:
            pop = offspring[:]
    return pop, log

# pop = create_random_population()
# pprint.pprint(pop)
# print(f"I1: {pop[0]}")
# print(f"I2: {pop[1]}")
# print(f"CO: ")
# pprint.pprint(cross(pop[0], pop[1]))
# print(f"I1: {pop[0]}")
# print(f"MU: {mutate(pop[0])}")

import numpy as np
import matplotlib.pyplot as plt 

x = list(range(MAX_GEN))

logs = []
for _ in range(100):
    pop, log = evolutionary_algorithm(elitism=False)
    logs.append(log)

logs = np.array(logs)
plt.plot(logs.mean(axis=0))
low = np.percentile(logs, axis=0, q=25)
high = np.percentile(logs, axis=0, q=75)
plt.fill_between(x, low, high, alpha=0.4)


logs = []
for _ in range(100):
    pop, log = evolutionary_algorithm(elitism=True)
    logs.append(log)

logs = np.array(logs)
plt.plot(logs.mean(axis=0))
low = np.percentile(logs, axis=0, q=25)
high = np.percentile(logs, axis=0, q=75)
plt.fill_between(x, low, high, alpha=0.4)

plt.show()