library(rpart)
library(survival)
library(ipred)

source("func2.R")
options(contrasts=c("contr.treatment","contr.poly"))

## Bring in data -
# need covariates in a matrix
# survival time in days or months
# censoring time
# covariates to include in \bar{G}
# cross-validation level (vfold)
# independent test sample size

# As an example - the library ipred is called and the dataset GBSG2 is used to illustrate the code
data(GBSG2)
attach(GBSG2)
# put all data in a list to call function
comb.dat <- list()
comb.dat[[1]] <- as.numeric(time)
comb.dat[[2]] <- as.numeric(cens)
comb.dat[[3]] <- data.frame(GBSG2[,1:8])
comb.dat[[4]] <- data.frame(horTh=horTh, tgrade=tgrade)
vfold <- 5
ind.test <- T

  
# BEGIN CROSS-VALIDATION LOOP

#rename variables for use within loop
w <- comb.dat[[3]]
delta <- comb.dat[[2]]
ttilde <- comb.dat[[1]]
deltac <- 1 - delta
g.set <- comb.dat[[4]]
                                       

### Set up cross validation retaining censoring percentage
v <- ifelse(ind.test=="TRUE",vfold+1,vfold)
test.set.ind <- vfold+1
get.sam <- rep(c(1:v),sum(delta==1)/v)
if(length(get.sam)<sum(delta==1)) get.sam <- c(get.sam,c(1:(sum(delta==1)-length(get.sam))))
delta.1 <- sample(get.sam,sum(delta==1),replace=F)
get.cont <-rep(c(1:v),sum(delta==0)/v)
if(length(get.cont)<sum(delta==0)) get.cont <- c(get.cont,c(1:(sum(delta==0)-length(get.cont))))
delta.0 <- sample(get.cont,sum(delta==0),replace=F)

k <- 1
l <- 1
grp.delt <- NULL
for(m in 1:nrow(w)){
  if(delta[m]==0){
    grp.delt[m] <- delta.0[k]
    k <- k+1
  }
  if(delta[m]==1){
    grp.delt[m] <- delta.1[l]
    l <- l+1
  }
}

## take out the TEST group

if(ind.test){
  test.grp <- grp.delt==test.set.ind
  ttilde.ds <- ttilde[grp.delt!=test.set.ind]
  deltac.ds <- deltac[grp.delt!=test.set.ind]
  delta.ds <- delta[grp.delt!=test.set.ind]
  w.ds <- w[grp.delt!=test.set.ind,]
  grp.delt.ds <- grp.delt[grp.delt!=test.set.ind]
  g.set.ds <- g.set[grp.delt!=test.set.ind,]
}

                                      
bin <- 100 # Keep this large in case the trees become big
keep.sel.risk <- matrix(NA,nrow=bin,ncol=v-1) # matrix to retain the CV results

## Cross Validation Loop
for(i in 1:(v-1)){

  set.w <- (w.ds[grp.delt.ds!=i,])  # Training Set data
  tr <- grp.delt.ds!=i              # Indicator of training set 
  
######################################################
##   Univariate Cox Models for Survival              #
######################################################

# Get contrast matrix for \bar{G}
  mm <- model.matrix(~horTh+tgrade ,g.set.ds[tr,])  # Need to specify the matrix
  keep.for.cox <- apply(apply(g.set.ds[tr,],1,is.na),2,sum)
  ttilde.cox <- ttilde.ds[tr][keep.for.cox==0]
  deltac.cox <- deltac.ds[tr][keep.for.cox==0]
  delta.cox <- delta.ds[tr][keep.for.cox==0]
  cov.tr <- as.matrix(mm[,-1])
  surv.cox <- coxph(Surv(ttilde.cox, deltac.cox) ~ cov.tr)	
  surv.cox
                                        #	Build IPCW weights
		coeff.cox <- surv.cox$coefficients	
	#	Get baseline survival
		surv.cens <- survfit(surv.cox, newdata = data.frame(cov.tr = 0), 
			type = "kaplan-meier")	
	#	Vector of baseline survival
		basesurv <- get.at.surv.times(surv.cens, ttilde.cox)
	#	UG(D) for each subject
		ic0 <- get.init.am(basesurv, coeff.cox, cov.tr, delta.cox, ttilde.cox)

  ## what to do with missing data in \bar{G} now give them a one or 0
  cox.wt1 <- keep.for.cox
  cox.wt1[c(row.names(ic0))] <- ic0
  cox.wt1[keep.for.cox==1] <- delta.ds[tr][keep.for.cox==1]

  cpval <- 0



  r.x.cens <- rpart(log(ttilde.ds[tr])~ ., data=set.w  , method=alist, control=rpart.control(
                      cp = cpval, xval=0,minbucket=15,maxcompete=0,maxsurrogate=10),weight= cox.wt1)
  inde.c <-  r.x.cens$cptable[,"CP"]  # test samples



# Now build evaluation set, IPCW weights and evaluate training set's tree
  set.w <- w.ds[grp.delt.ds==i,]
  ts <- grp.delt.ds==i
 



######################################################
##   Univariate Cox Models for Survival              #
######################################################

# Get contrast matrix for \bar{G}
  mm <- model.matrix(~horTh + tgrade ,g.set.ds[ts,])  #SPECIFY SAME matrix!!!!
  keep.for.cox <- apply(apply(g.set.ds[ts,],1,is.na),2,sum)
  ttilde.cox <- ttilde.ds[ts][keep.for.cox==0]
  deltac.cox <- deltac.ds[ts][keep.for.cox==0]
  delta.cox <- delta.ds[ts][keep.for.cox==0]
  cov.ts <- as.matrix(mm[,-1])
  surv.cox <- coxph(Surv(ttilde.cox, deltac.cox) ~ cov.ts)	
  surv.cox

# BUILD IPCW weights
	#	Get coefficients from cox model
		coeff.cox <- surv.cox$coefficients	
	#	Get baseline survival
		surv.cens <- survfit(surv.cox, newdata = data.frame(cov.ts = 0), 
			type = "kaplan-meier")	
	#	Vector of baseline survival
		basesurv <- get.at.surv.times(surv.cens, ttilde.cox)
	#	UG(D) for each subject
		ic0 <- get.init.am(basesurv, coeff.cox, cov.ts, delta.cox, ttilde.cox)

  cox.wt1.ts <- keep.for.cox
  cox.wt1.ts[c(row.names(ic0))] <- ic0
  cox.wt1.ts[keep.for.cox==1] <- delta.ds[ts][keep.for.cox==1]

  cpval <- 0


# Calculate the cross-validated risk for each possible pruned tree keeping track of the number of nodes - to average overall

  pred.risk <- NULL
  for(k in (1:length(inde.c))){
    get.k <- inde.c[k]
    new.cB <- prune(r.x.cens,cp = (mean(c(inde.c[k],ifelse(k==1,inde.c[k]+1,inde.c[k-1])))))
    pred.out <- predict(new.cB,set.w)
    real.out <- (ttilde.ds[ts])
    pred.risk[k] <- sum(((real.out-pred.out)^2)*cox.wt1.ts)/length(real.out)
  }

  pred.risk
 
  # Store the error rate according to nsplit (not always the same for each rpart.tr)
  for(l in 1:length(r.x.cens$cptable[,"nsplit"])){
    right.row <- r.x.cens$cptable[l,"nsplit"]+1
    keep.sel.risk[right.row,i] <- pred.risk[l] 
  }
  for(m in 3:nrow(keep.sel.risk)){
    if(is.na(keep.sel.risk[(m-1),i])) keep.sel.risk[(m-1),i] <- keep.sel.risk[(m-2),i]
  }

} # end for loop


apply(keep.sel.risk.da[1:10,],1,mean,na.rm=TRUE)

# Pick level of tree with minimum cross-validated risk
sr.mean <- apply(keep.sel.risk[1:10,],1,mean,na.rm=TRUE)

if(length(c(1:length(sr.mean))[sr.mean==min(sr.mean)])==1){
  ind.tr <- c(1:length(sr.mean))[sr.mean==min(sr.mean)]
}
else { ind.tr <- min(c(1:length(sr.mean))[sr.mean==min(sr.mean)])}

ind.tr # level of tree with the minimimum risk 


####################################################################
### Now evaluate with the independent test set #####################
####################################################################

  
i <- test.set.ind      
set.w <- (w[grp.delt!=i,]) # data with the entire learning set included
tr <- grp.delt !=i         # indicator of learning set
  
######################################################
##   Univariate Cox Models for Survival              #
######################################################

# Get contast matrix for \bar{G}

mm <- model.matrix(~horTh + tgrade,g.set[tr,])  #Need to include correct variables for \bar{G}
keep.for.cox <- apply(apply(g.set[tr,],1,is.na),2,sum)
ttilde.cox <- ttilde[tr][keep.for.cox==0]
deltac.cox <- deltac[tr][keep.for.cox==0]
delta.cox <- delta[tr][keep.for.cox==0]
cov.tr <- as.matrix(mm[,-1])
surv.cox <- coxph(Surv(ttilde.cox, deltac.cox) ~ cov.tr)	
surv.cox


# Build IPCW weights with Cox PH
                                        
surv.cox <- coxph(Surv(ttilde.cox, deltac.cox) ~ cov.tr)	
	#	Get coefficients from cox model
		coeff.cox <- surv.cox$coefficients	
	#	Get baseline survival
		surv.cens <- survfit(surv.cox, newdata = data.frame(cov.tr = 0), 
			type = "kaplan-meier")	
	#	Vector of baseline survival
		basesurv <- get.at.surv.times(surv.cens, ttilde.cox)
	#	UG(D) for each subject
		ic0 <- get.init.am(basesurv, coeff.cox, cov.tr, delta.cox, ttilde.cox)

cox.wt1 <- keep.for.cox
cox.wt1[c(row.names(ic0))] <- ic0
cox.wt1[keep.for.cox==1] <- delta[tr][keep.for.cox==1]


# Run RPART with IPCW on learning set
cpval <- 0
r.x.cens <- rpart(log(ttilde[tr])~ ., data=set.w  , method=alist, control=rpart.control(
                  cp = cpval, xval=0,minbucket=15,maxcompete=0,maxsurrogate=10),weight= cox.wt1)
inde.c <-  r.x.cens$cptable[,"CP"]  # test samples


# Now just include the test set to evaluate tree on learning set
set.w <- w[grp.delt==i,]
ts <- grp.delt==i
 
######################################################
##   Univariate Cox Models for Survival              #
######################################################

mm <- model.matrix(~horTh + tgrade ,g.set[ts,])  
keep.for.cox <- apply(apply(g.set[ts,],1,is.na),2,sum)
ttilde.cox <- ttilde[ts][keep.for.cox==0]
deltac.cox <- deltac[ts][keep.for.cox==0]
delta.cox <- delta[ts][keep.for.cox==0]
cov.ts <- as.matrix(mm[,-1])
surv.cox <- coxph(Surv(ttilde.cox, deltac.cox) ~ cov.ts)	
surv.cox

surv.cox <- coxph(Surv(ttilde.cox, deltac.cox) ~ cov.ts)	
	#	Get coefficients from cox model
		coeff.cox <- surv.cox$coefficients	
	#	Get baseline survival
		surv.cens <- survfit(surv.cox, newdata = data.frame(cov.ts = 0), 
			type = "kaplan-meier")	
	#	Vector of baseline survival
		basesurv <- get.at.surv.times(surv.cens, ttilde.cox)
	#	UG(D) for each subject
		ic0 <- get.init.am(basesurv, coeff.cox, cov.ts, delta.cox, ttilde.cox)
cox.wt1.ts <- keep.for.cox
cox.wt1.ts[c(row.names(ic0))] <- ic0
cox.wt1.ts[keep.for.cox==1] <- delta[ts][keep.for.cox==1]

cpval <- 0


# Use L2 loss function to evalute the loss for each level of the tree

pred.risk <- NULL
for(k in (1:length(inde.c))){
  get.k <- inde.c[k]
  new.cB <- prune(r.x.cens,cp = (mean(c(inde.c[k],ifelse(k==1,inde.c[k]+1,inde.c[k-1])))))
  pred.out <- predict(new.cB,set.w)
  real.out <- log(ttilde[ts])
  pred.risk[k] <- sum(((real.out-pred.out)^2)*cox.wt1.ts)/length(real.out)
}

# Want level that minimizes the pred.risk
pred.risk