using LinearAlgebra, Statistics, StatsBase, Distributions, JuMP, DelimitedFiles
using Gurobi, Random, Ipopt, Mosek, MosekTools, Convex, SCS, CSV, DataFrames

function calculate_analytic_center_diet(B,η,ρ,q,C,d)
    K, L = size(B)
    m1 = Model(Ipopt.Optimizer)
    @variable(m1, r[1:K]>=0)
    @variable(m1, t[1:K])
    @variable(m1, v[1:K])

    @constraint(m1, C*r .<= d)
    @constraint(m1, [l in 2:L], B[:,l]'*r >= η[l])
    @constraint(m1, sum(t[k] for k in 1:K) == (1/ρ)*(B[:,1]'*r - η[1]))
    for k in 1:K
        @NLconstraint(m1, v[k] == t[k]*((1/ρ)*(sum(B[j,1]*r[j] for j in 1:K) - η[1]))^(q-1) - r[k]^q)
        @constraint(m1, v[k] >= 0)
    end
    @NLobjective(m1, Max, sum(log(d[i] - sum(C[i,j]*r[j] for j in 1:K)) for i in 1:size(C)[1]) +
                          sum(log(sum(B[k,l]*r[k] for k in 1:K) - η[l]) for l in 2:L) +
                          sum(log(v[k]) for k in 1:K) +
                          sum(log(r[k]) for k in 1:K ))
    optimize!(m1)
    return JuMP.value.(r)
end

function generate_hyperplane_mip_diet(X,B,η,ρ,q,C,d,δ1,use_ac)
    M = 1e5
    n1, L1 = size(X)
    if mod(L1,2) == 0
        q = Int64(L1/2)
    else
        q = Int64((L1+1)/2)
    end
    m1 = Model(Gurobi.Optimizer)
    @variable(m1, f[1:n1])
    @variable(m1, l)
    @variable(m1, z[1:L1], Bin)
    @constraint(m1, sum(f[i] for i in 1:n1) >= δ1)
    if use_ac
        xac = calculate_analytic_center_diet(B,η,ρ,q,C,d)
        @constraint(m1, f'*xac == l)
    end
    @constraint(m1, [i in 1:L1], f'*X[:,i] <= l + M*(1-z[i]))
    @constraint(m1, [i in 1:L1], f'*X[:,i] >= l - M*z[i])
    @constraint(m1, sum(z[i] for i in 1:L1) == q)
    @objective(m1, Min, f'*f)
    optimize!(m1)
    if termination_status(m1) == MOI.OPTIMAL
        return false, JuMP.value.(f), JuMP.value.(l)
    else
        return true, zeros(n), 0.0, zeros(n)
    end
end

function calculate_different_vectors(X,x)
    Y = x
    for i in 1:length(X)
        if norm(X[i] .- x) > 1e-3
            Y = hcat(Y,X[i])
        end
    end
    return Y
end

function get_feasible_theta(q)
    m1 = Model(Gurobi.Optimizer)
    @variable(m1, θ[1:2,1:2]>=0.1)
    @constraint(m1, θ[1,1] + θ[2,1] == 1/q)
    @constraint(m1, θ[1,2] + θ[2,2] == 1 - 1/q)
    @constraint(m1, θ[1,1] + θ[1,2] == 1/q)
    @constraint(m1, θ[2,1] + θ[2,2] == 1 - 1/q)
    @objective(m1, Min, 0)
    optimize!(m1)
    return JuMP.value.(θ)
end


function is_feasible_diet(x,ρ,q,B,η,rob_nutr)
    K, L = size(B)
    rest_nutrients = [i for i in 1:L if i != rob_nutr]
    B_sub = B[:,rest_nutrients]
    η_sub = η[rest_nutrients]
    cnt_1 = (ρ*norm(x,q) <= B[:,rob_nutr]'*x - η[rob_nutr])*1
    cnt_2 = sum(B_sub'*x .>= η_sub)
    return cnt_1 + cnt_2 == L
end


function calculate_diet_ub_ipopt(c,ρ,q,B,η,rob_nutr,r0)
    K, L = size(B)
    rest_nutrients = [i for i in 1:L if i != rob_nutr]
    B_sub = B[:,rest_nutrients]
    η_sub = η[rest_nutrients]
    m1 = Model(Ipopt.Optimizer)
    @variable(m1, sqrt(10)>=r[1:K]>=0)
    @variable(m1, t[1:K])
    @variable(m1, v[1:K])
    for k in 1:K
        JuMP.set_start_value(r[k], r0[k])
    end
    @constraint(m1, B_sub'*r .>= η_sub)
    @constraint(m1, sum(t[k] for k in 1:K) == (1/ρ)*(B[:,rob_nutr]'*r - η[rob_nutr]))
    for k in 1:K
        @NLconstraint(m1, v[k] == t[k]*((1/ρ)*(sum(B[j,rob_nutr]*r[j] for j in 1:K) - η[rob_nutr]))^(q-1) - r[k]^q)
        @constraint(m1, v[k] >= 0)
    end
    @NLobjective(m1, Min, sum(c[k]*sqrt(r[k]) for k in 1:K))
    optimize!(m1)
    return objective_value(m1)
end


function calculate_candidate_vectors_diet(r,t,R,V,ρ,q,B,η,rob_nutr)
    S = [r]
    for j in 1:size(R)[2]
        s = R[:,j]/r[j]
        if sum(s .>= 0) == size(s)[1]
            if is_feasible_diet(s,ρ,q,B,η,rob_nutr)
                push!(S,s)
            end
        end
    end
    for j in 1:size(V)[2]
        s = V[:,j]/t[j]
        if sum(s .>= 0) == size(s)[1]
            if is_feasible_diet(s,ρ,q,B,η,rob_nutr)
                push!(S,s)
            end
        end
    end
    λ, U = eigen(R)
    for j in 1:size(U)[2]
        s = U[:,j]
        if sum(s .>= 0) == size(s)[1]
            if is_feasible_diet(s,ρ,q,B,η,rob_nutr)
                push!(S,s)
            end
        end
    end
    return S
end


# Read Data
Nutrients = ["Energy(kcal)", "Protein(g)", "Fat(g)", "Calcium(mg)", "Iron(mg)", "VitaminA(ug)", "ThiamineB1(mg)",
             "RiboflavinB2(mg)", "NicacinB3(mg)", "Folate(ug)", "VitaminC(mg)"]
η = [2100, 52.5, 89.25, 1100, 22, 500, 0.9, 1.4, 12, 160, 28]
# c = [2781.6, 1140.78, 4856.08, 8031.74, 14637.37, 3507.66, 1243.79,
#      1682.86, 1545.27, 1719.78, 1330.82, 3209.79, 1048.02]
# ind_keep = [1,2,3,4,5,7,9,10,13,14,20,21,23]
B_mat_df = CSV.read("/Users/thodoris/Downloads/WFP_B_Matrix.csv", DataFrame)
B = convert(Matrix, B_mat_df[:,2:12])
B[:,2] = replace.(B[:,2], r"," => ".")
B[:,3] = replace.(B[:,3], r"," => ".")
B[:,5] = replace.(B[:,5], r"," => ".")
B[:,7] = replace.(B[:,7], r"," => ".")
B[:,8] = replace.(B[:,8], r"," => ".")
B[:,9] = replace.(B[:,9], r"," => ".")
B[:,2] = map(x->parse(Float64,x),B[:,2])
B[:,3] = map(x->parse(Float64,x),B[:,3])
B[:,5] = map(x->parse(Float64,x),B[:,5])
B[:,7] = map(x->parse(Float64,x),B[:,7])
B[:,8] = map(x->parse(Float64,x),B[:,8])
B[:,9] = map(x->parse(Float64,x),B[:,9])
B = Float64.(B)
# B = B[ind_keep,:]
c_df = CSV.read("/Users/thodoris/Downloads/WFP_Costs.csv", DataFrame; header=false)
c = Float64.(c_df[:,3])
C, d = [], []

B[:,1] = B[:,1]/100
η[1] = η[1]/100

B[:,2] = B[:,2]/100
η[2] = η[2]/100

B[:,3] = B[:,3]/100
η[3] = η[3]/100

B[:,4] = B[:,4]/100
η[4] = η[4]/100

B[:,5] = B[:,5]/10
η[5] = η[5]/10

B[:,6] = B[:,6]/100
η[6] = η[6]/100

B[:,9] = B[:,9]/10
η[9] = η[9]/10

B[:,10] = B[:,10]/100
η[10] = η[10]/100

B[:,11] = B[:,11]/100
η[11] = η[11]/100

c = c/1000




function solve_rpt_sdp_diet_quad(A,b,h,c,B,η,ρ,θ,q,λ,C,d,use_lmi,use_new,rob_nutr)
    K, L = size(B)
    s1 = θ[1,1] + θ[1,2] + θ[2,1]
    s2 = θ[1,1]/s1 + θ[1,2]/s1
    s3 = (θ[1,1]/s1)/s2

    rest_nutrients = [i for i in 1:L if i != rob_nutr]

    m1 = Model(Mosek.Optimizer)
    @variable(m1, sqrt(1000)>=r[1:K]>=0)
    @variable(m1, t[1:K])

    @variable(m1, 1000>=R[1:K,1:K]>=0, Symmetric)
    @variable(m1, T[1:K,1:K], Symmetric)
    @variable(m1, 1000>=V[1:K,1:K])

    @constraint(m1, [k in 1:K], T[k,k] >= 0)

    @constraint(m1, [l in rest_nutrients], B[:,l]'*r >= η[l])
    @constraint(m1, sum(t[k] for k in 1:K) == (1/ρ)*(B[:,rob_nutr]'*r - η[rob_nutr]))
    @constraint(m1, [k in 1:K], [t[k], (1/ρ)*(B[:,rob_nutr]'*r - η[rob_nutr]), r[k]] in MOI.PowerCone(1/q))

    @constraint(m1, c'*r <= λ)
    @constraint(m1, [k in 1:K], c'*R[:,k] <= λ*r[k])
    @constraint(m1, [l in rest_nutrients], η[l]*c'*r + λ*B[:,l]'*r - c'*R*B[:,l] - λ*η[l] >= 0)
    @constraint(m1, λ^2 - 2*λ*c'*r + c'*R*c >= 0)
    @constraint(m1, [k in 1:K], [λ*t[k] - c'*V[:,k],
                                (1/ρ)*(λ*(B[:,rob_nutr]'*r - η[rob_nutr]) + η[rob_nutr]*c'*r - c'*R*B[:,rob_nutr]),
                                λ*r[k] - c'*R[:,k]] in MOI.PowerCone(1/q))

    @constraint(m1, [l in rest_nutrients, k in 1:K], B[:,l]'*R[:,k] >= η[l]*r[k])
    @constraint(m1, [l in rest_nutrients, k in 1:K], [B[:,l]'*V[:,k] - η[l]*t[k],
                                                      (1/ρ)*(B[:,rob_nutr]'*R*B[:,l] - η[l]*B[:,rob_nutr]'*r -
                                                       η[rob_nutr]*B[:,l]'*r + η[rob_nutr]*η[l]),
                                                       B[:,l]'*R[:,k] - η[l]*r[k]] in MOI.PowerCone(1/q))

    @constraint(m1, sum(V[:,k] for k in 1:K) .== (1/ρ)*(R*B[:,rob_nutr] .- η[rob_nutr]*r))
    @constraint(m1, sum(T[:,k] for k in 1:K) .== (1/ρ)*(V'*B[:,rob_nutr] .- η[rob_nutr]*t))

    @constraint(m1, [k in 1:K, j in 1:K], [V[j,k],
                                           (1/ρ)*(B[:,rob_nutr]'*R[:,j] - η[rob_nutr]*r[j]),
                                           R[j,k]] in MOI.PowerCone(1/q))

    if use_new
        @variable(m1, aux1[1:K,1:K])
        @variable(m1, aux2[1:K,1:K])
        for k in 1:K
            for j in 1:K
                @constraint(m1, [aux1[k,j], (1/ρ^2)*(B[:,rob_nutr]'*R*B[:,rob_nutr] - 2*η[rob_nutr]*B[:,rob_nutr]'*r + η[rob_nutr]^2),
                                 R[k,j]] in MOI.PowerCone(s1))
                @constraint(m1, [aux2[k,j], (1/ρ)*(B[:,rob_nutr]'*V[:,j] - η[rob_nutr]*t[j]), aux1[k,j]] in MOI.PowerCone(s2))
                @constraint(m1, [T[k,j], (1/ρ)*(B[:,rob_nutr]'*V[:,k] - η[rob_nutr]*t[k]), aux2[k,j]] in MOI.PowerCone(s3))
            end
        end
    end

    if length(C) > 0
        @constraint(m1, C*r .<= d)
        @constraint(m1, [k in 1:K], C*R[:,k] .<= r[k]*d)
        @constraint(m1, d*r'*C' .+ C*r*d' .<= d*d' .+ C*R*C')
        @constraint(m1, [l in 2:L], C*R*B[:,l] .- η[l]*C*r .<= d*(B[:,l]'*r - η[l]))
        @constraint(m1, [k in 1:K, j in 1:size(C)[1]], [d[j]*t[k] - C[j,:]'*V[:,k],
                                                       (1/ρ)*(d[j]*B[:,1]'*r - C[j,:]'*R*B[:,1] -
                                                        η[1]*d[j] + η[1]*C[j,:]'*r),
                                                        d[j]*r[k] - C[j,:]'*R[:,k]] in MOI.PowerCone(1/q))
    end

    if use_lmi
        @constraint(m1, [R V r; V' T t; r' t' 1] in PSDCone())
    end
    @objective(m1, Min, -tr(A*R) - b'*r - h)
    optimize!(m1)
    if termination_status(m1) == MOI.OPTIMAL || termination_status(m1) == MOI.SLOW_PROGRESS
        return true, JuMP.value.(r), JuMP.value.(t), JuMP.value.(R), JuMP.value.(V), objective_value(m1)
    else
        return false, zeros(K), zeros(K), zeros(K,K), zeros(K,K), 1e6
    end
end

function calculate_diet_ub_quad(X,A,b,h,ρ,q,B,η,c,λ,rob_nutr)
    best_ub = 1e6
    K, L = size(B)
    rest_nutrients = [i for i in 1:L if i != rob_nutr]
    if length(X) > size(c)[1]
        for j in 1:size(X)[2]
            r = X[:,j]
            # for i in 1:size(r)[1]
            #     if r[i] < 0
            #         r[i] = 0
            #     end
            # end
            if c'*r <= λ && sum(r .>= 0) == K
                if ρ*norm(r,q) <= B[:,rob_nutr]'*r - η[rob_nutr]
                    if sum(B[:,rest_nutrients]'*r .>= η[rest_nutrients]) == L - 1
                        cur_ub = -r'*A*r - b'*r - h
                        if cur_ub < best_ub
                            best_ub = cur_ub
                        end
                    end
                end
            end
            # cur_ub = calculate_diet_ub_ipopt(c,ρ,q,B,η,rob_nutr,r)
        end
    else
        r = X
        if c'*r <= λ && sum(r .>= 0) == K
            if ρ*norm(r,q) <= B[:,rob_nutr]'*r - η[rob_nutr]
                if sum(B[:,rest_nutrients]'*r .>= η[rest_nutrients]) == L - 1
                    cur_ub = -r'*A*r - b'*r - h
                    if cur_ub < best_ub
                        best_ub = cur_ub
                    end
                end
            end
        end
        # cur_ub = calculate_diet_ub_ipopt(c,ρ,q,B,η,rob_nutr,r)
    end
    return best_ub
end

function rpt_branch_and_split_diet_quad(A,b,h,c,B,η,ρ,θ,q,λ,δ,δ1,use_ac,use_lmi,use_new,rob_nutr)
    C, d = [], []
    gen_hyper = 0

    # Root Node
    res_root, rlb, tlb, Rlb, Vlb, lb = solve_rpt_sdp_diet_quad(A,b,h,c,B,η,ρ,θ,q,λ,C,d,use_lmi,use_new,rob_nutr)
    X = calculate_candidate_vectors_diet(rlb,tlb,Rlb,Vlb,ρ,q,B,η,rob_nutr)
    Y_split = calculate_different_vectors(X,rlb)
    ub_root = calculate_diet_ub_quad(Y_split,A,b,h,ρ,q,B,η,c,λ,rob_nutr)

    UB, LB, opt_sol, opt_val = ub_root, lb, rlb, ub_root
    # unvisited_nodes = Queue{Vector{}}()
    nodes_list = []
    ub_list, lb_list, hyp_list_f, hyp_list_l = [], [], [], []
    push!(ub_list, UB)
    push!(lb_list, LB)

    t0_1 = time_ns()
    total_time = 0.0

    while UB - LB > δ && total_time < 3600 && length(Y_split) > size(b)[1]
        # Split
        use_ac = true
        if gen_hyper < 1
            use_ac = false
        end
        res, fopt, lopt = generate_hyperplane_mip_diet(Y_split,B,η,ρ,q,C,d,δ1,use_ac)
        if res
            L1 = size(Y_split)[2]
            q = Int64(round(L1/2))
            res_1, fopt, lopt = generate_hyperplane_mip_diet(Y_split[:,1:q],B,η,ρ,q,C,d,δ1,use_ac)
        end
        fr, lr, fl, ll  = fopt, lopt, -fopt, -lopt
        if gen_hyper < 1
            Cr, dr, Cl, dl = Matrix(fr'), [lr], Matrix(fl'), [ll]
        else
            Cr, dr, Cl, dl = vcat(C,fr'), vcat(d,lr), vcat(C,fl'), vcat(d,ll)
        end
        gen_hyper += 1
        push!(hyp_list_f, fopt)
        push!(hyp_list_l, lopt)


        # Right child
        res_right, rlb_right, tlb_right, Rlb_right, Vlb_right, lb_right = solve_rpt_sdp_diet_quad(A,b,h,c,B,η,ρ,θ,q,λ,Cr,dr,use_lmi,use_new,rob_nutr)
        if res_right
            X_right = calculate_candidate_vectors_diet(rlb_right,tlb_right,Rlb_right,Vlb_right,ρ,q,B,η,rob_nutr)
            Y_right = calculate_different_vectors(X_right,rlb_right)
            ub_right = calculate_diet_ub_quad(Y_right,A,b,h,ρ,q,B,η,c,λ,rob_nutr)
            if lb_right < UB
                push!(nodes_list,[ub_right,lb_right,rlb_right,Cr,dr,Y_right])
            end
        end
        # Left child
        res_left, rlb_left, tlb_left, Rlb_left, Vlb_left, lb_left = solve_rpt_sdp_diet_quad(A,b,h,c,B,η,ρ,θ,q,λ,Cl,dl,use_lmi,use_new,rob_nutr)
        if res_left
            X_left = calculate_candidate_vectors_diet(rlb_left,tlb_left,Rlb_left,Vlb_left,ρ,q,B,η,rob_nutr)
            Y_left = calculate_different_vectors(X_left,rlb_left)
            ub_left = calculate_diet_ub_quad(Y_left,A,b,h,ρ,q,B,η,c,λ,rob_nutr)
            if lb_left < UB
                push!(nodes_list,[ub_left,lb_left,rlb_left,Cl,dl,Y_left])
            end
        end

        if isempty(nodes_list)
            if res_right
                push!(ub_list, ub_right)
                push!(lb_list, lb_right)
            end
            if res_left
                push!(ub_list, ub_left)
                push!(lb_list, lb_left)
            end
            break
        else
            ind = argmin([nodes_list[i][2] for i in 1:length(nodes_list)])
            cur_node = nodes_list[ind]
            deleteat!(nodes_list, ind)
            cur_ub, cur_lb = cur_node[1], cur_node[2]
            C, d, Y_split = cur_node[4], cur_node[5], cur_node[6]
            LB = cur_lb
            if cur_ub < UB
                UB = cur_ub
                opt_sol, opt_val = cur_node[3], cur_node[1]
            end
            push!(ub_list, cur_ub)
            push!(lb_list, cur_lb)
        end
        t0_2 = time_ns()
        total_time = (t0_2-t0_1)*10^(-9)
    end
    return opt_sol, opt_val, gen_hyper, ub_list, lb_list, hyp_list_f, hyp_list_l
end






df = CSV.read("/Users/thodoris/Downloads/WFP_dataset.csv", DataFrame; header=true)
df = select!(df, Not(:Salt))
X = convert(Matrix, df[:,1:24])
y = convert(Vector,df[:,25])
n = size(X)[2]
m1 = Model(Gurobi.Optimizer)
@variable(m1, A[1:n,1:n])
@variable(m1, b[1:n])
@variable(m1, h)
@objective(m1, Min, (1/n)*sum(
                    (X[i,:]'*A*X[i,:] + b'*X[i,:] + h - y[i])^2
                    for i in 1:n ) )
optimize!(m1)

A_opt, b_opt, h_opt = JuMP.value.(A), JuMP.value.(b), JuMP.value(h)




ρ = 0.1
p = 3
# q = 1/(1-1/p)
q = 3/2
θ = get_feasible_theta(q)
rob_nutr = 10
λ = 5

use_lmi = false

use_new = false
res_nolmi_wo, r_opt, t_opt, R_opt, V_opt, lb_no_lmi_wo = solve_rpt_sdp_diet_quad(A_opt,b_opt,h_opt,c,B,η,ρ,θ,q,λ,C,d,use_lmi,use_new,rob_nutr)
X = calculate_candidate_vectors_diet(r_opt,t_opt,R_opt,V_opt,ρ,q,B,η,rob_nutr)
Y_split = calculate_different_vectors(X,r_opt)
ub_no_lmi_wo = calculate_diet_ub_quad(Y_split,A_opt,b_opt,h_opt,ρ,q,B,η,c,λ,rob_nutr)

use_new = true
res_nolmi_w, r_opt, t_opt, R_opt, V_opt, lb_no_lmi_w = solve_rpt_sdp_diet_quad(A_opt,b_opt,h_opt,c,B,η,ρ,θ,q,λ,C,d,use_lmi,use_new,rob_nutr)
X = calculate_candidate_vectors_diet(r_opt,t_opt,R_opt,V_opt,ρ,q,B,η,rob_nutr)
Y_split = calculate_different_vectors(X,r_opt)
ub_no_lmi_w = calculate_diet_ub_quad(Y_split,A_opt,b_opt,h_opt,ρ,q,B,η,c,λ,rob_nutr)

use_lmi = true

use_new = false
res_lmi_wo, r_opt, t_opt, R_opt, V_opt, lb_lmi_wo = solve_rpt_sdp_diet_quad(A_opt,b_opt,h_opt,c,B,η,ρ,θ,q,λ,C,d,use_lmi,use_new,rob_nutr)
X = calculate_candidate_vectors_diet(r_opt,t_opt,R_opt,V_opt,ρ,q,B,η,rob_nutr)
Y_split = calculate_different_vectors(X,r_opt)
ub_lmi_wo = calculate_diet_ub_quad(Y_split,A_opt,b_opt,h_opt,ρ,q,B,η,c,λ,rob_nutr)

use_new = true
res_lmi_w, r_opt, t_opt, R_opt, V_opt, lb_lmi_w = solve_rpt_sdp_diet_quad(A_opt,b_opt,h_opt,c,B,η,ρ,θ,q,λ,C,d,use_lmi,use_new,rob_nutr)
X = calculate_candidate_vectors_diet(r_opt,t_opt,R_opt,V_opt,ρ,q,B,η,rob_nutr)
Y_split = calculate_different_vectors(X,r_opt)
ub_lmi_w = calculate_diet_ub_quad(Y_split,A_opt,b_opt,h_opt,ρ,q,B,η,c,λ,rob_nutr)


println("------------------------------------")
println("Without the LMI:")
println("------------------------------------")
println("Lower Bound without Additions:")
println(lb_no_lmi_wo)
println("Lower Bound with Additions:")
println(lb_no_lmi_w)
println("Upper Bound without Additions:")
println(ub_no_lmi_wo)
println("Upper Bound with Additions:")
println(ub_no_lmi_w)
println("------------------------------------")
println("With the LMI:")
println("------------------------------------")
println("Lower Bound without Additions:")
println(lb_lmi_wo)
println("Lower Bound with Additions:")
println(lb_lmi_w)
println("Upper Bound without Additions:")
println(ub_lmi_wo)
println("Upper Bound with Additions:")
println(ub_lmi_w)

δ = 1e-3
δ1 = 1
use_ac = true
use_new = true
use_lmi = true
t1 = time_ns()
opt_sol, opt_val, gen_hyper, ub_list, lb_list, hyp_list_f, hyp_list_l  = rpt_branch_and_split_diet_quad(A_opt,b_opt,h_opt,c,B,η,ρ,θ,q,λ,δ,δ1,use_ac,use_lmi,use_new,rob_nutr)
t2 = time_ns()
total_time = (t2-t1)*10^(-9)
