diff --git a/src/LedController/LedController.c b/src/LedController/LedController.c
index e21ae6b..a50d3ca 100644
--- a/src/LedController/LedController.c
+++ b/src/LedController/LedController.c
@@ -12,14 +12,9 @@
 #include "LedController.h"
 #include "avr/io.h"
 
-#define PA_B1 (1 << 2)
-#define PA_B2 (1 << 3)
-#define PA_B3 (1 << 6)
-#define PA_B4 (1 << 7)
+#define LED_BM_PORTA (1 << 6)|(1 << 1)|(1 << 2)|(1 << 3)
+#define LED_BM_PORTB (1 << 3)|(1 << 2)|(1 << 0)|(1 << 1)
 
-#define BITS_IN_BYTE (1 << 2) | (1 << 3) | (1 << 6) | (1 << 7)
-
-#define HALF_BYTE_BM
 
 LedByte LedController_New(uint8_t *port)
 {
diff --git a/src/LedController/LedController.h b/src/LedController/LedController.h
index 510b8a7..607984f 100644
--- a/src/LedController/LedController.h
+++ b/src/LedController/LedController.h
@@ -18,7 +18,7 @@
  */
 typedef struct Led
 {
-    uint8_t *port;
+    volatile uint8_t *port;
     uint8_t  pin_num;
     bool     state;
 } Led;
@@ -28,6 +28,7 @@ typedef struct LedByte
     Led leds[8];
 } LedByte;
 
+
 /**
  * Returns a instance of the LedByte structure.
  */
diff --git a/src/main.c b/src/main.c
index 49c2f0b..6a9a0c2 100644
--- a/src/main.c
+++ b/src/main.c
@@ -25,8 +25,11 @@
 
 #define SW1PIN (1 << 5)
 #define SW2PIN (1 << 4)
-#define RELAYPIN (1 << 1)
-#define RELAYREADINGPIN (1 << 2) // It would be better to use PA7 so USART worked
+#define RELAYPIN (1 << 7)
+#define RELAYREADINGPIN (1 << 0) // It would be better to use PA7 so USART worked
+
+#define LED_BM_PORTA (1 << 6)|(1 << 1)|(1 << 2)|(1 << 3)
+#define LED_BM_PORTB (1 << 3)|(1 << 2)|(1 << 0)|(1 << 1)
 
 // Set the function pointer for the delay func
 void (*Delay_MicroSeconds)(double us) = _delay_us;
@@ -69,6 +72,44 @@ void WaitForRelayDisconnect(void)
     }
 }
 
+void LedSetup(LedByte *led_byte)
+{
+    PORTA.DIR |= LED_BM_PORTA;
+    PORTB.DIR |= LED_BM_PORTB;
+
+    //pin 13, bit 0, PA3
+    led_byte->leds[0].port = (void *) &PORTA.OUT; 
+    led_byte->leds[0].pin_num = 3;
+
+    //pin 12, bit 1, PA2
+    led_byte->leds[1].port = &PORTA.OUT; 
+    led_byte->leds[1].pin_num = 2;
+
+    //pin 11, bit 2, PA1
+    led_byte->leds[2].port = &PORTA.OUT; 
+    led_byte->leds[2].pin_num = 1;
+
+    //pin 9, bit 3, PB0
+    led_byte->leds[3].port = &PORTB.OUT; 
+    led_byte->leds[3].pin_num = 0;
+
+    //pin 8, bit 4, PB1
+    led_byte->leds[4].port = &PORTB.OUT; 
+    led_byte->leds[4].pin_num = 1;
+
+    //pin 2, bit 5, PB2
+    led_byte->leds[5].port = &PORTB.OUT; 
+    led_byte->leds[5].pin_num = 2;
+
+    //pin 3, bit 6, PB3
+    led_byte->leds[6].port = &PORTB.OUT; 
+    led_byte->leds[6].pin_num = 3;
+
+    //pin 3, bit 7, PA6
+    led_byte->leds[7].port = &PORTA.OUT; 
+    led_byte->leds[7].pin_num = 6;
+}
+
 int main(int argc, char **argv)
 {
     PORTA.DIR |= RELAYPIN;
@@ -81,6 +122,9 @@ int main(int argc, char **argv)
     uint8_t  temp       = 0;
 
     LedByte led_byte = LedController_New(&PORTA.OUT);
+    LedSetup(&led_byte);
+
+    //wait for a second then disable UPDI
 
     while (true)
     {
@@ -94,7 +138,7 @@ int main(int argc, char **argv)
         make_time = Timer_GetOverflowCount();
 
         // Output the Make time via LEDS
-        temp = (uint8_t)(make_time & 0x0F);
+        temp = (uint8_t)(make_time & 0xFF);
         LedController_ShowByte(&led_byte, temp);
 
         SW2_Wait();
@@ -107,7 +151,7 @@ int main(int argc, char **argv)
         break_time = Timer_GetOverflowCount();
 
         // Output the Break time via LEDS
-         temp = (uint8_t)(break_time & 0x0F);
+        temp = (uint8_t)(break_time & 0xFF);
         LedController_ShowByte(&led_byte, temp);