using ForwardDiff, Distributions, FiniteDifferences, Tracker, Random, LinearAlgebra, PDMats
using Turing: Turing, gradient_logp_reverse, invlink, link, SampleFromPrior
using Turing.Core.RandomVariables: getval
using Turing.Core: TuringMvNormal, TuringDiagNormal
using ForwardDiff: Dual
using StatsFuns: binomlogpdf, logsumexp
using Test, LinearAlgebra
const FDM = FiniteDifferences

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

_to_cov(B) = B * B' + Matrix(I, size(B)...)
@testset "ad.jl" begin
    @turing_testset "adr" begin
        ad_test_f = gdemo_default
        vi = Turing.VarInfo()
        ad_test_f(vi, SampleFromPrior())
        svn = vi.vns[1]
        mvn = vi.vns[2]
        _s = getval(vi, svn)[1]
        _m = getval(vi, mvn)[1]

        x = map(x->Float64(x), vi[SampleFromPrior()])
        ∇E = gradient_logp_reverse(x, vi, ad_test_f)[2]
        grad_Turing = sort(∇E)

        dist_s = InverseGamma(2,3)

        # Hand-written logp
        function logp(x::Vector)
          s = x[2]
          # s = invlink(dist_s, s)
          m = x[1]
          lik_dist = Normal(m, sqrt(s))
          lp = logpdf(dist_s, s) + logpdf(Normal(0,sqrt(s)), m)
          lp += logpdf(lik_dist, 1.5) + logpdf(lik_dist, 2.0)
          lp
        end

        # Call ForwardDiff's AD
        g = x -> ForwardDiff.gradient(logp, x);
        # _s = link(dist_s, _s)
        _x = [_m, _s]
        grad_FWAD = sort(g(_x))

        # Compare result
        @test grad_Turing ≈ grad_FWAD atol=1e-9
    end
    @turing_testset "passing duals to distributions" begin
        float1 = 1.1
        float2 = 2.3
        d1 = Dual(1.1)
        d2 = Dual(2.3)

        @test logpdf(Normal(0, 1), d1).value ≈ logpdf(Normal(0, 1), float1) atol=0.001
        @test logpdf(Gamma(2, 3), d2).value ≈ logpdf(Gamma(2, 3), float2) atol=0.001
        @test logpdf(Beta(2, 3), (d2 - d1) / 2).value ≈ logpdf(Beta(2, 3), (float2 - float1) / 2) atol=0.001

        @test pdf(Normal(0, 1), d1).value ≈ pdf(Normal(0, 1), float1) atol=0.001
        @test pdf(Gamma(2, 3), d2).value ≈ pdf(Gamma(2, 3), float2) atol=0.001
        @test pdf(Beta(2, 3), (d2 - d1) / 2).value ≈ pdf(Beta(2, 3), (float2 - float1) / 2) atol=0.001
    end
    @turing_testset "general AD tests" begin
        # Tests gdemo gradient.
        function logp1(x::Vector)
            dist_s = InverseGamma(2, 3)
            s = x[2]
            m = x[1]
            lik_dist = Normal(m, sqrt(s))
            lp = Turing.logpdf_with_trans(dist_s, s, false) + Turing.logpdf_with_trans(Normal(0,sqrt(s)), m, false)
            lp += logpdf(lik_dist, 1.5) + logpdf(lik_dist, 2.0)
            return lp
        end

        test_model_ad(gdemo_default, logp1, [:m, :s])

        # Test Wishart AD.
        @model wishart_ad() = begin
            v ~ Wishart(7, [1 0.5; 0.5 1])
            v
        end

        # Hand-written logp
        function logp3(x)
            dist_v = Wishart(7, [1 0.5; 0.5 1])
            v = [x[1] x[3]; x[2] x[4]]
            lp = Turing.logpdf_with_trans(dist_v, v, false)
            return lp
        end

        test_model_ad(wishart_ad(), logp3, [:v])
    end
    @turing_testset "Tracker + logdet" begin
        rng, N = MersenneTwister(123456), 7
        ȳ, B = randn(rng), randn(rng, N, N)
        test_tracker_ad(B->logdet(cholesky(_to_cov(B))), ȳ, B; rtol=1e-8, atol=1e-8)
    end
    @turing_testset "Tracker + fill" begin
        rng = MersenneTwister(123456)
        test_tracker_ad(x->fill(x, 7), randn(rng, 7), randn(rng))
        test_tracker_ad(x->fill(x, 7, 11), randn(rng, 7, 11), randn(rng))
        test_tracker_ad(x->fill(x, 7, 11, 13), rand(rng, 7, 11, 13), randn(rng))
    end
    @turing_testset "Tracker + MvNormal" begin
        rng, N = MersenneTwister(123456), 11
        B = randn(rng, N, N)
        m, A = randn(rng, N), B' * B + I

        # Generate from the TuringMvNormal
        d, back = Tracker.forward(TuringMvNormal, m, A)
        x = Tracker.data(rand(d))

        # Check that the logpdf agrees with MvNormal.
        d_ref = MvNormal(m, PDMat(A))
        @test logpdf(d, x) ≈ logpdf(d_ref, x)

        test_tracker_ad((m, B, x)->logpdf(MvNormal(m, _to_cov(B)), x), randn(rng), m, B, x)
    end
    @turing_testset "Tracker + Diagonal Normal" begin
        rng, N = MersenneTwister(123456), 11
        m, σ = randn(rng, N), exp.(0.1 .* randn(rng, N)) .+ 1

        d = TuringDiagNormal(m, σ)
        x = rand(d)

        # Check that the logpdf agrees with MvNormal.
        d_ref = MvNormal(m, σ)
        @test logpdf(d, x) ≈ logpdf(d_ref, x)

        test_tracker_ad((m, σ, x)->logpdf(MvNormal(m, σ), x), randn(rng), m, σ, x)
    end
    @turing_testset "Tracker + MvNormal Interface" begin
        # Note that we only test methods where the `MvNormal` ctor actually constructs
        # a TuringMvNormal.

        rng, N = MersenneTwister(123456), 7
        m, b, B, x = randn(rng, N), randn(rng, N), randn(rng, N, N), randn(rng, N)
        ȳ = randn(rng)

        # zero mean, dense covariance
        test_tracker_ad((B, x)->logpdf(MvNormal(_to_cov(B)), x), randn(rng), B, x)
        test_tracker_ad(B->logpdf(MvNormal(_to_cov(B)), x), randn(rng), B)

        # zero mean, diagonal covariance
        test_tracker_ad((b, x)->logpdf(MvNormal(exp.(b)), x), randn(rng), b, x)
        test_tracker_ad(b->logpdf(MvNormal(exp.(b)), x), randn(rng), b)

        # dense mean, dense covariance
        test_tracker_ad((m, B, x)->logpdf(MvNormal(m, _to_cov(B)), x),
            randn(rng),
            randn(rng, N), randn(rng, N, N), randn(rng, N),
        )
        test_tracker_ad((m, B)->logpdf(MvNormal(m, _to_cov(B)), x),
            randn(rng),
            randn(rng, N), randn(rng, N, N),
        )
        test_tracker_ad((m, x)->logpdf(MvNormal(m, _to_cov(B)), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((B, x)->logpdf(MvNormal(m, _to_cov(B)), x),
            randn(rng),
            randn(rng, N, N), randn(rng, N),
        )
        test_tracker_ad(m->logpdf(MvNormal(m, _to_cov(B)), x), randn(rng), randn(rng, N))
        test_tracker_ad(B->logpdf(MvNormal(m, _to_cov(B)), x), randn(rng), randn(rng, N, N))

        # dense mean, diagonal covariance
        test_tracker_ad((m, b, x)->logpdf(MvNormal(m, Diagonal(exp.(b))), x),
            randn(rng),
            randn(rng, N), randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((m, b)->logpdf(MvNormal(m, Diagonal(exp.(b))), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((m, x)->logpdf(MvNormal(m, Diagonal(exp.(b))), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((b, x)->logpdf(MvNormal(m, Diagonal(exp.(b))), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad(m->logpdf(MvNormal(m, Diagonal(exp.(b))), x),
            randn(rng),
            randn(rng, N),
        )
        test_tracker_ad(b->logpdf(MvNormal(m, Diagonal(exp.(b))), x),
            randn(rng),
            randn(rng, N),
        )

        # dense mean, diagonal variance
        test_tracker_ad((m, b, x)->logpdf(MvNormal(m, exp.(b)), x),
            randn(rng),
            randn(rng, N), randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((m, b)->logpdf(MvNormal(m, exp.(b)), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((m, x)->logpdf(MvNormal(m, exp.(b)), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((b, x)->logpdf(MvNormal(m, exp.(b)), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad(m->logpdf(MvNormal(m, exp.(b)), x), randn(rng), randn(rng, N))
        test_tracker_ad(b->logpdf(MvNormal(m, exp.(b)), x), randn(rng), randn(rng, N))

        # dense mean, constant covariance
        b_s = randn(rng)
        test_tracker_ad((m, b, x)->logpdf(MvNormal(m, exp(b)), x),
            randn(rng),
            randn(rng, N), randn(rng), randn(rng, N),
        )
        test_tracker_ad((m, b)->logpdf(MvNormal(m, exp(b)), x),
            randn(rng),
            randn(rng, N), randn(rng),
        )
        test_tracker_ad((m, x)->logpdf(MvNormal(m, exp(b_s)), x),
            randn(rng),
            randn(rng, N), randn(rng, N)
        )
        test_tracker_ad((b, x)->logpdf(MvNormal(m, exp(b)), x),
            randn(rng),
            randn(rng), randn(rng, N),
        )
        test_tracker_ad(m->logpdf(MvNormal(m, exp(b_s)), x), randn(rng), randn(rng, N))
        test_tracker_ad(b->logpdf(MvNormal(m, exp(b)), x), randn(rng), randn(rng))

        # dense mean, constant variance
        b_s = randn(rng)
        test_tracker_ad((m, b, x)->logpdf(MvNormal(m, exp(b) * I), x),
            randn(rng),
            randn(rng, N), randn(rng), randn(rng, N),
        )
        test_tracker_ad((m, b)->logpdf(MvNormal(m, exp(b) * I), x),
            randn(rng),
            randn(rng, N), randn(rng),
        )
        test_tracker_ad((m, x)->logpdf(MvNormal(m, exp(b_s) * I), x),
            randn(rng),
            randn(rng, N), randn(rng, N),
        )
        test_tracker_ad((b, x)->logpdf(MvNormal(m, exp(b) * I), x),
            randn(rng),
            randn(rng), randn(rng, N),
        )
        test_tracker_ad(m->logpdf(MvNormal(m, exp(b_s) * I), x), randn(rng), randn(rng, N))
        test_tracker_ad(b->logpdf(MvNormal(m, exp(b) * I), x), randn(rng), randn(rng))

        # zero mean, constant variance
        test_tracker_ad((b, x)->logpdf(MvNormal(N, exp(b)), x),
            randn(rng),
            randn(rng), randn(rng, N),
        )
        test_tracker_ad(b->logpdf(MvNormal(N, exp(b)), x), randn(rng), randn(rng))
    end
    @testset "Simplex Tracker AD" begin
        @model dir() = begin
            theta ~ Dirichlet(1 ./ fill(4, 4))
        end
        Turing.setadbackend(:reverse_diff)
        sample(dir(), HMC(0.01, 1), 1000);
    end
    @testset "PDMatDistribution Tracker AD" begin
        @model wishart() = begin
            theta ~ Wishart(4, Matrix{Float64}(I, 4, 4))
        end
        Turing.setadbackend(:reverse_diff)
        sample(wishart(), HMC(0.01, 1), 1000);

        @model invwishart() = begin
            theta ~ InverseWishart(4, Matrix{Float64}(I, 4, 4))
        end
        Turing.setadbackend(:reverse_diff)
        sample(invwishart(), HMC(0.01, 1), 1000);
    end
end


################################################################################
#
#  Pure and Kummer extensions
#
################################################################################

@doc Markdown.doc"""
    isradical_extension(L::SimpleNumField) -> Bool

Tests if $L/K$ is pure, that is, if the defining polynomial is of the form
$x^n - b$ for some $b \in K$.
"""
function isradical_extension(K::SimpleNumField)
  if !ismonic(K.pol)
    return false
  end
  return all(i -> iszero(coeff(K.pol, i)), 1:degree(K)-1)
end

@doc Markdown.doc"""
    iskummer_extension(L::SimpleNumField) -> Bool

Tests if $L/K$ is a Kummer extension, that is, if the defining polynomial is
of the form $x^n - b$ for some $b \in K$ and if $K$ contains the $n$-th roots
of unity.
"""
function iskummer_extension(K::SimpleNumField)
  if !isradical_extension(K)
    return false
  end

  k = base_field(K)
  Zk = maximal_order(k)
  _, o = torsion_units_gen_order(Zk)
  if o % degree(K) != 0
    return false
  end
  return true
end

function iskummer_extension(K::AnticNumberField)
  if degree(K) != 2
    return false
  end
  return isradical_extension(K)
end

function radical_extension(n::Int, a::FacElem;
                        cached::Bool = true, check::Bool = true)
  return radical_extension(n, evaluate(a), cached = cached, check = check)
end

@doc Markdown.doc"""
    radical_extension(n::Int, a::NumFieldElem; s = "_$",
                   check = true, cached = true) -> NumField, NumFieldElem

Given an element $a$ of a number field $K$ and an integer $n$, create the simple
extension of $K$ with the defining polynomial $x^n - a$.
"""
function radical_extension(n::Int, a::NumFieldElem; s::String = "_\$",
                        cached::Bool = true, check::Bool = true)
  k = parent(a)
  kx, x = PolynomialRing(k, cached = false)
  return number_field(x^n - a, s, check = check, cached = cached)
end

function radical_extension(n::Int, a::fmpq; s::String = "_\$",
                        cached::Bool = true, check::Bool = true)
  k = parent(a)
  kx, x = PolynomialRing(k, cached = false)
  return number_field(x^n - a, s, check = check, cached = cached)
end

################################################################################
#
#  Basis
#
################################################################################

@doc Markdown.doc"""
    basis(L::SimpleNumField) -> Vector{NumFieldElem}

Returns the canonical basis of a simple extension $L/K$, that is, the elements
$1,a,\dotsc,a^{d - 1}$, where $d$ is the degree of $K$ and $a$ the primitive
element.
"""
basis(::SimpleNumField)

################################################################################
#
#  Defining polynomial
#
################################################################################

export defining_polynomial

@doc Markdown.doc"""
    defining_polynomial(L::SimpleNumField) -> PolyElem

Given a simple number field $L/K$, constructed as $L = K[x]/(f)$, this function
returns $f$.
"""
defining_polynomial(::SimpleNumField)

defining_polynomial(K::NfRel) = K.pol

defining_polynomial(K::AnticNumberField) = K.pol

################################################################################
#
#  Discriminant
#
################################################################################

@doc Markdown.doc"""
    discriminant(L::SimpleNumField) -> NumFieldElem

The discriminant of the defining polynomial of $L$, *not* the discriminant of
the maximal order of $L$.
"""
function discriminant(K::SimpleNumField)
  return discriminant(defining_polynomial(K))
end

################################################################################
#
#  Absolute discriminant
#
################################################################################

@doc Markdown.doc"""
    absolute_discriminant(L::SimpleNumField, QQ) -> fmpq

The absolute discriminant of the defining polynomial of $L$, *not* the
discriminant of the maximal order of $L$. This is the norm of the discriminant
times the $d$-th power of the discriminant of the base field, where $d$ is the
degree of $L$.
"""
absolute_discriminant(::SimpleNumField)

function absolute_discriminant(K::AnticNumberField)
  return discriminant(K)
end

function absolute_discriminant(K::NfRel)
  d = norm(discriminant(K)) * absolute_discriminant(base_field(K))^degree(K)
  return d
end

function discriminant(K::FlintRationalField)
  return one(K)
end

################################################################################
#
#  Is subfield
#
################################################################################

@doc Markdown.doc"""
    issubfield(K::SimpleNumField, L::SimpleNumField) -> Bool, Map

Returns `true` and an injection from $K$ to $L$ if $K$ is a subfield of $L$.
Otherwise the function returns `false` and a morphism mapping everything to
$0$.
"""
issubfield(::SimpleNumField, ::SimpleNumField)

@doc Markdown.doc"""
    isisomorphic(K::SimpleNumField, L::SimpleNumField) -> Bool, Map

Returns `true` and an isomorphism from $K$ to $L$ if $K$ and $L$ are isomorphic.
Otherwise the function returns `false` and a morphism mapping everything to $0$.
"""
isisomorphic(::SimpleNumField, ::SimpleNumField)

################################################################################
#
#  Linear disjointness
#
################################################################################

# TODO (easy): Do this for Non-Simple number fields
@doc Markdown.doc"""
    islinear_disjoint(K::SimpleNumField, L::SimpleNumField) -> Bool

Given two number fields $K$ and $L$ with the same base field $k$, this function
returns whether $K$ and $L$ are linear disjoint over $k$.
"""
islinear_disjoint(K::SimpleNumField, L::SimpleNumField)


using FactCheck, Base.Test, Images, Colors, FixedPointNumbers

facts("Corner") do
    A = zeros(41,41)
    A[16:26,16:26] = 1

	context("Harris") do
		Ac = Images.imcorner(A,method="harris")
		# check corner intensity
		@fact Ac[16,16] --> roughly(maximum(Ac))
		@fact Ac[16,26] --> roughly(Ac[16,16])
		@fact Ac[26,16] --> roughly(Ac[16,16])
		@fact Ac[26,26] --> roughly(Ac[16,16])

		# check edge intensity
		@fact Ac[16,21] --> roughly(minimum(Ac))
		@fact Ac[21,16] --> roughly(Ac[16,21])
		@fact Ac[21,26] --> roughly(Ac[16,21])
		@fact Ac[26,21] --> roughly(Ac[16,21])

		# check background intensity
		@fact Ac[5,5] --> roughly(0)
		@fact Ac[5,36] --> roughly(0)
		@fact Ac[36,5] --> roughly(0)
		@fact Ac[36,36] --> roughly(0)

	end

	context("Shi-Tomasi") do
		Ac = Images.imcorner(A,method="shi-tomasi")
		# check corner intensity
		@fact Ac[16,16] --> roughly(maximum(Ac))
		@fact Ac[16,26] --> roughly(Ac[16,16])
		@fact Ac[26,16] --> roughly(Ac[16,16])
		@fact Ac[26,26] --> roughly(Ac[16,16])

		# check edge intensity
		@fact Ac[16,21] --> roughly(0)
		@fact Ac[21,16] --> roughly(0)
		@fact Ac[21,26] --> roughly(0)
		@fact Ac[26,21] --> roughly(0)

		# check background intensity
		@fact Ac[5,5] --> roughly(0)
		@fact Ac[5,36] --> roughly(0)
		@fact Ac[36,5] --> roughly(0)
		@fact Ac[36,36] --> roughly(0)

	end

end


#=
export DirectProductOperation
export isidentity
export directproduct
export ×ˢ

"""
    DirectProductOperation{Ops}

Direct product of symmetry operations.
math```
  g = (g_1, g_2, \\ldots)
```
"""
struct DirectProductOperation{Ops<:Tuple{Vararg{AbstractSymmetryOperation}}}<:AbstractSymmetryOperation
    operations::Ops
    function DirectProductOperation(ops::Vararg{AbstractSymmetryOperation})
        T = typeof(ops)
        return new{T}(ops)
    end
    function DirectProductOperation(ops::T) where {T<:Tuple{Vararg{AbstractSymmetryOperation}}}
        return new{T}(ops)
    end
end

Base.hash(x::P, h::UInt) where {P<:DirectProductOperation} = Base.hash(P, Base.hash(x.operations, h))

Base.:(*)(lhs::P, rhs::P) where {P<:DirectProductOperation} = DirectProductOperation(lhs.operations .* rhs.operations)
Base.:(^)(obj::DirectProductOperation, n::Integer) = DirectProductOperation(obj.operations.^n)
Base.:(==)(lhs::P, rhs::P) where {P<:DirectProductOperation} = all(lhs.operations .== rhs.operations)

Base.inv(obj::DirectProductOperation) = DirectProductOperation(Base.inv.(obj.operations))

function Base.isapprox(lhs::P, rhs::P; atol::Real=0, rtol::Real=Base.rtoldefault(Float64)) where {P<:DirectProductOperation}
    return all(isapprox.(lhs.operations, rhs.operations; atol=atol, rtol=rtol))
end

isidentity(obj::DirectProductOperation) = all(isidentity, obj.operations)

×ˢ(lhs::AbstractSymmetryOperation, rhs::AbstractSymmetryOperation) = DirectProductOperation(lhs, rhs)
×ˢ(lhs::DirectProductOperation, rhs::AbstractSymmetryOperation) = DirectProductOperation(lhs.operations..., rhs)
×ˢ(lhs::AbstractSymmetryOperation, rhs::DirectProductOperation) = DirectProductOperation(lhs, rhs.operations...)
×ˢ(lhs::DirectProductOperation, rhs::DirectProductOperation) = DirectProductOperation(lhs.operations..., rhs.operations...)
function ×ˢ(lhs::AbstractVector{<:AbstractSymmetryOperation}, rhs::AbstractVector{<:AbstractSymmetryOperation})
    return [×ˢ(l, r) for l in lhs, r in rhs]
end

directproduct(lhs::AbstractSymmetryOperation, rhs::AbstractSymmetryOperation) = DirectProductOperation(lhs, rhs)
directproduct(lhs::DirectProductOperation, rhs::AbstractSymmetryOperation) = DirectProductOperation(lhs.operations..., rhs)
directproduct(lhs::AbstractSymmetryOperation, rhs::DirectProductOperation) = DirectProductOperation(lhs, rhs.operations...)
directproduct(lhs::DirectProductOperation, rhs::DirectProductOperation) = DirectProductOperation(lhs.operations..., rhs.operations...)
function directproduct(lhs::AbstractVector{<:AbstractSymmetryOperation}, rhs::AbstractVector{<:AbstractSymmetryOperation})
    return [directproduct(l, r) for l in lhs, r in rhs]
end

function directproduct(::Type{E}, products::Function...) where {E<:DirectProductOperation}
    function product(lhs::E, rhs::E)
        return DirectProductOperation([p(l, r) for (p, l, r) in zip(products, lhs.operations, rhs.operations)]...)
    end
    return product
end
=#


using LinearAlgebra
using StaticArrays
using Plots; pyplot()
using FastClosures
using NPZ
using LaTeXStrings

include("MagnusExpansion.jl")
include("AuxFunctions.jl")
include("HarmonicCorrectionSnap.jl")
include("snap.jl")

function run(tf)
	σ = pauli_matrices()
	χ_GHz = (-8281.0, 0.0, 0.0)
	χ = χ_GHz./abs(χ_GHz[1])
	N = 10
	levels = SA[0, 4]
	α = (pi, pi)./2
	Ncorr = 10

	system_params = Dict("chi"=>χ, "N"=>N, "levels"=>levels, "alpha"=>α,
						 "tf"=>tf, "Ncorr"=>Ncorr)

	Identity = Matrix{Complex{Float64}}(I, 2*N, 2*N)
	Uideal = U(system_params, σ; ideal_dynamics=true, matrix=true)
	Ur = U(system_params, σ; ideal_dynamics=false, matrix=true)
	F0 = fidelity(Ur, Uideal, 2*N)
	# "First" order correction
	magnus_order = 1
	max_order = 2
	coeffs = get_coeffs(magnus_order, max_order, Ncorr, system_params)

	Ucorr = U(system_params, σ, coeffs, max_order; frame="qubit", matrix=true)
	F1 = fidelity(Ucorr, Uideal, 2*N)
	# "Second" order correction
	magnus_order = 2
	coeffs2 = get_coeffs(coeffs, magnus_order, max_order, Ncorr, system_params)

	coeffs2 = coeffs + coeffs2
	Ucorr = U(system_params, σ, coeffs2, max_order; frame="qubit", matrix=true)
	F2 = fidelity(Ucorr, Uideal, 2*N)

	return F0, F1, F2
end

npoints = 100
tf0, tf1 = 15, 200
r = (tf1/tf0)^(1/(npoints-1))
tf_array = tf0*r.^collect(0: npoints-1)
data = zeros(npoints, 4)
for i = 1: npoints
	tf = tf_array[i]
	println(tf)
	data[i, 1] = tf
	data[i, 2: end] .= run(tf)
	println(data[i, 2: end])
end

x = data[:,1]
y = log10.(1 .- data[:, [2, 3, 4]])
plot(x, y,
	 linewidth = 2,
	 labels = ["No\ncorrection" "2nd order\nMagnus" "4th order\nMagnus"])
plot!(tickfont = font(12, "Serif"),
      xlabel = "gate time "*L"|\chi| t_\mathrm{f}",
      ylabel = "Fidelity error "*L"\varepsilon",
      guidefont = font(14, "Serif"),
	  legendfont = font(12, "Serif"),
	  background_color_legend = false,
      margin = 5Plots.mm)
savefig("fidelity_error.pdf")
npzwrite("data.npy", data)


using Match

input = joinpath(@__DIR__, "input")
instructions = readlines(input)

instructions = map(instructions) do ins
    ins = split(ins)
    (first(ins), parse(Int, last(ins)))
end

function navigate(instructions, depth, pos)
    for i in instructions
        (op, val) = i
        @match op begin
            "forward" => begin pos += val end
            "up" => begin depth -= val end
            "down" => begin depth += val end
        end
    end
    return depth * pos
end

function navigate(instructions, depth, pos, aim)
    for i in instructions
        (op, val) = i
        @match op begin
            "forward" => begin
                pos += val
                depth += (val * aim)
            end
            "up" => begin aim -= val end
            "down" => begin aim += val end
        end
    end
    return depth * pos
end

p1 = navigate(instructions, 0, 0)
p2 = navigate(instructions, 0, 0, 0)

println("-----------------------------------------------------------------------")
println("dive! -- part one :: $p1")
println("dive! -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 1714680)
@assert(p2 == 1963088820)


bugs = Int8.(map(==("#"), hcat(split.(readlines(ARGS[1]), "")...)))'

function adjacent(m, row, col)
    get(m, (row,col+1), 0) +
        get(m, (row,col-1), 0) +
        get(m, (row+1,col), 0) +
        get(m, (row-1,col), 0)
end

function life(m)
    m2 = Array{Int,2}(undef, 5, 5)

    for row=1:5
        for col=1:5
            adj = adjacent(m, row, col)
            if m[row,col] == 1
                m2[row,col] = adj == 1 ? 1 : 0
            else
                m2[row,col] = (1 == adj || 2 == adj) ? 1 : 0
            end
        end
    end

    return m2
end

# Every possible state has one of 33554432 possible states.
# They can be efficiently serialized into an integer.
# This is column-major, meaning you will get one column
# at a time. It doesn't really matter as long as the "signature"
# algorithm does not change at runtime.
# Check this out:
#   a = rand(collect('a':'z'), 5, 5)
#   foldl((x,y) -> x * y, a)
biodiversity(m) =foldr((x,y) -> x+2y, m')

function render(m)
    s = ""
    for row=1:size(m)[1]
        for col=1:size(m)[2]
            s *= m[row,col] == 1 ? "#" : "."
        end
        s *= "\n"
    end
    return s
end

function test(m)
    println("Initial state:")
    print(render(m))
    println("Biodiversity rating: $(biodiversity(m))\n")

    for i=1:4
        m = life(m)
        println("After $(i) minute$(i > 1 ? "s" : ""):")
        print(render(m))
        println("Biodiversity rating: $(biodiversity(m))\n")
    end
end

#test(bugs)

seen = Set()
while true
    d = biodiversity(bugs)
    if d ∈ seen
        println("Day 24 Part 1: $(d)")
        print(render(bugs))
        break
    end
    push!(seen, d)
    global bugs = life(bugs)
end


module GSTest
using GenericSchur
using Test
using Random
using LinearAlgebra
include("TMGlib.jl")
using .TMGlib

import LinearAlgebra.BLAS.BlasFloat

const verbosity = Ref(0)

cplxord = t -> (real(t), imag(t))
csort(v) = sort(v, by = cplxord)

"""
Generate Godunov's strongly non-normal matrix with known eigenvalues.

One needs to retain at least 70 bits of precision to have any confidence in
computing the eigenvalues here.
"""
function godunov(T)
    A = convert.(T,[289 2064 336 128 80 32 16;
                    1152 30 1312 512 288 128 32;
                    -29 -2000 756 384 1008 224 48;
                    512 128 640 0 640 512 128;
                    1053 2256 -504 -384 -756 800 208;
                    -287 -16 1712 -128 1968 -30 2032;
                    -2176 -287 -1565 -512 -541 -1152 -289])
    vals = [-4,-2,-1,0,1,2,4]
    econd = 7e16 # condition of the worst ones
    A,vals,econd
end

include("wrappers.jl")
include("symtridiag.jl")
include("complex.jl")
include("balance.jl")
include("real.jl")
#include("complex.jl")

include("ordschur.jl")

include("generalized.jl")
include("gordschur.jl")

include("errors.jl")

end # module


# # Elipse

# ## Model

module elipse   #jl

using ReachabilityAnalysis

A = [3.0 -9.0;
     4.0 -3.0]

function model(X0)
    S = @system(x' = Ax)
    return IVP(S, X0)
end

end #jl

# ## References


#
# Solve the 2d wave equation using finite differencing
#
# @author arbennett
#

# Begin timing
#t_0 = time()

# Parameters for solving
const N = 50
const M = 50
const T = 6000
const dt = 0.1
const plot_frequency = 10
u = zeros(3,N,M)
f = open("solution.dat", "w")

# Set the initial conditions
for i=1:N
    for j=1:M
        u[2:3,i,j] = 1 + cos(2 * pi * (i/N)) + cos(2 * pi * (j/M))
        #u[3,i,j] = sin(2 * pi * ((i+1)/N)) + cos(2 * pi * ((j+1)/M))
    end
end

# Enforce boundary conditions
u[:,:,1] = u[:,:,2]
u[:,:,end] = u[:,:,end-1]
u[:,1,:] = u[:,2,:]
u[:,end,:] = u[:,end-1,:]

# Do some timestepping
for t=1:int(T)
    # Calculate the new u
    u[1,2:end-1,2:end-1] = (2*u[2,2:end-1,2:end-1] - u[3,2:end-1,2:end-1] +
        dt^2 * (u[2,2:end-1,1:end-2] + u[2,2:end-1,3:end] + u[2,1:end-2,2:end-1]
        + u[2,3:end,2:end-1] - 4*u[2,2:end-1,2:end-1]))

    # Enforce boundary conditions
    u[:,:,1] = u[:,:,2]
    u[:,:,end] = u[:,:,end-1]
    u[:,1,:] = u[:,2,:]
    u[:,end,:] = u[:,end-1,:]

    # Move the time forwards
    u = circshift(u,[1,0,0])
    if mod(t, plot_frequency) == 0
        writedlm("out/solution.dat."*lpad(string(t),length(string(int(T))),"0"), u[1,:,:])
    end
end

# Report the time taken
#t_1 = time()
#println("Time elapsed: ", t_1 - t_0, "s")


module Functional

using LinearAlgebra: Matrix, zero, similar
using LinearAlgebra, Printf
# using Roots
# using Optim
using Quadmath
# using ProfileView
# using InteractiveUtils

using ..Discrete

# const Float = Float64
# const Float = BigFloat
const Float = Float128

include("kernel.jl")
include("findmax.jl")

# using Plots
# function plotResidual(basis, proj, gmin, gmax, candidate=nothing, residual=nothing)
#     ω = LinRange(gmin, gmax, 1000)
#     y = [Residual(basis, proj, w) for w in ω]
#     p = plot(ω, y, xlims=(gmin, gmax))
#     if isnothing(candidate) == false
#         plot!(p, candidate, residual, seriestype=:scatter)
#     end
#     display(p)
#     readline()
# end

mutable struct Basis
    Λ::Float
    rtol::Float

    N::Int # number of basis
    grid::Vector{Float} # grid for the basis
    residual::Vector{Float} # achieved error by each basis
    Q::Matrix{Float} # K = Q*R
    proj::Matrix{Float} # the overlap of basis functions <K(g_i), K(g_j)>

    candidate::Vector{Float}
    candidateResidual::Vector{Float}

    function Basis(Λ, rtol)
        _Q = Matrix{Float}(undef, (0, 0))
        return new(Λ, rtol, 0, [], [], _Q, similar(_Q), [], [])
    end
end

function addBasis!(basis, proj, g0::Float)
    basis.N += 1
    if basis.N == 1
        idx = 1
        basis.grid = [g0,]
        basis.Q = zeros(Float, (basis.N, basis.N))
        basis.Q[1, 1] = 1 / sqrt(proj(basis.Λ, g0, g0))
        basis.proj = projKernel(basis, proj)
    else
        idxList = findall(x -> x > g0, basis.grid)
        # if ω is larger than any existing freqencies, then idx is an empty list
        idx = length(idxList) == 0 ? basis.N : idxList[1] # the index to insert the new frequency

        insert!(basis.grid, idx, g0)
        basis.proj = projKernel(basis, proj)
        _Q = copy(basis.Q)
        basis.Q = zeros(Float, (basis.N, basis.N))
        basis.Q[1:idx-1, 1:idx-1] = _Q[1:idx-1, 1:idx-1]
        basis.Q[1:idx-1, idx+1:end] = _Q[1:idx-1, idx:end]
        basis.Q[idx+1:end, 1:idx-1] = _Q[idx:end, 1:idx-1]
        basis.Q[idx+1:end, idx+1:end] = _Q[idx:end, idx:end]
        # println(maximum(abs.(GramSchmidt(basis, idx, g0) .- mGramSchmidt(basis, idx, g0))))
        basis.Q[idx, :] = mGramSchmidt(basis, idx, g0)
    end

    scanResidual!(basis, proj, g0, idx)
    insert!(basis.residual, idx, maximum(basis.candidateResidual)) # record error after the new grid is added
    return idx
end

function scanResidual!(basis, proj, g0, idx)
    grids = copy(basis.grid)
    if basis.grid[1] > Float(0)
        insert!(grids, 1, Float(0))
    end
    if basis.grid[end] < basis.Λ
        append!(grids, basis.Λ)
    end
    resize!(basis.candidate, length(grids) - 1)
    resize!(basis.candidateResidual, length(grids) - 1)
    # println(g0, " and ", idx)
    # println(grids)

    for i = 1:length(grids)-1 # because of the separation of scales, the grids far away from idx is rarely affected
        g = findCandidate(basis, proj, grids[i], grids[i+1])
        basis.candidate[i] = g
        basis.candidateResidual[i] = Residual(basis, proj, g)
    end
end

function printCandidate(basis, idx)
    lower = (idx == 1) ? 0 : basis.grid[idx-1]
    upper = (idx == basis.N) ? basis.Λ : basis.grid[idx+1]

    @printf("%3i : ω=%24.8f ∈ (%24.8f, %24.8f) -> error=%24.16g\n", basis.N, basis.grid[idx], lower, upper, basis.residual[idx])
end

function QR(Λ, rtol, proj, g0; N = nothing)
    basis = Basis(Λ, rtol)
    # println(g0)
    for g in g0
        idx = addBasis!(basis, proj, Float(g))
        # @printf("%3i : ω=%24.8f ∈ (%24.8f, %24.8f) -> error=%24.16g\n", 1, g, 0, Λ, basis.residual[idx])
        printCandidate(basis, idx)
    end

    # @code_warntype Residual(basis, proj, Float(1.0))
    # exit(0)
    maxResidual, ωi = findmax(basis.candidateResidual)
    # plotResidual(basis, proj, Float(0), Float(100), basis.candidate, basis.candidateResidual)

    while isnothing(N) ? maxResidual > rtol / 10 : basis.N < N

        newω = basis.candidate[ωi]
        idx = addBasis!(basis, proj, newω)
        printCandidate(basis, idx)
        # println(length(basis.grid))
        # println(idx)
        # lower = (idx == 1) ? 0 : basis.grid[idx - 1]
        # upper = (idx == basis.N) ? Λ : basis.grid[idx + 1]

        # @printf("%3i : ω=%24.8f ∈ (%24.8f, %24.8f) -> error=%24.16g\n", basis.N, newω, lower, upper, basis.residual[idx])
        # println("$(length(freq)) basis: ω=$(Float64(newω)) between ($(Float64(freq[idx - 1])), $(Float64(freq[idx + 1])))")
        # plotResidual(basis, proj, Float(0), Float(100), candidate, residual)
        maxResidual, ωi = findmax(basis.candidateResidual)
    end
    testOrthgonal(basis)
    # @printf("residual = %.10e, Fnorm/F0 = %.10e\n", residual, residualF(freq, Q, Λ))
    @printf("residual = %.10e\n", maximum(basis.candidateResidual))
    # plotResidual(basis, proj, Float(0), Float(100), basis.candidate, basis.candidateResidual)
    return basis
end

"""
q1=sum_j c_j K_j
q2=sum_k d_k K_k
return <q1, q2> = sum_jk c_j*d_k <K_j, K_k>
"""
projqq(basis, q1::Vector{Float}, q2::Vector{Float}) = q1' * basis.proj * q2

"""
<K(g_i), K(g_j)>
"""
function projKernel(basis, proj)
    K = zeros(Float, (basis.N, basis.N))
    for i = 1:basis.N
        for j = 1:basis.N
            K[i, j] = proj(basis.Λ, basis.grid[i], basis.grid[j])
        end
    end
    return K
end

"""
modified Gram-Schmidt process
"""
function mGramSchmidt(basis, idx, g::Float)
    qnew = zeros(Float, basis.N)
    qnew[idx] = 1

    for qi = 1:basis.N
        if qi == idx
            continue
        end
        q = basis.Q[qi, :]
        qnew -= projqq(basis, q, qnew) .* q  # <q, qnew> q
    end
    return qnew / sqrt(projqq(basis, qnew, qnew))
end

# """
# Gram-Schmidt process
# """
# function GramSchmidt(basis, idx, g::Float)
#     q0 = zeros(Float, basis.N)
#     q0[idx] = 1
#     qnew = copy(q0)

#     for qi in 1:basis.N
#         if qi == idx
#     continue
#     end
#         q = basis.Q[qi, :]
#         qnew -=  projqq(basis, q, q0) .* q
#     end

#     norm = sqrt(projqq(basis, qnew, qnew))
#     return qnew / norm
# end

function Residual(basis, proj, g::Float)
    # norm2 = proj(g, g) - \sum_i (<qi, K_g>)^2
    # qi=\sum_j Q_ij K_j ==> (<qi, K_g>)^2 = (\sum_j Q_ij <K_j, K_g>)^2 = \sum_jk Q_ij*Q_ik <K_j, K_g>*<K_k, Kg>

    KK = [proj(basis.Λ, gj, g) for gj in basis.grid]
    norm2 = proj(basis.Λ, g, g) - (norm(basis.Q * KK))^2

    # norm2 = proj(basis.Λ, g, g)
    # for j in 1:basis.N
    #     norm2 -= basisQ[j, :]
    # end
    return norm2 < 0 ? Float(0) : sqrt(norm2)
end

function testOrthgonal(basis)
    println("testing orthognalization...")
    II = basis.Q * basis.proj * basis.Q'
    maxerr = maximum(abs.(II - I))
    println("Max Orthognalization Error: ", maxerr)
end

"""
function build(dlrGrid, print::Bool = true)
    Construct discrete Lehmann representation

#Arguments:
- `dlrGrid`: struct that contains the information to construct the DLR grid. The following entries are required:
   Λ: the dimensionless scale β*Euv, rtol: the required relative accuracy, isFermi: fermionic or bosonic, symmetry: particle-hole symmetry/antisymmetry or none
- `print`: print the internal information or not
"""
function build(dlrGrid, print::Bool = true)
    print && println("Using the functional algorithm to build DLR ...")
    Λ = Float(dlrGrid.Λ)
    rtol = Float(dlrGrid.rtol)
    symmetry = dlrGrid.symmetry
    if symmetry == :ph
        print && println("Building ω grid ... ")
        ωBasis = QR(Λ, rtol, projPH_ω, [Float(0), Float(Λ)])
        ωGrid = ωBasis.grid
        rank = ωBasis.N
    elseif symmetry == :pha
        print && println("Building ω grid ... ")
        ωBasis = QR(Λ, rtol, projPHA_ω, [Float(Λ),])
        ωGrid = ωBasis.grid
        rank = ωBasis.N
    else
        error("Functional algorithm for the symmetry $symmetry has not yet been implemented!")
        # elseif type == :fermi
        #     println("Building ω grid ... ")
        #     ωBasis = QR(Λ, rtol, projPH_ω, [Float(0), Float(Λ)])
        #     ωGrid = vcat(-ωBasis.grid[end:-1:2], ωBasis.grid)
        #     rank = length(ωGrid)
        #     println("rank: $rank")
        #     println("Building τ grid ... ")
        #     τBasis = tauGrid(ωGrid, rank, Λ, rtol, :fermi)
        #     # τBasis = QR(Λ / 2, rtol / 10, projPH_τ, Float(0), N=ωBasis.N)
        #     println("Building n grid ... ")
        #     nBasis = MatFreqGrid(ωGrid, rank, Λ, :fermi)
    end

    ωGrid = Float64.(ωGrid)
    degree = 128
    τ = Discrete.τChebyGrid(dlrGrid, degree, print)
    kernel = Discrete.preciseKernelT(dlrGrid, τ, ωGrid, print)
    Discrete.testInterpolation(dlrGrid, τ, ωGrid, kernel, print)

    τIndex = Discrete.τnQR(kernel, rank, print)
    τGrid = sort(τ.grid[τIndex])

    nFineGrid, nFermiKernel, nBoseKernel = Discrete.preciseKernelΩn(dlrGrid, ωGrid, print)

    nFermiIndex = Discrete.τnQR(nFermiKernel, rank, print)
    nFermiGrid = sort(nFineGrid[nFermiIndex])

    nBoseIndex = Discrete.τnQR(nBoseKernel, rank, print)
    nBoseGrid = sort(nFineGrid[nBoseIndex])

    # τGrid = τBasis / Λ
    # τGrid = τBasis
    # nGrid = nBasis
    ########### output  ############################
    # @printf("%5s  %32s  %32s  %8s\n", "index", "real freq", "tau", "ωn")
    # for r = 1:rank
    #     @printf("%5i  %32.17g  %32.17g  %16i\n", r, ωGrid[r], τGrid[r], nGrid[r])
    # end
    ########### output  ############################
    print && @printf("%5s  %32s  %32s  %11s  %11s\n", "index", "real freq", "tau", "fermi ωn", "bose ωn")
    for r = 1:rank
        print && @printf("%5i  %32.17g  %32.17g  %16i %16i\n", r, ωGrid[r], τGrid[r], nFermiGrid[r], nBoseGrid[r])
    end

    return ωGrid, τGrid, nFermiGrid, nBoseGrid
    # return Dict([(:ω, ωGrid), (:τ, τGrid), (:ωn, nGrid)])
end

end

if abspath(PROGRAM_FILE) == @__FILE__
    # freq, Q = findBasis(1.0e-3, Float(100))
    # basis = QR(100, 1e-3)
    Λ = 1e10
    # Λ = 100
    # @time ωBasis = QR(Λ, 1e-13, projPH_ω, [Float(0), Float(Λ)])
    @time ωBasis = Functional.QR(Λ, 1e-12, projPHA_ω, [Float(Λ),])
    # @time τBasis = QR(Λ / 2, 1e-11, projPHA_τ, Float(0), N=ωBasis.N)
    # nBasis = MatFreqGrid(ωBasis.grid, ωBasis.N, Λ, :acorr)

    # @time basis = QR(100, 1e-10)
    # readline()
    # basis = QR(100, 1e-3)

end


@testset "Q" begin
    @testset "q2m" begin
        exp = [0.607843137254902 0.27450980392156854 0.7450980392156862;
               0.6666666666666666 0.33333333333333326 -0.6666666666666666;
               -0.43137254901960775 0.9019607843137255 0.019607843137254832]
        q = [0.5, 0.4, 0.3, 0.1]
        act = q2m(q)
        @testset for i in eachindex(act, exp)
            @test act[i] ≈ exp[i]
        end
    end
    @testset "qdq2av" begin
        angle = deg2rad.([-20.0, 50.0, -60.0])
        m = eul2m(angle[3], angle[2], angle[1], 3, 1, 3)
        q = m2q(m)
        exp = [1.0, 2.0, 3.0]
        qav = [0.0, 1.0, 2.0, 3.0]
        dq = -0.5 * qxq(q, qav)
        act = qdq2av(q, dq)
        @testset for i in eachindex(act, exp)
            @test act[i] ≈ exp[i]
        end
    end
    @testset "qxq" begin
        q_id = [1.0, 0.0, 0.0, 0.0]
        nq_id = [-1.0, 0.0, 0.0, 0.0]
        q_i = [0.0, 1.0, 0.0, 0.0]
        q_j = [0.0, 0.0, 1.0, 0.0]
        q_k = [0.0, 0.0, 0.0, 1.0]
        q_ij = qxq(q_i, q_j)
        @testset for i in eachindex(q_ij, q_k)
            @test q_ij[i] ≈ q_k[i]
        end
        q_jk = qxq(q_j, q_k)
        @testset for i in eachindex(q_jk, q_i)
            @test q_jk[i] ≈ q_i[i]
        end
        q_ki = qxq(q_k, q_i)
        @testset for i in eachindex(q_ki, q_j)
            @test q_ki[i] ≈ q_j[i]
        end
        q_ii = qxq(q_i, q_i)
        @testset for i in eachindex(q_ii, nq_id)
            @test q_ii[i] ≈ nq_id[i]
        end
        q_jj = qxq(q_j, q_j)
        @testset for i in eachindex(q_jj, nq_id)
            @test q_jj[i] ≈ nq_id[i]
        end
        q_kk = qxq(q_k, q_k)
        @testset for i in eachindex(q_kk, nq_id)
            @test q_kk[i] ≈ nq_id[i]
        end
        q_idi = qxq(q_id, q_i)
        @testset for i in eachindex(q_idi, q_i)
            @test q_idi[i] ≈ q_i[i]
        end
        q_iid = qxq(q_i, q_id)
        @testset for i in eachindex(q_iid, q_i)
            @test q_iid[i] ≈ q_i[i]
        end
    end
end


#!/usr/bin/env julia

export pe3

# need Primes.jl
# https://github.com/JuliaMath/Primes.jl
using Primes

pe3(n::Integer=600851475143)=factor(Vector,n)|>maximum

if !haskey(ENV,"PROJECT_EULER_WITHOUT_RESULT")
    pe3()|>print
end


#=

  Scheduling speakers in Julia ConstraintSolver.jl

  From Rina Dechter, Constraint Processing, page 72
  Scheduling of 6 speakers in 6 slots.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function scheduling_speakers(available_slots,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = length(available_slots)
    @variable(model, 1 <= xs[1:n] <= n, Int)

    @constraint(model, xs in CS.AllDifferent())

    for (x,a) in zip(xs,available_slots)
        len = length(a)
        b = @variable(model,[1:len],Bin)
        for i in 1:len
            @constraint(model,b[i] := {a[i] == x})
        end
        @constraint(model,sum(b) >= 1)

        # This version is too slow for showin all (2) solutions:
        # ix = @variable(model,[1:len],CS.Integers(1:len))
        # my_element(model,ix[1],a,x)
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                xs_val = convert.(Integer,JuMP.value.(xs; result=sol))
                println("xs:$xs_val")
            end
        end
    else
        println("status:$status")
    end

    return status
end

available_slots = [
                       # Reasoning:
         [3,4,5,6],    # 2) the only one with 6 after speaker F -> 1
         [3,4],        # 5) 3 or 4
         [2,3,4,5],    # 3) only with 5 after F -> 1 and A -> 6
         [2,3,4],      # 4) only with 2 after C -> 5 and F -> 1
         [3,4],        # 5) 3 or 4
         [1,2,3,4,5,6] # 1) the only with 1

]
@time scheduling_speakers(available_slots,true,true)


struct BudgetedBipartiteMatchingInstance{T, U} <: CombinatorialInstance
    matching::BipartiteMatchingInstance{T}
    weight::Dict{Edge{T}, U}
    budget::U

    function BudgetedBipartiteMatchingInstance(
        graph::AbstractGraph{T},
        reward::Dict{Edge{T}, Float64},
        weight::Dict{Edge{T}, U},
        budget::U,
    ) where {T, U}
        matching = BipartiteMatchingInstance(graph, reward)
        return new{T, U}(matching, weight, budget)
    end
end

function _budgeted_bipartite_matching_compute_value(
    i::BudgetedBipartiteMatchingInstance{T, U},
    solution::Vector{Edge{T}},
) where {T, U}
    return _budgeted_bipartite_matching_compute_value(i.matching, solution)
end

function _budgeted_bipartite_matching_compute_value(
    i::BipartiteMatchingInstance{T},
    solution::Vector{Edge{T}},
) where {T}
    return sum(
        i.reward[(e in keys(i.reward)) ? e : reverse(e)] for e in solution
    )
end

function _budgeted_bipartite_matching_compute_weight(
    i::BudgetedBipartiteMatchingInstance{T, U},
    solution::Vector{Edge{T}},
) where {T, U}
    return sum(
        i.weight[(e in keys(i.weight)) ? e : reverse(e)] for e in solution
    )
end

abstract type BudgetedBipartiteMatchingSolution{T, U} <: CombinatorialSolution
    # instance::BudgetedBipartiteMatchingInstance{T, U}
    # solution::Vector{Edge{T}}
    # value::Float64 # TODO: remove me, only useful for Lagrangian.
end


using Oceananigans
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Oceananigans.Utils: prettysummary
using Printf
using GLMakie

κ = ν = 1
Nz = 32
Lz = 1

underlying_grid = RectilinearGrid(size=Nz, z=(0, Lz), topology = (Flat, Flat, Bounded))

@inline bottom_height(x, y) = 0.1
grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))

c_immersed_bc = ValueBoundaryCondition(1)
c_top_bc = ValueBoundaryCondition(-1)
c_bcs = FieldBoundaryConditions(immersed=c_immersed_bc, top=c_top_bc)

u_drag_func(i, j, k, grid, clock, model_fields) = - 1e-3 * model_fields.u[i, j, k]^2
u_immersed_bc = ValueBoundaryCondition(-1)
u_top_bc = ValueBoundaryCondition(1)
u_bcs = FieldBoundaryConditions(immersed=u_immersed_bc, top=u_top_bc)

model = NonhydrostaticModel(; grid,
                            advection = nothing,
                            timestepper = :RungeKutta3,
                            tracers = :c,
                            closure = VerticalScalarDiffusivity(; ν, κ),
                            boundary_conditions = (u=u_bcs, c=c_bcs))

#set!(model, u = 0(x, y, z) -> z - 0.1)

simulation = Simulation(model, Δt=1e-4, stop_iteration=1000)

outputs = merge(model.velocities, model.tracers)
simulation.output_writers[:jld2] = JLD2OutputWriter(model, outputs,
                                                    filename = "immersed_couette_flow.jld2",
                                                    schedule = IterationInterval(10),
                                                    overwrite_existing = true)

progress(sim) = @info string("Iteration: ", iteration(sim), ", time: ", time(sim))
simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))

run!(simulation)

ct = FieldTimeSeries("immersed_couette_flow.jld2", "c")
ut = FieldTimeSeries("immersed_couette_flow.jld2", "u")
Nt = length(ct.times)

z = znodes(ut)

fig = Figure()
axc = Axis(fig[2, 1], xlabel="c", ylabel="z")
axu = Axis(fig[2, 2], xlabel="u", ylabel="z")
slider = Slider(fig[3, 1:2], range=1:Nt, startvalue=1)
n = slider.value

title = @lift string("Immersed Couette flow at t = ", prettysummary(ct.times[$n]))
Label(fig[1, 1:2], title)

cn = @lift interior(ct[$n], 1, 1, :)
un = @lift interior(ut[$n], 1, 1, :)

lines!(axc, cn, z)
lines!(axu, un, z)

display(fig)

record(fig, "immersed_couette_flow.mp4", 1:Nt, framerate=24) do nn
    n[] = nn
end


# fft_convolve.jl

using BenchmarkTools: @btime
using SPECTrecon: plan_psf
using SPECTrecon: fft_conv!, fft_conv_adj!

function fft_conv_time()
    M = 200
    N = 64
    T = Float32
    img = zeros(T, M, N)
    img[20:150, 20:40] .= rand(131,21)
    output = similar(img)
    nx_psf = 5
    ker = ones(T, nx_psf, nx_psf) / nx_psf^2
    plan = plan_psf(M, N, nx_psf; nthread = 1, T = T)[1]

    println("fft_conv")
    @btime fft_conv!($output, $img, $ker, $plan)
    # 565.631 μs (0 allocations: 0 bytes)
    println("fft_conv_adj")
    @btime fft_conv_adj!($output, $img, $ker, $plan)
    # 561.371 μs (0 allocations: 0 bytes)
    nothing
end

# run all functions, time may vary on different machines, but should be all zero allocation.
fft_conv_time()


@testset "MulticlassLDA" begin
    ## Data
    Xfull, y = @load_smarket
    X = selectcols(Xfull, [:Lag1,:Lag2])
    train = selectcols(Xfull, :Year) .< Dates.Date(2005)
    test = .!train
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest = selectrows(X, test)
    ytest = selectrows(y, test)

    lda_model = LDA()

    ## Check model `fit`
    fitresult, = fit(lda_model, 1, Xtrain, ytrain)
    class_means, projection_matrix = fitted_params(lda_model, fitresult)
    @test round.(class_means', sigdigits = 3) == [0.0428 0.0339; -0.0395 -0.0313]
    ## Check model `predict`
    preds = predict(lda_model, fitresult, Xtest)
    mce = cross_entropy(preds, ytest) |> mean
    @test 0.685 ≤ mce ≤ 0.695
    ## Check model `transform`
    # MultivariateStats Linear Discriminant Analysis transform
    proj = fitresult[1].proj
    XWt = matrix(X) * proj
    tlda_ms = table(XWt, prototype=X)
    # MLJ Linear Discriminant Analysis transform
    tlda_mlj = transform(lda_model, fitresult, X)
    @test tlda_mlj == tlda_ms
    ## Check model traits
    d = info_dict(LDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "LDA"
end

@testset "MLDA-2" begin
    ## Data
    Random.seed!(1125)
    X1 = -2 .+ randn(100, 2)
    X2 = randn(100, 2)
    X3 = 2 .+ randn(100, 2)
    y1 = ones(100)
    y2 = 2ones(100)
    y3 = 3ones(100)
    X = vcat(X1, X2, X3)
    y = vcat(y1, y2, y3)
    p = Random.randperm(300)
    X = X[p, :]
    y = y[p]
    X = table(X)
    y = categorical(y)
    train, test = partition(eachindex(y), 0.7)
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest = selectrows(X, test)
    ytest = selectrows(y, test)

    lda_model = LDA()
    ## Check model `fit`/`predict`
    fitresult, = fit(lda_model, 1, Xtrain, ytrain)
    preds = predict_mode(lda_model, fitresult, Xtest)
    mcr = misclassification_rate(preds, ytest)
    @test mcr ≤ 0.15
end

@testset "BayesianMulticlassLDA" begin
    ## Data
    Xfull, y = @load_smarket
    X = selectcols(Xfull, [:Lag1,:Lag2])
    train = selectcols(Xfull, :Year) .< Dates.Date(2005)
    test = .!train
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest = selectrows(X, test)
    ytest = selectrows(y, test)

    BLDA_model = BayesianLDA()
    ## Check model `fit`
    fitresult, = fit(BLDA_model, 1, Xtrain, ytrain)
    class_means, projection_matrix, priors = fitted_params(BLDA_model, fitresult)
    @test round.(class_means', sigdigits = 3) == [0.0428 0.0339; -0.0395 -0.0313]
    ## Check model `predict`
    preds = predict(BLDA_model, fitresult, Xtest)
    mce = cross_entropy(preds, ytest) |> mean
    @test 0.685 ≤ mce ≤ 0.695
    ## Check model traits
    d = info_dict(BayesianLDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "BayesianLDA"
end

@testset "BayesianSubspaceLDA" begin
    ## Data
    X, y = @load_iris
    LDA_model = BayesianSubspaceLDA()
    ## Check model `fit`
    fitresult, _, report = fit(LDA_model, 1, X, y)
    class_means, projection_matrix, prior_probabilities = fitted_params(
        LDA_model, fitresult
    )
    @test mean(
        abs.(
            class_means' - [
                5.006 3.428 1.462 0.246;
                5.936 2.770 4.260 1.326;
                6.588 2.974 5.552 2.026
            ]
        )
    ) < 0.01
    @test mean(
        abs.(
            projection_matrix ≈ [
                0.0675721   0.00271023;
  			    0.127666    0.177718;
 				-0.180211   -0.0767255;
                -0.235382    0.231435
            ]
        )
    ) < 0.05
    @test round.(prior_probabilities, sigdigits=7) == [0.3333333, 0.3333333, 0.3333333]
    @test round.(report.explained_variance_ratio, digits=4) == [0.9915, 0.0085]

    ## Check model `predict`
    preds=predict(LDA_model, fitresult, X)
    predicted_class = predict_mode(LDA_model, fitresult, X)
    mcr = misclassification_rate(predicted_class, y)
    mce = cross_entropy(preds, y) |> mean
    @test round.(mcr, sigdigits=1) == 0.02
    @test 0.04 ≤ mce ≤ 0.045
    ## Check model traits
    d = info_dict(BayesianSubspaceLDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "BayesianSubspaceLDA"
end

@testset "SubspaceLDA" begin
    ## Data
    Random.seed!(1125)
    X1 = -2 .+ randn(100, 2)
    X2 = randn(100, 2)
    X3 = 2 .+ randn(100, 2)
    y1 = ones(100)
    y2 = 2ones(100)
    y3 = 3ones(100)
    X = vcat(X1, X2, X3)
    y = vcat(y1, y2, y3)
    p = Random.randperm(300)
    X = X[p, :]
    y = y[p]
    X = table(X)
    y = categorical(y)
    train, test = partition(eachindex(y), 0.7)
    Xtrain = selectrows(X, train)
    ytrain = selectrows(y, train)
    Xtest = selectrows(X, test)
    ytest = selectrows(y, test)

    lda_model = SubspaceLDA()
    ## Check model `fit`/ `transform`
    fitresult, = fit(lda_model, 1, Xtrain, ytrain)
    preds = predict_mode(lda_model, fitresult, Xtest)
    mcr = misclassification_rate(preds, ytest)
    @test mcr ≤ 0.15
    # MultivariateStats Linear Discriminant Analysis transform
    proj = fitresult[1].projw
    projLDA = fitresult[1].projLDA
    proj = proj * projLDA
    XWt = matrix(X) * proj
    tlda_ms = table(XWt, prototype=X)
    # MLJ Linear Discriminant Analysis transform
    tlda_mlj = transform(lda_model, fitresult, X)
    @test tlda_mlj == tlda_ms
    ## Check model traits
    d = info_dict(SubspaceLDA)
    @test d[:input_scitype] == Table(Continuous)
    @test d[:target_scitype] == AbstractVector{<:Finite}
    @test d[:name] == "SubspaceLDA"
end

@testset "discriminant models checks" begin
## Data to be used for tests
y = categorical(["apples", "oranges", "carrots", "mango"])
X = (x1 =rand(4), x2 = collect(1:4))

## Note: The following test depend on the order in which they are written.
## Hence do not change the ordering of the tests.

## Check to make sure error is thrown if we only have a single
## unique class during training.
model = LDA()
# categorical array with same pool as y but only containing "apples"
y1 = y[[1,1,1,1]]
@test_throws ArgumentError fit(model, 1, X, y1)

## Check to make sure error is thrown if we don't have more samples
## than unique classes during training.
@test_throws ArgumentError fit(model, 1, X, y)

## Check to make sure error is thrown if `out_dim` exceeds the number of features in
## sample matrix used in training.
model = LDA(out_dim=3)
# categorical array with same pool as y but only containing "apples" & "oranges"
y2 = y[[1,2,1,2]]
@test_throws ArgumentError fit(model, 1, X, y2)

## Check to make sure error is thrown if length(`priors`) !=  number of classes
## in common pool of target vector used in training.
model = BayesianLDA(priors=[0.1, 0.5, 0.4])
@test_throws ArgumentError fit(model, 1, X, y)

## Check to make sure error is thrown if sum(`priors`) isn't approximately equal to 1.
model = BayesianLDA(priors=[0.1, 0.5, 0.4, 0.2])
@test_throws ArgumentError fit(model, 1, X, y)

## Check to make sure error is thrown if `priors .< 0` or `priors .> 1`.
model = BayesianLDA(priors=[-0.1, 0.0, 1.0, 0.1])
@test_throws ArgumentError fit(model, 1, X, y)
model = BayesianLDA(priors=[1.1, 0.0, 0.0, -0.1])
@test_throws ArgumentError fit(model, 1, X, y)
end


@testset "Vector spaces" begin
    @testset "VectorSpace" begin
        𝒫 = ParameterSpace()
        𝒯 = Taylor(0)
        # different types of vector spaces cannot be equal
        @test !(𝒫 == 𝒯)
        # different types of vector spaces cannot be subset of one another
        @test !(𝒫 ⊆ 𝒯)
        # intersection and union between different types of vector spaces is not allowed
        @test_throws MethodError intersect(𝒫, 𝒯)
        @test_throws MethodError union(𝒫, 𝒯)
    end

    @testset "ParameterSpace" begin
        𝒫 = ParameterSpace()
        @test 𝒫 ⊆ 𝒫
        @test 𝒫 ∩ 𝒫 == 𝒫 ∪ 𝒫 == 𝒫
        @test dimension(𝒫) == 1
        @test RadiiPolynomial._firstindex(𝒫) == 1
        @test RadiiPolynomial._lastindex(𝒫) == 1
        @test indices(𝒫) == Base.OneTo(1)
        @test RadiiPolynomial._findposition(1, 𝒫) == 1
    end

    @testset "Taylor" begin
        𝒯₂, 𝒯₃ = Taylor(2), Taylor(3)
        @test !(𝒯₂ == 𝒯₃)
        @test 𝒯₂ ⊆ 𝒯₃
        @test 𝒯₂ ∩ 𝒯₃ == 𝒯₂
        @test 𝒯₂ ∪ 𝒯₃ == 𝒯₃
        @test order(𝒯₂) == 2
        @test dimension(𝒯₂) == 3
        @test RadiiPolynomial._firstindex(𝒯₂) == 0
        @test RadiiPolynomial._lastindex(𝒯₂) == 2
        @test indices(𝒯₂) == 0:2
        @test RadiiPolynomial._findindex_constant(𝒯₂) == 0
        @test RadiiPolynomial._findposition(1, 𝒯₂) == 2
        @test RadiiPolynomial._findposition(0:2, 𝒯₂) == 1:3
        @test RadiiPolynomial._findposition([0, 2], 𝒯₂) == [1, 3]
        @test RadiiPolynomial._findposition(:, 𝒯₂) == Colon()
    end

    @testset "Fourier" begin
        ℱ₂, ℱ₃ = Fourier(2, 1.0), Fourier(3, 1.0)
        @test !(ℱ₂ == ℱ₃)
        @test ℱ₂ ⊆ ℱ₃
        @test ℱ₂ ∩ ℱ₃ == ℱ₂
        @test ℱ₂ ∪ ℱ₃ == ℱ₃
        @test order(ℱ₂) == 2
        @test frequency(ℱ₂) == 1.0
        @test dimension(ℱ₂) == 5
        @test RadiiPolynomial._firstindex(ℱ₂) == -2
        @test RadiiPolynomial._lastindex(ℱ₂) == 2
        @test indices(ℱ₂) == -2:2
        @test RadiiPolynomial._findindex_constant(ℱ₂) == 0
        @test RadiiPolynomial._findposition(-1, ℱ₂) == 2
        @test RadiiPolynomial._findposition(0:1, ℱ₂) == 3:4
        @test RadiiPolynomial._findposition([-2, 2], ℱ₂) == [1, 5]
        @test RadiiPolynomial._findposition(:, ℱ₂) == Colon()
        @test convert(Fourier{Float64}, ℱ₂) == ℱ₂
        @test convert(Fourier{Int}, ℱ₂) == Fourier(2, 1)
        @test promote_type(Fourier{Float64}, Fourier{Float64}) == Fourier{Float64}
        @test promote_type(Fourier{Float64}, Fourier{Int}) == Fourier{Float64}
    end

    @testset "Chebyshev" begin
        𝒞₂, 𝒞₃ = Chebyshev(2), Chebyshev(3)
        @test !(𝒞₂ == 𝒞₃)
        @test 𝒞₂ ⊆ 𝒞₃
        @test 𝒞₂ ∩ 𝒞₃ == 𝒞₂
        @test 𝒞₂ ∪ 𝒞₃ == 𝒞₃
        @test order(𝒞₂) == 2
        @test dimension(𝒞₂) == 3
        @test RadiiPolynomial._firstindex(𝒞₂) == 0
        @test RadiiPolynomial._lastindex(𝒞₂) == 2
        @test indices(𝒞₂) == 0:2
        @test RadiiPolynomial._findindex_constant(𝒞₂) == 0
        @test RadiiPolynomial._findposition(1, 𝒞₂) == 2
        @test RadiiPolynomial._findposition(0:2, 𝒞₂) == 1:3
        @test RadiiPolynomial._findposition([0, 2], 𝒞₂) == [1, 3]
        @test RadiiPolynomial._findposition(:, 𝒞₂) == Colon()
    end

    @testset "TensorSpace" begin
        𝒯₂, ℱ₂, 𝒞₂ = Taylor(2), Fourier(2, 1.0), Chebyshev(2)
        𝑇 = TensorSpace((𝒯₂, ℱ₂, 𝒞₂))
        @test spaces(𝑇) == (𝒯₂, ℱ₂, 𝒞₂)
        @test 𝒯₂ ⊗ ℱ₂ ⊗ 𝒞₂ == (𝒯₂ ⊗ ℱ₂) ⊗ 𝒞₂ == 𝒯₂ ⊗ (ℱ₂ ⊗ 𝒞₂) == 𝑇
        @test 𝑇 ⊗ 𝑇 == TensorSpace((𝒯₂, ℱ₂, 𝒞₂, 𝒯₂, ℱ₂, 𝒞₂))
        @test 𝑇[1] == 𝒯₂
        @test 𝑇[1:3] == 𝑇[[1, 2, 3]] == 𝑇[:] == 𝑇
        @test Base.tail(𝑇) == TensorSpace((ℱ₂, 𝒞₂))
        @test Base.front(𝑇) == TensorSpace((𝒯₂, ℱ₂))
        @test 𝑇 ⊆ Taylor(3) ⊗ ℱ₂ ⊗ 𝒞₂
        @test 𝑇 ∩ (Taylor(3) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(3)) == 𝒯₂ ⊗ Fourier(1, 1.0) ⊗ 𝒞₂
        @test 𝑇 ∪ (Taylor(3) ⊗ Fourier(1, 1.0) ⊗ Chebyshev(3)) == Taylor(3) ⊗ ℱ₂ ⊗ Chebyshev(3)
        @test dimension(𝑇) == 45
        @test dimension(𝑇, 1) == 3
        @test dimensions(𝑇) == (3, 5, 3)
        @test RadiiPolynomial._firstindex(𝑇) == (0, -2, 0)
        @test RadiiPolynomial._lastindex(𝑇) == (2, 2, 2)
        @test indices(𝑇) == Base.Iterators.product(0:2, -2:2, 0:2)
        @test RadiiPolynomial._findindex_constant(𝑇) == (0, 0, 0)
        @test RadiiPolynomial._findposition((1, 2, 0), 𝑇) == 14
        @test RadiiPolynomial._findposition((:, -2:2, :), 𝑇) == collect(1:45)
        @test RadiiPolynomial._findposition([(1, 2, 0), (2, 2, 2)], 𝑇) == [14, 45]
        @test RadiiPolynomial._findposition((:, :, :), 𝑇) == Colon()
        @test RadiiPolynomial._findposition(:, 𝑇) == Colon()
        @test order(𝑇) == (2, 2, 2)
        @test order(𝑇, 1) == 2
        @test_throws MethodError frequency(𝑇)
        @test frequency(𝑇, 2) == 1.0
        @test convert(TensorSpace{Tuple{Taylor,Fourier{Float64},Chebyshev}}, 𝑇) == 𝑇
        @test convert(TensorSpace{Tuple{Taylor,Fourier{Int},Chebyshev}}, 𝑇) == Taylor(2) ⊗ Fourier(2, 1) ⊗ Chebyshev(2)
    end

    @testset "CartesianSpace" begin
        𝒯₁¹ = CartesianPower(Taylor(1), 1)
        𝒞₂² = CartesianPower(Chebyshev(2), 2)
        ℱ₃³ = CartesianPower(Fourier(3, 1.0), 3)
        𝒯₁¹𝒞₂²ℱ₃³ = CartesianProduct((𝒯₁¹, 𝒞₂², ℱ₃³))
        𝒯₂¹𝒞₃²ℱ₃³ = CartesianProduct((CartesianPower(Taylor(2), 1), CartesianPower(Chebyshev(3), 2), ℱ₃³))
        @test space(𝒯₁¹) == Taylor(1)
        @test spaces(𝒯₁¹) == [Taylor(1)]
        @test spaces(𝒯₁¹𝒞₂²ℱ₃³) == (𝒯₁¹, 𝒞₂², ℱ₃³)
        @test nb_cartesian_product(ℱ₃³) == nb_cartesian_product(𝒯₁¹𝒞₂²ℱ₃³) == 3
        @test 𝒯₁¹^1 == CartesianPower(𝒯₁¹, 1)
        @test 𝒯₁¹ × 𝒞₂² × ℱ₃³ == 𝒯₁¹ × (𝒞₂² × ℱ₃³) == (𝒯₁¹ × 𝒞₂²) × ℱ₃³ == (𝒯₁¹ × 𝒞₂²) × CartesianProduct(tuple(ℱ₃³)) == 𝒯₁¹𝒞₂²ℱ₃³
        @test Base.front(𝒯₁¹𝒞₂²ℱ₃³) == 𝒯₁¹ × 𝒞₂²
        @test Base.tail(𝒯₁¹𝒞₂²ℱ₃³) == 𝒞₂² × ℱ₃³
        @test 𝒞₂²[1] == Chebyshev(2)
        @test 𝒞₂²[1:2] == 𝒞₂²[[1, 2]] == 𝒞₂²[:] == 𝒞₂²
        @test 𝒯₁¹𝒞₂²ℱ₃³[1] == 𝒯₁¹
        @test 𝒯₁¹𝒞₂²ℱ₃³[1:3] == 𝒯₁¹𝒞₂²ℱ₃³[[1, 2, 3]] == 𝒯₁¹𝒞₂²ℱ₃³[:] == 𝒯₁¹𝒞₂²ℱ₃³
        @test 𝒯₁¹𝒞₂²ℱ₃³ ⊆ 𝒯₂¹𝒞₃²ℱ₃³
        @test 𝒯₁¹𝒞₂²ℱ₃³ ∩ 𝒯₂¹𝒞₃²ℱ₃³ == 𝒯₁¹𝒞₂²ℱ₃³
        @test 𝒯₁¹𝒞₂²ℱ₃³ ∪ 𝒯₂¹𝒞₃²ℱ₃³ == 𝒯₂¹𝒞₃²ℱ₃³
        @test dimension(𝒯₁¹𝒞₂²ℱ₃³) == 29
        @test dimension(𝒯₁¹𝒞₂²ℱ₃³, 1) == dimension(𝒯₁¹, 1) == 2
        @test dimensions(𝒯₁¹) == [2]
        @test dimensions(𝒯₁¹𝒞₂²ℱ₃³) == (2, 6, 21)
        @test RadiiPolynomial._firstindex(𝒯₁¹𝒞₂²ℱ₃³) == RadiiPolynomial._firstindex(𝒞₂²) == 1
        @test RadiiPolynomial._lastindex(𝒞₂²) == 6
        @test RadiiPolynomial._lastindex(𝒯₁¹𝒞₂²ℱ₃³) == 29
        @test indices(𝒞₂²) == Base.OneTo(6)
        @test indices(𝒯₁¹𝒞₂²ℱ₃³) == Base.OneTo(29)
        @test order(𝒯₁¹𝒞₂²ℱ₃³) == ([1], [2, 2], [3, 3, 3])
        @test order(𝒯₁¹𝒞₂²ℱ₃³, 1) == [1]
        @test order(𝒯₁¹, 1) == 1
        @test_throws MethodError frequency(𝒯₁¹𝒞₂²ℱ₃³)
        @test frequency(𝒯₁¹𝒞₂²ℱ₃³, 3) == frequency(ℱ₃³) == [1.0, 1.0, 1.0]
        @test frequency(ℱ₃³, 1) == 1.0
        @test RadiiPolynomial._component_findposition(1, 𝒯₁¹) == RadiiPolynomial._component_findposition(1:1, 𝒯₁¹) == 1:2
        @test RadiiPolynomial._component_findposition(2, 𝒞₂²) == RadiiPolynomial._component_findposition(2:2, 𝒞₂²) == 4:6
        @test RadiiPolynomial._component_findposition(1, 𝒯₁¹𝒞₂²ℱ₃³) == RadiiPolynomial._component_findposition(1:1, 𝒯₁¹𝒞₂²ℱ₃³) == 1:2
        @test RadiiPolynomial._component_findposition(2, 𝒯₁¹𝒞₂²ℱ₃³) == RadiiPolynomial._component_findposition(2:2, 𝒯₁¹𝒞₂²ℱ₃³) == 3:8
        @test convert(CartesianProduct{Tuple{CartesianPower{Taylor}, CartesianPower{Chebyshev}, CartesianPower{Fourier{Float64}}}}, 𝒯₁¹𝒞₂²ℱ₃³) == 𝒯₁¹𝒞₂²ℱ₃³
        @test convert(CartesianProduct{Tuple{CartesianPower{Taylor}, CartesianPower{Chebyshev}, CartesianPower{Fourier{Int}}}}, 𝒯₁¹𝒞₂²ℱ₃³) == 𝒯₁¹ × 𝒞₂² × Fourier(3, 1)^3
    end
end


using Test
using ForwardDiff
using IndirectTrajOpt

# Initialize parameters
ps = initCR3BPIndirectWithSTMParams("Low Thrust 10 CR3BP")

# Initial condition vector
y0 = [-0.019488511458668, -0.016033479812051, 0.0,
       8.918881923678198, -4.081793688818725, 0.0, 1.0,
       15.616017, 32.875896, -0.094522,
       -0.101606, 0.044791, -0.000150, 0.133266]

# Initial condition vector with STM
z0 = [y0; zeros(196)]
for i in 1:14
    z0[14 + i + (i-1)*14] = 1.0
end

ps.ϵ = 1.0
for utype in [0,1,2]
    ps.utype = utype
    # Evaluate eoms with Jacobian
    dz0 = zeros(210)
    IndirectTrajOpt.cr3bpEomIndirectWithSTM!(dz0, z0, ps, 0.0)

    # Evaluate Jacobian of eoms with ForwardDiff
    dy0 = zeros(14)
    jac = ForwardDiff.jacobian(
        (y,x)->IndirectTrajOpt.cr3bpEomIndirect!(y,x,ps,0.0),
        dy0, y0)
    ps.utype
    # Evaluate jacobian difference
    jacDiff = abs.(reshape(dz0[15:end], (14,14)) .- jac)

    # Tests
    for i in 1:196
        @test jacDiff[i] < 1e-9
    end
end

ps.ϵ = 0.0
for utype in [0,2]
    ps.utype = utype
    # Evaluate eoms with Jacobian
    dz0 = zeros(210)
    IndirectTrajOpt.cr3bpEomIndirectWithSTM!(dz0, z0, ps, 0.0)

    # Evaluate Jacobian of eoms with ForwardDiff
    dy0 = zeros(14)
    jac = ForwardDiff.jacobian(
        (y,x)->IndirectTrajOpt.cr3bpEomIndirect!(y,x,ps,0.0),
        dy0, y0)
    ps.utype
    # Evaluate jacobian difference
    jacDiff = abs.(reshape(dz0[15:end], (14,14)) .- jac)

    # Tests
    for i in 1:196
        @test jacDiff[i] < 1e-9
    end
end

# Define function for test
function cr3bpAnalyticJacTestFunc(x, tspan, ps)
    out, extra = cr3bpOptIntegrate(x, tspan, ps)
    return out
end

# Integrate state and co-states with STM
for ϵ in [1.0, 0.5, 0.0]
    ps.ϵ = ϵ
    tspan = (0.0, 8.6404*24*3600/ps.crp.TU)
    zf = cr3bpOptWithSTMIntegrate(z0, tspan, ps)

    # Compute STM with ForwardDiff
    stm = ForwardDiff.jacobian(
        (x)->cr3bpAnalyticJacTestFunc(x, tspan, ps), y0)

    # Evaluate STM difference
    jacDiff = abs.(reshape(zf[15:end], (14,14)) .- stm)
    display(jacDiff)

    # Tests
    for i in 1:196
        @test jacDiff[i] < 1e-2
    end
end


q = Quaternion(1, 0, 0, 0)
r = σmap(q)

@test isapprox(r, ℝ³(1, 0, 0))


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

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

# ╔═╡ 489b13da-7d23-11eb-39de-57a12469721f
using Plots, LaTeXStrings

# ╔═╡ 82c5386a-7d2f-11eb-3992-e985ee16933c
using ColorSchemes

# ╔═╡ 4bd45848-7d23-11eb-1a8b-b36f7461f368
begin
	import WGLMakie
	using JSServe, PlutoUI
end

# ╔═╡ fa6e4914-7d24-11eb-0ff2-1974e4ec8df2
using StatsBase, Statistics, LinearAlgebra

# ╔═╡ 2ef712f2-7d2e-11eb-228c-2d023dc6ab10
using Random

# ╔═╡ a1b684a0-7d40-11eb-2a4c-e990f12c6cbb
using Distributions

# ╔═╡ 4256e210-7d23-11eb-2f7a-1d8aceab0682
function ingredients(path::String)
	# this is from the Julia source code (evalfile in base/loading.jl)
	# but with the modification that it returns the module instead of the last object
	name = Symbol(basename(path))
	m = Module(name)
	Core.eval(m,
        Expr(:toplevel,
             :(eval(x) = $(Expr(:core, :eval))($name, x)),
             :(include(x) = $(Expr(:top, :include))($name, x)),
             :(include(mapexpr::Function, x) = $(Expr(:top, :include))(mapexpr, $name, x)),
             :(include($path))))
	m
end

# ╔═╡ 3d2bf622-7d23-11eb-33fb-3f8a4f0f309e
function makie(f::Function)
    scene = WGLMakie.Scene(resolution = (600, 400))
	f(scene)

    return scene
end

# ╔═╡ 5e1fa2de-7d23-11eb-18da-ddf0ec5714cc
M = ingredients("../src/main.jl")

# ╔═╡ a0c3a118-7d3b-11eb-0309-a9540887541f
I = ingredients("../src/infer.jl")

# ╔═╡ fce24760-7d22-11eb-2474-a5efdd176182
scrna, genes = M.SeqSpace.expression();

# ╔═╡ 330fd846-7d25-11eb-3933-b5f8871eedc5
size(scrna)

# ╔═╡ 4ddffd5a-7d24-11eb-2c70-675569787f53
"""
This section generates figures for linear dimensional reduction
"""

# ╔═╡ 72a6c262-7d25-11eb-1a05-756c421933f5
pvals(matrix) = (λ=reverse(eigvals(cov(matrix,dims=2)))[1:1297])/sum(λ)

# ╔═╡ 85650856-7d24-11eb-055d-75a99e549182
function sample_from!(null, matrix)
	@assert size(null) == size(matrix)
	for g ∈ 1:size(null,1)
		sample!(matrix[g,:], view(null,g,:); replace=false)
	end
end

# ╔═╡ c3fdc5ce-7d26-11eb-235a-33e432adcd91
function null_pvals(matrix; niter=10)
	null = zeros(size(matrix))
	λ 	 = zeros(minimum(size(matrix)), niter)
	for i ∈ 1:niter
		sample_from!(null, matrix)
		λ[:,i] = pvals(null)
	end

	return λ
end

# ╔═╡ 7bf6e4e4-7d27-11eb-0d11-a343a5ead1fe
λ = (
	data = pvals(scrna),
	null = null_pvals(scrna),
);

# ╔═╡ 564474de-7d2a-11eb-107a-094fc5b153f6
λ.null

# ╔═╡ de3c3612-7d25-11eb-3233-393941f590e5
begin
	plot(λ.data[1:end-1], label="data", color=:cyan3, linewidth=3)
	for i ∈ 1:size(λ.null,2)
		plot!(λ.null[1:end-1,i], label="", color=:orangered2, alpha=0.5, linewidth=1)
	end
	plot!(λ.null[1:end-1,1], label="null", color=:orangered2, alpha=1, linewidth=1)

	xaxis!("principal value", :log10)
	yaxis!("fraction of variance", :log10)

	#savefig("../figs/principal_values.svg")
end

# ╔═╡ 59156b62-7d2c-11eb-1f14-3139307e312e
"""
This section generates figures analyzing state localization
"""

# ╔═╡ 713cebde-7d2c-11eb-3918-31d392278150
ipr(ψ) = sum(ψ.^4, dims=1) ./ (sum(ψ.^2, dims=1)).^2

# ╔═╡ 151d45be-7d2d-11eb-00ae-495dff2916f2
function scan_ipr(x, Ns; niter=5)
	IPR = zeros(size(x,2), niter, length(Ns))
	for (n,N) ∈ enumerate(Ns)
		for i ∈ 1:niter
			ι = randperm(size(x,1))[1:N]
			IPR[:,i,n] = ipr(x[ι,:])
		end
	end
	IPR
end

# ╔═╡ f940b7ee-7d2d-11eb-3ecf-930c0eb1b938
F = svd(scrna);

# ╔═╡ 1ab88ad2-7d2e-11eb-2327-3f2015c5220c
Ns = [100,250,500,750,1000,2000,3000,4000,5000,6000,7000,8000,size(F.U,1)]; IPR = scan_ipr(F.U, Ns; niter=50);

# ╔═╡ 476449e0-7d2e-11eb-2d91-11e5a358d060
ĪP̄R̄ = dropdims(mean(IPR,dims=2);dims=2)

# ╔═╡ 4a089ff2-7d2e-11eb-0e3b-9d580d220c45
begin
	cmap = ColorSchemes.plasma
	plot(Ns, 1 ./ ĪP̄R̄[1,:],color=get(cmap,0.0),linewidth=2,label="")
	for i in 2:49
		plot!(Ns, 1 ./ ĪP̄R̄[i,:],color=get(cmap,(i-1)/50),linewidth=2,label="")
	end
	plot!(Ns, 1 ./ ĪP̄R̄[50,:],color=get(cmap,1.0),linewidth=2,label="")

	xaxis!("number of genes subsampled", :log10)
	yaxis!(L"1 / \sum \psi^4", :log10)

	#savefig("../figs/participation_ratio_subsample_scaling.svg")
end

# ╔═╡ 49dc2956-7d31-11eb-03c8-f5e29d9d0798
ψ = ipr(F.U)

# ╔═╡ 58a0f57a-7d31-11eb-3b59-293c95b4aa07
begin
	plot(1 ./ ψ')

	xaxis!("principal component", :log10)
	yaxis!("participation ratio", :log10)

	#savefig("../figs/participation_ratio.svg")
end

# ╔═╡ b499f6bc-7d34-11eb-106d-53126cd7a5c6
"""
This section generates figures analyzing point cloud scaling.
"""

# ╔═╡ c115efae-7d34-11eb-0d58-bba15e9922bb
G = F.U[:, 1:25] * Diagonal(F.S[1:25]) * F.Vt[1:25,:];

# ╔═╡ 1a663bc4-7d35-11eb-20e2-99f298281450
Dₑ = M.SeqSpace.PointCloud.distance(G);

# ╔═╡ 3838406e-7d35-11eb-31fb-256eea2a72cb
Dₚ = M.SeqSpace.PointCloud.geodesics(G, 6);

# ╔═╡ 629180aa-7d35-11eb-2bd8-47deefc64f91
function scaling(D, N)
	Rₘᵢₙ = minimum(D[D .> 0])
	Rₘₐₓ = maximum(D)
	Rs   = range(Rₘᵢₙ,Rₘₐₓ,length=N)

	ϕ = zeros(size(D,1), N)
	for (i,R) ∈ enumerate(Rs)
		ϕ[:,i] = sum(D .<= R, dims=1)
	end

	return ϕ, Rs
end

# ╔═╡ 4c497cfc-7d36-11eb-013a-07b8d183909f
ρₑ, Rₑ = scaling(Dₑ, 1000);

# ╔═╡ f3da8bec-7d37-11eb-06c3-e58f9312c561
ρₚ, Rₚ = scaling(Dₚ, 1000);

# ╔═╡ 61689ce4-7d36-11eb-206f-a362c5f1b24b
begin
	plot(Rₑ, ρₑ', alpha=0.03, color=:cyan3, label="", title="Euclidean")
	plot!(Rₑ, mean(ρₑ, dims=1)', linewidth=5, alpha=1, color=:cyan3, label="mean")
	plot!(Rₑ, 5e-10*Rₑ.^7, linestyle=:dashdot, linewidth=2, alpha=1, color=:black, label=L"\alpha R^7")

	xaxis!("radius", :log10, (25, 400))
	yaxis!("number of points", :log10, (10, 1297))

	#savefig("../figs/euclidean_ball_scaling.svg")
end

# ╔═╡ f1e88e9a-7d37-11eb-0529-a3ab42f90864
begin
	plot(Rₚ, ρₚ', alpha=0.02, color=:cyan3, label="", title="Geodesic")
	plot!(Rₚ, mean(ρₚ, dims=1)', linewidth=5, alpha=1, color=:cyan3, label="mean")
	plot!(Rₚ, 9e-5*Rₚ.^3.2, linestyle=:dashdot, linewidth=2, alpha=1, color=:black, label=L"\alpha R^3")

	xaxis!("radius", :log10, (25, 400))
	yaxis!("number of points", :log10, (10, 1297))

	#savefig("../figs/geodesic_ball_scaling.svg")
end

# ╔═╡ e2a9cb1c-7d3a-11eb-088a-7d51b807d714
"""
This section generates figures analyzing isomap
"""

# ╔═╡ 8c321e78-7d3b-11eb-375b-5dd081921104
invert, embryo = I.Inference.inversion(); Ψ = invert(0.5);

# ╔═╡ a7c62f6e-7d3b-11eb-1ef0-f3743d0e58e1
begin
	ψᵣ = Ψ  ./ sum(Ψ, dims=2)
	ψₗ = (Ψ ./ sum(Ψ, dims=1))'
end;

# ╔═╡ b0bc877e-7d3b-11eb-1110-913567d545da
r̂ = (ψₗ * embryo)'; AP = r̂[1,:]; DV = atan.(r̂[2,:], r̂[3,:])

# ╔═╡ ed76e2e6-7d3a-11eb-04ed-69f2bdd9f68a
ξ = M.SeqSpace.PointCloud.mds(Dₚ.^2, 10)

# ╔═╡ 0c2a0f80-7d3e-11eb-17b6-71185c47896d
begin

residual = zeros(size(ξ,2))
for d ∈ 1:size(ξ,2)
	D̄ = .√M.SeqSpace.PointCloud.distance²(ξ[:,1:d]')

	residual[d] = cor(M.SeqSpace.PointCloud.upper_tri(D̄), M.SeqSpace.PointCloud.upper_tri(Dₚ))
end

end;

# ╔═╡ 694eb1e4-7d42-11eb-0ac6-8fe4f3896971
begin
	d̄ = M.SeqSpace.PointCloud.upper_tri(.√M.SeqSpace.PointCloud.distance²(ξ[:,1:6]'));
	d = M.SeqSpace.PointCloud.upper_tri(Dₚ)

	ι = sample(1:length(d), 5000)

	scatter(d[ι], d̄[ι], legend=false)

	xaxis!("estimated geodesic distance")
	yaxis!("pairwise euclidean distance in latent space")
	#savefig("../figs/distance_isomap_correlation_eg_d3.svg")
end

# ╔═╡ 8cc4e066-7d3e-11eb-37b7-312b8d2a4a73
begin
	plot(residual, linewidth=2, label="")

	xaxis!("latent dimension")
	yaxis!("correlation of distance matrices", (0.7,0.95))

	#savefig("../figs/distance_correlation_vs_isomap_dimension.svg")
end

# ╔═╡ 43b9ec00-7d3d-11eb-2973-cbb9af2197a8
@bind THETA PlutoUI.Slider(range(0,90,length=100))

# ╔═╡ 6cca0e84-7d3d-11eb-1a00-77473b588de7
@bind PHI PlutoUI.Slider(range(0,90,length=100))

# ╔═╡ ffdcaae6-7d3b-11eb-3ddb-558d8e1638d6
begin
	scatter(ξ[:,2],-ξ[:,1],ξ[:,3],
		marker_z=AP,
		markersize=5,
		xlims=(-150,150),
		ylims=(-150,150),
		zlims=(-125,75),
		title="predicted AP",
		camera=(THETA,PHI))

	#savefig("../figs/isomap_3d_AP.svg")
end

# ╔═╡ a3fd873e-7d3d-11eb-391e-cfc5d3c80797
begin
	scatter(ξ[:,2],-ξ[:,1],ξ[:,3],
		marker_z=DV,
		markersize=5,
		xlims=(-150,150),
		ylims=(-150,150),
		zlims=(-125,75),
		title="predicted DV",
		camera=(THETA,PHI))

	#savefig("../figs/isomap_3d_DV.svg")
end

# ╔═╡ 66270b84-7d46-11eb-2e09-c97d61f6425b
begin
scatter(F.Vt[3,:],F.Vt[4,:],F.Vt[5,:],
		marker_z=AP,
		markersize=5,
		title="predicted AP",
		camera=(THETA,PHI)
)

#savefig("../figs/svd_3d_AP.svg")

end

# ╔═╡ 927cd290-7d46-11eb-057c-c14cffce4ce5
begin

scatter(F.Vt[2,:],F.Vt[3,:],F.Vt[4,:],
		marker_z=DV,
		markersize=5,
		title="predicted DV",
		camera=(THETA,PHI)
)

#savefig("../figs/svd_3d_DV.svg")

end

# ╔═╡ d1c09536-7d46-11eb-2d39-79a39782580b
Dₓ = M.SeqSpace.PointCloud.geodesics(r̂,6)

# ╔═╡ ead3da76-7d46-11eb-0d7c-cb4496b2c27a
begin
	dx = M.SeqSpace.PointCloud.upper_tri(Dₓ)
	dp = M.SeqSpace.PointCloud.upper_tri(Dₚ)

	ι₂ = sample(1:length(dx), 5000)

	scatter(dx[ι₂], dp[ι₂])
	xlabel!("spatial geodesic distance")
	ylabel!("expression geodesic distance")

	#savefig("../figs/spatial_vs_expression_geodesic.svg")
end

# ╔═╡ Cell order:
# ╠═489b13da-7d23-11eb-39de-57a12469721f
# ╠═82c5386a-7d2f-11eb-3992-e985ee16933c
# ╠═4bd45848-7d23-11eb-1a8b-b36f7461f368
# ╠═4256e210-7d23-11eb-2f7a-1d8aceab0682
# ╠═3d2bf622-7d23-11eb-33fb-3f8a4f0f309e
# ╠═5e1fa2de-7d23-11eb-18da-ddf0ec5714cc
# ╠═a0c3a118-7d3b-11eb-0309-a9540887541f
# ╠═fce24760-7d22-11eb-2474-a5efdd176182
# ╠═330fd846-7d25-11eb-3933-b5f8871eedc5
# ╠═4ddffd5a-7d24-11eb-2c70-675569787f53
# ╠═fa6e4914-7d24-11eb-0ff2-1974e4ec8df2
# ╠═72a6c262-7d25-11eb-1a05-756c421933f5
# ╠═85650856-7d24-11eb-055d-75a99e549182
# ╠═c3fdc5ce-7d26-11eb-235a-33e432adcd91
# ╠═7bf6e4e4-7d27-11eb-0d11-a343a5ead1fe
# ╠═564474de-7d2a-11eb-107a-094fc5b153f6
# ╠═de3c3612-7d25-11eb-3233-393941f590e5
# ╠═59156b62-7d2c-11eb-1f14-3139307e312e
# ╠═713cebde-7d2c-11eb-3918-31d392278150
# ╠═151d45be-7d2d-11eb-00ae-495dff2916f2
# ╠═f940b7ee-7d2d-11eb-3ecf-930c0eb1b938
# ╠═2ef712f2-7d2e-11eb-228c-2d023dc6ab10
# ╠═1ab88ad2-7d2e-11eb-2327-3f2015c5220c
# ╠═476449e0-7d2e-11eb-2d91-11e5a358d060
# ╠═4a089ff2-7d2e-11eb-0e3b-9d580d220c45
# ╠═49dc2956-7d31-11eb-03c8-f5e29d9d0798
# ╠═58a0f57a-7d31-11eb-3b59-293c95b4aa07
# ╠═b499f6bc-7d34-11eb-106d-53126cd7a5c6
# ╠═c115efae-7d34-11eb-0d58-bba15e9922bb
# ╠═1a663bc4-7d35-11eb-20e2-99f298281450
# ╠═3838406e-7d35-11eb-31fb-256eea2a72cb
# ╠═629180aa-7d35-11eb-2bd8-47deefc64f91
# ╠═4c497cfc-7d36-11eb-013a-07b8d183909f
# ╠═f3da8bec-7d37-11eb-06c3-e58f9312c561
# ╠═61689ce4-7d36-11eb-206f-a362c5f1b24b
# ╠═f1e88e9a-7d37-11eb-0529-a3ab42f90864
# ╠═e2a9cb1c-7d3a-11eb-088a-7d51b807d714
# ╠═8c321e78-7d3b-11eb-375b-5dd081921104
# ╠═a7c62f6e-7d3b-11eb-1ef0-f3743d0e58e1
# ╠═b0bc877e-7d3b-11eb-1110-913567d545da
# ╠═ed76e2e6-7d3a-11eb-04ed-69f2bdd9f68a
# ╠═a1b684a0-7d40-11eb-2a4c-e990f12c6cbb
# ╠═0c2a0f80-7d3e-11eb-17b6-71185c47896d
# ╠═694eb1e4-7d42-11eb-0ac6-8fe4f3896971
# ╠═8cc4e066-7d3e-11eb-37b7-312b8d2a4a73
# ╠═43b9ec00-7d3d-11eb-2973-cbb9af2197a8
# ╠═6cca0e84-7d3d-11eb-1a00-77473b588de7
# ╠═ffdcaae6-7d3b-11eb-3ddb-558d8e1638d6
# ╠═a3fd873e-7d3d-11eb-391e-cfc5d3c80797
# ╠═66270b84-7d46-11eb-2e09-c97d61f6425b
# ╠═927cd290-7d46-11eb-057c-c14cffce4ce5
# ╠═d1c09536-7d46-11eb-2d39-79a39782580b
# ╠═ead3da76-7d46-11eb-0d7c-cb4496b2c27a


module FlexLinearAlgebra
using LinearAlgebra

import Base: (+), (-), (*), (==), sum, #(.*),
    getindex, setindex!, hash, show, keys, values, size,
    keytype, valtype, length, haskey, hash, Vector, Matrix

#    dot, adjoint

#import LinearAlgebra: dot, adjoint

include("FlexVector.jl")
include("FlexMatrix.jl")
end # module


function ADScalarObjective(f,x0::AbstractArray)
    Hres = DiffResults.HessianResult(x0)
    function _g(df,x,Hresult)
        ForwardDiff.gradient!(Hresult,f,x)
        df .= DiffResults.gradient(Hresult)
        df
    end
    function _fg(df,x,Hresult)
        ForwardDiff.gradient!(Hresult,f,x)
        df .= DiffResults.gradient(Hresult)
        fx = DiffResults.value(Hresult)
        return fx,df
    end

    function _fgh(df,d2f,x,Hresult)
        ForwardDiff.hessian!(Hresult,f,x)
        d2f .= DiffResults.hessian(Hresult)
        df .= DiffResults.gradient(Hresult)
        fx = DiffResults.value(Hresult)
        return fx,df,d2f
    end

    function h(d2f,x)
        ForwardDiff.hessian!(d2f,f,x)
        d2f
    end
    g(df,x) = _g(df,x,Hres)
    fg(df,x) = _fg(df,x,Hres)
    fgh(df,d2f,x) =  _fgh(df,d2f,x,Hres)
    return ScalarObjective(f=f,
    g=g,
    fg=fg,
    fgh=fgh,
    h=h)
end

function ADScalarObjective(f,x0::Number)
    function g(x)
        return ForwardDiff.derivative(f,x)
    end

    function fg(∂fx,x)
        ∂fx,x = f∂f(f,x)
        return ∂fx,x
    end

    function fgh(∂fx,∂2fx,x)
        fx,∂fx,∂2fx = f∂f∂2f(f,x)
        return fx,∂fx,∂2fx
    end

    function h(∂2fx,x)
        ∂2fx = ForwardDiff.derivative(g,x)
        return ∂2fx
    end

    return ScalarObjective(f=f,
    g=g,
    fg=fg,
    fgh=fgh,
    h=h)
end
"""
    function optimize(f,x0,method=LineSearch(Newton()), options=OptimizationOptions())
"""
function optimize(f,x0::Number,method=LineSearch(Newton()),options=OptimizationOptions())
    scalarobj = ADScalarObjective(f,x0)
    optprob = OptimizationProblem(scalarobj; inplace=false)
    return NLSolvers.solve(optprob, x0, method,options)
end

#this should allow to use another schemes to update
struct ModifiedActiveBox{M,T}
    scheme::M
    ϵ::T
end

Base.summary(ab::ModifiedActiveBox) = "Active Box, using $(Base.summary(ab.scheme))"

NLSolvers.modelscheme(ab::ModifiedActiveBox) = ab.scheme
ModifiedActiveBox(scheme=Newton(); epsilon=1e-12) =  ModifiedActiveBox(scheme,epsilon)

function box_optimize(f,x0,lb,ub,options=OptimizationOptions();tolbounds = 1e-12)
    scalarobj = ADScalarObjective(f,x0)
    prob = OptimizationProblem(obj=scalarobj,
    bounds=(lb,ub); inplace=true)
    NLSolvers.solve(prob, x0, ModifiedActiveBox(Newton(linsolve = cholesky_linsolve);epsilon=tolbounds), options)
end

#see https://github.com/JuliaNLSolvers/NLSolvers.jl/issues/21
function NLSolvers.solve(prob::OptimizationProblem, x0, scheme::ModifiedActiveBox, options::OptimizationOptions)
    t0 = time()

    x0, B0 = x0, (false*x0*x0' +I)
    lower, upper = NLSolvers.bounds(prob)
    ϵbounds = mapreduce(b->(b[2] - b[1])/2, min, zip(lower, upper)) # [1, pp. 100: 5.41]

    !any(clamp.(x0, lower, upper) .!= x0) || error("Initial guess not in the feasible region")

    linesearch = NLSolvers.ArmijoBertsekas()
    mstyle = NLSolvers.OutOfPlace()
    direction_scheme = NLSolvers.modelscheme(scheme)
    objvars = NLSolvers.prepare_variables(prob, scheme, x0, copy(x0), B0)
    f0, ∇f0 = objvars.fz, norm(objvars.∇fz, Inf) # use user norm
    fz, ∇fz = objvars.fz, objvars.∇fz # use user norm
    fx, ∇fx = fz, copy(∇fz)
    B = B0
    x, z = copy(x0), copy(x0)
    Tf = typeof(fz)
    is_first=false
    Ix = Diagonal(z.*0 .+ 1)

    unconstrained_d = copy(∇fz)
    d = copy(∇fz)
    dmin = copy(∇fz)
    y = nothing
    s = copy(∇fz)
    xnorm = copy(∇fz)
    activeset = similar(∇fz,Bool)
    Ĥ = copy(B)
    Ĥchol = cholesky(Ĥ)
    for iter = 1:options.maxiter
        x .= z
        fx = fz
        ∇fx .= ∇fz

        ϵ = min(norm(clamp.(x.-∇fx, lower, upper).-x), ϵbounds) # Kelley 5.41 and just after (83) in [1]
        activeset .= NLSolvers.is_ϵ_active.(x, lower, upper, ∇fx, ϵ)
        Ĥ .= NLSolvers.diagrestrict.(B, activeset, activeset', Ix)
        #d = clamp.(x.-HhatChol\∇fx, lower, upper).-x
        # Update current gradient and calculate the search direction
        NLSolvers.find_direction!(unconstrained_d,Ĥ,objvars.Pg,∇fx,direction_scheme)
        d .= clamp.(x .+ unconstrained_d, lower, upper).-x #clamp unconstrained direction
        dmin .= d./unconstrained_d
        d_scale = minimum(dmin)
        #scale the direction, the original direction is the same.
        #d .=  @. unconstrained_d*d_scale + (one(d_scale)-d_scale)*d
        !iszero(d_scale) && (d .= unconstrained_d .* d_scale)
        #if the scale is zero, we are actually in a boundary, and we can only grow parallel along the box.
        φ = NLSolvers._lineobjective(mstyle, prob, ∇fz, z, x, d, fz, dot(∇fz, d))

        # Perform line search along d
        # Also returns final step vector and update the state
        α, f_α, ls_success = NLSolvers.find_steplength(mstyle, linesearch, φ, Tf(1), ∇fz, activeset, lower, upper, x, d, ∇fx, activeset)
        # # Calculate final step vector and update the state
        if !ls_success
            return NLSolvers.ConvergenceInfo(scheme, (prob=prob, B=B, ρs=norm(x.-z), ρx=norm(x), minimizer=z, fx=fx, minimum=fz, ∇fz=∇fz, f0=f0, ∇f0=∇f0, iter=iter, time=time()-t0), options)
        end

        s .= @. α * d
        z .= @. x + s
        s .= clamp.(z, lower, upper) .- x
        z .= x .+ s
        # Update approximation. it only works with newton at the moment
        fz, ∇fz, B, s, y = NLSolvers.update_obj!(prob.objective,s,y, ∇fx, z, ∇fz, B, direction_scheme)
        #fz, ∇fz, B, s, y = NLSolvers.update_obj(prob.objective, s, ∇fx, z, ∇fz, B, direction_scheme,is_first)
        #@show fz, ∇fz, B, s, y
        #@show Ĥ
        xnorm .= x.-clamp.(x.-∇fz, lower, upper)
        normx = norm(xnorm, Inf)
        if normx < options.g_abstol
            return NLSolvers.ConvergenceInfo(scheme, (prob=prob, B=B, ρs=norm(x.-z), ρx=norm(x), minimizer=z, fx=fx, minimum=fz, ∇fz=∇fz, f0=f0, ∇f0=∇f0, iter=iter, time=time()-t0), options)
        end
    end
  iter =  options.maxiter
  return NLSolvers.ConvergenceInfo(scheme, (prob=prob, B=B, ρs=norm(x.-z), ρx=norm(x), minimizer=z, fx=fx, minimum=fz, ∇fz=∇fz, f0=f0, ∇f0=∇f0, iter=iter, time=time()-t0), options)
end

#=
julia> Optim.minimizer(results)
2-element Vector{Float64}:
 1.2500000000000002
 1.5626548683420372

julia> Optim.minimum(results)
0.06250239842033663
=#


module TestMultiobjective

using Gurobi
using Test

function runtests()
    for name in names(@__MODULE__; all = true)
        if startswith("$(name)", "test_")
            @testset "$(name)" begin
                getfield(@__MODULE__, name)()
            end
        end
    end
    return
end

const MOI = Gurobi.MOI

const GRB_ENV = isdefined(Main, :GRB_ENV) ? Main.GRB_ENV : Gurobi.Env()

function test_multiobjective()
    model = Gurobi.Optimizer(GRB_ENV)
    MOI.set(model, MOI.Silent(), true)
    MOI.Utilities.loadfromstring!(
        model,
        """
variables: x, y
minobjective: 2x + y
c1: x + y >= 1.0
c2: 0.5 * x + 1.0 * y >= 0.75
c3: x >= 0.0
c4: y >= 0.25
""",
    )
    x = MOI.get(model, MOI.VariableIndex, "x")
    y = MOI.get(model, MOI.VariableIndex, "y")

    f = MOI.ScalarAffineFunction(
        [MOI.ScalarAffineTerm(1.0, x), MOI.ScalarAffineTerm(3.0, y)],
        0.0,
    )

    MOI.set(model, Gurobi.MultiObjectiveFunction(2), f)

    @test MOI.get(model, Gurobi.MultiObjectiveWeight(1)) == 1.0
    @test MOI.get(model, Gurobi.MultiObjectiveWeight(2)) == 1.0
    @test MOI.get(model, Gurobi.MultiObjectivePriority(1)) == 0
    @test MOI.get(model, Gurobi.MultiObjectivePriority(2)) == 0

    MOI.optimize!(model)
    @test MOI.get(model, MOI.ObjectiveValue()) ≈ 1.5
    @test MOI.get(model, MOI.VariablePrimal(), x) ≈ 0.5
    @test MOI.get(model, MOI.VariablePrimal(), y) ≈ 0.5

    BFS = [
        (x = 1.0, y = 0.25, f1 = 2.25, f2 = 1.75),
        (x = 0.5, y = 0.5, f1 = 1.5, f2 = 2.0),
        (x = 0.0, y = 1.0, f1 = 1.0, f2 = 3.0),
    ]
    for (i, λ) in enumerate([0.2, 0.5, 0.8])
        MOI.set(model, Gurobi.MultiObjectiveWeight(1), λ)
        MOI.set(model, Gurobi.MultiObjectiveWeight(2), 1 - λ)
        MOI.optimize!(model)
        @test MOI.get(model, MOI.VariablePrimal(), x) ≈ BFS[i].x
        @test MOI.get(model, MOI.VariablePrimal(), y) ≈ BFS[i].y
        @test MOI.get(model, Gurobi.MultiObjectiveValue(1)) ≈ BFS[i].f1
        @test MOI.get(model, Gurobi.MultiObjectiveValue(2)) ≈ BFS[i].f2
    end

    MOI.set(model, Gurobi.MultiObjectiveWeight(1), 1.0)
    MOI.set(model, Gurobi.MultiObjectiveWeight(2), 1.0)
    MOI.set(model, Gurobi.MultiObjectivePriority(1), 1)
    MOI.set(model, Gurobi.MultiObjectivePriority(2), 2)
    @test MOI.get(model, Gurobi.MultiObjectivePriority(1)) == 1
    @test MOI.get(model, Gurobi.MultiObjectivePriority(2)) == 2

    MOI.optimize!(model)

    @test MOI.get(model, MOI.VariablePrimal(), x) ≈ BFS[1].x
    @test MOI.get(model, MOI.VariablePrimal(), y) ≈ BFS[1].y
    @test MOI.get(model, Gurobi.MultiObjectiveValue(1)) ≈ BFS[1].f1
    @test MOI.get(model, Gurobi.MultiObjectiveValue(2)) ≈ BFS[1].f2

    MOI.set(model, Gurobi.MultiObjectivePriority(1), 2)
    MOI.set(model, Gurobi.MultiObjectivePriority(2), 1)
    @test MOI.get(model, Gurobi.MultiObjectivePriority(1)) == 2
    @test MOI.get(model, Gurobi.MultiObjectivePriority(2)) == 1

    MOI.optimize!(model)

    @test MOI.get(model, MOI.VariablePrimal(), x) ≈ BFS[3].x
    @test MOI.get(model, MOI.VariablePrimal(), y) ≈ BFS[3].y
    @test MOI.get(model, Gurobi.MultiObjectiveValue(1)) ≈ BFS[3].f1
    @test MOI.get(model, Gurobi.MultiObjectiveValue(2)) ≈ BFS[3].f2
end

end

TestMultiobjective.runtests()