# Function to create orthogonal P-spline basis
#
# Author: ___UITZONDERINGSGROND_2___ - RIVM
#
# Returns unpenalized null space (polynomial) and reduced rank orthogonal basis Z
# Based on the work by Scheipl (2011)
# See https://www.jstatsoft.org/index.php/jss/article/view/v043i14/v43i14.pdf, page 8 - 9

ortho_Pspline <- function(x, x.min = min(x), x.max = max(x),
  k = 20, spline.degree = 3, diff.ord = 2, cyclic = FALSE,
  rankZ = 0.999, tol = 1e-10) {

  # Helper functions
  scaleMat <- function(x, factor = 2) {
    return(x/(factor*frob(x)))
  }
  # Frobenius norm
  frob <- function(x) {
    trace <- sum(sapply(1:ncol(x), function(i) crossprod(x[, i])))
    return(sqrt(trace/nrow(x)))
  }
  # Shift vector p places in forward direction (default)
  shift <- function(x, p, forward = TRUE) {
    n <- length(x)
    if (forward) {
      x[c(seq_len(n)[-seq_len(n - p)], seq_len(n - p))]
    }
    else {
      x[c(seq_len(n)[-seq_len(p)], seq_len(p))]
    }
  }

  # Difference order should be > 0
  if (diff.ord <= 0) stop("Difference order should be > 0")

  # B-spline basis
  if (!cyclic) {
    # Model matrix for regular B-spline basis
    dx <- (x.max - x.min)/(k - spline.degree)
    knots <- seq(from = x.min - spline.degree*dx, to = x.max + spline.degree*dx, by = dx)
    B <- splines::splineDesign(x = x, knots = knots, ord = spline.degree + 1, outer.ok = TRUE)
  } else {
    # Model matrix for cyclic B-spline basis
    knots <- seq(from = x.min, to = x.max, length = k - spline.degree + 1)
    B <- mgcv::cSplineDes(x = x, knots = knots, ord = spline.degree + 1)
  }

  # Differences of order d penalize deviations from polynomials of order d − 1
  # The unpenalized effects, e.g. the constant (d = 1) or the straight line (d = 2), are called the null space of the penalty
  # The null space remains unpenalized
  # Cylic P-splines have a null space that is constant
  if (diff.ord - 1 > 0 & !cyclic) {
    X <- poly(x, degree = diff.ord - 1)
    X <- scaleMat(X)
  } else {
    X <- NULL
  }

  # Center B-spline around null space
  X1 <- cbind(rep(1, length(x)), X)
  qr.X <- qr(X1)
  B <- qr.resid(qr = qr.X, y = B)

  # Penalty matrix
  k <- ncol(B)
  D <- diff(diag(k), diff = diff.ord)
  if (cyclic) {
    for (i in 1:diff.ord) {
      D <- rbind(D, shift(D[k - diff.ord, ], p = i))
    }
  }
  P <- crossprod(D)

  # Orthogonal decomposition of f(x)
  P.inv <- MASS::ginv(P)
  C <- B %*% (P.inv %*% t(B))
  eigen.C <- eigen(C, symmetric = TRUE)
  nullvals <- eigen.C$values < tol
  ncol.Z <- max(3, min(which(cumsum(eigen.C$values[!nullvals])/sum(eigen.C$values[!nullvals]) > rankZ)))
  use <- 1:ncol.Z
  Z <- t(sqrt(eigen.C$values[use]) * t(eigen.C$vectors[, use]))
  Z <- scaleMat(Z, factor = 2)

  # Return matrices
  return(list(X = X, Z = Z))
}