
/* C Header */

/*
    Copyright (C) 2022- Torsten Hothorn

    This file is part of the 'mvtnorm' R add-on package.

    'mvtnorm' 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, version 2.

    'mvtnorm' 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 'mvtnorm'.  If not, see <http://www.gnu.org/licenses/>.


    DO NOT EDIT THIS FILE

    Edit 'lmvnorm_src.w' and run 'nuweb -r lmvnorm_src.w'
*/

#ifndef  USE_FC_LEN_T
# define USE_FC_LEN_T
#endif
#include <Rconfig.h>
#include <R_ext/Lapack.h> /* for dtptri */
#ifndef FCONE
# define FCONE
#endif
#include <R.h>
#include <Rmath.h>
#include <Rinternals.h>
#include <Rdefines.h>
/* colSumsdnorm */

SEXP R_ltMatrices_colSumsdnorm (SEXP z, SEXP N, SEXP J) {
    /* number of columns */
    int iN = INTEGER(N)[0];
    /* number of rows */
    int iJ = INTEGER(J)[0];
    SEXP ans;
    double *dans, Jl2pi, *dz;

    Jl2pi = iJ * log(2 * M_PI);
    PROTECT(ans = allocVector(REALSXP, iN));
    dans = REAL(ans);
    dz = REAL(z);

    for (int i = 0; i < iN; i++) {
        dans[i] = 0.0;
        for (int j = 0; j < iJ; j++)
            dans[i] += pow(dz[j], 2);
        dans[i] = - 0.5 * (Jl2pi + dans[i]);
        dz += iJ;
    }
    
    UNPROTECT(1);
    return(ans);
}

/* solve */

SEXP R_ltMatrices_solve (SEXP C, SEXP y, SEXP N, SEXP J, SEXP diag, SEXP transpose)
{

    SEXP ans;
    double *dans;
    int i, ONE = 1;

    /* RC input */
    
    /* pointer to C matrices */
    double *dC = REAL(C);
    /* number of matrices */
    int iN = INTEGER(N)[0];
    /* dimension of matrices */
    int iJ = INTEGER(J)[0];
    /* C contains diagonal elements */
    Rboolean Rdiag = asLogical(diag);
    /* p = J * (J - 1) / 2 + diag * J */
    int len = iJ * (iJ - 1) / 2 + Rdiag * iJ;
    
    /* diagonal elements are always present */
    if (!Rdiag) len += iJ;
    /* C length */
    
    int p;
    if (LENGTH(C) == len)
        /* C is constant for i = 1, ..., N */
        p = 0;
    else 
        /* C contains C_1, ...., C_N */
        p = len;
    
    /* lapack options */
    
    char di, uplo;
    if (Rdiag) {
        /* non-unit diagonal elements */
        di = 'N';
    } else {
        /* unit diagonal elements; NOTE: these diagonals 1s ARE always present but
           ignored in the computations */
        di = 'U';
    }

    SEXP abyrow = getAttrib(C, Rf_install("byrow"));
    if (TYPEOF(abyrow) != LGLSXP) error("non-logical byrow attribute");
    if (LOGICAL(abyrow)[0]) {
        uplo = 'U';
    } else {
        uplo = 'L';
    }

    


    char tr = 'N';
    Rboolean Rtranspose = asLogical(transpose);
    /* t(C) instead of C for column-major*/
    if (Rtranspose && !LOGICAL(abyrow)[0])
        tr = 'T';
    /* C instead for row-major*/
    if (!Rtranspose && LOGICAL(abyrow)[0])
        tr = 'T';

    PROTECT(ans = duplicate(y));
    dans = REAL(ans);
    
    /* loop over matrices, ie columns of C  / y */    
    for (i = 0; i < iN; i++) {

        /* solve linear system */
        F77_CALL(dtpsv)(&uplo, &tr, &di, &iJ, dC, dans, &ONE FCONE FCONE FCONE);
        dans += iJ;
        dC += p;
    }

    UNPROTECT(1);
    return(ans);
}

/* solve C */

SEXP R_ltMatrices_solve_C (SEXP C, SEXP N, SEXP J, SEXP diag, SEXP transpose)
{

    SEXP ans;
    int i, info;

    /* RC input */
    
    /* pointer to C matrices */
    double *dC = REAL(C);
    /* number of matrices */
    int iN = INTEGER(N)[0];
    /* dimension of matrices */
    int iJ = INTEGER(J)[0];
    /* C contains diagonal elements */
    Rboolean Rdiag = asLogical(diag);
    /* p = J * (J - 1) / 2 + diag * J */
    int len = iJ * (iJ - 1) / 2 + Rdiag * iJ;
    
    /* diagonal elements are always present */
    if (!Rdiag) len += iJ;
    /* lapack options */
    
    char di, uplo;
    if (Rdiag) {
        /* non-unit diagonal elements */
        di = 'N';
    } else {
        /* unit diagonal elements; NOTE: these diagonals 1s ARE always present but
           ignored in the computations */
        di = 'U';
    }

    SEXP abyrow = getAttrib(C, Rf_install("byrow"));
    if (TYPEOF(abyrow) != LGLSXP) error("non-logical byrow attribute");
    if (LOGICAL(abyrow)[0]) {
        uplo = 'U';
    } else {
        uplo = 'L';
    }

    

    PROTECT(ans = duplicate(C));
    /* dC is defined in RC input but otherwise not used */
    dC = REAL(ans);
    
    /* loop over matrices, ie columns of C  / y */    
    for (i = 0; i < iN; i++) {

        /* compute inverse */
        F77_CALL(dtptri)(&uplo, &di, &iJ, dC, &info FCONE FCONE);
        if (info != 0)
            error("Cannot solve ltMatrices");

        dC += len;
    }

    UNPROTECT(1);
    /* note: ans always includes diagonal elements */
    return(ans);
}

/* logdet */

SEXP R_ltMatrices_logdet (SEXP C, SEXP N, SEXP J, SEXP diag, SEXP byrow) {

    SEXP ans;
    double *dans;
    int i, j, k;

    /* RC input */
    
    /* pointer to C matrices */
    double *dC = REAL(C);
    /* number of matrices */
    int iN = INTEGER(N)[0];
    /* dimension of matrices */
    int iJ = INTEGER(J)[0];
    /* C contains diagonal elements */
    Rboolean Rdiag = asLogical(diag);
    /* p = J * (J - 1) / 2 + diag * J */
    int len = iJ * (iJ - 1) / 2 + Rdiag * iJ;
    
    Rboolean Rbyrow = asLogical(byrow);
    /* C length */
    
    int p;
    if (LENGTH(C) == len)
        /* C is constant for i = 1, ..., N */
        p = 0;
    else 
        /* C contains C_1, ...., C_N */
        p = len;
    

    PROTECT(ans = allocVector(REALSXP, iN));
    dans = REAL(ans);

    for (i = 0; i < iN; i++) {
        dans[i] = 0.0;
        if (Rdiag) {
            k = 1;
            for (j = 0; j < iJ; j++) {
                dans[i] += log(dC[k - 1]);
                k += (Rbyrow ? j + 2 : iJ - j);
            }
            dC += p;
        }
    }
    
    UNPROTECT(1);
    return(ans);
}

/* mult */

SEXP R_ltMatrices_Mult (SEXP C, SEXP y, SEXP N, SEXP J, SEXP diag, SEXP transpose) {

    SEXP ans;
    double *dans;
    int i, j, nry, lty, thisJ, incr;
    int iONE = 1;

    /* RC input */
    
    /* pointer to C matrices */
    double *dC = REAL(C);
    /* number of matrices */
    int iN = INTEGER(N)[0];
    /* dimension of matrices */
    int iJ = INTEGER(J)[0];
    /* C contains diagonal elements */
    Rboolean Rdiag = asLogical(diag);
    /* p = J * (J - 1) / 2 + diag * J */
    int len = iJ * (iJ - 1) / 2 + Rdiag * iJ;
    
    /* diagonal elements are always present */
    if (!Rdiag) len += iJ;
    /* C length */
    
    int p;
    if (LENGTH(C) == len)
        /* C is constant for i = 1, ..., N */
        p = 0;
    else 
        /* C contains C_1, ...., C_N */
        p = len;
    
    /* lapack options */
    
    char di, uplo;
    if (Rdiag) {
        /* non-unit diagonal elements */
        di = 'N';
    } else {
        /* unit diagonal elements; NOTE: these diagonals 1s ARE always present but
           ignored in the computations */
        di = 'U';
    }

    SEXP abyrow = getAttrib(C, Rf_install("byrow"));
    if (TYPEOF(abyrow) != LGLSXP) error("non-logical byrow attribute");
    if (LOGICAL(abyrow)[0]) {
        uplo = 'U';
    } else {
        uplo = 'L';
    }

    
    /* row column */
    
    char tr = 'N';
    Rboolean Rtranspose = asLogical(transpose);
    /* t(C) instead of C for column-major*/
    if (Rtranspose && !LOGICAL(abyrow)[0])
        tr = 'T';
    /* C instead for row-major*/
    if (!Rtranspose && LOGICAL(abyrow)[0])
        tr = 'T';
    

    PROTECT(ans = duplicate(y));
    dans = REAL(ans);

    /* mult workhorse */
    
    /* use case (1) */
    if (LENGTH(y) == iJ * iN) {
        nry = 1;
        lty = 0;
    /* use case (2) */
    } else if (LENGTH(y) == iJ * iJ * iN) {
        nry = iJ;
        lty = 0;
    /* use case (3) */
    } else if (LENGTH(y) == iJ * (iJ + 1) / 2 * iN) {
        if (!(Rtranspose && !LOGICAL(abyrow)[0]))
            error("matrix multiplication not implemented");
        nry = iJ;
        lty = 1;
    } else {
        error("incorrect dimension of y argument");
    }

    thisJ = iJ;
    for (i = 0; i < iN; i++) {
        incr = 0;
        for (j = 0; j < nry; j++) {       
            if (lty) thisJ = iJ - j;
            F77_CALL(dtpmv)(&uplo, &tr, &di, &thisJ,
                            dC + incr, dans, &iONE FCONE FCONE FCONE);
            dans += (lty ? iJ - j : iJ);
            if (lty) 
                incr += thisJ;
        }
        dC += p;
    }
    

    UNPROTECT(1);
    return(ans);
}

/* chol */

SEXP R_syMatrices_chol (SEXP Sigma, SEXP N, SEXP J) {

    SEXP ans;
    double *dans;
    int iJ = INTEGER(J)[0];
    int pJ = iJ * (iJ + 1) / 2;
    int iN = INTEGER(N)[0];
    int i, info = 0;
    char uplo = 'L';

    SEXP abyrow = getAttrib(Sigma, Rf_install("byrow"));
    if (TYPEOF(abyrow) != LGLSXP) error("non-logical byrow attribute");
    if (LOGICAL(abyrow)[0])
        uplo = 'U';

    /* duplicate preserves classes, so ltMatrices is returned */
    ans = PROTECT(duplicate(Sigma));
    dans = REAL(ans);

    for (i = 0; i < iN; i++) {

        F77_CALL(dpptrf)(&uplo, &iJ, dans, &info FCONE);

        if (info != 0) {
            if (info > 0)
                error("the leading minor of order %d is not positive definite",
                      info);
            error("argument %d of Lapack routine %s had invalid value",
                  -info, "dpptrf");
        }

        dans += pJ;
    }
    UNPROTECT(1);
    return(ans);
}

/* invchol workhorse */

void C_invchol (int J, double* ans, int* info) {
        
    int i, j, k;
    int start = 1, str;
    int end = 0;
    double sd = 0.0;
    double *sigma, *x;

    info[0] = 0;

    if (ans[0] < DBL_EPSILON) {
       info[0] = 1;
    } else {
        ans[0] = 1 / sqrt(ans[0]);

        for (j = 1; j < J; j++) {
            sigma = ans + start;
            sd = 0.0;
            /* compute L mS + its 2norm */
            /* <FIXME> use dtpmv <FIXME> */
            for (i = j - 1; i >= 0; i--) {
                sigma[i] *= ans[end];
                for (k = 1; k <= i; k++)
                    sigma[i] += ans[end - k] * sigma[i - k];
                sd += pow(sigma[i], 2);
                end -= i + 1;
            }
            /* compute lambda_jj^-1 */
            sd = sigma[j] - sd;
            if (sd < DBL_EPSILON) {
                info[0] = j;
                break;
            }
            sd = sqrt(sd);
            str = 0;
            /* compute L^\top (L mS) */
            /* <FIXME> use dtpmv <FIXME> */
            for (i = 0; i < j; i++) {
                x = ans + str + i;
                sigma[i] *= x[0]; 
                for (k = i + 1; k < j; k++) {
                    x += k;
                    sigma[i] += x[0] * sigma[k];
                }
                str += i + 1;
            }
            /* jth row of L */
            for (i = 0; i < j; i++)
                sigma[i] = - sigma[i] / sd;
            sigma[j] = 1 / sd;
            start += j + 1;
            end = start - 1;
        }
    }
}

/* invchol */

SEXP R_syMatrices_invchol (SEXP Sigma, SEXP N, SEXP J) {

    SEXP ans;
    double *dans;
    int iJ = INTEGER(J)[0];
    int pJ = iJ * (iJ + 1) / 2;
    int iN = INTEGER(N)[0];
    int i, info = 0;

    ans = PROTECT(duplicate(Sigma));
    dans = REAL(ans);

    for (i = 0; i < iN; i++) {

        C_invchol(iJ, dans, &info);

        if (info != 0) {
            if (info > 0)
                error("the leading minor of order %d is not positive definite",
                      info);
        }

        dans += pJ;
    }
    UNPROTECT(1);
    return(ans);
}

