problem_setup.jl

using FrankWolfe
using LinearAlgebra
using Random
#The following are only used if we want to doublecheck f_opt with
#CVX.
using Gurobi
using JuMP
using Convex
using SCS


#This script is used to generate a problem instance for running
#CGALM.

function generate_psd_matrix(n,diag_range)
    #This is a subfunction of generate_cq_constraint.
    #
    #Create orthogonal matrix first by (a) Generating a matrix with
    #uniform random entries in [0,1], then (b) grabbing the "Q"
    #from QR decomposition.
    U = qr(rand(n,n)).Q
    # Prescribe some eigenvalues to have values within 'diag_range'
    e = zeros(n)
    for i in range(1,n)
        e[i] = Float64(rand(diag_range))
    end
    #Return matrix with desired spectral decomposition
    A = U*Diagonal(e)*(U')
    #Symmetrize just for  good measure:
    A = 0.5*(A + A')
    return A
end

function generate_cq_constraint(n, max_trials =(5,20))
    #Generates variables for the Convex Quadratic inequality
    #constraint,
    #   x'Qx + <r, x> + d <= 0
    #where:
    #   Q (n by n) PSD matrix
    #   r (n by n matrix)
    #   d (real scalar)
    # such that:
    # (A) the barycenter of the Birkhoff polytope is on the
    # interior of the inequality constraint
    # and 
    # (B) At least one Birkhoff vertex is excluded.
    #
    #This is done by randomly generating r matrices and d values
    #that try to simultaneously satisyfy (A) and (B). Based on some
    #testing, this appears to be highly likely since ZW has never
    #had to increase the number of trials in order to meet the
    #requirements.
    #
    #Begin by generating a quadratic matrix of the appropriate size
    #Q = generate_psd_matrix(n,1:10)   
    Q = generate_psd_matrix(n,0.1:0.1:1)   
    #Barycenter of n-by-n Birkhoff polytope is matrix of all 1/n's
    c_barycenter = ones(n,n)./n
    #Partial evaluation of quadratic part of constraint at
    #barycenter:
    QC = dot(vec(c_barycenter),vec(Q*c_barycenter))
    #AB_flag will be only set to 'true' if both properties (A)
    #and (B) outlined above are satisfied. 
    AB_flag = false
    r_trial = 1
    #Generate candidate matrix r to potentially be used in
    #constraint as described above.
    r = rand(n,n)
    #Initialize constant term d for constraint as described above
    #(this will be changed).
    d = 0
    #Construct a random Birkhoff vertex:
    P = I(n)[randperm(n), :]
    while (AB_flag == false &&  r_trial <= max_trials[1])
        r = rand(n,n)
        #Selecting the constant term d to be strictly greater than
        #d_min value ensures that the barycenter is on the interior.
        d_min = - ( QC + dot(vec(r),vec(c_barycenter)) )
        #This ensures that (A) will always be satisfied:
        d = d_min - 1.0
        #To satisfy part (B), generate a random permutation matrix
        #(= random vertex of the Birkhoff polytope)
        P = I(n)[randperm(n), :]
        #Evaluate the CQ inequality at P:
        test_val =dot(vec(P),vec(Q*P)) + dot(vec(r),vec(P)) + d
        #Initialize trial count
        perm_trial = 1
        #Only enter loop if P satisfies constraint. Exit if P
        #violates constraint, or maximum trials occurs.
        while (perm_trial <= max_trials[2] && test_val <= 0 )
            #If this permutation doesn't work,  try a different one
            #until it works
            P = I(n)[randperm(n), :]
            test_val = dot(vec(P),vec(Q*P)) + dot(vec(r),vec(P)) + d
            perm_trial += 1
        end
        #If we verified at least one partition is excluded, set the
        #flag to true to exit the loop.
        if test_val > 0
            AB_flag = true
        end
        #If all permutations checked happened to be feasible for
        #this quadratic, try re-generating a new r matrix.
        r_trial += 1
    end
    #Only return constants for the quadratic constraint if (A) and
    #(B) above were verified to be true.
    if AB_flag == false
        return false
    else
        return Q, r, d, P, c_barycenter
    end
end

function generate_problem_instance(n; params = false, return_matrices = false)
    #Generates matrices and vectors such that the following opt
    #problem:
    #
    #Minimize (x - b)'A(x - b)
    #s.t.
    #   x\in Birkhoff polytope \\
    #   x'Q_1x + <r_1, x> + d_1 <= 0 \\
    #   x'Q_2x + <r_2, x> + d_2 <= 0 \\
    #
    #has the following properties:
    #   1) Global minimizer of objective function is not centered
    #   within the Birkhoff polytope
    #   2) Each inequality constraint is not redundant: at least one
    #   vertex of the Birkhoff polytope is excluded by each
    #   inequality constraint.
    #   3) The problem is feasible: ensure that the barycenter of
    #   the Birkhoff polytope is strictly feasible for all
    #   inequality constraints.
    #
    #This problem has an n-by-n matrix variable x
    #The following calls generate a convex quadratic kernel (Q),
    #matrix variable (r), scalar (d), an excluded vertex (P), and
    #the barycenter of the Birkhoff polytope (c_bary)
    #
    #If the optional keyword argument 'params' is true, then
    #this will also generate the relevant constants required to run
    #CGALM within short-step mode and CoexDurCG.
    #
    num_constraints = 2
    Q1, r1, d1, P1, c_bary = generate_cq_constraint(n)
    Q2, r2, d2, P2, _ = generate_cq_constraint(n)

    #Constraint functions:
    #if grad = true, returns gradient of g at x.
    #Otherwise, by default, returns primal evaluation g(x)
    function g_1(x; grad = false)
        if grad
            return 2*(Q1*x) + r1 
        else
            #Return primal value
            return dot(vec(x),vec(Q1*x)) + dot(vec(r1),vec(x)) + d1
        end
    end
    function g_2(x; grad = false)
        if grad
            return 2*(Q2*x) + r2 
        else
            #Return primal value
            return dot(vec(x),vec(Q2*x)) + dot(vec(r2),vec(x)) + d2
        end
    end
    #Doublecheck assertions to make claimed permutations are
    #excluded, and barycenter is strictly feasible for both
    #constraints.
    @assert g_1(P1) > 0
    @assert g_2(P2) > 0
    @assert g_1(c_bary) < 0
    @assert g_2(c_bary) < 0
    if P1 == P2
        println("Unlucky: both inequality constraints exclude the same vertex")
    end

    #Generate objective function (eigenvalue range is strictly
    #positive, so there is no nontrivial nullspace of A):
    A = generate_psd_matrix(n,1:10)
    #Scale b so far that it won't be in the Birkhoff polytope;
    #scaled to be in the direction of an excluded
    #vertex, so the problem becomes geometrically interesting.
    b = c_bary + 10*(P1 - c_bary)

    #Construct function; this one is different from constraints: if
    #compute_grad is true, will return both primal and gradient
    function f(x, compute_grad = false)
        temp = x - b
        Amult = A*temp
        primal = 0.5*dot(vec(temp),vec(Amult))
        if compute_grad
            return primal, Amult
        else
            return primal
        end
    end

    function f_grad!(x, storage)
        #Used for only computing the gradient of f (excludes the
        #primal computation). This directly modifies the memory in
        #"storage": writes gradient of f at x.
        storage .= A*(x - b) 
    end

    #g is the vector-valued array of constraint functions:
    g = Array{Function}(undef, num_constraints);
    g[1], g[2] = g_1, g_2
    #Frank-Wolfe constraint for this problem will be the Birkhoff
    #Polytope. FrankWolfe.jl relies on the Hungarian algorithm to
    #implement this LMO.
    lmo = FrankWolfe.TrackingLMO(FrankWolfe.BirkhoffPolytopeLMO())
    if return_matrices
        matrix_data = ( A = A,
                            b = b,
                            Q1 = Q1, r1 = r1, d1 = d1,
                            Q2 = Q2, r2 = r2, d2 = d2,
                            P1 = P1, P2 = P2,
                            c_bary = c_bary,
                            n = n
                        )
    end
    if params
        #L_g should be a vector of Lipschitz constants of nabla g_i
        L_g = [2*opnorm(Q1), 2*opnorm(Q2)]
        #L_f should be Lipschitz constant of nabla f.
        L_f = opnorm(A)
        #B should be a vector of upper bounds on ||nabla g_i(x)||,
        #for all x in domain of h.
        #Recall nabla g[i](x) = 2Qx + r,
        #and for this problem, x is always a doubly-stochastic
        #matrix, so ||x||_F <= sqrt(n). Hence
        #B[i] = 2sqrt(n)||Q|| + ||r||
        #is a valid upper bound. Reusing the fact that we just
        #computed 2norm(Q) above:
        B = [sqrt(n)*L_g[1] + norm(r1), sqrt(n)*L_g[2] + norm(r2)]
        if return_matrices
             return f, f_grad!, g, lmo, c_bary, Prob_Params(L_g,L_f,B,2,sqrt(2*n)), matrix_data
        else
            return f, f_grad!, g, lmo, c_bary, Prob_Params(L_g,L_f,B,2,sqrt(2*n))
        end
    else
        #If running in open-loop, no need to report constants L_g
        #or B.
        if return_matrices
            return f, f_grad!, g, lmo, c_bary, matrix_data 
        else
            return f, f_grad!, g, lmo, c_bary
        end
    end
end


function test_matrix_psd(A, trials)
    #Only used for testing
    n = size(A)[1]
    x = rand(n)
    for i = 1:trials
        temp = dot(vec(x),vec(A*x))
        if temp < -1e-14
            println("Error: non-PSD detected")
            return false
        elseif i==trials
            return true
        end
    end
end

#Temporarily commented this due to errors
#
function solve_problem_with_Gurobi(matrix_data, silent=true)
   
    A, b = matrix_data.A, matrix_data.b
    Q1, r1, d1 = matrix_data.Q1, matrix_data.r1, matrix_data.d1
    Q2, r2, d2 = matrix_data.Q2, matrix_data.r2, matrix_data.d2

    begin
  
        model = Model(Gurobi.Optimizer)
        @variable(model, X[1:n,1:n] >= 0)
        @constraint(model, dot( vec(X), vec(Q1*X) ) + dot( vec(r1),  vec(X) ) + d1 <= 0 )
        @constraint(model, dot( vec(X), vec(Q2*X) ) + dot( vec(r2),  vec(X) ) + d2 <= 0 )
        for i in 1:n
            @constraint(model, sum(X[i,:]) == 1 )
            @constraint(model, sum(X[:,i]) == 1 )
        end
         @objective( model, Min , 0.5*dot(vec(X-b), vec(A*(X-b))) ) 

    end

    optimize!(model)
    termination_status(model)
 
    return (
        status = termination_status(model),
        optval = objective_value(model),
        Xstar = value.(X),
        elapsed_time = solve_time(model),
        data = matrix_data,
        model = model
    )

end


function solve_problem_with_cvx(matrix_data; silent=true, tol=1e-6,
        eps_abs=tol, eps_rel=(tol^2), max_iters=300000)
    n = matrix_data.n
    A, b = matrix_data.A, matrix_data.b
    Q1, r1, d1 = matrix_data.Q1, matrix_data.r1, matrix_data.d1
    Q2, r2, d2 = matrix_data.Q2, matrix_data.r2, matrix_data.d2


    if (!silent)
        println("Constructing problem for CVX...")
    end
    # Convex variable: n-by-n matrix
    X = Variable(n, n)

    constraints = Any[]

    # Birkhoff polytope constraints
    push!(constraints, X >= 0)

    for i in 1:n
        push!(constraints, sum(X[i, :]) == 1)   # row sums
        push!(constraints, sum(X[:, i]) == 1)   # column sums
    end

    # Objective:
    # 0.5 * sum_j (X[:,j] - b[:,j])' A (X[:,j] - b[:,j])
    objective = 0.5 * sum(quadform(X[:, j] - b[:, j], A) for j in 1:n)

    # Quadratic inequality constraints
    g1_expr = sum(quadform(X[:, j], Q1) for j in 1:n) + sum(r1 .* X) + d1
    g2_expr = sum(quadform(X[:, j], Q2) for j in 1:n) + sum(r2 .* X) + d2

    push!(constraints, g1_expr <= 0)
    push!(constraints, g2_expr <= 0)

    problem = minimize(objective, constraints)
    if (!silent)
        println("Solving problem with CVX...")
    end



    solver = Convex.MOI.OptimizerWithAttributes(
        SCS.Optimizer,
        "eps_abs"   => eps_abs,
        "eps_rel"   => eps_rel,
        "max_iters" => max_iters,
        "verbose"   => (!silent),
    )
    start = time()
    solve!(problem, solver; 
           silent=silent)
    elapsed_time = time() - start

    return (
        status = problem.status,
        optval = problem.optval,
        Xstar = evaluate(X),
        elapsed_time = elapsed_time,
        data = matrix_data,
        problem = problem,
    )
end

function check_feas(x,g, tol=1e-6)
    n, _ = size(x)
    feas = true
    g_vals = [g[i](x) for i in eachindex(g)]
    for i in eachindex(g)
        feas = feas && ( g_vals[i] <= tol)
    end
    poly_val = 0
    for i in range(1,n)
        poly_val = max(poly_val, abs(sum(x[i, :]) - 1) )
        feas = feas && ( poly_val <= tol)
        poly_val = max(poly_val, abs(sum(x[:, i]) - 1) )
        feas = feas && ( poly_val <= tol )
    end
    return feas, g_vals, poly_val
end

#Another version of check_feas that adds all infeasibilities
#into one scalar:
#function check_feas(x,g)
#    n, _ = size(x)
#    feas = 0;
#    for i in eachindex(g)
#         feas = feas + max(g[i](x), feas)
#     end
#    for i in range(1,n)
#        feas = feas + abs(sum(x[i, :])-1)
#        feas = feas + abs(sum(x[:, i])-1)
#    end
#    return feas
#end
#

#############################################
##Uncomment for testing to ensure we can evaluate functions from generate_problem_instance(n) correctly.
#g_primals_test = zeros(m)
#for i = 1:length(g)
#    #Test computing primal and gradient for all constraint
#    #functions
#    g_primals_test[i] = g[i](x)
#    g[i](x; grad=true)
#end
#f(x)
#f_grad!(x,st)
#F_grad!(x, z, 1.0, f_grad!, g, g_primals_test, st)
#############################################