
/**
    Node splitting and prediction
    *\file Predict.c
    *\author $Author: thothorn $
    *\date $Date: 2024-08-15 13:57:20 +0200 (Thu, 15 Aug 2024) $
*/
                
#include "party.h"


/**
    Split a node according to a splitting rule \n
    *\param node the current node with primary split specified
    *\param learnsample learning sample
    *\param control an object of class `TreeControl'
    *\todo outplace the splitting since there are at least 3 functions
           with nearly identical code
*/
                
void C_splitnode(SEXP node, SEXP learnsample, SEXP control) {

    SEXP weights, leftnode, rightnode, split;
    SEXP responses, inputs, whichNA;
    double cutpoint, *dx, *dweights, *leftweights, *rightweights;
    double sleft = 0.0, sright = 0.0;
    int *ix, *levelset, *iwhichNA;
    int nobs, i, nna;
                    
    weights = S3get_nodeweights(node);
    dweights = REAL(weights);
    responses = GET_SLOT(learnsample, PL2_responsesSym);
    inputs = GET_SLOT(learnsample, PL2_inputsSym);
    nobs = get_nobs(learnsample);
            
    /* set up memory for the left daughter */
    SET_VECTOR_ELT(node, S3_LEFT, leftnode = allocVector(VECSXP, NODE_LENGTH));
    C_init_node(leftnode, nobs, 
        get_ninputs(learnsample), get_maxsurrogate(get_splitctrl(control)),
        ncol(get_predict_trafo(responses)));
    leftweights = REAL(S3get_nodeweights(leftnode));

    /* set up memory for the right daughter */
    SET_VECTOR_ELT(node, S3_RIGHT, 
                   rightnode = allocVector(VECSXP, NODE_LENGTH));
    C_init_node(rightnode, nobs, 
        get_ninputs(learnsample), get_maxsurrogate(get_splitctrl(control)),
        ncol(get_predict_trafo(responses)));
    rightweights = REAL(S3get_nodeweights(rightnode));

    /* split according to the primary split */
    split = S3get_primarysplit(node);
    if (has_missings(inputs, S3get_variableID(split))) {
        whichNA = get_missings(inputs, S3get_variableID(split));
        iwhichNA = INTEGER(whichNA);
        nna = LENGTH(whichNA);
    } else {
        nna = 0;
        whichNA = R_NilValue;
        iwhichNA = NULL;
    }
    
    if (S3is_ordered(split)) {
        cutpoint = REAL(S3get_splitpoint(split))[0];
        dx = REAL(get_variable(inputs, S3get_variableID(split)));
        for (i = 0; i < nobs; i++) {
            if (nna > 0) {
                if (i_in_set(i + 1, iwhichNA, nna)) continue;
            }
            if (dx[i] <= cutpoint) 
                leftweights[i] = dweights[i]; 
            else 
                leftweights[i] = 0.0;
            rightweights[i] = dweights[i] - leftweights[i];
            sleft += leftweights[i];
            sright += rightweights[i];
        }
    } else {
        levelset = INTEGER(S3get_splitpoint(split));
        ix = INTEGER(get_variable(inputs, S3get_variableID(split)));

        for (i = 0; i < nobs; i++) {
            if (nna > 0) {
                if (i_in_set(i + 1, iwhichNA, nna)) continue;
            }
            if (levelset[ix[i] - 1])
                leftweights[i] = dweights[i];
            else 
                leftweights[i] = 0.0;
            rightweights[i] = dweights[i] - leftweights[i];
            sleft += leftweights[i];
            sright += rightweights[i];
        }
    }
    
    /* for the moment: NA's go with majority */
    if (nna > 0) {
        for (i = 0; i < nna; i++) {
            if (sleft > sright) {
                leftweights[iwhichNA[i] - 1] = dweights[iwhichNA[i] - 1];
                rightweights[iwhichNA[i] - 1] = 0.0;
            } else {
                rightweights[iwhichNA[i] - 1] = dweights[iwhichNA[i] - 1];
                leftweights[iwhichNA[i] - 1] = 0.0;
            }
        }
    }
}


/**
    Get the terminal node for obs. number `numobs' of `newinputs' \n
    *\param subtree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param numobs observation number
    *\param varperm which variable shall be permuted?
    *\todo handle surrogate splits
*/

SEXP C_get_node(SEXP subtree, SEXP newinputs, 
                double mincriterion, int numobs, int varperm) {

    SEXP split, whichNA, ssplit, surrsplit;
    double cutpoint, x, swleft, swright;
    int level, *levelset, i, ns;

    if (S3get_nodeterminal(subtree) || 
        REAL(S3get_maxcriterion(subtree))[0] < mincriterion) 
        return(subtree);
    
    split = S3get_primarysplit(subtree);

    /* Maybe store the proportions left / right in each node? */
    swleft = S3get_sumweights(S3get_leftnode(subtree));
    swright = S3get_sumweights(S3get_rightnode(subtree));

    /* splits based on variable varperm are random */    
    if (S3get_variableID(split) == varperm) {
        if (unif_rand() < swleft / (swleft + swright)) {
            return(C_get_node(S3get_leftnode(subtree),
                       newinputs, mincriterion, numobs, varperm));
        } else {
            return(C_get_node(S3get_rightnode(subtree),
                       newinputs, mincriterion, numobs, varperm));
        }
    }
                   
    /* missing values */
    if (has_missings(newinputs, S3get_variableID(split))) {
        whichNA = get_missings(newinputs, S3get_variableID(split));
    
        /* numobs 0 ... n - 1 but whichNA has 1:n */
        if (C_i_in_set(numobs + 1, whichNA)) {
        
            surrsplit = S3get_surrogatesplits(subtree);
            ns = 0;
            i = numobs;      

            /* try to find a surrogate split */
            while(TRUE) {
    
                if (ns >= LENGTH(surrsplit)) break;
            
                ssplit = VECTOR_ELT(surrsplit, ns);
                /* any missings in the surrogate variable ? */
                if (has_missings(newinputs, S3get_variableID(ssplit))) {
                    /* i in 0 ... n - 1 but get_missings in 1:n */
                    /* if i is missing, move to next surrogate split */
                    if (C_i_in_set(i + 1, 
                                   get_missings(newinputs, S3get_variableID(ssplit)))) {
                        ns++;
                        continue;
                    }
                }

                cutpoint = REAL(S3get_splitpoint(ssplit))[0];
                x = REAL(get_variable(newinputs, S3get_variableID(ssplit)))[i];
                     
                if (S3get_toleft(ssplit)) {
                    if (x <= cutpoint) {
                        return(C_get_node(S3get_leftnode(subtree),
                                          newinputs, mincriterion, numobs, varperm));
                    } else {
                        return(C_get_node(S3get_rightnode(subtree),
                               newinputs, mincriterion, numobs, varperm));
                    }
                } else {
                    if (x <= cutpoint) {
                        return(C_get_node(S3get_rightnode(subtree),
                                          newinputs, mincriterion, numobs, varperm));
                    } else {
                        return(C_get_node(S3get_leftnode(subtree),
                               newinputs, mincriterion, numobs, varperm));
                    }
                }
                break;
            }

            /* if this was not successful, we go with the majority */
            if (swleft > swright) {
                return(C_get_node(S3get_leftnode(subtree), 
                                  newinputs, mincriterion, numobs, varperm));
            } else {
                return(C_get_node(S3get_rightnode(subtree), 
                                  newinputs, mincriterion, numobs, varperm));
            }
        }
    }
    
    if (S3is_ordered(split)) {
        cutpoint = REAL(S3get_splitpoint(split))[0];
        x = REAL(get_variable(newinputs, 
                     S3get_variableID(split)))[numobs];
        if (x <= cutpoint) {
            return(C_get_node(S3get_leftnode(subtree), 
                              newinputs, mincriterion, numobs, varperm));
        } else {
            return(C_get_node(S3get_rightnode(subtree), 
                              newinputs, mincriterion, numobs, varperm));
        }
    } else {
        levelset = INTEGER(S3get_splitpoint(split));
        level = INTEGER(get_variable(newinputs, 
                            S3get_variableID(split)))[numobs];
        /* level is in 1, ..., K */
        if (levelset[level - 1]) {
            return(C_get_node(S3get_leftnode(subtree), newinputs, 
                              mincriterion, numobs, varperm));
        } else {
            return(C_get_node(S3get_rightnode(subtree), newinputs, 
                              mincriterion, numobs, varperm));
        }
    }
}


/**
    R-Interface to C_get_node \n
    *\param subtree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param numobs observation number
*/

SEXP R_get_node(SEXP subtree, SEXP newinputs, SEXP mincriterion, 
                SEXP numobs, SEXP varperm) {
    return(C_get_node(subtree, newinputs, REAL(mincriterion)[0],
                      INTEGER(numobs)[0] - 1, INTEGER(varperm)[0]));
}


/**
    Get the node with nodeID `nodenum' \n
    *\param subtree a tree
    *\param nodenum a nodeID
*/

SEXP C_get_nodebynum(SEXP subtree, int nodenum) {
    
    if (nodenum == S3get_nodeID(subtree)) return(subtree);

    if (S3get_nodeterminal(subtree)) 
        error("no node with number %d\n", nodenum);

    if (nodenum < S3get_nodeID(S3get_rightnode(subtree))) {
        return(C_get_nodebynum(S3get_leftnode(subtree), nodenum));
    } else {
        return(C_get_nodebynum(S3get_rightnode(subtree), nodenum));
    }
}


/**
    R-Interface to C_get_nodenum \n
    *\param subtree a tree
    *\param nodenum a nodeID
*/

SEXP R_get_nodebynum(SEXP subtree, SEXP nodenum) {
    return(C_get_nodebynum(subtree, INTEGER(nodenum)[0]));
}


/**
    Get the prediction of a new observation\n
    *\param subtree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param numobs observation number
    *\param varperm which variable shall be permuted?
*/

SEXP C_get_prediction(SEXP subtree, SEXP newinputs, 
                      double mincriterion, int numobs, int varperm) {
    return(S3get_prediction(C_get_node(subtree, newinputs, 
                            mincriterion, numobs, varperm)));
}


/**
    Get the weights for a new observation \n
    *\param subtree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param numobs observation number
*/

SEXP C_get_nodeweights(SEXP subtree, SEXP newinputs, 
                       double mincriterion, int numobs) {
    return(S3get_nodeweights(C_get_node(subtree, newinputs, 
                             mincriterion, numobs, -1)));
}


/**
    Get the nodeID for a new observation \n
    *\param subtree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param numobs observation number
    *\param varperm which variable shall be permuted?
*/

int C_get_nodeID(SEXP subtree, SEXP newinputs,
                  double mincriterion, int numobs, int varperm) {
     return(S3get_nodeID(C_get_node(subtree, newinputs, 
            mincriterion, numobs, varperm)));
}


/**
    R-Interface to C_get_nodeID \n
    *\param tree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
*/

SEXP R_get_nodeID(SEXP tree, SEXP newinputs, SEXP mincriterion, SEXP varperm) {

    SEXP ans;
    int nobs, i, *dans;
            
    nobs = get_nobs(newinputs);
    PROTECT(ans = allocVector(INTSXP, nobs));
    dans = INTEGER(ans);
    for (i = 0; i < nobs; i++)
         dans[i] = C_get_nodeID(tree, newinputs, REAL(mincriterion)[0], i, INTEGER(varperm)[0]);
    UNPROTECT(1);
    return(ans);
}


/**
    Get all predictions for `newinputs' \n
    *\param tree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param varperm which variable shall be permuted?
    *\param ans return value
*/

void C_predict(SEXP tree, SEXP newinputs, double mincriterion, 
               int varperm, SEXP ans) {
    
    int nobs, i;
    
    nobs = get_nobs(newinputs);    
    if (LENGTH(ans) != nobs) 
        error("ans is not of length %d\n", nobs);
        
    for (i = 0; i < nobs; i++)
        SET_VECTOR_ELT(ans, i, C_get_prediction(tree, newinputs, 
                       mincriterion, i, varperm));
}


/**
    R-Interface to C_predict \n
    *\param tree a tree
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param varperm which variable shall be permuted? -1 for no permutation
*/

SEXP R_predict(SEXP tree, SEXP newinputs, SEXP mincriterion,
               SEXP varperm) {

    SEXP ans;
    int nobs;

    nobs = get_nobs(newinputs);
    PROTECT(ans = allocVector(VECSXP, nobs));
    GetRNGstate();
    C_predict(tree, newinputs, REAL(mincriterion)[0], 
              INTEGER(varperm)[0], ans);
    PutRNGstate();
    UNPROTECT(1);
    return(ans);
}

/**
    Get the predictions from `where' nodes\n
    *\param tree a tree
    *\param where vector of nodeID's
    *\param ans return value
*/

void C_getpredictions(SEXP tree, SEXP where, SEXP ans) {

    int nobs, i, *iwhere;
    
    nobs = LENGTH(where);
    iwhere = INTEGER(where);
    if (LENGTH(ans) != nobs)
        error("ans is not of length %d\n", nobs);
        
    for (i = 0; i < nobs; i++)
        SET_VECTOR_ELT(ans, i, S3get_prediction(
            C_get_nodebynum(tree, iwhere[i])));
}


/**
    R-Interface to C_getpredictions\n
    *\param tree a tree
    *\param where vector of nodeID's
*/
            
SEXP R_getpredictions(SEXP tree, SEXP where) {

    SEXP ans;
    int nobs;
            
    nobs = LENGTH(where);
    PROTECT(ans = allocVector(VECSXP, nobs));
    C_getpredictions(tree, where, ans);
    UNPROTECT(1);
    return(ans);
}                        

/**
    Predictions weights from RandomForest objects
    *\param forest a list of trees
    *\param where list (length b) of integer vectors (length n) containing terminal node numbers    
    *\param weights list (length b) of bootstrap case weights
    *\param newinputs an object of class `VariableFrame'
    *\param mincriterion overwrites mincriterion used for tree growing
    *\param oobpred a logical indicating out-of-bag predictions
*/

SEXP R_predictRF_weights(SEXP forest, SEXP where, SEXP weights, 
                         SEXP newinputs, SEXP mincriterion, SEXP oobpred, SEXP scale) {

    SEXP ans, tree, bw;
    int ntrees, nobs, i, b, j, iwhere, oob = 0, count = 0, ntrain, *tnsize;
    
    if (LOGICAL(oobpred)[0]) oob = 1;
    
    nobs = get_nobs(newinputs);
    ntrees = LENGTH(forest);

    if (oob) {
        if (LENGTH(VECTOR_ELT(weights, 0)) != nobs)
            error("number of observations don't match");
    }    
    
    tree = VECTOR_ELT(forest, 0);
    ntrain = LENGTH(VECTOR_ELT(weights, 0));
    
    PROTECT(ans = allocVector(VECSXP, nobs));
    
    tnsize = R_Calloc(ntrain, int);
    for (j = 0; j < ntrain; j++) tnsize[j] = 1;
    
    for (i = 0; i < nobs; i++) {
        count = 0;
        SET_VECTOR_ELT(ans, i, bw = allocVector(REALSXP, ntrain));
        for (j = 0; j < ntrain; j++)
            REAL(bw)[j] = 0.0;
        for (b = 0; b < ntrees; b++) {
            tree = VECTOR_ELT(forest, b);

            if (oob && 
                REAL(VECTOR_ELT(weights, b))[i] > 0.0) 
                continue;

            iwhere = C_get_nodeID(tree, newinputs, REAL(mincriterion)[0], i, -1);

            /* possibly scale by sum of weights in each terminal node */
            if (LOGICAL(scale)[0]) {
                for (j = 0; j < ntrain; j++)
                    tnsize[j] = 0;
                for (j = 0; j < ntrain; j++)
                    tnsize[INTEGER(VECTOR_ELT(where, b))[j] - 1] += REAL(VECTOR_ELT(weights, b))[j];
            }

            for (j = 0; j < ntrain; j++) {
                if (iwhere == INTEGER(VECTOR_ELT(where, b))[j])
                    REAL(bw)[j] += REAL(VECTOR_ELT(weights, b))[j] / 
                                   tnsize[INTEGER(VECTOR_ELT(where, b))[j] - 1];
            }
            count++;
        }
        if (count == 0) 
            error("cannot compute out-of-bag predictions for observation number %d", i + 1);
    }
    R_Free(tnsize);
    UNPROTECT(1);
    return(ans);
}


/**
    Proximity matrix for random forests
    *\param where list (length b) of integer vectors (length n) containing terminal node numbers
*/

SEXP R_proximity(SEXP where) {

    SEXP ans, bw, bin;
    int ntrees, nobs, i, b, j, iwhere;
    
    ntrees = LENGTH(where);
    nobs = LENGTH(VECTOR_ELT(where, 0));
    
    PROTECT(ans = allocVector(VECSXP, nobs));
    PROTECT(bin = allocVector(INTSXP, nobs));
     
    for (i = 0; i < nobs; i++) {
        SET_VECTOR_ELT(ans, i, bw = allocVector(REALSXP, nobs));
        for (j = 0; j < nobs; j++) {
            REAL(bw)[j] = 0.0;
            INTEGER(bin)[j] = 0;
        }
        for (b = 0; b < ntrees; b++) {
            /* don't look at out-of-bag observations */
            if (INTEGER(VECTOR_ELT(where, b))[i] == 0)
                continue;
            iwhere = INTEGER(VECTOR_ELT(where, b))[i];
            for (j = 0; j < nobs; j++) {
                if (iwhere == INTEGER(VECTOR_ELT(where, b))[j])
                    /* only count the number of trees; no weights */
                    REAL(bw)[j]++;
                if (INTEGER(VECTOR_ELT(where, b))[j] > 0)
                    /* count the number of bootstrap samples
                    containing both i and j */
                    INTEGER(bin)[j]++;
            }
        }
        for (j = 0; j < nobs; j++)
            REAL(bw)[j] = REAL(bw)[j] / INTEGER(bin)[j];
    }
    UNPROTECT(2);
    return(ans);
}
