# sqs version 3.2 by John Alroy
# performs shareholder quorum subsampling on an array of specimen counts
# can be used to perform classical rarefaction instead of SQS
# written 29 July 2010; version 2.0 completed 14 February 2011; versions 3.0,
#  3.1, 3.2, and 3.3 written 3 June, 12 July, 19-23 August, and 2 December 2011
# change in version 3.3: fixed bug in computation of maximum estimated coverage
# changes in version 3.2: specimens drawn and subsampled U are based on
#  geometric means of counts; ignore.singletons option added
# change in version 3.1: an even better frequency adjustment involving
#  singletons and doubletons
# changes in version 3.0: an even better subsampling algorithm involving a new
#  adjustment to u combined with a new throwback criterion
# changes in version 2.0: improved subsampling algorithm; including the dominant
#  taxon is now the default; improved reporting of errors and basic statistics
# warning: do not use this program with taxonomic occurrence data drawn from
#  multiple published references because it is not designed to count
#  single-reference taxa or adjust for long taxonomic lists
# warning: version 1.0 yields estimates that are downwards-biased when q < 0.6
#  and abundance distributions are highly uneven

# Read in data
data<-read.csv("file.csv",header=TRUE)

#Convert occurrence list to occurrence table

data.occs<-table(data[,"genus"])

sqs<-function(ab,q,trials,method,ignore.singletons,dominant)  {

  params <- array(data=NA,dim=0)
  if (missing(trials))  {
    trials <- 100
  }
  if (missing(method))  {
    method <- ""
  } else if (method != "" && method != "rarefaction" && method != "CR")  {
    return(print('If the method is rarefaction enter method="rarefaction" or "CR"',quote=F))
  }
  if ((q <= 0 || q >= 1) && method != "rarefaction" && method != "CR")  {

    return(print("If the method is SQS the quota must be greater than zero and less than one",quote=F))
  } else if (q < 1 && (method == "rarefaction" || method == "CR"))  {

    return(print("If the method is rarefaction the quota must be an integer",quote=F))
  }
  ignore <- 0
  if (! missing(ignore.singletons) && (ignore.singletons == T || ignore.singletons == "yes" || ignore.singletons == "y"))  {
    ignore <- 1
  }
  if (missing(dominant))  {
    dominant <- 0
  } else if (dominant != "" && dominant != "exclude" && dominant != "no")  {
    return(print('To exclude the dominant taxon, enter dominant="exclude" or "no"',quote=F))
  }

  # compute basic statistics
  specimens <- sum(ab)
  singletons <- 0
  doubletons <- 0
  highest <- 0
  for (i in 1:length(ab))  {
    if (ab[i] == 1)  {
      singletons <- singletons + 1
    } else if (ab[i] == 2)	{
      doubletons <- doubletons + 1
    }
    if (ab[i] > highest)	{
      highest <- ab[i]
      mostfrequent <- i
    }
  }
  u <- 1 - singletons / specimens
  # optionally, exclude the dominant taxon
  if (dominant == "exclude" || dominant == "no")	{
    u <- 1 - singletons / (specimens - highest)
  }

  if (u == 0)	{
    return(print("Coverage is zero because all taxa are singletons",quote=F))
  }

  # compute raw taxon frequencies (temporarily)
  freq <- ab / specimens

  # standard recursive equation for Fisher's alpha
  alpha <- 10
  oldalpha <- 0
  while (abs(alpha - oldalpha) > 0.0000001)	{
    oldalpha <- alpha
    alpha <- length(ab) / log(1 + specimens/alpha)
  }

  params["raw richness"] <- length(ab)
  params["subsampled richness"] <- u
  params["subsampled richness"] <- NA
  params["subsampled u"] <- NA
  params["Chao 1"] <- length(ab) + singletons**2/(2* doubletons)
  params["subsampled Chao 1"] <- NA
  # governing parameter of the geometric series distribution
  params["k"] <- abs(lm(log(sort(freq)) ~ c(1:length(freq)))$coefficients[2])
  params["Fisher's alpha"] <- alpha
  params["Shannon's H"] <- -1 * sum(freq * log(freq))
  params["Hurlbert's PIE"] <- (1 - sum(freq**2)) * length(ab) / (length(ab) - 1)
  params["dominance"] <- highest / specimens
  params["specimens"] <- specimens
  params["singletons"] <- singletons
  params["doubletons"] <- doubletons
  params["specimens drawn"] <- 0

  if (dominant != "exclude" && dominant != "no")	{
    highest <- 0
    mostfrequent <- 0
  }

  # return if the rarefaction quota is equal to or higher than the
  #  specimen count
  if (method == "rarefaction" && q >= specimens - highest)	{
    return(params)
  }

  # compute taxon frequencies (tweak added in version 3.1)
  freq <- ab - (singletons + doubletons / 2) / length(ab)
  freq <- freq / (specimens - highest)

  # return if the quorum target is higher than estimated coverage
  if ((q > sum(freq) && method != "rarefaction" && method != "CR") || (q >= sum(ab)))	{
    return(params)
  }

  # create an array in which each cell corresponds to one specimen
  ids <- array()
  n <- 0
  for (i in 1:length(ab))	{
    for (j in 1:ab[i])	{
      n <- n + 1
      ids[n] <- i
    }
  }

  # subsampling trial loop
  # s will be the subsampled taxon count
  s <- array(rep(0,trials))
  drawn <- array(rep(0,trials))
  mostfrequentdrawn <- array(rep(0,trials))
  subsingle <- array(rep(0,trials))
  subdouble <- array(rep(0,trials))
  subchao <- array(rep(0,trials))

  for (trial in 1:trials)	{
    pool <- ids
    left <- length(pool)
    seen <- array(data=rep(0,length(ab)))
    subfreq <- array(rep(0,length(ab)))
    if (method != "rarefaction" && method != "CR")	{
      udrawn <- 0
      # equation new to version 3.0
      # the exponent corrects for downwards bias
      while (udrawn < q)	{
        # draw a specimen
        x <- floor(runif(1,min=1,max=left+1))
        # add to frequency and taxon sums if species has
        #  not been drawn previously
        subfreq[pool[x]] <- subfreq[pool[x]] + 1
        if (seen[pool[x]] == 0)	{
          if (pool[x] != mostfrequent && (ignore == 0 || ab[pool[x]] > 1))	{
            udrawn <- udrawn + freq[pool[x]]
          }
          seen[pool[x]] <- 1
          # randomly throw back some draws that put the sum over q
          #  (an even better algorithm added in version 3.0)
          if (runif(1) <= freq[pool[x]] || udrawn < q)	{
            s[trial] <- s[trial] + 1
          } else	{
            subfreq[pool[x]] <- subfreq[pool[x]] - 1
          }
        }
        # decrease pool of specimens not yet drawn
        pool[x] <- pool[left]
        left <- left - 1
      }
    } else	{
      i <- 0
      draws <- 0
      while (i < q)	{
        draws <- draws + 1
        x <- floor(runif(1,min=1,max=length(ids)-draws+2))
        subfreq[pool[x]] <- subfreq[pool[x]] + 1
        if (pool[x] != mostfrequent)	{
          i <- i + 1
        }
        if (seen[pool[x]] == 0)	{
          seen[pool[x]] <- 1
          s[trial] <- s[trial] + 1
        }
        pool[x] <- pool[length(ids)-draws+1]
      }
    }
    for (i in 1:length(ab))	{
      if (subfreq[i] == 1 && i != mostfrequent)	{
        subsingle[trial] <- subsingle[trial] + 1
      } else if (subfreq[i] == 2 && i != mostfrequent)	{
        subdouble[trial] <- subdouble[trial] + 1
      }
    }
    if (subsingle[trial] > 0 && subdouble[trial] > 0)	{
      subchao[trial] <- s[trial] + subsingle[trial]**2/(2*subdouble[trial])
    } else	{
      subchao[trial] <- s[trial]
    }
    drawn[trial] <- sum(subfreq)
    if (mostfrequent != 0)	{
      mostfrequentdrawn[trial] <- subfreq[mostfrequent]
    }
  }

  # compute vectors of non-zero counts
  options(warn=-1)
  s2 <- sort(sqrt(s-1))^2+1
  d2 <- sort(sqrt(drawn-1))^2+1
  m2 <- sort(sqrt(mostfrequentdrawn-1))^2+1
  ss2 <- sort(sqrt(subsingle-1))^2+1
  options(warn=0)

  # compute geometric means
  params["subsampled richness"] <- exp(mean(log(s2))) * length(s2)/length(s)
  params["specimens drawn"] <- exp(mean(log(d2))) * length(d2)/length(drawn)
  meanmost <- 0
  if (sum(mostfrequentdrawn) > 0)	{
    meanmost <- exp(mean(log(m2))) * length(m2)/length(mostfrequentdrawn)
  }
  meansubsingle <- exp(mean(log(ss2))) * length(ss2)/length(subsingle)

  params["subsampled u"] <- 1 - meansubsingle / (params["specimens drawn"] - meanmost)
  params["subsampled Chao 1"] <- exp(mean(log(subchao)))
  return(params)

}

# Check real value with a quorum of 0.4
sqs(data.occs,0.4)['subsampled richness']

# Bootstrap
s<-rep(NA,1000)
for(i in 1:length(s)){
  bootstrapped<-data.occs[sample(1:length(data.occs),replace=TRUE)]
  s[i]<-sqs(bootstrapped,0.4)['subsampled richness']
}

# To remove any NAs (bootstrapping can produce a lot of pseudo-singletons)
s<-na.omit(s)
s<-as.numeric(s)

# Returns the median and 95% confidence intervals
quantile(s,probs=c(0.5,0.025,0.975))
hist(s)


#  Functions used in the Settlement Vegetation analysis from Simon Goring:

#  The Morisita density estimator:
morisita <- function(processed.data, correction.factor = NULL, veil=FALSE) {
  #  Function to calculate stem density using the morista function.  The input
  #  is 'processed.data', which should be the file 'used.data'.  'correction
  #  factor' is the modified Cottam Correction factor determined in
  #  'load.estimate.correction.R' using a generalized linear model (with a Gamma
  #  distribution).

  azim <- processed.data[,c('az1', 'az2', 'az3', 'az4')]
  diam <- processed.data[,c('diam1', 'diam2', 'diam3', 'diam4')]
  dist <- processed.data[,c('dist1', 'dist2', 'dist3', 'dist4')]
  spec <- processed.data[,c('species1', 'species2', 'species3', 'species4')]
  corn <- processed.data[,'corner']

  if(veil){
    diam[diam < 8] <- NA #double check that this is corrected veil line, and do we want to use it?
  }

  diam[diam == 0 & !spec == 'No tree'] <- NA

  #m.diam <- diam/100 #diameters are in cm already
  m.diam <- (diam * 2.54) / 100 # convert diameter from inches to meters

  dist <- floor(apply(dist, 2, function(x)as.numeric(as.character(x))))
  azim <- floor(apply(azim, 2, function(x)as.numeric(as.character(x))))

  #  This tells us how many quadrats are used.  I'd prefer to use all points
  #  where samples are drawn from two quadrats, but in some cases it seems that
  #  there are NAs in the data.
  #  If a point has recorded azimuths we state that they must be in two different
  #  quadrats:

  two.quads <- apply(azim[,1:2], 1, function(x) sum(!is.na(unique(floor(x/90)))))

  #  There are 10,155 points for which the first two trees were sampled in the
  #  same quadrat.  In general these are randomly distributed, but interestingly
  #  there's a big clump of them in Wisconsin.  Even so, there are lots of other
  #  points around.  We can accept that these points are categorically wrong.
  #
  #
  #  sum((two.quads == 1 & !(is.na(azim[,1]) | is.na(azim[,2]))))

  #  we need to change:

  two.quads[((two.quads < 2 & (is.na(azim[,1]) | is.na(azim[,2]))) &
               !(is.na(dist[,1]) | is.na(dist[,2])))] <- 2

  #  Exclusions include:
  #  Plots with a tree as plot center:
  two.quads[dist[,1] == 0] <- 0

  #  Plots where one of the trees has no measured diameter:
  two.quads[is.na(diam[,1]) | is.na(diam[,2])] <- 0

  #  Plots where a distance to tree is missing:
  two.quads[is.na(dist[,1]) | is.na(dist[,2])] <- 0
  #  This is the same as k in Charlie's spreadsheet:

  q <- two.quads

  #  Tree dist is measured in links in the dataset, I am converting to
  #  meters and adding one half a dimater (in cm), on Charlie's advice.

     m.dist <- dist * 0.201168 + 0.5 * m.diam # convert distances from chains (links) to meters
     #m.dist <- dist + 0.5 * m.diam
  #  rsum is the sum of the squared radii, in cases where there are two trees in
  #  the same quadrant I'm going to drop the site, as I will with any corner
  #  with only one tree since the morista density estimator can't calculate
  #  density with less than two trees, and requires two quadrats.

  #  I'm going to let the NAs stand in this instance.
  rsum <- rowSums((m.dist[,1:2])^2, na.rm=T)
  #rmax <- max(m.dist[,1:2], na.rm=T)
  #rmax<- apply(m.dist[,1:2], 1, max, na.rm= T)

  #  A set of conditions to be met for the rsum to be valid:
  rsum[rowSums(is.na(m.dist[,1:2])) == 2 |  q < 2 | rsum == 0 | rowSums(m.dist[,1:2], na.rm=T) < 0.6035] <- NA
  #rmax[rowSums(is.na(m.dist[,1:2])) == 2 |  q < 2 | rmax == 0 | rowSums(m.dist[,1:2], na.rm=T) < 0.6035] <- NA

  #  From the formula,
  #  lambda = kappa * theta * (q - 1)/(pi * n) * (q / sum_(1:q)(r^2))
  #  here, n is equal to 1.
  #  units are in stems / m^2

  #charilies morisita.est have separate corrections for each type of corner
  morisita.est <- ((q - 1) / (pi * 1)) * (2 / rsum) *
    correction.factor$kappa  * correction.factor$theta* correction.factor$zeta * correction.factor$phi

  morisita.est[q < 2] <- NA

  #  Now they're in stems / hectare
  morisita.est <- morisita.est * 10000

  #  Basal area is the average diameter times the stem density.
  #  The stem density is measured in trees / ha.
  met.rad <- (diam / 2) * 2.54 / 100
  #met.rad <- (diam / 2)/100
  basal.area <- morisita.est * rowSums(pi * met.rad^2, na.rm=TRUE)

  basal.area[ q < 2 ] <- NA

 # radius <- max(m.dist, na.rm = TRUE)
  #plot.area <- pi*rmax/2 # calculate average plot area
  return(list(morisita.est, basal.area))

}


#' Expectation step.
#'
#' @param p List of parameters.
#' @param item_data Matrix or dataframe of item responses.
#' @param pred_data Matrix or dataframe of DIF and/or impact predictors.
#' @param item_type Vector of character values indicating the item type.
#' @param mean_predictors Possibly different matrix of predictors for the mean
#' impact equation.
#' @param var_predictors Possibly different matrix of predictors for the
#' variance impact equation.
#' @param theta Vector of fixed quadrature points.
#' @param samp_size Sample size in dataset.
#' @param num_items Number of items in dataset.
#' @param num_responses Number of responses for each item.
#' @param adaptive_quad Logical indicating whether to use adaptive quadrature.
#' @param num_quad Number of quadrature points used for approximating the
#' latent variable.
#' @param get_eap Logical indicating whether to compute EAP scores.
#' @param NA_cases Logical vector indicating missing observations.
#'
#' @return a \code{"list"} of posterior values from the expectation step
#'
#' @keywords internal
#'
Estep <-
  function(p,
           item_data,
           pred_data,
           item_type,
           mean_predictors,
           var_predictors,
           theta,
           samp_size,
           num_items,
           num_responses,
           adapt_quad,
           num_quad,
           get_eap,
           NA_cases) {

    # Make space for the trace lines and the E-tables.
    observed_ll <- 0
    itemtrace <- rep(list(NA),num_items)
    etable <- matrix(0, nrow = samp_size, ncol = num_quad)
    eap_scores <- if(get_eap) matrix(NA, nrow = length(NA_cases), ncol = 1)
    eap_sd <- if(get_eap) matrix(NA, nrow = length(NA_cases), ncol = 1)

    # Impact.
    alpha <- mean_predictors %*% p[[num_items+1]]
    phi <- exp(var_predictors %*% p[[num_items+2]])

    if(adapt_quad == TRUE) {
      theta <- mean(alpha) +
        sqrt(2*mean(phi))*statmod::gauss.quad(n = num_quad,
                                              kind = "hermite")$nodes
    }

    # Compute the trace lines.
    for (item in 1:num_items) {
      if(item_type[item] == "cfa") {
        itemtrace[[item]] <- gaussian_traceline_pts(p[[item]],
                                                    theta,
                                                    item_data[,item],
                                                    pred_data,
                                                    samp_size)
      } else if (item_type[item] == "2pl") {
        itemtrace[[item]] <- bernoulli_traceline_pts(p[[item]],
                                                     theta,
                                                     pred_data,
                                                     samp_size)
      } else {
        itemtrace[[item]] <- cumulative_traceline_pts(p[[item]],
                                                      theta,
                                                      pred_data,
                                                      samp_size,
                                                      num_responses[item],
                                                      num_quad)
      }
    }

    # Obtain E-table.
    for(i in 1:samp_size) {

      # Get weights for computing adaptive or fixed-point quadrature.
      posterior <- dnorm(theta,
                         mean = alpha[i],
                         sd = sqrt(phi[i]))

      # For each individual i, loop over items (j) and compute posterior
      # probability of response pattern.
      for(j in 1:num_items) {

        x <- item_data[i,j]
        if(is.na(x)) next

        if(item_type[j] == "cfa") { # Continuous responses.
          posterior <- posterior*itemtrace[[j]][i,]
        } else if(item_type[j] == "2pl") { # Binary responses.
          if(x == 1) {
            posterior <- posterior*(1-itemtrace[[j]][i,])
          } else {
            posterior <- posterior*itemtrace[[j]][i,]
          }
        } else if(item_type[j] == "graded") { # Ordered categorical responses.
          if(x == 1) {
            posterior <- posterior*(1-itemtrace[[j]][[1]][i,])
          } else if(x == num_responses[j]) {
            posterior <- posterior*itemtrace[[j]][[num_responses[j]-1]][i,]
          } else {
            posterior <- posterior*(itemtrace[[j]][[x-1]][i,]-
                                        itemtrace[[j]][[x]][i,])
          }
        }

      }

      # Normalize posterior.
      marginal <- sum(posterior, na.rm = TRUE)
      observed_ll <- observed_ll + log(marginal)
      if(marginal == 0) marginal <- 1
      etable[i,] <- posterior/marginal

      # EAPs.
      if(get_eap) {
        eap_scores[which(!NA_cases)[i]] <- sum(posterior*theta)/marginal
        eap_sd[which(!NA_cases)[i]] <-
          sqrt(sum(posterior*(theta-eap_scores[which(!NA_cases)[i]])**2)/marginal)
      }

    }

    # E-table matrix to be used in Q function and (possibly adaptive)
    # theta values.
    return(list(etable=etable,eap_scores=eap_scores,eap_sd=eap_sd,
                theta=theta,observed_ll=observed_ll))

  }

#' Expectation step with proxy data.
#'
#' @param p List of parameters.
#' @param item_data Matrix or dataframe of item responses.
#' @param pred_data Matrix or dataframe of DIF and/or impact predictors.
#' @param item_type Vector of character values indicating the item type.
#' @param mean_predictors Possibly different matrix of predictors for the mean
#' impact equation.
#' @param var_predictors Possibly different matrix of predictors for the
#' variance impact equation.
#' @param prox_data Vector of observed proxy scores.
#' @param samp_size Sample size in dataset.
#' @param num_items Number of items in dataset.
#' @param num_responses Number of responses for each item.
#' @param get_eap Logical indicating whether to compute EAP scores.
#' @param NA_cases Logical vector indicating missing observations.
#'
#' @return a \code{"list"} of posterior values from the expectation step
#'
#' @keywords internal
#'
Estep_proxy <-
  function(p,
           item_data,
           pred_data,
           item_type,
           mean_predictors,
           var_predictors,
           prox_data,
           samp_size,
           num_items,
           num_responses,
           get_eap,
           NA_cases) {

    # Make space for the trace lines and the E-tables.
    observed_ll <- 0
    itemtrace <- rep(list(NA),num_items)
    etable <- matrix(0, nrow = samp_size, ncol = 1)
    eap_scores <- if(get_eap) matrix(NA, nrow = length(NA_cases), ncol = 1)
    eap_sd <- if(get_eap) matrix(NA, nrow = length(NA_cases), ncol = 1)

    # Impact.
    alpha <- mean_predictors %*% p[[num_items+1]]
    phi <- exp(var_predictors %*% p[[num_items+2]])

    # Compute the trace lines.
    for (item in 1:num_items) {
      if(item_type[item] == "cfa") {
        itemtrace[[item]] <- gaussian_traceline_pts_proxy(p[[item]],
                                                          prox_data,
                                                          item_data[,item],
                                                          pred_data,
                                                          samp_size)
      } else if (item_type[item] == "2pl") {
        itemtrace[[item]] <- bernoulli_traceline_pts_proxy(p[[item]],
                                                           prox_data,
                                                           pred_data)
      } else {
        itemtrace[[item]] <- cumulative_traceline_pts_proxy(p[[item]],
                                                      prox_data,
                                                      pred_data,
                                                      samp_size,
                                                      num_responses[item])
      }
    }

    # Obtain E-table.
    for(i in 1:samp_size) {

      # Get weights for computing adaptive or fixed-point quadrature.
      posterior <- dnorm(prox_data[i],
                         mean = alpha[i],
                         sd = sqrt(phi[i]))

      # For each individual i, loop over items (j) and compute posterior
      # probability of response pattern.
      for(j in 1:num_items) {

        x <- item_data[i,j]
        if(is.na(x)) next

        if(item_type[j] == "cfa") { # Continuous responses.
          posterior <- posterior*itemtrace[[j]][i,]
        } else if(item_type[j] == "2pl") { # Binary responses.
          if(x == 1) {
            posterior <- posterior*(1-itemtrace[[j]][i,])
          } else {
            posterior <- posterior*itemtrace[[j]][i,]
          }
        } else if(item_type[j] == "graded") { # Ordered categorical responses.
          if(x == 1) {
            posterior <- posterior*(1-itemtrace[[j]][[1]][i,])
          } else if(x == num_responses[j]) {
            posterior <- posterior*itemtrace[[j]][[num_responses[j]-1]][i,]
          } else {
            posterior <- posterior*(itemtrace[[j]][[x-1]][i,]-
                                      itemtrace[[j]][[x]][i,])
          }
        }

      }

      # Normalize posterior.
      observed_ll <- observed_ll + log(posterior)
      etable[i,] <- posterior

      # EAPs.
      if(get_eap) {
        eap_scores[which(!NA_cases)[i]] <- posterior*prox_data[i]
        eap_sd[which(!NA_cases)[i]] <-
          sqrt(posterior*(prox_data[i]-eap_scores[which(!NA_cases)[i]])**2)
      }

    }

    # E-table matrix to be used in Q function and (possibly adaptive)
    # theta values.
    return(list(etable=etable,eap_scores=eap_scores,eap_sd=eap_sd,observed_ll=observed_ll))

  }


#!/usr/bin/env Rscript

#Andrew Burt - a.burt@ucl.ac.uk

getFieldResults <- function(caldata,ffname,alpha,runs,localreg,nn=200)
{
	model <- fitNLS(caldata)
	nlqrmodels <- nlQuantileRegression(model,caldata,alpha)
	bmodels <- bootstrap(model,caldata,runs)
	results <- list()
	j <- 1
	for(i in 1:length(ffname))
	{
		###field data
		fielddata <- read.table(ffname[i],col.names=c("id","x","y","d","h","rho"))
		tmp1 <- strsplit(ffname[i],'/')
		tmp2 <- tmp1[[1]][length(tmp1[[1]])]
		tmp3 <- strsplit(tmp2,'\\.')
		pname <- tmp3[[1]][1]
		##filter
		#fielddata <- fielddata[fielddata$d >= 0.9,]
		##
		fdata <- data.frame(fielddata$d*fielddata$d*fielddata$h*fielddata$rho)
		colnames(fdata) <- c("X")
		stem_count <- nrow(fielddata)
		basal_area <- sum((fielddata$d / 2) ^ 2 * pi)
		lorey_height <- sum(((fielddata$d / 2) ^ 2 * pi) * fielddata$h ) / sum((fielddata$d / 2) ^ 2 * pi)
		wd_baw <- sum((fielddata$d / 2) ^ 2 * pi * fielddata$rho ) / sum((fielddata$d / 2) ^ 2 * pi)
		###ols
		ols_t_agb <- yhatBIOMASS(fielddata) * 1000
		ols_intervals <- uncertaintyIntervalsBIOMASS(fielddata,0.05)
		ols_t_agb_li <- ols_intervals[[1]][,1] * 1000
		ols_t_agb_ui <- ols_intervals[[1]][,2] * 1000
		ols_t_u <- 0.5 * (ols_t_agb_ui - ols_t_agb_li) / ols_t_agb
		ols_p_agb <- sum(ols_t_agb)
		ols_p_agb_li <- ols_intervals[[2]][1] * 1000
		ols_p_agb_ui <- ols_intervals[[2]][2] * 1000
		ols_p_u <- 0.5 * (ols_p_agb_ui - ols_p_agb_li) / ols_p_agb
		###nls
		nlmle_sigma_meas <- measSigma(model,fielddata)
		nlmle_t_agb <- yhat(model,fdata)
		nlmle_t_agb_ali <- getPredictionIntervals(nlqrmodels,fdata)[,1]
		nlmle_t_agb_aui <- getPredictionIntervals(nlqrmodels,fdata)[,2]
		nlmle_t_agb_mli <- nlmle_t_agb - (nlmle_sigma_meas * qnorm(alpha/2,lower.tail=FALSE))
		nlmle_t_agb_mui <- nlmle_t_agb + (nlmle_sigma_meas * qnorm(alpha/2,lower.tail=FALSE))
		nlmle_t_agb_li <- nlmle_t_agb_ali - (nlmle_sigma_meas * qnorm(alpha/2,lower.tail=FALSE))
		nlmle_t_agb_ui <- nlmle_t_agb_aui + (nlmle_sigma_meas * qnorm(alpha/2,lower.tail=FALSE))
		nlmle_t_u_a <- 0.5 * (nlmle_t_agb_aui - nlmle_t_agb_ali) / nlmle_t_agb
		nlmle_t_u_m <- 0.5 * (nlmle_t_agb_mui - nlmle_t_agb_mli) / nlmle_t_agb
		nlmle_t_u <- 0.5 * (nlmle_t_agb_ui - nlmle_t_agb_li) / nlmle_t_agb
		#
		nlmle_p_agb <- sum(nlmle_t_agb)
		nlmle_p_agb_ali <- sum(getConfidenceIntervals(bmodels,fdata,alpha)[,1])
		nlmle_p_agb_aui <- sum(getConfidenceIntervals(bmodels,fdata,alpha)[,2])
		nlmle_p_agb_mli <- nlmle_p_agb - (1/sqrt(nrow(fdata)) * sum(nlmle_sigma_meas*qnorm(alpha/2,lower.tail=FALSE)))
		nlmle_p_agb_mui <- nlmle_p_agb + (1/sqrt(nrow(fdata)) * sum(nlmle_sigma_meas*qnorm(alpha/2,lower.tail=FALSE)))
		nlmle_p_agb_li <- nlmle_p_agb_ali - (1/sqrt(nrow(fdata)) * sum(nlmle_sigma_meas*qnorm(alpha/2,lower.tail=FALSE)))
		nlmle_p_agb_ui <- nlmle_p_agb_aui + (1/sqrt(nrow(fdata)) * sum(nlmle_sigma_meas*qnorm(alpha/2,lower.tail=FALSE)))
		nlmle_p_u_a <- 0.5 * (nlmle_p_agb_aui - nlmle_p_agb_ali) / nlmle_p_agb
		nlmle_p_u_m <- 0.5 * (nlmle_p_agb_mui - nlmle_p_agb_mli) / nlmle_p_agb
		nlmle_p_u <- 0.5 * (nlmle_p_agb_ui - nlmle_p_agb_li) / nlmle_p_agb
		#
		if(localreg == FALSE)
		{
			###sort tree results
			tr <- data.frame(pname,fielddata$id,fielddata$x,fielddata$y,round(fielddata$d,3),round(fielddata$h,3),round(fielddata$rho,3),
			round(ols_t_agb,3),round(ols_t_agb_li,3),round(ols_t_agb_ui,3),round(ols_t_u,3),
			round(nlmle_t_agb,3),round(nlmle_t_agb_ali,3),round(nlmle_t_agb_aui,3),round(nlmle_t_agb_mli,3),round(nlmle_t_agb_mui,3),round(nlmle_t_agb_li,3),round(nlmle_t_agb_ui,3),round(nlmle_t_u_a,3),round(nlmle_t_u_m,3),round(nlmle_t_u,3))
			colnames(tr) <- c("pid","tid","x","y","d","h","rho",
			"ols_t_agb","ols_t_agb_li","ols_t_agb_ui","ols_t_u",
			"nlmle_t_agb","nlmle_t_agb_ali","nlmle_t_agb_aui","nlmle_t_agb_mli","nlmle_t_agb_mui","nlmle_t_agb_li","nlmle_t_agb_ui","nlmle_t_u_a","nlmle_t_u_m","nlmle_t_u")
			###sort plot results
			pr <- data.frame(pname,round(stem_count,0),round(basal_area,3),round(lorey_height,3),round(wd_baw,3),
			round(ols_p_agb,3),round(ols_p_agb_li,3),round(ols_p_agb_ui,3),round(ols_p_u,3),
			round(nlmle_p_agb,3),round(nlmle_p_agb_ali,3),round(nlmle_p_agb_aui,3),round(nlmle_p_agb_mli,3),round(nlmle_p_agb_mui,3),round(nlmle_p_agb_li,3),round(nlmle_p_agb_ui,3),round(nlmle_p_u_a,3),round(nlmle_p_u_m,3),round(nlmle_p_u,3))
			colnames(pr) <- c("pid","sc","ba","lh","wd_baw",
			"ols_p_agb","ols_p_agb_li","ols_p_agb_ui","ols_p_u",
			"nlmle_p_agb","nlmle_p_agb_ali","nlmle_p_agb_aui","nlmle_p_agb_mli","nlmle_p_agb_mui","nlmle_p_agb_li","nlmle_p_agb_ui","nlmle_p_u_a","nlmle_p_u_m","nlmle_p_u")
			###
			results[[j]] <- tr
			j <- j+1
			results[[j]] <- pr
			j <- j+1
		}
		if(localreg == TRUE)
		{
			lnlmle_t_agb <- numeric(nrow(fdata))
			lnlmle_t_agb_ali <- numeric(nrow(fdata))
			lnlmle_t_agb_aui <- numeric(nrow(fdata))
			lnlmle_t_agb_mli <- numeric(nrow(fdata))
			lnlmle_t_agb_mui <- numeric(nrow(fdata))
			lnlmle_t_agb_li <- numeric(nrow(fdata))
			lnlmle_t_agb_ui <- numeric(nrow(fdata))
			lnlmle_t_u_a <- numeric(nrow(fdata))
			lnlmle_t_u_m <- numeric(nrow(fdata))
			lnlmle_t_u <- numeric(nrow(fdata))
			lnlmle_p_agb <- 0
			lnlmle_p_agb_ali <- 0
			lnlmle_p_agb_aui <- 0
			lnlmle_p_agb_mli <- 0
			lnlmle_p_agb_mui <- 0
			lnlmle_p_agb_li <- 0
			lnlmle_p_agb_ui <- 0
			for(k in 1:nrow(fielddata))
			{
				lcaldata <- head(caldata[order(abs(caldata$X-fdata$X[k])),],nn)
				lfielddata <- data.frame(fielddata[k,])
				lfdata <- data.frame(fdata$X[k])
				colnames(lfdata) <- c("X")
				lmodel <- fitNLS(lcaldata)
				lnlqrmodels <- nlQuantileRegression(lmodel,lcaldata,alpha)
				lbmodels <- bootstrap(lmodel,lcaldata,runs)
				lnlmle_t_agb[k] <- yhat(lmodel,lfdata)
				lnlmle_t_agb_ali[k] <- getPredictionIntervals(lnlqrmodels,lfdata)[,1]
				lnlmle_t_agb_aui[k] <- getPredictionIntervals(lnlqrmodels,lfdata)[,2]
				lnlmle_t_agb_mli[k] <- lnlmle_t_agb[k] - (nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_t_agb_mui[k] <- lnlmle_t_agb[k] + (nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_t_agb_li[k] <- lnlmle_t_agb_ali[k] - (nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_t_agb_ui[k] <- lnlmle_t_agb_aui[k] + (nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_t_u_a[k] <- 0.5 * (lnlmle_t_agb_aui[k] - lnlmle_t_agb_ali[k]) / lnlmle_t_agb[k]
				lnlmle_t_u_m[k] <- 0.5 * (lnlmle_t_agb_mui[k] - lnlmle_t_agb_mli[k]) / lnlmle_t_agb[k]
				lnlmle_t_u[k] <- 0.5 * (lnlmle_t_agb_ui[k] - lnlmle_t_agb_li[k]) / lnlmle_t_agb[k]
				#
				lnlmle_p_agb <- lnlmle_p_agb + lnlmle_t_agb[k]
				lnlmle_p_agb_ali <- lnlmle_p_agb_ali + getConfidenceIntervals(lbmodels,lfdata,alpha)[,1]
				lnlmle_p_agb_aui <- lnlmle_p_agb_aui + getConfidenceIntervals(lbmodels,lfdata,alpha)[,2]
				lnlmle_p_agb_mli <- lnlmle_p_agb_mli + lnlmle_t_agb[k] - (1/sqrt(nrow(fdata)) * nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_p_agb_mui <- lnlmle_p_agb_mui + lnlmle_t_agb[k] + (1/sqrt(nrow(fdata)) * nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_p_agb_li <- lnlmle_p_agb_li + getConfidenceIntervals(lbmodels,lfdata,alpha)[,1] - (1/sqrt(nrow(fdata)) * nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
				lnlmle_p_agb_ui <- lnlmle_p_agb_ui + getConfidenceIntervals(lbmodels,lfdata,alpha)[,2] + (1/sqrt(nrow(fdata)) * nlmle_sigma_meas[k] * qnorm(alpha/2,lower.tail=FALSE))
			}
			lnlmle_p_u_a <- 0.5 * (lnlmle_p_agb_aui - lnlmle_p_agb_ali) / lnlmle_p_agb
			lnlmle_p_u_m <- 0.5 * (lnlmle_p_agb_mui - lnlmle_p_agb_mli) / lnlmle_p_agb
			lnlmle_p_u <- 0.5 * (lnlmle_p_agb_ui - lnlmle_p_agb_li) / lnlmle_p_agb
			###sort tree results
			tr <- data.frame(pname,fielddata$id,fielddata$x,fielddata$y,round(fielddata$d,3),round(fielddata$h,3),round(fielddata$rho,3),
			round(ols_t_agb,3),round(ols_t_agb_li,3),round(ols_t_agb_ui,3),round(ols_t_u,3),
			round(nlmle_t_agb,3),round(nlmle_t_agb_ali,3),round(nlmle_t_agb_aui,3),round(nlmle_t_agb_mli,3),round(nlmle_t_agb_mui,3),round(nlmle_t_agb_li,3),round(nlmle_t_agb_ui,3),round(nlmle_t_u_a,3),round(nlmle_t_u_m,3),round(nlmle_t_u,3),
			round(lnlmle_t_agb,3),round(lnlmle_t_agb_ali,3),round(lnlmle_t_agb_aui,3),round(lnlmle_t_agb_mli,3),round(lnlmle_t_agb_mui,3),round(lnlmle_t_agb_li,3),round(lnlmle_t_agb_ui,3),round(lnlmle_t_u_a,3),round(lnlmle_t_u_m,3),round(lnlmle_t_u,3))
			colnames(tr) <- c("pid","tid","x","y","d","h","rho",
			"ols_t_agb","ols_t_agb_li","ols_t_agb_ui","ols_t_u",
			"nlmle_t_agb","nlmle_t_agb_ali","nlmle_t_agb_aui","nlmle_t_agb_mli","nlmle_t_agb_mui","nlmle_t_agb_li","nlmle_t_agb_ui","nlmle_t_u_a","nlmle_t_u_m","nlmle_t_u",
			"lnlmle_t_agb","lnlmle_t_agb_ali","lnlmle_t_agb_aui","lnlmle_t_agb_mli","lnlmle_t_agb_mui","lnlmle_t_agb_li","lnlmle_t_agb_ui","lnlmle_t_u_a","lnlmle_t_u_m","lnlmle_t_u")
			###sort plot results
			pr <- data.frame(pname,round(stem_count,0),round(basal_area,3),round(lorey_height,3),round(wd_baw,3),
			round(ols_p_agb,3),round(ols_p_agb_li,3),round(ols_p_agb_ui,3),round(ols_p_u,3),
			round(nlmle_p_agb,3),round(nlmle_p_agb_ali,3),round(nlmle_p_agb_aui,3),round(nlmle_p_agb_mli,3),round(nlmle_p_agb_mui,3),round(nlmle_p_agb_li,3),round(nlmle_p_agb_ui,3),round(nlmle_p_u_a,3),round(nlmle_p_u_m,3),round(nlmle_p_u,3),
			round(lnlmle_p_agb,3),round(lnlmle_p_agb_ali,3),round(lnlmle_p_agb_aui,3),round(lnlmle_p_agb_mli,3),round(lnlmle_p_agb_mui,3),round(lnlmle_p_agb_li,3),round(lnlmle_p_agb_ui,3),round(lnlmle_p_u_a,3),round(lnlmle_p_u_m,3),round(lnlmle_p_u,3))
			colnames(pr) <- c("pid","sc","ba","lh","wd_baw",
			"ols_p_agb","ols_p_agb_li","ols_p_agb_ui","ols_p_u",
			"nlmle_p_agb","nlmle_p_agb_ali","nlmle_p_agb_aui","nlmle_p_agb_mli","nlmle_p_agb_mui","nlmle_p_agb_li","nlmle_p_agb_ui","nlmle_p_u_a","nlmle_p_u_m","nlmle_p_u",
			"lnlmle_p_agb","lnlmle_p_agb_ali","lnlmle_p_agb_aui","lnlmle_p_agb_mli","lnlmle_p_agb_mui","lnlmle_p_agb_li","lnlmle_p_agb_ui","lnlmle_p_u_a","lnlmle_p_u_m","lnlmle_p_u")
			###
			results[[j]] <- tr
			j <- j+1
			results[[j]] <- pr
			j <- j+1
		}
	}
	return(results)
}

writeResults <- function(results)
{
	i <- 1
	while(i < length(results))
	{
		pname <- toString(results[[i+1]][1,"pid"])
		tfile <- paste(pname,'_tree.txt',sep='')
		pfile <- paste(pname,'_plot.txt',sep='')
		write.table(results[[i]],file=tfile,col.names=TRUE,row.names=FALSE,quote=FALSE)
		write.table(results[[i+1]],file=pfile,col.names=TRUE,row.names=FALSE,quote=FALSE)
		i <- i + 2
	}
	tsum <- data.frame()
	psum <- data.frame()
	j <- 1
	while(j < length(results))
	{
		tsum <- rbind(tsum,results[[j]])
		psum <- rbind(psum,results[[j+1]])
		j <- j + 2
	}
	write.table(tsum,file="tree_results.txt",col.names=TRUE,row.name=FALSE,quote=FALSE)
	write.table(psum,file="plot_results.txt",col.names=TRUE,row.name=FALSE,quote=FALSE)
}


# To add a new cell, type '# %%'
# To add a new markdown cell, type '# %% [markdown]'
# %% [markdown]
# # Network-based *in silico* drug efficacy screening
#
# This notebook reproduces the proximity-based analysis in the the manuscript. Known drug-disease associations, proximity and relative efficacy values are given in Supplementary Data 1 (the information is combined in "proximity.dat").

# %%
# Read the data file containing information on analyzed drug-disease pairs
# flag: True (known) | CT (unknown & in trials) | False (unknown & not in trials)
# d: average shortest path length from drug targets to closest disease gene
# z: proximity value
# re: relative efficacy value
setwd('/scratch/hdd001/home/haotian/Covid19Datasets/resources/proximity')
file.name = "proximity.dat"
d = read.table(file.name, header=T)
draw.plots = T #F

# %% [markdown]
# ## Area under ROC Curve
# %% [markdown]
# ***AUC using proximity***

# %%
#require(pROC)
f = data.frame(V1=-d$z, V2=ifelse(d$flag=="True", 1, 0))
#auc = auc(f$V2, f$V1)
#print(sprintf("AUC: %.3f", auc))
require(ROCR)
pred = prediction(f$V1, f$V2)
perf = performance(pred, "auc")
print(sprintf("AUC: %.3f", perf@y.values[[1]]))
# Draw the ROC curve (not in the main text)
if(draw.plots) {
   perf = performance(pred, "tpr", "fpr")
   plot(perf)
}

# %% [markdown]
# ***AUC using proximity based similarity***

# %%
file.name = "proximity_similarity.dat"
e = read.table(file.name, header=T)
f = data.frame(V1=e$s, V2=ifelse(e$flag=="True", 1, 0))
#auc = auc(f$V2, f$V1)
#print(sprintf("AUC: %.3f", auc))
pred = prediction(f$V1, f$V2)
perf = performance(pred, "auc")
print(sprintf("AUC: %.3f", perf@y.values[[1]]))
# Draw the ROC curve (not in the main text)
if(draw.plots) {
   perf = performance(pred, "tpr", "fpr")
   plot(perf)
}

# %% [markdown]
# ## Drug target degree vs shortest path and proximity
# ***Correlation between min shortest path length and drug target degree***

# %%
e = d[d$flag == "True",]
selected = "d"
variable = "d.target"
a = cor.test(e[,variable], e[,selected], method="spearman")
print(sprintf("%s vs %s: %.3f %e", selected, variable, a$estimate, a$p.value))

# %% [markdown]
# ***Correlation between proximity and drug target degree***

# %%
e = d[d$flag == "True",]
selected = "z"
variable = "d.target"
a = cor.test(e[,variable], e[,selected], method="spearman")
print(sprintf("%s vs %s: %.3f %e", selected, variable, a$estimate, a$p.value))

# %% [markdown]
# ## Proximity cutoff

# %%
require(ROCR)
f = data.frame(V1=-d$z, V2=ifelse(d$flag=="True", 1, 0))
pred = prediction(f$V1, f$V2)
perf = performance(pred, "sens")
x = -perf@x.values[[1]]
a = perf@y.values[[1]]
# Draw sensitivity - specificity curve
if(draw.plots) {
    plot(x, a, col="blue", type="l", bty="n", lwd=3, xlab="Proximity (z)", ylab="Performance (%)", yaxt='n', cex.lab=1.5, cex.axis=1.5)
    axis(2, labels=F)
    text(par("usr")[1]-0.1, y = seq(0, 1, by=0.2), labels = seq(0, 1, by=0.2), pos=2, cex=1.5, xpd=T)
    perf = performance(pred, "spec")
    b = perf@y.values[[1]]
    lines(x, b, col="green", lwd=3)
    legend("left", legend=c("Sensitivity", "Specificity"), col=c("blue", "green"), lty=c(1,1), lwd=c(3,3), bty="n", cex=1.5)
}
i = which.max(abs(a - b) <= 0.01)
print(sprintf("Cutoff: %.4f", x[i]))
print(sprintf("Sens: %.4f Spec: %.4f", a[i], b[i]))

# Set z cutoff for further analysis
Z.CUTOFF = -0.15
cutoff = Z.CUTOFF

# %% [markdown]
# ## Enrichment of known and clinical trial drug-disease pairs

# %%
e = d
# Draw the bar plots
if(draw.plots) {
    require(ggplot2)
    # Enrichment plot
    e$close = ifelse(e$z<=cutoff, "Proximal", "Distant")
    e$close = factor(e$close, c("Proximal", "Distant"))
    e$known = ifelse(e$flag == "True", "Known", ifelse(e$flag=="CT", "In trials", "Not in trials"))
    e$known = factor(e$known, c("Known", "In trials", "Not in trials"))
    e$ct = ifelse(e$flag == "CT", "in CT", "not in CT")
    p = ggplot(data=e, aes(known)) + geom_bar(aes(fill=close), position="dodge", alpha=0.8) + labs(x="", y="Number of drug-disease pairs") + guides(fill=guide_legend("")) + scale_fill_manual(values=c("blue", "green", "orange"))
    print(p)
}

# %% [markdown]
# ***Known vs Unknown (in trails and not in trials)***

# %%
# Enrichment of known among proximal and distant at proximity cutoff
method = "z"
cutoff = Z.CUTOFF
x=e[e$flag=="True", ]
y=e[e$flag!="True", ]
a = sum(x[[method]]<=cutoff, na.rm=T)
b = sum(y[[method]]<=cutoff, na.rm=T)
contigency = matrix(c(a, nrow(x)-a, b, nrow(y)-b), 2, 2)
s = fisher.test(contigency)
print(sprintf("z <= %.2f: OR = %.3f P = %s", cutoff, s$estimate, s$p.value))
print(contigency)
print(sum(contigency))

# %% [markdown]
# ***In trials vs not in trials***

# %%
# CT enrichment
method = "z"
cutoff = Z.CUTOFF
x=e[e$flag=="CT", ]
y=e[e$flag=="False", ]
a = sum(x[[method]]<=cutoff, na.rm=T)
b = sum(y[[method]]<=cutoff, na.rm=T)
contigency = matrix(c(a, nrow(x)-a, b, nrow(y)-b), 2, 2)
s = fisher.test(contigency)
print(sprintf("CT vs Uknown z <= %.2f: OR = %.3f P = %s", cutoff, s$estimate, s$p.value))
print(contigency)
print(sum(contigency))

# %% [markdown]
# ***Known vs Unknown of significantly proximal***

# %%
# Enrichment of significantly proximal
method = "z"
cutoff = -2
x=e[e$flag=="True", ]
y=e[e$flag!="True", ]
a = sum(x[[method]]<=cutoff, na.rm=T)
b = sum(y[[method]]<=cutoff, na.rm=T)
contigency = matrix(c(a, nrow(x)-a, b, nrow(y)-b), 2, 2)
s = fisher.test(contigency)
print(sprintf("z <= %.2f: OR = %.3f P = %s", cutoff, s$estimate, s$p.value))
print(contigency)
print(sum(contigency))

# %% [markdown]
# ***Pairs in which drug targets are disease genes***

# %%
# Enirchment of pairs in which the drug targets are disease genes
method = "d"
cutoff = 0
a = sum(x[[method]]<=cutoff, na.rm=T) # n.overlap
b = sum(y[[method]]<=cutoff, na.rm=T)
contigency = matrix(c(a, nrow(x)-a, b, nrow(y)-b), 2, 2)
s = fisher.test(contigency)
print(sprintf("n.overlap >= %.2f: OR = %.3f P = %s", cutoff, s$estimate, s$p.value))
print(contigency)
print(sum(contigency))

# %% [markdown]
# ***Pairs with at least one overlapping drug target and disease gene***

# %%
# Enirchment of pairs with at least one overlapping drug target and disease gene
method = "n.overlap"
cutoff = 0
a = sum(x[[method]]>cutoff, na.rm=T) # n.overlap
b = sum(y[[method]]>cutoff, na.rm=T)
contigency = matrix(c(a, nrow(x)-a, b, nrow(y)-b), 2, 2)
s = fisher.test(contigency)
print(sprintf("n.overlap >= %.2f: OR = %.3f P = %s", cutoff, s$estimate, s$p.value))
print(contigency)
print(sum(contigency))

# %% [markdown]
# ## Proximity and drug efficacy

# %%
method = "z"
cutoff = Z.CUTOFF

# Read drug label / palliativeness info
g = read.table("palliative.csv", header=T)
g = g[-(g$label==-1),]

e = d[d$flag=="True",]
e$re = 100*e$re
e = e[order(e$disease, e$group.name),]
e$close = ifelse(e$z<=cutoff, "Proximal", "Distant")
e$palliative = ifelse(g$palliative>0.5, "Palliative", "Non-palliative")
e$label = ifelse(g$label==1, "Label", "Off-label")
e$type = ifelse(e$label=="Label", ifelse(e$palliative=="Palliative", "Palliative", "Non-palliative"), "Off-label")

# %% [markdown]
# *** Relative efficacy (RE) of palliative pairs vs non-palliative and off-label***

# %%
# RE among palliative vs non-palliaitve and off-label
x = e[e$type=="Non-palliative",]
y = e[e$type=="Palliative",]
a = wilcox.test(x$re, y$re, alternative="greater")
print(sprintf("RE non-palliative ~ palliative: %.3f %.3f %s", median(x$re, na.rm=T), median(y$re, na.rm=T), a$p.value))
x = e[e$type=="Off-label",]
a = wilcox.test(x$re, y$re, alternative="greater")
print(sprintf("RE off-label ~ palliative: %.3f %.3f %s", median(x$re, na.rm=T), median(y$re, na.rm=T), a$p.value))

f = e
f$type = factor(f$type, levels=rev(c("Palliative", "Non-palliative", "Off-label")))
if(draw.plots) {
    # Type vs RE
    #p = draw.violinplot(f, "type", "re", "", "Relative Efficacy (%)")
    variable = "type"
    value = "re"
    y.lab = "Relative Efficacy (%)"
    p = ggplot(data=f, aes_string(x=variable, y=value)) + geom_violin(alpha=0.5, color="gray")
    p = p + geom_jitter(alpha=0.5, aes_string(color=variable), position = position_jitter(width = 0.1))
    p = p + stat_summary(fun.y="median", geom='point', color='black', size=10, alpha=0.5) + coord_flip()
    p = p + scale_color_manual(values=c("blue", "green", "orange"))
    p = p + labs(y = y.lab, x="") + guides(color=F)
    print(p)
}

# %% [markdown]
# *** Proximity of palliative pairs vs non-palliative and off-label***

# %%
# Proximity and RE among palliative vs non-palliaitve and off-label
y = e[e$type=="Palliative",]
x = e[e$type=="Non-palliative",]
a = wilcox.test(x$z, y$z, alternative="less")
print(sprintf("z non-palliative ~ palliative: %.3f %.3f %s", median(x$z, na.rm=T), median(y$z, na.rm=T), a$p.value))
x = e[e$type=="Off-label",]
a = wilcox.test(x$z, y$z, alternative="less")
print(sprintf("z off-label ~ palliative: %.3f %.3f %s", median(x$z, na.rm=T), median(y$z, na.rm=T), a$p.value))

if(draw.plots) {
    # Type vs z
    #p = draw.violinplot(f, "type", "z", "", "Proximity (z)") in R_utilities at toolbox
    variable = "type"
    value = "z"
    y.lab = "Proximity (z)"
    p = ggplot(data=f, aes_string(x=variable, y=value)) + geom_violin(alpha=0.5, color="gray")
    p = p + geom_jitter(alpha=0.5, aes_string(color=variable), position = position_jitter(width = 0.1))
    p = p + stat_summary(fun.y="median", geom='point', color='black', size=10, alpha=0.5) + coord_flip()
    p = p + scale_color_manual(values=c("blue", "green", "orange"))
    p = p + labs(y = y.lab, x="") + guides(color=F)
    print(p)
}

# %% [markdown]
# *** RE of proximal vs distant pairs***

# %%
# RE among proximal vs distant
x = e[e$close=="Proximal",]
y = e[e$close=="Distant",]
a = wilcox.test(x$re, y$re, alternative="greater")
print(sprintf("RE proximal ~ distant: %.3f %.3f %s", median(x$re, na.rm=T), median(y$re, na.rm=T), a$p.value))

if(draw.plots) {
    # z vs RE
    e$close = factor(e$close, c("Proximal", "Distant"))
    #p = draw.violinplot(e, "close", "re", "", "Relative Efficacy (%)")
    variable = "close"
    value = "re"
    y.lab = "Relative Efficacy (%)"
    p = ggplot(data=e, aes_string(x=variable, y=value)) + geom_violin(alpha=0.5, color="gray")
    p = p + geom_jitter(alpha=0.5, aes_string(color=variable), position = position_jitter(width = 0.1))
    p = p + stat_summary(fun.y="median", geom='point', color='black', size=10, alpha=0.5) + coord_flip()
    p = p + scale_color_manual(values=c("blue", "green", "orange"))
    p = p + labs(y = y.lab, x="") + guides(color=F)
    print(p)
}

# %% [markdown]
# ***Proximal drugs per disease***

# %%
method = "z"
cutoff = Z.CUTOFF
n.drug.cutoff = 5

e = d[d$flag=="True",]
a = split(e$group, e$disease)
y = sapply(a, length)
f = e[e[[method]]<=cutoff,]
b = split(f$group, f$disease)
x = sapply(b, length)
f = data.frame(disease=names(y), n.known=y, n.close=x, ratio=x/y)

f = f[f$n.known >= n.drug.cutoff,]
n.max = ceiling(max(f$n.known) / 10) * 10
f = f[rev(order(f$ratio, f$n.close, f$n.known)),]

if(draw.plots) {
    #par(mar=c(12, 4, 4, 5) + 0.1)
    a<-barplot(rbind(f$n.close, f$n.known-f$n.close), beside=F, ylab="Number of drugs", xaxt="n", width=1, ylim=c(0,n.max), col=c("blue", "green"))
    x.labels = lapply(f$disease, function(x){ return(gsub("[.]", " ", x)) })

    text(a,  par("usr")[3], labels=x.labels, adj=c(1, 1), srt=45, cex=0.6, xpd=T)
    par(new=T)

    plot(a, 100*f$n.close/f$n.known, col=2, xaxt="n", yaxt="n", xlab="", ylab="", type='b', pch=20, bty="n", lty=2, lwd=2, ylim=c(0,100))
    axis(4, xpd=T, col=2, col.axis=2)
    mtext("Ratio (%)", side=4, line=3, col=2)
    legend("topright", c("Proximal", "Distant", "Ratio"), pch=c(15, 15, 19), pt.cex=2, col=c(c("blue", "green"), 2), lty=c(0,0,1), bty="n")
}

# %% [markdown]
# ***Relative efficacy per disease***

# %%
# Efficacy per disease violin plots
e = d[d$flag=="True",]
e$re = 100*e$re
a = split(e$group, e$disease)
y = sapply(a, length)
f = e[e$z<=cutoff,]
b = split(f$group, f$disease)
x = sapply(b, length)

idx = names(y[y>=n.drug.cutoff])
y = y[idx]
x = x[idx]

idx = rev(order(x/y, x, y))
e$disease = factor(e$disease, levels = factor(names(y[idx])))

if(draw.plots) {
    require(beanplot)
    #par(mar=c(12, 4, 4, 5) + 0.1)
    a = beanplot(re~disease, data=e, col=c("blue", "green"), ylim=c(0,100), ylab="Relative Efficacy (%)", las=2, ll=0.1, cex.axis=0.8, xaxt="n", bw='nrd0')
    x.labels = lapply(names(y[idx]), function(x){ return(gsub("[.]", " ", x)) })
    text(1:length(a$stats), par("usr")[3], labels=x.labels, adj=c(1, 1), srt=45, cex=0.6, xpd=T)
}


#' @export
#' @title pie.add
#' @description couldn't accurately describe
#' @param \code{z}
#' @param \code{x} defaults to 0
#' @param \code{y} defaults to 0
#' @param \code{labels} defaults to names(z)
#' @param \code{radius} defaults to 1
#' @param \code{edges} defaults to 200
#' @param \code{clockwise} defaults to TRUE
#' @param \code{init.angle} defaults to 90
#' @param \code{density} defaults to NULL
#' @param \code{angle} defaults to 45
#' @param \code{col} defaults to NULL
#' @param \code{border} defaults to NULL
#' @param \code{lty} defaults to NULL
#' @param \code{label.dist} defaults to 1.1
#' @family  plotting
#' @author  unknown, \email{<unknown>@@dfo-mpo.gc.ca}
#' @export
pie.add <- function (z, x = 0, y = 0, labels = names(z), radius = 1, edges = 200,
                     clockwise = TRUE, init.angle = 90, density = NULL, angle = 45,
                     col = NULL, border = NULL, lty = NULL, label.dist = 1.1,
                     ...) {
  if (!is.numeric(z) || any(is.na(z) | z < 0))
    stop("'z' values must be positive.")
  if (is.null(labels))
    labels <- as.character(seq_along(z))
  else labels <- as.graphicsAnnot(labels)
  z <- c(0, cumsum(z)/sum(z))
  dz <- diff(z)
  nz <- length(dz)
  asp <- get.asp()
  if (is.null(col))
    col <- if (is.null(density))
      c(rainbow(8,alpha=0.5))
  else par("fg")
  if (!is.null(col))
    col <- rep_len(col, nz)
  if (!is.null(border))
    border <- rep_len(border, nz)
  if (!is.null(lty))
    lty <- rep_len(lty, nz)
  angle <- rep(angle, nz)
  if (!is.null(density))
    density <- rep_len(density, nz)
  twopi <- if (clockwise)
    -2 * pi
  else 2 * pi
  t2xy <- function(t) {
    t2p <- twopi * t + init.angle * pi/180
    list(x = asp * radius * cos(t2p) + x, y = radius * sin(t2p) +
           y)
  }
  for (i in 1L:nz) {
    n1 <- max(2, floor(edges * dz[i]))
    P <- t2xy(seq.int(z[i], z[i + 1], length.out = n1))

    polygon(c(P$x, 0 + x), c(P$y, 0 + y), density = density[i],
            angle = angle[i], border = border[i], col = col[i],
            lty = lty[i])
    P <- t2xy(mean(z[i + 0:1]))
    lab <- as.character(labels[i])
    if (!is.na(lab) && nzchar(lab)) {
      text(label.dist * (P$x - x) + x, label.dist * (P$y -
                                                       y) + y, labels[i], xpd = TRUE, adj = ifelse(P$x -
                                                                                x < 0, 1, 0), ...)
    }
  }
}


#!/usr/bin/env Rscript
#5a_NbClust.r
#
#Author: Abelardo Aguilar Camara
#
#Task performed:
#    This script takes the CPFSCC 28 dimensional vectors of a family,
#    performs 13 different clustering indices, and returns a family
#    membership data frame with consensus clustering
#
#
#NOTE:"CPFSCC" stands for central moments and covariance vectors of cumulative fourier transform power and phase spectra
#
#INPUT: CPFSCC vectors of anytaxon from  ../results/CPFSCC_vectors/anytaxon_CPFSCC_vectors.txt
#
#OUTPUT: Membership vectors        in    ../results/NbClust_membership_vectors/
#        Distance matrix           in    ../results/Distance_Matrices/
#        Point plot of distances   in    ../results/Clustering_graphics
#        Clusters PCA              in    ../results/Clustering_graphics
#
#################################################
#
#5a_NbClust.r
#
#################################################
# Load needed packages
suppressPackageStartupMessages(library("NbClust"))
suppressPackageStartupMessages(library("factoextra"))
suppressPackageStartupMessages(library("ggplot2"))
#################################################
# Reading input
setwd("../results/CPFSCC_vectors") #Location of CPFSCC vectors files
PosArgs <- as.character(commandArgs(trailingOnly = TRUE))
family = PosArgs[1]  #Reads positional arguments to define family (e.g. family = "Geminiviridae")
mincluster = as.numeric(PosArgs[2]) # defines minimum number of clusters
maxcluster = as.numeric(PosArgs[3]) # defines maximum number of clusters
file_suffix <- ("CPFSCC_vectors.txt") #Suffix to build input filename
family_CPFSCC_file <- paste(family,file_suffix, sep = "_") #Read filename
tmp <- as.matrix(read.csv(family_CPFSCC_file, header = TRUE, sep = ",", dec = ".")) #Transform to matrix
tmp2 <- t(tmp)
colnames(tmp2) <- sprintf("D%d", 1:28) #Name columns dimensions
datos <- as.data.frame(tmp2) # Object datos is the final input for the upcoming computations
#################################################
#indices to estimate
indexes <- c("kl",
             "ch",
             "hartigan",
             "cindex",
             "db",
             "silhouette",
             "ratkowsky",
             "ball",
             "ptbiserial",
             "mcclain",
             "dunn",
             "sdindex",
             "sdbw")
diss_matrix<- dist(datos, method = "euclidean", diag=FALSE) #Pairwise distance matrix
#################################################
#Estimating the best number of clusters
res <- c() # Vector to contain results
Best_num_clust <- c() #Vector to contain best number of cluster proposed by each index
#Loop over indices to compute
for (i in indexes){
  res <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "ward.D", index = i)
  Best_num_clust[i] <- res$Best.nc[1]
}
Bestnc <- as.data.frame(Best_num_clust) #to dataframe
nc <- names(sort(summary(as.factor(Bestnc$Best_num_clust)), decreasing = TRUE)[1]) #Saves the most observed number of clusters
ncsubset <- subset(Bestnc, Best_num_clust==nc) #Defines a subset of those indexes indicating matching nc
first_index <- rownames(ncsubset)[1] #Selects index to operate partitions
#################################################
#Computing partitions
#Can be changed to for loop but I prefer this way to easily find when a partition method aborts
Partition_wardD <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "ward.D", index = first_index)
Partition_wardD2 <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "ward.D2", index = first_index)
Partition_single <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "single", index = first_index)
Partition_complete <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "complete", index = first_index)
Partition_average <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "average", index = first_index)
Partition_mcquitty <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "mcquitty", index = first_index)
Partition_median <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "median", index = first_index)
Partition_centroid <- NbClust(datos, diss=diss_matrix, distance = NULL, min.nc=mincluster, max.nc=maxcluster, method = "centroid", index = first_index)
#Save partitions to data frame
taxids <- rownames(datos)
Pertenencia.df <- data.frame(Partition_wardD$Best.partition,
                             Partition_wardD2$Best.partition,
                             Partition_single$Best.partition,
                             Partition_complete$Best.partition,
                             Partition_average$Best.partition,
                             Partition_mcquitty$Best.partition,
                             Partition_median$Best.partition,
                             Partition_centroid$Best.partition)
#################################################
#Define a function "moda" to recover the most popular clusters membership proposition for each genome
function.moda<-function(n)
{
  names(sort(summary(as.factor(n)), decreasing = TRUE)[1])
}
#################################################
#Use function.moda to estimate consensus membership vector
Pertenencia.df$Consenso <- apply(Pertenencia.df, 1, function.moda)
#################################################
#################################################
#Save outputs
#Membership vectors
setwd("..")
outdir <- ("mkdir -p NbClust_membership_vectors") #Make membership vectors outdir
system(outdir) #Make outdir
setwd("NbClust_membership_vectors")
outfile <- paste(family,"membership_vectors.csv", sep="_")
write.csv(Pertenencia.df, outfile, row.names = TRUE) #Write membership vectors file
#################################################
#Distance matrix
outdir2 <- ("mkdir -p ../Distance_Matrices/")
system(outdir2) #Make distance matrices outdir
setwd("../Distance_Matrices/")
Distancias <- dist(datos, method = "euclidean", diag=TRUE) #Estimate pairwise distances
Distancias.matrix <- as.matrix(Distancias) #Transform to matrix
outfile2 <- paste(family,"distance_matrix.csv", sep="_")
write.csv(Distancias.matrix, outfile2, row.names = TRUE) #Write distance matrix as csv
#################################################
#Point plot of ordered distances (pplot)
ene <- nrow(Distancias.matrix) #number of genomes
SumDist <- (rowSums(Distancias.matrix)) #get sum of distances by genome
SumDist.df <- as.data.frame(SumDist) #transform to data frame
mean_Distance <- (SumDist.df$SumDist)/ene #get average sum of distances per genome
SumDist.df$MeanDif <- mean_Distance #save average sum to data frame
SumDist.df <- cbind(SumDist.df, Pertenencia.df)
colnames(SumDist.df) <- c("SumDist",
                          "MeanDif",
                          "ward.D_method",
                          "ward.D2_method",
                          "single_method",
                          "complete_method",
                          "average_method",
                          "mcquitty_method",
                          "median_method",
                          "centroid_method",
                          "consensus_method")
methd <- c(3:11)
SumDist.df[,methd] <- lapply(SumDist.df[,methd] , factor)
OrderedDist.df <- SumDist.df[order(SumDist.df$MeanDif),] #sort sum of distances
IndexByDistSum <- seq(from = 1, to = ene)
OrderedDist.df$IndexByDistSum <- IndexByDistSum
setwd("..")
outdir3 <- ("mkdir -p Clustering_graphics") #make clustering graphics outdir
system(outdir3)
setwd("Clustering_graphics")
GraphMethods <- list(OrderedDist.df$ward.D_method,
                  OrderedDist.df$ward.D2_method,
                  OrderedDist.df$single_method,
                  OrderedDist.df$complete_method,
                  OrderedDist.df$average_method,
                  OrderedDist.df$mcquitty_method,
                  OrderedDist.df$median_method,
                  OrderedDist.df$centroid_method,
                  OrderedDist.df$consensus_method)
for (i in 1:length(GraphMethods)){
  outfilename <- paste(family, "_", colnames(OrderedDist.df)[2+i],"_distances_pplot.tiff", sep = "")
  legendtitle <- paste(family, paste("(", colnames(OrderedDist.df)[2+i], ")", sep = ""),"ordered distances by cluster")
  tiff(outfilename, units="in", width=7, height=5, res=600)
  print(ggplot(data=OrderedDist.df, aes(IndexByDistSum, MeanDif)) +
    geom_point(aes(colour = GraphMethods[[i]]), size = 2, alpha = 0.2) +
    theme(legend.position="top") +
    labs(color= legendtitle ))
  dev.off()
}
#################################################
#Clusters PCA
res.pca <- prcomp(datos, scale = TRUE) #make PCA
groups <- as.factor(Pertenencia.df$Consenso) #recover clusters
outfile4 <- paste(family,"_PCA_clusters.tiff", sep = "")
#Draw PCA to tiff file
#Clusters PCA
res.pca <- prcomp(datos, scale = TRUE) #make PCA
grupos <- list(as.factor(Pertenencia.df$Partition_wardD.Best.partition),
               as.factor(Pertenencia.df$Partition_wardD2.Best.partition),
               as.factor(Pertenencia.df$Partition_single.Best.partition),
               as.factor(Pertenencia.df$Partition_complete.Best.partition),
               as.factor(Pertenencia.df$Partition_average.Best.partition),
               as.factor(Pertenencia.df$Partition_mcquitty.Best.partition),
               as.factor(Pertenencia.df$Partition_median.Best.partition),
               as.factor(Pertenencia.df$Partition_centroid.Best.partition),
               as.factor(Pertenencia.df$Consenso))
for (i in 1:length(grupos)){
  outfilename <- paste(family, "_", colnames(OrderedDist.df)[2+i],"_PCA_clusters.tiff", sep = "")
  PCA_title <- as.character(paste(family, "PCA", paste("(", colnames(OrderedDist.df)[2+i], ")", sep = "")))
  tiff(outfilename, units="in", width=7, height=5, res=600) #tiff resolution parameters
  current_pca <- fviz_pca_ind(res.pca,
               col.ind = grupos[[i]], # color by groups
               palette = c("#00AFBB",
                           "#FC4E07",
                           "#50f0d0",
                           "#f05070",
                           "#d050f0",
                           "#70f050",
                           "#f0d050",
                           "#ffd88d",
                           "#5086a7",
                           "#dd9082",
                           "#093d3a",
                           "#196f69",
                           "#786045",
                           "#d59c2c",
                           "#336936",
                           "#d050f0",
                           "#70f050",
                           "#f0d050",
                           "#50f0d0",
                           "#f05070"),
               addEllipses = TRUE,
               ellipse.level=0.95,
               legend.title = "cluster",
               label = "none",
               title = PCA_title)
  print(current_pca)
  dev.off()
}


#' @title Generate a nested case-control study
#'
#' @description
#' \code{ncc_selection_prob} calculates the probability of selection into a
#' nested case control study. Given a nested case control study with
#' time of entry, time of exit, exit status, and the number of potentially
#' eligible controls for each case from the full cohort, the probability of
#' being selected as a control for one or more cases is calculated.
#'
#' @param entry time of entry to follow-up
#' @param exit time of exit from follow-up
#' @param fail indicator of status on exit from follow-up (censored=0, fail=1)
#' @param origin the origin of the analysis time-scale. For instance, date of
#' birth, if age is the desired time-scale.
#' @param controls the number of controls  that were sampled for each failure
#' @param match a list of categorical variables that were matched on
#' @param caseset a caseset identifier (common within matched case-control sets)
#' @param elig_controls an integer indicating the number of individuals from
#' the full cohort who were eligible be selected as a control for each case.
#' This should be constant within \code{caseset}.
#' @param data a \code{\link{data.frame}} which contains the follow-up,
#' matching, caseset identifier, and number of eligible controls variables.
#' @param silent if \code{FALSE}, provides entertainment by echoing a fullstop
#' to the screen as each risk set is generated. If TRUE, output to the
#' console is suppressed.
#'
#' @return
#' A \code{data.frame} comprising the original data, with a variable indicating
#' the probability of being selected in the nested case control sample (ncc_pr)
#'
#' @author David C Muller
#'
#' @references
#' Samuelsen S. O. (1997). A psudolikelihood approach to analysis of nested
#' case-control studies. Biometrika, 84(2), 379-394. doi:10.1093/biomet/84.2.379
#'
#' @import data.table
#' @export
#'
ncc_selection_prob <- function(entry = 0, exit, fail, origin = 0, controls = 1,
                       match = list(), caseset, elig_controls, data = NULL,
                       silent = FALSE) {
  entry <- eval(substitute(entry), data)
  exit <- eval(substitute(exit), data)
  fail <- eval(substitute(fail), data)
  origin <- eval(substitute(origin), data)
  elig_controls <- eval(substitute(elig_controls), data)
  n <- length(fail)
# individual id for identifying sets of split records
  ncc_id <- 1:n
# add this id to the supplied data frame for later merging
  data <- data.table(data)
  data[, "ncc_id" := ncc_id]
  if (length(exit) != n) {
    stop("All vectors must have length")
  }
  if (length(entry) == 1) {
    entry <- rep(entry, n)
  }
  else {
    if (length(entry) != n) {
      stop("All vectors must have same length")
    }
  }
  if (length(origin) == 1) {
    origin <- rep(origin, n)
  }
  else {
    if (length(origin) != n)
      stop("All vectors must have same length")
  }
  t_entry <- as.numeric(entry - origin)
  t_exit <- as.numeric(exit - origin)

  # process matching vectors
  arg_match <- substitute(match)
  match <- lapply(arg_match, eval, data)
  names(match) <- as.character(arg_match)
  if (names(match[1]) == "list" && class(match[[1]]) == "function") {
    match[1] <- NULL
  }
  match_names <- names(match)

  grp <- rep(1, n)
  if (length(match) > 0) {
    grp <- as.numeric(factor(do.call(paste0, match)))
  }
  fg <- unique(grp[fail != 0])
  gresult <- vector("list", length(fg))
  # index for groups
  gii <- 0
  set <- 0
  ties <- 0
  for (g in fg) {
    gii <- gii + 1
    ft <- unique(t_exit[(grp == g) & (fail != 0)])
    ft <- ft[order(ft)]
    tresult <- vector("list", length(ft))

    # index for times
    for (tii in 1:length(ft)) {
      t_end <- ft[tii]
      case <- (grp == g) & (t_exit == t_end) & (fail != 0)
      ncase <- sum(case)
      if (ncase > 1) {
        ties <- ties + 1
      }
      # split tied failures
      sresult <- vector("list", ncase)
      case_index <- cumsum(case)*case
      for (tjj in 1:ncase) {
        set <- set + 1
        if (!silent) {
          dots <- ifelse((set %% 50 == 0), ". 50\n", ".")
          cat(dots)
        }
        failure <- case & case_index==tjj
        noncase <- (grp == g) & (t_entry <= t_end) & (t_exit >= t_end) & !failure
        n_eligible <- elig_controls[failure==TRUE]
        ineligible <- !failure & !noncase
        df <- data.frame(ncc_id   = ncc_id,
                         ncc_set  = set,
                         ncc_fail = failure,
                         nfail    = rep(sum(failure), length(failure)),
                         ncc_elig_co = rep(n_eligible, length(failure)),
                         ncc_time     = rep(t_end, length(failure)),
                         ineligible   = ineligible)
        sresult[[tjj]] <- df[!ineligible, ]
      }
      tresult[[tii]] <- rbindlist(sresult)
    }
    gresult[[gii]] <- rbindlist(tresult)
    rm(tresult)
  }
  tsplit <- rbindlist(gresult)

# Calculate probability of inclusion
# (Samuelsen, Biometrika, 1997, 84(2) p379)
  tsplit[, "pr_not" := 1 - ((nfail * controls) / ncc_elig_co)]
  tsplit[, "ncc_pr" := 1- prod(pr_not), by=ncc_id]
  tsplit[, "pr_not" := NULL]
  tsplit[, "ncc_elig_co" := NULL]
# cases have probability of inclusion of 1 (comment this out, we still want
# the probability that they were ever included as a control)
#  tsplit[ncc_fail==TRUE, ncc_pr := 1]
  # if risk sets are smaller than requested number of cases, the calculated
# probability will be greater than 1. The correct value is 1.
  tsplit[ncc_pr > 1, ncc_pr := 1]
  # we no longer need the number of failures per set
  #tsplit[, nfail := NULL]

  ncc_frame <- tsplit[, head(.SD, 1), by = list(ncc_id)]

  ncc_frame[, ncc_fail := NULL]

  # merge ncc frame with original data
  setkey(ncc_frame, ncc_id)
  setkey(data, ncc_id)
  res <- data[ncc_frame]

  return(res)
}


#' rand_measure
#'
#' Returns the Rand measure, which is a statistic for comparing two data
#' clusterings.
#'
#' @param l1,l2
#' Two vectors or shaqs (with 1 column each) of the same length. If shaqs, they
#' must be distributed in the exact same way.
#'
#' @return
#' The rand measure (global).
#'
#' @references
#' Rand Index \url{https://en.wikipedia.org/wiki/Rand_index}
#'
#' @export
rand_measure = function(l1, l2)
{
  if (is.shaq(l1) && is.shaq(l2))
    R = rand_measure_shaq(l1, l2)
  else if (isavec(l1) && isavec(l2))
    R = rand_measure_atomic(l1, l2)
  else
    comm.stop("arguments 'l1' and 'l2' should either both be numeric vectors or both be shaqs")

  R
}

isavec = function(x) (is.atomic(x) && is.vector(x))

rand_measure_ab = function(n, l1, l2)
{
  a = b = 0.0

  for (i in 1:n)
  {
    for (j in (i+1L):n)
    {
      if (j > n) # R indexing is stupid
        break

      same1 = (l1[i] == l1[j])
      same2 = (l2[i] == l2[j])

      if (same1 && same2)
        a = a + 1.0
      else if (!same1 && !same2)
        b = b + 1.0
    }
  }

  c(a, b)
}

rand_measure_atomic = function(l1, l2)
{
  n = length(l1)
  if (n != length(l2))
    comm.stop("'l1' and 'l2' must be the same length")

  ab = rand_measure_ab(n, l1, l2)

  a = ab[1]
  b = ab[2]

  R = (a + b) / choose(n, 2)
  R
}

rand_measure_shaq = function(l1, l2)
{
  n = nrow(l1)
  if (nrow(l2) != n || ncol(l1) != 1 || ncol(l2) != 2)
    comm.stop("")

  ab = rand_measure_ab(n, l1, l2)
  ab = allreduce_dbl(ab)

  a = ab[1]
  b = ab[2]

  R = (a + b) / choose(n, 2)
  R
}


# Estimate spatial null models for mean family age and fuse across a global grid of 2x2 (resolution as in the original estimation of the presence absence matrix).
# The estimation of the null models is performed by shuffling entries in the PAM in two ways:
# - NULL 1 randomly shuffles occurrences within sites (i.e., per-row shuffling), keeping the observed species richness.
# - NULL 2 randomly shuffles occurrences within families (i.e., per-column shuffling), keeping the observed family prevalence.
# Continues with R objects generated in previous steps.

MDT_crown_NULL2 <- matrix(nrow = dim(prPAM_g2)[1], ncol = 1002)
for (j in 3:ncol(MDT_crown_NULL2)){
  cat("SIMULATION NULL2",j-2,"\r")
  prPAM_random <- prPAM_g2
  crown <- which(!is.na(resB_fams$CG_age_Acal))
  names_crown <- rownames(resB_fams)[crown]
  mm <- match(names_crown,colnames(prPAM_random))
  prPAM_random[,mm] <- randomizeMatrix(prPAM_random[,mm],null.model="freq")
  prPAM_random_gt <- prPAM_random
  for (i in 3:438){
    prPAM_random_gt[which(prPAM_random_gt[,i]>=1),i] <- resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"CG_age_Acal"]
  }
  richness_tot <- apply(prPAM_random[,-c(1:2)], MARGIN = 1,sum,na.rm=T)
  MDT_mat <- prPAM_random
  MDT_mat[,-c(1:2)] <- NA
  MDT_mat[,-c(1:2)] <- prPAM_random_gt[,-c(1:2)] * prPAM_random[,-c(1:2)]
  MDT_crown_random <- apply(MDT_mat[,-c(1:2)],MARGIN = 1, sum,na.rm=T) /richness_tot
  MDT_crown_NULL2[,j] <- MDT_crown_random
}
save(MDT_crown_NULL2,file=paste(ruta_write,"MDT_crown_NULL2.Rdata"))

MDT_crown_NULL1 <- matrix(nrow = dim(prPAM_g2)[1], ncol = 1002)
for (j in 3:ncol(MDT_crown_NULL1)){
  cat("SIMULATION NULL1:",j-2,"\r")
  prPAM_random <- prPAM_g2
  crown <- which(!is.na(resB_fams$CG_age_Acal))
  names_crown <- rownames(resB_fams)[crown]
  mm <- match(names_crown,colnames(prPAM_random))
  prPAM_random[,mm] <- randomizeMatrix(prPAM_random[,mm],null.model="richness")
  prPAM_random_gt <- prPAM_random
  for (i in 3:438){
    prPAM_random_gt[which(prPAM_random_gt[,i]>=1),i] <- resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"CG_age_Acal"]
  }
  richness_tot <- apply(prPAM_random[,-c(1:2)], MARGIN = 1,sum,na.rm=T)
  MDT_mat <- prPAM_random
  MDT_mat[,-c(1:2)] <- NA
  MDT_mat[,-c(1:2)] <- prPAM_random_gt[,-c(1:2)] * prPAM_random[,-c(1:2)]
  MDT_crown_random <- apply(MDT_mat[,-c(1:2)],MARGIN = 1, sum,na.rm=T) /richness_tot
  MDT_crown_NULL1[,j] <- MDT_crown_random
}
save(MDT_crown_NULL1,file=paste(ruta_write,"MDT_crown_NULL1.Rdata"))

MDT_stem_NULL2 <- matrix(nrow = dim(prPAM_g2)[1], ncol = 1002)
for (j in 3:ncol(MDT_stem_NULL2)){
  cat("SIMULATION NULL2",j-2,"\r")
  prPAM_random <- prPAM_g2
  back_up<-colnames(prPAM_random)[-c(1:2)]
  crown <- which(!is.na(resB_fams$SG_age_Acal))
  names_crown <- rownames(resB_fams)[crown]
  mm <- match(names_crown,colnames(prPAM_random))
  prPAM_random[,mm] <- randomizeMatrix(prPAM_random[,mm],
                                       null.model="freq")
  prPAM_random_gt <- prPAM_random
  for (i in 3:438){
    prPAM_random_gt[which(prPAM_random_gt[,i]>=1),i] <- resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"SG_age_Acal"]
  }
  back_up->colnames(prPAM_random_gt)[-c(1:2)]
  richness_tot <- apply(prPAM_random[,-c(1:2)], MARGIN = 1,sum,na.rm=T)
  MDT_mat <- prPAM_random
  MDT_mat[,-c(1:2)] <- NA
  MDT_mat[,-c(1:2)] <- prPAM_random_gt[,-c(1:2)] * prPAM_random[,-c(1:2)]
  MDT_stem_random <- apply(MDT_mat[,-c(1:2)],MARGIN = 1, sum,na.rm=T) /richness_tot
  MDT_stem_NULL2[,j] <- MDT_stem_random
}
save(MDT_stem_NULL2,file=paste(ruta_write,"MDT_stem_NULL2.Rdata"))

MDT_stem_NULL1 <- matrix(nrow = dim(prPAM_g2)[1], ncol = 1002)
for (j in 3:ncol(MDT_stem_NULL1)){
  cat("SIMULATION NULL1:",j-2,"\r")
  prPAM_random <- prPAM_g2
  back_up<-colnames(prPAM_random)[-c(1:2)]
  crown <- which(!is.na(resB_fams$SG_age_Acal))
  names_crown <- rownames(resB_fams)[crown]
  mm <- match(names_crown,colnames(prPAM_random))
  prPAM_random[,mm] <- randomizeMatrix(prPAM_random[,mm],
                                       null.model="richness")
  prPAM_random_gt <- prPAM_random
  for (i in 3:438){
    prPAM_random_gt[which(prPAM_random_gt[,i]>=1),i] <- resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"SG_age_Acal"]
  }
  back_up->colnames(prPAM_random_gt)[-c(1:2)]
  richness_tot <- apply(prPAM_random[,-c(1:2)], MARGIN = 1,sum,na.rm=T)
  MDT_mat <- prPAM_random
  MDT_mat[,-c(1:2)] <- NA
  MDT_mat[,-c(1:2)] <- prPAM_random_gt[,-c(1:2)] * prPAM_random[,-c(1:2)]
  MDT_stem_random <- apply(MDT_mat[,-c(1:2)],MARGIN = 1, sum,na.rm=T) /richness_tot
  MDT_stem_NULL1[,j] <- MDT_stem_random
}
save(MDT_stem_NULL1,file=paste(ruta_write,"MDT_stem_NULL1.Rdata"))

MDT_fuse_non_NULL2 <- matrix(nrow = dim(prPAM_g2)[1], ncol = 1002)
for (j in 3:ncol(MDT_fuse_non_NULL2)){
  cat("SIMULATION NULL2",j-2,"\r")
  prPAM_random <- prPAM_g2
  back_up<-colnames(prPAM_random)[-c(1:2)]
  crown <- which(!is.na(resB_fams$SG_age_Acal))
  names_crown <- rownames(resB_fams)[crown]
  mm <- match(names_crown,colnames(prPAM_random))
  prPAM_random[,mm] <- randomizeMatrix(prPAM_random[,mm],
                                       null.model="freq")
  prPAM_random_gt <- prPAM_random
  for (i in 3:438){
    prPAM_random_gt[which(prPAM_random_gt[,i]>=1),i] <- (resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"SG_age_Acal"]-resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"CG_age_Acal"])
  }
  back_up->colnames(prPAM_random_gt)[-c(1:2)]
  richness_tot <- apply(prPAM_random[,-c(1:2)], MARGIN = 1,sum,na.rm=T)
  MDT_mat <- prPAM_random
  MDT_mat[,-c(1:2)] <- NA
  MDT_mat[,-c(1:2)] <- prPAM_random_gt[,-c(1:2)] * prPAM_random[,-c(1:2)]
  MDT_fuse_random <- apply(MDT_mat[,-c(1:2)],MARGIN = 1, sum,na.rm=T) /richness_tot
  MDT_fuse_non_NULL2[,j] <- MDT_fuse_random
}
save(MDT_fuse_non_NULL2,file=paste(ruta_write,"MDT_fuse_non_NULL2.Rdata"))
MDT_fuse_non_NULL1 <- matrix(nrow = dim(prPAM_g2)[1], ncol = 1002)
for (j in 3:ncol(MDT_fuse_non_NULL1)){
  cat("SIMULATION NULL1:",j-2,"\r")
  prPAM_random <- prPAM_g2
  back_up<-colnames(prPAM_random)[-c(1:2)]
  crown <- which(!is.na(resB_fams$SG_age_Acal))
  names_crown <- rownames(resB_fams)[crown]
  mm <- match(names_crown,colnames(prPAM_random))
  prPAM_random[,mm] <- randomizeMatrix(prPAM_random[,mm],
                                       null.model="richness")
  prPAM_random_gt <- prPAM_random
  for (i in 3:438){
    prPAM_random_gt[which(prPAM_random_gt[,i]>=1),i] <- (resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"SG_age_Acal"]-resB_fams[rownames(resB_fams)==colnames(prPAM_random_gt)[i],"CG_age_Acal"])
  }
  back_up->colnames(prPAM_random_gt)[-c(1:2)]
  richness_tot <- apply(prPAM_random[,-c(1:2)], MARGIN = 1,sum,na.rm=T)
  MDT_mat <- prPAM_random
  MDT_mat[,-c(1:2)] <- NA
  MDT_mat[,-c(1:2)] <- prPAM_random_gt[,-c(1:2)] * prPAM_random[,-c(1:2)]
  MDT_fuse_random <- apply(MDT_mat[,-c(1:2)],MARGIN = 1, sum,na.rm=T) /richness_tot
  MDT_fuse_non_NULL1[,j] <- MDT_fuse_random
}
save(MDT_fuse_non_NULL1,file=paste(ruta_write,"MDT_fuse_non_NULL1.Rdata"))

pdf(paste(ruta_write,"SES-MAF_NULL2.pdf",sep=""))
par(mfrow = c(1, 1), pty = "s")
null_crown <- MDT_crown_NULL2
null_stem <- MDT_stem_NULL2
null_fuse <- MDT_fuse_non_NULL2
iden <- which(rowSums(prPAM_g2[,-c(1:2)]) > 5)
ssp<-coordinates(prPAM_g2[iden,1:2])
SES_crownMAF <- (MDT_crown - apply(null_crown[,-c(1:2)],MARGIN=1,FUN=mean,na.rm=T))/(apply(null_crown[,-c(1:2)],MARGIN=1,FUN=sd,na.rm=T))
Data<-data.frame(Order=1:dim(prPAM_g2[iden,])[1],z=SES_crownMAF)
Data<-Data[order(Data$z),]
f <- clean_coordinates(as.data.frame(ssp), lon = "LON",lat = "LAT",
                       tests="seas",species=NULL,value="flagged",seas_scale=110)
k=20
breaks <- getJenksBreaks(Data$z,k)
cols_breaks <- (colorRampPalette(c("aliceblue","lightblue","dodgerblue3","darkblue"),bias=1.1)(k))
Data$col <- cols_breaks[k]
for (i in k:1){
  Data$col[which(Data$z <= breaks[i])] <- cols_breaks[i]}
orderedcolors<-Data[order(Data$Order),"col"]
maps::map("world",interior=F)
points(ssp[f,],pch=22,lwd=0.1,bg=orderedcolors[f],cex=0.5,col=NULL)
rect(xleft =-180,xright = 190,ybottom = -86,ytop = 85,border = "black")
color.legend(xl=-35,xr=35,yb=-83,yt=-67,align="lt",gradient="x",
             legend=round(range(SES_crownMAF,na.rm = T),2),rect.col=cols_breaks)
title("SES-crown MAF")

SES_stemMAF <- (MDT_stem - apply(null_stem[,-c(1:2)],MARGIN=1,FUN=mean,na.rm=T))/(apply(null_stem[,-c(1:2)],MARGIN=1,FUN=sd,na.rm=T))
hist(SES_stemMAF)
Data<-data.frame(Order=1:dim(prPAM_g2[iden,])[1],z=SES_stemMAF)
Data<-Data[order(Data$z),]
k=20
breaks <- getJenksBreaks(Data$z,k)
cols_breaks <- (colorRampPalette(c("linen","lightpink","red3","darkred"),bias=1.1)(k))
Data$col <- cols_breaks[k]
for (i in k:1){
  Data$col[which(Data$z <= breaks[i])] <- cols_breaks[i]}
orderedcolors<-Data[order(Data$Order),"col"]
maps::map("world",interior=F)
points(ssp[f,],pch=22,lwd=0.1,bg=orderedcolors[f],cex=0.5,col=NULL)
rect(xleft =-180,xright = 190,ybottom = -86,ytop = 85,border = "black")
color.legend(xl=-35,xr=35,yb=-83,yt=-67,align="lt",gradient="x",
             legend=round(range(SES_stemMAF,na.rm = T),2),rect.col=cols_breaks)
title("SES-stem MAF")

SES_fuseMAF <- (MDT_fuse_non - apply(null_fuse[,-c(1:2)],MARGIN=1,FUN=mean,na.rm=T))/(apply(null_fuse[,-c(1:2)],MARGIN=1,FUN=sd,na.rm=T))
Data<-data.frame(Order=1:dim(prPAM_g2[iden,])[1],z=SES_fuseMAF)
Data<-Data[order(Data$z),]
k=20
breaks <- getJenksBreaks(Data$z,k)
cols_breaks <- (colorRampPalette(c("honeydew","darkseagreen1","forestgreen","black"),bias=1)(k))
Data$col <- cols_breaks[k]
for (i in k:1){
  Data$col[which(Data$z <= breaks[i])] <- cols_breaks[i]}
orderedcolors<-Data[order(Data$Order),"col"]
maps::map("world",interior=F)
points(ssp[f,],pch=22,lwd=0.1,bg=orderedcolors[f],cex=0.5,col=NULL)
rect(xleft =-180,xright = 190,ybottom = -86,ytop = 85,border = "black")
color.legend(xl=-35,xr=35,yb=-83,yt=-67,align="lt",gradient="x",
             legend=round(range(SES_fuseMAF,na.rm = T),2),rect.col=cols_breaks)
title("SES-fuse")
dev.off()

pdf(paste(ruta_write,"SES-MAF_NULL1.pdf",sep=""))
par(mfrow = c(1, 1), pty = "s")
null_crown <- MDT_crown_NULL1
null_stem <- MDT_stem_NULL1
null_fuse <- MDT_fuse_non_NULL1
iden <- which(rowSums(prPAM_g2[,-c(1:2)]) > 5)
ssp<-coordinates(prPAM_g2[iden,1:2])
SES_crownMAF <- (MDT_crown - apply(null_crown[,-c(1:2)],MARGIN=1,FUN=mean,na.rm=T))/(apply(null_crown[,-c(1:2)],MARGIN=1,FUN=sd,na.rm=T))
Data<-data.frame(Order=1:dim(prPAM_g2[iden,])[1],z=SES_crownMAF)
Data<-Data[order(Data$z),]
f <- clean_coordinates(as.data.frame(ssp), lon = "LON",lat = "LAT",
                       tests="seas",species=NULL,value="flagged",seas_scale=110)
k=20
breaks <- getJenksBreaks(Data$z,k)
cols_breaks <- (colorRampPalette(c("aliceblue","lightblue","dodgerblue3","darkblue"),bias=1.1)(k))
Data$col <- cols_breaks[k]
for (i in k:1){
  Data$col[which(Data$z <= breaks[i])] <- cols_breaks[i]}
orderedcolors<-Data[order(Data$Order),"col"]
maps::map("world",interior=F)
points(ssp[f,],pch=22,lwd=0.1,bg=orderedcolors[f],cex=0.5,col=NULL)
rect(xleft =-180,xright = 190,ybottom = -86,ytop = 85,border = "black")
color.legend(xl=-35,xr=35,yb=-83,yt=-67,align="lt",gradient="x",
             legend=round(range(SES_crownMAF,na.rm = T),2),rect.col=cols_breaks)
title("SES-crown MAF")

SES_stemMAF <- (MDT_stem - apply(null_stem[,-c(1:2)],MARGIN=1,FUN=mean,na.rm=T))/(apply(null_stem[,-c(1:2)],MARGIN=1,FUN=sd,na.rm=T))
hist(SES_stemMAF)
Data<-data.frame(Order=1:dim(prPAM_g2[iden,])[1],z=SES_stemMAF)
Data<-Data[order(Data$z),]
k=20
breaks <- getJenksBreaks(Data$z,k)
cols_breaks <- (colorRampPalette(c("linen","lightpink","red3","darkred"),bias=1.1)(k))
Data$col <- cols_breaks[k]
for (i in k:1){
  Data$col[which(Data$z <= breaks[i])] <- cols_breaks[i]}
orderedcolors<-Data[order(Data$Order),"col"]
maps::map("world",interior=F)
points(ssp[f,],pch=22,lwd=0.1,bg=orderedcolors[f],cex=0.5,col=NULL)
rect(xleft =-180,xright = 190,ybottom = -86,ytop = 85,border = "black")
color.legend(xl=-35,xr=35,yb=-83,yt=-67,align="lt",gradient="x",
             legend=round(range(SES_stemMAF,na.rm = T),2),rect.col=cols_breaks)
title("SES-stem MAF")

SES_fuseMAF <- (MDT_fuse_non - apply(null_fuse[,-c(1:2)],MARGIN=1,FUN=mean,na.rm=T))/(apply(null_fuse[,-c(1:2)],MARGIN=1,FUN=sd,na.rm=T))
Data<-data.frame(Order=1:dim(prPAM_g2[iden,])[1],z=SES_fuseMAF)
Data<-Data[order(Data$z),]
k=20
breaks <- getJenksBreaks(Data$z,k)
cols_breaks <- (colorRampPalette(c("honeydew","darkseagreen1","forestgreen","black"),bias=1)(k))
Data$col <- cols_breaks[k]
for (i in k:1){
  Data$col[which(Data$z <= breaks[i])] <- cols_breaks[i]}
orderedcolors<-Data[order(Data$Order),"col"]
maps::map("world",interior=F)
points(ssp[f,],pch=22,lwd=0.1,bg=orderedcolors[f],cex=0.5,col=NULL)
rect(xleft =-180,xright = 190,ybottom = -86,ytop = 85,border = "black")
color.legend(xl=-35,xr=35,yb=-83,yt=-67,align="lt",gradient="x",
             legend=round(range(SES_fuseMAF,na.rm = T),2),rect.col=cols_breaks)
title("SES-fuse")
dev.off()


#!/usr/bin/env Rscript
#
# File Name  : asequan.r
# Author     : Zhenhua Zhang
# E-mail     : zhenhua.zhang217@gmail.com
# Created on : Fri 25 Jan 2019 10:51:30 AM CET
# Version    : v0.1.0
# License    : MIT
#

# Example
# Rscript asequan.r input.tsv output.tsv

rm(list = ls())

library(dplyr)

# Customized paste function
ppaste <- function(...){ paste(..., sep = "/") }

# Log-likelihood function under Binomial distribution
bn_llik <- function(p, alts, refs) {
  totals <- alts + refs
  r <- -sum(lchoose(totals, alts) + alts * log(p) + refs * log(1 - p))

  if (r == -Inf){
    return(-.Machine$integer.max)
  } else if (r == Inf){
    return(.Machine$integer.max)
  } else {
    return(r)
  }
}

# Likelihood ratio test under binomial likelihood
bn_lrt <- function(alt_counts, ref_counts) {

  nul_p <- 0.5
  nul_llik <- bn_llik(nul_p, alt_counts, ref_counts)  # Filter out alt == ref

  alt_p <- exp(mean(log(alt_counts / ref_counts)))
  alt_llik_opt <- optim(
    par = alt_p, fn = bn_llik, NULL,
      method = "L-BFGS-B", lower = c(1e-12), upper = c(1),
      control = list(maxit = 10000),
      alts = alt_counts, refs = ref_counts
  )

  alt_p <- alt_llik_opt$par
  alt_llik <- alt_llik_opt$value

  dof <- length(alt_counts) - 1

  if (dof == 0) {
    chisq <- NA
    chisq_p <- NA
  } else {
    chisq <- - 2 * (alt_llik - nul_llik)
    chisq_p <- 1 - pchisq(chisq, dof)
  }

  r <- list(nul_bn_llik=nul_llik, alt_bn_llik=alt_llik, p_value = chisq_p)
  return(r)
}

# Log-likelihood function under Beta-Binomial distribution
bb_llik <- function(p_od, alts, refs){
  alts_len <- length(alts)
  refs_len <- length(refs)
  if(alts_len != refs_len) {
    stop("alts and refs should have identical length...")
  }

  if (length(p_od) == 1){
    p <- 0.5
    od <- p_od[[1]]
  } else if (length(p_od) == 2){
    p <- p_od[[1]]
    od <- p_od[[2]]
  } else {
    stop("p_od should be no more than 2 elements")
  }

  if (p > 1 || p < 0){
    stop("The first element of p_od should be a decimal between 0 to 1")
  }

  if (od <= 0) {
	stop("The second element of p_od should be positive, but get ", od)
  }

  a <- od * p
  b <- od * (1 - p)
  r <- -sum(lbeta(alts + a, refs + b) - lbeta(a, b) + lchoose(alts + refs, alts))

  if (r == -Inf){
    return(-.Machine$integer.max)
  } else if (r == Inf){
    return(.Machine$integer.max)
  } else {
    return(r)
  }
}

# Likelihood ratio test under Beta-Binomial distribution
bb_lrt <- function(alt_counts, ref_counts){
  alts_len <- length(alt_counts)
  refs_len <- length(ref_counts)

  if(alts_len != refs_len) {
    stop("alt_counts and ref_counts should have identical length...")
  }

  alt_od_init <- sum(alt_counts + ref_counts)  # Using total reads counts
  alt_p_init <- exp(mean(log(alt_counts / ref_counts)))
  alt_opt <- optim(
    par = c(alt_p_init, alt_od_init), fn = bb_llik, NULL,
    method = "L-BFGS-B", lower = c(1e-12, 1e-10), upper = c(1, 1e10),
    control = list(maxit = 10000),
    alts = alt_counts, refs = ref_counts
  )

  alt_par <- alt_opt$par
  alt_llik <- alt_opt$val

  # Using optim to find the over-dispersion parameter under null model
  # nul_opt <- optim(
  #   par = c(10), fn = bb_llik, NULL,
  #     method = "L-BFGS-B", lower = c(1e-12), upper = c(1e12),
  #     control = list(maxit = 10000),
  #     alts = alt_counts, refs = ref_counts
  # )
  # nul_par <- nul_opt$par
  # nul_llik <- nul_opt$val

  # nul_par <- alt_par[[2]]  # Using the over-dispersion parameter derived from samples
  nul_par <- sum(alt_counts + ref_counts)  # Using average total reads counts
  nul_llik <- bb_llik(c(0.5, nul_par), alt_counts, ref_counts)
  nul_llik <- nul_llik

  dof <- alts_len - 1
  if (dof == 0) {
    chisq <- NA
    chisq_p <- NA
  } else {
    chisq <- - 2 * (alt_llik - nul_llik)
    chisq_p <- 1 - pchisq(chisq, dof)
  }

  r <- list(
    nul_p=0.5, nul_od=nul_par, nul_llik_bb=nul_llik,
    alt_p=alt_par[[1]], alt_od=alt_par[[2]], alt_llik_bb=alt_llik,
    p_value = chisq_p
  )

  return(r)
}

trans_into_bin <- function(rtb, cr, pv = 0.05, min_dep = 10, min_dep_per = 3){
  cat("FILTER: BY refCountsBios and altCountsBios ...\n")
  gp <- rtb %>%
    filter(
    (refCountsBios >= min_dep_per)
    & (altCountsBios >= min_dep_per)
    & (refCountsBios != altCountsBios)
    & (refCountsBios + altCountsBios >= min_dep)
    & (!is.na(GeneID))
    & (!is.na(FeatureID))
    & (!is.na(GeneName))
  )

  cat("MUTATE: ADD bn_p, bb_p, group_size, and log2FC ...\n")
  gp <- gp %>%
    group_by(Chrom, Pos, Ref, Alt, GeneID) %>%
    mutate(
      log2FC = log2(sum(altCountsBios) / sum(refCountsBios)),
    bn_p = bn_lrt(altCountsBios, refCountsBios)$p_value,
    bb_p = bb_lrt(altCountsBios, refCountsBios)$p_value,
    group_size = n_distinct(sampleBios)
    )

  cat("MUTATE: ADD bn_p_adj, bn_ASE, bb_p_adj, bb_ASE ...\n")
  gp <- gp %>%
    ungroup() %>%
    mutate(
      bn_p_adj = p.adjust(bn_p, method = "fdr"),
      bb_p_adj = p.adjust(bb_p, method = "fdr"),
      bn_ASE = ifelse(bn_p_adj < pv, ifelse(log2FC < 0, -1, 1), 0),
      bb_ASE = ifelse(bb_p_adj < pv, ifelse(log2FC < 0, -1, 1), 0)
    ) %>%
    select(rn)

  cat("ARRANGE: Filter ...\n")
  gp <- gp %>%
    arrange(Chrom, Pos, Ref, Alt) %>%
    distinct() %>%
    as.data.frame()

  return(gp)
}

# Deal with input and output
args <- commandArgs(trailingOnly = TRUE)
if(length(args) == 2){
  ipf <- args[1]
  opf <- args[2]
} else {
  stop("Wrong number of arguments. Require 2...\n")
}

# Deal with `tar` compressed files
if (endsWith(ipf, '.tar.gz') || endsWith(ipf, '.tgz')){ ipf <- untar(ipf) }
rtb <- read.csv(ipf, header = TRUE, sep = "\t")
rn <- c(
  colnames(rtb)[1:107], "log2FC", "bn_p", "bn_p_adj", "bb_p", "bb_p_adj",
  "group_size", "bn_ASE", "bb_ASE"
)

# Calculate
odf <- trans_into_bin(rtb, rn)
# cat(dim(odf[which(abs(odf$bn_ASE)==1), ]), " in bn_ASE(1, -1)\n")
# cat(dim(odf[which(abs(odf$bb_ASE)==1), ]), " in bb_ASE(1, -1)\n")

# Dump out results
write.table(odf, file = opf, quote = FALSE, sep = "\t", row.names = FALSE)