// Matrix utilities

// [[Rcpp::depends(RcppEigen)]]

#include "matrix.h"
#include <RcppEigen.h>
using namespace Rcpp;
using namespace Eigen;

// find columns that exactly match previous columns
// returns numeric vector with -1 indicating no match to an earlier column and
//                             >0 indicating matches that earlier column
//                                (indexes starting at 1)
// [[Rcpp::export]]
NumericVector find_matching_cols(const NumericMatrix& mat, const double tol=1e-12)
{
    const int ncol = mat.cols();
    const int nrow = mat.rows();
    NumericVector result(ncol);

    if(ncol < 1) Rcpp::stop("Matrix has 0 columns");

    result[0] = -1;
    if(ncol==1) return(result);

    for(int i=1; i<ncol; i++) {
        result[i] = -1;
        for(int j=0; j<i; j++) {
            double max_diff=0.0;
            for(int k=0; k<nrow; k++) {
                const bool na_i = NumericVector::is_na(mat(k,i));
                const bool na_j = NumericVector::is_na(mat(k,j));
                // if both missing, return 0.0
                // if one missing but not other, return 1.0
                // otherwise, return difference
                double d = na_i != na_j ? 1.0 : ((na_i && na_j) ? 0.0 : fabs(mat(k,i) - mat(k,j)));
                if(d > max_diff) max_diff = d;
            }
            if(max_diff < tol) {
                result[i] = j+1;
                break;
            }
        }
    }

    return result;
}

// find set of linearly independent columns in a matrix
// returns a vector of column indices (starting at 1)
// [[Rcpp::export]]
IntegerVector find_lin_indep_cols(const NumericMatrix& mat, const double tol=1e-12)
{
    const int ncol=mat.cols();

    // QR decomp with column pivoting
    const MatrixXd XX(as<Map<MatrixXd> >(mat));
    typedef Eigen::ColPivHouseholderQR<MatrixXd> CPivQR;
    typedef CPivQR::PermutationType Permutation;
    CPivQR PQR ( XX );
    PQR.setThreshold(tol);

    // pivot matrix, treated as regular matrix
    Permutation Pmat = PQR.colsPermutation();
    MatrixXd PPmat(Pmat);

    // rank of input matrix
    const int rank=PQR.rank();
    IntegerVector result(rank);

    // for each column, find the row with a 1
    for(int j=0; j<rank; j++) {
        for(int i=0; i<ncol; i++) {
            if(fabs(PPmat(i,j) - 1.0) < tol) {
                result[j] = i+1;
                break;
            }
        }
    }

    return result;
}

// form X matrix with intcovar
// has_intercept = true indicates that addcovar has an intercept
//                 and so probs reduced by one column
//               = false means probs has full set of columns
//
// This is maybe a bit confusing.
//
// In the has_intercept=true case, an intercept is included in the
// addcovar matrix, and probs is missing the first column.
// The matrix formed is [A P (P.I)] where A=addcovar, P=probs, I=intcovar
//
// In the has_intercept=false case, no intercept is included in the
// addcovar matrix, and the probs have all columns (so each row sums
// to 1). The matrix formed is [P A (P*.I)] where in the P*.I bit we
// drop the first column of the probs when getting interactions with
// intercovar.
//
// [[Rcpp::export]]
NumericMatrix formX_intcovar(const NumericVector& probs,
                             const NumericMatrix& addcovar,
                             const NumericMatrix& intcovar,
                             const int position, // with indexes starting at 0
                             const bool has_intercept=true)
{

    if(Rf_isNull(probs.attr("dim")))
        throw std::invalid_argument("probs should be a matrix or 3d array but has no dim attribute");
    const Dimension d = probs.attr("dim");
    if(d.size() < 2)
        throw std::invalid_argument("probs should be a matrix or 3d array");
    if(position != 0 && d.size() != 3)
        throw std::invalid_argument("probs should be a 3d array");
    const int nrow  = d[0];
    const int ngen = d[1];
    const int recsize = nrow*ngen;
    const int offset = recsize*position;
    const int nadd  = addcovar.cols();
    const int nint  = intcovar.cols();

    int totcol;
    if(has_intercept) totcol = nadd + ngen*(nint+1);
    else totcol = ngen + nadd + (ngen-1)*nint;

    NumericMatrix result(nrow, totcol);

    if(position < 0 || (position != 0 && position >= d[2])) // if position == 0, don't look at d[2]
        throw std::range_error("position out of range of 0 .. (n_pos-1)");
    if(addcovar.rows() != nrow)
        throw std::range_error("nrow(addcovar) != nrow(probs)");
    if(intcovar.rows() != nrow)
        throw std::range_error("nrow(intcovar) != nrow(probs)");

    if(has_intercept) {
        std::copy(addcovar.begin(), addcovar.end(), result.begin());
        std::copy(probs.begin()+offset,
                  probs.begin()+offset+recsize,
                  result.begin() + nrow*nadd);

        for(int i=0, rescol=ngen+nadd; i<nint; i++) {
            for(int j=0; j<ngen; j++, rescol++) {
                for(int k=0; k<nrow; k++) {
                    result(k,rescol) = probs[offset + k + j*nrow] * intcovar(k,i);
                }
            }
        }
    }
    else {
        std::copy(probs.begin()+offset,
                  probs.begin()+offset+recsize,
                  result.begin());
        std::copy(addcovar.begin(), addcovar.end(), result.begin() + nrow*ngen);

        for(int i=0, rescol=ngen+nadd; i<nint; i++) {
            for(int j=1; j<ngen; j++, rescol++) {
                for(int k=0; k<nrow; k++) {
                    result(k,rescol) = probs[offset + k + j*nrow] * intcovar(k,i);
                }
            }
        }
    }
    return result;
}


// expand genotype probabilities with intcovar
// [[Rcpp::export]]
NumericVector expand_genoprobs_intcovar(const NumericVector& probs, // 3d array ind x prob x pos
                                        const NumericMatrix& intcovar)
{
    if(Rf_isNull(probs.attr("dim")))
        throw std::invalid_argument("probs should be a 3d array but has no dim attribute");
    Dimension d = probs.attr("dim");
    if(d.size() != 3)
        throw std::invalid_argument("probs should be a 3d array");
    const int nrow  = d[0];
    const int ngen = d[1];
    const int npos = d[2];
    const int nint  = intcovar.cols();

    if(intcovar.rows() != nrow)
        throw std::range_error("nrow(intcovar) != nrow(probs)");

    const int ngen_result = d[1]*(nint+1); // no. cols in result
    const int recsize = nrow*ngen; // ind x geno rectangle
    const int recsize_result = nrow*ngen_result; // ind x geno rectangle in result

    NumericVector result(recsize_result*npos);

    for(int i=0; i<npos; i++) {
        // paste probs into first batch
        std::copy(probs.begin()+i*recsize,
                  probs.begin()+(i+1)*recsize,
                  result.begin()+i*recsize_result);
        for(int j=0; j<nint; j++) {
            for(int k=0; k<ngen; k++) {
                for(int s=0; s<nrow; s++)
                    result[i*recsize_result + (j+1)*recsize + k*nrow + s] =
                        probs[i*recsize + k*nrow + s] * intcovar(s,j);
            }
        }
    }

    // add dimension attribute
    d[1] = ngen_result;
    result.attr("dim") = d;
    rownames(result) = rownames(probs);
    return result;
}

// multiply each column of a matrix by a set of weights
// [[Rcpp::export]]
NumericMatrix weighted_matrix(const NumericMatrix& mat,
                              const NumericVector& weights)
{
    const int nrow = mat.rows();
    const int ncol = mat.cols();
    if(nrow != weights.size())
        throw std::range_error("nrow(mat) != length(weights)");

    NumericMatrix result(nrow,ncol);

    for(int j=0; j<ncol; j++)
        for(int i=0; i<nrow; i++)
            result(i,j) = mat(i,j)*weights[i];

    return result;
}

// multiply each column of a 3-dimensional array by a set of weights
// [[Rcpp::export]]
NumericVector weighted_3darray(const NumericVector& array,
                               const NumericVector& weights)
{
    if(Rf_isNull(array.attr("dim")))
        throw std::invalid_argument("array should be a 3d array but has no dim attribute");
    const Dimension d = array.attr("dim");
    if(d.size() != 3)
        throw std::invalid_argument("array should be a 3d array");
    const int n = d[0];
    const int ncol = d[1]*d[2];
    if(n != weights.size())
        throw std::range_error("nrow(array) != length(weights)");

    NumericVector result(n*ncol);
    result.attr("dim") = d;

    for(int j=0, k=0; j<ncol; j++)
        for(int i=0; i<n; i++, k++)
            result[k] = array[k]*weights[i];

    return result;
}

// matrix multiplication
// [[Rcpp::export]]
NumericMatrix matrix_x_matrix(const NumericMatrix& X,
                              const NumericMatrix& Y)
{
    const MatrixXd XX(as<Map<MatrixXd> >(X));
    const MatrixXd YY(as<Map<MatrixXd> >(Y));

    if(XX.cols() != YY.rows())
        throw std::range_error("ncol(X) != nrow(Y)");

    NumericMatrix result(wrap(XX * YY));
    return(result);
}

// multiply matrix by vector
// [[Rcpp::export]]
NumericVector matrix_x_vector(const NumericMatrix& X,
                              const NumericVector& y)
{
    const MatrixXd XX(as<Map<MatrixXd> >(X));
    const VectorXd yy(as<Map<VectorXd> >(y));

    if(XX.cols() != yy.size())
        throw std::range_error("ncol(X) != length(y)");

    NumericVector result(wrap(XX * yy));
    return(result);
}

// multiply matrix by array
// [[Rcpp::export]]
NumericVector matrix_x_3darray(const NumericMatrix& X,
                               NumericVector& A)
{
    if(Rf_isNull(A.attr("dim")))
        throw std::invalid_argument("A should be a 3d array but has no dim attribute");
    Dimension d = A.attr("dim");
    if(d.size() != 3)
        throw std::invalid_argument("A should be 3-dimensional array");
    const int Xrow = X.rows();
    const int Xcol = X.cols();
    const int Arow = d[0];
    const int Acol = d[1];
    const int Apos = d[2];
    if(Xcol != Arow)
        throw std::invalid_argument("ncol(X) != nrow(A)");

    // treat as a matrix
    A.attr("dim") = Dimension(Arow, Acol*Apos);

    // cast for Eigen
    const MatrixXd XX(as<Map<MatrixXd> >(X));
    const MatrixXd AA(as<Map<MatrixXd> >(A));

    // matrix multiplication
    NumericVector result = wrap(XX*AA);
    result.attr("dim") = Dimension(Xrow, Acol, Apos);

    // fix dimension attribute of input matrix
    A.attr("dim") = d;

    return result;
}
