#include <assert.h>
#include <errno.h>
#include "util.h"
#include "convolve.h"

/* Return the lowest power of 2 that is >= x (and >= 1) */
static int next_pow2(const int x) {
  int y = 1;
  while (y < x) {
    assert(can_mul_int(y, 2));
    y *= 2;
  }
  return y;
}

/* Multiplication of complex numbers */
/* Allows output to overwrite input */
static void multiply_complex(fftw_complex z, const fftw_complex x, const fftw_complex y) {
  const double a = x[0] * y[0] - x[1] * y[1];
  const double b = x[0] * y[1] + x[1] * y[0];
  z[0] = a;
  z[1] = b;
}

/* Dot product of complex vectors */
/* Can't use const for pointer-to-pointer */
static void dot_complex(fftw_complex *z, fftw_complex *x, fftw_complex *y, int size) {
  int i;
  for (i = 0; i < size; ++i)
    multiply_complex(z[i], x[i], y[i]);
}

/* malloc or die */
static void *fftw_xmalloc(size_t size) {
  void *p = fftw_malloc(size);
  if (p == NULL)
    die("%s: error allocating %lu bytes: %s\n",
        prog_name, (unsigned long)size, strerror(errno));
  return p;
}

void fft_init(fft_convolver *f, const int max_size) {
  int zero_pad, size, mem;
  assert(max_size > 0);
  zero_pad = max_size - 1;  /* padding to avoid wraparound convolution */
  assert(can_add_int(max_size, zero_pad));
  size = next_pow2(max_size + zero_pad);  /* power-of-2 size is faster??? */
  assert(can_mul_int(2, size / 2 + 1));
  mem = 2 * (size / 2 + 1);  /* padding for in-place transforms */

  f->size = size;

  f->x = fftw_xmalloc(mem * sizeof(double));
  f->y = fftw_xmalloc(mem * sizeof(double));
  f->z = fftw_xmalloc(mem * sizeof(double));

  f->x_plan = fftw_plan_dft_r2c_1d(size, f->x, (fftw_complex *)f->x,
				   FFTW_MEASURE);
  f->y_plan = fftw_plan_dft_r2c_1d(size, f->y, (fftw_complex *)f->y,
				   FFTW_MEASURE);
  f->z_plan = fftw_plan_dft_c2r_1d(size, (fftw_complex *)f->z, f->z,
				   FFTW_MEASURE);
}

void fft_convolve(double *z, const double *x, const double *y, const int size, fft_convolver *f) {
  const int half_size = f->size / 2 + 1;
  const int pad_size = f->size - size;
  assert(size <= f->size / 2);

  COPY(f->x, x, size);
  COPY(f->y, y, size);
  set_dbl(f->x + size, pad_size, 0);  /* pad with zeros */
  set_dbl(f->y + size, pad_size, 0);  /* pad with zeros */

  fftw_execute(f->x_plan);
  fftw_execute(f->y_plan);
  dot_complex((fftw_complex *)f->z, (fftw_complex *)f->x, (fftw_complex *)f->y,
	      half_size);
  fftw_execute(f->z_plan);
  mul_dbl(f->z, f->size, 1.0 / f->size);

  COPY(z, f->z, size);
}
