//
//  pke.c
//  c
//

#include <stdint.h>
#include "poly.h"
#include "pke.h"
#include "fips202.h"
#include "parameter.h"
#include <string.h>
#include "rand.h"

const uint16_t QD[4]={416,1248,2080,2912};


int pke_keygen(unsigned char *pk,unsigned long long *pklen, unsigned char *sk,unsigned long long *sklen)
  
  {
       
      uint16_t s[N],e_0[N],bhat[N],s_ahat[N],a[N];
      unsigned char seed[2*PKE_SEED_BYTES],seed1[PKE_SEED_BYTES];
      int i;
	  int iq, ir;

	  //generate a
      rand_byts(PKE_SEED_BYTES,seed1);
      shake256(seed,2*PKE_SEED_BYTES,seed1, PKE_SEED_BYTES);
      poly_gen_a(a,seed);
	 
    
	  /*generate 2s,2e_0, 中心二项分布采样k=1*/
	  poly_sample(s,seed+PKE_SEED_BYTES,0);
	  poly_sample(e_0,seed+PKE_SEED_BYTES,1);
      //2s+1
      s[0]=s[0]+1;


	  kntt(s);	
	  poly_tobyte(sk,s); 
	  kntt(e_0);

	  //b=as+e
	  k_pointwise(s_ahat,a,s);
	  poly_add(bhat,s_ahat,e_0);
	   
	  
      poly_tobyte(pk,bhat);
      for(i=0;i<PKE_SEED_BYTES;i++)
          pk[PKE_POLY_BYTES+i]=seed[i];
    
	 /* *pklen = PKE_PUBLICKEYBYTES;
	  *sklen = PKE_SECRETKEYBYTES;*/
     return 0;//success
    
    
}


//pke with seed
int pke_enc(unsigned char *pk,unsigned long long pklen, unsigned char  *m, unsigned long long mlen,
            unsigned char *ct, unsigned long long *clen, unsigned char *coins)
{
    uint16_t r[N],e_1[N],e_r[N],c[N],c_1[N],c_bar[N],ar[N],a[N],br[N],mm[N],b[N],brr[N],mmm[N];
    int i,temp,j;
	
	//generate a
    poly_gen_a(a,pk+PKE_POLY_BYTES);

	//generate 2r,2e_1,2e_r
    poly_sample(r,coins,0);
	poly_sample(e_1,coins,1);
	poly_sample(e_r,coins,2);

    for(i=0;i<N;i++)
        r[i]=r[i]>>1;
    

	//c=br+e_r
    byte_topoly(pk,b);
	kntt(r);
    k_pointwise(br,b,r);
	kinv_ntt(br);

	//backward permutation
	for(i=0;i<8;i++)
	   for(j=0;j<128;j++)
	   { 
		  brr[(j<<3)+i]=br[j+(i<<7)];
        }
	poly_add(c,brr,e_r);

      
    for(i=0;i<N;i++)
    {
		mm[i] = (m[(i >> 3)] >> (i & 0x7)) & 1;
        mm[i]=(mm[i]^(c[i]&0x1))& 0x1+e_1[i]; //mm=(m+alpha) mod 2 +e_1
        c_bar[i]=(c[i]-1)/832;
    }

	
	//forward  permutation
	for(i=0;i<8;i++)
	   for(j=0;j<128;j++)
	     { mmm[(i<<7)+j]=mm[(j<<3)+i];
         }


	//c_1=ar+(m+alpha)+e_1
	kntt(mmm);
    k_pointwise(ar,a,r);
	poly_add(c_1,ar,mmm);

    poly_tobyte(ct,c_1);
     for(i=0;i<N;i=i+4)
     {
        temp=i>>2;
        ct[PKE_POLY_BYTES+temp]=c_bar[i];
        ct[PKE_POLY_BYTES+temp]+=((c_bar[i+1]&3)<<2);
        ct[PKE_POLY_BYTES+temp]+=((c_bar[i+2]&3)<<4);
        ct[PKE_POLY_BYTES+temp]+= ((c_bar[i+3]&3)<<6);

      }   

	/* *clen = PKE_CIP_BYTES;*/
    return 0;   
}





int pke_dec(unsigned char *sk,unsigned long long sklen,
            unsigned char *ct, unsigned long long clen,
            unsigned char *m, unsigned long long *mlen)

{
    uint16_t mprime[N],c_1[N],s[N],cs[N],cprime[N],ccs[N];
    int16_t mmprime[N];
    int i,j;
    unsigned char temp;

    for(i=0;i<N4;i=i+1)
    {
        temp=ct[i+PKE_POLY_BYTES];
        for(j=0;j<4;j++)

       {  
		   cprime[(i<<2)+j]=(temp & 3);
  	       temp=temp>>2;
        }

     }



	byte_topoly(sk,s);	
    byte_topoly(ct,c_1);

	k_pointwise(cs,c_1,s);
	kinv_ntt(cs);



	 for(i=0;i<8;i++)
	   for(j=0;j<128;j++)
	     { 
			 ccs[(j<<3)+i]=cs[j+(i<<7)];
	      }
  
  
    for(i=0;i<N;i++)
    {         
        cprime[i]=QD[cprime[i]];//QD[c]={(2c+1)*416}
        mmprime[i]=ccs[i]-cprime[i];
		mprime[i]=( mmprime[i]-(mmprime[i]/1665)*Q)&0x1; //[-(Q-1)/2,(Q-1)/2]      
    }

    memset(m,0,sizeof(m));	

	for(i=0;i<N8;i++)
    { 
		 j = (i<<3);
  		 m[i]=mprime[j];
         m[i]= m[i]+(mprime[j+1]<<1);
         m[i]= m[i]+(mprime[j+2]<<2);
         m[i]= m[i]+(mprime[j+3]<<3);
         m[i]= m[i]+(mprime[j+4]<<4);
         m[i]= m[i]+(mprime[j+5]<<5);
         m[i]= m[i]+(mprime[j+6]<<6);
         m[i]= m[i]+(mprime[j+7]<<7);

    }
   
 /*  *mlen = PKE_MESS_BYTES;*/
    return 0;
    
    
    
    
}
