##
## Huang & al (2009) design
##

source("/Users/pm/s/first.s")
dir <- "/Users/pm/book/ctrial/pm/4/"
setwd(dir)

require("bayesm")

##################################################################
# initialize
##################################################################


## prior
beta <- c(40,300,750,1100) ## mu_xk ~ IG(a_k,b_k)
alpha <-  rep(11,4)
gamma <- rep(0.5,4)        ## (p_x1..p_x4) ~ Dir(gamma_1..gamma_4)
alphaC <- 2                ## prior for "common" design - for comparison
betaC <- 60
prior <- list(beta=rbind(beta,beta),alpha=rbind(alpha,alpha),
              gamma=rbind(gamma,gamma),
              alphaC=c(alphaC,alphaC), betaC=c(betaC,betaC))

## assumed scenarios (simulation truth)
## scenario 3
p3 <- rbind(c(.2, .4, .1, .3),   ## p_x1..p_x4, x=0,1
            c(.1, .1, .2, .6))
m3 <- rbind(c(4,30,75,110),      ## mu_x1..mu_x4, x=0,1
            c(6,45,112,165))

## design paramteres
nmax <- 120   # max n patients
mxweek <- 40  # n weeks beyond last patient
k <-  1       # n patients recruited per week
pL <- 0.025   # thresholds for early stopping
pU <- 1-pL
M0 <- matrix(0,nrow=2,ncol=4)  # 0-matrix for summaries

##################################################################
# Simulation
##################################################################

p0 <- p3; m0 <- m3; nmax <- 120; mxweek <- 160
sim.trial <- function(nmax=120,mxweek=160,p0=p3, m0=m3, comm=F)
  { ## simulate one realization of the trial
    ## nmax:   max n patients
    ## mxweek: max n weeks (assuming 1 pat/week, mxweek >= nmax
    ## p0:     simulation truth for resposne
    ## m0:     simulation truth for PFS (exponential mean pars)
    ## comm:   indicator for using common design
    
    ## initialize data
    dta <- list(n=0,        # sample size
                ## dta${S,delta,t} are summary counts and sums
                S=M0,       # response counts
                delta=M0,   # censoring indicators
                t=M0,       # censoring (if delta=0) or obs (delta=1) times
                ## all dta$XXi are vectors with XX[i] for the i-th patient
                Si=NULL,    # vector of responses (1..n)
                Ti=NULL,    # vector of obs times
                deltai=NULL,# vector of censoring indicators
                t0i=NULL,   # vector of recruitment times
                xi=NULL,    # treatment allocation
                pi=NULL,    # alloc prob at time t=1...n,...tmx
                muxi=NULL,
                piC=NULL,   # same for "common design" (for plotting)
                muxiC=NULL)
                
    post <- list(g=gamma,a=alpha,b=beta,p=.5,
                 mux=c(0,0)) # posterior parameters
                 ## p(p_x | y)  = Dir(g[x,]),      x=0,1
                 ## p(mu_x | y) = IG(a[x,],b[x,]), x=0,1
    week  <- 0     # running calendar time in months
    
    
    repeat{  ## loop over weeks
      post <- post.update(dta,post)
      dta$pi <- c(dta$pi,post$p)     # record p for plotting
      dta$muxi <- rbind(dta$muxi,post$mux)
      dta$pCi <- c(dta$pCi,post$pC)  # record p for plotting
      dta$muxCi <- rbind(dta$muxCi,post$muxC)
      if (comm)
        p <- post$pC
      else
        p <- post$p
      if ( (p < pL) | (p > pU) ){
        if (p<pL)
          cat(" Stopping for inferiority. \n")
        else
          cat(" Stopping for superiority. \n")
        break
      }
      week <- week+1
      if (week > mxweek){
        cat(" Stopping for max time. \n")
        break
      }
      if (dta$n < nmax)
        dta <- next.cohort(week,post$p,p0,m0,dta)
      dta <- dta.summaries(week,dta)
      cat(" Week ", week, ", mua,mub,p=",c(post$mux,p),"\n")
    }
    return(list(dta=dta,week=week,post=post))
  }

                      
next.cohort <- function(week,p,p0,m0,dta)
  { ## generate next batch of k patients
    ## week: calendar time (in weeks)
    ## p:  alloc prob for trtmt 1 ("A")
    ## p0: simulation truth for p(S=j | x), j=1..4, x=1,2
    ## m0: sim truth for p(T | S=j,x) = Exp(1/mu_xj)

    ## treatment assignments
    u <- runif(k)                  # treatment assignment 
    xi <- ifelse(u<p,1,2)          #    p(x=A)=p
    ## responses for next k patients
    S0 <- sample(1:4,k,replace=T,prob=p0[1,]) # Sj | x=A
    S1 <- sample(1:4,k,replace=T,prob=p0[2,]) # Sj | x=B
    Si <- ifelse(xi==1,S0,S1)
    ## PFS for next k patients
    Ti <- rep(0,k)                 # PFS times, initialize
    for(i in 1:k){
      x <- xi[i]
      S <- Si[i]
      Tm <- rexp(1,1/m0[x,S])      # PFS times, generate
      Ti[i] <- Tm
    }
    ## record new data
    dta$n <- dta$n + k             # increment sample size
    dta$Ti <- c(dta$Ti,Ti)         # PFS times, save
    dta$Si <- c(dta$Si,Si)
    dta$deltai <- c(dta$deltai,rep(0,k)) # censoring indicators
    dta$t0i <- c(dta$t0i,rep(week,k))  # record enrollment time
    dta$xi <- c(dta$xi,xi)
    return(dta)
  }

dta.summaries <- function(week,dta)
  { # update data summaries
    # NOTE: called without additional recruitment for tm=121..160
    #       therefore separate from next.cohort
    dta$t <- M0
    dta$delta <- M0
    dta$S <- M0
    for(i in 1:dta$n){
      x <- dta$xi[i]
      S <- dta$Si[i]
      if (week > dta$t0i[i]+dta$Ti[i]){ # fully observed
        dta$deltai[i] <- 1
        Tm <- dta$Ti[i]
      } else 
        Tm <- (week-dta$t0i[i])
      dta$S[x,S] <- dta$S[x,S]+1    # update counts for responses
      dta$t[x,S] <- dta$t[x,S]+Tm   # update sum
      dta$delta[x,S] <-  dta$delta[x,S]+dta$deltai[i] # 
    }
    return(dta)
  }

M <- 1000
post.update <- function(dta,post,M=1000)
  { ## posterior updating
    ## M: Monte Carlo sample size to compute p 
    post$g <- prior$gamma + dta$S
    post$a <- prior$alpha + dta$delta
    post$b <- prior$beta  + dta$t
    ## common design (for comparison)
    post$aC <- prior$alphaC + apply(dta$delta,1,sum)
    post$bC <- prior$betaC  + apply(dta$t,    1,sum)
    
    ## compute p = Pr(mu_a > mu_b | y)
    ma <- rep(0,M)
    mb <- rep(0,M)
    maC <- rep(0,M) # same for common design
    mbC <- rep(0,M) #
    for(i in 1:M){
      px1 <- rdirichlet(post$g[1,])
      px2 <- rdirichlet(post$g[2,])
      lambda1 <- rgamma(4,post$a[1,],rate=post$b[1,])
      lambda2 <- rgamma(4,post$a[2,],rate=post$b[2,])
      mu1 <- 1/lambda1 ## IG(a1,b1)
      mu2 <- 1/lambda2 ## IG(a2,b2)
      ma[i] <- px1 %*% mu1
      mb[i] <- px2 %*% mu2
      ## same for common design
      lambdaC <- rgamma(2,post$aC,rate=post$bC)
      muC <- 1/lambdaC ## IG(aC,bC)
      maC[i] <- muC[1]
      mbC[i] <- muC[2]
    }
    post$p <- sum(ma>mb)/M
    post$mux <- c(mean(ma),mean(mb))
    post$pC <- sum(maC>mbC)/M
    post$muxC <- c(mean(maC),mean(mbC))
    return(post)
  }


##################################################################
# example run
##################################################################

plt.p <- function(trial,pc=T)
  { ##
    pi <- trial$dta$pi
    pCi <- trial$dta$pCi
    wk <- 1:length(pi)
    plot(wk,pi,type="l",bty="l",xlab="WEEK",ylab="p(mu[A] > mu[B] | y)")
    if (pc)
      lines(wk,pCi,type="l",lty=2)
}    

plt.mx <- function(trial,p0,m0,comm=T,yA= -5)
  { ## yA: offset for p
    muA <- p0[1,]%*% m0[1,]
    muB <- p0[2,]%*% m0[2,]
    mx <- trial$dta$muxi
    mxC <- trial$dta$muxCi
    y1 <- max(c(mx,muA,muB))+5
    wk <- 1:length(trial$dta$pi)
    matplot(wk,mx,type="l",bty="l",xlab="WEEK",
            ylab="E(mu[A], mu[B] | y)", ylim=c(0,y1))

    abline(h=c(muA,muB),col=c(1,2),lty=c(1,2))
    text(0,muB+4,"mu[B] (truth)",adj=0,col=2)
    text(max(wk),muA+yA,"mu[A] (truth)",adj=1)
    if (comm)
      matlines(wk,mxC,type="l",bty="l",lwd=2)
  }    


example <- function()
  {# run this as example
    trial <- sim.trial(nmax,mxweek,p3,m3)

    ps("443-p")
    ## plt.p(trial,pc=F)
    plt.p(trial,pc=T)
    devoff()
    ps("443-mux")
    plt.mx(trial,p3,m3,comm=F)
    devoff()
    ps("443-muxC")
    plt.mx(trial,p3,m3,comm=T,yA=5)
    devoff()
    
    trialC <- sim.trial(nmax,mxweek,p3,m3,comm=T)
  }