##########################################################################################################################################################
# S. mansoni - ZWE IA 2021 analysis

# section 1: geostatistical model fitting and prediction
# section 2: design-based estimates from empirical data

pkgs <- c("tidyverse", "tmap", "sf","PrevMap", "lme4", "geoR", "raster", "terra", "survey")
pacman::p_load(pkgs, character.only = T)
source("outputs/ms_shared/functions.R")

#########################################################################################################################################################
# read in and transform data

smans_data <- read_csv("inputs/smans_analysis_data.csv")

# empirical logit transformation
smans_data <- smans_data %>%
  mutate(logit_smans = log( (Smns_tt + 0.5) / ((Smns_xm - Smns_tt) + 0.5) )) 

ggplot(smans_data, aes(x = logit_smans)) +
  geom_histogram(fill = "tomato", col = "gray1", alpha = 0.7, binwidth = 0.5) +
  labs(x = "S.mans empirical logit", y = "Count") +
  theme_bw() 

# transform variables as needed
smans_data <- smans_data %>%
  dplyr::select(-(contains(c("pop", "rwi")))) %>%
  mutate(t_cf_ppp_log = log(zwe_stack_cf_ppp),
         t_cf_wat_occ = log(zwe_stack_cf_wat_occ + 0.1),
         t_crop_log = log(zwe_stack_crops + 1), 
         t_crop_bin = if_else(zwe_stack_crops <= 7.5,1,0),
         t_flow_acc_log = log(zwe_stack_flow_acc), 
         t_flow_acc_log_bin = if_else(t_flow_acc_log <= 0.1,1,0),
         t_grass = log(zwe_stack_grass + 1),
         t_ndvi_min_log = log(zwe_stack_ndvi_min),
         t_perm_water_log = log(zwe_stack_perm_water_dist + 0.1),
         t_precip_driest_log = log(zwe_stack_precip_driest_month + 0.0001), 
         t_shrubs_log = log(zwe_stack_shrubs + 0.1),
         t_shrubs_bin = if_else(t_shrubs_log <=7.5,1,0)) %>% 
  dplyr::select(!c("zwe_stack_cf_ppp", "zwe_stack_cf_wat_occ", "zwe_stack_crops",
                   "zwe_stack_flow_acc", "zwe_stack_grass",
                   "zwe_stack_ndvi_min",  "zwe_stack_perm_water_dist",
                   "zwe_stack_precip_driest_month", "zwe_stack_shrubs")) # remove un-transformed versions


# convert binary preds to factors
smans_data <- smans_data %>%
  mutate(across(.cols = ends_with("_bin"), ~as_factor(.x)))

# scale and centre
id_covs <- grep("^zwe_|^t_", names(which(sapply(smans_data, is.numeric))), value = T)
scaled_covs <- scale(smans_data[id_covs])
mean_covs <- attr(scaled_covs, "scaled:center")
sd_covs <- attr(scaled_covs, "scaled:scale")
smans_data[id_covs]  <- scaled_covs

# make coorrdinates in utm projection for PrevMap
smans_data <- smans_data %>%
  sf::st_as_sf(., coords = c("X", "Y"), crs = 3857) %>%
  sf::st_transform(., crs = epsgKM(20936)) %>%
  sf::st_coordinates(.) %>%
  as.data.frame(.) %>%
  rename(x_utm = X,
         y_utm = Y) %>%
  cbind(., smans_data)

#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*
# section 1: geostatistical estimates 
#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*

###########################################################################################################################################
# covariate selection
###########################################################################################################################################

# check univariate associations (linear fits)
covariates <- smans_data %>%
  dplyr::select(starts_with(c("zwe_", "t_")))

# set null intercept only model
m0 <- glmer(cbind(Smns_tt, Smns_xm - Smns_tt) ~ 1 + (1 | siteID),  
            control = glmerControl(optimizer = "bobyqa"), family = "binomial",  data = smans_data)

# generate glm estimates for all covariates
vars <- colnames(covariates)
for(i in 1:length(vars)){
  
  # set formula
  fmla <- paste0("cbind(Smns_tt, Smns_xm - Smns_tt) ~ 1 + ", vars[i], "+ (1 | siteID)")
  
  # run model 
  m1 <- glmer(fmla, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
  
  # extract results
  se <- sqrt(diag(vcov(m1)))
  cof <- cbind(Est = fixef(m1), LL = fixef(m1) - 1.96 * se, UL = fixef(m1) + 1.96 *se)
  cof <- data.frame(exp(cof))
  a <- anova(m0, m1)
  pv <- a$`Pr(>Chisq)`[2]
  cof$pvalue <- pv
  cof$aic <- AIC(m1)
  cof$bic <- BIC(m1)
  
  # name variables and remove intercept
  cof$variable <- row.names(cof)
  row.names(cof) <- NULL
  cof <- subset(cof, variable !="(Intercept)")
  
  # add to main dataset
  if(!exists("dataset")){
    dataset <- cof
  }
  
  if(exists("dataset")){
    dataset <- rbind(dataset, cof)
    rm(cof)
  }
}

rm(fmla, i, pv, se, m1, vars, a)

# remove the 1 suffix introduced in the loop above
dataset$variable <- str_remove(dataset$variable, pattern = "1")

################################################################################################################################################################
# check collinearity and remove if correlation >0.75

# convert all to numeric
numeric_covars <- covariates %>%
  mutate(across(.cols = ends_with("_bin"), ~as.numeric(.x)))
# keep all those with pvalue <= 0.2
significant_vars <- subset(dataset, pvalue <= 0.2)
significant_var_names <- significant_vars$variable

# define a correlation matrix
correlationMatrix <- cor(numeric_covars[significant_var_names])
# find correlations
hc_covars <- caret::findCorrelation(correlationMatrix, cutoff = 0.75)
# remove correlated variables
hc_covars <- sort(hc_covars)
hc_vars <-  numeric_covars[significant_var_names]
hc_vars <- hc_vars[,-c(hc_covars)]
head(hc_vars)
final_vars <- names(hc_vars)

# collate non-correlated variables and those with p-values <0.2 for main outcomes
set_nocorr_sig <- dataset[dataset$variable %in% final_vars,] 
(set_nocorr_sig <- set_nocorr_sig %>% arrange(aic))

rm(hc_covars,hc_vars, numeric_covars, covariates, dataset, correlationMatrix, significant_vars, significant_var_names, final_vars)

################################################################################################################################################################
# forwards model building using remaining vars

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1], "+ (1 | siteID)")
m1 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
AIC(m0,m1)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+ (1 | siteID)")
m2 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m2)
AIC(m1,m2)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[3], "+ (1 | siteID)")
m3 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m3)
AIC(m2,m3)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+ (1 | siteID)")
m4 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m4)
AIC(m2,m4)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+",
           set_nocorr_sig$variable[5], "+ (1 | siteID)")
m5 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m5)
AIC(m4,m5)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+",
           set_nocorr_sig$variable[5], "+",
           set_nocorr_sig$variable[6], "+ (1 | siteID)")
m6 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m6)
AIC(m5,m6)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+",
           set_nocorr_sig$variable[5], "+",
           set_nocorr_sig$variable[7], "+ (1 | siteID)")
m7 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m7)
AIC(m5,m7)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+",
           set_nocorr_sig$variable[5], "+",
           set_nocorr_sig$variable[7], "+",
           set_nocorr_sig$variable[8], "+ (1 | siteID)")
m8 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m8)
AIC(m7,m8)

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+",
           set_nocorr_sig$variable[5], "+",
           set_nocorr_sig$variable[7], "+",
           set_nocorr_sig$variable[9], "+ (1 | siteID)")
m9 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)
summary(m9)
AIC(m8,m9)

######################################################################################################################################################################################################
# final model selected

f <- paste("cbind(Smns_tt, Smns_xm - Smns_tt) ~", set_nocorr_sig$variable[1],"+",
           set_nocorr_sig$variable[2], "+",
           set_nocorr_sig$variable[4], "+",
           set_nocorr_sig$variable[5], "+",
           set_nocorr_sig$variable[7], "+",
           set_nocorr_sig$variable[9], "+ (1 | siteID)")

n1 <- glmer(f, control = glmerControl(optimizer = "bobyqa"), family = "binomial", data = smans_data)

summary(n1)
sjPlot::plot_model(n1)
performance::check_collinearity(n1)
lattice::dotplot(ranef(n1, condVar = TRUE))

#########################################################################################################################################################
# fit geostatistical model with covariates selected in previous step
#########################################################################################################################################################

f_logit <- as.formula(logit_smans ~ zwe_stack_exc_probs + zwe_stack_precip_ann_mean + t_grass + 
                        t_cf_ppp_log + zwe_stack_temp_ann_range + zwe_stack_vaccine)

f_glm <- as.formula(cbind(Smns_tt, Smns_xm - Smns_tt) ~ zwe_stack_exc_probs + zwe_stack_precip_ann_mean + t_grass + 
                      t_cf_ppp_log + zwe_stack_temp_ann_range + zwe_stack_vaccine)

# plot empirical semivariograms

# fit a linear model using the logit prevalence
fit_lm <- lm(f_logit, data = smans_data)
summary(fit_lm)
resid_lm <- rstandard(fit_lm)
ggvario(coords = smans_data[c("x_utm", "y_utm")], data = resid_lm, nsim = 1000)

# fit binomial model 
fit_binom <- glm(f_glm,family = "binomial", data = smans_data)
summary(fit_binom)
resid_binom <- rstandard(fit_binom)
ggvario(coords = smans_data[c("x_utm", "y_utm")], data = resid_binom, nsim = 1000)

###########################################################################################################################################
# fit Gaussian model from logit to generate initial parameters

geo_linear <- linear.model.MLE(formula = f_logit,
                               coords =  ~ x_utm + y_utm,
                               data = smans_data, 
                               start.cov.pars = c(20), # estimated range (phi) and nugget effect (not estimated) from variogram
                               kappa = 0.5, 
                               messages = T, 
                               fixed.rel.nugget = 0, # fix nugget to zero as fitting less stable when nugget estimated
                               method = "nlminb")
summary(geo_linear)
geoR::practicalRange("matern", phi = 0.21, kappa = 0.5) # see change of phi estimate below regarding this
PrevMap::variog.diagnostic.lm(geo_linear, uvec = seq(0,60,3)) # diagnostics

###########################################################################################################################################
# fitting binomial geostatistical model

# control MCML 
cmcmc <- list()
cmcmc[[1]] <- control.mcmc.MCML(n.sim = 10000, burnin = 2000, thin = 8)
cmcmc[[2]] <- control.mcmc.MCML(n.sim = 10000, burnin = 2000, thin = 8)
cmcmc[[3]] <- control.mcmc.MCML(n.sim = 65000, burnin = 5000, thin = 6)

# define initial parameter estimates
par0 <- as.numeric(coef(fit_binom))
p <- length(par0)
par0 <- c(par0, as.numeric(coef(geo_linear)[(p + 1):(p + 2)])) 
#par0 <- c(par0, as.numeric(coef(geo_linear)[(p + 1):(p + 3)])) # use this if estimating nugget

# set model formula for PrevMap
f <- as.formula(Smns_tt ~ zwe_stack_exc_probs + zwe_stack_precip_ann_mean + t_grass + 
                  t_cf_ppp_log + zwe_stack_temp_ann_range + zwe_stack_vaccine)

#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*
# set more informative initial estimate of the MCML 
# larger phi estimate as model becomes unstable with smaller initial estimate of phi

par0[p + 2] <- 20

#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*

# run the model 3 times with two short simulations and one long final one
for (i in 1:3) {
  #init_cov_pars <- c(par0[p + 2], par0[p + 3] / par0[p + 1]) # use if estimating nugget
  init_cov_pars <- par0[p + 2] # use with fixed nugget
  geo_binomial <- binomial.logistic.MCML(formula = f, 
                                             units.m = ~ Smns_xm,
                                             coords = ~ x_utm + y_utm,
                                             data = smans_data,
                                             par0 = par0, # parameter estimates of sampling distribution
                                             control.mcmc = cmcmc[[i]],
                                             kappa = 0.5,
                                             fixed.rel.nugget = 0, # fix nugget so only one estimate needed for start.cov.pars
                                             start.cov.pars = init_cov_pars, # starting points of phi (only nu2 as well if estimating nugget)
                                             method = "nlminb", messages = T)
  par0 <- as.numeric(coef(geo_binomial)) # this stores the outputs of each run and then is passed as updated parameter estimates
  print(summary(geo_binomial, l = F))
}

# evaluate model fit
summary(geo_binomial) 
geoR::practicalRange("matern", phi = 30.7, kappa = 0.5)
PrevMap::variog.diagnostic.glgm(geo_binomial, uvec = seq(0,160,8)) # diagnostic checks

###########################################################################################################################################
# plot predictions and exceedance probs

# read-in 1km grid of extracted covariates
zim_1km_covars <- read_csv("inputs/zim_covars_1km.csv")

# select and transform variables to same scale as survey data
zim_1km_covars <- zim_1km_covars %>%
  dplyr::select(-(contains(c("pop", "rwi")))) %>%
  mutate(t_cf_ppp_log = log(zwe_stack_cf_ppp),#no zeros,
         t_grass = log(zwe_stack_grass + 1)) %>% # no zeros
  dplyr::select(c("zwe_stack_exc_probs", "zwe_stack_precip_ann_mean",
                  "t_grass","t_cf_ppp_log", "zwe_stack_temp_ann_range",
                  "zwe_stack_vaccine", "X", "Y"))

# convert binary preds to factors
zim_1km_covars <- zim_1km_covars %>%
  mutate(across(.cols = ends_with("_bin"), ~as_factor(.x)))

# scale covariates to scale of model covariates (mean and SD)
id_covs2 <- intersect(id_covs, names(zim_1km_covars))
zim_1km_covars[id_covs2] <- scale(zim_1km_covars[id_covs2],
                                  center = mean_covs[id_covs2], 
                                  scale = sd_covs[id_covs2])

# define prediction coordinates from extracted variables at 1km resolution
pred_coords <- zim_1km_covars %>% 
  dplyr::select(X, Y) %>% 
  st_as_sf(coords = c("X", "Y"), crs = 3857) %>% 
  sf::st_transform(., crs = epsgKM(20936))

pred_coords <- st_coordinates(pred_coords)

predictors <- zim_1km_covars %>%
  dplyr::select(id_covs2)

# check missingness
apply(predictors, 2, function(x) sum(is.na(x)))

# replace with 0s; as over water and will be masked later
predictors[is.na(predictors)] <- 0

# set as a dataframe for running
predictors <- as.data.frame(predictors)
input_df <- data.frame(pred_coords, predictors)

# set MCML for conditional S in prediction
cmcmc <- control.mcmc.MCML(n.sim = 40000, burnin = 10000, thin = 30)

# marginal predictions of full model
pred_list <-spatial_pred(geo_binomial,
                         grid.pred = pred_coords,
                         predictors = input_df[id_covs2],
                         control.mcmc = cmcmc,
                         type = "marginal",
                         messages = T,
                         plot.correlogram = F)

geo_preds <- list(samples = pred_list$samples, 
                  coords = pred_list$coords)


prev_samples <- plogis(geo_preds$samples)
prev_out <- data.frame(geo_preds$coords, mean_prev = rowMeans(prev_samples),
                       ex50 = rowMeans(prev_samples > 0.50),
                       ex20 = rowMeans(prev_samples > 0.20),
                       ex10 = rowMeans(prev_samples > 0.10),
                       ex5 = rowMeans(prev_samples > 0.05),
                       ex2 = rowMeans(prev_samples > 0.02),
                       ex1 = rowMeans(prev_samples > 0.01)
)


#######################################################################################################################################################
# generate 1km marginal predictions

# get coordinates in 3857
coords_3857 <- zim_1km_covars[c("X", "Y")]
coords_3857$id <- 1:nrow(coords_3857)

# convert to sf and transform to 3857
pred_pixels_3857 <- prev_out %>% 
  st_as_sf(coords = c("X", "Y"), crs = epsgKM(20936)) %>% 
  st_transform(3857)

# make raster of coords in 20936
rast_3857 <- rasterFromXYZ(coords_3857)
#extract ids from 20936 and pred pixels
id <- extract(rast_3857, pred_pixels_3857)
# remove geometry
preds_nogeom <- st_drop_geometry(pred_pixels_3857)

# make raster
out <- rasterFromXYZ(data.frame(coords_3857[id,1:2], preds_nogeom[,1:7]))
crs(out) <- "EPSG:3857"
plot(out[[1]])

#######################################################################################################################################################
# calculate exceedance probs at IU level (wards) from joint distribution
#######################################################################################################################################################

pred_coords <- zim_1km_covars %>% 
  dplyr::select(X, Y) %>% 
  st_as_sf(coords = c("X", "Y"), crs = 3857) %>%
  sf::st_transform(., crs = epsgKM(20936))

IUs <- st_read("data/wards/ZWE_ward_boundaries.shp") %>%
  st_transform(crs = st_crs(pred_coords)) 

id_ius <- pred_coords %>% 
  st_join(IUs["ADM3_PCODE"]) %>% 
  pull("ADM3_PCODE")

# set input dataframe
input <- data.frame(st_coordinates(pred_coords), predictors, id_ius)

# remove the pred points that fall outside the IUs boundaries
keep <- !is.na(input$id_ius)
input <- input[keep, ]

# split the input into a list of IUs for joint estimation
input_list <- split(input, f = input$id_ius)

# set the MCL parameters for S
cmcmc <- control.mcmc.MCML(n.sim = 40000, burnin = 10000, thin = 30)

# read-in laplace samples of S (or generate own by changing save.sim to TRUE)
S.sim <- readRDS("inputs/smans_laplace_sampling.rds")

# joint estimation per IU
pred_list <- lapply(input_list,
                    function(x) {
                      type <- ifelse(nrow(x) == 1, "marginal", "joint")
                      out <- spatial_pred(geo_binomial,
                                          grid.pred = x[c("X", "Y")],
                                          predictors = x[id_covs2],
                                          control.mcmc = cmcmc,
                                          S.sim = S.sim,
                                          save.sim = FALSE,
                                          type = type,
                                          messages = T,
                                          plot.correlogram = F)
                      return(out$samples)
                    })


samples <- do.call(rbind, pred_list)
predictions <- list(samples = samples, 
                    coords = do.call(rbind, input_list)[c("X", "Y")],
                    IU = do.call(rbind, input_list)$id_ius)

#######################################################################################################################################################
# population weight joint estimates per IU
#######################################################################################################################################################

# read in 1km unconstrained pop surface
world_pop_zwe <- raster("inputs/worldpop_zwe_unconstrained_2020.tif")

# get coordinates in 4326 to match WorldPop
coords_ll <- predictions$coords %>% 
  as.data.frame() %>% 
  st_as_sf(coords = c("X", "Y"), crs = st_crs(pred_coords)) %>% 
  st_transform(4326) %>% 
  st_coordinates()

# extract population
pop <- raster::extract(world_pop_zwe, coords_ll)

# aggregate with population weighting
samples <- data.frame(admin = predictions$IU, pop, plogis(predictions$samples)) %>% 
  na.omit() %>% 
  joint_weighted_prev() 

summary_prev <- samples %>% 
  group_by(admin) %>% 
  summarise(mean_prev = mean(prev),
            ex50 = mean(prev > 0.5), 
            ex20 = mean(prev > 0.2), 
            ex10 = mean(prev > 0.1), 
            ex5 = mean(prev > 0.05),
            ex2 = mean(prev > 0.02),
            ex1 = mean(prev > 0.01))

# join to shapefile for plotting
prova <- inner_join(IUs, summary_prev, by = c("ADM3_PCODE" = "admin"))

ggplot(prova) +
  geom_sf(aes(fill = mean_prev * 100), col = NA) +
  scale_fill_distiller(type = "div", palette = 9)

#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*
# section 2: design-based estimates adjusted for sampling design 
#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*

# needs weighting by inclusion probability which is 1/exceedance probability from 2018 impact assessment survey

sf_use_s2(FALSE)

# read back in so exceedance probability is un-scaled
data <- read_csv("inputs/smans_analysis_data.csv")

wards <- st_read("inputs/ZWE_ward_boundaries.shp") %>%
  st_transform(3857)

data_sf <- st_as_sf(data, coords = c("X", "Y"), crs = 3857) 

###  shapefile and admn data may not align so join and adapt here
data_bernoulli <- data_sf %>%
  st_join(., wards) 

# check missing values
sum(is.na(data_bernoulli$ADM1_EN))

# generate reciprocal weights
data_bernoulli$inc_probs <- 1/ data_bernoulli$zwe_stack_exc_probs

# make data bernoulli
bernouli_dat <- data_bernoulli %>%
  dplyr::select(ADM1_EN, ADM2_EN, ADM3_EN, siteID, Smns_tt, Smns_xm, inc_probs) %>%
  st_drop_geometry()

bernouli_dat <-  purrr::pmap_dfr(bernouli_dat, 
                                 function(ADM1_EN, ADM2_EN, ADM3_EN, siteID, Smns_tt, Smns_xm, inc_probs) {
                                   data.frame(admin_1 = ADM1_EN,
                                              admin_2 = ADM2_EN,
                                              admin_3 = ADM3_EN,
                                              village = siteID,
                                              inc_probs = inc_probs,
                                              outcome = c(rep(1, Smns_tt),
                                                          rep(0, Smns_xm - Smns_tt) ) )
                                 }
)

# make weighted design object and analyse
design <- svydesign(id = ~village, 
                            weights= ~inc_probs, # to check results without weighting change to ~1
                            nest = TRUE, 
                            survey.lonely.psu = "adjust", 
                            data = bernouli_dat)

# population mean table
pop_mean <- svymean(~outcome, design)
# stata svyprop version
pop_confint <- svyciprop(~I(outcome == 1), design, method="lo", df=degf(design))
pop_confint <- attr(pop_confint,"ci")
pop_out <- t(as.data.frame(c(pop_mean, pop_confint)))
colnames(pop_out) <- c("outcome", "lci", "uci")
rownames(pop_out) <- c("prev")

# calculate numbers
total_numbers <- bernouli_dat %>%
  summarise(infected = sum(outcome),
            examined = n(),
            crude_prev = infected / examined)

# create table 
means_out <- data.frame(rbind(pop_out))
totals_out <- data.frame(rbind(total_numbers[,1:2]))
summary_out <- data.frame(cbind(means_out, totals_out))

#################################################################################################################################
# province analysis

adm1_means <- svyby(~outcome, ~admin_1, design, na = TRUE, svymean)
# fails to generate CIs for districts with 0 cases and few observations
adm1_cis <- svyby(~outcome, ~admin_1, design, na = TRUE, 
                  svyciprop, vartype="ci", method = "logit")

region_numbers <- bernouli_dat %>%
  group_by(admin_1) %>%
  summarise(infected = sum(outcome),
            examined = n())

region_numbers <- adm1_cis %>%
  left_join(., region_numbers, by = c("admin_1"))

#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*

# to repeat for S. haematobium analysis: read-in urine filtration survey data and replace infection outcome (Smns_tt) and total examined (Smns_xm) per site variables with those of s. haem UF data
# repeat process but change i) include estimation of nugget into fitted geospatial models and ii) only use zwe_stack_exc_probs as single covariate (reduced variable selection)

#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*#*