#This script will take the Tg data (temp and density) and plot it.
#It will find the Tg temp by fitting two lines to the data.
#You will need to update the breakpoint guess for your system. I recommend setting the initial guess as the Tg value for your system.
#
#License information:
#
#MIT License
#
#Copyright (c) 2019 Will Pisani
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

if (Sys.info()['sysname']=="Linux"){
  setwd("/path/to/directory/on/linux")
} else {
# Paths in Windows in R require \\ instead of \
 setwd("C:\\path\\to\\directory\\on\\Windows")

}

#Read in csv file
tgfile <-"your_file_name_here_TgData.csv"
tgdat <- read.csv(tgfile)
tgfilename <- strsplit(tgfile,"\\.")[[1]]

#create vectors for temp and dens
tgtemp <- tgdat[,1]
tgdens <- tgdat[,2]

#fit a regular linear model to data, this is for the segmented command
lin.mod <- lm(tgdens ~ tgtemp)

#use segmented package to calculated breakpoint (Tg)
segmented.mod <- segmented::segmented(lin.mod, seg.Z = ~tgtemp, psi=416)
#Define moving average function
mav <- function(x,n=1000){filter(x,rep(1/n,n),sides=2)}

tgdens_ma <- mav(tgdens,8500)

# Get point at tg_value
tg_value <- segmented.mod$psi[2]
tg_value_stderror <- segmented.mod$psi[3]
tg_y_value <- segmented.mod$coefficients[2]*tg_value+segmented.mod$coefficients[1]

# Save regular plot to pdf
pdf(paste(tgfilename[1],"pdf",sep="."))
plot(tgtemp,tgdens,col="blue",xlim=c(300,520),ylim=c(1.24,1.34),xlab = 'Temperature (Kelvin)',ylab = 'Density (g/cc)')
#plot(tgtemp,tgdens)
par(new = TRUE)
plot(segmented.mod,col="red",xlim=c(300,520),ylim=c(1.24,1.34),xlab = '',ylab = '',rug=FALSE)
par(new = TRUE)
plot(tg_value,tg_y_value,pch=24,col="white",bg="white",xlim=c(300,520),ylim=c(1.24,1.34),xlab = '',ylab = '')
dev.off()

# Save moving average plot to pdf
ma_name <- paste(tgfilename[1],"ma",sep="_")
pdf(paste(ma_name,"pdf",sep="."))
plot(tgtemp,tgdens_ma,col="lightskyblue",xlim=c(300,520),ylim=c(1.26,1.32),xlab = 'Temperature (K)',ylab = 'Density (g/cc)',cex.lab=1.5, cex.axis=1.5, cex.main=1.5, cex.sub=1.5)
#plot(tgtemp,tgdens)
#Comment out the next part if not using PEEK1013...S4_cool
par(new = TRUE)
clip(300,tg_value+30,-100,100)
abline(coef=segmented.mod$coefficients,col="black")
par(new = TRUE)
clip(tg_value-30,520,-100,100)
abline(coef=c(1.381,-0.0002156),col="black")
#And uncomment the next line
#plot(segmented.mod,col="red",xlim=c(300,500),ylim=c(1.24,1.34),xlab = '',ylab = '',rug=FALSE)
par(new = TRUE)
#Triangle is pch=24
plot(tg_value,tg_y_value,pch=21,col="black",bg="black",xlim=c(300,520),ylim=c(1.26,1.32),xlab = '',ylab = '',cex.lab=1.5, cex.axis=1.5, cex.main=1.5, cex.sub=1.5)
dev.off()

# Write Tg value to file
tg_data <- data.frame(tg_value,tg_value_stderror)
write.table(tg_data, file=paste(tgfilename[1],"txt",sep="."))


sdP <- 8
sdS <- 4
r2 <- 0.4

answer <- sqrt(r2) * sdP / sdS
write(round(answer, 2), stdout())


library(ANTsRCore)
library(ANTsR)
library(extrantsr)
library(fslr)
library(Rcpp)3686＊
library(RcppParallel)
library(parallelDist)
library(rlist)
library(optrees)
Sys.setenv(ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS = 48)
setwd("/home/sugon/Group_Reg/files")
source("../code/find_distance.R")
# source("../code/createParentTable.r")
source("../code/patch_voxel.r")
source("../code/notlargerthan.R")
source("../code/cal_weight.R")
source("../code/cal_phai.R")
options(digits = 9)

# Node Class
node <- setClass(
  "node",
  slots = c(
    pos = "numeric",
    parent = "numeric",
    child = "vector",
    level = "numeric",
    viewed = "logical"
  ),
  prototype = list(
    pos = 0,
    parent = 0,
    child = vector(),
    level = 0,
    viewed = FALSE
  )
)

# Init Params
r<-1
iter<-20
# Bottomright point coor
assign("N", 100, envir = .GlobalEnv)

# Img index of rank of dist to others, increasing.
#distance <- find_distance(0)
distance <- read.csv("distance.csv")
d_sum <- rowSums(distance)
dist_order <- order(d_sum)
t <- 0
# Init delta_r = max distance 2 other images

delta_r<-max(distance[dist_order[1],])

# Init closet mean image
mean_node <- node(pos = dist_order[1])
M_init<-antsImageRead(paste("ext", mean_node@pos, ".nii.gz", sep = ""))

# Image size
assign("size", dim(M_init), envir = .GlobalEnv)
# patch size
coor <- size
# ?
patch_size <- coor
b <- min(size)
b <- round(b*(1-1/iter))+1
r <- r+delta_r*(1/iter)

# Init
# d <- array()
# d[mean_node@pos] <- 0

# patches: all individual patches matrix values & coordinate range, coor: center coordinate of patches
# patches: 8*99=792 elements, coor: logical patches @ current b&r
omega_list <- cal_weight(M_init, mean_node@pos, r, t)
M <- makeImage(voxval = cal_mean(omega_list$omega, omega_list$patches, omega_list$omega_coor)
          , imagesize = size, origin = antsGetOrigin(M_init), direction = antsGetDirection(M_init), pixeltype = "float", spacing = antsGetSpacing(M_init))
#M <- makeImage(voxval = cal_mean(M, omega_list$omega, omega_list$patches, omega_list$omega_coor)
#               , imagesize = size, origin = antsGetOrigin(M), direction = antsGetDirection(M), pixeltype = "float", spacing = antsGetSpacing(M))
antsImageWrite(M, "M0.nii.gz")

rm("omega_list")

# Iteration
for(t in seq(iter))
{
  # t-1
  # Build reg MST
  vec<-vector()
  for(a in 1:N)
    for(b in (a+1):(N+1))
      vec<-rbind(vec, c(a,b))
  distance <- find_distance(t, M)
  distance[lower.tri(distance)]=0
  vec<-cbind(vec, as.vector(distance[which(distance!=0)]))
  MST<-msTreeKruskal(1:101, vec)
  tree <- MST$tree.arcs

  # Cal weight & r
  omega_list <- cal_weight(M, mean_node@pos, r, t)
  r<-r+delta_r*(t/iter)

  omega_list <- c()
    for(s in 1:N)
    {
      if(t==0)
        img<-antsImageRead(paste("ext", s, ".nii.gz", sep = ""))
      else
        img<-antsImageRead(paste("mov", t-1, "-", s, ".nii.gz", sep = ""))
      d_sum<-0
      # Patches
      patch_num<-0
      total_patch<-(255-coor[1])*(255-coor[2])*(255-coor[3])
      omega<-matrix(0, nrow = total_patch, ncol = s)
      patch_size <- coor-1 #patch size is a vector
      for(x in coor[1]:255)
          for(y in coor[2]:255)
              for(z in coor[3]:255)
              {
                  mov_patch<-patch_voxel(patch_size, coor, img)
                  fix_patch<-patch_voxel(patch_size, coor, M)
                  d <- parallelDist(as.vector(mov_patch), as.vector(fix_patch))
                  omega[patch_num, s]<-exp(-d/r)
                  d <- as.matrix(d)
                  patch_num<-patch_num+1
              }
    }
    for(patch_num in seq(total_patch))
        omega_sum[patch_num]<-sum(omega[patch_num, seq(N)])
    for(s in seq(N))
        for(patch_num in seq(total_patch))
            omega[patch_num, s]<-omega[patch_num, s]/omega_sum[patch_num]

  # Compute mean image
  M <- makeImage(voxval = cal_mean(M, omega_list$omega, omega_list$patches, omega_list$omega_coor)
               , imagesize = size, origin = antsGetOrigin(M), direction = antsGetDirection(M), pixeltype = "float", spacing = antsGetSpacing(M))

    M_data<-M_data+sum(((sum(omega[seq(total_patch), s]))/prod(coor))*img[])
    M<-makeImage(voxval = M_data, origin = antsGetOrigin(M), direction = antsGetDirection(M), pixeltype = float, spacing = antsGetSpacing(M))
    antsImageWrite(M, paste("mean", t, ".nii.gz", sep = ""))

    # MST BFS registration
    while(T)
    {
      tar@child <- getChild(mean_node)
      for(i in tar@child)
      {

      }
    }

    # Warp

    # go to leaf
    while(tar)
    {
      searched[tar] <- TRUE
      child <- getChild(tree, tar)
      parent <- NULL
      # Go end left
      while (child) {
        parent <- tar
        tar <- child
        child <- getChild(tree, tar)
      }

      # Reg
      mov <- antsImageRead(paste("mov", iter-1, "-", tar, ".nii.gz", sep = ""))
      fix <- antsImageRead(paste("mov", iter-1, "-", parent, ".nii.gz", sep = ""))
      # Set as variable to directly reach
      antsImageWrite(antsRegistration(moving = mov, fixed = fix, typeofTransform = "SyNOnly"))

    }

    # Update M

}


require(Hmisc)
set.seed(135)
maj <- factor(c(rep('North',13),rep('South',13)))
g <- paste('Category',rep(letters[1:13],2))
n <- sample(1:15000, 26, replace=TRUE)
y1 <- runif(26)
y2 <- pmax(0, y1 - runif(26, 0, .1))
png('/tmp/summaryD.png', width=550, height=800)
spar(mfrow=c(3,2))
f <- function(x) sprintf('%4.2f', x)
summaryD(y1 ~ maj + g, xlab='Mean', auxtitle='', fmtvals=f)
summaryD(y1 ~ maj + g, groupsummary=FALSE)
summaryD(y1 ~ g, fmtvals=f, auxtitle='')
Y <- cbind(y1, y2)
summaryD(Y  ~ maj + g, fun=function(y) y[1,], symbol=c(1,17))
rlegend(.1, 26, c('y1','y2'), pch=c(1,17))

summaryD(y1 ~ maj, fun=function(y) c(Mean=mean(y), n=length(y)),
         auxvar='n')
dev.off()

# options(grType='plotly')
sym <- if(grType() == 'plotly') c('circle', 'line-ns-open') else c(21, 3)
h <- function(x) c(mean=mean(x), Q1=unname(quantile(x, .25)),
                 Q3=unname(quantile(x, .75)), N=length(x))
summaryD(Y  ~ maj + g, fun=h, auxvar='N', symbol=sym[c(1,2,2)],
         col=colorspace::rainbow_hcl(2)[c(1,2,2)],
         legendgroup=c('Mean', 'Quartiles', 'Quartiles'))

png('/tmp/summaryD2.png', width=300, height=100)
# Or: pdf('/tmp/z.pdf', width=3.5, height=1.25)
spar()
summaryD(y1 ~ maj, fmtvals=function(x) round(x,4),
         xlab=labelPlotmath('Velocity', 'm/s'))
dev.off()


#+setup, include=FALSE, cache=FALSE
library(knitr)
opts_chunk$set(fig.path='flows-figures/', fig.align='center', fig.show='hold')
options(replace.assign=TRUE,width=80)

#' ## Load Packages

set.seed(4253)
library(lme4)
library(car)
library(Hmisc)
library(pscl)
library(glmmADMB)
library(ggplot2)
print(sessionInfo())

#' ## Data prep
iqr <- function(x) diff(quantile(x, probs=c(0.25, 0.75), na.rm=TRUE))
scaleiqr <- function(x) scale(x, scale=iqr(x))
clog <- function(x) log(x + 0.5)

ObservedMatchingInternalFlows <- function(flows, case.data) {
    unwanted <- c('week', 'totalNumberSwineAccessions', 'Unk')
    ind <- which(!colnames(case.data) %in% unwanted)
    observed <- case.data[, ind]
    avail <- rownames(flows)[!is.na(flows[, 'impInternalFlow'])]
    zeros <- setdiff(avail, colnames(observed))
    observed[, zeros] <- 0
    observed
}

ReshapeAugmentObserved <- function(observed, fiAll, flows, stateCty) {
  om <- reshape(observed, v.names='cases', varying=colnames(observed),
                direction='long', timevar='state', times=colnames(observed),
                idvar='week')
  splt <- split(om, om$state)
  tmpf <- function(x) {
      ord <- order(x$week)
      x <- x[ord, ]
      x$cumCases <- cumsum(x$cases)
      x
  }
  proc <- lapply(splt, tmpf)
  om <- do.call(rbind, args=proc)
  om$nFarms <- fiAll[om$state, ]$X2002.Farms
  om$nSusc <- om$nFarms - om$cumCases
  omprev <- om
  omprev$week <- omprev$week + 1
  omprev$nFarms <- NULL
  omprev$cumCases <- NULL
  c1 <- which(colnames(omprev) == 'cases')
  c2 <- which(colnames(omprev) == 'nSusc')
  colnames(omprev)[c1] <- 'cases1WksAgo'
  colnames(omprev)[c2] <- 'nSusc1WksAgo'
  om <- merge(om, omprev, by=c('week', 'state'), all.x=TRUE)
  omprev$week <- omprev$week + 1
  colnames(omprev)[c1] <- 'cases2WksAgo'
  colnames(omprev)[c2] <- 'nSusc2WksAgo'
  om <- merge(om, omprev, by=c('week', 'state'), all.x=TRUE)
  omprev$week <- omprev$week + 1
  colnames(omprev)[c1] <- 'cases3WksAgo'
  colnames(omprev)[c2] <- 'nSusc3WksAgo'
  om <- merge(om, omprev, by=c('week', 'state'), all.x=TRUE)
  omprev$week <- omprev$week + 1
  colnames(omprev)[c1] <- 'cases4WksAgo'
  colnames(omprev)[c2] <- 'nSusc4WksAgo'
  om <- merge(om, omprev, by=c('week', 'state'), all.x=TRUE)
  om$internalFlow <- flows[om$state, "impInternalFlow"]
  rownames(om) <- paste(om$week, om$state, sep='.')
  test <- !is.na(om$internalFlow)
  om <- om[test, ]
  om <- merge(om, stateCty, by='state')
  rownames(om) <- paste(om$week, om$state, sep='.')
  om$logInternalFlowScaled <- scaleiqr(log(2 * om$internalFlow))
  om$weekCent <- scale(om$week, center=TRUE, scale=FALSE)
  om$logCmedDenseScaled <- scaleiqr(log(om$cmedDense * om$nFarms * om$nFarms))
  om$logMDenseScaled <- scaleiqr(log(om$mDense * om$nFarms * om$nFarms))
  om$clogCases1wa <- clog(om$cases1WksAgo)
  om$clogCases2wa <- clog(om$cases2WksAgo)
  om$clogCases3wa <- clog(om$cases3WksAgo)
  om$clogCases4wa <- clog(om$cases4WksAgo)
  om$clogCases1wa05 <- log(om$cases1WksAgo + 0.05)
  om$clogCases1wa005 <- log(om$cases1WksAgo + 0.005)
  om
}

data('real.case.data', package='sds')
data('internal.flows', package='sds')
data('farms.and.inventory', package='sds')
data('state.summaries', package='sds')

observed <- ObservedMatchingInternalFlows(flows=internal.flows, case.data=real.case.data)
om <- ReshapeAugmentObserved(observed=observed, fiAll=farms.and.inventory,
                             flows=internal.flows, stateCty=state.summaries)

#' ## Model fitting
#'
#' Let's start with a least square fit on the logged response. We'll
#' center the predictors around means to reduce the correlation
#' between slopes and intercepts.

m <- list()
f <- list()
f$lm <- as.formula(clog(cases) ~ clogCases1wa + logInternalFlowScaled + logCmedDenseScaled + weekCent + offset(log(nSusc1WksAgo) - 2*log(nFarms)))
m$lm <- lm(f$lm, data=om)
summary(m$lm)

#' We've included week to account for the general trend of increasing
#' cases over the period studied, possibly due to seasonal changes in
#' manure spreading, virus survivablility, or pig production.

tmpf <- function() {
  allCases <- with(om, tapply(cases, week, sum))
  week <- as.integer(names(allCases))
  plot(week, allCases, log='y')
}
tmpf()

#' Next we use a random effect to account for lack of dependence within states.
f$lme <- update(f$lm, . ~ . + (1|state))
f$lmeN <- update(f$lme,  . ~ . - logInternalFlowScaled)
m$lme <- lmer(f$lme, data=om, REML=FALSE)
m$lmeN <- lmer(f$lmeN, data=om, REML=FALSE)
m$lme

likRatio <- function(null, alt){
    null <- logLik(null)
    alt <- logLik(alt)
    dfn <- attr(null, "df")
    dfa <- attr(alt, "df")
    if(dfn >= dfa) stop("null model has greater or same df than alternative")
    df <- dfa - dfn
    likRatio <- -2 * (as.numeric(null) - as.numeric(alt))
    p <- pchisq(likRatio, df=df, lower.tail = FALSE)/2
    res <- c(lr=likRatio, df=df, p=p)
    signif(res, 3)
}
likRatio(m$lmeN, m$lme)
#' The effect of flows is still statistically significant.

tmpf <- function() {
    df <- model.frame(m$lme)
    iqry <- quantile(df$"clog(cases)", probs=c(.1, .9))
    effectSize <- fixef(m$lme)['logInternalFlowScaled']
    effectSize/diff(iqry)
}
tmpf()
#' The effect size is non-negligible.
#'
#' A count model would be more appropriate for the discrete outcome.

f$ct <- update(f$lm, cases ~ .)
m$pois <- glm(f$ct, family=poisson, data=om)
getDispStat <- function(mod) {
    df <- mod$df.residual
    E <- residuals(mod, type = "pearson")
    sum(E^2)/df
}
getDispStat(m$pois)

#' The dispersion parameter is fairly large. Let's check for outliers.

(otest <- outlierTest(m$pois))

#' There's a fairly large number, given the size of the data set,
#' which suggests that a negative binomial model may be more
#' appropriate. Such a model would also have a theoretical
#' justification, as in the TSIR model.

m$nb <- glm.nb(f$ct, data=om)
summary(m$nb)
likRatio(m$pois, m$nb)

#' The likelihod ratio test supports the negative binomial model.
#' Is this support sensitive to outliers?

tmpf <- function() {
    orows <- names(otest$rstudent)
    test <- !rownames(om) %in% orows
    omno <- om[test, ]
    mpno <- glm(f$ct, family='poisson', data=omno)
    print(getDispStat(mpno))
    print(outlierTest(mpno))
    mnbno <- glm.nb(f$ct, data=omno)
    print(summary(mnbno))
    print(likRatio(mpno, mnbno))
}
tmpf()
rm(otest)
#' No.

#' Does the negative binomial model support the includion of the flow term?

m$nbN <- update(m$nb, . ~ . - logInternalFlowScaled)
likRatio(m$nbN, m$nb)

#' Yes. What if we use a different predictor for geographic density?
#' The current choice is highly correlated with flows on the log scale
#' as shown next, but we'll try another one to be safe.

tmpf <- function(stateCty, fiAll, flows) {
    sel <- unique(om$state)
    rownames(stateCty) <- stateCty$state
    den <- stateCty[sel, c('mDense', 'cmDense', 'medDense', 'cmedDense')]
    nf <- fiAll[sel,]$X2002.Farms
    fl <- flows[sel, 'impInternalFlow']
    tmpff <- function(x) {
        list(flowVdense=cor(log(cbind(flow=fl, den)), method=x),
             flowVdenseNf=cor(log(cbind(flow=fl, den*nf)), method=x),
             flowVdesneNfNf=cor(log(cbind(flow=fl, den*nf*nf)), method=x))
    }
    lapply(c(pearson='pearson', spearman='spearman'), tmpff)
}
tmpf(stateCty=state.summaries, fiAll=farms.and.inventory, flows=internal.flows)

m$nbmDense <- update(m$nb, . ~ . - logMedDenseScaled + logMDenseScaled)
m$nbmDenseN <- update(m$nbmDense, . ~ . - logInternalFlowScaled, control=glm.control(maxit=60))
summary(m$nbmDense)
likRatio(m$nbmDenseN, m$nbmDense)

#' There coefficients and significance level are similar.

#' Let's test different lags.

f$ct2 <- as.formula(cases ~ clogCases2wa + logCmedDenseScaled + logInternalFlowScaled + weekCent + offset(log(nSusc2WksAgo) - 2 * log(nFarms)))
m$nb2 <- glm.nb(f$ct2, data=om)
logLik(m$nb2)
logLik(glm.nb(f$ct, data=om[om$week > 2, ]))

f$ct3 <- as.formula(cases ~ clogCases3wa + logCmedDenseScaled + logInternalFlowScaled + weekCent + offset(log(nSusc3WksAgo) - 2 * log(nFarms)))
m$nb3 <- glm.nb(f$ct3, data=om, control=glm.control(maxit=50))
logLik(m$nb3)
logLik(glm.nb(f$ct, data=om[om$week > 3, ]))

f$ct4 <- as.formula(cases ~ clogCases4wa + logCmedDenseScaled + logInternalFlowScaled + weekCent + offset(log(nSusc4WksAgo) - 2 * log(nFarms)))
m$nb4 <- glm.nb(f$ct4, data=om, control=glm.control(maxit=50))
logLik(m$nb4)
logLik(glm.nb(f$ct, data=om[om$week > 4, ]))

#' A lag of 1 is by far the best by likelihood.

#' Let's try different baseline hazards.

f$nb05 <- update(f$ct, . ~ . - clogCases1wa + clogCases1wa05)
m$nb05 <- glm.nb(f$nb05, data=om, control=glm.control(maxit=100))

f$nb005 <- update(f$ct, . ~ . - clogCases1wa + clogCases1wa005)
m$nb005 <- glm.nb(f$nb005, data=om, control=glm.control(maxit=100))

lapply(m[c('nb', 'nb05', 'nb005')], logLik)

# A baseline hazard of 0.05 is best, although the model seems fairly
# insensitive to this parameter.

#' Let's compare the observed and expected distribution of cases after
#' summing over the variation in the mean.

makeCompPlot <- function(xvals = 0:10, mod, exclude = integer(0), ...) {
    ypred <- predict(mod, type = "response")
    tmpf <- function(x) {
        ret <- dnbinom(xvals, mu = x, size=mod$theta)
        ret/sum(ret)
    }
    ds <- sapply(ypred, tmpf)
    expectedProbs <- rowSums(ds)/ncol(ds)
    resp <- mod$model$cases
    if (length(exclude) > 0) {
        resp <- resp[-exclude]
    }
    obsResp <- table(resp)
    expResp <- expectedProbs * sum(obsResp)
    ye <- expResp
    yo <- obsResp[as.character(xvals)]
    plot(xvals, yo, pch = "o", xlab = "response value", ylab = "number", ...,
         type = "b")
    points(xvals, ye, pch = "e", type = "b")
}
makeCompPlot(mod=m$nb, xvals=0:30, log='y')

#' There's no strong signs of zero inflation, or any other major
#' defects.

m$nbzi <- zeroinfl(f$ct, data=om, dist="negbin")
AIC(m$nb, m$nbzi)

#' One could justify a zero-inflated model based on AIC, but it's not
#' clear to me such a model is better here. What mechanism would cause
#' the zero-inflation, and do we have a good reason to spend degrees
#' of freedom on it?
#'
#' Let's check for evidence of non-linearity in partial residual plots.

crPlots(m$nb)

#' Toward the extreme ends of the data range, there is some sign of
#' non-linearity, but generally linearity assumptions seem reasonable.

#' Next let's check for influential observations.

influencePlot(m$nb)

tmpf <- function() {
    filterInf <- function(mod, df, cutoff=0.2) {
        dfb <- dfbetas(mod)
        test <- rowSums(dfb > cutoff) > 0
        if(any(test)) {
            inf <- dfb[test, ]
            print(inf)
            infNames <- names(which(test))
             df[!rownames(df) %in% infNames, ]
        } else {
            NULL
        }
    }
    i <- 1
    mods <- list(m$nb)
    omni <- filterInf(m$nb, om)
    while(!is.null(omni)) {
        omnip <- omni
        mnbni <- glm.nb(f$ct, data=omnip)
        omni <- filterInf(mnbni, omni)
        i <- i + 1
        mods[[i]] <- mnbni
    }
    print(sapply(mods, coef))
    mnbniNull <- update(mnbni, . ~ . - logInternalFlowScaled, control=glm.control(maxit=100))
    print(likRatio(mnbniNull, mnbni))
}
tmpf()

#' Removing the influential observations changes the coefficients a
#' bit, but the significance of the flows term is not affected.
#'
#' Next, let's account for correlations of data from the same state by
#' using a clustered bootstrap variance estimate.

clusteredBoot <- function(data, statistic, B, cluster){
    ## based on rms::bootcov() and boot::boot()
    n <- nrow(data)
    if (length(cluster) != n) stop("length of cluster does not match # rows used in fit")
    if (any(is.na(cluster))) stop("cluster contains NAs")
    cluster <- as.character(cluster)
    clusters <- unique(cluster)
    nc <- length(clusters)
    Obsno <- split(1:n, cluster)
    pb <- txtProgressBar(min=1,max=B)
    drawStat <- function(x) {
        setTxtProgressBar(pb, x)
        j <- sample(clusters, nc, replace=TRUE)
        obs <- unlist(Obsno[j])
        statistic(data, obs)
    }
    res <- sapply(1:B, drawStat, simplify='array')
    close(pb)
    res
}

tmpf <- function(B) {
    fun <- function(data, i) {
        m <- glm.nb(f$ct, data=data[i, ], control=glm.control(maxit=200))
        coef(summary(m))[, 1:2]
    }
    res <- clusteredBoot(om, statistic=fun, B=B, cluster=om$state)
    ests <- res['logInternalFlowScaled', 'Estimate', ]
    print(quantile(ests, probs=c(0.025, 0.975)))
    print(quantile(ests, prob=1:5/100))
    qqPlot(ests)
    print(splom(t(res[, 'Estimate', ])))
    res
}
system.time(clusterBootDf <- tmpf(B=300))
save.image(file='flows-checkpoint1.RData')
date()

#' It seems like a good idea to also estimate a parameter for a
#' non-linear effect of the number of susceptibles, since we do that
#' for the number of infectives.

tmpf <- function() {
    form <- as.formula(cases ~ clogCases1wa + logInternalFlowScaled + weekCent + log(nSusc1WksAgo) + offset(-2 * log(nFarms)))
    func <- function(data, i) {
        m <- glm.nb(form, data=data[i, ], control=glm.control(maxit=100))
        coef(summary(m))[, 1:2]
    }
    res2 <- clusteredBoot(om, statistic=func, B=20, cluster=om$state)
    print(splom(t(res2[, 'Estimate', ])))
    print(cor(t(res2[c('log(nSusc1WksAgo)', '(Intercept)'), 'Estimate', ])))
}
tmpf()

#' However, the estimate turns out to be highly correlated with the
#' intercept, probably because the number of suceptibles varies little
#' in relation to the number of farms. We leave this parameter out to
#' avoid inflation of the variance of our other estimates.

#' Next we a random effects model using the lme4 package to see if our
#' results are robust.

#' Check that we understand the parameterization of glmer.nb
tmpf <- function() {
    x <- runif(400) - 0.5
    z <- gl(n=40, k=10)
    m <- model.matrix(~x + z)
    u <- rnorm(40, sd=0.5)
    eta <- m %*% c(0, 3, u[-1] - u[1])
    y <- rnbinom(n=length(eta), size=3, mu=exp(eta))
    data.frame(y, x, z)
}
simdf <- tmpf()
glmer.nb(y~x + (1|z), nAGQ=1L, data=simdf)

#' It would seem that nAGQ>1 should increase accuracy, but rather we
#' find it shrinks the estimate of the size parameter for the negative
#' binomial:

glmer.nb(y~x + (1|z), nAGQ=10L, data=simdf)
rm(simdf)

#' Now with the data:

om$off <- with(om, log(nSusc1WksAgo) - 2 * log(nFarms))
om$offCen <- om$off - mean(om$off, na.rm=TRUE)
# centering the offset prevents a scaling error: https://github.com/lme4/lme4/issues/173
om$stateF <- factor(om$state)

f$glme <- as.formula(cases ~ clogCases1wa + logCmedDenseScaled + logInternalFlowScaled + weekCent + (1|stateF) + offset(offCen))
m$nbme <- glmer.nb(f$glme, data=om, nAGQ=1L)
f$glmeN <- update(f$glme, . ~ . -logInternalFlowScaled)
m$nbmeN <- glmer.nb(f$glmeN, data=om, nAGQ=1L)
likRatio(m$nbmeN, m$nbme)
summary(m$nbme)
summary(m$nbmeN)

#' Now we use glmmADMB for comparison

f$admb <- f$glme
omcc <- om[rownames(model.frame(m$nbme)), ]
m$admb <- glmmadmb(f$admb, data=omcc, family='nbinom',
                   admb.opts=admbControl(shess=FALSE, impSamp=100),
                   extra.args='-gh 10')
m$admb200 <- glmmadmb(f$admb, data=omcc, family='nbinom',
                      admb.opts=admbControl(shess=FALSE, impSamp=200))
m$admb0 <- glmmadmb(f$admb, data=omcc, family='nbinom')
summary(m$admb0)
signif(sapply(list(m$nbme, m$admb, m$admb0, m$admb200), fixef), 3)

#' The results look nearly identical, so importance sampling is
#' probably not needed to ensure accuracy of the integration.

f$admbN <- update(f$admb, . ~ . - logInternalFlowScaled)
m$admbN <- glmmadmb(f$admbN, data=omcc, family='nbinom')
likRatio(m$admbN, m$admb0)

#' glmmADMB also supports the inclusion of the flow terms.
#'
#' Output the descriptive stats for this test

omcc$logInternalFlowScaled <- as.numeric(omcc$logInternalFlowScaled)
omcc$clogCases1wa <- as.numeric(omcc$clogCases1wa)
omcc$logCmedDenseScaled <- as.numeric(omcc$logCmedDenseScaled)
omcc$weekCent <- as.numeric(omcc$weekCent)
desc <- describe(omcc[, c('cases', 'clogCases1wa', 'logCmedDenseScaled',
                          'logInternalFlowScaled', 'weekCent', 'stateF',
                          'offCen')])
invisible(latex(desc, file='tsir-describe.tex'))

save.image('flows-checkpoint2.RData')
date()
rm(omcc)

#' Next we plot the response distribution simplify to get some feeling
#' for our fitted model's behavior.

plotResponseDist <- function(mod, size){
    tmpf <- function(y) {
        x <- 1:40
        dnbinom(x, mu=y, size=size)
    }
    q <- quantile(predict(mod, type='response'), probs=c(.05, .25, .5, .75, .95))
    m <- sapply(q, tmpf)
    main <- paste('Negative Binomial, Theta == ', signif(size, 3))
    matplot(m, type='l', xlab='Cases', ylab='Probability', main=main)
    leg <- signif(q, 2)
    col <- seq_len(length(q))
    lty <- col
    legend(x='topright', legend=leg, lty=lty, title='Mean', col=col)
}
plotResponseDist(m$admb0, summary(m$admb0)$alpha)
plotResponseDist(m$nb, m$nb$theta)

#' Next we check for autocorrelation of the residuals in time.

tmpf <- function(mnb) {
    E <-resid(mnb, type='deviance')
    df <- data.frame(wk=mnb$model$week, st=om[names(E), ]$state  , E=E)
    df <- reshape(df, timevar='st', idvar='wk', direction='wide')
    rownames(df) <- df$wk
    df <- df[order(df$wk), ]
    dft <- as.ts(df[, -1])
    maxInd <- ceiling(ncol(dft) / 4)
    makeInd <- function(x) {
        i <- (x - 1)*4 + 1;
        i:(i + 3)
    }
    ind <- lapply(1:maxInd, makeInd)
    plot0 <- function(x) {
        x <- x[x <= ncol(dft)]
        plot.ts(dft[, x], main='', nc=2, plot.type='multiple', ylab=colnames(dft)[x])
    }
    lapply(ind, plot0)
    ac <- lapply(df[, -1], acf, lag.max=4, plot=FALSE)
    par(mfrow=c(2,2))
    plot1 <- function(x, y) if(!is.na(y)) plot(x, main=y)
    makePanel <- function(x) mapply(plot1, ac[x], names(ac)[x])
    lapply(ind, makePanel)
    par(mfrow=c(1,1))
}
tmpf(m$nb)

#' No major autocorrelations are present for those time series having cases.
#'
#' Now let's see if the data support models with of spread between
#' states. First we try incorporating this by weighting cases in the
#' previous week by the amount of flow between that state and other
#' states.

tmpf <- function(rel='directed', flowMat, flows, fiAll, observed) {
    fmo <- flowMat[state.abb, state.abb]
    diag(fmo) <- flows[state.abb, 'impInternalFlow']
    Fund <- fmo + t(fmo)
    Fdir <- fmo
    ## Fdir[i, j] == head sent to i from j
    if(rel == 'directed') {
        F <- fmo
    } else {
        F <- fmo + t(fmo)
    }
    Fsum <- rowSums(F)
    Fnorm <- F/Fsum
    n <- fiAll[state.abb, ]$X2002.Farms
    W <- Fnorm * n
    W <- t(t(W) / n)
    sel <- colnames(observed)
    W <- W[sel, sel]
    Wobs <- W %*% t(observed)
    list(ts=t(Wobs), Fsum=Fsum)
}
data('flows.matrix', package='sds')
wcDir <- tmpf('directed', flowMat=t(flows.matrix), flows=internal.flows,
              fiAll=farms.and.inventory, observed=observed)
wcUnd <- tmpf('undirected', flowMat=t(flows.matrix), flows=internal.flows,
              fiAll=farms.and.inventory, observed=observed)

par(mfrow=c(3,1))
plot.ts(wcDir$ts, plot.type='single')
plot.ts(wcUnd$ts, plot.type='single')
plot.ts(observed, plot.type='single')
par(mfrow=c(1,1))

tmpf <- function() {
    wc <- wcDir$ts
    wc <- reshape(data.frame(wc), v.names='casesDirectedFlowWeighted', varying=colnames(wc),
                  direction='long', timevar='state', times=colnames(wc), idvar='week')
    wc$week <- wc$week + 1
    colnames(wc)[2] <- 'casesLastWeekDirectedFlowWeighted'
    om <- merge(om, wc, by=c('week', 'state'), all.x=TRUE)

    om$directedFlow <- wcDir$Fsum[om$state]
    om$logDirectedFlowScaled <- scaleiqr(log(om$directedFlow))
    om$clogCases1waDir <- clog(om$casesLastWeekDirectedFlowWeighted)

    wc <- wcUnd$ts
    wc <- reshape(data.frame(wc), v.names='casesDirectedFlowWeighted', varying=colnames(wc),
                  direction='long', timevar='state', times=colnames(wc), idvar='week')
    wc$week <- wc$week + 1
    colnames(wc)[2] <- 'casesLastWeekUndirectedFlowWeighted'
    om <- merge(om, wc, by=c('week', 'state'), all.x=TRUE)

    om$undirectedFlow <- wcUnd$Fsum[om$state]
    om$logUndirectedFlowScaled <- scaleiqr(log(om$undirectedFlow))
    om$clogCases1waUnd <- clog(om$casesLastWeekUndirectedFlowWeighted)
    rownames(om) <- paste(om$week, om$state, sep='.')
    om
}
om <- tmpf()

f$ctdir <- update(f$ct, . ~ . + clogCases1waDir + logDirectedFlowScaled - clogCases1wa - logInternalFlowScaled)
m$nbdir <- glm.nb(f$ctdir, data=om)

f$ctundir <- update(f$ct, . ~ . + clogCases1waUnd + logUndirectedFlowScaled - clogCases1wa - logInternalFlowScaled)
m$nbundir <- glm.nb(f$ctundir, data=om)

m[c('nb', 'nbdir', 'nbundir')]

#' The internal flow model has the lowest AIC among the fixed effects
#' models. Both of the other models have negligible time components,
#' but otherwise the coefficients are similar. The undirected model
#' has the highest AIC.

f$glmedir <- update(f$glme, . ~ . + clogCases1waDir + logDirectedFlowScaled - clogCases1wa - logInternalFlowScaled)
m$nbmedir <- glmer.nb(f$glmedir, data=om, nAGQ=1L)
m$admbdir <- glmmadmb(f$glmedir, data=om[!is.na(om$clogCases1waDir), ], family='nbinom')

sapply(m[c('nbmedir', 'admbdir')], fixef)

#' ADMB and lme4 again give similar fits.

f$glmeundir <- update(f$glme, . ~ . + clogCases1waUnd + logUndirectedFlowScaled - clogCases1wa - logInternalFlowScaled)
m$nbmeundir <- glmer.nb(f$glmeundir, data=om, nAGQ=1L)
m$admbundir <- glmmadmb(f$glmeundir, data=om[!is.na(om$clogCases1waUnd), ], family='nbinom')

sapply(m[c('nbmeundir', 'admbundir')], fixef)

#' Another good match.

sort(sapply(m[c('nbme', 'nbmedir', 'nbmeundir')], AIC))

#' The undirected flow matrix has the lowest AIC.
#'
#' One shortcoming of the above comparison is that we are not
#' optimizing the baseline hazard, and so the result may be sensitive
#' to our initial and somewhat arbitrary setting. So next we optimize
#' this parameter.

fint <- function(x,data) {
    data$logInf <- log2(data$cases1WksAgo + x)
    fm <- as.formula(cases ~ logInf + logInternalFlowScaled + weekCent + (1 | stateF) + offset(offCen))
    fit <- glmmadmb(fm, data=data[!is.na(data$logInf), ], family='nbinom')
    ll <- logLik(glmer.nb(fit, data=data))
    print(paste(signif(c(x, ll), 5), collapse=","))
    ll
}
(hint <- optimize(f=fint, interval=c(0.01, 5), data=om, tol=0.1, maximum=TRUE))

fund <- function(x,data) {
    data$logInf <- log(data$casesLastWeekUndirectedFlowWeighted + x)
    fm <- as.formula(cases ~ logInf + logUndirectedFlowScaled + weekCent + (1 | stateF) + offset(offCen))
    fit <- glmmadmb(fm, data=data[!is.na(data$logInf), ], family='nbinom')
    ll <- as.numeric(logLik(fit))
    print(paste(signif(c(x, ll), 5), collapse=","))
    ll
}
(hund <- optimize(f=fund, interval=c(0.01, 5), data=om, tol=0.1, maximum=TRUE))

fdir <- function(x,data) {
    data$logInf <- log(data$casesLastWeekDirectedFlowWeighted + x)
    fm <- as.formula(cases ~ logInf + logDirectedFlowScaled + weekCent + (1 | stateF) + offset(offCen))
    fit <- glmmadmb(fm, data=data[!is.na(data$logInf), ], family='nbinom')
    ll <- as.numeric(logLik(fit))
    print(paste(signif(c(x, ll), 5), collapse=","))
    ll
}
(hdir <- optimize(f=fdir, interval=c(0.01, 5), data=om, tol=0.1, maximum=TRUE))

om$logInfInt <- log(om$cases1WksAgo + hint$maximum)
om$logInfUnd <- log(om$casesLastWeekUndirectedFlowWeighted + hund$maximum)
om$logInfDir <- log(om$casesLastWeekDirectedFlowWeighted + hdir$maximum)

f$glmeundirhmax <- update(f$glmeundir, . ~ . - clogCases1waUnd + logInfUnd)
m$glmeundirhmax <- glmer.nb(f$glmeundirhmax, data=om)

f$glmehmax <- update(f$glme, . ~ . - clogCases1wa + logInfInt)
m$glmehmax <- glmer.nb(f$glmehmax, data=om)

f$glmedirhmax <- update(f$glmedir, . ~ . - clogCases1waDir + logInfDir)
m$glmedirhmax <- glmer.nb(f$glmedirhmax, data=om)

sapply(m[c('glmehmax', 'glmeundirhmax', 'glmedirhmax')], logLik)
abs(unname(diff(sapply(m[c('glmehmax', 'glmeundirhmax')], AIC))))

#' We have a difference of about 11 in the AIC now, with the undirected
#' flow model preferred.

save.image('flows-checkpoint3.RData')
date()

#' Next we make diagnostic plots to verify that no assumptions are
#' badly violated and to see what parts of the data seem to lead to
#' differences in the fits.

fmint <- fortify(m$glmehmax, data=model.frame(m$glmehmax))
fund <- fortify(m$glmeundirhmax, data=model.frame(m$glmeundirhmax))
fdir <- fortify(m$glmedirhmax, data=model.frame(m$glmedirhmax))

cols <- c('stateF', '.fitted', '.resid', '.scresid', 'weekCent', 'logCmedDenseScaled')

D <- rbind(cbind(fmint[, cols], foi=fmint$logInfInt,
                 flow.measure=fmint$logInternalFlowScaled, flow.type='internal'),
           cbind(fund[, cols], foi=fund$logInfUnd,
                 flow.measure=fund$logUndirectedFlowScaled, flow.type='undirected'),
           cbind(fdir[, cols], foi=fdir$logInfDir,
                 flow.measure=fdir$logDirectedFlowScaled, flow.type='directed'))

#'
#+ fig.width=14, fig.height=14, dev.args=list(pointsize=18)
g <- ggplot(D, aes(.fitted, .scresid, colour=flow.type))
g <- g + geom_point() + facet_wrap(~stateF, scales='free') + geom_hline(yintercept=0)
g

g <- ggplot(D, aes(foi, .scresid))
g <- g + geom_point(colour='blue') + facet_grid(.~flow.type)
g <- g + geom_line(aes(group=stateF), alpha=0.4) + geom_smooth(method="loess", span=.9)
g

g <- ggplot(D, aes(flow.measure, .scresid))
g <- g + geom_point(colour='blue') + facet_grid(.~flow.type)
g <- g + geom_line(aes(group=stateF), alpha=0.4) + geom_smooth(method="loess", span=.9)
g

g <- ggplot(D, aes(logCmedDenseScaled, .scresid))
g <- g + geom_point(colour='blue') + facet_grid(.~flow.type)
g <- g + geom_line(aes(group=stateF), alpha=0.4) + geom_smooth(method="loess", span=.9)
g

g <- ggplot(D, aes(weekCent, .scresid))
g <- g + geom_point(colour='blue') + facet_grid(.~flow.type)
g <- g + geom_line(aes(group=stateF), alpha=0.4) + geom_smooth(method="loess", span=.9)
g

#' None of the models have large trends in the Pearson residuals
#' againt the predictors.
#'
#+ fig.width=14, fig.height=14, dev.args=list(pointsize=18)
ggplot(D, aes(stateF, .scresid, colour=flow.type)) + geom_boxplot() + geom_hline(yintercept=0)

ggplot(D, aes(flow.type, .scresid, colour=flow.type)) + geom_violin()

#' The distributions of Pearson residuals also look similar.

tmpf <- function(x) ranef(x)$state[, '(Intercept)']
re <- sapply(m[c('glmehmax', 'glmeundirhmax', 'glmedirhmax')], tmpf)
rownames(re) <- rownames(ranef(m$glmehmax)$stateF)
dotplot(re, pch=c('i', 'u', 'd'))

#' The conditional modes for the random effects are generally similar.

lapply(m[c('glmehmax', 'glmeundirhmax', 'glmedirhmax')], summary)

#' The estimated dispersion parameters and state group variance
#' provides one way of summarizing the difference in the models.


#' Multivariate metropolis (i.e., metropolis within gibbs) sampler
#' @param target Target function (returning log unnormalised posterior density);
#' 	this function should take the parameter as its first argument and a data list as its second
#' @param initial Initial parameter vector
#' @param data Data to pass to the target
#' @param iter Number of iterations
#' @param scale Scale vector for the proposal distributions, if NA, the algorithm will adaptively
#' 		choose scales
#' @param n_adapt Number of iterations for adaptation, if needed
#' @param ... Additional named parameters to pass to m_adapt
#'
#' @return A list, with three components: 'chain' is the markov chain, 'scale'
#' 		is the scale parameter used, and 'accept' is the acceptance rate
metropolis = function(target, initial, data, iter = 1000, scale = NA, ...) {

	# check that the target distribution is defined at the initial value
	if(!is.finite(target(initial, data)))
		stop("Target distribution undefined at the initial value, choose better inits")

	# adaptation, if needed
	if(any(is.na(scale)))
		scale = m_adapt(target, initial, data, ...)

	# do the sampling
	result = m_sample(target, initial, data, scale, iter)

	# save the scale as well, because this allows resuming the chain
	result$scale = scale
	return(result)
}

#' Worker function to actually do the sampling
#' @param target Target function (returning log unnormalised posterior density);
#' 	this function should take the parameter as its first argument and a data list as its second
#' @param state Initial state vector (parameter values)
#' @param data Data to pass to the target
#' @param scale Scale for the proposal distributions
#' @param iter Number of iterations
m_sample = function(target, state, data, scale, iter) {
	# number of parameters
	k = length(state)

	# set up the markov chain
	chain = matrix(NA, nrow=iter, ncol=k)
	colnames(chain) = names(state)

	# keep track of how many times we accept the proposals
	# the acceptance rate is an important diagnostic
	accept = rep(0, k)

	for(i in 1:(iter)) {
		# for multivariate problems, it's helpful to keep track of the previous state as well
		prev_state = state

		# for each iteration, we need to randomize the order of parameters
		jj = sample(k)

		# loop over parameters one at a time
		for(j in jj) {
			# propose a value and accept/reject, updating the current state with the new value
			state = m_propose(state, scale, target, data, j)
			if(state[j] != prev_state[j]) {
				# if the value changed, that was an acceptance, increment accept variable
				accept[j] = accept[j] + 1
			}
		}
		# after we have proposed for each parameter, update the chain
		chain[i, ] = state
	}
	return(list(chain = chain, accept = accept/iter))
}

#' Runs the adaptation phase of the metropolis sampler
#' @param target Target function
#' @param state Initial vector of k parameter values
#' @param data Data to pass to the target
#' @param chunk_size Number of iterations in each chunk before adapting
#' @param max_chunk Maximum number of chunks before stopping adaptation
#' 		or divide if the acceptance rate is too small; should be a vector of length k
#' @param accept_range Target range of acceptance rates
#' @param scale Optional initial guess for the scale vector, default 1
#' @param scale_step Optional, how large the adaptation steps are, default 1.1
#' @return Scale parameter after adaptation
m_adapt = function(target, state, data, chunk_size = 100, max_chunk = 100,
			accept_range = c(0.2, 0.4), scale, scale_step) {
	k = length(state)
	accept_rate = rep(0, k)
	if(missing(scale))
		scale = rep(1, k)
	if(missing(scale_step)) {
		scale_step = rep(1.1, k)
	} else if(any(scale_step <= 1)) {
		stop("Scale step must be greater than one")
	}
	nchunk = 0
	while(nchunk < max_chunk & (any(accept_rate < min(accept_range)) |
							any(accept_rate > max(accept_range)))) {
		chunk = m_sample(target, state, data, scale, iter = chunk_size)
		accept_rate = chunk$accept
		state = chunk$chain[nrow(chunk$chain),]

		for(j in 1:k) {
			if(accept_rate[j] < min(accept_range)) {
				scale[j] = scale[j] / scale_step[j]
			} else if(accept_rate[j] > min(accept_range)) {
				scale[j] = scale[j] * scale_step[j]
			}
		}
		nchunk = nchunk + 1
	}
	cat("Finished adaptation after", nchunk, "chunks of", chunk_size,
		"iterations; acceptance rate:\n", "   ", accept_rate, "\n")
	return(scale)
}

#' Propose and select candidate values for the metropolis-hastings algorithm
#' @param state Current state of the chain
#' @param scale Scale of the proposal distribution
#' @param target Target distribution density function
#' @param data Data for the target distribution
#' @param j Index of the state/scale vectors to propose, others will be held constant
#' @return New value in the chain
m_propose <- function(state, scale, target, data, j) {
	# choose a candidate value from the proposal distribution (rnorm)
	candidate = state
	candidate[j] = rnorm(1, state[j], scale[j])

	# compute the log posterior density at the current state and the candidate value
	ld_cand = target(candidate, data)
	ld_current = target(state, data)

	# compute the acceptance probability r, making sure the density is finite
	# r is the ratio of densities: (d_cand / d_current)
	# because r should be the exponential of the difference of log densities
	if(!is.finite(ld_cand)) {
		r = 0
	} else {
		r = exp(ld_cand - ld_current)
	}

	# now choose a value such that the chance of choosing d_cand is equal to r
	# the easiest way is to draw a random uniform number U on (0,1) and compare it to r
	# pr(U < r) = r (acceptance), pr(U >= r) = 1 - r (rejection)
	# could also use sample() to do this
	U = runif(1)
	if(U < r) {
		val = candidate
	} else {
		val = state
	}
	return(val)
}

#' Computes a multivariate hpdi interval
#'
#' Assumes a unimodal posterior distribution
#' @param samples A matrix of mcmc samples
#' @param posterior A function returning the log unnormalized posterior density
#' @param data The data used to compute the samples
#' @param density The (minimum) probability density contained in the interval
hpdi = function(samples, posterior, data, density = 0.9) {

	# compute the indices of all possible intervals
	n = nrow(samples)
	n_include = ceiling(density * n)
	lower = 1:(n - n_include)
	upper = lower + n_include

	# sort the samples by their posterior density
	dens = apply(samples, 1, posterior, data = data)
	ind = order(dens)
	dens = dens[ind]
	samples = samples[ind, ]

	# compute the width for each variable, then the area in the parameter dimension
	area = mapply(function(l, u, samp)
		prod(apply(samp, 2, function(x) x[u[i]] - x[l[i]])),
		l = lower, u = upper, moreArgs = list(samp = samples))

	i = which.min(area)
	cbind(samples[lower[i], ], samples[upper[i], ])
}


#' ---
#' author: "Jenny Bryan"
#' output: github_document
#' ---

#+ setup, include = FALSE
library(knitr)
opts_chunk$set(fig.path = 'figure/uni-factor-', error = TRUE)

#' Note: this is rendered by applying `knitr::spin()` to an R script. So the
#' narrative is very minimal.

#' load the data and ggplot2 (part of the tidyverse)
library(tidyverse)
library(gapminder)
gapminder

#' bar charts
#' consider: no. of observations for each continent
table(gapminder$continent)

#' this works because default stat for geom_bar() is "bin"
ggplot(gapminder, aes(x = continent)) + geom_bar()

#' let's reorder the continents based on frequency
p <- ggplot(gapminder, aes(x = reorder(continent, continent, length)))
p + geom_bar()

#' would you rather the bars run horizontally?
p + geom_bar() + coord_flip()

#' how about a better data:ink ratio?
p + geom_bar(width = 0.05) + coord_flip()

#' consider a scenario where you DON'T want the default "bin" stat, i.e. the bar
#' length or height already exists as a variable
(continent_freq <- gapminder %>% count(continent))

#' this simple call no longer works, because we have pre-tabulated
ggplot(continent_freq, aes(x = continent)) + geom_bar()

#' THIS works when bar length or height already exists
ggplot(continent_freq, aes(x = continent, y = n)) + geom_bar(stat = "identity")

sessionInfo()


rm(list = ls())

# Initial Parameters given by the user.
args <- commandArgs(trailingOnly = TRUE)
dataMatrix <- args[1]
transposition <- as.logical(args[2])
Envdataset <- args[3]
datasetpath <- args[4]
method <- args[5]
index <- args[6]
upto <- as.numeric(args[7])
trace <- as.logical(args[8])
metric <- args[9]

library(vegan)
library(cluster)

setwd(datasetpath)

comm <- read.csv(dataMatrix, header = TRUE, row.names = 1, sep = ",")
env <- read.csv(Envdataset, header = TRUE, row.names = 1, sep = ",")

if(transposition) {
  dat <- t(comm)
} else if (tranposition != TRUE) {
  dat <- comm
} else {
  stop("Wrong choice. Choose to transpose the matrix (TRUE) or not (FALSE).")
}

if((upto <= 0) || (upto > ncol(env))) {
  stop("Wrong choice of upto parameter, transcends available upto choices.")
}

env <- env[1:nrow(comm), ]

partial = NULL
parallel = getOption("mc.cores")

# load libraries
library(pbdMPI, quiet = TRUE)

# Start the clock
fullProgramTimer <- proc.time()

init()

# TP: Get the size of the cluster
processors <- comm.size()
myrank <- comm.rank()

if (is.null(partial)) { # TRUE
  corfun <- function(dx, dy, dz, method) {
    cor(dx, dy, method = method)
  }
} else {
  corfun <- function(dx, dy, dz, method) {
    rxy <- cor(dx, dy, method = method)
    rxz <- cor(dx, dz, method = method)
    ryz <- cor(dy, dz, method = method)
    (rxy - rxz * ryz)/sqrt(1 - rxz * rxz)/sqrt(1 - ryz *
                                                 ryz)
  }
}

if (!is.null(partial)) { # FALSE
  partpart <- deparse(substitute(partial))
} else partpart <- NULL
if (!is.null(partial) && !inherits(partial, "dist")) # FALSE
  partial <- dist(partial)
if (!is.null(partial) && !pmatch(method, c("pearson", "spearman"),
                                 nomatch = FALSE)) # FALSE
  stop("method ", method, " invalid in partial bioenv")

# Main calculation start
mainCalcTimer <- proc.time()

 n <- ncol(env)
ntake <- 2^n - 1
ndone <- 0
if (n > 8 || trace) {
  if (upto < n)
    cat("Studying", nall <- sum(choose(n, 1:upto)), "of ")
  cat(ntake, "possible subsets (this may take time...)\n")
  flush.console()
}
if (metric == "euclidean") {
  x <- scale(env, scale = TRUE)
  distfun <- function(x) dist(x)
} else if (metric == "mahalanobis") {
  x <- as.matrix(scale(env, scale = FALSE))
  distfun <- function(x) dist(veganMahatrans(x))
} else if (metric == "gower") {
  x <- env
  distfun <- function(x) daisy(x, metric = "gower")
} else if (metric == "manhattan") {
  x <- decostand(env, "range")
  distfun <- function(x) dist(x, "manhattan")
} else {
  stop("unknown metric")
}

best <- list()

if (inherits(comm, "dist")) {
  comdis <- comm
  index <- attr(comdis, "method")
  if (is.null(index))
    index <- "unspecified"
} else if (is.matrix(comm) && nrow(comm) == ncol(comm) && isTRUE(all.equal(comm,
                                                                         t(comm)))) { # FALSE
  comdis <- as.dist(comm)
  index <- "supplied square matrix"
} else {
  comdis <- vegdist(comm, method = index)
}

if(processors >= upto) {
   processors <-  upto
}

step <- floor(upto/processors)
step2 <- upto - (processors - 1)*step

if (myrank != (processors - 1)) {
  for (i in (myrank*step + 1):((myrank + 1)*step)) {
    if (trace) { # FALSE
      nvar <- choose(n, i)
      cat("No. of variables ", i, ", No. of sets ", nvar,
          "...", sep = "")
      flush.console()
    }
    sets <- t(combn(1:n, i))
    if (!is.matrix(sets))
      sets <- as.matrix(t(sets))
    est <- numeric(nrow(sets))
    for (j in 1:nrow(sets)) est[j] <- corfun(comdis, dist(x[,
                                                            sets[j, ]]), partial, method = method)

#    best[[i]] <- list(best = sets[which.max(est), ], est = max(est))

    best[[i - myrank*step]] <- list(best = sets[which.max(est), ], est = max(est))

#    best[is.na(best)] <- 0

    if (trace) {
      ndone <- ndone + nvar
      cat(" done (", round(100 * ndone/ntake, 1), "%)\n",
          sep = "")
      flush.console()
    }
  }
} else {
  for (i in (upto - step2 + 1):upto) {
    if (trace) { # FALSE
      nvar <- choose(n, i)
      cat("No. of variables ", i, ", No. of sets ", nvar,
          "...", sep = "")
      flush.console()
    }
    sets <- t(combn(1:n, i))
    if (!is.matrix(sets))
    sets <- as.matrix(t(sets))
    est <- numeric(nrow(sets))
    for (j in 1:nrow(sets)) est[j] <- corfun(comdis, dist(x[,
                                                          sets[j, ]]), partial, method = method)
#    best[[i]] <- list(best = sets[which.max(est), ], est = max(est))

    best[[i - upto + step2]] <- list(best = sets[which.max(est), ], est = max(est))

#    best[is.na(best)] <- 0

    if (trace) {
      ndone <- ndone + nvar
      cat(" done (", round(100 * ndone/ntake, 1), "%)\n",
          sep = "")
      flush.console()
    }
  }
}

best.all <- gather(best)
best.este <- list()

if(myrank == 0){
  best.este <- list()
  j = length(best.all[[1]])
  for(i in 1:j) {
    best.este[[i]] <- best.all[[1]][[i]]
  }
for(k in 2:length(best.all)) {
    j <- length(best.all[[k]])
    for(i in 1:j){
      best.este[[i]] <- best.all[[k]][[i]]
    }
  }
}

whichbest <- which.max(lapply(best.este, function(tmp) tmp$est))
out <- list(names = colnames(env), method = method, index = index,
            metric = metric, upto = upto, models = best.este, whichbest = whichbest,
            partial = partpart, x = x, distfun = distfun)

out$call <- as.call(list(quote(bioenv), quote(comm), quote(env),
                   method = method, index = index, metric = metric,
                   upto = upto, trace = trace, partial = partial))

out$call[[1]] <- as.name("bioenv")
class(out) <- "bioenv"
comm.print(out)

# Stop the clock
comm.print("Full program execution time for each processor", rank.source = 0)
comm.print(proc.time() - fullProgramTimer, all.rank = T)

comm.print("Full program execution time for each processor", rank.source = 0)
comm.print(proc.time() - mainCalcTimer, all.rank = T)

comm.print("Names:", rank.source = 0)
comm.print(out[[1]], rank.source = T)

comm.print("Method:", rank.source = 0)
comm.print(out[[2]], rank.source = T)

comm.print("Index", rank.source = 0)
comm.print(out[[3]], rank.source = T)

comm.print("Metric", rank.source = 0)
comm.print(out[[4]], rank.source = T)

comm.print("upto:", rank.source = 0)
comm.print(out[[5]], rank.source = T)

comm.print("Model:", rank.source = 0)
comm.print(out[[6]], rank.source = T)

comm.print("whichbest:", rank.source = 0)
comm.print(out[[7]], rank.source = T)

comm.print("trace:", rank.source = 0)
comm.print(out[[8]], rank.source = T)

comm.print("x:", rank.source = 0)
comm.print(out[[9]], rank.source = T)

comm.print("Distance function:", rank.source = 0)
comm.print(out[[10]], rank.source = T)

comm.print("Call:", rank.source = 0)
comm.print(out[[11]], rank.source = 0)

comm.print("Final result's size", rank.source = 0)
comm.print((object.size(out)/(1024^2)), rank.source = T)

finalize()


args <- commandArgs(trailingOnly = TRUE)

# Call in data for plotting with argument line (regression same as Python exercise)
dataset <- read.csv(args[1])
dataset

y1 <- dataset[,1]
x1 <- dataset[,2]

# Plot data points, opening png device.
# This method for plotting the scatter plot calls up each x,y coordinate for plotting.
png('R_Scatter Data.png')
plot(y1,
     x1,
     ylab='Y',
     xlab='X',
     pch = 16,
     col = 'black',
     main = 'Scatter Data')
dev.off() # closing png device

# Create linear regression model, plot points, and add regression
# Opening png device
# This method for plotting the scatter plot and adding the regression is more straight forward because it assumes the data points are in x,y format.
png('R_Regression Data.png')
lmRegrex = lm(x~y, data = dataset)
plot(dataset,
    pch = 16,
    col = 'black',
    main = 'Regression Data')
abline(lmRegrex)
dev.off()  # closing png device


## --------------------
## --------------------
## --------------------
## This file contains some basic functions for implementing Adaptive Penalized Tensor Decomposition.
## Most function are only suitable for three dimension tensor. For higher-order tensor, codes should be modified adaptively.
## --------------------
## --------------------
## --------------------

## --------------------
## inital values
## k @int
## --------------------
inits <- function(k) {
    return(t(as.matrix(rnorm(k, 0, 1))))
}

## --------------------
## tensor product
## --------------------
product <- function(X, ind, u, v, w) {
    u <- t(as.matrix(u))
    v <- t(as.matrix(v))
    w <- t(as.matrix(w))

    if (length(ind) == 2) {
        lizt <- list("mat2" = u, "mat3" = v)
        u <- ttl(X, lizt, ms = ind)
        return(as.vector(u@data))
    }

    lizt <- list("mat" = u, "mat2" = v, "mat3" = w)
    u <- ttl(X, lizt, ms = ind)
    return(as.vector(u@data))
}

## --------------------
## generate inital value for tensor decompostion if initial value is null.
## --------------------
gen_startvals <- function(num_factors, tnsr) {
    u <- matrix(inits(dim(tnsr)[1] * num_factors), nrow = num_factors)
    v <- matrix(inits(dim(tnsr)[2] * num_factors), nrow = num_factors)
    w <- matrix(inits(dim(tnsr)[3] * num_factors), nrow = num_factors)
    for (i in 1:num_factors) {
        u[i, ] <- u[i, ] / norm(as.matrix(u[i, ]), "F")
    }
    for (i in 1:num_factors) {
        v[i, ] <- v[i, ] / norm(as.matrix(v[i, ]), "F")
    }
    for (i in 1:num_factors) {
        w[i, ] <- w[i, ] / norm(as.matrix(w[i, ]), "F")
    }

    return(list(u = u, v = v, w = w))
}

## --------------------
## generate mortality tensor.
## --------------------
generate_mortality_tensor <- function(x, log = TRUE, std = TRUE) {
    x <- as.data.table(x)
    len <- ncol(x)
    circu_m <- data.matrix(x[Cause == "Circulatory system"][, (len - 18):len])
    c_m <- data.matrix(x[Cause == "Cancer"][, (len - 18):len])
    r_m <- data.matrix(x[Cause == "Respiratory system"][, (len - 18):len])
    e_m <- data.matrix(x[Cause == "External causes"][, (len - 18):len])
    i_m <- data.matrix(x[Cause == "Infectious and parasitic diseases"][, (len - 18):len])
    o_m <- data.matrix(x[Cause == "Others"][, (len - 18):len])
    ## define the tensor and apply log transformation
    if (log) {
        m_tensor <- array(NA, dim = c(6, 19, (nrow(x) / 6)))
        m_tensor[1, , ] <- t(log(circu_m))
        m_tensor[2, , ] <- t(log(c_m))
        m_tensor[3, , ] <- t(log(r_m))
        m_tensor[4, , ] <- t(log(e_m))
        m_tensor[5, , ] <- t(log(i_m))
        m_tensor[6, , ] <- t(log(o_m))
    } else {
        m_tensor <- array(NA, dim = c(6, 19, (nrow(x) / 6)))
        m_tensor[1, , ] <- t(circu_m)
        m_tensor[2, , ] <- t(c_m)
        m_tensor[3, , ] <- t(r_m)
        m_tensor[4, , ] <- t(e_m)
        m_tensor[5, , ] <- t(i_m)
        m_tensor[6, , ] <- t(o_m)
    }
    ###
    if (std) {
        tensor1 <- (m_tensor - array(mean(m_tensor), dim = dim(m_tensor))) / sd(m_tensor)
        return(as.tensor(tensor1))
    } else {
        tensor1 <- m_tensor
        return(as.tensor(tensor1))
    }
}


#####################  ANALYSIS OF GLOBAL PREFERENCES  #########################

# =============================== #
#### 0. LOAD DATA AND SET PATH ####
# =============================== #

# Set the path
setwd("~/Desktop/Projects/Global-Preferences-Survey/ReproductionAnalysis/")

# Source helper functions
source("functions/helper_functions/SourceFunctions.r")
SourceFunctions(path = "functions/")
SourceFunctions(path = "functions/helper_functions/")

# Load libraries
LoadRequiredLibraries()

# Load the data
data_all <- LoadData()

# ========================= #
#### 1. PREPARE THE DATA ####
# ========================= #

data_all <- PrepareData(data_all)

# Standardize preferences at country level
data_all$data <- Standardize(data    = data_all$data,
                             columns = c(5:10),
                             level   = "country")

# Use only the complete dataset
dataComplete <- data_all$data[complete.cases(data_all$data)]

# ========================== #
#### 2. CREATE THE MODELS ####
# ========================== #

# Model on country level of the preferences
models <- CreateModelsForPreferencesCountryLevel(dataComplete)

# Summarize the preferences for each country
dataCoeff <- SummaryCoeffPerPreferencePerCountry(models)

# Adjust data for plotting
dataCoeff[data_all$data, `:=` (isocode     = i.isocode,
                               logAvgGDPpc = log(i.avgGDPpc)),
          on = "country"]
setnames(dataCoeff, old = "gender1", new = "gender")

# ===================================== #
#### 3. PRINCIPAL COMPONENT ANALYSIS ####
# ===================================== #

# PCA on the preferences
summaryIndex <- AvgGenderDiffPreferencesPCA(dataCoeff)

# Prepare summary index
summaryIndex <- CreateSummaryIndex(summaryIndex, data_all)

# Perform the principal component analysis imputing missing values
summaryIndex$GenderIndex <- GenderIndexPCA(summaryIndex[, c(5:8)])

# Standardize the predictors (mean 0 and std dev 1)
summaryIndex <- Standardize(data    = summaryIndex,
                            columns = c(2, 4:8, 13),
                            newName = TRUE)
# Set the Gender Index on a scale between 0 and 1
summaryIndex[, GenderIndexRescaled := Rescale(GenderIndex)]

# Prepare summary histograms
dataSummary <- SummaryHistograms(dataCoeff, summaryIndex)

# ================================ #
#### 4. VARIABLES CONDITIONING  ####
# ================================ #

# TODO: The authors took the log GDP and the Gender Equality (Index, meaning
# the one already in a scale from 0 to 1) and standardise them. After this, they
# perform the conditional analysis. Still I miss a factor 10 in the y-axis...

# Add residuals to the summary index
summaryIndex <- AddResiduals(summaryIndex)

# Invert trend for two variables
summaryIndex[, `:=` (ValueUNStd = -1 * ValueUNStd,
                     DateStd    = -1 * DateStd)]

# =============================== #
#### 5. SUPPLEMENTARY MATERIAL ####
# =============================== #

colsToKeep_coeff <- c("gender", "preference", "country", "isocode", "logAvgGDPpc")

colsToKeep_summary <- c("GenderIndex",
                        "country")
dataCoeff_summary <- merge(dataCoeff[, ..colsToKeep_coeff],
                           summaryIndex[, ..colsToKeep_summary],
                           by = "country")

# Invert the trend of those preferences with opposite direction of the difference
dataCoeff_summary <- InvertPreference(dataCoeff_summary)

dataCoeff_summary <- AddResidualsSinglePreference(dataCoeff_summary)

# Use the original gender coefficient (not inverted) to calculate the mean for
# each preference, and the 95% confidence interval of the standard error
dataCoeff_summary[, meanGender := mean(genderOrig), by = "preference"]
dataCoeff_summary[, stdGender := 1.96 * sqrt(sd(genderOrig)^2 / uniqueN(country)),
                  by = "preference"]

# -------------- #

## Alternative model

modelsAlternative <- CreateAlternativeModels(dataComplete)
dataCoeffAlternative <- SummaryCoeffPerPreferencePerCountry(modelsAlternative)

# Adjust data for plotting
dataCoeffAlternative[data_all$data, `:=` (isocode     = i.isocode,
                                          logAvgGDPpc = log(i.avgGDPpc)),
                     on = "country"]
setnames(dataCoeffAlternative, old = "gender1", new = "gender")

# Invert the trend of those preferences with opposite direction of the difference
dataCoeffAlternative <- InvertPreference(dataCoeffAlternative)

# -------------- #

## Gender differences standardize at global level

dataStdGlobal <- Standardize(data = data_all$data,
                             columns = c(5:10))

# Use only the complete dataset
dataCompleteGlobal <- dataStdGlobal[complete.cases(dataStdGlobal)]

# Model on country level of the preferences
modelsGlobal <- CreateModelsForPreferencesCountryLevel(dataCompleteGlobal)

# Summarize the preferences for each country
dataCoeffGlobal <- SummaryCoeffPerPreferencePerCountry(modelsGlobal)

# Adjust data for plotting
dataCoeffGlobal[data_all$data, `:=` (isocode     = i.isocode,
                                     logAvgGDPpc = log(i.avgGDPpc)),
                on = "country"]
setnames(dataCoeffGlobal, old = "gender1", new = "gender")

# Invert the trend of those preferences with opposite direction of the difference
dataCoeffGlobal <- InvertPreference(dataCoeffGlobal)

# ================================ #
#### 6. MODEL EVALUATION  ####
# ================================ #

# Add the data from the article
dt_article <- fread("files/input/data_extracted_from_article.csv")

# Create a data table for a quick comparison
dt_compare <- dt_article[summaryIndex, .(avgDiffArticle = AverageGenderDifference,
                                         avgDiffOurs = i.avgGenderDiffRescaled,
                                         isocode = i.isocode), on = "isocode"]
dt_compare <- dt_compare[complete.cases(dt_compare)]

# Plot the results
PlotSummary(data = dt_compare, var1 = "avgDiffArticle", var2 = "avgDiffOurs",
            labs = c("Average Gender Differences from the Article",
                     "Average Gender Differences from our Model"),
            display = TRUE
            )

# ============================= #
#### 7. WRITE DATA SUMMARIES ####
# ============================= #

## ---------------------- Write csv data summaries -------------------------- ##
fwrite(dataSummary,
       file = "files/output/main_data_for_histograms.csv")
fwrite(summaryIndex,
       file = "files/output/main_data_aggregatedByCountry_preferencePCA_genderIndexPCA.csv")
fwrite(dataCoeff_summary,
       file = "files/output/supplementary_data_aggregatedByCountry_singlePreference_genderCoefficients.csv")
fwrite(dataCoeffGlobal[, c(1, 3, 7, 8, 9)],
       file = "files/output/supplementary_data_aggregatedByCountry_singlePreference_genderCoefficientsGlobal.csv")
fwrite(dataCoeffAlternative[, -2],
       file = "files/output/supplementary_data_aggregatedByCountry_singlePreference_genderCoefficients_alternativeModel.csv")
#------------------------------------------------------------------------------#


args <- commandArgs(trailingOnly = TRUE)
intab <- args[1]
outtab <- args[2]
data <- read.table(intab, sep="\t", header=TRUE, row.names=1)
#remove transcript IDs if present
if(any(grepl("transcript", colnames(data)))) {
   data=data[, -grep("transcript", colnames(data))]
}
a=min(data[data!=0])
dim(data)
#filter genes with all zero
data=data[rowSums(data[,-1])>0,]
dim(data)
#take log2 of value+min
ldata=log2(data + a)
#quantile
library(limma)
limmaRes = normalizeQuantiles(ldata)
limmaRes = round(limmaRes, digits=3)
write.table(limmaRes, file=outtab, row.names = TRUE ,col.names = TRUE,quote = FALSE, sep="\t")