spokes 0.4.0-rc.4

use std::{
    collections::{HashMap, VecDeque},
    fmt::Debug,
    hash::Hash,
    iter::Sum,
    ops::{AddAssign, SubAssign},
};

use num::Zero;

use crate::{
    ArcInfo, ArcStorage, Network,
    algorithms::{Distance, next_admissible_arc},
    node_storage::NodeStorage,
    search::Direction,
};

/// Preflow push Algorithm for determing max-flow.
///
/// # Arguments
/// * `network` - A network with capacities on the arcs.
/// * `source` - Id of source node
/// * `sink` - Id of sink node
///
/// # Example
///```rust
///  use spokes::{ArcStorage, Network, algorithms::preflow_push};
///  let mut network: Network<usize, (), usize> = Network::new();
///
///  network.add_nodes((0..=5).map(|i| (i, ())));
///
///  network.add_arcs([
///      (0, 1, 15),
///      (0, 3, 4),
///      (1, 2, 12),
///      (2, 3, 3),
///      (2, 5, 7),
///      (3, 4, 10),
///      (4, 1, 5),
///      (4, 5, 10),
///  ]);
///
///  let (max_flow_value, max_flows) = preflow_push(&network, &0, &5);
///
///  assert_eq!(max_flow_value, 14);
///  assert_eq!(max_flows.arc(&0, &1), Some(&10));
///  assert_eq!(max_flows.arc(&0, &3), Some(&4));
///  assert_eq!(max_flows.arc(&1, &2), Some(&10));
///  assert_eq!(max_flows.arc(&2, &3), Some(&3));
///  assert_eq!(max_flows.arc(&2, &5), Some(&7));
///  assert_eq!(max_flows.arc(&3, &4), Some(&7));
///  assert_eq!(max_flows.arc(&4, &1), None);
///  assert_eq!(max_flows.arc(&4, &5), Some(&7));
///  ```
///
/// # Panics
/// This will panic if the `ArcStore` or `NodeStore` are broken.
pub fn preflow_push<I, T, NodeStore, ArcStore>(
    network: &Network<I, (), T, NodeStore, ArcStore>,
    source: &I,
    sink: &I,
) -> (T, Network<I, (), T, NodeStore, ArcStore>)
where
    ArcStore:
        ArcStorage<I, T> + FromIterator<ArcInfo<I, T>> + Clone + IntoIterator<Item = ArcInfo<I, T>>,
    NodeStore: NodeStorage<I, ()> + Clone + IntoIterator<Item = (I, ())> + FromIterator<(I, ())>,
    I: Hash + Eq + Copy,
    T: Zero + Sum + Copy + Ord + Eq + AddAssign + SubAssign,
{
    let n = network.n_nodes();
    let mut flows: Network<I, PfpNode<T>, T, HashMap<I, PfpNode<T>>, ArcStore> = (*network)
        .clone()
        .map_nodes(|(idx, _attrs)| {
            if &idx == source {
                (idx, PfpNode::new(Distance::Finite(n), Distance::Infinite))
            } else {
                (idx, PfpNode::new(Distance::Infinite, Distance::zero()))
            }
        })
        .expect("By construction, this is valid");

    // Assign distances using bfs
    flows.node_mut(sink).expect("Should exist").distance_label = Distance::zero();

    let bfs_nodes: Vec<I> = flows
        .bfs(sink, Direction::Reverse)
        .skip(1)
        .copied()
        .collect();
    for node in bfs_nodes {
        let distance_to_node = flows
            .forward_arcs(&node)
            .filter_map(|a| flows.node(&a.head))
            .map(PfpNode::label)
            .map(|d| d + 1)
            .min()
            .unwrap_or(Distance::Infinite);
        flows.node_mut(&node).expect("Should exist").distance_label = distance_to_node;
    }

    // Set source to a distance of n nodes
    flows
        .node_mut(source)
        .expect("Node should exist")
        .distance_label = Distance::Finite(flows.n_nodes());

    let mut next_nodes: VecDeque<I> = VecDeque::new();

    // Push saturating flows out of source and add destitions to next_nodes
    let arcs_from_source: Vec<ArcInfo<I, T>> = flows.forward_arcs(source).cloned().collect();
    for arc in arcs_from_source {
        pfp_push(&mut flows, &arc, &mut next_nodes, sink, source);
    }

    while let Some(node_id) = next_nodes.pop_front() {
        while !flows.node(&node_id).expect("Should exist").excess.is_zero() {
            if let Some(arc) = next_admissible_arc(&flows, PfpNode::label, &node_id).cloned() {
                pfp_push(&mut flows, &arc, &mut next_nodes, sink, source);
            } else {
                // Relabel
                let new_label = flows
                    .forward_arcs(&node_id)
                    .filter_map(|arc| flows.node(&arc.head))
                    .map(|n| n.distance_label + 1)
                    .min()
                    .unwrap_or(Distance::Infinite);

                flows
                    .node_mut(&node_id)
                    .expect("Should exist")
                    .distance_label = new_label;

                if !flows.node(&node_id).expect("Should exist").excess.is_zero() {
                    next_nodes.push_back(node_id);
                }
                break;
            }
        }
    }

    //swap arc direction
    flows = flows
        .map_arcs(|arc| ArcInfo::new(arc.head, arc.tail, arc.attributes))
        .expect("By construction, this is valid");

    // Remove extra arcs not in original network
    let arcs_to_remove: Vec<(I, I)> = flows
        .arc_iter()
        .filter_map(|arc| {
            if network.contains_arc(&arc.tail, &arc.head) {
                None
            } else {
                Some((arc.tail, arc.head))
            }
        })
        .collect();

    flows.remove_arcs(arcs_to_remove.iter().map(|(a, b)| (a, b)));

    // Compute total flow into sink
    let total_flow: T = *flows
        .node(sink)
        .expect("should exist")
        .excess
        .finite_value()
        .expect("Total flow should be finite");

    let flows = flows
        .map_nodes(|(idx, _attrs)| (idx, ()))
        .expect("By construction, this is valid");

    (total_flow, flows)
}

#[inline]
fn pfp_push<I, T, NodeStore, ArcStore>(
    flows: &mut Network<I, PfpNode<T>, T, NodeStore, ArcStore>,
    arc: &ArcInfo<I, T>,
    next_nodes: &mut VecDeque<I>,
    sink: &I,
    source: &I,
) where
    NodeStore: NodeStorage<I, PfpNode<T>>,
    ArcStore: ArcStorage<I, T>,
    I: Hash + Eq + Copy,
    T: Zero + Sum + Copy + Ord + Eq + AddAssign + SubAssign,
{
    // Push flow along arc
    let node = flows.node(&arc.tail).expect("Should exist");
    let delta = *node
        .excess
        .min(Distance::Finite(arc.attributes))
        .finite_value()
        .expect("Capacity is finite so this is finite");

    // Update forward arcs
    *flows
        .arc_mut(&arc.tail, &arc.head)
        .expect("Arc should exist") -= delta;

    // Remove the arc if the capacity is zero
    if flows
        .arc(&arc.tail, &arc.head)
        .expect("Should exist")
        .is_zero()
    {
        flows.remove_arc(&arc.tail, &arc.head);
    }

    // And or update existing reverse arc with flow
    if flows.contains_arc(&arc.head, &arc.tail) {
        *flows
            .arc_mut(&arc.head, &arc.tail)
            .expect("Arc should exist") += delta;
    } else {
        let _ = flows.add_arc(ArcInfo::new(arc.head, arc.tail, delta));
    }

    // Update excesses
    flows.node_mut(&arc.tail).expect("Should exist").excess -= Distance::Finite(delta);
    flows.node_mut(&arc.head).expect("Should exist").excess += Distance::Finite(delta);

    // Add forward node to active nodes unless it's the sink
    if &arc.head != sink && &arc.head != source {
        next_nodes.push_back(arc.head);
    }
}

#[derive(Clone, Debug)]
struct PfpNode<T> {
    pub distance_label: Distance<usize>,
    pub excess: Distance<T>,
}

impl<T> PfpNode<T> {
    const fn new(distance_label: Distance<usize>, excess: Distance<T>) -> Self {
        Self {
            distance_label,
            excess,
        }
    }

    /// Extrac the distance label
    pub const fn label(&self) -> Distance<usize> {
        self.distance_label
    }
}

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

    #[test]
    #[ignore = "Preflow Push isn't working, so we can ignore this test for now"]
    fn example() {
        let mut network: Network<usize, (), usize> = Network::new();

        network.add_nodes((0..=5).map(|i| (i, ())));

        network
            .add_arcs([
                (0, 1, 15),
                (0, 3, 4),
                (1, 2, 12),
                (2, 3, 3),
                (2, 5, 7),
                (3, 4, 10),
                (4, 1, 5),
                (4, 5, 10),
            ])
            .expect("should work");

        let (max_flow_value, max_flows) = preflow_push(&network, &0, &5);

        assert_eq!(max_flow_value, 14);
        assert_eq!(max_flows.arc(&0, &1), Some(&10));
        assert_eq!(max_flows.arc(&0, &3), Some(&4));
        assert_eq!(max_flows.arc(&1, &2), Some(&10));
        assert_eq!(max_flows.arc(&2, &3), Some(&3));
        assert_eq!(max_flows.arc(&2, &5), Some(&7));
        assert_eq!(max_flows.arc(&3, &4), Some(&7));
        assert_eq!(max_flows.arc(&4, &1), None);
        assert_eq!(max_flows.arc(&4, &5), Some(&7));
    }
}