#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Many auxiliary functions for ground repeating, Sun-synchronous (GRSS) orbit
#   computations.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export adjacent_track_distance_grss, adjacent_track_angle_grss
export compute_ss_orbit_by_num_rev_per_day
export list_ss_orbits_by_rep_period, sort_list_ss_orbits_by_height

"""
    adjacent_track_distance_grss(T::Number, i::Number, To::Int, lat::Number)

Compute the distance between adjacent ground tracks [m] at a given latitude
`lat` [rad] for a ground repeating, Sun-synchronous orbit with period `T` [s],
inclination `i` [rad], and orbit cycle `To` [days].

# Remarks

The functions **does not** check if the orbit is a GRSS orbit.

"""
function adjacent_track_distance_grss(T::Number, i::Number, To::Int, lat::Number)
    # Angle between two adjacent traces.
    theta = (T/To)*pi/43200.0

    # Angle of swath.
    beta = asin(sin(theta)*sin(i))

    # Minimum swath.
    beta*R0*cos(lat)
end

"""
    adjacent_track_distance_grss(a::Number, e::Number, i::Number, To::Int, lat::Number)

Compute the distance between adjacent ground tracks [m] at a given latitude
`lat` [rad] for a ground repeating, Sun-synchronous orbit with semi-major axis
`a` [m], eccentricity `e`, inclination `i` [rad], and orbit cycle `To` [days].

# Remarks

The functions *does not* check if the orbit is a GRSS orbit.

"""
function adjacent_track_distance_grss(a::Number, e::Number, i::Number, To::Int,
                                      lat::Number)
    # Orbit period.
    T = period(a, e, i, :J2)

    # Compute the distance between adjacent ground tracks.
    adjacent_track_distance_grss(T, i, To, lat)
end

"""
    adjacent_track_angle_grss(h::Number, T::Number, i::Number, To::Int, lat::Number)

Compute the angle between two adjacent ground tracks [rad] in a given latitude
`lat` [rad] measured from the satellite position for a ground repeating,
Sun-synchronous orbit with altitude in the Equator `h` [m], period `T` [s],
inclination `i` [rad], and orbit cycle `To` [days].

# Remarks

The functions **does not** check if the orbit is a GRSS orbit.

"""
function adjacent_track_angle_grss(h::Number, T::Number, i::Number, To::Int,
                                   lat::Number)
    # Angle between two adjacent traces.
    theta = (T/To)*pi/43200.0

    # Angle of swath.
    beta = asin(sin(theta)*sin(i))

    # Get the Earth radius in the plane perpendicular to the Equatorial plane.
    Rp = R0*cos(lat)

    # Compute the angle between the two ground tracks.
    b = sqrt(Rp^2+(Rp+h)^2-2*Rp*(Rp+h)*cos(beta/2.0))
    asin(Rp/b*sin(beta/2.0))
end

"""
    adjacent_track_angle_grss(h::Number, a::Number, e::Number, i::Number, To::Int, lat::Number)

Compute the angle between two adjacent ground tracks [rad] in a given latitude
`lat` [rad] measured from the satellite position for a ground repeating,
Sun-synchronous orbit with altitude in the Equator `h` [m], semi-major axis `a`
[m], eccentricity `e`, inclination `i` [rad], and orbit cycle `To` [days].

# Remarks

The functions *does not* check if the orbit is a GRSS orbit.

"""
function adjacent_track_angle_grss(h::Number, a::Number, e::Number, i::Number,
                                   To::Int, lat::Number)
    # Period (J2).
    T = period(a, e, i, :J2)

    # Compute the minimum half FOV.
    adjacent_track_angle_grss(h, T, i, To, lat)
end

"""
    compute_ss_orbit_by_num_rev_per_day(numRevPD::Number, e::Number)

Compute the Sun-synchronous orbit given the number of revolutions per day
`numRevPD` and the eccentricity `e`.

# Returns

* The semi-major axis [m].
* The inclination [rad].
* The residues of the two functions.
* A boolean variable that indicates if the numerical algorithm converged.

"""
function compute_ss_orbit_by_num_rev_per_day(numRevPD::Number, e::Number)
    compute_ss_orbit_by_ang_vel(numRevPD*2*pi/86400.0, e)
end

"""
    list_ss_orbits_by_rep_period(minRep::Int, maxRep::Int, minAlt::Number=-1.0, maxAlt::Number=-1.0, e::Number=0.0)

Compute a list of repeating Sun-synchronous orbits.

# Args

* `minRep`: Minimum repetition time of the orbit [days].
* `maxRep`: Maximum repetition time of the orbit [days].
* `minAlt`: Minimum altitude of the orbits on the list [m].
* `maxAlt`: Minimum altitude of the orbits on the list [m].
* `e`: Eccentricity.

# Returns

A matrix containing the orbits found with the following format:

    Semi-major axis [m] | Altitude [m] | Inclination [rad] | Period [s] | Int | Num | Den
    --------------------|--------------|-------------------|------------|-----|-----|-----

in which the period is Int + Num/Den.

# Remarks

If `minAlt` or `maxAlt` is < 0.0, then the altitude will not be checked when a
orbit is added to the list.

"""
function list_ss_orbits_by_rep_period(minRep::Int,         maxRep::Int,
                                      minAlt::Number=-1.0, maxAlt::Number=-1.0,
                                      e::Number=0.0)
    # Check if the arguments are valid.
    minRep <= 0       && throw(ArgumentError("The minimum repetition time must be greater than 0."))
    maxRep <= 0       && throw(ArgumentError("The maximum repetition time must be greater than 0."))
    minRep > maxRep   && throw(ArgumentError("The minimum repetition time must be smaller or equal to the maximum repetition time."))
    !( 0. <= e < 1. ) && throw(ArgumentError("The eccentricity must be within the interval 0 <= e < 1."))

    # Matrix to store the available orbits.
    ss_orbits = Matrix{Float64}(undef, 0, 7)

    # Integer part of the number of orbits per day.
    intNumOrb = (13, 14, 15, 16, 17)

    # Check if the altitude interval must be verified.
    check_altitude = (minAlt > 0) && (maxAlt > 0)
    check_altitude && (minAlt > maxAlt) && throw(ArgumentError("The minimum altitude must be lower than the maximum altitude."))

    # Loop for the possible repetition times.
    for den = minRep:maxRep
        for num = 0:den-1

            # Check if the fraction num/den is irreducible.
            if gcd(num, den) == 1

                # Loop through the integer parts.
                for ino in intNumOrb
                    addOrbit = false

                    # Number of revolutions per day.
                    numRevPD = ino + num/den

                    (a, i, f1r, f2r, converged) =
                        compute_ss_orbit_by_num_rev_per_day(numRevPD, e)

                    # Check if the orbit is valid.
                    !(@check_orbit(a,e)) && continue

                    # Check if the altitude interval must be verified.
                    if check_altitude
                        if (minAlt < a-R0) && (a-R0 < maxAlt) && (converged)
                            addOrbit = true
                        end
                    else
                        addOrbit = converged
                    end

                    # Check if the orbit must be added to the list.
                    if addOrbit
                        # Compute the period of the orbit considering
                        # perturbations (J2).
                        T = period(a, e, i, :J2)

                        ss_orbits = vcat(ss_orbits, [a a-R0 i T ino num den])
                    end
                end
            end
        end
    end

    # Return the list of orbits.
    ss_orbits
end

"""
    sort_list_ss_orbits_by_height(ss_orbits::Matrix)

Sort the list of Sun-synchronous orbits `ss_orbits` (see
`list_ss_orbits_by_rep_period`) by height and return a new matrix.

"""
sort_list_ss_orbits_by_height(ss_orbits::Matrix) =
    sortslices(ss_orbits, dims=1, by=x->x[1])


using HPAT
using MPI

@acc hpat function calcKmeans(k::Int64, file_name)
    points = DataSource(Matrix{Float64},HDF5,"/points",file_name)
    clusters = Kmeans(points, k, 10)
    return clusters
end

K = 3

if (length(ARGS) > 0)
	K = parse(Int, ARGS[1])
end

c2 = calcKmeans(K, HPAT.getDefaultDataPath()*"/kmeans_large.hdf5")

if MPI.Comm_rank(MPI.COMM_WORLD)==0	println("Centroids: ", c2) end


using Statistics

function fuelcost(steps)
    cost = 0
    stepcost = 1
    for i in 1:steps
        cost += stepcost
        stepcost += 1
    end
    return cost
end

function solve(filename, timesteps = 80)
    #load txt file
    values = readlines(filename*".txt")
    values = split(values[1], ',')
    positions = parse.(Int, values)
    goalpos = median(positions)
    avgpos = round(mean(positions))
    cost = 0
    costavg = 0
    costavgsmall = 0
    costavgbig = 0
    for p in positions
        cost += abs(p - goalpos)
        costavg += fuelcost(abs(p - avgpos))
        costavgsmall += fuelcost(abs(p-(avgpos-1)))
        costavgbig += fuelcost(abs(p-(avgpos+1)))
    end
    costavg = min(costavg, costavgsmall, costavgbig)
    println("Median: "*string(goalpos))
    println("Accumulated fuel cost: "*string(cost))
    println("Average: "*string(avgpos))
    println("Avg fuel cost accumulated: "*string(costavg))
end

solve("smalltest")
solve("input")


function fevd(estimate::VARRepresentation, H)
	Psi(estimate, H)
	FEVD = [sum(estimate.Psi[i,j,:].^2) for i=1:estimate.vars, j=1:estimate.vars]
	FEVD = hcat([FEVD[:,i]/sum(FEVD[:,i]) for i=1:estimate.vars]...)
	return FEVD
end

function genFEVD(estimate::VARRepresentation, H::Int; nocorr = false)
	Phi(estimate, H)

	A = estimate.Phi
	Σ = deepcopy(estimate.Σ)
	K = estimate.vars

	den = zeros(size(A)[1])
	num = zeros(A[:,:,1])

	if nocorr
		Σ = diagm(diag(deepcopy(estimate.Σ)))
	end

	for i=1:H
		t = A[:,:,i]
		z = t'*Σ
		den += diag(z*t)
		num += z.^2
	end

	θ = [num[i,j]/(den[i]*Σ[j,j]) for i=1:K, j=1:K]
	for i=1:K
		θ[i,:] = θ[i,:]/sum(θ[i,:])
	end
	θ = convert(Array{Float64}, θ)

	return θ
end

# Fourier decomposition of the FEVDs

function fftFEVD(estimate::VARRepresentation, H; nocorr = false, range = nothing)
	k = estimate.vars

	Phi(estimate, H)

	if nocorr
		Σ = diagm(diag(deepcopy(estimate.Σ)))
	else
		Σ = estimate.Σ
	end

	P = chol(Σ)

	if !(typeof(range) <: Tuple)
		range[1] = 1
		range[2] = H
	end

	ft = mapslices(fft, estimate.Phi, [3])
	decomp = [(abs(ft[:,i,z]'*P'[:,j]).^2)[1] / H for i=1:k, j = 1:k, z = 1:H]
	denom = mapslices(sum, decomp[:,:,range[1]:range[2]], [3, 2])

	return  mapslices((x) -> x./denom[:,:,1], decomp, [1,2])
end

function fftGenFEVD(estimate::VARRepresentation, H; nocorr = false, range = nothing)
	k = estimate.vars

	Phi(estimate, H)

	if nocorr
		Σ = diagm(diag(estimate.Σ))
	else
		Σ = estimate.Σ
	end

	P = convert(Array{Float64,2}, chol(Σ))
	ft = mapslices(fft, estimate.Phi, [3])
	decompNew = zeros(k,k,H)
	for i=1:k
		temp = ft[:,i,:]
		for j=1:k
			temp2 = Σ[j,:]
			for h=1:H
				decompNew[i,j,h] = (abs(temp[:,h]'*temp2).^2)[1] / H
			end
		end
	end

	decomp = zeros(k,k,H)

	for i=1:k
		for j=1:k
			for h=1:H
				decomp[i,j,h] = (abs(ft[:,i,h]'*P'[:,j]).^2)[1] / H
			end
		end
	end

	if !(typeof(range) <: Tuple)
		range = [1,H]
	end

	denom = vec(mapslices(sum, decomp[:,:,range[1]:range[2]], [3, 2]))

	θ = [decompNew[i,j,h]/(denom[i]*Σ[j,j]) for i=1:k, j = 1:k, h = 1:H]

	for i=1:k
		s = sum(θ[i,:,range[1]:range[2]])
		for j=1:k
			for h=1:H
				θ[i,j,h] = θ[i,j,h]/s
			end
		end
	end

	return θ
end


#
module Spectral

using Reexport

@reexport using SciMLBase

using LinearAlgebra
using LinearSolve

import UnPack: @unpack
import Setfield: @set!
import Base.ReshapedArray
import SciMLBase: AbstractDiffEqOperator
import Lazy: @forward
import SparseArrays: sparse
import NNlib: gather, scatter
import FastGaussQuadrature: gausslobatto, gausslegendre, gausschebyshev
import FFTW: plan_rfft, plan_irfft

# AbstractVector subtyping
import Base: summary, show, similar, zero
import Base: size, getindex, setindex!, IndexStyle
import Base.Broadcast: BroadcastStyle

# overload maths
import Base: +, -, *, /, \, adjoint, ∘, inv, one
import LinearAlgebra: mul!, ldiv!, lmul!, rmul!

###
# Abstract Supertypes
###

""" Scalar function field in D-Dimensional space """
abstract type AbstractField{T,D} <: AbstractVector{T} end
""" Operators acting on fields in D-Dimensional space """
abstract type AbstractOperator{T,D} <: AbstractDiffEqOperator{T} end
""" D-Dimensional physical domain """
abstract type AbstractDomain{T,D} end
""" Function space in D-Dimensional space """
abstract type AbstractSpace{T,D} end
""" Boundary condition on domain in D-Dimensional space """
abstract type AbstractBonudaryCondition{T,D} end

AbstractSupertypes{T,D} = Union{
                                AbstractField{T,D},
                                AbstractOperator{T,D},
                                AbstractSpace{T,D},
                                AbstractDomain{T,D},
                                AbstractBonudaryCondition{T,D}
                               }

Base.eltype(::AbstractSupertypes{T,D}) where{T,D} = T
dims(::AbstractSupertypes{T,D}) where{T,D} = D

###
# AbstractOperator subtypes
###

""" Tensor product operator in D-Dimensional space """
abstract type AbstractTensorProductOperator{T,D} <: AbstractOperator{T,D} end

""" Gather-Scatter operator in D-Dimensional space """
abstract type AbstractGatherScatterOperator{T,D} <: AbstractOperator{T,D} end

###
# AbstractSpace subtypes
###

""" D-Dimensional spectral space """
abstract type AbstractSpectralSpace{T,D} <: AbstractSpace{T,D} end

""" D-Dimensional tensor-product space """
abstract type AbstractTensorProductSpace{T,D} <: AbstractSpace{T,D} end

###
# other abstract types
###

""" Boundary Condition on D-Dimensional domain """
abstract type AbstractBoundaryCondition{T,D} end

#""" Deform D-Dimensional domain """
#abstract type AbstractDeformation{D} end

# misc
include("utils.jl")
include("NDgrid.jl")

# field
include("Field.jl")

# operators
include("OperatorBasics.jl")
include("Operators.jl")

# domain
include("Domain.jl")

# space
include("Space.jl")
include("DeformSpace.jl")

# polynomial spaces
include("LagrangeMatrices.jl")
include("LagrangePolynomialSpace.jl")
#include("GalerkinOperators.jl")

# fourier spaces
#include("FourierSpace.jl")

# boundary value problem
#include("BoundaryValueProblem.jl")

export
       # Domains
       IntervalDomain, BoxDomain,
       deform, end_points, isperiodic,
       unit_box, annulus_2D,

       # space functionality
       grid, gradOp, massOp, laplaceOp, advectionOp, interpOp,
#      BoundaryCondition,

       # spaces
       LagrangePolynomialSpace,
       GaussLobattoLegendre1D, GaussLegendre1D, GaussChebychev1D,
       GaussLobattoLegendre2D, GaussLegendre2D, GaussChebychev2D

       # boundary value problem
#      BoundaryValuePDEProblem

end # module


# This file is a part of MGVI.jl, licensed under the MIT License (MIT).

using Distributions
using DistributionsAD
using LinearAlgebra
import Zygote

"""
    _fi_cov_only_upper_params(s; offset=0)

Return indices of the covariance part of the FI matrix that
correspond to upper triangular part of the covariance matrix.

Specify offset to get indices shifted to the right and bottom by
the same offset. (useful to remove irrelevant cols from full FI)
"""
function _fi_cov_only_upper_params(s; offset=0)
    slices = []
    for i in 1:s
        start = (i-1)*s+1
        push!(slices, (start:start+i-1))
    end
    cov_slices = union(slices...) .+ offset
    relevant_slices = union(1:offset, cov_slices)
    relevant_slices
end

# fisher information MC implementation

function _grad_logpdf(model::Function, params::AbstractVector, data_point)
    g = Zygote.gradient(dp -> logpdf(model(dp), data_point), params)[1]
    g_flat = collect(g)
    g_flat * g_flat'
end

_cut_params(res, dist::Distribution) = res

"""
    _cut_params(res, dist::Distribution)

noop for most of the distributions.

for MvNormal it removes rows and cols from FI that correspond to lower triangular
part of the covariance matrix. Covariance matrix is symmetric, so lower triangular
parameters are redundand.
"""
function _cut_params(res, dist::MvNormal)
    mean_size = length(dist.μ)
    slices = _fi_cov_only_upper_params(mean_size; offset=mean_size)
    res[slices, slices]
end

function fisher_information_mc(model::Function, params::AbstractVector, n::Integer)
    dist = model(params)
    sample() = _grad_logpdf(model, params, rand(dist))/n

    res = [sample() for _ in 1:Threads.nthreads()]
    Threads.@threads for i in 1:n-Threads.nthreads()
        res[Threads.threadid()] .+= sample()
    end

    _cut_params(sum(res), dist)
end

# end fisher information mc

"""
    explicit_mv_normal_fi(variance::AbstractMatrix)

Compute fisher information of MvNormal based on the wiki article:

https://en.wikipedia.org/wiki/Fisher_information#Multivariate_normal_distribution
"""
function explicit_mv_normal_fi(variance::AbstractMatrix)
    var_dims = length(variance)
    mean_dims = size(variance, 1)
    dims = mean_dims + var_dims

    res = zeros(dims, dims)
    inv_var = inv(variance)
    res[1:mean_dims, 1:mean_dims] .= inv_var
    for m in 1:var_dims
        for n in 1:var_dims
            m_mat = zeros(mean_dims, mean_dims)
            m_mat[m] = 1
            m_mat = Symmetric(m_mat)
            n_mat = zeros(mean_dims, mean_dims)
            n_mat[n] = 1
            n_mat = Symmetric(n_mat)
            res[mean_dims + m, mean_dims + n] = tr(inv_var * m_mat * inv_var * n_mat)/2
        end
    end

    slices = _fi_cov_only_upper_params(mean_dims; offset=mean_dims)
    res[slices, slices]
end


# vortex particle definition

# import modules
import LinearAlgebra
la = LinearAlgebra

"""
`Particle(args)`

Description here.
"""
mutable struct Particle

    # properties here
    gamma::Array{Number, 1}       # vortex strength vector
    x::Array{Number, 1}           # particle position
    v::Array{Number, 1}           # particle velocity

    Particle(gamma,x,v) = new(gamma,x,v)

end

"""
`vinduced(particle::Particle,X)`

Calculates the induced velocity produced by the vortex particle argument at the specified location.
"""
function vinduced(particle,x)
    r = x .- particle.x
    vinduced = -la.cross(r/4/pi/la.norm(r)^3, particle.gamma)

    return vinduced
end

"""
`place(coordinates::Array{Array{Number,1},1},gammas::Array{Array{Number,1},1},vs::Array{Array{Number,1},1})`

Places vortex particles at the specified locations.
"""
function place(coordinates, gammas, vs)
    particles = Particle[]
    for i = range(1,length=length(coordinates))
        push!(particles,Particle(gammas[i],coordinates[i],vs[i]))
    end

    return particles
end

"""
`netv(particles::Array{Particle,1},x::Array{Particle,1})`

Returns the net velocity induced by the vector of Particles at location x.
"""
function netv(particles, x; err = 0.001)
    # find velocity induced at a point due to all vortex particles
    velocity = [0.0, 0.0, 0.0]
    for particle in particles
        if la.norm(particle.x .- x) < err
            continue
        else
            velocity .+= vinduced(particle, x)
        end
    end

    return velocity
end

"""
`step(particles::Array{Particle,1},timestep::Number,Uinf)`

Calculates induced velocities and advances forward the specified timestep.
"""
function advance(particles,timestep,Uinf)
    # update velocity of each particle
    for particle in particles
        # yogi = netv(particles, particle.x) .+ Uinf(particle.x)
        particle.v = netv(particles, particle.x) .+ Uinf(particle.x)#yogi # WHY DO I NEED YOGI????
    end

    # move each particle
    for particle in particles
        particle.x .= particle.x + particle.v * timestep
    end
end

"""
`vfield(particles::Array{Particle,1},xs::Array{Number,1},ys::Array{Number,1},zs::Array{Number,1})`

Evaluates the velocity field at discrete points on a cartesian grid defined by the arguments:

    * `xs` defines the X values of the cartesian grid
    * `ys` defines the Y values of the cartesian grid
    * `zs` defines the Z values of the cartesian grid

Returns six vectors:

    * `gridx` defines the X coordinate of each grid point
    * `gridy` defines the Y coordinate of each grid point
    * `gridz` defines the Z coordinate of each grid point
    * `vx` defines the X component of the vorticity-induced velocity at each grid point
    * `vy` defines the Y component of the vorticity-induced velocity at each grid point
    * `vz` defines the Z component of the vorticity-induced velocity at each grid point
"""
function vfield(particles,xs, ys, zs)
    # build grid
    gridx = Number[]
    gridy = Number[]
    gridz = Number[]
    vx = Number[]
    vy = Number[]
    vz = Number[]

    for x in xs
        for y in ys
            for z in zs
                # add coordinate
                push!(gridx,x)
                push!(gridy,y)
                push!(gridz,z)
                v = netv(particles,[x,y,z]) # get induced velocity
                if la.norm(v) > 2.0
                    v = v./la.norm(v)*2.0
                end
                # unpack velocity components
                push!(vx,v[1])
                push!(vy,v[2])
                push!(vz,v[3])
            end
        end
    end

    return gridx, gridy, gridz, vx, vy, vz
end

"""
`gfield(particles::Array{Particle,1})`

Useful for plotting the circulation vector of each particle

Returns six vectors:

    * `gridx` defines the X coordinate of each particle
    * `gridy` defines the Y coordinate of each particle
    * `gridz` defines the Z coordinate of each particle
    * `gx` defines the X component of the circulation of each particle
    * `gy` defines the Y component of the circulation of each particle
    * `gz` defines the Z component of the circulation of each particle
"""
function gfield(particles)
    # set up vectors
    gridx = Number[]
    gridy = Number[]
    gridz = Number[]
    gx = Number[]
    gy = Number[]
    gz = Number[]
    # get coordinates and circulations
    for particle in particles
        push!(gridx,particle.x[1])
        push!(gridy,particle.x[2])
        push!(gridz,particle.x[3])
        push!(gx,particle.gamma[1])
        push!(gy,particle.gamma[2])
        push!(gz,particle.gamma[3])
    end

    return gridx, gridy, gridz, gx, gy, gz
end

"""
`vorticity(particles::Array{Particle,1}, x::Vector{Number})`

Returns the particle-induced vorticity at coordinates x.
"""
function vorticity(particles, x)
    # update velocity of each particle
    for particle in particles

    end

end


function mtdmin(
    latt_vars::LatticeVars, atm_vars::AtomicVars, rmt
)
    v1 = zeros(Float64,3)
    v2 = zeros(Float64,3)
    is = 1
    js = 1
    dmin = 1.0e6

    avec = latt_vars.avec
    epslat = latt_vars.epslat
    atposc = atm_vars.atposc
    nspecies = atm_vars.nspecies
    natoms = atm_vars.natoms

    for i1 in -1:1, i2 in -1:1, i3 in -1:1
        #
        @views v1[:] = i1*avec[:,1] + i2*avec[:,2] + i3*avec[:,3]
        #
        for ks in 1:nspecies, ka in 1:natoms[ks]
            #
            @views v2[:] = v1[:] + atposc[:,ka,ks]
            #
            for ls in 1:nspecies
                #
                t1 = rmt[ks] + rmt[ls]
                #
                for la in 1:natoms[ls]
                    if ( (i1 != 0) || (i2 != 0) || (i3 != 0) || (ks != ls) ||  (ka != la) )
                        t2 = sqrt( (v2[1] - atposc[1,la,ls])^2 +
                                   (v2[2] - atposc[2,la,ls])^2 +
                                   (v2[3] - atposc[3,la,ls])^2 )
                        t3 = t2 - t1
                        if t3 < (dmin - epslat)
                            is = ks
                            js = ls
                            dmin = t3
                        end
                    end
                end
            end
        end
    end

    return dmin, is, js

end


#! /usr/bin/env julia

using Arcus
using Base.Test

import Base.isapprox
isapprox(a::Float64, b::Float64) = isapprox(a, b, atol = 2eps(typeof(a))) || isapprox(1/a, 1/b, atol = 2eps())

@testset begin

@testset "creation" begin

  @test Arc(0.0) != nothing
  @test Arc(pi/2) != nothing
  @test Arcd(44) != nothing
  @test Arc(1.0, 1.0) != nothing
  @test_throws MethodError Arc(1.0, 2.0, 3.0)

end

@testset "dedicated values" begin
  @test Arc(0.0) == zero(Arc)
  @test Float64(Arcd(90)) == pi / 2
  @test Float64(Arcd(180)) == 1pi
  @test Float64(Arcd(270)) == -pi / 2
  @test Arcd(360) == Arcd(0)
end

# @testset "arithmetic" begin
  @test -Arc(pi*3/8) == Arc(-pi*3/8)
  @test -Arc(pi*3/4) == Arc(-pi*3/4)
  @test -Arc(pi*5/4) == Arc(-pi*5/4)
  @test -Arc(pi*7/4) == Arc(-pi*7/4)
  @test -Arc(pi/3) == Arc(5pi/3)
  @test Arc(1) + Arc(1) == 2Arc(1)
  @test Arc(1) + Arc(2) ≈ Arc(3)
  @test Arc(3) - Arc(2) ≈ Arc(1)
  @test -Arc(1) + Arc(3) ≈ Arc(2)
  @test Arcd(190) / 170 ≈ Arcd(-1.0)
  @test Arc(0.3) * 30 ≈ Arc(mod(9.0, 2pi))
  @test 20Arcd(260) == Arcd(260) * 20
# end

@testset "representation as sine 0" begin
  @test reinterpret(Float64, Arc(0)) == 0
end

@testset "representation as sine $(w)" for w in linspace(-pi/4, pi/4, 11)
  @test reinterpret(Float64, Arc(w)) ≈ sin(w)
end

const logsp = logspace(-16, 306, 5) * (nextfloat(0.0)*1e16)

@testset "representation as sine $(w)" for w in logsp
  @test reinterpret(Float64, Arc(w)) == sin(w)
end

const ep = eps(2.0)
@testset "trigonometric functions $w" for w in 0:45:360
  @test sin(Arcd(w)) ≈ sind(w)
  @test cos(Arcd(w)) ≈ cosd(w)
  @test tan(Arcd(w)) ≈ tand(w)
  @test csc(Arcd(w)) ≈ cscd(w)
  @test sec(Arcd(w)) ≈ secd(w)
  @test cot(Arcd(w)) ≈ cotd(w)
end

@testset "trigonometric functions $w" for w in 2/7:2/3:2pi
  @test sin(Arc(w)) ≈ sin(w)
  @test cos(Arc(w)) ≈ cos(w)
  @test tan(Arc(w)) ≈ tan(w)
  @test csc(Arc(w)) ≈ csc(w)
  @test sec(Arc(w)) ≈ sec(w)
  @test cot(Arc(w)) ≈ cot(w)
end

end


using InfiniteOpt, Ipopt, Plots
#using KNITRO

exponential_relaxation(z, δ) = exp(one(typeof(z)) - z / δ) - one(typeof(z)) - log(δ)
function relaxed_log_barrier(z; δ=0.3f0)
    return max(z > δ ? -log(z) : exponential_relaxation(z, δ), zero(typeof(z)))
end
# Set the SEIR parameters
γ = 0.303
β = 0.727
N = 1e4
ξ = 0.2

# Set the domain information
i_max = 0.001
t0 = 0
tf = 50

# Set the intial condition values
e0 = 1 / N
i0 = 0
r0 = 0
s0 = 1 - 1 / N

model = InfiniteModel(Ipopt.Optimizer)

@infinite_parameter(model, t ∈ [t0, tf], num_supports = 201,
    derivative_method = OrthogonalCollocation(2))

@variable(model, s ≥ 0, Infinite(t))
@variable(model, e ≥ 0, Infinite(t))
@variable(model, i ≥ 0, Infinite(t))
@variable(model, r ≥ 0, Infinite(t))
@variable(model, 0.2 ≤ u ≤ 0.8, Infinite(t), start = 0.2)

@objective(model, Min, ∫(u^2, t))

# Define the initial conditions
@constraint(model, s(0) == s0)
@constraint(model, e(0) == e0)
@constraint(model, i(0) == i0)
@constraint(model, r(0) == r0)

# Define the SEIR equations
@constraint(model, s_constr, ∂(s, t) == -(1 - u) * β * s * i)
@constraint(model, e_constr, ∂(e, t) == (1 - u) * β * s * i - ξ * e)
@constraint(model, i_constr, ∂(i, t) == ξ * e - γ * i)
@constraint(model, r_constr, ∂(r, t) == γ * i)

# Define the infection rate limit
@constraint(model, imax_constr, i ≤ i_max)

print(model)

optimize!(model)

r_opt = value(r, ndarray=true) * 100 # make the population fractions into percentages
s_opt = value(s, ndarray=true) * 100
i_opt = value(i, ndarray=true) * 100
e_opt = value(e, ndarray=true) * 100
u_opt = value(u)
@show obj_opt = objective_value(model)
ts = value(t)

p = plot(ts, i_opt, label="i(t)")
plot!(ts, e_opt, label="e(t)")
ylabel!("Pop. (%)")
xlabel!("Time (Days)")
display(p)

p1 = plot(ts, u_opt, label="u(t)")
xlabel!("Time (Days)")
ylabel!("Distancing Ratio")
display(p1)

using DiffEqFlux, DifferentialEquations
const solveeq = DifferentialEquations.solve

tspan = (0.0f0, 50.0f0)
ann = FastChain(FastDense(1, 32, relu), FastDense(32, 1), (x, p) -> 0.6 * sigmoid.(x) .+ 0.2)
θ = initial_params(ann)
function sir_nn!(du, u, p, t)
    S, E, I, R, z = u
    du[1] = -(1 - ann([t], p)[1]) * β * S * I
    du[2] = (1 - ann([t], p)[1]) * β * S * I - ξ * E
    du[3] = ξ * E - γ * I
    du[4] = γ * I
    du[5] = (ann([t], p)[1])^2
    du[6] = (relu(I - i_max))^2
end
u0 = [s0, e0, i0, r0, 0, 0]
ts = Float32.(collect(0.0:1:tspan[2]))
prob = ODEProblem(sir_nn, u0, tspan, θ)
sol_init = solveeq(prob, Vern9())
plot(sol_init)
function predict_adjoint(θ)
    Array(solveeq(prob, Vern9(), p=θ, saveat=ts, sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true))))
end

# function quadratic_relaxation(z, δ)
#     return one(typeof(z)) / 2 * (((z - 2δ) / δ)^2 - one(typeof(z))) - log(δ)
# end
#
λ = 10000
for ii = 1:2
    λ = 10* λ
    function loss_adjoint(θ)
        x = predict_adjoint(θ)
        # I_barrier_loss = map(y -> relaxed_log_barrier(i_max - y; δ = 100 / N), x[3, :])
        λ * x[6, end] + x[5, end]
    end
    @show l = loss_adjoint(θ)
    cb = function (θ, l)
        println(l)
        p = plot(Array(solveeq(remake(prob, p=θ), Tsit5(), saveat=0.5))[3, :], lw=3)
        p1 = plot(ts, [first(ann([t], θ)) for t in ts], label="u(t)", lw=3)
        display(p)
        display(p1)
        return false
    end
    # Display the ODE with the current parameter values.
    #cb(θ, l)
    #loss1 = loss_adjoint(θ)
    res1 = DiffEqFlux.sciml_train(loss_adjoint, θ, ADAM(0.05), cb=cb, maxiters=100)
    θ = res1.u
    res2 = DiffEqFlux.sciml_train(loss_adjoint, θ,
        BFGS(initial_stepnorm=0.01), cb=cb, maxiters=100,
        allow_f_increases=false)
    θ = res2.u
end
p = plot(Array(solveeq(remake(prob, p=θ), Tsit5(), saveat=0.5))[3, :], lw=3)
# Part Three: Saving the results
display(p)
using DataFrames
using CSV
p_value = res1.u
results_saving_opt = DataFrame()
results_saving_opt.t = tspan
results_saving_opt.I = I_value
results_saving_opt.u = u_value
results_saving_opt.I_nn = Array(solveeq(remake(prob, p=p_value), Tsit5(), saveat=0.2))[3, :]
results_saving_opt.u_nn = [first(ann([t], p_value)) for t in ts]
CSV.write("results_saving_SEIRconstraintmodel.csv", results_saving_opt)

plot(results_saving_opt.t, results_saving_opt.I, lw=2, label="I(t) by direct collection")
plot!(results_saving_opt.t, results_saving_opt.I_nn, lw=2, label="I(t) by deep learning")
xlabel!("t(day)")
ylabel!("I(t)")
savefig("SEIRconstraint_I.png")

plot(results_saving_opt.t[2:end-1], results_saving_opt.u[2:end-1], lw=2, label="u(t) by direct collection")
plot!(results_saving_opt.t[2:end-1], results_saving_opt.u_nn[2:end-1], lw=2, label="u(t) by deep learning")
xlabel!("t(day)")
ylabel!("u(t)")
savefig("SEIRconstraint_control.png")

# Loss saving

opt_method_loss = 13.90369069230782
nn_method_loss = 13.711496454176164


function forward_step(body::Body, dt)
	acc = accel(body)
	body.pos += body.vel * dt
	body.vel += acc * dt
end

function leapfrog_step(body::Body, dt)
	body.vel += accel(body) * 0.5 * dt
	body.pos += body.vel * dt
	body.vel += accel(body) * 0.5 * dt
end

function rk2_step(body::Body, dt)
	old_pos = body.pos
	half_vel = body.vel + accel(body)*0.5*dt
	body.pos += body.vel * 0.5 * dt
	body.vel += accel(body) * dt
	body.pos = old_pos + half_vel*dt
end

function rk4_step(body::Body, dt)
	old_pos = body.pos
	a0 = accel(body)
	body.pos = old_pos + body.vel*0.5*dt + a0*0.125*dt*dt
	a1 = accel(body)
	body.pos = old_pos + body.vel*dt + a1*0.5*dt*dt
	a2 = accel(body)
	body.pos = old_pos + body.vel*dt + (a0 + a1*2)*(1/6.0)*dt*dt
	body.vel += (a0 + a1*4 + a2) * (1/6.0) * dt
end

function yo6_step(body::Body, dt)
	d = [0.784513610477560e0, 0.235573213359357e0,
				-1.17767998417887e0, 1.31518632068391e0 ]

	for i in 1:3; leapfrog_step(body, dt*d[i]); end;
	leapfrog_step(body, dt*d[4])
	for i in 0:2; leapfrog_step(body, dt*d[3-i]); end;
end

function yo8_step(body::Body, dt)
	d = [0.104242620869991e1, 0.182020630970714e1, 0.157739928123617e0,
				0.244002732616735e1, -0.716989419708120e-2, -0.244699182370524e1,
				-0.161582374150097e1, -0.17808286265894516e1]

	for i in 1:7; leapfrog_step(body, dt*d[i]); end;
	leapfrog_step(body, dt*d[8])
	for i in 0:6; leapfrog_step(body, dt*d[7-i]); end;
end

######## Multistep ########

function ms2_step(body::Body, dt)
	if body.step == 0
		set_prev_accel!(body, 1, accel(body))
		rk2_step(body, dt)
	else
		old_acc = accel(body)
		jdt = old_acc - prev_accel(body, 1)
		body.pos += body.vel*dt + 0.5*old_acc*dt*dt
		body.vel += old_acc*dt + 0.5*jdt*dt
		set_prev_accel!(body, 1, old_acc)
	end
end

function ms4_step(body::Body, dt)
	if body.step <= 2
		body.prev_accel[:,3-body.step] = accel(body)
		rk4_step(body, dt)
	else
		ap0 = accel(body)
		jdt = ap0*(11.0/6.0) - 3*body.prev_accel[:,1] + 1.5*body.prev_accel[:,2] - body.prev_accel[:,3]/3
		sdt2 = 2*ap0 - 5*body.prev_accel[:,1] + 4*body.prev_accel[:,2] - body.prev_accel[:,3]
		cdt3 = ap0 - 3*body.prev_accel[:,1] + 3*body.prev_accel[:,2] - body.prev_accel[:,3]
		body.pos += body.vel*dt + (ap0/2.0 + jdt/6.0 + sdt2/24.0)*dt*dt
		body.vel += ap0*dt + (jdt/2.0 + sdt2/6.0 + cdt3/24.0)*dt
		# Update prev acc
		body.prev_accel[:,3] = body.prev_accel[:,2]
		body.prev_accel[:,2] = body.prev_accel[:,1]
		body.prev_accel[:,1] = ap0
	end
end

function ms4pc_step(body::Body, dt)
	if body.step == 0
		set_prev_accel!(body, 4, accel(body))
		rk4_step(body, dt)
	elseif body.step == 1
		set_prev_accel!(body, 3	, accel(body))
		rk4_step(body, dt)
	elseif body.step == 2
		set_prev_accel!(body, 2, accel(body))
		rk4_step(body, dt)
		set_prev_accel!(body, 1, accel(body))
	else
		old_pos = body.pos
		old_vel = body.vel
		# ap0 = accel(body)
		jdt = prev_accel(body, 1)*(11.0/6.0) - 3*prev_accel(body, 2) + 1.5*prev_accel(body, 3) - prev_accel(body, 4)/3
		sdt2 = 2*prev_accel(body, 1) - 5*prev_accel(body, 2) + 4*prev_accel(body, 3) - prev_accel(body, 4)
		cdt3 = prev_accel(body, 1) - 3*prev_accel(body, 2) + 3*prev_accel(body, 3) - prev_accel(body, 4)
		body.pos += body.vel*dt + (prev_accel(body, 1)/2.0 + jdt/6.0 + sdt2/24.0)*dt*dt

		set_prev_accel!(body, 4, prev_accel(body, 3))
		set_prev_accel!(body, 3, prev_accel(body, 2))
		set_prev_accel!(body, 2, prev_accel(body, 1))
		set_prev_accel!(body, 1, accel(body))

		jdt = prev_accel(body, 1)*(11.0/6.0) - 3*prev_accel(body, 2) + 1.5*prev_accel(body, 3) - prev_accel(body, 4)/3
		sdt2 = 2*prev_accel(body, 1) - 5*prev_accel(body, 2) + 4*prev_accel(body, 3) - prev_accel(body, 4)
		cdt3 = prev_accel(body, 1) - 3*prev_accel(body, 2) + 3*prev_accel(body, 3) - prev_accel(body, 4)

		body.vel = old_vel + prev_accel(body, 1)*dt + (-jdt/2.0 + sdt2/6.0 - cdt3/24.0)*dt
		body.pos = old_pos + body.vel*dt + (-prev_accel(body, 1)/2.0 + jdt/6.0 - sdt2/24.0)*dt*dt
	end
end

function hermite_step(body::Body, dt)
	old_pos = body.pos
	old_vel = body.vel
	old_acc = accel(body)
	old_jerk = jerk(body)

	body.pos += body.vel*dt + old_acc*(dt*dt/2.0) + old_jerk*(dt*dt*dt/6.0)
	body.vel += old_acc*dt + old_jerk*(dt*dt/2.0)
	body.vel = old_vel + (old_acc + accel(body))*(dt/2.0) + (old_jerk - jerk(body))*(dt*dt/12.0)
	body.pos = old_pos + (old_vel + body.vel)*(dt/2.0) + (old_acc - accel(body))*(dt*dt/12.0)
end


abstract type AbstractDiscretizer end

struct QuantileSlicingDiscretizer <: AbstractDiscretizer
    nbins::Int64
end

function learn_discretizations(
    disc::QuantileSlicingDiscretizer,
    Ps_train,
    Xs_train::AbstractVector{<:Real},
    Xs_test::AbstractVector{<:Real},
    α,
    multiple_testing_method,
)

    @unpack nbins = disc
    qs = [-Inf; quantile([Xs_train; Xs_test], (1:(nbins-1)) ./ nbins); +Inf]
    Xs_train_d = cut(Xs_train, qs, allowempty = true)
    Xs_test_d = cut(Xs_test, qs, allowempty = true)
    (Xs_train_d, Xs_test_d, qs)
end

struct DiscretizationWeightLearner{CW<:CategoricalWeightLearner,D<:AbstractDiscretizer} <:
       WeightLearner
    wlearner::CW
    discretizer::D
end

function learn_weights(
    dw::DiscretizationWeightLearner,
    Ps_train,
    Xs_train,
    Xs_test,
    α,
    multiple_testing_method,
)
    # TODO: check levels are the same

    Xs_train_d, Xs_test_d, discr_scheme = learn_discretizations(
        dw.discretizer,
        Ps_train,
        Xs_train,
        Xs_test,
        α,
        multiple_testing_method,
    )

    ws, ws_fit = learn_weights(
        dw.wlearner,
        Ps_train,
        Xs_train_d,
        Xs_test_d,
        α,
        multiple_testing_method,
    )
    discr_ws_fit = (discr_fit = discr_scheme, ws_fit = ws_fit)
    (ws, discr_ws_fit)
end


#Based upon alg 2 in https://ieeexplore.ieee.org/document/8573778

"""
$(TYPEDEF)
`SR3` is an optimizer framework introduced [by Zheng et. al., 2018](https://ieeexplore.ieee.org/document/8573778) and used within
[Champion et. al., 2019](https://arxiv.org/abs/1906.10612). `SR3` contains a sparsification parameter `λ`, a relaxation `ν`.
It solves the following problem
```math
\\argmin_{x, w} \\frac{1}{2} \\| Ax-b\\|_2 + \\lambda R(w) + \\frac{\\nu}{2}\\|x-w\\|_2
```
Where `R` is a proximal operator and the result is given by `w`.

# Fields
$(FIELDS)

# Example
```julia
opt = SR3()
opt = SR3(1e-2)
opt = SR3(1e-3, 1.0)
opt = SR3(1e-3, 1.0, SoftThreshold())
```
## Note
Opposed to the original formulation, we use `ν` as a relaxation parameter,
as given in [Champion et. al., 2019](https://arxiv.org/abs/1906.10612). In the standard case of
hard thresholding the sparsity is interpreted as `λ = threshold^2 / 2`, otherwise `λ = threshold`.
"""
mutable struct SR3{T, V, P <: AbstractProximalOperator} <: AbstractOptimizer{T}
   """Sparsity threshold"""
   λ::T
   """Relaxation parameter"""
   ν::V
   """Proximal operator"""
   R::P

   function SR3(threshold::T = 1e-1, ν::V = 1.0, R::P = HardThreshold()) where {T,V <: Number,P <: AbstractProximalOperator}
      @assert all(threshold .> zero(eltype(threshold))) "Threshold must be positive definite"
      @assert ν > zero(V) "Relaxation must be positive definite"

       λ = isa(R, HardThreshold) ? threshold.^2 /2 : threshold
       return new{typeof(λ), V, P}(λ, ν, R)
   end

   function SR3(threshold::T , R::P) where {T,P <: AbstractProximalOperator}
      @assert all(threshold .> zero(eltype(threshold))) "Threshold must be positive definite"
       λ = isa(R, HardThreshold) ? threshold.^2 /2 : threshold
       ν = one(eltype(λ))
       return new{typeof(λ), eltype(λ), P}(λ, ν, R)
   end
end

Base.summary(::SR3) = "SR3"

function (opt::SR3{T,V,R})(X, A, Y, λ::V = first(opt.λ);
    maxiter::Int64 = maximum(size(A)), abstol::V = eps(eltype(T)), progress = nothing, kwargs...)  where {T, V, R}

   n, m = size(A)
   ν = opt.ν
   W = copy(X)

   # Init matrices
   H = A'*A+I(m)*ν
   H = cholesky!(H)
   X̂ = A'*Y

   w_i = similar(W)
   w_i .= W
   iters = 0

   iters = 0
   converged = false

   xzero = zero(eltype(X))
   obj = xzero
   sparsity = xzero
   conv_measure = xzero

   _progress = isa(progress, Progress)
   initial_prog = _progress ? progress.counter : 0

   @views while (iters < maxiter) && !converged
       iters += 1

       # Solve ridge regression
       X .= H \ (X̂ .+ W*ν)
       #ldiv!(X, H, X̂ .+ W*ν)
       # Proximal
       opt.R(W, X, λ)

       conv_measure = norm(w_i .- W, 2)

       if _progress
           obj = norm(Y - A*X, 2)
           sparsity = norm(X, 0, λ)

           ProgressMeter.next!(
           progress;
           showvalues = [
               (:Threshold, λ), (:Objective, obj), (:Sparsity, sparsity),
               (:Convergence, conv_measure)
           ]
           )
       end

       if conv_measure < abstol
           converged = true

           #if _progress

            #   ProgressMeter.update!(
            #   progress,
            #   initial_prog + maxiter
            #   )
           #end

       else
           w_i .= W
       end
   end
   # We really search for W here
   X .= W
   @views clip_by_threshold!(X, λ)
   return
end


macro test_binary_operator(op, vals)
    quote
        @test $(op)(false, false) == $(vals)[1]
        @test $(op)(true, false) == $(vals)[2]
        @test $(op)(false, true) == $(vals)[3]
        @test $(op)(true, true) == $(vals)[4]
    end
end

@testset "Expressions" begin
    @testset "Built-in Operators" begin
        @test isdefined(Deductive, :¬)
        @test isdefined(Deductive, :∧)
        @test isdefined(Deductive, :∨)
        @test isdefined(Deductive, :→)
        @test isdefined(Deductive, :⟷)

        @test ¬(true) == false
        @test ¬(false) == true

        @test_binary_operator ∨ (false, true, true, true)
        @test_binary_operator ∧ (false, false, false, true)

        @test_binary_operator → (true, false, true, true)
        @test_binary_operator ⟷ (true, false, false, true)
    end

    @testset "LogicalSymbol Core" begin
        @test (LogicalSymbol <: Deductive.AbstractExpression)

        a = LogicalSymbol(:a)

        my_metadata = "AbCdEfG"
        with_metadata = LogicalSymbol(:b, my_metadata)

        # conversion to julia symbol
        @test Symbol(a) == :a
        @test Symbol(with_metadata) == :b

        # tree / node tools
        @test !istree(a)
        @test isnode(a)
        @test !istree(with_metadata)
        @test isnode(with_metadata)

        # metadata getters
        @test isnothing(metadata(a))
        @test !isnothing(metadata(with_metadata))
        @test metadata(with_metadata) == my_metadata

        # constituent variables
        @test length(variables(a)) == 1
        @test first(variables(a)) === a
        @test length(variables(with_metadata)) == 1
        @test first(variables(with_metadata)) == with_metadata

        # constituent symbols (always none)
        @test isempty(operations(a))
        @test isempty(operations(with_metadata))

        # symbol comparison
        @test isequal(a, a)
        @test isequal(with_metadata, with_metadata)
        @test isequal(a, LogicalSymbol(:a))
        @test isequal(with_metadata, LogicalSymbol(:b, my_metadata))
        @test !isequal(a, LogicalSymbol(:b))
        @test !isequal(with_metadata, LogicalSymbol(:b))
        @test !isequal(a, LogicalSymbol(:a, 0))
        @test !isequal(a, LogicalSymbol(:a, (nothing,)))
    end

    @testset "LogicalOperation Core" begin
        neg_op = Deductive.LogicalOperation(x -> !x, :¬, 1, false, false)
        and_op = Deductive.LogicalOperation(:∧, 2, true, true) do x, y
            x && y
        end
        ternary_op = Deductive.LogicalOperation(:Δ, 3, true, true) do x, y, z
            x && y && z
        end

        @testset "Operation Names" begin
            @test Symbol(neg_op) == :¬
            @test Symbol(and_op) == :∧
            @test Symbol(ternary_op) == :Δ
        end

        @testset "isunary and isbinary" begin
            @test isunary(neg_op)
            @test !isbinary(neg_op)
            @test !isunary(and_op)
            @test isbinary(and_op)
            @test !isunary(ternary_op)
            @test !isbinary(ternary_op)
        end

        @testset "Equality Checking" begin
            # two operators are equal if their names, arguments, and associative / commutative properties are the same
            @test isequal(neg_op, Deductive.LogicalOperation((x) -> !x, :¬, 1, false, false))
            @test !isequal(neg_op, Deductive.LogicalOperation((x) -> !x, :¬, 2, false, false))
            @test !isequal(neg_op, Deductive.LogicalOperation((x) -> !x, :-, 1, false, false))
            @test isequal(neg_op, Deductive.LogicalOperation((x) -> x, :¬, 1, false, false))
            @test !isequal(neg_op, Deductive.LogicalOperation((x) -> x, :¬, 1, true, false))
            @test !isequal(neg_op, Deductive.LogicalOperation((x) -> x, :¬, 1, false, true))
        end

        @testset "Material Operation" begin
            @test neg_op(false)
            @test !neg_op(true)
            @test and_op(true, true)
            @test !and_op(true, false)
            @test ternary_op(true, true, true)
            @test !ternary_op(false, true, true)
            @test !ternary_op(true, false, true)
            @test !ternary_op(true, true, false)
            @test ternary_op(BitVector((1, 1, 1)))
            @test !ternary_op(BitVector((1, 0, 1)))
        end

        @testset "isassociative and iscommutative" begin
            # isassociative for builtins
            @test isassociative(∧)
            @test isassociative(∨)
            @test isassociative(⟷)
            @test !isassociative(¬)
            @test !isassociative(→)

            # iscommutative for builtins
            @test iscommutative(∧)
            @test iscommutative(∨)
            @test iscommutative(⟷)
            @test !iscommutative(¬)
            @test !iscommutative(→)
        end
    end

    @testset "LogicalExpression Core" begin
        @test (Deductive.LogicalExpression <: Deductive.AbstractExpression)

        a, b, c, d = LogicalSymbol.([:a, :b, :c, :d])
        st = a ∧ b

        @testset "Basics" begin
            @test st isa Deductive.LogicalExpression
            @test istree(st)
            @test !isnode(st)
            @test operation(st) === ∧
            @test all(arguments(st) .== (a, b))
            @test isnothing(metadata(st))
        end

        @testset "Equality Checking" begin
            @test isequal(st, st)
            @test isequal(st, a ∧ b)
            @test isequal(a ∧ b, a ∧ b)
            @test !isequal(a ∧ b, a ∨ b)
            @test !isequal(¬a, ¬¬a)
            @test !isequal((a ∧ b) ∧ c, a ∧ (b ∧ c))
            @test !isequal(a ∧ b, b ∧ a)
            @test !isequal(a ∧ b, a ∧ c)
        end

        @testset "Variables Tree" begin
            @test variables(¬a) == Set([a])
            @test variables(¬¬a) == Set([a])
            @test variables(a ∧ b) == Set([a, b])
            @test variables(¬(a ∨ b)) == Set([a, b])
            @test variables(¬(a ∨ b) → a) == Set([a, b])
            @test variables(¬(a ∨ b) → c) == Set([a, b, c])
        end

        @testset "Operations Tree" begin
            @test operations(¬a) == Set([¬])
            @test operations(¬¬a) == Set([¬])
            @test operations(a ∧ b) == Set([∧])
            @test operations(a ∧ ¬b) == Set([∧, ¬])
            @test operations(a ∧ ¬(b ∨ c)) == Set([∧, ∨, ¬])
            @test operations(¬(a ∨ b) → a) == Set([∨, ¬, →])
        end

        @testset "Binary Operator Methods" begin
            @test left(a ∧ b) == a
            @test right(a ∧ b) == b
            @test left((a ∧ b) ∨ (b ∧ c)) == (a ∧ b)
            @test right((a ∧ b) ∨ (b ∧ c)) == (b ∧ c)
        end

        @testset "Associative Expressions" begin
            @test isassociative(a ∧ b)
            @test isassociative(a ∨ b)
            @test isassociative(a ∨ b ∨ c)
            @test isassociative(a ⟷ b ⟷ c)
            @test !isassociative(¬a)
            @test !isassociative(a → b)
            @test !isassociative(a → b → c)
        end

        @testset "Commutative Expressions" begin
            @test iscommutative(a ∧ b)
            @test iscommutative(a ∨ b)
            @test iscommutative(a ∨ b ∨ c)
            @test iscommutative(a ⟷ b ⟷ c)
            @test !iscommutative(¬a)
            @test !iscommutative(a → b)
            @test !iscommutative(a → b → c)
        end

        @testset "Associative Orderings" begin
            @test associative_ordering(a) == [a]
            @test associative_ordering(b) == [b]
            @test associative_ordering(a ∧ b) == [a, b]
            @test associative_ordering(b ∧ a) == [b, a]
            @test associative_ordering(¬b ∧ a) == [b, a]
            @test associative_ordering((¬b ∨ c) ∧ a) == [b, c, a]
            @test associative_ordering((¬a ∨ (b ∨ c)) ∧ d) == [a, b, c, d]
        end

        @testset "Associative Equality" begin
            @test isequal_associative(a, a)
            @test isequal_associative(b, b)
            @test !isequal_associative(a, b)
            @test isequal_associative(a ∧ b, a ∧ b)
            @test !isequal_associative(a ∧ b, b ∧ a)
            @test !isequal_associative(a ∧ b, a ∨ b)
            @test !isequal_associative(a ∧ ¬b, a ∧ ¬b)
        end

        @testset "Associative Tree Count" begin
            @test Deductive.associative_tree_count(0) == 1
            @test Deductive.associative_tree_count(1) == 1
            @test Deductive.associative_tree_count(2) == 2
            @test Deductive.associative_tree_count(3) == 5
            @test Deductive.associative_tree_count(4) == 14
            @test Deductive.associative_tree_count(5) == 42
            @test Deductive.associative_tree_count(6) == 132
        end

        # mutability tests
        @testset "Copy Expression" begin
            # Base.copy tests
            copied_st = copy(st)
            @test copied_st == st   # still the same expression
            @test copied_st !== st  # but not the same reference
            @test first(arguments(copied_st)) == first(arguments(copied_st))
            @test first(arguments(copied_st)) === first(arguments(copied_st))

            # Base.deepcopy tests
            st2 = (a ∧ b) ∧ c
            dc_st2 = deepcopy(st2)
            @test dc_st2 == st2
            @test dc_st2 !== st2
            @test first(arguments(dc_st2)) == first(arguments(st2))
            @test first(arguments(dc_st2)) !== first(arguments(st2))
            @test arguments(dc_st2)[2] == arguments(st2)[2]
            @test arguments(dc_st2)[2] === arguments(st2)[2]  # the second arg is a symbol so it shouldn't be copied
        end

        @testset "Parent Expressions" begin
            st2 = a ∧ b
            @test isempty(parents(st2))
            st3 = a ∧ st2
            @test st3 ∈ parents(st2)
            st4 = a ∧ st3
            @test st4 ∈ parents(st3)

            # parent removal upon argument reassignment
            st4.arguments[2] = b
            @test st4 ∉ parents(st3)
            st3.arguments[2] = st4
            @test st3 ∈ parents(st4)
        end

        @testset "Mutate Operation" begin
            # simple expression
            st2 = a ∧ b
            @test operation(st2) == ∧
            @test operations(st2) == Set{Deductive.LogicalOperation}([∧])
            st2.operation = ∨
            @test operation(st2) == ∨
            @test operations(st2) == Set{Deductive.LogicalOperation}([∨])
            st2.operation = →
            @test operation(st2) == →
            @test operations(st2) == Set{Deductive.LogicalOperation}([→])
            try
                st2.operation = ¬  # cannot set operation of different argument count
                @test false
            catch e
                @test true
            end

            # nested expression
            st3 = (a ∧ b) → (c ∨ d)
            @test operation(st3) == →
            @test operations(st3) == Set{Deductive.LogicalOperation}([→, ∧, ∨])
            st3.operation = ⟷
            @test operation(st3) == ⟷
            @test operations(st3) == Set{Deductive.LogicalOperation}([⟷, ∧, ∨])
            st3.operation = ∧
            @test operation(st3) == ∧
            @test operations(st3) == Set{Deductive.LogicalOperation}([∧, ∨])
            st3.operation = →
            arguments(st3)[1].operation = →
            @test operation(first(arguments(st3))) == →
            @test operations(st3) == Set{Deductive.LogicalOperation}([→, ∨])
            arguments(st3)[2].operation = →
            @test operation(last(arguments(st3))) == →
            @test operations(st3) == Set{Deductive.LogicalOperation}([→])
        end

        @testset "Mutate Arguments" begin
            # simple expression
            st2 = a ∧ b
            @test arguments(st2) == Vector{Deductive.AbstractExpression}([a, b])
            @test variables(st2) == Set{Deductive.LogicalSymbol}([a, b])
            st2.arguments[1] = b
            @test arguments(st2) == Vector{Deductive.AbstractExpression}([b, b])
            @test variables(st2) == Set{Deductive.LogicalSymbol}([b])
            st2.arguments[2] = a
            @test arguments(st2) == Vector{Deductive.AbstractExpression}([b, a])
            @test variables(st2) == Set{Deductive.LogicalSymbol}([a, b])
            st2.arguments[1:2] = [a, a]
            @test arguments(st2) == Vector{Deductive.AbstractExpression}([a, a])
            @test variables(st2) == Set{Deductive.LogicalSymbol}([a])

            # nested expression
            st3 = (a ∧ b) → (c ∨ d)
            arguments(st3)[1].arguments[1] = b
            @test arguments(st3) == Vector{Deductive.AbstractExpression}([b ∧ b, c ∨ d])
            @test variables(st3) == Set{Deductive.LogicalSymbol}([b, c, d])
            arguments(st3)[2].arguments[2] = a
            @test arguments(st3) == Vector{Deductive.AbstractExpression}([b ∧ b, c ∨ a])
            @test variables(st3) == Set{Deductive.LogicalSymbol}([b, c, a])
        end
    end
end


### first, a simple bayesian estimator ###

confidence(χ) = χ/(1+χ)

function update_confidence(event, p_q, p_a, χ)

    χ *= p_q(event)/p_a(event)
    χ

end

function compute_χ(events, p_q, p_a)
    χ = 1.
    for event in events
        χ = update_confidence(event, p_q, p_a, χ)
    end
    χ
end

"""
    compute_confidence(events,p_q, p_a)

A bayesian confidence estimator: return the probability that the null hypothesis
Q is right compared to the alternative hypothesis A.
"""
function compute_confidence(events,p_q, p_a)

    confidence(compute_χ(events,p_q, p_a))
end

"""
    compute_confidence_array(events, p_q, p_a)

Return an array of the probabilities of H being true as we process more and
more events.
"""
function compute_confidence_array(events, p_q, p_a)

    χ_array = [1.]

    for event in events
        push!(χ_array, update_confidence(event, p_q, p_a, χ_array[end]))
    end

    confidence.(χ_array)

end

### tests with standard boson sampling ###

# we will test that we are indeed in the bosonic case
# compared to the distinguishable one

# first we generate a series of bosonic events

n_events = 10
n = 3
m = 4
interf = RandHaar(m)
input_state = Input{Bosonic}(first_modes(n,m))

events = []

for i in 1:n_events

    # generate a random output pattern
    output_state = FockDetection(random_mode_occupation_collisionless(n,m))

    # compute the event probability
    this_event = Event(input_state, output_state, interf)
    compute_probability!(this_event)
    push!(events, this_event)

end

# now we have the vector of observed events with probabilities

events

# next, from events, recover the probabilities under both
# hypothesis

function p_B(event::Event)

    interf = event.interferometer
    r = event.input_state.r
    input_state = Input{Bosonic}(r)
    output_state = event.output_measurement

    event_H = Event(input_state, output_state, interf)
    compute_probability!(event_H)

    event_H.proba_params.probability

end

function p_D(event::Event)

    interf = event.interferometer
    r = event.input_state.r
    input_state = Input{Distinguishable}(r)
    output_state = event.output_measurement

    event_A = Event(input_state, output_state, interf)
    compute_probability!(event_A)

    event_A.proba_params.probability

end

# hypothesis : the events were from a bosonic distribution

p_q = p_B
p_a = p_D

confidence(compute_χ(events, p_q, p_a))

# hypothesis : the events were from a distinguishable distribution

p_q = p_D
p_a = p_B

confidence(compute_χ(events, p_q, p_a))

@test confidence(compute_χ(events, p_q, p_a)) + confidence(compute_χ(events, p_a, p_q)) ≈ 1 atol = 1e-6

##### the only thing I see is a problem in the probabilities themselves?


### Singletons ###
abstract type ALogicConstant <: AConstantSymbol end

# starts with capital so as not to conflict with `true`
@new_constant ALogicConstant ATrue True Statements
@new_constant ALogicConstant AFalse False Statements
# not yet computed
@new_constant ALogicConstant AUnknown Unknown Statements
# computed and the result could not be obtained within the constraints
@new_constant ALogicConstant AUnknowable Unknowable Statements

### Functions ###
abstract type ALogicFunction <: AFunctionSymbol end

@new_function(ALogicFunction, AAnd, And, (Statements, Statements), Statements, 2, ∧)
@new_function(ALogicFunction, AOr, Or, (Statements, Statements), Statements, 2, ∨)
@new_function(ALogicFunction, ANot, Not, (Statements,), Statements, 2, ~, false)
Base.show(io::IO, self::Term{<:ANot}) = print(io, "$(Meta.parse("$(data(self))($(args(self)[1]))"))")

@new_function(ALogicFunction, AImplies, Implies, (Statements, Statements), Statements, 2, ⟹)
@new_function(ALogicFunction, AEquivalent, Equivalent, (Statements, Statements), Statements, 2, ⩶)

@new_function(ALogicFunction, AElementOf, ElementOf, (Ω, Sets), Statements, 2, ∈, false)
@new_function(ALogicFunction, ASubset, Subset, (Sets, Sets), Statements, 2, ⊂)

@new_function(ALogicFunction, AExists, Exists, (Ω, Sets, Statements), Statements, 3)
Exists(x::Term{ElementOf}, y::Term) = Exists(args(x)[1], args(x)[2], y)
Base.show(io::IO, self::Term{<:AExists}) = print(io, "Exists($(args(self)[1]) ∈ $(args(self)[2]), $(args(self)[3]))")

@new_function(ALogicFunction, AForAll, ForAll, (Ω, Sets, Statements), Statements, 3)
ForAll(x::Term{ElementOf}, y::Term) = ForAll(args(x)[1], args(x)[2], y)
Base.show(io::IO, self::Term{<:AForAll}) = print(io, "ForAll($(args(self)[1]) ∈ $(args(self)[2]), $(args(self)[3]))")


using Gen
import Random

@gen function model()
    if ({:z} ~ bernoulli(0.5))
        m1 = ({:m1} ~ gamma(1, 1))
        m2 = ({:m2} ~ gamma(1, 1))
    else
        m = ({:m} ~ gamma(1, 1))
        (m1, m2) = (m, m)
    end
    {:y1} ~ normal(m1, 0.1)
    {:y2} ~ normal(m2, 0.1)
end

@gen function mean_random_walk_proposal(trace)
    if trace[:z]
        {:m1} ~ normal(trace[:m1], 0.1)
        {:m2} ~ normal(trace[:m2], 0.1)
    else
        {:m} ~ normal(trace[:m], 0.1)
    end
end

@gen function split_merge_proposal(trace)
    if trace[:z]
        # currently two means, switch to one
    else
        # currently one mean, switch to two
        {:u} ~ uniform_continuous(0, 1)
    end
end

function merge_mean(m1, m2)
    m = sqrt(m1 * m2)
    u = m1 / (m1 + m2)
    (m, u)
end

function split_mean(m, u)
    m1 = m * sqrt((u / (1 - u)))
    m2 = m * sqrt(((1 - u) / u))
    (m1, m2)
end

@transform split_merge_involution (model_in, aux_in) to (model_out, aux_out) begin
    if @read(model_in[:z], :discrete)
        # currently two means, switch to one
        @write(model_out[:z], false, :discrete)
        m1 = @read(model_in[:m1], :continuous)
        m2 = @read(model_in[:m2], :continuous)
        (m, u) = merge_mean(m1, m2)
        @write(model_out[:m], m, :continuous)
        @write(aux_out[:u], u, :continuous)
    else
        # currently one mean, switch to two
        @write(model_out[:z], true, :discrete)
        m = @read(model_in[:m], :continuous)
        u = @read(aux_in[:u], :continuous)
        (m1, m2) = split_mean(m, u)
        @write(model_out[:m1], m1, :continuous)
        @write(model_out[:m2], m2, :continuous)
    end
end

is_involution!(split_merge_involution)

function do_inference_simple(y1, y2)
    trace, = generate(model, (), choicemap((:y1, y1), (:y2, y2), (:z, false), (:m, 1.2)))
    zs = Bool[]
    m = Float64[]
    m1 = Float64[]
    m2 = Float64[]
    for iter=1:100
        trace, = mh(trace, select(:z))
        trace, = mh(trace, mean_random_walk_proposal, ())
        push!(zs, trace[:z])
        push!(m, trace[:z] ? NaN : trace[:m])
        push!(m1, trace[:z] ? trace[:m1] : NaN)
        push!(m2, trace[:z] ? trace[:m2] : NaN)
    end
    (zs, m, m1, m2)
end

function do_inference_rjmcmc(y1, y2)
    trace, = generate(model, (), choicemap((:y1, y1), (:y2, y2), (:z, false), (:m, 1.2)))
    zs = Bool[]
    m = Float64[]
    m1 = Float64[]
    m2 = Float64[]
    for iter=1:100
        trace, = mh(
            trace, split_merge_proposal, (), split_merge_involution; check=true)
        trace, = mh(trace, mean_random_walk_proposal, ())
        push!(zs, trace[:z])
        push!(m, trace[:z] ? NaN : trace[:m])
        push!(m1, trace[:z] ? trace[:m1] : NaN)
        push!(m2, trace[:z] ? trace[:m2] : NaN)
    end
    (zs, m, m1, m2)
end

y1, y2 = (1.0, 1.3)
(zs, m, m1, m2) = do_inference_rjmcmc(y1, y2)
println(zs)

using Plots

function make_plots()
    Random.seed!(2)

    y1, y2 = (1.0, 1.3)
    (zs, m, m1, m2) = do_inference_rjmcmc(y1, y2)
    p1 = plot(title="Involution MH (RJMCMC)", m, label="m")
    plot!(m1, label="m1")
    plot!(m2, label="m2")
    ylims!(0.5, 1.5)

    p2 = plot(zs, label="z", color="black")
    xlabel!("# MCMC moves")
    yticks!([0, 1], ["F", "T"])
    ylims!(-0.1, 1.1)

    y1, y2 = (1.0, 1.3)
    (zs, m, m1, m2) = do_inference_simple(y1, y2)
    p3 = plot(title="Selection MH", m, label="m")
    plot!(m1, label="m1")
    plot!(m2, label="m2")
    ylims!(0.5, 1.5)

    p4 = plot(zs, label="z", color="black")
    xlabel!("# MCMC moves")
    yticks!([0, 1], ["F", "T"])
    ylims!(-0.1, 1.1)

    plot(p1, p3, p2, p4)
    savefig("rjmcmc.png")
end

make_plots()


@reexport using Distributions
import Distributions: rand
export MultivariateMix, NonMixedMultivariateDistribution, cov, cor, simulate_nobs_independent_vectors, covariance_matrix, correlation_matrix
# ## create new type not in Distributions for multivariate mixture

"""
    MultivariateMix(vecd::AbstractVector{<:UnivariateDistribution}, Γ)
An N dimensional `MultivariateMix` constructed from a vector of N `UnivariateDistribution`s, and the user-specified covariance matrix, Γ.
Since Distributions.jl does not currently allow for multivariate vectors with mixed continuous and discrete components, this type is not of super type Distributions.MultivariateDistribution.
We will pre-allocate a Vector of `ContinuousUnivariateCopula`s and `DiscreteUnivariateCopula`s to then fill in the appropriate constants c0, c1, c2, recursively.
"""
struct MultivariateMix{
    V<:AbstractVector{UnivariateDistribution},
    T<: BlasReal
}
    vecd::V
    Γ::Matrix{T}
    trΓ::Float64
    gc_obs::Vector{Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}}
    function MultivariateMix(vecd::V, Γ::Matrix{T}) where {T <: BlasReal,
        V<:AbstractVector{UnivariateDistribution}}
        n = length(vecd)
        trΓ = tr(Γ)
        gc_obs = Vector{Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}}(undef, n)
        return new{V, T}(vecd, Γ, trΓ, gc_obs)
    end
end

"""
    NonMixedMultivariateDistribution(vecd::AbstractVector{<:UnivariateDistribution}, Γ)
An N dimensional `MultivariateDistribution` constructed from a vector of N `UnivariateDistribution`s, and the user-specified covariance matrix, Γ.
We will pre-allocate a Vector of `ContinuousUnivariateCopula`s or `DiscreteUnivariateCopula`s to then fill in the appropriate constants c0, c1, c2, recursively.
"""
struct NonMixedMultivariateDistribution{
    S<:ValueSupport,
    T<:UnivariateDistribution{S},
    V<:AbstractVector{T},
}
    vecd::V
    Γ::Matrix{Float64}
    trΓ::Float64
    gc_obs::Vector{Union{DiscreteUnivariateCopula, ContinuousUnivariateCopula}}
    function NonMixedMultivariateDistribution(vecd::V, Γ::Matrix{Float64}) where
        V<:AbstractVector{T} where
        T<:UnivariateDistribution{S} where
        S<:ValueSupport
        n = length(vecd)
        trΓ = tr(Γ)
        gc_obs = Vector{Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}}(undef, n)
        return new{S, T, V}(vecd, Γ, trΓ, gc_obs)
    end
end

####
function rand(gc_vec::NonMixedMultivariateDistribution{S, T, V},
    Y::Vector{Float64},
    res::Vector{Float64}) where {S<: ValueSupport, T <: UnivariateDistribution{S}, V<:AbstractVector{T}}
    for i in 1:length(gc_vec.vecd)
        # form constants for conditional density of i given 1, ..., i-1
        gc_vec.gc_obs[i] = pdf_constants(gc_vec.Γ, res, i, gc_vec.vecd[i])
        # generate y_i given y_1, ..., y_i-1
        Y[i] = rand(gc_vec.gc_obs[i])
        # update residuals 1,..., i-1
        res[i] = update_res!(Y[i], res[i], gc_vec.gc_obs[i])
    end
    res[end] = update_res!(Y[end], res[end], gc_vec.gc_obs[end])
    Y
end

# ### we need a function that is like rand, but cannot use the distributions framework.
function rand(gc_vec::MultivariateMix{V, T},
    Y::Vector{T},
    res::Vector{T}) where {T <: BlasReal, V<:AbstractVector{UnivariateDistribution}}
    for i in 1:length(gc_vec.vecd)
        # form constants for conditional density of i given 1, ..., i-1
        gc_vec.gc_obs[i] = pdf_constants(gc_vec.Γ, res, i, gc_vec.vecd[i])
        # generate y_i given y_1, ..., y_i-1
        Y[i] = rand(gc_vec.gc_obs[i])
        # update residuals 1,..., i-1
        res[i] = update_res!(Y[i], res[i], gc_vec.gc_obs[i])
    end
    res[end] = update_res!(Y[end], res[end], gc_vec.gc_obs[end])
    Y
end

"""
    simulate_nobs_independent_vectors(multivariate_distribution::Union{NonMixedMultivariateDistribution, MultivariateMix}, n_obs::Integer)
Simulate n_obs independent realizations from multivariate copula-like distribution.
"""
function simulate_nobs_independent_vectors(
    multivariate_distribution::Union{NonMixedMultivariateDistribution, MultivariateMix},
    n_obs::Integer)
    dimension = length(multivariate_distribution.vecd)
    Y = [Vector{Float64}(undef, dimension) for i in 1:n_obs]
    res = [Vector{Float64}(undef, dimension) for i in 1:n_obs]
    for i in 1:n_obs
        rand(multivariate_distribution, Y[i], res[i])
    end
    Y
end

"""
    logpdf(d::Union{ContinuousUnivariateCopula, DiscreteUnivariateCopula}, x::Real)
Theoretical log pdf under the copula model.
"""
function logpdf(multivar::Union{MultivariateMix, NonMixedMultivariateDistribution}, Y::Vector{T}) where T <: BlasReal
    logl = zero(T)
    for i in 1:length(multivar.vecd)
        logl += logpdf(multivar.gc_obs[i], Y[i])
    end
    logl
end

"""
    cov(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix}, k::Int64, l::Int64)
Theoretical covariance of the kth and lth element in the random vector.
"""
function cov(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix}, k::Int64, l::Int64)
    cov = sqrt(var(gc_vec.vecd[k])) * sqrt(var(gc_vec.vecd[l])) * gc_vec.Γ[k, l] * inv(1 + 0.5 * gc_vec.trΓ)
    return cov
end

"""
    cor(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix}, k::Int64, l::Int64)
Theoretical correlation of the kth and lth element in the random vector.
"""
function cor(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix}, k::Int64, l::Int64)
    cor = gc_vec.Γ[k, l] / (sqrt(1 + 0.5 * gc_vec.trΓ + gc_vec.Γ[k, k]) * sqrt(1 + 0.5 * gc_vec.trΓ + gc_vec.Γ[l, l]) )
    return cor
end

"""
    covariance_matrix(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix})
Theoretical covariance matrix of a random vector.
"""
function covariance_matrix(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix})
    n = length(gc_vec.gc_obs)
    Covariance = zeros(n, n)
    for i in 1:n
        Covariance[i, i] = QuasiCopula.var(gc_vec.gc_obs[i])
        for j = i+1:n
            Covariance[j, i] = QuasiCopula.cov(gc_vec, j, i)
            Covariance[i, j] = Covariance[j, i]
        end
    end
    Covariance
end

"""
    correlation_matrix(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix})
Theoretical covariance matrix of a random vector.
"""
function correlation_matrix(gc_vec::Union{NonMixedMultivariateDistribution, MultivariateMix})
    n = length(gc_vec.gc_obs)
    Corr = zeros(n, n)
    for i in 1:n
        Corr[i, i] = 1.0
        for j = i+1:n
            Corr[j, i] = QuasiCopula.cor(gc_vec, j, i)
            Corr[i, j] = Corr[j, i]
        end
    end
    Corr
end


function generateall(input)
    cubes = Set{Tuple{Int,Int,Int}}()
    hypercubes = Set{Tuple{Int,Int,Int,Int}}()
    for (i, r) in enumerate(input)
        for (j, c) in enumerate(r)
            if c == '#'
                push!(cubes, (i, j, 0))
                push!(hypercubes, (i, j, 0, 0))
            end
        end
    end
    return cubes, hypercubes
end

function generatecuberange()
    r = (-1:1)
    return collect((x, y, z) for x in r for y in r for z in r if x != 0 || y != 0 || z != 0)
end

function generatehypercuberange()
    r = (-1:1)
    return collect((x, y, z, w) for x in r for y in r for z in r for w in r if x != 0 || y != 0 || z != 0 || w != 0)
end

function nextcubes(cubes)
    newcubes = deepcopy(cubes)
    marked = Dict()
    for (x, y, z) in cubes
        neighbors = 0
        for (dx, dy, dz) in generatecuberange()
            cube = (x + dx, y + dy, z + dz)
            if ∈(cube, cubes)
                neighbors += 1
            else
                push!(marked, cube => (get!(marked, cube, 0) + 1))
            end
        end
        if ∉(neighbors, [2, 3])
            pop!(newcubes, (x, y, z))
        end
    end
    for (cube, count) in marked
        if count == 3
            push!(newcubes, cube)
        end
    end
    return newcubes
end

function nexthypercubes(hypercubes)
    newhypercubes = deepcopy(hypercubes)
    marked = Dict()
    for (x, y, z, w) in hypercubes
        neighbors = 0
        for (dx, dy, dz, dw) in generatehypercuberange()
            cube = (x + dx, y + dy, z + dz, w + dw)
            if ∈(cube, hypercubes)
                neighbors += 1
            else
                push!(marked, cube => (get!(marked, cube, 0) + 1))
            end
        end
        if ∉(neighbors, [2, 3])
            pop!(newhypercubes, (x, y, z, w))
        end
    end
    for (cube, count) in marked
        if count == 3
            push!(newhypercubes, cube)
        end
    end
    return newhypercubes
end

function cycle(something, next, cycle = 6)
    i = 0
    while i < cycle
        something = next(something)
        i += 1
    end
    return something
end

input = readlines("d17.in")
cubes, hypercubes = generateall(input)
cubes = cycle(cubes, nextcubes)
hypercubes = cycle(hypercubes, nexthypercubes)

println("P1: ", length(cubes))
println("P2: ", length(hypercubes))


for N in [Float64, Rational{Int}, Float32]
    # ==============================
    # Testing box approximation
    # ==============================
    # Box approximation of a 2D square
    b = BallInf(N[1, 1], N(1//10))
    h = box_approximation(b)
    hexp = Hyperrectangle(N[1, 1], N[1//10, 1//10])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 2D unit ball in the 1-norm
    b = Ball1(N[1, -2], N(0.2))
    h = box_approximation(b)
    hexp = Hyperrectangle(N[1, -2], N[0.2, 0.2])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 3D unit ball in the 1-norm
    b = Ball1(N[1, 2, 0], N(1))
    h = box_approximation(b)
    hexp = Hyperrectangle(N[1, 2, 0], N[1, 1, 1])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a Hyperrectangle
    hexp = Hyperrectangle(N[0, 0], N[1, 1])
    h = box_approximation(hexp)
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a VPolytope
    P = VPolytope([N[0, 0], N[2, 1], N[1, 2]])
    H = box_approximation(P)
    @test H == Hyperrectangle(N[1, 1], N[1, 1])

    # empty set
    E = EmptySet{N}(2)
    @test box_approximation(E) == E

    # ===================================================================
    # Testing box_approximation_symmetric (= symmetric interval hull)
    # ===================================================================
    # Box approximation of a 2D square
    b = BallInf(N[1, 1], N(1//10))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(N[0, 0], N[11//10, 11//10])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 2D unit ball in the 1-norm
    b = Ball1(N[1, -2], N(2//10))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(N[0, 0], N[12//10, 22//10])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 3D unit ball in the 1-norm
    b = Ball1(N[1, 2, 0], N(1//10))
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(N[0, 0, 0], N[11//10, 21//10, 1//10])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Box approximation of a 4D hyperrectangle
    b = Hyperrectangle(N[-15//10, -25//10, 24//10, -4//10],
                       N[1//10, 2//10, 3//10, 4//10])
    h = box_approximation_symmetric(b)
    hexp = Hyperrectangle(zeros(N, 4), N[16//10, 27//10, 27//10, 8//10])
    @test h.center ≈ hexp.center
    @test h.radius ≈ hexp.radius

    # Testing alias symmetric_interval_hull
    h = symmetric_interval_hull(b)

    # empty set
    E = EmptySet{N}(2)
    @test box_approximation_symmetric(E) == E
end

for N in [Float64]
    # empty intersection due to line search
    X = HalfSpace(N[-1], N(0)) ∩ HalfSpace(N[1], N(-1e-15))
    @test box_approximation(X) isa EmptySet{N}
end


using Interpolations

export interpolatable_field

"""
    interpolatable_field(x,y,f::ScalarGridData)

Generates an interpolatable version of grid data `f`, based on coordinates in
`x` and `y` (which should be in range form). The output can be called as a function
with coordinate pairs as arguments.
"""
function interpolatable_field(x::AbstractArray,y::AbstractArray,f::ScalarGridData)

  (length(x) == size(f,1) && length(y) == size(f,2)) ||
    error("incompatible sizes of field and coordinate data")

  return CubicSplineInterpolation((x, y),f, extrapolation_bc = (Flat(),Flat()))

end

"""
    interpolatable_field(xu,yu,xv,yv,f::EdgeData/NodePair)

Generates an interpolatable version of grid data `f`, based on coordinates in
`xu`, `yu`, `xv`, `yv` (which should be in range form). The output can be called as a function
with coordinate pairs as arguments.
"""
function interpolatable_field(xu::AbstractRange,yu::AbstractRange,
                              xv::AbstractRange,yv::AbstractRange,q::Union{Edges,NodePair})

  return interpolatable_field(xu,yu,q.u), interpolatable_field(xv,yv,q.v)

end

"""
    interpolatable_field(f::GridData,g::PhysicalGrid)

Generates an interpolatable version of grid data `f`, based on grid `g`. The output can be called as a function
with coordinate pairs as arguments.
"""
interpolatable_field(f::GridData,g::PhysicalGrid) = interpolatable_field(coordinates(f,g)...,f)


export SemialgebraicSetsHCSolver

function ModelKit.System(V::SemialgebraicSets.AbstractAlgebraicSet; kwargs...)
    ModelKit.System(SemialgebraicSets.equalities(V); kwargs...)
end

"""
    SemialgebraicSetsHCSolver(; compile = false, options...)

Construct a `SemialgebraicSets.AbstractAlgebraicSolver` to be used in `SemialgebraicSets`.
`options` are all valid options for [`solve`](@ref).

## Example

```
julia> using HomotopyContinuation, SemialgebraicSets;

julia> solver = SemialgebraicSetsHCSolver(; show_progress = false)
SemialgebraicSetsHCSolver(; compile = :none, show_progress = false)

julia> @polyvar x y
(x, y)

julia> V = @set x^2 == 1 && y^2 == 2 solver
Algebraic Set defined by 2 equalities
 x^2 - 1.0 = 0
 y^2 - 2.0 = 0

julia> collect(V)
4-element Array{Array{Float64,1},1}:
 [1.0, 1.414213562373095]
 [1.0, -1.414213562373095]
 [-1.0, 1.414213562373095]
 [-1.0, -1.414213562373095]
```
"""
struct SemialgebraicSetsHCSolver <: SemialgebraicSets.AbstractAlgebraicSolver
    options::Any
end
SemialgebraicSetsHCSolver(; compile = :none, options...) =
    SemialgebraicSetsHCSolver((compile = compile, options...))

function Base.show(io::IO, solver::SemialgebraicSetsHCSolver)
    print(io, "SemialgebraicSetsHCSolver(; ")
    join(
        io,
        [
            "$(k) = $(v isa Symbol ? Expr(:quote, v) : v)" for
            (k, v) in pairs(solver.options)
        ],
        ", ",
    )
    print(io, ")")
end

function SemialgebraicSets.solvealgebraicequations(
    V::SemialgebraicSets.AbstractAlgebraicSet,
    solver::SemialgebraicSetsHCSolver,
)
    F = System(V)
    m, n = size(F)
    # Check that we can have isolated solutions
    m ≥ n || return nothing
    # Only return real, non-singular solutions
    return real_solutions(solve(F; solver.options...); only_nonsingular = true)
end


#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 01-2018
#--------------------------------------------------------------------

function testdistribute(bnd1::Function, bnd2::Function, Nx::Int, Ny::Int, M::Int)::Float64
    dbase = distribute(bnd1, bnd2, (x,y)->0, Nx, Ny, M)
    sumL2 = 0.0
    for m in 1:M, n in 1:M
        sPatch = interpolatePatch(dbase[[m,n]], 2Nx, 2Ny).value
        fPatch = Float64[bnd1(x) + bnd2(y) for x in chebgrid(2Nx, M, m), y in chebgrid(2Ny, M, n)]
        sumL2 += (L2norm(fPatch, sPatch, chebweights(2Nx)/M, chebweights(2Ny)/M))^2
    end
    return sqrt(sumL2)
end

@test testdistribute(x->sin(pi*x), y->sin(pi*y), 20, 20, 2) < 1e-8


function day10()
    part = [0, 0]
    day10lines = filter(!isempty, strip.(readlines("AoCdata/AoC_2021_day10.txt")))
    rightop, vals, scores = [')', ']', '}', '>'], [3, 57, 1197, 25137], Int[]
    for line in day10lines
        corrupted = false
        rights = Char[]
        for c in line
            if c == '('
                push!(rights, ')')
                continue
            elseif c == '['
                push!(rights, ']')
                continue
            elseif c == '{'
                push!(rights, '}')
                continue
            elseif c == '<'
                push!(rights, '>')
                continue
            else
                if isempty(rights) || c != rights[end]
                    part[1] += vals[findfirst(==(c), rightop)]
                    corrupted = true
                else
                    pop!(rights)
                end
            end
        end
        !corrupted && push!(
            scores,
            reduce((x, y) -> 5x + findfirst(==(y), rightop), reverse(rights); init=0),
        )
    end
    part[2] = sort(scores)[(length(scores) + 1) ÷ 2]
    return part
end

part = day10()
println("Part 1:", part[1])
println("Part 2:", part[2])


using Distances, LinearAlgebra
using Test
using KernelFunctions

dims = [10,5]
X = rand(dims...)
k = SqExponentialKernel()
@testset "Kernel Matrix Approximations" begin
    @testset "Nystrom" begin
        for obsdim in [1, 2]
            @test kernelmatrix(k, X; obsdim=obsdim) ≈ kernelmatrix(nystrom(k, X, 1.0; obsdim=obsdim))
            @test kernelmatrix(k, X; obsdim=obsdim) ≈ kernelmatrix(nystrom(k, X, collect(1:dims[obsdim]); obsdim=obsdim))
        end
    end
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
;ls D:\\texlive\\2021\\bin\\win32\\lualatex.exe

# %%
using PGFPlotsX

p1 = range(0, 360; length = 100)' # parameter 1
p2 = range(0, 360; length = 70)   # parameter 2
w1 = @. sind(3*(p1 + 2 * p2)) + 1.25 # wave 1
w2 = @. 6 + w1 * cosd(p1)            # wave 2
x = vec(@. w2 * cosd(p2))
y = vec(@. w2 * sind(p2))
z = vec(@. w1 * sind(p1))
logo_colors = [(77,100,174), (57,151,79), (255,255,255), (146,89,163), (202,60,50)]

@pgf Axis(
    {
        axis_equal,
        axis_lines = "none"
    },
    Plot3(
        {
            surf,
            z_buffer = "sort",
            colormap = "{Julia}{$(join(map(c -> "rgb255=$c", logo_colors), ", "))}",
            "mesh/rows" = length(p1)
        },
        Table(x, y, z)))

# %%


@testset "6.5 Hyperbolic secant" begin
    include("6.5.1 (c+d x)^m (a+b sech)^n.jl")
    include("6.5.2 (e x)^m (a+b sech(c+d x^n))^p.jl")
    include("6.5.3 Hyperbolic secant functions.jl")
    include("6.5.7 (d hyper)^m (a+b (c sech)^n)^p.jl")
end

test/6 Hyperbolic functions/6.5 Hyperbolic secant/.jl


using GPFlux
using Zygote
using Calculus
using Test, Random

@testset "power function (scalar)" begin
	Random.seed!(4)
	xs = 10*rand(4)
	αs = [-0.5, 3.5]

	for x in xs, α in αs
		dfdx, dfdα = gradient(GPFlux.power, x, α)
		Finit_dfdx, Finit_dfdα = Calculus.gradient(z->GPFlux.power(z[1], z[2]), [x, α])
		@test dfdx ≈ Finit_dfdx
		@test dfdα ≈ Finit_dfdα
	end
end

@testset "power function (Array)" begin
	Random.seed!(4)
	x = rand(3, 3)
	w = rand(3)
	f(x, α) = w'*GPFlux.power(x, α)*w

	dfdx, dfdα = gradient(f, x, 2.5)
	Finit_dfdα = Calculus.gradient(z->f(x, first(z)), [2.5])
	@test dfdα ≈ Finit_dfdα[1]
end

@testset "positive constraint" begin
	Random.seed!(4)
	x = randn(3, 5)
	w = rand(3)
	v = rand(5)
	f(x) = w'*GPFlux.positive(x)*v
	g(y) = w'*GPFlux.positive(reshape(y, 3, 5))*v

	dfdx, = gradient(f, x)
	Finit_dfdx = reshape(Calculus.gradient(g, vec(x)), 3, 5)
	@test dfdx ≈ Finit_dfdx
end


"""
    assemble_flux_matrix(facets, nodes[, weights])

Assemble 3D flux matrix using P1 finite elements.

This function is based on the Matlab function `flux_matrixP1_3D.m` from the Matlab
nodal matrix assembly toolbox by Jan Valdman and Talal Rahman.

https://uk.mathworks.com/matlabcentral/fileexchange/27826-fast-fem-assembly-nodal-elements

Talal Rahman and Jan Valdman: Fast MATLAB assembly of FEM matrices in 2D and 3D: nodal
elements, Applied Mathematics and Computation 219, 7151–7158 (2013).
"""
function assemble_flux_matrix(facets, nodes, weights = 1)
    nnode = size(nodes, 1)
    areas = compute_areas(facets, nodes)

    x = kron(ones(1, 3), facets)
    y = kron(facets, ones(1, 3))
    if length(weights) == 1 || length(weights) == size(facets, 1)
        # P0 weights
        z = kron(weights .* areas, reshape((ones(3, 3) + I(3)) / 12, 1, 9))
    else
        # P1 weights
        M1 = [6 2 2; 2 2 1; 2 1 2] / 60
        M2 = M1[[3, 1, 2], [3, 1, 2]]
        M3 = M2[[3, 1, 2], [3, 1, 2]]

        z = (
            kron(areas .* weights[facets[:, 1]], reshape(M1, 1, 9)) +
            kron(areas .* weights[facets[:, 2]], reshape(M2, 1, 9)) +
            kron(areas .* weights[facets[:, 3]], reshape(M3, 1, 9))
        )
    end

    # Enforce symmetry
    sym(x) = (x + x') / 2

    sym(sparse(x[:], y[:], z[:], nnode, nnode))
end

"""
    compute_areas(facets, nodes)

Compute facet areas.
"""
function compute_areas(facets, nodes)
    nfacet = size(facets, 1)

    v1 = nodes[facets[:, 2], :] - nodes[facets[:, 1], :]
    v2 = nodes[facets[:, 3], :] - nodes[facets[:, 1], :]

    matrix_3D = zeros(2, 2, size(facets, 1))
    matrix_3D[1, 1, :] = sum(v1 .* v1, dims = 2)
    matrix_3D[1, 2, :] = sum(v1 .* v2, dims = 2)
    matrix_3D[2, 1, :] = sum(v1 .* v2, dims = 2)
    matrix_3D[2, 2, :] = sum(v2 .* v2, dims = 2)

    [√abs(det(matrix_3D[:, :, i])) / 2 for i = 1:nfacet]
end


#######################################################################################
# GPU-crutch solve method
# Makes the simplest possible method for GPU-compatibility
# Out of place only
#######################################################################################
struct GPUSimpleRK4 <: DiffEqBase.AbstractODEAlgorithm end
export GPUSimpleRK4

@muladd function DiffEqBase.solve(prob::ODEProblem,
                          alg::GPUSimpleRK4;
                          dt = error("dt is required for this algorithm"))
  @assert !isinplace(prob)
  u0 = prob.u0
  tspan = prob.tspan
  f = prob.f
  p = prob.p
  t = tspan[1]
  tf = prob.tspan[2]
  ts = tspan[1]:dt:tspan[2]
  us = MVector{length(ts),typeof(u0)}(undef)
  us[1] = u0
  u = u0
  half = convert(eltype(u0),1//2)
  sixth = convert(eltype(u0),1//6)

  for i in 2:length(ts)
      uprev = u; t = ts[i]
      k1 = f(u,p,t)
      tmp = uprev+dt*half*k1
      k2 = f(tmp, p, t+half*dt)
      tmp = uprev+dt*half*k2
      k3 = f(tmp, p, t+half*dt)
      tmp = uprev+dt*k3
      k4 = f(tmp, p, t+dt)
      u = uprev+dt*sixth*(k1+2k2+2k3+k4)
      us[i] = u
  end

  sol = DiffEqBase.build_solution(prob,alg,ts,SArray(us),
                                  k = nothing, destats = nothing,
                                  calculate_error = false)
  DiffEqBase.has_analytic(prob.f) && DiffEqBase.calculate_solution_errors!(sol;timeseries_errors=true,dense_errors=false)
  sol
end


function FHN(du,u,p,t) # FitzHugh-Nagumo oscillator
    du[1] = 30*u[1] - 10*u[1]^3 - 30*u[2]
    du[2] = 7.5*u[1]
end

function spiral(du,u,p,t) # Spiral
    true_A = [-0.1 -2.0 0; 2.0 -0.1 0; 0 0 -0.3]
    du .= true_A * (u .^ 3 + u)
end

function NDM(du,u,p,t) # Mutual Inhibition
    du[1] = 10 * (-u[1] + 1 ./ ( 1 + exp( 16 * ( u[2] - 0.5 ) ) ))
    du[2] = 10 * (-u[2] + 1 ./ ( 1 + exp( 16 * ( u[1] - 0.5 ) ) ))
end


#Benchmarking of RBC and FVGQ variants
using DifferenceEquations, BenchmarkTools
using DelimitedFiles, Distributions, Zygote, LinearAlgebra

# for benchmarking construction itself
function make_problem_1(A, B, C, u0, noise, observables, D; kwargs...)
    prob = LinearStateSpaceProblem(A, B, u0, (0, size(observables, 2)); C, observables_noise = D,
                                   noise, observables, kwargs...)
    return prob.A[1, 1] + prob.B[1, 1]
end
function make_problem_kalman(A, B, C, u0_prior, observables, D; kwargs...)
    prob = LinearStateSpaceProblem(A, B, u0_prior, (0, size(observables, 2)); C, u0_prior,
                                   observables_noise = D, noise = nothing, observables, kwargs...)
    return prob.A[1, 1] + prob.B[1, 1]
end

# for benchmarking likelihoods
function joint_likelihood_1(A, B, C, u0, noise, observables, D; kwargs...)
    prob = LinearStateSpaceProblem(A, B, u0, (0, size(observables, 2)); C, observables_noise = D,
                                   noise, observables, kwargs...)
    return solve(prob).logpdf
end
function kalman_likelihood(A, B, C, u0_prior, observables, D; kwargs...)
    prob = LinearStateSpaceProblem(A, B, u0_prior, (0, size(observables, 2)); C, u0_prior,
                                   observables_noise = D, noise = nothing, observables, kwargs...)
    return solve(prob).logpdf
end

function simulate_model_no_noise_1(A, B, C, u0, observables, D; kwargs...)
    prob = LinearStateSpaceProblem(A, B, u0, (0, size(observables, 2)); C, observables_noise = D,
                                   observables, kwargs...)
    sol = solve(prob)
    return sol.retcode
end

function simulate_model_no_observations_1(A, B, C, u0, T; kwargs...)
    prob = LinearStateSpaceProblem(A, B, u0, (0, T); C, kwargs...)
    sol = solve(prob)
    return sol.retcode
end

# Matrices from RBC
const A_rbc = [0.9568351489231076 6.209371005755285;
               3.0153731819288737e-18 0.20000000000000007]
const B_rbc = reshape([0.0; -0.01], 2, 1) # make sure B is a matrix
const C_rbc = [0.09579643002426148 0.6746869652592109; 1.0 0.0]
const D_rbc = MvNormal(Diagonal(abs2.([0.1, 0.1])))
const u0_rbc = zeros(2)
const u0_prior_rbc = MvNormal(diagm(ones(length(u0_rbc))))

const observables_rbc = readdlm(joinpath(pkgdir(DifferenceEquations),
                                         "test/data/RBC_observables.csv"), ',')' |> collect
const noise_rbc = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/RBC_noise.csv"), ',')' |>
                  collect
const T_rbc = size(observables_rbc, 2)
# Matrices from FVGQ
# Load FVGQ data for checks
const A_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_A.csv"), ',')
const B_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_B.csv"), ',')
const C_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_C.csv"), ',')
const D_FVGQ = MvNormal(Diagonal(abs2.(ones(6) * 1e-3)))

const observables_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations),
                                          "test/data/FVGQ20_observables.csv"), ',')' |> collect

const noise_FVGQ = readdlm(joinpath(pkgdir(DifferenceEquations), "test/data/FVGQ20_noise.csv"),
                           ',')' |> collect
const u0_FVGQ = zeros(size(A_FVGQ, 1))
const u0_prior_FVGQ = MvNormal(diagm(ones(length(u0_FVGQ))))
const T_FVGQ = size(observables_FVGQ, 2)
# executing gradients once to avoid compilation time in benchmarking

make_problem_1(A_rbc, B_rbc, C_rbc, u0_rbc, noise_rbc, observables_rbc, D_rbc)
gradient((args...) -> make_problem_1(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc, u0_rbc,
         noise_rbc)

make_problem_kalman(A_rbc, B_rbc, C_rbc, u0_prior_rbc, observables_rbc, D_rbc)
gradient((args...) -> make_problem_kalman(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
         u0_prior_rbc)

kalman_likelihood(A_rbc, B_rbc, C_rbc, u0_prior_rbc, observables_rbc, D_rbc)
gradient((args...) -> kalman_likelihood(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
         u0_prior_rbc)
joint_likelihood_1(A_rbc, B_rbc, C_rbc, u0_rbc, noise_rbc, observables_rbc, D_rbc)
gradient((args...) -> joint_likelihood_1(args..., observables_rbc, D_rbc), A_rbc, B_rbc, C_rbc,
         u0_rbc, noise_rbc)

kalman_likelihood(A_FVGQ, B_FVGQ, C_FVGQ, u0_prior_FVGQ, observables_FVGQ, D_FVGQ)
gradient((args...) -> kalman_likelihood(args..., observables_FVGQ, D_FVGQ), A_FVGQ, B_FVGQ, C_FVGQ,
         u0_prior_FVGQ)
joint_likelihood_1(A_FVGQ, B_FVGQ, C_FVGQ, u0_FVGQ, noise_FVGQ, observables_FVGQ, D_FVGQ)
gradient((args...) -> joint_likelihood_1(args..., observables_FVGQ, D_FVGQ), A_FVGQ, B_FVGQ, C_FVGQ,
         u0_FVGQ, noise_FVGQ)

####### Becnmarks

const LINEAR = BenchmarkGroup()

const LINEAR["rbc"] = BenchmarkGroup()

const LINEAR["rbc"]["make_problem_1"] = @benchmarkable make_problem_1($A_rbc, $B_rbc, $C_rbc,
                                                                      $u0_rbc, $noise_rbc,
                                                                      $observables_rbc, $D_rbc)
const LINEAR["rbc"]["make_problem_1_gradient"] = @benchmarkable gradient((args...) -> make_problem_1(args...,
                                                                                $observables_rbc,
                                                                                $D_rbc),
                                                                         $A_rbc, $B_rbc, $C_rbc,
                                                                         $u0_rbc, $noise_rbc)

const LINEAR["rbc"]["simulate_model_no_noise_1"] = @benchmarkable simulate_model_no_noise_1($A_rbc,
                                                                                $B_rbc,
                                                                                $C_rbc,
                                                                                $u0_rbc,
                                                                                $observables_rbc,
                                                                                $D_rbc)
const LINEAR["rbc"]["simulate_model_no_observations_1"] = @benchmarkable simulate_model_no_observations_1($A_rbc,
                                                                                $B_rbc,
                                                                                $C_rbc,
                                                                                $u0_rbc,
                                                                                T_rbc)
const LINEAR["rbc"]["joint_1"] = @benchmarkable joint_likelihood_1($A_rbc, $B_rbc, $C_rbc, $u0_rbc,
                                                                   $noise_rbc, $observables_rbc,
                                                                   $D_rbc)
const LINEAR["rbc"]["joint_1_gradient"] = @benchmarkable gradient((args...) -> joint_likelihood_1(args...,
                                                                                $observables_rbc,
                                                                                $D_rbc),
                                                                  $A_rbc, $B_rbc, $C_rbc, $u0_rbc,
                                                                  $noise_rbc)
const LINEAR["rbc"]["make_problem_kalman"] = @benchmarkable make_problem_kalman($A_rbc, $B_rbc,
                                                                                $C_rbc,
                                                                                $u0_prior_rbc,
                                                                                $observables_rbc,
                                                                                $D_rbc)
const LINEAR["rbc"]["make_problem_kalman_gradient"] = @benchmarkable gradient((args...) -> make_problem_kalman(args...,
                                                                                $observables_rbc,
                                                                                $D_rbc),
                                                                              $A_rbc, $B_rbc,
                                                                              $C_rbc, $u0_prior_rbc)
const LINEAR["rbc"]["kalman"] = @benchmarkable kalman_likelihood($A_rbc, $B_rbc, $C_rbc,
                                                                 $u0_prior_rbc, $observables_rbc,
                                                                 $D_rbc)
const LINEAR["rbc"]["kalman_gradient"] = @benchmarkable gradient((args...) -> kalman_likelihood(args...,
                                                                                $observables_rbc,
                                                                                $D_rbc),
                                                                 $A_rbc, $B_rbc, $C_rbc,
                                                                 $u0_prior_rbc)

# FVGQ
const LINEAR["FVGQ"] = BenchmarkGroup()
const LINEAR["FVGQ"]["make_problem_1"] = @benchmarkable make_problem_1($A_FVGQ, $B_FVGQ, $C_FVGQ,
                                                                       $u0_FVGQ, $noise_FVGQ,
                                                                       $observables_FVGQ, $D_FVGQ)
const LINEAR["FVGQ"]["make_problem_1_gradient"] = @benchmarkable gradient((args...) -> make_problem_1(args...,
                                                                                $observables_FVGQ,
                                                                                $D_FVGQ),
                                                                          $A_FVGQ, $B_FVGQ, $C_FVGQ,
                                                                          $u0_FVGQ, $noise_FVGQ)
const LINEAR["FVGQ"]["simulate_model_no_noise_1"] = @benchmarkable simulate_model_no_noise_1($A_FVGQ,
                                                                                $B_FVGQ,
                                                                                $C_FVGQ,
                                                                                $u0_FVGQ,
                                                                                $observables_FVGQ,
                                                                                $D_FVGQ)
const LINEAR["FVGQ"]["simulate_model_no_observations_1"] = @benchmarkable simulate_model_no_observations_1($A_FVGQ,
                                                                                $B_FVGQ,
                                                                                $C_FVGQ,
                                                                                $u0_FVGQ,
                                                                                $T_FVGQ)
const LINEAR["FVGQ"]["joint_1"] = @benchmarkable joint_likelihood_1($A_FVGQ, $B_FVGQ, $C_FVGQ,
                                                                    $u0_FVGQ, $noise_FVGQ,
                                                                    $observables_FVGQ, $D_FVGQ)
const LINEAR["FVGQ"]["joint_1_gradient"] = @benchmarkable gradient((args...) -> joint_likelihood_1(args...,
                                                                                $observables_FVGQ,
                                                                                $D_FVGQ),
                                                                   $A_FVGQ, $B_FVGQ, $C_FVGQ,
                                                                   $u0_FVGQ, $noise_FVGQ)
const LINEAR["FVGQ"]["make_problem_kalman"] = @benchmarkable make_problem_kalman($A_FVGQ, $B_FVGQ,
                                                                                $C_FVGQ,
                                                                                $u0_prior_FVGQ,
                                                                                $observables_FVGQ,
                                                                                $D_FVGQ)
const LINEAR["FVGQ"]["make_problem_kalman_gradient"] = @benchmarkable gradient((args...) -> make_problem_kalman(args...,
                                                                                $observables_FVGQ,
                                                                                $D_FVGQ),
                                                                               $A_FVGQ, $B_FVGQ,
                                                                               $C_FVGQ,
                                                                               $u0_prior_FVGQ)
const LINEAR["FVGQ"]["kalman"] = @benchmarkable kalman_likelihood($A_FVGQ, $B_FVGQ, $C_FVGQ,
                                                                  $u0_prior_FVGQ, $observables_FVGQ,
                                                                  $D_FVGQ)
const LINEAR["FVGQ"]["kalman_gradient"] = @benchmarkable gradient((args...) -> kalman_likelihood(args...,
                                                                                $observables_FVGQ,
                                                                                $D_FVGQ),
                                                                  $A_FVGQ, $B_FVGQ, $C_FVGQ,
                                                                  $u0_prior_FVGQ)

# return for the test suite
LINEAR