# simple random ---------------------------------------------------------------
base_survey_design_simple <- svydesign(ids = ~1, # 1 for no cluster ids
                   weights = NULL,               # No weight added
                   strata = NULL,                # sampling was simple (no strata)
                   data = survey_data            # have to specify the dataset
                  )

## stratified ------------------------------------------------------------------
base_survey_design_strata <- svydesign(ids = ~1,  # 1 for no cluster ids
                   weights = ~surv_weight_strata, # weight variable created above
                   strata = ~health_district,     # sampling was stratified by district
                   data = survey_data             # have to specify the dataset
                  )

# cluster ---------------------------------------------------------------------
base_survey_design_cluster <- svydesign(ids = ~village_name, # cluster ids
                   weights = ~surv_weight_cluster, # weight variable created above
                   strata = NULL,                 # sampling was simple (no strata)
                   data = survey_data              # have to specify the dataset
                  )

# stratified cluster ----------------------------------------------------------
base_survey_design <- svydesign(ids = ~village_name,      # cluster ids
                   weights = ~surv_weight_cluster_strata, # weight variable created above
                   strata = ~health_district,             # sampling was stratified by district
                   data = survey_data                     # have to specify the dataset
                  )

## simple random ---------------------------------------------------------------
survey_design_simple <- survey_data %>% 
  as_survey_design(ids = 1, # 1 for no cluster ids 
                   weights = NULL, # No weight added
                   strata = NULL # sampling was simple (no strata)
                  )
## stratified ------------------------------------------------------------------
survey_design_strata <- survey_data %>%
  as_survey_design(ids = 1, # 1 for no cluster ids
                   weights = surv_weight_strata, # weight variable created above
                   strata = health_district # sampling was stratified by district
                  )
## cluster ---------------------------------------------------------------------
survey_design_cluster <- survey_data %>%
  as_survey_design(ids = village_name, # cluster ids
                   weights = surv_weight_cluster, # weight variable created above
                   strata = NULL # sampling was simple (no strata)
                  )

## stratified cluster ----------------------------------------------------------
survey_design <- survey_data %>%
  as_survey_design(ids = village_name, # cluster ids
                   weights = surv_weight_cluster_strata, # weight variable created above
                   strata = health_district # sampling was stratified by district
                  )

## counts and props of the study population
ag <- survey_data %>% 
  group_by(age_group) %>% 
  drop_na(age_group) %>% 
  tally() %>% 
  mutate(proportion = n / sum(n), 
         n_total = sum(n))

## counts and props of the source population
propcount <- population %>% 
  group_by(age_group) %>%
    tally(population) %>%
    mutate(proportion = n / sum(n))

## bind together the columns of two tables, group by age, and perform a 
## binomial test to see if n/total is significantly different from population
## proportion.
  ## suffix here adds to text to the end of columns in each of the two datasets
left_join(ag, propcount, by = "age_group", suffix = c("", "_pop")) %>%
  group_by(age_group) %>%
  ## broom::tidy(binom.test()) makes a data frame out of the binomial test and
  ## will add the variables p.value, parameter, conf.low, conf.high, method, and
  ## alternative. We will only use p.value here. You can include other
  ## columns if you want to report confidence intervals
  mutate(binom = list(broom::tidy(binom.test(n, n_total, proportion_pop)))) %>%
  unnest(cols = c(binom)) %>% # important for expanding the binom.test data frame
  mutate(proportion_pop = proportion_pop * 100) %>%
  ## Adjusting the p-values to correct for false positives 
  ## (because testing multiple age groups). This will only make 
  ## a difference if you have many age categories
  mutate(p.value = p.adjust(p.value, method = "holm")) %>%
                      
  ## Only show p-values over 0.001 (those under report as <0.001)
  mutate(p.value = ifelse(p.value < 0.001, 
                          "<0.001", 
                          as.character(round(p.value, 3)))) %>% 
  
  ## rename the columns appropriately
  select(
    "Age group" = age_group,
    "Study population (n)" = n,
    "Study population (%)" = proportion,
    "Source population (n)" = n_pop,
    "Source population (%)" = proportion_pop,
    "P-value" = p.value
  )

## define x-axis limits and labels ---------------------------------------------
## (update these numbers to be the values for your graph)
max_prop <- 35      # choose the highest proportion you want to show 
step <- 5           # choose the space you want beween labels 

## this part defines vector using the above numbers with axis breaks
breaks <- c(
    seq(max_prop/100 * -1, 0 - step/100, step/100), 
    0, 
    seq(0 + step / 100, max_prop/100, step/100)
    )

## this part defines vector using the above numbers with axis limits
limits <- c(max_prop/100 * -1, max_prop/100)

## this part defines vector using the above numbers with axis labels
labels <-  c(
      seq(max_prop, step, -step), 
      0, 
      seq(step, max_prop, step)
    )


## create plots individually  --------------------------------------------------

## plot the source population 
## nb: this needs to be collapsed for the overall population (i.e. removing health districts)
source_population <- population %>%
  ## ensure that age and sex are factors
  mutate(age_group = factor(age_group, 
                            levels = c("0-2", 
                                       "3-14", 
                                       "15-29",
                                       "30-44", 
                                       "45+")), 
         sex = factor(sex)) %>% 
  group_by(age_group, sex) %>% 
  ## add the counts for each health district together 
  summarise(population = sum(population)) %>% 
  ## remove the grouping so can calculate overall proportion
  ungroup() %>% 
  mutate(proportion = population / sum(population)) %>% 
  ## plot pyramid 
  age_pyramid(
            age_group = age_group, 
            split_by = sex, 
            count = proportion, 
            proportional = TRUE) +
  ## only show the y axis label (otherwise repeated in all three plots)
  labs(title = "Source population", 
       y = "", 
       x = "Age group (years)") + 
  ## make the x axis the same for all plots 
  scale_y_continuous(breaks = breaks, 
    limits = limits, 
    labels = labels)
  
  
## plot the unweighted sample population 
sample_population <- plot_age_pyramid(survey_data, 
                 age_group = "age_group", 
                 split_by = "sex",
                 proportion = TRUE) + 
  ## only show the x axis label (otherwise repeated in all three plots)
  labs(title = "Unweighted sample population", 
       y = "Proportion (%)", 
       x = "") + 
  ## make the x axis the same for all plots 
  scale_y_continuous(breaks = breaks, 
    limits = limits, 
    labels = labels)


## plot the weighted sample population 
weighted_population <- survey_design %>% 
  ## make sure the variables are factors
  mutate(age_group = factor(age_group), 
         sex = factor(sex)) %>%
  plot_age_pyramid(
    age_group = "age_group",
    split_by = "sex", 
    proportion = TRUE) +
  ## only show the x axis label (otherwise repeated in all three plots)
  labs(title = "Weighted sample population", 
       y = "", 
       x = "")  + 
  ## make the x axis the same for all plots 
  scale_y_continuous(breaks = breaks, 
    limits = limits, 
    labels = labels)

## combine all three plots  ----------------------------------------------------
## combine three plots next to eachother using + 
source_population + sample_population + weighted_population + 
  ## only show one legend and define theme 
  ## note the use of & for combining theme with plot_layout()
  plot_layout(guides = "collect") & 
  theme(legend.position = "bottom",                    # move legend to bottom
        legend.title = element_blank(),                # remove title
        text = element_text(size = 18),                # change text size
        axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1) # turn x-axis text
       )

## summarize data
flow_table <- survey_data %>%
  count(startcause, endcause, sex) %>%  # get counts 
  gather_set_data(x = c("startcause", "endcause")) %>%     # change format for plotting
  mutate(x = fct_relevel(x, c("startcause", "endcause")),  # set startcause as first level
         x = fct_recode(x, 
                        "Start \n cause" = "startcause",   # add line break (\n) after start
                        "End \n cause"   = "endcause")
        )


## plot your dataset 
  ## on the x axis is the start and end causes
  ## gather_set_data generates an ID for each possible combination
  ## splitting by y gives the possible start/end combos
  ## value as n gives it as counts (could also be changed to proportion)
ggplot(flow_table, aes(x, id = id, split = y, value = n)) +
  ## colour lines by sex 
  geom_parallel_sets(aes(fill = sex), alpha = 0.5, axis.width = 0.2) +
  ## fill in the label boxes grey
  geom_parallel_sets_axes(axis.width = 0.15, fill = "grey80", color = "grey80") +
  ## change text colour and angle (needs to be adjusted)
  geom_parallel_sets_labels(color = "black", angle = 0, size = 5) +
  ## adjusted y and x axes (probably needs more vertical space)
  scale_x_discrete(name = NULL, expand = c(0, 0.2)) + 
  ## remove axis labels
  theme(
    title = element_text(size = 26),
    text = element_text(size = 26),
    axis.line = element_blank(),
    axis.ticks = element_blank(),
    axis.text.y = element_blank(),
    panel.background = element_blank(),
    legend.position = "bottom",                    # move legend to bottom
    legend.title = element_blank(),                # remove title
  )

## produce weighted counts 
svytable(~died, base_survey_design)

## produce weighted proportions
svyciprop(~died, base_survey_design, na.rm = T)

## get the design effect 
svymean(~died, base_survey_design, na.rm = T, deff = T) %>% 
  deff()

# Define function to calculate weighted counts, proportions, CI and design effect
# x is the variable in quotation marks 
# design is your survey design object

svy_prop <- function(design, x) {
  
  ## put the variable of interest in a formula 
  form <- as.formula(paste0( "~" , x))
  ## only keep the TRUE column of counts from svytable
  weighted_counts <- svytable(form, design)[[2]]
  ## calculate proportions (multiply by 100 to get percentages)
  weighted_props <- svyciprop(form, design, na.rm = TRUE) * 100
  ## extract the confidence intervals and multiply to get percentages
  weighted_confint <- confint(weighted_props) * 100
  ## use svymean to calculate design effect and only keep the TRUE column
  design_eff <- deff(svymean(form, design, na.rm = TRUE, deff = TRUE))[[TRUE]]
  
  ## combine in to one data frame
  full_table <- cbind(
    "Variable"        = x,
    "Count"           = weighted_counts,
    "Proportion"      = weighted_props,
    weighted_confint, 
    "Design effect"   = design_eff
    )
  
  ## return table as a dataframe
  full_table <- data.frame(full_table, 
             ## remove the variable names from rows (is a separate column now)
             row.names = NULL)
  
  ## change numerics back to numeric
  full_table[ , 2:6] <- as.numeric(full_table[, 2:6])
  
  ## return dataframe
  full_table
}

## iterate over several variables to create a table 
purrr::map(
  ## define variables of interest
  c("left", "died", "arrived"), 
  ## state function using and arguments for that function (design)
  svy_prop, design = base_survey_design) %>% 
  ## collapse list in to a single data frame
  bind_rows() %>% 
  ## round 
  mutate(across(where(is.numeric), round, digits = 1))

## use the srvyr design object
survey_design %>% 
  summarise(
    ## produce the weighted counts 
    counts = survey_total(died), 
    ## produce weighted proportions and confidence intervals 
    ## multiply by 100 to get a percentage 
    props = survey_mean(died, 
                        proportion = TRUE, 
                        vartype = "ci") * 100, 
    ## produce the design effect 
    deff = survey_mean(died, deff = TRUE)) %>% 
  ## only keep the rows of interest
  ## (drop standard errors and repeat proportion calculation)
  select(counts, props, props_low, props_upp, deff_deff)

# Define function to calculate weighted counts, proportions, CI and design effect
# design is your survey design object
# x is the variable in quotation marks 


srvyr_prop <- function(design, x) {
  
  summarise(
    ## using the survey design object
    design, 
    ## produce the weighted counts 
    counts = survey_total(.data[[x]]), 
    ## produce weighted proportions and confidence intervals 
    ## multiply by 100 to get a percentage 
    props = survey_mean(.data[[x]], 
                        proportion = TRUE, 
                        vartype = "ci") * 100, 
    ## produce the design effect 
    deff = survey_mean(.data[[x]], deff = TRUE)) %>% 
  ## add in the variable name
  mutate(variable = x) %>% 
  ## only keep the rows of interest
  ## (drop standard errors and repeat proportion calculation)
  select(variable, counts, props, props_low, props_upp, deff_deff)
  
}
  

## iterate over several variables to create a table 
purrr::map(
  ## define variables of interest
  c("left", "died", "arrived"), 
  ## state function using and arguments for that function (design)
  ~srvyr_prop(.x, design = survey_design)) %>% 
  ## collapse list in to a single data frame
  bind_rows()

## using the survey design object
survey_design %>% 
  ## pass the names of variables of interest unquoted
  tab_survey(arrived, left, died, education_level,
             deff = TRUE,   # calculate the design effect
             pretty = TRUE  # merge the proportion and 95%CI
             )

confidence_intervals <- function(data, variable, by, ...) {
  
  ## extract the confidence intervals and multiply to get percentages
  props <- svyciprop(as.formula(paste0( "~" , variable)),
              data, na.rm = TRUE)
  
  ## extract the confidence intervals 
  as.numeric(confint(props) * 100) %>% ## make numeric and multiply for percentage
    round(., digits = 1) %>%           ## round to one digit
    c(.) %>%                           ## extract the numbers from matrix
    paste0(., collapse = "-")          ## combine to single character
}

## using the survey package design object
tbl_svysummary(base_survey_design, 
               include = c(arrived, left, died),   ## define variables want to include
               statistic = list(everything() ~ c("{n} ({p}%)"))) %>% ## define stats of interest
  add_n() %>%  ## add the weighted total 
  add_stat(fns = everything() ~ confidence_intervals) %>% ## add CIs
  ## modify the column headers
  modify_header(
    list(
      n ~ "**Weighted total (N)**",
      stat_0 ~ "**Weighted Count**",
      add_stat_1 ~ "**95%CI**"
    )
    )

ratio <- svyratio(~died, 
         denominator = ~obstime, 
         design = base_survey_design)

ci <- confint(ratio)

cbind(
  ratio$ratio * 10000, 
  ci * 10000
)

survey_design %>% 
  ## survey ratio used to account for observation time 
  summarise(
    mortality = survey_ratio(
      as.numeric(died) * 10000, 
      obstime, 
      vartype = "ci")
    )