library(magrittr)
library(data.table)
library(GGally)
library(ggplot2)
library(grid)
library(gridExtra)
library(doBy)


doWilcoxon <- function(d, byVersion = F, byVariant = F){
  if(byVersion){
    d = dcast(d, Location+Week~Version, value.var='Score')
  } else if(byVariant){
    d = dcast(d, Location+Week~Variant, value.var='Score')
  }
  
  d = as.matrix(subset(d, select = -c(Location, Week)))

  
  two <- matrix(NA, nrow = ncol(d)-1, ncol = ncol(d)-1, dimnames = list(tail(colnames(d), -1), head(colnames(d), -1)))
  greater <- matrix(NA, nrow = ncol(d)-1, ncol = ncol(d)-1, dimnames = list(tail(colnames(d), -1), head(colnames(d), -1)))
  lesser <- matrix(NA, nrow = ncol(d)-1, ncol = ncol(d)-1, dimnames = list(tail(colnames(d), -1), head(colnames(d), -1)))
  
  for(i in 1:(ncol(d)-1)){
    for(j in (i+1):ncol(d)){
      x <- d[,i]
      y <- d[,j]
      
      two[j-1, i] <- wilcox.test(x, y, paired = T, alternative = 'two.sided')$p.value %>% round(2)
      greater[j-1, i] <- wilcox.test(x, y, paired = T, alternative = 'greater')$p.value %>% round(2)
      lesser[j-1, i] <- wilcox.test(x, y, paired = T, alternative = 'less')$p.value %>% round(2)
    }
  }
  
  list(g = greater)
}


inDir <- 'C:/cpid/ili/Papers/Self/s1617/JRSI/rev1/SI/Archive/'
baseline <- data.table(REGION = c('US National', paste0('HHS Region ', 1:10)),
                       Value = c(2.1, 1.4, 3.0, 2.2, 1.7, 1.9, 4.1, 1.8, 1.4, 2.5, 1.1))




load(paste0(inDir, 'ScoresErrors.RData'))


# Table 1a - probabilistic forecasts
tbl1a <- dScore[Variant == 'Realtime' & ScoreType == 'Log-Score', 
                j = list(Score = round(sum(Score))), 
                by = list(Target, Version)]
dcast(tbl1a, Target~Version, value.var = 'Score')



# Table 1b - point forecasts
tbl1b <-  dScore[Variant == 'Realtime' & ScoreType != 'Log-Score', 
                 j = list(Score = round(mean(Score, na.rm = T), 3)), 
                 by = list(Target, Version, ScoreType)]
dcast(tbl1b, Target~Version, value.var = 'Score')




# For submitted version, wilcoxon p of errors in point prediction
tbl2 <- dScore[ScoreType != 'Log-Score' & Variant == 'Realtime',
                j = list(Location, Target, Week, Version, Score)]
splitBy(~Target, tbl2[grepl('wk ahead', Target)]) %>%
  lapply(FUN = doWilcoxon, byVersion = T, byVariant = F)




# Table 3 - log scores by variant
tbl3 <- dScore[ScoreType == 'Log-Score', 
                j = list(Score = round(sum(Score))), 
               by = list(Target, Version, Variant)]
dcast(tbl3, Version+Target~Variant, value.var = 'Score')




# Table 4  - errors in point forecasts by variant
tbl4 <- dScore[ScoreType != 'Log-Score', 
                   j = list(Score = round(mean(Score, na.rm = T), 3)), 
               by = list(Target, Version, Variant, ScoreType)]
dcast(tbl4, Version+Target~Variant, value.var = 'Score')




# Figure 1 - Log score of the three methods
temp <- dScore[Variant == 'Realtime' &  ScoreType == 'Log-Score' & Target != '1 wk ahead',
               j = list(Score = sum(Score)),
               by = list(Target, Version, Week)]
ggplot(temp) +
  geom_line(aes(x = as.integer(Week), y = Score, color = Version), size = 1.02) + 
  facet_wrap(~Target, ncol = 3) +
  labs(x = 'Week of season (1=MMWR43)', y = 'Log score', color = '', 
       title = 'Scores of forecasts in real-time, by method and target') +
  theme_bw() + 
  theme(legend.position = 'bottom', 
        legend.text = element_text(size = 16), legend.title = element_text(size = 18),
        strip.text = element_text(size = 18), title =  element_text(size = 18), 
        axis.text = element_text(size = 16), axis.title = element_text(size = 18))




# Figure 2 - Boxplot of cumulative score by method, target and region
fig2 <- dScore[Variant == 'Realtime' &  ScoreType == 'Log-Score' & Target != '1 wk ahead',
                       j = list(Score = sum(Score)),
                       by = list(Target, Version, Location)]
fig2$Version.jitter <- jitter(as.numeric(factor(fig2$Version)))
fig2$Location <- factor(fig2$Location, levels = c('US National', paste0('HHS Region ', 1:10)), 
                    labels = c('US', paste0('R', 1:10)))
finalSumScoreMean <- fig2[, j = list(Score = mean(Score)), by = list(Version, Target)]


ggplot(fig2) +
  geom_boxplot(aes(x = Version, y = Score, color = Version), width = 0.25, outlier.color = NA, size = 0.5, alpha = 0.25) + 
  geom_point(aes(x = Version.jitter, y = Score, color = Version), size = 2) +
  geom_text(aes(x = Version.jitter, y = Score, color = Version, label = Location), size = 5, nudge_x = 0.1, nudge_y = 0.1) +
  geom_point(data = finalSumScoreMean[Target != 'ALL' & Target != '1 wk ahead' ],
             aes(x = Version, y = Score, color = Version), shape = 17, size = 3) +
  geom_text(data = finalSumScoreMean[Target != 'ALL' & Target != '1 wk ahead' ],
            aes(x = Version, y = Score, label = round(Score, 2)), size = 6, nudge_x = 0.25, nudge_y = 0.1) +
  facet_wrap(~Target, ncol = 3) + labs(title  = 'Cumulative score at end of season', y = 'Log score', x = '', color = '') +
  theme_bw() + theme(legend.position = 'bottom', 
                     legend.text = element_text(size = 16), legend.title = element_text(size = 18),
                     strip.text = element_text(size = 18), title =  element_text(size = 18), 
                     axis.text = element_text(size = 16), axis.title = element_text(size = 18)) +
  guides(color = F)





# Figure 3 - Rolling sum of log score of *all 3 methods*, by variant and target
fig3 <- dScore[Variant != 'Stable-Nowcast' &  ScoreType == 'Log-Score',
               j = list(Score = sum(Score)),
               by = list(Target, Version, Variant, Week)]
fig3 <- fig3[, j = list(Week, Score = cumsum(Score)), by = list(Target, Version, Variant)]
fig3$Variant <- factor(fig3$Variant, levels = c('Baseline', 'Baseline-Nowcast', 'Baseline+Post',  'Stable', 'Stable-Nowcast', 'Realtime'),
                        labels = c('Baseline', 'Baseline w/o nowcast', 'Baseline w/ post-processing',  'Stable ILI', 'Stable-Nowcast', 'Realtime'))
fig3$Target <- factor(fig3$Target, levels = c('Season onset', 'Season peak week', 'Season peak percentage', paste0(1:4, ' wk ahead')),
                       labels = c('Onset', 'Peak week', 'Peak percentage', paste0(1:4, ' week ahead')))
fig3$Linetype <- ifelse(fig3$Variant %in% c('Baseline w/o nowcast', 'Baseline w/ post-processing'), '0', '1')

ggplot(fig3[Variant != 'Baseline']) +
  geom_line(aes(x = as.integer(Week), y = Score, color = Variant, linetype = Linetype), size = 1.01, alpha = 0.8) + 
  geom_point(aes(x = as.integer(Week), y = Score, color = Variant), alpha = 0.8, shape = 1, size = 0.9) + 
  geom_line(data = fig3[Variant == 'Baseline'], aes(x = as.integer(Week), y = Score, color = Variant, linetype = Linetype), 
            size = 1.01) + 
  scale_color_manual(values = c('black', 'red', 'orange', 'green', 'blue')) +
  scale_linetype_manual(values = c('dashed', 'solid')) +
  facet_grid(Version~Target, scales = 'free_y') +
  labs(x = 'Week of season (1=MMWR43)', y = 'Log score, cumulative', 
       title = 'Scores of DYN/STAT/SE forecasts, by variant and target', color = '') +
  theme_bw() + 
  theme(legend.position = 'bottom', 
        legend.text = element_text(size = 16), legend.title = element_text(size = 16),
        strip.text = element_text(size = 18), title =  element_text(size = 18), 
        axis.text = element_text(size = 16), axis.title = element_text(size = 18)) +
  guides(color=guide_legend(nrow=2, byrow=TRUE), linetype = F)




# Figure 4 - 1-week ahead scores Baseline vs X
fig4 <- dScore[ScoreType == 'Log-Score' & Variant != 'Stable-Nowcast']
fig4$Variant <- factor(fig4$Variant, levels = c('Baseline', 'Baseline-Nowcast',  'Stable', 'Realtime', 'Baseline+Post'),
                       labels = c('Baseline', 'Baseline w/o nowcast',  'Stable ILI', 'Realtime', 'Baseline w/ post-processing'))
fig4$Target <- factor(fig4$Target, levels = c('Season onset', 'Season peak week', 'Season peak percentage', paste0(1:4, ' wk ahead')),
                      labels = c('Onset', 'Peak week', 'Peak percentage', paste0(1:4, ' week ahead')))

fig4 <- merge(fig4[Variant != 'Baseline', j = list(Location, Variant, Version, Week, Target, Score)],
              fig4[Variant == 'Baseline', j = list(Location, Version, Week, Target, Ref = Score)],
              by = c('Location', 'Version', 'Week', 'Target'))


ggplot(fig4[Target == '1 week ahead' & Variant != 'Baseline w/ post-processing' & Variant != 'Realtime']) + 
  geom_abline(slope = 1, intercept = 0, linetype = 'dashed', size=1.025) +
  geom_point(aes(x = exp(Ref), y = exp(Score), color = as.integer(Week), shape = Version), size = 4, alpha = 0.9) +
  scale_shape_manual(values = c(16, 17, 18)) +
  scale_color_gradient2(low = 'darkblue', mid = 'darkgrey', high = 'red', midpoint = 15) +
  facet_grid(Variant~Version) + 
  labs(color = 'Week', shape = '', x = 'Baseline score, exp', y = 'Variant score, exp') +
  theme_bw() + theme(legend.position = 'bottom',
                     legend.text = element_text(size = 12), legend.title = element_text(size = 16),
                     strip.text = element_text(size = 18), title =  element_text(size = 18), 
                     axis.text = element_text(size = 16), axis.title = element_text(size = 18))





# Table S1 - Wilcoxon on errors by variant
tblS1 <- dScore[ScoreType != 'Log-Score', j = list(Location, Target, Week, Variant, Version, Score)]
splitBy(~Version+Target, tblS1[grepl('wk ahead', Target) & Variant != 'Baseline+Post' & Variant != 'Stable-Nowcast']) %>%
  lapply(FUN = doWilcoxon, byVersion = F, byVariant = T)