radiate-gp 1.3.1

use super::TreeChromosome;
use crate::TreeNode;
use radiate_core::{AlterContext, Crossover, Expr, RateSet, random_provider};
use radiate_core::{SmallStr, genome::*};

const DEFAULT_MAX_SIZE: usize = 30;
const MAX_ATTEMPTS: usize = 3;
const TN_X_ATTEMPTS: SmallStr = SmallStr::from_static("tn_x_att");

#[derive(Clone, Debug)]
pub struct TreeCrossover {
    rate: Expr,
    max_size: usize,
}

impl TreeCrossover {
    pub fn new(rate: impl Into<Expr>) -> Self {
        TreeCrossover {
            rate: rate.into(),
            max_size: DEFAULT_MAX_SIZE,
        }
    }

    pub fn with_max_size(mut self, max_size: usize) -> Self {
        self.max_size = max_size;
        self
    }

    pub fn cross_nodes<T>(
        node_one: &mut TreeNode<T>,
        node_two: &mut TreeNode<T>,
        max_size: usize,
        ctx: &mut AlterContext,
    ) -> usize {
        let one_size = node_one.size();
        let two_size = node_two.size();

        if one_size == 1 || two_size == 1 {
            return 0;
        }

        let mut attempts = 0;
        while attempts < MAX_ATTEMPTS {
            let one_rand_index = random_provider::range(1..one_size);
            let two_rand_index = random_provider::range(1..two_size);

            let one_sub_node = node_one.get_mut(one_rand_index);
            let two_sub_node = node_two.get_mut(two_rand_index);

            if let (Some(one_sub_node), Some(two_sub_node)) = (one_sub_node, two_sub_node) {
                let one_sub_size = one_sub_node.size();
                let two_sub_size = two_sub_node.size();

                let one_crossover_size = one_size - one_sub_size + two_sub_size;
                let two_crossover_size = two_size - two_sub_size + one_sub_size;

                if one_crossover_size <= max_size && two_crossover_size <= max_size {
                    std::mem::swap(one_sub_node, two_sub_node);
                    ctx.upsert(TN_X_ATTEMPTS, attempts + 1);
                    return 2;
                }
            }

            attempts += 1;
        }

        0
    }
}

impl<T> Crossover<TreeChromosome<T>> for TreeCrossover
where
    T: Clone + PartialEq,
{
    fn rates(&self) -> RateSet {
        RateSet::new(self.rate.clone())
    }

    #[inline]
    fn cross_chromosomes(
        &self,
        chrom_one: &mut TreeChromosome<T>,
        chrom_two: &mut TreeChromosome<T>,
        ctx: &mut AlterContext,
    ) -> usize {
        let swap_one_index = random_provider::range(0..chrom_one.len());
        let swap_two_index = random_provider::range(0..chrom_two.len());

        let one_node = chrom_one.get_mut(swap_one_index);
        let two_node = chrom_two.get_mut(swap_two_index);

        if let Some((one_node, two_node)) = one_node.zip(two_node) {
            return Self::cross_nodes(one_node, two_node, self.max_size, ctx);
        }

        0
    }
}