library(scoop)

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){
  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(X), seq(0,1,length = (num.knots+2)))[-c(1,(num.knots+2))]
    B <- t(apply(as.matrix(X),1,"isplineDesign",knot.vec=knot.vec))
  } 
  res <- list("B"=B, "knot.vec"=knot.vec)
  return(res)
}

monotone.splines <- function(Xf,num.knots){
  print("Making monotone splines")
  pf <- ncol(Xf)
  Z <- NULL
  for(j in 1:pf){
    x <- Xf[,j]
    spline <- monotone.basis(x, (num.knots+2), "ispline", xl = min(x), xr = max(x), FALSE)
    z1 <- spline$B
    Z=cbind(Z,z1)
  }
  Z <- scale(Z, center = TRUE, scale = FALSE)
  return(Z)
}


#########################################################################################################
monotone.lasso <- function(Xf, Yf, familyf, num.knotsf, lambda = NULL){
  groups <- as.vector(t(matrix(rep((1):(ncol(Xf)),(num.knotsf+2)),ncol(Xf),(num.knotsf+2))))
  Yc <- Yf-mean(Yf)
  Z <- monotone.splines(Xf, num.knotsf)  
  print("Fitting monotone lasso")
  coopfit <- coop.lasso(Z,Yc, groups, family = familyf, intercept = FALSE, lambda = lambda)
  msfit <- coopfit	
  return(msfit)
}

adaptive.monotone <- function(Xf, Yf, familyf, num.knotsf, w, lambda = NULL){
 groups <- as.vector(t(matrix(rep((1):(ncol(Xf)),(num.knotsf+2)),ncol(Xf),(num.knotsf+2))))
 Yc <- Yf-mean(Yf)
 Z <- monotone.splines(Xf, num.knotsf)
 print("Fitting adaptive monotone lasso")
 adaptivecoopfit <- coop.lasso(Z, Yc, groups, family = familyf, wk = w, intercept = FALSE, lambda = lambda)
 adaptivemsfit <- adaptivecoopfit
 return(adaptivemsfit)
}

############################################################################################################
cvmslasso <- function(Xfcv, Yfcv, K, method, num.knots,w, lambda){
#Cross-validation function for ms-lasso and adaptive ms-lasso. Must be provided with a lambda sequence. 
  error <- matrix(NA, nrow = K, ncol = 100)
  print("cv")
  n = nrow(Xfcv)
  antf <- floor(n/K)
  
  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){
    print(k)
    test.ind <- which(fold == k)
    trainX <- Xfcv[-test.ind,]
    testX <- Xfcv[test.ind,]
    trainY <- Yfcv[-test.ind]
    testy <- Yfcv[test.ind]
    if(method == "monotone.lasso"){  
      cv.fit <- monotone.lasso(trainX, trainY, "gaussian", num.knots, lambda)
      Ztest <- monotone.splines(testX, num.knots)
    }else if(method == "adaptive.monotone"){
      cv.fit = adaptive.monotone(trainX, trainY, "gaussian", num.knots,w, lambda)
      Ztest <- monotone.splines(testX, num.knots)
    }  
    error.k <- (testy-predict(cv.fit, newx = Ztest))^2
    error[k,] <- colSums(error.k)
  }
  MSE <- colMeans(error)
  return(MSE)
}

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