#include <math.h>
#include <gmp.h>
#include <mpfr.h>
#include <mpc.h>
#include "params.h"
#include "fft.h"

mpc_t zetas[N/2];
mpc_t zetas_inv[N/2];
static mpc_t u;

static unsigned int revbin(unsigned int n, unsigned int k) {
  unsigned int r = 0;

  while(n >>= 1) {
    r <<= 1;
    r += (k & 1U);
    k >>= 1;
  }

  return r;
}

void mpfft_init(void) {
  unsigned int level, i, j, k;

  // FIXME: Nicer!

  k = 1;
  for(level = 1; level < 12; ++level) {
    for(i = 0; i < (1U << (level - 1)); ++i) {
      j = (1 << level) + i;
      j = revbin(N, j);
      //printf("%d\n", j);
      mpc_init2(zetas[k], PREC);
      mpc_rootofunity(zetas[k], 2*N, j, MPC_RNDNN);
      k += 1;
    }
  }

  k = 0;
  for(level = 11; level >= 1; --level) {
    for(i = 0; i < (1U << (level - 1)); ++i) {
      j = (1 << level) + i;
      j = revbin(N, j);
      //printf("%d\n", j);
      mpc_init2(zetas_inv[k], PREC);
      mpc_rootofunity(zetas_inv[k], 2*N, j, MPC_RNDNN);
      mpc_conj(zetas_inv[k], zetas_inv[k], MPC_RNDNN);
      k += 1;
    }
  }

  mpc_init2(u, PREC);
}

void mpfft(mpc_t r[N/2]) {
  unsigned int len, start, j, k;

  k = 1;
  for(len = N/4; len >= 1; len >>= 1) {
    for(start = 0; start < N/2; start = j + len) {
      for(j = start; j < start + len; ++j) {
        mpc_mul(u, zetas[k], r[j + len], MPC_RNDNN);
        mpc_sub(r[j + len], r[j], u, MPC_RNDNN);
        mpc_add(r[j], r[j], u, MPC_RNDNN);
      }
      k += 1;
    }
  }
}

void mpfft_inv(mpc_t r[N/2]) {
  unsigned int start, len, j, k;

  k = 0;
  for(len = 1; len <= N/4; len <<= 1) {
    for(start = 0; start < N/2; start = j + len) {
      for(j = start; j < start + len; ++j) {
        mpc_sub(u, r[j], r[j + len], MPC_RNDNN);
        mpc_add(r[j], r[j], r[j + len], MPC_RNDNN);
        mpc_mul(r[j + len], zetas_inv[k], u, MPC_RNDNN);
      }
      k += 1;
    }
  }

  for(j = 0; j < N/2; ++j)
    mpc_div_2ui(r[j], r[j], log2(N/2), MPC_RNDNN);
}
