@testset "Lennard-Jones potential" begin
    @testset "Two interacting particles" begin

        m1 = 1.0
        m2 = 1.0
        r1 = 1.3

        t1 = 0.0
        t2 = 1.0
        τ = (t2 - t1) / 1000

        p1 = MassBody(SVector(-r1 / 2, 0.0, 0.0), SVector(0.0, 0.0, 0.0), m1)
        p2 = MassBody(SVector(r1 / 2, 0.0, 0.0), SVector(0.0, 0.0, 0.0), m2)

        σ = 1.0
        ϵ = 1.0
        parameters = LennardJonesParameters(ϵ, σ, Inf)
        system = PotentialNBodySystem([p1, p2], Dict(:lennard_jones => parameters))
        simulation = NBodySimulation(system, (t1, t2))
        sim_result = run_simulation(simulation, VelocityVerlet(), dt=τ)

        r2 = get_position(sim_result, t2, 2) - get_position(sim_result, t2, 1)
        v_expected = sqrt(4ϵ / m1 * ( ((σ / norm(r1))^12 - (σ / norm(r2))^12) - ((σ / norm(r1))^6 - (σ / norm(r2))^6 ) ))
        v_actual = norm(get_velocity(sim_result, t2, 2))

        ε = 0.001 * v_expected
        @test v_expected ≈ v_actual atol = ε

        io = IOBuffer()

        @test sprint(io -> show(io, system)) ==
    "Potentials: \nLennard-Jones:\n\tϵ:1.0\n\tσ:1.0\n\tR:Inf\n"

        @test sprint(io -> show(io, simulation)) ==
    "Timespan: (0.0, 1.0)\nBoundary conditions: InfiniteBox{Float64}([-Inf, Inf, -Inf, Inf, -Inf, Inf])\nPotentials: \nLennard-Jones:\n\tϵ:1.0\n\tσ:1.0\n\tR:Inf\n"

        @test sprint(io -> show(io, sim_result)) ==
    "N: 2\nTimespan: (0.0, 1.0)\nBoundary conditions: InfiniteBox{Float64}([-Inf, Inf, -Inf, Inf, -Inf, Inf])\nPotentials: \nLennard-Jones:\n\tϵ:1.0\n\tσ:1.0\n\tR:Inf\nTime steps: 1000\nt: 0.0, 1.0\n"
    end

    @testset "Three particles of liquid argon and their \"temperature\"" begin
        T = 120.0 # °K
        kb = 8.3144598e-3 # kJ/(K*mol)
        ϵ = T * kb
        σ = 0.34 # nm
        ρ = 1374/1.6747# Da/nm^3
        m = 39.95# Da
        L = 5σ # 10.229σ
        N = 3 # floor(Int, ρ * L^3 / m)
        R = 2.25σ
        v_dev = sqrt(3 * kb * T / m)
        r1 = SVector(L / 3, L / 3, 2 * L / 3)
        r2 = SVector(L / 3, 2 * L / 3, L / 3)
        r3 = SVector(2 * L / 3, L / 3, L / 3)
        v1 = SVector(0, 0, -v_dev)
        v2 = SVector(0, -v_dev, 0)
        v3 = SVector(-v_dev, 0, 0)
        p1 = MassBody(r1, v1, m)
        p2 = MassBody(r2, v2, m)
        p3 = MassBody(r3, v3, m)

        τ = 0.5e-3 # ps or 1e-12 s
        t1 = 0.0
        t2 = 100τ

        parameters = LennardJonesParameters(ϵ, σ, R)
        lj_system = PotentialNBodySystem([p1, p2, p3], Dict(:lennard_jones => parameters));
        simulation = NBodySimulation(lj_system, (t1, t2), PeriodicBoundaryConditions(L), kb);
        result = run_simulation(simulation, VelocityVerlet(), dt=τ)

        T1 = temperature(result, t1)
        ε = 1e-6
        @test T1 ≈ 120.0 atol = ε

        e_kin_1 = m * (dot(v1, v1) + dot(v2, v2) + dot(v3, v3)) / 2
        @test e_kin_1 == kinetic_energy(result, t1)

        e_tot_1 = total_energy(result, t1)
        ε = 0.1 * e_tot_1
        e_tot_2 = total_energy(result, t2)
        @test e_tot_1 ≈ e_tot_2 atol = ε

        simulation = NBodySimulation(lj_system, (t1, t2), CubicPeriodicBoundaryConditions(L));
        result = run_simulation(simulation, VelocityVerlet(), dt=τ)
        e_tot_1 = total_energy(result, t1)
        ε = 0.1 * e_tot_1
        e_tot_2 = total_energy(result, t2)
        @test e_tot_1 ≈ e_tot_2 atol = ε

        io = IOBuffer()
        pdb_data = sprint(io -> NBodySimulator.write_pdb_data(io, result))
        splitted_data = split(pdb_data, '\n')

        hetatm_count = 0
        timestep_count = 0

        for s in splitted_data
            if length(s) >= 10 && s[1:10] == "REMARK 250"
                timestep_count += 1
            elseif length(s) >= 6 && s[1:6] == "HETATM"
                hetatm_count += 1
            end
        end

        molecule_number = hetatm_count / timestep_count

        @test length(result.solution.t) == timestep_count
        @test molecule_number == length(lj_system.bodies)

    end

    @testset "Testing RDF and MSD calculation" begin
        T = 120.0 # °K
        kb = 8.3144598e-3 # kJ/(K*mol)
        ϵ = T * kb
        σ = 0.34 # nm
        ρ = 1374/1.6747# Da/nm^3
        m = 39.95# Da
        L = 5σ # 10.229σ
        N = 125 # floor(Int, ρ * L^3 / m)
        R = 0.5L
        τ = 0.5e-3 # ps or 1e-12 s
        t1 = 0.0
        t2 = 400τ
        v_dev = sqrt(kb * T / m)
        bodies = generate_bodies_in_cell_nodes(N, m, v_dev, L)

        parameters = LennardJonesParameters(ϵ, σ, R)
        lj_system = PotentialNBodySystem(bodies, Dict(:lennard_jones => parameters));
        simulation = NBodySimulation(lj_system, (t1, t2), CubicPeriodicBoundaryConditions(L), kb);
        result = run_simulation(simulation, VelocityVerlet(), dt=τ)

        (ts, mean_square_displacement) = msd(result)
        @test mean_square_displacement[1] < mean_square_displacement[end]

        (rs, grs) = rdf(result)
        (val, ind) = findmax(grs)
        @test (rs[ind] / σ) ≈ 1.0 atol = 1.0
    end

    @testset "Constructing electorstatic potential parameters entity" begin
        default_potential = LennardJonesParameters()
        @test 1.0 == default_potential.ϵ
        @test 1.0 == default_potential.σ
        @test 2.5 == default_potential.R
        @test 1.0 == default_potential.σ2
        @test 6.25 == default_potential.R2

        io = IOBuffer()
        potential1 = LennardJonesParameters()
        potential2 = LennardJonesParameters(2, 5, 10)
        @test sprint(io -> show(io, potential1)) == "Lennard-Jones:\n\tϵ:1.0\n\tσ:1.0\n\tR:2.5\n"
        @test sprint(io -> show(io, potential2)) == "Lennard-Jones:\n\tϵ:2\n\tσ:5\n\tR:10\n"
    end

    @testset "Periodic boundar conditions" begin
        pbc = PeriodicBoundaryConditions(15e-6)
        boundary = [0.0, 15e-6, 0.0, 15e-6, 0.0, 15e-6]
        ind = 1
        for b in pbc
            b = boundary[ind]
            ind += 1
        end
        for i = 1:6
            @test boundary[i] == pbc[i]
        end
    end
end


"""
    Nilplexification{T <: Number} <: Plexification{T}

A nilplexification as a pair.
"""
struct Nilplexification{T <: Number} <: Plexification{T}
    l::T
    r::T
end

function (==)(x::Nilplexification, y::Nilplexification)
    x.l == y.l && x.r == y.r
end

function real(z::Nilplexification)
    real(z.l)
end

function zero(::Type{Nilplexification{T}}) where T <: Number
    Nilplexification{T}(zero(T), zero(T))
end

function zero(::Nilplexification{T}) where T <: Number
    zero(Nilplexification{T})
end

function one(::Type{Nilplexification{T}}) where T <: Number
    Nilplexification{T}(one(T), zero(T))
end

function one(::Nilplexification{T}) where T <: Number
    one(Nilplexification{T})
end

function (-)(z::Nilplexification)
    Nilplexification(-(z.l), -(z.r))
end

function (+)(x::Nilplexification{T}, y::Nilplexification{T}) where T <: Number
    Nilplexification{T}(x.l + y.l, x.r + y.r)
end

function (+)(a::T, z::Nilplexification{T}) where T <: Number
    Nilplexification{T}(a + z.l, z.r)
end

function (+)(z::Nilplexification{T}, a::T) where T <: Number
    Nilplexification{T}(z.l + a, z.r)
end

function (+)(a::Real, z::Nilplexification)
    Nilplexification(a + z.l, z.r)
end

function (+)(z::Nilplexification, a::Real)
    Nilplexification(z.l + a, z.r)
end

function (-)(x::Nilplexification{T}, y::Nilplexification{T}) where T <: Number
    Nilplexification{T}(x.l - y.l, x.r - y.r)
end

function (-)(a::T, z::Nilplexification{T}) where T <: Number
    Nilplexification{T}(a - z.l, -(z.r))
end

function (-)(z::Nilplexification{T}, a::T) where T <: Number
    Nilplexification{T}(z.l - a, z.r)
end

function (-)(a::Real, z::Nilplexification)
    Nilplexification(a - z.l, -(z.r))
end

function (-)(z::Nilplexification, a::Real)
    Nilplexification(z.l - a, z.r)
end

function (*)(x::Nilplexification{T}, y::Nilplexification{T}) where T <: Number
    Nilplexification{T}(
        (x.l * y.l),
        (x.l * y.r) + (x.r * y.l)
    )
end

function (*)(a::T, z::Nilplexification{T}) where T <: Number
    Nilplexification{T}(a * z.l, a * z.r)
end

function (*)(z::Nilplexification{T}, a::T) where T <: Number
    Nilplexification{T}(z.l * a , z.r * a)
end

function (*)(a::Real, z::Nilplexification)
    Nilplexification(a * z.l, a * z.r)
end

function (*)(z::Nilplexification, a::Real)
    Nilplexification(z.l * a, z.r * a)
end

function (/)(z::Nilplexification, a::Real)
    if iszero(a)
        error(ZeroDenominator)
    end

    Nilplexification(z.l / a, z.r / a)
end

function (\)(a::Real, z::Nilplexification)
    if iszero(a)
        error(ZeroDenominator)
    end

    Nilplexification(a \ z.l, a \ z.r)
end


#using Plots
using FFTW

include("Domain.jl")
include("Field.jl")
include("Distributor.jl")

dom = RGDomain(512,512)

hdom = get_codomain(dom)
pdom = PowerDomain(hdom)

f = randn(hdom)
fcorr = similar(f, hdom)
kernel = field(pdom, 1 ./ (1 .+ pdom._kvec).^1.5)

fcorr = PowerDistribute(f, kernel);

corr = real(fft(fcorr))
#heatmap(real(fft(fcorr.val)))


using Combinatorics
using StatsBase

using Base.Iterators

function njet(::Val{N}, ::typeof(sin), x₀) where {N}
    (s, c) = sincos(x₀)
    Jet(x₀, s, tuple(take(cycle((c, -s, -c, s)), N)...))
end

function njet(::Val{N}, ::typeof(cos), x₀) where {N}
    (s, c) = sincos(x₀)
    Jet(x₀, s, tuple(take(cycle((-s, -c, s, c)), N)...))
end

function njet(::Val{N}, ::typeof(exp), x₀) where {N}
    exped = exp(x₀)
    Jet(x₀, exped, tuple(take(repeated(exped), N)...))
end

jeval(j, x) = j(x)
for f in (sin, cos, exp)
    function (∂☆ₙ::∂☆{N})(fb::ZeroBundle{N, typeof(f)}, x::TaylorBundle{N}) where {N}
        njet(Val{N}(), primal(fb), primal(x))(x)
    end
    function (∂⃖ₙ::∂⃖{N})(∂☆ₘ::∂☆{M}, fb::ZeroBundle{M, typeof(f)}, x::TaylorBundle{M}) where {N, M}
        ∂⃖ₙ(jeval, njet(Val{N+M}(), primal(fb), primal(x)), x)
    end
end

function (::Diffractor.∂☆new{N})(B::ATB{N, Type{T}}, args::ATB{N}...) where {N, T<:SArray}
    error("Should have intercepted the constructor")
end


immutable VonMisesSampler <: Sampleable{Univariate,Continuous}
    μ::Float64
    κ::Float64
    r::Float64
end

# algorithm from
#     DJ Best & NI Fisher (1979). Efficient Simulation of the von Mises
#     Distribution. Journal of the Royal Statistical Society. Series C
#     (Applied Statistics), 28(2), 152-157.
function rand(s::VonMisesSampler)
    f = 0.0
    while true
        t, u = 0.0, 0.0
        while true
            const v, w = rand() - 0.5, rand() - 0.5
            const d, e = v ^ 2, w ^ 2
            if d + e <= 0.25
                t = d / e
                u = 4 * (d + e)
                break
            end
        end
        const z = (1.0 - t) / (1.0 + t)
        f = (1.0 + s.r * z) / (s.r + z)
        const c = s.κ * (s.r - f)
        if c * (2.0 - c) > u || log(c / u) + 1 >= c
            break
        end
    end
    mod(s.μ + (rand() > 0.5 ? acos(f) : -acos(f)), twoπ)
end

immutable VonMisesNormalApproxSampler <: Sampleable{Univariate,Continuous}
    μ::Float64
    σ::Float64
end

rand(s::VonMisesNormalApproxSampler) = mod(s.μ + s.σ * randn(), twoπ)

# normal approximation for large concentrations
VonMisesSampler(μ::Float64, κ::Float64) =
    κ > 700.0 ? VonMisesNormalApproxSampler(μ, sqrt(1.0 / κ)) :
                  begin
                      τ = 1.0 + sqrt(1.0 + 4 * κ ^ 2)
                      ρ = (τ - sqrt(2.0 * τ)) / (2.0 * κ)
                      VonMisesSampler(μ, κ, (1.0 + ρ ^ 2) / (2.0 * ρ))
                  end


"""
    PS.jl

This is a Julia port of a MATLAB implementation of batch and incremental
Partition Separation (PS) Cluster Validity Index.

Authors:
MATLAB implementation: Leonardo Enzo Brito da Silva
Julia port: Sasha Petrenko <sap625@mst.edu>

REFERENCES
[1] Miin-Shen Yang and Kuo-Lung Wu, "A new validity index for fuzzy clustering,"
10th IEEE International Conference on Fuzzy Systems. (Cat. No.01CH37297), Melbourne,
Victoria, Australia, 2001, pp. 89-92, vol.1.
[2] E. Lughofer, "Extensions of vector quantization for incremental clustering," Pattern
Recognit., vol. 41, no. 3, pp. 995–1011, 2008.
"""

"""
    PS

The stateful information of the Partition Separation (PS) Cluster Validity Index.

# References
1. Miin-Shen Yang and Kuo-Lung Wu, "A new validity index for fuzzy clustering," 10th IEEE International Conference on Fuzzy Systems. (Cat. No.01CH37297), Melbourne, Victoria, Australia, 2001, pp. 89-92, vol.1.
2. E. Lughofer, "Extensions of vector quantization for incremental clustering," Pattern Recognit., vol. 41, no. 3, pp. 995–1011, 2008.
"""
mutable struct PS <: CVI
    label_map::LabelMap
    dim::Integer
    n_samples::Integer
    n::IntegerVector        # dim
    v::RealMatrix           # dim x n_clusters
    D::RealMatrix           # n_clusters x n_clusters
    v_bar::RealVector       # dim
    beta_t::Float
    PS_i::RealVector        # n_clusters
    n_clusters::Integer
    criterion_value::Float
end # PS <: CVI

"""
    PS()

Default constructor for the Partition Separation (PS) Cluster Validity Index.
"""
function PS()
    PS(
        LabelMap(),                     # label_map
        0,                              # dim
        0,                              # n_samples
        Array{Integer, 1}(undef, 0),    # n
        Array{Float, 2}(undef, 0, 0),   # v
        Array{Float, 2}(undef, 0, 0),   # D
        Array{Float, 1}(undef, 0),      # v_bar
        0.0,                            # beta_t
        Array{Float, 1}(undef, 0),      # PS_i
        0,                              # n_clusters
        0.0                             # criterion_value
    )
end # PS()

"""
    setup!(cvi::PS, sample::Vector{T}) where {T<:RealFP}
"""
function setup!(cvi::PS, sample::Vector{T}) where {T<:RealFP}
    # Get the feature dimension
    cvi.dim = length(sample)
    # Initialize the 2-D arrays with the correct feature dimension
    cvi.v = Array{T, 2}(undef, cvi.dim, 0)
end # setup!(cvi::PS, sample::Vector{T}) where {T<:RealFP}

function param_inc!(cvi::PS, sample::RealVector, label::Integer)
    # Get the internal label
    i_label = get_internal_label!(cvi.label_map, label)

    n_samples_new = cvi.n_samples + 1
    if isempty(cvi.v)
        setup!(cvi, sample)
    end

    if i_label > cvi.n_clusters
        n_new = 1
        v_new = sample
        if cvi.n_clusters == 0
            D_new = zeros(1, 1)
        else
            D_new = zeros(cvi.n_clusters + 1, cvi.n_clusters + 1)
            D_new[1:cvi.n_clusters, 1:cvi.n_clusters] = cvi.D
            d_column_new = zeros(cvi.n_clusters + 1)
            for jx = 1:cvi.n_clusters
                d_column_new[jx] = sum((v_new - cvi.v[:, jx]).^2)
            end
            D_new[:, i_label] = d_column_new
            D_new[i_label, :] = transpose(d_column_new)
        end
        # Update 1-D parameters with a push
        cvi.n_clusters += 1
        push!(cvi.n, n_new)
        # Update 2-D parameters with appending and reassignment
        cvi.v = [cvi.v v_new]
        cvi.D = D_new
    else
        n_new = cvi.n[i_label] + 1
        v_new = (1 - 1/n_new) .* cvi.v[:, i_label] + (1/n_new) .* sample
        d_column_new = zeros(cvi.n_clusters)
        for jx = 1:cvi.n_clusters
            if jx == i_label
                continue
            end
            d_column_new[jx] = sum((v_new - cvi.v[:, jx]).^2)
        end
        # Update parameters
        cvi.n[i_label] = n_new
        cvi.v[:, i_label] = v_new
        cvi.D[:, i_label] = d_column_new
        cvi.D[i_label, :] = transpose(d_column_new)
    end
    cvi.n_samples = n_samples_new
end # param_inc!(cvi::PS, sample::RealVector, label::Integer)

function param_batch!(cvi::PS, data::RealMatrix, labels::IntegerVector)
    cvi.dim, cvi.n_samples = size(data)
    # Take the average across all samples, but cast to 1-D vector
    u = unique(labels)
    cvi.n_clusters = length(u)
    cvi.n = zeros(Integer, cvi.n_clusters)
    cvi.v = zeros(cvi.dim, cvi.n_clusters)
    cvi.D = zeros(cvi.n_clusters, cvi.n_clusters)
    for ix = 1:cvi.n_clusters
        subset = data[:, findall(x->x==u[ix], labels)]
        cvi.n[ix] = size(subset, 2)
        cvi.v[1:cvi.dim, ix] = mean(subset, dims=2)
    end
    for ix = 1 : (cvi.n_clusters - 1)
        for jx = ix + 1 : cvi.n_clusters
            cvi.D[jx, ix] = sum((cvi.v[:, ix] - cvi.v[:, jx]).^2)
        end
    end
    cvi.D = cvi.D + transpose(cvi.D)
end # param_batch!(cvi::PS, data::RealMatrix, labels::IntegerVector)

function evaluate!(cvi::PS)
    if cvi.n_clusters > 1
        cvi.v_bar = vec(mean(cvi.v, dims=2))
        cvi.beta_t = 0.0
        cvi.PS_i = zeros(cvi.n_clusters)
        for ix = 1:cvi.n_clusters
            delta_v = cvi.v[:, ix] - cvi.v_bar
            cvi.beta_t = cvi.beta_t + transpose(delta_v) * delta_v
        end
        cvi.beta_t /= cvi.n_clusters
        n_max = maximum(cvi.n)
        for ix = 1:cvi.n_clusters
            d = cvi.D[:, ix]
            deleteat!(d, ix)
            cvi.PS_i[ix] = (cvi.n[ix] / n_max) - exp(-minimum(d) / cvi.beta_t)
        end
        cvi.criterion_value = sum(cvi.PS_i)
    end
end # evaluate(cvi::PS)


export FidMtx,
       DispStable!,
       InitFidMtx,
      #  InitFidMtx_back,
       vp2lambda,
       lambda2vp,
       PMLCoef,
       InitCpml,
       modExpand,
       modSmooth,
       Mvz,
       Mvx,
       Mpz,
       Mpx,
      #  Mvz_back,
      #  Mvx_back,
      #  Mpz_back,
      #  Mpx_back,
       SnapShot,
       InitSnapShot,
       InitSnapShot,
       copy,
       -,
       +,
       write,
       read,
       writeSnapShot,
       readSnapShot,
       numberOfSnapShots,
       reverseSnapShotsOrder,
       Source,
       Ricker,
       InitSources,
       AddSources!,
       AddSource!,
       SrcRange,
       src2spt,
       src2SnapShots,
       SrcSpt,
       InitSrcSpt,
       src2srcSpts,
       readSrcSpt,
       InfoSrcSpt,
       AddSrcSpt2Spt,
       saveSpt2SrcSpt,
       reverseSrcSptOrder,
       OneStepForward!,
       MultiStepForward!,
       MultiStepForward,
       Shot3,
       InitShot3,
       Shot,
       InitShot,
       spt2shot!,
       AddShot2SnapShot!,
       InfoWfd,
       InfoFullWfd,
       wfd2shot,
       fullWfd2Shot,
       writeShot!,
       readShot,
       OneStepAdjoint!,
       MultiStepAdjoint

include("PML.jl")
include("Mvz.jl")
include("Mvx.jl")
include("Mpz.jl")
include("Mpx.jl")
include("fidMtx.jl")
include("SnapShot.jl")
include("Source.jl")
include("Forward.jl")
include("Shot.jl")
include("Adjoint.jl")


using StatsBase: countmap

rand(DCHSBM_sampler(x->1, [1,1,1,2,2], fill(1, 5), 3), 10)
rand(DCHSBM_sampler(x->1/FBD.sorted_unique_count(x)^2, repeat(1:3, inner=10000000), fill(1, 30000000), 2))

Z = [1,1,1,2,2]
θ = 1:5
kmax = 3
m = 1_000_000
intensity = x -> hash(x)/typemax(UInt)+1 # determinitic random ∈ [1,2]
sampler = DCHSBM_sampler(intensity, Z, θ, kmax)

graph = rand(sampler, m)

nodes = []
for e in graph
    append!(nodes, e)
end
degrees = countmap(nodes)

@testset "shape" begin
    @test length(graph) == m
    @test all(length.(graph) .== kmax)
    @test Set(keys(degrees)) == Set(axes(Z, 1))
end

tolerance = 6
max_rtol = ceil(tolerance/sqrt(minimum(values(degrees)))*10_000)/10_000
@testset "degree distribution (rtol ≤ $max_rtol)" begin
    for group = [1:3, 4:5]
        for i in group
            for j in group
                rtol = tolerance/sqrt(min(degrees[i], degrees[j]))
                @test rtol ≤ max_rtol
                @test degrees[i]/degrees[j] ≈ θ[i]/θ[j] rtol=rtol
            end
        end
    end
end


# Some sample data

"""
    ol() -> C, rho

Return the normalised elastic constants `C` and density `rho` for olivine,
handy for testing purposes
"""
ol() = EC{DEFAULT_FLOAT}((320.5,  68.1,  71.6,   0.0,   0.0,   0.0,
                           68.1, 196.5,  76.8,   0.0,   0.0,   0.0,
                           71.6,  76.8, 233.5,   0.0,   0.0,   0.0,
                            0.0,   0.0,   0.0,  64.0,   0.0,   0.0,
                            0.0,   0.0,   0.0,   0.0,  77.0,   0.0,
                            0.0,   0.0,   0.0,   0.0,   0.0,  78.7).*1.e9./3355.0), 3355.0


using StatsBase
using MultipleTesting
using IndependentHypothesisWeighting
using Test

pvs = rand(100)
for method in [Bonferroni(), BenjaminiHochberg()]
	unit_priority_ws = PriorityWeights(ones(100))
	unit_ws = UnitWeights{Float64}(100)
	adjust_unwt = adjust(pvs, method)
	adjust_wt = adjust(pvs, unit_priority_ws, method)
	adjust_unit_wt = adjust(pvs, unit_ws, method)

	@test adjust_unwt == adjust_wt
	@test adjust_unwt == adjust_unit_wt

	half_ws = PriorityWeights([2 .* ones(50); zeros(50)])
	adjust_half_wt = adjust(pvs, half_ws, method)
	adjust_half_unwt = adjust(pvs[1:50], method)
	@test  adjust_half_wt[1:50] == adjust_half_unwt
end


"""
    KSDist(n)

Distribution of the (two-sided) Kolmogorov-Smirnoff statistic

```math
D_n = \\sup_x | \\hat{F}_n(x) -F(x)| \\sqrt(n)
```

``D_n`` converges a.s. to the Kolmogorov distribution.
"""
struct KSDist <: ContinuousUnivariateDistribution
    n::Int
end

@distr_support KSDist 1 / (2 * d.n) 1.0

#### Evaluation

# TODO: implement Simard and L'Ecuyer (2011) meta-algorithm
# requires Pomeranz and Pelz-Good algorithms
function cdf(d::KSDist,x::Float64)
    n = d.n
    b = x*n
    # known exact values
    if b <= 1/2
        return 0.0
    elseif b <= 1
        # accuracy could be improved
        return exp(lfact(n)+n*(log(2*b-1)-log(n)))
    elseif x >= 1
        return 1.0
    elseif b >= n-1
        return 1 - 2*(1-x)^n
    end

    a = b*x
    if a >= 18
        return 1.0
    elseif n <= 10_000
        if a <= 4
            return cdf_durbin(d,x)
        else
            return 1 - ccdf_miller(d,x)
        end
    else
        return cdf(Kolmogorov(),sqrt(a))
    end
end

function ccdf(d::KSDist,x::Float64)
    n = d.n
    b = x*n
    # Ruben and Gambino (1982) known exact values
    if b <= 0.5
        return 1.0
    elseif b <= 1
        return 1-exp(lfact(n)+n*(log(2*b-1)-log(n)))
    elseif x >= 1
        return 0.0
    elseif b >= n-1
        return 2*(1-x)^n
    end

    a = b*x
    if a >= 370
        return 0.0
    elseif a >= 4 || (n > 140 && a >= 2.2)
        return ccdf_miller(d,x)
    else
        return 1-cdf(d,x)
    end
end

# Durbin matrix CDF method, based on Marsaglia, Tsang and Wang (2003)
# modified to avoid need for exponent tracking
function cdf_durbin(d::KSDist,x::Float64)
    n = d.n
    k, ch, h = ceil_rems_mult(n,x)

    m = 2*k-1
    H = Matrix{Float64}(m, m)
    for i = 1:m, j = 1:m
        H[i,j] = i-j+1 >= 0 ? 1 : 0
    end
    r = 1.0
    for i = 1:m
        # (1-h^i) = (1-h)(1+h+...+h^(i-1))
        H[i,1] = H[m,m-i+1] = ch*r
        r += h^i
    end
    H[m,1] += h <= 0.5 ? -h^m : -h^m+(h-ch)
    for i = 1:m, j = 1:m
        for g = 1:max(i-j+1,0)
            H[i,j] /= g
        end
        # we can avoid keeping track of the exponent by dividing by e
        # (from Stirling's approximation)
        H[i,j] /= ℯ
    end
    Q = H^n
    s = Q[k,k]
    s*stirling(n)
end

# Miller (1956) approximation
function ccdf_miller(d::KSDist, x::Real)
    2*ccdf(KSOneSided(d.n),x)
end

## these functions are used in durbin and pomeranz algorithms
# calculate exact remainders the easy way
function floor_rems_mult(n,x)
    t = big(x)*big(n)
    fl = floor(t)
    lrem = t - fl
    urem = (fl+one(fl)) - t
    return convert(typeof(n),fl), convert(typeof(x),lrem), convert(typeof(x),urem)
end
function ceil_rems_mult(n,x)
    t = big(x)*big(n)
    cl = ceil(t)
    lrem = t - (cl - one(cl))
    urem = cl - t
    return convert(typeof(n),cl), convert(typeof(x),lrem), convert(typeof(x),urem)
end

# n!*(e/n)^n
function stirling(n)
    if n < 500
        s = 1.0
        for i = 1:n
            s *= i/n*ℯ
        end
        return s
    else
        # 3rd-order Stirling's approximation more accurate for large n
        twn = 12*n
        return sqrt(2*pi*n)*(1 + twn\(1 + (2*twn)\(1 - (15*twn)\139)))
    end
end


"""
    compact_mnm_hyper(α::AbstractArray{Float32,1})

[Dirichlet Distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution)
"""
struct compact_mnm_hyper <: distribution_hyper_params
    α::AbstractArray{Float32,1}
end

mutable struct compact_mnm_sufficient_statistics <: sufficient_statistics
    N::Float32
    points_sum::AbstractArray{Float32,1}
end

function calc_posterior(prior:: compact_mnm_hyper, suff_statistics::Vector{Tuple{sufficient_statistics,Number}})
    if length(suff_statistics) == 0
        return prior
    end
    points_num = sum([post_kernel(x[2],global_time)*x[1].N for x in suff_statistics])
    if points_num == 0
        return prior
    end

    suff_sum = sum([post_kernel(x[2],global_time)*x[1].points_sum for x in suff_statistics],dims=1)[1]
    return compact_mnm_hyper(prior.α + suff_sum)
end

function sample_distribution(hyperparams::compact_mnm_hyper)
    return compact_mnm_dist(log.(rand(Dirichlet(Float64.(hyperparams.α)))))
end

function create_sufficient_statistics(hyper::compact_mnm_hyper,posterior::compact_mnm_hyper,points::AbstractArray{Float32,2}, pts_to_group = 0)
    if length(points) == 0
        return compact_mnm_sufficient_statistics(size(points,2),zeros(Int64,length(hyper.α)))
    end
    pt_count = counts(Int.(points),length(hyper.α))
    return compact_mnm_sufficient_statistics(size(points,2),pt_count)
end

function log_marginal_likelihood(hyper::compact_mnm_hyper, posterior_hyper::compact_mnm_hyper, suff_stats::Vector{Tuple{sufficient_statistics,Number}})
    D = length(suff_stats[1][1].points_sum)
    logpi = log(pi)
    val = logabsgamma(sum(hyper.α))[1] -logabsgamma(sum(posterior_hyper.α))[1] + sum((x-> logabsgamma(x)[1]).(posterior_hyper.α) - (x-> logabsgamma(x)[1]).(hyper.α))
    return val
end

function aggregate_suff_stats(suff_l::compact_mnm_sufficient_statistics, suff_r::compact_mnm_sufficient_statistics)
    return compact_mnm_sufficient_statistics(suff_l.N+suff_r.N, suff_l.points_sum + suff_r.points_sum)
end

function posterior_predictive!(r::AbstractArray,x::AbstractMatrix,posterior_hyper::compact_mnm_hyper)
    if length(posterior_hyper.α) == 0
        return
    end
    post_sum = float(sum(posterior_hyper.α))
    @inbounds for i in eachindex(r)
        r[i] = 0.0
        @inbounds for j=1:size(x,1)
            if x[j,i] > 0
                r[i] += log(posterior_hyper.α[Int64.(x[j,i])]/post_sum)
            end
        end
    end
end


# # Examples

## load packages
using BenchmarkTools
using WaterPhysics
FT = Float32;
#------------------------------------------------------------------------------

## define the variables
rand_r = (rand(FT) + 20) * FT(1e-6);
rand_T = rand(FT) + 298;
rand_α = rand(FT) * 50;
rand_Ψ = rand(FT) - 3;

gas_air = TraceGasAir{FT}();
gas_CO₂ = TraceGasCO₂{FT}();
#------------------------------------------------------------------------------

# ## Capillary pressure
@info "Benchmarking capillary_pressure functions...";
@btime capillary_pressure($rand_r, $rand_T);
@btime capillary_pressure($rand_r, $rand_T, $rand_α);
#------------------------------------------------------------------------------

# ## Diffusive coefficient
@info "Benchmarking diffusive_coefficient functions...";
@btime diffusive_coefficient($rand_T, $gas_CO₂, $gas_air);
@btime relative_diffusive_coefficient($rand_T);
#------------------------------------------------------------------------------

# ## Latent heat of evaporation
@info "Benchmarking latent_heat_vapor functions...";
@btime latent_heat_vapor($rand_T);
#------------------------------------------------------------------------------

# ## Saturation vapor pressure
@info "Benchmarking saturation_vapor_pressure functions...";
@btime pressure_correction($rand_T, $rand_Ψ);
@btime saturation_vapor_pressure($rand_T);
@btime saturation_vapor_pressure($rand_T, $rand_Ψ);
@btime saturation_vapor_pressure_slope($rand_T);
@btime saturation_vapor_pressure_slope($rand_T, $rand_Ψ);
#------------------------------------------------------------------------------

# ## Surface tension
@info "Benchmarking surface_tension functions...";
@btime surface_tension($rand_T);
@btime relative_surface_tension($rand_T);
#------------------------------------------------------------------------------

# ## Viscosity
@info "Benchmarking viscosity functions...";
@btime viscosity($rand_T);
@btime relative_viscosity($rand_T);
#------------------------------------------------------------------------------


module TestGPA
using GaussianProcesses, Calculus, Distributions
using Test, Random

Random.seed!(1)

@testset "GPA" begin
    d, n = 3, 20
    ll = rand(d)
    X = 2π * rand(d, n)
    f = [sum(sin, view(X, :, i)) / d for i in 1:n]

    # Test Bernoulli, binomial, exponential, Gaussian, Poisson, and Student-t likelihood
    liks = (BernLik(), BinLik(n), ExpLik(), GaussLik(-1.0), PoisLik(), StuTLik(3, 0.1))
    ys = ([Bool(rand(Distributions.Bernoulli(abs(f[i])))) for i in 1:n],
          [rand(Distributions.Binomial(n,exp(f[i]) / (1 + exp(f[i])))) for i in 1:n],
          [rand(Distributions.Exponential(f[i]^2)) for i in 1:n],
          f,
          [rand(Distributions.Poisson(exp(f[i]))) for i in 1:n],
          f .+ rand(Distributions.TDist(3), n))

    # Mean and kernel function
    mZero = MeanZero()
    kern = SE(0.0, 0.0)

    @testset "Likelihood $(typeof(lik))" for (lik, y) in zip(liks, ys)
        println("\tTesting ", nameof(typeof(lik)), "...")

        # Fit GP
        gp = GP(X, y, mZero, kern, lik)

        # Sample random parameters
        params = 0.5 * randn(GaussianProcesses.num_params(gp))
        set_params!(gp, params)

        # Exact gradient
        GaussianProcesses.update_target_and_dtarget!(gp)
        exact_grad = gp.dtarget

        # Numerical approximation
        num_grad = Calculus.gradient(params) do params
            set_params!(gp, params)
            update_target!(gp)
            gp.target
        end

        @test num_grad ≈ exact_grad
    end

end
end


function solve()
    a, b, c, d = [parse(Int, x) for x in split(readline())]

    takahashi = (c+b-1) ÷ b
    aoki = (a+d-1) ÷ d
    takahashi <= aoki ? "Yes" : "No"
end

println(solve())


###############################################################
######################GAMMA####################################
###############################################################

immutable PGamma <: ContinuousUnivariateDistribution
    params::Array{Float64}
end

PGamma() = PGamma([1.])

for op in [:pdf, :logpdf, :cdf, :quantile]
    @eval ($op)(pdist::PGamma, t::Real) = ($op)(Gamma(exp(pdist.params[1]),exp(-pdist.params[1])), t)
end

rand(pdist::PGamma) = rand(Gamma(exp(pdist.params[1]),exp(-pdist.params[1])))

ds(pdist::PGamma, t, int_coefs, i)    =  exp(pdist.params[1])*(log(t)-t) + int_coefs.aux[1]
d²s(pdist::PGamma, t, int_coefs,i, j)  = ds(pdist, t, int_coefs,i)+ int_coefs.aux[2]

function dl!(gl, pdist::PGamma, t, int_coefs)
    gl[2] = exp(pdist.params[1])
    gl[1] = muladd(gl[2],(log(t)-t), int_coefs.aux[1])
    gl[2] *= (t-1)
end

function d²l!(gl, hl, pdist::PGamma, t, int_coefs)
    hl[2,2] = gl[2] = exp(pdist.params[1])
    gl[1] = muladd(gl[2],(log(t)-t), int_coefs.aux[1])
    gl[2] *= (t-1)
    hl[1,1] = gl[1] + int_coefs.aux[2]
    hl[1,2] = hl[2,1] = gl[2]
    hl[2,2] *= -t
end

auxvec(pdist::PGamma) = [exp(pdist.params[1])*(-polygamma(0,exp(pdist.params[1]))+1+pdist.params[1]),
-exp(2*pdist.params[1])*polygamma(1,exp(pdist.params[1]))+exp(pdist.params[1])]

function initialize_aft(S::AbstractVector,X::AbstractArray, pdist::PGamma, N)
    St = [(s.t₀+s.t₁)/2 for s in S]
    valid = isfinite.(St)
    res = glm(X[valid,:], St[valid], Gamma(), LogLink())
    T = 1/GLM.dispersion(res, true)
    return vcat(log(T),coef(res))
end


@testset "VAE" begin

    using Distributions

    function vae_cost(We1, be1, We2, be2, We3, be3, We4, be4, Wd1, bd1, Wd2, bd2, Wd3, bd3, eps, x)
        dummy = 42.0
        # encoder
        he1 = tanh.(We1 * x .+ be1)
        he2 = tanh.(We2 * he1 .+ be2)
        mu = We3 * he2 .+ be3
        log_sigma2 = We4 * he2 .+ be4
        z = mu .+ sqrt.(exp.(log_sigma2)) .* eps
        # decoder
        hd1 = tanh.(Wd1 * z .+ bd1)
        hd2 = tanh.(Wd2 * hd1 .+ bd2)
        x_rec = logistic.(Wd3 * hd2 .+ bd3)
        # loss
        rec_loss = sum(x .* log.(1e-10 .+ x_rec) + (1 .- x) .* log.(1e-10 .+ 1 .- x_rec), dims=1)
        latent_loss = -0.5 * sum(1 .+ log_sigma2 .- mu .^ 2 .- exp.(log_sigma2), dims=1)
        cost = mean(rec_loss .+ latent_loss)
    end

    function xavier_init(dim_in, dim_out; c=1)
        low = -c * sqrt(6.0 / (dim_in + dim_out))
        high = c * sqrt(6.0 / (dim_in + dim_out))
        return rand(Uniform(low, high), dim_in, dim_out)
    end

    let
        We1 = xavier_init(500, 784); be1 = randn(500);
        We2 = xavier_init(500, 500); be2 = randn(500);
        We3 = xavier_init(20, 500); be3 = randn(20);
        We4 = xavier_init(20, 500); be4 = randn(20);
        Wd1 = xavier_init(500, 20); bd1 = randn(500);
        Wd2 = xavier_init(500, 500); bd2 = randn(500);
        Wd3 = xavier_init(784, 500); bd3 = randn(784);
        x = rand(784, 100); eps = rand(Normal(0, 1),  20)
        inputs = [:We1 => We1, :be1 => be1, :We2 => We2, :be2 => be2,
                  :We3 => We3, :be3 => be3, :We4 => We4, :be4 => be4,
                  :Wd1 => Wd1, :bd1 => bd1, :Wd2 => Wd2, :bd2 => bd2,
                  :Wd3 => Wd3, :bd3 => bd3, :eps => eps, :x => x]

        test_compare(vae_cost; inputs...)
    end
end


using JLD2
using Distributions
using StatsBase
using LinearAlgebra
using SparseArrays

"""
    COOccurence(coo_matrix, labels)
"""
mutable struct COOccurence
    coo_matrix::SparseMatrixCSC{Float64,Int64}
    ind2labels::Dict{Int, String}
end

# Note: Sample input occurrance files are those generated with PubMedMiner.jl
# https://github.com/bcbi/PubMedMiner.jl.git
function occur2coo(occur, label2ind)
    #invert map
    ind2label = Dict()
    for key  in keys(label2ind)
        val = label2ind[key]
        ind2label[val] = key
    end
    # labels = Array{String}(length(ind2label))
    # for i in range(1,length(ind2label))
    #     labels[i] = ind2label[i]
    # end

    coo_m = occur * transpose(occur)
    COO = COOccurence(coo_m, ind2label)
    return COO
end

function read_sp_occur(file="./occur_sp.jdl")
    fin  = jldopen(file, "r")
    occur = read(fin, "occur")
    return occur
end

function read_occur_label_dict(file="/mesh2ind.jdl")
    file  = jldopen(file, "r")
    label2ind = read(file, "mesh2ind")
    return label2ind
end

"""
    corrcoef(occur)

Compute correlation coeffiecients for a occurence data matrix
###Arguments

* `occur`: Occurrence data matrix
"""
function corrcoef(occur; vardim = 1)
    # compute the covariance matrix
    C = cov(occur, dims=vardim)

    # the correlation coefficients are given by
    # C_{i,j} / sqrt(C_{ii} * C_{jj})
    d = diag(C)
    CiiCjj = zeros(size(C))
    BLAS.ger!(1.0, d, d, CiiCjj) #CiiCjj += d*d'

    #should we devectorize?
    coeffs = C./sqrt.(CiiCjj)

    return coeffs

end

"""
    pmi_mat(coo_matrix, N)

Point-wise mutual information from co-occurrence  matrix

Point-wise mutual information
log p(x,y)/p(x)p(y)

A measure the takes into account chance by comparing the probability of
joint event to the product of the individual events

###Outcome

* `pmi`:PMI matrix (symmetric matrix)
"""
function pmi_mat(coo_matrix, N)
    d = diag(coo_matrix)

    #Compute conditionals P(X,Y)/P(X)
    cooccurrence_diagonal = diag(coo_matrix)
    conditional = coo_matrix./cooccurrence_diagonal

    # PMI (point-wise mutual information) P(X,Y)/P(X)P(Y)
    pmi_full = log.((N .* conditional') ./ cooccurrence_diagonal)
    # pmi = LowerTriangular(pmi_full)
end

"""
    chi2(coo_matrix)

Point-wise Chi square statistic from co-occurrence  matrix

Chi square statistic to test whether the occurence of one
feature is independent of another, e.i
Null hypothesis = p(x,y) = p(x)p(y)

###Arguments
* `occur`: Occurrance data matrix: Observations are rows, columns are features

###Outcome

* `chi2`:Matrix of Chi Square Statistics  (symmetric matrix)
"""
function chi2_mat(occur; min_freq=5)

    nvars = size(occur, 2)
    chi2 = LowerTriangular(zeros(nvars, nvars))
    for j1=1:nvars
        x = occur[:,j1]
        for j2=j1+1:nvars
            y = occur[:,j2]
            chi2[j2, j1] = chi2_statistic(x,y,span(occur); min_freq=min_freq)
        end
    end
    return chi2
end


using Test
using Distributions
using Random

using EnsembleKalmanProcesses.DataStorage

@testset "DataStorage" begin
    seed = 2021
    Random.seed!(seed)

    parameter_samples = rand(MvNormal(2, 0.1), 10) #10 samples of 4D params
    data_samples = rand(MvNormal(12, 2), 10) #10 samples of 12D data
    data_samples_short = data_samples[:, 1:(end - 1)]

    new_parameter_samples = rand(MvNormal(2, 0.1), 10) #10 samples of 4D params
    new_data_samples = rand(MvNormal(12, 2), 10) #10 samples of 12D data
    new_data_samples_short = new_data_samples[:, 1:(end - 1)]

    #test DataContainer
    parameter_samples_T = permutedims(parameter_samples, (2, 1))
    idata = DataContainer(parameter_samples)
    idata_T = DataContainer(parameter_samples_T, data_are_columns = false)

    @test get_data(idata) == parameter_samples
    @test get_data(idata_T) == parameter_samples
    @test size(idata) == size(parameter_samples)

    #test Paired Container
    idata = DataContainer(parameter_samples)
    odata = DataContainer(data_samples)
    odata_short = DataContainer(data_samples_short)

    @test_throws DimensionMismatch PairedDataContainer(parameter_samples, data_samples_short)
    iopairs = PairedDataContainer(parameter_samples, data_samples)

    @test_throws DimensionMismatch PairedDataContainer(idata, odata_short)
    iopairs2 = PairedDataContainer(idata, odata)

    @test get_data(iopairs) == get_data(iopairs2)
    @test get_inputs(iopairs) == parameter_samples
    @test get_outputs(iopairs) == data_samples

end


function drawcircle!(mat :: Matrix, xc, yc, r)

    x = 0;
    y = r;
    p = 1 - r;

     while x < y

         mat[round(xc+x), round(yc+y)] = 1
         mat[round(xc-x), round(yc+y)] = 1
         mat[round(xc+x), round(yc-y)] = 1
         mat[round(xc-x), round(yc-y)] = 1
         mat[round(xc+y), round(yc+x)] = 1
         mat[round(xc-y), round(yc+x)] = 1
         mat[round(xc+y), round(yc-x)] = 1
         mat[round(xc-y), round(yc-x)] = 1

         x += 1

         if p < 0
            p = p + 2 * x + 1
         else
            y = y - 1
            p = p + 2 * (x - y) + 1
         end

         mat[round(xc+x), round(yc+y)] = 1
         mat[round(xc-x), round(yc+y)] = 1
         mat[round(xc+x), round(yc-y)] = 1
         mat[round(xc-x), round(yc-y)] = 1
         mat[round(xc+y), round(yc+x)] = 1
         mat[round(xc-y), round(yc+x)] = 1
         mat[round(xc+y), round(yc-x)] = 1
         mat[round(xc-y), round(yc-x)] = 1
     end
end


using Plots

include("../src/parzenwindow.jl")

# Data generation
data1 = ones(20)*2+randn(20)
data2 = ones(10)*6+randn(10)
data = vcat(data1, data2)
data = reshape(data, (1, size(data)[1]))

# Plot data
plt = plot(title = "Parzen window estimation (window size 'h')", xlabel="Data value", ylabel="Likelihood")
scatter!(data1, zeros(20), color="black", label="Data")
scatter!(data2, zeros(10), color="black", label = "")

# Plot parzen window function
colors = ["red", "green", "blue", "black"]
kernel(x) = k(x)
for step in 0.1:0.5:1.6
    Y = []
    model(x) = create_parzen_window(step, data, kernel, x)
    for i in minimum(data):0.01:maximum(data)
        append!(Y, model(Vector{Float64}([i])))
    end
    plot!(plt, minimum(data):0.01:maximum(data), Y, color=colors[Int((step+0.4)*2)], label=string("h = ", string(step)))
end
plt


mutable struct Fact
    uin::Vector{Float64}
    uout::Vector{Ptr{Float64}}
    neigs::Vector{Int}
    w::Vector{Float64}
end

Fact() = Fact(Float64[], Ptr{Float64}[], Int[], Float64[])

deg(f::Fact) = length(f.uin)

mutable struct FGMatching
    N::Int
    γ::Float64
    fnodes::Vector{Fact}
    adjlist::Vector{Vector{Int}}
end

function FGMatching(net::Network; γ=Inf)
    @assert has_eprop(net, "w")

    N = nv(net)
    fnodes = [Fact() for i=1:N]

    # Prune Graph
    w = [Float64[] for i=1:N]
    adjlist = [Int[] for i=1:N]
    for i=1:N
        for e in edges(net, i)
            wij = eprop(net, e)["w"]
            if wij < γ
                push!(w[i], wij)
                push!(adjlist[i], dst(e))
            end
        end
    end

    for (i, v) in enumerate(fnodes)
        resize!(v.uin, length(adjlist[i]))
        resize!(v.uout, length(adjlist[i]))
        resize!(v.w, length(adjlist[i]))
    end

    for (i, f) in enumerate(fnodes)
        for (ki, j) in enumerate(adjlist[i])
            f.uin[ki] = 0
            kj = findfirst(==(i), adjlist[j])
            fnodes[j].uout[kj] = getref(f.uin, ki)
            f.w[ki] = w[i][ki]
        end
    end

    FGMatching(N, γ, fnodes, adjlist)
end

function initrand!(g::FGMatching)
    for f in g.fnodes
        for k=1:deg(f)
            f.uin[k] = randn()
        end
    end
end

function update!(f::Fact)
    @unpack w, uin, uout = f
    m1 = m2 = Inf
    i1 = 0
    for i=1:deg(f)
        h = w[i] - uin[i]
        if h < m1
            m2 = m1
            m1 = h
            i1 = i
        elseif h < m2
            m2 = h
        end
    end

    Δ = abs(uout[i1][] - m2)
    for i=1:deg(f)
        uout[i][] = m1
    end
    uout[i1][] = m2
    return Δ
end

function findmatch(f::Fact)
    @unpack w, uin = f
    m1 = Inf
    i1 = 0
    for i=1:deg(f)
        h = w[i] - uin[i]
        if h < m1
            m1 = h
            i1 = i
        end
    end
    return i1, w[i1]
end

function oneBPiter!(g::FGMatching)
    Δ = 0.
    for a in randperm(g.N)
        d = update!(g.fnodes[a])
        Δ = max(Δ, d)
    end
    return Δ
end

function converge!(g::FGMatching; maxiters=100, ϵ=1e-8, verbose=true)
    Eold = 0.
    tstop = 0
    it = 0
    while it < maxiters
        it += 1
        Δ = oneBPiter!(g)
        E, matchmap, nfails = energy(g)
        verbose && @printf("it=%d  E=%.5f  nfails=%d \tΔ=%f \n", it, E, nfails, Δ)

        if abs(Eold - E) < ϵ && nfails == 0
            tstop += 1
            if tstop == 10
                verbose && println("Found ground state")
                break
            end
        else
            tstop = 0
        end

        Eold = E
    end
    return Eold, it
end

function energy(g::FGMatching)
    @unpack fnodes, N, adjlist = g
    E = 0.
    matchmap = zeros(Int, N)
    for i=1:N
        f = fnodes[i]
        k, wij = findmatch(f)
        E += wij
        matchmap[i] = adjlist[i][k]
    end
    E /= 2
    nfails = 0
    for i=1:N
        j = matchmap[i]
        nfails += matchmap[j] != i
    end
    return E, matchmap, nfails
end

"""
Return the optimal cost and the matching map:
`matchmap[i] = j`  if (i,j) is in the optimal matching.

The cutoff on the costs is  γ.
"""
function run_bp(net::Network;
                γ = Inf,
                maxiters = 10000,
                ϵ = 1e-4,
                seed = -1,
                verbose=true)
    seed > 0 && Random.seed!(seed)
    g = FGMatching(net; γ)
    initrand!(g)
    iters = converge!(g; maxiters, ϵ, verbose)
    E, matchmap, nfails = energy(g)
    res = (energy = E,
           match = matchmap,
           bpgraph = g,
           num_violations = nfails,
           iters = iters,
           ok = (nfails == 0))
    return res
end


using DiffEqFlux, Flux
using LinearAlgebra, Distributions
using Optim, GalacticOptim
using Test

function run_test(f, layer, atol)

    data_train_vals = [rand(length(layer.model)) for k in 1:500]
    data_train_fn = f.(data_train_vals)

    function loss_function(component)
        data_pred = [layer(x,component) for x in data_train_vals]
        loss = sum(norm.(data_pred.-data_train_fn))/length(data_train_fn)
        return loss
    end

    function cb(p,l)
        @show l
        return false
    end

    optfunc = GalacticOptim.OptimizationFunction((x, p) -> loss_function(x), GalacticOptim.AutoZygote())
    optprob = GalacticOptim.OptimizationProblem(optfunc, layer.p)
    res = GalacticOptim.solve(optprob, ADAM(0.1), cb=cb, maxiters = 100)
    optprob = GalacticOptim.OptimizationProblem(optfunc, res.minimizer)
    res = GalacticOptim.solve(optprob, ADAM(0.01), cb=cb, maxiters = 100)
    optprob = GalacticOptim.OptimizationProblem(optfunc, res.minimizer)
    res = GalacticOptim.solve(optprob, BFGS(), cb=cb, maxiters = 200)
    opt = res.minimizer

    data_validate_vals = [rand(length(layer.model)) for k in 1:100]
    data_validate_fn = f.(data_validate_vals)

    data_validate_pred = [layer(x,opt) for x in data_validate_vals]

    return sum(norm.(data_validate_pred.-data_validate_fn))/length(data_validate_fn) < atol
end

##test 01: affine function, Chebyshev and Polynomial basis
A = rand(2,2)
b = rand(2)
f = x -> A*x + b
layer = TensorLayer([ChebyshevBasis(10), PolynomialBasis(10)], 2)
@test run_test(f, layer, 0.05)

##test 02: non-linear function, Chebyshev and Legendre basis
A = rand(2,2)
b = rand(2)
f = x -> A*x*norm(x)+ b*sin(norm(x))
layer = TensorLayer([ChebyshevBasis(7), FourierBasis(7)], 2)
@test run_test(f, layer, 0.10)


# =================================================================================================
# =                                         Hyland83.jl                                           =
# =================================================================================================
@testset "Hyland83" begin
# Appendix of ref. [1]. First table
@testset "Specific volume of saturated ice" begin
@test 1.0768e-3 ≈ MoistAir.volumeice(173.15) atol=0.0001e-3
@test 1.0829e-3 ≈ MoistAir.volumeice(223.15) atol=0.0001e-3
@test 1.0909e-3 ≈ MoistAir.volumeice(273.16) atol=0.0001e-3
end
@testset "Specific volume of saturated liquid water" begin
@test 1.00021e-3 ≈ MoistAir.volumewater(273.16) atol=0.00001e-3
@test 1.01215e-3 ≈ MoistAir.volumewater(323.15) atol=0.00001e-3
@test 1.04346e-3 ≈ MoistAir.volumewater(373.15) atol=0.00001e-3 # In the paper there is a typo in the power of 10: 1.04346e10
@test 1.09050e-3 ≈ MoistAir.volumewater(423.15) atol=0.00001e-3
@test 1.15653e-3 ≈ MoistAir.volumewater(473.15) atol=0.00001e-3
end

@testset "Saturation enthalpy of saturated ice" begin
@test -507.215e3 ≈ MoistAir.enthalpyice(173.15) atol=0.001e3
@test -429.413e3 ≈ MoistAir.enthalpyice(223.15) atol=0.001e3
@test -333.429e3 ≈ MoistAir.enthalpyice(273.15) atol=0.001e3  # The table presents -333.409. Typo???
end

@testset "Saturation enthalpy of saturated liquid water" begin
@test 0.0 ≈ MoistAir.enthalpywater(273.16) atol=1.0
@test 209.330e3 ≈ MoistAir.enthalpywater(323.15) atol=5
@test 419.158e3 ≈ MoistAir.enthalpywater(373.15) atol=5
@test 632.210e3 ≈ MoistAir.enthalpywater(423.15) atol=9
@test 852.329e3 ≈ MoistAir.enthalpywater(473.15) atol=15
end

# Table 2 of reference [2]
@testset "Enhancement factor" begin
P = 0.1e6
@test 1.0105 ≈ MoistAir.efactor(173.15, P) atol=0.0001
@test 1.0039 ≈ MoistAir.efactor(273.15, P) atol=0.0001
@test 1.0039 ≈ MoistAir.efactor(363.15, P) atol=0.0001
P = 0.5e6
@test 1.054 ≈ MoistAir.efactor(173.15, P) atol=0.001
@test 1.0177 ≈ MoistAir.efactor(273.15, P) atol=0.0006  # I think the table has a typo!!?
@test 1.0180 ≈ MoistAir.efactor(363.15, P) atol=0.0001
@test 1.0188 ≈ MoistAir.efactor(373.15, P) atol=0.0001
@test 1.0022 ≈ MoistAir.efactor(423.15, P) atol=0.0001
P = 1.0e6
@test 1.113 ≈ MoistAir.efactor(173.15, P) atol=0.001
@test 1.0353 ≈ MoistAir.efactor(273.15, P) atol=0.0015  # I think the table has a typo!!?
@test 1.0284 ≈ MoistAir.efactor(363.15, P) atol=0.0001
@test 1.0295 ≈ MoistAir.efactor(373.15, P) atol=0.0001
@test 1.0288 ≈ MoistAir.efactor(423.15, P) atol=0.0001
@test 1.0235 ≈ MoistAir.efactor(433.15, P) atol=0.0001
@test 1.0142 ≈ MoistAir.efactor(443.15, P) atol=0.0001

P = 5.0e6
@test 1.820 ≈ MoistAir.efactor(173.15, P) atol=0.001
@test 1.191 ≈ MoistAir.efactor(273.15, P) atol=0.008  # I think the table has a typo!!?
@test 1.1102 ≈ MoistAir.efactor(363.15, P) atol=0.0001
@test 1.1082 ≈ MoistAir.efactor(373.15, P) atol=0.0001
@test 1.111 ≈ MoistAir.efactor(423.15, P) atol=0.001
@test 1.114 ≈ MoistAir.efactor(433.15, P) atol=0.001
@test 1.116 ≈ MoistAir.efactor(443.15, P) atol=0.001
@test 1.117 ≈ MoistAir.efactor(453.15, P) atol=0.001
@test 1.116 ≈ MoistAir.efactor(473.15, P) atol=0.001
end

@testset "Saturation pressure of water vapor over ice" begin
# Second table in the appendix of reference [1]
@test 1.40510e-3 ≈ MoistAir.Pws_s(173.15) atol=0.00001e-3
@test 6.11153e2 ≈ MoistAir.Pws_s(273.15) atol=0.00001e2
end

@testset "Saturation pressure of water vapor over liquid water" begin
# Second table in the appendix of reference [1]
@test 6.11213e2 ≈ MoistAir.Pws_l(273.15) atol=0.00001e2
@test 1.01419e5 ≈ MoistAir.Pws_l(373.15) atol=0.00001e5
@test 1.55507e6 ≈ MoistAir.Pws_l(473.15) atol=0.00001e6
end

@testset " saturation temperature function" begin
@test MoistAir.Tws(MoistAir.Pws_s(173.15)) ≈ 173.15 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_s(223.15)) ≈ 223.15 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_s(273.15)) ≈ 273.15 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_s(273.16)) ≈ 273.16 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_l(273.16)) ≈ 273.16 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_l(323.15)) ≈ 323.15 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_l(373.15)) ≈ 373.15 atol=1e-5
@test MoistAir.Tws(MoistAir.Pws_l(473.15)) ≈ 473.15 atol=1e-5
end

@testset "First virial coefficient of saturated vapor B'" begin
# Second table in the appendix of reference [1]
@test -3.2939e-5 ≈ MoistAir.Blin(173.15) atol=0.0001e-5
@test -8.3497e-7 ≈ MoistAir.Blin(273.15) atol=0.0001e-7
@test -1.4658e-7 ≈ MoistAir.Blin(373.15) atol=0.0001e-5
@test -5.0508e-8 ≈ MoistAir.Blin(473.15) atol=0.0001e-5
end

@testset "Second virial coefficient of saturated vapor C'" begin
# Second table in the appendix of reference [1]
@test -4.6563e-9 ≈ MoistAir.Clin(173.15) atol=0.0001e-9
@test -2.0928e-12 ≈ MoistAir.Clin(273.15) atol=0.0001e-12
@test -5.7548e-14 ≈ MoistAir.Clin(373.15) atol=0.0001e-14
@test -6.3933e-15 ≈ MoistAir.Clin(473.15) atol=0.0001e-15
end

@testset "Specific enthalpy of saturated water vapor" begin
# First table of the appendix of ref. [1]
@test 2315.87 ≈ MoistAir.enthalpyvapor(173.15)/1000 atol=0.01
@test 2408.41 ≈ MoistAir.enthalpyvapor(223.15)/1000 atol=0.01
@test 2500.81 ≈ MoistAir.enthalpyvapor(273.16)/1000 atol=0.01
@test 2591.29 ≈ MoistAir.enthalpyvapor(323.15)/1000 atol=0.01
@test 2675.46 ≈ MoistAir.enthalpyvapor(373.15)/1000 atol=0.01
@test 2746.15 ≈ MoistAir.enthalpyvapor(423.15)/1000 atol=0.01
@test 2793.11 ≈ MoistAir.enthalpyvapor(473.15)/1000 atol=0.01
end

@testset "Specific volume of saturated water vapor" begin
# First table of the appendix of ref. [1]
@test 5.6873e7 ≈ MoistAir.volumevapor(173.15) atol=0.0001e7
@test 2.6146e4 ≈ MoistAir.volumevapor(223.15) atol=0.0001e4
@test 2.0601e2 ≈ MoistAir.volumevapor(273.16) atol=0.0001e2
@test 1.2030e1 ≈ MoistAir.volumevapor(323.15) atol=0.0001e1
@test 1.6718e0 ≈ MoistAir.volumevapor(373.15) atol=0.0001e0
@test 3.9253e-1 ≈ MoistAir.volumevapor(423.15) atol=0.0001e-1
@test 1.2722e-1 ≈ MoistAir.volumevapor(473.15) atol=0.0001e0
end

@testset "Specific entropy of saturated water vapor" begin
# First table of the appendix of ref. [1]
@test 14.30387 ≈ MoistAir.entropyvapor(173.15)/1000 atol=0.00005
@test 11.10965 ≈ MoistAir.entropyvapor(223.15)/1000 atol=0.00005
@test  9.15510 ≈ MoistAir.entropyvapor(273.16)/1000 atol=0.00005
@test  8.07477 ≈ MoistAir.entropyvapor(323.15)/1000 atol=0.00005
@test  7.35365 ≈ MoistAir.entropyvapor(373.15)/1000 atol=0.00005
@test  6.83731 ≈ MoistAir.entropyvapor(423.15)/1000 atol=0.00005
@test  6.43218 ≈ MoistAir.entropyvapor(473.15)/1000 atol=0.00005
end
end


# Val{true} is passed at the last mp iteration and Val{false} during the previous iterations.

function approxTranslation( cmatrix::Array{T,N}, alg::String, ::Val{true} ) where {T<:AbstractFloat,N}
    ccenter = div.( size( cmatrix ), 2 );
    peak1, maxVal = firstPeak( cmatrix );
    return subPixel( cmatrix, alg, peak1, maxVal ) .- ccenter;
end

function approxTranslation( cmatrix::Array{T,N}, alg::String, ::Val{false} ) where {T<:AbstractFloat,N}
    ccenter = div.( size( cmatrix ) , 2 );
    return firstPeak( cmatrix )[1] .- ccenter
end

function subPixel( cmatrix::Array{T,N}, alg::String, peak1::NTuple{N,I}, maxVal::T ) where {T<:AbstractFloat,N}

    refinement = zeros( T, N );

    if alg == "" || indexOnBorder( peak1, size( cmatrix ) )
        # no sub-pixel refinement
    else
        neighs = getNeighbours( peak1, cmatrix );

        # ensuring values are positive
        minVal = minimum( neighs );
        for idx in 1:length(neighs)
            neighs[idx] = 1 + neighs[idx] - minVal
        end

        if     alg == "gaussian"
            refinement = point3Gaussian( neighs, 1+maxVal-minVal );
        elseif alg == "centroid"
            refinement = centroid( neighs, maxVal );
        end

        # nan occur when neighbours contains 0's, resulting in 0/0
        refinement = [ isnan(x) ? 0.0 : x for x in refinement ]
    end

    return peak1 .+ refinement
end

# 3-POINT GAUSSIAN SUBPIXEL

# 3-point Gaussian subpixel 2D

function point3Gaussian( ns::Array{T,2}, mx::T ) where {T<:AbstractFloat}
    return point3_2D( log(ns[4]), log(ns[6]), log(ns[2]), log(ns[8]), log(mx) );
end

function point3_2D( up::T, down::T, left::T, right::T, mx::T ) where {T<:AbstractFloat}
    return [ ( up  - down )/( 2* up  - 4*mx + 2*down  ),
             (left - right)/( 2*left - 4*mx + 2*right ) ];
end

# 3-point Gaussian subpixel 3D
function point3Gaussian( ns::Array{T,3}, mx::T ) where {T<:AbstractFloat}
    return point3_3D( log(ns[13]), log(ns[15]), log(ns[11]), log(ns[17]),
                      log(ns[ 5]), log(ns[23]), log(mx) );
end

function point3_3D( up::T, down::T, left::T, right::T, front::T, back::T, mx::T ) where {T<:AbstractFloat}
    return [ (  up  - down )/( 2*  up  - 4*mx + 2*down  ),
             ( left - right)/( 2* left - 4*mx + 2*right ),
             (front - back )/( 2*front - 4*mx + 2*back  ) ];
end

# CENTROID SUBPIXEL

# Centroid subpixel 2D

function centroid( ns::Array{T,2}, mx::T ) where {T<:AbstractFloat}
	return centroid_2D( ns[4], ns[6], ns[2], ns[8], mx )
end

function centroid_2D( up::T, down::T, left::T, right::T, mx::T ) where {T<:AbstractFloat}
    return [ ( up  - down )/(  up  - mx + down  ),
             (left - right)/( left - mx + right ) ];
end

# Centroid subpixel 3D

function centroid( ns::Array{T,3}, mx::T ) where {T<:AbstractFloat}
	return centroid_3D( ns[13], ns[15], ns[11], ns[17], ns[5], ns[23], mx )
end

function centroid_3D( up::T, down::T, left::T, right::T, front::T, back::T, mx::T ) where {T<:AbstractFloat}
    return [ (  up  - down )/(   up  - mx + down  ),
             ( left - right)/(  left - mx + right ),
			 (front - back )/( front - mx + back  ) ];
end


#Code written by Skipper af Rosenborg and Esben Bager
#File with functions generating LASSO results
using JuMP
using Gurobi
env = Gurobi.Env()

function generateLassoModel(Xtrain, Ytrain, gamma)
	bCols = size(Xtrain)[2]
	M = JuMP.Model(solver = GurobiSolver(env, OutputFlag = 0, Threads=1))
	@variables M begin
			b[1:bCols]
			t[1:bCols]
			w
	end
	@objective(M,Min,w+gamma*ones(bCols)'*t)
	@constraint(M, soc, norm( [1-w; sqrt(2)*(Xtrain*b-Ytrain)] ) <= 1+w) #second order cone constraint
	@constraint(M,  b .<= t)
	@constraint(M, -t .<= b)

	solve(M)

#	println("solved by worker $(myid())")
	return getvalue(b)

#	return M
end

function lassoValidationAndTest(gamma, xTrain, yTrain, xVali, yVali, xTest, yTest)
	beta = generateLassoModel(xTrain, yTrain, gamma)

	#Shrink LASSO values
	for i in 1:length(beta)
		if beta[i] <= 1e-6
			if beta[i] >= -1e-6
				beta[i] = 0
			end
		end
	end

	k = countnz(beta)
	if k == 0
		maxPairwise = 0
	else
		maxPairwise = findMaxCor(xTrain, beta)
	end

	valiRsqr = getRSquared(xVali, yVali, beta)
	testRsqr = getRSquared(xTest, yTest, beta)

	return maxPairwise, testRsqr, valiRsqr, k, beta
end

function processOutput(Xtrain, Ytrain, Xpred, Ypred, bSolved)
	for i in 1:length(bSolved)
		if bSolved[i] <= 1e-6
			if bSolved[i] >= -1e-6
				bSolved[i] = 0
			end
		end
	end

	#In-Sample R-squared value
	errors = (Ytrain-Xtrain*bSolved)
	errorTotal = sum(errors[i]^2 for i=1:length(errors))
	errorsMean = Ytrain-mean(Ytrain)
	errorMeanTotal = sum(errorsMean[i]^2 for i=1:length(errorsMean))
	ISRsquared = 1-(errorTotal/errorMeanTotal)

	#Indicator Results
	YpredValue = Ypred[1]
	Yestimate = Xpred'*bSolved
	YestimateValue = Yestimate[1]
	if YestimateValue > 0
		Indicator = 1
	else
		Indicator = 0
	end

	return ISRsquared, Indicator, YestimateValue, bSolved
end

function generatSolveAndProcess(Xtrain, Ytrain, Xpred, Ypred, gamma)
	bSolved = generateLassoModel(Xtrain, Ytrain, gamma)

	ISRsquared, Indicator, YestimateValue, bSolved = processOutput(Xtrain, Ytrain, Xpred, Ypred, bSolved)

	return ISRsquared, Indicator, YestimateValue, bSolved
end

ML/Code to hand in/ParallelModelGeneration.jl


@testset "Quire" begin

    # Test 32+1-32 == 1
    # as for Posit8 32+1 = 32
    q = zero(Quire8)
    q = fma(q,Posit8(32),one(Posit8))
    q = fma(q,Posit8(1),one(Posit8))
    q = fms(q,Posit8(32),one(Posit8))
    @test one(Posit8) == Posit8(q)

    # Test 1e6 + 1 - 1e6 == 1
    # as for Posit16 1e6+1 = 1e6
    q = zero(Quire16)
    q = fma(q,Posit16(1e6),one(Posit16))
    q = fma(q,Posit16(1),one(Posit16))
    q = fms(q,Posit16(1e6),one(Posit16))
    @test one(Posit16) == Posit16(q)

    # Test 1e30 + 1 - 1e30 == 1
    # as for Posit32 1e30+1 = 1e6
    q = zero(Quire32)
    q = fma(q,Posit32(1e30),one(Posit32))
    q = fma(q,Posit32(1),one(Posit32))
    q = fms(q,Posit32(1e30),one(Posit32))
    @test one(Posit32) == Posit32(q)

    @test zero(Posit8) == Posit8(zero(Quire8))
    @test zero(Posit16) == Posit16(zero(Quire16))
    @test zero(Posit32) == Posit32(zero(Quire32))

    @test one(Posit8) == Posit8(one(Quire8))
    @test one(Posit16) == Posit16(one(Quire16))
    @test one(Posit32) == Posit32(one(Quire32))
end


immutable Cubic{BC<:Flag} <: Degree{3} end
Cubic{BC<:Flag}(::BC) = Cubic{BC}()

"""
Assuming uniform knots with spacing 1, the `i`th piece of cubic spline
implemented here is defined as follows.

    y_i(x) = cm p(x-i) + c q(x-i) + cp q(1- (x-i)) + cpp p(1 - (x-i))

where

    p(δx) = 1/6 * (1-δx)^3
    q(δx) = 2/3 - δx^2 + 1/2 δx^3

and the values `cX` for `X ∈ {m, _, p, pp}` are the pre-filtered coefficients.

For future reference, this expands out to the following polynomial:

    y_i(x) = 1/6 cm (1+i-x)^3 + c (2/3 - (x-i)^2 + 1/2 (x-i)^3) +
             cp (2/3 - (1+i-x)^2 + 1/2 (1+i-x)^3) + 1/6 cpp (x-i)^3

When we derive boundary conditions we will use derivatives `y_0'(x)` and
`y_0''(x)`
"""
Cubic

"""
`define_indices_d` for a cubic b-spline calculates `ix_d = floor(x_d)` and
`fx_d = x_d - ix_d` (corresponding to `i` `and `δx` in the docstring for
`Cubic`), as well as auxiliary quantities `ixm_d`, `ixp_d` and `ixpp_d`
"""
function define_indices_d{BC}(::Type{BSpline{Cubic{BC}}}, d, pad)
    symix, symixm, symixp = Symbol("ix_",d), Symbol("ixm_",d), Symbol("ixp_",d)
    symixpp, symx, symfx = Symbol("ixpp_",d), Symbol("x_",d), Symbol("fx_",d)
    quote
        # ensure that all of ix_d, ixm_d, ixp_d, and ixpp_d are in-bounds no
        # matter the value of pad
        $symix = clamp(floor(Int, $symx), $(2-pad), size(itp,$d)+$(pad-2))
        $symfx = $symx - $symix
        $symix += $pad # padding for oob coefficient
        $symixm = $symix - 1
        $symixp = $symix + 1
        $symixpp = $symixp + 1
    end
end

"""
`define_indices_d`  for a cubic, periodic b-spline calculates `ix_d = floor(x_d)`
and `fx_d = x_d - ix_d` (corresponding to `i` and `δx` in the docstring entry
for `Cubic`), as well as auxiliary quantities `ixm_d`, `ixp_d` and `ixpp_d`.

If any `ixX_d` for `x ∈ {m, p, pp}` (note: not `c_d`) should fall outside of
the data interval, they wrap around.
"""
function define_indices_d(::Type{BSpline{Cubic{Periodic}}}, d, pad)
    symix, symixm, symixp = Symbol("ix_",d), Symbol("ixm_",d), Symbol("ixp_",d)
    symixpp, symx, symfx = Symbol("ixpp_",d), Symbol("x_",d), Symbol("fx_",d)
    quote
        $symix = clamp(floor(Int, $symx), 1, size(itp,$d))
        $symfx = $symx - $symix
        $symixm = mod1($symix - 1, size(itp,$d))
        $symixp = mod1($symix + 1, size(itp,$d))
        $symixpp = mod1($symix + 2, size(itp,$d))
    end
end

padding{BC<:Flag}(::Type{BSpline{Cubic{BC}}}) = Val{1}()
padding(::Type{BSpline{Cubic{Periodic}}}) = Val{0}()

"""
In `coefficients` for a cubic b-spline we assume that `fx_d = x-ix_d`
and we define `cX_d` for `X ⋹ {m, _, p, pp}` such that

    cm_d  = p(fx_d)
    c_d   = q(fx_d)
    cp_d  = q(1-fx_d)
    cpp_d = p(1-fx_d)

where `p` and `q` are defined in the docstring entry for `Cubic`, and
`fx_d` in the docstring entry for `define_indices_d`.
"""
function coefficients{C<:Cubic}(::Type{BSpline{C}}, N, d)
    symm, sym =  Symbol("cm_",d), Symbol("c_",d)
    symp, sympp = Symbol("cp_",d) ,Symbol("cpp_",d)
    symfx = Symbol("fx_",d)
    symfx_cub = Symbol("fx_cub_", d)
    sym_1m_fx_cub = Symbol("one_m_fx_cub_", d)
    quote
        $symfx_cub = cub($symfx)
        $sym_1m_fx_cub = cub(1-$symfx)
        $symm = SimpleRatio(1,6)*$sym_1m_fx_cub
        $sym  = SimpleRatio(2,3) - sqr($symfx) + SimpleRatio(1,2)*$symfx_cub
        $symp = SimpleRatio(2,3) - sqr(1-$symfx) + SimpleRatio(1,2)*$sym_1m_fx_cub
        $sympp = SimpleRatio(1,6)*$symfx_cub
    end
end

"""
In `gradient_coefficients` for a cubic b-spline we assume that `fx_d = x-ix_d`
and we define `cX_d` for `X ⋹ {m, _, p, pp}` such that

    cm_d  = p'(fx_d)
    c_d   = q'(fx_d)
    cp_d  = q'(1-fx_d)
    cpp_d = p'(1-fx_d)

where `p` and `q` are defined in the docstring entry for `Cubic`, and
`fx_d` in the docstring entry for `define_indices_d`.
"""
function gradient_coefficients{C<:Cubic}(::Type{BSpline{C}}, d)
    symm, sym, symp, sympp = Symbol("cm_",d), Symbol("c_",d), Symbol("cp_",d), Symbol("cpp_",d)
    symfx = Symbol("fx_",d)
    symfx_sqr = Symbol("fx_sqr_", d)
    sym_1m_fx_sqr = Symbol("one_m_fx_sqr_", d)
    quote
        $symfx_sqr = sqr($symfx)
        $sym_1m_fx_sqr = sqr(1 - $symfx)

        $symm  = -SimpleRatio(1,2) * $sym_1m_fx_sqr
        $sym   =  SimpleRatio(3,2) * $symfx_sqr     - 2 * $symfx
        $symp  = -SimpleRatio(3,2) * $sym_1m_fx_sqr + 2 * (1 - $symfx)
        $sympp =  SimpleRatio(1,2) * $symfx_sqr
    end
end

"""
In `hessian_coefficients` for a cubic b-spline we assume that `fx_d = x-ix_d`
and we define `cX_d` for `X ⋹ {m, _, p, pp}` such that

    cm_d  = p''(fx_d)
    c_d   = q''(fx_d)
    cp_d  = q''(1-fx_d)
    cpp_d = p''(1-fx_d)

where `p` and `q` are defined in the docstring entry for `Cubic`, and
`fx_d` in the docstring entry for `define_indices_d`.
"""
function hessian_coefficients{C<:Cubic}(::Type{BSpline{C}}, d)
    symm, sym, symp, sympp = Symbol("cm_",d), Symbol("c_",d), Symbol("cp_",d), Symbol("cpp_",d)
    symfx = Symbol("fx_",d)
    quote
        $symm  = 1 - $symfx
        $sym   = 3 * $symfx - 2
        $symp  = 1 - 3 * $symfx
        $sympp = $symfx
    end
end

function index_gen{C<:Cubic,IT<:DimSpec{BSpline}}(::Type{BSpline{C}}, ::Type{IT}, N::Integer, offsets...)
    if length(offsets) < N
        d = length(offsets)+1
        symm, sym, symp, sympp =  Symbol("cm_",d), Symbol("c_",d), Symbol("cp_",d), Symbol("cpp_",d)
        return :($symm * $(index_gen(IT, N, offsets...,-1)) + $sym * $(index_gen(IT, N, offsets..., 0)) +
                 $symp * $(index_gen(IT, N, offsets..., 1)) + $sympp * $(index_gen(IT, N, offsets..., 2)))
    else
        indices = [offsetsym(offsets[d], d) for d = 1:N]
        return :(itp.coefs[$(indices...)])
    end
end

# ------------ #
# Prefiltering #
# ------------ #

"""
`Cubic`: continuity in function value, first and second derivatives yields

    2/3 1/6
    1/6 2/3 1/6
        1/6 2/3 1/6
           ⋱  ⋱   ⋱
"""
function inner_system_diags{T,C<:Cubic}(::Type{T}, n::Int, ::Type{C})
    du = fill(convert(T, SimpleRatio(1, 6)), n-1)
    d = fill(convert(T, SimpleRatio(2, 3)), n)
    dl = copy(du)
    dl, d, du
end

"""
`Cubic{Flat}` `OnGrid` amounts to setting `y_1'(x) = 0` at `x = 1`.
Applying this condition yields

    -cm + cp = 0
"""
function prefiltering_system{T,TC}(::Type{T}, ::Type{TC}, n::Int,
                                   ::Type{Cubic{Flat}}, ::Type{OnGrid})
    dl, d, du = inner_system_diags(T, n, Cubic{Flat})
    d[1] = d[end] = -one(T)
    du[1] = dl[end] = zero(T)

    # Now Woodbury correction to set `[1, 3], [n, n-2] ==> 1`
    specs = WoodburyMatrices.sparse_factors(T, n, (1, 3, one(T)), (n, n-2, one(T)))

    Woodbury(lufact!(Tridiagonal(dl, d, du), Val{false}), specs...), zeros(TC, n)
end

"""
`Cubic{Flat}`, `OnCell` amounts to setting `y_1'(x) = 0` at `x = 1/2`.
Applying this condition yields

    -9/8 cm + 11/8 c - 3/8 cp + 1/8 cpp = 0

or, equivalently,

    -9 cm + 11 c -3 cp + 1 cpp = 0

(Note that we use `y_1'(x)` although it is strictly not valid in this domain; if we
were to use `y_0'(x)` we would have to introduce new coefficients, so that would not
close the system. Instead, we extend the outermost polynomial for an extra half-cell.)
"""
function prefiltering_system{T,TC}(::Type{T}, ::Type{TC}, n::Int,
                                   ::Type{Cubic{Flat}}, ::Type{OnCell})
    dl, d, du = inner_system_diags(T,n,Cubic{Flat})
    d[1] = d[end] = -9
    du[1] = dl[end] = 11

    # now need Woodbury correction to set :
    #    - [1, 3] and [n, n-2] ==> -3
    #    - [1, 4] and [n, n-3] ==> 1
    specs = WoodburyMatrices.sparse_factors(T, n,
                                  (1, 3, T(-3)),
                                  (n, n-2, T(-3)),
                                  (1, 4, one(T)),
                                  (n, n-3, one(T))
                                  )

    Woodbury(lufact!(Tridiagonal(dl, d, du), Val{false}), specs...), zeros(TC, n)
end

"""
`Cubic{Line}` `OnCell` amounts to setting `y_1''(x) = 0` at `x = 1/2`.
Applying this condition yields

    3 cm -7 c + 5 cp -1 cpp = 0

(Note that we use `y_1'(x)` although it is strictly not valid in this domain; if we
were to use `y_0'(x)` we would have to introduce new coefficients, so that would not
close the system. Instead, we extend the outermost polynomial for an extra half-cell.)
"""
function prefiltering_system{T,TC}(::Type{T}, ::Type{TC}, n::Int,
                                   ::Type{Cubic{Line}}, ::Type{OnCell})
    dl,d,du = inner_system_diags(T,n,Cubic{Line})
    d[1] = d[end] = 3
    du[1] = dl[end] = -7

    # now need Woodbury correction to set :
    #    - [1, 3] and [n, n-2] ==> -3
    #    - [1, 4] and [n, n-3] ==> 1
    specs = WoodburyMatrices.sparse_factors(T, n,
                                  (1, 3, T(5)),
                                  (n, n-2, T(5)),
                                  (1, 4, -one(T)),
                                  (n, n-3, -one(T))
                                  )

    Woodbury(lufact!(Tridiagonal(dl, d, du), Val{false}), specs...), zeros(TC, n)
end

"""
`Cubic{Line}` `OnGrid` amounts to setting `y_1''(x) = 0` at `x = 1`. Applying this
condition gives:

    1 cm -2 c + 1 cp = 0
"""
function prefiltering_system{T,TC}(::Type{T}, ::Type{TC}, n::Int,
                                   ::Type{Cubic{Line}}, ::Type{OnGrid})
    dl,d,du = inner_system_diags(T,n,Cubic{Line})
    d[1] = d[end] = 1
    du[1] = dl[end] = -2

    # now need Woodbury correction to set :
    #    - [1, 3] and [n, n-2] ==> 1
    specs = WoodburyMatrices.sparse_factors(T, n,
                                  (1, 3, one(T)),
                                  (n, n-2, one(T)),
                                  )

    Woodbury(lufact!(Tridiagonal(dl, d, du), Val{false}), specs...), zeros(TC, n)
end

"""
`Cubic{Periodic}` `OnGrid` closes the system by looking at the coefficients themselves
as periodic, yielding

    c0 = c(N+1)

where `N` is the number of data points.
"""
function prefiltering_system{T,TC,GT<:GridType}(::Type{T}, ::Type{TC}, n::Int,
                                                ::Type{Cubic{Periodic}}, ::Type{GT})
    dl, d, du = inner_system_diags(T,n,Cubic{Periodic})

    specs = WoodburyMatrices.sparse_factors(T, n,
                                  (1, n, du[1]),
                                  (n, 1, dl[end])
                                  )

    Woodbury(lufact!(Tridiagonal(dl, d, du), Val{false}), specs...), zeros(TC, n)
end

"""
`Cubic{Free}` `OnGrid` and `Cubic{Free}` `OnCell` amount to requiring an extra
continuous derivative at the second-to-last cell boundary; this means
`y_1'''(2) = y_2'''(2)`, yielding

    1 cm -3 c + 3 cp -1 cpp = 0
"""
function prefiltering_system{T,TC,GT<:GridType}(::Type{T}, ::Type{TC}, n::Int,
                                                ::Type{Cubic{Free}}, ::Type{GT})
    dl, d, du = inner_system_diags(T,n,Cubic{Periodic})

    specs = WoodburyMatrices.sparse_factors(T, n,
                                  (1, n, du[1]),
                                  (n, 1, dl[end])
                                  )

    Woodbury(lufact!(Tridiagonal(dl, d, du), Val{false}), specs...), zeros(TC, n)
end


# load packages
using CSV
using Random

include(pwd()*"/src/SDEMEM OU neuron data/load_neuron_data.jl")

"""
    set_up(M::Int=312, seed::Int=100)

Returns data and prior distributions.
"""
function set_up(M::Int=312, seed::Int=100)

    # define priors
    μ_0_1 = log(1/10); M_0_1 = 1; α_1 = 2; β_1 = 1
    μ_0_2 = log(1.5); M_0_2 = 1; α_2 = 2; β_2 = 1
    μ_0_3 = log(0.5); M_0_3 = 1; α_3 = 2; β_3 = 1

    prior_parameters_η = [μ_0_1 M_0_1 α_1 β_1;μ_0_2 M_0_2 α_2 β_2;μ_0_3 M_0_3 α_3 β_3]
    prior_parameters_σ_ϵ = [-1; 1/1]

    # load data
    ISI_all, nbr_ISE = load_neuron_data()
    y = ISI_all[1:M]

    # return: data and parameteres for priors
    return y,prior_parameters_η,prior_parameters_σ_ϵ

end


"""
    TriangularPrism(a, b, c, n, height, offset)
    TriangularPrism(a, b, c, height, offset = 0)

Create a `TriangularPrism` that can be used for spatial lookup.

`a`, `b`, and `c` are the points of the triangle, `n` if provided is the normal
for the plane defined by `a`, `b`, and `c`.

The `height` is along the normal of the plane defined by the points of the triangle.
`offset` is a distance along the negative of the normal to `offset` the `height`.

For example, a prism with `height` 1 and `offset` 0 extends from the plane a
distance of 1 along the normal.  A prism with `height` 1 and `offset` 0.5 extends
from a plane 0.5 along the negative of the normal and a distance of 1 along
the normal, in essence centering the volume for the prism evenly on both sides
of the plane defined by the three points of the triangle.

See also `AbstractRegion`
"""
struct TriangularPrism{T <: Real} <: AbstractRegion{T}
    a::SVector{3,T}
    b::SVector{3,T}
    c::SVector{3,T}
    n::SVector{3,T}
    height::T
    offset::T
end

@inline function TriangularPrism(a::StaticVector{3,T}, b::StaticVector{3,T}, c::StaticVector{3,T}, height::T, offset::T = 0.0) where T <: Real
    v1 = b - a
    v2 = c - a
    n = cross(v1, v2)
    n = n / norm(n)

    return TriangularPrism{T}(a, b, c, n, height, offset)
end

"""
    volume(prism::TriangularPrism)

Returns the volume of the `prism`.
"""
@inline volume(t::TriangularPrism) = t.height * dot(t.n, cross(t.b - t.a, t.c - t.a)) / 2.0

function boundingbox(t::TriangularPrism{T}) where T
    points = vcat(
        map(x -> x - t.offset * t.n, [t.a, t.b, t.c]),
        map(x -> x + (t.height - t.offset) * t.n, [t.a, t.b, t.c])
    )

    xmin, xmax = extrema(map(x -> x[1], points))
    ymin, ymax = extrema(map(x -> x[2], points))
    zmin, zmax = extrema(map(x -> x[3], points))
    return BoundingBox{T}(xmin, ymin, zmin, xmax, ymax, zmax)
end

function Base.in(p::StaticVector{3,T}, t::Region) where T <: Real where Region <: TriangularPrism
    a, b, c, n, h = t.a, t.b, t.c, t.n, t.height

    distance  = dot(n, p - a)

    if distance < -t.offset || distance > (t.height - t.offset)
        return false
    end

    # project p to the plane
    p = p - t.n * distance

    # determine if the point is in the triangle when projected to the plane
    # using barycentric coordinates http://mathworld.wolfram.com/BarycentricCoordinates.html

    areaABC = dot(n, cross(b - a, c - a))

    areaPBC = dot(n, cross(b - p, c - p))
    areaPCA = dot(n, cross(c - p, a - p))
    areaPAB = dot(n, cross(a - p, b - p))

    u = areaPBC / areaABC
    v = areaPCA / areaABC
    w = areaPAB / areaABC

    return (u + v + w) ≈ 1.0 && u > 0 && v > 0 && w > 0 && u < 1 && v < 1 && w < 1
end


@testset "N-Body" begin
    @testset "Libration Points" begin
        l1ref = 0.15093428265127021
        l2ref = 0.16783274367911588
        l3ref = 0.99291206108848129
        @test libration_dist(μ(EARTH), μ(MOON), :L1) == l1ref
        @test libration_dist(μ(EARTH), μ(MOON), :L2) == l2ref
        @test libration_dist(μ(EARTH), μ(MOON), :L3) == l3ref

        rv = ones(6)
        rv[1:3] *= 1000
        epoch = TDBEpoch(2000,1,1)
        ref = [-1.1717170062386328, -1.1212220608364672E-003, 1.6674369835388907E-003,
               -1.6420893636515004, -0.45661886467983320, 0.68484198823320419]
        rv_lib = gcrf_to_libration_norm(rv, epoch, EARTH, MOON, :L2)
        @test rv_lib ≈ ref
        rv_eci = libration_norm_to_gcrf(rv_lib, epoch, EARTH, MOON, :L2)
        @test rv_eci ≈ rv
    end
end


function detrend!(x::Vector, order=0, t = 1:length(x))
    x[:] .-= mean(x)
    if order == 1
        k = x\t
        x[:] .-= k*t
    end
    x
end

"""detrend(x, order=0, t = 1:length(x))
Removes the trend of order `order`, i.e, mean and, if order=1, the slope of the signal `x`
If `order` = 1, then the sampling points of `x` can be supplied as `t` (default: `t = 1:length(x)`)
"""
function detrend(x,args...)
    y = copy(x)
    detrend!(y,args...)
end

"""basis_activation_func(V,Nv,normalize,coulomb)

Returns a func v->ϕ(v) ∈ ℜ(Nv) that calculates the activation of `Nv` basis functions spread out to cover V nicely. If coulomb is true, then we get twice the number of basis functions, 2Nv
"""
function basis_activation_func(V,Nv,normalize,coulomb)
    if coulomb # If Coulomb setting is activated, double the number of basis functions and clip the activation at zero velocity (useful for data that exhibits a discontinuity at v=0, like coulomb friction)
        vc      = range(0, stop=maximum(abs.(V)), length=Nv+2)
        vc      = vc[2:end-1]
        vc      = [-vc[end:-1:1]; vc]
        Nv      = 2Nv
        gamma   = Nv/(abs(vc[1]-vc[end]))
        K       = normalize ? V -> _Kcoulomb_norm(V,vc,gamma) : V -> _Kcoulomb(V,vc,gamma) # Use coulomb basis function instead
    else
        vc      = range(minimum(V), stop=maximum(V), length=Nv)
        gamma   = Nv/(abs(vc[1]-vc[end]))
        K       = normalize ? V -> _K_norm(V,vc,gamma) : V -> _K(V,vc,gamma)
    end
end

"""`ridgereg(A,b,λ)`
Accepts `λ` to solve the ridge regression problem using the formulation `[A;λI]\\[b;0]. λ should be given with the same dimension as the columns of A, i.e. if λ represents an inverse standard deviation, then 1/λ = σ, not 1/λ = σ²`"""
function ridgereg(A,b,λ)
    n = size(A,2)
    [A; λ*I]\[b;zeros(n)]
end

"""`real_complex_bs(A,b, λ=0)`
Replaces the backslash operator For complex arguments. Expands the A-matrix into `[real(A) imag(A)]` and performs the computation using real arithmetics. Optionally accepts `λ` to solve the ridge regression problem using the formulation `[A;λI]\\[b;0]. λ should be given with the same dimension as the columns of A, i.e. if λ represents a standard deviation, then λ = σ, not λ = σ²`
"""
function real_complex_bs(A,b, λ=0)
    n  = size(A,2)
    Ar = [real(A) imag(A)]
    xr = λ > 0 ? [Ar; λ*I]\[b;zeros(2n)] : Ar\b
    x  = complex.(xr[1:n], xr[n+1:end])
end

function fourier_solve(A,y,zerofreq,λ=0)
    n  = size(A,2)
    x = λ > 0 ? [A; λ*I]\[y;zeros(n)] : A\y
    fourier2complex(x,zerofreq)
end

function fourier2complex(x,zerofreq)
    n = length(x)÷2
    if zerofreq === nothing
        return complex.(x[1:n], x[n+1:end])
    else
        x0 = x[zerofreq]
        x = deleteat!(copy(x), zerofreq)
        x = complex.(x[1:n], x[n+1:end])
        insert!(x,zerofreq,x0)
        return x
    end
end

""" Returns params as a [nω × N] matrix"""
reshape_params(x,Nf) = reshape(x, Nf,round(Int,length(x)/Nf))

## Complex Normal
import Base.rand

mutable struct ComplexNormal{T<:Complex}
    m::AbstractVector{T}
    Γ::Cholesky
    C::Symmetric{T}
end

function ComplexNormal(X::AbstractVecOrMat,Y::AbstractVecOrMat)
    @assert size(X) == size(Y)
    mc  = complex.(mean(X,dims=1)[:], mean(Y,dims=1)[:])
    V   = Symmetric(cov([X Y]))
    Γ,C = cn_V2ΓC(V)
    ComplexNormal(mc,Γ,C)
end

function ComplexNormal(X::AbstractVecOrMat{T}) where T<:Complex
    ComplexNormal(real.(X),imag.(X))
end

function ComplexNormal(m::AbstractVector{T},V::AbstractMatrix{T}) where T<:Real
    n   = Int(length(m)/2)
    mc  = complex.(m[1:n], m[n+1:end])
    Γ,C = cn_V2ΓC(V)
    ComplexNormal(mc,Γ,C)
end

function ComplexNormal(mc::AbstractVector{Tc},V::AbstractMatrix{Tr}) where {Tr<:Real, Tc<:Complex}
    Γ,C = cn_V2ΓC(V)
    ComplexNormal(mc,Γ,C)
end

function cn_V2ΓC(V::Symmetric{T}) where T<:Real
    n   = size(V,1)÷2
    Vxx = V[1:n,1:n]
    Vyy = V[n+1:end,n+1:end]
    Vxy = V[1:n,n+1:end]
    Vyx = V[n+1:end,1:n]
    Γ   = cholesky(complex.(Vxx + Vyy, Vyx - Vxy))
    C   = Symmetric(complex.(Vxx - Vyy, Vyx + Vxy))
    Γ,C
end

cn_V2ΓC(V::AbstractMatrix{T}) where {T<:Real} = cn_V2ΓC(Symmetric(V))

@inline cn_Vxx(Γ,C) = cholesky(real.(Matrix(Γ)+C)/2)
@inline cn_Vyy(Γ,C) = cholesky(real.(Matrix(Γ)-C)/2)
@inline cn_Vxy(Γ,C) = cholesky(imag.(-Matrix(Γ)+C)/2)
@inline cn_Vyx(Γ,C) = cholesky(imag.(Matrix(Γ)+C)/2)

@inline cn_fVxx(Γ,C) = real.(Matrix(Γ)+C)/2
@inline cn_fVyy(Γ,C) = real.(Matrix(Γ)-C)/2
@inline cn_fVxy(Γ,C) = imag.(-Matrix(Γ)+C)/2
@inline cn_fVyx(Γ,C) = imag.(Matrix(Γ)+C)/2

@inline cn_Vs(Γ,C) = cn_Vxx(Γ,C),cn_Vyy(Γ,C),cn_Vxy(Γ,C),cn_Vyx(Γ,C)
@inline cn_fV(Γ,C) = [cn_fVxx(Γ,C) cn_fVxy(Γ,C); cn_fVyx(Γ,C) cn_fVyy(Γ,C)]
@inline cn_V(Γ,C) = cholesky(cn_fV(Γ,C))
@inline Σ(cn::ComplexNormal) = Matrix(cn.Γ) # TODO: check this

for f in [:cn_Vxx,:cn_Vyy,:cn_Vxy,:cn_Vyx,:cn_fVxx,:cn_fVyy,:cn_fVxy,:cn_fVyx,:cn_Vs,:cn_V,:cn_fV]
    @eval ($f)(cn::ComplexNormal) = ($f)(cn.Γ,cn.C)
end

"""
`f(cn::ComplexNormal, z)`

Probability density Function `f(z)` for a complex normal distribution.
This can probably be more efficiently implemented
"""
function pdf(cn::ComplexNormal, z)
    k = length(cn.m)
    R = conj(cn.C)'*inv(cn.Γ)
    P = Matrix(cn.Γ)-R*cn.C # conj(Γ) = Γ for Γ::Cholesky
    cm = conj(cn.m)
    cz = conj(z)
    zmm = z-cn.m
    czmm = cz-cm
    ld = [czmm' zmm']
    rd = [zmm; czmm]
    S = [Matrix(cn.Γ) cn.C;conj(cn.C) Matrix(cn.Γ)] # conj(Γ) = Γ for Γ::Cholesky
    1/(π^k*sqrt(det(cn.Γ)*det(P))) * exp(-0.5* ld*(S\rd))
end

affine_transform(cn::ComplexNormal, A,b) = ComplexNormal(A*cn.m+b, cholesky(A*Matrix(cn.Γ)*conj(A')), Symmetric(A*cn.C*A'))

function rand(cn::ComplexNormal,s::Integer)
    L = cn_V(cn).U
    m = [real(cn.m); imag(cn.m)]
    n = length(cn.m)
    z = (m' .+ randn(s,2n)*L)
    return complex.(z[:,1:n],z[:,n+1:end])
end


function part1(input)
    expenses = falses(2021)
    for expense_string in readlines(input)
        expense = parse(Int, expense_string)
        if expenses[1 + 2020 - expense]
            return expense * (2020 - expense)
        end
        expenses[1 + expense] = true
    end
    error("No solution found.")
end

function part2(input)
    expenses = Int[]
    two_expenses = Dict{Int, Int}()
    for expense_string in readlines(input)
        expense = parse(Int, expense_string)
        if haskey(two_expenses, 2020 - expense)
            return two_expenses[2020 - expense] * expense
        end
        for expense2 in expenses
            two_expenses[expense + expense2] = expense * expense2
        end
        push!(expenses, expense)
    end
    error("No solution found.")
end


using SymPy: N, Sym, diff, expand, solve, subs

"""
Axisymmetric Solov'ev equilibra in (R/R₀,Z/R₀,phi) coordinates.
Based on Cerfon & Freidberg, Physics of Plasmas 17, 032502, 2010.

Parameters:
 * `R₀`: position of magnetic axis
 * `B₀`: B-field at magnetic axis
 * `ϵ`:  inverse aspect ratio
 * `κ`:  elongation
 * `δ`:  triangularity
 * `a`:  free constant, determined to match a given beta value
"""
struct Solovev{T <: Number} <: AbstractSolovevEquilibrium
    name::String
    R₀::T
    B₀::T
    ϵ::T
    κ::T
    δ::T
    a::T
    c::Vector{T}

    function Solovev{T}(R₀::T, B₀::T, ϵ::T, κ::T, δ::T, a::T, c::Vector{T}) where T <: Number
        new("SolovevEquilibrium", R₀, B₀, ϵ, κ, δ, a, c)
    end
end

function Solovev(R₀::T, B₀::T, ϵ::T, κ::T, δ::T, a::T) where T <: Number

    n = 7

    x = [Sym("x" * string(i)) for i in 1:3]
    c = [Sym("c" * string(i)) for i in 1:n]

    ψ = (ψ₀(x,a) + c[1] * ψ₁(x)
                 + c[2] * ψ₂(x)
                 + c[3] * ψ₃(x)
                 + c[4] * ψ₄(x)
                 + c[5] * ψ₅(x)
                 + c[6] * ψ₆(x)
                 + c[7] * ψ₇(x) )

    eqs = [
        expand(subs(subs(ψ, x[1], 1+ϵ), x[2], 0)),
        expand(subs(subs(ψ, x[1], 1-ϵ), x[2], 0)),
        expand(subs(subs(ψ, x[1], 1-δ*ϵ), x[2], κ*ϵ)),
        expand(subs(subs(diff(ψ, x[1]), x[1], 1-δ*ϵ), x[2], κ*ϵ)),
        expand(subs(subs(diff(ψ, x[2], 2), x[1], 1+ϵ), x[2], 0) - (1 + asin(δ))^2 / (ϵ * κ^2) * subs(subs(diff(ψ, x[1]), x[1], 1+ϵ), x[2], 0)),
        expand(subs(subs(diff(ψ, x[2], 2), x[1], 1-ϵ), x[2], 0) + (1 - asin(δ))^2 / (ϵ * κ^2) * subs(subs(diff(ψ, x[1]), x[1], 1-ϵ), x[2], 0)),
        expand(subs(subs(diff(ψ, x[1], 2), x[1], 1-δ*ϵ), x[2], κ*ϵ) - κ / (ϵ * (1 - δ^2)) * subs(subs(diff(ψ, x[2]), x[1], 1-δ*ϵ), x[2], κ*ϵ))
    ]

    csym = solve(eqs, c)
    cnum = [N(csym[c[i]]) for i in 1:n]

    Solovev{T}(R₀, B₀, ϵ, κ, δ, a, cnum)
end

SolovevITER() = Solovev(6.2, 5.3, 0.32, 1.7, 0.33, -0.155)
SolovevNSTX() = Solovev(0.85, 0.3, 0.78, 2.0, 0.35, 1.0)

function Base.show(io::IO, equ::Solovev)
    print(io, "Solovev Equilibrium with\n")
    print(io, "  R₀ = ", equ.R₀, "\n")
    print(io, "  B₀ = ", equ.B₀, "\n")
    print(io, "  ϵ  = ", equ.ϵ, "\n")
    print(io, "  κ  = ", equ.κ, "\n")
    print(io, "  δ  = ", equ.δ, "\n")
    print(io, "  a  = ", equ.a)
end

function A₃(x::AbstractArray{T,1}, equ::Solovev) where {T <: Number}
    (ψ₀(x, equ.a) + equ.c[1] * ψ₁(x)
                  + equ.c[2] * ψ₂(x)
                  + equ.c[3] * ψ₃(x)
                  + equ.c[4] * ψ₄(x)
                  + equ.c[5] * ψ₅(x)
                  + equ.c[6] * ψ₆(x)
                  + equ.c[7] * ψ₇(x) )
end

macro solovev_equilibrium(R₀, B₀, ϵ, κ, δ, a)
    generate_equilibrium_code(Solovev(R₀, B₀, ϵ, κ, δ, a); output=false)
end


module Branch_Test

using Compat
using Compat.Test
using IntervalArithmetic
using EAGOBranchBound

S = BnBSolver()
B = BnBModel([Interval(1.0,2.0)])
tL = 3.0
tU = 5.0
X1 = [Interval(1.0,2.0)]
X2 = [Interval(3.0,5.0)]
X = [Interval(3.0,5.0)]
pos = 4
EAGOBranchBound.BM_breadth!(S,B,tL,tU,X1,X2,pos)
@test (B.pos[1:3] == [5; 5; 1])
@test (B.LBD[1:3] == [3.0; 3.0; -Inf])
@test (B.UBD[1:3] == [5.0; 5.0; Inf])
@test (B.id[1:3] == [2; 3; 1])

EAGOBranchBound.BM_Single!(S,B,tL,tU,X,pos)
@test (B.pos[4] == 4)
@test (B.LBD[4] == 3.0)
@test (B.UBD[4] == 5.0)
@test (B.id[4] == 3)

end


@testset "mesh from field" begin
    srand(123)
    df = gaussian_displacement(3, 3, .3, 2)
    dm = DisplacementMesh(df, testimg)
    show(dm); println()
    showcompact(dm); println()
    @test typeof(dm) <: DisplacementMesh
    @test typeof(dm) <: ImageTransformation
    dm2 = DisplacementMesh(df, size(testimg, "x"), size(testimg, "y"))
    @test dm.input_vertices == dm2.input_vertices
    @test dm.output_vertices == dm2.output_vertices
    @test dm.indices == dm2.indices
    @test size(dm.input_vertices, 1) == 9
    @test size(dm.output_vertices, 1) == 9
    @test size(dm.indices, 1) == 8
    @test dm.indices ==
        [ 1  5  2;
          1  4  5;
          2  5  3;
          3  5  6;
          4  7  5;
          5  7  8;
          5  9  6;
          5  8  9 ]

    srand(123)
    df = gaussian_displacement(4, 8, .3, 2)
    dm = DisplacementMesh(df, testimg)
    # @plottest "displacemen_mesh_1" plot(dm)
    # @plottest "displacemen_mesh_2" plot(dm, testimg)
end

@testset "transform" begin
    srand(123)
    df = gaussian_displacement(3, 3, .3, 2)
    dm = DisplacementMesh(df, testimg)

    @imagetest "DisplacementMesh" transform(dm, testimg)
end


using Memoize
@memoize josephus(n::Integer, k::Integer, m::Integer=1) = n == m ? collect(0:m .- 1) : mod.(josephus(n - 1, k, m) + k, n)

@show josephus(41, 3)
@show josephus(41, 3, 5)


function score(x, y)
    distance = (x^2 + y^2)^0.5
    if distance > 10
        return 0
    elseif distance > 5
        return 1
    elseif distance > 1
        return 5
    end

return 10

end


using jInvSeismic.FWI
using jInvSeismic.Utils
using Helmholtz
using jInv.Mesh
using jInv.Utils
using jInv.ForwardShare
using Multigrid.ParallelJuliaSolver
using Distributed

using DelimitedFiles

if nworkers() == 1
	addprocs(16);
elseif nworkers() < 16
	addprocs(16 - nworkers());
end

@everywhere begin
using jInv.InverseSolve
using jInv.LinearSolvers
using jInvSeismic.FWI
using jInv.Mesh
using jInv.Utils
using DelimitedFiles
using jInv.ForwardShare
using KrylovMethods
end

dim     = 2;
pad     = 30;
newSize = [600,300];

(m,Mr,mref,boundsHigh,boundsLow) = readModelAndGenerateMeshMref("examples","SEGmodel2Dsalt.dat",dim,pad,[0.0,13.5,0.0,4.2],newSize,1.752,2.7);
m = 1 ./ (m.^2);
mref = 1 ./ (mref.^2);

# attenuation for BC
padx = 34; padz = 34
a    = 2.0;
xc = getCellCenteredGrid(Mr);
gamma = getHelmholtzABL(Mr,true,[padx;padz],a);
# gamma = Helmholtz.getABL(Mr,true,[padx;padz],a);
println("Size gamma: ", size(gamma))
println("Size Mr: " , Mr.n)
println("Size m: " , size(m))
# parameters for the Helmholtz (units in km)
h = Mr.h;
n = Mr.n;
# omega = 0.3*2*pi*[2.0;5.0; 10.0;20.0;25.0]
omega = 2*pi*[2.0;2.5;3.5;4.5;6.0;]
nfreq = length(omega)
# generate sources
q = zeros(tuple(n.+1...)); q[padx+1:4:end-padx-1,1] .= 1e4
print(size(q))
Q = sdiag(vec(q))
print(size(Q))
Q = Q[:,(LinearIndices(sum(Q,dims=2) .!= 0))[findall(sum(Q,dims=2) .!= 0)]]
print(size(Q))
nsrc = size(Q,2)
# receivers
p = zeros(tuple(n.+1...)); p[padx+1:end-padx-1,1] .= 1
P = sdiag(vec(p))
P = P[:,(LinearIndices(sum(P,dims=2) .!= 0))[findall(sum(P,dims=2) .!= 0)]]
nrec = size(P,2)
Ainv  = getParallelJuliaSolver(ComplexF64,Int64,numCores=4,backend=1);

(pForp,contDivFWI,SourcesSubIndFWI)  = getFWIparam(omega,ones(ComplexF64,size(omega)),vec(gamma),Q,P,Mr,Ainv,workers())
# pForp = getBasicFWIparam(omega,gamma,Q,P,Mr,true)

# P = nothing
# Q = nothing
# inversion mesh and forward mesh are the same here
M2Mp = ones(length(pForp))


# import Particles package
using Particles

# parameters
Δt = 0.1    # time step size
nt = 200    # number of time steps
np = 10000  # number of particles
nx = 16     # number of grid points

# random initial conditions
x₀ = randn(np)
v₀ = randn(np)

# shift x₀ to the interval [0,1]
xmax = ceil(maximum(abs.(x₀)))
x₀ .+= xmax
x₀ ./= 2*xmax

# concatenate x₀ and v₀
z₀ = [ collect(z) for z in zip(x₀, v₀) ]

# vector field
function lorentz_force!(ż, z, p::PoissonSolver)
    ż[1] = z[2]
    ż[2] = eval_field(p, z[1])
end

# create Poisson solver
p = PoissonSolver{eltype(x₀)}(nx)
solve!(p, x₀)

# create an Equation instance with callbacks
equ = Equation((ż, z) -> lorentz_force!(ż, z, p), z₀; f_preproc = z -> solve!(p, reduce(hcat, z)[1,:]))

# create Simulation instance
sim = Simulation(equ, ExplicitEuler, Δt, nt)

# run simulation
z = sim()

# load Plots package
using Plots

# compute plot ranges
vmax = ceil(maximum(abs.(v₀)))
xlim = (0, 1)
vlim = (-vmax, +vmax)

# plot initial condition
z₀plot = reduce(hcat, z₀)
scatter(mod.(z₀plot[1,:], 1), z₀plot[2,:],
        marker = 3,
        xlim = xlim,
        ylim = vlim,
        title = "Vlasov-Poisson",
        legend = false,
        size = (800, 600)
)

# save figure to file
savefig("vlasov_poisson_z₀.png")

# create animation
# anim = @animate for n in 0:nt
#     scatter(mod.(z[1,:,n], 1), z[2,:,n],
#         marker = 3,
#         xlim = xlim,
#         ylim = vlim,
#         title = "Vlasov-Poisson",
#         legend = false,
#         size = (800, 600)
#     )
# end

# save animation to file
# gif(anim, "vlasov_poisson_anim.gif", fps=10)


# https://adventofcode.com/2020/day/3

# You start on the open square (.) in the top-left corner and need to reach the
# bottom (below the bottom-most row on your map).

# The toboggan can only follow a few specific slopes (you opted for a cheaper model
# that prefers rational numbers); start by counting all the trees you would encounter
# for the slope right 3, down 1:

# From your starting position at the top-left, check the position that is right 3
# and down 1. Then, check the position that is right 3 and down 1 from there, and
# so on until you go past the bottom of the map.

function part1(patch, right, down=1)
    trees_count = 0
    position = 1
    time_to_move = 1

    for row in split(patch, '\n', keepempty=false)
        if time_to_move == 1
            if row[position] == '#'
                trees_count += 1
            end
            position = mod1(position + right, length(row))
            time_to_move = down
        else
            time_to_move -= 1
        end
    end
    return trees_count
end

example = "
..##.......
#...#...#..
.#....#..#.
..#.#...#.#
.#...##..#.
..#.##.....
.#.#.#....#
.#........#
#.##...#...
#...##....#
.#..#...#.#
"

@assert part1(example, 3) == 7

# Determine the number of trees you would encounter if, for each of the following
# slopes, you start at the top-left corner and traverse the map all the way
# to the bottom:
#
#  Right 1, down 1.
#  Right 3, down 1. (This is the slope you already checked.)
#  Right 5, down 1.
#  Right 7, down 1.
#  Right 1, down 2.
#
# In the above example, these slopes would find 2, 7, 3, 4, and 2 tree(s) respectively;
# multiplied together, these produce the answer 336.

function part2(patch, right, down)
    result = 1
    for (r, d) in zip(right, down)
        result *= part1(patch, r, d)
    end
    return result
end

@assert part2(example, [1 3 5 7 1], [1 1 1 1 2]) == 336

test = read("data/day-03.txt", String)
println("Part 1: $(result1 = part1(test, 3))")
println("Part 2: $(result2 = part2(test, [1 3 5 7 1], [1 1 1 1 2]))")

@assert result1 == 203
@assert result2 == 3316272960


function mm_spec_kyparisis3_dimspec()::MPCCDimSpec
    dimspec = MPCCDimSpec(2, 0, 0, 0, 3, 2, 0)
    return dimspec
end

function mm_spec_kyparisis3_defn(x, pr, ps)

    f = (1/2)*(x[1]-pr[1])^2 + (1/2)*(x[2]-pr[2])^2 + x[1] + x[2]

    ce = Vector{Num}()

    ci = Vector{Num}(undef, 3)
    ci[1] = x[1]
    ci[2] = x[2]
    ci[3] = x[1] + x[2] - pr[1] - pr[2]

    F = Matrix{Num}(undef, 0, 0)
    # F[1,1] = x[1]
    # F[2,1] = x[2]

    return MPCCDefinition(f, ce, ci, F)
end

function mm_spec_kyparisis3_nzmask()::MPCCModelNZMask
    nz_mask = MPCCModelNZMask(
        # ce
        [],
        # ci
        [   Set{Int64}([1, 2]),
            Set{Int64}([1]) ],
        # F
        [
            # F1
            [   Set{Int64}([1]) ],
            # F2
            [   Set{Int64}([2]) ]
        ] )
    return nz_mask
end

function mm_spec_kyparisis3_testvectors()
    return Vector{MPCCModelTestVector}(undef, 0)
end

function mm_spec_kyparisis3_knownsols()

    function local_kyparisis3_knownsol(pr::Vector{R}, ps::Vector{Int64}) where {R <: Real}
        @assert ( 2 == length(pr) && 0 == length(ps) )
        e1 = pr[1]
        e2 = pr[2]

        if ( e1 >= 0 && e2 >= 0 )
            x_sol = (e1, e2)
        elseif ( e1 >= 0 && e2 < 0 )
            x_sol = (max(e1-1, e1 + e2, 0), 0)
        elseif ( e1 < 0 && e2 >= 0 )
            x_sol = (0, max(e2-1, e1 + e2, 0))
        else
            x_sol = (0, 0)
        end
        f_sol = (1/2)*(x_sol[1]-e1)^2 + (1/2)*(x_sol[2]-e2)^2 + x_sol[1] + x_sol[2]

        return (x_sol, f_sol)
    end

    return [ local_kyparisis3_knownsol ]
end

function mm_spec_kyparisis3_parameterisations(t)

    defn1 = MPCCParameterisationDefn(
            Vector{Num}([ 3.0, t ]),
            (-3.0, 3.0),
            "Vertical, rising"
        )

    return Vector{MPCCParameterisationDefn}( [ defn1 ] )
end


using Printf

a = readlines()

ROWS = lastindex(a)
COLS = lastindex(a[begin])

digits = []
for row in 1:ROWS
    for col in 1:COLS
        append!(digits, parse(Int, a[row][col]))
    end
end

basin = []

function contains(what, lst)
    # @printf("contains %d in %s\n", what, lst)
    what in Set(lst)
end

function candidates(i, previous)
    out = []

    col = (i - 1) % COLS
    row = div((i - 1), COLS)

    if col > 0
        append!(out, i - 1)
    end
    if col < (COLS-1)
        append!(out, i + 1)
    end
    if row > 0
        append!(out, i - COLS)
    end
    if row < (ROWS-1)
        append!(out, i + COLS)
    end
    ret = []
    for x in out
        if ! (x in previous)
            append!(ret, x)
        end
    end
    ret
end

for i = 1:lastindex(digits)
    digit = digits[i]
    cands = candidates(i, [])
    low = true
    for c in cands
        if digit >= digits[c]
            low = false
        end
    end
    if low
        append!(basin, i)
    end
end

total = 0
for x in basin
    global total += digits[x] + 1
end
println(total)

basin_sizes = []
additional = []
basin = Set(basin)
seen = Set()

lengths = []

for start in basin
    b = Set([start])
    q = []
    for c in candidates(start, [])
        append!(q, c)
    end
    while 0 < length(q)
        c = popfirst!(q)
        if digits[c] < 9
            ns = candidates(c, b)
            good = true
            for n in ns
                if digits[n] < digits[c]
                    good = false
                end
            end
            if good
                push!(b, c)
                for n in ns
                    append!(q, n)
                end
            end
        end
    end
    append!(lengths, length(b))
end

sorted_lengths = sort(lengths)
# println(sorted_lengths)

println(sorted_lengths[end] * sorted_lengths[end-1] * sorted_lengths[end-2])


export differentiate

"""
    differentiate(p::AbstractPolynomialLike, v::AbstractVariable, deg::Union{Int, Val}=1)

Differentiate `deg` times the polynomial `p` by the variable `v`.

    differentiate(p::AbstractPolynomialLike, vs, deg::Union{Int, Val}=1)

Differentiate `deg` times the polynomial `p` by the variables of the vector or
tuple of variable `vs` and return an array of dimension `deg`. It is recommended
to pass `deg` as a `Val` instance when the degree is known at compile time, e.g.
`differentiate(p, v, Val{2}())` instead of `differentiate(p, x, 2)`, as this
will help the compiler infer the return type.

    differentiate(p::AbstractArray{<:AbstractPolynomialLike, N}, vs, deg::Union{Int, Val}=1) where N

Differentiate the polynomials in `p` by the variables of the vector or tuple of variable `vs` and return an array of dimension `N+deg`.
If `p` is an `AbstractVector` this returns the Jacobian of `p` where the i-th row containts the partial
derivaties of `p[i]`.

### Examples

```julia
p = 3x^2*y + x + 2y + 1
differentiate(p, x) # should return 6xy + 1
differentiate(p, x, Val{1}()) # equivalent to the above
differentiate(p, (x, y)) # should return [6xy+1, 3x^2+1]
differentiate( [x^2+y, z^2+4x], [x, y, z]) # should return [2x 1 0; 4 0 2z]
```
"""
function differentiate end

# Fallback for everything else
differentiate(α::T, v::AbstractVariable) where T = zero(T)
differentiate(v1::AbstractVariable, v2::AbstractVariable) = v1 == v2 ? 1 : 0
differentiate(t::AbstractTermLike, v::AbstractVariable) = coefficient(t) * differentiate(monomial(t), v)
# The polynomial function will take care of removing the zeros
differentiate(p::APL, v::AbstractVariable) = polynomial(differentiate.(terms(p), v), SortedState())
differentiate(p::RationalPoly, v::AbstractVariable) = (differentiate(p.num, v) * p.den - p.num * differentiate(p.den, v)) / p.den^2

const ARPL = Union{APL, RationalPoly}

function differentiate(ps::AbstractArray{PT}, xs::AbstractArray) where {PT <: ARPL}
    differentiate.(reshape(ps, (size(ps)..., 1)), reshape(xs, 1, :))
end

function differentiate(ps::AbstractArray{PT}, xs::Tuple) where {PT <: ARPL}
    differentiate(ps, collect(xs))
end

# TODO: this signature is probably too wide and creates the potential
# for stack overflows
differentiate(p::ARPL, xs) = [differentiate(p, x) for x in xs]

# differentiate(p, [x, y]) with TypedPolynomials promote x to a Monomial
differentiate(p::ARPL, m::AbstractMonomial) = differentiate(p, variable(m))

# The `R` argument indicates a desired result type. We use this in order
# to attempt to preserve type-stability even though the value of `deg` cannot
# be known at compile time. For scalar `p` and `x`, we set R to be the type
# of differentiate(p, x) to give a stable result type regardless of `deg`. For
# vectors p and/or x this is impossible (since differentiate may return an array),
# so we just set `R` to `Any`
function (_differentiate_recursive(p, x, deg::Int, ::Type{R})::R) where {R}
    if deg < 0
        throw(DomainError(deg, "degree is negative"))
    elseif deg == 0
        return p
    else
        return differentiate(differentiate(p, x), x, deg-1)
    end
end

differentiate(p, x, deg::Int) = _differentiate_recursive(p, x, deg, Base.promote_op(differentiate, typeof(p), typeof(x)))
differentiate(p::AbstractArray, x,                              deg::Int) = _differentiate_recursive(p, x, deg, Any)
differentiate(p,                x::Union{AbstractArray, Tuple}, deg::Int) = _differentiate_recursive(p, x, deg, Any)
differentiate(p::AbstractArray, x::Union{AbstractArray, Tuple}, deg::Int) = _differentiate_recursive(p, x, deg, Any)

# This is alternative, Val-based interface for nested differentiation.
# It has the advantage of not requiring an conversion or calls to
# Base.promote_op, while maintaining type stability for any argument
# type.
differentiate(p, x, ::Val{0}) = p
differentiate(p, x, ::Val{1}) = differentiate(p, x)

@static if VERSION < v"v0.7.0-"
    # Marking this @pure helps julia v0.6 figure this out
    Base.@pure _reduce_degree(::Val{N}) where {N} = Val{N - 1}()
    function differentiate(p, x, deg::Val{N}) where N
        if N < 0
            throw(DomainError(deg))
        else
            differentiate(differentiate(p, x), x, _reduce_degree(deg))
        end
    end
else
    # In Julia v0.7 and above, we can remove the _reduce_degree trick
    function differentiate(p, x, deg::Val{N}) where N
        if N < 0
            throw(DomainError(deg))
        else
            differentiate(differentiate(p, x), x, Val{N - 1}())
        end
    end
end