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

require(gurobi); 

######################################################################
# min sum(rho_{M}(y - x * b)) + lambda * |D * b|_1
# 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
#    M      --- scalar, tuning parameter in Huber loss, default = 1.35 
#    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 
#    ...    --- 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 
######################################################################

huber.genlasso.fit.gurobi = function(
    x, y, lambda = 0.0, 
    Cmat = NULL, dvec = NULL, Emat = NULL, fvec = NULL, 
    Dmat = diag(rep(1.0, NCOL(x))), M = 1.35, ...)
{
    stopifnot(NROW(x) == length(y)); 
    stopifnot(lambda >= 0.0, M >= 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 * (bv)^T * (X^TX X^T -X^T; X I -I; -X; -I; I) * bv - bv^T * (X^Ty, y) 
    #                  + M * (v^+ + v^-) + lambda * sum(z+ + z-)
    # arguments: 
    #   b   --  length p
    #   v+  --  length n,   v+ >= 0
    #   v-  --  length n,   v- >= 0
    #   z+  --  length nD,  z+ >= 0
    #   z-  --  length nD,  z- >= 0
    ######################################################################

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

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

    model$A = cbind(Dmat, zeros(m, n + n),  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, n + n + 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, n + n + 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];
    v = QPfit$x[p + (1 : n)] - QPfit$x[(p + n) + (1 : n)];
    z = QPfit$x[(p + n + n) + (1 : m)] - QPfit$x[(p + n + n + m) + (1 : m)];

    list(x = x, y = y, lambda = lambda, M = M,
         Cmat = Cmat, dvec = dvec, Emat = Emat, fvec = fvec, Dmat = Dmat,
         b = b, v = v, 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
######################################################################