diff --git a/code/Imm_gen.py b/code/Imm_gen.py
index e211483..f289227 100644
--- a/code/Imm_gen.py
+++ b/code/Imm_gen.py
@@ -1,23 +1,15 @@
-def Imm_gen(inst_code):
-    test=inst_code[25:32]
-    if test=='0010011':
-        imm_out=20*inst_code[0]+inst_code[0:12]
-    elif test=='0000011':
-        pass
-    elif test=='0100011':
-        pass
-    elif test=='1100011':
-        pass
-    elif test=='1101111':
-        pass
-    elif test=='1100111':
-        pass
-    elif test=='0110111':
-        pass
-def imm_gen(inst_code,imm_out):
-    test=inst_code[25:32]
-    if test=="0010011":
-        imm_out=inst_code[0]*20+inst_code[0:12]
-    else:
-        imm_out='0'*32
-    return imm_out
\ No newline at end of file
+def Imm_gen(inst_code: str) -> str:
+    test = inst_code[25:32]
+    if test in ['0010011', '0000011', '1100111']:  # i-Type, l, jalr
+        imm_out = inst_code[0:12].rjust(32, inst_code[0])
+    elif test == '0100011':  # s
+        imm_out = (inst_code[0:7] + inst_code[20:25]).rjust(32, inst_code[0])
+    elif test == '1100011':  # branch
+        imm_out = (inst_code[0] + inst_code[24] + inst_code[1:7] +
+                   inst_code[20:24]).rjust(32, inst_code[0])
+    elif test == '1101111':  # jal
+        imm_out = (inst_code[0] + inst_code[12:20] + inst_code[11] +
+                   inst_code[1:11]).rjust(32, inst_code[0])
+    elif test in ['0110111', '0010111']:  # lui, auipc
+        imm_out = inst_code[0:20].rjust(32, inst_code[0])
+    return imm_out