#---------------------------------------
# Install and load required packages
#---------------------------------------
# install.packages(c("devtools","plyr","ARPobservation"))
# library(devtools)
# install_github("jepusto/Pusto")

library(Pusto)
library(ARPobservation)
library(plyr)
rm(list=ls())

#---------------------------------------
# generate data for ES calculation
#---------------------------------------

m = 10; n = 10; L = 10; phi = 0.5; zeta = 1; shape = 1; delta = 1; behavior = "state"; increase = FALSE; trunc_scale = 2

mu <- rep(phi / zeta, m + n)
lambda <- c(rep((1 - phi) / zeta, m), rep((1 - phi * delta) / (zeta * delta), n))

BS <- r_behavior_stream(m + n, mu = mu, lambda = lambda,
                        F_event = F_exp(), F_interim = F_gam(shape),
                        stream_length = L)

if (behavior == "state") {
  CR <- continuous_duration_recording(BS)
  c_int <- c(10, 20, 30) / 60
  MTS <- sapply(c_int, momentary_time_recording, BS = BS, summarize = TRUE)
  PIR <- sapply(c_int, interval_recording, BS = BS, rest_length = 0, partial = TRUE, summarize = TRUE)
  Y_dat <- data.frame(CR, MTS, PIR)
  colnames(Y_dat) <- c("CR","MTS-10","MTS-20","MTS-30","PIR-10","PIR-20","PIR-30")
  trunc_vals <- c(1 / 60, c_int, c_int) / (trunc_scale * L)
} else {
  EC <- event_counting(BS)
  c_int <- c(10, 20, 30) / 60
  PIR <- sapply(c_int, interval_recording, BS = BS, rest_length = 0, partial = TRUE, summarize = TRUE)
  Y_dat <- data.frame(EC, PIR)
  colnames(Y_dat) <- c("EC","PIR-10","PIR-20","PIR-30")
  trunc_vals <- c(1, c_int / L) / trunc_scale
}

base_phase <- 0
phase <- c(rep(0,m), rep(1,n))
data <- Y_dat[,"PIR-30"]
trunc_val <- trunc_vals[1]
nObs <- table(phase)
measures = c("PND","PEM","PAND","IRD","NAP","SMD","LRR")

#----------------------------------
# effect size measures
#----------------------------------

PND <- function(data, phase, base_phase, increase = TRUE, ...) {
  if (!increase) data <- -1 * data
  100 * mean(data[phase != base_phase] > max(data[phase == base_phase]))
}


PEM <- function(data, phase, base_phase, increase = TRUE, ...) {
  if (!increase) data <- -1 * data
  med <- median(data[phase == base_phase])
  100 * mean((data[phase != base_phase] > med) + 0.5 * (data[phase != base_phase] == med))
}


PAND <- function(data, phase, base_phase, increase = TRUE, ...) {
  if (!increase) data <- -1 * data
  m <- sum(phase == base_phase)
  n <- sum(phase != base_phase)

  X <- c(-Inf, sort(data[phase == base_phase]))
  Y <- c(sort(data[phase != base_phase]), Inf)
  ij <- expand.grid(i = 1:(m + 1), j = 1:(n + 1))
  ij$no_overlap <- mapply(function(i, j) X[i] < Y[j], i = ij$i, j = ij$j)
  ij$overlap <- with(ij, i + n - j)
  overlaps <- with(ij, max(overlap * no_overlap))
  
  100 * overlaps / length(data)
}


IRD <- function(data, phase, base_phase, increase = TRUE, ...) {
  pand <- PAND(data = data, phase = phase, base_phase = base_phase, increase = increase) / 100
  m <- sum(phase == base_phase)
  n <- sum(phase != base_phase)
  
  ((m+n)^2 * pand - m^2 - n^2) / (2 * m * n)
}


NAP <- function(data, phase, base_phase, increase = TRUE, ...) {
  if (!increase) data <- -1 * data
  XY <- expand.grid(x = data[phase==base_phase], y = data[phase!=base_phase])
  100 * mean(with(XY, (y > x) + 0.5 * (y == x)))
}


Tau <- function(data, phase, base_phase, increase = TRUE, ...) {
  nap <- NAP(data = data, phase = phase, base_phase = base_phase, increase = TRUE)
  2 * nap / 100 - 1
}

SMD <- function(data, phase, trunc_val = 0, nObs = NULL, ...) {
  if (is.null(nObs)) nObs <- table(phase)
  y_bar <- tapply(data, phase, mean)
  s_sq <- tapply(data, phase, var)
  s_tilde_sq <- pmax(s_sq, trunc_val^2)
  d <- (y_bar[[2]] - y_bar[[1]]) / sqrt(s_tilde_sq[[1]])
  g <- (1 - 3 / (4 * nObs[[1]] - 5)) * d
  c(d = d, g = g)
}

LRR <- function(data, phase, trunc_val = 0, nObs = NULL, ...) {
  
  if (is.null(nObs)) nObs <- table(phase)
  y_bar <- tapply(data, phase, mean)
  y_tilde <- pmax(y_bar, trunc_val / nObs)
  s_sq <- tapply(data, phase, var)
  
  lRR <- log(y_tilde)[[2]] - log(y_tilde)[[1]]
  BC <- log(y_tilde) + s_sq / (2 * nObs * y_tilde^2)
  lRR_C <- BC[[2]] - BC[[1]]
  
  c(R1 = lRR, R2 = lRR_C)
}


effect_sizes <- function(data, phase, base_phase = 0, increase = TRUE, trunc_val = 0, nObs = NULL,
                        measures = c("PND","PEM","PAND","IRD","NAP","SMD","LRR")) {
  unlist(sapply(measures, do.call, 
         args = list(data = data, phase = phase, base_phase = base_phase, 
                     increase = increase, trunc_val = trunc_val, nObs = nObs)))
}

ES <- mapply(function(dat, trunc_val) effect_sizes(dat = dat, phase = phase, trunc_val = trunc_val, nObs = c(m, n)), 
             dat = Y_dat, trunc_val = trunc_vals)
ES

#----------------------------------
# Data generation
#----------------------------------

r_ES <- function(m, n, L, phi, zeta, shape, delta, behavior, increase, trunc_scale = 2) {
  mu <- rep(phi / zeta, m + n)
  lambda <- c(rep((1 - phi) / zeta, m), rep((1 - phi * delta) / (zeta * delta), n))
  
  BS <- r_behavior_stream(m + n, mu = mu, lambda = lambda,
                          F_event = F_exp(), F_interim = F_gam(shape),
                          stream_length = L)
  
  if (behavior == "state") {
    CR <- continuous_duration_recording(BS)
    c_int <- c(10, 20, 30) / 60
    MTS <- sapply(c_int, momentary_time_recording, BS = BS, summarize = TRUE)
    PIR <- sapply(c_int, interval_recording, BS = BS, rest_length = 0, partial = TRUE, summarize = TRUE)
    Y_dat <- data.frame(CR, MTS, PIR)
    colnames(Y_dat) <- c("CR","MTS-10","MTS-20","MTS-30","PIR-10","PIR-20","PIR-30")
    trunc_vals <- c(1 / 60, c_int, c_int) / (trunc_scale * L)
  } else {
    EC <- event_counting(BS)
    c_int <- c(10, 20, 30) / 60
    PIR <- sapply(c_int, interval_recording, BS = BS, rest_length = 0, partial = TRUE, summarize = TRUE)
    Y_dat <- data.frame(EC, PIR)
    colnames(Y_dat) <- c("EC","PIR-10","PIR-20","PIR-30")
    trunc_vals <- c(1, c_int / L) / trunc_scale
  }
  
  phase <- c(rep(0,m), rep(1,n))
  ES <- mapply(function(dat, trunc_val) 
                effect_sizes(dat = dat, phase = phase, increase = increase, trunc_val = trunc_val, nObs = c(m, n)), 
               dat = Y_dat, trunc_val = trunc_vals)
  ES
}

r_ES(m, n, L, phi, zeta, shape, delta, behavior, increase, trunc_scale = 2)

#----------------------------------
# Simulation driver
#----------------------------------

drive_sim <- function(reps, m, n, L, phi, zeta, shape, delta, behavior, seed = NULL) {
  set.seed(seed)
  ES_sims <- replicate(reps, r_ES(m, n, L, phi, zeta, shape, delta, behavior, increase = FALSE))
  ES_mean <- as.data.frame(t(apply(ES_sims, 1:2, mean, na.rm=TRUE)))
  ES_mean$procedure <- rownames(ES_mean)
  ES_mean
}

drive_sim(reps = 100, m, n, L, phi, zeta, shape, delta, behavior)

source_obj <- ls()

#-------------------------------------------------------------
# Design parameters - operational sensitivities simulation
#-------------------------------------------------------------

set.seed(20170510)

L <- c(5, 10, 15, 20)
m <- c(5, 10, 15, 20)
n <- c(5, 10, 15, 20)
phi <- c(0.2, 0.5, 0.8)
zeta_state <- c(1.0, 2.0)
zeta_event <- c(0.5, 1.0, 2.0)
shape <- c(1, 2)
delta <- c(1.0, 0.5, 0.2)
reps <- 2000

params_check <- within(expand.grid(phi = phi, zeta = zeta_state, delta = delta), {
  mu_B <- phi / zeta
  lambda_B <- (1 - phi) / zeta
  phi_T <- phi * delta
  zeta_T <- phi_T / mu_B
  lambda_T <- (1 - phi * delta) / (zeta * delta)
})
params_check

params_state <- expand.grid(L = L, m = m, n = n, phi = phi, zeta = zeta_state, 
                            shape = 1, delta = delta, behavior = "state", reps = reps,
                            KEEP.OUT.ATTRS = FALSE)
params_event <- expand.grid(L = L, m = m, n = n, phi = 0, zeta = zeta_event, 
                            shape = shape, delta = delta, behavior = "event", reps = reps,
                            KEEP.OUT.ATTRS = FALSE)
params <- rbind(params_state, params_event)

params$seed <- 1:nrow(params) + round(runif(1, max = 2^20))

nrow(params)
str(params)


#--------------------------------------------------------
# run simulations in parallel
#--------------------------------------------------------

cluster <- start_parallel(source_obj = source_obj, packages = "ARPobservation", register = TRUE)
system.time(results <- mdply(params, drive_sim, .parallel = TRUE))
parallel::stopCluster(cluster)

save(L, m, n, phi, zeta_state, zeta_event, shape, delta, reps, params, results, 
     file = "R/EffectSizeMeasures.Rdata")

#--------------------------------------------------------
# Format for export to csv 
#--------------------------------------------------------

results_export <- within(results, {
  change <- factor(paste0((1 - delta) * 100, "%"))
  prevalence <- factor(paste0(phi * 100, "%"))
  incidence <- factor(paste0(zeta, "/min"))
  distribution <- factor(shape, levels = 1:2, labels = c("Exponential", "Gamma(2)"))
})

results_export <- subset(results_export, 
                         select = c(behavior, prevalence, incidence, distribution, change, 
                                    procedure, L, m, n, PND:NAP, SMD.d, SMD.g, LRR.R1, LRR.R2))
names(results_export)[7:9]
names(results_export)[7:9] <- c("Session_length","Baseline_length","Treatment_length")

write.csv(results_export, file = "R/EffectSizeMeasures.csv", row.names = FALSE)