using JuMP, Ipopt, LinearAlgebra, Distributions, Gurobi

## create a dictionary with optimization data
n = 10 # number of generators 
data = Dict(:C => diagm(ones(n)), :c => [i*n for i in 1:n], :d => 100, :p̲ => zeros(n), :p̅ => [i*n/5 for i in 1:n])


# solve the primal ED problem 
function ed_opt_primal(data)
    # model size 
    n = length(data[:c])
    # create a JuMP model 
    model = Model(() -> Ipopt.Optimizer())
    set_silent(model)
    # optimization variables 
    @variable(model, p[1:n])
    # objective function 
    @objective(model, Min, p'*data[:C]*p + data[:c]'*p)
    # constraints 
    @constraint(model, ν, ones(n)'*p == data[:d])
    @constraint(model, λ̲, p .>= data[:p̲])
    @constraint(model, λ̅, -p .>= -data[:p̅])
    # solve optimization 
    optimize!(model)
    return Dict(:p => JuMP.value.(p), :obj_val => JuMP.objective_value(model), :ν => JuMP.dual.(ν), :λ̲ => JuMP.dual.(λ̲), :λ̅ => JuMP.dual.(λ̅))
end 
sol_pr  = ed_opt_primal(data)

# solve the ED problem via its KKTs
function KKT_ed(data)
    # model size 
    n = length(data[:c])
    # create a JuMP model 
    model = Model(() -> Gurobi.Optimizer())
    set_silent(model)
    # optimization variables 
    @variable(model, p[1:n])
    @variable(model, ν)
    @variable(model, λ̲[1:n]>=0)
    @variable(model, λ̅[1:n]>=0)
    @variable(model, u̲[1:n], Bin)
    @variable(model, u̅[1:n], Bin)

    # primal constraints 
    @constraint(model, ones(n)'*p == data[:d])
    @constraint(model, p .>= data[:p̲])
    @constraint(model, -p .>= -data[:p̅])

    # Lagrange optimality 
    @constraint(model, 2*data[:C]*p + data[:c] + λ̅ .- λ̲ - ones(n)*ν .== 0)

    # complementarity slackness 
    for i in 1:n 
        @constraint(model, p[i] - data[:p̲][i] <= 1000 * (1 - u̲[i]))
        @constraint(model, λ̲[i] <= 1000 * u̲[i])
        @constraint(model, data[:p̅][i] - p[i] <= 1000 * (1 - u̅[i]))
        @constraint(model, λ̅[i] <= 1000 * u̅[i])
    end

    # solve optimization 
    optimize!(model)
    return Dict(:p => JuMP.value.(p), :obj_val => JuMP.objective_value(model), :ν => JuMP.value.(ν), :λ̲ => JuMP.value.(λ̲), :λ̅ => JuMP.value.(λ̅))
end 


sol_primal = ed_opt_primal(data)
sol_dual = KKT_ed(data)

# if the priaml solution and KKT point are the same 
# then the KKTs are formulated correctly 
@show norm(sol_primal[:p] .- sol_dual[:p])
@show norm(sol_primal[:ν] .- sol_dual[:ν])
@show norm(sol_primal[:λ̅] .- sol_dual[:λ̅])
@show norm(sol_primal[:λ̲] .- sol_dual[:λ̲])


### create a dataset of k observations 
k = 300
d = zeros(k)
p_opt = zeros(n,k)
ν_opt = zeros(k)
for i in 1:k
    data_copy = deepcopy(data)
    data_copy[:d] = rand(Uniform(0,110),1)[1]
    d[i] = data_copy[:d]
    sol_primal = ed_opt_primal(data_copy)
    p_opt[:,i] = sol_primal[:p]
    ν_opt[i] = sol_primal[:ν]
end

### Inverse optimization problem 
function IO_ed(data)
    # model size 
    n = length(data[:c])
    # create a JuMP model 
    model = Model(() -> Gurobi.Optimizer())
    # set_silent(model)
    # optimization variables 
    @variable(model, c[1:n])
    @variable(model, p[1:n,1:k])
    @variable(model, ν[1:k])
    @variable(model, λ̲[1:n,1:k]>=0)
    @variable(model, λ̅[1:n,1:k]>=0)
    @variable(model, u̲[1:n,1:k], Bin)
    @variable(model, u̅[1:n,1:k], Bin)

    # objective function 
    @objective(model, Min, sum((p[i,j]-p_opt[i,j])^2 for i in 1:n, j in 1:k) + sum((ν[j]-ν_opt[j])^2 for j in 1:k))

    # upper level constraints 
    @constraint(model, 0 .<= c .<= 100)

    # primal constraints of the lower-level problem 
    @constraint(model, con_1[i=1:k], ones(n)'*p[:,i] == d[i])
    @constraint(model, con_2[i=1:k], p[:,i] .>= data[:p̲])
    @constraint(model, con_3[i=1:k], -p[:,i] .>= -data[:p̅])

    # Lagrange optimality of the lower-level problem 
    @constraint(model, con_4[i=1:k], 2*data[:C]*p[:,i] .+ c .+ λ̅[:,i] .- λ̲[:,i] .- ones(n)*ν[i] .== 0)

    # complementarity slackness of the lower-level problem 
    for i in 1:n
        for j in 1:k
            @constraint(model, p[i,j] - data[:p̲][i] <= 1000 * (1 - u̲[i,j]))
            @constraint(model, λ̲[i,j] <= 1000 * u̲[i,j])
            @constraint(model, data[:p̅][i] - p[i,j] <= 1000 * (1 - u̅[i,j]))
            @constraint(model, λ̅[i,j] <= 1000 * u̅[i,j])
        end
    end

    # solve optimization 
    optimize!(model)
    return Dict(:c => JuMP.value.(c), :p => JuMP.value.(p), :obj_val => JuMP.objective_value(model), :ν => JuMP.value.(ν), :λ̲ => JuMP.value.(λ̲), :λ̅ => JuMP.value.(λ̅))
end 


io_sol = IO_ed(data)
[io_sol[:c] data[:c] io_sol[:c].-data[:c]]