use super::align::Align;
use crate::ebi_traits::ebi_trait_finite_stochastic_language::EbiTraitFiniteStochasticLanguage;
use ebi_objects::{
    BusinessProcessModelAndNotation, LabelledPetriNet, StochasticBusinessProcessModelAndNotation,
    StochasticLabelledPetriNet,
    anyhow::{Context, Result, anyhow},
    ebi_arithmetic::{Fraction, Zero},
    ebi_bpmn::Token,
    ebi_objects::language_of_alignments::Move,
};

pub trait AlignmentMiner {
    type T;

    fn mine_stochastic_alignment(
        self,
        language: Box<dyn EbiTraitFiniteStochasticLanguage>,
    ) -> Result<Self::T>;
}

impl AlignmentMiner for BusinessProcessModelAndNotation {
    type T = StochasticBusinessProcessModelAndNotation;

    fn mine_stochastic_alignment(
        mut self,
        language: Box<dyn EbiTraitFiniteStochasticLanguage>,
    ) -> Result<Self::T> {
        if self.get_initial_marking()?.is_none() {
            //empty language stays empty
            return Ok(StochasticBusinessProcessModelAndNotation { bpmn: self });
        }

        //reset the weights
        {
            let ids = self
                .sequence_flows()
                .iter()
                .map(|sequence_flow| sequence_flow.global_index())
                .collect::<Vec<_>>();
            for sequence_flow_global_index in ids {
                self.global_index_2_sequence_flow_mut(sequence_flow_global_index)
                    .ok_or_else(|| anyhow!("sequence flow not found"))?
                    .weight = Some(Fraction::zero());
            }
        }

        let alignments = self.align_stochastic_language(language)?;
        for index in 0..alignments.len() {
            let probability = alignments
                .get_probability(index)
                .ok_or_else(|| anyhow!("should not happen"))?;

            if let Some(mut marking) = self.get_initial_marking()? {
                for movee in alignments
                    .get(index)
                    .ok_or_else(|| anyhow!("should not happen"))?
                {
                    // println!("move {:?}", movee);
                    match movee {
                        Move::LogMove { .. } => {}
                        Move::ModelMove { transition, .. }
                        | Move::SynchronousMove { transition, .. }
                        | Move::SilentMove { transition, .. } => {
                            for token in self.transition_2_produced_tokens(*transition, &marking)? {
                                // println!("\ttoken {:?}", token);
                                if let Token::SequenceFlow(sequence_flow_index) = token {
                                    let sequence_flow = self
                                        .global_index_2_sequence_flow_mut(sequence_flow_index)
                                        .ok_or_else(|| anyhow!("sequence flow not found"))?;

                                    *sequence_flow
                                        .weight
                                        .get_or_insert_with(|| Fraction::zero()) += probability;
                                }
                            }

                            self.execute_transition(&mut marking, *transition)?;
                        }
                    }
                }
            }
        }

        self.try_into().with_context(|| anyhow!("Creating SBPMN."))
    }
}

impl AlignmentMiner for LabelledPetriNet {
    type T = StochasticLabelledPetriNet;

    fn mine_stochastic_alignment(
        mut self,
        language: Box<dyn EbiTraitFiniteStochasticLanguage>,
    ) -> Result<StochasticLabelledPetriNet> {
        if self.initial_marking.is_none() {
            return Ok(StochasticLabelledPetriNet::new_empty_language());
        }

        let mut probabilities: Vec<Fraction> =
            vec![Fraction::zero(); self.get_number_of_transitions()];

        let alignments = self.align_stochastic_language(language)?;
        for index in 0..alignments.len() {
            let probability = alignments
                .get_probability(index)
                .ok_or_else(|| anyhow!("should not happen"))?;

            for movee in alignments
                .get(index)
                .ok_or_else(|| anyhow!("should not happen"))?
            {
                match movee {
                    Move::LogMove { .. } => {}
                    Move::ModelMove { transition, .. }
                    | Move::SynchronousMove { transition, .. }
                    | Move::SilentMove { transition } => {
                        probabilities[*transition] += probability;
                    }
                }
            }
        }

        Ok((self, probabilities).into())
    }
}

#[cfg(test)]
mod tests {
    use std::fs;

    use ebi_objects::{
        BusinessProcessModelAndNotation, FiniteStochasticLanguage, HasActivityKey,
        LabelledPetriNet, TranslateActivityKey, ebi_bpmn::Token,
    };

    use crate::techniques::alignment_stochastic_miner::AlignmentMiner;

    #[test]
    fn lpn_uniform() {
        let fin = fs::read_to_string("testfiles/aa-ab-ba.lpn").unwrap();
        let lpn = Box::new(fin.parse::<LabelledPetriNet>().unwrap());

        let fin2 = fs::read_to_string("testfiles/ba.slang").unwrap();
        let slang: Box<FiniteStochasticLanguage> =
            Box::new(fin2.parse::<FiniteStochasticLanguage>().unwrap());

        let slpn = lpn.mine_stochastic_alignment(slang).unwrap();
        let fout = fs::read_to_string("testfiles/aa-ab-ba_ali.slpn").unwrap();

        assert_eq!(fout, slpn.to_string())
    }

    #[test]
    fn bpmn_alignment_estimator() {
        let fin2 = fs::read_to_string("testfiles/aa.slang").unwrap();
        let mut lang = Box::new(fin2.parse::<FiniteStochasticLanguage>().unwrap());

        let fin1 = fs::read_to_string("testfiles/flower.bpmn").unwrap();
        let mut bpmn = fin1.parse::<BusinessProcessModelAndNotation>().unwrap();
        bpmn.translate_using_activity_key(lang.activity_key_mut());

        {
            let mut marking = bpmn.get_initial_marking().unwrap().unwrap();

            println!("marking: {}", marking);

            assert!(!bpmn.is_final_marking(&marking).unwrap());
            assert_eq!(bpmn.get_enabled_transitions(&marking).unwrap(), vec![0]);
            assert_eq!(
                bpmn.transition_2_produced_tokens(0, &marking).unwrap(),
                vec![Token::SequenceFlow((7, ()))]
            );
            bpmn.execute_transition(&mut marking, 0).unwrap();

            println!("marking: {}", marking);

            assert!(!bpmn.is_final_marking(&marking).unwrap());
            assert_eq!(bpmn.get_enabled_transitions(&marking).unwrap(), vec![2]);
            assert_eq!(
                bpmn.transition_2_produced_tokens(2, &marking).unwrap(),
                vec![Token::SequenceFlow((9, ()))]
            );
            bpmn.execute_transition(&mut marking, 2).unwrap();
        }

        let _sbpmn = bpmn.mine_stochastic_alignment(lang).unwrap();
        // dbg!(sbpmn);
    }
}
