//! Algorithms for various network related tasks.

use std::{
    collections::{HashSet, VecDeque},
    fmt::Debug,
    hash::Hash,
    marker::PhantomData,
};

use crate::{ArcInfo, arc_storage::ArcStorage};

/// Direction for seaching with a searching algorithm
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
    /// Only include forward arcs
    Forward,
    /// Only include reverse arcs
    Reverse,
    /// Include both forward and reverse arcs
    Undirected,
}

/// Iterator for Depth first traversal of a network starting at node `starting_node`.
///
/// # Example
/// ```rust
/// use spokes::search::{DfsIter, Direction};
/// use spokes::arc_storage::{ForwardAndReverseStar, ArcStorage};
///
/// // A B C D E F G
/// // 0 1 2 3 4 5 6
/// let (a, b, c, d, e, f, g) = (0, 1, 2, 3, 4, 5, 6);
/// let mut arcs = ForwardAndReverseStar::new();
/// arcs.add_arcs([
///     (a, b), (b, a),
///     (a, c), (c, a),
///     (a, e), (e, a),
///     (b, d), (d, b),
///     (b, f), (f, b),
///     (c, g), (g, c),
///     (e, f), (f, e),
/// ]);
///
/// assert_eq!(
///     DfsIter::new(&arcs, &0, Direction::Forward).copied().collect::<Vec<usize>>(),
///     vec![a, e, f, b, d, c, g],
/// );
/// ```
pub struct DfsIter<'a, I, A, ArcStore>
where
    ArcStore: ArcStorage<I, A>,
{
    arcs: &'a ArcStore,
    to_visit: VecDeque<&'a I>,
    labeled: HashSet<&'a I>,
    direction: Direction,
    _phantom_a: PhantomData<A>,
}

impl<'a, I, A, AS> DfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
{
    /// Create a new Depth-First Iterator of the arcs in `arcs` starting from `starting_node`.
    #[must_use]
    pub fn new(arcs: &'a AS, starting_node: &'a I, direction: Direction) -> Self {
        let mut to_visit = VecDeque::new();
        to_visit.push_front(starting_node);
        Self {
            arcs,
            to_visit,
            direction,
            labeled: HashSet::new(),
            _phantom_a: PhantomData,
        }
    }
}

impl<'a, I, A, AS> Iterator for DfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
    I: Hash + Eq,
    A: 'a,
{
    type Item = &'a I;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(v) = self.to_visit.pop_front() {
            if !self.labeled.insert(v) {
                let edges =
                    match self.direction {
                        Direction::Forward => &mut self.arcs.forward_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Reverse => &mut self.arcs.reverse_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Undirected => {
                            &mut self.arcs.all_arcs(v) as &mut dyn Iterator<Item = &ArcInfo<I, A>>
                        }
                    };

                for edge in edges {
                    match self.direction {
                        Direction::Forward => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_front(&edge.head);
                            }
                        }
                        Direction::Reverse => {
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_front(&edge.tail);
                            }
                        }
                        Direction::Undirected => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_front(&edge.head);
                            }
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_front(&edge.tail);
                            }
                        }
                    }
                }
                return Some(v);
            }
        }

        None
    }
}

/// Iterator for Breadth first traversal of a network starting at node `starting_node`.
///
/// # Example
/// ```rust
/// use spokes::search::{BfsIter, Direction};
/// use spokes::arc_storage::{ForwardAndReverseStar, ArcStorage};
///
/// let (a, b, c, d, e, f, g) = (0, 1, 2, 3, 4, 5, 6);
/// let mut arcs = ForwardAndReverseStar::new();
/// arcs.add_arcs([
///     (a, b), (b, a),
///     (a, c), (c, a),
///     (a, e), (e, a),
///     (b, d), (d, b),
///     (b, f), (f, b),
///     (c, g), (g, c),
///     (e, f), (f, e),
/// ]);
///
/// assert_eq!(
///     BfsIter::new(&arcs, &a, Direction::Forward).copied().collect::<Vec<usize>>(),
///     vec![a, b, c, e, d, f, g],
/// );
/// ```
pub struct BfsIter<'a, I, A, ArcStore>
where
    ArcStore: ArcStorage<I, A>,
{
    arcs: &'a ArcStore,
    to_visit: VecDeque<&'a I>,
    labeled: HashSet<&'a I>,
    direction: Direction,
    _phantom_a: PhantomData<A>,
}

impl<'a, I, A, AS> BfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
{
    /// Create a new Breadth-First Iterator of the arcs in `arcs` starting from `starting_node`.
    #[must_use]
    pub fn new(arcs: &'a AS, starting_node: &'a I, direction: Direction) -> Self {
        let mut to_visit = VecDeque::new();
        to_visit.push_front(starting_node);
        Self {
            arcs,
            to_visit,
            labeled: HashSet::new(),
            direction,
            _phantom_a: PhantomData,
        }
    }
}

impl<'a, I, A, AS> Iterator for BfsIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
    I: Hash + Eq,
    A: 'a,
{
    type Item = &'a I;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(v) = self.to_visit.pop_front() {
            if !self.labeled.insert(v) {
                let edges =
                    match self.direction {
                        Direction::Forward => &mut self.arcs.forward_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Reverse => &mut self.arcs.reverse_arcs(v)
                            as &mut dyn Iterator<Item = &ArcInfo<I, A>>,
                        Direction::Undirected => {
                            &mut self.arcs.all_arcs(v) as &mut dyn Iterator<Item = &ArcInfo<I, A>>
                        }
                    };

                for edge in edges {
                    match self.direction {
                        Direction::Forward => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_back(&edge.head);
                            }
                        }
                        Direction::Reverse => {
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_back(&edge.tail);
                            }
                        }
                        Direction::Undirected => {
                            if !self.labeled.contains(&edge.head) {
                                self.to_visit.push_back(&edge.head);
                            }
                            if !self.labeled.contains(&edge.tail) {
                                self.to_visit.push_back(&edge.tail);
                            }
                        }
                    }
                }

                return Some(v);
            }
        }

        None
    }
}

/// Create an iterator over all all non-cycle paths in the network.
pub struct PathIter<'a, I, A, ArcStore>
where
    ArcStore: ArcStorage<I, A>,
{
    arcs: &'a ArcStore,
    curent_path: Vec<&'a I>,
    branches: VecDeque<(usize, &'a I)>,
    seen: HashSet<&'a I>,
    direction: Direction,
    _phantom_a: PhantomData<A>,
}

impl<'a, I, A, AS> PathIter<'a, I, A, AS>
where
    A: 'a,
    AS: ArcStorage<I, A>,
    I: Eq + Hash,
{
    /// Create a new Depth-First Iterator of the arcs in `arcs` starting from `starting_node`.
    #[must_use]
    pub fn new(arcs: &'a AS, starting_node: &'a I, direction: Direction) -> Self {
        let branches = match direction {
            Direction::Forward => arcs
                .forward_arcs(starting_node)
                .map(|arc| (1, arc.head()))
                .collect(),
            Direction::Reverse => arcs
                .reverse_arcs(starting_node)
                .map(|arc| (1, arc.tail()))
                .collect(),
            Direction::Undirected => arcs
                .reverse_arcs(starting_node)
                .map(|arc| (1, arc.tail()))
                .chain(arcs.forward_arcs(starting_node).map(|arc| (1, arc.head())))
                .collect(),
        };

        Self {
            arcs,
            curent_path: vec![starting_node],
            branches,
            direction,
            seen: HashSet::from([starting_node]),
            _phantom_a: PhantomData,
        }
    }
}

impl<'a, I, A, AS> Iterator for PathIter<'a, I, A, AS>
where
    AS: ArcStorage<I, A>,
    I: Hash + Eq,
    A: 'a,
{
    type Item = Vec<&'a I>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some((len, next_node)) = self.branches.pop_front() {
            self.curent_path.drain(len..).for_each(|x| {
                self.seen.remove(x);
            });
            self.curent_path.push(next_node);
            self.seen.insert(next_node);

            match self.direction {
                Direction::Forward => self.arcs.forward_arcs(next_node).for_each(|arc| {
                    if !self.seen.contains(arc.head()) {
                        self.branches.push_front((len + 1, arc.head()));
                    }
                }),
                Direction::Reverse => self.arcs.reverse_arcs(next_node).for_each(|arc| {
                    if !self.seen.contains(arc.tail()) {
                        self.branches.push_front((len + 1, arc.tail()));
                    }
                }),
                Direction::Undirected => {
                    self.arcs.forward_arcs(next_node).for_each(|arc| {
                        if !self.seen.contains(arc.head()) {
                            self.branches.push_front((len + 1, arc.head()));
                        }
                    });
                    self.arcs.reverse_arcs(next_node).for_each(|arc| {
                        if !self.seen.contains(arc.tail()) {
                            self.branches.push_front((len + 1, arc.tail()));
                        }
                    });
                }
            }

            Some(self.curent_path.clone())
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::Network;

    use super::*;

    #[test]
    fn simple_bfs() {
        let mut network: Network<i32, (), ()> = Network::new();
        network.add_nodes((0..4).map(|i| (i, ())));

        network
            .add_arcs([(0, 1, ()), (1, 2, ()), (2, 3, ())])
            .expect("should construct");
        assert_eq!(
            BfsIter::new(network.arc_store(), &0, Direction::Forward)
                .copied()
                .collect::<Vec<i32>>(),
            vec![0, 1, 2, 3]
        );

        assert_eq!(
            BfsIter::new(network.arc_store(), &3, Direction::Reverse)
                .copied()
                .collect::<Vec<i32>>(),
            vec![3, 2, 1, 0]
        );

        assert_eq!(
            BfsIter::new(network.arc_store(), &2, Direction::Undirected)
                .copied()
                .collect::<Vec<i32>>(),
            vec![2, 3, 1, 0]
        );
    }

    #[test]
    fn all_paths() {
        let mut network: Network<u8, (), ()> = Network::new();
        network.add_nodes((0..8).map(|i| (i, ())));

        network
            .add_arcs([
                (0, 1),
                (0, 2),
                (0, 3),
                (1, 4),
                (2, 4),
                (3, 4),
                (0, 5),
                (4, 6),
                (1, 7),
            ])
            .expect("should construct");

        let mut paths: Vec<Vec<u8>> = PathIter::new(network.arc_store(), &0, Direction::Forward)
            .map(|x| x.into_iter().copied().collect())
            .collect();
        paths.sort_by_key(|x| {
            x.iter()
                .enumerate()
                .map(|(i, v)| (*v as usize) << i)
                .sum::<usize>()
        });

        let mut expected = vec![
            vec![0_u8, 1],
            vec![0, 1, 7],
            vec![0, 1, 4],
            vec![0, 1, 4, 6],
            vec![0, 5],
            vec![0, 3],
            vec![0, 3, 4],
            vec![0, 3, 4, 6],
            vec![0, 2],
            vec![0, 2, 4],
            vec![0, 2, 4, 6],
        ];
        expected.sort_by_key(|x| {
            x.iter()
                .enumerate()
                .map(|(i, v)| (*v as usize) << i)
                .sum::<usize>()
        });

        assert_eq!(paths, expected);
    }
}
