/*
    libmaus2
    Copyright (C) 2017-2020 German Tischler

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#if ! defined(LIBMAUS2_MATH_MATRIX_HPP)
#define LIBMAUS2_MATH_MATRIX_HPP

#include <libmaus2/autoarray/AutoArray.hpp>
#include <cmath>

namespace libmaus2
{
	namespace math
	{
		template<typename _N> struct MatrixBase
		{
		};
		template<> struct MatrixBase<double>
		{
			static double frexp(double const v, int * exp)
			{
				return ::std::frexp(v,exp);
			}

			static double ldexp(double const v, int const exp)
			{
				return ::std::ldexp(v,exp);
			}
		};

		template<typename _N> struct Matrix;
		template<typename N> std::ostream & operator<<(std::ostream & out, Matrix<N> const & M);

		template<typename _N> struct Matrix
		{
			typedef _N N;
			typedef Matrix<N> this_type;
			typedef std::unique_ptr<this_type> unique_ptr_type;
			typedef std::shared_ptr<this_type> shared_ptr_type;

			private:
			Matrix & operator=(Matrix const & M) const = delete;

			public:
			uint64_t r; // rows
			uint64_t c; // columns
			libmaus2::autoarray::AutoArray<N> A; // data

			N const & operator()(uint64_t const i, uint64_t const j) const
			{
				return A [ i * c + j ];
			}

			N & operator()(uint64_t const i, uint64_t const j)
			{
				return A [ i * c + j ];
			}

			Matrix()
			{
			}

			Matrix & operator=(Matrix && M)
			{
				r = M.r;
				c = M.c;
				A = M.A;
				return *this;
			}

			Matrix(uint64_t const rr, uint64_t const rc)
			: r(rr), c(rc), A(r*c)
			{
			}

			Matrix(Matrix<N> const & O)
			: r(O.r), c(O.c), A(r*c)
			{
				std::copy(O.A.begin(),O.A.end(),A.begin());
			}

			Matrix(Matrix<N> && O)
			: r(O.r), c(O.c), A(O.A)
			{

			}

			void swapRows(uint64_t const i, uint64_t const j)
			{
				for ( uint64_t k = 0; k < c; ++k )
					std::swap(
						(*this)(i,k),
						(*this)(j,k)
					);
			}

			void addRows(uint64_t const i, uint64_t const j, N const s)
			{
				for ( uint64_t k = 0; k < c; ++k )
					(*this)(i,k) += s * (*this)(j,k);
			}

			void multRow(uint64_t const i, N const s)
			{
				for ( uint64_t k = 0; k < c; ++k )
					(*this)(i,k) *= s;
			}

			// inverse matrix by Gauss elimination
			Matrix<N> inverse(N const t = N())
			{
				assert ( r == c );

				Matrix<N> U = *this;
				Matrix<N> I(r,c);

				for ( uint64_t i = 0; i < r; ++i )
					I(i,i) = N(1);

				for ( uint64_t i = 0; i < c; ++i )
				{
					N m = U(i,i);
					uint64_t mj = i;

					for ( uint64_t j = i; j < r; ++j )
						if ( abs(m) < abs(U(j,i)) )
						{
							m = U(j,i);
							mj = j;
						}

					if ( abs(m) <= t )
					{
						libmaus2::exception::LibMausException lme;
						lme.getStream() << "Matrix::inverse: Matrix is not invertible" << std::endl;
						lme.finish();
						throw lme;
					}

					U.swapRows(i,mj);
					I.swapRows(i,mj);

					for ( uint64_t j = i+1; j < r; ++j )
					{
						N const s = -U(j,i)/U(i,i);

						U.addRows(j,i,s);
						I.addRows(j,i,s);
					}
				}

				for ( uint64_t ii = 0; ii < c; ++ii )
				{
					uint64_t const i = c-ii-1;

					for ( uint64_t jj = ii+1; jj < c; ++jj )
					{
						uint64_t const j = c-jj-1;

						N const s = -U(j,i) / U(i,i);

						U.addRows(j,i,s);
						I.addRows(j,i,s);
					}
				}

				for ( uint64_t i = 0; i < c; ++i )
				{
					N const s = N(1) / U(i,i);

					U.multRow(i,s);
					I.multRow(i,s);
				}

				for ( uint64_t i = 0; i < r; ++i )
					for ( uint64_t j = 0; j < c; ++j )
						if ( abs(I(i,j)) <= t )
							I(i,j) = N();

				return I;
			}

			static N abs(N s)
			{
				if ( s < N() )
					return -s;
				else
					return s;
			}

			Matrix<N> & operator*=(N const & v)
			{
				for ( uint64_t i = 0; i < A.size(); ++i )
					A[i] *= v;

				return *this;
			}

			Matrix<N> & operator+=(Matrix<N> const & O)
			{
				assert ( r == O.r );
				assert ( c == O.c );

				for ( uint64_t i = 0; i < A.size(); ++i )
					A[i] += O.A[i];

				return *this;
			}

			Matrix<N> & operator-=(Matrix<N> const & O)
			{
				assert ( r == O.r );
				assert ( c == O.c );

				for ( uint64_t i = 0; i < A.size(); ++i )
					A[i] -= O.A[i];

				return *this;
			}

			static shared_ptr_type identity(uint64_t const n)
			{
				shared_ptr_type ptr(new this_type(n,n));
				for ( uint64_t i = 0; i < n; ++i )
					for ( uint64_t j = 0; j < n; ++j )
						(*ptr)(i,j) = (i == j);
				return ptr;
			}

			shared_ptr_type sclone() const
			{
				shared_ptr_type ptr(new this_type(*this));
				return ptr;
			}

			int getMaximumExponent() const
			{
				if ( r&&c )
				{
					int maxexp = std::numeric_limits<int>::min();
					for ( uint64_t i = 0; i < r*c; ++i )
					{
						int lexp;
						MatrixBase<N>::frexp(A[i],&lexp);
						maxexp = std::max(maxexp,lexp);
					}
					return maxexp;
				}
				else
					return 0;
			}
		};

		template<typename N> Matrix<N> operator*(Matrix<N> const & A, Matrix<N> const & B)
		{
			assert ( A.c == B.r );

			Matrix<N> R(A.r,B.c);

			for ( uint64_t i = 0; i < R.r; ++i )
				for ( uint64_t j = 0; j < R.c; ++j )
					for ( uint64_t k = 0; k < A.c; ++k )
						R(i,j) += A(i,k) * B(k,j);

			return R;
		}

		template<typename N> Matrix<N> multiply(Matrix<N> const & A, Matrix<N> const & B, uint64_t const threads)
		{
			assert ( A.c == B.r );

			Matrix<N> R(A.r,B.c);

			#if defined(_OPENMP)
			#pragma omp parallel for num_threads(threads)
			#endif
			for ( uint64_t i = 0; i < R.r; ++i )
			{
				for ( uint64_t j = 0; j < R.c; ++j )
				{
					N v = N();
					for ( uint64_t k = 0; k < A.c; ++k )
					{
						v += A(i,k) * B(k,j);
					}
					R(i,j) = v;
				}
			}

			return R;
		}

		template<typename N> std::pair<typename Matrix<N>::shared_ptr_type,uint64_t> power(libmaus2::math::Matrix<N> const & A, uint64_t const e, int const check = 64, uint64_t const threads = 1)
		{
			uint64_t const r = A.r;
			uint64_t const c = A.c;

			if ( e == 0 )
				return std::make_pair(libmaus2::math::Matrix<N>::identity(r),0);
			if ( e == 1 )
				return std::make_pair(A.sclone(),0);

			std::vector< typename libmaus2::math::Matrix<N>::shared_ptr_type > V;
			V.push_back(A.sclone());
			std::vector< int > E;
			E.push_back(0);

			uint64_t s = 1;
			int l = 1;
			while ( 2*s <= e )
			{
				typename libmaus2::math::Matrix<N>::shared_ptr_type ptr(new libmaus2::math::Matrix<N>);
				libmaus2::math::Matrix<N> const & back = *(V.back());

				if ( threads <= 1 )
					*ptr = back * back;
				else
					*ptr = multiply(back,back,threads);
				V.push_back(ptr);
				int const eback = E.back();

				int const lexp = ptr->getMaximumExponent();
				if ( lexp >= check )
				{
					int const frac = lexp / check;
					int const eadd = frac * check;
					for ( uint64_t i = 0; i < r; ++i )
						for ( uint64_t j = 0; j < c; ++j )
							(*ptr)(i,j) = MatrixBase<N>::ldexp((*ptr)(i,j),-eadd);
					E.push_back(2*eback+eadd);
				}
				else
				{
					E.push_back(2*eback);
				}

				s *= 2;
				l++;
			}

			assert ( 2*s > e );
			assert ( s == (1ull<<(l-1)) );

			typename libmaus2::math::Matrix<N>::shared_ptr_type ptr;
			bool init = false;
			uint64_t esum = 0;

			while ( l > 0 )
			{
				--l;

				if ( (e & (1ull<<l)) != 0 )
				{
					esum += E[l];

					if ( !init )
					{
						ptr = V[l];
						init = true;
					}
					else
					{
						if ( threads <= 1 )
							*ptr = *ptr * *V[l];
						else
							*ptr = multiply(*ptr,*V[l],threads);

						// get exponent of centre number
						int const lexp = ptr->getMaximumExponent();
						if ( lexp >= check )
						{
							int const frac = lexp / check;
							int const eadd = frac * check;
							esum += eadd;
							for ( uint64_t i = 0; i < r; ++i )
								for ( uint64_t j = 0; j < c; ++j )
									(*ptr)(i,j) = MatrixBase<N>::ldexp((*ptr)(i,j),-eadd);
						}
					}
				}
			}

			assert ( init );

			return std::make_pair(ptr,esum);
		}

		template<typename N> std::ostream & operator<<(std::ostream & out, Matrix<N> const & A)
		{
			for ( uint64_t i = 0; i < A.r; ++i )
			{
				for ( uint64_t j = 0; j < A.c; ++j )
				{
					out << A(i,j);

					if ( j+1 < A.c )
						out.put('\t');
				}

				out.put('\n');
			}

			return out;
		}
	}
}
#endif
