#################
## Parameters and functions
#################
require(data.table)
require(RColorBrewer)
require(ggplot2) # for Fig. 1 hexplots. also requires package hexbin
require(gridExtra) # for Fig. 1

# takes values 0-1 and turns them into colors. takes ColorBrewer pallete as argument.
colramp <- function(x, pal='RdBu', alpha=255){
	cr <- colorRamp(brewer.pal(9, pal))
	i <- !is.na(x)
	cols <- rgb(cr(x[i]), alpha=rep(alpha, sum(i)), maxColorValue=255)
	out <- rep(NA, length(x))
	out[i] <- cols
	return(out)
}

cols <- c('#a6cee333', '#1f78b433') # light blue, blue, partially transparent: for alternating LGs
cols2 <- c('#a6cee3', '#1f78b4') # light blue, blue: for alternating LGs outliers
colout <- '#b2182b' # red, part of RdGy colorbrewer, for outlier loci
#cols3 <- c('#F7941E', '#F2592A', '#BF1E2D', '#FFDE17', '#BEA512') # Lof07, Lof11, Lof14, Can40, CanMod (reds, yellows, from Bastiaan)
cols3 <- c('#e7d4e8', '#af8dc3', '#762a83', '#d9f0d3', '#1b7837') # Lof07, Lof11, Lof14, Can40, CanMod (PRGn colorbrewer)

# takes p-values and returns a color
# uses colout for p<thresh, otherwise a colorbrewer palette
# alternates color by chromosome
lgcolsramp <- function(x, lg = 1, thresh = 0.05){
  if(lg == 1) rmp <- colorRamp(brewer.pal(11, 'BrBG')[5:1])
  if(lg == 2) rmp <- colorRamp(brewer.pal(11, 'BrBG')[7:11])
  if(!(lg %in% c(1,2))) error('lg must be 1 or 2')
  out <- rep(colout, length(x))
  out[x <= thresh] <- colout
  rgbs <- rmp(-log10(x[x > thresh])/-log10(thresh))
  out[x > thresh] <- rgb(rgbs[,1], rgbs[,2], rgbs[,3], maxColorValue = 256)
  return(out)
}

# draw a vertical colorbar at xposmin to xposmax, from yposmin to yposmax
# cols in hex, x are values to go with each col
color.bar <- function(cols, x, axis = TRUE, cex = 1, nticks=11,
                      xposmin, xposmax, yposmin, yposmax, title = '', titley = yposmax) {
  scale = (length(cols)-1)/(yposmax - yposmin)

  for (i in 1:(length(cols)-1)) {
    y = (i-1)/scale + yposmin
    rect(xposmin,y,xposmax,y+1/scale, col=cols[i], border=NA)
  }

  if(axis){
    line(c(xposmax, xposmax), c(yposmin, yposmax))
    line(c(xposmax, xposmax + (xposmax-xposmin)/2), c(yposmin, yposmin))
    line(c(xposmax, xposmax + (xposmax-xposmin)/2), c(yposmax, yposmax))
    at = seq(yposmin, yposmax, length.out = nticks)
    labs <- signif(x[round(seq(1, length(x), length.out = nticks))], 2)
    text(rep(xposmax + (xposmax-xposmin)/2, nticks), y = at, labels = labs, adj = 0, cex = cex)
  }

 text(xposmin, titley, title, adj = 0.5, cex = cex)
}

#####################################################
# Part of Fig. 1 and S6: allele frequency change heatmap
#####################################################
colLof <- '#F26523' # matching Bastiaan
colCan <- '#BEA512'

# read in and merge data
datCan40 <- fread('data_31_01_20/Can_40_freq.mafs.gz')
datCan14 <- fread('data_31_01_20/Can_14_freq.mafs.gz')
datLof07 <- fread('data_31_01_20/Lof_07_freq.mafs.gz')
datLof11 <- fread('data_31_01_20/Lof_11_freq.mafs.gz')
datLof14 <- fread('data_31_01_20/Lof_14_freq.mafs.gz')

datCan <- merge(datCan40[, .(chromo, position, freq1 = knownEM)], datCan14[, .(chromo, position, freq2 = knownEM)])
datLof0711 <- merge(datLof07[, .(chromo, position, freq1 = knownEM)], datLof11[, .(chromo, position, freq2 = knownEM)])
datLof0714 <- merge(datLof07[, .(chromo, position, freq1 = knownEM)], datLof14[, .(chromo, position, freq2 = knownEM)])
datLof1114 <- merge(datLof11[, .(chromo, position, freq1 = knownEM)], datLof14[, .(chromo, position, freq2 = knownEM)])

# trim to unlinked loci
ulnkCan <- fread('analysis/ld.unlinked.Can.gatk.nodam.csv.gz')
ulnkLof <- fread('analysis/ld.unlinked.Lof.gatk.nodam.csv.gz')

datCan <- merge(datCan, ulnkCan[, .(chromo = CHROM, position = POS)])
datLof0711 <- merge(datLof0711, ulnkLof[, .(chromo = CHROM, position = POS)])
datLof0714 <- merge(datLof0714, ulnkLof[, .(chromo = CHROM, position = POS)])
datLof1114 <- merge(datLof1114, ulnkLof[, .(chromo = CHROM, position = POS)])

# r2 values
datCan[, cor(freq1, freq2)]^2
datLof0711[, cor(freq1, freq2)]^2
datLof0714[, cor(freq1, freq2)]^2
datLof1114[, cor(freq1, freq2)]^2

# plots
bks <- c(0, 1, 10, 100, 1000, 10000)
p1 <- ggplot(datCan, aes(freq1, freq2)) +
  geom_hex() +
  scale_fill_gradient(low="grey", high=colCan, trans = 'log', breaks = bks, labels = bks) +
  labs(title = 'Canada', x = 'Historical', y = 'Modern') +
  geom_abline(slope = 1, intercept = 0, linetype = 'dashed') +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        panel.background = element_blank(), axis.line = element_line(colour = "grey"),
        legend.position = 'none') +
  coord_cartesian(xlim =c(0, 1), ylim = c(0, 1))
p2 <- ggplot(datLof0711, aes(freq1, freq2)) +
  geom_hex() +
  scale_fill_gradient(low="grey", high=colLof, trans = 'log', breaks = bks, labels = bks) +
  labs(title = 'Norway 07-11', x = 'Historical', y = 'Modern') +
  geom_abline(slope = 1, intercept = 0, linetype = 'dashed') +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        panel.background = element_blank(), axis.line = element_line(colour = "grey"),
        legend.position = 'none') +
  coord_cartesian(xlim =c(0, 1), ylim = c(0, 1))
p3 <- ggplot(datLof0714, aes(freq1, freq2)) +
  geom_hex() +
  scale_fill_gradient(low="grey", high=colLof, trans = 'log', breaks = bks, labels = bks) +
  labs(title = 'Norway 07-14', x = 'Historical', y = 'Modern') +
  geom_abline(slope = 1, intercept = 0, linetype = 'dashed') +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        panel.background = element_blank(), axis.line = element_line(colour = "grey"),
        legend.position = 'none') +
  coord_cartesian(xlim =c(0, 1), ylim = c(0, 1))
p4 <- ggplot(datLof1114, aes(freq1, freq2)) +
  geom_hex() +
  scale_fill_gradient(low="grey", high=colLof, trans = 'log', breaks = bks, labels = bks) +
  labs(title = 'Norway 11-14', x = 'Historical', y = 'Modern') +
  geom_abline(slope = 1, intercept = 0, linetype = 'dashed') +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        panel.background = element_blank(), axis.line = element_line(colour = "grey"),
        legend.position = 'none') +
  coord_cartesian(xlim =c(0, 1), ylim = c(0, 1))

p5 <- grid.arrange(p1, p2, p3, p4, ncol = 4)

ggsave(p5, filename = 'figures/figure1.png', width = 7.5, height = 2)
ggsave(p5, filename = 'figures/figure1.pdf', width = 7.5, height = 2)

########################################
## Fig. 2 Manhattan plot FSTs by region
########################################

# read in data: sliding window fst from ANGSD (GATK nodam2 sites)
dat1 <- fread('analysis/Can_40.Can_14.gatk.slide', col.names = c('region', 'CHROM', 'midPos', 'nloci', 'fst'), skip = 1) # output by angsd_fst.sh. skip headers.
dat3 <- fread('analysis/Lof_07.Lof_14.gatk.slide', col.names = c('region', 'CHROM', 'midPos', 'nloci', 'fst'), skip = 1)

dat1[, pop := 'can']
dat3[, pop := 'lof0714']

dat <- rbind(dat1, dat3)

# trim to non-overlapping windows
dat <- dat[(midPos - 1) %% 25000 == 0, ]

# read in shared outlier regions
outl <- fread('tables/outlier_annotation.csv') # annotated list of outliers, from annotate_outliers.r
outl[, midPos := as.numeric(midPos) + 1] # add one to match dat centers
outl <- outl[region == 1 & test == '99th percentile across pops', .(CHROM, midPos, outl = 1)] # trim to just regions
dat <- merge(dat, outl, all.x = TRUE)
dat[is.na(outl), outl := 0]

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

setkey(dat, CHROM)
dat <- dat[chrmax[, .(CHROM = chr, start)], ]
dat[, posgen := midPos + start]
dat[,start := NULL]

# trim out Unplaced or windows with <2 loci
dat <- dat[!(CHROM %in% c('Unplaced')) & nloci > 1, ]
	nrow(dat)

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := cols[1]]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := cols[2]]

### set up plot
adjlet <- -0.14 # horizontal adjustment for subplot letter
cexlet <- 1
linelet <- 0.35
cexsc <- 1/5
xlims <- dat[, range(posgen, na.rm=TRUE)]
ylims <- c(0,0.4)

# quartz(height=8, width=6)
#png(height=3.3, width=6, units='in', res=300, file='figures/figure2.png')
pdf(height=3.3, width=6, file='figures/figure2.pdf')
par(mfrow = c(2,1), las=1, mai=c(0.3, 0.6, 0.25, 0.1), omi = c(0.3, 0, 0, 0))

dat[pop == 'can', plot(posgen, fst, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims,
                       ylim = ylims, xlab = '', ylab = expression(F[ST]), bty = 'l',
                       cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
dat[pop == 'can' & outl == 1, points(posgen, fst, col = colout, cex=log(nloci)*cexsc)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8, mgp = c(1, 0, 0)) # plot x-axis
mtext(side=3, 'A. Canada 1940-2013', adj=adjlet, line=linelet, cex=cexlet)

legend('topleft', legend = c(2, 5, 10, 100), pch = 1, pt.cex = log(c(2,5,10,100))*cexsc, title = '# of SNPs', bty = 'n', cex = 0.7)

dat[pop == 'lof0714', plot(posgen, fst, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims,
                           ylim = ylims, xlab = '', ylab = expression(F[ST]), bty = 'l',
                           cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
dat[pop == 'lof0714' & outl == 1, points(posgen, fst, col = colout, cex=log(nloci)*cexsc)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8, mgp = c(1, 0, 0))
mtext(side=3, 'B. Norway 1907-2014', adj=adjlet, line=linelet, cex=cexlet)
mtext(side=1, 'Linkage group', outer = TRUE)

dev.off()

###################################################
## Fig. S5: read balance at heterozygote genotypes
## (check for reference bias)
###################################################
# read in
can13 <- fread('analysis/allele_balance_popCan13.csv.gz', header= FALSE)
can40 <- fread('analysis/allele_balance_popCan40.csv.gz', header= FALSE)
lof07 <- fread('analysis/allele_balance_popLof07.csv.gz', header= FALSE)
lof11 <- fread('analysis/allele_balance_popLof11.csv.gz', header= FALSE)
lof14 <- fread('analysis/allele_balance_popLof14.csv.gz', header= FALSE)

# plot
linex <- 2
cexx <- 0.6

png(height=4, width=6.5, units='in', res=300, file='figures/figureS5.png')
par(mfrow = c(2,3), las=1, mai=c(0.5, 0.7, 0.25, 0.1), mgp = c(4,0.8,0))

hist(lof07$V1, main = 'Norway 1907')
mtext('REF allele read proportion', 1, cex = cexx, line = linex)

hist(lof11$V1, main = 'Norway 2011')
mtext('REF allele read proportion', 1, cex = cexx, line = linex)

hist(lof14$V1, main = 'Norway 2014')
mtext('REF allele read proportion', 1, cex = cexx, line = linex)

hist(can40$V1, main = 'Canada 1940')
mtext('REF allele read proportion', 1, cex = cexx, line = linex)

hist(can13$V1, main = 'Canada 2013')
mtext('REF allele read proportion', 1, cex = cexx, line = linex)

dev.off()

####################################################
## Fig. S7 Manhattan plot FSTs 1907-2011 and 2011-2014
####################################################

# read in data: sliding window fst from ANGSD (GATK nodam2 sites)
dat2 <- fread('analysis/Lof_07.Lof_11.gatk.slide', col.names = c('region', 'CHROM', 'midPos', 'nloci', 'fst'), skip = 1)
dat4 <- fread('analysis/Lof_11.Lof_14.gatk.slide', col.names = c('region', 'CHROM', 'midPos', 'nloci', 'fst'), skip = 1)

dat2[, pop := 'lof0711']
dat4[, pop := 'lof1114']

dat <- rbind(dat2, dat4)

# trim to non-overlapping windows
dat <- dat[(midPos - 1) %% 25000 == 0, ]

# read in outlier regions
outl <- fread('tables/outlier_annotation.csv') # annotated list of outliers, from annotate_outliers.r
outl[, midPos := as.numeric(midPos) + 1] # add one to match dat centers
outl <- outl[region == 1, .(CHROM, midPos, outl = 1)] # trim to just regions
dat <- merge(dat, outl, all.x = TRUE)
dat[is.na(outl), outl := 0]

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

setkey(dat, CHROM)
dat <- dat[chrmax[, .(CHROM = chr, start)], ]
dat[, posgen := midPos + start]
dat[,start := NULL]

# trim out Unplaced or windows with <2 loci
dat <- dat[!(CHROM %in% c('Unplaced')) & nloci > 1, ]
nrow(dat)

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := cols[1]]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := cols[2]]

### set up plot
adjlet <- -0.14 # horizontal adjustment for subplot letter
cexlet <- 0.7
linelet <- 0.2
cexsc <- 1/5

# quartz(height=8, width=6)
png(height=3, width=6, units='in', res=300, file='figures/figureS7.png')
par(mfrow = c(2,1), las=1, mai=c(0.3, 0.6, 0.3, 0.1))

xlims <- dat[, range(posgen, na.rm=TRUE)]

dat[pop == 'lof0711', plot(posgen, fst, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims, xlab = '', ylab = expression(F[ST]), bty = 'l', cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
dat[pop == 'lof0711' & outl == 1, points(posgen, fst, col = colout, cex=log(nloci)*cexsc)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8, mgp = c(1, 0, 0))
mtext(side=3, 'A. Norway 1907-2011', adj=adjlet, line=linelet, cex=cexlet)

legend('topleft', legend = c(2, 5, 10, 100), pch = 1, pt.cex = log(c(2,5,10,100))*cexsc, title = '# of SNPs', bty = 'n', cex = 0.5)

dat[pop == 'lof1114', plot(posgen, fst, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims, xlab = '', ylab = expression(F[ST]), bty = 'l', cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8, mgp = c(1, 0, 0))
mtext(side=3, 'B. Norway 2011-2014', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S8 Manhattan plot FST site-shuffle p-value
####################################################

# read in data: sliding window fst and site-shuffle p-values from ANGSD (GATK nodam2 unlinked sites)
dat <- fread('output/fst_siteshuffle.angsd.gatk.csv.gz') # output by angsd_fst_siteshuffle_null_stats.r

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

setkey(dat, CHROM)
dat <- dat[chrmax[, .(CHROM = chr, start)], ]
dat[, posgen := midPos + start]
dat[,start := NULL]

# trim out Unplaced or windows with <2 loci
dat <- dat[!(CHROM %in% c('Unplaced')) & nloci > 1, ]
nrow(dat)

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := lgcolsramp(0.2, lg = 1, thresh = 0.1)]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := lgcolsramp(0.2, lg = 2, thresh = 0.1)]

### set up plot
adjlet <- -0.14 # horizontal adjustment for subplot letter
cexlet <- 0.8
linelet <- 0.2
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS8.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.6, 0.2, 0.1))

ymin <- min(c(dat[, min(p, na.rm=TRUE)], 0.05))
xlims <- dat[, range(posgen, na.rm=TRUE)]

dat[pop == 'can', plot(posgen, p, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims, xlab = '',
                       ylim = c(1, ymin), ylab = 'p', bty = 'l', cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8) # plot x-axis
mtext(side=3, 'A. Canada 1940-2013', adj=adjlet, line=linelet, cex=cexlet)

legend('topleft', legend = c(2, 5, 10, 100), pch = 1, pt.cex = log(c(2,5,10,100))*cexsc, title = '# of SNPs', bty = 'n', cex = 0.5)

dat[pop == 'lof0711', plot(posgen, p, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims, xlab = '',
                           ylim = c(1, ymin), ylab = 'p', bty = 'l', cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8) # plot x-axis
mtext(side=3, 'B. Norway 1907-2011', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0714', plot(posgen, p, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims, xlab = '',
                           ylim = c(1, ymin), ylab = 'p', bty = 'l', cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8) # plot x-axis
mtext(side=3, 'C. Norway 1907-2014', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof1114', plot(posgen, p, type='p', cex=log(nloci)*cexsc, col=lgcol, xlim = xlims, xlab = '',
                           ylim = c(1, ymin), ylab = 'p', bty = 'l', cex.lab = 1.5, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8) # plot x-axis
mtext(side=3, 'D. Norway 2011-2014', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S9 Manhattan plot change in pi by region
####################################################
winsz = 5e4 # for scaling the windowed pi and thetaW values

# read in data: sliding window pi from ANGSD (GATK nodam2 unlinked sites)
dat1 <- fread('analysis/theta_change_region_50000.Can_40.Can_14.gatk.csv.gz', drop = 1) # from angsd_theta_siteshuffle_null.r
dat2 <- fread('analysis/theta_change_region_50000.Lof_07.Lof_11.gatk.csv.gz', drop = 1)
dat3 <- fread('analysis/theta_change_region_50000.Lof_07.Lof_14.gatk.csv.gz', drop = 1)
dat4 <- fread('analysis/theta_change_region_50000.Lof_11.Lof_14.gatk.csv.gz', drop = 1)

dat1[, pop := 'can']
dat2[, pop := 'lof0711']
dat3[, pop := 'lof0714']
dat4[, pop := 'lof1114']

dat <- rbind(dat1, dat2, dat3, dat4)

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

dat <- merge(dat[, .(pop, CHROM = Chromo, midPos = WinCenter, tPd)], chrmax[, .(CHROM = chr, start)])
dat[, POSgen := midPos + start]
dat[, start := NULL]

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := cols[1]]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := cols[2]]

### set up plot
adjlet <- -0.11 # horizontal adjustment for subplot letter
cexlet <- 1
linelet <- -0.5
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS9.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.8, 0.1, 0.1), mgp = c(4, 1, 0))

ylims <- dat[, range(tPd/winsz, na.rm=TRUE)]
xlims <- dat[, range(POSgen, na.rm=TRUE)]

dat[pop == 'can', plot(POSgen, tPd/winsz, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = expression("Change in" ~ pi), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'A', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0711', plot(POSgen, tPd/winsz, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = expression("Change in" ~ pi), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'B', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0714', plot(POSgen, tPd/winsz, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = expression("Change in" ~ pi), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'C', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof1114', plot(POSgen, tPd/winsz, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = expression("Change in" ~ pi), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'D', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S10 Manhattan plot change in D by region
####################################################

# read in data: sliding window D from ANGSD (GATK nodam2 unlinked sites)
dat1 <- fread('analysis/theta_change_region_50000.Can_40.Can_14.gatk.csv.gz', drop = 1) # from angsd_theta_siteshuffle_null.r
dat2 <- fread('analysis/theta_change_region_50000.Lof_07.Lof_11.gatk.csv.gz', drop = 1)
dat3 <- fread('analysis/theta_change_region_50000.Lof_07.Lof_14.gatk.csv.gz', drop = 1)
dat4 <- fread('analysis/theta_change_region_50000.Lof_11.Lof_14.gatk.csv.gz', drop = 1)

dat1[, pop := 'can']
dat2[, pop := 'lof0711']
dat3[, pop := 'lof0714']
dat4[, pop := 'lof1114']

dat <- rbind(dat1, dat2, dat3, dat4)

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

dat <- merge(dat[, .(pop, CHROM = Chromo, midPos = WinCenter, tDd)], chrmax[, .(CHROM = chr, start)])
dat[, POSgen := midPos + start]
dat[, start := NULL]

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := cols[1]]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := cols[2]]

### set up plot
adjlet <- -0.11 # horizontal adjustment for subplot letter
cexlet <- 1
linelet <- -0.5
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS10.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.8, 0.1, 0.1), mgp = c(4, 1, 0))

ylims <- dat[, range(tDd, na.rm=TRUE)]
xlims <- dat[, range(POSgen, na.rm=TRUE)]

dat[pop == 'can', plot(POSgen, tDd, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = "Change in D", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'A', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0711', plot(POSgen, tDd, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = "Change in D", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'B', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0714', plot(POSgen, tDd, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = "Change in D", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'C', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof1114', plot(POSgen, tDd, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = "Change in D", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'D', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S11 Manhattan plot change in LD by region
####################################################

# read in data: sliding window D from ngsLD (GATK nodam2 sites)
dat <- fread('analysis/ld_change_region_5e4_ngsLD.gatk.csv.gz'); width='5e4' # output by ngsLD_region_change.r

# Remove unplaced
dat <- dat[chr != 'Unplaced', ]

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]
dat <- merge(dat, chrmax[, .(chr, start)], by = 'chr')
dat[, POSgen := BIN_START + start + as.numeric(width)/2]

# long format for plotting
datl <- melt(dat[, .(chr, POSgen, ld_diff_Can, ld_diff_Lof0711, ld_diff_Lof0714, ld_diff_Lof1114)], id.vars = c('chr', 'POSgen'), variable.name = 'pop', value.name = 'ldchange')
datl[, pop := gsub('ld_diff_', '', pop)]

# add a vector for color by LG
lgs <- datl[, sort(unique(chr))]
datl[,lgcol := lgcolsramp(0.2, lg = 1, thresh = 0.1)]
datl[chr %in% lgs[seq(2, length(lgs),by=2)], lgcol := lgcolsramp(0.2, lg = 2, thresh = 0.1)]

### set up plot
adjlet <- -0.11 # horizontal adjustment for subplot letter
cexlet <- 1
linelet <- -0.5
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS11.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.8, 0.1, 0.1), mgp = c(4, 1, 0))

ylims <- datl[, range(ldchange, na.rm=TRUE)]
xlims <- datl[, range(POSgen, na.rm=TRUE)]

datl[pop == 'Can', plot(POSgen, ldchange, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = "Change in LD", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.75)
mtext(side=3, 'A', adj=adjlet, line=linelet, cex=cexlet)

datl[pop == 'Lof0711', plot(POSgen, ldchange, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = "Change in LD", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.75)
mtext(side=3, 'B', adj=adjlet, line=linelet, cex=cexlet)

datl[pop == 'Lof0714', plot(POSgen, ldchange, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = "Change in LD", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'C', adj=adjlet, line=linelet, cex=cexlet)

datl[pop == 'Lof1114', plot(POSgen, ldchange, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = "Change in LD", bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'D', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S12 Manhattan plot p-value for change in pi by region
####################################################

# read in data: sliding window pi from ANGSD (GATK nodam2 unlinked sites)
dat <- fread('analysis/theta_siteshuffle.angsd.gatk.csv.gz') # output by angsd_theta_siteshuffle_null_stats.r

# LG mid-points
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

# add a vector for color by LG
lgs <- dat[, sort(unique(Chromo))]
dat[,lgcol := lgcolsramp(0.2, lg = 1, thresh = 0.1)]
dat[Chromo %in% lgs[seq(2, length(lgs),by=2)], lgcol := lgcolsramp(0.2, lg = 2, thresh = 0.1)]

### set up plot
adjlet <- -0.14 # horizontal adjustment for subplot letter
cexlet <- 0.8
linelet <- 0.2
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS12.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.6, 0.2, 0.1))

ylims <- c(1, dat[, min(c(tPd.p, 0.05), na.rm=TRUE)])
xlims <- dat[, range(POSgen, na.rm=TRUE)]

dat[pop == 'can', plot(POSgen, tPd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'A. Canada 1940-2013', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0711', plot(POSgen, tPd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'B. Norway 1907-2011', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0714', plot(POSgen, tPd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'C. Norway 1907-2014', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof1114', plot(POSgen, tPd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'D. Norway 2011-2014', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S13 Manhattan plot p-value for change in D by region
####################################################

# read in data: sliding window pi from ANGSD (GATK nodam2 unlinked sites)
dat <- fread('analysis/theta_siteshuffle.angsd.gatk.csv.gz') # output by angsd_theta_siteshuffle_null_stats.r

# LG mid-points
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$len)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

# add a vector for color by LG
lgs <- dat[, sort(unique(Chromo))]
dat[,lgcol := lgcolsramp(0.2, lg = 1, thresh = 0.1)]
dat[Chromo %in% lgs[seq(2, length(lgs),by=2)], lgcol := lgcolsramp(0.2, lg = 2, thresh = 0.1)]

### set up plot
adjlet <- -0.14 # horizontal adjustment for subplot letter
cexlet <- 0.8
linelet <- 0.2
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS13.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.6, 0.2, 0.1))

ylims <- c(1, dat[, min(c(tDd.p, 0.05), na.rm=TRUE)])
xlims <- dat[, range(POSgen, na.rm=TRUE)]

dat[pop == 'can', plot(POSgen, tDd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'A. Canada 1940-2013', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0711', plot(POSgen, tDd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'B. Norway 1907-2011', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof0714', plot(POSgen, tDd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'C. Norway 1907-2014', adj=adjlet, line=linelet, cex=cexlet)

dat[pop == 'lof1114', plot(POSgen, tDd.p, type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = 'p', bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
abline(h = 0.05, lty = 2, col = 'grey')
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'D. Norway 2011-2014', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

############################################################
## Fig. S14 Power to detect selection from FST site-shuffle
############################################################
require(ggplot2)

# read in data: outlier test from pcangsd (GATK nodam2 unlinked sites)
dat <- fread('analysis/slim_fst_siteshuffle.summary.csv.gz') # output by slim_fst_siteshuffle_plot.R

# plot
fs13 <- ggplot(dat[comb == 1, ], aes(s, fpl05, group = f, color = as.factor(f))) +
  geom_point() +
  geom_smooth(method = 'lm') +
  facet_grid(~ ne, labeller = label_both) +
  coord_cartesian(ylim = c(0,1)) +
  theme(panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        axis.line = element_line(colour = "black"),
        panel.spacing = unit(1, "lines")) +
  labs(y = 'Proportion that\ndetect outlier loci', color = "Initial frequency")
ggsave(plot = fs13, filename = 'figures/figureS14.png', width = 7, height = 2, dpi = 300)

#########################################
# Fig. S15 Fst around putative outliers
#########################################
library(data.table)
ncol = 3 # number of columns in graph
rowinch <- 2 # inches per row in graph
rng <- 100000 # how many bp to go in either direction from SNP

# read in outliers
dat <- fread('tables/outlier_annotation.csv') # annotated list of outliers, from annotate_outliers.r
dat[, midPos := as.numeric(dat$midPos)]

# read in Fsts
fstLof1 <- fread('analysis/Lof_07.Lof_11.fst.AB.gz', col.names = c('CHROM', 'POS', 'A', 'B')) # from angsd_fst.sh
fstLof2 <- fread('analysis/Lof_07.Lof_14.fst.AB.gz', col.names = c('CHROM', 'POS', 'A', 'B'))
fstCan <- fread('analysis/Can_40.Can_14.fst.AB.gz', col.names = c('CHROM', 'POS', 'A', 'B'))

# trim fsts to nodam2 loci
gatk <- fread('data_2020.05.07/GATK_filtered_SNP_no_dam2.tab', col.names = c('CHROM', 'POS', 'REF', 'ALT'))
fstLof1 <- merge(fstLof1, gatk[, .(CHROM, POS)])
fstLof2 <- merge(fstLof2, gatk[, .(CHROM, POS)])
fstCan <- merge(fstCan, gatk[, .(CHROM, POS)])

# calc and merge fsts
fstLof1[, fst := A/B]
fstLof2[, fst := A/B]
fstLof <- merge(fstLof1[, .(CHROM, POS, fst1 = A/B)], fstLof2[, .(CHROM, POS, fst2 = A/B)])
fstLof[, fst := rowMeans(cbind(fst1, fst2))]
fstCan[, fst := A/B]

# read in and merge sample size data
datCan40 <- fread('data_31_01_20/Can_40_freq.mafs.gz')
datCan14 <- fread('data_31_01_20/Can_14_freq.mafs.gz')
datLof07 <- fread('data_31_01_20/Lof_07_freq.mafs.gz')
datLof11 <- fread('data_31_01_20/Lof_11_freq.mafs.gz')
datLof14 <- fread('data_31_01_20/Lof_14_freq.mafs.gz')

datCan <- merge(datCan40[, .(CHROM = chromo, POS = position, n1 = nInd)], datCan14[, .(CHROM = chromo, POS = position, n2 = nInd)])
datLof <- merge(datLof07[, .(CHROM = chromo, POS = position, n1 = nInd)], datLof11[, .(CHROM = chromo, POS = position, n2 = nInd)])
datLof <- merge(datLof, datLof14[, .(CHROM = chromo, POS = position, n3 = nInd)])

# harmonic mean sample size
datCan[, n := 2/(1/n1 + 1/n2)]
datLof[, n12 := 2/(1/n1 + 1/n2)]
datLof[, n13 := 2/(1/n1 + 1/n3)]
datLof[, n123 := 3/(1/n1 + 1/n2 + 1/n3)]

datCan[, n.sc := (n - min(n))/(max(n) - min(n))] # scale 0-1
datLof[, n12.sc := (n12 - min(n12))/(max(n12) - min(n12))] # scale 0-1
datLof[, n13.sc := (n13 - min(n13))/(max(n13) - min(n13))] # scale 0-1
datLof[, n123.sc := (n123 - min(n123))/(max(n123) - min(n123))] # scale 0-1

fstCan <- merge(fstCan, datCan[, .(CHROM, POS, n, n.sc)], by = c('CHROM', 'POS'))
fstLof1 <- merge(fstLof1, datLof[, .(CHROM, POS, n12, n12.sc)], by = c('CHROM', 'POS'))
fstLof2 <- merge(fstLof2, datLof[, .(CHROM, POS, n13, n13.sc)], by = c('CHROM', 'POS'))
fstLof <- merge(fstLof, datLof[, .(CHROM, POS, n12, n13, n123, n12.sc, n13.sc, n123.sc)], by = c('CHROM', 'POS'))

# merge adjacent regions of the same comparison & test
dat2 <- dat # the version to modify
dat2[, todelete := 0]
for(i in 1:nrow(dat)){
  j <- dat[, which(CHROM == dat$CHROM[i] & comp == dat$comp[i] & test == dat$test[i] & abs(midPos - dat$midPos[i]) < 100000)]
  if(length(j) > 1){
    print(i)
    dat2$POS[j[1]] <- paste0(dat$POS[j], collapse = ',')
    dat2$midPos[j[1]] <- mean(dat$midPos[j])
    dat2$todelete[j[-1]] <- 1
  }
}
dat2[, .(CHROM, POS, midPos, comp, test, todelete)]
dat2 <- dat2[todelete == 0, ]

# make plot
nrow = ceiling(nrow(dat)/ncol)
bty <- 'l'
outlcol <- 'red'

png(height= 9, width=6.5, units='in', res=300, file='figures/figureS15.png')
par(mfrow = c(nrow, ncol), mai = c(0.3, 0.3, 0.4, 0.1), omi = c(0.3, 0.3, 0, 0))
for(i in 1:nrow(dat2)){
  if(dat2$comp[i] %in% c('can', 'Canada')){
    thisdat <- fstCan[CHROM == dat2$CHROM[i] & abs(POS - dat2$midPos[i]) < rng, ]
    outl <- fstCan[CHROM == dat2$CHROM[i] & POS == dat2$midPos[i], ]
    if(grepl(',', dat2$POS[i])) outl <- fstCan[CHROM == dat2$CHROM[i] & POS %in% unlist(strsplit(dat2$POS[i], split = ',')), ] # if a joined region
    pop <- 'Canada'
    thisdat[, plot(POS/1e6, fst, cex = n.sc, xlab = '', ylab = '',
                   main = paste0(dat2$CHROM[i], ' ', pop), bty = bty)]
    outl[, points(POS/1e6, fst, col = outlcol)]
  }
  if(dat2$comp[i] == 'Norway 1907-2011'){
    thisdat <- fstLof1[CHROM == dat2$CHROM[i] & abs(POS - dat2$midPos[i]) < rng, ]
    outl <- fstLof1[CHROM == dat2$CHROM[i] & POS == dat2$midPos[i], ]
    if(grepl(',', dat2$POS[i])) outl <- fstCan[CHROM == dat2$CHROM[i] & POS %in% unlist(strsplit(dat2$POS[i], split = ',')), ] # if a joined region
    pop <- 'Norway 1907-2011'
    thisdat[, plot(POS/1e6, fst, cex = n12.sc, xlab = '', ylab = '',
                   main = paste0(dat2$CHROM[i], ' ', pop), bty = bty)]
    outl[, points(POS/1e6, fst, col = outlcol)]
  }
  if(dat2$comp[i] == 'Norway 1907-2014'){
    thisdat <- fstLof2[CHROM == dat2$CHROM[i] & abs(POS - dat2$midPos[i]) < rng, ]
    outl <- fstLof2[CHROM == dat2$CHROM[i] & POS == dat2$midPos[i], ]
    if(grepl(',', dat2$POS[i])) outl <- fstCan[CHROM == dat2$CHROM[i] & POS %in% unlist(strsplit(dat2$POS[i], split = ',')), ] # if a joined region
    pop <- 'Norway 1907-2014'
    thisdat[, plot(POS/1e6, fst, cex = n13.sc, xlab = '', ylab = '',
                   main = paste0(dat2$CHROM[i], ' ', pop), bty = bty)]
    outl[, points(POS/1e6, fst, col = outlcol)]
  }
  if(dat2$comp[i] == 'Norway 1907-2011-2014'){
    thisdat <- fstLof[CHROM == dat2$CHROM[i] & abs(POS - dat2$midPos[i]) < rng, ]
    outl <- fstLof[CHROM == dat2$CHROM[i] & POS == dat2$midPos[i], ]
    if(grepl(',', dat2$POS[i])) outl <- fstCan[CHROM == dat2$CHROM[i] & POS %in% unlist(strsplit(dat2$POS[i], split = ',')), ] # if a joined region
    pop <- 'Norway 1907-2011-2014'
    thisdat[, plot(POS/1e6, fst, cex = n123.sc, xlab = '', ylab = '',
                   main = paste0(dat2$CHROM[i], ' ', pop), cex.main = 0.95, bty = bty)]
    outl[, points(POS/1e6, fst, col = outlcol)]
  }
  if(dat2$comp[i] == 'Canada-Norway'){
    thisdat1 <- fstLof[CHROM == dat2$CHROM[i] & abs(POS - dat2$midPos[i]) < rng, ]
    thisdat2 <- fstCan[CHROM == dat2$CHROM[i] & abs(POS - dat2$midPos[i]) < rng, ]
    pop <- 'All'
    thisdat1[, plot(POS/1e6, fst, cex = n123.sc, xlab = '', ylab = '',
                    main = paste0(dat2$CHROM[i], ' ', pop), bty = bty)]
    thisdat2[, points(POS/1e6, fst, cex = n.sc, col = 'grey')]
  }
  title(main = dat2$test[i], line = 0.4, cex.main = 0.7)
}

mtext('Position (Mb)', side = 1, outer = TRUE)
mtext('Fst', side = 2, outer = TRUE)

dev.off()

#########################################
## Fig. S16 Manhattan plot FSTs by SNP
#########################################

# WFS results, from wfs_nullmodel_analysis.r
wfs <- fread('analysis/wfs_nullmodel_padj.csv.gz')

# 1907-2011-2014 fsts
fst <- fread('analysis/gatk.lof07-11-14.weir.fst', header=TRUE) # from vcftools --gzvcf data_2020.05.07/Historic_dataset_no_clip.vcf.gz_HF_GQ_HWE_MISS_IND_Kmer_VAR_Binom_No_Dam2.vcf.gz --weir-fst-pop data_2020.05.07/popLof07.txt --weir-fst-pop data_2020.05.07/popLof11.txt --weir-fst-pop data_2020.05.07/popLof14.txt --out analysis/gatk.lof07-11-14
datLof <- merge(fst, wfs[pop == 'Lof', .(CHROM, POS, pop, p.adj)], by = c('CHROM', 'POS'))

# Canada fsts
fstCan <- fread('analysis/gatk.can.weir.fst', header=TRUE) # from vcftools --gzvcf data_2020.05.07/Historic_dataset_no_clip.vcf.gz_HF_GQ_HWE_MISS_IND_Kmer_VAR_Binom_No_Dam2.vcf.gz --weir-fst-pop data_2020.05.07/popCan40.txt --weir-fst-pop data_2020.05.07/popCan13.txt --out analysis/gatk.can
datCan <- merge(fstCan, wfs[pop == 'Can', .(CHROM, POS, pop, p.adj)], by = c('CHROM', 'POS'))
nrow(datCan)

# combine
dat <- rbind(datLof, datCan)

# add genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$length)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

dat <- merge(dat, chrmax[, .(CHROM = chr, start)], by = c('CHROM'))
dat[, POSgen := POS + start]
dat[,start := NULL]

# trim out NAs
dat <- dat[!is.na(WEIR_AND_COCKERHAM_FST), ]

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := lgcolsramp(p.adj, lg = 1, thresh = 0.1)]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := lgcolsramp(p.adj, lg = 2, thresh = 0.1)]

# and for plotting a colorbar
cbar <- data.frame(logx = seq(0, 1.30103, length.out=50))
cbar$x <- 10^(-cbar$logx)
cbar$col1 <- lgcolsramp(cbar$x, lg = 1, thresh = 0.1)
cbar$col2 <- lgcolsramp(cbar$x, lg = 2, thresh = 0.1)

# order
setorder(dat, pop, -p.adj)

### set up plot
adjlet <- -0.11 # horizontal adjustment for subplot letter
cexlet <- 1
linelet <- -0.5
cexsc <- 1/5

ylims <- dat[, range(WEIR_AND_COCKERHAM_FST, na.rm=TRUE)]
xlims <- dat[, range(POSgen, na.rm=TRUE)]

png(height=3.5, width=6, units='in', res=300, file='figures/figureS16.png')
par(mfrow = c(2,1), las=1, mai=c(0.3, 0.8, 0.1, 0.1))

dat[pop == 'Can', plot(POSgen, WEIR_AND_COCKERHAM_FST, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = expression(F[ST]), bty = 'l', cex.lab = 1, cex.axis = 0.8, cex = 0.3, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.5, line = -0.6) # plot x-axis
mtext(side=3, 'A', adj=adjlet, line=linelet, cex=cexlet)

color.bar(cbar$col1, cbar$x, axis = FALSE, nticks = 5, xposmin =20e6, xposmax = 26e6, yposmin = -0.95, yposmax = -0.2)
color.bar(cbar$col2, cbar$x, cex = 0.7, axis = TRUE, nticks = 5,
          xposmin =26e6, xposmax = 32e6, yposmin = -0.95, yposmax = -0.2, title = 'p-value', titley = -0.05)

dat[pop == 'Lof', plot(POSgen, WEIR_AND_COCKERHAM_FST, type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = expression(F[ST]), bty = 'l', cex.lab = 1, cex.axis = 0.8, cex = 0.3, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.5, line = -0.6) # plot x-axis
mtext(side=3, 'B', adj=adjlet, line=linelet, cex=cexlet)

dev.off()

####################################################
## Fig. S17 Manhattan plot p-value for pcangsd
####################################################

# read in data: outlier test from pcangsd (GATK nodam2 unlinked sites)
dat <- fread('analysis/pcangsd_outlier.gatk.nodam.unlinked.csv.gz') # output by angsd_pcangsd_plot_selection.r

# LG mid-points and genome position
chrmax <- fread('data/lg_length.csv')
chrmax[, start := c(0,cumsum(chrmax$length)[1:(nrow(chrmax)-1)])]
chrmax[, mid := rowSums(cbind(start, length/2))]

dat <- merge(dat, chrmax[, .(CHROM = chr, start)], by = c('CHROM'))
dat[, POSgen := POS + start]
dat[,start := NULL]

# add a vector for color by LG
lgs <- dat[, sort(unique(CHROM))]
dat[,lgcol := lgcolsramp(0.2, lg = 1, thresh = 0.1)]
dat[CHROM %in% lgs[seq(2, length(lgs),by=2)], lgcol := lgcolsramp(0.2, lg = 2, thresh = 0.1)]

### set up plot
adjlet <- -0.14 # horizontal adjustment for subplot letter
cexlet <- 0.8
linelet <- 0.2
cexsc <- 1/5

png(height=5, width=6, units='in', res=300, file='figures/figureS17.png')
par(mfrow = c(4,1), las=1, mai=c(0.3, 0.6, 0.2, 0.1))

ylims <- c(0, dat[, max(c(-log10(pfdr), -log10(0.05)), na.rm=TRUE)])
xlims <- dat[, range(POSgen, na.rm=TRUE)]

dat[pop == 'can', plot(POSgen, -log10(pfdr), type='p', col=lgcol, xlim = xlims, xlab = '',
                       ylim = ylims, ylab = expression(-log[10](p)), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'A. Canada 1940-2013', adj=adjlet, line=linelet, cex=cexlet)
abline(h = -log10(0.05), lty = 2, col = 'grey')

dat[pop == 'lof0711', plot(POSgen, -log10(pfdr), type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = expression(-log[10](p)), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'B. Norway 1907-2011', adj=adjlet, line=linelet, cex=cexlet)
abline(h = -log10(0.05), lty = 2, col = 'grey')

dat[pop == 'lof0714', plot(POSgen, -log10(pfdr), type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = expression(-log[10](p)), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'C. Norway 1907-2014', adj=adjlet, line=linelet, cex=cexlet)
abline(h = -log10(0.05), lty = 2, col = 'grey')

dat[pop == 'lof1114', plot(POSgen, -log10(pfdr), type='p', col=lgcol, xlim = xlims, xlab = '',
                           ylim = ylims, ylab = expression(-log[10](p)), bty = 'l', cex.lab = 1, xaxt = 'n', xaxs = 'i', tcl = -0.3)]
axis(side=1, at = chrmax$mid, labels = gsub('LG|LG0', '', chrmax$chr), tick = FALSE, cex.axis = 0.8)
mtext(side=3, 'D. Norway 2011-2014', adj=adjlet, line=linelet, cex=cexlet)
abline(h = -log10(0.05), lty = 2, col = 'grey')

dev.off()

############################################################
## Fig. S18 Power to detect selection from PCAngsd
############################################################
require(ggplot2)

# read in data: outlier test from pcangsd (GATK nodam2 unlinked sites)
dat <- fread('analysis/slim_pcangsd.summary.csv.gz') # output by slim_pcangsdoutlier_plot.R

# plot
fs14 <- ggplot(dat[comb == 1, ], aes(s, prop, group = f, color = as.factor(f))) +
  geom_point() +
  geom_smooth(method = 'lm') +
  facet_grid(~ ne, labeller = label_both) +
  coord_cartesian(ylim = c(0,1)) +
  theme(panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.background = element_blank(),
        axis.line = element_line(colour = "black"),
        panel.spacing = unit(1, "lines")) +
  labs(y = 'Proportion that\ndetect outlier loci', color = "Initial frequency")
ggsave(plot = fs14, filename = 'figures/figureS18.png', width = 7, height = 2, dpi = 300)

######################
## Table S5 Outliers
######################
library(data.table)

# read in data
dat <- fread('tables/outlier_annotation.csv') # annotated list of outliers, from annotate_outliers.r
dat[, midPos := as.integer(midPos)]
freqCan40 <- fread('data_31_01_20/Can_40_freq.mafs.gz') # allele frequencies
freqCan14 <- fread('data_31_01_20/Can_14_freq.mafs.gz') # actually 2013, not 2014
freqLof07 <- fread('data_31_01_20/Lof_07_freq.mafs.gz')
freqLof11 <- fread('data_31_01_20/Lof_11_freq.mafs.gz')
freqLof14 <- fread('data_31_01_20/Lof_14_freq.mafs.gz')

setkey(freqCan40, chromo, position)
setkey(freqCan14, chromo, position)
setkey(freqLof07, chromo, position)
setkey(freqLof11, chromo, position)
setkey(freqLof14, chromo, position)

# p-value for 99th percentile is an fst, not a p-value, so erase it
dat[testvaltype == 'Average Fst', testval := NA_real_]
dat[testvaltype == 'Average Fst', testvaltype := '']

# adjust text slightly
dat[test == '99th percentile across pops', test := 'Shared 99th percentile']
dat[, test := gsub('pcangsd', 'PCAngsd', test)]
dat[testvaltype == 'Genome-wide p-value for the region', testvaltype := 'Genome-wide p']

# add column for functional location
dat[,annotation:='']
dat[NearGene != '', annotation:= paste(annotation, '<25kb from gene', sep = ', ')]
dat[grepl('CDS', feature), annotation:='coding']
dat[grepl('three_prime_UTR', feature), annotation:= paste(annotation, "3' UTR", sep = ', ')]
dat[grepl('five_prime_UTR', feature), annotation:= paste(annotation, "5' UTR", sep = ', ')]
dat[grepl('mRNA', feature) & !grepl('CDS|three_prime_UTR|five_prime_UTR', feature), annotation:= paste(annotation, "transcript", sep = ', ')]
dat[, annotation := gsub('^, ', '', annotation)]

# Add allele frequencies
dat[, Canada1940 := NA_character_]
dat[, Canada2013 := NA_character_]
dat[, Norway1907 := NA_character_]
dat[, Norway2011 := NA_character_]
dat[, Norway2014 := NA_character_]
dat[, sort(unique(comp))] # make sure all cases are handled below
for(i in 1:nrow(dat)){
  if(dat$region[i] == 0){ # can only do this for snps
    if(dat$comp[i] %in% c('can', 'Canada', 'Canada-Norway')){
      dat$Canada1940[i] <- freqCan40[.(dat$CHROM[i], dat$midPos[i]), round(knownEM, 2)]
      dat$Canada2013[i] <- freqCan14[.(dat$CHROM[i], dat$midPos[i]), round(knownEM, 2)]
    }
    if(dat$comp[i] %in% c('Norway 1907-2011-2014', 'Norway 1907-2011', 'Canada-Norway', 'Norway 1907-2014')){
      dat$Norway1907[i] <- freqLof07[.(dat$CHROM[i], dat$midPos[i]), round(knownEM, 2)]
    }
    if(dat$comp[i] %in% c('Norway 1907-2011-2014', 'Norway 1907-2011', 'Canada-Norway')){
      dat$Norway2011[i] <- freqLof11[.(dat$CHROM[i], dat$midPos[i]), round(knownEM, 2)]
    }
    if(dat$comp[i] %in% c('Norway 1907-2011-2014', 'Norway 1907-2014', 'Canada-Norway')){
      dat$Norway2014[i] <- freqLof14[.(dat$CHROM[i], dat$midPos[i]), round(knownEM, 2)]
    }
  }
}

# make Table S5
out <- dat[,.(LG=CHROM, Position=POS, Pop=gsub('can', 'Canada', comp), Canada1940, Canada2013, Norway1907, Norway2011, Norway2014, Test=paste0(test, ' ', gsub(' p-value', '', testvaltype)), p = signif(testval,2), annotation,
	Gene=lapply(strsplit(paste(WithinAnno, NearAnno), split=':'), FUN=function(x) return(x[1])))]

# convert to text and turn NA to ''
for (j in names(out)) set(out, j = j, value = as.character(out[[j]]))
for (j in names(out)) set(out, i=which(is.na(out[[j]])), j = j, value = '')

# write out
write.csv(out, file='tables/tableS5.csv', row.names=FALSE)


library(tradepolicy)
library(eflm)
library(bench)

# data ----

ch1_application2 <- agtpa_applications %>%
  select(exporter, importer, pair_id, year, trade, dist, cntg, lang, clny) %>%
  # this filter covers both OLS and PPML
  filter(year %in% seq(1986, 2006, 4)) %>%
  mutate(
    # variables for both OLS and PPML
    exp_year = paste0(exporter, year),
    imp_year = paste0(importer, year),
    year = paste0("log_dist_", year),
    log_trade = log(trade),
    log_dist = log(dist),
    smctry = ifelse(importer != exporter, 0, 1),

    # PPML specific variables
    log_dist_intra = log_dist * smctry,
    intra_pair = ifelse(exporter == importer, exporter, "inter")
  ) %>%
  spread(year, log_dist, fill = 0) %>%
  mutate(across(log_dist_1986:log_dist_2006, ~ .x * (1 - smctry)))

# ols ----

form <- log_trade ~ 0 + log_dist_1986 + log_dist_1990 + log_dist_1994 +
  log_dist_1998 + log_dist_2002 + log_dist_2006 + cntg +
  lang + clny + exp_year + imp_year

d <- filter(ch1_application2, importer != exporter, trade > 0)

bench_ols <- mark(
  lm(
    form,
    data = d,
    y = FALSE,
    model = FALSE
  )$coefficients,
  elm(
    form,
    data = d,
    y = FALSE,
    model = FALSE,
    reduce = TRUE
  )$coefficients
)

rm(d)

# ppml ----

form <- trade ~ 0 + log_dist_1986 + log_dist_1990 +
  log_dist_1994 + log_dist_1998 + log_dist_2002 + log_dist_2006 +
  cntg + lang + clny + exp_year + imp_year

d <- filter(ch1_application2, importer != exporter)

bench_ppml <- mark(
  glm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE
  )$coefficients,
  eglm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE,
    reduce = TRUE
  )$coefficients
)

rm(d)

# internal distance ----

form <- trade ~ 0 + log_dist_1986 + log_dist_1990 +
  log_dist_1994 + log_dist_1998 + log_dist_2002 + log_dist_2006 +
  cntg + lang + clny + exp_year + imp_year + log_dist_intra

d <- ch1_application2

bench_internal_distance <- mark(
  glm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE
  )$coefficients,
  eglm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE,
    reduce = TRUE
  )$coefficients
)

rm(d)

# internal distance and home bias ----

form <- trade ~ 0 + log_dist_1986 + log_dist_1990 +
  log_dist_1994 + log_dist_1998 + log_dist_2002 + log_dist_2006 +
  cntg + lang + clny + exp_year + imp_year + log_dist_intra + smctry

d <- ch1_application2

bench_home_bias <- mark(
  glm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE
  )$coefficients,
  eglm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE,
    reduce = TRUE
  )$coefficients
)

rm(d)

# fe ----

form <- trade ~ 0 + log_dist_1986 + log_dist_1990 +
  log_dist_1994 + log_dist_1998 + log_dist_2002 + log_dist_2006 +
  cntg + lang + clny + exp_year + imp_year + intra_pair

d <- ch1_application2

bench_fe <- mark(
  glm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE
  )$coefficients,
  eglm(
    form,
    family = stats::quasipoisson(link = "log"),
    data = d,
    y = FALSE,
    model = FALSE,
    reduce = TRUE
  )$coefficients
)

rm(d, form, ch1_application2)

save.image("benchmarks/02-chapter1-distance-puzzle-bench.RData", compress = "xz")

rm(list = ls())
gc()


#!/usr/bin/R
# Analyze raw data (ecr.xlsx) from C&A ECR study

# required packages
# - readxl
# - dplyr
# - Hmisc (requires acepack, requires gcc-fortran on *nix)

# library
library("readxl")
library("dplyr")
library("Hmisc")

# import raw excel to data frame
raw_data <- read_excel("raw_data.xlsx", sheet="Raw Data")

# convert data to numerical matrix
df <- as.matrix(as.data.frame(lapply(raw_data, as.numeric)))

# output file
sink("ecr_descriptive_statistics.txt", split=TRUE, append=FALSE)

# descriptive statistics
summary(raw_data)

# output file
sink("ecr_corr_sig.txt", split=TRUE, append=FALSE)

# correlations (r) and significance (p)
rcorr(df, type=c("pearson"))


junk <-matrix(scan("Data/intensity.dat"),byrow=T,ncol=3)
junk1<-junk[,1]
junk2<-junk[,2]
junk3<-junk[,3]
junk1<-rev(junk1);junk2<-rev(junk2);junk3<-rev(junk3)
junkn<-table(junk1)
y<-rep(1,junkn[1])
u<-junk3[junk1==0];v<-junk2[junk1==0]
B<-rep(0,junkn[1])
A<-1+0*u
for (i in 2:7){
   y<-c(y,rep(0,(i-1)*junkn[i]),rep(1,junkn[i]))
   B<-c(B,rep(i+1,i*junkn[i]))
   junku<-junk3[junk1==i]
   junkv<-junk2[junk1==i]
   u<-c(u,rep(junku,i))
   v<-c(v,rep(junkv,i))
   junkA<-rep(1:i,rep(junkn[i],i))
   A<-c(A,junkA)
   }
y<-c(y,rep(0,7*junkn[8]))
B<-c(B,rep(8,7*junkn[8]))
BB<-as.factor(B)
junku<-junk3[junk1==8]
junkv<-junk2[junk1==8]
u<-c(u,rep(junku,7))
v<-c(v,rep(junkv,7))
junkA<-rep(1:7,rep(junkn[8],7))
A<-c(A,junkA)
AA<-as.factor(A)

res <- glm(y~AA-1,family="binomial")


library(tidyverse)
theme_set(theme_bw(18))
source("helpers.r")

getNormalizedProbability = function(d) {
  dd = as.data.frame(d)
  total = sum(dd$MeanResponse)
  smooth = dd %>%
    group_by(slider_id) %>%
    mutate(NormalizedProb = MeanResponse / total)
  return(smooth$NormalizedProb)
}

# load("data/r.RData")
# load("../../1_sinking-marbles-prior/results/data/priors.RData")
r1 = read.csv("../data/sinking_marbles.csv", header=T)
r1$workerid = r1$workerid + 60

r2 = read.csv("../../../20_sinking-marbles-prior-sliders/results/data/sinking_marbles.csv", header=T)

r = rbind(r1,r2)
nrow(r)
head(r)

r$trial = r$slide_number_in_experiment - 2
r = r[,c("assignmentid","workerid", "rt", "effect", "cause","language","gender.1","age","gender","other_gender", "object_level", "response", "object","slider_id","num_objects","trial","enjoyment","asses","comments","Answer.time_in_minutes","slider_id")]
r$object_level = factor(r$object_level, levels=c("object_high", "object_mid", "object_low"))
r$Item = as.factor(paste(r$effect,r$object))
table(r$Item)

# compute normalized probabilities
nr = r %>%
  select(Item,slider_id,response) %>%
  group_by(Item,slider_id) %>%
  summarize(MeanResponse = mean(response)) %>%
  ungroup() %>%
  group_by(Item) %>%
  nest() %>%
  mutate(SmoothedProportion = map(data,getNormalizedProbability)) %>%
  unnest()

ggplot(nr, aes(x=slider_id,y=SmoothedProportion)) +
  geom_point() +
  geom_line() +
  facet_wrap(~Item)

# write probabilities (no smoothing required, but we call it smoothed anyway for comparison with the other priors datasets) for model comparison
smoothed_spread = nr %>%
  select(-MeanResponse) %>%
  spread(slider_id,SmoothedProportion)

write.csv(smoothed_spread,file="../data/priors_binnedhistogram_smoothed.csv",row.names=F,quote=F)
write.csv(smoothed_spread,file="../../../../bayesian_model_comparison/data/priors_binnedhistogram_smoothed.csv",row.names=F,quote=F)

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

ggplot(aes(x=gender.1), data=r) +
  geom_histogram()

ggplot(aes(x=rt), data=r) +
  geom_histogram() +
  scale_x_continuous(limits=c(0,50000))

ggplot(aes(x=age), data=r) +
  geom_histogram()

ggplot(aes(x=age,fill=gender.1), data=r) +
  geom_histogram()

ggplot(aes(x=enjoyment), data=r) +
  geom_histogram()

ggplot(aes(x=asses), data=r) +
  geom_histogram()

ggplot(aes(x=Answer.time_in_minutes), data=r) +
  geom_histogram()

ggplot(aes(x=age,y=Answer.time_in_minutes,color=gender.1), data=unique(r[,c("assignmentid","age","Answer.time_in_minutes","gender.1")])) +
  geom_point() +
  geom_smooth()

unique(r$comments)


# Import required libraries

library(dplyr)

# Import required data from dataset files

importDataset <- function() {
    activity_labels <<- read.table("specdata/activity_labels.txt", col.names = c("activity_id", "activity_name"))
    feature_labels <<- read.table("specdata/features.txt", col.names = c("feature_id", "feature_name"))

    train_set <<- read.table("specdata/train/X_train.txt", col.names = feature_labels$feature_name)
    train_labels <<- read.table("specdata/train/y_train.txt", col.names = c("activity_id"))
    train_subject <<- read.table("specdata/train/subject_train.txt", col.names = c("subject_id"))

    test_set <<- read.table("specdata/test/X_test.txt", col.names = feature_labels$feature_name)
    test_labels <<- read.table("specdata/test/y_test.txt", col.names = c("activity_id"))
    test_subject <<- read.table("specdata/test/subject_test.txt", col.names = c("subject_id"))

    print("Datasets imported successfully!")
}

# Merge train and test dataset

mergeDataset <- function() {
    if (!(exists("activity_labels") &&
        exists("feature_labels") &&
        exists("train_set") &&
        exists("train_labels") &&
        exists("train_subject") &&
        exists("test_set") &&
        exists("test_labels") &&
        exists("test_subject"))) {
        importDataset()
    }

    merged_subjects <- rbind(train_subject, test_subject)
    merged_sets <- rbind(train_set, test_set)
    merged_labels <- rbind(train_labels, test_labels)
    merged_labels$activity_name <- activity_labels$activity_name[match(merged_labels$activity_id, activity_labels$activity_id)]

    merged_dataset <<- cbind(merged_subjects, merged_labels, merged_sets)

    print("Train and test datasets merged successfully!")
}

# Select measurement columns labeled with mean or std

extractMeansStd <- function() {
    if (!exists("merged_dataset")) {
        mergeDataset()
    }

    merged_dataset %>%
        select("subject_id", "activity_name", contains("mean") | contains("std"))
}

# Create a separate dataset summarizing the means of each measurement from the previous function

generateTidyMeans <- function() {
    tidy_means <<- extractMeansStd() %>%
        group_by(subject_id, activity_name) %>%
        summarize_at(base_data, -c(1:2), mean)

    write.table(tidy_means, file = "tidy_means.txt", row.names = FALSE)
}


readallmcmc <- function(file){
    years <- scan(file,skip=1,nlines=1)
    nyrs <- length(years)
    ages <- scan(file,skip=3,nlines=1)
    nages <- length(ages)
    variables <- scan(file,skip=5,nlines=1) #.  1 N etc
    names(variables) <- c("n","cw","sw","mat","f","cno")
    n1 <- length(variables[variables==1])
    print(paste("n1",n1))
    nvar <- 1:length(variables)
    nvar <- nvar[variables==1]
    dat <- read.table(file,skip=8,header=F)
    niter <- nrow(dat)/(nyrs*n1)
    iter <- 1:niter
    result <- expand.grid(list(year=years,variable=nvar,iter=iter,age=ages))
    result$value <- c(unlist(dat))
    return(result)
}


# Add the library statement here. For some reason this isnt working currently.
library(helpRFunctions)

context('Test bucketize.cont')

test_that('bucketize.cont returns the correct amount of values.', {
  vals = c(-1.5, 1000.1, .7, NA)
  expect_equal(bucketize.cont(vals, bucket.size = 1), c(-1, 1001, 1, NA))
  expect_equal(bucketize.cont(vals, bucket.size = 10000), c(0, 10000, 10000, NA))
  expect_equal(bucketize.cont(vals, bucket.size = .025), c(-1.5, 1000.1, .7, NA))
  }
  )


# Load ggplot2
library(ggplot2)
library(dplyr)

# The mtcars dataset is natively available
# head(mtcars)

# A really basic boxplot.

#

#### PART4 : Visualization with ggplot2
#--------------------------------------------------------------

# This libaray was created by Hadley Wickham based on graphical method developed by Leland Wilkison. This defines a standard grammar used to create graphs. Since 2014, ggplots is the most downloadable libaries in R.
#
# GGplot2 works based on the concept of Layering. Each part of the graph is visualized using a specific layer such as layer for coordinate, layer for size, layer for lable, layer for specific type of graph..
#
# Ok then how can we practice, now if we want to visualize the relationship between low and high temperature, change color according to season and choose size, how can we do.
#
#  Steps in using ggplot2
#
# A graph is created using ggplot() function, and right after this, we can add layer.
#
# data <- data.frame(data1)
#
#
# ggplot(data = dataframe, aes(x = var1, y= var2, colour = var3, size = var4))
#
# geom_xxx()   ---> this layer is added after ggplot() and used to create type of graph such as boxplot, histogram, etc. Note that if we do not specify anything inside the bracket, R will create a default graph, but if we want to decorate, we can declare inside the bracket what attributes we want for our graph. When we add attribute, new layer will be added on top of the previous declared attribute. that is how it works.
#
# ggplot(data, coordinate)+
#   layer1+
#   layer2+
#   layer3

mtcars <-data.frame(mtcars)

ggplot(mtcars, aes(x="", y=mpg))+
  geom_boxplot()

ggplot(mtcars, aes(x="", y=mpg)) +
  geom_boxplot() +
  xlab("")+
  ylab("Mile per gallon")

ggplot(mtcars, aes(x="", y=mpg)) +
  geom_boxplot(alpha=0.5, colour = "red", fill = "blue") +
  xlab("")

ggplot(mtcars, aes(x="", y=mpg)) +
  geom_boxplot(fill="yellow", alpha=0.3, colour = "blue") +
  xlab("")+
  labs(title="Boxplot for variable mpg",
       subtitle="Example of adding collor to the graph",
       y = "Distance km/1 gallon ")

#http://applied-r.com/colors-in-ggplot/
#
ggplot(mtcars, aes(x="", y=mpg)) +
  geom_boxplot(fill="slateblue", alpha=0.2) +
  xlab("")+
  ylab("Distance km/1 gallon")

# same graphs but adding more variable

#http://www.sthda.com/english/wiki/ggplot2-point-shapes
#http://applied-r.com/rcolorbrewer-palettes/

library(RColorBrewer)

ggplot(mtcars, aes(x=as.factor(cyl), y=mpg)) +
  geom_boxplot(fill="red", alpha=0.6)+
  xlab("cylinder")+
  ylab("Distance km/1 gallon")+
  stat_summary(fun=mean, geom="point", shape=23, size=10, color="yellow", fill="yellow")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="Set1")

ggplot(mtcars, aes(x=as.factor(cyl), y=mpg, fill = as.factor(cyl))) +
  geom_boxplot(alpha=0.7) +
  xlab("cylinder")+
  ylab("Distance km/1 gallon")+
  stat_summary(fun=mean, geom="point", shape=20, size=4, color="red", fill="red")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="Set1")

ggplot(mtcars, aes(x=as.factor(cyl), y=mpg, fill = as.factor(cyl))) +
  geom_boxplot(alpha=0.7) +
  xlab("Cylinder")+
  ylab("Distance km/1 gallon")+
  stat_summary(fun=mean, geom="point", shape=20, size=4, color="red", fill="red")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="Dark2")

#about palette http://rstudio-pubs-static.s3.amazonaws.com/5312_98fc1aba2d5740dd849a5ab797cc2c8d.html

cbPalette <- c("#999999", "#E69F00", "#56B4E9")

# The palette with black:
cbbPalette <- c("#000000", "#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7")

# To use for fills, add
scale_fill_manual(values=cbPalette)

# To use for line and point colors, add
scale_colour_manual(values=cbPalette)

ggplot(mtcars, aes(x=as.factor(cyl), y=mpg, fill = as.factor(cyl))) +
  geom_boxplot(alpha=0.7) +
  xlab("Cylinder")+
  ylab("Distance/Gallon")+
  stat_summary(fun=mean, geom="point", shape=20, size=4, color="red", fill="red")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="cbPalette")

#
p <- ggplot(mtcars, aes(x=as.factor(cyl), y=mpg, fill = as.factor(cyl))) +
  geom_boxplot(alpha=0.7, color ="red")+
  xlab("Cylinder")+
  ylab("Distance/Gallon")+
  stat_summary(fun=mean, geom="point", shape=20, size=4, color="red", fill="red")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="Dark2")

# background

p+theme_gray(base_size = 14)

p+theme_bw()

p + theme_minimal()

p + theme_classic()

p + theme_void()

p + theme_dark()

q<-ggplot(mtcars, aes(x=as.factor(cyl), y=mpg, fill = as.factor(cyl))) +
  geom_boxplot(alpha=0.7, color ="red") +
  stat_summary(fun=mean, geom="point", shape=20, size=4, color="red", fill="red")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="Dark2")

ggplot(mtcars%>%
         filter(mpg>20),aes(x=as.factor(cyl), y=mpg, fill = as.factor(cyl))) +
  geom_boxplot(alpha=0.7, color ="red") +
  stat_summary(fun=mean, geom="point", shape=20, size=4, color="red", fill="red")+
  theme(legend.position="none") +
  scale_fill_brewer(palette="Dark2")

q

q+
  labs(title="Boxplot Example",
       subtitle = "Another example uusing xlab and ylab",
       x ="Cylinder",
       y ="Distance/Gallon",
       caption = "(Data source from ...)",fill="Legends")+
      theme(legend.position="bottom")+
      scale_fill_manual(values = c("#d8b365", "#f5f5f5", "#5ab4ac"))

#http://www.cookbook-r.com/Graphs/Legends_(ggplot2)/

# display values on the graph

means <- aggregate(mpg ~  cyl, mtcars, mean)
means

q+
  geom_text(data = means, aes(label = mpg, y = mpg))

#### trang trí thêm một chút

means <-round(means,2)

means

q+ geom_text(data = means, aes(label = mpg, y = mpg + 0.08))

q+ geom_text(data = means, aes(label = mpg, y = mpg + 1))

# Annotation

#http://www.sthda.com/english/wiki/ggplot2-texts-add-text-annotations-to-a-graph-in-r-software

#https://ggplot2.tidyverse.org/reference/annotate.html
q+ geom_text(data = means, aes(label = mpg, y = mpg + 1))+
  annotate(geom="text", x=3, y=30, label="This is an annotation",
           color="red", size =5)

# Some additional library

library(tidyverse)
library(hrbrthemes)
library(viridis)
library(dplyr)

# create a dataset
data <- data.frame(
  name=c( rep("A",500), rep("B",500), rep("B",500), rep("C",20), rep('D', 100)  ),
  value=c( rnorm(500, 10, 5), rnorm(500, 13, 1), rnorm(500, 18, 1), rnorm(20, 25, 4), rnorm(100, 12, 1) )
)

# biểu đồ đàn vĩ cầm, dùng biểu đồ này khi dữ liệu của chúng ta khá nhiều và phân tán

data %>%
  ggplot(aes(x=name, y=value, fill=name)) +
  geom_violin() +
  scale_fill_viridis(discrete = TRUE, alpha=0.6, option="A") +
  theme_ipsum() +
  theme(
    legend.position="none",
    plot.title = element_text(size=11)
  ) +
  ggtitle("Violin chart") +
  xlab("")

# back to car data

# method
mtcars %>%
  ggplot( aes(x=as.factor(cyl), y=mpg, fill=as.factor(cyl))) +
  geom_violin() +
  scale_fill_viridis(discrete = TRUE, alpha=0.6, option="A") +
  theme_ipsum() +
  theme(
    legend.position="none",
    plot.title = element_text(size=11)
  ) +
  ggtitle("Violin chart") +
  xlab("")

# Reprensing data on boxplot

# Plot
data %>%
  ggplot( aes(x=name, y=value, fill=name)) +
  geom_boxplot() +
  scale_fill_viridis(discrete = TRUE, alpha=0.6) +
  geom_jitter(color="black", size=0.4, alpha=0.9) +
  theme_ipsum() +
  theme(
    legend.position="none",
    plot.title = element_text(size=11)
  ) +
  ggtitle("A boxplot with jitter") +
  xlab("")

# data mtcars

z<-mtcars %>%
  ggplot(aes(x=as.factor(cyl), y=mpg, fill=as.factor(cyl))) +
  geom_boxplot() +
  scale_fill_viridis(discrete = TRUE, alpha=0.6) +
  geom_jitter(color="black", size=0.4, alpha=0.9) +
  theme_ipsum() +
  theme(
    legend.position="none",
    plot.title = element_text(size=11)
  ) +
  ggtitle("Example with jitter") +
  xlab("")

z

# another example - Density

# Load dataset from github
data <- read.table("https://raw.githubusercontent.com/holtzy/data_to_viz/master/Example_dataset/1_OneNum.csv", header=TRUE)
data <- data.frame(data)
view(data)
# Make the histogram

ggplot(data, aes(x=price)) +
  geom_boxplot()

# method 1
ggplot(data, aes(x=price)) +
  geom_density(fill="yellow", color="blue", alpha=0.8)

## method 2
data %>%
  # filter( price<300 ) %>%
  ggplot( aes(x=price)) +
  geom_density(fill="#69b3a2", color="#e9ecef", alpha=0.8)

ggplot(data, aes(x=price)) +
  geom_boxplot()

data %>%
  ggplot( aes(x=price)) +
  geom_boxplot()

data %>%
   filter(price<300) %>%
  ggplot( aes(x=price)) +
  geom_boxplot()

quantile(data$price,probs=seq(0,1,0.05)) #quantile from 0 to 1 with a step of 5%
quantile(data$price,probs=seq(0.95,1,0.01)) #quantilte from 0.95 to 1 with a step of 1%

data %>%
  # filter( price<300 ) %>%
  ggplot( aes(x=price)) +
  geom_density(fill="#69b3a2", color="#e9ecef", alpha=0.8)

p<-data %>%
  filter( price<300 ) %>%
  ggplot( aes(x=price)) +
  geom_density(fill="#69b3a2", color="#e9ecef", alpha=0.8)

p

p+labs(title = "Density Plot",
       y = "Density",
       x = "Price in USD"

)

q<-data %>%
  filter( price<300 ) %>%
  ggplot( aes(x=price))+
  geom_histogram(fill="red",alpha=0.4) #fill="#69b3a2", color="#e9ecef", alpha=0.8

q

p  +
  geom_histogram(aes(y = ..density..),
                 colour = 1, fill = "white")

p
q

p  +
  geom_histogram(aes(y = ..density..),
                 colour = 1, fill = "yellow", alpha = 0.2)

q +
  annotate("text", x=250, y=500, label= "Example 1", cex=5, colour = "blue", angle = 0)
# draw a line

q +
  annotate("text", x=250, y=750, label= "Ví dụ 1", cex=5, colour = "red", angle = 90)+
  geom_vline(xintercept = 200, colour="blue", lwd = 1)

library(plotly)

p

a1<-ggplotly(p)
a1

a2<-ggplotly(q)

a1

a2


x <- rnorm(100)
res <- mean(x)
print("Running the R job again!!")
save(res, file="/data/results.rdata")


library(ggplot2)
library(reshape2)
shinyUI(navbarPage(
  "Population Genetics",

  tabPanel("Make Your Population",
  tags$head(tags$link(rel = "icon", type = "image/x-icon", href =
       "https://webresource.its.calpoly.edu/cpwebtemplate/5.0.1/common/images_html/favicon.ico")),

  fluidRow(
    column(4,
      numericInput("start", tags$div(
        HTML(paste("Starting %")),
        HTML(paste("A", tags$sup(1), sep = ""))
         ), 50),
      numericInput("gen", "Number of Generations", 100),
      numericInput("pop", "Starting Population", 100),
      numericInput("AA",
                   #withMathJax(helpText('$$A^1$$')),
                   tags$div(
                     HTML(paste("A", tags$sup(1), sep = "")),
                     HTML(paste("A", tags$sup(1), sep = "")),
                     HTML(paste("Fitness %"))
                   ),
                   100),
      numericInput("AB",
                   tags$div(
                     HTML(paste("A", tags$sup(1), sep = "")),
                     HTML(paste("A", tags$sup(2), sep = "")),
                     HTML(paste("Fitness %"))
                   )
                   , 100),
      numericInput("BB",
                   tags$div(
                     HTML(paste("A", tags$sup(2), sep = "")),
                     HTML(paste("A", tags$sup(2), sep = "")),
                     HTML(paste("Fitness %"))
                   ), 100),
      actionButton("begin","Click To Begin"),
      uiOutput("uiButton1"),
      uiOutput("uiButton2"),
      uiOutput("uiButton3"),
      br(), br(), br(), br(),

      div("Shiny app by", a(href="https://www.linkedin.com/in/helen-hilton-39aba269",target="_blank",
                            "Helen Hilton"),align="left", style = "font-size: 8pt"),

      div("Base R code by", a(href="https://www.linkedin.com/in/helen-hilton-39aba269/",target="_blank",
                              "Helen Hilton"),align="left", style = "font-size: 8pt"),

      div("Shiny source files:", a(href="https://gist.github.com/calpolystat/f9dd35b71dc3e73b845160018ecacea2",
                                   target="_blank","GitHub Gist"),align="left", style = "font-size: 8pt"),

      div(a(href="http://www.statistics.calpoly.edu/shiny",target="_blank",
            "Cal Poly Statistics Dept Shiny Series"),align="left", style = "font-size: 8pt")

    ),

    column(4,
           absolutePanel(
             top=125, bottom=0, left=320, right=0, fixed=TRUE,
             #              h3("Population 1"),
             imageOutput("Population1")

           )
    ),

    column(4,
           absolutePanel(
             top=425, bottom=0, left=320, right=0, fixed=TRUE,
             #              h3("Population 2"),
             imageOutput("Population2")

           )
    ),

    column(4,
           absolutePanel(  #Prints information for Population #1
             top=170, bottom=0, left=670, right=0, fixed=TRUE,
             tags$style(type="text/css",
                        ".shiny-output-error { visibility: hidden; }",
                        ".shiny-output-error:before { visibility: hidden; }"),
             uiOutput("start1"),
             textOutput("gen1"),
             textOutput("pop1"),
              textOutput("AA1"),
              textOutput("AB1"),
              textOutput("BB1")
           ), #absolutePanel
           absolutePanel(  #Prints information for Population #2
             top=470, bottom=0, left=670, right=0, fixed=TRUE,
              textOutput("start2"),
              textOutput("gen2"),
              textOutput("pop2"),
              textOutput("AA2"),
              textOutput("AB2"),
              textOutput("BB2")
           )

    ) #column

  ) #fluidRow
), #tabPanel

tabPanel("Simulation",
         sidebarLayout(

           sidebarPanel(

             numericInput("drawmorepops", "Number of Simulations", 1),
             actionButton("drawpop", "Simulate"),
             br(), br(), br(), br(),

             div("Shiny app by", a(href="https://www.linkedin.com/in/helen-hilton-39aba269/",target="_blank",
                                   "Helen Hilton"),align="right", style = "font-size: 8pt"),

             div("Base R code by", a(href="https://www.linkedin.com/in/helen-hilton-39aba269/",target="_blank",
                                     "Helen Hilton"),align="right", style = "font-size: 8pt"),

             div("Shiny source files:", a(href="https://gist.github.com/calpolystat/f9dd35b71dc3e73b845160018ecacea2",
                                          target="_blank","GitHub Gist"),align="right", style = "font-size: 8pt"),

             div(a(href="http://www.statistics.calpoly.edu/shiny",target="_blank",
                   "Cal Poly Statistics Dept Shiny Series"),align="right", style = "font-size: 8pt")

           ),#sidebarPanel
           mainPanel(
             plotOutput("distPlot")
           )#mainPanel
         )#sidebarLayout

) #tabPanel
) #navbarPage
) #shinyUI


args = commandArgs(trailingOnly=TRUE)
my.theta <- as.numeric(args[1])  ## Level of allelic imbalance
simruns <- as.numeric(args[2])   ## Number of simulations
out.file <- args[3]
nreps <- as.numeric(args[4])     ## Number of biological replicates
n_allele_specific_reads <- as.numeric(args[5]) ## Sum of unambiguously mapped reads (i.e. map better to one allele)
                                               ## in the experiment across all bioreps

my.qtest <- 0.8
my.qline <- 0.8
myreads <- n_allele_specific_reads/2/nreps/mean(my.qtest, my.qline) ## This is coverage from the G3 paper
my.alpha <- sqrt((1/my.theta)-1) ## In the stan2 model, theta = 1/(alpha[i]^2+1)
out.file <- paste0(paste(out.file, 'theta', my.theta, 'rsim-g1', my.qtest, 'rsim-g2',
                  my.qline, 'nbiorep', nreps, 'allelicreads', n_allele_specific_reads,
                  'simruns', simruns, sep='_'), '.tsv')

#Removed set.seed(0) because always gave the same result!!!!!
#set.seed(0)

seqcond <- "c1"

repsuffix <- paste0("_rep", seq(1, nreps))

## The three possible alignment states: better to allele g1 (formerly tester),
## better to allele g2 (formerly line), equally good to both.
counttester <- as.vector(sapply(paste("counts_", seqcond, "_g1_total", sep=""), paste0, repsuffix))
countline <- as.vector(sapply(paste("counts_", seqcond, "_g2_total", sep=""), paste0, repsuffix))
countboth <- as.vector(sapply(paste("counts_", seqcond, "_both_total", sep=""), paste0, repsuffix))
countheader <- paste(counttester, countline, countboth, sep="\t", collapse="\t")
priorstester <- paste("prior_", seqcond, "_g1", sep="")
priorsline <- paste("prior_", seqcond, "_g2", sep="")
allpriors <- sort(c(priorstester, priorsline))
seqreps <- paste(seqcond, "_num_reps", sep="")
activeflag <- paste(seqcond, "flag_analyze", sep="_")
headers <- paste(c("FEATURE_ID", seqreps, countheader, allpriors, activeflag), collapse="\t")
print(unlist(strsplit(headers,"\t")))
cat(headers, file=out.file, append=FALSE, sep="\n")

## End of arguments and names copied from the scripts used for the analysis of real data

##################################
##
## Simulation
##
##################################
print(simruns)
for (aaa in 1:simruns)
{
  ## Parameters of true model
  # true_betas_row1 <- c(1, 1.2, 0.7)*myreads   #These are the biorep effets, the beta_i s
  true_betas_row1 <- rep(1, nreps)*myreads
  true_alpha <- my.alpha     ## if different from 1, AI in the environment
  true_gamma <- 1            ## The current model assumes 1, but it used to be useful before
  true_tau <- 1              ## The current model assumes 1, but at some point we considered it different from 1

  allq <- c(my.qtest, my.qline)

  true_phi <- 0.02           #Neg binomial dispersion parameter, the bigger it is the higher the variance=mu+phi mu^2 is

  flaganalyze <- 1

  means <- c(allq[1]/true_alpha, allq[2]*true_alpha, ((1-allq[1])/true_alpha+(1-allq[2])*true_alpha)*true_tau)
	for(i in 1:nreps)
  {
    xs <- rnbinom(1, size=1/true_phi, mu=means[1]*true_betas_row1[i])
		ys <- rnbinom(1, size=1/true_phi, mu=means[2]*true_betas_row1[i])
		zs <- rnbinom(1, size=1/true_phi, mu=means[3]*true_betas_row1[i])
    if(i==1) mycounts <- paste(xs, ys, zs, sep="\t") else mycounts <- paste(mycounts, xs, ys, zs, sep="\t")
	}

	cat("mycounts is", mycounts, "\n")

	out <- paste("fusion_id", nreps,
	          mycounts,
	          paste(as.vector(t(allq)), collapse="\t"), paste(flaganalyze, collapse="\t"), sep="\t")
	cat(out, file=out.file, append=TRUE, sep="\n")
}


library(tidyverse)
library(matrixStats)
library(data.table)
library(lubridate)
library(gdata)
library(tidyr)
library(EnvStats)
library(scales)
library(stringr)
library(abind)
library(optparse)

source("utils/mobility/read-data-subnational.r")

make_table <- function(filename, date_till_percentage){
  print(sprintf("Running %s up until %s",filename, date_till_percentage))
  load(filename)
  out <- rstan::extract(fit)

  ifrs <- read_ifr_data()
  #df_pop= read.csv("data/popt_ifr3.28.csv", stringsAsFactors = FALSE)
  #df_pop$country[df_pop$country == "United Kingdom"] = "United_Kingdom"

  dates_italy <- dates[[which(regions == "Lombardy")]]
  len_dates <- length(dates_italy)

  cases <- vector("list", length = length(regions))
  total_cases <- vector("list", length = length(regions))
  total_cases_ui <- vector("list", length = length(regions))
  total_cases_li <- vector("list", length = length(regions))
  deaths <- vector("list", length = length(regions))
  total_deaths <- vector("list", length = length(regions))
  rt <- vector("list", length = length(regions))
  fraction_infected <- vector("list", length = length(regions))
  fraction_infected_li <- vector("list", length = length(regions))
  fraction_infected_ui <- vector("list", length = length(regions))
  fraction_obs_infected <- vector("list", length = length(regions))
  fraction_total_obs_infected <- vector("list", length = length(regions))
  y <- vector("list", length = length(regions))

  for(i in 1:length(regions)) {
    Country = regions[i]
    print(Country)
    x = dates[[i]]
    N = length(x)
    forecast = 7
    x  = c(x,x[length(x)]+1:forecast)
    padding <- len_dates - length(dates[[i]])
    y[[i]] = c(rep(0, padding),reported_cases[[i]], rep(NA, forecast))

    cases[[i]] = c(rep(0, padding), round(colMeans(estimated_cases_raw[,1:length(x),i])))

    total_cases[[i]] = c( round(cumsum(colMeans(estimated_cases_raw[,1:length(x),i]))))
    # chk = c(round((colMeans(rowCumsums(estimated_cases_raw[,1:length(x),i])))))

    total_cases_li[[i]] = c(
      round((colQuantiles(rowCumsums(estimated_cases_raw[,1:length(x),i]),probs=.025))))
    total_cases_ui[[i]] = c(
      round((colQuantiles(rowCumsums(estimated_cases_raw[,1:length(x),i]),probs=.975))))

    deaths[[i]] =  c(rep(0, padding), round(colMeans(estimated_deaths_raw[,1:length(x),i])))
    total_deaths[[i]] =  c(rep(0, padding), round(cumsum(colMeans(estimated_deaths_raw[,1:length(x),i]))))
    rt[[i]] = c(rep(NA, padding), colMeans(out$Rt_adj[,1:length(x),i]))

    fraction_infected[[i]] = c(rep(0, padding), total_cases[[i]]/ ifrs$popt[ifrs$country==Country])
    fraction_infected_li[[i]] = c(rep(0, padding),
                                  total_cases_li[[i]]/ ifrs$popt[ifrs$country==Country])
    fraction_infected_ui[[i]] = c(rep(0, padding),
                                  total_cases_ui[[i]]/ ifrs$popt[ifrs$country==Country])
    fraction_obs_infected[[i]] = c(rep(0, padding), y[[i]] / cases[[i]])
    fraction_total_obs_infected[[i]] = c(rep(0, padding), cumsum(y[[i]]) / cases[[i]])

    total_cases[[i]]  = c(rep(0, padding),total_cases[[i]])
  }

  dates_italy  = c(dates_italy,dates_italy[length(dates_italy)]+1:forecast)
  cases <- do.call(rbind, cases)
  cases_df <- as.data.frame(cases)
  names(cases_df) <- dates_italy
  cases_df$regions <- regions
  # write.csv(cases_df, "figures/cases.csv")

  total_cases <- do.call(rbind, total_cases)
  total_cases_df <- as.data.frame(total_cases)
  names(total_cases_df) <- dates_italy
  total_cases_df$regions <- regions
  # write.csv(total_cases_df, "figures/total_cases.csv")

  deaths <- do.call(rbind, deaths)
  deaths_df <- as.data.frame(deaths)
  names(deaths_df) <- dates_italy
  deaths_df$regions <- regions
  # write.csv(deaths_df, "figures/deaths.csv")

  total_deaths <- do.call(rbind, total_deaths)
  total_deaths_df <- as.data.frame(total_deaths)
  names(total_deaths_df) <- dates_italy
  total_deaths_df$regions <- regions
  # write.csv(total_deaths_df, "figures/total_deaths.csv")

  rt <- do.call(rbind, rt)
  rt_df <- as.data.frame(rt)
  names(rt_df) <- dates_italy
  rt_df$regions <- regions
  # write.csv(rt_df, "figures/rt.csv")

  fraction_infected <- do.call(rbind, fraction_infected)
  fraction_infected_df <- as.data.frame(fraction_infected)
  names(fraction_infected_df) <- dates_italy
  fraction_infected_df$regions <- regions
  # write.csv(fraction_infected_df, "figures/fraction_infected.csv")

  fraction_infected_li <- do.call(rbind, fraction_infected_li)
  fraction_infected_li_df <- as.data.frame(fraction_infected_li)
  names(fraction_infected_li_df) <- dates_italy
  fraction_infected_li_df$regions <- regions
  # write.csv(fraction_infected_li_df, "figures/fraction_infected_li.csv")

  fraction_infected_ui <- do.call(rbind, fraction_infected_ui)
  fraction_infected_ui_df <- as.data.frame(fraction_infected_ui)
  names(fraction_infected_ui_df) <- dates_italy
  fraction_infected_ui_df$regions <- regions
  # write.csv(fraction_infected_ui_df, "figures/fraction_infected_ui.csv")

  total_infected = data.frame(regions=regions,mean=fraction_infected[,dates_italy == date_till_percentage],
                              li=fraction_infected_li[,dates_italy == date_till_percentage],ui=fraction_infected_ui[,dates_italy == date_till_percentage])

  ## Calculation for whole of italy.
  date = date_till_percentage
  infections = NULL
  for(i in 1:21) {
    today.ii = which(dates[[i]] == date)
    infections = cbind(infections, rowSums(out$prediction[, 1:today.ii,i]))
  }
  total.infections = NULL
  total.infections = c(total.infections,rowSums(infections))

  infections.by.date = mean(total.infections)
  lower.infections.by.date = quantile(total.infections,prob=.025)
  upper.infections.by.date = quantile(total.infections,prob=.975)

  # Read deaths data for regions
  d <- read_obs_data()
  regions<-unique(as.factor(d$country))

  # Read ifr
  ifr.by.country <- read_ifr_data(unique(d$country))
  popt <- ifr.by.country[which(ifr.by.country$country == "Italy"),]$popt

  ar <- infections.by.date/popt
  ar.lower <- lower.infections.by.date/popt
  ar.upper <- upper.infections.by.date/popt

  italy_total_infected = data.frame(regions = "Italy", mean = ar, li = ar.lower, ui = ar.upper)
  print(italy_total_infected)
  total_infected <- bind_rows(total_infected, italy_total_infected)

  total_infected$text = sprintf("%.02f%% [%.02f%%-%.02f%%]",
                                total_infected$mean*100,total_infected$li*100,total_infected$ui*100)

  print(total_infected[order(total_infected$regions),c("regions","text")])
  print(total_infected[order(-total_infected$mean),c("regions","text")])

  table<-total_infected[order(-total_infected$mean),c("regions","text")]
  return(table)
}

make_table_simulation <- function(filename, date_till_percentage){
  print(sprintf("Running %s up until %s",filename, date_till_percentage))
  load(filename)

  out <- rstan::extract(fit)
  ifrs <- read_ifr_data()
  #df_pop= read.csv("data/popt_ifr3.28.csv", stringsAsFactors = FALSE)
  #df_pop$country[df_pop$country == "United Kingdom"] = "United_Kingdom"

  dates_italy <- dates[[which(regions == "Lombardy")]]
  len_dates <- length(dates_italy)

  cases <- vector("list", length = length(regions))
  total_cases <- vector("list", length = length(regions))
  total_cases_ui <- vector("list", length = length(regions))
  total_cases_li <- vector("list", length = length(regions))
  deaths <- vector("list", length = length(regions))
  total_deaths <- vector("list", length = length(regions))
  rt <- vector("list", length = length(regions))
  fraction_infected <- vector("list", length = length(regions))
  fraction_infected_li <- vector("list", length = length(regions))
  fraction_infected_ui <- vector("list", length = length(regions))
  fraction_obs_infected <- vector("list", length = length(regions))
  fraction_total_obs_infected <- vector("list", length = length(regions))
  y <- vector("list", length = length(regions))

  for(i in 1:length(regions)) {
    Country = regions[i]
    print(Country)
    x = dates[[i]]
    padding <- len_dates - length(dates[[i]])

    cases[[i]] = c(round(colMeans(estimated_cases_raw[,1:length(x),i])))

    total_cases[[i]] = c( round(cumsum(colMeans(estimated_cases_raw[,1:length(x),i]))))
    # chk = c(round((colMeans(rowCumsums(estimated_cases_raw[,1:length(x),i])))))

    total_cases_li[[i]] = c(
      round((colQuantiles(rowCumsums(estimated_cases_raw[,1:length(x),i]),probs=.025))))
    total_cases_ui[[i]] = c(
      round((colQuantiles(rowCumsums(estimated_cases_raw[,1:length(x),i]),probs=.975))))

    deaths[[i]] =  c(round(colMeans(estimated_deaths_raw[,1:length(x),i])))
    total_deaths[[i]] =  c(round(cumsum(colMeans(estimated_deaths_raw[,1:length(x),i]))))
    rt[[i]] = c( colMeans(out$Rt_adj[,1:length(x),i]))

    fraction_infected[[i]] = c(total_cases[[i]]/ ifrs$popt[ifrs$country==Country])
    fraction_infected_li[[i]] = c(total_cases_li[[i]]/ ifrs$popt[ifrs$country==Country])
    fraction_infected_ui[[i]] = c(total_cases_ui[[i]]/ ifrs$popt[ifrs$country==Country])
    total_cases[[i]]  = c(total_cases[[i]])
  }

  cases <- do.call(rbind, cases)
  cases_df <- as.data.frame(cases)
  names(cases_df) <- dates_italy
  cases_df$regions <- regions
  # write.csv(cases_df, "figures/cases.csv")

  total_cases <- do.call(rbind, total_cases)
  total_cases_df <- as.data.frame(total_cases)
  names(total_cases_df) <- dates_italy
  total_cases_df$regions <- regions
  # write.csv(total_cases_df, "figures/total_cases.csv")

  deaths <- do.call(rbind, deaths)
  deaths_df <- as.data.frame(deaths)
  names(deaths_df) <- dates_italy
  deaths_df$regions <- regions
  # write.csv(deaths_df, "figures/deaths.csv")

  total_deaths <- do.call(rbind, total_deaths)
  total_deaths_df <- as.data.frame(total_deaths)
  names(total_deaths_df) <- dates_italy
  total_deaths_df$regions <- regions
  # write.csv(total_deaths_df, "figures/total_deaths.csv")

  rt <- do.call(rbind, rt)
  rt_df <- as.data.frame(rt)
  names(rt_df) <- dates_italy
  rt_df$regions <- regions
  # write.csv(rt_df, "figures/rt.csv")

  fraction_infected <- do.call(rbind, fraction_infected)
  fraction_infected_df <- as.data.frame(fraction_infected)
  names(fraction_infected_df) <- dates_italy
  fraction_infected_df$regions <- regions
  # write.csv(fraction_infected_df, "figures/fraction_infected.csv")

  fraction_infected_li <- do.call(rbind, fraction_infected_li)
  fraction_infected_li_df <- as.data.frame(fraction_infected_li)
  names(fraction_infected_li_df) <- dates_italy
  fraction_infected_li_df$regions <- regions
  # write.csv(fraction_infected_li_df, "figures/fraction_infected_li.csv")

  fraction_infected_ui <- do.call(rbind, fraction_infected_ui)
  fraction_infected_ui_df <- as.data.frame(fraction_infected_ui)
  names(fraction_infected_ui_df) <- dates_italy
  fraction_infected_ui_df$regions <- regions
  # write.csv(fraction_infected_ui_df, "figures/fraction_infected_ui.csv")

  total_infected = data.frame(regions=regions,mean=fraction_infected[,dates_italy == date_till_percentage],
                              li=fraction_infected_li[,dates_italy == date_till_percentage],ui=fraction_infected_ui[,dates_italy == date_till_percentage])
  total_infected$text = sprintf("%.02f%% [%.02f%%-%.02f%%]",
                                total_infected$mean*100,total_infected$li*100,total_infected$ui*100)
  print(total_infected[order(total_infected$regions),c("regions","text")])

  # write.csv(total_infected,pasteo("results/total_infected_",date_till_percentage,".csv"),row.names=F)

  return (total_infected )
}