/*
  Copyright (c) 2015, Rafat Hussain
*/
/*
This code is a C translation ( with some modifications) of Wavelet Software provided by
C. Torrence and G. Compo, and is available at URL: http://atoc.colorado.edu/research/wavelets/''.
*/

#include "cwt.h"

/*
static double factorial3(int N)
{
	double factorial = 1;
	for (int i = 1; i <= N; ++i) { factorial *= i; }
	return factorial;
}
*/
double factorial(int N)
{
	if (N > 40)
	{
		printf("This program is only valid for N <= 40 \n");
		return -1.0;
	}
	double fact[41] = {
		1, 1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800, 39916800, 479001600, 6227020800, 87178291200, 1307674368000,
		20922789888000, 355687428096000, 6402373705728000, 121645100408832000.0, 2432902008176640000.0, 51090942171709440000.0, 1124000727777607680000.0,
		25852016738884976640000.0, 620448401733239439360000.0, 15511210043330985984000000.0, 403291461126605635584000000.0, 10888869450418352160768000000.0,
		304888344611713860501504000000.0, 8841761993739701954543616000000.0, 265252859812191058636308480000000.0, 8222838654177922817725562880000000.0,
		263130836933693530167218012160000000.0, 8683317618811886495518194401280000000.0, 295232799039604140847618609643520000000.0,
		10333147966386144929666651337523200000000.0,
		371993326789901217467999448150835200000000.0, 13763753091226345046315979581580902400000000.0, 523022617466601111760007224100074291200000000.0,
		20397882081197443358640281739902897356800000000.0, 815915283247897734345611269596115894272000000000.0
	};

	return fact[N];
}

static void wave_function(int nk, double dt, int mother, double param, double scale1, double* kwave, double pi, double* period1, double* coi1,
						  fft_data* daughter)
{
	double norm, expnt, fourier_factor;
	int k, m;
	double temp;
	int sign, re;


	if (mother == 0)
	{
		//MORLET
		if (param < 0.0) { param = 6.0; }
		norm = sqrt(2.0 * pi * scale1 / dt) * pow(pi, -0.25);

		for (k = 1; k <= nk / 2 + 1; ++k)
		{
			temp               = (scale1 * kwave[k - 1] - param);
			expnt              = -0.5 * temp * temp;
			daughter[k - 1].re = norm * exp(expnt);
			daughter[k - 1].im = 0.0;
		}
		for (k = nk / 2 + 2; k <= nk; ++k) { daughter[k - 1].re = daughter[k - 1].im = 0.0; }
		fourier_factor = (4.0 * pi) / (param + sqrt(2.0 + param * param));
		*period1       = scale1 * fourier_factor;
		*coi1          = fourier_factor / sqrt(2.0);
	}
	else if (mother == 1)
	{
		// PAUL
		if (param < 0.0) { param = 4.0; }
		m    = (int)param;
		norm = sqrt(2.0 * pi * scale1 / dt) * (pow(2.0, (double)m) / sqrt(m * factorial(2 * m - 1)));
		for (k = 1; k <= nk / 2 + 1; ++k)
		{
			temp               = scale1 * kwave[k - 1];
			expnt              = - temp;
			daughter[k - 1].re = norm * pow(temp, (double)m) * exp(expnt);
			daughter[k - 1].im = 0.0;
		}
		for (k = nk / 2 + 2; k <= nk; ++k) { daughter[k - 1].re = daughter[k - 1].im = 0.0; }
		fourier_factor = (4.0 * pi) / (2.0 * m + 1.0);
		*period1       = scale1 * fourier_factor;
		*coi1          = fourier_factor * sqrt(2.0);
	}
	else if (mother == 2)
	{
		if (param < 0.0) { param = 2.0; }
		m = (int)param;

		if (m % 2 == 0) { re = 1; }
		else { re = 0; }

		if (m % 4 == 0 || m % 4 == 1) { sign = -1; }
		else { sign = 1; }


		norm = sqrt(2.0 * pi * scale1 / dt) * sqrt(1.0 / cwt_gamma(m + 0.50));
		norm *= sign;

		if (re == 1)
		{
			for (k = 1; k <= nk; ++k)
			{
				temp               = scale1 * kwave[k - 1];
				daughter[k - 1].re = norm * pow(temp, (double)m) * exp(-0.50 * pow(temp, 2.0));
				daughter[k - 1].im = 0.0;
			}
		}
		else if (re == 0)
		{
			for (k = 1; k <= nk; ++k)
			{
				temp               = scale1 * kwave[k - 1];
				daughter[k - 1].re = 0.0;
				daughter[k - 1].im = norm * pow(temp, (double)m) * exp(-0.50 * pow(temp, 2.0));
			}
		}
		fourier_factor = (2.0 * pi) * sqrt(2.0 / (2.0 * m + 1.0));
		*period1       = scale1 * fourier_factor;
		*coi1          = fourier_factor / sqrt(2.0);
	}
}

int cwavelet(double* y, int N, double dt, int mother, double param, double s0, double dj, int jtot, int npad, double* wave, double* scale, double* period,
			 double* coi)
{
	double period1, coi1;

	const double pi = 4.0 * atan(1.0);

	if (npad < N)
	{
		printf("npad must be >= N \n");
		return 1;
	}

	const fft_object obj  = fft_init(npad, 1);
	const fft_object iobj = fft_init(npad, -1);

	fft_data* ypad     = (fft_data*)malloc(sizeof(fft_data) * npad);
	fft_data* yfft     = (fft_data*)malloc(sizeof(fft_data) * npad);
	fft_data* daughter = (fft_data*)malloc(sizeof(fft_data) * npad);
	double* kwave      = (double*)malloc(sizeof(double) * npad);

	double ymean = 0.0;

	for (int i = 0; i < N; ++i) { ymean += y[i]; }

	ymean /= N;

	for (int i = 0; i < N; ++i)
	{
		ypad[i].re = y[i] - ymean;
		ypad[i].im = 0.0;
	}

	for (int i = N; i < npad; ++i) { ypad[i].re = ypad[i].im = 0.0; }


	// Find FFT of the input y (ypad)

	fft_exec(obj, ypad, yfft);

	for (int i = 0; i < npad; ++i)
	{
		yfft[i].re /= (double)npad;
		yfft[i].im /= (double)npad;
	}


	//Construct the wavenumber array

	const double freq1 = 2.0 * pi / ((double)npad * dt);
	kwave[0]           = 0.0;

	for (int i = 1; i < npad / 2 + 1; ++i) { kwave[i] = i * freq1; }

	for (int i = npad / 2 + 1; i < npad; ++i) { kwave[i] = -kwave[npad - i]; }
	
	// Main loop

	for (int j = 1; j <= jtot; ++j)
	{
		const double scale1 = scale[j - 1];// = s0*pow(2.0, (double)(j - 1)*dj);
		wave_function(npad, dt, mother, param, scale1, kwave, pi, &period1, &coi1, daughter);
		period[j - 1] = period1;
		for (int k = 0; k < npad; ++k)
		{
			const double tmp1 = daughter[k].re * yfft[k].re - daughter[k].im * yfft[k].im;
			const double tmp2 = daughter[k].re * yfft[k].im + daughter[k].im * yfft[k].re;
			daughter[k].re    = tmp1;
			daughter[k].im    = tmp2;
		}
		fft_exec(iobj, daughter, ypad);
		const int iter = 2 * (j - 1) * N;
		for (int i = 0; i < N; ++i)
		{
			wave[iter + 2 * i]     = ypad[i].re;
			wave[iter + 2 * i + 1] = ypad[i].im;
		}
	}


	for (int i = 1; i <= (N + 1) / 2; ++i)
	{
		coi[i - 1] = coi1 * dt * ((double)i - 1.0);
		coi[N - i] = coi[i - 1];
	}


	free(kwave);
	free(ypad);
	free(yfft);
	free(daughter);

	free_fft(obj);
	free_fft(iobj);

	return 0;
}

void psi0(int mother, double param, double* val, int* real)
{
	int sign;

	const int m     = (int)param;
	const double pi = 4.0 * atan(1.0);

	if (mother == 0)
	{
		// Morlet
		*val  = 1.0 / sqrt(sqrt(pi));
		*real = 1;
	}
	else if (mother == 1)
	{
		//Paul
		if (m % 2 == 0) { *real = 1; }
		else { *real = 0; }

		if (m % 4 == 0 || m % 4 == 1) { sign = 1; }
		else { sign = -1; }
		*val = sign * pow(2.0, (double)m) * factorial(m) / (sqrt(pi * factorial(2 * m)));
	}
	else if (mother == 2)
	{
		// D.O.G
		*real = 1;

		if (m % 2 == 0)
		{
			if (m % 4 == 0) { sign = -1; }
			else { sign = 1; }
			const double coeff = sign * pow(2.0, (double)m / 2) / cwt_gamma(0.5);
			*val               = coeff * cwt_gamma(((double)m + 1.0) / 2.0) / sqrt(cwt_gamma(m + 0.50));
		}
		else { *val = 0; }
	}
}

static int maxabs(double* array, int N)
{
	double maxval = 0.0;
	int index     = -1;

	for (int i = 0; i < N; ++i)
	{
		const double temp = fabs(array[i]);
		if (temp >= maxval)
		{
			maxval = temp;
			index  = i;
		}
	}

	return index;
}


double cdelta(int mother, double param, double psi0)
{
	int N     = 0;
	double s0 = 0;

	double subscale = 8.0;
	const double dt = 0.25;
	if (mother == 0)
	{
		N  = 16;
		s0 = dt / 4;
	}
	else if (mother == 1)
	{
		N  = 16;
		s0 = dt / 4.0;
	}
	else if (mother == 2)
	{
		s0 = dt / 8.0;
		N  = 256;
		if (param == 2.0)
		{
			subscale = 16.0;
			s0       = dt / 16.0;
			N        = 2048;
		}
	}

	const double dj = 1.0 / subscale;
	const int jtot  = 16 * (int)subscale;

	double* delta  = (double*)malloc(sizeof(double) * N);
	double* wave   = (double*)malloc(sizeof(double) * 2 * N * jtot);
	double* coi    = (double*)malloc(sizeof(double) * N);
	double* scale  = (double*)malloc(sizeof(double) * jtot);
	double* period = (double*)malloc(sizeof(double) * jtot);
	double* mval   = (double*)malloc(sizeof(double) * N);


	delta[0] = 1;

	for (int i = 1; i < N; ++i) { delta[i] = 0; }

	for (int i = 0; i < jtot; ++i) { scale[i] = s0 * pow(2.0, (double)(i) * dj); }

	cwavelet(delta, N, dt, mother, param, s0, dj, jtot, N, wave, scale, period, coi);

	for (int i = 0; i < N; ++i) { mval[i] = 0; }

	for (int j = 0; j < jtot; ++j)
	{
		const int iter   = 2 * j * N;
		const double den = sqrt(scale[j]);
		for (int i = 0; i < N; ++i) { mval[i] += wave[iter + 2 * i] / den; }
	}


	const int maxarr  = maxabs(mval, N);
	const double cdel = sqrt(dt) * dj * mval[maxarr] / psi0;

	free(delta);
	free(wave);

	free(scale);
	free(period);
	free(coi);
	free(mval);

	return cdel;
}

void icwavelet(double* wave, int N, double* scale, int jtot, double dt, double dj, double cdelta, double psi0, double* oup)
{
	const double coeff = sqrt(dt) * dj / (cdelta * psi0);

	for (int i = 0; i < N; ++i) { oup[i] = 0.0; }

	for (int j = 0; j < jtot; ++j)
	{
		const int iter   = 2 * j * N;
		const double den = sqrt(scale[j]);
		for (int i = 0; i < N; ++i) { oup[i] += wave[iter + 2 * i] / den; }
	}

	for (int i = 0; i < N; ++i) { oup[i] *= coeff; }
}
