library(OpasnetUtils)
library(ggplot2)

#########################
# Lifetable function	#
#########################

lifetable <- function(
		pop = data.frame(Birth = 2000, Result = 1000), # population by birth year.
		mort_frac = data.frame(Age = 0:100, Result = 0.02), # mortality by age.
		followup = 100, # follow-up time.
		keep = (0:20) * 5, # Which years to save? 
		starttime = 2000, # What year is the starting time or reference? Ignored if present in data.
		... # Ovariable evaluation extra parameters
) {
	# Check for ovariable input; Instead of having 4 different methods for all combinations we can check
	# individually in one function.
	
	if(class(mort_frac)=="ovariable"){
		if(nrow(mort_frac@output) == 0) mort_frac <- EvalOutput(mort_frac, ...)
		# Proceed to remove non-marginal columns
		temp <- result(mort_frac)
		mort_frac <- mort_frac@output[mort_frac@marginal]
		mort_frac$Time <- NULL
		mort_frac$Mort <- temp
	}
	
	if(class(pop)=="ovariable"){
		if(nrow(pop@output) == 0) pop <- EvalOutput(pop, ...)
		# Proceed to remove non-marginal columns
		temp <- result(pop)
		pop <- pop@output[pop@marginal]
		pop$Result <- temp
	}
	
	# Prepare input data for yearly loops of follow-up.
	
	if(!"Time" %in% colnames(pop)) {
		pop$Time <- starttime # create Time index if it does not exist
	} else {
		starttime <- min(pop$Time) # else take lowest value as starting point
	}
	
	temp_pop <- data.frame() # init
	
	out <- pop
	
	# Follow the population year by year.
	
	for(i in 0:followup) {
		
		# Take any newborn population.
		temp_pop_add <- pop[pop$Time == starttime + i,] 
		
		# Rbind newborn with leftover population from last iteration
		temp_pop <- rbind(temp_pop_add, temp_pop)
		
		temp_pop$Age <- temp_pop$Time - temp_pop$Birth # update age
		
		temp_pop <- merge(temp_pop, mort_frac, all.x = TRUE) # combine population and age-specific mortality
		temp_pop$Mort[is.na(temp_pop$Mort)] <- 1 # All die beyond the mortality table
		
		temp_pop$Result <- temp_pop$Result * (1 - temp_pop$Mort) # update the population size
		temp_pop <- temp_pop[colnames(pop)]
		temp_pop$Time <- temp_pop$Time + 1 # update time
		
		# Rbind wanted years with output
		if(i %in% keep) {
			out <- rbind(out, temp_pop)
		}
	}
	return(out)
}

#################################
# Ovariable Definitions		#
#################################

Population <- Ovariable(name = "Population", ddata = "Op_en5994.population")
Population@data$Age <- as.integer(levels(Population@data$Age)[Population@data$Age])
Population@data$Time <- as.integer(levels(Population@data$Time)[Population@data$Time])
Population@data$Birth <- Population@data$Time - Population@data$Age

mort <- Ovariable(name = "mort", ddata = "Op_en5994.mortality")
mort@data$Age <- as.integer(levels(mort@data$Age)[mort@data$Age])

mort_frac <- Ovariable("mort_frac", dependencies = data.frame(Name = c("mort", "Population")), 
		formula = function(...){
			temp <- mort / Population
			temp <- temp@output[temp@output$Time == "2011" , !colnames(temp@output) %in% c("Birth")]
			return(temp)
		}
)

#################################
# Evaluation and presentation	#
#################################

out <- lifetable(pop = Population, mort_frac = mort_frac, N = 100)

cat("Mortality, Population and mortality rate by age in 2011.\n")
oprint(mort_frac)

#cat("Population data for the start years of subcohorts.\n")
#oprint(Population)

cat("Lifetable\n")
oprint(out[1:100,])
#oprint(out)

ggplot(out, aes(x = Time, y = Result, group = Birth)) + geom_line(aes(colour = Birth)) + theme_grey(base_size = 24)