"""
kozax: Genetic programming framework in JAX

Copyright (c) 2024 sdevries0

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

import jax
import jax.numpy as jnp
import jax.random as jr
from jax.random import PRNGKey
from functools import partial
from jax import Array
from typing import Tuple, Callable

def sample_node(i: int, 
                carry: Tuple[PRNGKey, Array, int, int, int, Array, Tuple]) -> Tuple[PRNGKey, Array, int, int, int, Array, Tuple]:
    """Samples nodes sequentially in breadth-first order, storing them depth-first.

    Parameters
    ----------
    i : int
        Index of the node.
    carry : Tuple[PRNGKey, Array, int, int, int, Array, Tuple]
        Tuple containing the random key, tree, open slots, max init depth, max nodes, variable array, and other arguments.

    Returns
    -------
    Tuple[PRNGKey, Array, int, int, int, Array, Tuple]
        Updated tuple with the random key, tree, open slots, max init depth, max nodes, variable array, and other arguments.
    """
    key, tree, open_slots, max_init_depth, max_nodes, variable_array, args = carry
    variable_indices, operator_indices, operator_probabilities, slots, coefficient_sd, map_b_to_d = args
    coefficient_key, leaf_key, variable_key, node_key, operator_key = jr.split(key, 5)
    _i = map_b_to_d[i].astype(int)  # Get depth first index

    depth = (jnp.log(i + 1 + 1e-10) / jnp.log(2)).astype(int)  # Compute depth of node
    coefficient = jr.normal(coefficient_key) * coefficient_sd
    leaf = jax.lax.select(jr.uniform(leaf_key) < 0.5, 1, jr.choice(variable_key, variable_indices, shape=(), p=variable_array))  # Sample coefficient or variable

    index = jax.lax.select((open_slots < max_nodes - i - 1) & (depth + 1 < max_init_depth),  # Check if max depth has been reached, or if the number of open slots reached the max number of nodes
                           jax.lax.select(jr.uniform(node_key) < (0.7 ** depth),  # At higher depth, a leaf node is more probable
                                           jr.choice(operator_key, a=operator_indices, shape=(), p=operator_probabilities), 
                                           leaf), 
                           leaf)
    
    index = jax.lax.select(open_slots == 0, 0, index)  # If there are no open slots, the node should be empty

    # If parent node is a leaf, the node should be empty
    index = jax.lax.select(i > 0, jax.lax.select((slots[jnp.maximum(tree[map_b_to_d[(i + (i % 2) - 2) // 2].astype(int), 0], 0).astype(int)] + i % 2) > 1, index, 0), index)

    # Set index references
    tree = jax.lax.select(slots[index] > 0, tree.at[_i, 1].set(map_b_to_d[2 * i + 1]), tree.at[_i, 1].set(-1))
    tree = jax.lax.select(slots[index] > 1, tree.at[_i, 2].set(map_b_to_d[2 * i + 2]), tree.at[_i, 2].set(-1))

    tree = jax.lax.select(index == 1, tree.at[_i, 3].set(coefficient), tree)  # Set coefficient value
    tree = tree.at[_i, 0].set(index)

    open_slots = jax.lax.select(index == 0, open_slots, jnp.maximum(0, open_slots + slots[index] - 1))  # Update the number of open slots

    return (jr.fold_in(key, i), tree, open_slots, max_init_depth, max_nodes, variable_array, args)

def prune_row(i: int, 
              carry: Tuple[Array, int, int], 
              old_tree: Array) -> Tuple[Array, int, int]:
    """Sequentially adds nodes to the new tree if it is not empty.

    Parameters
    ----------
    i : int
        Index of the node.
    carry : Tuple[Array, int, int]
        Tuple containing the tree, counter, and tree size.
    old_tree : Array
        Tree with empty nodes that have to be pruned.

    Returns
    -------
    Tuple[Array, int, int]
        Updated tuple with the tree, counter, and tree size.
    """
    tree, counter, tree_size = carry
    _i = tree_size - i - 1
    row = old_tree[_i]

    # If node is not empty, add node and update index references
    tree = jax.lax.select(row[0] != 0, tree.at[counter].set(row), tree.at[:, 1:3].set(jnp.where(tree[:, 1:3] > _i, tree[:, 1:3] - 1, tree[:, 1:3])))
    counter = jax.lax.select(row[0] != 0, counter - 1, counter)

    return (tree, counter, tree_size)
    
def prune_tree(tree: Array, 
               tree_size: int, 
               max_nodes: int) -> Array:
    """Removes empty nodes from a tree. The new tree is filled with empty nodes at the end to match the max number of nodes.

    Parameters
    ----------
    tree : Array
        Tree to be pruned.
    tree_size : int
        Max size of the old tree.
    max_nodes : int
        Max number of nodes in the new tree.

    Returns
    -------
    Array
        Tree with empty nodes pruned.
    """
    tree, counter, _ = jax.lax.fori_loop(0, tree_size, partial(prune_row, old_tree=tree), (jnp.tile(jnp.array([0.0, -1.0, -1.0, 0.0]), (max_nodes, 1)), max_nodes - 1, tree_size))
    tree = tree.at[:, 1:3].set(jnp.where(tree[:, 1:3] > -1, tree[:, 1:3] + counter + 1, tree[:, 1:3]))  # Update index references after pruning
    return tree

def sample_tree(key: PRNGKey, 
                depth: int, 
                variable_array: Array, 
                tree_size: int, 
                max_nodes: int, 
                args: Tuple) -> Array:
    """Initializes a tree.

    Parameters
    ----------
    key : PRNGKey
        Random key.
    depth : int
        Max depth in a tree at initialization.
    variable_array : Array
        The valid variables for this tree.
    tree_size : int
        Max size of the tree.
    max_nodes : int
        Max number of nodes in a tree.
    args : Tuple
        Miscellaneous parameters required for initialization.

    Returns
    -------
    Array
        Initialized tree.
    """
    # First sample tree at full size given depth
    tree = jax.lax.fori_loop(0, tree_size, sample_node, (key, jnp.zeros((tree_size, 4)), 1, depth, max_nodes, variable_array, args))[1]  # Sample nodes in a tree sequentially

    # Prune empty rows in tree
    pruned_tree = prune_tree(tree, tree_size, max_nodes)
    return pruned_tree

def sample_population(key: PRNGKey, 
                      population_size: int, 
                      num_trees: int, 
                      max_init_depth: int, 
                      variable_array: Array,
                      sample_function: Callable) -> Array:
    """Initializes a population of candidates.

    Parameters
    ----------
    key : PRNGKey
        Random key.
    population_size : int
        Number of candidates that have to be sampled.
    num_trees : int
        Number of trees in a candidate.
    max_init_depth : int
        Max depth in a tree at initialization.
    variable_array : Array
        The valid variables for each tree.
    sample_function : Callable
        Function to sample a tree.

    Returns
    -------
    Array
        Population of candidates.
    """
    sample_candidate = lambda keys: jax.vmap(sample_function, in_axes=[0, None, 0])(keys, max_init_depth, variable_array)
    return jax.vmap(sample_candidate)(jr.split(key, (population_size, num_trees)))