#!/usr/bin/env Rscript

#===========================================================
# R script for analyzing the statistical significance of interactions
# this is an implementation of the paper Ay et. al. 2014 (FitHiC)
# incorporates spline fit for raw contact count
# also implements the bias calculation and the bias based filtering and spline fit procedure

#Author: Sourya Bhattacharyya
#Vijay-Ay lab, LJI
#===========================================================

library(splines)
library(fdrtool)
library(parallel)	# library for parallel processing

library(optparse)

# #===========================
# # parallel computation of binomial distribution on input vector
# parallel.BinomDistr <- function(inpvec) {
# 	idx <- inpvec[1]
# 	cc <- inpvec[2]
# 	tc <- inpvec[3]
# 	currprob <- inpvec[4]

# 	curr_dbinom <- dbinom(cc, size=tc, prob=currprob)
# 	curr_pbinom <- pbinom(cc, size=tc, prob=currprob, lower.tail=FALSE)
# 	return c(idx, curr_dbinom, (curr_dbinom + curr_pbinom))
# }

# #===========================

option_list = list(
  	make_option(c("--InpFile"), type="character", default=NULL, help="File with interactions + normalization features among individual genomic bins, preferably sorted with respect to interaction distance"),
  	make_option(c("--headerInp"), type="logical", action="store_true", default=FALSE, help="If TRUE, input interaction file has the header. Default FALSE."),
	make_option(c("--OutFile"), type="character", default=NULL, help="Output file name storing interactions + probability, P value and Q value"),
	make_option(c("--P2P"), type="integer", action="store", default=0, help="Can be 0 or 1. If 1, spline fit is modeled using only peak to peak interactions. Default 0."),
  	make_option(c("--EqLenBin"), type="logical", action="store_true", default=FALSE, help="If TRUE, Equal length bins are used. Default FALSE, means that equal occupancy bins are used. Users need not alter this parameter"),
  	make_option(c("--BiasCorr"), type="integer", action="store", default=0, help="If 0, raw contact count is used. Else, bias corrected contact count is used before applying FitHiC. Default 0."),
	make_option(c("--BiasLowThr"), type="numeric", default=0.2, help="Lower threshold of bias. Default 0.2"),
	make_option(c("--BiasHighThr"), type="numeric", default=5, help="Higher threshold of bias. Default 5"),
  	make_option(c("--nbins"), type="integer", action="store", default=200, help="Number of bins employed for FitHiC.", metavar="nbins"),
  	make_option(c("--Draw"), type="logical", action="store_true", default=FALSE, help="If TRUE, spline fit of FitHiC is plotted. Default FALSE."),
  	make_option(c("--cccol"), type="integer", action="store", default=7, help="Column number of the file storing the contact count"),
  	make_option(c("--TimeFile"), type="character", default=NULL, help="If specified, denotes the file which will contain time profiling information"),
  	make_option(c("--BiasFilt"), type="integer", action="store", default=0, help="If 1, interactions are filtered according to BiasLowThr and BiasHighThr, before processing. Default 0"),
	make_option(c("--ProbBias"), type="integer", action="store", default=0, help="If 1, probability values are multiplied with the bias values. Default 0.")
);

opt_parser = OptionParser(option_list=option_list);
opt = parse_args(opt_parser);

nbins <- as.integer(opt$nbins) # number of bins used in FitHiC
if (is.null(opt$TimeFile)) {
	timeprof <- 0
} else {
	timeprof <- 1
}

# peak to peak background usage option
Peak2PeakBackg <- as.integer(opt$P2P)
cat(sprintf("\n Peak to peak background usage for spline fit: %s ", Peak2PeakBackg))

# number of processors within the system
ncore <- detectCores()
cat(sprintf("\n Number of cores in the system: %s ", ncore))

# directory containing the input interaction file
inpdir <- dirname(opt$InpFile)

# directory to contain the spline fitted output interaction file
OutIntDir <- dirname(opt$OutFile)

# this directory will store the spline graph according to the model
outdir <- paste0(OutIntDir,'/Results')
system(paste('mkdir -p', outdir))

# this file stores the spline fitted model
if (opt$EqLenBin) {
	plotfile <- paste0(outdir,'/','EqLenBin_SplinePass1.pdf')
} else {
	plotfile <- paste0(outdir,'/','EqOccBin_SplinePass1.pdf')
}

if (timeprof == 1) {
	starttime <- Sys.time()
}

# load the interaction data matrix
# Note: the data has a header information
interaction.data <- read.table(opt$InpFile, header=opt$headerInp)
if (opt$headerInp == FALSE) {
	colnames(interaction.data) <- c("chr1", "s1", "e1", "chr2", "s2", "e2", "cc", "d1", "isPeak1", "Bias1", "mapp1", "gc1", "cut1", "d2", "isPeak2", "Bias2", "mapp2", "gc2", "cut2")
}

if (timeprof == 1) {
	endtime <- Sys.time()
	fp <- file(opt$TimeFile, open="a")
	outstr <- paste('\n Time to load the interaction data file in the R structure: ', (endtime - starttime))
	write(outstr, file=fp, append=T)
	close(fp)
}

cat(sprintf("\n Number of unfiltered interactions: %s ", nrow(interaction.data)))

# return if the number of interaction is 0
if (nrow(interaction.data) == 0) {
	cat(sprintf("\n No of interactions : 0 - exit !!!"))
	return()
}

if (ncol(interaction.data) < opt$cccol) {
	cat(sprintf("\n There is no contact count column or the formatting of the input file has problems - exit !!!"))
	return()
}

# vector of initial (prior) probabilities for contact counts
Prob_Val <- c()

# probability of the observed contact count for a single locus pair, from the spline fit
Spline_Binom_Prob_CC <- c()

# P-value of the observed contact count for a single locus pair
# binomial distribution + spline based estimation
Spline_Binom_P_Val_CC <- c()

# for each bin, count the no of distinct locus pairs
no_distinct_loci <- c()

# for each bin, no of observed contacts
NumContact <- c()

# for each bin, stores the average contact count per locus pair
avg_contact <- c()

# prior contact probability for a specific locus pair falling within a bin
prior_contact_prob <- c()

# average interaction distance for all locus pairs falling within a bin
avg_int_dist <- c()

#======================================================
# if normalization using bias values is enabled
# then discard the interactions where either loci has bias values outside specified thresholds
if ((opt$BiasCorr == 1) & (opt$BiasFilt == 1)) {

	# specifying the columns in the interaction data
	# where the bias information is provided
	bias1.col <- 10
	bias2.col <- 16

	if (ncol(interaction.data) >= bias2.col) {
		interaction.data <- interaction.data[which((interaction.data[,bias1.col] >= opt$BiasLowThr) & (interaction.data[,bias1.col] <= opt$BiasHighThr) & (interaction.data[,bias2.col] >= opt$BiasLowThr) & (interaction.data[,bias2.col] <= opt$BiasHighThr)), ]
		cat(sprintf("\n Bias specific filtering is enabled - Number of interactions where both loci satisfy bias criterion: %s ", nrow(interaction.data)))
	}
}

#======================================================
# absolute genomic distance for an interaction instance
gene.dist <- abs(interaction.data[,2] - interaction.data[,5])

# no of interactions pairs
numPairs <- length(gene.dist)
cat(sprintf("\n *** Total number of interactions: %s ", numPairs))

# total number of contacts for all the interactions
TotContact <- sum(interaction.data[,opt$cccol])

if (timeprof == 1) {
	starttime <- Sys.time()
}

#======================================================
# add - sourya
# we form the training data out of the total number of interactions
# if --P2P is enabled, training data consists of only peak to peak interactions
# else training data consists of the complete interaction set
if ((Peak2PeakBackg == 1) & (ncol(interaction.data) >= 15)) {
	TrainingData <- interaction.data[which((interaction.data[,9] == 1) & (interaction.data[,15] == 1)), ]
	# append the interaction distance as the last column of the training data
	TrainingData <- cbind(TrainingData, abs(TrainingData[,2] - TrainingData[,5]))
	# now sort the data with respect to increasing interaction distance
	TrainingData <- TrainingData[order( TrainingData[,ncol(TrainingData)] ), ]
	# remove the last column
	TrainingData <- TrainingData[,1:(ncol(TrainingData)-1)]
} else {
	TrainingData <- interaction.data
}

numTrainingPairs <- nrow(TrainingData)
if (numTrainingPairs == 0) {
	cat(sprintf("\n No training data for FitHiC spline fit - error - return !!! "))
	return()
}

cat(sprintf("\n *** Total number of training interactions: %s ", numTrainingPairs))

# interaction distance vector for the training data only
Trainingdata_GeneDist <- abs(TrainingData[,2] - TrainingData[,5])

# contact count sum for the training data
TotTrainingDataContact <- sum(TrainingData[,opt$cccol])

#=====================================================
# divide the genomic distance into b quantiles (number of bins)

# Note: for equal length bins (opt$EqLenBin is TRUE), each bin would have equal no of entries (pairs of loci)
# from the training data
# on the other hand, for equal occupancy bins (opt$EqLenBin is FALSE), each bin would have equal no of contact count
# from the training data
# In both cases, however, corresponding bin intervals will be variable

# error condition - sourya
# if the number of interactions is less than the argument 'nbins'
# then re-adjust the values of nbins and then assign the
# no of entries in each of the bins

if (opt$EqLenBin) {
	# no of entries (occupancies) for each bin
	# each bin have same no of locus pairs
	if (numTrainingPairs < nbins) {
		nbins <- numTrainingPairs
		nentry <- 1
	} else {
		nentry <- floor(numTrainingPairs / nbins)
	}
} else {
	# no of contacts for each bin
	# each bin have same no of contacts
	if (TotTrainingDataContact < nbins) {
		nbins <- numTrainingPairs
	}
	ncontactbin <- floor(TotTrainingDataContact / nbins)
}

#==============================================

if (opt$EqLenBin) {
	# equal length bin case (equal locus entries)
	for (bin_idx in (1:nbins)) {
		si <- (bin_idx - 1) * nentry + 1
		if (bin_idx < nbins) {
			ei <- si + nentry - 1
		} else {
			ei <- numTrainingPairs	# last read
		}
		no_distinct_loci[bin_idx] <- (ei-si+1)
		NumContact[bin_idx] <- sum(TrainingData[si:ei,opt$cccol])
		avg_contact[bin_idx] <- NumContact[bin_idx] / no_distinct_loci[bin_idx]
		prior_contact_prob[bin_idx] <- (avg_contact[bin_idx] / TotTrainingDataContact)
		avg_int_dist[bin_idx] <- mean(Trainingdata_GeneDist[si:ei])
	}
} else {
	# equal occupancy (contact count) bin - preferred

	# keeps track of total no of contact counts for a particular bin
	cumContactCount <- 0
	nelem <- 0
	ei <- 0
	bin_idx <- 0

	# keeps track of the no of different genomic distance values encountered within this bin
	# and also the no of elements (locus pairs) having that particular genomic distance
	ValGeneDist <- c()
	nelemSameGeneDist <- c()

	while (ei < numTrainingPairs) {
		si <- ei + 1
		curr_gene_dist <- abs(TrainingData[si,2] - TrainingData[si,5])
		idx_list <- which(Trainingdata_GeneDist == curr_gene_dist)
		nelem <- nelem + length(idx_list)
		ei <- ei + length(idx_list)
		cumContactCount <- cumContactCount + sum(TrainingData[si:ei, opt$cccol])

		ValGeneDist <- c(ValGeneDist, curr_gene_dist)
		nelemSameGeneDist <- c(nelemSameGeneDist, length(idx_list))

		if (cumContactCount >= ncontactbin) {
			# current cumulative contact exceeds the expected average contact per bin
			# and also, all the contacts with this specified genomic distance is covered
			# so, we fix this bin
			bin_idx <- bin_idx + 1
			no_distinct_loci[bin_idx] <- nelem
			NumContact[bin_idx] <- cumContactCount
			avg_contact[bin_idx] <- cumContactCount / nelem
			prior_contact_prob[bin_idx] <- (avg_contact[bin_idx] / TotTrainingDataContact)

			avg_int_dist[bin_idx] <- 0
			for (j in (1:length(ValGeneDist))) {
				avg_int_dist[bin_idx] <- avg_int_dist[bin_idx] + ((nelemSameGeneDist[j] * 1.0) / nelem) * ValGeneDist[j]
			}

			# reset the couners
			cumContactCount <- 0
		 	ValGeneDist <- c()
		 	nelemSameGeneDist <- c()
		 	nelem <- 0
		}
	}
	# now dump the data in a text file
	OutBinfile <- paste0(OutIntDir, '/', 'Bin_Info.log')
	write.table(cbind(avg_int_dist, no_distinct_loci, NumContact, avg_contact, prior_contact_prob), OutBinfile, row.names = FALSE, col.names = c("AvgDist","NumLoci","NumContact","AvgContact","PriorProb"), sep = "\t", quote=FALSE, append=FALSE)
}

cat(sprintf("\n after printing the bin log file"))

#=====================================
# now prepare a spline fit where x axis is the 'avg_int_dist'
# and y axis is the 'prior_contact_prob'
#=====================================

# # the first spline works according to the specified degree of freedom (16 - for testing)
# fit <- smooth.spline(avg_int_dist, prior_contact_prob, df=16)

# if (timeprof == 1) {
# 	endtime <- Sys.time()
# 	fp <- file(opt$TimeFile, open="a")
# 	outstr <- paste('\n Time for spline (df=16): ', (endtime - starttime))
# 	write(outstr, file=fp, append=T)
# 	close(fp)
# 	starttime <- Sys.time()
# }

# the spline works according to the cross validation principle
fit2 <- smooth.spline(avg_int_dist, prior_contact_prob, cv=TRUE)

if (timeprof == 1) {
	endtime <- Sys.time()
	fp <- file(opt$TimeFile, open="a")
	outstr <- paste('\n Time for spline (CV): ', (endtime - starttime))
	write(outstr, file=fp, append=T)
	close(fp)
	starttime <- Sys.time()
}

# # perform anti-tonic regression on the first spline
# fit.mr <- monoreg(fit$x, fit$y, type="antitonic")

# if (timeprof == 1) {
# 	endtime <- Sys.time()
# 	fp <- file(opt$TimeFile, open="a")
# 	outstr <- paste('\n Time for antitonic regression on spline (df = 16): ', (endtime - starttime))
# 	write(outstr, file=fp, append=T)
# 	close(fp)
# 	starttime <- Sys.time()
# }

# perform anti-tonic regression on the second spline
# antitonic = monotonic decreasing (as the probability decreases with increasing distance)
fit2.mr <- monoreg(fit2$x, fit2$y, type="antitonic")

if (timeprof == 1) {
	endtime <- Sys.time()
	fp <- file(opt$TimeFile, open="a")
	outstr <- paste('\n Time for antitonic regression on spline (CV): ', (endtime - starttime))
	write(outstr, file=fp, append=T)
	close(fp)
}

# # plot the data, if plotting is enabled via the parameter opt$Draw
# if (opt$Draw) {
# 	pdf(plotfile, width=14, height=10)
# 	plot(avg_int_dist, prior_contact_prob, cex=0.5, col="darkgrey", xlab="Average interaction distance", ylab="Prior contact probability")
# 	title("Smooth spline - antitonic regression - pass 1")
# 	lines(fit.mr, col="red", lwd=2)
# 	lines(fit2.mr, col="blue", lwd=2)
# 	legend("topright",legend=c("16 DF", paste(as.integer(fit2$df), "df (cv)")), col=c("red", "blue"), lty=1, lwd=2, cex=0.8)
# 	dev.off()
# }

#=========================
# now use this fitted spline
# to get predicted probability values for all distance values of the interactions
# in the unit of bin size employed in the current file)
binsize <- abs(interaction.data[1,2] - interaction.data[1,5])
TestDistVal <- seq(min(gene.dist), max(gene.dist), binsize)

# probability obtained from the spline
pp.fit2 <- predict(fit2, TestDistVal)

# use this computed probability and the input distance to fit a spline again
fit2_new <- smooth.spline(pp.fit2$x, pp.fit2$y, cv=TRUE)

# perform antitonic regression on this spline
fit2_new.mr <- monoreg(fit2_new$x, fit2_new$y, type="antitonic")

# plot the fitted spline and the smoothing regression
if (opt$Draw) {
	plotfile1 <- paste0(outdir,'/','EqOccBin_SplinePass1.pdf')
	pdf(plotfile1, width=10, height=8)
	plot(avg_int_dist, prior_contact_prob, cex=0.5, col="black", xlab="Average interaction distance", ylab="Prior contact probability", xlim=c(0, max(gene.dist)))
	lines(fit2$x, fit2$y, col="yellow", lwd=0.5)
	lines(fit2.mr$x, fit2.mr$yf, col="blue", lwd=0.5)
	lines(fit2_new$x, fit2_new$y, col="green", lwd=0.5)
	lines(fit2_new.mr$x, fit2_new.mr$yf, col="red", lwd=0.5)
	legend("topright",legend=c(paste("spline (fit2) - ", as.integer(fit2$df), "df (cv)"), "MR on spline (fit2.mr)", paste("5 Kb uniform spline (fit2_new) - ", as.integer(fit2_new$df), "df (cv)"), "MR on 5 Kb uniform spline (fit2_new.mr)"), col=c("yellow", "blue", "green", "red"), lty=1, lwd=1, cex=0.8)
	title("Smooth spline - antitonic regression - pass 1")
	dev.off()
}

# plot the fitted spline and the smoothing regression
# in log scale
if (opt$Draw) {
	plotfile1 <- paste0(outdir,'/','EqOccBin_SplinePass1_LOGScale.pdf')
	pdf(plotfile1, width=10, height=8)
	plot(avg_int_dist, log10(prior_contact_prob), cex=0.5, col="black", xlab="Average interaction distance", ylab="Prior contact probability (log)", xlim=c(0, max(gene.dist)))
	lines(fit2$x, log10(fit2$y), col="yellow", lwd=0.5)
	lines(fit2.mr$x, log10(fit2.mr$yf), col="blue", lwd=0.5)
	lines(fit2_new$x, log10(fit2_new$y), col="green", lwd=0.5)
	lines(fit2_new.mr$x, log10(fit2_new.mr$yf), col="red", lwd=0.5)
	legend("topright",legend=c(paste("spline (fit2) - ", as.integer(fit2$df), "df (cv)"), "MR on spline (fit2.mr)", paste("5 Kb uniform spline (fit2_new) - ", as.integer(fit2_new$df), "df (cv)"), "MR on 5 Kb uniform spline (fit2_new.mr)"), col=c("yellow", "blue", "green", "red"), lty=1, lwd=1, cex=0.8)
	title("Smooth spline - antitonic regression - pass 1 - log scale")
	dev.off()
}

# create a data frame to write the old spline and old MR regression results
old_sample.df <- cbind.data.frame(fit2$x, fit2$y, fit2.mr$x, fit2.mr$yf)
colnames(old_sample.df) <- c('fit2$x', 'fit2$y', 'fit2.mr$x', 'fit2.mr$yf')
write.table(old_sample.df, paste0(outdir,'/EqOccBin_OldSpline_Regression.txt'), row.names = FALSE, col.names = TRUE, sep = "\t", quote=FALSE, append=FALSE)

# create a data frame to write the old spline and old MR regression results
new_sample.df <- cbind.data.frame(fit2_new$x, fit2_new$y, fit2_new.mr$x, fit2_new.mr$yf)
colnames(new_sample.df) <- c('fit2_new$x', 'fit2_new$y', 'fit2_new.mr$x', 'fit2_new.mr$yf')
write.table(new_sample.df, paste0(outdir,'/EqOccBin_NewSpline_Regression.txt'), row.names = FALSE, col.names = TRUE, sep = "\t", quote=FALSE, append=FALSE)

#--------------------------------------------
# Paper: Ferhat Ay, 2014
# the probability is computed with respect to the spline plot
# which is then substituted to the binomial distribution, for P-value estimation

if (is.unsorted(gene.dist) == TRUE) {

	cat(sprintf("\n **** The interaction file is unsorted --- check ****"))

	# here the genomic distance is unsorted
	# so we employ the predict function and probability distribution computation
	# for the complete array

	# get the spline interpolation prediction from the already modeled smoothing spline object (CV spline)
	# with respect to the input genomic distance vector
	# the returned value is a component of 2 fields: (x,y)
	# where x is the input data and y is the predicted data

	# input will be the genomic distance
	# output is the corresponding contact probability

	# note: the probability was computed with respect to 'numPairs' placed as the denominator
	# accordingly we modify the binomial distribution equation

	# predict from the new spline fit - sourya
	# pp <- predict(fit2, gene.dist)
	pp <- predict(fit2_new, gene.dist)

	for (k in (1:numPairs)) {
		# probability for this particular interaction is in the y field
		p <- pp$y[k]
		if (opt$BiasCorr == 0) {
			# model binomial distribution with this probability
			Spline_Binom_Prob_CC[k] <- dbinom(interaction.data[k, opt$cccol], size=TotContact, prob=p)
			# compute the p value with respect to the modified distribution
			Spline_Binom_P_Val_CC[k] <- pbinom(interaction.data[k, opt$cccol], size=TotContact, prob=p, lower.tail=FALSE) + Spline_Binom_Prob_CC[k]
			# store the prior probability
			Prob_Val[k] <- p
		} else {
			# model binomial distribution with the product of probability, and the bias values of two segments
			if (opt$ProbBias == 1) {
				curr_prob <- p * interaction.data[k, bias1.col] * interaction.data[k, bias2.col]
			} else {
				curr_prob <- p
			}
			Spline_Binom_Prob_CC[k] <- dbinom(interaction.data[k, opt$cccol], size=TotContact, prob=curr_prob)
			# compute the p value with respect to the modified distribution
			Spline_Binom_P_Val_CC[k] <- pbinom(interaction.data[k, opt$cccol], size=TotContact, prob=curr_prob, lower.tail=FALSE) + Spline_Binom_Prob_CC[k]
			# store the prior probability
			Prob_Val[k] <- p
		}
	}
} else {

	if (timeprof == 1) {
		starttime <- Sys.time()
	}

	# here the genomic distance is sorted
	# so we selectively apply the predict function to only the distinct elements of the list
	uniq.idx <- order(gene.dist)[!duplicated(gene.dist)]

	if (timeprof == 1) {
		endtime <- Sys.time()
		fp <- file(opt$TimeFile, open="a")
		outstr <- paste('\n Time to find unique indices in genomic distance: ', (endtime - starttime))
		write(outstr, file=fp, append=T)
		close(fp)
	}

	if (timeprof == 1) {
		starttime <- Sys.time()
	}

	for (k in (1:length(uniq.idx))) {
		# all the locations from si to ei will have the same probability
		# since the gene distance value is the same
		si <- uniq.idx[k]
		if (k < length(uniq.idx)) {
			ei <- uniq.idx[k+1] - 1
		} else {
			ei <- numPairs	# last read
		}
		# compute the probability according to the genomic distance (of the start location - single element)
		# this probability value will be used for all the values within the interval (si to ei)

		# predict from the new spline fit - sourya
		# pp <- predict(fit2, gene.dist[si])
		pp <- predict(fit2_new, gene.dist[si])
		p <- pp$y

		#==========================
		# compute the probability distribution for the current interval
		#==========================
		if (opt$BiasCorr == 0) {
			curr_dbinom <- dbinom(interaction.data[si:ei, opt$cccol], size=TotContact, prob=p)
			curr_pbinom <- pbinom(interaction.data[si:ei, opt$cccol], size=TotContact, prob=p, lower.tail=FALSE)
			# append the distribution to the final vector
		 	Spline_Binom_Prob_CC <- c(Spline_Binom_Prob_CC, curr_dbinom)
		 	Spline_Binom_P_Val_CC <- c(Spline_Binom_P_Val_CC, (curr_dbinom + curr_pbinom))
		 	# store the prior probability as well
		 	Prob_Val[si:ei] <- p
		} else {
			# bias based probability correction

			# comment - sourya
			# for (idx in (si:ei)) {
			# 	# model binomial distribution with the product of probability, and the bias values of two segments
			# 	curr_prob <- p * interaction.data[idx, (opt$cccol + 1)] * interaction.data[idx, (opt$cccol + 3)]
			# 	curr_dbinom <- dbinom(interaction.data[idx, opt$cccol], size=TotContact, prob=curr_prob)
			# 	curr_pbinom <- pbinom(interaction.data[idx, opt$cccol], size=TotContact, prob=curr_prob, lower.tail=FALSE)
			# 	# append the distribution to the final vector
			#  	Spline_Binom_Prob_CC[idx] <- curr_dbinom
			#  	Spline_Binom_P_Val_CC[idx] <- (curr_dbinom + curr_pbinom)
			#  	# store the prior probability as well
		 		# Prob_Val[idx] <- curr_prob
			# }

			# add - sourya

			# cat(sprintf("\n *** Before executing Parallel from si: %s to ei: %s ", si, ei))

			# now apply this matrix on the parallel version of binomial distribution computation
			# the second argument 1 means that individual rows of matrix is processed
			# in the parallel computation
			result_binomdistr <- as.data.frame(parallel:::mclapply( si:ei , mc.cores = ncore , function(idx){
				cc <- interaction.data[idx, opt$cccol]
				if (opt$ProbBias == 1) {
					pr <- p * interaction.data[idx, bias1.col] * interaction.data[idx, bias2.col]
				} else {
					pr <- p
				}
				db <- dbinom(cc, size=TotContact, prob=pr)
				pb <- pbinom(cc, size=TotContact, prob=pr, lower.tail=FALSE)
				return(c(idx,pr,db,(db+pb)))
				} ))

			# cat(sprintf("\n *** rows: %s   columns: %s ", nrow(result_binomdistr), ncol(result_binomdistr)))

			# copy the results
			# columns : si to ei
			# rows: 4
			# for (i in (1:ncol(result_binomdistr))) {
			# 	idx <- result_binomdistr[1,i]
			# 	Prob_Val[idx] <- result_binomdistr[2,i]
			# 	Spline_Binom_Prob_CC[idx] <- result_binomdistr[3,i]
			# 	Spline_Binom_P_Val_CC[idx] <- result_binomdistr[4,i]
			# }
			Prob_Val[si:ei] <- as.double(result_binomdistr[2,1:ncol(result_binomdistr)])
			Spline_Binom_Prob_CC[si:ei] <- as.double(result_binomdistr[3,1:ncol(result_binomdistr)])
			Spline_Binom_P_Val_CC[si:ei] <- as.double(result_binomdistr[4,1:ncol(result_binomdistr)])

			# end add - Sourya
		}
	}

	if (timeprof == 1) {
		endtime <- Sys.time()
		fp <- file(opt$TimeFile, open="a")
		outstr <- paste('\n Time for Spline based distribution compute (p value) (sorted interaction file): ', (endtime - starttime))
		write(outstr, file=fp, append=T)
		close(fp)
	}
}

if (timeprof == 1) {
	starttime <- Sys.time()
}

# from the generated P values, obtain the Q value using BH correction
Spline_Binom_QVal <- p.adjust(Spline_Binom_P_Val_CC, method = "BH")

if (timeprof == 1) {
	endtime <- Sys.time()
	fp <- file(opt$TimeFile, open="a")
	outstr <- paste('\n Time to compute Q value (from P value) in spline: ', (endtime - starttime))
	write(outstr, file=fp, append=T)
	close(fp)
}

if (timeprof == 1) {
	starttime <- Sys.time()
}

# accumulate all results - also add header information
FinalData <- cbind(interaction.data, Prob_Val, Spline_Binom_Prob_CC, Spline_Binom_P_Val_CC, Spline_Binom_QVal)
colnames(FinalData) <- c(colnames(interaction.data), "p", "dbinom", "P-Value", "Q-Value")

#===================================
# add - sourya - debug
# scan through individual genomic distance
# and plot the number of contacts / interactions for each distance
qval_thr <- 0.01
uniq_gene_dist <- sort(unique(gene.dist))

# distance specific contact count and no of interactions
unfilt_cc_dist <- c()
filt_cc_dist <- c()
unfilt_p2p_cc_dist <- c()
filt_p2p_cc_dist <- c()
unfilt_int_dist <- c()
filt_int_dist <- c()
unfilt_p2p_int_dist <- c()
filt_p2p_int_dist <- c()

for (i in 1:length(uniq_gene_dist)) {
	curr_dist <- uniq_gene_dist[i]
	curr_dist_unfilt_int_idx <- which(abs(FinalData[,2] - FinalData[,5]) == curr_dist)
	curr_dist_filt_int_idx <- which((abs(FinalData[,2] - FinalData[,5]) == curr_dist) & (FinalData[,ncol(FinalData)] < qval_thr))

	if (ncol(interaction.data) >= 15) {
		curr_dist_unfilt_p2p_idx <- which((abs(FinalData[,2] - FinalData[,5]) == curr_dist) & (FinalData[,9] == 1) & (FinalData[,15] == 1))
		curr_dist_filt_p2p_idx <- which((abs(FinalData[,2] - FinalData[,5]) == curr_dist) & (FinalData[,9] == 1) & (FinalData[,15] == 1) & (FinalData[,ncol(FinalData)] < qval_thr))
	}

	unfilt_int_dist <- c(unfilt_int_dist, length(curr_dist_unfilt_int_idx))
	unfilt_cc_dist <- c(unfilt_cc_dist, sum(FinalData[curr_dist_unfilt_int_idx, 7]))
	filt_int_dist <- c(filt_int_dist, length(curr_dist_filt_int_idx))
	filt_cc_dist <- c(filt_cc_dist, sum(FinalData[curr_dist_filt_int_idx, 7]))

	if (ncol(interaction.data) >= 15) {
		unfilt_p2p_int_dist <- c(unfilt_p2p_int_dist, length(curr_dist_unfilt_p2p_idx))
		unfilt_p2p_cc_dist <- c(unfilt_p2p_cc_dist, sum(FinalData[curr_dist_unfilt_p2p_idx, 7]))
		filt_p2p_int_dist <- c(filt_p2p_int_dist, length(curr_dist_filt_p2p_idx))
		filt_p2p_cc_dist <- c(filt_p2p_cc_dist, sum(FinalData[curr_dist_filt_p2p_idx, 7]))
	}
}

x <- sum(unfilt_cc_dist)
unfilt_cc_dist_frac <- unfilt_cc_dist / (x * 1.0)

x <- sum(filt_cc_dist)
filt_cc_dist_frac <- filt_cc_dist / (x * 1.0)

x <- sum(unfilt_int_dist)
unfilt_int_dist_frac <- unfilt_int_dist / (x * 1.0)

x <- sum(filt_int_dist)
filt_int_dist_frac <- filt_int_dist / (x * 1.0)

if (ncol(interaction.data) >= 15) {

	x <- sum(unfilt_p2p_cc_dist)
	unfilt_p2p_cc_dist_frac <- unfilt_p2p_cc_dist / (x * 1.0)

	x <- sum(filt_p2p_cc_dist)
	filt_p2p_cc_dist_frac <- filt_p2p_cc_dist / (x * 1.0)

	x <- sum(unfilt_p2p_int_dist)
	unfilt_p2p_int_dist_frac <- unfilt_p2p_int_dist / (x * 1.0)

	x <- sum(filt_p2p_int_dist)
	filt_p2p_int_dist_frac <- filt_p2p_int_dist / (x * 1.0)

	out.df <- cbind.data.frame(uniq_gene_dist, unfilt_int_dist, unfilt_int_dist_frac, unfilt_cc_dist, unfilt_cc_dist_frac, unfilt_p2p_int_dist, unfilt_p2p_int_dist_frac, unfilt_p2p_cc_dist, unfilt_p2p_cc_dist_frac, filt_int_dist, filt_int_dist_frac, filt_cc_dist, filt_cc_dist_frac, filt_p2p_int_dist, filt_p2p_int_dist_frac, filt_p2p_cc_dist, filt_p2p_cc_dist_frac)

	colnamevec <- c("dist", "U-Num", "%U-Num", "U-CC", "%U-CC", "U-P2P-Num", "%U-P2P-Num", "U-P2P-CC", "%U-P2P-CC", "F-Num", "%F-Num", "F-CC", "%F-CC", "F-P2P-Num", "%F-P2P-Num", "F-P2P-CC", "%F-P2P-CC")
} else {
	out.df <- cbind.data.frame(uniq_gene_dist, unfilt_int_dist, unfilt_int_dist_frac, unfilt_cc_dist, unfilt_cc_dist_frac, filt_int_dist, filt_int_dist_frac, filt_cc_dist, filt_cc_dist_frac)
	colnamevec <- c("dist", "U-Num", "%U-Num", "U-CC", "%U-CC", "F-Num", "%F-Num", "F-CC", "%F-CC")
}

colnames(out.df) <- colnamevec
write.table(out.df, paste0(outdir, '/Count_Summary.txt'), row.names = FALSE, col.names = TRUE, sep = "\t", quote=FALSE, append=FALSE)

colorvec <- c("black", "yellow", "blue", "green", "red", "brown", "magenta", "cyan")

plotfile1 <- paste0(outdir,'/Count_Summary.pdf')
pdf(plotfile1, width=10, height=8)
plot(out.df[,1], out.df[,3], cex=0.5, col=colorvec[1], type="l", xlab="Distance", ylab="Percentage")
lines(out.df[,1], out.df[,5], col=colorvec[2], lwd=0.5)
lines(out.df[,1], out.df[,7], col=colorvec[3], lwd=0.5)
lines(out.df[,1], out.df[,9], col=colorvec[4], lwd=0.5)

if (ncol(interaction.data) >= 15) {
	lines(out.df[,1], out.df[,11], col=colorvec[5], lwd=0.5)
	lines(out.df[,1], out.df[,13], col=colorvec[6], lwd=0.5)
	lines(out.df[,1], out.df[,15], col=colorvec[7], lwd=0.5)
	lines(out.df[,1], out.df[,17], col=colorvec[8], lwd=0.5)
}

if (ncol(interaction.data) >= 15) {
	legendvec <- c(colnamevec[3], colnamevec[5], colnamevec[7], colnamevec[9], colnamevec[11], colnamevec[13], colnamevec[15], colnamevec[17])
} else {
	legendvec <- c(colnamevec[3], colnamevec[5], colnamevec[7], colnamevec[9])
}

legend("topright",legend=legendvec, col=colorvec, lty=1, lwd=1, cex=0.8)
title("% of different quantities")
dev.off()
# end add - sourya - debug
#===================================

# append the Spline distribution probability and corresponding P value as separate columns
# and write in a separate text file
temp_outfile <- paste0(OutIntDir, '/', 'temp_out.bed')
write.table(FinalData, temp_outfile, row.names = FALSE, col.names = TRUE, sep = "\t", quote=FALSE, append=FALSE)

# now sort the file contents and write that in the final specified output file
system(paste('sort -k1,1 -k2,2n -k5,5n', paste0('-k',opt$cccol,',',opt$cccol,'nr'), temp_outfile, '>', opt$OutFile))
# delete the temporary output file
system(paste('rm', temp_outfile))

# # sort the data according to the chromosome name, start positions of both the interacting segments
# # and finally with the decreasing contact count
# FinalData <- FinalData[order( FinalData[,1], FinalData[,2], FinalData[,5], -FinalData[,opt$cccol] ), ]

# # write the data with the header information
# write.table(FinalData, opt$OutFile, row.names = FALSE, col.names = TRUE, sep = "\t", quote=FALSE, append=FALSE)

if (timeprof == 1) {
	endtime <- Sys.time()
	fp <- file(opt$TimeFile, open="a")
	outstr <- paste('\n Time to write the interaction data (with Q value) for spline based distribution: ', (endtime - starttime))
	write(outstr, file=fp, append=T)
	close(fp)
}


psa <- function(A, B) {

  string_to_vector <- function(x) {
      len <- nchar(x)
      return(substring(x, 1:len, 1:len))
  }

  move_cost <- function() {
      x <- pointer[1]
      y <- pointer[2]
      vertical_gap_penalty <- if (semi_global_mode && pointer[1] == (length(a))) 1 else -1
      vertical_move <- psa_array[x,y - 1] +  if (y > 0) vertical_gap_penalty else -Inf
      horizontal_gap_penalty <- if (semi_global_mode && y == (length(b))) 1 else -1
      horizontalMove <- psa_array[x - 1,y] + if (x > 0) horizontal_gap_penalty else -Inf
      if (x > 1 && y > 1 && a[x - 1] == b[y - 1]) {
        diagonalScore <- 1
        diagonalMove <- psa_array[x - 1,y - 1] + diagonalScore
      }  else {
        diagonalScore <- -Inf
        diagonalMove <- diagonalScore
      }
      moves_cost <- c(vertical_move, horizontalMove, diagonalMove)
      move_types <- rbind(c(0, -1),   # vertical
                          c(-1, 0),   # horizontal
                          c(-1, -1))  # diagonal
      return(c(max(moves_cost), move_types[which.max(moves_cost),]))
  }

    a <-  string_to_vector(A)
    b <- string_to_vector(B)
    a_len <- length(a)
    b_len <- length(b)
    semi_global_mode <-  min(a_len, b_len)/max(a_len, b_len) < 0.6

    psa_array <- array(0, dim=c(a_len+1, b_len+1))
    psa_array[1,] <- if (!semi_global_mode) 0:-b_len
    psa_array[,1] <- if (!semi_global_mode) 0:-a_len

    for (i in 1:a_len+1) {
        for (j in 1:b_len+1) {
            pointer <- c(i, j)
            psa_array[pointer[1], pointer[2]] <- move_cost()[1]
        }
    }

    pointer <- c(a_len+1,b_len+1)
    optimal_move_cost <- psa_array[pointer[1], pointer[2]]
    alignment_path <- c()
    while ((pointer[1] > 2) || (pointer[2] > 2)) {
        best_move <- move_cost()
        move <- best_move[2:3]
        alignment_path <- rbind(alignment_path, move)
        pointer <- pointer + move
    }
    return(tail(c(psa_array),1))
}

psa("ACACT", "ACT")
psa("ATGCTAAGCATA", "TACGATTCGTAT")
psa("ATGCTAAGCATA", "TTAGCTAAGCATA")


#### Theoretical simulation for learning parking (such as parallel or perpendicular parking)
#### Usage: Just load and run parking_simulation.20220109.r in R
####   by dragging the R script file into R console,
####   or by executing source("parking_simulation.20220109.r") in R console.

### Author: Adam Yongxin Ye
### Version: 0.1.0 (2022-01-09)

library(grid)

`%.%` = function(x,y) paste0(x,y)

carGrob = function(x=0, y=0, wheel_angle=0, width=1, length=2, wheel_to_front=0.3, wheel_to_back=wheel_to_front, wheel_length=0.3, frame_lwd=1, wheel_lwd=3, col="black", default.units="cm"){
	if (!is.unit(x))  x <- unit(x, default.units)
	if (!is.unit(y))  y <- unit(y, default.units)
	if (!is.unit(width))  width <- unit(width, default.units)
	if (!is.unit(length))  length <- unit(length, default.units)
	if (!is.unit(wheel_to_front))  wheel_to_front <- unit(wheel_to_front, default.units)
	if (!is.unit(wheel_to_back))  wheel_to_back <- unit(wheel_to_back, default.units)
	if (!is.unit(wheel_length))  wheel_length <- unit(wheel_length, default.units)

	x1 = x - width*0.5
	x2 = x + width*0.5
	y1 = y - wheel_to_back
	y2 = y1 + length
	y3 = y2 - wheel_to_front
	frameGrob = rectGrob(x=x1 + width*0.5, y=y1 + length*0.5, width=width, height=length, gp=gpar(col=col, lwd=frame_lwd))
	wheel_axis_F = segmentsGrob(x1, y3, x2, y3, gp=gpar(col=col, lwd=frame_lwd))
	wheel_axis_R = segmentsGrob(x1, y, x2, y, gp=gpar(col=col, lwd=frame_lwd))
	wheel_dx = wheel_length*0.5*sinpi(wheel_angle/180)
	wheel_dy = wheel_length*0.5*cospi(wheel_angle/180)
	wheel_FL = segmentsGrob(x1-wheel_dx, y3+wheel_dy, x1+wheel_dx, y3-wheel_dy, gp=gpar(col=col, lwd=wheel_lwd))
	wheel_FR = segmentsGrob(x2-wheel_dx, y3+wheel_dy, x2+wheel_dx, y3-wheel_dy, gp=gpar(col=col, lwd=wheel_lwd))
	wheel_RL = segmentsGrob(x1, y+wheel_length*0.5, x1, y-wheel_length*0.5, gp=gpar(col=col, lwd=wheel_lwd))
	wheel_RR = segmentsGrob(x2, y+wheel_length*0.5, x2, y-wheel_length*0.5, gp=gpar(col=col, lwd=wheel_lwd))
	gTree(children = gList(frameGrob, wheel_axis_F, wheel_axis_R, wheel_FL, wheel_FR, wheel_RL, wheel_RR))
}

drawCarVP = function(x=0, y=0, angle=0, wheel_angle=0, width=1, length=2, wheel_to_front=0.3, wheel_to_back=wheel_to_front, wheel_length=0.3, frame_lwd=1, wheel_lwd=3, col="black", default.units="cm"){
	if (!is.unit(x))  x <- unit(x, default.units)
	if (!is.unit(y))  y <- unit(y, default.units)
	if (!is.unit(width))  width <- unit(width, default.units)
	if (!is.unit(length))  length <- unit(length, default.units)
	if (!is.unit(wheel_to_front))  wheel_to_front <- unit(wheel_to_front, default.units)
	if (!is.unit(wheel_to_back))  wheel_to_back <- unit(wheel_to_back, default.units)
	if (!is.unit(wheel_length))  wheel_length <- unit(wheel_length, default.units)

	dy_to_center = length*0.5 - wheel_to_back
	vp_x = x - dy_to_center * sinpi(angle/180)
	vp_y = y + dy_to_center * cospi(angle/180)
	carVP = viewport(x=vp_x, y=vp_y, width=width, height=length, angle=angle)
	carGr = carGrob(x=width*0.5, y=wheel_to_back, wheel_angle=wheel_angle, width=width, length=length, wheel_to_front=wheel_to_front, wheel_to_back=wheel_to_back, wheel_length=wheel_length, frame_lwd=frame_lwd, wheel_lwd=wheel_lwd, col=col, default.units=default.units)
	pushViewport(carVP)
	grid.draw(carGr)
	popViewport()
	invisible(carVP)
}

createCar = function(x=0, y=0, angle=0, wheel_angle=0, wheel_angle_limit=30, width=1, length=2, wheel_to_front=0.3, wheel_to_back=wheel_to_front, wheel_length=0.3, frame_lwd=1, wheel_lwd=3, col="black", default.units="cm"){
	ans = list(x=x, y=y, angle=angle, wheel_angle=wheel_angle, wheel_angle_limit=wheel_angle_limit, width=width, length=length, wheel_to_front=wheel_to_front, wheel_to_back=wheel_to_back, wheel_length=wheel_length, frame_lwd=frame_lwd, wheel_lwd=wheel_lwd, col=col, default.units=default.units)
	class(ans) = "car"
	invisible(ans)
}
drawCar = function(car){
	with(car, drawCarVP(x=x, y=y, angle=angle, wheel_angle=wheel_angle, width=width, length=length, wheel_to_front=wheel_to_front, wheel_to_back=wheel_to_back, wheel_length=wheel_length, frame_lwd=frame_lwd, wheel_lwd=wheel_lwd, col=col, default.units=default.units))
}

carStearWheel = function(car, new_wheel_angle){
	car$wheel_angle = new_wheel_angle
	if(car$wheel_angle > car$wheel_angle_limit){
		car$wheel_angle = car$wheel_angle_limit
	}
	if(car$wheel_angle < -car$wheel_angle_limit){
		car$wheel_angle = -car$wheel_angle_limit
	}
	invisible(car)
}
carStearWheelDelta = function(car, delta_wheel_angle){
	carStearWheel(car, car$wheel_angle + delta_wheel_angle)
}
carMove = function(car, distance){
	if(car$wheel_angle==0){
		car$x = car$x - distance*sinpi(car$angle/180)
		car$y = car$y + distance*cospi(car$angle/180)
	}else{
		radius = (car$length-car$wheel_to_front-car$wheel_to_back) / tanpi(car$wheel_angle/180)
		center_x = car$x - radius*cospi(car$angle/180)
		center_y = car$y - radius*sinpi(car$angle/180)
		angle_distance = distance / radius / pi * 180
		car$angle = car$angle + angle_distance
		car$x = center_x + radius*cospi(car$angle/180)
		car$y = center_y + radius*sinpi(car$angle/180)
	}
	invisible(car)
}

#carDemo = createCar(5, 5)
#carDemo %>% drawCar
#carDemo %>% carMove(5) %>% drawCar
#carDemo %>% carMove(5) %>% carStearWheel(30) %>% carMove(5) %>% drawCar

ToyotaCorolla_scale = c(4630, 1780, 2700)   # length, width, wheel base

ToyotaCorolla_scale = ToyotaCorolla_scale / 1500

restart_simulation = function(userCar_x=5, userCar_y=5, parkedCar_x=6.5, parkedCar_y=c(3,11), parkedCar_angle=0, roadEdge_x=7.3, wheel_length=0.4){
	parkedCar_list = list()
	num_parkedCar = max(c(length(parkedCar_x), length(parkedCar_y)))
	parkedCar_x = rep(parkedCar_x, length.out=num_parkedCar)
	parkedCar_y = rep(parkedCar_y, length.out=num_parkedCar)
	parkedCar_angle = rep(parkedCar_angle, length.out=num_parkedCar)
	for(k in 1:num_parkedCar){
		parkedCar_list[[k]] = createCar(parkedCar_x[k], parkedCar_y[k], angle=parkedCar_angle[k], length=ToyotaCorolla_scale[1], width=ToyotaCorolla_scale[2], wheel_to_front=(ToyotaCorolla_scale[1]-ToyotaCorolla_scale[3])/2, wheel_length=wheel_length)
	}

	userCar = createCar(userCar_x, userCar_y, col="brown", length=ToyotaCorolla_scale[1], width=ToyotaCorolla_scale[2], wheel_to_front=(ToyotaCorolla_scale[1]-ToyotaCorolla_scale[3])/2, wheel_length=wheel_length)
	speed = 0

	redraw = function(){
		grid.newpage()
		for(k in 1:num_parkedCar){
			drawCar(parkedCar_list[[k]])
		}
		drawCar(userCar)
		grid.text("speed=" %.% speed %.% " , wheel_angle=" %.% userCar$wheel_angle %.% " , car_angle=" %.% round(userCar$angle, 3), 0, 0, just=c(0,0))
		for(now_roadEdge_x in roadEdge_x){
			grid.segments(now_roadEdge_x, 0, now_roadEdge_x, 100, default.units="cm", gp=gpar(lwd=2))
		}
	}

	dt = 0.5
	t = 0
	while(TRUE){
		redraw()
		### ref: https://stackoverflow.com/questions/15272916/how-to-wait-for-a-keypress-in-r
		prompt_txt = "Press arrow key to adjust, or any other key to maintain speed, or ESC to end"
		should_exit = FALSE
		getGraphicsEvent(prompt=prompt_txt, consolePrompt=prompt_txt, onKeybd=function(key){
			if(key=="Up"){
				speed <<- round(speed + 0.1, 2)
			}else if(key=="Down"){
				speed <<- round(speed - 0.1, 2)
			}else if(key=="Left"){
				userCar <<- carStearWheelDelta(userCar, + 30/1.5*0.5)
			}else if(key=="Right"){
				userCar <<- carStearWheelDelta(userCar, - 30/1.5*0.5)
			}else if(key=="ctrl-[" || key=="\033" || key=="ctrl-C"){   # ESC
				should_exit <<- TRUE
			}else{
#				print(key)
			}
			1
		})
		if(should_exit){
			break
		}
		userCar = carMove(userCar, speed*dt)
		t = t + dt
#		Sys.sleep(0.01)
	}
}

### Add X11() for compatibility in MacOS
X11(type="cairo")

simulation_option = 1

if(simulation_option==1){

	## parallel parking
	restart_simulation(userCar_x=5, userCar_y=5, parkedCar_x=6.5, parkedCar_y=c(3,11), parkedCar_angle=0, roadEdge_x=7.3, wheel_length=0.4)

}else if(simulation_option==2){

	## perpendicular parking
	restart_simulation(userCar_x=5, userCar_y=5, parkedCar_x=6.5, parkedCar_y=c(5,8.1), parkedCar_angle=-90, roadEdge_x=c(9.2,2), wheel_length=0.4)

}
### ref: https://www.acko.com/car-guide/how-to-park-a-car-perfectly-easy-guide-to-parallel-parking/


Step2plus_Celltype_marker_inference_new<-function(tg_R1_lists,data_CORS_cancer,tg_R1_cut=8,tg_R1_list_stat,cell_type_enrich_cut=0.5,resolution_level=resolution_level0,hcutn0=40)
{
        #have two sets of parameters for BCV t tests
        #1. For testing total rank and determining adding the base of a new cell type (strong condition)
        #2. For the linear model based cell type specific marker selected (weak condition)
        #3. Parameter set includes: number of BCV bins, number of rounds for t test and msep cut
	  #4. Add a new step, select top K R1 bases by using hclust, and add R1 bases by a mixed score of (1) three combination linear dependency and (2) cell type uniqueness
        tg_R1_lists_st<-tg_R1_lists
        for(i in 1:length(tg_R1_lists_st))
        {
                tg_R1_lists_st[[i]]<-tg_R1_lists_st[[i]][1:min(length(tg_R1_lists_st[[i]]),tg_R1_cut)]
        }

        print("Compute_total_rank!")
        tg_all_genes<-c()
        for(i in 1:length(tg_R1_lists_st))
        {
                tg_all_genes<-c(tg_all_genes,tg_R1_lists_st[[i]])
        }
        tg_all_genes<-unique(tg_all_genes)
        tg_data_ccc<-data_CORS_cancer[tg_all_genes,]
        BCV_stat_c<-BCV_ttest3(tg_data_ccc,rounds=100,slice0=4,maxrank0=30,msep_cut=0.0001)
        dim_tt<-sum(BCV_stat_c[[1]]<0.01)
        print("Total Cell Dim")
	print( dim_tt)
        #remove the leave one out step and add a hclust pre-filtering step
	data_c<-data_CORS_cancer
	Base_all<-c()
	for(i in 1:length(tg_R1_lists_st))
	{
        tg_data_c<-data_c[tg_R1_lists_st[[i]],]
        cc<-svd(tg_data_c)$v[,1]
        ccc<-cor(cc,t(tg_data_c))
        if(mean(ccc)<0)
      {
                cc<--cc
        }
        Base_all<-rbind(Base_all,cc)
	}
	rownames(Base_all)<-names(tg_R1_lists_st)
	bbb_all0<-Base_all
	rownames(bbb_all0)<-1:nrow(bbb_all0)

	hcutn0<-min(hcutn0,nrow(bbb_all0))
	data_genes_all<-data_CORS_cancer[tg_all_genes,]
	hcutn1<-min(hcutn0+3,nrow(bbb_all0))
	hclust_R1_screen<-hclust_screen_top_bases(bbb_all0,data_genes_all,hcutn=hcutn1)
	ccc<-c()
	for(i in 1:length(hclust_R1_screen[[1]]))
	{
		aaa<-cor(t(hclust_R1_screen[[1]][[i]]))
		diag(aaa)<-0
		ccc<-c(ccc,max(aaa))
	}
	R_bases_hclust_top_correlations<-ccc
	hh<-hclust_R1_screen[[3]][[hcutn0]]
	Base_hclust_screen_result<-Base_screen(bbb_all0,hh,tg_R1_list_stat)
	Base_hclust_screen_selected<-as.numeric(Base_hclust_screen_result[[1]])

	tg_R1_lists_selected<-list()
	tg_R1_list_stat_selected<-list()
	nn<-c()
	for(i in 1:length(Base_hclust_screen_selected))
	{
		tg_R1_lists_selected[[i]]<- tg_R1_lists_st[[Base_hclust_screen_selected[[i]]]]
		tg_R1_list_stat_selected[[i]]<-tg_R1_list_stat[[Base_hclust_screen_selected[[i]]]][tg_R1_lists_selected[[i]],]
		cc<-apply(tg_R1_list_stat_selected[[i]],2,mean)
		nn<-c(nn,names(which(cc==max(cc))[1]))
	}
	names(tg_R1_lists_selected)<-nn
	tgs<-tg_R1_lists_selected

        print("Linking graph based cell type selection")
        ddd<-combn(length(tgs),3)
        comp_ids<-c()
        stat_all<-rep(0,length(tgs))
        stat_2total<-rep(0,length(tgs))
        names(stat_all)<-1:length(tgs)
        data_c<-data_CORS_cancer

	Base_all<-c()
	for(i in 1:length(tgs))
	{
        tg_data_c<-data_c[tgs[[i]],]
        cc<-svd(tg_data_c)$v[,1]
        ccc<-cor(cc,t(tg_data_c))
        if(mean(ccc)<0)
      {
                cc<--cc
        }
        Base_all<-rbind(Base_all,cc)
	}
	rownames(Base_all)<-names(tgs)

for(i in 1:ncol(ddd))
{
        tg_genes<-unique(c(tgs[[ddd[1,i]]],tgs[[ddd[2,i]]],tgs[[ddd[3,i]]]))
     	  pp<-sum(BCV_ttest2(data_c[tg_genes,],rounds=20,maxrank0=5)<0.001)
       if(pp==2)
        {
                        comp_ids<-c(comp_ids,i)
                        #print(i)
                  lm0<-lm(Base_all[ddd[1,i],]~Base_all[ddd[2,i],]+Base_all[ddd[3,i],]+0)
                  if(sum(is.na(coefficients(lm0)))==0)
                        {
                                stat_2total[ddd[,i]]<-stat_2total[ddd[,i]]+1
                        if((sign(coefficients(lm0))[1]==1)&(sign(coefficients(lm0))[2]==1))
                        {
                                stat_all[ddd[2,i]]<-stat_all[ddd[2,i]]+1
                                stat_all[ddd[3,i]]<-stat_all[ddd[3,i]]+1
                        }
                        if((sign(coefficients(lm0))[1]==-1)&(sign(coefficients(lm0))[2]==1))
                        {
                                stat_all[ddd[2,i]]<-stat_all[ddd[2,i]]+1
                                stat_all[ddd[1,i]]<-stat_all[ddd[1,i]]+1
                        }
                        if((sign(coefficients(lm0))[1]==1)&(sign(coefficients(lm0))[2]==-1))
                        {
                                stat_all[ddd[3,i]]<-stat_all[ddd[3,i]]+1
                                stat_all[ddd[1,i]]<-stat_all[ddd[1,i]]+1
                        }
                        }
       }
         if(i%%100==1)
        {
                print(i)
        }
}

stat_all0<-stat_all/stat_2total
scores10<-stat_all0
s1<-scores10[order(-scores10)]
ss<-rownames(Base_all)[order(-scores10)]
score_plus<-rep(0,length(ss))
names(score_plus)<-names(s1)
ss0<-unique(ss)
ss0_counts<-rep(1,length(ss0))
names(ss0_counts)<-ss0
for(i in 1:length(ss))
{
	score_plus[i]<-ss0_counts[ss[i]]
	ss0_counts[ss[i]]<-ss0_counts[ss[i]]+1
}
score_plus0<-(1/score_plus)^2
s2<-score_plus0+s1
s2<-s2[order(-s2)]

#names(tgs)[as.numeric(names(s2))]
#names(tgs)[selected_id]

total_dim<-dim_tt
N<-0
tt_selected_genes<-c()
selected_p<-c()
selected_c<-c()
tg_d_stat<-c()
selected_id<-c()
for(i in 1:length(s2))
{
        if(N<=total_dim)
        {
                tg_id_c<-as.numeric(names(s2)[i])
                print(c(tg_id_c,tg_id_c))
                if(length(selected_p)==0)
                {
                        selected_p<-tgs[[tg_id_c]]
                        tg_d_stat<-1
                        selected_id<-c(tg_id_c)
                        N<-N+1
                        print(N)
                        print(selected_id)
                }
                else
                {
                        selected_c<-c(selected_p,tgs[[tg_id_c]])
                        selected_c<-unique(selected_c)
                        tg_data_ccc0<-data_CORS_cancer[selected_c,]
                        BCV_stat_c0<-BCV_ttest3(tg_data_ccc0,rounds=100,slice0=2,maxrank0=30,msep_cut=0.0001)
                        tg_d_c<-sum(BCV_stat_c0[[1]]<0.001)#which(BCV_stat_c0[[1]]>0.001)[1]-1
                        print(tg_d_c)
                        if(tg_d_c>tg_d_stat[length(tg_d_stat)])
                        {
                                tg_d_stat<-c(tg_d_stat,tg_d_c)
                                selected_id<-c(selected_id,tg_id_c)
                                selected_p<-selected_c
                                N<-N+1
                                print(N)
                                print(selected_id)
                        }
                        print("")
                }
        }
}

tg_2badded_cells<-setdiff(unique(names(tgs)),unique(names(tgs)[selected_id]))
add_R1_base_cut<-list()
#add_R1_base_full<-list()
if(length(tg_2badded_cells)>0)
{
        ccc<-rep(0,length(tg_2badded_cells))
        names(ccc)<-tg_2badded_cells
        ttt<-s2
        for(i in 1:length(ttt))
        {
                if(sum(names(tgs)[as.numeric(names(ttt[i]))]==tg_2badded_cells))
                {
                        tg<-names(tgs)[as.numeric(names(ttt[i]))]
                        if(ccc[tg]==0)
                        {
                                add_R1_base_cut[[which(tg_2badded_cells==tg)]]<-tgs[[as.numeric(names(ttt[i]))]]
                                #add_R1_base_full[[which(tg_2badded_cells==tg)]]<-tg_R1_lists[[as.numeric(names(ttt[i]))]]
                                ccc[tg]<-1
                        }
                }
        }
}
names(add_R1_base_cut)<-tg_2badded_cells

#############################
nn<-rep("",length(tg_R1_list_stat_selected))
for(i in 1:length(tg_R1_list_stat_selected))
{
        if(max(tg_R1_list_stat_selected[[i]][min(tg_R1_cut,nrow(tg_R1_list_stat_selected[[i]])),resolution_level])>cell_type_enrich_cut)
        {
                ccc<-tg_R1_list_stat_selected[[i]][min(tg_R1_cut,nrow(tg_R1_list_stat_selected[[i]])),resolution_level]
                nn[i]<-names(which(ccc==max(ccc))[1])
        }
}
selected_cell_types<-intersect(resolution_level,unique(nn))

cellreso_R1_base_cut<-list()
#cellreso_R1_base_full<-list()

for(i in 1:length(selected_cell_types))
{
        cc1<-c()
        #cc2<-c()
        for(j in 1:length(nn))
        {
                if(nn[j]==selected_cell_types[i])
                {
                        cc1<-c(cc1,tgs[[j]])
                        #cc2<-c(cc2,tg_R1_lists[[j]])
                }
        }
        cellreso_R1_base_cut[[i]]<-unique(cc1)
        #cellreso_R1_base_full[[i]]<-unique(cc2)
}
names(cellreso_R1_base_cut)<-selected_cell_types
#names(cellreso_R1_base_full)<-selected_cell_types

#############################
supp_results<-list(tg_R1_lists_selected,tg_R1_list_stat_selected,selected_id,tg_d_stat,dim_tt,stat_all,stat_2total,hclust_R1_screen,Base_hclust_screen_selected,R_bases_hclust_top_correlations)
names(supp_results)<-c("tg_R1_lists_selected","tg_R1_list_stat_selected","selected_id","tg_d_stat","total_cell_type_number","stat_all","stat_2total","hclust_R1_screen","Base_hclust_screen_selected","R_bases_hclust_top_correlations")

selected_R1_base_cut<-list()
#selected_R1_base_full<-list()

for(i in 1:length(selected_id))
{
        selected_R1_base_cut[[i]]<-tgs[[selected_id[i]]]
        #selected_R1_base_full[[i]]<-tg_R1_lists[[selected_id[i]]]
}
names(selected_R1_base_cut)<-names(tgs)[selected_id]
#names(selected_R1_base_full)<-names(tgs)[selected_id]

R2_selected_cell_type_markers<-list(selected_R1_base_cut,add_R1_base_cut,cellreso_R1_base_cut,supp_results)
names(R2_selected_cell_type_markers)<-c("selected_R1_base_cut","add_R1_base_cut","cellreso_R1_base_cut","supp_results")
return(R2_selected_cell_type_markers)
}

hclust_screen_top_bases<-function(bbb_all0,data_genes_all,hcutn=80)
{
aaa<-data_genes_all
h<-hclust(dist(bbb_all0))
tg_trunc_list<-list()
tg_trunc_list_old<-list()
st<-0
trunc_RS_bases_all<-list()
RMSE_IM<-list()
hclust_info<-list()
hcutn0<-min(hcutn,nrow(bbb_all0))
for(i in 1:hcutn0)
{
	fff<-cutree(h,i)
	hclust_info[[i]]<-fff
	fff0<-unique(fff)
	tg_trunc_list_old<-tg_trunc_list
	tg_trunc_list<-list()
	#tg_trunc_list_info<-list()
	tg_trunc_RS_bases<-c()
	RMSE_table_c<-c()
	if(st==0)
	{
		for(j in 1:length(fff0))
		{
			tg_ccc<-names(which(fff==fff0[j]))
			tg_trunc_list[[j]]<-tg_ccc
			if(length(tg_ccc)==1)
        		{
                		tg_trunc_RS_bases<-rbind(tg_trunc_RS_bases,bbb_all0[tg_ccc,])
        		}
        		else
        		{
               		bbb_all0_c<-bbb_all0[tg_ccc,]
                		bbb_all0_cc<-t(svd(bbb_all0_c)$v)[1,]
				aaa0<-cor(t(t(bbb_all0_cc)),t(bbb_all0_c))
				aaa0<-aaa0[which(abs(aaa0)>0.6)]
				if(mean(aaa0)<0)
				{
					bbb_all0_cc<--bbb_all0_cc
				}
               		tg_trunc_RS_bases<-rbind(tg_trunc_RS_bases,bbb_all0_cc)
        		}

		}
		rownames(tg_trunc_RS_bases)<-1:nrow(tg_trunc_RS_bases)
		trunc_RS_bases_all[[i]]<-tg_trunc_RS_bases
		ccc<-c()
		for(j in 1:nrow(tg_trunc_RS_bases))
		{
        		ttt_c<-tg_trunc_RS_bases[j,]%*%t(tg_trunc_RS_bases[j,])/sum((tg_trunc_RS_bases[j,])^2)
        		ccc<-cbind(ccc,RMSE_row(aaa%*%ttt_c)/RMSE_row(aaa))
        	}
		RMSE_IM[[i]]<-ccc
		st<-1
	}
	else
	{
		ccc_old<-ccc
		for(j in 1:length(fff0))
		{
			tg_ccc<-names(which(fff==fff0[j]))
			tg_trunc_list[[j]]<-tg_ccc
			if(length(tg_ccc)==1)
        		{
                		tg_trunc_RS_bases<-rbind(tg_trunc_RS_bases,bbb_all0[tg_ccc,])
        		}
        		else
        		{
               		bbb_all0_c<-bbb_all0[tg_ccc,]
                		bbb_all0_cc<-t(svd(bbb_all0_c)$v)[1,]
				aaa0<-cor(t(t(bbb_all0_cc)),t(bbb_all0_c))
				aaa0<-aaa0[which(abs(aaa0)>0.6)]
				if(mean(aaa0)<0)
				{
					bbb_all0_cc<--bbb_all0_cc
				}
               		tg_trunc_RS_bases<-rbind(tg_trunc_RS_bases,bbb_all0_cc)
              	}
		}
		rownames(tg_trunc_RS_bases)<-1:nrow(tg_trunc_RS_bases)
		trunc_RS_bases_all[[i]]<-tg_trunc_RS_bases
		perturb_id<-hcut_update(tg_trunc_list,tg_trunc_list_old)
		ccc<-c()
		for(j in 1:nrow(tg_trunc_RS_bases))
		{
			if(sum(j==perturb_id[[2]])==0)
			{
				ccc<-cbind(ccc,ccc_old[,perturb_id[[3]][j]])
			}
			else
			{
				ttt_c<-tg_trunc_RS_bases[j,]%*%t(tg_trunc_RS_bases[j,])/sum((tg_trunc_RS_bases[j,])^2)
				ccc<-cbind(ccc,RMSE_row(aaa%*%ttt_c)/RMSE_row(aaa))
				#print(c(i,j))
			}
        }
		RMSE_IM[[i]]<-ccc
	}
}
if(hcutn0<nrow(bbb_all0))
{
N<-hcutn
tg_ids<-seq(hcutn+1,nrow(bbb_all0),by=30)
tg_ids<-c(tg_ids,nrow(bbb_all0))
for(i in 1:length(tg_ids))
{
	N<-N+1
	#print(c(N,tg_ids[i]))
	fff<-cutree(h,tg_ids[i])
	hclust_info[[N]]<-fff
	fff0<-unique(fff)
	tg_trunc_list<-list()
	tg_trunc_RS_bases<-c()
	RMSE_table_c<-c()
	st<-0
	if(st==0)
	{
		for(j in 1:length(fff0))
		{
			tg_ccc<-names(which(fff==fff0[j]))
			tg_trunc_list[[j]]<-tg_ccc
			if(length(tg_ccc)==1)
        		{
                		tg_trunc_RS_bases<-rbind(tg_trunc_RS_bases,bbb_all0[tg_ccc,])
        		}
        		else
        		{
               		bbb_all0_c<-bbb_all0[tg_ccc,]
                		bbb_all0_cc<-t(svd(bbb_all0_c)$v)[1,]
				aaa0<-cor(t(t(bbb_all0_cc)),t(bbb_all0_c))
				aaa0<-aaa0[which(abs(aaa0)>0.6)]
				if(mean(aaa0)<0)
				{
					bbb_all0_cc<--bbb_all0_cc
				}
               		tg_trunc_RS_bases<-rbind(tg_trunc_RS_bases,bbb_all0_cc)
        		}
		}
		rownames(tg_trunc_RS_bases)<-1:nrow(tg_trunc_RS_bases)
		trunc_RS_bases_all[[N]]<-tg_trunc_RS_bases
		ccc<-c()
		for(j in 1:nrow(tg_trunc_RS_bases))
		{
        		ttt_c<-tg_trunc_RS_bases[j,]%*%t(tg_trunc_RS_bases[j,])/sum((tg_trunc_RS_bases[j,])^2)
        		ccc<-cbind(ccc,RMSE_row(aaa%*%ttt_c)/RMSE_row(aaa))
        	}
		RMSE_IM[[N]]<-ccc
	}
}
}
ccc<-list(trunc_RS_bases_all,RMSE_IM,hclust_info)
names(ccc)<-c("Bases_all","RMSE_all","clust_info_all")
return(ccc)
}

Base_screen<-function(bbb_all0,hh,tg_R1_list_stat)
{
	tg_R1_list_stat0<-tg_R1_list_stat
	names(tg_R1_list_stat0)<-1:length(tg_R1_list_stat)
	ttt<-unique(hh)
	merge_list_info<-list()
	R1_selected_c<-c()
	for(i in 1:length(ttt))
	{
		#print(i)
		cc<-names(which(hh==ttt[i]))
		merge_list_info[[i]]<-list(cc,1)
		if(length(cc)==1)
		{
			R1_selected_c<-c(R1_selected_c,cc)
		}
		if(length(cc)==2)
		{
			ddd<-cor(t(bbb_all0[cc,]))
			#print(ddd)
			m1<-max(apply(tg_R1_list_stat0[[cc[1]]],2,mean))
			m2<-max(apply(tg_R1_list_stat0[[cc[2]]],2,mean))
			if(m1>=m2)
			{
				R1_selected_c<-c(R1_selected_c,cc[1])
			}
			else
			{
				R1_selected_c<-c(R1_selected_c,cc[2])
			}
			merge_list_info[[i]][[2]]<-ddd
		}
		if(length(cc)>2)
		{
			ddd<-cor(t(bbb_all0[cc,]))
			fff<-apply(ddd,1,mean)
			ggg<-which(fff==max(fff))
			if(length(ggg)==1)
			{
				R1_selected_c<-c(R1_selected_c,cc[ggg])
			}
			else
			{
				dd<-cc[ggg]
				mm<-c()
				for(j in 1:length(dd))
				{
					mm<-c(mm,max(apply(tg_R1_list_stat0[[dd[j]]],2,mean)))
				}
				ee<-dd[which(mm==max(mm))[1]]
				R1_selected_c<-c(R1_selected_c,ee)
			}
			merge_list_info[[i]][[2]]<-ddd
		}
	}
      return(list(R1_selected_c,merge_list_info))
}

hcut_update<-function(tg_trunc_list,tg_trunc_list_old)
{
	out_list<-rep(0,length(tg_trunc_list))
	for(i in 1:length(tg_trunc_list))
	{
		for(j in 1:length(tg_trunc_list_old))
		{
			if(length(intersect(tg_trunc_list[[i]],tg_trunc_list_old[[j]]))>0)
			{
				out_list[i]<-j
			}
		}
	}
	tg_old_id<-as.numeric(names(which(table(out_list)==2)))
	tg_new_id<-which(out_list==tg_old_id)
	return(list(tg_old_id,tg_new_id,out_list))
}


#author: Thomas Coleman
# Difference in Differences Regressions - Error Analysis and Graphs for data from Snow 1855

# See "Causality in the Time of Cholera" working paper at https://papers.ssrn.com/abstract=3262234
# and my John Snow project website, http://www.hilerun.org/econ/papers/snow

# This code is licensed under the BSD 2-Clause License, https://opensource.org/licenses/BSD-2-Clause

# This collect together some simple functions that produce a standard format graph, and are 'source'd
# into various notebooks

# **preperrdata** Before graphing, to prepare the data

#  * Takes in *fittedmodel* - a regression that has been already run. From this it extracts the necessary
#  parameters. Also *single*, a string which for "single" says that there is a single treatment effect.
#  * Calculates the 1849 and 1854 predicted counts and rates
#  * Calculates *approximate* 95% error bars around the predicted rates, based on whether the fitted model
#  is Poisson or Negative Binomial
#  * Produces an adjusted 1854 predicted rate, adjusting for the 1854 time effect and treatment effect, so
#  that it is comparable to the 1849 predicted rate (for purposes of plotting with error bars)

#  This function changes global data (the x1849 & x1854 dataframes) using the "<<-" instead of "<-"
#  assignment. This is poor programming style but I could not find another easy way of doing what I wanted.

# **plot2_worker** Plots actual vs predicted, with error bars around the predicted

#  * Actually does the plotting, given all the data as input arguments (sequence no. for sub-districts;
#  the actual mean or rate; predicted; the 2.5% and 97.5% points; title)
#  * btw, the hack for plotting error bars is from
#  https://stackoverflow.com/questions/13032777/scatter-plot-with-error-bars

# **plot2** Is a cover function which unpacks the actual versus predicted mean from the appropriate dataframe

# **plot3** Plots actual 1849, 1854 (adjusted for time & treatment effects), predicted, with error bars

# **plotcomp** Plots actual 1849 versus 1854, with error bars around actual 1849

# **ploterrbars** is NOT a function you should use - I use it to print out .pdf versions of the graphs I want to use

# 22-jul-2020 modify to create "adjusted rates" for both 1849 and 1854

preperrdata <- function(fittedmodel,single = "single",link="log",conf=.95,population="population") {       # This is not a good way to do this because I am
	                                                  # changing globals from within the function (using <<-
	                                                  # instead of <-)
# Function to prepare the data (error bars) for graphing
# The 2.5% and 97.5% confidence bands are calculated assuming either Poisson or Negative Binomial
# The rates are calculated by generating the counts up and down from the "expected" (from the fitted model)

	# Set the confidence levels (inputting 0.95 means lower leve .025, upper .975, total .05 / .95)
	xconfidl <- (1-conf)/2
	xconfidu <- 1-xconfidl
	expected <- predict(fittedmodel)
	xfamily <- family(fittedmodel)$family
	if (link == "log") {
		expected <- exp(expected)       # expected values
	}
	if (substr(xfamily,1,8) == "gaussian") {   # For OLS "rate" model need to convert from rates to counts
		expected[1:28] <- expected[1:28] * x1849[,population]
		expected[29:56] <- expected[29:56] * x1854[,population]
	}
	theta <- fittedmodel$theta
	x1849$predcount <<- expected[1:28]
	x1854$predcount <<- expected[29:56]
	x1849$predrate <<- 10000 * expected[1:28] / x1849[,population]
	x1854$predrate <<- 10000 * expected[29:56] / x1854[,population]
	if (xfamily == "poisson") {         # Get 95% confidence bands depending on model used (Poiss vs Neg Binom)
		x1849$limdn <<- 10000 * qpois(xconfidl,lambda=x1849$predcount) / x1849[,population]
		x1849$limup <<- 10000 * qpois(xconfidu,lambda=x1849$predcount) / x1849[,population]
		x1854$limdn <<- 10000 * qpois(xconfidl,lambda=x1854$predcount) / x1854[,population]
		x1854$limup <<- 10000 * qpois(xconfidu,lambda=x1854$predcount) / x1854[,population]
		x1849$limdnact <<- 10000 * qpois(xconfidl,lambda=x1849$deaths) / x1849[,population]
		x1849$limupact <<- 10000 * qpois(xconfidu,lambda=x1849$deaths) / x1849[,population]
	} else if (substr(xfamily,1,8) == "Negative"){
		x1849$limdn <<- 10000 * qnbinom(xconfidl,size=theta,mu=x1849$predcount) / x1849[,population]
		x1849$limup <<- 10000 * qnbinom(xconfidu,size=theta,mu=x1849$predcount) / x1849[,population]
		x1854$limdn <<- 10000 * qnbinom(xconfidl,size=theta,mu=x1854$predcount) / x1854[,population]
		x1854$limup <<- 10000 * qnbinom(xconfidu,size=theta,mu=x1854$predcount) / x1854[,population]
		x1849$limdnact <<- 10000 * qnbinom(xconfidl,size=theta,mu=x1849$deaths) / x1849[,population]
		x1849$limupact <<- 10000 * qnbinom(xconfidu,size=theta,mu=x1849$deaths) / x1849[,population]
	} else if (substr(xfamily,1,8) == "gaussian"){
		x3 <- predict(fittedmodel,interval="prediction")   # I think "confidence" is what I want and not "prediction"
		x1849$limdn <<- x3[1:28,2]
		x1849$limup <<- x3[1:28,3]
		x1854$limdn <<- x3[29:56,2]
		x1854$limup <<- x3[29:56,3]
		if (link == "log") {    # There is no "link" item for lm, so the only way I can tell if this model is run
		                      # in level or log form is to look at the predicted, and if it is negative presume it's in logs
			x1849$limdn <<- exp(x1849$limdn)
			x1849$limup <<- exp(x1849$limup)
			x1854$limdn <<- exp(x1854$limdn)
			x1854$limup <<- exp(x1854$limup)
		}
		x1849$limdn <<- 10000 * x1849$limdn
		x1849$limup <<- 10000 * x1849$limup
		x1854$limdn <<- 10000 * x1854$limdn
		x1854$limup <<- 10000 * x1854$limup
	}

	# Now adjust the 1854 predicted counts (and rates) for the time and treatment fixed effects -
	# effectively netting out the 1854 effect and making it 1849-equivalent
	# Adjust for the year effect only - this should make 1849 & 1854 comparable net of time effect
	if (link == "log") {
		xyr1854 <- (coef(fittedmodel)["year1854"])
		x49 <- exp(xyr1854)
		x54 <- exp(-xyr1854)
		x1849$rateadjyr <<- 10000 * (x1849$deaths*x49) / x1849[,population]
		x1854$rateadjyr <<- 10000 * (x1854$deaths*x54) / x1854[,population]
	} else {
		x49 <- (coef(fittedmodel)["year1854"])
		x54 <- -(coef(fittedmodel)["year1854"])
		x1849$rateadjyr <<- x1849$rate + x49
		x1854$rateadjyr <<- x1854$rate + x54
	}
	# Adjust the 1854 actual for the estimated Treatment Effect and the estimated Year Effect
	if (single == "single")  {   # model with single treatment effect
		xdegless <- (coef(fittedmodel)["supplierSouthwarkVauxhall_Lambeth:year1854"])
		xdegmore <- xdegless               # Make both treatment effects the same
		if (link == "log") {
			x49 <- x49 * exp((x1849$lambethdegree == "less_Lambeth")*xdegless)
			x49 <- x49 * exp((x1849$lambethdegree == "more_Lambeth")*xdegmore)
			x54 <- x54 * exp(-(x1854$lambethdegree == "less_Lambeth")*xdegless)
			x54 <- x54 * exp(-(x1854$lambethdegree == "more_Lambeth")*xdegmore)
			x1849$rateadj <<- 10000 * (x1849$deaths*x49) / x1849[,population]
			x1854$rateadj <<- 10000 * (x1854$deaths*x54) / x1854[,population]
		} else {
			x49 <- x49 + (x1849$lambethdegree == "less_Lambeth")*xdegless
			x54 <- x49 + (x1849$lambethdegree == "more_Lambeth")*xdegmore
			x54 <- x54 -(x1854$lambethdegree == "less_Lambeth")*xdegless
			x54 <- x54 -(x1854$lambethdegree == "more_Lambeth")*xdegmore
			x1849$rateadj <<- x1849$rate + x49
			x1854$rateadj <<- x1854$rate + x54
		}
	} else if (single == "two") {      # model with two treatment effects
		xdegless <- coef(fittedmodel)["lambethdegreeless_Lambeth:year1854"]
		xdegmore <- coef(fittedmodel)["lambethdegreemore_Lambeth:year1854"]
		if (link == "log") {
			x49 <- x49 * exp((x1849$lambethdegree == "less_Lambeth")*xdegless)
			x49 <- x49 * exp((x1849$lambethdegree == "more_Lambeth")*xdegmore)
			x54 <- x54 * exp(-(x1854$lambethdegree == "less_Lambeth")*xdegless)
			x54 <- x54 * exp(-(x1854$lambethdegree == "more_Lambeth")*xdegmore)
			x1849$rateadj <<- 10000 * (x1849$deaths*x49) / x1849[,population]
			x1854$rateadj <<- 10000 * (x1854$deaths*x54) / x1854[,population]
		} else {
			x49 <- x49 + (x1849$lambethdegree == "less_Lambeth")*xdegless
			x54 <- x49 + (x1849$lambethdegree == "more_Lambeth")*xdegmore
			x54 <- x54 -(x1854$lambethdegree == "less_Lambeth")*xdegless
			x54 <- x54 -(x1854$lambethdegree == "more_Lambeth")*xdegmore
			x1849$rateadj <<- x1849$rate + x49
			x1854$rateadj <<- x1854$rate + x54
		}
	} else { # model with continuous treatment (population proportions)
		xperc_lambeth54 <- coef(fittedmodel)["perc_lambeth54"]
		if (link == "log") {
			x49 <-  exp((x1849$perc_lambeth * xperc_lambeth54)) / x49
			x54 <- x54 * exp(-(x1854$perc_lambeth * xperc_lambeth54))
			x1849$rateadj <<- 10000 * (x1849$deaths*x49) / x1849[,population]
			x1854$rateadj <<- 10000 * (x1854$deaths*x54) / x1854[,population]
		} else {
			x49 <- (x1849$perc_lambeth * xperc_lambeth54) - x49
			x54 <- x54 -(x1854$perc_lambeth * xperc_lambeth54)
			x1849$rateadj <<- x1849$rate + x49
			x1854$rateadj <<- x1854$rate + x54
		}
	}

	return(xfamily)
}

# "Worker" function to plot mean, predicted, and error bars
plot2_worker <- function(yseq, xmean,xlimdn,xpred,xlimup,title,legposition="bottomright") {
	xplot <- plot(xmean, yseq,
	    xlim=range(c(xmean, xpred,xlimdn,xlimup)),
	    ylim=rev(range(yseq)), col="red",
	    main=title,xlab="Mortality rate actual (red filled) vs predicted (empty circle)",ylab="sub-district",
	    pch=19)
	lines(xpred, yseq, type="p",
	    xlim=range(c(xmean, xpred,xlimdn,xlimup)),
	    ylim=rev(range(yseq)),
	    pch=1)
	# horizontal error bars
	xplot <- arrows(xlimdn, yseq, xlimup, yseq, length=0.05, angle=90, code=3,lty=3)
# Ooops - this legend is not for "plot2" but for "plotcomp"
#	xplot <- legend(legposition, c("1849 red circle", "1854 blue triangle"), col = c(2, 4),
#       text.col = c("red","blue"), lty = c(0, 0), pch = c(19, 17),
#       merge = TRUE,bty="o",bg="white")    # The bty="o" should overwrite the background

	xplot
}
# "Cover" function that takes in the dataframe and unpacks
plot2 <- function(confidata,xsupplier,title,legposition="bottomright") {
	xmean <- subset(confidata,supplier==xsupplier)[,"rate"]
	xlimdn <- subset(confidata,supplier==xsupplier)[,"limdn"]
	xlimup <- subset(confidata,supplier==xsupplier)[,"limup"]
	xpred <- subset(confidata,supplier==xsupplier)[,"predrate"]
	yseq <- subset(confidata,supplier==xsupplier)[,"seq"]
#	xtitle <- paste(xsupplier,title)
	xplot <- plot2_worker(yseq,xmean,xlimdn,xpred,xlimup,title,legposition)
}

# 29-jul-20 Split out "plot3_worker"
plot3_worker <- function(yseq,xmean, xpred,xlimdn,xlimup,x1854adj,title,legposition="bottomright") {
	xplot <- plot(xmean, yseq,
	    xlim=range(c(xmean, xpred,xlimdn,xlimup,x1854adj)),
	    ylim=rev(range(yseq)), col="red",
	    main=title,xlab="Mortality per 10,000 population",ylab="sub-district",
	    pch=19)
	lines(xpred, yseq, type="p",
	    xlim=range(c(xmean, xpred,xlimdn,xlimup)),
	    ylim=rev(range(yseq)),
	    pch=1)
	lines(x1854adj, yseq, type="p",
	    xlim=range(c(xmean, xpred,xlimdn,xlimup)),
	    ylim=rev(range(yseq)), col="blue",
	    pch=17)

	# horizontal error bars
	xplot <- arrows(xlimdn, yseq, xlimup, yseq, length=0.05, angle=90, code=3,lty=3)
	xplot <- legend(legposition, c("1849 red circle", "1854 blue triangle", "predicted"), col = c(2, 4, 1),
       text.col = c("red","blue","black"), lty = c(0, 0, 3), pch = c(19, 17, 1),
       merge = TRUE,bty="o",bg="white")    # The bty="o" should overwrite the background
}

# 26-feb-20 add (optional) argument "rateadj" which can be set to either
#  - "rateadj" = 1854 rate adjusted for both year effect and treatment effect
#  - "rateadjyr" = 1854 adjusted only for year effect (not treaement effect)
# 22-jul-20 add (optional) argument "yearadj" which can be set to either 1849 or 1854
#  - "1849" - display base 1849, adjust 1854 back to 1849 by year effect and treatment
#  - "1854" - display base 1854, adjust 1849 forward to 1854by year (no adjustment) and treatment)
plot3 <- function(confidata1849,confidata1854,xsupplier,title,legposition="bottomright",rateadj="rateadj",yearadj="1849") {
	if (yearadj == "1849") {
		xmean <- subset(confidata1849,supplier==xsupplier)[,"rate"]
		xlimdn <- subset(confidata1849,supplier==xsupplier)[,"limdn"]
		xlimup <- subset(confidata1849,supplier==xsupplier)[,"limup"]
		xpred <- subset(confidata1849,supplier==xsupplier)[,"predrate"]
		yseq <- subset(confidata1849,supplier==xsupplier)[,"seq"]
		x1854adj <- subset(confidata1854,supplier==xsupplier)[,rateadj]
	}
	else {
		xmean <- subset(confidata1849,supplier==xsupplier)[,rateadj]
		xlimdn <- subset(confidata1854,supplier==xsupplier)[,"limdn"]
		xlimup <- subset(confidata1854,supplier==xsupplier)[,"limup"]
		xpred <- subset(confidata1854,supplier==xsupplier)[,"predrate"]
		yseq <- subset(confidata1854,supplier==xsupplier)[,"seq"]
		x1854adj <- subset(confidata1854,supplier==xsupplier)[,"rate"]
	}
	xplot <- plot3_worker(yseq,xmean, xpred,xlimdn,xlimup,x1854adj,title,legposition)
	# xplot <- plot(xmean, y,
	#     xlim=range(c(xmean, xpred,xlimdn,xlimup,x1854adj)),
	#     ylim=rev(range(y)), col="red",
	#     main=title,xlab="Mortality per 10,000 population",ylab="sub-district",
	#     pch=19)
	# lines(xpred, y, type="p",
	#     xlim=range(c(xmean, xpred,xlimdn,xlimup)),
	#     ylim=rev(range(y)),
	#     pch=1)
	# lines(x1854adj, y, type="p",
	#     xlim=range(c(xmean, xpred,xlimdn,xlimup)),
	#     ylim=rev(range(y)), col="blue",
	#     pch=17)

	# # horizontal error bars
	# xplot <- arrows(xlimdn, y, xlimup, y, length=0.05, angle=90, code=3,lty=3)
	# xplot <- legend(legposition, c("1849 red circle", "1854 blue triangle", "predicted"), col = c(2, 4, 1),
 #       text.col = c("red","blue","black"), lty = c(0, 0, 3), pch = c(19, 17, 1),
 #       merge = TRUE,bty="o",bg="white")    # The bty="o" should overwrite the background

}

# Plotting joint region for 1849 & 1854 with error bars for each
plotcomp <- function(confidata1849,confidata1854,xsupplier,title,legposition="bottomright") {
	xmean <- subset(confidata1849,supplier==xsupplier)[,"rate"]
	xlimdn <- subset(confidata1849,supplier==xsupplier)[,"limdnact"]
	xlimup <- subset(confidata1849,supplier==xsupplier)[,"limupact"]
	xpred <- subset(confidata1849,supplier==xsupplier)[,"predrate"]
	y <- subset(confidata1849,supplier==xsupplier)[,"seq"]
	xmean1854 <- subset(confidata1854,supplier==xsupplier)[,"rateadjyr"]
#	xlimdn1854 <- subset(confidata1854,supplier==xsupplier)[,"limdn"]
#	xlimup1854 <- subset(confidata1854,supplier==xsupplier)[,"limup"]
	xplot <- plot(xmean, y,
	    xlim=range(c(xmean, xmean1854,xlimdn,xlimup)),
	    ylim=rev(range(y)), col="red",
	    main=title,xlab="Mortality: 1849 red circle, 1854 blue diamond",ylab="sub-district",
	    pch=19)
	lines(xmean1854, y, type="p",
	    xlim=range(c(xmean, xmean1854,xlimdn,xlimup)),
	    ylim=rev(range(y)), col="blue",
	    pch=17)

	# horizontal error bars
	xplot <- arrows(xlimdn, y, xlimup, y, length=0.05, angle=90, code=3,,lty=3)
	xplot <- legend(legposition, c("1849 red circle", "1854 blue triangle"), col = c(2, 4),
       text.col = c("red","blue"), lty = c(0, 0), pch = c(19, 17),
       bty="o",bg="white")    # The bty="o" should overwrite the background

#	xplot <- arrows(xlimdn1854, y, xlimup1854, y, length=0.05, angle=90, code=3)
}

ploterrbars <- function(fittedmodel,plotname,single = "single") {       # This is not a good way to do this because I am
	                                                  # changing globals from within the function (using <<-
	                                                  # instead of <-)
	xfamily <- preperrdata(fittedmodel,single)  # this function modifies global data

	pdf(paste("../paper/figures/errbar_",plotname,"a.pdf",sep=""))
		plot2(x1849,"SouthwarkVauxhall",paste("First-12 Southwark-only ",xfamily," 1849 "))
	dev.off()
	pdf(paste("../paper/figures/errbar_",plotname,"b.pdf",sep=""))
		plot2(x1849,"SouthwarkVauxhall_Lambeth",paste("Next-16 Jointly-Supplied ",xfamily," 1849 "))
	dev.off()
	pdf(paste("../paper/figures/errbar_",plotname,"c.pdf",sep=""))
		plot3(x1849,x1854,"SouthwarkVauxhall",paste("First-12 Southwark-only ",xfamily," 1849vs1854 "))
	dev.off()
	pdf(paste("../paper/figures/errbar_",plotname,"d.pdf",sep=""))
		plot3(x1849,x1854,"SouthwarkVauxhall_Lambeth",paste("Next-16 Jointly-Supplied ",xfamily," 1849vs1854 "))
	dev.off()
	if (xfamily != "poisson") {         # Plot comparison for joint region only for Negative Binomial
		pdf(paste("../paper/figures/errbar_",plotname,"e.pdf",sep=""))
			plotcomp(x1849,x1854,"SouthwarkVauxhall",paste("First-12 Southwark-only ",xfamily," 1849vs1854 "))
		dev.off()
		pdf(paste("../paper/figures/errbar_",plotname,"f.pdf",sep=""))
			plotcomp(x1849,x1854,"SouthwarkVauxhall_Lambeth",paste("Next-16 Jointly-Supplied ",xfamily," 1849vs1854 "))
		dev.off()
	}

}


# *****************************************************************
# *************** Population Dynamics Modeling ********************
# ******************* UTR Lake Sturgeon ***************************
# ****************** Author: Dan Walker ***************************
# ****************** Began on 27/06/2017 **************************
# *****************************************************************

# FINAL CODE - see popdy_scratch for development and testing

setwd("C:/Users/danwa/OneDrive/Research/Lake Sturgeon/Chapter3/LST Data/LST ABundance")
set.seed(865)

# Data from RMark POPAN Jolly-Seber MLE solution:
# Total UTR population in 2011: 5643
# Weight class - above/below 2500 g
# Length class - <650 mm TL, 650 - 950 mm TL, >950 mm TL
# Phi(survival) for heavy.avg.length = ~100%
# Phi for skinny, average length = 0.42 (0.15 - 0.75 95% CI)

# Annual releases
release = read.csv("UTR.yearly.release.csv")
plot(release)
summary(release) # releases ~ N(8862, 3616.236^2)
# 124,069 fish released 2000-2013

# 20 year classes >15 y.o. (SLSWG) = stock until 2035
addtl.releases = 22

sim.releases = as.integer(rnorm(addtl.releases, 8862, 3616.236))
year = 2014:2035
sim.releases.df = data.frame(year, sim.releases)
colnames(sim.releases.df) = c("Year", "count.released")
full.sim.release = rbind(release, sim.releases.df)

full.sim.release[15,2]= 14615
full.sim.release[16,2]= 16308

# ************************************************************
# end of data entry necessary for Simulation procedures
# ************************************************************

write.csv(full.sim.release, "full.sim.release.csv")

plot(full.sim.release,
     pch = ifelse(full.sim.release$Year >= 2016, 1, 16))
title(main = "Number of Reintroduced Lake Sturgeon")

mean(full.sim.release$count.released) # 9,163 fish released per year on average
sum(full.sim.release$count.released) # 329,872 fish released total

# Catch data
library(readxl)
catch = read_excel("CaptureData2.xlsx", sheet = 3)
View(catch)
l_a = na.omit(data.frame(catch$TL.mm, catch$Age.caught))
colnames(l_a)<-c("TL", "Age")
View(l_a)

boxplot(TL~Age, data = l_a)
plot(TL~Age, data = l_a)
title(main = "Lake Sturgeon Length at Age")

# fit a cubic polynomial to the data NEEDS WORK
l_a_cube = lm(TL~poly(Age, 3), data = l_a)
summary(l_a_cube)

lm.test = data.frame(l_a$Age)
pred.cubic = predict(l_a_cube, x = lm.test, interval = 'confidence', level = 0.99)

plot(TL~Age, data = l_a)
title(main = "Lake Sturgeon Length at Age")
abline(l_a_cube)

lines(pred.cubic[, 1], col = "black", lwd = 3)

cubic = function(x){
  781.7 + 1047.7*(x) + 123.6*(x^2) - 445.9*(x^3)
}
x = l_a$Age

y = cubic(x)

plot(x, y)

# VB growth curve

# *************************************************
# ****************** PopDy ************************
# *************************************************

# initialize storage matrix, constants
sim.mat = matrix(nrow = 1000, ncol = 36)
S.full = as.vector(seq(from = 0.15, to = 0.75, by = 0.01))
Z.full = -log(S.full) #remove step if Z's are changed
i = seq(from = 1 , to = 1000, by = 1000)
m = matrix(0, nrow = 36, ncol = 36)
# m
stock.numbers = full.sim.release[ ,2]
diag(m)<- stock.numbers
m[upper.tri(m)]<- NA #upper triangular matrix = NA, lower = 0, diag = stocking #s

# *************************************************
# *************** Simulation  *********************
# *************************************************
setwd("C:/Users/danwa/OneDrive/Research/Lake Sturgeon/Chapter3/LST Data/LST ABundance")
set.seed(865)
release = read.csv("UTR.yearly.release.csv")

# Necessary functions:
# 1
fish.at.age = function(x){

  non.na.index = as.integer(min(which(!is.na(x$N)))) #find first year of stocking

  if(non.na.index < 36){
    x$t = x$t - non.na.index #adjust t to reflect first year
    x = x[complete.cases(x), ] #remove NAs
    for(d in 1:(nrow(x)-1)){
      x$N[d+1] = round(x$N[d]*exp((-x$Z[d]*x$t[d+1])))
    }
  } else {
    x$N = as.integer(rnorm(1, 8862, 3616.236))
    }
    return(x)
}

# 2
na.pad = function(x, lent){
  x[1:lent]
}

# 3
rotate = function(x){
  t(apply(x, 2, rev))
}

# Simulation - fixed N0, variable Z
# initialize storage matrix, constants
sim.mat = matrix(nrow = 1000, ncol = 36)
S.full = as.vector(seq(from = 0.15, to = 0.75, by = 0.01))
Z.full = -log(S.full) #remove step if Z's are changed
i = seq(from = 1, to = 1000, by = 1)
m = matrix(0, nrow = 36, ncol = 36)
# m
stock.numbers = full.sim.release[ ,2]
diag(m)<- stock.numbers
m[upper.tri(m)]<- NA #upper triangular matrix = NA, lower = 0, diag = stocking #s

for(i in 1:1000){ #iterate 1000x simulations

  single.sim.result = matrix(NA, nrow = 36, ncol = 36) # initialize storage dataframe for output
  lent = 36

  for(n in 1:ncol(m)){ #iterate over all columns

    test.col = m[ , n]
    N = as.integer(test.col)
    df = as.data.frame(N)
    df$Z = sample(Z.full, 36) #tweak value here to lower mortality rate
    df$t = seq(from = 1, to = 36, by = 1)
    x = df

    if(n <= 35){
      x1 = fish.at.age(x)
      x1.exp = as.vector(na.pad(x1$N, lent))
      single.sim.result[ , n] = x1.exp
    } else {
      x$N[36] = as.integer(rnorm(1, 8862, 3616.236))
      x$N[1] <- x$N[36]
      x$N[36] <- NA
      single.sim.result[ , n] = x$N
    }
  }

  rott = rotate(single.sim.result)
  sim.resul = diag(rott)
  sim.resul.t = t(sim.resul)

  sim.mat[i, ] = sim.resul.t

}

View(sim.mat)

write.csv(sim.mat, "Fixed_N0_Varying_Z.csv")

# **********************************************************
# ************ first complete run 20/07/2017 ***************
# **********************************************************
# **********************************************************

# RESET ENVIRONMENT BEFORE RUNNING SECOND SIM!!!

# **********************************************************
# Simulation with fixed N0, fixed Z
# **********************************************************
setwd("C:/Users/danwa/OneDrive/Research/Lake Sturgeon/Chapter3/LST Data/LST ABundance")
set.seed(865)
release = read.csv("UTR.yearly.release.csv")

# New central function:

fish.age.fixedZ = function(x){

  if(non.na.index < 36){
    x = x[complete.cases(x), ] #remove NAs
    for(d in 1:(nrow(x)-1)){
      x$N[d+1] = round(x$N[d]*exp((-x$Z[d]*x$t[d+1])))
    }
    # N.export = as.vector(x$N)
    return(x)
  }
  else {
    x$N = as.integer(rnorm(1, 8862, 3616.236))
  }
}
# 2
na.pad = function(x, lent){
  x[1:lent]
}

# 3
rotate = function(x){
  t(apply(x, 2, rev))
}

# initialize storage matrix, constants
sim.mat = matrix(nrow = 1000, ncol = 36)
S.full = as.vector(seq(from = 0.15, to = 0.75, by = 0.01))
Z.full = -log(S.full) #remove step if Z's are changed
i = seq(from = 1, to = 1000, by = 1)
m = matrix(0, nrow = 36, ncol = 36)
# m
stock.numbers = full.sim.release[ ,2]
diag(m)<- stock.numbers
m[upper.tri(m)]<- NA #upper triangular matrix = NA, lower = 0, diag = stocking #s

for(i in 1:1000){ #iterate 1000x simulations

  single.sim.result = matrix(NA, nrow = 36, ncol = 36) # initialize storage dataframe for output
  lent = 36

  for(n in 1:ncol(m)){ #iterate over all columns

    test.col = m[ , n]
    N = as.integer(test.col)
    df = as.data.frame(N)
    df$t = seq(from = 1, to = 36, by = 1)
    x = df
    non.na.index = as.integer(min(which(!is.na(x$N)))) #find first year of stocking
    x$t = x$t - non.na.index #adjust t to reflect first year
    if(n < 36){
      x = x[complete.cases(x), ]
    }
    x$Z = ifelse(x$t <= 3, -log(0.75), -log(0.99))

    if(n <= 35){
      x1 = fish.age.fixedZ(x)
      x1.exp = as.vector(na.pad(x1$N, lent))
      single.sim.result[ , n] = x1.exp
    } else {
      x$N[36] = as.integer(rnorm(1, 8862, 3616.236))
      x$N[1] <- x$N[36]
      x$N[36] <- NA
      single.sim.result[ , n] = x$N
    }

  }

  rott = rotate(single.sim.result)
  sim.resul = diag(rott)
  sim.resul.t = t(sim.resul)

  sim.mat[i, ] = sim.resul.t

}

View(sim.mat)
write.csv(sim.mat, file = "Fixed_Z_Fixed_N0.csv")

# *************************************************************
# simulation with varying N0, fixed Z
# *************************************************************
setwd("C:/Users/danwa/OneDrive/Research/Lake Sturgeon/Chapter3/LST Data/LST ABundance")
set.seed(865)
release = read.csv("UTR.yearly.release.csv")

fish.age.fixedZ = function(x){

  if(non.na.index < 36){
    x = x[complete.cases(x), ] #remove NAs
    for(d in 1:(nrow(x)-1)){
      x$N[d+1] = round(x$N[d]*exp((-x$Z[d]*x$t[d+1])))
    }
    # N.export = as.vector(x$N)
    return(x)
  }
  else {
    x$N = as.integer(rnorm(1, 8862, 3616.236))
  }
}
# 2
na.pad = function(x, lent){
  x[1:lent]
}

# 3
rotate = function(x){
  t(apply(x, 2, rev))
}

sim.mat = matrix(nrow = 1000, ncol = 36)
addtl.releases = 22
for(j in 1:1000){
  # create new simulated release numbers
  sim.releases = as.integer(rnorm(addtl.releases, 8862, 3616.236))
  year = 2014:2035
  sim.releases.df = data.frame(year, sim.releases)
  colnames(sim.releases.df) = c("Year", "count.released")
  sim.releases.df = rbind(release, sim.releases.df)

  sim.releases.df[15,2]= 14615
  sim.releases.df[16,2]= 16308

  m = matrix(0, nrow = 36, ncol = 36)
  stock.numbers = sim.releases.df[ ,2]
  diag(m)<- stock.numbers
  m[upper.tri(m)]<- NA #upper triangular matrix = NA, lower = 0, diag = stocking #s

  single.sim.result = matrix(NA, nrow = 36, ncol = 36) # initialize storage dataframe for output
  lent = 36

  for(n in 1:ncol(m)){ #iterate over all columns

    test.col = m[ , n]
    N = as.integer(test.col)
    df = as.data.frame(N)
    df$t = seq(from = 1, to = 36, by = 1)
    x = df
    non.na.index = as.integer(min(which(!is.na(x$N)))) #find first year of stocking
    x$t = x$t - non.na.index #adjust t to reflect first year
    if(n < 36){
      x = x[complete.cases(x), ]
    }
    x$Z = ifelse(x$t <= 3, -log(0.75), -log(0.99))

    if(n <= 35){
      x1 = fish.age.fixedZ(x)
      x1.exp = as.vector(na.pad(x1$N, lent))
      single.sim.result[ , n] = x1.exp
    } else {
      x$N[36] = as.integer(rnorm(1, 8862, 3616.236))
      x$N[1] <- x$N[36]
      x$N[36] <- NA
      single.sim.result[ , n] = x$N
    }

  }

  rott = rotate(single.sim.result)
  sim.resul = diag(rott)
  sim.resul.t = t(sim.resul)

  sim.mat[j, ] = sim.resul.t

}
View(sim.mat)
sim.mat[sim.mat<0] <- 0
write.csv(sim.mat, file = "Fixed_Z_Varying_N0.csv")

# *************************************************************
# simulation with varying N0, fixed Z, doubled reintroductions
# *************************************************************
setwd("C:/Users/danwa/OneDrive/Research/Lake Sturgeon/Chapter3/LST Data/LST ABundance")
set.seed(865)
release = read.csv("UTR.yearly.release.csv")

fish.age.fixedZ = function(x){

  if(non.na.index < 36){
    x = x[complete.cases(x), ] #remove NAs
    for(d in 1:(nrow(x)-1)){
      x$N[d+1] = round(x$N[d]*exp((-x$Z[d]*x$t[d+1])))
    }
    # N.export = as.vector(x$N)
    return(x)
  }
  else {
    x$N = as.integer(rnorm(1, (2*8862), 3616.236))
  }
}
# 2
na.pad = function(x, lent){
  x[1:lent]
}

# 3
rotate = function(x){
  t(apply(x, 2, rev))
}

sim.mat = matrix(nrow = 1000, ncol = 36)
addtl.releases = 22
for(j in 1:1000){
  # create new simulated release numbers
  sim.releases = as.integer(rnorm(addtl.releases, (8862*2), 3616.236))
  year = 2014:2035
  sim.releases.df = data.frame(year, sim.releases)
  colnames(sim.releases.df) = c("Year", "count.released")
  sim.releases.df = rbind(release, sim.releases.df)

  sim.releases.df[15,2]= 14615
  sim.releases.df[16,2]= 16308

  m = matrix(0, nrow = 36, ncol = 36)
  stock.numbers = sim.releases.df[ ,2]
  diag(m)<- stock.numbers
  m[upper.tri(m)]<- NA #upper triangular matrix = NA, lower = 0, diag = stocking #s

  single.sim.result = matrix(NA, nrow = 36, ncol = 36) # initialize storage dataframe for output
  lent = 36

  for(n in 1:ncol(m)){ #iterate over all columns

    test.col = m[ , n]
    N = as.integer(test.col)
    df = as.data.frame(N)
    df$t = seq(from = 1, to = 36, by = 1)
    x = df
    non.na.index = as.integer(min(which(!is.na(x$N)))) #find first year of stocking
    x$t = x$t - non.na.index #adjust t to reflect first year
    if(n < 36){
      x = x[complete.cases(x), ]
    }
    x$Z = ifelse(x$t <= 3, -log(0.75), -log(0.99))

    if(n <= 35){
      x1 = fish.age.fixedZ(x)
      x1.exp = as.vector(na.pad(x1$N, lent))
      single.sim.result[ , n] = x1.exp
    } else {
      x$N[36] = as.integer(rnorm(1, 8862, 3616.236))
      x$N[1] <- x$N[36]
      x$N[36] <- NA
      single.sim.result[ , n] = x$N
    }

  }

  rott = rotate(single.sim.result)
  sim.resul = diag(rott)
  sim.resul.t = t(sim.resul)

  sim.mat[j, ] = sim.resul.t

}
View(sim.mat)
sim.mat[sim.mat<0] <- 0
write.csv(sim.mat, file = "Fixed_Z_Varying_N0_2xReintro.csv")