get.design.simulation <- function(t.max, N.stations, z.time, z.dyad){
  
  # add the latent station
  N.stations <- N.stations + 1
  
  # indices for time and station
  ind.day <- rep(rep(1:t.max, each = N.stations), N.stations)
  ind.station.1 <- rep(1:N.stations, each = N.stations*t.max)
  ind.station.2 <- rep(1:N.stations, N.stations*t.max)
  
  # design matrices for outgoing and incoming trips
  Z.out <- cbind(1,
                 Matrix(z.time[ind.day], sparse = TRUE),
                 Matrix(z.dyad[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))
  Z.in <- cbind(1,
                 Matrix(z.time[ind.day], sparse = TRUE),
                 Matrix(z.dyad[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))
  return(list(Z.out = Z.out, Z.in = Z.in, H = 1:t.max))
}

simulate.network <- function(s, N.stations, t.max, beta, Sigma, sim.parallel){
  
  # simulate time-dependent and dyadic covariate
  z.time <- rnorm(t.max)
  z.dyad <- matrix(0, N.stations, N.stations)
  for (i in 1:N.stations) {
    for (j in i:N.stations) {
      z.dyad[i, j] <- z.dyad[j, i] <- rnorm(1)
    }
  }
  z.dyad <- cbind(z.dyad, 0)
  z.dyad <- rbind(z.dyad, 0)
  
  # save as design matrix
  design <- get.design.simulation(t.max, N.stations, z.time, z.dyad)
  Z.out <- design$Z.out
  Z.in <- design$Z.in
  
  # simulate random effects
  u <- mvrnorm(N.stations, c(0, 0), Sigma) 
  u.out <- u[ ,1]
  u.in <- u[, 2] 
  
  # compute expected values
  factor <- c(log(0.9), 0)
  mu <- array(0, dim = c(2, t.max, N.stations, N.stations))
  for (h in 1:2) {
    for (t in 1:t.max) {
      for (i in 1:N.stations) {
        for (j in 1:N.stations) {
          mu[h, t, i, j] <- exp(beta[1] + beta[2]*z.time[t] + beta[3]*z.dyad[i, j] + 
                                 u.out[i] + u.in[j] + factor[h])
        }
      }
    }  
  }
  mu.out <- mu.in <- matrix(0, t.max, N.stations)
  for (t in 1:t.max) {
    mu.out[t, ] <- rowSums(mu[2, t, , ])
    mu.in[t, ] <- 1/3*colSums(mu[1, t, , ]) + 2/3*colSums(mu[2, t, , ])
  }
  
  # simulate flows or return the design
  if (sim.parallel){
    fit <- simulate.flows(s, Sigma, N.stations, t.max, design, mu, sim.parallel)
    return(fit)
  } else {
    return(list(design = design, z.time = z.time, z.dyad = z.dyad, mu = mu, u = u))
  }
}

simulate.flows = function(s, Sigma, N.stations, t.max, design, mu, sim.parallel){
  
  # simulate the flows
  N <- array(0, dim = c(2, t.max, N.stations, N.stations))
  N.out <- N.in  <- array(0, dim = c(3, t.max, N.stations))
  for (h in 1:2) {
    for (t in 1:t.max) {
      for (i in 1:N.stations) {
        for (j in 1:N.stations) {
          N[h, t, i, j] <- rpois(1, mu[h, t, i, j])
          if (N[h, t, i, j] > 0){
            N.out[h, t, i] <- N.out[h, t, i] + N[h, t, i, j]
            for (r in 1:N[h, t, i, j]) {
              end <- rbinom(1, 1, 1/3)
              N.in[h + end, t, j] <- N.in[h+end, t, j] + 1
            }
          }
        }
      }
    }  
  }
  # compute the differences
  D = cbind(N.in[2, , ] - N.out[2, , ], -rowSums(N.in[2, , ] - N.out[2, , ]))
  
  # fit the model to the simulated data or return the simulated flows
  if (sim.parallel){
    fit <- NULL
    while (is.null(fit)) {
      fit <- fit.model.simulation(s, N.stations, as.vector(D), design, Sigma)
    }
    return(fit)
  } else {
    return(list(D = D, N = N, N.out = N.out, N.in = N.in))
  }
}

fit.model.simulation <- function(s, N.stations, D, design, Sigma){
  
  N.stations <- N.stations + 1
  p <- ncol(design$Z.in)
  p.fixed <- p - 2*N.stations
  
  # find a reasonable start value
  logL <- rep(0, 100)
  theta <- matrix(runif(100*p, -1, 1), nrow = 100, ncol = p)
  for (k in 1:100) {
    logL[k] <- logL.skellam(theta = theta[k, ], design = design, mode = "dyadic", distribution = "skellam", 
                            N.stations = N.stations, D = D, C = NULL, Sigma = Sigma, lambda = NULL)
  }
  ind.inf <- which(!is.finite(logL))
  if (length(ind.inf) > 0){
    theta <- theta[- ind.inf, ]
    logL <- logL[-ind.inf]
  } 
  if (length(ind.inf) == 100){
    return(NULL)
  }
  theta <- theta[which.max(logL), ]

  # start approximate EM-algorithm
  Delta  <- Inf
  it <- 0
  while(Delta > eps) {
    it <- it + 1
    cat(paste("Iteration ", it, "\n"))
    
    fit <- withTimeout(optim(par = theta, fn = logL.skellam, gr = score.skellam, method = "BFGS", 
                            control = list(fnscale = -1, reltol = 10^(-18), maxit = 400, trace = trace), hessian = TRUE,
                            design = design, mode = "dyadic", distribution = "skellam", N.stations = N.stations,
                            D = D, Sigma = Sigma), 
                       timeout = 500, onTimeout = "silent")
    # refit the model if model fit takes too long or maximum number of iterations has been reached
    if (is.null(fit) | it > max.it){
      return(NULL)
    }
    theta <- fit$par
    u.hat <- matrix(theta[(p - 2*N.stations + 1):p], ncol = 2)

    # refit the model if fitted random effects are clearly not centered around zero
    if (sum(abs(colSums(u.hat))) > 1){
      return(NULL)
    }
    
    # refit the model if observed fisher information does not have full rank
    fisher <- -fit$hessian
    if (rankMatrix(fisher) == p){
      V.hat <- solve(fisher)
    } else {
      return(NULL)
    }
    
    # update the covariance matrix of random effects
    Sigma.new <- Sigma.star <- matrix(0, 2, 2)
    for (i in 1:N.stations) {
      Sigma.new <- Sigma.new + V.hat[(p.fixed + 2*i - 1):(p.fixed + 2*i), 
                                 (p.fixed + 2*i - 1):(p.fixed + 2*i)] + 
        u.hat[i, ]%*%t(u.hat[i, ])
      if (i < N.stations){
        Sigma.star <- Sigma.new
      }
    }
    Sigma.new <- Sigma.new/N.stations
    Sigma.star <- Sigma.star/(N.stations - 1)

    # compute norm of the relative difference between the old and the new Sigma
    Delta <- norm(Sigma.new - Sigma, type = "F")/norm(Sigma, type = "F")
    Sigma <- Sigma.new
  }
  
  return(list(theta.hat = theta, Sigma.hat = Sigma.new, Sigma.star = Sigma.star))
}

simulation.truncated.skellam <- function(s, beta, t.max, range.C.max){
  
  # design of the model
  z.t <- rnorm(t.max)
  mu.out <- exp(beta[1]*z.t + beta[2])
  mu.in <- exp(beta[1]*z.t + beta[3])
  Z.out <- cbind(z.t, rep(1, t.max), rep(0, t.max))
  Z.in <- cbind(z.t, rep(0, t.max), rep(1, t.max))
  design <- list(Z.in = Z.in, Z.out = Z.out)

  # station feeds
  if (length(range.C.max) > 1){
    C.max <- sample(range.C.max, size = t.max, replace = TRUE)
  } else {
    C.max <- rep(range.C.max, t.max)
  }
  C.0 <- rep(NA, t.max)
  C.1 <- rep(NA, t.max)
  
  
  # simulate as described in Appendix D
  N.in <- N.out <- rep(0, t.max)
  for (i in 1:t.max) {
    k <- 0
    C.0[i] <- sample(0:C.max[i], size = 1)
    C.1[i] <- C.0[i]
    while(k < 1){
      k <- k + rexp(1, mu.in[i] + mu.out[i])
      if (k < 1){
        in.out <- sample(c(1,2), size = 1, prob = c(mu.in[i]/(mu.in[i]+mu.out[i]), mu.out[i]/(mu.in[i]+mu.out[i])))
        if (in.out == 1 & C.1[i] < C.max[i]){
          N.in[i] <- N.in[i] + 1
          C.1[i] <- C.1[i] + 1
        }
        if (in.out == 2 & C.1[i] > 0){
          N.out[i] <- N.out[i] + 1
          C.1[i] <- C.1[i] - 1
        }
      }
    }
  }
  D <- C.1 - C.0
  C <- tibble(C.max = C.max, C.0 = C.0)
  
  # fit the models
  theta <- rep(0, 3)
  fit.skellam <- optim(par = theta, fn = logL.skellam, gr = score.skellam, D = D, method = "L-BFGS-B", mode = "station",
                      control = list(fnscale = -1, factr = 1e4, maxit = 1000), hessian = FALSE,
                      design = design, Sigma = NULL, distribution = "skellam")

  fit.tskellam <- optim(par = theta, fn = logL.skellam, gr = score.skellam, D = D, method = "L-BFGS-B", mode = "station",
                    control = list(fnscale = -1, factr = 1e4, maxit = 1000), hessian = FALSE,
                    C = C, design = design, Sigma = NULL, distribution = "tskellam")

  # results
  mu.in.skellam <- exp(design$Z.in%*%fit.skellam$par)
  mu.in.tskellam <- exp(design$Z.in%*%fit.tskellam$par)

  mu.out.skellam <- exp(design$Z.out%*%fit.skellam$par)
  mu.out.tskellam <- exp(design$Z.out%*%fit.tskellam$par)

  delta.in.skellam <- (sum(N.in) - sum(mu.in.skellam))/sum(N.in)
  delta.in.tskellam <- (sum(N.in) - sum(mu.in.tskellam))/sum(N.in)
  
  delta.out.skellam <- (sum(N.out) - sum(mu.out.skellam))/sum(N.out)
  delta.out.tskellam <- (sum(N.out) - sum(mu.out.tskellam))/sum(N.out)
  
  return(list(theta.skellam = fit.skellam$par, theta.tskellam = fit.tskellam$par,
              delta.out.skellam = delta.out.skellam, delta.out.tskellam = delta.out.tskellam,
              delta.in.skellam = delta.in.skellam, delta.in.tskellam = delta.in.tskellam))
}
