\name{extreme_deconvolution}
\alias{extreme_deconvolution}
\title{Density estimation using Gaussian mixtures in the presence of noisy, heterogeneous and incomplete data}
\description{We present a general algorithm to infer a d-dimensional distribution function given a set of heterogeneous, noisy observations or samples. This algorithm reconstructs the error-deconvolved or 'underlying' distribution function common to all samples, even when the individual samples have unique error and missing-data properties. The underlying distribution is modeled as a mixture of Gaussians, which is completely general. Model parameters are chosen to optimize a justified, scalar objective function: the logarithm of the probability of the data under the error-convolved model, where the error convolution is different for each data point. Optimization is performed by an Expectation Maximization (EM) algorithm, extended by a regularization technique and 'split-and-merge' procedure. These extensions mitigate problems with singularities and local maxima, which are often encountered when using the EM algorithm to estimate Gaussian density mixtures.}

\usage{
extreme_deconvolution(ydata,ycovar,
                      xamp,xmean,xcovar,
                      projection=NULL,
                      weight=NULL,
                      fixamp=NULL,fixmean=NULL,fixcovar=NULL,
                      tol=1.e-6,maxiter=1e9,w=0,logfile=NULL,
                      splitnmerge=0,maxsnm=FALSE,likeonly=FALSE,
                      logweight=FALSE)
}

\arguments{
  \item{ydata}{[ndata,dy] matrix of observed quantities}
  \item{ycovar}{[ndata,dy] / [ndata,dy,dy] / [dy,dy,ndata] matrix, list or 3D array
    of observational error covariances
    (if [ndata,dy] then the error correlations are assumed to vanish)}
  \item{xamp}{[ngauss] array of initial amplitudes (*not* [1,ngauss])}
  \item{xmean}{[ngauss,dx] matrix of initial means}
  \item{xcovar}{[ngauss,dx,dx] list of matrices of initial covariances}
  \item{projection}{[ndata,dy,dx] list of projection matrices}
  \item{weight}{[ndata] array of weights to be applied to the data points}
  \item{logweight}{(bool, default=False) if True, weight is actually
                   log(weight)}
  \item{fixamp}{(default=None) None, True/False, or list of bools}
  \item{fixmean}{(default=None) None, True/False, or list of bools}
  \item{fixcovar}{(default=None) None, True/False, or list of bools}
  \item{tol}{(double, default=1.e-6) tolerance for convergence}
  \item{maxiter}{(long, default= 10**9) maximum number of iterations to perform}
  \item{w}{(double, default=0.) covariance regularization parameter
            (of the conjugate prior)}
  \item{logfile}{basename for several logfiles (_c.log has output from
                 the c-routine; _loglike.log has the log likelihood path of
                 all the accepted routes, i.e. only parts which increase
                 the likelihood are included, during splitnmerge)}
  \item{splitnmerge}{(int, default=0) depth to go down the splitnmerge path}
  \item{maxsnm}{(Bool, default=False) use the maximum number of split 'n'
                 merge steps, K*(K-1)*(K-2)/2}
  \item{likeonly}{(Bool, default=False) only compute the total log
                   likelihood of the data}
}

\value{
  \item{avgloglikedata}{avgloglikedata after convergence}
  \item{xamp}{updated xamp}
  \item{xmean}{updated xmean}
  \item{xcovar}{updated xcovar}
}

\details{
  ...
}

\author{Jo Bovy, David W. Hogg, & Sam T. Roweis}

\references{
Inferring complete distribution functions from
noisy, heterogeneous and incomplete observations Jo Bovy, David
W. Hogg, & Sam T. Roweis, Submitted to AOAS (2009) [arXiv/0905.2979]
}

\examples{
library("ExtremeDeconvolution")
?extreme_deconvolution
ydata <- c(2.62434536,0.38824359,0.47182825,-0.07296862,1.86540763,-1.30153870,2.74481176,0.23879310,1.31903910,0.75062962,2.46210794,-1.06014071,0.67758280,0.61594565,2.13376944,-0.09989127,0.82757179,0.12214158,1.04221375,1.58281521,-0.10061918,2.14472371,1.90159072,1.50249434,1.90085595,0.31627214,0.87710977,0.06423057,0.73211192,1.53035547,0.30833925,0.60324647,0.31282730,0.15479436,0.32875387,0.98733540,-0.11731035,1.23441570,2.65980218,1.74204416,0.80816445,0.11237104,0.25284171,2.69245460,1.05080775,0.36300435,1.19091548,3.10025514,1.12015895,1.61720311,1.30017032,0.64775015,-0.14251820,0.65065728,0.79110577,1.58662319,1.83898341,1.93110208,1.28558733,1.88514116,0.24560206,2.25286816,1.51292982,0.70190717,1.48851815,0.92442829,2.13162939,2.51981682,3.18557541,-0.39649633,-0.44411380,0.49553414,1.16003707,1.87616892,1.31563495,-1.02220122,0.69379599,1.82797464,1.23009474,1.76201118,0.77767186,0.79924193,1.18656139,1.41005165,1.19829972,1.11900865,0.32933771,1.37756379,1.12182127,2.12948391,2.19891788,1.18515642,0.62471505,0.36126959,1.42349435,1.07734007,0.65614632,1.04359686,0.37999916,1.69803203,0.55287144,2.22450770,1.40349164,1.59357852,-0.09491185,1.16938243,1.74055645,0.04629940,0.73378149,1.03261455,-0.37311732,1.31515939,1.84616065,0.14048406,1.35054598,-0.31228341,0.96130449,-0.61577236,2.12141771,1.40890054,0.97538304,0.22483838,2.27375593,2.96710175,-0.85798186,2.23616403,2.62765075,1.33801170,-0.19926803,1.86334532,0.81907970,0.39607937,-0.23005814,1.55053750,1.79280687,0.37646927,1.52057634,-0.14434139,1.80186103,1.04656730,0.81343023,0.89825413,1.86888616,1.75041164,1.52946532,1.13770121,1.07782113,1.61838026,1.23249456,1.68255141,0.68988323,-1.43483776,2.03882460,3.18697965,1.44136444,0.89984477,0.86355526,0.88094581,1.01740941,-0.12201873,0.48290554,0.00297317,1.24879916,0.70335885,1.49521132,0.82529684,1.98633519,1.21353390,3.19069973,-0.89636092,0.35308331,1.90148689,3.52832571,0.75136522,1.04366899,0.77368576,2.33145711,0.71269214,1.68006984,0.68019840,-0.27255875,1.31354772,1.50318481,2.29322588,0.88955297,0.38263794,1.56276110,1.24073709,1.28066508,0.92688730,2.16033857,1.36949272,2.90465871,2.11105670,1.65904980,-0.62743834,1.60231928,1.42028220,1.81095167,2.04444209)
ydata <- matrix(ydata, length(ydata), 1)
N <- dim(ydata)[1]
ycovar <- ydata * 0 + 0.01
xamp <- c(0.5, 0.5)
xmean <- matrix(c(0.86447943, 0.67078879, 0.322681, 0.45087394), 2, 2)
xcovar <- list(matrix(c(0.03821028, 0.04014796, 0.04108113, 0.03173839), 2, 2),
               matrix(c(0.06219194,0.09738021, 0.04302473, 0.06778009), 2, 2))
projection <- list()
for (i in 1:N) projection[[i]] = matrix(c(i\%\%2,(i+1)\%\%2), 1, 2)
res <- extreme_deconvolution(ydata, ycovar, xamp, xmean, xcovar, projection=projection, logfile="ExDeconDemo")
print(res)
stopifnot((res$avgloglikedata - (-1.3114744655258121)) **2. < 10.**-8)
stopifnot((res$xmean[1,1]-2.30368235)**2. < 10.**-5)
stopifnot((res$xmean[1,2]-1.70701517)**2. < 10.**-5)
stopifnot((res$xmean[2,1]-1.08009397)**2. < 10.**-5)
stopifnot((res$xmean[2,2]-0.8888667)**2. < 10.**-5)
stopifnot((res$xcovar[[1]][1,1]-0.445645987259)**2. < 10.**-5.)
stopifnot((res$xamp[1]-0.11968415)**2. < 10.**-5)
stopifnot((res$xamp[2]-0.880315852981)**2. < 10.**-5)
}


library(openxlsx)
library(matlib)
library(Matrix)

this.dir<- dirname(parent.frame(2)$ofile)
parent.dir <- strsplit(this.dir,'scripts')[[1]]
setwd(parent.dir)

load_file <- 'efa/missing-loadings-f3-f12-sig-2.xlsx' #only statistically signficant loadings included
#To compute LSI using all loadings set load_file to 'efa/missing-loadings-f3-f12.csv'

f12_load_mat <- as.matrix(read.xlsx(load_file, sheet = 1, colNames = FALSE))
f11_load_mat <- as.matrix(read.xlsx(load_file, sheet = 2, colNames = FALSE))
f10_load_mat <- as.matrix(read.xlsx(load_file, sheet = 3, colNames = FALSE))
f9_load_mat <- as.matrix(read.xlsx(load_file, sheet = 4, colNames = FALSE))
f8_load_mat <- as.matrix(read.xlsx(load_file, sheet = 5, colNames = FALSE))
f7_load_mat <- as.matrix(read.xlsx(load_file, sheet = 6, colNames = FALSE))
f6_load_mat <- as.matrix(read.xlsx(load_file, sheet = 7, colNames = FALSE))
f5_load_mat <- as.matrix(read.xlsx(load_file, sheet = 8, colNames = FALSE))
f4_load_mat <- as.matrix(read.xlsx(load_file, sheet = 9, colNames = FALSE))
f3_load_mat <- as.matrix(read.xlsx(load_file, sheet = 10, colNames = FALSE))

eps <- 1.0e-8
w_fun <- function(x) {(x^2 + eps)^(10*x^2)}
"%^%" <- function(x, n) with(eigen(x), vectors %*% (values^n * t(vectors)))
lsi <- function(load_mat) {
  nfac <- ncol(load_mat)
  c_mat <- diag(t(load_mat) %*% load_mat)
  c_mat <- diag(c_mat)
  h_mat <- diag(load_mat %*% inv(c_mat) %*% t(load_mat))
  h_mat <- diag(h_mat)
  b_mat <- (h_mat %^% (-0.5)) %*% load_mat %*% (c_mat %^% (-0.5))
  ww <- sum(sapply( b_mat, w_fun,simplify="array"))*(1/64.0)*(1.0/nfac)
  ee <- ((1.0/nfac) + eps)^(10.0/nfac)
  ind <- (ww-ee)/(1.0-ee)
  return(ind)
}

print(paste0(ncol(f12_load_mat)," factors: LSI =", lsi(f12_load_mat)))
print(paste0(ncol(f11_load_mat)," factors: LSI =", lsi(f11_load_mat)))
print(paste0(ncol(f10_load_mat)," factors: LSI =", lsi(f10_load_mat)))
print(paste0(ncol(f9_load_mat)," factors: LSI =", lsi(f9_load_mat)))
print(paste0(ncol(f8_load_mat)," factors: LSI =", lsi(f8_load_mat)))
print(paste0(ncol(f7_load_mat)," factors: LSI =", lsi(f7_load_mat)))
print(paste0(ncol(f6_load_mat)," factors: LSI =", lsi(f6_load_mat)))
print(paste0(ncol(f5_load_mat)," factors: LSI =", lsi(f5_load_mat)))
print(paste0(ncol(f4_load_mat)," factors: LSI =", lsi(f4_load_mat)))
print(paste0(ncol(f3_load_mat)," factors: LSI =", lsi(f3_load_mat)))


# movies is pre-loaded in your workspace

# List all the graphical parameters
par()

# Specify the mfrow parameter

par(mfrow = c(2,1))
# Build two plots
plot(movies$votes,movies$rating)
hist(movies$votes)


ifelse(sqrt(9)<2,sqrt(9),0)

ifelse(sqrt(100)>9,sqrt(100),0)

x=12
if(is.numeric(x))
  y=x*2
y

z=6
if(z<0)
{
  y=z*3
}else
{
  y=z*5
}
y

x=15
y=3
if(is.numeric(x))
  if(is.numeric(y) & y!=0)
    z=x/y
z

x=letters[20]
if(is.numeric(x))
  print("Is numeric");
if(is.character(x))
  print("Is character");

z=='i'
if(z%in% letters)
{
  if(z=='a')
  {
     n=1
  }else if(z=='e')
   { n=2
  }else if(z=='b')
   { n=3
  }else if(z=='i')
   { n=4
  }else
  {
    n=5
  }
}
    n


model_conductance <- function (vonKarman = 0.42,
         heightWeatherMeasurements = 2.0,
         zm = 0.13,
         zh = 0.013,
         d = 0.67,
         plantHeight = 0.0,
         wind = 124000.0){
    #'- Name: Conductance -Version: 1.0, -Time step: 1
    #'- Description:
    #'            * Title: Conductance Model
    #'            * Author: Pierre Martre
    #'            * Reference: Modelling energy balance in the wheat crop model SiriusQuality2:
    #'            Evapotranspiration and canopy and soil temperature calculations
    #'
    #'            * Institution: INRA/LEPSE Montpellier
    #'            * Abstract: The boundary layer conductance is expressed as the wind speed profile above the
    #'            canopy and the canopy structure. The approach does not take into account buoyancy
    #'            effects.
    #'
    #'- inputs:
    #'            * name: vonKarman
    #'                          ** description : von Karman constant
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 1
    #'                          ** default : 0.42
    #'                          ** unit : dimensionless
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    #'                          ** inputtype : parameter
    #'                          ** parametercategory : constant
    #'            * name: heightWeatherMeasurements
    #'                          ** description : reference height of wind and humidity measurements
    #'                          ** parametercategory : soil
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 10
    #'                          ** default : 2
    #'                          ** unit : m
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    #'                          ** inputtype : parameter
    #'            * name: zm
    #'                          ** description : roughness length governing momentum transfer, FAO
    #'                          ** parametercategory : constant
    #'                          ** inputtype : parameter
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 1
    #'                          ** default : 0.13
    #'                          ** unit : m
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    #'            * name: zh
    #'                          ** description : roughness length governing transfer of heat and vapour, FAO
    #'                          ** parametercategory : constant
    #'                          ** inputtype : parameter
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 1
    #'                          ** default : 0.013
    #'                          ** unit : m
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    #'            * name: d
    #'                          ** description : corresponding to 2/3. This is multiplied to the crop heigth for calculating the zero plane displacement height, FAO
    #'                          ** inputtype : parameter
    #'                          ** parametercategory : constant
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.67
    #'                          ** min : 0
    #'                          ** max : 1
    #'                          ** unit : dimensionless
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547rl
    #'            * name: plantHeight
    #'                          ** description : plant Height
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0
    #'                          ** min : 0
    #'                          ** max : 1000
    #'                          ** unit : mm
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    #'                          ** inputtype : variable
    #'                          ** variablecategory : auxiliary
    #'            * name: wind
    #'                          ** description : wind
    #'                          ** variablecategory : auxiliary
    #'                          ** datatype : DOUBLE
    #'                          ** default : 124000
    #'                          ** min : 0
    #'                          ** max : 1000000
    #'                          ** unit : m/d
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    #'                          ** inputtype : variable
    #'- outputs:
    #'            * name: conductance
    #'                          ** description : the boundary layer conductance
    #'                          ** variablecategory : state
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0
    #'                          ** max : 10000
    #'                          ** unit : m/d
    #'                          ** uri : http://www1.clermont.inra.fr/siriusquality/?page_id=547
    h <- max(10.0, plantHeight) / 100.0
    conductance <- wind * vonKarman ^ 2 / (log((heightWeatherMeasurements - (d * h)) / (zm * h)) * log((heightWeatherMeasurements - (d * h)) / (zh * h)))
    return (list('conductance' = conductance))
}


# Replicate of calcpath subroutine

calcpath <- function(nohrs,slope,aspect,path){
  for (i in 1:nohrs){
    if (zenang[i] < pid2*0.998){
      print(paste('zenang of i = '),zenang[i])
      return()
    }
    path[i] = cos(zenang[i])*cos(slope)+sin(zenang[i])*sin(slope)*cos(aspect-sunazm[i])
    if (path[i] == 0){
      path[i] = 1.0e-11
      path[i] = 1.0/path[i]
    }
    if (path[i] <= 0.0 | path[i] > 1.0e10){
      if (path[i] <= 0.0){
        path[i] = -1.0
      }
    } else {
      # Sun below horizon
      path[i] = -1.0
    }
  }
}


#' 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.


library(pbdSBASE, quietly=TRUE)
library(pbdDMAT, quietly=TRUE)
petsc_init()

x <- matrix(1:30, 10)

dx <- as.dmat(x)
dsx <- as.dsmatrix(dx)
dsy <- t(dsx) %*% dsx

print(dsy)
print(dsy, all=TRUE, fmt="matlab")

petsc_finalize()
finalize()


# coding club 2018-05-22
#
# S. Van Hoey
#

library(tidyverse)

surveys <- read_csv("../data/20180222/20180222_surveys.csv") %>%
  filter(!is.na(weight),           # remove missing weight
         !is.na(hindfoot_length),  # remove missing hindfoot_length
         !is.na(sex))                # remove missing sex

# ----------------
# challenge 1
# ----------------

ggplot(surveys, aes(x = weight,
                    y = hindfoot_length,
                    color = sex)) +
  facet_wrap(facets = factor("year")) +
  geom_point(alpha = 0.5) +
  ylab("hindfoot length") +
  scale_x_log10() +
  scale_color_brewer(2) +
  theme_light()

# ----------------
# challenge 2
# ----------------

tidy_bevolking <- read_csv("../data/20180522/20180522_gent_groeiperwijk_tidy.csv")

# heatmap of the data
# https://learnr.wordpress.com/2010/01/26/ggplot2-quick-heatmap-plotting/
ggplot(tidy_bevolking, aes(factor(year), wijk, fill = growth)) +
  geom_tile(colour = "white") +
  theme_minimal(base_size = 14) +
  scale_fill_distiller(palette = "RdBu", type = "div",
                       direction = 1) +
  labs(x = "", y = "")

# ----------------
# challenge 3
# ----------------

# make sure the 'magick' package is installed

library(cowplot)
library(INBOtheme)

weight_scatter <- ggplot(surveys, aes(x = weight,
                    y = hindfoot_length)) +
  geom_point()

weight_density <- ggplot(surveys, aes(x = weight)) +
  geom_histogram(bins = 30, aes(y = ..density..),
                 alpha = 0.3, color = "gray") +
  geom_density(color = "black", size = 0.75)

# combine the plots with cowplot
my_plot <- plot_grid(weight_scatter, weight_density,
          labels = c("A", "B"), align = "h")

# add logo to the graph with cowplot
ggdraw() +
  draw_plot(my_plot) +
  draw_image("../images/coding_club_logo_1.png",
             x = 0.6, y = 0.6, scale = 0.2)


library(VIM)
library(mice)

library(MixAll)
library(rtkore)

library(arm)

library(jtools)
library(ggstance)
library(ggplot2)

library(interactions)

library(corrplot)

library(maps)
library(RColorBrewer)
library(tidyverse)

library(klaR)
library(naivebayes)
library(dplyr)
library(psych)

library(lme4)
library(lattice)

library(plotly)

library(readr)
library(dplyr)

library(leaflet)
library(htmltools)
library(ggplot2)
library(gridExtra)
library(lubridate)
library(VIM)
library(usmap)

library(gganimate)
library(ggrepel)
library(gapminder)

data <- read.csv("incarceration_trends.csv")
head(data)

columns <- c("year","fips","state","county_name", "total_pop",
             "total_jail_pop","total_jail_pop_rate","total_prison_pop","total_prison_pop_rate")
selected_data = data[columns]
colnames(selected_data) <- c("year","fips","state","county_name", "total_pop",
             "total_jail_pop","total_jail_pop_rate","total_prison_pop","total_prison_pop_rate")
str(selected_data)

selected_data[is.na(selected_data)] = 0
str(selected_data)

selected_data_year <- selected_data %>%
  group_by(year) %>%
  summarize(sum_total_pop= sum(as.numeric(total_pop)),
            sum_total_jail_pop= sum(as.numeric(total_jail_pop)),
            sum_total_jail_pop_rate = sum_total_jail_pop/sum_total_pop,
            sum_total_prison_pop= sum(as.numeric(total_prison_pop)),
            sum_total_prison_pop_rate = sum_total_prison_pop/sum_total_pop
            )
head(selected_data_year)

selected_data_year_type <- pivot_longer(
    selected_data_year,
    cols = c(sum_total_prison_pop_rate,sum_total_jail_pop_rate),
    names_to = "type"
)
head(selected_data_year_type)

selected_data_county <- selected_data %>%
  group_by(fips) %>%
  summarize(sum_total_pop= sum(as.numeric(total_pop)),
            sum_total_jail_pop= sum(as.numeric(total_jail_pop)),
            sum_total_jail_pop_rate = sum_total_jail_pop/sum_total_pop,
            sum_total_prison_pop= sum(as.numeric(total_prison_pop)),
            sum_total_prison_pop_rate = sum_total_prison_pop/sum_total_pop
            )
head(selected_data_county)

selected_data_state <- selected_data %>%
# filter(year == 2000) %>%
  group_by(state) %>%
  summarize(sum_total_pop= sum(as.numeric(total_pop)),
            sum_total_jail_pop= sum(as.numeric(total_jail_pop)),
            sum_total_jail_pop_rate = sum_total_jail_pop/sum_total_pop,
            sum_total_prison_pop= sum(as.numeric(total_prison_pop)),
            sum_total_prison_pop_rate = sum_total_prison_pop/sum_total_pop
            )
head(selected_data_state)

ggplot(selected_data_year, aes(x = year)) +
  geom_area(aes(y=sum_total_jail_pop_rate+sum_total_prison_pop_rate,
                fill="sum_total_prison_pop_rate")) +
  geom_area(aes(y=sum_total_jail_pop_rate,
                fill="sum_total_jail_pop_rate")) +
  theme_bw()+
  theme(
    legend.title = element_blank()
  )+
  scale_fill_brewer(palette = "Paired")+
  labs(title='Nationwide Jail and Prison Population Rate through Time',
       subtitle='Calculated from the sum of related population by counties ÷ the sum of population by counties.')+
  ylab('')+
  xlab('Year')

cors <- cor(selected_data_year)
png(height=2800, width=3200, pointsize=32, file="cors.png")
corrplot(cors, tl.col = 'black')
corrplot(cors, add=TRUE, type="lower", method="number",
         diag=FALSE, tl.pos="n", cl.pos="n")
corrplot(cors, add=TRUE, type="upper", tl.pos="n", tl.col = 'black')
dev.off()
corrplot(cors,tl.col = 'black')

qplot(sum_total_prison_pop, sum_total_jail_pop,
      data=selected_data_year,
      size=sum_total_pop, col=year, main="Jail-Prison") + theme_bw()

qplot(sum_total_jail_pop_rate, sum_total_prison_pop_rate,
      data=selected_data_state,
      size=sum_total_pop, col=sum_total_pop, main="Jail-Prison") + theme_bw()

plot_usmap (
  data = selected_data_county, values = "sum_total_jail_pop_rate", color = "red"
) +
scale_fill_continuous(
  low = "white", high = "red", name = "Population Rate", label = scales::comma
) +
labs(
  title = "Total Historical Jail Population Rate by Counties",
  subtitle = "Calculated by sum of total jail population overtime ÷ sum of total population overtime."
) +
theme(legend.position = "right")

plot_usmap (
  data = selected_data_county, values = "sum_total_prison_pop_rate", color = "red"
) +
scale_fill_continuous(
  low = "white", high = "red", name = "Population Rate", label = scales::comma
) +
labs(
  title = "Total Historical Prison Population Rate by Counties",
  subtitle = "Calculated by sum of total prison population overtime ÷ sum of total population overtime."
) +
theme(legend.position = "right")

plot_usmap(
  data = selected_data_state, values = "sum_total_jail_pop_rate", color = "red"
) +
scale_fill_continuous(
  low = "white", high = "red", name = "Population Rate", label = scales::comma
) +
labs(
  title = "Total Historical Jail Population Rate by States",
  subtitle = "Calculated by sum of total jail population by counties ÷ sum of total population by counties."
) +
theme(legend.position = "right")

plot_usmap(
  data = selected_data_state, values = "sum_total_prison_pop_rate", color = "red"
) +
scale_fill_continuous(
  low = "white", high = "red", name = "Population Rate", label = scales::comma
) +
labs(
  title = "Total Historical Prison Population Rate by States",
  subtitle = "Calculated by sum of total Prison population by counties ÷ sum of total population by counties."
) +
theme(legend.position = "right")


mul <- function(a, b) a * b


#' 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)
}


#' Get the sequence of tuning parameters and then make the call to the C code that implements the adaptive group lasso
#'
grouplassoRun <-
    function(data, index, weights, adaweights, family, thresh=0.0001, nlam=20, lambda=NULL, inner.iter=100, outer.iter=100, outer.thresh=0.0001, optim.step.size=0.8, reset=10, min.frac=0.05, verbose=FALSE) {
        lambdaSeek <- grouplassoLambdas(data=data, index=index, weights=weights, family=family, min.frac=min.frac, nlam=nlam, type="linear", adaweights=adaweights)
        if (is.null(lambda)) {
            lambda = lambdaSeek$lambda
        }

        X <- data$x
        y <- as.vector(data$y)
        n <- nrow(X)
        p <- ncol(X)

        #put the groups and adaweights in numerical order
        groups <- unique(index)
        ord.g = order(groups)
        groups = groups[ord.g]
        adaweights = adaweights[ord.g]

        #Reorder columns of X so that groups are contiguous
        ord <- order(index)
        index <- index[ord]
        X <- X[,ord]
        unOrd <- match(1:length(ord),ord)

        ## Coming up with other C++ info ##
        inner.iter = as.integer(inner.iter)
        outer.iter = as.integer(outer.iter)
        n = as.integer(nrow(X))
        p = as.integer(ncol(X))
        reset = as.integer(reset)

        num.groups = as.integer(length(groups))
        range.group.ind = rep(0,(num.groups+1))
        for(i in 1:num.groups){
            range.group.ind[i] <- min(which(index == groups[i])) - 1
        }
        range.group.ind[num.groups+1] <- ncol(X)
        range.group.ind = as.integer(range.group.ind)

        group.length <- as.integer(diff(range.group.ind))

        ## DONE SETTING UP C STUFF ##

        nlam = length(lambda)
        beta <- lambdaSeek$beta.unpen
        eta <- lambdaSeek$eta.unpen

        junk <- rcppLinNest(X = X,
                            y = y,
                            w = weights,
                            linkinv=family$linkinv,
                            devfun=family$dev.resids,
                            adaweights = adaweights,
                            nrow = n,
                            ncol = p,
                            numGroup = num.groups,
                            rangeGroupInd = range.group.ind,
                            groupLen = group.length,
                            lambda = lambda,
                            beta = beta,
                            innerIter = inner.iter,
                            outerIter = outer.iter,
                            thresh = thresh,
                            outerThresh = outer.thresh,
                            eta = eta,
                            gamma = optim.step.size,
                            betaIsZero = as.integer(rep(1, num.groups)),
                            reset = reset
        )

        return(list(beta=beta[unOrd,], lambda=lambda))
    }


library('DMwR')
dataset=read.csv("kdd_cup_test.csv")
dataset=dataset[1:100,]
outlier.scores<-lofactor(dataset,5)
jpeg("plot_densities.jpg");
plot(density(log(outlier.scores)))
outliers<-order(outlier.scores,decreasing=T)[1:5]
print(outliers)
n<-nrow(dataset)
labels<-1:n
labels[-outliers]<-"."
jpeg("biplot.jpg");
biplot(prcomp(dataset),cex=2.1,xlabs=labels)
pch<-rep(".",n)
pch[outliers]<-"+"
col<-rep("black",n)
col[outliers]<-"red"
png("scatter_matrix_1.jpg");
pairs(dataset[,1:7],pch=pch,col=col)
png("scatter_matrix_2.jpg");
pairs(dataset[,8:14],pch=pch,col=col)
png("scatter_matrix_3.jpg");
pairs(dataset[,15:21],pch=pch,col=col)
png("scatter_matrix_4.jpg");
pairs(dataset[,22:28],pch=pch,col=col)
png("scatter_matrix_5.jpg");
pairs(dataset[,22:28],pch=pch,col=col)
png("scatter_matrix_6.jpg");
pairs(dataset[,29:35],pch=pch,col=col)
png("scatter_matrix_7.jpg");
pairs(dataset[,36:41],pch=pch,col=col)
png("scatter_matrix.jpg");
pairs(dataset,pch=pch,col=col)
dev.off()


library(R.matlab)
get_cat_H_matrix <- function(cat_mat, cat_number){
	get_R <- function(cat_mat, cat_number) cat_mat$Cats[[cat_number]][[1]][[1]]
	get_J <- function(cat_mat, cat_number) cat_mat$Cats[[cat_number]][[1]][[2]]
	get_J_pinv <- function(cat_mat, cat_number) pracma::pinv(get_J(cat_mat,cat_number))
	scalings_for_muscles <- as.numeric(cat_mat$afl95) * (cat_mat$fmax * cat_mat$cosa95)
	muscle_scaling_diag <- diag(as.numeric(scalings_for_muscles))
	RFm <- get_R(cat_mat, 1) %*% muscle_scaling_diag
	H_matrix_cat <- pracma::pinv(get_J(cat_mat,cat_number))%>%t %*% RFm

	cat_muscle_names <- as.character(unlist(cat_mat$muscles))

	cat_wrench_names <- c("x","y","z","tx","ty","tz")
	colnames(H_matrix_cat) <- cat_muscle_names
	rownames(H_matrix_cat) <- cat_wrench_names
	bounds_tuple_of_numeric_cat <- rep(list(list(lower = 0, upper = 1)), 31)
return(list(H_matrix = H_matrix_cat%>%as.matrix, bounds_tuple_of_numeric = bounds_tuple_of_numeric_cat)
	)}

get_muscle_bounds_for_rhs_task <- function(constraint){
	n_muscles <- ncol(constraint$constr)
	muscle_identities <- diag(n_muscles)
	bound_solutions <- apply(muscle_identities, 1, function(muscle_focus){
	lb_solution <- simplex(muscle_focus, nulltask$constr,nulltask$rhs)
	ub_solution <- simplex(-1*muscle_focus, nulltask$constr,nulltask$rhs)
	return(data.table(lb=lb_solution$value,ub=ub_solution$value*-1.0))
	})
	return(rbindlist(bound_solutions))
}

get_sorted_csv_filenames_for_taskA <- function(){
	base <- "/Volumes/GoogleDrive/My\ Drive/st/task_A"
	raw_filenames <- dir(base)
	first_underscores <- sub("^[^_]*_", "", raw_filenames)
	second_underscores <- sub("^[^_]*_", "", first_underscores)
	numbers <- as.integer(gsub("_.*","",second_underscores))
	sorted_indices <- order(numbers)
	return(file.path(base,raw_filenames[sorted_indices]))
}

mat_path <- "~/Documents/GitHub/bc/stfeasibility/data/Sohn2013_hinlimb_models.mat"
cat_mat <- readMat(mat_path)
cat1 <- get_cat_H_matrix(cat_mat, 1)


 #this is the code used to generate the plot in Box 1
 library(plyr)
library(fields)

#Model parameters:
Npop=200 #initial pop size (not important)
Tmax=25 #max years
Tfishing=1 #when we start fishing
Trecovery=Tmax #stop fishing

   #Life history parameters

 	  	sizeatbirth=1
    Amax=4 #max age

 	amat=2 #age at maturation - this varies depending on the matration function chosen
	Linf=15 #max size
	fishingsize = .75*Linf #size-selective fishery threshold size
       Lmat=1 #length at maturity
       k=0.1 # Growth function param
 	  kappa=0.2 #natural mortality coef #params from Lorenzen et al. 2005
	 eps=0.18 #natural mortality exp

	  c=1   #egg prod coef
	  d=3 #egg &sperm prod exp
	  g = 7000000#sperm prod coef
	  v = 50 #mass-at-size coef
	  wexp = .3  #mass-at-size exp
	  h=1 #sperm fertilization effectiveness

#params that can be scalar or vector:
 natmort=0.2 #instantaneous natural mortality rate
 alpha = 0.0001#recruitment function param
 beta = 0.000000001#recruitment function param - strenght of density dependence

       q=0.8 #maturation ogive steepness (1 = very steep)

    Fvec = seq(0, 1.2, by = 0.3) #Fishing mortality vector

    #set up matrices
N=matrix(nrow=Amax, ncol=Tmax, data=0)

W=matrix(nrow=Amax, ncol=1, data=0)
L=matrix(nrow=Amax, ncol=1, data=0)
E=matrix(nrow=Amax, ncol=Tmax, data=0)
S=matrix(nrow=1, ncol=Tmax, data=0)
P=matrix(nrow=1, ncol=Tmax)
B=matrix(nrow=1, ncol=Tmax)
Y=matrix(nrow=1, ncol=Tmax)

 #set up vectors to hold age-specific variables
Pmale=rep(0, Amax) #probability of sex change
fert=rep(0, Tmax) #fert probability
pmat=rep(0, Amax) #maturation probability
eggs=rep(0, Amax) #fecundity
mu=rep(0, Amax) #natural mortality
select=rep(0, Amax) #fishery selectivity

N[1, 1]=Npop #initial population size
L[1]=sizeatbirth #larval size at recruitment
 sums=matrix(nrow=Tmax, ncol=length(Fvec))

for(Fish in 1:length(Fvec)) {

  mu_f=Fvec[Fish]

#First: Define the age-specific growth, mortality, and maturation function

  for (a in 1:(Amax-1)) {

#GROWTH

	L[a+1]= Linf*(1-exp(-k)) + (L[a])*exp(-k) #length at age
     W[a]= v*L[a]^wexp #weight at age
      #natural mortality

  mu[a]=0.05 #for simplicity, fixed at 0.05

     pmat[a]=ifelse( a < amat, 0, 1)

 #FISHING SELECTIVITY
     #JUST AFTER AGE AT MATURATION

    select[a] = ifelse( a < amat+1, 0, 1)  #knife edge gear selectivity right after maturation

 #FECUNDITY
      # eggs[a]= c*W[a]^2 #fecundity is a function of individual mass-at-age
       eggs[a]=100000

   } #end 1st age loop

plot(1:Amax, eggs, type="l")

#Now for these age specific rates, simulate pop dynamics through time
 for(t in 1:(Tmax-1)) {

E[,t]=N[, t]*pmat*eggs  #assuming spawning occurs between 1 t and the next

P[t]= sum(E[,t])  #assuming fertilization is 100%

 N[1,t+1]= alpha*P[t]/(1+beta*P[t]) #this is the N_0 class that is born and enters in the next time step..

     for (age in 1:(Amax-1)) {

   Fishing= if (Tfishing < t & Trecovery > t) {
   				Fishing = select[age]*mu_f } else {
   				Fishing = 0
                     }

 	N[age+1,t+1] = N[age, t]*exp(-mu[age]-Fishing)  #assuming mortality happens after spawning

 	B[t]=sum(N[,t]*W)

    } #end second age loop
        } #end t loop
      sums[,Fish]=apply(N, 2, sum) #calculate popsize (numbers) at each time
        } #end Fish loop

 #plot steady state for each value of fishing:
 quartz()
 matplot(1:Tmax, sums/100, type="l", lwd=2, xlab="Time", ylab="", ylim=c(0, 4500), lty = 1, col = c(1, "gray30", "gray50", "gray60", "gray80"), las=1)


smart.round <- function(x, digits = 0) {
# function to round while preserving the sum
  up <- 10 ^ digits
  x <- x * up
  y <- floor(x)
  indices <- tail(order(x-y), round(sum(x)) - sum(y))
  y[indices] <- y[indices] + 1
  y / up
}

projPops <- function(pops, names, nYears, decline, genYears, demogV, envtlV, envtlCorr, nSims) {
# function to build population simulations
# pops is a vector of starting population sizes for each population at t = 0
# names is a vector of the names of the populations
# nYears is how many years to run the simulation for
# decline in the simulated percent decline over 3 generations
# genYears is the number of years for three generations
# demogV is demographic stochasticity (variance in the growth rate among individuals)
# envtlV is the variance for environmental stochasticity
# envtlCorr is the correlation in environmental stochasticiity between populations
# nSims is how many independent simulations to run

  sim <- function(pops, names, nYears, r, demogV, envtlVC) {
  # function to run one iteration of the simulation
    Output <- round(matrix(pops))
    Rs <- matrix(NA, nrow = length(pops), ncol = nYears)

    for (t in 1:nYears) {
      # draw random deviates and growth rate for this year
      # note here we adjust for the expectation of a log normal distribution: exp((sigma^2)/2)
      DemR <- rnorm(n = length(pops), mean = 0, sd = ifelse(Output[, ncol(Output)] > 0, sqrt(demogV / Output[, ncol(Output)]), 0))
      EnvR <- mvrnorm(n = 1, mu = rep(0, length(pops)), Sigma = envtlVC)
      R <- rep(r, length(pops)) + DemR + EnvR
      Rs[, t] <- R
      Output <- cbind(Output, round(Output[, ncol(Output)] * exp(R)))
    }

    out <- as_tibble(cbind(as_tibble(names), as_tibble(Output))) %>% rename(POP = value)
    return(out)
  }

  # construct the variance covariance matrix for environmental stochasiticity
  envtlVC <- diag(envtlV, nrow = length(pops), ncol = length(pops))
  envtlVC[lower.tri(envtlVC)] <- envtlV * envtlCorr
  envtlVC[upper.tri(envtlVC)] <- envtlV * envtlCorr

  # get deterministic r for the given decline
  r <- log((1 - (decline / 100)) ^ (1 / genYears))

  # replicate simulations nSims times and return list
  return(replicate(n = nSims, expr = sim(pops = pops, names = names, nYears = nYears, r = r, demogV = demogV, envtlVC = envtlVC), simplify = FALSE))
}

getSurveys <- function(pops, survParams) {
# function to generate a survey sites and years
# pops is the population size data and the strata
# survParams is the survey parameters consisting of: f0, f0se, tranDayLine, tranDayArea, siteSize,
# budget, survIntens, method, monInter, strat, and monitRep

  # get area per site in ha
  areaSite <- as.numeric(survParams["survIntens"]) * as.numeric(survParams["siteSize"]) / 100
  # get line length per site in metres
  lengthSite <- (as.numeric(survParams["survIntens"]) * as.numeric(survParams["siteSize"]) * 100) / (2 / as.numeric(survParams["f0"]))
  # get days needed for each site
  daysPerSiteArea <- areaSite / as.numeric(survParams["tranDayArea"])
  daysPerSiteLine <- lengthSite / as.numeric(survParams["tranDayLine"])
  # get maximum number of sites per stratum
  pops_unn <- unnest(pops, cols = c(data)) %>% ungroup() %>% mutate(MAXSITES = AREA / as.numeric(survParams["siteSize"]))
  #get sums of weights for the different stratification strategies
  numStrat <- nrow(pops_unn)
  numStratUrb <- nrow(pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),])
  numStratNonUrb <- numStrat - numStratUrb
  areaAll <- sum(pops_unn[, "AREA"])
  areaUrb <- sum(pops_unn[which(pops_unn$LUCategory == "Urban Footprint"), "AREA"])
  areaNonUrb <- areaAll - areaUrb
  DensAll <- sum(pops_unn[, "KD_MEAN"])
  DensUrb <- sum(pops_unn[which(pops_unn$LUCategory == "Urban Footprint"), "KD_MEAN"])
  DensNonUrb <- DensAll - DensUrb
  InvDensAll <- sum(1 / pops_unn[, "KD_MEAN"])
  InvDensUrb <- sum(1 / pops_unn[which(pops_unn$LUCategory == "Urban Footprint"), "KD_MEAN"])
  InvDensNonUrb <- InvDensAll - InvDensUrb

  # allocate sampling effort per stratum based on stratification procedure
  if (survParams["strat"] == "equal") {
    if (survParams["method"] == "line") {
      pops_unn <- pops_unn %>% mutate(SITES = as.numeric(survParams["budget"]) / (numStrat * daysPerSiteLine)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      # re-allocate sites for strata with insufficient maximum number of sites
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn$SITES <- smart.round(pops_unn$SITES)
    } else if ((survParams["method"] == "area")) {
      pops_unn <- pops_unn %>% mutate(SITES = as.numeric(survParams["budget"]) / (numStratNonUrb * daysPerSiteLine + numStratUrb * daysPerSiteArea)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"SIZE"] <- areaSite
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"TYPE"] <- "area"
      # re-allocate sites for strata with insufficient maximum number of sites (line transects)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # re-allocate sites for strata with insufficient maximum number of sites (area searches)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn[which(pops_unn$TYPE == "line"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "line"), "SITES"]))
      pops_unn[which(pops_unn$TYPE == "area"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "area"), "SITES"]))
    } else {
      stop("incorrect survey type")
    }
  } else if (survParams["strat"] == "area") {
    if (survParams["method"] == "line") {
      pops_unn <- pops_unn %>% mutate(SITES = AREA * as.numeric(survParams["budget"]) / (areaAll * daysPerSiteLine)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      # re-allocate sites for strata with insufficient maximum number of sites
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn$SITES <- smart.round(pops_unn$SITES)
    } else if ((survParams["method"] == "area")) {
      pops_unn <- pops_unn %>% mutate(SITES = AREA * as.numeric(survParams["budget"]) / (areaNonUrb * daysPerSiteLine + areaUrb * daysPerSiteArea)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"SIZE"] <- areaSite
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"TYPE"] <- "area"
      # re-allocate sites for strata with insufficient maximum number of sites (line transects)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # re-allocate sites for strata with insufficient maximum number of sites (area searches)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn[which(pops_unn$TYPE == "line"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "line"), "SITES"]))
      pops_unn[which(pops_unn$TYPE == "area"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "area"), "SITES"]))
    } else {
      stop("incorrect survey type")
    }
  } else if (survParams["strat"] == "density") {
    if (survParams["method"] == "line") {
      pops_unn <- pops_unn %>% mutate(SITES = KD_MEAN * as.numeric(survParams["budget"]) / (DensAll * daysPerSiteLine)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      # re-allocate sites for strata with insufficient maximum number of sites
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn$SITES <- smart.round(pops_unn$SITES)
    } else if ((survParams["method"] == "area")) {
      pops_unn <- pops_unn %>% mutate(SITES = KD_MEAN * as.numeric(survParams["budget"]) / (DensNonUrb * daysPerSiteLine + DensUrb * daysPerSiteArea)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"SIZE"] <- areaSite
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"TYPE"] <- "area"
      # re-allocate sites for strata with insufficient maximum number of sites (line transects)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # re-allocate sites for strata with insufficient maximum number of sites (area searches)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn[which(pops_unn$TYPE == "line"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "line"), "SITES"]))
      pops_unn[which(pops_unn$TYPE == "area"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "area"), "SITES"]))
    } else {
      stop("incorrect survey type")
    }
  } else if (survParams["strat"] == "invdensity") {
    if (survParams["method"] == "line") {
      pops_unn <- pops_unn %>% mutate(SITES = (1 / KD_MEAN) * as.numeric(survParams["budget"]) / (InvDensAll * daysPerSiteLine)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      # re-allocate sites for strata with insufficient maximum number of sites
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES)) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn$SITES <- smart.round(pops_unn$SITES)
    } else if ((survParams["method"] == "area")) {
      pops_unn <- pops_unn %>% mutate(SITES = (1 / KD_MEAN) * as.numeric(survParams["budget"]) / (InvDensNonUrb * daysPerSiteLine + InvDensUrb * daysPerSiteArea)) %>% mutate(SIZE = lengthSite, TYPE = "line")
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"SIZE"] <- areaSite
      pops_unn[which(pops_unn$LUCategory == "Urban Footprint"),"TYPE"] <- "area"
      # re-allocate sites for strata with insufficient maximum number of sites (line transects)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "line"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # re-allocate sites for strata with insufficient maximum number of sites (area searches)
      if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
        toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
      } else {
        toAllocate <- 0
      }
      i <- 0 # note bails out if tries to allocate sites using more than 100 iterations
      while (toAllocate > 0 & i < 100) {
        allocSite <- toAllocate / length(which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"))
        pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES < pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] + allocSite
        pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] <- pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"]
        if (length(which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area")) > 0) {
          toAllocate <- sum(pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "SITES"] - pops_unn[which(pops_unn$SITES > pops_unn$MAXSITES & pops_unn$TYPE == "area"), "MAXSITES"])
        } else {
          toAllocate <- 0
        }
        i <- i + 1
        if (i == 100) {warning("more than 100 iterations to allocate survey effort")}
      }
      # round number of sites
      pops_unn[which(pops_unn$TYPE == "line"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "line"), "SITES"]))
      pops_unn[which(pops_unn$TYPE == "area"), "SITES"] <- smart.round(as.matrix(pops_unn[which(pops_unn$TYPE == "area"), "SITES"]))
    } else {
      stop("incorrect survey type")
    }
  } else {
    stop("incorrect stratification protocol")
  }

  # generate the survey sites and years in a #sites (M) * #time steps (T) matrix
  getSites <- function(x, monInter, monitRep) {
    maxInter <- ifelse(x["MAXSITES"] / x["SITES"] < 1, 1, floor(x["MAXSITES"] / x["SITES"]))[[1]]
    monInter <- min(monInter, maxInter)
    out <- matrix(NA, nrow = x["SITES"] * monInter, ncol = monitRep)
    for (i in 1:monitRep) {
      out[(((i %% monInter) * x["SITES"]) + 1):(((i %% monInter) * x["SITES"]) + x["SITES"]), i] <- i
    }
    return(out)
  }

  # generate survey
  pops_unn <- pops_unn %>% add_column(SURVEY = apply(dplyr::select(pops_unn, MAXSITES, SITES), MARGIN = 1, FUN = function (x, monInter, monitRep) {if(x["SITES"] > 0) {return(getSites(x, monInter, monitRep))} else {return(NULL)}},
                        monInter = as.numeric(survParams["monInter"]), monitRep = as.numeric(survParams["monitRep"])))

  return(pops_unn)
}

simData <- function(survey, popSim, f0, f0se, stripMissLow, stripMissMean, stripMissHigh) {
# function to generate survey data
# survey is the sites selection for a single monitoring parameter combination
# popSim is the replicates for a single population dynamics parameter combination
# f0 is the mean of the detection function
# f0se is the standard error of the detection function
# stripMissLow is the lower bound probability of missing a koala in a strip/area search
# stripMissMean is the mean probability of missing a koala in a strip/area search
# stripMissHigh is the upper bound probability of missing a koala in a strip/area search

  # calculate mean and se of f0 on the log scale (since we assume f0 is log-normally distributed)
  f0log <- log(f0) - (0.5 * (f0se^2))
  f0selog <- sqrt(log(exp(2 * log(f0) - log(f0se^2)) + 1))
  sim <- function(popSim, survey, f0, f0se, stripMissLow, stripMissMean, stripMissHigh) {
    # join simulations to survey for this iteration
    join <- left_join(x = survey, y = popSim, by = "POP")
    # output for densities
    out <- list()
    # output for counts
    out2 <- list()
    # output for perpendicular distances
    out3 <- list()
    # output for model data
    out4 <- list()

    # loop through strata
    for (i in 1:nrow(join)) {
      if (is.null(join$SURVEY[[i]])) {
        out[[i]] <- NA
        out2[[i]] <- NA
        out3[[i]] <- NA
        out4[[i]] <- NA
      } else {
        # get change in density
        out[[i]] <- join$SURVEY[[i]]
        out[[i]][which(!is.na(out[[i]]))] <- as.numeric(join[i, paste("V", join$SURVEY[[i]][which(!is.na(out[[i]]))], sep="")]) / as.numeric(join[i,"V1"])
        # get counts and perpendicular distances
        if (join[i, "TYPE"] == "line") {
          f0Rand <- f0  # exp(rnorm(n = 1, mean = f0log, sd = f0selog)) note here we assume f0 is the same among strata and known with certainty, but uncertainty can be introduced via the standard error
          sigma <- sqrt(2) / (f0Rand * sqrt(pi))
          M <- nrow(out[[i]])
          T <- ncol(out[[i]])
          J <- 4
          db <- qhnorm(p = c(0, 0.25, 0.5, 0.75, 0.99), sigma = sigma)
          # Half-normal, line transect
          g <- function(x, sig) exp(-x^2/(2*sig^2))
          f <- function(x, sig) exp(-x^2/(2*sig^2)) / sqrt(pi*sigma^2/2)
          cp <- u <- a <- numeric(J)
          L <-  1
          a[1] <- L * db[2]
          cp[1] <- integrate(g, db[1], db[2], sig = sigma)$value
          for(j in 2:J) {
            a[j] <-  db[j+1]  - sum(a[1:j])
            cp[j] <- integrate(g, db[j], db[j + 1], sig = sigma)$value
          }
          u <- a / sum(a)
          cp <- (cp / a) * u
          cp[j + 1] <- 1 - sum(cp)
          primPer <- matrix(as.integer(rep(1:T, M)), nrow = M, ncol = T, byrow = TRUE)
          # get true N assuming variation initial N among sites is Poisson distributed
          out2[[i]] <-  round(out[[i]] * matrix(rep(rpois(n = M, lambda = as.numeric(join[i, "KD_MEAN"]) * 2 * as.numeric(join[i, "SIZE"]) * db[J + 1] / 10000), T), nrow = M, ncol = T))
          # get observed counts in perpendicular distance bands
          out3[[i]] <- array(NA, c(M, J, T))
          for (k in 1:M) {
            if (!is.na(out2[[i]][k, 1])) {
              out3[[i]][k, 1:J, 1] <- rmultinom(1, out2[[i]][k, 1], cp)[1:J]
            }
            for(t in 1:(T - 1)) {
                if (!is.na(out2[[i]][k, t + 1])) {
                  out3[[i]][k, 1:J, t + 1] <- rmultinom(1, out2[[i]][k, t + 1], cp)[1:J]
                }
            }
          }
          out3[[i]] <- matrix(out3[[i]], M)
          # get model data
          out4[[i]] <- unmarkedFrameDSO(y = out3[[i]], numPrimary = T, primaryPeriod = primPer, dist.breaks = db,
                                  survey = "line", unitsIn = "m", tlength = rep(as.numeric(join[i, "SIZE"]), M))
        } else if ((join[i, "TYPE"]) == "area") {
          missRand <- stripMissMean # runif(1, stripMissLow, stripMissHigh)  note here we assume detection error is the same among strata and known with certainty,
                                      # but uncertainty can be introduced via the range of values
          M <- nrow(out[[i]])
          T <- ncol(out[[i]])
          J <- 1
          primPer <- matrix(as.integer(rep(1:T, M)), nrow = M, ncol = T, byrow = TRUE)
          # get observed counts assuming variation in initial N among sites is Poisson distributed
          out2[[i]] <- matrix(suppressWarnings(rbinom(n = M * T, size = round(out[[i]] * matrix(rep(rpois(n = M, lambda = as.numeric(join[i, "SIZE"]) * as.numeric(join[i, "KD_MEAN"])), T), nrow = M, ncol = T)),
                          prob = 1 - missRand)), nrow = M, ncol = T)
          # get perpendicular distances (not needed in this case)
          out3[[i]] <- NA
          # get model data
          out4[[i]] <- unmarkedFramePCO(y = out2[[i]], numPrimary = T, primaryPeriod = primPer)
        }
      }
    }

    out <- join %>% add_column(MODDAT = out4) %>% dplyr::select(ID, SIZE, MODDAT)

    return(out)
  }

  # get simulated data
  data <- lapply(popSim, FUN = sim, survey = survey, f0 = f0, f0se = f0se, stripMissLow = stripMissLow, stripMissMean = stripMissMean, stripMissHigh = stripMissHigh)

  return(data)
}

simDataAdapt <- function(survey, popSim, f0, f0se, stripMissLow, stripMissMean, stripMissHigh, yearAdapt) {
# function to generate survey data but where an adaptive strategy is adopted whereby
# if a koala is not seen for yearAdapt years then the site is abandoned and another site is chosen
# yearAdapt must be > 1
# survey is the sites selection for a single monitoring parameter combination
# popSim is the replicates for a single population dynamics parameter combination
# f0 is the mean of the detection function
# f0se is the standard error of the detection function
# stripMissLow is the lower bound probability of missing a koala in a strip/area search
# stripMissMean is the mean probability of missing a koala in a strip/area search
# stripMissHigh is the upper bound probability of missing a koala in a strip/area search
# yearAdapt is how frequently the monitoring strategy is adapted

  # calculate mean and se of f0 on the log scale (since we assume f0 is log-normally distributed)
  f0log <- log(f0) - (0.5 * (f0se^2))
  f0selog <- sqrt(log(exp(2 * log(f0) - log(f0se^2)) + 1))
  sim <- function(popSim, survey, f0, f0se, stripMissLow, stripMissMean, stripMissHigh, yearAdapt) {
    # join simulations to survey for this iteration
    join <- left_join(x = survey, y = popSim, by = "POP")
    # output for densities
    out <- list()
    # output for counts
    out2 <- list()
    # output for perpendicular distances
    out3 <- list()
    # output for model data
    out4 <- list()

    # loop through strata
    for (i in 1:nrow(join)) {
      if (is.null(join$SURVEY[[i]])) {
        out[[i]] <- NA
        out2[[i]] <- NA
        out3[[i]] <- NA
        out4[[i]] <- NA
      } else {
        # get change in density
        out[[i]] <- join$SURVEY[[i]]
        out[[i]][which(!is.na(out[[i]]))] <- as.numeric(join[i, paste("V", join$SURVEY[[i]][which(!is.na(out[[i]]))], sep="")]) / as.numeric(join[i,"V1"])
        # get counts and perpendicular distances
        if (join[i, "TYPE"] == "line") {
          f0Rand <- f0  # exp(rnorm(n = 1, mean = f0log, sd = f0selog)) note here we assume f0 is the same among strata and known with certainty, but uncertainty can be introduced via the standard error
          sigma <- sqrt(2) / (f0Rand * sqrt(pi))
          M <- nrow(out[[i]])
          T <- ncol(out[[i]])
          J <- 4
          db <- qhnorm(p = c(0, 0.25, 0.5, 0.75, 0.99), sigma = sigma)
          # Half-normal, line transect
          g <- function(x, sig) exp(-x^2/(2*sig^2))
          f <- function(x, sig) exp(-x^2/(2*sig^2)) / sqrt(pi*sigma^2/2)
          cp <- u <- a <- numeric(J)
          L <-  1
          a[1] <- L * db[2]
          cp[1] <- integrate(g, db[1], db[2], sig = sigma)$value
          for(j in 2:J) {
            a[j] <-  db[j+1]  - sum(a[1:j])
            cp[j] <- integrate(g, db[j], db[j + 1], sig = sigma)$value
          }
          u <- a / sum(a)
          cp <- (cp / a) * u
          cp[j + 1] <- 1 - sum(cp)
          primPer <- matrix(as.integer(rep(1:T, M)), nrow = M, ncol = T, byrow = TRUE)
          # get true N assuming variation initial N among sites is Poisson distributed
          out2[[i]] <-  round(out[[i]] * matrix(rep(rpois(n = M, lambda = as.numeric(join[i, "KD_MEAN"]) * 2 * as.numeric(join[i, "SIZE"]) * db[J + 1] / 10000), T), nrow = M, ncol = T))
          # get observed counts in perpendicular distance bands
          out3[[i]] <- array(NA, c(M, J, T))
          k <- 1
          while (k <= M) {
            if (!is.na(out2[[i]][k, 1])) {
              out3[[i]][k, 1:J, 1] <- rmultinom(1, out2[[i]][k, 1], cp)[1:J]
            }
            for(t in 1:(T - 1)) {
                if (!is.na(out2[[i]][k, t + 1])) {
                  out3[[i]][k, 1:J, t + 1] <- rmultinom(1, out2[[i]][k, t + 1], cp)[1:J]
                  # count how many surveys in a row no koalas have been seen
                  numZero <- 0
                  for (z in t:1) {
                    if (!any(is.na(out3[[i]][k, 1:J, z + 1]))) {
                      if (sum(out3[[i]][k, 1:J, z + 1]) == 0) {
                        numZero <- numZero + 1
                      } else {
                        break
                      }
                    }
                  }
                  # adapt strategy if necessary
                  if ((numZero >= yearAdapt) & (t < (T - 1))) {
                      # add a new site
                      M <- M + 1
                      out[[i]] <- rbind(out[[i]], out[[i]][k, ])
                      out2[[i]] <- rbind(out2[[i]],
                                    round(out[[i]][nrow(out[[i]]), ] * matrix(rep(rpois(n = 1, lambda = as.numeric(join[i, "KD_MEAN"]) * 2 * as.numeric(join[i, "SIZE"]) * db[J + 1] / 10000), T),
                                    nrow = 1, ncol = T)))
                     # add new slice to the observed perpendicular distance data
                     out3[[i]] <- abind(out3[[i]], array(NA, c(1, J, T)), along = 1)
                     # set times that will now not be surveyed to NA
                     out2[[i]][k, (t + 2):T] <- NA
                     out2[[i]][M, 1:(t + 1)] <- NA
                     # recreate primary period matrix
                     primPer <- matrix(as.integer(rep(1:T, M)), nrow = M, ncol = T, byrow = TRUE)
                  }
                }
            }
            k <- k + 1
          }
          out3[[i]] <- matrix(out3[[i]], M)
          # get model data
          out4[[i]] <- unmarkedFrameDSO(y = out3[[i]], numPrimary = T, primaryPeriod = primPer, dist.breaks = db,
                                  survey = "line", unitsIn = "m", tlength = rep(as.numeric(join[i, "SIZE"]), M))
        } else if ((join[i, "TYPE"]) == "area") {
          missRand <- stripMissMean # runif(1, stripMissLow, stripMissHigh)  note here we assume detection error is the same among strata and known with certainty,
                                      # but uncertainty can be introduced via the range of values
          M <- nrow(out[[i]])
          T <- ncol(out[[i]])
          J <- 1
          primPer <- matrix(as.integer(rep(1:T, M)), nrow = M, ncol = T, byrow = TRUE)
          # get observed counts assuming variation in initial N among sites is Poisson distributed
          out2[[i]] <- matrix(suppressWarnings(rbinom(n = M * T, size = round(out[[i]] * matrix(rep(rpois(n = M, lambda = as.numeric(join[i, "SIZE"]) * as.numeric(join[i, "KD_MEAN"])), T), nrow = M, ncol = T)),
                          prob = 1 - missRand)), nrow = M, ncol = T)
          # do adaptation
          k <- 1
          while (k <= M) {
            for(t in 1:(T - 1)) {
                if (!is.na(out2[[i]][k, t + 1])) {
                  # count how many surveys in a row no koalas have been seen
                  numZero <- 0
                  for (z in t:1) {
                    if (!is.na((out2[[i]][k, z + 1]))) {
                      if(out2[[i]][k, z + 1] == 0) {
                        numZero <- numZero + 1
                      } else {
                        break
                      }
                    }
                  }
                  # adapt strategy if necessary
                  if ((numZero >= yearAdapt) & (t < (T - 1))) {
                      # add a new site
                      M <- M + 1
                      out[[i]] <- rbind(out[[i]], out[[i]][k, ])
                      out2[[i]] <- rbind(out2[[i]], matrix(suppressWarnings(rbinom(n = 1 * T, size = round(out[[i]][nrow(out[[i]]), ] * matrix(rep(rpois(n = 1, lambda = as.numeric(join[i, "SIZE"]) *
                                      as.numeric(join[i, "KD_MEAN"])), T), nrow = 1, ncol = T)), prob = 1 - missRand)), nrow = 1, ncol = T))
                     # set times that will now not be surveyed to NA
                     out2[[i]][k, (t + 2):T] <- NA
                     out2[[i]][M, 1:(t + 1)] <- NA
                     # recreate primary period matrix
                     primPer <- matrix(as.integer(rep(1:T, M)), nrow = M, ncol = T, byrow = TRUE)
                  }
                }
            }
            k <- k + 1
          }
          # get perpendicular distances (not needed in this case)
          out3[[i]] <- NA
          # get model data
          out4[[i]] <- unmarkedFramePCO(y = out2[[i]], numPrimary = T, primaryPeriod = primPer)
        }
      }
    }

    out <- join %>% add_column(MODDAT = out4) %>% dplyr::select(ID, SIZE, MODDAT)

    return(out)
  }

  # get simulated data
  data <- lapply(popSim, FUN = sim, survey = survey, f0 = f0, f0se = f0se, stripMissLow = stripMissLow, stripMissMean = stripMissMean, stripMissHigh = stripMissHigh, yearAdapt)

  return(data)
}

fitModels <- function(survData, pops, expR) {
# survData is a list of survey data for all replicates for a given population scenario and given monitoring strategy

  #functions
  fitOneModel <- function (data, expR) {
  # function to fit one model in one stratum
    out <- tryCatch(
    {
      if (!(class(data) == "unmarkedFrameDSO" | class(data) == "unmarkedFramePCO")) {
        out_nerr <- NA
        type <- NA
      } else {
        if (class(data) == "unmarkedFrameDSO") {
          out_nerr <- distsampOpen(lambdaformula = ~1, gammaformula = ~1, omegaformula = ~1, pformula = ~1, data = data, K = max(getY(data), na.rm = T) * 2 + 5, keyfun = "halfnorm",
          output = "density", unitsOut = "ha", dynamics = "trend", method = "BFGS", starts = c(log(0.15), expR, log(sqrt(2) / (0.0285 * sqrt(pi)))))
          type <- slot(out_nerr, "fitType") #distsampOpen
        } else if (class(data) == "unmarkedFramePCO") {
          out_nerr <- pcountOpen(lambdaformula = ~1, gammaformula = ~1, omegaformula = ~1, pformula = ~1, data = data, mixture = "P", K = max(getY(data), na.rm = T) * 2 + 5,
                  dynamics = "trend", method = "BFGS", starts = c(log(0.15 * 30), expR, log(0.84 / (1-0.84))))
          type <- slot(out_nerr, "fitType") #"pcountOpen"
        } else {
          stop("incorrect survey type")
        }
      }
      # create output
      if (class(out_nerr) == "logical") {
        out_nerr <- matrix(NA, nrow = 2, ncol = 3)
        dimnames(out_nerr)[[1]] <- c("mu", "sigma")
      } else {
        out_nerr <- rbind(coef(out_nerr), SE(out_nerr))
        dimnames(out_nerr)[[1]] <- c("mu", "sigma")
      }
      #write("OK\n", "errorlogfile.txt", append = TRUE)
      list(type, out_nerr)
    },
    error = function(cond) {
      out_err <- matrix(NA, nrow = 2, ncol = 3)
      dimnames(out_err)[[1]] <- c("mu", "sigma")
      type <- NA
      # Add error message to the error log file
      #write(toString(cond), "errorlogfile.txt", append = TRUE)
      return(list(type, out_err))
    },
    warning = function(cond) {
      out_err <- matrix(NA, nrow = 2, ncol = 3)
      dimnames(out_err)[[1]] <- c("mu", "sigma")
      type <- NA
      # Add error message to the error log file
      #write(toString(cond), "errorlogfile.txt", append = TRUE)
      return(list(type, out_err))
    })
    return(out)
  }

  aggFits <- function(fits) {
  # function to get weighted average trends estimates, standard errors and whether trend detected at p = 0.05 level
    extractTrend <- function(fit) {
      if (!is.na(fit[[1]])) {
        if (is.finite(fit[[2]][1, 2])) {
          return(fit[[2]][1, 2])
        } else {
          return(NA)
          #write("Trend estimate not finite, setting to NA\n", "errorlogfile.txt", append = TRUE)
        }
      } else {
        return(NA)
      }
    }
    extractSE <- function(fit) {
      if (!is.na(fit[[1]])) {
        if (is.finite(fit[[2]][2, 2])) {
          return(fit[[2]][2, 2])
        } else {
          return(NA)
          #write("Standard error of trend estimate not finite, setting to NA\n", "errorlogfile.txt", append = TRUE)
        }
      } else {
        return(NA)
      }
    }
    extractDens <- function(fit, size) {
      if (!is.na(fit[[1]])) {
        if (fit[[1]] == "distsampOpen") {
          out <- exp(fit[[2]][1, 1] + (0.5 * fit[[2]][2, 1] ^ 2))
          if (is.finite(out)) {
            return(out)
          } else {
            return(NA)
            #write("Density estimate not finite, setting to NA\n", "errorlogfile.txt", append = TRUE)
          }
        } else if (fit[[1]] == "pcountOpen") {
          out <- exp(fit[[2]][1, 1] + (0.5 * fit[[2]][2, 1] ^ 2)) / size
          if (is.finite(out)) {
            return(out)
          } else {
            return(NA)
            #write("Density estimate not finite, setting to NA\n", "errorlogfile.txt", append = TRUE)
          }
        } else {
          stop("wrong model type")
        }
      } else {
        return(NA)
      }
    }

    result <- fits %>% mutate(TREND = map_dbl(.x = FITS, .f = extractTrend), SE = map_dbl(.x = FITS, .f = extractSE), DENS = map2_dbl(.x = FITS, .y = SIZE, .f = extractDens)) %>%
                        mutate(DAREA = DENS * AREA)

    if (all(is.na(result$TREND))) {
      out <- list()
      out$r_est <- NA
      out$r_decline <- 0
    } else {
      weights <- result$DAREA / sum(result$DAREA, na.rm = TRUE)
      out <- list()
      outtrend <-  sum(result$TREND * weights, na.rm = TRUE) # / sum(weights, na.rm = TRUE)
      if (any(is.infinite(out$trend))) {
        out$trend[which(is.infinite(out$trend))] <- NA
        #write("Trend not finite, setting to NA\n", "errorlogfile.txt", append = TRUE)
      }
      outse <- sqrt(sum(result$SE ^ 2 * weights ^ 2, na.rm = TRUE)) # / (sum(weights, na.rm = TRUE) ^ 2))
      outp <- pnorm(outtrend, mean = 0, sd = outse)
      out$r_est <- outtrend
      #out$r_se <- outse
      #out$p <- outp
      out$r_decline <- if (outp < 0.05) {1} else {0}
    }

    return(out)
  }

  fitStrataModel <- function (stratData, pops_unn, expR) {
  # function to fit a model to each population

    # fit models and return list
    fits <- lapply(stratData$MODDAT, FUN = fitOneModel, expR = expR)

    # fit models and add to data frame
    stratData <- stratData %>% dplyr::select(ID, SIZE) %>% add_column(FITS = fits)

    unnestData <- left_join(x = pops_unn, y = stratData, by = "ID")
    nestData <- unnestData %>% group_by(POP) %>% nest()
    numPops <- pops_unn %>% group_by(POP) %>% nest() %>% nrow()

    # get whether a decline is detected for the region
    out1 <- unlist(lapply(list(unnestData), FUN = aggFits))
    out1a <- out1[1] # trend
    out1b <- out1[2] # decline detected

    # get the number of declines detected for each population
    out2 <- unlist(lapply(nestData$data, FUN = aggFits)) # which populations detected decline
    out2b <- out2[seq(2, length(out2), 2)] # select only the detections (not trends)
    out2a <- sum(out2b, na.rm = TRUE) # number of populations with declines

    output <- c(out1a, out1b, out2a, out2b)
    names(output) <- c("r_est", "region_decline", "pops_decline", sprintf("pop%s", seq(1:numPops)))

    return(output)
  }

  # unnest populations data
  pops_unn <- unnest(pops, cols = c(data)) %>% ungroup() %>% dplyr::select(POP, ID, AREA)

  # fit models across all replicates
  out <- lapply(survData, FUN = fitStrataModel, pops_unn = pops_unn, expR = expR)

  numPops <- pops_unn %>% group_by(POP) %>% nest() %>% nrow()

  # prepare output
  out <- matrix(unlist(out), nrow = 3 + numPops, ncol = length(survData))
  dimnames(out) <- list(c("r_est", "region_decline", "pops_decline", sprintf("pop%s", seq(1:numPops))), NULL)
  trendMean <- mean(out["r_est",], na.rm = TRUE)
  trendSE <- sd(out["r_est",], na.rm = TRUE) / (length(out["r_est", which(!is.na(out["r_est", ]))]) - 1)
  trendBias <- trendMean - expR
  regionDet <- mean(out["region_decline", ], na.rm = TRUE)
  regionDetLow <- binom.confint(sum(out["region_decline", ], na.rm = TRUE), length(out["region_decline", which(!is.na(out["region_decline", ]))]), method = "exact")$lower
  regionDetHigh <- binom.confint(sum(out["region_decline", ], na.rm = TRUE), length(out["region_decline", which(!is.na(out["region_decline", ]))]), method = "exact")$upper
  popsNumDet <- mean(out["pops_decline",], na.rm = TRUE)
  popsDet <- apply(out[sprintf("pop%s", seq(1:numPops)), ], MARGIN = 1, FUN = mean, na.rm =TRUE)
  output <- c(trendMean, trendSE, expR, trendBias, regionDet, regionDetLow, regionDetHigh, popsNumDet, popsDet)
  names(output) <- c("r_est", "r_se", "r_exp", "r_bias", "region_decline_mean", "region_decline_low", "region_decline_high", "num_pop_decline", sprintf("pop%s_decline", seq(1:numPops)))

  return(output)
}


# 4. faza: Analiza podatkov

source("lib/libraries.r", encoding="UTF-8")

## Predikcija

podatki.napoved <- osnova %>% group_by(Leto) %>% mutate("bdp" = mean(BDPpp), Leto=Leto) %>% select(3,6) %>% unique()
odst <- lm(data=podatki.napoved, formula = bdp ~ I(Leto) + I(Leto^2))
vm <- data.frame(Leto=seq(2020, 2022, 1))
napoved <- mutate(vm, "bdp"=predict(odst, vm)) %>% data.frame() %>% group_by(Leto)
podatki.napoved <- data.frame(rbind(podatki.napoved, napoved), stringsAsFactors=FALSE) %>% group_by(Leto)

graf.napoved <- ggplot(podatki.napoved, aes(x=Leto, y=bdp)) +
  geom_point() +
  geom_smooth(method=lm, color="black", size=0.1, formula = y ~ I(x) + I(x^2)) +
  scale_x_continuous(breaks = c(2010, 2013, 2016, 2019, 2022))  +
  ylab("BDP per capita ($)") +
  ggtitle("Napoved gibanja povprečnega BDP per capita v ZDA")

## Metoda voditeljev

podatki.vod <- osnova_2019 %>% ungroup %>% select(2)
koncni <- podatki.vod %>% as.matrix() %>% scale()
vdtl <- kmeans(koncni, 3, nstart = 1e3)
razdelitev <- data.frame(Zvezna_drzava=osnova_2019$Zvezna_drzava, Skupina=factor(vdtl$cluster))

map3 <- tm_shape(merge(zemljevid, razdelitev, by.x="STATE_NAME", by.y="Zvezna_drzava"), simplify = 0.2) +
  tm_polygons("Skupina", palette = "Set2") +
  tm_layout(legend.outside=TRUE) + tm_text("STATE_ABBR", size=0.3)


require(rms)
y <- 1 : 100
units(y) <- 'year'
S <- Surv(y)
g <- rep('a', 100)
g[seq(1, 100, by=5)] <- 'b'

tapply(1:100, g, function(x) sum(x >= 50))  # a=41 b=10

tapply(1:100, g, function(x) length(x) / sum(x))
fs <- survfit(S ~ g)
fs
i <- fs$time %in% c(46, 50)
fs$n.risk[i]   # a=41 b=11
z <- qnorm(.975)
sur <- fs$surv[i]
seS <- fs$std.err[i] * sur   # se on S(t) scale instead of log S(t)
with(fs, cbind(n.risk[i], surv[i], lower[i], upper[i], std.err[i], seS, sur - z * seS, sur + z * seS))
# Last 2 columns not supposed to agree with summary.survfit since summary use log S(t)
# as basis for CLs

# summary(f, times=1:100)$time

s <- summary(fs, times=50)
s    # a=41 b=10

# Manually compute lower and upper half CL at t=50
mean(sur) + c(-1, 1) * 0.5 * z * sqrt(seS[1]^2 +seS[2]^2) # .3774 .6224
# Compare to width of CL for smallest stratum above       # ,2898 .7191

f <- npsurv(S ~ g)

survplot(f)
survplot(f, conf='diffbands')
survdiffplot(f)   # modern art

survplotp(f)
survplotp(f, aehaz=TRUE)
survplotp(f, times=c(50,60))
survplotp(f, aehaz=TRUE, times=c(5,60))

h <- function(y) 1 - y
survplotp(f, fun=h, ylab='Cumulative Incidence')
survplotp(f, fun=h, aehaz=TRUE, times=c(5, 60))