package org.batfish.datamodel;

import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonValue;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedSet;
import java.io.Serializable;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;
import org.batfish.datamodel.collections.NodeInterfacePair;

public final class Topology implements Serializable {

  private static final long serialVersionUID = 1L;

  @JsonCreator
  private static Topology jacksonCreateTopology(SortedSet<Edge> edges) {
    return new Topology(edges);
  }

  private final SortedSet<Edge> _edges;

  private final Map<NodeInterfacePair, SortedSet<Edge>> _interfaceEdges;

  private final Map<String, SortedSet<Edge>> _nodeEdges;

  public Topology(SortedSet<Edge> edges) {
    _edges = new TreeSet<>(edges);
    _nodeEdges = new HashMap<>();
    _interfaceEdges = new HashMap<>();
    rebuildFromEdges();
  }

  @JsonIgnore
  public SortedSet<Edge> getEdges() {
    return _edges;
  }

  @JsonIgnore
  public Map<NodeInterfacePair, SortedSet<Edge>> getInterfaceEdges() {
    return _interfaceEdges;
  }

  public Set<NodeInterfacePair> getNeighbors(NodeInterfacePair iface) {
    return getInterfaceEdges()
        .getOrDefault(iface, ImmutableSortedSet.of())
        .stream()
        .filter(e -> e.getFirst().equals(iface))
        .map(Edge::getSecond)
        .collect(ImmutableSet.toImmutableSet());
  }

  @JsonIgnore
  public Map<String, SortedSet<Edge>> getNodeEdges() {
    return _nodeEdges;
  }

  /** Removes the specified blacklists from the topology */
  public void prune(
      Set<Edge> blacklistEdges,
      Set<String> blacklistNodes,
      Set<NodeInterfacePair> blacklistInterfaces) {
    if (blacklistEdges != null) {
      SortedSet<Edge> edges = getEdges();
      edges.removeAll(blacklistEdges);
    }
    if (blacklistNodes != null) {
      for (String blacklistNode : blacklistNodes) {
        removeNode(blacklistNode);
      }
    }
    if (blacklistInterfaces != null) {
      for (NodeInterfacePair blacklistInterface : blacklistInterfaces) {
        removeInterface(blacklistInterface);
      }
    }
    rebuildFromEdges();
  }

  private void rebuildFromEdges() {
    _nodeEdges.clear();
    _interfaceEdges.clear();
    for (Edge edge : getEdges()) {
      String node1 = edge.getNode1();
      String node2 = edge.getNode2();
      NodeInterfacePair int1 = edge.getInterface1();
      NodeInterfacePair int2 = edge.getInterface2();

      SortedSet<Edge> node1Edges = _nodeEdges.computeIfAbsent(node1, k -> new TreeSet<>());
      node1Edges.add(edge);

      SortedSet<Edge> node2Edges = _nodeEdges.computeIfAbsent(node2, k -> new TreeSet<>());
      node2Edges.add(edge);

      SortedSet<Edge> interface1Edges = _interfaceEdges.computeIfAbsent(int1, k -> new TreeSet<>());
      interface1Edges.add(edge);

      SortedSet<Edge> interface2Edges = _interfaceEdges.computeIfAbsent(int2, k -> new TreeSet<>());
      interface2Edges.add(edge);
    }
  }

  private void removeInterface(NodeInterfacePair iface) {
    SortedSet<Edge> interfaceEdges = _interfaceEdges.get(iface);
    if (interfaceEdges != null) {
      _edges.removeAll(interfaceEdges);
    }
  }

  private void removeNode(String hostname) {
    SortedSet<Edge> nodeEdges = _nodeEdges.get(hostname);
    if (nodeEdges != null) {
      _edges.removeAll(nodeEdges);
    }
  }

  @JsonValue
  public SortedSet<Edge> sortedEdges() {
    return new TreeSet<>(_edges);
  }
}

@@ -25,16 +25,80 @@
 model;%0A%0A
+import static com.google.common.base.MoreObjects.firstNonNull;%0A%0A
 import c
@@ -559,16 +559,152 @@
 ePair;%0A%0A
+/**%0A * Represents a set of %7B@link Edge Edges%7D and provides methods to prune the edges with edge, node,%0A * and interface blacklists.%0A */%0A
 public f
@@ -914,21 +914,60 @@
 opology(
-edges
+firstNonNull(edges, ImmutableSortedSet.of())
 );%0A  %7D%0A%0A
@@ -1013,79 +1013,242 @@
 %0A%0A  
-private final Map%3CNodeInterfacePair, SortedSet%3CEdge%3E%3E _interfaceEdges;%0A
+// Mapping of interface -%3E set of all edges whose source or dest is that interface%0A  private final Map%3CNodeInterfacePair, SortedSet%3CEdge%3E%3E _interfaceEdges;%0A%0A  // Mapping of node -%3E set of all edges whose source or dest is on that node
 %0A  p
@@ -2324,50 +2324,9 @@
     
-SortedSet%3CEdge%3E edges = getEdges();%0A      
+_
 edge
@@ -2458,32 +2458,88 @@
     
-removeNode(blacklistNode
+_edges.removeAll(_nodeEdges.getOrDefault(blacklistNode, ImmutableSortedSet.of())
 );%0A 
@@ -2675,42 +2675,98 @@
     
-removeInterface(blacklistInterface
+_edges.removeAll(_interfaceEdges.getOrDefault(blacklistInterface, ImmutableSortedSet.of())
 );%0A 
@@ -2917,26 +2917,22 @@
  edge : 
-getE
+_e
 dges
-()
 ) %7B%0A    
@@ -3028,18 +3028,20 @@
 cePair i
-nt
+face
 1 = edge
@@ -3083,18 +3083,20 @@
 cePair i
-nt
+face
 2 = edge
@@ -3123,37 +3123,8 @@
     
- SortedSet%3CEdge%3E node1Edges =
  _no
@@ -3179,26 +3179,8 @@
 %3E())
-;%0A      node1Edges
 .add
@@ -3191,43 +3191,13 @@
 e);%0A
-%0A
      
- SortedSet%3CEdge%3E node2Edges =
  _no
@@ -3252,26 +3252,8 @@
 %3E())
-;%0A      node2Edges
 .add
@@ -3264,48 +3264,13 @@
 e);%0A
-%0A
      
- SortedSet%3CEdge%3E interface1Edges =
  _in
@@ -3299,18 +3299,20 @@
 Absent(i
-nt
+face
 1, k -%3E 
@@ -3331,31 +3331,8 @@
 %3E())
-;%0A      interface1Edges
 .add
@@ -3343,48 +3343,13 @@
 e);%0A
-%0A
      
- SortedSet%3CEdge%3E interface2Edges =
  _in
@@ -3378,18 +3378,20 @@
 Absent(i
-nt
+face
 2, k -%3E 
@@ -3410,426 +3410,17 @@
 %3E())
-;%0A      interface2Edges.add(edge);%0A    %7D%0A  %7D%0A%0A  private void removeInterface(NodeInterfacePair iface) %7B%0A    SortedSet%3CEdge%3E interfaceEdges = _interfaceEdges.get(iface);%0A    if (interfaceEdges != null) %7B%0A      _edges.removeAll(interfaceEdges);%0A    %7D%0A  %7D%0A%0A  private void removeNode(String hostname) %7B%0A    SortedSet%3CEdge%3E nodeEdges = _nodeEdges.get(hostname);%0A    if (nodeEdges != null) %7B%0A      _edges.removeAll(nodeEdges
+.add(edge
 );%0A 

Simplify code and add clarifying comments in Topology (#2483)