### A Pluto.jl notebook ###
# v0.17.5

using Markdown
using InteractiveUtils

# ╔═╡ 4e046622-7300-11ec-0a1d-974094a4752e
# hideall
let
    docs_dir = dirname(dirname(@__DIR__))
    pkg_dir = dirname(docs_dir)

    using Pkg: Pkg
    Pkg.activate(docs_dir)
    Pkg.develop(; path=pkg_dir)
    Pkg.instantiate()

    # Putting the include here to avoid Pluto getting confused about cell order.
    include(joinpath(docs_dir, "src", "tutorials_utils.jl"))
end;

# ╔═╡ 544dfb44-b8ae-4c2c-a03b-3e14c88c1d66
using CSV

# ╔═╡ d6fa288d-acba-49a9-ba58-c7a871cca5ad
using DataFrames

# ╔═╡ 3061ae4d-49c4-4275-9505-f75aa8be23a7
using TuringGLM

# ╔═╡ 89586742-fbe8-4213-90d4-204f6b15d514
md"""
For our tutorial on **Logistic Regression**, let's use a famous dataset called `wells` (Gelman & Hill, 2007),
which is data from a survey of 3,200 residents in a small area of Bangladesh suffering from arsenic contamination of groundwater.
Respondents with elevated arsenic levels in their wells had been encouraged to switch their water source to a safe public or private well in the nearby area and the survey was conducted several years later to learn which of the affected residents had switched wells.
It has 3,200 observations and the following variables:

* `switch` -- binary/dummy (0 or 1) for well-switching.
* `arsenic` -- arsenic level in respondent's well.
* `dist` -- distance (meters) from the respondent's house to the nearest well with safe drinking water.
* `association` -- binary/dummy (0 or 1) if member(s) of household participate in community organizations.
* `educ` -- years of education (head of household).
"""

# ╔═╡ e0cf019d-2c2e-470b-a62a-25fbb9a22b4d
url = "https://github.com/TuringLang/TuringGLM.jl/raw/main/data/wells.csv";

# ╔═╡ 0de64701-70cf-4b1c-ae5a-e288e34a5e3c
wells = CSV.read(download(url), DataFrame)

# ╔═╡ ffb09bcb-3d02-455c-b153-345f4857fe9d
md"""
Using `switch` as dependent variable and `dist`, `arsenic`, `assoc`, and `educ` as independent variables:
"""

# ╔═╡ 07175d0e-7293-4095-95cc-91fe493a1aef
fm = @formula(switch ~ dist + arsenic + assoc + educ)

# ╔═╡ 2ebfa422-f8a5-44d3-8f2c-a34d7832d3f2
md"""
Now we instantiate our model with `turing_model` passing a third argument `Logistic()` to indicate that the model is a logistic regression:
"""

# ╔═╡ 1f1158ce-73f0-49fd-a48a-af3b36376030
model = turing_model(fm, wells, Logistic());

# ╔═╡ ebe074cb-7bad-4b52-9c9d-b9752af4bedd
chn = sample(model, NUTS(), 2_000);

# ╔═╡ d6499c1d-c2a1-4cc0-a8f4-c4ba86d5cf30
# hide
plot_chains(chn)

# ╔═╡ 2361e758-1b5b-4cb2-b7d5-8a03760befa6
md"""
## References

Gelman, A., & Hill, J. (2007). Data analysis using regression and multilevel/hierarchical models. Cambridge university press.
"""

# ╔═╡ Cell order:
# ╠═4e046622-7300-11ec-0a1d-974094a4752e
# ╠═89586742-fbe8-4213-90d4-204f6b15d514
# ╠═544dfb44-b8ae-4c2c-a03b-3e14c88c1d66
# ╠═d6fa288d-acba-49a9-ba58-c7a871cca5ad
# ╠═3061ae4d-49c4-4275-9505-f75aa8be23a7
# ╠═e0cf019d-2c2e-470b-a62a-25fbb9a22b4d
# ╠═0de64701-70cf-4b1c-ae5a-e288e34a5e3c
# ╠═ffb09bcb-3d02-455c-b153-345f4857fe9d
# ╠═07175d0e-7293-4095-95cc-91fe493a1aef
# ╠═2ebfa422-f8a5-44d3-8f2c-a34d7832d3f2
# ╠═1f1158ce-73f0-49fd-a48a-af3b36376030
# ╠═ebe074cb-7bad-4b52-9c9d-b9752af4bedd
# ╠═d6499c1d-c2a1-4cc0-a8f4-c4ba86d5cf30
# ╠═2361e758-1b5b-4cb2-b7d5-8a03760befa6


@testset "Capillary pressure" begin
    f = reshape(collect(1.0:25),5,5)
    f_float = reshape(collect(1.0f0:25.0f0),5,5)
    res = zeros(5,5)
    @testset "No contact angle" begin
        Swalbe.filmpressure!(res, f, 1.0, 0.0, 3, 2, 0.1, 0.1)
        sol = [-30.0 -5.0 -5.0 -5.0 20;
               -25.0 0.0 0.0 0.0 25.0;
               -25.0 0.0 0.0 0.0 25.0;
               -25.0 0.0 0.0 0.0 25.0;
               -20.0 5.0 5.0 5.0 30.0;
               ]
        for i in eachindex(sol)
            @test res[i] .≈ sol[i] atol=1e-10
        end
        Swalbe.filmpressure!(res, f, 0.0)
        for i in eachindex(sol)
            @test res[i] .≈ 0.01 * sol[i] atol=1e-10
        end
    end
    @testset "No height gradient" begin
        nograd = ones(5,5)
        Swalbe.filmpressure!(res, nograd, 1.0, 1/2, 3, 2, 0.1, 0.0)
        for i in eachindex(res)
            @test res[i] .≈ -2(0.1^2-0.1) atol=1e-10
        end
    end
    @testset "Gradient and contact angle" begin
        Swalbe.filmpressure!(res, f, 1.0, 1/2, 3, 2, 0.1, 0.0)
        sol = [30.0 5.0 5.0 5.0 -20;
               25.0 0.0 0.0 0.0 -25.0;
               25.0 0.0 0.0 0.0 -25.0;
               25.0 0.0 0.0 0.0 -25.0;
               20.0 -5.0 -5.0 -5.0 -30.0]
        for i in eachindex(res)
            # Now the result comprisses of two components the disjoining potential and the laplace term.
            @test res[i] .≈ -1(sol[i] + 20((0.1/f[i])^3-(0.1/f[i])^2)) atol=1e-10
        end
    end
    @testset "Gradient and contact angle Float32" begin
        Swalbe.filmpressure!(res, f_float, 1.0f0, 0.5f0, 3, 2, 0.1f0, 0.0f0)
        sol = [30.0f0 5.0f0 5.0f0 5.0f0 -20f0;
               25.0f0 0.0f0 0.0f0 0.0f0 -25.0f0;
               25.0f0 0.0f0 0.0f0 0.0f0 -25.0f0;
               25.0f0 0.0f0 0.0f0 0.0f0 -25.0f0;
               20.0f0 -5.0f0 -5.0f0 -5.0f0 -30.0f0]

        for i in eachindex(res)
            # Now the result comprisses of two components the disjoining potential and the laplace term.
            @test res[i] .≈ -1(sol[i] + 20((0.1f0/f_float[i])^3-(0.1f0/f_float[i])^2)) atol=1e-6
        end
    end
    dgrad = zeros(5,5,8)
    @testset "No contact angle circshift!" begin
        Swalbe.filmpressure!(res, f, dgrad, 1.0, 0.0, 3, 2, 0.1, 0.1)
        sol = [-30.0 -5.0 -5.0 -5.0 20;
               -25.0 0.0 0.0 0.0 25.0;
               -25.0 0.0 0.0 0.0 25.0;
               -25.0 0.0 0.0 0.0 25.0;
               -20.0 5.0 5.0 5.0 30.0;
               ]
        for i in eachindex(sol)
            @test res[i] .≈ sol[i] atol=1e-10
        end
    end
end

@testset "Capillary pressure 1D" begin
    f = collect(1.0:30)
    f_float = collect(1.0f0:30.0f0)
    res = zeros(30)
    dummy = zeros(30,3)
    @testset "No contact angle" begin
        Swalbe.filmpressure!(res, f, dummy, 1.0, 0.0, 3, 2, 0.1, 0.1)
        # println("My result: $res")
        sol = zeros(30)
        sol[1] = 30
        sol[end] = -30
        @test all(res .== -sol)
    end
    @testset "No height gradient" begin
        nograd = ones(30)
        Swalbe.filmpressure!(res, nograd, dummy, 1.0, 1/2, 3, 2, 0.1, 0.0)
        for i in eachindex(res)
            @test res[i] .≈ -2(0.1^2-0.1) atol=1e-10
        end
    end
    @testset "Gradient and contact angle" begin
        Swalbe.filmpressure!(res, f, dummy, 1.0, 1/2, 3, 2, 0.1, 0.0)
        sol = zeros(30)
        sol[1] = 30
        sol[end] = -30
        for i in eachindex(res)
            # Now the result comprisses of two components the disjoining potential and the laplace term.
            @test res[i] .≈ -1(sol[i] + 20((0.1/f[i])^3-(0.1/f[i])^2)) atol=1e-10
        end
    end
end

@testset "Capillary pressure 1D with Gamma" begin
    f = collect(1.0:30)
    res = zeros(30)
    rho = fill(0.1, 30)
    dummy = zeros(30,2)
    @testset "No Gamma no theta" begin
        Swalbe.filmpressure!(res, f, dummy, rho, 1.0, 0.0, 3, 2, 0.1, 0.1)
        # println("My result: $res")
        sol = zeros(30)
        sol[1] = 30
        sol[end] = -30
        @test all(res .== -sol)
    end
    @testset "No Gamma no laplacian" begin
        nograd = ones(30)
        Swalbe.filmpressure!(res, nograd, dummy, rho, 1.0, 1/2, 3, 2, 0.1, 0.0)
        for i in eachindex(res)
            @test res[i] .≈ -2(0.1^2-0.1) atol=1e-10
        end
    end
    @testset "Gamma but rho=0 no laplacian" begin
        nograd = ones(30)
        Swalbe.filmpressure!(res, nograd, dummy, zeros(30), 1.0, 1/2, 3, 2, 0.1, 0.0)
        for i in eachindex(res)
            @test res[i] .≈ -2(0.1^2-0.1) atol=1e-10
        end
    end
    @testset "Gamma, rho, gradient and contact angle" begin
        rho = fill(0.1, 30)
        Swalbe.filmpressure!(res, f, dummy, rho, 1.0, 1/2, 3, 2, 0.1, 0.0, Gamma=0.1)
        sol = zeros(30)
        sol[1] = 30
        sol[end] = -30
        for i in eachindex(res)
            # Now the result comprisses of two components the disjoining potential and the laplace term.
            # The prefactor (1 + 0.01) -> (gamma + Gamma * rho)
            @test res[i] .≈ -(1 + 0.01) * (sol[i] + 20((0.1/f[i])^3-(0.1/f[i])^2)) atol=1e-10
        end
    end
end

@testset "Power function" begin
    arg = 2
    argsimple = 5.0
    argsimplefp32 = 5.0f0
    argvec = [2.0 3.0 4.0]
    n = [3 6]
    r = [8 64]
    r1 = [125.0 15625.0]
    r2 = [125.0f0 15625.0f0]
    r3 = [8.0 27.0 64.0; 64.0 729.0 4096.0]
    for i in 1:2
        res0 = Swalbe.power_broad(arg, n[i])
        @test res0 == r[i]
        res0 = Swalbe.power_broad(argsimple, n[i])
        @test res0 == r1[i]
        res1 = Swalbe.power_broad(argsimplefp32, n[i])
        @test res1 == r2[i]
        res2 = Swalbe.power_broad.(argvec, n[i])
        for j in 1:3
            @test res2[j] .== r3[i,j]
        end
    end
end


using ChaosTools
using Test, StaticArrays, OrdinaryDiffEq

println("\nTesting orbit diagrams...")

@testset "Orbit Diagram" begin
    @testset "Discrete 1D" begin
        ds = Systems.logistic()
        i = 1
        parameter = 1
        pvalues = 2:0.01:4
        ics = [rand() for m in 1:10]
        n = 50
        Ttr = 5000
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test output[1][1] ≈ 0.5
        @test output[2][1] == 0.5024875621890547
        @test output[end][1] != output[end][2] != output[end][3]
    end

    @testset "Discrete 2D" begin
        ds = Systems.standardmap()
        i = 2
        parameter = 1
        pvalues = 0:0.005:2
        ics = [0.001rand(2) for m in 1:10]
        n = 50
        Ttr = 5000
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test length(output[1]) == n
    end

    @testset "Discrete IIP" begin
        ds = Systems.henon_iip()
        i = 2
        parameter = 1
        pvalues = [1.4, 1.5]
        ics = [rand(2), rand(2)]
        n = 50
        Ttr = 5000
        output = orbitdiagram(ds, i, parameter, pvalues; n = n, Ttr = Ttr)
        @test length(output[1]) == n
    end
end

@testset "Poincare SOS" begin
    @testset "Henon Helies" begin
        ds = Systems.henonheiles([0, .483000, .278980390, 0] )
        psos = poincaresos(ds, (2, 0.0), 1000.0)
        xcross = psos[:, 2]
        @test length(xcross) > 1
        for x in xcross
            @test abs(x) < 1e-3
        end

        # @inline Vhh(q1, q2) = 1//2 * (q1^2 + q2^2 + 2q1^2 * q2 - 2//3 * q2^3)
        # @inline Thh(p1, p2) = 1//2 * (p1^2 + p2^2)
        # @inline Hhh(q1, q2, p1, p2) = Thh(p1, p2) + Vhh(q1, q2)
        # @inline Hhh(u::AbstractVector) = Hhh(u...)
        #
        # E = [Hhh(p) for p in psos]
        #
        # @test std(E) < 1e-10
        # @test maximum(@. E - E[1]) < 1e-10
    end
    @testset "Gissinger crazy plane" begin
        gis = Systems.gissinger([2.32865, 2.02514, 1.98312])
        # Define appropriate hyperplane for gissinger system
        ν = 0.1
        Γ = 0.9 # default parameters of the system

        # I want hyperperplane defined by these two points:
        Np(μ) = SVector{3}(sqrt(ν + Γ*sqrt(ν/μ)), -sqrt(μ + Γ*sqrt(μ/ν)), -sqrt(μ*ν))
        Nm(μ) = SVector{3}(-sqrt(ν + Γ*sqrt(ν/μ)), sqrt(μ + Γ*sqrt(μ/ν)), -sqrt(μ*ν))

        # Create hyperplane using normal vector to vector connecting points:
        gis_plane(μ) = (d = (Np(μ) - Nm(μ)); [d[2], -d[1], 0, 0])

        μ = 0.12
        set_parameter!(gis, 1, μ)
        psos = poincaresos(gis, gis_plane(μ), 10000.0, Ttr = 200.0, direction = -1)
        @test length(psos) > 1
        @test generalized_dim(2, psos) < 1
    end
    @testset "beginning on the plane" begin

        @inbounds @inline function ż(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
        z0 = SVector{4}(2.3499921651423565, -9.801351029039825, 0.0, -2.7230316965872268)

        ds = ContinuousDynamicalSystem(ż, z0, (A=1, B=0.55, D=0.4))
        psos = poincaresos(ds, (3, 0), 10., direction=1, idxs=[2,4])
        @test length(psos) > 0
        @test size(psos)[2] == 2
    end
end

@testset "Produce OD" begin
  @testset "Shinriki" begin
    ds = Systems.shinriki([-2, 0, 0.2])

    pvalues = range(19,stop=22,length=11)
    parameter = 1
    i = 1
    j = 2
    tf = 200.0

    de = Dict(:abstol=>1e-9, :reltol => 1e-9)
    output = produce_orbitdiagram(ds, (j, 0.0), i, parameter, pvalues; tfinal = tf,
    Ttr = 100.0, printparams = false, de...)
    @test length(output) == length(pvalues)

    v = round.(output[1], digits = 4)
    s = collect(Set(v))
    @test length(s) == 1
    @test s[1] == -0.856

    v = round.(output[4], digits = 4)
    s = Set(v)
    @test length(s) == 2
    @test s == Set([-0.376, -1.2887])

  end
end

# # This test offers nothing to the Suite (does pass though)
# @testset "Stroboscopic" begin
#   ds = Systems.duffing(β = -1, ω = 1, f = 0.3)
#   a = trajectory(ds, 100000.0, dt = 2π)
#   D = information_dim(a)
#   @test 1.25 < D < 1.5
# end


export Contingency

using LinearAlgebra

struct Contingency{ T, P <: AbstractMatrix{T} } <: ContinuousMatrixDistribution
    p :: P
end

contingency_matrix(distribution::Contingency) = distribution.p

vague(::Type{ <: Contingency }, dims::Int) = Contingency(ones(dims, dims) ./ abs2(dims))

function entropy(distribution::Contingency)
    P = contingency_matrix(distribution)
    -sum(P .* log.(clamp.(P, tiny, Inf)))
end


using Test
using ForwardDiff;
using FinancialToolbox
Dual_=ForwardDiff.Dual
print_colored("Starting Forward Diff Dual Numbers Test\n",:green)

#Test Parameters
testToll=1e-14;
spot=10;K=10;r=0.02;T=2.0;sigma=0.2;d=0.01;
assert_(value,toll)=@test abs(value)<toll
#EuropeanCall Option
PriceCall=blsprice(spot,K,r,T,sigma,d);
DeltaCall=blsdelta(spot,K,r,T,sigma,d);
ThetaCall=blstheta(spot,K,r,T,sigma,d);
RhoCall=blsrho(spot,K,r,T,sigma,d);
LambdaCall=blslambda(spot,K,r,T,sigma,d);
VannaCall=blsvanna(spot,K,r,T,sigma,d);

#EuropeanPut Option
PricePut=blsprice(spot,K,r,T,sigma,d,false);
DeltaPut=blsdelta(spot,K,r,T,sigma,d,false);
ThetaPut=blstheta(spot,K,r,T,sigma,d,false);
RhoPut=blsrho(spot,K,r,T,sigma,d,false);
LambdaPut=blslambda(spot,K,r,T,sigma,d,false);
VannaPut=blsvanna(spot,K,r,T,sigma,d,false);

#Equals for both Options
Gamma=blsgamma(spot,K,r,T,sigma,d);
Vega=blsvega(spot,K,r,T,sigma,d);

########DUAL NUMBERS
DerToll=1e-13;
#Function definition
#Call
FcallDual(spot)=blsprice(spot,K,r,T,sigma,d);
GcallDual(r)=blsprice(spot,K,r,T,sigma,d);
HcallDual(T)=blsprice(spot,K,r,T,sigma,d);
LcallDual(sigma)=blsprice(spot,K,r,T,sigma,d);
PcallDual(spot)=blsprice(spot,K,r,T,sigma,d)*spot.value/blsprice(spot.value,K,r,T,sigma,d);
VcallDual(sigma)=blsdelta(spot,K,r,T,sigma,d);
#Put
FputDual(spot)=blsprice(spot,K,r,T,sigma,d,false);
GputDual(r)=blsprice(spot,K,r,T,sigma,d,false);
HputDual(T)=blsprice(spot,K,r,T,sigma,d,false);
PputDual(spot)=blsprice(spot,K,r,T,sigma,d,false)*spot.value/blsprice(spot.value,K,r,T,sigma,d,false);
VPutDual(sigma)=blsdelta(spot,K,r,T,sigma,d,false);
#Input
SpotDual=Dual_(spot,1.0);
rDual=Dual_(r,1.0);
TDual=Dual_(T,1.0);
SigmaDual=Dual_(sigma,1.0);

#Automatic Differentiation Test
#TEST
print_colored("--- European Call Sensitivities: DualNumbers\n",:yellow)
print_colored("-----Testing Delta\n",:blue);
assert_(FcallDual(SpotDual).partials[1]-DeltaCall,DerToll)
print_colored("-----Testing Rho\n",:blue);
assert_(GcallDual(rDual).partials[1]-RhoCall,DerToll)
print_colored("-----Testing Theta\n",:blue);
assert_(-HcallDual(TDual).partials[1]-ThetaCall,DerToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(PcallDual(SpotDual).partials[1]-LambdaCall,DerToll)
print_colored("-----Testing Vanna\n",:blue);
assert_(VcallDual(SigmaDual).partials[1]-VannaCall,DerToll)
#TEST
print_colored("--- European Put Sensitivities: DualNumbers\n",:yellow)
print_colored("-----Testing Delta\n",:blue);
assert_(FputDual(SpotDual).partials[1]-DeltaPut,DerToll)
print_colored("-----Testing Rho\n",:blue);
assert_(GputDual(rDual).partials[1]-RhoPut,DerToll)
print_colored("-----Testing Theta\n",:blue);
assert_(-HputDual(TDual).partials[1]-ThetaPut,DerToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(PputDual(SpotDual).partials[1]-LambdaPut,DerToll)
print_colored("-----Testing Vanna\n",:blue);
assert_(VPutDual(SigmaDual).partials[1]-VannaPut,DerToll)

print_colored("-----Testing Vega\n",:blue);
assert_(LcallDual(SigmaDual).partials[1]-Vega,DerToll)
print_colored("Dual Numbers Test Passed\n",:green)
println("")

#TEST OF INPUT VALIDATION
print_colored("Starting Input Validation Test Dual\n",:magenta)
print_colored("----Testing Negative  Spot Price \n",:cyan)
@test_throws(ErrorException, blsprice(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blkprice(-SpotDual,K,r,T,sigma))
@test_throws(ErrorException, blsdelta(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blstheta(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsrho(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blspsi(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blslambda(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsvega(-SpotDual,K,r,T,sigma,d))

print_colored("----Testing Negative  Strike Price \n",:cyan)
KK=Dual_(K,0);
@test_throws(ErrorException, blsprice(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blkprice(spot,-KK,r,T,sigma))
@test_throws(ErrorException, blsdelta(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blstheta(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsrho(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blspsi(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blslambda(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsvega(spot,-KK,r,T,sigma,d))

print_colored("----Testing Negative  Time to Maturity \n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blkprice(spot,K,r,-TDual,sigma))
@test_throws(ErrorException, blsdelta(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsgamma(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blstheta(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsrho(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blspsi(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blslambda(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsvanna(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsvega(spot,K,r,-TDual,sigma,d))

print_colored("----Testing Negative  Volatility \n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blkprice(spot,K,r,T,-SigmaDual))
@test_throws(ErrorException, blsdelta(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsgamma(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blstheta(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsrho(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blspsi(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blslambda(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsvanna(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsvega(spot,K,r,T,-SigmaDual,d))

print_colored("Dual Input Validation Test Passed\n",:magenta)


"""
Compute a scattered grid of M points randomly distributed in `Ω`, using
the uniform probability measure on `Ω`.
"""
function randomgrid(Ω::Domain, M::Int)
    points = [randompoint(Ω, boundingbox(Ω)) for m in 1:M]
    ScatteredGrid(sort(points))
end

"Generate a single random point inside the given box, with `eltype` `T`."
function randompoint(dom::ProductDomain)
    convert(eltype(dom),map(randompoint,elements(dom)))
end

# Don't return an SVector in 1d, just a value
function randompoint(dom::AbstractInterval)
    T = float(eltype(dom))
    val = rand(T)
    convert(T,val * infimum(dom) + (1-val) * supremum(dom))
end

"""
Generate a single random point inside the given domain, with `eltype` `T`.
Random points are generated inside the given box, until one is inside the domain.
"""
function randompoint(Ω::Domain, box = boundingbox(Ω))
    local point
    indomain = false
    while ~indomain
        point = randompoint(box)
        indomain = point ∈ Ω
    end
    point
end


###########################################################################################
#Forecast using regulated regression methods which include ridge models and lasso models.
#This is the inner function to compute the RMSE for each λ
###########################################################################################
@everywhere function glmnet_inner_RMSE(J, h, h_lag, J_cv, Y, Y_lag, X_lag, α, λ)
    ## Vector Index for oos forecast goes from J_cv - (h_lag + h) to J - (1 + h_lag + h)
    ŷ_cv = zeros(Float64, size(J_cv - (h_lag + h):J - (1 + h_lag + h)))
    for (n, vec_î) in enumerate(J_cv - (h_lag + h):J - (1 + h_lag + h))
        LHS = Y[1:vec_î - h]
        sample_size = length(LHS)
        RHS = [Y_lag[1:sample_size] X_lag[1:sample_size, :]]
        res = glmnet(RHS, LHS, alpha = α, lambda = [λ], tol = 10e-20, maxit = 10_000_000)
        β̂ = [res.a0; collect(res.betas)]
        RHS_crnt = [1; Y_lag[vec_î]; X_lag[vec_î, :]]
        ŷ_cv[n] = dot(β̂, RHS_crnt)
    end
    RMSE = sqrt(mean((Y[J_cv - (h_lag + h):J - (1 + h_lag + h)] - ŷ_cv) .^ 2))
    return [RMSE, λ]
end

###########################################################################################
#Forecast using regulated regression methods which include ridge models and lasso models.
#This is the inner function to compute the RMSE for each λ
###########################################################################################
function glmnet_oos_forc(J, h, h_lag, J_cv, Y, Y_lag, X_lag, λ_vec, α)
    β̂ = ones(size(X_lag, 2) + 2, 1)
    # Step1 : pick the best λ using parallel computing
    RMSE_cv = zeros(length(λ_vec), 2)
    result = pmap(λ -> glmnet_inner_RMSE(J, h, h_lag, J_cv, Y, Y_lag, X_lag, α, λ),
                  λ_vec)::Array{Array{Float64,1},1}
    for i in 1:length(λ_vec)
        RMSE_cv[i, :] = result[i]
    end
    index = findmin(RMSE_cv[:, 1])[2]
    λ = λ_vec[index]
    # Step 2: oos forecasting
    ## Vector Index for oos forecast goes from J - (h_lag + h) to length(Y)
    ŷ = zeros(Float64, size(J - (h_lag + h):length(Y)))
    for (n, vec_î) in enumerate(J - (h_lag + h):length(Y))
        LHS = Y[1:vec_î - h]
        sample_size = length(LHS)
        RHS = [Y_lag[1:sample_size] X_lag[1:sample_size, :]]
        res = glmnet(RHS, LHS, alpha = α, lambda = [λ], tol = 10e-20, maxit = 10_000_000)
        β̂ = [res.a0; collect(res.betas)]
        RHS_crnt = [1; Y_lag[vec_î]; X_lag[vec_î, :]]
        ŷ[n] = dot(β̂, RHS_crnt)
    end
    RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
    return β̂, λ, ŷ, RMSE, RMSE_cv
end

#α = 0.0
#λ_vec = 10.:0.1:20. or λ_vec = 0.1.:0.2:20.
#β̂, λ, ŷ, RMSE_cv = glmnet_oos_forc(J, h, h_lag, J_cv, Y, Y_lag, X_lag, λ_vec, α)
#plot(RMSE_cv[:, 1])
#plot(Y[J - (h_lag + h):length(Y)])
#plot!(ŷ)
#RMSE = sqrt(mean((Y[J - (h_lag + h):length(Y)] - ŷ) .^ 2))
#println("Tuning Parameter Chosen for ridge: $λ")
#println("RMSE from ridge including first order terms: \nRMSE = $RMSE")


#=

  Archery Puzzle in Julia ConstraintSolver.jl

  http://www.eonhq.com/m/puzzles/images/archery-puzzle.jpg
  Archery puzzle by Sam Loyd:
  """
  How close can the young archer come to scoring a total of
  100 - using as many arrows as she please.
  [The targets are: 16, 17, 23, 24, 39, 40]
  """
  Via: The Aperiodical: "Manchester MathsJam June 2012 Recap"
  http://aperiodical.com/2012/06/manchester-mathsjam-june-2012-recap/

  Solution (unique):
  x = [2,4,0,0,0,0]
  z = 100

  i.e. 2 arrows on score 16 and 4 arrows on score 17:
    2*16 + 4*17 = 100

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

=#

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

function archery_puzzle(print_solutions=true,all_solutions=true)

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

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

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

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

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

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

                                                            "time_limit" => 6,

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

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

    n = 6
    targets =  [16, 17, 23, 24, 39, 40]
    target = 100
    @variable(model, 0 <= x[1:n] <= 10, Int)
    @variable(model, 0 <= z <= n*maximum(targets), Int)
    @variable(model, 0 <= diff_z <= n*maximum(targets), Int)

    scalar_product(model,x,targets,:(==),z)
    my_abs(model,target,z,diff_z)

    @objective(model,Min,diff_z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                z_val = convert.(Integer,JuMP.value.(z; result=sol))
                diff_z_val = convert.(Integer,JuMP.value.(diff_z; result=sol))
                println("x:$x_val")
                println("z:$z_val diff_z:$diff_z_val")
                println(["target:$(targets[i]) x:$(x_val[i])  " for i in 1:n if x_val[i] > 0]|>join)
                println()

            end
        end
    else
        println("status:$status")
    end

    return status, num_sols
end

@time archery_puzzle(true,true)


function jacobi_preconditioned_gmres(A::SparseMatrixCSC, x0::Vector{Float64}, b::Vector{Float64}, Nz, pchunks; w=2/3, Na=10, Ns=100, m=20, err_thresh=1e-8)

    Nk = div(length(x0), Nz)

    J = JacobiSmoother(A, b::Vector{Float64}, w, pchunks, Nk, Nz, Na)
    x = x0[:]
    x[:] = J*x
    for t=1:Ns
        x[:] = J * x
    end

    err = norm(A*x - b) / norm(b)

    if err > err_thresh
        zslices = collect(Iterators.partition(1:Nk*Nz, Nz))
        c = pmap(zslices) do inds
            Aloc = A[inds, inds]
            bloc = b[inds]
            gmres_solve(Aloc, x[inds], bloc, m; err_thresh=err_thresh)
        end
        return vcat(c...)
    else
        return x
    end

end

#
#
# struct LinearLeftHandSide
#
#     A :: SparseMatrixCSC
#     M :: SparseMatrixCSC
#     D :: SparseMatrixCSC
#
# end
# LinearLeftHandSide(A::SparseMatrixCSC, w::Float64) = LinearLeftHandSide(A, compute_jacobi_matrices(A, w)...)
#
#
#
# function jacobi_preconditioned_gmres(L::LinearLeftHandSide, x0::Vector{Float64}, b::Vector{Float64}; n=100, m=10, err_thresh=1e-10)
#
#     # precodition using jacobi_relaxation()
#     x = jacobi_relaxation(n, L.A, L.M, L.D, x0, b)
#
#     # finish with gmres
#     x = gmres_solve(L.A, x, b, m; err_thresh=err_thresh)
#
#     return x
# end
#
#
# function parallel_jacobi_preconditioned_gmres(L::LinearLeftHandSide, x0::Vector{Float64}, b::Vector{Float64}, pchunks::Int64, Nz::Int64; nsynch=100, nasynch=5, m=10, err_thresh=1e-10)
#
#     # jacobi preconditioning
#     x = distribute(x0; dist=pchunks*Nz)
#     bdist = distribute(b; dist=pchunks*Nz)
#     for _=1:nsynch
#         x = asynch_jacobi_relaxation(nasynch, L.M, L.D, x, bdist)
#     end
#
#     # do gmres in parallel
#     x = DArray((size(x)[1],), workers()[1:Nz], Nz) do inds
#
#         xs = Array(x0[inds[1],1])
#         bs = b[inds[1]]
#         gmres_solve(L.A, xs, bs, m, err_thresh=err_thresh)
#     end
#
#     return Array(x)
# end


module TestGLM

# using Revise
using Test

using MLJBase
import MLJModels
import GLM
using MLJModels.GLM_

###
### OLSREGRESSOR
###

task = load_boston()
X, y = task();

train, test = partition(eachindex(y), 0.7)

atom_ols = OLSRegressor()

Xtrain = selectrows(X, train)
ytrain = selectrows(y, train)
Xtest  = selectrows(X, test)

fitresult, _, report = fit(atom_ols, 1, Xtrain, ytrain)
MLJBase.fitted_params(atom_ols, fitresult)

p = predict_mean(atom_ols, fitresult, Xtest)

# hand made regression to compare

Xa = MLJBase.matrix(X) # convert(Matrix{Float64}, X)
Xa1 = hcat(Xa, ones(size(Xa, 1)))
coefs = Xa1[train, :] \ y[train]
p2 = Xa1[test, :] * coefs

@test p ≈ p2

info(atom_ols)

p_distr = predict(atom_ols, fitresult, selectrows(X, test))

###
### GLMCOUNT
###

using Random: seed!
using Distributions

seed!(0)

X = randn(100, 3) .* randn(3)'
Xtable = table(X)

α = 0.1
β = [-0.3, 0.2, -0.1]
λ = exp.(α .+ X * β)

y = [rand(Poisson(λᵢ)) for λᵢ ∈ λ]

atom_glmcount = GLMCount()

fitresult, _, _ = fit(atom_glmcount, 1, Xtable, y)

p = predict_mean(atom_glmcount, fitresult, Xtable)

p_distr = predict(atom_glmcount, fitresult, Xtable)

end # module
true


using LazyArrays, LinearAlgebra, FillArrays, Test
import LazyArrays: InvMatrix, ApplyBroadcastStyle, LdivStyle, Applied, LazyLayout
import Base.Broadcast: materialize

@testset "Ldiv" begin
    @testset "Float64 \\ *" begin
        A = randn(5,5)
        b = randn(5)
        M = Ldiv(A,b)
        @test all(materialize(M) .=== (A\b) .=== materialize(applied(\,A,b)))

        @test applied(\,A,b) isa Applied{LdivStyle}
        @test applied(\,A,b)[1] ≈ (A\b)[1]

        @test parent(InvMatrix(A)) === inv(InvMatrix(A)) === pinv(InvMatrix(A)) === A
        @test parent(PInvMatrix(A)) === pinv(InvMatrix(A)) === A

        @test InvMatrix(A)[1,2] ≈ PInvMatrix(A)[1,2] ≈ inv(A)[1,2]

        @test all(copyto!(similar(b), Ldiv(A,b)) .===
                    (similar(b) .= Ldiv(A,b)) .=== InvMatrix(A) * b .===
                    materialize(Ldiv(A,b)) .===
                    apply(\,A,b) .===
                  (A\b) .=== (b̃ =  copy(b); LAPACK.gesv!(copy(A), b̃); b̃))

        @test all(copyto!(similar(b), Ldiv(UpperTriangular(A) , b)) .===
                    (similar(b) .= Ldiv(UpperTriangular(A),b)) .===
                    InvMatrix(UpperTriangular(A))*b .===
                    BLAS.trsv('U', 'N', 'N', A, b) )

        b = randn(5) + im*randn(5)
        @test InvMatrix(A) * b ≈ Matrix(A) \ b
        @test InvMatrix(A) \ b == A*b
    end

    @testset "ComplexF64 \\ *" begin
        T = ComplexF64
        A = randn(T,5,5)
        b = randn(T,5)

        b = randn(5)
        @test InvMatrix(A) * b ≈ A \ b
    end

    @testset "Rectangle PInv" begin
        A = randn(5,3)
        A_orig = copy(A)
        b = randn(5)
        b_orig = copy(b)
        @test_throws DimensionMismatch InvMatrix(A)
        @test PInvMatrix(A) * b == (A\b)
        @test all(b .=== b_orig)
    end

    @testset "inv and pinv"  begin
        A = randn(5,5)
        @test inv(Inv(A)) === inv(PInv(A)) === pinv(Inv(A)) === pinv(PInv(A)) === A

        Ai = PInvMatrix(A)
        @test Matrix(Ai) ≈ inv(A)
        Ai = InvMatrix(A)
        @test Matrix(Ai) ≈ inv(A)

        A = randn(5,3)
        @test pinv(PInv(A)) === A
        Ai = PInvMatrix(A)
        @test Ai ≈ pinv(A)
        @test_throws DimensionMismatch inv(PInv(A))
        @test_throws DimensionMismatch InvMatrix(A)

        A = randn(3,5)
        @test pinv(PInv(A)) === A
        Ai = PInvMatrix(A)
        @test Ai ≈ pinv(A)
        @test_throws DimensionMismatch inv(PInv(A))
        @test_throws DimensionMismatch InvMatrix(A)
    end

    @testset "Int" begin
        A = [1 2 ; 3 4]; b = [5,6];
        @test eltype(applied(inv, A)) == eltype(applied(pinv, A)) == eltype(applied(\, A, b)) == eltype(Ldiv(A, b)) ==
            eltype(ApplyArray(\, A, b)) == eltype(InvMatrix(A)) == eltype(PInvMatrix(A)) == Float64
        @test apply(\,A,b) == A\b
    end

    @testset "QR" begin
        A = randn(5,3)
        b = randn(5)
        B = randn(5,5)
        Q,R = qr(A)
        @test Q\b ≈ apply(\,Q,b)
        @test Q\B ≈ apply(\,Q,B)
        @test_throws DimensionMismatch apply(\, Q, randn(4))
        @test_throws DimensionMismatch apply(\, Q, randn(4,3))
        dest = fill(NaN,5)
        @test copyto!(dest, applied(\,Q,b)) == apply(\,Q,b)
    end

    @testset "Lazy *" begin
        A = randn(5,5)
        B = randn(5,5)
        M = ApplyArray(*, A, B)
        b = 1:5
        @test apply(\,B,ApplyArray(*,A,b)) ≈ B\A*b
        @test M \ b ≈ B\(A\b)
        @test M \ M ≈ Eye(5)
        @test M \ b isa Vector
        @test M\M isa Matrix

        @testset "Inv of Mul" begin
            @test inv(M) ≈ inv(A*B)
            @test_throws DimensionMismatch inv(ApplyArray(*,randn(5,6), rand(6,5)))
        end

        @testset "Inv * Eye" begin
            @test ApplyArray(inv,A) * Eye(5) == Eye(5) * ApplyArray(inv,A)
            @test ApplyArray(inv,A) * Diagonal(Fill(2,5)) == Diagonal(Fill(2,5)) * ApplyArray(inv,A)
        end
    end

    @testset "Triangular ldiv" begin
        A = randn(5,5)
        b = randn(5)
        L = applied(\, UpperTriangular(A), b)
        @test copyto!(similar(L), L) ≈ UpperTriangular(A) \ b ≈ materialize!(L)
    end

    @testset "Diagonal \\ Mul" begin
        D = Diagonal(randn(5))
        A = ApplyArray(*, randn(5,5), randn(5))
        @test D \ A ≈ D \ Vector(A)
    end

    @testset "LdivArray Mul" begin
        A = ApplyArray(\, Diagonal(1:5), randn(5,5))
        b = randn(5)
        @test A * b ≈ Matrix(A) * b
    end

    @testset "Diagonal \\ Lazy" begin
        D = Diagonal(randn(5))
        A = BroadcastArray(*, randn(5,5), randn(5))
        @test D \ A ≈ D \ Matrix(A)
    end
end

@testset "Rdiv" begin
    A = randn(5,5)
    R = ApplyArray(/, A, A)
    @test axes(R) == (axes(R,1),axes(R,2)) == axes(A)
    @test R ≈ apply(/, A, A) ≈ Eye(5)
end


immutable TwoPlayerGame
    # Number of actions
    na1::Int
    na2::Int

    # Discount factor
    δ::Float64

    # Payoffs
    g1::Array{Float64, 2}
    g2::Array{Float64, 2}

    # Best Response
    br1::Array{Float64, 1}
    br2::Array{Float64, 1}

    # Gains from deviations
    h1::Array{Float64, 2}
    h2::Array{Float64, 2}
end

# Output flow utility for action profile
@inline g(tpg::TwoPlayerGame, i::Int, j::Int) = [tpg.g1[i, j], tpg.g2[i, j]]
@inline g1(tpg::TwoPlayerGame, i::Int, j::Int) = tpg.g1[i, j]
@inline g2(tpg::TwoPlayerGame, i::Int, j::Int) = tpg.g2[i, j]

# Output period's value
@inline v(tpg::TwoPlayerGame, i::Int, j::Int, w::Vector{Float64}) =
    [v1(tpg, i, j, w[1]), v2(tpg, i, j, w[2])]
@inline v1(tpg::TwoPlayerGame, i::Int, j::Int, w1::Float64) =
    (1-tpg.δ)*g1(tpg, i, j) + tpg.δ*w1
@inline v2(tpg::TwoPlayerGame, i::Int, j::Int, w2::Float64) =
    (1-tpg.δ)*g2(tpg, i, j) + tpg.δ*w2

# Function to pull out each player's best response given action set
@inline bestresponse(tpg::TwoPlayerGame, i::Int, j::Int) = [tpg.br1[j], tpg.br2[i]]
@inline bestresponse1(tpg::TwoPlayerGame, i::Int, j::Int) = tpg.br1[j]
@inline bestresponse2(tpg::TwoPlayerGame, i::Int, j::Int) = tpg.br2[i]

function TwoPlayerGame(δ::Float64, g1::Array{Float64, 2}, g2::Array{Float64, 2})

    # Pull off size and then make sure matrices are same size
    na1, na2 = size(g1)
    @assert size(g2) == (na1, na2)

    # Get best responses
    br1 = Array(Int, na2)
    @inbounds for i=1:na2
        br1[i] = indmax(g1[:, i])
    end
    br2 = Array(Int, na1)
    @inbounds for i=1:na1
        br2[i] = indmax(g2[i, :])
    end

    # Get gains from deviating
    h1 = Array(Float64, na1, na2)
    h2 = Array(Float64, na1, na2)
    @inbounds for i=1:na1
        _br_2 = br2[i]

        @inbounds for j=1:na2
            _br_1 = br1[j]
            h1[i, j] = g1[_br_1, j] - g1[i, j]
            h2[i, j] = g2[i, _br_2] - g2[i, j]
        end
    end

    return TwoPlayerGame(na1, na2, δ, g1, g2, br1, br2, h1, h2)
end

# Function to initialize the value set
function intitialize_w0(tp::TwoPlayerGame)
    # Start with all extreme points
    w0 = [Point(tp.minu1, tp.minu2), Point(tp.minu1, tp.maxu2),
          Point(tp.maxu1, tp.maxu2), Point(tp.maxu1, tp.minu2)]

    return w0
end


include("utils.jl")

@testset "Sphere" begin
    M = Sphere(2)
    @testset "Sphere Basics" begin
        @test representation_size(M) == (3,)
        @test !is_manifold_point(M, [1., 0., 0., 0.])
        @test !is_tangent_vector(M, [1.,0.,0.], [0., 0., 1., 0.])
    end
    types = [Vector{Float64},
             SizedVector{3, Float64},
             MVector{3, Float64},
             Vector{Float32},
             SizedVector{3, Float32},
             MVector{3, Float32},
             Vector{Double64},
             MVector{3, Double64},
             SizedVector{3, Double64}]
    for T in types
        @testset "Type $T" begin
            pts = [convert(T, [1.0, 0.0, 0.0]),
                   convert(T, [0.0, 1.0, 0.0]),
                   convert(T, [0.0, 0.0, 1.0])]
            test_manifold(M,
                          pts,
                          test_reverse_diff = isa(T, Vector),
                          test_project_tangent = true,
                          test_musical_isomorphisms = true,
                          test_vector_transport = true,
                          test_mutating_rand = isa(T, Vector),
                          point_distributions = [Manifolds.uniform_distribution(M, pts[1])],
                          tvector_distributions = [Manifolds.normal_tvector_distribution(M, pts[1], 1.0)])

            @test isapprox(-pts[1], exp(M, pts[1], log(M, pts[1], -pts[1])))
        end
    end

    @testset "Distribution tests" begin
        usd_mvector = Manifolds.uniform_distribution(M, @MVector [1.0, 0.0, 0.0])
        @test isa(rand(usd_mvector), MVector)

        gtsd_mvector = Manifolds.normal_tvector_distribution(M, (@MVector [1.0, 0.0, 0.0]), 1.0)
        @test isa(rand(gtsd_mvector), MVector)
    end

    @testset "log edge case" begin
        n = manifold_dimension(M)
        x = normalize(randn(n + 1))
        v = log(M, x, -x)
        @test norm(v) ≈ π
        @test isapprox(dot(x, v), 0; atol=1e-12)
        vexp = normalize(project_tangent(M, x, [1, zeros(n)...]))
        @test v ≈ π * vexp

        x = [1, zeros(n)...]
        v = log(M, x, -x)
        @test norm(v) ≈ π
        @test isapprox(dot(x, v), 0; atol=1e-12)
        vexp = normalize(project_tangent(M, x, [0, 1, zeros(n - 1)...]))
        @test v ≈ π * vexp
    end

end


using ReachabilityAnalysis
using ReachabilityAnalysis: solve
using JuMP, Ipopt
using Plots
using LaTeXStrings

# IVP parameters
A = [0.0 1; -4π 0]
X = Universe(2)

# discretization parameter
δ = 0.005

# varying initial states
X0s = [
       (Singleton([0.0, 10.0]), "□(x₀, 0)"),
       (BallInf([0.0, 10.0], 0.01), "□(x₀, 0.01)"),
       (BallInf([0.0, 10.0], 0.1), "□(x₀, 0.1)"),
       (BallInf([0.0, 10.0], 1.0), "□(x₀, 1)")
      ]

# methods
algs = [
        (FirstOrderddt(), "d/dt"),
        (FirstOrderZonotope(), "Zonotope"),
        (CorrectionHull(order=4), "Correction hull"),
        (FirstOrder(), "First order"),
        (ForwardBackward(solver=Ipopt.Optimizer()), "Forward/backward"),
        (Forward(), "Forward")
       ]

# results
X02Ω0s = []
for (X0, name_X0) in X0s
    local Ω0s = []
    local P = @ivp(x' = A*x, x ∈ X, x(0) ∈ X0)
    for (alg, name_alg) in algs
        local Pd = discretize(P, δ, alg)
        local Ω0 = Pd.x0
        push!(Ω0s, Ω0)
    end
    push!(X02Ω0s, Ω0s)
end

# plot sets
colors = [:orange, :green, :blue, :purple, :red, :black]
plots = []
for (i, (X0, name_X0)) in enumerate(X0s)
    local fig = plot(leg=:bottomleft, xlab="X₀ = $name_X0")
    for (j, (alg, name_alg)) in enumerate(algs)
        local Ω0 = X02Ω0s[i][j]
        if name_alg == "Forward/backward"
            Ω0 = overapproximate(Ω0, PolarDirections(30))  # workaround because σ is not available
        end
        plot!(Ω0, alpha=0, color=colors[j], linecolor=colors[j], linealpha=1, lw=2, lab="Ω₀ $(name_alg)")
    end

    # reachable states at time t
    R(t) = linear_map(exp(A * t), X0)

    [plot!(R(δi), color=:gray, alpha=0.05, linealpha=0, lab="") for δi in range(0.0; stop=δ, length=20)]

    savefig("experiment_X0_$i.pdf")
    push!(plots, fig)
end

return plots


function part1(input)
    return maximum(decode.(readlines(input)))
end

function part2(input)
    seats = decode.(readlines(input))
    return only(setdiff(minimum(seats):maximum(seats), seats))
end

function decode(s)
    for pair in ["B" => "1", "R" => "1", "F" => "0", "L" => "0"]
        s = replace(s, pair)
    end
    return parse(Int, s, base = 2)
end


function dense_w(in_size, out_size, m; high_accuracy::Bool=false)
    if m == "Xavier"
        w = randn((out_size, in_size)) ./ sqrt(in_size)
    elseif m == "He"
        w = randn((out_size, in_size)) ./ sqrt(in_size) .* sqrt(2)
    else
        try
            w = m(out_size, in_size)
        catch
            throw(ArgumentError("`set_w` must be `Xavier`, `He` or a function to create weight."))
        end
    end
    if !high_accuracy
        return Float32.(w)
    else
        return w
    end
end

function conv_w(K, io, m; high_accuracy::Bool=false)
    if m == "Xavier"
        w = randn(K..., io...) ./ sqrt(io[1])
    elseif m == "He"
        w = randn(K..., io...) ./ sqrt(io[1]) .* sqrt(2)
    else
        throw(ArgumentError("Custom initialization functions aren't yet aupported!"))
    end
    if !high_accuracy
        return Float32.(w)
    else
        return w
    end
end

create_bias(w, dims...) = fill!(similar(w, dims...), 0)

expand(N, i::Tuple) = i
expand(N, i::Integer) = ntuple(_ -> i, N)

struct KeepSize end

expand_padding(padding, k::NTuple{N, T}, dilation, stride) where {N, T} = expand(2N, padding)
function expand_padding(::KeepSize, k::NTuple{N, T}, dilation, stride) where {N, T}
    k_eff = @. k + (k - 1) * (dilation - 1)
    pad_amt = @. k_eff - 1
    return Tuple(mapfoldl(i -> [cld(i, 2), fld(i,2)], vcat, pad_amt))
end

function params(model)
    pa = Array{Any}(undef, 0)
    for i in 1 : length(model.net)
        try
            layer = model[i]
            push!(pa, trainable(layer)...)
        catch
            continue
        end
    end
    return Zygote.Params(pa)
end


#= using UnicodePlots =#

# Set up model
function build_model()
    F = diagm([1.0])
    V = diagm([2.0])
    G = reshape([1, 2, -0.5], 3, 1)
    W = diagm([8.0, 2.5, 4.0])
    x0 = randn(1)
    P0 = diagm([1e7])
    StateSpaceModel(F, V, G, W, x0, P0)
end

# Time varying model
function sinusoid_model(omega::Real; fs::Int=256, x0=[0.5, -0.5], W::AbstractFloat=0.1)
    F(n)  = [1.0 0; 0 1.0]
    V(n)  = diagm([1e-10, 1e-10])
    G(n)  = [cos(2*pi*omega*(1/fs)*n) -sin(2*pi*omega*(1/fs)*n)]
    w(n)  = diagm([W])
    P0    = diagm([1e-1, 1e-1])
    StateSpaceModel(F, V, G, w, x0, P0)
end

@testset "Kalman Filter" begin

    srand(1)

    @testset "Time invariant models" begin

        @testset "Building model" begin
            build_model()
        end

        @testset "Simulations" begin
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)
        end

        @testset "Filtering" begin
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)
            filt = kalman_filter(y, mod1)
        end

        @testset "Smoothing" begin
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)

            smooth = kalman_filter(y, mod1) |> kalman_smooth
            smooth2 = kalman_smooth(y, mod1)

            @test smooth.smoothed ≈ smooth2.smoothed atol=1e-2
        end

        @testset "Missing data" begin
            mod1 = build_model()
            x, y = simulate(mod1, 100)
            y[1:9:end] = NaN
            y[100] = NaN
            filt = kalman_filter(y, mod1)
            smooth = kalman_smooth(filt)
            smooth2 = kalman_smooth(y, mod1)
            @test !any(isnan, filt.filtered)
            @test !any(isnan, smooth.filtered)
            @test !any(isnan, smooth2.smoothed)
        end

        @testset "Return sizes" begin
            mod1 = build_model()
            x, y = TimeModels.simulate(mod1, 100)
            filt = kalman_filter(y, mod1)
            smooth = kalman_smooth(filt)
            smooth2 = kalman_smooth(y, mod1)

            @test size(filt.filtered) == size(filt.predicted)
            @test size(filt.filtered) == size(smooth.filtered)
            @test size(filt.filtered) == size(smooth.smoothed)
            @test size(filt.error_cov) == size(smooth.error_cov)
            @test size(filt.filtered) == size(smooth2.smoothed)
            @test size(filt.error_cov[:,:,1]) == size(smooth2.error_cov[1])
            @test size(filt.error_cov, 3) == length(smooth2.error_cov)
        end

        @testset "Linear regression test" begin
            m, b, s, dt = 5, 2, 2, .1
            t = 0:dt:10
            y_true = m*t + b
            input = 100*[sin.(t./2) sin.(t./4) cos.(t./2) cos.(t./4)] .+ 10
            y_noisy = [y_true zeros(length(t)) -y_true] +
                        100*[sin.(t./2) + sin.(t./4) sin.(t./2) + cos.(t./2) cos.(t./2) + cos.(t./4)] .+ 10 .+ randn(length(t), 3)
            lm = StateSpaceModel([1 dt; 0 1], zeros(2,2),
                                  [1. 0; 0 0; -1 0], s*eye(3),
                                  zeros(2), 100*eye(2),
                                  B=zeros(2, 4), D=[1. 1 0 0; 1 0 1 0; 0 0 1 1])
            lm_filt = kalman_filter(y_noisy, lm, u=input)
            @test lm_filt.filtered[end,1] ≈ y_true[end, 1] atol=3*sqrt(lm_filt.error_cov[1,1,end])

            lm_smooth = kalman_smooth(lm_filt)
            stderr = sqrt.(lm_smooth.error_cov[1:1,1:1,:][:])
            @test lm_filt.filtered[end:end,:] == lm_smooth.smoothed[end:end,:]
            @test all(abs.(y_true - lm_smooth.smoothed[:,1]) .< 3*stderr)
            @test ones(t) * lm_smooth.smoothed[1,2] ≈ lm_smooth.smoothed[:, 2] atol=1e-12

            # Repeat with DK smoother
            lm_smooth = kalman_smooth(y_noisy, lm, u=input)
            stderr = sqrt.(Float64[P[1,1] for P in lm_smooth.error_cov])
            @test all(abs.(y_true - lm_smooth.smoothed[:,1]) .< 3*stderr)
            @test ones(t) * lm_smooth.smoothed[1,2] ≈ lm_smooth.smoothed[:, 2] atol=1e-12

        end

    end

    @testset "Time varying models" begin

        @testset "Building Model" begin
            sinusoid_model(40)
        end

        @testset "Simulations" begin
            srand(1)
            fs = 256
            mod2 = sinusoid_model(4, fs = 256)
            x, y = TimeModels.simulate(mod2, fs*2)

            #= display(lineplot(collect(1:size(x, 1)) / fs, vec(y), width = 120, title = "Original Data")) =#
        end

        @testset "Filtering" begin
            srand(1)
            fs = 256
            mod2 = sinusoid_model(4, fs = 256)
            x, y = TimeModels.simulate(mod2, fs*2)

            @testset "Correct initial guess" begin
                filt = kalman_filter(y, mod2)
                @test filt.predicted[end, :] ≈ [0.5, -0.5] atol=0.3
            end

            @testset "Incorrect initial guess" begin
                mod3 = sinusoid_model(4, fs = 256, x0=[1.7, -0.2])
                filt = kalman_filter(y, mod3)
                @test filt.predicted[end, :] ≈ [0.5, -0.5] atol= 0.3

            end

            @testset "Model error" begin
                mod4 = sinusoid_model(4, fs = 256, x0=[1.7, -0.2], W=3.0)
                filt = kalman_filter(y, mod4)
                @test filt.predicted[end, :] ≈ [0.5, -0.5] atol= 0.3
            end

            @testset "Standalone log-likelihood function works" begin
                loglik1 = kalman_filter(y, mod2).loglik
                loglik2 = loglikelihood(y, mod2)
                @test loglik1 ≈ loglik2
            end

        end

        @testset "Smoothing" begin
            srand(1)
            fs = 256
            mod2 = sinusoid_model(4, fs = 8192)
            x, y = TimeModels.simulate(mod2, fs*10)
            smooth = kalman_smooth(y, sinusoid_model(4, fs = 8192, x0=[1.7, -0.2]) )
            @test mean(smooth.smoothed, 1) ≈ [0.5 -0.5] atol= 0.1

            #= x_est = round(smooth.smoothed[end:end, :], 3) =#
            #= display(lineplot(collect(1:size(x, 1)) / fs, vec(smooth.smoothed[1:end, 1]), width = 120, title="Smoothed State 1: $(x_est[1])")) =#
            #= display(lineplot(collect(1:size(x, 1)) / fs, vec(smooth.smoothed[1:end, 2]), width = 120, title="Smoothed State 2: $(x_est[2])")) =#
        end

        @testset "Sparse regression test" begin
            # To be fair, not an exemplary example of sparse matrices in action
            # since the state covariance off-diagonals are small but nonzero
            n, m, s = 5000, 10, 5.
            x = 100*rand(n, m)
            coeffs = randn(m)
            y = x * coeffs + s*randn(n)

            I_m   = speye(m)
            I0_m  = spzeros(m,m)
            S     = diagm(s)
            mlm   = StateSpaceModel(
                _->I_m, # A
                _->I0_m, # V
                t->x[t:t,:], # C
                _->S, # W
                zeros(m), # x1
                speye(m) # P1
            )
            mlm_smooth = kalman_smooth(y, mlm)
            @test ones(n,m) .* mlm_smooth.smoothed[1:1,:] ≈ mlm_smooth.smoothed atol=1e-12
            @test vec(mlm_smooth.smoothed[1:1,:]) ≈ coeffs atol=1e-2

            @testset "Correct failure" begin
                @test_throws MethodError StateSpaceModel([1 0.1; 0 1], zeros(2,3), zeros(2,2),
                    [1. 0; 0 0; -1 0], [1. 1 0 0; 1 0 1 0; 0 0 1 1], eye(3), zeros(2), 100*eye(2))
            end
        end

    end

end


using Distributions

function dists(specs; N=50)
  X = Float64[]
  Y = Float64[]
  for (x, y, σ) in specs
    xd = rand(Normal(x, σ), N)
    yd = rand(Normal(y, σ), N)
    append!(X, xd)
    append!(Y, yd)
  end
  X, Y
end

N = 1000000
x, y = dists([(2,2,0.02), (2,-2,0.1), (-2,-2,0.5), (-2,2,1.0), (0,0,3.0)], N=N)

println("# of points: ", length(x))

using GR

setviewport(0.1, 0.95, 0.1, 0.95)
setwindow(-10, 10, -10, 10)
setcharheight(0.02)
axes2d(0.5, 0.5, -10, -10, 4, 4, -0.005)
setcolormap(GR.COLORMAP_HOT)

@time shadepoints(x, y, xform=GR.XFORM_EQUALIZED)

updatews()


function trinary_to_decimal(str::AbstractString)
    typeof(match(r"^[0-2]+$", str)) == Void && return 0
    mapreduce(i->(i[2]-'0')*3^i[1], +, zip(0:length(str), reverse(str)))
end


@doc raw"""
`DAE`: Differential Algebraic Equation

Defines a differential algebraic initial value problem
```math
\begin{align*}
\dot{q} (t) &= v(t, q(t)) + u(t, q(t), \lambda(t)) , & q(t_{0}) &= q_{0} , \\
0 &= \phi (t, q(t), \lambda(t)) , & \lambda(t_{0}) &= \lambda_{0} ,
\end{align*}
```
with vector field ``v``, projection ``u``, algebraic constraint ``\phi=0``,
initial conditions ``q_{0}`` and ``\lambda_{0}``, the dynamical variable ``q``
taking values in ``\mathbb{R}^{m}`` and the algebraic variable ``\lambda``
taking values in ``\mathbb{R}^{n}``.

### Fields

* `d`: dimension of dynamical variable ``q`` and the vector field ``v``
* `m`: dimension of algebraic variable ``\lambda`` and the constraint ``\phi``
* `n`: number of initial conditions
* `v`: function computing the vector field
* `u`: function computing the projection
* `ϕ`: algebraic constraint
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q``
* `λ₀`: initial condition for algebraic variable ``\lambda``

The function `v`, providing the vector field, takes three arguments,
`v(t, q, v)`, the functions `u` and `ϕ`, providing the projection and the
algebraic constraint take four arguments, `u(t, q, λ, u)` and `ϕ(t, q, λ, ϕ)`,
where `t` is the current time, `q` and `λ` are the current solution vectors,
and `v`, `u` and `ϕ` are the vectors which hold the result of evaluating the
vector field ``v``, the projection ``u`` and the algebraic constraint ``\phi``
on `t`, `q` and `λ`.

### Example

```julia
    function v(t, q, v)
        v[1] = q[1]
        v[2] = q[2]
    end

    function u(t, q, λ, u)
        u[1] = +λ[1]
        u[2] = -λ[1]
    end

    function ϕ(t, q, λ, ϕ)
        ϕ[1] = q[2] - q[1]
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ₀ = [0.]

    dae = DAE(v, u, ϕ, t₀, q₀, λ₀)

```
"""
struct DAE{dType <: Number, tType <: Number, vType <: Function, uType <: Function,
           ϕType <: Function, pType <: Union{Tuple,Nothing}, N} <: AbstractEquationDAE{dType, tType}

    d::Int
    m::Int
    n::Int
    v::vType
    u::uType
    ϕ::ϕType
    t₀::tType
    q₀::Array{dType, N}
    λ₀::Array{dType, N}
    parameters::pType
    periodicity::Vector{dType}

    function DAE(DT::DataType, N::Int, d::Int, m::Int, n::Int,
                 v::vType, u::uType, ϕ::ϕType, t₀::tType,
                 q₀::AbstractArray{dType}, λ₀::AbstractArray{dType};
                 parameters=nothing, periodicity=zeros(DT,d)) where {
                        dType <: Number, tType <: Number, vType <: Function,
                        uType <: Function, ϕType <: Function}

        @assert d == size(q₀,1)
        @assert m == size(λ₀,1)
        @assert n == size(q₀,2) == size(λ₀,2)
        @assert d ≥ m
        @assert ndims(q₀) == ndims(λ₀) == N ∈ (1,2)

        new{DT, tType, vType, uType, ϕType, typeof(parameters), N}(d, m, n, v, u, ϕ, t₀,
                convert(Array{DT}, q₀), convert(Array{DT}, λ₀), parameters, periodicity)
    end
end

function DAE(v, u, ϕ, t₀, q₀::AbstractArray{DT}, λ₀::AbstractArray{DT}; kwargs...) where {DT}
    DAE(DT, ndims(q₀), size(q₀,1), size(λ₀,1), size(q₀,2), v, u, ϕ, t₀, q₀, λ₀; kwargs...)
end

function DAE(v, u, ϕ, q₀, λ₀; kwargs...)
    DAE(v, u, ϕ, zero(eltype(q₀)), q₀, λ₀; kwargs...)
end

Base.hash(dae::DAE, h::UInt) = hash(dae.d, hash(dae.m, hash(dae.n, hash(dae.v,
        hash(dae.u, hash(dae.t₀, hash(dae.q₀, hash(dae.λ₀,
        hash(dae.periodicity, hash(dae.parameters, h))))))))))

Base.:(==)(dae1::DAE, dae2::DAE) = (
                                dae1.d == dae2.d
                             && dae1.m == dae2.m
                             && dae1.n == dae2.n
                             && dae1.v == dae2.v
                             && dae1.u == dae2.u
                             && dae1.t₀ == dae2.t₀
                             && dae1.q₀ == dae2.q₀
                             && dae1.λ₀ == dae2.λ₀
                             && dae1.parameters == dae1.parameters
                             && dae1.periodicity == dae1.periodicity)

function Base.similar(dae::DAE, q₀, λ₀=get_λ₀(q₀, dae.λ₀); kwargs...)
    similar(dae, dae.t₀, q₀, λ₀; kwargs...)
end

function Base.similar(dae::DAE, t₀::TT, q₀::AbstractArray{DT}, λ₀::AbstractArray{DT}=get_λ₀(q₀, dae.λ₀);
                      parameters=dae.parameters, periodicity=dae.periodicity) where {DT  <: Number, TT <: Number}
    @assert dae.d == size(q₀,1)
    @assert dae.m == size(λ₀,1)
    DAE(dae.v, dae.u, dae.ϕ, t₀, q₀, λ₀; parameters=parameters, periodicity=periodicity)
end

@inline Base.ndims(dae::DAE) = dae.d

@inline periodicity(equation::DAE) = equation.periodicity


using Distributions

function test_gaussian_recon_loss_layer(backend::Backend, T, eps)
    println("-- Testing GaussianReconLossLayer on $(typeof(backend)){$T}...")

    ############################################################
    # Prepare Data for Testing
    ############################################################
    dims = (3, 4, 5)
    tensor_dim = length(dims)
    p, n = prod(dims[1:end-1]), dims[end]
    println("        > $dims")

    ############################################################
    # Setup
    ############################################################
    layer = GaussianReconLossLayer(bottoms=[:mu, :sigma, :x])

    mu = rand(p, n)  # sigmoid output is in [0, 1]
    sigma = max(1e-2, 0.2rand(p, n)) # PReLU output is > 0
    x = max(1e-2, rand(dims))  # normalized pixel values are in [0, 1]

    inputs = Blob[
        make_blob(backend, convert(Array{T}, mu)),
        make_blob(backend, convert(Array{T}, sigma)),
        make_blob(backend, convert(Array{T}, x))]

    diffs = Blob[
        make_blob(backend, rand(T, size(inputs[1]))),  # gets overwritten
        make_blob(backend, rand(T, size(inputs[2]))),  # gets overwritten
        NullBlob()]  # the data is fixed

    state = setup(backend, layer, inputs, diffs)

    ############################################################
    # Forward Propagation
    ############################################################
    forward(backend, state, inputs)

    loss = 0.
    for i in 1:n
        Xi_given_zi = MvNormal(mu[:, i], sigma[:, i])  # a random var
        xi = slicedim(x, tensor_dim, i)[:]  # an instance
        loss -= logpdf(Xi_given_zi, xi)
    end
    loss /= n
    @test loss >= 0

    info("expected $loss; got $(state.loss)")
    @test -eps < loss - state.loss < eps

    ############################################################
    # Backward Propagation
    ############################################################
    backward(backend, state, inputs, diffs)

    for i in 1:2
        input_i = Array(T, size(inputs[i]))
        copy!(input_i, inputs[i])
        diff_i = Array(T, size(diffs[i]))
        copy!(diff_i, diffs[i])

        for k in 1:p, j in 1:n
            x0 = input_i[k, j]
            function f(x::FloatingPoint)
                input_i[k, j] = x
                copy!(inputs[i], input_i)
                forward(backend, state, inputs)
                input_i[k, j] = x0
                copy!(inputs[i], input_i)
                state.loss
            end
            test_deriv(f, x0, diff_i[k, j])
        end
    end

    ############################################################
    # Shutdown
    ############################################################
    shutdown(backend, state)
end

function test_gaussian_recon_loss_layer(backend::Backend)
    test_gaussian_recon_loss_layer(backend, Float64, 1e-8)
end

if test_cpu
    test_gaussian_recon_loss_layer(backend_cpu)
end
if test_gpu
    test_gaussian_recon_loss_layer(backend_gpu)
end


module TestPoints

using Test
using TreeParzen

objective(params) = params[:x]^2 + params[:y]^2
space = Dict(:x => HP.Uniform(:x, -10.0, 10.0), :y => HP.Uniform(:y, -10.0, 10.0))
points = [Dict(:x => 0.0, :y => 0.0), Dict(:x => 1.0, :y => 1.0)]

best = fmin(objective, space, 10, points)
@test iszero(best[:x])
@test iszero(best[:y])

# test to catch ArgumentError when number of iterations is lower than the number of points
@test_throws ArgumentError fmin(objective, space, 1, points)

end
true


export swish

doc"""
    Swish

Swish activation function.

```math
f(x) = x \cdot \sigma (\beta x)
```
where ``\beta`` is a leanable parameter.

# References
* Ramachandran et al. ["Searching for Activation Functions"](https://arxiv.org/abs/1710.05941), arXiv 2017.
"""
struct Swish
    beta::Var
end
Swish(::Type{T}) where T = Swish(zerograd(ones(T,1)))
(f::Swish)(x) = swish(x, f.beta)

swish(x::Var, beta::Var) = Var(swish.(x.data,beta.data), (swish,x,beta))
swish(x::T, beta::T) where T = x * sigmoid(beta*x)

function addgrad!(y::Var, ::typeof(swish), x::Var, beta::Var)
    isvoid(x.grad) || ∇swish_x!(y.data, y.grad, x.data, x.grad, beta.data)
    isvoid(beta.grad) || ∇swish_beta!(y.data, y.grad, x.data, beta.data, beta.grad)
end

function ∇swish_x!(y::Array{T}, gy::Array{T}, x::Array{T}, gx::Array{T}, beta::Vector{T}) where T
    @inbounds for i = 1:length(x)
        s = sigmoid(beta[1] * x[i])
        gx[i] += gy[i] * (s + beta[1]*x[i] * s * (1-s))
    end
end

function ∇swish_beta!(y::Array{T}, gy::Array{T}, x::Array{T}, beta::Vector{T}, gbeta::Vector{T}) where T
    @inbounds for i = 1:length(gy)
        s = sigmoid(beta[1] * x[i])
        gbeta[1] += gy[i] * x[i] * x[i] * s * (1-s)
    end
end


#
# General utils
#

#
# Utils for Akl Toussaint
#
"""
Creates the "T inverse" matrix described by wikipedia and
above that will right multiply points to get the Barycentric
coordinates
"""
function create_t_inv(a::Point{2}, b::Point{2}, c::Point{2})
    # Pull out points
    a1, a2 = a
    b1, b2 = b
    c1, c2 = c

    # Build t inverse matrix
    temp = (a1*b2 - a1*c2 - c1*b2) - (b1*a2 - b1*c2 - c1*a2)
    return [(b2-c2)/temp (c1-b1)/temp; (c2-a2)/temp (a1-c1)/temp]
end

# One liner to check if λ values are in (0, 1) and sum to less than 1
@inline check_λ(λ1, λ2) = (0 < λ1 < 1) && (0 < λ2 < 1) && (λ1 + λ2 < 1) ? false : true

"""
Determines which points in a vector of points lie inside or outside the
triangle determined by the triangle abc.
"""
function pt_out_tri{T}(x::Vector{Point{2, T}}, a::Point{2}, b::Point{2}, c::Point{2})
    # Get number of points in x
    npts = length(x)

    # Allocate space
    out = Array(Bool, npts)

    # Create our "T inverse" matrixs
    tinv = create_t_inv(a, b, c)

    @inbounds for n=1:npts
        # Get the vector we need to multiply by T inverse to get
        # the Barycentric coordinates
        temp = x[n] - c

        # Get Barycentric coordinates
        λ1, λ2 = tinv*temp
        out[n] = check_λ(λ1, λ2)
    end

    return out
end

#
# Wrap around set of points
#
"""
Takes a vector of points (which must be sorted either by
relative angle or x-value) and wraps around them by assuring
that they only make counter-clockwise turns.
"""
function wrappoints{T<:Real}(p::Vector{Point{2, T}})

    # Get number of points
    npts = length(p)

    # Put first two points inside the convex hull
    ch = [p[1], p[2]]
    sizehint!(ch, npts)

    # Starting value for number of extreme points
    m = 2

    # Will check each point
    @inbounds for i=3:npts
        # If it makes a clockwise turn then the
        # last element of the convex hull doesn't
        # belong, so pop it out. Make sure that at
        # least two elements are always in it
        while m>1 && cw(ch[m-1], ch[m], p[i])
            pop!(ch)
            m = m-1
        end
        push!(ch, p[i])
        m = m+1
    end

    return ch
end

#
# Utils for Monotone Chain
#
function split_xL_xU{T<:Real}(psort::Vector{Point{2, T}})
    # Allocate space to identify where points belong
    # We make endpoints true because we want them in both sets
    npts = length(psort)
    xL_bool = Array(Bool, npts)
    xL_bool[1] = true; xL_bool[end] = true
    xU_bool = Array(Bool, npts)
    xU_bool[1] = true; xU_bool[end] = true

    # Get the endpoints and slope
    pmin = psort[1]; pmax = psort[end]
    ls = LineSegment(pmin, pmax)

    # We check whether it is above or below the line between
    # the points, but could also check whether it was clockwise
    # or counter-clockwise (above/below seems like less operations)
    # Iterate through all points except endpoints
    @inbounds for i=2:npts-1
        # Pull out x and y values for ith point
        p_i = psort[i]
        xi, yi = p_i[1], p_i[2]

        # Compute where y on line would be
        yi_line = evaluatey(ls, xi)

        if yi_line <= yi
            xL_bool[i] = false
            xU_bool[i] = true
        else
            xL_bool[i] = true
            xU_bool[i] = false
        end
    end
    xL = psort[xL_bool]
    xU = psort[xU_bool]

    return xL, xU
end

#
# Additional functions for Point{2, T} that will be useful
#

# Matrix times point
function _unsafe_mult_2by2_point{T}(A::Matrix{T}, x::Point{2, T})
    x1 = x[1]
    x2 = x[2]
    out = (A[1, 1]*x1 + A[1, 2]*x2, A[2, 1]*x1 + A[2, 2]*x2)
    return out
end

function _unsafe_mult_Nby2_point{T}(A::Matrix{T}, x::Point{2, T})
    nr, nc = size(A)
    x1 = x[1]
    x2 = x[2]

    out = Vector(T, nr)
    for i=1:nr
        out[i] = A[i, 1]*x1+A[i, 2]*x2
    end
    return out
end

function *(A::Matrix, x::Point{2})
    nr, nc = size(A)

    # Depending on size call the right function
    if (nc==nr==2)
        out = _unsafe_mult_2by2_point(A, x)
    elseif (nc==2)
        out = _unsafe_mult_nby2(A, x)
    else
        error("Size is wrong. Number of cols should be 2 (currently is $nc)")
    end

    return out
end

function extremax{T}(p::Vector{Point{2, T}})
    # Get number of points
    npts = length(p)

    # Set initial values
    minval, minind = Inf, 0
    maxval, maxind = -Inf, 0

    for i=1:npts
        cx, cy = p[i]
        minval, minind = cx < minval ? (cx, i) : (minval, minind)
        maxval, maxind = cx > maxval ? (cx, i) : (maxval, maxind)
    end

    return minval, minind, maxval, maxind
end

function extremay{T}(p::Vector{Point{2, T}})
    # Get number of points
    npts = length(p)

    # Set initial values
    minval, minind = Inf, 0
    maxval, maxind = -Inf, 0

    for i=1:npts
        cx, cy = p[i]
        minval, minind = cy < minval ? (cy, i) : (minval, minind)
        maxval, maxind = cy > maxval ? (cy, i) : (maxval, maxind)
    end

    return minval, minind, maxval, maxind
end

@inline distance{T}(a::Point{2, T}, b::Point{2, T}) = sqrt((b[1] - a[1])^2 - (b[2] - a[2])^2)

# Comparisons
isless{T}(x::Point{2, T}, y::Point{2, T}) = x[1] < y[1] ? true : (x[1]==y[1] && x[2]<y[2]) ? true : false
isequal{T}(x::Point{2, T}, y::Point{2, T}) = (x[1]==y[1]) && (x[2]==y[2]) ? true : false

# Orientation
@inline orientation(a::Point{2}, b::Point{2}, c::Point{2}) =
    (b[1]-a[1])*(c[2]-a[2]) - (b[2]-a[2])*(c[1]-a[1])
cw(a::Point{2}, b::Point{2}, c::Point{2}) = orientation(a, b, c) < 0.0 ? true : false
ccw(a::Point{2}, b::Point{2}, c::Point{2}) = orientation(a, b, c) >= 0.0 ? true : false


using KernelAbstractions
using StaticArrays
using Test

const TILE_DIM = 32

@kernel function coalesced_matmul_kernel!(output, @Const(input1), @Const(input2), N, R, M,
                                             ::Val{BANK}=Val(1)) where BANK
     gi, gj = @index(Group, NTuple)
     i, j   = @index(Local, NTuple)

     TILE_DIM = @uniform groupsize()[1]

     # +1 to avoid bank conflicts on shared memory
     tile1 = @localmem eltype(output) (TILE_DIM+BANK, TILE_DIM)
     tile2 = @localmem eltype(output) (TILE_DIM+BANK, TILE_DIM)

     # private variable for tile output
     outval = @private eltype(output) 1
     @inbounds outval[1] = -zero(eltype(output))

     @uniform N = size(output, 1)
     # number of tiles depends on inner dimension
     @uniform NUM_TILES = div(R + TILE_DIM - 1, TILE_DIM)

     # loop over all tiles needed for this calculation
     for t in 0:NUM_TILES-1
         # Can't use @index(Global), because we use a smaller ndrange
         I = (gi-1) * TILE_DIM + i
         J = (gj-1) * TILE_DIM + j

         # load inputs into tiles, with bounds checking for non-square matrices
         if I <= N && t*TILE_DIM + j <= R
             @inbounds tile1[i, j] = input1[I, t*TILE_DIM + j]
         else
             @inbounds tile1[i, j] = 0.0
         end
         if t*TILE_DIM + i <= R && J <= M
             @inbounds tile2[i, j] = input2[t*TILE_DIM + i, J]
         else
             @inbounds tile2[i, j] = 0.0
         end

         # wait for all tiles to be loaded
         @synchronize

         # get global values again
         I = (gi-1) * TILE_DIM + i
         J = (gj-1) * TILE_DIM + j

         # calculate value of spot in output, use temporary value to allow for vectorization
         out = zero(eltype(output))
         @simd for k in 1:TILE_DIM
             @inbounds out += tile1[i, k] * tile2[k, j]
         end
         outval[1] += out

         @synchronize
     end

     # get global indices again
     I = (gi-1) * TILE_DIM + i
     J = (gj-1) * TILE_DIM + j

     # save if inbounds
     if I <= N && J <= M
         @inbounds output[I, J] = outval[1]
     end
end

N = 1024
R = 512
M = 2048
A = rand(N, R)
B = rand(R, M)
C = zeros(N, M)

kern = coalesced_matmul_kernel!(CPU(), (TILE_DIM, TILE_DIM))

wait(kern(C, A, B, N, R, M, ndrange=size(C)))
@test isapprox(A*B, C)


module Examples

include("example_base.jl")

include("arnold_cat_map.jl")
include("bakers_map.jl")
include("henon_map.jl")
include("linz_sprott_99.jl")
include("lorenz_63.jl")
include("lozi_map.jl")
include("rnn.jl")
include("standard_map.jl")
include("van_der_pol.jl")

end

using .Examples.ExampleBase: LEDemo, LEExample,
    ContinuousExample, DiscreteExample

# Make LyapunovExponents.lorenz_63 etc. available:
using .Examples.ArnoldCatMap
using .Examples.BakersMap
using .Examples.HenonMap
using .Examples.LinzSprott99
using .Examples.Lorenz63
using .Examples.LoziMap
using .Examples.RNN
using .Examples.StandardMap
using .Examples.VanDerPol

const DEMOS = [
    lorenz_63,
    linz_sprott_99,
    van_der_pol,
    beer_95,
    henon_map,
    standard_map,
    bakers_map,
    arnold_cat_map,
    lozi_map,
]


# Internal facilities for fast random number generation

immutable RandIntSampler  # for generating Int samples in [0, K-1]
	a::Int
	Ku::Uint
	U::Uint

	RandIntSampler(K::Int) = (Ku = uint(K); new(1, Ku, div(typemax(Uint), Ku) * Ku))
	RandIntSampler(a::Int, b::Int) = (Ku = uint(b-a+1); new(a, Ku, div(typemax(Uint), Ku) * Ku))
end

function rand(s::RandIntSampler)
	x = rand(Uint)
	while x >= s.U
		x = rand(Uint)
	end
	s.a + int(rem(x, s.Ku))
end

randi(K::Int) = rand(RandIntSampler(K))
randi(a::Int, b::Int) = rand(RandIntSampler(a, b))

# draw a number from a binomial distribution

rand_binom(n::Real, p::Real) = int(ccall((:rbinom, :libRmath), Float64, (Float64, Float64), n, p))


#############################################
# Van der Pol Oscillator controlled by VSSM #
#############################################

K=500
w=1

#########
# Time  #
#########
δt=1e-3
LONG=Int(2e4)

########################################
# Kinematic Block Parameter (2nd order)#
########################################
Λ=12

####################################
# Parameters for Nominal Trajectory#
####################################

ω=0.5
Ampl=2

##########################
# Exact Model Parameters #
##########################

μₑ=0.4
ωₑ=0.46
αₑ=1
λₑ=0.1
mₑ=1

################################
# Approximate Model Parameters #
################################

μₐ=0.5
ωₐ=0.42
αₐ=0.9
λₐ=0.09
mₐ=0.8

###############
# Exact Model #
###############
function Exact(q,q_p,u)
  q_pp=(u+μₑ*(1-q^2)*q_p-ωₑ^2*q-αₐ*q^3-λₑ^5)/mₑ
  return q_pp
end

################
# Approx Model #
################

function Approx(q,q_p,q_pp)
  u=mₐ*q_pp-μₐ*(1-q^2)*q_p+ωₐ^2*q+αₐ*q^3+λₐ^5
  return u
end

#ErrorMetrics(hint,h,h_p) = Λ^2*hint+2*Λ*h+h_p # "S" block for 2nd order

#KinBlock(S,h,h_p,qN_pp) = K*tanh(S/w)+Λ^2*h+2*Λ*h_p+qN_pp #q_ppDes -t adja vissza

ErrorMetrics(hint,h,h_p) = Λ^2*hint+2*Λ*h+h_p # "S" block for 2nd order

KinBlock(S,h,h_p,qN_pp) = K*tanh(S/w)+Λ^2*h+2*Λ*h_p+qN_pp #q_ppDes -t adja vissza

function nominalTraj(t)
  qN=Ampl*sin(ω*t*δt)
  q_pN=ω*Ampl*cos(ω*t*δt)
  q_ppN=-Ampl*ω^2*sin(ω*t*δt)
  return [qN,q_pN,q_ppN]
end

q_mem=zeros(LONG)
q_p_mem=zeros(LONG)
q_pp_mem=zeros(LONG)
q_ppDes_mem=zeros(LONG)
q_N_mem=zeros(LONG)
q_pN_mem=zeros(LONG)
q_ppN_mem=zeros(LONG)
time_mem=zeros(LONG)
S_mem=zeros(LONG)
u_mem=zeros(LONG)

hint=0
##############
# Simulation #
##############

for t=1:LONG-1
  #create time
  time_mem[t]=δt*t
  #Nominal trajectory for the actual time frame
#  q_temp=nominalTraj(time_mem[t])
  q_temp=nominalTraj(t)
  q_N_mem[t] = q_temp[1]
  q_pN_mem[t] = q_temp[2]
  q_ppN_mem[t] = q_temp[3]
  #errors
  h=q_N_mem[t]-q_mem[t]
  #println(h)
  #Error Integral (Euler integral)
  hint=hint+δt*h
  #The time derivative of the error
  h_p=q_pN_mem[t]-q_p_mem[t]
  #the error metric
  S_mem[t]=ErrorMetrics(hint,h,h_p)
  #the kinematic block
  q_ppDes_mem[t]=KinBlock(S_mem[t],h,h_p,q_ppN_mem[t])
  #println(q_ppDes_mem[t])
  #the aprox model (inverse model)
  u_mem[t]=Approx(q_mem[t],q_p_mem[t],q_ppDes_mem[t])
  #the exact model
  q_pp_mem[t]=Exact(q_mem[t],q_p_mem[t],u_mem[t])
  #	Integrate back with Euler integral
  q_p_mem[t+1]=q_p_mem[t]+δt*q_pp_mem[t]
  q_mem[t+1]=q_mem[t]+q_p_mem[t]*δt
end #for

###############################
# Load matplotlib for Figures #
###############################

using PyPlot
#println(time_mem)
figure(1)
plot(time_mem[3:LONG-1],q_N_mem[3:LONG-1],color="#7684FF")
plot(time_mem[3:LONG-1],q_mem[3:LONG-1],color="#FFAA41")
grid("on")
title("Nominal and Realized Trajectory\n Nominal Trajectory components: bluish\n Realized Trajectory components:orangish ")
xlabel("Time [s]")
ylabel("q [m]")
#plot(time_mem[3:LONG-1],qN[3:LONG-1,1],color="#7684FF")
#plot(time_mem[3:LONG-1],qN[3:LONG-1,2],color="#2A40FF")
#plot(time_mem[3:LONG-1],qN[3:LONG-1,3],color="#3B427F")
#plot(time3d[3:LONG-1],qR[3:LONG-1,:],color="green")
#plot(time_mem[3:LONG-1],qR[3:LONG-1,1],color="#FFAA41")
#plot(time_mem[3:LONG-1],qR[3:LONG-1,2],color="#FF9F28")
#plot(time_mem[3:LONG-1],qR[3:LONG-1,3],color="#B2690E")

figure(2)
plot(time_mem[3:LONG-1],q_N_mem[3:LONG-1]-q_mem[3:LONG-1])

examples/VSSM_SISO/van der Pol Oscillator/van_der_Pol_VSSM.jl


using Test
using Distributions
using StableRNGs
using StatsBase
using LinearAlgebra
using Random

using EnsembleKalmanProcesses.ParameterDistributions

@testset "ParameterDistributions" begin
    @testset "ParameterDistributionType" begin
        # Tests for the ParameterDistributionType
        d = Parameterized(Gamma(2.0, 0.8))
        @test d.distribution == Gamma(2.0, 0.8)

        d = Samples([1 2 3; 4 5 6])
        @test d.distribution_samples == [1 2 3; 4 5 6]
        d = Samples([1 4; 2 5; 3 6]; params_are_columns = false)
        @test d.distribution_samples == [1 2 3; 4 5 6]
        d = Samples([1, 2, 3, 4, 5, 6])
        @test d.distribution_samples == [1 2 3 4 5 6]
        d = Samples([1, 2, 3, 4, 5, 6]; params_are_columns = false)
        @test d.distribution_samples == reshape([1, 2, 3, 4, 5, 6], :, 1)

        # Vector container -
        d = VectorOfParameterized(repeat([Normal(0, 1)], 5))
        @test d.distribution == repeat([Normal(0, 1)], 5)
    end
    @testset "ConstraintType" begin
        # Tests for the ConstraintType
        # The predefined transforms
        c1 = bounded_below(0.2)
        @test isapprox(c1.constrained_to_unconstrained(1.0) - (log(1.0 - 0.2)), 0)
        @test isapprox(c1.unconstrained_to_constrained(0.0) - (exp(0.0) + 0.2), 0)

        c2 = bounded_above(0.2)
        @test isapprox(c2.constrained_to_unconstrained(-1.0) - (log(0.2 - -1.0)), 0)
        @test isapprox(c2.unconstrained_to_constrained(10.0) - (0.2 - exp(10.0)), 0)

        c3 = bounded(-0.1, 0.2)
        @test isapprox(c3.constrained_to_unconstrained(0.0) - (log((0.0 - -0.1) / (0.2 - 0.0))), 0)
        @test isapprox(c3.unconstrained_to_constrained(1.0) - ((0.2 * exp(1.0) + -0.1) / (exp(1.0) + 1)), 0)
        @test_throws DomainError bounded(0.2, -0.1)

        #an example with user defined invertible transforms
        c_to_u = (x -> 3 * x + 14)
        u_to_c = (x -> (x - 14) / 3)

        c4 = Constraint(c_to_u, u_to_c)
        @test isapprox(c4.constrained_to_unconstrained(5.0) - c_to_u(5.0), 0)
        @test isapprox(c4.unconstrained_to_constrained(5.0) - u_to_c(5.0), 0)

        #length, size
        @test length(c1) == 1
        @test size(c1) == (1,)

    end

    @testset "ParameterDistribution: Build and combine" begin

        # create single ParameterDistribution
        d = VectorOfParameterized([Normal(0, 1), Gamma(2.0, 0.8)])
        c_mismatch = no_constraint()
        c_wrongtype = [no_constraint(), 3.0]
        c = [no_constraint(), no_constraint()]
        name = "normal_and_gamma"

        @test_throws ArgumentError ParameterDistribution(d, c_wrongtype, name) #wrong type of constraint
        @test_throws DimensionMismatch ParameterDistribution(d, c_mismatch, name) #wrong number of constraints

        # test checks on stored information
        u = ParameterDistribution(d, c, name)
        @test u.distribution == [d]
        @test u.constraint == c #as c is already a vector
        @test u.name == [name]

        d = Parameterized(Normal(0, 1))
        c = [no_constraint()]
        u = ParameterDistribution(d, c, name)
        @test u.constraint == c #as c is already a vector

        # test concatenation
        d1 = Parameterized(Normal(0, 1))
        c1 = no_constraint()
        name1 = "unconstrained_normal"
        u1 = ParameterDistribution(d1, c1, name1)
        @test u1.constraint == [c1]

        d2 = VectorOfParameterized(repeat([MvNormal(3, 0.2)], 4))
        c2 = repeat([no_constraint()], 12) #3D distribution repeated 4 times has 12 constraints
        name2 = "three_unconstrained_MvNs"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c3 = [bounded(10, 15), no_constraint()]
        name3 = "constrained_sampled"
        u3 = ParameterDistribution(d3, c3, name3)

        @test_throws ArgumentError ParameterDistribution([u1, u2])

        u = combine_distributions([u1, u2, u3])
        @test u.distribution == [d1, d2, d3]
        @test u.constraint == cat([[c1], c2, c3]..., dims = 1)
        @test u.name == [name1, name2, name3]

    end

    @testset "ParameterDistribution: Dictionary interface" begin

        # dictionary - testing test checks on the distribution elements
        d = Parameterized(Normal(0, 1))
        c = no_constraint()
        name = "unconstrained_normal"

        param_dict_fail = [Dict("a" => 1), 1]
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not an array of Dicts

        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not a Dict

        param_dict_fail = Dict("distribution" => d)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not all the keys
        param_dict_fail = Dict("distribution" => 1, "constraint" => c, "name" => name)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right distribution type
        param_dict_fail = Dict("distribution" => d, "constraint" => 1, "name" => name)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right constraint type
        param_dict_fail = Dict("distribution" => d, "constraint" => [1, no_constraint(), 3], "name" => name)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right constraint type
        param_dict_fail = Dict("distribution" => d, "constraint" => c, "name" => 1)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right name type
        param_dict_fail = Dict("distribution" => d, "constraint" => [no_constraint(), no_constraint()], "name" => name)
        @test_throws DimensionMismatch ParameterDistribution(param_dict_fail) # wrong number of constraints

        param_dict = Dict("distribution" => d, "constraint" => c, "name" => name)
        u = ParameterDistribution(param_dict)
        @test u.distribution == [d]
        @test u.constraint == [c]
        @test u.name == [name]

        c_arr = [no_constraint()]
        param_dict = Dict("distribution" => d, "constraint" => c_arr, "name" => name)
        u = ParameterDistribution(param_dict)
        @test u.constraint == [c]

        # Tests for the ParameterDistribution
        d = Parameterized(MvNormal(4, 0.1))
        c = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name = "constrained_mvnormal"
        param_dict = Dict("distribution" => d, "constraint" => c, "name" => name)

        u = ParameterDistribution(param_dict)
        @test u.distribution == [d]
        @test u.constraint == c
        @test u.name == [name]
        param_dict_fail = param_dict
        param_dict_fail["constraint"] = c[1]
        @test_throws DimensionMismatch ParameterDistribution(param_dict_fail) # wrong number of constraints

        # Tests for the ParameterDistribution
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        param_dict1 = Dict("distribution" => d1, "constraint" => c1, "name" => name1)

        d2 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c2 = [bounded(10, 15), no_constraint()]
        name2 = "constrained_sampled"
        param_dict2 = Dict("distribution" => d2, "constraint" => c2, "name" => name2)

        param_dict_array = [param_dict1, param_dict2]
        u = ParameterDistribution(param_dict_array)
        @test u.distribution == [d1, d2]
        @test u.constraint == cat([c1, c2]..., dims = 1)
        @test u.name == [name1, name2]

    end

    @testset "getter functions" begin
        # setup for the tests:
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = VectorOfParameterized(repeat([Beta(2, 2)], 3))
        c3 = repeat([no_constraint()], 3)
        name3 = "vector_beta"
        u3 = ParameterDistribution(d3, c3, name3)

        u = combine_distributions([u1, u2, u3])

        # Test for get_dimension(s)
        @test ndims(u1) == 4
        @test ndims(u2) == 1
        @test ndims(u) == 8

        @test get_dimensions(u1) == [4]
        @test get_dimensions(u2) == [1]
        @test get_dimensions(u) == [4, 1, 3]

        # Tests for get_name
        @test get_name(u1) == [name1]
        @test get_name(u) == [name1, name2, name3]

        # Tests for get_n_samples
        @test typeof(get_n_samples(u)[name1]) <: String
        @test typeof(get_n_samples(u)[name3]) <: String
        @test get_n_samples(u)[name2] == 4

        # Tests for get_distribution
        @test get_distribution(d1) == MvNormal(4, 0.1)
        @test get_distribution(u1)[name1] == MvNormal(4, 0.1)
        @test typeof(get_distribution(d2)) == Array{Int64, 2}
        @test typeof(get_distribution(u2)[name2]) == Array{Int64, 2}

        @test get_distribution(d3) == repeat([Beta(2, 2)], 3)
        @test get_distribution(u3)[name3] == repeat([Beta(2, 2)], 3)

        d = get_distribution(u)
        @test d[name1] == MvNormal(4, 0.1)
        @test typeof(d[name2]) == Array{Int64, 2}

        # Test for get_all_constraints
        @test get_all_constraints(u) == cat([c1, c2, c3]..., dims = 1)
    end

    @testset "statistics functions" begin

        # setup for the tests:
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"

        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = VectorOfParameterized([Beta(2, 2), MvNormal(2, 0.1)])
        c3 = [no_constraint(), no_constraint(), no_constraint()] # d3 has 3 dimensions
        name3 = "unconstrained_beta_and_MvN"
        u3 = ParameterDistribution(d3, c3, name3)

        u = combine_distributions([u1, u2])

        d4 = Samples([1 2 3 4 5 6 7 8; 8 7 6 5 4 3 2 1])
        c4 = [no_constraint(), no_constraint()]
        name4 = "constrained_MVsampled"
        u4 = ParameterDistribution(d4, c4, name4)

        v = combine_distributions([u1, u2, u3, u4])

        # Tests for sample distribution
        seed = 2020
        Random.seed!(seed)
        s1 = rand(MvNormal(4, 0.1), 1)
        Random.seed!(seed)
        @test sample(u1) == s1

        Random.seed!(seed)
        s1 = rand(MvNormal(4, 0.1), 3)
        Random.seed!(seed)
        @test sample(u1, 3) == s1

        Random.seed!(seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        Random.seed!(seed)
        @test sample(u2) == s2

        Random.seed!(seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 3)
        s2 = d2.distribution_samples[:, idx]
        Random.seed!(seed)
        @test sample(u2, 3) == s2

        Random.seed!(seed)
        s3 = zeros(3, 1)
        s3[1] = rand(Beta(2, 2))
        s3[2:3] = rand(MvNormal(2, 0.1))
        Random.seed!(seed)
        @test sample(u3) == s3

        Random.seed!(seed)
        s1 = sample(u1, 3)
        s2 = sample(u2, 3)
        s3 = sample(u3, 3)
        s4 = sample(u4, 3)
        Random.seed!(seed)
        s = sample(v, 3)
        @test s == cat([s1, s2, s3, s4]..., dims = 1)

        #Test for get_logpdf
        @test_throws ErrorException get_logpdf(u, zeros(ndims(u)))
        x_in_bd = [0.5, 0.5, 0.5]
        Random.seed!(seed)
        lpdf3 = sum([logpdf(Beta(2, 2), x_in_bd[1])[1], logpdf(MvNormal(2, 0.1), x_in_bd[2:3])[1]]) #throws deprecated warning without "."

        Random.seed!(seed)
        @test isapprox(get_logpdf(u3, x_in_bd) - lpdf3, 0.0; atol = 1e-6)
        @test_throws DimensionMismatch get_logpdf(u3, [0.5, 0.5])

        #Test for cov, var
        block_cov = cat([cov(d1), var(d2), cov(d3), cov(d4)]..., dims = (1, 2))
        @test isapprox(cov(v) - block_cov, zeros(ndims(v), ndims(v)); atol = 1e-6)
        block_var = [block_cov[i, i] for i in 1:size(block_cov)[1]]
        @test isapprox(var(v) - block_var, zeros(ndims(v)); atol = 1e-6)

        #Test for mean
        means = cat([mean(d1), mean(d2), mean(d3), mean(d4)]..., dims = 1)
        @test isapprox(mean(v) - means, zeros(ndims(v)); atol = 1e-6)

    end

    @testset "statistics functions: explict RNG" begin

        # setup for the tests:
        rng_seed = 1234
        test_d = MvNormal(4, 0.1)
        test_d3a = Beta(2, 2)
        test_d3b = MvNormal(2, 0.1)
        d0 = Parameterized(test_d)

        d1 = Parameterized(test_d)
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = VectorOfParameterized([Beta(2, 2), MvNormal(2, 0.1)])
        c3 = repeat([no_constraint()], 3)
        name3 = "beta_and_normal"
        u3 = ParameterDistribution(d3, c3, name3)

        # Tests for sample distribution
        rng1 = Random.MersenneTwister(rng_seed)
        @test sample(copy(rng1), d0) == rand(copy(rng1), test_d, 1)
        @test sample(copy(rng1), d0, 3) == rand(copy(rng1), test_d, 3)

        @test sample(copy(rng1), u1) == rand(copy(rng1), test_d, 1)
        @test sample(copy(rng1), u1, 3) == rand(copy(rng1), test_d, 3)

        idx = StatsBase.sample(copy(rng1), collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        @test sample(copy(rng1), u2) == s2
        @test sample(copy(rng1), u2, 1) == s2
        @test sample(copy(rng1), d2) == s2
        @test sample(copy(rng1), d2, 1) == s2

        rng1_new = copy(rng1) # (copy here, as internally u3 samples rng1 sequentially)
        @test sample(copy(rng1), u3) == cat([rand(rng1_new, test_d3a, 1), rand(rng1_new, test_d3b, 1)]..., dims = 1)
        rng1_new = copy(rng1)
        @test sample(copy(rng1), u3, 3) ==
              cat([reshape(rand(rng1_new, test_d3a, 3), :, 3), rand(rng1_new, test_d3b, 3)]..., dims = 1)
        rng1_new = copy(rng1)
        @test sample(copy(rng1), d3) == cat([rand(rng1_new, test_d3a, 1), rand(rng1_new, test_d3b, 1)]..., dims = 1)
        rng1_new = copy(rng1)
        @test sample(copy(rng1), d3, 1) == cat([rand(rng1_new, test_d3a, 1), rand(rng1_new, test_d3b, 1)]..., dims = 1)

        # try it again with different RNG; use StableRNG since Random doesn't provide a
        # second seedable algorithm on julia <=1.7
        rng2 = StableRNG(rng_seed)
        @test sample(copy(rng2), d0) == rand(copy(rng2), test_d, 1)
        @test sample(copy(rng2), d0, 3) == rand(copy(rng2), test_d, 3)

        @test sample(copy(rng2), u1) == rand(copy(rng2), test_d, 1)
        @test sample(copy(rng2), u1, 3) == rand(copy(rng2), test_d, 3)

        idx = StatsBase.sample(copy(rng2), collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        @test sample(copy(rng2), u2) == s2

        rng2_new = copy(rng2)
        @test sample(copy(rng2), u3) == cat([rand(rng2_new, test_d3a, 1), rand(rng2_new, test_d3b, 1)]..., dims = 1)
        rng2_new = copy(rng2)
        @test sample(copy(rng2), u3, 3) ==
              cat([reshape(rand(rng2_new, test_d3a, 3), :, 3), rand(rng2_new, test_d3b, 3)]..., dims = 1)

        # test that optional parameter defaults to Random.GLOBAL_RNG, for all methods.
        # reset the global seed instead of copying the rng object's state
        rng_seed = 2468
        Random.seed!(rng_seed)
        test_lhs = sample(d0)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 1)

        Random.seed!(rng_seed)
        test_lhs = sample(d0, 3)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 3)

        Random.seed!(rng_seed)
        test_lhs = sample(u1)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 1)

        Random.seed!(rng_seed)
        test_lhs = sample(u1, 3)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 3)

        Random.seed!(rng_seed)
        test_lhs = sample(d0)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 1)

        Random.seed!(rng_seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 1)
        test_lhs = d2.distribution_samples[:, idx]
        Random.seed!(rng_seed)
        @test test_lhs == sample(u2)

        Random.seed!(rng_seed)
        test_lhs = sample(d3)
        Random.seed!(rng_seed)
        @test test_lhs == cat([rand(test_d3a, 1), rand(test_d3b, 1)]..., dims = 1)

    end

    @testset "transform functions" begin
        #setup for the tests
        tol = 1e-8
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        param_dict1 = Dict("distribution" => d1, "constraint" => c1, "name" => name1)
        u1 = ParameterDistribution(param_dict1)

        d2 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c2 = [bounded(10, 15), no_constraint()]
        name2 = "constrained_sampled"
        param_dict2 = Dict("distribution" => d2, "constraint" => c2, "name" => name2)
        u2 = ParameterDistribution(param_dict2)

        param_dict = [param_dict1, param_dict2]
        u = ParameterDistribution(param_dict)

        x_unbd = rand(MvNormal(6, 3), 1000)  #6 x 1000
        # Tests for transforms
        x_real_constrained1 = mapslices(x -> transform_unconstrained_to_constrained(u1, x), x_unbd[1:4, :]; dims = 1)
        @test isapprox(
            x_unbd[1:4, :] -
            mapslices(x -> transform_constrained_to_unconstrained(u1, x), x_real_constrained1; dims = 1),
            zeros(size(x_unbd[1:4, :]));
            atol = tol,
        )

        x_real_constrained2 = mapslices(x -> transform_unconstrained_to_constrained(u2, x), x_unbd[5:6, :]; dims = 1)
        @test isapprox(
            x_unbd[5:6, :] -
            mapslices(x -> transform_constrained_to_unconstrained(u2, x), x_real_constrained2; dims = 1),
            zeros(size(x_unbd[5:6, :]));
            atol = tol,
        )

        x_real = mapslices(x -> transform_unconstrained_to_constrained(u, x), x_unbd; dims = 1)
        x_unbd_tmp = mapslices(x -> transform_constrained_to_unconstrained(u, x), x_real; dims = 1)
        @test isapprox(x_unbd - x_unbd_tmp, zeros(size(x_unbd)); atol = tol)

        # Tests transforms for other input structures
        @test isapprox(transform_unconstrained_to_constrained(u1, x_unbd[1:4, :]), x_real_constrained1; atol = tol)
        @test isapprox(
            x_unbd[1:4, :] - transform_constrained_to_unconstrained(u1, x_real_constrained1),
            zeros(size(x_unbd[1:4, :]));
            atol = tol,
        )
        @test isapprox(transform_unconstrained_to_constrained(u2, x_unbd[5:6, :]), x_real_constrained2; atol = tol)
        @test isapprox(
            x_unbd[5:6, :] - transform_constrained_to_unconstrained(u2, x_real_constrained2),
            zeros(size(x_unbd[5:6, :]));
            atol = tol,
        )
        @test isapprox(
            transform_unconstrained_to_constrained(u2, [x_unbd[5:6, :], x_unbd[5:6, :]]),
            [x_real_constrained2, x_real_constrained2];
            atol = tol,
        )
    end
end


# Chimera_states_v4.jl
# ====================   u.p. 13.9.21

using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot

function chimera_ODE(du,u,p,t)
    ω, R, alpha, nosc = p
    @inbounds for i = 1:nosc
        term = 0.0
        for j = i-R:i+R
            jj = 1 + mod(j-1,nosc)
            term = term + sin( u[i] - u[jj] + alpha )
        end
        du[i] = ω - term / (2*R )
    end
    return
end

#  ------------------------------  main program  --------------------------------
fig, axs = subplots(1,3; figsize=(figx,figy))

for axi in 1:2
if axi == 2
    # case a:  with chaotic transient
    nosc = 40 # 50#  40 # 100 # 40    # number of oscillators
    R = 14 # 15 # 14 # 35  # 14  # coupling range
    ish = 102  # select snapshot
    tsamp = 50.  # sampling time
    tspan = 25000.  # total lengthof time series
else
    # case b:  no chaotic transient (or transient end much later), but more oscillators
    nosc = 100 # 50#  40 # 100 # 40    # number of oscillators
    R = 35 # 15 # 14 # 35  # 14  # coupling range
    ish = 102  # select snapshot
    tsamp = 50.  # sampling time
    tspan = 25000.  # total lengthof time series
end

indx = LinRange(1,nosc,nosc)
ω = 0.
alpha = 1.46

p = ω, R, alpha, nosc

iseed = 7117 # 5117
rng = MersenneTwister(iseed) # random initial conditiosn in [0 ,2*pi]
u0 = 2 .* pi .* (rand(rng,nosc) .- 0.5)

prob = ODEProblem(chimera_ODE, u0, tspan, p)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8, saveat = tsamp)
uout = sol.u
tvec = sol.t
nts = length(tvec)

if axi == 1

axs[1].scatter(indx,mod.(uout[ish], 2*pi),color=COLORS[1])
axs[1].set_xlim(indx[1],indx[end])
axs[1].set_xticks(indx[1]:30:indx[end])
axs[1].set_ylim(0,2pi)
axs[1].set_xlabel(L"i"; labelpad = -20)
axs[1].set_ylabel(L"\phi"; labelpad=-10)
end

# plot evolution of phase velocity
# compute phase velocities (could perhaps be replaced by call of ODE function)
qmat = zeros(nosc,nts)
for its = 1:nts
    u = mod.(uout[its], 2*pi)
    for i = 1:nosc
        term = 0.0
        for j = i-R:i+R
            jj = 1 + mod(j-1,nosc)
            term = term + sin( u[i] - u[jj] + alpha )
        end
        qmat[i,its] = ω - term / (2*R )
    end
end

im1 = axs[1+axi].pcolormesh(tvec,indx,qmat,cmap = "inferno")
axs[1+axi].set_xlabel(L"t"; labelpad = -20)
axs[axi+1].set_ylabel(L"i"; labelpad = -20)
axs[axi+1].set_yticks([indx[1], indx[end]])
axs[axi+1].set_xticks([0, tvec[end]])

if axi == 2
    cbar = plt.colorbar(im1)  # , ticks=[-0.2, 0.2, 0.6, 1.], orientation="horizontal")
    # cbar.set_clim(1., 4.)
    cbar.set_label(L"d \phi /  dt")
end

end

# %%
add_identifiers!(fig, axs; xloc = 0.05)

fig.tight_layout(pad=0.3)

wsave(plotsdir("10", "chaimera"), fig)