diff --git a/DRAMSys/library/src/controller/BankMachine.cpp b/DRAMSys/library/src/controller/BankMachine.cpp
index 012c76ed..572e92c0 100644
--- a/DRAMSys/library/src/controller/BankMachine.cpp
+++ b/DRAMSys/library/src/controller/BankMachine.cpp
@@ -5,16 +5,6 @@
 BankMachine::BankMachine(ControllerNew *controller, SchedulerNew *scheduler, CheckerDDR3New *checker, Bank bank)
     : controller(controller), scheduler(scheduler), checker(checker), bank(bank) {}
 
-void BankMachine::setCommandFinishedTime(sc_time time)
-{
-    commandFinishedTime = time;
-}
-
-sc_time BankMachine::getCommandFinishedTime()
-{
-    return commandFinishedTime;
-}
-
 tlm_generic_payload *BankMachine::getNextStateAndResult()
 {
     tlm_generic_payload *payloadToReturn = nullptr;
diff --git a/DRAMSys/library/src/controller/BankMachine.h b/DRAMSys/library/src/controller/BankMachine.h
index e699d711..8a857d65 100644
--- a/DRAMSys/library/src/controller/BankMachine.h
+++ b/DRAMSys/library/src/controller/BankMachine.h
@@ -27,8 +27,6 @@ class BankMachine
 {
 public:
     BankMachine(ControllerNew *, SchedulerNew *, CheckerDDR3New*, Bank);
-    void setCommandFinishedTime(sc_time);
-    sc_time getCommandFinishedTime();
     tlm_generic_payload *getNextStateAndResult();
     void startBankMachine();
     std::pair<Command, tlm_generic_payload *> getNextCommand();
@@ -40,7 +38,6 @@ private:
     Row currentRow = Row(0);
     Command nextCommand = Command::NOP;
     sc_time timeToSchedule = SC_ZERO_TIME;
-    sc_time commandFinishedTime = SC_ZERO_TIME;
     ControllerNew *controller;
     SchedulerNew *scheduler;
     CheckerDDR3New *checker;
diff --git a/DRAMSys/library/src/controller/CommandMux.cpp b/DRAMSys/library/src/controller/CommandMux.cpp
index 7aee2d15..3198c3a6 100644
--- a/DRAMSys/library/src/controller/CommandMux.cpp
+++ b/DRAMSys/library/src/controller/CommandMux.cpp
@@ -1,46 +1,29 @@
 #include "CommandMux.h"
 #include "core/configuration/Configuration.h"
 
-CommandMux::CommandMux(ControllerNew *controller, ControllerState *state, std::map<Bank, BankMachine *> &bankMachines) :
-    controller(controller), state(state), bankMachines(bankMachines) {}
-
 void CommandMux::insertPayload(tlm_generic_payload *payload)
 {
     payloadOrder.push(payload);
 }
 
-void CommandMux::selectCommand()
+std::pair<Command, tlm_generic_payload *>
+CommandMux::selectCommand(std::vector<std::pair<Command, tlm_generic_payload *>> &readyCommands)
 {
-    for (auto it : bankMachines)
+    for (auto it : readyCommands)
     {
-        std::pair<Command, tlm_generic_payload *> result = it.second->getNextCommand();
-        if (result.first != Command::NOP)
-            readyCommands.push_back(result);
+        if (it.first == Command::ACT || it.first == Command::PRE)
+            return it;
     }
-    if (!readyCommands.empty())
+    for (auto it : readyCommands)
     {
-        for (auto it : readyCommands)
+        if (it.first == Command::RD || it.first == Command::WR)
         {
-            if (it.first == Command::ACT || it.first == Command::PRE)
+            if (it.second == payloadOrder.front())
             {
-                controller->sendToDram(it.first, it.second);
-                readyCommands.clear();
-                return;
+                payloadOrder.pop();
+                return it;
             }
         }
-        for (auto it : readyCommands)
-        {
-            if (it.first == Command::RD || it.first == Command::WR)
-            {
-                if (it.second == payloadOrder.front())
-                {
-                    payloadOrder.pop();
-                    controller->sendToDram(it.first, it.second);
-                    readyCommands.clear();
-                    return;
-                }
-            }
-        }
-        readyCommands.clear();
     }
+    return std::pair<Command, tlm_generic_payload *>(Command::NOP, nullptr);
 }
diff --git a/DRAMSys/library/src/controller/CommandMux.h b/DRAMSys/library/src/controller/CommandMux.h
index ba74115c..46a6ce33 100644
--- a/DRAMSys/library/src/controller/CommandMux.h
+++ b/DRAMSys/library/src/controller/CommandMux.h
@@ -10,21 +10,13 @@
 
 using namespace tlm;
 
-class ControllerNew;
-class BankMachine;
-
 class CommandMux
 {
 public:
-    CommandMux(ControllerNew *, ControllerState*, std::map<Bank, BankMachine *> &);
     void insertPayload(tlm_generic_payload *);
-    void selectCommand();
+    std::pair<Command, tlm_generic_payload *> selectCommand(std::vector<std::pair<Command, tlm_generic_payload *>> &);
 private:
-    ControllerState *state;
-    ControllerNew *controller;
     std::queue<tlm_generic_payload *> payloadOrder;
-    std::map<Bank, BankMachine *> &bankMachines;
-    std::vector<std::pair<Command, tlm_generic_payload *>> readyCommands;
 };
 
 #endif // COMMANDMUX_H
diff --git a/DRAMSys/library/src/controller/ControllerNew.cpp b/DRAMSys/library/src/controller/ControllerNew.cpp
index f58bc9dc..073a9fa6 100644
--- a/DRAMSys/library/src/controller/ControllerNew.cpp
+++ b/DRAMSys/library/src/controller/ControllerNew.cpp
@@ -20,8 +20,11 @@ ControllerNew::ControllerNew(sc_module_name name, TlmRecorder *tlmRecorder) :
     checker = new CheckerDDR3New(Configuration::getInstance(), *state);
     scheduler = new SchedulerNew();
     for (unsigned bankID = 0; bankID < Configuration::getInstance().memSpec->NumberOfBanks; bankID++)
+    {
         bankMachines[Bank(bankID)] = new BankMachine(this, scheduler, checker, Bank(bankID));
-    commandMux = new CommandMux(this, state, bankMachines);
+        commandFinishedTime[Bank(bankID)] = SC_ZERO_TIME;
+    }
+    commandMux = new CommandMux();
 }
 
 ControllerNew::~ControllerNew()
@@ -30,7 +33,6 @@ ControllerNew::~ControllerNew()
     delete checker;
     for (auto it : bankMachines)
         delete it.second;
-    delete commandMux;
     delete scheduler;
 }
 
@@ -68,7 +70,7 @@ tlm_sync_enum ControllerNew::nb_transport_bw(tlm_generic_payload &trans,
     printDebugMessage("[bw] " + phaseNameToString(phase) + " notification in " +
                       delay.to_string());
     Bank bank = DramExtension::getExtension(trans).getBank();
-    bankMachines[bank]->setCommandFinishedTime(sc_time_stamp() + delay);
+    commandFinishedTime[bank] = sc_time_stamp() + delay;
     triggerEventQueueAfterDelay(delay);
     return TLM_ACCEPTED;
 }
@@ -130,7 +132,7 @@ void ControllerNew::controllerMethod()
 
         for (auto it : bankMachines)
         {
-            if (it.second->getCommandFinishedTime() == sc_time_stamp())
+            if (commandFinishedTime[it.first] == sc_time_stamp())
             {
                 tlm_generic_payload *result = it.second->getNextStateAndResult();
                 if (result != nullptr)
@@ -144,7 +146,20 @@ void ControllerNew::controllerMethod()
         for (auto it : bankMachines)
             it.second->startBankMachine();
 
-        commandMux->selectCommand();
+        std::vector<std::pair<Command, tlm_generic_payload *>> readyCommands;
+        std::pair<Command, tlm_generic_payload *> result;
+        for (auto it : bankMachines)
+        {
+            result = it.second->getNextCommand();
+            if (result.second != nullptr)
+                readyCommands.push_back(result);
+        }
+        if (!readyCommands.empty())
+        {
+            result = commandMux->selectCommand(readyCommands);
+            if (result.second != nullptr)
+                sendToDram(result.first, result.second);
+        }
 
         for (auto it : bankMachines)
             it.second->startBankMachine();
diff --git a/DRAMSys/library/src/controller/ControllerNew.h b/DRAMSys/library/src/controller/ControllerNew.h
index 824978d5..4eea4566 100644
--- a/DRAMSys/library/src/controller/ControllerNew.h
+++ b/DRAMSys/library/src/controller/ControllerNew.h
@@ -36,14 +36,6 @@ public:
     void triggerEventAfterDelay(sc_time);
     void triggerEventQueueAfterDelay(sc_time);
 
-    ControllerState *state;
-    std::map<Bank, BankMachine *> bankMachines;
-    CommandMux *commandMux;
-    SchedulerNew *scheduler;
-    CheckerDDR3New *checker;
-
-    void sendToDram(Command, tlm_generic_payload *);
-
 private:
     tlm_sync_enum nb_transport_fw(tlm_generic_payload &trans,
                                   tlm_phase &phase, sc_time &delay);
@@ -64,13 +56,22 @@ private:
     TlmRecorder *tlmRecorder;
     DebugManager *debugManager;
 
-    void controllerMethod();
-    sc_event triggerEvent;
-    sc_event_queue triggerEventQueue;
+    ControllerState *state;
+    std::map<Bank, BankMachine *> bankMachines;
+    CommandMux *commandMux;
+    SchedulerNew *scheduler;
+    CheckerDDR3New *checker;
 
     void releasePayload();
     void acquirePayload();
     void sendToFrontend();
+    void sendToDram(Command, tlm_generic_payload *);
+
+    void controllerMethod();
+    sc_event triggerEvent;
+    sc_event_queue triggerEventQueue;
+
+    std::map<Bank, sc_time> commandFinishedTime;
 };
 
 #endif // CONTROLLERNEW_H