diff --git a/src/main/scala/CPU.scala b/src/main/scala/CPU.scala
index 690c6ea..ceb014d 100644
--- a/src/main/scala/CPU.scala
+++ b/src/main/scala/CPU.scala
@@ -2,6 +2,7 @@ package FiveStage
 
 import chisel3._
 import chisel3.core.Input
+import chisel3.util.ShiftRegister
 import chisel3.experimental.MultiIOModule
 import chisel3.experimental._
 
@@ -32,9 +33,6 @@ class CPU extends MultiIOModule {
   val MEM = Module(new MemoryFetch)
   // val WB  = Module(new Execute) (You may not need this one?)
 
-  val data1 = RegInit(UInt(32.W), 0.U)
-  val data2 = RegInit(UInt(32.W), 0.U)
-
   /**
     * Setup. You should not change this code
     */
@@ -56,17 +54,30 @@ class CPU extends MultiIOModule {
   /**
     TODO: Your code here
     */
+  // Stage 2
   ID.io.PC := IF.io.PC
   ID.io.instruction := IF.io.instruction
-  ID.io.regWriteData := EX.io.result
-  data1 := ID.io.data1
-  data2 := ID.io.data2
+  //printf(p"S2: PC=${IF.io.PC}, Opcode=${IF.io.instruction.opcode}, rd=${IF.io.instruction.registerRd}, rs1=${IF.io.instruction.registerRs1}, rs2=${IF.io.instruction.registerRs2} || ")
+
+  // Stage 3
+  EX.io.ALUop := ID.io.ALUop
+  EX.io.data1 := ID.io.data1
+  EX.io.data2 := ID.io.data2
+  //printf(p"S3: ALUop=${ID.io.ALUop}, data1=${ID.io.data1.asSInt}, data2=${ID.io.data2.asSInt} || ")
 
+  // Stage 4
   MEM.io.dataIn := 0.U
-  MEM.io.dataAddress := 0.U
-  MEM.io.writeEnable := false.B
+  MEM.io.dataAddress := EX.io.result
+  MEM.io.writeEnable := ShiftRegister(ID.io.controlSignals.memWrite, 1)
+  //printf(p"S4: res=${EX.io.result} || ")
 
-  EX.io.ALUop := ID.io.ALUop
-  EX.io.data1 := data1
-  EX.io.data2 := data2
+  // Stage 5
+  val memOrEx = ShiftRegister(ID.io.controlSignals.memToReg, 2)        // From stage 3
+  val exResult = ShiftRegister(EX.io.result, 1)                        // From stage 4
+  ID.io.writeEnable := ShiftRegister(ID.io.controlSignals.regWrite, 2) // From stage 3
+  ID.io.writeAddr := ShiftRegister(IF.io.instruction.registerRd, 3)    // From stage 2
+  ID.io.writeData := Mux(memOrEx, MEM.io.dataOut, exResult)
+  //printf(p"S5: Mem=${MEM.io.dataOut} Ex=${exResult}")
+
+  //printf("\n\n")
 }
diff --git a/src/main/scala/ID.scala b/src/main/scala/ID.scala
index 4405fa7..5729ecc 100644
--- a/src/main/scala/ID.scala
+++ b/src/main/scala/ID.scala
@@ -23,7 +23,9 @@ class InstructionDecode extends MultiIOModule {
         */
       val PC = Input(UInt(32.W))
       val instruction = Input(new Instruction)
-      val regWriteData = Input(UInt(32.W))
+      val writeEnable = Input(UInt(32.W))
+      val writeAddr = Input(UInt(32.W))
+      val writeData = Input(UInt(32.W))
 
       val controlSignals = Output(new ControlSignals)
       val data1 = Output(UInt(32.W))
@@ -40,17 +42,6 @@ class InstructionDecode extends MultiIOModule {
   val ALUop = RegInit(UInt(4.W), 0.U)
   val controlSignals = Reg(new ControlSignals)
 
-  // The data will be written back in 4 cycles,
-  // so we move the data between registers while we wait
-  class RegInput extends Bundle {
-    val writeEnable = Bool()
-    val writeAddress = UInt(32.W)
-  }
-
-  val regInput0 = Reg(new RegInput)
-  val regInput1 = Reg(new RegInput)
-  val regInput2 = Reg(new RegInput)
-
 
   /**
     * Setup. You should not change this code
@@ -67,9 +58,9 @@ class InstructionDecode extends MultiIOModule {
 
   registers.io.readAddress1 := io.instruction.registerRs1
   registers.io.readAddress2 := io.instruction.registerRs2
-  registers.io.writeEnable  := regInput2.writeEnable
-  registers.io.writeAddress := regInput2.writeAddress
-  registers.io.writeData    := io.regWriteData
+  registers.io.writeEnable  := io.writeEnable
+  registers.io.writeAddress := io.writeAddr
+  registers.io.writeData    := io.writeData
 
   ALUop := decoder.ALUop
   io.ALUop := ALUop
@@ -79,12 +70,6 @@ class InstructionDecode extends MultiIOModule {
   io.data1 := data1
   io.data2 := data2
 
-
-  regInput0.writeEnable := decoder.controlSignals.regWrite
-  regInput0.writeAddress := io.instruction.registerRd
-  regInput1 := regInput0
-  regInput2 := regInput1
-
   switch (decoder.op1Select) {
     is (Op1Select.rs1) {
       data1 := registers.io.readData1