diff --git a/include/controller.h b/include/controller.h
new file mode 100644
index 0000000..573081a
--- /dev/null
+++ b/include/controller.h
@@ -0,0 +1,18 @@
+#ifndef CONTROLLER_H
+#define CONTROLLER_H
+
+class Controller
+{
+public:
+    virtual ~Controller()
+    {
+    }
+
+    virtual double update(double ref, double actual) = 0;
+
+    virtual void reset()
+    {
+    }
+};
+
+#endif
\ No newline at end of file
diff --git a/include/initialization_state.h b/include/initialization_state.h
index 01cc28c..6b674b8 100644
--- a/include/initialization_state.h
+++ b/include/initialization_state.h
@@ -7,7 +7,7 @@
 class InitializationState : public State
 {
 private:
-    static const unsigned long WAIT_MS = 2000;
+    static const unsigned long WAIT_MS = 2000UL;
 
     unsigned long enter_time_ms = 0;
 
diff --git a/include/motor_controller.h b/include/motor_controller.h
index e62a481..f175fd0 100644
--- a/include/motor_controller.h
+++ b/include/motor_controller.h
@@ -4,7 +4,8 @@
 #include "rotary_encoder_driver.h"
 #include "motor_driver.h"
 //#include "p_controller.h"
-#include "pi_controller.h"
+//#include "pi_controller.h"
+#include "controller.h"
 #include "digital_out.h"
 
 class State;
@@ -26,7 +27,8 @@ private:
     RotaryEncoderDriver& encoder;
     MotorDriver& motor;
     //P_controller& controller;
-    PI_controller& controller;
+    //PI_controller& controller;
+    Controller& controller;
 
     Digital_out led;
 
@@ -37,7 +39,8 @@ public:
     MotorController(
         RotaryEncoderDriver& encoder,
         MotorDriver& motor,
-        P_controller& controller,
+        //P_controller& controller,
+        Controller& controller,
         State& initialization,
         State& operational,
         State& stopped
@@ -75,11 +78,15 @@ public:
 
     void reset_speed_measurement();
 
+    void reset_controller();
+
+
+
     void led_on();
 
     void led_off();
 
-    void reset_controller();
+    
 };
 
 
diff --git a/include/p_controller.h b/include/p_controller.h
index a7d5813..4992ebc 100644
--- a/include/p_controller.h
+++ b/include/p_controller.h
@@ -1,7 +1,9 @@
 #ifndef P_CONTROLLER_H
 #define P_CONTROLLER_H
 
-class P_controller
+#include "controller.h"
+
+class P_controller : public Controller
 {
 private:
     double Kp;
@@ -9,7 +11,7 @@ private:
 public:
     P_controller(double Kp);
 
-    double update(double ref, double actual);
+    double update(double ref, double actual) override;
 
 
     
diff --git a/include/pi_controller.h b/include/pi_controller.h
index da19f73..f18e09c 100644
--- a/include/pi_controller.h
+++ b/include/pi_controller.h
@@ -1,12 +1,15 @@
 #ifndef PI_CONTROLLER_H
 #define PI_CONTROLLER_H
 
-class PI_controller
+#include "controller.h"
+
+class PI_controller : public Controller
 {
 private:
     double Kp;
     double Ti;
     double T;
+
     double integral;
 
 public:
@@ -16,9 +19,9 @@ public:
         double T
     );
 
-    double update(double ref, double actual);
+    double update(double ref, double actual) override;
 
-    void reset();
+    void reset() override;
 };
 
 #endif
\ No newline at end of file
diff --git a/src/motor_controller.cpp b/src/motor_controller.cpp
index 3a65f31..e2f179f 100644
--- a/src/motor_controller.cpp
+++ b/src/motor_controller.cpp
@@ -8,7 +8,8 @@ MotorController::MotorController(
     RotaryEncoderDriver& encoder,
     MotorDriver& motor,
     //P_controller& controller,
-    PI_controller& controller;
+    //PI_controller& controller;
+    Controller& controller,
     State& initialization,
     State& operational,
     State& stopped
@@ -59,6 +60,8 @@ void MotorController::transition_to(State& new_state)
 
     current_state = &new_state;
 
+    current_state->set_context(this);
+
     current_state->enter();
 }
 
@@ -131,6 +134,11 @@ bool MotorController::fault_active()
     return motor.fault_active();
 }
 
+void MotorController::reset_controller()
+{
+    controller.reset();
+}
+
 void MotorController::led_on()
 {
    led.set_hi();
@@ -142,7 +150,3 @@ void MotorController::led_off()
    led.set_lo();
 }
 
-void MotorController::reset_controller()
-{
-    controller.reset();
-}
\ No newline at end of file
diff --git a/src/pi_controller.cpp b/src/pi_controller.cpp
index de793be..e78ebe3 100644
--- a/src/pi_controller.cpp
+++ b/src/pi_controller.cpp
@@ -1,5 +1,6 @@
 #include "pi_controller.h"
 
+
 PI_controller::PI_controller(
     double Kp,
     double Ti,
@@ -17,22 +18,29 @@ double PI_controller::update(double ref, double actual)
 {
     double error = ref - actual;
 
-    double new_integral = integral + error * T;
+    double new_integral =
+        integral + error * T;
 
     double output =
         Kp * (error + new_integral / Ti);
 
-    // Saturation + anti-windup
+
+    // Positive saturation
     if (output > 100.0)
     {
         return 100.0;
     }
 
+
+    // Negative saturation
     if (output < -100.0)
     {
         return -100.0;
     }
 
+
+    // Only save the integral if the controller
+    // is not saturated
     integral = new_integral;
 
     return output;