rm(list = ls()) # Remove all defined objects



# The following three packages are required to use more than one core


library(bindata)
library(blockrand)
library(data.table)
library(plyr)
library(sqldf)



##################################
## subject allocation
##################################
falloc = function(n, c, round, bssr){
  if (missing(bssr)){ bssr = 0 }
  # equal allocation if nothing else specified
  alloc = rep(1,c)/c
  # equal allocation
  if (round=="exact"){
    corder = runif(n=c)
    rorder = rank(corder)
    rest = n %% c
    vec = (rorder <= rest)
    qest = ceiling(n * alloc)*vec + floor(n * alloc)*(1-vec)
  }
  # unequal 1 allocation with equal probabilities
  if ((round=="multinom")&(bssr==0)){
    # same rate for each centre
    rate = n / c
    # distribute (n-c) subjects to c centres with multinom
    # add single subject to each centre to avoid 0
    qest = as.vector(rmultinom(1, (n-c), rep(rate, c))+1)
  }
  # unequal 2 allocation with random probabilities
  if ((round=="multinom2")&(bssr==0)){
    # random rate for each centre
    p0 = runif(c, min=0, max=1)
    psum = sum(p0)
    rate = p0 / psum
    # distribute (n-c) subjects to c centres with multinom
    # add single subject to each centre to avoid 0
    qest = as.vector(rmultinom(1, (n-c), rate)+1)
  }
  return(qest)
}



##################################
## Moments for block randomisation
##################################
block = function(b, balloc){
  if (missing(balloc)){ balloc = 1 } 
  if ( (b %% (balloc+1)) >0 ){ 
    b = 2
    balloc = 1
    print("balloc must be mult. factor of block length.")
    print("Returned values refer to b=2, balloc=1.")
  } 
  # create set of all randomisation tuples with block length b
  # for 2 treatment groups A, B and allocation parameter balloc
  PI = df1 = data.frame(trt1=c("A", "B"))
  # create all possible tuples with treatments A and B
  for (i in (2:b )){
    PI =  sqldf(paste("select PI.*, df1.trt1 as trt", i, " from PI,  df1" , sep=""))
  }
  # count number of treatments A and B per tuple
  PI$a=apply(X=PI,1,FUN=function(x) length(which(x=='A')))
  PI$b=apply(X=PI,1,FUN=function(x) length(which(x=='B')))
  PI$n = PI$a + PI$b
  # reduce data set to those tupels matching 1-to-balloc restiction
  PI = PI[which(PI$a==b/(1+balloc)),]
  # create data frame for summary statistics
  df = data.frame(id = (1:(b)), mdiff=NA, mdiffk = NA, mdiff2k=NA)
  # calculate summary statistics for each number of subjects randomized within the "last" block
  for (i in (2:(b+1))){
    if (i==(b+1)){ 
      PIsum = PI 
    } else { PIsum = PI[-(i:(b+1))] }
    PIsum$nA=apply(X=PIsum ,1,FUN=function(x) length(which(x=='A')))
    PIsum$nB=apply(X=PIsum ,1,FUN=function(x) length(which(x=='B')))
    PIsum$sum = PIsum$nA + PIsum$nB
    PIsum$diff = abs(PIsum$nA - PIsum$nB)
    PIsum$diffk = abs(PIsum$nA - PIsum$nB/balloc)
    PIsum$diff2k = (PIsum$nA - PIsum$nB/balloc)^2
    # calculate summary statistics
    df$mdiff[(i-1)] = mean(PIsum$diff)
    df$mdiffk[(i-1)] = mean(PIsum$diffk)
    df$mdiff2k[(i-1)] =  mean(PIsum$diff2k)
  }
  result = list(df=df, combo = PIsum)
  return(result)
}


##################################
## Sample size calculation
##################################
samsi = function(approx, alpha, beta, mu, sigma, sigma.c, c, b, df_b, balloc){
  # calculate quantiles
  if (missing(alpha)){ alpha = 0.05 }
  if (missing(beta)){ beta = 0.2 }
  al = 1-alpha/2
  be = 1-beta
  t.a = qnorm(p=al)
  t.b = qnorm(p=be)
  # calculate summary statistics of block randomisation
  if (missing(balloc)){ balloc = 1 }
  if (missing(df_b)){ 
    delta2 = block(b=b, balloc=balloc)[[1]]
  } else delta2 = df_b[[1]]
  erw2 = mean(delta2$mdiff2k)
  erw_max = max(delta2$mdiff2k)
  
  ##################################
  # sample size formulas
  if (approx == "nv_uni"){
    # lower boundary
    n = ceiling(  (balloc+1)^2 / (balloc) * sigma^2 * (t.a+t.b)^2 / mu^2 )
  } else if (approx == "nv_slope"){
    # new formula equal centres
    n_rk = rep(NA, b)
    rk = delta2$id
    for (i in 1:b){
      # calculate expected imbalance for rk == i
      berw = delta2$mdiff2k[i]
      # calculate sample size for each rk
      n_rk[i] = ceiling(  (balloc+1)^2 / (2*balloc) * sigma^2 * (t.a+t.b)^2 / mu^2  + 
                            abs( sqrt(  (balloc+1)^2 * sigma.c^2 * c * berw + (balloc+1)^4 / (4*balloc^2) *
                                          sigma^4 * (t.a+t.b)^2 / mu^2) * (t.a+t.b) / mu )  )
    }
    # find optimal sample size according to min (n/c) mod b
    n = ceiling(mean(n_rk[((round(n_rk/c) %% b - 1:b)^2 == min((round(n_rk/c) %% b - 1:b)^2))]))
  } else if (approx == "nv_ewb"){
    # new formula unequal centres
    n = ceiling(  (balloc+1)^2 / (2*balloc) * sigma^2 * (t.a+t.b)^2 / mu^2  + 
                    abs( sqrt(  (balloc+1)^2 * sigma.c^2 * c * erw2 + (balloc+1)^4 / (4*balloc^2) *
                                  sigma^4 * (t.a+t.b)^2 / mu^2) * (t.a+t.b) / mu )  )
  } else if (approx == "nv_max"){
    # new formula upper boundary
    n = ceiling(  (balloc+1)^2 / (2*balloc) * sigma^2 * (t.a+t.b)^2 / mu^2  + 
                    abs( sqrt(  (balloc+1)^2 * sigma.c^2 * c * erw_max + (balloc+1)^4 / (4*balloc^2) *
                                  sigma^4 * (t.a+t.b)^2 / mu^2) * (t.a+t.b) / mu )  )
  }
  return(n)
}


##################################
## Funktion zur Powerberechnung
##################################
power.fkt = function(N, c, round, b, balloc, alpha, beta, 
                     mu0, mu, sigma, sigma.init, sigma.c, sigma.c.init, 
                     nsim, approx, df_b){
  
  start = proc.time()[3]
  
  # whatch out for missing values
  if (missing(N))         { N = 0 }
  if (missing(balloc))    { balloc = 1 } 
  if (missing(df_b)){ df_b = block(b=b, balloc=balloc) } 
  
  if (missing(sigma.init))  { sigma.init   = sigma}
  if (missing(sigma.c.init)){ sigma.c.init = sigma.c}
  if (missing(mu0))         { mu0 = 0 }
  
  dist = "normal"
  
  
  ##################################
  ## Start initial sample size calculation
  if (N==0){
    if (mu == 0){
      n = 100
      print("Unter H0 bitte Fallzahl spezifizieren, n wird auf 100 gesetzt")
    } else {
      n=samsi(approx=approx, alpha=alpha, beta=beta, mu=mu, sigma=sigma.init, sigma.c=sigma.c.init, c=c, b=b, df_b=df_b, balloc=balloc)
    }
  } else if (N>0){ n=N }
  
  ## End initial sample size calculation 
  ##################################
  
  
  ##################################
  ## Start data generation
  
  # store simulation run results
  simrun= data.frame(nsim=1:nsim, c=c, b=b, balloc=balloc, alpha=alpha, power=(1-beta), dist=dist, 
                     mu0 = mu0, mu = mu, sigma2 = (sigma^2), sigma2.init = (sigma.init^2), sigma2.c= (sigma.c^2), sigma2.c.init = (sigma.c.init^2),
                     approx=approx, round=round)
  
  ##################################
  ## start nsim-slope
  
  for (j in 1:nsim){
    # create random effects on centres
    dframe = data.frame(n.init=NA, n.bssr=NA, stratum=as.factor(1:c), sig.c = sigma.c, sc =NA, sd=sigma)
    dframe$sc = rnorm(n=c, mean=0, sd=sigma.c)
    
    # subject allocation (n_1, ..., n_c)
    # initial subject allocation
    qest.init = falloc(n=n, c=c, round=round, bssr=0)
    cqest = cumsum(qest.init)
    dframe$n.init = qest.init
    
    # create more complete randomisation blocks than needed
    qest.complete = 1 * ceiling(qest.init/b) * b
    scqest1 = cumsum(qest.complete)
    # indicator function for start and end-positions of centres in complete list
    qest.complete.a = scqest1 - qest.complete + 1
    qest.complete.e = scqest1 
    
    # indicator function for start and end-positions for centres of initially recruited subjects
    qest.init.a = qest.complete.a
    qest.init.e = qest.init.a + qest.init -1
    
    # create block-randomisation for all subjects of qest.complete
    n.complete = sum(qest.complete)
    dframe.complete =   blockrand(n=n.complete, id.prefix=, block.sizes=c(b/(1+balloc),b/(1+balloc)), num.levels=(1+balloc), block.prefix=)
    dframe.complete$stratum = NA
    dframe.complete$init = NA
    dframe.complete$final = NA
    
    ##################################
    ## indicator functions for subsets
    
    ## ... for initial data
    ind.init = rep(FALSE, n.complete) 
    for (i in 1:c){
      # add centre id
      dframe.complete$stratum[  qest.complete.a[i] : qest.complete.e[i]  ] = i
      # indicator which subjects are initially recruited (N_0)
      dframe.complete$init[  qest.complete.a[i] : qest.init.e[i]  ] = 1
      ind.init[ qest.complete.a[i] : qest.init.e[i] ] = TRUE 
    }
    
    ##################################
    ## create data
    
    data0 = data.frame()
    data0 = merge(x=dframe.complete, y=dframe, by="stratum")
    data0$mu = (data0$treatment=="A")*mu
    
    # create standard gaussian errors
    data0$y.sig = sigma * rnorm(n=n.complete, mean=0, sd=1) 
      
    # create observations based on the model
    data0$outcome = data0$mu + data0$sc + data0$y.sig
    data0$stratum = as.factor(data0$stratum)
    
    ## reduce complete data to
    # ... initial data
    data.init = data0[ind.init,]
    n.init = length(data.init$mu)
    # analysis data set
    final_data=data.init
    n_final = length(final_data$mu)
    final_data$n.final = final_data$n.init
    simrun$n.final[j] = n_final
  
    ##################################
    ## data analysis
    
    # calculate number of patients by centre and allocation imbalance
    n_all = plyr::ddply(final_data, ~stratum+treatment, summarise, n=length(id), .drop=FALSE)
    n_A = n_all[which(n_all$treatment=="A"),]
    n_B = n_all[which(n_all$treatment=="B"),]
    length_A = length(n_A$n[which(n_A$n >0)])
    length_B = length(n_B$n[which(n_B$n >0)])
    n_center = n_A$n + n_B$n
    N1=sum(n_A$n)
    N2=sum(n_B$n)
    simrun$N1[j] = N1
    simrun$N2[j] = N2
    # calculate sum delta^2|r
    moment = sum((n_A$n/N1-n_B$n/N2)^2)
    simrun$moment[j] = moment
    
    # calculate LS variance estimators
    deskr_final = plyr::ddply(final_data, ~stratum+treatment, summarise, n=mean(n.init),
                              mean=mean(outcome), sd=sd(outcome), var=var(outcome), .drop=FALSE)
    sigma2.LS = mean(deskr_final$var, na.rm=TRUE)
    simrun$sigma2.LS[j] = sigma2.LS
    if (c==1){
      sigma2.c.LS = 0
    } else sigma2.c.LS = mean( c(var(deskr_final$mean[which(deskr_final$treatment=="A")], na.rm=TRUE), var(deskr_final$mean[which(deskr_final$treatment=="B")], na.rm=TRUE)), na.rm=TRUE)
    simrun$sigma2.c.LS[j] = sigma2.c.LS

    # calculate overall treatment group means
    deskr.mud.all = plyr::ddply(final_data, ~treatment,summarise , n=mean(n), 
                                mean=mean(outcome), sd=sd(outcome), var=var(outcome), .drop=FALSE)
    mud.all = deskr.mud.all$mean[1] - deskr.mud.all$mean[2]
    simrun$mud.all[j] = mud.all
    
    # linear model without centre
    erg_lm = summary( lm(outcome~1+treatment, data=final_data, method="qr"))
    simrun$mud.lm[j] = abs(erg_lm$coef[2,1])
    simrun$var.lm[j]  =  erg_lm$coef[2,2]
    simrun$t.lm[j] = erg_lm$coef[2,3]
    simrun$df.lm[j] = erg_lm$df[2]
    simrun$p.lm[j] = erg_lm$coef[2,4]
    simrun$sigma2.lm[j] = erg_lm$sigma^2

    # test statistics
    t_true = mud.all / sqrt( sigma^2 * n_final/(N1 * N2)  + sigma.c^2 * moment )
    t_LS = mud.all / sqrt( sigma2.LS * n_final/(N1 * N2)  + sigma2.c.LS * moment )
    simrun$t_true[j] = t_true
    simrun$t_LS[j] = t_LS

    # p-values
    p_NV.true = 2*pnorm(abs(t_true), mean=0, sd=1, lower.tail=FALSE)
    p_NV.LS = 2*pnorm(abs(t_LS), mean=0, sd=1, lower.tail=FALSE)
    simrun$p_NV.true[j] = p_NV.true  
    simrun$p_NV.LS[j] = p_NV.LS 
  }

  ## Ende Schleife nsim 
  ##################################
  
  ##################################
  ## Output generation

  pd_NV.true = sum(simrun$p_NV.true<0.05)*1/nsim
  pd_NV.LS = sum(simrun$p_NV.LS<0.05)*1/nsim
  pd_t.lm = sum(simrun$p.lm<0.05)*1/nsim

  result = data.frame(n=n, n_final=mean(simrun$n.final), c=c, b=b, k=balloc, sm=approx, round=round, dist=dist, mu=mu, 
                      sigma=sigma, sigma.init=sigma.init, sigma.c= sigma.c, sigma.c.init = sigma.c.init, nsim=nsim, 
                      pd_NV.true=pd_NV.true, pd_NV.LS=pd_NV.LS, pd_t.lm=pd_t.lm )
  
  end = proc.time()[3]
  duration = end - start
  result$dur = duration
  
  res_list = list(data=simrun, erg = result)
  return(res_list)
}