diff --git a/boot.py b/boot.py
index 001f817..f6dad95 100755
--- a/boot.py
+++ b/boot.py
@@ -1,12 +1,20 @@
-import ujson
-import network
-from hal import *
-from utime import sleep_ms
-from machine import Timer, Pin
+# Import functions from time library
+from utime import ticks_us, sleep_us, sleep_ms
 
+# Give time to cancel boot script
 print("Press Ctrl-C to stop boot script...")
 sleep_ms(200)
 
+# Import libraries
+import os
+import ujson
+import network
+import _thread
+import uwebsockets
+from machine import Timer, reset
+from hal import *
+# Loading libraries takes ca 400ms
+
 # Open and parse the config file
 with open("config.json", "r") as config_file:
     config = ujson.load(config_file)
@@ -16,9 +24,11 @@ sumorobot = Sumorobot(config)
 
 # Indiacte booting with blinking status LED
 timer = Timer(0)
-sumorobot.toggle_led()
 timer.init(period=2000, mode=Timer.PERIODIC, callback=sumorobot.toggle_led)
 
+# Connected Wi-Fi SSID
+ssid = None
+
 # Connect to WiFi
 wlan = network.WLAN(network.STA_IF)
 # Activate the WiFi interface
@@ -30,9 +40,10 @@ if not wlan.isconnected():
     # Go trough all scanned WiFi networks
     for network in networks:
         # Extract the networks SSID
-        ssid = network[0].decode("utf-8")
+        temp_ssid = network[0].decode("utf-8")
         # Check if the SSID is in the config file
-        if ssid in config["wifis"].keys():
+        if temp_ssid in config["wifis"].keys():
+            ssid = temp_ssid
             # Start to connect to the pre-configured network
             wlan.connect(ssid, config["wifis"][ssid])
             break
diff --git a/hal.py b/hal.py
index b2132bc..7e82eb5 100755
--- a/hal.py
+++ b/hal.py
@@ -1,7 +1,7 @@
 import os
 import ujson
 from utime import sleep_us, sleep_ms
-from machine import Pin, PWM, ADC, time_pulse_us, deepsleep
+from machine import Pin, PWM, ADC, time_pulse_us
 
 # LEDs
 STATUS = 0
diff --git a/main.py b/main.py
index 95922b8..cc3dcd9 100755
--- a/main.py
+++ b/main.py
@@ -1,7 +1,4 @@
-import _thread
-import ubinascii
-import uwebsockets
-
+# The code processing thread
 def step():
     while True:
         # Execute to see LED feedback for sensors
@@ -22,24 +19,36 @@ def step():
         # Leave time to process WebSocket commands
         sleep_ms(50)
 
+# The WebSocket processing thread
 def ws_handler():
+    global conn
+
     while True:
         # When WiFi has just been reconnected
         if wlan.isconnected() and not sumorobot.is_wifi_connected:
-            #conn = uwebsockets.connect(url)
+            print("main.py reconnected to Wi-Fi")
+            timer.deinit()
             sumorobot.set_led(STATUS, True)
             sumorobot.is_wifi_connected = True
         # When WiFi has just been disconnected
         elif not wlan.isconnected() and sumorobot.is_wifi_connected:
+            print("main.py lost Wi-Fi connection, reconnecting...")
             sumorobot.set_led(STATUS, False)
             sumorobot.is_wifi_connected = False
+            timer.init(period=2000, mode=Timer.PERIODIC, callback=sumorobot.toggle_led)
         elif not wlan.isconnected():
             # Continue to wait for a WiFi connection
             continue
 
+        data = None
         try: # Try to read from the WebSocket
             data = conn.recv()
         except: # Socket timeout, no data received
+            # TODO: implement watchdog
+            # Try reconnecting to the websocket if Wi-Fi is connected
+            if wlan.isconnected():
+                print("main.py WebSocket error, reconnecting")
+                conn = uwebsockets.connect(uri)
             # Continue to try to read data
             continue
 
@@ -100,21 +109,33 @@ def ws_handler():
         else:
             print("main.py: unknown cmd=", data)
 
-# Try to load the user code
-try:
-    with open("code.py", "r") as code:
-        temp = code.read()
-        sumorobot.python_code = temp
-        sumorobot.compiled_python_code = compile(temp, "snippet", "exec")
-except:
-    print("main.py: error loading code.py file")
+# When user code (copy.py) exists
+if 'code.py' in os.listdir():
+    print("main.py: loading user code")
+    # Try to load the user code
+    try:
+        with open("code.py", "r") as code:
+            temp = code.read()
+            sumorobot.python_code = temp
+            sumorobot.compiled_python_code = compile(temp, "snippet", "exec")
+    except:
+        print("main.py: error loading code.py file")
 
 # Start the code processing thread
 _thread.start_new_thread(step, ())
 
+# Wifi connection counter
+wifi_connection_counter = 0
 # Wait for WiFi to get connected
 while not wlan.isconnected():
     sleep_ms(100)
+    wifi_connection_counter += 1
+    # When Wi-Fi didn't connect in X seconds
+    if wifi_connection_counter == 30:
+        print("main.py reinitiating Wi-Fi connection")
+        # Re initiate the Wi-Fi connection
+        wlan.connect(ssid, config["wifis"][ssid])
+        wifi_connection_counter = 0
 
 # Connect to the websocket
 uri = "ws://%s/p2p/sumo-%s/browser/" % (config['sumo_server'], config['sumo_id'])
diff --git a/uwebsockets.py b/uwebsockets.py
index 5ce1b3c..0b7dd41 100755
--- a/uwebsockets.py
+++ b/uwebsockets.py
@@ -236,6 +236,7 @@ def connect(uri):
     # We don't (currently) need these headers
     # FIXME: should we check the return key?
     while header:
+        print("uwebsockets.py header:", header)
         header = sock.readline()[:-2]
 
     return WebsocketClient(sock)