# Load the packages
using JuMP, MosekTools, Mosek, LinearAlgebra, SCS, COSMO, Literate, OffsetArrays

# %% Some helper functions

# %% construct e_i in R^n
function e_i(n, i)
    e_i_vec = zeros(n, 1)
    e_i_vec[i] = 1
    return e_i_vec
end

# %% construct symmetric outer product

function ⊙(a,b)
    return ((a*b')+(b*a')) ./ 2
end

# %% Parameters to be tuned
N = 5
L = 1
μ = 0
c_w = 1
c_f = 0

# define all the bold vectors
# --------------------------

# define 𝐰_0

𝐰_0 = e_i(N+2, 1)

𝐰_star = zeros(N+2, 1)

𝐠_star = zeros(N+2,1)

# define 𝐠_0, 𝐠_1, …, 𝐠_N

# first we define 𝐠_Julia vectors and then 𝐠 vectors

𝐠 = OffsetArray(zeros(N+2, N+1), 1:N+2, 0:N)
# 𝐠= [𝐠_0 𝐠_1 𝐠_2 ... 𝐠_N]
for i in 0:N
    𝐠[:,i] = e_i(N+2, i+2)
end


𝐟_star = zeros(N+1,1)

# time to define 𝐟_0, 𝐟_1, …, 𝐟_N

𝐟 = OffsetArray(zeros(N+1, N+1), 1:N+1, 0:N)
# 𝐟 = [𝐟_0, 𝐟_1, …, 𝐟_N]

for i in 0:N
    𝐟[:,i] = e_i(N+1, i+1)
end

pep_model = Model(optimizer_with_attributes(Mosek.Optimizer, "MSK_DPAR_INTPNT_CO_TOL_PFEAS" => 1e-10))

# time to define all the variables
# -------------------------------

# define α′ (can be typed as \alpha[TAB]\prime[TAB])
@variable(pep_model, α′[1:N, 0:N-1])

# define λ variables

@variable(pep_model, λ_i_ip1[0:N-1] >= 0)
# this defines (λ_{i,i+1})_{i∈[0:N-1]} in Julia indexing

@variable(pep_model, λ_star_i[0:N] >= 0)
# this defines (λ_{⋆,i})_{i∈[0:N]} in Julia indexing

# define τ
@variable(pep_model, τ >= 0)

# define objective
@objective(
    pep_model,
    Min,
    τ
    )

# Add the linear equality constraint
# ----------------------------------
@constraint(pep_model,
    τ * c_f .* 𝐟[:,0] + sum(λ_i_ip1[i] .* (𝐟[:,i+1]-𝐟[:,i]) for i in 0:N-1)
    + sum(λ_star_i[i] .* (𝐟[:,i] - 𝐟_star) for i in 0:N)
    .== 𝐟[:,N] - 𝐟_star
)

# Add the giant LMI constraint step by step
# ----------------------------------------

term_1 = c_w * τ * ⊙(𝐰_0,𝐰_0)

term_2_part_1 = sum(λ_i_ip1[i] .* ⊙(𝐠[:,i]-𝐠[:,i+1],𝐠[:,i]-𝐠[:,i+1]) for i in 0:N-1)

term_2_part_2 = sum(λ_star_i[i] .* ⊙(𝐠_star - 𝐠[:,i],𝐠_star - 𝐠[:,i]) for i in 0:N)

term_2 = (term_2_part_1 + term_2_part_2) ./ (2*L)

term_3 = - λ_star_i[0] .* ⊙(𝐠[:,0],𝐰_0)

term_4_part_1 = - sum( λ_i_ip1[i] .* ⊙(𝐠[:,i],𝐰_0) for i in 1:N-1)

term_4_part_2 = (1/L) .*
sum( sum(α′[i,j] .* ⊙(𝐠[:,i],𝐠[:,j]) for j in 0:i-1) for i in 1:N-1)

term_4 = term_4_part_1 + term_4_part_2

term_5 = - ( ⊙(𝐠[:,N],𝐰_0) - (1/L) .* sum(α′[N,j] .* ⊙(𝐠[:,N],𝐠[:,j]) for j in 0:N-1))

term_6_part_1 = sum(λ_i_ip1[i] .* ⊙(𝐠[:,i+1],𝐰_0) for i in 0:N-1)

term_6_part_2 = - (1/L) .*
sum( sum( α′[i,j] .* ⊙(𝐠[:,i+1],𝐠[:,j]) for j in 0:i-1) for i in 1:N-1)

term_6 = term_6_part_1 + term_6_part_2

# oof, okay constructed the terms for the LMI constraint, let us hope that there is no mistake and add them together

S_mat = term_1 + term_2 + term_3 + term_4 + term_5 + term_6

# add the LMI constraint now

@constraint(pep_model,
S_mat in PSDCone()
)

# time to optimize!
# -----------------
optimize!(pep_model)

println("termination status =", termination_status(pep_model) )

# Collect the decision variables
# -----------------------------
α′_opt = OffsetArray(value.(α′), 1:N, 0:N-1)
λ_opt_i_ip1 = OffsetVector(value.(λ_i_ip1), 0:N-1)
λ_opt_star_i = OffsetVector(value.(λ_star_i), 0:N)
τ_opt = value(τ)

# Recover α_{i,j} from α′_{i,j}
# -----------------------------

α_opt = OffsetArray(zeros(size(α′_opt)), 1:N, 0:N-1)

for i in 1:N-1
    for j in 0:i-1
            α_opt[i,j] = (α′_opt[i,j] / λ_opt_i_ip1[i])
    end
end

for j in 0:N-1
    α_opt[N,j] = α′_opt[N,j]
end



# Recover h_{i,j} from α_{i,j}
h_opt =  OffsetArray(zeros(size(α_opt)), 1:N, 0:N-1)

# set h(1,0)
h_opt[1,0] = α_opt[1,0]

for i in 2:N
    for j in 0:i-1
        if j == i-1
            h_opt[i,j] = α_opt[i,j]
        else
            h_opt[i,j] = α_opt[i,j] - α_opt[i-1,j]
        end
    end
end

function full_pep_solver(N, L, μ, c_w, c_f)

    # define all the bold vectors
    # --------------------------

    # define 𝐰_0

    𝐰_0 = e_i(N+2, 1)

    𝐰_star = zeros(N+2, 1)

    𝐠_star = zeros(N+2,1)

    # define 𝐠_0, 𝐠_1, …, 𝐠_N

    # first we define 𝐠_Julia vectors and then 𝐠 vectors

    𝐠 = OffsetArray(zeros(N+2, N+1), 1:N+2, 0:N)
    # 𝐠= [𝐠_0 𝐠_1 𝐠_2 ... 𝐠_N]
    for i in 0:N
        𝐠[:,i] = e_i(N+2, i+2)
    end


    𝐟_star = zeros(N+1,1)

    # time to define 𝐟_0, 𝐟_1, …, 𝐟_N

    𝐟 = OffsetArray(zeros(N+1, N+1), 1:N+1, 0:N)
    # 𝐟 = [𝐟_0, 𝐟_1, …, 𝐟_N]

    for i in 0:N
        𝐟[:,i] = e_i(N+1, i+1)
    end

    # Define JuMP model
    # ----------------

    pep_model = Model(optimizer_with_attributes(Mosek.Optimizer, "MSK_DPAR_INTPNT_CO_TOL_PFEAS" => 1e-10))

    #define all the variables
    # -------------------------------

    # define α′ (can be typed as \alpha[TAB]\prime[TAB])
    @variable(pep_model, α′[1:N, 0:N-1])

    # define λ variables

    @variable(pep_model, λ_i_ip1[0:N-1] >= 0)
    # this defines (λ_{i,i+1})_{i∈[0:N-1]} in Julia indexing

    @variable(pep_model, λ_star_i[0:N] >= 0)
    # this defines (λ_{⋆,i})_{i∈[0:N]} in Julia indexing

    # define τ
    @variable(pep_model, τ >= 0)

    # define objective
    # ------------------

    @objective(
    pep_model,
    Min,
    τ
    )

    # Add the linear equality constraint
    # ----------------------------------
    @constraint(pep_model,
    τ * c_f .* 𝐟[:,0] + sum(λ_i_ip1[i] .* (𝐟[:,i+1]-𝐟[:,i]) for i in 0:N-1)
    + sum(λ_star_i[i] .* (𝐟[:,i] - 𝐟_star) for i in 0:N)
    .== 𝐟[:,N] - 𝐟_star
    )

    # Add the giant LMI constraint step by step
    # ----------------------------------------

    # Define all the terms one by one

    term_1 = c_w * τ * ⊙(𝐰_0,𝐰_0)

    term_2_part_1 = sum(λ_i_ip1[i] .* ⊙(𝐠[:,i]-𝐠[:,i+1],𝐠[:,i]-𝐠[:,i+1]) for i in 0:N-1)

    term_2_part_2 = sum(λ_star_i[i] .* ⊙(𝐠_star - 𝐠[:,i],𝐠_star - 𝐠[:,i]) for i in 0:N)

    term_2 = (term_2_part_1 + term_2_part_2) ./ (2*L)

    term_3 = - λ_star_i[0] .* ⊙(𝐠[:,0],𝐰_0)

    term_4_part_1 = - sum( λ_i_ip1[i] .* ⊙(𝐠[:,i],𝐰_0) for i in 1:N-1)

    term_4_part_2 = (1/L) .*
    sum( sum(α′[i,j] .* ⊙(𝐠[:,i],𝐠[:,j]) for j in 0:i-1) for i in 1:N-1)

    term_4 = term_4_part_1 + term_4_part_2

    term_5 = - ( ⊙(𝐠[:,N],𝐰_0) - (1/L) .* sum(α′[N,j] .* ⊙(𝐠[:,N],𝐠[:,j]) for j in 0:N-1))

    term_6_part_1 = sum(λ_i_ip1[i] .* ⊙(𝐠[:,i+1],𝐰_0) for i in 0:N-1)

    term_6_part_2 = - (1/L) .*
    sum( sum( α′[i,j] .* ⊙(𝐠[:,i+1],𝐠[:,j]) for j in 0:i-1) for i in 1:N-1)

    term_6 = term_6_part_1 + term_6_part_2

    # oof, okay constructed the terms for the LMI constraint, let us hope that there is no mistake and add them together

    S_mat = term_1 + term_2 + term_3 + term_4 + term_5 + term_6

    # add the LMI constraint now

    @constraint(pep_model,
    S_mat in PSDCone()
    )

    # time to optimize!
    # -----------------
    optimize!(pep_model)

    println("termination status =", termination_status(pep_model) )

    α′_opt = OffsetArray(value.(α′), 1:N, 0:N-1)
    λ_opt_i_ip1 = OffsetVector(value.(λ_i_ip1), 0:N-1)
    λ_opt_star_i = OffsetVector(value.(λ_star_i), 0:N)
    τ_opt = value(τ)

    # Recover α_{i,j} from α′_{i,j}
    # -----------------------------

    α_opt = OffsetArray(zeros(size(α′_opt)), 1:N, 0:N-1)

    for i in 1:N-1
        for j in 0:i-1
            α_opt[i,j] = (α′_opt[i,j] / λ_opt_i_ip1[i])
        end
    end

    for j in 0:N-1
        α_opt[N,j] = α′_opt[N,j]
    end

    # Recover h_{i,j} from α_{i,j}
    # ----------------------------

    h_opt =  OffsetArray(zeros(size(α_opt)), 1:N, 0:N-1)

    # set h(1,0)
    h_opt[1,0] = α_opt[1,0]

    # set the rest of the variables

    for i in 2:N
        for j in 0:i-1
            if j == i-1
                h_opt[i,j] = α_opt[i,j]
            else
                h_opt[i,j] = α_opt[i,j] - α_opt[i-1,j]
            end
        end
    end

    # return all the outputs
    # ----------------------

    return   α′_opt, λ_opt_i_ip1, λ_opt_star_i, τ_opt, α_opt, h_opt

end

α′_opt, λ_opt_i_ip1, λ_opt_star_i, τ_opt, α_opt, h_opt = full_pep_solver(N, L, μ, c_w, c_f)