diff --git a/AES.v b/AES.v
index 9c1d3b9..055375c 100644
--- a/AES.v
+++ b/AES.v
@@ -1,38 +1,68 @@
-module AES(LED, HEX0, HEX1, HEX2, clk, reset);
-	localparam Nk = 4;
-	localparam Nr = 10;
-
+module AES(LED, HEX0, HEX1, HEX2, sel, clk);
+    input [1:0] sel; // 00 -> 128-bit AES, 01 -> 192-bit AES, 10/11 -> 256-bit AES
     input clk;
-	input reset;
+
     output LED;
     output [6:0] HEX0;
     output [6:0] HEX1;
     output [6:0] HEX2;
-    output [6:0] HEX3;
-    output [6:0] HEX4;
-    output [6:0] HEX5;
 
-    // Key
-    wire [127:0] key = 128'h00_01_02_03_04_05_06_07_08_09_0a_0b_0c_0d_0e_0f;
+    // Keys
+    wire [127:0] key128 = 128'h000102030405060708090a0b0c0d0e0f;
+    wire [191:0] key192 = 192'h000102030405060708090a0b0c0d0e0f1011121314151617;
+    wire [255:0] key256 = 256'h000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f;
 
-    // Key Expansion
-    wire [(11 * 128) - 1:0] allKeys;
-    KeyExpansion keysGetter(key, allKeys);
+    // Nr / Nk
+    wire [3:0] Nr = sel == 0 ? 4'b1010 : sel == 1 ? 4'b1100 : 4'b1110;
 
     // Data
     wire [127:0] data = 128'h00_11_22_33_44_55_66_77_88_99_aa_bb_cc_dd_ee_ff;
 
-    // AES
-    wire [127:0] tempEncryptedOutput;
-    wire [127:0] tempDecryptedOutput;
-
     // Current round count for encryption and decryption
-    reg [5:0] count;
-	initial
-		count = 0;
+    reg [5:0] count = 0;
+
+    // AES Decrypt Enable
+    reg AESDecryptEnable = 1'b0;
+
+    // Selection Reg
+    reg [1:0] selReg = 0;
+
+    // 128-bit AES
+    wire aesReset128 = selReg == 0 ? 1'b0 : 1'b1;
+    wire [127:0] tempEncryptedOutput128;
+    wire [127:0] tempDecryptedOutput128;
+    wire [1407:0] allKeys128;
+
+    KeyExpansion #(4, 10) keysGetter128(key128, allKeys128);
+    AESEncrypt #(4, 10) aese128(data, allKeys128, tempEncryptedOutput128, clk, aesReset128);
+    AESDecrypt #(4, 10) aesd128(tempEncryptedOutput128, allKeys128, tempDecryptedOutput128, clk, AESDecryptEnable, aesReset128);
+
+    // 192-bit AES
+    wire aesReset192 = selReg == 1 ? 1'b0 : 1'b1;
+    wire [127:0] tempEncryptedOutput192;
+    wire [127:0] tempDecryptedOutput192;
+    wire [1663:0] allKeys192;
+
+    KeyExpansion #(6, 12) keysGetter192(key192, allKeys192);
+    AESEncrypt #(6, 12) aese192(data, allKeys192, tempEncryptedOutput192, clk, aesReset192);
+    AESDecrypt #(6, 12) aesd192(tempEncryptedOutput192, allKeys192, tempDecryptedOutput192, clk, AESDecryptEnable, aesReset192);
+
+    // 256-bit AES
+    wire aesReset256 = (selReg == 2 || selReg == 3) ? 1'b0 : 1'b1;
+    wire [127:0] tempEncryptedOutput256;
+    wire [127:0] tempDecryptedOutput256;
+    wire [1919:0] allKeys256;
+
+    KeyExpansion #(8, 14) keysGetter(key256, allKeys256);
+    AESEncrypt #(8, 14) aese256(data, allKeys256, tempEncryptedOutput256, clk, aesReset256);
+    AESDecrypt #(8, 14) aesd256(tempEncryptedOutput256, allKeys256, tempDecryptedOutput256, clk, AESDecryptEnable, aesReset256);
+
+    // Encrypted and Decrypted Data
+    wire [127:0] tempEncryptedOutput = selReg == 0 ? tempEncryptedOutput128 : selReg == 1 ? tempEncryptedOutput192 : tempEncryptedOutput256;
+    wire [127:0] tempDecryptedOutput = selReg == 0 ? tempDecryptedOutput128 : selReg == 1 ? tempDecryptedOutput192 : tempDecryptedOutput256;
 
     // Assign bcdInput based on count
-    // count = 0 -> Encrypted Data
+    // count = 0 -> Original Data
     // count = 1 to Nr -> Encrypted Data
     // count = Nr + 2 -> Decrypted Data
     wire [7:0] bcdInput = (count == 0) ? data[7:0] : (count <= Nr + 1) ? tempEncryptedOutput[7:0] : tempDecryptedOutput[7:0];
@@ -41,54 +71,70 @@ module AES(LED, HEX0, HEX1, HEX2, clk, reset);
     wire [11:0] bcdOutput;
 	Binary2BCD b2b(bcdInput, bcdOutput);
 
-	// 7-Segment Logic
+    // 7-Segment Logic
 	DisplayDecoder dd1(bcdOutput[3:0], HEX0);
     DisplayDecoder dd2(bcdOutput[7:4], HEX1);
     DisplayDecoder dd3(bcdOutput[11:8], HEX2);
 
-    // Encrypt
-    // reg AESEncryptEnable = 1'b1;
-    AESEncrypt AESE(data, allKeys, tempEncryptedOutput, clk,reset);
-
-    // Decrypt
-    reg AESDecryptEnable = 1'b0;
-    AESDecrypt AESD(tempEncryptedOutput, allKeys, tempDecryptedOutput, clk, AESDecryptEnable,reset);
-
     // LED = 1 if Decrypted Data is same as original data
     assign LED = (tempDecryptedOutput == data && count > Nr + 1);
 
+    // Previous Selection
+    reg [1:0] prevSel = 2'b0;
+
     always @(negedge clk) begin
-		if (reset) begin
-			count = 0;
-			AESDecryptEnable = 1'b0;
-		end
-		else begin
-			if (count == Nr)
-				AESDecryptEnable = 1'b1;
-			else if (count == ((Nr + 1) * 2))
-				AESDecryptEnable = 1'b0;
+        prevSel = selReg;
+        selReg = sel;
+
+        if (prevSel != selReg) begin
+            prevSel = selReg;
+            count = 0;
+            AESDecryptEnable = 1'b0;
+        end
+        else begin
+            if (count == Nr)
+                AESDecryptEnable = 1'b1;
+            else if (count == ((Nr + 1) * 2))
+                AESDecryptEnable = 1'b0;
 
             if (count <= (Nr + 1) * 2)
-                count <= count + 6'b000001;
-		end
+                count = count + 6'b000001;
+        end
     end
 endmodule
 
 module AES_DUT();
+    reg [1:0] sel = 2'b00;
     reg clk = 1'b1;
     wire LED;
     wire [6:0] HEX0, HEX1, HEX2;
 
-    AES aes(LED, HEX0, HEX1, HEX2, clk, 1'b0);
+    AES aes(LED, HEX0, HEX1, HEX2, sel, clk);
 
-    reg [5:0] count = 0;
+    integer count = 0;
     initial begin
         clk = 1'b1;
+
         forever begin
             #10 clk = ~clk;
             if (!clk) begin
-                count = count + 1;
-                $display("Current Round Count: %d, LED Status: %b, Encrypted State: %h (%0d), Decrypted State: %h (%0d)", count, LED, aes.tempEncryptedOutput, aes.tempEncryptedOutput[7:0], aes.tempDecryptedOutput, aes.tempDecryptedOutput[7:0]);
+                if (count < 80)
+                    count = count + 1;
+                
+                $display("Current Round Count: %0d, LED Status: %b, Encrypted State: %h (%0d), Decrypted State: %h (%0d)", count, LED, aes.tempEncryptedOutput, aes.tempEncryptedOutput[7:0], aes.tempDecryptedOutput, aes.tempDecryptedOutput[7:0]);
+
+                if (count == 23) begin
+                    $display("==============================================");
+                    $display("Switching to 192-bit AES Encryption and Decryption");
+                    $display("==============================================");
+                    sel = 2'b01;
+                end
+                else if (count == 50) begin
+                    $display("==============================================");
+                    $display("Switching to 256-bit AES Encryption and Decryption");
+                    $display("==============================================");
+                    sel = 2'b10;
+                end
             end
         end
     end
@@ -96,5 +142,7 @@ module AES_DUT();
     initial begin
         $display("AES Encryption and Decryption");
         $display("================================");
+        $display("128-bit AES Encryption and Decryption");
+        $display("================================");
     end
 endmodule
\ No newline at end of file
diff --git a/AESDecrypt.v b/AESDecrypt.v
index 91a35f2..97738c0 100644
--- a/AESDecrypt.v
+++ b/AESDecrypt.v
@@ -1,6 +1,6 @@
 module AESDecrypt #(parameter Nk = 4, parameter Nr = 10) (data, allKeys, state, clk, enable, reset);
 	input [127:0] data;
-	input [(11 * 128) - 1:0] allKeys;
+	input [((Nr + 1) * 128) - 1:0] allKeys;
 	input clk;
 	input enable;
 	input reset;
@@ -16,9 +16,9 @@ module AESDecrypt #(parameter Nk = 4, parameter Nr = 10) (data, allKeys, state,
 	wire [127:0] stateOut;
 
 	// Instantiate AES modules needed for decryption
-	InvShiftRows shft(state, shiftRowsWire);	
+	InvShiftRows shft(state, shiftRowsWire);
 	InvSubBytes sub(shiftRowsWire, subByteWire);
-	AddRoundKey addkey(keyInput, allKeys[((12 - roundCount) * 128 - 1) -: 128], afterRoundKey);
+	AddRoundKey addkey(keyInput, allKeys[(roundCount * 128) - 1 -: 128], afterRoundKey);
 	InvMixColumns mix(afterRoundKey, mixColumnsWire);
 
 	// Assign keyInput based on roundCount
@@ -42,16 +42,57 @@ module AESDecrypt #(parameter Nk = 4, parameter Nr = 10) (data, allKeys, state,
 		if (reset)
 			roundCount = 1;
 		else if (enable && roundCount <= Nr + 1) begin
-			state = stateOut;				
+			state = stateOut;
 			roundCount = roundCount + 6'b000001;
 		end
 	end
 endmodule
 
-module AESDecrypt_DUT();
+module AESDecrypt128_DUT();
+	localparam Nk = 4;
+	localparam Nr = 10;
+
 	wire [127:0] data = 128'h69c4e0d86a7b0430d8cdb78070b4c55a;
-	wire [127:0] key = 128'h000102030405060708090a0b0c0d0e0f;
-	wire [(11 * 128) - 1:0] allKeys;
+	wire [Nk * 32 - 1:0] key = 128'h000102030405060708090a0b0c0d0e0f;
+	wire [((Nr + 1) * 128) - 1:0] allKeys;
+	wire [127:0] out;
+	reg clk;
+
+	KeyExpansion #(Nk, Nr) ke(key, allKeys);
+	AESDecrypt #(Nk, Nr) aes(data, allKeys, out, clk, 1'b1);
+
+	initial begin
+		clk = 1'b1;
+		forever #10 clk = ~clk;
+	end
+endmodule
+
+module AESDecrypt192_DUT();
+	localparam Nk = 6;
+	localparam Nr = 12;
+
+	wire [127:0] data = 128'hdda97ca4864cdfe06eaf70a0ec0d7191;
+	wire [Nk * 32 - 1:0] key = 192'h000102030405060708090a0b0c0d0e0f1011121314151617;
+	wire [((Nr + 1) * 128) - 1:0] allKeys;
+	wire [127:0] out;
+	reg clk;
+
+	KeyExpansion #(Nk, Nr) ke(key, allKeys);
+	AESDecrypt #(Nk, Nr) aes(data, allKeys, out, clk, 1'b1);
+
+	initial begin
+		clk = 1'b1;
+		forever #10 clk = ~clk;
+	end
+endmodule
+
+module AESDecrypt256_DUT();
+	localparam Nk = 8;
+	localparam Nr = 14;
+
+	wire [127:0] data = 128'h8ea2b7ca516745bfeafc49904b496089;
+	wire [Nk * 32 - 1:0] key = 256'h000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f;
+	wire [((Nr + 1) * 128) - 1:0] allKeys;
 	wire [127:0] out;
 	reg clk;
 
@@ -59,7 +100,7 @@ module AESDecrypt_DUT();
 	AESDecrypt aes(data, allKeys, out, clk, 1, 0);
 
 	initial begin
-		clk = 0;
+		clk = 1'b1;
 		forever #10 clk = ~clk;
 	end
 endmodule
\ No newline at end of file
diff --git a/AESEncrypt.v b/AESEncrypt.v
index 083f860..0633962 100644
--- a/AESEncrypt.v
+++ b/AESEncrypt.v
@@ -1,6 +1,6 @@
 module AESEncrypt #(parameter Nk = 4, parameter Nr = 10) (data, allKeys, state, clk, reset);
 	input [127:0] data;
-	input [(11 * 128) - 1:0] allKeys;
+	input [((Nr + 1) * 128) - 1:0] allKeys;
 	input clk;
 	input reset;
 	output reg [127:0] state = 0; // Holds the state of the AES encryption
@@ -15,15 +15,15 @@ module AESEncrypt #(parameter Nk = 4, parameter Nr = 10) (data, allKeys, state,
 
 	// Instantiate AES modules needed for encryption
 	SubBytes sub(state, subByteWire);
-	ShiftRows shft(subByteWire, shiftRowsWire);	
-	MixColumns mix(shiftRowsWire, mixColumnsWire);	
-	AddRoundKey addkey(roundKeyInput , allKeys[(128 * roundCount) - 1 -: 128], stateOut);
+	ShiftRows shft(subByteWire, shiftRowsWire);
+	MixColumns mix(shiftRowsWire, mixColumnsWire);
+	AddRoundKey addkey(roundKeyInput , allKeys[((Nr + 1) * 128) - (roundCount - 1) * 128 - 1 -: 128], stateOut);
 
 	// Assign roundKeyInput based on roundCount
 	// roundCount = 1 -> Data
 	// roundCount = 2 to Nr -> mixColumnsWire
 	// roundCount = Nr + 1 -> shiftRowsWire
-	assign roundKeyInput = (roundCount == 1) ? data : (roundCount < Nr + 1) ? mixColumnsWire: shiftRowsWire;
+	assign roundKeyInput = (roundCount == 1) ? data : (roundCount < Nr + 1) ? mixColumnsWire : shiftRowsWire;
 
 	// Assign state to data on data change and reset roundCount
 	initial @(data) begin 
@@ -38,21 +38,62 @@ module AESEncrypt #(parameter Nk = 4, parameter Nr = 10) (data, allKeys, state,
 			roundCount = 1;
 		end 
 		else if (roundCount <= Nr + 1) begin
-			state = stateOut;				
+			state = stateOut;
 			roundCount = roundCount + 6'b000001;
 		end
 	end
 endmodule
 
-module AESEncrypt_DUT();
+module AESEncrypt128_DUT();
+	localparam Nk = 4;
+	localparam Nr = 10;
+
+	wire [127:0] data = 128'h00112233445566778899aabbccddeeff;
+	wire [Nk * 32 - 1:0] key = 128'h000102030405060708090a0b0c0d0e0f;
+	wire [((Nr + 1) * 128) - 1:0] allKeys;
+	wire [127:0] out;
+	reg clk;
+
+	KeyExpansion #(Nk, Nr) ke(key, allKeys);
+	AESEncrypt #(Nk, Nr) aes(data, allKeys, out, clk, 1'b0);
+
+	initial begin
+		clk = 0;
+		forever #10 clk = ~clk;
+	end
+endmodule
+
+module AESEncrypt192_DUT();
+	localparam Nk = 6;
+	localparam Nr = 12;
+
+	wire [127:0] data = 128'h00112233445566778899aabbccddeeff;
+	wire [Nk * 32 - 1:0] key = 192'h000102030405060708090a0b0c0d0e0f1011121314151617;
+	wire [((Nr + 1) * 128) - 1:0] allKeys;
+	wire [127:0] out;
+	reg clk;
+
+	KeyExpansion #(Nk, Nr) ke(key, allKeys);
+	AESEncrypt #(Nk, Nr) aes(data, allKeys, out, clk, 1'b0);
+
+	initial begin
+		clk = 0;
+		forever #10 clk = ~clk;
+	end
+endmodule
+
+module AESEncrypt256_DUT();
+	localparam Nk = 8;
+	localparam Nr = 14;
+
 	wire [127:0] data = 128'h00112233445566778899aabbccddeeff;
-	wire [127:0] key = 128'h000102030405060708090a0b0c0d0e0f;
-	wire [(11 * 128) - 1:0] allKeys;
+	wire [Nk * 32 - 1:0] key = 256'h000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f;
+	wire [((Nr + 1) * 128) - 1:0] allKeys;
 	wire [127:0] out;
 	reg clk;
 
-	KeyExpansion ke(key, allKeys);
-	AESEncrypt aes(data, allKeys, out, clk, 0);
+	KeyExpansion #(Nk, Nr) ke(key, allKeys);
+	AESEncrypt #(Nk, Nr) aes(data, allKeys, out, clk, 1'b0);
 
 	initial begin
 		clk = 0;
diff --git a/DecryptRound.v b/DecryptRound.v
index 220d191..bed4cc8 100644
--- a/DecryptRound.v
+++ b/DecryptRound.v
@@ -7,10 +7,10 @@ module DecryptRound(stateIn, key, stateOut);
 	wire [127:0] shiftRowsWire;
 	wire [127:0] afterRoundKey;
 
-	InvShiftRows shft(stateIn, shiftRowsWire);	
+	InvShiftRows shft(stateIn, shiftRowsWire);
 	InvSubBytes sub(shiftRowsWire, subByteWire);
 	AddRoundKey addkey(subByteWire, key, afterRoundKey);
-	InvMixColumns mix(afterRoundKey, stateOut);	
+	InvMixColumns mix(afterRoundKey, stateOut);
 endmodule 
 
 module LastDecryptRound(stateIn, key, stateOut);
@@ -21,7 +21,7 @@ module LastDecryptRound(stateIn, key, stateOut);
 	wire [127:0] subByteWire;
 	wire [127:0] shiftRowsWire;
 
-	InvShiftRows shft(stateIn, shiftRowsWire);	
+	InvShiftRows shft(stateIn, shiftRowsWire);
 	InvSubBytes sub(shiftRowsWire, subByteWire);
 	AddRoundKey addkey(subByteWire, key, stateOut);
 endmodule 
\ No newline at end of file
diff --git a/EncryptRound.v b/EncryptRound.v
index 6f5104b..083683a 100644
--- a/EncryptRound.v
+++ b/EncryptRound.v
@@ -8,8 +8,8 @@ module EncryptRound(stateIn, key, stateOut);
 	wire [127:0] mixColumnsWire;
 
 	SubBytes sub(stateIn, subByteWire);
-	ShiftRows shft(subByteWire, shiftRowsWire);	
-	MixColumns mix(shiftRowsWire, mixColumnsWire);	
+	ShiftRows shft(subByteWire, shiftRowsWire);
+	MixColumns mix(shiftRowsWire, mixColumnsWire);
 	AddRoundKey addkey(mixColumnsWire, key, stateOut);
 endmodule 
 
@@ -22,6 +22,6 @@ module LastEncryptRound(stateIn, key, stateOut);
 	wire [127:0] shiftRowsWire;
 
 	SubBytes sub(stateIn, subByteWire);
-	ShiftRows shft(subByteWire, shiftRowsWire);	
+	ShiftRows shft(subByteWire, shiftRowsWire);
 	AddRoundKey addkey(shiftRowsWire, key, stateOut);
 endmodule 
\ No newline at end of file
diff --git a/KeyExpansion.v b/KeyExpansion.v
index 79a0868..cffb2e3 100644
--- a/KeyExpansion.v
+++ b/KeyExpansion.v
@@ -1,86 +1,174 @@
-module KeyExpansionRound(roundCount, keyIn, keyOut);
+module KeyExpansionRound #(parameter Nk = 4, parameter Nr = 10) (roundCount, keyIn, keyOut);
     input [3:0] roundCount;
-    input [127:0] keyIn;
+    input [32 * Nk - 1:0] keyIn;
 
-    output [127:0] keyOut;
+    output [32 * Nk - 1:0] keyOut;
 
     genvar i;
 
-    // Split the key into 4 words
-    wire [31:0] words[3:0];
+    // Split the key into Nk words
+    wire [31:0] words[Nk - 1:0];
 
     generate
-        for (i = 0; i < 4; i = i + 1) begin: KeySplitLoop
-            assign words[i] = keyIn[127 - i * 32 -: 32];
+        for (i = 0; i < Nk; i = i + 1) begin: KeySplitLoop
+            assign words[i] = keyIn[(32 * Nk - 1) - i * 32 -: 32];
         end
     endgenerate
 
     // Rotate the words (rotWord)
-    wire [31:0] w3Rot = {words[3][23:0], words[3][31:24]};
+    wire [31:0] w3Rot = {words[Nk - 1][23:0], words[Nk - 1][31:24]};
 
     // Perform the substitution of the words (subWord)
     wire [31:0] w3Sub;
 
-    generate
-        for (i = 0; i < 4; i = i + 1) begin: SBoxLoop
-            SubTable sBox(w3Rot[i * 8 +: 8], w3Sub[i * 8 +: 8]);
+    generate 
+        for (i = 0; i < 4; i = i + 1) begin: SubWordLoop
+            SubTable subTable(w3Rot[8 * i +: 8], w3Sub[8 * i +: 8]);
         end
     endgenerate
 
     // Perform the XOR operation with the round constant (roundConstant)
-    reg [31:0] roundConstant;
-    always @(*)
-    begin
-        case (roundCount)
-            4'h1: roundConstant = 32'h01_000000;
-            4'h2: roundConstant = 32'h02_000000;
-            4'h3: roundConstant = 32'h04_000000;
-            4'h4: roundConstant = 32'h08_000000;
-            4'h5: roundConstant = 32'h10_000000;
-            4'h6: roundConstant = 32'h20_000000;
-            4'h7: roundConstant = 32'h40_000000;
-            4'h8: roundConstant = 32'h80_000000;
-            4'h9: roundConstant = 32'h1b_000000;
-            4'ha: roundConstant = 32'h36_000000;
-            default: roundConstant = 32'h00_000000;
-        endcase
-    end
+    wire [7:0] roundConstantStart = roundCount == 1 ? 8'h01
+                                        : roundCount == 2 ? 8'h02
+                                        : roundCount == 3 ? 8'h04
+                                        : roundCount == 4 ? 8'h08
+                                        : roundCount == 5 ? 8'h10
+                                        : roundCount == 6 ? 8'h20
+                                        : roundCount == 7 ? 8'h40
+                                        : roundCount == 8 ? 8'h80
+                                        : roundCount == 9 ? 8'h1b
+                                        : roundCount == 10 ? 8'h36
+                                        : roundCount == 11 ? 8'h6c
+                                        : roundCount == 12 ? 8'hd8
+                                        : roundCount == 13 ? 8'hab
+                                        : roundCount == 14 ? 8'h4d
+                                        : roundCount == 15 ? 8'h9a
+                                        : roundCount == 16 ? 8'h2f
+                                        : 8'h00;
+    wire [31:0] roundConstant = {roundConstantStart, 24'h00};
+
+    assign keyOut[32 * Nk - 1 -: 32] = words[0] ^ w3Sub ^ roundConstant; // XOR the first word with the round constant
+
+    // Perform SubWord transformation for i % Nk work (256 bits key only)
+    wire [31:0] wSub;
+    generate 
+        for (i = 0; i < 4; i = i + 1) begin: SubWordLoopForWSub
+            SubTable subTable(keyOut[(32 * Nk - 1) - 3 * 32 - i * 8 -: 8], wSub[(3 - i) * 8 +: 8]);
+        end
+    endgenerate
 
-    assign keyOut[127:96] = words[0] ^ w3Sub ^ roundConstant; // XOR the first word with the round constant
-    assign keyOut[95:64] = words[1] ^ keyOut[127:96]; // XOR the second word with the first word
-    assign keyOut[63:32] = words[2] ^ keyOut[95:64]; // XOR the third word with the second word
-    assign keyOut[31:0] = words[3] ^ keyOut[63:32]; // XOR the fourth word with the third word
+    generate
+        for (i = 1; i < Nk; i = i + 1) begin: KeyExpansionLoop
+            assign keyOut[(32 * Nk - 1) - i * 32 -: 32] = words[i] ^ (Nk == 8 && i == 4 ? wSub : keyOut[(32 * Nk - 1) - (i - 1) * 32 -: 32]); // XOR word i with word i - 1
+        end
+    endgenerate
 endmodule
 
-module KeyExpansion(keyIn, keysOut);    
-    input [127:0] keyIn;
-    output [(11 * 128) - 1:0] keysOut;
+module KeyExpansion #(parameter Nk = 4, parameter Nr = 10) (keyIn, keysOut);    
+    localparam rounds = (Nr == 10 ? 9 : (Nr == 12 ? 7 : 6));
+
+    input [(Nk * 32) - 1:0] keyIn;
+    output  [((Nr + 1) * 128) - 1:0] keysOut;
+
+    assign keysOut[((Nr + 1) * 128) - 1 -: (Nk * 32)] = keyIn;
 
-    assign keysOut[127:0] = keyIn;
-    
     // Perform the key expansion rounds (KeyExpansionRound)
     genvar i;
     generate
-        for (i = 0; i < 10; i = i + 1) begin : KeyExpansionLoop
-            KeyExpansionRound keyExpansionRound(i[3:0] + 4'b0001, keysOut[127 + i * 128 -: 128], keysOut[255 + i * 128 -: 128]);
+        for (i = 0; i < rounds; i = i + 1) begin: KeyExpansionRoundLoop
+            KeyExpansionRound #(Nk, Nr) keyExpansionRound(i[3:0] + 4'b0001, keysOut[((Nr + 1) * 128) - 1 - i * (Nk * 32) -: (Nk * 32)], keysOut[((Nr + 1) * 128) - 1 - (i + 1) * (Nk * 32) -: (Nk * 32)]);
         end
     endgenerate
+
+    // Perform the last key expansion round (LastKeyExpansionRound)
+    wire [Nk * 32 - 1:0] lastkey;
+    KeyExpansionRound #(Nk, Nr) lastKeyExpansionRound(rounds[3:0] + 4'b0001, keysOut[128 +: (Nk * 32)], lastkey);
+
+    assign keysOut[127:0] = lastkey[Nk * 32 - 1 -: 128];
 endmodule
 
-module KeyExpansion_DUT();
-    reg [127:0] keyIn;
-    wire [(11 * 128) - 1:0] keysOut;
+module KeyExpansion128_DUT();
+    localparam Nr = 10;
+    localparam Nk = 4;
+    localparam rounds = (Nr == 10 ? 9 : (Nr == 12 ? 7 : 6));
+
+    reg [(Nk * 32) - 1:0] keyIn;
+    wire [((Nr + 1) * 128) - 1:0] keysOut;
 
-    KeyExpansion keyExpansion(keyIn, keysOut);
+    KeyExpansion #(Nk, Nr) keyExpansion(keyIn, keysOut);
 
     initial begin
         keyIn = 128'h2b_7e_15_16_28_ae_d2_a6_ab_f7_15_88_09_cf_4f_3c; // 128-bit key from the AES standard document
     end
 
+    integer i;
+
     initial begin
-        $display("KeyExpansion_DUT");
+        $display("KeyExpansion128_DUT");
         $display("==================================");
         $display("keyIn = %h", keyIn);
-        $monitor("keysOut = %h, %h, %h, %h, %h, %h, %h, %h, %h, %h, %h", keysOut[127:0], keysOut[255:128], keysOut[383:256], keysOut[511:384], keysOut[639:512], keysOut[767:640], keysOut[895:768], keysOut[1023:896], keysOut[1151:1024], keysOut[1279:1152], keysOut[1407:1280]);
+        #10;
+
+        for (i = 0; i <= rounds; i = i + 1) begin: DisplayKeys
+            $display("keysOut[%0d] = %h", i, keysOut[((Nr + 1) * 128) - 1 - i * (Nk * 32) -: (Nk * 32)]);
+        end
+    end
+endmodule
+
+module KeyExpansion192_DUT();
+    localparam Nr = 12;
+    localparam Nk = 6;
+    localparam rounds = (Nr == 10 ? 9 : (Nr == 12 ? 7 : 6));
+
+    reg [(Nk * 32) - 1:0] keyIn;
+    wire [((Nr + 1) * 128) - 1:0] keysOut;
+
+    KeyExpansion #(Nk, Nr) keyExpansion (keyIn, keysOut);
+
+    initial begin
+        keyIn = 192'h8e_73_b0_f7_da_0e_64_52_c8_10_f3_2b_80_90_79_e5_62_f8_ea_d2_52_2c_6b_7b; // 192-bit key from the AES standard document
+    end
+
+    integer i;
+
+    initial begin
+        $display("KeyExpansion196_DUT");
+        $display("==================================");
+        $display("keyIn = %h", keyIn);
+        #10;
+
+        for (i = 0; i <= rounds; i = i + 1) begin: DisplayKeys
+            $display("keysOut[%0d] = %h", i, keysOut[((Nr + 1) * 128) - 1 - i * (Nk * 32) -: (Nk * 32)]);
+        end
+        $display("keysOut[%0d] = %h", rounds + 1, keysOut[127:0]);
+    end
+endmodule
+
+module KeyExpansion256_DUT();
+    localparam Nr = 14;
+    localparam Nk = 8;
+    localparam rounds = (Nr == 10 ? 9 : (Nr == 12 ? 7 : 6));
+
+    reg [(Nk * 32) - 1:0] keyIn;
+    wire [((Nr + 1) * 128) - 1:0] keysOut;
+
+    KeyExpansion #(Nk, Nr) keyExpansion (keyIn, keysOut);
+
+    initial begin
+        keyIn = 256'h60_3d_eb_10_15_ca_71_be_2b_73_ae_f0_85_7d_77_81_1f_35_2c_07_3b_61_08_d7_2d_98_10_a3_09_14_df_f4; // 256-bit key from the AES standard document
+    end
+
+    integer i;
+
+    initial begin
+        $display("KeyExpansion256_DUT");
+        $display("==================================");
+        $display("keyIn = %h", keyIn);
+        #10;
+
+        for (i = 0; i <= rounds; i = i + 1) begin: DisplayKeys
+            $display("keysOut[%0d] = %h", i, keysOut[((Nr + 1) * 128) - 1 - i * (Nk * 32) -: (Nk * 32)]);
+        end
+        $display("keysOut[%0d] = %h", rounds + 1, keysOut[127:0]);
     end
 endmodule
\ No newline at end of file