library(tidyverse)
source("./utils.R")


# Sample plot
df <- generate.dat(n = 50, p = 2, seed = 1)
df %>%
  mutate(Location = paste0("(", s1, ",", s2, ")")) %>%
  ggplot() +
  theme_bw() +
  geom_line(aes(x = t, y = Yt, color = Location), size = 1) +
  xlab("Time") +
  ylab("Yt")



# =========================
# Testing
# ========================

df <- generate.dat(n = 20, p = 4, seed = 1)
head(df, 5)

df.grouped <- df %>%
  add.basis() %>%
  group_by(t) %>%
  summarise(
    Xt = mean(Xt), # this is same as Xt
    Ystar1 = mean(b1 * Yt), # grouped average within each time
    Ystar2 = mean(b2 * Yt),
    Ystar3 = mean(b3 * Yt),
    Ystar4 = mean(b4 * Yt),
    Ystar5 = mean(b5 * Yt),
    Ystar6 = mean(b6 * Yt)
  ) %>%
  ungroup() %>%
  arrange(t)
head(df.grouped)

out <- est.mean(df.grouped, x = 0, 20, p = 5)
out

# ===============
# Simulation (Prediction)
# ==============
prediction.loop <- function(n, p, seed, pred_only = 10, type = 1) {
  # pred_only = number of datapoints to predict for at last
  # prediction type
  #   1 = When predicting Yt, we know Xt.
  #   2 = When predicting Yt, only X(t-1) is known.
  #   3 = When predicting Yt, only Y(t-1) is known.
  #       In this case, an estimate of X(t-1) = avg(Y(t-1)), as they are linearly related except the error eps

  df <- generate.dat(n, p, seed) # generate data
  df.grouped <- df %>%
    add.basis() %>%
    group_by(t) %>%
    summarise(
      Xt = mean(Xt), # the Xt values will remain same
      Ystar1 = mean(b1 * Yt), # grouped average within each time
      Ystar2 = mean(b2 * Yt),
      Ystar3 = mean(b3 * Yt),
      Ystar4 = mean(b4 * Yt),
      Ystar5 = mean(b5 * Yt),
      Ystar6 = mean(b6 * Yt)
    ) %>%
    ungroup() %>%
    arrange(t) # find the grouped data

  timepoints <- sort(unique(df$t)) # timepoints
  pred_timepoints <- rev(rev(timepoints)[1:pred_only]) # last 10 timepoints
  pred_prev_timepoints <- rev(rev(timepoints)[2:(pred_only + 1)]) # previous time to last 10 timepoints

  predlist <- as.list(numeric(pred_only)) # contains the prediction for the last 10 timepoints
  for (i in 1:pred_only) {
    cur_t <- pred_timepoints[i] # predicting for timepoint t
    prev_t <- pred_prev_timepoints[i]

    if (type == 1) {
      # Xt is known
      x <- df.grouped$Xt[df.grouped$t == cur_t][1] # current x
      df.grouped_history <- df.grouped %>% filter(t < cur_t) # previous Xt values
    } else if (type == 2) {
      # only X(t-1) is known, so current x where to predict is unknown
      x <- df.grouped$Xt[df.grouped$t == prev_t][1]
      df.grouped_history <- df.grouped %>%
        filter(t < cur_t) %>%
        mutate(Xt = lag(Xt, 1)) %>% # this means, for each t, only X(t-1) is available
        drop_na()
    } else if (type == 3) {
      # only Y(t-1) is known
      if (i > 1) {
        # based on the last prediction, find the grouped data
        df.grouped_total <- bind_rows(
          df.grouped,
          bind_rows(predlist[1:(i - 1)]) %>%
            add.basis() %>%
            group_by(t) %>%
            summarise(
              Xt = mean(Xt), # the Xt values will remain same
              Ystar1 = mean(b1 * mu_hat), # grouped average within each time
              Ystar2 = mean(b2 * mu_hat),
              Ystar3 = mean(b3 * mu_hat),
              Ystar4 = mean(b4 * mu_hat),
              Ystar5 = mean(b5 * mu_hat),
              Ystar6 = mean(b6 * mu_hat)
            ) %>%
            ungroup() %>%
            arrange(t) # find the grouped data
        )
      } else {
        df.grouped_total <- df.grouped
      }

      # now that I have grouped total (combined with previous predictions)
      # I can take x = average of historical Y(t-1)

      x <- df.grouped_total$Ystar1[df.grouped_total$t == prev_t][1] # Ystar1 = Yt
      df.grouped_history <- df.grouped_total %>%
        filter(t < cur_t) %>%
        mutate(Xt = lag(Ystar1, 1)) %>% # for historical data, Xt was simply Y(t-1)
        drop_na()
    }

    res <- est.mean(df.grouped_history, x, n, p)
    predlist[[i]] <- res$muhat # store the muhat
    predlist[[i]]$t <- cur_t # add a timepoint
  }

  preddf <- bind_rows(predlist) # contains all predictions of mu_hat

  mergedf <- df %>%
    inner_join(preddf, by = c("s1", "s2", "t")) %>%
    mutate(error = mu - mu_hat)

  metricsdf <- mergedf %>%
    group_by(t) %>%
    summarise(
      Bias = mean(error),
      MAE = mean(abs(error)),
      RMSE = sqrt(mean(error^2)),
      MAPE = mean(abs(error / mu_hat)) * 100
    )

  return(metricsdf) # return the bias, RMSE, and MAPE for each timepoint
}

# Main simulation function
do.sim <- function(B, n, p, simtype = 3, parallel = FALSE, init.seed = 1234) {
  set.seed(init.seed)
  seeds <- sample(1:10000, size = B)

  if (parallel) {
    # make the cluster
    ncore <- parallel::detectCores() - 1
    cl <- parallel::makeCluster(ncore)

    # load tidyverse on all workers
    parallel::clusterEvalQ(cl, library(tidyverse))

    # export variables to cluster
    parallel::clusterExport(cl, varlist = c(
      "add.basis", "est.mean", "generate.dat",
      "kernel.fun", "prediction.loop"
    ))

    results <- pblapply(X = 1:B, FUN = function(b) {
      return(prediction.loop(
        n,
        p,
        seeds[b],
        pred_only = 10,
        type = simtype
      ))
    }, cl = cl)
    parallel::stopCluster(cl) # this is important!
  } else {
    results <- as.list(numeric(B))
    pb <- txtProgressBar(min = 0, max = B, style = 3)
    for (b in 1:B) {
      results[[b]] <- prediction.loop(
        n,
        p,
        seeds[b],
        pred_only = 10,
        type = simtype
      )
      setTxtProgressBar(pb, value = b)
    }
    close(pb)
  }

  # combine the result over multiple dataframes
  combdf <- bind_rows(results) %>%
    group_by(t) %>%
    summarise(
      Bias = mean(Bias),
      MAE = mean(MAE),
      RMSE = sqrt(mean(RMSE^2)),
      MAPE = mean(MAPE)
    ) %>%
    ungroup() %>%
    mutate(
      n = n,
      p = p
    )

  combdf %>% write_csv("./simdata-grid-v2.csv", append = TRUE) # comment / uncomment as needed

  combdf
}


# simulation for a single case
do.sim(
  B = 100, n = 100,
  p = 50, # modify this to modify # of locations
  simtype = 3, # modify this as needed
  parallel = T # set to true for parallization
)


# plot over p
df2 <- read_csv("./simdata-grid-v2.csv")
df2 %>%
  mutate(t = as.factor(round(t, 3))) %>%
  ggplot(aes(x = p, y = MAPE, color = t)) +
  theme_bw() +
  geom_line(size = 1) +
  geom_point(size = 2)
