use egui_graphs::{DefaultNodeShape, Graph, default_edge_transform, default_node_transform, to_graph_custom};
use petgraph::{stable_graph::{NodeIndex, StableGraph}, visit::{VisitMap, Visitable}, Directed};
use egui::Pos2;
use crate::layout::CustomEdgeShape;
use egui::Color32;

pub trait ResidualGraph {
    fn residual(&self) -> StableGraph<(f32, f32), (u64, u64)>;
}

impl ResidualGraph for StableGraph<(f32, f32), (u64, u64)> {
    fn residual(&self) -> StableGraph<(f32, f32), (u64, u64)> {
        let mut res_graph = self.clone();

        for edge in res_graph.edge_indices().map(|e| e).collect::<Vec<_>>() {
            let (source, target) = res_graph.edge_endpoints(edge).expect("edge not connected");
            if let Some(reverse_edge) = res_graph.find_edge(target, source) {
                let (flow, capacity) = *res_graph.edge_weight(edge).expect("forward edge not found");
                let (reverse_flow, reverse_capacity) = *res_graph.edge_weight(reverse_edge).expect("reverse edge not found");
                // update the forward edge
                // TODO: this seems to overflow in certain cases

                res_graph.update_edge(source, target, (flow, capacity + reverse_flow));
                // update the reverse edge
                //res_graph.update_edge(target, source, (reverse_flow, reverse_capacity + flow));
            } else {
                // add a residual edge with a flow of 0 if the reverse edge doesn't exist
                res_graph.add_edge(target, source, (0, 0));
            }
        }
        res_graph
    }
}

pub trait GuiGraph {
    fn to_gui_graph(&self, s: NodeIndex, t: NodeIndex, active_flows_only: bool) -> Graph<(f32, f32), (u64, u64), Directed, u32, DefaultNodeShape, CustomEdgeShape>;
}

impl GuiGraph for StableGraph<(f32, f32), (u64, u64)> {
    fn to_gui_graph(&self, s: NodeIndex, t: NodeIndex, active_flows_only: bool) -> Graph<(f32, f32), (u64, u64), Directed, u32, DefaultNodeShape, CustomEdgeShape> {
        let mut graph = to_graph_custom(&self,
            |n| {
                let (x, y) = *n.payload();
                default_node_transform(n);
                n.set_location(Pos2::new(x, y));
                n.set_label(String::from(""));
            },
            |e| {
                let (flow, capacity): (u64, u64) = *e.payload();
                default_edge_transform(e);
                e.set_label(format!("{flow}:{capacity}"));
                if flow > 0 {
                    e.set_selected(true);
                }
            }
        );
        graph.node_mut(s).expect("node index not found").set_label(String::from("s"));
        graph.node_mut(t).expect("node index not found").set_label(String::from("t"));
        graph.node_mut(s).expect("node index not found").set_color(Color32::RED);
        graph.node_mut(t).expect("node index not found").set_color(Color32::GREEN);
        graph
    }
}

pub trait FlowGraph {
    fn filter_empty_flows(&self) -> StableGraph<(f32, f32), (u64, u64)>;
}

impl FlowGraph for StableGraph<(f32, f32), (u64, u64)> {
    fn filter_empty_flows(&self) -> StableGraph<(f32, f32), (u64, u64)> {
        self.filter_map( |_, n| {
            Some(*n)
        },
        |_, &e| {
            match e.0 {
                0 => None,
                _ => Some(e),
            }
        })
    }
}