use maxflow_rs::graph::FlowGraph;
use petgraph::prelude::StableGraph;
use petgraph::stable_graph::NodeIndex;
use maxflow_rs::random_generator::MaxflowProblem;
use maxflow_rs::algorithms::{MaxflowAlgorithm, FordFulkerson, EdmondsKarp, Dinic, GoldbergTarjan};

// Example taken from: https://en.wikipedia.org/wiki/Dinic%27s_algorithm#Example
fn generate_small_maxflow_example() -> (StableGraph<(f32, f32), (u64, u64)>, StableGraph<(f32, f32), (u64, u64)>, NodeIndex, NodeIndex) {
    let mut g: StableGraph<(f32, f32), (u64, u64)> = StableGraph::new();
    let s = g.add_node((0., 5.));
    let one = g.add_node((10., 0.));
    let two = g.add_node((10., 10.));
    let three = g.add_node((20., 0.));
    let four = g.add_node((20., 10.));
    let t = g.add_node((30., 5.));
    g.add_edge(s, one, (0, 10));
    g.add_edge(s, two, (0, 10));
    g.add_edge(one, three, (0, 4));
    g.add_edge(one, four, (0, 8));
    g.add_edge(one, two, (0,2));
    g.add_edge(two, four, (0,9));
    g.add_edge(three, t, (0,10));
    g.add_edge(four, three, (0,6));
    g.add_edge(four, t, (0,10));

    // TODO:
    // let mut m: StableGraph<(f32, f32), (u64, u64)> = StableGraph::new();
    // let s = m.add_node((0., 5.));
    // let one = m.add_node((10., 0.));
    // let two = m.add_node((10., 10.));
    // let three = m.add_node((20., 0.));
    // let four = m.add_node((20., 10.));
    // let t = m.add_node((30., 5.));
    // m.add_edge(s, one, (10, 10));
    // m.add_edge(s, two, (9, 10));
    // m.add_edge(one, three, (4, 4));
    // m.add_edge(one, four, (6, 8));
    // m.add_edge(one, two, (0,2));
    // m.add_edge(two, four, (9,9));
    // m.add_edge(three, t, (9,10));
    // m.add_edge(four, three, (5,6));
    // m.add_edge(four, t, (10,10));

    let mut m = g.clone();
    m.update_edge(s, one, (10, 10));
    m.update_edge(s, two, (9, 10));
    m.update_edge(one, three, (4, 4));
    m.update_edge(one, four, (6, 8));
    m.update_edge(one, two, (0,2));
    m.update_edge(two, four, (9,9));
    m.update_edge(three, t, (9,10));
    m.update_edge(four, three, (5,6));
    m.update_edge(four, t, (10,10));

    return (g, m, s, t);
}

fn graph_equal(one: StableGraph<(f32, f32), (u64, u64)>, other: StableGraph<(f32, f32), (u64, u64)>) -> bool {
    // ensure all edges have the same weights
    println!("Graph 1: {:?}", one);
    println!("Graph 2: {:?}", other);
    let edges = one.edge_indices().map(|e| e).collect::<Vec<_>>();
    for edge in edges {
        if one.edge_weight(edge).expect("edge index not found") != other.edge_weight(edge).expect("edge index not found") {
            return false;
        }
    }

    // ensure all nodes have the same weights
    let nodes = one.node_indices().map(|n| n).collect::<Vec<_>>();
    for node in nodes {
        if one.node_weight(node).expect("node index not found") != other.node_weight(node).expect("node index not found") {
            return false;
        }
    }

    true
}

#[test]
fn test_ford_fulkerson() {
    let (g, m, s, t) = generate_small_maxflow_example();

    let problem = MaxflowProblem::from(g, s, t);
    let mut  algo = FordFulkerson::from_problem(&problem);
    let solution = algo.run();

    assert!(graph_equal(solution, m.residual()));
}