---
title: "Block Shuffle Model Training"
output: html_document
date: "2024-07-15"
---

```{r setup, include=FALSE}
suppressPackageStartupMessages({
library(tictoc)
library(this.path)
library(Rcpp)
library(dplyr)
library(tidyr)
library(data.table)
library(qs2)
library(stringplus)
library(parallel)
library(xgboost)
library(caret)
library(future.apply)
})

# dynamic blocksize set to 1048576
# qs2:::internal_set_blocksize(1048576)

DATA_PATH <- "~/datasets/processed"

# holdout testing datasets
datasets <- DATA_PATH & "/" & c("1000genomes_noncoding_vcf.csv.gz", "B_cell_petshopmouse3.tsv.gz",
                                "ip_location_2023.csv.gz", "Netflix_Ratings.csv.gz")

read_dataset <- function(d) {
  if(d %like% "json.gz") {
    DATA <- RcppSimdJson::fload(d)
  } else if(d %like% ".csv.gz") {
    DATA <- fread(d, sep = ",", data.table=F)
  } else if(d %like% ".tsv.gz") {
    DATA <- fread(d, sep = "\t", data.table=F)
  } else {
    DATA <- readRDS(d)
  }
}

rmse <- function(x, y) sqrt(mean((x-y)^2))

bst <- xgboost::xgb.load("data/blockshuffle_xgboost_model.json")
xgb.parameters(bst) <- list(nthread = 1) # use 1 thread for comparison to standalone implementation
sourceCpp("xgboost_model_cpp_standalone.cpp")
```

```{r}
outfile <- "data/blockshuffle_testing_data.csv.gz"
if(!file.exists(outfile)) {
  libs <- system("pkg-config --libs libzstd", intern = TRUE)
  cflags <- system("pkg-config --cflags libzstd", intern = TRUE)
  Sys.setenv(PKG_CPPFLAGS = "-mavx2 %s %s" | c(cflags, libs))
  Sys.setenv(PKG_LIBS = "-mavx2 %s %s" | c(cflags, libs))
  sourceCpp("blockshuffle_heuristic.cpp", verbose=TRUE, rebuild = TRUE)
  min_shuffleblock_size <- 262144
  
  blocks_df <- lapply(datasets, function(d) {
    tmp <- tempfile()
    data <- read_dataset(d)
    dname<- basename(d) %>% gsub("\\..+", "", .)
    qs2::qd_save(data, file = tmp)
    x <- qs2::qx_dump(tmp)
    r1 <- tibble(dataset = dname, blocks = x$blocks, algo = "qdata")
    qs2::qs_save(data, file = tmp)
    x <- qs2::qx_dump(tmp)
    r2 <- tibble(dataset = dname, blocks = x$blocks, algo = "qs2")
    rbind(r1, r2)
  }) %>% rbindlist
  blocks_df$blocksize <- sapply(blocks_df$blocks, length)
  blocks_df <- filter(blocks_df, blocksize >= min_shuffleblock_size)
  
  gc(full=TRUE)
  compress_levels <- 22:1
  results <- mclapply(compress_levels, function(cl) {
    print(cl)
    output <- shuffle_heuristic(blocks_df$blocks)
    output$no_shuffle_zblocksize <- og_compress(blocks_df$blocks, cl)$size
    output$shuffle_zblocksize <- shuffle_compress(blocks_df$blocks, 8, cl)$size
    output <- output %>% mutate(compress_level = cl)
  }, mc.cores=8, mc.preschedule=FALSE) %>% rbindlist
  results2 <- blocks_df %>% dplyr::select(dataset, algo, blocksize) %>%
    {lapply(1:length(compress_levels), function(i) .)} %>% rbindlist
  results <- cbind(results2, results)
  
  # add block index per dataset
  results <- results %>%
    group_by(dataset, compress_level, algo) %>%
    mutate(index = 1:n()) %>%
    as.data.frame
  fwrite(results, outfile, sep = ",")
} else {
  results <- fread(outfile, data.table=FALSE)
}
```

```{r, compare-cpp-implementation}
# compare C++ implementation
test_data <- results %>% mutate(improvement = log(no_shuffle_zblocksize/shuffle_zblocksize))

timing_data <- test_data %>% dplyr::select(h1,h2,h3,h4,h5,h6,h7,h8,compress_level) %>% {lapply(1:5, function(i) .)} %>% rbindlist
tic(msg = "R package prediction time")
dtest <- xgb.DMatrix(data = timing_data %>% data.matrix)
r_pred <- predict(bst, dtest)
toc()

tic(msg = "Cpp package prediction time")
cpp_pred <- predict_xgboost_cpp(timing_data)
toc()

dtest <- xgb.DMatrix(data = test_data %>% dplyr::select(h1,h2,h3,h4,h5,h6,h7,h8,compress_level) %>% data.matrix)
test_data <- test_data %>%
  mutate(r_prediction = predict(bst, dtest)) %>%
  mutate(cpp_prediction = predict_xgboost_cpp(test_data %>% dplyr::select(h1,h2,h3,h4,h5,h6,h7,h8,compress_level)))

# compare predictions
abs(test_data$r_prediction - test_data$cpp_prediction) %>% summary
```

```{r qdata-prediction-accuracy, fig.width=11.5, fig.height=10.5}
# compare prediction to actual

MAX_COMPRESS <- 16 # don't plot blocks with >16x compression, noisy
test_data2 <- test_data %>%
  filter(algo == "qdata") %>%
  filter(32768*4/(h2 + h4 + h6 + h8) < MAX_COMPRESS)
pal <- palette.colors(palette = "Okabe-Ito")
ggplot(test_data2, aes(x = cpp_prediction, y = improvement, color = dataset)) + 
  geom_abline(aes(slope=1, intercept = 0), lty = 2) + 
  geom_vline(aes(xintercept=0), lty=2, color = "orange") +
  geom_point(shape=21, alpha=0.75) + 
  facet_wrap(~compress_level, ncol=4) +
  scale_color_manual(values = pal) +
  theme_bw(base_size = 12)

```
```{r qdata-efficiency-comparison}
# plot cumulative savings
test_data3 <- test_data %>%
  filter(algo == "qdata") %>%
  mutate(predicted_compression = 32768*4/(h2 + h4 + h6 + h8)) %>%
  mutate(do_blockshuffle = predicted_compression < MAX_COMPRESS & cpp_prediction > 0) %>%
  group_by(dataset, compress_level) %>%
  mutate(optimal = cumsum(pmin(no_shuffle_zblocksize, shuffle_zblocksize))) %>%
  mutate(shuffle_heuristic = cumsum(ifelse(do_blockshuffle, shuffle_zblocksize, no_shuffle_zblocksize))) %>%
  mutate(no_shuffle = cumsum(no_shuffle_zblocksize)) %>%
  mutate(heuristic_was_optimal = do_blockshuffle == (shuffle_zblocksize < no_shuffle_zblocksize) ) %>%
  ungroup

test_data3 <- test_data3 %>%
  filter(compress_level %in% c(3,9)) %>%
  dplyr::select(dataset, compress_level, index, optimal, shuffle_heuristic, no_shuffle) %>%
  pivot_longer(c(-index, -dataset, -optimal, -compress_level), names_to = "shuffle_selection", values_to = "cumulative_bytes") %>%
  mutate(inefficiency = (cumulative_bytes - optimal)/max(optimal) )
  

ggplot(test_data3, aes(x = index, y = inefficiency, color = shuffle_selection, lty = factor(compress_level))) + 
  geom_line() + 
  scale_y_continuous(labels = scales::percent) + 
  facet_wrap(~dataset, scales = "free") +
  theme_bw(base_size=12) + 
  labs(x = "Block Index", y = "Inefficiency", lty = "Compress Level", color = "Shuffle Selection")
  
```


```{r qs2-prediction-accuracy, fig.width=11.5, fig.height=10.5}
# compare prediction to actual

MAX_COMPRESS <- 16 # don't plot blocks with >16x compression, noisy
test_data2 <- test_data %>%
  filter(algo == "qs2") %>%
  filter(32768*4/(h2 + h4 + h6 + h8) < MAX_COMPRESS)
pal <- palette.colors(palette = "Okabe-Ito")
ggplot(test_data2, aes(x = cpp_prediction, y = improvement, color = dataset)) + 
  geom_abline(aes(slope=1, intercept = 0), lty = 2) + 
  geom_vline(aes(xintercept=0), lty=2, color = "orange") +
  geom_point(shape=21, alpha=0.75) + 
  facet_wrap(~compress_level, ncol=4) +
  scale_color_manual(values = pal) +
  theme_bw(base_size = 12)

```
```{r qs2-efficiency-comparison}
# plot cumulative savings
test_data3 <- test_data %>%
  filter(algo == "qs2") %>%
  mutate(predicted_compression = 32768*4/(h2 + h4 + h6 + h8)) %>%
  mutate(do_blockshuffle = predicted_compression < MAX_COMPRESS & cpp_prediction > 0) %>%
  group_by(dataset, compress_level) %>%
  mutate(optimal = cumsum(pmin(no_shuffle_zblocksize, shuffle_zblocksize))) %>%
  mutate(shuffle_heuristic = cumsum(ifelse(do_blockshuffle, shuffle_zblocksize, no_shuffle_zblocksize))) %>%
  mutate(no_shuffle = cumsum(no_shuffle_zblocksize)) %>%
  mutate(heuristic_was_optimal = do_blockshuffle == (shuffle_zblocksize < no_shuffle_zblocksize) ) %>%
  ungroup

test_data3 <- test_data3 %>%
  filter(compress_level %in% c(3,9)) %>%
  dplyr::select(dataset, compress_level, index, optimal, shuffle_heuristic, no_shuffle) %>%
  pivot_longer(c(-index, -dataset, -optimal, -compress_level), names_to = "shuffle_selection", values_to = "cumulative_bytes") %>%
  mutate(inefficiency = (cumulative_bytes - optimal)/max(optimal) )
  

ggplot(test_data3, aes(x = index, y = inefficiency, color = shuffle_selection, lty = factor(compress_level))) + 
  geom_line() + 
  scale_y_continuous(labels = scales::percent) + 
  facet_wrap(~dataset, scales = "free") +
  theme_bw(base_size=12) + 
  labs(x = "Block Index", y = "Inefficiency", lty = "Compress Level", color = "Shuffle Selection")
  
```
