use egui_graphs::{DefaultNodeShape, Graph, default_edge_transform, default_node_transform, to_graph_custom};
use petgraph::{algo::has_path_connecting, stable_graph::{EdgeIndex, NodeIndex, StableGraph}, visit::{EdgeRef, VisitMap, Visitable}, Directed, Direction};
use egui::{Pos2};
use std::collections::{HashMap, HashSet, VecDeque};
use std::cmp::min;
use crate::layout::CustomEdgeShape;
use egui::Color32;

// All methods which convert the graph into a GUI representations
pub trait GuiGraph {
    fn to_gui_graph(&self, s: NodeIndex, t: NodeIndex) -> Graph<(f32, f32), (u64, u64), Directed, u32, DefaultNodeShape, CustomEdgeShape>;
    fn to_gui_graph_distance(&self, s: NodeIndex, t: NodeIndex, distance_labels: HashMap<NodeIndex, u64>) -> 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) -> 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::BLUE);
        graph.node_mut(t).expect("node index not found").set_color(Color32::ORANGE);
        graph
    }

    fn to_gui_graph_distance(&self, s: NodeIndex, t: NodeIndex, distance_labels: HashMap<NodeIndex, u64>) -> 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));
                if let Some(d) = distance_labels.get(&n.id()) {
                    if n.id() == s {
                        n.set_label(format!("s ({})", d));
                    } else if n.id() == t {
                        n.set_label(format!("t ({})", d));
                    } else {
                        n.set_label(format!("{}", d));
                    }
                } else {
                    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_color(Color32::BLUE);
        graph.node_mut(t).expect("node index not found").set_color(Color32::ORANGE);
        graph
    }
}

// all methods specific to MaxFlow problems
pub trait FlowGraph {
    fn filter_empty_flows(&self) -> StableGraph<(f32, f32), (u64, u64)>;

    fn reset_flow(&mut self);

    fn prune_zero(&self) -> StableGraph<(f32, f32), (u64, u64)>;

    fn valid_flow(&self, source: NodeIndex, sink: NodeIndex) -> bool;
    
    fn saturated_cut(&self, source: NodeIndex, sink: NodeIndex) -> bool;

    fn total_flow(&self, source: NodeIndex, sink: NodeIndex) -> u64;

    fn available_capacity(&self, edge: EdgeIndex) -> u64;

    fn layer_graph(&self, source: NodeIndex, sink: NodeIndex) -> StableGraph<(f32, f32), (u64, u64)>;

    fn add_flow(&mut self, edge: EdgeIndex, flow: u64);

    fn residual(&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),
            }
        })
    }

    fn layer_graph(&self, source: NodeIndex, sink: NodeIndex) -> StableGraph<(f32, f32), (u64, u64)> {
        // filter graph for edges with no remaining capacity
        let mut layer_graph = self.filter_map(|_, n| {
            Some(*n)
        }, |_, &e| {
            match e.1 - e.0 {
                0 => None,
                _ => Some(e),
            }
        });

        // filter graph for s-t paths
        let mut visited = layer_graph.visit_map();
        let mut queue = VecDeque::from([(source, 0)]);
        while let Some((node, layer)) = queue.pop_front() {
            //let outgoing_edges = layer_graph.edges_directed(node, Direction::Outgoing).collect::<Vec<_>>();
            let outgoing_edges = layer_graph.edges_directed(node, Direction::Outgoing).map(|e| e.id()).collect::<Vec<_>>();
            visited.visit(node);

            for edge in outgoing_edges {
                let neighbor = layer_graph.edge_endpoints(edge).expect("edge index not found").1;
                // if neighbor is unvisited, this is the shortest path to neighbor -> keep the edge
                if !visited.is_visited(&neighbor) {
                    visited.visit(neighbor);
                    // stop traversing this path if destination is reached
                    if neighbor == sink {
                        continue;
                    } else {
                        // add neighbor to queue if destination is not reached
                        queue.push_back((neighbor, layer+1));
                    }
                } else {
                    // if neighbor was visited before this edge is not on the shortest path -> remove it
                    layer_graph.remove_edge(edge);
                }
            }
        }

        layer_graph
    }

    // set all flows back to zero
    fn reset_flow(&mut self) {
        for edge in self.edge_indices().collect::<Vec<_>>() {
            let weight = self.edge_weight_mut(edge).expect("edge not found");
            (*weight).0 = 0;
        }
    }

    // removes all edges where the current flow is zero & returns a pruned graph
    fn prune_zero(&self) -> StableGraph<(f32, f32), (u64, u64)> {
        let mut zero_graph = self.clone();
        for edge in zero_graph.edge_indices().collect::<Vec<_>>() {
            let weight = zero_graph.edge_weight(edge).expect("edge not found");
            if weight.0 == 0 {
                zero_graph.remove_edge(edge);
            }
        }

        zero_graph
    }

    // check if flow is valid:
    // - for every edge, the flow doesn't exceed the capacity
    // - for every node, the amount of incoming flows is equal to the amount of outgoing flows
    fn valid_flow(&self, source: NodeIndex, sink: NodeIndex) -> bool {
        // capacity isn't exceeded
        for edge in self.edge_indices().map(|e| e).collect::<Vec<_>>() {
            if self.edge_weight(edge).expect("edge index not found").0 > self.edge_weight(edge).expect("edge index not found").1 {
                return false;
            }
        }
        // sum incoming flows = sum outgoing flows
        for node in self.node_indices().map(|n| n).collect::<Vec<_>>() {
            let outgoing_flows: u64 = self
                .edges_directed(node, Direction::Outgoing)
                .map(|e| e.weight().0)
                .sum();
            let incoming_flows: u64 = self
                .edges_directed(node, Direction::Incoming)
                .map(|e| e.weight().0)
                .sum();
            if incoming_flows != outgoing_flows {
                if node != source && node != sink {
                    // println!("invalid flow at {:?} incoming flows: {} outgoing flows: {}", node, incoming_flows, outgoing_flows);
                    return false;
                }
            }
        }
        true
    }

    // checks if there exists a saturated cut in the graph (a bip)
    fn saturated_cut(&self, source: NodeIndex, sink: NodeIndex) -> bool {
        let mut s = HashSet::new();
        let mut t = HashSet::new();

        for node in self.node_indices().into_iter().collect::<Vec<_>>() {
            if has_path_connecting(self, source, node, None) {
                s.insert(node);
            } else {
                t.insert(node);
            }
        }

        return true;
    }

    fn total_flow(&self, source: NodeIndex, sink: NodeIndex) -> u64 {
        let outgoing_flows_source: u64 = self.edges_directed(source, Direction::Outgoing).map(|e| e.weight().0).sum();
        let incoming_flows_source: u64 = self.edges_directed(source, Direction::Incoming).map(|e| e.weight().0).sum();
        let outgoing_flows_sink: u64 = self.edges_directed(sink, Direction::Outgoing).map(|e| e.weight().0).sum();
        let incoming_flows_sink: u64 = self.edges_directed(sink, Direction::Incoming).map(|e| e.weight().0).sum();
        let total_flow_source = outgoing_flows_source - incoming_flows_source;
        let total_flow_sink = incoming_flows_sink - outgoing_flows_sink;
        if total_flow_sink != total_flow_source {
            return 0;
        }
        return total_flow_source;
    }

    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 edge
                res_graph.update_edge(source, target, (flow, capacity + reverse_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
    }
    
    fn available_capacity(&self, edge: EdgeIndex) -> u64 {
        self.edge_weight(edge).expect("edge not found").1 - self.edge_weight(edge).expect("edge not found").0
    }
    
    // adds the specified flow to the edge of a residual graph
    fn add_flow(&mut self, forward_edge: EdgeIndex, flow: u64) {
        // search for reverse edge
        let (source, target) = self.edge_endpoints(forward_edge).expect("edge not connected");
        let reverse_edge = self.find_edge(target, source).expect("reverse edge not found");
        
        // calculate how much flow needs to be added on the reverse edge
        let current_reverse_flow = self.edge_weight(reverse_edge).expect("reverse erge not found").0;
        // maximum flow we can subtract from reverse edge
        let reverse_flow = min(current_reverse_flow, flow);
        let forward_flow = flow - reverse_flow;
        if reverse_flow > 0 {
            // decrease reverse flow
            let reverse_weight: &mut (u64, u64) = self.edge_weight_mut(reverse_edge).expect("reverse edge not found");
            (*reverse_weight).0 -= reverse_flow;
            assert!(reverse_weight.0 <= reverse_weight.1);
            // decrease forward capacity 
            let forward_weight: &mut (u64, u64) = self.edge_weight_mut(forward_edge).expect("edge not found");
            (*forward_weight).1 -= reverse_flow;
            assert!(forward_weight.0 <= forward_weight.1);
        }

        // add remaining flow to forward edge
        if forward_flow > 0 {
            // increase forward flow
            let forward_weight: &mut (u64, u64) = self.edge_weight_mut(forward_edge).expect("edge not found");
            (*forward_weight).0 += forward_flow;
            assert!(forward_weight.0 <= forward_weight.1);
            // increase reverse capacity
            let reverse_weight: &mut (u64, u64) = self.edge_weight_mut(reverse_edge).expect("reverse edge not found");
            (*reverse_weight).1 += forward_flow;
            assert!(reverse_weight.0 <= reverse_weight.1);
        }
    }
}