# Factor Analysis

"""Factor Analysis type"""
immutable FactorAnalysis{T<:AbstractFloat}
    mean::Vector{T}       # sample mean: of length d (mean can be empty, which indicates zero mean)
    W::Matrix{T}          # factor loadings matrix: of size d x p
    Ψ::Vector{T}          # noise covariance: diagonal of size d x d
end

indim(M::FactorAnalysis) = size(M.W, 1)
outdim(M::FactorAnalysis) = size(M.W, 2)

Base.mean(M::FactorAnalysis) = fullmean(indim(M), M.mean)
projection(M::FactorAnalysis) = svdfact(M.W)[:U] # recover principle components from the weight matrix
Base.cov(M::FactorAnalysis) = M.W*M.W'+ Diagonal(M.Ψ)
Base.var(M::FactorAnalysis) = M.Ψ
loadings(M::FactorAnalysis) = M.W

## use

function transform{T<:AbstractFloat}(m::FactorAnalysis{T}, x::AbstractVecOrMat{T})
    xn = centralize(x, mean(m))
    W = m.W
    WᵀΨ⁻¹ = W'*diagm(1./m.Ψ)  # (q x d) * (d x d) = (q x d)
    return inv(I+WᵀΨ⁻¹*W)*(WᵀΨ⁻¹*xn)  # (q x q) * (q x d) * (d x 1) = (q x 1)
end

function reconstruct{T<:AbstractFloat}(m::FactorAnalysis{T}, z::AbstractVecOrMat{T})
    W  = m.W
    # ΣW(W'W)⁻¹z+μ = ΣW(W'W)⁻¹W'Σ⁻¹(x-μ)+μ = Σ(WW⁻¹)((W')⁻¹W')Σ⁻¹(x-μ)+μ = ΣΣ⁻¹(x-μ)+μ = (x-μ)+μ = x
    return cov(m)*W*inv(W'W)*z .+ mean(m)
end

## show

function Base.show(io::IO, M::FactorAnalysis)
    print(io, "Factor Analysis(indim = $(indim(M)), outdim = $(outdim(M)))")
end

## core algorithms

""" Expectation-maximization (EM) algorithms for maximum likelihood factor analysis.

    Rubin, Donald B., and Dorothy T. Thayer. "EM algorithms for ML factor analysis." Psychometrika 47.1 (1982): 69-76.
"""
function faem{T<:AbstractFloat}(S::DenseMatrix{T}, mv::Vector{T}, n::Int;
             maxoutdim::Int=size(X,1)-1,
             tol::Real=1.0e-6,   # convergence tolerance
             tot::Integer=1000)  # maximum number of iterations

    d = size(S,1)
    q = maxoutdim
    W = eye(T,d,q)
    Ψ = fill(T(0.01),d)

    L_old = 0.
    for c in 1:tot
        # EM-steps
        Ψ⁻¹W = diagm(1./Ψ)*W
        SΨ⁻¹W = S*Ψ⁻¹W
        H = SΨ⁻¹W*inv(I + W'*Ψ⁻¹W)
        W⁺ = SΨ⁻¹W*inv(I + H'*Ψ⁻¹W)
        Ψ⁺ = diag(S - H*(W⁺)')
        # new parameters
        W = W⁺
        Ψ = Ψ⁺
        # log likelihood
        Ψ⁻¹ = diagm(1./Ψ)
        WᵀΨ⁻¹ = W'*Ψ⁻¹
        detΣ = prod(Ψ)*det(I + WᵀΨ⁻¹*W)
        Σ⁻¹ = Ψ⁻¹ - Ψ⁻¹*W*inv(I + WᵀΨ⁻¹*W)*WᵀΨ⁻¹
        L = (-n/2)*(d*log(2π) + log(detΣ) + trace(Σ⁻¹*S))
        # println("$c] ΔL: $(abs(L_old - L)), L: $L")
        if abs(L_old - L) < tol
            break
        end
        L_old = L
    end

    return FactorAnalysis(mv, W, Ψ)
end

""" Fast conditional maximization (CM) algorithm for factor analysis.

    Zhao, J-H., Philip LH Yu, and Qibao Jiang. "ML estimation for factor analysis: EM or non-EM?." Statistics and computing 18.2 (2008): 109-123.
"""
function facm{T<:AbstractFloat}(S::DenseMatrix{T}, mv::Vector{T}, n::Int;
             maxoutdim::Int=size(X,1)-1,
             tol::Real=1.0e-6,   # convergence tolerance
             tot::Integer=1000,  # maximum number of iterations
             η = tol)            # variance low bound
    d = size(S,1)

    q = maxoutdim
    W = eye(T,d,q)
    Ψ = fill(T(0.01),d)
    V = zeros(T,q)
    eᵢeᵢ = zeros(T,d,d)
    Bᵢ⁻¹ = zeros(T,d,d)
    addconst = d*log(2π)

    L_old = 0.
    for c in 1:tot
        # CM-step 1
        Ψ⁻ʰ = Diagonal(1./sqrt.(Ψ))
        S⁺ = Ψ⁻ʰ*S*Ψ⁻ʰ

        F = eigfact(S⁺)
        λ = real(F[:values])
        ord = sortperm(λ, rev=true)
        λ = λ[ord]

        q′ = λ[q] > 0 ? q : findlast(λ .> 0)

        λq = 0.0
        @inbounds for i in 1:q′
            λq += log(λ[i]) - λ[i] + 1.
            V[i] = sqrt(λ[i] - 1.)
        end

        L = (addconst + log(prod(Ψ)) + trace(S⁺) + λq)*(-n/2)

        U = convert(Matrix{T}, F[:vectors][:,ord[1:q′]])
        W = U*diagm(V[1:q′]) # set new parameter

        # CM-step 2
        @inbounds for i in 1:q′
            V[i] = 1./λ[i] - 1.
        end
        eBe⁻¹ = ωᵢᵗ⁺¹ = 0.0
        for i in 1:d
            if i == 1
                Bᵢ⁻¹ = U*diagm(V[1:q′])*U' + I
            else
                eᵢeᵢ[i-1,i-1] = 1.
                Bᵢ⁻¹ = Bᵢ⁻¹ -  ωᵢᵗ⁺¹*Bᵢ⁻¹*eᵢeᵢ*Bᵢ⁻¹ / (1. + ωᵢᵗ⁺¹ * Bᵢ⁻¹[i-1,i-1])
                eᵢeᵢ[i-1,i-1] = 0.
            end
            eBe⁻¹ = 1./Bᵢ⁻¹[i,i]
            ωᵢᵗ⁺¹ = (Bᵢ⁻¹*S⁺*Bᵢ⁻¹)[i,i] * eBe⁻¹ * eBe⁻¹ - eBe⁻¹
            Ψ[i] = max(η, (ωᵢᵗ⁺¹ + 1.)*Ψ[i]) # set new parameter
        end

        # println("$c] ΔL: $(abs(L_old - L)), L: $L")
        if abs(L_old - L) < tol
            break
        end
        L_old = L
    end

    return FactorAnalysis(mv, Diagonal(sqrt.(Ψ))*W, Ψ)
end

## interface functions
function fit{T<:AbstractFloat}(::Type{FactorAnalysis}, X::DenseMatrix{T};
             method::Symbol=:cm,
             maxoutdim::Int=size(X,1)-1,
             mean=nothing,
             tol::Real=1.0e-6,   # convergence tolerance
             tot::Integer=1000,  # maximum number of iterations
             η = tol)            # variance low bound

    d, n = size(X)

    # process mean
    mv = preprocess_mean(X, mean)
    S = Base.covm(X, isempty(mv) ? 0 : mv, 2)
    if method == :em
        M = faem(S, mv, n, maxoutdim=maxoutdim, tol=tol, tot=tot)
    elseif method == :cm
        M = facm(S, mv, n, maxoutdim=maxoutdim, tol=tol, tot=tot, η = η)
    else
        throw(ArgumentError("Invalid method name $(method)"))
    end

    return M::FactorAnalysis
end


using VLOptionsModeling
using Plots

# create an emprt contract set for now -
assetSymbol = "XYZ"
strikePrice = 50.0
underlying_price = 50.0
expirationDate = nothing
premimumValue = 0.0
numberOfContracts = 1
sense = :buy
contractMultiplier = 1.0
put_contract = VLPutOptionContract(assetSymbol, strikePrice, premimumValue; sense=sense, contractMultiplier=contractMultiplier)

# setup the lattice (Example 19.1/Fig. 19.3 Hull)
𝝙t = 30.42 / (365.0) # years
μ = log(1.0084) * (1 / 𝝙t)
σ = (1 / √𝝙t) * log(1.1224)
numberOfLevels = 36

# Create the lattice model -
latticeModel_1 = VLBinomialLattice(μ, σ, 𝝙t, numberOfLevels)
latticeModel_2 = VLBinomialLattice(μ, σ, 0.5 * 𝝙t, numberOfLevels)
latticeModel_3 = VLBinomialLattice(μ, σ, 0.25 * 𝝙t, numberOfLevels)
latticeModel_4 = VLBinomialLattice(μ, σ, 0.125 * 𝝙t, numberOfLevels)
latticeModel_5 = VLBinomialLattice(μ, σ, 0.0625 * 𝝙t, numberOfLevels)

# setup strike array -
strike_price_array = range(35.0, stop=65.0, step=0.1) |> collect

# sim -
sim_array_1 = binomial_price(put_contract, latticeModel_1, underlying_price, strike_price_array) |> check
sim_array_2 = binomial_price(put_contract, latticeModel_2, underlying_price, strike_price_array) |> check
sim_array_3 = binomial_price(put_contract, latticeModel_3, underlying_price, strike_price_array) |> check
sim_array_4 = binomial_price(put_contract, latticeModel_4, underlying_price, strike_price_array) |> check
sim_array_5 = binomial_price(put_contract, latticeModel_5, underlying_price, strike_price_array) |> check

# plots -
plot(sim_array_1[:,1],sim_array_1[:,2])
plot!(sim_array_2[:,1],sim_array_2[:,2])
plot!(sim_array_3[:,1],sim_array_3[:,2])
plot!(sim_array_4[:,1],sim_array_4[:,2])
plot!(sim_array_5[:,1],sim_array_5[:,2])


############################################################################
### Expectation Maximization Algorithm
############################################################################

# An EM Algorithm for MVN-distributed Data with missing values
# Adapted from supplementary Material to the book Machine Learning: A Probabilistic Perspective
# Copyright (2010) Kevin Murphy and Matt Dunham
# found at https://github.com/probml/pmtk3/blob/master/toolbox/BasicModels/gauss/sub/gaussMissingFitEm.m
# and at https://github.com/probml/pmtk3/blob/master/toolbox/Algorithms/optimization/emAlgo.m

# what about random restarts?

# outer function -----------------------------------------------------------------

function em_mvn(
    observed::SemObsMissing;
    start_em = start_em_observed,
    max_iter_em = 100,
    rtol_em = 1e-4,
    kwargs...)

    n_obs, n_man = observed.n_obs, Int(observed.n_man)

    # preallocate stuff?
    𝔼x_pre = zeros(n_man)
    𝔼xxᵀ_pre = zeros(n_man, n_man)

    ### precompute for full cases
    if length(observed.patterns[1]) == observed.n_man
        for row ∈ observed.rows[1]
            row = observed.data_rowwise[row]
            𝔼x_pre += row;
            𝔼xxᵀ_pre += row*row';
        end
    end

    # ess = 𝔼x, 𝔼xxᵀ, ismissing, missingRows, n_obs
    # estepFn = (em_model, data) -> estep(em_model, data, EXsum, EXXsum, ismissing, missingRows, n_obs)

    # initialize
    em_model = start_em(observed; kwargs...)
    em_model_prev = EmMVNModel(zeros(n_man, n_man), zeros(n_man), false)
    iter = 1
    done = false
    𝔼x = zeros(n_man)
    𝔼xxᵀ = zeros(n_man, n_man)

    while !done

        em_mvn_Estep!(𝔼x, 𝔼xxᵀ, em_model, observed, 𝔼x_pre, 𝔼xxᵀ_pre)
        em_mvn_Mstep!(em_model, n_obs, 𝔼x, 𝔼xxᵀ)

        if iter > max_iter_em
            done = true
            @warn "EM Algorithm for MVN missing data did not converge. Likelihood for FIML is not interpretable.
            Maybe try passing different starting values via 'start_em = ...' "
        elseif iter > 1
            # done = isapprox(ll, ll_prev; rtol = rtol)
            done = isapprox(em_model_prev.μ, em_model.μ; rtol = rtol_em) & isapprox(em_model_prev.Σ, em_model.Σ; rtol = rtol_em)
        end

        # print("$iter \n")
        iter = iter + 1
        em_model_prev.μ, em_model_prev.Σ = em_model.μ, em_model.Σ

    end

    # update EM Mode in observed
    observed.em_model.Σ .= em_model.Σ
    observed.em_model.μ .= em_model.μ
    observed.em_model.fitted = true

    return nothing

end

# E and M step ------------------------------------------------------------------------------

function em_mvn_Estep!(𝔼x, 𝔼xxᵀ, em_model, observed, 𝔼x_pre, 𝔼xxᵀ_pre)

    𝔼x .= 0.0
    𝔼xxᵀ .= 0.0

    𝔼xᵢ = copy(𝔼x)
    𝔼xxᵀᵢ = copy(𝔼xxᵀ)

    μ = em_model.μ
    Σ = em_model.Σ

    # Compute the expected sufficient statistics
    for i in 2:length(observed.pattern_n_obs)

        # observed and unobserved vars
        u = observed.patterns_not[i]
        o = observed.patterns[i]

        # precompute for pattern
        V = Σ[u, u] - Σ[u, o] * (Σ[o, o]\Σ[o, u])

        # loop trough data
        for row in observed.rows[i]
            m = μ[u] + Σ[u, o] * ( Σ[o, o] \ (observed.data_rowwise[row]-μ[o]) )

            𝔼xᵢ[u] = m
            𝔼xᵢ[o] = observed.data_rowwise[row]
            𝔼xxᵀᵢ[u, u] = 𝔼xᵢ[u] * 𝔼xᵢ[u]' + V
            𝔼xxᵀᵢ[o, o] = 𝔼xᵢ[o] * 𝔼xᵢ[o]'
            𝔼xxᵀᵢ[o, u] = 𝔼xᵢ[o] * 𝔼xᵢ[u]'
            𝔼xxᵀᵢ[u, o] = 𝔼xᵢ[u] * 𝔼xᵢ[o]'

            𝔼x .+= 𝔼xᵢ
            𝔼xxᵀ .+= 𝔼xxᵀᵢ
        end

    end

    𝔼x .+= 𝔼x_pre
    𝔼xxᵀ .+= 𝔼xxᵀ_pre

end

function em_mvn_Mstep!(em_model, n_obs, 𝔼x, 𝔼xxᵀ)

    em_model.μ = 𝔼x/n_obs;
    Σ = Symmetric(𝔼xxᵀ/n_obs - em_model.μ*em_model.μ')

    # ridge Σ
    # while !isposdef(Σ)
    #     Σ += 0.5I
    # end

    em_model.Σ = Σ

    # diagonalization
    #if !isposdef(Σ)
    #    print("Matrix not positive definite")
    #    em_model.Σ .= 0
    #    em_model.Σ[diagind(em_model.Σ)] .= diag(Σ)
    #else
        # em_model.Σ = Σ
    #end

    return nothing
end

# generate starting values --------------------------------------------------------------

# use μ and Σ of full cases
function start_em_observed(observed::SemObsMissing; kwargs...)

    if (length(observed.patterns[1]) == observed.n_man) & (observed.pattern_n_obs[1] > 1)
        μ = copy(observed.obs_mean[1])
        Σ = copy(Symmetric(observed.obs_cov[1]))
        if !isposdef(Σ)
            Σ = Matrix(Diagonal(Σ))
        end
        return EmMVNModel(Σ, μ, false)
    else
        return start_em_simple(observed, kwargs...)
    end

end

# use μ = O and Σ = I
function start_em_simple(observed::SemObsMissing; kwargs...)
    n_man = Int(observed.n_man)
    μ = zeros(n_man)
    Σ = rand(n_man, n_man); Σ = Σ*Σ'
    # Σ = Matrix(1.0I, n_man, n_man)
    return EmMVNModel(Σ, μ, false)
end

# set to passed values
function start_em_set(observed::SemObsMissing; model_em, kwargs...)
    return em_model
end


using Primes

function count_prime_sums(limit::Int)
    p_list = primes(limit)
    counts = Dict(0=>1)

    for p in p_list
        for i in p:limit
            counts[i] = get(counts, i, 0) + get(counts, i - p, 0)
        end
    end

    for i in 2:limit
        if counts[i] > 5000
            return i
        end
    end
end

println(count_prime_sums(100))


using ExploratoryDataAnalysis
using DataFrames
using CategoricalArrays
# using Plots

df = DataFrame(t = [fill(0, 500); fill(1, 500)], x1 = collect(1:1000), x2 = randn(1000))

df.tc = categorical(df.t)

eda(df, :t)

eda(df, :x1)

f = [9800 9925 11825; 190200 190075 188175]'
infovalue(f)
mutualinfo(f)


function calc_rho( pw::PWGrid, Focc, psi::Array{ComplexF64,2} )
    Ω  = pw.Ω
    Ns = pw.Ns
    Npoints = pw.Npoints
    Nstates = size(psi)[2]
    #
    rho = zeros(ComplexF64,Npoints)

    # Transform to real space
    psiR = G_to_R(Ns,psi)

    # orthonormalization in real space
    ortho_gram_schmidt!( Nstates, psiR )
    psiR = sqrt(Npoints/Ω)*psiR

    for is = 1:Nstates
        for ip = 1:Npoints
            rho[ip] = rho[ip] + conj(psiR[ip,is])*psiR[ip,is]*Focc[is]
        end
    end
    return real(rho)
end


export rsvd_fnkz

using LinearAlgebra: SVD

struct OuterProduct{T}
    X :: Matrix{T}
    Y :: Matrix{T}
end

X ⊗ Y = OuterProduct{promote_type(eltype(X), eltype(Y))}(X, Y)

"""
    rsvd_fnkz(A, k)

Compute the randomized SVD by iterative refinement from randomly selected
columns/rows [^Friedland2006].

# Arguments

- `A`: matrix whose SVD is desired;
- `k::Int`: desired rank of approximation (`k ≤ min(m, n)`).

## Keywords

- `l::Int = k`: number of columns/rows to sample at each iteration (`1 ≤ l ≤ k`);
- `N::Int = minimum(size(A))`: maximum number of iterations;
- `ϵ::Real = prod(size(A))*eps()`: relative threshold for convergence, as
  measured by growth of the spectral norm;
- `method::Symbol = :eig`: problem to solve.
  1. `:eig`: eigenproblem.
  2. `:svd`: singular problem.
- `verbose::Bool = false`: print convergence information at each iteration.

# Return value

SVD object of `rank ≤ k`.

[^Friedland2006]:
    Friedland, Shmuel, et al. "Fast Monte-Carlo low rank approximations for
    matrices." System of Systems Engineering, 2006 IEEE/SMC International
    Conference on. IEEE, 2006.
"""
function rsvd_fnkz(A, k::Int;
    l::Int=k, N::Int=minimum(size(A)), verbose::Bool=false,
    method::Symbol=:eig,
    ϵ::Real=prod(size(A))*eps(real(float(one(eltype(A))))))

    m, n = size(A)
    dosvd = method == :svd
    @assert 1 ≤ l ≤ k ≤ min(m, n)

    #Initialize k-rank approximation B₀ according to (III.9)
    tallandskinny = m ≥ n
    B₀ = if tallandskinny
        #Select k columns of A
        π = randperm(n)[1:k]
        A[:, π]
    else
        warn("Not implemented")
        #Select k rows of A
        π = randperm(m)[1:k]
        A[π, :]
    end
    QQQ, RRR = qr(B₀)
    X = Matrix(QQQ)[:, abs.(diag(RRR)) .> ϵ] #Remove linear dependent columns
    B = tallandskinny ? X ⊗ A'X : X ⊗ X'A

    oldnrmB = 0.0
    Λ = 0
    X̃ = 0
    V = 0
    for t=1:N
        π = randperm(k)[1:l]
        #Update B using Theorem 2.4
        QQQ, RRR = qr([B.X A[:, π]]) #We are doing more work than needed here
        X = Matrix(QQQ)[:, abs.(diag(RRR)) .> ϵ] #Remove linearly dependent columns
        Y = A'X
        if dosvd
            S = svd!(Y)
            Λ = S.S.^2
            B = (X * S.V) ⊗ (S.U*Diagonal(S.S))
        else
            E = eigen!(Symmetric(Y'Y))
            π = sortperm(E.values, rev=true)[1:min(length(E.values), k)]
            Λ = E.values[π]
            O = E.vectors[:, π] #Eq. 2.6
            X̃ = X * O
            B = X̃ ⊗ A'X̃
        end

        length(Λ)==0 && break

        nrmB = √Λ[1]
        verbose && println("iteration $t: Norm: $nrmB")
        oldnrmB > (1-ϵ)*nrmB && break
        oldnrmB = nrmB
    end
    for i=1:size(B.Y, 2)
        rmul!(view(B.Y, :, i), 1/√Λ[i])
    end
    SVD(B.X, sqrt.(Λ), B.Y')
end


using DelimitedFiles

function solve(input)
    raw = readdlm(input,'\n', Int16)
    serial_number=raw[1]
    @show search(serial_number,300,3:3)
    @show search(serial_number,300,1:300)
end

function compute_power_level(x,y,serial_number)
    rack_id = x + 10
    power_level = rack_id * y
    power_level += serial_number
    power_level *= rack_id
    power_level = floor(power_level%1000/100)
    power_level -= 5
end

@assert compute_power_level(3,5,8) == 4
@assert compute_power_level(122,79,57) == -5
@assert compute_power_level(217,196,39) == 0
@assert compute_power_level(101,153,71) == 4

function search(serial_number, grid_size, square_sizes)
    power_sum = zeros(Int32,grid_size,grid_size)
    x_max, y_max, square_size_max = 0,0,0
    for x in 2:grid_size
        for y in 2:grid_size
            power=compute_power_level(x,y,serial_number)
            power_sum[x,y] = power + power_sum[x,y-1] + power_sum[x-1,y] - power_sum[x-1,y-1];
        end
    end
    largest_total_power = 0
    for s in square_sizes
        for x in s+1:grid_size
            for y in s+1:grid_size
                #current = sum(grid[x:x+s-1,y:y+s-1])
                current = power_sum[x,y] - power_sum[x,y-s] - power_sum[x-s,y] + power_sum[x-s,y-s]
                if current > largest_total_power
                    largest_total_power = current
                    x_max = x
                    y_max = y
                    square_size_max = s
                end
            end
        end
    end
    return (x_max-square_size_max+1,y_max-square_size_max+1,square_size_max)
end
@assert search(18,300,3:3) == (33,45,3)
@assert search(42,300,3:3) == (21,61,3)
@assert search(18,300,1:300) == (90,269,16)
@assert search(42,300,1:300) == (232,251,12)

solve("input.txt")


function discount!(rs::AbstractVector{T}, γ::T) where T <: AbstractFloat
    @inbounds for i in length(rs)-1:-1:1
        rs[i] += γ * rs[i+1]
    end
end

function nstepr!(rs::AbstractVector{T}, brs::AbstractVector{T}, γ::T, n::Int64) where T <: AbstractFloat
    @inbounds for t in 1:length(rs)
        runningadd = zero(T)
        l = min(n, length(rs)-t)
        for i in 0:l
            runningadd += γ^(i) * (i == n ? brs[t+i] : rs[t+i])
        end
        rs[t] = runningadd
    end
end

function GAEλ!(vs::AbstractVector{T}, rs::AbstractVector{T}, γ::T, λ::T) where T <: AbstractFloat
    vs[1:end-1] .= rs[1:end-1] .+ γ .* vs[2:end] .- vs[1:end-1]
    vs[end] = rs[end] - vs[end]
    discount!(vs, γ*λ)
end

function norm!(vs::AbstractVector{T}) where T <: AbstractFloat
    θmean = mean(vs)
    θstd = std(vs, mean=θmean)
    vs .= (vs .- θmean) ./ (θstd + T(1f-12))
end


import GeometryBasics: Sphere, decompose, normals

#-----------------------------Generated geometries------------------------#
function generate_sphere(complexity = 2)
    X = 0.525731112f0
    Z = 0.850650808f0
    vertices = Point3f0[(-X, 0, Z), (X, 0, Z), (-X, 0, -Z), (X, 0, -Z), (0, Z, X),
                        (0, Z, -X), (0, -Z, X), (0, -Z, -X), (Z, X, 0), (-Z, X, 0),
                        (Z, -X, 0), (-Z, -X, 0)]

    faces = map(x -> x .+ 1,
                TriangleFace{Int32}[(1, 4, 0), (4, 9, 0), (4, 5, 9), (8, 5, 4), (1, 8, 4),
                                    (1, 10, 8), (10, 3, 8), (8, 3, 5), (3, 2, 5), (3, 7, 2),
                                    (3, 10, 7), (10, 6, 7), (6, 11, 7), (6, 0, 11),
                                    (6, 1, 0), (10, 1, 6), (11, 0, 9), (2, 11, 9),
                                    (5, 2, 9), (11, 2, 7)])

    outverts, outfaces = subdivide(vertices, faces, complexity)
    normals = Vec3f0.(copy(outverts))
    return outverts, normals, outfaces
end

function subdivide(vertices, faces, level)
    outvertices = Point3f0[]
    outfaces    = TriangleFace{Int32}[]
    newfaces    = TriangleFace{Int32}[]
    newvertices = Point3f0[]
    if level > 0
        for face in faces
            v1 = vertices[face[1]]
            v2 = vertices[face[2]]
            v3 = vertices[face[3]]

            v4 = normalize((v1 + v2) / 2.0)
            v5 = normalize((v2 + v3) / 2.0)
            v6 = normalize((v1 + v3) / 2.0)

            vertlen = length(newvertices)
            push!(newfaces, TriangleFace{Int32}(vertlen + 1, vertlen + 4, vertlen + 6))
            push!(newfaces, TriangleFace{Int32}(vertlen + 4, vertlen + 2, vertlen + 5))
            push!(newfaces, TriangleFace{Int32}(vertlen + 5, vertlen + 3, vertlen + 6))
            push!(newfaces, TriangleFace{Int32}(vertlen + 4, vertlen + 5, vertlen + 6))
            push!(newvertices, v1, v2, v3, v4, v5, v6)
        end
        vertices, faces = subdivide(newvertices, newfaces, level - 1)
    end
    return vertices, faces
end

function sphere(dio::Diorama, pos, radius, complexity = 2, attributes_...; uniforms...)
    sphpos = convert(Point3f0, pos)
    sphrad = convert(Float32, radius)
    modelmat = translmat(sphpos) * scalemat(Point3f0(sphrad, sphrad, sphrad))

    verts, norms, faces = Sphere(complexity)
    atdict = SymAnyDict(attributes_)
    unidict = SymAnyDict(uniforms)
    if haskey(unidict, :color)
        colors = fill(pop!(unidict, :color), length(verts))
    elseif haskey(atdict, :color)
        colors = pop!(atdict, :color)
    elseif haskey(atdict, :colors)
        colors = pop!(atdict, :colors)
    else
        colors = fill(RGB{Float32}(0, 0, 0), length(verts))
    end
    unidict[:modelmat] = modelmat
    mesh = AttributeMesh((color = colors,), BasicMesh(verts, faces, norms))
    sphrend = InstancedMeshRenderable(:sphere, mesh, unidict, Dict(:default => false),
                                      false)
    add!(dio, sphrend)
    return sphrend
end

#-----------------------------OBJ based geometries------------------------#

loadobj(filename::String) = load(joinpath(@__DIR__, "../../assets/obj", filename))

function cylinder(dio::Diorama, startpos, endpos, radius, name = "cylinder", attributes...;
                  uniforms...)
    startp = convert(Point3f0, startpos)
    endp   = convert(Point3f0, endpos)
    cylrad = convert(Float32, radius)
    rotmat = rotate(startp, endp)
    scalm  = scalemat(Point3f0(cylrad, cylrad, norm(endp - startp)))
    tmat   = translmat(startp)

    modelmat = tmat * rotmat * scalm

    cylmesh = loadobj("cylinder.obj")

    atdict = SymAnyDict(attributes)
    unidict = SymAnyDict(uniforms)
    if haskey(unidict, :color)
        atdict[:color] = fill(pop!(unidict, :color), length(cylmesh.vertices))
    else
        atdict[:color] = fill(RGB{Float32}(0, 0, 0), length(cylmesh.vertices))
    end
    unidict[:modelmat] = modelmat

    cylrend = MeshRenderable(0, name, cylmesh, atdict...; unidict...)
    add!(dio, cylrend)
    return cylrend
end

function rectangle(dio::Diorama, startpos, endpos, widths, name = "rectangle",
                   attributes...; uniforms...)
    startp = convert(Point3f0, startpos)
    endp   = convert(Point3f0, endpos)
    rotmat = rotate(startp, endp)
    scalm  = scalemat(Point3f0(widths[1], widths[2], norm(endp - startp)))
    tmat   = translmat(startp)

    modelmat = tmat * rotmat * scalm

    cubmesh = loadobj("cube.obj")

    atdict = SymAnyDict(attributes)
    unidict = SymAnyDict(uniforms)
    if haskey(unidict, :color)
        atdict[:color] = fill(pop!(unidict, :color), length(cubmesh.vertices))
    else
        atdict[:color] = fill(RGB{Float32}(0, 0, 0), length(cubmesh.vertices))
    end
    unidict[:modelmat] = modelmat
    cubrend = MeshRenderable(0, name, cubmesh, atdict...; unidict...)
    add!(dio, cubrend)
    return cubrend
end

function cone(dio::Diorama, startpos, endpos, radius, name = "cone", attributes...;
              uniforms...)
    startp = convert(Point3f0, startpos)
    endp   = convert(Point3f0, endpos)
    cylrad = convert(Float32, radius)
    rotmat = rotate(startp, endp)
    scalm  = scalemat(Point3f0(cylrad, cylrad, norm(endp - startp)))
    tmat   = translmat(startp)

    modelmat = tmat * rotmat * scalm

    cylmesh = loadobj("cone.obj")

    atdict = SymAnyDict(attributes)
    unidict = SymAnyDict(uniforms)
    if haskey(unidict, :color)
        atdict[:color] = fill(pop!(unidict, :color), length(cylmesh.vertices))
    else
        atdict[:color] = fill(RGB{Float32}(0, 0, 0), length(cylmesh.vertices))
    end
    unidict[:modelmat] = modelmat

    cylrend = MeshRenderable(0, name, cylmesh, atdict...; unidict...)
    add!(dio, cylrend)
    return cylrend
end

function arrow(dio::Diorama, startpos, endpos, rad1, rad2, name = "arrow",
               headratio = 1 / 4, attributes...; uniforms...)
    startp = convert(Point3f0, startpos)
    endp   = convert(Point3f0, endpos)

    rotmat   = rotate(startp, endp)
    scalm    = scalemat(Point3f0(rad1, rad1, norm(endp - startp)))
    tmat     = translmat(startp)
    modelmat = tmat * rotmat * scalm
    conemesh = loadobj("cone.obj")
    cylmesh  = loadobj("cylinder.obj")

    cyllen = norm(endp - startp) * (1 - headratio)
    conelen = norm(endp - startp) * headratio

    conetrans = translmat(Point3f0(0, 0, cyllen))
    conescale = scalemat(Point3f0(rad2 / rad1, rad2 / rad1, conelen))
    conemat = conetrans * conescale
    coneverts = [Point3f0((conemat*Vec4f0(v[1], v[2], v[3], 1.0f0))[1:3]...)
                 for v in conemesh.vertices]
    conefaces = [v .+ Int32(length(cylmesh.vertices) + 1) for v in conemesh.faces]
    allverts = Point3f0.([cylmesh.vertices; coneverts])
    allnorms = [cylmesh.normals; conemesh.normals]
    allfaces = TriangleFace{Int32}.([cylmesh.faces; conefaces])

    atdict = SymAnyDict(attributes)
    unidict = SymAnyDict(uniforms)
    if haskey(unidict, :color)
        atdict[:color] = fill(pop!(unidict, :color), length(allverts))
    else
        atdict[:color] = fill(RGB{Float32}(0, 0, 0), length(allverts))
    end
    unidict[:modelmat] = modelmat
    atdict[:vertices] = allverts
    atdict[:normals] = allnorms
    atdict[:faces] = allfaces
    mesh = homogenousmesh(atdict)
    arrowrend = MeshRenderable(0, name, mesh; unidict...)
    add!(dio, arrowrend)
    return arrowrend
end


using LinearAlgebra
using Test

using ProximalOperators
using ProximalAlgorithms

@testset "Lasso small ($T)" for T in [Float32, Float64, ComplexF32, ComplexF64]
    A = T[
        1.0 -2.0 3.0 -4.0 5.0
        2.0 -1.0 0.0 -1.0 3.0
        -1.0 0.0 4.0 -3.0 2.0
        -1.0 -1.0 -1.0 1.0 3.0
    ]
    b = T[1.0, 2.0, 3.0, 4.0]

    m, n = size(A)

    R = real(T)

    lam = R(0.1) * norm(A' * b, Inf)
    @test typeof(lam) == R

    f = Translate(SqrNormL2(R(1)), -b)
    f2 = LeastSquares(A, b)
    g = NormL1(lam)

    x_star = T[-3.877278911564627e-01, 0, 0, 2.174149659863943e-02, 6.168435374149660e-01]

    TOL = R(1e-4)

    @testset "ForwardBackward" begin

        ## Nonfast/Nonadaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.ForwardBackward(tol = TOL)
        x, it = solver(x0, f = f, A = A, g = g, Lf = opnorm(A)^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 150
        @test x0 == x0_backup

        # Nonfast/Adaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.ForwardBackward(tol = TOL, adaptive = true)
        x, it = solver(x0, f = f, A = A, g = g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 300
        @test x0 == x0_backup

        # Fast/Nonadaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.ForwardBackward(tol = TOL, fast = true)
        x, it = solver(x0, f = f, A = A, g = g, Lf = opnorm(A)^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 100
        @test x0 == x0_backup

        # Fast/Adaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver =
            ProximalAlgorithms.ForwardBackward(tol = TOL, adaptive = true, fast = true)
        x, it = solver(x0, f = f, A = A, g = g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 200
        @test x0 == x0_backup
    end

    @testset "ZeroFPR" begin

        # ZeroFPR/Nonadaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.ZeroFPR(tol = TOL)
        x, it = solver(x0, f = f, A = A, g = g, Lf = opnorm(A)^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20
        @test x0 == x0_backup

        # ZeroFPR/Adaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.ZeroFPR(adaptive = true, tol = TOL)
        x, it = solver(x0, f = f, A = A, g = g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20
        @test x0 == x0_backup

    end

    @testset "PANOC" begin

        # PANOC/Nonadaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.PANOC(tol = TOL)
        x, it = solver(x0, f = f, A = A, g = g, Lf = opnorm(A)^2)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20
        @test x0 == x0_backup

        ## PANOC/Adaptive

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.PANOC(adaptive = true, tol = TOL)
        x, it = solver(x0, f = f, A = A, g = g)
        @test eltype(x) == T
        @test norm(x - x_star, Inf) <= TOL
        @test it < 20
        @test x0 == x0_backup

    end

    @testset "DouglasRachford" begin

        # Douglas-Rachford

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver =
            ProximalAlgorithms.DouglasRachford(gamma = R(10) / opnorm(A)^2, tol = TOL)
        y, z, it = solver(x0, f = f2, g = g)
        @test eltype(y) == T
        @test eltype(z) == T
        @test norm(y - x_star, Inf) <= TOL
        @test norm(z - x_star, Inf) <= TOL
        @test it < 30
        @test x0 == x0_backup

    end

    @testset "DouglasRachford line search" begin

        # Douglas-Rachford line search

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.DRLS(tol = 10 * TOL)
        y, z, it = solver(x0, f = f2, g = g, Lf = opnorm(A)^2)
        @test eltype(y) == T
        @test eltype(z) == T
        @test norm(y - x_star, Inf) <= 10 * TOL
        @test norm(z - x_star, Inf) <= 10 * TOL
        @test it < 26
        @test x0 == x0_backup

    end

    @testset "AFBA" begin

        x0 = zeros(T, n)
        x0_backup = copy(x0)

        solver = ProximalAlgorithms.AFBA(theta = R(1), mu = R(1), tol = R(1e-6))
        x_afba, y_afba, it_afba = solver(x0, zeros(T, n), f = f2, g = g, beta_f = opnorm(A)^2)
        @test eltype(x_afba) == T
        @test eltype(y_afba) == T
        @test norm(x_afba - x_star, Inf) <= 1e-4
        @test it_afba <= 80
        @test x0 == x0_backup

        solver = ProximalAlgorithms.AFBA(theta = R(1), mu = R(1), tol = R(1e-6))
        x_afba, y_afba, it_afba = solver(x0, zeros(T, n), f = f2, h = g, beta_f = opnorm(A)^2)
        @test eltype(x_afba) == T
        @test eltype(y_afba) == T
        @test norm(x_afba - x_star, Inf) <= 1e-4
        @test it_afba <= 100
        @test x0 == x0_backup

        solver = ProximalAlgorithms.AFBA(theta = R(1), mu = R(1), tol = R(1e-6))
        x_afba, y_afba, it_afba = solver(x0, zeros(T, m), h = f, L = A, g = g)
        @test eltype(x_afba) == T
        @test eltype(y_afba) == T
        @test norm(x_afba - x_star, Inf) <= 1e-4
        @test it_afba <= 150
        @test x0 == x0_backup

    end

    @testset "SFISTA" begin

        # SFISTA

        x0 = zeros(T, n)
        x0_backup = copy(x0)
        solver = ProximalAlgorithms.SFISTA(tol = 10 * TOL)
        y, it = solver(x0, f = f2, h = g, Lf = opnorm(A)^2)
        @test eltype(y) == T
        @test norm(y - x_star, Inf) <= 10 * TOL
        @test it < 100
        @test x0 == x0_backup

    end

end


"Abstract Lyapunov function type"
abstract type LyapunovFunction <: CertificateFunction end

"""
    ControlLyapunovFunction

Control Lyapunov function V for a control affine system
# Fields
- `V`: function V(x) that represents the CLF
- `∇V`: function ∇V(x) defining the gradient of the CLF
- `α`: class K function α(V(x))
"""
struct ControlLyapunovFunction <: LyapunovFunction
    V
    ∇V
    α
end

"""
    ControlLyapunovFunction(V, α)

Construct a CLF given V(x) and class K function α(V(x)).
"""
function ControlLyapunovFunction(V, α)
    ∇V(x) = gradient(V, x)[1]'
    return ControlLyapunovFunction(V, ∇V, α)
end

"""
	(V::ControlLyapunovFunction)(x)

Evaluate CLF at state x.
"""
function (V::ControlLyapunovFunction)(x)
    return V.V(x)
end


using DifferentialEquations
using Plots; gr()
H = <HAMILTONIAN> :: Array{Complex{Float64},2}
f(u,p,t) = im*(u*H - H*u)
u0 = <START> :: Array{Complex{Float64},2}
tspan = (0.0,1.0)
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)
t = sol.t
u = sol.u
#pl1 = plot(sol)
#gui()


##############################################################################
# JuMP
# An algebraic modelling langauge for Julia
# See http://github.com/JuliaOpt/JuMP.jl
#############################################################################
# cluster.jl
#
# From "Approximating K-means-type clustering via semidefinite programming"
# By Jiming Peng and Yu Wei
#
# Given a set of points a_1, ..., a_m  in R_n, allocate them to k clusters
#############################################################################

using JuMP
using Mosek

# Data points
n = 2
m = 6
a = Any[[2.0, 2.0], [2.5, 2.1], [7.0, 7.0],
        [2.2, 2.3], [6.8, 7.0], [7.2, 7.5]]
k = 2

# Weight matrix
W = zeros(m,m)
for i = 1:m
    for j = i+1:m
        dist = exp(-norm(a[i] - a[j])/1.0)
        W[i,j] = dist
        W[j,i] = dist
    end
end

mod = Model()

# Z >= 0, PSD
@variable(mod, Z[1:m,1:m], SDP)
@constraint(mod, Z .>= 0)

# min Tr(W(I-Z))
@objective(mod, Min, trace(W * (eye(m,m) - Z)))

# Z e = e
@constraint(mod, Z*ones(m) .== ones(m))
#for i = 1:m
#    addConstraint(mod, sum([ Z[i,j] for j = 1:m]) == 1)
#end

# Tr(Z) = k
@constraint(mod, trace(Z) == k)

solve(mod)

Z_val = getvalue(Z)[:,:]
println("Raw solution")
println(round(Z_val,4))

# A simple rounding scheme
which_cluster = zeros(Int,m)
num_clusters = 0
for i = 1:m
    Z_val[i,i] <= 1e-6 && continue

    if which_cluster[i] == 0
        num_clusters += 1
        which_cluster[i] = num_clusters
        for j = i+1:m
            if norm(Z_val[i,j] - Z_val[i,i]) <= 1e-6
                which_cluster[j] = num_clusters
            end
        end
    end
end

# Print results
for cluster = 1:k
    println("Cluster $cluster")
    for i = 1:m
        if which_cluster[i] == cluster
            println(a[i])
        end
    end
end


## This also tests some vanilla pricing engines ##
using Base.Test
using QuantLib

cal = QuantLib.Time.TargetCalendar()
todaysDate = Date(1998, 5, 15)
settlementDate = Date(1998, 5, 17)

set_eval_date!(settings, todaysDate)

optionType = Put()
underlying = 36.0
strike = 40.0
dividendYield = 0.00
riskFreeRate = 0.06
vol = 0.20
mat = Date(1999, 5, 17)
dc = QuantLib.Time.Actual365()

exerciseDates = Date[settlementDate + Dates.Month(3 * i) for i = 1:4]

europeanExercise = EuropeanExercise(mat)
bermudanExercise = BermudanExercise(exerciseDates)
americanExercise = AmericanExercise(settlementDate, mat)

underlyingH = Quote(underlying)

flatTermStructure = FlatForwardTermStructure(settlementDate, cal, Quote(riskFreeRate), dc)
flatDividendTS = FlatForwardTermStructure(settlementDate, cal, Quote(dividendYield), dc)
flatVolTS = BlackConstantVol(settlementDate, cal, vol, dc)

payoff = PlainVanillaPayoff(optionType, strike)

bsmProcess = BlackScholesMertonProcess(underlyingH, flatTermStructure, flatDividendTS, flatVolTS)

# test a European Option
bsPE = AnalyticEuropeanEngine(bsmProcess) # black scholes PE
europeanOption = VanillaOption(payoff, europeanExercise, bsPE)

# American Option
baroneAdesiPE = BaroneAdesiWhaleyApproximationEngine(bsmProcess)
americanOption = VanillaOption(payoff, americanExercise, baroneAdesiPE)

# Bermudan Option
timeSteps = 801
fdBermudanPE = FDBermudanEngine(bsmProcess, CrankNelson, timeSteps, timeSteps - 1)
bermudanOption = VanillaOption(payoff, bermudanExercise, fdBermudanPE)

@test npv(europeanOption) == 3.8443077915968398
@test npv(americanOption) == 4.459627613776478
@test npv(bermudanOption) == 4.3608068625088645


using Poisson2D
using JLD2
using Printf

# set problem parameters
a = 0;
b = 1;
c = 0;
d = 1;
n = 999;

kx = 5;
ky = 4;
f = (x,y)->((kx*π/(b-a))^2 + (ky * π/(d-c))^2) * sin(kx*π*(x-a)/(b-a)) * sin(ky*π*(y-c)/(d-c));

# set up and solve problem
problem = FDPoisson2DProblem(a,b,c,d,n,f);
assemble_system!(problem);
U = direct_solve_poisson2d(problem);

# name of file
fname = @sprintf("soln1_n%d.jld2", n);

# save to disk
jldsave(fname; problem, U);


### A Pluto.jl notebook ###
# v0.12.4

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ d2bf60aa-1a49-11eb-3706-5f7507fc511c
begin
	using Images
	using PlutoUI
	using ImageFiltering
	using Statistics
	using LinearAlgebra
end

# ╔═╡ b2a8e540-19af-11eb-0214-9d5f04ac2486
md"""
# Seam Carving

1. We use convolution with Sobel filters for "edge detection".
2. We use that to write an algorithm that removes "uninteresting" bits of an image in order to shrink it.
"""

# ╔═╡ 455ad8f4-1a4f-11eb-38aa-b9da92a25c17
md"""
## Utility functions
"""

# ╔═╡ 988f2d30-1a50-11eb-0a3d-a15827ed7248
function hbox(x, y, gap=16; sy=size(y), sx=size(x))
	w,h = (max(sx[1], sy[1]),
		   gap + sx[2] + sy[2])

	slate = fill(RGB(1,1,1), w,h)
	slate[1:size(x,1), 1:size(x,2)] .= RGB.(x)
	slate[1:size(y,1), size(x,2) + gap .+ (1:size(y,2))] .= RGB.(y)
	slate
end

# ╔═╡ 14fc8df0-1a64-11eb-13af-cdd3cfc72f2e
vbox(x,y, gap=16) = hbox(x', y')'

# ╔═╡ d5fa6566-1a4f-11eb-0528-a35576836b80
function shrink_image(image, ratio=5)
	height, width = size(image)
	new_height = height ÷ ratio - 1
	new_width = width ÷ ratio - 1
	list = [
		mean(image[
				ratio * i:ratio * (i + 1),
				ratio * j:ratio * (j + 1)
		])
		for i in 1:new_height, j in 1:new_width
	]
	#reshape(list, new_height, new_width)
end

# ╔═╡ ba4af658-1a4f-11eb-217a-7f9165d30dbd
function convolve(M, kernel)
	height, width = size(kernel)

	half_height = height ÷ 2
	half_width = width ÷ 2

	new_image = similar(M)

	# (i, j) loop over the original image
	m, n = size(M)
	@inbounds for i in 1:m
		for j in 1:n
			# (k, l) loop over the neighbouring pixels
			accumulator = 0 * M[1, 1]
			for k in -half_height:-half_height + height - 1
				for l in -half_width:-half_width + width - 1
					Mi = i - k
					Mj = j - l
					# First index into M
					if Mi < 1
						Mi = 1
					elseif Mi > m
						Mi = m
					end
					# Second index into M
					if Mj < 1
						Mj = 1
					elseif Mj > n
						Mj = n
					end

					accumulator += kernel[k, l] * M[Mi, Mj]
				end
			end
			new_image[i, j] = accumulator
		end
	end

	return new_image
end

# ╔═╡ 571690f6-1a4f-11eb-2cb7-8d75d2dfc5eb
function show_colored_array(array)
	pos_color = RGB(0.36, 0.82, 0.8)
	neg_color = RGB(0.99, 0.18, 0.13)
	to_rgb(x) = max(x, 0) * pos_color + max(-x, 0) * neg_color
	to_rgb.(array) / maximum(abs.(array))
end

# ╔═╡ 53bf61d0-19b4-11eb-2212-d91065dae28f
@bind image_url Select([
"https://cdn.shortpixel.ai/spai/w_1086+q_lossy+ret_img+to_webp/https://wisetoast.com/wp-content/uploads/2015/10/The-Persistence-of-Memory-salvador-deli-painting.jpg",
"https://upload.wikimedia.org/wikipedia/commons/thumb/1/17/Gustave_Caillebotte_-_Paris_Street%3B_Rainy_Day_-_Google_Art_Project.jpg/1014px-Gustave_Caillebotte_-_Paris_Street%3B_Rainy_Day_-_Google_Art_Project.jpg",
"https://web.mit.edu/facilities/photos/construction/Projects/stata/1_large.jpg",
])

# ╔═╡ 53d99a46-19b4-11eb-21fa-859486e6289b
img = load(download(image_url))

# ╔═╡ 53f647ea-19b4-11eb-2356-755781371a3a
# arbitrarily choose the brightness of a pixel as a mean of rgb
brightness(c::AbstractRGB) = 0.5 * c.r + 0.59 * c.g + 0.11 * c.b

# ╔═╡ 65932dac-1a67-11eb-2b1d-a74cd795c111
function edgeness(img)
	Sy, Sx = Kernel.sobel()
	b = brightness.(img)
	∇y = convolve(b, Sy)
	∇x = convolve(b, Sx)
	sqrt.(∇x.^2 + ∇y.^2)
end

# ╔═╡ 5413928c-19b4-11eb-1890-230a8077f4ae
Gray.(img)

# ╔═╡ 542f1244-19b4-11eb-207c-131b5699ad24
Gray.(brightness.(img))

# ╔═╡ ec7a7c5e-1a52-11eb-0b3f-49a1c0c2ab4b
shrink_image(img)

# ╔═╡ 2b5e3d8c-1a64-11eb-3a64-cb4e1c314db2
show_colored_array(brightness.(img))

# ╔═╡ ad7c6f50-1a4b-11eb-35fe-a7331d599b45
md"""
## Edge detection filter

We use the Sobel edge detection filter.

```math
\begin{align}

G_x &= \begin{bmatrix}
1 & 0 & -1 \\
2 & 0 & -2 \\
1 & 0 & -1 \\
\end{bmatrix}*A\\
G_y &= \begin{bmatrix}
1 & 2 & 1 \\
0 & 0 & 0 \\
-1 & -2 & -1 \\
\end{bmatrix}*A
\end{align}
```

Where: $A$ is the target image. We can think of these as derivates in the $x$ and $y$ directions.

Then we combine them by finding the magnitude of the **gradient** (in the sense of the multivatiate calculus) be defining.

$$G_\text{total} = \sqrt{G_x^2 + G_y^2}.$$
"""

# ╔═╡ ad6289f0-1a4b-11eb-1c1e-b1158e269574
Sy, Sx = Kernel.sobel();

# ╔═╡ ad4684c6-1a4b-11eb-3c63-5f888a1d834b
[show_colored_array(Sx) show_colored_array(Sy)]

# ╔═╡ ad2c6d34-1a4b-11eb-3f16-7332941d79fa
begin
	img_brightness = brightness.(img)
	∇x = convolve(img_brightness, Sx)
	∇y = convolve(img_brightness, Sy)
	hbox(show_colored_array(∇x), show_colored_array(∇y))
end

# ╔═╡ ad0d9044-1a4b-11eb-1a23-b7160b7793ed
vbox(
	hbox(img[200:end, 50:250], img[200:end, 50:250]),
	hbox(show_colored_array.((∇x[200:end, 50:250], ∇y[200:end, 50:250]))...)
)

# ╔═╡ ac812a78-1a4b-11eb-216d-91a51c6b4742
begin
	edged = edgeness(img)
	hbox(img, Gray.(edged) / maximum(abs.(edged)))
end

# ╔═╡ ac652cc4-1a4b-11eb-2fee-dd2887e80796
md"""
## Seam carving idea

The idea of seam carving is to find the path from the top of the image to the bottom of the imaga where the path minimized the edgness.

In other words, this path minimized the number of edges it crosses.

"""

# ╔═╡ ac48d830-1a4b-11eb-3296-4d12c4153138

# ╔═╡ ac28a5f6-1a4b-11eb-16e8-e18a949846cb

# ╔═╡ ac002a2c-1a4b-11eb-3c4e-cb541f081ff1

# ╔═╡ abe1468e-1a4b-11eb-00e2-a394e9ae1504

# ╔═╡ abbef57a-1a4b-11eb-2c48-b12aa928410f

# ╔═╡ Cell order:
# ╠═d2bf60aa-1a49-11eb-3706-5f7507fc511c
# ╟─b2a8e540-19af-11eb-0214-9d5f04ac2486
# ╟─455ad8f4-1a4f-11eb-38aa-b9da92a25c17
# ╟─988f2d30-1a50-11eb-0a3d-a15827ed7248
# ╟─14fc8df0-1a64-11eb-13af-cdd3cfc72f2e
# ╟─d5fa6566-1a4f-11eb-0528-a35576836b80
# ╟─ba4af658-1a4f-11eb-217a-7f9165d30dbd
# ╟─571690f6-1a4f-11eb-2cb7-8d75d2dfc5eb
# ╟─65932dac-1a67-11eb-2b1d-a74cd795c111
# ╟─53bf61d0-19b4-11eb-2212-d91065dae28f
# ╠═53d99a46-19b4-11eb-21fa-859486e6289b
# ╠═53f647ea-19b4-11eb-2356-755781371a3a
# ╠═5413928c-19b4-11eb-1890-230a8077f4ae
# ╠═542f1244-19b4-11eb-207c-131b5699ad24
# ╠═ec7a7c5e-1a52-11eb-0b3f-49a1c0c2ab4b
# ╠═2b5e3d8c-1a64-11eb-3a64-cb4e1c314db2
# ╟─ad7c6f50-1a4b-11eb-35fe-a7331d599b45
# ╠═ad6289f0-1a4b-11eb-1c1e-b1158e269574
# ╠═ad4684c6-1a4b-11eb-3c63-5f888a1d834b
# ╟─ad2c6d34-1a4b-11eb-3f16-7332941d79fa
# ╟─ad0d9044-1a4b-11eb-1a23-b7160b7793ed
# ╟─ac812a78-1a4b-11eb-216d-91a51c6b4742
# ╠═ac652cc4-1a4b-11eb-2fee-dd2887e80796
# ╠═ac48d830-1a4b-11eb-3296-4d12c4153138
# ╠═ac28a5f6-1a4b-11eb-16e8-e18a949846cb
# ╠═ac002a2c-1a4b-11eb-3c4e-cb541f081ff1
# ╠═abe1468e-1a4b-11eb-00e2-a394e9ae1504
# ╠═abbef57a-1a4b-11eb-2c48-b12aa928410f


#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

using HPAT
using MPI
using DocOpt

#CompilerTools.OptFramework.set_debug_level(3)
#CompilerTools.CFGs.set_debug_level(3)
#CompilerTools.Loops.set_debug_level(3)
#ParallelAccelerator.DomainIR.set_debug_level(3)
#ParallelAccelerator.ParallelIR.set_debug_level(3)
#ParallelAccelerator.CGen.set_debug_level(3)
#HPAT.Checkpointing.set_debug_level(3)
#HPAT.CGenPatternMatch.set_debug_level(3)
#HPAT.set_debug_level(3)
#HPAT.Checkpointing.setCheckpointDebug(50)    # Do checkpoints every 50 seconds.
#CompilerTools.LivenessAnalysis.set_debug_level(5)

@acc hpat_checkpoint function logistic_regression(iterations, file_name)
    points = DataSource(Matrix{Float32},HDF5,"/points", file_name)
    responses = DataSource(Vector{Float32},HDF5,"/responses", file_name)
    D = size(points,1) # number of features
    N = size(points,2) # number of instances
    labels = reshape(responses,1,N)
    w = reshape(2.0f0.*rand(Float32,D)-1.0f0,1,D)

    for i in 1:iterations
       w -= ((1.0f0./(1.0f0.+exp(-labels.*(w*points))).-1.0f0).*labels)*points'
    end
    w
end

function main()
    doc = """Logistic regression statistical method.

Usage:
  logistic_regression.jl -h | --help
  logistic_regression.jl [--iterations=<iterations>] [--file=<file>] [--restart]

Options:
  -h --help                  Show this screen.
  --iterations=<iterations>  Specify number of iterations; defaults to 20.
  --file=<file>              Specify input file; defaults to HPAT's default generated data file.
  --restart                  Restart the program from the last checkpoint taken.

"""
    arguments = docopt(doc)
    iterations = 20
    file_name = HPAT.getDefaultDataPath()*"logistic_regression.hdf5"

    if (arguments["--iterations"] != nothing)
        iterations = parse(Int, arguments["--iterations"])
    end

    if (arguments["--file"] != nothing)
        file_name = arguments["--file"]
    end

    do_restart = arguments["--restart"]

    if do_restart
        W = HPAT.restart(logistic_regression, iterations, file_name)
    else
        W = logistic_regression(iterations, file_name)
    end
    if MPI.Comm_rank(MPI.COMM_WORLD)==0 println("result = ", W) end
end

main()


# Shape inference

#### preceding or succeeding length

function prec_length(s::Dims, d::Int)
    d == 1 ? 1 :
    d == 2 ? s[1] :
    d == 3 ? s[1] * s[2] : prod(s[1:d-1])
end

function succ_length{N}(s::NTuple{N,Int}, d::Int)
    d == N ? 1 :
    d == N-1 ? s[N] : prod(s[d+1:N])
end

### Shape of mapping

# on size tuples

mapshape(s::Dims) = s

mapshape{N}(s1::NTuple{N,Int}, s2::NTuple{N,Int}) = (s1 == s2 || error("Argument dimensions are not map-compatible."); s1)

function mapshape{N1,N2}(s1::NTuple{N1,Int}, s2::NTuple{N2,Int})
    if N1 < N2
        for i = 1 : N1
            s1[i] == s2[i] || error("Argument dimensions are not map-compatible.")
        end
        for i = N1+1 : N2
            s2[i] == 1 || error("Argument dimensions are not map-compatible.")
        end
        return s2
    else
        for i = 1 : N2
            s1[i] == s2[i] || error("Argument dimensions are not map-compatible.")
        end
        for i = N2+1 : N1
            s1[i] == 1 || error("Argument dimensions are not map-compatible.")
        end
        return s1
    end
end

mapshape(s1::(), s2::()) = ()
mapshape(s1::Dims, s2::()) = s1
mapshape(s1::(), s2::Dims) = s2

mapshape(s1::(Int,Int), s2::(Int,)) = ((s1[1] == s2[1] && s1[2] == 1) || error("Argument dimensions are not map-compatible."); s1)
mapshape(s1::(Int,), s2::(Int,Int)) = ((s1[1] == s2[1] && s2[2] == 1) || error("Argument dimensions are not map-compatible."); s2)

mapshape(s1::Dims, s2::Dims, s3::Dims, rs::Dims...) = mapshape(mapshape(s1, s2), mapshape(s3, rs...))

# on arrays / numbers

mapshape(a1::Number) = ()
mapshape(a1::AbstractArray) = size(a1)

mapshape(a1::Number, a2::Number) = ()
mapshape(a1::AbstractArray, a2::Number) = size(a1)
mapshape(a1::Number, a2::AbstractArray) = size(a2)
mapshape(a1::AbstractArray, a2::AbstractArray) = mapshape(size(a1), size(a2))

mapshape(a1::Number, a2::Number, a3::Number, ra::Number...) = ()
mapshape(a1::Union(AbstractArray,Number),
         a2::Union(AbstractArray,Number),
         a3::Union(AbstractArray,Number),
         ra::Union(AbstractArray,Number)...) = mapshape(mapshape(a1, a2), mapshape(a3, ra...))

### Length of mapping

maplength(a1::Number) = 1
maplength(a1::AbstractArray) = length(a1)

maplength(a1::Number, a2::Number) = 1
maplength(a1::AbstractArray, a2::Number) = length(a1)
maplength(a1::Number, a2::AbstractArray) = length(a2)
maplength(a1::AbstractArray, a2::AbstractArray) = prod(mapshape(a1, a2))

maplength(a1::Number, a2::Number, a3::Number, ra::Number...) = 1
maplength(a1::Union(AbstractArray,Number),
          a2::Union(AbstractArray,Number),
          a3::Union(AbstractArray,Number),
          ra::Union(AbstractArray,Number)...) = prod(mapshape(a1, a2, a3, ra...))


# ---
# title: 174. Dungeon Game
# id: problem174
# author: Indigo
# date: 2021-01-20
# difficulty: Hard
# categories: Binary Search, Dynamic Programming
# link: <https://leetcode.com/problems/dungeon-game/description/>
# hidden: true
# ---
#
# The demons had captured the princess ( **P** ) and imprisoned her in the
# bottom-right corner of a dungeon. The dungeon consists of M x N rooms laid out
# in a 2D grid. Our valiant knight ( **K** ) was initially positioned in the
# top-left room and must fight his way through the dungeon to rescue the
# princess.
#
# The knight has an initial health point represented by a positive integer. If
# at any point his health point drops to 0 or below, he dies immediately.
#
# Some of the rooms are guarded by demons, so the knight loses health (
# _negative_ integers) upon entering these rooms; other rooms are either empty (
# _0 's_) or contain magic orbs that increase the knight's health ( _positive_
# integers).
#
# In order to reach the princess as quickly as possible, the knight decides to
# move only rightward or downward in each step.
#
#
#
# **Write a function to determine the knight 's minimum initial health so that
# he is able to rescue the princess.**
#
# For example, given the dungeon below, the initial health of the knight must be
# at least **7** if he follows the optimal path `RIGHT-> RIGHT -> DOWN -> DOWN`.
#
# -2 (K) | -3 | 3
# ---|---|---
# -5 | -10 | 1
# 10 | 30 | -5 (P)
#
#
#
# **Note:**
#
#   * The knight's health has no upper bound.
#   * Any room can contain threats or power-ups, even the first room the knight enters and the bottom-right room where the princess is imprisoned.
#
#
## @lc code=start
using LeetCode

function calculate_minimum_hp(dungeon::Vector{Vector{Int}})
    m, n = length(dungeon), length(dungeon[1])
    dp = fill(typemax(Int) , m + 1, n + 1)
    dp[m+1, n] = dp[m, n+1] = 1
    for j in n:-1:1
        for i in m:-1:1
            dp[i, j] = max(min(dp[i+1, j], dp[i, j+1]) - dungeon[i][j], 1)
        end
    end
    dp[1, 1]
end
## @lc code=end


using Test
using CompScienceMeshes

h = 2π / 51

for T in [Float32, Float64]
    Γ = meshcircle(T(1.0), T(h),2) ## creating the mesh first
    points=vertices(Γ)     ## get the vertices
    num_points=length(points)
    segments=cellarray(Γ)  ## get the faces
    Γ2 = barycentric_refinement(Γ) ## create the refinment of that one
    points2=vertices(Γ2)     ## get the vertices
    segments2=cellarray(Γ2)  ## get the faces

    @test points[1] == points2[1]               # test if the first point in the original and refined mesh are the same
    @test points[num_points]==points2[num_points]# test if the last point in the original has the same index in the refined one
    @test segments[1,1] ==segments2[1,1]        # test if the face are starting with the same point in test
    @test segments[1,2] ==segments2[2,2]        # test if the original second point in the first cell got separated by a single point

    # test barycentric refinement of surface mesh
    local m = meshrectangle(T(1.0), T(1.0), T(0.25), 3)
    f = barycentric_refinement(m)

    @test CompScienceMeshes.refines(f,m)
    @test numcells(f) == 6*numcells(m)

    local m1 = skeleton(m,1)
    f1 = skeleton(f,1)
    @test numcells(f1) == 2*numcells(m1) + 6*numcells(m)

    m0 = skeleton(m,0)
    @test numvertices(f) == numcells(m0) + numcells(m1) + numcells(m)

    ## test bisecting referinment of surfacic meshes
    b = bisecting_refinement(m)
    @test numcells(b) == 4*numcells(m)

    b1 = skeleton(b,1)
    @test numcells(b1) == 2*numcells(m1) + 3*numcells(m)
    @test numvertices(b) == numcells(m0) + numcells(m1)
end


#    _|_|    _|                    _|                                      _|      _|_|_|    _|_|_|    _|_|_|_|
#  _|    _|  _|_|_|      _|_|_|  _|_|_|_|  _|  _|_|    _|_|_|    _|_|_|  _|_|_|_|  _|    _|  _|    _|  _|
#  _|_|_|_|  _|    _|  _|_|        _|      _|_|      _|    _|  _|          _|      _|_|_|    _|    _|  _|_|_|
#  _|    _|  _|    _|      _|_|    _|      _|        _|    _|  _|          _|      _|        _|    _|  _|
#  _|    _|  _|_|_|    _|_|_|        _|_|  _|          _|_|_|    _|_|_|      _|_|  _|        _|_|_|    _|

abstract type AbstractPDF{N} <: AbstractFunctionWithParameters end  # N is the dimension

# calls
function (d::AbstractPDF)(x; p=freepars(d))
    normalization = normalizationintegral(d; p)
    if normalization ≈ 0.0
        println("Error: normalization ≈ 0!")
        normalization = 1.0
    end
    allp = p+fixedpars(d)
    return func(d,x; p=allp) / normalization
end

# 1 dims
function normalizationintegral(d::AbstractPDF{1}; p=freepars(d))
    allp = p+fixedpars(d)
    quadgk(x->func(d, x; p=allp), lims(d)...)[1]
end
#
# 2 dims
# function normalizationintegral(d::AbstractPDF{2}; p=freepars(d))
#     xmap = x->mapx_to_unit(x,lims(d))
#     curhe((x,f)->f[1]=func(d, xmap(x); p=p), lims(d)...)[1]
# end

"""
noparsnormf(d::AbstractPDF; p=pars(d))

Returns a single-argument lambda-function with parameters fixed to `p` and normalization computed.
"""
noparsnormf(d::AbstractPDF; p=pars(d)) = (ns=normalizationintegral(d;p=p); (x;kw...)->func(d,x;p=p)/ns)

# by default, I assume that the fields "lims" is present
lims(d::AbstractPDF) = getfield(d, :lims)
updatevalueorflag(d::AbstractPDF, s::Symbol, isfree::Bool, v=getproperty(pars(d),s)) =
    typeof(d)(updatevalueorflag(d.p, s, isfree, v), d.lims)

# other methods
function integral(d::AbstractPDF{1}, lims; p=freepars(d))
    allpars = p+fixedpars(d)
    quadgk(x->func(d,x; p=allpars), lims...)[1] / normalizationintegral(d; p=allpars)
end

####################################################################

struct Normalized{T<:AbstractFunctionWithParameters,L} <: AbstractPDF{1}
    lineshape::T
    lims::L
end
lineshape(d::Normalized) = getfield(d, :lineshape)

# two methods to be defined
import Base: getproperty
getproperty(d::Normalized, sym::Symbol) = sym==:p ? pars(lineshape(d)) : getfield(d, sym)
func(d::Normalized, x::NumberOrTuple; p=pars(d)) = func(lineshape(d), x; p)
pars(d::Normalized, isfree::Bool) = pars(lineshape(d), isfree)
updatevalueorflag(d::Normalized, s::Symbol, isfree::Bool, v=getproperty(pars(d),s)) =
    Normalized(updatevalueorflag(lineshape(d), s, isfree, v), d.lims)

# short cuts
# 1 argument
Normalized(f;p,lims) = Normalized(FunctionWithParameters(f; p), lims)

######################################################################

fixedshapepdf(f, lims) = Normalized((x;p)->f(x); lims=lims, p=∅)

######################################################################

#    _|                                                    _|      _|_|
#  _|_|_|_|  _|    _|  _|_|_|      _|_|    _|_|_|      _|_|_|    _|
#    _|      _|    _|  _|    _|  _|_|_|_|  _|    _|  _|    _|  _|_|_|_|
#    _|      _|    _|  _|    _|  _|        _|    _|  _|    _|    _|
#      _|_|    _|_|_|  _|_|_|      _|_|_|  _|_|_|      _|_|_|    _|
#                  _|  _|                  _|
#              _|_|    _|                  _|

"""
    @makepdftype MyPDF(x;p) = unnormdensity(x, p.a, p.b)

    Expected form of the expression is `f(x;p)` on the left
"""
macro makepdftype(ex)
    #
    fpx = ex.args[1].args
    name = fpx[1]
    p = fpx[2].args[1]
    (p != :p) && error("expected format f(x;p) = ... " )
    x = fpx[3]
    body = ex.args[2]
    #
    quote
        struct $name{T,N} <: AbstractPDF{1}
            p::T
            lims::N
        end
        $(esc(name))(;p,lims) = $(esc(name))(p, lims)

        import AlgebraPDF: func, pars, updatevalueorflag
        #
        $(esc(:func))(d::$(esc(name)), $(esc(x))::NumberOrTuple; p=$(esc(:pars))(d)) = $(esc(body))
    end
end


#=

@author : Spencer Lyon <spencer.lyon@nyu.edu>

=#
using QuantEcon
using Plots
pyplot()

# model parameters
a_0 = 100
a_1 = 0.5
ρ = 0.9
σ_d = 0.05
β = 0.95
c = 2
γ = 50.0
θ = 0.002
ac = (a_0 - c) / 2.0

# Define LQ matrices
R = [ 0   ac    0;
     ac -a_1  0.5;
     0.  0.5    0]
R = -R  # For minimization
Q = Matrix([γ / 2.0]')
A = [1. 0. 0.;
     0. 1. 0.;
     0. 0.  ρ]
B = [0. 1. 0.]'
C = [0. 0. σ_d]'

## Functions

function evaluate_policy(θ::AbstractFloat, F::AbstractArray)
    rlq = RBLQ(Q, R, A, B, C, β, θ)
    K_F, P_F, d_F, O_F, o_F = evaluate_F(rlq, F)
    x0 = [1.0 0.0 0.0]'
    value = - x0' * P_F * x0 - d_F
    entropy = x0' * O_F * x0 + o_F
    return value[1], entropy[1]  # return scalars
end

function value_and_entropy{TF<:AbstractFloat}(emax::AbstractFloat,
                                              F::AbstractArray{TF},
                                              bw::String,
                                              grid_size::Integer=1000)
    if lowercase(bw) == "worst"
        θs = 1 ./ linspace(1e-8, 1000, grid_size)
    else
        θs = -1 ./ linspace(1e-8, 1000, grid_size)
    end

    data = Array{TF}(grid_size, 2)

    for (i, θ) in enumerate(θs)
        data[i, :] = collect(evaluate_policy(θ, F))
        if data[i, 2] >= emax  # stop at this entropy level
            data = data[1:i, :]
            break
        end
    end
    return data
end

## Main

# compute optimal rule
optimal_lq = LQ(Q, R, A, B, C, zero(B'A), bet=β)
Po, Fo, Do = stationary_values(optimal_lq)

# compute robust rule for our θ
baseline_robust = RBLQ(Q, R, A, B, C, β, θ)
Fb, Kb, Pb = robust_rule(baseline_robust)

# Check the positive definiteness of worst-case covariance matrix to
# ensure that θ exceeds the breakdown point
test_matrix = eye(size(Pb, 1)) - (C' * Pb * C ./ θ)[1]
eigenvals, eigenvecs = eig(test_matrix)
@assert all(eigenvals .>= 0)

emax = 1.6e6

# compute values and entropies
optimal_best_case = value_and_entropy(emax, Fo, "best")
robust_best_case = value_and_entropy(emax, Fb, "best")
optimal_worst_case = value_and_entropy(emax, Fo, "worst")
robust_worst_case = value_and_entropy(emax, Fb, "worst")

# we reverse order of "worst_case"s so values are ascending
data_pairs = ((optimal_best_case, optimal_worst_case),
              (robust_best_case, robust_worst_case))

egrid = linspace(0, emax, 100)
egrid_data = Array{Float64}[]
for data_pair in data_pairs
    for data in data_pair
        x, y = data[:, 2], data[:, 1]
        curve = LinInterp(x, y)
        push!(egrid_data, curve.(egrid))
    end
end
plot(egrid, egrid_data, color=[:red :red :blue :blue])
plot!(egrid, egrid_data[1], fillrange=egrid_data[2],
      fillcolor=:red, fillalpha=0.1, color=:red, legend=:none)
plot!(egrid, egrid_data[3], fillrange=egrid_data[4],
      fillcolor=:blue, fillalpha=0.1, color=:blue, legend=:none)
plot!(xlabel="Entropy", ylabel="Value")


"""
    arrhenius(A,Eₐ,Tₖ,Tᵣₖ,R = Constants().R)

The Arrhenius function for dependence of the rate constant of a chemical reaction.

# Arguments

- `A`: pre-exponential factor, a constant for each chemical reaction
- `Eₐ`: activation energy for the reaction (``J\\ mol^{-1}``)
- `Tₖ`: temperature (Kelvin)
- `Tᵣₖ`: reference temperature (Kelvin) at which A was measured
- `R`: universal gas constant (``J\\ mol^{-1}\\ K^{-1}``)

# Examples

```julia
# Importing physical constants
constants = Constants()
# Using default values for the model:
A = Fvcb()

# Computing Jmax:
arrhenius(A.JMaxRef,A.Eₐⱼ,28.0-constants.K₀,A.Tᵣ-constants.K₀,constants.R)
# ! Warning: temperatures must be given in Kelvin

# Computing Vcmax:
arrhenius(A.VcMaxRef,A.Eₐᵥ,28.0-constants.K₀,A.Tᵣ-constants.K₀,constants.R)
```
"""
function arrhenius(A,Eₐ,Tₖ,Tᵣₖ,R = Constants().R)
    A * exp(Eₐ * (Tₖ - Tᵣₖ) / (R * Tₖ * Tᵣₖ))
end

"""
    arrhenius(A,Eₐ,Tₖ,Tᵣₖ,Hd,Δₛ,R = Constants().R)

The Arrhenius function for dependence of the rate constant of a chemical reaction,
modified following equation (17) from Medlyn et al. (2002) to consider the negative effect of
very high temperatures.

# Arguments

- `A`: the pre-exponential factor, a constant for each chemical reaction
- `Eₐ`: activation energy (``J\\ mol^{-1}``), or the exponential rate of rise
 of the function (Ha in the equation of Medlyn et al. (2002))
- `Tₖ`: current temperature (Kelvin)
- `Tᵣₖ`: reference temperature (Kelvin) at which A was measured
- `Hd`: rate of decrease of the function above the optimum (called EDVJ in
[MAESPA](http://maespa.github.io/) and [plantecophys](https://remkoduursma.github.io/plantecophys/))
- `Δₛ`: entropy factor
- `R`: is the universal gas constant (``J\\ mol^{-1}\\ K^{-1}``)

References

Medlyn, B. E., E. Dreyer, D. Ellsworth, M. Forstreuter, P. C. Harley, M. U. F. Kirschbaum,
X. Le Roux, et al. 2002. « Temperature response of parameters of a biochemically based model
of photosynthesis. II. A review of experimental data ». Plant, Cell & Environment 25 (9): 1167‑79.
https://doi.org/10.1046/j.1365-3040.2002.00891.x.

# Examples

```julia
# Importing physical constants
constants = Constants()
# Using default values for the model:
A = Fvcb()

# Computing Jmax:
PlantBiophysics.arrhenius(A.JMaxRef,A.Eₐⱼ,28.0-constants.K₀,A.Tᵣ-constants.K₀,A.Hdⱼ,A.Δₛⱼ)
# ! Warning: temperatures must be given in Kelvin

# Computing Vcmax:
PlantBiophysics.arrhenius(A.VcMaxRef,A.Eₐᵥ,28.0-constants.K₀,A.Tᵣ-constants.K₀,A.Hdᵥ,A.Δₛᵥ)

```
"""
function arrhenius(A,Eₐ,Tₖ,Tᵣₖ,Hd,Δₛ,R = Constants().R)
    # Equation split in 3 parts for readability:
    ftk1 = arrhenius(A,Eₐ,Tₖ,Tᵣₖ,R)
    ftk2 = (1.0 + exp((Tᵣₖ * Δₛ - Hd) / (Tᵣₖ * R)))
    ftk3 = (1.0 + exp((Tₖ * Δₛ - Hd) / (Tₖ * R)))

    ftk = ftk1 * ftk2 / ftk3

    return ftk
end

"""
    Γ_star(Tₖ,Tᵣₖ,R = Constants().R)

CO₂ compensation point ``Γ^⋆`` (``μ mol\\ mol^{-1}``) according to equation (12)
from Medlyn et al. (2002).

> ``Γ^⋆`` is the [CO₂] at which oxygenation proceeds at
twice the rate of carboxylation causing photosynthetic uptake of CO2 to be exactly
compensated by photorespiratory CO₂ release (Sharkey et al., 2007).

# Notes

Could be replaced by equation (38) from Farquhar et al. (1980), but Medlyn et al. (2002)
states that ``Γ^⋆`` as a relatively low effect on the model outputs.

# Arguments

- `Tₖ` (Kelvin): current temperature
- `Tᵣₖ` (Kelvin): reference temperature at which A was measured
- `R` (``J\\ mol^{-1}\\ K^{-1}``): is the universal gas constant

# Examples

```julia
# Importing the physical constants:
constants = Constants()
# computing the temperature dependence of γˢ:
Γ_star(28-constants.K₀,25-constants.K₀,constants.R)
```

# References

Farquhar, G. D., S. von von Caemmerer, et J. A. Berry. 1980. « A biochemical model of
photosynthetic CO2 assimilation in leaves of C3 species ». Planta 149 (1): 78‑90.

Medlyn, B. E., E. Dreyer, D. Ellsworth, M. Forstreuter, P. C. Harley, M. U. F. Kirschbaum,
X. Le Roux, et al. 2002. « Temperature response of parameters of a biochemically based model
of photosynthesis. II. A review of experimental data ». Plant, Cell & Environment 25 (9): 1167‑79.
https://doi.org/10.1046/j.1365-3040.2002.00891.x.

Sharkey, Thomas D., Carl J. Bernacchi, Graham D. Farquhar, et Eric L. Singsaas. 2007.
« Fitting Photosynthetic Carbon Dioxide Response Curves for C3 Leaves ». Plant, Cell &
Environment 30 (9): 1035‑40. https://doi.org/10.1111/j.1365-3040.2007.01710.x.

"""
function Γ_star(Tₖ,Tᵣₖ,R = Constants().R)
    arrhenius(oftype(Tₖ,42.75),oftype(Tₖ,37830.0),Tₖ,Tᵣₖ,R)
end

"""
Compute the effective Michaelis–Menten coefficient for CO2 ``Km`` (``μ mol\\ mol^{-1}``) according to
Medlyn et al. (2002), equations (5) and (6).

# References

Medlyn, B. E., E. Dreyer, D. Ellsworth, M. Forstreuter, P. C. Harley, M. U. F. Kirschbaum,
X. Le Roux, et al. 2002. « Temperature response of parameters of a biochemically based model
of photosynthesis. II. A review of experimental data ». Plant, Cell & Environment 25 (9): 1167‑79.
https://doi.org/10.1046/j.1365-3040.2002.00891.x.

# Examples

```julia
# computing the temperature dependence of γˢ:
get_km(28,25,210.0)
```
"""
function get_km(Tₖ,Tᵣₖ,O₂,R = Constants().R)
    KC = arrhenius(oftype(Tₖ,404.9),oftype(Tₖ,79430.0),Tₖ,Tᵣₖ,R)
    KO = arrhenius(oftype(Tₖ,278.4),oftype(Tₖ,36380.0),Tₖ,Tᵣₖ,R)
    return KC * (1.0 + O₂/KO)
end


using Distributions

# Set Confidence Intervals
η_v = .05 # 1 - Confidence for voltage limit cc
η_g = .05 # 1 - Confidence for generation limit cc
ξ = .005# 1 - Confidence for variance estimator

# SND Quantiles for global use
snd = Normal(0,1)
z_v = quantile(snd, 1-η_v)
z_g = quantile(snd, 1-η_g)
z_g2 = quantile(snd, 1-(η_g/2))

# Set True Load Standard Deviation
ld = 0.2

# Save Historic Data Globally
global ERROR_HIST = 0

# Calculate variance estimation confidence
function chisq_interval(s, n, alpha)
    # println("s = $s, n = $n, alpha = $alpha")
    # s... sample variance
    # n... degrees of freedom
    # alpha... confidence
    # Returns lower and upper bound of the 1-alpha-confidence interval
    # of the estmated variance
    ch = Chisq(n)
    lower = (n * s)/quantile(ch, 1- alpha)
    upper = (n * s)/quantile(ch, alpha)
    return [lower, upper]
end

function create_data(var_true, N)
    obs = zeros(length(var_true), N)
    for (i,v) in enumerate(var_true)
        if v > 0
            dist = Normal(0,sqrt(v))
            obs[i,:] = rand(dist, N)
        end
    end
    return obs
end

function init_stochasticity(;create_new_data=true, sample_N=100, variance_opt="implicit")

# Set the ture error distribution
# Same Variance for each bus
n_buses = length(BUSES)
if variance_opt == "implicit"
# Create individual node variance from load
    loads = [BUSES[b].d_P for b in 1:n_buses]
    var_true_vector = (loads.*ld).^2
    # Variance = (ld% of load)^2
elseif variance_opt == "explicit"
# Same variance for all nodes
    if DATA_DIR == "basecase"
        f = 0.01
    elseif DATA_DIR == "simplecase"
        f = 0.1
    elseif DATA_DIR == "33buscase_pu"
        f = 4
    else
        f = 1
        warn("Data Dir. $(DATA_DIR) unkown for uncertainty setting")
    end
    var_true_vector = ones(n_buses)*f
elseif variance_opt == "manual"
# Manually create an array of variances
    var_true_vector = []
    if length(var_true_vector) < n_buses
        warn("Variance Vector to short")
        var_true_vector = ones(n_buses).*1e-8
    end
else
    warn("Unknown Variance Setting")
    var_true_vector = ones(n_buses).*1e-8
end

# Create or load data for samples
if create_new_data || ERROR_HIST == 0
    info("Creating new Historical Observations for DR Implementation")
    N = sample_N
    error_hist = create_data(var_true_vector, N)
    global ERROR_HIST = error_hist
else
    N = length(ERROR_HIST[1,:])
    error_hist = ERROR_HIST
end

# Work with the data
# Calculate sample variances
var_sample_vector =[]
for i in 1:n_buses
    sn = 1/N * sum(error_hist[i,:].^2)
    push!(var_sample_vector, sn)
end

# var_sample_covmat = 1/N * error_hist * error_hist'
# var_sample_vector =  diag(var_sample_covmat) # Drop covariance entries

intervals = [chisq_interval(s_n, N, ξ) for s_n in var_sample_vector]
var_sample_lower = [intervals[i][1] for i in 1:n_buses]
var_sample_upper = [intervals[i][2] for i in 1:n_buses]

# return var_true_vector, error_dist_true, var_sample_upper, var_sample_vector
return var_true_vector, var_sample_upper, var_sample_vector
end


pos = [6, 2]  # input position here
rolloutcomes = [[3, 1], [4, 3], [5, 6], [6, 7], [7, 6], [8, 3], [9, 1]]

function roll1000()
    detdie(idx) = 3 * (3 * mod1(idx, 100) - 1)
    scorespart1 = [0, 0]
    positionspart1 = deepcopy(pos)
    won = [false, false]
    round = 0
    while !won[1] && !won[2]
        for (i, s) in enumerate(scorespart1)
            round += 1
            score = detdie(round)
            positionspart1[i] = mod1(positionspart1[i] + score, 10)
            scorespart1[i] += positionspart1[i]
            if scorespart1[i] >= 1000
                won[i] = true
                break
            end
        end
    end
    return 3 * round * minimum(scorespart1)
end

println("Part 1: ", roll1000())

function winners(player1, r1, player2, r2)
    if r2 <= 0
        return (0, 1)
    end
    wins1, wins2 = 0, 0
    for (roll, result) in rolloutcomes
        unis2, unis1 = winners(
            player2, r2, (player1 + roll) % 10, r1 - 1 - (player1 + roll) % 10
        )
        wins1 += result * unis1
        wins2 += result * unis2
    end
    return wins1, wins2
end

wins1, wins2 = winners(pos[1] - 1, 21, pos[2] - 1, 21)
println("Part 2: ", max(wins1, wins2))


# Test that the volumes of the supercells are correct compared to the lattice.
# Mostly a sanity check to make sure the lattices are correc.

@testset "Supercell volumes" begin
@testset "fcc" begin
    Vfcc = abs(det(cellfcc))
    @test abs(det(fcc1*cellfcc))/Vfcc   ≈ 1.0
    @test abs(det(fcc4*cellfcc))/Vfcc   ≈ 4.0
    @test abs(det(fcc7*cellfcc))/Vfcc   ≈ 7.0
    @test abs(det(fcc13*cellfcc))/Vfcc  ≈ 13.0
    @test abs(det(fcc23*cellfcc))/Vfcc  ≈ 23.0
    @test abs(det(fcc26*cellfcc))/Vfcc  ≈ 26.0
    @test abs(det(fcc32*cellfcc))/Vfcc  ≈ 32.0
    @test abs(det(fcc38*cellfcc))/Vfcc  ≈ 38.0
    @test abs(det(fcc50*cellfcc))/Vfcc  ≈ 50.0
    @test abs(det(fcc70*cellfcc))/Vfcc  ≈ 70.0
    @test abs(det(fcc92*cellfcc))/Vfcc  ≈ 92.0
    @test abs(det(fcc96*cellfcc))/Vfcc  ≈ 96.0
    @test abs(det(fcc100*cellfcc))/Vfcc ≈ 100.0
    @test abs(det(fcc108*cellfcc))/Vfcc ≈ 108.0
end;
@testset "cubic" begin
    Vcubic = abs(det(cellcubic))
    @test abs(det(cubic1*cellcubic))/Vcubic   ≈ 1.0
    @test abs(det(cubic3*cellcubic))/Vcubic   ≈ 3.0
    @test abs(det(cubic7*cellcubic))/Vcubic   ≈ 7.0
    @test abs(det(cubic8*cellcubic))/Vcubic   ≈ 8.0
    @test abs(det(cubic13*cellcubic))/Vcubic  ≈ 13.0
    @test abs(det(cubic14*cellcubic))/Vcubic  ≈ 14.0
    @test abs(det(cubic24*cellcubic))/Vcubic  ≈ 24.0
    @test abs(det(cubic27*cellcubic))/Vcubic  ≈ 27.0
    @test abs(det(cubic33*cellcubic))/Vcubic  ≈ 33.0
    @test abs(det(cubic38*cellcubic))/Vcubic  ≈ 38.0
    @test abs(det(cubic52*cellcubic))/Vcubic  ≈ 52.0
    @test abs(det(cubic56*cellcubic))/Vcubic  ≈ 56.0
    @test abs(det(cubic64*cellcubic))/Vcubic  ≈ 64.0
    @test abs(det(cubic73*cellcubic))/Vcubic  ≈ 73.0
    @test abs(det(cubic76*cellcubic))/Vcubic  ≈ 76.0
    @test abs(det(cubic93*cellcubic))/Vcubic  ≈ 93.0
    @test abs(det(cubic100*cellcubic))/Vcubic ≈ 100.0

end;
@testset "bcc" begin
    Vbcc = abs(det(cellbcc))
    @test abs(det(bcc1*cellbcc))/Vbcc   ≈ 1.0
    @test abs(det(bcc2*cellbcc))/Vbcc   ≈ 2.0
    @test abs(det(bcc6*cellbcc))/Vbcc   ≈ 6.0
    @test abs(det(bcc9*cellbcc))/Vbcc   ≈ 9.0
    @test abs(det(bcc16*cellbcc))/Vbcc  ≈ 16.0
    @test abs(det(bcc21*cellbcc))/Vbcc  ≈ 21.0
    @test abs(det(bcc25*cellbcc))/Vbcc  ≈ 25.0
    @test abs(det(bcc28*cellbcc))/Vbcc  ≈ 28.0
    @test abs(det(bcc35*cellbcc))/Vbcc  ≈ 35.0
    @test abs(det(bcc48*cellbcc))/Vbcc  ≈ 48.0
    @test abs(det(bcc53*cellbcc))/Vbcc  ≈ 53.0
    @test abs(det(bcc54*cellbcc))/Vbcc  ≈ 54.0
    @test abs(det(bcc65*cellbcc))/Vbcc  ≈ 65.0
    @test abs(det(bcc84*cellbcc))/Vbcc  ≈ 84.0
    @test abs(det(bcc91*cellbcc))/Vbcc  ≈ 91.0
    @test abs(det(bcc103*cellbcc))/Vbcc ≈ 103.0
    @test abs(det(bcc107*cellbcc))/Vbcc ≈ 107.0
    @test abs(det(bcc112*cellbcc))/Vbcc ≈ 112.0
    @test abs(det(bcc128*cellbcc))/Vbcc ≈ 128.0
end;
@testset "hcp" begin
    Vhcp = abs(det(cellhcp))
    @test abs(det(hcp1*cellhcp))/Vhcp   ≈ 1.0
    @test abs(det(hcp3*cellhcp))/Vhcp   ≈ 3.0
    @test abs(det(hcp6*cellhcp))/Vhcp   ≈ 6.0
    @test abs(det(hcp8*cellhcp))/Vhcp   ≈ 8.0
    @test abs(det(hcp11*cellhcp))/Vhcp  ≈ 11.0
    @test abs(det(hcp14*cellhcp))/Vhcp  ≈ 14.0
    @test abs(det(hcp16*cellhcp))/Vhcp  ≈ 16.0
    @test abs(det(hcp21*cellhcp))/Vhcp  ≈ 21.0
    @test abs(det(hcp27*cellhcp))/Vhcp  ≈ 27.0
    @test abs(det(hcp37*cellhcp))/Vhcp  ≈ 37.0
    @test abs(det(hcp39*cellhcp))/Vhcp  ≈ 39.0
    @test abs(det(hcp42*cellhcp))/Vhcp  ≈ 42.0
    @test abs(det(hcp53*cellhcp))/Vhcp  ≈ 53.0
    @test abs(det(hcp62*cellhcp))/Vhcp  ≈ 62.0
    @test abs(det(hcp65*cellhcp))/Vhcp  ≈ 65.0
    @test abs(det(hcp69*cellhcp))/Vhcp  ≈ 69.0
    @test abs(det(hcp76*cellhcp))/Vhcp  ≈ 76.0
    @test abs(det(hcp80*cellhcp))/Vhcp  ≈ 80.0
    @test abs(det(hcp88*cellhcp))/Vhcp  ≈ 88.0
    @test abs(det(hcp90*cellhcp))/Vhcp  ≈ 90.0
    @test abs(det(hcp100*cellhcp))/Vhcp ≈ 100.0
end;
end;


#=
Cohen:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-10
=#

using SuiteSparseGraphBLAS
using SparseArrays
using LinearAlgebra

##A is the input matrix, Output: Triangle Count
function cohen(A)

    #extracting lower matrix from input matrix
    L = tril(A)

    #extracting upper matrix from input matrix
    U = triu(A')

    B = mul(L, U, Semirings.PLUS_TIMES)
    C = emul(B, A,  BinaryOps.TIMES)
    result = reduce(Monoids.PLUS_MONOID[Float64], C, dims=:)
    result = result /2
    print("RESULT ")
    print(result)

    return result

end


@doc doc"""
    Stiefel{n,k,T} <: Manifold

The Stiefel manifold consists of all $n\times k$, $n\geq k$ orthonormal matrices, i.e.

````math
\mathcal M = \{ x \in \mathbb F^{n\times k} : x^{\mathrm{H}}x = I_k \},
````

where $\mathbb F \in \{\mathbb R, \mathbb C\}$,
$\cdot^{\mathrm{H}}$ denotes the complex conjugate transpose or Hermitian, and
$I_n \in \mathbb R^{n\times n}$ denotes the $k \times k$ identity matrix.

The tangent space at a point $x\in\mathcal M$ is given by

````math
T_x\mathcal M = \{ v \in \mathbb{F}^{n\times k} : x^{\mathrm{H}}v + v^{\mathrm{H}}x=0_n\},
````

where $0_n$ is the $k\times k$ zero matrix.

The metric is either inherited from $\mathbb R^{n,k}$ for the real-valued case
or the one inherited from interpreting the complex valued entries in the Gaussian
plane $\mathbb R^2$ and then over all entries as before, i.e. the latter
may be called an Hermitian metric in the complex-valued matrices.

The manifold is named after
[Eduard L. Stiefel](https://en.wikipedia.org/wiki/Eduard_Stiefel) (1909–1978).

# Constructor
    Stiefel(n,k,F=ℝ)

Generate the (real-valued) Stiefel manifold of $n\times k$ dimensional orthonormal matrices.
"""
struct Stiefel{n,k,F} <: Manifold end

Stiefel(n::Int, k::Int, F::AbstractNumbers = ℝ) = Stiefel{n,k,F}()

@doc doc"""
    check_manifold_point(M::Stiefel, x; kwargs...)

Check whether `x` is a valid point on the [`Stiefel`](@ref) `M`=$\operatorname{St}(n,k)$,
i.e. that it has the right [`AbstractNumbers`](@ref) type and $x^{\mathrm{H}}x$
is (approximatly) the identity, where $\cdot^{\mathrm{H}}$ is the complex conjugate
transpose. The settings for approximately can be set with `kwargs...`.
"""
function check_manifold_point(M::Stiefel{n,k,T}, x; kwargs...) where {n,k,T}
    if (T === ℝ) && !(eltype(x) <: Real)
        return DomainError(
            eltype(x),
            "The matrix $(x) is not a real-valued matrix, so it does noe lie on the Stiefel manifold of dimension ($(n),$(k)).",
        )
    end
    if (T === ℂ) && !(eltype(x) <: Real) && !(eltype(x) <: Complex)
        return DomainError(
            eltype(x),
            "The matrix $(x) is neiter real- nor complex-valued matrix, so it does noe lie on the complex Stiefel manifold of dimension ($(n),$(k)).",
        )
    end
    if any(size(x) != representation_size(M))
        return DomainError(
            size(x),
            "The matrix $(x) is does not lie on the Stiefel manifold of dimension ($(n),$(k)), since its dimensions are wrong.",
        )
    end
    c = x' * x
    if !isapprox(c, one(c); kwargs...)
        return DomainError(
            norm(c - one(c)),
            "The point $(x) does not lie on the Stiefel manifold of dimension ($(n),$(k)), because x'x is not the unit matrix.",
        )
    end
end

@doc doc"""
    check_tangent_vector(M::Stiefel, x, v; kwargs...)

Check whether `v` is a valid tangent vector at `x` on the [`Stiefel`](@ref)
`M`=$\operatorname{St}(n,k)$, i.e. the [`AbstractNumbers`](@ref) fits and
it (approximtly) holds that $x^{\mathrm{H}}v + v^{\mathrm{H}}x = 0$, where
`kwargs...` is passed to the `isapprox`.
"""
function check_tangent_vector(M::Stiefel{n,k,T}, x, v; kwargs...) where {n,k,T}
    mpe = check_manifold_point(M, x)
    mpe === nothing || return mpe
    if (T === ℝ) && !(eltype(v) <: Real)
        return DomainError(
            eltype(v),
            "The matrix $(v) is not a real-valued matrix, so it can not be a tangent vector to the Stiefel manifold of dimension ($(n),$(k)).",
        )
    end
    if (T === ℂ) && !(eltype(v) <: Real) && !(eltype(v) <: Complex)
        return DomainError(
            eltype(v),
            "The matrix $(v) is neiter real- nor complex-valued matrix, so it can not bea tangent vectorto the complex Stiefel manifold of dimension ($(n),$(k)).",
        )
    end
    if any(size(v) != representation_size(M))
        return DomainError(
            size(v),
            "The matrix $(v) is does not lie in the tangent space of $(x) on the Stiefel manifold of dimension ($(n),$(k)), since its dimensions are wrong.",
        )
    end
    if !isapprox(x' * v + v' * x, zeros(k, k); kwargs...)
        return DomainError(
            norm(x' * v + v' * x),
            "The matrix $(v) is does not lie in the tangent space of $(x) on the Stiefel manifold of dimension ($(n),$(k)), since x'v + v'x is not the zero matrix.",
        )
    end
end

@doc doc"""
    exp(M, x, v)

Compute the exponential map on the [`Stiefel`](@ref)`{n,k,T}`() manifold `M`
eminating from `x` into tangent direction `v`.

$\operatorname{exp}_{x} v = \begin{pmatrix}
   x\\v
 \end{pmatrix}
 \exp
 \left(
 \begin{pmatrix} x^{\mathrm{H}}v & - v^{\mathrm{H}}v\\
 I_n & x^{\mathrm{H}}v\end{pmatrix}
 \right)
\begin{pmatrix}  \exp( -x^{\mathrm{H}}v) \\ 0_n\end{pmatrix}$

where $\exp$ denotes matrix exponential,
$\cdot^{\mathrm{H}}$ denotes the complex conjugate transpose or Hermitian, and $I_k$ and
$0_k$ are the identity matrix and the zero matrix of dimension $k \times k$, respectively.
"""
exp(::Stiefel, ::Any...)

function exp!(M::Stiefel{n,k}, y, x, v) where {n,k}
    return copyto!(
        y,
        [x v] *
        exp([x'v -v' * v; one(zeros(eltype(x), k, k)) x' * v]) *
        [exp(-x'v); zeros(eltype(x), k, k)],
    )
end

@doc doc"""
    inner(M::Stiefel, x, v, w)

Compute the inner product for two tangent vectors `v`, `w` from the
tangent space of `x` on the [`Stiefel`](@ref) manifold `M`. The formula reads

````math
(v,w)_x = \operatorname{trace}(v^{\mathrm{H}}w),
````
i.e. the [`EuclideanMetric`](@ref) from the embedding restricted to the tangent
space. For the complex-valued case this is the Hermitian metric, to be precise.
"""
inner(::Stiefel, x, v, w) = dot(v, w)

@doc doc"""
    inverse_retract(M::Stiefel, x, y, ::PolarInverseRetraction)

Compute the inverse retraction based on a singular value decomposition
for two points `x`, `y` on the [`Stiefel`](@ref) manifold `M` and return
the resulting tangent vector in `v`. This follows the folloing approach:
From the Polar retraction we know that

````math
\operatorname{retr}_x^{-1}(y) = ys - q
````

if such a symmetric positive definite $k\times k$ matrix exists. Since $ys-q$ is
also a tangent vector at $x$ we obtain

````math
x^{\mathrm{H}}ys + s(x^{\mathrm{H}}y)^{\mathrm{H}} + 2I_k = 0,
````
which can either be solved by a Lyapunov approach or a continuous-time
algebraic Riccati equation as described in [^KanekoFioriTanaka2013]

This implementation follows the Lyapunov approach.

    inverse_retract(M, x, y, ::QRInverseRetraction)

Compute the inverse retraction based on a qr decomposition
for two points `x`, `y` on the [`Stiefel`](@ref) manifold `M` and return
the resulting tangent vector in `v`. The computation follows Algorithm 1
in [^KanekoFioriTanaka2013].

[^KanekoFioriTanaka2013]:
    > T. Kaneko, S. Fiori, T. Tanaka: "Empirical Arithmetic Averaging over the
    > Compact Stiefel Manifold", IEEE Transactions on Signal Processing, 2013,
    > doi: [10.1109/TSP.2012.2226167](https://doi.org/10.1109/TSP.2012.2226167).
"""
inverse_retract(::Stiefel, ::Any...)

function inverse_retract!(::Stiefel, v, x, y, ::PolarInverseRetraction)
    A = x' * y
    H = -2 * one(x' * x)
    B = lyap(A, H)
    return copyto!(v, y * B - x)
end
function inverse_retract!(::Stiefel{n,k}, v, x, y, ::QRInverseRetraction) where {n,k}
    A = x' * y
    R = zeros(typeof(one(eltype(x)) * one(eltype(y))), k, k)
    for i = 1:k
        b = zeros(i)
        b[i] = 1
        b[1:(end-1)] = -transpose(R[1:(i-1), 1:(i-1)]) * A[i, 1:(i-1)]
        R[1:i, i] = A[1:i, 1:i] \ b
    end
    return copyto!(v, y * R - x)
end

function isapprox(M::Stiefel, x, v, w; kwargs...)
    return isapprox(sqrt(inner(M, x, zero_tangent_vector(M, x), v - w)), 0; kwargs...)
end
isapprox(M::Stiefel, x, y; kwargs...) = isapprox(norm(x - y), 0; kwargs...)

@doc doc"""
    manifold_dimension(M::Stiefel)

Return the dimension of the [`Stiefel`](@ref) manifold `M`=$\operatorname{St}(n,k,𝔽)$.
The dimension is given by

````math
\dim \mathrm{Stiefel}(n, k, ℝ) &= nk - \frac{1}{2}k(k+1)\\
\dim \mathrm{Stiefel}(n, k, ℂ) &= 2nk - k^2\\
\dim \mathrm{Stiefel}(n, k, ℍ) &= 4nk - k(2k-1)
````
"""
manifold_dimension(::Stiefel{n,k,ℝ}) where {n,k} = n * k - div(k * (k + 1), 2)
manifold_dimension(::Stiefel{n,k,ℂ}) where {n,k} = 2 * n * k - k * k
manifold_dimension(::Stiefel{n,k,ℍ}) where {n,k} = 4 * n * k - k * (2k - 1)

@doc doc"""
    project_tangent(M, x, v)

Project `v` onto the tangent space of `x` to the [`Stiefel`](@ref) manifold `M`.
The formula reads

````math
\operatorname{proj}_{\mathcal M}(x,v) = v - x \operatorname{Sym}(x^{\mathrm{H}}v),
````

where $\operatorname{Sym}(y)$ is the symmetrization of $y$, e.g. by
$\operatorname{Sym}(y) = \frac{y^{\mathrm{H}}+y}{2}$.
"""
project_tangent(::Stiefel, ::Any...)

project_tangent!(::Stiefel, w, x, v) = copyto!(w, v - x * Symmetric(x' * v))

@doc doc"""
    retract(M, x, v, ::PolarRetraction)

Compute the SVD-based retraction [`PolarRetraction`](@ref) on the
[`Stiefel`](@ref) manifold `M`. With $USV = x + v$ the retraction reads
````math
\operatorname{retr}_x(v) = U\bar{V}^\mathrm{H}.
````

    retract(M, x, v, ::QRRetraction )

Compute the QR-based retraction [`QRRetraction`](@ref) on the
[`Stiefel`](@ref) manifold `M`. With $QR = x + v$ the retraction reads
````math
\operatorname{retr}_x(v) = QD,
````
where D is a $n\times k$ matrix with
````math
D = \operatorname{diag}\bigl(\operatorname{sgn}(R_{ii}+0,5)_{i=1}^k \bigr),
````
where $\operatorname{sgn}(x) = \begin{cases}
1 & \text{ for } x > 0,\\
0 & \text{ for } x = 0,\\
-1& \text{ for } x < 0.
\end{cases}$
"""
retract(::Stiefel, ::Any...)

function retract!(::Stiefel, y, x, v, ::PolarRetraction)
    s = svd(x + v)
    return mul!(y, s.U, s.Vt)
end
function retract!(::Stiefel, y, x, v, ::QRRetraction)
    qrfac = qr(x + v)
    d = diag(qrfac.R)
    D = Diagonal(sign.(sign.(d .+ 0.5)))
    return copyto!(y, Matrix(qrfac.Q) * D)
end

@doc doc"""
    representation_size(M::Stiefel)

Returns the representation size of the [`Stiefel`](@ref) `M`=$\operatorname{St}(n,k)$,
i.e. `(n,k)`, which is the matrix dimensions.
"""
@generated representation_size(::Stiefel{n,k}) where {n,k} = (n, k)

@doc doc"""
    zero_tangent_vector(M::Stiefel, x)

Returns the zero tangent vector from the tangent space at `x` on the [`Stiefel`](@ref)
`M`=$\operatorname{St}(n,k)$, i.e. an `(n,k)` zero matrix.
"""
zero_tangent_vector(::Stiefel, ::Any...)

zero_tangent_vector!(::Stiefel, v, x) = fill!(v, 0)


using TrussMorph: SectionProperties, MaterialProperties
using LinearAlgebra: norm, cholesky, I
using SparseArrays

function local_coord_rotation(u::Vector{Float64}, v::Vector{Float64})
    local L = norm(u - v)
    c = (v - u) ./ L
    R = zeros(Float64, 3, 3)
    R[1,1] = c[1]
    R[1,2] = c[2]
    R[2,1] = -c[2]
    R[2,2] = c[1]
    R[3,3] = 1
    return R
end

"""
local_stiffness_matrix: compute elemental FRAME stiffness matrix
in the local coordinates.

Return: 6 x 6 matrix K_le (local elemental)

Note: only support 2d right now.
"""
function local_stiffness_matrix(L::Float64, mp::MaterialProperties, sp::SectionProperties)
    E = mp.E
    G = mp.G
    A = sp.A
    I = sp.Iz
    fs = 1.0 # shear deformation const
    beta = (12*E*I*fs)/(G*A*L^2)

    K_beam = zeros(Float64, 4, 4)
    K_beam[1,2] = 6*L;
    K_beam[1,3] = -12;
    K_beam[1,4] = 6*L;
    K_beam[2,3] = -6*L;
    K_beam[2,4] = L^2*(2-beta);
    K_beam[3,4] = -6*L;
    K_beam = K_beam' + K_beam;

    K_beam += Diagonal([12, L^2*(4+beta), 12, L^2*(4+beta)])
    K_beam *= (E*I)/(L^3*(1+beta))

    K_le = zeros(Float64, 6, 6)
    K_le[1,1] = E * A / L
    K_le[1,4] = -E * A / L
    K_le[4,1] = -E * A / L
    K_le[4,4] = E * A / L
    K_le[2:3,2:3] = K_beam[1:2,1:2];
    K_le[2:3,5:6] = K_beam[1:2,3:4];
    K_le[5:6,2:3] = K_beam[3:4,1:2];
    K_le[5:6,5:6] = K_beam[3:4,3:4];
    return K_le
end

# TODO: can be optimized using symbolic substitution, only X is changing here
# NOTE: r might be set to be uniform, if we are using ∑ f_e * l_e formula?
function assemble_global_stiffness_matrix(X::Matrix{Float64}, T::Matrix{Int64},
                                          r::Vector{Float64}, mp::MaterialProperties,
                                          node_dof::Int, full_node_dof::Int)
    n_elements = size(T,1)
    n_nodes = size(X,1)

    # element -> dof id map
    sys_dof::Int = node_dof * n_nodes
    id_map = zeros(Int, n_elements, 2*node_dof)
    back_dof_lin = collect(node_dof-1:-1:0)
    for i=1:n_elements
        uid = T[i,1]
        vid = T[i,2]
        id_map[i, 1:node_dof] = uid * node_dof * ones(node_dof) - back_dof_lin
        id_map[i, node_dof+1:2*node_dof] = vid * node_dof * ones(node_dof) - back_dof_lin
    end

    # cross sec properties
    sp_vec = Array{SectionProperties}(undef, n_elements)
    for e=1:n_elements
        sp_vec[e] = compute_round_section_properties(r[e])
    end

    # for elemental force calc
    KR_es = Array{Matrix{Float64},1}(undef, n_elements)

    if 2 == node_dof
        ex_id = [1, 4]
        xy_id = [1, 2, 4, 5]
    else
        ex_id = collect(1:2*node_dof)
        xy_id = collect(1:2*node_dof)
    end

    # TODO: array size node_dof^2 * n_elements
    I = Int[]
    J = Int[]
    V = Float64[]
    # K ∈ (node_dof*n_node)^2
    for e=1:n_elements
        u = X[T[e,1],:]
        v = X[T[e,2],:]
        L = norm(u - v)

        # element local stiffness matrix
        K_le = local_stiffness_matrix(L, mp, sp_vec[e])

        # local -> global rotation matrix
        R_b = local_coord_rotation(u, v)
        R = zeros(Float64, 2*full_node_dof, 2*full_node_dof)
        for k=1:(Int).(full_node_dof/3) * 2
            R[k*3-3+1:k*3, k*3-3+1:k*3] = R_b
        end

        K_Ge = R[ex_id, xy_id]' * K_le[ex_id, ex_id] * R[ex_id, xy_id]
        KR_es[e] = K_le[ex_id, ex_id] * R[ex_id, xy_id]

        for i=1:2*node_dof
            for j=1:2*node_dof
                push!(I, id_map[e,i])
                push!(J, id_map[e,j])
                push!(V, K_Ge[i,j])
            end
        end
    end
    return sparse(I, J, V, sys_dof, sys_dof), KR_es, id_map
end

function assemble_load_vector(Load::Matrix{Float64}, n_nodes::Int, node_dof::Int)
    @assert(size(Load,2) == node_dof + 1)
    sys_dof::Int = node_dof * n_nodes
    n_load_nodes = size(Load, 1)
    I = Int[]
    V = Float64[]
    for i=1:n_load_nodes
        v_id = Int.(Load[i,1])
        append!(I, node_dof*(v_id-1)+1 : node_dof*v_id)
        append!(V, Load[i, 2:1:node_dof+1])
    end
    return sparsevec(I, V, sys_dof)
end

function dof_permutation(S::Matrix{Int}, n_nodes::Int, node_dof::Int)
    sys_dof::Int = node_dof * n_nodes
    n_fix = size(S,1)
    n_dof_fix = sum(sum(S[:,2:1+node_dof] .== 1))

    full_dof_fix = zeros(Bool, sys_dof)
    for i=1:n_fix
        v_id = Int.(S[i,1])
        full_dof_fix[node_dof*(v_id-1)+1 : node_dof*v_id] = S[i, 2:1+node_dof]
        #TODO: 3D node_dof here might be incompatible, 2D is okay
    end

    perm = zeros(Int, sys_dof)
    n_dof_free = sys_dof - n_dof_fix
    free_tail = 1
    fix_tail = n_dof_free + 1
    for i=1:sys_dof
        if !full_dof_fix[i]
            perm[free_tail] = i
            free_tail += 1
        else
            perm[fix_tail] = i
            fix_tail += 1
        end
    end
    I = collect(1:sys_dof)
    perm_spm = sparse(I, perm, ones(sys_dof), sys_dof, sys_dof)
    return perm, perm_spm, n_dof_free
end

function get_weight_calculation_fn(X_full::Matrix{Float64}, r::Vector{Float64}, T::Matrix{Int64},
    F_perm_m::Vector{Float64}, perm::SparseMatrixCSC{Float64, Int}, n_dof_free::Int, node_dof::Int, full_node_dof::Int, mp::MaterialProperties, design_var_id::Vector{Int}; opt_compliance::Bool=false, kwargs...)::Function

    # assemble stiffness matrix
    n_v = size(X_full, 1)
    n_e = size(T, 1)
    sys_dof = n_v * node_dof

    # split design variable + the rest
    X_template = copy(X_full)
    design_var_id_map = zeros(Int, length(design_var_id), 2)

    function id_convert(id::Int)::Vector{Int}
        return [Int.((id + id%2) / 2), 2 - id%2]
    end

    for i=1:length(design_var_id)
        design_var_id_map[i,:] = id_convert(design_var_id[i])
    end

    # sym_var_ids = undef
    # if :sym_var_ids in keys(kwargs) && kwargs[:sym_var_ids] != undef
    #     sym_var_ids = kwargs[:sym_var_ids]
    #     @assert(size(sym_var_ids, 2) == 2)
    # end

    function calc_weight(X_var::Vector{Float64})
        for i=1:length(design_var_id)
            X_template[id_convert(design_var_id[i])[1], id_convert(design_var_id[i])[2]] = X_var[i]
        end
        # if sym_var_ids != undef
        #     for i=1:size(sym_var_ids,1)
        #         X_template[id_convert(sym_var_ids[i,2])[1], id_convert(sym_var_ids[i,2])[2]] = X_template[id_convert(sym_var_ids[i,1])[1], id_convert(sym_var_ids[i,1])[2]]
        #     end
        # end

        K, KR_es, id_map = assemble_global_stiffness_matrix(X_template, T, r, mp, node_dof, full_node_dof)
        K_perm = perm * K * perm'
        K_mm = K_perm[1:n_dof_free, 1:n_dof_free]

        # solve linear system
        # U_m = K_mm\F_perm_m
        K_mm = (K_mm + K_mm')/2
        Kpf = cholesky(K_mm, check=false)
        if !LinearAlgebra.issuccess(Kpf)
            return 1e10
        end
        U_m = Kpf\F_perm_m

        # U_I, U_V = findnz(U_m)
        U_perm = vcat(U_m, zeros(sys_dof - n_dof_free))

        # compliance
        if opt_compliance
            return F_perm_m' * U_m
        end

        U = perm' * U_perm

        # TODO: buckling sizing is ignored for now
        # calculate ∑ F_e * l_e
        weight = 0.0
        if node_dof == 2
            e_react_dof = 1 # truss
        else
            e_react_dof = 3 # frame
        end

        eF = zeros(n_e, e_react_dof*2)
        for e=1:n_e
            eF[e,:] = KR_es[e] * U[id_map[e,:]]
            eL = norm(X_template[T[e,1], :] - X_template[T[e,2], :])
            if 2 == node_dof
                weight += abs(eF[e,:][1]) * eL
            else
                # axial force only
                # weight += abs(eF[e,:][1]) * eL

                # moment only
                # weight += 4*max(abs(eF[e,:][3]), abs(eF[e,:][6]))/r[e]) * eL

                # 4M + F
                weight += (abs(eF[e,:][1]) + 4*max(abs(eF[e,:][3]), abs(eF[e,:][6]))) * eL
                # weight += (abs(eF[e,:][1])/(pi * r[e]^2) + abs(eF[e,:][3])/(pi * r[e]^3 / 4)) * eL
            end
        end
        # @show eF
        # @show U
        return weight
    end

    return calc_weight
end