loglik = function(X, y, beta, family) {
  K = dim(beta)[2]
  link = cbind(1, X) %*% beta
  yrep = repmat(y, 1, K)
  if (family == "gaussian")
    return(apply((yrep - link)^2, 2, sum))
  if (family == "poisson")
    return(apply(exp(link) - yrep * link, 2, sum))
  if (family == "binomial")
    return(apply(log(1 + exp(link)) - yrep * link, 2, sum))
}


repmat = function(X, m, n) {
  ## R equivalent of repmat (matlab)
  X = as.matrix(X)
  mx = dim(X)[1]
  nx = dim(X)[2]
  matrix(t(matrix(X, mx, nx * n)), mx * m, nx * n, byrow = T)
}


getdf = function(coef.beta) {
  apply(abs(coef.beta) > 1e-10, 2, sum)
}

margcoef <- function(x, y, condind = NULL, family, null.model = FALSE, iterind, method='SIS') {
  n = dim(x)[1]
  p = dim(x)[2]
  ones = rep(1, n)
  candind = setdiff(1:p, condind)
  if (iterind == 0) {
    if (family == "cox")
      margcoef = abs(cor(x, y[, 1])) else {
        margcoef = case_when(
          method=='SIS'~abs(cor(x, y)),
          # method=='RoSIS'~abs(mean(x*fy(y))),
          # method=='RRCS'~abs(cor(x, y,method = 'kendall')),
          TRUE~abs(cor(x, y))
        )
      }
  } else {
      if (null.model == TRUE) {
        if (is.null(condind) == TRUE) {
          x = x[sample(1:n), ]
        }
        if (is.null(condind) == FALSE) {
          x[, candind] = x[sample(1:n), candind]
        }
      }
      margcoef = abs(sapply(candind, mg, x, y, ones, family, condind,method))
    }


  return(margcoef)
}

mg <- function(index, x = x, y = y, ones = ones, family = family, condind = condind, method=method) {
  if(method=='SIS'){
    margfit = switch(family, gaussian = coef(glm.fit(cbind(ones, x[, index], x[, condind]), y, family = gaussian()))[2],
                     binomial = coef(glm.fit(cbind(ones, x[, index], x[, condind]), y, family = binomial()))[2],
                     poisson = coef(glm.fit(cbind(ones, x[, index], x[, condind]), y, family = poisson()))[2],
                     probit = coef(glm.fit(cbind(ones, x[, index], x[, condind]), y, family = binomial("probit")) )[2],
                     negbin = coef(glm.nb(y~cbind(ones, x[, index], x[, condind]), init.theta = stop))[2],
                     cox = coef(coxph(y ~ cbind(x[, index], x[, condind])))[1],
                     coxp = coef(coxph(y ~ cbind(x[, index], x[, condind])))[1]/
                            sqrt(vcov(coxph(y ~ cbind(x[, index], x[, condind])))[1,1])
    )
  }else{
    fy <- with(density(y), splinefun(x, cumsum(y) / sum(y), method="monoH.FC"))
    margfit = switch(method, RoSIS=mean(x[,c(index,condind)]*fy(y)),
                     RRCS = cor(x[,index], y,method = 'kendall'))
  }
}


## partially penalized function for SIS
#' @importFrom ncvreg cv.ncvsurv ncvsurv
tune.fit.adj <- function(x, y, family = c("gaussian", "binomial", "poisson", "cox"), penalty = c("SCAD", "MCP", "lasso"),
                         concavity.parameter = switch(penalty, SCAD = 3.7, 3),
                         tune = c("cv", "aic", "bic", "ebic"), nfolds = 10,
                         lambda = NULL, nlambda = 100,
                     type.measure = c("deviance", "class", "auc", "mse", "mae"), gamma.ebic = 0.5,
                     penalty.factor = rep(1, ncol(x))
                    ) {

  if (is.null(x) || is.null(y))
    stop("The data is missing!")

  x = as.matrix(x)

  this.call = match.call()
  family = match.arg(family)
  penalty = match.arg(penalty)
  # if (class(concavity.parameter) != "numeric")
  #   stop("concavity.parameter must be numeric!")
  # tune = match.arg(tune)
  # if (class(nfolds) != "numeric")
  #   stop("nfolds must be numeric!")
  type.measure = match.arg(type.measure)
  nobs = nrow(x)
  nvars = ncol(x)

  if (tune == "cv") {
    if (penalty == "lasso" ) {
      cv.fit = cv.glmnet(x, y, family = family, type.measure = type.measure, nfolds = nfolds)
      coef.beta = coef(cv.fit, s = "lambda.1se")
      reg.fit = cv.fit$glmnet.fit
      lambda = cv.fit$lambda.1se
      lambda.ind = which(cv.fit$lambda == cv.fit$lambda.1se)
    } else if (family != 'cox') {
      cv.fit = cv.ncvreg(x, y, family = family, penalty = penalty, gamma = concavity.parameter, nfolds = nfolds)
      cv.1se.ind = min(which(cv.fit$cve<cv.fit$cve[ cv.fit$min]+cv.fit$cvse[ cv.fit$min]))
      coef.beta = cv.fit$fit$beta[, cv.1se.ind]  # extract coefficients at a single value of lambda, including the intercept
      reg.fit = cv.fit$fit

      lambda = cv.fit$lambda[cv.1se.ind]
      lambda.ind = cv.1se.ind
    } else {
      cv.fit = cv.ncvsurv(x, y, family = family, penalty = penalty, gamma = concavity.parameter, nfolds = nfolds)
      cv.1se.ind = min(which(cv.fit$cve<cv.fit$cve[ cv.fit$min]+cv.fit$cvse[ cv.fit$min]))
      coef.beta = cv.fit$fit$beta[, cv.1se.ind]  # extract coefficients at a single value of lambda
      reg.fit = cv.fit$fit

      lambda = cv.fit$lambda[cv.1se.ind]
      lambda.ind = cv.1se.ind
    }
  } else {
    n = nrow(x)
    if (penalty == "lasso" ) {
      if(is.null(lambda)){
        reg.fit = glmnet(x, y, family = family, nlambda = nlambda,
                         penalty.factor = penalty.factor)
      }else{
        reg.fit = glmnet(x, y, family = family, lambda = lambda, nlambda = nlambda,
                         penalty.factor = penalty.factor)
      }
      coef.beta = rbind(reg.fit$a0,as.matrix(reg.fit$beta))  # extract coefficients at all values of lambda,  including the intercept
      dev = deviance(reg.fit)
      reg.df = reg.fit$df
    } else {
      if(family != 'cox'){
        if(is.null(lambda)){
        reg.fit = ncvreg(x, y, family = family, penalty = penalty, gamma = concavity.parameter,
                        nlambda = nlambda,
                         penalty.factor = penalty.factor)
        }else{
          reg.fit = ncvreg(x, y, family = family, penalty = penalty, gamma = concavity.parameter,
                           lambda=lambda, nlambda = nlambda,
                           penalty.factor = penalty.factor)
        }
        coef.beta = reg.fit$beta  # extract coefficients at all values of lambda, including the intercept
        dev = loglik(x, y, coef.beta, family = family)
        reg.df = getdf(coef.beta[-1, , drop = FALSE])
      } else {
        if(is.null(lambda)){
          reg.fit = ncvsurv(x, y, family = family, penalty = penalty,
                            gamma = concavity.parameter, nlambda = nlambda,
                            penalty.factor = penalty.factor)
        }else{
          reg.fit = ncvsurv(x, y, family = family, penalty = penalty,
                            gamma = concavity.parameter, lambda = lambda,
                            penalty.factor = penalty.factor)
        }
        coef.beta = reg.fit$beta  # extract coefficients at all values of lambda, including the intercept
        dev = 2*reg.fit$loss
        reg.df = getdf(coef.beta)
      }
    }

    if (tune == "aic") {
      obj = dev + 2 * reg.df
    }
    if (tune == "bic") {
      obj = dev + log(n) * reg.df
    }
    if (tune == "ebic") {
      obj = dev + log(n) * reg.df + 2 * gamma.ebic * log(choose(dim(x)[2], reg.df))
    }
    lambda.ind = which.min(obj)
    coef.beta = coef.beta[, lambda.ind]
    lambda = reg.fit$lambda[lambda.ind]
    obj = obj
  }

  if(family != 'cox'){
    a0 = coef.beta[1]
    coef.beta = coef.beta[-1]
  } else{
    a0 = NULL
    coef.beta = as.vector(coef.beta)
  }
  ix = which(coef.beta != 0)
  beta = coef.beta[ix]
  return(list(ix = ix, a0 = a0, beta = beta, fit = reg.fit, lambda = lambda, lambda.ind = lambda.ind, obj = obj, lambda.seq = reg.fit$lambda))
}



## ---------------------------------------------------------------------------
## Sure Independence Screening for the Cox proportional hazards model
## ---------------------------------------------------------------------------

## One screening step for the Cox model.
##
## Ranks the candidate features by the magnitude of their (conditional)
## marginal Cox utility and, for the iterative version, selects the subset that
## minimises the (extended) BIC of the resulting partial-likelihood fit.
##
## Arguments
##   x0          : the n x p feature matrix.
##   y0          : a 'Surv' object with the observed times and event indicators.
##   nsis        : maximum number of new features to consider in this step.
##   family      : marginal utility used for ranking, 'cox' (marginal
##                 coefficient) or 'coxp' (marginal Wald z-statistic).
##   candind     : indices of the candidate features (typically setdiff(1:p, ix0)).
##   BIC_old     : the current best (extended) BIC to improve upon.
##   ix0         : indices of the features already retained by previous steps.
##   iterind     : the current screening-iteration index (0 uses a correlation
##                 screen, >= 1 uses conditional marginal Cox fits).
##   sis_method  : 'isis' for the BIC-driven iterative selection, 'sis' for a
##                 single threshold-based pass on the marginal statistic.
##   gamman      : threshold on the marginal statistic used when sis_method='sis'.
#' @importFrom survival coxph Surv
#' @importFrom stats BIC sd
SIS_cox <- function(x0, y0, nsis, family, candind, BIC_old, ix0,
                    iterind, sis_method = 'isis', gamman = NA) {
  new.ind = 0
  p = ncol(x0)

  ## (conditional) marginal utility of every candidate feature
  mcoef = margcoef(x0, y0, ix0, family = family, null.model = FALSE,
                   iterind = iterind, method = 'SIS')
  rankcoef = sort(mcoef, decreasing = TRUE, index.return = TRUE)

  if (sis_method == 'isis') {

    nsis = min(nsis, length(candind))          ## never look past the candidates
    nseq = seq(1, nsis, by = 1)
    bic = rep(0, length(nseq) + 1)
    var_id0 = c()
    bic[1] = BIC_old

    for (sk in 1:length(nseq)) { ## grow the model and track the extended BIC
      newix = candind[rankcoef$ix[1:nseq[sk]]]
      xf = as.matrix(x0[, c(ix0, newix)])
      fit0 <- coxph(y0 ~ xf)
      eBIC = BIC(fit0) + log(choose(p, sk)) # extended BIC with gamma = 0.5
      bic[sk + 1] = eBIC
      if (BIC(fit0) < BIC_old) {
        new.ind = 1
        var_id0_min = newix
        BIC_old = BIC(fit0)
      }
    }

    if (new.ind == 1) {
      sk = min(which(bic < min(bic) + sd(bic) / sqrt(length(nseq))))
      if (sk == 1) var_id0 = c() else {
        var_id0 = sort(candind[rankcoef$ix[1:(sk - 1)]])
      }
      BIC_min = bic[sk]
    } else {
      var_id0 = c()
      BIC_min = BIC_old
    }

  } else { ## single threshold-based screening pass
    var_id0 = sort(candind[rankcoef$ix[rankcoef$x >= gamman]])
    if (length(var_id0) == 0) {
      BIC_min = BIC_old
    } else {
      xf = as.matrix(x0[, c(ix0, var_id0)])
      fit0 <- coxph(y0 ~ xf)
      BIC_min = BIC(fit0)
    }
  }

  return(list(BIC = BIC_min, ix0 = var_id0))
}


## (Iterative) Sure Independence Screening for the Cox model.
##
## Wraps 'SIS_cox' into the (iterative) screening loop used for dimension
## reduction before the partially penalised inference step.  Setting
## iter = FALSE performs a single BIC-driven SIS pass; iter = TRUE performs the
## full ISIS procedure.
##
## Arguments
##   X         : the n x p feature matrix.
##   y         : a 'Surv' object with the observed times and event indicators.
##   iter      : logical; whether to run the iterative (ISIS) procedure.
##   iter.max  : maximum number of screening iterations.
##   trace     : logical; whether to print the selection at each iteration.
##
## Value: a sorted integer vector of the retained feature indices.
#' @importFrom survival Surv
screen_cox <- function(X, y, iter = TRUE, iter.max = 20, trace = FALSE) {
  n = nrow(X)
  p = ncol(X)
  ix0 = c()
  iterind = 0
  BIC_old = 1e4

  repeat {
    iterind = iterind + 1
    if (length(setdiff(1:p, ix0)) == 0) break          ## no candidates left
    nsis = max(ceiling(n / log(n)) - length(ix0), 2)
    SIS_i = SIS_cox(x0 = X, y0 = y, nsis = nsis, family = 'cox',
                    candind = setdiff(1:p, ix0), BIC_old = BIC_old,
                    ix0 = ix0, iterind = iterind, sis_method = 'isis')
    ix1 = SIS_i$ix0
    BIC_old = SIS_i$BIC

    if (trace) cat("Iter", iterind, ", selection: ", c(ix0, ix1), "\n")

    ix0 = c(ix0, ix1)

    if (!iter) break                                   ## single SIS pass
    if (length(ix1) >= nsis || iterind >= iter.max) {  ## stop growing
      if (trace && length(ix1) >= nsis)
        cat("Maximum number of variables selected \n")
      if (trace && iterind >= iter.max)
        cat("Maximum number of iterations reached \n")
      break
    }
    if (iterind > 1 && length(ix1) == 0) {             ## model stabilised
      if (trace) cat("Model already selected \n")
      break
    }
  }

  return(sort(ix0))
}




