#' generate a batch stat table for a given single cell umi matrix
#' @param scmat - the matrix to work on
#' @return batch stats table
#'
#' @export

mcell_batch_stats =  function(scmat)
{
	md = scmat@cell_metadata
	umis = scmat@mat
	# report batch stats
	batch_stats = data.frame()
	batch_nms = sort(unique(md$amp_batch_id))
	for(b in batch_nms) {
		batch_stats = rbind(batch_stats, mcell_calc_one_batch_stats(scmat, b))
	}
	rownames(batch_stats) = batch_nms

	return(batch_stats)
}
#	write.table(batch_stats, sprintf("%s/batch_stats.txt", get_param("outdir")), quote=F, sep="\t")
#	.plot_batch_stats(batch_stats, batch_meta , sprintf("%s/batch_stats.png", get_param("outdir")))

#' calc batch stats - essentially umi distribution
#'
#' @param scmat - a matrix object
#' @param batch_id - the batch name to summarize
#'
#' @return one row data frame for merging
#' @export
#'
mcell_calc_one_batch_stats = function(scmat, batch_id)
{
	md = scmat@cell_metadata
	mat = scmat@mat

	f_batch = md$amp_batch_id == batch_id

	n_c = colSums(mat[,f_batch])

	data.frame(
		total = sum(n_c),
		total_ercc = sum(md$spike_count[f_batch]),
		n_lt50 = sum(n_c < 50),
		n_50_100 = sum(n_c >=  50 & n_c < 100),
		n_100_200 = sum(n_c >=  100 & n_c < 200),
		n_200_400 = sum(n_c >=  200 & n_c < 400),
		n_400_800 = sum(n_c >=  400 & n_c < 800),
		n_800_1600 = sum(n_c >=  800 & n_c < 1600),
		n_gt1600 = sum(n_c >= 1600),
		n_valid = length(n_c),
		valid_q25 = quantile(n_c, 0.25),
		valid_q50 = quantile(n_c, 0.5),
		valid_q75 = quantile(n_c, 0.75)
		)
}

#' plot batches stats
#'
#' @param batches_stats collected by .calc_batch_stats
#' @param fig_nm output file name
#' @param cells_breakdown_field color #cells per batch by this field
#' @param cell_col_dict color dictionary for cells_breakdown_field values
#'
#' @export
#'
mcell_plot_batch_stats = function(mat_id, fig_nm = NULL)
{
	scmat = scdb_mat(mat_id)
	if(is.null(fig_nm)) {
		fig_nm = scfigs_fn(mat_id, "batchstat")
	}

	cells_breakdown_field=get_param("scm_cells_breakdown_field")
	cell_col_dict=get_param("scm_cell_col_dict")

	batch_stats = mcell_batch_stats(scmat)
	md = scmat@cell_metadata
	batch_meta = unique(md[, grep("spike_count", colnames(md), invert=T)])

	rownames(batch_meta) = batch_meta$amp_batch_id

	batch_stats = batch_stats[order(gsub("AB", "", rownames(batch_stats)), decreasing = T), ]
	batch_meta = batch_meta[rownames(batch_stats), ]

	b_tot = batch_stats$total
	names(b_tot) = rownames(batch_stats)

	main_cex = 2
	bp_cex = 1

	.plot_start(fig_nm, w=960, h=max(800, 100 + 10 * length(b_tot)))
	layout(matrix(1:4, 1, 4), w=c(2,1,5,1))

	type_cols = 'blue'
	if (!is.null(cell_col_dict) & !is.null(cells_breakdown_field)) {
		type_cols = unlist(cell_col_dict[batch_meta[, cells_breakdown_field]])
	}

	par(mar=c(4,5,2,0))
	barplot(b_tot, horiz=T, las=2, cex.axis=bp_cex, cex.names=bp_cex, border=NA, col=type_cols)
	title(main='total UMIs', cex.main=main_cex)
	if (!is.null(cell_col_dict) & !is.null(cells_breakdown_field)) {
		legend("topleft", legend=names(type_cols), bty='n', border=NA, fill=unlist(cell_col_dict), cex=bp_cex)
	}
	box(lwd=0.5)

	par(mar=c(4,1,2,0))
	barplot(batch_stats$total_ercc/batch_stats$total, cex.axis=bp_cex, horiz=T, border=NA, col='black')
	title(main='%ERCC', cex.main=main_cex)
	box(lwd=0.5)

	bd_cols = RColorBrewer::brewer.pal(n=7, 'Blues')
	barplot(t(batch_stats[,3:9]), horiz=T, border=NA, cex.axis=bp_cex, col=bd_cols, yaxt='n')
	title(main='#cells by #UMIs', cex.main=main_cex)
	legend("topleft", legend=gsub("_", "-", gsub("n_", "", colnames(batch_stats)[3:9])), bty='n', fill=bd_cols, cex=1.2 * bp_cex, pt.cex=1.2 * bp_cex, ncol=length(bd_cols), border=NA, x.intersp=0.8)
	box(lwd=0.5)

	par(mar=c(4,1,2,1))
	barplot(batch_stats$n_valid, horiz=T, border=NA, cex.axis=bp_cex, col='black')
	title(main='#valid', cex.main=main_cex)
	box(lwd=0.5)

	dev.off()
}

#' Plot histogram of total number of umis per cell in the umis matrix
#'
#' @param mat_id id of the mat to use
#' @param min_umis_cutoff If not null, will mark this cutoff by a line. If NA, auto calculate by finding the first bin in the histogram with more cells then its predecessor.
#' @param bin_for_cutoff bin size for histogram used to auto find min_umis_cutoff
#'
#' @return min_umis_cutoff (either the given, the calculated, or null)
#'
#' @export
#'
mcell_plot_umis_per_cell = function(mat_id, min_umis_cutoff = NA, bin_for_cutoff = 50)
{
  mat = scdb_mat(mat_id)

  if (is.null(mat)) {
    stop(sprintf("MCERR: mat with id %s not found", mat_id))
  }

  uc = Matrix::colSums(mat@mat)
  .plot_start(scfigs_fn(mat_id, "total_umi_distr"), w=450, h=450)

  if (is.na(min_umis_cutoff)) {
    h = hist(uc, breaks=seq(0, max(uc) + bin_for_cutoff, by=bin_for_cutoff), plot=F)
    min_umis_cutoff = min(which(diff(h$counts) > 0)) * bin_for_cutoff
  }
  h = hist(log2(uc), 200, xlab='total cell UMIs (log2)', main=mat_id)

  if (!is.null(min_umis_cutoff)) {
    abline(v=log2(min_umis_cutoff), col='red', lty=2)
    text(log2(min_umis_cutoff), max(h$count)/2, min_umis_cutoff, pos=2, col='red')
  }

  dev.off()

  return(min_umis_cutoff)
}
#


#!/usr/bin/env Rscript

#
# Runs the social network analysis using Peter Hoff's GBME
#
suppressMessages(library("optparse"))
suppressMessages(library("xtable"))
suppressMessages(library("R.utils"))

cargs = commandArgs(trailingOnly = FALSE)
source_dir = dirname(sub("^--file=", "", cargs[grep("^--file=", cargs)]))
source(file.path(source_dir, "gbme.r"))

option_list = list(
  make_option(
    c("-f", "--file"),
    default = "",
    type = "character",
    help = "Similarity matrix",
    metavar = "character"
  ),
  make_option(
    c("-o", "--out_dir"),
    default = "",
    type = "character",
    help = "outdir",
    metavar = "character"
  ),
  make_option(
    c("-s", "--sna"),
    default = "sna-gbme.pdf",
    type = "character",
    help = "GBME SNA output filename",
    metavar = "character"
  ),
  make_option(
    c("-n", "--number"),
    default = as.integer(20000),
    type = "integer",
    help = "Number of GBME iterations",
    metavar = "integer"
  ),
  make_option(
    c("-a", "--alias"),
    default = "",
    type = "character",
    help = "Sample name alias file",
    metavar = "alias"
  )
);

opt_parser   = OptionParser(option_list = option_list);
opt          = parse_args(opt_parser);
matrix.file  = opt$file
out.dir      = opt$out_dir
n_iter       = opt$number
alias.file   = opt$alias
sna.filename = opt$sna

Y = as.matrix(read.table(matrix.file, header = TRUE))

if (nchar(matrix.file) == 0) {
  stop("Missing matrix file argument")
}

if (!file.exists(matrix.file)) {
  stop(paste("Invalid --matrix", matrix.file))
}

if (nchar(out.dir) == 0) {
  out.dir = dirname(matrix.file)
}

if (!dir.exists(out.dir)) {
  printf("Creating outdir '%s'\n", out.dir)
  dir.create(out.dir)
}
out.dir = normalizePath(out.dir)
setwd(out.dir)

GBME_OUT = file.path(out.dir, "gbme.out")
if (file.exists(GBME_OUT)) {
  file.remove(GBME_OUT)
}

# Look for the "*.meta" files
meta_dir = file.path(out.dir, "meta")

printf("Reading %s matrix\n", matrix.file)
Y = as.matrix(read.table(matrix.file, header = TRUE))
n = nrow(Y)
Xss = NULL
k = 0

if (dir.exists(meta_dir)) {
  meta_files = list.files(path = meta_dir, pattern = "*.meta$")
  print(meta_files)
  k = length(meta_files)

  printf("Found %s file%s in meta dir '%s'\n", k,
         if (k == 1) { '' } else {'s'}, meta_dir)

  if (k > 0) {
    Xss = array(NA, dim = c(n,n,k))

    for (i in 1:k) {
      file = file.path(meta_dir, meta_files[i])
      d = as.matrix(read.table(file, header = TRUE))
      if (nrow(d) == n) {
        Xss[,,i] = d
      }
    }
  }
}

printf("Running GBME with %s scans\n", n_iter)
Z_path = file.path(out.dir, "Z")

if (is.null(Xss)) {
  gbme(Y = Y, fam = "gaussian", k = 2, direct = F, NS = n_iter, odens = 10, ofilename = GBME_OUT, zfilename = Z_path)
} else {
  tryCatch(gbme(Y = Y, Xss, fam = "gaussian", k = 2, direct = F, NS = n_iter, odens = 10, ofilename = GBME_OUT, zfilename = Z_path),
           error = function(c) {
               stop("Error running GBME with metadata -- TRY WITHOUT!")
           }
           )
}

printf("looking for %s", GBME_OUT)
if (!file.exists(GBME_OUT)) {
  stop("Failed to run GBME!")
}

printf("Using GBME_OUT '%s'\n", GBME_OUT)
OUT = read.table(GBME_OUT, header = T)

# posterior samples, dropping
# the first half of the chain
# to allow for burn in
PS <- OUT[OUT$scan > round(max(OUT$scan) / 2),-(1:3)]

#
# analysis of latent positions
#
Z <- read.table(Z_path)

#
# convert to an array
#
nss <- dim(OUT)[1]
n <- dim(Z)[1] / nss
k <- dim(Z)[2]
PZ <- array(dim = c(n,k,nss))
for (i in 1:nss) {
  PZ[,,i] <- as.matrix(Z[((i - 1) * n + 1):(i * n) ,])
}

PZ <- PZ[,,-(1:round(nss / 2))] #drop first half for burn in

#
# find posterior mean of Z %*% t(Z)
#
ZTZ <- matrix(0,n,n)
for (i in 1:dim(PZ)[3]) {
  ZTZ <- ZTZ + PZ[,,i] %*% t(PZ[,,i])
}
ZTZ <- ZTZ / dim(PZ)[3]

#
# a configuration that approximates posterior mean of ZTZ
#
tmp <- eigen(ZTZ)
Z.pm <- tmp$vec[, 1:k] %*% sqrt(diag(tmp$val[1:k]))

#now transform each sample Z to a common orientation
for (i in 1:dim(PZ)[3]) {
  PZ[,,i] <- proc.rr(PZ[,,i],Z.pm)
}

#
# a two dimensional plot of "mean" latent locations
# and marginal confidence regions
#
r <- atan2(Z.pm[,2], Z.pm[,1])
r <- r + abs(min(r))
r <- r / max(r)
g <- 1 - r
b <- (Z.pm[,2] ^ 2 + Z.pm[,1] ^ 2)
b <- b / max(b)

par(mfrow = c(1,1))
sna_path = file.path(out.dir, sna.filename)
printf('Creating "%s"\n', sna_path)
pdf(sna_path, width = 6, height = 6)
plot(
  Z.pm[,1],
  Z.pm[,2],
  xlab = "",
  ylab = "",
  type = "n",
  xlim = range(PZ[,1,]),
  ylim = range(PZ[,2,])
)

abline(h = 0,lty = 2);abline(v = 0,lty = 2)

for (i in 1:n) {
  points(PZ[i,1,],
         PZ[i,2,],
         pch = 46,
         col = rgb(r[i], g[i], b[i]))
}

# add labels here
labels = rownames(Y)
if (length(alias.file) > 0 & file.exists(alias.file)) {
  printf("Using alias file '%s'\n", alias.file)
  aliases = read.table(alias.file, header=T, as.is=T)
  for (i in 1:length(labels)) {
    label = labels[i]
    alias = aliases[aliases$name == label, "alias"]
    if (length(alias) > 0) {
      printf("Alias '%s' -> '%s'\n", label, alias)
      labels[i] = alias
    }
  }
}

text(Z.pm[,1],Z.pm[,2], cex = 0.3, labels = labels)
invisible(dev.off())

printf("Done, see output dir '%s'\n", out.dir)


  ggplot(df, aes(x = EV_DiffAbs, y = rt)) +
         geom_smooth(method="lm", alpha = .25, size = 1.5) +
         labs(title = "Expected Value Difference Predicts Decision Reaction Time",
              subtitle = "Absolutle Difference in Expected Value Has No Effect on Reaction Time.",
              x = "Absolute Difference in Expected Value",
              y ="Reaction time (in seconds)",
              caption = "I'm, like, a total R whiz") +
         coord_cartesian(ylim=c(0, 2)) +
         theme_classic() +
         theme(plot.title = element_text(face="bold", size=13, hjust = 1.0)) +
         theme(plot.subtitle = element_text(face = "italic", size = 10, hjust = 1.0)) +
         theme(plot.caption = element_text(face = "italic", size = 8, hjust = 0.0)) +
         theme(axis.title = element_text(size = 12)) +
         theme(axis.text.x = element_text(size = 14, color = "Blue")) +
         theme(axis.text.y = element_text(size = 14, color = "Blue"))


# Pattern Matching Problem

library("seqinr")
genome <- read.fasta(file = "vibrio_cholerae.fasta",
                     as.string = TRUE)

genome1 <- as.character(genome$NC_002505.1)
genome2 <- as.character(genome$NC_002506.1)

genome1 <- toupper(genome1)
genome2 <- toupper(genome2)

k_mer1 <- gregexpr("(?=ATGATCAAG)" , genome1, perl = TRUE)
k_mer2 <- gregexpr("(?=ATGATCAAG)" , genome2, perl = TRUE)

result1 <- k_mer1[[1]]
result2 <- k_mer2[[1]]


context("compute_model_prediction")

test_that("compute_model_prediction preserves datetimes", {
  # Data frame with POSIXct, and zigzag values
  dat <- data.frame(
    d = as.POSIXct('2001-06-11 21:00', tz = 'UTC') + seq(1, 1000, by = 10),
    value = 1:100 + rep(c(-3, 3), 50)
  )
  res <- dat %>% compute_model_prediction(value ~ d, n = 10, model = "lm")
  expect_equal(range(dat$d), range(res$pred_))
})

test_that("compute_model_prediction preserves dates", {
  # Data frame with Date, and zigzag values
  dat <- data.frame(
    d = as.Date('2001-06-11') + seq(1, 1000, by = 10),
    value = 1:100 + rep(c(-3, 3), 50)
  )
  res <- dat %>% compute_model_prediction(value ~ d, n = 10, model = "lm")
  expect_equal(range(dat$d), range(res$pred_))
})

test_that("compute_model_prediction works with datetimes", {
  # Perfectly linear data
  dat <- data.frame(
    d = as.POSIXct('2001-06-11 21:00', tz = 'America/New_York') + 1:10 * 100,
    value = 1:10
  )

  # Tests with various models
  res <- dat %>% compute_model_prediction(value ~ d, n = 10, model = "loess")
  expect_equal(range(dat$d), range(res$pred_))
  expect_equal(attr(dat$d, "tzone"), attr(res$pred_, "tzone"))
  expect_equal(range(dat$value), range(res$resp_))

  res <- dat %>% compute_model_prediction(value ~ d, n = 10, model = "lm")
  expect_equal(range(dat$d), range(res$pred_))
  expect_equal(attr(dat$d, "tzone"), attr(res$pred_, "tzone"))
  expect_equal(range(dat$value), range(res$resp_))

  res <- dat %>% compute_model_prediction(value ~ d, n = 10, model = "glm")
  expect_equal(range(dat$d), range(res$pred_))
  expect_equal(attr(dat$d, "tzone"), attr(res$pred_, "tzone"))
  expect_equal(range(dat$value), range(res$resp_))
})

test_that("compute_model_prediction works with more complex formulas", {
  dat <- data.frame(x = 1:10, y = (1:10 - 5)^2 + 4 * 1:10 + 100)
  res <- dat %>%
          compute_model_prediction(y ~ I(x^2) + x, n = 10, model = "lm") %>%
          setNames( c("x", "y"))
  expect_equal(dat, res)

  dat <- data.frame(x = 1:10, y = 2.5*(1:10)^3 + 7*(1:10)^2 + 4*(1:10) + 100)
  res <- dat %>%
          compute_model_prediction(y ~ poly(x, 3), n = 10, model = "lm") %>%
          setNames( c("x", "y"))
  expect_equal(dat, res)
})

test_that("Can control domain", {
  dat <- data.frame(x = 1:10, y = 5 * (1:10))
  res <- dat %>%
          compute_model_prediction(y ~ x, n = 10, model = "lm", domain = c(11, 20)) %>%
          setNames( c("x", "y"))
  expect_equal(res, data.frame(x = 11:20, y = 5 * (11:20)))
})

test_that("Zero-row inputs", {
  res <- mtcars[0,] %>% compute_model_prediction(wt ~ mpg, model = "lm")
  expect_equal(nrow(res), 0)
  expect_true(setequal(names(res), c("pred_", "resp_")))

  res <- mtcars[0,] %>% compute_model_prediction(wt ~ mpg, model = "lm", se = TRUE)
  expect_equal(nrow(res), 0)
  expect_true(setequal(
    names(res),
    c("pred_", "resp_", "pred_lwr_", "pred_upr_", "pred_se_" )
  ))

  # Smooth
  res <- mtcars[0,] %>% compute_smooth(wt ~ mpg)
  expect_equal(nrow(res), 0)
  expect_true(setequal(names(res), c("pred_", "resp_")))

  # Grouped
  res <- mtcars[0,] %>% group_by(cyl) %>%
    compute_model_prediction(wt ~ mpg, model = "lm", se = FALSE)
  expect_true(setequal(names(res), c("cyl", "pred_", "resp_")))

  res <- mtcars[0,] %>% group_by(cyl) %>%
    compute_model_prediction(wt ~ mpg, model = "lm", se = TRUE)
  expect_true(setequal(
    names(res),
    c("cyl", "pred_", "resp_", "pred_lwr_", "pred_upr_", "pred_se_" )
  ))
})


# variables
wd = "/Storage/data1/jorge.munoz/NRGSC"
# set working directory
setwd(wd)

# libraries
library(tidyverse, lib.loc = "./libraries")
library(DESeq2, lib.loc = "./libraries")
library(pheatmap, lib.loc = "./libraries")
library(viridis, lib.loc = "./libraries")

# function to graph heatmaps from contrast data

heatmaps <- function (vst_path, title, contrast_table) {

  # read vst file
  load(vst_path)

  sample_table <-read.table("./metadata_complete.csv", sep = ",", header = T)
  sample_table$Condition <- as.factor((sample_table$Condition))
  sample_table$Genotype <- as.factor((sample_table$Genotype))
  sample_table$DevStage <- as.factor((sample_table$DevStage))
  sample_table$Individual <- as.factor((sample_table$Individual))
  sample_table$Group <- as.factor(paste(sample_table$Genotype, '_', sample_table$Condition, '_', sample_table$DevStage, sep=''))
  sample_table <- sample_table[vst$Sample.Number,]

  dea_contrast <- read.table(contrast_table, header = T, sep = ",")
  df <- as.data.frame((assay(vst)))
  annotation_col <- sample_table
  anot <- select(annotation_col, "Group","DevStage","Genotype", "Condition")

  dat <- df[dea_contrast$names,]

  rownames(anot) <- colnames(dat)
  #heat map complete with all genes
  png(paste("./DEA/by_groups/hm_", title, ".png", sep = ""),  width = 15*1.3, height = 15, res = 320, units = "cm")
  pheatmap(dat, main = title, annotation_col = anot, show_rownames = F , cellwidth = 20, col = inferno(299))
  dev.off()

  # upregulated genes
  upregulated <- dea_contrast[tail(order(dea_contrast$log2FoldChange),20),]
  dat <- as.data.frame(df[upregulated$names,])
  png(paste("./DEA/by_groups/upregulated_", title, ".png", sep = ""),  width = 15*1.3, height = 15, res = 320, units = "cm")
  pheatmap(dat, main =  paste( title," upregulated", sep = ""), annotation_col = anot, show_rownames = T, fontsize_row = 8, cellwidth = 20, col = inferno(299))
  dev.off()

  # downregulated genes
  downregulated <- dea_contrast[head(order(dea_contrast$log2FoldChange),20),]
  dat <- as.data.frame((df[downregulated$names,]))
  png(paste("./DEA/by_groups/downregulated_", title, ".png", sep = ""),  width = 15*1.3, height = 15, res = 320, units = "cm")
  pheatmap(dat, main = paste( title," downregulated", sep = ""), annotation_col = anot, show_rownames = T, fontsize_row = 8, cellwidth = 20, col = inferno(299))
  dev.off()
}

# call function with different contrast

# NR_270_B0_vs_NR_10_B0
heatmaps("./DEA/by_groups/NR_270_B0_vs_NR_10_B0_vst.Rdata", "NR_270_B0 vs NR_10_B0",  "DEA/by_groups/filter_3NR_270_B0_vs_NR_10_B0.csv")

# NR_270_B vs NR_10_B
heatmaps("./DEA/by_groups/NR_270_B_vs_NR_10_B_vst.Rdata", "NR_270_B vs NR_10_B",  "DEA/by_groups/filter_3NR_270_B_vs_NR_10_B.csv")

# NR_270_M_vs_NR_10_M
heatmaps("./DEA/by_groups/NR_270_M_vs_NR_10_M_vst.Rdata", "NR_270_M vs NR_10_M",  "DEA/by_groups/filter_3NR_270_M_vs_NR_10_M.csv")

# NR_270_P_vs_NR_10_P
heatmaps("./DEA/by_groups/NR_270_P_vs_NR_10_P_vst.Rdata", "NR_270_P vs NR_10_P",  "DEA/by_groups/filter_3NR_270_P_vs_NR_10_P.csv")

# R_270_B0_vs_R_10_B0
heatmaps("./DEA/by_groups/R_270_B0_vs_R_10_B0_vst.Rdata", "R_270_B0 vs R_10_B0",  "DEA/by_groups/filter_3R_270_B0_vs_R_10_B0.csv")

# R_270_B_vs_R_10_B
heatmaps("./DEA/by_groups/R_270_B_vs_R_10_B_vst.Rdata", "R_270_B vs R_10_B",  "DEA/by_groups/filter_3R_270_B_vs_R_10_B.csv")

# R_270_M_vs_R_10_M
heatmaps("./DEA/by_groups/R_270_M_vs_R_10_M_vst.Rdata", "R_270_M vs R_10_M",  "DEA/by_groups/filter_3R_270_M_vs_R_10_M.csv")

# R_270_P_vs_R_10_P
heatmaps("./DEA/by_groups/R_270_P_vs_R_10_P_vst.Rdata", "R_270_P vs R_10_P",  "DEA/by_groups/filter_3R_270_P_vs_R_10_P.csv")


#########################################################################################
#
# Mark Cembrowski, Janelia Research Campus, April 15 2015
#
# This script plots, in detail, the expression of a given gene.
#
# INPUT:
# theGene: character.  Symbol or ID of gene to plot.
# replicates=F: logical.  Visualise individual replicate FPKM values as data points.
# mask=c(): character.  Sample names to hide from final visualisation.
# rMax=-1: numeric.  Maximum y value.
#
#########################################################################################

fpkmBarplot <- function(theGene,replicates=F,mask=c(),unmask=c(),rMax=-1){
	if(length(mask)*length(unmask)>0.1){
		stop('Cannot use mask and unmask options at same time.')
	}

	if(!isGeneId(theGene)){
		# Transform to Ensembl ID.
		theGene <- symToId(theGene)[[1]] # Successful hits.
	}

	# Get FPKM value for gene.
	df <- subset(fpkmPool,gene_id==theGene)

	# If requested replicates, get replicates.
	if(replicates){
		dfRep <- subset(fpkmRep,gene_id==theGene)
	}

	# Mask samples, if selected.
	if(length(mask)>0.1){
		toMask <- grep(paste(mask,collapse='|'),df$sample_name)
		if(length(toMask)<0.1){
			stop('All requested masked sample(s) not found.')
		}

		df <- df[-toMask,]

		if(replicates){
			toMask <- grep(paste(mask,collapse='|'),dfRep$sample_name)
			dfRep <- dfRep[-toMask,]
		}
	}

	# Show only unmasked samples, if selected.  This operation is the opposite
	# of mask; ie, only the unmasked samples will be shown.
	if(length(unmask)>0.1){
		toUnmask <- grep(paste(unmask,collapse='|'),df$sample_name)
		if(length(toUnmask)<0.1){
			stop('All requested unmasked sample(s) not found.')
		}

		df <- df[toUnmask,]

		if(replicates){
			toUnmask <- grep(paste(unmask,collapse='|'),dfRep$sample_name)
			dfRep <- dfRep[toUnmask,]
		}
	}
	# Assign ordering to specific groups.
	theLimits <- c('dg_d','dg_v','ca3_d','ca3_v','ca2','ca1_d','ca1_i','ca1_v','pv','sst')

	# Account for masked samples.
	if(length(mask)>0.1){
		for (ii in 1:length(mask)){
			curMask <- mask[ii]
			toMask <- theLimits==curMask
			theLimits <- theLimits[!toMask]
		}
	}

	# Account for unmasked samples.
	if(length(unmask)>0.1){
		theLimits <- unmask
	}

	# Plot.
	p <- ggplot(df,aes(x=sample_name,y=fpkm))
	p <- p + geom_bar(stat='identity',aes(fill=sample_name,colour=sample_name))
	if(replicates){
		p <- p + geom_point(aes(x=sample_name,y=fpkm),size=3,
			shape=18,colour='black',data=dfRep)
	}

	p <- p + geom_errorbar(aes(ymin=conf_lo,ymax=conf_hi,group=1),colour='black',width=0.5)

	if(rMax>0){
		p <- p + coord_cartesian(ylim=c(-rMax/100,rMax)) # -rMax/100 to avoid clipping.
	}

	p <- p + xlab('Cell type') + ylab('FPKM')
	p <- p + ggtitle(idToSym(theGene))
	p <- p + theme_bw() + theme(legend.position='none')
	p <- p + scale_x_discrete(limits=theLimits)

	print(p)
}


# https://benalexkeen.com/creating-a-timeline-graphic-using-r-and-ggplot2/
library(ggplot2)
library(scales)
library(lubridate)

df <- read.csv('/Users/andrea/Documents/Works/GitHub/RExercises/Samples/milestones.csv')

df$date <- with(df, ymd(sprintf('%04d%02d%02d', year, month, 1)))
df <- df[with(df, order(date)), ]

status_levels <- c("Complete", "On Target", "At Risk", "Critical")
status_colors <- c("#0070C0", "#00B050", "#FFC000", "#C00000")

df$status <- factor(df$status, levels=status_levels, ordered=TRUE)

positions <- c(0.5, -0.5, 1.0, -1.0, 1.5, -1.5)
directions <- c(1, -1)

line_pos <- data.frame(
    "date"=unique(df$date),
    "position"=rep(positions, length.out=length(unique(df$date))),
    "direction"=rep(directions, length.out=length(unique(df$date)))
)

df <- merge(x=df, y=line_pos, by="date", all = TRUE)
df <- df[with(df, order(date, status)), ]

text_offset <- 0.05

df$month_count <- ave(df$date==df$date, df$date, FUN=cumsum)
df$text_position <- (df$month_count * text_offset * df$direction) + df$position
head(df)

month_buffer <- 2

month_date_range <- seq(min(df$date) - months(month_buffer), max(df$date) + months(month_buffer), by='month')
month_format <- format(month_date_range, '%b')
month_df <- data.frame(month_date_range, month_format)

year_date_range <- seq(min(df$date) - months(month_buffer), max(df$date) + months(month_buffer), by='year')
year_date_range <- as.Date(
    intersect(
        ceiling_date(year_date_range, unit="year"),
        floor_date(year_date_range, unit="year")
    ),  origin = "1970-01-01"
)
year_format <- format(year_date_range, '%Y')
year_df <- data.frame(year_date_range, year_format)

#### PLOT ####

timeline_plot<-ggplot(df,aes(x=date,y=0, col=status, label=milestone))
timeline_plot<-timeline_plot+labs(col="Milestones")
timeline_plot<-timeline_plot+scale_color_manual(values=status_colors, labels=status_levels, drop = FALSE)
timeline_plot<-timeline_plot+theme_classic()

# Plot horizontal black line for timeline
timeline_plot<-timeline_plot+geom_hline(yintercept=0,
                color = "black", size=0.3)

# Plot vertical segment lines for milestones
timeline_plot<-timeline_plot+geom_segment(data=df[df$month_count == 1,], aes(y=position,yend=0,xend=date), color='black', size=0.2)

# Plot scatter points at zero and date
timeline_plot<-timeline_plot+geom_point(aes(y=0), size=3)

# Don't show axes, appropriately position legend
timeline_plot<-timeline_plot+theme(axis.line.y=element_blank(),
                 axis.text.y=element_blank(),
                 axis.title.x=element_blank(),
                 axis.title.y=element_blank(),
                 axis.ticks.y=element_blank(),
                 axis.text.x =element_blank(),
                 axis.ticks.x =element_blank(),
                 axis.line.x =element_blank(),
                 legend.position = "bottom"
                )

# Show text for each month
timeline_plot<-timeline_plot+geom_text(data=month_df, aes(x=month_date_range,y=-0.1,label=month_format),size=2.5,vjust=0.5, color='black', angle=90)
# Show year text
timeline_plot<-timeline_plot+geom_text(data=year_df, aes(x=year_date_range,y=-0.2,label=year_format, fontface="bold"),size=2.5, color='black')
# Show text for each milestone
timeline_plot<-timeline_plot+geom_text(aes(y=text_position,label=milestone),size=2.5)
print(timeline_plot)


library(randomForest)
library(caret)
library(doMC)
library(mmadsenr)
library(futile.logger)
library(dplyr)
library(ggthemes)

# Train and tune random forest classifiers for each of the three data sets coming out of the experiment
# "equifinality-3", for binary analysis.
#
# Assumes that data-preparation.r has previously loaded CSV files and created binary
# data files, stored in a local data directory.
#
# Reduces the four class data into a two class problem for basic classifier analysis, and it splits off
# a test data set, placing it in the environment.
#
# NOTE:  This analysis takes a LONG time to run, so it is kept separate from the Rmarkdown
# analysis script.  At the completion of the analysis, it saves an environment image, which
# other scripts, such as RMarkdown documents, can load, and access the fitted model popsampled_results.
#
# Only the model training and fitting is done in this script.  The test data set is not analyzed
# here, since it can be done quickly enough that I want to be working in RMarkdown to examine different
# options.

# Set up logging
log_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3", filename = "popsampled-classification.log")
flog.appender(appender.file(log_file), name='cl')

clargs <- commandArgs(trailingOnly = TRUE)
if(length(clargs) == 0) {
  pop_data_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3", filename = "equifinality-3-population-data.rda")
  sampled_data_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3", filename = "equifinality-3-sampled-data.rda")
} else {
  pop_data_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3", filename = "equifinality-3-population-data.rda", args = clargs)
  sampled_data_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3", filename = "equifinality-3-sampled-data.rda", args = clargs)
}

load(pop_data_file)
load(sampled_data_file)

flog.info("Loaded data file: %s", pop_data_file, name='cl')
flog.info("Loaded data file: %s", sampled_data_file, name='cl')

flog.info("Beginning classification analysis of equifinality-3 data sets with full 4 classes", name='cl')

# set up parallel processing - use all the cores (unless it's a dev laptop under OS X) - from mmadsenr
num_cores <- get_parallel_cores_given_os(dev=TRUE)
flog.info("Number of cores used in analysis: %s", num_cores, name='cl')
registerDoMC(cores = num_cores)

########### Training and Tuning Variables ##############

#
# Common training and tuning parameters for ctmixtures analysis
#

gbm_grid <- expand.grid(.interaction.depth = (1:6)*2,
                        .n.trees = (1:10)*25,
                        .shrinkage = 0.05)

training_control <- trainControl(method="repeatedcv",
                                 number=10, repeats=5)

# make this repeatable - comment this out or change it to get a fresh analysis result
seed_value <- 58132133
set.seed(seed_value)
flog.info("RNG seed to replicate this analysis: %s", seed_value, name='cl')

# Set up sampling of train and test data sets
training_set_fraction <- 0.8
test_set_fraction <- 1.0 - training_set_fraction

# set up popsampled_results data frames
experiment_names <- c("Four Class Census", "Four Class Sample Size 10%", "Four Class Sample Size 20%")
fourclass_popsampled_results <- data.frame()

fourclass_popsampled_results_roc <- NULL
fourclass_popsampled_results_model <- NULL
fourclass_popsampled_results_cm <- NULL

###### Population Data ######

flog.info("Starting analysis of Four Class population census data", name='cl')

# first row of fourclass_popsampled_results
i <- 1
exp_name <- experiment_names[i]

# remove fields from analysis that aren't predictors, and the detailed label with 4 classes
exclude_columns <- c("simulation_run_id", "innovation_rate")

#######

#model <- train_randomforest(df, training_set_fraction, fit_grid, fit_control, exclude_columns)
model <- train_gbm_classifier(eq3_pop_df, training_set_fraction, "model_class_label", gbm_grid, training_control, exclude_columns, verbose=FALSE)

fourclass_popsampled_results_model[["population_census"]] <- model$tunedmodel

# use the test data split by the train_randomforest function and calculate tuned model predictions
# and then get the confusion matrix and fitting metrics
predictions <- predict(model$tunedmodel, newdata=model$test_data)
cm <- confusionMatrix(predictions, model$test_data$model_class_label)
results <- get_parsed_binary_confusion_matrix_stats(cm)
results$experiments <- exp_name
results$elapsed <- model$elapsed
fourclass_popsampled_results_cm[["population_census"]] <- cm

# calculate a ROC curve
population_roc <- calculate_roc_binary_classifier(model$tunedmodel, model$test_data, "model_class_label", exp_name)
results$auc <- unlist(population_roc$auc@y.values)
fourclass_popsampled_results_roc[["population_census"]]  <- population_roc

fourclass_popsampled_results <- rbind(fourclass_popsampled_results, results)

########################### sampled data ##########################

flog.info("Starting analysis of sampled data set 10%", name='cl')

## sample size 10 ##
# second row of fourclass_popsampled_results
i <- 2
exp_name <- experiment_names[i]

eq3_sampled_10 <- dplyr::filter(eq3_sampled_df, sample_size == 10)

# remove fields from analysis that aren't predictors, and the detailed label with 4 classes
exclude_columns <- c("simulation_run_id", "innovation_rate", "sample_size")

#######

#model <- train_randomforest(df, training_set_fraction, fit_grid, fit_control, exclude_columns)
model <- train_gbm_classifier(eq3_sampled_10, training_set_fraction, "model_class_label", gbm_grid, training_control, exclude_columns, verbose=FALSE)
fourclass_popsampled_results_model[["sampled_10"]] <- model$tunedmodel

# use the test data split by the train_randomforest function and calculate tuned model predictions
# and then get the confusion matrix and fitting metrics
predictions <- predict(model$tunedmodel, newdata=model$test_data)
cm <- confusionMatrix(predictions, model$test_data$model_class_label)
results <- get_parsed_binary_confusion_matrix_stats(cm)
results$experiments <- exp_name
results$elapsed <- model$elapsed
fourclass_popsampled_results_cm[["sampled_10"]] <- cm

# calculate a ROC curve
population_roc <- calculate_roc_binary_classifier(model$tunedmodel, model$test_data, "model_class_label", exp_name)
results$auc <- unlist(population_roc$auc@y.values)
fourclass_popsampled_results_roc[["sampled_10"]]  <- population_roc

fourclass_popsampled_results <- rbind(fourclass_popsampled_results, results)

## sample size 20 ##
# second row of fourclass_popsampled_results
i <- 3
exp_name <- experiment_names[i]

eq3_sampled_20 <- dplyr::filter(eq3_sampled_df, sample_size == 20)

# remove fields from analysis that aren't predictors, and the detailed label with 4 classes
exclude_columns <- c("simulation_run_id", "innovation_rate", "sample_size")

#######

#model <- train_randomforest(df, training_set_fraction, fit_grid, fit_control, exclude_columns)
model <- train_gbm_classifier(eq3_sampled_20, training_set_fraction, "model_class_label", gbm_grid, training_control, exclude_columns, verbose=FALSE)

fourclass_popsampled_results_model[["sampled_20"]] <- model$tunedmodel

# use the test data split by the train_randomforest function and calculate tuned model predictions
# and then get the confusion matrix and fitting metrics
predictions <- predict(model$tunedmodel, newdata=model$test_data)
cm <- confusionMatrix(predictions, model$test_data$model_class_label)
results <- get_parsed_binary_confusion_matrix_stats(cm)
results$experiments <- exp_name
results$elapsed <- model$elapsed
fourclass_popsampled_results_cm[["sampled_20"]] <- cm

# calculate a ROC curve
population_roc <- calculate_roc_binary_classifier(model$tunedmodel, model$test_data, "model_class_label", exp_name)
results$auc <- unlist(population_roc$auc@y.values)
fourclass_popsampled_results_roc[["sampled_20"]]  <- population_roc

fourclass_popsampled_results <- rbind(fourclass_popsampled_results, results)

############## Complete Processing and Save fourclass_popsampled_results ##########3

# we can now use plot_multiple_roc() to plot all the ROC curves on the same plot, etc.
# as well as graph various of the metrics as they vary across sample size and TA duratio
#plot_multiple_roc_from_list(fourclass_popsampled_results_roc)

if(length(clargs) == 0) {

  image_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3",
                              filename = "classification-pop-sampled_results-gbm.RData")
  image_file_results <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3",
                                      filename = "classification-pop-sampled_results-gbm-dfOnly.RData")

} else {

  image_file <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3",
                              filename = "classification-pop-sampled_results-gbm.RData", args = clargs)
  image_file_results <- get_data_path(suffix = "experiment-ctmixtures/equifinality-3",
                                      filename = "classification-pop-sampled_results-gbm-dfOnly.RData", args = clargs)
}

flog.info("Saving fourclass_popsampled_results of analysis to R environment snapshot: %s", image_file, name='cl')
save(fourclass_popsampled_results, fourclass_popsampled_results_model, fourclass_popsampled_results_roc, fourclass_popsampled_results_cm, file=image_file)

flog.info("Saving just data frame of fourclass_popsampled_results of analysis to R environment snapshot: %s", image_file_fourclass_popsampled_results, name='cl')
save(fourclass_popsampled_results, file=image_file_popsampled_results)

# End
flog.info("Analysis complete", name='cl')


library(EpiNow2)
library(data.table)
library(optparse)

option_list <- list(
            make_option(
                    "--first_day_modelled",
                    action = "store",
                    default = NA,
                    type = "character",
                    help = "First date of date to use"
                ),
            make_option(
                    "--weeks_modelled",
                    action = "store",
                    default = NA,
                    type = "integer",
                    help = "Number of weeks to model"
                ),
            make_option(
                    "--forecast_horizon",
                    action = "store",
                    default = NA,
                    type = "integer",
                    help = "Number of days for which to forecast"
                ),
            make_option(
                    "--area",
                    action = "store",
                    default = NA,
                    type = "character",
                    help = "Name of the area to model, case sensitive!"
                ),
            make_option(
                    "--case_counts",
                    action = "store",
                    default = NA,
                    type = "character",
                    help = "Path to preprocessed input data"
                ),
            make_option(
                    "--ncores",
                    action = "store",
                    default = NA,
                    type = "integer",
                    help = "Number of cores to tell stan to use"
                ),
            make_option(
                    "--output_folder",
                    action = "store",
                    default = NA,
                    type = "character",
                    help = "Where to store raw results"
                ),
            make_option(
                    "--region_codes",
                    action = "store",
                    default = NA,
                    type = "character",
                    help = "Path to JSON file with region codes. Format area_name -> code."
                ),
            make_option(
                    "--verbose",
                    action = "store_true",
                    default = FALSE,
                    type = "logical",
                    help = "Whether stan should be verbose"
                ),
            make_option(
                    "--area_index",
                    action="store",
                    type="integer",
                    default=0,
                    help="Area index; used if area not provided."
                )
        )

opt <- parse_args(OptionParser(option_list = option_list))

start_date <- as.IDate(opt$first_day_modelled)
df <- fread(opt$case_counts)
region_codes <- rjson::fromJSON(file = opt$region_codes)
if(is.na(opt$area)){
    if (opt$area_index==0) {
        stop("Area index 0.")
    }
    opt$area = names(region_codes)[opt$area_index]
}
area_code <- region_codes[[opt$area]]
end_date <- start_date + opt$weeks_modelled * 7
region_data <- subset(df[df$region == opt$area], select = c("date", "confirm"))
input_data <- region_data[region_data$date >= start_date & region_data$date < end_date]
input_data$date <- as.Date(input_data$date) # for compat with epinow2
print(tail(input_data, 1))

# as close as we can reasonably get to the paper
reporting_delay <- list(mean=log(6.5^2 / sqrt(17^2 + 6.5^2)),
                        mean_sd=0.05,
                        sd=log(17^2 / 6.5^2 + 1),
                        sd_sd=0.05,
                        max=30)

# generation interval, incubation time and rt prior are as in the paper
generation_time <- get_generation_time(disease = "SARS-CoV-2", source = "ganyani")
incubation_period <- get_incubation_period(disease = "SARS-CoV-2", source = "lauer")
rt_options <- rt_opts(prior=list(mean=1, sd=1))

stan_options <- stan_opts(cores = opt$ncores)

# the setting of zero_threshold is to stop backfilling
estimates <- epinow(reported_cases = input_data,
                    generation_time = generation_time,
                    delays = delay_opts(incubation_period, reporting_delay),
                    rt = rt_options,
                    stan = stan_options,
                    horizon = opt$forecast_horizon,
                    verbose = opt$verbose,
                    zero_threshold = Inf
                    )

standata <- rstan::extract(estimates$estimates$fit)
fitted_r_samples <- standata$R
fitted_case_samples <- standata$imputed_reports

write.table(
            fitted_r_samples,
            file = file.path(opt$output_folder, sprintf("%d_r_samples.txt", area_code)),
            row.names = FALSE,
            col.names = FALSE
        )

write.table(
            fitted_case_samples,
            file = file.path(opt$output_folder, sprintf("%d_case_samples.txt", area_code)),
            row.names = FALSE,
            col.names = FALSE
        )


./circle2.verify: test performed on 2001-03-05(basunus i686 running Linux)

#### Test: ./circle2.verify running circle.py 1 0.21
-1.8 1.8 -1.8 1.8
1.0  1.3195e+00
-2.7802e-01  1.6476e+00
-9.1367e-01  4.9135e-01
 4.8177e-02 -4.1189e-01
 1.1622e+00  2.9512e-01
end
CPU time of circle.py: 0.1 seconds on basunus i686, Linux


#  James Rekow

identifyNumSubgroupsTestSuite = function(iStrength = 1, sigmaMax = 0.1, thresholdMult = 10 ^ (-1),
                                         maxSteps = 10 ^ 4, tStep = 10 ^ (-2), intTime = 1,
                                         interSmplMult = 0.01, conGraph = NULL, rollRange = 0.025,
                                         numReplicates = 3){

  #  ARGS:
  #
  #  RETURNS: percentErrorDF - data frame containing columns for all the input data and the average
  #                            percent of samples placed in the wrong subgroup in each treatment

  source("identifyNumSubgroupsTest.r")

  #  source("identifySubgroupsVer1Test_Exp_Weight.r")
  #  numClust = KNOWN
  #
  #  inputs:
  #    1: M = 40, N = 20, lambda = 0, univ = 1, numSubgroups = 2   : standard case
  #    2: M = 100, N = 20, lambda = 0, univ = 1, numSubgroups = 2  : large M
  #    3: M = 40, N = 20, lambda = 0, univ = 1, numSubgroups = 6   : lot's of subgroups
  #    4: M = 100, N = 20, lambda = 100, univ = 1, numSubgroups = 6: large M, many subgrps, interactions
  #    5: M = 40, N = 20, lambda = 0, univ = 0, numSubgroups = 2   : no universality

  #  default setting is that number of clusters in each treatment is unknown
  numClust = NULL

  # input = vector("list", 5)
  # input[[1]] = c(40, 20, 2, 0, 1)
  # input[[2]] = c(100, 20, 2, 0, 1)
  # input[[3]] = c(40, 20, 6, 0, 1)
  # input[[4]] = c(100, 20, 6, 100, 1)
  # input[[5]] = c(40, 20, 2, 0, 0)

  input = vector("list", 5)
  input[[1]] = c(40, 20, 2, 0, 1)
  input[[2]] = c(40, 20, 4, 100, 1)
  input[[3]] = c(40, 20, 4, 200, 1)
  input[[4]] = c(100, 20, 4, 100, 1)
  input[[5]] = c(100, 20, 4, 200, 0)

  #  create a matrix of all unique combinations of elements from Mvec, NVec, lambdaVec, and univVec
  inputMat = Reduce(rbind, input)

  performIdentifyNumSubgroupsTest = function(inputVec){

    #  ARGS: inputVec - vector of input values from inputMat which are fed to the algorithm to be tested
    #
    #  RETURNS: meanPercentError - the mean percent error of the algorithm with the given inputs

    #  extract input values
    MVec = inputVec[1]
    NVec = inputVec[2]
    numSubgroupsVec = inputVec[3]
    lambdaVec = inputVec[4]
    univVec = inputVec[5]

    meanPercentError = identifyNumSubgroupsTest(MVec = MVec, NVec = NVec, lambdaVec = lambdaVec,
                                                numSubgroupsVec = numSubgroupsVec, iStrength = iStrength,
                                                univVec = univVec, sigmaMax = sigmaMax,
                                                thresholdMult = thresholdMult, maxSteps = maxSteps,
                                                tStep = tStep, intTime = intTime,
                                                interSmplMult = interSmplMult, conGraph = conGraph,
                                                rollRange = rollRange, numReplicates = numReplicates)[[6]]

    return(meanPercentError)

  } #  end performIdentifyNumSubgroupsTest function

  meanPercentErrorVec = apply(inputMat, 1, performIdentifyNumSubgroupsTest)

  meanPercentErrorDF = data.frame(M = inputMat[ , 1], N = inputMat[ , 2], numSubgroups = inputMat[ , 3],
                                  lambda = inputMat[ , 4], univVec = inputMat[ , 5],
                                  meanPercentError = meanPercentErrorVec)

  return(meanPercentErrorDF)

} #  end identifyNumSubgroupsTestSuite function


#scenario<-'T1-T2_stoc_rec_assess_run01'
#scenario<-'T1-T2_stoc_rec_assess_run01_0pctConstraints'
scenario.dir<-file.path(data.path,scenario)
scenario.dir<-data.path

my.dev<-'png'   # output device:  'screen', 'wmf', 'png', 'pdf'
nox<-2; noy<-2;

if ("scenario" %in% ls()) run.ID<-scenario else run.ID<-"ass_uncer"         # file id used for paper output
cleanup()

species.in.title<-T
first.year.on.plot<-2000
#first.year.on.plot<-2008
last.year.on.plot<-2050

#burn.in.delete<-5
include.assess.forcast.line<-T
include.F.reference.points<-T
include.SSB.reference.points<-TRUE

#percentile<-c(0.50,0.25,0.75,0.025,0.975)   #first value must be 50 and last value the highest
#percentile<-c(0.50,0.25,0.75)   #first value must be 50 and last value the highest
percentile<-c(0.50,0.25,0.75,0.05,0.95)   #first value must be 50 and last value the highest
#percentile<-c(0.50)   #first value must be 50 and last value the highest

incl.sp<-c(1,2,3)  # species number to be included
incl.sp<-"all"

my.lwd.median<-2.5
my.lwd.uncertainty<-2

cleanup()
plotfile<-function(dev='screen',out) {
  if (dev=='screen') X11(width=8, height=8, pointsize=12)
  if (dev=='wmf') win.metafile(filename = file.path(scenario.dir,paste(out,'.wmf',sep='')), width=8, height=10, pointsize=12)
  if (dev=='png') png(filename =file.path(scenario.dir,paste(out,'.png',sep='')), width = 1200, height = 1600,units = "px", pointsize = 30, bg = "white")
  if (dev=='pdf') pdf(file =file.path(scenario.dir,paste(out,'.pdf',sep='')), width = 8, height = 10,pointsize = 12,onefile=FALSE)
}

do.plot.forecast1<-function(sp.plot=c(1),incl.prob=FALSE,incl.closure=FALSE,incl.closure2=FALSE,
                             incl.constraints=TRUE,incl.0.group.M2=FALSE,incl.eaten=T,
                            plot.yield=T,myDir=data.path,minus.year=0,incl.percieved.F=F) {
#########################################
    control<-read.FLSMS.control(file=file.path(myDir,'SMS.dat'))
    sp.names<<-control@species.names
    HCR<-read.FLSMS.predict.control(control=control,file=file.path(myDir,'HCR_options.dat'))
    years<-c(control@first.year,control@last.year.model)
    ssb.r<-Read.MCMC.SSB.rec.data(dir=myDir)
    ssb.r<-subset(ssb.r,(Year>=first.year.on.plot & Year<=last.year.on.plot) ,drop=T)

    F.yield<-Read.MCMC.F.yield.data(dir=myDir)
    print('F.yield is done')
    F.yield<-subset(F.yield,Year>=first.year.on.plot & Year<=last.year.on.plot-minus.year,drop=T)
    if (HCR@age.group.output[1,"M2"] >0) {
      eaten.M2<-Read.MCMC.eaten_M2(dir=myDir)
      print('Read.MCMC.eaten_M2 is done')
      eaten.M2<-subset(eaten.M2,(Year>=first.year.on.plot & Year<=last.year.on.plot) ,drop=T)
      ms.incl<-sum(eaten.M2$eaten.M2)>1
      #ms.incl<-F
    } else ms.incl<-F;
    if (include.F.reference.points | include.SSB.reference.points) ref.points<-Read.reference.points()

    if (incl.closure) {
      closure <-Read.MCMC.closure.data(dir=myDir)
      print('Read.MCMC.closure.data is done')
      closure<-subset(closure,Year>=first.year.on.plot & Year<=last.year.on.plot,drop=T)
    }

    if (incl.closure2 | incl.constraints) {
      closure.constraint <-Read.MCMC.closure.constraints.data(dir=myDir)
      closure.constraint<-subset(closure.constraint,Year>=first.year.on.plot & Year<=last.year.on.plot,drop=T)
    }

    if (ms.incl) {
      M2<-Read.MCMC.detailed.M2(dir=myDir)
      print('Read.MCMC.detailed.M2 is done')
      M2<-subset(M2,(Year>=first.year.on.plot & Year<=last.year.on.plot) ,drop=T)
    }

    len.per<-length(percentile)

    if (is.character(incl.sp)) if (incl.sp=="all") {
       info<-control@species.info
       n.pred<-0; n.oth<-0
       for (s in (1:control@no.species)) {
         if(info[s,'predator']>=1) n.pred<-n.pred+1;
         if(info[s,'predator']==2) n.oth<-n.oth+1;
       }
       first.VPA<-n.oth+1
       n.VPA<-control@no.species-first.VPA+1

       sp.plot<-seq(first.VPA,length(unique(ssb.r$Species))+n.oth)
    }

for (sp in (sp.plot)) {

  plotfile(dev=my.dev,out=file.path(paste(run.ID,'_',sp.names[sp],sep='')));
  par(mar=c(3,5,3,2))
  par(mfcol=c(nox,noy))
  loc.name<-sp.names[sp]

  if (!species.in.title) loc.name[]<-""

  a<-subset(ssb.r,Species.n==sp)
  b<-tapply(a$SSB,list(a$Year), function(x) quantile(x,probs = percentile))
  res.ssb<-b

  y<-as.numeric(names(b))
  q<-matrix(unlist(b),nrow=len.per)/1000
  plot(y,q[1,],main=paste(loc.name, "SSB"),xlab="",ylab='SSB (1000t)',
        type='l',lwd=my.lwd.median,ylim=c(0,max(q[len.per,])))
  for (i in (1:len.per)) lines(y,q[i,],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)

  if (include.assess.forcast.line) abline(v=years[2])
  if (include.SSB.reference.points) {
    if (ref.points[sp,3] >0) abline(h=ref.points[sp,3]/1000,lty=2)
    if (ref.points[sp,4] >0) abline(h=ref.points[sp,4]/1000,lty=3)
  }

  if (incl.prob & ref.points[sp,3] >0 & ref.points[sp,4] >0) {
    a<-subset(ssb.r,Species.n==sp)
    q<-tapply(a$SSB,list(a$Year,a$Repetion,a$Iteration),sum)
    #all.res<-cbind(all.res,q)
    T1<-ref.points[sp,3]
    T2<-ref.points[sp,4]

    q[q>T2]<-0
    q[q>0]<-1
    q<-apply(q,c(1),sum)/(dim(q)[2]*dim(q)[3])

    y<-as.numeric(names(q))
    plot(y,q,main=paste("Prob., SSB below", T2/1000,"&", T1/1000),xlab="",ylab='probability',
            type='l',lwd=my.lwd.median,ylim=c(0,0.2))
    q<-tapply(a$SSB,list(a$Year,a$Repetion,a$Iteration),sum)
    q[q>T1]<-0
    q[q>0]<-1
    q<-apply(q,c(1),sum)/(dim(q)[2]*dim(q)[3])

    y<-as.numeric(names(q))
    lines(y,q,lty=2)
    abline(h=0.05,lty=3)
    if (include.assess.forcast.line) abline(v=years[2])
  }

  a<-subset(F.yield,Species.n==sp)
  b<-tapply(a$mean.F,list(a$Year), function(x) quantile(x,probs =percentile,na.rm=TRUE))
  res.F<-b

  y<-as.numeric(names(b))
  q<-matrix(unlist(b),nrow=len.per)
  plot(y,q[1,],main=paste("Mean F"),xlab="",ylab='',
        type='l',lwd=my.lwd.median,ylim=c(0,max(q[len.per,])))
  for (i in (1:len.per)) lines(y,q[i,],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)
  if (include.assess.forcast.line) abline(v=years[2])
  if (include.F.reference.points) {
    if (ref.points[sp,1] >0) abline(h=ref.points[sp,1],lty=2)
    if (ref.points[sp,2] >0) abline(h=ref.points[sp,2],lty=3)
  }
  if (incl.percieved.F) {
    b<-tapply(a$mean_F_percieved,list(a$Year), function(x) quantile(x,probs =0.50,na.rm=TRUE))
    lines(y,b,lty=1,col=1,lwd=3)
  }

  #a<-subset(F.yield,Species.n==sp)
  b<-tapply(a$Yield,list(a$Year), function(x) quantile(x,probs =percentile,na.rm=TRUE))
  res.yield<-b
  y<-as.numeric(names(b))
  q<-matrix(unlist(b),nrow=len.per)/1000
  if (plot.yield) {
  plot(y,q[1,],main=paste("Yield"),xlab="",ylab='Yield (1000 t)',
        type='l',lwd=my.lwd.median,ylim=c(0,max(q[len.per,])))
  for (i in (1:len.per)) lines(y,q[i,],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)
  if (include.assess.forcast.line) abline(v=years[2])
  #grid()
  }

  a<-subset(ssb.r,Species.n==sp & Year>=first.year.on.plot & Year<=last.year.on.plot-minus.year)
  b<-tapply(a$rec,list(a$Year), function(x) quantile(x,probs = percentile,na.rm=TRUE))
  res.recruit<-b
  y<-as.numeric(names(b))
  q<-matrix(unlist(b),nrow=len.per)/1000
  plot(y,q[1,],main=paste("Recruits"),xlab="",ylab='Recruits (millions)',
        type='l',lwd=my.lwd.median,ylim=c(0,max(q[len.per,])))
  for (i in (1:len.per)) lines(y,q[i,],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)
  if (include.assess.forcast.line) abline(v=years[2])

 if (ms.incl) {
  a<-subset(eaten.M2,Species.n==sp)
  b<-tapply(a$eaten.M2,list(a$Year), function(x) quantile(x,probs =percentile,na.rm=TRUE))
  y<-as.numeric(names(b))

  q<-matrix(unlist(b),nrow=len.per)/1000
  if (sum(q[1,])>=0.1 & incl.eaten) {
    plot(y,q[1,],main=paste("Eaten biomass"),xlab="",ylab='Eaten biomass(1000 t)',
            type='l',lwd=my.lwd.median,ylim=c(0,max(q[len.per,])))
    for (i in (1:len.per)) lines(y,q[i,],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)
    if (include.assess.forcast.line) abline(v=years[2])
  }

  a<-subset(M2,Species.n==sp)

  if (incl.0.group.M2) {
    a1<-subset(a,Age<=1)
    b<-tapply(a1$M2,list(a1$Year,a1$Age,a1$Repetion,a1$Iteration),sum)
    b1<-apply(b,c(1,2), function(x) quantile(x,probs =percentile,na.rm=TRUE))

    if (sum(b1)>=0.1) {
        y<-as.numeric(dimnames(b1)[[2]])
        plot(y,b1[1,,1],main=paste("M2 at age 0"),xlab="",ylab='M2',
                type='l',lwd=my.lwd.median,col=1,ylim=c(0,max(b1)))
        for (i in (1:len.per)) lines(y,b1[i,,1],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)
        if (include.assess.forcast.line) abline(v=years[2])

        if (F) {
          plot(y,b1[1,,2],main=paste("M2 at age 1"),xlab="",ylab='M2',
            type='l',lwd=my.lwd.median,col=1,ylim=c(0,max(b1[,,2])))
          for (i in (1:len.per)) lines(y,b1[i,,2],lty=1+i%/%2,col=1+i%/%2,lwd=my.lwd.uncertainty)
          if (include.assess.forcast.line) abline(v=years[2])
        }

    }

  }

  if (T) {
    b<-tapply(a$M2,list(a$Year,a$Age,a$Repetion,a$Iteration),sum)
    b<-apply(b,c(1,2),median)

    if (sum(b)>=0.1) {
      y<-as.numeric(dimnames(b)[[1]])
      plot(y,b[,1],main=paste("median M2 at age"),xlab="",ylab='M2',
              type='l',lwd=my.lwd.uncertainty,ylim=c(0,max(b)))
      for (a in (2:(dim(b)[2]))) if(max(b[,a]>0.001)) lines(y,b[,a],lty=a,col=a,lwd=my.lwd.uncertainty)
      if (include.assess.forcast.line) abline(v=years[2])
    }
  }
 }  #ms.incl

  if (incl.closure) {
    a<-subset(closure,Species.n==sp)
    max.closure<-max(a$closure)
    if (max.closure<=1) {
       b<-tapply(a$closure,list(a$Year), mean)
       y<-as.numeric(names(b))
      plot(y,b,main=paste("Probability of closure"),
            xlab="",ylab='probability',
            type='s',lwd=my.lwd.median,ylim=c(0,max(b)),xlim=c(first.year.on.plot,max(y)))
      if (include.assess.forcast.line) abline(v=years[2])
    } else {
      a<-data.frame(Year=a$Year,closure=a$closure,closure1=a$closure%/%2,closure2=a$closure%%2)
      b<-tapply(a$closure1,list(a$Year), mean)
      y<-as.numeric(names(b))
      plot(y,b,main=paste("Probability of closure"),
            xlab="",ylab='probability',
            type='s',lwd=my.lwd.median,ylim=c(0,max(b)),xlim=c(first.year.on.plot,max(y)))
     b<-tapply(a$closure2,list(a$Year), mean)
     lines(y,b,type='s',lwd=1.5,col=2,lty=4)
     if (include.assess.forcast.line) abline(v=years[2])
    }
  }
   colors=c("green","yellow","red")
   if (incl.closure2) {
       a<-subset(closure.constraint,Species.n==sp & closure>0)
       lab<-data.frame(closure=c(1,2,3),
             label=c("a) F<Fmsy & SSB>Bpa","b) Fmsy<F<Fpa & SSB>Bpa","c) F>Fpa |SSB<Bpa"))

       a<-merge(a,lab)
       b<-tapply(a$closure,list(a$label,a$Year), length)
       b[is.na(b)]<-0
       #print(b)
       #barplot(b,main="TAC Decision Rule",ylab='frequency')
       c<-b/apply(b,2,sum)
               barplot(c,main="TAC Decision Rule",ylab='probability',col=colors)
               legend(,x="left", legend=rownames(b),fill=colors,bg="white")
    }

   colors=c(0,2,3)
   if (incl.constraints) {
      a<-subset(closure.constraint,Species.n==sp & closure>0)

       b<-tapply(a$constraints,list(a$constraint,a$Year), length)
       combi<-data.frame(constraints=unique(a$constraints))
       lab<-data.frame(constraints=c(1,3,5),label=c('no constraints','lower','upper'))
       lab<-merge(lab,combi)
       lab<-(lab$label)
       b[is.na(b)]<-0
       c<-b/apply(b,2,sum)
               barplot(c,main="TAC constraints",ylab='probability',col=colors)
               legend(,x="left", legend=lab,fill=colors,bg="white")

    }
  all.res.glo<<-list(SSB=res.ssb,FI=res.F,Recruit=res.recruit,Yield=res.yield)
  if (my.dev!='screen') cleanup()

} #sp

}

do.plot.forecast1(sp.plot=incl.sp,incl.prob=F,incl.closure=F,incl.closure2=F,incl.constraints=F, incl.eaten=T,
               incl.0.group.M2=F,plot.yield=T,minus.year=1,myDir=scenario.dir)


## Import

atop <- readr::read_csv("atop4_01m.csv")

## drop non-aggression pacts
atop <- dplyr::filter(atop, !(nonagg ==1 & consul == 0 &
    neutral == 0 & offense == 0 & defense == 0))

## Clean data
bc_ext <- dplyr::tibble(
    "atopid" = atop$atopid,
    "member" = atop$member,
    "yrent" = atop$yrent,
    "yrexit" = atop$yrexit,
    ## Military action conditional war environment
    "macowe" = with(atop, ifelse(
        defcoadv == 1 | defcoloc == 1 | defcocon == 1|
        defconum == 1| offcoadv == 1| offcoloc == 1 |
        offcocon == 1 | offconum == 1,
        1, 0
    )),
    ## Military action conditional on non-compliance
    "maconc" = with(atop, ifelse(
        defcodem == 1 | offcodem == 1,
        1, 0
    )),
    ## Military action conditional on non-provocation
    "maconp" = atop$defconpr,
    ## Non-military action conditional on war environment
    "nmacowe" = with(atop, ifelse(
        neucoadv == 1 | neucoloc == 1 | neucocon == 1 |
        neuconum == 1 | neucoatt == 1 | concoadv == 1 |
        concoloc == 1 | concocon == 1,
        1, 0
    )),
    ## Non-military action required only if requested
    "nmaroir" = atop$concoreq,
    ## Other conditions
    "condother" = atop$olimob,
    ## Renunciation allowed
    "renounceok" = with(atop, ifelse(
        renounce == 1,
        1, 0
    )),
    ## Renunciation prohibited
    "renounceproh" = with(atop, ifelse(
        renounce == 2,
        1, 0
    )),
    ## Renunciation conditional
    "renouncecond" = with(atop, ifelse(
        renounce == 3,
        1, 0
    )),
    ## Offense alliance
    "offenseatop" = atop$offense,
    ## Defense alliance
    "defenseatop" = atop$defense,
    ## Neutrality pact
    "neutatop" = atop$neutral,
    ## Consultation
    "consultatop" = atop$consul,
    ## Military contact
    "milcontact" = with(atop, ifelse(
        milcon == 2,
        1, 0
    )),
    ## Common defense pact
    "commondef" = with(atop, ifelse(
        milcon == 3,
        1, 0
    )),
    ## Integrated command
    "intcom" = atop$intcom,
    ## Military aid
    "milaid" = atop$milaid,
    ## Military base
    "base" = atop$base,
    ## Specific contribution
    "contrib" = atop$contrib,
    ## Requires the creation of any organization
    "organ1" = atop$organ1,
    ## Economic aid
    "ecaid" = with(atop, ifelse(
        ecaid > 0,
        1, 0
    )),
    ## Secret pact
    "secrecy" = atop$pubsecr
)

all_alliances <- dplyr::tibble(
    "atopid" = as.integer(),
    "macowe" = as.integer(),
    "maconc" = as.integer(),
    "maconp" = as.integer(),
    "nmacowe" = as.integer(),
    "nmaroir" = as.integer(),
    "condother" = as.integer(),
    "renounceok" = as.integer(),
    "renounceproh" = as.integer(),
    "renouncecond" = as.integer(),
    "offenseatop" = as.integer(),
    "defenseatop" = as.integer(),
    "neutatop" = as.integer(),
    "consultatop" = as.integer(),
    "milcontact" = as.integer(),
    "commondef" = as.integer(),
    "intcom" = as.integer(),
    "milaid" = as.integer(),
    "base" = as.integer(),
    "contrib" = as.integer(),
    "organ1" = as.integer(),
    "ecaid" = as.integer(),
    "secrecy" = as.integer()
)

## First, list of unique alliance
alliances <- unique(bc_ext$atopid)

for(i in 1:length(alliances)){
    ## subset data on all states in alliance
    temp <- dplyr::filter(bc_ext, atop$atopid == alliances[i])
    all_alliances[i, 1] <- temp$atopid[1]
    all_alliances[i, 2] <- sum(temp$macowe)/nrow(temp)
    all_alliances[i, 3] <- sum(temp$maconc)/nrow(temp)
    all_alliances[i, 4] <- sum(temp$maconp)/nrow(temp)
    all_alliances[i, 5] <- sum(temp$nmacowe)/nrow(temp)
    all_alliances[i, 6] <- sum(temp$nmaroir)/nrow(temp)
    all_alliances[i, 7] <- sum(temp$condother)/nrow(temp)
    all_alliances[i, 8] <- sum(temp$renounceok)/nrow(temp)
    all_alliances[i, 9] <- sum(temp$renounceproh)/nrow(temp)
    all_alliances[i, 10] <- sum(temp$renouncecond)/nrow(temp)
    all_alliances[i, 11] <- sum(temp$offenseatop)/nrow(temp)
    all_alliances[i, 12] <- sum(temp$defenseatop)/nrow(temp)
    all_alliances[i, 13] <- sum(temp$neutatop)/nrow(temp)
    all_alliances[i, 14] <- sum(temp$consultatop)/nrow(temp)
    all_alliances[i, 15] <- sum(temp$milcontact)/nrow(temp)
    all_alliances[i, 16] <- sum(temp$commondef)/nrow(temp)
    all_alliances[i, 17] <- sum(temp$intcom)/nrow(temp)
    all_alliances[i, 18] <- sum(temp$milaid)/nrow(temp)
    all_alliances[i, 19] <- sum(temp$base)/nrow(temp)
    all_alliances[i, 20] <- sum(temp$contrib)/nrow(temp)
    all_alliances[i, 21] <- sum(temp$organ1)/nrow(temp)
    all_alliances[i, 22] <- sum(temp$ecaid)/nrow(temp)
    all_alliances[i, 23] <- sum(temp$secrecy)/nrow(temp)
}

all_alliances <- tidyr::drop_na(all_alliances)

scope <- MCMCpack::MCMCfactanal(
    ~ macowe + maconc + maconp + nmacowe + nmaroir +
    condother + renounceok + renounceproh + renouncecond +
    offenseatop + defenseatop + neutatop + consultatop,
    data = all_alliances,
    factors = 1,
    lambda.constraints = list(compel = list(1, "+")),
    std.mean = TRUE,
    std.var = TRUE,
    verbose = 100000,
    mcmc = 1000000,
    burnin = 100000,
    thin = 1000,
    store.scores = TRUE)

factor.load3 <- scope[, 1:26]
apply(factor.load3, 2, mean)
rbd <- scope[ , -seq(1:26)]
scope.score <- apply(rbd, 2, mean)

depth <- MCMCpack::MCMCfactanal(
    ~ milcontact + commondef + intcom + milaid + base +
    contrib + organ1 + ecaid + secrecy,
    data = all_alliances,
    factors = 1,
    lambda.constraints = list(compel = list(1, "+")),
    std.mean = TRUE,
    std.var = TRUE,
    verbose = 100000,
    mcmc = 1000000,
    burnin = 100000,
    thin = 1000,
    store.scores = TRUE)

factor.load2 <- depth[, 1:18]
apply(factor.load2, 2, mean)
rbd <- depth[,-seq(1:18)]
depth.score <- apply(rbd, 2, mean)

AllianceDataScore <- as.data.frame(cbind(all_alliances$atopid,
    depth.score, scope.score))

AllianceDataScore <- round(AllianceDataScore, digits = 3)

write.csv(AllianceDataScore, file = "AllianceDataScore_Kent.csv")


#' @include geom-.r
#' @include geom-raster.r
NULL

#' Annotation: high-performance rectangular tiling
#'
#' This is a special version of [geom_raster()] optimised for static
#' annotations that are the same in every panel. These annotations will not
#' affect scales (i.e. the x and y axes will not grow to cover the range
#' of the raster, and the raster must already have its own colours). This
#' is useful for adding bitmap images.
#'
#' @param raster raster object to display, may be an `array` or a `nativeRaster`
#' @param xmin,xmax x location (in data coordinates) giving horizontal
#'   location of raster
#' @param ymin,ymax y location (in data coordinates) giving vertical
#'   location of raster
#' @param interpolate If `TRUE` interpolate linearly, if `FALSE`
#'   (the default) don't interpolate.
#' @export
#' @examples
#' # Generate data
#' rainbow <- matrix(hcl(seq(0, 360, length.out = 50 * 50), 80, 70), nrow = 50)
#' ggplot(mtcars, aes(mpg, wt)) +
#'   geom_point() +
#'   annotation_raster(rainbow, 15, 20, 3, 4)
#' # To fill up whole plot
#' ggplot(mtcars, aes(mpg, wt)) +
#'   annotation_raster(rainbow, -Inf, Inf, -Inf, Inf) +
#'   geom_point()
#'
#' rainbow2 <- matrix(hcl(seq(0, 360, length.out = 10), 80, 70), nrow = 1)
#' ggplot(mtcars, aes(mpg, wt)) +
#'   annotation_raster(rainbow2, -Inf, Inf, -Inf, Inf) +
#'   geom_point()
#' rainbow2 <- matrix(hcl(seq(0, 360, length.out = 10), 80, 70), nrow = 1)
#' ggplot(mtcars, aes(mpg, wt)) +
#'   annotation_raster(rainbow2, -Inf, Inf, -Inf, Inf, interpolate = TRUE) +
#'   geom_point()
annotation_raster <- function(raster, xmin, xmax, ymin, ymax,
                              interpolate = FALSE) {
  if (!inherits(raster, 'nativeRaster'))
    raster <- grDevices::as.raster(raster)

  layer(
    data = dummy_data(),
    mapping = NULL,
    stat = StatIdentity,
    position = PositionIdentity,
    geom = GeomRasterAnn,
    inherit.aes = FALSE,
    params = list(
      raster = raster,
      xmin = xmin,
      xmax = xmax,
      ymin = ymin,
      ymax = ymax,
      interpolate = interpolate
    )
  )

}

#' @rdname ggplot2-ggproto
#' @format NULL
#' @usage NULL
#' @export
GeomRasterAnn <- ggproto("GeomRasterAnn", Geom,
  extra_params = "",
  handle_na = function(data, params) {
    data
  },

  draw_panel = function(data, panel_params, coord, raster, xmin, xmax,
                        ymin, ymax, interpolate = FALSE) {
    if (!inherits(coord, "CoordCartesian")) {
      cli::cli_abort("{.fn annotation_raster} only works with {.fn coord_cartesian}")
    }
    corners <- new_data_frame(list(x = c(xmin, xmax), y = c(ymin, ymax)), n = 2)
    data <- coord$transform(corners, panel_params)

    x_rng <- range(data$x, na.rm = TRUE)
    y_rng <- range(data$y, na.rm = TRUE)

    rasterGrob(raster, x_rng[1], y_rng[1],
      diff(x_rng), diff(y_rng), default.units = "native",
      just = c("left","bottom"), interpolate = interpolate)
  }
)


######################### SCRIPT 2: MODEL-BUILDING #########################

# This script runs the three main NSUM models
# 1. Basic model
# 2. Transmission error model
# 3. Barrier effect (+ transmission error) model

# dependencies: JAGS, R packages (coda, rjags)

# Load required packages
library(rjags)

# load the functions from the supporting_function.r script
source('codes/supporting_functions.r')

# read in data
knownpop = read.csv('data/knownpop.csv',as.is = TRUE)
nsum = read.csv('data/nsum_sg.csv',as.is = TRUE)
validation = read.csv('data/validation.csv',as.is = TRUE)

names_pop = c('Medical doctors','Primary school teachers','Full-time NS men',
              'Licensed taxi drivers','Hawker stalls owners',
              'Licensed property agents','Male clients of sex workers',
              'Women who had a baby in 2016','NUS graduates 2016',
              'Stroke 2016','Heart attack 2016','Bought an HDB in 2016',
              'Driving license 2016','PSLE 2016','NDP 2016','MSM',
              'O-Levels 2016','Couples who got married in 2016',
              'Dengue 2016','Female sex workers','Men >70 years old',
              'Women >70 years old','Single men >50 years old','IVDU')
hiddenpop = c(0,0,1,0,0,0,1,0,0,1,1,0,0,0,0,1,0,0,0,1,0,0,0,1)
popindex = grep(pattern = '_total',x = names(nsum))

data0 = nsum[,popindex]
te = nsum[,c('socialaccept_mcfsw','socialaccept_msm','socialaccept_fsw','socialaccept_pwid')]
demo = nsum[,c('age','sex','ethnicity_x')]
demo$age[is.na(demo$age)] = mean(demo$age,na.rm = TRUE)
demo$malay = 1*(demo$ethnicity_x %in% "Malay")
demo$indian = 1*(demo$ethnicity_x %in% "Indian")

validation$known_population[validation$rank<11]
indexk = validation$index[validation$rank<11]
indexu = c(7,16,20,24)

# The models were fit using a Markov chain Monte Carlo algorithm with 50 000 iterations
# with a burn-in of 5 000, storing 1 out of 10 iterations.
# Convergence was assessed visually with trace plots. The data analyses and visualisations
# were performed in R18 and the model building was done in JAGS.
# Deviance Information Criterion (DIC) was used to compare the models.

set.seed(666)
results = runModel(indexu = c(7,16,20,24),indexk = validation$index[validation$rank<11],NITERATION=50000)
set.seed(666)
results_te = runModel_te(indexu = c(7,16,20,24),indexk = validation$index[validation$rank<11],NITERATION=50000,x = te)
set.seed(666)
results_te_be = runModel_te_be(indexu = c(7,16,20,24),indexk = validation$index[validation$rank<11],NITERATION=500,x = te)

rs = results
rs = results_te
rs = results_te_be
G = geweke.diag(rs$posterior) # Note: you may want to run these diagnostics on the other results too: this just does it for the final model
H = heidel.diag(rs$posterior)
H
if(FALSE)
{
  # Note: these lines create plots to assess convergence
  #       However, if you are using RStudio and your plots panel is too small, the GUI may stop the code here
  #       So you should manually run this if you want trace plots, by switching the if statement to TRUE, or
  #       ...just by running these lines manually
  hist(G[[1]][[1]])

  plot(rs$posterior[[1]][,'lambda'])

  plot(rs$posterior[[1]][,c('beta[1]','beta[2]','beta[3]','beta[4]')])

  plot(rs$posterior[[1]][,c('Su[1]','Su[2]','Su[3]','Su[4]')])

  plot(rs$posterior[[1]][,c('bu1','bu2','bu3','bu4')])
}

posterior = list()

posterior$basic = as.matrix(results$posterior)
posterior$transmission = as.matrix(results_te$posterior)
posterior$barrier = as.matrix(results_te_be$posterior)

size_posterior = list()

size_posterior$basic = posterior$basic[,grep('Su',colnames(posterior$basic))]
size_posterior$transmission = posterior$transmission[,grep('Su',colnames(posterior$transmission))]
size_posterior$barrier = posterior$barrier[,grep('Su',colnames(posterior$barrier))]

alpha = list()

alpha$basic = posterior$basic[,grep('alpha',colnames(posterior$basic))]
alpha$transmission = posterior$transmission[,grep('alpha',colnames(posterior$transmission))]
alpha$barrier = posterior$barrier[,grep('alpha',colnames(posterior$barrier))]

lambda = list()

lambda$basic = posterior$basic[,grep('lambda',colnames(posterior$basic))]
lambda$transmission = posterior$transmission[,grep('lambda',colnames(posterior$transmission))]
lambda$barrier = posterior$barrier[,grep('lambda',colnames(posterior$barrier))]

size_estimates = list()

size_estimates$basic = (signif(apply(posterior$basic[,grep('Su',colnames(posterior$basic))],2,function(x)quantile(x,probs = c(0.5,0.025,0.975))),2))
size_estimates$transmission = (signif(apply(posterior$transmission[,grep('Su',colnames(posterior$transmission))],2,function(x)quantile(x,probs = c(0.5,0.025,0.975))),2))
size_estimates$barrier = (signif(apply(posterior$barrier[,grep('Su',colnames(posterior$barrier))],2,function(x)quantile(x,probs = c(0.5,0.025,0.975))),2))

colnames(size_estimates[[1]]) = c('MCFSW','MSM','FSW','IVDU')
colnames(size_estimates[[2]]) = c('MCFSW','MSM','FSW','IVDU')
colnames(size_estimates[[3]]) = c('MCFSW','MSM','FSW','IVDU')

sizeestimate_summary = rbind(cbind(model = 1,size_estimates[[1]]),
                             cbind(model = 2,size_estimates[[2]]),
                             cbind(model = 3,size_estimates[[3]]))

write.csv(x = sizeestimate_summary,file = 'results/sizeestimates.csv')

postmat_adjusted = as.matrix(results_te$posterior)
sizehiddenmat_adjusted = postmat_adjusted[,grep('Su',colnames(postmat_adjusted))]
lambdas_adjusted = postmat_adjusted[,grep('lambda',colnames(postmat_adjusted))]
alphas_adjusted = postmat_adjusted[,grep('alpha',colnames(postmat_adjusted))]
betas = postmat_adjusted[,grep('beta',colnames(postmat_adjusted))]
#transmission error
betasummary = t(signif(apply(betas,2,function(x)quantile(x,probs = c(0.5,0.025,0.975))),2))
rownames(betasummary) = c('MCFSW','MSM','FSW','IVDU')
write.csv(x = betasummary,file = 'results/transmissionerror.csv')
#size estimates
sizeestimates = t(signif(apply(sizehiddenmat_adjusted,2,function(x)quantile(x,probs = c(0.5,0.025,0.975))),2))
rownames(sizeestimates) = c('MCFSW','MSM','FSW','IVDU')
write.csv(x = sizeestimates,file = 'results/sizeestimates.csv')

# network size
mod = 3
alphas_mean = apply(alpha[[mod]],2,function(x)quantile(x,probs = c(0.5)))
alphas_cil = apply(alpha[[mod]],2,function(x)quantile(x,probs = c(0.025)))
alphas_ciu = apply(alpha[[mod]],2,function(x)quantile(x,probs = c(0.975)))
networksize_crude = mean(lambda[[mod]])*5E6
networksize_crude
networksize = mean(lambda[[mod]])*mean(alphas_mean)*5E6
networksize_cil = mean(lambda[[mod]])*mean(alphas_cil)*5E6
networksize_ciu = mean(lambda[[mod]])*mean(alphas_ciu)*5E6
network  = data.frame(networksize=networksize,networksize_cil=networksize_cil,networksize_ciu=networksize_ciu)
network
write.csv(round(network,1),'results/network.csv',row.names = FALSE)


# Obtain mean stem and crown ages (and 95-HPD) for orders, families, and subfamilies.
# Dated phylogeny produced in Beast + treeAnnotator.
# The MCC trees are stored as part of the Supplementary Data of Ramírez-Barahona et al. (2020)
# Uses additional files

tree <- "SB_MCC.v_1.tre" # file in data folder (the six MCCs trees are available as a ZIP file)
tre<-phyloch::read.beast(tree)
##### CORRECT SOME TIP NAMES, ALTHOUGH NOT NECESSARY AT THIS POINT
tre$tip.label[which(tre$tip.label=="Celastrales_Celastraceae_Euonymus_spp_Celastraceae")]<-"Celastrales_Celastraceae_Euonymus_spp_Celastraceae.sstr"
tre$tip.label[which(tre$tip.label=="Laurales_Lauraceae_Hypodaphniszenkeri_Hypodaphnideae")]<-"Laurales_Lauraceae_Hypodaphnis_zenkeri_Hypodaphnideae"
tre$tip.label[which(tre$tip.label=="Brassicales_Batidaceae_Batis_maritima")]<-"Brassicales_Bataceae_Batis_maritima"
tre$tip.label[which(tre$tip.label=="Caryophyllales_Didieraceae_Alluaudia_spp")]<-"Caryophyllales_Didiereaceae_Alluaudia_spp"
tre$tip.label[which(tre$tip.label=="Caryophyllales_Didieraceae_Portulacaria_afra")]<-"Caryophyllales_Didiereaceae_Portulacaria_afra"
tre$tip.label[which(tre$tip.label=="Huerteales_Gerrardiniaceae_Gerrardina_foliosa")]<-"Huerteales_Gerrardinaceae_Gerrardina_foliosa"
tre$tip.label[which(tre$tip.label=="Huerteales_Petenaceae_Petenaea_cordata")]<-"Huerteales_Petenaeaceae_Petenaea_cordata"
tre$tip.label[which(tre$tip.label=="Liliales_Petermaniaceae_Petermannia_cirrosa")]<-"Liliales_Petermanniaceae_Petermannia_cirrosa"
tre$tip.label[which(tre$tip.label=="Liliales_Rhipogonaceae_Rhipogonum_elseyanum")]<-"Liliales_Ripogonaceae_Ripogonum_elseyanum"
tre$tip.label[which(tre$tip.label=="Liliales_Rhipogonaceae_Ripogonum_scandens")]<-"Liliales_Ripogonaceae_Ripogonum_scandens"
tre$tip.label[which(tre$tip.label=="Liliales_Rhipogonaceae_Rhipogonum_elseyanum")]<-"Liliales_Ripogonaceae_Ripogonum_elseyanum"
tre$tip.label[which(tre$tip.label=="Zingiberales_Strelitiziaceae_Ravenala_madagascariensis")]<-"Zingiberales_Strelitziaceae_Ravenala_madagascariensis"
tre$tip.label[which(tre$tip.label=="Zingiberales_Strelitiziaceae_Strelitizia_spp")]<-"Zingiberales_Strelitziaceae_Strelitzia_spp"
tre$tip.label[which(tre$tip.label=="Asterales_Stylidaceae_Donatia_spp")]<-"Asterales_Stylidiaceae_Donatia_spp"
tre$tip.label[which(tre$tip.label=="Asterales_Stylidaceae_Stylidium_spp")]<-"Asterales_Stylidiaceae_Stylidium_spp"
tre <- ladderize(tre,TRUE); list<-tre$tip.label
list_ords <- foreach(i=1:length(list),.combine=c) %do% {strsplit(list,"_")[[i]][1]}
list_fams <- foreach(i=1:length(list),.combine=c) %do% {strsplit(list,"_")[[i]][2]}
list_subfams <- foreach(i=1:length(list),.combine=c) %do% {strsplit(list,"_")[[i]][5]}
list_ords_un<-sort(unique(list_ords))
list_fams_un<-sort(unique(list_fams))
list_subfams_un<-sort(unique(list_subfams))
list_un <- (c(list_ords_un,list_fams_un,list_subfams_un))
list_table<-cbind(list_ords,list_fams,list_subfams);list_table<-list_table[!duplicated(list_table),]
res<-matrix(nrow=length(list_fams_un)+length(list_ords_un)+length(list_subfams_un),ncol=10)
rownames(res)<-list_un; colnames(res)<-c("CG_age_Acal","SG_age_Acal","No_tips","CG_minHPD","CG_maxHPD","SG_minHPD","SG_maxHPD","Order","Family","Subfamily")
for(i in 1:length(list_un)){
  iden<-grep(list_un[i],x = tre$tip.label)
  res[i,3]<-length(iden)
  cat(i,"----",list_un[i],"\n")
  if(length(iden) == 1) {
    res[i,1]<-NA
    ssNODE <-tre$edge[which(tre$edge[,2]==iden),1]
    subtrea<-extract.clade(tre,ssNODE)
    res[i,2]<-max(branching.times(subtrea))
    stopifnot(round(tre$height_median[ssNODE-1215],digits = 2) == round(res[i,2],digits = 2))
    res[i,6]<-tre$`height_95%_HPD_MIN`[ssNODE-1215];res[i,7]<-tre$`height_95%_HPD_MAX`[ssNODE-1215]
  }
  else{
    subtre<-extract.clade(tre,getMRCA(tre,tre$tip.label[iden]))
    if(length(subtre$tip.label)-length(iden)!=0) {cat("-----------","non-monophyletic","\n"); next}
    res[i,3]<-length(iden)
    res[i,1]<-max(branching.times(subtre))
    NODE <- getMRCA(tre,tre$tip.label[iden])
    stopifnot(round(tre$height_median[NODE-1215],digits = 2)==round(res[i,1],digits = 2))
    res[i,4]<-tre$`height_95%_HPD_MIN`[NODE-1215];res[i,5]<-tre$`height_95%_HPD_MAX`[NODE-1215]
    sNODE <-tre$edge[which(tre$edge[,2]==NODE),1]
    subtre<-extract.clade(tre,sNODE)
    res[i,2]<-max(branching.times(subtre))
    stopifnot(round(tre$height_median[sNODE-1215],digits = 2)==round(res[i,2],digits = 2))
    res[i,6]<-tre$`height_95%_HPD_MIN`[sNODE-1215];res[i,7]<-tre$`height_95%_HPD_MAX`[sNODE-1215]
  }}
for (i in 1:length(list_un)){
  if(i <= 68){res[i,8] <- as.character(list_un[i])}
  if(i > 68 & i<= 510) {res[i,8]<-as.character(list_table[match(rownames(res),list_table[,2])[i],1]); res[i,9]<-as.character(list_table[match(rownames(res),list_table[,2])[i],2])}
  if(i > 510) {res[i,8]<-as.character(list_table[match(rownames(res),list_table[,3])[i],1]); res[i,9]<-as.character(list_table[match(rownames(res),list_table[,3])[i],2]); res[i,10]<-as.character(list_table[match(rownames(res),list_table[,3])[i],3])
  }
}
res<-as.data.frame(res)
write.tree(tre,file=paste(ruta_write,"0.MCC_parsed.tree",sep=""))
write.table(res,file=paste(ruta_write,"2.Ages_complete.csv",sep=""),sep=",",row.names=T,col.names=T,quote = F)


#------------------------------
# Execute with: Rscript convert.r
# Used software: R version 3.2.2 (2015-08-14)
# https://www.r-project.org/
#------------------------------

print("Preparing contestdata values...")

voltage_factor <- 1
current_factor <- .01

voltage_data <- read.table("voltage_allrep.txt", sep = "" , header = F , na.strings ="", stringsAsFactors= F)[1]
current_data <- read.table("current.txt", sep = "" , header = F , na.strings ="", stringsAsFactors= F)

voltage_data <- voltage_data * voltage_factor
current_data <- current_data * current_factor

svg("voltage_all.svg")
matplot(voltage_data, lty=1, type="l")
svg("current_all.svg")
matplot(current_data, lty=1, type="l")

writePlotData <- function(folder, min_row, length, eachtimestep) {
    print(paste("Process folder:", folder))
    dir.create(folder)

    max_row <- min_row + length

    voltage_data <- voltage_data[min_row:max_row, ]
    current_data <- current_data[min_row:max_row, ]

    voltage_data <- voltage_data[seq(1, length(voltage_data), eachtimestep)]
    current_data <- current_data[seq(1, length(current_data), eachtimestep)]

    voltagefile = paste(folder, "/voltage.csv", sep="")
    inputfile = paste(folder, "/input.csv", sep="")

    write.table(voltage_data, file = voltagefile, row.names=FALSE, col.names=FALSE, sep="\n")
    write.table(current_data, file = inputfile, row.names=FALSE, col.names=FALSE, sep="\n")

    svg(paste(folder, "/voltage.svg", sep=""))
    matplot(voltage_data, lty=1, type="l")

    svg(paste(folder, "/input.svg", sep=""))
    matplot(current_data, lty=1, type="l")
}

## Old data sets
#writePlotData("../input-short", 186000, 16400, 1)
#writePlotData("../input-long", 186000, 47900, 1)

writePlotData("../input-calibration-01", 40000, 80000, 1)
writePlotData("../input-calibration-05", 40000, 80000, 5)
writePlotData("../input-real-01", 300000, 90000, 1)
writePlotData("../input-real-05", 300000, 90000, 5)


#!/usr/local/bin/R
library(tidyr)
library(ggplot2)
library(dplyr)

t <- as.tbl(read.csv("foo.tsv", header=T,
                     sep='\t',
                     na.strings=c("None", "", "?")))
t %>% glimpse()

princomp(t[,3-9], na.action=na.omit, scores=TRUE)$scores