library(heatmap3)
library(RColorBrewer)
library(grDevices)

args <- commandArgs(TRUE)
filename = as.character(args[1])
isCat = as.numeric(args[2])

# isCat = 0
# directory = "../resources/toR/"
# filename = "mushrooms-Scale(27x1)-Original"
# filename = paste(directory, filename, sep="")

input = paste(filename, ".rformat", sep="")

df = read.table(input, header=FALSE, row.names=NULL)

removeZeroVar <- function(dat) {
	# https://stackoverflow.com/questions/8805298/quickly-remove-zero-variance-variables-from-a-data-frame
    dat[sapply(dat, function(x) length(levels(factor(x)))>1)]
}

labelCode = as.numeric(as.character(df[1,1]))
rowSide = NULL
colSide = NULL

if(labelCode > 3){
	# Unrecognized code
	print("Unrecognized label code")
	print(labelCode)
	quit(save="no", status=1)
}

if(labelCode == 1 || labelCode == 3) {
	# Only need columns
	colLabels = as.character(df[1,seq(2, dim(df)[2])]) # first row is label
	nUniqueColLabels = dim(table(colLabels))[1]
	col2 <- brewer.pal(nUniqueColLabels, "Set3")
	colorIndicesCol = match(colLabels, levels(factor(colLabels)))
	colSide = sapply(colorIndicesCol, function(x) col2[x])
	if(labelCode == 1){
		rowSide = NULL
	}
	stopifnot(length(colSide) == dim(df)[2]-1)
}

if(labelCode == 2 || labelCode == 3) {
	# Only need rows
	rowLabels = as.character(df[seq(2, dim(df)[1]),1]) # first column is label
	nUniqueRowLabels = dim(table(rowLabels))[1]
	col1 <- brewer.pal(nUniqueRowLabels, "Set3")
	colorIndicesRow = match(rowLabels, levels(factor(rowLabels)))
	rowSide = sapply(colorIndicesRow, function(x) col1[x])
	if(labelCode == 2){
		colSide = NULL
	}
	stopifnot(length(rowSide) == dim(df)[1]-1)
}

mydata = df[-1,-1] # Take all but first column (labels)

m = data.matrix(mydata)
# png(filename=paste(filename,"-heatmap.png",sep=""), 5000,5000)

method = "complete"

if(isCat == 1) {
	palette = colorRampPalette(c("navy", "white","firebrick3"))(1024)
} else {
	palette = gray.colors(2, start = 1, end = 0, gamma = 2.2, alpha = NULL)
}

if(is.null(rowSide) && is.null(colSide)){
	timeTaken = system.time(hm <- heatmap3(m, distfun = function(x) dist(x, method = "manhattan", diag = FALSE, upper = FALSE, p = 2), scale="none", method=method, col=palette))
} else if(is.null(rowSide)){
	timeTaken = system.time(hm <- heatmap3(m, distfun = function(x) dist(x, method = "manhattan", diag = FALSE, upper = FALSE, p = 2), ColSideColors=colSide, scale="none", method=method,col=palette))
} else if(is.null(colSide)) {
	timeTaken = system.time(hm <- heatmap3(m, distfun = function(x) dist(x, method = "manhattan", diag = FALSE, upper = FALSE, p = 2), RowSideColors=rowSide, scale="none", method=method,col=palette))
} else {
	timeTaken = system.time(hm <- heatmap3(m, distfun = function(x) dist(x, method = "manhattan", diag = FALSE, upper = FALSE, p = 2), RowSideColors=rowSide, ColSideColors=colSide, scale="none", method=method, col=palette))
}

print(filename)
print(timeTaken)

write.csv(hm["rowInd"],paste(filename,"rowInd.order",sep=""), row.names=FALSE)
write.csv(hm["colInd"],paste(filename,"colInd.order",sep=""), row.names=FALSE)
dev.off()


library(ape)
library(coda)
library(phangorn)

my.gelman.preplot <- function (x, bin.width = bin.width, max.bins = max.bins, confidence = confidence,
    transform = transform, autoburnin = autoburnin)
{
    ## based on code in coda:::gelman.preplot
    ## this version avoids calculating the multivariate PSRF, which lead to errors for some bins
    x <- as.mcmc.list(x)
    if (niter(x) <= 50)
        stop("Less than 50 iterations in chain")
    nbin <- min(floor((niter(x) - 50)*thin(x)/bin.width), max.bins)
    binw <- floor((niter(x) - 50)/nbin)
    last.iter <- c(seq(from = start(x) + 50 * thin(x), by = binw *
        thin(x), length = nbin), end(x))
    shrink <- array(dim = c(nbin + 1, nvar(x), 2))
    dimnames(shrink) <- list(last.iter, varnames(x), c("median",
        paste(50 * (confidence + 1), "%", sep = "")))
    for (i in 1:(nbin + 1)) {
        shrink[i, , ] <- gelman.diag(window(x, end = last.iter[i]),
            confidence = confidence, transform = transform, autoburnin = autoburnin, multivariate=F)$psrf
    }
    all.na <- apply(is.na(shrink[, , 1, drop = FALSE]), 2, all)
    if (any(all.na)) {
        cat("\n******* Error: *******\n")
        cat("Cannot compute Gelman & Rubin's diagnostic for any chain \n")
        cat("segments for variables", varnames(x)[all.na], "\n")
        cat("This indicates convergence failure\n")
    }
    return(list(shrink = shrink, last.iter = last.iter))
}

my.gelman.plot <- function (x, bin.width = 10, max.bins = 50, confidence = 0.95,
    transform = FALSE, autoburnin = TRUE, auto.layout = TRUE,
    ask, col = 1:2, lty = 1:2, xlab = "last iteration in chain",
    ylab = "shrink factor", type = "l", ...)
{
    ## same as coda::gelman.plot but uses my.gelman.preplot
    if (missing(ask)) {
        ask <- if (is.R()) {
            dev.interactive()
        }
        else {
            interactive()
        }
    }
    x <- as.mcmc.list(x)
    oldpar <- NULL
    on.exit(par(oldpar))
    if (auto.layout)
        oldpar <- par(mfrow = coda:::set.mfrow(Nchains = nchain(x),
            Nparms = nvar(x)))
    y <- my.gelman.preplot(x, bin.width = bin.width, max.bins = max.bins,
        confidence = confidence, transform = transform, autoburnin = autoburnin)
    all.na <- apply(is.na(y$shrink[, , 1, drop = FALSE]), 2,
        all)
    if (!any(all.na))
        for (j in 1:nvar(x)) {
            matplot(y$last.iter, y$shrink[, j, ], col = col,
                lty = lty, xlab = xlab, ylab = ylab, type = type,
                ...)
            abline(h = 1)
            ymax <- max(c(1, y$shrink[, j, ]), na.rm = TRUE)
            leg <- dimnames(y$shrink)[[3]]
            xmax <- max(y$last.iter)
            legend(xmax, ymax, legend = leg, lty = lty, bty = "n",
                col = col, xjust = 1, yjust = 1)
            title(main = varnames(x)[j])
            if (j == 1)
                oldpar <- c(oldpar, par(ask = ask))
        }
    return(invisible(y))
}

burnin <- 0.5
runstems <- paste0('beast/run', 1:4, '/pedv')

files <- paste(runstems, '.log', sep='')
runs <- lapply(files, read.table, header=TRUE, sep="\t")

tmpf <- function(x) {
    diff(x$state[1:2])
}
thin <- sapply(runs, tmpf)
all.elements.equal <- function(x) all.equal(x, rep(x[1], length.out=length(x)))
stopifnot(isTRUE(all.elements.equal(thin)))
thin <- thin[1]

tmpf <- function(x) {
    x[, 'state'] <- NULL
#    x[, 'clock.rate'] <- NULL
#    x[, 'posterior'] <- NULL
#    x[, 'coalescent'] <- NULL
#    x[, 'treeLikelihood'] <- NULL
#    x[, 'prior'] <- NULL
    x
}
runsp <- lapply(runs, tmpf)

mobs <- lapply(runsp, mcmc, thin=thin)
ml <- mcmc.list(mobs)
wstart <- burnin*end(ml)
wml <- window(ml, start=wstart)

plot(wml)

(essl <- lapply(wml, effectiveSize))
(ess <- effectiveSize(wml))

heidel.diag(wml)

(gd <- gelman.diag(wml, multivariate=FALSE))
gelman.diag(wml, transform=TRUE, multivariate=FALSE)

#' ## Gelman plots
my.gelman.plot(wml, ylim=c(1, 1.1))

#' ## Geweke plots
geweke.plot(wml)

#' ##Now for the trees

phymcmc <- function(x, start=1, end=numeric(0), thin=1){
    ## extension of mcmc to multiPhylo objects, based on code in coda
    if (!inherits(x, 'multiPhylo')) {
        stop("Not an object of class multiPhylo")
    } else {
        niter <- length(x)
        nms <- names(x)[1:2]
        s12 <- as.integer(sapply(strsplit(nms, split='_'), '[[', 2))
        if(missing(thin))
            thin <- diff(s12)
        if(missing(start))
            start=s12[1]
        if(missing(end))
            end <- start + thin * (niter - 1)
        nobs <- floor((end - start)/thin + 1)
        if (niter != nobs)
            stop("Start, end and thin incompatible with data")
        attr(x, 'mcpar') <- c(start, end, thin)
        attr(x, 'class') <- c('phymcmc', 'multiPhylo')
        x
    }
}

start.phymcmc <- function (x, ...)
    {
        mcpar(x)[1]
    }

end.phymcmc <- function (x, ...)
    {
        mcpar(x)[2]
    }

frequency.phymcmc <- function (x, ...)
    {
        1/thin(x)
    }

thin.phymcmc <- function (x, ...)
    {
        mcpar(x)[3]
    }

time.phymcmc <- function (x, ...)
    {
        ts(seq(from = start(x), to = end(x), by = thin(x)), start = start(x),
           end = end(x), deltat = thin(x))
    }

window.phymcmc <- function (x, start, end, thin, ...)
{
    ts.eps <- getOption("ts.eps")
    xmcpar <- mcpar(x)
    xstart <- xmcpar[1]
    xend <- xmcpar[2]
    xthin <- xmcpar[3]
    if (missing(thin))
        thin <- xthin
    else if (thin%%xthin != 0) {
        thin <- xthin
        warning("Thin value not changed")
    }
    xtime <- as.vector(time(x))
    if (missing(start))
        start <- xstart
    else if (length(start) != 1)
        stop("bad value for start")
    else if (start < xstart) {
        start <- xstart
        warning("start value not changed")
    }
    if (missing(end))
        end <- xend
    else if (length(end) != 1)
        stop("bad value for end")
    else if (end > xend) {
        end <- xend
        warning("end value not changed")
    }
    if (start > end)
        stop("start cannot be after end")
    if (all(abs(xtime - start) > abs(start) * ts.eps)) {
        start <- xtime[(xtime > start) & ((start + xthin) > xtime)]
    }
    if (all(abs(end - xtime) > abs(end) * ts.eps)) {
        end <- xtime[(xtime < end) & ((end - xthin) < xtime)]
    }
    use <- 1:length(x)
    use <- use[use >= trunc((start - xstart)/xthin + 1.5) & use <=
        trunc((end - xstart)/xthin + 1.5) & (use - trunc((start -
        xstart)/xthin + 1.5))%%(thin%/%xthin) == 0]
    y <- x[use]
    return(phymcmc(y, start = start, end = end, thin = thin))
}

print.phymcmc <- function (x, ...)
    {
        x.orig <- x
        cat("Markov Chain Monte Carlo (MCMC) output:\nStart =", start(x),
            "\nEnd =", end(x), "\nThinning interval =", thin(x), "\n")
        attr(x, "mcpar") <- NULL
        attr(x, "class") <- 'multiPhylo'
        NextMethod("print", ...)
        invisible(x.orig)
    }

is.phymcmc <- function(x) inherits(x, 'phymcmc')

phymcmc.list <- function (...)
    {
        x <- list(...)
        if (length(x) == 1 && is.list(x[[1]]))
            x <- x[[1]]
        if (!all(unlist(lapply(x, is.phymcmc))))
            stop("Arguments must be phymcmc objects")
        nargs <- length(x)
        if (nargs >= 2) {
            xmcpar <- lapply(x, mcpar)
            if (!all(unlist(lapply(xmcpar, "==", xmcpar[[1]]))))
                stop("Different start, end or thin values in each chain")
            xtlabels <- lapply(x, attr, 'TipLabel')
            if (!all(unlist(lapply(xtlabels, "==", xtlabels[[1]]))))
                stop("Different taxa among chains")
        }
        class(x) <- "phymcmc.list"
        return(x)
    }

start.phymcmc.list <- function (x, ...)
    {
        start(x[[1]])
    }

end.phymcmc.list <- function (x, ...)
    {
        end(x[[1]])
    }

thin.phymcmc.list <- function (x, ...)
    {
        thin(x[[1]])
    }

window.phymcmc.list <- function (x, ...)
    {
        structure(lapply(x, window.phymcmc, ...), class = "phymcmc.list")
    }

getSDSF <- function(x, minCladeFz=0.1){
    niter <- length(x[[1]])
    prt <- lapply(x, prop.part)
    labs <- lapply(prt, attr, 'labels')
    allSame <- function(v) all(sapply( v[-1], FUN=function(z) {identical(z, v[[1]])}))
    stopifnot(allSame(labs))
    tmpf <- function(x) {
        count <- attr(x, 'number')
        tmpf <- function(x){
            stopifnot(!is.unsorted(x, strictly=TRUE))
            paste(x, collapse=',')
        }
        clade <- sapply(x, tmpf)
        data.frame(clade=clade, count=count, stringsAsFactors=FALSE)
    }
    foo <- lapply(prt, tmpf)
    tmpf <- function(x, y) {
        n <- ncol(x)
        suff <- as.character(c(n-1, n))
        merge(x, y, by='clade', all=TRUE, suffixes=suff)
    }
    mg <- Reduce(f=tmpf, x=foo)
    cols <- grep('count', colnames(mg))
    counts <- mg[, cols]
    rownames(counts) <- mg$clade
    counts[is.na(counts)] <- 0
    fz <- counts/niter
    test <- apply(fz, 1, function(x) max(x) > minCladeFz)
    fz <- fz[test,]
    sdcf <- apply(fz, 1, sd)
    sdcf
}

asdsf.preplot <- function(x, nbin){
    niter <- length(x[[1]])
    if(niter<=50){
        stop("Less than 50 iterations in chain")
    }
    binw <- floor((niter - 50)/nbin)
    last.iter <- c(seq(from=start(x) + 50*thin(x), by=binw * thin(x) , length=nbin), end(x))
    tmpf <- function(y) {
        z <- window(x, end=y)
        getSDSF(z)
    }
    sdsf <- lapply(last.iter, tmpf)
    return(list(sdsfl=sdsf, last.iter=last.iter))
}

asdsf.plot <- function(...){
    ans <- asdsf.preplot(...)
    means <- sapply(ans$sdsfl, mean)
    plot(c(ans$last.iter[1], ans$last.iter), c(0, means), type='n',
         xlab='Last iteration', ylab='ASDCF')
    lines(ans$last.iter, means)
    abline(h=0, col='grey')
    abline(h=0.01, col='blue', lty=2)
    return(invisible(ans))
}

## Now to actually make use of all those functions

tnames <- c('nonsIndel-aligned.fasta-gb', 'sIndel-aligned.fasta-gb')
tmpf <- function(x) {
    paste(runstems, x,  'trees', sep='.')
}
tfiles <- lapply(tnames, tmpf)
tmpf <- function(x) {
    lapply(x, read.nexus)
}
truns <- lapply(tfiles, tmpf)
tmpf <- function(x) {
    lapply(x, phymcmc)
}
tmobs <- lapply(truns, tmpf)
tml <- lapply(tmobs, phymcmc.list)

tmpf <- function(x) {
    burnin*end(x)
}
wstart <- lapply(tml, tmpf)

tmpf <- function(x, y) {
    window(x, start=y)
}
wtml <- mapply(tmpf, tml, wstart, SIMPLIFY=FALSE)

nbin <- 8
tmpf <- function(x) {
    asdsf.plot(x, nbin=nbin)
}
ans <- lapply(wtml, tmpf)

tmpf <- function(x) {
    hist(log10(10e-8 + x$sdsfl[[nbin]]))
}
lapply(ans, tmpf)

tmpf <- function(x, y) {
    file <- paste0(y, '.combined.trees')
    tsamp <- Reduce('c', x=x)
    class(tsamp) <- 'multiPhylo'
    attr(tsamp, 'mcpar') <- NULL
    write.nexus(tsamp, file=file)
}
mapply(tmpf, wtml, tnames)

tmpff <- function(xx, stem, thin=50000){
    tmpf <- function(x) {
        x <- window(x, thin=thin)
        class(x) <- 'multiPhylo'
        x
    }
    xx <- lapply(xx, tmpf)
    con <- consensus(xx[[1]], p=0.5)
    file <- paste0('densiTree.', stem, '.pdf')
    pdf(file=file)
    lapply(xx, densiTree, consensus=con, alpha=0.1)
    dev.off()
}
mapply(tmpff, xx=wtml, stem=tnames)

tmpf <- function(x, y, mar, text) {
    file <- paste0(y, '.plotted.trees')
    tsamp <- window(x, thin=thin(x)*50)
    tsamp <- Reduce('c', x=tsamp)
    class(tsamp) <- 'multiPhylo'
    attr(tsamp, 'mcpar') <- NULL
    tl <- attr(tsamp, 'TipLabel')
    tl <- gsub("_.*_", " ", tl)
    n <- max(nchar(tl))
    if(n > maxLabLen) maxLabLen <<- n
    tmpff <- function(xx) {
        while(nchar(xx) < maxLabLen) xx <- paste0(xx, ' ')
        xx
    }
    attr(tsamp, 'TipLabel') <- sapply(tl, tmpff)
    con <- consensus(tsamp, p=0.5)
    par(mar=mar)
    par(xpd=NA)
    par(yaxs='i') ## avoid extending y-range by 4%
    densiTree(tsamp, consensus=con, alpha=0.01)
    mtext(text, side=2, las=1, at=length(tl), cex=1.5)
    write.nexus(tsamp, file=file)
}

maxLabLen <- 0
res <- 400
png('trees.png', pointsize=8, height=8.75*res, width=3.25*res, res=res)
layout(matrix(c(1,2,3), ncol=1), heights=c(7,1,0.25))
mar <- list(c(1,2,0,0.1), c(1,2,2,0.1))
text <- list('A', 'B')
mapply(tmpf, wtml, tnames, mar, text)
mtext('Years before Februrary 2014', side=1, line=2.75, at=0.5)
dev.off()

save.image('convergence-check.RData')


# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set.seed(4)

## Pad with zeros on both ends
pad = function(x, pad.size=500) {
    c(rep(0,pad.size), x, rep(0,pad.size))
}

## Create a randomized impulse by creating white noise, down-filtering and windowing
impulse = function(n=32, f=n/2, pad.size=500) {
    t = 1:(2*n) - n - 1 + 1e-8
    v = sin(t/f * pi)/t * cos(t/2/n*pi)^2
    n = pad(rnorm(length(t)), n * 3)
    convolve(n,v, type="filter")
}

## Encode a windowed time-series using a specified dictionary
encode = function(win, dict) {
    ## largest dot product is most aligned centroid, size gives window scale
    dot = win %*% dict
    ## return nearest centroid in column 1, scale in column 2
    t(apply(dot, 1, function(v) {
        mx = max(abs(v))
        k = which(abs(v)==mx)[1]
        c(k, v[k])
    }))
}

## Reconstructs a signal from an encoded form
decode = function(data, dict) {
    n = dim(data)[1]
    r = rep(0, (n+1)*16)
    ## step a window at a time
    for (i in 1:n) {
        k = data[i,1]
        v = data[i,2]
        off = (i-1)*16
        ## add in current centroid after scaling and shifting it
        r[off + 1:32] = r[off + 1:32] + dict[,k] * v
    }
    r
}

## Computes the reconstruction error by reconstructing and measuring
reconstruction.error = function(data, dict) {
    rx = window(data)
    kx = encode(rx, dict)
    rz = decode(kx, dict)
    n = length(rz)
    mean(abs(data[15 + 1:n] - rz))
}

## Generates a randomized time series
generate = function(windows, v.1, v.2) {
    ## signal is original all zeros
    r = rep(0, 16 * windows)
    t = 0
    while (t < length(r) - length(v.1)) {
        ## next pulse start point is exponentially distributed
        t = t + -0.5 * length(v.1) * log(runif(1))
        ## flip a coin to pick which pulse we want
        flip = rnorm(1) > 0
        n1 = floor(t)
        n2 = n1 + length(v.1) - 1
        if (n2 > length(r)) {
            return(r)
        }
        r[n1:n2] = r[n1:n2] + v.1*flip + v.2*(1-flip)
    }
    r
}

## Vector magnitude
mag = function(v) {
    sqrt(sum(v^2))
}

## Stack up a time series as overlapping window segments
window = function(data, window.size=32) {
    ## how many windows?
    n = floor(2 * length(data)/window.size)
    ## the window shape
    w = sin((0:(window.size-1))/window.size * pi)^2
    ## for each row index, apply the window function to a sub-sequence
    t(apply(matrix(1:(n-window.size/2), ncol=1), 1, function(i) {
        n = i*16
        w * data[n:(n+window.size-1)]
    }))
}

## Generate the data
v.1 = impulse(64, 8)
v.2 = impulse(64, 8)

training.data = generate(10000, v.1, v.2)
test.data = generate(10000, v.1, v.2)

## Go ahead and window once to help with all the clustering
rx = t(apply(window(training.data), 1, function(v){v/(1e-100+mag(v))}))

k.10 = kmeans(rx, centers=10, nstart=10)
c.10 = apply(k.10$centers,1,function(v){v/mag(v)})
x.10 = reconstruction.error(training.data, c.10)
y.10 = reconstruction.error(test.data, c.10)
print(10)

k.100 = kmeans(rx, centers=100, nstart=10, iter.max=20)
c.100 = apply(k.100$centers,1,function(v){v/(1e-100 + mag(v))})
x.100 = reconstruction.error(training.data, c.100)
y.100 = reconstruction.error(test.data, c.100)
print(100)

k.200 = kmeans(rx, centers=200, nstart=10, iter.max=30)
c.200 = apply(k.200$centers,1,function(v){v/(1e-100 + mag(v))})
x.200 = reconstruction.error(training.data, c.200)
y.200 = reconstruction.error(test.data, c.200)
print(200)

k.500 = kmeans(rx, centers=500, nstart=10, iter.max=30)
c.500 = apply(k.500$centers,1,function(v){v/(1e-100 + mag(v))})
x.500 = reconstruction.error(training.data, c.500)
y.500 = reconstruction.error(test.data, c.500)
print(500)

k.1000 = kmeans(rx, centers=1000, nstart=5, iter.max=10)
c.1000 = apply(k.1000$centers,1,function(v){v/(1e-100 + mag(v))})
x.1000 = reconstruction.error(training.data, c.1000)
y.1000 = reconstruction.error(test.data, c.1000)
print(1000)

k.2000 = kmeans(rx, centers=2000, nstart=5, iter.max=10)
c.2000 = apply(k.2000$centers,1,function(v){v/(1e-100 + mag(v))})
x.2000 = reconstruction.error(training.data, c.2000)
y.2000 = reconstruction.error(test.data, c.2000)
print(2000)

reconstruct.plot = function () {
    plot(test.data[1:2000], type='l', xlab="Time", ylab=NA, main="Time series training data (first 2000 samples)",
         lwd=3)
    rex = window(test.data)
    re.data = decode(encode(rex, c.1000), c.1000)
    lines(15 + 1:2000, re.data[1:2000], col='red')
    lines(15 + 1:2000, test.data[15+1:2000]-re.data[1:2000]-1, col='blue')
    legend(800, 2.2, legend=c("Test data", "Reconstruction", "Error"), col=c("black", "red", "blue"), pch=21, lty=1, cex=0.8)
}

time.series.error.plot = function() {
    plot(c(10,100,200,500,1000,2000), c(x.10, x.100, x.200, x.500, x.1000, x.2000),
         type='b', lwd=2, xlab="Centroids", ylab="MAV Error",
         ylim=c(0,0.15),
         main="Reconstruction error for time-series data")

    lines(c(10,100,200,500,1000,2000), c(y.10, y.100, y.200, y.500, y.1000, y.2000),
          type='b', lwd=2, col='red')

    legend(1000, 0.15, legend=c("Training data", "Held-out data"), col=c("black", "red"), pch=21, lwd=2)
}

time.series.plot = function() {
    plot(training.data[1:2000], type='l', xlab="Time", ylab=NA, main="Time series training data (first 2000 samples)")
}

pdf(file="time-series-errors.pdf", width=5, height=5)
time.series.error.plot()
dev.off()

png(file="images/time-series-errors.png", width=400, height=400)
time.series.error.plot()
dev.off()

pdf(file="time-series.pdf", width=5, height=5)
time.series.plot()
dev.off()

png(file="images/time-series.png", width=400, height=400)
time.series.plot()
dev.off()

pdf(file="time-series-reconstruction.pdf", width=5, height=5)
reconstruct.plot()
dev.off()

png(file="images/time-series-reconstruction.png", width=400, height=400)
reconstruct.plot()
dev.off()


#' scaler
#'
#' A function to center (subtract mean) and/or scale (divide by
#' standard deviation) data column-wise in a computationally
#' efficient way.
#'
#' @details
#' Unlike its R counterpart \code{scale()}, the arguments
#' \code{center} and \code{scale} can only be logical values
#' (and not vectors).
#'
#' @param x
#' The input matrix.
#' @param center,scale
#' Logical; determine if the data should be centered and/or scaled.
#'
#' @return
#' The centered/scaled data, with attributes as in R's \code{scale()}.
#'
#' @export
scaler <- function(x, center=TRUE, scale=TRUE)
{
  check.is.flag(center)
  check.is.flag(scale)
  if (!is.numeric(x))
    stop("argument 'x' must be numeric")

  if (!is.double(x))
    storage.mode(x) <- "double"

  if (!center && !scale)
    return(x)
  else
    .Call(R_scaler, center, scale, x)
}


birth_death_fishing = function( selection="stan_code", fit=NULL, vn=NULL, sppoly=NULL, poly_match=NULL, time_match=NULL, catches=NULL, wgts=NULL, ... ) {

  message(" ")
  message("todo :: use carstm estimated numbers as IOA")
  message("todo :: add zero-inflation")
  message(" ")

  if (grepl("parameters", selection) ) {

    return(params)
  }

  if (grepl("stan_code", selection) ) {

    # simple discrete ode form, spatially explicit (au), birth/death separated
    # n[au,t] = n [au,t−1] + b[au] n[au,t−1] (1 − b[au,t−1]  ) − c [au,t−1]

    out = "
      data {
        int<lower=0> N; // no. years
        int<lower=0> U; // no. regions
        int CAT[U,N];
        int IOA[U,N];
      }

      transformed data {
        real eps = 1e-9;
        real IOAmin[U];
        real IOAmax[U];
        real IOAmed[U];
        real IOAerror[U];
        for (u in 1:U) {
          IOAmax[u] = 1e9;
          IOAmin[u] = 1e-9;
          for (n in 1:N) {
            real test = IOA[u,n] + CAT[u,n] ;
            if (test > IOAmax[u]) IOAmax[u] = test;
            if (test <= IOAmin[u]) IOAmin[u] = test;
          }
          IOAmed[u] = (IOAmax[u] + IOAmin[u])/2.0 ;
          IOAerror[u] = IOAmax[u] * 0.05 ; // 5% error..  estim of sd
        }
      }

      parameters {
        real <lower=eps, upper=1.0> X[U,N];
        real <lower=0.0, upper=10 > g[U,N];
        real <lower=0.0, upper=0.5> m[U];
        real <lower=1.0> K[U];
        real <lower=0.0, upper=0.5>  psd[U];  // process error
        real <lower=0.0, upper=1.0> res[U,N];
        real <lower=0.0, upper=0.5>  ressd[U];  /////
        real <lower=0.0, upper=1.0>  gsd[U]; /////
        real <lower=0.5, upper=1.5>   q[U];
        real <lower=0.0, upper=0.5>   msd; ////
        real <lower=0.0, upper=0.2>   qsd;
      }

      model {
        psd ~ normal(0.0, 0.2) ;
        qsd ~ normal(0.0, 0.05);  // ranges from 0.5 to 1.5
        ressd ~ normal(0.0, 0.05) ;
        msd ~ normal(0.0, 0.05) ;
        gsd ~ normal(0.0, 0.05) ; // 0.1
        m ~ normal(0.0, msd);
        q ~ normal(1.0, qsd );

        for (u in 1:U) {
          g[u,] ~ normal( 0.0, gsd[u] );
          K[u] ~ normal( IOAmax[u], IOAerror[u] ) ;
          X[u,1] ~ normal( 0.5, psd[u] ) ;
          for (n in 1:N) {
            res[u,n] ~ normal( q[u] * (X[u,n] - CAT[u,n]/K[u] ), ressd[u] );  // no catch observation error  .. keep separate to force positive value
            IOA[u,n] ~ poisson( K[u] * res[u,n] ); // survey index observation model
            if (n < N) {
              X[u,n+1] ~ normal( X[u,n] * (1.0 + g[u,n] - m[u]*X[u,n] ) - CAT[u,n]/K[u], psd[u] ) ; //process model
            }
          }
        }
      }
    "

    if (grepl("stan_code_testing_hurdle", selection) ) {

      # simple discrete ode form, spatially explicit (au), birth/death separated
      # n[au,t] = n [au,t−1] + b[au] n[au,t−1] (1 − b[au,t−1]  ) − c [au,t−1]

      out = "
        functions {
          int num_zero(int[] y) { // count no of zero-values
            int nz = 0;
            for (n in 1:size(y))
              if (y[n] == 0)
                nz += 1;
            return nz;
          }
        }

        data {
          int<lower=0> N; // no. years
          int<lower=0> U; // no. regions
          int CAT[U,N];
          int IOA[U,N];
        }

        transformed data {
          real eps = 1e-9;
          real IOAmin[U];
          real IOAmax[U];
          real IOAmed[U];
          real IOAerror[U];
          int<lower=0 > N0[U]; // no of zero-values in each au
          int<lower=0 > Ngt0[U];  // no of non-zero values in each au
          int<lower=1> IOA_nz[U], N - num_zero(y);  // values of IOA > 0
          for (u in 1:U) {
            N0[u] = num_zero( IOA[u,] );
            Ngt0[u] = N - N0[u];
            IOAmax[u] = 1e9;
            IOAmin[u] = 1e-9;
            for (n in 1:N) {
              real test = IOA[u,n] + CAT[u,n] ;
              if (test > IOAmax[u]) IOAmax[u] = test;
              if (test <= IOAmin[u]) IOAmin[u] = test;
            }
            IOAmed[u] = (IOAmax[u] + IOAmin[u])/2.0 ;
            IOAerror[u] = IOAmax[u] * 0.05 ; // 5% error..  estim of sd
          }
        }

        parameters {
          real <lower=eps, upper=1.0> X[U,N];
          real <lower=0.0> g[U,N];
          real <lower=0.0, upper=0.5> m[U];
          real <lower=1.0> K[U];
          real <lower=0.0, upper=0.5>  psd[U];  // process error
          real <lower=0.0, upper=1.0> res[U,N];
          real <lower=0.0, upper=0.5>  ressd[U];  /////
          real <lower=0.0, upper=1.0>  gsd[U]; /////
          real <lower=0.5, upper=1.5>   q[U];
          real <lower=0.0, upper=0.5>   msd; ////
          real <lower=0.0, upper=0.2>   qsd;
        }

        model {
          psd ~ normal(0.0, 0.2) ;
          qsd ~ normal(0.0, 0.05);  // ranges from 0.5 to 1.5
          ressd ~ normal(0.0, 0.05) ;
          msd ~ normal(0.0, 0.05) ;
          gsd ~ normal(0.0, 0.05) ; // 0.1
          m ~ normal(0.0, msd);
          q ~ normal(1.0, qsd );

          for (u in 1:U) {
            g[u,] ~ normal( 0.0, gsd[u] );
            K[u] ~ normal( IOAmax[u], IOAerror[u] ) ;
            X[u,1] ~ normal( 0.5, psd[u] ) ;
            for (n in 1:N) {
              res[u,n] ~ normal( q[u] * (X[u,n] - CAT[u,n]/K[u] ), ressd[u] );  // no catch observation error  .. keep separate to force positive value
              IOA[u,n] ~ poisson( K[u] * res[u,n] ); // survey index observation model
              if (n < N) {
                X[u,n+1] ~ normal( X[u,n] * (1.0 + g[u,n] - m[u]*X[u,n] ) - CAT[u,n]/K[u], psd[u] ) ; //process model
              }
            }
          }
        }
      "
    }

    return(out)
  }

  if (grepl("stochastic_simulation", selection) ) {

    ellp = list(...)
    attach(ellp)

    Xdim = dim(X)
    nau = Xdim[1]
    ntu = Xdim[2]
    nsim = Xdim[3]

    if (!exists("ST")) ST = c(
      "@ -> b*X  -> X" ,
      "X -> d*(X/K)*X -> @",
      "X -> f*X -> catch"
    )

    if (!exists("SC"))  SC = c( "X",  "catch" )

    if (!exists("nthreads")) nthreads = 4
    if (!exists("nprojections")) nprojections = 30
    if (!exists("istart"))  istart = ntu

    nssims = min( 10, nsim )
    iss = sample.int( nsim, nssims )

    sim = array( NA, dim=c( nssims, length(SC), nprojections, nau ) )
    tspan = 1:nprojections

    for (iau in 1:nau) {

      # posterior subset
      aup = data.frame(
        X = as.integer( floor(abundance[iss, iau, istart]) ),
        catch = as.integer( rep(0, length(X))),
        f = f[iss, iau, istart],
        b = g[iss, iau, istart],
        m = m[iss, iau, istart],
        K = as.integer( K[iss, iau] )
      )

      # simulate over each posterior subset
      for (i in 1:nssims) {
        sim[i,,iau]  = slot( run( model=mparse(
          transitions = ST,
          compartments = SC,
          gdata = c( K=aup$K[i], g=aup$g[i], m=aup$m[i], f=aup$f[i] ),
          u0 = aup[i, SC],
          tspan = tspan
          ), threads=nthreads
        ), "U")[]
      }

    }

    detach(ellp)
    out = list( sim=sim, iss =iss)

    return(out)
  }

  if (grepl("posteriors", selection) ) {

    ellp = list(...)
    attach(ellp)

    posteriors =  stan_extract( as_draws_df( fit$draws() ) )  # extract all posteriors = mcmc posteriors from STAN

    Xdim = dim(posteriors$X)
    nsim = Xdim[1]
    nau = Xdim[2]
    ntu = Xdim[3]

    posteriors$numbers = posteriors$X[] * 0
    for ( i in 1:ntu ) posteriors$numbers[,,i] = posteriors$X[,,i] * posteriors$K[]  # X is fraction of K

    posteriors$fraction.fished = posteriors$X[] * 0
    for ( i in 1:nsim ) posteriors$fraction.fished[i,,] = catches[] / posteriors$numbers[i,,]

    posteriors$fishing.mortality =  1.0 - posteriors$fraction.fished[]
    posteriors$fishing.mortality[ which(posteriors$fishing.mortality < 1e-6) ] = 1e-6
    posteriors$fishing.mortality = -log( posteriors$fishing.mortality )
    posteriors$fishing.mortality[ !is.finite(posteriors$fishing.mortality)] = 0

    # K = array( K, dim=c(dim(K), nsim) )
    # MSY    = r * K / 4.0 ; # maximum height of of the latent productivity (yield)
    # aMSY   = K /2.0 ; # numbers at MSY
    # FMSY   = 2.0 * MSY / K ; # fishing mortality at MSY

    # add other params computed post-mcmc from conditional distributions
    # convert num to biomass
    posteriors$biomass = posteriors$numbers[] * NA  # in kg
    for (i in 1:nsim) posteriors$biomass[i,,] = posteriors$numbers[i,,] * wgts[]

    detach(ellp)
    return( posteriors )
  }

  if (grepl("summary", selection)) {

    sims = lapply( apply( posteriors$biomass, 1,
      function(biom) {
        biom = as.matrix(biom)
        biom[!is.finite(biom)] = NA
        o = list()
        o$cfaall    = colSums( biom * sppoly$au_sa_km2/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfanorth  = colSums( biom * sppoly$cfanorth_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfasouth  = colSums( biom * sppoly$cfasouth_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfa23     = colSums( biom * sppoly$cfa23_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfa24     = colSums( biom * sppoly$cfa24_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfa4x     = colSums( biom * sppoly$cfa4x_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        return(o)
      }
    ), data.frame )

    RES = data.frame( index = 1:dim(posteriors$biomass)[3] )

    RES$cfaall = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, mean)
    RES$cfaall_sd = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, sd )
    RES$cfaall_median = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, median )
    RES$cfaall_lb = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfaall_ub = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfanorth = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, mean )
    RES$cfanorth_sd = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, sd )
    RES$cfanorth_median = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, median )
    RES$cfanorth_lb = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfanorth_ub = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfasouth = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, mean )
    RES$cfasouth_sd = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, sd )
    RES$cfasouth_median = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, median )
    RES$cfasouth_lb = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfasouth_ub = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfa23 = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, mean )
    RES$cfa23_sd = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, sd )
    RES$cfa23_median = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, median )
    RES$cfa23_lb = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfa23_ub = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfa24 = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, mean )
    RES$cfa24_sd = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, sd )
    RES$cfa24_median = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, median )
    RES$cfa24_lb = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfa24_ub = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfa4x = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, mean )
    RES$cfa4x_sd = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, sd )
    RES$cfa4x_median = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, median )
    RES$cfa4x_lb = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfa4x_ub = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, quantile, probs=0.975 )

    return(RES)

  }

  if (grepl("map", selection) ) {

    require(sf)

    ellp =list(...)   # if plotting, all ellipsis contents are expected to be sppolt args

    if (is.null(vn)) stop("must have a  vn object")
    if (is.null(spmatrix)) stop("must have a  spmatrix object")
    if (is.null(sppoly)) stop("must have an sppoly object")

    sppoly[,vn] = NA

    # first index is spatial strata

    res_dim = dim( spmatrix )

    if (is.null(time_match)) {
      if (res_dim == 1 ) time_match = NULL
      if (res_dim == 2 ) time_match = list(year="2000")
      if (res_dim == 3 ) time_match = list(year="2000", dyear="0.8" )
    }

    if (is.null(poly_match)) poly_match = sppoly[["AUID"]]

    i_poly = match( poly_match, sppoly[["AUID"]] )  # should match exactly but in case a subset is sent as sppoly

    data_dimensionality = length( dim(spmatrix[[vn]]) )
    if (data_dimensionality==1) {
      sppoly[, vn] = spmatrix[[vn]] [ i_poly ]  # year only
    }

    if (!is.null(time_match)) {
      n_indexes = length( time_match )
      if (data_dimensionality==2) {
        if (n_indexes==1) sppoly[, vn] = spmatrix[[vn]] [ i_poly, time_match[[1]] ]  # year only
      }
      if (data_dimensionality==3) {
        if (n_indexes==1) sppoly[, vn] = spmatrix[[vn]] [ i_poly, time_match[[1]] , ]  # year only
        if (n_indexes==2) sppoly[, vn] = spmatrix[[vn]] [ i_poly, time_match[[1]], time_match[[2]] ] # year/subyear
      }
    }

    if (length(poly_match) > 1 ) {

      dev.new();

      if ( !exists("main", ellp ) )  ellp[["main"]]=vn
      ellp$obj = sppoly
      ellp$zcol=vn
      ellp$col="transparent"

      do.call(spplot, ellp )

    }

  }

}


#' @importFrom plgp randomForest lme4
# '
#' inverse canonical link
#'
#' a simple function for inverse canonical link
#'
#' calculate the inverse canonical link \code{\link{unsae}}
#'
#' @param x a numeric vector
#'
#' @return calculated value
#'
#' @examples
#' hc(1)
#'
#' @export
hc <- function(x) as.numeric(1 /(1 + exp(-x))) # inverse canonical link

#' spatial likelihood
#'
#' a simple (log) spatial likelihood
#' used within multilevel_EM function of \code{\link{unsae}}
#'
#' @param par parameter values
#'
#' @param D distance matrix from plgp::distance
#'
#' @param Y "observed" response vector
#'
#' @param mu current estimate of mu
#'
#' @return negative log likelihood
#'
spatial_lik <- function(par, D, Y, mu) {
  # par: parameters (range and nugget)
  # D: coordinates
  # Y: responses
  # mu: current estimate of mu
  theta <- par[1]
  g <- par[2]
  n <- length(Y)
  K <- exp(-D/theta) + g*diag(n)
  K_inv <- solve(K)
  ldetK <- determinant(K, logarithm=TRUE)$modulus
  ll <- -(n/2)*log(t(Y - mu) %*% K_inv %*% Y - mu) - (1/2)*ldetK
  return(-ll)
}

#' multilevel EM
#'
#' a main function to fit multilevel em with gaussian spatial correlation. Some chunk of the code has been dedicated to handle the syntax. glmer function is called to get the initial value of the estimates.
#'
#' @param formula formula argument following the glmer syntax.
#'
#' @param tol tolerence level to control the em convergence. default value 1e-3.
#'
#' @param data name of the data set.
#'
#' @param coordinates specifying the name of coordinates within the data. These columns *should* exist within the data set provided.
#'
#' @return list of element that are needed for the prediction.
#'
#' @export

multilevel_EM <-
  function(formula, data, tol = 1e-3, coordinates){
    re_pattern <- "\\|(.*?)\\)"
    random_effect <- str_squish(str_match(formula, re_pattern)[,
                                                               2])
    if (!is.factor(data[[random_effect]])) {
      stop("area index must be a factor")
    }
    data$area_index <- droplevels(data[[random_effect]]) %>%
      as.numeric
    M <- data$area_index %>% unique %>% length
    mod_glmer <- lme4::glmer(formula = as.formula(formula), data = data,
                             family = "binomial")
    beta_1_t <- mod_glmer@beta
    sigma2.a_t <- as.numeric(VarCorr(mod_glmer)[1])
    a_i_vec <- ranef(mod_glmer)[[random_effect]][, 1]
    a_i.init_vec <- a_i_vec
    M <- length(a_i_vec)
    area_index_long <- data$area_index
    area_index_unique <- unique(area_index_long)
    a_i_ini_tbl <- tibble(a_i = a_i_vec, area_index = area_index_unique)
    a_i_long_tbl <- tibble(area_index = area_index_long)
    a_i_tbl <- left_join(a_i_long_tbl, a_i_ini_tbl, by = "area_index")
    re_form_pattern <- "\\((.*?)\\)"
    re_part_form <- str_match(formula, re_form_pattern)[1]
    fe_form <- str_remove_all(string = formula, pattern = coll(re_part_form)) %>%
      str_remove_all(pattern = "\\)") %>% str_remove_all(pattern = "\\(") %>%
      str_remove_all(pattern = "\\|") %>% str_squish()
    error_pattern <- "\\~(.*?)\\+"
    fixed_form <- ifelse(str_ends(string = fe_form, pattern = "\\+"),
                         str_sub(string = fe_form, start = 1L, end = (str_length(fe_form) -
                                                                        1)), str_replace(string = fe_form, pattern = error_pattern,
                                                                                replace = "~")) %>% str_squish()
    response_var <- str_squish(str_split(fixed_form, "\\~",
                                         simplify = TRUE)[1])
    x <- model.matrix(as.formula(fixed_form), data = data)
    y <- data[[response_var]]
    if (length(unique(y)) != 2) {
      stop("too many values in response")
    }
    if (!is.numeric(y)) {
      y <- factor(y)
    }
    if (is.factor(y)) {
      y <- as.numeric(y == levels(factor(y))[2])
    }
    data$y <- y
    pred_part <- str_squish(str_split(fixed_form, "\\~",
                                      simplify = TRUE)[2])
    beta1.init <- rep(0, ncol(x))
    a_i.old_vec_long <- rep(0, length(y))
    change <- Inf
    beta1.old <- beta1.init
    a_i.old_vec <- a_i.init_vec
    outer_change <- Inf
    a_i.old_vec <- rep(0, M)
    cnt <- 0
    mu_hat <- 0
    tau2_hat <- 1
    params <- c(0.2, 2)
    while (outer_change > tol & cnt <= 20) {
      change <- Inf
      while (change > tol) {
        x_beta1 <- x %*% beta1.old %>% as.numeric
        eta <- x_beta1 + a_i.old_vec_long
        y.hat <- hc(eta)
        h.prime_eta <- ifelse(near(y.hat * (1 - y.hat), 0),
                              1e-04, y.hat * (1 - y.hat))
        z <- (y - y.hat)/h.prime_eta
        temp_tbl <- data
        temp_tbl$z <- z
        wls_form <- paste("z~", pred_part)
        beta1.new <- beta1.old + lm(formula = as.formula(wls_form),
                                    weights = h.prime_eta, data = data)$coef
        change <- sqrt(sum((beta1.new - beta1.old)^2))
        beta1.old <- beta1.new
      }
      beta_1.t <- beta1.old
      wk_tbl <- data %>% rename(a_index = area_index)
      wk_tbl$w_resid <- y - x %*% beta_1.t %>% as.numeric
      wk_tbl <- wk_tbl %>% mutate(pi_w = 1 - hc(w_resid))
      a_i.new_vec <- rep(NA, length(a_i.old_vec))
      v_i_vec <- rep(NA, length(a_i.old_vec))
      for (m in 1:M) {
        wk_tbl_i <- wk_tbl %>% filter(a_index == m)
        x_i <- model.matrix(as.formula(fixed_form), data = wk_tbl_i)
        y_i <- wk_tbl_i$y
        if (length(unique(y_i))>1){a_i.old <- a_i.old_vec[m]
        change <- Inf
        inner_cnt <- 0
        while (change > tol & inner_cnt <= 100) {
          x_beta1 <- x_i %*% beta1.old
          eta <- x_beta1 + a_i.old %>% as.numeric
          y.hat_i <- hc(eta)
          h.prime_eta_i <- ifelse(near(y.hat_i * (1 - y.hat_i),
                                       0), 1e-04, y.hat_i * (1 - y.hat_i))
          z_i <- (y_i - y.hat_i)/h.prime_eta_i
          a_i.new <- a_i.old + lm(z_i ~ 1, weights = h.prime_eta_i)$coef %>%
            as.numeric
          change <- sqrt(sum((a_i.new - a_i.old)^2))
          a_i.old <- a_i.new
          inner_cnt <- inner_cnt + 1
          a_i.new_vec[m] <- a_i.old
          v_i_vec[m] <- 1/sum(h.prime_eta_i)
        }
        }else{
          a_i.new_vec[m] <- a_i.old
          v_i_vec[m] <- 1e6
        }
      }
      V <- diag(v_i_vec)

      spat_coord <- data %>% dplyr::select(all_of(coordinates)) %>%
        unique
      D <- plgp::distance(spat_coord)
      K <- exp(-D/params[1]) + diag(params[2], nrow(spat_coord))
      Ki <- solve(K)
      mu_hat <- drop(t(rep(1, nrow(Ki))) %*% Ki %*% a_i.new_vec/sum(Ki))
      a_star <- drop(mu_hat + solve(tau2_hat * K + V) %*% (tau2_hat *
                                                             K) %*% (a_i.new_vec - mu_hat))
      V_star <- tau2_hat * K - tau2_hat * K %*% solve(tau2_hat *
                                                        K + V) %*% (tau2_hat * K)
      tau2_hat <- drop(sum(diag(solve(tau2_hat * K) %*% V_star)) +
                         t(a_star) %*% solve(tau2_hat * K) %*% a_star)/nrow(K)
      out <- optim(c(0.1, 0.01), spatial_lik, method = "L-BFGS-B",
                   lower = 1e-07, upper = c(10, 10), D = D, Y = a_star,
                   mu = mu_hat)
      params <- out$par
      a_i_ini_tbl_new <- tibble(a_i = a_star, area_index = unique(data$area_index))
      a_i_long_tbl_new <- tibble(area_index = data$area_index)
      a_i_tbl_new <- left_join(a_i_long_tbl_new, a_i_ini_tbl_new,
                               by = "area_index")
      a_i.old_vec_long <- a_i_tbl_new$a_i
      outer_change <- mean((a_star - a_i.old_vec)^2)
      # plot(a_i.old_vec, a_star)
      a_i.old_vec <- a_star
      cnt <- cnt + 1
      cat("...EM iteration:", cnt, "\n")
    }

    outcome <- list(cnt = cnt, area_tbl = a_i_tbl_new, area_re = a_i.old_vec,
                    params = c(out$par), mu_hat = mu_hat, spat_coord = spat_coord,
                    fomula = formula, V = v_i_vec, beta_hat = beta1.new,
                    tau2_hat = tau2_hat, D = D, coordinates = coordinates)

  return(outcome)
}

#' spatial prediction
#'
#' to calculate the spatial random effect for a new data set.
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set). The data set *should* have the predictor columns as well as the coordinate columns, all with the same names with the training data set.
#'
#' @return predicted spatial random effect at new data set

spatial_pred <- function(em_output, test_set){
  coordinates <- em_output$coordinates
  spat_coord <- em_output$spat_coord
  new_site <- test_set %>% dplyr::select(all_of(coordinates))
  if (is.null(nrow(new_site))){ # to make sure it works for a single site
    new_site <- as.data.frame(matrix(new_site, nrow = 1))
  }
  par <- em_output$params
  DXX <- plgp::distance(new_site)
  D <-  em_output$D
  mu_hat <- em_output$mu_hat
  a_star <- em_output$area_re
  K <- exp(- D/par[1]) + par[2]*diag(nrow(spat_coord))
  spat_coord <- em_output$spat_coord
  KXX <- exp(-DXX/par[1]) + par[2]*diag(ncol(DXX))
  DX <- plgp::distance(new_site, spat_coord)
  KX <- exp(-DX/par[1])
  Ki <- solve(K)
  a_hat <- mu_hat + KX %*% Ki %*% (a_star - mu_hat)
  a_hat <- drop(a_hat)
  return(a_hat)
}

#' covariate prediction
#'
#' to calculate the spatial random effect for a new data set.
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set). The data set *should* have the predictor columns as well as the coordinate columns, all with the same names with the training data set.
#'
#' @return predicted covariates parts Xb at new data set

covariate_pred <- function(em_output, test_set){
  beta_hat <- em_output$beta_hat
  formula <- em_output$fomula
  re_form_pattern <- "\\((.*?)\\)"
  re_part_form <- str_match(formula, re_form_pattern)[1]
  fe_form <-
    str_remove_all(string = formula, pattern = coll(re_part_form)) %>%
    # coll is to force the matching the "full expression", not as a regex
    str_remove_all(pattern = "\\)") %>%
    str_remove_all(pattern = "\\(") %>%
    str_remove_all(pattern = "\\|") %>%
    str_squish()

  error_pattern <- "\\~(.*?)\\+"
  fixed_form <-
    ifelse(str_ends(string = fe_form, pattern = "\\+"),
           str_sub(string = fe_form, start = 1L, end = (str_length(fe_form)-1)),
           str_replace(string = fe_form, pattern = error_pattern, replace = "~")) %>%
    str_squish()

  mm_form <- str_squish(str_split(fixed_form, "\\~", simplify = TRUE)[2])

  x <- model.matrix(as.formula(paste("~", mm_form)), data = test_set)
  out <- as.numeric(x %*% beta_hat)
  return(out)
}

#' main prediction
#'
#' a wrapper to calculate the spatial random effect for a new data set.
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set). The data set *should* have the predictor columns as well as the coordinate columns, all with the same names with the training data set.
#'
#' @return predicted outcome combining covariates parts Xb and spatial random effects at new data set
#'
#' @export

predict_em <- function(em_output, test_set){
  yhat <- covariate_pred(em_output, test_set)
  spat_hat <- spatial_pred(em_output, test_set)
  combined_hat <- yhat + spat_hat
  phat <- hc(combined_hat)
  return(phat)
}

#' calculate the MSE for each area
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set).
#' The data set *should* have the predictor columns as well
#' as the coordinate columns,
#' all with the same names with the training data set.
#'
#' @return simulated spatial random effect of size 100 (100 is default)
#' at new data set
#'
#'

cal_mse <- function(em_output, test_set, size = 100){
  coordinates <- em_output$coordinates
  spat_coord <- em_output$spat_coord
  tau2_hat <- em_output$tau2_hat
  new_site <- test_set %>% dplyr::select(all_of(coordinates))
  if (is.null(nrow(new_site))){ # to make sure it works for a single site
    new_site <- as.data.frame(matrix(new_site, nrow = 1))
  }
  par <- em_output$params
  DXX <- plgp::distance(new_site)
  D <-  em_output$D
  mu_hat <- em_output$mu_hat
  a_star <- em_output$area_re
  K <- exp(- D/par[1]) + par[2]*diag(nrow(spat_coord))
  spat_coord <- em_output$spat_coord
  KXX <- exp(-DXX/par[1]) + par[2]*diag(ncol(DXX))
  DX <- plgp::distance(new_site, spat_coord)
  KX <- exp(-DX/par[1])
  Ki <- solve(K)
  a_hat <- mu_hat + KX %*% Ki %*% (a_star - mu_hat)
  a_hat <- drop(a_hat)

  Sigmap <- tau2_hat*(KXX - KX %*% Ki %*% t(KX))

  DXX <- plgp::distance(new_site)
  KXX <- exp(-DXX/par[1]) + diag(par[2], nrow(DXX))

  Sigmap <- tau2_hat*(KXX - KX %*% Ki %*% t(KX))

  Sigma.int <- tau2_hat*(exp(-DXX) + diag(par[2], nrow(DXX))
                        - KX %*% Ki %*% t(KX))
  a_k <- rmvnorm(size, a_hat, Sigma.int)

  # routine to calculate the MSE part 1
  yhat <- covariate_pred(em_output, test_set)
  rep_yhat_mat <- matrix(rep(yhat, size), nrow = size, byrow = TRUE)
  combined_mat <- a_k + rep_yhat_mat

  v_2 <- var((apply(1 / (1 + exp(-combined_mat)), 1, mean)))
  v_1 <- var(apply(1 / (1 + exp(-combined_mat)), 2, mean)) / ncol(combined_mat)

  total_v <- v_1 + v_2

  return(total_v)
}


################################################################################
# Dependencies
# library("mboost")

MboostBlackboostTrain <- function(train.name="train",target=NULL,args=list(test.name=NULL,R=1,pls=FALSE,
		prune=FALSE,std=FALSE,features=c("static","alsfrs","fvc","svc","vital"),
		slope.features=NULL)) {
	# Trains a MboostBlackboost model given training data.
	# Args:
	#	train.name: base name of training dataset
	#	target: target vector; if NULL, future slopes are used as target
	#	args: list of additional arguments
	#	- pls: TRUE means predicting the last score rather than the slope
	#	- prune: TRUE means keeping only features with data for at least 918/2
	#		patients
	#	- std: TRUE means standardize columns by subtracting away column means
	#	- features: a character vector of feature sets to get for the training
	#		data.  Choose from "static","alsfrs","fvc","svc","vital","lab"
	#	- slope.features: a character vector of datasets to get slope features
	#		for.  Choose from NULL if no slopes are desired, or among "static",
	#		"alsfrs","fvc","svc","vital","lab"
	#
	# Returns:
	#	a function that given a test set, returns predictions
	#
	# Side effects:
	#	Saves predictions to file
	#
	# Example usage:
	if(FALSE) {
		features = c("static","alsfrs","fvc","svc","vital")
	   	slope.features = c("alsfrs","fvc","svc","vital")
	   	args=list(test.name="test",R=1,pls=FALSE,prune=TRUE,std=FALSE,
	   		features=features,slope.features=slope.features)
		model <- MboostBlackboostTrain(train.name="train",target=NULL,args=args)
		predictions<-model("test")
	}

	# Preprocess training data according to arguments
	train.preprocess<-PreprocessFeatures(train.name=train.name,test.name=
		args[["test.name"]],args=args)

	# Get training data and target
	train.data<-as.data.frame(train.preprocess$feature.data)
	target<-train.preprocess$target

	# Drop constant columns
	train.data<-train.data[,apply(train.data,2,function(c)length(unique(c))>1)]

	# Fit multiple models to training data
	GetModel <- function(i) {
		# Boosting with P-splines with a B-spline basis
		# object<-mboost(target~.,data.frame(target,train.data[,-1]),baselearner="bbs")
		# gamboost(target~.,data.frame(target,train.data[,-1]),
				# control = boost_control(mstop = ...)),
				# baselearner = c("bbs", "bols", "btree", "bss", "bns"), dfbase = 4)
		# Boosting with linear models
		# object<-mboost(target~.,data.frame(target,train.data[,-1]),baselearner="bols")
		# glmboost(formula, data = list(),
				# na.action = na.fail, contrasts.arg = NULL,
				# center = TRUE, control = boost_control(), ...)
		# Boosting with conditional trees
		#object<-mboost(target~.,data.frame(target,train.data[,-1]),baselearner="btree")
		object<-blackboost(target~.,data.frame(target,train.data[,-1]),
						   control = boost_control(mstop = 200),
						   # Default tree control
						    # tree_controls = ctree_control(
							 # teststat = "max",
							 # testtype = "Teststatistic",
							 # mincriterion = 0,
							 # maxdepth = 2, savesplitstats = FALSE))
							# Stump tree control
							tree_controls = ctree_control(stump = TRUE,
							mincriterion = 0,
							 savesplitstats = FALSE))

		# Select optimal number of boosting iterations using CV
		cv.err<-cvrisk(object, folds = cv(model.weights(object),type="kfold"))
		object[mstop(cv.err)]

		# Report summary results
		print(sprintf("Model: %d, CV error: %g, Best iter: %g",i,
			sqrt(min(colMeans(cv.err[,]))),mstop(cv.err)))

		# predictions <- predict(object,data.frame(test.data[,-1]))
		# Rmsd(predictions, test.target)
		#test.target <- GroundTruthTrain()()
		#Rmsd(predictions, test.target[,2])

		return(list(object = object))
	}
	models <- lapply(1:args[["R"]], GetModel)

	MboostBlackboostPredict <- function(test.name="test",save.preds=TRUE) {
		# Returns predictions based on the trained model given a test set
		#
		# Args:
		#	test.name: base name of test dataset
		#   save.preds: save predictions to file? TRUE by default
		#
		# Returns:
		#	Two column matrix, ordered by the first column, with first column
		#   "subject.id" containing test set subject ids and the second column
		#   "prediction" containing predicted slopes.
		#
		# Side effects:
		#	Saves predictions to file

		# Preprocess test data according to arguments
		test.preprocess<-PreprocessFeatures(train.name=NULL,test.name=test.name,
			args=args,train.cols=train.preprocess$train.cols,train.data.std=
			train.preprocess$train.data.std,train.data.na=
			train.preprocess$train.data.na)

		test.data<-as.data.frame(test.preprocess$feature.data)

		# Match train columns
		test.data<-test.data[,colnames(train.data)]

		# Form test predictions
		predictions <- rowMeans(sapply(models, function(model) predict(model$object,
			test.data[,-1])))

		# If predicting last score, transform predictions back to slopes
		if (args[["pls"]]) {
			predictions<-predictions-
				scaling$score.first*scaling$time.scale
		}
		predictions <- cbind(test.data[,1],predictions)

		if(save.preds) {
			# Write predictions to file
			out.file = GetPredictionFilename("MboostBlackboostTrain",train.name,test.name,args)
			WritePredictions(predictions, out.file)
		}
		return(predictions)
	}
	return(MboostBlackboostPredict)
}


##################################################################################################################
####  Purpose: Final Project - Classification - c50
####  Group: Ravi P., Derek P., Aneesh K., Ninad K.
####  Date: 11/28/2018
####  Comment:
####    setwd("~/Stevens/Fall2018/CS513_Knowledge_Discovery/Final Project/Classification")
####
##################################################################################################################

#### 0. Clean up environment ####
rm(list = ls())

#### 1. Import Libraries and Datasets ####
library('C50')
library(rpart)
library(rpart.plot)
library(partykit)

df_train <- read.csv("./data/train_new.csv")
df_test <- read.csv("./data/test_new.csv")

#### 2. Remove unwanted data ####
df_train <- df_train[,-1]
df_test <- df_test[,-1]
levels(df_test$SimpleColor) <- levels(df_train$SimpleColor)

#### 3.C5.0 Tree ####

mytree <- C5.0( OutcomeType~ AnimalType + AgeinDays + HasName + IsNeutered + IsMix + Gender + SimpleColor, data =df_train )
summary(mytree)
plot(mytree)

# myTree2 <- C50:::as.party.C5.0(mytree)
# fit<-fitted(myTree2)
# plot(myTree2[2])
# plot(myTree2[10])

#### 4. Error Checking ####
prediction<-predict( mytree ,df_test , type="class" )
table(actual=df_test[,2],prediction)
wrong<- (df_test[,2]!=prediction)
rate<-sum(wrong)/length(wrong)
print(rate)


cat("#### Test estimateV at, fa, rr functions for MET data with asreml4\n")
test_that("MET_estimateV_asreml4", {
  skip_if_not_installed("asreml")
  skip_on_cran()
  library(dae)
  library(asreml)
  library(asremlPlus)

  data(MET)
  asreml.options(design = TRUE, keep.order=TRUE)
  #Test at
  asreml.obj <-asreml(fixed = GY.tha ~  + at(expt, c(1:5)):rep + at(expt, c(1)):vrow +
                        at(expt, c(2,3,6,7)):colblocks +
                        at(expt, c(1:5,7)):vcol + Genotype*Condition*expt,
                      random = ~  at(expt, c(1)):dev(vrow) + at(expt, c(2)):spl(vcol) +
                        at(expt, c(3,5,7)):dev(vcol) + at(expt, c(7)):units,
                      data=comb.dat, maxiter = 100, workspace = "1Gb")

  summary(asreml.obj)$varcomp
  ranterms <- names(asreml.obj$G.param)
  n <- nrow(comb.dat)
  V.g <- matrix(0, nrow = n, ncol = n)
  for (term in ranterms)
  {
    cols <- grep(term, colnames(asreml.obj$design), fixed = TRUE)
    V.g <- V.g + asreml.obj$vparameters[term] * (asreml.obj$design[, cols] %*%
                                                   t(as.matrix(asreml.obj$design[, cols])))
  }
  V.g <- asreml.obj$sigma2 * (V.g + mat.I(n))
  Vat <- estimateV(asreml.obj)
  testthat::expect_true(all(abs(Vat - V.g) < 1e-06))

  #Test fa
  asreml.obj <-asreml(fixed = GY.tha ~  + at(expt, c(1:5)):rep + at(expt, c(1)):vrow +
                        at(expt, c(2,3,6,7)):colblocks +
                        at(expt, c(1:5,7)):vcol + Condition*expt,
                      random = ~  fa(exptCond, k = 2):Genotype +
                        at(expt, c(1)):dev(vrow) + at(expt, c(2)):spl(vcol) +
                        at(expt, c(3,5,7)):dev(vcol) + at(expt, c(7)):units,
                      data=comb.dat, maxiter = 100, workspace = "1Gb")

  summary(asreml.obj)$varcomp
  ranterms <- names(asreml.obj$G.param)
  n <- nrow(comb.dat)
  V.g <- matrix(0, nrow = n, ncol = n)
  for (term in ranterms[2:7])
  {
    cols <- grep(term, colnames(asreml.obj$design), fixed = TRUE)
    V.g <- V.g + asreml.obj$vparameters[term] * (asreml.obj$design[, cols] %*%
                                                   t(as.matrix(asreml.obj$design[, cols])))
  }
  term <- ranterms[1]
  cols <- grep(term, colnames(asreml.obj$design), fixed = TRUE)[1:1364]
  vp <- asreml.obj$vparameters[names(asreml.obj$vparameters)[grep(term,
                                           names(asreml.obj$vparameters), fixed = TRUE)]]
  spec.var <- diag(vp[grepl("!var", names(vp), fixed = TRUE)])
  loads <- matrix(vp[grepl("!fa", names(vp), fixed = TRUE)], ncol = 2)
  Gfa <- loads %*% t(loads) + spec.var
  V.g <- V.g + (asreml.obj$design[, cols] %*% kronecker(Gfa, mat.I(62)) %*%
                  t(as.matrix(asreml.obj$design[, cols])))
  V.g <- asreml.obj$sigma2 * (V.g + mat.I(n))
  Vfa <- estimateV(asreml.obj)
  testthat::expect_true(all(abs(Vfa - V.g) < 1e-06))

  #Test rr
  asreml.obj <-asreml(fixed = GY.tha ~  + at(expt, c(1:5)):rep + at(expt, c(1)):vrow +
                        at(expt, c(2,3,6,7)):colblocks +
                        at(expt, c(1:5,7)):vcol + Condition*expt,
                      random = ~  rr(exptCond, k = 2):Genotype +
                        at(expt, c(1)):dev(vrow) + at(expt, c(2)):spl(vcol) +
                        at(expt, c(3,5,7)):dev(vcol) + at(expt, c(7)):units,
                      data=comb.dat, maxiter = 100, workspace = "1Gb")

  summary(asreml.obj)$varcomp
  ranterms <- names(asreml.obj$G.param)
  n <- nrow(comb.dat)
  V.g <- matrix(0, nrow = n, ncol = n)
  for (term in ranterms[2:7])
  {
    cols <- grep(term, colnames(asreml.obj$design), fixed = TRUE)
    V.g <- V.g + asreml.obj$vparameters[term] * (asreml.obj$design[, cols] %*%
                                                   t(as.matrix(asreml.obj$design[, cols])))
  }
  term <- ranterms[1]
  cols <- grep(term, colnames(asreml.obj$design), fixed = TRUE)[1:1364]
  vp <- asreml.obj$vparameters[names(asreml.obj$vparameters)[grep(term,
                                                                  names(asreml.obj$vparameters), fixed = TRUE)]]
  loads <- matrix(vp[grepl("!fa", names(vp), fixed = TRUE)], ncol = 2)
  Gfa <- loads %*% t(loads)
  V.g <- V.g + (asreml.obj$design[, cols] %*% kronecker(Gfa, mat.I(62)) %*%
                  t(as.matrix(asreml.obj$design[, cols])))
  V.g <- asreml.obj$sigma2 * (V.g + mat.I(n))
  Vrr <- estimateV(asreml.obj)
  testthat::expect_true(all(abs(Vrr - V.g) < 1e-06))

  asreml.options(design = FALSE)

})


################################################################################
#This file uses the function calcRa.r to calculate extraterrestrial solar
#radiation (solar radiation at the top of the atmosphere) for a daily time
#step that is then averaged into monthly.
#NB: This does not need to be done for each year of the time series because Ra
#is a function of day of year and latitude, neither of which vary between years.
#It then produces a text file of Ra values for each month.
#NB:  The text files produced here have no header for easier processing in R but
#may need to be modified for usein ArcMap.
################################################################################

#PLEASE READ THE INTERNAL DOCUMENTATION IN calcRa.r FOR MORE INFORMATION ON THE CALCULATION
library(raster)
library(sp)
library(maptools)
library(rgeos)

#SET THE WORKING DIRECTORY
setwd("/workspace/Shared/Tech_Projects/ALFRESCO_Inputs/project_data/TEM_Data/girr_radiation_cmip3_process/")

outDir<-"/workspace/Shared/Tech_Projects/ALFRESCO_Inputs/project_data/TEM_Data/girr_radiation_cmip3_process/AKCAN_2km/"

#ADD THE FUNCTION calcRa.r
source("/workspace/Shared/Tech_Projects/ALFRESCO_Inputs/project_data/CODE/tem_ar5_inputs/downscale_cmip5/bin/calcRa.r")

# use this as the template.map
template.map <- raster("/Data/malindgren/cru_november_final/ar5/GFDL-CM3/cld/downscaled/cld_metric_GFDL-CM3_historical_r1i1p1_03_1955.tif")

# here we make only the cells with actual values a new value of zero.  the rest remain out of bounds.
ind <- which(getValues(template.map) > -9999)
values(template.map)[ind] <- 0

#project raster
#template.map.projExt <- projectExtent(template.map, crs="+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs")
#template.map <- projectRaster(template.map, template.map.projExt, method="ngb")

# create a new map where the cells values are the latitudes of the cell center
template.map.lat <- raster(matrix(coordinates(template.map)[,1],nrow=nrow(template.map),ncol=ncol(template.map)),xmn=xmin(template.map),xmx=xmax(template.map),ymn=ymin(template.map),ymx=ymax(template.map),crs=projection(template.map))

# take that template map and turn all values that are NA in the template.map back to NA
ind <- which(is.na(getValues(template.map)))
values(template.map.lat)[ind] <- NA

# calculate the radians here:
#  radians = degrees*pi/180
template.map.rad <- (template.map.lat*pi)/180

# this is a testing line that will write out the current file to an ascii grid to see if reading the damn thing in as a table will solve the issues I am having...
writeRaster(template.map.rad, filename=paste(outDir,"template_map_akalb_rad.asc",sep=""), overwrite=TRUE)

# here we read that same file back in as a table where we skip the 6 lines at the top of the arc ascii grid which is a header
template.map.rad <- t(read.table(paste(outDir,"template_map_akalb_rad.asc",sep=""), sep=" ", skip=6))

#MAKE A VECTOR THAT DESCRIBES THE NUMBER OF DAYS IN EACH MONTH
ndm <- matrix(c(31,28,31,30,31,30,31,31,30,31,30,31))

#USE THAT MATRIX TO PULL TOGETHER A MATRIX LISTING THE FIRST AND LAST DAY OF EACH MONTH
n <- matrix(NA,12,2)
for (m in 1:12) {
n[m,2] <- sum(c(ndm[1:m]))   #Sum up number of days in months to find last day of month
} #close that loop
n[,1] <- (n[,2]-ndm)+1 #Subtract number of days in a month from the end date and add one to get Julian Day of the first day of month
rm(ndm,m) #Delete ndm b/c not needed any more; clean up looping variable

#DEFINE A VECTOR OR MONTH DESIGNATIONS FOR READING IN FILES AND NAMING OUTPUT FILES
month<-c(1:12)

#RUN THE FUNCTION AND SAVE OUTPUT TO TEXT FILES
#This version of the file calculates Ra for each day within the month and then averages daily radiation to get a monthly value, rather than using a "representative" day.
for (m in 1:12) {  #month loop
    #Identify the Julian Days associated with that month
    md <-n[m,1]:n[m,2]
    #Pre-allocate the Ra matrix
    ra <- matrix(NA,nrow(template.map.rad),ncol(template.map.rad))
    #Also pre-allocate a temporary array to hold daily values of Ra
    temp <- array(NA,c(length(md),nrow(template.map.rad),ncol(template.map.rad)))
    for (t in 1:length(md))    {  #loop over the number of days in each month
       r <- matrix(NA,nrow(template.map.rad),ncol(template.map.rad)) #Make a matrix for  each day
       #Run the function  for each day
       r <- calcRa(md[t],template.map.rad,nrow(template.map.rad),ncol(template.map.rad))
      #Slot the daily matrix r into array temp and delete r
       temp[t,,] = r
        rm(r)
    }      #End looping over number of days per month
    #Average the daily Ra values into a single month value
    ra <- apply(temp,c(2,3),mean, na.rm = TRUE)
    ra <- round(ra,3)
    rm(temp)   #clean up the temporary matrix
    #Set up a geotiff file to store data
    out <- raster(ra,xmn=xmin(template.map),xmx=xmax(template.map),ymn=ymin(template.map), ymx=ymax(template.map), crs=projection(template.map))
    writeRaster(out,filename=paste(outDir,"rad_",month[m],"_1km_iem_akalb.tif",sep=""),options=c('COMPRESS=LZW'), overwrite=TRUE)

    # ra.akalb.projExt <- projectExtent(out, crs=projection(template.map))
    # ra.akalb <- projectRaster(out, ra.akalb.projExt, method='bilinear', filename=paste(outDir,"rad_",month[m],"_1km_iem.tif",sep=""), options=c('COMPRESS=LZW'),overwrite=TRUE)
    #Write data to the file
    #writeRaster(out,filename=paste('ak_Ra_',month[m],'.tif',sep=''),format = 'GTiff',options='COMPRESS=LZW',datatype='FLT4S',overwrite=T)
    rm(ra,out)  #delete files that will be regenerated for the next month
    print(paste(" >>> completed: ",m,sep=""))

} #close month loop


designBasedContNullBatch2 = function(gamma_vec, outcome_mat, pcf_cat_vec, instid_vec){
	#Input an N x iterations matrix of null data

	##Computes D-G z scores less cost prohabitivly....
	#There is a lot of redundancy in the computation

	###Draper Gittoes
	#M student types - pcf categories
	#N universities
	#n_ga,,a gamma types
	#iters, number of iterations
	#N_total - overall number of students

	###First, turn everything into a factor
	# create lambda matrix
	#sparse indicators for membership``
	#loop over iters
	#	use aggregate to compute mean and sd
	# will have to insert into a sparse matrix after computing w/ aggregate...shouldnt be too hard
	#  	sparse linear algebra for to transformations, try to avoid using too  much storage space

	#Helper Function
	insertAggregate = function(Mat, agg){
  #We do this a few times
  #Aggregated on 2 dimensions, then insert into a sparse matrix
  Mat[ cbind( as.numeric(agg[ ,1]),
              as.numeric(agg[ ,2]))] = as.numeric(agg[ ,3])
	return(Mat)
	}

	#Extract variables
pcf_cat_vec = as.factor(pcf_cat_vec)
pcf_cat_vec_i = as.integer(pcf_cat_vec)
instid_vec = as.factor(instid_vec)
instid_vec_i = as.integer(instid_vec)

#Summary statistics
n_gamma = length(gamma_vec)
n_replicates = dim(outcome_mat)[2]
n_total = length(pcf_cat_vec)
iters = dim(outcome_mat)[2]
N = length(levels(instid_vec))
M = length(levels(pcf_cat_vec))
insts = levels(instid_vec)
cats = levels(pcf_cat_vec)

#Count number in each category (lower half of HEFCE grid (table 2))
n_counts = aggregate(rep(1, n_total) ~
              instid_vec * pcf_cat_vec, FUN = sum)
n_mat = matrix(0, nrow = N, ncol = M)
n_mat = insertAggregate(n_mat, n_counts)

	#Marginal counts
	n_u_vec = apply(n_mat, 1, sum)
	n_cat_vec = apply(n_mat, 2, sum)

	#This takes a while, but only needs to be done once for a batch of null simulations
    #Compute 3d array of lambdas
	#THIS PART TAKES FOREVER
	lambdaFun = function(i,k,j){
		if (i == k){
			(n_mat[i,j] / n_u_vec[i]) *
				(1 - (n_mat[i,j] / (n_cat_vec[j])))
		} else {
			-(n_mat[i,j] * n_mat[k,j])/(n_u_vec[i] * n_cat_vec[j])
		}
	}

	#Generating the array like this is a little faster
	#Tried using foreach, wasn't faster. I think this is memory access, rather than compute time limited
	lambdaFun = Vectorize(lambdaFun, vectorize.args = c('i','k'))
	print('computing lambda')
	lambda_array = array(NA, c(N, N, M) )
	#lambda_array = array(rnorm(N*N*M), c(N, N, M) )
	 for (j in 1:M){
		 lambda_array[ , ,j] = outer(1:N, 1:N, lambdaFun, j = j)
	 }
	lambda_array_squared = lambda_array^2
	 print('done computing lambda')

	d_mat = matrix(NA, nrow = N, ncol = n_replicates)
	z_array = se_array = array(NA, c(N, n_gamma, iters))
	#cl = makeCluster(2)
	#registerDoParallel(cl)
	n_array = outer(n_mat, rep(1, n_gamma))
	#DO par doesnt speed it up that much due to memory constraints :(
	for (i in 1:n_replicates){#Loop over minimal_data iterations #Turn this parallel?
		outcome_vec = outcome_mat[ ,i]
		var_global = var(outcome_vec)
		#Compute mean and variance matrix
		#N x M
		mean_mat = var_mat = matrix(NA, nrow = N, ncol = M)
		aggregated_mean = aggregate(outcome_vec ~
				instid_vec * pcf_cat_vec, FUN = mean)
		mean_mat = insertAggregate(mean_mat, aggregated_mean)

		aggregated_var = aggregate(outcome_vec ~
				instid_vec * pcf_cat_vec, FUN = var)
		var_mat = insertAggregate(var_mat, aggregated_var)
		var_mat[n_mat==1] = var_global

		#Compute D vec using linear algebra
		d_mat[ ,i] = apply( lambda_array * outer(rep(1,N), mean_mat), 1, sum, na.rm = TRUE)

		varArray = function(gamma){
			#Creates array of dimention N x N x M to compute SE vec using same linear algebra as above
			v_temp = matrix(var_global, nrow = N, ncol = M) * gamma + (1-gamma) * var_mat
			v_temp[n_mat > 0] = v_temp[n_mat > 0] / n_mat[n_mat > 0]
			return( outer(rep(1,N), v_temp))
		}

		##Do it w/ linear algebra instead
		#Matrix algebra to compute the gamma matrix
		#And then matrix algebrea inside foreach (this is instead of calling a function w/ indexing)
		#Foreach doesnt save much time...
		# var_gamma_array = ((outer(var_mat, rep(1, n_gamma)) *
				  # outer(matrix(1,nrow = N, ncol = M), 1-gamma_vec )) + outer( matrix(var_global, nrow = N, ncol = M), gamma_vec) )
		# var_gamma_array[n_array > 0] = var_gamma_array[n_array > 0] / n_array[n_array > 0]
		# a = foreach(j = 1:n_gamma,.combine ='cbind') %dopar% apply(lambda_array_squared * outer( rep(1, N), var_gamma_array[, ,j]), 1, sum, na.rm = TRUE)
		#n_gamma X N x N x M
		a = foreach(j=1:n_gamma, .combine = 'cbind') %do%
				apply( lambda_array_squared * varArray(gamma_vec[j]), 1, sum, na.rm = TRUE)

		se_array[ , ,i] = sqrt(a)
		z_array[ , ,i] = matrix( d_mat[ ,i], nrow = N, ncol = n_gamma, byrow = FALSE) / se_array[ , ,i]
		if ((i %% 10) == 0){print(i)}
	}
		#stopCluster(cl)
	return( #We don't compute O and E for time reasons, use designBasedCont
	list(D = d_mat, Z = z_array, SE = se_array))
}


# this function reformats a .csv file, replacing its multi-line header
# with a single line, and converting the date-time field from Excel
# format to 2 columns, 1 for date in dd/mm/YYYY format, and 1 for time
# in HH:MM:SS format

csv_reformat <- function(infile, outfile)
{
	# remove redundant header info from file with skip, and strip white space
	data <- read.csv(infile, sep=",", dec=".", skip=8, strip.white=TRUE)

	# date conversion from Windows Excel format: see ?format.Date
    dates <- format(as.Date("1899-12-30") + floor(data[[2]]), "%d/%m/%Y")

    # time conversion from Windows Excel format: see ?as.POSIXct
    # and: http://stackoverflow.com/questions/12161984/how-to-elegantly-convert-datetime-from-decimal-to-d-m-y-hms
    times <- format(as.POSIXct(data[[2]]*60*60*24, origin=as.Date("1899-12-30"), tz="UTC"), "%H:%M:%S")

	# remove old date.time col (2) and add new Date and Time columns
	newtable <- data.frame(data[1], dates, times, data[3:ncol(data)])

	# add column names as new header row
	names(newtable) <- c("Scan Number", "Date", "Time", "Sensor 1 Raw Value", "Sensor 1 Water mm", "Sensor 2 Raw Value", "Sensor 2 Water mm", "Sensor 3 Raw Value", "Sensor 3 Water mm", "Sensor 4 Raw Value", "Sensor 4 Water mm", "Sensor 5 Raw Value", "Sensor 5 Water mm")

	# write the reformatted data back out to a csv file
	write.table(newtable, file=outfile, sep=",", row.names=FALSE)
}


#' @export
#' @rdname geom_hex
#' @inheritParams stat_bin_2d
#' @section Computed variables:
#' \describe{
#'   \item{count}{number of points in bin}
#'   \item{density}{density of points in bin, scaled to integrate to 1}
#'   \item{ncount}{count, scaled to maximum of 1}
#'   \item{ndensity}{density, scaled to maximum of 1}
#' }
stat_bin_hex <- function(mapping = NULL, data = NULL,
                         geom = "hex", position = "identity",
                         ...,
                         bins = 30,
                         binwidth = NULL,
                         na.rm = FALSE,
                         show.legend = NA,
                         inherit.aes = TRUE) {
  layer(
    data = data,
    mapping = mapping,
    stat = StatBinhex,
    geom = geom,
    position = position,
    show.legend = show.legend,
    inherit.aes = inherit.aes,
    params = list(
      bins = bins,
      binwidth = binwidth,
      na.rm = na.rm,
      ...
    )
  )
}

#' @export
#' @rdname geom_hex
#' @usage NULL
stat_binhex <- stat_bin_hex

#' @rdname ggplot2-ggproto
#' @format NULL
#' @usage NULL
#' @export
StatBinhex <- ggproto("StatBinhex", Stat,
  default_aes = aes(weight = 1, fill = after_stat(count)),

  required_aes = c("x", "y"),

  compute_group = function(data, scales, binwidth = NULL, bins = 30,
                           na.rm = FALSE) {
    check_installed("hexbin", reason = "for `stat_binhex()`")

    binwidth <- binwidth %||% hex_binwidth(bins, scales)
    wt <- data$weight %||% rep(1L, nrow(data))
    out <- hexBinSummarise(data$x, data$y, wt, binwidth, sum)
    out$density <- as.vector(out$value / sum(out$value, na.rm = TRUE))
    out$ndensity <- out$density / max(out$density, na.rm = TRUE)
    out$count <- out$value
    out$ncount <- out$count / max(out$count, na.rm = TRUE)
    out$value <- NULL

    out
  }
)


load.libraries <- c('gridExtra', 'corrplot', 'GGally', 'ggplot2', 'e1071', 'tidyverse', 'ggthemes', 'reshape2')
# install.lib <- load.libraries[!load.libraries %in% installed.packages()]
# for(libs in install.lib) install.packages(libs, dependences = TRUE)

sapply(load.libraries, require, character = TRUE)

bank <- read.csv("./data/bank-full.csv",  sep = ";", na.strings = "unknown")

dim(bank)
str(bank)
summary(bank)

bank %>%
  group_by(y) %>%
  summarise(count = n())

## function to convert to month number

abb2m <- function(x) match(tolower(x), tolower(month.abb))
bank$month <- as.factor(abb2m(bank$month))

# Lines of report
#
# Plot missing data

plot_Missing <- function(data_in, title = NULL){
  temp_df <- as.data.frame(ifelse(is.na(data_in), 0, 1))
  temp_df <- temp_df[,order(colSums(temp_df))]
  data_temp <- expand.grid(list(x = 1:nrow(temp_df), y = colnames(temp_df)))
  data_temp$m <- as.vector(as.matrix(temp_df))
  data_temp <- data.frame(x = unlist(data_temp$x), y = unlist(data_temp$y), m = unlist(data_temp$m))
  ggplot(data_temp) +
    geom_tile(aes(x = x, y = y, fill = factor(m))) +
    scale_fill_manual(values = c("white", "black"), name = "Missing\n(0=Yes, 1=No)") +
    theme_light() +
    ylab("") + xlab("") +
    ggtitle(title)
}

plot_Missing(bank[,colSums(is.na(bank)) > 0])

#
#
#
#
#
# Show distributions of numerical variables.

'my_plots <- lapply(names(bank), function(var_x){
  p <-
    ggplot(bank) +
    aes_string(var_x) +
    labs(subtitle = var_x)

  if(is.numeric(bank[[var_x]])) {
    p <- p + geom_density()

  } else {
    p <- p + geom_bar() +
      coord_flip()
  }

  p  +
    theme(plot.title = element_text(hjust = .5),
    axis.ticks = element_blank(),
    panel.grid.minor.x = element_line(colour = "gray80")) +
    theme_fivethirtyeight()

})

plot_grid(plotlist = my_plots)'

# Show distribution of target variable by month.
# What percentage of potential clients accepted to subscribe?
# Check interesting traits of month age and job. Plot distributions.
# Plot density of each target group.

lngDataYes <-
  bank %>%
  filter(y == 'yes') %>%
  select_if(is.numeric) %>%
  gather()

lngDataNo <-
  bank %>%
  filter(y == 'no') %>%
  select_if(is.numeric) %>%
  gather()

ggplot(data = data.frame(), aes(x = value)) +
  geom_density(data = lngDataYes, fill = '#377EB8', color = '#377EB8', alpha = .3) +
  geom_density(data = lngDataNo, fill = '#E41A1C', color = '#E41A1C', alpha = .3) +
  facet_wrap( ~ key, scales = 'free') +
  labs(title = "Density",
       subtitle = 'of each numeric variable') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '',values = c(yes = '#377EB8', no = '#E41A1C')) +
  theme_fivethirtyeight()

## This variables, except day and age, have extremley longtailed distibutions

## plot non numerical variables similar to above

## random test for facet grid

bank %>%
  select_if(negate(is.numeric)) %>%
  ggplot() +
  geom_bar(aes(x = default, fill = y)) +
  facet_grid(housing ~ loan, scales = 'free') +
  coord_flip() +
  labs(title = "Distribution of target variable",
       subtitle = 'individual proportion over jobs') +
  theme(plot.title = element_text(hjust = .5),
        panel.grid.minor.x = element_line(colour = "gray80"),
        strip.text.y = element_blank()) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

# Plot scatter of explanatory variable against count of clients, color by target variable.

# Plot jitter of campaign calls against duration of call, color by tv.

bank %>%
  ggplot(aes(x = duration, y = campaign, color = y)) +
  geom_jitter() +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x= element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = 'blue', no ='red')) +
  theme_fivethirtyeight()

bank %>%
  ggplot(aes(x = campaign, y = duration, color = y)) +
  geom_jitter(stat = "identity") +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_color_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

# Plot histogram of count of clients by tv distributed by number of calls, add loess.

bank %>%
  ggplot(aes(x = campaign, fill = y)) +
  geom_histogram(position = 'dodge') +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight() +
  coord_cartesian(xlim = c(0, 20))

# Plot density of job status by tv.

# Plot distribution over loan status.

# Age pyramid?

bank %>%
  ggplot(aes(x = age)) +
  geom_density(aes(x = age, fill = y), alpha = .6) +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.text.x = element_text(angle = 90, hjust = 1),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

## older and younger clients tend to subscribe more often than midde aged.
## correlated with job?

## other stuff

bank %>%
  ggplot(aes(y)) +
  geom_bar(aes(fill = y)) +
  xlab("Frequency of target variable") +
  ylab("Absolute frequency") +
  theme_fivethirtyeight() +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

myCorr <- function(data, mapping, method = "loess", ...){
  p <- ggplot(data = data, mapping = mapping) +
    geom_jitter(size = .25) +
    geom_smooth(method = method, color = 'red', size = .5)
  p
}

# ggpairs(bank %>%
#           filter(y == 'yes') %>%
#           select(age, job, marital, education, default, y),
#           lower = list(continuous = myCorr)) + theme_fivethirtyeight(10)

bank %>%
  group_by(y, job) %>%
  summarise(count = n()) %>%
  arrange(count) %>%
  mutate(freq = count / sum(count)) %>%
  ggplot(aes(x = reorder(job, -freq), y = freq, fill = y)) +
  geom_bar(stat = "identity", width = .5) + coord_flip() +
  labs(title = "Distribution of target variable",
       subtitle = 'individual proportion over jobs') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

## interesting that retired and unemployed seem to have higher percentage of suscription

bank %>%
  group_by(y, loan) %>%
  summarise(count = n()) %>%
  arrange(count) %>%
  ggplot() +
  geom_bar(x = loan)

bank %>%
  ggplot(aes(x = balance)) +
  geom_density(aes(fill = y)) +
  labs(title = "Balance density",
       subtitle = 'analyzed over the full dataset') +
  theme(plot.title = element_text(hjust = .5),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight() +
  coord_cartesian(xlim = c(0, 50000))

##↑ seems that there is a higher tendency to reject a deposit when the balance is lower

bank %>%
  select(month, balance, y) %>%
  ggplot(aes(month, balance)) +
  geom_boxplot(aes(fill = y)) +
  coord_flip() +
  coord_cartesian(ylim = c(-3000, 11000)) +
  xlab('Campaign month') +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.text.x = element_text(angle = 90, hjust = 1),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

bank %>%
  group_by(month, y) %>%
  summarise(count = n()) %>%
  mutate(prop = count / sum(count)) %>%
  ggplot(aes(x = month, y = prop * 100)) +
  geom_bar(aes(fill = y), stat = 'identity') +
  # geom_density(aes(fill = y), alpha = .6) +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.text.x = element_text(angle = 90, hjust = 1),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_fill_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight()

## enormous quantity of outliers

correlations <-
  bank %>%
  select_if(is.numeric) %>%
  na.omit() %>%
  cor()

# correlations are very low
corrplot(correlations, method = "square")

bank %>%
  group_by(month, y) %>%
  summarise(count = n()) %>%
  ggplot() +
  geom_point(aes(x = month, y = count, color = y)) +
  labs(title = "placeholder",
       subtitle = 'placeholder') +
  theme(plot.title = element_text(hjust = .5),
        axis.text.x = element_text(angle = 90, hjust = 1),
        axis.ticks = element_blank(),
        panel.grid.minor.x = element_line(colour = "gray80")) +
  scale_color_manual(name = '', values = c(yes = '#377EB8', no ='#E41A1C')) +
  theme_fivethirtyeight() +
  coord_flip()


library(assertthat)
library(collections)
library(MASS)
library(NlcOptim)
library(mltools)
library(data.table)
library(tictoc)

source("data_structures.r")
source("estimators.r")
source("losses.r")

cross_validate<- function(data,synth_effect,
                          n_trees,max_depth=3,regularization = 2,
                          min_samples_leaf=1,nfolds=5){
  assert_that(is(data,"data"))
  assert_that(is.numeric(synth_effect))
  assert_integer(n_trees)
  assert_integer(max_depth)
  assert_that(is.numeric(regularization))
  assert_integer(min_samples_leaf)
  assert_integer(nfolds)

  #training_error<-vector(mode = "list",length = nfolds)
  test_error<-vector(mode = "list",length = nfolds)

  n_fold_elements<-data$N*(1-1/nfolds)

  for(i in 1:nfolds){
    #There are two nfolds parameters in this function.The nfold of the
    #function below states that the data will be splitted into [nfolds] sets.
    #The first of them has N/nfolds elements, the second - N*(nfolds-1)/nfolds
    #In this loop we do nfold(from my function) such splits and always take the
    #first fold for training and the second for test.

    fold <- folds(data$W*1,nfolds = data$N/n_fold_elements,stratified = TRUE)
    tr<-which(fold==1)
    te<-which(fold==2)

    F<-constrained_boost(data,synth_effect,
                         n_trees,max_depth,regularization,
                         min_samples_leaf,tr)

    #training_error[[i]]<-vector(mode ="numeric",length = (n_trees+1))
    test_error[[i]]<-vector(mode ="numeric",length = (n_trees+1))

    for(j in 1:(n_trees+1)){
      #training_error[[i]][[j]]<-sum(evaluate(data$Y[tr],F[[j]]$observed$`TRUE`[tr]))/length(tr)
      test_error[[i]][[j]]<-sum(evaluate(data$Y[te],F[[j]]$observed$`TRUE`[te]))/length(te)
    }
  }
  #return (list(training_error=training_error,test_error=test_error))
  return(test_error)
}

constrained_boost <- function(data,synth_effect,
                              n_trees,max_depth=3,regularization = 2,
                              min_samples_leaf=1,tr=TRUE){

  tic("constrained boosting")
  assert_that(is(data,"data"))
  assert_that(is.numeric(synth_effect))
  assert_integer(n_trees)
  assert_integer(max_depth)
  assert_that(is.numeric(regularization))
  assert_integer(min_samples_leaf)
  assert_that(is.logical(tr)||is.numeric(tr))

  const_pair = fit_const_pair(get_index(data,tr),synth_effect)

  F<-vector(mode="list",length = n_trees)
  F[[1]]<-predict_counterfactuals(const_pair,data)

  residuals<-obs_gradient(data$Y,data$W,F[[1]],ind = tr)

  for (i in 1:n_trees){
    tree_pair<-fit_tree_pair(get_index(data,tr),residuals,max_depth,min_samples_leaf)

    leaf_index_dict<-build_leaf_index(tree_pair,data$X)

    v<-step_search(data,F[[i]],leaf_index_dict,regularization)

    F[[i+1]]<-add_counterfactuals(F[[i]],predict_counterfactuals(tree_pair,data,v))

    residuals<-obs_gradient(data$Y,data$W,F[[i+1]],ind=tr)
  }
  toc()
  return(F)
}

fit_tree_pair<-function(data,residuals,max_depth,min_samples_leaf){
  assert_that(is(data,"data"))
  assert_that(is(residuals,"numeric_treatment_dictionary"))
  assert_integer(max_depth)
  assert_integer(min_samples_leaf)

  tree_pair<-TreatmentDictionary(0,0)
  for(tr in names(tree_pair)){
    tr_logical<-as.logical(tr)
    tree<-fit_regression_tree(get_elements_by_treatment(data$X,data$W,t=tr_logical),
                              residuals[[tr]],
                              min_samples_leaf=min_samples_leaf,
                              max_depth=max_depth)
    tree_pair[[tr]]<-tree
  }
  tree_pair
}

fit_const_pair <- function(data,synth_effect){
  assert_that(is(data,"data"))
  assert_that(is.numeric(synth_effect))

  mean <- mean(data$Y)
  dict<-TreatmentDictionary(0,0)
  for(tr in names(dict)){
    tr_logical<-as.logical(tr)
    value<-mean+(2*tr_logical-1)*synth_effect*data$N_treated[[tr]]/data$N
    dict[[tr]]<-Leaf(value)
  }
  dict
}

predict_counterfactuals <- function(estimator_pair, data,v=NULL){
  assert_that(is(estimator_pair,"treatment_dictionary"))
  assert_that(is(estimator_pair$`TRUE`,"leaf")||is(estimator_pair$`TRUE`,"node"))
  assert_that(is(data,"data"))
  assert_that(is.null(v)||is(v,"treatment_dictionary"))

  if(is.null(v)){
    return (Counterfactuals(predict_treated_dictionary(estimator_pair,data$X),data$W))
  }
  else{
    treated<-NumericTreatmentDictionary(vector(mode = "numeric",length = length(data$W)),
                                        vector(mode = "numeric",length = length(data$W)))
    for(tr in names(estimator_pair)){
      assignments<-predict_rt_matrix(estimator_pair[[tr]],data$X,TRUE)

      for(i in 1:length(assignments)){
        treated[[tr]][[i]]<-v[[tr]]$get(as.character(assignments[i]))
      }
    }
    return (Counterfactuals(treated,data$W))
  }

}

obs_gradient <- function(Y,W,F,ind=TRUE){
  assert_that(is.numeric(Y))
  assert_that(is.logical(W))
  assert_that(is(F,"counterfactuals"))
  assert_that(is.logical(ind)||is.numeric(ind))

  dict<-NumericTreatmentDictionary(0,0)
  for(tr in names(dict)){
    tr_logical<-as.logical(tr)

    value<- neg_grad(get_elements_by_treatment(Y[ind],W[ind],tr_logical),
                     get_elements_by_treatment(get_index(F,ind),W[ind],tr_logical,TRUE))
    dict[[tr]]<-value
  }
  dict
}

build_leaf_index<-function(tree_pair,X){
  assert_that(is(tree_pair,"treatment_dictionary"))
  assert_that(is.matrix(X)&&is.numeric(X))

  assignments<-predict_treated_dictionary(tree_pair,X,TRUE)

  assignment_index_treatment_dictionary <-TreatmentDictionary(Dict$new(),Dict$new())
  for(tr in names(assignments)){
    #Unique leaves' indeces for treatment
    leaves<-unique(assignments[[tr]])
    for(l in leaves){
      indeces<-which(assignments[[tr]]==l)
      assignment_index_treatment_dictionary[[tr]]$set(as.character(l),indeces)
    }
  }
  assignment_index_treatment_dictionary
}

step_search<-function(data, F, assignment_indeces_dict,regularization,ind=TRUE){
  assert_that(is(data,"data"))
  assert_that(is(F,"counterfactuals"))
  assert_that(is(assignment_indeces_dict,"treatment_dictionary"))
  assert_that(is.numeric(regularization))
  assert_that(is.logical(ind)||is.numeric(ind))

  F1 <- F$treated$`TRUE`
  F0 <- F$treated$`FALSE`
  Y <- data$Y
  W <- data$W

  idx1<-assignment_indeces_dict$`TRUE`
  idx0<-assignment_indeces_dict$`FALSE`

  filter1<-which(W==1)
  filter0<-which(W==0)

  if(!is.logical(ind)){
    filter1<-intersect(filter1,ind)
    filter0<-intersect(filter0,ind)
  }

  idx1obs <- filter_index(idx1, filter1)
  idx0obs <-filter_index(idx0, filter0)

  n1<-length(idx1$keys())
  n0<-length(idx0$keys())

  v<-rep(0.1,times=n1+n0)

  objfun=function(x){
    fun<-0
    for (i in 1:length(idx1obs$keys())){
      k<-idx1obs$keys()[i]
      sum_loss_key<-0
      for(j in idx1obs$get(k)){
        sum_loss_key<-sum_loss_key+(Y[j] - F1[j] - x[i])^2
      }
      fun<- fun+sum_loss_key
    }
    for (i in 1:length(idx0obs$keys())){
      k<-idx0obs$keys()[i]
      sum_loss_key<-0
      for(j in idx0obs$get(k)){
        sum_loss_key<-sum_loss_key+(Y[j] - F0[j] - x[n1+i])^2
      }
      fun<- fun+sum_loss_key
    }
    for (i in 1:length(idx1$keys())){
      k<-idx1$keys()[i]
      sum_loss_key<-0
      for(j in idx1$get(k)){
        sum_loss_key<-sum_loss_key + regularization*x[i]^2
      }
      fun<- fun+sum_loss_key
    }
    for (i in 1:length(idx0$keys())){
      k<-idx0$keys()[i]
      sum_loss_key<-0
      for(j in idx0$get(k)){
        sum_loss_key<-sum_loss_key + regularization*x[n1+i]^2
      }
      fun<- fun+sum_loss_key
    }

    return(fun)
  }

  confun=function(x){
    fun1<-0
    for (i in 1:length(idx1$keys())){
      k<-idx1$keys()[i]
      sum_loss_key<-0
      for(j in idx1$get(k)){
        sum_loss_key<-sum_loss_key + x[i]
      }
      fun1<- fun1+sum_loss_key
    }
    fun0<-0
    for (i in 1:length(idx0$keys())){
      k<-idx0$keys()[i]
      sum_loss_key<-0
      for(j in idx0$get(k)){
        sum_loss_key<-sum_loss_key +x[n1+i]
      }
      fun0<- fun0+sum_loss_key
    }

    f=NULL
    f=rbind(f,fun1-fun0)
    return(list(ceq=f,c=NULL))
  }

  v<-solnl(v,objfun=objfun,confun=confun)$par[,1]
  vopt1<-v[1:n1]
  vopt0<-v[(n1+1):(n1+n0)]

  dict<-TreatmentDictionary(0,0)
  for(tr in names(dict)){
    dict[[tr]]<-Dict$new()
    if(as.logical(tr)){
      for (i in 1:length(idx1$keys())){
        k<-idx1$keys()[i]
        dict[[tr]]$set(k,vopt1[i])
      }
    }
    else{
      for (i in 1:length(idx0$keys())){
        k<-idx0$keys()[i]
        dict[[tr]]$set(k,vopt0[i])
      }
    }
  }
  dict
}

filter_index<-function(indices, include_subjects){
  filtered_index = Dict$new()
  for (k in indices$keys()){
    v_in <- intersect(indices$get(k),include_subjects)
    filtered_index$set(k,v_in)
  }
  filtered_index
}

get_n_trees_min_error<-function(test_error){
  error_matrix<-NULL

  for(i in 1:length(test_error)){
    error_matrix<-rbind(error_matrix,test_error[[i]])
  }

  means<-vector(mode = "numeric",length = length(error_matrix[1,]))

  for(i in 1:length(means)){
    means[[i]]<-mean(error_matrix[,i])
  }

  return(which.min(means))
}


library(data.table)
library(stringr)

recursive_mean <- function(lastmean, k, x){
  return(((k-1)*lastmean+x)/k)
}

cmd_args <- commandArgs()
CURRENTDIR <- dirname(regmatches(cmd_args, regexpr("(?<=^--file=).+", cmd_args, perl=TRUE)))
ROOTDIR <- dirname(dirname(CURRENTDIR))
OUTDIR <- file.path(ROOTDIR, "output", "two_planets")

# OUTDIR <- "../../output/two_planets"

SAMPLERDIRS <- c("MALA", "AM", "SMMALA", "GAMC")

nsamplerdirs <- length(SAMPLERDIRS)

true_pars <- c(
  3.71357,
  3.43399,
  0.141421,
  0.141421,
  1.5708,
  4.40428,
  3.43399,
  0.141421,
  0.141421,
  1.5708,
  0.0
)

nchains <- 10
nmcmc <- 110000
nburnin <- 10000
npostburnin <- nmcmc-nburnin

nmeans <- npostburnin

ci <- c(3, 2, 4, 4)
pi <- 2

submeans <- matrix(data=NA, nrow=nmeans, ncol=nsamplerdirs)
curmeans <- rep(0, nsamplerdirs)

for (j in 1:nsamplerdirs) {
  chains <- t(fread(
    file.path(OUTDIR, SAMPLERDIRS[j], paste("chain", str_pad(ci[j], 2, pad="0"), ".csv", sep="")), sep=",", header=FALSE
  ))

  for (i in 1:nmeans) {
    curmeans[j] <- recursive_mean(curmeans[j], i, chains[i, pi])
    submeans[i, j] <- curmeans[j]
  }
}

# apply(submeans, 2, mean)-true_pars[pi]

cols <- c("green", "blue", "orange", "red")

pdf(file=file.path(OUTDIR, "rv_two_planets_meanplot.pdf"), width=10, height=6)

oldpar <- par(no.readonly=TRUE)

par(fig=c(0, 1, 0, 1), mar=c(2.25, 4, 3.5, 1)+0.1)

plot(
  1:nmeans,
  submeans[, 1],
  type="l",
  ylim=c(3.4, 3.5),
  col=cols[1],
  lwd=2,
  xlab="",
  ylab="",
  cex.axis=1.8,
  cex.lab=1.7,
  yaxt="n"
)

axis(
  2,
  at=seq(3.4, 3.5, by=0.05),
  labels=seq(3.4, 3.5, by=0.05),
  cex.axis=1.8,
  las=1
)

lines(
  1:nmeans,
  submeans[, 2],
  type="l",
  col=cols[2],
  lwd=2
)

lines(
  1:nmeans,
  submeans[, 3],
  type="l",
  col=cols[3],
  lwd=2
)

lines(
  1:nmeans,
  submeans[, 4],
  type="l",
  col=cols[4],
  lwd=2
)

# abline(h=true_pars[pi], lwd=2, col="black")

par(fig=c(0, 1, 0.89, 1), mar=c(0, 0, 0, 0), new=TRUE)

plot.new()

legend(
  "center",
  SAMPLERDIRS,
  lty=c(1, 1, 1, 1),
  lwd=c(5, 5, 5, 5),
  col=cols,
  cex=1.5,
  bty="n",
  text.width=0.125,
  ncol=4
)

par(oldpar)

dev.off()


library(stringr)
library(reshape2)
library(plyr)
source("tidy-data/data/xtable.r")
source("tidy-data/data/read-fwf.r")
options(stringsAsFactors = FALSE)

# Define format for fixed width file
cols <- data.frame(
  name =  c("id", "year", "month", "element"),
  start = c(1,     12,    16,      18),
  end =   c(11,    15,    17,      21))

names <- str_c(c("value", "mflag", "qflag", "sflag"), rep(1:31, each = 4), sep = "_")
starts <- cumsum(c(22, rep(c(5, 1, 1, 1), 31)))
starts <- starts[-length(starts)]
ends <- c(starts[-1], starts[length(starts)] + 1) - 1

values <- data.frame(name = names, start = starts, end = ends)
cols <- rbind(cols, values)

# Load data and subset to small example
raw <- read.fwf2("tidy-data/data/weather.txt",  cols)
raw <- subset(raw, year == 2010 & element %in% c("TMIN", "TMAX"))
raw <- raw[, c(1:4, which(str_detect(names(raw), "value")))]
raw$id <- str_c(str_sub(raw$id, 1, 2), str_sub(raw$id, -5, -1))

names(raw)[-(1:4)] <- str_c("d", 1:31)
raw[raw == -9999] <- NA
raw[-(1:4)] <- raw[-(1:4)] / 10
rownames(raw) <- NULL
raw$element <- tolower(raw$element)

xtable(raw[1:10, 1:12], file = "tidy-data/data/weather-raw.tex", digits = 1)

# Melt and tidy

clean1 <- melt(raw, id = 1:4, na.rm = T)
clean1$day <- as.integer(str_replace(clean1$variable, "d", ""))
clean1$date <- as.Date(ISOdate(clean1$year, clean1$month, clean1$day))

clean1 <- clean1[c("id", "date", "element", "value")]
clean1 <- arrange(clean1, date, element)
clean1$date <- as.character(clean1$date) # work around xtable bug
xtable(clean1[1:10, ], file = "tidy-data/data/weather-clean-1.tex", digits = 1)

# Cast

clean2 <- dcast(clean1, ... ~ element)
xtable(clean2[1:10, ], file = "tidy-data/data/weather-clean-2.tex", digits = 1)


#load library
library(MethComp)

#import data
df<-read.csv("cani.csv", header = TRUE, sep = ";", dec= ",", stringsAsFactors = FALSE)
#
# #ignore columns
# df<-df[,1:22]
#
# head(df)
# str(df)

#delete record number 99
df<-df[-c(393:396),]
df_Meth<-Meth(df, 2,1,3,20)
df_Meth$meth<-factor(df_Meth$meth, levels=c("Procyte Idexx", "DF-50"))
summary(df_Meth)
plot(df_Meth)

BA.plot(df_Meth, repl.conn = FALSE)

PBdf <- PBreg(df_Meth)
# par(mfrow=c(1,2))
plot(PBdf)

plot(df[,-c(1,2,3, 13:25)])

#convert data to Meth object type and save pdf with meth plot

pdf()

for(i in c(4:25)){

df_Meth<-Meth(df, 2,1,3,i)
df_Meth$meth<-factor(df_Meth$meth, levels=c("Procyte Idexx", "DF-50"))
par(mfrow=c(1,2))
plot(df_Meth)
title(main=list(colnames(df[i]), col="red"), line=3, adj=0.45)
# reg<-PBreg(df_Meth)
# print(reg)

print(i)
}
dev.off()

layout(matrix(c(1,1,2),nrow=1,ncol=3, byrow=TRUE))
par(ps=16)
#Bland-Altman plot
mycol<-rgb(193,193,193, maxColorValue = 255, alpha=255/3)
p<-BA.plot(df_Meth, repl.conn=FALSE, meth.names = TRUE, eqax = TRUE, pch.points = 21, bg= mycol, main=2)

par(mfrow=c(2,1))
#passing  bablok regression

reg<-PBreg(df_Meth)
print(reg)

plot(reg,subtype= 1, main=1)