##### Code for estimating country specific effective reproduction number with different methods  ####
#####  parts of the code reutilizes already partly written code by othe researcrhes, the sources for these are:
#####https://www.datacamp.com/community/tutorials/replicating-in-r-covid19
##### https://www.r-bloggers.com/effective-reproduction-number-estimation/
#####
######
######
#####

### instructions to run the code  ####
### 1. set the homedirectory in your computer where to save the results
##   ADD HOMEDIRECTORY, for example "c:/DATA/COVID/RESULTS_CALCULATION"
Homedir<-"ADD DIRECTORY TO SAVE RESULTS HERE"

### 2. Uncomment the package installations if they are not installed

# install.packages("purrr")
# install.packages("knitr")
#install.packages("kableExtra")
# install.packages("viridis")
# install.packages("accelerometry")
# install.packages("tidyverse")
# install.packages("R0")
#install.packages("IRdisplay")
#install.packages("smoother")
#install.packages("HDInterval")

####### 3.   Specify the country the analysis should be performed for   #####
Country_Name<-"France"

##### 4.  Specify the date the analysis should start
starting_date <- "2020-03-22"

### 5. Execute all the code, outputs will be saved in the homedirectory as a csv file with name "Reproduction_number_estimates.csv"

#### Possible errors    #####
##### some error messages are displayed during the analysis, these are no real errors but mostly messages from different
##### estimation tools on assumptions taken or if some parts of the analysis takes longer than expected
##### only if the analysis stops and it not finished there can be an error. This error is due to the country specific data
##### this can either depend tht the starting date is taken too early, some methods do not work in this situation. Other possibility
#### is that in some countries zero cases can be reported for a longer timeperiod.
############################

suppressPackageStartupMessages(library(purrr))
suppressPackageStartupMessages(library(knitr))
suppressPackageStartupMessages(library(kableExtra))
suppressPackageStartupMessages(library(viridis))
suppressPackageStartupMessages(library(accelerometry))
suppressPackageStartupMessages(library(tidyverse))
suppressPackageStartupMessages(library(IRdisplay))
suppressPackageStartupMessages(library(smoother))
suppressPackageStartupMessages(library(HDInterval))

starting_date <- "2020-03-22"
Country_Name<-"France"

### Set up distribution on generation time
GT_pmf <- structure( c(0, 0.1, 0.1, 0.2, 0.2, 0.2, 0.1, 0.1), names=0:7)
GT_obj <- R0::generation.time("empirical", val=GT_pmf)

GT_obj<-R0::generation.time("gamma", c(6.6, 1.5))

Homedir<-"H:/COVID/R0_technical_paper/Graphs_test"

## Read in ECDC data   ##
data_ECDC <- read.csv("https://opendata.ecdc.europa.eu/covid19/casedistribution/csv", na.strings = "", fileEncoding = "UTF-8-BOM")

#####  Define functions to calculate effective reproduction number #######

####### Functions to use for effective reproduction number calculations   #######

##### Kevin systrom method   ########

est_rt_Systrom <- function(covid_cases, state_selected) {

  R_T_MAX = 12
  r_t_range = seq(0, R_T_MAX, length = R_T_MAX*100 + 1)
  GAMMA = 1/7

  #' Compute new cases and smooth them
  smooth_new_cases_1 <- function(cases){
    cases %>%
      arrange(date) %>%
      mutate(new_cases = c(cases[1], diff(cases))) %>%
      mutate(new_cases_smooth = round(
        smoother::smth.gaussian(new_cases, window = 7, alpha=7/4, tails = FALSE)
      )) %>%
      select(state, date, new_cases, new_cases_smooth)
  }

  #' Start using data first when smoothed value is above 25
  smooth_new_cases_25 <- function(x){
    start_index <- which(x$new_cases_smooth > 25)[1]
    length_vector <-dim(x)[1]
    return(x[start_index:length_vector,])
  }

  compute_likelihood <- function(cases){
    likelihood <- cases %>%
      filter(new_cases_smooth > 0) %>%
      mutate(
        r_t = list(r_t_range),
        lambda = map(lag(new_cases_smooth, 1), ~ .x * exp(GAMMA * (r_t_range - 1))),
        likelihood_r_t = map2(new_cases_smooth, lambda, dpois, log = TRUE)
      ) %>%
      slice(-1) %>%
      select(-lambda) %>%
      unnest(c(likelihood_r_t, r_t))
  }

  compute_posterior <- function(likelihood){
    likelihood %>%
      arrange(date) %>%
      group_by(r_t) %>%
      mutate(posterior = exp(
        zoo::rollapplyr(likelihood_r_t, 7, sum, partial = TRUE)
      )) %>%
      group_by(date) %>%
      mutate(posterior = posterior / sum(posterior, na.rm = TRUE)) %>%
      # HACK: NaNs in the posterior create issues later on. So we remove them.
      mutate(posterior = ifelse(is.nan(posterior), 0, posterior)) %>%
      ungroup() %>%
      select(-likelihood_r_t)
  }

  # Estimate R_t and a 95% highest-density interval around it
  estimate_rt <- function(posteriors){
    posteriors %>%
      group_by(state, date) %>%
      summarize(
        r_t_simulated = list(sample(r_t_range, 10000, replace = TRUE, prob = posterior)),
        r_t_most_likely = r_t_range[which.max(posterior)]
      ) %>%
      mutate(
        r_t_lo = map_dbl(r_t_simulated, ~ hdi(.x)[1]),
        r_t_hi = map_dbl(r_t_simulated, ~ hdi(.x)[2])
      ) %>%
      select(-r_t_simulated)
  }

  estimates <- covid_cases %>%
    filter(state == state_selected) %>%
    smooth_new_cases_1() %>%
    smooth_new_cases_25()  %>%
    compute_likelihood() %>%
    compute_posterior() %>%
    estimate_rt()

  return(estimates)

}

##### cislaghi method

est_rt_Cislaghi <- function(ts, window_interval) {
  data1 <- accelerometry::movingaves(ts, window=5)
  names(data1) <- names(ts)[3:(length(ts)-2)]
  res <- sapply( (1+window_interval):length(data1), function(t) {
    data1[t]/data1[t-window_interval]
  })
  return(res)

}

###### JRC method ###########################

est_rt_JRC <- function(ts, window_interval) {
  res <- sapply( 1:(length(ts)-window_interval), function(t) {
    ((log(ts[t+window_interval])-log(ts[t]))/window_interval)*7 + 1
  })
  return(res)

}

#### RKI method   ##############################

est_rt_rki <- function(ts, GT=4L) {
  # Sanity check
  if (!is.integer(GT) | !(GT>0)) stop("GT has to be postive integer.")
  # Estimate, if s=1 is t-7
  res <- sapply( (2*GT):length(ts), function(t) {
    sum(ts[t-(0:(GT-1))]) / sum(ts[t-2*GT+1:GT])
  })
  names(res) <- names(ts)[(2*GT):length(ts)]
  return(res)
}

###### Exponential growth, Wallinga & Lipsitch   ########

R.from.r <- function (r, GT) {
  Tmax = length(GT$GT)
  R = 1/sum(GT$GT * (exp(-r * (0:(Tmax - 1)))))
}

est_rt_exp <- function(ts, GT_obj, half_window_width=3L) {
  # Loop over all time points where the sliding window fits in
  res <- sapply((half_window_width+1):(length(ts)-half_window_width), function(t) {
    # Define the sliding window
    idx <- (t-half_window_width):(t+half_window_width)
    data <- data.frame(Date=1:length(ts), y=as.numeric(ts))
    # Fit a Poisson GLM
    m <- glm( y ~ 1 + Date, family=poisson, data= data  %>% slice(idx))
    # Extract the growth rate
    r <- as.numeric(coef(m)["Date"])

    # Equation 2.9 from Wallinga & Lipsitch
    R <- R.from.r(r, GT_obj)
    return(R)
  })
  names(res) <- names(ts)[(half_window_width+1):(length(ts)-half_window_width)]
  return(res)
}

######### Wallinga & Teunis   ##############

est_rt_wt <- function(ts1, GT_obj) {
  end1 <- length(ts1)

  R0::est.R0.TD(ts1, GT=GT_obj, begin=1, end=as.numeric(end1),correct=TRUE, nsim=1000)
}

#### function to calculate country specific value

Est_R0 <- function(start_date, country_name) {

  data2<-data_ECDC[,c("dateRep", "cases", "countriesAndTerritories")]
  # Set cases with minus value to 0
  data2$cases[which(data2$cases <0)]<-0
  # change format of Dates
  data2$dateRep <- as.Date(as.character(data2$dateRep), format="%d/%m/%Y")
  names(data2) <- c("Date", "y", "CountryName")

  data2<-data_ECDC[,c("dateRep", "cases", "countriesAndTerritories")]
  # Set cases with minus value to 0
  data2$cases[which(data2$cases <0)]<-0
  # change format of Dates
  data2$dateRep <- as.Date(as.character(data2$dateRep), format="%d/%m/%Y")
  names(data2) <- c("Date", "y", "CountryName")
  eval(parse(text=paste("index1 <- which(data2$CountryName ==\"",country_name, "\")", sep="" )))
  out <- data2[index1,]
  # Get starting date for the country
  out<-out[dim(out)[1]:1,]

  out <- out[out$Date %in% seq(from=as.Date(start_date) , to=out$Date[length(out$Date)] , by=1),]

out3 <- out
    names(out3)<-c("date", "cases1", "state")

  out3$cases[1] <-out3$cases1[1]
  for (j in 2:dim(out3)[1]) {
    out3$cases[j] <-   out3$cases[j-1]+  out3$cases1[j]
  }

  #### cislaghi method  ######

  rt_Cislaghi_3 <- est_rt_Cislaghi(out$y  %>% setNames(out$Date), window_interval=3)
  rt_Cislaghi_4 <- est_rt_Cislaghi(out$y  %>% setNames(out$Date), window_interval=4)
  rt_Cislaghi_5 <- est_rt_Cislaghi(out$y  %>% setNames(out$Date), window_interval=5)
  rt_Cislaghi_6 <- est_rt_Cislaghi(out$y  %>% setNames(out$Date), window_interval=6)

  rt_Cislaghi_3_df <- data.frame(Date=as.Date(names(rt_Cislaghi_3)), R_hat=rt_Cislaghi_3, Method=" 3")
  rt_Cislaghi_4_df <- data.frame(Date=as.Date(names(rt_Cislaghi_4)), R_hat=rt_Cislaghi_4, Method=" 4")
  rt_Cislaghi_5_df <- data.frame(Date=as.Date(names(rt_Cislaghi_5)), R_hat=rt_Cislaghi_5, Method=" 5")
  rt_Cislaghi_6_df <- data.frame(Date=as.Date(names(rt_Cislaghi_6)), R_hat=rt_Cislaghi_6, Method=" 6")

  #### Kevins Systrom method    ######

  # r_t_range is a vector of possible values for R_t
  R_T_MAX = 12
  r_t_range = seq(0, R_T_MAX, length = R_T_MAX*100 + 1)

  # Gamma is 1/serial interval
  # https://wwwnc.cdc.gov/eid/article/26/6/20-0357_article
  GAMMA = 1/7

  eval(parse(text=paste(" rt_Systrom <- est_rt_Systrom(covid_cases=out3, state_selected=\"",country_name, "\")", sep="" )))

  rt_Systrom_df <- data.frame(Date=rt_Systrom$date, R_hat=rt_Systrom$r_t_most_likely, upper=rt_Systrom$r_t_hi, lower=rt_Systrom$r_t_lo, Method="Systrom")

  ####   JRC method    #######

  rt_jrc <- est_rt_JRC(out$y  %>% setNames(out$Date), window_interval=7)

  rt_jrc_df <- data.frame(Date=as.Date(names(rt_jrc)), R_hat=rt_jrc, Method="JRC, GT=7")

  #############  RKI method ################

  # RKI method as specified in 2020-04-24 version of the article
  rt_rki <- est_rt_rki(out$y  %>% setNames(out$Date), GT=4L)

  rt_rki_df <- data.frame(Date=as.Date(names(rt_rki)), R_hat=rt_rki, Method="RKI, GT=4")

  ###### Wallinga & Lipsitch method    #############

  # Exponential growth with discrete Laplace transform of a GT \equiv = 4 distribution
  rt_exp2 <- est_rt_exp( out$y %>% setNames(out$Date), GT_obj=R0::generation.time("empirical", c(0,0,0,0,0,0,0,1)), half_window_width=3)

  # Exponential growth with discrete Laplace transform of the correct GT distribution
  rt_exp <- est_rt_exp( out$y %>% setNames(out$Date), GT_obj=GT_obj, half_window_width=3)

  rt_exp_df <- data.frame(Date=as.Date(names(rt_exp)), R_hat=rt_exp, Method="Stochastic generation time")
  #rt_exp_df <- data.frame(Date=as.Date(names(rt_exp_int[1:dim(rt_exp_int)[2]])), R_hat=rt_exp_int[1,], lower= rt_exp_int[2,], upper=rt_exp_int[3,], Method="Exp-Growth, random GT")

  rt_exp2_df <- data.frame(Date=as.Date(names(rt_exp2)), R_hat=rt_exp2, Method="Generation time 7 days (constant)")

  ### Wallinga and Teuniss############

  rt_wt <- est_rt_wt( out$y, GT=GT_obj)

  rt_wt_df <- cbind(Date=out$Date[as.numeric(names(rt_wt$R))], R_hat=rt_wt$R, rt_wt$conf.int, Method="W & T, ramdom GT")

  ### Create dataset with all estmates

  out_data <- data.frame(Date = as.Date(out$Date),Cases = out$y)

  # Merge the data

names(rt_Cislaghi_3_df)[2] <- "Cislaghi_3"
out_data <-  merge(x=out_data, y=rt_Cislaghi_3_df[,1:2], by = "Date", all.x=TRUE)
names(rt_Cislaghi_4_df)[2] <- "Cislaghi_4"
out_data <-  merge(x=out_data, y=rt_Cislaghi_4_df[,1:2], by = "Date", all.x=TRUE)
names(rt_Cislaghi_5_df)[2] <- "Cislaghi_4"
out_data <-  merge(x=out_data, y=rt_Cislaghi_5_df[,1:2], by = "Date", all.x=TRUE)
names(rt_Cislaghi_6_df)[2] <- "Cislaghi_4"
out_data <-  merge(x=out_data, y=rt_Cislaghi_6_df[,1:2], by = "Date", all.x=TRUE)
names(rt_Systrom_df)[2:4] <- c("Systrom", "Systrom_upper", "Systrom_lower")
out_data <-  merge(x=out_data, y=rt_Systrom_df[,1:4], by = "Date", all.x=TRUE)
names(rt_jrc_df)[2] <- "JRC"
out_data <-  merge(x=out_data, y=rt_jrc_df[,1:2], by = "Date", all.x=TRUE)
names(rt_rki_df)[2] <- "RKI"
out_data <-  merge(x=out_data, y=rt_rki_df[,1:2], by = "Date", all.x=TRUE)
names(rt_exp_df)[2] <- "Exp_growth_random_GT"
out_data <-  merge(x=out_data, y=rt_exp_df[,1:2], by = "Date", all.x=TRUE)
names(rt_exp2_df)[2] <- "Exp_growth_constant_GT"
out_data <-  merge(x=out_data, y=rt_exp2_df[,1:2], by = "Date", all.x=TRUE)
names(rt_wt_df)[2:4] <- c("Wallinga_Teunis", "Wallinga_Teunis_Upper", "Wallinga_Teunis_lower")
out_data <-  merge(x=out_data, y=rt_wt_df[,1:4], by = "Date", all.x=TRUE)

return(out_data)

}

R_estimates <- Est_R0(start_date=starting_date, country_name=Country_Name)

write.csv(R_estimates, file.path(Homedir, "Reproduction_number_estimates.csv"))


#' Generic M-step function to re-estimate the paramters of an emission probability distribution.
#'
#' @param x numeric matrix, size T x p, where T is the number of observations, p the dimension of the obs data
#' @param wt weight matrix
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return the five types of M-step results
#' @description Re-estimate the parameters of an emission distribution as part of the EM algorithm for HMM.
#'    This is called by the hext function.
#' @usage mstep.generic(x, wt, model)
#' @details see example code for usage
#' @author Horace W Tso
#' @export
mstep.generic = function(x, wt, model) {
  # type codes (mtype):
  #   discrete/binomial = 1
  #   multivariate normal = 2
  #   gamma = 3
  #   beta = 4
  #   gaussian mixture = 5
  D = model$D
  mtype = model$mtype
  ix1 = which(mtype == 1)
  if ( length(ix1) == 1 ) {
    # if there is only one column in obs, then no need to use ix1
    if ( length(mtype) == 1 ) {
      M1 = mstep.disc(x, wt, D)
    } else {
      # this is the case where there is only one discrete column
      # and one or more other types of observation, thus need ix1.
      M1 = mstep.disc(x[,ix1, drop=FALSE], wt, D)
    }
  } else if ( length(ix1) > 1 ) {
    M1 = mstep.mvdisc(x[,ix1], wt, D)
  } else {
    M1 = list(prob.l=list(), labels=NULL)
  }

  ix2 = which(mtype == 2)
  if ( length(ix2) == 1 ) {
    M2 = mstep.norm.hext(x[,ix2, drop=FALSE], wt)
  } else if ( length(ix2) > 1 ) {
    # =================================
    M2 = mstep.mvnorm.wish(x[,ix2], wt)
    # =================================
  } else {
    M2 = list(mu.l=list(), sigma.l=list())
  }

  ix3 = which(mtype == 3)
  if ( length(ix3) == 1 ) {
    M3 = mstep.gamma(x[,ix3, drop=FALSE], wt)
  } else if ( length(ix3) > 1 ) {
    M3 = mstep.mvgamma(x[,ix3], wt)
  } else {
    M3 = list(shape.l=list(), scale.l=list())
  }

  ix4 = which(mtype == 4)
  if ( length(ix4) == 1 ) {
    M4 = mstep.beta(x[,ix4, drop=FALSE], wt)
  } else if ( length(ix4) > 1 ) {
    M4 = mstep.mvbeta(x[,ix4], wt)
  } else {
    M4 = list(shape1.l=list(), shape2.l=list())
  }

  ix5 = which(mtype == 5)
  if ( length(ix5) > 0 ) {
    #cat("ix5:", ix5, "\n")
    # ====================================================================
    M5 = mstep.gaussian.mix(x[,ix5, drop=FALSE], wt, model)
    # ====================================================================
  } else {
    M5 = list(c.l=list(), mu.ll=rep(list(list()),K), sigma.ll=rep(list(list()),K))
  }
  return(list(M1=M1, M2=M2, M3=M3, M4=M4, M5=M5))
}

#' M-step re-estimation function for Gaussian Mixture emission distribution with inverse Wishart regularizer
#'
#' @param x numeric matrix, size T x p, where T is the number of observations, p the dimension of the obs data
#' @param w weight matrix, gamma from EStep C function
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of the gaussian mixture emission distribution as part of the EM algorithm for HMM.
#' @details Ref : Snoussi, Mohammad-Djafari (2005), Degeneracy and likelihood penalization
#'      in multivariate gaussian mixture models, preprint.
#'      NOTE : gam.jk is *not* in log domain.
#' @author Horace W Tso
#' @export
mstep.gaussian.mix = function (x, w, model) {
    idx = apply(is.na(x), 1, any)
    x = x[!idx, , drop = FALSE]
    w = w[!idx, , drop = FALSE]
    K = ncol(w)  # should agree with model$K
    Ttot = NROW(x) # = sum(model$TT)
    p = NCOL(x)  # = model$p
    M = model$M
    wdf = model$wdf

    if ( any(is.na(w)) | any(is.infinite(w))) {
      stop("na/infty found in w[].")
    }
    if ( all(w < 1e-100) == TRUE ) {
      stop("all w[] are near zero.")
    }
    w.tot = colSums(w)  # denominator for the reestimation equations, 1 x K

    emission = list(c.l=list(), mu.ll=rep(list(list()),K), sigma.ll=rep(list(list()),K))
    for ( i in 1:K ) {
      emission$c.l[[i]] = rep(NA, M)
    }

    gam.jk = array(NA, dim=c(Ttot, K, M))
    for ( i in 1:K )  {
       mu.l = model$parms.emission$M5$mu.ll[[i]]
       #if ( length(mu.l) != M ) cat("length of mu.l <> M\n")
       sigma.l = model$parms.emission$M5$sigma.ll[[i]]
       #if ( any(is.na(sigma.l))) {
        # stop("na in sigma.l:")
       #}
       #if ( length(sigma.l) != M ) cat("length of sigma.l <> M\n")
       c.v = model$parms.emission$M5$c.l[[i]]  # c.v is a vector of length M
       #if ( length(c.v) != M ) cat("length of c.v <> M\n")
       # Compute the conditional probabilities of the gaussian components in log domain :
       logbb = matrix(NA, nrow=Ttot, ncol=M)
       for ( m in 1:M ) {
         logbb[,m] = log(c.v[m]) + Dmvnorm(x, mean=mu.l[[m]], sigma=sigma.l[[m]], log=TRUE)
       }
       # Sum the contribution from the individual gaussian components
       # <==> sum by rows ==> a vector of length Ttot
       denom = double(Ttot)
       for ( tt in 1:Ttot) {
         denom[tt] = Elnsum.vec(logbb[tt,]) # in the non-log version: denom = rowSums(bb)
       }
       for ( m in 1:M ) {
          tmp = logbb[,m] + log(w[,i]) - denom
          gam.jk[,i,m] = exp(tmp)
       }
       if ( any(is.na(gam.jk[,i,m])) | any(is.infinite(gam.jk[,i,m]))) {
         cat("na/inf in gam.jk: i, m ", i, m)
         stop("...")
       }
    }

    for ( i in 1:K ) {
      # sum over mixture state k and time t, for every state i
      denom.j = sum(gam.jk[,i,])  # sum over mixture and sum over time, a scalar
      # cat("denom.j:", denom.j, "\n")
      for ( k in 1:M ) {
        if ( all(gam.jk[,i,k] < 1e-100) == FALSE ) {
          # c_jk =======================================================================
          emission$c.l[[i]][k] = sum(gam.jk[,i,k]) / denom.j  # sum over t; scalar / scalar = scalar
          #if ( any(is.na(emission$c.l[[i]][k])) ) {
          #  warning("na in c.l at m-step.")
          #}
          # mu_jk ====================================================================================

          xx = gam.jk[,i,k] * x   # (1 x T) (*) (T x p) = T x p, component by component multiplication
          denom.jk = sum(gam.jk[,i,k])  # summing over t for every i and k, result is a scalar
          mu.bar = colSums(xx) / denom.jk  # (1 x p) / scalar = 1 x p
          #if ( any(is.na(mu.bar)) ) {
          #  warning("na in mu.bar at m-step.")
          #}
          emission$mu.ll[[i]][[k]] = mu.bar

          # sigma_jk =============================================================================================
          x.centered = sweep(x, 2, mu.bar)  # T x p, 'sweep'' subtracts mu.bar ( 1 x p vector) from each row of x,
                                            # in effect, it is subtracting the mean from each column of
                                            # the multivariate observation vector (matrix), or (O_t - mu_jk)
          mat.sum = matrix(0, nrow=p, ncol=p)
          for ( tt in 1:Ttot ) {
            crossp = x.centered[tt,] %*% t(x.centered[tt,]) # x %*% t(x) = p x p matrix
            mat.w = crossp * gam.jk[tt,i,k]  # p x p matrix * scalar = p x p
            mat.sum = mat.sum + mat.w   # element by element addition : p x p
          }
          # Regularization with Inverse Wishart distribution ======================
          # Fraser, p. 55, regularized covariance matrix.
          # see p. 8 of Snoussi(2005), where nu = p + wdf, in numerator
          # nu + n + 1 = (p + wdf) + p + 1 = wdf + 2*p + 1, in denominator
          cov.mat = ( mat.sum + (p + wdf)*diag(p) ) / ( denom.jk + wdf + 2*p + 1 )
          # =======================================================================
          # cat("na in cov.mat:", sum(is.na(cov.mat)), "\n")
          # =============================================================
          # matrix's columns and rows must have the same names, otherwise
          # isSymmetric in dmvnorm complains
          rownames(cov.mat) = paste("X", 1:p, sep="")
          colnames(cov.mat) = rownames(cov.mat)
          #if ( any(is.na(cov.mat)) ) {
          #  warning("na in cov.mat at m-step.")
          #}
          emission$sigma.ll[[i]][[k]] = cov.mat
        } else {
          emission$c.l[[i]][k] = 0
          emission$mu.ll[[i]][[k]] = rep(0, p)
          emission$sigma.ll[[i]][[k]] = diag(p)
        }
      }
    }
    emission
}

#' M-step re-estimation function for multivariate gaussian emission distribution with inverse Wishart regularizer
#'
#' @param x numeric matrix, size T x p, where T is the number of observations, p the dimension of the obs data (columns are features)
#' @param w weight matrix, size T x K, where each row is the weight for a time step, and columns are states.
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of the multivariate normal emission distribution with inverse Wishart regularization.
#' @details Ref : Snoussi, Mohammad-Djafari (2005), Degeneracy and likelihood penalization
#'      in multivariate gaussian mixture models, preprint.
#'      wdf : Wishart degree of freedom, i.e. wishart df = wdf + p, where p is the dimension of the data. Thus, wdf > -1.
#' @author Horace W Tso
#' @export
mstep.mvnorm.wish = function (x, w, wdf=1) {
    ix = apply(is.na(x), 1, any)
    x = x[!ix, , drop = FALSE]
    w = w[!ix, , drop = FALSE]
    K = ncol(w)
    TT = NROW(x)
    p = NCOL(x)

    w.tot = colSums(w)  # denominator for the reestimation equations, 1 x K
    #if ( any(w.tot < 1e-100) ) {
      # cat("w.tot :", w.tot, "\n")
      #stop("w.tot too small!")
    #}
    emission = list(mu.l=list(), sigma.l=list())
    mat.sum = matrix(0, nrow=p, ncol=p)
    for ( i in 1:K ) {
        xx = w[,i] * x   # (1 x T) x (T x p) = T x p
        mu.bar = colSums(xx) / w.tot[i]  # (1 x p) x scalar = 1 x p
        x.centered = sweep(x, 2, mu.bar)  # T x p
        for ( tt in 1:TT ) {
          crossp = x.centered[tt,] %*% t(x.centered[tt,])
          mat.w = crossp * w[tt, i]  # p x p matrix times a scalar  = p x p
          mat.sum = mat.sum + mat.w   # element by element addition : p x p
        }
        # Regularization using the Inverse Wishart distribution ==============
        # Ref : Snoussi (2005), p. 8, where the inverse scale matrix is to be taken
        # as a diagonal matrix, wish.df is nu, the degree of freedom, p is the
        # dimension of the data n.
        cov.mat = (mat.sum + (p + wdf)*diag(p)) / ( w.tot[i] + wdf + 2*p + 1 )
        # ====================================================================
        # matrix's columns and rows must have the same names, otherwise
        # isSymmetric() complains
        rownames(cov.mat) = paste("X", 1:p, sep="")
        colnames(cov.mat) = rownames(cov.mat)
        emission$mu.l[[i]] = mu.bar
        emission$sigma.l[[i]] = cov.mat
    }
    emission
}

#' M-step re-estimation function for multivariate gaussian emission distribution, no regularization
#'
#' @param x numeric matrix, size T x p, where T is the number of observations, p the dimension of the obs data
#' @param w weight matrix, gamma from EStep C function
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of the multivariate gaussian emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @author Horace W Tso
#' @export
mstep.mvnorm.hext = function (x, wt) {
    ix = apply(is.na(x), 1, any)
    x = x[!ix, , drop = FALSE]
    wt = wt[!ix, , drop = FALSE]
    emission = list(mu.l=list(), sigma.l=list())
    for (i in 1:ncol(wt)) {
        tmp <- cov.wt(x, wt[,i])
        emission$mu.l[[i]] <- tmp$center
        emission$sigma.l[[i]] <- tmp$cov
    }
    emission
}

#' M-step re-estimation function for univarite gaussian emission distribution, no regularization
#'
#' @param x numeric matrix, size T x p, where T is the number of observations, p the dimension of the obs data
#' @param w weight matrix, gamma from EStep C function
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of the univariate gaussian emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @author Horace W Tso
#' @export
mstep.norm.hext = function (x, wt) {
    k = ncol(wt)
    mu = numeric(k)
    sigma = numeric(k)
    for (i in 1:k) {
        tmp = cov.wt(data.frame(x[!is.na(x)]), wt[!is.na(x),i])
        mu[i] = tmp$center
        sigma[i] = tmp$cov
    }
    list(mu.l = mu, sigma.l = sigma)
}

#' M-step re-estimation function for univariate discrete emission distribution
#'
#' @param x numeric matrix, size sum(T) x 1. All sequences stacked onto one column
#' @param w weight matrix, size sum(T) x K, gamma(i, t)
#' @param D number of alphabets in x
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of discrete emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @author Horace W Tso
#' @export
mstep.disc = function(x, w, D) {
  K = NCOL(wt)
  TT = NROW(x)
  w.sum = colSums(w)  # a vector of length K
  if ( length(which(w.sum < 1e-300)) > 0 ) {
    stop("some of the wt column sums are too close to zero.")
  }
  # prob.l is a list of matrices
  prob.l = list()
  # In this univariate case, b.vec is really just a vector; but to be consistent with
  # the usage in generic.density, i have to make it a single column matrix.
  b.vec = matrix(0, nrow=D, ncol=1)
  for ( j in 1:K ) {
    for ( fa in 1:D ) {
      ix = which( x == fa )
      b.vec[fa, 1] = sum(w[ix, j])
    }
    prob.l[[j]] = b.vec / w.sum[j]
  }
  list(prob.l=prob.l, labels=NULL)
}

#' M-step re-estimation function for multivariate discrete emission distribution
#'
#' @param x numeric matrix, size sum(T) x p
#' @param w weight matrix, size T x K
#' @param D number of alphabets in x, a vector of length equal to ncol(x)
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of discrete emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @author Horace W Tso
#' @export
mstep.mvdisc = function(x, w, D)  {
  K = NCOL(wt)
  TT = NROW(x)
  p = NCOL(x)
  wt.sum = colSums(wt)  # a vector of length K
  if ( length(which(wt.sum < 1e-300)) > 0 ) {
    stop("some of the wt column sums are too close to zero.")
  }
  prob.l = list()
  lookup.mat = matrix(0, nrow=max(D), ncol=p)
  for ( j in 1:K ) {
    for ( k in 1:p ) {
      for ( fa in 1:D[k] ) {
        # IMPORTANT NOTE : here i assume the discrete obs take on numeric value
        # from 1 to D[k] for column k of the obs. The next line looks for all
        # occurences of the factor (fa) in obs x[,k],
        ix = which( x[,k] == fa )
        # then sum up the soft weights, i.e. probability of the occurence of fa
        # in x for state j.
        lookup.mat[fa, k] = sum(wt[ix, j])
        # RECALL : wt = P( x_t=fa | s_t=j )
        # it has as many rows as x and as many columns as the no of states
      }
    }
    prob.l[[j]] = lookup.mat / wt.sum[j]
  }
  list(prob.l=prob.l, labels=NULL)
}

#' M-step re-estimation function for univariate beta emission distribution
#'
#' @param x numeric matrix, size sum(T) x p
#' @param w weight matrix, size T x K
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of univariate beta emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @details Needs the solnp function in the Rsolnp package.
#' @author Horace W Tso
#' @export
mstep.beta = function(x, w) {
  require(Rsolnp)
  ix = which(is.na(x))
  if ( length(ix) > 0 ) {
    x = x[!ix]
    w = w[!ix, ,drop = FALSE]
  }
  para.start = c(runif(1,min=0.5,max=7), runif(1,min=0.5,max=7))   # c(shape1, shape2)
  lower.b = c(0.01, 0.01)
  upper.b = c(Inf, Inf)
  K = ncol(wt)
  Shape1 = list()
  Shape2 = list()
  i = 1
  nfail = 0
  while ( i <= K ) {
    sol = try(solnp(para.start, Qaux.beta, x=x, wt=w[,i],
             LB=lower.b, UB=upper.b,
             control=list(trace=FALSE, inner.iter=1000, outer.iter=1000,
             tol=1e-13, delta=1e-13)), FALSE)
    if ( !inherits(sol, "try-error") ) {
      Shape1[[i]] = sol$par[1]
      Shape2[[i]] = sol$par[2]
      i = i + 1
    } else {
      nfail = nfail + 1
      if ( nfail >= 5 ) {
        stop("Too many problems with solnp!")
      }
    }
  }
  list(shape1.l=Shape1, shape2.l=Shape2)
}

#' M-step re-estimation function for multivariate beta emission distribution
#'
#' @param x numeric matrix, size sum(T) x p
#' @param w weight matrix, size T x K
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of multivariate beta emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @details Needs the solnp function in the Rsolnp package.
#' @author Horace W Tso
#' @export
mstep.mvbeta = function(x, wt) {
  require(Rsolnp)
  p = NCOL(x)
  K = NCOL(wt)
  ix = apply(is.na(x), 1, any)
  if ( length(ix) > 0 ) {
    x = x[!ix, ,drop = FALSE]
    wt = wt[!ix, ,drop = FALSE]
  }
  emission = list(shape1.l=list(), shape2.l=list())
  lower.b = c(0.01, 0.01)
  upper.b = c(Inf, Inf)
  h1.mat = matrix(NA, ncol=K, nrow=p)
  h2.mat = matrix(NA, ncol=K, nrow=p)
  for ( j in 1:p ) {
    i = 1
    nfail = 0
    while ( i <= K ) {
      para.start = c(runif(1,min=0.5,max=7), runif(1,min=0.5,max=7))
      sol <- try(solnp(para.start, Qaux.beta, x=x[,j], wt=wt[,i],
               LB=lower.b, UB=upper.b,
               control=list(trace=FALSE, inner.iter=1000, outer.iter=1000,
               tol=1e-13, delta=1e-13)), FALSE)
      if (!inherits(sol,"try-error")) {
        h1.mat[j,i] = sol$par[1]
        h2.mat[j,i] = sol$par[2]
        i = i + 1
      } else {
        nfail = nfail + 1
        if ( nfail >= 5 ) {
          stop("Too many problems with solnp!")
        }
      }
    }
  }
  for ( i in 1:K ) {
    # columns are states
    emission$shape1.l[[i]] = h1.mat[,i]
    emission$shape2.l[[i]] = h2.mat[,i]
  }
  return(emission)
}

#' M-step re-estimation function for univariate gamma emission distribution
#'
#' @param x numeric vector
#' @param w weight matrix, size T x K
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of univariate gamma emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @details Needs the solnp function in the Rsolnp package.
#' @author Horace W Tso
#' @export
mstep.gamma = function(x, w) {
  require(Rsolnp)
  ix = which(is.na(x))
  if ( length(ix) > 0 ) {
    x = x[!ix]
    w = w[!ix, ,drop = FALSE]
  }
  lower.b = c(0.01, 0.01)
  upper.b = c(Inf, Inf)
  K = ncol(w)
  Shape = list()
  Scale = list()
  i = 1
  nfail = 0
  while( i <= K ) {
    para.start = c(runif(1,min=0.5,max=10), runif(1,min=0.1,max=2.5))   # c(shape, scale)
    sol <- try(solnp(para.start, Qaux.gamma, x=x, wt=w[,i],
             LB=lower.b, UB=upper.b,
             control=list(trace=FALSE, inner.iter=1000, outer.iter=1000,
             tol=1e-13, delta=1e-13)), FALSE)
    if ( !inherits(sol, "try-error")) {
       Shape[[i]] = sol$par[1]
       Scale[[i]] = sol$par[2]
       i = i + 1
    } else {
      nfail = nfail + 1
      if ( nfail >= 5 ) {
        stop("Too many problems with solnp!")
      }
    }
  }
  list(shape.l=Shape, scale.l=Scale)
}

#' M-step re-estimation function for multivariate gamma emission distribution
#'
#' @param x numeric vector
#' @param w weight matrix, size T x K
#' @param model HMM model as defined in hmmspec.gen(..)
#' @return emission slot of Hext object
#' @description Re-estimates the parameters of multivariate gamma emission distribution as part of the EM algorithm for HMM. No regularizer is used.
#' @details Needs the solnp function in the Rsolnp package.
#' @author Horace W Tso
#' @export
mstep.mvgamma = function(x, w) {
  require(Rsolnp)
  p = NCOL(x)
  K = NCOL(w)
  ix = apply(is.na(x), 1, any)
  if ( length(ix) > 0 ) {
    x = x[!ix, ,drop = FALSE]
    w = w[!ix, ,drop = FALSE]
  }
  emission = list(shape.l=list(), scale.l=list())
  lower.b = c(0.01, 0.01)
  upper.b = c(Inf, Inf)
  a.mat = matrix(NA, ncol=K, nrow=p)
  s.mat = matrix(NA, ncol=K, nrow=p)
  for ( j in 1:p ) {
    i = 1
    nfail = 0
    while ( i <= K ) {
      para.start = c(runif(1,min=0.5,max=10), runif(1,min=0.1,max=2.5))
      sol <- try(solnp(para.start, Qaux.gamma, x=x[,j], wt=w[,i],
               LB=lower.b, UB=upper.b,
               control=list(trace=FALSE, inner.iter=1000, outer.iter=1000,
               tol=1e-13, delta=1e-13)), FALSE)
      if (!inherits(sol,"try-error")) {
        a.mat[j,i] = sol$par[1]
        s.mat[j,i] = sol$par[2]
        i = i + 1
      } else {
        nfail = nfail + 1
        if ( nfail >= 5 ) {
          stop("Too many problems with solnp!")
        }
      }
    }
  }
  for ( i in 1:K ) {
    # columns are states
    emission$shape.l[[i]] = a.mat[,i]
    emission$scale.l[[i]] = s.mat[,i]
  }
  return(emission)
}


decay_instant <- function(t, hl, C_max){
     C_max * 2^(-t/hl)
 }

decay_instant_integral <- function(t, hl, C_max){
    -hl * C_max * 2^(-t/hl)/log(2)
 }

return_druginf <- function(){

    druginf <- list()

    druginf[['3TC']] <-
        c(IC_50 = 0.0298, slope = 1.15, c_max = 15.3, hl = 10, dosing = 2)
    druginf[['AZT']] <-
        c(IC_50 = 0.1823, slope = 0.85, c_max = 4.5, hl = 8.5, dosing = 2)
    druginf[['D4T']] <-
        c(IC_50 = 0.5524, slope = 1.13, c_max = 2.3, hl = 3.5, dosing = 2)
    druginf[['EFV']] <-
        c(IC_50 = 0.0035, slope = 1.69, c_max = 12.9, hl = 35.8, dosing = 1)
    druginf[['NFV']] <-
        c(IC_50 = 0.2360, slope = 1.88, c_max = 5.1, hl = 4.0, dosing = 3)
    druginf[['LPV']] <-
        c(IC_50 = 0.0380, slope = 2.1, c_max = 15.6, hl = 9.9, dosing = 2)
    druginf[['NVP']] <-
        c(IC_50 = 0.0490, slope = 1.49, c_max = 25.2, hl = 21.5, dosing = 1)

    return(druginf)
}

return_mutinf <- function(){
    mutinf <- list()

    #For 3TC, mutation options are:
    #K65R, M184V
    mutinf[['3TC']] <-
        c(
            #s is taken from K65R
            s = 0.41,
            #mu is equivalent
            mu =  1.1 * (10^-5),
            #rho and sigma are from M184V
            rho =  963,
            sigma = -0.58)

    #For AZT, mutation options are:
    #M184V, M41L, T215Y* (excluded, because multi-nuc)
    mutinf[['AZT']] <-
        c(
            #s is taken from M41L
            s = 0.17,
            #mu is taken from M184V
            mu = 1.1 * 10^(-5),
            #rho and sigma are taken from M41L
            rho = 2.2,
            sigma = 0.07)

    #For D4T, mutation options are:
    #M41L, T215Y* (excluded, because multi-nuc)
    #all parameters taken from M41L
    mutinf[['D4T']] <-
        c(s = 0.17,
          mu = 1.3 * 10^(-6),
#          rho = 1.0,
#          sigma = 0.07)
          rho = 1.08,
          sigma = -0.12)

    #For EFV, mutation options are:
    #G190S, K103N, Y181C
    mutinf[['EFV']] <-
        #s is taken from Y181C
        c(s = 0.26,
          #mu is taken from G190S
          mu = 2.2 * 10^(-5),
          #rho is taken from K103N
          rho = 85,
          sigma = -0.17)

    #For NFV, mutation options are
    #D30N, L90M,
    mutinf[['NFV']] <-
        #s is taken from D30N
        c(s = 0.27,
          #mut is taken from D30N
          mu = 5.5 * 10^(-5),
          #rho and sigma are taken from D30N
          rho = 2.3,
          sigma = -0.29)

    #For LPV/r, mutation options are
    #I47A*, I47V, V32I, V82A, V82F,
    mutinf[['LPV']] <-
        #s is taken from I47V
        c(s = 0.05,
          #mut is taken from V32I
          mu = 4.1 * 10^(-5),
          #rho and sigma are taken from I47V
          rho = 1.8,
          sigma = -0.29)

    #For NVP, mutation options are
    #G190S, K103N, Y181C, #Y181I
    mutinf[['NVP']] <-
        #s is taken from K103N
        c(s = 0.3,
          #mut is taken from G190S
          mu = 2.2 * 10^(-5),
          #rho and sigma are taken from K103N
          rho = 94,
          sigma = -0.15)

    return(mutinf)

}


#################################################################################
# Monte Carlo method to calculate Multi-Species Indices (MSI) 3.0
# CBS (Statistics Netherlands), june 2017
# 	- generate MSIs with confidence interval, accounting for sampling error
#   - baseyear for simulated MSIs is flexible
#   - baseyear for plots can be set to 100 for index or for smoothed trend value
#	  - linear trend classification
#   - smoothed trend classification
#   - compare linear trends before and after changepoint
#   - assess overall change and change in last years
#   - truncation of low indices (and their SE set to 0)
#################################################################################
# LOAD PACKAGES
library(ggplot2)

# DATA IMPORT AND SPECIFICATIONS
setwd("E:/Ontwikkel/Lslt/Graadmeter/BASISTOOLS") # set working directory for input and output files
rdata <- read.table("Padhout.txt", header=TRUE) # name of input file
jobname <- as.character ("Padhout") # generic name for output fileS
nsim <- 1000		    # number of Monte Carlo simulations
SEbaseyear <- 2013 # year to set SE to 0. SE of other years is relative to the SE of SEbaseyear
plotbaseyear <- 1965 # year to set MSI or smoothed trend (dependend on 'index_smooth') to 100 in final MSI-plot
index_smooth <- "SMOOTH" # INDEX / SMOOTH: "INDEX" will cause the MSI in plotbaseyear set to 100; "SMOOTH" will set the smoothed trend value in the plotbaseyear to 100
lastyears <- 10		# last X years of time series for which separate trend (short-term trends) should be calculated
maxCV <- 3       # maximum allowed mean Coefficient of Variation (CV) of species indices (0.5 = 50%). Species with higher mean CV are excluded
changepoint <- 1994 # compare trends before and after this year
truncfac <- 10    # truncation factor (=maximum allowed year-to-year index ratio). Default for Living Planet Index = 10
TRUNC <- 1 # set all indices below TRUNC to this value and their SE to 0. TRUNC <- 0 means no truncation
# testyear <- 1994 # additional year smoothed trend values should be compared with
# END OF SPECIFICATIONS

# Define and calculate model parameters
species <- rdata[,"species"]
year <- rdata[,"year"]
index <- rdata[,"index"]
se <- rdata[,"se"]
uspecies <- sort(unique(species))
nspecies <- length(uspecies)
uyear <- sort(unique(year))
nyear <- length(unique(year))
meanindex <- tapply(index, species, mean, na.rm=TRUE)/100
mnindex1 <- as.data.frame(rep(meanindex,each=nyear))
minyear <- min(year)
maxyear <- max(year)
plotbaseyear <- plotbaseyear-minyear+1
baseyear <- max(1,SEbaseyear-minyear+1)
lastyears <- min(lastyears,nyear)
# testyear <- nyear-(maxyear-testyear)
INP <- data.frame(cbind(species, year, index, se))
SPEC <- as.matrix(species)
species <- sort(rep(uspecies,each=nyear))
year <- rep(uyear,nspecies)
INP1 <- data.frame(cbind(species,year))
INP2 <- merge(INP, INP1, by=c("species","year"), sort=TRUE, all=TRUE)

# Calculate and plot mean CV for indices per species
CVtemp <- INP2[INP2$index >= 10, ] # select records with index >= 10 for CV calculation
CVtemp <- CVtemp[!is.na(CVtemp$index), ] # reject missing values
CV1 <- CVtemp$se/CVtemp$index
CV1[CV1== 0] <- NA
CV1[CV1== Inf] <- NA
species2 <- CVtemp$species
mnCV <- tapply(CV1, species2, mean, na.rm=TRUE)
plot(uspecies, mnCV, type="p", pch=19, col="black", main=jobname, xlab="species code", ylab="meanCV")
CV <- as.data.frame(rep(mnCV, each=nyear))

# replace small indices by TRUNC and their SE by zero
INP3 <- data.frame(cbind(species, year, CV, mnindex1))
colnames(INP3) <- c("species","year", "CV", "mnindex1")
INP4 <- merge(INP2, INP3, by=c("species","year"), sort=FALSE, all=TRUE)
INP5 <- subset(INP4, CV < maxCV, select = c(species, year, index, se, mnindex1))
INP5$index[INP5$index < TRUNC & !is.na(INP5$index)] <- TRUNC
INP5$se[INP5$index == TRUNC & !is.na(INP5$index)] <- 0

# reset parameters
index <- as.vector(INP5["index"])
se <- as.vector(INP5["se"])
nobs <- NROW(INP5)
uspecies <- sort(unique(INP5$species))
nspecies <- length(uspecies)
year <- rep(uyear, nspecies)

# Transform indices and standard deviations to log scale (Delta method)
LNindex <- as.vector(log(index))
LNse <- as.vector(se/index)

# Monte Carlo simulations of species indices
MC <- matrix(NA, nobs, nsim)
for (s in 1:nsim) {
       MC[,s] <- rnorm(nobs, LNindex[,1], LNse[,1])
 }
MC[MC < log(TRUNC)] <- log(TRUNC)

# impute missing values using chain method
CHAIN <- matrix(NA, nobs, nsim)
for (s in 1:nsim ){
  for (o in 1:nobs-1) {
    CHAIN[o,s] <- MC[o+1,s]-MC[o,s]
  }
  for (sp in 1:nspecies) {
    CHAIN[sp*nyear,s] <- NA
  }
}

CHAIN[CHAIN > log(truncfac)] <- log(truncfac)
CHAIN[CHAIN < log(1/truncfac)] <- log(1/truncfac)

mnCHAIN <- matrix(NA, nyear, nsim)
for (s in 1:nsim) {
  mnCHAIN[,s] <- tapply(CHAIN[,s], year, mean, na.rm=TRUE)
}

simMSI <- matrix(NA, nyear, nsim)
bmin <- baseyear-1
bplus <- baseyear+1
for (s in 1:nsim){
simMSI[baseyear,s] <- log(100)
  for (y in bmin:min(bmin,1)){
    simMSI[y,s] <- simMSI[y+1,s] - mnCHAIN[y,s]
  }
  for (y in min(bplus,nyear):nyear){
    simMSI[y,s] <- simMSI[y-1,s] + mnCHAIN[y-1,s]
  }
}

# calculate MSI and SE
meanMSI <- array(NA, dim=c(1, nyear))
stdevMSI <- array(NA, dim=c(1, nyear))
for (y in 1:nyear) {
  meanMSI[y] <- mean(simMSI[y,])
  stdevMSI[y] <- sd(simMSI[y,])
}

# Monte Carlo simulations of MSI (for trend calculation)
simMSI <- matrix(NA, nyear, nsim)
for (s in 1:nsim) {
  for (y in 1:nyear) {
    simMSI[y,s] <- rnorm(1, meanMSI[y], stdevMSI[y])
  }
}

# Back-transformation of MSI to index scale (Delta method)
meanMSI <- array(NA, dim=c(1, nyear))
stdevMSI <- array(NA, dim=c(1, nyear))
for (y in 1:nyear) {
  meanMSI[y] <- round(exp(mean(simMSI[y,])), digits=2)
  stdevMSI[y] <- round(sd(simMSI[y,])*meanMSI[y], digits=2)
}

# Confidence interval of MSI on index-scale
CI <- array(NA, dim=c(nyear, 2))
for (y in 1:nyear){
  CI[y,1] <- exp(mean(simMSI[y,])-1.96*sd(simMSI[y,]))
  CI[y,2] <- exp(mean(simMSI[y,])+1.96*sd(simMSI[y,]))
}

# Overall linear trend per simulation
lrMSI <- array(NA, dim=c(2, nsim))
for (s in 1:nsim) {
  summ <- summary(lm(simMSI[,s] ~ uyear, na.action=na.exclude))
  lrMSI[1,s] <- round(summ$coefficients[2], digits=4)	# coeff[2] = slope
  lrMSI[2,s] <- round(summ$coefficients[2,2], digits=4) 	# coeff[2,2] = se of slope
}
SLOPE <- mean(lrMSI[1,])
sdSLOPE <- sd(as.vector(lrMSI[1,]))

# Back transform trends to index scale:
SLOPE_mult <- round(exp(SLOPE), digits=4)
sdSLOPE_mult <- round(sdSLOPE*SLOPE_mult, digits=4)

# Overall trend classification (see Soldaat et al. 2007 (J. Ornithol. DOI 10.1007/s10336-007-0176-7))
TrendClass <- "X"
if (SLOPE_mult - 1.96*sdSLOPE_mult > 1.05) {TrendClass <- "strong increase" } else
  if (SLOPE_mult + 1.96*sdSLOPE_mult < 0.95) { TrendClass <- "steep decline" } else
    if (SLOPE_mult - 1.96*sdSLOPE_mult > 1.00) { TrendClass <- "moderate increase"} else
      if (SLOPE_mult + 1.96*sdSLOPE_mult < 1.00) { TrendClass <- "moderate decline" } else
        if ((SLOPE_mult - 1.96*sdSLOPE_mult - 0.95)*(1.05 - SLOPE_mult + 1.96*sdSLOPE_mult) < 0.00) { TrendClass <- "uncertain"} else
          if ((SLOPE_mult + 1.96*sdSLOPE_mult) - (SLOPE_mult - 1.96*sdSLOPE_mult) > 0.10) { TrendClass <- "uncertain" } else
          {TrendClass <- "stable"}
TrendClass

# Short term linear trend per simulation
lrMSI_short <- array(NA, dim=c(2, nsim))
lyear <- c((maxyear-lastyears+1):maxyear)
simMSI_short <- array(NA, dim=c(lastyears, nsim))
for (s in 1:nsim) {
  for (y in 1:lastyears) {
    simMSI_short[y, s] <- simMSI[(nyear-lastyears+y),s]
  }
}
for (s in 1:nsim) {
  summ <- summary(lm(simMSI_short[,s] ~ lyear, na.action=na.exclude))
  lrMSI_short[1,s] <- round(summ$coefficients[2], digits=4)	# coeff[2] = slope
  lrMSI_short[2,s] <- round(summ$coefficients[2,2], digits=4) 	# coeff[2,2] = sd of slope
}
SLOPE_short <- mean(lrMSI_short[1,])
sdSLOPE_short <- sd(as.vector(lrMSI_short[1,]))

# Back-transform short-term trends to index scale
SLOPE_short_mult <- round(exp(SLOPE_short), digits=4)
sdSLOPE_short_mult <- round(sdSLOPE_short*SLOPE_short_mult, digits=4)

# Short term trend classification (Soldaat et al. 2007 (J. Ornithol. DOI 10.1007/s10336-007-0176-7))
TrendClass_short <- "X"
if (SLOPE_short_mult - 1.96*sdSLOPE_short_mult > 1.05) {TrendClass_short <- "strong increase" } else
  if (SLOPE_short_mult + 1.96*sdSLOPE_short_mult < 0.95) { TrendClass_short <- "steep decline" } else
    if (SLOPE_short_mult - 1.96*sdSLOPE_short_mult > 1.00) { TrendClass_short <- "moderate increase"} else
      if (SLOPE_short_mult + 1.96*sdSLOPE_short_mult < 1.00) { TrendClass_short <- "moderate decline" } else
        if ((SLOPE_short_mult - 1.96*sdSLOPE_short_mult - 0.95)*(1.05 - SLOPE_short_mult + 1.96*sdSLOPE_short_mult) < 0.00) { TrendClass_short <- "uncertain"} else
          if ((SLOPE_short_mult + 1.96*sdSLOPE_short_mult) - (SLOPE_short_mult - 1.96*sdSLOPE_short_mult) > 0.10) { TrendClass_short <- "uncertain" } else
          {TrendClass_short <- "stable"}
TrendClass_short

# linear trend before changepoint per simulation
lrMSI_before <- array(NA, dim=c(2, nsim))
byear <- c(minyear:changepoint)
byears <- length(byear)
simMSI_before <- array(NA, dim=c(byears, nsim))
for (s in 1:nsim) {
  for (y in 1:byears) {
    simMSI_before[y, s] <- simMSI[y,s]
  }
}
for (s in 1:nsim) {
  summb <- summary(lm(simMSI_before[,s] ~ byear, na.action=na.exclude))
  lrMSI_before[1,s] <- round(summb$coefficients[2], digits=4)  # coeff[2] = slope
  lrMSI_before[2,s] <- round(summb$coefficients[2,2], digits=4) 	# coeff[2,2] = sd of slope
}
SLOPE_before <- round(mean(lrMSI_before[1,]), digits=5)
sdSLOPE_before <- round(sd(as.vector(lrMSI_before[1,])), digits=5)

# Back-transform trends before changepoint to index scale
SLOPE_before_mult <- round(exp(SLOPE_before), digits=4)
sdSLOPE_before_mult <- round(sdSLOPE_before*SLOPE_before_mult, digits=4)

# Trend classification before changepoint
TrendClass_before <- "X"
if (SLOPE_before_mult - 1.96*sdSLOPE_before_mult > 1.05) {TrendClass_before <- "strong increase" } else
  if (SLOPE_before_mult + 1.96*sdSLOPE_before_mult < 0.95) { TrendClass_before <- "steep decline" } else
    if (SLOPE_before_mult - 1.96*sdSLOPE_before_mult > 1.00) { TrendClass_before <- "moderate increase"} else
      if (SLOPE_before_mult + 1.96*sdSLOPE_before_mult < 1.00) { TrendClass_before <- "moderate decline" } else
        if ((SLOPE_before_mult - 1.96*sdSLOPE_before_mult - 0.95)*(1.05 - SLOPE_before_mult + 1.96*sdSLOPE_before_mult) < 0.00) { TrendClass_before <- "uncertain"} else
          if ((SLOPE_before_mult + 1.96*sdSLOPE_before_mult) - (SLOPE_before_mult - 1.96*sdSLOPE_before_mult) > 0.10) { TrendClass_before <- "uncertain" } else
          {TrendClass_before <- "stable"}
TrendClass_before

# linear trend after changepoint per simulation
lrMSI_after <- array(NA, dim=c(2, nsim))
ayear <- c(changepoint:maxyear)
ayears <- length(ayear)
simMSI_after <- array(NA, dim=c(ayears, nsim))
for (s in 1:nsim) {
  for (y in 1:ayears) {
    simMSI_after[y, s] <- simMSI[(byears-1+y),s]
  }
}
for (s in 1:nsim) {
  summa <- summary(lm(simMSI_after[,s] ~ ayear, na.action=na.exclude))
  lrMSI_after[1,s] <- round(summa$coefficients[2], digits=4)  # coeff[2] = slope
  lrMSI_after[2,s] <- round(summa$coefficients[2,2], digits=4)   # coeff[2,2] = sd of slope
}
SLOPE_after <- round(mean(lrMSI_after[1,]), digits=5)
sdSLOPE_after <- round(sd(as.vector(lrMSI_after[1,])),digits=5)

# Trend classification before changepoint
SLOPE_after_mult <- round(exp(SLOPE_after), digits=4)
sdSLOPE_after_mult <- round(sdSLOPE_after*SLOPE_after_mult, digits=4)

# Trend classification after changepoint
TrendClass_after <- "X"
if (SLOPE_after_mult - 1.96*sdSLOPE_after_mult > 1.05) {TrendClass_after <- "strong increase" } else
  if (SLOPE_after_mult + 1.96*sdSLOPE_after_mult < 0.95) { TrendClass_after <- "steep decline" } else
    if (SLOPE_after_mult - 1.96*sdSLOPE_after_mult > 1.00) { TrendClass_after <- "moderate increase"} else
      if (SLOPE_after_mult + 1.96*sdSLOPE_after_mult < 1.00) { TrendClass_after <- "moderate decline" } else
        if ((SLOPE_after_mult - 1.96*sdSLOPE_after_mult - 0.95)*(1.05 - SLOPE_after_mult + 1.96*sdSLOPE_after_mult) < 0.00) { TrendClass_after <- "uncertain"} else
          if ((SLOPE_after_mult + 1.96*sdSLOPE_after_mult) - (SLOPE_after_mult - 1.96*sdSLOPE_after_mult) > 0.10) { TrendClass_after <- "uncertain" } else
          {TrendClass_after <- "stable"}
TrendClass_after

# compare linear trends before and after changepoint
compare <- array(NA, dim=c(nsim, 1))
for (s in 1:nsim){
  compare[s,1] <- lrMSI_after[1,s]-lrMSI_before[1,s]
}
meanTrendDiff <- mean(compare)
seTrendDiff <- sd(as.vector(compare))
significance <- NA
if (abs(meanTrendDiff)-2.58*seTrendDiff > 0) { significance <- "p<0.01" } else
  if (abs(meanTrendDiff)-1.96*seTrendDiff > 0) { significance <- "p<0.05" } else
    {significance <- "n.s."}

# Smoothing
# span = 0.75 is default in R and usually fits well. But trying other values may give better results, depending on your data
loessMSI <- array(NA, dim=c(nyear, nsim))
Diff <- array(NA, dim=c(nyear, nsim))
# Difftest <- array(NA, dim=c(nyear, nsim))
for (s in 1:nsim) {
  smooth <- predict(loess(simMSI[,s]~uyear, span=0.75, degree=2, na.action=na.exclude),
                    data.frame(uyear), se=TRUE)
  loessMSI[,s] <- round(smooth$fit, digits=4)

  for (y in 1:nyear) {
    Diff[y,s] <- loessMSI[nyear,s] - loessMSI[y,s]
#    Difftest[y,s] <- loessMSI[testyear,s] - loessMSI[y,s]
  }
}

# calculate overall change (first year set to 100) and change in last years
pct_CHANGE <- array(NA, dim=c(nsim, 2))
for (s in 1:nsim) {
pct_CHANGE[s,1] <- exp(loessMSI[nyear,s])*100/exp(loessMSI[1,s])-100
pct_CHANGE[s,2] <- exp(loessMSI[nyear,s])*100/exp(loessMSI[1,s])-exp(loessMSI[(nyear-lastyears+1),s])*100/exp(loessMSI[1,s])
}
pct_CHANGE_long <- round(mean(pct_CHANGE[,1]),digits=3)
sd_pct_CHANGE_long <- round(sd(pct_CHANGE[,1]), digits=3)
pct_CHANGE_short <- round(mean(pct_CHANGE[,2]),digits=3)
sd_pct_CHANGE_short <- round(sd(pct_CHANGE[,2]), digits=3)

# significance of percentage change
significance_PCT_long <- NA
if (abs(pct_CHANGE_long)-2.58*sd_pct_CHANGE_long > 0) { significance_PCT_long <- "p<0.01" } else
  if (abs(pct_CHANGE_long)-1.96*sd_pct_CHANGE_long > 0) { significance_PCT_long <- "p<0.05" } else
  {significance_PCT_long <- "n.s."}
significance_PCT_short <- NA
if (abs(pct_CHANGE_short)-2.58*sd_pct_CHANGE_short > 0) { significance_PCT_short <- "p<0.01" } else
  if (abs(pct_CHANGE_short)-1.96*sd_pct_CHANGE_short > 0) { significance_PCT_short <- "p<0.05" } else
  {significance_PCT_short <- "n.s."}

# Test for difference in smoothed trend value between each year and 'testyear'
# YEARTEST <- array(NA, dim=c(nyear, 4))
# for (y in 1:nyear) {
#   YEARTEST[y,1] <- round(mean(Difftest[y,]), digits=3)
#  YEARTEST[y,2] <- round(sd(Difftest[y,]), digits=3)
#  YEARTEST[y,3] <- round(YEARTEST[y,1]-1.96*YEARTEST[y,2], digits=3)
#  YEARTEST[y,4] <- round(YEARTEST[y,1]+1.96*YEARTEST[y,2], digits=3)
# }

# create output for flexible trend estimates
smoothMSI <- array(NA, dim=c(nyear, 13))
for (y in 1:nyear) {
smoothMSI[y,1] <- round(mean(loessMSI[y,]), digits=4) # smooth MSI on log scale
smoothMSI[y,2] <- round(sd(loessMSI[y,]), digits=4) # SE of smooth MSI on log scale
smoothMSI[y,3] <- round(mean(Diff[y,]), digits=4) # Difference smooth MSI with last year on log scale
smoothMSI[y,4] <- round(sd(Diff[y,]), digits=4) # SE of difference with ast year
smoothMSI[y,12] <- round(smoothMSI[y,1]-1.96*smoothMSI[y,2], digits=2) # lower CI of smooth MSI on log scale
smoothMSI[y,13] <- round(smoothMSI[y,1]+1.96*smoothMSI[y,2], digits=2) # upper CI of smooth MSI on log scale
}

# Trendclassification, based on smoothed trends (Soldaat et al. 2007 (J. Ornithol. DOI 10.1007/s10336-007-0176-7))
for (y in 1:nyear) {
  smoothMSI[y,5] <- round(smoothMSI[nyear,1]/smoothMSI[y,1], digits=4)		# = TCR
  smoothMSI[y,6] <- round(smoothMSI[y,5]-1.96*(smoothMSI[y,4]/smoothMSI[y,1]), digits=4)	# = CI- TCR
  smoothMSI[y,7] <- round(smoothMSI[y,5]+1.96*(smoothMSI[y,4]/smoothMSI[y,1]), digits=4)	# = CI+ TCR
  smoothMSI[y,8] <- round(exp(log(smoothMSI[y,5])/(nyear-y)), digits=4)		# = YCR
  smoothMSI[y,9] <- round(exp(log(smoothMSI[y,6])/(nyear-y)), digits=4)		# = CI- YCR
  smoothMSI[y,10] <- round(exp(log(smoothMSI[y,7])/(nyear-y)), digits=4)		# = CI+ YCR
}
for (y in 1:(nyear-1)) {
  if (smoothMSI[y,9] > 1.05) {smoothMSI[y,11] <- 1} else
    if (smoothMSI[y,10] < 0.95) {smoothMSI[y,11] <- 6} else
      if (smoothMSI[y,9] > 1.00) {smoothMSI[y,11] <- 2} else
        if (smoothMSI[y,10] < 1.00) {smoothMSI[y,11] <- 5} else
          if ((smoothMSI[y,9] - 0.95)*(1.05 - smoothMSI[y,10]) < 0.00) {smoothMSI[y,11] <- 3} else
            if ((smoothMSI[y,10]) - (smoothMSI[y,9]) > 0.10) {smoothMSI[y,11] <- 3} else
            {smoothMSI[y,11] <- 4}
}
TrendClass_flex <- matrix(NA, nrow= nyear, ncol=1)
for (y in 1:(nyear-1)) {
  if (smoothMSI[y,9] > 1.05) {TrendClass_flex[y] <- "strong_increase" } else
    if (smoothMSI[y,10] < 0.95) { TrendClass_flex[y] <- "steep_decline" } else
      if (smoothMSI[y,9] > 1.00) { TrendClass_flex[y] <- "moderate_increase"} else
        if (smoothMSI[y,10] < 1.00) { TrendClass_flex[y] <- "moderate_decline" } else
          if ((smoothMSI[y,9] - 0.95)*(1.05 - smoothMSI[y,10]) < 0.00) { TrendClass_flex[y] <- "uncertain"} else
            if ((smoothMSI[y,10]) - (smoothMSI[y,9]) > 0.10) { TrendClass_flex[y] <- "uncertain" } else
            {TrendClass_flex[y] <- "stable"}
}

# rescaling (for presentation of plot)
rescale <- NA
 if (index_smooth =="INDEX") {rescale <- 100/meanMSI[plotbaseyear]} else
  if (index_smooth =="SMOOTH") {rescale <- 100/exp(smoothMSI[plotbaseyear,1])} else
   {rescale <- NA}
simMSImean <- round(as.vector(rescale*meanMSI), digits=2)
simMSIsd <- round(as.vector(rescale*stdevMSI), digits=2)
uppCI_MSI <- round(rescale*CI[,2], digits=2)
lowCI_MSI <- round(rescale*CI[,1], digits=2)
trend_flex <- round(rescale*exp(smoothMSI[,1]), digits=2)
lowCI_trend_flex <- round(rescale*exp(smoothMSI[,12]), digits=2)
uppCI_trend_flex <- round(rescale* exp(smoothMSI[,13]), digits=2)

# create output file for MSI + smoothed trend
jobnameRESULTS <- paste (jobname, "_RESULTS.csv", sep = "")
RES <- as.data.frame(cbind(uyear, simMSImean, simMSIsd, lowCI_MSI, uppCI_MSI, trend_flex, lowCI_trend_flex, uppCI_trend_flex))
RES$trend_class <- TrendClass_flex
names(RES) <- c("year", "MSI", "sd_MSI", "lower_CL_MSI", "upper_CL_MSI", "Trend", "lower_CL_trend", "upper_CL_trend", "trend_class")
write.csv2(RES, file=jobnameRESULTS, row.names=FALSE, quote=FALSE)

# create output file with all linear trends
SIMTRENDS <- t(rbind(lrMSI[1,],lrMSI_short[1,]))
jobnameSIMTRENDS <- paste (jobname, "_SIMTRENDS.csv", sep = "")
write.csv2(SIMTRENDS, file=jobnameSIMTRENDS, row.names=FALSE)

# create output for linear trend estimates and % change
rownames <- rbind("overall trend","SE overall trend", "trend last years", "SE trend last years", "changepoint",
                  "trend before changepoint", "SE trend before changepoint", "trend after changepoint", "SE trend after changepoint",
                  "% change", "SE % change", "% change last years", "SE % change last years")
output <- cbind(SLOPE_mult, sdSLOPE_mult, SLOPE_short_mult, sdSLOPE_short_mult,
                changepoint, SLOPE_before_mult, sdSLOPE_before_mult, SLOPE_after_mult, sdSLOPE_after_mult,
                pct_CHANGE_long, sd_pct_CHANGE_long, pct_CHANGE_short, sd_pct_CHANGE_short)
TRENDS <- as.data.frame(t(output), row.names=rownames)
TRENDS$significance <- ""
TRENDS["changepoint", "significance"] <- significance
TRENDS[1, "significance"] <- TrendClass
TRENDS[3, "significance"] <- TrendClass_short
TRENDS[6, "significance"] <- TrendClass_before
TRENDS[8, "significance"] <- TrendClass_after
TRENDS[10, "significance"] <- significance_PCT_long
TRENDS[12, "significance"] <- significance_PCT_short
names(TRENDS) <- c("value", "significance")
jobnameTRENDS <- paste (jobname, "_TRENDS.csv", sep = "")
write.csv2(TRENDS, file=jobnameTRENDS)

# graph without error bars
limits <- aes(ymax = simMSImean+simMSIsd, ymin=simMSImean-simMSIsd)
jobnameGRAPH <- paste (jobname, "_GRAPH.jpg", sep = "")
x_position <- ifelse(SLOPE>0, maxyear-nyear/1.5, minyear+1)
y_position <- 50
legend10 <- paste(nspecies, "species")
legend11 <- paste("overall trend:", round(SLOPE_mult*100-100, digits=1), "% / yr,", TrendClass)
legend12 <- paste("last", lastyears, "years:", round(SLOPE_short_mult*100-100, digits=1), "% / yr,", TrendClass_short)
ylab <- paste("MSI (",plotbaseyear+minyear-1,"= 100 )")
lastyear01 <- c(rep(0,nyear-lastyears), rep(1,lastyears))
ggplot(RES, aes(x=year, y=MSI))+
  scale_color_manual(values=c("blue", "dark green"))+
  geom_point(colour = "dark green",size=3)+
  ylim(0, NA)+
  ylab(ylab)+
  ggtitle(jobname)+
        theme(plot.title = element_text(size = 25, face = "bold"),
        axis.title.x=element_blank(),
        axis.title.y=element_text(size=20,face="bold"),
        axis.text = element_text(size = 20),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill=NA, size=1))+
  geom_ribbon(aes(ymin=lowCI_trend_flex, ymax=uppCI_trend_flex), alpha=0.2)+
  geom_line(aes(y=trend_flex), colour="dark green", size=1)+
  annotate("text", x=x_position, y=y_position, label = legend10, size=7, fontface=2, hjust=0)+
  annotate("text", x=x_position, y=y_position-15, label = legend11, size=7, hjust=0)+
  annotate("text", x=x_position, y=y_position-30, label = legend12, size=7, hjust=0)
ggsave(jobnameGRAPH)

# graph with error bars
# limits <- aes(ymax = simMSImean+simMSIsd, ymin=simMSImean-simMSIsd)
jobnameGRAPH2 <- paste (jobname, "_GRAPH2.jpg", sep = "")
ggplot(RES, aes(x=year, y=MSI))+
  scale_color_manual(values=c("blue", "dark green"))+
  geom_point(colour = "dark green",size=3)+
  ylim(0, NA)+
  ylab(ylab)+
  ggtitle(jobname)+
  theme(plot.title = element_text(size = 25, face = "bold"),
        axis.title.x=element_blank(),
        axis.title.y=element_text(size=20,face="bold"),
        axis.text = element_text(size = 20),
        panel.background = element_blank(),
        panel.border = element_rect(colour = "black", fill=NA, size=1))+
  geom_ribbon(aes(ymin=lowCI_trend_flex, ymax=uppCI_trend_flex), alpha=0.2)+
  geom_line(aes(y=trend_flex), colour="dark green", size=1)+
  geom_pointrange(limits, colour="dark green")+
annotate("text", x=x_position, y=y_position, label = legend10, size=7, fontface=2, hjust=0)+
annotate("text", x=x_position, y=y_position-15, label = legend11, size=7, hjust=0)+
annotate("text", x=x_position, y=y_position-30, label = legend12, size=7, hjust=0)
ggsave(jobnameGRAPH2)


library(related)
library(ggplot2)

# read in genotypes  (modify this name as necessary for different colonies or datasets)
x <- read.table("filtered_100kbp.related", stringsAsFactors=F)

# input species name (e.g., herculeanus, laevigatus, or modoc)
species <- "laevigatus"

# input colony number (e.g., 1, 2, or 3)
colony_number <- 1

# define number of simulations
nsims <- 100

# simulations
# first, obtain overall allele freqs for each locus
# allow 5% variation of allele frequencies (i.e., a little randomness)
n_loci <- (ncol(x) - 1) / 2
alleles <- list()
allele_freqs <- list()
for(a in 1:n_loci) {
	alleles[[a]] <- c(x[,a * 2], x[,a * 2 + 1])
	af1 <- length(alleles[[a]][alleles[[a]] == 1]) / length(alleles[[a]])
	# add +/- 5% variation to allele frequencies (some stochasticity for small sample sizes)
	af1 <- af1 + sample(seq(from=-0.05, to=0.05, by=0.001), 1)
	# make sure the frequency is between 0.01 and 0.99
	if(af1 < 0.01) {
		af1 <- 0.01
	} else if (af1 > 0.99) {
		af1 <- 0.99
	}
	# change to integer on zero to one hundred scale
	af1 <- round(af1, digits=2) * 100
	af2 <- 100 - af1
	allele_freqs[[a]] <- c(rep(1, af1), rep(5, af2))
}

# output sim lists
offspring_sims <- list()
half_sibs <- list()
unrelated <- list()
# loop for simulations (this may take up to 30-60 minutes)
for(a in 1:nsims) {
	if(a %% 10 == 0) {
		writeLines(paste("Simulation ", a, " started.", sep=""))
	}

	# generate a haploid male
	male <- list()
	for(b in 1:n_loci) {
		male[[b]] <- sample(allele_freqs[[b]], 1)
	}
	male <- unlist(male)

	# sample a diploid female
	female_1 <- list()
	female_2 <- list()
	for(b in 1:n_loci) {
		female_1[[b]] <- sample(allele_freqs[[b]], 1)
		female_2[[b]] <- sample(allele_freqs[[b]], 1)
	}
	female_1 <- unlist(female_1)
	female_2 <- unlist(female_2)

	# create 1st offspring of the two parents
	offspring1 <- list()
	for(b in 1:length(female_1)) {
		offspring1[[b]] <- sort(c(male[b], sample(c(female_1[b], female_2[b]), 1)))
	}
	offspring1 <- unlist(offspring1)

	# create 2nd offspring of the two parents
	offspring2 <- list()
	for(b in 1:length(female_1)) {
		offspring2[[b]] <- sort(c(male[b], sample(c(female_1[b], female_2[b]), 1)))
	}
	offspring2 <- unlist(offspring2)

	# combine two offspring's data into one dataframe
	sim <- as.data.frame(cbind(c(paste("offspring1__", a, sep=""), paste("offspring2__", a, sep="")), rbind(offspring1, offspring2)))
	# make sure the first column of sim is character
	sim[,1] <- as.character(sim[,1])
	# add to offspring sims
	offspring_sims[[a]] <- sim

	# for half sib relationships with queen and two males mated
	# generate a second haploid male
	male2 <- list()
	for(b in 1:n_loci) {
		male2[[b]] <- sample(allele_freqs[[b]], 1)
	}
	male2 <- unlist(male2)

	# create 1st offspring of the one female and male 1
	halfsib1 <- list()
	for(b in 1:length(female_1)) {
		halfsib1[[b]] <- sort(c(male[b], sample(c(female_1[b], female_2[b]), 1)))
	}
	halfsib1 <- unlist(halfsib1)

	# create 2nd offspring of the one female, but mated to male 2
	halfsib2 <- list()
	for(b in 1:length(female_1)) {
		halfsib2[[b]] <- sort(c(male2[b], sample(c(female_1[b], female_2[b]), 1)))
	}
	halfsib2 <- unlist(halfsib2)

	# combine two halfsib's data into one dataframe
	sim <- as.data.frame(cbind(c(paste("halfsib1__", a, sep=""), paste("halfsib2__", a, sep="")), rbind(halfsib1, halfsib2)))
	# make sure the first column of sim is character
	sim[,1] <- as.character(sim[,1])
	# add to halfsib sims
	half_sibs[[a]] <- sim

	# unrelated individuals
	# sample a random genotype (x2) from the allele freqs
	random_1 <- list()
	random_2 <- list()
	for(b in 1:n_loci) {
		random_1[[b]] <- sample(allele_freqs[[b]], 2)
		random_2[[b]] <- sample(allele_freqs[[b]], 2)
	}
	random_1 <- unlist(random_1)
	random_2 <- unlist(random_2)

	# combine two randoms' data into one dataframe
	sim <- as.data.frame(cbind(c(paste("unrelated1__", a, sep=""), paste("unrelated2__", a, sep="")), rbind(random_1, random_2)))
	# make sure the first column of sim is character
	sim[,1] <- as.character(sim[,1])
	# add to unrelated sims
	unrelated[[a]] <- sim
}

# add all real and simulated data into a single data frame for relatedness estimates
total <- x
offspring_sims2 <- do.call("rbind", offspring_sims)
half_sibs2 <- do.call("rbind", half_sibs)
unrelated2 <- do.call("rbind", unrelated)
total <- rbind(total, offspring_sims2, half_sibs2, unrelated2)

# clean up objects no longer needed and clean up garbage
rm(offspring_sims2, half_sibs2, unrelated2, offspring_sims, half_sibs, unrelated)

# measure relatedness of all real and simulated individuals (this may take up to 30-60 minutes)
total_relate <- coancestry(total, lynchli=2, wang=2)

# extract point estimates of relatedness for each grouping
# real data
real_data_relatedness <- total_relate$relatedness[total_relate$relatedness$group == "C-C-",]
real_data_95CI <- total_relate$relatedness.ci95[total_relate$relatedness.ci95$group == "C-C-",]

# full sister simulations
full_sister_relatedness <- total_relate$relatedness[total_relate$relatedness$group == "ofof",]
full_sister_relatedness <- full_sister_relatedness[sapply(strsplit(full_sister_relatedness[,2], "__"), "[[", 2) == sapply(strsplit(full_sister_relatedness[,3], "__"), "[[", 2), ]

# half sister simulations
half_sister_relatedness <- total_relate$relatedness[total_relate$relatedness$group == "haha",]
half_sister_relatedness <- half_sister_relatedness[sapply(strsplit(half_sister_relatedness[,2], "__"), "[[", 2) == sapply(strsplit(half_sister_relatedness[,3], "__"), "[[", 2), ]

# unrelated individuals simulations
unrelated_relatedness <- total_relate$relatedness[total_relate$relatedness$group == "unun",]
unrelated_relatedness <- unrelated_relatedness[sapply(strsplit(unrelated_relatedness[,2], "__"), "[[", 2) == sapply(strsplit(unrelated_relatedness[,3], "__"), "[[", 2), ]

# save the work
save.image(paste(species, "__", colony_number, "__relatedness.RData", sep=""))

# make a merged data frame for boxplots of simulations
full_sister_related <- data.frame(Relationship=as.character(rep("Full Sister", nrow(full_sister_relatedness)*2)), Measurement=as.character(c(rep("Wang (2002)", nrow(full_sister_relatedness)), rep("Li et al. (1993)", nrow(full_sister_relatedness)))), Relatedness=as.numeric(c(full_sister_relatedness$wang, full_sister_relatedness$lynchli)))
half_sister_related <- data.frame(Relationship=as.character(rep("Half Sister", nrow(half_sister_relatedness)*2)), Measurement=as.character(c(rep("Wang (2002)", nrow(half_sister_relatedness)), rep("Li et al. (1993)", nrow(half_sister_relatedness)))), Relatedness=as.numeric(c(half_sister_relatedness$wang, half_sister_relatedness$lynchli)))
unrelated_sister_related <- data.frame(Relationship=as.character(rep("Unrelated", nrow(unrelated_relatedness)*2)), Measurement=as.character(c(rep("Wang (2002)", nrow(unrelated_relatedness)), rep("Li et al. (1993)", nrow(unrelated_relatedness)))), Relatedness=as.numeric(c(unrelated_relatedness$wang, unrelated_relatedness$lynchli)))

# combine the three into one
total_related <- rbind(full_sister_related, half_sister_related, unrelated_sister_related)

# make a boxplot
ggplot(total_related, aes(x=Relationship, y=Relatedness, fill=Measurement)) + geom_boxplot()

# save a pdf of the boxplot (change the name as necessary)
pdf(file=paste(species, "__", colony_number, "__simboxplot.pdf", sep=""), width=5, height=7)
ggplot(total_related, aes(x=Relationship, y=Relatedness, fill=Measurement)) + geom_boxplot()
dev.off()

# combine the real data into a nice data matrix and save the table
real_data_output <- data.frame(Ind1=as.character(real_data_relatedness$ind1.id),
							   Ind2=as.character(real_data_relatedness$ind2.id),
							   Species=as.character(species),
							   Colony=as.character(colony_number),
							   Wang=as.numeric(real_data_relatedness$wang),
							   Wang_low=as.numeric(real_data_95CI$wang.low),
							   Wang_high=as.numeric(real_data_95CI$wang.high),
							   Li=as.numeric(real_data_relatedness$lynchli),
							   Li_low=as.numeric(real_data_95CI$lynchli.low),
							   Li_high=as.numeric(real_data_95CI$lynchli.high))
# print the matrix to look at
real_data_output
# write the matrix to table
write.table(real_data_output, file=paste(species, "__", colony_number, "__relatedness.txt", sep=""), sep="\t",
			quote=F, col.names=T, row.names=F)


#!/usr/bin/env Rscript
#autor:      Joao Sollari Lopes
#local:      University of Reading, Reading, UK
#Rversion:   3.2.3
#criado:     01.03.2010
#modificado: 20.11.2017

# @arg data_file - file with summary of 'real' data (.trg)
# @arg rej_file  - file with rejection step results (.rej)
# @arg pri_file  - file with sample of priors (.pri)
reg_step <- function(data_file , rej_file , pri_file){

	#import libraries
	library(locfit)
	library(nnet)
    library(abc)

	#import Mark Beaumont's scripts
	source("make_pd2.r")

	#demographic parameters' priors (minimum and maximum values)
	mintev <- 0
	maxtev <- 10000
	minNe1 <- 0
	maxNe1 <- 12500
	minNe2 <- 0
	maxNe2 <- 40000
	minNeA <- 0
	maxNeA <- 10000
	minmig1 <- 0
	maxmig1 <- 0.0005
	minmig2 <- 0
	maxmig2 <- 0.0005

    truetev <- 5000
    trueNe1 <- 10000
    trueNe2 <- 20000
    trueNeA <- 5000
    truemig1 <- 0.0001
    truemig2 <- 0.0001

	#import the .rej file
	abc.rej <- data.matrix(read.table(rej_file))

	#import the .pri files
	priors <- data.matrix(read.table(pri_file))

	#import the .trg files
	target <- data.matrix(read.table(data_file))

	# Plot line for splitting time and save it in a .eps file
	abc.reg.tev <- makepd4(target,abc.rej[,7],abc.rej[,13:21],tol=1,rej=F,transf="log")
	abc.nn.tev <- abc(target,abc.rej[,7],abc.rej[,13:21],tol=1,method="neuralnet",transf="log")
    fname <- "../results/tev_reg.eps"
    postscript(file=fname,width=7,height=7,colormodel="rgb",horizontal=FALSE,onefile=FALSE,paper="special")
	plot(locfit(~abc.reg.tev$x,weight=abc.reg.tev$wt,xlim=c(mintev,maxtev)),main="prior (black) and posterior (blue) distributions",xlab="tev",col="blue")
	lines(locfit(~abc.nn.tev$adj.values,weight=abc.nn.tev$weights,xlim=c(mintev,maxtev)),col="purple")
	lines(locfit(~priors[,7],xlim=c(mintev,maxtev)),col="black")
    abline(v=truetev,lty=2,col="black")
	dev.off()
	print("tev done.")

	# Plot line for effective size of population 1 and save it in a .eps file
	abc.reg.ne1 <- makepd4(target,abc.rej[,8],abc.rej[,13:21],tol=1,rej=F,transf="log")
	abc.nn.ne1 <- abc(target,abc.rej[,8],abc.rej[,13:21],tol=1,method="neuralnet",transf="log")
    fname <- "../results/Ne1_reg.eps"
    postscript(file=fname,width=7,height=7,colormodel="rgb",horizontal=FALSE,onefile=FALSE,paper="special")
	plot(locfit(~abc.reg.ne1$x,weight=abc.reg.ne1$wt,xlim=c(minNe1,maxNe1)),main="prior (black) and posterior (blue) distributions",xlab="Ne1",col="blue")
	lines(locfit(~abc.nn.ne1$adj.values,weight=abc.nn.ne1$weights,xlim=c(minNe1,maxNe1)),col="purple")
	lines(locfit(~priors[,8],xlim=c(minNe1,maxNe1)),col="black")
    abline(v=trueNe1,lty=2,col="black")
	dev.off()
	print("Ne1 done.")

	# Plot line for effective size of population 2 and save it in a .eps file
	abc.reg.ne2 <- makepd4(target,abc.rej[,9],abc.rej[,13:21],tol=1,rej=F,transf="log")
	abc.nn.ne2 <- abc(target,abc.rej[,9],abc.rej[,13:21],tol=1,method="neuralnet",transf="log")
    fname <- "../results/Ne2_reg.eps"
    postscript(file=fname,width=7,height=7,colormodel="rgb",horizontal=FALSE,onefile=FALSE,paper="special")
	plot(locfit(~abc.reg.ne2$x,weight=abc.reg.ne2$wt,xlim=c(minNe2,maxNe2)),main="prior (black) and posterior (blue) distributions",xlab="Ne2",col="blue")
	lines(locfit(~abc.nn.ne2$adj.values,weight=abc.nn.ne2$weights,xlim=c(minNe2,maxNe2)),col="purple")
	lines(locfit(~priors[,9],xlim=c(minNe2,maxNe2)),col="black")
    abline(v=trueNe2,lty=2,col="black")
	dev.off()
	print("Ne2 done.")

	# Plot line for effective size of ancestor population and save it in a .eps file
	abc.reg.nea <- makepd4(target,abc.rej[,10],abc.rej[,13:21],tol=1,rej=F,transf="log")
	abc.nn.nea <- abc(target,abc.rej[,10],abc.rej[,13:21],tol=1,method="neuralnet",transf="log")
    fname <- "../results/NeA_reg.eps"
    postscript(file=fname,width=7,height=7,colormodel="rgb",horizontal=FALSE,onefile=FALSE,paper="special")
	plot(locfit(~abc.reg.nea$x,weight=abc.reg.nea$wt,xlim=c(minNeA,maxNeA)),main="prior (black) and posterior (blue) distributions",xlab="NeA",col="blue")
	lines(locfit(~abc.nn.nea$adj.values,weight=abc.nn.nea$weights,xlim=c(minNeA,maxNeA)),col="purple")
	lines(locfit(~priors[,10],xlim=c(minNeA,maxNeA)),col="black")
    abline(v=trueNeA,lty=2,col="black")
	dev.off()
	print("NeA done.")

	# Plot line for migration rate of population 1 and save it in a .eps file
	abc.reg.mig1 <- makepd4(target,abc.rej[,11],abc.rej[,13:21],tol=1,rej=F,transf="log")
	abc.nn.mig1 <- abc(target,abc.rej[,11],abc.rej[,13:21],tol=1,method="neuralnet",transf="log")
    fname <- "../results/mig1_reg.eps"
    postscript(file=fname,width=7,height=7,colormodel="rgb",horizontal=FALSE,onefile=FALSE,paper="special")
	plot(locfit(~abc.reg.mig1$x,weight=abc.reg.mig1$wt,xlim=c(minmig1,maxmig1)),main="prior (black) and posterior (blue) distributions",xlab="mig1",col="blue")
	lines(locfit(~abc.nn.mig1$adj.values,weight=abc.nn.mig1$weights,xlim=c(minmig1,maxmig1)),col="purple")
	lines(locfit(~priors[,11],xlim=c(minmig1,maxmig1)),col="black")
    abline(v=truemig1,lty=2,col="black")
	dev.off()
	print("mig1 done.")

	# Plot line for migration rate of population 2 and save it in a .eps file
	abc.reg.mig2 <- makepd4(target,abc.rej[,12],abc.rej[,13:21],tol=1,rej=F,transf="log")
	abc.nn.mig2 <- abc(target,abc.rej[,12],abc.rej[,13:21],tol=1,method="neuralnet",transf="log")
    fname <- "../results/mig2_reg.eps"
    postscript(file=fname,width=7,height=7,colormodel="rgb",horizontal=FALSE,onefile=FALSE,paper="special")
	plot(locfit(~abc.reg.mig2$x,weight=abc.reg.mig2$wt,xlim=c(minmig2,maxmig2)),main="prior (black) and posterior (blue) distributions",xlab="mig2",col="blue")
	lines(locfit(~abc.nn.mig2$adj.values,weight=abc.nn.mig2$weights,xlim=c(minmig2,maxmig2)),col="purple")
	lines(locfit(~priors[,12],xlim=c(minmig2,maxmig2)),col="black")
    abline(v=truemig2,lty=2,col="black")
	dev.off()
	print("mig2 done.")

}

args = commandArgs(trailingOnly=TRUE)
reg_step(args[1],args[2],args[3])


# Demo 3: T tests and ANOVA
# Example script for handout
# Heini Saarimäki 2.10.2014

# -----

# 1. Data preparations and descriptive statistics

# ---

# 1.2 Load and prepare your data

setwd('Z:/Desktop/Demo3')

deprivation <- read.csv('http://becs.aalto.fi/~heikkih3/deprivation.csv')

# Look at the data:

head(deprivation)		# 10 variables
dim(deprivation)		# 47 rows (i.e. 47 participants, one per row)

summary(deprivation)

# categorical variables ID, age and education are not yet coded as factors
# --> change ID, age and education to factors

deprivation$ID <- factor(deprivation$ID)
deprivation$age <- factor(deprivation$age)
deprivation$education <- factor(deprivation$education)

summary(deprivation)	# looks allright now

# look for missing values:

summary(deprivation)
boxplot(deprivation[5:7])	# take boxplot for digit symbol scores at all time points
boxplot(deprivation[8:10])	# similar boxplot for benton error scores

# no very extreme outliers (at least no wrongly coded missing values)
# so no missing values in the data, no need to correct these

# when you're happy with how the data looks like, use attach:

attach(deprivation)

# ---

# 1.3 Describe your data

# check frequencies for each group in all categorical variables:

summary(deprivation[2:4])

# you can also check for contingency tables:

table(age,education)
table(age, ht_group)
table(education, ht_group)

# table of descriptive statistics for continuous variables:

library('psych')			# load package 'psych' to get describe function
describe(deprivation[5:10])	# take descriptive statistics

# if you want to save the descriptive stats table for further use:

deprivation.descriptives <- describe(deprivation[5:10])
write.csv(deprivation.descriptives, "deprivation_descriptives.csv")

# plot continuous variables:

layout(matrix(c(1,2,3,4,5,6),3,2))		# create 6-cell display first
boxplot(digitsymbol_1, ylim=c(0,80), main="Digit symbol score, time 1")
boxplot(digitsymbol_2, ylim=c(0,80), main="Digit symbol score, time 2")
boxplot(digitsymbol_3, ylim=c(0,80), main="Digit symbol score, time 3")
boxplot(bentonerror_1, ylim=c(0,14), main="Benton error score, time 1")
boxplot(bentonerror_2, ylim=c(0,14), main="Benton error score, time 2")
boxplot(bentonerror_3, ylim=c(0,14), main="Benton error score, time 3")

# -----

# 2. Distributions

# 2.1 Normal distribution

# the original distribution of 10000 samples of random integers between 0 and 10 is uniform:

hist(runif(10000)*10, main="")  # here, runif function creates your simulated values

# calculate mean for 10000 samples of 5 integers between 0 and 10 using for-loop:

means <- numeric(10000)
for (i in 1:10000) {means[i] <- mean(runif(5)*10) }

# finally, see how these means are normally distributed:

hist(means, ylim=c(0,1600))

# let's add a normal curve:

h <- hist(means, freq=F, ylim=c(0,0.35))
x <- seq(min(h$breaks),max(h$breaks), by=0.1)
y <- dnorm(x, mean=mean(means), sd=sd(means))
lines(x,y,col="red", lty="dashed", lwd=3)

# ---

# 2.2 Student's t-distribution

# create a vector that includes simulated values between -4 and 4:

simu <- seq(-4,4,0.01)

# draw a normal distribution:

plot(simu, dnorm(simu),  type="l", lty=2, ylab="Probability density", xlab="Deviates")

# add t-distributions for different sample sizes

lines(simu, dt(simu, df=5), col="red")		# sample size 5
lines(simu, dt(simu, df=10), col="green")		# sample size 10
lines(simu, dt(simu, df=20), col="yellow")	# sample size 20
lines(simu, dt(simu, df=30), col="blue")		# sample size 30
lines(simu, dt(simu, df=100), col="grey")		# sample size 100

# notice that when sample size increases to 30 and beyond, t-distribution resembles normal distribution

# ---

# 2.3 F-distribution

# create a vector that includes simulated values between 0 and 4

simu2 <- seq(0,4,0.01)

# draw an F-distribution with degrees of freedom 1 and 1:

plot(simu2, df(simu2, 1,1), type="l", lty=2, ylab="Probability density", xlab="Deviates")

# try changing each degrees of freedom:

lines(simu2, df(simu2, 1,10), col="red")		# df1 = 1, df2 = 10
lines(simu2, df(simu2, 1,40), col="magenta")	# df1 = 1, df2 = 40

lines(simu2, df(simu2, 2,40), col="blue")		# df1 = 2, df2 = 40
lines(simu2, df(simu2, 3,40), col="cyan")		# df1 = 3, df2 = 40
lines(simu2, df(simu2, 4,40), col="green")	# df1 = 4, df2 = 40

# -----

# 3. Tests for normality

# descriptive statistics:

describe(deprivation[5:10])		# look for mean, median, skewness
# for normality, mean should be close to median and skewness close to 0

# histograms:

layout(matrix(c(1,2,3,4,5,6), 3,2))			# set layout to 6-cell
hist(digitsymbol_1, main="Digit symbol 1")
hist(digitsymbol_2, main="Digit symbol 2")
hist(digitsymbol_3, main="Digit symbol 3")
hist(bentonerror_1, main="Benton error 1")
hist(bentonerror_2, main="Benton error 2")
hist(bentonerror_3, main="Benton error 3")

# note the skewness in Benton error scores

# use quantile-quantile plots (observed values = dots should be close to line)

qqnorm(digitsymbol_1)
qqline(digitsymbol_1)

# tests for normality:

# Kolmogorov-Smirnov for Digit symbol score time 1:
ks.test(digitsymbol_1, "pnorm", mean=mean(digitsymbol_1), sd=sd(digitsymbol_1))
# suggests normality

# Shapiro-Wilk:
shapiro.test(digitsymbol_1)
# suggests normality

# ---

# 3.1 Exercises

# Question 1

layout(matrix(c(1,2,3,4,5,6), 3,2))			# set layout to 6-cell
qqnorm(digitsymbol_1, main="Digit symbol 1")
qqline(digitsymbol_1)
qqnorm(digitsymbol_2, main="Digit symbol 2")
qqline(digitsymbol_2)
qqnorm(digitsymbol_3, main="Digit symbol 3")
qqline(digitsymbol_3)
qqnorm(bentonerror_1, main="Benton error 1")
qqline(bentonerror_1)
qqnorm(bentonerror_2, main="Benton error 2")
qqline(bentonerror_2)
qqnorm(bentonerror_3, main="Benton error 3")
qqline(bentonerror_3)

# all variables look relatively normally distributed

# -

# Question 2

# Digit symbol time point 1:
ks.test(digitsymbol_1, "pnorm", mean=mean(digitsymbol_1), sd=sd(digitsymbol_1))
shapiro.test(digitsymbol_1)
# --> normally distributed

# Digit symbol time point 2:
ks.test(digitsymbol_2, "pnorm", mean=mean(digitsymbol_2), sd=sd(digitsymbol_2))
shapiro.test(digitsymbol_2)
# --> normally distributed

# Digit symbol time point 3:
ks.test(digitsymbol_3, "pnorm", mean=mean(digitsymbol_3), sd=sd(digitsymbol_3))
shapiro.test(digitsymbol_3)
# --> normally distributed

# Benton error time point 1:
ks.test(bentonerror_1, "pnorm", mean=mean(bentonerror_1), sd=sd(bentonerror_1))
shapiro.test(bentonerror_1)
# --> normally distributed

# Benton error time point 2:
ks.test(bentonerror_2, "pnorm", mean=mean(bentonerror_2), sd=sd(bentonerror_2))
shapiro.test(bentonerror_2)
# --> KS suggests normality SW does not

# Benton error time point 3:
ks.test(bentonerror_3, "pnorm", mean=mean(bentonerror_3), sd=sd(bentonerror_3))
shapiro.test(bentonerror_3)
# --> KS suggests normality SW does not

# Benton error scores at times 2 and 3 might not be normally distributed.
# For t tests and ANOVA we assume normality, so this assumption might not
# be fulfilled in these two variables.
# You might want to consider non-parametric tests.

# However, here for practice purpose, we consider all variables normally distributed.

# -----

# 4 T tests

# --

# 4.1 One-sample t test

# compare sample mean to population mean:

summary(digitsymbol_1)		# sample mean 42.23, population mean 40

# is this difference statistically significant?
# use one-sample t-test:

t.test(digitsymbol_1, mu=40)

# --> mean of our sample does not differ from population mean (t(46)=1.33, p>.05)

# ---

# 4.2 Independent samples t-test

# take t test for age differences in Digit symbol score time 1:

t.test(digitsymbol_1 ~ age)	# categorical variable after ~

# --> mean in the two age groups does not differ (t(35)=0.27, p>.05)

# without correction for equality of variances:

t.test(digitsymbol_1 ~ age, var.equal=T)

# in this case the results are the same regardless of the equality variances assumption

# visualize results:

boxplot(digitsymbol_1 ~ age, xlab="Age group", ylab="Digit symbol score")

# ---

# 4.3 Repeated-measures t-test

# take t test for Digit score differences between time points 1 and 2:

t.test(digitsymbol_1, digitsymbol_2, paired=T)

# --> no differences between time points 1 and 2 (t(46)=-0.79, p>.05)

# visualize results:

boxplot(digitsymbol_1, digitsymbol_2, ylab="Digit symbol score", xlab="Time point")

# ---

# 4.4 Exercises

# Question 3

# (i) Differences in Digit symbol baseline performance between hormone treatment groups?

t.test(digitsymbol_1 ~ ht_group)

# --> no differences (t(44)=-0.17, p>.05)

# (ii) Differences in Benton error baseline performance between hormone treatment groups?

t.test(bentonerror_1 ~ ht_group)

# --> control group has higher number of errors (t(45)=3.01, p<.01)

# use boxplot to visualize this result:

boxplot(bentonerror_1 ~ ht_group, xlab="Hormone treatment", ylab="Benton error score")

# -

# Question 4

# (i) Differences in Digit symbol scores between times 1 and 3

t.test(digitsymbol_1, digitsymbol_3, paired=T)
boxplot(digitsymbol_1, digitsymbol_3, xlab="Time point", ylab="Digit symbol score")

# --> test scores are higher at time 3 (t(46)=-6.94, p<.001)
# interpretation: test scores improve due to learning? (from 1st to 3rd measurement)

# (ii) Differences in Digit symbol scores between times 2 and 3

t.test(digitsymbol_2, digitsymbol_3, paired=T)
boxplot(digitsymbol_2, digitsymbol_3, xlab="Time point", ylab="Digit symbol score")

# --> test scores are higher at time 3 (t(46)=-10.60, p<.001)
# interpretation: test scores improve after a night of normal sleep

# -

# Question 5

# (i) Differences in Benton error scores between times 1 and 2

t.test(bentonerror_1, bentonerror_2, paired=T)

# --> no differences (t(46)=1.53, p>.05)
# interpretation: no effect of sleep deprivation on error rate (maybe learning effect cancels this effect?)

# (ii) Differences in Benton error scores between times 1 and 3

t.test(bentonerror_1, bentonerror_3, paired=T)
boxplot(bentonerror_1,bentonerror_3, xlab="Time point", ylab="Benton error score")

# --> error rates are lower at time 3 (t(46)=5.83, p<.001)
# interpretation: error rates improve due to learning? (from 1st to 3rd measurement)

# (iii) Differences in Benton error scores between times 2 and 3

t.test(bentonerror_2, bentonerror_3, paired=T)
boxplot(bentonerror_2, bentonerror_3, xlab="Time point", ylab="Benton error score")

# --> error rates are lower at time 3 (t(46)=5.99, p<.001)
# interpretation: test scores improve after a night of normal sleep

# -----

# 5. Analysis of variance

# --

# 5.1 One-way ANOVA

# Does baseline performance in Digit symbol task vary between education groups?

A1 <- aov(digitsymbol_1 ~ education)
summary(A1)

# --> no differences between groups (F(2,44)=2.16, p>.05)

model.tables(A1, "means")	# mean score for each group

# visualize results:

boxplot(digitsymbol_1 ~ education, xlab = "Education group", ylab="Digit symbol score",	main="Baseline Digit symbol score by education group", col="grey")

# get diagnostic plots:

layout(matrix(c(1,2,3,4),2,2))
plot(A1)

# --

# 5.1.1 ANOVA assumptions

# test for homogeneity of variances:

install.packages('car')
library('car')
leveneTest(digitsymbol_1 ~ education)
bartlett.test(digitsymbol_1 ~ education)

# variances are equal in Digit symbol task at time point 1

# --

# 5.1.2 Visualizing results

# familiar box plot:

boxplot(digitsymbol_1 ~ education)

# means with confidence intervals:

install.packages('gplots')
library('gplots')
plotmeans(digitsymbol_1 ~ education, ylim=c(10,70), n.label=FALSE)

# ---

# 5.2 Two-way between subject ANOVA

# Does Digit symbol baseline performance differ between age and hormone treatment groups?

A2 <- aov(digitsymbol_1 ~ age*ht_group)
summary(A2)
model.tables(A2, "means")
layout(matrix(c(1,2,3,4),2,2))
plot(A2)

# --> no age main effect (F(1,43)=0.08, p>.05)
# --> no hormone treatment main effect (F(1,43)=0.89, p>.05)
# --> age and hormone treatment have a significant interaction effect (F(1,43)=6.81, p<.05)

# visualize results:

boxplot(digitsymbol_1 ~ age*ht_group, col=c("green", "grey"))	# not very clear..

interaction.plot(age, ht_group, digitsymbol_1, type="b", col=c(1:3),
	leg.bty="o", leg.bg="beige", lwd=2, pch=c(18,24,22),
	xlab="Age group", ylab="Digit symbol score", main="Interaction plot")

# --

# 5.2.1 Type of Sums of Squares

library(car)
Anova(A2, type=3)			# ANOVA with type 3 sums of squares

# --

# 5.2.2 Post-hoc comparisons

# Test where are the differences you found in ANOVA between age and hormone treatment effects:

pairwise.t.test(digitsymbol_1, education, p.adj="none")
pairwise.t.test(digitsymbol_1, education, p.adj="bonf")	# Bonferroni correction
pairwise.t.test(digitsymbol_1, education, p.adj="holm")	# Holm correction

TukeyHSD(A1)								# Tukey's test

# ---

# 5.3 Exercises

# Question 6

# Does hormone treatment affect Digit symbol baseline differently in different education groups?

anova1 <- aov(digitsymbol_1~ht_group*education)
summary(anova1)

# --> no it does not, because interaction effect is nonsignificant (F(2,41)=1.24, p>.05)
# --> education and hormone treatment do not affect Digit symbol baseline performance (non-significant main effects)

# -

# Question 7

# Do education and age together affect the Digit symbol baseline performance?

anova2 <- aov(digitsymbol_1 ~ age*education)
summary(anova2)

# --> no they don't (only a trend towards significance in interaction effect) (F(2,41)=2.73, p>.05)
# --> age and education do not affect Digit symbol baseline performance (non-significant main effects)

# -

# Question 8

# Age, education and hormone treatment effects in Benton task baseline performance?

anova3 <- aov(bentonerror_1 ~ age*education*ht_group)
summary(anova3)

# --> only hormone treatment has a main effect on Benton task score (F(1,35)=10.21, p<.01)
# --> age and education do not have main effects
# --> age and education have an interaction effect (F(2,35)=3.66, p<.05)
# --> education and hormone treatment group have an interaction effect (F(2,35)=3.46, p>.05)

# visualize results:

boxplot(bentonerror_1 ~ ht_group, xlab="Hormone treatment group", ylab="Benton errors")
# --> hormone treatment group makes less errors than control group at baseline

interaction.plot(education,age,bentonerror_1)
# --> higher education improves scores (=less errors) in younger group, opposite trend in older group (!?)

interaction.plot(education,ht_group, bentonerror_1)
# --> higher education improves scores in hormone treatment group, slightly opposite in control group

# -----

# 6. Exercises

# Load in data:
naming <- read.csv('http://becs.aalto.fi/~heikkih3/naming_wide.csv')

# Get familiar with the data:
head(naming)
summary(naming)
dim(naming)

# It seems that the factors (sex, reading_time) are already coded as factors.
# No missing values (actually NA's from previous weeks have been removed totally!)
# Note that now 'ms' variable is split to 'ms.regular' and 'ms.regular
# - this explains why data has only half the rows from previous weeks:
# now one participant has all the variables in a single row.

detach()			# detach all previous data frames
attach(naming)		# attach naming data frame

# -

# Question 9

t.test(iq, mu=100)

# --> intelligence does not differ from the population mean (t(997)=0, p>.05).

# -

# Question 10

t.test(iq ~ sex)

# --> no sex differences in intelligence (t(961)=1.07, p>.05).

# -

# Question 11

t.test(hrs ~ sex)

# --> no sex differences in hours spent on reading (t(972)=0.29, p>.05).

# -

# Question 12

anova.1 <- aov(iq ~ reading_time)
summary(anova.1)

# --> no differences between reading categories in intelligence (F(2,995)=0.03, p>.05).

# -

# Question 13

t.test(ms.regular, ms.exception, paired=T)
mean(ms.regular)
mean(ms.exception)
boxplot(ms.regular, ms.exception, xlab="Word type (1=regular, 2=exceptional))", ylab="Time (ms)")

# --> reading times are shorter for regular words (t(997)=-68.8, p<.001).

# -

# Question 14

# Regular words:

anova.2 <- aov(ms.regular ~ reading_time)
summary(anova.2)
boxplot(ms.regular ~ reading_time)

# take a closer look at the differences using post-hoc comparison:
pairwise.t.test(ms.regular, reading_time, p.adj="bonf")
tapply(ms.regular, reading_time, mean)
tapply(ms.regular, reading_time, sd)

# --> there are differences between reading categories (based on hours spent on reading)
# in time taken to read regular words (F(2,995)=3.52, p<.05).
# --> post-hoc comparisons using a Bonferroni corrected pairwise t-test
# revealed that the differences are between the highest reading category (mean=993, sd=70.39)
# and the lowest one (mean=1009, sd=79).

# Exceptional words:

anova.3 <- aov(ms.exception ~ reading_time)
summary(anova.3)
boxplot(ms.exception ~ reading_time)

# take a closer look at the differences using post-hoc comparison:
pairwise.t.test(ms.exception, reading_time, p.adj="bonf")
tapply(ms.exception, reading_time, mean)
tapply(ms.exception, reading_time, sd)

# --> there are differences between reading categories (based on hours spent on reading)
# in time taken to read exceptional words (F(2,995)=38.99, p<.001).
# --> post-hoc comparisons using a Bonferroni corrected pairwise t-test
# revealed that all groups differ from each other:
# highest group (mean=1173, sd=76) has higher scores than the two other groups,
# and middle group (mean=1203, sd=79) has higher scores than the lowest group (mean=1228, sd=76).

# -

# Question 15

# Regular words:

anova.4 <- aov(ms.regular ~ sex * reading_time)
summary(anova.4)

# interaction plot:
interaction.plot(reading_time, sex, ms.regular)

# --> sex does not modify the effect of reading time in
# time taken to read regular words

# Exceptional words:

anova.5 <- aov(ms.exception ~ sex * reading_time)
summary(anova.5)

# interaction plot:
interaction.plot(reading_time, sex, ms.exception)

t.test(ms.exception ~ sex)
boxplot(ms.exception ~ sex)

# neither does sex modify the effect of reading time in
# time taken to read exceptional words

# -----


subroutine initad(j,nadj,madj,x,y,ntot,eps,ntri,nerror)

# Initial adding-in of a new point j.
# Called by addpt.

implicit double precision(a-h,o-z)
dimension nadj(-3:ntot,0:madj), x(-3:ntot), y(-3:ntot)
integer tau(3)
integer tadj(1000)

# Find the triangle containing vertex j.
call trifnd(j,tau,nedge,nadj,madj,x,y,ntot,eps,ntri,nerror)
if(nerror > 0) return

# If the new point is on the edge of a triangle, detach the two
# vertices of that edge from each other.  Also join j to the vertex
# of the triangle on the reverse side of that edge from the `found'
# triangle (defined by tau) -- given that there ***is*** such a triangle.
if(nedge!=0) {
        ip = nedge
        i  = ip-1
        if(i==0) i = 3 # Arithmetic modulo 3.
        call pred(k,tau(i),tau(ip),nadj,madj,ntot,nerror)
	if(nerror > 0) return
        call succ(kk,tau(ip),tau(i),nadj,madj,ntot,nerror)
	if(nerror > 0) return
        call delet(tau(i),tau(ip),nadj,madj,ntot,nerror)
	if(nerror > 0) return
        if(k==kk) call insrt(j,k,nadj,madj,x,y,ntot,nerror,eps)
	if(nerror > 0) return
}

# Join the new point to each of the three vertices.
do i = 1,3 {
	call insrt(j,tau(i),nadj,madj,x,y,ntot,nerror,eps)
	if(nerror > 0) return
}
if(j==580) {
    nj = nadj(j,0)
    do jc = 1,nj {
        tadj(jc) = nadj(j,jc)
    }
    call intpr("Initial adjacency list of point 580:",-1,tadj,nj)
}

return
end


# Listing 17.7 은 가급적 돌리지 말것. 많은 시간이 소요되며 잘 반영되지 않음.
# 본 코드를 수행하기 위해서는
# 원본 데이터를 아래와 같이 변형해야 함 head(kbo)참조.

# 원 데이터는 두팀이 한row에 있으며 모든 자료가 같은 row를 차지하고 있음.
# 팀명으로 홈 어웨이를 구분(hOrA)하고 같은 row의 자료를 아래로 Rbind하여 위 데이터셋을 만들어야 함.
# 마지막에 있는 코드는 각 분류방법에 대한 정확도와 민감도를 나타냄. 가장 좋은 것은 역시 랜덤포레스트.

# install.packages(c("rpart", "party", "randomForest", "e1071", "rpart.plot", "readr"))
# install.packages("partykit")
library(rpart)
library(party)
library(randomForest)
library(e1071)
library(rpart.plot)
library(readr)
library(readxl)

kbo <- read_csv("kbolist_re.csv")
head(kbo)

df <- kbo[-1]
df$winOrloss <- factor(df$winOrloss, levels=c(1,0),
                   labels=c("win", "loss"))

set.seed(1234)
train <- sample(nrow(df), 0.7*nrow(df))
df.train <- df[train,]
df.validate <- df[-train,]
table(df.train$winOrloss)
# win loss
# 2443 2422
table(df.validate$winOrloss)
# win loss
# 1032 1053

# Listing 17.2 - Logistic regression with glm()
fit.logit <- glm(winOrloss~., data=df.train, family=binomial())
summary(fit.logit)
##############################################################################################
# Call:
#   glm(formula = winOrloss ~ ., family = binomial(), data = df.train)
#
# Deviance Residuals:
#   Min       1Q   Median       3Q      Max
# -2.0658  -1.0863  -0.5907   1.0797   2.0736
#
# Coefficients:
#   Estimate Std. Error z value Pr(>|z|)
# (Intercept)  0.307633   0.160909   1.912 0.055897 .
# hOrA        -0.414391   0.060146  -6.890 5.59e-12 ***
#   inning      -0.232867   0.098284  -2.369 0.017821 *
#   tajasu       0.214432   0.054491   3.935 8.31e-05 ***
#   tugusu      -0.000244   0.003561  -0.069 0.945372
# tasu        -0.160913   0.047018  -3.422 0.000621 ***
#   anta         0.013630   0.033197   0.411 0.681376
# fourball    -0.181790   0.057509  -3.161 0.001572 **
#   homerun      0.083115   0.044310   1.876 0.060690 .
# samjin      -0.009144   0.016582  -0.551 0.581337
# siljum       0.130653   0.049832   2.622 0.008744 **
#   jacheck     -0.041126   0.048831  -0.842 0.399676
# score       -0.096739   0.008490 -11.395  < 2e-16 ***
#   ---
#   Signif. codes:  0 ‘***’ 0.001 ‘**’ 0.01 ‘*’ 0.05 ‘.’ 0.1 ‘ ’ 1
#
# (Dispersion parameter for binomial family taken to be 1)
#
# Null deviance: 6744.2  on 4864  degrees of freedom
# Residual deviance: 6326.0  on 4852  degrees of freedom
# AIC: 6352
#
# Number of Fisher Scoring iterations: 4
##############################################################################################
prob <- predict(fit.logit, df.validate, type="response")
logit.pred <- factor(prob > .5, levels=c(FALSE, TRUE),
                     labels=c("win", "loss"))
logit.perf <- table(df.validate$winOrloss, logit.pred,
                    dnn=c("Actual", "Predicted"))
logit.perf
#prob
#     Predicted
# Actual win loss
#  win  686  346
#  loss 359  694
##############################################################################################

# Listing 17.3 - Creating a classical decision tree with rpart()
library(rpart)
set.seed(1234)
dtree <- rpart(winOrloss ~ ., data=df.train, method="class",
               parms=list(split="information"))
dtree$cptable
plotcp(dtree)

dtree.pruned <- prune(dtree, cp=.0125)

library(rpart.plot)
#?prp()
prp(dtree.pruned, type = 2, extra = 104,
    fallen.leaves = TRUE, main="Decision Tree")

dtree.pred <- predict(dtree.pruned, df.validate, type="class")
dtree.perf <- table(df.validate$winOrloss, dtree.pred,
                    dnn=c("Actual", "Predicted"))
dtree.perf

# Listing 17.4 - Creating a conditional inference tree with ctree()
library(party)
fit.ctree <- ctree(winOrloss~., data=df.train)
plot(fit.ctree, main="Conditional Inference Tree")

ctree.pred <- predict(fit.ctree, df.validate, type="response")
ctree.perf <- table(df.validate$winOrloss, ctree.pred,
                    dnn=c("Actual", "Predicted"))
ctree.perf

library(partykit)
plot(as.party(dtree.pruned))

# Listing 17.5 - Random forest
library(randomForest)
set.seed(1234)
fit.forest <- randomForest(winOrloss~., data=df.train,
                           na.action=na.roughfix,
                           importance=TRUE)
fit.forest
importance(fit.forest, type=2)
forest.pred <- predict(fit.forest, df.validate)
forest.perf <- table(df.validate$winOrloss, forest.pred,
                     dnn=c("Actual", "Predicted"))
forest.perf

varImpPlot(fit.forest)

# Listing 17.6 - A support vector machine
library(e1071)
set.seed(1234)
fit.svm <- svm(winOrloss~., data=df.train)
fit.svm
svm.pred <- predict(fit.svm, na.omit(df.validate))
svm.perf <- table(na.omit(df.validate)$winOrloss,
                  svm.pred, dnn=c("Actual", "Predicted"))
svm.perf

# Listing 17.7 Tuning an RBF support vector machine (this can take a while)
# many time spend..... if you have a dull, please cancel this process.
#set.seed(1234)
#tuned <- tune.svm(winOrloss~., data=df.train,
#                  gamma=10^(-6:1),
#                  cost=10^(-10:10))
#tuned
#fit.svm <- svm(winOrloss~., data=df.train, gamma=.01, cost=1)
#svm.pred <- predict(fit.svm, na.omit(df.validate))
#svm.perf <- table(na.omit(df.validate)$winOrloss,
#                  svm.pred, dnn=c("Actual", "Predicted"))
#svm.perf

# Listing 17.8 Function for assessing binary classification accuracy
performance <- function(table, n=2){
  if(!all(dim(table) == c(2,2)))
    stop("Must be a 2 x 2 table")
  tn = table[1,1]
  fp = table[1,2]
  fn = table[2,1]
  tp = table[2,2]
  sensitivity = tp/(tp+fn)
  specificity = tn/(tn+fp)
  ppp = tp/(tp+fp)
  npp = tn/(tn+fn)
  hitrate = (tp+tn)/(tp+tn+fp+fn)
  result <- paste("Sensitivity = ", round(sensitivity, n) ,
                  "\nSpecificity = ", round(specificity, n),
                  "\nPositive Predictive Value = ", round(ppp, n),
                  "\nNegative Predictive Value = ", round(npp, n),
                  "\nAccuracy = ", round(hitrate, n), "\n", sep="")
  cat(result)
}

# Listing 17.9 - Performance of breast cancer data classifiers
performance(dtree.perf)
performance(ctree.perf)
performance(forest.perf)
performance(svm.perf)