/*********************************************************************
 * This file is part of the PRAPI library.
 *
 * Copyright (C) 2001 Topi Mäenpää and Jaakko Viertola
 * All rights reserved.
 *
 * This program is free software. You can redistribute and/or modify
 * it under the terms of the free software licence found in the
 * accompanying file "COPYING". The licence terms must always be
 * redistributed with this source file. The above copyright notice
 * must be reproduced in all modified and unmodified copies of this
 * source file.
 *
 * $Revision: 1.6 $
 *********************************************************************/

#ifndef _FFT_H
#define _FFT_H

#include "../ImageTransform.h"
#include <Matrix.h>

#include <Complex.h>

#include <math.h>

using namespace util;

namespace prapi { namespace transforms {
	/**
	 * This enumeration is for the direction of the Fast Fourier transform.
	 * <ul>
	 * <li>FORWARD   - The direct fourier transform is made.
	 * <li>REVERSED  - The inversed fourier transform is made.
	 * </ul>
	 **/
	enum Direction { FORWARD, REVERSED };
	
 	/**
	 * This class includes FastFourierTransform algoritms for
	 * 2D transforms. This FFT class includes the high-speed radix-4 fast Fourier transform and basic fft
	 * (which calculates the evend and odds separately). If its possible to use radix4 Fourier transform
	 * it will be used otherwise the slower basic fft is used. If wanted to use fast Radix4 transform the
	 * size of the matrix must be power of 2 (or the points of transform) and the minimum number of
	 * points has to be 4. 
	 **/
	template <class T, class U> class FFT : public ImageTransform<Complex<T>,U>
	{
	public:
		/**
		 * The Constructor of FFT.
		 *
		 * @param direction The direction of transform forward / reversed.
		 **/
		FFT(Direction direction=FORWARD, int pointsInRow=-1, int pointsInCol=-1) :
			_Direction(direction),_iNumberOfPoints(0),_iPointsInRow(pointsInRow),_iPointsInCol(pointsInCol),
			_iPowerOfTwoPoints(0),_dScaleFactor(0.0),_lstSinTable(0),_lstBitTable(0),_lstPowers(0) {}
		/**
		 * Set the direction.
		 **/
		void setDirection(Direction dir){_Direction=dir;}
		/**
		 * Set the points.
		 **/
		void setPointsInRow(int points){_iPointsInRow=points;}
		void setPointsInColumn(int points){_iPointsInCol=points;}
		/**
		 * Get the direction.
		 **/
		Direction getDirection(){return _Direction;}
		/**
		 * Get the points.
		 **/
		int getPointsInRow(){return _iPointsInRow;}
		int getPointsInColumn(){return _iPointsInCol;}
		/**
		 * The fuction makes the FFT transform for the matrix given in parameter.
		 * If Matrix is nsmaller than points the zeros will be added to borders. And if the
		 * matrix is bigger than points the matrix will be truncated.
		 *
		 * @param mat The Matrix wanted to transform.
		 * @exception ImageTransformException if the pointsInRow/pointsInCol is not dividible for 2.
		 **/
		Matrix<Complex<T> > getTransformedImage(const Matrix<U>& mat)
			throw (ImageTransformException&);

	private:
		/**
		 * Makes the basic fft for even and odd values apart.
		 *
		 * @param data The list of complex numbers.
		 **/
		List<Complex<T> >& fft(List<Complex<T> >& data);
		/**
		 * Fuction makes the fft (using radix 4) for List of complex numbers by using
		 * function fftCalculation.
		 *
		 * @param data The list of complex numbers.
		 **/
		List<Complex<T> >& fftRadix4(List<Complex<T> >& data);
		/**
		 * This function makes the fft calculation for given data.
		 *
		 * @param data The data where the values are.
		 * @param buffer The precalculated values.
		 **/
		void fftCalculation(List<Complex<double> >& data,List<Complex<double> >& buffer);
		/**
		 * This fuction makes the true fft calculation recursively.
		 *
		 * @param data The list of complex numbers.
		 * @param level indicates how recursively deep the algorithm currently is.
		 * @param chunck tells us which is the starting node of the sub-FFT. 
		 **/
		void fftCalculationRadix4(List<Complex<T> >& data, int level, int chunk);
		/**
		 * The fuction initializes the fft transform making
		 * sin tables , bit tables and power pt 2.
		 *
		 * @param points The amount of points in transform.
		 **/
		void fftInitialization(int points);
		/**
		 * The fuction reverses the bits of the value.
		 **/
		int bitReverse(int bits);
		/**
		 * Direction of the transform FORWARD/REVERSED,
		 **/
		Direction _Direction;
		/**
		 * The nummber of points in the transform.
		 **/
		int _iNumberOfPoints;
		int _iPointsInRow;
		int _iPointsInCol;
		/**
		 * The 2^_iPowerOfTwoPoits = _iNumberOfPoints.
		 **/
		int _iPowerOfTwoPoints;
		/**
		 * The scale factor = 1/_iNumberOfPoits.
		 **/
		double _dScaleFactor;
		/**
		 * The Precalculated sin table.
		 **/
		List<double> _lstSinTable;
		/**
		 * Bit table for unscrambling.
		 **/
		List<int> _lstBitTable;
		/**
		 * The list for power of 2.
		 **/
		List<int> _lstPowers;
	};

	template <class T,class U> Matrix<Complex<T> > FFT<T,U>::getTransformedImage(const Matrix<U>& mat)
			throw (ImageTransformException&)
	{
		bool useFastRadix4 = true;
		const int rows=mat.getRows();
		const int cols=mat.getColumns();
		if(_iPointsInRow < 0)_iPointsInRow = rows;
		if(_iPointsInCol < 0)_iPointsInCol = cols;

		if(_iPointsInRow%2 !=0 || _iPointsInCol%2 !=0)
			throw ImageTransformException("FFT<T,U>::getTransformedImage(const Matrix<U>&) : The Matrix size must be divisible by 2.");
		
		// give the result matrix the size of fft
		Matrix<Complex<T> > result(_iPointsInRow,_iPointsInCol);
		
		//chek if zeros have to be added to borders.
		if(rows < _iPointsInRow || cols < _iPointsInCol)
			{ //add rows to bottom
				// because zeros are added automatically to the border
				// we dont have to care about them
				for(int r=0;r<rows;r++)
					for(int c=0;c<cols;c++)result(r,c)=mat(r,c);
			}//for rows
		else // if the size of columns and rows is the same or less
			for(int r=0;r<_iPointsInRow;r++)
				for(int c=0;c<_iPointsInCol;c++)result(r,c)=mat(r,c);

		// use fast radix 4 if the poits is power of 2
		// first check if its possible to use radix4 to cols
		int index=0,fact=1;
		for(fact=1,index=1;index<32;index++)if((fact<<=1)==_iPointsInRow)break;
		if(index == 32)useFastRadix4=false;
		
		//then check if fast radix 4 can be used.
		if(useFastRadix4 && (_iPointsInRow>=4))
			{ // Initialize the FFT if needed.
				if (_iNumberOfPoints != _iPointsInRow)fftInitialization(_iPointsInRow);
				// Transform Columns
				for(int i=0;i<_iPointsInCol;i++)
					{
						List<Complex<T> > lst(result.getColumn(i));
						result.setColumn(i,fftRadix4(lst));
					}
			}
		else // otherwise use the slower fft
			{
				
				// Perform the FFT in two dimensions 
				// Transform Columns
				_dScaleFactor = 1.0/(double)_iPointsInCol;
				for(int i=0;i<_iPointsInCol;i++)
					{
						List<Complex<T> > lst(result.getColumn(i));
						result.setColumn(i,fft(lst));
					}
			}
				
		//and then again check if the radix4 can be used for rows
		useFastRadix4=true;
		for(fact=1,index=1;index<32;index++)if((fact<<=1)==_iPointsInCol)break;
		if(index == 32)useFastRadix4=false;
		
		if(useFastRadix4 && (_iPointsInCol >= 4))
			{
				// Initialize the FFT again if not square Matrix
				if (_iNumberOfPoints != _iPointsInCol)fftInitialization(_iPointsInCol);
				// Transform Rows
				for(int i=0;i<_iPointsInRow;i++)
					{
						List<Complex<T> > lst(result.getRow(i));
						result.setRow(i,fftRadix4(lst));
					}
			}
		else // otherwise use the slower fft
			{
				_dScaleFactor = 1.0/(double)_iPointsInRow;
				// Transform Rows
				for(int i=0;i<_iPointsInRow;i++)
					{
						List<Complex<T> > lst(result.getRow(i));
						result.setRow(i,fft(lst));
					}
			}
		
		return result;
	}
	
	template <class T,class U> List<Complex<T> >& FFT<T,U>::fftRadix4(List<Complex<T> >& data)
	{	   
		fftCalculationRadix4(data, 1, 0 ); // Call recursive FFT routine to transform data 
		                   // Start at recursion level 1, begin at node 0 
		
		//Unscramble final values
		Complex<T> temp;
		for(int i=0;i<_iNumberOfPoints;i++)
			if(_lstBitTable[i] <= i)
				{  //If not yet de-scrambled 
					temp = data[i];
					data[i] = data[_lstBitTable[i]];
					data[_lstBitTable[i]] = temp;
				}
		
		// Multiply by scale factor 
		if(_Direction == REVERSED)data *= _dScaleFactor;
		
		return data;
	}
	
	template <class T,class U> void FFT<T,U>::fftCalculationRadix4(List<Complex<T> >& data, int level, int chunk)
	{
		int nodes; // nodes indicates how many nodes are in the sub-FFT under consideration.
		int  sinIndex, cosIndex; // Indices into sin table built by fftInitialization 
		int  dual1, dual2;
		double dir = -1.0; //for the direction of trnsform direct/inversed
		Complex<T>  dual1val,dual2val,dualprod,wp;
		
		nodes = _lstPowers[_iPowerOfTwoPoints-level];
		sinIndex = _lstBitTable[chunk/nodes]; // Get index into trig table for sin & cos values needed now 
		cosIndex = (sinIndex + _iNumberOfPoints/4) % _iNumberOfPoints;
		wp.real = _lstSinTable[cosIndex];
		if(_Direction == REVERSED)dir = 1.0; //change the direction if freversed
		wp.imag = dir*_lstSinTable[sinIndex];
		
    for(int i=0;i<nodes;i++)
			{
				dual1 = chunk + i;
				dual2 = dual1 + nodes;
				dual1val = data[dual1];
				dual2val = data[dual2];
				dualprod = dual2val * wp;
				data[dual1] = dual1val + dualprod;
				data[dual2] = dual1val - dualprod;
			}
		
		// are we ready or not
    if(level < _iPowerOfTwoPoints)
			{  
				fftCalculationRadix4(data, level+1, chunk); // Do top dual node pair sub-FFT 
				fftCalculationRadix4(data, level+1, chunk+nodes); // Do bottom dual node pair sub-FFT 
			}
	}
	
	template <class T,class U> void FFT<T,U>::fftInitialization(int points)
	{
		_iNumberOfPoints = points; // save the number of points
		_iPowerOfTwoPoints = int((log((double)points )/log( 2.0 )) + 0.5);//Compute nn s.t. 2^nn = _iNumberOfPoints
		_dScaleFactor = 1.0/(double)points;// Compute scale factor 
		
		// clear the tables
		_lstSinTable.setLength(0);
		_lstBitTable.setLength(0);
		_lstPowers.setLength(0);
		
		// build the power of two
		//for(int i=0;i<=_iPowerOfTwoPoints;i++)_lstPowers += pow( 2.0,(double)i) + 0.5;
		int base = 1;
		for (int i=0;i<=_iPowerOfTwoPoints;i++)
			{
				_lstPowers += base;
				base <<= 1;
			}
		// Build sine and bit reverse tables
		double pi2=2*M_PI;
		for(int i=0;i<_iNumberOfPoints;i++)
			{  
				_lstSinTable += sin(pi2*double(i)/double(_iNumberOfPoints)); 
				_lstBitTable += bitReverse(i);   
			}
	}
	
	template <class T,class U> int FFT<T,U>::bitReverse(int bits)
	{
		int lookmask = 1;   // Mask to look at bits in input index, start with bit 0
		int setmask = _iNumberOfPoints; // Look at MSBit of significance 
		setmask >>= 1; // But recall that indices from 0 to _iNumberOfPoints - 1 
		int tempbit = 0; // Initialize reversed value to 0 
		
		for(int i=0;i<_iPowerOfTwoPoints;i++)
			{  
				if( (bits & lookmask) == lookmask )tempbit = tempbit | setmask;
				lookmask <<= 1;     
				setmask >>= 1;      
			}
		return(tempbit);
	}
	
	template <class T,class U> List<Complex<T> >& FFT<T,U>::fft(List<Complex<T> >& data)
	{
		int n=data.getLength();
		//this are for the temp containers.  
		List<Complex<double> > sum(n);
		List<Complex<double> > diff(n);
		// and then the Even and Odds
		int m = n>>1; // divide by two
		List<Complex<double> > even(m);
		List<Complex<double> > odd(m);
		// remember to set the lengths
		sum.setLength(n);
		diff.setLength(n);
		even.setLength(m);
		odd.setLength(m);
		
		double init= atan(1.0)/double(n)*-8.0;
		double dir = -1.0; //for the direction of trnsform direct/inversed
		if(_Direction == REVERSED)dir = 1.0; //change the direction if freversed
		Complex<double> tmp(cos(init),-dir*sin(init));
		Complex<double> z1(1.0,0.0);
		
		for (int i=0; i<m; i++)
			{
				sum[i] = Complex<double>(data[i] + data[i+m]);
				diff[i] = z1*Complex<double>(data[i] - data[i+m]);
				z1 *= tmp;
			}
		
		fftCalculation(sum, even);
		fftCalculation(diff, odd);
		
		for(int i=0; i<m; i++)
			{
				data[(i<<1)] = Complex<T>(even[i]);
				data[((i<<1) + 1)] = Complex<T>(odd[i]);
			}
		
		// Multiply by scale factor 
		if(_Direction == REVERSED)data *= _dScaleFactor;
		
		return data;
	}
	
	template <class T, class U> void FFT<T,U>::fftCalculation(List<Complex<double> >& data,List<Complex<double> >& buffer)
	{
		// Constant factor -2 pi
		int n = buffer.getLength();
		double init= atan(1.0)/double(n)*-8.0;
		double dir = -1.0; //for the direction of trnsform direct/inversed
		if(_Direction == REVERSED)dir = 1.0; //change the direction if freversed
		Complex<double> tmp(cos(init),-dir*sin(init));
		
		// Pre-compute most of the exponential
		List<Complex<double> > pre(n);
		pre.setLength(n);
		pre[0].real=1.0;
		for (int i=1; i<n; i++)pre[i] = pre[i-1]*tmp;
		
		// Double loop to compute all Y entries 
		for (int m = 0; m<n; m++)
			{
				buffer[m]=data[0];		
				for (int k=1; k<=n-1; k++)buffer[m] += data[k]*pre[k*m % n];
		}
	}
}}
#endif