### SHOULD BE RUN ON CLUSTER ###
## This requires as much memory as possible, so we ran it on only 10 cores but higher mem.

rm(list=ls())
library(foreign)
library(rstanarm)
library(rstan)
library(Hmisc)
library(data.table)
library(parallel)

options("stringsAsFactors" = FALSE) 

setwd('~/research/gTiger/Article_citsci/open_code')
source('settings/parameters.r')
source('settings/paths.r')
source('settings/functions.r')

setwd(path_to_data)

ncores = 10

rstan_options(auto_write = TRUE)
options(mc.cores = ncores)

Apr1Julian = as.POSIXlt("01-04-2015", format = "%d-%m-%Y")$yday
Dec1Julian = as.POSIXlt("01-12-2015", format = "%d-%m-%Y")$yday

Sep1Julian = as.POSIXlt("01-09-2015", format = "%d-%m-%Y")$yday
Sep30Julian = as.POSIXlt("30-09-2015", format = "%d-%m-%Y")$yday

std2 = function(x) (x-mean(x,na.rm=T))/(2*sd(x, na.rm=T))
invstd2 = function(x, d) (x*2*sd(d, na.rm=T))+mean(d,na.rm=T)
p2o = function(x) x/(1-x)


std2_d = function(x, d) (x-mean(d, na.rm=T))/(2*sd(d, na.rm=T))

d = readRDS('ovitrap_model_dataset.Rds')

d$trials = 1

M = readRDS('untracked/ovitrap_model.Rds')

modeled_cells = unique(d$TigacellID)
day_range = 1:365
dp = data.table(TigacellID = rep(modeled_cells, length(day_range)), sea_days = unlist(lapply(day_range, function(d) rep(d, length(modeled_cells)))))

dp$sea_days_std2 = std2_d(dp$sea_days, d$sea_days)
dp$sea_days_std2_sq = dp$sea_days_std2^2
dp$sea_days_std2_cu = dp$sea_days_std2^3

Expos_days_std2 = std2_d(14, d$Expos_days)

dp$Expos_days_std2 = Expos_days_std2

ntrials = 1000
sample_size = 1000

dp$trials = ntrials

dp$Egg_pres = 0

dp$Trap_ID = "_NEW_Trap_ID"

mf = readRDS('tigacell_muni_fracs_ine.Rds')

### daily predictions 1 April - 1 December
selected_sea_days = Apr1Julian:Dec1Julian



M_preds_dt = readRDS('untracked/O_preds_dt.Rds')


this_M_preds_dt = M_preds_dt[sea_days==selected_sea_days[1]]
mf_provmuns = unique(mf$PROVMUN)
full_sampling_frame = merge(mf, this_M_preds_dt, by.x="TigacellID", by.y="rn", all.x=TRUE, all.y=FALSE, allow.cartesian=TRUE)

full_sampling_frame$rn = 1:nrow(full_sampling_frame)

muni_masks = mclapply(1:length(mf_provmuns), function(i){
	print(i)
	flush.console()
	this_provmun = mf_provmuns[i]
	this_sampling_frame = full_sampling_frame[PROVMUN==this_provmun]
	if(is.null(this_sampling_frame) | nrow(this_sampling_frame)==0){
		return(NULL)}
		if(nrow(this_sampling_frame)>1){
			this_sample_rows = sample(x=this_sampling_frame$rn, size=sample_size, replace=TRUE, prob=this_sampling_frame$frac)
		} else{
			this_sample_rows = rep(1, sample_size)
		}
		this_sample_cols = sapply(1:sample_size, function(j) sample(x=1:sample_size, size=1))
		sample_mask = data.table(rn=this_sample_rows, V=this_sample_cols, PROVMUN=this_provmun)
		return(sample_mask)
	}, mc.cores=ncores)

mm = rbindlist(muni_masks)

full_mask = as.matrix(mm[, 1:2, with=F])

muni_post_list = mclapply(selected_sea_days, function(i){
	print(i)
	flush.console()
	this_M_preds_dt = M_preds_dt[sea_days==i]
	mf_provmuns = unique(mf$PROVMUN)
	full_sampling_frame = merge(mf, this_M_preds_dt, by.x="TigacellID", by.y="rn", all.x=TRUE, all.y=FALSE)
	full_sampling_frame$rn = 1:nrow(full_sampling_frame)
	full_mat = as.matrix(full_sampling_frame[ , grep("^V", names(full_sampling_frame)), with=FALSE])
	vector_result = full_mat[full_mask]
	mat_result = matrix(vector_result, 	length(mf_provmuns), sample_size, byrow=TRUE) # this builds a matrix with a row of 1000 draws from the sampling cell posteriors for each muni
	result_means = rowMeans(mat_result, na.rm=TRUE)/ntrials
	result_dt = data.table(tiga_prob=result_means, PROVMUN=mf_provmuns, sea_days=i)
	return(result_dt)
}, mc.cores=ncores)

muni_probs = rbindlist(muni_post_list)

saveRDS(muni_probs, file='ovi_muni_daily_probs.Rds')
