#include <gmp.h>
#include <mpfr.h>
#include <mpc.h>
#include "params.h"
#include "mpvec.h"
#include "gaussian.h"
#include "fft.h"
#include "ffldl.h"
#include "ffnp.h"

static void (*sampler)(mpz_t z, const mpfr_t mu, const mpfr_t sigma);

static void ffnpi(mpc_t *t, const mpfr_t sigma) {
  mpfr_t *t1, *t2;
  mpz_t z;

  t1 = &mpc_realref(*t);
  t2 = &mpc_imagref(*t);
  mpz_init2(z, PREC);

  sampler(z, *t2, sigma);
  mpfr_sub_z(*t2, *t2, z, MPFR_RNDN);
  sampler(z, *t1, sigma);
  mpfr_sub_z(*t1, *t1, z, MPFR_RNDN);

  mpz_clear(z);
}

static void ffnp2(mpc_t t[],
                  const mpc_t L[],
                  unsigned int level,
                  mpcvec tmp)
{
  unsigned int j, k, len, off;

  len = N/4 >> level;
  off = 1 << level;
  k = N/2 - 2*len;
  for(j = 0; j < N/2; j += 2*off) {
    mpc_sub(tmp[0], t[j], t[j + off], MPC_RNDNN);
    mpc_add(t[j], t[j], t[j + off], MPC_RNDNN);
    mpc_div_2ui(t[j], t[j], 1, MPC_RNDNN);
    mpc_div_2ui(t[j + off], tmp[0], 1, MPC_RNDNN);
    mpc_mul(t[j + off], t[j + off], zetas_inv[k++], MPC_RNDNN);
  }

  if(level < 10)
    ffnp2(t + off, L + len + (11 - level)*len/2, level + 1, tmp);
  else
    ffnpi(t + off, mpc_imagref(L[1]));

  k = 0;
  for(j = 0; j < N/2; j += 2*off) {
    mpc_mul(tmp[k], t[j + off], L[k], MPC_RNDNN);
    mpc_add(t[j], t[j], tmp[k], MPC_RNDNN);
    k++;
  }

  if(level < 10)
    ffnp2(t, L + len, level + 1, tmp + len);
  else
    ffnpi(t, mpc_realref(L[1]));

  k = 0;
  for(j = 0; j < N/2; j += 2*off)
    mpc_sub(t[j], t[j], tmp[k++], MPC_RNDNN);

  k = len;
  for(j = 0; j < N/2; j += 2*off) {
    mpc_mul(tmp[0], t[j + off], zetas[k++], MPC_RNDNN);
    mpc_sub(t[j + off], t[j], tmp[0], MPC_RNDNN);
    mpc_add(t[j], t[j], tmp[0], MPC_RNDNN);
  }
}

void mpffnp(mpcvec t[K],
            const mpcvec L0[K][K],
            const mpc_t L1[K][12*N/4],
            unsigned int sampler_flag)
{
  unsigned int i;
  int j;
  mpcvec tt, tmp;

  sampler = gaussian_sampler;
  mpcvec_init(tt);
  mpcvec_init(tmp);

  ffnp2(t[K-1], L1[K-1], 0, tmp);

  for(j = K-2; j >= 0; --j) {
    mpcvec_mul(tt, t[j+1], L0[j+1][j]);
    for(i = j + 2; i < K; ++i)
      mpcvec_fma(tt, t[i], L0[i][j], tt);

    mpcvec_add(t[j], t[j], tt);
    ffnp2(t[j], L1[j], 0, tmp);
    mpcvec_sub(t[j], t[j], tt);
  }

  mpcvec_clear(tt);
  mpcvec_clear(tmp);
}
