import jax
from jax import vmap, numpy as jnp
from . import BaseMutation
from tensorneat.common import (
    fetch_first,
    fetch_random,
    I_INF,
    check_cycles,
)
from ...utils import (
    unflatten_conns,
    add_node,
    add_conn,
    delete_node_by_pos,
    delete_conn_by_pos,
    extract_gene_attrs, 
    set_gene_attrs
)


class DefaultMutation(BaseMutation):
    def __init__(
        self,
        conn_add: float = 0.2,
        conn_delete: float = 0.2,
        node_add: float = 0.1,
        node_delete: float = 0.1,
    ):
        self.conn_add = conn_add
        self.conn_delete = conn_delete
        self.node_add = node_add
        self.node_delete = node_delete

    def __call__(
        self, state, genome, randkey, nodes, conns, new_node_key, new_conn_key
    ):
        assert (
            new_node_key.shape == ()
        )  # scalar, as there is max one new node in each mutation
        assert new_conn_key.shape == (
            3,
        )  # there are max 3 new connections (mutate add node + mutate add conn)

        k1, k2 = jax.random.split(randkey)

        nodes, conns = self.mutate_structure(
            state, genome, k1, nodes, conns, new_node_key, new_conn_key
        )
        nodes, conns = self.mutate_values(state, genome, k2, nodes, conns)

        return nodes, conns

    def mutate_structure(
        self, state, genome, randkey, nodes, conns, new_node_key, new_conn_key
    ):
        def mutate_add_node(key_, nodes_, conns_):
            """
            add a node while do not influence the output of the network
            """

            remain_node_space = jnp.isnan(nodes_[:, 0]).sum()
            remain_conn_space = jnp.isnan(conns_[:, 0]).sum()
            i_key, o_key, idx = self.choose_connection_key(
                key_, conns_
            )  # choose a connection

            def successful_add_node():
                # remove the original connection and record its attrs
                original_attrs = extract_gene_attrs(genome.conn_gene, conns_[idx])
                new_conns = delete_conn_by_pos(conns_, idx)

                # add a new node with identity attrs
                new_nodes = add_node(
                    nodes_, jnp.array([new_node_key]), genome.node_gene.new_identity_attrs(state)
                )

                # whether to use historical marker in connection
                if "historical_marker" in genome.conn_gene.fixed_attrs:
                    fix_attrs1 = jnp.array([i_key, new_node_key, new_conn_key[0]])
                    fix_attrs2 = jnp.array([new_node_key, o_key, new_conn_key[1]])
                else:
                    fix_attrs1 = jnp.array([i_key, new_node_key])
                    fix_attrs2 = jnp.array([new_node_key, o_key])

                # add two new connections
                # first is with identity attrs
                new_conns = add_conn(
                    new_conns,
                    fix_attrs1,
                    genome.conn_gene.new_identity_attrs(state),
                )
                # second is with the origin attrs
                new_conns = add_conn(
                    new_conns,
                    fix_attrs2,
                    original_attrs,
                )

                return new_nodes, new_conns

            return jax.lax.cond(
                (idx == I_INF) | (remain_node_space < 1) | (remain_conn_space < 2),
                lambda: (nodes_, conns_),  # do nothing
                successful_add_node,
            )

        def mutate_delete_node(key_, nodes_, conns_):
            """
            delete a node
            """
            # randomly choose a node
            key, idx = self.choose_node_key(
                key_,
                nodes_,
                genome.input_idx,
                genome.output_idx,
                allow_input_keys=False,
                allow_output_keys=False,
            )

            def successful_delete_node():
                # delete the node
                new_nodes = delete_node_by_pos(nodes_, idx)

                # delete all connections
                new_conns = jnp.where(
                    ((conns_[:, 0] == key) | (conns_[:, 1] == key))[:, None],
                    jnp.nan,
                    conns_,
                )

                return new_nodes, new_conns

            return jax.lax.cond(
                idx == I_INF,  # no available node to delete
                lambda: (nodes_, conns_),  # do nothing
                successful_delete_node,
            )

        def mutate_add_conn(key_, nodes_, conns_):
            """
            add a connection while do not influence the output of the network
            """

            remain_conn_space = jnp.isnan(conns_[:, 0]).sum()

            # randomly choose two nodes
            k1_, k2_ = jax.random.split(key_, num=2)

            # input node of the connection can be any node
            i_key, from_idx = self.choose_node_key(
                k1_,
                nodes_,
                genome.input_idx,
                genome.output_idx,
                allow_input_keys=True,
                allow_output_keys=True,
            )

            # output node of the connection can be any node except input node
            o_key, to_idx = self.choose_node_key(
                k2_,
                nodes_,
                genome.input_idx,
                genome.output_idx,
                allow_input_keys=False,
                allow_output_keys=True,
            )

            conn_pos = fetch_first((conns_[:, 0] == i_key) & (conns_[:, 1] == o_key))
            is_already_exist = conn_pos != I_INF

            def nothing():
                return nodes_, conns_

            def successful():
                # add a connection with zero attrs
                if "historical_marker" in genome.conn_gene.fixed_attrs:
                    new_fix_attrs = jnp.array([i_key, o_key, new_conn_key[2]])
                else:
                    new_fix_attrs = jnp.array([i_key, o_key])
                return nodes_, add_conn(
                    conns_, new_fix_attrs, genome.conn_gene.new_zero_attrs(state)
                )

            if genome.network_type == "feedforward":
                u_conns = unflatten_conns(nodes_, conns_)
                conns_exist = u_conns != I_INF
                is_cycle = check_cycles(nodes_, conns_exist, from_idx, to_idx)

                return jax.lax.cond(
                    is_already_exist | is_cycle | (remain_conn_space < 1),
                    nothing,
                    successful,
                )

            elif genome.network_type == "recurrent":
                return jax.lax.cond(
                    is_already_exist | (remain_conn_space < 1),
                    nothing,
                    successful,
                )

            else:
                raise ValueError(f"Invalid network type: {genome.network_type}")

        def mutate_delete_conn(key_, nodes_, conns_):
            # randomly choose a connection
            i_key, o_key, idx = self.choose_connection_key(key_, conns_)

            return jax.lax.cond(
                idx == I_INF,
                lambda: (nodes_, conns_),  # nothing
                lambda: (nodes_, delete_conn_by_pos(conns_, idx)),  # success
            )

        k1, k2, k3, k4 = jax.random.split(randkey, num=4)
        r1, r2, r3, r4 = jax.random.uniform(k1, shape=(4,))

        def nothing(_, nodes_, conns_):
            return nodes_, conns_

        if self.node_add > 0:
            nodes, conns = jax.lax.cond(
                r1 < self.node_add, mutate_add_node, nothing, k1, nodes, conns
            )

        if self.node_delete > 0:
            nodes, conns = jax.lax.cond(
                r2 < self.node_delete, mutate_delete_node, nothing, k2, nodes, conns
            )

        if self.conn_add > 0:
            nodes, conns = jax.lax.cond(
                r3 < self.conn_add, mutate_add_conn, nothing, k3, nodes, conns
            )

        if self.conn_delete > 0:
            nodes, conns = jax.lax.cond(
                r4 < self.conn_delete, mutate_delete_conn, nothing, k4, nodes, conns
            )

        return nodes, conns

    def mutate_values(self, state, genome, randkey, nodes, conns):
        k1, k2 = jax.random.split(randkey)
        nodes_randkeys = jax.random.split(k1, num=genome.max_nodes)
        conns_randkeys = jax.random.split(k2, num=genome.max_conns)

        node_attrs = vmap(extract_gene_attrs, in_axes=(None, 0))(
            genome.node_gene, nodes
        )
        new_node_attrs = vmap(genome.node_gene.mutate, in_axes=(None, 0, 0))(
            state, nodes_randkeys, node_attrs
        )
        new_nodes = vmap(set_gene_attrs, in_axes=(None, 0, 0))(
            genome.node_gene, nodes, new_node_attrs
        )

        conn_attrs = vmap(extract_gene_attrs, in_axes=(None, 0))(
            genome.conn_gene, conns
        )
        new_conn_attrs = vmap(genome.conn_gene.mutate, in_axes=(None, 0, 0))(
            state, conns_randkeys, conn_attrs
        )
        new_conns = vmap(set_gene_attrs, in_axes=(None, 0, 0))(
            genome.conn_gene, conns, new_conn_attrs
        )

        # nan nodes not changed
        new_nodes = jnp.where(jnp.isnan(nodes), jnp.nan, new_nodes)
        new_conns = jnp.where(jnp.isnan(conns), jnp.nan, new_conns)

        return new_nodes, new_conns

    def choose_node_key(
        self,
        key,
        nodes,
        input_idx,
        output_idx,
        allow_input_keys: bool = False,
        allow_output_keys: bool = False,
    ):
        """
        Randomly choose a node key from the given nodes. It guarantees that the chosen node not be the input or output node.
        :param key:
        :param nodes:
        :param input_idx:
        :param output_idx:
        :param allow_input_keys:
        :param allow_output_keys:
        :return: return its key and position(idx)
        """

        node_keys = nodes[:, 0]
        mask = ~jnp.isnan(node_keys)

        if not allow_input_keys:
            mask = jnp.logical_and(mask, ~jnp.isin(node_keys, input_idx))

        if not allow_output_keys:
            mask = jnp.logical_and(mask, ~jnp.isin(node_keys, output_idx))

        idx = fetch_random(key, mask)
        key = jnp.where(idx != I_INF, nodes[idx, 0], jnp.nan)
        return key, idx

    def choose_connection_key(self, key, conns):
        """
        Randomly choose a connection key from the given connections.
        :return: i_key, o_key, idx
        """

        idx = fetch_random(key, ~jnp.isnan(conns[:, 0]))
        i_key = jnp.where(idx != I_INF, conns[idx, 0], jnp.nan)
        o_key = jnp.where(idx != I_INF, conns[idx, 1], jnp.nan)

        return i_key, o_key, idx
