use crate::{
    ebi_framework::displayable::Displayable,
    semantics::{labelled_petri_net_semantics::LPNMarking, semantics::Semantics},
};
use ebi_objects::{
    AutomatonSemantics, AutomatonState, DeterministicFiniteAutomaton, DirectlyFollowsGraph,
    DirectlyFollowsModel, LabelledPetriNet, ProcessTree, StochasticDeterministicFiniteAutomaton,
    StochasticDirectlyFollowsModel, StochasticLabelledPetriNet,
    StochasticNondeterministicFiniteAutomaton,
    anyhow::{Result, anyhow},
    ebi_objects::process_tree::TreeMarking,
};
use pastey::paste;
use std::collections::{HashMap, HashSet, hash_map::Entry};

pub trait IsPartOfLivelock {
    type LivState: Displayable;

    /**
     * If more than a couple of states need to be checked for livelocks, then use the livelock cache instead.
     */
    fn is_state_part_of_livelock(&self, state: &Self::LivState) -> Result<bool>;

    /**
     * If only a few livelocks need to be checked, then individual calls to is_state_part_of_livelock may be faster.
     * If more than a few livelock calls need to be made, use this cache object.
     */
    fn get_livelock_cache(&self) -> Box<dyn LiveLockCache<LivState = Self::LivState> + '_>;
}

pub trait LiveLockCache {
    type LivState;

    fn is_state_part_of_livelock(&mut self, state: &Self::LivState) -> Result<bool>;
}

impl IsPartOfLivelock for ProcessTree {
    type LivState = TreeMarking;

    fn is_state_part_of_livelock(&self, _state: &Self::LivState) -> Result<bool> {
        Ok(false)
    }

    fn get_livelock_cache(&self) -> Box<dyn LiveLockCache<LivState = Self::LivState> + '_> {
        Box::new(LiveLockCacheProcessTree {})
    }
}

pub struct LiveLockCacheProcessTree {}

impl LiveLockCache for LiveLockCacheProcessTree {
    type LivState = TreeMarking;

    fn is_state_part_of_livelock(&mut self, _: &Self::LivState) -> Result<bool> {
        Ok(false)
    }
}

macro_rules! lpn {
    ($t:ident, $u:ident) => {
        impl IsPartOfLivelock for $t {
            type LivState = LPNMarking;

            fn is_state_part_of_livelock(&self, state: &Self::LivState) -> Result<bool> {
                //for now, the following only works if the model is bounded
                let mut queue = vec![];
                queue.push(state.clone());
                let mut visited = HashSet::new();
                visited.insert(state.clone());

                while let Some(state) = queue.pop() {
                    if self.is_final_state(&state) {
                        return Ok(false);
                    }

                    for transition in self.get_enabled_transitions(&state) {
                        let mut child_state = state.clone();
                        self.execute_transition(&mut child_state, transition)?;

                        if visited.insert(child_state.clone()) {
                            queue.push(child_state);
                        }
                    }
                }

                return Ok(true);
            }

            fn get_livelock_cache(&self) -> Box<dyn LiveLockCache<LivState = Self::LivState> + '_> {
                Box::new($u::new(&self))
            }
        }

        pub struct $u<'a>(&'a $t, HashMap<LPNMarking, bool>, bool);

        impl<'a> $u<'a> {
            pub fn new(net: &'a $t) -> Self {
                //if there is a transition without an input place, the net is always livelocked.
                let trivial_livelock = net
                    .transition2input_places
                    .iter()
                    .any(|input_places| input_places.is_empty());

                Self(net, HashMap::new(), trivial_livelock)
            }
        }

        impl<'a> LiveLockCache for $u<'a> {
            type LivState = LPNMarking;

            fn is_state_part_of_livelock(&mut self, state: &Self::LivState) -> Result<bool> {
                if self.2 {
                    return Ok(true);
                }

                match self.1.entry(state.clone()) {
                    Entry::Occupied(occupied_entry) => Ok(*occupied_entry.get()),
                    Entry::Vacant(vacant_entry) => {
                        let answer = self.0.is_state_part_of_livelock(state)?;
                        vacant_entry.insert(answer);
                        Ok(answer)
                    }
                }
            }
        }
    };
}

macro_rules! aut {
    ($t:ident) => {
        paste! {
        impl IsPartOfLivelock for $t {
            type LivState = AutomatonState;

            fn is_state_part_of_livelock(&self, state: &Self::LivState) -> Result<bool> {
                let mut queue = vec![];
                queue.push(state.clone());
                let mut visited = HashSet::new();
                visited.insert(state.clone());

                while let Some(state) = queue.pop() {
                    if self.is_final_state(&state) {
                        return Ok(false);
                    }

                    for transition in self.outgoing_transitions(state) {
                        let child_state = self
                            .transition_2_target(transition)
                            .ok_or_else(|| anyhow!("Transition does not exist."))?;

                        if visited.insert(child_state.clone()) {
                            queue.push(child_state);
                        }
                    }
                }

                return Ok(true);
            }

            fn get_livelock_cache(&self) -> Box<dyn LiveLockCache<LivState = Self::LivState>> {
                Box::new([< $t LivelockCache >]::new(self))
            }
        }

        pub struct [< $t LivelockCache >](Vec<bool>);

        impl [< $t LivelockCache >] {
            pub fn new(automaton: &$t) -> Self {
                let mut result = vec![true; AutomatonSemantics::number_of_states(automaton)];
                let mut result_last = vec![true; AutomatonSemantics::number_of_states(automaton)];

                //stage 1: set final states
                for state in automaton.states() {
                    if automaton.is_final_state(&state) {
                        result[state] = false;
                    }
                }

                //stage 2: waterfall
                while result != result_last {
                    result_last.clone_from(&result);

                    for (_, source, target, _) in automaton.transitions() {
                        if !result[target] {
                            result[source] = false;
                        }
                    }
                }

                Self(result)
            }
        }

        impl LiveLockCache for [< $t LivelockCache >] {
            type LivState = AutomatonState;

            fn is_state_part_of_livelock(&mut self, state: &Self::LivState) -> Result<bool> {
                Ok(self.0[state])
            }
        }
    }
    };
}

lpn!(LabelledPetriNet, LiveLockCacheLabelledPetriNet);
lpn!(
    StochasticLabelledPetriNet,
    LiveLockCacheStochasticLabelledPetriNet
);
aut!(DeterministicFiniteAutomaton);
aut!(StochasticDeterministicFiniteAutomaton);
aut!(StochasticNondeterministicFiniteAutomaton);
aut!(StochasticDirectlyFollowsModel);
aut!(DirectlyFollowsModel);
aut!(DirectlyFollowsGraph);

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

    use crate::{semantics::semantics::Semantics, techniques::livelock::IsPartOfLivelock};
    use ebi_objects::{
        DeterministicFiniteAutomaton, DirectlyFollowsModel, LabelledPetriNet,
        StochasticDeterministicFiniteAutomaton, StochasticLabelledPetriNet,
    };

    #[test]
    fn livelock_lpn() {
        let fin = fs::read_to_string("testfiles/a-a-livelock.lpn").unwrap();
        let object = fin.parse::<LabelledPetriNet>().unwrap();

        assert!(
            object
                .get_livelock_cache()
                .is_state_part_of_livelock(&object.get_initial_state().unwrap())
                .unwrap()
        );
    }

    #[test]
    fn livelock_slpn() {
        let fin = fs::read_to_string("testfiles/infinite_bs.slpn").unwrap();
        let object = fin.parse::<StochasticLabelledPetriNet>().unwrap();

        assert!(
            !object
                .get_livelock_cache()
                .is_state_part_of_livelock(&object.get_initial_state().unwrap())
                .unwrap()
        );
    }

    #[test]
    fn livelock_dfa() {
        let fin = fs::read_to_string("testfiles/a-loop.dfa").unwrap();
        let object = fin.parse::<DeterministicFiniteAutomaton>().unwrap();

        let state = object.get_initial_state().unwrap();
        assert!(!object.is_state_part_of_livelock(&state).unwrap());

        assert!(
            !object
                .get_livelock_cache()
                .is_state_part_of_livelock(&state)
                .unwrap()
        );
    }

    #[test]
    fn livelock_sdfa() {
        let fin = fs::read_to_string("testfiles/a-loop.sdfa").unwrap();
        let object = fin
            .parse::<StochasticDeterministicFiniteAutomaton>()
            .unwrap();

        let state = object.get_initial_state().unwrap();
        assert!(!object.is_state_part_of_livelock(&state).unwrap());

        assert!(
            !object
                .get_livelock_cache()
                .is_state_part_of_livelock(&state)
                .unwrap()
        );
    }

    #[test]
    fn livelock_dfm() {
        let fin = fs::read_to_string("testfiles/a-b_star_empty_trace.dfm").unwrap();
        let object = fin.parse::<DirectlyFollowsModel>().unwrap();

        let state = object.get_initial_state().unwrap();
        assert!(!object.is_state_part_of_livelock(&state).unwrap());

        assert!(
            !object
                .get_livelock_cache()
                .is_state_part_of_livelock(&state)
                .unwrap()
        );
    }
}
