// Suppress stripping R API prefixes, all API functions have the form `Rf_<name>`
// and public variables are called `R_<name>`.
#define R_NO_REMAP
#include <R.h>
#include <Rinternals.h>

/**
 * Matrix Times a Vectorized Kronecker product
 *
 *      C = A vec(B_1 %x% ... %x% B_r)
 *
 * Note the reverse order of the `B_k` in the Kronecker product.
 *
 * @param A matrix of dimensions `n` by `pp * qq`
 * @param Bs list of matrices such that the product there Kronecker product has
 *  dimensions `pp` by `qq`.
 */
extern SEXP mtvk(SEXP A, SEXP Bs) {

    if ((TYPEOF(A) != REALSXP) || !Rf_isMatrix(A)) {
        Rf_error("First argument must be a numeric matrices");
    }

    // extract dimensions from `A` and `Bs`
    size_t nrow = Rf_nrows(A);
    size_t ncol = Rf_ncols(A);
    size_t r = Rf_length(Bs);
    size_t pp = 1;  // `nrow(Reduce("%x%", Bs))` see below
    size_t qq = 1;  // `ncol(Reduce("%x%", Bs))` see below

    // retrieve dimensions and direct memory access for each component in
    // reverse order
    // (Allocated memory freed at function exit and errors)
    double** bs = (double**)R_alloc(r, sizeof(double*));
    // dimensions of reshaped `A` columns, we treat the matrix `A` as an
    // `nrow` by `p[1]` by ... by `p[2 r]` tensor where the `p[1:r]` dimensions
    // as the row and `p[(r + 1):(2 r)]` the column dimensions of the Kronecker
    // product components.
    size_t* p = (size_t*)R_alloc(2 * r + 1, sizeof(size_t));
    // Multi-Index into the reshaped `A`:
    // `A[i, j] == do.call(array(A, dim = c(nrow, p), c(i, J)))`
    size_t* J = (size_t*)R_alloc(2 * r + 1, sizeof(size_t));

    // reversed order of Kronecker components
    for (size_t k = 0; k < r; ++k) {
        SEXP Bk = VECTOR_ELT(Bs, k);
        if ((TYPEOF(Bk) == REALSXP) && Rf_isMatrix(Bk)) {
            bs[r - 1 - k] = REAL(Bk);
            pp *= (p[    r - 1 - k] = Rf_nrows(Bk));
            qq *= (p[2 * r - 1 - k] = Rf_ncols(Bk));
            // set all zero
            J[k] = J[k + r] = 0;
        } else {
            Rf_error("Second argument must be a list of numeric matrices");
        }
    }

    // Additional last element such that the condition `p[2 * r] <= (++J[2 * r])`
    // (evaluated exactly once) is false. Allows to skip the check `k < 2 * r`
    // in the loop for updating the Multi-Index `J`.
    p[2 * r] = 127; // biggest value of smallest signed type
    J[2 * r] = 0;

    // check if dimensions match
    if ((ncol != pp * qq) || (r < 1)) {
        Rf_error("Dimension Missmatch");
    }

    // Create new R result vector
    SEXP C = PROTECT(Rf_allocVector(REALSXP, nrow));
    double* c = REAL(C);
    memset(c, 0, nrow * sizeof(double));

    // main operation (single iteration over `A`)
    double* a = REAL(A);
    for (size_t j = 0; j < ncol; ++j) {
        // Compute `prod_k (B_k)_{J_k, J_k+r}` (identical for each row)
        double prod_BJ = 1.0;
        for (size_t k = 0; k < r; ++k) {
            prod_BJ *= bs[k][J[k] + p[k] * J[k + r]];
        }

        // Add `j`th summand `A_i,J prod_k (B_k)_{J_k, J_k+r}` to each component
        for (size_t i = 0; i < nrow; ++i) {
            c[i] += a[i + j * nrow] * prod_BJ;
        }

        // Compute matching Multi-Index `J` for reshaped `A` columns
        // In `R` the relation between `j` and `J` is
        // `A[i, j] == do.call(array(A, dim = c(nrow, p), c(i, J)))`
        // Note: the check `k < 2 * r` is skipped as described above.
        for (size_t k = 0; p[k] <= (++J[k]); ++k) {
            J[k] = 0;
        }
    }

    // Onle the result object `C` needs to be unprotected
    UNPROTECT(1);

    return C;
}
