library(survival)
source('sensitivity/R/sensitivitySGD_followup.R')

rnewdist2 <- function(N, a, b, tau, beta, alpha) {
  ## If
  ##           - beta*I(beta < 0)*tau - alpha
  ##     Pi*(%e                               + 1)
  ## c = -----------------------------------------
  ##                         p
  ##
  ## and
  ##
  ##           - beta*min(y, tau) - alpha
  ##     Pi*(%e                           + 1)
  ## r = -------------------------------------
  ##                      c*p
  ##
  ## If c is substituted into r then
  ##
  ##          - beta*min(y, tau) - alpha
  ##        %e                           + 1
  ##  r = ------------------------------------
  ##        - beta*I(beta < 0)*tau - alpha
  ##      %e                               + 1
  ##

  
  ## preallocate vectors
  y <- r <- numeric(1000L)
  y.min <- accept <- logical(1000L)
  dist <- numeric(N)
  count <- 0L

  ## calculate the denominator of r and the value of r when y >= tau
  if(beta < 0) {
    r.denom <- 1 + exp(-alpha - beta * tau)
    r.y.max <- 1
  } else {
    r.denom <- 1 + exp(-alpha)
    r.y.max <- (1 + exp(-alpha - beta * tau))/r.denom
  }

  ## Loop until specified number of values have been collected
  repeat {
    y[] <- rweibull(1000L, a, b)

    y.min[] <- y < tau
    r[y.min] <- (1+exp(-alpha-beta*y[y.min]))/r.denom
    r[!y.min] <- r.y.max
    ## generate acceptace boolean vector for y and subset
    y.accept <- y[as.logical(rbinom(1000L, 1L, r))]
    y.len <- length(y.accept)

    if(y.len > N)
      dist[seq.int(to=N, length.out=N)] <- y.accept[seq_len(length.out=N)]
    else
      dist[seq(to=N, length.out=y.len)] <- y.accept

    if(y.len >= N)
      break
    
    N <- N - y.len
  }

  return(dist)
}


## Function to generate a sample set of length N
generate.sample <- function(N, a, b, tau, psi,
                            p0, beta0, alpha0,
                            p1, beta1, alpha1,
                            event=FALSE, inf.time) {
  N <- as.integer(N)
  
  if(event == TRUE) {
    s <- NULL
    z <- NULL
    N.samp <- trunc(ceiling(N/(p0 + p1) * 1.05))
    while(sum(s) < N) {
      new.z <- matrix(rep(c(FALSE,TRUE), each=N), ncol=2L)
      mat <- matrix(as.logical(rbinom(length(new.z), 1L, ifelse(new.z, p1, p0))), ncol=2L)

      rsum <- rowSums(mat)

      nTotal <- sum(s) + cumsum(rsum)
      if(any(nTotal >= N)) {
        loc <- which(nTotal >= N)[1L]

        indx <- seq_len(loc)
        z <- c(z, new.z[indx,])
        s <- c(s, mat[indx,])
      } else {
        z <- c(z, new.z)
        s <- c(s, mat)
      }
    }

    N <- length(z)
    N0 <- sum(!z)
    N1 <- N - N0
  } else {
    N0 <- as.integer(trunc(floor(N/2L)))
    N1 <- N - N0

    ## assign half of group to placebo and the rest to treatment
    z <- rep(c(FALSE,TRUE), each=c(N0,N1))

    ## expoential disrib on length of follow up for each patient
    u <- rexp(N, rate=ifelse(z, 0.0014248137, 0.002926681))

    ## time until censoring
    c<-runif(N,0L,100L)

    ## observed follow-up time
    v<-ifelse(u<c,u,c)    

    ## were they infected during follow-up (1=yes)
    s.star<-ifelse(u<c,TRUE,FALSE)

    s <- ifelse(v>=inf.time | !s.star, FALSE, TRUE)
  }
  
  n0 <- sum(s[!z], na.rm=TRUE)
  n1 <- sum(s[z], na.rm=TRUE)

  t <- rep(NA, N)

  ## for those with z = FALSE and s = TRUE generate Y from a modified wyble distribution
  t[!z&s] <- rnewdist2(N=n0, a=a, b=b, tau=tau,
                       beta=beta0, alpha=alpha0)

  t[z&s] <- rnewdist2(N=n1, a=a, b=b, tau=tau,
                      beta=beta1, alpha=alpha1)
  
  c <- pmin(24L, runif(N, 0L, 100L))

  y <- pmin(c, t)

  d <- y == t

  return(list(z=z,s=s,d=d,y=y, v=v, inf.time=inf.time))
}

calc.sen <- function(data.set, tau, psi, beta0, beta1, time.points, N.boot,
                     Nevent) {
  cat('.',file='')
  ans <- sensitivitySGDFollowup(z=data.set$z, y=data.set$y, v=data.set$v,
                                s=data.set$s, d=data.set$d,tau=tau,
                                psi=psi, followup.time=data.set$inf.time,
                                beta0=beta0, beta1=beta1,
                                time.points=time.points,
                                selection=TRUE, groupings=c(FALSE,TRUE),
                                ci.method='boot',
                                empty.principle.stratum=c(FALSE,TRUE),
                                trigger=TRUE,
                                N.boot=N.boot, ci=0.95,
                                oneSidedTest=TRUE, twoSidedTest=TRUE)
  ## return(c(unlist(ans),coverage=ifelse(ans[["ci.SCE.min"]] > ans[["SCE"]] | ans[["ci.SCE.max"]] < ans[["SCE"]], 0, 1)))
  results <- list(coverage.ci = ifelse(ans[["SCE.ci"]][1L,,,,] <= 0 & ans[["SCE.ci"]][2L,,,,] >= 0, 1L, 0L),
                  coverage.var=ifelse(ans[["SCE"]] + sqrt(ans[["SCE.var"]])*qnorm(0.975) >= 0L & ans[["SCE"]] + sqrt(ans[["SCE.var"]])*qnorm(0.025) <= 0L, 1L, 0L))
  return(c(lapply(ans, drop), lapply(results, drop),
           ActualBootReps=attr(ans, 'ActualBootReps')))
}

my.sapply <- function(X, FUN) {
  FUN <- match.fun(FUN)
  answer <- lapply(X, FUN)

  if(length(answer) && length(common.len <- unique(unlist(lapply(answer,
        length)))) == 1L) {
    if(common.len == 1L)
      unlist(answer, recursive=FALSE)
    else if(common.len > 1L)
      matrix(unlist(answer, recursive=FALSE), byrow=TRUE, ncol=common.len, nrow=length(X),
             dimnames=if (!(is.null(n1 <- names(answer[[1L]])) &
               is.null(n2 <- names(answer))))
               list(n2,n1)
             )
    else answer
  } else answer
}

my.replicate <- function(n, expr, simplify=TRUE) {
  do.call(mapply, c(lapply(integer(n), eval.parent(substitute(function(...) expr))),
          FUN=rbind, SIMPLIFY=FALSE))
}

tointegrate<-function(x, beta, tau, a, b) {
  (a/b)*(x/b)^(a-1)*exp(-beta*pmin(x,tau)-(x/b)^a)
}

calc.alpha <- function(p, beta, Pi, tau, a, b) {
  int.val <- integrate(tointegrate, 0, 24, beta=beta, tau=tau, a=a, b=b,
                       subdivisions=2^25, rel.tol=.Machine$double.eps^0.5)$value
  ans <- -log((p-Pi)/(Pi*(int.val + exp(tau * -beta) * (1 - pweibull(tau, a, b)))))
  return(ans)
}

args <- commandArgs(TRUE)
project <- args[1L]
comp.num <- args[2L]
setwd(args[3L])
load(paste(project, "_input_params_", comp.num, ".rda", sep=''))
output <- paste(project, "_output_slave", comp.num, '.txt', sep='')

cat("begining calculation run at ", format(Sys.time()), "\n", sep='', file=output)
cat("\n", append=TRUE, file=output)


calc.result <- function(varName, p0, p1, beta0, beta1, N, psi, alpha0, alpha1,
                        nRep, nBoot, event=FALSE, inf.time) {
  if(file.exists(paste(varName, 'rda', sep='.'))) {
    load(file=paste(varName, 'rda', sep='.'), envir=.GlobalEnv)
  } else {
    title <- paste(ifelse(event, "Events", "Records"), " = ", N, ", p0 = ", p0, ", p1 = ", p1,
                   ", beta0 = ", beta0, ", beta1 = ", beta1, " Num replications = ", nRep, " Num bootstraps = ", nBoot, "\n", sep='')
    cat("Calculating results for ", title, "\n", sep='')
    cat("Calculating results for ", title, "\n", sep='', file=output, append=TRUE)

    assign(varName,
           lapply(my.replicate(nRep,
                               calc.sen(generate.sample(N=N,
                                                        p0=p0,
                                                        p1=p1,
                                                        a=a,
                                                        b=b,
                                                        psi=psi,
                                                        tau=tau,
                                                        alpha0=alpha0,
                                                        alpha1=alpha1,
                                                        beta0=beta0,
                                                        beta1=beta1,
                                                        event=event,
                                                        inf.time=inf.time),
                                        beta0=beta0, beta1=beta1, psi=psi,
                                        tau=tau, time.points=tau, N.boot=nBoot)),
                  FUN=drop),
           envir = .GlobalEnv)
    
    cat(" saving ... ")
    save(list=varName, file=paste(varName, 'rda', sep='.'))
    cat("done\n")
  }
  get(varName, envir=.GlobalEnv)
}

do.call(mapply, args=c(FUN=calc.result, dat, list(SIMPLIFY=FALSE)))
warnings()
cat("Processing Done\n")