# Helper to transform time series data
tf_dat <- function(x, type = 1){
  if (type == 1){
    400*(log(x[-1]) - log(x[-length(x)]))
  } else if (type == 2){
    (x[-1] - x[-length(x)])
  }
}

# Helper to process real-time data
rev_helper <- function(origin_date, horizon, spf){
  out <- rep(NA, length(origin_date))
  ind <- origin_date != "1968 Q4" & horizon != 5
  for (jj in 1:length(out)){
    if (ind[jj]){
      sel_tmp <- which( (origin_date == origin_date[jj] - .25) & (horizon == horizon[jj] + 1) )
      if (length(sel_tmp) > 1){
        stop()
      } else if (length(sel_tmp) == 1){
        out[jj] <- spf[sel_tmp]  
      }
    }
  }
  out
}

# Another helper to prepare real-time data
prepare_rt <- function(dat, type = 1, sub_nm = "ROUTPUT"){
  dat %>%
  (function(x) cbind(x[-1,1], 
                     apply(x[,-1], 2, function(z) tf_dat(z, type = type)))) %>%
    melt(id.vars = c("DATE")) %>% 
    mutate(variable = gsub(sub_nm, "", variable), 
           vint_yr = substr(variable, 1, 2), 
           vint_q = substr(variable, 4, 4), 
           vint_yr = if_else(vint_yr < 30, paste0("20", vint_yr),
                             paste0("19", vint_yr))) %>%
    transmute(date = as.yearqtr(gsub(":Q", "-", DATE)), 
              vint = as.yearqtr(paste0(vint_yr, "-", vint_q)), 
              value)
}

# Helper to prepare SPF forecast data
prepare_spf <- function(dat, type = 1, sub_nm = "RGDP"){
  dat %>% mutate(date = as.yearqtr(paste0(YEAR, "-", QUARTER))) %>%
    select(one_of(c("date", "ID", paste0(sub_nm, 1:6)))) %>%
    (function(x) cbind(x[,1:2], 
                       400*(log(x[,4:8])-log(x[,3:7])))) %>%
    melt(id.vars = c("date", "ID")) %>% 
    mutate(aux = as.character(variable), 
           horizon = as.integer(substr(aux, nchar(aux), nchar(aux))) - 2) %>%
    transmute(date, horizon, id = ID, value) %>%
    na.omit %>% 
    group_by(date, horizon) %>% summarise(spf = mean(value)) %>%
    ungroup %>%
    transmute(origin_date = date, 
              date = date + .25*horizon, horizon = horizon + 1, spf)
}

# Helper to find grid for plotting a density
find_support <- function(f, eps = 1e-5, n = 1001){
  grid0 <- seq(from = -500, to = 500, length.out = 10001)
  data.frame(grid0) %>% mutate(fg = f(grid0)) %>%
    filter(fg > eps) %>% select(grid0) %>% unlist %>%
    unname %>% range
}

# Helper to edit output of knitr:kable
rm_tabular <- function(x){
  l1 <- gregexpr("midrule", x)[[1]][1]
  l2 <- gregexpr("bottomrule", x)
  l2 <- l2[[length(l2)]][1]
  substr(x, l1+7,l2-2)
}

# Helper to print tex table
table_helper <- function(df, digits = 3, bold = TRUE){
  c1 <- df[,1]
  mat <- df[,-1]
  out <- paste0(c1, "&")
  for (jj in 1:ncol(mat)){
    tmp <- format(round(mat[,jj], digits = digits), nsmall = digits)
    if (bold){
      ind <- which.min(tmp)
      tmp[ind] <- paste0("\\textbf{", tmp[ind], "}")  
    }
    if (jj < ncol(mat)){
      out <- paste0(out, tmp, "&")  
    } else {
      out <- paste0(out, tmp, "\\\\")
    }
  }
  out
}

# Helper to write tex output
writeLines2 <- function(x, path){
  con <- file(path)
  on.exit(close(con))
  writeLines(x, con)
}

# Function to compute MSE-optimal weight (n = 2 forecasts)
get_weight <- function(e1, e2){
  ff <- function(w) sum((w*e1+(1-w)*e2)^2)
  optimize(ff, interval = c(0, 1))$minimum
}

# Function to relocate density at a new mean
relocate_f <- function(f, new_mean = 0){
  old_mean <- integrate(function(z) f(z)*z, lower = -Inf, 
                        upper = Inf)$value
  shift <- new_mean - old_mean
  new_f <- function(z) f(z - shift)
  new_mean_check <- integrate(function(z) new_f(z)*z, -Inf, Inf)
  if (abs(new_mean_check$value - new_mean) > 1e-6){
    stop("relocation did not work")
  }
  list(shift = shift, new_f = new_f, 
       new_mean_check = new_mean_check)
}

# Function to plot densities 
plot_dens <- function(f1, f2, fcomb, supp, l_wd = 1.2, n_pt = 1001,
                      col = "green4", ft_size = 14){
  ggplot(data.frame(x = 0), aes(x = x)) + 
    stat_function(fun = f1, col = "darkgrey", lty = 2, lwd = l_wd, n = n_pt) + 
    stat_function(fun = f2, col = "darkgrey", lty = 2, lwd = l_wd, n = n_pt) + 
    stat_function(fun = fcomb, col = col, lty = 1, lwd = l_wd, n = n_pt) + 
    xlim(supp[1], supp[2]) + theme_minimal(base_size = ft_size) + 
    ylab("Density")
}

# Generate density for (simplified) two-component mixture
# Components have means that are symmetric around cent, and have the same variance; 
# mixture weight is 0.5
gen_dmix <- function(m, s, cent){
  f <- function(x){
    .5*(dnorm(x, mean = cent - m, sd = s) + 
          dnorm(x, mean = cent + m, sd = s))
  }
  f
}

# Function for Diebold-Mariano test
dm_test <- function(sc1, sc2){
  d <- sc1 - sc2
  l <- lm(d~1)
  coeftest(l, vcov. = NeweyWest)
}

# Generator function for density and CDF (based on MCMC draws)
gen_Ff <- function(m_draws, v_draws){
  force(m_draws)
  force(v_draws)
  F <- (function(z, c) mean(pnorm(z, mean = m_draws,
                               s = c*sqrt(v_draws)))) %>%
    Vectorize
  f <- (function(z, c) mean(dnorm(z, mean = m_draws,
                               s = c*sqrt(v_draws)))) %>%
    Vectorize
  list(F = F, f = f)
}
var2 <- function(x){
  n <- length(x)
  if (n == 1){
    return(0)
  } else {
    return(((n-1)/n)*var(x))
  } 
}
    
# Function to compute mean and variance of mixture
get_mv <- function(draws1, draws2, w = .5, c = 1){
  # Means of two components
  means <- sapply(list(draws1, draws2), function(z) mean(z$m_draws))
  # Overall mean
  m <- sum(c(w, 1-w)*means)
  # Disagreement among two means
  d <- w*(means[1]-m)^2 + (1-w)*(means[2]-m)^2
  # Variances of two components
  variances <- sapply(list(draws1, draws2), function(z) var2(z$m_draws)+ (c^2)*mean(z$v_draws))
  # Overall variance
  v <- d + sum(c(w, 1-w)*variances)
  list(m = m, v = v)
}

# Helper function to handle MCMC draws
loop_helper <- function(w, draws1, draws2, rlz, ff = pnorm, c = 1){
  n <- length(rlz)
  out <- rep(NA, n)
  for (jj in 1:n){
    out[jj] <- w*mean(ff(rlz[jj], mean = draws1$m_draws[,jj],
                         sd = c*sqrt(draws1$v_draws[,jj]))) + 
      (1-w)*mean(ff(rlz[jj], mean = draws2$m_draws[,jj],
                    sd = c*sqrt(draws2$v_draws[,jj])))
  }
  out
}

# Function to estimate Spread-adjusted Linear pool (SLP) parameters
est_SLP <- function(draws1, draws2, rlz, w = NULL, 
                    new_draws1 = NULL, new_draws2 = NULL, new_rlz = NULL, 
                    center = FALSE){
  # Check whether all inputs have same length
  aux <- c(sapply(draws1, ncol), 
           sapply(draws2, ncol), 
           length(rlz))
  if (any(aux != aux[1])){
    stop()
  } else {
    n <- aux[1]
  }
  if (center & is.null(w)){
    stop("Centering only implemented for fixed weight")
  }
  # Center draws (if selected)
  if (center){
    aux1 <- center_draws(draws1$m_draws, draws2$m_draws, w)
    draws1$m_draws <- aux1[[1]]
    draws2$m_draws <- aux1[[2]]
    if (!is.null(new_draws1)){
      aux2 <- center_draws(new_draws1$m_draws, new_draws2$m_draws, w)
      new_draws1$m_draws <- aux2[[1]]
      new_draws2$m_draws <- aux2[[2]]
    }
  }
  if (!is.null(w)){
    # Create objective function based on inputs
    ff <- function(c){
      ec <- exp(c)
      # Get PDF values at realizations
      rf <- loop_helper(w = w, draws1 = draws1, draws2 = draws2,
                        rlz = rlz, ff = dnorm, c = ec)
      -sum(log(rf))
    }
    opt <- optimize(Vectorize(ff), interval = c(-5, 3))
    # Estimated value of c (spread parameter to scale sd of input dists)
    c <- exp(opt$minimum)  
  } else {
    ff <- function(c){
      ec <- exp(c[1])
      w <- pnorm(c[2])
      # Get PDF values at realizations
      rf <- loop_helper(w = w, draws1 = draws1, draws2 = draws2,
                        rlz = rlz, ff = dnorm, c = ec)
      -sum(log(rf))
    }
    opt <- optim(rep(0, 2), ff)
    # Estimated value of c
    c <- exp(opt$par[1]) 
    # Estimated weight on first component
    w <- pnorm(opt$par[2]) 
  }
  # If new forecast draws are provided: Compute mean and variance of fitted comb
  if (!is.null(new_draws1) & !is.null(new_draws2) & !is.null(new_rlz)){
    mv <- get_mv(new_draws1, new_draws2, w = w, c = c)  
    m_c <- mv$m
    v_c <- mv$v
    e2_c <- (new_rlz - m_c)^2
    dss_c <- dss(v_c, e2_c)
    # Density at realization: First component
    aux1 <- mean(dnorm(new_rlz, mean = new_draws1$m_draws, 
                       sd = c*sqrt(new_draws1$v_draws)))
    # Density at realization: Second component
    aux2 <- mean(dnorm(new_rlz, mean = new_draws2$m_draws, 
                       sd = c*sqrt(new_draws2$v_draws)))
    # Log score
    logs_c <- log(w*aux1 + (1-w)*aux2)
  } else {
    m_c <- v_c <- e2_c <- dss_c <- logs_c <- NA
  }
  list(w = w, par1 = c, par2 = NA, mean = m_c, variance = v_c, e2 = e2_c, 
       dss = dss_c, logs = logs_c)
}

# Function for Dawid-Sebastiani score
dss <- function(v, e2){
  .5*(log(2*pi) + e2/v + log(v))
}


# Function to estimate Beta-adjusted Linear pool (BLP) parameters
est_BLP <- function(draws1, draws2, rlz, w = NULL, 
                    new_draws1 = NULL, new_draws2 = NULL, new_rlz = NULL, 
                    center = FALSE){
  # Check whether all inputs have same length
  aux <- c(sapply(draws1, ncol), 
           sapply(draws2, ncol), 
           length(rlz))
  if (any(aux != aux[1])){
    stop()
  } else {
    n <- aux[1]
  }
  if (center & is.null(w)){
    stop("Centering only implemented for fixed weight")
  }
  # Center draws (if selected)
  if (center){
    aux1 <- center_draws(draws1$m_draws, draws2$m_draws, w)
    draws1$m_draws <- aux1[[1]]
    draws2$m_draws <- aux1[[2]]
    if (!is.null(new_draws1)){
      aux2 <- center_draws(new_draws1$m_draws, new_draws2$m_draws, w)
      new_draws1$m_draws <- aux2[[1]]
      new_draws2$m_draws <- aux2[[2]]
    }
  }
  # Proceed with combination 
  if (!is.null(w)){
    # Get PDF and CDF values at realizations
    rF <- loop_helper(w = w, draws1 = draws1, draws2 = draws2,
                      rlz = rlz, ff = pnorm, c = 1)
    # Create objective function based on inputs
    ff <- function(c){
      ec <- exp(c)
      out <- log(dbeta(rF, ec[1], ec[2]))
      -sum(out)
    }
    opt <- optim(rep(0, 2), ff)
    ab <- exp(opt$par)  
  } else {
    ff <- function(c){
      ec <- exp(c[1:2])
      w <- pnorm(c[3])
      # Get CDF and PDF values at realizations
      rF <- loop_helper(w = w, draws1 = draws1, draws2 = draws2,
                        rlz = rlz, ff = pnorm, c = 1)
      rf <- loop_helper(w = w, draws1 = draws1, draws2 = draws2,
                        rlz = rlz, ff = dnorm, c = 1)
      out <- log(rf) + log(dbeta(rF, ec[1], ec[2]))
      -sum(out)
    }
    opt <- optim(rep(0, 3), ff)
    ab <- exp(opt$par[1:2])
    w <- pnorm(opt$par[3]) 
  }
  # If new forecast draws are provided: Compute mean and variance of fitted comb
  if (!is.null(new_draws1) & !is.null(new_draws2) & !is.null(new_rlz)){
    # Density of BLP
    f_blp0 <- function(z){
      f <- w*mean(dnorm(z, mean = new_draws1$m_draws, 
                        sd = sqrt(new_draws1$v_draws))) + 
        (1-w)*mean(dnorm(z, mean = new_draws2$m_draws, 
                         sd = sqrt(new_draws2$v_draws)))
      F <- w*mean(pnorm(z, mean = new_draws1$m_draws, 
                        sd = sqrt(new_draws1$v_draws))) + 
        (1-w)*mean(pnorm(z, mean = new_draws2$m_draws, 
                         sd = sqrt(new_draws2$v_draws)))
      dens <- dbeta(F, ab[1], ab[2])*f
      okay <- is.finite(dens) & !(is.na(dens)) & (dens >= 0)
      if (!okay){
        dens <- 0
      } 
      dens
    }
    f_blp <- Vectorize(f_blp0)
    m_c <- integrate(function(z) f_blp(z)*z, lower = -Inf, upper = Inf)$value
    v_c <- integrate(function(z) f_blp(z)*(z^2), lower = -Inf, upper = Inf)$value -
      m_c^2
    e2_c <- (new_rlz - m_c)^2
    dss_c <- dss(v_c, e2_c)
    logs_c <- log(f_blp(new_rlz))
  } else {
    m_c <- v_c <- e2_c <- dss_c <- logs_c <- NA
  }
  list(w = w, par1 = ab[1], par2 = ab[2], mean = m_c, variance = v_c, e2 = e2_c, 
       dss = dss_c, logs = logs_c)
}
lp_draws <- function(new_draws1, new_draws2, rlz, w = .5, 
                     center = FALSE){
  # Center draws (if selected)
  if (center){
    # Get individual means and combined mean
    aux <- center_draws(new_draws1$m_draws, new_draws2$m_draws, w = w)
    # Update mean draws to remove disagreement across models
    # (but keep disagreement within each model)
    new_draws1$m_draws <- aux[[1]]
    new_draws2$m_draws <- aux[[2]]
  }
  # Now apply linear pool
  mv <- get_mv(draws1 = new_draws1, draws2 = new_draws2, w = w, c = 1)
  e2_c <- (rlz - mv$m)^2
  dss_c <- dss(mv$v, e2_c)
  f_c <- w*mean(dnorm(rlz, mean = new_draws1$m_draws, 
                      sd = sqrt(new_draws1$v_draws))) + 
    (1-w)*mean(dnorm(rlz, mean = new_draws2$m_draws, 
                     sd = sqrt(new_draws2$v_draws)))
  logs_c <- log(f_c)
  list(w = w, par1 = NA, par2 = NA, mean = mv$m, variance = mv$v, e2 = e2_c, 
       dss = dss_c, logs = logs_c)
}

# Function to center forecast draws around common mean (as needed for CLP)
center_draws <- function(mean_draws1, mean_draws2, w){
  # Get individual means and combined mean
  means <- c(mean(mean_draws1), 
             mean(mean_draws2))
  comb_mean <- sum(c(w, 1-w)*means)  
  # Update mean draws to remove disagreement across models
  # (but keep disagreement within each model)
  list(mean_draws1_c = mean_draws1 + (comb_mean - means[1]), 
       mean_draws2_c = mean_draws2 + (comb_mean - means[2]))
}
