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

static double leaf;

static void ffldl2(mpc_t L[], mpfr_t d[], unsigned int level, mpfr_t tmp) {
  unsigned int j, k, len;

  len = (1 << level); // offset

  k = N/4 >> level;
  for(j = 0; j < N/2; j += 2*len) {
    mpfr_sub(tmp, d[j], d[j + len], MPFR_RNDN);
    mpfr_div_2ui(tmp, tmp, 1, MPFR_RNDN);
    mpfr_add(d[j], d[j], d[j + len], MPFR_RNDN);
    mpfr_div_2ui(d[j], d[j], 1, MPFR_RNDN);
    mpfr_div(tmp, tmp, d[j], MPFR_RNDN);
    mpc_mul_fr(*(L++), zetas[k++], tmp, MPC_RNDNN);
    mpfr_sqr(tmp, tmp, MPFR_RNDN);
    mpfr_ui_sub(tmp, 1, tmp, MPFR_RNDN);
    mpfr_mul(d[j + len], d[j], tmp, MPFR_RNDN);
  }

  if(level < 10) {
    ffldl2(L, d, level + 1, tmp);
    ffldl2(L + (11 - level)*(N/8 >> level), d + len, level + 1, tmp);
  }
  else {
    mpfr_div_2ui(tmp, d[0], 1, MPFR_RNDN); //FIXME
    mpfr_sqrt(tmp, tmp, MPFR_RNDN);
    mpfr_d_div(tmp, leaf, tmp, MPFR_RNDN);
    mpfr_set(mpc_realref(L[0]), tmp, MPFR_RNDN);

    mpfr_div_2ui(tmp, d[len], 1, MPFR_RNDN);
    mpfr_sqrt(tmp, tmp, MPFR_RNDN);
    mpfr_d_div(tmp, leaf, tmp, MPFR_RNDN);
    mpfr_set(mpc_imagref(L[0]), tmp, MPFR_RNDN);
  }
}

void mpffldl(mpcvec L0[K][K],
             mpc_t L1[K][12*N/4],
             const mpcvec B[K][K],
             double sigma)
{
  unsigned int i, j, k;
  mpfrvec vfr, d[K];
  mpcvec vc;

  leaf = sigma;
  mpfrvec_init(vfr);
  mpcvec_init(vc);
  for(i = 0; i < K; ++i)
    mpfrvec_init(d[i]);

  /* Cholesky decomposition */
  for(j = 0; j < K; ++j) {
    /* g_{j,j} */
    mpcvec_abssq(d[j], B[j][0]);
    for(k = 1; k < K; ++k) {
      mpcvec_abssq(vfr, B[j][k]);
      mpfrvec_add(d[j], d[j], vfr);
    }

    /* d_j */
    for(k = 0; k < j; ++k) {
      mpcvec_abssq(vfr, L0[j][k]);
      mpfrvec_mul(vfr, vfr, d[k]);
      mpfrvec_sub(d[j], d[j], vfr);
    }

    /* l_{i,j} */
    for(i = j+1; i < K; ++i) {
      /* g_{i,j} */
      mpcvec_conj(vc, B[j][0]);
      mpcvec_mul(L0[i][j], vc, B[i][0]);
      for(k = 1; k < K; ++k) {
        mpcvec_conj(vc, B[j][k]);
        mpcvec_mul(vc, vc, B[i][k]);
        mpcvec_add(L0[i][j], L0[i][j], vc);
      }

      for(k = 0; k < j; ++k) {
        mpcvec_conj(vc, L0[j][k]);
        mpcvec_mul(vc, vc, L0[i][k]);
        mpcvec_mul_frvec(vc, vc, d[k]);
        mpcvec_sub(L0[i][j], L0[i][j], vc);
      }

      mpcvec_div_frvec(L0[i][j], L0[i][j], d[j]);
    }
  }

  for(i = 0; i < K; ++i)
    ffldl2(L1[i], d[i], 0, vfr[0]);

  mpfrvec_clear(vfr);
  mpcvec_clear(vc);
  for(i = 0; i < K; ++i)
    mpfrvec_clear(d[i]);
}
