#----------------------------------------------------------
# Some useful functions
#----------------------------------------------------------
# Load library
library(doBy)
library(MASS)

#----------------------------------------------------------
# Estimate Cox Model with factors effects for each transitions
#  msdata : data with the long format as defined in the mstate package
#  covs : the names of covariates
#  expand : if TRUE, the covariates will be expanded for each transition.
#  time_scale : time or a duration multistate model
#  method : method for the Cox model.
#----------------------------------------------------------
estimcox <- function(msdata, covs, expand=T, time_scale = "clock.forward", method = "breslow"){ 

  # Extend the covariates
  if(expand){
    m <- ncol(msdata)
    tmat <- attributes(msdata)$trans
    msdata <- expand.covs(msdata,covs)
    
    # Extand covariates 
    Covar <- names(msdata[(m+1):length(msdata)])
    attr(msdata,"trans") <- tmat
    attr(msdata,"class") <- c("msdata","data.frame")
  }
  else{
    Covar=covs
  }
  
  # Cox formula
  if(time_scale =="clock.forward"){
    FormulaCox  <- paste("Surv(","Tstart",",","Tstop",",","status",")~", Covar[1], sep="")
  }
  if((time_scale =="clock.reset")){
    FormulaCox <- paste("Surv(","time",",","status",")~", Covar[1], sep="")
  }
  # Add covariates
  for (i in 2:length(Covar)){
    FormulaCox=paste(FormulaCox, Covar[i], sep="+")
  }
  
  FormulaCox <- paste(FormulaCox,"+strata(trans)", sep="")
  # Creaete Cox objet
  cox <- coxph(as.formula(FormulaCox), data = msdata,method = method)
  return(cox)
}

#----------------------------------------------------------
# Estimate Cox Model and transition intensities with factors effects for each transitions
#  msdata : data with the long format as defined in the mstate package
#  covs : the names of covariates
#  occur : the value taken by covariates
#  time_scale : time or a duration multistate model
#  variance : if TRUE the variance is calculated
#----------------------------------------------------------
estimintens <- function(msdata, tmat, covs, occur, time_scale = "clock.forward", variance = T){

  m <-  length(covs)
  # Create condition for selecting data 
  LogicalFiltre <- msdata[,covs[1]]==occur[1]
  if(m > 1){
    for(i in 2:m){
      LogicalFiltre <- LogicalFiltre & msdata[, covs[i]] == occur[i]
    }
  }
  # Estimate newdata
  Filtre <-  which(LogicalFiltre)
  NbTrans <- max(tmat[which(tmat>1)])
  newdata <- msdata[rep(Filtre[1],NbTrans),c("id",covs)]
  newdata$trans  <-  1:NbTrans
  newdata <- newdata[,-1]
  attr(newdata, "trans") = tmat
  newdata = expand.covs(newdata, covs)
  newdata$strata = newdata$trans
  #Estimate Cox
  cox <- estimcox(msdata,covs,T,time_scale)
  # Calculate intensities
  msf <- msfit(object = cox,newdata=newdata, vartype = "aalen", trans = tmat,variance=variance)
  return(msf)
}

#----------------------------------------------------------
# Get the incidence rates from a survfit object 
#  surv : a survfit object
#  AgeMin : lower bound for the age
#  AgeMax : upper bound for the age
#  strata : the names of strata
#  pas : the timescale
#----------------------------------------------------------
getqxmarg=function(surv, AgeMin, AgeMax, strata, pas=1){
  ##########################
  x=seq(AgeMin,AgeMax,pas)
  res=expand.grid(x,strata)
  names(res)=c("x","trans")
  #Alimentation du tableau
  for(i in strata){
      xdate=sapply(c(x,AgeMax+pas),function(y){max(surv[i]$time[surv[i]$time<=y])})
      temp=surv[i]$surv[surv[i]$time %in% xdate]
      res$qxmarg[res$trans==i]=(temp[1:(length(temp)-1)]-temp[2:length(temp)])/temp[1:(length(temp)-1)]
    }
  return(res)
}


##########################
#  A modified version of the function probtrans in the mstate package for calculating incidence rates
#----------------------------------------------------------
probtrans2=function (object, predt, pas=1, direction = c("forward", "fixedhorizon","ratecurve"), 
          method = c("aalen", "greenwood"), variance = TRUE, covariance = FALSE) 
{
  if (!inherits(object, "msfit")) 
    stop("'object' must be a 'msfit' object")
  method <- match.arg(method)
  direction <- match.arg(direction)
  trans <- object$trans
  transit <- to.trans2(trans)
  numtrans <- nrow(transit)
  stackhaz <- object$Haz
  stackvarhaz <- object$varHaz
  for (i in 1:numtrans) stackhaz$dhaz[stackhaz$trans == i] <- diff(c(0, 
                                                                     stackhaz$Haz[stackhaz$trans == i]))
  if (direction == "forward" | direction =="ratecurve") 
    stackhaz <- stackhaz[stackhaz$time > predt, ]
  else stackhaz <- stackhaz[stackhaz$time <= predt, ]
  untimes <- sort(unique(stackhaz$time))
  #Ajout vecteur des dates de ré-initialisation
  inittime<-seq(predt,max(stackhaz$time),pas)
  numrangetime=1
  dateinit=NULL
  if(inittime[length(inittime)]<max(stackhaz$time))
    inittime=c(inittime,max(stackhaz$time))
  TT <- length(untimes)
  S <- nrow(trans)
  if (covariance) 
    variance <- TRUE
  if (direction == "forward" | direction =="ratecurve") {
    if (variance == TRUE) 
      res <- array(0, c(TT + 1, 2 * S + 1, S))
    else res <- array(0, c(TT + 1, S + 1, S))
    res[1, 1, ] <- predt
    for (j in 1:S) res[1, 1 + j, ] <- rep(c(0, 1, 0), c(j - 
                                                          1, 1, S - j))
    if (variance) 
      res[1, (S + 2):(2 * S + 1), ] <- 0
  }
  else {
    if (predt %in% untimes) {
      if (variance) 
        res <- array(0, c(TT + 1, 2 * S + 1, S))
      else res <- array(0, c(TT + 1, S + 1, S))
      res[TT + 1, 1, ] <- predt
      for (j in 1:S) res[TT + 1, 1 + j, ] <- rep(c(0, 1, 
                                                   0), c(j - 1, 1, S - j))
      if (variance) 
        res[TT + 1, (S + 2):(2 * S + 1), ] <- 0
    }
    else {
      if (variance) 
        res <- array(0, c(TT + 2, 2 * S + 1, S))
      else res <- array(0, c(TT + 2, S + 1, S))
      res[TT + 1, 1, ] <- max(untimes)
      for (j in 1:S) res[TT + 1, 1 + j, ] <- rep(c(0, 1, 
                                                   0), c(j - 1, 1, S - j))
      if (variance) 
        res[TT + 1, (S + 2):(2 * S + 1), ] <- 0
      res[TT + 2, 1, ] <- predt
      for (j in 1:S) res[TT + 2, 1 + j, ] <- rep(c(0, 1, 
                                                   0), c(j - 1, 1, S - j))
      if (variance) 
        res[TT + 2, (S + 2):(2 * S + 1), ] <- 0
    }
  }
  P <- diag(S)
  if (covariance) 
    varParr <- array(0, c(TT, S^2, S^2))
  if (variance) {
    varP <- matrix(0, S^2, S^2)
    if (direction == "forward" | direction =="ratecurve") {
      varAnew <- array(0, c(S, S, S, S))
      if (predt != 0) {
        tmin <- max(stackvarhaz$time[stackvarhaz$time <= 
                                       predt])
        varHaz <- stackvarhaz[stackvarhaz$time == tmin, 
                              ]
        lHaz <- nrow(varHaz)
        for (j in 1:lHaz) {
          from1 <- transit$from[transit$transno == varHaz$trans1[j]]
          to1 <- transit$to[transit$transno == varHaz$trans1[j]]
          from2 <- transit$from[transit$transno == varHaz$trans2[j]]
          to2 <- transit$to[transit$transno == varHaz$trans2[j]]
          varAnew[from1, to1, from2, to2] <- varAnew[from2, 
                                                     to2, from1, to1] <- varHaz$varHaz[j]
        }
      }
    }
    else {
      varA <- array(0, c(S, S, S, S))
      varHaz <- stackvarhaz[stackvarhaz$time == untimes[TT], 
                            ]
      lHaz <- nrow(varHaz)
      for (j in 1:lHaz) {
        from1 <- transit$from[transit$transno == varHaz$trans1[j]]
        to1 <- transit$to[transit$transno == varHaz$trans1[j]]
        from2 <- transit$from[transit$transno == varHaz$trans2[j]]
        to2 <- transit$to[transit$transno == varHaz$trans2[j]]
        varA[from1, to1, from2, to2] <- varA[from2, to2, 
                                             from1, to1] <- varHaz$varHaz[j]
      }
    }
  }
  for (i in 1:TT) {
    idx <- ifelse(direction == "forward" | direction =="ratecurve", i, TT + 1 - i)
    tt <- untimes[idx]
    Haztt <- stackhaz[stackhaz$time == tt, ]
    lHaztt <- nrow(Haztt)
    IplusdA <- diag(S)
    for (j in 1:lHaztt) {
      from <- transit$from[transit$transno == Haztt$trans[j]]
      to <- transit$to[transit$transno == Haztt$trans[j]]
      IplusdA[from, to] <- Haztt$dhaz[j]
      IplusdA[from, from] <- IplusdA[from, from] - Haztt$dhaz[j]
    }
    if (any(diag(IplusdA) < 0)) 
      warning("Warning! Negative diagonal elements of (I+dA); the estimate may not be meaningful. \n")
    #Ré-initialisation au date de ré-initialisation
    if (tt>=inittime[numrangetime+1] & direction == "ratecurve"){
      numrangetime=numrangetime+1
      P=diag(S)
      if (variance) {varP=matrix(0, S^2, S^2)}
      dateinit=c(dateinit,untimes[max(1,idx-1)])
    }
    if (variance) { 
      if (direction == "forward" | direction =="ratecurve") {
        varA <- varAnew
        varAnew <- array(0, c(S, S, S, S))
        varHaztt <- stackvarhaz[stackvarhaz$time == tt, 
                                ]
        lHaztt <- nrow(varHaztt)
        for (j in 1:lHaztt) {
          from1 <- transit$from[transit$transno == varHaztt$trans1[j]]
          to1 <- transit$to[transit$transno == varHaztt$trans1[j]]
          from2 <- transit$from[transit$transno == varHaztt$trans2[j]]
          to2 <- transit$to[transit$transno == varHaztt$trans2[j]]
          varAnew[from1, to1, from2, to2] <- varAnew[from2, 
                                                     to2, from1, to1] <- varHaztt$varHaz[j]
        }
        vardA <- varAnew - varA
      }
      else {
        varAttmin <- array(0, c(S, S, S, S))
        varHazttmin <- stackvarhaz[stackvarhaz$time == 
                                     untimes[idx - 1], ]
        lHazttmin <- nrow(varHazttmin)
        for (j in 1:lHazttmin) {
          from1 <- transit$from[transit$transno == varHazttmin$trans1[j]]
          to1 <- transit$to[transit$transno == varHazttmin$trans1[j]]
          from2 <- transit$from[transit$transno == varHazttmin$trans2[j]]
          to2 <- transit$to[transit$transno == varHazttmin$trans2[j]]
          varAttmin[from1, to1, from2, to2] <- varAttmin[from2, 
                                                         to2, from1, to1] <- varHazttmin$varHaz[j]
        }
        vardA <- varA - varAttmin
        varA <- varAttmin
      }
      for (from in 1:S) {
        for (from2 in 1:S) {
          for (to2 in 1:S) {
            if (to2 != from2) 
              vardA[from, from, from2, to2] <- vardA[from2, 
                                                     to2, from, from] <- -sum(vardA[from, 
                                                                                    -from, from2, to2])
          }
        }
      }
      for (from in 1:S) {
        for (from2 in 1:S) vardA[from, from, from2, from2] <- vardA[from2, 
                                                                    from2, from, from] <- -sum(vardA[from, from, 
                                                                                                     from2, -from2])
      }
      vardA <- matrix(vardA, S^2, S^2)
    }
    if (method == "aalen") {
      if (direction == "forward" | direction =="ratecurve") {
        P <- P %*% IplusdA
        if (variance) {
          tmp1 <- kronecker(t(IplusdA), diag(S)) %*% 
            varP %*% kronecker(IplusdA, diag(S))
          tmp2 <- kronecker(diag(S), P) %*% vardA %*% 
            kronecker(diag(S), t(P))
          varP <- tmp1 + tmp2
        }
      }
      else {
        if (variance) {
          tmp1 <- kronecker(diag(S), IplusdA) %*% varP %*% 
            kronecker(diag(S), t(IplusdA))
          tmp2 <- kronecker(t(P), IplusdA) %*% vardA %*% 
            kronecker(P, t(IplusdA))
          varP <- tmp1 + tmp2
        }
        P <- IplusdA %*% P
      }
    }
    if (method == "greenwood") {
      if (direction == "forward" | direction =="ratecurve") {
        if (variance) {
          tmp1 <- kronecker(t(IplusdA), diag(S)) %*% 
            varP %*% kronecker(IplusdA, diag(S))
          tmp2 <- kronecker(diag(S), P) %*% vardA %*% 
            kronecker(diag(S), t(P))
          varP <- tmp1 + tmp2
        }
        P <- P %*% IplusdA
      }
      else {
        if (variance) {
          tmp1 <- kronecker(diag(S), IplusdA) %*% varP %*% 
            kronecker(diag(S), t(IplusdA))
          tmp2 <- kronecker(t(P), diag(S)) %*% vardA %*% 
            kronecker(P, diag(S))
          varP <- tmp1 + tmp2
        }
        P <- IplusdA %*% P
      }
    }
    if (variance) {
      seP <- sqrt(diag(varP))
      seP <- matrix(seP, S, S)
    }
    if (covariance) 
      varParr[i, , ] <- varP
    if (direction == "forward" | direction =="ratecurve") {
      res[idx + 1, 1, ] <- tt
      res[idx + 1, 2:(S + 1), ] <- t(P)
      if (variance) 
        res[idx + 1, (S + 2):(2 * S + 1), ] <- t(seP)
    }
    else {
      res[idx, 1, ] <- ifelse(i == TT, 0, untimes[TT - 
                                                    i])
      res[idx, 2:(S + 1), ] <- t(P)
      if (variance) 
        res[idx, (S + 2):(2 * S + 1), ] <- t(seP)
    }
  }
  if (covariance == FALSE) 
    res2 <- vector("list", S)
  else {
    res2 <- vector("list", S + 1)
    res2[[S + 1]] <- varParr
  }
  for (s in 1:S) {
    tmp <- as.data.frame(res[, , s])
    if (min(dim(tmp)) == 1) 
      tmp <- res[, , s]
    if (variance) 
      names(tmp) <- c("time", paste("pstate", 1:S, sep = ""), 
                      paste("se", 1:S, sep = ""))
    else names(tmp) <- c("time", paste("pstate", 1:S, sep = ""))
    res2[[s]] <- tmp
    #Selection des résultats uniquement sur les dates de réinitialisation
    if (direction == "ratecurve"){
      res2[[s]]=subset(res2[[s]],res2[[s]]$time %in% c(predt,dateinit))
    }
  }
  res2$trans <- trans
  class(res2) <- "probtrans"
  return(res2)
}
#----------------------------------------------------------
to.trans2=function (trans) 
{
  dm <- dim(trans)
  mx <- max(trans, na.rm = TRUE)
  res <- matrix(NA, mx, 3)
  res[, 1] <- 1:mx
  transvec <- as.vector(trans)
  for (i in 1:mx) {
    idx <- which(transvec == i)
    res[i, 2:3] <- c(idx%%dm[2], idx%/%dm[2] + 1)
  }
  res <- data.frame(res)
  names(res) <- c("transno", "from", "to")
  if (!is.null(dimnames(trans))) {
    states <- dimnames(trans)[[1]]
    res$fromname <- states[res$from]
    res$toname <- states[res$to]
  }
  return(res)
}
#----------------------------------------------------------
# The main function which estimate the model and regularize incidence rates. This function is used in the bootstrap procedure. 
process.qx=function(data,AgeMin,AgeMax,crit_WH,crit_WH_marg,poids){
  # Parameter
  tmat=attr(data, "trans")
  names_trans=c("Pathologie 1","Pathologie 2","Pathologie 3","Pathologie 4","Deces","Resiliation")
  #Estimation modèle multi-états
  c=coxph(Surv(Tstart,Tstop,status)~strata(trans),data=data,method = "breslow")
  msf=msfit(object = c, vartype = "aalen", trans = tmat,variance=T)
  qx=probtrans2(msf,predt=AgeMin,pas=1,method="aalen",variance=F,direction="ratecurve") 
  #Dataset des résultats multi-états
  Age=rep(qx[[1]]$time,6)
  qx.brut=c(qx[[1]]$pstate2,
         qx[[1]]$pstate3,
         qx[[1]]$pstate4,
         qx[[1]]$pstate5,
         qx[[1]]$pstate6,
         qx[[1]]$pstate7)
  trans=c(rep(1,nrow(qx[[1]])),
          rep(2,nrow(qx[[1]])),
          rep(3,nrow(qx[[1]])),
          rep(4,nrow(qx[[1]])),
          rep(5,nrow(qx[[1]])),
          rep(6,nrow(qx[[1]])))
  trans=as.factor(trans)
  trans=factor(trans,labels=names_trans)
  qx_nhmm=data.frame(Age,qx.brut,trans)
  rm(Age,qx.brut,trans)

  
  qx_nhmm$Age=qx_nhmm$Age-1
  qx_nhmm=qx_nhmm[qx_nhmm$Age!=AgeMin-1,]
  qx_nhmm=subset(qx_nhmm,qx_nhmm$Age<=AgeMax)
  qx_nhmm$Age=rep(AgeMin:AgeMax,length(names_trans))
  # Estimation with KM 
  surv=survfit(Surv(Tstart,Tstop,status)~trans,data=data,type = "kaplan-meier",se.fit=FALSE)
  qx_marg=getqxmarg(surv,AgeMin,AgeMax,strata=1:length(names_trans),pas=1)
  qx_marg=data.frame(Age=qx_marg$x,qx.brut=qx_marg$qxmarg,trans=qx_marg$trans)
  qx_marg$trans=as.factor(qx_marg$trans)
  qx_marg$trans=factor(qx_marg$trans,labels=names_trans)
  # Smooth the rates
  qx_nhmm$qx.fit=0
  qx_marg$qx.fit_marg=0
  # Use the best criteria W-H
  for (i in names_trans){
    y=qx_nhmm$qx.brut[qx_nhmm$trans==i]
    z=qx_marg$qx.brut[qx_marg$trans==i]
    ordre=crit_WH$ordre[crit_WH$trans==i]
    alpha=crit_WH$alpha[crit_WH$trans==i]
    ordre_marg=crit_WH_marg$ordre[crit_WH_marg$trans==i]
    alpha_marg=crit_WH_marg$alpha[crit_WH_marg$trans==i]   
    temp.WH=W_H1(y,poids,ordre,alpha)
    temp.WH_marg=W_H1(z,poids,ordre_marg,alpha_marg)
    qx_nhmm$qx.fit[qx_nhmm$trans==i]=temp.WH$fit
    qx_marg$qx.fit[qx_marg$trans==i]=temp.WH_marg$fit
    rm(y,ordre,alpha,ordre_marg,alpha_marg,temp.WH,temp.WH_marg)
  }
  # The algorithme to added dependance between the rates estimated using KM estimators
  qx_marg=marg.to.trans.borne(qx_marg)
  qx_nhmm$marg.up=qx_marg$qx.fit
  qx_nhmm$marg.down=qx_marg$qx.borne
  return(qx_nhmm)
}

#----------------------------------------------------------
boot.process.qx=function (theta, data, B = 3, id = "id", ...) 
{
  if (!inherits(data, "msdata")) 
    stop("'data' must be a 'msdata' object")
  trans <- attr(data, "trans")
  ids <- unique(data[[id]])
  res <- NULL
  ptm <- proc.time()
  dd=data.frame(id=data[[id]],num=1:length(data[[id]]))
  for (b in 1:B) {
    #Texte d'information
    cat("\nBootstrap replication", b, "\n")
    print(date())
    flush.console()
    bootdata <- NULL
    bids <- data.frame(x=sample(ids, replace = TRUE))
    bidxs <- merge(bids,dd,by.x=1,by.y=1,all.x=T)  
    #unlist(sapply(bids, function(x) which(x == data[[id]]))) 
    bootdata <- data[bidxs$num, ]
    cat("estimate Qx's for the bootstrapped sample...")
    thstar <- theta(bootdata, ...)
    res=c(res,list(thstar))
    print(proc.time()-ptm)
    rm(bootdata,bids,bidxs,thstar)
  }
  cat("\n")
  return(res)
}


#----------------------------------------------------------
marg.to.trans.borne=function(qx.data){
  ##########################
  qx.data$qx.borne=0
  trans.list=levels(qx.data$trans)
  n=length(trans.list)
  for (i in 1:n){
    temp.list=trans.list
    temp.list[i]=temp.list[n]
    temp.list[n]=trans.list[i]
    borne=marg.to.trans(qx.data,temp.list)
    qx.data$qx.borne[qx.data$trans==trans.list[i]]=borne$qx.ajust[borne$trans==trans.list[i]]
  }
  return(qx.data)
}
