library(tidyverse)
library(lubridate)
library(pbapply)

##################
# Main estimation functions
model_fit <- function(
    data,   # the dataset has a response variable, and covariate columns
    response_var, # the string name of the response variable
    timepoints, # the vector of timepoints to be performed fitting for (in POSIXct format)
    covariate_cols, # a vector of covariate column names
    covariate_lag = 6,  # 6 lagged covariate values are used as complete covariate setup
    covariate_phi = 0.9,  # this is the discounting factor used for the covariate effect
    kernel_type = "gaussian", # the type of the kernel function to use
    lambda_bandwidth = 49, # lambda * b_n for the kernel function
    kernel_direction = "past", # kernel only expands to the past timepoints
    location_id_column = "LocationId", # the name of the locationid column
    latitude_column = "lat", # the name of the latitude column
    longitude_column = "long", # the name of the lognitude column
    verbose = FALSE # whether to print out the progress
) {

  # precompute Ytk using modified monte carlo (time index -> Yt_basis)
  Ytk <- data %>%
    rename(all_of(c(
      LocationId = location_id_column,
      Y = response_var,
      lat = latitude_column,
      long = longitude_column
    ))) %>%  # rename so that we can use lat, long values
    select(LocationId, Datetime, Y, lat, long) %>%
    mutate(
      basis_1 = 1,
      basis_2 = 2 * sqrt(3) * (lat - 1 / 2),
      basis_3 = 2 * sqrt(3) * (long - 1 / 2),
      basis_4 = 12 * (lat - 1 / 2) * (long - 1 / 2),
      basis_5 = 6 * sqrt(5) * (lat^2 - lat + 1 / 6),
      basis_6 = 6 * sqrt(5) * (long^2 - long + 1 / 6)
    ) %>%
    group_by(Datetime) %>%
    summarise(
      across(starts_with("basis_"), ~ mean(Y * ., na.rm = T))
    ) %>%
    ungroup() %>%
    arrange(Datetime)  

  # get the list of unique location ids
  location_ids <- unique(data[[location_id_column]])

  # creating the covariate datasets (time index -> covariate_t)
  covardf <- data %>%
    pivot_longer(cols = all_of(covariate_cols), names_to = "covariate_name", values_to = "value") %>%
    mutate(covariate_loc_name = paste0(covariate_name, "_", .data[[location_id_column]])) %>%
    select(Datetime, covariate_loc_name, value) %>%
    pivot_wider(names_from = covariate_loc_name, values_from = "value") %>%
    arrange(Datetime)
  
  
  # convert respective dataframes to a matrices for faster lookup
  Ytk_mat <- Ytk %>%
    select(-Datetime) %>%
    as.matrix()
  rownames(Ytk_mat) <- format(Ytk$Datetime, "%Y-%m-%d %H:%M:%S")
  covardf_mat <- covardf %>%
    select(-Datetime) %>%
    as.matrix()
  rownames(covardf_mat) <- format(covardf$Datetime, "%Y-%m-%d %H:%M:%S")
  
  # the kernel function to use
  kernel_fun <- function(cov1, cov2, kernel_type = "gaussian") {
    cov_diff <- cov1 - cov2  # this is (covariate_lag, p) matrix
    x_norm <- mean(
      rowMeans(cov_diff^2, na.rm = T) *
        (covariate_phi^rev(seq_len(nrow(cov_diff)))), # discounted phi
      na.rm = TRUE
    )
    
    if (kernel_type == "gaussian") {
      return(exp(-x_norm))
    } else if (kernel_type == "uniform") {
      return(1/nrow(cov_diff))  # return constant value, adjusted by kernel normalization
    } else if (kernel_type == "inverse") {
      return(1 / (1 + x_norm))
    } else {
      stop("Invalid kernel type")
    }
  }
  
  # now start estimating as weighted averages
  timedel <- timepoints[2] - timepoints[1] # time delta
  Tlen <- length(timepoints)
  
  # main function to compute estimate at i-th timepoint
  compute_trends <- function(i) {
    # calculate Yt_basis(hat) at timepoints[i]
    tt <- timepoints[i]
    tt_char <- format(tt, "%Y-%m-%d %H:%M:%S")
    tt_covar_char <- format(tt - timedel * c((covariate_lag - 1): 0), "%Y-%m-%d %H:%M:%S")
    
    # based on the kernel direction and lambda_bandwidth, pick up indices that matter
    if (kernel_direction == "both") {
      indlen <- (lambda_bandwidth - 1) / 2
      tt_main <- tt + timedel * c((-indlen):indlen)
      tt_indices <- format(tt_main, "%Y-%m-%d %H:%M:%S")
      tt_subindices <- c(rep(FALSE, indlen / 2), rep(TRUE, indlen + 1), rep(FALSE, indlen / 2))  # logical index
    } else if (kernel_direction == "past") {
      indlen <- lambda_bandwidth
      subindlen <- round(indlen / 2)
      tt_main <- tt - timedel * c(indlen:1)
      tt_indices <- format(tt_main, "%Y-%m-%d %H:%M:%S")
      tt_subindices <- c(rep(FALSE, indlen - subindlen), rep(TRUE, subindlen))  # logical index
    } else {
      stop("Invalid kernel direction")
    }
    
    # get the Ytk response for this timepoint within the bandwidth
    Ytk_sub <- Ytk_mat[tt_indices, , drop = FALSE] # should be (lambda_bandwidth x 6)
    covmat_current <- covardf_mat[tt_covar_char,] # should be (covariate_lag, p)
    
    # apply the kernels
    kvec <- sapply(1:length(tt_main), FUN = function(j) {
      tt_j_char <- format(tt_main[j] - timedel * c((covariate_lag - 1):0), "%Y-%m-%d %H:%M:%S")
      covmat_sub <- covardf_mat[tt_j_char,] # should be (covariate_lag, p)
      kernel_fun(covmat_sub, covmat_current, kernel_type)
    })
    kvec[is.na(kvec)] <- 0 # set 0 to avoid NA
    
    # handle missing data
    misindex <- rowSums(is.na(Ytk_sub)) > 0
    Ytk_sub[misindex, ] <- 0 # set 0 to avoid NA
    kvec[misindex] <- 0 # set 0 to avoid NA
    
    # compute the trend
    mu_k <- as.vector((t(Ytk_sub) %*% kvec) / sum(kvec)) # yields a 6x1 vector
    names(mu_k) <- paste0("muk", 1:6)
    
    # trend estimate for subindices for debiasing purposes
    mu_k2 <- as.vector(t(Ytk_sub[tt_subindices, ]) %*% kvec[tt_subindices]) / sum(kvec[tt_subindices]) # yields a 6 x 1 vector
    names(mu_k2) <- paste0("muk2", 1:6)
    
    # sigma computation
    Ek <- Ytk_sub - matrix(mu_k, nrow = nrow(Ytk_sub), ncol = ncol(Ytk_sub), byrow = TRUE) # lambda_bandwidth x 6
    sig <- (t(Ek) %*% diag(kvec) %*% Ek) / sum(kvec) # computes the Sigma(k,k) matrix (6 x 6)

    return(
      list(
        "mu" = mu_k, # (6 x 1)
        "mu_debias" = 2 * mu_k2 - mu_k, # bias corrected mu (6 x 1)
        "sigma" = sig # (6 x 6)
      )
    )
  }
  
  if (verbose) {
    trends <- pblapply(1:Tlen, compute_trends) # create the estimates
  } else {
    trends <- lapply(1:Tlen, compute_trends) # without progress bar
  }
  
  return(list(
    "trends" = trends,
    "timepoints" = timepoints,
    "covariate_cols" = covariate_cols,
    "covariate_lag" = covariate_lag,
    "covariate_phi" = covariate_phi,
    "lambda_bandwidth" = lambda_bandwidth,
    "kernel_type" = kernel_type,
    "kernel_direction" = kernel_direction,
    "location_id_column" = location_id_column,
    "latitude_column" = latitude_column,
    "longitude_column" = longitude_column
  ))
    
}


model_predict <- function(
    fitted_model, # the fitted model object from model_fit function
    pred_locations, # the locations to predict for as a dataframe
    confidence_level = 0.95, # the confidence level for the intervals
    scale_factor = NULL 
    # the scale factor for the integral
    # this can be obtained variance scaling
) {
  GUMBEL_Z <- -log(log(1 / sqrt(confidence_level)))
  GUMBEL_MN <- length(fitted_model$timepoints)
  GUMBEL_FACTOR <- sqrt(2 * log(GUMBEL_MN)) -
    (0.5 * log(log(GUMBEL_MN)) + log(2 * sqrt(pi))) / sqrt(2 * log(GUMBEL_MN)) +
    GUMBEL_Z / sqrt(2 * log(GUMBEL_MN))
  
  # get the trend objects as dataframe
  location_list <- unique(pred_locations[[fitted_model$location_id_column]])
  predlistdf <- lapply(location_list, function(loc) {
    # get the lat and long
    loc_details <- pred_locations %>%
      filter(!!sym(fitted_model$location_id_column) == loc) %>%
      head(1)
    pred_lat <- loc_details$lat
    pred_long <- loc_details$long
    basis <- c(
      b1 = 1,
      b2 = 2 * sqrt(3) * (pred_lat - 1 / 2),
      b3 = 2 * sqrt(3) * (pred_long - 1 / 2),
      b4 = 12 * (pred_lat - 1 / 2) * (pred_long - 1 / 2),
      b5 = 6 * sqrt(5) * (pred_lat^2 - pred_lat + 1 / 6),
      b6 = 6 * sqrt(5) * (pred_long^2 - pred_long + 1 / 6)
    )
    bnorm <- sum(basis^2)
    basis <- basis / bnorm
    
    # compute the mu, sigma, lower, upper CI for that location
    muY <- sapply(1:length(fitted_model$timepoints), FUN = function(i) {
      sum(fitted_model$trends[[i]]$mu_debias * basis)
    }, simplify = TRUE)
    sigmaY <- sapply(1:length(fitted_model$timepoints), FUN = function(i) {
      sqrt(crossprod(basis, fitted_model$trends[[i]]$sigma %*% basis)[1, 1])
    }, simplify = TRUE)
    if (!is.null(scale_factor) & length(muY) > 2) {
      # this adjustment allows it match the mean and variance of the data
      # necessary for the stationarity assumption
      muY_med <- quantile(muY, 0.5, na.rm = TRUE)
      muY_iqr <- quantile(muY, 0.75, na.rm = TRUE) - quantile(muY, 0.25, na.rm = TRUE)
      muY <- scale_factor[1] + scale_factor[2] * (muY - muY_med) / muY_iqr
      sigmaY <- scale_factor[2] * sigmaY / muY_iqr
    }
    preddf <- tibble(
      Datetime = fitted_model$timepoints,
      muY = muY,
      sigmaY = sigmaY,
    ) %>%
      mutate(
        lowerCI = muY + 1.13 * qnorm((1 - confidence_level) / 2) * sigmaY / sqrt(fitted_model$lambda_bandwidth), # 1.13 ~ xi_2/sqrt(xi_1)
        upperCI = muY - 1.13 * qnorm((1 - confidence_level) / 2) * sigmaY / sqrt(fitted_model$lambda_bandwidth),
        lowerCI2 = muY - 1.13 * GUMBEL_FACTOR * sigmaY / sqrt(fitted_model$lambda_bandwidth),
        upperCI2 = muY + 1.13 * GUMBEL_FACTOR * sigmaY / sqrt(fitted_model$lambda_bandwidth)
      )
    preddf[[fitted_model$location_id_column]] <- loc # name of the location
    return(preddf)
  }) %>% bind_rows()
  
  return(predlistdf)
}


##################
# Additional Utility functions

compute_distance_orders <- function(
    locdf, # a dataframe with locationid, latitude and longitude
    location_id_col = "LocationId",
    latitude_col = "latitude",
    longitude_col = "longitude",
    n_closest = -1) {
  # calculate the distance matrix first
  dmat <- locdf %>%
    select(all_of(c(latitude_col, longitude_col))) %>%
    dist() %>% # euclidean distance based approximation
    as.matrix()
  
  if (n_closest <= 0) {
    n_closest <- nrow(locdf) - 1 # exclude the self
  }
  
  closest_indices <- t(apply(dmat, 1, order)[2:(n_closest + 1), ]) # each row = locations, each column, closest items
  rownames(closest_indices) <- locdf[[location_id_col]]
  
  closest_distances <- t(apply(dmat, 1, sort)[2:(n_closest + 1), ]) # each row = locations
  rownames(closest_distances) <- locdf[[location_id_col]]
  
  return(list(
    "locations" = locdf[[location_id_col]],
    "index" = closest_indices,
    "distances" = closest_distances
  ))
}



# utility function to compare two predictions
compare_prediction_errors <- function(mergeddf) {
  if (nrow(mergeddf) == 0) {
    return(list("accuracy" = NA, "test" = NA))
  }
  test_result <- pretest::dm_cw(mergeddf$e1, mergeddf$e2)
  metrics_mat <- matrix(
    c(
      sqrt(mean(mergeddf$e1^2)), sqrt(mean(mergeddf$e2^2)),
      mean(abs(mergeddf$e1)), mean(abs(mergeddf$e2)),
      mean(abs(mergeddf$e1) / mergeddf$Y), mean(abs(mergeddf$e2) / mergeddf$Y)
    ),
    nrow = 3, ncol = 2, byrow = TRUE,
  )
  colnames(metrics_mat) <- c("model1", "model2")
  rownames(metrics_mat) <- c("RMSE", "MAE", "MAPE")
  
  metrics_test <- matrix(data = c(
    test_result$dm, test_result$pv_dm,
    test_result$dm_nw, test_result$pv_dm_nw,
    test_result$cw, test_result$pv_cw,
    test_result$cw_nw, test_result$pv_cw_nw
  ), nrow = 4, ncol = 2, byrow = TRUE)
  colnames(metrics_test) <- c("statistic", "pvalue")
  rownames(metrics_test) <- c("dm", "dm_nw", "cw", "cw_nw")
  
  return(list(
    "accuracy" = metrics_mat,
    "test" = metrics_test
  ))
}


# create gapped vector
create_gapped_vector <- function(v, take_every, skip_every) {
  start_indices <- seq(1, length(v), by = (take_every + skip_every))
  take_indices <- unlist(lapply(start_indices, FUN = function(i) { i + (0:(take_every - 1)) } ))
  return(v[take_indices[take_indices <= length(v)]])
}
