//! Network flow algorithms.
//!
//! Ported from classical algorithm sources, adapted to NetworKit's Graph structure.
//!
//! Algorithms:
//! - edmonds_karp: Maximum flow using BFS augmenting paths (Edmonds-Karp)
//! - min_cut: Minimum s-t cut (derived from max-flow via max-flow min-cut theorem)

use super::super::graph::{Graph, NodeId};
use std::collections::{HashMap, HashSet, VecDeque};

// ─────────────────────────────────────────────────────────────────────────────
// Edmonds-Karp (BFS-based Ford-Fulkerson)

/// Maximum flow result.
pub struct FlowResult {
    /// Total flow from source to sink
    pub max_flow: f64,
    /// Flow on each edge: (source, target) → flow
    pub edge_flows: HashMap<(NodeId, NodeId), f64>,
}

/// Minimum cut result.
pub struct CutResult {
    /// Minimum cut capacity
    pub min_cut: f64,
    /// Nodes reachable from source after max-flow (source side of cut)
    pub source_side: HashSet<NodeId>,
    /// Nodes not reachable from source (sink side of cut)
    pub sink_side: HashSet<NodeId>,
    /// Cut edges: (source_side_node, sink_side_node, capacity)
    pub cut_edges: Vec<(NodeId, NodeId, f64)>,
}

/// Maximum flow using the Edmonds-Karp algorithm.
///
/// Implements Ford-Fulkerson with BFS augmenting paths, guaranteeing
/// polynomial time O(V * E²).
///
/// Reference:
/// Edmonds, J.; Karp, R.M. (1972). "Theoretical improvements in algorithmic efficiency
/// for network flow problems." Journal of the ACM.
///
/// # Arguments
/// * `graph` - The flow network (edge weights = capacities)
/// * `source` - Source node
/// * `sink` - Sink node
///
/// # Returns
/// `FlowResult` with max flow value and flow on each edge.
///
/// # Example
/// ```ignore
/// let result = edmonds_karp(&graph, source, sink);
/// println!("Max flow: {}", result.max_flow);
/// for ((u, v), flow) in &result.edge_flows {
///     println!("Edge {}->{}: {} units", u, v, flow);
/// }
/// ```
pub fn edmonds_karp(graph: &Graph, source: NodeId, sink: NodeId) -> FlowResult {
    // Build residual capacity map from graph
    let mut capacity: HashMap<(NodeId, NodeId), f64> = HashMap::new();

    for (u, v, weight, _) in graph.edges() {
        let w = weight.unwrap_or(1.0);
        *capacity.entry((u, v)).or_insert(0.0) += w;
        // Ensure reverse edge exists in residual graph (with 0 capacity initially)
        capacity.entry((v, u)).or_insert(0.0);
    }

    let mut flow: HashMap<(NodeId, NodeId), f64> = HashMap::new();
    let mut total_flow = 0.0;

    loop {
        // BFS to find augmenting path
        let path = bfs_augmenting_path(&capacity, &flow, source, sink, graph);

        match path {
            None => break, // No more augmenting paths
            Some(augmenting_path) => {
                // Find bottleneck capacity along path
                let bottleneck = augmenting_path.windows(2)
                    .map(|window| {
                        let (u, v) = (window[0], window[1]);
                        let cap = capacity.get(&(u, v)).copied().unwrap_or(0.0);
                        let f = flow.get(&(u, v)).copied().unwrap_or(0.0);
                        cap - f // Residual capacity
                    })
                    .fold(f64::INFINITY, f64::min);

                if bottleneck <= 0.0 {
                    break;
                }

                // Update flow along path
                for window in augmenting_path.windows(2) {
                    let (u, v) = (window[0], window[1]);
                    *flow.entry((u, v)).or_insert(0.0) += bottleneck;
                    *flow.entry((v, u)).or_insert(0.0) -= bottleneck;
                }

                total_flow += bottleneck;
            }
        }
    }

    // Extract actual edge flows (non-negative flows on original edges)
    let edge_flows: HashMap<(NodeId, NodeId), f64> = flow.into_iter()
        .filter(|(_, f)| *f > 0.0)
        .collect();

    FlowResult {
        max_flow: total_flow,
        edge_flows,
    }
}

/// BFS to find an augmenting path from source to sink in the residual graph.
fn bfs_augmenting_path(
    capacity: &HashMap<(NodeId, NodeId), f64>,
    flow: &HashMap<(NodeId, NodeId), f64>,
    source: NodeId,
    sink: NodeId,
    graph: &Graph,
) -> Option<Vec<NodeId>> {
    let mut visited: HashSet<NodeId> = HashSet::new();
    let mut parent: HashMap<NodeId, NodeId> = HashMap::new();
    let mut queue = VecDeque::new();

    queue.push_back(source);
    visited.insert(source);

    'bfs: while let Some(node) = queue.pop_front() {
        if node == sink {
            break 'bfs;
        }

        // Explore all possible nodes in residual graph
        // (both original neighbors and reverse edges)
        let neighbors: HashSet<NodeId> = graph.out_neighbors(node)
            .iter()
            .map(|n| n.target)
            .chain(graph.in_neighbors(node).iter().map(|n| n.target))
            .collect();

        for next in neighbors {
            if visited.contains(&next) {
                continue;
            }

            let cap = capacity.get(&(node, next)).copied().unwrap_or(0.0);
            let f = flow.get(&(node, next)).copied().unwrap_or(0.0);
            let residual = cap - f;

            if residual > 0.0 {
                parent.insert(next, node);
                visited.insert(next);
                queue.push_back(next);
            }
        }
    }

    if !visited.contains(&sink) {
        return None; // No augmenting path found
    }

    // Reconstruct path
    let mut path = vec![sink];
    let mut current = sink;
    while current != source {
        let p = *parent.get(&current)?;
        path.push(p);
        current = p;
    }
    path.reverse();
    Some(path)
}

/// Minimum s-t cut derived from maximum flow (max-flow min-cut theorem).
///
/// After computing max flow, the minimum cut consists of edges from the
/// source-reachable set to the sink-reachable set in the residual graph.
///
/// # Returns
/// `CutResult` with min cut value, source/sink partition, and cut edges.
///
/// # Example
/// ```ignore
/// let cut = min_cut(&graph, source, sink);
/// println!("Min cut value: {}", cut.min_cut);
/// println!("Cut edges: {:?}", cut.cut_edges);
/// ```
pub fn min_cut(graph: &Graph, source: NodeId, sink: NodeId) -> CutResult {
    // Run max flow first
    let flow_result = edmonds_karp(graph, source, sink);

    // Build residual graph capacities
    let mut capacity: HashMap<(NodeId, NodeId), f64> = HashMap::new();
    for (u, v, weight, _) in graph.edges() {
        let w = weight.unwrap_or(1.0);
        *capacity.entry((u, v)).or_insert(0.0) += w;
        capacity.entry((v, u)).or_insert(0.0);
    }

    // Find source-reachable nodes in residual graph after max flow
    let mut source_side: HashSet<NodeId> = HashSet::new();
    let mut queue = VecDeque::new();
    queue.push_back(source);
    source_side.insert(source);

    while let Some(node) = queue.pop_front() {
        let neighbors: HashSet<NodeId> = graph.out_neighbors(node)
            .iter()
            .map(|n| n.target)
            .chain(graph.in_neighbors(node).iter().map(|n| n.target))
            .collect();

        for next in neighbors {
            if source_side.contains(&next) {
                continue;
            }
            let cap = capacity.get(&(node, next)).copied().unwrap_or(0.0);
            let f = flow_result.edge_flows.get(&(node, next)).copied().unwrap_or(0.0);
            let residual = cap - f;
            if residual > 0.0 {
                source_side.insert(next);
                queue.push_back(next);
            }
        }
    }

    // Sink side = all nodes not in source side
    let sink_side: HashSet<NodeId> = graph.nodes()
        .filter(|n| !source_side.contains(n))
        .collect();

    // Find cut edges
    let mut cut_edges = Vec::new();
    let mut cut_capacity = 0.0;

    for (u, v, weight, _) in graph.edges() {
        if source_side.contains(&u) && sink_side.contains(&v) {
            let cap = weight.unwrap_or(1.0);
            cut_edges.push((u, v, cap));
            cut_capacity += cap;
        }
    }

    CutResult {
        min_cut: cut_capacity,
        source_side,
        sink_side,
        cut_edges,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::super::graph::GraphConfig;

    fn make_flow_network() -> (Graph, NodeId, NodeId) {
        let mut g = Graph::new(GraphConfig::directed().with_edge_index());
        let s = g.add_node(); // source
        let a = g.add_node();
        let b = g.add_node();
        let c = g.add_node();
        let t = g.add_node(); // sink

        // Classic max-flow example
        //   s -10-> a -8-> t
        //   s -5 -> b -7-> t
        //   a -3 -> b
        g.add_edge(s, a, Some(10.0));
        g.add_edge(s, b, Some(5.0));
        g.add_edge(a, t, Some(8.0));
        g.add_edge(b, t, Some(7.0));
        g.add_edge(a, b, Some(3.0));
        let _ = c; // Unused in this topology

        (g, s, t)
    }

    #[test]
    fn test_edmonds_karp_basic() {
        let (g, source, sink) = make_flow_network();
        let result = edmonds_karp(&g, source, sink);

        // Max flow should be 15 (limited by sink capacity 8+7=15)
        assert_eq!(result.max_flow, 15.0);
    }

    #[test]
    fn test_edmonds_karp_no_path() {
        let mut g = Graph::new(GraphConfig::directed());
        let s = g.add_node();
        let t = g.add_node();
        // No edges
        let result = edmonds_karp(&g, s, t);
        assert_eq!(result.max_flow, 0.0);
    }

    #[test]
    fn test_min_cut() {
        let (g, source, sink) = make_flow_network();
        let cut = min_cut(&g, source, sink);

        // Min cut should equal max flow (max-flow min-cut theorem)
        assert_eq!(cut.min_cut, 15.0);

        // Source must be in source_side
        assert!(cut.source_side.contains(&source));

        // Sink must be in sink_side
        assert!(cut.sink_side.contains(&sink));
    }

    #[test]
    fn test_simple_path_flow() {
        let mut g = Graph::new(GraphConfig::directed().with_edge_index());
        let s = g.add_node();
        let m = g.add_node();
        let t = g.add_node();

        g.add_edge(s, m, Some(5.0));
        g.add_edge(m, t, Some(3.0)); // Bottleneck

        let result = edmonds_karp(&g, s, t);
        assert_eq!(result.max_flow, 3.0); // Limited by bottleneck
    }
}

// ─── Minimum Cost Flow ────────────────────────────────────────────────────────

/// Result of a minimum cost flow computation.
#[derive(Debug, Clone)]
pub struct MinCostFlowResult {
    /// Total flow pushed from source to sink.
    pub flow: f64,
    /// Total cost of the flow.
    pub cost: f64,
    /// Flow on each edge: `flow_on[(u, v)]`.
    pub flow_on: HashMap<(NodeId, NodeId), f64>,
}

/// Compute minimum cost maximum flow using the Successive Shortest Paths algorithm.
///
/// Each edge `(u, v)` must have a non-negative weight (used as cost per unit of flow)
/// and a capacity.  The function pushes `demand` units of flow from `source` to `sink`
/// at minimum cost.  If `demand` is `None`, the full max-flow is computed.
///
/// Uses Bellman-Ford to find shortest (minimum cost) augmenting paths so that
/// negative-cost edges are handled correctly.
///
/// Runtime: O(demand * n * m) — practical for small-to-medium instances.
///
/// Reference: Ford-Fulkerson + Bellman-Ford successive shortest paths.
pub fn min_cost_flow(
    graph: &Graph,
    source: NodeId,
    sink: NodeId,
    demand: Option<f64>,
) -> MinCostFlowResult {
    let n = graph.upper_node_id_bound() as usize;

    // Build residual graph:
    // Each original edge (u,v,cap,cost) → forward arc + backward arc with -cost
    // We represent edges as: Vec<(from, to, cap, cost, rev_idx)>
    // adj[u] = list of edge indices in `edges`
    struct Arc { to: NodeId, cap: f64, cost: f64, rev: usize }
    let mut adj: Vec<Vec<usize>> = vec![Vec::new(); n];
    let mut arcs: Vec<Arc> = Vec::new();

    let add_arc = |adj: &mut Vec<Vec<usize>>, arcs: &mut Vec<Arc>,
                   u: NodeId, v: NodeId, cap: f64, cost: f64| {
        let fwd = arcs.len();
        let bwd = fwd + 1;
        adj[u as usize].push(fwd);
        adj[v as usize].push(bwd);
        arcs.push(Arc { to: v, cap, cost,  rev: bwd });
        arcs.push(Arc { to: u, cap: 0.0, cost: -cost, rev: fwd });
    };

    for u in graph.nodes() {
        for e in graph.out_neighbors(u).iter() {
            let cap  = e.weight.unwrap_or(1.0) as f64;
            let cost = e.weight.unwrap_or(1.0) as f64; // weight = cost
            add_arc(&mut adj, &mut arcs, u, e.target, cap, cost);
        }
    }

    let mut total_flow = 0.0_f64;
    let mut total_cost = 0.0_f64;
    let max_demand = demand.unwrap_or(f64::INFINITY);

    loop {
        if total_flow >= max_demand { break; }

        // Bellman-Ford shortest path in residual graph (cost = distance)
        let mut dist = vec![f64::INFINITY; n];
        let mut prev_arc: Vec<Option<usize>> = vec![None; n];
        dist[source as usize] = 0.0;

        for _ in 0..n {
            let mut updated = false;
            for u in 0..n {
                if dist[u].is_infinite() { continue; }
                for &idx in &adj[u] {
                    let arc = &arcs[idx];
                    if arc.cap > 1e-9 {
                        let nd = dist[u] + arc.cost;
                        if nd < dist[arc.to as usize] - 1e-12 {
                            dist[arc.to as usize] = nd;
                            prev_arc[arc.to as usize] = Some(idx);
                            updated = true;
                        }
                    }
                }
            }
            if !updated { break; }
        }

        if dist[sink as usize].is_infinite() { break; } // no augmenting path

        // Find bottleneck capacity along the path
        let mut path_cap = max_demand - total_flow;
        let mut cur = sink as usize;
        while cur != source as usize {
            let idx = prev_arc[cur].unwrap();
            path_cap = path_cap.min(arcs[idx].cap);
            cur = arcs[arcs[idx].rev].to as usize;
        }
        if path_cap < 1e-9 { break; }

        // Augment
        let mut cur = sink as usize;
        while cur != source as usize {
            let idx = prev_arc[cur].unwrap();
            let rev = arcs[idx].rev;
            arcs[idx].cap -= path_cap;
            arcs[rev].cap += path_cap;
            cur = arcs[rev].to as usize;
        }

        total_flow += path_cap;
        total_cost += path_cap * dist[sink as usize];
    }

    // Reconstruct flow on original edges
    let mut flow_on: HashMap<(NodeId, NodeId), f64> = HashMap::new();
    for u in graph.nodes() {
        for &idx in &adj[u as usize] {
            let arc = &arcs[idx];
            // Original forward arcs have non-negative cost and were added first (even indices)
            // We identify them by idx being even AND the reverse arc has 0 original capacity
            let orig_cap = arcs[arc.rev].cap; // backward arc cap = flow pushed
            if orig_cap > 1e-9 && arc.cost >= 0.0 {
                flow_on.insert((u, arc.to), orig_cap);
            }
        }
    }

    MinCostFlowResult { flow: total_flow, cost: total_cost, flow_on }
}

#[cfg(test)]
mod min_cost_flow_tests {
    use super::*;
    use crate::graph::backends::networkit_rust::graph::GraphConfig;

    #[test]
    fn test_mcf_simple_path() {
        // 0→1→2, capacity and cost 1.0 on each edge; demand 1 unit
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(1, 2, Some(1.0));
        let r = min_cost_flow(&g, 0, 2, Some(1.0));
        assert!((r.flow - 1.0).abs() < 1e-6, "flow={}", r.flow);
        assert!(r.cost > 0.0, "cost should be positive");
    }

    #[test]
    fn test_mcf_no_path() {
        // No edge from 0 to 2
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        let r = min_cost_flow(&g, 0, 2, Some(1.0));
        assert_eq!(r.flow, 0.0);
    }

    #[test]
    fn test_mcf_two_parallel_paths() {
        // Two paths: 0→1→3 (cost 2) and 0→2→3 (cost 2); max demand 2
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(1, 3, Some(1.0));
        g.add_edge(0, 2, Some(1.0));
        g.add_edge(2, 3, Some(1.0));
        let r = min_cost_flow(&g, 0, 3, Some(2.0));
        assert!((r.flow - 2.0).abs() < 1e-6, "flow={}", r.flow);
    }
}
