#' @title
#' Plot heatmap for the functional regression surface
#'
#' @description
#' This function creates heatmaps for the functional regression surface in a
#' multivariate functional linear regression. Based on the fitting results from the
#' nested reduced-rank regression, different kinds of regression surfaces
#' (at the original scale or the latent scale) can be visualized to give a
#' clear illustration of the functional correlation between the user-specified
#' predictor (or latent predictor) trajectory and response
#' (or latent response) trajectory.
#'
#'
#' @usage
#' NRRR.plot.reg(Ag, Bg, Al, Bl, rx, ry, sseq, phi, tseq, psi,
#'               x_ind, y_ind, x_lab = NULL, y_lab = NULL,
#'               tseq_index = NULL, sseq_index = NULL,
#'               method = c("latent", "x_original",
#'               "y_original", "original")[1])
#'
#'
#' @param Ag,Bg,Al,Bl,rx,ry the estimated U, V, A, B, rx and ry from a NRRR fitting.
#' @param sseq the sequence of time points at which the predictor trajectory is observed.
#' @param phi the set of basis functions to expand the predictor trajectory.
#' @param tseq the sequence of time points at which the response trajectory is observed.
#' @param psi the set of basis functions to expand the response trajectory.
#' @param x_ind,y_ind two indices to locate the regression surface for which the heat map is to be drawn.
#'                    If \code{method = "original"}, then \eqn{0 < x_ind <= p, 0 < y_ind <= d}
#'                    and the function plots \eqn{C_{x_ind,y_ind}(s,t)} in Eq. (1) of the NRRR paper.
#'                    If \code{method = "latent"}, then \eqn{0 < x_ind <= rx, 0 < y_ind <= ry}
#'                    and the function plots \eqn{C^*_{x_ind,y_ind}(s,t)} in Eq. (2) of the NRRR paper.
#'                    If \code{method = "y_original"}, then \eqn{0 < x_ind <= rx, 0 < y_ind <= d}.
#'                    If \code{method = "x_original"}, then \eqn{0 < x_ind <= p, 0 < y_ind <= ry}.
#' @param x_lab,y_lab the user-specified x-axis (with x_lab for predictor) and
#'                    y-axis (with y_lab for response) label,
#'                    and it should be given as a character string, e.g., x_lab = "Temperature".
#' @param tseq_index,sseq_index the user-specified x-axis (with sseq_index for predictor)
#'                              and y-axis (with tseq_index for response) tick marks, and it should be
#'                              given as a vector of character strings of the same length as sseq or tseq, respectively.
#' @param method 'original': the function plots the correlation heatmap between the original
#'               functional response \eqn{y_i(t)} and the original functional predictor \eqn{x_j(s)};
#'               'latent': the function plots the correlation heatmap between
#'               the latent functional response \eqn{y^*_i(t)} and the latent functional predictor \eqn{x^*_j(s)};
#'               'y_original': the function plots the correlation heatmap between \eqn{y_i(t)} and \eqn{x^*_j(s)};
#'               'x_original': the function plots the correlation heatmap between \eqn{y^*_i(t)} and \eqn{x_j(s)}.
#'
#' @return A ggplot2 object.
#'
#' @details
#' More details and the examples of its usage can be found in the vignette of electricity demand analysis.
#'
#' @references
#' Liu, X., Ma, S., & Chen, K. (2020). Multivariate Functional Regression via Nested Reduced-Rank Regularization.
#' arXiv: Methodology.
#'
#' @import ggplot2
#' @importFrom reshape2 melt
#' @export

NRRR.plot.reg <- function(Ag, Bg, Al, Bl, rx, ry,
                                 sseq, phi, tseq, psi,
                                 x_ind, y_ind,
                                 x_lab = NULL, y_lab = NULL,
                                 tseq_index = NULL, sseq_index = NULL,
                                 method = c("latent", "x_original", "y_original","original")[1]
){
  ry <- ry
  rx <- rx
  Al <- Al
  Bl <- Bl
  Ag <- Ag
  Bg <- Bg
  p <- dim(Bg)[1]
  d <- dim(Ag)[1]
  ns <- dim(phi)[1]
  nt <- dim(psi)[1]
  jx <- dim(phi)[2]
  jy <- dim(psi)[2]

  if (method == "original" & any(c(0 > x_ind, x_ind > p, 0 > y_ind, y_ind > d))) stop("when 'original' is selected, 0 < x_ind <= p, 0 < y_ind <= d")
  if (method == "latent" & any(c(0 > x_ind, x_ind > rx, 0 > y_ind, y_ind > ry))) stop("when 'latent' is selected, 0 < x_ind <= rx, 0 < y_ind <= ry")
  if (method == "y_original" & any(c(0 > x_ind, x_ind > rx, 0 > y_ind, y_ind > d))) stop("when 'y_original' is selected, 0 < x_ind <= rx, 0 < y_ind <= d")
  if (method == "x_original" & any(c(0 > x_ind, x_ind > p, 0 > y_ind, y_ind > ry))) stop("when 'x_original' is selected, 0 < x_ind <= p, 0 < y_ind <= ry")

  Jpsi <- matrix(nrow = jy, ncol = jy, 0)
  tdiff <- (tseq - c(0, tseq[-nt]))
  for (t in 1:nt) Jpsi <- Jpsi + psi[t, ] %*% t(psi[t, ]) * tdiff[t]
  eJpsi <- eigen(Jpsi)
  Jpsihalf <- eJpsi$vectors %*% diag(sqrt(eJpsi$values)) %*% t(eJpsi$vectors)
  Jpsihalfinv <- eJpsi$vectors %*% diag(1 / sqrt(eJpsi$values)) %*% t(eJpsi$vectors)

  alindex <- rep(1:ry,jy)
  Alstar <- Al[order(alindex),]
  blindex <- rep(1:rx,jx)
  Blstar <- Bl[order(blindex),]

  Alstar <- kronecker(diag(ry),Jpsihalfinv)%*%Alstar
  Core <- Alstar%*%t(Blstar)

  comp <- array(dim = c(ry, rx, nt, ns), NA)

  for (i in 1:ry){
    for (j in 1:rx){
      comp[i,j,,] <- psi%*%Core[c(((i-1)*jy + 1):(i*jy)),
                                    c(((j-1)*jx + 1):(j*jx))]%*%t(phi) # nt \times ns
    }
  }

  if (method == "original"){
    comp.ori <- array(dim = c(d, p, nt, ns), NA)
    for (i in 1:nt){
      for (j in 1:ns){
        comp.ori[ , , i, j] <- Ag%*%comp[ , , i, j]%*%t(Bg)
      }
    }
    comp <- comp.ori
  } else if (method == "y_original"){
    comp.ori <- array(dim = c(d, rx, nt, ns), NA)
    for (i in 1:nt){
      for (j in 1:ns){
        comp.ori[ , , i, j] <- Ag%*%comp[ , , i, j]
      }
    }
    comp <- comp.ori
  } else if (method == "x_original"){
    comp.ori <- array(dim = c(ry, p, nt, ns), NA)
    for (i in 1:nt){
      for (j in 1:ns){
        comp.ori[ , , i, j] <- comp[ , , i, j]%*%t(Bg)
      }
    }
    comp <- comp.ori
  }

  # heatmap for coefficient function
  # library(ggplot2)
  # library(reshape2)

    dat <- reshape2::melt(as.matrix(comp[y_ind, x_ind, , ]))

    dat$Var1 <- factor(dat$Var1)
    dat$Var2 <- factor(dat$Var2)
    if (is.null(tseq_index)) {
      levels(dat$Var1) <- factor(tseq)
      y_breaks <- tseq[seq(1,nt,2)]
    } else {
      levels(dat$Var1) <- tseq_index
      y_breaks <- tseq_index[seq(1,nt,2)]
    }

    if (is.null(sseq_index)) {
      levels(dat$Var2) <- factor(sseq)
      x_breaks <- sseq[seq(1,ns,2)]
    } else {
      levels(dat$Var2) <- sseq_index
      x_breaks <- sseq_index[seq(1,ns,2)]
    }

    # Var1 <- dat$Var1
    # Var2 <- dat$Var2
    # value <- dat$value

    ggplot2::ggplot(dat, aes(Var2, Var1, fill = value)) + geom_tile() +
      scale_fill_gradient2(low = "blue", mid = "white",
                           high = "red", midpoint = 0, limits= range(comp),
                           space = "Lab", name = "",
                           na.value = "grey50", guide = "colourbar", aesthetics = "fill")+
      scale_x_discrete(breaks = x_breaks) +
      scale_y_discrete(breaks = y_breaks) +
      ylab(ifelse(is.null(y_lab), "response (t)", y_lab)) +
      xlab(ifelse(is.null(x_lab), "predictor (s)", x_lab)) +
      theme_classic() +
      theme(axis.text.x = element_text(angle = 45, hjust = 1)) +
      theme(axis.text=element_text(size=10),axis.title=element_text(size=13),
            plot.title = element_text(hjust = 0.5, size = 15))
}


# Explained Variance (R2) calculation for left-handedness (binomial trait)
# V.Odintsova, J.van Dongen, C.Dolan
# Based on Lee et al.
# Genetic Epidemiology 36 : 214–224 (2012)
# A Better Coefficient of Determination for Genetic Profile Analysis

########################### ############## ##############
# Load data
########################### ############## ##############
# polygenic risk scores (PGS)
# methylation scores (MS)
# EWAS covariates
# PGS covariates

########################### ############## ##############
# Select variables for the prediction models
########################### ############## ##############
# FISNumber
# FamilyNumber
# sex
# age (age at DNA bloodsampling, scaled)
# hand01 (left-handedness variable: right-handed = 0, left-handed = 1)

# PGS and GWAS covariates
# PGS
# PC1 - PC10 (principal components)
# platform2 - platform6 (platform dummy variables)

# Methylation scores and EWAS covariates from adjusted model
# MS1 (methylation score calculated with weights at p-value thresholds 10^-1)
# MS3 (methylation score calculated with weights at p-value thresholds 10^-3)
# MS5 (methylation score calculated with weights at p-value thresholds 10^-5)
# smoking
# BMI
# Neut_Perc (neutrophils percentage)
# Mono_Perc (monocytes percentage)
# Eos_Perc (eosinophils percentage)
# Array_rownum (array rownumber)
# Sample_Plate (sample plate)
# Sample_Plate1 - Sample_PlateN (dummy variables for sample plate)

########################### ############## ##############
# Logistic regression to calculate R2 for PGS and MS
############################ ############## ##############
# function for logistic regression model with PGS and covariates prints PGS effect size, se, p-value, R2
PGSlogit_function = function(NTRdata) {
                  dat = NTRdata
                  hand01=as.numeric(dat$hand01)
                  sex=as.numeric(dat$sex)
                  PGS=as.numeric(dat$PGS)
                  PC1 = as.numeric(dat$PC1)
                  PC2 = as.numeric(dat$PC2)
                  PC3 = as.numeric(dat$PC3)
                  PC4 = as.numeric(dat$PC4)
                  PC5 = as.numeric(dat$PC5)
                  PC6 = as.numeric(dat$PC6)
                  PC7 = as.numeric(dat$PC7)
                  PC8 = as.numeric(dat$PC8)
                  PC9 = as.numeric(dat$PC9)
                  PC10 = as.numeric(dat$PC10)
                  platform2=as.numeric(dat$platform2)
                  platform3=as.numeric(dat$platform3)
                  platform5=as.numeric(dat$platform5)
                  platform6=as.numeric(dat$platform6)
                  # MODEL 1: handedness ~ PGS + GWAS covariates
                  vare1 = pi^2 / 3     # residual (homoskedastic) variance
                  r1=glm(hand01~PGS+sex+ platform2+platform3+platform5+platform6+PC1+PC2+PC3+PC4+PC5+PC6+PC7+PC8+PC9+PC10,
                         family=binomial(link='logit'))  # logistic regression
                  vp1=var( r1$coefficients["PGS"]*PGS+
                          r1$coefficients["sex"]*sex+
                          r1$coefficients["platform2"]*platform2+
                          r1$coefficients["platform3"]*platform3+
                          r1$coefficients["platform5"]*platform5+
                          r1$coefficients["platform6"]*platform6+
                          r1$coefficients["PC1"]*PC1+
                          r1$coefficients["PC2"]*PC2+
                          r1$coefficients["PC3"]*PC3+
                          r1$coefficients["PC4"]*PC4+
                          r1$coefficients["PC5"]*PC5+
                          r1$coefficients["PC6"]*PC6+
                          r1$coefficients["PC7"]*PC7+
                          r1$coefficients["PC8"]*PC8+
                          r1$coefficients["PC9"]*PC9+
                          r1$coefficients["PC10"]*PC10)    # variance explained
                  R2logist1 = vp1 / (vp1+vare1)   # explained var  / total var
                  # MODEL 2: handedness ~ GWAS covariates
                  r2=glm(hand01~sex+platform2+platform3+platform5+platform6+PC1+PC2+PC3+PC4+PC5+PC6+PC7+PC8+PC9+PC10,
                         family=binomial(link='logit'))
                  vp2=var(
                            r2$coefficients["sex"]*sex+
                            r2$coefficients["platform2"]*platform2+
                            r2$coefficients["platform3"]*platform3+
                            r2$coefficients["platform5"]*platform5+
                            r2$coefficients["platform6"]*platform6+
                            r1$coefficients["PC1"]*PC1+
                            r1$coefficients["PC2"]*PC2+
                            r1$coefficients["PC3"]*PC3+
                            r1$coefficients["PC4"]*PC4+
                            r1$coefficients["PC5"]*PC5+
                            r1$coefficients["PC6"]*PC6+
                            r1$coefficients["PC7"]*PC7+
                            r1$coefficients["PC8"]*PC8+
                            r1$coefficients["PC9"]*PC9+
                            r1$coefficients["PC10"]*PC10)    # variance explained
                  R2logist2 = vp2 / (vp2+vare1)   # explained var  / total var
                  t = matrix(0, 1, 6)
                  t[1,1] = summary(r1)$coefficients[2,1]
                  t[1,2] = summary(r1)$coefficients[2,2]
                  t[1,3] = summary(r1)$coefficients[2,4]
                  t[1,4] = R2logist1 - R2logist2
                  print(t)
}

# function for logistic regression model with MS, PGS and covariates prints MS effect size, se, p-value, R2
MSlogit_function = function(NTRdata, MScore) {
                  dat = NTRdata
                  MS = as.numeric(MScore)
                    hand01=as.numeric(dat$hand01)
                    age=as.numeric(dat$age)
                    sex=as.numeric(dat$sex)
                    PGS=as.numeric(dat$PGS)
                    PC1 = as.numeric(dat$PC1)
                    PC2 = as.numeric(dat$PC2)
                    PC3 = as.numeric(dat$PC3)
                    PC4 = as.numeric(dat$PC4)
                    PC5 = as.numeric(dat$PC5)
                    PC6 = as.numeric(dat$PC6)
                    PC7 = as.numeric(dat$PC7)
                    PC8 = as.numeric(dat$PC8)
                    PC9 = as.numeric(dat$PC9)
                    PC10 = as.numeric(dat$PC10)
                    platform1=as.numeric(dat$platform1) # not used in models
                    platform2=as.numeric(dat$platform2)
                    platform3=as.numeric(dat$platform3)
                    platform5=as.numeric(dat$platform5)
                    platform6=as.numeric(dat$platform6)
                    smoking = as.numeric(dat$smoking)
                    BMI = as.numeric(dat$BMI)
                    Neut_Perc = as.numeric(dat$Neut_Perc)
                    Mono_Perc = as.numeric(dat$Mono_Perc)
                    Eos_Perc = as.numeric(dat$Eos_Perc)
                    Sample_Plate = as.factor(dat$Sample_Plate)
                    Sample_Plate1 = ifelse(Sample_Plate== 1, 1, 0)
                    Sample_Plate2 = ifelse(Sample_Plate== 2, 1, 0)
                    Sample_Plate3 = ifelse(Sample_Plate== 3, 1, 0)
                    Sample_Plate4 = ifelse(Sample_Plate== 4, 1, 0)
                    Sample_Plate5 = ifelse(Sample_Plate== 5, 1, 0)
                    Sample_Plate6 = ifelse(Sample_Plate== 6, 1, 0)
                    Sample_Plate7 = ifelse(Sample_Plate== 7, 1, 0)
                    Sample_Plate8 = ifelse(Sample_Plate== 8, 1, 0)
                    Sample_Plate9 = ifelse(Sample_Plate== 9, 1, 0)
                    Sample_Plate10 = ifelse(Sample_Plate== 10, 1, 0)
                    Sample_Plate11 = ifelse(Sample_Plate== 11, 1, 0)
                    Sample_Plate12 = ifelse(Sample_Plate== 12, 1, 0)
                    Sample_Plate13 = ifelse(Sample_Plate== 13, 1, 0)
                    Sample_Plate14 = ifelse(Sample_Plate== 14, 1, 0)
                    Sample_Plate15 = ifelse(Sample_Plate== 15, 1, 0)
                    Sample_Plate16 = ifelse(Sample_Plate== 16, 1, 0)
                    Sample_Plate17 = ifelse(Sample_Plate== 17, 1, 0)
                    Sample_Plate18 = ifelse(Sample_Plate== 18, 1, 0)
                    Sample_Plate19 = ifelse(Sample_Plate== 19, 1, 0)
                    Sample_Plate20 = ifelse(Sample_Plate== 20, 1, 0)
                    Sample_Plate21 = ifelse(Sample_Plate== 21, 1, 0)
                    Sample_Plate22 = ifelse(Sample_Plate== 22, 1, 0)
                    Sample_Plate23 = ifelse(Sample_Plate== 23, 1, 0)
                    Sample_Plate24 = ifelse(Sample_Plate== 24, 1, 0)
                    Sample_Plate25 = ifelse(Sample_Plate== 25, 1, 0)
                    Sample_Plate26 = ifelse(Sample_Plate== 26, 1, 0)
                    Sample_Plate27 = ifelse(Sample_Plate== 27, 1, 0)
                    Sample_Plate28 = ifelse(Sample_Plate== 28, 1, 0)
                    Sample_Plate29 = ifelse(Sample_Plate== 29, 1, 0)
                    Sample_Plate30 = ifelse(Sample_Plate== 30, 1, 0)
                    Sample_Plate31 = ifelse(Sample_Plate== 31, 1, 0)
                    Sample_Plate32 = ifelse(Sample_Plate== 32, 1, 0)
                    Sample_Plate33 = ifelse(Sample_Plate== 33, 1, 0)
                    Sample_Plate34 = ifelse(Sample_Plate== 34, 1, 0)
                    # MODEL 1: handedness ~ MS+PGS + covariates
                    vare1 = pi^2 / 3     # residual (homoskedastic) variance
                    r1=glm(hand01~PGS+age+sex+
                             platform2+platform3+platform5+platform6+PC1+PC2+PC3+PC4+PC5+PC6+PC7+PC8+PC9+PC10+
                             smoking+BMI+Neut_Perc+Mono_Perc+Eos_Perc+Sample_Plate+MS,
                             family=binomial(link='logit'))  # logistic regression
                    vp1=var( r1$coefficients["PGS"]*PGS+
                            r1$coefficients["age"]*age+
                            r1$coefficients["sex"]*sex+
                            r1$coefficients["platform2"]*platform2+
                            r1$coefficients["platform3"]*platform3+
                            r1$coefficients["platform5"]*platform5+
                            r1$coefficients["platform6"]*platform6+
                            r1$coefficients["PC1"]*PC1+
                            r1$coefficients["PC2"]*PC2+
                            r1$coefficients["PC3"]*PC3+
                            r1$coefficients["PC4"]*PC4+
                            r1$coefficients["PC5"]*PC5+
                            r1$coefficients["PC6"]*PC6+
                            r1$coefficients["PC7"]*PC7+
                            r1$coefficients["PC8"]*PC8+
                            r1$coefficients["PC9"]*PC9+
                            r1$coefficients["PC10"]*PC10+
                            r1$coefficients["smoking"]*smoking+
                            r1$coefficients["BMI"]*BMI+
                            r1$coefficients["Neut_Perc"]*Neut_Perc+
                            r1$coefficients["Mono_Perc"]*Mono_Perc+
                            r1$coefficients["Eos_Perc"]*Eos_Perc+
                            r1$coefficients["Sample_Plate2"]*Sample_Plate2+
                            r1$coefficients["Sample_Plate3"]*Sample_Plate3+
                            r1$coefficients["Sample_Plate4"]*Sample_Plate4+
                            r1$coefficients["Sample_Plate5"]*Sample_Plate5+
                            r1$coefficients["Sample_Plate6"]*Sample_Plate6+
                            r1$coefficients["Sample_Plate7"]*Sample_Plate7+
                            r1$coefficients["Sample_Plate8"]*Sample_Plate8+
                            r1$coefficients["Sample_Plate9"]*Sample_Plate9+
                            r1$coefficients["Sample_Plate10"]*Sample_Plate10+
                            r1$coefficients["Sample_Plate11"]*Sample_Plate11+
                            r1$coefficients["Sample_Plate12"]*Sample_Plate12+
                            r1$coefficients["Sample_Plate13"]*Sample_Plate13+
                            r1$coefficients["Sample_Plate14"]*Sample_Plate14+
                            r1$coefficients["Sample_Plate15"]*Sample_Plate15+
                            r1$coefficients["Sample_Plate16"]*Sample_Plate16+
                            r1$coefficients["Sample_Plate17"]*Sample_Plate17+
                            r1$coefficients["Sample_Plate18"]*Sample_Plate18+
                            r1$coefficients["Sample_Plate19"]*Sample_Plate19+
                            r1$coefficients["Sample_Plate20"]*Sample_Plate20+
                            r1$coefficients["Sample_Plate21"]*Sample_Plate21+
                            r1$coefficients["Sample_Plate22"]*Sample_Plate22+
                            r1$coefficients["Sample_Plate23"]*Sample_Plate23+
                            r1$coefficients["Sample_Plate24"]*Sample_Plate24+
                            r1$coefficients["Sample_Plate25"]*Sample_Plate25+
                            r1$coefficients["Sample_Plate26"]*Sample_Plate26+
                            r1$coefficients["Sample_Plate27"]*Sample_Plate27+
                            r1$coefficients["Sample_Plate28"]*Sample_Plate28+
                            r1$coefficients["Sample_Plate29"]*Sample_Plate29+
                            r1$coefficients["Sample_Plate30"]*Sample_Plate30+
                            r1$coefficients["Sample_Plate31"]*Sample_Plate31+
                            r1$coefficients["Sample_Plate32"]*Sample_Plate32+
                            r1$coefficients["Sample_Plate33"]*Sample_Plate33+
                            r1$coefficients["Sample_Plate34"]*Sample_Plate34+
                            r1$coefficients["MS"]*MS)   # variance explained
                    R2logist1 = vp1 / (vp1+vare1)   # explained var  / total var
                    # MODEL 2: handedness ~ PGS +  covariates
                    r2=glm(hand01~PGS5+age+sex+platform2+platform3+platform5+platform6+
                                PC1+PC2+PC3+PC4+PC5+PC6+PC7+PC8+PC9+PC10+
                                smoking+BMI+Neut_Perc+Mono_Perc+Eos_Perc+Sample_Plate,
                                family=binomial(link='logit'))

                    vp2=var(            r2$coefficients["PGS"]*PGS+
                                        r2$coefficients["age"]*age+
                                        r2$coefficients["sex"]*sex+
                                        r2$coefficients["platform2"]*platform2+
                                        r2$coefficients["platform3"]*platform3+
                                        r2$coefficients["platform5"]*platform5+
                                        r2$coefficients["platform6"]*platform6+
                                        r2$coefficients["PC1"]*PC1+
                                        r2$coefficients["PC2"]*PC2+
                                        r2$coefficients["PC3"]*PC3+
                                        r2$coefficients["PC4"]*PC4+
                                        r2$coefficients["PC5"]*PC5+
                                        r2$coefficients["PC6"]*PC6+
                                        r2$coefficients["PC7"]*PC7+
                                        r2$coefficients["PC8"]*PC8+
                                        r2$coefficients["PC9"]*PC9+
                                        r2$coefficients["PC10"]*PC10+
                                        r2$coefficients["smoking"]*smoking+
                                        r2$coefficients["BMI"]*BMI+
                                        r2$coefficients["Neut_Perc"]*Neut_Perc+
                                        r2$coefficients["Mono_Perc"]*Mono_Perc+
                                        r2$coefficients["Eos_Perc"]*Eos_Perc+
                                        r2$coefficients["Sample_Plate2"]*Sample_Plate2+
                                        r2$coefficients["Sample_Plate3"]*Sample_Plate3+
                                        r2$coefficients["Sample_Plate4"]*Sample_Plate4+
                                        r2$coefficients["Sample_Plate5"]*Sample_Plate5+
                                        r2$coefficients["Sample_Plate6"]*Sample_Plate6+
                                        r2$coefficients["Sample_Plate7"]*Sample_Plate7+
                                        r2$coefficients["Sample_Plate8"]*Sample_Plate8+
                                        r2$coefficients["Sample_Plate9"]*Sample_Plate9+
                                        r2$coefficients["Sample_Plate10"]*Sample_Plate10+
                                        r2$coefficients["Sample_Plate11"]*Sample_Plate11+
                                        r2$coefficients["Sample_Plate12"]*Sample_Plate12+
                                        r2$coefficients["Sample_Plate13"]*Sample_Plate13+
                                        r2$coefficients["Sample_Plate14"]*Sample_Plate14+
                                        r2$coefficients["Sample_Plate15"]*Sample_Plate15+
                                        r2$coefficients["Sample_Plate16"]*Sample_Plate16+
                                        r2$coefficients["Sample_Plate17"]*Sample_Plate17+
                                        r2$coefficients["Sample_Plate18"]*Sample_Plate18+
                                        r2$coefficients["Sample_Plate19"]*Sample_Plate19+
                                        r2$coefficients["Sample_Plate20"]*Sample_Plate20+
                                        r2$coefficients["Sample_Plate21"]*Sample_Plate21+
                                        r2$coefficients["Sample_Plate22"]*Sample_Plate22+
                                        r2$coefficients["Sample_Plate23"]*Sample_Plate23+
                                        r2$coefficients["Sample_Plate24"]*Sample_Plate24+
                                        r2$coefficients["Sample_Plate25"]*Sample_Plate25+
                                        r2$coefficients["Sample_Plate26"]*Sample_Plate26+
                                        r2$coefficients["Sample_Plate27"]*Sample_Plate27+
                                        r2$coefficients["Sample_Plate28"]*Sample_Plate28+
                                        r2$coefficients["Sample_Plate29"]*Sample_Plate29+
                                        r2$coefficients["Sample_Plate30"]*Sample_Plate30+
                                        r2$coefficients["Sample_Plate31"]*Sample_Plate31+
                                        r2$coefficients["Sample_Plate32"]*Sample_Plate32+
                                        r2$coefficients["Sample_Plate33"]*Sample_Plate33+
                                        r2$coefficients["Sample_Plate34"]*Sample_Plate34
                    )    # variance explained
                    R2logist2 = vp2 / (vp2+vare1)   # explained var  / total var
                    t = matrix(0, 1, 6)
                    t[1,1] = summary(r1)$coefficients[57,1]
                    t[1,2] = summary(r1)$coefficients[57,2]
                    t[1,3] = summary(r1)$coefficients[57,4]
                    t[1,6] = R2logist1 - R2logist2
                    print(t)
}


#'@title calcKristPwr
#'
#' @description Calculate instantaneous main engine power (kW) using the
#' Kristensen 2016 updates of Harvald equations.
#'
#' @param totalInstalledPwr Total installed main engine power (vector of
#'  numericals, kW) (maximum
#' continuous rated power)
#' @param shipSpeed Ship actual speed (vector of numericals, m/s) (see
#'  \code{\link{calcSpeedUnitConversion}})
#' @param actualDraft Actual draft (vector of numericals, m)
#' @param maxDraft Maximum summer load line draft (vector of numericals, m)
#' @param shipType Ship type (vector of strings, see \code{\link{calcShipType}}),
#' determined by Stat 5 code: \itemize{
#' \item"container.ship"
#' \item"bulk.carrier"
#' \item"tanker"
#' \item"general.cargo"
#' \item"vehicle.carrier"
#' \item"reefer"
#' \item"ro.ro"
#' \item"passenger"
#' \item"tug"
#' \item"misc"
#'}
#' @param lwl Waterline length (vector of numericals, m) (see
#' \code{\link{calclwl}})
#' @param breadth Moulded breadth (vector of numericals, m)
#' @param maxDisplacement Maximum ship displacement (vector of numericals, m^3)
#' @param Cb Maximum block coefficient (vector of numericals, dimensionless)
#' (see \code{\link{calcCb}})
#' @param nProp Number of propellers (vector of numericals, see
#' \code{\link{calcPropNum}})
#' @param dwt Ship maximum deadweight tonnage (vector of numericals, tonnage)
#' @param serviceMargin A service margin to account for weather and sea effects:
#' \itemize{\item At-sea operations = 15 (Default) \item Coastal operations = 10}
#' Can supply either a vector of numericals, a single number, or rely on the
#' default
#' @param shaftEff Shaft efficiency (dimensionless). Default = 0.98.
#'  Ratio of power delivered to the propeller and the brake power delivered by
#'  the engine. Can supply either a vector of numericals, a single number, or
#'  rely on the default
#' @param relRotationEff Relative rotational efficiency (dimensionless).
#'  Default = 1. Accounts for effect of rotational flow of water around propeller.
#'  Can supply either a vector of numericals, a single number, or rely on the
#'  default
#' @param seawaterTemp Sea water temperature. Default = 15 (degrees Celsius). Can
#' supply either a vector of numericals, a single number, or rely on the default
#' @param seawaterDensity Sea water density. Default = 1.025 (g/cm^3). Can
#' supply either a vector of numericals, a single number, or rely on the default
#' @param pwrUpperBoundPercent Percent of total installed power at which
#' required power is capped. Default = 1, which indicates required power cannot
#' exceed \code{totalInstalledPwr}. Can supply either a vector of numericals, a
#' single number, or rely on the default
#' @param pwrLowerBoundPercent Percent of total installed power to act as lower
#' bound for required power. Default = 0.02, which indicates required power
#' cannot go below 2\% of \code{totalInstalledPwr}. Can supply either a vector
#' of numericals, a single number, or rely on the default
#' @param CmEquationType Type of equation to estimate the midship section
#' coefficient (see \code{\link{calcCm}}): \itemize{
#' \item"kristensen"
#' \item"benford"
#' \item"schneekluth"}
#' @param tankerBulkCarrierShipTypes Ship types specified in input
#' \code{shipTypes} to be modeled as tankers and bulk carriers.
#' @param tugShipTypes Ship types specified in input \code{shipTypes} to be
#' modeled as tugs
#' @param roroPaxShipTypes Ship types specified in input \code{shipTypes} to be
#' modeled as RORO and passenger ships
#' @param gCargoShipTypes Ship types specified in input \code{shipTypes} to be
#' modeled as general cargo
#' @param containerShipTypes Ship types specified in input \code{shipTypes} to
#' be modeled as container ships
#'
#' @details
#' Primary method from Kristensen (2013). Estimation of some inputs use
#' methodology from Rakke (2016).
#'
#' This method this requires ship types to be grouped. Use the
#' \code{tankerBulkCarrierShipTypes}, \code{tugShipTypes}, \code{roroPaxShipTypes},
#' \code{gCargoShipTypes}, \code{containerShipTypes} grouping parameters to
#' provide these ship type groupings. Any ship types not included in these groupings
#' will be considered as miscellaneous vessels.
#'
#' Ship speed and actual draft are typically obtained from sources such as AIS
#' messages or ship records.
#'
#' @return power (vector of numericals, kW)
#'
#' @references
#'Kristensen, H. O. and Lutzen, M. 2013. "Prediction of Resistance and Propulsion
#'Power of Ships."
#'
#'Kristensen, H. O. 2016. "Revision of statistical analysis and determination of
#'regression formulas for main dimensions of container ships based on data from
#'Clarkson."
#'
#'\href{https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}{Kristensen, H. O.
#'"Ship-Desmo-Tool." https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}
#'
#'\href{http://hdl.handle.net/11250/2410741}{Rakke, S. G. 2016. "Ship Emissions
#'Calculation from AIS." NTNU.}
#'
#'@seealso \itemize{
#'\item \code{\link{calclwl}}
#'\item \code{\link{calcSpeedUnitConversion}}
#'\item \code{\link{calcCb}}
#'\item \code{\link{calcPropNum}}
#'\item \code{\link{calcCm}}
#'\item \code{\link{calcShipType}}
#'\item \code{vignette("OverviewOfPowerModels", package="ShipPowerModel")}
#'\item \code{vignette("Kristensen.Example", package="ShipPowerModel")}
#'}
#'
#' @family Kristensen Calculations
#'
#' @examples
#' calcKristPwr(
#' totalInstalledPwr=9363,
#' shipSpeed=10.8,
#' actualDraft=12.5,
#' maxDraft=13.6,
#' shipType="bulk.carrier",
#' lwl=218,
#' breadth=32.25,
#' maxDisplacement=80097,
#' Cb=0.8162717,
#' nProp=1,
#' dwt=70000,
#' serviceMargin=15,
#' shaftEff=0.98,
#' relRotationEff=1,
#' seawaterTemp=15,
#' seawaterDensity=1.025,
#' CmEquationType="kristensen"
#' )
#'
#' @export

calcKristPwr<-function(
  totalInstalledPwr,
  shipSpeed,
  actualDraft,
  maxDraft,
  shipType,
  lwl,
  breadth,
  maxDisplacement,
  Cb,
  nProp,
  dwt,
  serviceMargin=15,
  shaftEff=0.98,
  relRotationEff=1,
  seawaterTemp=15,
  seawaterDensity=1.025,
  pwrUpperBoundPercent=1,
  pwrLowerBoundPercent=0.02,
  CmEquationType="kristensen",
  tankerBulkCarrierShipTypes=c("tanker","chemical.tanker","liquified.gas.tanker","oil.tanker","other.tanker","bulk.carrier"),
  tugShipTypes=c("service.tug","tug"),
  roroPaxShipTypes=c("passenger","ferry.pax","ferry.ro.pax","cruise","cruise.ed","yacht","ro.ro"),
  gCargoShipTypes=c("general.cargo"),
  containerShipTypes=c("container.ship")
){
#=============================================================
  #Inputs
  Cbw<-calcCbw(Cb, actualDraft, maxDraft)
  actualDisplacement<-calcActualDisp(Cb,
                                     Cbw,
                                     actualDraft,
                                     maxDraft,
                                     maxDisplacement)

  Cm<-calcCm(shipType,Cbw,maxDraft,actualDraft,CmEquationType,tankerBulkCarrierShipTypes,tugShipTypes,roroPaxShipTypes)

  Cp<-calcCp(Cm, Cbw,shipType, bounds="none", roroPaxContainerShipTypes=union(roroPaxShipTypes,containerShipTypes),gCargoShipTypes,tankerBulkCarrierShipTypes)

  propDiam<-calcPropDia(shipType, maxDraft,tankerBulkCarrierGCargoShipTypes=union(tankerBulkCarrierShipTypes,gCargoShipTypes),containerShipTypes)

  M<- calcShipM(actualDisplacement,lwl)

  froudeNum<- calcFroudeNum(shipSpeed,lwl)
#=============================================================
#Frictional Resistance
Cf<-calcCf(shipSpeed, lwl, seawaterTemp, seawaterDensity)
#=============================================================
#Wetted Surface Area
wettedSA<-calcKristWettedSA(shipType,
                            maxDisplacement,
                            maxDraft,
                            actualDraft,
                            lwl,
                            breadth,
                            Cbw,
                            seawaterDensity,
                            tankerBulkCarrierGCargoShipTypes=union(tankerBulkCarrierShipTypes,gCargoShipTypes),
                            containerShipTypes,
                            paxTugShipTypes=union(roroPaxShipTypes, tugShipTypes))

#==========================================
#Incremental Resistance
Ca<-calcKristCa(shipType,actualDisplacement,
                tankerBulkCarrierGCargoShipTypes=union(tankerBulkCarrierShipTypes,gCargoShipTypes),
                containerShipTypes)
#==========================================
#Air Resistance
Caa<-calcKristCaa(shipType,dwt,
                  tankerBulkCarrierGCargoShipTypes=union(tankerBulkCarrierShipTypes,gCargoShipTypes),
                  containerShipTypes)
#==========================================
#Residual Resistance
Cr<-calcKristCr(shipType,
                M,
                froudeNum,
                actualDraft,
                breadth,
                Cp,
                tankerBulkCarrierGCargoShipTypes=union(tankerBulkCarrierShipTypes,gCargoShipTypes))

#=============================================================
#Thrust Deduction Factor
t<-calcKristThrustFactor(
                    shipType,
                    breadth,
                    lwl,
                    Cbw,
                    propDiam,
                    M,
                    nProp,
                    tankerBulkCarrierShipTypes)
#=============================================================
#Wake Fraction
#*(1-t)/(1-w)= hull efficiency
w<-calcKristWakeFrac(shipType,
                breadth,
                lwl,
                Cbw,
                propDiam,
                M,
                nProp,
                tankerBulkCarrierShipTypes)
#==========================================================
#Total Resistance
R<-calcKristTotalRes(wettedSA,
                     Cf,
                     Cr,
                     Ca,
                     Caa,
                     seawaterDensity,
                     shipSpeed,
                     serviceMargin)
#==========================================================
# Open Water Efficiency
no<-calcOpenWaterEff(R,
                     t,
                     nProp,
                     w,
                     propDiam,
                     shipSpeed,
                     seawaterDensity)
#=============================================================
power<-calcResistanceShipPwr(R,shipSpeed, (1-t)/(1-w), no,totalInstalledPwr, shaftEff,relRotationEff, pwrUpperBoundPercent, pwrLowerBoundPercent)

return(power)
}


########################### BayesKin ###########################
## Code for a comparison of Bayesian models for inverse kinematic problems
## with Maximum likelihood estimation

## Written by: Andy Pohl
## UofC - Faculty of Kinesiology
## June-Dec 2020
## Revision 1: May 2020
################################################################

########################### library.r  #########################
## Contains source functions for BayesKin package - modified for performing
## sensitivity to measurement noise.
################################################################
##### Preliminaries
library(coda)
library(numDeriv)
library(rjags)

# set number of parallel cores
options(mc.cores=4)

########################### Section A  #########################
## Miscellaneous helper functions
################################################################
rt_ls <- function(n, df, mu, a){
  # Generates a random sample from the location scale student-t distribution with mean mu, degrees of freedom
  # df and scale a

  return(rt(n,df)*a + mu)
}

mod_180 = function(angles){
  # Rescales an angle on the (-inf, inf) scale to the interval [-180, 180)
  temp = sapply(angles, function(ang){
    if(abs(ang)<=180){
      return(ang)
    }else if(ang >180){
      remainder = ang-180
      return(-180 + remainder)
    }else if (ang< -180){
      remainder = ang + 180
      return(180 + remainder)
    }
  })
  return(temp)
}
########################### Section B  #########################
## Links, posture and observations
################################################################

gen_link = function(seg.length,
                    plate.center,
                    marker.plate.dim = c(0.1, 0.05)){
  # Creates a rigid link of length seg.length with a marker plate at plate.center
  # with dimensions plate.placement.length*seg.length and markers at each corner
  # of the marker plate.
  plate.loc = c(plate.center,0)

  markers = cbind(c(plate.loc[1]-marker.plate.dim[1]/2, marker.plate.dim[2]/2),
                  c(plate.loc[1]+marker.plate.dim[1]/2, marker.plate.dim[2]/2),
                  c(plate.loc[1]+marker.plate.dim[1]/2, -marker.plate.dim[2]/2),
                  c(plate.loc[1]-marker.plate.dim[1]/2, -marker.plate.dim[2]/2))

  link = list(length = seg.length,
              x = markers)
  return(link)

}

gen_posture = function(links,
                       r,
                       theta,
                       degrees = TRUE){
  # Applies rigid body transformations translation r and rotation by theta to a
  # collection of links.  Note that links must be specified in order (link 1 -
  # link 3) and the length of theta should be the same as the number of links

  if(degrees){
    theta = theta * pi/180
  }

  nlinks = length(links)

  #  Generate rotation matrix for eac link
  Gam = vector('list', nlinks)
  for(i in 1:nlinks){
    Gam[[i]] = matrix(c(cos(theta[i]), sin(theta[i]),
                        -sin(theta[i]), cos(theta[i])), 2, 2)
  }

  # Compute the locations of joint centers for each link in the kinematic chain
  J = vector('list', nlinks+1)
  J[[1]] = matrix(r, 2, 1)
  for(i in 2:(nlinks+1)){
    J[[i]] =  J[[i-1]] + Gam[[i-1]] %*% c(links[[i-1]]$length, 0)
  }

  # Compute actual position of markers alpha in the global coordinate system
  alpha = vector(mode = 'list', length = nlinks)
  for(i in 1:nlinks){
    nmarkers = ncol(links[[i]]$x)
    alpha[[i]] = matrix(NA, 2, nmarkers)
    for(j in 1:nmarkers){
      x_ij = links[[i]]$x[,j]
      alpha[[i]][,j] = J[[i]] + Gam[[i]]%*%x_ij
    }
  }

  # Return posture object (list with origin r, joint position J and expected marker position alpha)
  return(list(r = r, J = J, alpha = alpha))
}

gen_obs = function(posture, sigma){
  # Generates an observation of a given posture subject to random 0 mean gaussain
  # noise with standard deviation sigma
  nlinks = length(posture$alpha)

  y = vector(mode = 'list', length = nlinks)
  for(i in 1:nlinks){
    nmarkers = ncol(links[[i]]$x)
    y[[i]] = matrix(NA, ncol = nmarkers, nrow = 2)
    for(j in 1:nmarkers){
      y[[i]][1, j] = rnorm(1, mean = posture$alpha[[i]][1, j], sd= sigma)
      y[[i]][2, j] = rnorm(1, mean = posture$alpha[[i]][2, j], sd= sigma)
    }
  }
  return(y)
}

plot_system= function(posture, y=NA, ...){
  # Plots the Kinematic chain and observed markers y
  plot(NULL, NULL,
       xlim = c(0, 1), ylim = c(-1, 0.1),
       xlab = 'x', ylab = 'y')
  points(t(do.call(cbind,posture$J)), type = "o", pch = 20,...) # mu_tmp
  orange = rgb(255/255, 140/255, 0, 0.5)
  blue = rgb(0, 0, 1, 0.5)
  points(t(do.call(cbind,posture$alpha)), pch = 20, col = orange)
  if(!is.na(y)){
    points(t(do.call(cbind,y)), pch = 20, col = blue)
  }
}

########################### Section C  #########################
## Maximum likelihood/Least Squares solutions.
################################################################

cost = function(params, y, links){
  # The cost function (equation ... in paper) given a vector
  # params = vector of parameters (rx, ry, theta1, theta2, theta3 ), observations
  # y and known positions of markers within the LCS of each link links.
  nlinks = length(links)
  r.hat = params[1:2]
  theta.hat = params[3:(nlinks+2)]

  obs.hat = gen_posture(links, r = r.hat, theta = theta.hat)
  y.hat = obs.hat$alpha

  cost = 0
  for(i in 1:nlinks){
    nmarkers = ncol(y.hat[[i]])
    for(j in 1:nmarkers){
      diff = y.hat[[i]][,j] - y[[i]][,j]
      cost = cost+ (sqrt(sum(diff*diff)))
    }
  }
  return(cost)
}

costgrad_analytical = function(x, y, links){
  # Generates the analytical gradient of the cost function.  Originally generated
  # via MATLAB and the symbolic toolbox.
  # x is parameter vector (r1, r2, theta1, theta2, theta3)
  # y[[nlinks]] of observations
  # links contains the known LCS information of each link

  nlinks = length(links)
  r1 = x[1]
  r2 = x[2]
  x111 = links[[1]]$x[1,1]
  x112 = links[[1]]$x[2,1]
  x121 = links[[1]]$x[1,2]
  x122 = links[[1]]$x[2,2]
  x131 = links[[1]]$x[1,3]
  x132 = links[[1]]$x[2,3]
  x141 = links[[1]]$x[1,4]
  x142 = links[[1]]$x[2,4]

  y111 = y[[1]][1,1]
  y112 = y[[1]][2,1]
  y121 = y[[1]][1,2]
  y122 = y[[1]][2,2]
  y131 = y[[1]][1,3]
  y132 = y[[1]][2,3]
  y141 = y[[1]][1,4]
  y142 = y[[1]][2,4]

  if(nlinks ==1){
    theta1 = x[3]
    gradient = rep(NA, 3)
    for(i in 1:length(x)){
        gradtxtfile = sprintf("./AnalyticalCostGrad/SingleLink/CostGradientChar%.0f.txt", i)
        gradtxt = read.table(file = gradtxtfile, sep=',' , stringsAsFactors = FALSE)
        if((dim(gradtxt)[1] != 1) & (dim(gradtxt)[2] !=1)){stop("Error in reading gradient text file!")}
        gradient[i] = eval(parse(text = gradtxt[1,1]))
      }
    }else if(nlinks ==2){
    theta1 = x[3]
    theta2 = x[4]

    L1 = links[[1]]$length

    x211 = links[[2]]$x[1,1]
    x212 = links[[2]]$x[2,1]
    x221 = links[[2]]$x[1,2]
    x222 = links[[2]]$x[2,2]
    x231 = links[[2]]$x[1,3]
    x232 = links[[2]]$x[2,3]
    x241 = links[[2]]$x[1,4]
    x242 = links[[2]]$x[2,4]

    y211 = y[[2]][1,1]
    y212 = y[[2]][2,1]
    y221 = y[[2]][1,2]
    y222 = y[[2]][2,2]
    y231 = y[[2]][1,3]
    y232 = y[[2]][2,3]
    y241 = y[[2]][1,4]
    y242 = y[[2]][2,4]

    gradient = rep(NA, 4)
    for(i in 1:length(x)){
      gradtxtfile = sprintf("./AnalyticalCostGrad/DoubleLink/CostGradientChar%.0f.txt", i)
      gradtxt = read.table(file = gradtxtfile, sep=',' , stringsAsFactors = FALSE)
      if((dim(gradtxt)[1] != 1) & (dim(gradtxt)[2] !=1)){stop("Error in reading gradient text file!")}
      gradient[i] = eval(parse(text = gradtxt[1,1]))
    }

  }else if(nlinks ==3){
    theta1 = x[3]
    theta2 = x[4]
    theta3 = x[5]

    L1 = links[[1]]$length
    L2 = links[[2]]$length

    x211 = links[[2]]$x[1,1]
    x212 = links[[2]]$x[2,1]
    x221 = links[[2]]$x[1,2]
    x222 = links[[2]]$x[2,2]
    x231 = links[[2]]$x[1,3]
    x232 = links[[2]]$x[2,3]
    x241 = links[[2]]$x[1,4]
    x242 = links[[2]]$x[2,4]

    y211 = y[[2]][1,1]
    y212 = y[[2]][2,1]
    y221 = y[[2]][1,2]
    y222 = y[[2]][2,2]
    y231 = y[[2]][1,3]
    y232 = y[[2]][2,3]
    y241 = y[[2]][1,4]
    y242 = y[[2]][2,4]

    x311 = links[[3]]$x[1,1]
    x312 = links[[3]]$x[2,1]
    x321 = links[[3]]$x[1,2]
    x322 = links[[3]]$x[2,2]
    x331 = links[[3]]$x[1,3]
    x332 = links[[3]]$x[2,3]
    x341 = links[[3]]$x[1,4]
    x342 = links[[3]]$x[2,4]

    y311 = y[[3]][1,1]
    y312 = y[[3]][2,1]
    y321 = y[[3]][1,2]
    y322 = y[[3]][2,2]
    y331 = y[[3]][1,3]
    y332 = y[[3]][2,3]
    y341 = y[[3]][1,4]
    y342 = y[[3]][2,4]
    gradient = rep(NA, 5)
    for(i in 1:length(x)){
      gradtxtfile = sprintf("./AnalyticalCostGrad/TripleLink/CostGradientChar%.0f.txt", i)
      gradtxt = read.table(file = gradtxtfile, sep=',' , stringsAsFactors = FALSE)
      if((dim(gradtxt)[1] != 1) & (dim(gradtxt)[2] !=1)){stop("Error in reading gradient text file!")}
      gradient[i] = eval(parse(text = gradtxt[1,1]))
    }
  }
  return(gradient)
}

hess_analytical = function(x, y, links){
  # Generates the analytical hessian of the cost function.  Originally generated
  # via MATLAB and the symbolic toolbox.
  # x is parameter vector (r1, r2, theta1, theta2, theta3)
  # y[[nlinks]] of observations
  # links contains the known LCS information of each link

  nlinks = length(links)
  r1 = x[1]
  r2 = x[2]
  lsigma = x[length(x)]

  x111 = links[[1]]$x[1,1]
  x112 = links[[1]]$x[2,1]
  x121 = links[[1]]$x[1,2]
  x122 = links[[1]]$x[2,2]
  x131 = links[[1]]$x[1,3]
  x132 = links[[1]]$x[2,3]
  x141 = links[[1]]$x[1,4]
  x142 = links[[1]]$x[2,4]

  y111 = y[[1]][1,1]
  y112 = y[[1]][2,1]
  y121 = y[[1]][1,2]
  y122 = y[[1]][2,2]
  y131 = y[[1]][1,3]
  y132 = y[[1]][2,3]
  y141 = y[[1]][1,4]
  y142 = y[[1]][2,4]

  if(nlinks ==1){
    theta1 = x[3]
    hess = matrix(NA, nrow = length(x), ncol = length(x))
    for(i in 1:length(x)){
      for(j in 1:length(x)){
        hesstxtfile = sprintf("./AnalyticalHessian/SingleLink/hesschar%.0f_%.0f.txt", i,j)
        hesstxt = read.table(file = hesstxtfile, sep=',' , stringsAsFactors = FALSE)
        if((dim(hesstxt)[1] != 1) & (dim(hesstxt)[2] !=1)){stop("Error in reading hessian text file!")}
        hess[i,j] = eval(parse(text = hesstxt[1,1]))
      }
    }

  }else if(nlinks ==2){
    theta1 = x[3]
    theta2 = x[4]
    L1 = links[[1]]$length

    x211 = links[[2]]$x[1,1]
    x212 = links[[2]]$x[2,1]
    x221 = links[[2]]$x[1,2]
    x222 = links[[2]]$x[2,2]
    x231 = links[[2]]$x[1,3]
    x232 = links[[2]]$x[2,3]
    x241 = links[[2]]$x[1,4]
    x242 = links[[2]]$x[2,4]

    y211 = y[[2]][1,1]
    y212 = y[[2]][2,1]
    y221 = y[[2]][1,2]
    y222 = y[[2]][2,2]
    y231 = y[[2]][1,3]
    y232 = y[[2]][2,3]
    y241 = y[[2]][1,4]
    y242 = y[[2]][2,4]

    hess = matrix(NA, nrow = length(x), ncol = length(x))
    for(i in 1:length(x)){
      for(j in 1:length(x)){
        hesstxtfile = sprintf("./AnalyticalHessian/DoubleLink/hesschar%.0f_%.0f.txt", i,j)
        hesstxt = read.table(file = hesstxtfile, sep=',' , stringsAsFactors = FALSE)
        if((dim(hesstxt)[1] != 1) & (dim(hesstxt)[2] !=1)){stop("Error in reading hessian text file!")}
        hess[i,j] = eval(parse(text = hesstxt[1,1]))
      }
    }
  }else if(nlinks ==3){
    theta1 = x[3]
    theta2 = x[4]
    theta3 = x[5]

    L1 = links[[1]]$length
    L2 = links[[2]]$length

    x211 = links[[2]]$x[1,1]
    x212 = links[[2]]$x[2,1]
    x221 = links[[2]]$x[1,2]
    x222 = links[[2]]$x[2,2]
    x231 = links[[2]]$x[1,3]
    x232 = links[[2]]$x[2,3]
    x241 = links[[2]]$x[1,4]
    x242 = links[[2]]$x[2,4]

    y211 = y[[2]][1,1]
    y212 = y[[2]][2,1]
    y221 = y[[2]][1,2]
    y222 = y[[2]][2,2]
    y231 = y[[2]][1,3]
    y232 = y[[2]][2,3]
    y241 = y[[2]][1,4]
    y242 = y[[2]][2,4]

    x311 = links[[3]]$x[1,1]
    x312 = links[[3]]$x[2,1]
    x321 = links[[3]]$x[1,2]
    x322 = links[[3]]$x[2,2]
    x331 = links[[3]]$x[1,3]
    x332 = links[[3]]$x[2,3]
    x341 = links[[3]]$x[1,4]
    x342 = links[[3]]$x[2,4]

    y311 = y[[3]][1,1]
    y312 = y[[3]][2,1]
    y321 = y[[3]][1,2]
    y322 = y[[3]][2,2]
    y331 = y[[3]][1,3]
    y332 = y[[3]][2,3]
    y341 = y[[3]][1,4]
    y342 = y[[3]][2,4]

    hess = matrix(NA, nrow = length(x), ncol = length(x))
    for(i in 1:length(x)){
      for(j in 1:length(x)){
        hesstxtfile = sprintf("./AnalyticalHessian/TripleLink/hesschar%.0f_%.0f.txt", i,j)
        hesstxt = read.table(file = hesstxtfile, sep=',' , stringsAsFactors = FALSE)
        if((dim(hesstxt)[1] != 1) & (dim(hesstxt)[2] !=1)){stop("Error in reading hessian text file!")}
        hess[i,j] = eval(parse(text = hesstxt[1,1]))
      }
    }
  }
  return(hess)
}

LS_soln = function(y, links,
                   init_type = 'random',
                   inits = NA){
  # Computes the least squares solution given observations y.
  # y = observaitons
  # links = known LCS geometry
  # init_type = type of initial values to specify (random = randomly generated
  #             true_vals = true values of the simulation or speicified = custom
  #             specification of initial values.)
  # inits = specified initis if true_vals or specified is used for init_type
  nlinks = length(links)

  if(init_type =='random'){
    inits.r = c(runif(1, -0.1, 0.1),
                runif(1, -0.1, 0.1))
    inits.theta = runif(nlinks, -180, 180)
    print(paste('Computing LS solution with initial Values set to the RANDOM VALUES:', paste(c(inits.r, inits.theta), collapse = ', ' )))
  }else if(init_type == 'true_vals'){
    if (is.na(inits)){stop("If init_type is not random then inits must be specified")}
    print('Computing LS solution with initial Values set to the TRUE VALUES')
    inits.r = inits[1:2]
    inits.theta = inits[3:(nlinks +2)]
  }else if (init_type == "specified"){
    if (is.na(inits)){stop("If init_type is not random then inits must be specified")}
    print('Computing LS solution with initial Values set to the SPECIFIED VALUES')
    inits.r = inits[1:2]
    inits.theta = inits[3:(nlinks +2)]
  }
  inits = c(inits.r, inits.theta)

  # run optimisation using BFGS method.
  t1 = Sys.time()
  temp = optim(par=inits,fn = cost, gr = costgrad_analytical,
               y = y, links = links,
               method = "BFGS", control = list(trace=FALSE, maxit = 10000, reltol = 1e-20))
  t2 = Sys.time()
  return_list = list(r.hat = temp$par[1:2],
                     theta.hat = temp$par[3:(nlinks+2)],
                     value = temp$val, time = as.numeric(t2-t1))

  # generate residuals and use for corresponding estimate of sigma_hat
  nmarkers = 0
  obs.hat = gen_posture(links, return_list$r.hat, return_list$theta.hat)
  residual_sum = 0
  for(i in 1:nlinks){
    nmarkers = ncol(y[[i]])
    for(j in 1:nmarkers){
      residual_sum = residual_sum +
        (y[[i]][1, j]- obs.hat$alpha[[i]][1,j])^2 + # xresidual
        (y[[i]][2, j]- obs.hat$alpha[[i]][2, j])^2 # yresidual
      nmarkers = nmarkers +1
    }

  }
  p = length(return_list$r.hat) + length(return_list$theta.hat)
  return_list$sigma.hat = sqrt(residual_sum / (2*nmarkers))

  # Generate CI estimate via hessian matrix
  hess = hess_analytical(x = c(return_list$r.hat, return_list$theta.hat, log(return_list$sigma.hat)),
                         links = links,
                         y = y)
  fisher_info = solve(-hess)
  prop_sigma = diag(diag(sqrt(diag(fisher_info))))
  intervals = matrix(NA, nrow=nlinks+3, ncol =2)
  intervals[,1] = c(return_list$r.hat, return_list$theta.hat, return_list$sigma.hat) - 1.96*prop_sigma
  intervals[,2] = c(return_list$r.hat, return_list$theta.hat, return_list$sigma.hat) + 1.96*prop_sigma
  intervals[nrow(intervals), ] = exp(intervals[nrow(intervals),])
  return_list$intervals = intervals

  return(return_list)
}

########################### Section D  #########################
## Compute Bayesian solution.
################################################################

Bayes_soln = function(y, links, mdlfile,
                      true_vals,
                      init_type = 'random',
                      ls_est = NA,
                      ...){
  # Applies the Bayesian model to observations y and known data links.
  # y = observaitons
  # links = known LCS geometry
  # mdlfile = specified jags model file
  # true_vals = true values of the specified simulation.
  # init_type = type of initial values to specify (random = randomly generated
  #             true_vals = true values of the simulation or speicified = custom
  #             specification of initial values.)
  # inits = specified initis if true_vals is specified.

    nlinks = length(links)
    nmarkers = sapply(links, function(x){ncol(x$x)})
    M = 2 # dimension
    seg_lengths = sapply(links, function(x){x$length})
    x = lapply(links, function(x){x$x})

    # initilise x and y data for jags models
    xjags = yjags = array(NA, dim = c(nlinks, max(nmarkers), 2))
    for(i in 1:nlinks){
      xjags[i,,] = t(x[[i]])
      yjags[i,,] = t(y[[i]])
    }

    # input data for jags model
    jags_input = list(nlinks = nlinks,
                      nmarkers = nmarkers,
                      ndim = 2,
                      x = xjags,
                      y = yjags,
                      leng = cbind(seg_lengths, rep(0,nlinks)))

    # Specify initial values
    if(grepl('p4|p5', mdlfile)){
        # LS centered priors require the ls solution
      jags_input$ls_r = ls_est[1:2]
      jags_input$ls_theta = ls_est[3:(2+nlinks)]
      jags_input$ls_sigma = ls_est[length(ls_est)]
    }

    if(init_type =='true_vals'){
      r_true = true_vals[1:2]
      theta_true = true_vals[3:(2+nlinks)]
      sigma_true = true_vals[length(ls_est)]
      if(grepl('p3', mdlfile)){
        # if weakly informative prior specify inits in terms of hip, knee and ankle angles.
        print('Computing Bayes P3 solution with initial Values set to the TRUE VALUES')
        inits = list(hip = theta_true[1],
                     knee = theta_true[2] - theta_true[1],
                     ankle = -90 -theta_true[2] + theta_true[3],
                     r = r_true,
                     sigma_mm = sigma_true*1000)
      }else if(grepl('p2', mdlfile)){
        print('Computing Bayes P2 solution with initial Values set to the TRUE VALUES')
        inits = list(thetar = theta_true * pi/180,
                     r = r_true,
                     tau = 1/sigma_true/sigma_true)
        jags_input$true_vals = true_vals
      }else{
        print('Computing Bayes solution with initial Values set to the TRUE VALUES')
        inits = list(thetar = theta_true * pi/180,
                     r = r_true,
                     sigma = sigma_true)

      }
    }
    else if(init_type == 'random'){
      if(grepl('p3', mdlfile)){
          # if weakly informative prior specify inits in terms of hip, knee and ankle angles.
        print('Computing Bayes P3 solution with initial Values set to the RANDOM VALUES')
        r_random = rt_ls(2, df = 5, mu = 0, a = 1)
        hip_random = rt_ls(1, df = 5, mu=-45, a = 15)
        knee_random = rt_ls(1, df=5, mu=30, a=6)
        ankle_random = rt_ls(1, df=5, mu=0, a=7)
        sigma_mm_random = rgamma(1, shape = 1.2, scale=0.1)
        inits = list(hip = hip_random,
                     knee = knee_random,
                     ankle = ankle_random,
                     r = r_random,
                     sigma_mm = sigma_mm_random)
      }else if(grepl('p2', mdlfile)){
        print('Computing Bayes P2 solution with initial Values set to the RANDOM VALUES')
        r_random = c(rnorm(1, true_vals[1], 0.001), dnorm(1, true_vals[2], 0.001))
        theta_random = runif(nlinks, -180,180)
        tau_random = rnorm(1, 1/(true_vals[nlinks+3]*true_vals[nlinks+3]), 1)
        inits = list(thetar = theta_random * pi/180,
                     r = r_random,
                     tau = tau_random)
        jags_input$true_vals = true_vals
      }else if (grepl('p4', mdlfile)){
        print('Computing Bayes P4 solution with initial Values set to the RANDOM VALUES')
        r_random = c(runif(1, jags_input$ls_r[1]-0.1,jags_input$ls_r[1] + 0.1), runif(1, jags_input$ls_r[2]-0.1, jags_input$ls_r[2] + 0.1))
        theta_random = rep(0, nlinks)
        for(nl in 1:nlinks){
            theta_random[nl] = runif(1, jags_input$ls_theta[nl]-180, jags_input$ls_theta[nl]+180)
        }
        inits = list(theta = theta_random,
                     r = r_random,
                     sigma = runif(1, 0.1*jags_input$ls_sigma, 10*jags_input$ls_sigma))
      }else if (grepl('p5', mdlfile)){
        print('Computing Bayes P5 solution with initial Values set to the RANDOM VALUES')
        r_random = c(rnorm(1, jags_input$ls_r[1], 0.005), rnorm(1, jags_input$ls_r[2], 0.005))
        theta_random = rep(0, nlinks)
        for(nl in 1:nlinks){
          theta_random[nl] = rnorm(1, jags_input$ls_theta[nl], 1)
        }
        inits = list(theta = theta_random,
                     r = r_random,
                     sigma = abs(rnorm(1, jags_input$ls_sigma, 0.002)))
      }else{
        print('Computing Bayes P1 solution with initial Values set to the RANDOM VALUES')
        theta_random = runif(nlinks, -180,180)
        r_random = runif(2, -0.1, 0.1)
        inits = list(thetar = theta_random * pi/180,
                     r = r_random,
                     sigma = runif(1, 0.1*true_vals[nlinks+3], 10*true_vals[nlinks+3]))
      }
    }else if(init_type == 'ls_est'){
      if(sum(is.na(ls_est))>0){stop("If using 'ls_est' as initials please supply LS estimates")}
      if(grepl('p3', mdlfile)){
          # if wealky informative prior specify inits in terms of hip, knee and ankle angles.
        print('Computing Bayes P3 solution with initial Values set to the LS VALUES')
        hip_random = rt_ls(1, df = 5, mu=-45, a = 15)
        knee_random = rt_ls(1, df=5, mu=30, a=6)
        ankle_random = rt_ls(1, df=5, mu=0, a=7)

        inits = list(hip = hip_random,
                     knee = knee_random,
                     ankle = ankle_random,
                     r = c(ls_est[1], ls_est[2]),
                     sigma_mm = tail(ls_est,1)*1000)
      }else if(grepl('p2', mdlfile)){
        print('Computing Bayes P2 solution with initial Values set to the LS VALUES')
        jags_input$true_vals = true_vals
        inits = list(thetar = c(ls_est[3:(nlinks+2)]) * pi/180,
                     r = c(ls_est[1], ls_est[2]),
                     tau = 1/tail(ls_est,1)/tail(ls_est,1))
      }else if(grepl('p4|p5', mdlfile)){
          print('Computing Bayes P4 or P5 solution with initial Values set to the LS VALUES')
        inits = list(theta = c(ls_est[3:(nlinks+2)]),
                     r = c(ls_est[1], ls_est[2]),
                     sigma = tail(ls_est,1))
      }else if (grepl('p1_2', mdlfile)){
        print('Computing Bayes P1_2  solution with initial Values set to the LS VALUES')
        jags_input$true_vals = true_vals
        inits = list(thetar = c(ls_est[3:(nlinks+2)]) * pi/180,
                     r = c(ls_est[1], ls_est[2]),
                     tau = 1/tail(ls_est,1)/tail(ls_est,1))
      }
      else{
        print('Computing Bayes P1  solution with initial Values set to the LS VALUES')
        inits = list(thetar = c(ls_est[3:(nlinks+2)]) * pi/180,
                     r = c(ls_est[1], ls_est[2]),
                     sigma = tail(ls_est,1))
      }
    }

  # Compile jags model
  t1 = Sys.time()
  m1 = jags.model(file = mdlfile,
                  data = jags_input,
                  inits = inits,
                  n.chains = 4,
                  n.adapt = 10000
  )

  # Sample from model
  out1 = coda.samples(model = m1,
                      variable.names = c("theta", "r", "sigma"),
                      n.iter = 50000,
                      thin = 5)
  t2 = Sys.time()

  return(list(fit = out1, time = as.numeric(t2-t1)))
}


# data_in is a comma separted list of data points
# groups_in is a comma separted list of group indeces for the points in data_in

MGRAST_do_stats <<- function (data_file,
                              groups_file,
                              data_type = c("raw", "normalized"),
                              sig_test = c(
                                "t-test-paired", "Wilcoxon-paired",
                                "t-test-un-paired", "Mann-Whitney_un-paired-Wilcoxon",
                                "ANOVA-one-way", "Kruskal-Wallis"
                                ),
                              file_out = "my_stats")

{

  ### SUB TO WRITE OUTPUT
  write_log <<- function (sig_test, data_type){
    write(paste("test is     : ", sig_test),    file = "do_stats.log")
    write(paste("data-type   : ", data_type),   file = "do_stats.log", append = TRUE)
    write(paste("num-samples : ", num_samples), file = "do_stats.log", append = TRUE)
    write(paste("num-groups  : ", num_groups),  file = "do_stats.log", append = TRUE)
  }

  ### SUB TO PREP DATA FOR TWO GROUP ANALYSES
  prep_two_groups <<- function(row_data){
    group_1 = levels(row_data[,2])[1]
    group_2 = levels(row_data[,2])[2]
    group_1_size = 0
    group_2_size = 0
    for (i in 1:as.matrix(dim(row_data)[1])){
      if(identical(as.character(row_data[i,2]), as.character(group_1))){
        group_1_size = group_1_size + 1
      }
      if(identical(as.character(row_data[i,2]), as.character(group_2))){
        group_2_size = group_2_size + 1
      }
    }
    group_1_data <<- matrix(,group_1_size,1)
    group_2_data <<- matrix(,group_2_size,1)
    group_1_index = 0
    group_2_index = 0
    for (i in 1:as.matrix(dim(row_data)[1])){
      if (identical(as.character(row_data[i,2]), as.character(group_1))){
        group_1_index = group_1_index + 1
        group_1_data[group_1_index,1] <<- row_data[i,1]
      }
      if (identical(as.character(row_data[i,2]), as.character(group_2))){
        group_2_index = group_2_index + 1
        group_2_data[group_2_index,1] <<- row_data[i,1]
      }
    }
  }

  ### ### MAIN SCRIPT STARTS HERE ### ###

  ### TEST TO MAKE SURE ALL OF THE ARGUMENTS ARE SUPPLIED
  #if()

  # LOAD NECESSARY PACKAGES
  library(stats)
  library(nlme)

  ### IMPORT AND FORMAT THE DATA ###
  all_data = data.frame(read.table(data_file, row.names=1, header=TRUE, sep="\t", comment.char="", quote=""))
  groups = data.matrix(scan(file = groups_file, what = "character", sep = "\t", quiet=TRUE))
  num_samples = dim(all_data)[2]
  num_rows = dim(all_data)[1]
  num_groups <<- nlevels(as.factor(groups[,1]))    # <----- Rats
  group_names <<- matrix(levels(as.factor(groups[,1])))

  # CREATE OUTPUT TABLE AND GIVE IT APPROPRIATE HEADERS
  output_table <<- data.frame(matrix(0, num_rows, num_groups+2))
  row.names(output_table) <<- row.names(all_data)
  names(output_table) <<- c(paste("group_(", levels(as.factor(groups[,1])),")_stddev",sep =""),
                            paste(sig_test, "_stat",sep=""),
                            paste(sig_test, "_p_value",sep="")
                            )

  ### PERFORM THE ANALYSES
  for (my_row in 1:num_rows){

    numeric_row_data <<- as.numeric(all_data[my_row,])
    factor_row_groups <<- as.character(groups[,1]) # is ok here
    row_data <<- cbind(numeric_row_data, factor_row_groups)
    row_data <<- data.frame(row_data)
    row_data[,1] <<- numeric_row_data # some values get changed if you don;t reload them here
    row_data[,2] <<- as.factor(row_data[,2])
    names(row_data)<<-c("values","ind")

    for (group in 1:num_groups){

      group_ind <<- levels(row_data[,2])[group]

      ### CALCULATE THE STANDARD DEVIATION FOR EACH GROUP
      num_samples_in_group <<- 0
      for(sample in 1:num_samples){
        if(identical(((matrix(row_data[sample,2]))[1,1]), group_ind)){
          if (num_samples_in_group==0){
            group_sample_counts <<- row_data[sample,1]
            num_samples_in_group <<- num_samples_in_group + 1
          }else{
            group_sample_counts <<- c(group_sample_counts, row_data[sample,1])
            num_samples_in_group <<- num_samples_in_group + 1
          }
        }
      }
      output_table[my_row, group] <<- as.real(sd(group_sample_counts))
    }

    if(identical(sig_test, "t-test-un-paired")){   # <-- new
      prep_two_groups(row_data)
      ttest_unpaired_output = t.test(group_1_data, group_2_data)
      t_unpaired_stat_value = as.real(ttest_unpaired_output["statistic"])
      t_unpaired_p_value = as.real(ttest_unpaired_output["p.value"])
      output_table[my_row, num_groups+1] <<- t_unpaired_stat_value
      output_table[my_row, num_groups+2] <<- t_unpaired_p_value
    }

    else if(identical(sig_test, "t-test-paired")){   # <-- new
      prep_two_groups(row_data)
      ttest_paired_output = t.test(group_1_data, group_2_data, paired = TRUE)
      t_paired_stat_value = as.real(ttest_paired_output["statistic"])
      t_paired_p_value = as.real(ttest_paired_output["p.value"])

      output_table[my_row, num_groups+1] <<- t_paired_stat_value
      output_table[my_row, num_groups+2] <<- t_paired_p_value
    }

    else if (identical(sig_test, "ANOVA-one-way")){
      anova_output = anova(aov(values~ind, data=row_data))
      anova_F = anova_output["F value"]
      F_value = as.real(anova_F[1,1])
      anova_p = anova_output["Pr(>F)"]
      anova_p_value = as.real(anova_p[1,1])
      output_table[my_row, num_groups+1] <<- F_value
      output_table[my_row, num_groups+2] <<- anova_p_value
    }

    else if (identical(sig_test, "Mann-Whitney_un-paired-Wilcoxon")){
      prep_two_groups(row_data)
      MWhitney_output = wilcox.test(group_1_data, group_2_data, exact=TRUE) # x -> s in text
      MWhitney_stat_value = as.real(MWhitney_output["statistic"])
      MWhitney_p_value = as.real(MWhitney_output["p.value"])
      output_table[my_row, num_groups+1] <<- MWhitney_stat_value
      output_table[my_row, num_groups+2] <<- MWhitney_p_value
    }

    else if (identical(sig_test, "Wilcoxon-paired")){   # <-- new
      prep_two_groups(row_data)
      wilcox_output = wilcox.test(group_1_data, group_2_data, exact=TRUE, paired=TRUE) # x -> s in text
      wilcox_stat_value = as.real(wilcox_output["statistic"])
      wilcox_p_value = as.real(wilcox_output["p.value"])
      output_table[my_row, num_groups+1] <<- wilcox_stat_value
      output_table[my_row, num_groups+2] <<- wilcox_p_value
    }

    else if (identical(sig_test, "Kruskal-Wallis")){
      kruskal_output = kruskal.test(row_data[,1], row_data[,2])
      kruskal_K_value = as.real(kruskal_output["statistic"])
      kruskal_p_value = as.real(kruskal_output["p.value"])
      output_table[my_row, num_groups+1] <<- kruskal_K_value
      output_table[my_row, num_groups+2] <<- kruskal_p_value
    }

    else{
      stop("no significance test was chosen")
    }

    write.table(output_table, file = file_out, sep = "\t", col.names=NA, row.names = TRUE)

  }

}


#' @title Regression analysis of prevalence
#' @param table_file (optional) file to save parameter table to
#' @return a list of two plots (effect size and posterior predicitons) and a
#' table of effect size estimates
#' @importFrom dplyr select left_join mutate filter group_by ungroup inner_join bind_rows recode_factor summarise arrange if_else
#' @importFrom tidyr pivot_longer pivot_wider separate
#' @importFrom stats as.formula median quantile
#' @importFrom brms brm negbinomial
#' @importFrom tidybayes spread_draws median_qi
#' @importFrom forcats fct_rev
#' @importFrom ggplot2 ggplot aes xlab ylab xlim geom_hline coord_flip theme_minimal geom_point geom_ribbon theme scale_color_brewer scale_y_log10
#' @importFrom ggdist geom_pointinterval
#' @importFrom rstantools posterior_predict
#' @importFrom tidyselect starts_with matches
#' @importFrom tibble as_tibble
#' @importFrom kableExtra kable save_kable
#' @importFrom tools file_ext
#' @export
regression <- function(table_file = NULL) {
  unemp <- covariates$unemp
  age <- covariates$age
  pop_dens <- covariates$pop_dens
  roma <- covariates$roma

  ## all covariates and outcome
  prev <- ms.tst %>%
    dplyr::select(county, region, pop, dplyr::starts_with("attendance_"),
                  dplyr::starts_with("positive_"), pilot) %>%
    dplyr::left_join(unemp, by = "county") %>%
    dplyr::left_join(age, by = "county") %>%
    dplyr::left_join(pop_dens, by = "county") %>%
    dplyr::mutate(xcounty = sub(" [IV]+$",  "", county)) %>%
    dplyr::left_join(roma %>% select(xcounty = county, proportion_roma),
              by = "xcounty") %>%
    simplify_names() %>%
    dplyr::left_join(Rt.county %>%
                     rename(simple_name = county), by = "simple_name") %>%
    dplyr::select(-xcounty, -simple_name) %>%
    dplyr::mutate(unemp_rate = unemployed / active) %>%
    dplyr::select(county, region, pop, tidyselect::matches("^attendance_[123]"),
                  tidyselect::matches("^positive_[123]"), mean_age,
                  pop_dens, unemp_rate, proportion_roma, R)

  ## extract pilot variables
  pilot <- prev %>%
    dplyr::filter(!is.na(attendance_1)) %>%
    dplyr::mutate(round_attendance_prec = attendance_1 / pop,
                  round_prev_prec = positive_1 / attendance_1,
                  round_R_prec = R) %>%
    dplyr::select(county, dplyr::starts_with("round_")) %>%
    tidyr::pivot_longer(-county) %>%
    dplyr::group_by(name) %>%
    dplyr::mutate(value = (value - mean(value)) / sd(value)) %>%
    dplyr::ungroup() %>%
    tidyr::pivot_wider() %>%
    dplyr::mutate(`0` = 0, `1` = 1, `2` = 1) %>%
    tidyr::pivot_longer(tidyselect::matches("^[012]"),
                        names_to = "round", values_to = "multiplier") %>%
    dplyr::mutate(round = as.integer(round)) %>%
    tidyr::pivot_longer(c(-county, -round, -multiplier)) %>%
    dplyr::mutate(value = value * multiplier) %>%
    tidyr::pivot_wider() %>%
    dplyr::select(-multiplier)

  ## extract round 1 variables
  round1 <- prev %>%
    dplyr::filter(is.na(attendance_1)) %>%
    dplyr::mutate(round_attendance_prec = attendance_2 / pop,
                  round_prev_prec = positive_2 / attendance_2,
                  round_R_prec = R) %>%
    dplyr::select(county, tidyselect::starts_with("round_")) %>%
    tidyr::pivot_longer(-county) %>%
    dplyr::group_by(name) %>%
    dplyr::mutate(value = (value - mean(value)) / sd(value)) %>%
    dplyr::ungroup() %>%
    pivot_wider() %>%
    dplyr::mutate(`0` = 0, `1` = 1) %>%
    tidyr::pivot_longer(tidyselect::matches("^[01]"),
                        names_to = "round", values_to = "multiplier") %>%
    dplyr::mutate(round = as.integer(round)) %>%
    tidyr::pivot_longer(c(-county, -round, -multiplier)) %>%
    dplyr::mutate(value = value * multiplier) %>%
    tidyr::pivot_wider() %>%
    dplyr::select(-multiplier)

  prev_long <- prev %>%
    tidyr::pivot_longer(tidyselect::matches("^(positive|attendance)_[123]$")) %>%
    tidyr::separate(name, c("name", "round"), sep = "_") %>%
    dplyr::select(-R) %>%
    dplyr::filter(!is.na(value)) %>%
    tidyr::pivot_wider() %>%
    tidyr::pivot_longer(c(-county, -region, -positive, -attendance, -round)) %>%
    dplyr::group_by(name) %>%
    dplyr::mutate(value = (value - mean(value, na.rm = TRUE)) /
                    sd(value, na.rm = TRUE)) %>%
    dplyr::ungroup() %>%
    tidyr::pivot_wider() %>%
    dplyr::group_by(county) %>%
    dplyr::mutate(round = as.integer(round),
                  round = dplyr::if_else(rep(any(round == 1), n()),
                                         round - 1L, round - 2L)) %>%
    dplyr::ungroup()

  pilot_long <- prev_long %>%
    dplyr::inner_join(pilot, by = c("county", "round"))
  round1_long <- prev_long %>%
    dplyr::inner_join(round1, by = c("county", "round"))

  prev_long <- pilot_long %>%
    dplyr::bind_rows(round1_long)

  model <-
    stats::as.formula(positive ~
                        1 + (1 | county) + round + mean_age + pop_dens +
                        unemp_rate + proportion_roma + round_attendance_prec +
                        round_prev_prec + round_R_prec +
                        offset(log(attendance)))

  ## Fit model --------------------------------------------------------------
  fit <- brms::brm(model, data = prev_long, family = brms::negbinomial(),
                   iter = 3000)

  effects <- fit %>%
    tidybayes::spread_draws(`^b_.*`, regex = TRUE) %>%
    tidybayes::median_qi(.width = c(.95, .50)) %>%
    tidyr::pivot_longer(dplyr::starts_with("b_")) %>%
    dplyr::mutate(name =
                    dplyr::if_else(grepl("\\.", name), name,
                                   paste(name, .point, sep = "."))) %>%
    tidyr::separate(name, c("variable", "name"), sep = "\\.") %>%
    tidyr::pivot_wider() %>%
    dplyr::mutate(affects = dplyr::if_else(variable %in% c("b_Intercept",
                                                           "b_mean_age",
                                                           "b_pop_dens",
                                                           "b_unemp_rate",
                                                           "b_proportion_roma"),
                                           "Prevalence", "Reduction"),
                  variable =
                    dplyr::recode_factor(
                             variable,
                             b_Intercept = "Intercept",
                             b_mean_age = "Mean age",
                             b_pop_dens =  "Population density",
                             b_unemp_rate = "Unemployment rate",
                             b_proportion_roma = "Size of Marginalised Roma Community",
                             b_round = "Round",
                             b_round_attendance_prec = "Previous attendance",
                             b_round_prev_prec = "Previous prevalence",
                             b_round_R_prec = "Reproduction number"
                           ),
                  variable = forcats::fct_rev(variable))

  ep <-
    ggplot2::ggplot(effects %>%
                    dplyr::filter(variable != "Intercept") %>%
                    dplyr::mutate_at(vars(median, lower, upper), exp),
                    ggplot2::aes(x = variable, y = median,
                                 ymin = lower, ymax = upper,
                                 colour = affects)) +
    ggdist::geom_pointinterval() +
    ggplot2::xlab("") + ggplot2::ylab("Posterior ratio") +
    ggplot2::geom_hline(yintercept = 1, linetype = "dashed") +
    ggplot2::ylim(c(0.25, 1.75)) +
    ggplot2::coord_flip() +
    ggplot2::theme_minimal() +
    ggplot2::scale_color_brewer("", palette = "Set1") +
    ggplot2::theme(legend.position = "none")

  ## Posterior predictive plot -----------------------------------------------
  yhat <- rstantools::posterior_predict(fit)
  colnames(yhat) <- paste(prev_long$county, prev_long$round, sep = "_")

  plotpp <- tibble::as_tibble(yhat) %>%
    tidyr::pivot_longer(everything(), names_to = "county_round") %>%
    tidyr::separate(county_round, c("county", "round"), sep = "_") %>%
    dplyr::mutate(round = as.integer(round)) %>%
    dplyr::group_by(county, round) %>%
    dplyr::summarise(median = stats::median(value),
                     lower = stats::quantile(value, 0.25),
                     upper = stats::quantile(value, 0.75),
                     lowest = stats::quantile(value, 0.05),
                     uppest = stats::quantile(value, 0.95),
                     .groups = "drop") %>%
    dplyr::ungroup() %>%
    dplyr::left_join(prev_long, by = c("county", "round")) %>%
    dplyr::arrange(median) %>%
    dplyr::mutate(id = seq_len(n())) %>%
    dplyr::group_by(county) %>%
    dplyr::mutate(round = dplyr::if_else(rep(any(round == 2), n()),
                                         round + 1, round + 2)) %>%
    dplyr::ungroup() %>%
    dplyr::mutate(round = dplyr::if_else(round == 1, "Pilot",
                                         paste("Round", round - 1)))

  pp <- ggplot2::ggplot(plotpp, aes(x = id, y = median, color = round)) +
    ggplot2::geom_point(aes(y = positive), size = 1) +
    ggplot2::geom_ribbon(aes(ymin = lower, ymax = upper), alpha = 0.3,
                         color = NA) +
    ggplot2::geom_ribbon(aes(ymin = lowest, ymax = uppest), alpha = 0.15,
                         color = NA) +
    ggplot2::xlab("County") +
    ggplot2::ylab("Estimated # positive") +
    ggplot2::theme_minimal() +
    ggplot2::theme(axis.text.x = element_blank(),
                   axis.ticks.x = element_blank()) +
    ggplot2::scale_color_brewer("", palette = "Dark2") +
    ggplot2::scale_y_log10()

  if (!is.null(table_file)) {
    align <- c("l|", "r", "r", "r")
    if (file_ext(table_file) == "pdf") {
      format <- "latex"
    } else if (file_ext(table_file) == "html") {
      format <- "html"
      align <- sub("\\|", "", align)
    }
    col.names <- c("Variable", "Lower 95%", "Median", "Upper 95%")
    k <- kableExtra::kable(effects %>%
                           dplyr::filter(.width == 0.95) %>%
                           dplyr::select(variable, lower, median, upper) %>%
                           dplyr::mutate_if(is.numeric, exp) %>%
                           dplyr::mutate_if(is.numeric, signif, digits = 2) %>%
                           dplyr::mutate_if(is.numeric, as.character),
                           format = format,
                           col.names = col.names,
                           align = align,
                           booktabs = FALSE)

    kableExtra::save_kable(k, table_file)
  }

  return(list(plots = list(effects = ep,
                           predictions = pp),
              table = effects,
              covariates = prev))
}


  shannon.diversity = function(x, base=2, getid=TRUE) {
    rs = rowSums(x, na.rm=T)
    good = which(rs>0)
    if (length(good)==0) {
      if (getid) {
        out = rep(NA,4)
      } else {
        out = rep(NA,3)
      }
    }

    if (length(good)>0) {
      rs = rs[good]
      x = x[ good , ]
      id = rownames( x )
      pr = x / rs

      if (is.array(x) ) {
        shannon = - rowSums( pr * log( pr, base=base ), na.rm=T )
      } else {
        shannon = - sum( pr * log( pr, base=base ), na.rm=T )
      }

      Hmax = - log( rs, base=base )
      evenness = shannon / Hmax

      if (getid) {
        out = data.frame( id=id, shannon=shannon, evenness=evenness, Hmax=Hmax )
      } else {
        out = data.frame( shannon=shannon, evenness=evenness, Hmax=Hmax )
      }
    }

    return(out)
  }


#!/usr/bin/env Rscript

args = commandArgs(trailingOnly = TRUE)

if (length(args) != 1) {
    stop("Provide Motor Type (cim, 775) as Command Line Argument.")
}

if (args[1] == "cim") {
    rpm = 2670
} else if (args[1] == "775") {
    rpm = 9370
} else {
    rpm = 1
}

peakFile = paste(args[1], "-peak-power.csv", sep="")
stallFile = paste(args[1], "-stall-6v.csv", sep="")

peak <- read.csv(peakFile, header = TRUE)
stall <- read.csv(stallFile, header = TRUE)

x11()

peak$Torque = (9.5488 * peak$Power) / rpm

plot(peak$Time,
     peak$Torque,
     col="red",
     type="l",
     ylim=c(0, 1.3),
     xlim=c(0, 300),
     xlab="Time (s)",
     ylab="Torque (Nm)",
     main="Torque Output - Slipping vs Stalling")
points(stall$Time,
       stall$Torque,
       col="blue",
       type="l")
legend(x = "bottomright",
       legend = c("Slipping (Peak Power)", "Stalling (6 Volts)"),
       col = c("red", "blue"),
       lty = 1,
       bty = "n",
       cex = 0.8)

dev.print(png, "torque.png", width = 1000, height = 1000)

x11()

plot(peak$Time,
     peak$Current,
     col="red",
     type="l",
     ylim=c(0,70),
     xlim=c(0,300),
     ylab="Current (A)",
     xlab="Time(s)",
     main="Current Draw, Slipping vs Stalling")
points(stall$Time,
       stall$Amps,
       col="blue",
       type="l")

dev.print(png, "current.png", width = 1000, height = 1000)

x11()

peak$TorquePerAmp = 1000 * peak$Torque / peak$Current
stall$TorquePerAmp = 1000 * stall$Torque / stall$Amps

plot(peak$Time,
     peak$TorquePerAmp,
     col="red",
     type="l",
     ylim=c(0,30),
     ylab="Torque per Amp (mNm/A)",
     xlab="Time (s)",
     main="Torque per Amp, Slipping vs Stalling")
points(stall$Time,
       stall$TorquePerAmp,
       col="blue",
       type="l")

dev.print(png, "torque-per-amp.png", width = 1000, height = 1000)

Sys.sleep(1000)

# Do both draw the same current? We could plot this, or also Torque per Amp.


# -----------------------------------------------------------------------------
# logs
# -----------------------------------------------------------------------------

#' Logarithms and Exponentials
#'
#' exp/log functions.
#'
#' @param x
#' A float vector/matrix.
#' @param base
#' The logarithm base.
#'
#' @return
#' A float vector/matrix of the same dimensions as the input.
#'
#' @examples
#' \dontrun{
#' library(float)
#'
#' x = flrunif(10)
#' log(x)
#' }
#'
#' @useDynLib float R_exp_spm R_expm1_spm R_log_spm R_log10_spm R_log2_spm
#' @name log
#' @rdname log
NULL

exp_float32 = function(x)
{
  ret = .Call(R_exp_spm, DATA(x))
  float32(ret)
}

expm1_float32 = function(x)
{
  ret = .Call(R_expm1_spm, DATA(x))
  float32(ret)
}

log_float32 = function(x, base=exp(1))
{
  if (is.float(base))
    base = dbl(base[1])

  ret = .Call(R_log_spm, DATA(x), as.double(base))
  float32(ret)
}

log10_float32 = function(x)
{
  ret = .Call(R_log10_spm, DATA(x))
  float32(ret)
}

log2_float32 = function(x)
{
  ret = .Call(R_log2_spm, DATA(x))
  float32(ret)
}

#' @rdname log
#' @export
setMethod("exp", signature(x="float32"), exp_float32)

#' @rdname log
#' @export
setMethod("expm1", signature(x="float32"), expm1_float32)

#' @rdname log
#' @export
setMethod("log", signature(x="float32"), log_float32)

#' @rdname log
#' @export
setMethod("log10", signature(x="float32"), log10_float32)

#' @rdname log
#' @export
setMethod("log2", signature(x="float32"), log2_float32)

# -----------------------------------------------------------------------------
# trig
# -----------------------------------------------------------------------------

#' Trigonometric functions
#'
#' Basic trig functions.
#'
#' @param x
#' A float vector/matrix.
#'
#' @return
#' A float vector/matrix of the same dimensions as the input.
#'
#' @examples
#' \dontrun{
#' library(float)
#'
#' x = flrunif(10)
#' sin(x)
#' }
#'
#' @useDynLib float R_sin_spm R_cos_spm R_tan_spm R_asin_spm R_acos_spm R_atan_spm
#' @name trig
#' @rdname trig
NULL

sin_float32 = function(x)
{
  ret = .Call(R_sin_spm, DATA(x))
  float32(ret)
}

cos_float32 = function(x)
{
  ret = .Call(R_cos_spm, DATA(x))
  float32(ret)
}

tan_float32 = function(x)
{
  ret = .Call(R_tan_spm, DATA(x))
  float32(ret)
}

asin_float32 = function(x)
{
  ret = .Call(R_asin_spm, DATA(x))
  float32(ret)
}

acos_float32 = function(x)
{
  ret = .Call(R_acos_spm, DATA(x))
  float32(ret)
}

atan_float32 = function(x)
{
  ret = .Call(R_atan_spm, DATA(x))
  float32(ret)
}

#' @rdname trig
#' @export
setMethod("sin", signature(x="float32"), sin_float32)

#' @rdname trig
#' @export
setMethod("cos", signature(x="float32"), cos_float32)

#' @rdname trig
#' @export
setMethod("tan", signature(x="float32"), tan_float32)

#' @rdname trig
#' @export
setMethod("asin", signature(x="float32"), asin_float32)

#' @rdname trig
#' @export
setMethod("acos", signature(x="float32"), acos_float32)

#' @rdname trig
#' @export
setMethod("atan", signature(x="float32"), atan_float32)

# -----------------------------------------------------------------------------
# hyperbolic
# -----------------------------------------------------------------------------

#' Hyperbolic functions
#'
#' Hyperbolic functions.
#'
#' @param x
#' A float vector/matrix.
#'
#' @return
#' A float vector/matrix of the same dimensions as the input.
#'
#' @examples
#' \dontrun{
#' library(float)
#'
#' x = flrunif(10)
#' sinh(x)
#' }
#'
#' @useDynLib float R_sinh_spm R_cosh_spm R_tanh_spm R_asinh_spm R_acosh_spm R_atanh_spm
#' @name hyperbolic
#' @rdname hyperbolic
NULL

sinh_float32 = function(x)
{
  ret = .Call(R_sinh_spm, DATA(x))
  float32(ret)
}

cosh_float32 = function(x)
{
  ret = .Call(R_cosh_spm, DATA(x))
  float32(ret)
}

tanh_float32 = function(x)
{
  ret = .Call(R_tanh_spm, DATA(x))
  float32(ret)
}

asinh_float32 = function(x)
{
  ret = .Call(R_asinh_spm, DATA(x))
  float32(ret)
}

acosh_float32 = function(x)
{
  ret = .Call(R_acosh_spm, DATA(x))
  float32(ret)
}

atanh_float32 = function(x)
{
  ret = .Call(R_atanh_spm, DATA(x))
  float32(ret)
}

#' @rdname hyperbolic
#' @export
setMethod("sinh", signature(x="float32"), sinh_float32)

#' @rdname hyperbolic
#' @export
setMethod("cosh", signature(x="float32"), cosh_float32)

#' @rdname hyperbolic
#' @export
setMethod("tanh", signature(x="float32"), tanh_float32)

#' @rdname hyperbolic
#' @export
setMethod("asinh", signature(x="float32"), asinh_float32)

#' @rdname hyperbolic
#' @export
setMethod("acosh", signature(x="float32"), acosh_float32)

#' @rdname hyperbolic
#' @export
setMethod("atanh", signature(x="float32"), atanh_float32)

# -----------------------------------------------------------------------------
# misc
# -----------------------------------------------------------------------------

#' Miscellaneous mathematical functions
#'
#' Miscellaneous mathematical functions.
#'
#' @param x
#' A float vector/matrix.
#'
#' @return
#' A float vector/matrix of the same dimensions as the input.
#'
#' @examples
#' \dontrun{
#' library(float)
#'
#' x = flrunif(10)
#' sqrt(x)
#' }
#'
#' @useDynLib float R_abs_spm R_sqrt_spm
#' @name miscmath
#' @rdname miscmath
NULL

abs_float32 = function(x)
{
  ret = .Call(R_abs_spm, DATA(x))
  float32(ret)
}

sqrt_float32 = function(x)
{
  ret = .Call(R_sqrt_spm, DATA(x))
  float32(ret)
}

#' @rdname miscmath
#' @export
setMethod("abs", signature(x="float32"), abs_float32)

#' @rdname miscmath
#' @export
setMethod("sqrt", signature(x="float32"), sqrt_float32)

# -----------------------------------------------------------------------------
# special
# -----------------------------------------------------------------------------

#' Special mathematical functions
#'
#' Special mathematical functions.
#'
#' @param x
#' A float vector/matrix.
#'
#' @return
#' A float vector/matrix of the same dimensions as the input.
#'
#' @examples
#' \dontrun{
#' library(float)
#'
#' x = flrunif(10)
#' lgamma(x)
#' }
#'
#' @useDynLib float R_gamma_spm R_lgamma_spm
#' @name specialmath
#' @rdname specialmath
NULL

gamma_float32 = function(x)
{
  ret = .Call(R_gamma_spm, DATA(x))
  float32(ret)
}

lgamma_float32 = function(x)
{
  ret = .Call(R_lgamma_spm, DATA(x))
  float32(ret)
}

#' @rdname specialmath
#' @export
setMethod("gamma", signature(x="float32"), gamma_float32)

#' @rdname specialmath
#' @export
setMethod("lgamma", signature(x="float32"), lgamma_float32)

# -----------------------------------------------------------------------------
# mathis
# -----------------------------------------------------------------------------

#' Finite, infinite, and NaNs
#'
#' Finite, infinite, and NaNs.
#'
#' @param x
#' A float vector/matrix.
#'
#' @return
#' An integer vector/matrix of the same dimensions as the input.
#'
#' @examples
#' \dontrun{
#' library(float)
#'
#' x = flrnorm(10)
#' is.nan(sqrt(x))
#' }
#'
#' @useDynLib float R_isfinite_spm R_isinfinite_spm R_isnan_spm
#' @name mathis
#' @rdname mathis
NULL

is.finite_float32 = function(x)
{
  .Call(R_isfinite_spm, DATA(x))
}

is.infinite_float32 = function(x)
{
  .Call(R_isinfinite_spm, DATA(x))
}

is.nan_float32 = function(x)
{
  .Call(R_isnan_spm, DATA(x))
}

#' @rdname mathis
#' @export
setMethod("is.finite", signature(x="float32"), is.finite_float32)

#' @rdname mathis
#' @export
setMethod("is.infinite", signature(x="float32"), is.infinite_float32)

#' @rdname mathis
#' @export
setMethod("is.nan", signature(x="float32"), is.nan_float32)

# -----------------------------------------------------------------------------
# rounding
# -----------------------------------------------------------------------------

#' Round
#'
#' Rounding functions.
#'
#' @param x
#' A float vector/matrix.
#' @param digits
#' The number of digits to use in rounding.
#' @param ...
#' ignored
#'
#' @return
#' A float vector/matrix of the same dimensions as the input.
#'
#' @examples
#' library(float)
#'
#' x = flrnorm(10)
#' floor(x)
#'
#' @useDynLib float R_ceiling_spm R_floor_spm R_trunc_spm R_round_spm
#' @name round
#' @rdname round
NULL

ceiling_float32 = function(x)
{
  ret = .Call(R_ceiling_spm, DATA(x))
  float32(ret)
}

floor_float32 = function(x)
{
  ret = .Call(R_floor_spm, DATA(x))
  float32(ret)
}

trunc_float32 = function(x, ...)
{
  ret = .Call(R_trunc_spm, DATA(x))
  float32(ret)
}

round_float32 = function(x, digits=0)
{
  ret = .Call(R_round_spm, DATA(x), as.double(digits))
  float32(ret)
}

#' @rdname round
#' @export
setMethod("ceiling", signature(x="float32"), ceiling_float32)

#' @rdname round
#' @export
setMethod("floor", signature(x="float32"), floor_float32)

#' @rdname round
#' @export
setMethod("trunc", signature(x="float32"), trunc_float32)

#' @rdname round
#' @export
setMethod("round", signature(x="float32"), round_float32)


# 5- Clump Finding Problem

genome <- "AAAATAAAA"

k_mer <- 3
L <- 4
t <- 2
all_sub_genome <- substring(genome, first = 1:(nchar(genome)-L+1) ,
                        last = L:nchar(genome))
all_occuerance <- data.frame()
for (i_sub in all_sub_genome){
sub_dna <- substring(i_sub, first = 1:(nchar(i_sub)-k_mer + 1) ,
                   last = k_mer:nchar(i_sub))

table_count <- table(sub_dna)
all_occuerance <- rbind(all_occuerance, data.frame(table_count) )
}
result <- subset(all_occuerance, Freq >= t)
result_unique <- unique(result)
result_unique


################################################################################
# Part of the R/EpiILM package
#
# AUTHORS:
#         Vineetha Warriyar. K. V. <vineethawarriyar.kod@ucalgary.ca>,
#         Waleed Almutiry <wkmtierie@qu.edu.sa>, and
#         Rob Deardon <robert.deardon@ucalgary.ca>
#
# Free software under the terms of the GNU General Public License, version 2,
# a copy of which is available at http://www.r-project.org/Licenses/.
################################################################################

epiBR0 <- function(x = NULL, y = NULL, contact = NULL, sus.par,  trans.par = NULL, beta,
                   spark = NULL, infperiod, Sformula = NULL, Tformula = NULL, tmax, niter) {

  # Error checks for input arguments
  if (all(is.null(contact) &  (is.null(x) | is.null(y))) == TRUE) {
    stop('epiBR0: Specify contact network or x, y coordinates')
  }

  if (is.null(spark)) {
    spark <- 0
  }

  ns     <- length(sus.par)
  if (is.null(trans.par)) {
    nt <- 1
    trans.par <- 1
  } else {
    nt     <- length(trans.par)
  }
  ni     <- length(beta)

  if (!is.null(x)) {
    n <- length(x)

    if (any((length(y) != n) | (length(x) != n)) == TRUE) {
      stop('epiBR0: Length of x or y is not compatible')
    }
  }

  if (!is.null(contact)) {
    if (is.matrix(contact)) {
      if (dim(contact)[1] != dim(contact)[2]) {
        stop('epiBR0:  The contact network matrix is not a square matrix.')
      }
      n <- dim(contact)[1]
      network <- array(contact, c(n, n, ni))
    } else if (is.array(contact)) {
      if (dim(contact)[1] != dim(contact)[2] | length(contact)/(sqrt(dim(contact)[1]) * sqrt(dim(contact)[2])) != dim(contact)[3]) {
        stop('epiBR0:  One or all of the contact network matrix are not a square matrix.')
      }
      n <- sqrt(dim(contact)[1])
      network <- contact #array(contact, c(n, n, ni))

    } else {
      stop('epiBR0:  The contact network must be specified as an n by n square matrix or an array of n by n square matrices.')
    }
  }

  if (length(infperiod) != n) {
    stop('epiBR0: Length of infperiod is not compatible')
  }

  # formula for susceptibility function
  if (!is.null(Sformula)) {
    covmat.sus <- model.matrix(Sformula)

    if (all((ncol(covmat.sus) == length(all.vars(Sformula))) & (ns != length(all.vars(Sformula)))) == TRUE) {
      stop('epiBR0: Check Sformula (no intercept term) and the dimension of sus.par')
    }

    if (all((ncol(covmat.sus) > length(all.vars(Sformula))) & (ns != ncol(covmat.sus))) == TRUE) {
      stop('epiBR0: Check Sformula (intercept term) and the dimension of sus.par')
    }
  } else {
    if (ns == 1) {
      covmat.sus <- matrix(1.0, nrow = n, ncol = ns)
    }
    if (ns > 1) {
      stop('epiBR0: Please specify the susceptibility covariate')
    }
  }

  # formula for transmissibility function

  if (!is.null(Tformula)) {
    covmat.trans <- model.matrix(Tformula)

    if (all((ncol(covmat.trans) == length(all.vars(Tformula))) & (nt != length(all.vars(Tformula)))) == TRUE) {
      stop('epiBR0: Check Tformula (no intercept term) and the dimension of trans.par')
    }

  } else {
    if (nt == 1) {
      covmat.trans <- matrix(1.0, nrow = n, ncol = 1)
    }
    if (nt > 1) {
      stop('epiBR0: Please specify the transmissibility covariate')
    }
  }

  # Calling Fortran subroutines : contact network model
  if (!is.null(contact)) {
    tmp <- .Fortran("rconsir",
                    n = as.integer(n),
                    tmax = as.integer(tmax),
                    ns = as.integer(ns),
                    nt = as.integer(nt),
                    ni = as.integer(ni),
                    lambda = as.integer(infperiod),
                    suspar = as.numeric(sus.par),
                    transpar = as.numeric(trans.par),
                    beta = as.numeric(beta),
                    spark = as.numeric(spark),
                    covmatsus =  matrix(as.double(covmat.sus), ncol = ncol(covmat.sus), nrow = n),
                    covmattrans = matrix(as.double(covmat.trans), ncol = ncol(covmat.trans), nrow = n),
                    network = as.vector(network),
                    sim = as.integer(niter),
                    val = as.double(0),
                    countinf = as.integer(rep(0,niter))
                    )

    result1 <- list(BasicR0 = tmp$val,
                    simulated_BR0 = tmp$countinf)
  } else {
    # Calling Fortran subroutines : spatial model
    tmp <- .Fortran("rxysir",
                    n = as.integer(n),
                    tmax = as.integer(tmax),
                    ns = as.integer(ns),
                    nt = as.integer(nt),
                    ni = as.integer(ni),
                    suspar = as.numeric(sus.par),
                    transpar = as.numeric(trans.par),
                    beta = as.numeric(beta),
                    spark = as.numeric(spark),
                    covmatsus =  matrix(as.double(covmat.sus), ncol = ncol(covmat.sus), nrow = n),
                    covmattrans = matrix(as.double(covmat.trans), ncol = ncol(covmat.trans), nrow = n),
                    lambda = as.integer(infperiod),
                    x = as.double(x),
                    y = as.double(y),
                    sim = as.integer(niter),
                    val = as.double(0),
                    countinf = as.integer(rep(0,niter)))

    result1 <- list(BasicR0 = tmp$val,
                    simulated_BR0 = tmp$countinf)
  }
  return(result1)
}