#include "R_api.h"
#include "ttm.h"

int mlm(
    /*  options */  const int* trans, const int* modes, const int nrhs,
    /*     dims */  const int* dimA, const int* dimC, const int ord,
    /*  scalars */  const double alpha,
    /*   tensor */  const double* A,
    /* matrices */  const double** Bs, const int* ldBs,
    /*   scalar */  const double beta,
    /*   tensor */  double* C,
                    double* work_mem
) {
    // Compute total size of `A`, `C` and required working memory
    int sizeA = 1;  // `prod(dim(A))`
    int sizeC = 1;  // `prod(dim(C))`
    int work_size = 2;  // `2 * prod(pmax(dim(A), dim(C)))` (`+ ord` NOT included)
    for (int i = 0; i < ord; ++i) {
        sizeA *= dimA[i];
        sizeC *= dimC[i];
        work_size *= dimA[i] < dimC[i] ? dimC[i] : dimA[i];
    }

    // In requesting working memory size stop here and return size
    if (work_mem == NULL) {
        return work_size + ord;
    }

    // Copy dimensions of A to intermediate temp. dim
    int* dimT = (int*)(work_mem + work_size);   // `work_mem` is `ord` longer
    memcpy(dimT, dimA, ord * sizeof(int));

    // Setup work memory hooks (two swapable blocks)
    double* tmp1 = work_mem;
    double* tmp2 = work_mem + (work_size >> 1);

    // Multi Linear Multiplication is an iterated application of TTM
    for (int i = 0; i < nrhs; ++i) {
        // Get current `B` dimensions (only implicitly given)
        int nrowB = dimC[modes[i]];
        int ncolB = dimA[modes[i]];
        if (trans && trans[i]) {
            nrowB = dimA[modes[i]];
            ncolB = dimC[modes[i]];
        }

        // Tensor Times Matrix (`modes[i]` mode products)
        ttm(trans ? trans[i] : 0, modes[i], dimT, ord, nrowB, ncolB,
            1.0, i ? tmp2 : A, Bs[i], ldBs ? ldBs[i] : dimC[modes[i]], 0.0,
            tmp1);

        // Update intermediate temp dim
        dimT[modes[i]] = dimC[modes[i]];

        // Swap tmp1 <-> tmp2
        double* tmp3 = tmp1; tmp1 = tmp2; tmp2 = tmp3;
    }

    // dAXPY (a x + y) with `x = tmp2` and `y = beta C`
    if (beta != 0.0) {
        for (int i = 0; i < sizeC; ++i) {
            C[i] *= beta;
        }
    } else {
        memset(C, 0, sizeC * sizeof(double));
    }
    axpy(sizeC, alpha, tmp2, 1, C, 1);

    return 0;
}

/**
 * Multi Linear Multiplication (`R` binding)
 */
extern SEXP R_mlm(SEXP A, SEXP Bs, SEXP ms, SEXP ops) {

    // get dimension attribute of A
    SEXP dimA = Rf_getAttrib(A, R_DimSymbol);
    int ord = Rf_length(dimA);

    // Check if `A` is a real valued tensor
    if (!Rf_isReal(A) || Rf_isNull(dimA)) {
        Rf_error("Param. `A` need to be a real valued array");
    }

    // Validate that `Bs` is a list
    if (!Rf_isNewList(Bs)) {
        Rf_error("Param. `Bs` need to be a list of matrices");
    }
    if (!Rf_isInteger(ms) || !Rf_isLogical(ops)) {
        Rf_error("Param. type missmatch, expected modes as ints and ops logical");
    } 

    // Number of `B` matrices (Nr of Right Hand Side objects)
    int nrhs = Rf_length(Bs);

    // further parameter validations
    if ((nrhs != Rf_length(ms)) || (nrhs != Rf_length(ops))) {
        Rf_error("Dimension missmatch (Params length differ)");
    }

    // In case of an empty RHS (nr. of B's) is zero, this reduces to the identity
    if (nrhs == 0) {
        return A;
    }

    // Get modes and operations for Bs (transposed or not)
    int* modes = (int*)R_alloc(nrhs, sizeof(int)); // 0-indexed version of `ms`
    int* trans = INTEGER(ops);

    // Compute result dimensions `dim(C)` while extracting `B` data pointers
    SEXP dimC = PROTECT(Rf_duplicate(dimA));
    const double** bs = (const double**)R_alloc(nrhs, sizeof(double*));
    for (int i = 0; i < nrhs; ++i) {
        // Extract right hand side matrices (`B` matrices from list)
        SEXP B = VECTOR_ELT(Bs, i);
        if (!(Rf_isMatrix(B) && Rf_isReal(B))) {
            UNPROTECT(1);
            Rf_error("Param. `Bs` need to be a list of real matrices");
        }
        int* dimB = INTEGER(Rf_getAttrib(B, R_DimSymbol));
        bs[i] = REAL(B);

        // Convert 1-indexed modes to be 0-indexed
        modes[i] = INTEGER(ms)[i] - 1;

        // Check if mode is out or range (indexing a non-existing mode of `A`)
        if (!((0 <= modes[i]) && (modes[i] < ord))) {
            UNPROTECT(1);
            Rf_error("%d'th mode (%d) out of range", i + 1, modes[i] + 1);
        }

        // Check if `i`th mode of `A` matches corresponding `B` dimension
        if (INTEGER(dimA)[modes[i]] != dimB[!trans[i]]) {
            UNPROTECT(1);
            Rf_error("%d'th mode (%d) dimension missmatch", i + 1, modes[i] + 1);
        }

        INTEGER(dimC)[modes[i]] = dimB[!!trans[i]];
    }

    // Now, compute `C`s size `prod(dim(C))`
    int sizeC = 1;
    for (int i = 0; i < ord; ++i) {
        sizeC *= INTEGER(dimC)[i];
    }

    // Allocate response `C`
    SEXP C = PROTECT(Rf_allocVector(REALSXP, sizeC));
    Rf_setAttrib(C, R_DimSymbol, dimC);

    // allocate working memory size
    int work_size = mlm(trans, modes, nrhs, INTEGER(dimA), INTEGER(dimC), ord,
        1.0, REAL(A), bs, NULL, 0.0, REAL(C), NULL);
    double* work_memory = (double*)R_alloc(work_size, sizeof(double));

    // Compute Multi-Linear Multiplication (call subroutine)
    (void)mlm(trans, modes, nrhs, INTEGER(dimA), INTEGER(dimC), ord,
        1.0, REAL(A), bs, NULL, 0.0, REAL(C), work_memory);

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

    return C;
}