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

#------------------------------------------ Forest risk mapping models: Fig 1D-F code - 05/11/21 ----------------------------------------------

#-----------------------------------------------------------------------------------------------------------------------------------------------------------
# Author: William Anderegg (anderegg@utah.edu), University of Utah

library(rworldmap)
library(MASS)
library(SDMTools)
library(raster)
library(RNetCDF)
library(scales)
library(car)
library(RColorBrewer)
library(betareg)
library(pgirmess)
library(geosphere)
library(ape)

#---------------------------------- Pull in FIA mort data -------------------------------------
dir <- ""
fianew <- read.csv(paste(dir, "FIA-TerraClim-Wide-v17-04-14-2021.csv", sep = ""), header = T)
fialonga <- read.csv(paste(dir, "FIA-TerraClim-Long-1990.1999-v17-04-14-2021.csv", sep = ""), header = T)
fialongb <- read.csv(paste(dir, "FIA-TerraClim-Long-2000.2009-v17-04-14-2021.csv", sep = ""), header = T)
fialongc <- read.csv(paste(dir, "FIA-TerraClim-Long-2010.2019-v17-04-14-2021.csv", sep = ""), header = T)
FTpred_ins <- read.csv(paste(dir, "Insect_ForTypToPredict_04-21-2021.csv", sep = ""), header = T)
FTpred_drt <- read.csv(paste(dir, "Drought_ForTypToPredict_04-21-2021.csv", sep = ""), header = T)
# FTs that meet the minimum 20 mortality threshold and cross-validated AUC>0.6

# Processing projection and a useful logical function
proj2 <- CRS("+proj=longlat +datum=WGS84 +ellps=WGS84")
"%!in%" <- function(x, y) !("%in%"(x, y))

# Get functional trait data from Trugman et al. (2020) PNAS
p50a <- open.nc(paste(dir, "CWM_P50_025Deg.nc", sep = ""))
p50m <- var.get.nc(p50a, "CWM_P50")
dim(p50m)
p50r <- var.get.nc(p50a, "mean_range_within_site")
dim(p50r)
lat.tr <- var.get.nc(p50a, "lat")
dim(lat.tr)
lon.tr <- var.get.nc(p50a, "lon")
dim(lon.tr)

# Screen FIA data to remove fire, human disturbance, timber cutting, CONDPROP>0.3
fianew1 <- fianew[which(fianew[, 8] >= 0.3 & fianew[, 59] != "True" & fianew[, 63] != "True" & fianew[, 71] != "True" & fianew[, 18] > 1), ]

# Extract trait data for all FIA CONDs
traitdata <- array(dim = c(102453, 2))
for (i in 1:102453) {
  lati <- fianew1[i, 1]
  loni <- fianew1[i, 2]
  latcell <- which.min(abs(lati - lat.tr))
  loncell <- which.min(abs(loni - lon.tr))
  traitdata[i, 1] <- p50m[latcell, loncell]
  traitdata[i, 2] <- p50r[latcell, loncell]
}

traitlong <- array(dim = c(498410, 2))
for (i in 1:498410) {
  lati <- fialonga[i, 1]
  loni <- fialonga[i, 2]
  latcell <- which.min(abs(lati - lat.tr))
  loncell <- which.min(abs(loni - lon.tr))
  traitlong[i, 1] <- p50m[latcell, loncell]
  traitlong[i, 2] <- p50r[latcell, loncell]
}

#------------------------------------------- Processing FIA data --------------------------------------------
# For Obs, build a clean data-frame with dependent variable, predictor variables, FORTYP, lat/lon
fianew1[is.na(fianew1[, 23]) == "TRUE", 23] <- 0 # Impute 0s where no mort was measured
fianew1[is.na(fianew1[, 27]) == "TRUE", 27] <- 0 # Impute 0s where no mort was measured for insect mort
fianew1[which(fianew1[, 151] < -16), 151] <- -16 # Set lower PDSI bounds
fianew1[which(fianew1[, 151] > 16), 151] <- 16 # Set upper PDSI bounds
fianew1[which(fianew1[, 152] < -16), 152] <- -16 # Set lower PDSI bounds
fianew1[which(fianew1[, 152] > 16), 152] <- 16 # Set upper PDSI bounds
fianew1[which(fianew1[, 153] < -16), 153] <- -16 # Set lower PDSI bounds
fianew1[which(fianew1[, 153] > 16), 153] <- 16 # Set upper PDSI bounds

# Combine all cond data and variables into a dataframe with columns: MortDrt, MortIns, 18 clim vars, 2 age vars, 8 traits, lat,lon, FORTYP
fia.newm <- as.data.frame(cbind(
  (fianew1[, 23] / fianew1[, 18]) / (fianew1[, 15] - fianew1[, 14]), fianew1[, 27] * (fianew1[, 23] / fianew1[, 18]) / (fianew1[, 15] - fianew1[, 14]),
  scale(fianew1[, 127]), scale(fianew1[, 128]), scale(fianew1[, 129]), scale(fianew1[, 139]), scale(fianew1[, 140]), scale(fianew1[, 141]),
  scale(fianew1[, 151]), scale(fianew1[, 152]), scale(fianew1[, 153]), scale(fianew1[, 163]), scale(fianew1[, 164]), scale(fianew1[, 165]),
  scale(fianew1[, 175]), scale(fianew1[, 176]), scale(fianew1[, 177]), scale(fianew1[, 187]), scale(fianew1[, 188]), scale(fianew1[, 189]),
  scale(fianew1[, 3]^2), scale(fianew1[, 3]),
  scale(traitdata[, 1]), scale(traitdata[, 2]),
  (fianew1[, 19] - fianew1[, 18]), (fianew1[, 15] - fianew1[, 14]),
  fianew1[, 1], fianew1[, 2], fianew1[, 4], fianew1[, 18]
))

# Second, screen out NaN values that will mess up the model
fia2 <- fia.newm[which(fia.newm[, 1] != "NaN" & fia.newm[, 5] != "NaN" & fia.newm[, 3] != "NaN" & fia.newm[, 9] != "NaN" & fia.newm[, 29] != "NaN" & fia.newm[, 21] != "NaN" & fia.newm[, 23] != "NaN" & fia.newm[, 1] < 1.01 & is.na(fia.newm[, 2]) == "FALSE"), ]

# Third, scale up to 0.25x0.25 degree for model fitting
dummy <- raster(ncols = 244, nrows = 100, xmn = -126, xmx = -65, ymn = 25, ymx = 50, crs = proj2) # 0.25 degree dummy raster
fia.g1 <- array(dim = c(0, 30))
for (i in 1:112) {
  ftype <- unique(fia2[, 29])[i]
  d1 <- data.frame(fia2[which(fia2[, 29] == ftype), ])
  v1 <- rasterize(d1[, c(28, 27)], dummy, d1[, 1], fun = mean, background = NA, mask = FALSE, na.rm = T)
  v1i <- rasterize(d1[, c(28, 27)], dummy, d1[, 2], fun = mean, background = NA, mask = FALSE, na.rm = T)
  v1b <- rasterToPoints(v1)
  v1bi <- (as.matrix(as.vector(v1i)))[is.na(as.matrix(as.vector(v1i))) == "FALSE"]
  stack <- array(dim = c(length(v1bi), 0))
  for (j in 1:24) {
    varh <- rasterize(d1[, c(28, 27)], dummy, d1[, (j + 2)], fun = mean, background = NA, mask = FALSE, na.rm = T)
    var1 <- (as.matrix(as.vector(varh)))[is.na(as.matrix(as.vector(varh))) == "FALSE"]
    stack <- cbind(stack, var1)
  }
  ba <- rasterize(d1[, c(28, 27)], dummy, d1[, 30], fun = sum, background = NA, mask = FALSE, na.rm = T)
  ba0 <- (as.matrix(as.vector(ba)))[is.na(as.matrix(as.vector(ba))) == "FALSE"]
  fia.out <- cbind(v1b[, 3], v1bi, stack, v1b[, 1:2], rep(ftype, length(v1bi)), ba0)
  fia.g1 <- rbind(fia.g1, fia.out)
}

#------------------------------------------- Fit mortality functions --------------------------------------------------
# GLM control parameters - bump up maxit to help with convergence
gm.ctl <- glm.control(epsilon = 1e-8, maxit = 500, trace = FALSE)

mod_insect_proj <- function(fitdata, minthres, ftypcol, preddata) {
  colnames(preddata) <- c("V1", "V2", "V3", "V4", "V5", "V6", "V7", "V8", "V9", "V10", "V11", "V12", "V13")
  colnames(fitdata) <- c("V1", "V2", "V3", "V4", "V5", "V6", "V7", "V8", "V9", "V10")
  mod1 <- array(dim = c(0, 7))
  for (i in 1:112) {
    ftype <- unique(fitdata[, ftypcol])[i]
    non_zero <- ifelse(fitdata[which(fitdata[, ftypcol] == ftype), 1] > 0, 1, 0)
    d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero) # Historical mortality data
    preddata1 <- data.frame(preddata[which(preddata[, 10] == ftype), ]) # Projection data
    # Set up the output matrix
    drt.out <- array(dim = c(dim(preddata1)[1], 7))
    drt.out[, 1] <- ftype # FORTYPCD
    drt.out[, 2] <- preddata1[, 12] # lat
    drt.out[, 3] <- preddata1[, 11] # lon
    drt.out[, 7] <- preddata1[, 13] # BAlive0 for the FORTYP

    # If FORTYP has non-meaningful historical model, set future projections == mean historical
    if (ftype %!in% FTpred_ins[, 1]) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Predicted mort = mean historical where Nmort<minthres
    if (ftype %!in% FTpred_ins[, 1]) mod1 <- rbind(mod1, drt.out)
    if (ftype %!in% FTpred_ins[, 1]) print(paste("FORTYP prediction skipped due inadequate historical model: ", ftype, sep = ""))
    if (ftype %!in% FTpred_ins[, 1]) next

    # If FORTYP has Nmort<minthres, set future projections == mean historical
    if (sum(d1[, 11], na.rm = T) < minthres) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Pred mort = mean hist where Nmort<minthres
    if (sum(d1[, 11], na.rm = T) < minthres) mod1 <- rbind(mod1, drt.out)
    if (sum(d1[, 11], na.rm = T) < minthres) print(paste("FORTYP prediction skipped due Nmort<minthres: ", ftype, sep = ""))
    if (sum(d1[, 11], na.rm = T) < minthres) next

    # For FORTYPs with adequate historical models
    if (sum(non_zero, na.rm = T) > 0) d1 <- d1[which(d1[, 1] < quantile(d1[, 1], 0.995)), ]
    if (dim(d1)[1] == 0) d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero)

    # Construct historical mortality model
    m1 <- glm(non_zero ~ V2 + V3 + V4 + V5 + V6 + V7 + V8, data = d1, family = binomial(link = logit), control = gm.ctl)
    m2 <- betareg(V1 ~ V2 + V3 + V4 + V5 + V6 + V7 + V8, data = subset(d1, non_zero == 1))

    # Predict with future climate projections
    pred1 <- predict(m1, newdata = preddata1[, 2:8], se = TRUE, type = "response")
    pred2 <- predict(m2, newdata = preddata1[, 2:8], se = TRUE, type = "response")
    pred1a <- ifelse(pred1$fit >= 0.5, 1, 0)
    pred3 <- pred2 * pred1a
    pred4 <- pred2 * pred1$fit
    # 		Store model output
    drt.out[, 5] <- pred3
    drt.out[, 6] <- pred4 # Predicted (expected value) mort frac
    mod1 <- rbind(mod1, drt.out)
  }
  print(paste("FORTYPs projected: ", length(unique(mod1[which(mod1[, 5] != "NA"), 1])), sep = ""))
  print(paste("FORTYPs completed: ", i, sep = ""))
  return(mod1)
}

insect_histmodel <- function(fitdata, minthres, ftypcol) {
  colnames(fitdata) <- c("V1", "V2", "V3", "V4", "V5", "V6", "V7", "V8", "V9", "V10", "V11", "V12", "V13")
  mod1 <- array(dim = c(0, 7))
  for (i in 1:112) {
    ftype <- unique(fitdata[, ftypcol])[i]
    non_zero <- ifelse(fitdata[which(fitdata[, ftypcol] == ftype), 1] > 0, 1, 0)
    d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero)

    drt.out <- array(dim = c(dim(d1)[1], 7))
    drt.out[, 1] <- d1[, ftypcol] # FORTYPCD
    drt.out[, 2] <- d1[, 12] # lat
    drt.out[, 3] <- d1[, 11] # lon
    drt.out[, 4] <- d1[, 1] # Obs mort frac
    drt.out[, 7] <- d1[, 13] # BAlive0 for the FORTYP

    # If FORTYP has non-meaningful historical model, set future projections == mean historical
    if (ftype %!in% FTpred_ins[, 1]) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Predicted mort = mean historical where Nmort<minthres
    if (ftype %!in% FTpred_ins[, 1]) mod1 <- rbind(mod1, drt.out)
    if (ftype %!in% FTpred_ins[, 1]) print(paste("FORTYP prediction skipped due inadequate historical model: ", ftype, sep = ""))
    if (ftype %!in% FTpred_ins[, 1]) next

    # If FORTYP has Nmort<minthres, set future projections == mean historical
    if (sum(d1[, 14], na.rm = T) < minthres) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Predicted mort = mean historical where Nmort<minthres
    if (sum(d1[, 14], na.rm = T) < minthres) mod1 <- rbind(mod1, drt.out)
    if (sum(d1[, 14], na.rm = T) < minthres) print(paste("FORTYP prediction skipped due Nmort<minthres: ", ftype, sep = ""))
    if (sum(d1[, 14], na.rm = T) < minthres) next

    # For FORTYPs with adequate historical models, remove outliers
    if (sum(non_zero, na.rm = T) > 0) d1 <- d1[which(d1[, 1] < quantile(d1[, 1], 0.995)), ]
    if (dim(d1)[1] == 0) d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero)

    m1 <- glm(non_zero ~ V2 + V3 + V4 + V5 + V6 + V7 + V8, data = d1, family = binomial(link = logit), control = gm.ctl)
    m2 <- betareg(V1 ~ V2 + V3 + V4 + V5 + V6 + V7 + V8, data = subset(d1, non_zero == 1))

    pred1 <- predict(m1, newdata = d1[, 2:8], se = TRUE, type = "response")
    pred2 <- predict(m2, newdata = d1[, 2:8], se = TRUE, type = "response")
    pred1a <- ifelse(pred1$fit >= 0.5, 1, 0)
    pred3 <- pred1$fit
    pred4 <- pred2 * pred1$fit

    if (dim(drt.out)[1] != dim(d1)[1]) drt.out <- array(dim = c(dim(d1)[1], 7))
    drt.out[, 1] <- d1[, ftypcol] # FORTYPCD
    drt.out[, 2] <- d1[, 12] # lat
    drt.out[, 3] <- d1[, 11] # lon
    drt.out[, 4] <- d1[, 1] # Obs mort frac
    drt.out[, 7] <- d1[, 13] # BAlive0 for the FORTYP

    drt.out[, 5] <- pred3
    drt.out[, 6] <- pred4 # Predicted (expected value) mort frac
    mod1 <- rbind(mod1, drt.out)
  }
  print(paste("FORTYPs modeled: ", length(unique(mod1[which(mod1[, 5] != "NA"), 1])), sep = ""))
  print(paste("FORTYPs completed: ", i, sep = ""))
  return(mod1)
}

#----------------- FIAlong climate data preprocessing function
prepclim <- function(histclim, climdata, histcols, futcols) {
  climdata[which(climdata[, 6] < -16), 6] <- -16 # Set lower PDSI bounds
  climdata[which(climdata[, 6] > 16), 6] <- 16 # Set upper PDSI bounds
  climdata[which(climdata[, 13] < -16), 13] <- -16 # Set lower PDSI bounds
  climdata[which(climdata[, 13] > 16), 13] <- 16 # Set upper PDSI bounds
  histclim[which(histclim[, 25] < -16), 25] <- -16
  histclim[which(histclim[, 25] > 16), 25] <- 16
  histclim[which(histclim[, 26] < -16), 26] <- -16
  histclim[which(histclim[, 26] > 16), 26] <- 16
  print(histclim[1, histcols])
  print(climdata[1, futcols])
  # Get historical mean and SD for drought models
  hist1 <- array(dim = c(6, 2))
  for (i in 1:6) {
    hist1[i, 1] <- mean(histclim[, histcols[i]], na.rm = T)
    hist1[i, 2] <- sd(histclim[, histcols[i]], na.rm = T)
  }
  fia.futa <- as.data.frame(cbind(rep(1, 498410), (climdata[, futcols[1]] - hist1[1, 1]) / hist1[1, 2], (climdata[, futcols[2]] - hist1[2, 1]) / hist1[2, 2], (climdata[, futcols[3]] - hist1[3, 1]) / hist1[3, 2], (climdata[, futcols[4]] - hist1[4, 1]) / hist1[4, 2], (climdata[, futcols[5]] - hist1[5, 1]) / hist1[5, 2], (climdata[, futcols[6]] - hist1[6, 1]) / hist1[6, 2], scale(fialong[, 3]), scale(traitlong[, 2]), fialong[, 7], fialong[, 1], fialong[, 2], fialong[, 4]))
  finalclim <- fia.futa[which(fia.futa[, 1] != "NaN" & fia.futa[, 5] != "NaN" & fia.futa[, 3] != "NaN" & fia.futa[, 7] != "NaN" & fia.futa[, 10] != "NaN" & fia.futa[, 8] != "NaN" & is.na(fia.futa[, 2]) == "FALSE"), ]
  return(finalclim)
}

# Function when both inputs are FIAlong
prepclim1 <- function(histclim, climdata, histcols, futcols) {
  climdata[which(climdata[, 25] < -16), 25] <- -16 # Set lower PDSI bounds
  climdata[which(climdata[, 25] > 16), 25] <- 16 # Set upper PDSI bounds
  climdata[which(climdata[, 26] < -16), 26] <- -16 # Set lower PDSI bounds
  climdata[which(climdata[, 26] > 16), 26] <- 16 # Set upper PDSI bounds
  histclim[which(histclim[, 25] < -16), 25] <- -16
  histclim[which(histclim[, 25] > 16), 25] <- 16
  histclim[which(histclim[, 26] < -16), 26] <- -16
  histclim[which(histclim[, 26] > 16), 26] <- 16
  print(histclim[1, histcols])
  print(climdata[1, futcols])
  # Get historical mean and SD for drought models
  hist1 <- array(dim = c(6, 2))
  for (i in 1:6) {
    hist1[i, 1] <- mean(histclim[, histcols[i]], na.rm = T)
    hist1[i, 2] <- sd(histclim[, histcols[i]], na.rm = T)
  }
  fia.futa <- as.data.frame(cbind(rep(1, 498410), (climdata[, futcols[1]] - hist1[1, 1]) / hist1[1, 2], (climdata[, futcols[2]] - hist1[2, 1]) / hist1[2, 2], (climdata[, futcols[3]] - hist1[3, 1]) / hist1[3, 2], (climdata[, futcols[4]] - hist1[4, 1]) / hist1[4, 2], (climdata[, futcols[5]] - hist1[5, 1]) / hist1[5, 2], (climdata[, futcols[6]] - hist1[6, 1]) / hist1[6, 2], scale(fialonga[, 3]), scale(traitlong[, 2]), fialonga[, 7], fialonga[, 1], fialonga[, 2], fialonga[, 4]))
  finalclim <- fia.futa[which(fia.futa[, 1] != "NaN" & fia.futa[, 5] != "NaN" & fia.futa[, 3] != "NaN" & fia.futa[, 7] != "NaN" & fia.futa[, 10] != "NaN" & fia.futa[, 8] != "NaN" & is.na(fia.futa[, 2]) == "FALSE"), ]
  return(finalclim)
}

#-------------------- Function to rasterize and weight model mortality projections by observed/current BA
weight.proj1 <- function(raster1, grid, mortcol) {
  mort.bafi <- rasterize(raster1[which(raster1[, 1] == unique(raster1[, 1])[1]), c(2, 3)], grid, raster1[which(raster1[, 1] == unique(raster1[, 1])[1]), 7], fun = mean, background = NA, mask = FALSE, na.rm = TRUE)
  mort.predfi <- rasterize(raster1[which(raster1[, 1] == unique(raster1[, 1])[1]), c(2, 3)], grid, raster1[which(raster1[, 1] == unique(raster1[, 1])[1]), mortcol], fun = mean, background = NA, mask = FALSE, na.rm = TRUE)
  for (i in 2:length(unique(raster1[, 1]))) {
    mort.predfi <- addLayer(mort.predfi, rasterize(raster1[which(raster1[, 1] == unique(raster1[, 1])[i]), c(2, 3)], grid, raster1[which(raster1[, 1] == unique(raster1[, 1])[i]), mortcol], fun = mean, background = NA, mask = FALSE, na.rm = TRUE))
    mort.bafi <- addLayer(mort.bafi, rasterize(raster1[which(raster1[, 1] == unique(raster1[, 1])[i]), c(2, 3)], grid, raster1[which(raster1[, 1] == unique(raster1[, 1])[i]), 7], fun = mean, background = NA, mask = FALSE, na.rm = TRUE))
  }
  mort.predwi <- weighted.mean(mort.predfi, mort.bafi, na.rm = T)
  return(mort.predwi)
}

#-------------------------------- Run model and create figures
# Fig 1E
insect.hist <- insect_histmodel(fia.g1[, c(2, 5, 6, 11, 13, 15, 20, 22, 24, 29, 28, 27, 30)], 20, 10)

# Scale up by basal area
insect.obs <- weight.proj1(insect.hist, dummy, 4)

# writeRaster(insect.obs, paste(dir, "Fig1E_InsectModel_FIAwide-ObsMort_05-08-2021", sep=""),  format='GTiff')

# Fig 1F models
histcols6 <- c(19, 23, 25, 30, 32, 34)
insect.outa <- mod_insect_proj(fia.g1[, c(2, 5, 6, 11, 13, 15, 20, 22, 24, 29)], 20, 10, prepclim1(fialonga, fialonga, histcols6, histcols6))
insect.outb <- mod_insect_proj(fia.g1[, c(2, 5, 6, 11, 13, 15, 20, 22, 24, 29)], 20, 10, prepclim1(fialongb, fialongb, histcols6, histcols6))
insect.outc <- mod_insect_proj(fia.g1[, c(2, 5, 6, 11, 13, 15, 20, 22, 24, 29)], 20, 10, prepclim1(fialongc, fialongc, histcols6, histcols6))

# Make Fig 1F
insect1 <- insect.outa
insect1[, 6] <- rowMeans(cbind(insect.outa[, 6], insect.outb[, 6], insect.outc[, 6])) # Average modeled mortality across 3 baselines
mort.insects.modlong1 <- weight.proj1(insect1, dummy, 6)

# writeRaster(mort.insects.modlong1, paste(dir, "Fig1F_InsectModel_ModeledFIAlongEnsembleHistMort_05-08-2021", sep=""),  format='GTiff')

#-------------------------------------------- DROUGHT MODELS (Fig 1C/D)
mod_drought_proj <- function(fitdata, minthres, ftypcol, preddata) {
  colnames(preddata) <- c("V1", "V2", "V3", "V4", "V5", "V6", "V7", "V8", "V9", "V10", "V11", "V12", "V13")
  colnames(fitdata) <- c("V1", "V2", "V3", "V4", "V5", "V6", "V7", "V8", "V9", "V10")
  mod1 <- array(dim = c(0, 7))
  for (i in 1:112) {
    ftype <- unique(fitdata[, ftypcol])[i]
    non_zero <- ifelse(fitdata[which(fitdata[, ftypcol] == ftype), 1] > 0, 1, 0)
    d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero) # Historical mortality data
    preddata1 <- data.frame(preddata[which(preddata[, 10] == ftype), ]) # Projection data

    # Set up the output matrix
    drt.out <- array(dim = c(dim(preddata1)[1], 7))
    drt.out[, 1] <- ftype # FORTYPCD
    drt.out[, 2] <- preddata1[, 12] # lat
    drt.out[, 3] <- preddata1[, 11] # lon
    drt.out[, 7] <- preddata1[, 13] # Biomass for the FORTYP

    # If FORTYP has non-meaningful historical model, set future projections == mean historical
    if (ftype %!in% FTpred_drt[, 1]) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Predicted mort = mean historical where Nmort<minthres
    if (ftype %!in% FTpred_drt[, 1]) mod1 <- rbind(mod1, drt.out)
    if (ftype %!in% FTpred_drt[, 1]) print(paste("FORTYP prediction skipped due inadequate historical model: ", ftype, sep = ""))
    if (ftype %!in% FTpred_drt[, 1]) next

    # If FORTYP has Nmort<minthres, set future projections == mean historical
    if (sum(d1[, 11], na.rm = T) < minthres) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Pred mort = mean hist where Nmort<minthres
    if (sum(d1[, 11], na.rm = T) < minthres) mod1 <- rbind(mod1, drt.out)
    if (sum(d1[, 11], na.rm = T) < minthres) print(paste("FORTYP prediction skipped due Nmort<minthres: ", ftype, sep = ""))
    if (sum(d1[, 11], na.rm = T) < minthres) next

    # For FORTYPs with adequate historical models, outlier screen
    if (sum(non_zero, na.rm = T) > 0) d1 <- d1[which(d1[, 1] < quantile(d1[, 1], 0.995)), ]
    if (dim(d1)[1] == 0) d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero)

    # Construct historical mortality model
    m1 <- glm(non_zero ~ V2 + V3 + V4 + V5 + V6 + V7 + V8 + V9, data = d1, family = binomial(link = logit), control = gm.ctl)
    m2 <- betareg(V1 ~ V2 + V3 + V4 + V5 + V6 + V7 + V8 + V9, data = subset(d1, non_zero == 1))

    # Predict with future climate projections
    pred1 <- predict(m1, newdata = preddata1[, 2:9], se = TRUE, type = "response")
    pred2 <- predict(m2, newdata = preddata1[, 2:9], se = TRUE, type = "response")
    pred1a <- ifelse(pred1$fit >= 0.5, 1, 0)
    pred3 <- pred2 * pred1a
    pred4 <- pred2 * pred1$fit

    # 		Store model output
    drt.out[, 5] <- pred3
    drt.out[, 6] <- pred4 # Predicted (expected value) mort frac
    mod1 <- rbind(mod1, drt.out)
  }
  print(paste("FORTYPs projected: ", length(unique(mod1[which(mod1[, 5] != "NA"), 1])), sep = ""))
  print(paste("FORTYPs completed: ", i, sep = ""))
  return(mod1)
}

drought_histmodel <- function(fitdata, minthres, ftypcol) {
  colnames(fitdata) <- c("V1", "V2", "V3", "V4", "V5", "V6", "V7", "V8", "V9", "V10", "V11", "V12", "V13")
  mod1 <- array(dim = c(0, 7))
  for (i in 1:112) {
    ftype <- unique(fitdata[, ftypcol])[i]
    non_zero <- ifelse(fitdata[which(fitdata[, ftypcol] == ftype), 1] > 0, 1, 0)
    d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero)

    drt.out <- array(dim = c(dim(d1)[1], 7))
    drt.out[, 1] <- d1[, ftypcol] # FORTYPCD
    drt.out[, 2] <- d1[, 12] # lat
    drt.out[, 3] <- d1[, 11] # lon
    drt.out[, 4] <- d1[, 1] # Obs mort frac
    drt.out[, 7] <- d1[, 13] # BAlive0 for the FORTYP

    # If FORTYP has non-meaningful historical model, set modeled values == mean historical
    if (ftype %!in% FTpred_drt[, 1]) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Predicted mort = mean historical where Nmort<minthres
    if (ftype %!in% FTpred_drt[, 1]) mod1 <- rbind(mod1, drt.out)
    if (ftype %!in% FTpred_drt[, 1]) print(paste("FORTYP prediction skipped due inadequate historical model: ", ftype, sep = ""))
    if (ftype %!in% FTpred_drt[, 1]) next

    # If FORTYP has Nmort<minthres, set modeled values == mean historical
    if (sum(d1[, 14], na.rm = T) < minthres) drt.out[, 6] <- mean(d1[, 1], na.rm = T) # Predicted mort = mean historical where Nmort<minthres
    if (sum(d1[, 14], na.rm = T) < minthres) mod1 <- rbind(mod1, drt.out)
    if (sum(d1[, 14], na.rm = T) < minthres) print(paste("FORTYP prediction skipped due Nmort<minthres: ", ftype, sep = ""))
    if (sum(d1[, 14], na.rm = T) < minthres) next

    # For FORTYPs with adequate historical models, remove outliers
    if (sum(non_zero, na.rm = T) > 0) d1 <- d1[which(d1[, 1] < quantile(d1[, 1], 0.995)), ]
    if (dim(d1)[1] == 0) d1 <- data.frame(fitdata[which(fitdata[, ftypcol] == ftype), ], non_zero)

    # Build models
    m1 <- glm(non_zero ~ V2 + V3 + V4 + V5 + V6 + V7 + V8 + V9, data = d1, family = binomial(link = logit), control = gm.ctl)
    m2 <- betareg(V1 ~ V2 + V3 + V4 + V5 + V6 + V7 + V8 + V9, data = subset(d1, non_zero == 1))

    pred1 <- predict(m1, newdata = d1[, 2:9], se = TRUE, type = "response")
    pred2 <- predict(m2, newdata = d1[, 2:9], se = TRUE, type = "response")
    pred1a <- ifelse(pred1$fit >= 0.5, 1, 0)
    pred3 <- pred1$fit
    pred4 <- pred2 * pred1$fit

    if (dim(drt.out)[1] != dim(d1)[1]) drt.out <- array(dim = c(dim(d1)[1], 7))
    drt.out[, 1] <- d1[, ftypcol] # FORTYPCD
    drt.out[, 2] <- d1[, 12] # lat
    drt.out[, 3] <- d1[, 11] # lon
    drt.out[, 4] <- d1[, 1] # Obs mort frac
    drt.out[, 7] <- d1[, 13] # BAlive0 for the FORTYP

    drt.out[, 5] <- pred3
    drt.out[, 6] <- pred4 # Predicted (expected value) mort frac
    mod1 <- rbind(mod1, drt.out)
  }
  print(paste("FORTYPs modeled: ", length(unique(mod1[which(mod1[, 5] != "NA"), 1])), sep = ""))
  print(paste("FORTYPs completed: ", i, sep = ""))
  return(mod1)
}

#-------------------------------- Run model and create figures
# Fig 1C
drought.hist <- drought_histmodel(fia.g1[, c(1, 5, 7, 11, 13, 17, 18, 22, 24, 29, 28, 27, 30)], 20, 10)

# Scale up by basal area
drought.obs <- weight.proj1(drought.hist, dummy, 4)

# writeRaster(drought.obs, paste(dir, "Fig1C_DroughtModel_FIAwide-ObsMort_05-08-2021", sep=""),  format='GTiff')

# Fig 1D
histcols1a <- c(19, 24, 25, 30, 31, 35)
drought.outa <- mod_drought_proj(fia.g1[, c(1, 5, 7, 11, 13, 17, 18, 22, 24, 29)], 20, 10, prepclim1(fialonga, fialonga, histcols1a, histcols1a))
drought.outb <- mod_drought_proj(fia.g1[, c(1, 5, 7, 11, 13, 17, 18, 22, 24, 29)], 20, 10, prepclim1(fialongb, fialongb, histcols1a, histcols1a))
drought.outc <- mod_drought_proj(fia.g1[, c(1, 5, 7, 11, 13, 17, 18, 22, 24, 29)], 20, 10, prepclim1(fialongc, fialongc, histcols1a, histcols1a))

# Make Fig 1D
drought1 <- drought.outa
drought1[, 6] <- rowMeans(cbind(drought.outa[, 6], drought.outb[, 6], drought.outc[, 6])) # Average modeled mortality across 3 baselines
mort.drought.modlong1 <- weight.proj1(drought1, dummy, 6)

# writeRaster(mort.drought.modlong1, paste(dir, "Fig1D_DroughtModel_ModeledFIAlongEnsembleHistMort_05-08-2021", sep=""),  format='GTiff')

#-------------------------------------- First-order reversal estimates
# Historical reversals for drought and insects in FIA data
drt.rev <- rep(0, 102453)
drt.rev[which(fianew1[, 27] < 0.01 & fianew1[, 19] < fianew1[, 18] & fianew1[, 31] < 0.01 & fianew1[, 35] < 0.01 & fianew1[, 39] < 0.01 & fianew1[, 47] < 0.01 & fianew1[, 55] == "False" & fianew1[, 59] == "False" & fianew1[, 63] == "False" & fianew1[, 71] == "False" & fianew1[, 75] == "False" & fianew1[, 79] == "False" & fianew1[, 83] == "False" & fianew1[, 151] < -2 & fianew1[, 67] == "False" & ((fianew1[, 19] - fianew1[, 18]) / -fianew1[, 18]) > 0.1)] <- 1

ins.rev <- rep(0, 102453)
ins.rev[which(fianew1[, 27] > 0.5 & fianew1[, 19] < fianew1[, 18] & ((fianew1[, 19] - fianew1[, 18]) / -fianew1[, 18]) > 0.1)] <- 1

# Calculated annualized drought and insect p(rev)
revhist1 <- array(dim = c(102453, 2))
for (i in 1:102453) {
  ftrev <- fianew1[i, 4]
  revhist1[i, 1] <- drt.rev[i] / mean(fianew1[which(fianew1[, 4] == ftrev), 15] - fianew1[which(fianew1[, 4] == ftrev), 14], na.rm = T)
  revhist1[i, 2] <- ins.rev[i] / mean(fianew1[which(fianew1[, 4] == ftrev), 15] - fianew1[which(fianew1[, 4] == ftrev), 14], na.rm = T)
}

# Calculate 100-year integrated risk %s of constant historical reversal annual probabilities
100 * dbinom(1, 100, mean(revhist1[, 1], na.rm = T)) # Drought
100 * dbinom(1, 100, mean(revhist1[, 2], na.rm = T)) # Insects

# Simulate 100-year integrated risk %s with Fig 2 multi-model mean CMIP6 projections
out <- array(dim = c(1000, 0))
perc1 <- c(1, 1.022, 1.063, 1.148, 1.151, 1.230, 1.380, 1.384, 1.799, 1.846) # Percent changes in BAmortfrac/yr from drought models compared to 1990-2010 baselines
for (i in 1:10) {
  dec1 <- rbinom(1000, 10, 0.002457 * perc1[i]) # Draw 1000 10-year draws with observed probability
  out <- cbind(out, dec1)
  sumout <- rowSums(out)
  sumout1 <- sum(ifelse(sumout > 0, 1, 0))
}
100 * sumout1 / 1000

# Insects
out <- array(dim = c(1000, 0))
perc2 <- c(1, 1.049, 1.042, 1.193, 1.0231, 1.136, 1.289, 1.197, 1.622, 1.701) # Percent changes in BAmortfrac/yr from insect models compared to 1990-2010 baselines
for (i in 1:10) {
  dec1 <- rbinom(1000, 10, 0.00250982 * perc2[i]) # Draw 1000 10-year draws with observed probability
  out <- cbind(out, dec1)
  sumout <- rowSums(out)
  sumout1 <- sum(ifelse(sumout > 0, 1, 0))
}
100 * sumout1 / 1000

# Fire
out <- array(dim = c(1000, 0))
hist1 <- read.csv(paste(dir, "historical_fire.csv", sep = ""), header = T)
histave <- mean(hist1[17:35, 2])
fut1 <- read.csv(paste(dir, "future_fire.csv", sep = ""), header = T)
perc3 <- c(rep(1, 19), (fut1[50:130, 3] - histave) / histave + 1)
for (i in 1:100) {
  dec1 <- rbinom(1000, 1, histave * perc3[i]) # Draw 1000 1-year draws with observed probability
  out <- cbind(out, dec1)
  sumout <- rowSums(out)
  sumout1 <- sum(ifelse(sumout > 0, 1, 0))
}
100 * sumout1 / 1000


# ####
# ####
# #### AIM: Use PCT data to create scenario of change where commuting cycling levels increase and car journeys decrease which can be imported
# ####      into A/B Street city simulation software. This method should be fully reproducible for all other pct_regions.
# ####
# ####
#
# ####  LIBRARYS ####
# library(pct)
# library(sf)
# library(tidyverse)
# library(abstr)
#
# # for plotting
# # library(tmap)
# # library (mapview)
#
# ####    READ DATA ####
# devon_zones = pct::get_pct_zones(region = "devon", geography = "msoa") # get zone data
# exeter_zones = devon_zones %>% filter(lad_name == "Exeter") %>% select(geo_code) # filter for exeter
#
# exeter_commute_od = pct::get_pct_lines(region = "devon", geography = "msoa") %>% # get commute od data
#   filter(lad_name1 == "Exeter" &
#            lad_name2 == "Exeter") # filter for exeter
#
#
# ####    CLEAN DATA , CALCULATE EUCLIDEAN DISTANCE  & GENERATE SCENARIOS OF CHANGE ####
# exeter_commute_od = exeter_commute_od %>%
#   mutate(cycle_base = bicycle) %>%
#   mutate(walk_base = foot) %>%
#   mutate(transit_base = bus + train_tube) %>% # bunch of renaming -_-
#   mutate(drive_base = car_driver + car_passenger + motorbike + taxi_other) %>%
#   mutate(all_base = all) %>%
#   mutate(
#     # create new columns
#     pcycle_godutch_uptake = pct::uptake_pct_godutch_2020(distance = rf_dist_km, gradient = rf_avslope_perc),
#     cycle_godutch_additional = pcycle_godutch_uptake * drive_base,
#     cycle_godutch = cycle_base + cycle_godutch_additional,
#     pcycle_godutch = cycle_godutch / all_base,
#     drive__godutch = drive_base - cycle_godutch_additional,
#     across(c(drive__godutch, cycle_godutch), round, 0),
#     all_go_dutch = drive__godutch + cycle_godutch + transit_base + walk_base
#   ) %>%
#   select(
#     # select variables for new df
#     geo_code1,
#     geo_code2,
#     cycle_base,
#     drive_base,
#     walk_base,
#     transit_base,
#     all_base,
#     all_go_dutch,
#     drive__godutch,
#     cycle_godutch,
#     cycle_godutch_additional,
#     pcycle_godutch
#   )
#
# identical(exeter_commute_od$all_base, exeter_commute_od$all_go_dutch) # sanity check: make sure total remains the same (not a dynamic model where population change is factored in)
#
# ####    DOWNLOAD OSM BUILDING DATA ####
# osm_polygons = osmextract::oe_read(
#   "https://download.geofabrik.de/europe/great-britain/england/devon-latest.osm.pbf",
#   # download osm buildings for region using geofabrik
#   layer = "multipolygons"
# )
#
# building_types = c(
#   "yes",
#   "house",
#   "detached",
#   "residential",
#   "apartments",
#   "commercial",
#   "retail",
#   "school",
#   "industrial",
#   "semidetached_house",
#   "church",
#   "hangar",
#   "mobile_home",
#   "warehouse",
#   "office",
#   "college",
#   "university",
#   "public",
#   "garages",
#   "cabin",
#   "hospital",
#   "dormitory",
#   "hotel",
#   "service",
#   "parking",
#   "manufactured",
#   "civic",
#   "farm",
#   "manufacturing",
#   "floating_home",
#   "government",
#   "bungalow",
#   "transportation",
#   "motel",
#   "manufacture",
#   "kindergarten",
#   "house_boat",
#   "sports_centre"
# )
# osm_buildings  = osm_polygons %>%
#   dplyr::filter(building %in% building_types) %>%
#   dplyr::select(osm_way_id, name, building)
#
# osm_buildings_valid = osm_buildings[sf::st_is_valid(osm_buildings), ]
#
# exeter_osm_buildings_all = osm_buildings_valid[exeter_zones, ]
#
# #mapview(exeter_osm_buildings_all)
#
# # Filter down large objects for package -----------------------------------
# exeter_osm_buildings_all_joined = exeter_osm_buildings_all %>%
#   sf::st_join(exeter_zones)
#
# set.seed(2021)
# exeter_osm_buildings_sample = exeter_osm_buildings_all_joined %>%
#   dplyr::filter(!is.na(osm_way_id))
#
# exeter_osm_buildings_tbl = exeter_osm_buildings_all %>%
#   dplyr::filter(osm_way_id %in% exeter_osm_buildings_sample$osm_way_id)
#
#
# ####  LOGIC GATE ####
# # Logic gate for go_dutch scenario of change, where cycling levels increase to a proportion reflecting the Netherlands.
# #Switch to FALSE if you want census commuting OD
# go_dutch = TRUE
# if (go_dutch == TRUE) {
#   exeter_od = exeter_commute_od %>%
#     mutate(All = all_go_dutch) %>%
#     mutate(Bike = cycle_godutch) %>%
#     mutate(Transit = transit_base) %>%
#     mutate(Drive = drive_base) %>%
#     mutate(Walk = walk_base) %>%
#     select(geo_code1, geo_code2, All, Bike, Transit, Drive, Walk,geometry)
# } else {
#   exeter_od = exeter_commute_od %>%
#     mutate(All = all_base) %>%
#     mutate(Bike = cycle_base) %>%
#     mutate(Drive = drive_base) %>%
#     mutate(Transit = transit_base) %>%
#     mutate(Walk = walk_base) %>%
#     select(geo_code1, geo_code2, All, Bike, Transit, Drive, Walk, geometry)
# }
#
# ####  GENERATE A/B STREET SCENARIO ####
# output_sf = ab_scenario(
#   od = exeter_od,
#   zones = exeter_zones,
#   zones_d = NULL,
#   origin_buildings = exeter_osm_buildings_tbl,
#   destination_buildings = exeter_osm_buildings_tbl,
#   pop_var = 3,
#   time_fun = ab_time_normal,
#   output = "sf",
#   modes = c("Walk", "Bike", "Drive", "Transit")
# )
#
# # make map using tmap
# # tm_shape(output_sf) + tmap::tm_lines(col = "mode", lwd = .8, lwd.legeld.col = "black") +
# #   tm_shape(exeter_zones) + tmap::tm_borders(lwd = 1.2, col = "gray") +
# #   tm_text("geo_code", size = 0.6, col = "black") +
# #   tm_style("cobalt")
#
# #### SAVE JSON FILE ####
# output_json = ab_json(output_sf, time_fun = ab_time_normal, scenario_name = "Exeter Example")
# ab_save(output_json, f = "../../Desktop/exeter.json")
#
# #### COMMANDS FOR AB STREET
# # $ cargo run
# # $ cargo run --bin import_traffic -- --map=PATH/TO/MAP --input=/PATH/TO/JSON.json


# Module for graph based algorithms

# Graph representation:
#  - An edges frame is a data frame with columns 'source.id', 'target.id' and 'weight',
#    where 'source' and 'target' each reference a vertex v_i in V (indexed from 1 to |V|),
#    and 'weight' gives the edge weight (often just 1 to measure degree).
#  - An edges representation is an object with a get_edges() method.
#  - A graph representation is an edges representation with the following methods:
#     * get_node_ids()           - returns a vector of all vertex identifiers.
#     * get_out_edges(source.id) - returns an edges frame of all edges directed from the source vertex.
#     * get_in_edges(target.id)  - returns an edges frame of all edges directed to the target vertex.

# Description:
#   Wraps an edges frame into an edges representation.
# Input:
#   edges - The edges frame.
# Output:
#   edges.repr - The edges representation.
graph.get_edges_repr = function(edges) {
    edges.repr = list()
    edges.repr$get_edges = function() {
        return(edges)
    }
    return(edges.repr)
}

# Description:
#   Wraps an edges edges representation into a graph representation.
# Input:
#   edges.repr - The edges representation.
# Output:
#   graph.repr - The graph representation.
graph.from_edges_repr = function(edges.repr) {
    edges = edges.repr$get_edges()
    graph.repr = list()
    graph.repr$get_edges = function() {
        return(edges)
    }
    graph.repr$get_node_ids = function() {
        return(sort(unique(c(edges$source.id, edges$target.id))))
    }
    graph.repr$get_out_edges = function(source.id) {
        return(edges[edges$source.id == source.id,])
    }
    graph.repr$get_in_edges = function(target.id) {
        return(edges[edges$target.id == target.id,])
    }
    return(graph.repr)
}

# Description:
#   Internal function to encapsulate all clustering information into a single object.
# Input:
#   num.nodes - The total number |V| of vertices in the graph, on the assumption that
#               every vertex id has been uniquely mapped to an integer from 1 to num.nodes.
#   trace     - An optional flag indicating whether (TRUE) or not (FALSE) to
#               track incremental score changes; defaults to FALSE.
# Output:
#   clusters - An object representing the internal cluster functions.
graph.init_clusters = function(num.nodes, trace=FALSE) {
    # Information about all nodes:
    # Sum of all edge weights (half of total vertex weight)
    total_edge_weights = 0
    # Map node index to cluster id
    node_cluster_map = 1:num.nodes
    # Count weight of internal edges in each cluster (both nodes in cluster)
    internal_cluster_edges_map = rep(0, num.nodes)
    # Count weight of external edges pointing into our out of each cluster (only one node in cluster)
    external_cluster_edges_map = rep(0, num.nodes)
    # Initial modularity score
    initial_score = 0
    # Cummulative change in modularity score
    change_score = 0
    if (trace) {
        change_scores = NULL
    }
    # Control initial score computation
    score_counter = 0

    # Information about about current (pivot) node:
    # Total edge weight of a node
    node_weight = 0
    # Weight of self edge to and from node
    self_weight = 0
    # Weights of edges from node to each cluster
    local_cluster_edges_map = rep(0, num.nodes)

    # Create object
    obj = list()
    # Add node v_i to new cluster if not already clustered
    obj$add_cluster_node = function(i) {
        # No action needed as nodes already indexed.
    }
    # Put node v_i into existing cluster c_i
    obj$set_cluster_node = function(i, c_i) {
        node_cluster_map[i] <<- c_i
    }
    # Add directed edge between node clusters
    obj$add_cluster_edge = function(i, j, w_ij) {
        if (i == j) w_ij = 0.5 * w_ij
        total_edge_weights <<- total_edge_weights + w_ij
        c_i = node_cluster_map[[i]]
        c_j = node_cluster_map[[j]]
        if (c_i == c_j) {
            internal_cluster_edges_map[c_i] <<- internal_cluster_edges_map[c_i] + w_ij
        } else {
            external_cluster_edges_map[c_i] <<- external_cluster_edges_map[c_i] + w_ij
            external_cluster_edges_map[c_j] <<- external_cluster_edges_map[c_j] + w_ij
        }
    }
    # Initialise local cluster edge counts for a single node
    obj$init_local_cluster_edges = function() {
        node_weight <<- 0
        self_weight <<- 0
        local_cluster_edges_map <<- rep(0, num.nodes)
    }
    # Add undirected edge between pivot node i and node j
    obj$add_local_cluster_edge = function(i, j, w_ij) {
        if (i == j) {
            self_weight <<- w_ij
            w_ij = 0.5 * w_ij
        }
        node_weight <<- node_weight + w_ij
        c_j = node_cluster_map[[j]]
        local_cluster_edges_map[c_j] <<- local_cluster_edges_map[c_j] + w_ij
    }
    # If possible, move pivot node to another cluster with highest positive change of score
    obj$move_to_best_cluster = function(i) {
        c_i = node_cluster_map[[i]]
        A_dd = 2 * total_edge_weights # A..
        A_id = node_weight # A_i.
        score_counter <<- score_counter + 1
        if (score_counter <= num.nodes) {
            # Continue computing initial score
            Q_i = (self_weight - A_id * A_id / A_dd) / A_dd
            initial_score <<- initial_score + Q_i
        }
        # Compute score change in removing node i from its cluster
        A_i_gpi = internal_cluster_edges_map[[c_i]] # A_i,g+i
        A_ig =  A_i_gpi - self_weight # A_i,g
        sigma_gpi_tot = 2 * internal_cluster_edges_map[[c_i]] + external_cluster_edges_map[[c_i]] # Sigma_g+i^tot
        sigma_tot = sigma_gpi_tot - A_id # Sigma_g^tot
        delta_Q_gpimi = -2 * (A_ig  - sigma_tot * A_id / A_dd) / A_dd # Delta Q_(g+i)-i
        max_delta_Q = -Inf
        max_cluster_id = 0
        local_cluster_edges_map[c_i] = 0 # Ignore current cluster
        for (c_j in which(local_cluster_edges_map > 0)) {
            sigma_tot = 2 * internal_cluster_edges_map[[c_j]] + external_cluster_edges_map[[c_j]] # Sigma_g'^tot
            A_igd = local_cluster_edges_map[[c_j]] # A_i,g'
            delta_Q_gdpi = 2 * (A_igd - sigma_tot * A_id / A_dd) / A_dd # Delta Q_g'+i
            delta_Q = delta_Q_gpimi + delta_Q_gdpi
            if (delta_Q > 1e-8 && delta_Q > max_delta_Q) {
                max_delta_Q = delta_Q
                max_cluster_id = c_j
                #print(paste0("DEBUG: Best delta-score=", delta_Q, ". Suggest move node from cluster ", c_i, " to cluster ", c_j))
            }
        }
        if (max_cluster_id > 0) {
            c_j = max_cluster_id
            change_score <<- change_score + max_delta_Q
            if (trace) {
                change_scores <<- c(change_scores, max_delta_Q)
            }
            # Remove node i from cluster c_i
            internal_cluster_edges_map[[c_i]] <<- internal_cluster_edges_map[[c_i]] - A_ig - 0.5 * self_weight
            external_cluster_edges_map[[c_i]] <<- external_cluster_edges_map[[c_i]] - A_id + 2 * A_ig + self_weight
            # Add node i to cluster c_j
            A_igd = local_cluster_edges_map[[c_j]] # A_i,g'
            internal_cluster_edges_map[[c_j]] <<- internal_cluster_edges_map[[c_j]] + A_igd + 0.5 * self_weight
            external_cluster_edges_map[[c_j]] <<- external_cluster_edges_map[[c_j]] + A_id - 2 * A_igd - self_weight
            node_cluster_map[i] <<- c_j
            return(TRUE)
        } else {
            return(FALSE)
        }
    }
    # Provide initial score for singleton clusters
    obj$get_inital_score = function() {
        return(initial_score)
    }
    # Provide final score for clustered graph
    obj$get_final_score = function() {
        return(initial_score + change_score)
    }
    # Provide incremental score changes
    obj$get_score_changes = function() {
        if (trace) return(change_scores)
        return(c())
    }
    # Remap the remaining clusters to indices in 1, 2, ..., K for smallest K
    obj$get_cluster_map = function() {
        cluster.nums = sort(unique(node_cluster_map))
        K = length(cluster.nums)
        new_cluster_map = rep(0, K)
        for (k in 1:K) {
            idx = which(node_cluster_map == cluster.nums[[k]])
            new_cluster_map[idx] = k
        }
        return(new_cluster_map)
    }
    # Expose object
    return(obj)
}

# Description:
#   An implementation of phase 1 of the Louvain modularity community detection (i.e. node clustering).
# Input:
#   graph.repr - A graph representation.
#   trace      - An optional flag indicating whether (TRUE) or not (FALSE) to
#                track incremental score changes; defaults to FALSE.
# Output:
#   clusters - A |V| length vector giving the cluster index c_i for each vertex v_i.
graph.cluster = function(graph.repr, trace=FALSE) {
    node.ids = graph.repr$get_node_ids()
    num.nodes = length(node.ids)
    # Initialisation:
    clusters = graph.init_clusters(num.nodes, trace)
    for (i in 1:num.nodes) {
        v_i = node.ids[i]
        clusters$add_cluster_node(i)
        # To avoid double counting, only traverse edges in source -> target direction.
        out_edges = graph.repr$get_out_edges(v_i)
        num.out_edges = nrow(out_edges)
        if (num.out_edges > 0) {
            for (edge.idx in 1:num.out_edges) {
                edge = out_edges[edge.idx,]
                v_j = edge$target.id
                j = which(node.ids == v_j)
                clusters$add_cluster_node(j)
                w_ij = edge$weight
                clusters$add_cluster_edge(i, j, w_ij)
            }
        }
    }

    # Cluster reassignments:
    while (TRUE) {
        changed = FALSE
        for (i in 1:num.nodes) {
            v_i = node.ids[i]
            clusters$init_local_cluster_edges()
            out_edges = graph.repr$get_out_edges(v_i)
            num.out_edges = nrow(out_edges)
            if (num.out_edges > 0) {
                for (edge.idx in 1:num.out_edges) {
                    edge = out_edges[edge.idx,]
                    v_j = edge$target.id
                    j = which(node.ids == v_j)
                    w_ij = edge$weight
                    clusters$add_local_cluster_edge(i, j, w_ij) # i -> j becomes i -- j.
                }
            }
            in_edges = graph.repr$get_in_edges(v_i)
            num.in_edges = nrow(in_edges)
            if (num.in_edges > 0) {
                for (edge.idx in 1:num.in_edges) {
                    edge = in_edges[edge.idx,]
                    v_j = edge$source.id
                    j = which(node.ids == v_j)
                    w_ji = edge$weight
                    clusters$add_local_cluster_edge(i, j, w_ji) # Yes, i <- j becomes i -- j.
                }
            }
            moved = clusters$move_to_best_cluster(i)
            changed = changed || moved
        }
        if (!changed) break
    }
    # Map node id to cluster id
    cluster.ids = clusters$get_cluster_map()
    names(cluster.ids) = node.ids
    if (trace) {
        scores = cumsum(c(clusters$get_inital_score(), clusters$get_score_changes()))
    } else {
        scores = c(clusters$get_inital_score(), clusters$get_final_score())
    }
    return(list(
        cluster.ids = cluster.ids,
        initial.score = clusters$get_inital_score(),
        final.score = clusters$get_final_score(),
        incremental.scores = scores
    ))
}

graph.matrix = function(graph.repr) {
    node.ids = graph.repr$get_node_ids()
    num.nodes = length(node.ids)
    D = matrix(data=0, nrow=num.nodes, ncol=num.nodes)
    for (i in 1:num.nodes) {
        v_i = node.ids[[i]]
        out.edges = graph.repr$get_out_edges(v_i)
        num.out.edges = nrow(out.edges)
        if (num.out.edges > 0) {
            for (edge.idx in 1:num.out.edges) {
                edge = out.edges[edge.idx,]
                v_j = edge$target.id
                j = which(node.ids == v_j)
                w_ij = edge$weight
                D[i, j] = D[i, j] + w_ij
            }
        }
    }
    rownames(D) = node.ids
    colnames(D) = node.ids
    return(D)
}

test.graph.cluster = function() {
    edges = data.frame(source.id=c("X", "Y", "Z"), target.id=c("Z", "X", "Y"), weight=1, stringsAsFactors=FALSE)
    edges.repr = graph.get_edges_repr(edges)
    graph.repr = graph.from_edges_repr(edges.repr)
    D = graph.matrix(graph.repr)
    A = D + t(D)
    diag(A) = diag(D)
    A_id = rowSums(A)
    A_dd = sum(A_id)
    Q_mat = (A - A_id %*% t(A_id) / A_dd) / A_dd
    # Step 0 clusters: 1, 2, 3
    delta0 = matrix(0, nrow=3, ncol=3)
    diag(delta0) = 1
    step0.score = sum(Q_mat * delta0)
    # Step 1 clusters: 2, 2, 3
    delta1 = delta0
    delta1[1, 2] = 1; delta1[2, 1] = 1
    step1.score = sum(Q_mat * delta1)
    # Step 2 clusters: 2, 2, 2
    delta2 = delta1
    delta2[1, 3] = 1; delta2[2, 3] = 1; delta2[3, 1] = 1; delta2[3, 2] = 1
    step2.score = sum(Q_mat * delta2)
    scores = c(step0.score, step1.score, step2.score)
    # Step 3 clusters: 1, 1, 1 (remapped)
    res = graph.cluster(graph.repr, trace=TRUE)
    print(paste0("Expected incremental scores: ", paste(scores, collapse=", ")))
    print(paste0("Observed incremental scores: ", paste(res$incremental.scores, collapse=", ")))
    print(paste0("Expected clusters: ", paste(rep(1, 3), collapse=", ")))
    print(paste0("Observed clusters: ", paste(res$cluster.ids, collapse=", ")))
}


two_way_tanimoto_predict <- function(Kc, Kr, S0, S1, MW, similarity.consumer, similarity.resource, minimum_threshold) {
    # Two-way Tanimoto Algorithm
    # ===========================

    # Parameters:
    #   Kc  Integer, how many neighbors to select for consumers
    #   Kr  Integer, how many neighbors to select for resources
    #   S0  A large set of species and their preys, with column structure ['taxon', 'taxonomy', 'resource', 'non-resource']
    #   S1  The subset of S0 where we want to predict new preys, string vector with taxa name
    #   MW  Mimimum weight to accept a candidate as a prey
    #

    # // TODO: I think MW should be a function of Kc & Kr and perhaps of the number of candidate resources. For example, a similar consumer could have multiple prey, which would artificially inflate the weight added to each prey in the candidate list. For exemple, Atlantic cod has over 600 prey species listed in the interaction catalogue... Hence, the longer the candidate list, the more likely a very small similarity will be turned into a predicted interaction
    # // REVIEW: Multiply similar.consumer[similarity] * similar.resource[similarity]? It's a similarity of a similarity in a sense...

    # // TODO: Different similarity measurement for resources and consumers

    # // REVIEW: Remove cannibalism from empirical data, or allow for it, or add parameter that allows or prevents cannibalism in the predictions. There are lots of predicted cannibalism interations in the catalogue and empirical webs. Accuracy would increase if it was allowed.

    # Output
    #   A vector of sets (the preys for each species)

    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    # List of things to adjust - make it
    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    # Process steps:
      # A prior process to this is to get the similarity matrix between all combinations of catalogue taxa and species in S1
      # For each species in S1:
      # 1. Identify resources already known in interaction catalogue (S0) for S1 species
        # 1.1 If resoures are in S1, automatically add them to the predictions as empirically valid interactions
        # 1.2 If resources are not in S1, find Kr similar resources in S1 and add them to candidate list with weight equal to their similarity

      # 2. Identify Kc similar consumers to S1 in S0
        # 2.1 Extract set of candidate resources from each similar consumer, if any
        # 2.2 If candidate resource is in S1, add it to candidate list with weight 1
        # 2.3 If candidate resource not in S1, find Kr similar resources in S1 and add them to candidate list with weight equal to their similarity

      # 3. Make predictions:
        # 3.1 Remove taxa with weight < to minimum weight (MW) from prediction list
        # 3.2 Sort prediction list according to weight. Higher weights mean higher likelihood for resource being consumed

    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    # !!! étendre predator et non predator pour resource... il faudrait aussi calculer la similarité des proies sur la base de leurs prédateurs partagés !!!
    # !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    predictions <- matrix(nrow = length(S1), ncol = 3, data = "", dimnames = list(c(S1), c('consumer','resource_empirical','resource_predictions'))) # empty object for resource predictions
    predictions[, 'consumer'] <- S1

    # pb <- txtProgressBar(min = 0,max = length(S1), style = 3)
    for(i in 1:length(S1)) { # loop through each taxon in S1
        candidates <- matrix(nrow = 0, ncol = 2, dimnames = list(c(), c('resource', 'weight')), data = NA) # empty matrix for resource candidate list for S1[i], with taxon name and weight
        resources.S1 <- unlist(strsplit(S0[S1[i], 'resource'], " \\|\\ ")) # resources of S1[i]

        # Add resources that are already listed as resources for S1[i] in predictions[, 'resource_empirical'] or
        # Find similar resources to resources for S1[i] in S1
        if(length(resources.S1) > 0) {
            empirical <- character()
            for(j in 1:length(resources.S1)) { #loop through empirical resources for S1
                if(resources.S1[j] %in% S1) {
                    empirical <- c(empirical, resources.S1[j]) # observed resource found in S1 are automatically added to the column resource_empirical
                } else { # selecting Kr most similar resources in S1
                    # Let's assume for this part that we are not compiling a different similarity measure for predators and preys.
                    similar.resource <- matrix(nrow = length(S1)-1, ncol = 2, dimnames = list(c(), c('resource','similarity')), data = NA) # importing K nearest neighbors resources
                    similar.resource[, 'resource'] <- names(sort(similarity.resource[S1[-which(S1 == S1[i])], resources.S1[j]], decreasing = TRUE))
                    similar.resource[, 'similarity'] <- sort(similarity.resource[S1[-which(S1 == S1[i])], resources.S1[j]], decreasing = TRUE)

                    # If multiple taxa with same similarity, randomly select those that will be used as similar resources.
                    if(similar.resource[Kr+1, 'similarity'] == similar.resource[Kr, 'similarity']) {
                        same.similarity <- which(similar.resource[, 'similarity'] == similar.resource[Kr, 'similarity'])
                        similar.resource[same.similarity, ] <- similar.resource[sample(same.similarity), ]
                        similar.resource <- similar.resource[1:Kr, ]
                    } else {
                        similar.resource <- similar.resource[1:Kr, ]
                    }# if for random draw

                    for(l in 1:Kr) { # extracting resource candidates
                        if(all.equal(similar.resource[, 'similarity'], rep('0',Kr)) == TRUE) { # if similarities all == 0, break
                            break
                        } else if(similar.resource[l, 'similarity'] == '0') { # if similarity l == 0, no candidates provided
                            NULL
                            # minimum threshold try.. adding it as a Parameters.. might not make sense, have to discuss it. If we keep it, previous else ifs can be removed
                        } else if(similar.resource[l, 'similarity'] < minimum_threshold) {
                            NULL
                        } else if((similar.resource[l, 'resource'] %in% candidates[, 'resource']) == TRUE) { # if candidate is already in candidate list, add resource' with wt to its weight
                          candidates[which(candidates[, 'resource'] == similar.resource[l]), 'weight'] <- as.numeric(candidates[which(candidates[, 'resource'] == similar.resource[l]), 'weight']) + as.numeric(similar.resource[l, 'similarity'])
                        } else {
                              candidates <- rbind(candidates, similar.resource[l, ]) # if candidate is not in the list, add it resource' with wt to its weight
                        }#if3
                    }#l
                }#if
            }#j
            predictions[S1[i], 'resource_empirical'] <- paste(empirical, collapse = ' | ')
        }#if1

        # Identify similar consumers to S1[i]
        similar.consumer <- matrix(nrow = nrow(similarity.consumer)-1, ncol = 2, dimnames = list(c(), c('consumer','similarity')), data = NA) # emporting K nearest neighbors for consumers
        similar.consumer[, 'consumer'] <- names(sort(similarity.consumer[-which(colnames(similarity.consumer) == S1[i]), S1[i]], decreasing = TRUE))
        similar.consumer[, 'similarity'] <- sort(similarity.consumer[-which(colnames(similarity.consumer) == S1[i]), S1[i]], decreasing = TRUE)

        # If multiple taxa with same similarity, randomly select those that will be used as similar resources.
        if(similar.consumer[Kc+1, 'similarity'] == similar.consumer[Kc, 'similarity']) {
            same.similarity <- which(similar.consumer[, 'similarity'] == similar.consumer[Kc, 'similarity'])
            similar.consumer[same.similarity, ] <- similar.consumer[sample(same.similarity), ]
            similar.consumer <- similar.consumer[1:Kc, ]
        } else {
            similar.consumer <- similar.consumer[1:Kc, ]
        }# if for random draw

        # Est-ce que la valeur de similarité a de l'importance pour l'attribution des proies?
        # If yes, we could add an argument call wt_predator.
          # if(wt_predator == FALSE) {
          #   resources <- unique of all prey species of all similar predators
          # } else {}

        for(j in 1:Kc) { #loop through consumers

            if(all.equal(similar.consumer[, 'similarity'], rep('0',Kc)) == TRUE) { # if similarities all == 0, break
                break
            } else if(similar.consumer[j, 'similarity'] == '0') { # if similarity l == 0, no candidates provided
                NULL
            } else {

                # It's possible that consumers in the list have high taxonomic similarity, but no recorded resource
                candidate.resource <- unlist(strsplit(S0[similar.consumer[j, 'consumer'], 'resource'], " \\|\\ ")) # list of resources for consumer j
                # candidate.resource <- candidate.resource[(candidate.resource %in% resources.S1) == FALSE] # substracting candidate resources that are already listed as resources for S1[i] and hence considered in the preceding code segment

                for(k in 1:length(candidate.resource)) { # loop through resources of consumer j
                    if(length(candidate.resource) == 0) { # if candidate resource list is empty, break
                        break
                    } else if(candidate.resource[1] == "") { # if candidate list is an empty vector "", break
                        break
                    } else if(candidate.resource[k] == S1[i]) {
                    #   #// FIXME: if candidate resource is taxon for which predictions are being made, break (unless we want to allow CANIBALISM). Add argument for cannibalism allowed or not
                         NULL
                    } else if((candidate.resource[k] %in% S1) == TRUE) {
                        if((candidate.resource[k] %in% candidates[, 'resource']) == TRUE) {# if candidate is already in candidate list, add 1 to its weight
                            candidates[which(candidates[, 'resource'] == candidate.resource[k]), 'weight'] <- as.numeric(candidates[which(candidates[, 'resource'] == candidate.resource[k]), 'weight']) + 1
                        } else {
                            candidates <- rbind(candidates, c(candidate.resource[k], 1)) # if candidate is not in the list, add it with 1 to its weight
                        }#if2

                    } else {
                        # Let's assume for this part that we are not compiling a different similarity measure for predators and preys.
                        similar.resource <- matrix(nrow = length(S1)-1, ncol = 2, dimnames = list(c(), c('resource','similarity')), data = NA) # importing K nearest neighbors resources
                        similar.resource[, 'resource'] <- names(sort(similarity.resource[S1[-which(S1 == S1[i])], candidate.resource[k]], decreasing = TRUE))
                        similar.resource[, 'similarity'] <- sort(similarity.resource[S1[-which(S1 == S1[i])], candidate.resource[k]], decreasing = TRUE)

                        # If multiple taxa with same similarity, randomly select those that will be used as similar resources.
                        if(similar.resource[Kr+1, 'similarity'] == similar.resource[Kr, 'similarity']) {
                            same.similarity <- which(similar.resource[, 'similarity'] == similar.resource[Kr, 'similarity'])
                            similar.resource[same.similarity, ] <- similar.resource[sample(same.similarity), ]
                            similar.resource <- similar.resource[1:Kr, ]
                        } else {
                            similar.resource <- similar.resource[1:Kr, ]
                        }# if for random draw

                        for(l in 1:Kr) { # extracting resource candidates
                            if(all.equal(similar.resource[, 'similarity'], rep('0',Kr)) == TRUE) { # if similarities all == 0, break
                                break
                            } else if(similar.resource[l, 'similarity'] == '0') { # if similarity l == 0, no candidates provided
                                NULL
                                # minimum threshold try.. adding it as a Parameters.. might not make sense, have to discuss it. If we keep it, previous else ifs can be removed
                            } else if(similar.resource[l, 'similarity'] < minimum_threshold) {
                                NULL
                            } else if((similar.resource[l, 'resource'] %in% candidates[, 'resource']) == TRUE) { # if candidate is already in candidate list, add 1 to its weight
                              candidates[which(candidates[, 'resource'] == similar.resource[l]), 'weight'] <- as.numeric(candidates[which(candidates[, 'resource'] == similar.resource[l]), 'weight']) + as.numeric(similar.resource[l, 'similarity'])
                            } else {
                                  candidates <- rbind(candidates, similar.resource[l, ]) # if candidate is not in the list, add it with its weight = similarity
                            }#if3
                        }#l
                    } #if1
                }#k
            }#if
        }#j

        candidates <- candidates[which(candidates[, 'weight'] >= MW), ] # remove candidates with a weight below MW
        if(is.matrix(candidates) == TRUE) { #if it's a vector, there's only one predicted resource, no need to order
            candidates[order(candidates[, 'weight']), ] # sorts candidates according to their weight
            predictions[S1[i], 'resource_predictions'] <- paste(candidates[, 'resource'], collapse = ' | ')
        } else {
          predictions[S1[i], 'resource_predictions'] <- paste(candidates['resource'], collapse = ' | ')
        }#if
    # setTxtProgressBar(pb, i)
    }#i
    # close(pb)
    return(predictions)
}#two_way_tanimoto_predict function


#############################################################################
##
## Indirection - phase 2: output gating
## Role/filler recall task
## 2017 Mike Jovanovich
##
## In this version OG sees WM after it has been update by IG within a timestep.
## The action selection if this version simply checks if WM_out outputs only
## the correct filler
##
#############################################################################

## Usage
# Rscript outputgate.r 1 100000 3 10 3 C C D F F SG

#############################################################################
## Parameter initialization
#############################################################################

source('hrr.r')
source('sentenceSets.r')
require(methods)
require(Matrix)

## Build args dictionary
argnames <- list('seed','ntasks','nstripes','nfillers','nroles','state_cd','sid_cd','interstripe_cd',
  'use_sids_input','use_sids_output','gen_test')
args <- commandArgs(trailingOnly = TRUE)

## Set the seed for repeatability
set.seed(as.integer(args[which(argnames=='seed')]))

## Length of the HRRs
n <- 1024
## Number of possible fillers
nfillers <- as.integer(args[which(argnames=='nfillers')])
## Number of possible roles
nroles <- as.integer(args[which(argnames=='nroles')])
## Number of WM stripes
nstripes <- as.integer(args[which(argnames=='nstripes')])

## Define conj/disj scheme
state_cd <- args[which(argnames=='state_cd')]
sid_cd <- args[which(argnames=='sid_cd')]
interstripe_cd <- args[which(argnames=='interstripe_cd')]

## Identity vectors
hrr_z <- rep(0,n)
hrr_o <- rep(0,n)
hrr_o[1] <- 1

# Set this according to conj. disj. scheme;
if( interstripe_cd == 'D' ) {
    hrr_i <- hrr_z
} else {
    hrr_i <- hrr_o
}

## Filler vectors
## Fillers don't have friendly names, only indexes
## To index: [,filler]
fillers <- replicate(nfillers,hrr(n,normalized=TRUE))
wm_fillers <- replicate(nfillers,hrr(n,normalized=TRUE))
f_fillers <- paste('f',1:nfillers,sep='') # Friendly names

## Role vectors
## Roles don't have friendly names, only indexes
## To index: [,role]
roles <- replicate(nroles,hrr(n,normalized=TRUE))

## Stripe ID (SID) vectors
## To index: [,sid]
sids <- replicate(nstripes,hrr(n,normalized=TRUE))
use_sids_input <- args[which(argnames=='use_sids_input')] == 'T'
use_sids_output <- args[which(argnames=='use_sids_output')] == 'T'

## Op code vectors
## To index: [,1] for store, [,2] for retrieve
ops <- replicate(2,hrr(n,normalized=TRUE))

## Op code vectors
## To index: [,1] go, [,2] no go
gono <- replicate(2,hrr(n,normalized=TRUE))

## TD weight vectors
## A suffix of '_m' indicates relation to the maintenance layer NN
## To index: [,stripe]
W_m <- replicate(nstripes,hrr(n,normalized=TRUE))
W_o <- replicate(nstripes,hrr(n,normalized=TRUE))

## Action weight vectors
## There are 'nfillers' output units for this network
## Each output unit should come to represent the desired filler
## To index: [,output_unit]
W_a <- replicate(nfillers,hrr(n,normalized=TRUE))

## TD parameters
default_reward <- 0.0
success_reward <- 1.0
bias_m <- 1.0
bias_o <- 1.0
bias_a <- 1.0
gamma_m <- 1.0
gamma_o <- 1.0
gamma_a <- 1.0
lambda_m <- 0.9
lambda_o <- 0.9
lrate_m <- 0.1
lrate_o <- 0.1
lrate_a <- 0.1
epsilon_m <- .025
epsilon_o <- .025

## Task parameters
max_tasks <- as.integer(args[which(argnames=='ntasks')])
cur_task <- 1
cur_block_task <- 1
block_tasks_correct <- 0

## Get training and test sets
train_set_size <- 200
test_set_size <- 100

## Choose testing protocol
if( args[which(argnames=='gen_test')] == 'SG' ) {
    sets <- standardGeneralization(nroles,nfillers,train_set_size,test_set_size)
} else if( args[which(argnames=='gen_test')] == 'SA' ) {
    sets <- spuriousAnticorrelation(nroles,nfillers,train_set_size,test_set_size)
} else if( args[which(argnames=='gen_test')] == 'FC' ) {
    sets <- fullCombinatorial(nroles,nfillers,train_set_size,test_set_size)
} else if( args[which(argnames=='gen_test')] == 'NF' ) {
    sets <- novelFiller(nroles,nfillers,train_set_size,test_set_size)
}

#############################################################################
## selectAction:
##
## This function handles action selection. To keep things simple, there must
## be only one WM_out filler, otherwise a -1 action is returned.
##
#############################################################################
selectAction <- function() {
    action <- -1
    for( i in 1:nstripes ) {
        if( f_stripes_o[i] != 'I' )
            ## if there are more than one fillers return -1
            if( action != -1 )
                return(-1)
            else
                action <- as.integer(substr(f_stripes_o[i],2,nchar(f_stripes_o[i])))
    }
    return (action)
}

#############################################################################
## getState: Returns role/op combo
#############################################################################
getState <- function(o,r) {
    ## Encode state (role,op)
    if( state_cd == 'C' ) {
        state <- convolve(ops[,o],roles[,r])
    } else {
        state <- cnorm(ops[,o] + roles[,r])
    }
    #state <- roles[,r]
    return (state)
}

#############################################################################
## inputGate:
##
## This function handles input gating for a single timestep/operation in the task.
## At each timestep a stripe can either retain its contents or update with the
## provided filler.
##
## Multiple stripes update in a single timestep.
##
#############################################################################
inputGate <- function(state,f=-1) {

    ## Get NN output unit values
    ## Note that state_wm is for the previous timestep
    ## We will return this to use in the eligibility trace
    #state_wm <- convolve(state,cur_wm_m)

    ## This is state_wm convolved with both the go and no go hrrs
    #state_wm_gono <- apply(gono,2,convolve,state_wm)
    state_gono <- apply(gono,2,convolve,state)

    ## Build to matrix of eligility trace vectors that will be returned
    ## for TD updates; default to no go
    #elig <- replicate(nstripes,state_wm_gono[,2])
    elig <- replicate(nstripes,state_gono[,2])
    vals <- rep(0.0,nstripes)
    open <- rep(FALSE,nstripes)

    ## Determine if open or close value is better
    for( i in 1:nstripes ) {
        #temp_vals <- (apply(state_wm_gono,2,nndot,W_m[,i]) + bias_m)
        temp_vals <- (apply(state_gono,2,nndot,W_m[,i]) + bias_m)
        vals[i] <- max(temp_vals)

        ## Epsilon soft policy
        r <- runif(1)
        if( r < epsilon_m ) {
            ## Pick a random open or close state
            r <- runif(1)
            if( r > .5 ) {
                #elig[,i] <- state_wm_gono[,1]
                elig[,i] <- state_gono[,1]
                open[i] <- TRUE
                vals[i] <- temp_vals[1]
            } else {
                vals[i] <- temp_vals[2]
            }
        } else if( temp_vals[1] > temp_vals[2] ) {
            #elig[,i] <- state_wm_gono[,1]
            elig[,i] <- state_gono[,1]
            open[i] <- TRUE
        }
    }

    ## Update WM_m contents
    ## We are convolving the fill with the appropriate SID
    for( i in 1:nstripes ) {
        if( open[i] ) {
            if( f == -1 ) {
                stripes_m[,i] <- hrr_i
                stripes_mo[,i] <- hrr_i
                f_stripes_m[i] <- 'I'
            } else {
                if( sid_cd == 'C' ) {
                    if( use_sids_input )
                        stripes_m[,i] <- convolve(fillers[,f],sids[,i])
                    if( use_sids_output )
                        stripes_mo[,i] <- convolve(fillers[,f],sids[,i])
                } else {
                    if( use_sids_input )
                        stripes_m[,i] <- cnorm(fillers[,f]+sids[,i])
                    if( use_sids_output )
                        stripes_mo[,i] <- cnorm(fillers[,f]+sids[,i])
                }
                if( !use_sids_input )
                    stripes_m[,i] <- fillers[,f]
                if( !use_sids_output )
                    stripes_mo[,i] <- fillers[,f]
                f_stripes_m[i] <- f_fillers[f]
            }
        }
    }

    ## Get updated wm representation
    if( interstripe_cd == 'D' ) {
        wm <- cnorm(apply(stripes_m,1,sum))
        if( is.nan(wm[1]) )
            wm <- hrr_o
    } else {
        wm <- mconvolve(stripes_m)
    }

    return(list(
        wm = wm,
        elig = elig,
        vals = vals,
        stripes_m = stripes_m,
        stripes_mo = stripes_mo,
        f_stripes_m = f_stripes_m
    ))
}

#############################################################################
## outputGate:
##
## This function handles output gating for a single timestep/operation in the task.
## Each PFC stripe can either output its contents or not.
##
#############################################################################
outputGate <- function(state) {

    ## Start with a blank slate
    stripes_o <- replicate(nstripes,hrr_i)
    f_stripes_o <- replicate(nstripes,'I')

    ## Get NN output unit values
    ## Note that state_wm is for the previous timestep
    ## We will return this to use in the eligibility trace
    #state_wm <- convolve(state,cur_wm_m)

    ## This is state_wm convolved with both the go and no go hrrs
    #state_wm_gono <- apply(gono,2,convolve,state_wm)
    state_gono <- apply(gono,2,convolve,state)

    ## Build to matrix of eligility trace vectors that will be returned
    ## for TD updates; default to no go
    #elig <- replicate(nstripes,state_wm_gono[,2])
    elig <- replicate(nstripes,state_gono[,2])
    vals <- rep(0.0,nstripes)
    open <- rep(FALSE,nstripes)

    ## Determine if open or close value is better
    for( i in 1:nstripes ) {

        #temp_vals <- (apply(state_wm_gono,2,nndot,W_o[,i]) + bias_o)
        temp_vals <- (apply(state_gono,2,nndot,W_o[,i]) + bias_o)
        vals[i] <- max(temp_vals)

        ## Epsilon soft policy
        r <- runif(1)
        if( r < epsilon_o ) {
            ## Pick a random open or close state
            r <- runif(1)
            if( r > .5 ) {
                #elig[,i] <- state_wm_gono[,1]
                elig[,i] <- state_gono[,1]
                open[i] <- TRUE
                vals[i] <- temp_vals[1]
            } else {
                vals[i] <- temp_vals[2]
            }
        } else if( temp_vals[1] > temp_vals[2] ) {
            #elig[,i] <- state_wm_gono[,1]
            elig[,i] <- state_gono[,1]
            open[i] <- TRUE
        }
    }

    ## Update WM_o contents
    ## We are convolving the fill with the appropriate SID
    for( i in 1:nstripes ) {
        if( open[i] ) {
            stripes_o[,i] <- stripes_mo[,i]
            f_stripes_o[i] <- f_stripes_m[i]
        }
    }

    ## Get updated wm representation
    if( interstripe_cd == 'D' ) {
        wm <- cnorm(apply(stripes_o,1,sum))
        if( is.nan(wm[1]) )
            ## Change this depending on how we want 'nothing' to look for action selection
            #wm <- hrr_o
            wm <- hrr_z
    } else {
        wm <- mconvolve(stripes_o)
    }

    return(list(
        wm = wm,
        elig = elig,
        vals = vals,
        f_stripes_o = f_stripes_o
    ))
}

## Continue until we hit the max_tasks, or we have a block success rate of 95%
while( cur_task <= max_tasks ) {

    ## Setup 200 task blocks
    if( cur_task %% 200 == 1 ) {
        if( block_tasks_correct/200 >= .95 )
            break
        cur_block_task <- 1
        block_tasks_correct <- 0
    }

    #############################################################################
    ## Initialization and task setup
    #############################################################################

    reward <- default_reward
    elig_m <- replicate(nstripes,rep(0,n)) # There is an elig. trace for each BG input gate
    elig_o <- replicate(nstripes,rep(0,n)) # There is an elig. trace for each BG output gate
    prev_val_m <- rep(0.0,nstripes)     # prev. timestep values for maintenance NN output layer units
    prev_val_o <- rep(0.0,nstripes)     # prev. timestep values for output NN output layer units
    cur_wm_m <- hrr_o
    cur_wm_o <- hrr_o

    ## Fill stripes with identity vector
    stripes_m <- replicate(nstripes,hrr_i)  # Input layer stripes (hrrs)
    stripes_mo <- replicate(nstripes,hrr_i)  # Output layer stripes (hrrs); these are same as above, but with no SID
    f_stripes_m <- replicate(nstripes,'I')  # Friendly name for stripe contents of input WM layer
    f_stripes_o <- replicate(nstripes,'I')  # Friendly name for stripe contents of output WM layer

    #############################################################################
    ## Store fillers
    #############################################################################

    ## Choose a sentence at random from the sample set
    s_f <- sets$train_set[sample(train_set_size,1),]

    ## Permute the sentence so that roles are not always presented in the same order
    p <- sample(nroles,nroles,replace=FALSE)

    for( t in 1:nroles ) {
        #cat(sprintf('t=%d\n',t)) #debug

        state <- getState(1,p[t])

        #############################################################################
        ## Input gating
        #############################################################################

        ## Update WM input layer global variables
        ig <- inputGate(state,s_f[p[t]])
        cur_wm_m <- ig$wm
        stripes_m <- ig$stripes_m
        stripes_mo <- ig$stripes_mo
        f_stripes_m <- ig$f_stripes_m

        #############################################################################
        ## Output gating
        #############################################################################

        ## Update WM output layer global variables
        og <- outputGate(state)
        cur_wm_o <- og$wm
        f_stripes_o <- og$f_stripes_o

        #############################################################################
        ## Neural network and TD training for this trial
        #############################################################################

        ## INPUT GATE
        error <- (reward + gamma_m * ig$vals) - prev_val_m
        for( i in 1:nstripes ) {
            W_m[,i] <- W_m[,i] + lrate_m * error[i] * elig_m[,i]
            elig_m[,i] <- cnorm(lambda_m * elig_m[,i] + ig$elig[,i])
        }
        prev_val_m <- ig$vals

        ## OUTPUT GATE
        error <- (reward + gamma_o * og$vals) - prev_val_o
        for( i in 1:nstripes ) {
            W_o[,i] <- W_o[,i] + lrate_o * error[i] * elig_o[,i]
            elig_o[,i] <- cnorm(lambda_o * elig_o[,i] + og$elig[,i])
        }
        prev_val_o <- og$vals

    }
    #cat(sprintf('t=%d\n',t+1)) #debug

    #############################################################################
    ## Query for roles
    #############################################################################

    ## Permute the sentence so that roles are not always queried in the same order
    nqueries <- 1
    p <- sample(nroles,nqueries,replace=FALSE)
    correct_trial <- rep(0,nqueries)

    for( t in 1:nqueries ) {
        state <- getState(2,p[t])

        #############################################################################
        ## Input gating
        #############################################################################

        ## Update WM input layer global variables
        ig <- inputGate(state)
        cur_wm_m <- ig$wm
        stripes_m <- ig$stripes_m
        stripes_mo <- ig$stripes_mo
        f_stripes_m <- ig$f_stripes_m

        #############################################################################
        ## Output gating
        #############################################################################

        ## Update WM output layer global variables
        og <- outputGate(state)
        cur_wm_o <- og$wm
        f_stripes_o <- og$f_stripes_o

        #############################################################################
        ## Action selection
        #############################################################################

        ac <- selectAction()

        #############################################################################
        ## Neural network and TD training for this trial
        #############################################################################

        ## INPUT GATE
        error <- (reward + gamma_m * ig$vals) - prev_val_m
        for( i in 1:nstripes ) {
            W_m[,i] <- W_m[,i] + lrate_m * error[i] * elig_m[,i]
            elig_m[,i] <- cnorm(lambda_m * elig_m[,i] + ig$elig[,i])
        }
        prev_val_m <- ig$vals

        ## OUTPUT GATE
        error <- (reward + gamma_o * og$vals) - prev_val_o
        for( i in 1:nstripes ) {
            W_o[,i] <- W_o[,i] + lrate_o * error[i] * elig_o[,i]
            elig_o[,i] <- cnorm(lambda_o * elig_o[,i] + og$elig[,i])
        }
        prev_val_o <- og$vals

        ## Determine correctness
        ## The trial is correct of the selected action matches the filler that
        ## was paired with the requested role.
        if( ac == s_f[p[t]] )
            correct_trial[t] <- 1

        #############################################################################
        ## Absorb reward
        #############################################################################
        if( t == nqueries ) {

            ## Reward if entire sequence is correct
            ## Update block correct tally
            if( sum(correct_trial) == nqueries ) {
                block_tasks_correct <- block_tasks_correct + 1
                reward <- success_reward
            } else {
                reward <- default_reward
            }

            ## Input NN
            error <- reward - ig$vals
            for( i in 1:nstripes ) {
                W_m[,i] <- W_m[,i] + lrate_m * error[i] * elig_m[,i]
            }

            ## Output NN
            error <- reward - og$vals
            for( i in 1:nstripes ) {
                W_o[,i] <- W_o[,i] + lrate_o * error[i] * elig_o[,i]
            }
        }
    }

    #############################################################################
    ## Output prints
    #############################################################################

    #if( FALSE ) {
    if( cur_task %% 200 == 0 ) {
        cat(sprintf('Tasks Complete: %d\n',cur_task))
        cat(sprintf('Block Accuracy: %.2f\n',(block_tasks_correct/200)*100))

        ## Only printing final request state here
        cat('Input WM Layer: \t')
        cat(f_stripes_m)
        cat('\n')
        cat('Output WM Layer: \t')
        cat(f_stripes_o)
        cat('\n')
        cat(sprintf('Requested Role: %d\n',p[t]))
        cat(sprintf('Correct Action: %d\n',s_f[p[t]]))
        cat('\n')
    }

    #############################################################################
    ## Task wrapup
    #############################################################################

    ## Increment task tally
    cur_task <- cur_task + 1
}

## Print final results
#cat(sprintf('%d\n',cur_task))
cat(sprintf('Final Block Accuracy: %.2f\n',(block_tasks_correct/200)*100))

#############################################################################
## Generalization Test
#############################################################################
if(TRUE) {
    novel_tasks_correct <- 0
    for( i in 1:test_set_size ) {
        correct_trial <- rep(0,nqueries)

        #############################################################################
        ## Store fillers
        #############################################################################

        ## Retrieve novel sentence from the test set
        ## We aren't training here so no need to permute
        s_f <- sets$test_set[i,]

        ## Permute the sentence so that roles are not always presented in the same order
        p <- sample(nroles,nroles,replace=FALSE)

        ## Do a 'Store' for each filler
        ## Action selection can be skipped here
        ## Do not do any training
        for( t in 1:nroles ) {
            state <- getState(1,p[t])

            ig <- inputGate(state,s_f[p[t]])
            cur_wm_m <- ig$wm
            stripes_m <- ig$stripes_m
            stripes_mo <- ig$stripes_mo
            f_stripes_m <- ig$f_stripes_m

            og <- outputGate(state)
            cur_wm_o <- og$wm
            f_stripes_o <- og$f_stripes_o
        }

        ## Do a 'Retrieve' input gate, output gate, and select action
        p <- sample(nroles,nqueries,replace=FALSE)
        for( t in 1:nqueries ) {

            ## We can only query role 1 for the FC test
            ## consequently, nqueries must be set to 1
            state <- if(args[which(argnames=='gen_test')] == 'FC') getState(2,1) else getState(2,p[t])
            answer <- if(args[which(argnames=='gen_test')] == 'FC') s_f[1] else s_f[p[t]]

            ig <- inputGate(state)
            cur_wm_m <- ig$wm
            stripes_m <- ig$stripes_m
            stripes_mo <- ig$stripes_mo
            f_stripes_m <- ig$f_stripes_m

            og <- outputGate(state)
            cur_wm_o <- og$wm
            f_stripes_o <- og$f_stripes_o

            ac <- selectAction()
            if( ac == answer )
                correct_trial[t] <- 1
        }

        ## Determine if entire sequence is correct
        ## Not sure if we'll want to modify protocol above to train for this
        if( sum(correct_trial) == nqueries )
            novel_tasks_correct <- novel_tasks_correct + 1
    }

    ## Print final results
    cat(sprintf('Generalization Accuracy: %d\n',novel_tasks_correct))
    #cat(sprintf('%d\n',novel_tasks_correct))
} ## End generalization test