/*
Forward and inverse Modified Discrete Cosine Transform
alng with its usage on a signal.

Compilation: cc -Wall -Wextra -std=c99 -o dft dft.c -lm

Written by Paweł Sabat.
Web site: sabat.e90.biz

Licensed under Creative Commons Zero license.
See https://creativecommons.org/publicdomain/zero/1.0/ for full license text.
*/
#include <math.h>
#include <stdio.h>
#include <stdlib.h>

// C99 standard does not define constant PI
#define PI 3.1415926218032837

/*
Compute MDCT of x and store it in y.

x - input signal vector
x_length - length of vector x
y - output signal vector
	length(y) = x_length / 2;
*/
void mdct(float x[], int x_length, /*out*/ float y[])
{
	int N = x_length/2;
	float Nf = (float)N;

	// zero y
	for(int k=0; k<N; ++k)
	{
		y[k] = 0.0;
	}

	// compute mdct
	for(int k=0; k<N; ++k)
	{
		for(int n=0; n<2*N; ++n)
		{
			y[k] = y[k] + x[n] * cosf(2.0*PI * (k+1.0/2.0)/(2.0*Nf) * (n+1.0/2.0+Nf/2.0));
		}
	}
}

/*
Compute IMDCT of y and store it in z.

y - array of forward mdct values
y_length - number of elements in array y
z - output array
	length(z) = y_length * 2
*/
void imdct(float y[], int y_length, /*out*/ float z[])
{
	int N = y_length;
	float Nf = (float)N;

	// zero z
	for(int k=0; k<2*N; ++k)
	{
		z[k] = 0.0;
	}

	// compute imdct
	for(int n=0; n<2*N; ++n)
	{
		for(int k=0; k<N; ++k)
		{
			z[n] = z[n] + y[k] * cosf(2.0*PI * (k+1.0/2.0)/(2.0*Nf) * (n+1.0/2.0+Nf/2.0));
		}

		// divide by N
		z[n] = z[n] / Nf;
	}
}

/*
Pretty print for arrays with floats.
*/
void pprint(float v[], int length)
{
	printf("[%.4f", v[0]);
	for(int i=1; i<length; ++i)
	{
		printf(", %.4f", v[i]);
	}
	printf("]\n");
}

unsigned int signal_mdct_output_length(int signal_length, int block_size)
{
	/*
	overlap_size = block_size / 2;
	blocks_from_signal = signal_length / overlap_size - 1 = signal_length / (block_size/2) - 1 =
		signal_length * 2 / block_size - 1;
	samples_in_blocks = blocks_from_signal * block_size =
		((signal_length * 2) / block_size - 1) * block_size = signal_length * 2 - block_size

	Forward MDCT takes 2N samples and returns N samples so:
	total_output_samples = samples_in_blocks/2 = signal_length - block_size/2;
	*/
	return signal_length - block_size/2;
}

/*
Divide signal into overlapping blocks and compute forward MDCT for each block.

signal - values that should be converted by mdct
signal_length - how many values signal array stores
block_size - signal is divided into overlapping blocks of length block_size
	each block overlaps exactly half of the previous block
y_bx - output array, stores values computed by signal_mdct
	length(y_bx) = signal_length - block_size/2
	length of y_bx is given by function signal_mdct_output_length()
*/
int signal_mdct(float signal[], int signal_length, int block_size, /*out*/ float y_bx[])
{
/*
How this works:
	Signal is divided into blocks b0, b1, b2, .. where each block has length block_size
	and each block overlaps exactly half of previous block.
	For each block mdct() is computed
	y_b0 = mdct(b0), y_b1 = mdct(b1), ..
	and values y_bn are stored in a single array y_bx
	y_bx = [y_b0, y_b1, y_b2, ..]
	Remember that values y_bn are arrays themselves (they usually contain more than one value).

	Array bx is not explicity stored as it is not necessary.
	Blocks bn are taken directly from the signal.
*/

	// block size must be even
	if(block_size % 2 != 0)
	{
		fprintf(stderr, "signal_mdct() error: block_size must be even\n");
		return -1;
	}

	int overlap_size = block_size / 2;

	// signal_length must be multiplicity of overlap_size
	if(signal_length % overlap_size != 0)
	{
		fprintf(stderr, "signal_mdct() error: signal_length must be multiplicity of overlap_size\n");
		return -1;
	}

	// for each block
	for(int block_start=0; block_start+block_size-1<signal_length; block_start+=overlap_size)
	{
		// compute mdct and store it in y_bx
		mdct(signal+block_start, block_size, y_bx);
		// advance position of output buffer
		// mdct takes 2N samples but returns only N samples thus advance position of output buffer
		// only by half of block_size
		y_bx += block_size/2;
	}

	return 0;
}

/*
Restore signal from values of signal_mdct() (forward MDCT).

y_bx - array that contains output of signal_mdct() = values of vectors y_b0, y_b1, ...
	one after another
	each vector y_bn has length block_size/2
y_bx_length - how many values does y_bx store
block_size - size of the blocks which original signal was divided into
z_tmp - array for temporary values of computations
	length(z_tmp) = 2 * block_size
restored_signal - array that will contain values of original signal restored from y_bx
	length(restored_signal) = original_signal_length - block_size = y_bx_length - block_size/2
*/
void signal_imdct(float y_bx[], int y_bx_length, int block_size,
	float z_tmp[], /*out*/ float restored_signal[])
{
/*
How this works:
	Array y_bx stores values y_b0, y_b1, y_b2, ...
	Take y_b0, compute imdct(y_b0) = z_b0
	Take y_b1, compute imdct(y_b1) = z_b1
	Add overlapping parts of z_b0 and z_b1 to restore some values of the signal.
	Take y_b2, compute imdct(y_b2) = z_b2.
	Add overlapping parts of z_b1 and z_b2 to restore some values of the signal.
	etc.

	At any given time only two values are needed: current and previous block imdcts.
	previous block imdct = z_bn
	current block imdct = z_bm

	After z_bn is used, it is replaced by z_bm and z_bm is given new value in next
	iteration of the loop.
	To avoid copying values from z_bm to z_bn only pointers are swapped.
*/

	float *y_bx_end = y_bx + y_bx_length;
	float *z_bn = z_tmp;
	float *z_bm = z_tmp + block_size;
	int y_bx_item_size = block_size / 2;

	// take y_b0 and compute imdct of it
	float *y_bn = y_bx;
	imdct(y_bn, y_bx_item_size, z_bn);
	y_bx += y_bx_item_size;

	while(y_bx < y_bx_end)
	{
		// take y_bm and compute imdct of it
		float *y_bm = y_bx;
		imdct(y_bm, y_bx_item_size, z_bm);
		y_bx += y_bx_item_size;

		// restore values by adding overlapping parts of z_bn and z_bm
		// there will be overlap_size items restored = block_size / 2
		// restoration is by adding second half of z_bn to first half of z_bm
		const int half_block_size = block_size / 2;
		const int restored_items = half_block_size;
		for(int i=0; i<half_block_size; ++i)
		{
			restored_signal[i] = z_bn[half_block_size + i] + z_bm[i];
		}
		// advance output array position
		restored_signal += restored_items;

		// make current block imdct a previous one
		// put value of pointer y_bm in y_bn
		y_bn = y_bm;
		// swap pointers of z_bn and z_bm
		float *tmp = z_bn;
		z_bn = z_bm;
		z_bm = tmp;
	}
}

void compute_single_block_mdct()
{
	float x[] = {10.0, 2.0, 7.0, 9.0};
	const int x_length = 4;
	const int N = x_length/2;
	float y[N];
	float z[2*N];

	printf("Computing MDCT of a single block\n");

	mdct(x, 2*N, y);
	imdct(y, N, z);

	printf("x="); pprint(x, 2*N);
	printf("y="); pprint(y, N);
	printf("z="); pprint(z, 2*N);
}

/*
Pretty print for array y_bx displaying each y_b0, y_b1, ... separately
*/
void pprint_y_bx(float y_bx[], int y_bx_length, int block_size)
{
	int y_bx_item_length = block_size / 2;
	int items_in_y_bx = y_bx_length / y_bx_item_length;

	for(int i=0; i<items_in_y_bx; ++i)
	{
		printf("y_b%d=", i);
		pprint(y_bx+y_bx_item_length*i, y_bx_item_length);
	}
}

void compute_signal_mdct()
{
	float signal[] = {0,0,1,2,3,4,5,6,7,8,0,0};
	int signal_length = sizeof(signal) / sizeof(float);
	int block_size = 4;
	int y_bx_length = signal_mdct_output_length(signal_length, block_size);
	float *y_bx = (float*)malloc(sizeof(float) * y_bx_length);

	printf("Computing MDCT of a signal\n");
	int r = signal_mdct(signal, signal_length, block_size, y_bx);
	if(r == -1)
	{
		return;
	}

	int restored_signal_length = signal_length - block_size;
	float *restored_signal = (float*) malloc(sizeof(float) * restored_signal_length);
	float *z_tmp = (float*) malloc(sizeof(float) * (2*block_size));
	signal_imdct(y_bx, y_bx_length, block_size, z_tmp, restored_signal);

	printf("signal="); pprint(signal, signal_length);
	printf("y_bx="); pprint(y_bx, y_bx_length);
	pprint_y_bx(y_bx, y_bx_length, block_size);
	printf("restored_signal="); pprint(restored_signal, restored_signal_length);

	free(restored_signal);
	free(z_tmp);
	free(y_bx);
}

int main(void)
{
	compute_single_block_mdct();
	printf("\n");
	compute_signal_mdct();

	return 0;
}