diff --git a/DRAMSys/library/src/controller/BankMachine.cpp b/DRAMSys/library/src/controller/BankMachine.cpp
index 572e92c0..fe6606b0 100644
--- a/DRAMSys/library/src/controller/BankMachine.cpp
+++ b/DRAMSys/library/src/controller/BankMachine.cpp
@@ -1,6 +1,4 @@
 #include "BankMachine.h"
-#include "Command.h"
-#include "core/scheduling/checker/CheckerDDR3New.h"
 
 BankMachine::BankMachine(ControllerNew *controller, SchedulerNew *scheduler, CheckerDDR3New *checker, Bank bank)
     : controller(controller), scheduler(scheduler), checker(checker), bank(bank) {}
diff --git a/DRAMSys/library/src/controller/BankMachine.h b/DRAMSys/library/src/controller/BankMachine.h
index 8a857d65..c03b54ef 100644
--- a/DRAMSys/library/src/controller/BankMachine.h
+++ b/DRAMSys/library/src/controller/BankMachine.h
@@ -3,9 +3,12 @@
 
 #include <systemc.h>
 #include <tlm.h>
+#include <utility>
 #include "../common/dramExtensions.h"
+#include "ControllerNew.h"
 #include "Command.h"
 #include "SchedulerNew.h"
+#include "core/scheduling/checker/CheckerDDR3New.h"
 
 using namespace tlm;
 
diff --git a/DRAMSys/library/src/controller/CommandMux.cpp b/DRAMSys/library/src/controller/CommandMux.cpp
index 3198c3a6..866e38a7 100644
--- a/DRAMSys/library/src/controller/CommandMux.cpp
+++ b/DRAMSys/library/src/controller/CommandMux.cpp
@@ -1,5 +1,4 @@
 #include "CommandMux.h"
-#include "core/configuration/Configuration.h"
 
 void CommandMux::insertPayload(tlm_generic_payload *payload)
 {
diff --git a/DRAMSys/library/src/controller/CommandMux.h b/DRAMSys/library/src/controller/CommandMux.h
index 46a6ce33..67cb4f79 100644
--- a/DRAMSys/library/src/controller/CommandMux.h
+++ b/DRAMSys/library/src/controller/CommandMux.h
@@ -1,12 +1,12 @@
 #ifndef COMMANDMUX_H
 #define COMMANDMUX_H
 
-#include <queue>
 #include <systemc.h>
 #include <tlm.h>
-#include "ControllerState.h"
-#include "BankMachine.h"
-#include "ControllerNew.h"
+#include <utility>
+#include <vector>
+#include <queue>
+#include "Command.h"
 
 using namespace tlm;
 
diff --git a/DRAMSys/library/src/controller/ControllerNew.cpp b/DRAMSys/library/src/controller/ControllerNew.cpp
index 073a9fa6..1e99b479 100644
--- a/DRAMSys/library/src/controller/ControllerNew.cpp
+++ b/DRAMSys/library/src/controller/ControllerNew.cpp
@@ -34,6 +34,7 @@ ControllerNew::~ControllerNew()
     for (auto it : bankMachines)
         delete it.second;
     delete scheduler;
+    delete commandMux;
 }
 
 tlm_sync_enum ControllerNew::nb_transport_fw(tlm_generic_payload &trans,
@@ -123,13 +124,13 @@ void ControllerNew::controllerMethod()
     if (lastTimeCalled != sc_time_stamp())
     {
         lastTimeCalled = sc_time_stamp();
-
+        // Release payload if arbiter has accepted the result
         if (sc_time_stamp() == timeToRelease && payloadToRelease != nullptr)
             releasePayload();
-
+        // Accept new request from arbiter
         if (sc_time_stamp() >= timeToAcquire && payloadToAcquire != nullptr)
             acquirePayload();
-
+        // Update states of bank machines and get results if ready
         for (auto it : bankMachines)
         {
             if (commandFinishedTime[it.first] == sc_time_stamp())
@@ -139,13 +140,13 @@ void ControllerNew::controllerMethod()
                     responseQueue.push(result);
             }
         }
-
+        // Send result to arbiter
         if (payloadToRelease == nullptr && !responseQueue.empty())
             sendToFrontend();
-
+        // Start bank machines to issue new requests
         for (auto it : bankMachines)
             it.second->startBankMachine();
-
+        // Choose one request and send it to DRAM
         std::vector<std::pair<Command, tlm_generic_payload *>> readyCommands;
         std::pair<Command, tlm_generic_payload *> result;
         for (auto it : bankMachines)
@@ -160,7 +161,7 @@ void ControllerNew::controllerMethod()
             if (result.second != nullptr)
                 sendToDram(result.first, result.second);
         }
-
+        // Restart bank machines to issue new requests for the future
         for (auto it : bankMachines)
             it.second->startBankMachine();
     }
diff --git a/DRAMSys/library/src/controller/SchedulerNew.h b/DRAMSys/library/src/controller/SchedulerNew.h
index 1dd17fab..e5dcf6d6 100644
--- a/DRAMSys/library/src/controller/SchedulerNew.h
+++ b/DRAMSys/library/src/controller/SchedulerNew.h
@@ -3,9 +3,9 @@
 
 #include <systemc.h>
 #include <tlm.h>
-#include "../common/dramExtensions.h"
-#include "ControllerNew.h"
+#include <map>
 #include <deque>
+#include "../common/dramExtensions.h"
 
 using namespace tlm;