#include "square_mat.h"

namespace bdm {

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 ) const {
    mat IM = itpp::inv ( M );
    Inv = IM;
};
void fsqmat::clear() {
    M.clear();
};
fsqmat::fsqmat ( const mat &M0 ) : sqmat ( M0.cols() ) {
    bdm_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() ) {

    bdm_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 ) );
}

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();

    bdm_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 V;
}


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

    bdm_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 ;

    bdm_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 ) {
            bdm_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;
    }
}

}