### iid2.R --- 
#----------------------------------------------------------------------
## author: Brice Ozenne
## created: okt 12 2017 (13:16) 
## Version: 
## last-updated: feb 11 2019 (13:26) 
##           By: Brice Ozenne
##     Update #: 482
#----------------------------------------------------------------------
## 
### Commentary: 
## 
### Change Log:
#----------------------------------------------------------------------
## 
### Code:

## * Documentation - iid2
#' @title  Extract corrected i.i.d. decomposition
#' @description  Extract corrected i.i.d. decomposition from a gaussian linear model.
#' @name iid2
#'
#' @param object a linear model or a latent variable model
#' @param param [named numeric vector] the fitted parameters.
#' @param data [data.frame] the data set.
#' @param cluster [integer vector] the grouping variable relative to which the observations are iid.
#' @param bias.correct [logical] should the standard errors of the coefficients be corrected for small sample bias? Only relevant if the \code{sCorrect} function has not yet be applied to the object.
#' @param robust [logical] if \code{FALSE}, the i.i.d. decomposition is rescaled such its the squared sum equals the model-based standard error (instead of the robust standard error).
#' @param ... arguments to be passed to \code{sCorrect}.
#'
#' @details If argument \code{p} or \code{data} is not null, then the small sample size correction is recomputed to correct the influence function.
#'
#' @seealso \code{\link{sCorrect}} to obtain \code{lm2}, \code{gls2}, \code{lme2}, or \code{lvmfit2} objects.
#'
#' @return A matrix containing the 1st order influence function relative to each sample (in rows)
#' and each model coefficient (in columns).
#' 
#' @examples
#' n <- 5e1
#' p <- 3
#' X.name <- paste0("X",1:p)
#' link.lvm <- paste0("Y~",X.name)
#' formula.lvm <- as.formula(paste0("Y~",paste0(X.name,collapse="+")))
#'
#' m <- lvm(formula.lvm)
#' distribution(m,~Id) <- sequence.lvm(0)
#' set.seed(10)
#' d <- sim(m,n)
#'
#' ## linear model
#' e.lm <- lm(formula.lvm,data=d)
#' iid.tempo <- iid2(e.lm, bias.correct = FALSE)
#' range(iid.tempo[,1:4]-iid(e.lm))
#' 
#' ## latent variable model
#' e.lvm <- estimate(lvm(formula.lvm),data=d)
#' iid.tempo <- iid2(e.lvm, bias.correct = FALSE)
#' range(iid.tempo-iid(e.lvm))
#' ## difference due to the use of the observed info matrix vs. the expected one.
#'
#' ## rescale i.i.d using model-based standard error
#' iid.tempo <- iid2(e.lvm, robust = FALSE, bias.correct = FALSE)
#' diag(crossprod(iid.tempo))-diag(vcov(e.lvm))
#'
#' @concept small sample inference
#' @concept iid decomposition
#' @export
`iid2` <-
  function(object, ...) UseMethod("iid2")

## * iid2.lm
#' @rdname iid2
#' @export
iid2.lm <- function(object,
                    param = NULL, data = NULL, bias.correct = TRUE, ...){

    sCorrect(object, param = param, data = data, df = FALSE) <- bias.correct

    iid.tempo <- iid2(object, ...)

    ## ** keep track of the NA
    if(!is.null(object$na.action)){
        n.NA <- length(object$na.action)
        iid.withNA <- matrix(NA, nrow = NROW(iid.tempo) + n.NA, ncol = NCOL(iid.tempo),
                             dimnames = list(NULL, colnames(iid.tempo)))
        iid.withNA[-object$na.action,] <- iid.tempo
    }else{
        iid.withNA <- iid.tempo
    }
    
## ** export
    return(iid.withNA)
}

## * iid2.gls
#' @rdname iid2
#' @export
iid2.gls <- function(object, cluster = NULL,
                     param = NULL, data = NULL, bias.correct = TRUE, ...){

    sCorrect(object, cluster = cluster, param = param, data = data, df = FALSE) <- bias.correct

    if(!is.null(cluster)){
        cluster <- object$sCorrect$args$cluster
    }

### ** export
    return(iid2(object, ...))
}

## * iid2.lme
#' @rdname iid2
#' @export
iid2.lme <- iid2.lm

## * iid2.lvmfit
#' @rdname iid2
#' @export
iid2.lvmfit <- iid2.lm


## * iid2.lm2
#' @rdname iid2
#' @export
iid2.lm2 <- function(object, cluster = NULL,
                     param = NULL, data = NULL, robust = TRUE, ...){

    ## ** compute the score
    if(!is.null(param) || !is.null(data)){
        args <- object$sCorrect$args
        args$df <- FALSE
        object$sCorrect <- do.call(sCorrect,
                                   args = c(list(object, param = param, data = data),
                                            args))
    }else if(is.null(object$sCorrect$score)){
        stop("set argument \'score\' to TRUE when calling sCorrect \n")
    }

    ## ** compute iid
    if(is.null(cluster)){
        scoreCluster <- score2(object)
    }else{
        n.obs <- stats::nobs(object)

        if(length(cluster) != n.obs){
            stop("Argument \'cluster\' does not have the correct length (n=",n.obs,") \n")
        }
        scoreIndiv <- score2(object)
        scoreCluster <- do.call("rbind",tapply(1:n.obs,cluster, function(iIndex){
            colSums(scoreIndiv[iIndex,,drop=FALSE])
        }))
    }
    res <- scoreCluster %*% object$sCorrect$vcov.param
    if(robust == FALSE){
        vec.sigma <- sqrt(diag(object$sCorrect$vcov.param))
        vec.sigma.robust <- sqrt(apply(res^2,2,sum))
        res <- sweep(res, MARGIN = 2, FUN = "*", STATS = vec.sigma/vec.sigma.robust)
    }
    
    ## ** export
    return(res)

}

## * iid2.gls2
#' @rdname iid2
#' @export
iid2.gls2 <- iid2.lm2

## * iid2.lme2
#' @rdname iid2
#' @export
iid2.lme2 <- iid2.lm2

## * iid2.lvmfit2
#' @rdname iid2
#' @export
iid2.lvmfit2 <- function(object, cluster = NULL, data = NULL, ...){

    ## try to find cluster in data
    if(!is.null(object$cluster)){
        if(!is.null(cluster)){
            stop("Argument \'cluster\' must be NULL when the object already contain a cluster \n")
        }
        cluster <- object$cluster
    }else if(!is.null(cluster) && length(cluster) == 1){
        if(!is.null(data)){
            if(cluster %in%  names(data) == FALSE){
                stop("Could not find variable ",cluster," (argument \'cluster\') in argument \'data\' \n")
            }else{
                cluster <- data[[cluster]]
            }            
        }else{
            if(cluster %in% names(object$data$model.frame) == FALSE){
                stop("Could not find variable ",cluster," (argument \'cluster\') in argument object$data \n")
            }else{
                cluster <- object$data$model.frame[[cluster]]
            }            
        }
    }

    return(iid2.lm2(object = object, data = data, cluster = cluster, ...))
}



##----------------------------------------------------------------------
### iid2.R ends here
