diff --git a/openair1/PHY/INIT/nr_init.c b/openair1/PHY/INIT/nr_init.c
index c10a0071a7..9a83ed801b 100644
--- a/openair1/PHY/INIT/nr_init.c
+++ b/openair1/PHY/INIT/nr_init.c
@@ -113,58 +113,71 @@ int init_codebook_gNB(PHY_VARS_gNB *gNB) {
     //    X X X X ... X
     //   |<-----N1---->|
     int x_polarization = gNB->ap_XP;
+    AssertFatal(CSI_RS_antenna_ports == N1 * N2 * x_polarization,
+                "Nb of antenna ports at PHY %d does not correspond to what passed down with fapi %d\n",
+                 N1 * N2 * x_polarization, CSI_RS_antenna_ports);
+
     //Get the uniform planar array parameters
     // To be confirmed
-    int O2 = N2 > 1? 4 : 1; //Vertical beam oversampling (1 or 4)
+    int O2 = N2 > 1 ? 4 : 1; //Vertical beam oversampling (1 or 4)
     int O1 = CSI_RS_antenna_ports > 2 ? 4 : 1; //Horizontal beam oversampling (1 or 4)
-    AssertFatal(CSI_RS_antenna_ports == N1*N2*x_polarization,
-                "Nb of antenna ports at PHY %d does not correspond to what passed down with fapi %d\n",
-                 N1*N2*x_polarization, CSI_RS_antenna_ports);
+
+    // num of allowed k1 and k2 according to 5.2.2.2.1-3 and -4 in 38.214
+    int K1;
+    if(N2 == N1 || N1 == 2)
+      K1 = 2;
+    else if (N2 == 1)
+      K1 = 5;
+    else
+      K1 = 3;
+    int K2 = N2 > 1 ? 2 : 1;
 
     // Generation of codebook Type1 with codebookMode 1 (CSI_RS_antenna_ports < 16)
     if (CSI_RS_antenna_ports < 16) {
       //Generate DFT vertical beams
       //ll: index of a vertical beams vector (represented by i1_1 in TS 38.214)
-      double complex v[N1*O1][N1];
-      for (int ll=0; ll<N1*O1; ll++) { //i1_1
-        for (int nn=0; nn<N1; nn++) {
-          v[ll][nn] = cexp(I*(2*M_PI*nn*ll)/(N1*O1));
-          //printf("v[%d][%d] = %f +j %f\n", ll,nn, creal(v[ll][nn]),cimag(v[ll][nn]));
+      const int max_l = N1 * O1 + (K1 - 1) * O1;
+      double complex v[max_l][N1];
+      for (int ll = 0; ll < max_l; ll++) { //i1_1
+        for (int nn = 0; nn < N1; nn++) {
+          v[ll][nn] = cexp(I * (2 * M_PI * nn * ll) / (N1 * O1));
+          LOG_D(PHY,"v[%d][%d] = %f +j %f\n", ll, nn, creal(v[ll][nn]), cimag(v[ll][nn]));
         }
       }
       //Generate DFT Horizontal beams
       //mm: index of a Horizontal beams vector (represented by i1_2 in TS 38.214)
-      double complex u[N2*O2][N2];
-      for (int mm=0; mm<N2*O2; mm++) { //i1_2
-        for (int nn=0; nn<N2; nn++) {
-          u[mm][nn] = cexp(I*(2*M_PI*nn*mm)/(N2*O2));
-              //printf("u[%d][%d] = %f +j %f\n", mm,nn, creal(u[mm][nn]),cimag(u[mm][nn]));
+      const int max_m = N2 * O2 + (K2 - 1) * O2;
+      double complex u[max_m][N2];
+      for (int mm = 0; mm < max_m; mm++) { //i1_2
+        for (int nn = 0; nn < N2; nn++) {
+          u[mm][nn] = cexp(I * (2 * M_PI * nn * mm) / (N2 * O2));
+          LOG_D(PHY,"u[%d][%d] = %f +j %f\n", mm, nn, creal(u[mm][nn]), cimag(u[mm][nn]));
         }
       }
       //Generate co-phasing angles
       //i_2: index of a co-phasing vector
       //i1_1, i1_2, and i_2 are reported from UEs
       double complex theta_n[4];
-      for (int nn=0; nn<4; nn++) {
-        theta_n[nn] = cexp(I*M_PI*nn/2);
-        //printf("theta_n[%d] = %f +j %f\n", nn, creal(theta_n[nn]),cimag(theta_n[nn]));
+      for (int nn = 0; nn < 4; nn++) {
+        theta_n[nn] = cexp(I * M_PI * nn / 2);
+        LOG_D(PHY,"theta_n[%d] = %f +j %f\n", nn, creal(theta_n[nn]), cimag(theta_n[nn]));
       }
       //Kronecker product v_lm
-      double complex v_lm[N1*O1][N2*O2][N2*N1];
+      double complex v_lm[max_l][max_m][N2 * N1];
       //v_ll_mm_codebook denotes the elements of a precoding matrix W_i1,1_i_1,2
-      for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-        for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-          for (int nn1=0; nn1<N1; nn1++) {
-            for (int nn2=0; nn2<N2; nn2++) {
-              //printf("indx %d \n",nn1*N2+nn2);
-              v_lm[ll][mm][nn1*N2+nn2] = v[ll][nn1]*u[mm][nn2];
-              //printf("v_lm[%d][%d][%d] = %f +j %f\n",ll,mm, nn1*N2+nn2, creal(v_lm[ll][mm][nn1*N2+nn2]),cimag(v_lm[ll][mm][nn1*N2+nn2]));
+      for(int ll = 0; ll < max_l; ll++) { //i_1_1
+        for (int mm = 0; mm < max_m; mm++) { //i_1_2
+          for (int nn1 = 0; nn1 < N1; nn1++) {
+            for (int nn2 = 0; nn2 < N2; nn2++) {
+              v_lm[ll][mm][nn1 * N2 + nn2] = v[ll][nn1] * u[mm][nn2];
+              LOG_D(PHY,"v_lm[%d][%d][%d] = %f +j %f\n",ll, mm, nn1 * N2 + nn2, creal(v_lm[ll][mm][nn1*N2+nn2]), cimag(v_lm[ll][mm][nn1*N2+nn2]));
             }
           }
         }
       }
 
-      int max_mimo_layers = (CSI_RS_antenna_ports<NR_MAX_NB_LAYERS) ? CSI_RS_antenna_ports : NR_MAX_NB_LAYERS;
+      int max_mimo_layers = (CSI_RS_antenna_ports < NR_MAX_NB_LAYERS) ? CSI_RS_antenna_ports : NR_MAX_NB_LAYERS;
+      AssertFatal(max_mimo_layers <= 4, "Max number of layers supported is 4\n");
 
       gNB->nr_mimo_precoding_matrix = (int32_t ***)malloc16(max_mimo_layers * sizeof(int32_t **));
       int32_t ***mat = gNB->nr_mimo_precoding_matrix;
@@ -230,22 +243,11 @@ int init_codebook_gNB(PHY_VARS_gNB *gNB) {
       //Compute the code book size for generating 2 layers out of Tx antenna ports
 
       //pmi_size is computed as follows
-      gNB->pmiq_size[1] = 1;//1 for unity matrix
-      for(int llb=0; llb<N1*O1; llb++) { //i_1_1
-        for (int mmb=0; mmb<N2*O2; mmb++) { //i_1_2
-          for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-            for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-              for (int nn=0; nn<2; nn++) {
-                if((llb != ll) || (mmb != mm) || ((N1 == 1) && (N2 == 1))) gNB->pmiq_size[1] += 1;
-              }
-            }
-          }
-        }
-      }
-      mat[1] = (int32_t **)malloc16(gNB->pmiq_size[1]*sizeof(int32_t *));
+      gNB->pmiq_size[1] = 1 + (2 * N1 * O1 * N2 * O2 * K1 * K2);
+      mat[1] = (int32_t **)malloc16(gNB->pmiq_size[1] * sizeof(int32_t *));
 
       //pmi=0 corresponds to unit matrix
-      mat[1][0] = (int32_t *)calloc((2*N1*N2)*(2),sizeof(int32_t));
+      mat[1][0] = (int32_t *)calloc((2 * N1 * N2) * 2, sizeof(int32_t));
       for(int j_col=0; j_col<2; j_col++) { //2 layers
         for (int i_rows=0; i_rows<2*N1*N2; i_rows++) { //2-x polarized antenna
           if(j_col==i_rows) {
@@ -256,52 +258,50 @@ int init_codebook_gNB(PHY_VARS_gNB *gNB) {
 
       //pmi=1,...,pmi_size, we construct
       int pmiq = 0;
-      for(int llb=0; llb<N1*O1; llb++) { //i_1_1
-        for (int mmb=0; mmb<N2*O2; mmb++) { //i_1_2
-          for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-            for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-              for (int nn=0; nn<2; nn++) {
-                if((llb != ll) || (mmb != mm) || ((N1 == 1) && (N2 == 1))){
-                  pmiq += 1;
-                  mat[1][pmiq] = (int32_t *)malloc16((2*N1*N2)*(2)*sizeof(int32_t));
-                  LOG_D(PHY, "layer 2 Codebook pmiq = %d\n",pmiq);
-                  for(int j_col=0; j_col<2; j_col++) {
-                    if (j_col==0) {
-                      llc = llb;
-                      mmc = mmb;
-                      phase_sign = 1;
-                    }
-                    if (j_col==1) {
-                      llc = ll;
-                      mmc = mm;
-                      phase_sign = -1;
-                    }
-                    for (int i_rows=0; i_rows<N1*N2; i_rows++) {
-                      res_code=sqrt(1/(double)(2*CSI_RS_antenna_ports))*v_lm[llc][mmc][i_rows];
-                      if (creal(res_code)>0)
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
-                      else
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
-                      if (cimag(res_code)>0)
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
-                      else
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
-                      LOG_D(PHY, "2 Layer Precoding Matrix[1][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
-                            pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[1][pmiq][i_rows*2+j_col])[0],((short*) &mat[1][pmiq][i_rows*2+j_col])[1]);
-                    }
-                    for (int i_rows=N1*N2; i_rows<2*N1*N2; i_rows++) {
-                      res_code=sqrt(1/(double)(2*CSI_RS_antenna_ports))*(phase_sign)*theta_n[nn]*v_lm[llc][mmc][i_rows-N1*N2];
-                      if (creal(res_code)>0)
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
-                      else
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
-                      if (cimag(res_code)>0)
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
-                      else
-                        ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
-                      LOG_D(PHY, "2 Layer Precoding Matrix[1][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
-                            pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[1][pmiq][i_rows*2+j_col])[0],((short*) &mat[1][pmiq][i_rows*2+j_col])[1]);
-                    }
+      for(int ll = 0; ll < N1 * O1; ll++) { //i_1_1
+        for (int mm = 0; mm < N2 * O2; mm++) { // i_1_2
+          for(int k1 = 0; k1 < K1; k1++) {
+            for (int k2 = 0; k2 < K2; k2++) {
+              for (int nn = 0; nn < 2; nn++) {  // i_2
+                pmiq += 1;
+                mat[1][pmiq] = (int32_t *)malloc16((2 * N1 * N2) * 2 * sizeof(int32_t));
+                LOG_D(PHY, "layer 2 Codebook pmiq = %d\n", pmiq);
+                for(int j_col = 0; j_col < 2; j_col++) {
+                  if (j_col == 0) {
+                    llc = ll;
+                    mmc = mm;
+                    phase_sign = 1;
+                  }
+                  if (j_col == 1) {
+                    llc = ll + k1 * O1;
+                    mmc = mm + k2 * O2;
+                    phase_sign = -1;
+                  }
+                  for (int i_rows=0; i_rows<N1*N2; i_rows++) {
+                    res_code=sqrt(1/(double)(2*CSI_RS_antenna_ports))*v_lm[llc][mmc][i_rows];
+                    if (creal(res_code)>0)
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
+                    else
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
+                    if (cimag(res_code)>0)
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
+                    else
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
+                    LOG_D(PHY, "2 Layer Precoding Matrix[1][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
+                          pmiq, i_rows, j_col, creal(res_code), cimag(res_code), ((short*) &mat[1][pmiq][i_rows*2+j_col])[0], ((short*) &mat[1][pmiq][i_rows*2+j_col])[1]);
+                  }
+                  for (int i_rows=N1*N2; i_rows<2*N1*N2; i_rows++) {
+                    res_code=sqrt(1/(double)(2*CSI_RS_antenna_ports))*(phase_sign)*theta_n[nn]*v_lm[llc][mmc][i_rows-N1*N2];
+                    if (creal(res_code)>0)
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
+                    else
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
+                    if (cimag(res_code)>0)
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
+                    else
+                      ((short*) &mat[1][pmiq][i_rows*2+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
+                    LOG_D(PHY, "2 Layer Precoding Matrix[1][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
+                          pmiq, i_rows, j_col, creal(res_code), cimag(res_code), ((short*) &mat[1][pmiq][i_rows*2+j_col])[0], ((short*) &mat[1][pmiq][i_rows*2+j_col])[1]);
                   }
                 }
               }
@@ -312,85 +312,72 @@ int init_codebook_gNB(PHY_VARS_gNB *gNB) {
 
       //Table 5.2.2.2.1-7:
       //Codebook for 3-layer CSI reporting using antenna ports 3000 to 2999+PCSI-RS
-      if(max_mimo_layers>=3) {
+      if(max_mimo_layers >= 3) {
 
         //pmi_size is computed as follows
-        gNB->pmiq_size[2] = 1;//unity matrix
-        for(int llb=0; llb<N1*O1; llb++) { //i_1_1
-          for (int mmb=0; mmb<N2*O2; mmb++) { //i_1_2
-            for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-              for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-                for (int nn=0; nn<2; nn++) {
-                  if((llb != ll) || (mmb != mm)) gNB->pmiq_size[2] += 1;
-                }
-              }
-            }
-          }
-        }
-        mat[2] = (int32_t **)malloc16(gNB->pmiq_size[2]*sizeof(int32_t *));
+        gNB->pmiq_size[2] = 1 + (2 * N1 * O1 * N2 * O2 * K1 * K2);
+        mat[2] = (int32_t **)malloc16(gNB->pmiq_size[2] * sizeof(int32_t *));
         //pmi=0 corresponds to unit matrix
-        mat[2][0] = (int32_t *)calloc((2*N1*N2)*(3),sizeof(int32_t));
-        for(int j_col=0; j_col<3; j_col++) { //3 layers
-          for (int i_rows=0; i_rows<2*N1*N2; i_rows++) { //2-x polarized antenna
-            if(j_col==i_rows) {
-              mat[2][0][i_rows*3+j_col] = 0x7fff;
+        mat[2][0] = (int32_t *)calloc(2 * N1 * N2 * 3, sizeof(int32_t));
+        for(int j_col = 0; j_col < 3; j_col++) { //3 layers
+          for (int i_rows = 0; i_rows < 2 * N1 * N2; i_rows++) { //2-x polarized antenna
+            if(j_col == i_rows) {
+              mat[2][0][i_rows * 3 + j_col] = 0x7fff;
             }
           }
         }
 
         pmiq = 0;
         //pmi=1,...,pmi_size are computed as follows
-        for(int llb=0; llb<N1*O1; llb++) { //i_1_1
-          for (int mmb=0; mmb<N2*O2; mmb++) { //i_1_2
-            for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-              for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-                for (int nn=0; nn<2; nn++) {
-                  if((llb != ll) || (mmb != mm)){
-                    pmiq += 1;
-                    mat[2][pmiq] = (int32_t *)malloc16((2*N1*N2)*(3)*sizeof(int32_t));
-                    LOG_D(PHY, "layer 3 Codebook pmiq = %d\n",pmiq);
-                    for(int j_col=0; j_col<3; j_col++) {
-                      if (j_col==0) {
-                        llc = llb;
-                        mmc = mmb;
-                        phase_sign = 1;
-                      }
-                      if (j_col==1) {
-                        llc = ll;
-                        mmc = mm;
-                        phase_sign = 1;
-                      }
-                      if (j_col==2) {
-                        llc = ll;
-                        mmc = mm;
-                        phase_sign = -1;
-                      }
-                      for (int i_rows=0; i_rows<N1*N2; i_rows++) {
-                        res_code=sqrt(1/(double)(3*CSI_RS_antenna_ports))*v_lm[llc][mmc][i_rows];
-                        if (creal(res_code)>0)
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
-                        if (cimag(res_code)>0)
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
-                        LOG_D(PHY, "3 Layer Precoding Matrix[2][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
-                              pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[2][pmiq][i_rows*3+j_col])[0],((short*) &mat[2][pmiq][i_rows*3+j_col])[1]);
-                      }
-                      for (int i_rows=N1*N2; i_rows<2*N1*N2; i_rows++) {
-                        res_code=sqrt(1/(double)(3*CSI_RS_antenna_ports))*(phase_sign)*theta_n[nn]*v_lm[llc][mmc][i_rows-N1*N2];
-                        if (creal(res_code)>0)
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
-                        if (cimag(res_code)>0)
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
-                        LOG_D(PHY, "3 Layer Precoding Matrix[2][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
-                              pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[2][pmiq][i_rows*3+j_col])[0],((short*) &mat[2][pmiq][i_rows*3+j_col])[1]);
-                      }
+        for(int ll = 0; ll < N1 * O1; ll++) { //i_1_1
+          for (int mm = 0; mm < N2 * O2; mm++) { // i_1_2
+            for(int k1 = 0; k1 < K1; k1++) {
+              for (int k2 = 0; k2 < K2; k2++) {
+                for (int nn = 0; nn < 2; nn++) {  // i_2
+                  pmiq += 1;
+                  mat[2][pmiq] = (int32_t *)malloc16((2 * N1 * N2 * 3) * sizeof(int32_t));
+                  LOG_D(PHY, "layer 3 Codebook pmiq = %d\n",pmiq);
+                  for(int j_col = 0; j_col < 3; j_col++) {
+                    if (j_col == 0) {
+                      llc = ll;
+                      mmc = mm;
+                      phase_sign = 1;
+                    }
+                    if (j_col==1) {
+                      llc = ll + k1 * O1;
+                      mmc = mm + k2 * O2;
+                      phase_sign = 1;
+                    }
+                    if (j_col==2) {
+                      llc = ll;
+                      mmc = mm;
+                      phase_sign = -1;
+                    }
+                    for (int i_rows=0; i_rows<N1*N2; i_rows++) {
+                      res_code=sqrt(1/(double)(3*CSI_RS_antenna_ports))*v_lm[llc][mmc][i_rows];
+                      if (creal(res_code)>0)
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
+                      if (cimag(res_code)>0)
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
+                      LOG_D(PHY, "3 Layer Precoding Matrix[2][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
+                            pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[2][pmiq][i_rows*3+j_col])[0],((short*) &mat[2][pmiq][i_rows*3+j_col])[1]);
+                    }
+                    for (int i_rows=N1*N2; i_rows<2*N1*N2; i_rows++) {
+                      res_code=sqrt(1/(double)(3*CSI_RS_antenna_ports))*(phase_sign)*theta_n[nn]*v_lm[llc][mmc][i_rows-N1*N2];
+                      if (creal(res_code)>0)
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
+                      if (cimag(res_code)>0)
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[2][pmiq][i_rows*3+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
+                      LOG_D(PHY, "3 Layer Precoding Matrix[2][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
+                            pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[2][pmiq][i_rows*3+j_col])[0],((short*) &mat[2][pmiq][i_rows*3+j_col])[1]);
                     }
                   }
                 }
@@ -402,92 +389,76 @@ int init_codebook_gNB(PHY_VARS_gNB *gNB) {
 
       //Table 5.2.2.2.1-8:
       //Codebook for 4-layer CSI reporting using antenna ports 3000 to 2999+PCSI-RS
-      if(max_mimo_layers>=4) {
+      if(max_mimo_layers >= 4) {
         //pmi_size is computed as follows
-        gNB->pmiq_size[3] = 1;//unity matrix
-        for(int llb=0; llb<N1*O1; llb++) { //i_1_1
-          for (int mmb=0; mmb<N2*O2; mmb++) { //i_1_2
-            for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-              for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-                for (int nn=0; nn<2; nn++) {
-                  if((llb != ll) || (mmb != mm)) gNB->pmiq_size[3] += 1;
-                }
-              }
-            }
-          }
-        }
-
-        mat[3] = (int32_t **)malloc16(gNB->pmiq_size[3]*sizeof(int32_t *));
+        gNB->pmiq_size[3] = 1 + (2 * N1 * O1 * N2 * O2 * K1 * K2);
+        mat[3] = (int32_t **)malloc16(gNB->pmiq_size[3] * sizeof(int32_t *));
         //pmi=0 corresponds to unit matrix
-        mat[3][0] = (int32_t *)calloc((2*N1*N2)*(4),sizeof(int32_t));
-        for(int j_col=0; j_col<4; j_col++) { //4 layers
-          for (int i_rows=0; i_rows<2*N1*N2; i_rows++) { //2-x polarized antenna
-            if(j_col==i_rows) {
-              mat[3][0][i_rows*4+j_col] = 0x7fff;
+        mat[3][0] = (int32_t *)calloc((2 * N1 * N2 * 4), sizeof(int32_t));
+        for(int j_col = 0; j_col < 4; j_col++) { //4 layers
+          for (int i_rows = 0; i_rows < 2 * N1 * N2; i_rows++) { //2-x polarized antenna
+            if(j_col == i_rows) {
+              mat[3][0][i_rows * 4 + j_col] = 0x7fff;
             }
           }
         }
 
         pmiq = 0;
         //pmi=1,...,pmi_size are computed as follows
-        for(int llb=0; llb<N1*O1; llb++) { //i_1_1
-          for (int mmb=0; mmb<N2*O2; mmb++) { //i_1_2
-            for(int ll=0; ll<N1*O1; ll++) { //i_1_1
-              for (int mm=0; mm<N2*O2; mm++) { //i_1_2
-                for (int nn=0; nn<2; nn++) {
-
-                  if((llb != ll) || (mmb != mm)){
-                    pmiq += 1;
-                    mat[3][pmiq] = (int32_t *)malloc16((2*N1*N2)*4*sizeof(int32_t));
-                    LOG_D(PHY, "layer 4 pmiq = %d\n",pmiq);
-                    for(int j_col=0; j_col<4; j_col++) {
-                      if (j_col==0) {
-                        llc = llb;
-                        mmc = mmb;
-                        phase_sign = 1;
-                      }
-                      if (j_col==1) {
-                        llc = ll;
-                        mmc = mm;
-                        phase_sign = 1;
-                      }
-                      if (j_col==2) {
-                        llc = llb;
-                        mmc = mmb;
-                        phase_sign = -1;
-                      }
-                      if (j_col==3) {
-                        llc = ll;
-                        mmc = mm;
-                        phase_sign = -1;
-                      }
-                      for (int i_rows=0; i_rows<N1*N2; i_rows++) {
-                        res_code=sqrt(1/(double)(4*CSI_RS_antenna_ports))*v_lm[llc][mmc][i_rows];
-                        if (creal(res_code)>0)
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
-                        if (cimag(res_code)>0)
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
-                        LOG_D(PHY, "4 Layer Precoding Matrix[3][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
-                              pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[3][pmiq][i_rows*4+j_col])[0],((short*) &mat[3][pmiq][i_rows*4+j_col])[1]);
-                      }
-
-                      for (int i_rows=N1*N2; i_rows<2*N1*N2; i_rows++) {
-                        res_code=sqrt(1/(double)(4*CSI_RS_antenna_ports))*(phase_sign)*theta_n[nn]*v_lm[llc][mmc][i_rows-N1*N2];
-                        if (creal(res_code)>0)
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
-                        if (cimag(res_code)>0)
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
-                        else
-                          ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
-                        LOG_D(PHY, "4 Layer Precoding Matrix[3][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
-                              pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[3][pmiq][i_rows*4+j_col])[0],((short*) &mat[3][pmiq][i_rows*4+j_col])[1]);
-                      }
+        for(int ll = 0; ll < N1 * O1; ll++) { //i_1_1
+          for (int mm = 0; mm < N2 * O2; mm++) { // i_1_2
+            for(int k1 = 0; k1 < K1; k1++) {
+              for (int k2 = 0; k2 < K2; k2++) {
+                for (int nn = 0; nn < 2; nn++) {  // i_2
+                  pmiq += 1;
+                  mat[3][pmiq] = (int32_t *)malloc16((2*N1*N2)*4*sizeof(int32_t));
+                  LOG_D(PHY, "layer 4 pmiq = %d\n",pmiq);
+                  for(int j_col=0; j_col<4; j_col++) {
+                    if (j_col==0) {
+                      llc = ll;
+                      mmc = mm;
+                      phase_sign = 1;
+                    }
+                    if (j_col==1) {
+                      llc = ll + k1 * O1;
+                      mmc = mm + k2 * O2;
+                      phase_sign = 1;
+                    }
+                    if (j_col==2) {
+                      llc = ll;
+                      mmc = mm;
+                      phase_sign = -1;
+                    }
+                    if (j_col==3) {
+                      llc = ll + k1 * O1;
+                      mmc = mm + k2 * O2;
+                      phase_sign = -1;
+                    }
+                    for (int i_rows=0; i_rows<N1*N2; i_rows++) {
+                      res_code=sqrt(1/(double)(4*CSI_RS_antenna_ports))*v_lm[llc][mmc][i_rows];
+                      if (creal(res_code)>0)
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
+                      if (cimag(res_code)>0)
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
+                      LOG_D(PHY, "4 Layer Precoding Matrix[3][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
+                            pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[3][pmiq][i_rows*4+j_col])[0],((short*) &mat[3][pmiq][i_rows*4+j_col])[1]);
+                    }
+                    for (int i_rows=N1*N2; i_rows<2*N1*N2; i_rows++) {
+                      res_code=sqrt(1/(double)(4*CSI_RS_antenna_ports))*(phase_sign)*theta_n[nn]*v_lm[llc][mmc][i_rows-N1*N2];
+                      if (creal(res_code)>0)
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[0] = (short) ((creal(res_code)*32768)-0.5);//convert to Q15
+                      if (cimag(res_code)>0)
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)+0.5);//convert to Q15
+                      else
+                        ((short*) &mat[3][pmiq][i_rows*4+j_col])[1] = (short) ((cimag(res_code)*32768)-0.5);//convert to Q15
+                      LOG_D(PHY, "4 Layer Precoding Matrix[3][pmi %d][antPort %d][layerIdx %d]= %f+j %f -> Fixed Point %d+j %d \n",
+                            pmiq,i_rows,j_col, creal(res_code), cimag(res_code),((short*) &mat[3][pmiq][i_rows*4+j_col])[0],((short*) &mat[3][pmiq][i_rows*4+j_col])[1]);
                     }
                   }
                 }
@@ -497,6 +468,8 @@ int init_codebook_gNB(PHY_VARS_gNB *gNB) {
         }
       }
     }
+    else
+      AssertFatal(false, "Max number of antenna ports supported is 16\n");
   }
   return 0;
 }