source("/Users/pm/s/first.s")
source("/Users/pm/s/rand.s")
dir <- "/Users/pm/book/ctrial/pm/3/354/"
setwd(dir)
library("bayesm")
library("adapt")
library("stats4")

##################################################################
# implements
# Thall et al (2003) 2-agent dose escalation
##################################################################
## Model:
## p(d) = { a1*x1^b1+a2*x2^b2+a3[x1^b1 x2^b2]^b3 }/ (1 + ...)

pistar <- 0.30 # target toxicity level
debug <- T     # indicator for additional debugging info

## grid for L1 ############################################
## x1 values = L1x
## x2 values = L1x*rho
L1x <- c(.12,.25,.4, .52, .61, .7, .78, .93, 1)
rho <- 1          # slope of line segment for stage 1

    
## grid for stage 2 ############################################
M <-  50      # grid length
grid <- seq(from=min(c(L1x,rho*L1x)),
            to=max(c(L1x,rho*L1x)), length=M)
## hyperprior parameters
phi <- rbind( ## moi.gamma(m,v) finds Ga(a,b) parameters
              ## matching mean=m=a/b, var=v=a/b^2
         moi.gamma(0.4286,0.1054),  # a1
         moi.gamma(7.6494,5.7145),  # b1
         moi.gamma(0.4286,0.0791),  # a2
         moi.gamma(7.8019,3.9933),  # b2
         moi.lognorm(0.25,3),       # a3
         moi.lognorm(0.25,3))       # b3
#         moi.gamma(  1,3),          # a3
#         moi.gamma(  1,3))          # b3

## MC sample ############################################
N <- 100
TH <- matrix(nrow=6, ncol=N) # MC sample


## truth, using parameters in f(.) below
scenario0 <- c(0.4286, 7.6494, 0.4286, 7.8019, 0,0.5)   # no interaction
scenario1 <- c(0, 7.6494, 0, 7.8019, 1, 0.05)           # only interaction
scenario2 <- c(0.4286, 7.6494, 0.4286, 7.8019, 1, 0.25) # weak interaction
scenario3 <- c(0.4286, 7.6494, 0.4286, 7.8019, 1, 0.1)  # some interaction
scenario9 <- c(0.4286, 7.6494, 0.4286, 7.8019, 1, 0.05) # prior mean
scenario4 <- c(1.5, 7, 2, 7, 1, 0.05) # high tox

## Model ############################################
pi <- function(x,th)
  { ## mean toxicity at dose pair x, for par vector th
    ## th can be a (6 x 1) vector or (6 x N) matrix of pars
    th <- as.matrix(th)
    xt1 <- x[1]^th[2,]
    xt2 <- x[2]^th[4,]
    N <- th[1,]*xt1 + th[3,]*xt2 + th[5,]*(xt1*xt2)^th[6,]
    D <- 1+N
    return(N/D)
  }

lprior <- function(th)
  { ## log prior, th can be a (6 x 1) vecotr or (6 x N) matrix
    ## of par values
    ## evaluate p(th[1..4],eta[5,6]) for eta[j]=log(th[j])
    ## i.e., a3,b3 are parmaetrized on log scale
    ##    (but conveniently stored on absolute scale for likelihood eval)
    th <- as.matrix(th)
    lpj <- rep(0,ncol(th))
    for(j in 1:4)  # gammas
      lpj <- lpj + dgamma(th[j,],shape=phi[j,1],rate=phi[j,2],log=T)
    for(j in 5:6){ # log normals p(th) = N(log(th);m,s)
      etaj <- log(th[j,])
      lpj <- lpj + dnorm(etaj, mean=phi[j,1],sd=phi[j,2],log=T)
    }
    return(lpj)
  }

lprior.etaj <- function(etaj,j)
  { ## log marginal prior p(eta[j]) for for etaj=log(theta[j])
    ## NOTE: prior p(th) is indep across j
    ## deta/dtheta = 1/theta => p(eta) = p(theta(eta))*theta
    if (j<5){
      thj <- exp(etaj)
      lpj <- dgamma(thj,shape=phi[j,1],rate=phi[j,2],log=T) +
        etaj # last term is ".. + log(theta)" for Jacobian
    } else 
      lpj <- dnorm(etaj,mean=phi[j,1],sd=phi[j,2],log=T) 
    return(lpj)
  }

llik <- function(th,y,x)
  { # log likelihood for par th, data (y,x)
    ## th can be a (6 x 1) vector or (6 x N) matrix
    th <- as.matrix(th)
    logl <- rep(0,ncol(th))  # initialize log likelihood
    n <- nrow(y)
    for (i in 1:n){
      pix <- pi(x[i,],th)
      logli <- y[i,1]*log(pix) + y[i,2]*log(1-pix)
      logl <- logl + logli
    }
    return(logl)
  }


## dose range
d0 <- 2         # starting dose on L1
N1 <- 20        # n patients in stage 1
N2 <- 40        # n patients in stage 2
K <-  3         # cohort size
th0 <- scenario9
## Trial simulation ############################################
sim.trial <- function(th0=scenario9,d0=2,K=3,N1=20,N2=40,
                      L2.unif=T)
  { ## simulate one realization of the trial under scenario th0
    ## th0     = simulation truth;
    ## d0      = starting dose in L1
    ## K       = cohort size
    ## N1, N2  = max sample size in 1st and 2nd stage
    ## L2.unif = use uniform random choice of dose from L2

    if (debug)
    cat("TRUTH  th=", th0,"\n")
    mcmc.init()
    
    n <-  0    # initialize data
    y <- NULL
    x <- NULL
    ## stage 1 ######################
    if (debug) cat("\n Stage 1: \n")
    init.L1()
    d <- d0
    refined <- F                      # indicator for split steps
    tried <- rep(F,nrow(L1))          # tried[j] = I(tried j-th dose)
    pibar <- rep(0,nrow(L1))          # initialization..
    repeat{
      xi <- L1[d,]
      tried[d] <- T
      p <- pi(xi, th0)
      yi <- rbinom(1,K,p)             # simulate response
      if (!refined & (yi>0)){         # toxicity observed
        tried <- refine.L1(d,tried)   # refine L1, and update tried..
        refined <- T
      }
      ph <- pibar[d]
      y <- rbind(y,c(yi,K-yi,p,ph))   # update data matrix
           ## last two cols of y report true and estimated p(tox)
      x <- rbind(x,xi)
      n <- n+K
      mcmc(y,x)                       # update posterior MCMC sample
      pibar <- posterior.L1(d)
      y[nrow(y),4] <- pibar[d]        # fill in estimated p(tox) 
      d <- dose.assign.L1(pibar,tried)# dose assignment for next cohort
        ## returns 0 to indicate stopping if lowest dose is too toxic
      if ( (n>N1) | (d==0) )
        break
    }
    K1 <- nrow(y)                     # record n cohorts in stage 1
    ## stage 2 ######################
    L2 <- M:1  ## initialize 0.30 quantile contour
               ## see comments in posterior.L2()
    if (d>0) { # did not stop for excess toxicity in stage 1
      if (debug) cat("\n Stage 2: \n")
      left <- T  # indicator for trying left / right
      repeat{
        left <- !left                   # switch for next time
        L2 <- posterior.L2(L2)
        xi <-  dose.assign.L2(L2,left,y=y,x=x,unif=L2.unif)
        ## dose assignment for next cohort
        if (n>N1+N2)
          break
        p <- pi(xi, th0)                # sim truth
        yi <- rbinom(1,K,p)             # simulate response
        ph <- mean(apply(TH,2,function(th){pi(xi,th)}))
              ## estimated toxicity (should be around pistar)
        y <- rbind(y,c(yi,K-yi,p,ph))   # update data matrix
              ## note that estimated p(tox)=.3 by definition..
        x <- rbind(x,xi)
        n <- n+K
        mcmc(y,x)                        # update posterior MCMC sample
      }# repeat
      ## final reco
      xleft <- dose.assign.L2(L2,left=T,y=y,x=x,unif=L2.unif)
      xright <- dose.assign.L2(L2,left=F,y=y,x=x,unif=L2.unif)
      d0 <- which.min(abs(L2-(1:M)))         # center
      xmiddle <- grid[c(d0,d0)]
    } else { # stopping for excess toxicity in stage 1
      xleft <- L1[1,]; xright <- L1[1,]; xmiddle <- L1[1,]
      if (debug) cat("\n Stopping L1 for excess toxicity.\n")
    }
    ci95 <- c(0.025,0.975)
    ci.left <- quantile(pi(xleft,TH),ci95)
    ci.right <- quantile(pi(xright,TH),ci95)
    ci.middle <- quantile(pi(xmiddle,TH),ci95)
    return(list(y=y, x=x, xstar=rbind(xleft,xmiddle,xright),
                cistar=rbind(ci.left,ci.middle,ci.right),
                L2=L2,K1=K1))
  }


## Manipulating grid for L1 #################################

init.L1 <- function()
  { ## sets up grid for stage 1
    L1 <<- cbind(L1x, rho*L1x)
    mult <<- rep(1,6) # factor for step size in MCMC M-H proposal
  }

refine.L1 <- function(d,tried)
  { ## add half points
    ## NOTE: variation to paper. Thall & al define
    ##    half points only beyond dose d, and *all*
    ##    dose combinations on L1 below d
    ## could easily be changed here (that is why we include d)
    nL <- nrow(L1)
    L1new <- L1[1,]
    tried.new <- tried[1]
    for (i in 1:(nL-1)){
      xnew  <- .5*(L1[i,]+L1[i+1,]) # new half point
      L1new <- rbind(L1new,xnew,L1[i+1,])
      tried.new <- c(tried.new,tried[i+1],tried[i+1])
    }
    L1 <<- L1new
    return(tried.new) # return updated tried for use in algorithm 
  }

pix1 <- function(th,d)
  { # aux function for post.L1() below
    ## toxicity at L1[d,]
    pi(L1[d,],th)
  }

posterior.L1 <- function(d)
  { ## posterior mean toxicities on L1
    ## start with current d and go up/down until we cross pistar
    ## (avoids evaluation of pibar for doses far away from pistar)
    pibar <- rep(-1, nrow(L1))
    if (debug) cat("try doses d= ")
    repeat{
      if (debug) cat(d," ")
      pibar[d] <-  mean(apply(TH,2,pix1,d))
      if (pibar[d]>pistar)
        d <- d-1
      else 
        d <- d+1
      if ( (d > nrow(L1)) | (d < 1) )
        break
      if (pibar[d] != -1) # already done this one..
        break
    }
    if (debug) cat("\n")
    return(pibar)
  }


dose.assign.L1 <- function(pibar,tried)
  { ## find optimal dose for next cohort

    i <- which.min(abs(pistar-pibar))
    if (i==1)
      return(i)
    if (tried[i-1])
      return(i)
    istar <- max(which(tried))+1
    return(istar)
  }


## Manipulating grid for L2 ############################################

pix2 <- function(th,i,j)
  { # aux function for post.L2() below
    ## toxicity at x=(grid_i,grid_j)
    pi(grid[c(i,j)], th)
  }

## ------------------------------------------------------------------


posterior.L2 <- function(L2)
  { ## grid for stage 2
    ## for each x1=grid[i], find 0.30 quantile L2[i] of tox prob
    ## i.e., (x=grid, y=L2) defines the grid for stage 2
    ## NOTE: in contrast to L1 the grid L2 reports only x2 indices
    nx <- length(grid)
    for(i in 1:nx){       # get 0.30 quantile of tox prob for x1=grid[i] 
      pibar <- rep(-1,nx)
      if (i==1)
        j <- L2[i]   # start at old value
      else
        j <- L2[i-1] # start at value of previous dose
      repeat{
        pd <-  mean(apply(TH,2, pix2, i,j))
        pibar[j] <- pd
        j0 <- j
        if ((pd > pistar) & (j==1))  # .30 quantile is below  x2=grid[1]
          break
        if ((pd < pistar) & (j==nx)) # .30 quantile is beyond x2=grid[nx]
          break
        if (pd > pistar) # go down, .30 quantile is lower
          j <- j-1
        else             # go up, quantile is higher
          j <- j+1
        if (pibar[j]> -1){  # already tried this one
          ## i.e., we switched back up/down
          ## check if last value was closer?
          if (abs(pibar[j0]-pistar) < abs(pibar[j]-pistar))
            j <-  j0
          break
        } # found it
      }
      L2[i] <-  j
    } # for i
    
    return(L2)
  }


    
dose.assign.L2 <- function(L2, left=F, y,x, unif=F)
  {## find optimal dose in L2 for next cohort
    nx <- length(grid)
    if (left) # assign on left half curve
      side <- (L2 >= rho*(1:nx)) # recall rho=slop of L1 line
    else
      side <- (L2 <  rho*(1:nx))
    if (unif){ # uniform assignment on L2[side,]
      i <- sample(which(side),1)
      xstar <- grid[c(i,L2[i])]
    } else {   # decision theoretic assignment
      ## evaluate K(x) and I(x) on grid points in L2
      Kxx <- rep(0,nx)
      Ixx <- rep(0,nx)
      for(i in which(side)){
        j <- L2[i]
        Kxx[i] <- Kx(i,j)
        Ixx[i] <- Ix(i,j,y,x)
      }
      Kxx[!side] <- min(Kxx)-1 # make sure it is not chosen below..
      Ixx[!side] <- min(Ixx)-1 # make sure it is not chosen below..
      iK <- which.max(Kxx)
      iI <- which.max(Ixx)
      xK <- grid[c(iK,L2[iK])]
      xI <- grid[c(iI,L2[iI])]
      x0 <- .5*(xK+xI)
      istar <- which.min( (grid    -x0[1])^2 +
                          (grid[L2]-x0[2])^2 )
               ## closest point on .30 quantile curve L2
      xstar <- grid[c(istar,L2[istar])]
    }
    return(xstar)
  }

Kx <- function(i,j,lambda=1)
  { ## cell killing criterion
    return(grid[i]+lambda*grid[j])
  }
    
    
I.xth <- function(th,i,j)
  { ## Fisher information log(det I(th,x)), before *expectation*
    ##        w.r.t. p(th | data)
    ## evaluated for x=grid(i,j)
    x    <- grid[c(i,j)]
    a <- th[c(1,3)]
    b <- th[c(2,4)]
    a3   <- th[5]; b3 <- th[6]

    xt   <- c(x[1]^b[1], x[2]^b[2])
    xt12 <- (xt[1]*xt[2])^b3
    num  <- a[1]*xt[1] + a[2]*xt[2] + a3*xt12
    f    <- rep(0,6)
    for(k in 1:2){ # partial deriv w.r.t a[k],b[k], k=1,2
      offset <- 2*(k-1)
      f[offset+1] <- xt[k]
      f[offset+2] <- log(x[k])*(a[k]*xt[k]+ a3*b3*xt12)
    }
    f[5] <- xt12
    f[6] <- a3*(log(xt[1])+log(xt[2]))*xt12
    r1 <- 1/(num*(1+num)^2)     # pull out the 1/[ N*(1+N)^2 ]
    r2 <- sum(f^2)              # I(th,x) = r1*r2
    return(r1*r2)
  }
    
Ix <- function(i,j,y,x)
  {## expected information
    Iij <- mean(apply(TH,2, I.xth, i, j))
    return(Iij)
  }
    

## MCMC ############################################

mcmc.init <-  function()
  { # initialize posterior MCMC sample from prior
    ## Not currently strictly used, except for (optional)
    ##     starting value for mle routine in mcmc
    ## But better to always have a well defined MC sample around
    N <-  ncol(TH) # MC sample size (=n parallel chains)
    for(j in 1:4)   # th[1..4] have gamma priors
      TH[j,] <<- rgamma(N,shape=phi[j,1],rate=phi[j,2])
    for(j in 5:6){  # log a3, log b3 have normal priors
        ## but saved on absolute scale for easy liklihood eval later
      etaj <- rnorm(N,mean=phi[j,1],sd=phi[j,2])
      TH[j,] <<- exp(etaj)
    }
  }

mcmc <- function(y,x, n.it=100)
  { # udpates posterior MCMC sample

    ## initialize with MLE 
    ## using old values as starting values
    m <- log( apply(TH,1,mean) )
    names(m) <- c("a1","b1","a2","b2","a3","b3")
    mlpost <- function(a1=a1,b1=b1,a2=a2,b2=b2,a3=a3,b3=b3)
      { # log posterior for eta=log(th)
        th <- exp( c(a1,b1,a2,b2,a3,b3) )
        return( -lprior(th) - llik(th,y,x) )
      }
    ml <- mle(mlpost,  start=as.list(m))
    etahat <- ml@coef    # mle for log(theta)
    V <- ml@vcov
    L <- t(chol(V))      # factorize V
    sd <- sqrt(diag(V))  # will use for step size below
    N <- ncol(TH) # MC sample size
    Z <- matrix(rnorm(6*N),nrow=6,ncol=N) # inititialize at mle
    ETA <- etahat + L %*% Z
    Th <- exp(ETA)
    
    step <- mult*sd     # step size for M-H proposal below
    sum.acc <- rep(0,6) # keep track of acceptance rates
    sum.try <- rep(0,6)


    for(it in 1:n.it){  # n.it MCMC iterations
      for(j in 1:6){    # transition probs changing th[1]..th[6]
        etaj0  <- ETA[j,]
        etaj1  <- rnorm(N,mean=etaj0,sd=step[j]) # M-H proposal
        thj1 <- exp(etaj1)
        A <-  -lprior.etaj(etaj0,j)-llik(Th,y,x)
                        ## A = log(acceptance prob) for candidate etaj1
        Th[j,] <- thj1
        A <- A + lprior.etaj(etaj1,j)+llik(Th,y,x)
        accept <- (log(runif(N))<A)
        ETA[j,] <- ifelse(accept, etaj1, etaj0)
        Th[j,] <-  exp(ETA[j,])
        sum.acc[j] <- sum.acc[j]+sum(accept)
        sum.try[j] <- sum.try[j]+N
      }# j
    }# i
    for (j in 1:6){    # update step size for M-H proposal if needed
      if (sum.acc[j]/sum.try[j] > 0.9){
        mult[j] <<- mult[j]*2
        if (debug) cat("\n increasing mult[",j,"]\n")
      } else if (sum.acc[j]/sum.try[j] < 0.1){
        mult[j] <<- mult[j]/2
        if (debug) cat("\n decreasing mult[",j,"]\n")
      }
    }
    m <- apply(Th,1,mean)
    sd <- sqrt(apply(Th,1,mean))
    if (debug) cat(" m=",m," sd=",sd,"\t p(acc)=",sum(sum.acc)/sum(sum.try),"\n")
    TH <<- exp(ETA)    # upudate global copy of MC sample
  }


## Plotting of results ############################################
plt.post <- function(y,x,k=NULL,add.pts=T,xstar=NULL,
                     L2=NULL)
  { ## plot posterior up to cohort k, i.e. line k in data matrix
    if (is.null(k))
      k <- nrow(y)
    mcmc(y[1:k,],x[1:k,],n.it=100)
    z <- matrix(0,nrow=M,ncol=M)
    for(i in 1:50)
      for(j in 1:50)
        z[i,j] <-  mean(apply(TH,2, function(th){pix2(th,i,j)}))
    contour(x=grid,y=grid,z,xlab="X",ylab="Y",bty="l")
    if(add.pts){ # add observed data points
      pch <- ifelse(y==0,1,
                    ifelse(y==1,10,
                           ifelse(y==2,20,19)))
      points(x[,1],x[,2],pch=pch)
      lines(L1[c(1,nrow(L1)),1],L1[c(1,nrow(L1)),2],type="l")
    }
    if (!is.null(xstar)) # add final design points
      points(xstar,    pch=8,cex=2,col=2)
    if (!is.null(L2))
      lines(grid,grid[L2],type="l",lwd=3,col=4)
      
  }

plt.S <- function(th0,xstar=NULL)
  {## plot simulation truth th0
    z <- matrix(0,nrow=M,ncol=M)
    for(i in 1:50)
      for(j in 1:50)
        z[i,j] <-  pix2(th0,i,j)
    contour(x=grid,y=grid,z,xlab="X",ylab="Y",bty="l")
    if (!is.null(xstar)) # add final design points
      points(xstar,    pch=8,cex=2,col=2)
  }

plt.yp <- function(y,x,xstar=NULL,cistar=NULL)
  {
    nk <- nrow(y)
    p <- y[,3]
    pbar <- y[,4]
    ymx <- max(cistar,p)
    plot(1:nk,pbar,type="l",xlab="COHORT",ylab="P(TOX)",bty="l",
         ylim=c(0,ymx),xlim=c(1,nk+3))
    points(1:nk,p,pch=19)
    if (!is.null(xstar)){
      points((nk+1):(nk+3),rep(pistar,3),type="p",pch=8,col=2)
      for(i in 1:3)
        lines(c(nk+i,nk+i),cistar[i,],type="l")
    }
  }
                   
example <- function()
  { ## execute these lines for example
    st0  <<- sim.trial(th0=scenario0,L2.unif=F)
    st1  <<- sim.trial(th0=scenario1,L2.unif=F)
    st2  <<- sim.trial(th0=scenario2,L2.unif=F)
    st2u <<- sim.trial(th0=scenario2,L2.unif=T)
    st3  <<- sim.trial(th0=scenario3,L2.unif=F)
    st3u <<- sim.trial(th0=scenario3,L2.unif=T)
    st4  <<- sim.trial(th0=scenario4,L2.unif=F)
    st9  <<- sim.trial(th0=scenario9,L2.unif=F)
    st9u <<- sim.trial(th0=scenario9,L2.unif=T)

    ex.plots()
  }

ex.plots <- function()
  {
    ps("st0")
    plt.post(st0$y,st0$x,add.pts=T,xstar=st0$xstar,L2=st0$L2)
    devoff()
    ps("st1")
    plt.post(st1$y,st1$x,add.pts=T,xstar=st1$xstar,L2=st1$L2)
    devoff()
    ps("st2")
    plt.post(st2$y,st2$x,add.pts=T,xstar=st2$xstar,L2=st2$L2)
    devoff()
    ps("st2u")
    plt.post(st2u$y,st2u$x,add.pts=T,xstar=st2u$xstar,L2=st2u$L2)
    devoff()
    ps("st3")
    plt.post(st3$y,st3$x,add.pts=T,xstar=st3$xstar,L2=st3$L2)
    devoff()
    ps("st3u")
    plt.post(st3u$y,st3u$x,add.pts=T,xstar=st3u$xstar,L2=st3u$L2)
    devoff()
    ps("st4")
    plt.post(st4$y,st4$x,add.pts=T,xstar=st4$xstar,L2=st4$L2)
    devoff()
    ps("st9")
    plt.post(st9$y,st9$x,add.pts=T,xstar=st9$xstar,L2=st9$L2)
    devoff()
    ps("st9u")
    plt.post(st9u$y,st9u$x,add.pts=T,xstar=st9u$xstar,L2=st9u$L2)
    devoff()

    
    ps("th0")
    plt.S(scenario0,st0$xstar)
    devoff()
    ps("th1")
    plt.S(scenario1,st1$xstar)
    devoff()
    ps("th2")
    plt.S(scenario2,st2$xstar)
    devoff()
    ps("th2u")
    plt.S(scenario2,st2u$xstar)
    devoff()
    ps("th3")
    plt.S(scenario3,st3$xstar)
    devoff()
    ps("th3u")
    plt.S(scenario3,st3u$xstar)
    devoff()
    ps("th4")
    plt.S(scenario4,st4$xstar)
    devoff()
    ps("th9")
    plt.S(scenario9,st9$xstar)
    devoff()
    ps("th9u")
    plt.S(scenario9,st9u$xstar)
    devoff()
    
    ps("p0")
    plt.yp(st0$y,st0$x,st0$xstar,st0$cistar)
    devoff()
    ps("p1")
    plt.yp(st1$y,st1$x,st1$xstar,st1$cistar)
    devoff()
    ps("p2")
    plt.yp(st2$y,st2$x,st2$xstar,st2$cistar)
    devoff()
    ps("p3")
    plt.yp(st3$y,st3$x,st3$xstar,st3$cistar)
    devoff()
    ps("p4")
    plt.yp(st4$y,st4$x,st4$xstar,st4$cistar)
    devoff()
    ps("p9")
    plt.yp(st9$y,st9$x,st9$xstar,st9$cistar)
    devoff()
  }