library(scoop)

# Monotone functions necessary for PLAMM-1.
# This code is based on the implementation of mslasso, found at:
# https://www.mn.uio.no/math/english/people/aca/glad/r-scripts/mslasso/

isplineDesign <- function(x,knot.vec){
  #function  for computation of I-Splines of degree 2,cf Ramsay 88
  #modified from http://www.stat.uni-muenchen.de/institut/lehrstuhl/semsto/software/
  
  I <- numeric()
  deg <- 2
  for(i in 1:(length(knot.vec)-deg)){
    if(x<knot.vec[i]){
      I[i] <- 0
    }
    if (x>knot.vec[i+2]) {
      I[i] <- 1
    }else{
      if(x>=knot.vec[i] && x<=knot.vec[i+1]){
        I[i] <- (x-knot.vec[i])^2 / ((knot.vec[i+1]-knot.vec[i])*(knot.vec[i+2]-knot.vec[i]))
      }
      if(x>=knot.vec[i+1] && x<=knot.vec[i+2]){
        I[i] <- 1-(knot.vec[i+2]-x)^2 / ((knot.vec[i+2]-knot.vec[i])*(knot.vec[i+2]-knot.vec[i+1]))
      }
    }
  }
  return(I)
}

monotone.basis <- function(X,k,spline,xl,xr,intercept, knot.vec2 = NULL){
  n <- length(X)
  dx <- (xr-xl)/(k-1)
  num.knots <- k-2
  
  if(spline=="ispline"){
    deg <- 2
    knot.vec <- seq(xl-(deg-1)*dx,xr+(deg-1)*dx,by=dx)
    knot.vec[(deg+1):(num.knots+2)] <- quantile(unique(as.vector(X)), seq(0,1,length = (num.knots+2)))[-c(1,(num.knots+2))]
    if (is.null(knot.vec2)){
      B <- t(apply(as.matrix(X),1,"isplineDesign",knot.vec=knot.vec))
    }
    else{
      B <- t(apply(as.matrix(X),1,"isplineDesign",knot.vec=knot.vec2))
      knot.vec <- knot.vec2
    }
  }
  res <- list("B"=B, "knot.vec"=knot.vec)
  return(res)
}

monotone.splines <- function(Xf,num.knots, knot.vec2 = NULL){
  pf <- ncol(Xf)
  Z <- NULL
  kv <- NULL
  knotv <- NULL
  for(j in 1:pf){
    if (!is.null(knot.vec2)){
      knotv <- knot.vec2[, j]
    }
    x <- Xf[,j]
    spline <- monotone.basis(x, (num.knots+2), "ispline", xl = min(x), xr = max(x), FALSE, knotv)
    z1 <- spline$B
    knot.vec <- spline$knot.vec
    Z=cbind(Z,z1)
    kv = cbind(kv, knot.vec)
    Z <- scale(Z, center = TRUE, scale = FALSE)
  }
  res = list("Z" = Z, knot.vec = kv)
  return(res)
}


#########################################################################################################

partlin = function(X, Z, Y, family = "gaussian", num.knotsf, lambda = NULL, alpha, knot.vec = NULL){
  #Z matrix with non-linear effects, X matrix with linear effects
  groups <- as.vector(t(matrix(rep((1):(ncol(Z)),(num.knotsf+2)),ncol(Z),(num.knotsf+2))))
  groups <- c(groups, seq(ncol(Z)+1, ncol(Z)+ncol(X), by = 1))
  Z_I <- monotone.splines(Z, num.knotsf, knot.vec2 = knot.vec)$Z
  Desmat <- cbind(Z_I, X)
  coopfit <- coop.lasso(Desmat, Y, groups, wk = (tabulate(groups)) ^ alpha, intercept = FALSE, lambda = lambda)
  return(coopfit)
}

adaptivepartlin = function(X, Z, Y, family = "gaussian", num.knotsf,w, lambda = NULL, alpha, knot.vec = NULL){
  #Z matrix with non-linear effects, X matrix with linear effects
  #Vector with group categories for each column in the design matrix
  groups <- as.vector(t(matrix(rep((1):(ncol(Z)),(num.knotsf+2)),ncol(Z),(num.knotsf+2))))
  groups <- c(groups, seq(ncol(Z)+1, ncol(Z)+ncol(X), by = 1))
  Z_I <- monotone.splines(Z, num.knotsf, knot.vec2 = knot.vec)$Z
  Desmat <- cbind(Z_I, X)
  coopfit <- coop.lasso(Desmat, Y, groups, wk = w * (tabulate(groups)) ^ alpha, intercept = FALSE, lambda = lambda)
  return(coopfit)
}


cvmspartlin <- function(X, Z, Y, K, method, num.knots,w = NULL, lambda, alpha){
  #Z matrix with non-linear effects, X matrix with linear effects
  error <- matrix(NA, nrow = K, ncol = 100)
  n <- nrow(X)
  antf <- floor(n/K)
  Zres <- monotone.splines(Z, num.knots)
  knot.vec <- Zres$knot.vec
  if(n==(K*antf)){
    set.seed(1)
    fold <- sample(rep(1:K, antf))
  }else{
    fold <- sample(c(rep(1:K, antf), 1:(n-(K*antf))))
  }
  
  for(k in 1:K){
    test.ind <- which(fold == k)
    trainX <- X[-test.ind,]
    trainZ <- Z[-test.ind,]
    testX <- X[test.ind,]
    testZ <- Z[test.ind, ]
    trainY <- Y[-test.ind]
    testy <- Y[test.ind]
    
    if(method == "monotone.lasso"){
      cv.fit <- partlin(trainX, trainZ, trainY, "gaussian", num.knots, lambda, alpha, knot.vec)
      Ztest <- monotone.splines(testZ, num.knots, knot.vec)$Z
    }else if(method == "adaptive.monotone"){
      cv.fit <- adaptivepartlin(trainX, trainZ, trainY, "gaussian", num.knots,w, lambda, alpha, knot.vec)
      Ztest <- monotone.splines(testZ, num.knots, knot.vec)$Z
    }
    error.k <- (testy-predict(cv.fit, newx = cbind(Ztest, testX)))^2
    error[k,] <- colSums(error.k)
  }
  MSE <- colMeans(error)
  return(MSE)
}



Ispline2 = function(x, tk, tk1, tk2){
  if(x < tk){
    return(0)
  }
  
  if(tk <= x  & x <= tk1){
    return((x-tk)^2/((tk1-tk)*(tk2-tk)))
  }
  
  if(tk1 <= x & x <= tk2){
    return (1-(tk2-x)^2/((tk2-tk)*(tk2-tk1)))
  }
  
  if(x>=tk2){
    return(1)
  }
}

############################################################################################################