diff --git a/include/motor_controller.h b/include/motor_controller.h
index 8c5b502..770af32 100644
--- a/include/motor_controller.h
+++ b/include/motor_controller.h
@@ -18,9 +18,9 @@ private:
 
     State* current_state;
 
-    State* initialization_state;
-    State* operational_state;
-    State* stopped_state;
+    State& initialization_state;
+    State& operational_state;
+    State& stopped_state;
 
     RotaryEncoderDriver& encoder;
     MotorDriver& motor;
@@ -35,15 +35,18 @@ public:
     MotorController(
         RotaryEncoderDriver& encoder,
         MotorDriver& motor,
-        P_controller& controller
+        P_controller& controller,
+        State& initialization,
+        State& operational,
+        State& stopped
     );
 
 
-    void set_states(
-        State* initialization,
-        State* operational,
-        State* stopped
-    );
+    // void set_states(
+    //     State* initialization,
+    //     State* operational,
+    //     State* stopped
+    // );
 
 
     void begin();
diff --git a/src/main.cpp b/src/main.cpp
index a613548..261b8e0 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -33,7 +33,8 @@ MotorController motor_controller(
     motor,
     controller,
     initialization_state,
-    operational_state
+    operational_state,
+    stopped_state
 );
 
 
@@ -53,11 +54,11 @@ int main()
     sei();
 
 
-    motor_controller.set_states(
-        &initialization_state,
-        &operational_state,
-        &stopped_state
-    );
+    // motor_controller.set_states(
+    //     &initialization_state,
+    //     &operational_state,
+    //     &stopped_state
+    // );
 
 
     motor_controller.begin();
diff --git a/src/motor_controller.cpp b/src/motor_controller.cpp
index 7e192a9..3efb811 100644
--- a/src/motor_controller.cpp
+++ b/src/motor_controller.cpp
@@ -9,33 +9,35 @@ MotorController::MotorController(
     MotorDriver& motor,
     P_controller& controller,
     State& initialization,
-    State& operational
+    State& operational,
+    State& stopped
 )
     : current_state(nullptr),
-      initialization_state(nullptr),
-      operational_state(nullptr),
-      stopped_state(nullptr),
+      initialization_state(initialization),
+      operational_state(operational),
+      stopped_state(stopped),
       encoder(encoder),
       motor(motor),
       controller(controller),
       led(LED_PIN),
       reference_rpm(DEFAULT_REFERENCE_RPM)
 {
+    initialization_state.set_context(this);
+    operational_state.set_context(this);
+    stopped_state.set_context(this);
 }
 
 
-void MotorController::set_states(
-    State* initialization,
-    State* operational,
-    State* stopped
-)
-{
-    initialization_state = initialization;
-    operational_state = operational;
-    stopped_state = stopped;
-    //initialization_state.set_context(this);
-    //operational_state.set_context(this);
-}
+// void MotorController::set_states(
+//     State* initialization,
+//     State* operational,
+//     State* stopped
+// )
+
+    // initialization_state = initialization;
+    // operational_state = operational;
+    // stopped_state = stopped;
+
 
 
 void MotorController::begin()
@@ -128,13 +130,13 @@ bool MotorController::fault_active()
     return motor.fault_active();
 }
 
-//void MotorController::led_on()
-//{
-//    led.set_hi();
-//}
+void MotorController::led_on()
+{
+   led.set_hi();
+}
 
 
-//void MotorController::led_off()
-//{
-//    led.set_lo();
-//}
+void MotorController::led_off()
+{
+   led.set_lo();
+}
diff --git a/src/motor_driver.cpp b/src/motor_driver.cpp
index e298ef6..6b841a7 100644
--- a/src/motor_driver.cpp
+++ b/src/motor_driver.cpp
@@ -126,5 +126,5 @@ void MotorDriver::set(double pwm)
 
 bool MotorDriver::fault_active()
 {
-    return !(PIND & (1 << PIND4));p
+    return !(PIND & (1 << PIND4));
 }