# Chapter 8. MC models with CAV data. Fitted with optim()

# Ardo, UCL 2017.

# Run CH8_MC_CAV_msm.r first!

# Prelim:
digits <- 3

# Data info:
cat("\nCAV data:\n")
subjects <- as.numeric(names(table(cav$PTNUM)))
N <- length(subjects)
cat("Sample size =",N,"\n")
cat("\nState table:")
print(statetable.msm(state,PTNUM,data=cav))

# Choose model. 
# MODEL=1 for model in Jackson (2011)
# MODEL=2 for model in Chapter 1 with MC:
MODEL <- 2

# Prelim for models:
cav.split <- split(cav,cav$PTNUM)
Q <- oneway4.q
E <- matrix(0,4,4)
D <- 4
one <- rep(1,D)
logit  <- function(lp){ exp(lp)/(1+exp(lp)) }
probit <- function(lp){ pnorm(lp) }

#################################################
# Fit the msm model in Jackson (2011) using optim:
if(MODEL==1){
 
# Loglikelihood:
loglikelihood<-function(p){
 # Parameters:
 beta0 <- p[1:5]
 alpha <- p[6:9]
 #cat("beta0 and e: ", round(c(beta0,logit(alpha)),digits),"\n") 
 # Q matrix:
 Q[1,2] <- exp(beta0[1])
 Q[1,D] <- exp(beta0[2]);  Q[1,1] <- -(Q[1,2]+Q[1,D])
 Q[2,3] <- exp(beta0[3])
 Q[2,D] <- exp(beta0[4]);  Q[2,2] <- -(Q[2,3]+Q[2,D])
 Q[3,D] <- exp(beta0[5]);  Q[3,3] <- -Q[3,D]
 # MC:
 E[1,2] <- logit(alpha[1]); E[1,1] <- 1-E[1,2]
 E[2,1] <- logit(alpha[2]) 
 E[2,3] <- logit(alpha[3]); E[2,2] <- 1-E[2,1]-E[2,3]
 E[3,2] <- logit(alpha[4]); E[3,3] <- 1-E[3,2] 
 # Contribution per unit:
 loglik <- 0
 for(i in 1:N){
  data.i <- cav.split[[i]]
  O <- data.i$state
  t <- data.i$years
  T <- diag(D);T.j<-T
  # First state:
  f <- c(1,0,0,0)
  # Rest of the states
  for(j in 2:length(O)){
   # P matrix:
   P <- MatrixExp(mat=Q, t=t[j]-t[j-1])
   for(r in 1:D){
    for(s in 1:D){
     if(O[j]!=D){
       T.j[r,s] <- E[s,O[j]]*P[r,s]
     }else{
       T.j[r,s] <- P[r,s]*Q[s,D]
     }
    }
   }
   T <- T%*%T.j
  }
  loglik <- loglik+log(f%*%T%*%one)
  #cat("i = ",i, "and -2*Loglik = ", -2*loglik,"\n")
 }
 # Monitor:
 cat("-2*Loglik = ", -2*loglik,"\n")
 return(-loglik)
}

# Check with msm result:
cat("\nCheck loglikelihood function:\n")
cat("msm: -2loglik  =", misc.msm$minus2loglik,"\n")
E.msm <- ematrix.msm(misc.msm, ci="delta")
e <- c(E.msm$estimates[1,2],E.msm$estimates[2,1],
       E.msm$estimates[2,3],E.msm$estimates[3,2])
p0 <- c(misc.msm$estimates[1:5],log(e/(1-e)))
cat("-2loglikelihood function at msm estimate =", 
     2*loglikelihood(p0),"\n\n")

# Maximise:
p0  <- c(rep(-5,5),rep(-4,4))
max <- optim(par=p0, fn=loglikelihood, method = c("BFGS"),
           control=list(maxit=3000, fnscale=2000),hessian=TRUE)
cat("\n-2loglik =", 2*max$value,"\n")
cat("AIC =", 2*max$value+2*length(max$par),"\n")
conv <- max$convergence
cat("Convergence code =", conv,"\n")
p    <- max$par
p.se <- sqrt(diag(solve(max$hessian)))
print(cbind(q=c(qnames,rep("c",4)),p=round(p,digits),
            se=round(p.se,digits),"Wald ChiSq"=round((p/p.se)^2,digits),
            "Pr>ChiSq"=round(1-pchisq((p/p.se)^2,df=1),digits)),quote=FALSE)
}


#################################################
# Fit the model in Chapter 1 but now with MC:
if(MODEL==2){

# Loglikelihood:
loglikelihood <- function(p){
 # Parameters:
 beta0 <- p[1:5]
 alpha <- p[6:9]
 xi    <- p[10:11]
 beta1 <- p[12:13]
 beta2 <- p[14:18]
 #cat("beta0 and e: ", round(c(beta0,logit(alpha)),digits),"\n") 
 # MC:
 E[1,2] <- logit(alpha[1]); E[1,1] <- 1-E[1,2]
 E[2,1] <- logit(alpha[2]) 
 E[2,3] <- logit(alpha[3]); E[2,2] <- 1-E[2,1]-E[2,3]
 E[3,2] <- logit(alpha[4]); E[3,3] <- 1-E[3,2] 
 # Contribution per unit:
 loglik <- 0
 for(i in 1:N){
  data.i <- cav.split[[i]]
  O  <- data.i$state
  t  <- data.i$years
  x1 <- data.i$bage[1]
  x2 <- data.i$dage[1]
  T   <- diag(D)
  T.j <-T
  # First state (no MC):
  f <- c(1,0,0,0)
  # Rest of the states
  for(j in 2:length(O)){
   # Q matrix:
   Q[1,2] <- exp(beta0[1]+xi[1]*t[j-1]+beta1[1]*x1+beta2[1]*x2)
   Q[1,D] <- exp(beta0[2]+xi[2]*t[j-1]+beta1[2]*x1+beta2[2]*x2)
   Q[1,1] <- -(Q[1,2]+Q[1,D])
   Q[2,3] <- exp(beta0[3]+beta2[3]*x2)
   Q[2,D] <- exp(beta0[4]+beta2[4]*x2);  Q[2,2] <- -(Q[2,3]+Q[2,D])
   Q[3,D] <- exp(beta0[5]+beta2[5]*x2);  Q[3,3] <- -Q[3,D]
   # P matrix:
   P <- MatrixExp(mat=Q, t=t[j]-t[j-1])
   for(r in 1:D){
    for(s in 1:D){
     if(O[j]!=D){
       T.j[r,s] <- E[s,O[j]]*P[r,s]
     }else{
       T.j[r,s] <- P[r,s]*Q[s,D]
     }
    }
   }
   T <- T%*%T.j
  }
  loglik <- loglik+log(f%*%T%*%one)
  #cat("i = ",i, "and -2*Loglik = ", -2*loglik,"\n")
 }
 # Monitor:
 cat("-2*Loglik = ", -2*loglik,"\n")
 return(-loglik)
}

# Check with msm result:
cat("\nCheck loglikelihood function:\n")
cat("msm: -2loglik  =", model$minus2loglik,"\n")
E.msm <- ematrix.msm(misc.msm2, ci="delta")
e <- c(E.msm$estimates[1,2],E.msm$estimates[2,1],
     E.msm$estimates[2,3],E.msm$estimates[3,2])
p0 <- c(misc.msm2$estimates[c(1:5)],log(e/(1-e)),
      misc.msm2$estimates[c(6,7,11,12,16:20)])
cat("-2loglikelihood function at msm estimate =", 
     2*loglikelihood(p0),"\n\n")
 
# Maximise:
p0 <- p0 + 0.001
max<-optim(par=p0, fn=loglikelihood, method = c("BFGS"),
      control=list(maxit=3000, fnscale=2000),hessian=TRUE)
cat("\n-2loglik =", 2*max$value,"\n")
cat("AIC =", 2*max$value+2*length(max$par),"\n")
conv <- max$convergence
cat("Convergence code =", conv,"\n")
p    <- as.vector(max$par)
p.se <- as.vector(sqrt(diag(solve(max$hessian))))
print(cbind(p=round(p,digits),
            se=round(p.se,digits),"Wald ChiSq"=round((p/p.se)^2,digits),
            "Pr>ChiSq"=round(1-pchisq((p/p.se)^2,df=1),digits)),quote=FALSE)

# MC Matrix:
alpha <- p[6:9]
E <- diag(4)
E[1,2] <- logit(alpha[1]); E[1,1] <- 1-E[1,2]
E[2,1] <- logit(alpha[2]) 
E[2,3] <- logit(alpha[3]); E[2,2] <- 1-E[2,1]-E[2,3]
E[3,2] <- logit(alpha[4]); E[3,3] <- 1-E[3,2] 
cat("MC matrix:\n")
print(round(E,digits))
}