// The need for `USE_FC_LEN_T` and `FCONE` is due to a Fortran character string
// to C incompatibility. See: Writing R Extentions: 6.6.1 Fortran character strings
#define USE_FC_LEN_T
#include <R.h>
#include <Rinternals.h>
#include <R_ext/BLAS.h>
#ifndef FCONE
#define FCONE
#endif

/**
 * Tensor Times Matrix a.k.a. Mode Product
 *
 * @param A multi-dimensionl array
 * @param B matrix
 * @param m mode index (1-indexed)
 */
extern SEXP ttm(SEXP A, SEXP B, SEXP m) {

    // get zero indexed mode
    int mode = asInteger(m) - 1;

    // get dimension attribute of A
    SEXP dim = getAttrib(A, R_DimSymbol);

    // validate mode (mode must be smaller than the nr of dimensions)
    if (mode < 0 || length(dim) <= mode) {
        error("Illegal mode");
    }

    // and check if B is a matrix of non degenetate size
    if (!isMatrix(B)) {
        error("Expected a matrix as second argument");
    }
    if (!nrows(B) || !ncols(B)) {
        error("Zero dimension detected");
    }

    // check matching of dimensions
    if (INTEGER(dim)[mode] != ncols(B)) {
        error("Dimension missmatch (mode dim not equal to ncol)");
    }

    // calc nr of response elements `prod(dim(X)[-mode]) * ncol(X)`,
    int leny = 1;
    // and the strides
    //  `stride[0] <- prod(dim(X)[seq_len(mode - 1)])`
    //  `stride[1] <- dim(X)[mode]`
    //  `stride[2] <- prod(dim(X)[-seq_len(mode)])`
    int stride[3] = {1, INTEGER(dim)[mode], 1};
    for (int i = 0; i < length(dim); ++i) {
        int size = INTEGER(dim)[i];
        // check for non-degenetate dimensions
        if (!size) {
            error("Zero dimension detected");
        }
        leny *= (i == mode) ? nrows(B) : size;
        stride[0] *= (i < mode) ? size : 1;
        stride[2] *= (i > mode) ? size : 1;
    }

    // as well as the matrix rows (cols only needed once)
    int nrow = nrows(B);

    // create response object C
    SEXP C = PROTECT(allocVector(REALSXP, leny));

    // raw data access pointers
    double* a = REAL(A);
    double* b = REAL(B);
    double* c = REAL(C);

    // Tensor Times Matrix / Mode Product
    const double zero = 0.0;
    const double one  = 1.0;
    if (mode == 0) {
        // mode 1: (A x_1 B)_(1) = B A_(1) as a single Matrix-Matrix prod
        F77_CALL(dgemm)("N", "N", &nrow, &stride[2], &stride[1], &one,
            b, &nrow, a, &stride[1],
            &zero, c, &nrow FCONE FCONE);
    } else {
        // Other modes can be written as blocks of matrix multiplications
        for (int i2 = 0; i2 < stride[2]; ++i2) {
            F77_CALL(dgemm)("N", "T", &stride[0], &nrow, &stride[1], &one,
                &a[i2 * stride[0] * stride[1]], &stride[0], b, &nrow,
                &zero, &c[i2 * stride[0] * nrow], &stride[0] FCONE FCONE);
        }
    }
    /*
    // Tensor Times Matrix / Mode Product (reference implementation)
    memset(c, 0, leny * sizeof(double));
    for (int i2 = 0; i2 < stride[2]; ++i2) {
        for (int i1 = 0; i1 < stride[1]; ++i1) { // stride[1] == ncols(B)
            for (int j = 0; j < nrow; ++j) {
                for (int i0 = 0; i0 < stride[0]; ++i0) {
                    c[i0 + (j + i2 * nrow) * stride[0]] +=
                        a[i0 + (i1 + i2 * stride[1]) * stride[0]] * b[j + i1 * nrow];
                }
            }
        }
    }
    */

    // finally, set result dimensions
    SEXP newdim = PROTECT(allocVector(INTSXP, length(dim)));
    for (int i = 0; i < length(dim); ++i) {
        INTEGER(newdim)[i] = (i == mode) ? nrows(B) : INTEGER(dim)[i];
    }
    setAttrib(C, R_DimSymbol, newdim);

    // release C to the hands of the garbage collector
    UNPROTECT(2);

    return C;
}