From 1d624ea83457ba345760984b50c5ae89319fc0c8 Mon Sep 17 00:00:00 2001
From: corentinlger <corentin.lger@gmail.com>
Date: Wed, 14 Feb 2024 19:57:45 +0100
Subject: [PATCH] Add interaction between server, vis client and notebook
 client

---
 client.py                                     |  2 +-
 notebook_controller.ipynb                     | 15 ++++++++---
 simulate.py => run/simulate.py                |  0
 .../wrapp_simulation.py                       |  0
 server.py                                     | 27 +++++++++++++++++--
 simulationsandbox/utils/network.py            | 17 +++++++++++-
 simulationsandbox/{ => utils}/sim_types.py    |  0
 7 files changed, 54 insertions(+), 7 deletions(-)
 rename simulate.py => run/simulate.py (100%)
 rename wrapp_simulation.py => run/wrapp_simulation.py (100%)
 rename simulationsandbox/{ => utils}/sim_types.py (100%)

diff --git a/client.py b/client.py
index edf9337..ea27f9f 100644
--- a/client.py
+++ b/client.py
@@ -6,7 +6,7 @@
 
 from simulationsandbox.two_d_simulation import SimpleSimulation
 from simulationsandbox.utils.network import SERVER
-from simulationsandbox.sim_types import SIMULATIONS
+from simulationsandbox.utils.sim_types import SIMULATIONS
 
 
 PORT = 5050
diff --git a/notebook_controller.ipynb b/notebook_controller.ipynb
index dd26865..9351b66 100644
--- a/notebook_controller.ipynb
+++ b/notebook_controller.ipynb
@@ -21,8 +21,8 @@
     "import jax.numpy as jnp\n",
     "from flax import serialization\n",
     "\n",
-    "from MultiAgentsSim.two_d_simulation import SimpleSimulation\n",
-    "from MultiAgentsSim.utils.network import SERVER"
+    "from simulationsandbox.two_d_simulation import SimpleSimulation\n",
+    "from simulationsandbox.utils.network import SERVER"
    ]
   },
   {
@@ -89,7 +89,16 @@
     "    current_state = serialization.from_bytes(state_example, client.recv(state_bytes_size))\n",
     "    response_state = change_state_agent_color(current_state, agent_idx, color)\n",
     "    client.send(serialization.to_bytes(response_state))\n",
-    "    return "
+    "    return \n",
+    "\n",
+    "def pause(client):\n",
+    "    client.send(\"PAUSE\".encode())\n",
+    "\n",
+    "def resume(client):\n",
+    "    client.send(\"RESUME\".encode())\n",
+    "\n",
+    "def stop(client):\n",
+    "    client.send(\"STOP\".encode())"
    ]
   },
   {
diff --git a/simulate.py b/run/simulate.py
similarity index 100%
rename from simulate.py
rename to run/simulate.py
diff --git a/wrapp_simulation.py b/run/wrapp_simulation.py
similarity index 100%
rename from wrapp_simulation.py
rename to run/wrapp_simulation.py
diff --git a/server.py b/server.py
index 3a8b0d3..8df9bb4 100644
--- a/server.py
+++ b/server.py
@@ -9,7 +9,7 @@
 
 from simulationsandbox.wrapper import SimulationWrapper
 from simulationsandbox.utils.network import SERVER
-from simulationsandbox.sim_types import SIMULATIONS
+from simulationsandbox.utils.sim_types import SIMULATIONS
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--sim_type', type=str, default="two_d")
@@ -29,6 +29,7 @@
 MAX_AGENTS = 10
 NUM_OBS = 3 
 GRID_SIZE = 20 
+PRINT_DATA = False
 
 # Initialize server
 server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -50,7 +51,7 @@
 update_event = threading.Event()
 
 
-simulation = SimulationWrapper(sim, state, key, step_delay=STEP_DELAY, update_event=update_event)
+simulation = SimulationWrapper(sim, state, key, step_delay=STEP_DELAY, update_event=update_event, print_data=PRINT_DATA)
 print(f"{len(pickle.dumps(simulation.state))}")
 
 
@@ -102,10 +103,32 @@ def communicate_with_client(client, addr, connection_type):
 
                     elif request == "SET_STATE":
                         with sim_lock:
+                            simulation.pause()
                             client.send(serialization.to_bytes(simulation.state))  
                             updated_state = serialization.from_bytes(state, client.recv(state_byte_size))  
                             simulation.state = updated_state
+                            simulation.resume()
+
+                    elif request == "PAUSE":
+                        with sim_lock:
+                            simulation.pause()
+                        print("Simulation paused")
                     
+                    elif request == "RESUME":
+                        with sim_lock:
+                            simulation.resume()
+                        print("Simulation resumed")
+
+                    elif request == "STOP":
+                        with sim_lock:
+                            simulation.stop()
+                        print("Simulation stopped")
+
+                    elif request == "START":
+                        with sim_lock:
+                            simulation.start()
+                        print("Simulation started")
+
                     else:
                         print(f"Unknow request type {request}")
 
diff --git a/simulationsandbox/utils/network.py b/simulationsandbox/utils/network.py
index 48461b1..f7e48ba 100644
--- a/simulationsandbox/utils/network.py
+++ b/simulationsandbox/utils/network.py
@@ -4,4 +4,19 @@
 servers = ['localhost', '10.204.2.210', '10.204.2.229', '10.204.2.189', '192.168.1.24']
 server_idx = 0
 
-SERVER = servers[server_idx]
\ No newline at end of file
+SERVER = servers[server_idx]
+
+
+# Establish a connection with a client
+def establish_connection(client, addr, simulation, sim_type, data_size):
+    try:
+        client.send(sim_type.encode())
+        client.send(pickle.dumps(simulation.state))
+        connection_type = client.recv(data_size).decode()
+        print(f"{connection_type} connection established with {addr}")
+        return connection_type
+
+    except socket.error as e:
+        print(f"error: {e}")
+        client.close()
+        print(f"Client {client} disconnected")
diff --git a/simulationsandbox/sim_types.py b/simulationsandbox/utils/sim_types.py
similarity index 100%
rename from simulationsandbox/sim_types.py
rename to simulationsandbox/utils/sim_types.py