#include "R_api.h"
#include "ttm.h"
#include "mlm.h"
#include "solve.h"
#include "det.h"

/**
 * Generalized Multi Linear Model fitting for the Tensor Normal Distribution
 *
 * @todo TODO: doc this!
 */
void gmlm_tensor_normal(
    /*   options */ const int max_iter, const int max_line_iter,
    /*      dims */ const int* dimX, const int* dimF,
                    const int ord, const int sample_size,
    /*      data */ const double* X, const double* F,
    /*    params */ double* meanX, double** alphas, double** Omegas,
                    int* error
) {
    // product of X and F dimensions (dim of sample vec(X_i) and vec(F_i))
    int sizeX = 1;
    int sizeF = 1;
    for (int k = 0; k < ord; ++k) {
        sizeX *= dimX[k];
        sizeF *= dimF[k];
    }

    // Compute maximum temporary and working memory size for all called subroutines
    int work_size = mlm(FALSE, NULL, 0, dimX, dimF, ord, 0, NULL, NULL, NULL, 0, NULL, NULL);
    int tmp_size = work_size;
    // Update required working memory to also accomodate `solve` and `det`
    for (int k = 0, size = 0; k < ord; ++k) {
        solve(dimX[k], dimX[k], NULL, 0, NULL, 0, NULL, 0, NULL, &size);
        work_size = work_size < size ? size : work_size;

        (void)det(dimX[k], NULL, 0, NULL, &size);
        work_size = work_size < size ? size : work_size;
    }

    // allocate temporary and working memory
    double* work_mem = (double*)R_alloc(work_size, sizeof(double));
    double* tmp1_mem = (double*)R_alloc(tmp_size, sizeof(double));
    // tmp2_mem also needs to hold `sizeX^2` elements
    tmp_size = tmp_size < sizeX * sizeX ? sizeX * sizeX : tmp_size;
    double* tmp2_mem = (double*)R_alloc(tmp_size, sizeof(double));

    // initialize momentum and previous iteration terms for alphas and Omegas
    double** moment_alphas = (double**)R_alloc(ord, sizeof(double*));
    double** moment_Omegas = (double**)R_alloc(ord, sizeof(double*));
    double** prev_alphas = (double**)R_alloc(ord, sizeof(double*));
    double** prev_Omegas = (double**)R_alloc(ord, sizeof(double*));
    for (int k = 0; k < ord; ++k) {
        moment_alphas[k] = (double*)R_alloc(dimX[k] * dimF[k], sizeof(double));
        moment_Omegas[k] = (double*)R_alloc(dimX[k] * dimX[k], sizeof(double));
        prev_alphas[k] = (double*)R_alloc(dimX[k] * dimF[k], sizeof(double));
        prev_Omegas[k] = (double*)R_alloc(dimX[k] * dimX[k], sizeof(double));
        memcpy(moment_alphas[k], alphas[k], dimX[k] * dimF[k] * sizeof(double));
        memcpy(moment_Omegas[k], Omegas[k], dimX[k] * dimX[k] * sizeof(double));
        memcpy(prev_alphas[k], alphas[k], dimX[k] * dimF[k] * sizeof(double));
        memcpy(prev_Omegas[k], Omegas[k], dimX[k] * dimX[k] * sizeof(double));
    }

    // Allocate and compute initial values of `Deltas`, inverted `Omegas`
    //      Delta_k = Omega_k^-1 for k = 1, ..., ord
    double** Deltas = (double**)R_alloc(ord, sizeof(double*));
    for (int k = 0; k < ord; ++k) {
        Deltas[k] = (double*)R_alloc(dimX[k] * dimX[k], sizeof(double));
        solve(dimX[k], dimX[k], Omegas[k], dimX[k], NULL, 0,
            Deltas[k], dimX[k], work_mem, error);
    }

    // modes array for MLM
    int* modes = (int*)R_alloc(ord, sizeof(int));
    for (int k = 0; k < ord; ++k) { modes[k] = k; }

    ///// Allready computed by initial parameter estimate provided by calling R code
    // /* Step 1: compute mean(X) */
    // memset(meanX, 0, sizeX * sizeof(double));
    // for (int i = 0; i < sample_size; ++i) {
    //     F77_CALL(daxpy)(&sizeX, &d_one, X + i * sizeX, &i_one, meanX, &i_one);
    // }
    // const double inv_sample_size = 1.0 / (double)sample_size;
    // F77_CALL(dscal)(&sizeX, &inv_sample_size, meanX, &i_one);

    /* iteration 0, compute initial loss = log-likelihood */
    // TODO: figure out how to properly avoid the use of explicitly cast
    // see: https://stackoverflow.com/questions/12992407/warning-when-passing-non-const-parameter-to-a-function-that-expects-const-parame
    // see: https://c-faq.com/ansi/constmismatch.html
    double loss = 0.0;
    for (int i = 0; i < sample_size; ++i) {
        memcpy(tmp1_mem, meanX, sizeX * sizeof(double));    // tmp1_mem <- meanX
        (void)mlm(FALSE, modes, ord, dimF, dimX, ord,       // tmp1_mem <- mlm(F_i, alphas) + tmp1_mem
            1.0, F + i * sizeF, (const double**)alphas, dimX, 1.0, tmp1_mem, work_mem);
        (void)mlm(FALSE, modes, ord, dimX, dimX, ord,       // tmp1_mem <- mlm(tmp1_mem, Deltas) = mean(X | Y)
            1.0, tmp1_mem, (const double**)Deltas, dimX, 0.0, tmp1_mem, work_mem);
        // tmp1_mem <- tmp1_mem - X_i
        axpy(sizeX, -1.0, X + i * sizeX, 1, tmp1_mem, 1);
        // tmp2_mem <- mlm(tmp1_mem, Omegas)
        (void)mlm(FALSE, modes, ord, dimX, dimX, ord,
            1.0, tmp1_mem, (const double**)Omegas, dimX, 0.0, tmp2_mem, work_mem);
        // loss <- loss + <tmp1_mem, tmp2_mem>
        loss += dot(sizeX, tmp1_mem, 1, tmp2_mem, 1);
    }
    loss /= (double)(2 * sample_size);
    for (int k = 0; k < ord; ++k) {
        // loss <- loss - 1 / 2 (prod_{j != k} p_j) log(det(Omega_k))
        loss -= ((double)sizeX / (double)(2 * dimX[k]))
              * log(det(dimX[k], Omegas[k], dimX[k], work_mem, error));
    }

    /* NAGD parameters */
    int iter = 0;                   /* iteration counter */
    double m0 = 0.0, m1 = 1.0;      /* momentum extrapolation weights */
    double gamma = 0.61803398875;   /* line search step scaling */
    double step_size = 1e-2;        /* initial step size */

    /* Main parameter descent loop */
    for (iter = 0; iter < max_iter; ++iter) {

        /* momentum extrapolation */
        // `M <- theta_t + (m_{t-1} - 1) / m_t (theta_{t} - theta_{t-1})`
        for (int k = 0; k < ord; ++k) {
            lincomb(dimX[k] * dimF[k],
                1.0 + (m0 - 1.0) / m1, alphas[k], 1, (1.0 - m0) / m1, prev_alphas[k], 1,
                moment_alphas[k], 1);
            lincomb(dimX[k] * dimX[k],
                1.0 + (m0 - 1.0) / m1, Omegas[k], 1, (1.0 - m0) / m1, prev_Omegas[k], 1,
                moment_Omegas[k], 1);
        }

        /* Gradient at momentum extrapolated parameters */
        // Delta_k <- moment_Omega_k^-1, for k = 1, ..., ord
        for (int k = 0; k < ord; ++k) {
            solve(dimX[k], dimX[k], moment_Omegas[k], dimX[k], NULL, 0,
                Deltas[k], dimX[k], work_mem, error);
        }

        // 


        // line search (initialy try to increase the step size)
        step_size /= gamma;
        for (int line_iter = 0; line_iter < max_line_iter; ++line_iter) {
            /* gradient update for gradients at momentum extrapolated parameters */


            // iteration over observations
            for (int i = 0; i < sample_size; ++i) {
                // tmp1_mem <- meanX
                memcpy(tmp1_mem, meanX, sizeX * sizeof(double));
                // tmp1_mem <- mlm(F_i, moment_alphas) + tmp1_mem
                (void)mlm(FALSE, modes, ord, dimF, dimX, ord,
                    1.0, F + i * sizeF, (const double**)moment_alphas, dimX, 1.0, tmp1_mem, work_mem);
                // tmp1_mem <- mlm(tmp1_mem, Deltas)
                (void)mlm(FALSE, modes, ord, dimX, dimX, ord,
                    1.0, tmp1_mem, (const double**)Deltas, dimX, 0.0, tmp1_mem, work_mem);



                // TODO: continue




            }


            // decrease step size
            step_size *= gamma;
        }


    }





    for (int i = 0; i < ord; ++i) {
        memcpy(Omegas[i], Deltas[i], dimX[i] * dimX[i] * sizeof(double));
    }

    alphas[0][0] = loss;

}

/**
 * Generalized Multi Linear Model fitting for the Tensor Normal Distribution
 * (`R` binding)
 */
extern SEXP R_gmlm_tensor_normal(
    /*      data */ SEXP X, SEXP Fy,
    /*    params */ SEXP meanX, SEXP alphas, SEXP Omegas,
    /*   options */ SEXP max_iter, SEXP max_line_iter
) {
    // get dimension attribute of `X` and check if its a real valued tensor
    SEXP dimX = Rf_getAttrib(X, R_DimSymbol);
    if (!Rf_isReal(X) || Rf_isNull(dimX)) {
        Rf_error("Param. `X` need to be a real valued array");
    }
    // subtract 1 from `ord` as the last axis is the sample axis
    int ord = Rf_length(dimX) - 1;

    // the same for `Fy` and compare their order (nr. of axis)
    SEXP dimFy = Rf_getAttrib(Fy, R_DimSymbol);
    if (!Rf_isReal(Fy) || Rf_isNull(dimFy)) {
        Rf_error("Param. `Fy` need to be a real valued array");
    }
    if (ord + 1 != Rf_length(dimFy)) {
        Rf_error("Dimension mismatch, order of `X` and `Fy` must be equal");
    }

    // get direct access to dimension of `X` and `Fy`
    int* p = INTEGER(dimX);
    int* q = INTEGER(dimFy);

    // validate that `X` and `Fy` have same number of samples
    if (p[ord] != q[ord]) {
        Rf_error("Nr. samples in `X` is %d and in `Fy` is %d, they must be equal",
            p[ord], q[ord]);
    }
    int sample_size = p[ord];

    // Check if `alphas` and `Omegas` are lists of length `ord`
    if (!Rf_isNewList(alphas) || !Rf_isNewList(Omegas)
    || (Rf_length(alphas) != ord || Rf_length(Omegas) != ord)) {
        Rf_error("Params `alphas` and `Omegas` must be lists of length %d", ord);
    }

    // Get memory access to the meanX and validate its dimensions
    SEXP dimmeanX = Rf_getAttrib(meanX, R_DimSymbol);
    if (Rf_length(dimmeanX) != ord) {
        Rf_error("Dimension missmatch, meanX must have dim of `X` samples");
    }
    for (int i = 0; i < ord; ++i) {
        if (INTEGER(dimmeanX)[i] != p[i]) {
            Rf_error("Dimension missmatch, meanX must have dim of `X` samples");
        }
    }

    // Extract initial `alphas` and `Omegas` while checking their dimensions
    double** as = (double**)R_alloc(ord, sizeof(double*));
    double** Os = (double**)R_alloc(ord, sizeof(double*));
    for (int i = 0; i < ord; ++i) {
        // Extract alpha_i and Omega_i from `alphas` and `Omegas`
        SEXP alpha = VECTOR_ELT(alphas, i);
        SEXP Omega = VECTOR_ELT(Omegas, i);
        if (!(Rf_isMatrix(alpha) && Rf_isReal(alpha))
        ||  !(Rf_isMatrix(Omega) && Rf_isReal(Omega))) {
            Rf_error("Every alpha and Omega must be a real matrix");
        }

        // check dimensions
        int* dim = INTEGER(Rf_getAttrib(alpha, R_DimSymbol));
        if (dim[0] != p[i] || dim[1] != q[i]) {
            Rf_error("Dimension missmatch between dim(alpha_%d) = (%d, %d) "
                     "and dim(X)[%d] = %d and dim(Fy)[%d] = %d",
                i + 1, dim[0], dim[1], i + 1, p[i], i + 1, q[i]);
        }
        dim = INTEGER(Rf_getAttrib(Omega, R_DimSymbol));
        if (dim[0] != p[i] || dim[1] != p[i]) {
            Rf_error("Dimension missmatch between dim(Omega_%d) = (%d, %d) and dim(X)[%d] = %d",
                i + 1, dim[0], dim[1], i + 1, p[i]);
        }

        as[i] = REAL(alpha);
        Os[i] = REAL(Omega);
    }

    // Finally check algorithm options (configuration)
    if (!Rf_isInteger(max_iter) || Rf_length(max_iter) < 1
    ||  !Rf_isInteger(max_line_iter) || Rf_length(max_line_iter) < 1) {
        Rf_error("Unexpected max iterations, required to be pos int");
    }

    int error = 0;
    gmlm_tensor_normal(
        /*   options */ INTEGER(max_iter)[0], INTEGER0(max_line_iter)[0],
        /*      dims */ p, q, ord, sample_size,
        /*      data */ REAL(X), REAL(Fy),
        /*    params */ REAL(meanX), as, Os,
                        &error
    );

    // check if error occured which is indicated by negative number of iterations
    if (error) {
        Rf_error("Error in 'gmlm_tensor_normal' with error code %d", error);
    }

    return R_NilValue;
}
