diff --git a/cz.py b/cz.py
index edf1a36..d2d69f4 100644
--- a/cz.py
+++ b/cz.py
@@ -1,29 +1,24 @@
-from graph import Graph
 from matplotlib import pyplot as plt
+from graph import *
+from viz import draw
 
-class GraphState(Graph):
-
-    def __init__(self, n):
-        Graph.__init__(self, n)
-
-    def local_complementation(self, a):
-        for i in self.neighbours[a]:
-            for j in self.neighbours[a]:
-                if i<j:
-                    self.toggle_edge(i, j)
+def local_complementation(g, vops, a):
+    for i in g[a]:
+        for j in g[a]:
+            if i<j:
+                toggle_edge(g, i, j)
 
 if __name__ == '__main__':
-    g = GraphState(10)
-    g.add_edge(0, 1)
-    g.add_edge(1, 3)
-    g.add_edge(3, 2)
-    g.add_edge(3, 0)
-    g.add_edge(2, 0)
-    g.add_edge(0, 5)
-    g.vops[0]=1
-    g.draw()
-
+    g, vops = graph(10)
+    add_edge(g, 0, 1)
+    add_edge(g, 1, 3)
+    add_edge(g, 3, 2)
+    add_edge(g, 3, 0)
+    add_edge(g, 2, 0)
+    add_edge(g, 0, 5)
+    vops[0]=1
+    draw(g, vops)
     plt.clf()
-    g.local_complementation(0)
-    g.draw("out2.pdf")
+    local_complementation(g, vops, 0)
+    draw(g, vops, "out2.pdf")