######################################################################
# Least squares regression with generalized lasso penalty and linear constraints
# Peng Zeng @ Auburn University 
# 11-17-2021
######################################################################

require(gurobi); 

######################################################################
# min 0.5 * sum((y - x * b)^2) + lambda * sum(abs(D * b)) + ridge * sum(b^2)
# subject to C * b >= d, E * b = f
# solve it using Quadratic Programming (Gurobi)
# 
# Input:
#    x      --- n-by-p matrix, predictors 
#    y      --- n-by-1 vector, response
#    lambda --- scalar, nonnegative, tuning parameter in the penalty
#    Cmat   --- q-by-p matrix, inequality constraints  
#    dvec   --- q-by-1 vector, inequality constraints 
#    Emat   --- s-by-p matrix, equality constraints  
#    fvec   --- s-by-1 vector, equality constraints 
#    Dmat   --- m-by-p matrix, default: identity matrix 
#    ridge  --- scalar, ridge parameter 
#    ...    --- more parameters for gurobi
# 
# Output: 
#    b      --- p-by-1 vector, estimated coefficients 
#    v      --- n-by-1 vector   
#    z      --- m-by-1 vector, z = D %*% b 
#    Qvalue --- scalar, value of the objective function 
#    model  --- list, the QP problem to be solved by Gurobi 
#    gurobi --- list, results of QP problem by Gurobi 
######################################################################

genlasso.fit.gurobi = function(x, y, lambda = 0.0, 
    Cmat = NULL, dvec = NULL, Emat = NULL, fvec = NULL, 
    Dmat = diag(rep(1.0, NCOL(x))), ridge = 0.0, ...)
{
    stopifnot(NROW(x) == length(y)); 
    stopifnot(lambda >= 0.0, ridge >= 0.0); 
    stopifnot(NCOL(Dmat) == NCOL(x)); 

    n = NROW(x); p = NCOL(x); m = NROW(Dmat); 

    ######################################################################
    # prepare for Quadratic Programming 
    ######################################################################

    model = list(modelsense = 'min'); 

    ######################################################################
    # min 0.5 * b^T * X^TX * b - b^T * X^Ty + lambda * sum(z+ + z-) + ridge * sum(b^2)
    # arguments: 
    #   b   --  length p
    #   z+  --  length nD,  z+ >= 0
    #   z-  --  length nD,  z- >= 0
    ######################################################################

    model$Q = zeros(p + m + m, p + m + m);
    model$Q[1:p, 1:p] = 0.5 * t(x) %*% x + diag(rep(ridge, p));
    model$obj = c(-t(x) %*% y, rep(lambda, m + m));
    model$objcon = 0.5 * sum(y * y);
    model$lb = c(rep(-Inf, p), rep(0.0, m + m));
    model$ub = rep(Inf, p + m + m);

    ######################################################################
    # constraints: 1 to nD 
    #    D * b - z+ + z- = 0 or D * b = z
    ######################################################################

    model$A = cbind(Dmat, eye(m, -1.0), eye(m, 1.0));
    model$rhs = rep(0.0, m);
    model$sense = rep('=', m);

    ######################################################################
    # constraints: nD + (1 to nC) 
    #    C * b >= d
    ######################################################################

    if(!is.null(Cmat))
    {
        if(is.null(dvec)) dvec = rep(0.0, NROW(Cmat));
        stopifnot(NROW(Cmat) == length(dvec), NCOL(Cmat) == NCOL(x));
        q = NROW(Cmat);
    
        # update A, rhs, sense
        model$A = rbind(model$A, cbind(Cmat, zeros(q, m + m)));
        model$rhs = c(model$rhs, dvec);
        model$sense = c(model$sense, rep('>=', q));
    }

    ######################################################################
    # constraints: nD + nC + (1 to nD) 
    #    E * b >= f
    ######################################################################

    if(!is.null(Emat))
    {
        if(is.null(fvec)) fvec = rep(0.0, NROW(Emat));
        stopifnot(NROW(Emat) == length(fvec), NCOL(Emat) == NCOL(x));
        s = NROW(Emat);
    
        # update A, rhs, sense
        model$A = rbind(model$A, cbind(Emat, zeros(s, m + m)));
        model$rhs = c(model$rhs, fvec);
        model$sense = c(model$sense, rep('=', s));
    }

    ######################################################################
    # call gurobi() to solve the Quadratic Programming problem 
    ######################################################################

    params = list(OutputFlag = 0, ...);
    QPfit = gurobi(model, params);
    if(QPfit$status != 'OPTIMAL') 
        warning('\n\nThe status of gurobi(model, params) is ', QPfit$status, 
                '\n===== Something is wrong! =====\n');
  
    ######################################################################
    # get the solution from gurobi() 
    #   z = z+ - z- = D * b
    ######################################################################
    
    b = QPfit$x[1:p];
    z = QPfit$x[(p+1):(p+m)] - QPfit$x[(p+m+1):(p+m+m)];

    list(x = x, y = y, lambda = lambda, ridge = ridge,  
         Cmat = Cmat, dvec = dvec, Emat = Emat, fvec = fvec, Dmat = Dmat,
         b = b, z = z, Qvalue = QPfit$objval, 
         model = model, gurobi = QPfit);
}

######################################################################
# n-by-p matrix of zeros 
######################################################################

zeros = function(n, p)
{
    matrix(0.0, nrow = n, ncol = p); 
}

######################################################################
# p-by-p diagonal matrix with value as its diagonal elements 
######################################################################

eye = function(p, value = 1.0)
{
    if(p == 1) matrix(value, nrow = 1) else diag(rep(value, p)); 
}

######################################################################
# THE END
######################################################################