#--------------------------------------------------------------------------------------------------------------
# R functions accompanying "Boosting Method for Nonlinear Transformation Models with Censored Survival Data"
# Authors: Wenbin Lu and Lexin Li
#--------------------------------------------------------------------------------------------------------------


#--------------------------------------------------------------------------------------------------------------
# Functions for transformation model likelihood 
#--------------------------------------------------------------------------------------------------------------

# hazard and cumulative hazard functions for proportional hazards/odds models
Lambda = function(u,Fx,gamma){
  if(gamma == 0) return(exp(u+Fx))
  else return(log(1+gamma*exp(u+Fx))/gamma)
}

lambda = function(u,Fx,gamma){
  if(gamma == 0) return(exp(u+Fx))
  else return(exp(u+Fx)/(1+gamma*exp(u+Fx)))
}

dlambda = function(u,Fx,gamma){
  if(gamma == 0) return(exp(u+Fx))
  else return(exp(u+Fx)/(1+gamma*exp(u+Fx))^2)
}

ddlambda = function(u,Fx,gamma){
  if(gamma == 0) return(exp(u+Fx))
  else return(exp(u+Fx)*(1-(gamma^2)*exp(2*(u+Fx)))/(1+gamma*exp(u+Fx))^4)
}

# inverse CDF
invF = function(i,u,gamma){
  if(gamma == 0) return(log(-log(1-u[,i])))
  else return(log(((1/(1-u[,i])^gamma)-1)/gamma))
} 


# delta, z should be ordered according to the ranks of time
lik = function(i,sampH,delta,Fx,gamma){
  n = length(delta)
  Fxnull = numeric(n)
  L = ((lambda(sampH[,i],Fx,gamma)^delta)*exp(-Lambda(sampH[,i],Fx,gamma)))/
      ((lambda(sampH[,i],Fxnull,gamma)^delta)*exp(-Lambda(sampH[,i],Fxnull,gamma)))
          
  return(prod(L))
} 

dlik = function(i,sampH,delta,Fx,gamma){
  a = lik(i,sampH,delta,Fx,gamma)
  b = as.numeric(delta*dlambda(sampH[,i],Fx,gamma)/lambda(sampH[,i],Fx,gamma)- lambda(sampH[,i],Fx,gamma))
  return(a*b)
}  

wglik = function(i,w,sampH,delta,Fx,gx,gamma){
  n = length(delta)
  Fxnull = numeric(n)
  L = ((lambda(sampH[,i],Fx+w*gx,gamma)^delta)*exp(-Lambda(sampH[,i],Fx+w*gx,gamma)))/
      ((lambda(sampH[,i],Fxnull,gamma)^delta)*exp(-Lambda(sampH[,i],Fxnull,gamma)))
          
  return(prod(L))
} 

wgdlik = function(i,w,sampH,delta,Fx,gx,gamma){
  a = wglik(i,w,sampH,delta,Fx,gx,gamma)
  b = as.numeric(delta*dlambda(sampH[,i],Fx+w*gx,gamma)/lambda(sampH[,i],Fx+w*gx,gamma)- lambda(sampH[,i],Fx+w*gx,gamma))*gx
  c = sum(b)
  return(a*c)
} 

wgddlik = function(i,w,sampH,delta,Fx,gx,gamma){
  a = wglik(i,w,sampH,delta,Fx,gx,gamma)
  b = as.numeric(delta*dlambda(sampH[,i],Fx+w*gx,gamma)/lambda(sampH[,i],Fx+w*gx,gamma)- lambda(sampH[,i],Fx+w*gx,gamma))*gx
  c = sum(b)
  d = c*c
  e = as.numeric(delta*(ddlambda(sampH[,i],Fx+w*gx,gamma)*lambda(sampH[,i],Fx+w*gx,gamma) - dlambda(sampH[,i],Fx+w*gx,gamma)*dlambda(sampH[,i],Fx+w*gx,gamma))/(lambda(sampH[,i],Fx+w*gx,gamma)*lambda(sampH[,i],Fx+w*gx,gamma)) - dlambda(sampH[,i],Fx+w*gx,gamma))*gx*gx
  f = sum(e)
  return(a*(d+f))
}  


#sampling u according to ranks using progressive II censoring
spu = function(i,umat,n.cen,n,L){
  tempu = sort(umat[i,])
  rankumat = numeric(n)
  index = seq(1,n,1)
  
  i1=1
  i2=1+n.cen[1]
  if(n.cen[1] == 0){
    rankumat[i1] = tempu[1]
    index = index[-1]
  }
  if(n.cen[1] > 0){
    k = sample(index,n.cen[1])
    index = index[-k]
    rankumat[i1:(i2-1)] = 1.0e-8 ##should be 0, but numerically not feasible
    rankumat[i2] = tempu[index[1]]
    index = index[-1]
  } 
  j=2
  while(j<L){  
    i1 = i2 + 1
    i2 = i2 + n.cen[j] + 1
    if(n.cen[j] == 0){
      rankumat[i2] = tempu[index[1]]
      index = index[-1]
    }
    if(n.cen[j] > 0){
      rankumat[i1:(i2-1)] = rankumat[(i1-1)]
      k = sample(c(1:length(index)),n.cen[j])
      index = index[-k]
      rankumat[i2] = tempu[index[1]]
      index = index[-1]
    }
    j=j+1
  }
  if(n.cen[L] > 0){
    rankumat[(i2+1):n] = rankumat[i2]
  }
  return(rankumat)
}

maglik = function(M,sampH,delta,Fx,gamma){  
  s = sapply(1:M,lik,sampH=sampH,delta=delta,Fx=Fx,gamma=gamma)
  return(mean(s))
}

gradmaglik = function(M,sampH,delta,Fx,gamma){
  s = sapply(1:M,dlik,sampH=sampH,delta=delta,Fx=Fx,gamma=gamma)
  return(apply(s,1,mean)) 
}

# negative marginal likelihood wrt w
wgmaglik = function(w,M,sampH,delta,Fx,gx,gamma){  
  s = sapply(1:M,wglik,w=w,sampH=sampH,delta=delta,Fx=Fx,gx=gx,gamma=gamma)
  return(-1 * log(mean(s)))
  #return(-1 * mean(s))
}

gradwgmaglik = function(w,M,sampH,delta,Fx,gx,gamma){  
  s = sapply(1:M,wgdlik,w=w,sampH=sampH,delta=delta,Fx=Fx,gx=gx,gamma=gamma)
  return(mean(s))
}

hesswgmaglik = function(w,M,sampH,delta,Fx,gx,gamma){  
  s = sapply(1:M,wgddlik,w=w,sampH=sampH,delta=delta,Fx=Fx,gx=gx,gamma=gamma)
  return(mean(s))
}


# added by Lexin for nlm optimization
ltm1<-function(w,M,sampH,delta,Fx,gx,gamma)
{ 
   ans<-wgmaglik(w,M,sampH,delta,Fx,gx,gamma)  
   attr(ans, "gradient")<--1*gradwgmaglik(w,M,sampH,delta,Fx,gx,gamma)
   attr(ans, "hessian") <--1*hesswgmaglik(w,M,sampH,delta,Fx,gx,gamma)
   ans
}



#--------------------------------------------------------------------------------------------------------------
# Functions for boosting
#--------------------------------------------------------------------------------------------------------------

unispl<-function(y, x.tr, x.te=NULL) 
{
   p<-ncol(x.tr)
   crit<-NULL
   for(l in 1:p) {
      if(length(unique(x.tr[,l])) <  4) {
         fit.lm.l<-lm(y~x.tr[,l])
         yhat<-fitted(fit.lm.l)
      } else {
         fit.spl.l<-smooth.spline(x.tr[,l], y)
         yhat<-predict(fit.spl.l, x.tr[,l], deriv=0)$y
      }
      crit<-c(crit, mean((y - yhat)^2))
      #gx0<-predict(fit.spl.l, x=0, deriv=0)$y
      #crit<-c(crit, mean((y - yhat - gx0)^2))
   }
   l<-order(crit)[1]
  
   if(length(unique(x.tr[,l])) <  4) {
      fit.lm<-lm(y~x.tr[,l])
      gx0<-coef(fit.lm)[1]
      gx.tr<-fitted(fit.lm) - gx0
      dgx.tr<-coef(fit.lm)[2]
      if(!is.null(x.te)) {
         
      } else {
         gx.te<-NULL
      }      
   } else {
      fit.spl<-smooth.spline(x.tr[,l], y)
      gx0<-predict(fit.spl, x=0, deriv=0)$y
      gx.tr <-predict(fit.spl, x.tr[,l], deriv=0)$y - gx0
      dgx.tr<-predict(fit.spl, x.tr[,l], deriv=1)$y
      if(!is.null(x.te)) {
         gx.te <-predict(fit.spl, x.te[,l], deriv=0)$y - gx0
      } else {
         gx.te<-NULL
      }
   } 
   ans<-list(gx.tr=gx.tr, gx.te=gx.te, dgx.tr=dgx.tr, l=l)
}



surv.bst.tm<-function(data.x.tr, data.y.tr, data.x.te=NULL, data.y.te=NULL, max.iter=1000, lrate=0.05, rsp.rep=1000, gamma=0, line.approx=F) 
{
   # parameters
   n<-nrow(data.x.tr)
   p<-ncol(data.x.tr)
   stime<-data.y.tr[,1]
   delta<-data.y.tr[,2]
   o<-order(stime)
   r<-rank(stime)

   # resampling ranks based on progressive II censoring using uniform distribution
   n.fail<-sum(delta)                      # n.fail denotes number of failures in n subjects
   L<-n.fail + 1                           # L denotes number of intervals between failures
   n.cen<-numeric(L)                       # n.cen records the number of censored observations in each interval
   failrank<-sort(rank(stime)[delta==1])   # rank of ordered failure time
   temprank1<-c(0, failrank)
   temprank2<-c(failrank, n+1) 
   n.cen<-temprank2 - temprank1 - 1
   uvect<-runif(rsp.rep*n)
   umat<-matrix(uvect, nrow=rsp.rep, byrow=T)
   sampu<-sapply(1:rsp.rep, spu, umat=umat, n.cen=n.cen, n=n, L=L)
   sampH<-sapply(1:rsp.rep, invF, u=sampu, gamma=gamma)

   # initialization
   iter<-0
   mlik<-NULL
   Fx.tr<-rep(0, n)
   pred.list<-NULL
   fx.tr<-matrix(0, nrow=n, ncol=p)
   rinf.tr<-matrix(0, nrow=n, ncol=p)
   if(!is.null(data.x.te)) {
      Fx.te<-rep(0, nrow(data.x.te))
      fx.te<-matrix(0, nrow=nrow(data.x.te), ncol=p)
   } else {
      Fx.te<-fx.te<-NULL
   }

   gx.tr.a<-gx.te.a<-NULL
   dgx.tr.a<-NULL
   w.a<-NULL

   # loop
   while(iter < max.iter) {
      # compute the negative gradient of the negative log marginal likelihood using importance sampling
      U<-gradmaglik(M=rsp.rep, sampH=sampH, delta=delta[o], Fx=Fx.tr[o], gamma=gamma) / maglik(M=rsp.rep, sampH=sampH, delta=delta[o], Fx=Fx.tr[o], gamma=gamma)    
      U<-U[r]   # restore sample ordering according to rank of stime

      # fit the gradient using univariate spline
      out<-unispl(U, data.x.tr, data.x.te)
      gx.tr<-out$gx.tr
      gx.te<-out$gx.te
      pred.list<-c(pred.list, out$l)

      # line search
      if(!line.approx) {
         out.w<-nlm(wgmaglik, 1, M=rsp.rep, sampH=sampH, delta=delta[o], Fx=Fx.tr[o], gx=gx.tr[o], gamma=gamma)
         w<-out.w$estimate
      } else {
         fit.cox<-coxph(Surv(stime, delta)~gx.tr, weights=exp(Fx.tr))
         w<-unname(fit.cox$coef)
         #out.w<-nlm(wgmaglik, 1, M=rsp.rep, sampH=sampH, delta=delta[o], Fx=Fx.tr[o], gx=gx.tr[o], gamma=gamma)
         #wa.tmp<-c(wa.tmp, out.w$estimate)
      }

      # update
      Fx.tr<-Fx.tr + lrate * w * gx.tr
      Fx.te<-Fx.te + lrate * w * gx.te

      # record 
      gx.tr.a<-cbind(gx.tr.a, gx.tr)
      gx.te.a<-cbind(gx.te.a, gx.te)
      dgx.tr.a<-cbind(dgx.tr.a, out$dgx.tr)
      w.a<-c(w.a, w)

      # individual predictor contribution in F(x)
      fx.tr[, out$l]<-fx.tr[, out$l] + lrate * w * gx.tr
      fx.te[, out$l]<-fx.te[, out$l] + lrate * w * gx.te

      # individual predictor relative influence
      rinf.tr[, out$l]<-rinf.tr[, out$l] + lrate * w * out$dgx.tr

      # compute the marginal likelihood
      mlik<-c(mlik, maglik(M=rsp.rep, sampH=sampH, delta=delta[o], Fx=Fx.tr[o], gamma=gamma))       

      # update iter
      iter<-iter + 1
   }
   rinf<-sqrt(apply((rinf.tr^2), 2, mean) * apply(data.x.tr, 2, var))

   # return
   ans<-list(Fx.tr=Fx.tr, Fx.te=Fx.te, fx.tr=fx.tr, fx.te=fx.te, mlik=mlik, rinf=rinf, pred.list=pred.list, 
             gx.tr.a=gx.tr.a, gx.te.a=gx.te.a, dgx.tr.a=dgx.tr.a, w.a=w.a, rinf.tr=rinf.tr)
   #ans<-list(Fx.tr=Fx.tr, Fx.te=Fx.te, fx.tr=fx.tr, fx.te=fx.te, mlik=mlik, rinf=rinf, pred.list=pred.list)
   return(ans)
}



reconstruct<-function(object, max.iter=1000, lrate=0.05, varx.vec) 
{
   # parameters
   n.tr<-length(object$Fx.tr)
   n.te<-length(object$Fx.te)
   p<-length(object$rinf)

   # re-construction
   Fx.tr<-rep(0, n.tr)
   Fx.te<-rep(0, n.te)
   fx.tr<-matrix(0, nrow=n.tr, ncol=p)
   fx.te<-matrix(0, nrow=n.te, ncol=p)
   rinf.tr<-matrix(0, nrow=n.tr, ncol=p)
   for(i in 1:max.iter) {
      # contributing predictor
      l<-(object$pred.list)[i]

      # score function
      Fx.tr<-Fx.tr + lrate * object$w.a[i] * object$gx.tr.a[,i]
      Fx.te<-Fx.te + lrate * object$w.a[i] * object$gx.te.a[,i]

      # individual predictor contribution in F(x)
      fx.tr[, l]<-fx.tr[, l] + lrate * object$w.a[i] * object$gx.tr.a[,i]
      fx.te[, l]<-fx.te[, l] + lrate * object$w.a[i] * object$gx.te.a[,i]

      # individual predictor relative influence
      rinf.tr[, l]<-rinf.tr[, l] + lrate * object$w.a[i] * object$dgx.tr.a[,i]
   }
   rinf<-sqrt(apply((rinf.tr^2), 2, mean) * varx.vec)
   mlik<-object$mlik[1:max.iter]
   pred.list<-object$pred.list[1:max.iter]

   # return
   ans<-list(Fx.tr=Fx.tr, Fx.te=Fx.te, fx.tr=fx.tr, fx.te=fx.te, mlik=mlik, rinf=rinf, pred.list=pred.list)
   return(ans)
}