"""
    get_coefficients(α,θ,m)

α should be a vector of shapes of length n, θ should be a matrix of θ of size (n,d) for a MultivariateGammaConvolution with d marginals.

This function produce the coefficients of the multivariate (tensorised) expensions of the density. It is type stable and quite optimized.
it is based on some precomputations that are done in a certain precision.

"""
function get_coefficients(α, θ, m; P = nothing)
    # α must be an array with size (n,)
    # θ must be an array with size (n,d)
    # max_p must be a Tuple of integers with size (d,)
    T = Base.promote_eltype(α,θ,[1.0]) # At least float.
    α = T.(α)
    θ = T.(θ)

    # Trim down null αs values:
    are_pos = findall(α .> eps(T))
    if ndims(θ) == 1
        θ = θ[are_pos]
    else
        θ = θ[are_pos,:] # produces a matrix when θ is only a vector..
    end
    α = α[are_pos]
    if isnothing(P)
        P = get_precomp(T,sum(m))::PreComp{T}
    end
    coefs = zeros(T,m)
    build_coefficients!(coefs,α,θ,T(1.0),m,P)
    return coefs
end

function build_coefficients!(coefs,α,θ,cst1,m,P)

    # This is the default method, for higher dimensions.

    # Allocates ressources
    κ = deepcopy(coefs)
    μ = deepcopy(coefs)
    n = size(θ)[1]
    d = length(m)
    I = CartesianIndices(coefs)
    S = θ ./ (cst1 .+ sum(θ,dims=2))
    S_pow = [s^k for k in (0:Base.maximum(m)), s in S]

    # Starting the algorithm: there is an edge case for the Oth cumulant, 0th moment and 0th coef:
    to_log = cst1 .- sum(S,dims=2)
    to_log .= ifelse.(to_log .< 0, to_log .+ eps(Base.eltype(α)), to_log)
    κ[1] = sum(α .* log.(to_log))
    μ[1] = exp(κ[1])
    coefs[1] = μ[1]

    @inbounds for k in I[2:length(I)]
        # Indices and organisation
        k_arr = Tuple(k)
        degree = findfirst(k_arr .!= 1)
        sk = sum(k_arr)

        # Main Computations:
        for i in 1:n
            rez = α[i]
            for j in 1:d
                rez *= S_pow[k[j],i,j]
            end
            κ[k] += rez
        end
        κ[k] *= P.FACTS[sk-d]

        for j in CartesianIndices(k)
            rez_coefs = μ[j]
            if j[degree] < k[degree]
                rez_mu = μ[j] * κ[k - j + I[1]]
                for i in 1:d
                    rez_mu *= P.BINS[j[i],k[i]-Int(i==degree)]
                    rez_coefs *= P.LAGUERRE[j[i], k[i]]
                end
                μ[k] += rez_mu
            else
                for i in 1:d
                    rez_coefs *= P.LAGUERRE[j[i], k[i]]
                end
            end
            coefs[k] += rez_coefs
        end
    end
    coefs .*= sqrt(2*cst1)^d
end

function build_coefficients!(coefs::AbstractArray{T,1},α::AbstractArray{T,1},θ::AbstractArray{T,1},cst1::T,m::NTuple{1, T2},P::PreComp{T}) where {T2 <: Int,T <: Real}

    #kappa = Array{Base.eltype(α), 1}(undef, m[1])
    kappa = zeros(T,m[1])
    mu = zeros(T, m[1])

    @. θ /= 1 + θ

    # Edge case for the Oth cumulant, 0th moment and 0th coef:
    to_log = cst1 .- θ
    to_log .= ifelse.(to_log .< 0, to_log .+ eps(Base.eltype(α)), to_log) # correcting rounding errors.
    kappa[1] = sum(α .* log.(to_log))
    coefs[1] = mu[1] = exp(kappa[1]) # this is the only place kappa[1] is used.
    #coefs[1] = mu[1] = prod(to_log .^ α)
    @inbounds for k in 2:m[1]
        # Main Computations:
        α .*= θ
        kappa[k] = sum(α)
        @inbounds for j in 1:k-1
            mu[k] += mu[j] * kappa[k-j+1]
        end
        mu[k] /= k-1
        @views coefs[k] = P.LAGUERRE2[1:k, k]'mu[1:k]
    end
end

function build_coefficients!(coefs::AbstractArray{T,1},α::AbstractArray{T,1},θ::AbstractArray{T,2},cst1::T,m::NTuple{2, T2},P::PreComp{T}) where {T2 <: Int,T <: Real}
    # Allocates ressources
    κ = Array{Base.eltype(α), 2}(undef, m)
    μ = zeros(Base.eltype(α), m)

    #n = size(θ)[1]
    d = length(m)
    I = CartesianIndices(coefs)
    θ ./= (cst1 .+ sum(θ,dims=2))
    S_pow = [s^k for k in (0:Base.maximum(m)), s in θ]

    # reduce the number of multiplications :
    for i in 1:size(α,1)
        S_pow[:,i,:] .*= sqrt(α[i])
    end

    # Starting the algorithm: there is an edge case for the Oth cumulant, 0th moment and 0th coef:
    to_log = cst1 .- sum(θ,dims=2)
    to_log .= ifelse.(to_log .< 0, to_log .+ eps(Base.eltype(α)), to_log)
    κ[1] = sum(α .* log.(to_log)) # this log fails ifsum(S,dims=2) is greater than 1, which occurs because of rounding errors. We add an eps to avoid that.
    μ[1] = exp(κ[1])
    coefs[1] = μ[1]

    @inbounds for k in I[2:length(I)]
        # Indices and organisation
        degree = k[1]!=1 ? 1 : 2
        deg1 = Int(1==degree)
        deg2 = Int(2==degree)

        κ[k] = sum(S_pow[k[1],:,1] .* S_pow[k[2],:,2]) * P.FACTS[sum(Tuple(k))-d]
        for j in CartesianIndices(k)
            if j[degree] < k[degree]
                μ[k] += μ[j] * κ[k - j + I[1]] *  P.BINS[j[1],k[1]-deg1] * P.BINS[j[2],k[2]-deg2]
            end
            coefs[k] += μ[j] * P.LAGUERRE[j[1], k[1]] * P.LAGUERRE[j[2], k[2]]
        end
    end
    coefs .*= sqrt(2*cst1)^d
end

"""
    laguerre_L_2x(x,p)

Compute univariate laguerre polynomials Lₚ(2x).
"""
function laguerre_L_2x(x,p)
    P = get_precomp(Base.eltype(x),p+1)
    sum(P.LAGUERRE[1:(p+1),p+1] .* x .^ (0:p))
end

"""
    laguerre_phi(x,p)

Computes the function ϕₚ(x) = √2 e⁻ˣ Lₚ(2x). This functions form an orthonormal basis of R₊
"""
function laguerre_phi(x,p)
    exp(-x) * sqrt(2) * laguerre_L_2x(x,p)
end

"""
    laguerre_density(x,coefs)

Given some laguerre coefficients, computes the correpsonding function at the point x.
"""
function laguerre_density(x,coefs)
    rez = zero(Base.eltype(coefs))
    for p in CartesianIndices(coefs)
        rez += coefs[p] * prod(laguerre_phi.(x,Tuple(p) .-1))
    end
    return max(rez,0)
end

"""
    laguerre_phi_several_pts(x,max_p)

Computes the laguerre_phi for each row of x and each p in CartesianIndex(maxp)
This is a lot more efficient than broadcasting the laguerre_phi function,
but this is ugly in the sense that we re-wrote some of the code. This machanisme could be re-written more efficiently.
"""
function laguerre_phi_several_pts(x::AbstractArray{T},max_p) where {T <: Real}
    P = get_precomp(T,Base.maximum(max_p)+1)::PreComp{T}
    d,n = size(x)
    rez = ones(T,(n,max_p...))
    MP = Base.maximum(max_p)

    # println("Computing exponentials...")
    exponentials = dropdims(exp.(-sum(x,dims=1)),dims=1)

    laguerre_L = zeros(T,(d,MP,n))
    # println("Computing powers...")
    powers = x[:,na,:] .^ (0:(MP-1))[na,:,na]
    # println("Computing laguerre_L")
    Threads.@threads for p in 1:MP
        laguerre_L[:,p:p,:] = dropdims(sum(P.LAGUERRE[1:p,p:p][na,:,:,na] .* powers[:,1:p,na,:],dims=2),dims=2)
        print(p,"\n")
    end

    # println("Computing L(2x)")
    Threads.@threads for p in CartesianIndices(max_p)
        for i in 1:d
            rez[[CartesianIndex((i,Tuple(p)...)) for i in 1:n]] .*= laguerre_L[i,p[i],:]
        end
        print(p,"\n")
    end
    return rez .* sqrt(T(2))^d .* exponentials
end

"""
    empirical_coefs(x,maxp)

Efficiently Compute the empirical laguerre coefficients of the density of the random vector x, given as a Matrix with n row (number of samples) and d columns (number of dimensions)
"""
function empirical_coefs(x,max_p)
    T = Base.promote_eltype(x,[1.0])
    x = T.(x)
    y = laguerre_phi_several_pts(x,max_p)
    # println("Taking the means...")
    return dropdims(sum(y,dims=1)/size(y,1),dims=1)
end

function empirical_coefs(x,max_p::NTuple{2,M}) where M<:Int
    T = Base.promote_eltype(x,[1.0])
    x = T.(x)
    P = get_precomp(T,Base.maximum(max_p)+1)::PreComp{T}
    d,n = size(x)
    rez = ones(T,max_p)
    MP = Base.maximum(max_p)

    exponentials = dropdims(exp.(-sum(x,dims=1)),dims=1)

    laguerre_L = zeros(T,(d,MP,n))
    powers = x[:,na,:] .^ (0:(MP-1))[na,:,na]
    Threads.@threads for p in 1:MP
        laguerre_L[:,p:p,:] = dropdims(sum(P.LAGUERRE[1:p,p:p][na,:,:,na] .* powers[:,1:p,na,:],dims=2),dims=2)
    end

    Threads.@threads for p in CartesianIndices(max_p)
        rez[p] = sum(laguerre_L[1,p[1],:] .* laguerre_L[2,p[2],:] .* exponentials)
    end
    rez = (rez .* 2) ./ n
    return rez
end

"""
    L2Objective(par,emp_coefs)

A L2 distance to be minimized between laguerre coefficients of MultivariateGammaConvolutions and empirical laguerre coefficients.
This distance is some kind of very high degree polynomial * exponentials, so minimizing it is very hard.
We recomand using some global minimisation routine, we found the Particle Swarm optimizers from Optim.jl particularly efficient.
"""
function L2Objective(par, emp_coefs)
    m = size(emp_coefs)
    d = length(m)
    n = Int((size(par)[1])/(d+1))
    α = par[1:n] .^2 #make them positives
    rates = reshape(par[(n+1):(d+1)*n],(n,d)) .^ 2 # make them positives
    coefs = get_coefficients(α,rates,m)
    return sum((coefs .- emp_coefs) .^2)
end

function αθ_from_par(par)
    n = Int(length(par)/2)
    return par[1:n].^2, par[(n+1):2n].^2
end


using TransformVariables
using TransformVariables: AbstractTransform, CallableTransform, CallableInverse

export Pushforward
export Pullback

struct Pushforward{F,M} <: AbstractMeasure
    f::F
    μ::M
    logjac::Bool
end

Pushforward(f,μ) = Pushforward(f, μ, true)

struct Pullback{F,M} <: AbstractMeasure
    f::F
    ν::M
    logjac::Bool
end

Pullback(f,ν) = Pullback(f, ν, true)

function logdensity(pb::Pullback{F}, x) where {F <: CallableTransform}
    f = pb.f
    ν = pb.ν
    if pb.logjac
        y, logJ = transform_and_logjac(f.t, x)
        return logdensity(ν, y) + logJ
    else
        y = f(x)
        return logdensity(ν, y)
    end
end

function logdensity(pf::Pushforward{F}, y) where {F <: CallableTransform}
    f = pf.f
    μ = pf.μ
    x = inverse(f.t)(y)
    if pf.logjac
        _, logJ = transform_and_logjac(f.t, x)
        return logdensity(μ, x) - logJ
    else
        return logdensity(μ, x)
    end
end

Pullback(f::AbstractTransform, ν, logjac::Bool=true) = Pullback(transform(f), ν, logjac)
Pushforward(f::AbstractTransform, ν, logjac::Bool=true) = Pushforward(transform(f), ν, logjac)

Pullback(f::CallableInverse, ν, logjac::Bool=true) = Pushforward(transform(f.t), ν, logjac)

Pushforward(f::CallableInverse, ν, logjac::Bool=true) = Pullback(transform(f.t), ν, logjac)

Base.rand(rng::AbstractRNG, T::Type, ν::Pushforward) = ν.f(rand(rng, T, ν.μ))

Base.rand(rng::AbstractRNG, T::Type, μ::Pullback) = μ.f(rand(rng, T, μ.ν))

testvalue(ν::Pushforward) = transform(ν.f, testvalue(ν.μ))

testvalue(μ::Pullback) = transform(inverse(μ.f), testvalue(μ.ν))

basemeasure(μ::Pullback) = Pullback(μ.f, basemeasure(μ.ν), false)

basemeasure(ν::Pushforward) = Pushforward(ν.f, basemeasure(ν.μ), false)

TV.as(ν::Pushforward) = ν.f ∘ as(ν.μ)

TV.as(μ::Pullback) = inverse(μ.f) ∘ μ.ν

TV.as(::Lebesgue) = asℝ

# TODO: Make this work for affine embeddings
TV.as(d::Affine) = _as_affine(_firstval(d))

_firstval(d::Affine) = first(values(getfield(getfield(d, :f), :par)))
_as_affine(x::Real) = asℝ
_as_affine(x::AbstractArray) = as(Vector, size(x,1))

basemeasure(::Pushforward{TV.CallableTransform{T}, Lebesgue{ℝ}}) where {T <: TV.ScalarTransform} = Lebesgue(ℝ)
basemeasure(::Pullback{TV.CallableTransform{T}, Lebesgue{ℝ}}) where {T <: TV.ScalarTransform} = Lebesgue(ℝ)
# t = as𝕀
# μ = Normal()
# ν = Pushforward(t, μ)
# x = rand(μ)
# julia> logdensity(μ, x) ≈ logdensity(Pushforward(inverse(t), ν), x)
# true


using DataFrames

# Create toy data ---------------------------------
# const COUNTRIES = ["CZE", "DEU", "HUN", "IDN"]
# const PRODUCTS = ["1602", "1604", "1605", "1701", "1702", "1704", "1805", "1806", "1901", "1902"]

POP = [10474410;
       81776930;
       10000023;
       241834215]

POP = convert(Array{Float64, 1}, POP) # Some pop-models might have "half people"

EXP =  [5.66868e7 1.40287e7 1.76446e6 1.24263e8 4.21838e6 1.54845e8 3.0603e6 1.56225e8 3.14846e7 2.60488e7; # CZE
	1.231e9 6.55645e8 7.13669e7 8.70726e8 3.60293e8 8.40892e8 2.75906e8 3.22708e9 1.16899e9 1.82894e8; # DEU
	1.2379e8 1.00424e6 6.27152e6 1.3462e8 7.32589e7 7.9118e7 1.37126e6 7.98984e7 3.50975e7 3.30034e7; # HUN
	5.3835e5 2.63427e8 3.20992e8 1.91463e6 1.57506e7 9.4886e7 1.10084e8 4.17582e7 5.15586e7 1.32186e8] # IDN

# Test RCA ----------------------------------------
@testset "RCA tests" begin
	# Check inputs:
	@test_throws MethodError get_rca(DataFrame(EXP))
	@test_throws MethodError get_rca(EXP[1, :])
	@test_throws MethodError get_rca(["1" "2";
					  "3" "4"])

	# Check output types:
	@test typeof(get_rca(EXP, binary=true)) == Array{Int64, 2}
	@test typeof(get_rca(EXP, binary=false)) == Array{Float64, 2}

	# Check output result:
	@test get_rca(EXP, binary=false)[4, 3] == (EXP[4, 3] / sum(EXP[4,:])) / (sum(EXP[:,3]) / sum(EXP))
	@test get_rca(EXP, binary=false)[3, 2] == (EXP[3, 2] / sum(EXP[3,:])) / (sum(EXP[:,2]) / sum(EXP))

	@test get_rca(EXP, binary=true)[4, 3] == 1
	@test get_rca(EXP, binary=true)[3, 2] == 0
end

# Test RpcA ---------------------------------------
@testset "RPCA tests" begin
	# Check inputs:
	@test_throws MethodError get_rpca(EXP, convert(Array{Int64, 1}, POP))
	@test_throws MethodError get_rpca(DataFrame(EXP), POP)
	@test_throws MethodError get_rpca(EXP[1, :], POP)
	@test_throws MethodError get_rpca(["1" "2"
					   "3" "4"], POP)

	# Check output types:
	@test typeof(get_rpca(EXP, POP, binary=true)) == Array{Int64, 2}
	@test typeof(get_rpca(EXP, POP, binary=false)) == Array{Float64, 2}

	# Check output result:
	@test get_rpca(EXP, POP, binary=false)[4, 3] == (EXP[4, 3] / POP[4]) / (sum(EXP[:, 3]) / sum(POP))
	@test get_rpca(EXP, POP, binary=false)[3, 2] == (EXP[3, 2] / POP[3]) / (sum(EXP[:, 2]) / sum(POP))

	@test get_rpca(EXP, POP, binary=true)[4, 3] == 1
	@test get_rpca(EXP, POP, binary=true)[3, 2] == 0
end


using FCCQuad: Fct, Cheb, getweights, fccquad_cc, fccquad_cs, fccquad_sc, fccquad_ss, all_methods, supported_types, interval_methods
using QuadGK, Test

#=
Integral of cos(13cos(2x))*52sin(x)cos(x)dx from 0 to 17pi/4
  == integral of cos(u)*du from 0 to 13 == sin(13).
Integral of sin(13cos(2x))*52sin(x)cos(x)dx from 0 to 17pi/4
  == integral of sin(u)*du from 0 to 13 == 1-cos(13).
=#
@testset "analytic integrals" begin
    println("Integral of (cos or sin)(13cos(2x))*52sin(x)cos(x)dx from 0 to 17pi/4:")
    angle(x)=13cos(2x)
    ampcos(x)=52cos(x)
    ampsin(x)=52sin(x)
    println("\tmethod \tdata type \tintegrand evaluation counts")
    for mthd in all_methods
        for T in supported_types[mthd]
            print("\t$mthd \t$T \t")
            R=real(T)
            freqs=[one(R)]
            kwargs=(:xmin=>zero(R),:xmax=>17*R(pi)/4,:method=>mthd,:T=>R)

            exact=sin(R(13))
            result=fccquad_cc(ampsin,angle,freqs;kwargs...)
            print("$(result[2]) \t")
            @test isapprox(result[1][1,1],exact)

            result=fccquad_cs(ampcos,angle,freqs;kwargs...)
            print("$(result[2]) \t")
            @test isapprox(result[1][1,1],exact)

            exact=one(R)-cos(R(13))
            result=fccquad_sc(ampsin,angle,freqs;kwargs...)
            print("$(result[2]) \t")
            @test isapprox(result[1][1,1],exact)

            result=fccquad_ss(ampcos,angle,freqs;kwargs...)
            println("$(result[2])")
            @test isapprox(result[1][1,1],exact)
        end
    end
end

#slow type I DCT (for testing purposes):
#(dct(x_0,...,x_N))_n = sum_{0<=m<=N}x_m*cos(pi*m*n/N)
#with 1/2 weight given to first and last summands.
#N can be any positive integer
function dct(x::AbstractVector,T::Type=Complex{Float64})::Vector{T} #type I DCT
    n=length(x)-1
    y=Vector{T}(undef,n+1)
    p = real(T)(pi)
    for i in 0:n
        y[i+1]=0.5x[1]
        @simd for j in 1:n-1
            @inbounds y[i+1]+=x[j+1]*cos(p*i*j/n)
        end
        y[i+1]+=x[n+1]*ifelse(i%2==0,0.5,-0.5)
    end
    y
end

@testset "Comparing fast (I)DCT to simpler implementation" begin
    x0=rand(Complex{Float64},129)
    y1=Fct.fct(x0)
    y2=dct(x0)
    @test isapprox(y1,y2)
    x1=Fct.idct(y1)
    @test isapprox(x0,x1)
end

@testset "Chebyshev series evaluation, dilation, and multiplication" begin
    f=Dict(:cube=>x->x^3, :sin=>sin, :cos=>cos, :exp=>exp)
    c=Dict(:cube=>Cheb.OddChebSeries(x->x^3),
           :sin=>Cheb.OddChebSeries(sin),
           :cos=>Cheb.EvenChebSeries(cos),
           :exp=>Cheb.ChebSeries(exp))
    x0=-0.5
    kappa=0.9
    for k in keys(f)
        @test isapprox(Cheb.chebinterp(c[k],x0),f[k](x0))
        @test isapprox(Cheb.chebinterp(Cheb.dilate(c[k],kappa),x0),f[k](kappa*x0))
        for j in keys(f)
            @test isapprox(Cheb.chebinterp(c[j]*c[k],x0),f[j](x0)*f[k](x0))
        end
    end
end

#T_n(x)
function chebpoly(n::Integer,x)
    @assert n>=0
    n<=1 && return iszero(n) ? one(x) : x
    m = div(n,2)
    iseven(n) && return 2chebpoly(m,x)^2 - one(x) #T_{2m}(x)
    a,b = chebpair(m,x) #T_m(x), T_{m+1}(x)
    2a*b - x #T_{2m+1}(x)
end

#T_n(x),T_{n+1}(x)
function chebpair(n::Integer,x)
    @assert n>=0
    n==0 && return (one(x),x)
    m = div(n,2)
    e,o = chebpair(m,x) #T_m, T_{m+1}
    if iseven(n) #n == 2m
        2e^2 - one(x), 2e*o - x #T_{2m}, T_{2m+1}
    else # n == 2m+1
        2o*e - x, 2o^2 - one(x) #T_{2m+1}, T_{2m+2}
    end
end

@testset "Chebyshev weights" begin
    function checkweights(N,w,args...)
        a = getweights(N,w,args...)
        atol2 = N*eps()^2
        gk_atol=eps()
        gk_rtol=10^4 * eps()
        for n in 0:2:N
        gk,delta = collect(quadgk(x->chebpoly(n,x)*cos(w*x),-1,1,rtol=gk_rtol,atol=gk_atol))
            pass = abs2(a[n+1]-gk) <= max(atol2,abs2(delta))
            pass || println("N=$N,n=$n,omega=$w,weight=$(a[n+1]),gk=$gk,gk_delta=$delta")
            @test pass
        end
        for n in 1:2:N
            gk,delta = collect(quadgk(x->chebpoly(n,x)*sin(w*x),-1,1,rtol=gk_rtol,atol=gk_atol))
            pass = abs2(a[n+1]-gk) <= max(atol2,abs2(delta))
            pass || println("N=$N,n=$n,omega=$w,weight=$(a[n+1]),gk=$gk,gk_delta=$delta")
            @test pass
        end
    end
    println("Testing Chebyshev weights (0<=n<=100):")
    for freq in [0,-1e-9,1e-6,-1e-3,1e-2,-1e-1,1,2,3,4,5,-99.5,200]
        println("\ttesting angular frequency $freq...")
        checkweights(100,freq)
    end
end


# You've almost reached the exit of the cave, but the walls are getting closer
# together. Your submarine can barely still fit, though; the main problem is
# that the walls of the cave are covered in chitons, and it would be best not
# to bump any of them.
#
# The cavern is large, but has a very low ceiling, restricting your motion to
# two dimensions. The shape of the cavern resembles a square; a quick scan of
# chiton density produces a map of risk level throughout the cave (your puzzle
# input). For example:
#
# 1163751742
# 1381373672
# 2136511328
# 3694931569
# 7463417111
# 1319128137
# 1359912421
# 3125421639
# 1293138521
# 2311944581
#
# You start in the top left position, your destination is the bottom right
# position, and you cannot move diagonally. The number at each position is its
# risk level; to determine the total risk of an entire path, add up the risk
# levels of each position you enter (that is, don't count the risk level of
# your starting position unless you enter it; leaving it adds no risk to your
# total).
#
# Your goal is to find a path with the lowest total risk. In this example, a
# path with the lowest total risk is highlighted here:
#
# 1163751742
# 1381373672
# 2136511328
# 3694931569
# 7463417111
# 1319128137
# 1359912421
# 3125421639
# 1293138521
# 2311944581
#
# The total risk of this path is 40 (the starting position is never entered, so
# its risk is not counted).
#
# What is the lowest total risk of any path from the top left to the bottom
# right?

# load risk map
riskmap = hcat(map(l -> parse.(Int8, split(l, "")), eachline("input.txt"))...)
(width, height) = size(riskmap)

# some data structures to track things
cost = zeros(Int, size(riskmap))
queue = [(1,1)]
function update_cost(x, y, c)
  c += riskmap[x,y]
  if cost[x,y] == 0
    # first time setting a cost on this node - add to queue
    push!(queue, (x,y))
  elseif cost[x,y] <= c
    # we found a cheaper route to this node already
    return
  end
  cost[x,y] = c
end

# Dijkstra's
visited = falses(size(riskmap))
while !visited[width, height]
  # find the next node with the smallest tentative cost
  sort!(queue; by=((x,y),) -> cost[x,y], rev=true, alg=InsertionSort)
  (x,y) = pop!(queue)
  visited[x,y] = true

  # update cost of neighbors
  c = cost[x,y]
  x > 1 && !visited[x-1, y] && update_cost(x-1, y, c)
  y > 1 && !visited[x, y-1] && update_cost(x, y-1, c)
  x < width && !visited[x+1, y] && update_cost(x+1, y, c)
  y < height && !visited[x, y+1] && update_cost(x, y+1, c)
end

println("Result: ", cost[width, height])


@doc raw"""
    Bbox2

An object `b` of the type [`Bbox2`](@ref) is a bounding box in the
two-dimensional Euclidean plane ``Ε^2``.
"""
Bbox2

@doc raw"""
    Bbox2()

Introduces an *empty* bounding box with lower left corner point at ``(∞, ∞)``
and with upper right corner point at ``(-∞, -∞)``, ``∞`` being [`Inf`](@ref).
"""
Bbox2()

"""
    Bbox2(xmin::Real, ymin::Real, xmax::Real, ymax::Real)

Introduces a bounding box `b` with lower left corner at `(xmin, ymin)` and with
upper right corner at `(xmax, ymax)`.
"""
Bbox2(xmin::Real, ymin::Real, xmax::Real, ymax::Real) =
    Bbox2(convert(Float64, xmin), convert(Float64, ymin),
          convert(Float64, xmax), convert(Float64, ymax))

"""
    dimension(b::Bbox2)

Returns 2.
"""
dimension(b::Bbox2)

"""
    min(b::Bbox2, i::Integer)

Returns [`xmin(b)`](@ref) if `i == 0` or [`ymin(b)`](@ref) if `i == 1`.

!!! info "Precondition"

    `i == 0` or `i == 1`.
"""
min(b::Bbox2, i::Integer)

"""
    max(b::Bbox2, i::Integer)

Returns [`xmax(b)`](@ref) if `i == 0` or [`ymax(b)`](@ref) if `i == 1`.

!!! info "Precondition"

    `i == 0` or `i == 1`.
"""
max(b::Bbox2, i::Integer)

"""
    +(b::Bbox2, c::Bbox2)

Returns a bounding box of `b` and `c`.
"""
+(b::Bbox2, c::Bbox2)

"""
    dilate(b::Bbox2, dist::Integer)

Dilates the bounding box by a specified number of ULP.
"""
dilate(b::Bbox2, dist::Integer)


# vectors

a = 1.2v1 + 1.89v3 + 1.2v4

# Homogeneous multivectors

A = 1.3v1 + 2.7v12
B = 0.7v3 + 1.1v123 + 3.05v
C = 0.7v4 + 0.1v2 + 1.59v23 + 1.1v124 + 3.05v1234

A₀ = 0.9v
A₁ = 1.1v1
A₂ = 1.96v23
A₃ = 1.2v123
A₄ = 1.0v1234

B₀ = 2.3v
B₁ = 0.8v2
B₂ = -1.8v13
B₃ = 2.5v234
B₄ = 1.7v1234

C₀ = 0.3v
C₁ = 0.4v4
C₂ = 0.92v24
C₃ = 1.0v134
C₄ = 1.1v1234

As = (A₀, A₁, A₂, A₃, A₄)
Bs = (B₀, B₁, B₂, B₃, B₄)
Cs = (C₀, C₁, C₂, C₃, C₄)

Zs = tuple(As..., Bs..., Cs...)

# Non-homogeneous multivectors

SA = sum(As)
SB = sum(Bs)
SC = sum(Cs)

NHB₀ = 2.3v
NHB₁ = 0.3v1 + 0.8v2
NHB₂ = 0.47v12 + 0.92v23 - 1.8v13
NHB₃ = 1.0v123 + 2.5v234
NHB₄ = 1.7v1234

NHC₀ = 0.3v
NHC₁ = 0.4v4 + 0.3v2
NHC₂ = 0.47v13 + 0.92v24 - 1.8v14
NHC₃ = 1.0v134 + 2.5v124
NHC₄ = 1.1v1234

NHBs = (NHB₀, NHB₁, NHB₂, NHB₃, NHB₄)
NHCs = (NHC₀, NHC₁, NHC₂, NHC₃, NHC₄)

NHZs = tuple(NHBs..., NHCs...)


#=
day15:
- Julia version: 1.7.0
- Author: Paul.Mealor
- Date: 2021-12-15
=#

using AoC, Test

neighbourOffsets(dims) = CartesianIndex{dims}[
    [CartesianIndex(ntuple(i->i==j ? 1 : 0, dims)) for j in 1:dims];
    [CartesianIndex(ntuple(i->i==j ? -1 : 0, dims)) for j in 1:dims]
]

function findBestMove(cave)
    risks = fill(length(cave)*max(cave...), size(cave))
    risks[1] = 0

    moves = neighbourOffsets(ndims(cave))

    finished = false
    while !finished
        finished = true
        for I in CartesianIndices(cave)
            for move in moves
                J = I + move
                if checkbounds(Bool, cave, J)
                    risk = risks[I] + cave[J]
                    if risk < risks[J]
                        risks[J] = risk
                        finished = false
                    end
                end
            end
        end
    end

    return risks[end]
end
@test findBestMove([[1,2] [3,4]]) == 6
@test findBestMove([[1,3,1] [1,1,1] [2,1,1]]) == 4

function embiggen(cave)
    realCave = similar(cave, size(cave) .* 5)

    for I in CartesianIndices(cave)
        baseValue = cave[I]
        for j in I[2]:size(cave,2):size(realCave,2)
            value = baseValue
            for i in I[1]:size(cave,1):size(realCave,1)
                realCave[i,j] = value
                value = value < 9 ? value + 1 : 1
            end
            baseValue = baseValue < 9 ? baseValue + 1 : 1
        end
    end
    return realCave
end
@test embiggen(exampleIntMap(15,1))[1:10,1:10] == exampleIntMap(15,1)
@test embiggen(exampleIntMap(15,1)) == exampleIntMap(15,2)

part1(cave) = findBestMove(cave)

part2(cave) = findBestMove(embiggen(cave))

@test part1(exampleIntMap(15,1)) == 40
@test part2(exampleIntMap(15,1)) == 315

@show exampleIntMap(15,1) |> ll -> @time part1(ll)
@show exampleIntMap(15,1) |> ll -> @time part2(ll)

@show intMap(15) |> ll -> @time part1(ll)
@show intMap(15) |> ll -> @time part2(ll)


"""
    nonlinear stage-wise objective
"""
struct NonlinearStageObjective <: Objective
    l_stage
    l_terminal
end

function nonlinear_stage_objective(l_stage, l_terminal)
    return NonlinearStageObjective(l_stage, l_terminal)
end

function objective(Z, obj::NonlinearStageObjective, idx, T)
    J = 0.0

    for t = 1:T-1
        x = view(Z, idx.x[t])
        u = view(Z, idx.u[t])

        J += obj.l_stage(x, u, t)
    end

    x = view(Z, idx.x[T])
    J += obj.l_terminal(x)

    return J
end

function objective_gradient!(∇J, Z, obj::NonlinearStageObjective, idx, T)

    for t = 1:T-1
        x = view(Z, idx.x[t])
        u = view(Z, idx.u[t])

        lx(y) = obj.l_stage(y, u, t)
        lu(y) = obj.l_stage(x, y, t)

        ∇J[idx.x[t]] += ForwardDiff.gradient(lx, x)
        ∇J[idx.u[t]] += ForwardDiff.gradient(lu, u)

    end

    x = view(Z, idx.x[T])
    ∇J[idx.x[T]] += ForwardDiff.gradient(obj.l_terminal, x)

    return nothing
end


# indexnotation/product.jl
#
# A wrapper to store the contraction (product) of two indexed objects and
# evaluate lazily, i.e. evaluate upon calling `deindexify`.

immutable ProductOfIndexedObjects{IA,IB,CA,CB,OA,OB,TA,TB} <: AbstractIndexedObject
    A::IndexedObject{IA,CA,OA,TA}
    B::IndexedObject{IB,CB,OB,TB}
end

Base.conj(P::ProductOfIndexedObjects) = ProductOfIndexedObjects(conj(P.A),conj(P.B))

Base.eltype(P::ProductOfIndexedObjects) = promote_type(eltype(P.A),eltype(P.B))

*(P::ProductOfIndexedObjects, β::Number) = ProductOfIndexedObjects(β*P.A, P.B)
*(β::Number, P::ProductOfIndexedObjects) = *(P,β)
-(P::ProductOfIndexedObjects) = *(-1, P)

@generated function indices{IA,IB,CA,CB,OA,OB,TA,TB}(::Type{ProductOfIndexedObjects{IA,IB,CA,CB,OA,OB,TA,TB}})
    J = unique2([IA...,IB...])
    J = tuple(J...)
    meta = Expr(:meta, :inline)
    Expr(:block, meta, :($J))
end

@generated function *(A::AbstractIndexedObject,B::AbstractIndexedObject)
    IA = indices(A)
    IB = indices(B)
    IC = symdiff(IA,IB)
    oindA, cindA, oindB, cindB, = contract_indices(IA,IB,IC)
    meta = Expr(:meta,:inline)
    if A <: IndexedObject{IA}
        argA = :A
    else
        IA = IA[vcat(oindA,cindA)]
        IA = tuple(IA...)
        indA = :(Indices{$IA}())
        argA = :(indexify(deindexify(A,$indA),$indA))
    end
    if B <: IndexedObject{IB}
        argB = :B
    else
        IB = IB[vcat(cindB,oindB)]
        IB = tuple(IB...)
        indB = :(Indices{$IB}())
        argB = :(indexify(deindexify(B,$indB),$indB))
    end
    meta = Expr(:meta,:inline)
    Expr(:block,meta,:(ProductOfIndexedObjects($argA,$argB)))
end

@generated function deindexify{IA,IB,CA,CB,IC}(P::ProductOfIndexedObjects{IA,IB,CA,CB}, I::Indices{IC}, T::Type = eltype(P))
    meta = Expr(:meta, :inline)
    oindA, cindA, oindB, cindB, indCinoAB = contract_indices(IA, IB, IC)
    indCinAB = vcat(oindA,length(IA)+oindB)[indCinoAB]
    conjA = Val{CA}
    conjB = Val{CB}
    quote
        $meta
        deindexify!(similar_from_indices(T, $indCinAB, P.A.object, P.B.object, $conjA, $conjB), P, I)
    end
end

@generated function deindexify!{IA,IB,CA,CB,IC}(dst, P::ProductOfIndexedObjects{IA,IB,CA,CB}, ::Indices{IC}, β=0)
    oindA, cindA, oindB, cindB, indCinoAB = contract_indices(IA, IB, IC)
    conjA = Val{CA}
    conjB = Val{CB}
    meta = Expr(:meta,:inline)
    quote
        $meta
        contract!(P.A.α*P.B.α,P.A.object,$conjA,P.B.object,$conjB,β,dst,$oindA,$cindA,$oindB,$cindB,$indCinoAB)
    end
end


# Calculate and plot in and out degrees of generated Kronecker Graph
using JLD,PyPlot

# Associative arrays to accumulate in and out degrees
Aall = Assoc("","","")

# Iterate through files
(@isdefined Nfile) || (Nfile = 8)
for i = 1:Nfile
    global Aall
    ptime = @elapsed begin
        fname = joinpath(Base.source_dir(),"data", string(i))

        # Load associative array
        A = load(fname*".A.jld")["A"]
        #A = loadassoc(fname*".A.jld")

        # Create associative arrays and accumulate degrees
        Aall = A + Aall
    end

    println("Sum Time: ", ptime)
    println("Edges/sec: ", string(nnz(A)/ptime))
end

# Plot out degree distribution
figure()
loglog(convert(Array,OutDegree(Aall)'),"o")
xlabel("out degree")

# Plot out degree distribution
figure()
loglog(convert(Array,InDegree(Aall)'),"o")
xlabel("in degree")


#export Simplex

# U: the dimension of the universe
# D: the dimension of the manifold
# N: the number of vertices
# T: the type of the coordinates
# C: the complimentary dimension (should always be U-D)
struct Simplex{U,D,C,N,T}
    vertices::SVector{N,SVector{U,T}}
    tangents::SVector{D,SVector{U,T}}
    normals::SVector{C,SVector{U,T}}
    volume::T
end

"""
    coordtype(simplex)

Return coordinate type used by simplex.
"""
coordtype(::Type{Simplex{U,D,C,N,T}}) where {U,D,C,N,T} = T
coordtype(p::Simplex) = coordtype(typeof(p))

"""
    volume(simplex)

Return the volume of the simplex.
"""
volume(p::Simplex) = p.volume

"""
A tuple of points, aka an interval behaves trivially like a chart
"""
volume(x::Tuple{T,T}) where {T} = norm(x[2]-x[1])

"""
    dimension(simplex)

Return the manifold dimension of the simplex.
"""
dimension(::Type{Simplex{U,D,C,N,T}}) where {U,D,C,N,T} = D

"""
    dimension(simplex)

Return the manifold dimension of the simplex.
"""
dimension(p::Simplex) = dimension(typeof(p))

"""
    length(simplex)

Returns the number of vertices (equals dimension + 1)
"""
Base.length(p::Simplex) = dimension(typeof(p))+1

"""
  universedimension(p)

Return the dimension of the universe in which `p` is embedded.
"""
universedimension(::Type{Simplex{U,D,C,N,T}}) where {U,D,C,N,T} = U
universedimension(p::Simplex) = universedimension(typeof(p))

"""
    getindex(simplex, I)
    simplex[I]

Get the vertices at index I (scalar or array) defining the simplex
"""
getindex(p::Simplex, I::Union{Number,SVector,Array}) = p.vertices[I]

"""
    simplex(vertices)
    simplex(v1, v2, ...)
    simplex(vertices, Val{D})

Build a D-dimensional simplex. The vertices can be passed in
an array (static or dynamic), or supplied separately. If the
length of the array is not part of its type, the speed of the
construction can be improved by supplying an extra Val{D}
argument. In case it is not clear from the context whether
the vertex array is dynamically or statically sized, use the
third form as it will not incur notable performance hits.

Note that D is the dimension of the simplex, i.e. the number
of vertices supplied minus one.
"""
@generated function simplex(vertices::SVector{D1,P}) where {D1,P}
    U = length(P)
    D = D1 - 1
    C = U-D
    T = eltype(P)
    xp1 =:(())
    for i in 1:D
        push!(xp1.args, :(vertices[$i]-vertices[end]))
    end
    xp2 = :(SVector{$D,P}($xp1))
    quote
        tangents = $xp2
        normals, volume = _normals(tangents, Val{$C})
        Simplex(vertices, tangents, normals, volume)
    end
end

simplex(vertices...) = simplex(SVector((vertices...,)))

@generated function simplex(vertices, ::Type{Val{D}}) where D
    P = eltype(vertices)
    D1 = D + 1
    xp = :(())
    for i in 1:D1
        push!(xp.args, :(vertices[$i]))
    end
    :(simplex(SVector{$D1,$P}($xp)))
end

normal(s::Simplex{3,2,1,3,T}) where {T} = normalize(cross(s[1]-s[3],s[2]-s[3]))

function _normals(tangents, ::Type{Val{1}})
    PT = eltype(tangents)
    D  = length(tangents)
    T  = eltype(PT)

    n = zeros(T,D+1)
    b = Array{T}(undef,D,D)

    for i in 1:D+1
        fill!(b, zero(T))
        for j in 1:D
            for k in 1:i-1
                b[k,j] = tangents[j][k]
            end
            for k in i:D
                b[k,j] = tangents[j][k+1]
            end
        end
        n[i] = (-1)^(i-1) * det(b)
    end

    n *= (-1)^D / norm(n)
    normals = SVector{1,PT}([PT(n)])

    metric = T[dot(tangents[i], tangents[j]) for i in 1:D, j in 1:D]
    volume = sqrt(abs(det(metric))) /  factorial(D)

    return normals, volume

end

function _normals(tangents::SVector{2,SVector{3,T}}, ::Type{Val{1}}) where {T}

    n = tangents[1] × tangents[2]
    l = norm(n)

    P = SVector{3,T}
    SVector{1,P}(n/l), 0.5*l
end

function _normals(tangents, ::Type{Val{C}}) where C
    PT = eltype(tangents)
    D  = length(tangents)
    U = length(PT)
    T  = eltype(PT)

    metric = T[dot(tangents[i], tangents[j]) for i in 1:D, j in 1:D]
    volume = sqrt(abs(det(metric))) / factorial(D)

    # Fix this. This function needs to become gneerated
    normals = SVector{C,PT}(PT[zero(PT) for i in 1:C])

    return normals, volume
end

"""
    barytocart(simplex, uv)

Returns the point in the simplex with barycentric coordinates uv
"""
function barytocart(mani::Simplex, u)
    r = last(mani.vertices)
    for i in 1 : dimension(mani)
        ti = mani.tangents[i]
        ui = u[i]
        #r += mani.tangents[i] * u[i]
        r += ti * ui
    end
    return r
end

"""
    carttobary(simplex, point) -> barycoords

Compute the barycentric coordinates on 'simplex' of 'point'.
"""
function carttobary(p::Simplex{U,D,C,N,T}, cart) where {U,D,C,N,T}

    G = [dot(p.tangents[i], p.tangents[j]) for i in 1:D, j in 1:D]
    #w = [dot(p.tangents[i], cart - p.vertices[end]) for i in 1:D]

    o = p.vertices[end]
    w = [sum(t[j]*(cart[j]-o[j]) for j in 1:length(cart)) for t in p.tangents]

    u = G \ w

    return SVector{D}(u)
end

function edges(s::Simplex{3,3})
    C = [SVector{2}(c) for c in combinations(@SVector[1,2,3,4],2)]
    T = eltype(eltype(s.vertices))
    P = Simplex{3,1,2,2,T}
    Edges = Vector{P}(undef, length(C))
    # @show eltype(Edges)
    # @show eltype(s.vertices)
    for c in C
        q = relorientation(c, @SVector[1,2,3,4])
        Q = abs(q)
        Edges[Q] = q > 0 ?
            simplex(s.vertices[c[1]], s.vertices[c[2]]) :
            simplex(s.vertices[c[2]], s.vertices[c[1]])
    end
    return Edges
end

function edges(s::Simplex{3,2})
    return [
        simplex(s.vertices[2], s.vertices[3]),
        simplex(s.vertices[3], s.vertices[1]),
        simplex(s.vertices[1], s.vertices[2])]
end

"""
    ReferenceSimplex{Dimension, CoordType, NumVertices}

This domain is defined to bootstrap the quadrature generation strategy. The generic definition
of numquads on a chart pulls back to the domain. For a limit set of reference domains, explicit
quadrature rules are defined. The weights and points are then pushed forward to the configuaration
space element over which integration is desired.

For more details see the implementation in quadpoints.jl
"""
struct ReferenceSimplex{D,T,N}
    # N: number of defining points
    # D: dimension of the simplex
    # T: type of the coordinates
    simplex::Simplex{D,D,0,N,T}
end

function ReferenceSimplex{D,T,N}() where {D,T,N}
    P = SVector{D,T}[]
    for i in 1:D
        a = zeros(T,D)
        a[i] = 1
        p = SVector{D,T}(a)
        push!(P,p)
    end

    o = zero(SVector{D,T})
    push!(P,o)
    ReferenceSimplex{D,T,N}(simplex(P...))
end

barytocart(ch::ReferenceSimplex, u) = barytocart(ch.simplex, u)
carttobary(ch::ReferenceSimplex, p) = carttobary(ch.simplex, p)

domain(ch::Simplex{U,D,C,N,T}) where {U,D,C,T,N} = ReferenceSimplex{D,T,N}()
neighborhood(ch::ReferenceSimplex, u) = u


#=
Tools for representing and solving dynamic programs with continuous states.

Implement the Bellman equation collocation method as described in Mirand and
Fackler (2002), Chapter 9.

References
----------
* M. J. Miranda and P. L. Fackler, Applied Computational Economics and Finance,
  MIT press, 2002.

=#
using BasisMatrices
import Optim

#= Types and contructors =#

"""
    Interp{N,TS,TM,TL}

Type that contains information about interpolation

# Fields

- `basis::Basis{N}`: Object that contains interpolation basis information
- `S::TS<:VecOrMat`: Vector or Matrix that contains interpolation nodes
- `Scoord::NTuple{N,Vector{Float64}}` Tuple that contains transformed
  interpolation nodes
- `length::Int`: Degree of interpolation at tensor grid
- `size::NTuple{N,Int}`: Tuple that contains degree of interpolation at each
  dimension
- `lb::NTuple{N,Float64}`: Lower bound of domain
- `ub::NTuple{N,Float64}`: Upper bound of domain
- `Phi::TM<:AbstractMatrix`: Interpolation basis matrix
- `Phi_lu::TL<:Factorization`: LU factorized interpolation basis matrix
"""
struct Interp{N,TS<:VecOrMat,TM<:AbstractMatrix,TL<:Factorization}
    basis::Basis{N}
    S::TS
    Scoord::NTuple{N,Vector{Float64}}
    length::Int
    size::NTuple{N,Int}
    lb::NTuple{N,Float64}
    ub::NTuple{N,Float64}
    Phi::TM
    Phi_lu::TL
end

"""
    Interp(basis)

Constructor for `Interp`

# Arguments

-`basis::Basis`: Object that contains interpolation basis information
"""
function Interp(basis::Basis)
    S, Scoord = nodes(basis)
    grid_length = length(basis)
    grid_size = size(basis)
    grid_lb, grid_ub = min(basis), max(basis)
    Phi = BasisMatrix(basis, Expanded(), S).vals[1]
    Phi_lu = lu(Phi)
    interp = Interp(basis, S, Scoord, grid_length, grid_size, grid_lb, grid_ub,
                    Phi, Phi_lu)
end

"""
    ContinuousDP{N,TR,TS,Tf,Tg,Tlb,Tub}

Type that reperesents a continuous-state dynamic program

# Fields

- `f::Tf<:Function`: Reward function
- `g::Tg<:Function`: State transition function
- `discount::Float64`: Discount factor
- `shocks::TR<:AbstractVecOrMat`: Random variables' nodes
- `weights::Vector{Float64}`: Random variables' weights
- `x_lb::Tlb<:Function`: Lower bound of action variables
- `x_ub::Tub<:Function`: Upper bound of action variables
- `interp::Interp{N,TS<:VecOrMat}`: Object that contains information about
  interpolation
"""
mutable struct ContinuousDP{N,TR<:AbstractVecOrMat,TS<:VecOrMat,
                            Tf<:Function,Tg<:Function,
                            Tlb<:Function,Tub<:Function}
    f::Tf
    g::Tg
    discount::Float64
    shocks::TR
    weights::Vector{Float64}
    x_lb::Tlb
    x_ub::Tub
    interp::Interp{N,TS}
end

"""
    ContinuousDP(f, g, discount, shocks, weights, x_lb, x_ub, basis)

Constructor for `ContinuousDP`

# Arguments
- `f::Tf<:Function`: Reward function
- `g::Tg<:Function`: State transition function
- `discount::Float64`: Discount factor
- `shocks::TR<:AbstractVecOrMat`: Random variables' nodes
- `weights::Vector{Float64}`: Random variables' weights
- `x_lb::Tlb<:Function`: Lower bound of action variables
- `x_ub::Tub<:Function`: Upper bound of action variables
- `basis::Basis`: Object that contains interpolation basis information
"""
function ContinuousDP(f::Function, g::Function, discount::Float64,
                      shocks::Array{Float64}, weights::Vector{Float64},
                      x_lb::Function, x_ub::Function,
                      basis::Basis)
    interp = Interp(basis)
    cdp = ContinuousDP(f, g, discount, shocks, weights, x_lb, x_ub, interp)
    return cdp
end

"""
    CDPSolveResult{Algo,N,TR,TS}

Type that contains the solution result of continuous-state dynamic programming

# Fields

- `cdp::ContinuousDP{N,TR,TS}`: Object that contains model paramers
- `tol::Float64`: Convergence criteria
- `max_iter::Int`: Maximum number of iteration
- `C::Vector{Float64}`: Basis coefficients vector
- `converged::Bool`: Bool that shows whether model converges
- `num_iter::Int`: Number of iteration until model converges
- `eval_nodes::TS<:VecOrMat`: Evaluation vector or matrix
- `eval_nodes_coord::NTuple{N,Vector{Float64}}`: Tuple that contains evaluation
  transformed grids
- `V::Vector{Float64}`: Computed value function
- `X::Vector{Float64}`: Computed policy function
- `resid::Vector{Float64}`: Residuals of basis coefficients
"""
mutable struct CDPSolveResult{Algo<:DPAlgorithm,N,
                              TR<:AbstractVecOrMat,TS<:VecOrMat}
    cdp::ContinuousDP{N,TR,TS}
    tol::Float64
    max_iter::Int
    C::Vector{Float64}
    converged::Bool
    num_iter::Int
    eval_nodes::TS
    eval_nodes_coord::NTuple{N,Vector{Float64}}
    V::Vector{Float64}
    X::Vector{Float64}
    resid::Vector{Float64}

    function CDPSolveResult{Algo,N,TR,TS}(
            cdp::ContinuousDP{N,TR,TS}, tol::Float64, max_iter::Integer
        ) where {Algo,N,TR,TS}
        C = zeros(cdp.interp.length)
        converged = false
        num_iter = 0
        eval_nodes = cdp.interp.S
        eval_nodes_coord = cdp.interp.Scoord
        V = Float64[]
        X = Float64[]
        resid = Float64[]
        res = new{Algo,N,TR,TS}(cdp, tol, max_iter, C, converged, num_iter,
                                eval_nodes, eval_nodes_coord, V, X, resid)
        return res
    end
end

Base.ndims(::ContinuousDP{N}) where {N} = N
Base.ndims(::CDPSolveResult{Algo,N}) where {Algo,N} = N

"""
    evaluate!(res)

Evaluate the value function and the policy function at each point

# arguments

- `res::CDPSolveResult`: Object to store the result of dynamic programming
"""
function evaluate!(res::CDPSolveResult)
    cdp, C, s_nodes = res.cdp, res.C, res.eval_nodes
    res.V, res.X = s_wise_max(cdp, s_nodes, C)
    res.resid = res.V - vec(funeval(C, cdp.interp.basis, s_nodes))
    return res
end

function set_eval_nodes!(
        res::CDPSolveResult{Algo,1}, s_nodes_coord::NTuple{1,Vector{Float64}}
    ) where {Algo}
    s_nodes = s_nodes_coord[1]
    res.eval_nodes = s_nodes
    res.eval_nodes_coord = s_nodes_coord
    evaluate!(res)
end

function set_eval_nodes!(
        res::CDPSolveResult{Algo,N}, s_nodes_coord::NTuple{N,Vector{Float64}}
    ) where {Algo,N}
    s_nodes = gridmake(s_nodes_coord...)
    res.eval_nodes = s_nodes
    res.eval_nodes_coord = s_nodes_coord
    evaluate!(res)
end

function set_eval_nodes!(
        res::CDPSolveResult{Algo,N}, s_nodes_coord::NTuple{N,AbstractVector}
    ) where {Algo,N}
    T = Float64
    s_nodes_coord_vecs =
        ntuple(i -> collect(T, s_nodes_coord[i]), N)::NTuple{N,Vector{T}}
    set_eval_nodes!(res, s_nodes_coord_vecs)
end

function set_eval_nodes!(
        res::CDPSolveResult{Algo,N}, s_nodes_coord::Vararg{AbstractVector,N}
    ) where {Algo,N}
    set_eval_nodes!(res, s_nodes_coord)
end

@doc """
    set_eval_nodes!(res, s_nodes_coord)

Set evaluation nodes

# Arguments

- `res::CDPSolveResult`: Object that contains the result of dynamic programming
- `s_nodes_coord::NTuple{N,AbstractVector}`: Evaluation nodes
""" set_eval_nodes!

function (res::CDPSolveResult)(s_nodes::AbstractArray{Float64})
    cdp, C = res.cdp, res.C
    V, X = s_wise_max(cdp, s_nodes, C)
    resid = V - vec(funeval(C, cdp.interp.basis, s_nodes))
    return V, X, resid
end

#= Methods =#

"""
    _s_wise_max(cdp, s, C)

Find optimal value and policy for each grid point

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `s::AbstractArray{Float64}`: Interpolation nodes
- `C::Vector{Float64}`: Basis coefficients vector

# Returns

- `v::Vector{Float64}`: Updated value function vector
- `x::Vector{Float64}`: Updated policy function vector
"""
function _s_wise_max(cdp::ContinuousDP, s, C)
    sp = Array{Float64}(undef, size(cdp.shocks, 1), length(s))
    function objective(x)
        for i in 1:size(sp, 1)
            sp[i, :] .= cdp.g(s, x, cdp.shocks[i, axes(cdp.shocks)[2:end]...])
        end
        Vp = funeval(C, cdp.interp.basis, sp)
        cont = cdp.discount * dot(cdp.weights, Vp)
        flow = cdp.f(s, x)
        -1 * (flow + cont)
    end
    res = Optim.optimize(objective, cdp.x_lb(s), cdp.x_ub(s))
    v = -res.minimum::Float64
    x = res.minimizer::Float64
    return v, x
end

"""
    s_wise_max!(cdp, ss, C, Tv)

Find optimal value for each grid point

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `ss::AbstractArray{Float64}`: interpolation nodes
- `C::Vector{Float64}`: Basis coefficients vector
- `Tv::Vector{Float64}`: A buffer array to hold the updated value function

# Returns

- `Tv::Vector{Float64}`: Updated value function vector
"""
function s_wise_max!(cdp::ContinuousDP, ss::AbstractArray{Float64},
                     C::Vector{Float64}, Tv::Vector{Float64})
    n = size(ss, 1)
    t = Base.tail(axes(ss))
    for i in 1:n
        Tv[i], _ = _s_wise_max(cdp, ss[(i, t...)...], C)
    end
    return Tv
end

"""
    s_wise_max!(cdp, ss, C, Tv)

Find optimal value and policy for each grid point

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `ss::AbstractArray{Float64}`: interpolation nodes
- `C::Vector{Float64}`: Basis coefficients vector
- `Tv::Vector{Float64}`: A buffer array to hold the updated value function
- `X::Vector{Float64}`: A buffer array to hold the updeted policy function

# Returns

- `Tv::Vector{Float64}`: Updated value function vector
- `X::Vector{Float64}`: Updated policy function vector
"""
function s_wise_max!(cdp::ContinuousDP, ss::AbstractArray{Float64},
                     C::Vector{Float64}, Tv::Vector{Float64},
                     X::Vector{Float64})
    n = size(ss, 1)
    t = Base.tail(axes(ss))
    for i in 1:n
        Tv[i], X[i] = _s_wise_max(cdp, ss[(i, t...)...], C)
    end
    return Tv, X
end

"""
    s_wise_max(cdp, ss, C)

Find optimal value and policy for each grid point

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `ss::AbstractArray{Float64}`: Interpolation nodes
- `C::Vector{Float64}`: Basis coefficients vector

# Returns

- `Tv::Vector{Float64}`: Value function vector
- `X::Vector{Float64}`: Policy function vector
"""
function s_wise_max(cdp::ContinuousDP, ss::AbstractArray{Float64},
                    C::Vector{Float64})
    n = size(ss, 1)
    Tv, X = Array{Float64}(undef, n), Array{Float64}(undef, n)
    s_wise_max!(cdp, ss, C, Tv, X)
end

"""
    bellman_operator!(cdp, C, Tv)

Update basis coefficients. Values are stored in `Tv`

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `C::Vector{Float64}`: Basis coefficients vector
- `Tv::Vector{Float64}`: Vector to store values

# Returns

- `C::Vector{Float64}`: Updated basis coefficients vector
"""
function bellman_operator!(cdp::ContinuousDP, C::Vector{Float64},
                           Tv::Vector{Float64})
    Tv = s_wise_max!(cdp, cdp.interp.S, C, Tv)
    ldiv!(C, cdp.interp.Phi_lu, Tv)
    return C
end

"""
    compute_greedy!(cdp, C, X)
    compute_greedy!(cdp, ss, C, X)

Updates policy function vector

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `ss::AbstractArray{Float64}`: Interpolation nodes
- `C::Vector{Float64}`: Basis coefficients vector
- `X::Vector{Float64}`: A buffer array to hold the updated policy function.

# Returns

- `X::Vector{Float64}`: Updated policy function vector
"""
function compute_greedy!(cdp::ContinuousDP, ss::AbstractArray{Float64},
                         C::Vector{Float64}, X::Vector{Float64})
    n = size(ss, 1)
    t = Base.tail(axes(ss))
    for i in 1:n
        _, X[i] = _s_wise_max(cdp, ss[(i, t...)...], C)
    end
    return X
end

compute_greedy!(cdp::ContinuousDP, C::Vector{Float64}, X::Vector{Float64}) =
    compute_greedy!(cdp, cdp.interp.S, C, X)

"""
    evaluate_policy!(cdp, X, C)

Update basis coefficients

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `X::Vector{Float64}`: Policy function vector
- `C::Vector{Float64}`: A buffer array to hold the basis coefficients

# Returns

- `C::Vector{Float64}`: Updated basis coefficients vector
"""
function evaluate_policy!(cdp::ContinuousDP{N}, X::Vector{Float64},
                          C::Vector{Float64}) where N
    n = size(cdp.interp.S, 1)
    ts = Base.tail(axes(cdp.interp.S))
    te = Base.tail(axes(cdp.shocks))
    A = Array{Float64}(undef, n, n)
    A[:] = cdp.interp.Phi
    for i in 1:n
        s = cdp.interp.S[(i, ts...)...]
        for (j, w) in enumerate(cdp.weights)
            e = cdp.shocks[(j, te...)...]
            s_next = cdp.g(s, X[i], e)
            A[i, :] -= ckron(
                [vec(evalbase(cdp.interp.basis.params[k], s_next[k]))
                 for k in N:-1:1]...
            ) * cdp.discount * w
        end
    end
    A_lu = lu(A)
    for i in 1:n
        s = cdp.interp.S[(i, ts...)...]
        C[i] = cdp.f(s, X[i])
    end
    ldiv!(A_lu, C)
    return C
end

"""
    policy_iteration_operator!(cdp, C, X)

Update basis coefficients by policy function iteration

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `C::Vector{Float64}`: Basis coefficients vector
- `X::Vector{Float64}`: A buffer array to hold the updated policy function

# Returns

- `C::Vector{Float64}` Updated basis coefficients vector
"""
function policy_iteration_operator!(cdp::ContinuousDP, C::Vector{Float64},
                                    X::Vector{Float64})
    compute_greedy!(cdp, C, X)
    evaluate_policy!(cdp, X, C)
    return C
end

"""
    operator_iteration!(T, C, tol, max_iter; verbose=2, print_skip=50)

Updates basis coefficients until it converges.

# Arguments

- `T::Function`: Function that updates basis coefficients by VFI or PFI
- `C::Vector{Float64}`: initial basis coefficients vector
- `tol::Float64`: Tolerance to be used to update basis coefficients
- `max_iter::Int`: The maximum number of iteration
- `verbose::Int`: Level of feedback (0 for no output, 1 for warnings only, 2 for
   warning and convergence messages during iteration)
- `print_skip::Int`: if verbose == 2, how many iterations to apply between print
  messages

# Returns

- `converged::Bool`: Bool that shows whether basis coefficients vector converges
- `i::Int`: Number of iteration it took to converge
"""
function operator_iteration!(T::Function, C::TC, tol::Float64, max_iter;
                             verbose::Int=2, print_skip::Int=50) where TC
    converged = false
    i = 0

    if max_iter <= 0
        if verbose >= 1
            @warn("No computation performed with max_iter=$max_iter")
        end
        return converged, i
    end

    err = tol + 1
    C_old = similar(C)
    while true
        copyto!(C_old, C)
        C = T(C)::TC
        err = maximum(abs, C - C_old)
        i += 1
        (err <= tol) && (converged = true)

        (converged || i >= max_iter) && break

        if (verbose == 2) && (i % print_skip == 0)
            println("Compute iterate $i with error $err")
        end
    end

    if verbose == 2
        println("Compute iterate $i with error $err")
    end

    if verbose >= 1
        if !converged
            @warn("max_iter attained")
        elseif verbose == 2
            println("Converged in $i steps")
        end
    end

    return converged, i
end

#= Solve methods =#

"""
    solve(cdp, method=PFI; tol=sqrt(eps()), max_iter=500, verbose=2,
          print_skip=50)

Solve the continuous-state dynamic program

# Arguments

- `cdp::ContinuousDP`: Object that contains model parameters
- `method::Type{T<Algo}(PFI)`: Type name specifying solution method
   Acceptable arguments are 'VFI' for value function iteration or
   'PFI' for policy function iteration. Default solution method is 'PFI'.
- `tol::Real`: Value for epsilon-optimality
- `max_iter::Int`: Maximum number of iterations
- `verbose::Int`: Level of feedback (0 for no output, 1 for warnings only, 2 for
   warning and convergence messages during iteration)
- `print_skip::Int`: if verbose == 2, how many iterations to apply between print
  messages

# Returns

- `res::CDPSolveResult{Algo,N,TR,TS}`: Object to store the result of dynamic
  programming
"""
function solve(cdp::ContinuousDP{N,TR,TS}, method::Type{Algo}=PFI;
               tol::Real=sqrt(eps()), max_iter::Integer=500,
               verbose::Int=2,
               print_skip::Int=50) where {Algo<:DPAlgorithm,N,TR,TS}
    tol = Float64(tol)
    res = CDPSolveResult{Algo,N,TR,TS}(cdp, tol, max_iter)
    _solve!(cdp, res, verbose, print_skip)
    evaluate!(res)
    return res
end

# Policy iteration
"""
    _solve!(cdp, res, verbose, print_skip)

Implement Policy Iteration. See `solve` for further details.
"""
function _solve!(cdp::ContinuousDP, res::CDPSolveResult{PFI},
                 verbose, print_skip)
    C = res.C
    X = Array{Float64}(undef, cdp.interp.length)
    operator!(C) = policy_iteration_operator!(cdp, C, X)
    res.converged, res.num_iter =
        operator_iteration!(operator!, res.C, res.tol, res.max_iter,
                            verbose=verbose, print_skip=print_skip)
    return res
end

# Value iteration
"""
    _solve!(cdp, res, verbose, print_skip)

Implement Value Iteration. See `solve` for further details
"""
function _solve!(cdp::ContinuousDP, res::CDPSolveResult{VFI},
                 verbose, print_skip)
    C = res.C
    Tv = Array{Float64}(undef, cdp.interp.length)
    operator!(C) = bellman_operator!(cdp, C, Tv)
    res.converged, res.num_iter =
        operator_iteration!(operator!, res.C, res.tol, res.max_iter,
                            verbose=verbose, print_skip=print_skip)
    return res
end

#= Simulate methods =#

"""
    simulate!([rng=GLOBAL_RNG], s_path, res, s_init)

Generate a sample path of state variable(s)

# Arguments

- `rng::AbstractRNG`: Random number generator
- `s_path::VecOrMat`: Array to store the generated sample path
- `res::CDPSolveResult`: Object that contains result of dynamic programming
- `s_init`: Initial value of state variable(s)

# Return

- `s_path::VecOrMat`:: Generated sample path of state variable(s)
"""
function simulate!(rng::AbstractRNG, s_path::TS,
                   res::CDPSolveResult{Algo,N,TR,TS},
                   s_init) where {Algo,N,TR,TS<:VecOrMat}
    ts_length = size(s_path)[end]
    cdf = cumsum(res.cdp.weights)
    r = rand(rng, ts_length - 1)
    e_ind = Array{Int}(undef, ts_length - 1)
    for t in 1:ts_length - 1
        e_ind[t] = searchsortedlast(cdf, r[t]) + 1
    end

    basis = Basis(map(LinParams, res.eval_nodes_coord, ntuple(i -> 0, N)))
    X_interp = Interpoland(basis, res.X)

    s_ind_front = Base.front(axes(s_path))
    e_ind_tail = Base.tail(axes(res.cdp.shocks))
    s_path[(s_ind_front..., 1)...] = s_init
    for t in 1:ts_length - 1
        s = s_path[(s_ind_front..., t)...]
        x = X_interp(s)
        e = res.cdp.shocks[(e_ind[t], e_ind_tail...)...]
        s_path[(s_ind_front..., t + 1)...] = res.cdp.g(s, x, e)
    end

    return s_path
end

simulate!(s_path::VecOrMat{Float64}, res::CDPSolveResult, s_init) =
    simulate!(Random.GLOBAL_RNG, s_path, res, s_init)

"""
    simulate([rng=GLOBAL_RNG], res, s_init, ts_length)

Generate a sample path of state variable(s)

# Arguments

- `rng::AbstractRNG`: Random number generator
- `res::CDPSolveResult`: Object that contains result of dynamic programming
- `s_init`: Initial value of state variable(s)
- `ts_length::Integer`: Length of simulation

# Return

- `s_path::VecOrMat`:: Generated sample path of state variable(s)
"""
function simulate(rng::AbstractRNG, res::CDPSolveResult{Algo,1}, s_init::Real,
                  ts_length::Integer) where {Algo<:DPAlgorithm}
    s_path = Array{Float64}(undef, ts_length)
    simulate!(rng, s_path, res, s_init)
    return s_path
end

simulate(res::CDPSolveResult{Algo,1}, s_init::Real,
         ts_length::Integer) where {Algo<:DPAlgorithm} =
    simulate(Random.GLOBAL_RNG, res, s_init, ts_length)

function simulate(rng::AbstractRNG, res::CDPSolveResult, s_init::Vector,
                  ts_length::Integer)
    s_path = Array{Float64}(undef, length(s_init), ts_length)
    simulate!(rng, s_path, res, s_init)
    return s_path
end

simulate(res::CDPSolveResult, s_init::Vector, ts_length::Integer) =
    simulate(Random.GLOBAL_RNG, res, s_init, ts_length)


#---------------------------------------------------------------------#
#This routine advances the solution in time using Explicit RK methods
#written in Butcher tableau form.
#Written by F.X. Giraldo / P.R. Mugg on 9/22/19
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey, CA 93943-5216
#---------------------------------------------------------------------#
function erk_butcher(q0,c,αI,βI,I,Δt,time)

   Q = zeros(I,1)
   R = zeros(I,1)

   c𝑖 = zeros(I)
   for i = 1:I
      c𝑖[i] = sum(αI[i,:])
   end

   Q[1] = q0
   t = time
   (rhs,f,g) = rhs_function(Q[1],c,t)
   R[1] = rhs

   for i = 2:I
      R_sum = 0
      for j = 1:i-1
         R_sum += αI[i,j]*R[j]
      end
      Q[i] = q0 + Δt*R_sum
      t = time + Δt*c𝑖[i]
      (rhs,f,g) = rhs_function(Q[i],c,t)
      R[i] = rhs
   end

   R_sum = 0
   for i = 1:I
      R_sum += βI[i]*R[i]
   end

   return qp = q0 + Δt*R_sum
end #end function


function makeprimalstrictlyfeasible!{T<:Number}(sdp::SparseSDP{T}, epsilon::Number)
    for (ci, v) in rhs(sdp)
        v = zero(T)
        blockmatrix = cons(sdp)[ci]
        for (_, block) in blockmatrix
            for (indices, entry) in block
                i, j = indices
                if i == j
                    v += entry
                end
            end
        end
        setrhs!(sdp, ci, rhs(sdp, ci) + epsilon*v)
    end
end

function makedualstrictlyfeasible!{T<:Number}(sdp::SparseSDP{T}, epsilon::Number)
    for (_, block) in obj(sdp)
        for (indices, value) in block
            i, j = indices
            if i == j
                block[i, j] = block[i,j] - epsilon
            end
        end
    end
end


using ShuffleProofs
import ShuffleProofs: PrimeGenerator, Enc, Dec, shuffle, ElGamal

sk = 123 # Only authorithy who do decryption would know

g = PrimeGenerator(3, 23)
options = [g, g^2, g^3]
pk = g^sk

bbord = (;) # A representation of a buletinboard

######## Step 1. ########

bbord = let
    enc = Enc(pk, g)
    ciphertexts_in = ElGamal([enc(options[rand(1:3)], rand(1:10)) for i in 1:10])
    (; bbord..., ciphertexts_in)
end

######## Step 2. ########

bbord = let
    enc = Enc(pk, g)
    proposition, secret = shuffle(bbord.ciphertexts_in, enc) # reencrypts and sorts output
    ciphertexts_out = proposition.𝐞′
    (; bbord..., ciphertexts_out)
end

####### Step 3. ########

bbord = let
    dec = Dec(sk)
    votes = dec(bbord.ciphertexts_out)
    (; bbord..., votes)
end

######### As last step we can tally the votes ############

for (i, o) in enumerate(options)
    s = 0
    for v in bbord.votes
        if v == o
            s += 1
        end
    end

    println("Choice $i received $s votes")
end


############################
# Lorenz System  (MIMO)    #
# by VSSM (all directions) #
############################

###########################################
# Controll parameters for VSSM controller #
###########################################

K=500
w=1

#############################################
# Parameters to design a Nominal Trajectory #
#############################################

A₁=2
ω₁=0.5
A₂=3.8
ω₂=0.7
A₃=1.5
ω₃=0.9

#################
# Time variable #
#################

# cycle time in seconds
δt=1e-3
# The lenght of the simulation in seconds
LONG=Int(2e4)

Λ=12

##########################
# Exact model parameters #
##########################
σₑ=5
βₑ=8/3
ρₑ=40

############################
# Approx. model parameters #
############################
σₐ=5
βₐ=8/3
ρₐ=40

#Arrays

xN=zeros(LONG)
xN_p=zeros(LONG)
yN=zeros(LONG)
yN_p=zeros(LONG)
zN=zeros(LONG)
zN_p=zeros(LONG)

x_p_des=zeros(LONG)
y_p_des=zeros(LONG)
z_p_des=zeros(LONG)

u_x=zeros(LONG)
u_y=zeros(LONG)
u_z=zeros(LONG)

t=zeros(LONG)

S_x=zeros(LONG)
S_p_x=zeros(LONG)
S_y=zeros(LONG)
S_p_y=zeros(LONG)
S_z=zeros(LONG)
S_p_z=zeros(LONG)

x=zeros(LONG)
y=zeros(LONG)
z=zeros(LONG)
x_p=zeros(LONG)
y_p=zeros(LONG)
z_p=zeros(LONG)

#initial conditions
hint_x=0
hint_y=0
hint_z=0

x[1]=A₁*sin(ω₁*δt)
y[1]=A₂*sin(ω₂*δt)
z[1]=A₃*sin(ω₃*δt)

##############
# SIMULATION #
##############

for i=1:LONG-1 #i counts the cycles of the simulation
  t[i]=δt*i
  #create Nominal Trajectory
  xN[i]=A₁*sin(ω₁*t[i])
  xN_p[i]=A₁*ω₁*cos(ω₁*t[i])
  #xN_p[i]=-A₁*ω₁^2*sin(ω₁*t[i])

  yN[i]=A₂*sin(ω₂*t[i])
  yN_p[i]=A₂*ω₂*cos(ω₂*t[i])
  #yN_p[i]=-A₂*ω₂^2*sin(ω₂*t[i])

  zN[i]=A₃*sin(ω₃*t[i])
  zN_p[i]=A₃*ω₃*cos(ω₃*t[i])
  #zN_p[i]=-A₃*ω₃^2*sin(ω₃*t[i])
  # Compute  error (h) and error integral (hint) and the first derivative of the error (h_p)
  h_x=xN[i]-x[i]
  h_y=yN[i]-y[i]
  h_z=zN[i]-z[i]

  hint_x=hint_x+δt*h_x
  hint_y=hint_y+δt*h_y
  hint_z=hint_z+δt*h_z

  h_p_x=xN_p[i]-x_p[i]
  h_p_y=yN_p[i]-y_p[i]
  h_p_z=zN_p[i]-z_p[i]

  #Error metric (S), and the first derivative of the Error metric (S_p)
  S_x[i]=Λ*hint_x+h_x
  S_p_x[i]=Λ*h_x+h_p_x
  S_y[i]=Λ*hint_y+h_y
  S_p_y[i]=Λ*h_y+h_p_y
  S_z[i]=Λ*hint_z+h_z
  S_p_z[i]=Λ*h_z+h_p_z
  #Kinematic Block (PID type feedback)
  x_p_des[i]=xN_p[i]+Λ*h_x+K*tanh(S_x[i]/w)
  y_p_des[i]=yN_p[i]+Λ*h_y+K*tanh(S_y[i]/w)
  z_p_des[i]=zN_p[i]+Λ*h_z+K*tanh(S_z[i]/w)
  # Create the control signal
  u_x[i]=x_p_des[i]-σₐ*(y[i]-x[i])
  u_y[i]=y_p_des[i]-x[i]*(ρₐ-z[i])
  u_z[i]=z_p_des[i]-y[i]+βₐ*z[i]
  # Get the system "reaction" (the exact model of the system)
  x_p[i]=σₑ*(y[i]-x[i])+u_x[i]
  y_p[i]=x[i]*(ρₑ-z[i])-y[i]+u_y[i]
  z_p[i]=x[i]*y[i]-βₑ*z[i]+u_z[i]
  # Integrate back to x, y, z
  x[i+1]=x[i]+δt*x_p[i]
  y[i+1]=y[i]+δt*y_p[i]
  z[i+1]=z[i]+δt*z_p[i]
end #for

using PyPlot

figure(1)
grid("on")
title("Nominal and Realized trajectorys")
xlabel(L"Time $[s]$")
ylabel(L"q $[m]$")
plot(t[3:LONG-1],xN[3:LONG-1],color="#7684FF")
plot(t[3:LONG-1],yN[3:LONG-1],color="#2A40FF")
plot(t[3:LONG-1],zN[3:LONG-1],color="#3B427F")

plot(t[3:LONG-1],x[3:LONG-1],color="#FFAA41","r--")
plot(t[3:LONG-1],y[3:LONG-1],color="#FF9F28","r--")
plot(t[3:LONG-1],z[3:LONG-1],color="#B2690E","r--")

figure(2)
grid("on")
title("Tracking Error")
xlabel(L"Time $[s]$")
ylabel(L"error $[m]$")
plot(t[3:LONG-1],xN[3:LONG-1]-x[3:LONG-1],color="red")
plot(t[3:LONG-1],yN[3:LONG-1]-y[3:LONG-1],color="green")
plot(t[3:LONG-1],zN[3:LONG-1]-z[3:LONG-1],color="blue")

figure(3)
grid("on")
title("Phase Space of Error Metric (S_p vs S)")
plot(S_x[3:LONG-1],S_p_x[3:LONG-1],color="red")
plot(S_y[3:LONG-1],S_p_y[3:LONG-1],color="green","r--")
plot(S_z[3:LONG-1],S_p_z[3:LONG-1],color="blue","r-.")

figure(4)
grid("on")
title("Phase Space of motion (x_p vs x)")
plot(xN[1:LONG-1],xN_p[1:LONG-1],color="red")
plot(x[1:LONG-1],x_p[1:LONG-1],color="blue","r--")
plot(yN[1:LONG-1],yN_p[1:LONG-1],color="#7684FF")
plot(y[1:LONG-1],y_p[1:LONG-1],color="#FFAA41","r--")
plot(zN[1:LONG-1],zN_p[1:LONG-1],color="#2A40FF")
plot(z[1:LONG-1],z_p[1:LONG-1],color="#FF9F28","r-.")


using SnpArrays, Test, Statistics, Random, TraitSimulation
using DataFrames, LinearAlgebra, OrdinalMultinomialModels, GLM, DelimitedFiles
Random.seed!(1234)

# demo genetic data can be found in Snparrays package directory
filepath = SnpArrays.datadir()
filename = "EUR_subset"
EUR = SnpArray(SnpArrays.datadir(filename * ".bed"));

EUR_data = SnpData(SnpArrays.datadir(filename));
bimfile = EUR_data.snp_info # store the snp_info with the snp names
snpid  = bimfile[:, :snpid] # store the snp names in the snpid vector
causal_snp_index = findfirst(x -> x ==  "rs62057050" , snpid) # find the index of the snp of interest by snpid
locus = convert(Vector{Float64}, @view EUR[:, causal_snp_index])
famfile = EUR_data.person_info
sex = map(x -> strip(x) == "F" ? 0.0 : 1.0, famfile[!, :sex])
X = [sex locus]
β_full =
 [0.52
0.05]

# first ordinal power for the single snp
θ = [1.0, 1.2, 1.4]
trait_ordered = OrderedMultinomialTrait(X, β_full, θ, LogitLink())

n_sim = 10

writepheno(filename, filepath, trait_ordered, n_sim)

pheno_ordered = readdlm("pheno.txt", Any)

# test to see if the outputted phenotype file has the correct number of columns, the first two columns are family ID and Individual ID respectively. The remaining columns are the # of simulated univariate traits
@test size(pheno_ordered, 2) == n_sim + 2
# first column FID
@test famfile[:, 1] == string.(pheno_ordered[2:end, 1])
# second column IID
@test famfile[:, 2] == string.(pheno_ordered[2:end, 2])

## VCM trait bivariate test
β = [β_full β_full]
GRM = grm(EUR, minmaf = 0.05);
I_n = Matrix(I, size(GRM))
ΣE = [1.0 0.0; 0.0 1.0]
ΣA = [0.3 0.0; 0.0 0.3]
vc = @vc ΣA ⊗ 2GRM + ΣE ⊗ I_n;
trait_vcm = VCMTrait(X, β, vc)

n_sim = 5

writepheno(filename, filepath, trait_vcm, n_sim)

pheno_ordered = readdlm("pheno.txt", Any)

# test to see if the outputted phenotype file has the correct number of columns, the first two columns are family ID and Individual ID respectively. The remaining columns are the # of simulated univariate traits
n_traits = ntraits(trait_vcm)

@test size(pheno_ordered, 2) == n_traits * n_sim + 2
# first column FID
@test famfile[:, 1] == string.(pheno_ordered[2:end, 1])
# second column IID
@test famfile[:, 2] == string.(pheno_ordered[2:end, 2])


using Test
using KnetMetrics
using Random

@testset "regression-metrics" begin

    Random.seed!(42)
    u = rand(100)
    Random.seed!(123)
    v = rand(100)

    @testset "max-error" begin
        @test isapprox(max_error(u,v), 0.8624100780860133)
    end

    @testset "max-error" begin
        @test isapprox(max_error(u,v), 0.8624100780860133)
    end

    @testset "mean-absolute-error" begin
        @test isapprox(mean_absolute_error(u,v), 0.3666602951971807)
    end

    @testset "mean-squared-error" begin
        @test isapprox(mean_squared_error(u,v), 0.3666602951971807)
    end

    @testset "mean-squared-log-error" begin
        @test isapprox(mean_squared_log_error(u,v), 0.24971305530593674)
    end

    @testset "median-absolute-error" begin
        @test isapprox(median_absolute_error(u,v), 0.3183007014762824)
    end

    @testset "mean-absolute-percentage-error" begin
        @test isapprox(mean_absolute_percentage_error(u,v), 5.329695963484449)
    end

end


module StagedFilters

"""
Savitzky-Golay filter of window half-width M and degree
N. M is the number of points before and after to interpolate, i.e. the full
width of the window is 2M+1.
"""
abstract type AbstractStagedFilters end

struct SavitzkyGolayFilter{M,N} <: AbstractStagedFilters end
wrapL(i, n) = ifelse(1 ≤ i, i, i + n)
wrapR(i, n) = ifelse(i ≤ n, i, i - n)

"""
smooth!(filter,data, smoothed) -
apply `filter` to `data` writing result to `smoothed`.
Note that feeding `Int`s and not floats as data will result in a performance slowdown.
"""
@generated function smooth!(::Type{SavitzkyGolayFilter{M,N}}, data :: AbstractArray{T}, smoothed :: AbstractArray{S}) where {M,N,T,S}

  J = T[(i - M - 1 )^(j - 1) for i = 1:2M + 1, j = 1:N + 1]
  e₁ = [one(T); zeros(T,N)]
  C = J' \ e₁
  pre = :(for i = 1:$M end)
  main = :(for i = $(M + 1):n - $M end)
  post = :(for i = n - $(M - 1):n end)

  for loop in (pre, main, post)
      body = loop.args[2].args

      idx = loop !== pre ? :(i - $M) : :(wrapL(i - $M, n))   # Manually start the first iteration. See the "false" branch below.
      push!(body, :( x = muladd($(C[1]), data[$idx], $(zero(T))))) # Swap `muladd` instead of the additions. Note the index of 1.

      for j = reverse(1:M-1) # Because we bumped out the first iteration, we have to reduce the for loop index by one.
          idx = loop !== pre ? :(i - $j) : :(wrapL(i - $j, n))
          push!(body, :( x = muladd($(C[M + 1 - j]),data[$idx],x))) # muladd
      end

      push!(body, :( x = muladd($(C[M + 1]), data[i], x))) # muladd

      for j = 1:M
          idx = loop !== post ? :(i + $j) : :(wrapR(i + $j, n))
          push!(body, :( x = muladd($(C[M + 1 + j]), data[$idx], x))) # muladd
      end
      push!(body, :(smoothed[i] = x))
  end

 last_expr = quote
          n = length(data)
          n == length(smoothed) || throw(DimensionMismatch())
          @inbounds $pre; @inbounds  $main; @inbounds $post
          return smoothed
  end

  return last_expr = Base.remove_linenums!(last_expr)
end;

export SavitzkyGolayFilter, smooth!

end # module


# ---
# title: 799. Champagne Tower
# id: problem799
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/champagne-tower/description/>
# hidden: true
# ---
#
# We stack glasses in a pyramid, where the **first** row has `1` glass, the
# **second** row has `2` glasses, and so on until the 100th row.  Each glass
# holds one cup of champagne.
#
# Then, some champagne is poured into the first glass at the top.  When the
# topmost glass is full, any excess liquid poured will fall equally to the glass
# immediately to the left and right of it.  When those glasses become full, any
# excess champagne will fall equally to the left and right of those glasses, and
# so on.  (A glass at the bottom row has its excess champagne fall on the
# floor.)
#
# For example, after one cup of champagne is poured, the top most glass is full.
# After two cups of champagne are poured, the two glasses on the second row are
# half full.  After three cups of champagne are poured, those two cups become
# full - there are 3 full glasses total now.  After four cups of champagne are
# poured, the third row has the middle glass half full, and the two outside
# glasses are a quarter full, as pictured below.
#
# ![](https://s3-lc-upload.s3.amazonaws.com/uploads/2018/03/09/tower.png)
#
# Now after pouring some non-negative integer cups of champagne, return how full
# the `jth` glass in the `ith` row is (both `i` and `j` are 0-indexed.)
#
#
#
# **Example 1:**
#
#
#
#     Input: poured = 1, query_row = 1, query_glass = 1
#     Output: 0.00000
#     Explanation: We poured 1 cup of champange to the top glass of the tower (which is indexed as (0, 0)). There will be no excess liquid so all the glasses under the top glass will remain empty.
#
#
# **Example 2:**
#
#
#
#     Input: poured = 2, query_row = 1, query_glass = 1
#     Output: 0.50000
#     Explanation: We poured 2 cups of champange to the top glass of the tower (which is indexed as (0, 0)). There is one cup of excess liquid. The glass indexed as (1, 0) and the glass indexed as (1, 1) will share the excess liquid equally, and each will get half cup of champange.
#
#
# **Example 3:**
#
#
#
#     Input: poured = 100000009, query_row = 33, query_glass = 17
#     Output: 1.00000
#
#
#
#
# **Constraints:**
#
#   * `0 <= poured <= 109`
#   * `0 <= query_glass <= query_row < 100`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


lcm(m,n)


# Packages used in this chapter
# To check for installed packages in V1.0 we now need to use the PKG API

using Pkg
ipks = Pkg.installed()
pkgs = ["BenchmarkTools","UnicodePlots","PyPlot"]

for p in pkgs
    !haskey(ipks,p) && Pkg.add(p)
end

using BenchmarkTools, UnicodePlots, PyPlot

# Define a function to compute the sum of the squares of (x,y)

sumsq(x,y) = x*x + y*y;

# Set the number of trials and initialise the counter
# Using zero(Integer) will ensure the counter is an integer

N = 10^8;
K = zero(Integer);

using Printf

# Sum the number of pairs which lie within the inscribed circle

for i = 1:N
  if sumsq(rand(), rand()) < 1.0
    global K += 1
  end
end

# ... and output the value for PI

@Printf.printf "Estimate of PI for %d trials is %8.5f\n" N 4.0*(K / N);

# We can time this using the @time macro
# Note that his takes some time, will be different in the REPL

@time for i = 1:N
  if sumsq(rand(), rand()) < 1.0
    global K += 1
  end
end

# Now need the global K to access top level varables in the REPL
# The error message is less than informative
# i.e. UndefVarError: K not defined

# Try v0.7 to get a warning and understand the problem
# ┌ Warning: Deprecated syntax `implicit assignment to global variable `K``.
# │ Use `global K` instead.
# └ @ none:0

# This ONLY applies to the top level loops
# So wrapping in let/end, begin/end block is OK, since K is in local scope

# Has is by using using a function and passing K,N as parameters
#
# Here is a contrived example which demonstrates this

while (true)
    N = 1000000
    K = zero(Integer)
    for i = 1:N
        if sumsq(rand(), rand()) < 1.0
          K += 1
        end
    end
    println(4.0*(K/N))
    break
end

# We can rewrite as a generator expressions
# This does not require the variable 'K'
#
# Comprehensions can also be written without the enclosing square brackets,
# producing an object known as a generator.

((rand()^2 + rand()^2) < 1.0 for i = 1:N)

# The generator can be "evaluated" as the argument to an aggregate function

count((rand()^2 + rand()^2) < 1.0 for i = 1:N) * 4 / N

# And this runs much more quickly with a much better memory allocation

@time count((rand()^2 + rand()^2) < 1.0 for i = 1:N) * 4 / N

# JULIA_HOME is not longer defined in Version 1
# See: https://docs.julialang.org/en/v1/manual/environment-variables/
#
# JULIA_BINDIR sets Sys.BINDIR

Sys.BINDIR

# Rather than just using the @time macro we can get a more accuracy timing using the BenchmarkTools package

using BenchmarkTools

# Define the function to sum the series
# It converges very slowly

function basel(N::Integer)
  @assert N > 0
  s = 0.0
  for i = 1:N
    s += 1.0/float(i)^2
  end
  return s
end

basel(10^8)   # Evaluate it over 10^8 terms

π*π/6.0    # Compare this to the analytic solution

# Use the BenchmarkTools macro @benchmark to get more unbiased timings

@benchmark basel(10^8) samples=10

using UnicodePlots

# Generate an array of the numbers from 1 to 100
# The ordinate value is create using a list comprehensive

x = collect(1:100);
y = [x[i]*sin(0.3*x[i])*exp(-0.03*x[i]) for i = 1:length(x)];

myPlot = lineplot(x, y, title = "My Plot", name = "chp-01")

# Alternatively this can be done using a map() construction

t = collect(0.0:0.1:10.0);
y =  map(x -> x*sin(3.0*x)*exp(-0.3*x), t);
myPlot = lineplot(t, y, title = "Second Plot", name = "chp-01")

using PyPlot
plt = PyPlot   # Will get function name clashes with UnicodePlots.

S0  = 100;      # Spot price
K   = 102;      # Strike price
r   = 0.05;     # Risk free rate
q   = 0.0;      # Dividend yield
v   = 0.2;      # Volatility
tma = 0.25;     # Time to maturity
T   = 90;       # Number of time steps

dt  = tma/T;
N = T + 1;
x = collect(0:T);

# Qualify the function calls since UnicodePlot has exports same functions into Main.

plt.title("Asian Option trajectories");
plt.xlabel("Time");
plt.ylabel("Stock price");

for k = 1:5
  S = zeros(Float64,N)
  S[1] = S0;
  dW = randn(N)*sqrt(dt);
  [ S[t] = S[t-1] * (1 + (r - q - 0.5*v*v)*dt + v*dW[t] + 0.5*v*v*dW[t]*dW[t]) for t=2:N ]
  plt.plot(x,S)
end

# Notice that the reference to the PyPlot object is in scope inside this loop
# The same is true when passing an array into a loop
# Confusing YES?
#
K = [0]; N=10^6;
for i = 1:N
  ((rand()^2 + rand()^2) .<  1.0) && (K[1] += 1)
end
mypi = 4*(K[1]/N)


using Devectorize

function async_gs(m,b,x,r,eps)
###
# Solves the system m * x = b using distributed async Gauss-Seidel method
# m is a distributed matrix
# b,x,r are distributed vectors
# eps is the stopping criterium
###
    res = Inf
    I = localindexes(x)[1]
    localx = convert(Array,x)
    n = size(m,2)

    updatingx = false
    calculatingr = false

    iter = 0

    while res > eps

        # Updating x

        for i in I
            s = 0
            for j = 1:n
                if i != j
                    s += m[i,j] * localx[j]
                end
            end

            localx[i] = (b[i] - s) / m[i,i]
            localpart(x)[i - I.start + 1] = localx[i]
        end

        # Stopping criteria

        res = 0

        for i in I

            s = b[i]

            for j = 1:n

                s -= m[i,j] * localx[j]

            end

            localpart(r)[i - I.start + 1] = abs(s)

            res = max(res,s)

        end

        if res <= eps && !calculatingr

            calculatingr = true
            @async begin
                res = norm(r,Inf)
                calculatingr = false
            end

        end

        if !updatingx

            updatingx = true
            @async begin
                copy!(localx,x)
                updatingx = false
            end

        end

        iter += 1
    end
    return iter
end

function shared_gs(m,b,x,r,eps)
    ###
    # Solves the system m * x = b using a shared memory parallel
    # Gauss-Seidel method
    ###
    res = norm(r,Inf)
    I = localindexes(x)
    n = size(m,2)

    iter = 0
    while res > eps

        # Updating x

        for i in I
            s = 0
            for j = 1:n
                if i != j
                    s += m[i,j] * x[j]
                end
            end

            x[i] = (b[i] - s) / m[i,i]

        end

        # Stopping criterium

        res = 0

        for i in I

            s = b[i]

            for j = 1:n

                s -= m[i,j] * x[j]

            end

            r[i] = abs(s)

            res = max(res,s)

        end

        if res <= eps

            res = norm(r,Inf)

        end

        sleep(rand() / 1000)

        iter += 1
    end
    return iter
end

function serial_gs(m,b,x,r,eps)

    res = Inf
    n = size(m,1)
    iter = 0

    while res > eps
        for i in 1:n
            s = b[i]
            for j = 1:n
                if i != j
                    s -= m[i,j] * x[j]
                end
            end
            x[i] = s / m[i,i]
        end

        res = 0

        for i = 1:n
            s = b[i] - BLAS.dot(n,m[i,:],1,x,1)
#            for j = 1:n
#                s -= m[i,j] * x[j]
#            end
            res = max(res,abs(s))
        end

        iter += 1
    end

    return iter
end

function gs(m,b,x,N)

    for k = 1:N
        I = localindexes(x)[1]
        for i in I
            s = 0
            for j = 1:size(m,2)
                if i != j
                    s += m[i,j] * x[j]
                end
            end
            localpart(x)[i - I.start + 1] = (b[i] - s) / m[i,i]
        end
    end
end

function initializeRandomDistributedDDMatrix(I,l,u)
    m = l + (u - l) * rand(size(I[1],1),size(I[2],1))
    for i in 1:size(I[1],1)
        @devec m[i,I[1][i]] = sum(abs(m[i,:]))
    end
    return m
end

function generateRandomDistributedSystem(n,sol=ones((n,1)),w=workers(),
                                         l=-100,u=100)

    dist = [size(w,1),1]

    M = DArray((I)->initializeRandomDistributedDDMatrix(I,l,u),(n,n),w,dist);

    b = dzeros((size(M,1),1),w,dist)
    r = similar(b)

    for i in w
        @spawnat i localpart(b)[:,1] = localpart(M) * sol
        @spawnat i localpart(r)[:,1] = b[localindexes(b)[1],1]
    end

    x0 = dzeros((size(M,1),1),w,dist)

    return (M,b,x0,r)
end

# Para rodar
# require("gs_distribuido.jl")
#
# (M,b,x,r) = generateRandomDistributedSystem(20)
#
# for i in workers()
#     @spawnat i async_gs(M,b,x,r,0.1)
# end
#
# Ou, a versao de 1 linha (imprime o tempo e o total de iteracoes):
#
# @time reduce(+,pmap(fetch,{(@spawnat i async_gs(M,b,x,r,0.1)) for i in workers()}))


using FirstOrderLogic
using Base.Test

include("test_parser.jl")
include("test_renamed_quantifiers_form.jl")
include("test_quantified_variables_form.jl")
include("test_prenex_normal_form.jl")
include("test_skolem_normal_form.jl")
include("test_removed_quantifier_form.jl")
include("test_conjunctive_normal_form.jl")
include("test_resolution.jl")


#Extract variable names from a formula before creating schema
function varnames(f::FormulaTerm)
  rhs = varnames(f.rhs)
  if rhs isa Symbol rhs = [rhs] end
  return (lhs=varnames(f.lhs), rhs=rhs)
end

function varnames(f::Term)
  return Symbol(f)
end

function varnames(f::StatsModels.TupleTerm)
  raw = map(x -> varnames(x), f)
  return get_unique(raw)
end

function varnames(f::ConstantTerm)
  return nothing
end

function varnames(f::InteractionTerm)
  return map(x -> varnames(x), f.terms)
end

function get_unique(f::Tuple)
  out = Symbol[]
  fl(v) = for x in v
    if x isa Tuple fl(x) else push!(out, x) end
  end
  fl(f)
  return unique(out)
end

function combine_formulas(formulas::Union{Array,FormulaTerm}...)
  formulas = map(x -> x isa FormulaTerm ? [x] : x, formulas)
  form_combs = collect(Base.product(formulas...))
  return reshape(form_combs, length(form_combs))
end

"""

    allsub(f::FormulaTerm)

Takes a formula `f` and returns an array containing formulas for all
subsets of the formula terms.
"""
function allsub(f::FormulaTerm)
  if f.rhs isa Union{ConstantTerm,Term,Tuple{Term}} return [f] end
  combs = Tuple.(collect(combinations(Symbol.(f.rhs))))
  terms = map(x -> term.(x), combs)
  forms = map(x -> term(f.lhs.sym) ~ sum(x), terms)
  return [term(f.lhs.sym) ~ term(1); forms]
end


using DiffOpt
using Test
using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

using JuMP
using LinearAlgebra
import Ipopt
import OSQP
import Clp
import SCS

using DelimitedFiles
import MatrixOptInterface
const MatOI = MatrixOptInterface

@testset "Testing forward on trivial QP" begin
    # using example on https://osqp.org/docs/examples/setup-and-solve.html
    Q = [
        4.0 1.0
        1.0 2.0
    ]
    q = [1.0, 1.0]
    G = [
         1.0 1.0
         1.0 0.0
         0.0 1.0
        -1.0 -1.0
        -1.0 0.0
         0.0 -1.0
    ]
    h = [1, 0.7, 0.7, -1, 0, 0]

    model = JuMP.Model(() -> diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x[1:2])
    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model,
        G * x .<= h,
    )
    optimize!(model)

    @test JuMP.value.(x) ≈ [0.3, 0.7] atol=ATOL rtol=RTOL
end

@testset "Differentiating trivial QP 1" begin
    Q = [
        4.0 1.0
        1.0 2.0
    ]
    q = [1.0, 1.0]
    G = [1.0 1.0]
    h = [-1.0]

    model = JuMP.direct_model(diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    x = @variable(model, [1:2])
    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model,
        G * x .<= h,
    )
    optimize!(model)

    @test JuMP.value.(x) ≈ [-0.25, -0.75] atol=ATOL rtol=RTOL

    grad_wrt_h = backward(JuMP.backend(model), ["h"], ones(2))[1]

    @test grad_wrt_h ≈ [1.0] atol=2ATOL rtol=RTOL
end

@testset "Differentiating QP with inequality and equality constraints" begin
    # refered from: https://www.mathworks.com/help/optim/ug/quadprog.html#d120e113424
    # Find equivalent qpth program here - https://github.com/AKS1996/jump-gsoc-2020/blob/master/DiffOpt_tests_4_py.ipynb

    Q = [
         1.0 -1.0 1.0;
        -1.0  2.0 -2.0;
         1.0 -2.0 4.0
    ]
    q = [2.0, -3.0, 1.0]
    G = [0.0 0.0 1.0
         0.0 1.0 0.0
         1.0 0.0 0.0
         0.0 0.0 -1.0
         0.0 -1.0 0.0
         -1.0 0.0 0.0
    ]
    h = [1.0, 1.0, 1.0, 0.0, 0.0, 0.0,]
    A = [1.0 1.0 1.0;]
    b = [0.5]

    model = Model(() -> diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x[1:3])
    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model,
        G * x .<= h,
    )
    @constraint(model,
        A * x .== b,
    )
    optimize!(model)
    doptimizer = JuMP.backend(model).optimizer.model
    z = doptimizer.primal_optimal

    @test z ≈ [0.0, 0.5, 0.0] atol=ATOL rtol=RTOL

    grads = backward(doptimizer, ["Q","q","G","h","A","b"], ones(3))

    dl_dQ = grads[1]
    dl_dq = grads[2]
    dl_dG = grads[3]
    dl_dh = grads[4]
    dl_dA = grads[5]
    dl_db = grads[6]

    @test dl_dQ ≈ zeros(3,3)  atol=ATOL rtol=RTOL

    @test dl_dq ≈ zeros(3,1) atol=ATOL rtol=RTOL

    @test dl_dG ≈ zeros(6,3) atol=ATOL rtol=RTOL

    @test dl_dh ≈ zeros(6,1) atol=ATOL rtol=RTOL

    @test dl_dA ≈ [0.0 -0.5 0.0] atol=ATOL rtol=RTOL

    @test dl_db ≈ [1.0] atol=ATOL rtol=RTOL
end

# refered from https://github.com/jump-dev/MathOptInterface.jl/blob/master/src/Test/contquadratic.jl#L3
# Find equivalent CVXPYLayers and QPTH code here:
#               https://github.com/AKS1996/jump-gsoc-2020/blob/master/DiffOpt_tests_1_py.ipynb
@testset "Differentiating MOI examples 1" begin
    # homogeneous quadratic objective
    # Min x^2 + xy + y^2 + yz + z^2
    # st  x + 2y + 3z >= 4 (c1)
    #     x +  y      >= 1 (c2)
    #     x, y, z \in R

    model = JuMP.direct_model(diff_optimizer(OSQP.Optimizer))
    MOI.set(JuMP.backend(model).optimizer, MOI.RawParameter("eps_prim_inf"), 1e-7)
    MOI.set(JuMP.backend(model).optimizer, MOI.RawParameter("eps_dual_inf"), 1e-7)
    MOI.set(model, MOI.Silent(), true)
    @variables(model, begin
        x
        y
        z
    end)
    @constraint(model, c1,
        x + 2y + 3z >= 4
    )
    @constraint(model, c2,
        x +  y      >= 1
    )
    @objective(model, Min,
        x^2 + x*y + y^2 + y*z + z^2
    )

    optimize!(model)

    z = JuMP.backend(model).primal_optimal

    @test z ≈ [4/7, 3/7, 6/7] atol=ATOL rtol=RTOL

    # obtain gradients
    grads = backward(JuMP.backend(model), ["Q","q","G","h"], ones(3))

    dl_dQ = grads[1]
    dl_dq = grads[2]
    dl_dG = grads[3]
    dl_dh = grads[4]

    @test dl_dQ ≈ [-0.12244895  0.01530609 -0.11224488;
                    0.01530609  0.09183674  0.07653058;
                   -0.11224488  0.07653058 -0.06122449]  atol=ATOL rtol=RTOL

    @test dl_dq ≈ [-0.2142857;  0.21428567; -0.07142857] atol=ATOL rtol=RTOL

    @test_broken dl_dG ≈ [0.05102035   0.30612245  0.255102;
                   0.06122443   0.36734694  0.3061224] atol=ATOL rtol=RTOL

    @test dl_dh ≈ [-0.35714284; -0.4285714] atol=ATOL rtol=RTOL
end

@testset "Differentiating MOI examples 2 - non trivial backward pass vector" begin
    # non-homogeneous quadratic objective
    #    minimize 2 x^2 + y^2 + xy + x + y
    #       s.t.  x, y >= 0
    #             x + y = 1 (c1)

    model = JuMP.Model(() -> diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x ≥ 0);
    @variable(model, y ≥ 0);
    @constraint(model, c1, x + y == 1)
    @objective(model, Min, 2 * x^2 + y^2 + x * y + x + y)

    optimize!(model)

    doptimizer = JuMP.backend(model).optimizer.model
    z = doptimizer.primal_optimal

    @test z ≈ [0.25, 0.75] atol=ATOL rtol=RTOL

    # obtain gradients
    dl_dz = [1.3, 0.5]   # choosing a non trivial backward pass vector
    grads = backward(doptimizer, ["Q", "q", "G", "h", "A", "b"], dl_dz)

    dl_dQ = grads[1]
    dl_dq = grads[2]
    dl_dG = grads[3]
    dl_dh = grads[4]
    dl_dA = grads[5]
    dl_db = grads[6]

    @test dl_dQ ≈ [-0.05   -0.05;
                   -0.05    0.15]  atol=ATOL rtol=RTOL

    @test dl_dq ≈ [-0.2; 0.2] atol=ATOL rtol=RTOL

    @test dl_dG ≈ [1e-8  1e-8; 1e-8 1e-8] atol=ATOL rtol=RTOL

    @test dl_dh ≈ [1e-8; 1e-8] atol=ATOL rtol=RTOL

    @test dl_dA ≈ [0.375 -1.075] atol=ATOL rtol=RTOL

    @test dl_db ≈ [0.7] atol=ATOL rtol=RTOL
end

@testset "Differentiating non trivial convex QP JuMP" begin
    nz = 10
    nineq_le = 25
    neq = 10

    # read matrices from files
    names = ["P", "q", "G", "h", "A", "b"]
    matrices = []

    for name in names
        push!(matrices, readdlm(joinpath(dirname(dirname(pathof(DiffOpt))), "test", "data", name * ".txt"), ' ', Float64, '\n'))
    end

    Q, q, G, h, A, b = matrices
    q = vec(q)
    h = vec(h)
    b = vec(b)

    model = JuMP.Model(() -> diff_optimizer(Ipopt.Optimizer))
    MOI.set(model, MOI.Silent(), true)

    @variable(model, x[1:nz])

    @objective(model, Min, dot(Q * x, x) + dot(q, x))
    @constraint(model, G * x .<= h)
    @constraint(model, A * x .== b)

    optimize!(model)

    doptimizer = JuMP.backend(model).optimizer.model

    # obtain gradients
    grads = backward(doptimizer, ["Q", "q", "G", "h", "A", "b"], ones(nz))  # using dl_dz=[1,1,1,1,1,....]

    # read gradients from files
    param_names = ["dP", "dq", "dG", "dh", "dA", "db"]
    grads_actual = []

    for name in param_names
        push!(grads_actual,
            readdlm(
                joinpath(@__DIR__, "data", "$(name).txt"),
                ' ', Float64, '\n',
            )
        )
    end

    grads_actual[2] = vec(grads_actual[2])
    grads_actual[4] = vec(grads_actual[4])
    grads_actual[6] = vec(grads_actual[6])

    # testing differences
    for i in 1:size(grads)[1]
        @test grads[i] ≈  grads_actual[i] atol=1e-2 rtol=1e-2
    end
end

@testset "Differentiating LP; checking gradients for non-active contraints" begin
    # Issue #40 from Gurobi.jl
    # min  x
    # s.t. x >= 0
    #      x >= 3

    model = direct_model(diff_optimizer(Clp.Optimizer))
    MOI.set(model, MOI.Silent(), true)

    @variable(model, x[1:1])

    @objective(model, Min, 1.1 * x[1])
    @constraint(model, x[1] ≥ 0)
    @constraint(model, x[1] ≥ 3)

    optimize!(model)

    # obtain gradients
    grads = backward(JuMP.backend(model), ["G", "h"], [1.0])

    @test grads[1] ≈ [0.0, 3.0] atol=ATOL rtol=RTOL
    @test grads[2] ≈ [0.0, -1.0] atol=ATOL rtol=RTOL

    model = direct_model(diff_optimizer(Clp.Optimizer))
    MOI.set(model, MOI.Silent(), true)

    @variable(model, x[1:1])

    @objective(model, Min, x[1])
    @constraint(model, x[1] ≥ 0)
    @constraint(model, x[1] ≥ 3)

    optimize!(model)

    # obtain gradients
    # fixes https://github.com/jump-dev/DiffOpt.jl/issues/82
    grads = backward(JuMP.backend(model), ["G", "h"], [1.0])
    @test grads[1] ≈ [0.0, 3.0] atol=ATOL rtol=RTOL
    @test grads[2] ≈ [0.0, -1.0] atol=ATOL rtol=RTOL
end

@testset "Differentiating LP; checking gradients for non-active contraints" begin
    # refered from - https://en.wikipedia.org/wiki/Simplex_algorithm#Example

    # max 2x + 3y + 4z
    # s.t. 3x+2y+z <= 10
    #      2x+5y+3z <= 15
    #      x,y,z >= 0

    model = direct_model(diff_optimizer(SCS.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, v[1:3] ≥ 0)
    @objective(model, Min, dot([-2.0, -3.0, -4.0], v))
    (x, y, z) = v

    @constraint(model, 3x+2y+z <= 10)
    @constraint(model, 2x+5y+3z <= 15)
    optimize!(model)

    # obtain gradients
    grads = backward(JuMP.backend(model), ["Q", "q", "G", "h"], ones(3))  # using dl_dz=[1,1,1,1,1,....]

    @test grads[1] ≈ zeros(3,3) atol=ATOL rtol=RTOL
    @test grads[2] ≈ zeros(3) atol=ATOL rtol=RTOL
    @test grads[3] ≈ [0.0 0.0 0.0;
                    0.0 0.0 -5/3;
                    0.0 0.0 5/3;
                    0.0 0.0 -10/3;
                    0.0 0.0 0.0]   atol=ATOL rtol=RTOL
    @test grads[4] ≈ [0.0; 1/3; -1/3; 2/3; 0.0]   atol=ATOL rtol=RTOL
end

@testset "Differentiating LP with variable bounds" begin
    # max 2x + 3y + 4z
    # s.t. 3x+2y+z <= 10
    #      2x+5y+3z <= 15
    #      x ≤ 3
    #      0 ≤ y ≤ 2
    #      z ≥ 2
    #      x,y,z >= 0
    # variant of previous test with same solution

    model = direct_model(diff_optimizer(Clp.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, v[1:3] ≥ 0)
    (x, y, z) = v
    @objective(model, Max, 2x + 3y + 4z)
    @constraint(model, 3x+2y+z <= 10)
    @constraint(model, 2x+5y+3z <= 15)
    @constraint(model, x ≤ 3)
    @constraint(model, 1.0 * y ≤ 2)
    @constraint(model, z ≥ 2)

    optimize!(model)

    # obtain gradients
    grads = backward(JuMP.backend(model), ["Q", "q", "G", "h"], ones(3))  # using dl_dz=[1,1,1]

    @test grads[1] ≈ zeros(3,3) atol=ATOL rtol=RTOL
    @test grads[2] ≈ zeros(3) atol=ATOL rtol=RTOL
    @test grads[3] ≈ [0.0 0.0 0.0
                      0.0 0.0 -5/3
                      0.0 0.0 0.0
                      0.0 0.0 0.0
                      0.0 0.0 0.0
                      0.0 0.0 5/3
                      0.0 0.0 -10/3
                      0.0 0.0 0.0
                      ]   atol=ATOL rtol=RTOL
    @test grads[4] ≈ [0.0, 1/3, 0.0, 0.0, 0.0, -1/3, 2/3, 0.0] atol=ATOL rtol=RTOL

end

@testset "Differentiating simple SOCP" begin
    # referred from _soc2test, https://github.com/jump-dev/MathOptInterface.jl/blob/master/src/Test/contconic.jl#L1355
    # find equivalent diffcp python program here: https://github.com/AKS1996/jump-gsoc-2020/blob/master/diffcp_socp_1_py.ipynb

    # Problem SOC2
    # min  x
    # s.t. y ≥ 1/√2
    #      x² + y² ≤ 1
    # in conic form:
    # min  x
    # s.t.  -1/√2 + y ∈ R₊
    #        1 - t ∈ {0}
    #      (t,x,y) ∈ SOC₃

    model = JuMP.direct_model(diff_optimizer(SCS.Optimizer))
    MOI.set(model, MOI.Silent(), true)
    @variable(model, x)
    @variable(model, y)
    @variable(model, t)

    @objective(model, Min, x)
    @constraint(model, c1, [-1/√2 + y] in MOI.Nonnegatives(1))
    @constraint(model, c2, [1 - t] in MOI.Zeros(1))
    @constraint(model, c3, [1.0 * t, x, y] in SecondOrderCone())

    optimize!(model)

    v = JuMP.backend(model).primal_optimal
    # slack variables
    s = collect(Iterators.flatten(JuMP.value.([c1, c2, c3])))
    y = collect(Iterators.flatten(JuMP.dual.([c1, c2, c3])))

    # these matrices are benchmarked with the output generated by diffcp
    # refer the python file mentioned above to get equivalent python source code
    @test v ≈ [-1/√2; 1/√2; 1.0] atol=ATOL rtol=RTOL
    @test s ≈ [0.0, 0.0, 1.0, -1/√2, 1/√2] atol=ATOL rtol=RTOL
    @test y ≈ [1, √2, √2, 1, -1] atol=ATOL rtol=RTOL

    dA = zeros(5, 3)
    dA[1:3, :] .= Matrix(1.0I, 3, 3)
    db = zeros(5)
    dc = zeros(3)

    @test_broken backward(JuMP.backend(model), dA, db, dc)

end