# --- Definition compartments, population events and events affecting which compartments ---

## S..: Susceptible
## E...: Exposed
## I1...: Infected in first stage
## I2...: Infected in second stage
## I3...: Infected in final stage
## HS...: Harvested susceptible
## HE...: Harvested exposed
## ...
## D3_cf...: Dead calf female
## D3_cm...: Dead calf male
## ...

## .._cf_...: calf female
## .._cm_...: calf male
## .._yf_...: female yearling
## .._ym_...: male yearling
## .._af_...: adult female
## .._am_...: adult male

## ..._hs: homozygous susceptible
## ..._het: heterozygos 
## ..._hr: homozygous resistant

# Define constants
genetic_states = c("hs", "het", "hr")
base_compartments = c("S", "E", "I1", "I2", "I3")
age_groups <- c("cf", "cm", "yf", "ym", "af", "am")
tracking_variables = paste0("H", base_compartments)
dead_compartment = "D3"
process_compartments = c("activate_births", "activate_harvest")
n_compartments = length(base_compartments) # 5
n_age_groups = length(age_groups) # 6
n_genetic_states = length(genetic_states) # 3

# Generate compartments from tracking variables and age groups, including genetic states
compartments <- c(
  unlist(lapply(genetic_states, function(gs) paste0(base_compartments, "_", rep(age_groups, each = n_compartments), "_", gs))),
  unlist(lapply(genetic_states, function(gs) paste0(tracking_variables, "_", rep(age_groups, each = length(tracking_variables)), "_", gs))),
  paste0(dead_compartment, "_", age_groups),
  process_compartments)

#Extend the event matrix
#column 1    Reproduction
#column 2    Aging
#column 3-8  Harvest proportions specific to demographic class in the homozygous susceptible
#column 9-14 Harvest proportions specific to demographic class in the susceptible group
#column 15:20 Harvest proportions specific to demographic class in the homozygous resistant
#column  Harvest proportions specific to demographic class in the susceptible group
#column 21    Number of hunted set to zero at start of each year
#column 22   Extra summer calf mortality
#column 23 activate birth
#column 24 activate harvests

generate_event_indicators <- function(age_group, gene_group=NULL, events) {
  combined_state = paste0(age_group,"_",gene_group)
  select_events <- c(rep(0, 24))
  select_events[1] <- ifelse("reproduction" %in% events, 1, 0)
  select_events[2] <- ifelse("ageing" %in% events, 1, 0)
  
  if("harvest" %in% events) {
    age_group_harvest_columns <- c("cf"=3, "cm"=4, "yf"=5, "ym"=6, "af"=7, "am"=8)
    gene_group_harvest_columns <- c("hs"=0, "het"=6, "hr"=12)
    select_column = age_group_harvest_columns[age_group]+gene_group_harvest_columns[gene_group]
    select_events[select_column] <- 1
  }
  # Assuming hunted is the 4th event in the list
  if("hunted" %in% events) {
    select_events[21] <- 1
  }
  # Assuming extra summer calf mortality is the 5th event in the list
  if("extra_summer_calf_mortality" %in% events) {
    select_events[22] <- 1
  }
  # Assuming extra summer calf mortality is the 5th event in the list
  if("activate_births" %in% events) {
    select_events[23] <- 1
  }
  if("activate_harvest" %in% events) {
    select_events[24] <- 1
  }
  return(select_events)
}



E <- matrix(c(rep(generate_event_indicators(age_group="cf",gene_group="hs", events=c("ageing", "extra_summer_calf_mortality")),n_compartments),
              rep(generate_event_indicators(age_group="cm",gene_group="hs", events=c("ageing","extra_summer_calf_mortality")),n_compartments),
              rep(generate_event_indicators(age_group="yf",gene_group="hs",events=c("ageing")),n_compartments),
              rep(generate_event_indicators(age_group="ym", gene_group="hs",events=c( "ageing")),n_compartments),
              rep(generate_event_indicators(age_group="af", gene_group="hs",events=c("reproduction")),n_compartments),
              rep(generate_event_indicators(age_group="am", gene_group="hs",events=c("")),n_compartments),
              rep(generate_event_indicators(age_group="cf",gene_group="het", events=c("", "ageing", "extra_summer_calf_mortality")),n_compartments),
              rep(generate_event_indicators(age_group="cm",gene_group="het", events=c("", "ageing","extra_summer_calf_mortality")),n_compartments),
              rep(generate_event_indicators(age_group="yf",gene_group="het",events=c("", "ageing")),n_compartments),
              rep(generate_event_indicators(age_group="ym", gene_group="het",events=c("", "ageing")),n_compartments),
              rep(generate_event_indicators(age_group="af", gene_group="het",events=c("","reproduction")),n_compartments),
              rep(generate_event_indicators(age_group="am", gene_group="het",events=c("")),n_compartments),
              rep(generate_event_indicators(age_group="cf",gene_group="hr", events=c("", "ageing", "extra_summer_calf_mortality")),n_compartments),
              rep(generate_event_indicators(age_group="cm",gene_group="hr", events=c("", "ageing","extra_summer_calf_mortality")),n_compartments),
              rep(generate_event_indicators(age_group="yf",gene_group="hr",events=c("", "ageing")),n_compartments),
              rep(generate_event_indicators(age_group="ym", gene_group="hr",events=c("", "ageing")),n_compartments),
              rep(generate_event_indicators(age_group="af", gene_group="hr",events=c("","reproduction")),n_compartments),
              rep(generate_event_indicators(age_group="am", gene_group="hr",events=c("")),n_compartments),
              rep(generate_event_indicators(age_group=NULL, events=c("hunted")),n_compartments*n_age_groups*n_genetic_states), # 5*5 for 6 age groups
              rep(generate_event_indicators(age_group=NULL, events=c("")),n_age_groups),
              rep(generate_event_indicators(age_group=NULL, events=c("activate_births"))), # 6 tracking variables for "dead compartments"
              rep(generate_event_indicators(age_group=NULL, events=c("activate_harvest")))),
            byrow = TRUE,
            nrow = length(compartments),
            dimnames = list(compartments))

# nrow(E) 186


generate_compartment_indicators <- function(current_compartment, compartments = compartments) {

  compartment_name = gsub("_.*", "", current_compartment)
  compartment_starts_with = substr(compartment_name, 1, 1)
  age_group <- gsub("^[^_]*_([^_]*)_.*$", "\\1", current_compartment)
  gene_group <- gsub(".*_", "", current_compartment)
  gender <- substr(age_group, 2, 2)
   
  find_index_difference <- function(current_compartment, go_to_compartment) {
   if (go_to_compartment == 0) return(0)
   ind1 = which(compartments == current_compartment)
   ind2 = which(compartments == go_to_compartment)
   return(ind2 - ind1)
  }
  
  if (compartment_starts_with %in% c("H", "D","a")) { # a for activate births
    return(c(0, 0, 0, 0))
  }
  
  # The shift matrix columns setup
  if (age_group %in% c("cf", "cm")) {
    go_to_compartment = c("0", "0", paste0(compartment_name, "_y", gender, "_", gene_group), paste0("H", compartment_name, "_", age_group,"_", gene_group))
  } else if (age_group %in% c("yf", "ym")) {
    go_to_compartment = c("0", "0", paste0(compartment_name, "_a", gender,"_", gene_group), paste0("H", compartment_name, "_", age_group,"_", gene_group))
  } else if (age_group == "af") {
    go_to_compartment = c(paste0("S_cf","_",gene_group), paste0("S_cm","_",gene_group), "0", paste0("H", compartment_name, "_", age_group, "_", gene_group))
  } else if (age_group == "am") {
    go_to_compartment = c("0", "0", "0", paste0("H", compartment_name, "_", age_group,"_",gene_group))
  }
  
  shift_row = unlist(lapply(go_to_compartment, function(x) find_index_difference(current_compartment=current_compartment, x)))
  return(shift_row)
 }

## Shift matrix
#N: name of shift matrix in function mparse
#Each row corresponds to one compartment in the model.
N <- matrix(c(unlist(lapply(compartments, function(x) generate_compartment_indicators(x, compartments=compartments)))),## D3_am),
            byrow = TRUE,
            nrow = length(compartments),
            dimnames = list(compartments))


# Include harvest rates
# Monthly scale
population_events <- function(node,n_years,calf_mort_rate) {

  years <- n_years
  rbind(
        ageing_events(node = node,
                      month = 5, #after jan+feb+mar+apr
                      years = years),
        activate_birth_events(node = node, month=6,
                              n_years = n_years),
        end_birth_events(node=node, month=7,n_years = n_years),
        
        activate_harvest_events(node = node, month=9,
                                n_years = n_years), 
        end_harvest_events(node = node, month=10,
                           n_years = n_years), 
        calf_mort_events(node = node, month=7,
                         calf_mort_rate = calf_mort_rate))
}

##' Generate ageing events
##' Female and male yearlings become adult females and males after
##' approximately one year.

##' @param node the nodes to generate events for.
##' @param day an integer with the day of the year when to age calves.
##' @param years an integer with the number of years to generate
##'     events. #se Widgren et al. 2016 ??
ageing_events <- function(node, month, years) {
  time <- as.numeric(sapply(seq_len(years), function(year) {
    ##(year - 1) * 365 + day
    (year - 1) * 12 + month
  }))
  
  data.frame(
    ## In which node does the event occur
    node = rep(node, times = length(time)),
    
    ## When does the event occur
    time = rep(time, each = length(node)),
    
    event = "intTrans", #type of event
    dest = 0,           #destination node
    
    ## Set n to zero to use proportion
    n = 0,
    
    ## All calves (yearlings) are aged
    proportion = 1,
    select = 2,  #column in select matrix E
    shift = 3)  #column in shift matrix N
}



##' Activate birth events
activate_birth_events <- function(node,
                         #days = 31 + 29 + 31 + 30 + seq_len(31),
                         month = 6, #put at start June
                         n_years) {
  time <- as.numeric(sapply(seq_len(n_years), function(year) {
    (year - 1) * 12 + month
  }))
  
  data.frame(
    node = rep(node, times = length(time), each = 1),
    time = rep(time, each = length(node)),
    event = "enter",  #type of event
    dest = 0,         #destination node
    n = 1, #The number of individuals affected by the event
    proportion=0,
    select = 23, #column in select matrix E
    shift = 0)  #columns in shift matrix N
}


##' End birth events
end_birth_events <- function(node,
                                  #days = 31 + 29 + 31 + 30 + seq_len(31),
                                  month = 7,
                             n_years) {
  time <- as.numeric(sapply(seq_len(n_years), function(year) {
    (year - 1) * 12 + month
  }))
  
  data.frame(
    node = rep(node, times = length(time), each = 1),
    time = rep(time, each = length(node)),
    event = "exit",  #type of event
    dest = 0,         #destination node
    n = 1, #The number of individuals affected by the event
    proportion=0,
    select = 23, #column in select matrix E
    shift = 0)  #columns in shift matrix N
}


# Activate harvest
activate_harvest_events <- function(node,
                                  #days = 31 + 29 + 31 + 30 + seq_len(31),
                                  month = 9, #put at start June
                                  n_years) {
  time <- as.numeric(sapply(seq_len(n_years), function(year) {
    (year - 1) * 12 + month
  }))
  
  data.frame(
    node = rep(node, times = length(time), each = 1),
    time = rep(time, each = length(node)),
    event = "enter",  #type of event
    dest = 0,         #destination node
    n = 1, #The number of individuals affected by the event
    proportion=0,
    select = 24, #column in select matrix E
    shift = 0)  #columns in shift matrix N
}


##' End birth events
end_harvest_events <- function(node,
                             #days = 31 + 29 + 31 + 30 + seq_len(31),
                             month = 10,
                             n_years) {
  time <- as.numeric(sapply(seq_len(n_years), function(year) {
    (year - 1) * 12 + month
  }))
  
  data.frame(
    node = rep(node, times = length(time), each = 1),
    time = rep(time, each = length(node)),
    event = "exit",  #type of event
    dest = 0,         #destination node
    n = 1, #The number of individuals affected by the event
    proportion=0,
    select = 24, #column in select matrix E
    shift = 0)  #columns in shift matrix N
}


##' Generate extra summer mortality calves
##'
##' @param node the nodes (subpopulations) to generate events for.
##' @param month an integer sequence with the months within a year (monthly scale of model)
##' @param calf_mort_rates the summer mortality_rate of calves. A vector
##'     with number/proportion for the relevant period.
calf_mort_events <- function(node,
                             month = 7, 
                             calf_mort_rate) {
  
  time <- as.numeric(sapply(seq_len(length(calf_mort_rate)), function(year) {
    (year - 1) * 12 + month
  }))
  
  data.frame(
    ## In which node and does the event occur. 
    node = rep(node, times = length(time)), #2 demographic classes
    ## When does the event occur
    time = rep(time, each = length(unique(node))),#2 demographic classes
    event = 'exit',   #type of event
    dest = 0,         #destination node
    
    ## Set n to zero to use proportion
    n = 0, #The number of individuals affected by the event
    
    ## repeat the proportion for every node and every month with harvest event
    proportion=rep(as.numeric(calf_mort_rate),each=length(unique(node))),
    
    select = 22, #column in select matrix E
    shift = 0         #column in shift matrix N
  )
}