#################################################################################
##     R CODE FILE: plotEstimates.R
##
##         PROJECT: Evaluation of multiple-encounter population size
##                  estimators
##
##     DESCRIPTION: Contruct violin plots showing the distributions of the
##                  PSE estimates (Fig. 2)
##
##            NOTE: Estimates from M0 and the trend data-generating models are
##                  omitted
##
##           INPUT: ~/data/PSE_estimates_2021-11-11.rds
##
##          OUTPUT: ~/output/graphs/PSE_violin_plot_1K.png
##                  ~/output/graphs/PSE_violin_plot_10K.png
##                  ~/output/graphs/PSE_violin_plot_20K.png
##                  ~/output/graphs/Fig_2.png
##                  ~/output/graphs/Fig_2.tiff
##
##      WRITTEN BY: Steve Gutreuter                  E-mail:  sgutreuter@cdc.gov
##                  Statistics, Estimation and Modeling Team
##                  Division of Global HIV & TB
##                  Center for Global Health
##                  Centers for Disease Control & Prevention
#################################################################################
library(tidyverse)
library(ggplot2)
library(ggthemes)
library(scales)
library(svglite)
#################################################################################
## Define file paths
#################################################################################
basepath <- file.path(Sys.getenv("PROJ"), "PSE/PSEsim")
workpath <- file.path(basepath, "R")
datapath <- file.path(basepath, "data")
outpath <- file.path(basepath, "output")
graphs <- file.path(basepath, "output/graphs")
setwd(workpath)
#################################################################################
## Extract the estimates
#################################################################################
df0 <- readRDS(file.path(datapath, "PSE_estimates_2021-11-11.rds"))
df0 <- df0 %>%
    filter(!is.na(pmean) & Nest > 0 & !(gen.modl == "trend" |
                                        gen.modl == "M0")) %>%
    mutate(Nest = if_else(Ntrue == 1000 & Nest > 10^7, 10^7, Nest),
           Nest = if_else(Ntrue == 10000 & Nest > 10^7, 10^7, Nest),
           Nest = if_else(Ntrue == 20000 & Nest > 10^7, 10^7, Nest),) %>%
    mutate(ModLabel = case_when(
               ModelType == "LLM" ~ "LLM-AIC",
               ModelType == "BMA" ~ "LLM-BMA",
               ModelType == "LCM" ~ "BLCM"
           ),
           ModLabel =  factor(ModLabel,  ordered =  TRUE,
                              levels = c("LLM-AIC", "LLM-BMA", "BLCM" )),
           ModelType = factor(ModelType, ordered = TRUE,
                              levels = c("LLM", "BMA", "LCM")),

           Events = factor(Events, ordered = TRUE, levels = c(2, 3, 4, 5)),
           pmean = factor(pmean, ordered = TRUE,
                          levels = c("0.025", "0.050", "0.100", "0.150", "0.200")))
#################################################################################
## Construct violin plots of the estimates for Ntrue = 1,000
#################################################################################
df1 <- df0 %>%
    filter(Ntrue == 1000) %>%
    mutate(Nest = if_else(Nest < 1, 1, Nest))
p1 <- ggplot(df1, aes(Events, Nest)) +
    ylab(expression(hat(italic(N)))) + xlab("Encounter events") +
    facet_grid(pmean ~ ModLabel) +
    geom_hline(yintercept = 1000, size = 0.5, col = "gray") +
    geom_violin() +
    scale_y_log10(breaks = trans_breaks("log10", function(x) 10^x, n = 4),
                  limits = c(10^0,  10^7),
                  minor_breaks = waiver( ),
                  labels = trans_format("log10", math_format(10^.x))) +
    theme_bw( )
png(file.path(graphs, "PSE_violin_plot_1K.png"), width = 6.5, height = 6.5,
    units = "in", res = 600)
print(p1)
dev.off( )
#################################################################################
## Construct violin plots of the estimates for Ntrue = 10,000
#################################################################################
df2 <- df0 %>%
    filter(Ntrue == 10000) %>%
    mutate(Nest = if_else(Nest < 1, 1, Nest))
p2 <- ggplot(df2, aes(Events, Nest)) +
    ylab(expression(hat(italic(N)))) + xlab("Encounter events") +
    facet_grid(pmean ~ ModLabel) +
    geom_hline(yintercept = 10000, size = 0.5, col = "gray") +
    geom_violin() +
    scale_y_log10(breaks = trans_breaks("log10", function(x) 10^x, n = 4),
                  limits = c(10^0,  10^7),
                  minor_breaks = waiver( ),
                  labels = trans_format("log10", math_format(10^.x))) +
    theme_bw()
png(file.path(graphs, "PSE_violin_plot_10K.png"), width = 6.5, height = 6.5,
    units = "in", res = 600)
print(p2)
dev.off( )
#################################################################################
## Construct violin plots of the estimates for Ntrue = 20,000
#################################################################################
df3 <- df0 %>%
    filter(Ntrue == 20000) %>%
    mutate(Nest = if_else(Nest < 1, 1, Nest))
p3 <- ggplot(df3, aes(Events, Nest)) +
    ylab(expression(hat(italic(N)))) + xlab("Encounter events") +
    facet_grid(pmean ~ ModLabel) +
    geom_hline(yintercept = 20000, size = 0.5, col = "gray") +
    geom_violin() +
    scale_y_log10(breaks = trans_breaks("log10", function(x) 10^x, n = 4),
                  limits = c(10^0,  10^7),
                  minor_breaks = waiver( ),
                  labels = trans_format("log10", math_format(10^.x))) +
    theme_bw()
png(file.path(graphs, "PSE_violin_plot_20K.png"), width = 6.5, height = 6.5,
    units = "in", res = 600)
print(p3)
dev.off( )
################################################################################
## Two graphs on one page
################################################################################
vp1 <- grid::viewport(width = unit(3, "in"), height = unit(6, "in"), just = "left")
vp2 <- grid::viewport(width = unit(3, "in"), height = unit(6, "in"), just = "right")
png(file.path(graphs, "Fig_2.png"), width = 6.5, height = 6.5,
    units = "in", res = 600)
print(p1, vp = vp2)
print(p3, vp = vp1)
dev.off( )
tiff(file.path(graphs, "Fig_2.tiff"), width = 6.5, height = 6.5,
    units = "in", res = 600, compression = "lzw")
print(p1, vp = vp2)
print(p3, vp = vp1)
dev.off( )
################################################################################
## Log session information
################################################################################
cat("Execution ended: ", as.character(Sys.time()), "\n")
sessionInfo()
#################################  END of FILE  #################################
