rm(list = ls())

setwd("[enter directory where input data files have been stored (see readme.txt)]")

# Packages
library(readxl)
library(dplyr)
library(zoo)
library(reshape2)
library(knitr)
library(bvarsv3)

# Seed
set.seed(20200503)

# Load functions
source("procs_rev.R")

# Settings
# TRUE/FALSE switches for estimating the CMM model and the combination parameters. 
# If both switches are set to FALSE, the script merely collects the previous results
# and saves them in a joint file ("results_[date].RData")
run_bvar <- TRUE   # If TRUE, estimate CMM model anew
fit_combs <- TRUE  # If TRUE, estimate combination parameters anew
start_ind <- 20 # effective only if run_bvar is TRUE
n_hor_rw <- 5 # nr of horizons available for rw model
type_rw <- "first_release" # type of UCSV model
win_comb <- 10 # minimal window size for combinations
thin_mcmc <- 10 # use of thinning (only for estimating combination parameters)
w_all <- seq(from = 0, to = 1, by = .01) # grid sequence for weights (used for plots)

# File names for cmm and combinations
today <- Sys.Date()
cmm_name <- "[specify_name_here].RData" # file to be loaded (if run_bvar is FALSE). Note: 

# Load real-time data
dat <- read_excel("PQvQd.xlsx", na = "#N/A") %>%
  prepare_rt(sub_nm = "P")

# Load and prepare SPF data
spf0 <- read_excel("Individual_PGDP.xlsx", na = "#N/A") %>%
  prepare_spf(sub_nm = "PGDP") 

# Construct SPF forecast revisions as in CMM
spf_rev <- spf0 %>% 
  mutate(spf_previous = rev_helper(origin_date, horizon, spf), 
         spf_revision = spf - spf_previous, 
         horizon = paste0("rev", horizon)) %>% 
  dcast(origin_date~horizon, value.var = "spf_revision")

# Construct nowcast errors
spf_nc <- spf0 %>% filter(horizon == 1) %>% transmute(date, spf)
nc_e <- dat %>% na.omit %>% group_by(date) %>% 
  summarise(value_fv = value[vint == min(vint)],
            vint = min(vint)) %>%
  ungroup %>% merge(spf_nc) %>%
  # change date to one quarter later (which is when nowcast error becomes available)
  transmute(date, origin_date = date + .25, value_fv, spf, e_prev = value_fv - spf) %>%
  arrange(origin_date)

# merge both matrices to get CMM data format
dat_cmm <- merge(spf_rev, nc_e[, c("origin_date", "e_prev")]) %>%
  (function(x) x[, c("origin_date", "e_prev", paste0("rev", 1:4))]) %>%
  arrange(origin_date)

# Load existing bvar results (list "all")
# If run_bvar is TRUE: Remove existing bvar results in "all"
if (run_bvar){
  all <- list()  
} else {
  load(cmm_name)
}

# Origin dates for CMM model
origin_dates_cmm <- seq(from = 1973.75, to = 2018.5, by = .25)

# load UCSV results
load(paste0("ucsv_forecasts_", type_rw, ".RData"))

# get associated forecast (origin) dates
origin_dates_rw <- sapply(all_rw, function(x) x$origin_date)

# consider origin dates that are covered by both spf and rw
dts <- intersect(origin_dates_cmm, origin_dates_rw)
inds <- which(dat_cmm$origin_date %in% dts)  

# Nr of MCMC draws
n_mcmc <- c(sapply(all_rw, function(z) c(nrow(z$m_draws), nrow(z$v_draws)))) %>%
  unique %>% unlist
seq_mcmc <- (1:n_mcmc)[(1:n_mcmc) %% thin_mcmc == 0]
if (length(n_mcmc) > 1) stop("Nr of MCMC draws not coherent")

# Draws in another format (used for combinations)
comb_draws_cmm <- comb_draws_rw <- vector("list", n_hor_rw)
for (hh in 1:n_hor_rw){
  if (!run_bvar){
    comb_draws_cmm[[hh]]$m_draws <- sapply(all, function(x) x$m_draws[,hh])
    comb_draws_cmm[[hh]]$v_draws <- sapply(all, function(x) x$v_draws[,hh])  
  } else {
    comb_draws_cmm[[hh]]$m_draws <- comb_draws_cmm[[hh]]$v_draws <- 
      matrix(NA, n_mcmc, length(origin_dates_cmm))
  }
  comb_draws_rw[[hh]]$m_draws <- sapply(all_rw, 
                                          function(x) x$m_draws[1,hh]) %>%
    matrix(nrow = 1)
  comb_draws_rw[[hh]]$v_draws <- sapply(all_rw, 
                                          function(x) x$v_draws[1,hh]) %>%
    matrix(nrow = 1)
}

# Loop over time
pred_all <- data.frame()
comb_all <- data.frame()
for (vv in inds){
  # Origin date
  od <- dat_cmm$origin_date[vv]
  print(od)
  # Pick corresponding rw forecasts
  sel_ind <- which(origin_dates_rw == od)
  if (length(sel_ind) == 1){
    # Get USCV predictions
    rw_m_draws <- all_rw[[sel_ind]]$m_draws
    rw_v_draws <- all_rw[[sel_ind]]$v_draws
  } else {
    rw_m_draws <- rw_v_draws <- matrix(NA, 1, n_hor_rw)
  }
  # Get realizations
  rlz <- nc_e %>% filter(date %in% (od + (0:4)/4)) %>%
    arrange(date) %>% select(value_fv) %>% unname %>% unlist
  if (length(rlz) < 5){
    rlz <- c(rlz, rep(NA, 5 - length(rlz)))
  }
  # Expand data frame
  pred_all <- rbind(pred_all, 
                    data.frame(origin_date = rep(od, 3*n_hor_rw),
                               target_date = rep(od + (0:(n_hor_rw-1))/4, 3), 
                               method = rep("rw", 3*n_hor_rw), 
                               type = rep(c("mean", "e", "v"), each = n_hor_rw), 
                               value = c(colMeans(rw_m_draws), 
                                         rlz[1:n_hor_rw] - colMeans(rw_m_draws), 
                                         apply(rw_m_draws, 2, var2) + colMeans(rw_v_draws))))
  # Run BVAR (if selected)
  if (run_bvar){
    # Print progress
    print(paste("Now at origin date", od))
    # Fit model
    y_cmm <- dat_cmm[1:vv, -1] %>% as.matrix
    fit_cmm <- bvarsv3:::cmm(y_cmm)
    e_cmm <- fit_cmm$e_draws
    # Get SPF forecasts
    spf_mean <- spf0 %>% filter(origin_date == od) %>%
      arrange(horizon) %>% select(spf) %>% unname %>% unlist %>%
      matrix(nrow = nrow(e_cmm), ncol = 5, byrow = TRUE)
    # Get forecast draws
    fc_cmm <- spf_mean + e_cmm
    # Get forecast mean draws
    cmm_m_draws <- spf_mean + fit_cmm$e_m_draws
    # Get forecast variance draws
    cmm_v_draws <- fit_cmm$e_v_draws
    # Enter into big list
    all[[(vv-start_ind + 1)]] <- list(origin_date = od, 
                                      fc_draws = fc_cmm,
                                      m_draws = cmm_m_draws,
                                      v_draws = cmm_v_draws)
    # Draws in different format for combinations
    for (hh in 1:5){
      comb_draws_cmm[[hh]]$m_draws[, which(origin_dates_cmm == od)] <- cmm_m_draws[,hh]
      comb_draws_cmm[[hh]]$v_draws[, which(origin_dates_cmm == od)] <- cmm_v_draws[,hh]
    }
  } else {
    cmm_m_draws <- all[[which(origin_dates_cmm == od)]]$m_draws
    cmm_v_draws <- all[[which(origin_dates_cmm == od)]]$v_draws
  }
  # Add bvar predictions to data frame
  pred_all <- rbind(pred_all, 
                    data.frame(origin_date = rep(od, 15),
                               target_date = rep(od + (0:4)/4, 3), 
                               method = rep("cmm", 15), 
                               type = rep(c("mean", "e", "v"), each = 5), 
                               value = c(colMeans(cmm_m_draws), rlz - colMeans(cmm_m_draws),
                                         apply(cmm_m_draws, 2, var2) + colMeans(cmm_v_draws))))
  # Simple density combinations
  for (hh in 1:n_hor_rw){
    # Loop over weight choices
    for (ww in w_all){
      draws1 <- list(m_draws = cmm_m_draws[,hh], 
                     v_draws = cmm_v_draws[,hh])
      draws2 <- list(m_draws = rw_m_draws[,hh], 
                     v_draws = rw_v_draws[,hh])
      cmb_lp <- lp_draws(draws1, draws2, rlz[hh], w = ww, center = FALSE)
      comb_all <- rbind(comb_all, 
                        data.frame(origin_date = od, h = hh, method = "LP",
                                   as.data.frame(cmb_lp)))
      cmb_clp <- lp_draws(draws1, draws2, rlz[hh], w = ww, center = TRUE)
      comb_all <- rbind(comb_all, 
                        data.frame(origin_date = od, h = hh, method = "CLP",
                                   as.data.frame(cmb_clp)))  
    }
  }
  # Fitted density combinations
  if (fit_combs){
    for (hh in 1:n_hor_rw){
      draws1 <- draws2 <- list()
      # Past observations
      # note that date denotes target date here
      past_obs <- nc_e %>% filter((date %in% (dts + (hh-1)*.25)) & (date <= od - .25)) %>%
        select(date, rlz = value_fv)
      # Past origin dates
      past_ods <- intersect(dts, past_obs$date - (hh-1)*.25)
      # Get corresponding forecast draws
      dt_sel_cmm <- origin_dates_cmm %in% past_ods
      dt_sel_rw <- origin_dates_rw %in% past_ods
      # Fit combinations if enough data is available
      if (all(sapply(list(dt_sel_cmm, dt_sel_rw), sum) >= win_comb)){
        draws1$m_draws <- comb_draws_cmm[[hh]]$m_draws[seq_mcmc, dt_sel_cmm] 
        draws1$v_draws <- comb_draws_cmm[[hh]]$v_draws[seq_mcmc, dt_sel_cmm] 
        draws2$m_draws <- comb_draws_rw[[hh]]$m_draws[1, dt_sel_rw, drop = FALSE] 
        draws2$v_draws <- comb_draws_rw[[hh]]$v_draws[1, dt_sel_rw, drop = FALSE]
        # Run SLP and BLP
        for (cc in 1:4){
          if (cc == 1){
            ff <- est_SLP
            c_choice <- FALSE
            nm <- "SLP"
          } else if (cc == 2){
            ff <- est_BLP
            c_choice <- FALSE
            nm <- "BLP"
          } else if (cc == 3){
            ff <- est_SLP
            c_choice <- TRUE
            nm <- "SCLP"
          } else if (cc == 4){
            ff <- est_BLP
            c_choice <- TRUE
            nm <- "BCLP"
          }
          cmb <- ff(draws1, draws2, rlz = past_obs$rlz, w = .5, 
                    new_draws1 = list(m_draws = cmm_m_draws[,hh], 
                                      v_draws = cmm_v_draws[,hh]),
                    new_draws2 = list(m_draws = rw_m_draws[,hh], 
                                      v_draws = rw_v_draws[,hh]),
                    new_rlz = rlz[hh], center = c_choice)
          comb_all <- rbind(comb_all, 
                            data.frame(origin_date = od, h = hh, method = nm, 
                                       as.data.frame(cmb)))
        } # end loop over slp/blp
      } # end if check for data availability
    } # end loop over horizons
  } # end if check for comb fitting
} # end loop over time

# Save results that enter figures and tables
results <- list(pred_all = pred_all, comb_all = comb_all)
save(results, file = paste0("results_", today, ".RData"))