# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 0.6.4
#     language: julia
#     name: julia-0.6
# ---

# # Make PSD

include("../iLQR.jl")

using iLQR
using Base.Test
using BenchmarkTools
using Plots

# ## Test Dynamics

# +
n,m = 3,2
A = [0 1 0; 0 0 1; 0 0 0]
B = [0 0; 0 1; 1 1]
x = [1;2;3]
u = [4;5]

Qf = diagm([10, 20, 30])
R = eye(m)*2
Q = diagm([3, 2, 1])
# -

# ## Functions

# +
function makePSD_current(A)
    n = size(A,1)
    mu = 0
    for i = 1:100
        if any(eigvals(A).<0.)
            mu += 1
            A += mu*eye(n)
        else
            break;
        end
    end
    return A
end

function makePSD_eig(A)
    δ = 1e-3
    λ,Q = eig(A)
    λ[λ .< δ] = δ
    Q*diagm(λ)Q'
end

function makePSD_cur_mod(A,δ=1.)
    n = size(A,1)
    mu = 0
    for i = 1:100
        if ~isposdef(A)
            mu += δ
            A += mu*eye(n)
        else
            break;
        end
    end
    return A
end

function makePSD_33(A)
    n = size(A,1)
    I = eye(n)

    β = 1e-3
    min_diag = minimum(diag(A))
    if min_diag > 0
        constraint_decrease_ratio = 0
    else
        constraint_decrease_ratio = -min_diag + β
    end
    
    for k = 1:100
        v = eigvals(A)
        if isposdef(A+constraint_decrease_ratio*I)
#             println("iter: $k")
            return A+constraint_decrease_ratio*I
        else
            constraint_decrease_ratio = max(2*constraint_decrease_ratio,β)
        end
    end
    return A
end

function makePSD_33_mod(A)
    n = size(A,1)
    I = eye(n)
    mu = 0
    
    β = 1e-2
    min_diag = minimum(diag(A))
    if min_diag > 0
        mu = 0
        δ = 1
    else
        mu = -min_diag + β
        δ = -min_diag
    end
    A += mu*I
    min_diag = abs(min_diag)
    
    for k = 1:100
        if ~isposdef(A)
            mu += δ
            A += mu*I
        else
#             println("iter: $k")
            break;
        end
    end
    return A
end

function makePSD_chol(A::Array)
    n = size(A,1)
    d = zeros(n)
    C = zeros(A)
    L = eye(A)
    β = 1e-3
    δ = 1e-3
    for j = 1:n
        t = 0
        for s = 1:j-1
            t += d[s]*L[j,s]^2
        end
        C[j,j] = A[j,j] - t
#         d[j] = cjj
        θ = maximum(C[j:n,j])
        d[j] = max(abs(C[j,j]),(θ\β)^2,δ)
        
        for i = j+1:n
            t = 0
            for s = 1:j-1
                t += d[s]*L[i,s]*L[j,s]
            end
            C[i,j] = A[i,j] - t
            L[i,j] = C[i,j] / d[j]
        end
    end
    D = diagm(d)
    return L,D
end
# -

# ## Create Symmetric (non PSD) Matrix

n = 10
A = randn(n,n)
A = A'A-1
isposdef(A)

A1 = makePSD_current(A); println(norm(A1-A))
A2 = makePSD_cur_mod(A); println(norm(A2-A))
A3 = makePSD_33(A); println(norm(A3-A))
A4 = makePSD_33_mod(A); println(norm(A4-A))
@test isposdef(A1)
@test isposdef(A2)
@test isposdef(A3)
@test isposdef(A4)



# ## Timing Tests

@btime makePSD_cur_mod(A)
@btime makePSD_33_mod(A)

# ### Test PSD Matrix

n = 10
A = randn(n,n)
A = A'A
isposdef(A)

@btime makePSD_cur_mod(A)
@btime makePSD_33_mod(A)

N = 2.^(0:10)
T = zeros(length(N))
for (i,n) in enumerate(N)
    A = randn(n,n)
    A = A'A - 1
    t = @elapsed makePSD_cur_mod(A)
    T[i] = t
    println("n: $n, time=$t")
end

plot(N,T,xscale=:log10)



StarcoderdataJupyterScriptsDedupFiltered