#include <stdlib.h>
#include <gmp.h>
#include <mpfr.h>
#include <mpc.h>
#include "params.h"
#include "mpvec.h"
#include "polx.h"
#include "poly.h"
#include "polz.h"
#include "pdata.h"
#include "fft.h"
#include "ffldl.h"
#include "ffnp.h"
#include "keygen.h"
#if 1
#include <stdio.h>
#include <math.h>
#include "gaussian.h"
#endif

static void transpose(mpcvec A[4][4]) {
  int i,j,k;

  for(i=0;i<4;i++)
    for(j=i+1;j<4;j++)
      for(k=0;k<N/2;k++)
        mpc_swap(A[i][j][k],A[j][i][k]);
}

void keygen(poly *a1,
            poly *a2,
            polx *a3,
            polx a[2],
            polx b[2],
            polz *u,
            polx s[4])
{
  int i,j;
  mpz_t q;
  mpfr_t qinv;
  poly R[2][2];
  polz f;
  int32_t z[3] = {0,0,4};

  mpcvec Rhat[2][2];
  mpcvec ahat[2], uhat;
  mpcvec (*Bhat)[4], that[4], tmpvec;
  mpcvec (*L0)[4];
  mpc_t (*L1)[12*N/4];

  /* q, qinv */
  mpz_init2(q,96);
  mpfr_init2(qinv,PREC);
  mpz_set_ui(q,QH);
  mpz_mul_2exp(q,q,31);
  mpz_add_ui(q,q,QM);
  mpz_mul_2exp(q,q,31);
  mpz_add_ui(q,q,QL);
  mpfr_set_z(qinv,q,MPFR_RNDN);
  mpfr_ui_div(qinv,1,qinv,MPFR_RNDN);

  /* a1, a2, a, ahat, u, uhat */
  poly_uniform(a1,PDATA0);
  poly_uniform(a2,PDATA0);
  polz_uniform(u);
  polz_cmodq(u);
  polz_topolx(a3,u);
  polz_uniform(u);
  polz_topolx(&a[0],u);
  mpcvec_frompolz(ahat[0],u);
  polz_uniform(u);
  polz_topolx(&a[1],u);
  mpcvec_frompolz(ahat[1],u);
  polz_uniform(u);
  mpcvec_frompolz(uhat,u);
  poly_ntt(a1,PDATA0);
  poly_ntt(a2,PDATA0);
  polx_ntt(a3);
  polx_ntt(&a[0]);
  polx_ntt(&a[1]);
  mpfft(ahat[0]);
  mpfft(ahat[1]);
  mpfft(uhat);

  /* R, Rhat */
  poly_binary(&R[0][0]);
  mpcvec_frompoly(Rhat[0][0],&R[0][0]);
  mpfft(Rhat[0][0]);
  poly_binary(&R[0][1]);
  mpcvec_frompoly(Rhat[0][1],&R[0][1]);
  mpfft(Rhat[0][1]);
  poly_binary(&R[1][0]);
  mpcvec_frompoly(Rhat[1][0],&R[1][0]);
  mpfft(Rhat[1][0]);
  poly_binary(&R[1][1]);
  mpcvec_frompoly(Rhat[1][1],&R[1][1]);
  mpfft(Rhat[1][1]);

  /* b1, b2 */
  polx_frompoly(&s[0],&R[0][0]);
  polx_ntt(&s[0]);
  polx_pointwise(&b[0],&a[0],&s[0]);
  polx_frompoly(&s[0],&R[1][0]);
  polx_ntt(&s[0]);
  polx_pointwise(&s[0],&a[1],&s[0]);
  polx_add(&b[0],&b[0],&s[0]);
  //polx_scale(&b[0],&b[0],z);
  polx_invntt(&b[0]);
  polz_frompolx(&f,&b[0]);
  polz_topolx(&b[0],&f);
  polx_ntt(&b[0]);
  polx_frompoly(&s[0],&R[0][1]);
  polx_ntt(&s[0]);
  polx_pointwise(&b[1],&a[0],&s[0]);
  polx_frompoly(&s[0],&R[1][1]);
  polx_ntt(&s[0]);
  polx_pointwise(&s[0],&a[1],&s[0]);
  polx_add(&b[1],&b[1],&s[0]);
  //polx_scale(&b[1],&b[1],z);
  polx_invntt(&b[1]);
  polz_frompolx(&f,&b[1]);
  polz_topolx(&b[1],&f);
  polx_ntt(&b[1]);

#if 1
  mpcvec vhat[2];
  z[0] = 1; z[1] = 0; z[2] = 0;
  polx_scale(&s[0],&b[0],z);
  polx_invntt(&s[0]);
  polz_frompolx(&f,&s[0]);
  mpcvec_frompolz(vhat[0],&f);
  mpfft(vhat[0]);
  mpcvec_add_ui(vhat[0],vhat[0],1);
  polx_scale(&s[0],&b[1],z);
  polx_invntt(&s[0]);
  polz_frompolx(&f,&s[0]);
  mpcvec_frompolz(vhat[1],&f);
  mpfft(vhat[1]);
  mpcvec_add_ui(vhat[1],vhat[1],DELTA);
#endif

  /* B */
  mpcvec_init(tmpvec);
  Bhat = malloc(4*4*sizeof(mpcvec));
  for(i=0;i<4;i++)
    for(j=0;j<4;j++)
      mpcvec_init(Bhat[i][j]);

  mpcvec_mul_ui(Bhat[0][0],Rhat[0][0],DELTA);
  mpcvec_sub(Bhat[0][0],Rhat[0][1],Bhat[0][0]);
  mpcvec_mul_ui(Bhat[1][0],Rhat[1][0],DELTA);
  mpcvec_sub(Bhat[1][0],Rhat[1][1],Bhat[1][0]);
  mpcvec_mul_si(Bhat[0][1],Rhat[0][0],-Q0);
  mpcvec_mul_ui(tmpvec,Rhat[0][1],Q1);
  mpcvec_sub(Bhat[0][1],Bhat[0][1],tmpvec);
  mpcvec_mul_si(Bhat[1][1],Rhat[1][0],-Q0);
  mpcvec_mul_ui(tmpvec,Rhat[1][1],Q1);
  mpcvec_sub(Bhat[1][1],Bhat[1][1],tmpvec);

  mpcvec_set_ui(Bhat[2][0],DELTA);
  mpcvec_set_si(Bhat[3][0],-1);
  mpcvec_set_ui(Bhat[2][1],Q0);
  mpcvec_set_ui(Bhat[3][1],Q1);

  mpcvec_mul(Bhat[0][2],Rhat[0][0],ahat[0]);
  mpcvec_add_ui(Bhat[0][2],Bhat[0][2],1);
  mpcvec_mul(Bhat[1][2],Rhat[1][0],ahat[0]);
  mpcvec_mul(Bhat[0][3],Rhat[0][0],ahat[1]);
  mpcvec_mul(Bhat[1][3],Rhat[1][0],ahat[1]);
  mpcvec_add_ui(Bhat[1][3],Bhat[1][3],1);

  mpcvec_set_ui(Bhat[2][2],0);
  mpcvec_set_ui(Bhat[3][2],0);
  mpcvec_set_ui(Bhat[2][3],0);
  mpcvec_set_ui(Bhat[3][3],0);
  mpcvec_sub(Bhat[2][2],Bhat[2][2],ahat[0]);
  mpcvec_sub(Bhat[2][3],Bhat[2][3],ahat[1]);

#if 0
  for(i=0;i<4;++i) {
    mpcvec_mul(tmpvec,ahat[0],Bhat[0][i]);
    mpcvec_fma(tmpvec,ahat[1],Bhat[1][i],tmpvec);
    mpcvec_fma(tmpvec,vhat[0],Bhat[2][i],tmpvec);
    mpcvec_fma(tmpvec,vhat[1],Bhat[3][i],tmpvec);
    mpfft_inv(tmpvec);
    mpcvec_topolz(&f,tmpvec);
    polz_cmodq(&f); // FIXME
    //if(sqrt(polz_sprod(&f,&f)) > 1e-15)
    //  printf("ERROR: %d\n",i);
  }
#endif

  /* t */
  mpcvec_init(that[0]);
  mpcvec_init(that[1]);
  mpcvec_init(that[2]);
  mpcvec_init(that[3]);
  mpcvec_set_ui(that[3],0);
  mpcvec_div(that[2],uhat,ahat[0]);
  mpcvec_mul_fr(that[1],uhat,qinv);
  mpcvec_mul_ui(that[0],that[1],Q1);

#if 1
  mpcvec shat[4];
  mpcvec_init(shat[0]);
  mpcvec_init(shat[1]);
  mpcvec_init(shat[2]);
  mpcvec_init(shat[3]);
  for(i=0;i<4;++i) {
    mpcvec_mul(shat[i],Bhat[i][0],that[0]);
    mpcvec_fma(shat[i],Bhat[i][1],that[1],shat[i]);
    mpcvec_fma(shat[i],Bhat[i][2],that[2],shat[i]);
    mpcvec_fma(shat[i],Bhat[i][3],that[3],shat[i]);
  }
  mpcvec_mul(tmpvec,ahat[0],shat[0]);
  mpcvec_fma(tmpvec,ahat[1],shat[1],tmpvec);
  mpcvec_fma(tmpvec,vhat[0],shat[2],tmpvec);
  mpcvec_fma(tmpvec,vhat[1],shat[3],tmpvec);
  mpfft_inv(tmpvec);
  mpcvec_topolz(&f,tmpvec);
  polz_sub(&f,&f,u);
  polz_reduce(&f);
  polz_cmodq(&f); // FIXME
  //if(sqrt(polz_sprod(&f,&f)) > 1e-15)
  //  printf("ERROR\n");
#endif

  L0 = malloc(4*4*sizeof(mpcvec));
  L1 = malloc(4*12*N/4*sizeof(mpc_t));
  for(i=0;i<4;i++) {
    for(j=0;j<4;j++)
      mpcvec_init(L0[i][j]);
    for(j=0;j<12*N/4;j++)
      mpc_init2(L1[i][j],PREC);
  }

  transpose(Bhat); //FIXME
  mpffldl(L0,L1,Bhat,SIGMA);
  transpose(Bhat);
  //mpffnp(that,L0,L1,0);
  for(i=0;i<4;i++) {
    mpcvec_mul(tmpvec,Bhat[i][0],that[0]);
    for(j=1;j<4;j++)
      mpcvec_fma(tmpvec,Bhat[i][j],that[j],tmpvec);

    mpfft_inv(tmpvec);
    mpcvec_topolz(&f,tmpvec);
    polz_topolx(&s[i],&f);
    polx_ntt(&s[i]);
  }

#if 1
  mpz_t mpz;
  mpfr_t mu,sigma;
  mpz_init2(mpz,PREC);
  mpfr_init2(mu,PREC);
  mpfr_init2(sigma,PREC);
  mpfr_set_ui(mu,0,MPFR_RNDN);
  mpfr_set_d(sigma,SIGMA,MPFR_RNDN);
  for(i=0;i<4;i++) {
    for(j=0;j<N/2;j++) {
      gaussian_sampler(mpz,mu,sigma);
      mpfr_set_z(mpc_realref(shat[i][j]),mpz,MPFR_RNDN);
      gaussian_sampler(mpz,mu,sigma);
      mpfr_set_z(mpc_imagref(shat[i][j]),mpz,MPFR_RNDN);
    }
    mpcvec_topolz(&f,shat[i]);
    polz_topolx(&s[i],&f);
    polx_ntt(&s[i]);
    mpfft(shat[i]);
  }
  mpcvec_mul(tmpvec,ahat[0],shat[0]);
  mpcvec_fma(tmpvec,ahat[1],shat[1],tmpvec);
  mpcvec_fma(tmpvec,vhat[0],shat[2],tmpvec);
  mpcvec_fma(tmpvec,vhat[1],shat[3],tmpvec);
  mpfft_inv(tmpvec);
  mpcvec_topolz(u,tmpvec);
  mpz_clear(mpz);
  mpfr_clear(mu);
  mpfr_clear(sigma);
  mpcvec_clear(shat[0]);
  mpcvec_clear(shat[1]);
  mpcvec_clear(shat[2]);
  mpcvec_clear(shat[3]);
  mpcvec_clear(vhat[0]);
  mpcvec_clear(vhat[1]);
#endif

  /* clear */
  mpz_clear(q);
  mpfr_clear(qinv);
  mpcvec_clear(Rhat[0][0]);
  mpcvec_clear(Rhat[0][1]);
  mpcvec_clear(Rhat[1][0]);
  mpcvec_clear(Rhat[1][1]);
  mpcvec_clear(ahat[0]);
  mpcvec_clear(ahat[1]);
  mpcvec_clear(uhat);
  mpcvec_clear(that[0]);
  mpcvec_clear(that[1]);
  mpcvec_clear(that[2]);
  mpcvec_clear(that[3]);
  mpcvec_clear(tmpvec);
  for(i=0;i<4;i++)
    for(j=0;j<4;j++)
      mpcvec_clear(Bhat[i][j]);
  free(Bhat);
  for(i=0;i<4;i++) {
    for(j=0;j<4;j++)
      mpcvec_clear(L0[i][j]);
    for(j=0;j<12*N/4;j++)
      mpc_clear(L1[i][j]);
  }
  free(L0);
  free(L1);
}
