import random
from datetime import datetime
from copy import deepcopy

import neural_network as nn


# Function that will return True with a probability of (probability * 100)%
def chance(probability):
    return random.random() < probability


class EvolutionOptions:
    def __init__(self):
        self.activation_functions = list()

        self.gene_mutation_probability = 0
        self.weight_perturbation_probability = 0
        self.connection_disable_probability = 0
        self.weight_perturbation_max_delta = 0
        self.weight_random_mutation_range = 0
        self.node_insertion_chance = 0
        self.new_connection_chance = 0


# Brains have genotypes, and phenotypes(networks) generated by genotypes
class Brain:

    def __init__(self, input_keys=list(), output_nodes=dict(), evolution_options=EvolutionOptions()):
        self.birth_date = str(datetime.now())

        self.input_keys = input_keys
        self.output_nodes = output_nodes
        self.evolution_options = evolution_options

        self.genotype = dict()
        self.global_innov_counter = 0
        self.node_innov_counter = 0

        self.network = None

        # (Optional) To be set after instantiation
        self.allow_recurrency = False

    def randomize_gene(self, gene):
        gene_type = gene["type"]
        if gene_type == "connection":
            w = random.random() * self.evolution_options.weight_random_mutation_range
            gene["weight"] = random.choice([w, -w])
        elif gene_type == "node":
            gene["activation_function"] = random.choice(self.evolution_options.activation_functions)

    def mutate_genotype(self):
        enabled_genes = [gene for gene in self.genotype.values() if gene["enable"] is True]
        node_genes = [gene for gene in enabled_genes if gene["type"] == "node"]
        connection_genes = [gene for gene in enabled_genes if gene["type"] == "connection"]

        for node in node_genes:
            if chance(self.evolution_options.gene_mutation_probability):
                self.randomize_gene(gene=node)

        for connection in connection_genes:
            if chance(self.evolution_options.gene_mutation_probability):
                if chance(self.evolution_options.weight_perturbation_probability):
                    # Perturbate weight
                    max_value = self.evolution_options.weight_perturbation_max_delta
                    perturbation_delta = random.uniform(0, max_value)

                    connection["weight"] += random.choice([perturbation_delta, -perturbation_delta])
                else:
                    self.randomize_gene(gene=connection)

            if chance(self.evolution_options.connection_disable_probability):
                connection["enable"] = False

    def randomize_genotype(self):
        new_genotype = dict()

        # Start minimally connected
        for o in self.output_nodes.keys():
            i = random.choice(self.input_keys)
            new_genotype[self.global_innov_counter] = {
                "type": "connection",
                "conn": [i, o],
                "weight": 0,
                "enable": True
            }
            self.global_innov_counter += 1

        # Assign the randomized genotype to the phenotype
        self.genotype = new_genotype

        # Randomize initial gene weights
        self.mutate_genotype()

    def disable_connections(self, connections):
        connection_genes = [gene for gene in self.genotype.values() if gene["type"] == "connection"]

        for c in connections:
            for g in connection_genes:
                if g['conn'] == c:
                    g['enable'] = False
                    break

    def build_network(self):
        self.network = nn.Network()

        # Add the input and output nodes first
        for k in self.input_keys:
            self.network.add_input_node(
                node_id=k,
            )

        for key, function in zip(self.output_nodes.keys(), self.output_nodes.values()):
            self.network.add_output_node(
                node_id=key,
                activation_function=function
            )

        for gene in self.genotype.values():
            if gene["enable"]:
                gene_type = gene["type"]
                if gene_type == "connection":
                    # Both the input and output node need to exist in order to link them
                    in_key = gene["conn"][0]
                    out_key = gene["conn"][1]

                    # The connection won't be created if any node on either side is missing.
                    # This can happen if that node gene was disabled, or missing on the genotype.
                    if in_key in self.network.nodes and out_key in self.network.nodes:
                        self.network.add_connection(
                            weight=gene["weight"],
                            in_node_key=in_key,
                            out_node_key=out_key
                        )
                elif gene_type == "node":
                    self.network.add_hidden_node(node_id=gene["key"], activation_function=gene["activation_function"])

        # Optional recurrency
        if not self.allow_recurrency:
            for n in self.network.nodes.values():
                broken_connections = nn.break_loops(n)
                self.disable_connections(broken_connections)

    def random_insert_node(self):

        # A node can only be inserted if there are existing connections. Avoid  that with a try-finally block
        try:
            conn_keys = [k for k in self.genotype.keys() if self.genotype[k]["type"] == "connection"]
            key = random.choice([k for k in conn_keys if self.genotype[k]["enable"] is True])
            old_gene = self.genotype[key]

            # New node
            new_node = dict()
            new_node_genotype_key = self.global_innov_counter
            self.global_innov_counter += 1
            new_node["type"] = "node"
            new_node["key"] = self.node_innov_counter
            self.node_innov_counter += 1
            new_node["activation_function"] = random.choice(self.evolution_options.activation_functions)
            new_node["enable"] = True

            # New connection leading into the new node
            into_node = dict()
            into_node_genotype_key = self.global_innov_counter
            self.global_innov_counter += 1
            into_node["type"] = "connection"
            into_node["conn"] = []
            into_node["conn"].append(old_gene["conn"][0])
            into_node["conn"].append(new_node["key"])
            into_node["weight"] = 1
            into_node["enable"] = True

            # New connection leaving out of the new node
            out_of_node = dict()
            out_of_node_genotype_key = self.global_innov_counter
            self.global_innov_counter += 1
            out_of_node["type"] = "connection"
            out_of_node["conn"] = []
            out_of_node["conn"].append(new_node["key"])
            out_of_node["conn"].append(old_gene["conn"][1])
            out_of_node["weight"] = old_gene["weight"]
            out_of_node["enable"] = True

            # Disable old_gene
            old_gene["enable"] = False

            # append new genes to genome
            self.genotype[new_node_genotype_key] = new_node
            self.genotype[into_node_genotype_key] = into_node
            self.genotype[out_of_node_genotype_key] = out_of_node

        finally:
            pass

    def random_new_connection(self):

        conn_list = [gene for gene in self.genotype.values() if gene["type"] == "connection"]

        # Make a list of active node connections, to avoid duplicates
        en_conn_list = [g["conn"] for g in conn_list if g["enable"] is True]

        # Make a list of inactive node connections, to re-activate them instead of making a new one
        dis_conn_list = [g["conn"] for g in conn_list if g["enable"] is False]

        # Make a set of nodes that can be the input end of a connection (all nodes)
        in_node_keys = set()
        for g in conn_list:
            in_node_keys.add(g['conn'][0])
            in_node_keys.add(g['conn'][1])
        in_node_keys.update(self.input_keys)
        in_node_keys.difference_update(self.output_nodes.keys())

        # Make a set of nodes that can be the output end of a connection (all nodes, except input nodes)
        out_node_keys = set()
        for g in conn_list:
            out_node_keys.add(g['conn'][0])
            out_node_keys.add(g['conn'][1])
        out_node_keys.difference_update(self.input_keys)

        # Turn in_node_keys and out_node_keys into lists
        in_node_numbers = list(in_node_keys)
        out_node_numbers = list(out_node_keys)

        """
        Every node(M) can connect to any other non_input node(N). Non_input nodes are allowed to
        connect to themselves. Random numbers following the former criteria are selected to be the
        ends of new connections. M*N is used as the maximum number of allowed attempts, but
        keeping in mind that pairs are randomly selected, there's no guarantee that every possible
        pair will be tested.
        """
        # Initialize pair with None
        pair = [None, None]
        max_attempts = len(in_node_numbers) * len(out_node_numbers)

        while True:

            # Randomly chose an in_node and out_node
            pair[0] = random.choice(in_node_numbers)
            pair[1] = random.choice(out_node_numbers)

            # Found an available pair. Exit the loop
            if pair not in en_conn_list:
                break

            max_attempts -= 1
            if max_attempts <= 0:
                # Too many failed attempts. Return with no new connection
                return

        if pair in dis_conn_list:
            # Existing but disabled connection, re-enable it
            gene = next((g for g in conn_list if g["conn"] == pair))
            gene["enable"] = True

        else:
            # New connection, append it to genome
            new_conn = dict()

            new_conn_key = self.global_innov_counter
            self.global_innov_counter += 1
            new_conn["type"] = "connection"
            new_conn["conn"] = pair
            new_conn["weight"] = 0
            new_conn["enable"] = True

            # Append new gene to genome
            self.genotype[new_conn_key] = new_conn

            # Randomize new gene weight
            self.randomize_gene(
                gene=self.genotype[new_conn_key]
            )

    def clone(self):
        new_genotype = deepcopy(self.genotype)

        new_brain = Brain(
            input_keys=deepcopy(self.input_keys),
            output_nodes=deepcopy(self.output_nodes),
            evolution_options=self.evolution_options
        )

        new_brain.genotype = new_genotype
        new_brain.global_innov_counter = self.global_innov_counter
        new_brain.node_innov_counter = self.node_innov_counter

        new_brain.allow_recurrency = self.allow_recurrency

        new_brain.mutate_genotype()

        if chance(new_brain.evolution_options.node_insertion_chance):
            new_brain.random_insert_node()

        if chance(self.evolution_options.new_connection_chance):
            new_brain.random_new_connection()

        return new_brain
