#' ddirch_obs_uncertainty
#' This function takes dirichlet variances at the site, plot and core level.
#' It also requires a matrix of species abundances to represent "true" values at the site level.
#' It starts at the site level, and draws plots within site from the site mean, and within site-scale variance.
#' It then draws cores within plots, given the "true" plot means, and within plot-scale variance.
#' It then draws estimates of observed core values, given the "true" core mean and within core variance.
#' Once all this data is drawn, the function hierarchically estimates core, plot and site means (y_obs)
#' It then regresses y_obs at each scale against the "true" values that were drawn, and estimates an R2 value.
#' This is repeated some number of times to generate a distribution of potential R2 values.
#' This reflects the maximum predictive accuracy once could expect, given observation uncertainty, in principle a perfect model would predict the "true" values generated here.
#' Predictaibility likely varies as a function of abundance, as well as variance, and their interaction.
#' Predictiability will also vary as a function scale.
#' NOTE: larger variance values lead to less variance in the dirichlet.
#'
#' @param site_mu     #matrix of site-scale relative abundances. Number of rows = number of sites.
#' @param site.var    #within site-scale variance.
#' @param plot.var    #within plot-scale variance.
#' @param core.var    #within core-scale variance.
#' @param n.sim       #numbre of simulations to run.
#' @param n.plot      #Number of plots within site. Default 10 for NEON.
#' @param n.core      #Number of cores within plot. Default  3 for NEON.
#' @param n.threads   #Number of cores to run in parallel. Default will auto-detect.
#'
#' @return
#' @export
#'
#' @examples
ddirch_obs_uncertainty <- function(site_mu,
                                   site.var, plot.var, core.var,
                                   n.sim = 100, n.plot = 10, n.core = 3, n.threads = NA){
  #function tests to throw errors.----
  check <- "DirichletReg" %in% rownames(installed.packages())
  if(check == F){stop('The package DirichletReg is not installed. We need it to be.')}
  check <- "doParallel" %in% rownames(installed.packages())
  if(check == F){stop('The package doParallel is not installed. We need it to be.')}
  library(doParallel)

  #setup parallel.----
  if( is.na(n.threads)){
    n.cores <- detectCores()
    registerDoParallel(n.cores)
  }
  if(!is.na(n.threads)){
    registerDoParallel(n.threads)
  }

  #execute loop in parallel.----
  all.output <-
    foreach(j = 1:n.sim) %dopar% {
      #begin try-catch loop (sometimes we fail some type of hessian matrix calc...)----
      attempt = 0
      mod.plot <- NULL
      mod.site <- NULL
      while(is.null(mod.plot) && is.null(mod.site) && attempt <= 10){
        attempt = attempt + 1
        #1. Generate plot and core values.----
        #assign site labels.
        n.site <- nrow(site_mu)
        sites <- letters[1:n.site]

        #generate plot values.
        y.plot <- list()
        for(i in 1:nrow(site_mu)){
          plots <- DirichletReg::rdirichlet(n.plot, site_mu[i,]*site.var)
          #kill zeros.
          plots[plots == 0] <- min(plots[plots > 0])
          plot.lab <- paste0(sites[i],c(1:n.plot))
          plots <- data.frame(plot.lab, plots)
          colnames(plots) <- c('plotID',colnames(site_mu))
          y.plot[[i]] <- plots
        }
        y.plot <- do.call(rbind, y.plot)
        y.plot$siteID <- substr(y.plot$plotID, 1,1)

        #Generate core values.
        y.core <- list()
        for(i in 1:nrow(y.plot)){
          cores <- DirichletReg::rdirichlet(n.core, as.numeric(y.plot[i,grep('y',colnames(y.plot))]*plot.var))
          #kill zeros
          cores[cores == 0] <- min(cores[cores > 0])
          core.lab <- paste0(y.plot[i,'plotID'], '.', c(1:n.core))
          cores <- data.frame(core.lab, cores)
          colnames(cores) <- c('coreID',colnames(site_mu))
          y.core[[i]] <- cores
        }
        y.core <- do.call(rbind, y.core)
        y.core$coreID <- as.character(y.core$coreID)
        y.core$plotID <- substr(y.core$coreID, 1, nchar(y.core$coreID) - 2)
        y.core$siteID <- substr(y.core$coreID, 1,1)

        #Generate observed data by drawing from intra-core variance.
        y.obs <- DirichletReg::rdirichlet(nrow(y.core), as.matrix(y.core[,grep('y',colnames(y.core))]) * core.var)
        y.obs <- y.obs + min(y.obs[y.obs > 0])
        y.obs <- data.frame(y.obs)
        colnames(y.obs) <- colnames(y.core)[grep('y',colnames(y.core))]
        y.obs$coreID <- y.core$coreID
        y.obs$plotID <- substr(y.obs$coreID, 1, nchar(y.obs$coreID) - 2)

        #2. Aggregate to plot and site scale.----
        #I should maybe use JAGS here instead...
        #Plot scale aggregation.
        suppressWarnings(
          y.obs$Y <- DR_data(y.obs[,grep('y',colnames(y.obs))])
        )
        try(mod.plot <- DirichletReg::DirichReg(Y ~ plotID - 1, data = y.obs))
        plot.obs <- matrix(mod.plot$coefficients, ncol = ncol(site_mu))
        plot.obs <- exp(plot.obs)
        plot.obs <- data.frame(plot.obs / rowSums(plot.obs))
        colnames(plot.obs) <- colnames(site_mu)
        plot.obs$plotID <- unique(y.obs$plotID)
        plot.obs$siteID <- substr(plot.obs$plotID,1,1)

        #Site scale aggregation.
        suppressWarnings(
          plot.obs$Y <- DR_data(plot.obs[,grep('y',colnames(plot.obs))])
        )
        try(mod.site <- DirichletReg::DirichReg(Y ~ siteID - 1, data = plot.obs))
        site.obs <- matrix(mod.site$coefficients, ncol = ncol(site_mu))
        site.obs <- exp(site.obs)
        site.obs <- data.frame(site.obs / rowSums(site.obs))
        colnames(site.obs) <- colnames(site_mu)
        site.obs$siteID <- unique(plot.obs$siteID)
        #end try-catch loop.----
      }
      #3. Calculate r-squared values of 'true' vs. observed at core-plot-site scales.----
      #core-scale rsq values.
      core.rsq <- list()
      for(i in 1:ncol(site_mu)){
        pred <-  y.obs[,grep('y',colnames(y.obs))]
        obs <- y.core[,grep('y',colnames(y.core))]
        mod <- lm(obs[,i] ~ pred[,i])
        core.rsq[[i]] <- summary(mod)$r.squared
      }
      core.rsq <- unlist(core.rsq)
      names(core.rsq) <- colnames(site_mu)

      #plot-scale rsq values.
      plot.rsq <- list()
      for(i in 1:ncol(site_mu)){
        pred <-  plot.obs[,grep('y',colnames(plot.obs))]
        obs <- y.plot[,grep('y',colnames(y.plot))]
        mod <- lm(obs[,i] ~ pred[,i])
        plot.rsq[[i]] <- summary(mod)$r.squared
      }
      plot.rsq <- unlist(plot.rsq)
      names(plot.rsq) <- colnames(site_mu)

      #site-scale rsq values.
      site.rsq <- list()
      for(i in 1:ncol(site_mu)){
        pred <- site.obs[,grep('y',colnames(site.obs))]
        mod <- lm(site_mu[,i] ~ pred[,i])
        site.rsq[[i]] <- summary(mod)$r.squared
      }
      site.rsq <- unlist(site.rsq)
      names(site.rsq) <- colnames(site_mu)

      #wrap rsq output into list and return.
      rsq.out <- list(core.rsq, plot.rsq, site.rsq)
      names(rsq.out) <- c('core','plot','site')
      return(rsq.out)
    }

  #clean up all rsq output.----
  core.output <- list()
  plot.output <- list()
  site.output <- list()
  for(i in 1:length(all.output)){
    core.output[[i]] <- all.output[[i]]$core
    plot.output[[i]] <- all.output[[i]]$plot
    site.output[[i]] <- all.output[[i]]$site
  }
  core.output <- do.call(rbind, core.output)
  plot.output <- do.call(rbind, plot.output)
  site.output <- do.call(rbind, site.output)
  output <- list(core.output,plot.output,site.output)
  names(output) <- c('core','plot','site')

  #return output.----
  return(output)
} #end function.


#' Sorensen dissimilarity index
#'
#' @param comm Community matrix giving sites (rows) by species (columns).
#' @param partition Logical; if true, the index will be partitioned into turnover and
#' 			nestedness.
#' @references Baselga A. 2010. Partitioning the turnover and nestedness components of
#' 		beta diversity. Global Ecol. Biogeogr. 19: 134–143
#' @references Sørensen, T. (1948) A method of establishing groups of equal amplitude in
#' 		plant sociology based on similarity of species content. Kongelige Danske
#' 		Videnskabernes Selskabs Biologiske Skrifter, 5, 1–34.
#' @return If \code{partition == TRUE}, list of three matrices giving the turnover,
#' 		nestedness, and total components of dissimilarity. Otherwise, a matrix of
#' 		pairwise dissimilarity values between communities.
#' @export
sorensen <- function(comm, partition = FALSE)
{
	if(!is.matrix(comm)) comm <- as.matrix(comm)

	a <- comm %*% t(comm > 0)
	b <- comm %*% t(!comm)
	cc <- t(b)
	bsor <- (b + cc) / (2*a + b + cc)

	if(partition) {
		mbc <- pmin(b, cc)
		bturn <- mbc / (a + mbc)
		bnes <- bsor - bturn
		return(list(turnover = bturn, nestedness = bnes, sorensen = bsor))
	} else
		return(bsor)
}

#' Phylogenetic sorensen dissimilarity index
#'
#' @param comm Community matrix giving sites (rows) by species (columns).
#' @param tree Dendrogram of class \code{phylo}
#' @param partition Logical; if true, the index will be partitioned into turnover and
#' 			nestedness.
#' @references Pavoine, S. and Ricotta, C. 2014. Functional and phylogenetic similarity
#' 		among communities. Methods in Ecology and Evolution 5(7): 666–675.
#' @return If \code{partition == TRUE}, list of three matrices giving the turnover,
#' 		nestedness, and total components of dissimilarity. Otherwise, a matrix of
#' 		pairwise dissimilarity values between communities.
#' @export
phylosor <- function(comm, tree, partition = FALSE)
{
	if(is.null(tree$edge.length))
		stop("The tree must have branch lengths")

	if(any(! comm %in% c(0,1) )) {
		warning("Community matrix can only contain 0 and 1; assigning 1 to all nonzero
			entries")
		comm <- 1 * (comm > 0)
	}

	# node labels and community columns must match
	commDrop <- which(!colnames(comm) %in% tree$tip.label)
	if(length(commDrop) > 0)
	{
		if(length(commDrop) == ncol(comm))
			stop("No columns in comm match tip labels in tree")

		warning("Dropping ", length(commDrop), " columns from comm that were
			not found in tree")

		comm <- comm[,-commDrop]
	}
	treeDrop <- which(!tree$tip.label %in% colnames(comm))
	if(length(treeDrop) > 0)
	{
		if(length(treeDrop) == length(tree$tip.labe))
			stop("No tip labels in tree were found in column labels in comm")

		warning("Dropping ", length(treeDrop), " modes from tree that were
			not found in comm")

		tree <- ape::drop.tip(tree, treeDrop)
	}

	if(any(rowSums(comm) == 0))
	{
		drop <- which(rowSums(comm) == 0)
		warning(length(drop), " sites had 0 species and were dropped")
		comm <- comm[-drop,]
	}

	branch_sp <- branch_adjacency(tree, TRUE)
	# save weights
	wts <- attr(branch_sp, "branch.lengths")
	# match the margins
	branch_sp <- branch_sp[,match(colnames(comm), colnames(branch_sp))]

	site_branch <- comm %*% t(branch_sp)
	# convert site by branch to presence absence (instead of sp count)
	site_branch <- 1 * (site_branch > 0)
	# now weight by branch lengths
	site_branch <- sweep(site_branch, 2, wts, `*`)
	sorensen(site_branch, partition)
}

#' Sorensen dissimilarity index based on a dendrogram
#'
#' Produces dissimilarity matrices using a species by species distance matrix; useful
#' for example for functional diversity.
#'
#' @param comm Community matrix giving sites (rows) by species (columns).
#' @param distmat Distance matrix; margins must match exactly the columns in comm
#' @param partition Logical; if true, the index will be partitioned into turnover and
#' 			nestedness.
#' @param ... Additional arguments to pass to \code{\link{hclust}}.
#' @references Pavoine, S. and Ricotta, C. 2014. Functional and phylogenetic similarity
#' 		among communities. Methods in Ecology and Evolution 5(7): 666–675.
#' @return If \code{partition == TRUE}, list of three matrices giving the turnover,
#' 		nestedness, and total components of dissimilarity. Otherwise, a matrix of
#' 		pairwise dissimilarity values between communities.
#' @export
dendrosor <- function(comm, distmat, partition = TRUE, ...)
{
	if(!inherits(distmat, "dist"))
		distmat <- as.dist(distmat)

	if(!attr(distmat, "Size") == ncol(comm) &
		any(attr(distmat, "Labels") != colnames(comm)))
		stop("ncol(comm) must equal the number of rows/columns in distmat and the
			margin labels must be identical")

	tree <- hclust(distmat, ...)
	phylosor(comm, ape::as.phylo(tree), partition)
}

#' Branch adjacency
#'
#' Produce a branch adjacency matrix (i.e., branch by node)
#'
#' @param x Object of class phylo
#' @param branch.lengths Logical; if true, the values of the adjacency matrix will be the
#' 		branch lengths, otherwise they will be 1 or 0
#' @return A branch by node adjacency matrix
#' @keywords internal
branch_adjacency <- function(x, branch.lengths = !is.null(x$edge.length)) {
    m <- matrix(0, ncol=length(x$tip), nrow=nrow(x$edge))
    from <- x$edge[,1]
    to <- x$edge[,2]
    g <- seq_along(x$tip.label)
    while (any(!is.na(g))) {
        i <- match(g, to)
        m[cbind(i, seq_along(i))] <- 1
        g <- from[i]
    }
    rownames(m) <- paste0("Branch", seq.int(nrow(m)))
    colnames(m) <- x$tip.label
    if(branch.lengths)
    {
    	attr(m, "branch.lengths") <- x$edge.length
    }
    return(m)
}

#' Sorensen dissimilarity (without partitioning)
#'
#' Compute sorensen dissimilarity index; deprecated method (but faster than new one)
#'
#' @param comm Community matrix giving sites (rows) by species (columns).
#' @details The index is defined as 1 - (2 * (A ∩ B) / (A + B)).
#' @return Matrix of pairwise dissimilarity values between communities.
#' @keywords internal
sor <- function(comm)
{
	# convert to matrix of 0 and 1
	if(any(! comm %in% c(0,1) )) comm <- 1 * (comm > 0)
	# get the intersection of all sites - number of sp in common
	# this is the 0.5 * numerator of sorensen similarity
	site_similarity <- comm %*% t(comm)

	# compute species richness and repeat into a matrix
	richness <- rowSums(comm)
	ra <- matrix(richness, nrow=length(richness), ncol=length(richness))

	# denominator - ra is the number of spp at site1, t(ra) is the number at site 2
	denom <- ra + t(ra)
	dimnames(denom) <- dimnames(site_similarity)

	1 - ((2 * site_similarity) / denom)
}


runLimTREE <- function(line, paramFile, dat = NULL, ...) {
    if (is.null(dat))
	dat = loadInputData()

    params = read.csv(paramFile, stringsAsFactors=FALSE)[line,]
    params = unlist(params)

    out = LimTREE(dat[['MAP']], dat[['RainTerm_Drought']],
                  dat[['SW1']], dat[['SW2']],
                  dat[['BurntArea']], dat[['StressTerm_Drought']], dat[['MTWM']],
                  dat[['MTCM']],
                  dat[['PopDen']],
                  dat[['buffalo']], dat[['cattle']], dat[['goat']], dat[['sheep']],
                  dat[['urban']], dat[['crop']], dat[['pas']],
                  params['m_drought'],
                  params['min_mat'], params['max_mat'],
                  params['trans_d'], params['p_fire'],
                  params['k_pop'],
                  params['k_buffalo'], params['k_cattle'],
                  params['k_goat'], params['k_sheep'],
                  params['p_drought'],
                  params['min_maxTemp'], params['max_minTemp'],
                  params['q10_maxTemp'], params['q10_minTemp'],
                  params['v_drought'], params['v_maxTemp'], params['v_minTemp'],
                  params['v_pop'],
                  params['v_buffalo'], params['v_cattle'],
                  params['v_goat'], params['v_sheep'],
                  params['v_crop'] ,params['v_pas'],
                  params['MAP_x0'], params['MAP_k'], params['MAT_x0'], params['MAT_k'],
                  params['SW_x0'], params['SW_k'], params['mort_x0'], params['mort_k'],
                  params['ex_x0'], params['ex_k'],
                  params['max_T'], ...)

    return(out)
}

LimTREE <- function(MAP, rain_drought, SW1, SW2, fire, stress_drought,
                    maxTemp, minTemp,
                    popDen,
                    buffalo, cattle, goat, sheep,
                    urban, crop, pas,
		    m_drought, min_mat, max_mat,
                    d, p_fire,
                    k_popDen, k_buffalo, k_cattle, k_goat, k_sheep,
                    p_drought, min_maxTemp, max_minTemp,
                    q10_maxTemp, q10_minTemp,
		    v_drought, v_maxTemp, v_minTemp,
                    v_popDen, v_buffalo, v_cattle, v_goat, v_sheep,
                    v_crop, v_pas,
		    MAP0, MAPk, MAT0, MATk, SW0, SWk, Mort0, Mortk, Exc0, Exck, maxT,
		    includeSW = TRUE, ...) {

    f_MAP  = LimMAP  (MAP, rain_drought, m_drought, MAP0 , MAPk, ...)
    #f_MAT  = LimMAT  (MAT, min_mat, max_mat, MAT0 , MATk, ...)

    if (includeSW) f_SW   = LimSW(SW1, SW2  , d, SW0  , SWk, ...)
    else {
	f_SW = f_MAP
	f_SW[!is.na(f_SW)] = 1.0
    }

    f_MAT = f_MAP
    f_MAT[!is.na(f_MAT)] = 1.0

    f_Mort = LimMort (fire, stress_drought, maxTemp, minTemp, popDen,
                      buffalo, cattle, goat, sheep,
                      v_drought, v_maxTemp, v_minTemp,
                      v_popDen, v_buffalo, v_cattle, v_goat, v_sheep,
                      p_fire,
                      k_popDen, k_buffalo, k_cattle, k_goat, k_sheep,
                      p_drought,
                      min_maxTemp, max_minTemp, q10_maxTemp, q10_minTemp,
                      Mort0, -Mortk, ...)

    f_Exc  = LimExc  (urban, crop, pas, v_crop, v_pas, Exc0, -Exck, ...)

    f_SW  = raster::resample(f_SW , f_MAP)
    f_Exc = raster::resample(f_Exc, f_MAP)
    Tree = f_MAP * f_SW * f_Mort * f_Exc * maxT #  * f_MAT
    return(addLayer(Tree, f_MAP, f_MAT, f_SW, f_Mort, f_Exc))
}

logistic <- function(x, x0, k, sensitivity = FALSE) {
    FUN <- function(xi = x, ki = k)  1 / (1 + exp(-ki * (xi - x0)))
    dFUN <- function(xi) FUN(xi, k) * FUN(xi, -k)
    if (sensitivity)
	out = dFUN(x)/ dFUN(x0)
    else
	out = FUN()
    return(out)
}

LimMAP <- function(MAP, drought, m_drought,...){
    #MAP = MAP + (1-drought) * (1/m_drought) * (exp(-m_drought * MAP)-1)
    MAP = log(MAP)
    logistic(MAP, ...)
}

LimMAT <- function(MAT, min_mat, max_mat, ..., retutnVar = FALSE) {
    MAT = MAT - min_mat
    max_mat = max_mat - min_mat

    MAT[MAT <0] = 0

    MAT = logmin(MAT)
    if (retutnVar) return(MAT)
    logistic(MAT, ...)
}

LimList <- function(x, v, ..., retutnVar = FALSE) {
    v = c(1, v)
    x = mapply('*', x, v)
    xi = x[[1]]
    for (i in x[-1]) xi = xi + i

    xi = xi / sum(v)
    if (retutnVar) out = xi else out = try(logistic(xi, ...), silent = TRUE)
    if (class(out) =="try-error") out = xi
    return(out)
}

LimSW   <- function(SW1, SW2, d,...) {
    SW = (SW1 + d * SW2)/(1+d)
    #SW = logmin(SW)

    logistic(SW, ...)
}

LimTREE.convertUnity <- function(x, k)
    1 - exp(x * (-1/k))

LimTREE.Temp <- function(Temp, mn, q10) {
    Temp = (q10 + 1.0)^(Temp - mn)
    return(Temp)
}

LimMort <- function(fire, drought, maxTemp, minTemp, popDen,
                    buffalo, cattle, goat, sheep,
                    v_drought, v_maxTemp, v_minTemp,
                    v_popDen, v_buffalo, v_cattle, v_goat, v_sheep,
		    p_fire,
                    k_popDen, k_buffalo, k_cattle, k_goat, k_sheep,
                    p_drought, min_maxTemp, max_minTemp,
                    q10_maxTemp, q10_minTemp, ...) {

    fire = fire^p_fire
    popDen0 = popDen
    popDen   = LimTREE.convertUnity( popDen, k_popDen )
    buffalo  = LimTREE.convertUnity(buffalo, k_buffalo)
    cattle   = LimTREE.convertUnity( cattle, k_cattle )
    goat     = LimTREE.convertUnity(   goat, k_goat   )
    sheep    = LimTREE.convertUnity(  sheep, k_sheep  )

    maxTemp = LimTREE.Temp(maxTemp     ,  min_maxTemp, q10_maxTemp)
    minTemp = LimTREE.Temp(minTemp*(-1), -max_minTemp, q10_minTemp)
    drought = drought^p_drought

    LimList(c(fire, drought, maxTemp, minTemp,
              popDen, buffalo, cattle, goat, sheep),
            c(v_drought, v_maxTemp, v_minTemp,
              v_popDen, v_buffalo, v_cattle, v_goat, v_sheep), ...)
}

LimExc  <- function(urban, crop, pas, v_crop, v_pas, ...)
    LimList(c(urban, crop, pas), c(v_crop, v_pas), ...)


# Simulation study: drifting AR(1) model

round_timestamp <- function(ts, scalefactor.sec = 3600*12) {
  as.POSIXct(round(as.numeric(ts)/scalefactor.sec) * scalefactor.sec, origin = '1970-01-01', tz = 'UTC')
}

read.baro <- function(logger.name) {
  df <- gwloggeR.data::read(logger.name)$df
  if (nrow(df) == 0L) return(df)
  df <- df[!is.na(TIMESTAMP_UTC), ]
  df <- df[!is.na(PRESSURE_VALUE), ]
  df <- df[!duplicated(TIMESTAMP_UTC), ]
  df <- df[, .('PRESSURE_VALUE' = mean(PRESSURE_VALUE)),
           by = .('TIMESTAMP_UTC' = round_timestamp(TIMESTAMP_UTC))]

  # Meta data
  df[, 'FILE' := basename(logger.name)]
  df[, 'N' := .N]

  data.table::setkey(df, TIMESTAMP_UTC)
  data.table::setattr(df, 'logger.name', logger.name)

  df
}

# from analysis 04
compare <- function(df1, df2) {
  if (nrow(df1) == 0L || nrow(df2) == 0L) return(data.table::data.table())

  diff.df <- df1[J(df2), .('PRESSURE_DIFF' = x.PRESSURE_VALUE - i.PRESSURE_VALUE, TIMESTAMP_UTC)][!is.na(PRESSURE_DIFF), ]

  data.table::setkey(diff.df, TIMESTAMP_UTC)

  diff.df
}

ref.compare <- function(df) {
  compare(df, read.baro('KNMI_20200312_hourly'))
}

# plot(arima.sim(n = 10000, list(ar = c(0.9)), sd = sqrt(23.62)) + 1033.64)
# plot(forecast:::simulate.Arima(fit.arima, 10000))

df.drifter <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2010-01-01 00:00:00'), to = as.POSIXct('2020-01-01 00:00:00'), by = '12 hours'),
  PRESSURE_VALUE = arima.sim(n = 7305, list(ar = c(0.9)), sd = sqrt(23.62)) + 1033.64 + c(rep(0, 3000), 1:4305)*0.005,
  key = 'TIMESTAMP_UTC'
)

arima(df.drifter$PRESSURE_VALUE, order = c(1, 0, 0))
plot(df.drifter$PRESSURE_VALUE)

df.diff <- ref.compare(df.drifter)
plot(df.diff$PRESSURE_DIFF)

# AR(1) component in case of difference between two barometers -----------------
sim <- function(n, mu, sigma, phi1, betas = NULL, xreg = NULL,
                init = mu, xt = NULL, a = rnorm(n, 0, sd = sigma)) {

  R <- rep(0, n)
  if (!is.null(betas) || !is.null(xreg)) {
    reg <- xreg[-1,,drop=FALSE] - phi1*xreg[-n,,drop=FALSE]
    R <- c(0, reg %*% betas)
  }

  if (is.null(xt)) {
    # xp is the previous x: x_{t-1}
    Reduce(f = function(xp, t) phi1*xp + mu*(1-phi1) + R[t] + a[t], x = 2:n, init = init, accumulate = TRUE)
  } else {
    c(xt[1],
      phi1*xt[-n] + mu*(1-phi1) + R[-1] + a[-1])
  }
}

set.seed(2020)
correlated.errors <- data.table::as.data.table(
  MASS::mvrnorm(1e5, mu = c('a' = 0, 'b' = 0), Sigma = matrix(c(23, 20, 20, 23), ncol = 2))
)
cov(correlated.errors$a, correlated.errors$b)

## all equal: OK (error due to correlation is small --> -2*covariance)
df.diff <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2001-01-01 00:00:00'), by = '12 hours', length.out = 1e5),
  BARO_01 = sim(n = 1e5, phi1 = 0.9, mu = 1032, a = correlated.errors$a),
  BARO_02 = sim(n = 1e5, phi1 = 0.9, mu = 1032, a = correlated.errors$b)
)
df.diff[, DIFF := BARO_01 - BARO_02]
forecast::auto.arima(df.diff$DIFF, trace = TRUE, stationary = TRUE)

## uncorrelated: OK (error = the sum of individual errors)
set.seed(2020)
df.diff <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2001-01-01 00:00:00'), by = '12 hours', length.out = 1e5),
  BARO_01 = sim(n = 1e5, phi1 = 0.9, mu = 1020, sigma = sqrt(23)),
  BARO_02 = sim(n = 1e5, phi1 = 0.9, mu = 1035, sigma = sqrt(23))
)
df.diff[, DIFF := BARO_01 - BARO_02]
forecast::auto.arima(df.diff$DIFF, trace = TRUE, stationary = TRUE)

## mu different: OK (new mu = difference between the two)
df.diff <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2001-01-01 00:00:00'), by = '12 hours', length.out = 1e5),
  BARO_01 = sim(n = 1e5, phi1 = 0.9, mu = 1020, a = correlated.errors$a),
  BARO_02 = sim(n = 1e5, phi1 = 0.9, mu = 1035, a = correlated.errors$b)
)
df.diff[, DIFF := BARO_01 - BARO_02]
forecast::auto.arima(df.diff$DIFF, trace = TRUE, stationary = TRUE)

## sigma different: OK (new sigma: addition of the two minus 2*cov)
set.seed(2020)
correlated.errors.different.sigma <- data.table::as.data.table(
  MASS::mvrnorm(1e5, mu = c('a' = 0, 'b' = 0), Sigma = matrix(c(10, 15, 15, 23), ncol = 2))
)
cov(correlated.errors.different.sigma$a, correlated.errors.different.sigma$b)
df.diff <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2001-01-01 00:00:00'), by = '12 hours', length.out = 1e5),
  BARO_01 = sim(n = 1e5, phi1 = 0.9, mu = 1035, a = correlated.errors.different.sigma$a),
  BARO_02 = sim(n = 1e5, phi1 = 0.9, mu = 1035, a = correlated.errors.different.sigma$b)
)
df.diff[, DIFF := BARO_01 - BARO_02]
forecast::auto.arima(df.diff$DIFF, trace = TRUE, stationary = TRUE)

## phi different: MA component seems necessary, new sigma larger
df.diff <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2001-01-01 00:00:00'), by = '12 hours', length.out = 1e5),
  BARO_01 = sim(n = 1e5, phi1 = 0.5, mu = 1035, a = correlated.errors$a),
  BARO_02 = sim(n = 1e5, phi1 = 0.9, mu = 1035, a = correlated.errors$b)
)
df.diff[, DIFF := BARO_01 - BARO_02]
forecast::auto.arima(df.diff$DIFF, trace = TRUE, stationary = TRUE)

## phi different & uncorrelated: MA component seems necessary, new sigma larger
set.seed(2020)
df.diff <- data.table::data.table(
  TIMESTAMP_UTC = seq(from = as.POSIXct('2001-01-01 00:00:00'), by = '12 hours', length.out = 1e5),
  BARO_01 = sim(n = 1e5, phi1 = 0.5, mu = 1035, sigma = sqrt(23)),
  BARO_02 = sim(n = 1e5, phi1 = 0.9, mu = 1035, sigma = sqrt(23))
)
df.diff[, DIFF := BARO_01 - BARO_02]
forecast::auto.arima(df.diff$DIFF, trace = TRUE, stationary = TRUE)


# test wfs simulations

require(foreach)
require(doParallel)
source('scripts/wfs_simp.r')

# make cluster
cl <- makeCluster(3)
registerDoParallel(cl)

# probability of fixation of a neutral allele should equal the initial frequency
	ne <- 100
	gen=1000 # run for a long time: likely to fixaton (expected time is 4N)
	smin <- smax <- 0
	fs <- seq(1/100, 91/100, by=5/100) # initial allele frequencies
	nsims <- 1000

	sims <- data.frame(f1=rep(fs, rep(nsims, length(fs))))
	sims$f2 <- sims$f1samp <- sims$f2samp <- NA
	for(j in 1:length(fs)){
		print(j)
		thesesims <- foreach(i=1:nsims, .combine=rbind) %dopar% {
			wfs_simp(i, f1=fs[j], s=0, c1=50, c2=50, gen=gen, ne=ne)
		}
		sims[((j-1)*1000+1):(j*1000),c('f2', 'f1samp', 'f2samp')] <- as.data.frame(thesesims)
	}

	propfix <- aggregate(list(prop=sims$f2), by=list(f1=sims$f1), FUN=function(x) sum(x==1)/nsims)

	# plot of proportion fixed vs. initial frequency
	quartz(width=5, height=4)
	# png(width=5, height=4, units='in', res=300, file='analysis/figures/wfs_sim_test_s=0.png')
	par(las=1)
	with(propfix, plot(f1, prop, xlab='Initial allele frequency', ylab='Proportion fixed', main='Ne=100, s=0, nsims=1000'))
	abline(0,1)

	dev.off()

# probability of fixation of a non-neutral allele should equal s if initial frequency is 1/ne
# (usually written as 2s, but that's when fitness of waa is 1+s2. here it is 1+s)
# for large N and small s
	ne <- 100
	gen=1000 # run for a long time: likely to fixation in 4N generations with drift
	ss <- seq(0.05,1,by=0.05)
	f1 <- 1/ne # starting allele freq
	nsims <- 1000

	sims2 <- data.frame(s=rep(ss, rep(nsims, length(ss))))
	sims2$f2 <- sims2$f1samp <- sims2$f2samp <- NA
	length(ss)
	for(j in 1:length(ss)){
		print(j)
		thesesims <- foreach(i=1:nsims, .combine=rbind) %dopar% {
			if(i %% 100 == 0) cat(i)
			wfs_simp(i, f1=f1, s=ss[j], c1=50, c2=50, gen=gen, ne=ne)
		}
		sims2[((j-1)*1000+1):(j*1000),c('f2', 'f1samp', 'f2samp')] <- as.data.frame(thesesims)
	}

	sum(sims2$f2 %in% c(0,1))/nsims/length(ss) # fraction that reached fixation: make gen larger until this=1
	propfix2 <- aggregate(list(prop=sims2$f2), by=list(s=sims2$s), FUN=function(x) sum(x==1)/nsims)
	propfix2$s_exp1 <- (1-exp(-propfix2$s))/(1-exp(-2*ne*propfix2$s)) # Eq. 2.21 in Graur & Li
	propfix2$s_exp2 <- propfix2$s/(1-exp(-2*ne*propfix2$s)) # Eq. 2.22 in Graur & Li

	# plot of proportion fixed vs. s
	quartz(width=5, height=4)
	# png(width=5, height=4, units='in', res=300, file='analysis/figures/wfs_sim_test_s>0.png')
	par(las=1)
	with(propfix2, plot(s, prop, xlab='Selection coefficient (s)', ylab='Proportion fixed at p=1', main='Ne=100, p0=1/Ne, nsims=1000', ylim=c(0,0.7)))
	with(propfix2, lines(s, s_exp1, col='green'))
	legend('topleft', col='green', lty=1, legend='(1-exp(-s))/(1-exp(-2*ne*s))', bty='n', cex=0.7)

	dev.off()

## examine how much drift vs. sampling variance for large ne
	ne <- 10000 # like NEA cod
	gen=11 # like NEA cod
	smin <- smax <- 0
	f1min <- f1max <- ne/2/ne # starting allele freq
	nsims <- 1000

	sims <- foreach(i=1:nsims, .combine=rbind) %dopar% {
		if(i %% 100 == 0) cat(i)
		wfs(i, f1min=f1min, f1max=f1max, smin=smin, smax=smax, c1=50, c2=50, gen=gen, ne=ne)
	}
	sims <- as.data.frame(sims)

	# examine drift variance and sampling variance
	par(mfrow=c(3,1))
	bks <- seq(0,1,length.out=20)
	hist(sims$f1samp, main='f1samp', breaks=bks, col='grey'); abline(v=f1min, col='red', lwd=2)
	hist(sims$f2, main='f2', breaks=bks, col='grey'); abline(v=f1min, col='red', lwd=2)
	hist(sims$f2samp-sims$f2, main='delta f2samp', breaks=seq(-0.5,0.5,length.out=20), col='grey'); abline(v=0, col='red', lwd=2)

## examine how much drift to expect for large Ne
	# observations
	dat <- locnms <- fread('analysis/Frequency_table_Lof07_Lof14.txt', header=TRUE); setnames(locnms, 3:7, c('N_CHR_1', 'Freq_1', 'N_CHR_2', 'Freq_2', 'ABS_DIFF')) # the name and observed frequencies of all the loci, from output by Bastiaan Star

	# simulations
	ne <- 10000
	gen=11 # run for a long time: likely to fixaton (expected time is 4N)
	fs <- seq(1/100, 91/100, by=5/100) # initial allele frequencies
	nsims <- 100000 # takes 20 min or so

	sims <- data.frame(f1=rep(fs, rep(nsims, length(fs))))
	sims$f2 <- sims$f1samp <- sims$f2samp <- NA
	for(j in 1:length(fs)){
		print(j)
		thesesims <- foreach(i=1:nsims, .combine=rbind) %dopar% {
			wfs_simp(i, f1=fs[j], s=0, c1=44, c2=44, gen=gen, ne=ne)
		}
		sims[((j-1)*nsims+1):(j*nsims),c('f2', 'f1samp', 'f2samp')] <- as.data.frame(thesesims)
	}

	xlims <- range(c(dat[,Freq_2-Freq_1], sims$f2samp-sims$f1samp))
	cols <- c('black', '#7fcdbb', '#edf8b1', '#2c7fb8') # sims, obs 44, low obs, high obs

	# plot histogram of diff vs. initial observed frequency
	quartz(width=7, height=6)
	# png(width=7, height=7, filename='analysis/figures/wfs_sim_test_hist_obsdiff_by_obsf1_with_data.png', units='in', res=300)
	f1s <- seq(0,1.1,by=0.1)
	bks <- seq(-1,1,by=0.05)
	par(mfrow=c(3,4), las=1, tcl=-0.3, mgp=c(2.4,0.6,0), mai=c(0.5, 0.5, 0.3, 0.05))
	for(i in 2:length(f1s)){
		inds <- (sims$f1samp >= f1s[i-1]) & (sims$f1samp < f1s[i])
		hst <- hist(sims$f2samp[inds] - sims$f1samp[inds], plot=FALSE, breaks=bks) # returns some 0 values that get eliminated, but appear as holes in the histogram: doesn't change our interpretation
		j <- which(hst$density>0)
		plot(hst$mids[j], hst$density[j], type='o', xlab='Sample ∆freq', ylab='log(density)', main=paste('initial ', f1s[i-1], '-', f1s[i], sep=''), pch=16, xlim=xlims, log='y', ylim=c(0.0001,20), col=cols[1])
		abline(v=c(-0.06, 0.06), col='grey', lty=3) # genome-wide allele frequency change

			# low sample size
		inds3 <- dat[,Freq_1 >= f1s[i-1] & Freq_1 < f1s[i] & (N_CHR_1<44 | N_CHR_2<44)]
		hst3 <- hist(dat[inds3,Freq_2 - Freq_1], plot=FALSE, breaks=bks) # returns some 0 values that get eliminated, but appear as holes in the histogram: doesn't change our interpretation
		j3 <- which(hst3$density>0)
		lines(hst3$mids[j3], hst3$density[j3], type='o', pch=16, col=cols[3])

			# high sample size
		inds4 <- dat[,Freq_1 >= f1s[i-1] & Freq_1 < f1s[i] & (N_CHR_1>44 & N_CHR_2>44)]
		hst4 <- hist(dat[inds4,Freq_2 - Freq_1], plot=FALSE, breaks=bks) # returns some 0 values that get eliminated, but appear as holes in the histogram: doesn't change our interpretation
		j4 <- which(hst4$density>0)
		lines(hst4$mids[j4], hst4$density[j4], type='o', pch=16, col=cols[4])

			# sample size 44
		inds2 <- dat[,Freq_1 >= f1s[i-1] & Freq_1 < f1s[i] & N_CHR_1==44 & N_CHR_2==44]
		hst2 <- hist(dat[inds2,Freq_2 - Freq_1], plot=FALSE, breaks=bks) # returns some 0 values that get eliminated, but appear as holes in the histogram: doesn't change our interpretation
		j2 <- which(hst2$density>0)
		lines(hst2$mids[j2], hst2$density[j2], type='o', pch=16, col=cols[2])

	}

	legend('bottomright', col=cols, lty=1, legend=c('Sims n=44', 'Obs n=44', 'Obs low n', 'Obs high n'), cex=0.7, bty='n')

	dev.off()

#######################
# clean up cluster
#######################
stopCluster(cl)


    # copied from library carstm
    stmv_hyperparameters = function( reference_sd, alpha=0.5, reference_mean=0 ) {
      # some generic PC priors, scaled by sd of data
      # pc.prior to median .. minimally info. scale

      hyper = list(

        iid = list(
          prec = list(
            prior = "pc.prec",  # exponential decay
            param = c(reference_sd, alpha)
          )
        ),

        # means informative, sd marginally diffuse
        # see: inla.set.control.fixed.default() for defaults
        fixed = list(
            mean.intercept = reference_mean,
            prec.intercept = 1e-3,
            mean=0,
            prec=1e-2
        ),

        # param=c(u, alpha); u=sigma; alpha=prob;
        # see inla.doc("pc.prec")  ..prior sd attributable to rw2
        rw2 = list(
          prec = list(
            prior = "pc.prec",  # exponential decay
            param = c(reference_sd, alpha)
          )
        ),

        # see inla.doc("ar1") ; theta0, theta1 are expected
        # param=c(u, alpha); u=sigma; alpha=prob;
        # see inla.doc("pc.prec")  ..prior sd attributable to autocor rho
        # param=c(u, alpha); rho = 0.5; u=sqrt(1-rho); alpha=prob; see inla.doc("pc.cor1")
        ar1 = list(
          prec = list(
            prior = "pc.prec",  # exponential decay
            param = c(reference_sd, alpha)
          ),
          rho = list(
            prior = "pc.cor0", # inla.doc("pc.cor0") ..base model: rho = 0  --- expoential; will tend to 0 unless there is info
            param = c(sqrt(1-0.5), 0.1)  # rho=0.5; u=sqrt(1-rho)  ... 100-10% of probablity weight to rho 0.5 or less .. forces smooth and only goes high if really high
          )
        ),

        # naming convention is a bit different in groups .. Must be one of  theta theta1 rho logit correlation
        ar1_group = list(
          # theta1 = list(
          #   prior = "pc.prec",  # exponential decay
          #   param = c(reference_sd, alpha)
          # ),
          rho = list(
            prior = "pc.cor0", # inla.doc("pc.cor0") ..base model: rho = 0  --- expoential; will tend to 0 unless there is info
            param = c(sqrt(1-0.5), 0.1)  # rho=0.5; u=sqrt(1-rho)  ... 100-10% of probablity weight to rho 0.5 or less .. forces smooth and only goes high if really high
          )
        ),

        # param=c(u, alpha); u=phi (proportion spatial); alpha=prob
        bym2 = list(
          prec = list(
            prior = "pc.prec",
            param = c(reference_sd, alpha)
          ),
          phi = list(
            prior="pc",  # see bottom of inla.doc("bym2")
            param=c(0.5, 0.5) # c(phi=0.5, alpha=0.5)
          )
        )
      )

      return(hyper)
    }


library(animation)
x <- c(1,2,3,5,6,9)
ani.options(interval = 1,ani.height = 1000,ani.width = 1000)
saveGIF({
  for(i in 1:length(x)){
    plot(x,rep(0,length(x)), pch = 20,cex = 3, axes = FALSE, xlab = "", ylab = "")
    abline(h = 0)
    text(x = x, y = rep(0.1,length(x)), c(expression(x[1]),expression(x[2]),expression(x[3]),
                                                expression(x[4]),expression(x[5]),expression(x[6])),
         cex = 2)
    text(x = x, y = rep(-0.1,length(x)),x,cex = 2)
    arrows(mean(x),-0.75,mean(x),0,length = 0.2,lwd = 2)
    text(mean(x),-0.85,expression(paste(mu[x], "(central value)")),cex = 1.5)
    text(mean(x),-0.95,paste("=",bquote(.(round(mean(x),2)))),cex = 1.5)
    arrows(7,0.75,x[4:6],0.2,length = 0.2,lwd = 2,lty = c(2,1,1))
    text(7,0.9,"Possible values of X",cex = 1.5)
    ##legend("topleft",bty = "n", legend = expression(paste("Var(X) = E[(X -",mu[X],")",""^2,"]")),cex = 1.3)
    arrows(mean(x),-0.25,x[i],-0.25,length = 0.1,code = 3,lwd = 2)
    text(mean(c(x[i],mean(x))),-0.3,bquote(x[.(i)] - mu[X]),cex = 2)
    text(mean(c(x[i],mean(x))),-0.4, paste("=",bquote(.(round(x[i] - mean(x),2)))),cex = 1.5)
    text(1.5,0.6,bquote((x[.(i)] - mu[X])^2),cex = 2)
    text(2.5,0.6, paste("=",bquote(.(round((x[i] - mean(x))^2,2)))),cex = 2)
  }}, movie.name = "var.gif"
)


### RD design functions
###peter crosta pmcrosta at gmail dot com
## This file consists of only functions that are used in the
## RD analysis. It will be sourced

###########MOST FUNCTIONS ARE DEPENDENT ON maxbin AND binwidth (h) ###########################

## Create class binlist for numeric breaks
setClass("binlist", representation(breaks="numeric", right="numeric", left="numeric"))

getbins <- function(binwidth, maxbin) {
  rightbins <- seq(binwidth, maxbin, binwidth)
  leftbins <- -rev(seq(0, maxbin, binwidth))

  breakbins <- c(sort(leftbins), rightbins)

  z <- new("binlist")
  z@breaks <- breakbins[abs(breakbins)<=maxbin]
  z@right <- rightbins[rightbins<=maxbin]
  z@left <- leftbins[leftbins>=-maxbin]

  return(z)
}

# -244 to 121 might matter, days before and after Sep 1
mainhist <- function(binwidth, maxbin, yesplot=TRUE, emcp=emc) {
  breaks <- getbins(binwidth, maxbin)@breaks
  emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)
  maxbin.hist <- hist(emcp[emcdum==1], breaks=breaks, plot=yesplot, main=paste("Histogram of birth date (centered): Binwidth=",binwidth, ", ", length(breaks)-1, " bins", sep=""), xlim=c(min(breaks), max(breaks)), xlab="Centered Birth Date")
  if (yesplot) abline(v=0, lty=2, lwd=2, col="darkgrey")
  binmids <- maxbin.hist$mids
}

####this function is built from McCrary(2008); tests discontinuity of density
mccrary <- function(binwidth, maxbin, nn, emcp=emc) {
    f.triangle <- function(x) { ifelse(abs(x)<1, 1-abs(x), 0) }

    breaks <- getbins(binwidth, maxbin)@breaks
    emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)
    maxbin.hist <- hist(emcp[emcdum==1], breaks=breaks, plot=FALSE)
    binmids <- maxbin.hist$mids

    totaln <- sum(maxbin.hist$counts)
    Xj <- binmids
    Yj <- maxbin.hist$counts / (totaln*binwidth)

    ##set bandwidth (nearest neighbors); #nn <-20; this is argument in func call

    ##first step histogram
    plot(Xj, Yj, xlab="Bin midpoints", ylab="Normalized cell size", main=paste("First step histogram for McCrary(2008)\nbinwidth=",binwidth, " bandwidth=", nn, sep=""))

    ##local fit right and plot
    right.locfit <- locfit.raw(lp(Xj[Xj > 0], h=nn, deg=1), Yj[Xj > 0], kern='tria', family="gaussian")
    plot(right.locfit, add=TRUE, col="red")

    ###weighted reg around boundary point with triang kernel to get intercept term
    wt <- f.triangle((Xj[Xj > 0]-rep(0, length( Xj[Xj > 0])))/nn)
    wt <- (wt/(sum(wt)))*length(wt)
    templm <- lm(Yj[Xj > 0]~Xj[Xj > 0], weight=wt)
    fr <- coef(templm)[1]

    #local fit left and plot
    left.locfit <- locfit.raw(lp(Xj[Xj <= 0], h=nn, deg=1), Yj[Xj <= 0], kern='tria', family="gaussian")
    plot(left.locfit, add=TRUE, col="red")

    ###weighted reg around boundary point with trian kernel to get intercept term
    wt <- f.triangle((Xj[Xj <= 0]-rep(0,length(Xj[Xj <= 0])))/nn)
    wt <- (wt/(sum(wt)))*length(wt) #like aweights in stata
    templm <- lm(Yj[Xj <= 0]~Xj[Xj <= 0], weight=wt)
    fl <- coef(templm)[1]

    abline(v=0, lty=2, lwd=2, col="darkgrey")

    thetahat <- log(fr) - log(fl)
    sethetahat <- sqrt( (1/(totaln*nn)) * (24/5) * ((1/fr) + (1/fl)) )

    legend("topright", bg="white", cex=.9, pt.cex=.9, c("Cutoff", "Bin-mid regression with triangular kernel"), lty=c(2, 1), lwd=c(2, 1), col=c("darkgrey", "red"))
    text(-maxbin+30, max(Yj)*.92, paste("theta-hat =", round(thetahat, 4), "\nstd. err. = ", round(sethetahat,4)))
    return(cat(binwidth, "\t", maxbin, "\t", nn, "\t", thetahat, "\t", sethetahat, "\n"))
}

#####function to generate prop tests to look at subgroup covariate differences around cutoff (like gifted and atrisk)
prettyprop <- function(binwidth, maxbin, covariate, emcp=emc, lab="") {
    breaks <- getbins(binwidth, maxbin)@breaks
    emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

    covariate <- covariate[emcdum==1]
    cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)

    afters<-which(cutbybreak %in% tail(levels(cutbybreak), maxbin))
    befores<-which(cutbybreak %in% head(levels(cutbybreak), maxbin))

    x1 <- sum(covariate[befores], na.rm=TRUE)
    x2 <- sum(covariate[afters], na.rm=TRUE)
    n1 <- length(na.omit(covariate[befores]))
    n2 <- length(na.omit(covariate[afters]))

    pt <- prop.test(c(x1,x2), c(n1, n2))
    return(paste(lab, round(pt$estimate[1], 3), round(pt$estimate[2], 3), round(diff(pt$estimate), 3), round(pt$p.value, 3), n1, n2, sep=', '))
}

######## similar to prettyprop, but for variables that take a t-test
ttest <- function(binwidth, maxbin, covariate, emcp=emc, lab="") {
    breaks <- getbins(binwidth, maxbin)@breaks
    emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

    covariate <- covariate[emcdum==1]
    cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)

    afters<-which(cutbybreak %in% tail(levels(cutbybreak), maxbin))
    befores<-which(cutbybreak %in% head(levels(cutbybreak), maxbin))

    x1 <- covariate[befores]
    x2 <- covariate[afters]
    tt <- t.test(x1, x2)

    return(paste(lab, round(tt$estimate[1], 3), round(tt$estimate[2], 3), round(-diff(tt$estimate), 3), round(tt$p.value, 3), length(na.omit(x1)), length(na.omit(x2)), sep=","))
}

covfunc <- function(binwidth, maxbin, covariate, ylabcov, use.loess=FALSE, emcp=emc, YLIM=TRUE, nn=10, g) {

    breaks <- getbins(binwidth, maxbin)@breaks
    emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

    cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)
    binmids <- hist(emcp[emcdum==1], breaks=breaks, plot=FALSE)$mids

    binned.covariate <- tapply(covariate[emcdum==1], cutbybreak, function(x) mean(x, na.rm=TRUE))
    forcing <- ifelse(binmids>0, 1, 0)
    tx <- which(forcing==1)
    forcing.sample <- ifelse(emcp>0, 1, 0)

    if (YLIM) {
        plot(binned.covariate~binmids, col=NULL, xlab="Birth date (centered)", ylab=ylabcov, main=paste(ylabcov, " Grade ",g, "\nBinwidth=",binwidth, ", ", length(breaks)-1, " bins", sep=""), xaxt='n', ylim=c(0, 1))
    } else plot(binned.covariate~binmids, col=NULL, xlab="Birth date (centered)", ylab=ylabcov, main=paste(ylabcov, " Grade ",g, "\nBinwidth=",binwidth, ", ", length(breaks)-1, " bins", sep=""), axes=FALSE)

    points(binmids[tx], binned.covariate[tx], col="blue4", pch=4)
    points(binmids[-tx], binned.covariate[-tx], col="darkorange", pch=16)
    axis(1, at=seq(-maxbin, maxbin, as.integer(maxbin/10)))

    if (!YLIM) axis(2, at=round(seq(min(binned.covariate, na.rm=TRUE), max(binned.covariate, na.rm=TRUE), diff(range(binned.covariate, na.rm=TRUE))/10),2))

    abline(v=0, lty=2, lwd=2, col="darkgrey")

    if (ylabcov %in% covstlab) {
      pval <- ttest(1, binwidth, covariate, emcp)
    } else {
      pval <- prettyprop(1, binwidth, covariate, emcp)
    }
    pvalsplit <- as.numeric(unlist(strsplit(pval, split=",")))
    cat(pvalsplit, "\n")

    if (use.loess) {
      f.triangle <- function(x) { ifelse(abs(x)<1, 1-abs(x), 0) }
      Yj <- binned.covariate
      Xj <- binmids[which(!is.na(Yj))]
      Yj <- binned.covariate[which(!is.na(Yj))]

      #changed so this is now argument
      #nn<-binwidth*3

      ##local fit right and plot
      right.locfit <- locfit.raw(lp(Xj[Xj > 0], h=nn, deg=1), Yj[Xj > 0], kern="tria", family="gaussian")
      plot(right.locfit, add=TRUE, col="green")
      #local fit left and plot
      left.locfit <- locfit.raw(lp(Xj[Xj <= 0], h=nn, deg=1), Yj[Xj <= 0], kern='tria', family="gaussian")
      plot(left.locfit, add=TRUE, col="green")

      text(-maxbin/2, quantile(binned.covariate, na.rm=TRUE)[4], cex=0.8, paste("1 binwidth difference at cutoff =", round(pvalsplit[4],3), "\np-val=", round(pvalsplit[5],3)))
      #Add legend.
      legend("topright", bg="white", cex=.9, pt.cex=.9, c("Older", "Younger", "Cutoff", "Piecewise loess"), lty=c(NA, NA, 2, 1), lwd=c(NA, NA, 2, 1), col=c("blue4", "darkorange", "darkgrey", "green"), pch=c(4, 16, NA, NA, NA))
    }  else {
        ##run simple regression model and add curves to plot
        model <- lm(covariate~forcing.sample + emcp + I(emcp^2), subset=emcdum==1)
        coefs <- coefficients(model)
        curve(coefs[1] + coefs[3]*x + coefs[4]*x^2, -maxbin, 0, add=T, col="red")
        curve(coefs[1] + coefs[2] + coefs[3]*x + coefs[4]*x, 0, maxbin, add=T, col="red")
        legend("topright", bg="white", cex=.9, pt.cex=.9, c("Older", "Younger", "Cutoff", "Piecewise OLS"), lty=c(NA, NA, 2, 1), lwd=c(NA, NA, 2, 1), col=c("blue4", "darkorange", "darkgrey", "red"), pch=c(4, 16, NA, NA))
        text(-maxbin/2, quantile(binned.covariate, na.rm=TRUE)[4], cex=0.8, paste("OLS difference at cutoff = ", round(coefs[2], 3), "(", round(sqrt(vcov(model)[2,2]), 3), ")",sep=""))
    }

}

outvforce <- function(binwidth, maxbin, covariate, ylabcov, use.loess=FALSE, emcp=emc, YLIM=TRUE, nn=10, g, change=FALSE) {

  ###Plot of outcome vs. forcing variable
  breaks <- getbins(binwidth, maxbin)@breaks
  emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

  cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)
  binmids <- hist(emcp[emcdum==1], breaks=breaks, plot=FALSE)$mids

  ##group means with binwidth of ; midpoints for horiz axis supplied by hist
  binned.covariate <- tapply(covariate[emcdum==1], cutbybreak, function(x) mean(x, na.rm=TRUE))
  forcing <- ifelse(binmids>0, 1, 0)
  tx <- which(forcing==1)
  forcing.sample <- ifelse(emcp>0, 1, 0)

  if (change) {
    maintit=paste(ylabcov, " Grade ",g-1, " to ", g, "\nBinwidth=",binwidth, ", ", length(breaks)-1, " bins", sep="")
  } else maintit=paste(ylabcov, " Grade ",g, "\nBinwidth=",binwidth, ", ", length(breaks)-1, " bins", sep="")

  if (YLIM) {
        plot(binned.covariate~binmids, col=NULL, xlab="Birth date (centered)", ylab=ylabcov, main=maintit, xaxt='n', ylim=c(0, 1))
  } else plot(binned.covariate~binmids, col=NULL, xlab="Birth date (centered)", ylab=ylabcov, main=maintit, axes=FALSE)

  points(binmids[tx], binned.covariate[tx], col="blue4", pch=4)
  points(binmids[-tx], binned.covariate[-tx], col="darkorange", pch=16)
  axis(1, at=seq(-maxbin, maxbin, as.integer(maxbin/10)))

  if (!YLIM) axis(2, at=round(seq(min(binned.covariate, na.rm=TRUE), max(binned.covariate, na.rm=TRUE), diff(range(binned.covariate, na.rm=TRUE))/10),2))

  abline(v=0, lty=2, lwd=2, col="darkgrey")

 if (ylabcov %in% covstlab) {
      pval <- ttest(1, binwidth, covariate, emcp)
    } else {
      pval <- prettyprop(1, binwidth, covariate, emcp)
    }
    pvalsplit <- as.numeric(unlist(strsplit(pval, split=",")))
    cat(pvalsplit, "\n")

    ### LAST MOD IS HERE
    if (use.loess) {
    f.triangle <- function(x) { ifelse(abs(x)<1, 1-abs(x), 0) }
    Yj <- binned.covariate
    Xj <- binmids[which(!is.na(Yj))]
    Yj <- binned.covariate[which(!is.na(Yj))]

    #changed so this is now argument
    #nn<-binwidth*3

    ##local fit right and plot
    right.locfit <- locfit.raw(lp(Xj[Xj > 0], h=nn, deg=1), Yj[Xj > 0], kern='tria', family="gaussian")
    plot(right.locfit, add=TRUE, col="green")

    #local fit left and plot
    left.locfit <- locfit.raw(lp(Xj[Xj <= 0], h=nn, deg=1), Yj[Xj <= 0], kern='tria', family="gaussian")
    plot(left.locfit, add=TRUE, col="green")

    text(-maxbin/2, quantile(binned.covariate, na.rm=TRUE)[4], paste("1 binwidth difference at cutoff =", round(pvalsplit[4],3), "\np-val=", round(pvalsplit[5],3)))
    #Add legend.
      legend("topright", bg="white", cex=.9, pt.cex=.9, c("Older", "Younger", "Cutoff", "Piecewise loess"), lty=c(NA, NA, 2, 1), lwd=c(NA, NA, 2, 1), col=c("blue4", "darkorange", "darkgrey", "green"), pch=c(4, 16, NA, NA, NA))
    }  else {
    ##run simple regression model and add curves to plot
    model <- lm(binned.covariate~forcing + binmids + I(binmids^2), subset=emcdum==1)
    #model <- lm(covariate~forcing.sample + emcp + I(emcp^2), subset=emcdum==1)
    coefs <- coefficients(model)
    curve(coefs[1] + coefs[3]*x + coefs[4]*x^2, -maxbin, 0, add=T, col="red")
    curve(coefs[1] + coefs[2] + coefs[3]*x + coefs[4]*x, 0, maxbin, add=T, col="red")
    legend("topright", bg="white", cex=.9, pt.cex=.9, c("Older", "Younger", "Cutoff", "Piecewise OLS"), lty=c(NA, NA, 2, 1), lwd=c(NA, NA, 2, 1), col=c("blue4", "darkorange", "darkgrey", "red"), pch=c(4, 16, NA, NA))
        text(-maxbin/2, quantile(binned.covariate, na.rm=TRUE)[4], cex=0.8, paste("OLS difference\nat cutoff = ", round(coefs[2], 3), "(", round(sqrt(vcov(model)[2,2]), 3), ")",sep=""))
    }
}

##********************************************************************COVARIATES******************
#x11(height=10, width=10)

clx <- function(fm, dfcw, cluster){
         # R-codes (www.r-project.org) for computing
         # clustered-standard errors. Mahmood Arai, Jan 26, 2008.

	 # The arguments of the function are:
         # fitted model, cluster1 and cluster2
         # You need to install libraries `sandwich' and `lmtest'

         # reweighting the var-cov matrix for the within model
         #library(sandwich);library(lmtest)
         M <- length(unique(cluster))
         N <- length(cluster)
         K <- fm$rank
         dfc <- (M/(M-1))*((N-1)/(N-K))
         uj  <- apply(estfun(fm),2, function(x) tapply(x, cluster, sum));
         vcovCL <- dfc*sandwich(fm, meat=crossprod(uj)/N)*dfcw
         coeftest(fm, vcovCL)
}

mclx <- function(fm, dfcw, cluster1, cluster2){
         # R-codes (www.r-project.org) for computing multi-way
         # clustered-standard errors. Mahmood Arai, Jan 26, 2008.
         # See: Thompson (2006), Cameron, Gelbach and Miller (2006)
         # and Petersen (2006).
	 # reweighting the var-cov matrix for the within model

         # The arguments of the function are:
         # fitted model, cluster1 and cluster2
         # You need to install libraries `sandwich' and `lmtest'

         library(sandwich);library(lmtest)
         cluster12 = paste(cluster1,cluster2, sep="")
         M1  <- length(unique(cluster1))
         M2  <- length(unique(cluster2))
         M12 <- length(unique(cluster12))
         N   <- length(cluster1)
         K   <- fm$rank
         dfc1  <- (M1/(M1-1))*((N-1)/(N-K))
         dfc2  <- (M2/(M2-1))*((N-1)/(N-K))
         dfc12 <- (M12/(M12-1))*((N-1)/(N-K))
         u1j   <- apply(estfun(fm), 2, function(x) tapply(x, cluster1,  sum))
         u2j   <- apply(estfun(fm), 2, function(x) tapply(x, cluster2,  sum))
         u12j  <- apply(estfun(fm), 2, function(x) tapply(x, cluster12, sum))
         vc1   <-  dfc1*sandwich(fm, meat=crossprod(u1j)/N )
         vc2   <-  dfc2*sandwich(fm, meat=crossprod(u2j)/N )
         vc12  <- dfc12*sandwich(fm, meat=crossprod(u12j)/N)
         vcovMCL <- (vc1 + vc2 - vc12)*dfcw
         coeftest(fm, vcovMCL)
}

leecard <- function(binwidth, maxbin, outcome, testtype, emcp=emc) {
    ## outcome = const + B1*age + B2*Cutoff + e
    ## outcome = const + B1*emcp + B2*forcing.sample + e

  breaks <- getbins(binwidth, maxbin)@breaks
  emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

  cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)
  binmids <- hist(emcp[emcdum==1], breaks=breaks, plot=FALSE)$mids

  emcbinmids <- findInterval(emcp[emcdum==1], binmids, rightmost.closed=TRUE)
  emcmids <-  binmids[emcbinmids]

  #binned.outcome <- tapply(outcome[emcdum==1], cutbybreak, function(x) mean(x, na.rm=TRUE))
  #binned.variance <- tapply(outcome[emcdum==1], cutbybreak, function(x) var(x, na.rm=TRUE))

  #binned.test <- tapply(outcome[emcdum==1], cutbybreak, function(x) table(!is.na(x)))
  #Nj <- binned.takers <- unlist(lapply(binned.test, function(x) sum(x["TRUE"])))
  #binned.takers <- binned.takers / (sum(binned.takers) / length(binned.takers))
  #nomissbin <- sapply(names(unlist(binned.takers)), function(x) substr(x, start=1, stop=nchar(x)-5))

  forcing <- ifelse(binmids>0, 1, 0)
  tx <- which(forcing==1)
  forcing.sample <- ifelse(emcp>0, 1, 0)

  #leftcut <- sort(getbins(binwidth, maxbin)@left)[1]
  #rightcut <- sort(getbins(binwidth, maxbin)@right, decreasing=TRUE)[1]
  #rightobs <- ifelse(emc > 0 & emc <= rightcut, 1, 0)
  #leftobs <- ifelse(emc <= 0 & emc > leftcut, 1, 0)
  #obs <- ifelse(emc >= leftcut & emc <= rightcut, 1, 0)
  test <- outcome[emcdum==1][emcbinmids!=0]
  alpha <- forcing.sample[emcdum==1][emcbinmids!=0]

  ##for GOF test
  Unrestricted <- lm(test~alpha+as.factor(emcmids))
  #coef on forcing.sample is NA
  J <- Unrestricted$rank
  N <- dim(Unrestricted$model)[1]
  ESSur <- sum(Unrestricted$residuals^2)

  ####global polynomial models over microdata with cluster consistent standard errors and a GOF test (Card and Lee 2008)
  #p0.T <- lm(outcome~forcing.sample, subset=emcdum==1)
  #clx(p0.T, 1, clustersub)  #cluster consistent standard errors
  #sqrt(vcovHC(p0.T, "HC")) #robust Heteroskedasticity-consistent standard errors
  #K <- p0.T$rank
  #ESSr <- sum(p0.T$residuals^2)
  #G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  #print(pval <- pf(G, J-K, N-J))  GOODNESS OF FIT TEST

  cat("LeeCard GOF Tests", "P-values from F-test", "Polynomial order\n", sep=", ")
  p1.T <- lm(test~alpha+emcmids)
  clustersub <- emcp[attributes(p1.T$model)$row.names]
  c1<-clx(p1.T, 1, clustersub)
  K <- p1.T$rank
  ESSr <- sum(p1.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ") #

  p2.T <- p3.T <- update(p1.T, . ~ . + I(emcmids^2))
  clustersub <- emcp[attributes(p2.T$model)$row.names]
  c2<-clx(p2.T, 1, clustersub)
  K <- p2.T$rank
  ESSr <- sum(p2.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ")

  p3.T <- update(p2.T, . ~ . + I(emcmids^3))
  clustersub <- emcp[attributes(p3.T$model)$row.names]
  c3<-clx(p3.T, 1,clustersub)
  K <- p3.T$rank
  ESSr <- sum(p3.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ")

  p4.T <- update(p3.T, . ~ . + I(emcmids^4))
  clustersub <- emcp[attributes(p4.T$model)$row.names]
  clx(p4.T, 1, clustersub)
  K <- p4.T$rank
  ESSr <- sum(p4.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ")

  p5.T <- update(p4.T, . ~ . + I(emcmids^5))
  clustersub <- emcp[attributes(p5.T$model)$row.names]
  clx(p5.T, 1, clustersub)
  K <- p5.T$rank
  ESSr <- sum(p5.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ")

  p6.T <- update(p5.T, . ~ . + I(emcmids^6))
  clustersub <- emcp[attributes(p6.T$model)$row.names]
  clx(p6.T, 1, clustersub)
  K <- p6.T$rank
  ESSr <- sum(p6.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ")

  p7.T <- update(p6.T, . ~ . + I(emcmids^7))
  clustersub <- emcp[attributes(p7.T$model)$row.names]
  clx(p7.T, 1, clustersub)
  K <- p7.T$rank
  ESSr <- sum(p7.T$residuals^2)
  G = ((ESSr-ESSur)/(J-K)) / (ESSur / (N-J)) # ~F(J-K, N-J)
  cat(pval <- pf(G, J-K, N-J), K-2, "\n", sep=", ")

  ################################### Global fits end

  ## regs run on collapsed to cell level data with heteroskedasticity-consistent standard errors.
  ##Also include variance correction factor as in Card and Lee 2008 equations 12 and 13.
  ## Omit for now

  #p0.C <- lm(binned.outcome~forcing, weights=binned.takers)
  #sqrt(vcovHC(p0.C, "HC")) #robust Heteroskedasticity-consistent standard errors: about equal to clx(p0.T, 1, clustersub)
  #sigmahat2_a <- 2*(sum(p0.C$residuals^2)/sum(Nj) - sum(binned.variance, na.rm=T)/sum(Nj))  #maybe this one?
  #newstderror <- vcovHC(p0.C, "HC")[2,2] + sigmahat2_a

  #p1.C <- lm(binned.outcome~forcing*breaks[2:81], weights=binned.takers)
  #sqrt(vcovHC(p1.C, "HC"))
  #sigmahat2_a <- 2*(sum(p1.C$residuals^2 * Nj[1:79]) / sum(Nj)  - sum(binned.variance, na.rm=T)/sum(Nj)) #maybe this one?
  #newstderror <- vcovHC(p1.C, "HC")[2,2] + sigmahat2_a

  #p2.C <- lm(outcome~forcing.sample*emc + forcing.sample*I(emc^2), subset=obs==1)
  return(list(c1, c2, c3))
}

srdregs <- function(binwidth, maxbin, outcome, testtype, emcp=emc, coefs=1, use.coefs=FALSE) {
    ## outcome = const + B1*age + B2*Cutoff + e
    ## outcome = const + B1*emcp + B2*forcing.sample + e

  breaks <- getbins(binwidth, maxbin)@breaks
  emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

  cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)
  binmids <- hist(emcp[emcdum==1], breaks=breaks, plot=FALSE)$mids

  emcbinmids <- findInterval(emcp[emcdum==1], binmids, rightmost.closed=TRUE)
  emcmids <-  binmids[emcbinmids]

  forcing <- ifelse(binmids>0, 1, 0)
  #tx <- which(forcing==1)
  forcing.sample <- ifelse(emcp>0, 1, 0)

  test <- outcome[emcdum==1][emcbinmids!=0]
  alpha <- forcing.sample[emcdum==1][emcbinmids!=0]
  if (use.coefs){
    beta <- coefs[emcdum==1,][emcbinmids!=0,]
    sr.model <- lm(test~alpha+emcmids+I(emcmids^2)+beta)
  } else sr.model <- lm(test~alpha+emcmids+I(emcmids^2))

  clustersub <- emcp[attributes(sr.model$model)$row.names]
  adjrsq <- summary(sr.model)$adj.r.squared

  ##temporary fix for now until I figure out prob with sandwich estimator
  a<-try(print(clx(sr.model, 1, clustersub)), silent=TRUE)
  if (class(a)=="try-error") print(summary(sr.model))
  cat("N=",length(clustersub), " Adj. R-sq=", adjrsq, '\n\n', sep='')

}

frdregs <- function(binwidth, maxbin, outcome, testtype, emcp=emc, coefs=1, use.coefs=FALSE, instr) {
    ## outcome = const + B1*age + B2*Cutoff + e
    ## outcome = const + B1*emcp + B2*forcing.sample + e
    ## forcing.sample =  const+A1*age

  breaks <- getbins(binwidth, maxbin)@breaks
  emcdum <- ifelse(emcp >= min(breaks) & emcp <= max(breaks), 1, 0)

  cutbybreak <- cut(emcp[emcdum==1], breaks=breaks)
  binmids <- hist(emcp[emcdum==1], breaks=breaks, plot=FALSE)$mids

  emcbinmids <- findInterval(emcp[emcdum==1], binmids, rightmost.closed=TRUE)
  emcmids <-  binmids[emcbinmids]

  forcing <- ifelse(binmids>0, 1, 0)
  #tx <- which(forcing==1)
  forcing.sample <- ifelse(emcp>0, 1, 0)

  test <- outcome[emcdum==1][emcbinmids!=0]
  alpha <- forcing.sample[emcdum==1][emcbinmids!=0]
  instrument <- instr[emcdum==1][emcbinmids!=0]

  if (use.coefs){
    beta <- coefs[emcdum==1,][emcbinmids!=0,]
    mf <- data.frame(test, alpha, instrument, emcmids, beta)
    form1 <- as.formula(paste("test~emcmids+I(emcmids^2)+alpha+",
                   paste(colnames(mf[,5:dim(mf)[2]]), collapse="+"),
                         " | instrument+emcmids+I(emcmids^2)+",paste(colnames(mf[,5:dim(mf)[2]]), collapse="+"), sep=''))
    sr.model <- ivreg(form1, data=mf)
  } else {
    mf <- data.frame(test, alpha, instrument, emcmids)
    sr.model <- ivreg(test~emcmids+I(emcmids^2)+alpha | instrument+emcmids+I(emcmids^2), data=mf)
  }

  clustersub <- emcp[attributes(sr.model$model)$row.names]
  adjrsq <- summary(sr.model)$adj.r.squared

  ##temporary fix for now until I figure out prob with sandwich estimator
  a<-try(print(clx(sr.model, 1, clustersub)), silent=TRUE)
  if (class(a)=="try-error") print(summary(sr.model))
  cat("N=",length(clustersub), " Adj. R-sq=", adjrsq, '\n\n', sep='')

}

##$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ FUNCTIONS END $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$##


## ----echo=FALSE---------------------------------------------------------------
require(Hmisc)
knitrSet('continuousY', cache=TRUE)
knitr::read_chunk('~/doc/rms/continuousY/shared.R')

## ----lmassump,h=3.25,w=7,mfrow=c(1,2),ps=10,mgp=c(.5, .365, 0),left=1,cap='Assumptions of the linear model (left panel) and semiparametric ordinal probit or logit (proportional odds) models (right panel).  Ordinal models do not assume any shape for the distribution of $Y$ for a given $X$; they only assume parallelism.  The linear model can relax the parallelism assumption if $\\sigma$ is allowed to vary, but in practice it is difficult to know how to vary it except for the unequal variance two-sample $t$-test.',scap='Assumptions of linear vs.\\ semiparametric models',echo=FALSE----
pinv <- expression(paste(Phi^{-1},  '(F(y', '|', 'X))'))
plot(0, 0, xlim=c(0, 1), ylim=c(-2, 2), type='n', axes=FALSE,
     xlab=expression(y), ylab='')
mtext(pinv, side=2, line=1)
axis(1, labels=FALSE, lwd.ticks=0)
axis(2, labels=FALSE, lwd.ticks=0)
abline(a=-1.5, b=1)
abline(a=0, b=1)
arrows(.5, -1.5+.5, .5, 0+.5, code=3, length=.1)
text(.525, .5*(-1.5+.5+.5), expression(-Delta*X*beta/sigma), adj=0)
g <- function(x) -2.2606955+11.125231*x-37.772783*x^2+56.776436*x^3-
  26.861103*x^4
x <- seq(0, .9, length=150)
pinv <- expression(atop(paste(Phi^{-1},  '(F(y', '|', 'X))'),
    paste(logit, '(F(y', '|', 'X))')))
plot(0, 0, xlim=c(0, 1), ylim=c(-2, 2), type='n', axes=FALSE,
     xlab=expression(y), ylab='')
mtext(pinv, side=2, line=1)
axis(1, labels=FALSE, lwd.ticks=FALSE)
axis(2, labels=FALSE, lwd.ticks=FALSE)
lines(x, g(x))
lines(x, g(x)+1.5)
arrows(.5, g(.5), .5, g(.5)+1.5, code=3, length=.1)
text(.525, .5*(g(.55) + g(.55)+1.5), expression(-Delta*X*beta), adj=0)

## ----desc,results='asis',cache=FALSE------------------------------------------
require(rms)
options(prType='latex')   # for print, summary, anova
getHdata(nhgh)
w <- subset(nhgh, age >= 21 & dx==0 & tx==0, select=-c(dx,tx))
latex(describe(w), file='')
dd <- datadist(w); options(datadist='dd')

## ----lookdist,w=5.5,h=5.5,cap='Examination of normality and constant variance assumption, and assumptions for various ordinal models',scap='Examining normality and ordinal model assumptions'----
f <- ols(gh ~ rcs(age,5) + sex + re + rcs(bmi, 3), data=w)
pgh <- fitted(f)

p <- function(fun, row, col) {
  f <- substitute(fun); g <- function(F) eval(f)
  z <- Ecdf(~ gh, groups=cut2(pgh, g=6),
            fun=function(F) g(1 - F),
            ylab=as.expression(f), xlim=c(4.5, 7.75), data=w,
            label.curve=FALSE)
  print(z, split=c(col, row, 2, 2), more=row < 2 | col < 2)
}
p(log(F/(1-F)),   1, 1)
p(qnorm(F),       1, 2)
p(-log(-log(F)),  2, 1)
p(log(-log(1-F)), 2, 2)
# Get slopes of pgh for some cutoffs of Y
# Use glm complementary log-log link on Prob(Y < cutoff) to
# get log-log link on Prob(Y >= cutoff)
r <- NULL
for(link in c('logit','probit','cloglog'))
  for(k in c(5, 5.5, 6)) {
    co <- coef(glm(gh < k ~ pgh, data=w, family=binomial(link)))
    r <- rbind(r, data.frame(link=link, cutoff=k,
                             slope=round(co[2],2)))
}
print(r, row.names=FALSE)

## ----comparemany,cache=TRUE,h=6.5,w=6.75,cap='Three estimated quantiles and estimated mean using 6 methods, compared against caliper-matched sample quantiles/means (circles).  Numbers are mean absolute differences between predicted and sample quantities using overlapping intervals of age and caliper matching\\index{caliper matching}.  QR:quantile regression.',scap='Six methods for estimating quantiles or means.'----
ag <- 25:75
lag <- length(ag)
q2 <- q3 <- p90 <- means <- numeric(lag)
for(i in 1:lag) {
  s <- which(abs(w$age - ag[i]) < 5)
  y <- w$gh[s]
  a <- quantile(y, probs=c(.5, .75, .9))
  q2[i]    <- a[1]
  q3[i]    <- a[2]
  p90[i]   <- a[3]
  means[i] <- mean(y)
}
fams <- c('logistic', 'probit', 'loglog', 'cloglog')
fe   <- function(pred, target) mean(abs(pred$yhat - target))
mod  <- gh ~ rcs(age,6)
P    <- Er <- list()
for(est in c('q2', 'q3', 'p90', 'mean')) {
  meth <- if(est == 'mean') 'ols' else 'QR'
  p <- list()
  er <- rep(NA, 5)
  names(er) <- c(fams, meth)
  for(family in fams) {
    h  <- orm(mod, family=family, data=w)
    fun <- if(est == 'mean') Mean(h)
    else {
      qu <- Quantile(h)
      switch(est, q2  = function(x) qu(.5,  x),
                  q3  = function(x) qu(.75, x),
                  p90 = function(x) qu(.9,  x))
    }
    p[[family]] <- z <- Predict(h, age=ag, fun=fun, conf.int=FALSE)
    er[family] <- fe(z, switch(est, mean=means, q2=q2, q3=q3, p90=p90))
  }
  h <- switch(est,
              mean= ols(mod, data=w),
              q2  = Rq (mod, data=w),
              q3  = Rq (mod, tau=0.75, data=w),
              p90 = Rq (mod, tau=0.90, data=w))
  p[[meth]] <- z <- Predict(h, age=ag, conf.int=FALSE)
  er[meth] <- fe(z, switch(est, mean=means, q2=q2, q3=q3, p90=p90))

  Er[[est]] <- er
  pr <- do.call('rbind', p)
  pr$est <- est
  P <- rbind.data.frame(P, pr)
}

xyplot(yhat ~ age | est, groups=.set., data=P, type='l', # Figure (*\ref{fig:continuousY-comparemany}*)
       auto.key=list(x=.75, y=.2, points=FALSE, lines=TRUE),
       panel=function(..., subscripts) {
         panel.xyplot(..., subscripts=subscripts)
         est <- P$est[subscripts[1]]
         lpoints(ag, switch(est, mean=means, q2=q2, q3=q3, p90=p90),
                 col=gray(.7))
         er <- format(round(Er[[est]],3), nsmall=3)
         ltext(26, 6.15, paste(names(er), collapse='\n'),
               cex=.7, adj=0)
         ltext(40, 6.15, paste(er, collapse='\n'),
               cex=.7, adj=1)})

## ----predobs,w=4.75,h=4,cap='Observed (dashed lines, open circles) and predicted (solid lines, closed circles) exceedance probability distributions from a model using 6-tiles of OLS-predicted \\hba.  Key shows quantile group intervals of predicted mean \\hba.',scap='Observed and predicted distributions'----
w$pghg <- cut2(pgh, g=6)
f  <- orm(gh ~ pghg, family=loglog, data=w)
lp <- predict(f, newdata=data.frame(pghg=levels(w$pghg)))
ep <- ExProb(f)  # Exceedance prob. functn. generator in rms
z  <- ep(lp)
j  <- order(w$pghg)  # puts in order of lp (levels of pghg)
plot(z, xlim=c(4, 7.5), data=w[j,c('pghg', 'gh')]) # Fig. (*\ref{fig:continuousY-predobs}*)

## ----lookprobit,w=3,h=2.5,bty='l',cap='Estimated intercepts from probit model.  Linearity would have indicated Gaussian residuals.',scap='Estimated intercepts from probit model'----
f <- orm(gh ~ rcs(age,6), family=probit, data=w)
g <- ols(gh ~ rcs(age,6), data=w)
s <- g$stats['Sigma']
yu <- f$yunique[-1]
r <- quantile(w$gh, c(.005, .995))
alphas <- coef(f)[1:num.intercepts(f)]
plot(-yu / s, alphas, type='l', xlim=rev(- r / s), # Fig. (*\ref{fig:continuousY-lookprobit}*)
     xlab=expression(-y/hat(sigma)), ylab=expression(alpha[y]))

## ----htwtcoef,results='asis'--------------------------------------------------
f <- orm(gh ~ rcs(age,5) + log(ht) + log(wt),
         family=loglog, data=w)
f

## ----aichtwt------------------------------------------------------------------
aic <- NULL
for(mod in list(gh ~ rcs(age,5) + rcs(log(bmi),5),
                gh ~ rcs(age,5) + rcs(log(ht),5) + rcs(log(wt),5),
                gh ~ rcs(age,5) + rcs(log(ht),4) * rcs(log(wt),4)))
  aic <- c(aic, AIC(orm(mod, family=loglog, data=w)))
print(aic)

## ----coxhtwtcoef,results='asis'-----------------------------------------------
print(cph(Surv(gh) ~ rcs(age,5) + log(ht) + log(wt), data=w))

## ----redun,cap='Variable clustering for all potential predictors'-------------
v <- varclus(~ wt + ht + bmi + leg + arml + armc + waist +
             tri + sub + age + sex + re, data=w)
plot(v)   # Figure (*\ref{fig:continuousY-redun}*)
# Omit wt so it won't be removed before bmi
redun(~ ht + bmi + leg + arml + armc + waist + tri + sub,
      data=w, r2=.75)

## ----htchange,fig.align='right',cap="Estimated median height as a smooth function of age, allowing age to interact with sex, from a proportional odds model",scap="Median height vs.\\ age"----
f <- orm(ht ~ rcs(age,4)*sex, data=w)  # Prop. odds model
qu <- Quantile(f); med <- function(x) qu(.5, x)
ggplot(Predict(f, age, sex, fun=med, conf.int=FALSE),
       ylab='Predicted Median Height, cm')

## ----allocadf,w=4,cap='Generalized squared rank correlations',top=1-----------
s <- spearman2(gh ~ age + sex + re + wt + leg + arml + armc +
               waist + tri + sub, data=w, p=2)
plot(s)

## ----fitfullcasewise,results='asis'-------------------------------------------
f <- orm(gh ~ rcs(age,5) + sex + re + rcs(wt,3) + rcs(leg,3) + arml +
         rcs(armc,3) + rcs(waist,4) + tri + rcs(sub,3),
         family=loglog, data=w, x=TRUE, y=TRUE)
print(f, coefs=FALSE)
## Composite test:
anova(f, leg, arml, armc, waist, tri, sub)

## ----casewisemeanmed,h=4,w=5,cap='Estimated mean and 0.5 and 0.9 quantiles from the log-log ordinal model using casewise deletion, along with predictions of 0.5 and 0.9 quantiles from quantile regression (QR).  Age is varied and other predictors are held constant to medians/modes.',scap='Estimated mean and quantiles from casewise deletion model.'----
M      <- Mean(f)
qu     <- Quantile(f)
med    <- function(x) qu(.5, x)
p90    <- function(x) qu(.9, x)
fq     <- Rq(formula(f), data=w)
fq90   <- Rq(formula(f), data=w, tau=.9)
pmean  <- Predict(f,    age, fun=M,   conf.int=FALSE)
pmed   <- Predict(f,    age, fun=med, conf.int=FALSE)
p90    <- Predict(f,    age, fun=p90, conf.int=FALSE)
pmedqr <- Predict(fq,   age, conf.int=FALSE)
p90qr  <- Predict(fq90, age, conf.int=FALSE)
z <- rbind('orm mean'=pmean, 'orm median'=pmed, 'orm P90'=p90,
           'QR median'=pmedqr, 'QR P90'=p90qr)
ggplot(z, groups='.set.',
       adj.subtitle=FALSE, legend.label=FALSE)

## ----prbw---------------------------------------------------------------------
print(fastbw(f, rule='p'), estimates=FALSE)

## ----valbworm,cache=TRUE------------------------------------------------------
set.seed(13)  # so can reproduce results
v <- validate(f, B=100, bw=TRUE, estimates=FALSE, rule='p')

## ----prval,results='asis'-----------------------------------------------------
# Show number of variables selected in first 30 boots
latex(v, B=30, file='', size='small')

## ----sanova,cache=TRUE,results='asis',cap='ANOVA for reduced model, after multiple imputation, with addition of a combined effect for four size variables',scap='ANOVA for reduced model'----
a <- aregImpute(~ gh + wt + ht + bmi + leg + arml + armc + waist +
                tri + sub + age +re, data=w, n.impute=5, pr=FALSE)
g <- fit.mult.impute(gh ~ rcs(age,5) + re + rcs(leg,3) +
                     rcs(waist,4) + tri + rcs(sub,4),
                     orm, a, family=loglog, data=w, pr=FALSE)
print(g, needspace='1.5in')
an <- anova(g)
print(an, caption='ANOVA for reduced model after multiple imputation, with addition of a combined effect for four size variables')
b  <- anova(g, leg, waist, tri, sub)
# Add new lines to the plot with combined effect of 4 size var.
s <- rbind(an, size=b['TOTAL', ])
class(s) <- 'anova.rms'
plot(s)

## ----peffects,cache=TRUE,results='asis',cap='Partial effects (log hazard or log-log cumulative probability scale) of all predictors in reduced model, after multiple imputation',scap='Partial effects after multiple imputation',w=6.75,h=4.5,cache=TRUE----
ggplot(Predict(g), abbrev=TRUE, ylab=NULL)   # Figure (*\ref{fig:continuousY-peffects}*)

## ----cfmissmeth,cache=TRUE----------------------------------------------------
gc <- orm(gh ~ rcs(age,5) + re + rcs(leg,3) +
          rcs(waist,4) + tri + rcs(sub,4),
          family=loglog, data=w, x=TRUE, y=TRUE)
gb <- bootcov(gc, B=300)

## ----peffects2,cap='Partial effect for age from multiple imputation (center red line) and casewise deletion (center blue line) with symmetric Wald 0.95 confidence bands using casewise deletion (gray shaded area), basic bootstrap confidence bands using casewise deletion (blue lines), percentile bootstrap confidence bands using casewise deletion (dashed blue lines), and symmetric Wald confidence bands accounting for multiple imputation (red lines).',scap='Partial effect for age with bootstrap and Wald confidence bands',w=5,h=4,bot=1----
bootclb <- Predict(gb, age, boot.type='basic')
bootclp <- Predict(gb, age, boot.type='percentile')
multimp <- Predict(g,  age)
plot(Predict(gc, age), addpanel=function(...) {
  with(bootclb, {llines(age, lower, col='blue')
                 llines(age, upper, col='blue')})
  with(bootclp, {llines(age, lower, col='blue', lty=2)
                 llines(age, upper, col='blue', lty=2)})
  with(multimp, {llines(age, lower, col='red')
                 llines(age, upper, col='red')
                 llines(age, yhat, col='red')} ) },
     col.fill=gray(.9), adj.subtitle=FALSE)   # Figure (*\ref{fig:continuousY-peffects2}*)

## ----meanvs,cap='Predicted mean \\hba vs.\\ predicted median and 0.9 quantile along with their marginal distributions',scap='Predicted mean, median, and 0.9 quantile of \\hba',w=4.5,h=3.5----
M  <- Mean(g)
qu <- Quantile(g)
med <- function(lp) qu(.5, lp)
q90 <- function(lp) qu(.9, lp)
lp  <- predict(g)
lpr <- quantile(predict(g), c(.002, .998), na.rm=TRUE)
lps <- seq(lpr[1], lpr[2], length=200)
pmn <- M(lps)
pme <- med(lps)
p90 <- q90(lps)
plot(pmn, pme,   # Figure (*\ref{fig:continuousY-meanvs}*)
     xlab=expression(paste('Predicted Mean ',  HbA["1c"])),
     ylab='Median and 0.9 Quantile', type='l',
     xlim=c(4.75, 8.0), ylim=c(4.75, 8.0), bty='n')
box(col=gray(.8))
lines(pmn, p90, col='blue')
abline(a=0, b=1, col=gray(.8))
text(6.5, 5.5, 'Median')
text(5.5, 6.3, '0.9', col='blue')
nint <- 350
scat1d(M(lp),   nint=nint)
scat1d(med(lp), side=2, nint=nint)
scat1d(q90(lp), side=4, col='blue', nint=nint)

## ----nomogram,cap='Nomogram for predicting median, mean, and 0.9 quantile of glycohemoglobin, along with the estimated probability that \\hba $\\ge 6.5, 7$, or $7.5$, all from the log-log ordinal model',scap='Nomogram of log-log ordinal model for \\hba',w=6.75,h=5.75,ps=9----
g      <- Newlevels(g, list(re=abbreviate(levels(w$re))))
exprob <- ExProb(g)
nom <-
  nomogram(g, fun=list(Mean=M,
                'Median Glycohemoglobin' = med,
                '0.9 Quantile'           = q90,
                'Prob(HbA1c >= 6.5)'=
                     function(x) exprob(x, y=6.5),
                'Prob(HbA1c >= 7.0)'=
                     function(x) exprob(x, y=7),
                'Prob(HbA1c >= 7.5)'=
                     function(x) exprob(x, y=7.5)),
           fun.at=list(seq(5, 8, by=.5),
             c(5,5.25,5.5,5.75,6,6.25),
             c(5.5,6,6.5,7,8,10,12,14),
             c(.01,.05,.1,.2,.3,.4),
             c(.01,.05,.1,.2,.3,.4),
             c(.01,.05,.1,.2,.3,.4)))
plot(nom, lmgp=.28)   # Figure (*\ref{fig:continuousY-nomogram}*)


## load package deSolve for ode simulation
library(deSolve)

## types of noise:
## 1. process noise: complementary or essential resources in the main text model. core, tail, or general in the model in Appendix D
## 2. measurement noise
noisetypes=c('complementary','essential','core','tail','general','measurement')

## scaling factors: scale noise levels in the kernel
scaling.factors=data.frame(noise=c('complementary','essential','core','tail','general','measurement'),scaling=c(30,1,8,250,2,.4))

## number of species
S=200

## species niche/trait values
set.seed(0)
trait=sort(runif(S))
d=as.matrix(dist(trait))
d[d>max(d)/2]=max(d)-d[d>max(d)/2]

## kernel width
w=.1

## noise-free competition kernel
A0=exp(-(d/w)^4); A0=A0/mean(A0)*.2

## Function to generate the competition kernel.
## Inputs: species x-niche/trait values, type of noise, degree of noise, simulation number (run).
## Outputs: A = noisy kernel in the case of process noise, noise-free kernel in the case of measurement noise
## (cont) A0 = noise-free kernel in the case of process noise, noisy kernel in the case of measurement noise;
## (cont) d = distances on x-niche/trait axis (on a circular axis);
## (cont) prox = proxy trait values in the case of measurement noise; traity = y-niche/trait values;
## (cont) spear = Spearman's rho, correlation between kernel and distances on x-axis;
## (cont) pi = P(Atail > Amean) - P(Acore > Amean), quantifies core-tail structure in the kernel
Kernel=function(noiselevel,noisetype,traitvalues=trait,run=0){
	stopifnot(noisetype%in%noisetypes)

	## tie random number generator to run
	set.seed(run)

	## number of species
	S=length(traitvalues)

	## scaled noise level
	a0=noiselevel*with(scaling.factors,scaling[match(noisetype,noise)])

	## distances on x-niche/trait axis (axis is circular)
	d=as.matrix(dist(traitvalues)); d[d>max(d)/2]=max(d)-d[d>max(d)/2]

	## A0 = noise-free kernel in the case of process noise, noisy kernel in the case of measurement noise
	A0=exp(-(d/w)^4); A0=A0/mean(A0)*.2

	## threshold distance separating core from tail
	dcoretail=d[which.min(abs(A0-mean(A0)))]

	## traity = y-niche/trait values; will contribute to kernel in proportion to noise level
	traity=NULL

	## A = noisy kernel in the case of process noise, noise-free kernel in the case of measurement noise
	if(noisetype%in%c('complementary','essential')){ traity=runif(S); dy=as.matrix(dist(traity)); dy[dy>max(dy)/2]=max(dy)-dy[dy>max(dy)/2]}
	if(noisetype=='complementary'){ wyinverse=a0; A=exp(-sqrt((d/w)^2+(dy*wyinverse)^2)^4)}		## from 1 to 10
	if(noisetype=='essential'){ wy=a0; A=pmax(exp(-(d/w)^4),exp(-(dy/wy)^4))}				## from .01 to .15

	if(noisetype%in%c('core','tail','general')){
		sigma = if(noisetype=='core') a0*exp(-(d/.1)^4) else if(noisetype=='tail') a0*(exp(d^4)-1) else if(noisetype=='general') a0
		A=matrix(pmax(0,rnorm(S*S,mean=exp(-(d/w)^4),sd=sigma)),S,S)
		while(TRUE){ i=which(diag(A)==0); if(length(i)==0) break; diag(A)[i]=pmax(0,rnorm(length(i),mean=1,sd=a0))}
	}

	## prox = proxy trait values in the case of measurement noise
	prox=NULL
	if(noisetype=='measurement'){
		prox=rnorm(S,mean=traitvalues,sd=a0); d=as.matrix(dist(prox)); d[d>max(d)/2]=max(d)-d[d>max(d)/2]
		A=A0; A0=exp(-(d/w)^4); A0=A0/mean(A0)*.2
	}

	## normalize A
	A=A/mean(A)*.2

	## spear = Spearman's rho; pi = P(Atail > Amean) - P(Acore > Amean)
	spear=cor(as.numeric(d),as.numeric(A),method='spearman')											## -1 (perfect) to 0 (null)
	pi=sum(A[d>=dcoretail]>mean(A))/sum(d>=dcoretail)-(sum(A[d<=dcoretail]>mean(A))/sum(d<=dcoretail))	## -1 (perfect) to 0 (null)

	return(list(
		A=A,
		A0=A0,
		d=d,
		prox=prox,
		traity=traity,
		spear=spear,
		pi=pi
	))
}

Run=function(traits,compkernel,immig,maxtime=1e4,plot=1,logplot=1){
	## initial abundances
	x0=rep(1e-3,nrow(compkernel))+runif(nrow(compkernel),min=0,max=1e-3/100)

	## time series
	times=seq(0,maxtime,by=.1)

	## Lotka-Volterra dynamic equations
	LV=function(times,N,parms) list(pmax(N,0)*(1-as.numeric(parms$Kernel%*%pmax(N,0)))+parms$m)

	## Simulate model
	x=ode(y=x0,times=times,func=LV,parms=list(Kernel=compkernel,m=immig),method='lsoda'); n=x[nrow(x),-1]; n[n<1e-5]=0

	## Calculate richness at each time point
	s=apply(x[,-1],1,function(v) sum(v>1e-5))

	## Plot results
	if(plot) if(!logplot) plot(trait,n,t='h',las=1) else plot(trait[n>0],log(1+n[n>0]/min(n[n>0])),t='h',las=1)

	return(list(data=data.frame(trait=traits,abundance=n),richness=data.frame(time=x[,1],S=s)))
}


#' Function that optimizes sinusoid parameters to fit d18O data
#'
#' The second core function of the ShellChron growth model. Loops
#' through all data windows and uses the \code{growth_model} function
#' to create d18O series that match the input data. This step is
#' iterated and optimized (minimizing the Sum of Squared Residuals)
#' through the SCEUA algorithm (by Duan et al., 1992) which finds
#' the optimal input parameters to the growth rate and Sea Surface
#' Temperature (SST) sinusoids to simulate d18O data.
#' @param dat Matrix containing the input data
#' @param dynwindow Information on the position and length of modeling
#' windows
#' @param transfer_function Transfer function used to convert d18Oc to temperature
#' data.
#' @param d18Ow Either a single value (constant d18Ow) or a vector of length
#' equal to the period in SST data (365 days by default) containing information
#' about seasonality in d18Ow. Defaults to constant d18Ow of 0 permille VSMOW
#' (the modern mean ocean value)
#' @param T_per Period of SST sinusoid (in days; default = 365)
#' @param G_per Period of growth rate sinusoid (in days; default = 365)
#' @param t_int Time interval (in days; default = 1)
#' @param t_maxtemp Timing of the warmest day of the year (in julian day;
#' default = 182.5, or May 26th halfway through the year)
#' @param SCEUApar Parameters for SCEUA optimization (iniflg, ngs, maxn, kstop
#' pcento, peps). For details, refer to Duan et al. (1992) in references
#' @param sinfit Apply sinusoidal fitting to guess initial parameters for SCEUA
#' optimization? \code{TRUE/FALSE}
#' @param MC Number of Monte Carlo simulations to apply for error propagation
#' Default = 1000
#' @param plot Should results of modeling be plotted? \code{TRUE/FALSE}
#' @return A list containing the \code{resultarray} which contains the full
#' result of all simulations on each data window and the \code{parmat} listing
#' all optimized growth rate and SST parameters used to model d18O in each data
#' window
#' @seealso Duan, Qingyun, Soroosh Sorooshian, and Vijai Gupta. "Effective and
#' efficient global optimization for conceptual rainfall runoff models." Water
#' resources research 28.4 (1992): 1015-1031. https://doi.org/10.1029/91WR02985
#' @references package dependencies: ggplot2 3.2.1; rtop 0.5.14
#' Function dependencies: sinreg, d18O_model, growth_model
#'
#' \doi{10.1029/91WR02985}
#' @examples
#' # Create dummy input data column by column
#' dat <- as.data.frame(seq(1000, 40000, 1000))
#' colnames(dat) <- "D"
#' dat$d18Oc <- sin((2 * pi * (seq(1, 40, 1) - 8 + 7 / 4)) / 7)
#' dat$YEARMARKER <- c(0, rep(c(0, 0, 0, 0, 0, 0, 1), 5), 0, 0, 0, 0)
#' dat$D_err <- rep(100, 40)
#' dat$d18Oc_err <- rep(0.1, 40)
#' # Create dummy dynwindow data
#' dynwindow <- as.data.frame(seq(1, 29, 2))
#' colnames(dynwindow) <- "x"
#' dynwindow$y <- rep(12, 15)
#' # Run model function
#' \donttest{resultlist <- run_model(dat = dat,
#'     dynwindow = dynwindow,
#'     transfer_function = "KimONeil97",
#'     d18Ow = 0,
#'     T_per = 365,
#'     G_per = 365,
#'     t_int = 1,
#'     t_maxtemp = 182.5,
#'     SCEUApar = c(1, 25, 10000, 5, 0.01, 0.01),
#'     sinfit = TRUE,
#'     MC = 1000,
#'     plot = FALSE)}
#' @export
run_model <- function(dat, # Core function to run the entire model on the data (dat)
    dynwindow, # The window vetor resulting from reading in the data
    transfer_function = "KimONeil97",
    d18Ow = 0,
    T_per = 365, # Temperature sinusoid parameters
    G_per = 365, # Growth sinusoid parameters
    t_int = 1, # Default time interval = 1 day
    t_maxtemp = 182.5, # Default time (day) at which maximum temperature is reached is 182.5 (exactly halfway through the year, or 1st of June)
    SCEUApar = c(1, 25, 10000, 5, 0.01, 0.01), # Set parameters for SCEUA optimization (iniflg, ngs, maxn, kstop, pcento, peps)
    sinfit = TRUE, # Apply sinusoidal fitting to guess initial parameters for SCEUA optimization? (TRUE/FALSE)
    MC = 1000, # If errors = TRUE, give the number of iterations for Monte Carlo simulation used in error propagation (default = 1000, if MC = 0 no eror propagation is done)
    plot = FALSE # Should the progress of model fitting be plotted?
    ){

    d18Oc <- d18Oc_err <- Omod <- NULL # Predefine variables to circumvent global variable binding error

    # Prepare data arrays for storage of modeling results
    resultarray <- array( # Create array to contain all modeling results of overlapping windows
        rep(as.matrix(cbind(dat, matrix(NA, ncol = length(dynwindow$x), nrow = length(dat$D)))), 9), # Replicate matrix of dat + extra columns to contain all variables
        dim = c(length(dat$D), length(dynwindow$x) + length(dat[1,]), 9) # Six variables, being: Modeled d18O, residuals, Day of the Year, Growth between datapoints, Instantaneous growth rate at datapoint and Modeled temperature
    )

    parmat <- matrix(NA, nrow = 7, ncol = length(dynwindow$x)) # Matrix in which to store the modeling parameters
    colnames(parmat) <- dynwindow$x
    rownames(parmat) <- c("T_amp", "T_pha", "T_av", "G_amp", "G_pha", "G_av", "G_skw")

    # Prepare plot to show model progress

    if(plot == TRUE){
        dev.new()
        fitplot <- ggplot2::ggplot(dat, ggplot2::aes(D, d18Oc)) + # Create a plot showing the fit of each window on the original data, start with a plot of the original data
        ggplot2::geom_point() +
        ggplot2::geom_line() +
        ggplot2::geom_errorbar(ggplot2::aes(ymin = d18Oc - d18Oc_err, # Add error bars on model result (1 SD)
            ymax = d18Oc + d18Oc_err),
            width = dat$D_err) +
        ggplot2::geom_errorbarh(ggplot2::aes(xmin = D - D_err,
            xmax = D + D_err),
            height = 0.05) +
        ggplot2::ggtitle("Plot of d18Oc values vs. depth. Black = data, Red = model, Errorbars = 1 SD")
        plot(fitplot)
    }

    # Estimate growth rate variability and round up to nearest higher magnitude of 10 for conservative boundary
    GRavest <- max(diff(dat[dat$YEARMARKER == 1,1])) / 365 # Estimate maximum growth rate from yearmarkers
    GRavmax <- 10 ^ (ceiling(log(GRavest, 10))) # Round up to nearest higher magnitude of 10

    # Find tailored range of temperatures from data
    d18Oc_range <- range(dat$d18Oc) # Find d18Oc range in data
    if(transfer_function == "KimONeil97"){ # Find temperature range (to be superseded with inverse d18O_model function in later updates)
        T_range <- 18.03 * 1000 / (log((d18Oc_range - (0.97002 * rev(range(d18Ow)) - 29.98)) / 1000 + 1) * 1000 + 32.42) - 273.15 # Use Kim and O'Neil (1997) with conversion between VSMOW and VPDB by Brand et al. (2014)
    }else if(transfer_function == "GrossmanKu86"){
        T_range <-  20.6 - 4.34 * (d18Oc_range - rev(range(d18Ow)) - 0.2) # Use Grossmann and Ku (1986) modified by Dettmann et al. (1999)
    }else{
        print("ERROR: Supplied transfer function is not recognized")
    }
    T_max <- max(T_range)
    T_amp_max <- 2 * abs(diff(T_range))

    # Collate lower boundaries of parameters
    parl <- c(
        T_amp = 0, # Minimum T amplitude in degrees C
        T_pha = 0, # Minimum phase in days
        T_av = -4, # Minimum average T in degrees C
        G_amp = 0, # Minimum seasonal GR range in um/d
        G_pha = 0, # Minimum GR phase in days
        G_av = -1 * GRavmax, # Minimum average GR in um/d.
        G_skw = 0 # Minimum skew factor
    )

    # Collate upper boundaries of parameters
    paru <- c(
        T_amp = round(T_amp_max + 0.5, 0), # Maximum T amplitude in degrees C
        T_pha = 365, # Maximum phase in days
        T_av = round(T_max + 0.5, 0), # Maximum average T in degrees C
        G_amp = 2 * GRavmax, # Maximum seasonal GR range in um/d
        G_pha = 365, # Maximum GR phase in days
        G_av = GRavmax, # Maximum average GR in um/d (based on conservative boundaries of YEARMARKER indicators)
        G_skw = 100 # Maximum skew factor
    )

    # Set parameters for SCEUA optimization
    iniflg = SCEUApar[1] # Flag for initial parameter array (default = 1; included)
    ngs = SCEUApar[2] # Number of complexes (sub-populations, default = 25)
    maxn = SCEUApar[3] # Maximum number of function evaluations allowed during optimization (default = 10000)
    kstop = SCEUApar[4] # Maximum number of evolution loops before convergency (default = 5)
    pcento = SCEUApar[5] # Percentage change allowed in function value criterion before stop (default = 0.01)
    peps = SCEUApar[6] # Convergence level for parameter set (difference between parameters required for stop; default = 0.01)

    # Run the model on all windows

    for(i in 1:length(dynwindow$x)){ # Loop over shell record
        print(paste("Processing Datawindow:", i, "of", length(dynwindow$x))) # Keep track of progress

        # Isolate year of d18O data based on window data
        Dsam <- dat[dynwindow$x[i]:(dynwindow$x[i] + dynwindow$y[i] - 1), 1]
        Osam <- dat[dynwindow$x[i]:(dynwindow$x[i] + dynwindow$y[i] - 1), 2]
        if(MC > 0){
            D_err <- dat[dynwindow$x[i]:(dynwindow$x[i] + dynwindow$y[i] - 1), 4] # Optional: include error on D
            O_err <- dat[dynwindow$x[i]:(dynwindow$x[i] + dynwindow$y[i] - 1), 5] # Optional: include error on d18Oc
        }else{
            D_err <- rep(0, dynwindow$y)
            O_err <- rep(0, dynwindow$y)
            MC <- 0
        }

        if(sinfit){ # If sinusoidal fitting is enabled
            sinlist <- sinreg(Dsam, Osam) # Run sinusoidal regression to find initial parameter values
            # Estimate starting parameters from regression results
            O_av_start <- sinlist[[1]][1] # Export starting value for annual d18O average
            O_amp_start <- sinlist[[1]][2] # Export starting value for d18O amplitude
            O_pha_start <- sinlist[[1]][4] %% sinlist[[1]][3] # Estimate position (in depth of the first peak in d18O)
            O_per_start <- sinlist[[1]][3] # Export starting value for period in distance domain
        }else{
            O_av_start <- mean(Osam) # Estimate starting value for annual d18O average by mean of d18O in record
            O_amp_start <- diff(range(Osam)) / 2 # Estimate starting value for d18O amplitude by half the difference between minimum and maximum d18Oc
            O_per_start <- diff(range(Dsam)) # Estimate starting period as thickness of isolated year
            O_pha_start <- 0.25 * O_per_start # Set starting phase to one quarter of a cycle if it cannot be estimated from sinusoidal regression
        }

        if(transfer_function == "KimONeil97"){
            T_av_start <- 18.03 * 1000 / (1000 * log((O_av_start - (0.97002 * mean(d18Ow) - 29.98)) / 1000 + 1) + 32.42) - 273.15  # Estimate mean temperature. Use Kim and O'Neil (1997) with conversion between VSMOW and VPDB by Brand et al. (2014)
            T_amp_start <- 18.03 * 1000 / (1000 * log((O_av_start - O_amp_start - (0.97002 * mean(d18Ow) - 29.98)) / 1000 + 1) + 32.42) - 273.15 - T_av_start # Estimate temperature amplitude. Use Kim and O'Neil (1997) with conversion between VSMOW and VPDB by Brand et al. (2014)
        }else if(transfer_function == "GrossmanKu86"){
            T_av_start <- 20.6 - 4.34 * (O_av_start - mean(d18Ow) - 0.2) # Estimate mean temperature. Use Grossmann and Ku (1986) modified by Dettmann et al. (1999)
            T_amp_start <- 20.6 - 4.34 * (O_av_start - O_amp_start - mean(d18Ow) - 0.2) - T_av_start # Estimate mean temperature. Use Grossmann and Ku (1986) modified by Dettmann et al. (1999)
        }else{
            print("ERROR: Supplied transfer function is not recognized")
        }

        O_peak <- O_pha_start + Dsam[1] # Find position of d18O peak in distance domain
        T_pha_start <- ((O_pha_start - 0.5 * O_per_start) %% O_per_start) / O_per_start * T_per # Estimate position of first peak in temperature (low in d18O) relative to annual cycle (days)
        G_av_start <- O_per_start / G_per # Estimate average growth rate in distance/day

        years <- 3 # Set default number of years to 3

        # Collate starting parameters
        par0 <- c(
            T_amp = T_amp_start,
            T_pha = T_pha_start,
            T_av = T_av_start,

            G_amp = G_av_start / 2, # Start by estimating growth rate changes by half the average
            G_pha = T_pha_start, # Start by estimating that the peak in growth rate coincides with the peak in temperature
            G_av = G_av_start,
            G_skw = 50 # Start with a no skew
        )

        invisible(capture.output( # Suppress the details on converging SCEUA
            sceua_list <- rtop::sceua(growth_model,
                par0,
                T_per = T_per,
                G_per = G_per,
                years = years,
                t_int = t_int,
                transfer_function = transfer_function,
                d18Ow = d18Ow,
                Dsam = Dsam,
                Osam = Osam,
                t_maxtemp = t_maxtemp,
                parl,
                paru,
                maxn,
                kstop,
                pcento,
                ngs,
                iniflg = iniflg,
                peps = peps,
                implicit = function(pars){sum(pars[4]/2, pars[6]) < 1} # Make sure that the cumulative GR curve is not located below 0 (if G_av < - G_amp / 2)
                )
            ))

        par1 <- sceua_list[[1]] # Eport parameters of final model
        names(par1) <- names(par0)

        result <- growth_model(par1, T_per, G_per, years, t_int, transfer_function, d18Ow, Dsam, Osam, t_maxtemp, plot = FALSE, MC, D_err, O_err, return = "result") # Calculate the end result of the best fit

        if(plot == TRUE){
            fitplot <- fitplot + # Add the new model fit to the plot to track progress of the model
                ggplot2::geom_point(data = as.data.frame(result), ggplot2::aes(Dsam, Omod), colour = "red") +
                ggplot2::geom_line(data = as.data.frame(result), ggplot2::aes(Dsam, Omod), colour = "red")
            print(fitplot)
        }

        resultarray[, i + length(dat[1, ]), ] <- rbind(matrix(NA, nrow = i - 1, ncol = 9), result[, 3:11], matrix(NA, nrow = length(dat$D) - i - length(result[,1]) + 1, ncol = 9)) # Add results to result array
        parmat[, i] <- par1 # Add parameters to parameter array
    }

    # Provide names for all dimensions in the result array
    dimnames(resultarray) <- list(
        paste("sample", 1:length(resultarray[, 1, 3])),
        c(colnames(dat), paste("window", 1:length(dynwindow$x))),
        c("Modeled_d18O", "d18O_residuals", "Time_of_year", "Instantaneous_growth_rate", "Modeled temperature", "Modeled_d18O_SD", "Time_of_Year_SD", "Instantaneous_growth_rate_SD", "Modeled_temperature_SD")
    )

    colnames(parmat) <- paste("window", 1:length(parmat[1,]))
    return(list(resultarray, parmat))
}


#' LOO-CV, WAIC and Raw LPPD Calculations
#'
#' Implements LOO-CV, WAIC and Raw LPPD for the \code{\link{model_judgement}} function. Contains
#' the helper function "colVars" which calculates row-wise variances efficiently.
#'
#' @param stanfit A Stanfit object fitted on synthetic data.
#' @param current_model Name of the current model (Currently not in use.)
#' @param lik_name Name under which the log likelihoods have been saved in the models. Needs to be identical
#' across all Stanfit objects.
#' @param impute_inf A boolean which regulates if underflow values should be automatically imputed or not. If
#' \code{FALSE}, the models with such values will just be ignored. If \code{TRUE}, a report will be generated
#' on how many values were imputed and for which models.
#' @return A list with LOO-CV, WAIC and Raw LPPD calculations.
waic <- function(stanfit, current_model, lik_name, impute_inf){
    #http://kylehardman.com/BlogPosts/View/6 DIC code also from Gelman
    #Modified code of www.stat.columbia.edu/~gelman/research/unpublished/waic_stan.pdf
    #from gist.github.com/ihrke for underflow probs

    colVars <- function(a) {
        n <- dim(a)[[1]];
        c <- dim(a)[[2]];
        result <- (.colMeans(((a - matrix(.colMeans(a, n, c),
                                          nrow = n, ncol = c, byrow = TRUE)) ^ 2), n, c) * n / (n - 1))
        return(result)
    }

    log_lik <- rstan::extract(stanfit, lik_name)$log_lik

    if(impute_inf){
        lik_imp <- sum(is.infinite(log_lik))/length(log_lik)
        log_lik[is.infinite(log_lik)] <- mean(log_lik[!is.infinite(log_lik)])
    }
    dim(log_lik) <- if (length(dim(log_lik))==1) c(length(log_lik),1) else
        c(dim(log_lik)[1], prod(dim(log_lik)[2:length(dim(log_lik))]))
    S <- nrow(log_lik)
    n <- ncol(log_lik)
    #log pointwise
    lpd <- log(colMeans(exp(log_lik))) #only when posterior simulations in Stan are correctly made (and possible)

    #waic
    p_waic <- colVars(log_lik)
    elpd_waic <- lpd - p_waic
    waic <- -2*elpd_waic

    #loo
    loo_weights_raw <- 1/exp(log_lik-max(log_lik))

    if(impute_inf){
        loo_imp <- sum(is.infinite(loo_weights_raw))/length(loo_weights_raw)
        loo_weights_raw[is.infinite(loo_weights_raw)] <- mean(loo_weights_raw[!is.infinite(loo_weights_raw)])
    }

    loo_weights_normalized <- loo_weights_raw/matrix(colMeans(loo_weights_raw),nrow=S,ncol=n,byrow=TRUE)
    loo_weights_regularized <- pmin (loo_weights_normalized, sqrt(S))

    elpd_loo <- log(colMeans(exp(log_lik)*loo_weights_regularized)/colMeans(loo_weights_regularized))
    p_loo <- lpd - elpd_loo
    elpd_loo <- elpd_loo*-2
    pointwise <- cbind(waic,lpd,p_waic,elpd_waic,p_loo,elpd_loo)
    total <- colSums(pointwise)
    se <- sqrt(n*colVars(pointwise))
    ic <- list(waic=total["waic"], elpd_waic=total["elpd_waic"],
               p_waic=total["p_waic"], elpd_loo=total["elpd_loo"], p_loo=total["p_loo"],
               pointwise=pointwise, total=total, se=se)

    if(impute_inf){
        imps <- list('lik_imp' = lik_imp, 'loo_imp' = loo_imp)
        return(list('ic' = ic, 'imps' = imps))
    } else{
        return(list('ic' = ic))
    }
}


#' Function that propagates measurement uncertainty through model results
#'
#' Function to propagate combined errors on \code{x} (= \code{Dsam}) and
#' \code{y} (= \code{Osam}) on the modeled X (= \code{D}) and Y
#' (= \code{d18Oc}) values by means of direct projection of y–uncertainty
#' on \code{x} and then combine the errors on both in the \code{x} domain
#'
#' Note: projection y_err on x_err leads to large X errors on shallow
#' slopes due to numerical calculation of fist derivative.
#' @param x Vector of \code{x} values of input data
#' @param x_err Vector of uncertainties on \code{x} values
#' @param y Vector of \code{y} values of input data
#' @param y_err Vector of uncertainties on \code{y} values
#' @param X Vector of modeled \code{X} values on which the uncertainty is
#' to be projected
#' @param Y Matrix of modeled x and \code{Y} values
#' @param MC Number of Monte Carlo simulations to apply for error propagation
#' Default = 1000
#' @return A vector listing the standard deviations of propagated errors
#' propagated on all \code{X} values.
#' @examples
#' # Create dummy data for input
#' x <-  seq(1, 40, 1)
#' x_err <- rep(0.1, 40)
#' y <- sin((2 * pi * (seq(1, 40, 1) - 8 + 30 / 4)) / 30)
#' y_err <- rep(0.1, 40)
#' X <- seq(1.5, 39.5, 1)
#' Y <- cbind(seq(1, 39, 1), 0.9 * sin((2 * pi * (seq(1, 39, 1) - 9 +
#'     25 / 4)) / 25))
#' # Run function
#' result <- mc_err_proj(x, x_err, y, y_err, X, Y, 1000)
#' @export
mc_err_proj <- function(x,
    x_err,
    y,
    y_err,
    X,
    Y,
    MC = 1000){ # Function to propagate combined errors on x and y on the modeled X and Y values by means of local projection of y uncertainty on x and subsequent combination of uncertainties in X domain

    dYdX <- diff(Y[, 2]) / diff(X) # Create first derivative of Y by X
    dYdX[which(abs(dYdX) <= 10 ^ (floor(log(mean(abs(dYdX)), 10)) - 1))] <- sign(dYdX[which(abs(dYdX) <= 10 ^ (floor(log(mean(abs(dYdX)), 10)) - 1))]) * 10 ^ (floor(log(mean(abs(dYdX)), 10)) - 1) # Remove small absolute values (more than one order of magnitude smaller than the mean), preserving their sign
    dYdX <- append(dYdX, dYdX[length(dYdX)]) # repeat last value to increase the length of the vector to match X/Y (366 days)
    localslope <- dYdX[apply(abs(outer(x, X, FUN = "-")), 1, which.min)] # Find the local slope belonging to each sample position
    x_err_proj <- y_err / localslope # Project uncertainty on Y on x domain using slope
    x_err_comb <- sqrt(x_err_proj ^2 + x_err ^2) # Combine uncertainties on x and y in the X domain
    return(x_err_comb)
}