# set current path to this file (specify your own path here) and clear the global environment
setwd("~/RCodeBikeSharing")
rm(list = ls())

# choose the analyses to be conducted ----
section.4.1 <- FALSE 
section.4.2 <- FALSE 
section.4.3 <- FALSE
section.5 <- FALSE 
appendix.C.1 <- FALSE 
appendix.C.2.distributions <- FALSE 
appendix.C.2.interventions <- FALSE 
appendix.D <- FALSE 
appendix.E <- FALSE 

# save google API key needed for producing the Vienna maps (insert your personal google API key here) ----
API.key <- ""

# load required packages ----
library(dplyr)
library(readr)
library(ggplot2)
library(ggmap)
library(splines)
library(mgcv)
library(Matrix)
library(MASS)
library(matrixStats)
library(R.utils)
library(parallel)
library(gridExtra)

# load functions and plot functions
source("Functions.R")
source("FunctionsPlotting.R")
source("FunctionsSimulation.R")

# load public available data ----
stations <- read_csv("Data/Public/stations.csv")
distances <- read_rds("Data/Public/distances.rds")
feeds <- read_rds("Data/Public/feeds.rds")

# load artifical weather dataset if the original dataset is not available
#weather <- read_csv("Data/Public/weather_artificial.csv") 

# load confidential data ----
# without the confidential you can not run analyses with 
# - interventions = "corrected"
# - distribution = "poisson
source("Confidential.R")

# set global variables and global ggplot theme ----
d.max <- 362 # number of days from the first day
t.max <- d.max*24 # number of time indices
N.stations <- 120 # number of stations in the bikesharing network
we <- FALSE # include (TRUE) or not include (FALSE) days on the weekend
Sigma <- matrix(c(0.6, 0.5, 0.5, 0.6), ncol = 2) # starting value of the covariance matrix of random effects
smooths <- c("temp", "sun", "prec", "seas") # names of the variables which should be modelled non-linearly
lambda <- c(100, 100, 100, 10) # starting values of the penalization parameters
trace <- 3    # trace argument of the stats::optim function
max.it <- 10  # maximum number of iterations of the approximate EM-algorithm
eps <- 1e-3   # epsilon in the approximate EM-algorithm
theme_set(theme_bw() + theme(panel.grid = element_blank(), plot.title = element_text(hjust = 0.5))) # blank ggplot theme

# descriptive analyses (Section 4.1) ----
if (section.4.1) {
  source("Descriptives.R")
}

# station based model fitting for the hour from 5-6 pm (Section 4.2) ----
if (section.4.2){
  h <- 18 
  mode <- "station" 
  interventions <- "corrected"
  
  for (distribution in c("skellam", "tskellam", "poisson")) {  
    name <- paste(h, "_" , mode , "_",  interventions, "_", distribution, sep = "")
    
    # design of the model
    design <- get.design(h = h, mode = mode, distribution = distribution)
    saveRDS(design, file = paste0("Designs/", name, ".rds"))
    
    # get counts for the Poisson model or differences for the (T)-Skellam model
    if (distribution == "poisson"){
      data <- switch(interventions, 
                     "none" = N, 
                     N.wo.P)
      counts <- get.counts(data = data, mode = mode, design = design)
    } else {
      differences <- get.differences(design = design, mode = mode, interventions = interventions)
      D <- differences$D
      design <- differences$design.D
    }
    
    # if we fit the T-skellam model we also need the station feeds
    if (distribution == "tskellam"){
      C <- get.feeds(design = design, mode = mode, D = D)
    }
    
    # fit the model
    fit <- fit.model(design = design, mode = mode, distribution = distribution,
                     name = name, D = D, C = C, counts = counts)
    saveRDS(fit, file = paste0("Fits/", name, ".rds"))
    
    # Table 1
    tab <- tibble(names = design$names.theta[design$ind.beta],
                  estimate = round(fit$theta.hat[design$ind.beta], 3),
                  se = round(sqrt(diag(fit$V.hat[design$ind.beta, design$ind.beta])), 3))
    print(tab)
    
    # estimated covariance matrix of random effects
    print(round(fit$Sigma.hat, 3))
    
    # Figure 4
    plot.smooth.effects(fit = fit, design = design, name = name)
    
    # Figure 5
    plot.random.effects(fit = fit, h = h, mode = mode, distribution = distribution, name = name)
  }
}

# dyadic model fitting for the hour from 5-6 pm (Section 4.3) ----
if (section.4.3){
  h <- 18  
  mode <- "dyadic"  
  interventions <- "corrected" 
  
  for (distribution in c("skellam", "poisson")) {
    name <- paste(h, "_" , mode , "_",  interventions, "_", distribution, sep = "")
    
    # design of the model
    design <- get.design(h = h, mode = mode, distribution = distribution)
    saveRDS(design, file = paste0("Designs/", name, ".rds"))
    
    # get counts for the Poisson model or differences for the (T)-Skellam model
    if (distribution == "poisson"){
      data <- switch(interventions, 
                     "none" = N, 
                     N.wo.P)
      counts <- get.counts(data = data, mode = mode, design = design)
    } else {
      differences <- get.differences(design = design, mode = mode, interventions = interventions)
      D <- differences$D
      design <- differences$design.D
    }
    
    # fit the model
    fit <- fit.model(design = design, mode = mode, distribution = distribution,
                     name = name, D = D, C = C, counts = counts)
    saveRDS(fit, file = paste0("Fits/", name, ".rds"))
    
    # Table 2
    tab <- tibble(names = design$names.theta[design$ind.beta],
                  estimate = round(fit$theta.hat[design$ind.beta], 3),
                  se = round(sqrt(diag(fit$V.hat[design$ind.beta, design$ind.beta])), 3))
    print(tab)
  }
}

# simulation study (Section 5) ----
if (section.5){
  # set simulation specs
  beta <- list(c(-5, 1, -1), 
               c(-5, -1, 1), 
               c(-5, 1, 1), 
               c(-5, -1, -1))
  t.max.simulation <- 500
  N.stations.simulation <- 20
  S <- 250
  Sigma.simulation <- matrix(c(1, 0.9, 0.9, 1), ncol = 2)
  
  # uncomment the following lines to run two simulations for trial (2nd step by step)
  #sim.parallel <- TRUE
  #fit <- simulate.network(1, N.stations.simulation, t.max.simulation, beta[[1]], Sigma.simulation, sim.parallel)
  #print(fit$theta.hat[1:3])
  #print(fit$Sigma.star)
  #sim.parallel <- FALSE
  #network <- simulate.network(1, N.stations.simulation, t.max.simulation, beta[[1]], Sigma.simulation, sim.parallel)
  #design <- network$design
  #flows <- simulate.flows(1, Sigma.simulation, N.stations.simulation, t.max.simulation, design, network$mu, sim.parallel)
  #fit <- fit.model.simulation(1, N.stations.simulation, as.vector(flows$D), design, Sigma.simulation)
  #print(fit$theta.hat[1:3])
  #print(fit$Sigma.star)
  
  # run the simulation study parallely
  no.cores <- detectCores() - 4
  print(no.cores)
  cl <- makeCluster(no.cores)
  start <- Sys.time()
  clusterExport(cl, c("N.stations.simulation", "t.max.simulation", "Sigma.simulation", "beta", "max.it", "eps", "trace",
                      "simulate.network", "simulate.flows", "get.design.simulation", 
                      "fit.model.simulation", "logL.skellam", "score.skellam"))
  clusterEvalQ(cl, library(dplyr))
  clusterEvalQ(cl, library(Matrix))
  clusterEvalQ(cl, library(MASS))
  clusterEvalQ(cl, library(R.utils))
  
  
  results <- vector("list", 4)
  set.seed(1)
  results[[1]] <- parLapply(cl, 1:S, simulate.network, N.stations = N.stations.simulation, t.max = t.max.simulation, 
                            beta = beta[[1]], Sigma = Sigma.simulation, sim.parallel = TRUE)
  saveRDS(results, file = paste0('Fits/Simulation/N',  N.stations.simulation, '_S', S , '.rds'))
  set.seed(1)
  results[[2]] <- parLapply(cl, 1:S, simulate.network, N.stations = N.stations.simulation, t.max = t.max.simulation, 
                            beta = beta[[2]], Sigma = Sigma.simulation, sim.parallel = TRUE)
  saveRDS(results, file = paste0('Fits/Simulation/N',  N.stations.simulation, '_S', S , '.rds'))
  set.seed(1)
  results[[3]] <- parLapply(cl, 1:S, simulate.network, N.stations = N.stations.simulation, t.max = t.max.simulation, 
                            beta = beta[[3]], Sigma = Sigma.simulation, sim.parallel = TRUE)
  saveRDS(results, file = paste0('Fits/Simulation/N',  N.stations.simulation, '_S', S , '.rds'))
  set.seed(1)
  results[[4]] <- parLapply(cl, 1:S, simulate.network, N.stations = N.stations.simulation, t.max = t.max.simulation, 
                            beta = beta[[4]], Sigma = Sigma.simulation, sim.parallel = TRUE)
  saveRDS(results, file = paste0('Fits/Simulation/N',  N.stations.simulation, '_S', S , '.rds'))
  stopCluster(cl)
  
  # load results if task has been performed on an external machine
  #results <- read_rds("Fits/Simulation/N20_S250.rds")
  
  # Figure 6
  plot.simulation(results, S, N.stations.simulation)
}

# further analyses for the station based model (Appendix C) ----
if (appendix.C.1){

  # load results of the Skellam model and estimate the flows
  name.skellam <- "18_station_corrected_skellam"
  fit.skellam <- read_rds(paste0("Fits/", name.skellam, ".rds"))
  design.skellam <- read_rds(paste0("Designs/", name.skellam, ".rds"))
  estimations.skellam <- estimate.flows(fit = fit.skellam, design = design.skellam, mode = "station", distribution = "skellam")
  
  # load results of the T-Skellam model and estimate the flows
  name.tskellam <- "18_station_corrected_tskellam"
  fit.tskellam <- read_rds(paste0("Fits/", name.tskellam, ".rds"))
  design.tskellam <- read_rds(paste0("Designs/", name.tskellam, ".rds"))
  estimations.tskellam <- estimate.flows(fit = fit.tskellam, design = design.tskellam, mode = "station", distribution = "tskellam")
  
  # load results of the Poisson model and estimate the flows
  name.poisson <- "18_station_corrected_poisson"
  fit.poisson <- read_rds(paste0("Fits/", name.poisson, ".rds"))
  design.poisson <- read_rds(paste0("Designs/", name.poisson, ".rds"))
  estimations.poisson <- estimate.flows(fit = fit.poisson, design = design.poisson, mode = "station", distribution = "poisson")
  
  # Figure 7
  lim <- max(colSums(N.wo.P$outgoing[design.skellam$H, ]), 
             colSums(estimations.skellam$N.out.hat), 
             colSums(estimations.tskellam$N.out.hat))
  plot.cum.by.station(observed = N.wo.P, estimated = estimations.skellam, design = design.skellam, lim = lim, name = name.skellam)
  plot.cum.by.station(observed = N.wo.P, estimated = estimations.tskellam, design = design.tskellam, lim = lim, name = name.tskellam)

  # Figure 8
  lim <- max(rowSums(N.wo.P$outgoing[design.skellam$H, ]), 
             rowSums(estimations.skellam$N.out.hat), 
             rowSums(estimations.tskellam$N.out.hat))
  plot.cum.by.day(observed = N.wo.P, estimated = estimations.skellam, design = design.skellam, lim = lim, name = name.skellam)
  plot.cum.by.day(observed = N.wo.P, estimated = estimations.tskellam, design = design.tskellam, lim = lim, name = name.tskellam)
  
  # Figure 9
  plot.residuals(observed = N.wo.P, estimated = estimations.skellam, design = design.skellam, name = name.skellam)
  plot.residuals(observed = N.wo.P, estimated = estimations.poisson, design = design.poisson, name = name.poisson)
}

# station based model fitting for different distributions (Appendix C) ----
if (appendix.C.2.distributions){
  # specify the models
  models <- tibble(h = rep(9:20, 3), distribution = rep(c("skellam", "tskellam", "poisson"), each = 12),
                   interventions = "corrected", mode = "station")
  models <- setNames(split(models, seq(nrow(models))), rownames(models))
  
  # fit the models
  lapply(models, fit.models)
  
  # Figure 10
  plot.distributions()
}

# station based model fitting for different treatments of provider interventions (Appendix C) ----
if (appendix.C.2.interventions) {
  # specify the models
  models <- tibble(h = rep(9:20, 3), distribution = "skellam",
                   interventions = rep(c("corrected", "threshold", "none"), each = 12), mode = "station")
  models <- setNames(split(models, seq(nrow(models))), rownames(models))

  # fit the models
  lapply(models, fit.models)
  
  # Figure 11
  plot.interventions()
}

# simulation study (Appendix D) ----
if (appendix.D){
  # set simulation specs
  beta <- list(c(0.5, 0.25, 0.25), 
               c(0.5, 0.25, 0.5), 
               c(0.5, 0.5, 0.5))
  t.max.simulation <- 1000
  range.C.max <- 30
  S <- 250
  
  # initialize result table
  results <- vector("list", 3)
  tab.median <- matrix(0, 6, 5)
  colnames(tab.median) <- c("beta.t", "beta.out", "beta.in", "delta.out", "delta.in")
  rownames(tab.median) <- c("beta.1.skellam", "beta.1.tskellam", "beta.2.skellam", "beta.2.tskellam", 
                            "beta.3.skellam", "beta.3.tskellam")
  tab.lower <- tab.upper <- tab.median
  
  # run the simulation study
  for (i in 1:3) {
    set.seed(1)
    results <- lapply(1:S, simulation.truncated.skellam, beta = beta[[i]], t.max = t.max.simulation, range.C.max = range.C.max)
    saveRDS(results, paste0("Fits/Simulation/tskellam_beta_", i, "_S", S, ".rds"))
    
    beta.hat <- matrix(0, S, 6)
    delta <- matrix(0, S, 6)
    for (s in 1:S) {
      beta.hat[s, 1:3] <- results[[s]]$theta.skellam
      beta.hat[s, 4:6] <- results[[s]]$theta.tskellam
      delta[s, 1] <- results[[s]]$delta.out.skellam
      delta[s, 2] <- results[[s]]$delta.in.skellam
      delta[s, 3] <- results[[s]]$delta.out.tskellam
      delta[s, 4] <- results[[s]]$delta.in.tskellam
    }
    
    tab.median[2*i - 1, 1:3] <- round(colQuantiles(beta.hat, cols = 1:3, probs = 0.5), 2)
    tab.median[2*i, 1:3] <- round(colQuantiles(beta.hat, cols = 4:6, probs = 0.5), 2)
    tab.median[2*i - 1, 4:5] <- round(colQuantiles(delta, cols = 1:2, probs = 0.5), 2)
    tab.median[2*i, 4:5] <- round(colQuantiles(delta, cols = 3:4, probs = 0.5), 2)
    
    tab.lower[2*i - 1, 1:3] <- round(colQuantiles(beta.hat, cols = 1:3, probs = 0.05), 2)
    tab.lower[2*i, 1:3] <- round(colQuantiles(beta.hat, cols = 4:6, probs = 0.05), 2)
    tab.lower[2*i - 1, 4:5] <- round(colQuantiles(delta, cols = 1:2, probs = 0.05), 2)
    tab.lower[2*i, 4:5] <- round(colQuantiles(delta, cols = 3:4, probs = 0.05), 2)
    
    tab.upper[2*i - 1, 1:3] <- round(colQuantiles(beta.hat, cols = 1:3, probs = 0.95), 2)
    tab.upper[2*i, 1:3] <- round(colQuantiles(beta.hat, cols = 4:6, probs = 0.95), 2)
    tab.upper[2*i - 1, 4:5] <- round(colQuantiles(delta, cols = 1:2, probs = 0.95), 2)
    tab.upper[2*i, 4:5] <- round(colQuantiles(delta, cols = 3:4, probs = 0.95), 2)
  }
  
  # Table 3
  print(tab.median)
  print(tab.lower)
  print(tab.upper)
}

# further analyses for the dyadic model (Appendix E) ----
if (appendix.E){
  # load results of the dyadic Skellam model and estimate the flows
  name.skellam <- "18_dyadic_corrected_skellam"
  fit.skellam <- read_rds(paste0("Fits/", name.skellam, "_1000_4.rds"))
  design.skellam <- read_rds(paste0("Designs/", name.skellam, ".rds"))
  estimations.skellam <- estimate.flows(fit = fit.skellam, design = design.skellam, mode = "dyadic", distribution = "skellam")
  
  # load results of the dyadic Poisson model and estimate the flows
  name.poisson <- "18_dyadic_corrected_poisson"
  fit.poisson <- read_rds(paste0("Fits/", name.poisson, ".rds"))
  design.poisson <- read_rds(paste0("Designs/", name.poisson, ".rds"))
  estimations.poisson <- estimate.flows(fit = fit.poisson, design = design.poisson, mode = "dyadic", distribution = "poisson")
  
  # Figure 12
  plot.cum.dist(observed = N.wo.P, estimations.skellam = estimations.skellam, 
                estimations.poisson = estimations.poisson, design = design.skellam)
  plot.pois.vs.skellam(estimations.poisson = estimations.poisson, 
                       estimations.skellam = estimations.skellam, design = design.skellam)
  
  # Figure 13
  plot.check.poisson(observed = N.wo.P, estimated = estimations.skellam, design = design.skellam, name = name.skellam)
}
