/*
	File                 : nsl_dft.c
	Project              : LabPlot
	Description          : NSL discrete Fourier transform functions
	--------------------------------------------------------------------
	SPDX-FileCopyrightText: 2016 Stefan Gerlach <stefan.gerlach@uni.kn>
	SPDX-License-Identifier: GPL-2.0-or-later
*/

#include "nsl_dft.h"
#include "nsl_common.h"
#include <gsl/gsl_fft_halfcomplex.h>
#include <gsl/gsl_fft_real.h>
#ifdef HAVE_FFTW3
#include <fftw3.h>
#endif

const char* nsl_dft_result_type_name[] = {i18n("Magnitude"),
										  i18n("Amplitude"),
										  i18n("Real Part"),
										  i18n("Imaginary Part"),
										  i18n("Power"),
										  i18n("Phase"),
										  i18n("Amplitude in dB"),
										  i18n("Normalized Amplitude in dB"),
										  i18n("Magnitude Squared"),
										  i18n("Amplitude Squared"),
										  i18n("Raw")};
const char* nsl_dft_xscale_name[] = {i18n("Frequency"), i18n("Index"), i18n("Period")};

int nsl_dft_transform_window(double data[], size_t stride, size_t n, int two_sided, nsl_dft_result_type type, nsl_sf_window_type window_type) {
	/* apply window function */
	if (window_type != nsl_sf_window_uniform)
		nsl_sf_apply_window(data, n, window_type);

	/* transform */
	int status = nsl_dft_transform(data, stride, n, two_sided, type);

	return status;
}

int nsl_dft_transform(double data[], size_t stride, size_t n, int two_sided, nsl_dft_result_type type) {
	if (n < 2) // we need at least 2 points
		return 1;
	size_t i;
	double* result = (double*)malloc(2 * n * sizeof(double));
	size_t N = n / 2; /* number of resulting data points */
	if (two_sided)
		N = n;
#ifdef HAVE_FFTW3
	/* stride ignored */
	(void)stride;

	fftw_plan plan = fftw_plan_dft_r2c_1d((int)n, data, (fftw_complex*)result, FFTW_ESTIMATE);
	fftw_execute(plan);
	fftw_destroy_plan(plan);

	/* 2. unpack data */
	if (two_sided) {
		for (i = 1; i < n - i; i++) {
			result[2 * (n - i)] = result[2 * i];
			result[2 * (n - i) + 1] = -result[2 * i + 1];
		}
		if (i == n - i) {
			result[2 * i] = result[2 * (n - 1)];
			result[2 * i + 1] = 0;
		}
	}
#else
	/* 1. transform */
	gsl_fft_real_wavetable* real = gsl_fft_real_wavetable_alloc(n);
	gsl_fft_real_workspace* work = gsl_fft_real_workspace_alloc(n);

	gsl_fft_real_transform(data, stride, n, real, work);
	gsl_fft_real_wavetable_free(real);
	gsl_fft_real_workspace_free(work);

	/* 2. unpack data */
	gsl_fft_halfcomplex_unpack(data, result, stride, n);
#endif

	/* 3. write result */
	switch (type) {
	case nsl_dft_result_magnitude:
		for (i = 0; i < N; i++)
			data[i] = sqrt(gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1]));
		break;
	case nsl_dft_result_amplitude:
		for (i = 0; i < N; i++) {
			data[i] = sqrt(gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / n;
			if (i > 0)
				data[i] *= 2.;
		}
		break;
	case nsl_dft_result_real:
		for (i = 0; i < N; i++)
			data[i] = result[2 * i];
		break;
	case nsl_dft_result_imag:
		for (i = 0; i < N; i++)
			data[i] = result[2 * i + 1];
		break;
	case nsl_dft_result_power:
		for (i = 0; i < N; i++) {
			data[i] = (gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / n;
			if (i > 0)
				data[i] *= 2.;
		}
		break;
	case nsl_dft_result_phase:
		for (i = 0; i < N; i++)
			data[i] = -atan2(result[2 * i + 1], result[2 * i]);
		break;
	case nsl_dft_result_dB:
		for (i = 0; i < N; i++) {
			if (i == 0)
				data[i] = 20. * log10(sqrt(gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / (double)n);
			else
				data[i] = 20. * log10(2. * sqrt(gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / (double)n);
		}
		break;
	case nsl_dft_result_normdB: {
		double maxdB = 0;
		for (i = 0; i < N; i++) {
			if (i == 0)
				data[i] = 20. * log10(sqrt(gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / (double)n);
			else
				data[i] = 20. * log10(2. * sqrt(gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / (double)n);

			if (i == 0 || maxdB < data[i])
				maxdB = data[i];
		}
		for (i = 0; i < N; i++)
			data[i] -= maxdB;
		break;
	}
	case nsl_dft_result_squaremagnitude:
		for (i = 0; i < N; i++)
			data[i] = gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1]);
		break;
	case nsl_dft_result_squareamplitude:
		for (i = 0; i < N; i++) {
			data[i] = (gsl_pow_2(result[2 * i]) + gsl_pow_2(result[2 * i + 1])) / gsl_pow_2((double)n);
			if (i > 0)
				data[i] *= 4.;
		}
		break;
	case nsl_dft_result_raw:
#ifdef HAVE_FFTW3
		// write gsl_halfcomplex data
		data[0] = result[0];
		for (i = 1; i < N; i++)
			data[i] = result[i + 1];
#endif
		break;
	}

	free(result);

	return 0;
}
