/*
 * FFT/IFFT transforms.
 *
 * Extracted 2009 from mplayer 2009-02-10 libavcodec/fft.c and libavcodec/mdct.c
 *
 * Copyright (c) 2008 Loren Merritt
 * Copyright (c) 2002 Fabrice Bellard
 * Partly based on libdjbfft by D. J. Bernstein
 *
 * Licensed under the GNU Lesser General Public License, see file COPYING.LIB.
 */

/**
 * \file imdct.c Inverse modified discrete cosine transform.
 */

#include <math.h>

#include "para.h"
#include "error.h"
#include "string.h"
#include "imdct.h"
#include "wma.h"

/** Canonical representation of a complex number. */
struct fft_complex {
	/** Real part. */
	float re;
	/** Imaginary part. */
	float im;
};

/** FFT Lookup table. */
struct fft_context {
	/** Number of bits of this instance of the FFT. */
	int nbits;
	/** The lookup table for cosine values. */
	uint16_t *revtab;
};

struct mdct_context {
	/** Size of MDCT (number of input data * 2). */
	int n;
	/** n = 2^n bits. */
	int nbits;
	/** Cosine table for pre/post rotation. */
	float *tcos;
	/** Sine table for pre/post rotation. */
	float *tsin;
	/** The context for the underlying fast Fourier transform. */
	struct fft_context fft;
};

/** \cond cosine_tabs */

/* cos(2 * pi * x / n) for 0 <= x <= n / 4, followed by its reverse */
#define COSINE_TAB(n) static float cos_ ## n[n / 2] __a_aligned(16)

COSINE_TAB(16);
COSINE_TAB(32);
COSINE_TAB(64);
COSINE_TAB(128);
COSINE_TAB(256);
COSINE_TAB(512);
COSINE_TAB(1024);
COSINE_TAB(2048);
COSINE_TAB(4096);
COSINE_TAB(8192);
COSINE_TAB(16384);
COSINE_TAB(32768);
COSINE_TAB(65536);

static float *cos_tabs[] = {
	cos_16, cos_32, cos_64, cos_128, cos_256, cos_512, cos_1024, cos_2048,
	cos_4096, cos_8192, cos_16384, cos_32768, cos_65536,
};
/** \endcond cosine_tabs */

__a_const static int split_radix_permutation(int i, int n)
{
	int m;
	if (n <= 2)
		return i & 1;
	m = n >> 1;
	if ((i & m) == 0)
		return split_radix_permutation(i, m) * 2;
	m >>= 1;
	if ((i & m) == 0)
		return split_radix_permutation(i, m) * 4 + 1;
	else
		return split_radix_permutation(i, m) * 4 - 1;
}

#define BF(x, y, a, b) {\
	x = a - b;\
	y = a + b;\
}

#define BUTTERFLIES(a0, a1, a2, a3) {\
	BF(t3, t5, t5, t1);\
	BF(a2.re, a0.re, a0.re, t5);\
	BF(a3.im, a1.im, a1.im, t3);\
	BF(t4, t6, t2, t6);\
	BF(a3.re, a1.re, a1.re, t4);\
	BF(a2.im, a0.im, a0.im, t6);\
}

/*
 * Force loading all the inputs before storing any. This is slightly slower for
 * small data, but avoids store->load aliasing for addresses separated by large
 * powers of 2.
 */
#define BUTTERFLIES_BIG(a0, a1, a2, a3) {\
	float r0 = a0.re, i0 = a0.im, r1 = a1.re, i1 = a1.im;\
	BF(t3, t5, t5, t1);\
	BF(a2.re, a0.re, r0, t5);\
	BF(a3.im, a1.im, i1, t3);\
	BF(t4, t6, t2, t6);\
	BF(a3.re, a1.re, r1, t4);\
	BF(a2.im, a0.im, i0, t6);\
}

#define TRANSFORM(a0, a1, a2, a3, wre,wim) {\
	t1 = a2.re * wre + a2.im * wim;\
	t2 = a2.im * wre - a2.re * wim;\
	t5 = a3.re * wre - a3.im * wim;\
	t6 = a3.im * wre + a3.re * wim;\
	BUTTERFLIES(a0, a1, a2, a3)\
}

#define TRANSFORM_ZERO(a0, a1, a2, a3) {\
	t1 = a2.re;\
	t2 = a2.im;\
	t5 = a3.re;\
	t6 = a3.im;\
	BUTTERFLIES(a0, a1, a2, a3)\
}

/* z[0...8n - 1], w[1...2n - 1] */
static void pass(struct fft_complex *z, const float *wre, unsigned int n)
{
	float t1, t2, t3, t4, t5, t6;
	int o1 = 2 * n;
	int o2 = 4 * n;
	int o3 = 6 * n;
	const float *wim = wre + o1;

	n--;
	TRANSFORM_ZERO(z[0], z[o1], z[o2], z[o3]);
	TRANSFORM(z[1], z[o1 + 1], z[o2 + 1], z[o3 + 1], wre[1], wim[-1]);
	do {
		z += 2;
		wre += 2;
		wim -= 2;
		TRANSFORM(z[0], z[o1], z[o2], z[o3], wre[0], wim[0]);
		TRANSFORM(z[1], z[o1 + 1], z[o2 + 1], z[o3 + 1], wre[1], wim[-1]);
	} while (--n);
}

#undef BUTTERFLIES
#define BUTTERFLIES BUTTERFLIES_BIG

#define DECL_FFT(n, n2, n4)\
static void fft##n(struct fft_complex *z)\
{\
	fft ## n2(z);\
	fft ## n4(z + n4 * 2);\
	fft ## n4(z + n4 * 3);\
	pass(z, cos_ ## n, n4 / 2);\
}

static void fft4(struct fft_complex *z)
{
	float t1, t2, t3, t4, t5, t6, t7, t8;

	BF(t3, t1, z[0].re, z[1].re);
	BF(t8, t6, z[3].re, z[2].re);
	BF(z[2].re, z[0].re, t1, t6);
	BF(t4, t2, z[0].im, z[1].im);
	BF(t7, t5, z[2].im, z[3].im);
	BF(z[3].im, z[1].im, t4, t8);
	BF(z[3].re, z[1].re, t3, t7);
	BF(z[2].im, z[0].im, t2, t5);
}

static void fft8(struct fft_complex *z)
{
	float t1, t2, t3, t4, t5, t6, t7, t8;

	fft4(z);

	BF(t1, z[5].re, z[4].re, -z[5].re);
	BF(t2, z[5].im, z[4].im, -z[5].im);
	BF(t3, z[7].re, z[6].re, -z[7].re);
	BF(t4, z[7].im, z[6].im, -z[7].im);
	BF(t8, t1, t3, t1);
	BF(t7, t2, t2, t4);
	BF(z[4].re, z[0].re, z[0].re, t1);
	BF(z[4].im, z[0].im, z[0].im, t2);
	BF(z[6].re, z[2].re, z[2].re, t7);
	BF(z[6].im, z[2].im, z[2].im, t8);

	TRANSFORM(z[1], z[3], z[5], z[7], M_SQRT1_2, M_SQRT1_2);
}

static void fft16(struct fft_complex *z)
{
	float t1, t2, t3, t4, t5, t6;

	fft8(z);
	fft4(z + 8);
	fft4(z + 12);

	TRANSFORM_ZERO(z[0], z[4], z[8], z[12]);
	TRANSFORM(z[2], z[6], z[10], z[14], M_SQRT1_2, M_SQRT1_2);
	TRANSFORM(z[1], z[5], z[9], z[13], cos_16[1], cos_16[3]);
	TRANSFORM(z[3], z[7], z[11], z[15], cos_16[3], cos_16[1]);
}

DECL_FFT(32, 16, 8)
DECL_FFT(64, 32, 16)
DECL_FFT(128, 64, 32)
DECL_FFT(256, 128, 64)
DECL_FFT(512, 256, 128)

DECL_FFT(1024, 512, 256)
DECL_FFT(2048, 1024, 512)
DECL_FFT(4096, 2048, 1024)
DECL_FFT(8192, 4096, 2048)
DECL_FFT(16384, 8192, 4096)
DECL_FFT(32768, 16384, 8192)
DECL_FFT(65536, 32768, 16384)

static void (*fft_dispatch[]) (struct fft_complex *) = {
	fft4, fft8, fft16, fft32, fft64, fft128, fft256, fft512, fft1024,
	fft2048, fft4096, fft8192, fft16384, fft32768, fft65536,
};

/* complex multiplication: p = a * b */
#define CMUL(pre, pim, are, aim, bre, bim) \
{\
	float _are = (are);\
	float _aim = (aim);\
	float _bre = (bre);\
	float _bim = (bim);\
	(pre) = _are * _bre - _aim * _bim;\
	(pim) = _are * _bim + _aim * _bre;\
}

/*
 * Compute the middle half of the inverse MDCT, excluding the parts that can be
 * derived by symmetry.
 */
static void imdct_half(struct mdct_context *s, float *output,
		const float *input)
{
	int n = 1 << s->nbits, n4 = n >> 2, n8 = n >> 3;
	const uint16_t *revtab = s->fft.revtab;
	const float *tcos = s->tcos;
	const float *tsin = s->tsin;
	const float *in1 = input, *in2 = input + (n >> 1) - 1;
	struct fft_complex *z = (struct fft_complex *)output;

	/* pre rotation */
	for (int k = 0; k < n4; k++) {
		int j = revtab[k];
		CMUL(z[j].re, z[j].im, *in2, *in1, tcos[k], tsin[k]);
		in1 += 2;
		in2 -= 2;
	}
	fft_dispatch[s->fft.nbits - 2](z);

	/* post rotation + reordering */
	for (int k = 0; k < n8; k++) {
		float r0, i0, r1, i1;
		CMUL(r0, i1, z[n8 - k - 1].im, z[n8 - k - 1].re,
			tsin[n8 - k - 1], tcos[n8 - k - 1]);
		CMUL(r1, i0, z[n8 + k].im, z[n8 + k].re, tsin[n8 + k],
			tcos[n8 + k]);
		z[n8 - k - 1].re = r0;
		z[n8 - k - 1].im = i0;
		z[n8 + k].re = r1;
		z[n8 + k].im = i1;
	}
}

/**
 * Compute the inverse MDCT.
 *
 * \param ctx The initialized context structure.
 * \param output N samples.
 * \param input N/2 samples.
 *
 * \sa \ref imdct_init().
 */
void imdct(struct mdct_context *ctx, float *output, const float *input)
{
	int k;
	int n = 1 << ctx->nbits;
	int n2 = n >> 1;
	int n4 = n >> 2;

	imdct_half(ctx, output + n4, input);

	for (k = 0; k < n4; k++) {
		output[k] = -output[n2 - k - 1];
		output[n - k - 1] = output[n2 + k];
	}
}

static int fft_init(struct fft_context *s, int nbits)
{
	int i, j, n;

	if (nbits < 2 || nbits > 16)
		return -E_FFT_BAD_PARAMS;
	s->nbits = nbits;
	n = 1 << nbits;

	s->revtab = arr_alloc(n, sizeof(uint16_t));
	for (j = 4; j <= nbits; j++) {
		int k = 1 << j;
		double freq = 2 * M_PI / k;
		float *tab = cos_tabs[j - 4];
		for (i = 0; i <= k / 4; i++)
			tab[i] = cos(i * freq);
		for (i = 1; i < k / 4; i++)
			tab[k / 2 - i] = tab[i];
	}
	for (i = 0; i < n; i++)
		s->revtab[-split_radix_permutation(i, n) & (n - 1)] = i;
	return 0;
}

/**
 * Initialize the inverse modified cosine transform.
 *
 * \param nbits The number of bits to use (4 <= \a nbits <= 18).
 *
 * \param result Opaque structure that must be passed to \ref imdct().
 *
 * \return Standard.
 */
int imdct_init(int nbits, struct mdct_context **result)
{
	int ret, n, n4, i;
	double alpha;
	struct mdct_context *s;

	s = zalloc(sizeof(*s));
	n = 1 << nbits;
	s->nbits = nbits;
	s->n = n;
	n4 = n >> 2;
	s->tcos = arr_alloc(n4, sizeof(float));
	s->tsin = arr_alloc(n4, sizeof(float));

	for (i = 0; i < n4; i++) {
		alpha = 2 * M_PI * (i + 1.0 / 8.0) / n;
		s->tcos[i] = -cos(alpha);
		s->tsin[i] = -sin(alpha);
	}
	ret = fft_init(&s->fft, s->nbits - 2);
	if (ret < 0)
		goto fail;
	*result = s;
	return 0;
fail:
	freep(&s->tcos);
	freep(&s->tsin);
	free(s);
	return ret;
}

/**
 * Deallocate imdct resources.
 *
 * \param ctx The pointer obtained by imdct_init().
 */
void imdct_end(struct mdct_context *ctx)
{
	free(ctx->tcos);
	free(ctx->tsin);
	free(ctx->fft.revtab);
	free(ctx);
}