###########################

## convenience functions ##

###########################

## obtain the number/ID for all terminal nodes

terminal_nodeIDs <- function(node) {

  if(node$terminal) return(node$nodeID)

  ll <- terminal_nodeIDs(node$left)

  rr <- terminal_nodeIDs(node$right)

  return(c(ll, rr))

}

#########################

## workhorse functions ##

#########################

### determine which observations go left or right

mob_fit_childweights <- function(node, mf, weights) {

    partvar <- mf@get("part")

    xselect <- partvar[[node$psplit$variableID]]

    ## we need to coerce ordered factors to numeric

    ## this is what party C code does as well!

    if (class(node$psplit) == "orderedSplit") {

        leftweights <- (as.double(xselect) <= node$psplit$splitpoint) * weights

        rightweights <- (as.double(xselect) > node$psplit$splitpoint) * weights

    } else {

        leftweights <- (xselect %in%

            levels(xselect)[as.logical(node$psplit$splitpoint)]) * weights

        rightweights <- (!(xselect %in%

            levels(xselect)[as.logical(node$psplit$splitpoint)])) * weights

    }

    list(left = leftweights, right = rightweights)

}

### setup a new (inner or terminal) node of a tree

mob_fit_setupnode <- function(obj, mf, weights, control) {

    ### control parameters

    alpha <- control$alpha

    bonferroni <- control$bonferroni

    minsplit <- control$minsplit

    trim <- control$trim

    objfun <- control$objfun

    verbose <- control$verbose

    breakties <- control$breakties

    parm <- control$parm

    ### if too few observations: no split = return terminal node

    if (sum(weights) < 2 * minsplit) {

        node <- list(nodeID = NULL, weights = weights,

                     criterion = list(statistic = 0, criterion = 0, maxcriterion = 0),

                     terminal = TRUE, psplit = NULL, ssplits = NULL,

                     prediction = 0, left = NULL, right = NULL,

                     sumweights = as.double(sum(weights)))

        class(node) <- "TerminalModelNode"

        return(node)

    }

    ### variable selection via fluctuation tests

    test <- try(mob_fit_fluctests(obj, mf, minsplit = minsplit, trim = trim,

      breakties = breakties, parm = parm))

    if (!inherits(test, "try-error")) {

        if(bonferroni) {

          pval1 <- pmin(1, sum(!is.na(test$pval)) * test$pval)

          pval2 <- 1 - (1-test$pval)^sum(!is.na(test$pval))

          test$pval <- ifelse(!is.na(test$pval) & (test$pval > 0.01), pval2, pval1)

        }

        best <- test$best

        TERMINAL <- is.na(best) || test$pval[best] > alpha

        if (verbose) {

            cat("\n-------------------------------------------\nFluctuation tests of splitting variables:\n")

            print(rbind(statistic = test$stat, p.value = test$pval))

            cat("\nBest splitting variable: ")

            cat(names(test$stat)[best])

            cat("\nPerform split? ")

            cat(ifelse(TERMINAL, "no", "yes"))

            cat("\n-------------------------------------------\n")    

        }

    } else {

        TERMINAL <- TRUE

        test <- list(stat = NA, pval = NA)

    }

    ### splitting

    na_max <- function(x) {

      if(all(is.na(x))) NA else max(x, na.rm = TRUE)

    }

    if (TERMINAL) {

        node <- list(nodeID = NULL, weights = weights,

                 criterion = list(statistic = test$stat, 

                                      criterion = 1 - test$pval,

                      maxcriterion = na_max(1 - test$pval)),

                     terminal = TRUE, psplit = NULL, ssplits = NULL,

                     prediction = 0, left = NULL, right = NULL, 

                     sumweights = as.double(sum(weights)))

        class(node) <- "TerminalModelNode"

        return(node)

    } else {

        partvar <- mf@get("part")

        xselect <- partvar[[best]]

        thissplit <- mob_fit_splitnode(xselect, obj, mf, weights, minsplit = minsplit, 

                                       objfun = objfun, verbose = verbose)

        ## check if splitting was unsuccessful

        if (identical(FALSE, thissplit)) {

            node <- list(nodeID = NULL, weights = weights,

                         criterion = list(statistic = test$stat, 

                                          criterion = 1 - test$pval, 

                                          maxcriterion = na_max(1 - test$pval)),

                         terminal = TRUE, psplit = NULL, ssplits = NULL,

                         prediction = 0, left = NULL, right = NULL, 

                         sumweights = as.double(sum(weights)))

            class(node) <- "TerminalModelNode"  

            ### more confusion than information

        ### warning("no admissable split found", call. = FALSE)

        if(verbose)

          cat(paste("\nNo admissable split found in ", sQuote(names(test$stat)[best]), "\n", sep = ""))     

        return(node)

        }

        thissplit$variableID <- best

        thissplit$variableName <- names(partvar)[best]

        node <- list(nodeID = NULL, weights = weights, 

                     criterion = list(statistic = test$stat, 

                                      criterion = 1 - test$pval, 

                                      maxcriterion = na_max(1 - test$pval)),

                     terminal = FALSE,

                     psplit = thissplit, ssplits = NULL, 

                     prediction = 0, left = NULL, right = NULL, 

                     sumweights = as.double(sum(weights)))

        class(node) <- "SplittingNode"

    }

    node$variableID <- best

    if (verbose) {

        cat("\nNode properties:\n")

        print(node$psplit, left = TRUE)

        cat(paste("; criterion = ", round(node$criterion$maxcriterion, 3), 

              ", statistic = ", round(max(node$criterion$statistic), 3), "\n",

              collapse = "", sep = ""))

    }

    node

}

### variable selection:

### conduct all M-fluctuation tests of fitted obj 

### with respect to each variable from a set of

### potential partitioning variables in mf

mob_fit_fluctests <- function(obj, mf, minsplit, trim, breakties, parm) {

  ## Cramer-von Mises statistic might be supported in future versions

  CvM <- FALSE

  ## set up return values

  partvar <- mf@get("part")

  m <- NCOL(partvar)

  pval <- rep.int(0, m)

  stat <- rep.int(0, m)

  ifac <- rep.int(FALSE, m)

  ## extract estimating functions  

  process <- as.matrix(estfun(obj))

  k <- NCOL(process)

  ## extract weights

  ww <- weights(obj)

  if(is.null(ww)) ww <- rep(1, NROW(process))

  n <- sum(ww)

  ## drop observations with zero weight

  ww0 <- (ww > 0)

  process <- process[ww0, , drop = FALSE]

  partvar <- partvar[ww0, , drop = FALSE]

  ww <- ww[ww0]

  ## repeat observations with weight > 1

  process <- process/ww

  ww1 <- rep.int(1:length(ww), ww)

  process <- process[ww1, , drop = FALSE]

  stopifnot(NROW(process) == n)

  ## scale process

  process <- process/sqrt(n)

  J12 <- root.matrix(crossprod(process))

  process <- t(chol2inv(chol(J12)) %*% t(process))  

  ## select parameters to test

  if(!is.null(parm)) process <- process[, parm, drop = FALSE]

  k <- NCOL(process)

  ## get critical values for CvM statistic

  if(CvM) {

    if(k > 25) k <- 25 #Z# also issue warning

    critval <- get("sc.meanL2")[as.character(k), ]

  } else {

    from <- if(trim > 1) trim else ceiling(n * trim)

    from <- max(from, minsplit)

    to <- n - from

    lambda <- ((n-from)*to)/(from*(n-to))

    beta <- get("sc.beta.sup")

    logp.supLM <- function(x, k, lambda)

    {

      if(k > 40) {

        ## use Estrella (2003) asymptotic approximation

        logp_estrella2003 <- function(x, k, lambda)

          -lgamma(k/2) + k/2 * log(x/2) - x/2 + log(abs(log(lambda) * (1 - k/x) + 2/x))

        ## FIXME: Estrella only works well for large enough x

    ## hence require x > 1.5 * k for Estrella approximation and

    ## use an ad hoc interpolation for larger p-values

    p <- ifelse(x <= 1.5 * k, (x/(1.5 * k))^sqrt(k) * logp_estrella2003(1.5 * k, k, lambda), logp_estrella2003(x, k, lambda))

      } else {

        ## use Hansen (1997) approximation

        m <- ncol(beta)-1

        if(lambda<1) tau <- lambda

        else tau <- 1/(1+sqrt(lambda))

        beta <- beta[(((k-1)*25 +1):(k*25)),]

        dummy <- beta[,(1:m)]%*%x^(0:(m-1))

        dummy <- dummy*(dummy>0)

        pp <- pchisq(dummy, beta[,(m+1)], lower.tail = FALSE, log.p = TRUE)

        if(tau==0.5)

          p <- pchisq(x, k, lower.tail = FALSE, log.p = TRUE)

        else if(tau <= 0.01)

          p <- pp[25]

        else if(tau >= 0.49)

          p <- log((exp(log(0.5-tau) + pp[1]) + exp(log(tau-0.49) + pchisq(x,k,lower.tail = FALSE, log.p = TRUE)))*100)

        else

        {

          taua <- (0.51-tau)*50

          tau1 <- floor(taua)

          p <- log(exp(log(tau1 + 1 - taua) + pp[tau1]) + exp(log(taua-tau1) + pp[tau1+1]))

        }

      }

      return(as.vector(p))

    }

  }

  ## compute statistic and p-value for each ordering

  for(i in 1:m) {

    pvi <- partvar[,i]

    pvi <- pvi[ww1]

    if(is.factor(pvi)) {

      proci <- process[ORDER(pvi), , drop = FALSE]

      ifac[i] <- TRUE

      # re-apply factor() added to drop unused levels

      pvi <- factor(pvi[ORDER(pvi)])

      # compute segment weights

      segweights <- as.vector(table(pvi))/n ## tapply(ww, pvi, sum)/n      

      # compute statistic only if at least two levels are left

      if(length(segweights) < 2) {

        stat[i] <- 0

    pval[i] <- NA

      } else {      

        stat[i] <- sum(sapply(1:k, function(j) (tapply(proci[,j], pvi, sum)^2)/segweights))

        pval[i] <- pchisq(stat[i], k*(length(levels(pvi))-1), log.p = TRUE, lower.tail = FALSE)

      }

    } else {

      oi <- if(breakties) {

        mm <- sort(unique(pvi))

    mm <- ifelse(length(mm) > 1, min(diff(mm))/10, 1)

    ORDER(pvi + runif(length(pvi), min = -mm, max = +mm))

      } else {

        ORDER(pvi)

      }    

      proci <- process[oi, , drop = FALSE]

      proci <- apply(proci, 2, cumsum)

      stat[i] <- if(CvM) sum((proci)^2)/n 

        else if(from < to) {

      xx <- rowSums(proci^2)

      xx <- xx[from:to]

      tt <- (from:to)/n

      max(xx/(tt * (1-tt)))   

    } else {

      0

    }

      pval[i] <- if(CvM) log(approx(c(0, critval), c(1, 1-as.numeric(names(critval))), stat[i], rule=2)$y)

        else if(from < to) logp.supLM(stat[i], k, lambda) else NA

    }

  }

  ## select variable with minimal p-value

  best <- which.min(pval)

  if(length(best) < 1) best <- NA

  rval <- list(pval = exp(pval), stat = stat, best = best)

  names(rval$pval) <- names(partvar)

  names(rval$stat) <- names(partvar)

  if (!all(is.na(rval$best)))

      names(rval$best) <- names(partvar)[rval$best]

  return(rval)

}

### split in variable x, either ordered or nominal

mob_fit_splitnode <- function(x, obj, mf, weights, minsplit, objfun, verbose = TRUE) {

    ## process minsplit (to minimal number of observations)

    if (minsplit > 0.5 & minsplit < 1) minsplit <- 1 - minsplit

    if (minsplit < 0.5)

        minsplit <- ceiling(sum(weights) * minsplit)

    if (is.numeric(x)) {

    ### for numerical variables

        ux <- sort(unique(x))

        if (length(ux) == 0) stop("cannot find admissible split point in x")

        dev <- vector(mode = "numeric", length = length(ux))

        for (i in 1:length(ux)) {

            xs <- x <= ux[i]

            if (mob_fit_checksplit(xs, weights, minsplit)) {

                dev[i] <- Inf

            } else {

                dev[i] <- mob_fit_getobjfun(obj, mf, weights, xs, objfun = objfun)

            }

        }

        ## maybe none of the possible splits is admissible

        if (all(!is.finite(dev))) return(FALSE)

        split <- list(variableID = NULL, ordered = TRUE, 

                      splitpoint = as.double(ux[which.min(dev)]),

                      splitstatistic = dev, toleft = TRUE)

        class(split) <- "orderedSplit"

    } else {

    ### for categorical variables

        al <- mob_fit_getlevels(x)

        dev <- apply(al, 1, function(w) {

                   xs <- x %in% levels(x)[w]

                   if (mob_fit_checksplit(xs, weights, minsplit)) {

                       return(Inf)

                   } else {

                       mob_fit_getobjfun(obj, mf, weights, xs, objfun = objfun)

                   }

               })

        if (verbose) {

            cat(paste("\nSplitting ", if(is.ordered(x)) "ordered ",

                  "factor variable, objective function: \n", sep = ""))

            print(dev)

        }

        if (all(!is.finite(dev))) return(FALSE)

        ## ordered factors are of storage mode "numeric" in party!

        ## initVariableFrame coerces ordered factors to storage.mode "numeric"

        ## the following is consistent with party

        if (is.ordered(x)) {

            split <- list(variableID = NULL, ordered = TRUE,

                          splitpoint = as.double(which.min(dev)),

                          splitstatistic = dev, toleft = TRUE)

            class(split) <- "orderedSplit"

            attr(split$splitpoint, "levels") <- levels(x)

        }  else {

            tab <- as.integer(table(x[weights > 0]) > 0)

            split <- list(variableID = NULL, ordered = FALSE,

                          splitpoint = as.integer(al[which.min(dev),]), 

                          splitstatistic = dev, 

                          toleft = TRUE, table = tab)

            attr(split$splitpoint, "levels") <- levels(x)

            class(split) <- "nominalSplit"

        }

    }

    split

}

### get partitioned objective function for a particular split

mob_fit_getobjfun <- function(obj, mf, weights, left, objfun = deviance) {

  ## mf is the model frame

  ## weights are the observation weights

  ## left is 1 (if left of splitpoint) or 0

  weightsleft <- weights * left

  weightsright <- weights * (1 - left)

  ### fit left / right model 

  fmleft <- reweight(obj, weights = weightsleft)

  fmright <- reweight(obj, weights = weightsright)

  return(objfun(fmleft) + objfun(fmright))

}

### determine all possible splits for a factor, both nominal and ordinal

mob_fit_getlevels <- function(x) {

    nl <- nlevels(x)

    if (inherits(x, "ordered")) {

        indx <- diag(nl)

        indx[lower.tri(indx)] <- 1

        indx <- indx[-nl,]

    rownames(indx) <- levels(x)[-nl]

    } else {

        mi <- 2^(nl - 1) - 1

        indx <- matrix(0, nrow = mi, ncol = nl)

        for (i in 1:mi) { # go though all splits #

            ii <- i

            for (l in 1:nl) {

                indx[i, l] <- ii%%2;

                ii <- ii %/% 2   

            }

        }

        rownames(indx) <- apply(indx, 1, function(z) paste(levels(x)[z > 0], collapse = "+"))

    }

    colnames(indx) <- as.character(levels(x))

    storage.mode(indx) <- "logical"

    indx

}

### check split

mob_fit_checksplit <- function(split, weights, minsplit)

    (sum(split * weights) < minsplit || sum((1 - split) * weights) < minsplit)