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

f(x)  =   0.125*x'*x - 0.5*x[1] + 0.25*x[2]
h(x)  =   x[1]-x[2]

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

# optimal solution
x_opt, obj_opt = opt()

# plot objective level sets, feasible region and the optimal solution
# 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,xlims=(-1,2.05),ylims=(-1.5,1.5),legend=:bottomleft,aspect_ratio = :equal)
# plot rectangle corners of the feasible region
con(x) = x
plot!(-1:3,con,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 the optimall solution
scatter!([x_opt[1]],[x_opt[2]],label="",c=purple)
plot!(size=(500,500))


# Lagrangian function 
# L(x,λ) = f(x) + λ'*h(x) 
# write a function for the derivative w.r.t. x
∇L_x(x, λ) = 0.25 * x .+ [-0.5; 0.25] .+ λ * [1; -1]
# write a function for the derivative w.r.t. λ
∇L_λ(x, λ) = x[1] - x[2]


# saddle-flow alg parameters
η = 0.01; iter_max = 3000
# initialize variables 
x = [-0.5;-1]; λ = 1
# arrays to store solutions across iterations
x_sp_save = zeros(2,iter_max)
for i in 1:iter_max
    # write down the x update 
    x = x - η * ∇L_x(x, λ)
    # write down the dual update 
    λ = λ + η * ∇L_λ(x, λ)
    x_sp_save[:,i] = x
end

# add the saddle-point flow trajectory to the plot
plot!(x_sp_save[1,:],x_sp_save[2,:],label=false,lw=4,c=purple)


# Augmented Lagrangian function 
ρ = 100
# L(x,λ) = f(x) + λ'*h(x)  + ρ*norm(h(x))^2
# write a new function for the derivative w.r.t. x
∇L_x_ρ(x, λ) =  0.25 * x .+ [-0.5; 0.25] .+ λ * [1; -1] + 2*ρ*[x[1]-x[2];x[2]-x[1]]


# initialize variables 
x = [-0.5;-1]; λ = 1
# arrays to store solutions across iterations
x_sp_save = zeros(2,iter_max)
for i in 1:iter_max
    # write down the x update 
    x = x - η * ∇L_x_ρ(x, λ)
    # write down the dual update 
    λ = λ + η * ∇L_λ(x, λ)
    x_sp_save[:,i] = x
end

# add the regularized saddle-point flow trajectory to the plot
plot!(x_sp_save[1,:],x_sp_save[2,:],label=false,lw=4,c=blue)

# manipulate ρ to see how it can help or abstract convergence