hypergraphx 0.0.5

use std::{collections::VecDeque, ops::Sub};

// use crate::{
//     data::DataMap,
//     visit::{
//         EdgeCount, EdgeIndexable, IntoEdges, IntoEdgesDirected, NodeCount, NodeIndexable, VisitMap,
//         Visitable,
//     },
// };

use hashbrown::HashSet;

// use super::{EdgeRef, PositiveMeasure};
use crate::{algo::measures::PositiveMeasure, prelude::*};

fn residual_capacity<'a, G: ?Sized, K, F>(
    network: &'a G,
    edge: <G as GraphBasics<'a>>::EdgeIndex,
    vertex: G::NodeIndex,
    flow: K,
    edge_weight: F,
) -> Option<K>
where
    G: DiGraphProperties<'a>,
    K: PositiveMeasure + Sub<Output = K>,
    F: Fn(<G as GraphBasics<'a>>::EdgeIndex) -> K,
{
    if network.source(edge)?.contains(&vertex) {
        // backward edge
        Some(flow)
    } else if network.target(edge)?.contains(&vertex) {
        // forward edge
        return Some(edge_weight(edge) - flow);
    } else {
        // let end_point = NodeIndexable::to_index(&network, vertex);
        panic!("Illegal endpoint.");
    }
}

/// Gets the other endpoint of graph edge, if any, otherwise panics.
fn other_endpoint<'a, G: ?Sized>(
    network: &'a G,
    edge: <G as GraphBasics<'a>>::EdgeIndex,
    vertex: G::NodeIndex,
) -> Option<HashSet<G::NodeIndex>>
where
    G: DiGraphProperties<'a>,
{
    if network.source(edge)?.contains(&vertex) {
        // backward edge
        network.target(edge)
    } else if network.target(edge)?.contains(&vertex) {
        // forward edge
        network.source(edge)
    } else {
        // let end_point = NodeIndexable::to_index(&network, vertex);
        panic!("Illegal endpoint.");
    }
}

/// Tells whether there is an augmented path in the graph
fn has_augmented_path<'a, G: ?Sized, K, F>(
    network: &'a G,
    source: G::NodeIndex,
    destination: G::NodeIndex,
    edge_to: &mut [Option<(G::EdgeIndex, G::NodeIndex)>],
    flows: &[K],
    edge_weight: F,
) -> Option<bool>
where
    G: DiGraphProperties<'a> + GraphBasics<'a> + CommonProperties<'a, Directed>,
    K: Sub<Output = K> + PositiveMeasure,
    F: Fn(<G as GraphBasics<'a>>::EdgeIndex) -> K,
{
    let mut visited = network.visit_map();
    let mut queue = VecDeque::new();
    visited.set(source.into(), true);
    queue.push_back(source);

    while let Some(vertex) = queue.pop_front() {
        let out_edges = network.in_edges(vertex)?;
        let in_edges = network.out_edges(vertex)?;
        for edge in out_edges.union(&in_edges) {
            let nb = other_endpoint(network, *edge, vertex)?;
            // let edge_index: usize = EdgeIndexable::to_index(&network, edge.id());
            for next in nb {
                let residual_cap =
                    residual_capacity(network, *edge, next, flows[(*edge).into()], &edge_weight)?;
                if !visited.contains(next.into()) && (residual_cap > K::zero()) {
                    visited.set(next.into(), true);
                    edge_to[next.into()] = Some((*edge, vertex));
                    if destination == next {
                        return Some(true);
                    }
                    queue.push_back(next);
                }
            }
        }
    }
    Some(false)
}

fn adjust_residual_flow<'a, G: ?Sized, K>(
    network: &'a G,
    edge: <G as GraphBasics<'a>>::EdgeIndex,
    vertex: G::NodeIndex,
    flow: K,
    delta: K,
) -> Option<K>
where
    G: DiGraphProperties<'a>,
    K: Sub<Output = K> + PositiveMeasure,
{
    Some(if network.source(edge)?.contains(&vertex) {
        // backward edge
        flow - delta
    } else if network.target(edge)?.contains(&vertex) {
        // forward edge
        flow + delta
    } else {
        // let end_point = NodeIndexable::to_index(&network, vertex);
        panic!("Illegal endpoint.");
    })
    // if vertex == edge.source() {
    //     // backward edge
    //     flow - delta
    // } else if vertex == edge.target() {
    //     // forward edge
    //     flow + delta
    // } else {
    //     let end_point = NodeIndexable::to_index(&network, vertex);
    //     panic!("Illegal endpoint {}", end_point);
    // }
}

/// \[Generic\] Ford-Fulkerson algorithm.
///
/// Computes the [maximum flow][ff] of a weighted directed graph.
///
/// If it terminates, it returns the maximum flow and also the computed edge flows.
///
/// [ff]: https://en.wikipedia.org/wiki/Ford%E2%80%93Fulkerson_algorithm
///
/// # Example
/// ```rust
/// use petgraph::Graph;
/// use petgraph::algo::ford_fulkerson;
/// // Example from CLRS book
/// let mut graph = Graph::<u8, u8>::new();
/// let source = graph.add_node(0);
/// let _ = graph.add_node(1);
/// let _ = graph.add_node(2);
/// let _ = graph.add_node(3);
/// let _ = graph.add_node(4);
/// let destination = graph.add_node(5);
/// graph.extend_with_edges(&[
///    (0, 1, 16),
///    (0, 2, 13),
///    (1, 2, 10),
///    (1, 3, 12),
///    (2, 1, 4),
///    (2, 4, 14),
///    (3, 2, 9),
///    (3, 5, 20),
///    (4, 3, 7),
///    (4, 5, 4),
/// ]);
/// let (max_flow, _) = ford_fulkerson(&graph, source, destination);
/// assert_eq!(23, max_flow);
/// ```
pub fn ford_fulkerson<'a, G: ?Sized, K, F>(
    network: &'a G,
    source: G::NodeIndex,
    destination: G::NodeIndex,
    edge_weight: F,
) -> Option<(K, Vec<K>)>
where
    F: Fn(<G as GraphBasics<'a>>::EdgeIndex) -> K,
    G: DiGraphProperties<'a> + CommonProperties<'a, Directed>,
    K: Sub<Output = K> + PositiveMeasure,
{
    let mut edge_to = vec![None; network.node_count()];
    let mut flows = vec![K::zero(); network.edge_count()];
    let mut max_flow = K::zero();
    while has_augmented_path(
        network,
        source,
        destination,
        &mut edge_to,
        &flows,
        &edge_weight,
    )? {
        let mut path_flow = K::max();

        // Find the bottleneck capacity of the path
        let mut vertex = destination;
        // let mut vertex_index = NodeIndexable::to_index(&network, vertex);
        while let Some((edge, prev)) = edge_to[vertex.into()] {
            // let edge_index = EdgeIndexable::to_index(&network, edge.id());
            let residual_capacity =
                residual_capacity(network, edge, vertex, flows[edge.into()], &edge_weight)?;
            // Minimum between the current path flow and the residual capacity.
            path_flow = if path_flow > residual_capacity {
                residual_capacity
            } else {
                path_flow
            };
            vertex = prev;
            // vertex_index = NodeIndexable::to_index(&network, vertex);
        }

        // Update the flow of each edge along the path
        let mut vertex = destination;
        // let mut vertex_index = NodeIndexable::to_index(&network, vertex);
        while let Some((edge, prev)) = edge_to[vertex.into()] {
            // let edge_index = EdgeIndexable::to_index(&network, edge.id());
            flows[edge.into()] =
                adjust_residual_flow(network, edge, vertex, flows[edge.into()], path_flow)?;
            vertex = prev;
            // vertex_index = NodeIndexable::to_index(&network, vertex);
        }
        max_flow = max_flow + path_flow;
    }
    Some((max_flow, flows))
}