#include "square_mat.h"

using namespace itpp;

using std::endl;

void fsqmat::opupdt ( const vec &v, double w ) {M+=outer_product ( v,v*w );};
mat fsqmat::to_mat() const {return M;};
void fsqmat::mult_sym ( const mat &C) {M=C *M*C.T();};
void fsqmat::mult_sym_t ( const mat &C) {M=C.T() *M*C;};
void fsqmat::mult_sym ( const mat &C, fsqmat &U) const { U.M = ( C *(M*C.T()) );};
void fsqmat::mult_sym_t ( const mat &C, fsqmat &U) const { U.M = ( C.T() *(M*C) );};
void fsqmat::inv ( fsqmat &Inv ) {mat IM = itpp::inv ( M ); Inv=IM;};
void fsqmat::clear() {M.clear();};
fsqmat::fsqmat ( const mat &M0 ) : sqmat(M0.cols())
{
	it_assert_debug ( ( M0.cols() ==M0.rows() ),"M0 must be square" );
	M=M0;
};

//fsqmat::fsqmat() {};

fsqmat::fsqmat(const int dim0): sqmat(dim0), M(dim0,dim0) {};

std::ostream &operator<< ( std::ostream &os, const fsqmat &ld ) {
	os << ld.M << endl;
	return os;
}


ldmat::ldmat( const mat &exL, const vec &exD ) : sqmat(exD.length()) {
	D = exD;
	L = exL;
}

ldmat::ldmat() :sqmat(0) {}

ldmat::ldmat(const int dim0): sqmat(dim0), D(dim0),L(dim0,dim0) {}

ldmat::ldmat(const vec D0):sqmat(D0.length()) {
	D = D0;
	L = eye(dim);
}

ldmat::ldmat( const mat &V ):sqmat(V.cols()) {
//TODO check if correct!! Based on heuristic observation of lu()

	it_assert_debug( dim == V.rows(),"ldmat::ldmat matrix V is not square!" );
		
	// L and D will be allocated by ldform()
	
	//Chol is unstable
	this->ldform(chol(V),ones(dim));
//	this->ldform(ul(V),ones(dim));
}

void ldmat::opupdt( const vec &v,  double w ) {
	int dim = D.length();
	double kr;
	vec r = v;
	//beware! it is potentionally dangerous, if ITpp change _behaviour of _data()!
	double *Lraw = L._data();
	double *Draw = D._data();
	double *rraw = r._data();

	it_assert_debug( v.length() == dim, "LD::ldupdt vector v is not compatible with this ld." );

	for ( int i = dim - 1; i >= 0; i-- ) {
		dydr( rraw, Lraw + i, &w, Draw + i, rraw + i, 0, i, &kr, 1, dim );
	}
}

std::ostream &operator<< ( std::ostream &os, const ldmat &ld ) {
	os << "L:" << ld.L << endl;
	os << "D:" << ld.D << endl;
	return os;
}

mat ldmat::to_mat() const {
	int dim = D.length();
	mat V( dim, dim );
	double sum;
	int r, c, cc;

	for ( r = 0;r < dim;r++ ) { //row cycle
		for ( c = r;c < dim;c++ ) {
			//column cycle, using symmetricity => c=r!
			sum = 0.0;
			for ( cc = c;cc < dim;cc++ ) { //cycle over the remaining part of the vector
				sum += L( cc, r ) * D( cc ) * L( cc, c );
				//here L(cc,r) = L(r,cc)';
			}
			V( r, c ) = sum;
			// symmetricity
			if ( r != c ) {V( c, r ) = sum;};
		}
	}
	mat V2 = L.transpose()*diag( D )*L;
	return V2;
}


void ldmat::add( const ldmat &ld2, double w ) {
	int dim = D.length();

	it_assert_debug( ld2.D.length() == dim, "LD.add() incompatible sizes of LDs;" );

	//Fixme can be done more efficiently either via dydr or ldform
	for ( int r = 0; r < dim; r++ ) {
		// Add columns of ld2.L' (i.e. rows of ld2.L) as dyads weighted by ld2.D
		this->opupdt( ld2.L.get_row( r ), w*ld2.D( r ) );
	}
}

void ldmat::clear(){L.clear(); for ( int i=0;i<L.cols();i++ ){L( i,i )=1;}; D.clear();}

void ldmat::inv( ldmat &Inv ) const {
	Inv.clear();   //Inv = zero in LD
	mat U = ltuinv( L );

	Inv.ldform( U.transpose(), 1.0 / D );
}

void ldmat::mult_sym( const mat &C) {
	mat A = L*C.T();
	this->ldform(A,D);
}

void ldmat::mult_sym_t( const mat &C) {
	mat A = L*C;
	this->ldform(A,D);
}

void ldmat::mult_sym( const mat &C, ldmat &U) const {
	mat A=L*C.T(); //could be done more efficiently using BLAS
	U.ldform(A,D);
}

void ldmat::mult_sym_t( const mat &C, ldmat &U) const {
	mat A=L*C;
/*	vec nD=zeros(U.rows());
	nD.replace_mid(0, D); //I case that D < nD*/
	U.ldform(A,D);
}


double ldmat::logdet() const {
	double ldet = 0.0;
	int i;
// sum logarithms of diagobal elements
	for ( i=0; i<D.length(); i++ ){ldet+=log( D( i ) );};
	return ldet;
}

double ldmat::qform( const vec &v ) const {
	double x = 0.0, sum;
	int i,j;

	for ( i=0; i<D.length(); i++ ) { //rows of L
		sum = 0.0;
		for ( j=0; j<=i; j++ ){sum+=L( i,j )*v( j );}
		x +=D( i )*sum*sum;
	};
	return x;
}

double ldmat::invqform( const vec &v ) const {
	double x = 0.0;
	int i;
	vec pom(v.length());
	
	backward_substitution(L.T(),v,pom);
	
	for ( i=0; i<D.length(); i++ ) { //rows of L
		x +=pom(i)*pom(i)/D(i);
	};
	return x;
}

ldmat& ldmat::operator *= ( double x ) {
	D*=x;
	return *this;
}

vec ldmat::sqrt_mult( const vec &x ) const {
	int i,j;
	vec res( dim );
	//double sum;
	for ( i=0;i<dim;i++ ) {//for each element of result
		res( i ) = 0.0;
		for ( j=i;j<dim;j++ ) {//sum D(j)*L(:,i).*x
			res( i ) += sqrt( D( j ) )*L( j,i )*x( j );
		}
	}
//	vec res2 = L.transpose()*diag( sqrt( D ) )*x;
	return res;
}

void ldmat::ldform(const mat &A,const vec &D0 )
{
	int m = A.rows();
	int n = A.cols();
	int mn = (m<n) ? m :n ;

//	it_assert_debug( A.cols()==dim,"ldmat::ldform A is not compatible" );
	it_assert_debug( D0.length()==A.rows(),"ldmat::ldform Vector D must have the length as row count of A" );

	L=concat_vertical( zeros( n,n ), diag( sqrt( D0 ) )*A );
	D=zeros( n+m );

	//unnecessary big L and D will be made smaller at the end of file	
	vec w=zeros( n+m );
   
	double sum, beta, pom;

	int cc=0;
	int i=n; // indexovani o 1 niz, nez v matlabu
	int ii,j,jj;
	while ( (i>n-mn-cc) && (i>0) ) 
	{
		i--;
		sum = 0.0;

		int last_v = m+i-n+cc+1;
	
		vec v = zeros( last_v + 1 ); //prepare v
		for ( ii=n-cc-1;ii<m+i+1;ii++ ) 
		{
			sum+= L( ii,i )*L( ii,i );
			v( ii-n+cc+1 )=L( ii,i ); //assign v
		}

		if ( L( m+i,i )==0 ) 
			beta = sqrt( sum );		
		else
			beta = L( m+i,i )+sign( L( m+i,i ) )*sqrt( sum );		
     
		if ( std::fabs( beta )<eps )
		{
			cc++;
			L.set_row( n-cc, L.get_row( m+i ) );
			L.set_row( m+i,zeros(L.cols()) );
			D( m+i )=0; L( m+i,i )=1;
			L.set_submatrix( n-cc,m+i-1,i,i,0 );
			continue;
		}

		sum-=v(last_v)*v(last_v);
		sum/=beta*beta;
		sum++;

		v/=beta;
		v(last_v)=1;

		pom=-2.0/sum;
		// echo to venca   

		for ( j=i;j>=0;j-- ) 
		{
			double w_elem = 0;			
			for ( ii=n-	cc;ii<=m+i+1;ii++ ) 
				w_elem+= v( ii-n+cc )*L( ii-1,j );			
			w(j)=w_elem*pom;
		}

		for ( ii=n-cc-1;ii<=m+i;ii++ ) 
			for ( jj=0;jj<i;jj++ ) 
				L( ii,jj )+= v( ii-n+cc+1)*w( jj );

		for ( ii=n-cc-1;ii<m+i;ii++ )
			L( ii,i )= 0;

		L( m+i,i )+=w( i );
		D( m+i )=L( m+i,i )*L( m+i,i );

		for ( ii=0;ii<=i;ii++ ) 
			L( m+i,ii )/=L( m+i,i );		
	}

	if ( i>=0 )
		for ( ii=0;ii<i;ii++ ) 
		{
			jj = D.length()-1-n+ii;
			D(jj) = 0;
			L.set_row(jj,zeros(L.cols())); //TODO: set_row accepts Num_T
			L(jj,jj)=1;
		}

	L.del_rows(0,m-1);
	D.del(0,m-1);
	
	dim = L.rows();
}

//////// Auxiliary Functions

mat ltuinv( const mat &L ) {
	int dim = L.cols();
	mat Il = eye( dim );
	int i, j, k, m;
	double s;

//Fixme blind transcription of ltuinv.m
	for ( k = 1; k < ( dim );k++ ) {
		for ( i = 0; i < ( dim - k );i++ ) {
			j = i + k; //change in .m 1+1=2, here 0+0+1=1
			s = L( j, i );
			for ( m = i + 1; m < ( j ); m++ ) {
				s += L( m, i ) * Il( j, m );
			}
			Il( j, i ) = -s;
		}
	}

	return Il;
}

void dydr( double * r, double *f, double *Dr, double *Df, double *R, int jl, int jh, double *kr, int m, int mx )
/********************************************************************

   dydr = dyadic reduction, performs transformation of sum of
          2 dyads r*Dr*r'+ f*Df*f' so that the element of r pointed
          by R is zeroed. This version allows Dr to be NEGATIVE. Hence the name negdydr or dydr_withneg.

   Parameters :
     r ... pointer to reduced dyad
     f ... pointer to reducing dyad
     Dr .. pointer to the weight of reduced dyad
     Df .. pointer to the weight of reducing dyad
     R ... pointer to the element of r, which is to be reduced to
           zero; the corresponding element of f is assumed to be 1.
     jl .. lower index of the range within which the dyads are
           modified
     ju .. upper index of the range within which the dyads are
           modified
     kr .. pointer to the coefficient used in the transformation of r
           rnew = r + kr*f
     m  .. number of rows of modified matrix (part of which is r)
  Remark : Constant mzero means machine zero and should be modified
           according to the precision of particular machine

                                                 V. Peterka 17-7-89

  Added:
     mx .. number of rows of modified matrix (part of which is f)  -PN

********************************************************************/
{
	int j, jm;
	double kD, r0;
	double mzero = 2.2e-16;
	double threshold = 1e-4;

	if ( fabs( *Dr ) < mzero ) *Dr = 0;
	r0 = *R;
	*R = 0.0;
	kD = *Df;
	*kr = r0 * *Dr;
	*Df = kD + r0 * ( *kr );
	if ( *Df > mzero ) {
		kD /= *Df;
		*kr /= *Df;
	} else {
		kD = 1.0;
		*kr = 0.0;
		if ( *Df < -threshold ) {
			it_warning( "Problem in dydr: subraction of dyad results in negative definitness. Likely mistake in calling function." );
		}
		*Df = 0.0;
	}
	*Dr *= kD;
	jm = mx * jl;
	for ( j = m * jl; j < m*jh; j += m ) {
		r[j] -=  r0 * f[jm];
		f[jm] += *kr * r[j];
		jm += mx;
	}
}