diff --git a/lib/xlpmg/Communicator.hpp b/lib/xlpmg/Communicator.hpp
index 16bf9803..510c716c 100644
--- a/lib/xlpmg/Communicator.hpp
+++ b/lib/xlpmg/Communicator.hpp
@@ -14,6 +14,7 @@
 #include <arpa/inet.h>
 #include <sys/socket.h>
 #include <unistd.h>
+#include <poll.h>
 
 // time stamp
 #include <chrono>
@@ -222,6 +223,27 @@ namespace xlpmg
             return message;
         }
 
+        /// @brief Checks if the server responded with OK.
+        /// @return true if server responded
+        bool checkServerResponse()
+        {
+            char readBuffer[BUFF_SIZE_READ];
+            memset(readBuffer, 0, BUFF_SIZE_READ);
+            sockValread = read(sockClient_fd, readBuffer,
+                               BUFF_SIZE_READ - 1); // subtract 1 for the null
+            bool returnValue = true;
+            if(std::string(readBuffer).compare("OK") != 0)
+                returnValue = false;
+
+            //DONE
+            sockValread = read(sockClient_fd, readBuffer,
+                               BUFF_SIZE_READ - 1); // subtract 1 for the null
+            if(std::string(readBuffer).compare("DONE") != 0)
+                returnValue = false;
+
+            return returnValue;
+        }
+
         /// @brief Sends a message to the server.
         /// @param message String to send.
         int sendToServer(std::string message)
@@ -253,7 +275,7 @@ namespace xlpmg
                     logEvent(bytesSentStr.c_str(), DEBUG, true);
                 }
             }
-            sleep(1);
+
             send(sockClient_fd, "DONE", strlen("DONE"), MSG_NOSIGNAL);
 
             return 0;
@@ -409,6 +431,7 @@ namespace xlpmg
                     logEvent(bytesSentStr.c_str(), DEBUG, true);
                 }
             }
+
             send(new_socket, "DONE", 4, MSG_NOSIGNAL);
         }
     };
diff --git a/src/Server.cpp b/src/Server.cpp
index c61a4718..c22423eb 100644
--- a/src/Server.cpp
+++ b/src/Server.cpp
@@ -68,6 +68,7 @@ int main(int i_argc, char *i_argv[])
             {
                 continue;
             }
+            
             json l_parsedData = json::parse(l_rawData);
             xlpmg::Message l_message = xlpmg::jsonToMessage(l_parsedData);
             xlpmg::MessageType l_type = l_message.type;
@@ -89,7 +90,6 @@ int main(int i_argc, char *i_argv[])
                 else if (l_key == xlpmg::START_SIMULATION.key)
                 {
                     std::string l_config = l_parsedData.at(xlpmg::ARGS);
-                    std::cout << m_simulationThread.joinable() << std::endl;
                     if (!m_isSimulationRunning && !simulator->isPreparing())
                     {
                         if (m_simulationThread.joinable())
@@ -179,6 +179,7 @@ int main(int i_argc, char *i_argv[])
                     l_writeFile.write((char *)&l_byteVector[0], l_byteVector.size());
                     l_writeFile.close();
                     simulator->setBathymetryFilePath(m_bathTempFile);
+                    simulator->setPrepared(false);
                 }else if (l_key == xlpmg::SET_DISPLACEMENT_DATA.key)
                 {
                     std::vector<std::uint8_t> l_byteVector = l_args["bytes"];
@@ -186,6 +187,7 @@ int main(int i_argc, char *i_argv[])
                     l_writeFile.write((char *)&l_byteVector[0], l_byteVector.size());
                     l_writeFile.close();
                     simulator->setDisplacementFilePath(m_displTempFile);
+                    simulator->setPrepared(false);
                 }
                 else if(l_key == xlpmg::PAUSE_SIMULATION.key)
                 {
diff --git a/src/Simulator.cpp b/src/Simulator.cpp
index c511e36e..cb7039cb 100644
--- a/src/Simulator.cpp
+++ b/src/Simulator.cpp
@@ -271,7 +271,6 @@ void tsunami_lab::Simulator::constructSetup()
                                                       m_displacementFilePath.c_str(),
                                                       l_netCdfCustom,
                                                       m_nx);
-                                                      std::cout << "setup : " <<m_bathymetryFilePath<< std::endl;
     }else{
       m_setup = new tsunami_lab::setups::TsunamiEvent1d(m_bathymetryFilePath);
     }
@@ -743,6 +742,7 @@ void tsunami_lab::Simulator::prepareForCalculation()
     loadStations();
   }
   m_isPreparing = false;
+  m_isPrepared = true;
 }
 
 void tsunami_lab::Simulator::runCalculation()
@@ -861,7 +861,6 @@ int tsunami_lab::Simulator::start(std::string i_config)
   if (!m_isPrepared)
   {
     prepareForCalculation();
-    m_isPrepared = true;
   }
 
   // BREAKPOINT
diff --git a/src/Simulator.h b/src/Simulator.h
index 803441b9..811b1e0c 100644
--- a/src/Simulator.h
+++ b/src/Simulator.h
@@ -336,6 +336,10 @@ class tsunami_lab::Simulator
         m_displacementFilePath = i_filePath;
     }
 
+    void setPrepared(bool i_prepared){
+        m_isPrepared = i_prepared;
+    }
+
     //--------------------------------------------//
     //--------------PUBLIC DELETERS---------------//
     //--------------------------------------------//
diff --git a/src/ui/GUI.cpp b/src/ui/GUI.cpp
index 63b59e50..44e6eee0 100644
--- a/src/ui/GUI.cpp
+++ b/src/ui/GUI.cpp
@@ -268,8 +268,7 @@ int tsunami_lab::ui::GUI::launch()
                     if (ImGui::Button("Check connection"))
                     {
                         m_communicator.sendToServer(messageToJsonString(xlpmg::CHECK));
-                        std::string l_response = m_communicator.receiveFromServer();
-                        if (l_response == "OK")
+                        if (m_communicator.checkServerResponse())
                         {
                             m_connected = true;
                         }
@@ -278,6 +277,7 @@ int tsunami_lab::ui::GUI::launch()
                             m_connected = false;
                         }
                     }
+                    
                     ImGui::BeginDisabled(!m_connected);
                     ImGui::PushStyleColor(ImGuiCol_Button, (ImVec4)ImColor::HSV(1.0f, 0.6f, 0.6f));
                     ImGui::PushStyleColor(ImGuiCol_ButtonHovered, (ImVec4)ImColor::HSV(1.0f, 0.8f, 0.8f));