import org.jgrapht.graph.DefaultWeightedEdge;
import org.jgrapht.graph.SimpleWeightedGraph;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;

public abstract class GraphModificationFunctions {

    //remove over- and under-weight edges
    public static List<Integer[]> filterByOverlapThresholds(SimpleWeightedGraph<Integer, DefaultWeightedEdge> graph,
                                                            int low, int high) {
        List<Integer[]> removedEdges = new ArrayList<>();
        for(DefaultWeightedEdge e: graph.edgeSet()){
            if ((graph.getEdgeWeight(e) > high) || (graph.getEdgeWeight(e) < low)){
                Integer source = graph.getEdgeSource(e);
                Integer target = graph.getEdgeTarget(e);
                Integer weight = (int) graph.getEdgeWeight(e);
                Integer[] edge = {source, target, weight};
                removedEdges.add(edge);
            }
        }
        for (Integer[] edge : removedEdges) {
            graph.removeEdge(edge[0], edge[1]);
        }
        return removedEdges;
    }

    //Remove edges for pairs with large occupancy discrepancy
    public static List<Integer[]> filterByRelativeOccupancy(SimpleWeightedGraph<Integer, DefaultWeightedEdge> graph,
                                                  Map<Integer, Integer> alphaWellCounts,
                                                  Map<Integer, Integer> betaWellCounts,
                                                  Map<Integer, Integer> plateVtoAMap,
                                                  Map<Integer, Integer> plateVtoBMap,
                                                  Integer maxOccupancyDifference) {
        List<Integer[]> removedEdges = new ArrayList<>();
        for (DefaultWeightedEdge e : graph.edgeSet()) {
            Integer alphaOcc = alphaWellCounts.get(plateVtoAMap.get(graph.getEdgeSource(e)));
            Integer betaOcc = betaWellCounts.get(plateVtoBMap.get(graph.getEdgeTarget(e)));
            if (Math.abs(alphaOcc - betaOcc) >= maxOccupancyDifference) {
                Integer source = graph.getEdgeSource(e);
                Integer target = graph.getEdgeTarget(e);
                Integer weight = (int) graph.getEdgeWeight(e);
                Integer[] edge = {source, target, weight};
                removedEdges.add(edge);
            }
        }
        for (Integer[] edge : removedEdges) {
            graph.removeEdge(edge[0], edge[1]);
        }
        return removedEdges;
    }

    //Remove edges for pairs where overlap size is significantly lower than the well occupancy
    public static List<Integer[]> filterByOverlapPercent(SimpleWeightedGraph<Integer, DefaultWeightedEdge> graph,
                                                         Map<Integer, Integer> alphaWellCounts,
                                                         Map<Integer, Integer> betaWellCounts,
                                                         Map<Integer, Integer> plateVtoAMap,
                                                         Map<Integer, Integer> plateVtoBMap,
                                                         Integer minOverlapPercent) {
        List<Integer[]> removedEdges = new ArrayList<>();
        for (DefaultWeightedEdge e : graph.edgeSet()) {
            Integer alphaOcc = alphaWellCounts.get(plateVtoAMap.get(graph.getEdgeSource(e)));
            Integer betaOcc = betaWellCounts.get(plateVtoBMap.get(graph.getEdgeTarget(e)));
            double weight = graph.getEdgeWeight(e);
            double min = minOverlapPercent / 100.0;
            if ((weight / alphaOcc < min) || (weight / betaOcc < min)) {
                Integer source = graph.getEdgeSource(e);
                Integer target = graph.getEdgeTarget(e);
                Integer intWeight = (int) graph.getEdgeWeight(e);
                Integer[] edge = {source, target, intWeight};
                removedEdges.add(edge);
            }
        }
        for (Integer[] edge : removedEdges) {
            graph.removeEdge(edge[0], edge[1]);
        }
        return removedEdges;
    }

    public static void addRemovedEdges(SimpleWeightedGraph<Integer, DefaultWeightedEdge> graph,
                                       List<Integer[]> removedEdges) {
        for (Integer[] edge : removedEdges) {
            DefaultWeightedEdge e = graph.addEdge(edge[0], edge[1]);
            graph.setEdgeWeight(e, (double) edge[2]);
        }
    }

}