library(tidyverse)

# data generation
generate.dat <- function(n, p, seed = 1234) {
  # n = number of timepoints
  # p = grid length on each side of spatial horizon
  set.seed(seed) # set seed for reproducibility

  ts <- seq(0, 1, length.out = n) # time scaled from 0 to 1
  grids <- seq(0, 1, length.out = p) # grid axis points
  sp <- expand.grid(
    s1 = grids,
    s2 = grids
  ) # pxp grid for spatial locations
  corrmat <- exp(-as.matrix(dist(sp))^2) # create the spatial correlation matrix  (gr^2 x gr^2)
  eps <- MASS::mvrnorm(n = n, mu = rep(0, nrow(sp)), Sigma = 0.1^2 * corrmat) # the errors

  # covariate process
  Xt <- arima.sim(n = n, list(order = c(1, 0, 0), ar = 0.5), sd = 0.1)
  Xt <- as.numeric(Xt) + 2 # Xt = 1 + 0.5X(t-1) + et, so (Xt-2) = 0.5(X(t-1) - 2) + et

  # final data
  df <- expand.grid(
    s1 = grids,
    s2 = grids,
    t = ts
  ) %>%
    as_tibble() %>%
    inner_join(
      tibble(t = ts, Xt = Xt), # merge corresponding X values
      by = "t"
    ) %>%
    mutate(
      mu = Xt * (s1 + s2), # the mean values
      sigma = 1, # the standard deviation
      Yt = mu
    )
  for (i in 1:n) {
    index <- (df$t == ts[[i]])
    df$Yt[index] <- df$Yt[index] + df$sigma[index] * eps[i, ]
  }
  return(df)
}

# simple kernel function
kernel.fun <- function(x, xvec, span = 5) {
  # xvec <- vector of Xt to compare against
  knorm <- (xvec - x)^2 # calculate the vector of norms
  nth_smallest <- sort(knorm)[span] # take the span-th smallest norm

  return(ifelse(knorm > nth_smallest, 0, 1 / span)) # consider uniform kernel
}

# takes a data frame and add the basis function, expects two columns s1 and s2
add.basis <- function(df) {
  df %>% dplyr::mutate(
    b1 = 1,
    b2 = 2 * sqrt(3) * (s1 - 1 / 2),
    b3 = 2 * sqrt(3) * (s2 - 1 / 2),
    b4 = 12 * (s1 - 1 / 2) * (s2 - 1 / 2),
    b5 = 6 * sqrt(5) * (s1^2 - s1 + 1 / 6),
    b6 = 6 * sqrt(5) * (s2^2 - s2 + 1 / 6)
  )
}

# estimates the mean, its standard deviation components
est.mean <- function(df.grouped, x, n, p, span = 5) {
  # x = the x value for which I want to predict Yt

  # takes in the grouped data frame
  ts <- seq(0, 1, length.out = n) # time scaled from 0 to 1
  grids <- seq(0, 1, length.out = p) # grid axis points

  xvec <- df.grouped$Xt # df.grouped is assumed to be historical data only
  kvec <- kernel.fun(x, xvec, span = span) # get kernel outputs

  # estimate of mu_k(X_t) = weighted avg(Y_{t_i,k}^*)
  mu_k <- (t(as.matrix(df.grouped[, paste0("Ystar", 1:6)])) %*% kvec) / sum(kvec) # (H x 6)^T * (H x 1) => (6 x 1)
  errmat <- as.matrix(df.grouped[, paste0("Ystar", 1:6)]) - matrix(mu_k, nrow = nrow(df.grouped), ncol = 6, byrow = T) # produces (Y_{t_i,k} - mu_k(x) ), n x 6
  sigma_kl <- (t(errmat) %*% diag(kvec) %*% errmat) / sum(kvec) # (6 x H) * (H x H) * (H x 6) => 6 x 6

  sp <- expand.grid(s1 = grids, s2 = grids)
  preds <- as_tibble(sp) %>%
    add.basis() %>%
    mutate(
      Xt = x,
      mu_hat = mu_k[1] * b1 + mu_k[2] * b2 + mu_k[3] * b3 + mu_k[4] * b4 + mu_k[5] * b5 + mu_k[6] * b6
    ) %>% # mu is composed of individual basis functions
    select(c(s1, s2, Xt, mu_hat))

  return(list(
    "mu_k" = mu_k, # the value of component functions
    "sigma_kl" = sigma_kl, # the covariance of component functions
    "muhat" = preds # the prediction of mu_hats at each location
  ))
}
