abstract type MajFlowNorm end

struct OneNorm <: MajFlowNorm end

struct InfNorm <: MajFlowNorm end

export SortedProbVecMult

struct SortedProbVecMult{T, VT <: AbstractVector{T}} <: AbstractVector{T}
    distinct_entries::VT
    multiplicities::Vector{Int}
end

Base.size(r::SortedProbVecMult) = (sum(r.multiplicities),)
Base.eltype(r::SortedProbVecMult{T}) where {T} = T

Base.@propagate_inbounds function Base.getindex(r::SortedProbVecMult, i::Int)
    @boundscheck (1 <= i <= length(r)) || throw(BoundsError(r,i))
    ks = r.multiplicities
    seen_so_far = 0
    for j = eachindex(ks)
        seen_so_far += ks[j]
        if seen_so_far >= i
            return r.distinct_entries[j]
        end
    end
end

function SortedProbVecMult(v::AbstractVector)
    v_sorted = sort(v; rev = true)
    distinct_entries = eltype(v)[]
    multiplicities = Int[]
    for elt in v_sorted
        if isempty(distinct_entries) || distinct_entries[end] != elt
            push!(distinct_entries, elt)
            push!(multiplicities, 1)
        else
            multiplicities[end] += 1
        end

    end
    SortedProbVecMult{eltype(distinct_entries), typeof(distinct_entries)}(distinct_entries, multiplicities)
end

function Base.collect(spvm::SortedProbVecMult)
    @unpack distinct_entries, multiplicities = spvm
    v = eltype(distinct_entries)[]
    sizehint!(v, length(spvm))
    for (i, elt) in enumerate(distinct_entries)
        for j = 1:multiplicities[i]
            push!(v, elt)
        end
    end
    return v
end

# Default
majmin(v::SortedProbVecMult, ϵ) = majmin(OneNorm(), v, ϵ)

function majmin(nrm::MajFlowNorm, spvm::SortedProbVecMult, ϵ)
    majmin!(nrm, deepcopy(spvm), ϵ)
end

function majmin!(nrm::OneNorm, spvm::SortedProbVecMult, ϵ)
    @unpack distinct_entries, multiplicities = spvm
    @assert ϵ >= 0
    if length(distinct_entries) == 1
        return spvm
    elseif length(distinct_entries) == 2
        _majmin_onenorm_2entries!(spvm::SortedProbVecMult, ϵ)
    else
        kp = first(multiplicities)
        km = last(multiplicities)

        δp = kp * (distinct_entries[1] - distinct_entries[2])
        δm = km * (distinct_entries[end-1] - distinct_entries[end])
        gap = min(δp, δm)
        if ϵ < gap
            distinct_entries[1] -= ϵ / kp
            distinct_entries[end] += ϵ / km
            return spvm
        else
            if δp < δm
                # evolve for time δp and recurse
                deleteat!(distinct_entries, 1)
                deleteat!(multiplicities, 1)
                multiplicities[1] += kp
                distinct_entries[end] += δp / km
                return majmin!(nrm, spvm, ϵ - δp)
            elseif δm < δp
                # evolve for time δm and recurse
                pop!(distinct_entries)
                pop!(multiplicities)
                multiplicities[end] += km
                distinct_entries[1] -= δm / kp
                return majmin!(nrm, spvm, ϵ - δm)
            else
                # evolve for time gap=δm=δp and recurse
                deleteat!(distinct_entries, 1)
                deleteat!(multiplicities, 1)
                multiplicities[1] += kp
                pop!(distinct_entries)
                pop!(multiplicities)
                multiplicities[end] += km
                return majmin!(nrm, spvm, ϵ - gap)
            end
        end
    end
end

# With only 2 entries, the next crossing isn't just determined by
# gap = min(δp, δm).
function _majmin_onenorm_2entries!(spvm::SortedProbVecMult, ϵ)
    @unpack distinct_entries, multiplicities = spvm
    @assert length(distinct_entries) == length(multiplicities) == 2
    kp, km = multiplicities
    μp, μm = distinct_entries
    # the larger entry is decreasing at a rate t/kp
    # the smaller entry is increasing at a rate t/km
    # Need to solve   `μp - t/kp == μm + t/km` for `t`
    # t = (μp - μm) / ( 1/km + 1/kp )
    t = (μp - μm) * (km * kp) / (km + kp)
    if ϵ < t
        distinct_entries[1] -= ϵ / kp
        distinct_entries[2] += ϵ / km
        return spvm
    else
        pop!(distinct_entries)
        pop!(multiplicities)
        distinct_entries[] -= t / kp
        multiplicities[] += km
        return spvm
    end
end


#! /usr/bin/julia

# Rosetta Code, Non-continuous subsequences

iscontseq(n::Integer) = count_zeros(n) == leading_zeros(n) + trailing_zeros(n)
iscontseq(n::BigInt) = !ismatch(r"0", rstrip(bin(n), '0'))

function makeint2seq(n::Integer)
    const idex = collect(1:n)
    function int2seq(m::Integer)
        d = digits(m, 2, n)
        idex[d .== 1]
    end
    return int2seq
end

immutable NCSubSeq{T<:Integer}
    n::T
end

type NCSubState{T<:Integer}
    m::T
    m2s::Function
end

Base.length(a::NCSubSeq) = 2^a.n - div(a.n*(a.n+1), 2) - 1
Base.start(a::NCSubSeq) = NCSubState(5, makeint2seq(a.n))
Base.done(a::NCSubSeq, as::NCSubState) = 2^a.n-3 < as.m

function Base.next(a::NCSubSeq, as::NCSubState)
    s = as.m2s(as.m)
    as.m += 1
    while iscontseq(as.m)
        as.m += 1
    end
    return (s, as)
end

n = 4
print("Testing NCSubSeq for ", n, " items:\n   ")
for a in NCSubSeq(n)
    print(" ", a)
end
println()

s = "Rosetta"
cs = split(s, "")
m = 10
n = length(NCSubSeq(length(s))) - m
println()
println("The first and last ", m, " NC sub-sequences of \"", s, "\":")
for (i,a) in enumerate(NCSubSeq(length(cs)))
    i <= m || n < i || continue
    println(@sprintf "%6d %s" i join(cs[a], ""))
    i == m || continue
    println("    .. ......")
end

t = {}
append!(t, collect(1:10))
append!(t, collect(20:10:40))
append!(t, big(50):50:200)
println()
println("Numbers of NC sub-sequences of a given length:")
for i in t
    println(@sprintf("%7d => ", i), length(NCSubSeq(i)))
end


using Test
using Entropies
using DelayEmbeddings
using Wavelets
using StaticArrays
using Neighborhood: KDTree, BruteForce

@testset "Histogram estimation" begin
    x = rand(1:10, 100)
    D = Dataset([rand(1:10, 3) for i = 1:100])
    D2 = [(rand(1:10), rand(1:10, rand(1:10)) for i = 1:100)]
    @test Entropies._non0hist(x) isa Probabilities
    @test Entropies._non0hist(D) isa Probabilities
    @test Entropies._non0hist(D2) isa Probabilities

    @test Entropies._non0hist(x) |> sum ≈ 1.0
    @test Entropies._non0hist(D) |> sum ≈ 1.0
    @test Entropies._non0hist(D2)|> sum ≈ 1.0
    x = rand(100)
    @test genentropy(x, 100) ≠ NaN
end

@testset "Shorthand" begin
    D = Dataset([rand(1:10, 5) for i = 1:100])
    ps, bins = Entropies.binhist(D, 0.2)
    @test Entropies.binhist(D, 0.2) isa Tuple{Probabilities, Vector{<:SVector}}
    @test Entropies.binhist(D, RectangularBinning(0.2)) isa Tuple{Probabilities, Vector{<:SVector}}
    @test Entropies.binhist(D, RectangularBinning(5)) isa Tuple{Probabilities, Vector{<:SVector}}
    @test Entropies.binhist(D, RectangularBinning([5, 3, 4, 2, 2])) isa Tuple{Probabilities, Vector{<:SVector}}
    @test Entropies.binhist(D, RectangularBinning([0.5, 0.3, 0.4, 0.2, 0.2])) isa Tuple{Probabilities, Vector{<:SVector}}
end

@testset "Generalized entropy" begin
    x = rand(1000)
    xn = x ./ sum(x)
    xp = Probabilities(xn)
    @test genentropy(xp, q = 2) isa Real
    @test genentropy(xp, q = 1) isa Real
    @test_throws MethodError genentropy(xn, q = 2) isa Real
end

@testset "Probability/entropy estimators" begin
    @test CountOccurrences() isa CountOccurrences

    @test SymbolicPermutation() isa SymbolicPermutation
    @test SymbolicPermutation(lt = Base.isless) isa SymbolicPermutation
    @test SymbolicPermutation(lt = Entropies.isless_rand) isa SymbolicPermutation
    @test SymbolicWeightedPermutation() isa SymbolicWeightedPermutation
    @test SymbolicWeightedPermutation(lt = Base.isless) isa SymbolicWeightedPermutation
    @test SymbolicWeightedPermutation(lt = Entropies.isless_rand) isa SymbolicWeightedPermutation
    @test SymbolicAmplitudeAwarePermutation() isa SymbolicAmplitudeAwarePermutation
    @test SymbolicAmplitudeAwarePermutation(lt = Base.isless) isa SymbolicAmplitudeAwarePermutation
    @test SymbolicAmplitudeAwarePermutation(lt = Entropies.isless_rand) isa SymbolicAmplitudeAwarePermutation

    @test VisitationFrequency(RectangularBinning(3)) isa VisitationFrequency
    @test TransferOperator(RectangularBinning(3)) isa TransferOperator
    @test TimeScaleMODWT() isa TimeScaleMODWT
    @test TimeScaleMODWT(Wavelets.WT.Daubechies{8}()) isa TimeScaleMODWT
    @test Kraskov(k = 2, w = 1) isa Kraskov
    @test Kraskov() isa Kraskov
    @test KozachenkoLeonenko() isa KozachenkoLeonenko
    @test KozachenkoLeonenko(w = 5) isa KozachenkoLeonenko
    @test NaiveKernel(0.1) isa NaiveKernel

    @testset "Counting based" begin
        D = Dataset(rand(1:3, 1000, 3))
        ts = [(rand(1:4), rand(1:4), rand(1:4)) for i = 1:3000]
        @test Entropies.genentropy(D, CountOccurrences(), q = 2, base = 2) isa Real
    end

    @testset "NaiveKernel" begin
        N = 1000
        pts = Dataset([rand(2) for i = 1:N]);
        ϵ = 0.3
        est_direct = NaiveKernel(ϵ, KDTree)
        est_tree = NaiveKernel(ϵ, BruteForce)

        @test probabilities(pts, est_tree) isa Probabilities
        @test probabilities(pts, est_direct) isa Probabilities
        p_tree = probabilities(pts, est_tree)
        p_direct = probabilities(pts, est_direct)
        @test all(p_tree .== p_direct) == true

        @test Entropies.genentropy(pts, est_direct, base = 2) isa Real
        @test Entropies.genentropy(pts, est_tree, base = 2) isa Real
    end

    @testset "Permutation entropy" begin

        @testset "Encoding and symbolization" begin
            @test Entropies.encode_motif([2, 3, 1]) isa Int
            @test 0 <= Entropies.encode_motif([2, 3, 1]) <= factorial(3) - 1

            est = SymbolicPermutation(m = 5, τ = 1)
            N = 100
            x = Dataset(repeat([1.1 2.2 3.3], N))
            y = Dataset(rand(N, 5))
            z = rand(N)

            # Without pre-allocation
            D = genembed(z, [0, -1, -2])
            est = SymbolicPermutation(m = 5, τ = 2)

            @test Entropies.symbolize(z, est) isa Vector{<:Int}
            @test Entropies.symbolize(D, est) isa Vector{<:Int}

            # With pre-allocation
            N = 100
            x = rand(N)
            est = SymbolicPermutation(m = 5, τ = 2)
            s = fill(-1, N-(est.m-1)*est.τ)

            # if symbolization has occurred, s must have been filled with integers in
            # the range 0:(m!-1)
            @test all(Entropies.symbolize!(s, x, est) .>= 0)
            @test all(0 .<= Entropies.symbolize!(s, x, est) .< factorial(est.m))

            m = 4
            D = Dataset(rand(N, m))
            s = fill(-1, length(D))
            @test all(0 .<= Entropies.symbolize!(s, D, est) .< factorial(m))
        end

        @testset "Pre-allocated" begin
            est = SymbolicPermutation(m = 5, τ = 1)
            N = 500
            s = zeros(Int, N);
            x = Dataset(repeat([1.1 2.2 3.3], N))
            y = Dataset(rand(N, 5))
            z = rand(N)

            # Probability distributions
            p1 = probabilities!(s, x, est)
            p2 = probabilities!(s, y, est)
            @test sum(p1) ≈ 1.0
            @test sum(p2) ≈ 1.0

            # Entropies
            @test genentropy!(s, x, est, q = 1) ≈ 0  # Regular order-1 entropy
            @test genentropy!(s, y, est, q = 1) >= 0 # Regular order-1 entropy
            @test genentropy!(s, x, est, q = 2) ≈ 0  # Higher-order entropy
            @test genentropy!(s, y, est, q = 2) >= 0 # Higher-order entropy

            # For a time series
            sz = zeros(Int, N - (est.m-1)*est.τ)
            @test probabilities!(sz, z, est) isa Probabilities
            @test probabilities(z, est) isa Probabilities
            @test genentropy!(sz, z, est) isa Real
            @test genentropy(z, est) isa Real
        end

        @testset "Not pre-allocated" begin
            est = SymbolicPermutation(m = 5, τ = 1)
            N = 500
            x = Dataset(repeat([1.1 2.2 3.3], N))
            y = Dataset(rand(N, 5))

            # Probability distributions
            p1 = probabilities(x, est)
            p2 = probabilities(y, est)
            @test sum(p1) ≈ 1.0
            @test sum(p2) ≈ 1.0

            # Entropy
            @test genentropy(x, est, q = 1) ≈ 0  # Regular order-1 entropy
            @test genentropy(y, est, q = 2) >= 0 # Higher-order entropy
        end
    end

    @testset "Weighted permutation entropy" begin
        m = 4
        τ = 1
        τs = tuple([τ*i for i = 0:m-1]...)
        x = rand(100)
        D = genembed(x, τs)

        # Probability distributions
        est = SymbolicWeightedPermutation(m = m, τ = τ)
        p1 = probabilities(x, est)
        p2 = probabilities(D, est)
        @test sum(p1) ≈ 1.0
        @test sum(p2) ≈ 1.0
        @test all(p1.p .≈ p2.p)

        # Entropy
        e1 = genentropy(D, est)
        e2 = genentropy(x, est)
        @test e1 ≈ e2
    end

    @testset "Amplitude-aware permutation entropy" begin
        m = 4
        τ = 1
        τs = tuple([τ*i for i = 0:m-1]...)
        x = rand(25)
        D = genembed(x, τs)

        est = SymbolicAmplitudeAwarePermutation(m = m, τ = τ)
        # Probability distributions
        p1 = probabilities(x, est)
        p2 = probabilities(D, est)
        @test sum(p1) ≈ 1.0
        @test sum(p2) ≈ 1.0
        @test all(p1.p .≈ p2.p)

        # Entropy
        e1 = genentropy(D, est)
        e2 = genentropy(x, est)
        @test e1 ≈ e2
    end

    @testset "Permutation, custom sorting" begin

        @testset "isless_rand" begin
            # because permutations are partially random, we sort many times and check that
            # we get *a* (not *the one*) correct answer every time
            for i = 1:50
                s = sortperm([1, 2, 3, 2], lt = Entropies.isless_rand)
                @test s == [1, 2, 4, 3] || s == [1, 4, 2, 3]
            end
        end

        m = 4
        τ = 1
        τs = tuple([τ*i for i = 0:m-1]...)
        ts = rand(1:3, 100)
        D = genembed(ts, τs)

        @testset "SymbolicPermutation" begin
            est_isless = SymbolicPermutation(m = 5, τ = 1, lt = Base.isless)
            est_isless_rand = SymbolicPermutation(m = 5, τ = 1, lt = Entropies.isless_rand)
            @test Entropies.symbolize(ts, est_isless) isa Vector{<:Int}
            @test Entropies.symbolize(D, est_isless_rand) isa Vector{<:Int}
            @test probabilities(D, est_isless) isa Probabilities
            @test probabilities(D, est_isless_rand) isa Probabilities
        end

        @testset "SymbolicWeightedPermutation" begin
            est_isless = SymbolicWeightedPermutation(m = 5, τ = 1, lt = Base.isless)
            est_isless_rand = SymbolicWeightedPermutation(m = 5, τ = 1, lt = Entropies.isless_rand)
            @test probabilities(ts, est_isless) isa Probabilities
            @test probabilities(D, est_isless) isa Probabilities
        end

        @testset "SymbolicAmplitudeAwarePermutation" begin
            est_isless = SymbolicAmplitudeAwarePermutation(m = 5, τ = 1, lt = Base.isless)
            est_isless_rand = SymbolicAmplitudeAwarePermutation(m = 5, τ = 1, lt = Entropies.isless_rand)
            @test probabilities(ts, est_isless) isa Probabilities
            @test probabilities(D, est_isless) isa Probabilities
        end
    end

    @testset "VisitationFrequency" begin
        D = Dataset(rand(100, 3))

        @testset "Counting visits" begin
            @test Entropies.marginal_visits(D, RectangularBinning(0.2), 1:2) isa Vector{Vector{Int}}
            @test Entropies.joint_visits(D, RectangularBinning(0.2)) isa Vector{Vector{Int}}
        end

        binnings = [
            RectangularBinning(3),
            RectangularBinning(0.2),
            RectangularBinning([2, 2, 3]),
            RectangularBinning([0.2, 0.3, 0.3])
        ]

        @testset "Binning test $i" for i in 1:length(binnings)
            est = VisitationFrequency(binnings[i])
            @test probabilities(D, est) isa Probabilities
            @test genentropy(D, est, q=1, base = 3) isa Real # Regular order-1 entropy
            @test genentropy(D, est, q=3, base = 2) isa Real # Higher-order entropy
            @test genentropy(D, est, q=3, base = 1) isa Real # Higher-order entropy

        end
    end

    @testset "Wavelet" begin
        N = 200
        a = 10
        t = LinRange(0, 2*a*π, N)
        x = sin.(t .+  cos.(t/0.1)) .- 0.1;

        @testset "TimeScaleMODWT" begin
            wl = WT.Daubechies{4}()
            est = TimeScaleMODWT(wl)

            @test Entropies.get_modwt(x) isa AbstractArray{<:Real, 2}
            @test Entropies.get_modwt(x, wl) isa AbstractArray{<:Real, 2}

            W = Entropies.get_modwt(x)
            Nlevels = maxmodwttransformlevels(x)
            @test Entropies.energy_at_scale(W, 1) isa Real
            @test Entropies.energy_at_time(W, 1) isa Real

            @test_throws ErrorException Entropies.energy_at_scale(W, 0)
            @test_throws ErrorException Entropies.energy_at_scale(W, Nlevels + 2)
            @test_throws ErrorException Entropies.energy_at_time(W, 0)
            @test_throws ErrorException Entropies.energy_at_time(W, N+1)

            @test Entropies.relative_wavelet_energy(W, 1) isa Real
            @test Entropies.relative_wavelet_energies(W, 1:2) isa AbstractVector{<:Real}

            @test Entropies.time_scale_density(x, wl) isa AbstractVector{<:Real}
            @test genentropy(x, TimeScaleMODWT(), q = 1, base = 2) isa Real
            @test probabilities(x, TimeScaleMODWT()) isa Probabilities
        end
    end

    @testset "Nearest neighbor based" begin
        m = 4
        τ = 1
        τs = tuple([τ*i for i = 0:m-1]...)
        x = rand(250)
        D = genembed(x, τs)

        est_nn = KozachenkoLeonenko(w = 5)
        est_knn = Kraskov(k = 2, w = 1)

        @test genentropy(D, est_nn) isa Real
        @test genentropy(D, est_knn) isa Real
    end

    @testset "TransferOperator" begin
        D = Dataset(rand(1000, 3))

        binnings = [
            RectangularBinning(3),
            RectangularBinning(0.2),
            RectangularBinning([2, 2, 3]),
            RectangularBinning([0.2, 0.3, 0.3])
        ]

        @testset "Binning test $i" for i in 1:length(binnings)
            to = Entropies.transferoperator(D, binnings[i])
            @test to isa Entropies.TransferOperatorApproximationRectangular

            iv = invariantmeasure(to)
            @test iv isa InvariantMeasure

            p, bins = invariantmeasure(iv)
            @test p isa Probabilities
            @test bins isa Vector{<:SVector}

            @test probabilities(D, TransferOperator(binnings[i])) isa Probabilities
        end
    end
end


module CameraModels

using StaticArrays, CoordinateTransformations

export Ray, CameraModel, PixelCoordinate, Vector3, Vector2, Point3
export pixel2ray, point2pixel, canreproject, origin, direction, sensorsize, columns, rows, lookdirection, updirection

const PixelCoordinate = SVector{2, Float64}
const Vector2 = SVector{2, Float64}

const Point3 = SVector{3, Float64}
const Vector3 = SVector{3, Float64}

abstract type CameraModel end

origin3d = zeros(Point3)

struct Ray
    origin::Point3
    direction::Vector3
end

"""
    origin(ray)

Return the direction of the ray as a (Vector3)
"""
function origin end

origin(vector::Vector3) = origin3d
origin(ray::Ray) = vector.origin

"""
    direction(ray)

Return the origin of ray, typically just a zero $(Point3) for normal cameras.
"""
function direction end

direction(vector::Vector3) = vector
direction(ray::Ray) = vector.direction

"""
    columns(model::CameraModel)

Returns the width of the camera sensor.
"""
function columns end

"""
    rows(model::CameraModel)

Returns the height of the camera sensor.
"""
function rows end

"""
    sensorsize(model::CameraModel)

Return the size of the camera sensor. By default calling out to columns(model) and rows(model) to build an SVector{2}

`sensorsize(cameramodel::CameraModel) = SVector{2}(columns(cameramodel), rows(cameramodel))`
"""
function sensorsize end

sensorsize(cameramodel::CameraModel) = SVector{2}(columns(cameramodel), rows(cameramodel))

"""
    pixel2ray(cameramodel::CameraModel, pixelcoordinate::$(PixelCoordinate))

Returns the ray in space (origin + direction) corresponding to this `pixelcoordinate`.
"""
function pixel2ray end

"""
    point2pixel(camera::CameraModel, pointincamera::$(Point3))

Returns the pixel location onto which the 3D coordinate `pointincamera` is projected.
"""
function point2pixel end

"""
    lookdirection(camera::CameraModel)

Return the lookdirection of this camera model.
"""
function lookdirection end

"""
    updirection(camera::CameraModel)

Return the updirection of this camera model.
"""
function updirection end

"""
    canreproject(camera::CameraModel)

Confirms if point2pixel is implemented for this camera model.
"""
canreproject(camera::CameraModel) = true

include("Pinhole.jl")

end # module


using MixedLayerThermoclineDynamics, OffsetArrays, Test, Documenter

const rtol_interpolation = 8e-4
const rtol_derivatives = 8e-4

@time @testset "Grid tests" begin
    include("test_grids.jl")

    for Tx in (Periodic(), Bounded()), Ty in (Periodic(), Bounded())

        nx, ny = 10, 12
        hx, hy =  1, 2
        x0, y0 = 0.2, -0.4
        Lx, Ly = 2.0, 2.4

        dx, dy = Lx/nx, Ly/ny

        grid1D = Grid1D(Tx, nx, x0, x0 + Lx; hx=hx)
        grid2D = Grid2D(Tx, Ty, nx, ny, x0, x0 + Lx, y0, y0 + Ly; hx=hx, hy=hy)

        xF = isa(Tx, Periodic) ? range(x0 - hx*dx, stop = x0 + Lx - dx + hx*dx, length = nx + 2hx) :
                                 range(x0 - hx*dx, stop = x0 + Lx + hx*dx, length = nx + 1 + 2hx)

        xC = range(x0 + dx/2 - hx*dx, stop = x0 + Lx - dx/2 + hx*dx, length = nx + 2hx)

        yF = isa(Ty, Periodic) ? range(y0 - hy*dy, stop = y0 + Ly - dy + hy*dy, length = ny + 2hy) :
                                 range(y0 - hy*dy, stop = y0 + Ly + hy*dy, length = ny + 1 + 2hy)

        yC = range(y0 + dy/2 - hy*dy, stop = y0 + Ly - dy/2 + hy*dy, length = ny + 2hy)

        xF = OffsetArray(xF, -hx)
        xC = OffsetArray(xC, -hx)
        yF = OffsetArray(yF, -hy)
        yC = OffsetArray(yC, -hy)

        @test grid1D.dx == dx
        @test test_xF(grid1D, xF)
        @test test_xC(grid1D, xC)
        @test xdomain_length(grid1D, Lx)

        @test grid2D.dx == dx
        @test grid2D.dy == dy
        @test test_xF(grid2D, xF)
        @test test_xC(grid2D, xC)
        @test test_yF(grid2D, yF)
        @test test_yC(grid2D, yC)
        @test xdomain_length(grid2D, Lx)
        @test ydomain_length(grid2D, Ly)
    end

end

@time @testset "Field tests" begin

    include("test_fields.jl")

    nx, ny = 100, 120
    Lx, Ly = 1.0, 1.2
    hx, hy = 2, 3

    @test_throws ErrorException("Number of halo points in x cannot be zero") Grid1D(Periodic(), nx, 0, Lx; hx = 0)
    @test_throws ErrorException("Number of halo points in x cannot be zero") Grid2D(Periodic(), Periodic(), nx, 0, Lx, ny, 0, Ly; hx = 0, hy = 1)
    @test_throws ErrorException("Number of halo points in y cannot be zero") Grid2D(Periodic(), Periodic(), nx, 0, Lx, ny, 0, Ly; hx = 1, hy = 0)

    dx, dy = Lx/nx, Ly/ny # note: this is only correct for Periodic

    grid1D = Grid1D(Periodic(), nx, 0, Lx; hx = hx)
    grid2D = Grid2D(Periodic(), Periodic(), nx, ny, 0, Lx, 0, Ly; hx = hx, hy = hy)

    #####
    ##### 1D Fields
    #####

    Cdata = @. sin(2π * grid1D.xC / Lx)
    Fdata = @. cos(2π * grid1D.xF / Lx)

    Cdata_on_F = @. sin(2π * grid1D.xF / Lx)
    Fdata_on_C = @. cos(2π * grid1D.xC / Lx)

    ∂x_Cdata = @.   2π/Lx * cos(2π * grid1D.xF / Lx)
    ∂x_Fdata = @. - 2π/Lx * sin(2π * grid1D.xC / Lx)

    ∂x_Cdata_on_C = @.   2π/Lx * cos(2π * grid1D.xC / Lx)
    ∂x_Fdata_on_F = @. - 2π/Lx * sin(2π * grid1D.xF / Lx)

    C_Field1D = Field(Centre, Cdata, grid1D)
    F_Field1D = Field(Face, Fdata, grid1D)

    Cdata_on_F_Field1D = Field(Face, Cdata_on_F, grid1D)
    Fdata_on_C_Field1D = Field(Centre, Fdata_on_C, grid1D)

    Ftest_Field1D = Field(Face, zeros(nx), grid1D)
    Ctest_Field1D = Field(Centre, zeros(nx), grid1D)

    ∂x_Cdata_Field1D = Field(Face, ∂x_Cdata, grid1D)
    ∂x_Fdata_Field1D = Field(Centre, ∂x_Fdata, grid1D)
    ∂x_Cdata_on_C_Field1D = Field(Centre, ∂x_Cdata_on_C, grid1D)
    ∂x_Fdata_on_F_Field1D = Field(Face, ∂x_Fdata_on_F, grid1D)

    C_Field1D_from_array = Field1D(Centre, Cdata[1:nx], grid1D)
    F_Field1D_from_array = Field1D(Face, Fdata[1:nx], grid1D)

    @test C_Field1D.grid == C_Field1D_from_array.grid
    @test F_Field1D.grid == F_Field1D_from_array.grid
    @test C_Field1D.data ≈ C_Field1D_from_array.data
    @test F_Field1D.data ≈ F_Field1D_from_array.data

    @test typeof(C_Field1D) <: Field1D{Centre}
    @test typeof(F_Field1D) <: Field1D{Face}
    @test typeof(C_Field1D_from_array) <: Field1D{Centre}
    @test typeof(F_Field1D_from_array) <: Field1D{Face}

    @test C_Field1D.grid == grid1D
    @test F_Field1D.grid == grid1D

    @test test_𝐼x(Fdata_on_C_Field1D, Ctest_Field1D, F_Field1D)
    @test test_𝐼x(Cdata_on_F_Field1D, Ftest_Field1D, C_Field1D)

    @test test_∂x(∂x_Fdata_Field1D, Ctest_Field1D, F_Field1D)
    @test test_∂x(∂x_Cdata_Field1D, Ftest_Field1D, C_Field1D)
    @test test_∂x(∂x_Fdata_on_F_Field1D, Ftest_Field1D, F_Field1D)
    @test test_∂x(∂x_Cdata_on_C_Field1D, Ctest_Field1D, C_Field1D)

    Cdata_with_halos = OffsetArray(zeros(nx + 2hx), -hx)

    for i in 1:nx
        Cdata_with_halos[i] = Cdata[i]
    end

    for i in 1:hx
        Cdata_with_halos[nx+i] = Cdata[i]
        Cdata_with_halos[-i+1] = Cdata[nx-i+1]
    end

    Fdata_with_halos = OffsetArray(zeros(nx + 2hx), -hx)

    for i in 1:nx
        Fdata_with_halos[i] = Fdata[i]
    end

    for i in 1:hx
        Fdata_with_halos[nx+i] = Fdata[i]
        Fdata_with_halos[-i+1] = Fdata[nx-i+1]
    end

    @test C_Field1D.data ≈ Cdata_with_halos
    @test F_Field1D.data ≈ Fdata_with_halos

    #####
    ##### 2D Fields
    #####

    CCdata = @. [sin(2π * grid2D.xC[i] / Lx) * cos(2π * grid2D.yC[j] / Ly - π/5) for i in 1:nx, j in 1:ny]
    FCdata = @. [cos(2π * grid2D.xF[i] / Lx + π/3) * cos(2π * grid2D.yC[j] / Ly) for i in 1:nx, j in 1:ny]
    CFdata = @. [sin(2π * grid2D.xC[i] / Lx + π/4) * sin(2π * grid2D.yF[j] / Ly) for i in 1:nx, j in 1:ny]

    CCdata_on_FC = @. [sin(2π * grid2D.xF[i] / Lx) * cos(2π * grid2D.yC[j] / Ly - π/5) for i in 1:nx, j in 1:ny]
    CCdata_on_CF = @. [sin(2π * grid2D.xC[i] / Lx) * cos(2π * grid2D.yF[j] / Ly - π/5) for i in 1:nx, j in 1:ny]
    FCdata_on_CC = @. [cos(2π * grid2D.xC[i] / Lx + π/3) * cos(2π * grid2D.yC[j] / Ly) for i in 1:nx, j in 1:ny]
    CFdata_on_CC = @. [sin(2π * grid2D.xC[i] / Lx + π/4) * sin(2π * grid2D.yC[j] / Ly) for i in 1:nx, j in 1:ny]

    ∂x_CCdata_on_FC = @.   2π/Lx * [cos(2π * grid2D.xF[i] / Lx) * cos(2π * grid2D.yC[j] / Ly - π/5) for i in 1:nx, j in 1:ny]
    ∂y_CCdata_on_CF = @. - 2π/Ly * [sin(2π * grid2D.xC[i] / Lx) * sin(2π * grid2D.yF[j] / Ly - π/5) for i in 1:nx, j in 1:ny]
    ∂x_FCdata_on_CC = @. - 2π/Lx * [sin(2π * grid2D.xC[i] / Lx + π/3) * cos(2π * grid2D.yC[j] / Ly) for i in 1:nx, j in 1:ny]
    ∂y_CFdata_on_CC = @.   2π/Ly * [sin(2π * grid2D.xC[i] / Lx + π/4) * cos(2π * grid2D.yC[j] / Ly) for i in 1:nx, j in 1:ny]

    ∂x_CCdata_on_CC = @.   2π/Lx * [cos(2π * grid2D.xC[i] / Lx) * cos(2π * grid2D.yC[j] / Ly - π/5) for i in 1:nx, j in 1:ny]
    ∂y_CCdata_on_CC = @. - 2π/Ly * [sin(2π * grid2D.xC[i] / Lx) * sin(2π * grid2D.yC[j] / Ly - π/5) for i in 1:nx, j in 1:ny]

    ∂x_FCdata_on_FC = @. - 2π/Lx * [sin(2π * grid2D.xF[i] / Lx + π/3) * cos(2π * grid2D.yC[j] / Ly) for i in 1:nx, j in 1:ny]
    ∂y_CFdata_on_CF = @.   2π/Ly * [sin(2π * grid2D.xC[i] / Lx + π/4) * cos(2π * grid2D.yF[j] / Ly) for i in 1:nx, j in 1:ny]

    CC_Field2D = Field(Centre, Center, CCdata, grid2D)
    FC_Field2D = Field(Face, Center, FCdata, grid2D)
    CF_Field2D = Field(Centre, Face, CFdata, grid2D)

    CCdata_on_FC_Field2D = Field(Face, Centre, CCdata_on_FC, grid2D)
    CCdata_on_CF_Field2D = Field(Centre, Face, CCdata_on_CF, grid2D)
    FCdata_on_CC_Field2D = Field(Centre, Centre, FCdata_on_CC, grid2D)
    CFdata_on_CC_Field2D = Field(Centre, Centre, CFdata_on_CC, grid2D)

    FCtest_Field2D = Field(Face, Centre, zeros(nx, ny), grid2D)
    CFtest_Field2D = Field(Centre, Face, zeros(nx, ny), grid2D)
    CCtest_Field2D = Field(Centre, Centre, zeros(nx, ny), grid2D)

    ∂x_CCdata_on_FC_Field2D = Field(Face, Centre, ∂x_CCdata_on_FC, grid2D)
    ∂y_CCdata_on_CF_Field2D = Field(Centre, Face, ∂y_CCdata_on_CF, grid2D)
    ∂x_FCdata_on_CC_Field2D = Field(Centre, Centre, ∂x_FCdata_on_CC, grid2D)
    ∂y_CFdata_on_CC_Field2D = Field(Centre, Centre, ∂y_CFdata_on_CC, grid2D)
    ∂x_CCdata_on_CC_Field2D = Field(Centre, Centre, ∂x_CCdata_on_CC, grid2D)
    ∂y_CCdata_on_CC_Field2D = Field(Centre, Centre, ∂y_CCdata_on_CC, grid2D)
    ∂x_FCdata_on_FC_Field2D = Field(Face, Centre, ∂x_FCdata_on_FC, grid2D)
    ∂y_CFdata_on_CF_Field2D = Field(Centre, Face, ∂y_CFdata_on_CF, grid2D)

    CC_Field2D_from_array = Field2D(Centre, Centre, CCdata[1:nx, 1:ny], grid2D)
    FC_Field2D_from_array = Field2D(Face, Centre, FCdata[1:nx, 1:ny], grid2D)
    CF_Field2D_from_array = Field2D(Centre, Face, CFdata[1:nx, 1:ny], grid2D)

    @test typeof(CC_Field2D) <: Field2D{Centre, Centre}
    @test typeof(FC_Field2D) <: Field2D{Face, Centre}
    @test typeof(CF_Field2D) <: Field2D{Centre, Face}

    @test CC_Field2D.grid == grid2D
    @test FC_Field2D.grid == grid2D
    @test CF_Field2D.grid == grid2D

    @test CC_Field2D.grid == CC_Field2D_from_array.grid
    @test FC_Field2D.grid == FC_Field2D_from_array.grid
    @test CF_Field2D.grid == CF_Field2D_from_array.grid
    @test CC_Field2D.data ≈ CC_Field2D_from_array.data
    @test FC_Field2D.data ≈ FC_Field2D_from_array.data
    @test CF_Field2D.data ≈ CF_Field2D_from_array.data

    for field in [C_Field1D, F_Field1D, CC_Field2D, FC_Field2D, CF_Field2D]
        @test typeof(field) <: AbstractField
    end

    @test test_𝐼x(CCdata_on_FC_Field2D, FCtest_Field2D, CC_Field2D)
    @test test_𝐼y(CCdata_on_CF_Field2D, CFtest_Field2D, CC_Field2D)
    @test test_𝐼x(FCdata_on_CC_Field2D, CCtest_Field2D, FC_Field2D)
    @test test_𝐼y(CFdata_on_CC_Field2D, CCtest_Field2D, CF_Field2D)

    @test test_∂x(∂x_CCdata_on_FC_Field2D, FCtest_Field2D, CC_Field2D)
    @test test_∂y(∂y_CCdata_on_CF_Field2D, CFtest_Field2D, CC_Field2D)
    @test test_∂x(∂x_FCdata_on_CC_Field2D, CCtest_Field2D, FC_Field2D)
    @test test_∂y(∂y_CFdata_on_CC_Field2D, CCtest_Field2D, CF_Field2D)
    @test test_∂x(∂x_CCdata_on_CC_Field2D, CCtest_Field2D, CC_Field2D)
    @test test_∂y(∂y_CCdata_on_CC_Field2D, CCtest_Field2D, CC_Field2D)
    @test test_∂x(∂x_FCdata_on_FC_Field2D, FCtest_Field2D, FC_Field2D)
    @test test_∂y(∂y_CFdata_on_CF_Field2D, CFtest_Field2D, CF_Field2D)
end

@time @testset "Doctests" begin
    doctest(MixedLayerThermoclineDynamics)
end


abstract type AbstractCurve2D end
abstract type AbstractCircularLinearCurve <: AbstractCurve2D end

abstract type AbstractLinearCurve <: AbstractCircularLinearCurve end
abstract type AbstractLine <: AbstractLinearCurve end
abstract type AbstractRay <: AbstractLinearCurve end
abstract type AbstractLineSegment <: AbstractLinearCurve end

abstract type AbstractCircularCurve <: AbstractCircularLinearCurve end
abstract type AbstractCircle <: AbstractCircularCurve end
abstract type AbstractArc <: AbstractCircularCurve end


module UniformlyRefinedForestOfOctreesDiscreteModelsTests
  using MPI
  using Gridap
  using Gridap.Algebra
  using Gridap.FESpaces
  using PartitionedArrays
  using GridapDistributed
  using GridapP4est
  using Test
  using ArgParse
  using P4est_wrapper

  function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--subdomains", "-s"
        help = "Tuple with the # of coarse subdomains per Cartesian direction"
        arg_type = Int64
        default=[1,1]
        nargs='+'
        "--num-uniform-refinements", "-r"
        help = "# of uniform refinements"
        arg_type = Int64
        default=1
    end
    return parse_args(s)
  end

  function run(parts,subdomains,num_uniform_refinements)
    # Manufactured solution
    u(x) = x[1] + x[2]
    f(x) = -Δ(u)(x)

    if length(subdomains)==2
      domain=(0,1,0,1)
    else
      @assert length(subdomains)==3
      domain=(0,1,0,1,0,1)
    end

    coarse_discrete_model=CartesianDiscreteModel(domain,subdomains)
    model=UniformlyRefinedForestOfOctreesDiscreteModel(parts,
                                                       coarse_discrete_model,
                                                       num_uniform_refinements)

    # FE Spaces
    order=1
    reffe = ReferenceFE(lagrangian,Float64,order)
    V = TestFESpace(model,reffe,dirichlet_tags="boundary")
    U = TrialFESpace(u,V)

    trian=Triangulation(model)
    dΩ=Measure(trian,2*(order+1))

    function a(u,v)
      ∫(∇(v)⋅∇(u))dΩ
    end
    function l(v)
      ∫(v*f)dΩ
    end
    dv = get_fe_basis(V)
    du = get_trial_fe_basis(U)
    assem = SparseMatrixAssembler(U,V)

    dof_values = PVector(0.0,V.gids)
    uh = FEFunction(U,dof_values)
    data = collect_cell_matrix_and_vector(U,V,a(du,dv),l(dv),uh)
    A,b = assemble_matrix_and_vector(assem,data)

    x = A\b
    r = A*x -b
    uh = FEFunction(U,x)

    # Error norms and print solution
    trian=Triangulation(model)
    dΩ=Measure(trian,2*order)
    e = u-uh
    e_l2 = sum(∫(e*e)dΩ)
    tol = 1.0e-9
    @test e_l2 < tol
    map_parts(parts) do part
     if (part==1)
       println("$(e_l2) < $(tol)\n")
     end
    end
  end
  if !MPI.Initialized()
    MPI.Init()
  end
  parsed_args = parse_commandline()
  subdomains = Tuple(parsed_args["subdomains"])
  num_uniform_refinements = parsed_args["num-uniform-refinements"]
  parts = get_part_ids(mpi,(prod(subdomains)))
  run(parts,subdomains,num_uniform_refinements)
  MPI.Finalize()
end # module


# First passage time density (and stationary rate, and ISI statistics)
# for a single LIF neuron without refractoriness!

push!(LOAD_PATH, abspath(@__DIR__,"..","src"))
using LinearAlgebra,Statistics,StatsBase
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using FFTW
using ProgressMeter

function onesparsemat(w::Real)
  mat=Matrix{Float64}(undef,1,1) ; mat[1,1]=w
  return sparse(mat)
end

# Time parameters
dt = 1E-3
Ttot = 6.
ntimes = floor(Int64,Ttot/dt)+1
# LIF neuron parameters
myτ = 0.5
v_th = 10.
v_r = -5.0
τrefr = 0.0 # refractoriness
τpcd = 0.2 # post synaptic current decay
myinput = 9.5
myσ = 1.0
mynt = S.NTLIF(myτ,v_th,v_r,τrefr,τpcd)
myps = S.PSLIF(mynt,1)

## inputs here
# one periodic input (freq is the inverse of period)
ωin = inv(2.0)
infun = function(t)
  α = 1.
  return myinput + α*sin(ωin*t*2π)
end
in_state = S.PSSimpleInput(S.InputSomeFunction(infun))
conn_in = S.BaseConnection(onesparsemat(1.0))
# ... plus noise
in_state_noise = S.PSSimpleInput(S.InputIndependentNormal(myσ))
##
mypop = S.Population(myps,(conn_in,S.FakeConnection()),
    (in_state,in_state_noise))

## that's it, let's make the network

myntw = S.RecurrentNetwork(dt,(mypop,))

## Recorder
krec = 10
rec_state = S.RecStateNow(myps,krec,dt,Ttot)
rec_spikes = S.RecCountSpikes(myps,dt)

## Run

times = (0:myntw.dt:Ttot)
nt = length(times)
# initial conditions
myps.state_now[1] = v_r

S.reset!.([rec_state,rec_spikes])

for (k,t) in enumerate(times)
  rec_state(t,k,myntw)
  rec_spikes(t,k,myntw)
  S.dynamics_step!(t,myntw)
end
r0 = S.get_mean_rates(rec_spikes,dt,Ttot)[1]
##
_ = let plt=plot(;xlabel="time",ylabel="voltage",leg=false)
  plot!(plt,rec_state.times,rec_state.state_now[1,:],linewidth=2)
end
##

function simulation_linear_response_coef(ω,E;dt=1E-3,Ttot::Float64=1000.0)
  # network without extra input
  in_noise = S.PSSimpleInput(S.InputIndependentNormal(myσ))
  in_constant = S.PSSimpleInput(S.InputSimpleOffset(myinput))
  mypop = S.Population(myps,
      (S.FakeConnection(),S.FakeConnection()),
      (in_constant,in_noise))

  rec_spikes = S.RecCountSpikes(myps,dt)
  times = (0:dt:Ttot)
  # initial conditions
  myps.state_now[1] = v_r
  ntw = S.RecurrentNetwork(dt,(mypop,))
  S.reset!(rec_spikes)
  @showprogress "1/2 --> " for (k,t) in enumerate(times)
    rec_spikes(t,k,ntw)
    S.dynamics_step!(t,ntw)
  end
  r0 = S.get_mean_rates(rec_spikes,dt,Ttot)[1]
  # network with extra input
  infun = function(t)
    return myinput + E*sin(ω*t*2π)
  end
  in_func = S.PSSimpleInput(S.InputSomeFunction(infun))
  conn_in = S.BaseConnection(onesparsemat(1.0))
  mypop = S.Population(myps,
      (conn_in,S.FakeConnection()),
      (in_func,in_noise))

  myps.state_now[1] = v_r
  ntw = S.RecurrentNetwork(dt,(mypop,))
  S.reset!(rec_spikes)
  @showprogress "2/2 --> " for (k,t) in enumerate(times)
    rec_spikes(t,k,ntw)
    S.dynamics_step!(t,ntw)
  end
  r1 = S.get_mean_rates(rec_spikes,dt,Ttot)[1]
  return (r1-r0)/E
end

##

t_period = range(0.1,2;length=20)
lrs = map(t_period) do _t
  simulation_linear_response_coef(inv(_t),10.0;Ttot=1500.0)
end
lrs2 = map(t_period) do _t
  simulation_linear_response_coef(inv(_t),30.0;Ttot=1500.0)
end

plot(t_period,[lrs lrs2];label=["E=10" "E=30"], marker=:star)

############################
### now follow the approach of Richardson et al
##
# start with stationary case

dv = 1E-3
v_lb = v_r - 3.0
ndv = round(Int64,(v_th - v_lb)/dv)
k_re = round(Int64,(v_r - v_lb)/dv)
ks = collect(0:ndv)
Vks = @. v_lb + ks * dv

j0s = fill(NaN,ndv+1)
p0s = fill(NaN,ndv+1)

ksrev = collect(ndv+1:-1:2)

j0s[ndv+1]=1.0
p0s[ndv+1]=0.0

for k in ksrev
  krond = k==k_re ? 1.0 : 0.0
  j0s[k-1] = j0s[k] - krond
  Gk = (Vks[k]-myinput)/(myσ^2)
  Hk = myτ*j0s[k]/(myσ^2)
  Ak = exp(dv*Gk)
  Bk = Gk==0 ? dv*Hk : Hk/Gk*(Ak-1)
  p0s[k-1] = p0s[k]*Ak + Bk
end

r0 = inv(dv*sum(p0s))
P0s = r0 .* p0s

## now the non-stationary linear addendum

function theory_linear_response!(jr,je,pr,pe,ω::Real;E=1.0)
  jr[ndv+1] = 1.0
  je[ndv+1] = 0.
  pr[ndv+1] = 0.
  pe[ndv+1] = 0.
  for k in ksrev
    krond = k==k_re ? 1.0 : 0.0
    jr[k-1] = jr[k] + im*2π*ω*pr[k] - krond
    je[k-1] = je[k] + im*2π*ω*pe[k]
    G = (Vks[k]-myinput)/(myσ^2)
    Hr = myτ*jr[k]/(myσ^2)
    He = (myτ*je[k]-P0s[k])/(myσ^2)
    Ak = exp(dv*G)
    pr[k-1] = G==0 ? pr[k] + Hr*dv : pr[k]*Ak + Hr/G*(Ak-1)
    pe[k-1] = G==0 ? pe[k] + He*dv : pe[k]*Ak + He/G*(Ak-1)
    #@show pe[k]
    #@show pe[k-1]
    @assert isfinite(pe[k-1])
  end
  #return je
  AEω = (-je[1]/jr[1])
  Pfull = @. r0*p0s + E*AEω + E*pe
  return AEω ,Pfull
end

_jr,_je,_pr,_pe =ntuple(_->similar(p0s,ComplexF64),4)

_Aω,_Pfullω = theory_linear_response!(_jr,_je,_pr,_pe,inv(6.0))

_ = let y = real.(_Pfullω),
  x=Vks
  plot(x,y)
end


using StaticArrays
using OrdinaryDiffEq
using Test

@inbounds @inline function ż(z, p, t)
    A, B, D = p
    p₀, p₂ = z[1:2]
    q₀, q₂ = z[3:4]

    return SVector{4}(
        -A * q₀ - 3 * B / √2 * (q₂^2 - q₀^2) - D * q₀ * (q₀^2 + q₂^2),
        -q₂ * (A + 3 * √2 * B * q₀ + D * (q₀^2 + q₂^2)),
        A * p₀,
        A * p₂
  )
end

condition(u, t, integrator) = u
affect!(integrator) = nothing
cbf(idx) = ContinuousCallback(condition,
    affect!, nothing, save_positions=(false, true), idxs=idx)
z0 = SVector{4}(7.1989885061904335, -0.165912283356219, 0., -3.63534900748947)

tspan=(0.,300.)
prob=ODEProblem(ż, z0, tspan, (A=1, B=0.55, D=0.4), callback=cbf(3))
sol=solve(prob, Vern9(), abstol=1e-14, reltol=1e-14,
    save_everystep=false, save_start=false, save_end=false, maxiters=1e6)

@test length(sol) > 100

prob=ODEProblem(ż, z0, (0,400.), (A=1, B=0.55, D=0.4), callback=cbf(3))
sol=solve(prob, Vern9(), abstol=1e-14, reltol=1e-14, save_everystep=false, save_start=false, save_end=false, maxiters=2e4)

@test length(sol) > 100

prob=ODEProblem(ż, z0, (0,5000.), (A=1, B=0.55, D=0.4), callback=cbf(3))
sol=solve(prob, Vern9(), abstol=1e-14, reltol=1e-14, save_everystep=false, save_start=false, save_end=false, maxiters=1e6)

@test length(sol) > 1500

@info "Bouncing Ball"

f = function (du,u,p,t)
  du[1] = u[2]
  du[2] = -p[1]
end
function condition(u,t,integrator) # Event when event_f(u,t) == 0
  u[1]
end
function affect!(integrator)
  integrator.u[2] = -integrator.u[2]
end
cb2 = ContinuousCallback(condition,affect!)
tspan = (0.0,10000.0)
u0 = [50.0,0.0]
p = 9.8
prob = ODEProblem(f,u0,tspan,p)
sol = solve(prob,Tsit5(),callback=cb2)
@test minimum(Array(sol)) > -40

function fball(du,u,p,t)
  du[1] = u[2]
  du[2] = -p
  du[3] = u[4]
  du[4] = 0.0
end
u0 = [50.0,0.0,0.0,2.0]
tspan = (0.0,15.0)
p = 9.8
prob = ODEProblem(fball,u0,tspan,p)

x = Ref(0)
y = Ref(0)
z = Ref(0)

function condition(out,u,t,integrator)
  out[1] = u[1]
  out[2] = (10.0 - u[3])u[3]
end

function affect!(integrator, idx)
  if idx == 1
    x[] += 1
    integrator.u[2] = -0.9integrator.u[2]
  elseif idx == 2
    y[] += 1
    integrator.u[4] = -0.9integrator.u[4]
  end
end

function affect_neg!(integrator, idx)
  z[] += 1
  @show integrator.u[1]
  @show integrator.u[3]
end

cb = VectorContinuousCallback(condition,affect!,2)
cb2 = VectorContinuousCallback(condition,affect_neg!,affect!,2)

sol = solve(prob,Tsit5(),callback=cb,dt=1e-3,adaptive=false)
@test x[] == 3
@test y[] == 2

sol2 = solve(prob,Tsit5(),callback=cb2,dt=1e-3,adaptive=false)
@test sol.u == sol2.u
@test z[] == 0

# https://github.com/SciML/OrdinaryDiffEq.jl/issues/1273

function du!(du, u, p, t)
    du[1] = 1
end

callback = ContinuousCallback(
    (u, t, integrator) -> 1.0,
    (integrator) -> nothing
)

prob = ODEProblem(du!, [0], (0.0, 1.0), callback=callback)

solve(prob, Tsit5())
solve(prob, RadauIIA3())
solve(prob, RadauIIA5())

du(u, p, t) = [1]

prob = ODEProblem(du, [0], (0.0, 1.0), callback=callback)

solve(prob, Tsit5())
solve(prob, RadauIIA3())
solve(prob, RadauIIA5())


@testset "$TEST $G" begin

@test length(collect(edges(G()))) == 0

ga = G(10,20; seed=1)
gb = G(10,20; seed=1)

@test length(collect(edges(ga))) == 20
@test collect(edges(ga)) == collect(edges(gb))
for e in edges(ga)
    @test has_edge(ga, e)
end
@test length(unique(collect(edges(ga)))) == 20

ga = erdos_renyi(10,20, G; seed=1)
gb = erdos_renyi(10,20, DG; seed=1)

# TODO implement length(edges(g)) somehow
# @test length(collect(edges(ga))) == length(edges(ga)) == ne(ga)
# @test length(collect(edges(gb))) == length(edges(gb)) == ne(gb)
@test length(collect(edges(ga))) == ne(ga)
@test length(collect(edges(gb))) == ne(gb)

for e in edges(ga)
    @test has_edge(ga, e)
end
@test length(unique(collect(edges(ga)))) == 20

ga = G(10)
add_edge!(ga, 3, 2)
add_edge!(ga, 3, 10)
add_edge!(ga, 5, 10)
add_edge!(ga, 10, 3)

eit = edges(ga)
# TODO?
# @test Base.IteratorSize(typeof(eit)) == Base.HasLength()

i = 0
for e in eit
    @test src(e) <= dst(e)
    i += 1
end
@test i == ne(ga)

@test [e for e in eit] == [E(2, 3), E(3, 10), E(5,10)]

ga = DG(10)
add_edge!(ga, 3, 2)
add_edge!(ga, 3, 10)
add_edge!(ga, 5, 10)
add_edge!(ga, 10, 3)

eit = edges(ga)

@test [e for e in eit] == [E(3, 2), E(3, 10), E(5,10), E(10,3)]
@test collect(eit) == [e for e in eit]
# g = CompleteGraph(10)
# @test collect(edges(g,1:3)) == collect(edges(g,[1:3;]))
# @test length(collect(edges(g, 1:3))) == 3
# @test length(collect(edges(g, 1:0))) == 0
#
#
# g = CompleteDiGraph(10)
# @test collect(edges(g,1:3)) == collect(edges(g,[1:3;]))
# @test length(collect(edges(g, 1:3))) == 6
# @test length(collect(edges(g, 1:0))) == 0

end # testset


function normalize_ts(nrm_ts)
    mn=minimum(nrm_ts.price)
    nrm_ts.price=nrm_ts.price.-mn#set minimum to zero
    mx=maximum(nrm_ts.price)
    nrm_ts.price=nrm_ts.price./mx#normalize
    return nrm_ts
end

function rand_set_days(ukbe_data,dn,wn,mn)
    sampled_set=DataFrame()
    selected_months=[]; for j=1:mn; months=[i for i=1:12 if !issubset(i,selected_months)]; push!(selected_months,rand(months));end;sort!(selected_months)
    for mnth in selected_months;#println("month: "*string(mnth))
        hours_in_month=filter(e->month(e.time_stamp)==mnth,ukbe_data)
        days_in_month=round(Int64,length(hours_in_month[!,:time_stamp])/24)
        weeks_in_month=ceil(Int64,days_in_month/7);selected_weeks=[];
        #for w=1:wn; weeks=[i for i=1:weeks_in_month if !issubset(i,selected_weeks)]; push!(selected_weeks,rand(weeks));end;sort!(selected_weeks);selected_days=[];
        selected_weeks=wn;selected_days=[];
        for wek in selected_weeks;#println(wek)
            for d=1:dn; days=[i for i in (wek-1)*7+1:min((wek)*7,days_in_month) if !issubset(i,selected_days)]; if (length(days)>0); push!(selected_days,rand(days));end;end;sort!(selected_days);end
        for dy in selected_days;
            hours_in_day=filter(e->day(e.time_stamp)==dy,hours_in_month)
            sampled_set=vcat(sampled_set,hours_in_day)
        end
    end
    sort!(sampled_set)
    return sampled_set
end


@testset "480.sliding-window-median.jl" begin
    @test median_sliding_window([1,3,-1,-3,5,3,6,7], 3) == [1,-1,-1,3,5,6.]
end


mutable struct SumFun <: Function# OscLaurent{D<:SPeriodicLine,R}?
    funs::Vector{Fun}
end

evaluate(f::SumFun,x) = sum(map(y -> y(x),f.funs))

SumFun(f::Fun{T}) where T = SumFun([f])
SumFun(f::Fun{T}) where T<: SumSpace = SumFun(components(f))
SumFun(f::Fun{T}) where T<: PiecewiseSpace = SumFun(components(f))

conj(s::SumFun) = SumFun(map(conj,s.funs))
conj(s::Array{T}) where T<:SumFun = map(conj,s)

#sum(s::SumFun) = sum(map(sum,s.funs))
sum(s::SumFun) = s.funs == [] ? 0.0 : sum(map(sum,s.funs))
sum(s::Array{T}) where T<:SumFun = map(sum,s)

function combine!(s::SumFun)
    s.funs = combine!(s.funs)
    s
end
combine!(s::Vector{SumFun}) = map(combine!,s)
combine!(s::Array{SumFun}) = map(combine!,s)
components(f::Fun) = f

function combine!(f::Vector{T}) where T <: Fun
    f = vcat(map(components,f))
    j = 0
    @inbounds while j < length(f)
        j += 1
        if (length(f[j].coefficients) <= 4 && sum(abs.(f[j].coefficients)) < 1e-15 && length(f) > 2) || maximum(abs.(f[j].coefficients)) < 1e-15
            deleteat!(f,j)
            j -= 1
        end
    end
    i = 0
    @inbounds while i < length(f)
        i += 1
        j = i
        while j < length(f)
            j += 1
            if spacescompatible(f[i],f[j])
                f[i] += f[j]
                deleteat!(f,j)
                j -= 1
            end
        end
    end
    f
end

+(f::SumFun,g::SumFun) = SumFun(vcat(g.funs,f.funs))
-(f::SumFun,g::SumFun) = SumFun(vcat(-g.funs,f.funs))

+(f::SumFun,g::Fun) = f + SumFun(g)
+(g::Fun,f::SumFun) = f + SumFun(g)

-(f::SumFun,g::Fun) = f - SumFun(g)
-(g::Fun,f::SumFun) = -f + SumFun(g)

-(f::SumFun) = SumFun(map(y -> -y,f.funs))

*(f::Fun,g::SumFun) = SumFun(map(y -> f*y,g.funs))
*(f::SumFun,g::Fun) = SumFun(map(y -> g*y,f.funs))
*(f::SumFun,g::Number) = SumFun(map(y -> g*y,f.funs))
*(f::Number,g::SumFun) = SumFun(map(y -> f*y,g.funs))
*(f::SumFun,g::Array) = map(x -> f*x,g)

*(f::SumFun,g::SumFun) = combine!(+(map(y -> y*g,f.funs)...))

*(c::Float64,g::SumFun) = SumFun(map(y -> c*y,g.funs))
*(c::Complex{Float64},g::SumFun) = SumFun(map(y -> c*y,g.funs))

(f::SumFun)(x) = evaluate(f,x)
+(f::SumFun) = f

(f::Array)(x) = map(y -> evaluate(y,x),f)
chop!(f::SumFun) = SumFun(map(chop!,f.funs))
chop!(f::Array) = map(chop!,f)

chop!(f::SumFun,tol::Float64) = SumFun(map(x -> chop!(x,tol),f.funs))
chop!(f::Array,tol::Float64) = map(x -> chop!(x,tol),f)
copy(f::SumFun) = SumFun(copy(f.funs))


using PyPlot

function plot_on_circumference(angularposition, val; ylabel = "", title = "", ylims = [])
    fig, ax = PyPlot.subplots()
    ax.plot(angularposition, val)
    if length(ylims) > 0
        ax.set_ylim(ylims)
    end
    ax.set_ylabel(ylabel)
    ax.set_xlabel("Angular position (deg)")
    ax.grid()
    fig.tight_layout()
    ax.set_title(title)
    return fig
end

function plot_stress_strain(
    angularposition,
    stresscomponent,
    straincomponent;
    ylabels = ["", ""],
    title = "",
)
    fig, ax = PyPlot.subplots(2, 1, sharex = true)
    ax[1].plot(angularposition, stresscomponent)
    ax[1].set_ylabel(ylabels[1])
    ax[1].grid()
    ax[2].plot(angularposition, straincomponent)
    ax[2].set_ylabel(ylabels[2])
    ax[2].grid()
    ax[2].set_xlabel("Angular position (deg)")
    fig.tight_layout()
    fig.suptitle(title)
    return fig
end

function plot_all_stress_strain(angularposition, stress, symmdispgrad, filename)
    fig11 = plot_stress_strain(
        angularposition,
        stress[1, :],
        symmdispgrad[1, :],
        ylabels = [L"\sigma_{11}", L"\epsilon_{11}"],
    )
    fig22 = plot_stress_strain(
        angularposition,
        stress[2, :],
        symmdispgrad[2, :],
        ylabels = [L"\sigma_{22}", L"\epsilon_{22}"],
    )
    fig12 = plot_stress_strain(
        angularposition,
        stress[3, :],
        symmdispgrad[3, :],
        ylabels = [L"\sigma_{12}", L"2\epsilon_{12}"],
    )

    s33 = stress[4, :]
    lowlim = 1.1min(0.0, minimum(s33))
    uplim = 1.1max(0.0, maximum(s33))
    fig33 = plot_on_circumference(
        angularposition,
        s33,
        ylabel = L"\sigma_{33}",
        ylims = [lowlim, uplim],
    )

    fig33.savefig(filename*"33.png")
    fig11.savefig(filename * "11.png")
    fig22.savefig(filename * "22.png")
    fig12.savefig(filename * "12.png")
end

# using PyPlot
# fig,ax = PyPlot.subplots()
# ax.plot(angularposition,productpressure,label="product")
# ax.plot(angularposition,parentpressure,label="parent")
# ax.set_title("Pressure around circumference")
# ax.set_xlabel("Angular position (deg)")
# ax.set_ylabel("Pressure (Pa)")
# ax.grid()
# ax.legend()
# fig.savefig("Pressure-around-circumference.png")


# This test shows the interest to update the QR basis instead
# of computing from scratch at every time step

# function timing()
#     @btime begin
#         m = 1000
#         n = 50
#
#         A = randn(m, n)
#         qr(A)
#         for i=1:20
#             v = randn(m)
#             A = hcat(A, reshape(v,(m,1)))
#         end
#     end
#     @btime begin
#         m = 1000
#         n = 50
#
#         A = randn(m, n)
#         qr(A)
#         for i=1:20
#             v = randn(m)
#             A = hcat(A, reshape(v,(m,1)))
#             qr(A)
#         end
#     end
#     @btime begin
#                 m = 1000
#         n = 50
#
#         A = randn(m, n)
#         F = qr(A)
#         R = F.R
#         for i=1:20
#             v = randn(m)
#             R = qraddcol(A, R, v)
#
#             A = hcat(A, reshape(v,(m,1)))
#         end
#     end
# end
#
# 1.364 ms (86 allocations: 10.18 MiB)
# 17.929 ms (194 allocations: 20.08 MiB)
# 2.417 ms (368 allocations: 11.47 MiB)


# ---
# title: 1619. Mean of Array After Removing Some Elements
# id: problem1619
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/mean-of-array-after-removing-some-elements/description/>
# hidden: true
# ---
#
# Given an integer array `arr`, return _the mean of the remaining integers after
# removing the smallest`5%` and the largest `5%` of the elements._
#
# Answers within `10-5` of the **actual answer** will be considered accepted.
#
#
#
# **Example 1:**
#
#
#
#     Input: arr = [1,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,3]
#     Output: 2.00000
#     Explanation: After erasing the minimum and the maximum values of this array, all elements are equal to 2, so the mean is 2.
#
#
# **Example 2:**
#
#
#
#     Input: arr = [6,2,7,5,1,2,0,3,10,2,5,0,5,5,0,8,7,6,8,0]
#     Output: 4.00000
#
#
# **Example 3:**
#
#
#
#     Input: arr = [6,0,7,0,7,5,7,8,3,4,0,7,8,1,6,8,1,1,2,4,8,1,9,5,4,3,8,5,10,8,6,6,1,0,6,10,8,2,3,4]
#     Output: 4.77778
#
#
# **Example 4:**
#
#
#
#     Input: arr = [9,7,8,7,7,8,4,4,6,8,8,7,6,8,8,9,2,6,0,0,1,10,8,6,3,3,5,1,10,9,0,7,10,0,10,4,1,10,6,9,3,6,0,0,2,7,0,6,7,2,9,7,7,3,0,1,6,1,10,3]
#     Output: 5.27778
#
#
# **Example 5:**
#
#
#
#     Input: arr = [4,8,4,10,0,7,1,3,7,8,8,3,4,1,6,2,1,1,8,0,9,8,0,3,9,10,3,10,1,10,7,3,2,1,4,9,10,7,6,4,0,8,5,1,2,1,6,2,5,0,7,10,9,10,3,7,10,5,8,5,7,6,7,6,10,9,5,10,5,5,7,2,10,7,7,8,2,0,1,1]
#     Output: 5.29167
#
#
#
#
# **Constraints:**
#
#   * `20 <= arr.length <= 1000`
#   * `arr.length` ******is a multiple** of `20`.
#   * `0 <= arr[i] <= 105`
#
#
## @lc code=start
using LeetCode

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


"""
# RandomGeometricGraphs.jl

A small package for the generation of [random geometric graphs](https://en.wikipedia.org/wiki/Random_geometric_graph) in Julia. The package is heavily inspired by the generator used in NetworkX.

It's about 10-15x faster than `LightGraphs.euclidean_graph` because it uses KDTrees from `NearestNeighbors`. It's also about 10x faster than the generator in NetworkX because, y'know, Julia.

## Usage

```julia
using Pkg
Pkg.add(PackageSpec(url="https://github.com/bovine3dom/RandomGeometricGraphs.jl",rev="master"))

using RandomGeometricGraphs
nodes = 1000
radius = 0.2
random_geometric_graph(nodes,radius; dim=2)
```
"""
module RandomGeometricGraphs

    # Todo:
    # test that random_geometric_graph(n,r,dim=dim) ≈ LightGraphs.euclidean_graph(n,dim,cutoff=r)
    # provide seed
    # if positions matrix is not given, read from node attributes
    # consider typing?

    export random_geometric_graph

    import LightGraphs, NearestNeighbors, MetaGraphs

    getpos(g::T,n) where T <: MetaGraphs.AbstractMetaGraph = MetaGraphs.get_prop(g,n,:position)

    function addedgesinradius!(g::T,radius,positions=[]) where T <: MetaGraphs.AbstractMetaGraph
        tree = NearestNeighbors.KDTree(positions)
        for (source, neighbours) in enumerate(NearestNeighbors.inrange(tree, positions, radius))
            for n in neighbours
                if source != n
                    MetaGraphs.add_edge!(g,source,n)
                end
            end
        end
    end

    """
        random_geometric_graph(n,radius;dim=2)

    Generate a random geometric graph on the unit hypercube of dimension `dim` with `n` nodes and links added between nodes who are less than `radius` apart under the euclidean metric.

    """
    function random_geometric_graph(n,radius;dim=2)
        g = MetaGraphs.MetaGraph(n)
        positions = rand(dim,n)
        for i in 1:n
            MetaGraphs.set_prop!(g,i,:position,positions[:,i])
        end
        addedgesinradius!(g,radius,positions)
        g
    end

end


using .SFEM
using BenchmarkTools

A = "G://SMESH//4.cen"
B = "G://SMESH//4.ele"
C = "G://SMESH//4.facein"
D = "G://SMESH//4.faceout"
E = "G://SMESH//4.node"

const EE, NU, FF = 3E7, 0.3, -1000.0
Node, Ele, Cen, Face_out, Face_in = IOsfem(A, B, C, D, E)
@time begin
    KE, Q = Ksfem(Node, Ele, Cen, Face_out, Face_in, EE, NU, FF)
end
#=
@time begin
    Disp = MKL(KE, Q)
end
=#


using JuMP
using FlexJuMP
using Gurobi

# Setup the uncertainty set parameters
means = [620; 388; 583; 313]
a = 5
covar = [11.11 a a a; a 11.11 a a; a a 11.11 a; a a a 11.11]
box_dev = ones(4) * 10

# Setup the model
m = FlexibilityModel(solver = GurobiSolver(OutputFlag = 0))

# Define variables
@randomvariable(m, T[i = 1:4], mean = means[i])
@recoursevariable(m, Qc)
@variable(m, x)

# Define the constraints
@constraint(m, -100 - 0.67Qc + 2T[2] + x <= 0.0)
@constraint(m, -250 - T[2] == x)
@constraint(m, 1388.5 + 0.5Qc - 0.75T[1] - T[2] - T[3] <= 0.0)
@constraint(m, 2044 + Qc - 1.5T[1] - 2T[2] - T[3] <= 0.0)
@constraint(m, 2830 + Qc - 1.5T[1] - 2T[2] - T[3] - 2T[4] <= 0.0)
@constraint(m, -3153 - Qc + 1.5T[1] + 2T[2] + T[3] + 3T[4] <= 0.0)

# Define the uncertainty set
setcovariance(m, covar)
# setuncertaintyset(m, :Ellipsoid)
setuncertaintyset(m, :Hyperbox, [[box_dev]; [box_dev]])
# setuncertaintyset(m, :PNorm, 1)
# setuncertaintyset(m, :PNorm, 2)
# setuncertaintyset(m, :PNorm, Inf)

# Ensure the mean is valid
if ismeanfeasible(m)

    # Can estimate the SF without first solving the flexibility model
    SF = findstochasticflexibility(m, num_pts = 2000, use_vulnerability_model = true)

    # Solve
    status = solve(m, active_constr = true)

    if status == :Optimal
        # Retrieve optimized data
        flexibility_index = getflexibilityindex(m)
        temperatures = getvalue(T)
        cooling = getvalue(Qc)
        state = getvalue(x)
        actives = getactiveconstraints(m)

        # Parse data
        data = getflexibilitydata(m)

        # Get the confidence level
        if data.uncertainty_set.name == :Ellipsoid
            conf_lvl = getconfidencelevel(m)
        end

        # Rank the inequality constraints
        rank_data = rankinequalities(m, max_ranks = 3, active_constr = true)
    end
end


module RegisterUtilities

export Counter

#### Counter ####
#
# Stolen from Grid.jl. Useful when you want to do more math on the iterator.

struct Counter
    max::Vector{Int}
end
Counter(sz::Tuple) = Counter(Int[sz...])
function Base.iterate(c::Counter)
    N = length(c.max)
    (N == 0 || any(c.max.<=0)) && return nothing
    state = ones(Int,N)
    copy(state), state
end
function Base.iterate(c::Counter, state)
    state[1] += 1
    i = 1
    while state[i] > c.max[i] && i < length(state)
        state[i] = 1
        i += 1
        state[i] += 1
    end
    state[end] > c.max[end] && return nothing
    copy(state), state
end

# Below functions are from RegisterTestUtilities

using RegisterCore, LinearAlgebra

export quadratic, block_center, tighten

function quadratic(m, n, shift, Q)
    A = zeros(m, n)
    c = block_center(m, n)
    cntr = [shift[1]+c[1], shift[2]+c[2]]
    u = zeros(2)
    for j = 1:n, i = 1:m
        u[1], u[2] = i-cntr[1], j-cntr[2]
        A[i,j] = dot(u,Q*u)
    end
    A
end

quadratic(shift, Q, denom::Matrix) = MismatchArray(quadratic(size(denom)..., shift, Q), denom)

function block_center(sz...)
    ntuple(i->sz[i]>>1+1, length(sz))
end

function tighten(A::AbstractArray)
    T = typeof(first(A))
    for a in A
        T = promote_type(T, typeof(a))
    end
    At = similar(A, T)
    copyto!(At, A)
end

end


### A Pluto.jl notebook ###
# v0.11.10

using Markdown
using InteractiveUtils

# ╔═╡ b16a96ca-eba3-11ea-3a74-4fecb5dda838
begin
	using FileIO: load
	import LibSndFile
	using Plots
	using SampledSignals
	using FFTW
end

# ╔═╡ cf3ddc9c-ebab-11ea-1417-b7eb9e517b50
md"# Aligning video recordings with Julia"

# ╔═╡ 2c57c0fe-ed5b-11ea-2837-81321cf21495
md"Due to the restrictions imposed by COVID-19, dance teachers around the world are taking their classes online. As a dance student, it would be helpful to watch yourself against a recording of your teacher.

Something like this:"

# ╔═╡ 34a5c400-ebd9-11ea-04cf-3118b3902453
html"<iframe width=\"680\" height=\"383\" src=\"https://www.youtube-nocookie.com/embed/I8kuhv0FcTE\" frameborder=\"0\" allow=\"accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture\" allowfullscreen></iframe>"

# ╔═╡ ae29c670-ebd9-11ea-3995-77063c4c3f4d
md"It is very unlikely the song in both videos starts at exactly the same. Manually syncing them is not easy.

On the other hand, computers can compare the sound waves of each video file and quickly determine how long to wait before playing one video or the other so they are both aligned."

# ╔═╡ eb4bd808-ebc9-11ea-3c71-e9850bb24137
md"## Sound waves"

# ╔═╡ f7a9d13a-ebac-11ea-2abd-79345ce6fbef
md"[Sound waves](https://pudding.cool/2018/02/waveforms/) are represented in digital audio by very rapidly sampling the distortions of the medium (e.g. air) through a microphone and storing the resulting data in a vector."

# ╔═╡ cfdabcf8-ebae-11ea-307a-7728476c9825
md"Jazmine, the teacher, sent Julia a [recording of a jazz song](https://freesound.org/s/37750/) ([cc0](https://creativecommons.org/publicdomain/zero/1.0/)), sampled 16,000 times per second:"

# ╔═╡ 49b7f3f4-eba7-11ea-14f1-31d09ab786bb
x = load("../data/jazz.ogg")

# ╔═╡ 0342bfec-ebc2-11ea-15a2-c769c8039d17
x_duration_seconds = length(x) / x.samplerate

# ╔═╡ c10ff5da-ebae-11ea-3205-6d1f6666ca5c
typeof(x), x.samplerate, length(x)

# ╔═╡ 6e4d19da-ebaf-11ea-113d-55eca1da0769
md"Each sample in $$x$$ corresponds to a microphone reading. The first 5 samples at $$1s$$ are:"

# ╔═╡ c7eccb4a-ebad-11ea-16f5-21299446a8a2
x[16_000:16_005, 1]

# ╔═╡ becebb5c-ebaf-11ea-3e5d-4f91c13b24a1
md"Plotting the audio wave:"

# ╔═╡ 3ab06516-eba5-11ea-2316-732fe86075f6
plot(domain(x), x[:, 1], xlabel="time (s)", ylabel="amplitude", title="Jazz song", linealpha=0.85, legend=false, linecolor=palette(:default)[2], fmt = :png, dpi=300)

# ╔═╡ c0e56d6e-ebb4-11ea-1250-d598737b4049
md"Julia, then danced against the same song and recorded herself near [the sea](https://freesound.org/s/9332/) ([cc-by-nc](https://creativecommons.org/licenses/by-nc/3.0/)):"

# ╔═╡ fdb90c24-ebb6-11ea-1a47-63001e2b8819
x2 = load("../data/jazz-waves.ogg")

# ╔═╡ 997f7a16-ed4a-11ea-3a29-891b059c3667
x2_duration_seconds = length(x2) / x2.samplerate

# ╔═╡ a3c0c930-ed4a-11ea-15e7-cb31b4e9675c
typeof(x2), x2.samplerate, length(x2)

# ╔═╡ 57466fbe-ed5a-11ea-2da0-199907c6bece
md"Let's look at Julia's audio wave:"

# ╔═╡ fba9972c-ebb4-11ea-2a8d-31350dd2855b
plot(domain(x2), x2[:, 1], xlabel="time (s)", ylabel="amplitude", title="Jazz song at the Sea", linealpha=0.85, legend=false, fmt = :png, dpi=300)

# ╔═╡ fa3508de-ed50-11ea-3357-19fd3a4b0441
md"We would like to know how much lag there is between the 2 different recordings of the same song: $$x$$ and $$x2$$."

# ╔═╡ b3997960-ebb0-11ea-2f91-af92117a47a2
md"## Using Cross Correlation to align audio signals"

# ╔═╡ fa4385a0-ebdc-11ea-2683-4107d5129e7e
md"Visually, we could superimpose both signals and slide each other over the time axis until they match as closely as possible.

For example, take Jazmine's recording ($$x$$) and place it next to Julia's ($$x2$$). They are clearly not aligned."

# ╔═╡ 9701b1a8-ed51-11ea-3992-fb42d17e5d0f
begin
	x_right_padding_1 = zeros(Float32, length(x2) - length(x), 1)
	x_padded_1 = cat(x, x_right_padding_1, dims=1)
	mixed_1 = cat(x2, x_padded_1, dims=2)
	plot(domain(mixed_1), mixed_1, xlabel="time (s)", ylabel="amplitude", title="At the sea + Jazz song", linealpha=0.85, legend=false, fmt = :png, dpi=300)
end

# ╔═╡ 07457180-ed51-11ea-148b-5b8c40e3470f
md"Then slide Julia's recording ($$x2$$) $$0.1$$ seconds to the right, then again $$0.1$$ seconds, …

 $$1$$ second later:"

# ╔═╡ 2c083f2a-ed51-11ea-1884-33902424a474
begin
	x_left_padding_2 = zeros(Float32, convert(Int, 1 * x.samplerate), 1)
	x_right_padding_2 = zeros(Float32, length(x2) - length(x) - length(x_left_padding_2), 1)
	x_padded_2 = cat(x_left_padding_2, x, x_right_padding_2, dims=1)
	mixed_2 = cat(x2, x_padded_2, dims=2)
	plot(domain(mixed_2), mixed_2, xlabel="time (s)", ylabel="amplitude", title="At the sea + Jazz song. Offset 1 second", linealpha=0.85, legend=false, fmt = :png, dpi=300)
end

# ╔═╡ 05c3c5a0-ed56-11ea-24fd-2d6628bf80ba
md" $$2$$ seconds later:"

# ╔═╡ 2133853e-ed56-11ea-06a7-9bf1a022192c
begin
	x_left_padding_3 = zeros(Float32, convert(Int, 2 * x.samplerate), 1)
	x_right_padding_3 = zeros(Float32, length(x2) - length(x) - length(x_left_padding_3), 1)
	x_padded_3 = cat(x_left_padding_3, x, x_right_padding_3, dims=1)
	mixed_3 = cat(x2, x_padded_3, dims=2)
	plot(domain(mixed_2), mixed_3, xlabel="time (s)", ylabel="amplitude", title="At the sea + Jazz song. Offset 2 seconds", linealpha=0.85, legend=false, fmt = :png, dpi=300)
end

# ╔═╡ 5871916a-ed56-11ea-2db0-132308b8e7aa
md" $$2.5$$ seconds later both $$x$$ and $$x2$$ are perfectly aligned."

# ╔═╡ 594b2b64-ed56-11ea-104b-cbf7510a12b3
begin
	x_left_padding_4 = zeros(Float32, convert(Int, 2.5 * x.samplerate), 1)
	x_right_padding_4 = zeros(Float32, length(x2) - length(x) - length(x_left_padding_4), 1)
	x_padded_4 = cat(x_left_padding_4, x, x_right_padding_4, dims=1)
	mixed_4 = cat(x2, x_padded_4, dims=2)
	plot(domain(mixed_2), mixed_4, xlabel="time (s)", ylabel="amplitude", title="Ocean Wave + Jazz. Offset 2.5 seconds", linealpha=0.85, legend=false, fmt = :png, dpi=300)
end

# ╔═╡ 07e0ed2c-ed51-11ea-0dd4-cfca2a31506b
md"This sliding process is essentially what a cross-correlation measures. It tells us how similar 2 signals are when displaced relative to each other."

# ╔═╡ 1997e6ea-ebb1-11ea-0875-8ff96ec060dc
md"Mathematically, the [cross-correlation](https://en.wikipedia.org/wiki/Cross-correlation) of two temporal series $$u$$ and $$v$$ is defined as:"

# ╔═╡ 16b55988-ebb1-11ea-3b2b-6f432c3eb7d1
md"$$(u \star v)(t) \triangleq\ \int_{-\infty}^\infty \overline{u(\tau)} v(t + \tau) \, d\tau$$"

# ╔═╡ d4b1188e-ebb9-11ea-3487-bf00755dbd76
md"Where $$\overline{u(\tau)}$$ is the [complex conjugate](https://en.wikipedia.org/wiki/Complex_conjugate) of $$u(\tau)$$. And $$t$$ is the lag."

# ╔═╡ d8b9598a-ebb9-11ea-3f08-7d5e64c688b5
md"The cross-correlation ($$\star$$) of two functions is equivalent to the convolution ($$*$$) of one of the signals reversed in the time axis:

$$(u \star v)(t) =  u(−t) ∗ g (t)$$

Moreover, the Fourier transform ($$\mathcal{F}$$) of a convolution [is equivalent](https://en.wikipedia.org/wiki/Convolution_theorem) to multiplying the Fourier tranforms of the functions:

$$\mathcal{F}\{u * v\} = \mathcal{F}\{u\} \cdot \mathcal{F}\{v\}$$

$$(u \star v)(t) = \mathcal{F}^{-1}\big\{\mathcal{F}\{u(-t)\}\cdot\mathcal{F}\{v(t)\}\big\}$$

As such we can [implement the cross-correlation](https://dsp.stackexchange.com/questions/736/how-do-i-implement-cross-correlation-to-prove-two-audio-files-are-similar) with the following steps:"

# ╔═╡ 2786714e-ebba-11ea-2f24-9f3640dd4b29
md"1. zero padding the signals and reversing one of them:"

# ╔═╡ 608c49d2-ebba-11ea-00ba-c541b8589037
begin
	x_max_size = abs(length(x2) - length(x))

	x_pad_size = length(x) + 2 * length(x_right_padding_1)
	x_pad = SampledSignals.SampleBuf(zeros(Float32, x_pad_size, 1), x.samplerate)
	x_padded = vcat(x, x_pad)
	x_padded_reversed = reverse(x_padded, dims=1)

	x2_pad_size = length(x2)
	x2_pad = SampledSignals.SampleBuf(zeros(Float32, x2_pad_size, 1), x2.samplerate)
	x2_padded = vcat(x2, x2_pad)

	# Listen to this audio at 12 seconds. It is the same Jazz song (x) but reversed and with empty audio at the start
	x_padded_reversed
end

# ╔═╡ 6d193a4c-ebbb-11ea-2ba6-c1341160a063
md"2. Performing a convolution between both signals with their Fourier transforms: "

# ╔═╡ 5cbb311c-ebc2-11ea-3b5e-910d2b940555
convolution = irfft(rfft(x2_padded) .* rfft(x_padded_reversed), length(x_padded))

# ╔═╡ 7deba858-ebbb-11ea-2602-cffad18a927a
plot(domain(x2_padded), convolution[:, 1], xlabel="time (s)", ylabel="convolution(t)", legend=false, fmt = :png, dpi=300)

# ╔═╡ 25a1d500-ebbd-11ea-1c37-35974b7f928d
md"The point where the signals are most correlated corresponds to the peak of the convolution. This is the lag between $$x$$ and $$x2$$:"

# ╔═╡ 979fa4d2-ebbf-11ea-1139-b167cf216ac7
lag_index, peak_value = argmax(convolution[:, 1]), maximum(convolution[:, 1])

# ╔═╡ 7213820c-ebc3-11ea-3c87-d1dcb117d446
md"The lag is exactly 2.5 seconds. We can now use this value to align the audio/video files for Jazmine and Julia."

# ╔═╡ 8863e55c-ebc2-11ea-2495-a962207eacc0
lag_seconds = x_duration_seconds * lag_index / length(x)

# ╔═╡ Cell order:
# ╟─cf3ddc9c-ebab-11ea-1417-b7eb9e517b50
# ╟─2c57c0fe-ed5b-11ea-2837-81321cf21495
# ╟─34a5c400-ebd9-11ea-04cf-3118b3902453
# ╟─ae29c670-ebd9-11ea-3995-77063c4c3f4d
# ╟─eb4bd808-ebc9-11ea-3c71-e9850bb24137
# ╟─f7a9d13a-ebac-11ea-2abd-79345ce6fbef
# ╟─cfdabcf8-ebae-11ea-307a-7728476c9825
# ╟─b16a96ca-eba3-11ea-3a74-4fecb5dda838
# ╠═49b7f3f4-eba7-11ea-14f1-31d09ab786bb
# ╠═0342bfec-ebc2-11ea-15a2-c769c8039d17
# ╠═c10ff5da-ebae-11ea-3205-6d1f6666ca5c
# ╟─6e4d19da-ebaf-11ea-113d-55eca1da0769
# ╠═c7eccb4a-ebad-11ea-16f5-21299446a8a2
# ╟─becebb5c-ebaf-11ea-3e5d-4f91c13b24a1
# ╠═3ab06516-eba5-11ea-2316-732fe86075f6
# ╟─c0e56d6e-ebb4-11ea-1250-d598737b4049
# ╠═fdb90c24-ebb6-11ea-1a47-63001e2b8819
# ╠═997f7a16-ed4a-11ea-3a29-891b059c3667
# ╠═a3c0c930-ed4a-11ea-15e7-cb31b4e9675c
# ╟─57466fbe-ed5a-11ea-2da0-199907c6bece
# ╠═fba9972c-ebb4-11ea-2a8d-31350dd2855b
# ╟─fa3508de-ed50-11ea-3357-19fd3a4b0441
# ╟─b3997960-ebb0-11ea-2f91-af92117a47a2
# ╟─fa4385a0-ebdc-11ea-2683-4107d5129e7e
# ╟─9701b1a8-ed51-11ea-3992-fb42d17e5d0f
# ╟─07457180-ed51-11ea-148b-5b8c40e3470f
# ╟─2c083f2a-ed51-11ea-1884-33902424a474
# ╟─05c3c5a0-ed56-11ea-24fd-2d6628bf80ba
# ╟─2133853e-ed56-11ea-06a7-9bf1a022192c
# ╟─5871916a-ed56-11ea-2db0-132308b8e7aa
# ╟─594b2b64-ed56-11ea-104b-cbf7510a12b3
# ╟─07e0ed2c-ed51-11ea-0dd4-cfca2a31506b
# ╟─1997e6ea-ebb1-11ea-0875-8ff96ec060dc
# ╟─16b55988-ebb1-11ea-3b2b-6f432c3eb7d1
# ╟─d4b1188e-ebb9-11ea-3487-bf00755dbd76
# ╟─d8b9598a-ebb9-11ea-3f08-7d5e64c688b5
# ╟─2786714e-ebba-11ea-2f24-9f3640dd4b29
# ╠═608c49d2-ebba-11ea-00ba-c541b8589037
# ╟─6d193a4c-ebbb-11ea-2ba6-c1341160a063
# ╠═5cbb311c-ebc2-11ea-3b5e-910d2b940555
# ╠═7deba858-ebbb-11ea-2602-cffad18a927a
# ╟─25a1d500-ebbd-11ea-1c37-35974b7f928d
# ╠═979fa4d2-ebbf-11ea-1139-b167cf216ac7
# ╟─7213820c-ebc3-11ea-3c87-d1dcb117d446
# ╠═8863e55c-ebc2-11ea-2495-a962207eacc0


using Embeddings
countries = ["Afghanistan", "Algeria", "Angola", "Arabia", "Argentina", "Australia", "Bangladesh", "Brazil", "Britain", "Canada", "China", "Colombia", "Congo", "Egypt", "England", "Ethiopia", "France", "Germany", "Ghana", "India", "Indonesia", "Iran", "Iraq", "Ireland", "Italy", "Japan", "Kenya", "Korea", "Madagascar", "Malaysia", "Mexico", "Morocco", "Mozambique", "Myanmar", "Nepal", "Nigeria", "Pakistan", "Peru", "Philippines", "Poland", "Russia", "South", "Spain", "Sudan", "Tanzania", "Thailand", "Uganda", "Ukraine", "Usa", "Uzbekistan", "Venezuela", "Vietnam", "Wales", "Yemen"]
usa_cities = ["Albuquerque", "Atlanta", "Austin", "Baltimore", "Boston", "Charlotte", "Chicago", "Columbus", "Dallas", "Denver", "Detroit", "Francisco", "Fresno", "Houston", "Indianapolis", "Jacksonville", "Las", "Louisville", "Memphis", "Mesa", "Milwaukee", "Nashville", "Omaha", "Philadelphia", "Phoenix", "Portland", "Raleigh", "Sacramento", "San", "Seattle", "Tucson", "Vegas", "Washington"]
world_capitals = ["Accra", "Algiers", "Amman", "Ankara", "Antananarivo", "Athens", "Baghdad", "Baku", "Bangkok", "Beijing", "Beirut", "Berlin", "Bogotá", "Brasília", "Bucharest", "Budapest", "Cairo", "Caracas", "Damascus", "Dhaka", "Hanoi", "Havana", "Jakarta", "Kabul", "Kampala", "Khartoum", "Kinshasa", "Kyiv", "Lima", "London", "Luanda", "Madrid", "Manila", "Minsk", "Moscow", "Nairobi", "Paris", "Pretoria", "Pyongyang", "Quito", "Rabat", "Riyadh", "Rome", "Santiago", "Seoul", "Singapore", "Stockholm", "Taipei", "Tashkent", "Tehran", "Tokyo", "Vienna", "Warsaw", "Yaoundé"]
animals = ["alpaca","camel","cattle","dog","dove","duck","ferret","goldfish","goose","rat","llama","mouse","pigeon","yak"]
china_cities = ["Shanghai", "Tianjin", "Shenzhen", "Chengdu", "Changsha", "Nanjing", "Chongqing", "Hefei", "Qinghai", "Jilin", "Heilongjiang"]
sports = ["archery","badminton","basketball","boxing","cycling","diving","equestrian","fencing","field","football","golf","gymnastics","handball","hockey","judo","kayak","pentathlon","polo","rowing","rugby","sailing","shooting","soccer","swimming","taekwondo","tennis","triathlon","volleyball","weightlifting","wrestling"]

words_by_class = [countries, usa_cities, world_capitals, china_cities, animals, sports]
all_words = reduce(vcat, words_by_class)
embedding_table = load_embeddings(Word2Vec; keep_words = all_words)
@assert Set(all_words) == Set(embedding_table.vocab)

embeddings = embedding_table.embeddings
all_words = embedding_table.vocab
classes = map(all_words) do word
    findfirst(col -> word ∈ col, [countries, usa_cities, world_capitals, china_cities, animals, sports])
end;

##
using MultivariateStats
using Plots
plotly()

#Direct projection -- no DR -- just throw away the information in the other axies
xs=embeddings
scatter(xs[1,:], xs[2,:], xs[3,:]; hover=all_words, zcolor=classes, color = :viridis, legend = false)

##
M = fit(PCA, embeddings; maxoutdim=3)
xs = transform(M, embeddings)
scatter(xs[1,:], xs[2,:], xs[3,:]; hover=all_words, zcolor=classes, color = :viridis, legend = false)

M = fit(ICA, embeddings, 3; maxiter = 1000, tol = 1)
xs = transform(M, embeddings)
scatter(xs[1,:], xs[2,:], xs[3,:]; hover=all_words, zcolor=classes, color = :viridis, legend = false)

using TSne
xs = tsne(embeddings', 3, 3, 1000, 20.0)'
scatter(xs[1,:], xs[2,:], xs[3,:]; hover=all_words, zcolor=classes, color = :viridis, legend = false)


# taken from https://diffeq.sciml.ai/stable/features/noise_process/

using StochasticDiffEq,  DiffEqNoiseProcess
using Test

f1(u, p, t) = 1.01u
g1(u, p, t) = 1.01u
dt = 1//2^(4)
prob1 = SDEProblem(f1,g1,1.0,(0.0,1.0))
sol1 = solve(prob1,EM(),dt=dt,save_noise = true)

W2 = NoiseWrapper(sol1.W)
prob1 = SDEProblem(f1,g1,1.0,(0.0,1.0),noise=W2)
sol2 = solve(prob1,EM(),dt=dt)

@test sol1.u ≈ sol2.u

W3 = NoiseWrapper(sol1.W)
prob2 = SDEProblem(f1,g1,1.0,(0.0,1.0),noise=W3)

# smaller timestep
dt = 1//2^(5)
sol3 = solve(prob2,EM(),dt=dt)

using Plots
plot(sol1)
plot!(sol2)
plot!(sol3)


export examples, example

const Example = NamedTuple{(:name, :inv, :vars, :kwargs),Tuple{Symbol, Invariant, Vector{Symbol}, Any}}

_examples = Dict{Symbol,Example}()

example!(x::Example) = push!(_examples, x.name=>x)
examples() = values(_examples) |> collect
example(s) = _examples[s]

# Double
# x, y = 0, 1
# while true
#     x = 2x
#     y = 1/2 y + 1
# end
example!(Example((
    :double1,
    @invariant(x*y == 2x),
    [:x, :y],
    Dict()
)))
example!(Example((
    :double2,
    @invariant(x == 2y),
    [:x, :y],
    Dict()
)))

# Square
# a, b = 0, 0
# while true
#     a = a + 2b + 1
#     b = b + 1
# end
example!(Example((
    :square,
    @invariant(a == b^2),
    [:a, :b],
    Dict()
)))

example!(Example((
    :dblsquare,
    @invariant(a == 2*b^2),
    [:a, :b],
    Dict()
)))

# Sum1
# a, b, c = 0, 0, 1
# while true
#     a = a + 1
#     b = b + c
#     c = c + 2
# end
example!(Example((
    :sum1,
    @invariant(1 + 2a == c && 4b == (c-1)^2),
    [:a, :b, :c],
    Dict()
)))

# Sum2
# a, b, c, s = 0, 0, 1, 0
# while true
#     a = a + 1
#     b = b + c
#     c = c + 2
#     s = s + 2a + 1
# end
example!(Example((
    :sum2,
    @invariant(1 + 2a == c && b + c == 1 + s && c*(c+2) == 3 + 4s),
    [:s, :a, :b, :c],
    Dict()
)))

# eucliddiv
# r, q = x, 0
# while true
#     r = r - y
#     q = q + 1
# end
example!(Example((
    :eucliddiv,
    @invariant(x(0) == y(0)*q(n) + r(n)),
    [:r, :q, :x, :y],
    Dict(:params=>[:x, :y])
)))

# Integer Square Root - version 1
# k, j, m = 0, 1, 1
# while m<=n
#     k = k + 1
#     j = j + 2
#     m = m + j
# end
example!(Example((
    :intsqrt1,
    @invariant(j == 1 + 2k && (j+1)^2 == 4m),
    [:m, :k, :j],
    Dict()
)))

# Integer Square Root - version 2
# y, r = 1/2*a, 0
# while true
#     y = y - r
#     r = r + 1
# end
example!(Example((
    :intsqrt2,
    @invariant(a(0) + r == r^2 + 2y),
    [:y, :r, :a],
    Dict(:params=>[:a])
)))

# Integer Cubic Root
# x, r, s = a, 1, 13/4
# while true
#     x = x-s
#     s = s + 6r + 3
#     r = r + 1
# end
example!(Example((
    :intcbrt,
    @invariant(1/4 + 3r^2 == s && 1 + 4*a(0) + 6r^2 == 3r + 4r^3 + 4x),
    [:x, :s, :r, :a],
    Dict(:params=>[:a])
)))

# Consecutive Cubes
# n, x, y, z = 0, 0, 1, 6
# while true
#     x = x + y
#     y = y + z
#     z = z + 6
#     n = n + 1
# end
example!(Example((
    :cubes,
    @invariant(n^3 == x && 1 + 3n + 3n^2 == y && 6 + 6n == z),
    [:x, :y, :z, :n],
    Dict()
)))

# Petter 1
# x, y = 0, 0
# while true
#     x = x + y^1
#     y = y + 1
# end
example!(Example((
    :petter1,
    @invariant(y^2 == 2*x+y),
    [:x, :y],
    Dict()
)))

# Petter 2
# x, y = 0, 0
# while true
#     x = x + y^2
#     y = y + 1
# end
# 6 x == y * (-1 + 3 y - 2 y^2)

# Petter 3
# x, y = 0, 0
# while true
#     x = x + y^3
#     y = y + 1
# end
# 4 x - (-1 + y)^2 y^2

# Add (https://rise4fun.com/Dafny/Add)
# r = x
# n = y
# while n != 0
#     r = r + 1
#     n = n - 1
# end
example!(Example((
    :add1,
    @invariant(r(m) == x(0)+y(0)-n(m)),
    [:r, :n, :x, :y],
    Dict(:params=>[:x, :y])
)))

# r = 2x
# n = y
# while n != 0
#     r = r + 1
#     n = n - 1
# end
example!(Example((
    :add2,
    @invariant(r(m) == 2*x(0)+y(0)-n(m)),
    [:r, :n, :x, :y],
    Dict(:params=>[:x, :y])
)))

# ------------------------------------------------------------------------------

function Base.run(examples::Vector{Example}; solvers=[Yices,Z3], timeout=2, maxsol=1, shapes=[UnitUpperTriangular, UpperTriangular, FullSymbolic])
    progress = ProgressUnknown("Instances completed:")
    df = DataFrame(Solver = Type{<:NLSolver}[], Instance = Any[], Shape = MatrixShape[], Status = NLStatus[], Elapsed = Any[], Result = Union{RecSystem,Nothing}[])
    for ex in examples
        for shape in shapes
            strat = strategy_fixed(ex.inv, copy(ex.vars), shape, [length(ex.vars)+1]; constone=true, ex.kwargs...)
            for solver in solvers
                sols = solutions(strat; solver=solver, timeout=timeout, maxsol=maxsol)
                for res in sols
                    push!(df, (solver, ex.name, shape, res.status, res.elapsed, res.recsystem))
                    # push!(df, (solver, ex.name, shape, res.status, res.status == NLSat.timeout ? "-" : res.elapsed, res.recsystem))
                    next!(progress)
                end
            end
        end
    end
    finish!(progress)
    df
end

function wide(df)
    key = [:Instance, :Shape]
    tab = unstack(df, key, :Solver, :Elapsed)
    z3 = unstack(tab, :Instance, :Shape, :Z3Solver)
    yices = unstack(tab, :Instance, :Shape, :YicesSolver)
    rename!(z3,
        :UnitUpperTriangular => :Z3_UnitUpperTriangular,
        :UpperTriangular     => :Z3_UpperTriangular,
        :FullSymbolic        => :Z3_FullSymbolic
    )
    rename!(yices,
        :UnitUpperTriangular => :Yices_UnitUpperTriangular,
        :UpperTriangular     => :Yices_UpperTriangular,
        :FullSymbolic        => :Yices_FullSymbolic
    )
    join(yices, z3, on = :Instance, makeunique = true)
end


"""
Three greedy-based additive heuristics for the **Diffusion on Hypergraphs** problem,
i.e. finding the minimum influence target set *S in V* of a hypergraph *H=(V,E)*
able to influence the whole network.
"""

"""
    bisect(h, metaV, metaE)

Compute the target set exploiting a binary search approach.

Also referred to as *StaticGreedy*.
"""
function bisect(h, metaV, metaE; opt=false)
    degrees = length.(h.v2he)
    sortedVind = sortperm(degrees, rev=true)

    left = 1
    @assert length(sortedVind) == length(degrees)
    right = length(sortedVind)
    while left < right
        location = ceil(Int, (left + right)/2)
        actE = zeros(Bool, nhe(h))
        actV = zeros(Bool, nhv(h))
        # up to the first location true and remainder zeros
        actV[sortedVind[1:location]] .= true
        sres = simulate!(h, actV, actE, metaV, metaE; printme = false)
        if sres.actvs != nhv(h)
            left = location
        else
            right = location-1
        end
    end

    if opt
        new_seed = optimize_seed_set(h, Set(sortedVind[1:(left+1)]), metaV, metaE)
        return length(new_seed)
    end

    left+1
end

"""
    greedy_tss(H, metaV, metaE)

Compute the target set according to the following algorithm.

Initially, all nodes are added to the candidates set *U*.
At each stage, the node of maximum degree is added to the target set *S* and removed from *U*.
At this point, some nodes and/or hyperedges become infected.

The algorithm simulates the diffusion process, and the influenced edges are pruned from the network.
Consequentially, the degree of nodes δ(v) is updated.

Also referred to as *DynamicGreedy(H)*.
"""
function greedy_tss(h, metaV, metaE; opt=false)
    S = Dict{Int,Int}()
    U = Dict{Int,Int}()

    #init U set
    for v=1:nhv(h)
        heus = gethyperedges(h,v)
        d = 0
        for he in heus
            d += 1
        end
        push!(U, v => d)
    end

    while length(U) != 0
        #maxv = sort!(collect(U), by=x->x[2], rev = true)[1]
        maxv_val, maxv_key = findmax(U)
        delete!(U, maxv_key)

        S[maxv_key] = maxv_val

        actE = zeros(Bool, nhe(h))
        actV = zeros(Bool, nhv(h))

        for s in S
        actV[s.first] = true
        end

        simres = simulate!(h, actV, actE, metaV, metaE; printme = false)

        if simres.actvs == nhv(h)
            break
        end

         # comment this for to implement
         # a static greedy approach
        for he in gethyperedges(h,maxv_key)
            for nv in getvertices(h,he.first)

                !haskey(U, nv.first) && continue

                d = 0
                for he2 in gethyperedges(h, nv.first)
                    if ! actE[he2.first]
                        d+=1
                    end
                end

                push!(U, nv.first => d)
            end
        end
    end

    actE = zeros(Bool, nhe(h))
    actV = zeros(Bool, nhv(h))

    for s in S
        actV[s.first] = true
    end

    simres = simulate!(h, actV, actE, metaV, metaE; printme = false)

    if simres.actvs != nhv(h)
        println("ARGGGG something wrong", simres.actvs, " ", nhv(h))
    end

    if opt
        new_seed = optimize_seed_set(h, Set(keys(S)), metaV, metaE)
        return length(new_seed)
    end

    length(S)
end

"""
    greedy_tss_2section(H, metaV, metaE)

Compute the target set according to the following algorithm.

The degree of nodes is computed on the [H]₂ of the residual hypergraph Hⁱ of *H*.
Hⁱ is the hypergraph obtained by removing all hyperedges that are already influenced
by the nodes in *S* at stage *i*.

Also referred to as *DynamicGreedy([H]₂)*.
"""
function greedy_tss_2section(h, metaV, metaE; opt=false)
    S = Dict{Int,Int}()
    U = Dict{Int,Int}()

    #init U set
    for v=1:nhv(h)
        heus = gethyperedges(h,v)
        d = Set{Int}()
        for he in heus
            push!.(Ref(d),keys(getvertices(h,he.first)))
        end
        push!(U, v => length(d))
    end

    while length(U) != 0
        #maxv = sort!(collect(U), by=x->x[2], rev = true)[1]
        maxv_val, maxv_key = findmax(U)

        delete!(U, maxv_key)
        S[maxv_key] = maxv_val

        actE = zeros(Bool, nhe(h))
        actV = zeros(Bool, nhv(h))

        for s in S
            actV[s.first] = true
        end
        simres = simulate!(h, actV, actE, metaV, metaE; printme = false)

        if simres.actvs == nhv(h)
        break
        end

        visited = Set{Int}()
        for he in gethyperedges(h, maxv_key)

            for nv in getvertices(h,he.first)

                (!haskey(U,nv.first) || nv.first in visited) && continue

                push!(visited, nv.first)

                d = Set{Int}()
                for he2 in gethyperedges(h, nv.first)
                    if ! actE[he2.first]
                        push!.(Ref(d), keys(getvertices(h, he2.first)))
                    end
                end
                U[nv.first] = length(d)
            end
        end
    end

    if opt
        new_seed = optimize_seed_set(h, Set(keys(S)), metaV, metaE)
        return length(new_seed)
    end

    length(S)
end


using EvalMetrics
using Plots; pyplot()
using Random, MLBase; Random.seed!(42);
scores = sort(rand(10000));
targets = scores .>= 0.99;
targets[MLBase.sample(findall(0.98 .<= scores .< 0.99), 30; replace = false)] .= true;
targets[MLBase.sample(findall(0.99 .<= scores .< 0.995), 30; replace = false)] .= false;

savefig(prplot(targets, scores), "pr1.png")
savefig(prplot([targets, targets], [scores, scores .+ rand(10000) ./ 5]), "pr2.png")
savefig(rocplot(targets, scores), "roc1.png")
savefig(rocplot([targets, targets], [scores, scores .+ rand(10000) ./ 5]), "roc2.png")

savefig(prplot(targets, scores; xguide="RECALL", fill=:green, grid=false, xlims=(0.8, 1.0)), "pr3.png")
savefig(rocplot(targets, scores, title="Title", label="experiment", xscale=:log10), "roc3.png")
savefig(rocplot([targets, targets], [scores, scores .+ rand(10000) ./ 5], label=["a" "b";]), "roc4.png")

prplot(targets, scores; npoints=Inf, label="Original")
prplot!(targets, scores; npoints=10, label="Sampled (10 points)")
prplot!(targets, scores; npoints=100, label="Sampled (100 points)")
prplot!(targets, scores; npoints=1000, label="Sampled (1000 points)")
prplot!(targets, scores; npoints=5000, label="Sampled (5000 points)")
savefig("pr4.png")

savefig(rocplot(targets, scores; aucshow=false, label="a", diagonal=true), "roc5.png")


using LightGraphs
include("../src/utils.jl")

"""
    lattice_nearest_neighbor(Lx::Int,Ly::Int)

lattice with nearest neighbor interactions, where nearest neighbor here
also includes diagonal neighbors.
"""
function lattice_nearest_neighbor(Lx::Int,Ly::Int)
    N=Lx*Ly
    lattice=SimpleGraph(N)
    for y in 1:Ly
        for x in 1:Lx
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x-1,y+1,Lx,Ly))
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x  ,y+1,Lx,Ly))
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x+1,y+1,Lx,Ly))
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x-1,y  ,Lx,Ly))

            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x+1,y  ,Lx,Ly))
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x-1,y-1,Lx,Ly))
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x  ,y-1,Lx,Ly))
            add_edge!(lattice, pbc_index(x,y,Lx,Ly), pbc_index(x+1,y-1,Lx,Ly))
        end
    end
    return lattice
end

"""
    step!(rng, state, state_new, lattice, p_rec[, state_rev, val_ref])

recurrent update step with stochastic activations of `state_new` from `state`
based on interactions defined in `lattice` with probability `p_rec` using
random number generator `rng`.

External activations can be added to `state_new` before calling `step!`.

Optionally, one can pass the refractory state `state_rev` that is reinitalized
upon activation to `val_ref`, decreases one value per time step, and makes
neurons only excitable if `state_ref[i]>0`.
"""
function step!(
        rng::AbstractRNG,
        state::Array{Bool},
        state_new::Array{Bool},
        lattice::SimpleGraph{Int64},
        p_rec
    )
    prepare!(rng, state, state_new, lattice, p_rec)

    # swap
    state .= state_new
    state_new .= false
    return nothing
end
function step!(
        rng::AbstractRNG,
        state::Array{Bool},
        state_new::Array{Bool},
        lattice::SimpleGraph{Int64},
        p_rec,
        state_ref::Array{Int},
        val_ref::Int,
    )
    prepare!(rng, state, state_new, lattice, p_rec)

    # if still in refractory period then no update
    for i in 1:length(state)
        if state_ref[i] > 0
            state_new[i] = false
        end
    end

    # swap
    state .= state_new
    state_new .= false

    # update refractory period
    for i in 1:length(state)
        if state[i]
            state_ref[i] = val_ref
        else
            state_ref[i] -= 1
        end
    end
    return nothing
end
function prepare!(rng, state,state_new, lattice, p_rec)
    for i in 1:length(state)
        if state[i]==true
            for j in outneighbors(lattice, i)
                if rand(rng) < p_rec
                    state_new[j] = true
                end
            end
        end
    end
end


using Flux, DiffEqFlux, DiffEqSensitivity, StochasticDiffEq, Zygote, Test

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end
function lotka_volterra(u,p,t)
  x, y = u
  α, β, δ, γ = p
  dx = α*x - β*x*y
  dy = -δ*y + γ*x*y
  [dx,dy]
end
function lotka_volterra_noise(du,u,p,t)
  du[1] = 0.01u[1]
  du[2] = 0.01u[2]
end
function lotka_volterra_noise(u,p,t)
  [0.01u[1],0.01u[2]]
end
prob = SDEProblem(lotka_volterra,lotka_volterra_noise,[1.0,1.0],(0.0,10.0))
p = [2.2, 1.0, 2.0, 0.4]
function predict_fd_sde(p)
  solve(prob,SOSRI(),p=p,saveat=0.0:0.1:0.5,sensealg=ForwardDiffSensitivity())[1,:]
end
loss_fd_sde(p) = sum(abs2,x-1 for x in predict_fd_sde(p))
loss_fd_sde(p)

prob = SDEProblem{false}(lotka_volterra,lotka_volterra_noise,[1.0,1.0],(0.0,10.0))
p = [2.2, 1.0, 2.0, 0.4]
function predict_fd_sde(p)
  solve(prob,SOSRI(),p=p,saveat=0.0:0.1:0.5,sensealg=ForwardDiffSensitivity())[1,:]
end
loss_fd_sde(p) = sum(abs2,x-1 for x in predict_fd_sde(p))
loss_fd_sde(p)

@test !iszero(Zygote.gradient(loss_fd_sde,p)[1])

prob = SDEProblem(lotka_volterra,lotka_volterra_noise,[1.0,1.0],(0.0,0.5))
function predict_rd_sde(p)
  solve(prob,SOSRI(),p=p,saveat=0.0:0.1:0.5,sensealg=TrackerAdjoint())[1,:]
end
loss_rd_sde(p) = sum(abs2,x-1 for x in predict_rd_sde(p))
@test !iszero(Zygote.gradient(loss_rd_sde,p)[1])

prob = SDEProblem{false}(lotka_volterra,lotka_volterra_noise,[1.0,1.0],(0.0,0.5))
function predict_rd_sde(p)
  solve(prob,SOSRI(),p=p,saveat=0.0:0.1:0.5,sensealg=TrackerAdjoint())[1,:]
end
loss_rd_sde(p) = sum(abs2,x-1 for x in predict_rd_sde(p))
@test !iszero(Zygote.gradient(loss_rd_sde,p)[1])


using Gadfly

set_default_plot_size(6inch, 3inch)

spy(randn((10, 10)), Coord.cartesian(fixed=true))


function normalize_theta!(scales::AbstractArray, θ::AbstractArray)
    @assert length(scales) == size(θ, 1)
    @inbounds for (i, ti) in enumerate(eachrow(θ))
        scales[i] = norm(ti, 2)
        normalize!(ti, 2)
    end
    return
end

function rescale_xi!(Ξ::AbstractArray, scales::AbstractArray)
    @assert length(scales) == size(Ξ, 1)
    @inbounds for (si, ti) in zip(scales, eachrow(Ξ))
        ti .= ti ./ si
    end
    return
end

function rescale_theta!(θ::AbstractArray, scales::AbstractArray)
    @assert length(scales) == size(θ, 1)
    @inbounds for (i, ti) in enumerate(eachrow(θ))
        ti .= ti .* scales[i]
    end
    return
end

"""
    sparse_regression(X, Y, basis, p, t, maxiter, opt, denoise, normalize, convergence_error)
    sparse_regression!(Xi, X, Y, basis, p, t, maxiter, opt, denoise, normalize, convergence_error)
    spares_regression!(Xi, Theta, Y, maxiter, opt, denoise, normalize, convergence_error)

Performs a sparse regression via the algorithm `opt <: AbstractOptimizer`. `maxiter` specifies the upper bound of the iterations
of the optimizer, `convergence_error` the breaking condition due to convergence.
`denoise` defines if the matrix holding candidate trajectories should be thresholded via the [optimal threshold for singular values](http://arxiv.org/abs/1305.5870).
`normalize` normalizes the matrix holding candidate trajectories via the L2-Norm over each function.

If the data matrices `X`, `Y` are given with a `Basis` `basis` and the additional information for parameters `p` and time points of the
measurements `t`, it returns the coefficient matrix `Xi` and the iterations taken.
This function is also available in place, which returns just the iterations.

If `Xi`, `Theta` and `Y` are given, the sparse regression will find the coefficients `Xi`, which minimize the objective and return the iterations needed.

# Example

```julia
opt = STRRidge()
maxiter = 10
c_error = 1e-3

Xi, iters = sparse_regression(X, Y, basis, [], [], maxiter, opt, false, false, c_error)

iters = sparse_regression!(Xi,X, Y, basis, [], [], maxiter, opt, false, false, c_error)

Xi2 = zeros(size(Y, 1), size(X, 1))
iters = sparse_regression!(Xi2, X, Y, maxiter, opt, false, false, c_error)
```
"""
function sparse_regression(X::AbstractArray, Ẋ::AbstractArray, Ψ::Basis, p::AbstractArray, t::AbstractVector , maxiter::Int64 , opt::T, denoise::Bool, normalize::Bool, convergence_error) where T <: Optimize.AbstractOptimizer
    @assert size(X)[end] == size(Ẋ)[end]
    nx, nm = size(X)
    ny, nm = size(Ẋ)

    Ξ = zeros(eltype(X), length(Ψ), ny)
    scales = ones(eltype(X), length(Ψ))
    θ = zeros(eltype(X), length(Ψ), nm)
    Ψ(θ, X, p, t)

    denoise ? optimal_shrinkage!(θ') : nothing
    normalize ? normalize_theta!(scales, θ) : nothing

    Optimize.init!(Ξ, opt, θ', Ẋ')
    iters = Optimize.fit!(Ξ, θ', Ẋ', opt, maxiter = maxiter, convergence_error = convergence_error)

    normalize ? rescale_xi!(Ξ, scales) : nothing

    return Ξ, iters
end

function sparse_regression!(Ξ::AbstractArray, X::AbstractArray, Ẋ::AbstractArray, Ψ::Basis, p::AbstractArray , t::AbstractVector, maxiter::Int64 , opt::T, denoise::Bool, normalize::Bool, convergence_error) where T <: Optimize.AbstractOptimizer
    @assert size(X)[end] == size(Ẋ)[end]
    nx, nm = size(X)
    ny, nm = size(Ẋ)
    @assert size(Ξ) == (length(Ψ), ny)

    scales = ones(eltype(X), length(Ψ))
    θ = zeros(eltype(X), length(Ψ), nm)
    Ψ(θ, X, p, t)

    denoise ? optimal_shrinkage!(θ') : nothing
    normalize ? normalize_theta!(scales, θ) : nothing

    Optimize.init!(Ξ, opt, θ', Ẋ')
    iters = Optimize.fit!(Ξ, θ', Ẋ', opt, maxiter = maxiter, convergence_error = convergence_error)

    normalize ? rescale_xi!(Ξ, scales) : nothing

    return iters
end

# For pareto
function sparse_regression!(Ξ::AbstractArray, θ::AbstractArray, Ẋ::AbstractArray, maxiter::Int64 , opt::T, denoise::Bool, normalize::Bool, convergence_error) where T <: Optimize.AbstractOptimizer

    scales = ones(eltype(Ξ), size(θ, 1))

    denoise ? optimal_shrinkage!(θ') : nothing
    normalize ? normalize_theta!(scales, θ) : nothing

    Optimize.init!(Ξ, opt, θ', Ẋ')'
    iters = Optimize.fit!(Ξ, θ', Ẋ', opt, maxiter = maxiter, convergence_error = convergence_error)

    normalize ? rescale_xi!(Ξ, scales) : nothing
    normalize ? rescale_theta!(θ, scales) : nothing

    return iters
end

# One Variable on multiple derivatives
function SINDy(X::AbstractArray{S, 1}, Ẋ::AbstractArray, Ψ::Basis, opt::T = STRRidge(); kwargs...) where {T <: Optimize.AbstractOptimizer, S <: Number}
    return SINDy(X', Ẋ, Ψ, opt; kwargs...)
end

# Multiple on one
function SINDy(X::AbstractArray{S, 2}, Ẋ::AbstractArray{S, 1}, Ψ::Basis, opt::T = STRRidge(); kwargs...) where {T <: Optimize.AbstractOptimizer, S <: Number}
    return SINDy(X, Ẋ', Ψ, opt; kwargs...)
end

# General
"""
    SINDy(X, Y, basis, opt = STRRidge(); p, t, maxiter, convergence_error, denoise, normalize)
    SINDy(X, Y, basis, lambdas, opt = STRRidge(); f, g, p, t, opt, maxiter, convergence_error, denoise, normalize)

Performs Sparse Identification of Nonlinear Dynamics given the data matrices `X` and `Y` via the `AbstractBasis` `basis.`
Keyworded arguments include the parameter (values) of the basis `p` and the timepoints `t` which are passed in optionally.
`opt` is an `AbstractOptimizer` useable for sparse regression, `maxiter` the maximum iterations to perform and `convergence_error` the
bound which causes the optimizer to stop.
`denoise` defines if the matrix holding candidate trajectories should be thresholded via the [optimal threshold for singular values](http://arxiv.org/abs/1305.5870).
`normalize` normalizes the matrix holding candidate trajectories via the L2-Norm over each function.

Typically `X` represent the state measurements and `Y` the measurements of the differential state. Since only the number of measurements (column dimension of the matrices) have to be equal, it is possible to augment `X` with additional
data, e.g. external forcing or inputs.

If `SINDy` is called with an additional array of thresholds contained in `lambdas`, it performs a multi objective optimization over all thresholds.
The best candidate is determined via the mapping onto a feature space `f` and an (scalar, positive definite) evaluation `g`.
The signature of should be `f(xi, theta, yi)` where `xi` are the coefficients of the sparse optimization,`theta` is the evaluated candidate library and `yi` are the rows of the matrix `Y`.
Returns a `SINDyResult`. If the pareto optimization is used, the result combines the best candidate for each row of `Y`.
"""
function SINDy(X::AbstractArray{S, 2}, Ẋ::AbstractArray{S, 2}, Ψ::Basis, opt::T = STRRidge(); p::AbstractArray = [], t::AbstractVector = [], maxiter::Int64 = 10, denoise::Bool = false, normalize::Bool = true, convergence_error = eps()) where {T <: Optimize.AbstractOptimizer, S <: Number}
    Ξ, iters = sparse_regression(X, Ẋ, Ψ, p, t, maxiter, opt, denoise, normalize, convergence_error)
    convergence = iters < maxiter
    SparseIdentificationResult(Ξ, Ψ, iters, opt, convergence, Ẋ, X, p = p)
end

function SINDy(X::AbstractArray{S, 1}, Ẋ::AbstractArray, Ψ::Basis, thresholds::AbstractArray, opt::T = STRRidge(); kwargs...) where {T <: Optimize.AbstractOptimizer, S <: Number}
    return SINDy(X', Ẋ, Ψ, thresholds, opt; kwargs...)
end

function SINDy(X::AbstractArray{S, 2}, Ẋ::AbstractArray{S, 1}, Ψ::Basis, thresholds::AbstractArray, opt::T = STRRidge()::AbstractArray; kwargs...) where {T <: Optimize.AbstractOptimizer, S <: Number}
    return SINDy(X, Ẋ', Ψ, thresholds, opt; kwargs...)
end

function SINDy(X, DX, Ψ::Basis, thresholds::AbstractArray, opt::T = STRRidge(); f::Function = (xi, theta, dx)->[norm(xi, 0); norm(dx .- theta'*xi, 2)], g::Function = x->norm(x), p::AbstractArray = [], t::AbstractVector = [], maxiter::Int64 = 10, denoise::Bool = false, normalize::Bool = true, convergence_error = eps()) where {T <: Optimize.AbstractOptimizer}
    @assert size(X)[end] == size(DX)[end]
    nx, nm = size(X)
    ny, nm = size(DX)

    θ = Ψ(X, p, t)

    scales = ones(eltype(X), length(Ψ))

    ξ_tmp = zeros(eltype(X), length(Ψ), ny)
    Ξ_opt = zeros(eltype(X), length(Ψ), ny)

    iter_ = Inf
    _iter = Inf
    _thresh = Inf

    denoise ? optimal_shrinkage!(θ') : nothing
    normalize ? normalize_theta!(scales, θ) : nothing

    fg(xi, theta, dx) = (g∘f)(xi, theta, dx)

    @inbounds for j in 1:length(thresholds)
        set_threshold!(opt, thresholds[j])
        iter_ = sparse_regression!(view(ξ_tmp, :, :), view(θ, :, :), view(DX, :, :), maxiter, opt, false, false, convergence_error)
        normalize ? rescale_xi!(ξ_tmp, scales) : nothing
        for i in 1:ny
            if j == 1
                Ξ_opt .= ξ_tmp
                _iter = iter_
                _thresh = thresholds[j]
            else
                if evaluate_pareto!(view(Ξ_opt, :, i), view(ξ_tmp, :, i), fg, view(θ, :, :), view(DX, i, :))
                    _iter = iter_
                    _thresh = thresholds[j]
                end
            end
        end
    end

    set_threshold!(opt, _thresh)
    return SparseIdentificationResult(Ξ_opt, Ψ, _iter, opt, _iter < maxiter, DX, X, p = p)
end