diff --git a/template/tsp.py b/template/tsp.py
index 87ecd32..6ad480f 100644
--- a/template/tsp.py
+++ b/template/tsp.py
@@ -97,16 +97,10 @@ def TSP(tsp_file, model, timeout=60):
 
         data['edges'] = []
         for i in range(len(best_solution)):
-            if i == len(best_solution)-1:
-                data['edges'].append({
-                    'source': best_solution[i],
-                    'target': best_solution[0]
-                })
-            else:
-                data['edges'].append({
-                    'source': best_solution[i],
-                    'target': best_solution[i+1]
-                })
+            data['edges'].append({
+                'source': best_solution[i % len(best_solution)],
+                'target': best_solution[(i+1) % len(best_solution)]
+            })
 
         with open('output/' + os.path.splitext(os.path.basename(tsp_file))[0] + '.json', 'w') as outfile:
             json.dump(data, outfile)
diff --git a/template/utils/write_json.py b/template/utils/write_json.py
new file mode 100644
index 0000000..025ac48
--- /dev/null
+++ b/template/utils/write_json.py
@@ -0,0 +1,27 @@
+import os
+import json
+from utils.load_data import load_data
+
+def write_json(path):
+    for root, _, files in os.walk(path):
+        if(files):
+            for f in files:
+                # Get input file name
+                tsp_file = os.path.join(root, str(f))
+                print(tsp_file)
+                nodes = load_data(tsp_file)
+
+                # Write to JSON
+                data = {}
+
+                data['nodes'] = []
+                for i in range(len(nodes)):
+                    data['nodes'].append({
+                        'title': str(i),
+                        'id': i,
+                        'x': int(nodes[i][0]),
+                        'y': int(nodes[i][1])
+                    })
+
+                with open('./' + os.path.splitext(os.path.basename(tsp_file))[0] + '.json', 'w') as outfile:
+                    json.dump(data, outfile)