1 files changed, 62 insertions, 9 deletions

diff --git a/app.py b/app.py
index 121fc10..d862310 100644
--- a/app.py
+++ b/app.py
@@ -2,6 +2,7 @@ import os
 import subprocess
 import tempfile
 import json
+import csv
 from flask import Flask, request, jsonify, send_from_directory
 
 app = Flask(__name__)
@@ -10,6 +11,55 @@ app = Flask(__name__)
 TEST_ENVIRONMENT = os.getenv("BONSAI_TEST_MODE", "true").lower() == "true"
 
 
+def generate_matrix_from_edges(edges_file_path, matrix_file_path):
+    """Generate a square matrix file from edges.csv."""
+    try:
+        # First pass: find all unique nodes
+        nodes = set()
+        edges = []
+
+        with open(edges_file_path, "r") as f:
+            reader = csv.reader(f)
+            for row in reader:
+                if len(row) >= 3:
+                    source = int(row[0])
+                    target = int(row[1])
+                    weight = float(row[2])
+                    nodes.add(source)
+                    nodes.add(target)
+                    edges.append((source, target, weight))
+
+        # Create sorted list of nodes
+        node_list = sorted(list(nodes))
+        num_nodes = len(node_list)
+
+        # Create node index mapping
+        node_to_index = {node: i for i, node in enumerate(node_list)}
+
+        # Initialize matrix with zeros
+        matrix = [[0.0 for _ in range(num_nodes)] for _ in range(num_nodes)]
+
+        # Fill matrix with edge weights
+        for source, target, weight in edges:
+            source_idx = node_to_index[source]
+            target_idx = node_to_index[target]
+            matrix[source_idx][target_idx] = weight
+            matrix[target_idx][source_idx] = weight  # Assuming undirected graph
+
+        # Write matrix to file
+        with open(matrix_file_path, "w") as f:
+            for row in matrix:
+                f.write(" ".join(map(str, row)) + "\n")
+
+    except Exception as e:
+        print(f"Error generating matrix: {e}")
+        # Create a simple 3x3 dummy matrix for testing
+        with open(matrix_file_path, "w") as f:
+            f.write("0.0 0.5 0.0\n")
+            f.write("0.5 0.0 0.3\n")
+            f.write("0.0 0.3 0.0\n")
+
+
 @app.route("/")
 def index():
     """Serve the main HTML page."""
@@ -87,9 +137,7 @@ def generate_files():
                         print(
                             f"WARNING: {error_msg}. Creating dummy files for testing."
                         )
-                        edges_content = (
-                            "source,target,weight\nnode1,node2,0.5\nnode2,node3,0.3\n"
-                        )
+                        edges_content = "1,2,0.5\n2,3,0.3\n1,3,0.8\n"
                         nodes_content = "id,label,x,y\nnode1,Node 1,0,0\nnode2,Node 2,1,1\nnode3,Node 3,2,0\n"
 
                         with open(os.path.join(output_dir, "edges.csv"), "w") as f:
@@ -103,9 +151,7 @@ def generate_files():
                 error_msg = "Bonsai execution timed out after 30 seconds"
                 if TEST_ENVIRONMENT:
                     print(f"WARNING: {error_msg}. Creating dummy files for testing.")
-                    edges_content = (
-                        "source,target,weight\nnode1,node2,0.5\nnode2,node3,0.3\n"
-                    )
+                    edges_content = "1,2,0.5\n2,3,0.3\n1,3,0.8\n"
                     nodes_content = "id,label,x,y\nnode1,Node 1,0,0\nnode2,Node 2,1,1\nnode3,Node 3,2,0\n"
 
                     with open(os.path.join(output_dir, "edges.csv"), "w") as f:
@@ -119,9 +165,7 @@ def generate_files():
                 error_msg = "Bonsai tool not found. Please ensure Bonsai is installed and available."
                 if TEST_ENVIRONMENT:
                     print(f"WARNING: {error_msg}. Creating dummy files for testing.")
-                    edges_content = (
-                        "source,target,weight\nnode1,node2,0.5\nnode2,node3,0.3\n"
-                    )
+                    edges_content = "1,2,0.5\n2,3,0.3\n1,3,0.8\n"
                     nodes_content = "id,label,x,y\nnode1,Node 1,0,0\nnode2,Node 2,1,1\nnode3,Node 3,2,0\n"
 
                     with open(os.path.join(output_dir, "edges.csv"), "w") as f:
@@ -139,6 +183,15 @@ def generate_files():
                     with open(filepath, "r") as f:
                         files[filename] = f.read()
 
+            # Generate matrix file if edges.csv exists
+            edges_path = os.path.join(output_dir, "edges.csv")
+            matrix_path = os.path.join(output_dir, "matrix.txt")
+            if os.path.exists(edges_path):
+                generate_matrix_from_edges(edges_path, matrix_path)
+                if os.path.exists(matrix_path):
+                    with open(matrix_path, "r") as f:
+                        files["matrix.txt"] = f.read()
+
             if not files:
                 return (
                     jsonify({"error": "No output files were generated by Bonsai"}),