## NOTE: To reduce the long computation times in "Section 4.4" below
## you can uncomment lines 162-164, reducing the number of iterations/chains.
## Results will be less reliable but show how the software works in principle.

# load other packages
library("bayesplot")
library("loo")
library("ggplot2")
# load tipsae
library("tipsae")


###############
## Section 3 ##
###############

# load datasets
data("emilia")
head(emilia)
data("emilia_cs")

#################
## Section 4.1 ##
#################
# set parallel option
options(mc.cores = parallel::detectCores())

# Fit beta SAE model
fit_beta <- fit_sae(formula_fixed = hcr ~ x, data = emilia_cs, domains = "id",
  type_disp = "var", disp_direct = "vars", domain_size = "n", seed = 0)

# Model with Flexible Beta likelihood
# and Variance-Gamma prior for random effects
fit_FB <- fit_sae(formula_fixed = hcr ~ x, data = emilia_cs, domains = "id",
  type_disp = "var", disp_direct = "vars", domain_size = "n", 
  likelihood = "flexbeta", prior_reff = "VG", adapt_delta = 0.99, iter = 8000,
  seed = 0)


#################
## Section 4.2 ##
#################

# Traceplots
post_beta <- as.array(fit_beta$stanfit, pars = c("beta0", "beta"))
mcmc_trace(x = post_beta)

# Beta model summaries
summ_beta <- summary(fit_beta)
summ_beta

###################
## Section 4.2.1 ##
###################

# Flexible Beta model summaries
summ_FB <- summary(fit_FB)
# Comparing LOOIC
loo_compare(list("beta" = summ_beta$loo, "flexbeta" = summ_FB$loo))

# Posterior predictive checks
ppc_dens_overlay(y = summ_beta$direct_est, yrep = summ_beta$y_rep[1:50,]) +
  ggtitle("Beta likelihood")
ppc_dens_overlay(y = summ_FB$direct_est, yrep = summ_FB$y_rep[1:50,]) +
  ggtitle("Flexible Beta likelihood")

# Caterpillar plot for random effects
ggplot(summ_beta$raneff$unstructured, aes(x = reorder(Domains, mean))) +
  geom_point(aes(y = mean)) +
  geom_linerange(aes(ymin = `2.5%`, ymax = `97.5%`)) +
  geom_hline(yintercept = 0, lty = 2) +
  ylab("Random effect") + xlab("") +
  theme_bw(base_size = 12) +
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust = 1))


###################
## Section 4.2.2 ##
###################

# plot method for summary_fitsae S3 object
plot(summ_beta)
# density method for summary_fitsae S3 object
density(summ_beta)

# load shapefile and perform maps
data("emilia_shp")
map(x = summ_beta, spatial_df = emilia_shp, 
  spatial_id_domains = "NAME_DISTRICT")
map(x = summ_beta, spatial_df = emilia_shp, quantity = "SD", 
  spatial_id_domains = "NAME_DISTRICT")


###################
## Section 4.2.3 ##
###################

# Extracting area estimates
HB_estimates <- extract(summ_beta)
head(HB_estimates$in_sample)
# Creating CSV with estimates
export(HB_estimates, file = "results.csv", type = "all")

#################
## Section 4.3 ##
#################

# Example of new variance function
gini_variance <- function(x){ x^2 * (1 - x^2) }

# Smoothing procedure
smoo <- smoothing(data = emilia_cs, direct_estimates = "hcr", area_id = "id",
  raw_variance = "vars", areas_sample_sizes = "n", var_function = NULL,
  method = c("ols"))
smoo
plot(smoo)

# Adding smoothed dispersion parameters to data
emilia_cs$smoo_phi <- smoo$phi
emilia_cs$smoo_vars <- smoo$vars

# Population shares for benchmarking
shares <-  emilia_cs$pop / sum(emilia_cs$pop)

bmk <- benchmark(x = summ_beta, bench = 0.13, share = shares, method = "raking")
bmk
plot(bmk)

# Map benchmarked estimates
map(x = bmk, spatial_df = emilia_shp, spatial_id_domains = "NAME_DISTRICT")

# Benchmarking on a subset of areas
subset <- c("RIMINI", "RICCIONE", "RUBICONE", "CESENA - VALLE DEL SAVIO")
pop <- emilia_cs$pop[emilia_cs$id %in% subset]
shares_subset <- pop / sum(pop)
bmk_subset <- benchmark(x = summ_beta, bench = 0.13, share = shares_subset,
  method = "raking", areas = subset)
bmk_subset
plot(bmk_subset)

# Seeing results
bmk_subset$bench_est
bmk_subset$post_risk


#################
## Section 4.4 ##
#################

data("emilia")
data("emilia_shp")

# Fit Beta spatio-temporal model
fit_ST <- fit_sae(formula_fixed = hcr ~ x, domains = "id", disp_direct = "vars",
  type_disp = "var", domain_size = "n", data = emilia, spatial_error = TRUE,
  spatial_df = emilia_shp, domains_spatial_df = "NAME_DISTRICT", 
  temporal_error = TRUE, temporal_variable = "year", max_treedepth = 15, 
  seed = 0

  ## ## NOTE: To reduce computation times you can uncomment the two lines
  ## ## below, reducing the number of iterations/chains.
  ## ,
  ## iter = 100, 
  ## chain = 1
)

# Model summaries
summ_ST <- summary(fit_ST)
summ_ST

shares <- aggregate(emilia$pop, list(emilia$year), function(x) x / sum(x))
shares <- as.vector(t(shares[,-1]))
bmk_ST <- benchmark(x = summ_ST, bench = 0.09, share = shares[1:38], 
  method = "raking", time = "2014")
bmk_ST
plot(bmk_ST)


## R Session Information
sessionInfo()