// CONVOLUTION - CORRELATION module

#include <iostream.h>
#include <iomanip.h>
#include <fstream.h>
#include <stdlib.h>
#include <math.h>
#include "simutils.h"
#include "vecmat.h"

// mini header
void realft(double data[], unsigned long n, int isign);
void twofft(double data1[], double data2[], double fft1[], double fft2[],
	unsigned long n);
void four1(double data[], unsigned long nn, int isign);

/*=================================
		HIGH LEVEL code
=================================*/
vector convolve(const vector& v1, const vector& v2)
{
	if(features.convcorr_byFFT()) return convolve_FFT(v1,v2);
	else return convolve_BRUTE(v1,v2);
}
	vector operator&(const vector& v1, const vector& v2)
	{
		return convolve(v1, v2);
	}
	vector operator&(const vector& v, const int& c)
	{
		vector result = copy(v);
		if(c <= 1) error("convolve-> operator needs int > 1");
		for(int i = 2; i <= c; ++i)
			result = convolve(result, v);
		return result;
	}
vector correlate(const vector& v1, const vector& v2)
{
	if(features.convcorr_byFFT()) return correlate_FFT(v1,v2);
	else return correlate_BRUTE(v1,v2);
}
	vector operator|(const vector& v1, const vector& v2)
	{
		return correlate(v1, v2);
	}
	vector operator|(const vector& v, const int& c)
	{
		vector result = copy(v);
		if(c <= 1) error("correlate-> operator needs int > 1");
		for(int i = 2; i <= c; ++i)
			result = correlate(result, v);
		return result;
	}

/*=================================
		BRUTE FORCE code
=================================*/
vector convolve_BRUTE(const vector& vec1, const vector& vec2)
{
        const vector* v1;
        const vector* v2;
        vector v(vec1.size + vec2.size - 1, 0.0);

        if(vec1.size <= vec2.size)
        {
                v1 = &vec1;
                v2 = &vec2;
        }
        else
        {
                v1 = &vec2;
                v2 = &vec1;
        }

        long int i, j;
        // part 1
        for(i = 0; i <= (v1->size - 1); ++i)
                for(j = 0; j <= i; ++j)
                        v.data[i] += v1->data[j] * v2->data[i - j];
        // part 2
        if(v2->size > v1->size)
                for(i = v1->size; i <= (v2->size - 1); ++i)
                        for(j = 0; j <= (v1->size - 1); ++j)
                                v.data[i] += v1->data[j] * v2->data[i - j];
        // part 3
        for(i = 1; i <= (v1->size - 1); ++i)
                for(j = i; j <= (v1->size - 1); ++j)
                        v.data[(v2->size - 1) + i] +=
                        v1->data[j] * v2->data[(v2->size - 1) + i - j];

        return v;
}
vector correlate_BRUTE(const vector& v1, const vector& v2)
{
        vector v1copy = copy(v1);
        v1copy.involve();
        return convolve_BRUTE(v1copy, v2);
}

/*=================================
		FFT code
=================================*/
vector convolve_FFT(const vector& v1, const vector& v2)
{
	long int len = max(v1.size, v2.size);
	
	// FIRST CALCULATE THE FFT RESULT'S DIMENSION
	long int n = 2*len - 1;
	int expon;
	double value;
	value = frexp(n, &expon);
	n = (long int) pow(2, expon);

	// DO THE FFT CONVOLUTION
	// set up the zero-padded vectors
	double *s, *r;
	s = new double[n];
	r = new double[n];
	long int i;
	for(i = 0; i < v1.size; ++i)
		s[i] = v1.data[i];
	for(i = v1.size; i < n; ++i)
		s[i] = 0.0;
	for(i = 0; i < v2.size; ++i)
		r[i] = v2.data[i];
	for(i = v2.size; i < n; ++i)
		r[i] = 0.0;
	--s; --r;
	// use modified nrc code
	double *fft, *ans;
	fft = new double[2*n];
	ans = new double[2*n];
	--fft; --ans;
	long int no2 = n/2;
	twofft(s,r,fft,ans,n);

	double dum;
	for (i=2;i<=n+2;i+=2)
	{
		ans[i-1]=(fft[i-1]*(dum=ans[i-1])-fft[i]*ans[i])/no2;
      ans[i]=(fft[i]*dum+fft[i-1]*ans[i])/no2;
	}
   ans[2]=ans[n+1];
   realft(ans,n,-1);

	vector result;
	result.data = ans+1;
	result.size = v1.size + v2.size - 1;

	// CLEAN UP
	delete[] (s+1);
	delete[] (r+1);
	delete[] (fft+1);

	return result;
}

vector correlate_FFT(const vector& v1, const vector& v2)
{
	vector v1copy = copy(v1);
	v1copy.involve();
	return convolve(v1copy,v2);
}

/*===============================================================================
								NUMERICAL RECIPES ROUTINES
================================================================================*/
void twofft(double data1[], double data2[], double fft1[], double fft2[],
        unsigned long n)
{
        unsigned long nn3,nn2,jj,j;
        double rep,rem,aip,aim;

        nn3=1+(nn2=2+n+n);
        for (j=1,jj=2;j<=n;j++,jj+=2) {
                fft1[jj-1]=data1[j];
                fft1[jj]=data2[j];
        }
        four1(fft1,n,1);
        fft2[1]=fft1[2];
        fft1[2]=fft2[2]=0.0;
        for (j=3;j<=n+1;j+=2) {
                rep=0.5*(fft1[j]+fft1[nn2-j]);
                rem=0.5*(fft1[j]-fft1[nn2-j]);
                aip=0.5*(fft1[j+1]+fft1[nn3-j]);
                aim=0.5*(fft1[j+1]-fft1[nn3-j]);
                fft1[j]=rep;
                fft1[j+1]=aim;
                fft1[nn2-j]=rep;
                fft1[nn3-j] = -aim;
                fft2[j]=aip;
                fft2[j+1] = -rem;
                fft2[nn2-j]=aip;
                fft2[nn3-j]=rem;
        }
}
void realft(double data[], unsigned long n, int isign)
{
        unsigned long i,i1,i2,i3,i4,np3;
        double c1=0.5,c2,h1r,h1i,h2r,h2i;
        double wr,wi,wpr,wpi,wtemp,theta;

        theta=3.141592653589793/(double) (n>>1);
        if (isign == 1) {
                c2 = -0.5;
                four1(data,n>>1,1);
        } else {
                c2=0.5;
                theta = -theta;
        }
        wtemp=sin(0.5*theta);
        wpr = -2.0*wtemp*wtemp;
        wpi=sin(theta);
        wr=1.0+wpr;
        wi=wpi;
        np3=n+3;
        for (i=2;i<=(n>>2);i++) {
                i4=1+(i3=np3-(i2=1+(i1=i+i-1)));
                h1r=c1*(data[i1]+data[i3]);
                h1i=c1*(data[i2]-data[i4]);
                h2r = -c2*(data[i2]+data[i4]);
                h2i=c2*(data[i1]-data[i3]);
                data[i1]=h1r+wr*h2r-wi*h2i;
                data[i2]=h1i+wr*h2i+wi*h2r;
                data[i3]=h1r-wr*h2r+wi*h2i;
                data[i4] = -h1i+wr*h2i+wi*h2r;
                wr=(wtemp=wr)*wpr-wi*wpi+wr;
                wi=wi*wpr+wtemp*wpi+wi;
        }
        if (isign == 1) {
                data[1] = (h1r=data[1])+data[2];
                data[2] = h1r-data[2];
        } else {
                data[1]=c1*((h1r=data[1])+data[2]);
                data[2]=c1*(h1r-data[2]);
                four1(data,n>>1,-1);
        }
}
#define SWAP(a,b) tempr=(a);(a)=(b);(b)=tempr
void four1(double data[], unsigned long nn, int isign)
{
        unsigned long n,mmax,m,j,istep,i;
        double wtemp,wr,wpr,wpi,wi,theta;
        double tempr,tempi;

        n=nn << 1;
        j=1;
        for (i=1;i<n;i+=2) {
                if (j > i) {
                        SWAP(data[j],data[i]);
                        SWAP(data[j+1],data[i+1]);
                }
                m=n >> 1;
                while (m >= 2 && j > m) {
                        j -= m;
                        m >>= 1;
                }
                j += m;
        }
        mmax=2;
        while (n > mmax) {
                istep=mmax << 1;
                theta=isign*(6.28318530717959/mmax);
                wtemp=sin(0.5*theta);
                wpr = -2.0*wtemp*wtemp;
                wpi=sin(theta);
                wr=1.0;
                wi=0.0;
                for (m=1;m<mmax;m+=2) {
                        for (i=m;i<=n;i+=istep) {
                                j=i+mmax;
                                tempr=wr*data[j]-wi*data[j+1];
                                tempi=wr*data[j+1]+wi*data[j];
                                data[j]=data[i]-tempr;
                                data[j+1]=data[i+1]-tempi;
                                data[i] += tempr;
                                data[i+1] += tempi;
                        }
                        wr=(wtemp=wr)*wpr-wi*wpi+wr;
                        wi=wi*wpr+wtemp*wpi+wi;
                }
                mmax=istep;
        }
}
#undef SWAP