pdf_file<-"pdf/barcharts_multiple.pdf"
cairo_pdf(bg="grey98", pdf_file, width=13,height=10.5)

par(omi=c(0.65,0.75,1.25,0.75),mai=c(0.9,3.85,0.55,0),lheight=1.15,family="Lato Light",las=1)
source("scripts/inc_labels_za4753.r")
library(memisc)

# Read data and prepare chart

ZA4753<-spss.system.file("myData/ZA4753_v1-1-0.sav")
myData<-subset(ZA4753,select=c(v159,v160,v161,v162,v163,v164,v165))
attach(myData)
z<-NULL
y<-table(as.matrix(v165))
z<-c(z,100*(y["1"]+y["2"])/sum(y))
y<-table(as.matrix(v164))
z<-c(z,100*(y["1"]+y["2"])/sum(y))
y<-table(as.matrix(v163))
z<-c(z,100*(y["1"]+y["2"])/sum(y))
y<-table(as.matrix(v162))
z<-c(z,100*(y["1"]+y["2"])/sum(y))
y<-table(as.matrix(v161))
z<-c(z,100*(y["1"]+y["2"])/sum(y))
y<-table(as.matrix(v160))
z<-c(z,100*(y["1"]+y["2"])/sum(y))
y<-c(0,table(as.matrix(v159)))
z<-c(z,100*(y["1"]+y["2"])/sum(y))

# Create chart

bp<-barplot(z,names.arg=F,horiz=T,border=NA,xlim=c(0,100),
	col="grey",axes=F,family="Lato")
myColour<-rgb(255,0,210,maxColorValue=255)
rect(0,-0.1,20,8.6,col=rgb(191,239,255,80,maxColorValue=255),border=NA)
rect(20,-0.1,40,8.6,col=rgb(191,239,255,120,maxColorValue=255),border=NA)
rect(40,-0.1,60,8.6,col=rgb(191,239,255,80,maxColorValue=255),border=NA)
rect(60,-0.1,80,8.6,col=rgb(191,239,255,120,maxColorValue=255),border=NA)
rect(80,-0.1,100,8.6,col=rgb(191,239,255,80,maxColorValue=255),border=NA)
z2<-c(0,0,84.81928,0,0,0,0)
bp<-barplot(z2,names.arg=F,horiz=T,border=NA,xlim=c(0,100),
	col=myColour,axes=F,add=T)

# Other elements

for (i in 1:length(mynames))
{
if (i == 3) {myFont<-"Lato Bold"} else {myFont<-"Lato Light"}
text(-3,bp[i],mynames[i],xpd=T,adj=1,family=myFont,cex=1)
text(10,bp[i],format(round(z[i],1),nsmall=1),family=myFont,cex=1.25,
	col=ifelse(i==3,"white","black"))
}
arrows(50,-0.1,50,8.8,lwd=1.5,length=0,xpd=T,col="skyblue3")
arrows(50,-0.25,50,-0.1,lwd=5,length=0,xpd=T)
arrows(50,8.8,50,8.95,lwd=5,length=0,xpd=T)
text(48,8.9,"Majority",adj=1,xpd=T,cex=0.9,font=3)
text(52,8.9,"50%",adj=0,xpd=T,cex=0.9,family="Lato Bold",font=3)
text(100,8.9,"all values in percent",adj=1,xpd=T,cex=0.9,font=3)
mtext(c(0,20,40,60,80,100),at=c(0,20,40,60,80,100),1,line=0.75)

# Titling

mtext("It is often said that attitudes towards gender roles are changing",3,line=2.2,adj=0,cex=1.8,family="Lato Black",outer=T)
mtext("Agree strongly / agree ",3,line=0,adj=0,cex=1.5,outer=T)
mtext("Source: European Values Study 2008 Germany, ZA4753. www.gesis.org. Design: Stefan Fichtel, ixtract",1,line=0,adj=1,cex=0.95,outer=T,font=3)
dev.off()


#' topojson projections and extensions
#'
#' @export
#' @param proj Map projection name. One of albers, albersUsa, azimuthalEqualArea,
#' azimuthalEquidistant, conicEqualArea, conicConformal, conicEquidistant, equirectangular,
#' gnomonic, mercator, orthographic, stereographic, or transverseMercator.
#' @param rotate If rotation is specified, sets the projection's three-axis rotation to the
#' specified angles yaw, pitch and roll (or equivalently longitude, latitude and roll)
#' in degrees and returns the projection. If rotation is not specified, returns the current
#' rotation which defaults `[0, 0, 0]`. If the specified rotation has only two values, rather than
#' three, the roll is assumed to be 0.
#' @param center If center is specified, sets the projection's center to the specified location, a
#' two-element array of longitude and latitude in degrees and returns the projection. If center is
#' not specified, returns the current center which defaults to (0,0)
#' @param translate If point is specified, sets the projection's translation offset to the
#' specified two-element array `[x, y]` and returns the projection. If point is not specified,
#' returns the current translation offset which defaults to `[480, 250]`. The translation offset
#' determines the pixel coordinates of the projection's center. The default translation offset
#' places (0,0) at the center of a 960x500 area.
#' @param scale If scale is specified, sets the projection's scale factor to the specified value
#' and returns the projection. If scale is not specified, returns the current scale factor which
#' defaults to 150. The scale factor corresponds linearly to the distance between projected points.
#' However, scale factors are not consistent across projections.
#' @param clipAngle If angle is specified, sets the projection's clipping circle radius to the
#' specified angle in degrees and returns the projection. If angle is null, switches to
#' antimeridian cutting rather than small-circle clipping. If angle is not specified, returns the
#' current clip angle which defaults to null. Small-circle clipping is independent of viewport
#' clipping via clipExtent.
#' @param precision If precision is specified, sets the threshold for the projection's adaptive
#' resampling to the specified value in pixels and returns the projection. This value corresponds
#' to the Douglas-Peucker distance. If precision is not specified, returns the projection's current
#' resampling precision which defaults to Math.SQRT(1/2).
#' @param parallels Depends on the projection used! See
#' https://github.com/mbostock/d3/wiki/Geo-Projections#standard-projections for help
#' @param clipExtent If extent is specified, sets the projection's viewport clip extent to the
#' specified bounds in pixels and returns the projection. The extent bounds are specified as an
#' array `[[x0, y0], [x1, y1]]`, where x0 is the left-side of the viewport, y0 is the top, x1 is
#' the right and y1 is the bottom. If extent is null, no viewport clipping is performed. If extent
#' is not specified, returns the current viewport clip extent which defaults to null. Viewport
#' clipping is independent of small-circle clipping via clipAngle.
#' @param invert Projects backward from Cartesian coordinates (in pixels) to spherical coordinates
#' (in degrees). Returns an array `[longitude, latitude]` given the input array `[x, y]`.
#' @examples
#' projections(proj="albers")
#' projections(proj="albers", rotate='[98 + 00 / 60, -35 - 00 / 60]', scale=5700)
#' projections(proj="albers", scale=5700)
#' projections(proj="albers", translate='[55 * width / 100, 52 * height / 100]')
#' projections(proj="albers", clipAngle=90)
#' projections(proj="albers", precision=0.1)
#' projections(proj="albers", parallels='[30, 62]')
#' projections(proj="albers", clipExtent='[[105 - 87, 40], [105 + 87 + 1e-6, 82 + 1e-6]]')
#' projections(proj="albers", invert=60)
#' projections("orthographic")
projections <- function(proj, rotate=NULL, center=NULL, translate=NULL,
  scale=NULL, clipAngle=NULL, precision=NULL, parallels=NULL, clipExtent=NULL,
  invert=NULL) {

  if (missing(proj)) stop("You must provide a character string to 'proj'", call. = FALSE)
  vals <- list(
    albers = 'd3.geo.albers()%s',
    albersUsa = 'd3.geo.albersUsa()',
    azimuthalEqualArea = 'd3.geo.azimuthalEqualArea()',
    azimuthalEquidistant = 'd3.geo.azimuthalEquidistant()',
    conicEqualArea = 'd3.geo.conicEqualArea()',
    conicConformal = 'd3.geo.conicConformal()',
    conicEquidistant = 'd3.geo.conicEquidistant()',
    equirectangular = 'd3.geo.equirectangular()',
    gnomonic = 'd3.geo.gnomonic()',
    mercator = 'd3.geo.mercator()',
    orthographic = 'd3.geo.orthographic()',
    stereographic = 'd3.geo.stereographic()',
    transverseMercator = 'd3.geo.transverseMercator()'
  )
  got <- vals[[proj]]
  if (is.null(got)) stop("no match for 'proj' parameter input")
  args <- tg_compact(list(rotate = rotate, center = center,
    translate = translate, scale = scale,
    clipAngle = clipAngle, precision = precision, parallels = parallels,
    clipExtent = clipExtent, invert = invert))
  out <- list()
  for (i in seq_along(args)) {
    out[i] <- sprintf(".%s(%s)", names(args[i]), args[[i]])
  }
  argstogo <- paste(out, collapse = "")
  gotgo <- if (nzchar(argstogo)) sprintf(got, argstogo) else got
  gotgo <- gsub("%s", "", gotgo)
  if (is.null(gotgo)) {
    "That projection doesn't exist, check your spelling"
  } else {
    gotgo
  }
}


#####################################################################################################################
## WKRATIO: Portuguese data testing
## National discard raising procedure using RDBES data format (Hierarchy 3)
## Ana Cláudia Fernandes, 31may - 4jun
#####################################################################################################################

## No national CE table available - "Logbook effort data" prepared according to RDBES format
## Biological data: subset for blue whiting (WHB)

# Load our functions
	source("RDBES_Functions.R")

# IMPORTANT: Hack to stop write.csv changing numbers to scientific notation
	options(scipen=500) # big number of digits

## Loads data

### Effort data prepared	(CE table)
	load('.../Data/effTable.rdata')

### HIERARCHY 3 - Data subset for WHB
	load('.../Data/H3_WHB.rdata')

## data formats

	myExchangeFileH3$FM$FMnumberAtUnit <- as.numeric(as.character(myExchangeFileH3$FM$FMnumberAtUnit))
	myExchangeFileH3$SA$SAnumberSampled <- as.numeric(myExchangeFileH3$SA$SAnumberSampled)
	myExchangeFileH3$SA$SAnumberTotal <- as.numeric(myExchangeFileH3$SA$SAnumberTotal)
	myExchangeFileH3$SA$SAtotalWeightLive <- as.numeric(myExchangeFileH3$SA$SAtotalWeightLive)
	myExchangeFileH3$SA$SAsampleWeightLive <- as.numeric(myExchangeFileH3$SA$SAsampleWeightLive)
	myExchangeFileH3$SA$SAtotalWeightMeasured <- as.numeric(myExchangeFileH3$SA$SAtotalWeightMeasured)
	myExchangeFileH3$SA$SAsampleWeightMeasured <- as.numeric(myExchangeFileH3$SA$SAsampleWeightMeasured)
	myExchangeFileH3$SS$SSnumberSampled <- as.numeric(myExchangeFileH3$SS$SSnumberSampled)
	myExchangeFileH3$SS$SSnumberTotal <- as.numeric(myExchangeFileH3$SS$SSnumberTotal)

	myExchangeFileH3$FO$FOnumberSampled <- 1
	myExchangeFileH3$FO$FOnumberTotal <- 1

	myExchangeFileH3$FT$FTnumberSampled <- as.numeric(myExchangeFileH3$FT$FTnumberSampled)
	myExchangeFileH3$FT$FTnumberTotal <- as.numeric(myExchangeFileH3$FT$FTnumberTotal)

	myExchangeFileH3$VS$VSinclusionProb <- 0.1923077 ## ~ vessel coverage

# ## FROM LIZ CODE IN THE GITHUB ('WKRATIO estimates.R') #################################################################
## calculate inclusion probabilities assuming SRS within strata

	myExchangeFileH3$SA$SAinclusionProb <- myExchangeFileH3$SA$SAnumberSampled/myExchangeFileH3$SA$SAnumberTotal
	myExchangeFileH3$SA[is.na(myExchangeFileH3$SA$SAinclusionProb),]$SAinclusionProb <- 1
	myExchangeFileH3$SS$SSinclusionProb <- myExchangeFileH3$SS$SSnumberSampled/myExchangeFileH3$SS$SSnumberTotal
	myExchangeFileH3$FO$FOinclusionProb <- myExchangeFileH3$FO$FOnumberSampled/myExchangeFileH3$FO$FOnumberTotal
	myExchangeFileH3$FT$FTinclusionProb <- myExchangeFileH3$FT$FTnumberSampled/myExchangeFileH3$FT$FTnumberTotal

# caclulate quarter from date, and then make a domain of quarter combined with area
# allocate domains to FM as well
	myExchangeFileH3$SA$FOid <- myExchangeFileH3$SS$FOid[match(myExchangeFileH3$SA$SSid,myExchangeFileH3$SS$SSid)]
	myExchangeFileH3$SA$date <- myExchangeFileH3$FO$FOendDate[match(myExchangeFileH3$SA$FOid,myExchangeFileH3$FO$FOid)]
	myExchangeFileH3$SA$quarter <- paste("Q",(as.numeric(substr(myExchangeFileH3$SA$date,6,7))-1) %/% 3 + 1,sep="")
	myExchangeFileH3$SA$SAarea <- '27.9.a'
	myExchangeFileH3$SA$domain <- paste(myExchangeFileH3$SA$quarter,myExchangeFileH3$SA$SAarea)
	myExchangeFileH3$FM$domain <- myExchangeFileH3$SA$domain[match(myExchangeFileH3$FM$SAid,myExchangeFileH3$SA$SAid)]
	myExchangeFileH3$BV$domain <- myExchangeFileH3$FM$domain[match(myExchangeFileH3$BV$FMid,myExchangeFileH3$FM$FMid)]

# a function to calculate inclusion probabilities for the units at the final stage
# of sampling given all the other stages
	getIncProb <- function(RDB,stages){
	  #browser()
	  nStages <- length(stages)
	  if (any(stages %in% c("FM"))) {
		RDB[["FM"]][["FMinclusionProb"]] <- 1
	  }
	#browser()
	 RDB[[stages[[1]]]][["inclusionProb"]] <- RDB[[stages[[1]]]][[paste(stages[[1]],"inclusionProb",sep="")]]
	  for (i in 2:(nStages)) {
	  #browser()
		indx <- RDB[[stages[[i]]]][[paste(stages[[i-1]],"id",sep="")]]
		indxPrev <- RDB[[stages[[i-1]]]][[paste(stages[[i-1]],"id",sep="")]]
		RDB[[stages[[i]]]][["incProbPrev"]] <- RDB[[stages[[i-1]]]][[paste("inclusionProb",sep="")]][match(indx,indxPrev)]
		RDB[[stages[[i]]]][["inclusionProb"]] <- RDB[[stages[[i]]]][["incProbPrev"]]*RDB[[stages[[i]]]][[paste(stages[[i]],"inclusionProb",sep="")]]
	  }
	  return(RDB)
	}

# first do Number-at-Length (H5)
	stages <- list("VS","FT","FO","SS","SA","FM")
	myExchangeFileH3 <- getIncProb(myExchangeFileH3,stages)

###################################################################################################################################

### Discards estimation

## To calculate discards per hour in each trip (DpueTrp)
	trip_info <- myExchangeFileH3$FT %>% transmute(FTid, FTarrivalDate, FTquarter= quarters(as.Date(FTarrivalDate)),FishingDays=as.numeric((as.Date(FTarrivalDate)-as.Date(myExchangeFileH3$FT$FTdepartureDate))+1), FTnumberOfHauls=as.numeric(FTnumberOfHauls), FTnumberTotal=as.numeric(FTnumberTotal), FTnumberSampled=as.numeric(FTnumberSampled), FTinclusionProb, incProbPrev, inclusionProb, VSid)
	trip_info <- droplevels(trip_info[trip_info$FishingDays<4,])
	haul_info <- myExchangeFileH3$FO %>% transmute(FOid, FOendDate, FOduration=as.numeric(FOduration), FOarea, FOmetier5, FOsampled, incProbPrev, inclusionProb, FTid )
	haul_info$FishingDays <- trip_info$FishingDays[match(haul_info$FTid, trip_info$FTid)]

	samp <- myExchangeFileH3$SA[,c('SAid','SAspeciesCode','SAspeciesCodeFAO','SAcatchCategory','SAmetier5','SAtotalWeightLive','SAsampleWeightLive','SAnumberTotal','SAnumberSampled','SAinclusionProb','SAtotalWeightMeasured',
				'SAsampleWeightMeasured', 'quarter', 'incProbPrev', 'inclusionProb', 'SSid', 'FOid')]

	samp$FTid <- haul_info$FTid[match(samp$FOid, haul_info$FOid)]
	samp$SAtotalWeightMeasured <- samp$SAtotalWeightMeasured/1000 # to have weight in Kg

	haul_info_dis <- merge(haul_info[haul_info$FOsampled=='Y',], samp[samp$SAcatchCategory=='Dis',c('FOid','SAtotalWeightMeasured')], by='FOid', all.x=T)
	haul_info_dis$FOduration <- round(haul_info_dis$FOduration/60,2) ## duration in hours
	haul_info_dis$FishingDays <- trip_info$FishingDays[match(haul_info_dis$FTid, trip_info$FTid)]
	haul_info_dis$dpue <- round(haul_info_dis$SAtotalWeightMeasured/haul_info_dis$FOduration,2)
	haul_info_dis[is.na(haul_info_dis$dpue),]$dpue <- 0; haul_info_dis[is.na(haul_info_dis$SAtotalWeightMeasured),]$SAtotalWeightMeasured <- 0

## Estimation of DPUE at trip level
## DpueTrp, total duration and total discards for sampled hauls
	DpueTrp <- data.frame(haul_info_dis %>% group_by(FOmetier5,FishingDays,FTid) %>% summarise(DpueTrp=sum(SAtotalWeightMeasured)/sum(FOduration),
				totDuration=sum(FOduration), totDisc=sum(SAtotalWeightMeasured)))

## To raise discards
	totDurationTrp <- data.frame(haul_info_dis %>% group_by(FTid) %>% summarise(totDurationTrp=sum(FOduration)))
	trip_info_dis <- merge(trip_info, DpueTrp[,c('FTid','totDuration','totDisc','DpueTrp')], by='FTid', all.x=T)
	trip_info_dis$totDurationTrp <- totDurationTrp$totDurationTrp[match(trip_info_dis$FTid, totDurationTrp$FTid)]
## Get metier of the trip from FO table (P.S.- in OTB fleet all hauls from a trip have the same metier5...)
	trip_info_dis$FOmetier5 <- haul_info$FOmetier5[match(trip_info_dis$FTid, haul_info$FTid)]

## mean DPUE and variance estimation
	meanDpueTrp <- tapply(trip_info_dis$totDisc, list(trip_info_dis$FOmetier5, trip_info_dis$FishingDays), sum, na.rm=T)/tapply(trip_info_dis$totDuration, list(trip_info_dis$FOmetier5, trip_info_dis$FishingDays), sum, na.rm=T)

# variance of mean DPUE
	discTrp <- tapply(trip_info_dis$totDisc, list(trip_info_dis$FTid), sum, na.rm=T)
	durTrp <- tapply(trip_info_dis$totDuration, list(trip_info_dis$FTid), sum, na.rm=T)
	nDaysTrp <- tapply(trip_info_dis$FishingDays, list(trip_info_dis$FTid), unique)
	metierTrp <- tapply(trip_info_dis$FOmetier5, list(trip_info_dis$FTid), unique)
	meanDpueTrp2 <-  tapply((discTrp/durTrp)^2, list(metierTrp, nDaysTrp), na.rm=TRUE, sum)
	nTrp <- tapply(trip_info_dis$FTid, list(trip_info_dis$FOmetier5, trip_info_dis$FishingDays), function(x){length(unique(x))})

	varDpueTrp <- (meanDpueTrp2-nTrp*meanDpueTrp^2)/((nTrp-1)*nTrp)	#

## Discard estimates for total duration by fishing days: popDiscPerDaysWght
	## Total 'Fishing duration' from "CE table"
	popDurationPerDays <- tapply(effTable$CEoffVesFishHour, list(effTable$CEfishTech, effTable$CEoffFishDay), sum, na.rm=TRUE)

	## Discard estimates by fishing days
	popDiscPerDays <- (popDurationPerDays * meanDpueTrp)/1000

	## Variance estimates for total duration by fishing days (popVarPerDays)
	popVarPerDays <- (popDurationPerDays^2 * varDpueTrp)/1000000

	## Total discards and variance, by fleet
	popDiscTotalMetier <- apply(popDiscPerDays, c(1), sum, na.rm=TRUE)
	popVarTotalMetier <- apply(popVarPerDays, c(1), sum, na.rm=TRUE)


## For colocalisation analysis need mQTL P values for all SNPs not just significant ones
## This script extracts DNA methylation sites found within 250 kb of a ASD GWAS loci as defined by the clumped results of the
## PGC-iPSYCH meta-analysis and creates files for MatrixEQTL.

setwd("")
meth<-read.table("MatrixEQTL/Methylation.txt", header = TRUE, row.names = 1)

library(IlluminaHumanMethylation450kanno.ilmn12.hg19)
data(IlluminaHumanMethylation450kanno.ilmn12.hg19)
probeAnnot<-as.data.frame(IlluminaHumanMethylation450kanno.ilmn12.hg19@data$Locations)
probeAnnot<-probeAnnot[rownames(meth),]

### only first 3 are genome-wide significant
regions<-read.table("Analysis/AdditionalAutismRegions.txt", header = TRUE, row.names = 1)

for(i in 1:3){

chr<-regions$CHR[i]
start<-regions$BP[i]-250000
stop<-regions$BP[i]+250000
probes<-as.character(rownames(probeAnnot)[which(probeAnnot$chr == paste("chr", chr, sep = "") &
  probeAnnot$pos <= stop & probeAnnot$pos >= start)])
write.table(meth[intersect(probes, rownames(meth)),],
  paste("MatrixEQTL/Methylation_ProbesForASDColocAnalysis_Chr", chr, "_", start, "_", stop, ".txt", sep = ""),
  sep = "\t", quote = FALSE)

}


# View NEON out-of-sample validation density plots by spatial scale

rm(list=ls())
source('paths.r')
#source('NEFI_functions/zero_truncated_density.r')
# source dmulti_ddirch_forecast
library(RCurl)
script <- getURL("https://raw.githubusercontent.com/colinaverill/NEFI_microbe/master/NEFI_functions/zero_truncated_density.r", ssl.verifypeer = FALSE)
eval(parse(text = script))

#set output path.----
output.path <- 'test.png'

#Load calibration data.----
pl <- readRDS(bahram_16S_prior_dmulti.ddirch_all.group_JAGSfits) #all phylo and functional groups.
fg <- readRDS("/projectnb/talbot-lab-data/NEFI_data/16S/scc_gen/JAGS_output/bahram_16S_prior_dmulti.ddirch_fg_JAGSfits.rds")
pl <- c(pl[1:5], fg)
#re-order list, make function group first.
#pl <- pl[c('fg','phylum','class','order','family','genus')]

#Get calibration r2 values.----
pl.r2 <- list()
lev.r2.out <- list()
for(i in 1:length(pl)){
  lev <- pl[[i]]
  obs <- lev$observed
  pred <- lev$predicted
  lev.r2 <- list()
  for(j in 1:ncol(obs)){lev.r2[[j]] <- summary(lm(obs[,j] ~ pred[,j]))$r.squared}
  lev.r2 <- unlist(lev.r2)
  names(lev.r2) <- colnames(obs)
  lev.r2 <- lev.r2[names(lev.r2) != 'other']
  pl.r2[[i]] <- lev.r2
  #names(pl.r2)[[i]] <- names(pl)[i]
}

fg <- unname(pl.r2[6:17])
fg.r2.all <- list(unlist(fg))
names(fg.r2.all)<- "fg"
pl <- pl.r2[1:5]
names(pl) <- c('phylum','class','order','family','genus')
all.r2 <- c(fg.r2.all, pl)
cal.rsq <- unlist(all.r2)

#load forecasts predicted and observed.----
pl.cast <- readRDS(NEON_cps_fcast_dmulti.ddirch_16S.path)
pl.truth <- readRDS(NEON_phylo_fg_plot.site_obs_16S.path)
pl.core <- readRDS(NEON_16S_phylo_fg_abundances.path)
#
# #Re-order.
# pl.cast <- pl.cast[c('fg','phylum','class','order','family','genus')]
# pl.truth <- pl.truth[c('fg','phylum','class','order','family','genus')]
# pl.core <- pl.core[c('fg','phylum','class','order','family','genus')]

#names.
#names(pl.cast)[1] <- 'functional'

#get core, plot site R2 values out of sample.----
all.core.rsq <- list()
all.plot.rsq <- list()
all.site.rsq <- list()
for(i in 1:length(pl.cast)){
  fcast <- pl.cast[[i]]
  core.rsq <- list()
  plot.rsq <- list()
  site.rsq <- list()
  #core.level----
  y <- pl.truth[[i]]$core.fit
  #y <- (pl.core[[i]]$abundances + 1)/pl.core[[i]]$seq_total
  x <- fcast$core.fit$mean
  #make sure row and column orders match.
  rownames(y) <- gsub('-GEN','',rownames(y))
  truth <- y
  map <- readRDS(core_obs_16S.path)
  truth <- as.data.frame(truth)
  truth$deprecatedVialID <- rownames(truth)
  truth1 <- merge(truth, map[,c("deprecatedVialID", "geneticSampleID")], by = "deprecatedVialID")
  truth1 <- truth1[!duplicated(truth1$geneticSampleID),]
  rownames(truth1) <- gsub('-GEN','',truth1$geneticSampleID)
  truth <- truth1
  y <- truth

  y <- y[rownames(y) %in% rownames(x),]
  x <- x[rownames(x) %in% rownames(y),]
  y <- y[,colnames(y) %in% colnames(x)]
  x <- x[,colnames(x) %in% colnames(y)]
  x <- x[order(match(rownames(x),rownames(y))),]
  x <- x[,order(match(colnames(x),colnames(y)))]
  #fit model, grab r2.
  for(k in 1:ncol(x)){
    fungi_name <- colnames(x)[k]
    rsq <- summary(lm(y[,k] ~ x[,k]))$r.squared
    names(rsq) <- fungi_name
    core.rsq[[k]] <- rsq
  }
  #plot.level----
  x <- fcast$plot.fit$mean
  y <- pl.truth[[i]]$plot.fit$mean
  #make sure row and column order match.
  rownames(y) <- gsub('\\.','_',rownames(y))
  y <- y[rownames(y) %in% rownames(x),]
  x <- x[rownames(x) %in% rownames(y),]
  y <- y[,colnames(y) %in% colnames(x)]
  x <- x[,colnames(x) %in% colnames(y)]
  x <- x[order(match(rownames(x),rownames(y))),]
  x <- x[,order(match(colnames(x),colnames(y)))]
  #fit model, grab r2.
  for(k in 1:ncol(y)){
    fungi_name <- colnames(x)[k]
    rsq <- summary(lm(y[,k] ~ x[,k]))$r.squared
    names(rsq) <- fungi_name
    plot.rsq[[k]] <- rsq
  }
  #site.level----
  x <- fcast$site.fit$mean
  y <- pl.truth[[i]]$site.fit$mean
  #make sure row and column order match.
  rownames(y) <- gsub('-GEN','',rownames(y))
  y <- y[rownames(y) %in% rownames(x),]
  x <- x[rownames(x) %in% rownames(y),]
  y <- y[,colnames(y) %in% colnames(x)]
  x <- x[,colnames(x) %in% colnames(y)]
  x <- x[order(match(rownames(x),rownames(y))),]
  x <- x[,order(match(colnames(x),colnames(y)))]
  #fit model, grab r2.
  for(k in 1:ncol(y)){
    fungi_name <- colnames(x)[k]
    rsq <- summary(lm(y[,k] ~ x[,k]))$r.squared
    names(rsq) <- fungi_name
    site.rsq[[k]] <- rsq
  }
  #wrap up for return.----
  all.core.rsq[[i]] <- unlist(core.rsq)
  all.plot.rsq[[i]] <- unlist(plot.rsq)
  all.site.rsq[[i]] <- unlist(site.rsq)

  all.core.rsq[[i]] <-  all.core.rsq[[i]][-grep('other',names(all.core.rsq[[i]]))]
  all.plot.rsq[[i]] <-  all.plot.rsq[[i]][-grep('other',names(all.plot.rsq[[i]]))]
  all.site.rsq[[i]] <-  all.site.rsq[[i]][-grep('other',names(all.site.rsq[[i]]))]

}
core.rsq <- unlist(all.core.rsq)
plot.rsq <- unlist(all.plot.rsq)
site.rsq <- unlist(all.site.rsq)

# fix the functional-group grouping for each level
core.rsq.fg <- list(unlist(all.core.rsq[6:17]))
names(core.rsq.fg)<- "fg"
core.rsq.pl <- all.core.rsq[1:5]
names(core.rsq.pl) <- c('phylum','class','order','family','genus')
core.rsq <- c(core.rsq.fg, core.rsq.pl)
core.rsq <- unlist(core.rsq)

plot.rsq.fg <- list(unlist(all.plot.rsq[6:17]))
names(plot.rsq.fg)<- "fg"
plot.rsq.pl <- all.plot.rsq[1:5]
names(plot.rsq.pl) <- c('phylum','class','order','family','genus')
plot.rsq <- c(plot.rsq.fg, plot.rsq.pl)
plot.rsq <- unlist(plot.rsq)

site.rsq.fg <- list(unlist(all.site.rsq[6:17]))
names(site.rsq.fg)<- "fg"
site.rsq.pl <- all.site.rsq[1:5]
names(site.rsq.pl) <- c('phylum','class','order','family','genus')
site.rsq.list <- c(site.rsq.fg, site.rsq.pl)
site.rsq <- unlist(site.rsq.list)

# shouldn't we be getting these values *after* subsetting by calibration R2?
lev.mu <- unlist(lapply(site.rsq.list, mean  ))
lev.sd <- unlist(lapply(site.rsq.list, sd    ))
lev.N  <- unlist(lapply(site.rsq.list, length))
lev.se <- lev.sd / sqrt(lev.N)
names(lev.mu) <- names(site.rsq.list)

#Subset to observations that have a minimum calibration R2 value.----
pass <- cal.rsq[cal.rsq > .1]
core.rsq <- core.rsq[names(core.rsq) %in% names(pass)]
plot.rsq <- plot.rsq[names(plot.rsq) %in% names(pass)]
site.rsq <- site.rsq[names(site.rsq) %in% names(pass)]

# get densities
core.d <- zero_truncated_density(core.rsq)
plot.d <- zero_truncated_density(plot.rsq)
site.d <- zero_truncated_density(site.rsq)

#png save line.----
png(filename=output.path,width=8,height=5,units='in',res=300)

#global plot settings.----
par(mfrow = c(1,2))
limx <- c(0,1)
limy <- c(0, 8)
trans <- 0.2 #shading transparency.
o.cex <- 1.3 #outer label size.
cols <- c('purple','cyan','yellow')
par(mfrow = c(1,2), mar = c(5,4.2,1.5,1.5))

#Density plot.----
plot(site.d,xlim = c(0, 0.89), ylim = limy, bty = 'n', xlab = NA, ylab = NA, main = NA, yaxs='i', xaxs = 'i', las = 1, lwd = 0)
polygon(site.d, col = adjustcolor(cols[3],trans))
polygon(plot.d, col = adjustcolor(cols[2],trans))
polygon(core.d, col = adjustcolor(cols[1],trans))
mtext('Density', side = 2, line = 2.2, cex = o.cex)
mtext(expression(paste("Validation R"^"2")), side = 1, line = 2.5, cex = o.cex)
legend(x = 0.6, y = 6, legend = c('core','plot','site'), col ='black', pt.bg=adjustcolor(cols,trans), bty = 'n', pch = 22, pt.cex = 1.5)

#Validation rsq ~ function/phylo scale.----
x <- 1:length(lev.mu)
limy <- c(0,max(lev.mu + lev.se))
plot(lev.mu ~ x, cex = 2.5, ylim = limy, pch = 16, ylab = NA, xlab = NA, bty='n', xaxt = 'n')
arrows(x, lev.mu - lev.se, x1 = x, y1 = lev.mu + lev.se, length=0.00, angle=90, code=3, col = 'black')
lines(x, lev.mu, lty = 2)
mtext(expression(paste("Validation R"^"2")), side = 2, line = 2.2, cex = o.cex)
axis(1, labels = F)
text(x=x, y = -.04, labels= names(lev.mu), srt=45, adj=1, xpd=TRUE)

#end plot.----
dev.off()


# Code by Katie Devenish. ktd19ycv@bangor.ac.uk

# This code is used in the publication entitled "Madagascar's biggest mine is set to achieve No Net Loss of Forest"
# Authors: Katie Devenish, Sebastien Desbureaux, Simon Willcock and Julia Jones. 2021.

# 2- We demonstratate that adding a time trend in the DID regressions doesnt change the results

# Load libraries #

library("foreign")
library("dplyr")
library("xlsx")
library("ggplot2")
library("broom")
library("MatchIt")
library("ggrepel")
library("calibrate")
library("plm")
library("gridExtra")
library("tidyr")
library("gplots")
library("lmtest")
library("stargazer")

# Code Structure

# 1) Old code in bracket
# 2) Did with the additions

# --------------------- 1) Data Construction for Matching -----------------------------------------#

setwd("c:/Users/desbures/Documents/Sebastien/11_Autres/Katie/Ambatovy/")

# Function to read in input data for each offset, remove unnecessary columns and add column to indicate treatment status
# and offset of origin.

{
tidy_data <-function(path, name, number, label){
  name = read.dbf(path)
  name = subset(name, select= -c(1, 16))
  name$treated = number              # Number = 1 for pixels from an offset and 0 for control pixels.
  name$offset = label
  return(name)
}

TTF <- tidy_data("Input_data/Final_covariates/Sample_TTF3.dbf", TTF, 1, "TTF")
ANK <- tidy_data("Input_data/Final_covariates/Sample_ANK3.dbf", ANK, 1, "ANK")
CFAM <- tidy_data("Input_data/Final_covariates/Sample_CFAM3.dbf", CFAM, 1, "CFAM")
CZ <- tidy_data("Input_data/Final_covariates/Sample_CZ3.dbf", CZ, 1, "CZ")

# TTF = Torotorofotsy
# ANK = Ankerana
# CFAM = Corridor Forestier Analamay-Mantadia
# CZ = Conservation Zone

# Load Control separately because columns are different #

Control = read.dbf("Input_data/Final_covariates/Final_control.dbf")
Control$treated = 0
Control$offset = "Cont"
Control <- subset(Control, select = -c(4))        # Remove unwanted fire variable
Control <- Control[,c(2,3,1,4:16)]                # Re-order columns to match offset dataframes

# Rename columns #

cols <- c("X","Y","Tree_loss", "Pop_density", "Dist_sett", "Slope", "Elevation", "Aspect", "Annual_Rain",
          "Dist_track", "Dist_road", "Dist_river", "Dist_edge", "Dist_defor", "treated", "offset")

names(Control) <- paste0(cols)
names(ANK) <- paste0(cols)
names(CZ) <- paste0(cols)
names(CFAM) <- paste0(cols)
names(TTF) <- paste0(cols)

# Merge each offset with the control dataset
# These offset + control datasets will be the input for matching #

ANKCONT <- rbind(ANK, Control)
CZCONT <- rbind(CZ, Control)
CFAMCONT <- rbind(CFAM, Control)
TTFCONT <- rbind(TTF, Control)

# Data cleaning

# - remove erroneous -9999 values from missing values in the raster data layers
# - create calipers to drop extreme observations that will never be matched
# - caliper based on the distribution of values for the 5 essential
#   covariates within the treated sample.

Seb_dataclean <- function(data, x){

  # Replace all -9999 with NA
  data <- data %>% na_if(-9999)

  # We keep observations without na
  data <- data %>% drop_na()

  # We determine the caliper
  caliper_dist_def<-x*sd(data$Dist_defor[data$treated==1])
  caliper_slope<-x*sd(data$Slope[data$treated==1])
  caliper_elev<-x*sd(data$Elevation[data$treated==1])
  caliper_edge<-x*sd(data$Dist_edge[data$treated==1])
  caliper_road <- x*sd(data$Dist_road[data$treated==1])

  # We keep control observations that are on the same support
  data <- data[data$Dist_defor < max(data$Dist_defor[data$treated==1])+caliper_dist_def,]
  data <- data[data$Dist_defor > min(data$Dist_defor[data$treated==1])- caliper_dist_def,]

  data <- data[data$Slope < max(data$Slope[data$treated==1])+caliper_slope,]
  data <- data[data$Slope > min(data$Slope[data$treated==1])- caliper_slope,]

  data <- data[data$Elevation < max(data$Elevation[data$treated==1])+caliper_elev,]
  data <- data[data$Elevation > min(data$Elevation[data$treated==1])-caliper_elev,]

  data <- data[data$Dist_edge < max(data$Dist_edge[data$treated==1])+caliper_edge,]
  data <- data[data$Dist_edge > min(data$Dist_edge[data$treated==1])-caliper_edge,]

  data <- data[data$Dist_road < max(data$Dist_road[data$treated==1])+caliper_road,]
  data <- data[data$Dist_road > min(data$Dist_road[data$treated==1])-caliper_road,]

}

ANKCONT <- Seb_dataclean(data = ANKCONT, x= 1)
ANKCONT$ID <- seq(nrow(ANKCONT))              # Add column for observation ID
rownames(ANKCONT) <- ANKCONT$ID

CZCONT <- Seb_dataclean(CZCONT, 1)
CZCONT$ID <- seq(nrow(CZCONT))
rownames(CZCONT) <- CZCONT$ID

CFAMCONT <- Seb_dataclean(CFAMCONT, 1)
CFAMCONT$ID <- seq(nrow(CFAMCONT))
rownames(CFAMCONT) <- CFAMCONT$ID

TTFCONT <- Seb_dataclean(TTFCONT, 1)
TTFCONT$ID <- seq(nrow(TTFCONT))
rownames(TTFCONT) <- TTFCONT$ID

#----------------------------2) The Matching -----------------------------------------#

# The matching algorithm and selected arguments Ho et al (2007).
# Matching run separately for each offset.

variables <- c("Slope", "Elevation", "Dist_road", "Dist_edge", "Dist_defor")

# Set caliper for each covariate

cal <- rep(1, length(variables))
names(cal) <- variables

# The main matching specification. 1:1 nearest-neighbour matching without replacement using Mahalanobis distance measure
# and a caliper of 1sd. The other 4 specifications tested for balance can be found in the script entitled "Choosing_matching_spec"

matching <- function(offset){
  output = matchit(treated ~ Slope + Elevation + Dist_road + Dist_edge + Dist_defor,
                   data= offset, method= "nearest", distance = "mahalanobis", replace = FALSE, caliper=cal)
  return(output)
}

m.out.ANK <- matching(offset = ANKCONT)
m.out.CFAM <- matching(offset = CFAMCONT)
m.out.CZ <- matching(offset = CZCONT)
m.out.TTF <- matching(offset = TTFCONT)

# Summarise results #

List1 <- list(m.out.ANK, m.out.CFAM, m.out.CZ, m.out.TTF)

sum1_matching <- function(matched_data){
  sum1 = summary(matched_data, standardize = TRUE)      # Standardise = TRUE to get standardised mean difference in output
  return(sum1)
}

Output_names <- c("m.out.ANK", "m.out.CFAM", "m.out.CZ", "m.out.TTF")

allOutputs <- lapply(List1, sum1_matching)              # Apply the summary function to all matched datasets in List1

names(allOutputs) <- paste0(Output_names)

# Repeat but with standardise = FALSE to get the mean EQQ values. This is a second measure to assess covariate balance. #

sum2_matching <- function(matched_data){
  sum2 = summary(matched_data, standardize = FALSE)
  return(sum2)
}

allOutputs2 <- lapply(List1, sum2_matching)
names(allOutputs2) <- paste0(Output_names)

#---------------------------3) Extract matched pairs from matching output ------------------------------------#

# Extract matched pairs from input data

m.data.ANK <- match.data(m.out.ANK, group = "all", distance = "distance", weights = "weights", subclass = "subclass",
                         data = ANKCONT, drop.unmatched = TRUE)

m.data.CZ <- match.data(m.out.CZ, group = "all", distance = "distance", weights = "weights", subclass = "subclass",
                        data = CZCONT, drop.unmatched = TRUE)

m.data.CFAM <- match.data(m.out.CFAM, group = "all", distance = "distance", weights = "weights", subclass = "subclass",
                          data = CFAMCONT, drop.unmatched = TRUE)

m.data.TTF <- match.data(m.out.TTF, group = "all", distance = "distance", weights = "weights", subclass = "subclass",
                         data = TTFCONT, drop.unmatched = TRUE)

# ---------------------------4) Data Construction for DiD Regression using Matched Pairs -------------------------#

# Aggregate pixels into treated (offset) and control samples and tabulate observations within each sample by tree loss year.
# This gives the count of pixels within each sample deforested each year.

annual_defor_ANK <- data.frame(table(factor(m.data.ANK$Tree_loss, levels = 1:19), m.data.ANK$offset))

# levels = 1:19 removes observations with 0 value for Tree Loss year (which were not deforested over the study period).
# This is because we are interested in comparing deforestation outcomes between offsets and the matched controls.

annual_defor_CZ <- data.frame(table(factor(m.data.CZ$Tree_loss, levels = 1:19), m.data.CZ$offset))
annual_defor_CZ$Var2 <- factor(annual_defor_CZ$Var2, levels = c("CZ", "Cont"))    # For the plots the order of the factors needs to match the other offsets

annual_defor_CFAM <- data.frame(table(factor(m.data.CFAM$Tree_loss, levels = 1:19), m.data.CFAM$offset))

annual_defor_TTF <- data.frame(table(factor(m.data.TTF$Tree_loss, levels = 1:19), m.data.TTF$offset))
annual_defor_TTF$Var2 <- factor(annual_defor_TTF$Var2, levels = c("TTF", "Cont"))

label <- c("Year", "Sample", "Annual_Deforestation")
names(annual_defor_ANK) <- paste0(label)
names(annual_defor_CFAM) <- paste0(label)
names(annual_defor_CZ) <- paste0(label)
names(annual_defor_TTF) <- paste0(label)

# But offsets are different sizes - need to calculate annual deforestation as a percentage of total pixels in the sample to plot

# Calculating Percentage Annual Deforestation with 1:1 matching.

annual_defor_ANK$Perc_Annual_Defor <- (annual_defor_ANK$Annual_Deforestation/(nrow(m.data.ANK)/2))*100          # In the matched dataset the no. of control pixels = No. of treatment pixels because I matched 1:1.
annual_defor_CFAM$Perc_Annual_Defor <- (annual_defor_CFAM$Annual_Deforestation/(nrow(m.data.CFAM)/2))*100       # Therefore, nrow(m.data.CFAM)/2 is the total number of treatment and the total number of control pixels
annual_defor_CZ$Perc_Annual_Defor <- (annual_defor_CZ$Annual_Deforestation/(nrow(m.data.CZ)/2))*100             # in the matched dataset.
annual_defor_TTF$Perc_Annual_Defor <- (annual_defor_TTF$Annual_Deforestation/(nrow(m.data.TTF)/2))*100

}

#-------------------------------2) Difference in Differences with time trends -----------------------------------------#

# a) Data Construction

Data_construction_DiD <- function(offset, y){    # y corresponds to the year of protection - so Time = 0 before protection and 1 after protection.
  offset$Year <- as.numeric(rep(1:19,2))        # Have to make Year numeric for >= to work
  offset$TimeF <- factor(ifelse(offset$Year >= y, 1,0))
  offset$TreatedF <- factor(ifelse(offset$Sample != "Cont", 1,0))
  return(offset)
}

annual_defor_ANK <- Data_construction_DiD(annual_defor_ANK, 11)
annual_defor_CFAM <- Data_construction_DiD(annual_defor_CFAM, 13)
annual_defor_CZ <- Data_construction_DiD(annual_defor_CZ, 9)
annual_defor_TTF <- Data_construction_DiD(annual_defor_TTF, 14)

# b) Outcome variable transformation

# log(y+1) transformation of outcome variable required because non-normal properties of count data violate assumptions of
# homoscedascity of linear models.

annual_defor_ANK$log_annual_defor <- log(annual_defor_ANK$Annual_Deforestation + 1)
annual_defor_CZ$log_annual_defor <- log(annual_defor_CZ$Annual_Deforestation + 1)
annual_defor_CFAM$log_annual_defor <- log(annual_defor_CFAM$Annual_Deforestation +1)
annual_defor_TTF$log_annual_defor <- log(annual_defor_TTF$Annual_Deforestation +1)

# c) Test for parallel trends

# Parallel trends in outcomes between treated and control samples in the years before the intervention
# is a key assumption of difference-in-differences regressions.

# Use only data from the years before the offsets were protected #

ANK_data_before <- annual_defor_ANK[(annual_defor_ANK$Year <11),]
CFAM_data_before <- annual_defor_CFAM[(annual_defor_CFAM$Year <13),]
CZ_data_before <- annual_defor_CZ[(annual_defor_CZ$Year <9),]
TTF_data_before <- annual_defor_TTF[(annual_defor_TTF$Year <14),]

# ANK #

ANKa <- lm(log_annual_defor ~ Year*TreatedF, data= ANK_data_before)
summary(ANKa)       # If the interaction between Year and TreatedF is not significant, there is no
# significant difference in the relationship between Year and the log-transformed
# count of deforestation between treated and control samples --> parallel trends assumption holds.

# CFAM #

CFAMa <- lm(log_annual_defor ~ Year*TreatedF, data= CFAM_data_before)
summary(CFAMa)

# There is a significant difference in the trend in deforestation over time between treated and control samples
# No parallel trends.

# CZ #

CZa <- lm(log_annual_defor ~ Year*TreatedF, data= CZ_data_before)
summary(CZa)

# TTF #

TTFa <- lm(log_annual_defor ~ Year*TreatedF, data= TTF_data_before)
summary(TTFa)

# All showed parallel trends except CFAM which cannot be used in individual DiD regressions

# d) DiD Regression

# Formula = y ~ treatment + time + (treatment x time)

# Interaction between treated and time is the coefficent of interest. This represents the effect
# of an observation being in an offset, after protection on the log-transformed count of deforestation.
# If this is significant and negative it means protection significantly reduced deforestation within the offset,
# relative to the counterfactual.

# ANK #

modelANK <- lm(log_annual_defor ~ TreatedF*TimeF, data= annual_defor_ANK)
summary(modelANK)

modelANKtt <- lm(log_annual_defor ~ TreatedF*TimeF + Year, data= annual_defor_ANK)
summary(modelANKtt)

# CZ #

modelCZ <- lm(log_annual_defor ~ TreatedF*TimeF, data= annual_defor_CZ)
summary(modelCZ)

modelCZtt <- lm(log_annual_defor ~ TreatedF*TimeF + Year, data= annual_defor_CZ)
summary(modelCZtt)

# TTF #

modelTTF <- lm(log_annual_defor ~ TreatedF*TimeF, data= annual_defor_TTF)
summary(modelTTF)

modelTTFtt <- lm(log_annual_defor ~ TreatedF*TimeF + Year, data= annual_defor_TTF)
summary(modelTTFtt)

stargazer(modelANK,modelANKtt,modelCZ,modelCZtt,modelTTF,modelTTFtt,
          title = "DiD with time trends",
          out="did_tt.doc",
          type = "html")


createComposites <- function(data) {

  data$TotalActivityOccur <- data$SniffAnoEv + data$SniffHeadEv + data$DigEv + data$ActiveHEv +
    data$AvoidEv + data$SubEv + data$DefenseEv + data$OpenAggEv +
    data$RitAggEv + data$ActiveVEv + data$ApproachEv + data$StereoEv +
    data$DomEv + data$SniffBodyEv + data$FollowEv + data$StretchEv +
    data$AttackDelEv + data$AttackRecEv

  data$TotalActivityDuration <- data$SniffAnoDur + data$SniffHeadDur + data$DigDur + data$ActiveHDur +
    data$AvoidDur + data$SubDur + data$DefenseDur + data$OpenAggDur +
    data$RitAggDur + data$ActiveVDur + data$ApproachDur + data$StereoDur +
    data$DomDur + data$SniffBodyDur + data$FollowDur

  data$TotalSocialOccur <- data$SniffAnoEv + data$SniffHeadEv + data$AvoidEv + data$SubEv +
    data$DefenseEv + data$OpenAggEv + data$RitAggEv + data$ApproachEv + data$InactiveTogetherEv +
    data$DomEv + data$SniffBodyEv + data$FollowEv + data$StretchEv +
    data$AttackDelEv + data$AttackRecEv

  data$TotalSocialDuration <- data$SniffAnoDur + data$SniffHeadDur + data$AvoidDur + data$SubDur +
    data$DefenseDur + data$OpenAggDur + data$RitAggDur + data$InactiveTogetherDur +
    data$ApproachDur + data$DomDur + data$SniffBodyDur + data$FollowDur

  data$AgonisticDeliveredOccur <- data$FollowEv + data$DomEv + data$AttackDelEv
  data$AgonisticDeliveredDuration <- data$FollowDur + data$DomDur

  data$AgonisticReceivedOccur <- data$AvoidEv + data$SubEv + data$AttackRecEv + data$DefenseEv
  data$AgonisticReceivedDuration <- data$AvoidDur + data$SubDur + data$DefenseDur

  data$AgonisticTotalOccur <- data$AgonisticDeliveredOccur + data$AgonisticReceivedOccur
  data$AgonisticTotalDuration <- data$AgonisticDeliveredDuration + data$AgonisticReceivedDuration

  data$DominanceScoreOccur <- data$AgonisticDeliveredOccur - data$AgonisticReceivedOccur
  data$DominanceScoreDuration <- data$AgonisticDeliveredDuration - data$AgonisticReceivedDuration

  data$SocialInvestigationOccur <- data$SniffHeadEv + data$SniffBodyEv + data$SniffAnoEv + data$StretchEv + data$ApproachEv
  data$SocialInvestigationDuration <- data$SniffHeadDur + data$SniffBodyDur + data$SniffAnoDur + data$ApproachDur

  data$NonSocialOccur <- data$ActiveHEv + data$ActiveVEv + data$DigEv + data$StereoEv + data$InactiveEv +
    data$GroomEv
  data$NonSocialDuration <- data$ActiveHDur + data$ActiveVDur + data$DigDur + data$StereoDur +
    data$InactiveDur + data$GroomDur

  data$NonSocialLocomotorOccur <- data$ActiveHEv + data$ActiveVEv + data$DigEv
  data$NonSocialLocomotorDuration <- data$ActiveHDur + data$ActiveVDur + data$DigDur

  data$InactiveNonSocialOccur <- data$InactiveEv + data$GroomEv
  data$InactiveNonSocialDuration <- data$InactiveDur + data$GroomDur

  return(data)

}

lnTransform <- function(data) {

  data <- log(data + 1)

  return(data)

}

createFactors <- function(data) {
  data$Gonad <- factor(data$Gonad, levels = c("SHAM", "GDX", "REP"))
  data$Sex <- as.factor(data$Sex)
  data$Treatment <- as.factor(data$Treatment)

  return(data)
}

interactionGraph <- function(data, yVar, variableName, graphTitle, yMax, fileName, saveGraph = F) {
  interactionBar <- ggplot(data, aes_string(x = "Gonad", y = yVar, fill = "Treatment")) +
    stat_summary(fun.y = "mean", geom = "bar", position = "dodge") +
    stat_summary(fun.data = mean_se, geom = "errorbar", position = position_dodge(width = 0.9), width = 0.2, size  = 1) +
    coord_cartesian(ylim = c(0, yMax)) +
    scale_y_continuous(expand = c(0, 0)) +

    labs (title = graphTitle, x = NULL, y = variableName) +
    scale_x_discrete(breaks=c("SHAM", "GDX", "REP"), limits=c("SHAM", "GDX", "REP"), labels=c("Intact", "Gonadectomy + \nVehicle", "Gonadectomy + \nReplacement")) +
    scale_fill_manual("Prenatal Treatment", values = c("light gray", "black"), labels = c("Sesame Oil", "Testosterone")) +

    theme_bw() +
    theme(plot.title = element_text(colour="black", size=14, margin = margin(30, 0, 0, 0), face="bold"),
          axis.title.y = element_text(size=14, margin = margin(0, 10, 0, 0), face="bold"),
          axis.title.x = element_text(size=14, face="bold"),
          axis.text.y = element_text(size=12, color = "black"),
          axis.text.x = element_text(size=12, color = "black", face="bold"),
          axis.ticks = element_line(size=1, color = "black"),
          legend.text = element_text(size = 12, colour = "black"),
          legend.title = element_text(size = 14, colour = "black", face="bold"),
          legend.key = element_blank(),
          legend.key.size = unit(1, "cm"),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          panel.border = element_blank(),
          panel.background = element_blank(),
          axis.line.x = element_line(colour = 'black', size = 1),
          axis.line.y = element_line(colour = 'black', size = 1))

  if (saveGraph == T) {
    pdf(paste("output/PostPuberty/", fileName, sep = ""), 9, 6)
    interactionBar
    dev.off()
  } else {
    interactionBar
  }

}

interactionPanel <- function(data, yVar, variableName, graphTitle, yMax, fileName, saveGraph = F) {

  maleGraph <- interactionGraph(subset(data, Sex == "M"), yVar, variableName, "Male", yMax)
  femaleGraph <- interactionGraph(subset(data, Sex == "F"), yVar, variableName, "Female", yMax)
  graphLegend <- getLegend(femaleGraph);

  maleGraph <- maleGraph + theme(legend.position = "none")
  femaleGraph <- femaleGraph + theme(legend.position = "none")

  if (saveGraph == T) {
    pdf(paste("output/PostPuberty/", fileName, sep = ""), 12, 5)
    grid.arrange(maleGraph, femaleGraph, ncol = 2, nrow = 1, widths = c(6, 6), heights = c(5), top = textGrob(graphTitle, gp=gpar(fontsize=18, fontface = "bold")))
    dev.off()
  } else {
    grid.arrange(maleGraph, femaleGraph, ncol = 2, nrow = 1, widths = c(6, 6), heights = c(5), top = textGrob(graphTitle, gp=gpar(fontsize=18, fontface = "bold")))
  }

}

dominancePanel <- function(data, yVar, variableName, graphTitle, yMin, yMax, fileName, saveGraph = F) {

  maleGraph <- dominanceGraph(subset(data, Sex == "M"), yVar, variableName, "Male", yMin, yMax)
  femaleGraph <- dominanceGraph(subset(data, Sex == "F"), yVar, variableName, "Female", yMin, yMax)
  graphLegend <- getLegend(femaleGraph);

  maleGraph <- maleGraph + theme(legend.position = "none")
  femaleGraph <- femaleGraph + theme(legend.position = "none")

  # note, no legend!

  if (saveGraph == T) {
    pdf(paste("output/PostPuberty/", fileName, sep = ""), 12, 5)
    grid.arrange(maleGraph, femaleGraph, ncol = 2, nrow = 1, widths = c(7, 7), top = textGrob(graphTitle, gp=gpar(fontsize=18, fontface = "bold")))
    dev.off()
  } else {
    grid.arrange(maleGraph, femaleGraph, ncol = 2, nrow = 1, widths = c(7, 7), top = textGrob(graphTitle, gp=gpar(fontsize=18, fontface = "bold")))
  }

}

dominanceGraph <- function(data, yVar, variableName, graphTitle, yMin, yMax, fileName, saveGraph = F) {
  dominanceBar <- ggplot(data, aes_string(x = "Gonad", y = yVar, fill = "Treatment")) +
    stat_summary(fun.y = "mean", geom = "bar", position = "dodge") +
    stat_summary(fun.data = mean_se, geom = "errorbar", position = position_dodge(width = 0.9), width = 0.2, size  = 1) +
    coord_cartesian(ylim = c(yMin, yMax)) +
    scale_y_continuous(expand = c(0, 0)) +

    labs (title = graphTitle, x = NULL, y = variableName) +
    scale_x_discrete(breaks=c("SHAM", "GDX", "REP"), limits=c("SHAM", "GDX", "REP"), labels=c("Intact", "Gonadectomy + \nVehicle", "Gonadectomy + \nReplacement")) +
    scale_fill_manual("Prenatal Treatment", values = c("light gray", "black"), labels = c("Sesame Oil", "Testosterone")) +

    theme_bw() +
    theme(plot.title = element_text(colour="black", size=14, margin = margin(30, 0, 0, 0), face="bold"),
          axis.title.y = element_text(size=14, margin = margin(0, 10, 0, 0), face="bold"),
          axis.title.x = element_text(size=14, face="bold"),
          axis.text.y = element_text(size=12, color = "black"),
          axis.text.x = element_text(size=12, color = "black", face="bold"),
          axis.ticks = element_line(size=1, color = "black"),
          legend.text = element_text(size = 12, colour = "black"),
          legend.title = element_text(size = 14, colour = "black", face="bold"),
          legend.key = element_blank(),
          legend.key.size = unit(1, "cm"),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          panel.border = element_blank(),
          panel.background = element_blank(),
          axis.line.x = element_line(colour = 'black', size = 1),
          axis.line.y = element_line(colour = 'black', size = 1))

  if (saveGraph == T) {
    pdf(paste("output/PostPuberty/", fileName, sep = ""), 9, 6)
    dominanceBar
    dev.off()
  } else {
    dominanceBar
  }
}

interactionGraphPre <- function(data, yVar, variableName, graphTitle, yMax, nonParametric = FALSE, fileName, saveGraph = F, showLegend = T) {
  interactionBar <- ggplot(data, aes_string(x = "Sex", y = yVar, fill = "Treatment"))
  if (nonParametric == FALSE) {
    interactionBar <- interactionBar + stat_summary(fun.y = "mean", geom = "bar", position = "dodge") +
      stat_summary(fun.data = "mean_se", geom = "errorbar", position = position_dodge(width = 0.9), width = 0.2, size  = 1)
  } else {
    interactionBar <- interactionBar + stat_summary(fun.y = "median", geom = "bar", position = "dodge")
  }

  interactionBar <- interactionBar + coord_cartesian(ylim = c(0, yMax)) +
    scale_y_continuous(expand = c(0, 0)) +

    labs (title = graphTitle, x = NULL, y = variableName) +
    scale_x_discrete(limits = c("M", "F"), breaks=c("M", "F"), labels=c("Male", "Female")) +
    scale_fill_manual("Prenatal Treatment", values = c("light gray", "black"), labels = c("Sesame Oil", "Testosterone")) +

    theme_bw() +
    theme(plot.title = element_text(colour="black", size=18, margin = margin(60, 0, 0, 0), face="bold"),
          axis.title.y = element_text(size=18, margin = margin(0, 10, 0, 0), face="bold"),
          axis.title.x = element_text(size=18, face="bold"),
          axis.text.y = element_text(size = 14, color = "black", face = "bold"),
          axis.text.x = element_text(size=18, color = "black", face="bold"),
          axis.ticks = element_line(size=1, color = "black"),
          legend.text = element_text(size = 14, colour = "black", face = "bold"),
          legend.title = element_text(size = 18, colour = "black", face = "bold"),
          legend.key = element_blank(),
          legend.key.size = unit(1.5,"cm"),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          panel.border = element_blank(),
          panel.background = element_blank(),
          axis.line.x = element_line(colour = 'black', size = 1),
          axis.line.y = element_line(colour = 'black', size = 1))

  if (showLegend == F) {
    interactionBar <- interactionBar + theme(legend.position = "none")
  }

  if (saveGraph == T) {
    pdf(paste("output/PrePuberty/", fileName, sep = ""), 9, 6)
    interactionBar
    dev.off()
  } else {
    interactionBar
  }

}

dominanceGraphPre <- function(data, yVar, variableName, graphTitle, yMin, yMax, nonParametric = FALSE, showLegend = T, fileName, saveGraph = F) {
  dominanceBar <- ggplot(data, aes_string(x = "Sex", y = yVar, fill = "Treatment"))
  if (nonParametric == FALSE) {
    dominanceBar <- dominanceBar + stat_summary(fun.y = "mean", geom = "bar", position = "dodge") +
    stat_summary(fun.data = "mean_se", geom = "errorbar", position = position_dodge(width = 0.9), width = 0.2, size  = 1)
  } else {
    dominanceBar <- dominanceBar + stat_summary(fun.y = "median", geom = "bar", position = "dodge")
  }

  dominanceBar <- dominanceBar + coord_cartesian(ylim = c(yMin, yMax)) +
    scale_y_continuous(expand = c(0, 0)) +

    labs (title = graphTitle, x = NULL, y = variableName) +
    scale_x_discrete(limits = c("M", "F"), breaks=c("M", "F"), labels=c("Male", "Female")) +
    scale_fill_manual("Prenatal Treatment", values = c("light gray", "black"), labels = c("Sesame Oil", "Testosterone")) +

    theme_bw() +
    theme(plot.title = element_text(colour="black", size=18, margin = margin(60, 0, 0, 0), face="bold"),
          axis.title.y = element_text(size=18, margin = margin(0, 10, 0, 0), face="bold"),
          axis.title.x = element_text(size=18, face="bold"),
          axis.text.y = element_text(size = 14, color = "black", face = "bold"),
          axis.text.x = element_text(size=18, color = "black", face="bold"),
          axis.ticks = element_line(size=1, color = "black"),
          legend.text = element_text(size = 14, colour = "black", face = "bold"),
          legend.title = element_text(size = 18, colour = "black", face = "bold"),
          legend.key = element_blank(),
          legend.key.size = unit(1.5,"cm"),
          panel.grid.major = element_blank(),
          panel.grid.minor = element_blank(),
          panel.border = element_blank(),
          panel.background = element_blank(),
          axis.line.x = element_line(colour = 'black', size = 1),
          axis.line.y = element_line(colour = 'black', size = 1))

  if (showLegend == F) {
    dominanceBar <- dominanceBar + theme(legend.position = "none")
  }

  if (saveGraph == T) {
    pdf(paste("output/PrePuberty/", fileName, sep = ""), 9, 6)
    dominanceBar
    dev.off()
  } else {
    dominanceBar
  }
}

getLegend <- function(myPlot){
  tmp <- ggplot_gtable(ggplot_build(myPlot))
  leg <- which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
  legend <- tmp$grobs[[leg]]
  return(legend)
}


library(tidyverse)
library(data.table)
library(xgboost)

library(USAboundaries)
library(fuzzyjoin)
library(mgcv)
library(ggrepel)
library(cowplot)
library(sf)
library(RColorBrewer)
library(viridisLite)
library(scales)

##### pull in case data ####
johns_hopkins_cases = read_csv('https://raw.githubusercontent.com/CSSEGISandData/COVID-19/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_confirmed_global.csv') %>%
  pivot_longer(cols = matches('^([0-9])'), names_to = 'date', values_to = 'cases') %>%
  filter(
    # `Country/Region` %in% c('Korea, South', 'Italy', 'Germany')
  ) %>%
  mutate(
    date = as.Date(date, format = '%m/%d/%y'),
    is_country = is.na(`Province/State`)
  ) %>%
  rename(
    country_name = `Country/Region`
  ) %>%
  mutate(
    country_name = recode(country_name, `US` = 'United States'),
    entity_name = country_name
  )

#### get weather data from NOAA ####
# https://www1.ncdc.noaa.gov/pub/data/ghcn/daily/by_year
# https://www1.ncdc.noaa.gov/pub/data/ghcn/daily/by_year/readme.txt
# https://www1.ncdc.noaa.gov/pub/data/ghcn/daily/readme.txt

setwd("~/../Downloads")
download.file('https://www1.ncdc.noaa.gov/pub/data/ghcn/daily/by_year/2020.csv.gz', destfile = '2020_weather.csv.gz')

weather_2020 = fread('2020_weather.csv.gz') %>% setorder(V1, V2, V3)
names(weather_2020) = c('station', 'date', 'type', 'value', 'm_flag', 'q_flag', 's_flag', 'time')
weather_2020$date_upd = as.Date(weather_2020$date %>% as.character(), format = '%Y%m%d')
weather_2020$country = str_extract(weather_2020$station, '^([A-Z]{2})')

weather_country_list = read_table("ghcnd-countries.txt", col_names = F)
names(weather_country_list) = c('country', 'country_name')

weather_2020_fin = left_join(weather_2020, weather_country_list) %>% data.table()
setkey(weather_2020_fin, country_name, type)
cases_weather_country_list = intersect(weather_country_list$country_name, johns_hopkins_cases$country_name)

weather_sub = weather_2020_fin[J(cases_weather_country_list), ]
rm(weather_2020)
gc()

# ID            1-11   Character
# LATITUDE     13-20   Real
# LONGITUDE    22-30   Real
# ELEVATION    32-37   Real
# STATE        39-40   Character
# NAME         42-71   Character
# GSN FLAG     73-75   Character
# HCN/CRN FLAG 77-79   Character
# WMO ID       81-85   Character

stations = read_table('ghcnd-stations.txt', col_names = F)
names(stations) = c('station', 'latitude', 'longitude', 'ELEVATION', 'STATE', 'NAME', 'GSN FLAG', 'WMO ID')
stations$country = str_extract(stations$station, '^([A-Z]{2})')
# check for missing values
map_dbl(stations, function(x){sum(is.na(x))})

#### get locations of us cities and get closest weather stations ####
setwd("~/Public_Policy/Projects/COVID-19")

us_cities_shp = us_cities()

us_cities_coord = st_coordinates(us_cities_shp) %>% as.data.frame()
names(us_cities_coord) = c('longitude', 'latitude')
us_cities_shp_fin = bind_cols(us_cities_shp, us_cities_coord)

us_cities_shp_fin$geometry = NULL
selected_cities = filter(us_cities_shp_fin,
                         str_detect(city,
                  '^(New York City)|^(San Francisco)|(Seattle)|(Los Angeles)'))

# join on stations within 50 miles of cities
city_stations = geo_left_join(selected_cities, stations %>% filter(country == 'US'), max_dist = 50)

# average temp across entities
weather_sub_station_meta = left_join(weather_sub, stations) %>% select(-date) %>%
  mutate(
    temp_f = ifelse(type == 'TAVG', (value/10) * 9/5 + 32, value)
  ) %>%
  left_join(
    city_stations %>% select(station, city)
  ) %>%
  mutate(
    entity_name = ifelse(is.na(city), country_name, city)
  )

avg_temp_sub = weather_sub_station_meta %>% filter(type == 'TAVG')
avg_temp_by_entity = group_by(avg_temp_sub, country, country_name, entity_name, date_upd) %>%
  summarize(
    mean_temp = mean(temp_f)
  )

# use counties as proxies for city cases #
nyt_county_data = read_csv('https://raw.githubusercontent.com/nytimes/covid-19-data/master/us-counties.csv') %>%
  arrange(county, state, fips, date) %>%
  select(-deaths)

selected_counties_cases = filter(nyt_county_data,
       (county == 'King' & state == 'Washington') |
       county %in% c('San Francisco', 'Los Angeles', 'New York City')) %>%
  mutate(
    entity_name = recode(county, `King` = 'Seattle')
  )

setwd("~/Public_Policy/Projects/COVID-19")

### read oxford stringency data ###

oxford_stringency_index = read_csv("https://github.com/OxCGRT/covid-policy-tracker/raw/master/data/OxCGRT_latest.csv") %>%
  rename(
    entity_name = CountryName
  ) %>%
  mutate(
    entity_name = recode(entity_name, `South Korea` = 'Korea, South'),
    date = as.Date(Date %>% as.character(), format = '%Y%m%d')
  )

### disaggregate jh data -- some countries have provinces/states included ###
johns_hopkins_cases_dt = arrange(johns_hopkins_cases, date, country_name) %>% data.table()
johns_hopkins_cases_dt_agg = johns_hopkins_cases_dt[, {
  the_df = data.frame(date, cases) %>%
    group_by(date) %>%
    summarize(
      the_cases = sum(cases)
    )
  list(
    date = the_df$date,
    cases = the_df$the_cases
  )
}, by = list(country_name, entity_name)] %>%
  mutate(
    entity_name = recode(entity_name, `Taiwan*` = 'Taiwan')
  )

#### read in mobility data and join everything else on ####
mobility_dataset_df = tibble(
  dsn = list.files('data', pattern = 'applemobilitytrends', full.names = T),
  date = str_extract(dsn, '[0-9]{4}-[0-9]{2}-[0-9]{2}') %>% as.Date()
) %>%
  arrange(date) %>%
  tail(1)

apple_mobility_dat = read_csv(mobility_dataset_df$dsn) %>%
  pivot_longer(cols = matches('^([0-9+]{4})'), names_to = 'date') %>%
  mutate(
    date = as.Date(date),
    week_day = lubridate::wday(date),
    weekend_ind = ifelse(week_day %in% c(7, 1), 'Weekend', "Week Day"),
    entity_name = recode(region,
                         `UK` = 'United Kingdom',
                         `San Francisco - Bay Area` = 'San Francisco',
                         `Republic of Korea` = 'Korea, South')
  ) %>%
  left_join(
    avg_temp_by_entity, by = c('entity_name', 'date' = 'date_upd')
  ) %>%
  left_join(
    selected_counties_cases
  ) %>%
  left_join(
    johns_hopkins_cases_dt_agg , by = c('date', 'entity_name')
  ) %>%
  left_join(
    oxford_stringency_index
  ) %>%
  mutate(
    cases = coalesce(cases.x, cases.y, as.numeric(NA))
  ) %>%
  select(-cases.x, -cases.y) %>%
  mutate(
    entity_name = recode(entity_name, `Korea, South` = 'South Korea')
  )
write.csv(apple_mobility_dat, 'data/mobility_stringency.csv', row.names = F)

#### subset to walking data -- this seems more comparable across countries since driving/transit trades off ###
walking_data = filter(apple_mobility_dat, transportation_type == 'walking', !is.na(mean_temp), !is.na(value))

# make sure no entities are counted multiple times #
obs_by_entity = group_by(walking_data, entity_name) %>%
  summarize(
    obs = n(),
    n_dates = n_distinct(date)
  )
stopifnot(!any(obs_by_entity$obs != obs_by_entity$n_dates))

### fit a model for baseline period to capture the seasonality of mobility ####
baseline_period_end = as.Date('2020-02-15')

case_ten_dates = group_by(walking_data, entity_name) %>%
  summarize(
    date_case_ten = min(date[cases >= 10], na.rm = T)
  ) %>%
  arrange(date_case_ten)

baseline_data = left_join(walking_data, case_ten_dates) %>%
  filter(
    date <= pmin(date_case_ten, baseline_period_end)
  )
the_formula = as.formula('value ~ as.numeric(date) * region * mean_temp + week_day + I(week_day^2)  + weekend_ind')
pre_covid_walking_model = lm(the_formula, data = baseline_data)
the_anova = anova(pre_covid_walking_model)
the_anova$rsquared = the_anova$`Sum Sq`/sum(the_anova$`Sum Sq`)
the_anova$predictor = row.names(the_anova)
the_anova = arrange(the_anova, -rsquared)

train_dat = baseline_data
model_mat = model.matrix(the_formula, data = train_dat)
pre_covid_walking_model_boost = xgboost(data = model_mat, label = train_dat$value, nrounds = 100)

# walking_dat_sub = select(walking_data, region, date, value, mean_temp, week_day, weekend_ind)
# head(walking_dat_sub)
# na_rm_this = na.omit(walking_dat_sub)
# anti_join(walking_dat_sub, na_rm_this)
walking_data_mat = model.matrix(the_formula, data = walking_data)

walking_data_fin = mutate(
  walking_data,
  predicted_walking_baseline = predict(pre_covid_walking_model_boost, newdata = walking_data_mat),
  percent_of_predicted = value / predicted_walking_baseline
) %>%
  data.table() %>%
  arrange(region, date)

smoothed_pct_of_predicted = map(unique(walking_data_fin$entity_name), function(the_entity){
  region_sub = filter(walking_data_fin, entity_name == the_entity) %>%
    mutate(
      lag_cases = lag(cases, 1),
      new_cases = cases - lag_cases
    )

  smoother = loess(percent_of_predicted ~ as.numeric(date), data = region_sub, span = 0.2)
  region_sub$smoothed_percent_of_predicted = predict(smoother)
  return(region_sub)
}) %>%
  bind_rows() %>%
  rename(
    mobility = value
  )

write.csv(smoothed_pct_of_predicted, 'data/covid_cases_mobility_stringency.csv', row.names = F)

stats_by_region = group_by(smoothed_pct_of_predicted, entity_name) %>%
  summarize(
    obs = n(),
    date_80 = min(date[smoothed_percent_of_predicted <= 0.8]),
    date_80_val = smoothed_percent_of_predicted[date == date_80],
    date_80_cases = cases[date == date_80],
    last_date = max(date),
    last_val = smoothed_percent_of_predicted[date == last_date]
  )

# ggplot(smoothed_pct_of_predicted, aes(StringencyIndex, percent_of_predicted)) +
#   geom_point() +
#   stat_smooth()

avg_by_date = group_by(smoothed_pct_of_predicted, date) %>%
  summarize(
    avg_mobility = mean(mobility, na.rm = T),
    avg_smoothed_percent_of_predicted = mean(smoothed_percent_of_predicted, na.rm = T)
  )
ggplot(avg_by_date, aes(avg_smoothed_percent_of_predicted, avg_mobility)) +
  geom_point()

the_pal = brewer.pal(3, 'Set1')

last_vals_by_country = group_by(smoothed_pct_of_predicted, entity_name) %>%
  summarize(
    last_date = max(date),
    last_cases = cases[date == last_date],
    last_mobility = mobility[date == last_date],
    last_predicted_walking_baseline = predicted_walking_baseline[date == last_date],
    last_smoothed_percent_of_predicted = smoothed_percent_of_predicted[date == last_date]
  )

original_plot = ggplot(smoothed_pct_of_predicted, aes(date, mobility, group = entity_name, colour = StringencyIndex)) +
  theme_bw() +
  scale_colour_viridis_c(name = 'Stringency Index', option = 'C') +
  geom_line(size = 0.25, alpha = 0.6) +
  # geom_line(data = avg_by_date, aes(y = avg_mobility, group = NULL), colour = the_pal[1], size = 1) +
  labs(
    x = '', y = 'Mobility Index',
    title = 'Reported Mobility Index (Walking)'
  ) +
    theme(
      legend.position = c(0.85, 0.8),
      # plot.background = element_rect(fill = 'black'),
      # panel.background = element_rect(fill = 'black'),
      # panel.grid = element_line(colour = 'white'),
      panel.grid.minor = element_blank()) +
    scale_y_continuous(breaks = seq(0, 350, by = 50))

selected_entities = c('United States', 'Italy', 'South Korea', 'Germany')
linetype_data = data.frame(vals = c('Actual', 'Predicted'), date = rep(baseline_period_end, 2), mobility = rep(100, 2))

original_with_predictions =
  smoothed_pct_of_predicted %>% filter(entity_name %in% selected_entities) %>%
  ggplot(aes(date, colour = entity_name)) +
  theme_bw() +
  scale_linetype(name = '') +
  geom_line(aes(y = mobility), size = 0.5, show.legend = F) +
  scale_colour_hue(guide = F) +
  geom_line(data = linetype_data, aes(x = date, y = mobility, linetype = vals, colour = NULL), alpha = 0) +
  guides(linetype = guide_legend(override.aes = list(alpha = 1))) +
  labs(x = '', y = 'Mobility Index',
       title = 'Actual Mobility Index vs. Predicted Without COVID-19, Selected Countries'
       ) +
  theme(
    legend.position = c(0.1, 0.25),
    # plot.background = element_rect(fill = 'black'),
    # panel.background = element_rect(fill = 'black'),
    # panel.grid = element_line(colour = 'white'),
    panel.grid.minor = element_blank()) +
  scale_y_continuous(breaks = seq(0, 200, by = 50)) +
  geom_line(aes(y = predicted_walking_baseline), linetype = 'dashed', show.legend = T)
  # geom_text_repel(data = last_vals_by_country %>% filter(entity_name %in% selected_entities),
                  # aes(x = last_date + 4, y = last_predicted_walking_baseline, label = paste(entity_name, '(Predicted)'))) +
  # geom_text_repel(data = last_vals_by_country %>% filter(entity_name %in% selected_entities),
                  # aes(x = last_date + 4, y = last_mobility, label = paste(entity_name, '(Actual)')))

smoothed_plot = ggplot(smoothed_pct_of_predicted, aes(date, smoothed_percent_of_predicted, group = entity_name, colour = StringencyIndex)) +
  geom_line(size = 0.25, show.legend = F, alpha = 0.6) +
  # geom_line(data = avg_by_date, aes(y = avg_smoothed_percent_of_predicted, group = NULL), colour = the_pal[1], size = 1) +
  theme_bw() +
  scale_colour_viridis_c(name = 'Stringency Index', option = 'C') +
  labs(
    x = '', y = 'Percent of Predicted Baseline Mobility Index',
    title = 'Smoothed Mobility Index, All Available Countries'
  ) +

  theme(
    panel.grid.minor = element_blank()
        # plot.background = element_rect(fill = 'black'),
        # panel.background = element_rect(fill = 'black'),
        # panel.grid = element_line(colour = 'white')
        ) +
  scale_y_continuous(labels = percent, limits = c(0, 1.3), breaks = seq(0, 1.3, by = 0.2))

smoothed_sub =
  smoothed_pct_of_predicted %>% filter(entity_name %in% selected_entities) %>%
  ggplot(aes(date, colour = entity_name)) +
  theme_bw() +
  geom_line(aes(y = smoothed_percent_of_predicted), size = 0.75, show.legend = T) +
  labs(
    x = '', y = 'Percent of Predicted Baseline Mobility Index',
    title = 'Smoothed Mobility Index, Selected Countries'
  ) +
  theme(
    legend.position = c(0.1, 0.25),
    # plot.background = element_rect(fill = 'black'),
    # panel.background = element_rect(fill = 'black'),
    # panel.grid = element_line(colour = 'white'),
    panel.grid.minor = element_blank()) +
  scale_colour_hue(name = '') +
  guides(colour = guide_legend(override.aes = list(size = 1))) +
  scale_y_continuous(labels = percent, limits = c(0, 1.3), breaks = seq(0, 1.3, by = 0.2))
  # geom_text_repel(data = last_vals_by_country %>% filter(entity_name %in% selected_entities),
  #                 aes(x = last_date , y =
  #                       last_smoothed_percent_of_predicted, label = paste(entity_name, percent(last_smoothed_percent_of_predicted, accuracy = 1))),
  #                 show.legend = F, hjust=1, nudge_x = 4,
  #                 direction = 'y', segment.colour = 'gray', size = 2)

comb_plot = plot_grid(original_plot, original_with_predictions, smoothed_plot, smoothed_sub)
comb_plot_sub = add_sub(comb_plot, "Chart: Taylor G. White\nData: Apple (Mobility), Oxford (Stringency), Johns Hopkins CSSE, NYT, NOAA",
               x = 0.05, y = 0.35, hjust = 0, size = 10, fontface = 'italic')

title <- ggdraw() +
  draw_label(
    "The Month the Earth Stood Still: Reduction in Movement Associated with COVID-19 and Government Restrictions",
    x = 0,
    hjust = 0
  ) +
  theme(
    # add margin on the left of the drawing canvas,
    # so title is aligned with left edge of first plot
    plot.margin = margin(0, 0, 0, 7)
  )
final_plot = plot_grid(
  title, comb_plot_sub,
  ncol = 1,
  # rel_heights values control vertical title margins
  rel_heights = c(0.05, 1)
)

# final_plot
save_plot('output/smoothed_mobility_index_data.png', plot = final_plot,
          base_height = 12,
           base_width = 14, units = 'in', dpi = 100)

# smoothed_pct_of_predicted %>%
#   ggplot(aes(date, smoothed_percent_of_predicted, colour = entity_name)) +
#   geom_vline(data = filter(stats_by_region, region == 'New York City'), aes(xintercept = date_80, colour = region), show.legend = F) +
#   geom_line(size = 0.25, show.legend = F) +
#   geom_point(data = stats_by_region, aes(date_80, date_80_val, size = date_80_cases)) +
#   geom_point(aes(size = new_cases)) +
#   geom_point(data = stats_by_region, aes(date_80, date_80_val), size = 2, shape = 6) +
#   geom_segment(data = stats_by_region, aes(x = date_80, y = 0, yend = date_80_val, xend = date_80, colour = region), size = 0.25) +
#   geom_text_repel(data = stats_by_region, aes(date_80, date_80_val, label = paste0(format(date_80, '%b %d'), '\n', comma(date_80_cases, accuracy = 1))), size = 3.5, fontface = 'bold', min.segment.length = 0) +
#   scale_size(name = 'Daily New Cases',range = c(0, 3), label = comma) +
#   scale_y_continuous(labels = percent, breaks = seq(0, 1.5, by = 0.1)) +
#
#   geom_text_repel(data = stats_by_region, aes(last_date + 3, last_val, label = paste0(entity_name, ', ', percent(last_val, accuracy = 1))), size = 3) +
#   labs(
#     x = '', y = 'Percent of Baseline Activity',
#     title = 'Rate of Mobility Through the COVID-19 Epidemic',
#     subtitle = "Mobility data shows level of requests for walking directions from Apple maps for selected regions, which measures both potential COVID-19 exposure and economic activity. Baseline data is estimated using a boosted tree algorithm and then smoothed with LOESS. Triangles show date at which activity falls below 80% of baseline along with total cases." %>%
#       str_wrap(150),
#     caption = 'Chart: Taylor G. White\nData: Apple, Johns Hopkins CSSE, New York Times, NOAA'
#   ) +
#   theme_bw() +
#   theme(
#     panel.grid.minor = element_blank(),
#     legend.position = 'bottom',
#     plot.subtitle = element_text(size = 11, face = 'italic'),
#     plot.caption = element_text(size = 10, face = 'italic', hjust = 0)
#   ) +
#   scale_colour_hue(guide = F, name = 'Region') +
#   scale_x_date(date_breaks = '1 week', limits = c(as.Date('2020-01-20'), max(smoothed_pct_of_predicted$date) + 4), date_labels = '%b %d')
# ggsave('output/walking_mobility_data_by_date.png', height = 10, width = 12, units = 'in', dpi = 800)
#
#


define_factor_line <- function (df, is_weak = FALSE){
  function (factor_name) {
    variables <- df[[factor_name]]

    if (is_weak) {
      var_res = mapply(function (var_name, index){
        paste0("lambda", index, "*", var_name)
      }, tail(variables, -1), 2:length(variables))
      return(
        paste0(factor_name, " =~ 1*", variables[1], " + ", paste(var_res, collapse = " + "))
      )
    }
    paste0(factor_name, " =~ 1*", variables[1], " + ", paste(tail(variables, -1), collapse = " + "))
  }
}

define_factors <- function (factor_sets, is_weak = FALSE) {
  factor_names = colnames(factor_sets)
  define_factor_lineResult <- define_factor_line(factor_sets, is_weak)
  factor_lines = sapply(factor_names, define_factor_lineResult)
  paste(factor_lines, collapse = "\n")
}

define_intercept_line <- function (df,strong, unconstrained_parcel_indices = c()){
  function (factor_name) {
    if(!strong) {
      variables <- df[[factor_name]]
      tailResult <- sapply(tail(variables, -1), function(varr) {
        paste0(varr, " ~ 1")
      })
      return(
        paste0(variables[1], " ~ 1\n", paste(tailResult, collapse = "\n"))
      )
    }
    variables <- df[[factor_name]]
    tailResult <- mapply(function(varr,index) {
      if(index %in% unconstrained_parcel_indices) {
        return(paste0(varr, " ~ 1"))
      }
      paste0(varr, " ~ i", index, "*1")
    }, variables, 1:length(variables))
    paste(tailResult, collapse = "\n")
  }
}

define_intercepts <- function (factor_sets, strong = FALSE, unconstrained_parcel_indices = c()) {
  if(length(factor_sets[,1]) /2 < length(unconstrained_parcel_indices)) {
    stop("must constrain at least half of the parcels when defining a partial strong model")
  }
  factor_names = colnames(factor_sets)
  define_factor_lineResult <- define_intercept_line(factor_sets, strong, unconstrained_parcel_indices)
  factor_lines = sapply(factor_names, define_factor_lineResult)
  paste(factor_lines, collapse = "\n\n")
}

define_variance_line <- function (df, unique){
  function (factor_name) {
    if(!unique) {
      variables <- df[[factor_name]]
      tailResult <- sapply(variables, function(varr) {
        paste0(varr, " ~~ ", varr)
      })
      return(paste(tailResult, collapse = "\n"))
    }
    variables <- df[[factor_name]]
    tailResult <- mapply(function(varr, index) {
      paste0(varr, " ~~ u", index, "*", varr)
    }, variables, 1:length(variables))
    return(paste(tailResult, collapse = "\n"))
  }
}
define_variances <- function (factor_sets, unique = FALSE) {
  factor_names = colnames(factor_sets)
  define_factor_lineResult <- define_variance_line(factor_sets, unique)
  factor_lines = sapply(factor_names, define_factor_lineResult)
  paste(factor_lines, collapse = "\n\n")
}

define_covariance_line <- function (column_pairs, df){
  function (column_index) {
    col_names <- column_pairs[,column_index]
    v1 = df[[col_names[1]]]
    v2 = df[[col_names[2]]]
    tailResult <- mapply(function(varr, varr2) {
      paste0(varr, " ~~ ", varr2)
    },v1,v2 )
    paste(tailResult, collapse = "\n")
  }
}

define_covariances <- function (factor_sets) {
  factor_names = colnames(factor_sets)
  # get our column pairs
  column_pairs = combn(factor_names,2)
  indexes = 1:length(column_pairs[1,])
  define_factor_lineResult <- define_covariance_line(column_pairs, factor_sets)
  factor_lines = sapply(indexes, define_factor_lineResult)
  paste(factor_lines, collapse = "\n\n")
}

define_latent_variable_means <- function (factor_sets, constrain_all_means_to_zero) {
  factor_names = colnames(factor_sets)
  if(constrain_all_means_to_zero) {
    return(
      paste0(
        factor_names[1], " ~ 0*1\n",
        paste(tail(factor_names, -1), collapse = " ~ 0*1\n"),
        " ~ 0*1"
      )
    )
  }
  return (
    paste0(
      factor_names[1], " ~ 0*1\n",
      paste(tail(factor_names, -1), collapse = " ~ 1\n"),
      " ~ 1"
    )
  )
}

define_latent_variable_variances <- function (factor_sets) {
  factor_names = colnames(factor_sets)
  tailResult <- sapply(tail(factor_names, -1), function(p) { paste(p," ~~ ",p, sep = "")})
  paste0(
    factor_names[1],
    " ~~ ", factor_names[1],
    "\n",
    paste(tailResult, collapse = "\n")
  )
}

define_latent_variable_covariance_line <- function (column_pairs, constrain_to_zero){
  function (column_index) {
    col_names <- column_pairs[,column_index]
    if(constrain_to_zero) {
      return(paste0(col_names[1], " ~~ 0*", col_names[2]))
    }
    paste0(col_names[1], " ~~ ", col_names[2])
  }
}

define_latent_variable_covariances <- function (factor_sets, constrain_to_zero = FALSE) {
  factor_names = colnames(factor_sets)
  column_pairs = combn(factor_names,2)
  indexes = 1:length(column_pairs[1,])
  define_covariance_for_factor_pair <- define_latent_variable_covariance_line(column_pairs, constrain_to_zero)
  factor_lines = sapply(indexes, define_covariance_for_factor_pair)
  paste(factor_lines, collapse = "\n")
}