package layout;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import javafx.beans.property.ReadOnlySetWrapper;
import javafx.collections.FXCollections;
import javafx.collections.SetChangeListener;
import javafx.scene.Node;
import javafx.scene.layout.Pane;
import physics.ForceField;
import physics.Spring;
import util.UnorderedPair;

public class PhysLayout {

    private final Pane root;
    private final ReadOnlySetWrapper<Node> nodes;
    public final Set<ForceField> fields;

    private final Map<UnorderedPair<Node>, Spring> connections;
    private final Map<Node, Map<Node, Spring>> connectionsTo;
    private final Map<Node, Double> mass;

    public PhysLayout(Pane root) {
        this.root = root;
        nodes = new ReadOnlySetWrapper<>(FXCollections.observableSet());
        connections = new HashMap<>();
        connectionsTo = new HashMap<>();
        mass = new HashMap<>();
        fields = new HashSet<>();
    }

    public Spring getConnection(Node a, Node b) {
        return connections.get(new UnorderedPair(a, b));
    }

    /**
     * Set the mass of a node. If set to infinity, the node will be fixed in
     * place.
     *
     * @param a
     * @param m
     */
    public void setMass(Node a, double m) {
        mass.put(a, m);
    }

    public double getMass(Node a) {
        return mass.getOrDefault(a, 1.0);
    }

    public void addConnection(Node a, Node b, Spring s) {
        nodes.add(a);
        nodes.add(b);
        this.connections.put(new UnorderedPair(a, b), s);
        Map m = this.connectionsTo.get(a);
        Map n = this.connectionsTo.get(b);
        if (m == null) {
            m = new HashMap<>();
            this.connectionsTo.put(a, m);
        }
        if (n == null) {
            n = new HashMap<>();
            this.connectionsTo.put(b, n);
        }
        this.connectionsTo.get(a).put(b, s);
        this.connectionsTo.get(b).put(a, s);
    }

    public void removeConnection(Node a, Node b) {
        if (this.connections.remove(new UnorderedPair(a, b)) != null) {
            Map m = this.connectionsTo.get(a);
            Map n = this.connectionsTo.get(b);
            m.remove(b);
            n.remove(a);
            if (m.isEmpty()) {
                nodes.remove(a);
            }
            if (n.isEmpty()) {
                nodes.remove(b);
            }
        }
    }

    public Set<Node> getNeighbors(Node a) {
        Map<Node, Spring> m = this.connectionsTo.get(a);
        if (m == null) {
            return null;
        }
        return m.keySet();
    }

    public Set<Node> getNodes() {
        return nodes.getReadOnlyProperty();
    }

    public Set<Entry<UnorderedPair<Node>, Spring>> getAllConnections() {
        return connections.entrySet();
    }

    public void addNodeListener(SetChangeListener<? super Node> listener) {
        nodes.addListener(listener);
    }
}

@@ -650,40 +650,28 @@
 de, 
-Map
+Set
 %3CNode
-, Spring%3E%3E connectionsTo
+%3E%3E neighbors
 ;%0A  
@@ -892,29 +892,25 @@
         
-connectionsTo
+neighbors
  = new H
@@ -1574,24 +1574,25 @@
 b), s);%0A
+%0A
         
 Map m = 
@@ -1587,34 +1587,32 @@
     
-Map m = this.connectionsTo
+Set%3CNode%3E nA = neighbors
 .get
@@ -1628,34 +1628,32 @@
     
-Map
+Set%3CNode%3E
  n
+B
  = 
-this.connectionsTo
+neighbors
 .get
@@ -1661,33 +1661,34 @@
 b);%0A        if (
-m
+nA
  == null) %7B%0A    
@@ -1695,17 +1695,18 @@
         
-m
+nA
  = new H
@@ -1700,35 +1700,35 @@
    nA = new Hash
-Map
+Set
 %3C%3E();%0A          
@@ -1733,34 +1733,25 @@
     
-this.connectionsTo
+neighbors
 .put(a, 
 m);%0A
@@ -1742,25 +1742,26 @@
 bors.put(a, 
-m
+nA
 );%0A        %7D
@@ -1774,16 +1774,17 @@
    if (n
+B
  == null
@@ -1800,16 +1800,17 @@
        n
+B
  = new H
@@ -1804,35 +1804,35 @@
    nB = new Hash
-Map
+Set
 %3C%3E();%0A          
@@ -1833,34 +1833,25 @@
         
-this.connectionsTo
+neighbors
 .put(b, 
@@ -1851,16 +1851,17 @@
 put(b, n
+B
 );%0A     
@@ -1877,87 +1877,35 @@
     
-this.connectionsTo.get(a).put(b, s);%0A        this.connectionsTo.get(b).put(a, s
+nA.add(b);%0A        nB.add(a
 );%0A 
@@ -2049,34 +2049,32 @@
     
-Map m = this.connectionsTo
+Set%3CNode%3E nA = neighbors
 .get
@@ -2094,34 +2094,32 @@
     
-Map
+Set%3CNode%3E
  n
+B
  = 
-this.connectionsTo
+neighbors
 .get
@@ -2135,17 +2135,18 @@
         
-m
+nA
 .remove(
@@ -2162,16 +2162,17 @@
        n
+B
 .remove(
@@ -2191,17 +2191,18 @@
     if (
-m
+nA
 .isEmpty
@@ -2271,16 +2271,17 @@
    if (n
+B
 .isEmpty
@@ -2406,141 +2406,30 @@
     
-Map%3CNode, Spring%3E m = this.connectionsTo.get(a);%0A        if (m == null) %7B%0A            return null;%0A        %7D%0A        return m.keySet(
+return neighbors.get(a
 );%0A 

Store neighbours as node-set.