/**

    Linear statistics for conditional inference based on Strasser & Weber (1999)

    *\file LinearStatistic.c

    *\author $Author$

    *\date $Date$

*/

#include "party.h"

/**

    Computes the linear statistic, formula (1) in the paper\n

    *\param x values of the transformation

    *\param p dimension of the transformation

    *\param y values of the influence function

    *\param q dimension of the influence function

    *\param weights case weights

    *\param n number of observations

    *\param ans return value; a pointer to a REALSXP-vector of length pq

*/

void C_LinearStatistic (const double *x, const int p,

                        const double *y, const int q,

                        const double *weights, const int n,

                        double *ans) {

    int i, j, k, kp, kn;

    double tmp;

    for (k = 0; k < q; k++) {

        kn = k * n;

        kp = k * p;

        for (j = 0; j < p; j++) ans[kp + j] = 0.0;

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

            /* optimization: weights are often zero */

            if (weights[i] == 0.0) continue;

            tmp = y[kn + i] * weights[i];

            for (j = 0; j < p; j++)

                 ans[kp + j] += x[j*n + i] * tmp;

        }

    }

}

/**

    R-interface to C_LinearStatistic \n

    *\param x values of the transformation

    *\param y values of the influence function

    *\param weights case weights

*/

SEXP R_LinearStatistic(SEXP x, SEXP y, SEXP weights) {

    /* the return value; a vector of type REALSXP */

    SEXP ans;

    /* dimensions */

    int n, p, q;

    /* 

     *    only a basic check: we do not coerce objects since this

     *    function is for internal use only

     */

    if (!isReal(x) || !isReal(y) || !isReal(weights))

        error("LinStat: arguments are not of type REALSXP");

    n = nrow(y);

    if (nrow(x) != n || LENGTH(weights) != n)

        error("LinStat: dimensions don't match");

    q    = ncol(y);

    p    = ncol(x);

    PROTECT(ans = allocVector(REALSXP, p*q));

    C_LinearStatistic(REAL(x), p, REAL(y), q, REAL(weights), n, 

                      REAL(ans));

    UNPROTECT(1);

    return(ans);

}

/**

    Conditional expectation and covariance of the influence function\n

    *\param y values of the influence function

    *\param q dimension of the influence function

    *\param weights case weights

    *\param n number of observations

    *\param ans return value; an object of class `ExpectCovarInfluence'

*/

void C_ExpectCovarInfluence(const double* y, const int q,

                            const double* weights, const int n, 

                            SEXP ans) {

    int i, j, k, jq;

    /* pointers to the slots of object ans */

    double *dExp_y, *dCov_y, *dsweights, tmp;

    /*  return values: set to zero initially */

    dExp_y = REAL(GET_SLOT(ans, PL2_expectationSym));

    for (j = 0; j < q; j++) dExp_y[j] = 0.0;

    dCov_y = REAL(GET_SLOT(ans, PL2_covarianceSym));

    for (j = 0; j < q*q; j++) dCov_y[j] = 0.0;

    dsweights = REAL(GET_SLOT(ans, PL2_sumweightsSym));

    /*  compute the sum of the weights */

    dsweights[0] = 0;

    for (i = 0; i < n; i++) dsweights[0] += weights[i];

    if (dsweights[0] <= 1) 

        error("C_ExpectCovarInfluence: sum of weights is less than one");

    /*

     *    Expectation of the influence function

     */

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

        /*  observations with zero case weights do not contribute */

        if (weights[i] == 0.0) continue;

        for (j = 0; j < q; j++)

            dExp_y[j] += weights[i] * y[j * n + i];

    }

    for (j = 0; j < q; j++)

        dExp_y[j] = dExp_y[j] / dsweights[0];

    /*

     *    Covariance of the influence function

     */

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

        if (weights[i] == 0.0) continue;

        for (j = 0; j < q; j++) {

            tmp = weights[i] * (y[j * n + i] - dExp_y[j]);

            jq = j * q;

            for (k = 0; k < q; k++)

                dCov_y[jq + k] += tmp * (y[k * n + i] - dExp_y[k]);

        }

    }

    for (j = 0; j < q*q; j++)

        dCov_y[j] = dCov_y[j] / dsweights[0];

}

/**

    R-interface to C_ExpectCovarInfluence\n

    *\param y values of the influence function

    *\param weights case weights

*/

SEXP R_ExpectCovarInfluence(SEXP y, SEXP weights) {

    SEXP ans;

    int q, n;

    if (!isReal(y) || !isReal(weights))

        error("R_ExpectCovarInfluence: arguments are not of type REALSXP");

    n = nrow(y);

    q = ncol(y);

    if (LENGTH(weights) != n) 

        error("R_ExpectCovarInfluence: vector of case weights does not have %d elements", n);

    /*  allocate storage for return values */

    PROTECT(ans = NEW_OBJECT(MAKE_CLASS("ExpectCovarInfluence")));

    SET_SLOT(ans, PL2_expectationSym, 

             PROTECT(allocVector(REALSXP, q)));

    SET_SLOT(ans, PL2_covarianceSym, 

             PROTECT(allocMatrix(REALSXP, q, q)));

    SET_SLOT(ans, PL2_sumweightsSym, 

             PROTECT(allocVector(REALSXP, 1)));

    C_ExpectCovarInfluence(REAL(y), q, REAL(weights), n, ans);

    UNPROTECT(4);

    return(ans);

}

/**

    Conditional expectation and covariance of the a linear statistic\n

    *\param x values of the transformation

    *\param p dimension of the transformation

    *\param y values of the influence function

    *\param q dimension of the influence function

    *\param weights case weights

    *\param n number of observations

    *\param expcovinf an object of class `ExpectCovarInfluence'

    *\param ans return value; an object of class `ExpectCovar'

*/

void C_ExpectCovarLinearStatistic(const double* x, const int p, 

                                  const double* y, const int q,

                                  const double* weights, const int n,

                                  const SEXP expcovinf, SEXP ans) {

    int i, j, k, pq;

    double sweights = 0.0, f1, f2, tmp;

    double *swx, *CT1, *CT2, *Covy_x_swx, 

           *dExp_y, *dCov_y, *dExp_T, *dCov_T;

    pq   = p * q;

    /* the expectation and covariance of the influence function */

    dExp_y = REAL(GET_SLOT(expcovinf, PL2_expectationSym));

    dCov_y = REAL(GET_SLOT(expcovinf, PL2_covarianceSym));

    sweights = REAL(GET_SLOT(expcovinf, PL2_sumweightsSym))[0];

    if (sweights <= 1.0) 

        error("C_ExpectCovarLinearStatistic: sum of weights is less than one");

    /* prepare for storing the results */

    dExp_T = REAL(GET_SLOT(ans, PL2_expectationSym));

    dCov_T = REAL(GET_SLOT(ans, PL2_covarianceSym));

    /* allocate storage: all helpers, initially zero */

    swx = Calloc(p, double);               /* p x 1  */

    CT1 = Calloc(p * p, double);           /* p x p  */

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

        /*  observations with zero case weights do not contribute */

        if (weights[i] == 0.0) continue;

        for (k = 0; k < p; k++) {

            tmp = weights[i] * x[k * n + i];

            swx[k] += tmp;

            /* covariance part */

            for (j = 0; j < p; j++) {

                CT1[j * p + k] += tmp * x[j * n + i];

            }

        }

    }

    /*

    *   dExp_T: expectation of the linear statistic T

    */

    for (k = 0; k < p; k++) {

        for (j = 0; j < q; j++)

            dExp_T[j * p + k] = swx[k] * dExp_y[j];

    }

    /* 

    *   dCov_T:  covariance of the linear statistic T

    */

    f1 = sweights/(sweights - 1);

    f2 = (1/(sweights - 1));

    if (pq == 1) {

        dCov_T[0] = f1 * dCov_y[0] * CT1[0];

        dCov_T[0] -= f2 * dCov_y[0] * swx[0] * swx[0];

    } else {

        /* two more helpers needed */

        CT2 = Calloc(pq * pq, double);            /* pq x pq */

        Covy_x_swx = Calloc(pq * q, double);      /* pq x q  */

        C_kronecker(dCov_y, q, q, CT1, p, p, dCov_T);

        C_kronecker(dCov_y, q, q, swx, p, 1, Covy_x_swx);

        C_kronecker(Covy_x_swx, pq, q, swx, 1, p, CT2);

        for (k = 0; k < (pq * pq); k++)

            dCov_T[k] = f1 * dCov_T[k] - f2 * CT2[k];

        /* clean up */

        Free(CT2); Free(Covy_x_swx);

    }

    /* clean up */

    Free(swx); Free(CT1); 

}

/**

    R-interface to C_ExpectCovarLinearStatistic\n

    *\param x values of the transformation

    *\param y values of the influence function

    *\param weights case weights

    *\param expcovinf an object of class `ExpectCovarInfluence'

*/

SEXP R_ExpectCovarLinearStatistic(SEXP x, SEXP y, SEXP weights, 

                                  SEXP expcovinf) {

    SEXP ans;

    int n, p, q, pq;

    /* determine the dimensions and some checks */

    n  = nrow(x);

    p  = ncol(x);

    q  = ncol(y);

    pq = p * q;

    if (nrow(y) != n)

        error("y does not have %d rows", n);

    if (LENGTH(weights) != n) 

        error("vector of case weights does not have %d elements", n);

    PROTECT(ans = NEW_OBJECT(MAKE_CLASS("ExpectCovar")));

    SET_SLOT(ans, PL2_expectationSym, 

             PROTECT(allocVector(REALSXP, pq)));

    SET_SLOT(ans, PL2_covarianceSym, 

             PROTECT(allocMatrix(REALSXP, pq, pq)));

    C_ExpectCovarLinearStatistic(REAL(x), p, REAL(y), q, 

        REAL(weights), n, expcovinf, ans);

    UNPROTECT(3);

    return(ans);

}

/**

    Linear Statistic with permuted indices\n

    *\param x values of the transformation

    *\param p dimension of the transformation

    *\param y values of the influence function

    *\param q dimension of the influence function

    *\param n number of observations

    *\param nperm number of permutations

    *\param indx indices for the x-part

    *\param perm (permuted) indices for the y-part

    *\param ans return value; a pointer to a REALSXP-vector of length pq

*/

void C_PermutedLinearStatistic(const double *x, const int p,

                               const double *y, const int q,

                               const int n, const int nperm,

                               const int *indx, const int *perm, 

                               double *ans) {

    int i, j, k, kp, kn, knpi;

    for (k = 0; k < q; k++) {

        kn = k * n;

        kp = k * p;

        for (j = 0; j < p; j++) ans[kp + j] = 0.0;

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

            knpi = kn + perm[i];

            for (j = 0; j < p; j++)

                ans[kp + j] += x[j*n + indx[i]] * y[knpi];

        }

    }

}

/**

    Linear Statistic with permuted indices\n

    *\param x values of the transformation

    *\param y values of the influence function

    *\param indx indices for the x-part

    *\param perm (permuted) indices for the y-part

*/

SEXP R_PermutedLinearStatistic(SEXP x, SEXP y, SEXP indx, SEXP perm) {

    SEXP ans;

    int n, nperm, p, q, i, *iperm, *iindx;

    /* 

       only a basic check

    */

    if (!isReal(x) || !isReal(y))

        error("R_PermutedLinearStatistic: arguments are not of type REALSXP");

    if (!isInteger(perm))

        error("R_PermutedLinearStatistic: perm is not of type INTSXP");

    if (!isInteger(indx))

        error("R_PermutedLinearStatistic: indx is not of type INTSXP");

    n = nrow(y);

    nperm = LENGTH(perm);

    iperm = INTEGER(perm);

    if (LENGTH(indx)  != nperm)

        error("R_PermutedLinearStatistic: dimensions don't match");

    iindx = INTEGER(indx);

    if (nrow(x) != n)

        error("R_PermutedLinearStatistic: dimensions don't match");

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

        if (iperm[i] < 0 || iperm[i] > (n - 1) )

            error("R_PermutedLinearStatistic: perm is not between 1 and nobs");

        if (iindx[i] < 0 || iindx[i] > (n - 1) )

            error("R_PermutedLinearStatistic: indx is not between 1 and nobs");

    }

    q    = ncol(y);

    p    = ncol(x);

    PROTECT(ans = allocVector(REALSXP, p*q));

    C_PermutedLinearStatistic(REAL(x), p, REAL(y), q, n, nperm,

                 iindx, iperm, REAL(ans));

    UNPROTECT(1);

    return(ans);

}