# input arguments of functions in this file (in order of appearance) ----

# h:             hour of the day (integer from 9 to 20)
# mode:          station based mode ("station") or dyadic model ("dyadic")
# distribution:  Skellam distribution ("skellam"), T-Skellam distribution ("tskellam")
#                or Poisson distribution ("poisson")
# interventions: differences corrected for provider interventions ("corrected"),
#                differences corrected with the threshold rule ("threshold") 
#                or no correction of differences for provider interventions ("none")
# alpha:         value of alpha for the dyadic covariate
# m:             number of basis functions for each smooth function
# l:             order of the B-splines
# design:        a list which contains the design matrices, penalty matrices, etc.
# S:             number for the threshold rule in Appendix C
# data:          self explaining
# D:             vector of differences of stations feeds
# name           name of the model which is determined by h, mode, interventions, distribution, 
# C              a data frame which contains information about station capacities and station feeds
#                at the beginning of the time interval
# counts         vector of trip counts
# models         a data frame of models to be fitted with row names according to "name"
# theta          current estimate of the parameter vector
# N.stations     count of stations in the bike-sharing network (including the latent station if mode = "dyadic")
# Sigma          current estimate of the covariace matrix of random effects
# lambda         current estimates of the penalization parameters
# gamma          coefficients of the basis functions
# cov            estimated covariance matrix of the estimated gamma
# z              vector of points where basis functions are evaluated
# B              design matrix of the basis functions
# q              confidence band level of the smooth functions, default to q = 0.05
# S              number of repetitions to compute confidence bands by simulation


# pre fit ----

get.design <- function(h, mode, distribution, alpha = 1.7, m = 10L, l = 2L){
  
  # one additional station w in the dyadic skellam model
  N.stations <- N.stations + 1*(mode == "dyadic" & distribution != "poisson")
  
  # day indices of public holidays in Vienna in 2014
  public.holidays <- c(1, 6, 108, 111, 121, 149, 160, 170, 227, 342, 358, 359, 360)
  #initialize table (time point refers to end of each interval)
  weekday <- tibble(t = 1:t.max, Mon = 0, Tue = 0, Wed = 0, Thu = 0, Fri = 0, Sat = 0, Sun = 0, Ph = 0)
  # fill the table (1st weekday of the year is Wednesday)
  for (d in 1:d.max) {
    weekday[(24*(d-1)+1):(24*d), (d+1)%%7+2] = 1
    if (is.element(d, public.holidays)){
      weekday[(24*(d-1)+1):(24*d), 9] = 1
    }
  }
  
  # determine set of time points
  if (we){
    H <- seq(h, d.max*24, 24)
    weekday <- slice(weekday, H) %>% dplyr::select(weekday, -c(t, Mon))
  } else {
    H <- intersect(seq(h, d.max*24, 24), weekday$t[which(weekday$Sat == 0 & weekday$Sun == 0)])
    weekday <- slice(weekday, H) %>% dplyr::select(-c(t, Mon, Sat, Sun))
  }
  
  # weather covariates
  temp <- (weather$Temperature[which(is.element(weather$t, H))-1] + 
             weather$Temperature[which(is.element(weather$t, H))])/2
  sun.dur <-  (weather$SunshineDuration[which(is.element(weather$t, H))-1] + 
                weather$SunshineDuration[which(is.element(weather$t, H))])/120
  prec.dur <- (weather$PrecipDuration[which(is.element(weather$t, H))-1] + 
                weather$PrecipDuration[which(is.element(weather$t, H))])/120
  prec.amount <- weather$PrecipAmount[which(is.element(weather$t, H))-1] + 
    weather$PrecipAmount[which(is.element(weather$t, H))] 
  
  # semiparametric modelling
  M <- length(smooths) 
  B <- z <- k <- K <- vector("list", M)
  d.s <- rep(0, M)
  ind.gamma <- vector("list", M)
  gamma.names <- NULL
  ind <- 2
  for (s in 1:M) {
    min.z <- switch(smooths[s], 
                   "temp" = min(temp),
                   0)
    max.z <- switch(smooths[s],
                   "temp" = max(temp),
                   1)
    z[[s]] <- switch(smooths[s],
                    "temp" = temp,
                    "seas" = seq(min.z, max.z, length.out =  length(H)),
                    "sun" = sun.dur,
                    "prec" = prec.dur)
    delta <- (max.z - min.z)/(m-1)
    k[[s]] <- seq(min.z-l*delta, max.z+l*delta, delta)
    B.uncentered <- switch(smooths[s],
                          "temp" = splineDesign(knots = k[[s]], x = z[[s]], ord = l+1, outer.ok = TRUE),
                          "seas" = cSplineDes(x = z[[s]], knots = seq(0, 1, length.out = m+2), ord = l+1),
                          "sun" = splineDesign(knots = k[[s]], x = z[[s]], ord = l+1, outer.ok = TRUE),
                          "prec" = splineDesign(knots = k[[s]], x = z[[s]], ord = l+1, outer.ok = TRUE))
    B[[s]] <- sweep(B.uncentered, 2, colMeans(B.uncentered))[, -1]
    d.s[s] <- ncol(B[[s]])
    D <- diff(diag(d.s[s]+1), differences = 2)[, -1]
    K[[s]] <- t(D)%*%D
    ind.gamma[[s]] <- ind:(ind + d.s[s] - 1)
    ind <- ind + d.s[s]
    gamma.names <- c(gamma.names, paste0(smooths[s], ".", 1:d.s[s]))
  }
  
  # get nobikes and noboxes covariate
  nobikes <- noboxes <- matrix(0, N.stations, length(H))
  for (i in 1:N.stations) {
    d <- 0
    for (t in H) {
      d <- d + 1
      nobikes[i, d] <- 1*(feeds[[i]]$NoBikes[which(is.element(feeds[[i]]$t, t))] == 1 & 
                            feeds[[i]]$NoBikes[which(is.element(feeds[[i]]$t, t-1))] == 1)
      noboxes[i, d] <- 1*(feeds[[i]]$NoBoxes[which(is.element(feeds[[i]]$t, t))] == 1 & 
                            feeds[[i]]$NoBoxes[which(is.element(feeds[[i]]$t, t-1))] == 1)
    }
  }

  # Design matrix for the dyadic model
  if (mode == "dyadic"){
    
    # do not include the latent station in the dyadic Poisson model
    if (distribution == "poisson") {
      feeds <- feeds[1:N.stations]
      stations <- stations[1:N.stations, ]
    }
    
    # get indices for time and stations
    ind.d <- rep(rep(1:length(H), each = N.stations), N.stations)
    ind.station.1 <- rep(1:N.stations, each = N.stations*length(H))
    ind.station.2 <- rep(1:N.stations, N.stations*length(H))
    
    # part of the design matrix which is only time related
    Z.t <- cbind(B[[1]][ind.d, ], B[[2]][ind.d, ], B[[3]][ind.d, ], B[[4]][ind.d, ], 1*(prec.amount[ind.d] >= 1),
                 weekday %>% slice(ind.d) %>% as.matrix())
    
    # design matrix for outgoing bikes
    Z.out <- cbind(1,
                   Matrix(Z.t, sparse = TRUE),
                   Matrix(nobikes[cbind(ind.station.1, ind.d)], sparse = TRUE),
                   Matrix(stations$DistanceHub[ind.station.1]/100, sparse = TRUE),
                   Matrix(noboxes[cbind(ind.station.2, ind.d)], sparse = TRUE),
                   Matrix(stations$DistanceHub[ind.station.2]/100, sparse = TRUE), 
                   Matrix(distances[cbind(ind.station.1, ind.station.2)]^alpha*
                            exp(-distances[cbind(ind.station.1, ind.station.2)]), sparse = TRUE),
                   Matrix(diag(N.stations)[ind.station.1, ], sparse = TRUE),
                   Matrix(diag(N.stations)[ind.station.2, ], sparse = TRUE)
    )
    # design matrix for incoming bikes
    Z.in <- cbind(1,
                  Matrix(Z.t, sparse = TRUE),
                  Matrix(nobikes[cbind(ind.station.2, ind.d)], sparse = TRUE),
                  Matrix(stations$DistanceHub[ind.station.2]/100, sparse = TRUE),
                  Matrix(noboxes[cbind(ind.station.1, ind.d)], sparse = TRUE),
                  Matrix(stations$DistanceHub[ind.station.1]/100, sparse = TRUE), 
                  Matrix(distances[cbind(ind.station.2, ind.station.1)]^alpha*
                           exp(-distances[cbind(ind.station.2, ind.station.1)]), sparse = TRUE),
                  Matrix(diag(N.stations)[ind.station.2, ], sparse = TRUE),
                  Matrix(diag(N.stations)[ind.station.1, ], sparse = TRUE)
    )
    
    # do not include the covariates "nobikes" and "nobikes" in the T-Skellam model
    if (distribution == "tskellam"){
      ind <- which(names.theta == "nobikes" | names.theta == "noboxes")
      Z.out <- Z.out[, -ind]
      Z.in <- Z.in[, -ind]
      names.theta <- names.theta[-ind]
    }
    
    # variable names
    if (distribution == "poisson"){
      names.theta <- c("intercept", gamma.names, "precamount", "tue", "wed", "thu", "fri", "ph", 
                       "nobikes", "hubout", "noboxes", "hubin", "dist",
                       paste0("station.out.", 1:N.stations),
                       paste0("station.out.", 1:N.stations))
    } else {
      names.theta <- c("intercept", gamma.names, "precamount", "tue", "wed", "thu", "fri", "ph", 
                       "nobikes", "hubout", "noboxes", "hubin", "dist",
                       paste0("station.out.", 1:(N.stations-1)), "station.out.w",
                       paste0("station.out.", 1:(N.stations-1)), "station.in.w")
    }

    
    return(list(Z.out = Z.out, Z.in = Z.in, H = H,
                ind.gamma = ind.gamma, 
                ind.beta = c(1, (sum(d.s[1:M])+2):(ncol(Z.in) - 2*N.stations)),
                ind.dyadic = ncol(Z.in) - 2*N.stations,
                names.theta = names.theta,
                B = B, z = z, k = k, K = K, 
                m = rep(m, M), l = rep(l, M)))  
  } 
  # design matrix for the station based model
  if (mode == "station"){
    
    # get indices for station and day
    ind.day <- ind.station <- NULL
    ind <- 0
    for (i in 1:N.stations) {
      d <- min(which(H >= stations$ActiveFrom[i])) - 1 #0
      for (t in H[which(H >= stations$ActiveFrom[i])]) {
        ind <- ind + 1
        d <- d + 1
        ind.station <- c(ind.station, i)
        ind.day <- c(ind.day, d)
      }
    }
    
    # part of the design matrix which is only time related
    Z.t <- cbind(B[[1]][ind.day, ], B[[2]][ind.day, ], B[[3]][ind.day, ], 
                 B[[4]][ind.day, ], 1*(prec.amount[ind.day] >= 1),
                 weekday %>% slice(ind.day) %>% as.matrix())
    # design matrix for outgoing bikes
    Z.out <- cbind(1,
                   Matrix(Z.t, sparse = TRUE),
                   Matrix(nobikes[cbind(ind.station, ind.day)], sparse = TRUE),
                   Matrix(stations$DistanceHub[ind.station]/100, sparse = TRUE),
                   0,
                   0,
                   Matrix(diag(2*N.stations)[ind.station, ], sparse = TRUE)
    )
    # design matrix for incoming bikes
    Z.in <- cbind(1,
                  Matrix(Z.t, sparse = TRUE),
                  0,
                  0,
                  Matrix(noboxes[cbind(ind.station, ind.day)], sparse = TRUE),
                  Matrix(stations$DistanceHub[ind.station]/100, sparse = TRUE),
                  Matrix(diag(2*N.stations)[ind.station + N.stations, ], sparse = TRUE)
    )
    
    # variable names
    names.theta <- c("intercept", gamma.names, "precamount", "tue", "wed", "thu", "fri", "ph", 
                     "nobikes", "hubout", "noboxes", "hubin",
                     paste0("station.out.", 1:N.stations), paste0("station.out.", 1:N.stations))
    
    # do not include the covariates "nobikes" and "nobikes" in the T-Skellam model
    if (distribution == "tskellam"){
      ind <- which(names.theta == "nobikes" | names.theta == "noboxes")
      Z.out <- Z.out[, -ind]
      Z.in <- Z.in[, -ind]
      names.theta <- names.theta[-ind]
    }
    
    return(list(Z.in = Z.in, Z.out = Z.out, 
                ind.gamma = ind.gamma, 
                ind.beta = c(1, (sum(d.s[1:M]) + 2):(ncol(Z.out) - 2*N.stations)), 
                names.theta = names.theta,
                H = H, B = B, z = z, k = k, K = K, 
                m = rep(m, M), l = rep(l, M)))  
  } 
}

get.differences <- function(design, mode, interventions, S = 10){ 
  
  # check if provider interventions are saved in the memore if the differences
  # should be corrected for provider interventions
  if (interventions == "corrected" & !exists("P")){
    stop("Error. Please supply provider interventions.")
  }
  
  # for numerical reasons we ignore in the dyadic model that
  # some stations only get active during the course of the year
  if (mode == "dyadic") {
    stations$ActiveFrom <- 1
  }
  
  # determine differences of station feeds and save them as a matrix
  D <- matrix(NA, length(design$H), N.stations)
  for (i in 1:N.stations) {
    H <- design$H[which(design$H >= stations$ActiveFrom[i])]
    ind <- which(is.element(feeds[[i]]$t, H))
    D.i <- feeds[[i]]$C[ind] - feeds[[i]]$C[ind - 1] + 
      (interventions == "corrected")*(P$outgoing[H, i] - P$incoming[H, i])
    D[which(design$H >= stations$ActiveFrom[i]), i] <- D.i
    # threshold correction
    if (interventions == "threshold"){
      D.i.pre <- feeds[[i]]$C[ind - 1] - feeds[[i]]$C[ind - 2]
      D.i.post <- feeds[[i]]$C[ind + 1] - feeds[[i]]$C[ind]
      ind.interventions <- which(abs(D.i) > pmax(abs(D.i.pre), abs(D.i.post)) + 10)
      print(ind.interventions)
      D[which(design$H >= stations$ActiveFrom[i])[ind.interventions], i] <- Inf
    }
  }
  
  # compute the latent station's feed in the dyadic model
  if (mode == "dyadic") {
    D.w <- rep(0, length(design$H))
    for (d in 1:length(design$H)) {
      D.w[d] <- -sum(D[d, ][which(!is.na(D[d, ]))])
    }
    D <- cbind(D, D.w)
  }
  
  # save differences as a vector
  D <- D[which(!is.na(D))]
  
  # remove threshold values and adjust the design matrix
  if (length(which(is.infinite(D))) > 0){
    design$Z.in <- design$Z.in[-which(is.infinite(D)), ]
    design$Z.out <- design$Z.out[-which(is.infinite(D)), ]
    D <- D[-which(is.infinite(D))]
  }

  return(list(D = D, design.D = design))
}

get.counts <- function(data, mode, design){
  
  # get real counts in the dyadic model and the station based model
  if (mode == "dyadic"){
    N.h <- rep(NA, nrow(design$Z.out))
    ind <- 1
    for (i in 1:N.stations) {
      for (t in design$H) {
        N.h[ind:(ind + N.stations - 1)] <- data$dyadic[t, i, ]
        ind <- ind + N.stations
      }
    }
  } else {
    N.out.h <- N.in.h <- NULL
    for (i in 1:N.stations) {
      H <- design$H[which(design$H >= stations$ActiveFrom[i])]
      N.out.h <- c(N.out.h, data$outgoing[H, i])
      N.in.h <- c(N.in.h, data$incoming[H, i])
    }
    N.h <- c(N.out.h, N.in.h)
  }
  return(N.h)
}

get.feeds <- function(design, mode, D){ 
  
  # get a data frame which contain the maximum station feeds C.max and the 
  # station feeds C.0 at the beginning of each interval depending on 
  # station i and time point t
  N.stations <- N.stations + 1*(mode == "dyadic")
  C <- data.frame(matrix(NA, nrow = length(D), ncol = 5))
  colnames(C) <- c("C.max", "C.0", "C.1", "i", "t")
  
  ind <- 1
  for (i in 1:N.stations) { 
    H <- design$H[which(design$H >= stations$ActiveFrom[i])]
    C$C.max[ind:(ind + length(which(design$H >= stations$ActiveFrom[i])) - 1)] <- 
      feeds[[i]]$Cmax[which(is.element(feeds[[i]]$t, H)) - 1]
    
    C$C.0[ind:(ind + length(which(design$H >= stations$ActiveFrom[i])) - 1)] <- 
      feeds[[i]]$C[which(is.element(feeds[[i]]$t, H)) - 1]
    C$C.1[ind:(ind + length(which(design$H >= stations$ActiveFrom[i])) - 1)] <- 
      feeds[[i]]$C[which(is.element(feeds[[i]]$t, H))]
    
    C$i[ind:(ind + length(which(design$H >= stations$ActiveFrom[i])) - 1)] <- i
    
    C$t[ind:(ind + length(which(design$H >= stations$ActiveFrom[i])) - 1)] <- 
      feeds[[i]]$t[which(is.element(feeds[[i]]$t, H)) - 1]
    
    ind <- ind + length(which(design$H >= stations$ActiveFrom[i]))
  } 
  return(C)
} 


# model fit ----

fit.model <- function(design, mode, distribution, name, D = NULL, C = NULL, counts = NULL){
  
  message(paste0("Fit model ", name))
  
  # initialize
  N.stations <- N.stations + 1*(mode == "dyadic" & distribution != "poisson")
  p <- ncol(design$Z.in)
  theta <- rep(0, p)
  p.fixed <- p - 2*N.stations
  delta <- Inf
  it <- 0

  # start approximate EM algorithm
  while (delta > eps) {
    it <- it + 1
    fit <- NULL
    cat(paste("Iteration ", it, "\n"))
    # dyadic skellam model
    if (mode == "dyadic" & distribution != "poisson"){
      if (it == 1){
        theta[1] <- -5
        theta[design$ind.dyadic] <- 1
        lower <- rep(-Inf, p)
        lower[design$ind.dyadic] <- 0
      }
      fit <- optim(par = theta, fn = logL.skellam, gr = score.skellam, method = "L-BFGS-B", hessian = TRUE, 
                   lower = lower, control = list(fnscale = -1, maxit = 1500, factr = 1e4, trace = trace),
                   design = design, mode = mode, distribution = distribution, 
                   N.stations = N.stations, D = D, C = C, Sigma = Sigma, lambda = lambda)
      print(fit$par[design$ind.beta])
    } else {
      while(is.null(fit)){
        fit <- tryCatch(
          {
            switch (distribution,
                    "poisson" = optim(par = theta, fn = logL.poisson, gr = score.poisson, method = "L-BFGS-B", 
                                      control = list(fnscale = -1, maxit = 1000, factr = 1e4, pgtol = 1, trace = trace),
                                      design = design, mode = mode, 
                                      N.stations = N.stations, counts = counts, Sigma = Sigma, lambda = lambda),
                    optim(par = theta, fn = logL.skellam, gr = score.skellam, method = "L-BFGS-B", hessian = TRUE,
                          control = list(fnscale = -1, maxit = 1000, factr = 1e4, pgtol = 1, trace = trace),
                          design = design, mode = mode, distribution = distribution, 
                          N.stations = N.stations, D = D, C = C, Sigma = Sigma, lambda = lambda)
            )
          },
          error = function(cond) {
            message(cond)
            message("\nTry to refit the model with another starting value.")
            return(NULL)
          }
        )
        theta <- theta + runif(p, -0.5, 0.5)
      }
    } 
    theta <- fit$par
    sc <- switch(distribution,
                 "poisson" = score.poisson(theta = theta, design = design, mode = mode, 
                                           N.stations = N.stations, counts = counts, Sigma = Sigma, lambda = lambda),
                 score.skellam(theta = theta, design = design, mode = mode, distribution = distribution, 
                               N.stations = N.stations, D = D, C = C, Sigma = Sigma, lambda = lambda))
    
    if (distribution == "poisson") {
      fisher <- fisher.poisson(theta = theta, design= design, mode = mode, N.stations = N.stations, Sigma = Sigma, lambda = lambda)
    } else {
      fisher <- -fit$hessian
    }

    if (rankMatrix(fisher) == p){
      V <- solve(fisher)
    } else {
      print("Fisher-Information computationally singular, used Moore-Penrose-Inverse")
      V <- ginv(fisher)
    }
    
    # random effects
    if (!is.null(Sigma)){
      u <- matrix(theta[(p -2*N.stations + 1):p], ncol = 2)
      Sigma.new <- matrix(0, 2, 2)
      for (i in 1:N.stations) {
        Sigma.new = Sigma.new + V[c(p.fixed + i, p.fixed + N.stations + i), c(p.fixed + i, p.fixed + N.stations + i)] +
          u[i, ]%*%t(u[i, ])
      }
      Sigma.new <- Sigma.new/N.stations
      delta <- norm(Sigma.new - Sigma, type = "F")/norm(Sigma, type = "F")
      Sigma <- Sigma.new
    } else break
    
    # smooth functions
    if (length(design$B) > 0){
      lambda.new <- rep(0, length(lambda))
      for (m in 1:length(design$B)) {
        lambda.new[m] <- lambda[m]*max((sum(diag(ginv(lambda[m]*design$K[[m]])%*%design$K[[m]]))-
                                          sum(diag((V[design$ind.gamma[[m]], design$ind.gamma[[m]]]%*%design$K[[m]]))))/
                                         (theta[design$ind.gamma[[m]]]%*%design$K[[m]]%*%theta[design$ind.gamma[[m]]]), 0)
      }
      lambda <- lambda.new
    } else break
    
    if (it >= max.it){
      message("Maximum number of iterations has been reached. Optimization stopped.")
      break
    }
  }
  message(paste0("Model fit terminated after ", it, 
                 " iterations.\nNorm of the gradient at the ML estimator: ", 
                 round(sqrt(sum(sc^2)), 2), "."))
  return(list(theta.hat = theta, V.hat = V, Sigma.hat = Sigma))
}

# this is a wrapper around fit.model, which can be used for parallel fitting of the models
fit.models <- function(models){
  
  # model specification
  h <- models$h
  distribution <- models$distribution
  interventions <- models$interventions
  mode <- models$mode
  name <- paste(h, "_" , mode , "_",  interventions, "_", distribution, sep = "")
  
  # design of the model
  design <- get.design(h = h, mode = mode, distribution = distribution)
  saveRDS(design, 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 for this kind of handling service rides and that interval of the day
    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, mode, D)
  }
  
  # fit the model
  fit <- fit.model(design = design, mode = mode, distribution = distribution,
                   name = name, D = D, C = C, counts = counts)
  saveRDS(fit, paste0("Fits/", name, ".rds"))
}

logL.skellam <- function(theta, design, mode, distribution, N.stations, D, C, Sigma, lambda){
  
  # compute expected values of outgoing and incoming trip counts
  if (mode == "dyadic"){
    mu.out <- rowSums(matrix(exp(design$Z.out%*%theta), ncol = N.stations, byrow = TRUE)) 
    mu.in <- rowSums(matrix(exp(design$Z.in%*%theta), ncol = N.stations, byrow = TRUE)) 
  } else {
    mu.out <- as.vector(exp(design$Z.out%*%theta)) 
    mu.in <- as.vector(exp(design$Z.in%*%theta))
  }
  
  # compute logarithm of the modified Bessel function
  x <- 2*sqrt(mu.out*mu.in)
  log.bes <- x + log(besselI(x, D, expon.scaled = TRUE))
  # if x is very large (> 10^6) we use an approximation (Appendix A)
  log.bes[which(is.infinite(log.bes))] <- x[which(is.infinite(log.bes))]
  
  # compute the log likelihood depending on the distribution
  if (distribution == "skellam"){
    logL <- sum(-mu.out -mu.in + D/2*log(mu.in/mu.out) + log.bes)
  }
  if (distribution == "tskellam"){
    ind.w <- which(is.na(C$C.0))
    logL <- sum(D/2*log(mu.in/mu.out) + log.bes) - 
      sum(mu.in[ind.w] + mu.out[ind.w])
    
    for (c in sort(unique(C$C.max))) {
      ind <- which(C$C.max == c)
      d <- matrix(rep(0:c, length(ind)), ncol = c+1, byrow = TRUE) - matrix(rep(C$C.0[ind], c+1), ncol = c+1)
      x <- 2*sqrt(mu.in[ind]*mu.out[ind])
      logL <- logL - sum(x + log(rowSums((mu.in[ind]/mu.out[ind])^(d/2)*besselI(x, d, expon.scaled = TRUE))))
    }
  }
  
  # penalty random effects
  if (!is.null(Sigma)){
    p <- length(theta)
    u <- matrix(theta[(p-2*N.stations+1):p], ncol = 2) # first column u^out, second column u^in
    logL <- logL - 0.5*sum(diag(u%*%tcrossprod(solve(Sigma), u)))
  }
  
  # penalty smooth functions
  M <- length(design$K)
  if (M > 0){
    for (i in 1:M) {
      logL <- logL - 0.5*lambda[i]*
        as.vector(theta[design$ind.gamma[[i]]]%*%design$K[[i]]%*%theta[design$ind.gamma[[i]]]) 
    }
  }
  return(logL)
}

score.skellam <- function(theta, design, mode, distribution, N.stations, D, C, Sigma, lambda){
  
  # compute expected values of outgoing and incoming trip counts
  if (mode == "dyadic"){
    mu.out <- rowSums(matrix(exp(design$Z.out%*%theta), ncol = N.stations, byrow = TRUE)) 
    mu.in <- rowSums(matrix(exp(design$Z.in%*%theta), ncol = N.stations, byrow = TRUE)) 
  } else {
    mu.out <- as.vector(exp(design$Z.out%*%theta)) 
    mu.in <- as.vector(exp(design$Z.in%*%theta))
  }
  
  # compute the argument of the modified Bessel function and the compute the derivates of the log likelihood
  x <- 2*sqrt(mu.out*mu.in)
  frac.skellam <- besselI(x, D+1, expon.scaled = TRUE)/besselI(x, D, expon.scaled = TRUE)
  # if x is very large (> 10^6) we use an approximation (Appendix A)
  frac.skellam[which(!is.finite(frac.skellam))] <- 1
  if (distribution == "skellam"){
    delta.l.D.mu.out = -1 + sqrt(mu.in/mu.out)*frac.skellam
    delta.l.D.mu.in = -1 + D/mu.in + sqrt(mu.out/mu.in)*frac.skellam
  } 
  if (distribution == "tskellam"){
    delta.l.D.mu.out <- sqrt(mu.in/mu.out)*besselI(x, D+1, expon.scaled = TRUE)/besselI(x, D, expon.scaled = TRUE)
    delta.l.D.mu.in <- D/mu.in + sqrt(mu.out/mu.in)*besselI(x, D+1, expon.scaled = TRUE)/besselI(x, D, expon.scaled = TRUE)
    
    # skellam distriubtion for latent station
    ind.w = which(is.na(C$C.0))
    delta.l.D.mu.in[ind.w] <- delta.l.D.mu.in[ind.w] - 1
    delta.l.D.mu.out[ind.w] <- delta.l.D.mu.out[ind.w] - 1
    
    for (c in sort(unique(C$C.max))) {
      ind = which(C$C.max == c)
      d = matrix(rep(0:c, length(ind)), ncol = c+1, byrow = TRUE) - matrix(rep(C$C.0[ind], c+1), ncol = c+1)
      
      delta.l.D.mu.in[ind] <- delta.l.D.mu.in[ind] -
        1/rowSums((mu.in[ind]/mu.out[ind])^(d/2)*besselI(2*sqrt(mu.in[ind]*mu.out[ind]), d, expon.scaled = TRUE))* 
        rowSums(d/mu.in[ind]*(mu.in[ind]/mu.out[ind])^(d/2)*besselI(2*sqrt(mu.in[ind]*mu.out[ind]), d, expon.scaled = TRUE) + 
                  (mu.in[ind]/mu.out[ind])^((d-1)/2)*besselI(2*sqrt(mu.in[ind]*mu.out[ind]), d+1, expon.scaled = TRUE))
      
      delta.l.D.mu.out[ind] <- delta.l.D.mu.out[ind] -
        1/rowSums((mu.in[ind]/mu.out[ind])^(d/2)*besselI(2*sqrt(mu.in[ind]*mu.out[ind]), d, expon.scaled = TRUE))*
        rowSums((mu.in[ind]/mu.out[ind])^((d+1)/2)*besselI(2*sqrt(mu.in[ind]*mu.out[ind]), d+1, expon.scaled = TRUE))
    }
  }
  
  # compute the score function depending on the model
  if (mode == "dyadic"){
    score <- as.vector(colSums((delta.l.D.mu.in[rep(1:(N.stations*length(design$H)), each = N.stations)]*
                                  as.vector(exp(design$Z.in%*%theta)))*design$Z.in) + 
                        colSums((delta.l.D.mu.out[rep(1:(N.stations*length(design$H)), each = N.stations)]*
                                   as.vector(exp(design$Z.out%*%theta)))*design$Z.out))
  } else {
    score <- colSums(delta.l.D.mu.in*as.vector(exp(design$Z.in%*%theta))*design$Z.in) + 
      colSums(delta.l.D.mu.out*as.vector(exp(design$Z.out%*%theta))*design$Z.out) 
  }
  
  # penalty random effects
  if(!is.null(Sigma)){
    p <- length(theta)
    p.fixed <- p - 2*N.stations
    u <- matrix(theta[(p-2*N.stations+1):p], ncol = 2) # first column u^out, second column u^in
    for (i in 1:N.stations) {
      score[c(p.fixed+i, p.fixed+N.stations+i)] = score[c(p.fixed+i, p.fixed+N.stations+i)] - solve(Sigma)%*%u[i, ]
    }
  }
  
  # penalty smooth functions
  M <- length(design$K)
  if (M > 0){
    for (i in 1:M) {
      score[design$ind.gamma[[i]]] <- score[design$ind.gamma[[i]]] - 
        lambda[i]*design$K[[i]]%*%theta[design$ind.gamma[[i]]]
    }
  }
  return(score)
}

logL.poisson <- function(theta, design, mode, N.stations, counts, Sigma, lambda){

  # compute expected values of outgoing and incoming trip counts
  if (mode == "dyadic"){
    mu <- as.vector(exp(design$Z.out%*%theta)) 
  } else {
    mu.out <- as.vector(exp(design$Z.out%*%theta)) 
    mu.in <- as.vector(exp(design$Z.in%*%theta))
  }
  
  # compute the unpenalized log-likelihood
  logL <- switch(mode,
                 "station" = sum(counts*log(c(mu.out, mu.in)) - c(mu.out, mu.in)),
                 "dyadic" = sum(counts*log(mu) - mu)
                 ) 
  
  # penalty random effects
  if (!is.null(Sigma)){
    p <- length(theta)
    u <- matrix(theta[(p - 2*N.stations + 1):p], ncol = 2) # first column u^out, second column u^in
    logL <- logL - 0.5*sum(diag(u%*%tcrossprod(solve(Sigma), u)))
  }
  
  M <- length(design$K)
  # penalty splines
  if (M > 0){
    for (i in 1:M) {
      logL <- logL - 0.5*lambda[i]*
        as.vector(theta[design$ind.gamma[[i]]]%*%design$K[[i]]%*%theta[design$ind.gamma[[i]]]) 
    }
  }
  
  return(logL)
}

score.poisson <- function(theta, design, mode, N.stations, counts, Sigma, lambda){
  
  # compute expected values of outgoing and incoming trip counts
  if (mode == "dyadic"){
    mu <- as.vector(exp(design$Z.out%*%theta)) 
  } else {
    mu.out <- as.vector(exp(design$Z.out%*%theta)) 
    mu.in <- as.vector(exp(design$Z.in%*%theta))
  }
  
  # compute the unpenalized score function
  score <- switch(mode, 
                  "dyadic" = colSums((counts - mu)*design$Z.out),
                  "station" = colSums((counts - c(mu.out, mu.in))*rbind(design$Z.out, design$Z.in))
                  )
  
  #penalty random effects
  if(!is.null(Sigma)){
    p <- length(theta)
    p.fixed <- p - 2*N.stations
    u <- matrix(theta[(p - 2*N.stations + 1):p], ncol = 2) # first column u^out, second column u^in
    for (i in 1:N.stations) {
      score[c(p.fixed + i, p.fixed + N.stations + i)] <- score[c(p.fixed + i, p.fixed + N.stations + i)] - 
        solve(Sigma)%*%u[i, ]
    }
  }
  
  # penalty splines
  M <- length(design$K)
  if (M > 0){
    for (i in 1:M) {
      score[design$ind.gamma[[i]]] <- score[design$ind.gamma[[i]]] - 
        lambda[i]*design$K[[i]]%*%theta[design$ind.gamma[[i]]]
    }
  }
  
  return(score)
}

fisher.poisson <- function(theta, design, mode, N.stations, Sigma, lambda){
  
  # compute expected values of outgoing and incoming trip counts
  if (mode == "dyadic"){
    mu <- as.vector(exp(design$Z.out%*%theta)) 
  } else {
    mu.out <- as.vector(exp(design$Z.out%*%theta)) 
    mu.in <- as.vector(exp(design$Z.in%*%theta))
  }
  
  # compute the unpenalized Fisher information
  Z <- switch(mode,
              "dyadic" = design$Z.out,
              "station" = rbind(design$Z.out, design$Z.in)) 
  mu <- switch(mode,
               "dyadic" = mu,
               "station" = c(mu.out, mu.in))
  Mu <- Matrix(0, nrow(Z), nrow(Z))
  diag(Mu) <- mu
  fisher <- t(Z)%*%Mu%*%Z
 
  #penalty random effects
  if(!is.null(Sigma)){
    p <- length(theta)
    p.fixed <- p - 2*N.stations
    for (i in 1:N.stations) {
      fisher[c(p.fixed + i, p.fixed + N.stations + i), c(p.fixed + i, p.fixed + N.stations + i)] <- 
        fisher[c(p.fixed + i, p.fixed + N.stations + i), c(p.fixed + i, p.fixed + N.stations + i)] + solve(Sigma)
    }
  }
  
  # penalty splines
  M <- length(design$K)
  if (M > 0){
    for (i in 1:M) {
      fisher[design$ind.gamma[[i]], design$ind.gamma[[i]]] <- fisher[design$ind.gamma[[i]], design$ind.gamma[[i]]] + 
        lambda[i]*design$K[[i]]
    }
  }
  
  return(fisher)
}


# post fit ----

confidence.band <- function(gamma, cov, z, B, q = 0.05, S = 10000){
  set.seed(1)
  mu.sim <- matrix(0, S, length(z))
  # simulate from the posterior of the spline coefficients
  for (i in 1:S) {
    gamma.sim <- rmvn(1, gamma, cov)
    mu.sim[i, ] <- B%*%gamma.sim
  }
  lower <- upper <- rep(0, ncol(mu.sim))
  # compute the quantiles
  for (j in 1:ncol(mu.sim)) {
    lower[j] <- quantile(mu.sim[, j], probs = q/2)
    upper[j] <- quantile(mu.sim[, j], probs = 1-q/2)
  }
  return(list(lower = lower, upper = upper))
}

estimate.flows <- function(fit, design, mode, distribution){
  
  N.stations <- N.stations + 1*(mode == "dyadic" & distribution != "poisson")
  if (mode == "dyadic"){
    N.dyadic.hat <- array(NA, dim = c(length(design$H), N.stations, N.stations))
    for (i in 1:N.stations) {
      N.dyadic.hat[, i, ] <- 
        matrix(exp(design$Z.out%*%fit$theta.hat)[((i-1)*(N.stations*length(design$H)) + 1):(i*N.stations*length(design$H))], 
               nrow = length(design$H), 
               ncol = N.stations,
               byrow = TRUE)
    }
    # latent station in the dyadic (T)-Skellam model
    if (distribution != "poisson"){
      N.stations <- N.stations - 1
      N.dyadic.hat.w <- N.dyadic.hat
      N.dyadic.hat <- array(NA, dim = c(length(design$H), N.stations, N.stations))
      for (d in 1:length(design$H)) {
        for (i in 1:N.stations) {
          pi <- N.dyadic.hat.w[d, i, 1:N.stations]/sum(N.dyadic.hat.w[d, i, 1:N.stations])
          N.dyadic.hat[d, i, ] <- N.dyadic.hat.w[d, i, 1:N.stations] + N.dyadic.hat.w[d, i, N.stations + 1]*pi
        }
      }
    }
    return(N.dyadic.hat)
  } else {
    N.out.hat <- N.in.hat <- matrix(0, nrow = length(design$H), ncol = N.stations)
    ind <- 1
    for (i in 1:N.stations) {
      H <- design$H[which(design$H > stations$ActiveFrom[i])]
      N.out.hat[which(design$H >= stations$ActiveFrom[i]), i] <- 
        exp(design$Z.out%*%fit$theta.hat)[ind:(ind + length(H) - 1)]
      N.in.hat[which(design$H >= stations$ActiveFrom[i]), i] <- 
        exp(design$Z.in%*%fit$theta.hat)[ind:(ind + length(H) - 1)]
      ind <- ind + length(H)
    }
    return(list(N.out.hat = N.out.hat, N.in.hat = N.in.hat))
  }
}
