ebi 0.3.14

use crate::{
    ebi_framework::{displayable::Displayable, ebi_command::EbiCommand},
    ebi_traits::{
        ebi_trait_finite_stochastic_language::EbiTraitFiniteStochasticLanguage,
        ebi_trait_stochastic_partially_ordered_semantics::EbiTraitStochasticPartiallyOrderedSemantics,
        ebi_trait_stochastic_semantics::EbiTraitStochasticSemantics,
    },
    stochastic_partially_ordered_semantics::stochastic_partially_ordered_semantics::StochasticPartiallyOrderedSemantics,
    stochastic_semantics::stochastic_semantics::StochasticSemantics,
    techniques::livelock::LiveLockCache,
};
use ebi_objects::{
    FiniteStochasticLanguage, FiniteStochasticPartiallyOrderedLanguage, HasActivityKey,
    StochasticBusinessProcessModelAndNotation,
    anyhow::{Result, anyhow},
    ebi_arithmetic::{ChooseRandomly, Fraction, FractionRandomCache, One, Recip, Zero, f},
    ebi_bpmn::{BPMNMarking, partially_ordered_run::PartiallyOrderedRun},
};
use rand::RngExt;
use rayon::iter::{IntoParallelIterator, ParallelIterator};
use std::collections::hash_map::Entry;

pub trait Sampler {
    fn sample(&self, number_of_traces: usize) -> Result<FiniteStochasticLanguage>;
}

pub trait PartiallyOrderedSampler {
    fn sample_partially_ordered(
        &self,
        number_of_traces: usize,
    ) -> Result<FiniteStochasticPartiallyOrderedLanguage>;
}

pub trait Resampler {
    fn resample_cache_init(&self) -> Result<FractionRandomCache>;
    fn resample(&self, cache: &FractionRandomCache, number_of_traces: usize) -> Vec<Fraction>;
}

impl Sampler for EbiTraitStochasticSemantics {
    fn sample(&self, number_of_traces: usize) -> Result<FiniteStochasticLanguage> {
        match self {
            EbiTraitStochasticSemantics::Marking(s) => s.sample(number_of_traces),
            EbiTraitStochasticSemantics::Usize(s) => s.sample(number_of_traces),
            EbiTraitStochasticSemantics::AutomatonState(s) => s.sample(number_of_traces),
            EbiTraitStochasticSemantics::TreeMarking(s) => s.sample(number_of_traces),
        }
    }
}

impl Sampler for dyn EbiTraitFiniteStochasticLanguage {
    fn sample(&self, number_of_traces: usize) -> Result<FiniteStochasticLanguage> {
        let mut result = FiniteStochasticLanguage::new_hashmap();

        let cache = self.resample_cache_init()?;

        if self.number_of_traces().is_zero() {
            return Err(anyhow!("Cannot sample from empty language."));
        }

        for _ in 0..number_of_traces {
            let trace_index = Fraction::choose_randomly_cached(&cache);

            match result.entry(self.iter_traces().nth(trace_index).unwrap().clone()) {
                Entry::Occupied(mut e) => *e.get_mut() += 1,
                Entry::Vacant(e) => {
                    e.insert(Fraction::one());
                }
            };
        }

        Ok((self.activity_key().clone(), result).into())
    }
}

impl PartiallyOrderedSampler for StochasticBusinessProcessModelAndNotation {
    fn sample_partially_ordered(
        &self,
        number_of_traces: usize,
    ) -> Result<FiniteStochasticPartiallyOrderedLanguage> {
        let probabilities = vec![f!(number_of_traces).recip(); number_of_traces];

        let progress_bar = EbiCommand::get_progress_bar_ticks(number_of_traces);

        // Create thread pool with custom configuration
        let pool = rayon::ThreadPoolBuilder::new().build().unwrap();

        let mut traces = Vec::with_capacity(number_of_traces);
        pool.install(|| {
            traces = (0..number_of_traces)
                .into_par_iter()
                .map(|_| {
                    let run = PartiallyOrderedRun::new_random(self).unwrap().into();
                    progress_bar.inc(1);
                    run
                })
                .collect();
        });

        progress_bar.finish_and_clear();

        Ok((self.activity_key().clone(), traces, probabilities).into())
    }
}

impl<T, State> Sampler for T
where
    T: StochasticSemantics<StoSemState = State> + ?Sized,
    State: Displayable,
{
    fn sample(&self, number_of_traces: usize) -> Result<FiniteStochasticLanguage> {
        if let Some(initial_state) = self.get_initial_state() {

            //We need to ensure that we don't get stuck in a livelock.
            //Strategy: check it every 1000 states.
            let mut maybe_livelock_cache: Option<Box<dyn LiveLockCache<LivState = State>>> = None;

            let mut result = FiniteStochasticLanguage::new_hashmap();
            let mut sampled_traces = 0;

            'outer: while sampled_traces < number_of_traces {
                let mut current_state = initial_state.clone();
                let mut trace = vec![];

                let mut number_of_states_checked = 0;

                let mut outgoing_probabilities = vec![];

                while !self.is_final_state(&current_state) {
                    let total_weight = self
                        .get_total_weight_of_enabled_transitions(&current_state)
                        .unwrap();
                    let enabled_transitions = self.get_enabled_transitions(&current_state);

                    outgoing_probabilities.clear();
                    for transition in &enabled_transitions {
                        outgoing_probabilities.push(
                            self.get_transition_weight(&current_state, *transition)?
                                / &total_weight,
                        );
                    }

                    // get firing transition
                    let i = Fraction::choose_randomly(&outgoing_probabilities)?;
                    let chosen_transition = enabled_transitions[i];

                    let activity = self.get_transition_activity(chosen_transition, &current_state);

                    // execute transition
                    self.execute_transition(&mut current_state, chosen_transition)?;

                    match activity {
                        Some(activity) => trace.push(activity),
                        None => {}
                    }

                    // Ensure we do not end up in a livelock, but only every 1000 states.
                    number_of_states_checked += 1;
                    if number_of_states_checked == 1000 {
                        if let Some(cache) = &mut maybe_livelock_cache {
                            if cache.is_state_part_of_livelock(&current_state)? {
                                //we have hit a livelock; restart the sample
                                continue 'outer;
                            }
                        } else {
                            //create the cache
                            let mut cache = self.get_livelock_cache();

                            //check whether the initial state is a livelock
                            if cache.is_state_part_of_livelock(&initial_state)? {
                                return Err(anyhow!(
                                    "The initial state is a livelock, therefore there are no traces to sample."
                                ));
                            }

                            //check whether the current state is a livelock
                            if cache.is_state_part_of_livelock(&current_state)? {
                                //we have hit a livelock; restart the sample
                                continue 'outer;
                            }

                            maybe_livelock_cache = Some(cache);
                        }
                        number_of_states_checked = 0;
                    }
                }

                match result.entry(trace) {
                    Entry::Occupied(mut e) => *e.get_mut() += 1,
                    Entry::Vacant(e) => {
                        e.insert(Fraction::one());
                    }
                };
                sampled_traces += 1;
            }

            if result.is_empty() {
                return Err(anyhow!(
                    "Analysis resulted in an empty language; there are no traces in the model."
                ));
            }

            // log::debug!("Sampled {:?} traces", result);

            Ok((self.activity_key().clone(), result).into())
        } else {
            return Err(anyhow!("Language contains no traces, so cannot sample."));
        }
    }
}

impl Resampler for dyn EbiTraitFiniteStochasticLanguage {
    fn resample_cache_init(&self) -> Result<FractionRandomCache> {
        if self.number_of_traces().is_zero() {
            return Err(anyhow!("Cannot sample from empty language."));
        }

        let probabilities = self.iter_traces_probabilities().map(|(_, p)| p);
        let cache = Fraction::choose_randomly_create_cache(probabilities)?;

        Ok(cache)
    }

    fn resample(&self, cache: &FractionRandomCache, number_of_traces: usize) -> Vec<Fraction> {
        let mut result = vec![0usize; self.number_of_traces()];
        for _ in 0..number_of_traces {
            let trace_index = Fraction::choose_randomly_cached(&cache);
            result[trace_index] += 1;
        }

        //normalise vector
        result
            .into_iter()
            .map(|c| Fraction::from((c, number_of_traces)))
            .collect()
    }
}

impl Sampler for EbiTraitStochasticPartiallyOrderedSemantics {
    fn sample(&self, number_of_traces: usize) -> Result<FiniteStochasticLanguage> {
        match self {
            EbiTraitStochasticPartiallyOrderedSemantics::BPMNMarking(bpmn) => {
                bpmn.sample(number_of_traces)
            }
        }
    }
}

impl Sampler
    for dyn StochasticPartiallyOrderedSemantics<
            StoPOSemState = BPMNMarking,
            SemState = BPMNMarking,
            AliState = BPMNMarking,
        >
{
    fn sample(&self, number_of_traces: usize) -> Result<FiniteStochasticLanguage> {
        if let Some(initial_state) = self.get_initial_state() {
            let mut result = FiniteStochasticLanguage::new_hashmap();

            for _ in 0..number_of_traces {
                let mut current_state = initial_state.clone();
                let mut trace = vec![];

                let mut outgoing_probabilities = vec![];

                while !self.is_final_state(&current_state) {
                    let enabled_transitions = self.get_enabled_transitions(&current_state);

                    outgoing_probabilities.clear();
                    for transition in &enabled_transitions {
                        outgoing_probabilities.push(
                            self.get_transition_probabilistic_penalty(&current_state, *transition)
                                .ok_or_else(|| anyhow!("transition not found"))?,
                        );
                    }

                    // get firing transition
                    let i = Fraction::choose_randomly(&outgoing_probabilities)?;
                    let chosen_transition = enabled_transitions[i];

                    let activity = self.get_transition_activity(chosen_transition, &current_state);

                    // execute transition
                    self.execute_transition(&mut current_state, chosen_transition)?;

                    match activity {
                        Some(activity) => trace.push(activity),
                        None => {}
                    }
                }

                match result.entry(trace) {
                    Entry::Occupied(mut e) => *e.get_mut() += 1,
                    Entry::Vacant(e) => {
                        e.insert(Fraction::one());
                    }
                };
            }

            if result.is_empty() {
                return Err(anyhow!(
                    "Analysis resulted in an empty language; there are no traces in the model."
                ));
            }

            // log::debug!("Sampled {:?} traces", result);

            Ok((self.activity_key().clone(), result).into())
        } else {
            return Err(anyhow!("Language contains no traces, so cannot sample."));
        }
    }
}

/**
 * Fills the given vector with uniformly random numbers in the range 0..`number_of_indices`.
 */
pub fn sample_indices_uniform(number_of_indices: usize, result: &mut Vec<usize>) {
    let mut rng = rand::rng();
    for i in 0..result.len() {
        let trace_index = rng.random_range(0..number_of_indices);
        result[i] = trace_index;
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        ebi_framework::trait_importers::ToStochasticSemanticsTrait,
        ebi_traits::{
            ebi_trait_finite_stochastic_language::EbiTraitFiniteStochasticLanguage,
            ebi_trait_stochastic_semantics::EbiTraitStochasticSemantics,
        },
    };
    use ebi_objects::{FiniteStochasticLanguage, StochasticDirectlyFollowsModel};
    use std::fs;

    use super::Sampler;

    #[test]
    fn sample_finite_stochastic_language() {
        let fin = fs::read_to_string("testfiles/aa.slang").unwrap();
        let slpn = fin.parse::<FiniteStochasticLanguage>().unwrap();

        let slpn: Box<dyn EbiTraitFiniteStochasticLanguage> = Box::new(slpn);

        let sample = slpn.sample(10).unwrap();

        let slpn2 = fin.parse::<FiniteStochasticLanguage>().unwrap();

        assert_eq!(sample, slpn2);
    }

    #[test]
    fn sample_stochastic_semantics() {
        let fin = fs::read_to_string("testfiles/aa.slang").unwrap();
        let slpn = fin.parse::<FiniteStochasticLanguage>().unwrap();

        let slpn: EbiTraitStochasticSemantics = slpn.to_stochastic_semantics_trait();

        let sample = slpn.sample(10).unwrap();

        let slpn2 = fin.parse::<FiniteStochasticLanguage>().unwrap();

        assert_eq!(sample, slpn2);
    }

    #[test]
    fn sample_finite_stochastic_language_empty() {
        let fin = fs::read_to_string("testfiles/empty.slang").unwrap();
        let slpn = fin.parse::<FiniteStochasticLanguage>().unwrap();

        let slpn: Box<dyn EbiTraitFiniteStochasticLanguage> = Box::new(slpn);

        assert!(slpn.sample(10).is_err());
    }

    #[test]
    fn sample_stochastic_semantics_empty() {
        let fin = fs::read_to_string("testfiles/empty.slang").unwrap();
        let slpn = fin.parse::<FiniteStochasticLanguage>().unwrap();

        let slpn: EbiTraitStochasticSemantics = slpn.to_stochastic_semantics_trait();

        assert!(slpn.sample(10).is_err());
    }

    #[cfg(test)]
    pub mod tests {
        use crate::techniques::sample::Sampler;
        use ebi_objects::StochasticNondeterministicFiniteAutomaton;
        use std::fs;

        #[test]
        fn snfa_sample() {
            let fin = fs::read_to_string("testfiles/aa-ab-ba.snfa").unwrap();
            let snfa = fin
                .parse::<StochasticNondeterministicFiniteAutomaton>()
                .unwrap();

            println!("snfa {}", snfa);

            let sample = snfa.sample(10).unwrap();
            println!("{}", sample);
            for (trace, _) in sample {
                assert_eq!(trace.len(), 2);
            }
        }
    }

    #[test]
    fn sdfm_sample() {
        let fin = fs::read_to_string("./testfiles/aa-ab-ba.sdfm").unwrap();
        let sdfm = fin.parse::<StochasticDirectlyFollowsModel>().unwrap();

        let sem = sdfm.to_stochastic_semantics_trait();

        sem.sample(10).unwrap();
    }
}