import numpy as np
import random
from brain import Neural_Network
from params import MUTATION_RATE

def genetic_selection(brains):
    # tot_fitness = sum ([int(b.fitness) for b in brains])

    # does not seem very optimized... TBR
    # constitute a list where every brain is represented
    # proportionnally to its relative fitness
    wheel = []
    for b in brains :
        wheel += [b] * b.fitness

    
    tot_fitness = len(wheel)
    half_pop = int(len(brains)/2)

    # selection of pool/2 pair of parents to reproduce  
    parents_pool = []
    for _ in range(half_pop):
        parents_pool.append([
                wheel[round(random.random()*tot_fitness)], 
                wheel[round(random.random()*tot_fitness)]
                ])
    return parents_pool
    

def cross_mutate_genes(p1_gene, p2_gene):
    child = []
    p1_gene = list(p1_gene)
    p2_gene = list(p1_gene)
    for idx,x in enumerate(p2_gene):
        if random.random() > 0.5 :
            choice = p1_gene[idx]
        else :
            choice = p2_gene[idx]
        # Mutation
        if random.random() < MUTATION_RATE :
            choice[random.randint(0, len(choice) - 1)] = random.random()
            print("Mutation !")
        child.append(choice)
    return np.array(child)


def genetic_reproduction(parents_pool):
    # every pair of parents will produce a mixed child
    new_pop = []
    for [p1,p2] in parents_pool:
        W1_kid = cross_mutate_genes(p1.W1, p2.W1)
        W2_kid = cross_mutate_genes(p1.W2, p2.W2)
        c_brain1 = Neural_Network(W1=W1_kid, W2=W2_kid)
        c_brain2 = Neural_Network(W1=W1_kid, W2=W2_kid)
        new_pop.append(c_brain1)
        new_pop.append(c_brain2)
    return new_pop