radiate-core 1.3.1

use super::{Chromosome, Genotype, Population, random_provider};
use crate::{Ecosystem, MetricSet, Phenotype, error::RadiateResult, metric_names};
use radiate_error::radiate_bail;
use radiate_expr::Expr;
use radiate_utils::{AnyValue, DataType};
use std::{collections::HashSet, sync::Arc};

/// Trait for replacement strategies in the algorithms.
///
/// This trait defines a method for replacing a member of the [Population] with a new individual
/// after the current individual has been determined to be invalid. Typically, this is done by
/// replacing the individual with a new one generated by the encoder. But in some cases, it may
/// be desirable to replace the individual in a different way, such as by sampling from the
/// [Population].
pub trait ReplacementStrategy<C: Chromosome>: Send + Sync {
    fn replace(
        &self,
        population: &Population<C>,
        encoder: Arc<dyn Fn() -> Genotype<C> + Send + Sync>,
    ) -> Genotype<C>;

    fn replace_at(
        &self,
        _: usize,
        population: &Population<C>,
        encoder: Arc<dyn Fn() -> Genotype<C> + Send + Sync>,
    ) -> Genotype<C> {
        self.replace(population, encoder)
    }
}

/// Replacement strategy that replaces the individual with a new one generated by the encoder.
/// This is the default replacement strategy used in genetic algorithms.
pub struct EncodeReplace;

impl<C: Chromosome> ReplacementStrategy<C> for EncodeReplace {
    fn replace(
        &self,
        _: &Population<C>,
        encoder: Arc<dyn Fn() -> Genotype<C> + Send + Sync>,
    ) -> Genotype<C> {
        encoder()
    }
}

/// Replacement strategy that replaces the individual with a random member of the [Population].
/// This can be useful in cases where the [Population] is large and diverse or when the
/// chromosome grows or changes in size, thus encoding a new individual can result
/// in a member that that lacks significant diversity.
pub struct PopulationSampleReplace;

impl<C: Chromosome + Clone> ReplacementStrategy<C> for PopulationSampleReplace {
    fn replace(
        &self,
        population: &Population<C>,
        _: Arc<dyn Fn() -> Genotype<C> + Send + Sync>,
    ) -> Genotype<C> {
        let random_member = random_provider::range(0..population.len());
        population[random_member].genotype().clone()
    }
}

pub trait EcosystemFilter<C: Chromosome>: Send + Sync {
    fn filter(
        &mut self,
        generation: usize,
        ecosystem: &mut Ecosystem<C>,
        metrics: &mut MetricSet,
        replacer: Arc<dyn ReplacementStrategy<C>>,
        encoder: Arc<dyn Fn() -> Genotype<C> + Send + Sync>,
    ) -> RadiateResult<()>;
}

pub struct UniqueScoreFilter {
    filter_cond: Expr,
}

impl UniqueScoreFilter {
    pub fn new(max_stagnation: usize, threshold: f32) -> Self {
        Self {
            filter_cond: Expr::select(metric_names::BEST_SCORES)
                .stagnation(threshold)
                .cast(DataType::Usize)
                .gt(max_stagnation),
        }
    }
}

impl<C: Chromosome> EcosystemFilter<C> for UniqueScoreFilter {
    fn filter(
        &mut self,
        generation: usize,
        ecosystem: &mut Ecosystem<C>,
        metrics: &mut MetricSet,
        replacer: Arc<dyn ReplacementStrategy<C>>,
        encoder: Arc<dyn Fn() -> Genotype<C> + Send + Sync>,
    ) -> RadiateResult<()> {
        let should_filter = self.filter_cond.evaluate(metrics)?;

        let AnyValue::Bool(bool) = should_filter else {
            radiate_bail!(Engine: "Filter condition must evaluate to a boolean value, but got: {:?}", should_filter);
        };

        if !bool {
            metrics.upsert(metric_names::FILTER_UNIQUE_SCORES, 0);
            return Ok(());
        }

        let mut unique_scores = HashSet::new();
        let mut replaced_count = 0;

        for i in 0..ecosystem.population.len() {
            if let Some(phenotype) = ecosystem.get_phenotype(i) {
                let score = match phenotype.score() {
                    Some(score) => score.clone(),
                    None => continue,
                };

                if !unique_scores.insert(score) {
                    replaced_count += 1;
                    let new_genotype =
                        replacer.replace_at(i, ecosystem.population(), Arc::clone(&encoder));
                    if let Some(phenotype_mut) = ecosystem.get_phenotype_mut(i) {
                        *phenotype_mut = Phenotype::from((new_genotype, generation));
                    }
                }
            }
        }

        metrics.upsert(metric_names::FILTER_UNIQUE_SCORES, replaced_count);

        Ok(())
    }
}