require(Matrix)
holidays.Sweden <- as.Date(c("2020-04-10","2020-04-13","2020-05-01","2020-05-21","2020-06-19","2020-06-20"))

##
# sample reported deaths after rep.day (if Inf) predicit day
# assuming Multionimal distribution
#
#  samples  - (int) how many samples should be done for each N (cheap)
#  deaths   - (N x 1) true number of deaths each day
#  P        - (N x N) matrix of probabilites (only upper triangular part relevant)
#  Reported - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
#  alpha    - (N x 1) stepsizes
#  true.dat  - (int) how many days after recording should one sample
##
sample.deaths <- function(samples, deaths, P, Reported, alpha, true.day = 0,
                          use.prior=FALSE,
                          prior=NULL){

  N <- length(deaths)
  acc <- rep(0,N)
  for(i in 1:N){
    if(i > true.day){
      P_i         = P[i,i:N]
      Reported_i  = Reported[i,i:N]
      index = is.na(Reported_i)==F
      P_i         = P_i[index]
      Reported_i  = Reported_i[index]
      lik_i <- loglikDeathsGivenProb(deaths[i], P_i, Reported_i)
      if(use.prior==T)
        lik_i <- lik_i + prior(deaths[i], i)
      for(j in 1:samples){
        death_star <- sample((deaths[i]-alpha[i]):(deaths[i]+alpha[i]), 1)
        lik_star <- loglikDeathsGivenProb(death_star,  P_i, Reported_i)
        if(use.prior==T)
          lik_star <- lik_star + prior(death_star, i)
        if(is.nan(lik_star))
          next

        if(log(runif(1)) < lik_star-lik_i){
          lik_i = lik_star
          deaths[i] <- death_star
          acc[i] = acc[i] + 1
        }
      }
    }
  }
  return(list(deaths=deaths, acc = acc))
}

##
# sample reported deaths after rep.day (if Inf) predicit day using Beta binomial dist
#
#  samples  - (int) how many samples should be done for each N (cheap)
#  deaths   - (N x 1) true number of deaths each day
#  alpha     - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  beta      - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  Reported - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
#  alpha.MCMC    - (N x 1) stepsizes
#  true.dat  - (int) how many days after recording should one sample
#  prior     - (function) should  return log of prior density (use N and i)
##
sample.deathsBB <- function(samples, deaths, alpha, beta, Reported, alpha.MCMC, true.day = 0 ,
                            use.prior=FALSE,
                            prior=NULL){

  N <- length(deaths)
  acc <- rep(0,N)
  for(i in 1:N){
    if(i > true.day){
      alpha_i     = alpha[i,i:N]
      beta_i      = beta[i,i:N]
      Reported_i  = Reported[i,i:N]
      index = is.na(Reported_i)==F
      alpha_i  = alpha_i[index]
      beta_i  = beta_i[index]
      Reported_i  = Reported_i[index]
      if(length(Reported_i)>0){
        lik_i <- loglikDeathsGivenProbBB(deaths[i],alpha_i , beta_i, Reported_i)
      }else{
        lik_i <- 0
      }
      if(use.prior==T)
        lik_i <- lik_i + prior(deaths[i], i)
      for(j in 1:samples){
        death_star <- sample((deaths[i]-alpha.MCMC[i]):(deaths[i]+alpha.MCMC[i]), 1)
        if(length(Reported_i)>0){
          lik_star <- loglikDeathsGivenProbBB(death_star, alpha_i , beta_i, Reported_i)
        }else{
          lik_star <- 0
        }
        if(use.prior==T)
          lik_star <- lik_star + prior(death_star, i)
        if(is.nan(lik_star))
          next
        if(log(runif(1)) < lik_star-lik_i){
          lik_i = lik_star
          deaths[i] <- death_star
          acc[i] = acc[i] + 1
        }
      }
    }
  }
  return(list(deaths=deaths, acc = acc))
}

##
# sample reported deaths after rep.day (if Inf) predicit day using Beta binomial dist
#
#  samples  - (int) how many samples should be done for each N (cheap)
#  deaths   - (N x 1) true number of deaths each day
#  alpha     - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  beta      - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  Reported - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
#  alpha.MCMC    - (N x 1) stepsizes
#  true.dat  - (int) how many days after recording should one sample
#  prior     - (function) should  return log of prior density (use N and i)
##
sample.deathsBB_bi <- function(samples, deaths, alpha, beta, Reported, alpha.MCMC, true.day = 0 ,
                               p_null,
                               pi_null,
                            use.prior=FALSE,
                            prior=NULL){

  N <- length(deaths)
  acc <- rep(0,N)
  for(i in 1:N){
    if(i > true.day){
      alpha_i     = alpha[i,i:N]
      beta_i      = beta[i,i:N]
      Reported_i  = Reported[i,i:N]
      index = is.na(Reported_i)==F
      alpha_i  = alpha_i[index]
      beta_i  = beta_i[index]
      Reported_i  = Reported_i[index]
      if(length(Reported_i)>0){
        lik_i <- loglikDeathsGivenProbBB_bi(deaths[i],alpha_i , beta_i, pi_null,p_null, Reported_i)
      }else{
        lik_i <- 0
      }
      if(use.prior==T)
        lik_i <- lik_i + prior(deaths[i], i)
      for(j in 1:samples){
        death_star <- sample((deaths[i]-alpha.MCMC[i]):(deaths[i]+alpha.MCMC[i]), 1)
        if(length(Reported_i)>0){
          lik_star <- loglikDeathsGivenProbBB_bi(death_star, alpha_i , beta_i, pi_null,p_null, Reported_i)
        }else{
          lik_star <- 0
        }
        if(use.prior==T)
          lik_star <- lik_star + prior(death_star, i)
        if(is.nan(lik_star))
          next
        if(log(runif(1)) < lik_star-lik_i){
          lik_i = lik_star
          deaths[i] <- death_star
          acc[i] = acc[i] + 1
        }
      }
    }
  }
  return(list(deaths=deaths, acc = acc))
}

##
# sample reported deaths after rep.day (if Inf) predicit day using Beta binomial dist
#
#  samples  - (int) how many samples should be done for each N (cheap)
#  deaths   - (N x 1) true number of deaths each day
#  alpha     - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  beta      - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  Reported - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
#  alpha.MCMC    - (N x 1) stepsizes
#  true.dat  - (int) how many days after recording should one sample
#  prior     - (function) should  return log of prior density (use N and i)
##
sample.deathsBB_negbin <- function(samples, deaths, alpha, beta, Reported, alpha.MCMC, true.day = 0 ,
                                   theta,
                                   phi){

  N <- length(deaths)
  acc <- rep(0,N)
  for(i in 1:N){
    if(i > true.day){
      alpha_i     = alpha[i,i:N]
      beta_i      = beta[i,i:N]
      Reported_i  = Reported[i,i:N]
      index = is.na(Reported_i)==F
      alpha_i  = alpha_i[index]
      beta_i  = beta_i[index]
      Reported_i  = Reported_i[index]
      if(length(Reported_i)>0){
        lik_i <- loglikDeathsGivenProbBB(deaths[i],alpha_i , beta_i, Reported_i)
      }else{
        lik_i <- 0
      }
      lik_i <- lik_i + dnegbin(deaths[i], theta[i],phi)
      for(j in 1:samples){
        death_star <- sample((deaths[i]-alpha.MCMC[i]):(deaths[i]+alpha.MCMC[i]), 1)
        if(length(Reported_i)>0){
          lik_star <- loglikDeathsGivenProbBB(death_star, alpha_i , beta_i, Reported_i)
        }else{
          lik_star <- 0
        }
        lik_star <- lik_star + dnegbin(death_star, theta[i],phi)
        if(is.nan(lik_star))
          next
        if(log(runif(1)) < lik_star-lik_i){
          lik_i = lik_star
          deaths[i] <- death_star
          acc[i] = acc[i] + 1
        }
      }
    }
  }
  return(list(deaths=deaths, acc = acc))
}

##
# sample reported deaths after rep.day (if Inf) predicit day using Beta binomial dist
#
#  samples  - (int) how many samples should be done for each N (cheap)
#  deaths   - (N x 1) true number of deaths each day
#  alpha     - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  beta      - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  Reported - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
#  alpha.MCMC    - (N x 1) stepsizes
#  lag  - (int) how many days after recording should one sample
#  prior     - (function) should  return log of prior density (use N and i)
##
sample.deathsBB.lag <- function(samples, deaths, alpha, beta, Reported, alpha.MCMC, lag ,
                                dates,
                            use.prior=FALSE,
                            prior=NULL){

  N <- length(deaths)
  acc <- rep(0,N)
  holidays <- weekdays(dates)%in%c("Sunday","Saturday") | c(dates)%in%c(holidays.Sweden)
  for(i in 1:N){
    k <- which.max(cumsum(holidays[i:N]==F)==(lag+1))-1
    if(k==0){
      alpha_i     = alpha[i,i:N]
      beta_i      = beta[i,i:N]
      Reported_i  = Reported[i,i:N]
      index = is.na(Reported_i)==F
      alpha_i  = alpha_i[index]
      beta_i  = beta_i[index]
      Reported_i  = Reported_i[index]
      if(length(Reported_i)>0){
        lik_i <- loglikDeathsGivenProbBB(deaths[i],alpha_i , beta_i, Reported_i)
      }else{
        lik_i <- 0
      }
      if(use.prior==T)
        lik_i <- lik_i + prior(deaths[i], i)
      for(j in 1:samples){
        death_star <- sample((deaths[i]-alpha.MCMC[i]):(deaths[i]+alpha.MCMC[i]), 1)
        if(length(Reported_i)>0){
          lik_star <- loglikDeathsGivenProbBB(death_star, alpha_i , beta_i, Reported_i)
        }else{
          lik_star <- 0
        }
        if(use.prior==T)
          lik_star <- lik_star + prior(death_star, i)
        if(is.nan(lik_star))
          next
        if(log(runif(1)) < lik_star-lik_i){
          lik_i = lik_star
          deaths[i] <- death_star
          acc[i] = acc[i] + 1
        }
      }
    }
  }
  return(list(deaths=deaths, acc = acc))
}

##
# fill report using binimoal beta
#
#
##
fill.ReportBB <- function(deaths, Alpha, Beta, Reported, maxusage.day){
  N <- length(deaths)
  N_2 <- dim(Alpha)[2]
  for(i in 1:N){
      Alpha_i         = Alpha[i,i:N_2]
      Beta_i          = Beta[i,i:N_2]
      Reported_i  = Reported[i,i:N_2]
      index = is.na(Reported_i)==T
      for(j in min(which(index)):length(Reported_i)){
        p <- rbeta(1, Alpha_i[j], Beta_i[j])

        if(i> maxusage.day){
          if(j==1){
            Reported_i[j] <- rbinom(1, size=deaths[i] , prob = p)
          }else if(deaths[i] - Reported_i[j-1]>0){
            Reported_i[j] <- rbinom(1, size=deaths[i] - Reported_i[j-1], prob = p) + Reported_i[j-1]
          }else{Reported_i[j] <- Reported_i[j-1]}
        }else{
        Reported_i[j] <- Reported_i[j-1]
        }
      }
    Reported[i,i:N_2] = Reported_i
  }
  return(Reported)
}

##
# fill report using binimoal beta up to lag
#
# deaths    - (N x 1) number of estimated deaths
# Alpha     - (N x m) parameter for Beta Bin
# Beta      - (N x m) parameter for Beta Bin
# dates     - (m x 1) dates reported
# Reported  - (N x m) reported so far and na for non reported
# lag       - (int)   lag = 0 is only first non-holiday, lag = 1 two non holidays
##
fill.ReportBB.lag <- function(deaths, Alpha, Beta, Reported, dates, lag){
  N <- length(deaths)
  m <- dim(Alpha)[2]
  holidays <- weekdays(dates)%in%c("Sunday","Saturday") | c(dates)%in%c(holidays.Sweden)
  for(i in 1:N){
    k <- which.max(cumsum(holidays[i:m]==F)==(lag+1))-1
    if(k==0)
      k <- m - i
    Alpha_i         = Alpha[i,i:(i+k)]
    Beta_i          = Beta[i,i:(i+k)]
    Reported_i  = Reported[i,i:(i+k)]
    index = is.na(Reported_i)==T
    if(sum(index)==0)
      next

    for(j in min(which(index)):length(Reported_i)){
      p <- rbeta(1, Alpha_i[j], Beta_i[j])
      if(deaths[i] - Reported_i[j-1]>0){
          if(j==(k+1)){
            Reported_i[j] <- deaths[i]
          }else{
            Reported_i[j] <- rbinom(1, size=deaths[i] - Reported_i[j-1], prob = p) + Reported_i[j-1]
          }
        }else{Reported_i[j] <- Reported_i[j-1]}

    }
    Reported[i,i:(i+k)] = Reported_i
  }
  return(Reported)
}

##
# fill report using binimoal beta post lag
#
# deaths       - (N x 1) number of estimated deaths
# Alpha        - (N x m) parameter for Beta Bin
# Beta         -  (N x m) parameter for Beta Bin
# dates        - (m x 1) dates reported
# Reported     - (N x m) reported so far and na for non reported
# lag          - (int)   lag = 0 is only first non-holiday, lag = 1 two non holidays
##
fill.ReportBB.postlag <- function(deaths, Alpha, Beta, Reported, dates, lag){
  N <- length(deaths)
  m <- dim(Alpha)[2]
  holidays <- weekdays(dates)%in%c("Sunday","Saturday") | c(dates)%in%c(holidays.Sweden)
  for(i in 1:N){
    k <- which.max(cumsum(holidays[i:m]==F)==(lag+1))-1
    if(k==0)
      k <- m - i
    if(k+i >= m)
      next

    death_i <- deaths[i] +  Reported[i,(i+k)]
    # use i instead of i+1 to get the latest size
    Alpha_i         = Alpha[i,(i+k):m]
    Beta_i          = Beta[i,(i+k):m]
    Reported_i  = Reported[i,(i+k):m]
    index = is.na(Reported_i)==T
    if(sum(index)==0)
      next

    for(j in min(which(index)):length(Reported_i)){
      p <- rbeta(1, Alpha_i[j], Beta_i[j])
      rep_prev <- Reported_i[j-1]
      if( death_i - rep_prev>0 ){
        Reported_i[j] <- rbinom(1, size=death_i - rep_prev, prob = p) + rep_prev
      }else{Reported_i[j] <- rep_prev}

    }
    Reported[i,(i+k):m] = Reported_i
  }
  return(Reported)
}

##
# fill report using binimoal beta post lag
#
# deaths       - (N x 1) number of estimated deaths
# Alpha        - (N x m) parameter for Beta Bin
# Beta         -  (N x m) parameter for Beta Bin
# dates        - (m x 1) dates reported
# Reported     - (N x m) reported so far and na for non reported
# lag          - (int)   lag = 0 is only first non-holiday, lag = 1 two non holidays
##
fill.ReportBB.postlag_bi <- function(deaths,
                                     Alpha,
                                     Beta,
                                     pi_null,
                                     p_null,
                                     Reported,
                                     dates,
                                     lag){
  N <- length(deaths)
  m <- dim(Alpha)[2]
  holidays <- weekdays(dates)%in%c("Sunday","Saturday") | c(dates)%in%c(holidays.Sweden)
  for(i in 1:N){
    k <- which.max(cumsum(holidays[i:m]==F)==(lag+1))-1
    if(k==0)
      k <- m - i
    if(k+i >= m)
      next

    death_i <- deaths[i] +  Reported[i,(i+k)]
    # use i instead of i+1 to get the latest size
    Alpha_i         = Alpha[i,(i+k):m]
    Beta_i          = Beta[i,(i+k):m]
    Reported_i  = Reported[i,(i+k):m]
    index = is.na(Reported_i)==T
    if(sum(index)==0)
      next

    for(j in min(which(index)):length(Reported_i)){
      if(runif(1) >pi_null){
        p <- rbeta(1, Alpha_i[j], Beta_i[j])
      }else{
        p <- p_null
      }
      rep_prev <- Reported_i[j-1]
      if( death_i - rep_prev>0 ){
        Reported_i[j] <- rbinom(1, size=death_i - rep_prev, prob = p) + rep_prev
      }else{Reported_i[j] <- rep_prev}

    }
    Reported[i,(i+k):m] = Reported_i
  }
  return(Reported)
}

##
# fill report
#
#
##
fill.Report <- function(deaths, P, Reported, maxusage.day){
  N <- length(deaths)
  N_2 <- dim(P)[2]
  for(i in 1:N){
      P_i         = P[i,i:N_2]
      Reported_i  = Reported[i,i:N_2]
      index = is.na(Reported_i)==T
      for(j in min(which(index)):length(Reported_i)){
        Reported_i[j] <- rbinom(1, size=deaths[i] - Reported_i[j-1], prob =  P_i[j]) + Reported_i[j-1]
        if(i> maxusage.day){
          Reported_i[j] <- rbinom(1, size=deaths[i] - Reported_i[j-1], prob = P_i[j]) + Reported_i[j-1]
        }else{
          Reported_i[j] <- Reported_i[j-1]
        }
      }
      Reported[i,i:N_2] = Reported_i
  }
  return(Reported)
}
###
#
# posterior sampling of deaths given Prob vec
#
# P           - (N x N) probability matrix over probability of detecting reminder
# Reported    - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
# Predict.day - (int) which day to of reporting to predict
# sim         - (2 x 1) MCMC samples inner loop and outer loop
# alpha       - (N x 1) stepsizes
#
###
death.givenProb <- function(P, Reported ,Predict.day = Inf,sim=c(2000,10), alpha = NULL){

  N <- dim(P)[1]
  if(is.null(alpha))
    alpha <- rep(4, N)

  deaths <- matrix(NA, nrow=sim[1], ncol = N)
  deaths_est <- apply(Reported, 1, max, na.rm=T)

  burnin = ceiling(0.3*sim[1])
  for(i in 1:(sim[1] + burnin -1)){

    res <- sample.deaths(sim[2],deaths_est, P, Reported, alpha,rep.day=Predict.day)
    deaths_est <- res$deaths
    if(i < burnin){
      alpha[res$acc/sim[2] > 0.3] <- alpha[res$acc/sim[2] > 0.3] +1
      alpha[res$acc/sim[2] < 0.3] <- alpha[res$acc/sim[2] < 0.3] -1
      alpha[alpha<1] <- 1
    }else{
      deaths[i - burnin +1,] = deaths_est
    }
  }
  return(deaths)
}

##
# log liklihood of obseving report given death and prob
#  deaths  - (int) true number of deaths
#  p       - (n x 1) probability of report
#  report  - (n x 1) reported deaths culmative each date
##
loglikDeathsGivenProb <- function(death, p, report){

  if(death < max(report,na.rm=T))
    return(-Inf)
  n <- length(p)
  if(is.na(report[1])){
    #we dont have data from day one

    ndeaths <- diff(report[is.na(report)==F])
    report_adj <- report[1:(n-1)]
    report_adj <- report_adj[is.na(report_adj)==F]
  }else{
    ndeaths <- c(report[1],diff(report))
    if(n>1){
      report_adj <- c(0, report[1:(n-1)])
    }else{
      report_adj <- 0
    }
  }

  ndeaths[ndeaths <0 ] = 0

  return(sum(dbinom(ndeaths,  death - report_adj,  prob = p[is.na(p)==F], log=T )))
}

##
# build holiday covariates vector for holiday
#
##
#  holidays - (N x 1) true if day is holiday false else
##
buildXholiday <- function(N,holidays){
  ##
  # base matrix
  ##
  index_base <- t(matrix(rep(holidays,N),ncol = N))

  index_base <- index_base[upper.tri(index_base,diag=T)]
  index_base <- 1* index_base
  #sparse matrix
  i_base <- 1:length(index_base)
  i_base <- i_base[index_base==1]
  return(sparseMatrix(j=rep(1,length(i_base)),i=i_base,  dims=c(length(index_base), 1)))
}

buildXall <- function(N){
  ##
  # base matrix
  ##
  index_base <- t(matrix(1:N^2,ncol = N, nrow=N))

  index_base <- index_base[upper.tri(index_base,diag=T)]
  #sparse matrix
  i_base <- 1:length(index_base)
  j_base <- 1:length(index_base)
  return(sparseMatrix(j=j_base,i=j_base))
}

##
# build day mixed effect matrix
#  nDayEffects - number of speical days effect (1- first day, 2- first + second day)
#  N - number of days
#  nDayEffects - how many day covariate effect to create
##
buildXmixeddayeffect <- function(N, nDayEffects = 1){

  index = rep(0,N)
  index[nDayEffects] = 1
  index_days <- toeplitz(index)
  index_days <- index_days[upper.tri(index_days,diag=T)]
  j_ <-  1:sum(index_days[index_days>0])
  i_base <- 1:length(index_days)
  i_ <-  i_base[index_days>0]
  return(sparseMatrix(i=i_,j=j_ , dims=c(length(index_days), sum(index_days[index_days>0])))  )
}
##
# building an X matrix such that each day has a fixed effect
#
##
buildXday <- function(N){
  ##
  # base matrix
  ##
  index_base <- t(matrix(rep(1:N,N),ncol = N))
  index_base <- index_base[upper.tri(index_base,diag=T)]
  #sparse matrix
  j_base <- 1:length(index_base)
  #adding mean effect
  j_ <- c(index_base, rep(N+1,length(index_base)))
  i_ <- c(j_base, 1:length(index_base))
  #day effects
  if(nDayEffects > 0){
    index_days <- toeplitz(c( N+1 + (1:nDayEffects), rep(0,N -nDayEffects)))
    index_days <- index_days[upper.tri(index_days,diag=T)]
    j_ <- c(j_, index_days[index_days>0])
    i_ <- c(i_, j_base[index_days>0])
  }
  return(sparseMatrix(i=i_,j=j_))
}
##
#' split data before and after lag
#'
#'
#' Reported -  (N x N) number of reported deaths
#' dates    -  (N x 1) dates of reporting
#' lag      -  (int)   how long after to report (lag = 0 is only first non-holiday, lag = 1 two non holidays,
#' dates_not_reported - (N x 1) dayes with no reporting in
#'
##
splitlag <- function(Reported_T, dates, lag,
                     dates_not_reported){

  holidays <- weekdays(dates)%in%c("Sunday","Saturday") | c(dates)%in%c(holidays.Sweden)
  N <- dim(Reported_T)[1]
  N2 <- dim(Reported_T)[2]
  Reported_O = Reported_T
  day_completed <- rep(1,N)
  for(i in 1:N){
    # which day is lag +1 number of days before holiday
    k <- which.max(cumsum(holidays[i:N2]==F)==(lag+1))-1
    if(k==0)
      k <- N2-i
    Reported_O[i, i:min(i+k, N2)] <- NA
    if(k+i<N2){
      Reported_T_temp =  Reported_T[i,  (i+k)]
      if(is.na(Reported_T_temp)) # has been no reporting those days
        Reported_T_temp = 0
      Reported_O[i,  (i+k+1):N2] <- Reported_O[i,  (i+k+1):N2] -Reported_T_temp
      Reported_T[i, (i+k+1):N2]  <- NA
    }else{
      day_completed[i] <- 0
    }
  }
  return(list(Reported_O = Reported_O,
              Reported_T = Reported_T,
              day_completed = day_completed,
              dates_not_reported = dates_not_reported) )
}

##
#' removes observations on weekend, monday, and put diagonal to zero
#' putting the data into the style of current reporting (2020-11)
#'  @param reports      - (N x N) reported cumlative deaths
#'  @param report_dates - (N x 1) dates the data was reported
#'
report_clean <- function(reports, report_dates){
  diag(reports) <- 0
  holidays <- weekdays(report_dates)%in%c("Monday","Sunday","Saturday") | c(report_dates)%in%c(holidays.Sweden)
  reports[,holidays] <- NA
  return(reports)
}

##
#' cereate new_cases matrix
#'  @param reports - (N x N) reported cumlative deaths
##
newCases <- function(reports){
  newreport <- reports
  #reports_temp[is.na(reports_temp)]=0
  ##
  # ugly fix
  ##

  reports_temp <- reports

  for(i in 1:dim(reports_temp)[1]){
    reports_temp[i,1:(i-1)]=0
    if(is.na(reports_temp[i,i]))
      reports_temp[i,i] <- 0
    for(j in i:dim(reports_temp)[2]){
      if(is.na(reports_temp[i,j]))
        reports_temp[i,j]= reports_temp[i,j-1]
    }

  }
  reports_temp <- cbind(0,reports_temp)
  newreport <- t(diff(t(reports_temp)))
  newreport[is.na(reports)]=NA
  #newreport[upper.tri(newreport)] <- diff.report[upper.tri(diff.report,T)]
  newreport[newreport<0 & is.na(newreport)==F]=0 #fake
  return(newreport)
}

##
# how m
# transforms data,
#' we remove data if negative..
#'
#'  @param deaths  - (N x 1) how many has died (thruth)
#'  @param reports - (N x N) reported cumlative deaths
#'  @param maxusage.day - (int) only use data up to maxusage.days i.e.
#'                       reports[i,i + maxusage.day - 1]
#' @return list
#'             - death.remain how many remains to be reported
#'             - report.new   how many are reported
##
newDeaths <-function(deaths, reports,maxusage.day = -1){
  newreport <- reports
  #reports_temp[is.na(reports_temp)]=0
  ##
  # ugly fix
  ##

  reports_temp <- reports

  for(i in 1:dim(reports_temp)[1]){
    reports_temp[i,1:(i-1)]=0
    if(is.na(reports_temp[i,i]))
      reports_temp[i,i] <- 0
    for(j in i:dim(reports_temp)[2]){
      if(is.na(reports_temp[i,j]))
        reports_temp[i,j]= reports_temp[i,j-1]
    }

  }
  reports_temp <- cbind(0,reports_temp)
  newreport <- t(diff(t(reports_temp)))
  newreport[is.na(reports)]=NA
  #newreport[upper.tri(newreport)] <- diff.report[upper.tri(diff.report,T)]
  newreport[newreport<0 & is.na(newreport)==F]=0 #fake
  death.rem <- diag(deaths)

  cumsum_report <- t(apply(newreport, 1, function(x){x[is.na(x)==F]=cumsum(x[is.na(x)==F]); return(x)}))
  dr<-deaths - cumsum_report
  N <- length(deaths)
  dr <- dr[1:(N-1),1:(N-1)]
  death.rem[upper.tri(newreport)] <- dr[upper.tri(dr,T)]
  death.rem[lower.tri(death.rem)] <- NA
  diag(death.rem)[is.na(diag(reports))] <- NA

  ##
  # ugly fix 2
  ##
  for(k in 1:dim(death.rem)[1]){
    index <- is.na(death.rem[k,]) & is.na(newreport[k,])==F
    if(sum(index)>0){
      death.rem[k, index] <- deaths[k]
    }
  }
  if(maxusage.day>0){
    for(i in 1:length(death.rem)){
      if(i + maxusage.day - 1 < N ){
        death.rem[i,(i+maxusage.day):N] = NA
        newreport[i,(i+maxusage.day):N] = NA
      }
    }
  }
  #removing small bugs in reporting
  newreport[death.rem <0 & is.na(death.rem)==F] = 0
  death.rem[death.rem <0& is.na(death.rem)==F] = 0
  index <- (death.rem< newreport) &  is.na(death.rem) ==F & is.na(newreport)==F
  newreport[index] =death.rem[index]
  return(list(death.remain = death.rem, report.new = newreport))
}

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   Bin(deaths[upper.tri], logit(X\beta))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
#'               $hessian - Hessian of the likelihood
##
loglikProb <- function(beta, death.remain, report.new, X){

  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X <- X[index,]
  p = as.vector(1/(1+exp(-X%*%beta)))
  y <- y[index]
  n <- n[index]
  lik <- sum(dbinom(y,size = n, prob = as.vector(p), log = T))

  grad <- as.vector(t(X)%*%as.vector(y-n*p))

  Hessian <- -t(X)%*%diag(as.vector(n*p*(1-p)))%*%X
  Hessian <- diag(diag(Hessian))

  return(list(loglik = lik, grad = grad, Hessian = Hessian))
}
loglikPrior <- function(beta, n, sigma, res){
  n <- length(sigma)
  res$grad[1:n] <-  res$grad[1:n]-beta[1:n]/(sigma^2)
  diag(res$Hessian) <- diag(res$Hessian) - c(1/sigma^2, rep(0,length(beta)-n))
  res$loglik <- res$loglik -sum(beta[1:n]^2/(2*sigma^2))
  return(res)
}

loglik <- function(beta,death.remain, report.new, X, sigma){
  res <- loglikProb(beta, death.remain, report.new, X)
  res <- loglikPrior(beta, n, sigma, res)
  return(res)
}

#####
##
# beta binomial part
##
#####

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   BB(deaths[upper.tri], exp(X\beta_1), exp(X\beta_2))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
##
loglikProbBB_mixedeffect_bi <- function(beta, death.remain, report.new, X_mu, X_M, X_mixed, sigma_mixed){

  p <- dim(X_mu)[2]
  p1 <- dim(X_M)[2]
  p2 <- dim(X_mixed)[2]
  beta_mu    <- beta[1:p]
  beta_M     <- beta[p + (1:p1)]
  beta_mixed <- beta[p + p1 + (1:p2)]
  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X_mu    <- X_mu[index,]
  X_mixed <- X_mixed[index,]
  X_M     <- X_M[index,]
  mu      <- 1/(1+exp(-X_mu%*%beta_mu - X_mixed%*%beta_mixed))
  M       <- exp(X_M%*%beta_M)

  alpha <- M * mu
  beta  <- M * (1-mu)
  y <- y[index]
  n <- n[index]

  if(max(beta)>120)
    return(list(loglik = -Inf, grad = 0))

  lik <- sum(dBB(y, size = n, alpha = alpha, beta = beta, log.p=T))
  if(is.na(lik))
    return(list(loglik = -Inf, grad = 0))

  ##
  #
  # mixed effect
  ##
  lik <- lik - (0.5/sigma_mixed^2) * sum(beta_mixed^2)
  ##
  # prior
  ##
  #lik <- lik - sum(exp(beta_M))
  grad_lik <- grad_dBB(y, size = n, alpha = alpha, beta = beta)

  grad_mu <- (M*mu * (1-mu) * grad_lik$grad_alpha)
  grad_mu <- (grad_mu - (M * mu * (1-mu) * grad_lik$grad_beta))
  grad_mixed <- t(X_mixed)%*%grad_mu - (1/sigma_mixed^2) * beta_mixed
  grad_mu    <- t(X_mu)%*%grad_mu
  grad_M <- (M*mu * grad_lik$grad_alpha)
  grad_M <- grad_M + (M * (1-mu) * grad_lik$grad_beta)
  grad_M <- t(X_M)%*%grad_M #- exp(beta_M)
  grad =  c(as.vector(grad_mu),as.vector(grad_M),as.vector(grad_mixed))
  if(sum(is.na(grad))>0)
    return(list(loglik = -Inf, grad = 0))
  return(list(loglik = lik, grad = grad))
}

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   BB(deaths[upper.tri], exp(X\beta_1), exp(X\beta_2))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
##
loglikProbBB_bi <- function(beta, death.remain, report.new, X_mu, X_M, pi_null, p_null, z){

  p <- dim(X_mu)[2]
  p1 <- dim(X_M)[2]
  beta_mu    <- beta[1:p]
  beta_M     <- beta[p + (1:p1)]
  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X_mu    <- X_mu[index,]
  X_M     <- X_M[index,]
  mu      <- 1/(1+exp(-X_mu%*%beta_mu ))
  M       <- exp(X_M%*%beta_M)

  alpha <- M * mu
  beta  <- M * (1-mu)
  y <- y[index]
  n <- n[index]

  if(is.na(z[1]))
    z = rep(0,length(y))
  index_z   <- z==0
  y_lik     <- y[index_z]
  n_lik     <- n[index_z]
  alpha_lik <- alpha[index_z]
  beta_lik  <- beta[index_z]
  mu_lik     <- mu[index_z]
  M_lik     <- M[index_z]
  if(max(beta)>100)
    return(list(loglik = -Inf, grad = 0))

  lik_obs <- dBB(y, size = n, alpha = alpha, beta = beta, log.p=T)
  lik <- sum(lik_obs)
  if(is.na(lik))
    return(list(loglik = -Inf, grad = 0))

  ##
  # prior
  ##
  lik <- lik
  grad_lik <- grad_dBB(y_lik, size = n_lik, alpha = alpha_lik, beta = beta_lik)

  grad_mu <- (M_lik*mu_lik * (1-mu_lik) * grad_lik$grad_alpha)
  grad_mu <- (grad_mu - (M_lik * mu_lik * (1-mu_lik) * grad_lik$grad_beta))
  grad_mu    <- t(X_mu[index_z==T,])%*%grad_mu
  grad_M <- (M_lik*mu_lik * grad_lik$grad_alpha)
  grad_M <- grad_M + (M_lik * (1-mu_lik) * grad_lik$grad_beta)
  grad_M <- t(X_M[index_z==T,])%*%grad_M
  grad =  c(as.vector(grad_mu),as.vector(grad_M))
  if(sum(is.na(grad))>0)
    return(list(loglik = -Inf, grad = 0))

  ##
  # sample mixture
  ##
  lik_null <- dbinom(y, n, p_null, log=T)
  p_null   <- log(pi_null) + lik_null
  p_full  <-  log(1-pi_null)   + lik_obs
  p_max   <- apply(cbind(p_null,p_full),1, max)
  p0 <- exp(p_null - p_max)
  p0 <- p0/(p0 + exp(p_full-p_max))
  z <- as.vector(1*  (runif(length(p0)) <= p0))
  pi <- rbeta(1, 5+ sum(z==1), 5 + sum(z==0))
  p  <- rbeta(1, 1 + sum(y[z==1]), 20+ sum(n[z==1])-sum(y[z==1]))

  return(list(loglik = lik, grad = grad, p_null= p, pi_null = pi, z = z))
}

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   BB(deaths[upper.tri], exp(X\beta_1), exp(X\beta_2))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
##
loglikProbBB3 <- function(beta, death.remain, report.new, X_mu, X_M){

  p <- dim(X_mu)[2]
  p1 <- dim(X_M)[2]
  beta_mu    <- beta[1:p]
  beta_M     <- beta[p + (1:p1)]
  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X_mu    <- X_mu[index,]
  X_M     <- X_M[index,]
  mu      <- 1/(1+exp(-X_mu%*%beta_mu ))
  M       <- exp(X_M%*%beta_M)

  alpha <- M * mu
  beta  <- M * (1-mu)
  y <- y[index]
  n <- n[index]

  lik <- sum(dBB(y, size = n, alpha = alpha, beta = beta, log.p=T))
  if(is.na(lik))
    return(list(loglik = -Inf, grad = 0))

  ##
  # prior
  ##
  lik <- lik - sum(exp(beta_M))
  grad_lik <- grad_dBB(y, size = n, alpha = alpha, beta = beta)

  grad_mu <- (M*mu * (1-mu) * grad_lik$grad_alpha)
  grad_mu <- (grad_mu - (M * mu * (1-mu) * grad_lik$grad_beta))
  grad_mu    <- t(X_mu)%*%grad_mu
  grad_M <- (M*mu * grad_lik$grad_alpha)
  grad_M <- grad_M + (M * (1-mu) * grad_lik$grad_beta)
  grad_M <- t(X_M)%*%grad_M - exp(beta_M)
  grad =  c(as.vector(grad_mu),as.vector(grad_M))
  if(sum(is.na(grad))>0)
    return(list(loglik = -Inf, grad = 0))
  return(list(loglik = lik, grad = grad))
}

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   BB(deaths[upper.tri], exp(X\beta_1), exp(X\beta_2))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
##
loglikProbBB_param2 <- function(beta, death.remain, report.new, X, calcH=F){

  p <- dim(X)[2]
  beta_1 <- beta[1:p]
  beta_2 <- beta[(p+1):(2*p)]
  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X <- X[index,]
  mu <- 1/(1+exp(-X%*%beta_1))
  M  <- exp(X%*%beta_2)

  alpha <- M * mu
  beta  <- M * (1-mu)
  y <- y[index]
  n <- n[index]

  if(max(beta)>120)
    return(list(loglik = -Inf, grad = 0))

  lik <- sum(dBB(y, size = n, alpha = alpha, beta = beta, log.p=T))
  if(is.na(lik))
    return(list(loglik = -Inf, grad = 0))

  lik <- lik #- 2*sum(exp(beta_2))
  grad_lik <- grad_dBB(y, size = n, alpha = alpha, beta = beta)

  grad_mu <- (M*mu * (1-mu) * grad_lik$grad_alpha)
  grad_mu <- t(X)%*%(grad_mu - (M * mu * (1-mu) * grad_lik$grad_beta))
  grad_M <- (M*mu * grad_lik$grad_alpha)
  grad_M <- grad_M + (M * (1-mu) * grad_lik$grad_beta)
  grad_M <- t(X)%*%grad_M #- 2*exp(beta_2)
  grad =  c(as.vector(grad_mu),as.vector(grad_M))
  if(sum(is.na(grad))>0)
    return(list(loglik = -Inf, grad = 0))
  return(list(loglik = lik, grad = grad))
}

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   BB(deaths[upper.tri], exp(X\beta_1), exp(X\beta_2))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
##
loglikProbBB_mixedeffect <- function(beta, death.remain, report.new, X_mu, X_M, X_mixed, sigma_mixed){

  p <- dim(X_mu)[2]
  p1 <- dim(X_M)[2]
  p2 <- dim(X_mixed)[2]
  beta_mu    <- beta[1:p]
  beta_M     <- beta[p + (1:p1)]
  beta_mixed <- beta[p + p1 + (1:p2)]
  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X_mu    <- X_mu[index,]
  X_mixed <- X_mixed[index,]
  X_M     <- X_M[index,]
  mu      <- 1/(1+exp(-X_mu%*%beta_mu - X_mixed%*%beta_mixed))
  M       <- exp(X_M%*%beta_M)

  alpha <- M * mu
  beta  <- M * (1-mu)
  y <- y[index]
  n <- n[index]

  if(max(beta)>120)
    return(list(loglik = -Inf, grad = 0))

  lik <- sum(dBB(y, size = n, alpha = alpha, beta = beta, log.p=T))
  if(is.na(lik))
    return(list(loglik = -Inf, grad = 0))

  ##
  #
  # mixed effect
  ##
  lik <- lik - (0.5/sigma_mixed^2) * sum(beta_mixed^2)
  ##
  # prior
  ##
  #lik <- lik - sum(exp(beta_M))
  grad_lik <- grad_dBB(y, size = n, alpha = alpha, beta = beta)

  grad_mu <- (M*mu * (1-mu) * grad_lik$grad_alpha)
  grad_mu <- (grad_mu - (M * mu * (1-mu) * grad_lik$grad_beta))
  grad_mixed <- t(X_mixed)%*%grad_mu - (1/sigma_mixed^2) * beta_mixed
  grad_mu    <- t(X_mu)%*%grad_mu
  grad_M <- (M*mu * grad_lik$grad_alpha)
  grad_M <- grad_M + (M * (1-mu) * grad_lik$grad_beta)
  grad_M <- t(X_M)%*%grad_M #- exp(beta_M)
  grad =  c(as.vector(grad_mu),as.vector(grad_M),as.vector(grad_mixed))
  if(sum(is.na(grad))>0)
    return(list(loglik = -Inf, grad = 0))
  return(list(loglik = lik, grad = grad))
}

##
# likelihood observations given deaths and probabilites
# model is:
# newreport[upper.tri] \sim   BB(deaths[upper.tri], exp(X\beta_1), exp(X\beta_2))
#
#' @return obj - list
#'               $loglik  - logliklihood
#'               $grad    - gradient
##
loglikProbBB <- function(beta, death.remain, report.new, X, calcH=F){

  p <- dim(X)[2]
  beta_1 <- beta[1:p]
  beta_2 <- beta[(p+1):(2*p)]
  N <- dim(report.new)[1]

  index <- upper.tri(death.remain,diag = T)
  y = report.new[index]
  n = death.remain[index]
  index = is.na(y)==F & is.na(n)==F
  X <- X[index,]
  alpha <- exp(X%*%beta_1)
  beta  <- exp(X%*%beta_2) #unfortunate name
  if(min(1/(alpha+beta)) <10^-3)
    return(list(loglik = -Inf, grad = 0, Hessian = 0))
  y <- y[index]
  n <- n[index]

  lik <- sum(dBB(y, size = n, alpha = alpha, beta = beta, log.p=T))
  if(is.na(lik))
    return(list(loglik = -Inf, grad = 0, Hessian = 0))

  grad_lik <- grad_dBB(y, size = n, alpha = alpha, beta = beta)
  grad_alpha <- as.vector(t(X)%*%(alpha*grad_lik$grad_alpha))
  grad_beta  <- as.vector(t(X)%*%(beta*grad_lik$grad_beta))
  if(calcH){
    H_ <- Hessian_dBB(y, size = n, alpha = alpha, beta = beta)
    H_12 <-  t(X)%*%diag(as.vector(alpha*beta*H_$grad_alpha_beta))%*%X
    H_11 <-  t(X)%*%diag(as.vector(alpha^2*H_$grad_alpha_alpha))%*%X
    H_22 <-  t(X)%*%diag(as.vector(beta^2*H_$grad_beta_beta))%*%X
    Hessian <- rbind(cbind(H_11, H_12) ,
                     cbind(H_12, H_22) )
    Hessian  <- diag(diag(Hessian) )
    diag(Hessian) <- diag(Hessian)
  }else{
    Hessian =NULL
  }
  grad =  c(grad_alpha,grad_beta)
    if(sum(is.na(grad))>0)
      return(list(loglik = -Inf, grad = 0, Hessian = 0))
  return(list(loglik = lik, grad = grad, Hessian = Hessian))
}
##
# PRobability of beta binomial
##
dBB<- function(x, size, alpha, beta, log.p = F){
  const = lgamma(size + 1) - lgamma(x + 1) -
    lgamma(size - x + 1)
  ld <- const +
    lgamma(x + alpha) + lgamma(size - x + beta) -
    lgamma(size + alpha + beta) +
    lgamma(alpha + beta) -
    lgamma(alpha) - lgamma(beta)
  if(log.p==F){
    return(exp(ld))
  }
  return(ld)
}
##
# gradient of log of probability beta binomial
# return list,
# [[1]] grad alpha
# [[2]] grad beta
##
grad_dBB<- function(x, size, alpha, beta){
  grad_alpha <- digamma(x + alpha) -
    digamma(size + alpha + beta) +
    digamma(alpha + beta) -
    digamma(alpha)
  grad_beta <- digamma(size - x + beta) -
    digamma(size + alpha + beta) +
    digamma(alpha + beta) -
    digamma(beta)
  return(list(grad_alpha = grad_alpha,
              grad_beta  = grad_beta ))
}
##
# Hessian of log of probability beta binomial
# return list,
# [[1]] grad alpha
# [[2]] grad beta
##
Hessian_dBB<- function(x, size, alpha, beta){
  grad_alpha_alpha <- trigamma(x + alpha) -
    trigamma(size + alpha + beta) +
    trigamma(alpha + beta) -
    trigamma(alpha)
  grad_beta_beta <- trigamma(size - x + beta) -
    trigamma(size + alpha + beta) +
    trigamma(alpha + beta) -
    trigamma(beta)

  grad_alpha_beta <- -trigamma(size + alpha + beta) +
    trigamma(alpha + beta)
  return(list(grad_alpha_alpha = grad_alpha_alpha,
              grad_beta_beta  = grad_beta_beta,
              grad_alpha_beta = grad_alpha_beta))
}

##
# log liklihood of obseving report given death and prob
# density is Beta binomial
#  deaths  - (int) true number of deaths
#  alpha       - (n x 1) bb parameter 1
#  beta        - (n x 1) bb parameter 2
#  report  - (n x 1) reported deaths culmative each date
##
loglikDeathsGivenProbBB <- function(death, alpha, beta, report){

  if(death < max(report,na.rm=T))
    return(-Inf)
  n <- length(alpha)
  if(is.na(report[1])){
    #we dont have data from day one

    ndeaths <- diff(report[is.na(report)==F])
    report_adj <- report[1:(n-1)]
    report_adj <- report_adj[is.na(report_adj)==F]
  }else{
    ndeaths <- c(report[1],diff(report))
    if(n>1){
      report_adj <- c(0, report[1:(n-1)])
    }else{
      report_adj <- 0
    }
  }

  ndeaths[ndeaths <0 ] = 0

  return(sum(dBB(x =ndeaths,
                 size = death - report_adj,
                 alpha = alpha[is.na(alpha)==F],
                 beta = beta[is.na(alpha)==F],
                 log.p=T)))
}

##
# log liklihood of obseving report given death and prob
# density is Beta binomial
#  deaths  - (int) true number of deaths
#  alpha       - (n x 1) bb parameter 1
#  beta        - (n x 1) bb parameter 2
#  report  - (n x 1) reported deaths culmative each date
##
loglikDeathsGivenProbBB_bi <- function(death, alpha, beta, pi_null,p_null,report){

  if(death < max(report,na.rm=T))
    return(-Inf)
  n <- length(alpha)
  if(is.na(report[1])){
    #we dont have data from day one

    ndeaths <- diff(report[is.na(report)==F])
    report_adj <- report[1:(n-1)]
    report_adj <- report_adj[is.na(report_adj)==F]
  }else{
    ndeaths <- c(report[1],diff(report))
    if(n>1){
      report_adj <- c(0, report[1:(n-1)])
    }else{
      report_adj <- 0
    }
  }

  ndeaths[ndeaths <0 ] = 0
  p <- rep(0, length(ndeaths))

  log1 <- dBB(x =ndeaths,
              size = death - report_adj,
              alpha = alpha[is.na(alpha)==F],
              beta = beta[is.na(alpha)==F],
              log.p=T)
  #p_ <- rep(0,length(log1))
  #p_[1:min(2,length(ndeaths))] =log1[1:min(2,length(ndeaths))]
  #if(length(ndeaths)>2){
  #  p1 <- log(1-pi_null) + log1[3:length(ndeaths)]
  #  p2 <- log(pi_null) + dbinom(ndeaths[3:length(ndeaths)], death - report_adj[3:length(ndeaths)],prob=p_null,log = T)
  #  p_[3:length(ndeaths)] <- log( exp(p1) + exp(p2))
  #}
  p1 <- log(1-pi_null) + log1
  p2 <- log(pi_null) +  dbinom(ndeaths, death - report_adj,prob=p_null,log = T)
  p_ <- log( exp(p1) + exp(p2))
  return(sum(p_))
}

###
#
# posterior sampling of deaths given Prob vec
#
#  alpha     - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
#  beta      - (N x N) matrix of beta binom parameter (only upper triangular part relevant)
# Reported    - (N x N) matrix of reported deaths cumlative (only upper triangular relevant)
# Predict.day - (int) which day to of reporting to predict
# sim         - (2 x 1) MCMC samples inner loop and outer loop
# alpha       - (N x 1) stepsizes
# prior       - (function) prior for N (needs use N and i)
#
###
death.givenParamBB <- function(alpha, beta, Reported,Predict.day = Inf,sim=c(2000,10), alpha.MCMC = NULL,
                               use.prior=F,
                               prior=NULL){

  N <- dim(alpha)[1]
  if(is.null(alpha.MCMC))
    alpha.MCMC <- rep(4, N)

  deaths <- matrix(NA, nrow=sim[1], ncol = N)
  deaths_est <- apply(Reported, 1, max, na.rm=T)

  burnin = ceiling(0.3*sim[1])
  for(i in 1:(sim[1] + burnin -1)){
    res <- sample.deathsBB(sim[2],
                           deaths_est,
                           alpha,
                           beta,
                           Reported,
                           alpha.MCMC,
                           rep.day=Predict.day,
                           use.prior=use.prior,
                           prior= prior)
    deaths_est <- res$deaths
    if(i < burnin){
      alpha.MCMC[res$acc/sim[2] > 0.3] <- alpha.MCMC[res$acc/sim[2] > 0.3] +1
      alpha.MCMC[res$acc/sim[2] < 0.3] <- alpha.MCMC[res$acc/sim[2] < 0.3] -1
      alpha.MCMC[alpha.MCMC<1] <- 1
    }else{
      deaths[i - burnin +1,] = deaths_est
    }
  }
  return(deaths)
}

###
#' setup up data uk covarites data,
#' @param dates_report - (N x 1) which days are reports on
#' @param predict.day   - (int) how many days back are predictied
#' @param unique.prob   - (int) how many lags have there own probability
##

##
# buidling the X matrix
#
#
# unique.prob - (int) how many of days shall we have unique probabilites
##
setup_data <- function(N, Predict.day, dates_report, unique.prob=NULL){

  X <- buildXdayeffect(N,Predict.day)
  #holidays and weekends

  result$dates_report <- as.Date(dates_report)

  if(is.null(unique.prob)==F)
    X <- cbind(X[,1:unique.prob,drop=F],rowSums(X[,1:unique.prob,drop=F])==0)

  holidays <- weekdays(dates_report)%in%c("Sunday","Saturday") | c(dates_report)%in%c(holidays.Sweden)
  holiday.yesterday <- weekdays(dates_report - 1)%in%c("Sunday","Saturday") |
    (result$dates_report-1)%in%c(holidays.Sweden)
  Xhol <- buildXholiday(N,holidays - (holiday.yesterday*holidays==T))
  Xhol.sunday <- buildXholiday(N,holiday.yesterday*holidays==T)
  Xhol.yesterday <-buildXholiday(N,holiday.yesterday - (holiday.yesterday*holidays==T) )

  X_dim <- dim(X)[2]
  X <- cbind(X, Xhol, Xhol.sunday, Xhol.yesterday, (Xhol+Xhol.sunday)*X[,1])
  colnames(X) <- c(paste('lag_',0:(X_dim-1),sep=''),'holiday','sunday','holiday_yesterday','lag0_hol')
  return(X)
}

##
# buidling the X matrix
#
#
# unique.prob - (int) how many of days shall we have unique probabilites
##
setup_data_lag <- function(N, Predict.day, dates_report, unique.prob=NULL){

  X <- buildXdayeffect(N,Predict.day)
  #holidays and weekends

  result$dates_report <- as.Date(dates_report)

  if(is.null(unique.prob)==F)
    X <- cbind(X[,1:unique.prob,drop=F],rowSums(X[,1:unique.prob,drop=F])==0)

  holidays <- weekdays(dates_report)%in%c("Sunday","Saturday") | c(dates_report)%in%c(holidays.Sweden)
  holiday.yesterday <- weekdays(dates_report - 1)%in%c("Sunday","Saturday") |
    (result$dates_report-1)%in%c(holidays.Sweden)
  Xhol <- buildXholiday(N,holidays - (holiday.yesterday*holidays==T))
  Xhol.sunday <- buildXholiday(N,holiday.yesterday*holidays==T)
  Xhol.yesterday <-buildXholiday(N,holiday.yesterday - (holiday.yesterday*holidays==T) )
  X.Wednesday <-buildXholiday(N, weekdays(dates_report)%in%c('Wednesday') )
  #X <- cbind(X, Xhol, Xhol.sunday, Xhol.yesterday, Xhol*X[,1],X.Wednesday-X.Wednesday*Xhol)
  X <- cbind(X, Xhol, Xhol.sunday, Xhol.yesterday, Xhol*X[,1])

  return(X)
}
##
# buidling the X matrix
#
#
# unique.prob - (int) how many of days shall we have unique probabilites
##
setup_data_postlag <- function(N, lag, Predict.day, dates_report, unique.prob=NULL){
  if(Predict.day==0){
    X <- matrix(1, nrow=N*(N+1)/2,  ncol=1+lag+2)
    unique.prob= NULL
  }else{
    X <- buildXdayeffect(N,Predict.day+lag+2)
  }
  if(lag>0)
    X <- X[,-(1:(lag+2)),drop=F]
  #holidays and weekends

  result$dates_report <- as.Date(dates_report)

  if(is.null(unique.prob)==F)
    X <- cbind(X[,1:unique.prob,drop=F],rowSums(X[,1:unique.prob,drop=F])==0)

  holidays <- weekdays(dates_report)%in%c("Sunday","Saturday") | c(dates_report)%in%c(holidays.Sweden)
  holiday.yesterday <- weekdays(dates_report - 1)%in%c("Sunday","Saturday") |
    (result$dates_report-1)%in%c(holidays.Sweden)
  Xhol <- buildXholiday(N,holidays - (holiday.yesterday*holidays==T))
  Xhol.sunday <- buildXholiday(N,holiday.yesterday*holidays==T)
  Xhol.yesterday <-buildXholiday(N,holiday.yesterday - (holiday.yesterday*holidays==T) )
  #X.Wednesday <-buildXholiday(N, weekdays(dates_report)%in%c('Wednesday') )
  #X <- cbind(X, Xhol, Xhol.sunday, Xhol.yesterday, Xhol*X[,1],X.Wednesday-X.Wednesday*Xhol)
  X <- cbind(X, Xhol, Xhol.sunday, Xhol.yesterday)

  return(X)
}

##
# buidling the X matrix
#
#
# unique.prob - (int) how many of days shall we have unique probabilites
##
setup_data_postlag2 <- function(N, lag, params, dates_report){
  if(params==0){
    X <- matrix(1, nrow=N*(N+1)/2,  ncol=1)
    unique.prob= 0
  }else{
    X_ <- matrix(0, nrow=N, ncol = N)
    holidays <- weekdays(dates_report)%in%c("Sunday","Saturday") | c(dates_report)%in%c(holidays.Sweden)
    for(i in 1:N){
      k <- which.max(cumsum(holidays[i:N]==F)==(lag+1))-1
      for(p in 1:params){
        if((i+k +p) <= N)
          X_[i,(i+k +p)] <- p
      }
    }
    X_ <- X_[upper.tri(X_,diag=T)]
    j_ <-  X_[X_>0]
    i_base <- 1:length(X_)
    i_ <-  i_base[X_>0]
    X <- sparseMatrix(i=i_,j=j_ , dims=c(length(X_), max(j_))  )
    unique.prob = max(j_)
  }

  #holidays and weekends

  result$dates_report <- as.Date(dates_report)

  if(unique.prob)
    X <- cbind(X[,1:unique.prob,drop=F],rowSums(X[,1:unique.prob,drop=F])==0)

  holidays <- weekdays(dates_report)%in%c("Sunday","Saturday") | c(dates_report)%in%c(holidays.Sweden)
  holiday.yesterday <- weekdays(dates_report - 1)%in%c("Sunday","Saturday") |
    (result$dates_report-1)%in%c(holidays.Sweden)
  Xhol <- buildXholiday(N,holidays - (holiday.yesterday*holidays==T))
  Xhol.sunday <- buildXholiday(N,holiday.yesterday*holidays==T)
  Xhol.yesterday <-buildXholiday(N,holiday.yesterday - (holiday.yesterday*holidays==T) )
  X <- cbind(X, Xhol, Xhol.sunday, Xhol.yesterday)

  return(X)
}

setup_data_mixture_cov <- function(N, lag, dates_report){

    X_ <- matrix(0, nrow=N, ncol = N)
    holidays <- weekdays(dates_report)%in%c("Sunday")
    for(i in 1:N){
      k <- which.max(cumsum(holidays[i:N]==F)==(lag+1))-1
      if(i+k < N)
        X_[i,(i+k+1):N] = (i+k+1):N
    }
    X_[,holidays==T]=0
  return(X_)
}
setup_all_days <- function(N){

  X_ <- matrix(0, nrow=N, ncol = N)
  for(i in 1:N){
      X_[i,i:N] = i:N
  }
  X <- X_[upper.tri(X_,diag=T)]
  return(X)
}

setup_data_mixed_effect <- function(N, lag, dates_report){
  X_ <- setup_data_mixture_cov(N, lag, dates_report)
  X_ <- X_[upper.tri(X_,diag=T)]
  X_ <- as.integer(factor(X_))-1
  X <- matrix(0, nrow=length(X_), ncol = max(X_))
  for(i in 1:max(X_)){
    X[,i] <- X_==i
  }
  return(X)
}


## get redshift offsets from a data cube or stacked rss file

regrid <- function(lambda.out, lib) {
    lambda.in <- lib$lambda
    nc <- ncol(lib)
    lib.out <- matrix(NA, length(lambda.out), nc)
    lib.out[,1] <- lambda.out
    for (i in 2:nc) {
        lib.out[,i] <- approx(lambda.in, lib[,i], xout=lambda.out)$y
    }
    colnames(lib.out) <- colnames(lib)
    data.frame(lib.out)
}

getdz <- function(gdat, lib, snrthresh=5, nlthresh=2000, dzlim=0.003, searchinterval=1e-4) {
  dims <- dim(gdat$flux)
  nr <- dims[1]
  nc <- dims[2]
  dz <- matrix(NA, nr, nc)
  dz.err <- matrix(NA, nr, nc)
  z <- gdat$meta$z
  lambda <- gdat$lambda
  fmla <- as.formula(paste("f ~", paste(names(lib)[-1], collapse="+")))
  fn <- function(x) {
    lambda.out <- lambda/(1+z+x)
    lib.out <- regrid(lambda.out, lib)
    lfit <- lm(fmla, data=lib.out, weights=iv)
    deviance(lfit)
  }
  fn_v <- Vectorize(fn)
  z_search <- seq(-dzlim, dzlim, by=searchinterval)

  for (i in 1:nr) {
    for (j in 1:nc) {
      if (is.na(gdat$snr[i,j]) || gdat$snr[i,j]<=snrthresh) next
        f <- gdat$flux[i, j, ]
        if (length(which(!is.na(f))) < nlthresh) next
          iv <- gdat$ivar[i, j, ]
          dev_grid <- fn_v(z_search)
          z0 <- z_search[which.min(dev_grid)]
          bestz <- Rsolnp::solnp(pars=z0, fun=fn,
                                 LB=z0-searchinterval, UB=z0+searchinterval,
                                 control=list(trace=0))
          dz[i,j] <- bestz$pars
          dz.err[i,j] <- sqrt(2./bestz$hessian)
    }
  }
  list(dz=dz, dz.err=dz.err)
}

## peculiar velocity

vrel <- function(z, zbar) {
  299792.458*(z-zbar)/(1+zbar)
}

## simple 2D spatial interpolation

fillpoly <- function(ra, dec, zvals, dxy=0.5) {
    ny <- max(round((max(dec)-min(dec))*3600/dxy), 100)
    nx <- ny
    allok <- complete.cases(ra, dec, zvals)
    zsurf <- akima::interp(x=ra[allok], y=dec[allok], z=zvals[allok],
                           linear=FALSE, nx=nx, ny=ny)
    list(ra=zsurf$x, dec=zsurf$y, z=zsurf$z)
}

## an image function using ggplot2

ggimage <- function(zmat, x=NULL, y=NULL, col=viridis::viridis(256),
                    xlab=NULL, ylab=NULL, legend=NULL, title=NULL,
                    xrev=FALSE, yrev=FALSE, asp=1,
                    addContour=FALSE, binwidth=NULL
                   ) {
    require(ggplot2)
    if (is.null(x)) x <- 1:nrow(zmat)
    if (is.null(y)) y <- 1:ncol(zmat)
    if (!is.null(dim(x))) {
      x <- colMeans(x)
      y <- rowMeans(y)
    }
    xy <- expand.grid(x,y)
    df <- data.frame(cbind(xy, as.vector(zmat)))
    names(df) <- c("x","y","z")
    g1 <- ggplot(df, aes(x=x, y=y, z=z)) +
        geom_raster(aes(fill=z)) +
        scale_fill_gradientn(colors=col, na.value="#FFFFFF00") +
        coord_fixed(ratio = asp)
    if (addContour) {
        if (!is.null(binwidth)) {
            g1 <- g1 + geom_contour(binwidth=binwidth, na.rm=TRUE)
        } else {
            g1 <- g1 + geom_contour(na.rm=TRUE)
        }
    }
    if (xrev) g1 <- g1 + scale_x_reverse()
    if (yrev) g1 <- g1 + scale_y_reverse()
    if (!is.null(xlab)) {
        g1 <- g1 + xlab(xlab)
    }
    if (!is.null(ylab)) {
        g1 <- g1 + ylab(ylab)
    }
    if (!is.null(legend)) {
        g1 <- g1 + labs(fill=legend)
    }
    if (!is.null(title)) {
        g1 <- g1 + ggtitle(title)
    }
    g1
}

rainbow <- c("blue", "cyan", "green", "yellow", "red")

## disk galaxy rotation curve model

vrot <- function(gdat.stack, dz.stack, phi.guess, sd.phi.guess=10,
                      ci.guess=0.7, sd.ci.guess=0.1,
                      sd.kc0=0.5,
                      order=3, N_r=100,
                      r_eff=1, remax=1.5, v_norm=100, PLOTS=TRUE,
                      smodel=NULL,
                      stanfile="vrot.stan", stanfiledir=".",
                      seed=220755, iter=2000, warmup=1000,
                      chains=4, cores=4,
                      control=list(adapt_delta=0.975, max_treedepth=15)) {
  require(rstan)
  require(ggplot2)
  v <- vrel(dz.stack$dz+gdat.stack$meta$z, gdat.stack$meta$z)
  v_sample <- v[!is.na(v)]
  dv <- (dz.stack$dz.err/dz.stack$dz)*v
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v, dv=dv)
  df <- df[complete.cases(df),]
  rmax <- max(sqrt(gdat.stack$xpos^2+gdat.stack$ypos^2)[!is.na(dz.stack$dz)])
  if (rmax %% 2 != 0) rmax = rmax+1
  r_norm <- r_eff*remax
  r_post <- seq(0, 0.99, length=N_r)*rmax/r_norm

  vlos_data <- list(order=order, N=nrow(df), x=df$x/r_norm, y=df$y/r_norm,
                    v=df$v/v_norm, dv=df$dv/v_norm,
                    phi0=cos(phi.guess*pi/180), sd_phi0=sin(phi.guess*pi/180)*sd.phi.guess*pi/180,
                    si0 = sqrt(1-ci.guess^2), sd_si0=sd.ci.guess,
                    sd_kc0 = sd.kc0,
                    r_norm=r_norm, v_norm=v_norm,
                    N_r=N_r, r_post=r_post)

  inits <- function(chain_id) {
    p <- phi.guess*pi/180
    si <- sqrt(1-ci.guess^2)
    list(phi=p, si=si, x_c=0., y_c=0., v_sys=0.)
  }
  if (!is.null(smodel)) {
    stanfit <- sampling(smodel, data=vlos_data, init=inits, seed=seed,
                        iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  } else {
    stanfit <- stan(file=file.path(stanfiledir, stanfile), data=vlos_data, init=inits,
                  seed=seed, iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  }
  post <- extract(stanfit)
  vf_obs <- fillpoly(gdat.stack$ra.f, gdat.stack$dec.f, v)
  v_model <- rep(NA, length(v))
  v_model[!is.na(v)] <- colMeans(post$v_model)*v_norm
  vf_model <- fillpoly(gdat.stack$ra.f, gdat.stack$dec.f, v_model)
  r_q <- apply(post$r*remax, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(r_q) <- paste("r", rownames(r_q), sep="_")
  v_r_q <- apply(post$v_rot*v_norm, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(v_r_q) <- paste("v_rot", rownames(v_r_q), sep="_")
  v_e_q <- apply(post$v_exp*v_norm, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(v_e_q) <- paste("v_exp", rownames(v_e_q), sep="_")
  df_v <- data.frame(cbind(t(r_q), t(v_r_q), t(v_e_q)))
  graphs <- list()
  i <- 1
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v)
  graphs[[i]] <- ggplot(df) + geom_point(aes(x=x, y=y, color=v), size=7) +
                              scale_colour_gradientn(colors=rainbow, na.value="gray95") +
                              coord_fixed(ratio=1) +
                              ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": observed fiber velocities"))
  i <- i+1
  graphs[[i]] <- ggimage(vf_obs$z, vf_obs$ra, vf_obs$dec, col=rainbow, asp=1/cos(pi*gdat.stack$meta$dec/180), addContour=T) +
                          xlim(rev(range(gdat.stack$ra.f))) + ylim(range(gdat.stack$dec.f)) +
                          xlab(expression(alpha)) + ylab(expression(delta)) +
                          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": observed velocity field"))
  i <- i+1
  graphs[[i]] <- ggimage(vf_model$z, vf_model$ra, vf_model$dec, col=rainbow, asp=1/cos(pi*gdat.stack$meta$dec/180), addContour=T) +
                          xlim(rev(range(gdat.stack$ra.f))) + ylim(range(gdat.stack$dec.f)) +
                          xlab(expression(alpha)) + ylab(expression(delta)) +
                          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": model velocity field"))
  if (exists("v_res", post)) {
    v_res <- rep(NA, length(v))
    v_res[!is.na(v)] <- colMeans(post$v_res)*v_norm
    vf_res <- fillpoly(gdat.stack$ra.f, gdat.stack$dec.f, v_res)
  } else {
    vf_res <- list(z=vf_obs$z-vf_model$z, ra=vf_obs$ra, dec=vf_obs$dec)
  }
  i <- i+1
  graphs[[i]] <- ggimage(vf_res$z, vf_res$ra, vf_res$dec, asp=1/cos(pi*gdat.stack$meta$dec/180), addContour=T) +
                        xlim(rev(range(gdat.stack$ra.f))) + ylim(range(gdat.stack$dec.f)) +
                        xlab(expression(alpha)) + ylab(expression(delta)) +
                        ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": residual velocity field"))
  df1 <- data.frame(r=as.vector(abs(post$r))*remax, v_rot=as.vector(post$v_rot)*v_norm, v_exp=as.vector(post$v_exp)*v_norm)
  i <- i+1
  graphs[[i]] <- ggplot(df_v) + geom_point(aes(x=r_50., y=v_rot_50.)) +
          geom_errorbar(aes(x=r_50., ymin=v_rot_2.5., ymax=v_rot_97.5.)) +
          geom_errorbarh(aes(x=r_50., y=v_rot_50., xmin=r_2.5., xmax=r_97.5.)) +
          stat_density_2d(aes(x=r, y=v_rot, fill=..level..), data=df1, geom="polygon", alpha=0.5, show.legend=FALSE) +
          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": joint distribution of r, v_rot")) +
          xlab(expression(r/r[eff]))
  i <- i+1
  graphs[[i]] <- ggplot(df_v) + geom_point(aes(x=r_50., y=v_exp_50.)) +
          geom_errorbar(aes(x=r_50., ymin=v_exp_2.5., ymax=v_exp_97.5.)) +
          geom_errorbarh(aes(x=r_50., y=v_exp_50., xmin=r_2.5., xmax=r_97.5.)) +
          stat_density_2d(aes(x=r, y=v_exp, fill=..level..), data=df1, geom="polygon", alpha=0.5, show.legend=FALSE) +
          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": joint distribution of r, v_exp")) +
          xlab(expression(r/r[eff]))
  if(exists("vrot_post", post)) {
  	  vq <- t(apply(v_norm * post$vrot_post, 2, quantile, probs=c(0.025, 0.5, 0.975)))
  	  df2 <- data.frame(r=r_post*remax, ymin=vq[,"2.5%"], ymed=vq[,"50%"], ymax=vq[,"97.5%"])
          i <- i+1
  	  graphs[[i]] <- ggplot(df2) + geom_line(aes(x=r, y=ymed)) + geom_ribbon(aes(x=r, ymin=ymin, ymax=ymax), alpha=0.5) +
          xlab(expression(r/r[eff]))
  } else {
  	  df2 <- NULL
  }
  if (PLOTS) {
  	  for (i in seq_along(graphs)) {
  	  	  x11()
  	  	  plot(graphs[[i]])
  	  }
  }
  list(stanfit=stanfit, r_eff=r_eff, remax=remax, v_norm=v_norm, vlos_data=vlos_data,
       v_sample=v_sample, vf_obs=vf_obs, vf_model=vf_model, vf_res=vf_res, df_v=df_v, df1=df1, df2=df2,
       graphs=graphs
  )
}

## disk galaxy rotation curve model -- gaussian process version

vrot_gp <- function(gdat.stack, dz.stack, phi.guess, sd.phi.guess=10,
                      ci.guess=0.7, sd.ci.guess=0.1,
                      sd.kc0=0.5,
                      order=3, N_xy=25^2,
                      r_eff=1, remax=1.5, v_norm=100, PLOTS=TRUE,
                      smodel=NULL,
                      stanfile="vrot_gp.stan", stanfiledir="~/vrot_stanmodels",
                      seed=220755, iter=1000, warmup=500,
                      chains=4, cores=4,
                      control=list(adapt_delta=0.9, max_treedepth=10)) {
  require(rstan)
  require(spmutils)
  require(ggplot2)
  require(akima)
  v <- vrel(dz.stack$dz+gdat.stack$meta$z, gdat.stack$meta$z)
  v_sample <- v[!is.na(v)]
  dv <- (dz.stack$dz.err/dz.stack$dz)*v
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v, dv=dv)
  allok <- complete.cases(df)
  df <- df[allok,]

  n_r <- round(sqrt(N_xy))
  N_xy <- n_r^2
  rho <- seq(1/n_r, 1, length=n_r)
  theta <- seq(0, (n_r-1)*2*pi/n_r, length=n_r)
  rt <- expand.grid(theta, rho)
  x_pred <- rt[,2]*cos(rt[,1])
  y_pred <- rt[,2]*sin(rt[,1])
  r_norm <- r_eff*remax

  vlos_data <- list(order=order, N=nrow(df), x=df$x/r_norm, y=df$y/r_norm,
                    v=df$v/v_norm, dv=df$dv/v_norm,
                    phi0=phi.guess*pi/180, sd_phi0=sd.phi.guess*pi/180,
                    ci0 = ci.guess, sd_ci0=sd.ci.guess,
                    sd_kc0 = sd.kc0,
                    r_norm=r_norm, v_norm=v_norm,
                    N_xy=N_xy, N_r=n_r, x_pred=x_pred, y_pred=y_pred
                    )

  inits <- function(chain_id) {
    p <- phi.guess*pi/180
    ci <- ci.guess
    list(phi=p, cos_i=ci, x_c=0., y_c=0., v_sys=0.)
  }
  if (!is.null(smodel)) {
    stanfit <- sampling(smodel, data=vlos_data, init=inits, seed=seed,
                        iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  } else {
    stanfit <- stan(file=file.path(stanfiledir, stanfile), data=vlos_data, init=inits,
                  seed=seed, iter=iter, warmup=warmup, control=control, chains=chains, cores=cores)
  }
  post <- extract(stanfit)
  graphs <- list()
  vf_model <- rep(NA, length(gdat.stack$xpos))
  vf_res <- rep(NA, length(gdat.stack$ypos))
  vf_model[allok] <- colMeans(post$vf_model)*v_norm
  vf_res[allok] <- colMeans(post$vf_res)*v_norm
  df <- data.frame(x=gdat.stack$xpos, y=gdat.stack$ypos, v=v, vf_model=vf_model, vf_res=vf_res)
  g <- 1
  graphs[[g]] <- ggplot(df) + geom_point(aes(x=x, y=y, color=v), size=7) +
                              scale_colour_gradientn(colors=rainbow, na.value="gray95") +
                              coord_fixed(ratio=1) +
                              ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": observed fiber velocities"))
  g <- g+1
  df <- df[allok,]
  df1 <- akima::interp(x=df$x, y=df$y, z=df$vf_model, nx=200, ny=200)
  graphs[[g]] <- ggimage(df1$z, df1$x, df1$y, col=rainbow, xlab="x", ylab="y", legend="vf_model", addContour=TRUE)
  g <- g+1
  df1 <- akima::interp(x=df$x, y=df$y, z=df$vf_res, nx=200, ny=200)
  graphs[[g]] <- ggimage(df1$z, df1$x, df1$y, xlab="x", ylab="y", legend="vf_res", addContour=TRUE)

  v_pred <- colMeans(post$v_pred)
  sd_v_pred <- apply(post$v_pred, 2, sd)
  v_model <- colMeans(post$v_model)
  v_res <- colMeans(post$v_res)
  x <- x_pred*r_norm
  y <- y_pred*r_norm
  rho <- c(0, rho)*remax
  df <- akima::interp(x=x, y=y, z=v_pred, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, col=rainbow, xlab="x", ylab="y", legend="v_pred", addContour=TRUE)
  df <- akima::interp(x=x, y=y, z=v_model, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, col=rainbow, xlab="x", ylab="y", legend="v_model", addContour=TRUE)
  df <- akima::interp(x=x, y=y, z=v_res, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, xlab="x", ylab="y", legend="v_res", addContour=TRUE)
  df <- akima::interp(x=x, y=y, z=sd_v_pred, nx=200, ny=200)
  g <- g+1
  graphs[[g]] <- ggimage(df$z, df$x, df$y, xlab="x", ylab="y", legend="sd_v_pred", addContour=TRUE)
  vrot_pred <- post$vrot_pred
  vexp_pred <- post$vexp_pred
  v_r_q <- data.frame(cbind(rho, t(apply(vrot_pred, 2, quantile, probs=c(.025, 0.5, 0.975)))))
  names(v_r_q) <- c("rho", "vr_025", "vr_500", "vr_975")
  g <- g+1
  graphs[[g]] <- ggplot(v_r_q) + geom_line(aes(x=rho, y= vr_500)) + geom_ribbon(aes(x=rho, ymin=vr_025, ymax=vr_975), alpha=0.5) +
                  xlab(expression(r/r[eff]))+ylab(expression(V[rot]))
  v_e_q <- data.frame(cbind(rho, t(apply(vexp_pred, 2, quantile, probs=c(.025, 0.5, 0.975)))))
  names(v_e_q) <- c("rho", "ve_025", "ve_500", "ve_975")
  g <- g+1
  graphs[[g]] <- ggplot(v_e_q) + geom_line(aes(x=rho, y= ve_500)) + geom_ribbon(aes(x=rho, ymin=ve_025, ymax=ve_975), alpha=0.5) +
                  xlab(expression(r/r[eff]))+ylab(expression(V[exp]))

  r_q <- apply(post$r*remax, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(r_q) <- paste("r", rownames(r_q), sep="_")
  v_r_q <- apply(post$v_rot*v_norm, 2, quantile, probs=c(0.025, 0.5, 0.975))
  rownames(v_r_q) <- paste("v_rot", rownames(v_r_q), sep="_")
  df_v <- data.frame(cbind(t(r_q), t(v_r_q)))
  df <- data.frame(r=as.vector(abs(post$r))*remax, v_rot=as.vector(post$v_rot)*v_norm)

  g <- g+1
  graphs[[g]] <- ggplot(df_v) + geom_point(aes(x=r_50., y=v_rot_50.)) +
          geom_errorbar(aes(x=r_50., ymin=v_rot_2.5., ymax=v_rot_97.5.)) +
          geom_errorbarh(aes(x=r_50., y=v_rot_50., xmin=r_2.5., xmax=r_97.5.)) +
          ggtitle(paste("mangaid", gdat.stack$meta$mangaid, ": joint distribution of r, v_rot")) +
          xlab(expression(r/r[eff])) + ylab("v (km/sec)")
  if (PLOTS) {
    for (i in seq_along(graphs)) {
      x11()
      plot(graphs[[i]])
    }
  }

  list(stanfit=stanfit, r_eff=r_eff, remax=remax, v_norm=v_norm, vlos_data=vlos_data,
       v_sample=v_sample, x_pred=x_pred, y_pred=y_pred,
       graphs=graphs
  )
}


#' sweep
#'
#' Sweep a vector through a float matrix.
#'
#' @details
#' Note that if the length of \code{STATS} does not recycle exactly across
#' \code{MARGIN}, the results here will differ slightly from the results of
#' base R.
#'
#' @param x
#' A float vector/matrix.
#' @param MARGIN
#' 1 (rows) or 2 (columns)
#' @param STATS
#' Vector to sweep out.
#' @param FUN
#' Sweeping function; must be one of \code{"+"}, \code{"-"}, \code{"*"},  or
#' \code{"/"}.
#' @param check.margin
#' Should x/STATS margin lengths be checked?
#' @param ...
#' Theoretically these are additional arguments passed to an arbitrary function.
#' However, we only support basic arithmetic, so they are ignored.
#'
#' @return
#' A matrix of the same type as the highest precision input.
#'
#' @examples
#' library(float)
#'
#' s = flrunif(10, 3)
#' sweep(s, 2, fl(1))
#'
#' @useDynLib float R_sweep_spm
#' @name sweep
#' @rdname sweep
NULL

sweep_float32 = function(x, MARGIN, STATS, FUN="-", check.margin=TRUE, ...)
{
  if (is.double(STATS))
    return(sweep(dbl(x), MARGIN, STATS, FUN, check.margin, ...))
  else if (!is.logical(STATS) && !is.integer(STATS) && !is.float(STATS))
    stop("non-numeric argument to binary operator")

  if (isavec(x))
    stop("Error in array(STATS, dims[perm]) : 'dims' cannot be of length 0", call.=FALSE)

  if (!isTRUE(all.equal(MARGIN, 1)) && !isTRUE(all.equal(MARGIN, 2)))
    stop("missing value where TRUE/FALSE needed")

  if (!is.character(FUN))
    stop("only '+', '-', '*', and '/' are implemented for floats")

  FUN = match.arg(FUN, c('+', '-', '*', '/'))

  if (isTRUE(check.margin))
  {
    if ((MARGIN == 1 && NROW(x)%%length(STATS)) || (MARGIN == 2 && NCOL(x)%%length(STATS)))
      warning("STATS does not recycle exactly across MARGIN")
  }

  ret = .Call(R_sweep_spm, DATA(x), as.integer(MARGIN), STATS, FUN)
  float32(ret)
}

#' @rdname sweep
#' @export
setMethod("sweep", signature(x="float32"), sweep_float32)


# # This script was written by nicolas m. roberts, primarily using the geologyGeometry library of functions, written by Joshua R. Davis.
# # if you incounter bugs in the code, please inform nick at nmroberts@wisc.edu

### PRELIMINARY WORK ####

# Set the working directory.
setwd("~/Desktop/20170620geologyGeometry")

# [!!!] FOR WINDOWS: Set the working directory.
# setwd("C:/users/[Insert user name]/Desktop/20170620geologyGeometry")

# Load the necessary R libraries.
source("library/all.R")

# Load some custom functions
source("JSG_statsFunctions.r")

# #[!!!] Uncomment only if you have compiled C on your system and have the C libraries loaded at the top of this script. Uncomment and run the following line only if you have compiled C. # [!!!] Markov chain Monte Carlo and Kamb contouring in equal volume plots require C compiler. Skip MCMC and equal volume Kamb lines if you do not wish to install C. Load the necessary library
# source("libraryC/all.R")

#==========================LOAD THE DATA========================================

# Load the foliation-lineation data
Follins <- geoDataFromFile("data/Follins_Ahs.csv")

# Check how many measurements there are
nrow(Follins)

#==========================PLOT THE DATA========================================

# Plot foliation-lineation locations in map view
plot(
    Follins$easting,
    Follins$northing,
    xlab = "Easting (meters)",
    ylab = "Northing (meters)"
)

# Plot the data in Equal area plot. Pole to foliation (red), Lineation (cyan)
lineEqualAreaPlotTwo(lapply(Follins$rotation, function(s)
    s[1, ]),
    lapply(Follins$rotation, function(s)
        s[2, ]))

# Plot Follins in Equal Area plots with Kamb contours (numbers are 2 sigma values)
# Poles to foliation (circles)
lineKambPlot(lapply(Follins$rotation, function(s)
    s[1, ]), c(2, 6, 10, 14, 18))
# Lineations (squares)
lineKambPlot(lapply(Follins$rotation, function(s)
    s[2, ]),
    c(2, 6, 10, 14, 18),
    shapes = c("s"))

# Plot the Follins in an Equal Volume plot (after Davis and Titus, 2017)
oriEqualVolumePlot(
    Follins$rotation,
    group = oriLineInPlaneGroup,
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.1,
    colors = "black" ,
    axesColors = c("black", "black", "black"),
    fogStyle = "none"
)

# If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
afterMaximizingWindow("Ahsahka_EqualVol1.png", "Ahsahka_EqualVol2.png")

# [!!!] Uncomment only if you have compiled C on your system and have the C libraries loaded at the top of this script.
# Skip if you have not compiled C. Plot 6-sigma Kamb contours for the data in an equal volume plot (after Davis and Titus, 2017)
# oricKambPlot(
#     Follins$rotation,
#     group = oriLineInPlaneGroup,
#     multiple = 6,
#     backgroundColor = "white",
#     curveColor = "black",
#     boundaryAlpha = 0.1,
#     colors = "black",
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none"
# )
#
# # If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
# afterMaximizingWindow("Ahsahka_Kamb_equalVol1.png", "Ahsahka_Kamb_EqualVol2.png")

#==========================IDENTIFY GEOGRAPHIC TRENDS========================================

# 1) Visually inspect the possibility of geographic trends in the data

# Plot foliation-lineation locations in map view, colored in gray-scale by northing. This shading will be used in the following plots.
plot(
    Follins$easting,
    Follins$northing,
    col = shades(Follins$northing),
    pch = 19
)

# Plot foliation-lineation orientations in an equal volume plot (after Davis and Titus, 2017)
oriEqualVolumePlot(
    Follins$rotation,
    group = oriLineInPlaneGroup,
    simplePoints = FALSE,
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.1,
    colors = shades(Follins$northing),
    axesColors = c("black", "black", "black"),
    fogStyle = "none"
)

# If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
afterMaximizingWindow("Ahsahka_shadesNorthing_equalVol1.png",
                      "Ahsahka_shadesNorthing_EqualVol2.png")

# Plot foliation-lineation orientations colored by northing. Poles to foliation (circles), Lineation (squares)
lineEqualAreaPlot(
    c(
        lapply(Follins$rotation, function(s)
            s[1, ]),
        lapply(Follins$rotation, function(s)
            s[2, ])
    ),
    col = shades(Follins$northing),
    shapes = c(replicate(length(Follins$rotation), "c"), replicate(length(Follins$rotation), "s"))
)

#————————————————————————————————————————————————————————————

# 2) Run geodesic regressions to quantify any extant trends.

# Create a temporary version of the foliation lineation data and establish a geographic criteria that encompasses all the data.
FollinsAll <- Follins
AhsAllCrit <-
    FollinsAll$easting > 550000 & Follins$northing > 5140000
FollinsAll$domain <- replicate(nrow(FollinsAll), 1)
Follins$domain[AhsAllCrit] <- 1

# 2A) Perform a geodesic regression with respect to azimuths every 10 degrees on all data, using the geographic criteria defined above. Each regression asks the question: "does orientation change linearly with respect to an azimuth towards x degrees"
regressions <- regressionSweep(FollinsAll, 10, (AhsAllCrit), 0)

# Create a data frame that contains the summary information for the regressions
regressionsSum <- data.frame(regressions[[19]])

# Name the columns of the data frame
names(regressionsSum) = c("Azimuth", "Error", "Min_Eigen", "R_squared", "Pvalue")

# View the summary table for the output.
regressionsSum

# Plot the R^2 value as a function of Azimuth
plot(
    as.vector(regressionsSum$Azimuth),
    as.vector(regressionsSum$R_squared),
    xlab = "Azimuth",
    ylab = "R-squared"
)

#————————————————————————————————————————————————————————————

# 3) Run kernal regression to quantify any extant trends.

# Recategorize the symmetric copies of the data.
mu <- oriFrechetMean(Follins$rotation, group=oriLineInPlaneGroup)
Follins$rotation <- oriNearestRepresentatives(Follins$rotation, mu, group=oriLineInPlaneGroup)

# Define the bandwith for a kernal regression. If desired, uncomment the next line (it may take ~20 minutes to run). Otherwise, use the pre-calculated value, "bandwidth".
# bandwidth <- rotBandwidthForKernelRegression(Follins$northing,Follins$rotation)
bandwidth <- 0.02314421

# Perform a kernel regression with respect to northing. The kernel function is analogous to finding a best fit curve in 2D.
kernelReg <- lapply(
    seq(
        from = min(Follins$northing),
        to = max(Follins$northing),
        by = 100
    ),
    rotKernelRegression,
    Follins$northing,
    Follins$rotation,
    bandwidth,
    numSteps = 1000
)

# Filter the regression results to only keep those that give an error of 0 and minimum eigen value > 0
sapply(kernelReg, function(regr)
    regr$error)
sapply(kernelReg, function(regr) {
    regr$minEigenvalue > 0
})

# Plot the regression curve in an equal volume plot.
kernelRegCurve <- lapply(kernelReg, function(regr)
    regr$r)
oriEqualVolumePlot(
    kernelRegCurve,
    group=oriLineInPlaneGroup,
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.1,
    colors = "black",
    axesColors = c("black", "black", "black"),
    fogStyle = "none",
    simplePoints = TRUE
)

# Compute the R^2 for the kernel regression.
KernRsquared <-
    rotRsquaredForKernelRegression(Follins$northing, Follins$rotation, bandwidth, numSteps =
                                       1000)

# View the R^2 value
KernRsquared

# Do a permutation test for significance. This takes a significant amount of time (>>1 hr)
rSquareds <-
    rotKernelRegressionPermutations(Follins$northing, Follins$rotation, bandwidth, numPerms =
                                        1000)
length(rSquareds)
p <- sum(rSquareds > KernRsquared$rSquared)
p

#==========================DEFINE GEOGRAPHIC DOMAINS========================================

# Define geographic criteria, based on sampling area, with the proposed shear zone boundary into account.
domain1Crit <- Follins$easting < 560000 & Follins$northing < 5158000
domain2Crit <- (Follins$easting >= 560000 & Follins$northing < 5163000) | Follins$easting >= 565000
domain3Crit <- !(domain1Crit | domain2Crit)

# # Create a new column in the dataframe in which to store the domain information
Follins$domain <- replicate(nrow(Follins), 1)

# Classify the foliation-only dataset by domain
Follins$domain[domain1Crit] <- 1
Follins$domain[domain2Crit] <- 2
Follins$domain[domain3Crit] <- 3

# Plot the locations of the foliation-lineation data in map view, each domain a different color. Domain 1 (Red), Domain 2 (Green), Domain 3 (Blue)
plot(
    x = Follins$easting,
    y = Follins$northing,
    xlab = "Easting (meters)",
    ylab = "Northing (meters)",
    col = hues(Follins$domain),
    pch = 19
)

#==========================IDENTIFY GEOGRAPHIC TRENDS WITHIN DOMAINS========================================

# 1) Make sure each domain is roughly unimodal
#DOMAIN 1, n = 23.
oriEqualVolumePlot(
    Follins$rotation[domain1Crit],
    oriLineInPlaneGroup,
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.1,
    colors = "black",
    axesColors = c("black", "black", "black"),
    fogStyle = "none"
)

#[!!!] Uncomment only if you have compiled C on your system and have the C libraries loaded at the top of this script.
# oricKambPlot(
#     Follins$rotation[domain1Crit],
#     oriLineInPlaneGroup,
#     backgroundColor = "white",
#     curveColor = "black",
#     boundaryAlpha = 0.1,
#     colors = "black",
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none"
# )

#DOMAIN 2, n = 14.
oriEqualVolumePlot(
    Follins$rotation[domain2Crit],
    oriLineInPlaneGroup,
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.1,
    colors = "black",
    axesColors = c("black", "black", "black"),
    fogStyle = "none"
)

#[!!!] Uncomment only if you have compiled C on your system and have the C libraries loaded at the top of this script.
# oricKambPlot(
#     Follins$rotation[domain2Crit],
#     oriLineInPlaneGroup,
#     backgroundColor = "white",
#     curveColor = "black",
#     boundaryAlpha = 0.1,
#     colors = "black",
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none"
# )

#DOMAIN 3, n = 32.
oriEqualVolumePlot(
    Follins$rotation[domain2Crit],
    oriLineInPlaneGroup,
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.1,
    colors = "black",
    axesColors = c("black", "black", "black"),
    fogStyle = "none"
)
oriEqualVolumePlot(Follins$rotation[domain2Crit],
                   oriLineInPlaneGroup,
                   col = hues(Follins$easting[domain2Crit]))

#[!!!] Uncomment only if you have compiled C on your system and have the C libraries loaded at the top of this script.
# oricKambPlot(
#     Follins$rotation[domain2Crit],
#     oriLineInPlaneGroup,
#     backgroundColor = "white",
#     curveColor = "black",
#     boundaryAlpha = 0.1,
#     colors = "black",
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none"
# )

# Plot all the foliation data in an equal volume plot (after Davis and Titus, 2017), colored by domain. Domain 1 (red), Domain 2 (green), Domain 3 (blue)
oriEqualVolumePlot(
    Follins$rotation,
    oriLineInPlaneGroup,
    col = hues(Follins$domain),
    backgroundColor = "white",
    curveColor = "black",
    boundaryAlpha = 0.2,
    axesColors = c("black", "black", "black"),
    fogStyle = "none"
)

#————————————————————————————————————————————————————————————

# 2) Look for geographic patterns in plots.
#DOMAIN 1, n = 23. maybe a very faint rainbow--more like clumping of the dark blues together.
oriEqualVolumePlot(Follins$rotation[domain1Crit],
                   oriLineInPlaneGroup,
                   col = shades(Follins$northing[domain1Crit]))

#DOMAIN 2, n = 14. Part of the problem with small sample sizes--you can't really know whether the datapoints are independent...
oriEqualVolumePlot(Follins$rotation[domain2Crit],
                   oriLineInPlaneGroup,
                   col = shades(Follins$northing[domain2Crit]))

#DOMAIN 3, n = 32. No rainbow.
oriEqualVolumePlot(Follins$rotation[domain2Crit],
                   oriLineInPlaneGroup,
                   col = shades(Follins$easting[domain2Crit]))

#————————————————————————————————————————————————————————————

# 3) From visual inspection of the above plots, Domain 1 and Domain 3 potentially have geographic dependency. Perform regressions on these two domains.

#DOMAIN 1, n = 23. Perform a series of Geodesic Regressions for every 10° azimuth, with p-values not calculated.
regressionsD1 <- regressionSweep(Follins, 10, domain1Crit, 0)

# Create a dataframe with the summary information for the regressions
regressionsD1Sum <- data.frame(regressionsD1[[19]])

# Name the columns of the data frame
names(regressionsD1Sum) = c("Azimuth", "Error", "Min_Eigen", "R_squared", "Pvalue")

# View the summary table for the output.
regressionsD1Sum

# Plot the R^2 value as a function of Azimuth
plot(
    as.vector(regressionsD1Sum$Azimuth),
    as.vector(regressionsD1Sum$R_squared),
    xlab = "Azimuth",
    ylab = "R-squared"
)

#————————————————————————————————————————————————————————————

# Define the Azimuth with the highest R-squared value. (In this case, 130)
azimuthD1 <-
    Follins$easting[domain1Crit] * sin(130 * degree) + Follins$northing[domain1Crit] *
    cos(130 * degree)

# DOMAIN 1, n = 23. Define the bandwith for a kernal regression. If desired, uncomment the next line (it may take ~20 minutes to run). Otherwise, use the pre-calculated value, "bandwidthD1".
# bandwidthD1 <- rotBandwidthForKernelRegression(azimuthD1,Follins$rotation[domain1Crit])
bandwidthD1 <- 0.02202349

#Run the kernel regression
kernelRegD1 <- lapply(
    seq(
        from = min(azimuthD1),
        to = max(azimuthD1),
        by = 100
    ),
    rotKernelRegression,
    azimuthD1,
    Follins$rotation[domain1Crit],
    bandwidthD1,
    numSteps = 1000
)
sapply(kernelRegD1, function(regr)
    regr$error)
sapply(kernelRegD1, function(regr) {
    regr$minEigenvalue > 0
})

# Plot the regression curve
kernelRegCurveD1 <- lapply(kernelRegD1, function(regr)
    regr$r)
oriEqualVolumePlot(Follins$rotation[domain1Crit],
                   curves = list(kernelRegCurveD1),
                   simplePoints = TRUE, oriLineInPlaneGroup)

# Compute the R^2.
KernRsquaredD1 <-
    rotRsquaredForKernelRegression(azimuthD1, Follins$rotation[domain1Crit], bandwidthD1, numSteps =
                                       1000)
KernRsquaredD1

# Do a permutation test for significance.
rSquaredsD1 <-
    rotKernelRegressionPermutations(azimuthD1, Follins$rotation[domain1Crit], bandwidthD1, numPerms =
                                        1000)
length(rSquareds)
pD1 <- sum(rSquareds > KernRsquared$rSquared)

#the p-value for this regression.
pD1

#————————————————————————————————————————————————————————————

#Now do the same thing for Domain 3

#DOMAIN 3, n = 32. A series of Geodesic Regressions for every 10° azimuth, with p-values not calculated.
regressionsD3 <- regressionSweep(Follins, 10, domain3Crit, 0)

# Create a dataframe with the summary information for the regressions
regressionsD3Sum <- data.frame(regressionsD3[[19]])

# Name the columns of the data frame
names(regressionsD3Sum) = c("Azimuth", "Error", "Min_Eigen", "R_squared", "Pvalue")

# View the summary table for the output.
regressionsD3Sum

# Plot the R^2 value as a function of Azimuth
plot(
    as.vector(regressionsD3Sum$Azimuth),
    as.vector(regressionsD3Sum$R_squared),
    xlab = "Azimuth",
    ylab = "R-squared"
)

#————————————————————————————————————————————————————————————

# Define the Azimuth with the highest R-squared value. (In this case, 60)
azimuthD3 <-
    Follins$easting[domain3Crit] * sin(60 * degree) + Follins$northing[domain3Crit] *
    cos(60 * degree)

#DOMAIN 3, n = 32. Define the bandwith for a kernal regression. If desired, uncomment the next line (it may take ~20 minutes to run). Otherwise, use the pre-calculated value, "bandwidthD3".
# bandwidthD3 <-
#     rotBandwidthForKernelRegression(Follins$northing[domain3Crit], Follins$rotation[domain3Crit])
bandwidthD3 <- 0.1331092

#Run the kernel regression
kernelRegD3 <- lapply(
    seq(
        from = min(azimuthD3),
        to = max(azimuthD3),
        by = 100
    ),
    rotKernelRegression,
    azimuthD3,
    Follins$rotation[domain3Crit],
    bandwidthD3,
    numSteps = 1000
)
sapply(kernelRegD3, function(regr)
    regr$error)
sapply(kernelRegD3, function(regr) {
    regr$minEigenvalue > 0
})

# Plot the regression curve
kernelRegCurveD3 <- lapply(kernelRegD3, function(regr)
    regr$r)
oriEqualVolumePlot(Follins$rotation[domain3Crit],
                   curves = list(kernelRegCurveD3),
                   simplePoints = TRUE,
                   oriLineInPlaneGroup)

# Compute the R^2.
KernRsquaredD3 <-
    rotRsquaredForKernelRegression(azimuthD3, Follins$rotation[domain3Crit], bandwidthD3, numSteps = 1000)
KernRsquaredD3

# Do a permutation test for significance.
rSquaredsD3 <-
    rotKernelRegressionPermutations(Follins$northing[domain3Crit], Follins$rotation[domain3Crit], bandwidthD3, numPerms = 1000)
length(rSquareds)
pD3 <- sum(rSquareds > KernRsquaredD3$rSquared)

#the p-value for this regression.
pD3

##==========================STATISTICAL DESCRIPTORS========================================

# 1) Domain 1, n = 23

# Frechet mean (minimizes the Frechet variance).
FrechetMeanDom1 <-
    oriFrechetMean(Follins$rotation[domain1Crit], oriLineInPlaneGroup)
geoStrikeDipRakeDegFromRotation(FrechetMeanDom1)
FrechetVarDom1 <-
    oriVariance(Follins$rotation[domain1Crit], FrechetMeanDom1, oriLineInPlaneGroup)

# Plot the FrechetMean in an Equal Angle plot
FrechetCurvesDom1 <-
    lapply(Follins$rotation[domain1Crit], function(r)
        rotGeodesicPoints(FrechetMeanDom1, r, 10))
oriEqualAnglePlot(points = Follins$rotation[domain1Crit], curves = FrechetCurvesDom1, group = oriLineInPlaneGroup)

# Print the Strike, Dip, Rake of the Frechet mean.
geoStrikeDipRakeDegFromRotation(FrechetMeanDom1)

# Print the Frechet variance.
FrechetVarDom1

#————————————————————————————————————————————————————————————

# 2) Domain 2, n = 14
#Frechet mean minimizes the Frechet variance.
FrechetMeanDom2 <-
    oriFrechetMean(Follins$rotation[domain2Crit], oriLineInPlaneGroup)
geoStrikeDipRakeDegFromRotation(FrechetMeanDom2)
FrechetVarDom2 <-
    oriVariance(Follins$rotation[domain2Crit], FrechetMeanDom2, oriLineInPlaneGroup)

# Plot the FrechetMean in an Equal Angle plot
FrechetCurvesDom2 <-
    lapply(Follins$rotation[domain2Crit], function(r)
        rotGeodesicPoints(FrechetMeanDom2, r, 10))
oriEqualAnglePlot(points = Follins$rotation[domain2Crit], curves = FrechetCurvesDom2, oriLineInPlaneGroup)

# Print the Strike, Dip, Rake of the Frechet mean.
geoStrikeDipRakeDegFromRotation(FrechetMeanDom2)

# Print the Frechet variance
FrechetVarDom2

#————————————————————————————————————————————————————————————

# 3) Domain 3, n = 32
#Frechet mean minimizes the Frechet variance.
FrechetMeanDom3 <-
    oriFrechetMean(Follins$rotation[domain3Crit], oriLineInPlaneGroup)
geoStrikeDipRakeDegFromRotation(FrechetMeanDom3)
FrechetVarDom3 <-
    oriVariance(Follins$rotation[domain3Crit], FrechetMeanDom3, oriLineInPlaneGroup)

# Plot the FrechetMean in an Equal Angle plot
FrechetCurvesDom3 <-
    lapply(Follins$rotation[domain3Crit], function(r)
        rotGeodesicPoints(FrechetMeanDom3, r, 10))
oriEqualAnglePlot(points = Follins$rotation[domain3Crit], curves = FrechetCurvesDom3, oriLineInPlaneGroup)

# Print the Strike, Dip, Rake of the Frechet mean.
geoStrikeDipRakeDegFromRotation(FrechetMeanDom3)

# Print the Frechet variance.
FrechetVarDom3

#————————————————————————————————————————————————————————————

# 4) Dispersion of the data

#The standard way to get at dispersion is to compute the maximum matrix fisher likelihood. The matrix fisher distribution comprises a kind of "mean" and a positive definite symmetric matrix, which characterises the anisotropic dispersion.

# Redefine the rotations to be within one of the symmetric copies
mu <- oriFrechetMean(Follins$rotation, group = oriLineInPlaneGroup)
Follins$rotation <-
    oriNearestRepresentatives(Follins$rotation, mu, oriLineInPlaneGroup)

# Northern Domain, n = 16. Fisher maximum likelihood
mleDom1 <- rotFisherMLE(Follins$rotation[domain1Crit])
mleDom1
eigen(mleDom1$kHat, symmetric = TRUE, only.value = TRUE)$values

#Central Domain, n = 34. Fisher maximum likelihood
mleDom2 <- rotFisherMLE(Follins$rotation[domain2Crit])
mleDom2
eigen(mleDom2$kHat, symmetric = TRUE, only.value = TRUE)$values

#Southern Domain, n = 79. Fisher maximum likelihood
mleDom3 <- rotFisherMLE(Follins$rotation[domain3Crit])
mleDom3
eigen(mleDom3$kHat, symmetric = TRUE, only.value = TRUE)$values

#==========================INFERENCE========================================

# Here, we can perform some statistical techniques to use information about the samples (each domain) to compute a probability cloud of the mean of the population(s) from which those samples were taken.
# From numerical work in Davis and Titus (2017), MCMC will work the best for small sample sizes that have the matrix fisher anisotropy (eigenvalues) of (large, large, small). All four domains have this anisotropy, and have too few data points to use the (computationally quicker) bootstrapping method.

#We'll do both and compare.

# Compute the bootstrap cloud of means
Follins_domain1Boots <-
    oriBootstrapInference(Follins$rotation[domain1Crit], 10000, oriLineInPlaneGroup)
Follins_domain2Boots <-
    oriBootstrapInference(Follins$rotation[domain2Crit], 10000, oriLineInPlaneGroup)
Follins_domain3Boots <-
    oriBootstrapInference(Follins$rotation[domain3Crit], 10000, oriLineInPlaneGroup)

# [!!!] Requires C libraries to run.
# [!!!] Compute the Markov chain Monte Carlo cloud of means
# [!!!]Uncomment lines if you have compiled C on your computer and have loaded the C libraries at the top of this script.

# Follins_domain1MCMC <-
#     oricWrappedTrivariateNormalMCMCInference(Follins$rotation[domain1Crit],
#                                              group = oriLineInPlaneGroup,
#                                              numCollection = 100)
# Follins_domain2MCMC <-
#     oricWrappedTrivariateNormalMCMCInference(Follins$rotation[domain2Crit],
#                                              group = oriLineInPlaneGroup,
#                                              numCollection = 100)
# Follins_domain3MCMC <-
#     oricWrappedTrivariateNormalMCMCInference(Follins$rotation[domain3Crit],
#                                              group = oriLineInPlaneGroup,
#                                              numCollection = 100)

#————————————————————————————————————————————————————————————

# 1) Northern v. Central domains

# A) Using MCMC

# [!!!] Requires C libraries to run.
# Construct the 95% confidence ellipsoids from small triangles
# tris1_MCMC <-
#     rotEllipsoidTriangles(
#         Follins_domain1MCMC$mBar,
#         Follins_domain1MCMC$leftCovarInv,
#         Follins_domain1MCMC$q095,
#         numNonAdapt = 4
#     )
# tris2_MCMC <-
#     rotEllipsoidTriangles(
#         Follins_domain2MCMC$mBar,
#         Follins_domain2MCMC$leftCovarInv,
#         Follins_domain2MCMC$q095,
#         numNonAdapt = 4
#     )
#
# # Plot the MCMC comparison in an equal Volume plot, with 95% confidence ellipsoids
# oriEqualAnglePlot(
#     points = c(Follins_domain1MCMC$ms, Follins_domain2MCMC$ms),
#     boundaryAlpha = .1,
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none",
#     background = "white",
#     triangles = c(tris1_MCMC, tris2_MCMC),
#     simplePoints = TRUE,
#     colors = c(replicate(length(
#         Follins_domain1MCMC$ms
#     ), "black"), replicate(
#         length(Follins_domain2MCMC$ms), "orange"
#     )),
#     group = oriTrivialGroup
# )
#
# # If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
# afterMaximizingWindow("MCMC_Ahs_d1d2_1.png", "MCMC_WestMt_d1d2_2.png")
#
# # Plot the MCMC comparison in an Equal Area plot
# lineEqualAreaPlotTwo(c(
#     lapply(Follins_domain1MCMC$ms, function(s)
#         s[1, ]),
#     lapply(Follins_domain1MCMC$ms, function(s)
#         s[2, ])
# ),
# c(
#     lapply(Follins_domain2MCMC$ms, function(s)
#         s[1, ]),
#     lapply(Follins_domain2MCMC$ms, function(s)
#         s[2, ])
# ))

#....................................................................

# B) Using bootstrapping

# Construct the 95% confidence ellipsoids from small triangles
tris1_Boot <-
    rotEllipsoidTriangles(
        Follins_domain1Boots$center,
        Follins_domain1Boots$leftCovarInv,
        Follins_domain1Boots$q095,
        numNonAdapt = 4
    )
tris2_Boot <-
    rotEllipsoidTriangles(
        Follins_domain2Boots$center,
        Follins_domain2Boots$leftCovarInv,
        Follins_domain2Boots$q095,
        numNonAdapt = 4
    )

# Plot the bootstrap comparison in an equal Volume plot, with 95% confidence ellipsoids
rotEqualAnglePlot(
    points = c(
        Follins_domain1Boots$bootstraps,
        Follins_domain2Boots$bootstraps
    ),
    triangles = c(tris1_Boot, tris2_Boot),
    boundaryAlpha = .1,
    axesColors = c("black", "black", "black"),
    fogStyle = "none",
    background = "white",
    simplePoints = TRUE,
    colors = c(replicate(
        length(Follins_domain1Boots$bootstraps), "black"
    ), replicate(
        length(Follins_domain2Boots$bootstraps), "orange"
    ))
)

# If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
afterMaximizingWindow("Boots_Ahs_d1d2_1.png", "Boots_Ahs_d1d2_2.png")

# Plot the bootstrap comparison in an Equal Area plot
lineEqualAreaPlotTwo(c(
    lapply(Follins_domain1Boots$bootstraps, function(s)
        s[1, ]),
    lapply(Follins_domain1Boots$bootstraps, function(s)
        s[2, ])
),
c(
    lapply(Follins_domain2Boots$bootstraps, function(s)
        s[1, ]),
    lapply(Follins_domain2Boots$bootstraps, function(s)
        s[2, ])
))

#————————————————————————————————————————————————————————————

# 2) Northern vs. Southern domains

# A) Using MCMC

# [!!!] Requires C libraries to run.
# Construct the 95% confidence ellipsoids from small triangles
# tris1_MCMC <-
#     rotEllipsoidTriangles(
#         Follins_domain1MCMC$mBar,
#         Follins_domain1MCMC$leftCovarInv,
#         Follins_domain1MCMC$q095,
#         numNonAdapt = 5
#     )
# tris3_MCMC <-
#     rotEllipsoidTriangles(
#         Follins_domain3MCMC$mBar,
#         Follins_domain3MCMC$leftCovarInv,
#         Follins_domain3MCMC$q095,
#         numNonAdapt = 5
#     )
#
# # Plot the bootstrap comparison in an equal Volume plot, with 95% confidence ellipsoids
# oriEqualAnglePlot(
#     points = c(Follins_domain1MCMC$ms, Follins_domain3MCMC$ms),
#     triangles = c(tris1_MCMC, tris3_MCMC),
#     boundaryAlpha = 0.1,
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none",
#     background = "white",
#     simplePoints = TRUE,
#     colors = c(replicate(length(
#         Follins_domain1MCMC$ms
#     ), "black"), replicate(length(
#         Follins_domain3MCMC$ms
#     ), "blue")),
#     group = oriTrivialGroup
# )
#
# # If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
# afterMaximizingWindow("MCMC_Ahs_d1d3_1.png", "MCMC_Ahsd1d3_2.png")
#
# # Plot the MCMC comparison in an Equal Area plot
# lineEqualAreaPlotTwo(c(
#     lapply(Follins_domain1MCMC$ms, function(s)
#         s[1, ]),
#     lapply(Follins_domain1MCMC$ms, function(s)
#         s[2, ])
# ),
# c(
#     lapply(Follins_domain3MCMC$ms, function(s)
#         s[1, ]),
#     lapply(Follins_domain3MCMC$ms, function(s)
#         s[2, ])
# ))

#....................................................................

# B) Using bootstrapping
# Construct the 95% confidence ellipsoids from small triangles
tris1_Boot <-
    rotEllipsoidTriangles(
        Follins_domain1Boots$center,
        Follins_domain1Boots$leftCovarInv,
        Follins_domain1Boots$q095,
        numNonAdapt = 4
    )
tris3_Boot <-
    rotEllipsoidTriangles(
        Follins_domain3Boots$center,
        Follins_domain3Boots$leftCovarInv,
        Follins_domain3Boots$q095,
        numNonAdapt = 4
    )

# Plot the bootstrap comparison in an equal Volume plot, with 95% confidence ellipsoids.
rotEqualAnglePlot(
    points = c(
        Follins_domain1Boots$bootstraps,
        Follins_domain3Boots$bootstraps
    ),
    triangles = c(tris1_Boot, tris3_Boot),
    boundaryAlpha = .1,
    axesColors = c("black", "black", "black"),
    fogStyle = "none",
    background = "white",
    simplePoints = TRUE,
    colors = c(replicate(
        length(Follins_domain1Boots$bootstraps), "black"
    ),
    replicate(
        length(Follins_domain3Boots$bootstraps), "blue"
    ))
)

# If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
afterMaximizingWindow("Boots_Ahs_d1d3_1.png", "Boots_Ahs_d1d3_2.png")

# Plot the bootstrap comparison in an Equal Area plot
lineEqualAreaPlotTwo(c(
    lapply(Follins_domain1Boots$bootstraps, function(s)
        s[1, ]),
    lapply(Follins_domain1Boots$bootstraps, function(s)
        s[2, ])
),
c(
    lapply(Follins_domain3Boots$bootstraps, function(s)
        s[1, ]),
    lapply(Follins_domain3Boots$bootstraps, function(s)
        s[2, ])
))

#————————————————————————————————————————————————————————————

# 3) Central vs. Southern domains

# A) Using MCMC

# [!!!] Requires C libraries to run.
# Construct the 95% confidence ellipsoids from small triangles
# tris2_MCMC <-
#     rotEllipsoidTriangles(
#         Follins_domain2MCMC$mBar,
#         Follins_domain2MCMC$leftCovarInv,
#         Follins_domain2MCMC$q095,
#         numNonAdapt = 4
#     )
# tris3_MCMC <-
#     rotEllipsoidTriangles(
#         Follins_domain3MCMC$mBar,
#         Follins_domain3MCMC$leftCovarInv,
#         Follins_domain3MCMC$q095,
#         numNonAdapt = 4
#     )
#
# # Plot the bootstrap comparison in an equal Volume plot, with 95% confidence ellipsoids
# oriEqualAnglePlot(
#     points = c(Follins_domain2MCMC$ms, Follins_domain3MCMC$ms),
#     triangles = c(tris2_MCMC, tris3_MCMC),
#     boundaryAlpha = 0.1,
#     axesColors = c("black", "black", "black"),
#     fogStyle = "none",
#     background = "white",
#     simplePoints = TRUE,
#     colors = c(replicate(
#         length(Follins_domain2MCMC$ms), "orange"
#     ), replicate(length(
#         Follins_domain3MCMC$ms
#     ), "blue")),
#     group = oriTrivialGroup
# )
#
# # If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
# afterMaximizingWindow("MCMC_Ahs_d2d3_1.png", "MCMC_Ahs_d2d3_2.png")
#
# # Plot the MCMC comparison in an Equal Area plot
# lineEqualAreaPlotTwo(c(
#     lapply(Follins_domain2MCMC$ms, function(s)
#         s[1, ]),
#     lapply(Follins_domain2MCMC$ms, function(s)
#         s[2, ])
# ),
# c(
#     lapply(Follins_domain3MCMC$ms, function(s)
#         s[1, ]),
#     lapply(Follins_domain3MCMC$ms, function(s)
#         s[2, ])
# ))

#....................................................................

# B) Using bootstrapping

# Construct the 95% confidence ellipsoids from small triangles
tris2_Boot <-
    rotEllipsoidTriangles(
        Follins_domain2Boots$center,
        Follins_domain2Boots$leftCovarInv,
        Follins_domain2Boots$q095,
        numNonAdapt = 4
    )
tris3_Boot <-
    rotEllipsoidTriangles(
        Follins_domain3Boots$center,
        Follins_domain3Boots$leftCovarInv,
        Follins_domain3Boots$q095,
        numNonAdapt = 4
    )

# Plot the bootstrap comparison in an equal Volume plot, with 95% confidence ellipsoids
oriEqualAnglePlot(
    points = c(
        Follins_domain2Boots$bootstraps,
        Follins_domain3Boots$bootstraps
    ),
    triangles = c(tris2_Boot, tris3_Boot),
    boundaryAlpha = 0.1,
    axesColors = c("black", "black", "black"),
    fogStyle = "none",
    background = "white",
    simplePoints = TRUE,
    colors = c(replicate(
        length(Follins_domain1Boots$bootstraps), "orange"
    ), replicate(
        length(Follins_domain3Boots$bootstraps), "blue"
    )),
    group = oriTrivialGroup
)

# If you wish to save this figure, maximize the plot window on your screen before running this line. It will save to your working directory folder.
afterMaximizingWindow("Boots_Ahs_d2d3_1.png", "Boots_Ahs_d2d3_2.png")

# Plot the bootstrap comparison in an Equal Area plot
lineEqualAreaPlotTwo(c(
    lapply(Follins_domain2Boots$bootstraps, function(s)
        s[1, ]),
    lapply(Follins_domain2Boots$bootstraps, function(s)
        s[2, ])
),
c(
    lapply(Follins_domain3Boots$bootstraps, function(s)
        s[1, ]),
    lapply(Follins_domain3Boots$bootstraps, function(s)
        s[2, ])
))