/*********************************************************************
 * This file is part of the cpplibs suite.
 *
 * Copyright (C) 2001 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.4 $
 *********************************************************************/

#ifndef _WAVELET_H
#define _WAVELET_H

#include <List.h>

#include "Convolution.h"
#include "Daubechies.h"
#include <String.h>
#include <Exception.h>
#include <MatrixUtils.h>
#include <Math.h>


namespace prapi { namespace dsp {
	/**
	 * Methods for creating wavelets and performing the wavelet
	 * decomposition of 1-D or 2-D signals.
	 **/
	class Wavelet
	{
	public:
		/**
		 * Compute a quadrature mirror filter for a filter. The qmf is a
		 * reversed version of the input filter with every second item
		 * negated.
		 *
		 * @param filter the input filter
		 * @param odd if one, odd elements are negated. Otherwise, even
		 * elements are negated.
		 * @return the qmf filter
		 **/
		template <class T> static util::List<T> qmf(const util::List<T>& filter, int odd=1);

		/**
		 * Create an orthogonal wavelet filter set from the given scaling
		 * filter.
		 *
		 * @param filter a scaling filter for a wavelet
		 *
		 * @return a list of four filters: decomposition low-pass,
		 * decomposition high-pass, reconstruction low-pass and
		 * reconstruction high-pass in this order.
		 **/
		template <class T> static util::List<util::List<T> > orthfilt(const util::List<T>& filter);

		/**
		 * Performs a two-dimensional one-level discrete wavelet transform
		 * on the input matrix.
		 *
		 * @param mat the input matrix
		 * @param lo low-pass decomposition filter
		 * @param hi high-pass decomposition filter
		 * @param result a vector for storing the decomposition result
		 * @param type the type of matrix extension (for handling border effects)
		 *
		 * @return four matrices, the first one containing approximation
		 * coefficients and the last three containing horizontal, vertical
		 * and diagonal details in this order.
		 **/
		template <class T> static util::List<util::Matrix<T> >& dwt(const util::Matrix<T>& mat,
																																const util::List<T>& lo,
																																const util::List<T>& hi,
																																util::List<util::Matrix<T> >& result,
																																util::MatrixUtils::ExtendType type = util::MatrixUtils::EXTEND_ZEROPAD);
		/**
		 * Performs a two-dimensional one-level discrete wavelet transform
		 * on the input matrix.
		 *
		 * @param mat the input matrix
		 * @param wavelet the name of the wavelet to be used.
		 * @param result a vector for storing the decomposition result
		 * @param type the type of matrix extension (for handling border effects)
		 *
		 * @return four matrices, the first one containing approximation
		 * coefficients and the last three containing horizontal, vertical
		 * and diagonal details in this order.
		 *
		 * @exception InvalidArgumentException& if the name is not valid
		 *
		 * @see #getFilter(string)
		 **/
		template <class T> static util::List<util::Matrix<T> >& dwt(const util::Matrix<T>& mat,
																																std::string wavelet,
																																util::List<util::Matrix<T> >& result,
																																util::MatrixUtils::ExtendType type = util::MatrixUtils::EXTEND_ZEROPAD)
			throw (util::InvalidArgumentException&);

		/**
		 * Get the named wavelet scaling filter. Filters currently
		 * supported are:
		 * <ul>
		 * <li>Haar: "haar" (="db1")
		 * <li>Daubechies: "db1" ... "db10"
		 * </ul>
		 *
		 * @exception InvalidArgumentException& if the name is not valid
		 **/
		template <class T> static util::List<T> getFilter(std::string name) throw (util::InvalidArgumentException&);
		
		/**
		 * Perform a dyadic downsampling on the input matrix.
		 *
		 * @param param mat the input matrix
		 * @param rows if true, rows are removed
		 * @param cols if true, colums are removed
		 * @param odd if one, odd rows and/or columns are removed. If
		 * zero, even rows and/or columns are removed
		 *
		 * @return the input matrix with odd or even rows and/or columns
		 * removed
		 **/
		template <class T> static util::Matrix<T> dyaddown(const util::Matrix<T>& mat,
																											 bool rows,
																											 bool columns,
																											 int odd = 0);

		/**
		 * Take the central part of a matrix.
		 *
		 * @param mat the matrix whose center is to be taken.
		 * @param rows the number of rows in the returned matrix
		 * @param cols the number of columns in the returned matrix
		 **/
		template <class T> static util::Matrix<T> keep(const util::Matrix<T>& mat,
																									 int rows, int cols)
		{
			return mat((mat.getRows()-rows)>>1,(mat.getColumns()-cols)>>1,rows,cols);
		}

	private:
		template <class T> static util::Matrix<T> convolveAndDownSample(const util::Matrix<T>& mat,
																																		const util::List<T>& filter,
																																		util::MatrixUtils::ExtendType type);
	};

	template <class T> util::List<util::Matrix<T> >& Wavelet::dwt(const util::Matrix<T>& mat,
																																std::string name,
																																util::List<util::Matrix<T> >& result,
																																util::MatrixUtils::ExtendType type)
		throw (util::InvalidArgumentException&)
	{
		util::List<util::List<T> > filters(orthfilt(getFilter<T>(name)));
		return dwt(mat,filters[0],filters[1],result,type);
	}

	template <class T> util::List<T> Wavelet::getFilter(std::string name)
		throw (util::InvalidArgumentException&)
	{
		if (name.size() > 2 && name.substr(0,2) == "db")
			{
				int index = util::String::parse<int>(name.substr(2));
				return Daubechies::getFilter(index);
			}
		else if (name == "haar")
			return Daubechies::getFilter(1);
		else
			throw util::InvalidArgumentException("Wavelet::getFilter(string): Unknown wavelet.");
	}

	template <class T> util::List<util::Matrix<T> >& Wavelet::dwt(const util::Matrix<T>& mat,
																																const util::List<T>& lo,
																																const util::List<T>& hi,
																																util::List<util::Matrix<T> >& result,
																																util::MatrixUtils::ExtendType type)
	{
		int len = lo.getLength();

		util::Matrix<T> x(util::MatrixUtils::extend(mat,len-1,type,1+2)); //top and bottom for convolution
		util::Matrix<T> tmp(Convolution::conv(x,lo,true));

		result += convolveAndDownSample(tmp,lo,type);
		result += convolveAndDownSample(tmp,hi,type);

		tmp = Convolution::conv(x,hi,true);

		result += convolveAndDownSample(tmp,lo,type);
		result += convolveAndDownSample(tmp,hi,type);

		return result;
	}

	template <class T> util::Matrix<T> Wavelet::convolveAndDownSample(const util::Matrix<T>& mat,
																																		const util::List<T>& filter,
																																		util::MatrixUtils::ExtendType type)
	{
		int len = filter.getLength()-1;
		util::Matrix<T> result(Convolution::conv(util::MatrixUtils::extend(mat,len,type,4+8), //left-right
																						 filter,
																						 false));
		result = dyaddown(result,true,true,1);
		return result;
	}
	
	template <class T> util::List<T> Wavelet::qmf(const util::List<T>& filter, int odd)
	{
		util::List<T> result(filter);
		util::ListUtils::reverse(result);
		for (int i=odd&1;i<result.getLength();i+=2)
			result[i] = -result[i];
		return result;
	}
	
	template <class T> util::List<util::List<T> > Wavelet::orthfilt(const util::List<T>& filter)
	{
		util::List<util::List<T> > result(4);
		result.setLength(4);

		result[2] = filter / T(util::Math::sum(filter) / M_SQRT2);
		result[3] = qmf(result[2]);
		result[1] = result[3]; util::ListUtils::reverse(result[1]);
		result[0] = result[2]; util::ListUtils::reverse(result[0]);

		return result;
	}

	template <class T> util::Matrix<T> Wavelet::dyaddown(const util::Matrix<T>& mat,
																											 bool rows,
																											 bool columns,
																											 int odd)
	{
		int r = mat.getRows();
		int c = mat.getColumns();
		odd &= 1;
		util::Matrix<T> result(rows ? (r>>1) + (~odd&r&1) : r,
													 columns ? (c>>1) + (~odd&c&1) : c,
													 false);

		int cStep = columns ? 2 : 1;
		int rStep = (rows ? c : 0) + (columns && (c&1) ? (odd<<1)-1 : 0);

		T* resultData = result.getData();
		
		const T* sourceData = mat.getData() + odd * (c + 1);
		for (int i=result.getRows();i--;sourceData+=rStep)
			{
				for (int j=result.getColumns();j--;resultData++, sourceData+=cStep)
					*resultData = *sourceData;
			}

		return result;
	}
}}

#endif