/****************************************************************************
*
*                                 M U E S L I   v 1.8
*
*
*     Copyright 2020 IMDEA Materials Institute, Getafe, Madrid, Spain
*     Contact: muesli.materials@imdea.org
*     Author: Ignacio Romero (ignacio.romero@imdea.org)
*
*     This file is part of MUESLI.
*
*     MUESLI 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.
*
*     MUESLI 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 MUESLI.  If not, see <http://www.gnu.org/licenses/>.
*
****************************************************************************/

#ifndef WITHEIGEN
#define NUM_PER_LINE   5         /* numbers per line, for screen output */

#include <ctype.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <algorithm>
#include <cassert>
#include <iostream>
#include <iomanip>
#include "muesli/Math/mmatrix.h"
#include "muesli/Math/mrealvector.h"


using namespace muesli;


matrix::matrix() :
rows(0),
cols(0),
data(0),
start(0)
{
}




/* 'data' is a pointer to a list of pointers, each of these pointing
 * a row in the matrix. For compactness, all the data is allocated
 * at once, in a big block and then the pointers suitably directed.
 */
matrix::matrix(const size_t rows_, const size_t cols_) :
rows(rows_),
cols(cols_)
{
    data  =  new double* [rows];
    start =  new double[rows*cols];
    
    // set the pointers in data to each row
    for (size_t i=0; i<rows; i++)
    {
        data[i] =  &(start[i*cols]) ;
        for (int j=0; j<cols; j++)  data[i][j] = 0.0;
    }
}




matrix::matrix(const matrix& n) :
rows(n.rows),
cols(n.cols)
{
    data  =  new double* [rows];
    start =  new double  [rows*cols];
    
    // set the pointers in data to each row
    for (size_t i=0; i<rows; i++)
    {
        data[i] =  &(start[i*cols]) ;
        for (int j=0; j<cols; j++)  data[i][j] = n.data[i][j];
    }
}




matrix::~matrix()
{
    if (start != 0)     delete [] start;
    if (data != 0)      delete [] data;
}




void matrix::chsgn()
{
    matrix& m=*this;
    for (size_t i=0; i<m.rows; i++)
    {
        for (size_t j=0; j<m.cols; j++)
        {
            m(i,j) = -m(i,j);
        }
    }
}




matrix matrix::IdentityMatrix(size_t n)
{
    matrix m(n,n);
    
    for (size_t i=0; i<n; i++) m.data[i][i] = 1.0;
    
    return m;
}



extern "C"
{
    void dgesv_(int*    , int*, double*, int*, int*, double*, int*, int*);
}

/* we invert the matrix m by solving m * X = identity, with the lapack function dgesv.
 the inverse overwrites the output  */
void  matrix::invert()
{
    int    i, j;
    matrix id;
    int    *p , info;
    double idet, tmp;
    double old[3][3];
    matrix& m = *this;
    
    size_t n      = m.rows;
    int    in     = static_cast<int>(n);
    double **M = m.data;
    
    switch (n)
    {
        case 0:
            break;
            
        case 1:
            m(0,0) = 1.0/m(0,0);
            break;
            
        case 2:
            idet    = 1.0/(m(0,0)*m(1,1) - m(0,1)*m(1,0));
            tmp     = m(0,0);
            m(0,0)  = idet * m(1,1);
            m(1,1)  = idet * tmp;
            m(0,1) *= -idet;
            m(1,0) *= -idet;
            break;
            
        case 3:
            idet = 1.0/m.determinant();
            
            if (idet == 0.0)
                printf("\n Error in invert matrix. Singular matrix");
            
            else
            {
                for (i=0;i<3;i++) for(j=0;j<3;j++) old[i][j] = M[i][j];
                
                M[0][0] = (old[1][1]*old[2][2]-old[1][2]*old[2][1])*idet;
                M[0][1] = (old[0][2]*old[2][1]-old[0][1]*old[2][2])*idet;
                M[0][2] = (old[0][1]*old[1][2]-old[0][2]*old[1][1])*idet;
                
                M[1][0] = (old[2][0]*old[1][2]-old[1][0]*old[2][2])*idet;
                M[1][1] = (old[0][0]*old[2][2]-old[0][2]*old[2][0])*idet;
                M[1][2] = (old[0][2]*old[1][0]-old[0][0]*old[1][2])*idet;
                
                M[2][0] = (old[1][0]*old[2][1]-old[1][1]*old[2][0])*idet;
                M[2][1] = (old[0][1]*old[2][0]-old[0][0]*old[2][1])*idet;
                M[2][2] = (old[0][0]*old[1][1]-old[0][1]*old[1][0])*idet;
            }
            break;
            
        default:
            id = IdentityMatrix(n);
            
            
            p  = (int *) malloc(n *sizeof(int));
            dgesv_(&in, &in, m.start, &in, p, id.start , &in , &info);
            
            if (info > 0) printf("\n Error in InvertMatrix: singular matrix.");
            
            free(p);
            m = id;
    }
}




/* LU decomposition for full matrices with no pivoting.
 * Replaces matrix m with LU,
 * unit Lower triangular and U, upper triangular. Of course, m is lost.
 */
bool matrix::factorLU()
{
    double aii;
    double **M;
    bool   ret(true);
    matrix& m=*this;
    
    size_t c = cols;
    size_t r = rows;
    
    if (c != r )
    {
        printf("ERROR in FullLU_NoPiv. Only square matrices!");
        ret = false;
    }
    
    else
    {
        M = m.data;
        for (size_t i=0; i<= r-2 ;i++)
        {
            aii = M[i][i];
            if (aii == 0.0)
            {
                printf("FullLU_NoPiv. Zero pivot number %lu.", i+1);
                ret = false;
                break;
            }
            aii = 1.0/aii;
            for (size_t j=i+1; j<r; j++) M[j][i] *= aii;
            for (size_t j=i+1; j<r; j++)
            {
                for (size_t k=i+1; k<r; k++)
                {
                    M[j][k] -= M[j][i]*M[i][k];
                }
            }
            ret = true;
        }
    }
    return ret;
}




double matrix::determinant() const
{
    double det=0.0;
    const matrix& m = *this;
    
    assert(rows == cols);
    
    if (rows != cols)
    {
        std::cout << "\n Error in matrix::determinant";
    }
    
    else if (rows == 1)
    {
        det = m(0,0);
    }
    
    else if (rows == 2)
    {
        det = m(0,0)*m(1,1)-m(0,1)+m(1,0);
    }
    
    else if (rows == 3)
    {
        det = m(0,0)*m(1,1)*m(2,2) + m(0,1)*m(1,2)*m(2,0) + m(1,0)*m(2,1)*m(0,2)
        -     m(0,2)*m(1,1)*m(2,0) - m(0,1)*m(1,0)*m(2,2) - m(0,0)*m(1,2)*m(2,1);
    }
    
    else
    {
        // for larger matrices, perform an LU decomposition, and then multiply the diagonal elements
        // of U
        matrix tmp = m;
        tmp.factorLU();
        det = 1.0;
        for (unsigned i=0; i<rows; i++)
            det *= tmp(i,i);
    }
    
    return det;
}




extern "C" {void dgeev_(char*, char*, int*, double*, int*, double*, double*,
                        double*, int*, double*, int*,
                        double*, int*,
                        int*);}



complexvector matrix::eigenvalues() const
{
    matrix evectors(this->rows, this->cols);
    complexvector eval(this->rows);
    eigendata(*this, evectors, eval.real(), eval.imag());

    return eval;
}



/* computes the eigenvalues and eigenvectors of a general full, unsymmetric matrix
 with a lapack routine */
bool  matrix::eigendata(const matrix& K, matrix& evectors, realvector& evaluesR, realvector& evaluesI)
{
    int   info , dummyint=1;
    char  N='N', V='V';
    double *work=NULL, *dummy=NULL;
    
    evaluesR.resize(K.rows);
    evaluesI.resize(K.rows);
    
    double* vevaluesR = new double[K.rows];
    double* vevaluesI = new double[K.rows];
    
    
    int lwork = (int) (5*K.rows);
    work = (double *) malloc( (size_t) lwork * sizeof(double));
    
    int r = (int) K.rows;
    dgeev_(&N, &V          , &r ,
           K.start        , &r ,
           vevaluesR      , vevaluesI,
           dummy          , &dummyint ,
           evectors.start , &r,
           work           , &lwork,
           &info);
    free(work);
    
    
    for (size_t a=0; a<K.rows; a++)
    {
        evaluesR[a] = vevaluesR[a];
        evaluesI[a] = vevaluesI[a];
    }
    
    delete [] vevaluesI;
    delete [] vevaluesR;

    return info == 0;
}




double  matrix::matrix::norm(const char type)
{
    int i,j;
    matrix& m=*this;
    double n=0, **M=NULL, sum;
    
    M = m.data;
    
    if (type == 'F' || type == 'f')
    {
        // Frobenius norm
        for (i=0; i < m.rows; i++)
            for (j=0; j < m.cols; j++)
                n += M[i][j] * M[i][j];
        n = sqrt(n);
    }
    
    // 1-norm, maximum column sum
    else if (type == '1')
    {
        n = 0.0;
        for (i=0; i < m.cols; i++)
        {
            sum = 0.0;
            for (j=0; j < m.rows; j++) sum += fabs(M[j][i]);
            n = std::max(n, sum);
        }
    }
    
    
    /* infinite-norm, maximum row sum */
    else if (type == 'i' || type == 'I')
    {
        n = 0.0;
        for (i=0; i < m.rows; i++)
        {
            sum = 0.0;
            for (j=0; j < m.cols; j++) sum += fabs(M[i][j]);
            n = std::max(n, sum);
        }
    }
    
    return n;
}




matrix& matrix::operator=(const matrix& m)
{
    this->resize(m.rows, m.cols);
    for (size_t i=0; i< rows; i++)
    {
        for (size_t j=0; j<cols; j++)
        {
            data[i][j] = m.data[i][j];
        }
    }
    
    return *this;
}




matrix& matrix::operator+=(const matrix &m)
{
    for (size_t i=0; i< rows; i++)
        for (size_t j=0; j<cols; j++)
            data[i][j] += m(i,j);
    
    return *this;
}




matrix& matrix::operator-=(const matrix &m)
{
    for (size_t i=0; i< rows; i++)
        for (size_t j=0; j<cols; j++)
            data[i][j] -= m(i,j);
    
    return *this;
}



namespace muesli
{

    matrix operator*(const matrix& m1, const matrix& m2)
    {
        assert(m1.cols == m2.rows);
        matrix m3(m1.rows, m2.cols);

        for (size_t i=0; i<m3.rows; i++)
        {
            for (size_t j=0; j<m3.cols; j++)
            {
                m3(i,j) = 0.0;
                for (size_t k=0; k<m1.cols; k++)
                {
                    m3(i,j) += m1(i,k) * m2(k,j);
                }
            }
        }
        return m3;
    }
    
    


    matrix operator+(const matrix& m1, const matrix& m2)
    {
        int r,c;
        double **M1, **M2, **M3;
        matrix m3;

        if ( (m1.cols != m2.cols) || (m1.rows != m2.rows) )
            printf("Error in AddMatrices. Matrices of different size");

        else
        {
            m3.resize(m1.rows, m2.cols);
            M1 = m1.data;
            M2 = m2.data;
            M3 = m3.data;
            for (r=0; r < m1.rows; r++)
                for (c=0; c < m1.cols; c++)
                    M3[r][c] = M2[r][c] + M1[r][c];
        }
        return m3;
    }



    matrix operator-(const matrix& m1, const matrix& m2)
    {
        int r,c;
        double **M1, **M2, **M3;
        matrix m3;

        if ( (m1.cols != m2.cols) || (m1.rows != m2.rows) )
            printf("Error in AddMatrices. Matrices of different size");

        else
        {
            m3.resize(m1.rows, m2.cols);
            M1 = m1.data;
            M2 = m2.data;
            M3 = m3.data;
            for (r=0; r < m1.rows; r++)
                for (c=0; c < m1.cols; c++)
                    M3[r][c] = M1[r][c] - M2[r][c];
        }
        return m3;
    }

}

realvector matrix::operator*(const realvector& b) const
{
    realvector ax(rows);

    for (size_t i=0; i< rows; i++)
    {
        ax[i] = 0.0;
        for (size_t j=0; j<cols; j++)
        {
            ax[i] += data[i][j]*b[j];
        }
    }
    return ax;
}




void matrix::print(std::ostream& of) const
{
    const matrix& m=*this;
    size_t c = m.cols;
    
    if (m.rows > 0 && c > 0)
    {
        of << "\n";
        
        // row loop
        for (size_t i=0; i< m.rows; i++)
        {
            of << "\n" << std::setw(3) << i+1 << "  ";
            size_t pstart = 0;
            size_t pend   = std::min<size_t>(c, NUM_PER_LINE);


            of << std::scientific << std::setprecision(6);
            // column loop
            for (size_t j=pstart; j<pend; j++) of << std::setw(10) << m.data[i][j] << " ";

            pstart += NUM_PER_LINE;
            pend    = std::min<size_t>( c , pend + NUM_PER_LINE);
            
            while (pstart <= c )
            {
                of << "... \n     ";
                for (size_t j=pstart; j<pend; j++) of << std::setw(10) << m.data[i][j] << " ";
                pstart += NUM_PER_LINE;
                pend    = std::min<size_t >(m.cols , pend + NUM_PER_LINE);
            }
        }
    }
    of << std::flush;
}




void matrix:: resize(const size_t newrows, const size_t newcols)
{
    matrix &m=*this;
    
    if (m.rows != newrows || m.cols != newcols)
    {
        if (m.start != 0) delete [] m.start;
        if (m.data  != 0) delete [] m.data;
        
        m.rows  = newrows;
        m.cols  = newcols;
        
        m.data  = new double* [rows];
        m.start = new double[rows*cols];
        
        // set the pointers in data to each row
        for (size_t i=0; i < newrows; i++)
            m.data[i] =  &(m.start[i*newcols]) ;
    }
    
    // in any case, set to Zero
    setZero();
}





void matrix::round() const
{
    const matrix& m=*this;
    size_t tot = m.rows * m.cols;
    for (size_t i=0; i<tot; i++) m.start[i] *= ( fabs(m.start[i]) < 1e-30 ) ? 0.0 : 1.0;
}



extern "C"
{
    void dgesv_(int *n, int* nrhs, double* A, int* lda, int* ipiv,
                double* b, int* ldb, int* info);
}




int matrix::solveFull(realvector& b)
{
    
    if (rows == 1)
    {
        b[0] = b[0]/ (*this)(0,0);
    }
    
    else if (rows == 2)
    {
        matrix& m = (*this);
        const double det = m(0,0)*m(1,1) - m(1,0)*m(0,1);
        
        matrix im(2,2);
        im(0,0) =  m(1,1)/det;
        im(1,1) =  m(0,0)/det;
        im(0,1) = -m(0,1)/det;
        im(1,0) = -m(1,0)/det;
        
        realvector c = b;
        b[0] = im(0,0)*c[0] + im(0,1)*c[1];
        b[1] = im(1,0)*c[0] + im(1,1)*c[1];
    }
    
    else
    {
        double* A = new double[rows*cols];
        for (size_t i = 0; i<rows; i++)
        {
            for (size_t j = 0; j<rows; j++)
            {
                A[rows*j+i] = (*this)(i,j);
            }
        }
        
        int  one = 1;
        int  irows = (int)rows;
        int* ipiv = new int[rows];
        double* rhs = new double[rows];
        
        for (unsigned i=0; i<rows; i++) rhs[i] = b[i];
        
        int info;
        
        dgesv_(&irows, &one, A, &irows, ipiv, rhs, &irows, &info);
        
        for (unsigned i=0; i<rows; i++) b[i] = rhs[i];
        
        if (info != 0) std::cout << "\n In matrix solver info = " << info;
        
        delete [] ipiv;
        delete [] rhs;
        delete [] A;
    }
    
    return 1;
}




void matrix::solveLU(realvector& b)
{
    matrix &LU=*this;
    
    /* forward substitution */
    size_t c = LU.cols;
    for (size_t k=1; k<c; k++)
    {
        for (size_t i=0; i<k; i++)
        {
            b[k] -= LU.data[k][i]*b[i];
        }
    }
    
    /* back substitution */
    b[c-1] /= LU.data[c-1][c-1];
    for (int ik= (int) (c-2); ik>=0; ik--)
    {
        size_t k = (size_t) ik;
        
        for (size_t i=k+1; i<c; i++)
        {
            b[k] -= LU.data[k][i]*b[i];
        }
        b[k] /= LU.data[k][k];
    }
}




void matrix::symmetrize()
{
    assert( rows == cols);
    size_t r = rows;
    
    for (size_t i=1; i<r; i++)
        for (size_t j=i; j<r; j++)
        {
            double tmp = 0.5*(data[i][j] + data[j][i]);
            data[i][j] = tmp;
            data[j][i] = tmp;
        }
}




void matrix::transpose()
{
    matrix copy(*this);
    size_t oldrows = rows;
    size_t oldcols = cols;
    
    this->resize(oldcols, oldrows);
    
    for (unsigned i=0; i<rows; i++)
        for (unsigned j=0; j<rows; j++)
        {
            data[i][j] = copy(j,i);
        }
}




void matrix::setZero()
{
    for (size_t a=0; a<rows*cols; a++) start[a] = 0.0;
}




extern "C"
{
    void dsygv_(int *itype, char *jobz, char *uplo, int *n,
                double *a, int *lda, double *b, int *ldb,
                double *w, double *work, int *lwork,
                int *info);
}




bool  matrix::generalizedEigendata(const matrix& K, const matrix& M, realvector& evals, matrix& evecs)
{
    assert( K.cols == K.rows);
    assert( M.cols == M.rows);
    assert( K.cols == M.cols);
    
    const size_t n = K.cols;
    evals.resize(n);
    evecs.resize(n, n);
    int in = (int) n;
    
    char jobz, uplo;
    int itype, lda, ldb, lwork, info;
    double *work;
    
    itype = 1; /* This sets the type of generalized eigenvalue problem
                that we are solving.  We have the possible values
                1: Ax = lBx
                2: ABx = lx
                3: BAx = lx */
    
    jobz = 'V'; /* V/N indicates that eigenvectors should/should not
                 be calculated. */
    
    uplo = 'L'; /* U/L indicated that the upper/lower triangle of the
                 symmetric matrix is stored. */
    
    lda = in; // The leading dimension of the matrix A
    ldb = in; // The leading dimension of the matrix B
    
    lwork = 3*in-1;
    work = new double[lwork]; /* The work array to be used by dsygv and
                               its size. */
    
    matrix Kv(K);
    matrix Mv(M);
    double* veva = new double[n];
    
    dsygv_(&itype, &jobz, &uplo, &in, Kv.start, &lda, Mv.start, &ldb, veva, work, &lwork, &info);
    evecs = Kv;
    
    for (size_t a=0; a<n; a++) evals[a] = veva[a];

    delete [] veva;
    delete [] work;

    return info == 0;
}




bool  matrix::generalizedEigendata(const matrix& K, const matrix& M, realvector& evals)
{
    assert( K.cols == K.rows);
    assert( M.cols == M.rows);
    assert( K.cols == M.cols);
    
    const size_t n = K.cols;
    evals.resize(n);
    
    char jobz, uplo;
    int itype, info;
    
    itype = 1; /* This sets the type of generalized eigenvalue problem
                that we are solving.  We have the possible values
                1: Ax = lBx
                2: ABx = lx
                3: BAx = lx */
    
    jobz = 'N'; /* V/N indicates that eigenvectors should/should not
                 be calculated. */
    
    uplo = 'L'; /* U/L indicated that the upper/lower triangle of the
                 symmetric matrix is stored. */
    
    int lda = (int) n; // The leading dimension of the matrix A
    int ldb = (int) n; // The leading dimension of the matrix B
    
    int lwork = 3*((int)n)-1;
    double* work = new double[lwork]; /* The work array to be used by dsygv and its size. */
    double* veva = new double[n];
    
    matrix Kv(K);
    matrix Mv(M);
    
    
    int in = (int) n;
    dsygv_(&itype, &jobz, &uplo, &in, Kv.start, &lda, Mv.start, &ldb, veva, work, &lwork, &info);
    for (size_t a=0; a<n; a++) evals[a] = veva[a];

    delete [] veva;
    delete [] work;

    return info == 0;
}

#endif

