// Copyright Edward Rosten 2012

//All rights reserved.
//
//Redistribution and use in source and binary forms, with or without
//modification, are permitted provided that the following conditions
//are met:
//1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//2. Redistributions in binary form must reproduce the above copyright
//   notice, this list of conditions and the following disclaimer in the
//   documentation and/or other materials provided with the distribution.
//
//THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND OTHER CONTRIBUTORS ``AS IS''
//AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
//IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
//ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR OTHER CONTRIBUTORS BE
//LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
//CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
//SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
//INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
//CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
//ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
//POSSIBILITY OF SUCH DAMAGE.

#ifndef TOON_INCLUDE_QR_LAPACK_H
#define TOON_INCLUDE_QR_LAPACK_H


#include <TooN/TooN.h>
#include <TooN/lapack.h>
#include <utility>

namespace TooN{

/**
Performs %QR decomposition.

@warning this will only work if the number of columns is greater than 
the number of rows!

The QR decomposition operates on a matrix A. It can be performed with
or without column pivoting. In general:
\f[
AP = QR
\f]
Where \f$P\f$ is a permutation matrix constructed to permute the columns
of A. In practise, \f$P\f$ is stored as a vector of integer elements.

With column pivoting, the elements of the leading diagonal of \f$R\f$ will
be sorted from largest in magnitude to smallest in magnitude.

@ingroup gDecomps
*/
template<int Rows=Dynamic, int Cols=Rows, class Precision=double>
class QR_Lapack{

	private:
		static const int square_Size = (Rows>=0 && Cols>=0)?(Rows<Cols?Rows:Cols):Dynamic;

	public:	
		/// Construct the %QR decomposition of a matrix. This initialises the class, and
		/// performs the decomposition immediately.
		/// @param m The matrix to decompose
		/// @param p Whether or not to perform pivoting
		template<int R, int C, class P, class B> 
		QR_Lapack(const Matrix<R,C,P,B>& m, bool p=0)
		:copy(m),tau(square_size()), 
		 Q(square_size(), square_size()), 
		 do_pivoting(p), 
		 pivot(Zeros(m.num_cols()))
		{
			//pivot is set to all zeros, which means all columns are free columns
			//and can take part in column pivoting.

			compute();
		}
		
		///Return R
		const Matrix<Rows, Cols, Precision, ColMajor>& get_R()
		{
			return copy;
		}
		
		///Return Q
		const Matrix<square_Size, square_Size, Precision, ColMajor>& get_Q()
		{
			return Q;
		}	

		///Return the permutation vector. The definition is that column \f$i\f$ of A is
		///column \f$P(i)\f$ of \f$QR\f$.
		const Vector<Cols, int>& get_P()
		{
			return pivot;
		}

	private:

		void compute()
		{	
			FortranInteger M = copy.num_rows();
			FortranInteger N = copy.num_cols();
			
			FortranInteger LWORK=-1;
			FortranInteger INFO;
			FortranInteger lda = M;

			Precision size;
			
			//Set up the pivot vector
			if(do_pivoting)
				pivot = Zeros;
			else
				for(int i=0; i < pivot.size(); i++)
					pivot[i] = i+1;

			
			//Compute the working space
			geqp3_(&M, &N, copy.get_data_ptr(), &lda, pivot.get_data_ptr(), tau.get_data_ptr(), &size, &LWORK, &INFO);

			LWORK = (FortranInteger) size;

			Precision* work = new Precision[LWORK];
			
			geqp3_(&M, &N, copy.get_data_ptr(), &lda, pivot.get_data_ptr(), tau.get_data_ptr(), work, &LWORK, &INFO);


			if(INFO < 0)
				std::cerr << "error in QR, INFO was " << INFO << std::endl;

			//The upper "triangle+" of copy is R
			//The lower right and tau contain enough information to reconstruct Q
			
			//LAPACK provides a handy function to do the reconstruction
			Q = copy.template slice<0,0,square_Size, square_Size>(0,0,square_size(), square_size());
			
			FortranInteger K = square_size();
			M=K;
			N=K;
			lda = K;
			orgqr_(&M, &N, &K, Q.get_data_ptr(), &lda, tau.get_data_ptr(), work, &LWORK, &INFO);

			if(INFO < 0)
				std::cerr << "error in QR, INFO was " << INFO << std::endl;

			delete [] work;
			
			//Now zero out the lower triangle
			for(int r=1; r < square_size(); r++)
				for(int c=0; c<r; c++)
					copy[r][c] = 0;

			//Now fix the pivot matrix.
			//We need to go from FORTRAN to C numbering. 
			for(int i=0; i < pivot.size(); i++)
				pivot[i]--;
		}

		Matrix<Rows, Cols, Precision, ColMajor> copy;
		Vector<square_Size, Precision> tau;
		Matrix<square_Size, square_Size, Precision, ColMajor> Q;
		bool do_pivoting;
		Vector<Cols, FortranInteger> pivot;
		

		int square_size()
		{
			return std::min(copy.num_rows(), copy.num_cols());	
		}
};

}


#endif
