/*********************************************************************
 * This file is part of the PRAPI library.
 *
 * Copyright (C) 2001-2002 Topi Mäenpää
 * 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.5 $
 *********************************************************************/

#ifndef _CONVOLUTION_H
#define _CONVOLUTION_H

#include <List.h>
#include <Matrix.h>
#include <ListUtils.h>
#include <MatrixUtils.h>
#include <Complex.h>

namespace prapi { namespace dsp {
	/**
	 * Methods for performing operations related to convolution and
	 * correlation.
	 **/
	class Convolution
	{
	public:
		/**
		 * Two-dimensional convolution of two matrices. Convolution is
		 * performed on pixels on which it is fully defined. If you want
		 * to retain the original size of the input matrix, pad it with
		 * zeros. The size of the resulting matrix will be (M-O+1)x(N-P+1)
		 * if the input matrices are of sizes MxN and OxP. Make sure b is
		 * not larger than a.
		 *
		 * @param a source matrix
		 * @param b convolution filter
		 **/
		template <class T> static util::Matrix<T> conv(const util::Matrix<T>& a,
																									 const util::Matrix<T>& b);

		/**
		 * One-dimensional convolution of a matrix and a vector. The size
		 * of the resulting matrix will be either (M-L+1)xN or Mx(N-L+1)
		 * where the size of the input matrix is MxN and L is the length
		 * of the vector.
		 *
		 * @param a source matrix
		 * @param b convolution filter
		 * @param rows if true, convolution is calculated vertically
		 * (rows are convolved), otherwise it is calculated horizontally.
		 **/
		template <class T> static util::Matrix<T> conv(const util::Matrix<T>& a,
																									 const util::List<T>& b,
																									 bool rows = true);
		
		/**
		 * One-dimensional convolution of two vectors. The length of the
		 * resulting vector will be M-N+1 where M and N are the length of
		 * the input vectors a and b, respectively. If zero-padding is in
		 * use, the result will be M+N-1 items long.
		 *
		 * @param a source vector
		 * @param b convolution filter
		 * @param zeropad if true vectors add padded with zero if needeed.
		 *                (Convolution will be calculated for all points).
		 **/
		template <class T> static util::List<T> conv(const util::List<T>& a,
																								 const util::List<T>& b,
																								 bool zeropad = false);

		/**
		 * Two-dimensional cross-correlation of two matrices. Correlation
		 * is performed on pixels on which it is fully defined. If you
		 * want to retain the original size of the input matrix, pad it
		 * with zeros. The size of the resulting matrix will be
		 * (M-O+1)x(N-P+1) if the input matrices are of sizes MxN and OxP.
		 * Make sure that b is smaller than a. Autocorrelation can be
		 * calculated by padding a with zeros and performing the xcorr
		 * operation on the padded an the original matrix.
		 *
		 * @param a source matrix
		 * @param b convolution filter
		 **/
		template <class T> static util::Matrix<T> xcorr(const util::Matrix<T>& a,
																										const util::Matrix<T>& b);

		/**
		 * One-dimensional cross-correlation of a matrix and a vector. The
		 * size of the resulting matrix will be either (M-L+1)xN or
		 * Mx(N-L+1) where the size of the input matrix is MxN and L is
		 * the length of the vector.
		 *
		 * @param a source matrix
		 * @param b convolution filter
		 * @param rows if true, convolution is calculated vertically
		 * (rows are convolved), otherwise it is calculated horizontally.
		 **/
		template <class T> static util::Matrix<T> xcorr(const util::Matrix<T>& a,
																										const util::List<T>& b,
																										bool rows = true);

		/**
		 * One-dimensional cross-correlation of two vectors. The length of
		 * the resulting vector will be M-N+1 where M and N are the length
		 * of the input vectors a and b, respectively. If the length of
		 * the input vector is to be preserved, pad a with zeros.
		 *
		 * @param a source vector
		 * @param b convolution filter
		 **/
		template <class T> static util::List<T> xcorr(const util::List<T>& a,
																									const util::List<T>& b);
		
		/**
		 * Transform each complex value in a List into its complex
		 * conjugate.
		 **/
		template <class T> static void conjugate(util::List<util::Complex<T> >& lst);
		/**
		 * Transform each complex value in a Matrix into its complex
		 * conjugate.
		 **/
		template <class T> static void conjugate(util::Matrix<util::Complex<T> >& lst);

		template <class T> static void conjugate(util::List<T>& lst) {}
		template <class T> static void conjugate(util::Matrix<T>& lst) {}

	private:
		template <class T> static util::Matrix<T> correlate(const util::Matrix<T>& a,
																												const util::Matrix<T>& b);
		template <class T> static util::Matrix<T> correlate(const util::Matrix<T>& a,
																												const util::List<T>& b,
																												bool rows = true);
		template <class T> static util::List<T> correlate(const util::List<T>& a,
																											const util::List<T>& b);

		/**
		 * One-dimensional convolution of two vectors with zero-padding
		 *
		 * @param a source vector
		 * @param b convolution filter
		 **/
		template <class T> static util::List<T> convZeroPad(const util::List<T>& a,
																												const util::List<T>& b);

	};

	template <class T> void Convolution::conjugate(util::List<util::Complex<T> >& lst)
	{
		util::Complex<T>* data = lst.getData();
		for (int i=lst.getLength();i--;data++)
			data->conjugate();
	}
	
	template <class T> void Convolution::conjugate(util::Matrix<util::Complex<T> >& lst)
	{
		util::Complex<T>* data = lst.getData();
		for (int i=lst.getLength();i--;data++)
			data->conjugate();
	}

	template <class T> util::Matrix<T> Convolution::conv(const util::Matrix<T>& a,
																											 const util::List<T>& b,
																											 bool rows)
	{
		util::List<T> filter(b);
		util::ListUtils::reverse(filter);
		return correlate(a,filter,rows);
	}

	template <class T> util::Matrix<T> Convolution::conv(const util::Matrix<T>& a,
																											 const util::Matrix<T>& b)
	{
		util::Matrix<T> filter(b);
		util::MatrixUtils::turnAround(filter);
		return correlate(a,filter);
	}
	
	template <class T> util::List<T> Convolution::conv(const util::List<T>& a,
																										 const util::List<T>& b,
																										 bool zeropad)
	{
		if (zeropad)
			return convZeroPad(a,b);
		else
			{
				util::List<T> filter(b);
				util::ListUtils::reverse(filter);
				return correlate(a,filter);
			}
	}
	
	template <class T> util::Matrix<T> Convolution::xcorr(const util::Matrix<T>& a,
																												const util::Matrix<T>& b)
	{
		util::Matrix<T> func(a);
		conjugate(func);
		return correlate(func,b);
	}
																												
	template <class T> util::Matrix<T> Convolution::xcorr(const util::Matrix<T>& a,
																												const util::List<T>& b,
																												bool rows)
	{
		util::Matrix<T> func(a);
		conjugate(func);
		return correlate(func,b);
	}

	template <class T> util::List<T> Convolution::xcorr(const util::List<T>& a,
																											const util::List<T>& b)
	{
		util::List<T> func(a);
		conjugate(func);
		return correlate(func,b);
	}

	template <class T> util::Matrix<T> Convolution::correlate(const util::Matrix<T>& a,
																														const util::Matrix<T>& b)
	{
		int ar = a.getRows(), ac = a.getColumns(), br = b.getRows(), bc = b.getColumns();
		int rr = ar-br+1, rc = ac-bc+1;
		util::Matrix<T> result(rr,rc);
		T* rdata = result.getData();
		for (int row=0;row<rr;row++)
			for (int column=0;column<rc;column++,rdata++)
				{
					T sum(0);
					const T* adata = a.getData() + row*ac + column;
					const T* bdata = b.getData();
					for (int r=0;r<br;r++,adata += ac-bc)
						for (int c=0;c<bc;c++,adata++,bdata++)
							sum += *adata * *bdata;
					*rdata = sum;
				}
		return result;
	}

	template <class T> util::Matrix<T> Convolution::correlate(const util::Matrix<T>& a,
																														const util::List<T>& b,
																														bool rows)
	{
		int len = b.getLength();
		if (!rows)
			{
				util::Matrix<T> result(a.getRows(),a.getColumns()-len+1,false);
				T* rdata = result.getData();
				const T* sdata = a.getData();

				for (int r=result.getRows();r--;sdata+=len-1)
					for (int c=result.getColumns();c--;rdata++,sdata++)
						{
							T tmp(*sdata * b[0]);
							for (int i=1;i<len;i++)
								tmp += sdata[i] * b[i];
							*rdata = tmp;
						}
				return result;
			}
		else
			{
				int cols = a.getColumns();
				util::Matrix<T> result(a.getRows()-len+1,cols,false);
				T* rdata = result.getData();
				const T* sdata = a.getData();

				for (int r=result.getRows();r--;)
					for (int c=result.getColumns();c--;rdata++,sdata++)
						{
							T tmp(*sdata * b[0]);
							for (int i=1;i<len;i++)
								tmp += sdata[i*cols] * b[i];
							*rdata = tmp;
						}
				return result;
			}
	}

	template <class T> util::List<T> Convolution::correlate(const util::List<T>& a,
																													const util::List<T>& b)
	{
		int alen = a.getLength(), blen = b.getLength();
		int rlen = alen - blen + 1;
		util::List<T> result(rlen);
		result.setLength(rlen);
		T* rdata = result.getData();
		for (int i=0;i<rlen;i++,rdata++)
			{
				const T* adata = a.getData() + i, *bdata = b.getData();
				T sum(0);
				for (int j=blen;j--;adata++,bdata++)
					sum += *adata * *bdata;
				*rdata = sum;
			}
		return result;
	}
	
	template <class T> util::List<T> Convolution::convZeroPad(const util::List<T>& a,
																														const util::List<T>& b)
	{	
		int alen = a.getLength(), blen = b.getLength();
		int rlen = alen + blen - 1;
		util::List<T> result(rlen);
		result.setLength(rlen,T(0));
		T* rdata = result.getData();
		const T* bdata = b.getData();
		for (int i=0;i<blen;i++,rdata++,bdata++)
			{
				const T* adata = a.getData();
				for (int j=0;j<alen;j++,adata++)
						*(rdata+j) += *adata * *bdata;
			}
		return result;
	}
}}

#endif