using JuMP, Ipopt, LinearAlgebra, Plots
print("\033c")
"""
minimize    0.125 * ||x||^2 - 0.5x1 + 0.25 x2  
subject to  x1 >= 0
            x1 - x2 <= 0
"""

f(x)    =   0.125*x'*x - 0.5*x[1] + 0.25*x[2]
∇f(x)   =   # add function for f's derivative 
g(x)    =   [-x[1];x[1]-x[2]]
∇g(x)   =   # add function for g's derivative 

# function to solve optimization
function opt()
    model = Model(() -> Ipopt.Optimizer())
    set_silent(model)
    @variable(model, x[1:2])
    @objective(model, Min, f(x))
    @constraint(model, g(x) .<= 0)
    optimize!(model)
    return JuMP.value.(x)
end

# optimization function for dual variable update
function μ(x,α)
    model = Model(() -> Ipopt.Optimizer())
    set_silent(model)
    @variable(model, μ[1:2]>=0)
    @objective(model, Min, # write down the objective function  
    @constraint(model, # write down the admissible constraint set 
    optimize!(model)
    return JuMP.value.(μ)
end

# optimal solution
x_opt = opt()

# initial feasible point 
x₀ = [-1;0.5]

# constroller parameters 
η = 0.1;        # control gain
α = 1;          # constraint parameter
iter_max = 500 # iteration limit

# unconstrained optimization
function unc_opt(x) 
    x_save = zeros(2,iter_max)
    for i in 1:iter_max
        x = x - η*∇f(x) 
        x_save[:,i] = x
    end
    return x, x_save
end
x_unc_opt, x_unc_save = unc_opt(x₀)

# safe gradient flow optimization
function sgf_opt(x) 
    x_save = zeros(2,iter_max)
    for i in 1:iter_max
        x = x - η*(∇f(x) + ∇g(x)'*μ(x,α))
        x_save[:,i] = x
    end
    return x, x_save
end
x_sgf_opt, x_sgf_save = sgf_opt(x₀)

# make a plot
using Colors
purple=RGB(122/256,30/256,71/256)
blue=RGB(0/256, 39/256, 76/256)
plo = plot(frame=:box,ylabel="",xlabel="",xlims=(-1,2.05),ylims=(-1.5,1.5),legend=:bottomleft,aspect_ratio = :equal,xaxis=nothing,yaxis=nothing)
# Define rectangle corners of the feasible region
con2(x) = max(x,0)
plot!(-1:3,con2,fillrange = 2,color=purple,opacity=0.5,label=false,lw=2)
# plot objective level sets
f(x, y) = 0.125 * (x.^2 + y.^2) - 0.5 * x .+ 0.25 * y
contourf!(-1:0.1:2.5, -1.5:0.1:1.5, f, linewidth=1, c=blue, opacity=0.1, label="",colorbar=false)
# plot unconstrained optimization trajectory
plot!(x_unc_save[1,:],x_unc_save[2,:],label="unconstrained",lw=3,c=blue)
# plot safe gradient flow trajectory
plot!(x_sgf_save[1,:],x_sgf_save[2,:],label=false,lw=4,c=purple)
# plot the optimall solution
scatter!([x_opt[1]],[x_opt[2]],label="",c=purple)
scatter!([x_unc_opt[1]],[x_unc_opt[2]],label="",c=blue)
plot!(size=(500,500))