# common functions used in RStan Tiger analysis

#' Generate group data for use in multilevel models 
#' @group_vector atomic vector of type logical, integer, numeric, or character
make_group = function(group_vector){
	valid_input = c("logical", "integer", "numeric","character")
	if(class(group_vector) %in% valid_input){
		N = length(group_vector)
		unique_group_names = unique(group_vector)
		n_group = length(unique_group_names)
		group = rep(NA, N)
		for(i in 1:n_group){
			group[group_vector==unique_group_names[i]] <- i
		}
		return(list('group_assignments' = group, 'n_group' = n_group, 'unique_group_names' = unique_group_names))
	} else{
		cat('error: group_vector must be an atomic vector of type:', valid_input)
	}
}


make_stan_data = function(data, response_name, predictor_names, group_names, offset_name = NULL, slope_names=NULL){
	n_levels = length(group_names)
	column_names = unique(c(response_name, predictor_names, group_names, offset_name, slope_names))
	d = data[complete.cases(data[, column_names]), column_names]
	result = list(N = nrow(d), n_predictors=length(predictor_names))
	level_info = list()
	if(!is.null(offset_name)){
		result[['offset']] = d[, offset_name]
	}
	if(!is.null(slope_names)){
		for(name in slope_names){
			result[[name]] = d[, name]
		}
	}
	result[['x']] = as.matrix(d[,predictor_names])
	result[['y']] = d[,response_name]
	for(i in 1:n_levels){
		group_info = make_group(as.vector(d[,group_names[i]]))
		result[[paste('n_level_', i, sep='')]] = group_info$n_group
		result[[paste('level_', i, '_assignments', sep='')]] = group_info$group_assignments
		level_info[[paste('level_', i, '_names', sep='')]] = group_info$unique_group_names
	}
	final_result = list('for_stan' = result, 'level_info' = level_info)
	return(final_result)
}


stanml = function(data, response_name, predictor_names, group_names, offset_name=NULL, slope_names=NULL, model_code, model_name='', machine_name='', save=TRUE, print=TRUE, plot=TRUE, traceplot=TRUE, dir='', nchains=1, iterations=2000, warmup=1000, pars=NA, include=TRUE){
	this_data = make_stan_data(data, response_name, predictor_names, group_names, offset_name, slope_names)
	this_stan_model = model_code 
	compiled_model <- stan(file = this_stan_model, data = this_data$for_stan, chains = 0)
	fit = stan(fit = compiled_model, data = this_data$for_stan, iter=iterations, warmup=warmup, chains = nchains, pars=pars, include=include)
	this_base_name = paste(dir, '/', model_name, '_', fit@model_name, '_', machine_name, '_', format(Sys.time(), "%Y-%m-%d-%H-%M-%S"), sep='')
	this_parameter_table = data.frame("parameter" = c(response_name, predictor_names, group_names), "code" = c('y', paste('b', 1:length(predictor_names), sep=""), paste('level', 1:length(group_names), sep="")))
	if(save){
		save(fit, this_data, file=paste(this_base_name, ".RData", sep=''))
		write.csv(this_parameter_table, file=paste(this_base_name, "_parameter_table.csv"))
	}
	if(print){
		sink(file=paste(this_base_name, "_printed.txt", sep=''))
		print(fit)
		sink()
	}
	if(plot){
		pdf(file=paste(this_base_name, "_plots.pdf", sep=''))
		plot(fit, ask=FALSE)
		dev.off()
	}
	if(traceplot){	
		for(i in seq(1, length(names(fit)), 10)){
				end_index = min((i+9), length(names(fit)))
				pdf(paste0(this_base_name, 'traceplot_', i, '_', end_index, '.pdf'))
				traceplot(fit, inc_warmup = TRUE, pars = names(fit)[i:end_index], nrow=5)
				dev.off()
		}
	}
	return(list('fit'=fit, 'this_data'=this_data))
}

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)



