module LinSolvers
    using ..NEPCore
    using LinearAlgebra
    using SparseArrays
    using SuiteSparse
    using IterativeSolvers
    using LinearMaps
    using ArnoldiMethod

    # Linear system of equation solvers
    export LinSolver
    export FactorizeLinSolver
    export BackslashLinSolver
    export GMRESLinSolver
    export lin_solve

    # Eigenvalue solvers
    export EigSolver
    export DefaultEigSolver
    export EigenEigSolver
    export ArnoldiEigSolver
    export eig_solve

    import Base.eltype
    export eltype
    import Base.size
    export size
    import Base.*
    export *

##############################################################################
"""
    abstract type LinSolver

Structs inheriting from this type are able to solve linear systems associated with
a NEP, for a specific `λ`-value. The most common are direct solvers such as
[`FactorizeLinSolver`](@ref), [`BackslashLinSolver`](@ref) and iterative solvers
such as [`GMRESLinSolver`](@ref).

The LinSolver objects are usually created by the NEP-algorithms through
creator functions, which are passed as parameters.

# Example

The most common usecase is that you want to pass a `linsolvercreator`-function
as a parameter to the NEP-algorithm.
This example shows how you can use solvers based on backslash or `factorize()`.
In the example, `BackslashLinSolver` does not exploit that the system matrix
remains the same throughout the algorithm and is therefore slower.
```julia-repl
julia> nep=nep_gallery("qdep0");
julia> using BenchmarkTools
julia> v0=ones(size(nep,1));
julia> @btime λ,v=quasinewton(nep,λ=-1,v=v0, linsolvercreator=DefaultLinSolverCreator());
  181.017 ms (3779 allocations: 58.61 MiB)
julia> @btime λ,v=quasinewton(nep,λ=-1,v=v0, linsolvercreator=BackslashLinSolverCreator());
  2.040 s (4510 allocations: 553.24 MiB)
```

# Example

The `LinSolver`s are constructed for extendability. This example creates our
own `LinSolver` which uses an explicit formula for the inverse
if the NEP has dimension 2x2.

Create the types and a creator.
```julia
julia> using LinearAlgebra
julia> struct MyLinSolver <: LinSolver
   M::Matrix{ComplexF64}
end
julia> function my_linsolvercreator(nep,λ)
   M=compute_Mder(nep,λ);
   return MyLinSolver(M);
end
```
Explicit import `lin_solve` to show how to solve a linear system.
```julia-repl
julia> import NonlinearEigenproblems.LinSolvers.lin_solve;
julia> function lin_solve(solver::MyLinSolver,b::AbstractVecOrMat;tol=0)
   M=solver.M;
   invM=(1/(det(M)))*[M[2,2] -M[1,2];-M[2,1] M[1,1]]
   return invM*b
end
julia> nep=SPMF_NEP([[1.0 3.0; 4.0 5.0], [2.0 1.0; -1 2.0]], [S->S^2,S->exp(S)])
julia> λ,v=quasinewton(nep,λ=-1,v=[1;1],linsolvercreator=my_linsolvercreator);
```

See also: [`lin_solve`](@ref),
[`FactorizeLinSolver`](@ref), [`FactorizeLinSolver`](@ref),
[`DefaultLinSolverCreator`](@ref),
[`BackslashLinSolver`](@ref), [`BackslashLinSolverCreator`](@ref),
[`GMRESLinSolver`](@ref), [`GMRESLinSolverCreator`](@ref)

"""
    abstract type LinSolver end

##############################################################################
"""
    struct FactorizeLinSolver <: LinSolver

This represents the linear solver associated with julia `factorize()`.
See [`LinSolver`](@ref) and [`FactorizeLinSolverCreator`](@ref) for examples.
"""
    struct FactorizeLinSolver{T} <: LinSolver
        Afact::T
        umfpack_refinements::Int
    end

    function FactorizeLinSolver(nep::NEP, λ, umfpack_refinements)
        A=compute_Mder(nep,λ)
        Afact=factorize(A)
        return FactorizeLinSolver(Afact, umfpack_refinements)
    end

"""
    lin_solve(solver::LinSolver, b::AbstractVecOrMat; tol=0)

This function solves the linear system represented in `solver::LinSolver` with a
right-hand side `b`. The `tol` kwarg is controlling how accurate the linear
system needs to be solved. A NEP-algorithm will call this solver every
time a linear system associated with `M(λ)` needs to be solved.

This function must be overloaded if a user wants to define their own
way of solving linear systems. See [`LinSolver`](@ref) for examples.
"""
    function lin_solve(solver::FactorizeLinSolver, x::Array; tol = 0)
        with_umfpack_refinements(solver.umfpack_refinements) do
            solver.Afact \ x
        end
    end

    function with_umfpack_refinements(f::Function, refinements)
        current_refinements = SuiteSparse.UMFPACK.umf_ctrl[8]
        try
            SuiteSparse.UMFPACK.umf_ctrl[8] = refinements
            f()
        finally
            SuiteSparse.UMFPACK.umf_ctrl[8] = current_refinements
        end
    end

"""
    struct BackslashLinSolver <: LinSolver

This represents a linear solver corresponding to the backslash operator (no pre-factorization).

See also: [`LinSolver`](@ref) and [`BackslashLinSolverCreator`](@ref)
"""
    struct BackslashLinSolver{T_mat} <: LinSolver
        A::T_mat
    end
#    BackslashLinSolver(A::AbstractMatrix{T_num}) where T_num = new(A)

    function BackslashLinSolver(nep::NEP,λ)
        A=compute_Mder(nep,λ)
        return BackslashLinSolver{typeof(A)}(A)
    end

    function lin_solve(solver::BackslashLinSolver{T_mat}, x::Array; tol=0) where T_mat
        return solver.A\x
    end

##############################################################################
"""
    struct GMRESLinSolver <: LinSolver

This represents a solver done with the julia GMRES implementation.

See also: [`LinSolver`](@ref), [`GMRESLinSolverCreator`](@ref)
"""
    struct GMRESLinSolver{T_num<:Number, T_kwargs} <: LinSolver
        A::LinearMap{T_num}
        kwargs::T_kwargs

        function GMRESLinSolver{T_num}(nep, λ::T_num, kwargs::T_kwargs) where {T_num<:Number, T_kwargs}
            f = v -> compute_Mlincomb(nep,λ,v)
            A = LinearMap{T_num}(f, size(nep,1), ismutating=false)
            new{T_num, T_kwargs}(A, kwargs)
        end

    end

    function lin_solve(solver::GMRESLinSolver{T_num,T_kwargs}, b::Vector{T_num}; tol=eps(real(T_num))) where {T_num,T_kwargs}
        v = zero(b)
        retval=gmres!(v, solver.A, b; reltol=tol, solver.kwargs...)
        return v
    end

##############################################################################
"""
    abstract type EigSolver

Structs inheriting from this type are able to solve linear eigenvalue problems
arising in certain methods, such as, e.g., `mslp`, `sgiter`,
and `polyeig`.

The `EigSolver` objects are passed as types to the NEP-algorithms,
which uses it to dispatch the correct version of the function [`eig_solve`](@ref).

# Example

The most common usecase is that you do not want to specify anything in particular, since
the [`DefaultEigSolver`](@ref) will use a dense or a sparse method depending on you problem.
However, this example shows how you can force `mslp` to use the sparse solver.
```julia-repl
julia> nep=nep_gallery("qdep0");
julia> λ,v = mslp(nep, eigsolvertype=ArnoldiEigSolver);
julia> norm(compute_Mlincomb(nep,λ,v))
9.323110647141726e-16
```

# Example

The `EigSolver`s are constructed for extendability. As an illustartion this example
creates a naive `EigSolver` which casts the problem to a standard linear eigenproblem
and calls the built-in function to solve it.

Create the types and a creator.
```julia-repl
julia> struct MyEigSolver <: EigSolver
   A
   E
   function MyEigSolver(A,E)
      return new(A,E)
   end
end

julia> import NonlinearEigenproblems.LinSolvers.eig_solve;
julia> function eig_solve(solver::MyEigSolver;nev = 1, target = 0)
   M = solver.E \\ solver.A
   eig = eigen(M)
   i = argmin(abs.(eig.values))
   return eig.values[i], eig.vectors[:,i]
end
julia> nep=nep_gallery("dep0", 50);
julia> λ,v = mslp(nep, eigsolvertype=MyEigSolver, tol=1e-5);
julia> norm(compute_Mlincomb(nep,λ,v))
3.0777795031319117e-10
```

See also: [`eig_solve`](@ref),
[`DefaultEigSolver`](@ref), [`EigenEigSolver`](@ref),
[`ArnoldiEigSolver`](@ref), [`eig_solve`](@ref)

"""
    abstract type EigSolver end

##############################################################################
"""
    struct EigenEigSolver <: EigSolver

A linear eigenvalueproblem solver that calls Julia's in-built eigen()

Constructed as `EigenEigSolver(A, [B,])`, and solves the problem
```math
Ax = λBx
```
The paramter `B` is optional an default is indentity, for which a standard linear
eigenproblem is solved.

See also: [`EigSolver`](@ref) and [`eig_solve`](@ref)
"""
    struct EigenEigSolver{T_A,T_B} <: EigSolver
        A::T_A
        B::T_B

        function EigenEigSolver(A)
            return new{typeof(A),Missing}(A,missing)
        end
        function EigenEigSolver(A,B)
            return new{typeof(A),typeof(B)}(A,B)
        end
    end

    function eig_solve(solver::EigenEigSolver; nev = 1, target = 0)
        D,V = inner_eig_solve(solver)

        #Sort the eigenvalues wrt distance from target, and permute
        I = sortperm(abs.(target*ones(size(D,1))-D));
        D = D[I];V = V[:,I];
        return D[1:nev],V[:,1:nev];
    end

    function inner_eig_solve(solver::EigenEigSolver{T_A,T_B}) where {T_A, T_B}
        D,V = eigen(solver.A,solver.B)
    end
    function inner_eig_solve(solver::EigenEigSolver{T_A,T_B}) where {T_A, T_B<:Missing}
        D,V = eigen(solver.A)
    end

"""
    struct ArnoldiEigSolver <: EigSolver

A linear eigenproblem solver for large and sparse problems that calls
the Arnoldi method implemented in the Julia package ArnoldiMethod.jl.

Constructed as `ArnoldiEigSolver(A, [B,])`, and solves the problem
```math
Ax = λBx
```
The paramter `B` is optional an default is indentity, for which a standard linear
eigenproblem is solved.

See also: [`EigSolver`](@ref) and [`eig_solve`](@ref)
"""
    struct ArnoldiEigSolver{T_A,T_B} <: EigSolver
        A::T_A
        B::T_B

        function ArnoldiEigSolver(A)
            return new{typeof(A),Missing}(A,missing)
        end
        function ArnoldiEigSolver(A,B)
            return new{typeof(A),typeof(B)}(A,B)
        end

    end

    function eig_solve(solver::ArnoldiEigSolver; nev=6, target=0)
        D,V = inner_eigs_solve(solver, nev, target)
        return D,V
    end

    function inner_eigs_solve(solver::ArnoldiEigSolver{T_A,T_B}, nev, target) where {T_A, T_B}
        if (T_B <: Missing)
            C=target*I-solver.A;
            Cfact=factorize(C);
            Atransformed=LinearMap{eltype(Cfact)}(x->Cfact\x,
                                   size(solver.A,1),size(solver.A,1));
        else

            C=target*solver.B-solver.A;
            Cfact=factorize(C);
            Atransformed=LinearMap{eltype(Cfact)}(x->Cfact\(solver.B*x),
                                   size(solver.A,1),size(solver.A,1));
        end
        # Call restarted Arnoldi
        decomp, history = partialschur(Atransformed, nev=nev, tol=1e-10, which=LM());
        D0, V = partialeigen(decomp)
        # Sort. So we don't depend on partialschur to do be sorted already
        IJ=sortperm(-abs.(D0))
        # And reverse transformation
        D = target .- inv.(D0) # Reverse transformation
        return D[IJ[1:nev]],V[:,IJ[1:nev]]
    end

"""
    struct DefaultEigSolver <: EigSolver

A linear eigenvalueproblem solver that calls checks for sparsity and accordingly
assigns an appropriate solver.

See also: [`EigSolver`](@ref), [`eig_solve`](@ref), [`EigenEigSolver`](@ref), [`ArnoldiEigSolver`](@ref)
"""
    struct DefaultEigSolver{T_sub} <: EigSolver
        subsolver::T_sub

        function DefaultEigSolver(A,B)
            local subsolver
            if(issparse(A))
                subsolver = ArnoldiEigSolver(A,B)
            else
                subsolver = EigenEigSolver(A,B)
            end
            return new{typeof(subsolver)}(subsolver)
        end
        function DefaultEigSolver(A)
            local subsolver
            if(issparse(A))
                subsolver = ArnoldiEigSolver(A)
            else
                subsolver = EigenEigSolver(A)
            end
            return new{typeof(subsolver)}(subsolver)
        end
    end

"""
    eig_solve(solver::EigSolver; [nev,] [target,])

This function solves the linear eigenvalue problem represented in `solver::EigSolver`.
The `nev` kwarg is controlling the number of eigenvalues aimed for, and `target`
specifies around which point the eigenvalues are computed. The former has a defalut value
equalt to the seize of the problem, and the latter has a defalut value 0.

Return values are of the form (Vector, Matrix) where the former contains the eigenvalues
and the latter the eigenvectors.

This function must be overloaded if a user wants to define their own
way of solving linear eigenvalue problems. See [`EigSolver`](@ref) for examples.
"""
    function eig_solve(solver::DefaultEigSolver;nev=size(solver.subsolver.A,1),target=0)
        return eig_solve(solver.subsolver,nev=nev,target=target)
    end

    include("LinSolverCreators.jl");

end


struct ThermalParameters{T}
    κ::T
    α::T
    Cp::T
    dQdT::T
    H::T
end

function thermal_parameters(;
    κ = 1.0,
    α = 1e-6,
    Cp = 1.0,
    dQdT = 0.0,
    H = 0.0
)
    ThermalParameters(κ, α, Cp, dQdT, H)
end

function init_U_P(gr)
    nU = maximum(gr.e2n)
    P = fill(0.0, maximum(gr.e2nP))
    U = fill(0.0, nU .* 2)
    return U, P
end

function init_temperature(gr, IDs; type = :harmonic)

    θ = gr.θ
    r = gr.r

    Ttop, Tbot = 0.0, 1.0

    if type == :harmonic
        # Harmonic hermal perturbation
        s = @. (2.22-r)/(2.22-1.22)
        y44 = @. 1.0/8.0*√(35.0/π)*cos(4.0*θ)
        δT = @. 0.2*y44*sin(π*(1-s))
        T = @. 1.22*s/(2.22-s) + δT

    elseif type == :random
        # Linear temperature with random perturbation
        s = @. (2.22-r)/(2.22-1.22)
        T = s .* (1 .+ (rand(length(s)).-0.5).*0.01 )

    elseif type == :realistic
        transition = 2.22-0.2276
        T = @. (2.22-r)/(2.22-transition)
        idx = r .< transition
        T[idx] .= 1.0

    end

    fixT!(T, Ttop , Tbot, IDs)
    applybounds!(T, Tbot, Ttop)

    return T
end

function init_particle_temperature!(particle_fields, particle_info; type = :harmonic)

    θ = [particle_info[i].CPolar.x for i in eachindex(particle_info)]
    r = [particle_info[i].CPolar.z for i in eachindex(particle_info)]
    s = @. (2.22-r)/(2.22-1.22)

    if type == :harmonic
        # Harmonic hermal perturbation
        y44 = @. 1.0/8.0*√(35.0/π)*cos(4.0*θ)
        δT = @. 0.2*y44*sin(π*(1-s))
        particle_fields.T = @. 1.22*s/(2.22-s) + δT

    elseif type == :random
        # Linear temperature with random perturbation
        particle_fields.T = s .* (1 .+ (rand(length(s)).-0.5).*0.01 )

    end

end

function fixT!(T, Ttop, Tbot, IDs)
    @inbounds for i in axes(T,1)
        if IDs[i] == "inner"
            T[i] = Tbot
        elseif IDs[i] == "outter"
            T[i] = Ttop
        end
    end
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _airtovac{T<:AbstractFloat}(wave_air::T)
    if wave_air >= 2000.0
        wave_vac = wave_air
        for iter= 1:2
            sigma2 = (1e4/wave_vac)^2.0 # Convert to wavenumber squared
            # Computer conversion factor.
            fact = 1.0 + 5.792105e-2/(238.0185e0 - sigma2) +
                1.67917e-3/(57.362e0 - sigma2)
            wave_vac = wave_air*fact # Convert Wavelength
        end
        return wave_vac
    else
        return wave_air
    end
end

"""
    airtovac(wave_air) -> wave_vacuum

### Purpose ###

Converts air wavelengths to vacuum wavelengths.

### Explanation ###

Wavelengths are corrected for the index of refraction of air under standard
conditions.  Wavelength values below \$2000 Å\$ will not be altered.  Uses
relation of Ciddor (1996).

### Arguments ###

* `wave_air`: can be either a scalar or an array of numbers.  Wavelengths are
corrected for the index of refraction of air under standard conditions.
Wavelength values below \$2000 Å\$ will *not* be altered, take care within \$[1
Å, 2000 Å]\$.

### Output ###

Vacuum wavelength in angstroms, same number of elements as `wave_air`.

### Method ###

Uses relation of Ciddor (1996), Applied Optics 62, 958.

### Example ###

If the air wavelength is `w = 6056.125` (a Krypton line), then `airtovac(w)`
yields a vacuum wavelength of `6057.8019`.

### Notes ###

`vactoair` converts vacuum wavelengths to air wavelengths.

Code of this function is based on IDL Astronomy User's Library.
"""
airtovac(wave_air::Real) = _airtovac(float(wave_air))

@vectorize_1arg Real airtovac


# Submodule for imaging functions

#######################################
# Function to resize an array representation
# of an image to a certain number of pixeles
function resize_image(f, N)
    if size(f, 1) > size(f, 2)
		f = reshape(f, size(f, 2), size(f, 1), size(f, 3));
    end
	P = size(f, 1);
	# add 1 pixel boundary
	g = f;
	g = cat(2, g, reshape(g[:,1,:], size(g, 1), 1, size(g, 3)));
    g = cat(1, g, reshape(g[1,:,:], 1, size(g, 2), size(g, 3)));
	# interpolate
	t = linspace(1, P, N);
	ti = round.(Int64, floor.(t)) ; tj = round.(Int64, ceil.(t));
	fi = round.(Int64, t - floor.(t)); fj = 1 - fi;
	h = zeros(N, N, size(f, 3));
	for s in 1:size(f, 3)
	    h[:,:,s] = g[ti,ti,s] .* (fj * fj') + g[tj,tj,s] .* (fi * fi') + g[ti,tj,s] .* (fj * fi') + g[tj,ti,s] .* (fi * fj');
	end
	return h;
end # resize_image

#####################################
# Function that rescales date in a,b
"""
...
rescale(x::Array,a=0,b=1) rescales data in
[a,b]
...
"""
function rescale(x, a=0, b=1)
	m = minimum(x[:]);
	M = maximum(x[:]);

	if M - m < 1e-10
	    y = x;
	else
	    y = (b - a) * (x - m) / (M - m) + a;
	end
	return y;

end ## rescale

#####################################################
# Function that rescales a wavelet transform in an array A
function rescaleWav(A)
    v = maximum(abs.(A[:]));
    B = copy(A)
    if v > 0
        B = .5 + .5 * A / v;
    end
    return B;
end


const TO = TrajectoryOptimization

# Get the model
model = Dynamics.car
n,m = model.n, model.m
Q = (1e-2)*Diagonal(I,n)
Qf = 1000.0*Diagonal(I,n)
R = (1e-2)*Diagonal(I,m)
x0 = [0.;0.;0.]
xf = [0.;1.;0.]

obj = TrajectoryOptimization.LQRObjective(Q,R,Qf,xf,N)

# Test Expansion
Q = TO.Expansion{Float64}(n,m)
Q.x .= ones(n)
Q.xx .= ones(n,n)
2*Q
@test Q.x == ones(n)*2
Q/2
@test Q.xx == ones(n,n)

Q2 = copy(Q)
Q2.uu .+= 1
@test Q2.uu == ones(m,m)
@test Q.uu == zeros(m,m)
TO.reset!(Q2)
@test Q2.uu == zeros(m,m)

# Quadratic Cost
Q = Diagonal(1.0I,n)
R = Diagonal(1.0I,m)
H = zeros(m,n)
q = ones(n)
r = zeros(m)
Qf = Diagonal(I,n)*100.
qf = zeros(n)
quadcost = QuadraticCost(Q, R, H, q, r, 0.)
quadcost_f = QuadraticCost(Qf, qf, 0.)
@test quadcost.c == 0
@test quadcost.Q == Q
@test quadcost.R isa Diagonal
@test quadcost_f.c == 0
@test quadcost_f.Q == Qf
@test quadcost_f.q == qf

quadcost = QuadraticCost(Q,R)
@test quadcost.R isa Diagonal
@test quadcost.H == zeros(m,n)

# LQR Cost
xf = [0,1,0]
lqrcost = LQRCost(Q,R,xf)
@test lqrcost.Q == Q
@test lqrcost.q == -Q*xf

cost_term = TO.LQRCostTerminal(Qf, xf)
@test cost_term.Q == Qf
@test cost_term.q == -Qf*xf

# Test costs
x,u = rand(n), rand(m)
dt = rand(1)[1]
@test TO.stage_cost(quadcost,x,u,dt) == 0.5*(x'Q*x + u'R*u)*dt
@test TO.stage_cost(cost_term,x) ≈ 0.5*(x-xf)'Qf*(x-xf)

# Test expansion functions
E = TO.Expansion{Float64}(n,m)
TO.cost_expansion!(E, quadcost, x, u, dt)
@test E.x == Q*x*dt
@test E.xx == Q*dt
@test E.ux == zeros(m,n)*dt
@test u == R*u

TO.reset!(E)
TO.cost_expansion!(E, cost_term, x)
@test E.x ≈ Qf*(x-xf)
@test E.u == zeros(m)
@test E.xx == Qf

grad = PartedVector(zeros(n+m), create_partition((n,m),(:x,:u)))
TO.gradient!(grad, quadcost, x, u, dt)
@test grad.x == dt*Q*x
@test grad.u == dt*R*u

grad = zeros(n)
TO.gradient!(grad, cost_term, x)
@test grad ≈ Qf*(x-xf)

quadcost2 = copy(quadcost)
quadcost2.Q .+= Diagonal(I,n)
@test quadcost2.Q == Diagonal(2I,n)
@test quadcost.Q == Q

# Generic Costs
n,m = 2,1
x,u = rand(n), rand(m)
function mycost(x,u)
    R = Diagonal(0.1I,1)
    Q = 0.1
    return cos(x[1]) + u'R*u + Q*x[2]^2
end
function mycost(xN)
    return cos(xN[1]) + xN[2]^2
end
nlcost = GenericCost(mycost,mycost,n,m)
@test TO.stage_cost(nlcost,x,u) == mycost(x,u)
@test TO.stage_cost(nlcost,x) == mycost(x)

R = Diagonal(0.1I,m)
Q = 0.1
function hess(x,u)
    n,m = length(x),length(u)
    Qexp = Diagonal([-cos(x[1]), 2Q])
    Rexp = 2R
    Hexp = zeros(m,n)
    return Qexp,Rexp,Hexp
end
function hess(x)
    return Diagonal([-cos(x[1]), 2])
end
function gradient(x,u)
    q = [-sin(x[1]), 2Q*x[2]]
    r = 2R*u
    return q,r
end
function gradient(x)
    return [-sin(x[1]), 2*x[2]]
end
nlcost2 = GenericCost(mycost, mycost, gradient, hess, n, m)

E = TO.Expansion{Float64}(n,m)
TO.cost_expansion!(E, nlcost, x, u)
@test E.x == gradient(x,u)[1]
@test E.xx == hess(x,u)[1]
@test E.ux == hess(x,u)[3]
TO.cost_expansion!(E, nlcost2, x, u)
@test E.x == gradient(x,u)[1]
@test E.xx == hess(x,u)[1]
@test E.ux == hess(x,u)[3]

nlcost3 = copy(nlcost)
@test TO.stage_cost(nlcost3,x,u) == TO.stage_cost(nlcost,x,u)

# Minimum time
n = 4
m = 2
N = 10
Q = Diagonal(ones(n))
R = Diagonal(ones(m))
Qf = 0.7*Diagonal(ones(n))
xf = rand(n)

_x = rand(n+1)
_u = rand(m+1)
dt = 0.134

part_mt = (x=1:(n+1), u=(n+1) .+ (1:(m+1)))
part2_mt = (xx=(part_mt.x, part_mt.x), uu=(part_mt.u, part_mt.u), ux=(part_mt.u, part_mt.x), xu=(part_mt.x, part_mt.u))
_cost = LQRCost(Q,R,xf)
_cost_term = TO.LQRCostTerminal(Qf,xf)
_cost_mt = TO.MinTimeCost(_cost,1.6)
_cost_term_mt = TO.MinTimeCost(_cost_term,1.6)

EE = TO.Expansion{Float64}(n+1,m+1)
grad_mt = zeros(n+1+m+1)
grad_term_mt = zeros(n+1)
hess_mt = PartedMatrix(zeros(n+1+m+1,n+1+m+1),part2_mt)
hess_term_mt = zeros(n+1,n+1)
TO.cost_expansion!(EE, _cost_mt, _x, _u, dt)
TO.gradient!(grad_mt,_cost_mt,_x,_u,dt)
TO.gradient!(grad_term_mt,_cost_term_mt,_x)
TO.hessian!(hess_mt,_cost_mt,_x,_u,dt)
TO.hessian!(hess_term_mt,_cost_term_mt,_x)

@test isapprox(EE.x,grad_mt[1:(n+1)])
@test isapprox(EE.u,grad_mt[(n+1) .+ (1:(m+1))])
@test isapprox(EE.xx,hess_mt.xx)
@test isapprox(EE.uu,hess_mt.uu)
@test isapprox(EE.ux,hess_mt.ux)


using Base.Test
using SmoothingKernels

for testfile in ["kernels.jl"]
    include(testfile)
end


c0 = chain(:standard,1,10,Float64,Float64)
c1 = chain(:standard,3,4,Int16,Float64)
c2 = chain(:gradient,2,3,Float32,Float32)
@test numparas(c0) == 1 && numsamples(c0) == 10
@test numparas(c1) == 3 && numsamples(c1) == 4
@test numparas(c2) == 2 && numsamples(c2) == 3

s0 = samples(:base,1,4,Float64,Float64)
s1 = samples(:base,3,2,Int16,Float64)
s2 = samples(:tensor,2,2,Float32,Float32)

#test that you can't add a base sample to a gradient chain
@test_throws MethodError store!(c2,s0,1)

#test that the number of parameters is the same for sample and chain
@test_throws AssertionError store!(c0,s1,1)
@test_throws AssertionError store!(c1,s0,1)

copy!(s0.values,1.0:4.0)
store!(c0,s0,4)
store!(c0,s0,3)
store!(c0,s0,2)
@test c0.values[:,1:3] == s0.values[:,4:-1:2]
@test isequal(c0.loglikelihood[1:3],s0.loglikelihood[4:-1:2])

copy!(s1.values,[1.0 4.0;2.0 5.0;3.0 6.0])
store!(c1,s1,2)
store!(c1,s1,1)
store!(c1,s1,2)
store!(c1,s1,2)
@test c1.values == [4.0 1.0 4.0 4.0;5.0 2.0 5.0 5.0;6.0 3.0 6.0 6.0]
@test isequal(c1.loglikelihood,s1.loglikelihood[[2,1,2,2]])

copy!(s2.values,[1.0 3.0;2.0 4.0])
store!(c2,s2,1)
store!(c2,s2,1)
store!(c2,s2,2)
@test c2.values == [1.0 1.0 3.0;2.0 2.0 4.0]
@test isequal(c2.loglikelihood,s2.loglikelihood[[1,1,2]])
@test isequal(c2.gradloglikelihood,s2.gradloglikelihood[:,[1,1,2]])

#these should give warnings that the chain is full
println("Two warnings should be printed below this line")
store!(c1,s1,1)
store!(c2,s2,1)

@test c0.proposed == 3 && accepted!(c0,IndicatorStationary{Float64}(zeros(3),[1,4,3,2])) == 3
@test c1.proposed == 4 && accepted!(c1,IndicatorStationary{Float32}(zeros(Float32,4),[1,2,1,2,2])) == 3
@test c2.proposed == 3 && accepted!(c2,IndicatorStationary{Float64}(zeros(3),[1,1,1,2])) == 1

immutable ChainTestModel <: AbstractModel
    parameters
end

p0 = parameters([0.0],[2.0])
p1 = parameters([0.0,1.0,2.0],[10.0,5.0,3.0])
p2 = parameters([2.0,2.0],[10.0,10.0])
m0 = ChainTestModel(p0)
m1 = ChainTestModel(p1)
m2 = ChainTestModel(p2)

@test_approx_eq logprior(c0,m0) logprior(p0,c0.values,Float64)
@test_approx_eq logprior(c1,m1) logprior(p1,c1.values,Float64)
@test_approx_eq logprior(c2,m2) logprior(p2,c2.values,Float64)

@test_approx_eq logposterior(c0,m0) loglikelihood(c0) + logprior(c0,m0)
@test_approx_eq logposterior(c1,m1) loglikelihood(c1) + logprior(c1,m1)
@test_approx_eq logposterior(c2,m2) loglikelihood(c2) + logprior(c2,m2)

###Test show()
println()
println("====================")
println("Test show() function")
show(c0)
show(c1)
show(c2)
println("End  show() function")
println("====================")
println()
nothing


using IntcodeVM, DelimitedFiles, Memoize

@memoize function gravity(x::Int,y::Int)
    return last(last(IntcodeVM.run(code, inputs=[x,y])))
end

const code = vec(readdlm("input.txt", ',', Int, '\n'))
tractor_beam = [gravity(x,y) for x=0:49, y=0:49]

print("Day 19 Part 1: ")
println(count(==(1), tractor_beam))

# It would be computationally crazy to try to brute force this.
# This looks like something solvable through linear regression.

# We can perform a linear search for values xmin <= x <= xmax where gravity(x,y)==1
# and xmin + 100 == xmax.
# For such values, we don't actually need to search the entire square,
# just (xmin, y), (xmax, y), (xmin, y + 100).
# In fact, that might be lightweight enough to compute through
# brute force.

# We can look at the graphic and see that the solution is going to align
# to the right. So, we should probably find a pure function y=f(xmax) to
# predict where the right side of the triangle is.
# So there you have it: run linear regression along right side.

@memoize function get_bounds(y)
    y < 6 && return (0,0)

    # The current xmin value can never be less than xmin value at y-1.
    xmin = first(get_bounds(y-1))
    while gravity(xmin, y) == 0
        xmin += 1
    end

    # Same for xmax, we can never be to the left of the right bound at y-1.
    xmax = max(xmin, last(get_bounds(y-1)))
    while gravity(xmax + 1, y) == 1
        xmax += 1
    end
    return xmin, xmax
end

function print_beam(a)
    xmax,ymax = size(a)
    for y=1:ymax
        for x=1:xmax
            print(a[x,y])
        end
        # 1-indexed languages are so annoying sometimes.
        y >= 7 && print(" y=$(y-1): bounds are $(get_bounds(y - 1))")
        println()
    end
end

print_beam(tractor_beam)

# This is a linear programming question.
# We run a linear regression to predict what the left and right sides of our bounds are.
#using DataFrames, GLM;
#data_left = DataFrame(X=6:49, Y=[first(get_bounds(i)) for i=6:49])
#ols_left = lm(@formula(Y ~ X), data_left)
#data_right = DataFrame(X=6:49, Y=[last(get_bounds(i)) for i=6:49])
#ols_right = lm(@formula(Y ~ X), data_right)

# We can see from here that the right prediction is extremely accurate.
# To my disappointment, I now realize the formula is just right_bound = row - 1
#hcat([get_bounds(i) for i=6:49], predict(ols_left), predict(ols_right))
# D'oh, something goes wrong at y=216 and y=217! The right bound does not linearly increase between these rows.

function find_square(square_size::Int, search_space::UnitRange{Int})
    # Ahh, my old friend off-by-one errors.
    # We have to include the current location x or y in a square of size n.
    dy = square_size - 1
    for row=search_space
        b1 = get_bounds(row)
        b2 = get_bounds(row + dy)
        #println("Bounds(row $(row)) = $(b1), Bounds(row $(row + dy)) = $(b2)")
        #println(b1[2] - b2[1])
        if b1[2] - b2[1] == dy
            return (b2[1], row)
        end
    end
    return (-1,-1)
end

(x,y) = find_square(100, 1:10000)
println("Day 19 Part 2: $(x * 10000 + y)")


using GraphicalModelLearning
using Ipopt
using JuMP

using Test
using Random

import DelimitedFiles: readdlm
import LinearAlgebra: diag

const SOLVER = with_optimizer(Ipopt.Optimizer, print_level=0)

include("common.jl")

@testset "GraphicalModelLearning" begin

@testset "factor graphs" begin
    for (name, gm) in gms
        matrix = convert(Array{Float64,2}, gm)
        gm2 = FactorGraph(matrix)
        for key in keys(gm)
            @test isapprox(gm[key], gm2[key])
            if length(key) == 1
                @test isapprox(gm[key], matrix[key..., key...])
            else
                @test isapprox(gm[key], matrix[key...])
            end
        end
    end
end

@testset "gibbs sampler" begin
    for (name, gm) in gms
        gm_tmp = deepcopy(gm)
        gm_tmp.order = 3
        Random.seed!(0) # fix random number generator
        samples = sample(gm_tmp, gibbs_test_samples)
        #base_samples = readdlm("data/$(name)_samples.csv", ',', Int64)
        base_samples = readdlm("data/$(name)_samples.csv", ',')
        @test isapprox(samples, base_samples)
    end
end

@testset "gibbs sampler, 2nd order" begin
    for (name, gm) in gms
        Random.seed!(0) # fix random number generator
        samples = sample(gm, gibbs_test_samples)
        base_samples = readdlm("data/$(name)_samples.csv", ',')
        #println(name)
        #println(base_samples)
        #println(samples)
        #println(abs.(base_samples-samples))
        @test isapprox(samples, base_samples)
    end
end

@testset "gibbs sampler with replicates" begin
    for (name, gm) in gms
        Random.seed!(0) # fix random number generator
        samples_set = sample(gm, gibbs_test_samples, 3)
        for samples in samples_set
            @test sum(samples[:,1]) == gibbs_test_samples
        end
    end
end

@testset "inverse ising formulations" begin

    for (form_name, formulation) in formulations
        @testset "  $(form_name)" begin
            for (name, gm) in gms
                samples = readdlm("data/$(name)_samples.csv", ',')
                Random.seed!(0) # fix random number generator
                #learned_gm = inverse_ising(samples, method=formulation)
                learned_gm = learn(samples, formulation)
                base_learned_gm = readdlm("data/$(name)_$(form_name)_learned.csv", ',')
                #println(maximum(abs.(learned_gm - base_learned_gm)))
                @test isapprox(learned_gm, base_learned_gm)
            end
        end
    end

    samples = readdlm("data/mvt_samples.csv", ',')

    Random.seed!(0) # fix random number generator
    learned_gm_rise = learn(samples, RISE(0.2, false), NLP(SOLVER))
    base_learned_gm = readdlm("data/mvt_RISE_learned.csv", ',')
    #println(abs.(learned_gm_rise - base_learned_gm))
    @test isapprox(learned_gm_rise, base_learned_gm)

    Random.seed!(0) # fix random number generator
    learned_gm_logrise = learn(samples, logRISE(0.2, false), NLP(SOLVER))
    base_learned_gm = readdlm("data/mvt_logRISE_learned.csv", ',')
    #println(abs.(learned_gm_logrise - base_learned_gm))
    @test isapprox(learned_gm_logrise, base_learned_gm)

    Random.seed!(0) # fix random number generator
    learned_gm_rple = learn(samples, RPLE(0.2, false), NLP(SOLVER))
    base_learned_gm = readdlm("data/mvt_RPLE_learned.csv", ',')
    #println(abs.(learned_gm_rple - base_learned_gm))
    @test isapprox(learned_gm_rple, base_learned_gm)
end

accuracy_tests = [
    AccuracyTest(RISE(),     1000, 0.15)
    AccuracyTest(logRISE(),  1000, 0.15)
    AccuracyTest(RPLE(),     1000, 0.15)
    AccuracyTest(RISE(),    10000, 0.05)
    AccuracyTest(logRISE(), 10000, 0.05)
    AccuracyTest(RPLE(),    10000, 0.05)
]

Random.seed!(0) # fix random number generator
@testset "learned model accuracy" begin
    for act in accuracy_tests
        @testset "  $(act.formulation) $(act.samples) $(act.threshold)" begin
            for (gm_name, gm) in gms
                sample_histo = sample(gm, act.samples)
                #learned_gm = inverse_ising(sample_histo, method=act.formulation)
                learned_gm = learn(sample_histo, act.formulation)
                max_error = maximum(abs.(convert(Array{Float64,2}, gm) - learned_gm))
                @test max_error <= act.threshold
            end
        end
    end
end

@testset "inverse multi-body formulations" begin

    for (name, gm) in gms
        samples = readdlm("data/$(name)_samples.csv", ',')
        Random.seed!(0) # fix random number generator
        learned_ising = learn(samples, RISE(0.2, false))
        learned_two_body = learn(samples, multiRISE(0.2, false, 2))

        learned_ising_dict = convert(Dict, learned_ising)
        #println(learned_ising_dict)
        #println(learned_two_body)

        @test length(learned_ising_dict) == length(learned_two_body)
        for (key, value) in learned_ising_dict
            @test isapprox(learned_two_body[key], value, atol=1e-7)
        end
    end

    samples = readdlm("data/mvt_samples.csv", ',')

    rand(0) # fix random number generator
    learned_ising = learn(samples, RISE(0.2, false), NLP(SOLVER))
    learned_two_body = learn(samples, multiRISE(0.2, false, 2), NLP(SOLVER))

    learned_ising_dict = convert(Dict, learned_ising)
    @test length(learned_ising_dict) == length(learned_two_body)
    for (key, value) in learned_ising_dict
        @test isapprox(learned_two_body[key], value)
    end

    for (name, gm) in gms
        gm_tmp = deepcopy(gm)
        gm_tmp.order = 4
        Random.seed!(0) # fix random number generator
        samples = sample(gm_tmp, 10000)

        learned_gm = learn(samples, multiRISE(0.0, false, 4))

        for (key, value) in gm_tmp
            #println(learned_gm[key], value)
            @test isapprox(learned_gm[key], value, atol = 0.15)
        end

        samples = sample(learned_gm, 10000)
        learned_gm2 = learn(samples, multiRISE(0.0, false, 4))

        for (key, value) in gm_tmp
            #println(learned_gm2[key], " ", value)
            # this is a bug, the learned_gm2 value should not be twice as large
            @test isapprox(learned_gm2[key]/2.0, value, atol = 0.15)
        end

    end
end

Random.seed!(0) # fix random number generator
@testset "docs example" begin
    model = FactorGraph([0.0 0.1 0.2; 0.1 0.0 0.3; 0.2 0.3 0.0])
    samples = sample(model, 100000)
    learned = learn(samples)

    err = abs.(convert(Array{Float64,2}, model) - learned)
    @test maximum(err) <= 0.01
end

end


######################################################################################################
# Merge from                                                                                #
# * Wouter Koolen: https://bitbucket.org/wmkoolen/tidnabbil/src/master/purex_games_paper/, and       #
# * Xuedong Shang: https://github.com/xuedong/LinBAI.jl                                              #                                      #
######################################################################################################
struct FTL
    L :: Array{Float64, 1};

    function FTL(K)
        new(zeros(K));
    end
end
function act(h::FTL)
    u = h.L .== minimum(h.L);
    u ./ sum(u);
end
function incur!(h::FTL, l)
    h.L .+= l;
end

mutable struct AdaHedge
    L :: Array{Float64, 1};
    Δ :: Float64;

    function AdaHedge(K)
        new(zeros(K), 0.01);
    end
end
function act(h::AdaHedge)
    η = log(length(h.L))/h.Δ;
    u = exp.(-η.*(h.L .- minimum(h.L)));
    u ./ sum(u);
end
function incur!(h::AdaHedge, l)
    K = length(l);
    u = act(h);
    η = log(length(h.L))/h.Δ;
    M(L) = η==0 ? sum(L)/K : minimum(L) - 1/η*log(sum(exp.(-η.*(L .- minimum(L))))/K);
    Mpre = M(h.L);
    h.L .+= l;
    Mpst = M(h.L);
    m = Mpst-Mpre;

    @assert u'l ≥ m - 1e-7 "hedge loss $(u'l), Δ=$(h.Δ), η=$η, mix loss $m, u=$u, L=$(h.L), l=$(l)";
    h.Δ  += u'l - m;
end

######################################################################################################
# Implementation by Xuedong Shang: https://github.com/xuedong/LinBAI.jl                              #                                                                     #
######################################################################################################
mutable struct LinBAIAdaHedge
    L :: Array{Float64, 1};
    Δ :: Float64;

    function LinBAIAdaHedge(K)
        new(zeros(K), 0.01);
    end
end
function act(h::LinBAIAdaHedge)
    η = log(length(h.L))/h.Δ;
    u = exp.(-η.*(h.L .- minimum(h.L)));
    u ./ sum(u);
end
function incur!(h::LinBAIAdaHedge, l)
    u = act(h);
    η = log(length(h.L))/h.Δ;
    h.L .+= l;
    m = minimum(l) - 1/η*log(u'exp.(-η.*(l .- minimum(l))));
    # println(η, " ", u, " ", u'l, " ", -1/η*log(u'exp.(-η.*l)));
    @assert isfinite(m) && u'l ≥ m - 1e-7 "hedge loss $(u'l), Δ=$(h.Δ), η=$η, mix loss $m, u=$u, L=$(h.L)";
    h.Δ  += u'l - m;
end

# Menard calls this "Lazy Mirror Descent" As perhaps the main feature
# here is the mixing step, we call it "Fixed Share".  This is the
# version with slowly decreasing O(1/sqrt(t)) learning and switching
# rates.  The switching is used to force exploration
mutable struct FixedShare
    w :: Array{Float64, 1};
    t; # time
    S; # loss range: S ≥ \max_k l_k - \min_k l_k
    function FixedShare(K; S=1)
        new(fill(1/K,K), 0, S);
    end
end
function act(h::FixedShare)
    h.w
end
function incur!(h::FixedShare, l)
    K = length(l);
    h.t += 1;
    # Exponential Weights update
    η = sqrt(log(K)/h.t)/h.S;   # decaying learning rate
    h.w .*= exp.(-η*(l .- minimum(l)));
    h.w ./= sum(h.w);
    # Fixed Share (to unifor prior)
    γ = 1/(4*sqrt(h.t)); # decaying switching rate
    h.w .= (1-γ).*h.w .+ γ/K;
end


using vOptGeneric, MOBC, PyPlot, CPLEX

function parseSSCFLP(fname, solver = CplexSolver(CPX_PARAM_SCRIND = 0))
    line_to_vec(f) = parse.(Int, split(chomp(readline(f))))
    file = open(fname)
    nbfac, nbclient = line_to_vec(file)
    I = 1:nbfac
    J = 1:nbclient
    readline(file)
    s, f = Int[], Int[]
    for i = 1:nbfac
        si, fi = line_to_vec(file)
        push!(s, si)
        push!(f, fi)
    end
    readline(file)
    d = line_to_vec(file)
    readline(file)
    c = readdlm(file, Int)
    close(file)

    m = vModel(solver=solver)
    @variable(m, x[1:nbfac, 1:nbclient], Bin)
    @variable(m, y[1:nbfac], Bin)
    @addobjective(m, Min, vecdot(c, x))
    @addobjective(m, Min, dot(f, y))
    @constraint(m, [j=J], sum(x[i,j] for i ∈ I) == 1)
    @constraint(m, [i=I], sum(d[j]*x[i,j] for j ∈ J) <= s[i]*y[i])

    m
end

for folder in ["size_5_10", "size_10_20"]#, "size_15_30"]
    cd(joinpath(Pkg.dir("MOBC"), "test", "Instances_SSCFLP", folder)) do
        for file in readdir()
            m = parseSSCFLP(file)
            println(file)
            @time solve_stidsen(m, use_nsga=true, global_branch=false, docovercuts=true, lift_covers=true)
            @time solve_parragh(m, use_nsga=true, global_branch=false, docovercuts=true, lift_covers=true)
            @time solve(m, method=:epsilon, round_results=true, verbose=false, suppress_warnings=true)
        end
    end
end


using CayleyDickson
using Base.Test: @test, @test_throws

@test begin
    l = Exo2SplitBinion(Exo2Real(1,2,3,4))
    r = Exo2SplitBinion(1,0,2,0,3,0,4)
    l == r
end

@test begin
    a = CayleyDickson.randomBigFloat()
    isreal(Exo2SplitBinion(a))
end

@test begin
    a = 1
    !isreal(Exo2SplitBinion(a, a, a, a))
end

@test begin
    z = random(Exo2SplitBinion{Int})
    z == +(z)
end

@test begin
    a = CayleyDickson.randomBigFloat()
    real(Exo2SplitBinion(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, Exo2SplitBinion(1.0, 2, 3, 4, 5, 6, 7, 8))
    l = String(take!(io))
    r = "[1: 1.0, s: 2.0, a: 3.0, sa: 4.0, b: 5.0, sb: 6.0, ab: 7.0, (sa)b: 8.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2SplitBinion(1.0, 2, 3, 4, 5, 6, 7))
    l = String(take!(io))
    r = "[1: 1.0, s: 2.0, a: 3.0, sa: 4.0, b: 5.0, sb: 6.0, ab: 7.0, (sa)b: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2SplitBinion(1.0, 2, 3, 4, 5, 6))
    l = String(take!(io))
    r = "[1: 1.0, s: 2.0, a: 3.0, sa: 4.0, b: 5.0, sb: 6.0, ab: 0.0, (sa)b: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2SplitBinion(1.0, 2, 3, 4, 5))
    l = String(take!(io))
    r = "[1: 1.0, s: 2.0, a: 3.0, sa: 4.0, b: 5.0, sb: 0.0, ab: 0.0, (sa)b: 0.0]"
    l == r
end

@test begin
    io = IOBuffer()
    show(io, Exo2SplitBinion(1.0, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1.0, s: 2.0, a: 3.0, sa: 4.0, b: 0.0, sb: 0.0, ab: 0.0, (sa)b: 0.0]"
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    d = CayleyDickson.randomBigFloat()
    l = Exo2SplitBinion(Exo1SplitBinion(a, b, c, d))
    r = Exo2SplitBinion(a, b, c, d)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    l = Exo2SplitBinion(SplitBinion(a, b))
    r = Exo2SplitBinion(a, b)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    b = CayleyDickson.randomBigFloat()
    c = CayleyDickson.randomBigInt()
    l = Exo2SplitBinion(a, b, c)
    r = Exo2SplitBinion(a, b, c, 0)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigFloat})
    y = random(Exo2SplitBinion{BigFloat})
    x != y
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo2SplitBinion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = random(Exo1SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    z = random(Exo2SplitBinion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    a = random(Exo1SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = x * y
    r = y * x
    l != r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l != r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = alternatorL(x, y)
    r = zero(Exo2SplitBinion{BigInt})
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = alternatorR(x, y)
    r = zero(Exo2SplitBinion{BigInt})
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    conj(conj(x)) == x
end

@test begin
    x = random(Exo2SplitBinion{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = conj(x * y)
    r = conj(y) * conj(x)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = conj(x + y)
    r = conj(x) + conj(y)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = conj(x - y)
    r = conj(x) - conj(y)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{Rational{BigInt}})
    y = random(Exo2SplitBinion{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{Rational{BigInt}})
    y = random(Exo2SplitBinion{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Exo2SplitBinion{Rational{BigInt}})
    y = random(Exo2SplitBinion{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = CayleyDickson.randomBigInt()
    x = random(Exo2SplitBinion{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    a = random(Exo1SplitBinion{BigInt})
    x = random(Exo2SplitBinion{BigInt})
    l = commutator(a, x)
    r = zero(Exo2SplitBinion{BigInt})
    l != r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    z = random(Exo2SplitBinion{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    quadrance(x) > big(0) || quadrance(x) < big(0) ||quadrance(x) == big(0)
end

@test begin
    x = random(Exo2SplitBinion{BigInt})
    y = random(Exo2SplitBinion{BigInt})
    l = quadrance(x * y)
    r = quadrance(x) * quadrance(y)
    l == r
end

@test_throws ErrorException begin
    inv(Exo2SplitBinion(0, 0.0))
end

@test_throws ErrorException begin
    inv(Exo2SplitBinion(1, 1))
end

@test_throws ErrorException begin
    random(Exo2SplitBinion{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Exo2SplitBinion{Int})
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo2SplitBinion{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = CayleyDickson.randomBigRational()
    z = random(Exo2SplitBinion{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Exo2SplitBinion{BigInt})
    z == z + zero(Exo2SplitBinion{BigInt})
end

@test begin
    z = random(Exo2SplitBinion{BigInt})
    z == z * one(Exo2SplitBinion{BigInt})
end

@test begin
    z = random(Exo2SplitBinion{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Exo2SplitBinion{BigInt})
    z == z * one(z)
end


function moi_01(optimizer)

    # original problem
    # min -4x -3y
    # s.t.
    #      2x + y <= 4
    #       x +2y <= 4
    #       x     >= 0
    #           y >= 0
    #
    # sol x = y = 4/3
    # obj = 28/3

    MOI.empty!(optimizer)
    @test MOI.is_empty(optimizer)

    # add 10 variables - only diagonal is relevant
    X = MOI.add_variables(optimizer, 2)

    c1 = MOIU.normalize_and_add_constraint(optimizer,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(2.0, X[1]),
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    c2 = MOIU.normalize_and_add_constraint(optimizer,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1]),
            MOI.ScalarAffineTerm(2.0, X[2])
        ], 0.0), MOI.EqualTo(4.0))

    b1 = MOIU.normalize_and_add_constraint(optimizer,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[1])
        ], 0.0), MOI.GreaterThan(0.0))

    b2 = MOIU.normalize_and_add_constraint(optimizer,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X[2])
        ], 0.0), MOI.GreaterThan(0.0))

    MOI.set(optimizer,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([-4.0, -3.0], [X[1], X[2]]), 0.0)
        )
    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    MOI.optimize!(optimizer)

    obj = MOI.get(optimizer, MOI.ObjectiveValue())

    @test obj ≈ -9.33333 atol = 1e-2

    Xr = MOI.get(optimizer, MOI.VariablePrimal(), X)

    @test Xr ≈ [1.3333, 1.3333] atol = 1e-2

end

function moi_02(optimizer, mode = BilevelJuMP.SOS1Mode())

    # min -4x -3y
    # s.t.
    # y = argmin_y y
    #      2x + y <= 4
    #       x +2y <= 4
    #       x     >= 0
    #           y >= 0
    #
    # sol: x = 2, y = 0
    # obj_upper = -8
    # obj_lower =  0

    ### Build upper level model

    upper = MOIU.Model{Float64}()
    @test MOI.is_empty(upper)

    X = MOI.add_variables(upper, 2)

    MOI.set(upper, MOI.VariableName(), X[1], "x_u")
    MOI.set(upper, MOI.VariableName(), X[2], "y_u")

    MOI.set(upper,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([-4.0, -3.0], [X[1], X[2]]), 0.0)
        )
    MOI.set(upper, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    ### Build lower level model

    lower = MOIU.Model{Float64}()
    @test MOI.is_empty(lower)

    X2 = MOI.add_variables(lower, 2)
    MOI.set(lower, MOI.VariableName(), X2[1], "x_l")
    MOI.set(lower, MOI.VariableName(), X2[2], "y_l")

    c1 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(2.0, X2[1]),
            MOI.ScalarAffineTerm(1.0, X2[2])
        ], 0.0), MOI.LessThan(4.0))
    MOI.set(lower, MOI.ConstraintName(), c1, "l_c1")

    c2 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X2[1]),
            MOI.ScalarAffineTerm(2.0, X2[2])
        ], 0.0), MOI.LessThan(4.0))
    MOI.set(lower, MOI.ConstraintName(), c2, "l_c2")

    b1 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X2[1])
        ], 0.0), MOI.GreaterThan(0.0))
    MOI.set(lower, MOI.ConstraintName(), b1, "l_b1")

    b2 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, X2[2])
        ], 0.0), MOI.GreaterThan(0.0))
    MOI.set(lower, MOI.ConstraintName(), b2, "l_b12")

    MOI.set(lower,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0], [X2[2]]), 0.0)
        )
    MOI.set(lower, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    ### Build bilevel model combining both models

    links = Dict(X[1] => X2[1], X[2] => X2[2])

    parametric = [X2[1]]

    blp, _, _ = BilevelJuMP.build_bilevel(upper, lower, links, parametric, mode)

    # MOI.optimize!(bridged)

    # obj = MOI.get(bridged, MOI.ObjectiveValue())

    # @test obj ≈ -9.33333 atol = 1e-2

    # Xr = MOI.get(bridged, MOI.VariablePrimal(), X)

    # @test Xr ≈ [1.3333, 1.3333] atol = 1e-2

    # lp_model = MathOptFormat.MOF.Model()
    # MOI.copy_to(lp_model, blp)
    # MOI.write_to_file(lp_model, "my_model.LP")
    # @show pwd()

    MOI.empty!(optimizer)
    @test MOI.is_empty(optimizer)
    MOI.copy_to(optimizer, blp)

    MOI.optimize!(optimizer)

end

function moi_03(optimizer, mode = BilevelJuMP.SOS1Mode())

    upper = MOIU.Model{Float64}()
    @test MOI.is_empty(upper)

    x = MOI.add_variable(upper)
    y = MOI.add_variable(upper)

    MOI.set(upper, MOI.VariableName(), x, "x_u")
    MOI.set(upper, MOI.VariableName(), y, "y_u")

    MOI.set(upper,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([1.0, 3.0], [x, y]), 0.0)
        )
    MOI.set(upper, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    lower = MOIU.Model{Float64}()
    @test MOI.is_empty(lower)

    x2 = MOI.add_variable(lower)
    y2 = MOI.add_variable(lower)

    MOI.set(lower, MOI.VariableName(), x2, "x_l")
    MOI.set(lower, MOI.VariableName(), y2, "y_l")

    c1 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2),
            MOI.ScalarAffineTerm(1.0, y2)
        ], 0.0), MOI.LessThan(8.0))
    MOI.set(lower, MOI.ConstraintName(), c1, "l_c1")

    c2 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2),
            MOI.ScalarAffineTerm(4.0, y2)
        ], 0.0), MOI.GreaterThan(8.0))
    MOI.set(lower, MOI.ConstraintName(), c2, "l_c2")

    c3 = MOIU.normalize_and_add_constraint(lower,
    MOI.ScalarAffineFunction([
        MOI.ScalarAffineTerm(1.0, x2),
        MOI.ScalarAffineTerm(2.0, y2)
    ], 0.0), MOI.LessThan(13.0))
    MOI.set(lower, MOI.ConstraintName(), c3, "l_c3")

    b1 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2)
        ], 0.0), MOI.GreaterThan(1.0))
    MOI.set(lower, MOI.ConstraintName(), b1, "l_b1")

    b2 = MOIU.normalize_and_add_constraint(lower,
        MOI.ScalarAffineFunction([
            MOI.ScalarAffineTerm(1.0, x2)
        ], 0.0), MOI.LessThan(6.0))
    MOI.set(lower, MOI.ConstraintName(), b2, "l_b12")

    MOI.set(lower,
        MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(),
        MOI.ScalarAffineFunction(
            MOI.ScalarAffineTerm.([-1.0], [y2]), 0.0)
        )
    MOI.set(lower, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    links = Dict(x => x2, y => y2)

    parametric = [x2]

    blp, _, _ = BilevelJuMP.build_bilevel(upper, lower, links, parametric, mode)

    # MOI.optimize!(bridged)

    # obj = MOI.get(bridged, MOI.ObjectiveValue())

    # @test obj ≈ -9.33333 atol = 1e-2

    # Xr = MOI.get(bridged, MOI.VariablePrimal(), X)

    # @test Xr ≈ [1.3333, 1.3333] atol = 1e-2

    # lp_model = MathOptFormat.MOF.Model()
    # MOI.copy_to(lp_model, blp)
    # MOI.write_to_file(lp_model, "my_model.LP")
    # @show pwd()

    MOI.empty!(optimizer)
    @test MOI.is_empty(optimizer)
    MOI.copy_to(optimizer, blp)

    MOI.optimize!(optimizer)

end


# A Frisch-Newton algorithm as described in http://projecteuclid.org/euclid.ss/1030037960
# qreg_ip_coef is a modified lp_fnm of Daniel Morillo & Roger Koenker
# Translated from Ox to Matlab by Paul Eilers 1999
# Modified by Roger Koenker 2000
# More changes by Paul Eilers 2004
# QuantileRegressions.jl developers 2015
#
# The functions in this file is released under the BSD-3 license, although the original
# versions were not. Permissing was given to Patrick Kofod Mogensen by Roger Koenker and
# Paul Eilers on April 24 2015 via E-mail. Original mail correspondences can be provided
# upon request.

struct IP <: Solver
    cholesky::Bool
end
IP() = IP(false)

function bound!(d, x, dx)
# Fill vector with allowed step lengths
# Replace with -x/dx for negative dx
    d .= 1e20
    @inbounds for i = 1:length(dx)
        if dx[i] < 0.0
            d[i] = -x[i] / dx[i]
        end
    end
    return d
end

function qreg_coef(Y, X::Matrix, p, method::IP)
# input: X is an n x k matrix of exogenous regressors,
#        Y is an n x 1 vector of outcome variables
#        p \in (0,1) is the quantile of interest
# Output: p^{th} regression quantiles.
# Construct the dual problem of quantile regression
c = -Y
T = eltype(X)

m = size(X, 1)

x = fill((1 - p), m)
b = X'x
# Solve a linear program by the interior point method:
# min(c * u), s.t. A * x = b and 0 < x < u
# An initial feasible solution has to be provided as x

# Set some constants
beta = 0.99995
small = sqrt(eps(eltype(c)))
max_it = 500
n, m = size(X)
# Generate inital feasible point
s = 1 .- x
if method.cholesky
    y = -((X'X)\(X'Y))
else
    y = -X\Y
end
dy = copy(y)
r = c - X*y
# z and w are the "residuals" but on either
# side of the "check" we need to flip the sign
#    BLAS.axpy!(small, (r .== 0.0).*1.0, r)
#    z = r .* (r .> small)
#    w = abs.(z .- r)
z = copy(r)
w = copy(r)
for i ∈ eachindex(r)
    ri = r[i]
    z[i] = max(ri, 0)
    w[i] = max(-ri, 0)
    if abs(ri) <= small
        z[i] = z[i] + small
        w[i] = w[i] + small
    end
end

gap = dot(z, x) + dot(s, w)

# set up caches
Xtmp = copy(X)
if method.cholesky
    XtX  = similar(X, m, m)
end
xinv, xi = copy(x), copy(x)
sinv = copy(s)
q = copy(z)
dx, ds, dz, dw = copy(w), copy(w), copy(w), copy(w)
fx, fs, fz, fw = copy(w), copy(w), copy(w), copy(w)
dxdz, dsdw = copy(w), copy(w)
tmp = copy(w)
if method.cholesky
    rtmp = copy(r)
    Xtqr = similar(r, m)
else
    qr_tmp = copy(q)
    rhs = copy(r)
end
# Start iterations
for it = 1:max_it
    #   Compute affine step
    @. q = 1 / (z / x + w / s)
    @. r = z - w

    if method.cholesky
        @. Xtmp = q * X
        mul!(XtX, Xtmp', X)
        F = cholesky!(Symmetric(XtX))
        mul!(Xtqr, Xtmp', r)
        dy .= F\Xtqr
        #ldiv!(dy, F, Xtqr)
    else
        @. qr_tmp = sqrt(q)
        Q = Diagonal(qr_tmp) # Very efficient to do since Q diagonal
        AQtF = qr(mul!(Xtmp, Q, X)) # PE 2004
        @. rhs = qr_tmp*r
        dy .= AQtF\rhs
    end

    mul!(tmp, X, dy)
    @. dx = q*(tmp - r)
    @. ds = -dx
    @. dz = -z * (1 + dx / x)
    @. dw = -w * (1 + ds / s)
    # Compute maximum allowable step lengths
    bound!(fx, x, dx)
    bound!(fs, s, ds)
    bound!(fw, w, dw)
    bound!(fz, z, dz)
    @. tmp = min(fx, fs)
    fp = min(beta*minimum(tmp), 1)
    @. tmp = min(fw, fz)
    fd = min(beta*minimum(tmp), 1)
    # If full step is feasible, take it. Otherwise modify it
    if min(fp, fd) < 1.0
        # Update mu
        g0 = dot(z, x) + dot(w, s)
        gγPγD = T(0)
        for i = 1:length(z)
          @inbounds gγPγD += (z[i] + fd*dz[i])*(x[i] + fp*dx[i]) + (w[i] + fd*dw[i])*(s[i]+fp*ds[i])
        end
        mu = (gγPγD / g0)^3 * g0/(2n)
        # Compute modified step
        @. dxdz = dx * dz
        @. dsdw = ds * dw
        @. xinv = 1 / x
        @. sinv = 1 / s
        @. xi = mu * (xinv - sinv)
        if method.cholesky
            @. rtmp = r + dxdz - dsdw - xi
            mul!(Xtqr, Xtmp', rtmp)
            dy .= F\Xtqr
        else
            @. qr_tmp = qr_tmp * (dxdz .- dsdw .- xi)
            BLAS.axpy!(1.0, qr_tmp, rhs) # no gemv-wrapper gemv(Q, (dxdz - dsdw - xi), rhs,1,1,n)?
            dy .= AQtF\rhs
        end

        mul!(tmp, X, dy)
        @. dx = q * (tmp + xi - r - dxdz + dsdw)
        @. ds = -dx
        @. dz = (mu - z * dx)*xinv - z - dxdz
        @. dw = (mu - w * ds)*sinv - w - dsdw

        # Compute maximum allowable step lengths
        bound!(fx, x, dx)
        bound!(fs, s, ds)
        bound!(fw, w, dw)
        bound!(fz, z, dz)
        @. tmp = min(fx, fs)
        fp = min(beta*minimum(tmp), 1)
        @. tmp = min(fw, fz)
        fd = min(beta*minimum(tmp), 1)
    end

    # Take the steps
    BLAS.axpy!(fp, dx, x)
    BLAS.axpy!(fp, ds, s)
    BLAS.axpy!(fd, dy, y)
    BLAS.axpy!(fd, dw, w)
    BLAS.axpy!(fd, dz, z)

    gap = dot(z, x) + dot(s, w)

    if gap < small || isnan(gap)
        break
    end
end

return -y
end


using GeometricBase.Config
using GeometricBase.Utils
using GeometricIntegrators.Integrators
using GeometricIntegrators.Solutions
using GeometricIntegrators.Tableaus
using GeometricProblems.LotkaVolterra2d
using SimpleSolvers
using Test

SimpleSolvers.set_config(:nls_atol, 8eps())
SimpleSolvers.set_config(:nls_rtol, 2eps())

const Δt = 0.01
const nt = 10
const q₀ = [1.0, 1.0]
const parameters = (a₁=1.0, a₂=1.0, b₁=-1.0, b₂=-2.0)

ode  = lotka_volterra_2d_ode(q₀; params=parameters)
iode = lotka_volterra_2d_iode(q₀; params=parameters)
lode = lotka_volterra_2d_lode(q₀; params=parameters)

int  = IntegratorFIRK(ode, TableauGauss(8), Δt)
sol  = integrate(ode, int, nt)
refx = sol.q[end]

@testset "$(rpad("1st Order DVIs",80))" begin

    sol = Solution(iode, Δt, nt)
    int = IntegratorDVIA(iode, Δt)
    integrate!(int, sol)
    @test relative_maximum_error(sol.q, refx) < 1E-1

    sol = Solution(iode, Δt, nt)
    int = IntegratorDVIB(iode, Δt)
    integrate!(int, sol)
    @test relative_maximum_error(sol.q, refx) < 1E-1

end

@testset "$(rpad("2nd Order Centred DVIs",80))" begin

    sol = Solution(iode, Δt, nt)
    int = IntegratorCMDVI(iode, Δt)
    integrate!(int, sol)
    @test relative_maximum_error(sol.q, refx) < 4E-3

    sol = Solution(iode, Δt, nt)
    int = IntegratorCTDVI(iode, Δt)
    integrate!(int, sol)
    @test relative_maximum_error(sol.q, refx) < 4E-3

end


@testset "MixMatRep" begin
    change_fulldim_test(MixedMatHRep)
    change_fulldim_test(MixedMatVRep)

    isempty_vtest(vrep(zeros(0, 3)), 0, 0)
    isempty_vtest(vrep(zeros(1, 2)), 0, 1)
    isempty_vtest(vrep(zeros(1, 4), ones(2, 4)), 2, 1)
    isempty_vtest(vrep(zeros(2, 1), ones(1, 1)), 1, 2)

    isempty_htest(hrep(ones(0, 3), zeros(0)), 0, 0)
    isempty_htest(hrep(ones(1, 2), zeros(1)), 0, 1)
    isempty_htest(hrep(ones(2, 4), zeros(2), BitSet(1:2)), 2, 0)
    isempty_htest(hrep(ones(3, 1), zeros(3), BitSet([2])), 1, 2)

    @testset "with bad arguments" begin
        A = [1 1; -1 0; 0 -1]
        b = [1, 0, 0]
        linset = BitSet([1])

        @test_throws ErrorException hrep(A, [0, 0], linset)
        @test_throws ErrorException hrep(A, b, BitSet([4]))
        ine = hrep(A, b, linset)
        @test fulldim(ine) == 2
        @test (@inferred Polyhedra.FullDim(ine)) == 2
        @test Polyhedra.coefficient_type(ine) == Int
        @test translate(ine, [1, 0]).b == [2, -1, 0]

        V = [0 1; 1 0]
        @test_throws ErrorException vrep(zeros(0, 2), [1 0]) # V-consistency
        @test_throws ErrorException vrep(V, [1 0 0], BitSet())
        @test_throws ErrorException vrep(V, [1 1], BitSet([2]))
        ext = vrep(V)
        @test fulldim(ext) == 2
        @test (@inferred Polyhedra.FullDim(ine)) == 2
        @test Polyhedra.coefficient_type(ext) == Int
        @test translate(ext, [1, 0]).V == [1 1; 2 0]
    end

    @testset "Cartesian product" begin
        A = [1 2; 3 4; 5 6]
        a = [7, 8, 9]
        B = [10 11 12; 13 14 15]
        b = [16, 17]
        p1 = hrep(A, a, BitSet([2]))
        p2 = hrep(B, b, BitSet([1]))
        p = p1 * p2
        @test p.A == [A[2,:]' zeros(1, 3)
                      zeros(1, 2) B[1, :]'
                      A[[1,3],:] zeros(2, 3)
                      zeros(1, 2) B[2, :]']
        @test p.b == [a[2]; b[1]; a[[1,3]]; b[2]]
        @test p.linset == BitSet([1, 2])
    end
end


## LLE pressure solver
function x0_LLE_pressure(model::EoSModel,T,x)
    xx = 1 .-x

    pure = split_model(model)

    eachx = eachcol(Diagonal(ones(eltype(x),length(x))))

    sat = sat_pure.(pure,T)

    P_sat = [tup[1] for tup in sat]
    V_l_sat = [tup[2] for tup in sat]

    V0_l = sum(x.*V_l_sat)
    V0_ll = sum(xx.*V_l_sat)

    prepend!(xx,log10.([V0_l,V0_ll]))
    return xx
end

function LLE_pressure(model::EoSModel, T, x; v0 =nothing)
    TYPE = promote_type(eltype(T),eltype(x))
#     lb_v = lb_volume(model,x)
    ts = T_scales(model,x)
    pmix = p_scale(model,x)
    if v0 === nothing
        v0 = x0_LLE_pressure(model,T,x)
    end
    len = length(v0[1:end-1])
    #xcache = zeros(eltype(x0),len)
    Fcache = zeros(eltype(v0[1:end-1]),len)
    f! = (F,z) -> Obj_bubble_pressure(model, F, T, exp10(z[1]), exp10(z[2]), x,z[3:end],ts,pmix)
    r  =Solvers.nlsolve(f!,v0[1:end-1],LineSearch(Newton()))
    sol = Solvers.x_sol(r)
    v_l = exp10(sol[1])
    v_ll = exp10(sol[2])
    xx = FractionVector(sol[3:end])
    P_sat = pressure(model,v_l,T,x)
    return (P_sat, v_l, v_ll, xx)
end


import ...timeevolution: dschroedinger, dschroedinger_dynamic, check_schroedinger

"""
    stochastic.schroedinger(tspan, state0, H, Hs[; fout, ...])

Integrate stochastic Schrödinger equation.

# Arguments
* `tspan`: Vector specifying the points of time for which the output should
        be displayed.
* `psi0`: Initial state as Ket.
* `H`: Deterministic part of the Hamiltonian.
* `Hs`: Stochastic part(s) of the Hamiltonian (either an operator or a vector
        of operators).
* `fout=nothing`: If given, this function `fout(t, state)` is called every time
        an output should be displayed. ATTENTION: The given state is neither
        normalized nor permanent!
* `normalize_state=false`: Specify whether or not to normalize the state after
        each time step taken by the solver.
* `kwargs...`: Further arguments are passed on to the ode solver.
"""
function schroedinger(tspan, psi0::T, H::AbstractOperator{B,B}, Hs::Vector;
                fout::Union{Function,Nothing}=nothing,
                normalize_state::Bool=false,
                calback=nothing,
                kwargs...) where {B<:Basis,T<:Ket{B}}
    tspan_ = convert(Vector{Float64}, tspan)

    n = length(Hs)
    dstate = copy(psi0)
    x0 = psi0.data
    state = copy(psi0)

    # TODO: replace checks by dispatch
    for h=Hs
        @assert isa(h, AbstractOperator{B,B})
    end

    dschroedinger_determ(t::Float64, psi::T, dpsi::T) = dschroedinger(psi, H, dpsi)
    dschroedinger_stoch(dx::DiffArray,
            t::Float64, psi::T, dpsi::T, n::Int) = dschroedinger_stochastic(dx, psi, Hs, dpsi, n)

    if normalize_state
        norm_func(u::Vector{ComplexF64}, t::Float64, integrator) = normalize!(u)
        ncb = DiffEqCallbacks.FunctionCallingCallback(norm_func;
                 func_everystep=true,
                 func_start=false)
    else
        ncb = nothing
    end

    integrate_stoch(tspan_, dschroedinger_determ, dschroedinger_stoch, x0, state, dstate, fout, n;
                    ncb=ncb,
                    kwargs...)
end
schroedinger(tspan, psi0::Ket{B}, H::AbstractOperator{B,B}, Hs::AbstractOperator{B,B}; kwargs...) where B<:Basis = schroedinger(tspan, psi0, H, [Hs]; kwargs...)

"""
    stochastic.schroedinger_dynamic(tspan, state0, fdeterm, fstoch[; fout, ...])

Integrate stochastic Schrödinger equation with dynamic Hamiltonian.

# Arguments
* `tspan`: Vector specifying the points of time for which the output should
        be displayed.
* `psi0`: Initial state.
* `fdeterm`: Function `f(t, psi, u) -> H` returning the deterministic
        (time- or state-dependent) part of the Hamiltonian.
* `fstoch`: Function `f(t, psi, u, du) -> Hs` returning a vector that
        contains the stochastic terms of the Hamiltonian.
* `fout=nothing`: If given, this function `fout(t, state)` is called every time
        an output should be displayed. ATTENTION: The given state is neither
        normalized nor permanent!
* `noise_processes=0`: Number of distinct white-noise processes in the equation.
        This number has to be equal to the total number of noise operators
        returned by `fstoch`. If unset, the number is calculated automatically
        from the function output.
        NOTE: Set this number if you want to avoid an initial calculation of
        the function output!
* `normalize_state=false`: Specify whether or not to normalize the state after
        each time step taken by the solver.
* `kwargs...`: Further arguments are passed on to the ode solver.
"""
function schroedinger_dynamic(tspan, psi0::T, fdeterm::Function, fstoch::Function;
                fout::Union{Function,Nothing}=nothing, noise_processes::Int=0,
                normalize_state::Bool=false,
                kwargs...) where T<:Ket
    tspan_ = convert(Vector{Float64}, tspan)

    if noise_processes == 0
        fs_out = fstoch(0.0, psi0)
        n = length(fs_out)
    else
        n = noise_processes
    end

    dstate = copy(psi0)
    x0 = psi0.data
    state = copy(psi0)

    dschroedinger_determ(t::Float64, psi::T, dpsi::T) = dschroedinger_dynamic(t, psi, fdeterm, dpsi)
    dschroedinger_stoch(dx::DiffArray,
            t::Float64, psi::T, dpsi::T, n::Int) =
        dschroedinger_stochastic(dx, t, psi, fstoch, dpsi, n)

    if normalize_state
        norm_func(u::Vector{ComplexF64}, t::Float64, integrator) = normalize!(u)
        ncb = DiffEqCallbacks.FunctionCallingCallback(norm_func;
                 func_everystep=true,
                 func_start=false)
    else
        ncb = nothing
    end

    integrate_stoch(tspan, dschroedinger_determ, dschroedinger_stoch, x0, state,
            dstate, fout, n;
            ncb=ncb,
            kwargs...)
end

function dschroedinger_stochastic(dx::D, psi::T1, Hs::Vector{T2},
            dpsi::T1, index::Int) where {D<:Vector{ComplexF64},B<:Basis,T1<:Ket{B},T2<:AbstractOperator{B,B}}
    recast!(dx, dpsi)
    dschroedinger(psi, Hs[index], dpsi)
end
function dschroedinger_stochastic(dx::Array{ComplexF64, 2}, psi::T1, Hs::Vector{T2},
            dpsi::T1, n::Int) where {B<:Basis,T1<:Ket{B},T2<:AbstractOperator{B,B}}
    for i=1:n
        dx_i = @view dx[:, i]
        recast!(dx_i, dpsi)
        dschroedinger(psi, Hs[i], dpsi)
        recast!(dpsi, dx_i)
    end
end
function dschroedinger_stochastic(dx::DiffArray,
            t::Float64, psi::T, f::Function, dpsi::T, n::Int) where T<:Ket
    ops = f(t, psi)
    if QO_CHECKS[]
        @inbounds for h=ops
            check_schroedinger(psi, h)
        end
    end
    dschroedinger_stochastic(dx, psi, ops, dpsi, n)
end

recast!(psi::Ket, x::SubArray{ComplexF64, 1}) = (x .= psi.data)
recast!(x::SubArray{ComplexF64, 1}, psi::Ket) = (psi.data = x)


"""
    rvec(n, ρ::Matrix, margins::Vector{<:UnivariateDistribution})

Generate samples for a list of marginal distributions and a correaltion structure.

# Examples
```julia-repl
julia> using Distributions

julia> import LinearAlgebra: diagind

julia> margins = [Normal(3, 1), LogNormal(3, 1), Exponential(3)]

julia> R = fill(0.5, 3, 3); R[diagind(R)] .= 1.0;

julia> rvec(10, R, margins)
#>10×3 Matrix{Float64}:
 3.71109  82.0696   8.14332
 1.23022  38.8599   2.64595
 2.07222   3.76843  1.13465
 2.82434  11.6953   0.891066
 2.37599  10.0552   1.80555
 2.85431  25.4935   3.40865
 3.85253  21.2241   3.67532
 3.70605  59.2439   2.02502
 2.34066   1.89257  0.619948
 3.83507  16.787    0.66837
```
"""
function rvec end
for T in (Float64, Float32, Float16)
    @eval function rvec(n::Int, ρ::Matrix{$T}, margins::Vector{<:UD})
        d = length(margins)
        r,s = size(ρ)

        !(r == s == d) && throw(DimensionMismatch("The number of margins must match the size of the correlation matrix."))
        !iscorrelation(ρ) && throw(ValidCorrelationError())

        Z = SharedMatrix{$T}(_rmvn(n, ρ))
        @inbounds @threads for i in 1:d
            Z[:,i] = normal_to_margin(margins[i], Z[:,i])
        end
        sdata(Z)
    end
    @eval rvec(n::Real, ρ::Matrix{$T}, margins::Vector{<:UD}) = rvec(Int(n), ρ, margins)
end


# tools for model building

import Hypatia
import Hypatia.Cones: smat_to_svec!, svec_to_smat!, svec_side, svec_length

const rt2 = sqrt(2)

function svec(mat::AbstractMatrix{T}) where {T}
    vec = zeros(T, svec_length(size(mat, 1)))
    return smat_to_svec!(vec, mat, rt2)
end

function svec(mat::AbstractMatrix{Complex{T}}) where {T}
    vec = zeros(T, svec_length(ComplexF64, size(mat, 1)))
    return smat_to_svec!(vec, mat, rt2)
end

function smat(::Type{Float64}, vec::AbstractVector{T}) where {T}
    side = svec_side(length(vec))
    mat = zeros(T, side, side)
    return svec_to_smat!(mat, vec, rt2)
end

function smat(::Type{ComplexF64}, vec::AbstractVector{T}) where {T}
    side = svec_side(ComplexF64, length(vec))
    mat = zeros(Complex{T}, side, side)
    return svec_to_smat!(mat, vec, rt2)
end

function scale_svec(
    R::Type{<:Union{Float64, ComplexF64}},
    vec::AbstractVector{T},
    scal::Float64,
) where {T}
    incr = (R == Float64 ? 1 : 2)
    svec = copy(vec)
    k = 1
    for j in 1:svec_side(R, length(vec))
        for _ in 1:(incr * (j - 1))
            # scale off-diagonal
            svec[k] *= scal
            k += 1
        end
        k += 1
    end
    @assert k == length(vec) + 1
    return svec
end


"""
    house_fits_pattern(i::Int, p::Int, d::StudentHousingData)

Returns a Boolean denoting whether someone with preference pattern `p` would
be willing to live in the `i`^th house.
"""
function house_fits_pattern(i::Int, p::Int, d::StudentHousingData)
    h = d.houses[i]
    for c in d.patterns_allow[p]
        if house_fits_characteristic(h, d.all_characteristics[c], d.market_data)
            return true
        end
    end
    false
end

function house_allowedby(i::Int, d::StudentHousingData)
    ret = Int[]
    for p = 1:length(d.patterns_allow)
      if house_fits_pattern(i, p, d)
          push!(ret, p)
      end
    end
    ret
end

# ==============================================================================
# One stage

function onestagemodel(d::StudentHousingData)

    houses = d.houses
    nhouses = length(d.houses)
    # One stage and deterministic
    demand = d.demands[:, 1, 1]

    # The number of only legal patterns
    npatterns = length(d.patterns_allow)

    m = Model(solver = GurobiSolver(OutputFlag=0))
    # m = Model(solver = CplexSolver(CPX_PARAM_SCRIND = 0, CPX_PARAM_MIPDISPLAY=0))
    @variables(m, begin
        shortage[1:npatterns] >= 0
        investment[1:nhouses], Bin
        assignment[i=1:nhouses, p=house_allowedby(i,d)] >= 0, Int
    end)

    @constraints(m, begin
        [p = 1:npatterns], shortage[p] == demand[p] - sum(assignment[i, p] for i = 1:nhouses if house_fits_pattern(i, p, d))
        [i = 1:nhouses], sum(beds_needed(p, d) * assignment[i, p] for p = 1:npatterns if house_fits_pattern(i, p, d)) <= investment[i] * beds_avail(i, houses)
        sum(investment[i] * maintenance(i, houses) for i = 1:nhouses) <= d.budget
    end)

    @objective(m, Min, sum(shortage))
    m
end

# ==============================================================================
# Multistage
function multistagemodel(d::StudentHousingData)

    houses = d.houses
    nhouses = length(d.houses)
    demands = d.demands
    upperbound = sum(sum(sum(demands)))

    # The number of only legal patterns
    npatterns = length(d.patterns_allow)

    m = SDDPModel(stages = d.nstages,
            objective_bound = 0.0,
            sense = :Min,
            solver = GurobiSolver(OutputFlag=0)) do sp, stage

        @binarystate(sp, 0 <= investment[1:nhouses] <= 1,
                investment0 == 0, Bin)

        @variables(sp, begin
            shortage[1:npatterns] >= 0
            assignment[1:nhouses, 1:npatterns] >= 0, Int
        end)

        if stage == 1
            @constraint(sp, [p = 1:npatterns], shortage[p] + sum(house_fits_pattern(i, p, d) * assignment[i, p] for i = 1:nhouses) == demands[p, 1, 1])
        else
            for p = 1:npatterns
                @rhsnoise(sp, D = demands[p, :, stage],
                    shortage[p] + sum(house_fits_pattern(i, p, d) * assignment[i, p] for i = 1:nhouses) == D)
            end
        end

        @constraints(sp, begin
            # Can't uninvest
            investment0 .<= investment
            # Only legal assignments
            [i = 1:nhouses], sum(house_fits_pattern(i, p, d) * assignment[i, p] * beds_needed(p, d) for  p = 1:npatterns) <= investment[i] * beds_avail(i, houses)
            # Budget constraint
            sum(investment[i] * maintenance(i, houses) for i = 1:nhouses) <= d.budget
        end)

        # Minimize all unmet demand
        @stageobjective(sp, sum(shortage))

        setSDDiPsolver!(sp, method=KelleyMethod(), pattern = Pattern(benders = 0, lagrangian=1))
    end
    m

end


using BenchmarkTools
using Distances

const SUITE = BenchmarkGroup()

function create_distances(w, Q)
    dists = [
        SqEuclidean(),
        Euclidean(),
        PeriodicEuclidean(w),
        Cityblock(),
        TotalVariation(),
        Chebyshev(),
        Minkowski(3.0),
        Hamming(),

        CosineDist(),
        CorrDist(),
        ChiSqDist(),

        BhattacharyyaDist(),
        HellingerDist(),

        # Haversine(6371.),

        BrayCurtis(),

        WeightedSqEuclidean(w),
        WeightedEuclidean(w),
        WeightedCityblock(w),
        WeightedMinkowski(w, 3.0),
        WeightedHamming(w),

        SqMahalanobis(Q),
        Mahalanobis(Q),
    ]

    divs = [
        KLDivergence(),
        RenyiDivergence(0),
        RenyiDivergence(1),
        RenyiDivergence(2),
        RenyiDivergence(Inf),
        JSDivergence(),
    ]
    return dists, divs
end

###########
# Colwise #
###########

SUITE["colwise"] = BenchmarkGroup()

function evaluate_colwise(dist, x, y)
    n = size(x, 2)
    T = typeof(evaluate(dist, x[:, 1], y[:, 1]))
    r = Vector{T}(undef, n)
    for j = 1:n
        r[j] = @views evaluate(dist, x[:, j], y[:, j])
    end
    return r
end

function add_colwise_benchmarks!(SUITE)

    m = 200
    n = 10000

    x = rand(m, n)
    y = rand(m, n)

    p = x
    q = y
    for i = 1:n
        p[:, i] /= sum(x[:, i])
        q[:, i] /= sum(y[:, i])
    end

    w = rand(m)

    Q = rand(m, m)
    Q = Q' * Q

    _dists, divs = create_distances(w, Q)

    for (dists, (a, b)) in [(_dists, (x,y)), (divs, (p,q))]
        for dist in (dists)
            Tdist = typeof(dist)
            SUITE["colwise"][Tdist] = BenchmarkGroup()
            SUITE["colwise"][Tdist]["loop"]        = @benchmarkable evaluate_colwise($dist, $a, $b)
            SUITE["colwise"][Tdist]["specialized"] = @benchmarkable colwise($dist, $a, $b)
        end
    end
end

add_colwise_benchmarks!(SUITE)

############
# Pairwise #
############

SUITE["pairwise"] = BenchmarkGroup()

function evaluate_pairwise(dist, x, y)
    nx = size(x, 2)
    ny = size(y, 2)
    T = typeof(evaluate(dist, x[:, 1], y[:, 1]))
    r = Matrix{T}(undef, nx, ny)
    for j = 1:ny
        @inbounds for i = 1:nx
            r[i, j] = @views evaluate(dist, x[:, i], y[:, j])
        end
    end
    return r
end

function add_pairwise_benchmarks!(SUITE)
    m = 100
    nx = 200
    ny = 250

    x = rand(m, nx)
    y = rand(m, ny)

    p = x
    for i = 1:nx
        p[:, i] /= sum(x[:, i])
    end

    q = y
    for i = 1:ny
        q[:, i] /= sum(y[:, i])
    end

    w = rand(m)
    Q = rand(m, m)
    Q = Q' * Q

    _dists, divs = create_distances(w, Q)

     for (dists, (a, b)) in [(_dists, (x,y)), (divs, (p,q))]
        for dist in (dists)
            Tdist = typeof(dist)
            SUITE["pairwise"][Tdist] = BenchmarkGroup()
            SUITE["pairwise"][Tdist]["loop"]        = @benchmarkable evaluate_pairwise($dist, $a, $b)
            SUITE["pairwise"][Tdist]["specialized"] = @benchmarkable pairwise($dist, $a, $b; dims=2)
        end
    end
end

add_pairwise_benchmarks!(SUITE)


# The submarine has been making some odd creaking noises, so you ask it to
# produce a diagnostic report just in case.
#
# The diagnostic report (your puzzle input) consists of a list of binary
# numbers which, when decoded properly, can tell you many useful things about
# the conditions of the submarine. The first parameter to check is the power
# consumption.
#
# You need to use the binary numbers in the diagnostic report to generate two
# new binary numbers (called the gamma rate and the epsilon rate). The power
# consumption can then be found by multiplying the gamma rate by the epsilon
# rate.
#
# Each bit in the gamma rate can be determined by finding the most common bit
# in the corresponding position of all numbers in the diagnostic report. For
# example, given the following diagnostic report:
#
# 00100
# 11110
# 10110
# 10111
# 10101
# 01111
# 00111
# 11100
# 10000
# 11001
# 00010
# 01010
#
# Considering only the first bit of each number, there are five 0 bits and
# seven 1 bits. Since the most common bit is 1, the first bit of the gamma rate
# is 1.
#
# The most common second bit of the numbers in the diagnostic report is 0, so
# the second bit of the gamma rate is 0.
#
# The most common value of the third, fourth, and fifth bits are 1, 1, and 0,
# respectively, and so the final three bits of the gamma rate are 110.
#
# So, the gamma rate is the binary number 10110, or 22 in decimal.
#
# The epsilon rate is calculated in a similar way; rather than use the most
# common bit, the least common bit from each position is used. So, the epsilon
# rate is 01001, or 9 in decimal. Multiplying the gamma rate (22) by the
# epsilon rate (9) produces the power consumption, 198.
#
# Use the binary numbers in your diagnostic report to calculate the gamma rate
# and epsilon rate, then multiply them together. What is the power consumption
# of the submarine? (Be sure to represent your answer in decimal, not binary.)

# convert each line to a vector of bits
function mapper(l)
  map(d -> parse(Int, d), split(l, ""))
end

# sum all lines into a vector of sums - if any sum is >500 (since there are
# 1000 rows) then 1's are the most common value.
cnts = mapreduce(mapper, +, eachline("input.txt"))

# ok, we're going to use a trick here... cnts .> 500 will return a BitVector
# (aka BitArray with one dimension) where each element is a 1 if the condition
# is true (ie, >500) and a 0 otherwise. Under the hood, a BitVector stores the
# data as an Int64 but the bits are reversed. By reversing cnt, we can read out
# that int without writing any code to convert the BitVector to an int
gamma = (reverse(cnts) .> 500).chunks[1]

# the epsilon rate will always be the bitwise negation of gamma anded with
# 0xFFF since we're only interested in the 12 least significant bits
epsilon = (~gamma) & 0xfff

println("Result: ", gamma * epsilon)


mutable struct BlackBoxOptimizer <: HeuristicOptimizer
    lossfn
    optparams::AbstractDict
    optres
    BlackBoxOptimizer(lossfn::Function; kws...) = new(lossfn, kws, nothing)
end

BlackBoxOptimizer(m::InterpretableModel; lossfn = l2loss, kws...) =
    BlackBoxOptimizer(lossfn; kws...)

const DefaultBBOParams = Dict(
    :PopulationSize => 500,
    :MaxTime => 15.0,
)

function getp(key, d1, opt, dbackup = DefaultBBOParams)
    haskey(d1, key) && return d1[key]
    haskey(opt.optparams, key) && return opt.optparams[key]
    haskey(dbackup, key) && return dbackup[key]
    error("No value for parameter $key found!")
end

function fit!(m::InterpretableModel, o::BlackBoxOptimizer,
                X::AbstractMatrix{Float64}, y::AbstractVector{Float64}; kws...)

    @assert size(X, 1) >= length(y) - maxlag(m)

    fitnessfn(ps::Vector{Float64}) = begin
        yhat, lossvalue = predictandcalcloss(m, X, y, ps, o.lossfn)
        lossvalue # Loss value is the fitness so should be minimixed
    end

    o.optres = bboptimize(fitnessfn;
                SearchSpace = parambounds(m),
                PopulationSize = getp(:PopulationSize, kws, o),
                MaxTime = getp(:MaxTime, kws, o))

    setparams!(m, best_candidate(o.optres))
    m
end


@testset "Internal Sorting" begin
    data = [21, 12, 6, 26, 4, 15, 2, 8, 24, 22, 13, 30, 7, 9, 27, 1, 23, 19, 11, 14, 18, 28, 20, 25, 16, 17, 5, 3, 10, 29]
    @test insertionsort!(copy(data)) == 1:length(data)
    @test selectionsort!(copy(data)) == 1:length(data)
    @test shellsort!(copy(data)) == 1:length(data)

    @test mergesort!(copy(data)) == 1:length(data)
    @test quicksort!(copy(data)) == 1:length(data)

    @test heapsort!(copy(data)) == 1:length(data)
end


"""
    OILMM(fs, U, S, D)

An Orthogonal Instantaneous Linear Mixing Model (OILMM) -- a distribution over vector-
valued functions. Let `p` be the number of observed outputs, and `m` the number of latent
processes, then

# Arguments:
- fs: a length-`m` vector of Gaussian process objects from Stheno.jl.
- U: a `p x m` matrix with mutually orthonormal columns.
- S: an `m x m` `Diagonal` matrix. Same dim. as `fs`. Positive entries.
- D: an `m x m` `Diagonal` matrix, variance of noise on each latent process. Same size as
    `S`. Positive entries.

We recommend constructing `U` and `S` from the `svd` of some other matrix. e.g.
```julia
H = randn(p, m)
U, s, _ = svd(H)
S = Diagonal(H)
```
"""
struct OILMM{
    Tfs<:AbstractVector{<:AbstractGP},
    TU<:AbstractMatrix{<:Real},
    TS<:Diagonal{<:Real},
    TD<:Diagonal{<:Real},
} <: AbstractGP
    fs::Tfs
    U::TU
    S::TS
    D::TD
end

noise_var(Σ::Diagonal{<:Real, <:Fill}) = FillArrays.getindex_value(Σ.diag)

function unpack(fx::FiniteGP{<:OILMM, <:MOInput, <:Diagonal{<:Real, <:Fill}})
    fs = fx.f.fs
    S = fx.f.S
    U = fx.f.U
    D = fx.f.D
    σ² = noise_var(fx.Σy)
    x = fx.x.x

    # Check that the number of outputs requested agrees with the model.
    fx.x.out_dim == size(U, 1) || throw(error("out dim of x != out dim of f."))
    return fs, S, U, D, σ², x
end

reshape_y(y::AbstractVector{<:Real}, N::Int) = ColVecs(reshape(y, N, :)')

# Note that `cholesky` exploits the diagonal structure of `S`.
function project(
    S::Diagonal{T},
    U::AbstractMatrix{T},
    Y::ColVecs{T},
    σ²::T,
    D::Diagonal{T},
) where {T<:Real}

    # Compute the projection of the data.
    Yproj = cholesky(S).U \ U' * Y.X

    # Compute the projected noise, which is a matrix of size `size(Yproj)`.
    ΣT = repeat(diag(σ² * inv(S) + D), 1, size(Yproj, 2))

    return Yproj, ΣT
end

# Compute the regularisation term in the log marginal likelihood. See e.g. appendix A.4.
function regulariser(
    S::Diagonal{T},
    U::AbstractMatrix{T},
    σ²::T,
    Y::ColVecs{T},
) where {T<:Real}
    n = length(Y)
    p, m = size(U)
    return -(n * (logdet(cholesky(S)) + (p - m) * log(2π * σ²)) +
        sum(abs2, (I - U * U') * Y.X) / σ²) / 2
end

"""
    rand_latent(rng::AbstractRNG, fx::FiniteGP{<:OILMM})

Sample from the latent (noiseless) process.

See also `rand`.
"""
function rand_latent(rng::AbstractRNG, fx::FiniteGP{<:OILMM})
    fs, S, U, D, σ², x = unpack(fx)

    # Generate from the latent processes.
    X = hcat(map((f, d) -> rand(rng, f(x, d)), fs, D.diag)...)

    # Transform latents into observed space.
    return vec(U * cholesky(S).U * X')
end

"""
    rand(rng::AbstractRNG, fx::FiniteGP{<:OILMM})

Sample from the OILMM, including the observation noise.
Follows generative structure of model 2 from [1].
Follows the AbstractGPs.jl API.

See also `rand_latent`.

[1] - Bruinsma et al 2020.
"""
function AbstractGPs.rand(rng::AbstractRNG, fx::FiniteGP{<:OILMM})

    # Sample from the latent process.
    F = rand_latent(rng, fx)

    # Generate iid noise and add to each output.
    return F .+ sqrt(noise_var(fx.Σy)) .* randn(rng, size(F))
end

"""
    denoised_marginals(fx::FiniteGP{<:OILMM})

Returns the marginal distribution over the OILMM without the IID noise components.

See also `marginals`.
"""
function denoised_marginals(fx::FiniteGP{<:OILMM})
    fs, S, U, D, σ², x = unpack(fx)

    # Compute the marginals over the independent latents.
    fs_marginals = reduce(hcat, map(f -> marginals(f(x)), fs))
    M_latent = mean.(fs_marginals)'
    V_latent = var.(fs_marginals)'

    # Compute the latent -> observed transform.
    H = U * cholesky(S).U

    # Compute the means.
    M = H * M_latent

    # Compute the variances.
    V = abs2.(H) * V_latent

    # Package everything into independent Normal distributions.
    return AbstractGPs.Normal.(vec(M'), sqrt.(vec(V')))
end

# See AbstractGPs.jl API docs.
function AbstractGPs.mean_and_var(fx::FiniteGP{<:OILMM})
    fs, S, U, D, σ², x = unpack(fx)

    # Compute the marginals over the independent latents.
    fs_marginals = hcat(map(f -> marginals(f(x)), fs)...)
    M_latent = mean.(fs_marginals)'
    V_latent = var.(fs_marginals)'

    # Compute the latent -> observed transform.
    H = U * cholesky(S).U

    # Compute the means.
    M = H * M_latent

    # Compute the variances.
    V = abs2.(H) * (V_latent .+ D.diag) .+ σ²

    # Package everything into independent Normal distributions.
    return vec(M'), vec(V')
end

AbstractGPs.mean(fx::FiniteGP{<:OILMM}) = mean_and_var(fx)[1]

AbstractGPs.var(fx::FiniteGP{<:OILMM}) = mean_and_var(fx)[2]

# See AbstractGPs.jl API docs.
function AbstractGPs.logpdf(fx::FiniteGP{<:OILMM}, y::AbstractVector{<:Real})
    fs, S, U, D, σ², x = unpack(fx)

    # Projection step.
    Y = reshape_y(y, length(x))
    Yproj, ΣT = project(S, U, Y, σ², D)

    # Latent process log marginal likelihood calculation.
    y_rows = collect(eachrow(Yproj))
    ΣT_rows = collect(eachrow(ΣT))
    lmls_latents = map((f, s, y) -> logpdf(f(x, collect(s)), collect(y)), fs, ΣT_rows, y_rows)

    return regulariser(S, U, σ², Y) + sum(lmls_latents)
end

# See AbstractGPs.jl API docs.
function AbstractGPs.posterior(fx::FiniteGP{<:OILMM}, y::AbstractVector{<:Real})
    fs, S, U, D, σ², x = unpack(fx)

    # Projection step.
    Y = reshape_y(y, length(x))
    Yproj, ΣT = project(S, U, Y, σ², D)

    # Condition each latent process on the projected observations.
    y_rows = collect(eachrow(Yproj))
    ΣT_rows = collect(eachrow(ΣT))
    fs_posterior = map((f, s, y) -> posterior(f(x, collect(s)), collect(y)), fs, ΣT_rows, y_rows)
    return OILMM(fs_posterior, U, S, D)
end


using ITensors
using ITensorNetworkAD
using ITensorNetworkAD.ITensorNetworks:
  Models,
  inds_network,
  project_boundary,
  contract_approx,
  insert_projectors,
  sqnorm,
  sqnorm_approx
using Test

@testset "ITensorNetworks.jl" begin
  model = Models.Model"ising"()
  βc = Models.critical_point(model)
  β = 1.001 * βc

  N = (3, 4)
  ndims = length(N)

  space = 2
  tn_inds = inds_network(N...; linkdims=space, addtags="S=1/2")
  A = Models.local_boltzmann_weight("ising", Val(ndims); β=β)

  tn = itensor.((A,), tn_inds)

  # Project periodic boundary indices
  # onto the 1 state
  state = 1
  tn = project_boundary(tn, state)

  _cutoff = 1e-15
  _maxdim = 100

  # Contract in every directions
  boundary_mps = contract_approx(tn; maxdim=_maxdim, cutoff=_cutoff)

  #
  # Insert projectors horizontally (to measure e.g. properties
  # in a row of the network)
  #

  row = 2
  center = (row, :)
  tn_projected = insert_projectors(tn, boundary_mps; center=center)
  tn_split, Pl, Pr = tn_projected

  Pl_flat = reduce(vcat, Pl)
  Pr_flat = reduce(vcat, Pr)
  tn_projected_flat = mapreduce(vec, vcat, (tn_split, Pl_flat, Pr_flat))

  @test isempty(noncommoninds(tn_projected_flat...))

  @disable_warn_order begin
    @test contract(tn_projected_flat)[] ≈ contract(vec(tn))[]
  end

  #
  # Insert projectors vertically (to measure e.g. properties
  # in a column of the network)
  #

  column = 2
  center = (:, column)
  tn_projected = insert_projectors(tn, boundary_mps; center=center)
  tn_split, Pl, Pr = tn_projected

  Pl_flat = reduce(vcat, Pl)
  Pr_flat = reduce(vcat, Pr)
  tn_projected_flat = mapreduce(vec, vcat, (tn_split, Pl_flat, Pr_flat))

  @test isempty(noncommoninds(tn_projected_flat...))

  @disable_warn_order begin
    @test contract(tn_projected_flat)[] ≈ contract(vec(tn))[]
  end
end

function peps_tensor(; linkdim, sitedim)
  # left, right, top, bottom, site
  return randn(linkdim, linkdim, linkdim, linkdim, sitedim)
end

@testset "PEPS norm approx" begin
  N = (3, 3)
  ndims = length(N)

  site_inds = siteinds("S=1/2", N...)
  link_space = 2
  inds_net = inds_network(site_inds; linkdims=link_space)

  A = peps_tensor(; linkdim=link_space, sitedim=dim(first(site_inds)))
  ψ = itensor.((A,), inds_net)

  # Project periodic boundary indices
  # onto the 1 state
  state = 1
  ψ = project_boundary(ψ, state)

  row = 2
  center = (row, :)
  cutoff_ = 1e-15
  maxdim_ = 100

  sqnormψ = sqnorm(ψ)
  sqnormψ_approx = sqnorm_approx(ψ; center=center, cutoff=cutoff_, maxdim=maxdim_)
  @test isempty(noncommoninds(sqnormψ_approx...))

  @disable_warn_order begin
    @test contract(sqnormψ_approx)[] / contract(sqnormψ)[] ≈ 1.0
  end
end


# Sample from arbitrary arrays

################################################################
#
#  A variety of algorithms for sampling without replacement
#
#  They are suited for different cases.
#
#  Particularly,
#  - Fisher-Yates sampler is suited for general cases
#    where n is not overly large
#
#  - Self avoiding sampler is suited for cases where k << n
#
################################################################

function pick2!(a::AbstractArray, x::AbstractArray)
    # Pick a pair of values without replacement

    n0 = length(a)
    i1 = randi(n0)
    i2 = randi(n0 - 1)
    if i2 == i1
        i2 = n0
    end

    x[1] = a[i1]
    x[2] = a[i2]
end

## A sampler that implements without-replacement sampling
## via Fisher-Yates shuffling
##
immutable FisherYatesSampler
    n::Int
    seq::Vector{Int}   # Internal sequence for shuffling

    FisherYatesSampler(n::Int) = new(n, [1:n])
end

function rand!(s::FisherYatesSampler, a::AbstractArray, x::AbstractArray)
    # draw samples without-replacement to x

    n::Int = s.n
    k::Int = length(x)
    if k > n
        throw(ArgumentError("Cannot draw more than n samples without replacement."))
    end

    seq::Vector{Int} = s.seq
    for i = 1:k
        j = randi(i, n)
        sj = seq[j]
        x[i] = a[sj]
        seq[j] = seq[i]
        seq[i] = sj
    end
    x
end

fisher_yates_sample!(a::AbstractArray, x::AbstractArray) = rand!(FisherYatesSampler(length(a)), a, x)

function self_avoid_sample!{T}(a::AbstractArray{T}, x::AbstractArray)
    # This algorithm is suitable when length(x) << length(a)

    s = Set{T}()
    # sizehint(s, length(x))
    rgen = RandIntSampler(length(a))

    # first one
    idx = rand(rgen)
    x[1] = a[idx]
    add!(s, idx)

    # remaining
    for i = 2:length(x)
        idx = rand(rgen)
        while contains(s, idx)
            idx = rand(rgen)
        end
        x[i] = a[idx]
        add!(s, idx)
    end
    x
end

###########################################################
#
#   Interface functions
#
###########################################################

sample(a::AbstractArray) = a[randi(length(a))]

function sample!(a::AbstractArray, x::AbstractArray; replace=true)
    n = length(a)
    k = length(x)

    if !isempty(x)
        if replace   # with replacement
            s = RandIntSampler(n)
            for i = 1:k
                x[i] = a[rand(s)]
            end

        else  # without replacement
            if k > n
                throw(ArgumentError("n exceeds the length of x"))
            end

            if k == 1
                x[1] = sample(a)

            elseif k == 2
                pick2!(a, x)

            elseif n < k * max(k, 100)
                fisher_yates_sample!(a, x)

            else
                self_avoid_sample!(a, x)
            end
        end
    end
    x
end

function sample{T}(a::AbstractArray{T}, n::Integer; replace=true)
    sample!(a, Array(T, n); replace=replace)
end

function sample{T}(a::AbstractArray{T}, dims::Dims; replace=true)
    sample!(a, Array(T, dims); replace=replace)
end

################################################################
#
#  Weighted sampling
#
################################################################

function wsample(a::AbstractArray, w::AbstractArray{Float64}, wsum::Float64)
    n = length(w)
    t = rand() * wsum

    i = 1
    s = w[1]

    while i < n && s < t
        i += 1
        s += w[i]
    end
    a[i]
end

wsample(a::AbstractArray, w::AbstractArray{Float64}) = wsample(a, w, sum(w))

function wsample!(a::AbstractArray, w::AbstractArray{Float64}, x::AbstractArray;
    wsum::Float64=NaN)

    n = length(a)
    if length(w) != n
        throw(ArgumentError("Inconsistent argument dimensions."))
    end

    _wsum::Float64 = isnan(wsum) ? sum(w) : wsum
    for i = 1:length(x)
        x[i] = wsample(a, w, _wsum)
    end
    x
end

function wsample{T}(a::AbstractArray{T}, w::AbstractArray{Float64}, n::Integer; wsum::Float64=NaN)
    wsample!(a, w, Array(T, n); wsum=wsum)
end

function wsample{T}(a::AbstractArray{T}, w::AbstractArray{Float64}, dims::Dims; wsum::Float64=NaN)
    wsample!(a, w, Array(T, dims); wsum=wsum)
end


# test differentiable MPC policy on cartpole swingup
include_model("cartpole")
n = model.n
m = model.m
d = model.d

# fast dynamics
@variables x1[1:n], x2[1:n], u1[1:m], w1[1:d], hs[1:1], ts[1:1]

function fd(model::Cartpole{Midpoint, FixedTime}, x⁺, x, u, w, h, t)
    x⁺ - (x + h[1] * f(model, 0.5 * (x + x⁺), u, w))
end

_dynamics = fd(model, x2, x1, u1, w1, hs, ts)
dynamics = eval(build_function(_dynamics, x2, x1, u1, w1, hs, ts)[2])
∇dynamics_x2 = eval(build_function(Symbolics.jacobian(_dynamics, x2), x2, x1, u1, w1, hs, ts)[1])
∇dynamics_x1 = eval(build_function(Symbolics.jacobian(_dynamics, x1), x2, x1, u1, w1, hs, ts)[1])
∇dynamics_u1 = eval(build_function(Symbolics.jacobian(_dynamics, u1), x2, x1, u1, w1, hs, ts)[1])

# fast objective
function objective(x, u, t)
    J = 0.0
    if t < 10
        J += transpose(x) * x
        J += transpose(u) * u
    else
        J += 10.0 * transpose(x) * x
    end
    return J
end

_obj = objective(x1, u1, t)
obj = eval(build_function)
function objective(x)
    J = 0.0
    J += transpose(x) * x
    return J
end