import numpy as np
import random
from brain import Neural_Network
from params import MUTATION_RATE, SELECTION_ALG, KWAY_TOURNAMENT_PLAYERS


def kway_selection(brains, exclude=None):
    tourn_pool = []
    best_play = None
    if exclude:
        brains = [x for x in brains if x != exclude]
    for x in range(KWAY_TOURNAMENT_PLAYERS):
        new_play = random.choice(brains)
        while new_play in tourn_pool:
            new_play = random.choice(brains)
        if not best_play or best_play.fitness < new_play.fitness:
            best_play = new_play
    return best_play


def genetic_selection(brains):
    parents_pool = []
    half_pop = int(len(brains) / 2)

    if SELECTION_ALG == "kway":
        for x in range(half_pop):
            p1 = kway_selection(brains)
            p2 = kway_selection(brains, exclude=p1)
            parents_pool.append([p1, p2])

    elif SELECTION_ALG == "roulette":
        # 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)

        # selection of pool/2 pair of parents to reproduce
        for _ in range(half_pop):
            idx1 = round(random.random() * tot_fitness - 1)
            idx2 = round(random.random() * tot_fitness - 1)
            parents_pool.append([wheel[idx1], wheel[idx2]])
    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