#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "fft.h"

static void
do_fft (complex_t *buffer, int n, complex_t principal, int is_inverse)
{
    complex_t *temp, *even, *odd;
    complex_t alpha;
    complex_t p2;
    int i;
    double inv;
    
    if (n == 1)
	return;

    temp = malloc (n * sizeof (complex_t));
    even = temp;
    odd = temp + n/2;
    
    for (i = 0; i < n/2; ++i)
    {
	even[i] = buffer[2 * i];
	odd[i] = buffer[2 * i + 1];
    }

    p2 = complex_mul (principal, principal);
    
    do_fft (even, n/2, p2, FALSE);
    do_fft (odd, n/2, p2, FALSE);

    alpha.re = 1.0;
    alpha.im = 0;

    if (is_inverse)
	inv = 1.0 / n;
    
    for (i = 0; i < n/2; ++i)
    {
	buffer[i]	= complex_add (even[i], complex_mul (alpha, odd[i]));
	buffer[i + n/2] = complex_sub (even[i], complex_mul (alpha, odd[i]));

	if (is_inverse)
	{
	    buffer[i].re *= inv;
	    buffer[i].im *= inv;
	    buffer[i + n/2].re *= inv;
	    buffer[i + n/2].im *= inv;
	}
	
	alpha = complex_mul (alpha, principal);
    }

    free (temp);
}

void
fft (complex_t *buffer, int n)
{
    complex_t p;

    p.re = cos (2 * M_PI / n);
    p.im = sin (2 * M_PI / n);
    
    do_fft (buffer, n, p, FALSE);
}

void
ifft (complex_t *buffer, int n)
{
    complex_t p;

    p.re = cos (- (2 * M_PI / n));
    p.im = sin (- (2 * M_PI / n));

    do_fft (buffer, n, p, TRUE);
}

void
shift (complex_t *buffer, int n)
{
    int i;

    for (i = 0; i < n/2; ++i)
    {
	complex_t tmp = buffer[i];
	buffer[i] = buffer[i + n/2];
	buffer[i + n/2] = tmp;
    }
}

typedef void (* fft_1d_t) (complex_t *buffer, int n);

static void
do_fft_2d (complex_t *buffer, int n, fft_1d_t func1d)
{
    int i;

    /* Transform all the rows */
    for (i = 0; i < n; ++i)
    {
	complex_t *b = buffer + i * n;

	func1d (b, n);
    }

    /* Then the columns */
    for (i = 0; i < n; ++i)
    {
	int j;
	complex_t *tmp = malloc (n * sizeof (complex_t));

	for (j = 0; j < n; ++j)
	    tmp[j] = buffer[j * n + i];

	func1d (tmp, n);

	for (j = 0; j < n; ++j)
	    buffer[j * n + i] = tmp[j];

	free (tmp);
    }
}

void
fft_2d (complex_t *buffer, int n)
{
    do_fft_2d (buffer, n, fft);
}

void
ifft_2d (complex_t *buffer, int n)
{
    do_fft_2d (buffer, n, ifft);
}

void
shift_2d (complex_t *buffer, int n)
{
    int i, j;
    
    /* Swap quadrant 0 with quadrant 3 */
    for (i = 0; i < n/2; ++i)
    {
	for (j = 0; j < n/2; ++j)
	{
	    complex_t tmp;
	    
	    tmp = buffer[i * n + j];
	    buffer[i * n + j] = buffer[((n/2) + i) * n + (j + n/2)];
	    buffer[((n/2) + i) * n + (j + n/2)] = tmp;
	}
    }

    /* Swap quadrant 1 with quadrant 2 */
    for (i = 0; i < n/2; ++i)
    {
	for (j = 0; j < n/2; ++j)
	{
	    complex_t tmp;

	    tmp = buffer[i * n + j + n/2];
	    buffer[i * n + j + n/2] = buffer[(i + n/2) * n + j];
	    buffer[(i + n/2) * n + j] = tmp;
	}
    }
}

void
fft_shift_2d (complex_t *buffer, int n)
{
    fft_2d (buffer, n);
    shift_2d (buffer, n);
}

void
ifft_shift_2d (complex_t *buffer, int n)
{
    shift_2d (buffer, n);
    ifft_2d (buffer, n);
}