diff --git a/circuits/aes_256_encrypt.circom b/circuits/aes_256_encrypt.circom
index be01ac1..596fa9d 100644
--- a/circuits/aes_256_encrypt.circom
+++ b/circuits/aes_256_encrypt.circom
@@ -36,6 +36,7 @@ template AES256Encrypt()
     component bits2num_1[13][4][4];
     component num2bits_1[13][4][4];
     component xor_3[13][4][32];
+    component emulated_aesenc_enc_selector[13][4][4];
 
     for(i=0; i<13; i++)
     {
@@ -49,7 +50,15 @@ template AES256Encrypt()
                 
                 for(l=0; l<8; l++) bits2num_1[i][j][k].in[l] <== s_tmp[k*8+7-l];
 
-                num2bits_1[i][j][k].in <-- emulated_aesenc_enc_table(k, bits2num_1[i][j][k].out);
+                emulated_aesenc_enc_selector[i][j][k] = IndexSelector(256);
+
+                for(var iter = 0; iter<256; iter++){
+                    emulated_aesenc_enc_selector[i][j][k].in[iter] <== emulated_aesenc_enc_table(k, iter);
+                }
+
+                emulated_aesenc_enc_selector[i][j][k].index <== bits2num_1[i][j][k].out;
+
+                num2bits_1[i][j][k].in <== emulated_aesenc_enc_selector[i][j][k].out;
 
                 if(k==0)
                 {
diff --git a/circuits/aes_256_key_expansion.circom b/circuits/aes_256_key_expansion.circom
index 676420b..b2a8727 100644
--- a/circuits/aes_256_key_expansion.circom
+++ b/circuits/aes_256_key_expansion.circom
@@ -38,34 +38,49 @@ template AES256KeyExpansion()
     component bits2num_1[60][4];
     component num2bits_1[60][5];
     component xor_2[60][8];
-
+    component i_select[60][4];
+    component r_con_select[60];
 
     while(i<(Nb*(Nr+1)))
     {
         var tmp[32];
         tmp = ks[i-1];
+        i_select[i][0] = IndexSelector(256);
+        i_select[i][1] = IndexSelector(256);
+        i_select[i][2] = IndexSelector(256);
+        i_select[i][3] = IndexSelector(256);
+        for(var iter = 0; iter < 256; iter++){
+            i_select[i][0].in[iter] <== emulated_aesenc_rijndael_sbox(iter);
+            i_select[i][1].in[iter] <== emulated_aesenc_rijndael_sbox(iter);
+            i_select[i][2].in[iter] <== emulated_aesenc_rijndael_sbox(iter);
+            i_select[i][3].in[iter] <== emulated_aesenc_rijndael_sbox(iter);
+        }
 
         if(i%Nk == 0)
         {
             bits2num_1[i][0] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][0].in[j] <== tmp[7-j];
             num2bits_1[i][0] = Num2Bits(8);
-            num2bits_1[i][0].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][0].out);
+            i_select[i][0].index <== bits2num_1[i][0].out;
+            num2bits_1[i][0].in <== i_select[i][0].out;
 
             bits2num_1[i][1] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][1].in[j] <== tmp[7-j+8];
             num2bits_1[i][1] = Num2Bits(8);
-            num2bits_1[i][1].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][1].out);
+            i_select[i][1].index <== bits2num_1[i][1].out;
+            num2bits_1[i][1].in <== i_select[i][1].out;
 
             bits2num_1[i][2] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][2].in[j] <== tmp[7-j+16];
             num2bits_1[i][2] = Num2Bits(8);
-            num2bits_1[i][2].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][2].out);
+            i_select[i][2].index <== bits2num_1[i][2].out;
+            num2bits_1[i][2].in <== i_select[i][2].out;
 
             bits2num_1[i][3] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][3].in[j] <== tmp[7-j+24];
             num2bits_1[i][3] = Num2Bits(8);
-            num2bits_1[i][3].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][3].out);
+            i_select[i][3].index <== bits2num_1[i][3].out;
+            num2bits_1[i][3].in <== i_select[i][3].out;
 
             for(j=0; j<8; j++) tmp[j+24] = num2bits_1[i][0].out[7-j];
             for(j=0; j<8; j++) tmp[j] = num2bits_1[i][1].out[7-j];
@@ -73,7 +88,13 @@ template AES256KeyExpansion()
             for(j=0; j<8; j++) tmp[j+16] = num2bits_1[i][3].out[7-j];
 
             num2bits_1[i][4] = Num2Bits(8);
-            num2bits_1[i][4].in <== rcon[i/Nk-1];
+            
+            r_con_select[i] = IndexSelector(10);
+            for(var iter = 0; iter<10; iter++){
+                r_con_select[i].in[iter] <== rcon[iter];
+            }
+            r_con_select[i].index <== i/Nk-1;
+            num2bits_1[i][4].in <== r_con_select[i].out;
 
             for(j=0; j<8; j++)
             {
@@ -89,22 +110,30 @@ template AES256KeyExpansion()
             bits2num_1[i][0] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][0].in[j] <== tmp[7-j];
             num2bits_1[i][0] = Num2Bits(8);
-            num2bits_1[i][0].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][0].out);
+            
+            i_select[i][0].index <== bits2num_1[i][0].out;
+            num2bits_1[i][0].in <== i_select[i][0].out;
 
             bits2num_1[i][1] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][1].in[j] <== tmp[7-j+8];
             num2bits_1[i][1] = Num2Bits(8);
-            num2bits_1[i][1].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][1].out);
+            
+            i_select[i][1].index <== bits2num_1[i][1].out;
+            num2bits_1[i][1].in <== i_select[i][1].out;
 
             bits2num_1[i][2] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][2].in[j] <== tmp[7-j+16];
             num2bits_1[i][2] = Num2Bits(8);
-            num2bits_1[i][2].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][2].out);
+            
+            i_select[i][2].index <== bits2num_1[i][2].out;
+            num2bits_1[i][2].in <== i_select[i][2].out;
 
             bits2num_1[i][3] = Bits2Num(8);
             for(j=0; j<8; j++) bits2num_1[i][3].in[j] <== tmp[7-j+24];
             num2bits_1[i][3] = Num2Bits(8);
-            num2bits_1[i][3].in <-- emulated_aesenc_rijndael_sbox(bits2num_1[i][3].out);
+            
+            i_select[i][3].index <== bits2num_1[i][3].out;
+            num2bits_1[i][3].in <== i_select[i][3].out;
 
             for(j=0; j<8; j++) tmp[j] = num2bits_1[i][0].out[7-j];
             for(j=0; j<8; j++) tmp[j+8] = num2bits_1[i][1].out[7-j];
diff --git a/circuits/aes_emulation.circom b/circuits/aes_emulation.circom
index ca60768..234dd11 100644
--- a/circuits/aes_emulation.circom
+++ b/circuits/aes_emulation.circom
@@ -33,7 +33,21 @@ template EmulatedAesencSubstituteBytes()
 {
     signal input in[16];
     signal output out[16];
-
-    for(var i=0; i<16; i++) out[i] <-- emulated_aesenc_rijndael_sbox(in[i]);
     
+    component aesenc_select[16];
+
+    for(var iter = 0; iter<16; iter++){
+        aesenc_select[iter] = IndexSelector(256);
+    }
+
+    for(var i = 0; i<16; i++){
+        for(var j=0; j<256; j++){
+            aesenc_select[i].in[j] <== emulated_aesenc_rijndael_sbox(j);
+        }
+    }
+
+    for(var i=0; i<16; i++) {
+        aesenc_select[i].index <== in[i];
+        out[i] <== aesenc_select[i].out;
+    }
 }
diff --git a/circuits/helper_functions.circom b/circuits/helper_functions.circom
index 1c1d8fe..03b1872 100644
--- a/circuits/helper_functions.circom
+++ b/circuits/helper_functions.circom
@@ -230,22 +230,37 @@ template SumMultiple(n) {
     sum <== sums[n-1];
 }
 
+template IsZero2() {
+    signal input in;
+    signal output out;
+    signal inv;
+    inv <-- in!=0 ? 1/in : 0;
+    out <-- -in*inv + 1;// this is the issue
+    in*out === 0; 
+}
+
+template IsEqual2() {
+    signal input in[2];
+    signal output out;
+    component isz = IsZero2();
+    in[1] - in[0] --> isz.in;
+    isz.out ==> out;
+}
+
 template IndexSelector(total) {
     signal input in[total];
     signal input index;
     signal output out;
 
-    //maybe add (index<total) check later when we decide number of bits
-
     component calcTotal = SumMultiple(total);
     component equality[total];
 
     for(var i=0; i<total; i++){
-        equality[i] = IsEqual();
-        equality[i].in[0] <== i;
-        equality[i].in[1] <== index;
-        calcTotal.nums[i] <== equality[i].out * in[i];
+        equality[i] = IsEqual2();
+        equality[i].in[0] <-- i;
+        equality[i].in[1] <-- index;
+        calcTotal.nums[i] <-- equality[i].out * in[i];
     }
 
     out <== calcTotal.sum;
-}
+}
\ No newline at end of file