# This file is part of the TaylorIntegration.jl package; MIT licensed

using TaylorIntegration, DiffEqBase
using Test
using LinearAlgebra: norm
using Elliptic

@testset "Testing `taylorize.jl`" begin

    # Constants for the integrations
    local _order = 20
    local _abstol = 1.0E-20
    local t0 = 0.0
    local tf = 1000.0

    # Scalar integration
    @testset "Scalar case: xdot(x, p, t) = b-x^2" begin
        b1 = 3.0
        @taylorize xdot1(x, p, t) = b1-x^2
        @test length(methods(TaylorIntegration.jetcoeffs!)) == 3
        @test (@isdefined xdot1)

        x0 = 1.0
        tv1, xv1 = taylorinteg( xdot1, x0, t0, tf, _order, _abstol, maxsteps=1000,
            parse_eqs=false)
        tv1p, xv1p = taylorinteg( xdot1, x0, t0, tf, _order, _abstol, maxsteps=1000)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Now using `local` constants
        @taylorize xdot2(x, p, t) = (local b2 = 3; b2-x^2)
        @test length(methods(TaylorIntegration.jetcoeffs!)) == 4
        @test (@isdefined xdot2)

        tv2, xv2 = taylorinteg( xdot2, x0, t0, tf, _order, _abstol, maxsteps=1000,
            parse_eqs=false)
        tv2p, xv2p = taylorinteg( xdot2, x0, t0, tf, _order, _abstol, maxsteps=1000)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        # Passing a parameter
        @taylorize xdot3(x, p, t) = p-x^2
        @test length(methods(TaylorIntegration.jetcoeffs!)) == 5
        @test (@isdefined xdot3)

        tv3, xv3 = taylorinteg( xdot3, x0, t0, tf, _order, _abstol, b1, maxsteps=1000,
            parse_eqs=false)
        tv3p, xv3p = taylorinteg( xdot3, x0, t0, tf, _order, _abstol, b1, maxsteps=1000)

        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        # Comparing integrations
        @test length(tv1p) == length(tv2p) == length(tv3p)
        @test iszero( norm(tv1p-tv2p, Inf) )
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv2p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        xv4 = taylorinteg( xdot2, x0, t0:0.5:tf, _order, _abstol, maxsteps=1000,
            parse_eqs=false)
        xv4p = taylorinteg( xdot2, x0, t0:0.5:tf, _order, _abstol, maxsteps=1000)
        @test iszero( norm(xv4-xv4p, Inf) )

        # Compare to exact solution
        exact_sol(t, b, x0) = sqrt(b)*((sqrt(b)+x0)-(sqrt(b)-x0)*exp(-2sqrt(b)*t)) /
            ((sqrt(b)+x0)+(sqrt(b)-x0)*exp(-2sqrt(b)*t))
        @test norm(xv1p[end] - exact_sol(tv1p[end], b1, x0), Inf) < 1.0e-15
    end

    @testset "Scalar case: xdot(x, p, t) = -10" begin
        xdot1(x, p, t) = -10 + zero(t) # `zero(t)` is needed; cf #20
        @taylorize xdot1_parsed(x, p, t) = -10 + zero(t) # `zero(t)` can be avoided here

        @test (@isdefined xdot1_parsed)
        tv1, xv1   = taylorinteg( xdot1, 10, 1, 20.0, _order, _abstol)
        tv1p, xv1p = taylorinteg( xdot1_parsed, 10.0, 1.0, 20.0, _order, _abstol)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Now using `local`
        @taylorize function xdot2(x, p, t)
            local ggrav = -10 + zero(t)  # zero(t) is needed for the parse_eqs=false case
            tmp = ggrav  # needed to avoid an error when parsing
        end

        @test (@isdefined xdot2)

        tv2, xv2   = taylorinteg( xdot2, 10.0, 1.0, 20.0, _order, _abstol, parse_eqs=false)
        tv2p, xv2p = taylorinteg( xdot2, 10.0, 1, 20.0, _order, _abstol)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        # Passing a parameter
        @taylorize xdot3(x, p, t) = p + zero(t) # `zero(t)` can be avoided here

        @test (@isdefined xdot3)

        tv3, xv3   = taylorinteg( xdot3, 10.0, 1.0, 20.0, _order, _abstol, -10,
            parse_eqs=false)
        tv3p, xv3p = taylorinteg( xdot3, 10, 1.0, 20.0, _order, _abstol, -10)

        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        # Comparing integrations
        @test length(tv1p) == length(tv2p) == length(tv3p)
        @test iszero( norm(tv1p-tv2p, Inf) )
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv2p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        # Compare to exact solution
        exact_sol(t, g, x0) = x0 + g*(t-1.0)
        @test norm(xv1p[end] - exact_sol(20, -10, 10), Inf) < 10eps(exact_sol(20, -10, 10))
    end

    # Pendulum integrtf = 100.0
    @testset "Integration of the pendulum and DiffEqs interface" begin
        @taylorize function pendulum!(dx, x, p, t)
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
            nothing
        end

        @test (@isdefined pendulum!)
        q0 = [pi-0.001, 0.0]
        tv2, xv2 = taylorinteg(pendulum!, q0, t0, tf, _order, _abstol, parse_eqs=false,
            maxsteps=5000)
        tv2p, xv2p = taylorinteg(pendulum!, q0, t0, tf, _order, _abstol, maxsteps=5000)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        prob = ODEProblem(pendulum!, q0, (t0, tf), nothing) # no parameters
        sol1 = solve(prob, TaylorMethod(_order), abstol=_abstol, parse_eqs=true)
        sol2 = solve(prob, TaylorMethod(_order), abstol=_abstol, parse_eqs=false)
        sol3 = solve(prob, TaylorMethod(_order), abstol=_abstol)

        @test sol1.t == sol2.t == sol3.t == tv2p
        @test sol1.u[end] == sol2.u[end] == sol3.u[end] == xv2p[end,1:2]
    end

    # Complex dependent variables
    @testset "Complex dependent variable" begin
        cc = complex(0.0,1.0)
        @taylorize eqscmplx(x, p, t) = cc*x

        @test (@isdefined eqscmplx)
        cx0 = complex(1.0, 0.0)
        tv1, xv1 = taylorinteg(eqscmplx, cx0, t0, tf, _order, _abstol, maxsteps=1500,
            parse_eqs=false)
        tv1p, xv1p = taylorinteg(eqscmplx, cx0, t0, tf, _order, _abstol, maxsteps=1500)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Using `local` for the constant value
        @taylorize eqscmplx2(x, p, t) = (local cc1 = Complex(0.0,1.0); cc1*x)

        @test (@isdefined eqscmplx2)
        tv2, xv2 = taylorinteg(eqscmplx2, cx0, t0, tf, _order, _abstol, maxsteps=1500,
            parse_eqs=false)
        tv2p, xv2p = taylorinteg(eqscmplx2, cx0, t0, tf, _order, _abstol, maxsteps=1500)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        # Passing a parameter
        @taylorize eqscmplx3(x, p, t) = p*x
        @test (@isdefined eqscmplx3)
        tv3, xv3 = taylorinteg(eqscmplx3, cx0, t0, tf, _order, _abstol, cc, maxsteps=1500,
            parse_eqs=false)
        tv3p, xv3p = taylorinteg(eqscmplx3, cx0, t0, tf, _order, _abstol, cc, maxsteps=1500)

        @test length(tv3) == length(tv3p)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )

        # Comparing integrations
        @test length(tv1p) == length(tv2p) == length(tv3p)
        @test iszero( norm(tv1p-tv2p, Inf) )
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv2p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )

        # Compare to exact solution
        exact_sol(t, p, x0) = x0*exp(p*t)
        nn = norm.(abs2.(xv1p)-abs2.(exact_sol.(tv1p, cc, cx0)), Inf)
        @test maximum(nn) < 1.0e-14
    end

    @testset "Time-dependent integration (with and without `local` vars)" begin
        @taylorize function integ_cos1(x, p, t)
            y = cos(t)
            return y
        end
        @taylorize function integ_cos2(x, p, t)
            local y = cos(t)  # allows to calculate directly `cos(t)` *once*
            yy = y   # needed to avoid an error
            return yy
        end

        @test (@isdefined integ_cos1)
        @test (@isdefined integ_cos2)

        tv11, xv11 = taylorinteg(integ_cos1, 0.0, 0.0, pi, _order, _abstol, parse_eqs=false)
        tv12, xv12 = taylorinteg(integ_cos1, 0.0, 0.0, pi, _order, _abstol)
        @test length(tv11) == length(tv12)
        @test iszero( norm(tv11-tv12, Inf) )
        @test iszero( norm(xv11-xv12, Inf) )

        tv21, xv21 = taylorinteg(integ_cos2, 0.0, 0.0, pi, _order, _abstol, parse_eqs=false)
        tv22, xv22 = taylorinteg(integ_cos2, 0.0, 0.0, pi, _order, _abstol)
        @test length(tv21) == length(tv22)
        @test iszero( norm(tv21-tv22, Inf) )
        @test iszero( norm(xv21-xv22, Inf) )

        @test iszero( norm(tv12-tv22, Inf) )
        @test iszero( norm(xv12-xv22, Inf) )

        # Compare to exact solution
        @test norm(xv11[end] - sin(tv11[end]), Inf) < 1.0e-15
    end

    @testset "Time-dependent integration (vectorial case)" begin
        @taylorize function integ_vec(dx, x, p, t)
            dx[1] = cos(t)
            dx[2] = -sin(t)
            return dx
        end

        @test (@isdefined integ_vec)
        x0 = [0.0, 1.0]
        tv11, xv11 = taylorinteg(integ_vec, x0, 0.0, pi, _order, _abstol, parse_eqs=false)
        tv12, xv12 = taylorinteg(integ_vec, x0, 0.0, pi, _order, _abstol)
        @test length(tv11) == length(tv12)
        @test iszero( norm(tv11-tv12, Inf) )
        @test iszero( norm(xv11-xv12, Inf) )

        # Compare to exact solution
        @test norm(xv12[end, 1] - sin(tv12[end]), Inf) < 1.0e-15
        @test norm(xv12[end, 2] - cos(tv12[end]), Inf) < 1.0e-15
    end

    # Simple harmonic oscillator
    @testset "Simple harmonic oscillator" begin
        @taylorize function harm_osc!(dx, x, p, t)
            local ω = p[1]
            dx[1] = x[2]
            dx[2] = -ω^2 * x[1]
            return nothing
        end

        @test (@isdefined harm_osc!)
        q0 = [1.0, 0.0]
        p = [2.0]
        tv1, xv1 = taylorinteg(harm_osc!, q0, t0, tf, _order, _abstol,
            p, maxsteps=1000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(harm_osc!, q0, t0, tf, _order, _abstol,
            p, maxsteps=1000)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        # Comparing to exact solution
        @test norm(xv1p[end, 1] - cos(p[1]*tv1p[end]), Inf) < 1.0e-12
        @test norm(xv1p[end, 2] + p[1]*sin(p[1]*tv1p[end]), Inf) < 3.0e-12
    end

    @testset "Multiple pendula" begin
        NN = 3
        nnrange = 1:3
        @taylorize function multpendula1!(dx, x, p, t)
            for i in p[2]
                dx[i] = x[p[1]+i]
                dx[i+p[1]] = -sin( x[i] )
            end
            return nothing
        end

        @test (@isdefined multpendula1!)
        q0 = [pi-0.001, 0.0, pi-0.001, 0.0,  pi-0.001, 0.0]
        tv1, xv1 = taylorinteg(multpendula1!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(multpendula1!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000)

        @test length(tv1) == length(tv1p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )

        @taylorize function multpendula2!(dx, x, p, t)
            local NN = 3
            local nnrange = 1:NN
            for i in nnrange
                dx[i] = x[NN+i]
                dx[i+NN] = -sin( x[i] )
            end
            return nothing
        end

        @test (@isdefined multpendula2!)
        tv2, xv2 = taylorinteg(multpendula2!, q0, t0, tf, _order, _abstol,
            maxsteps=1000, parse_eqs=false)
        tv2p, xv2p = taylorinteg(multpendula2!, q0, t0, tf, _order, _abstol,
            maxsteps=1000)

        @test length(tv2) == length(tv2p)
        @test iszero( norm(tv2-tv2p, Inf) )
        @test iszero( norm(xv2-xv2p, Inf) )

        @taylorize function multpendula3!(dx, x, p, t)
            nn, nnrng = p    # `local` is not needed to reassign `p`; internally is treated as local)
            for i in nnrng
                dx[i] = x[nn+i]
                dx[i+nn] = -sin( x[i] )
            end
            return nothing
        end

        @test (@isdefined multpendula3!)
        tv3, xv3 = taylorinteg(multpendula3!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000, parse_eqs=false)
        tv3p, xv3p = taylorinteg(multpendula3!, q0, t0, tf, _order, _abstol,
            [NN, nnrange], maxsteps=1000)

        # Comparing integrations
        @test length(tv1) == length(tv2) == length(tv3)
        @test iszero( norm(tv1-tv2, Inf) )
        @test iszero( norm(tv1-tv3, Inf) )
        @test iszero( norm(xv1-xv2, Inf) )
        @test iszero( norm(xv1-xv3, Inf) )
    end

    # Kepler problem
    # Redefining some constants
    local _order = 28
    local tf = 2π*100.0
    @testset "Kepler problem (using `^`)" begin
        @taylorize function kepler1!(dq, q, p, t)
            μ = p
            r_p3d2 = (q[1]^2+q[2]^2)^1.5

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = μ * q[1]/r_p3d2
            dq[4] = μ * q[2]/r_p3d2

            return nothing
        end

        @taylorize function kepler2!(dq, q, p, t)
            nn, μ = p
            r2 = zero(q[1])
            for i = 1:nn
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r_p3d2 = r2^(3/2)
            for j = 1:nn
                dq[j] = q[nn+j]
                dq[nn+j] = μ*q[j]/r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler1!)
        @test (@isdefined kepler2!)
        pars = (2, -1.0)
        q0 = [0.2, 0.0, 0.0, 3.0]
        tv1, xv1 = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000)

        tv6p, xv6p = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, pars,
            maxsteps=500000, parse_eqs=false)
        tv7p, xv7p = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, pars,
            maxsteps=500000)

        @test length(tv1) == length(tv1p) == length(tv6p)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )
        @test iszero( norm(tv1-tv6p, Inf) )
        @test iszero( norm(xv1-xv6p, Inf) )
        @test iszero( norm(tv7p-tv6p, Inf) )
        @test iszero( norm(xv7p-xv6p, Inf) )

        @taylorize function kepler3!(dq, q, p, t)
            local mμ = -1.0
            r_p3d2 = (q[1]^2+q[2]^2)^1.5

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = mμ*q[1]/r_p3d2
            dq[4] = mμ*q[2]/r_p3d2

            return nothing
        end

        @taylorize function kepler4!(dq, q, p, t)
            local mμ = -1.0
            local NN = 2
            r2 = zero(q[1])
            for i = 1:NN
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r_p3d2 = r2^(3/2)
            for j = 1:NN
                dq[j] = q[NN+j]
                dq[NN+j] = mμ*q[j]/r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler3!)
        @test (@isdefined kepler4!)
        tv3, xv3 = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv3p, xv3p = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        tv4, xv4 = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv4p, xv4p = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        @test length(tv3) == length(tv3p) == length(tv4)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )
        @test iszero( norm(tv3-tv4, Inf) )
        @test iszero( norm(xv3-xv4, Inf) )
        @test iszero( norm(tv4p-tv4, Inf) )
        @test iszero( norm(xv4p-xv4, Inf) )

        # Comparing both integrations
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )
    end

    @testset "Kepler problem (using `sqrt`)" begin
        @taylorize function kepler1!(dq, q, p, t)
            local μ = -1.0
            r = sqrt(q[1]^2+q[2]^2)
            r_p3d2 = r^3

            dq[1] = q[3]
            dq[2] = q[4]
            dq[3] = μ * q[1] / r_p3d2
            dq[4] = μ * q[2] / r_p3d2

            return nothing
        end

        @taylorize function kepler2!(dq, q, p, t)
            local NN = 2
            μ = p
            r2 = zero(q[1])
            for i = 1:NN
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r = sqrt(r2)
            r_p3d2 = r^3
            for j = 1:NN
                dq[j] = q[NN+j]
                dq[NN+j] =  μ * q[j] / r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler1!)
        @test (@isdefined kepler2!)
        q0 = [0.2, 0.0, 0.0, 3.0]
        tv1, xv1 = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv1p, xv1p = taylorinteg(kepler1!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        tv2, xv2 = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000, parse_eqs=false)
        tv2p, xv2p = taylorinteg(kepler2!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000)

        @test length(tv1) == length(tv1p) == length(tv2)
        @test iszero( norm(tv1-tv1p, Inf) )
        @test iszero( norm(xv1-xv1p, Inf) )
        @test iszero( norm(tv1-tv2, Inf) )
        @test iszero( norm(xv1-xv2, Inf) )
        @test iszero( norm(tv2p-tv2, Inf) )
        @test iszero( norm(xv2p-xv2, Inf) )

        @taylorize function kepler3!(dq, q, p, t)
            μ = p
            x, y, px, py = q
            r = sqrt(x^2+y^2)
            r_p3d2 = r^3

            dq[1] = px
            dq[2] = py
            dq[3] = μ * x / r_p3d2
            dq[4] = μ * y / r_p3d2

            return nothing
        end

        @taylorize function kepler4!(dq, q, p, t)
            local NN = 2
            local μ = -1.0
            r2 = zero(q[1])
            for i = 1:NN
                r2_aux = r2 + q[i]^2
                r2 = r2_aux
            end
            r = sqrt(r2)
            r_p3d2 = r^3
            for j = 1:NN
                dq[j] = q[NN+j]
                dq[NN+j] = μ * q[j] / r_p3d2
            end

            nothing
        end

        @test (@isdefined kepler3!)
        @test (@isdefined kepler4!)
        tv3, xv3 = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000, parse_eqs=false)
        tv3p, xv3p = taylorinteg(kepler3!, q0, t0, tf, _order, _abstol, -1.0,
            maxsteps=500000)

        tv4, xv4 = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000, parse_eqs=false)
        tv4p, xv4p = taylorinteg(kepler4!, q0, t0, tf, _order, _abstol,
            maxsteps=500000)

        @test length(tv3) == length(tv3p) == length(tv4)
        @test iszero( norm(tv3-tv3p, Inf) )
        @test iszero( norm(xv3-xv3p, Inf) )
        @test iszero( norm(tv3-tv4, Inf) )
        @test iszero( norm(xv3-xv4, Inf) )
        @test iszero( norm(tv4p-tv4, Inf) )
        @test iszero( norm(xv4p-xv4, Inf) )

        # Comparing both integrations
        @test iszero( norm(tv1p-tv3p, Inf) )
        @test iszero( norm(xv1p-xv3p, Inf) )
    end

    local tf = 20.0
    @testset "Lyapunov spectrum and `@taylorize`" begin
        #Lorenz system parameters
        params = [16.0, 45.92, 4.0]

        #Lorenz system ODE:
        @taylorize function lorenz1!(dq, q, p, t)
            x, y, z = q
            σ, ρ, β = p
            dq[1] = σ*(y-x)
            dq[2] = x*(ρ-z) - y
            dq[3] = x*y - β*z
            nothing
        end

        #Lorenz system Jacobian (in-place):
        function lorenz1_jac!(jac, q, p, t)
            x, y, z = q
            σ, ρ, β = p
            jac[1,1] = -σ+zero(x)
            jac[2,1] = ρ-z
            jac[3,1] = y
            jac[1,2] = σ+zero(x)
            jac[2,2] = -one(y)
            jac[3,2] = x
            jac[1,3] = zero(z)
            jac[2,3] = -x
            jac[3,3] = -β+zero(x)
            nothing
        end

        q0 = [19.0, 20.0, 50.0] #the initial condition
        xi = set_variables("δ", order=1, numvars=length(q0))

        tv1, lv1, xv1 = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, maxsteps=2000, parse_eqs=false);

        tv1p, lv1p, xv1p = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, maxsteps=2000);

        @test tv1 == tv1p
        @test lv1 == lv1p
        @test xv1 == xv1p

        tv2, lv2, xv2 = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, lorenz1_jac!, maxsteps=2000, parse_eqs=false);

        tv2p, lv2p, xv2p = lyap_taylorinteg(lorenz1!, q0, t0, tf, _order, _abstol,
            params, lorenz1_jac!, maxsteps=2000,  parse_eqs=true);

        @test tv2 == tv2p
        @test lv2 == lv2p
        @test xv2 == xv2p

        # Comparing both integrations (lorenz1)
        @test tv1 == tv2
        @test lv1 == lv2
        @test xv1 == xv2

        #Lorenz system ODE:
        @taylorize function lorenz2!(dx, x, p, t)
            #Lorenz system parameters
            local σ = 16.0
            local β = 4.0
            local ρ = 45.92

            dx[1] = σ*(x[2]-x[1])
            dx[2] = x[1]*(ρ-x[3])-x[2]
            dx[3] = x[1]*x[2]-β*x[3]
            nothing
        end

        #Lorenz system Jacobian (in-place):
        function lorenz2_jac!(jac, x, p, t)
            #Lorenz system parameters
            local σ = 16.0
            local β = 4.0
            local ρ = 45.92

            jac[1,1] = -σ+zero(x[1])
            jac[2,1] = ρ-x[3]
            jac[3,1] = x[2]
            jac[1,2] = σ+zero(x[1])
            jac[2,2] = -1.0+zero(x[1])
            jac[3,2] = x[1]
            jac[1,3] = zero(x[1])
            jac[2,3] = -x[1]
            jac[3,3] = -β+zero(x[1])
            nothing
        end

        tv3, lv3, xv3 = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            maxsteps=2000, parse_eqs=false);

        tv3p, lv3p, xv3p = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            maxsteps=2000);

        @test tv3 == tv3p
        @test lv3 == lv3p
        @test xv3 == xv3p

        tv4, lv4, xv4 = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000, parse_eqs=false);

        tv4p, lv4p, xv4p = lyap_taylorinteg(lorenz2!, q0, t0, tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000,  parse_eqs=true);

        @test tv4 == tv4p
        @test lv4 == lv4p
        @test xv4 == xv4p

        # Comparing both integrations (lorenz2)
        @test tv3 == tv4
        @test lv3 == lv4
        @test xv3 == xv4

        # Comparing both integrations (lorenz1! vs lorenz2!)
        @test tv1 == tv3
        @test lv1 == lv3
        @test xv1 == xv3

        # Using ranges
        lv5, xv5 = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            maxsteps=2000, parse_eqs=false);

        lv5p, xv5p = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            maxsteps=2000, parse_eqs=true);

        @test lv5 == lv5p
        @test xv5 == xv5p

        lv6, xv6 = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000, parse_eqs=false);

        lv6p, xv6p = lyap_taylorinteg(lorenz2!, q0, t0:0.125:tf, _order, _abstol,
            lorenz2_jac!, maxsteps=2000,  parse_eqs=true);

        @test lv6 == lv6p
        @test xv6 == xv6p
    end

    @testset "Tests for throwing errors" begin
        # Wrong number of arguments
        ex = :(function f_p!(dx, x, p, t, y)
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # `&&` is not yet implemented
        ex = :(function f_p!(x, p, t)
            true && x
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # a is not an Expr; String
        ex = :(function f_p!(x, p, t)
            "a"
        end)
        @test_throws ArgumentError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # KeyError: key :fname not found
        ex = :(begin
            x=1
            x+x
        end)
        @test_throws KeyError TaylorIntegration._make_parsed_jetcoeffs(ex)

        # BoundsError; no return variable defined
        ex = :(function f_p!(x, p, t)
            local cos(t)
        end)
        @test_throws BoundsError TaylorIntegration._make_parsed_jetcoeffs(ex)
    end

    @testset "Jet transport with @taylorize macro" begin
        @taylorize function pendulum!(dx, x, p, t)
            dx[1] = x[2]
            dx[2] = -sin( x[1] )
            nothing
        end

        varorder = 2 #the order of the variational expansion
        p = set_variables("ξ", numvars=2, order=varorder) #TaylorN steup
        q0 = [1.3, 0.0] #the initial conditions
        q0TN = q0 + p #parametrization of a small neighbourhood around the initial conditions
        # T is the librational period == 4Elliptic.K(sin(q0[1]/2)^2)
        T = 4Elliptic.K(sin(q0[1]/2)^2) # equals 7.019250311844546
        integstep = 0.25*T #the time interval between successive evaluations of the solution vector

        #the time range
        tr = t0:integstep:T;
        #note that as called below, taylorinteg uses the parsed jetcoeffs! method by default
        xvp = taylorinteg(pendulum!, q0, tr, _order, _abstol, maxsteps=100)

        # "warmup" for jet transport integration
        xvTN = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=1, parse_eqs=false)
        xvTN = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=1)
        @test size(xvTN) == (5,2)
        #jet transport integration with parsed jetcoeffs!
        xvTNp = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=100)
        #jet transport integration with non-parsed jetcoeffs!
        xvTN = taylorinteg(pendulum!, q0TN, tr, _order, _abstol, maxsteps=100, parse_eqs=false)
        @test xvTN == xvTNp
        @test norm(xvTNp[:,:]() - xvp, Inf) < 1E-15

        dq = 0.0001rand(2)
        q1 = q0 + dq
        yv = taylorinteg(pendulum!, q1, tr, _order, _abstol, maxsteps=100)
        yv_jt = xvTNp[:,:](dq)
        @test norm(yv-yv_jt, Inf) < 1E-11

        dq = 0.001
        t = Taylor1([0.0, 1.0], 10)
        x0T1 = q0+[0t,t]
        q1 = q0+[0.0,dq]
        tv, xv = taylorinteg(pendulum!, q1, t0, 2T, _order, _abstol)
        tvT1, xvT1 = taylorinteg(pendulum!, x0T1, t0, 2T, _order, _abstol, parse_eqs=false)
        tvT1p, xvT1p = taylorinteg(pendulum!, x0T1, t0, 2T, _order, _abstol)
        @test tvT1 == tvT1p
        @test xvT1 == xvT1p
        xv_jt = xvT1p[:,:](dq)
        @test norm(xv_jt[end,:]-xv[end,:]) < 20eps(norm(xv[end,:]))
    end
end


# ---
# title: 473. Matchsticks to Square
# id: problem473
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Depth-first Search
# link: <https://leetcode.com/problems/matchsticks-to-square/description/>
# hidden: true
# ---
#
# Remember the story of Little Match Girl? By now, you know exactly what
# matchsticks the little match girl has, please find out a way you can make one
# square by using up all those matchsticks. You should not break any stick, but
# you can link them up, and each matchstick must be used **exactly** one time.
#
# Your input will be several matchsticks the girl has, represented with their
# stick length. Your output will either be true or false, to represent whether
# you could make one square using all the matchsticks the little match girl has.
#
# **Example 1:**
#
#
#
#     Input: [1,1,2,2,2]
#     Output: true
#
#     Explanation: You can form a square with length 2, one side of the square came two sticks with length 1.
#
#
# **Example 2:**
#
#
#
#     Input: [3,3,3,3,4]
#     Output: false
#
#     Explanation: You cannot find a way to form a square with all the matchsticks.
#
#
# **Note:**
#
#   1. The length sum of the given matchsticks is in the range of `0` to `10^9`.
#   2. The length of the given matchstick array will not exceed `15`.
#
#
## @lc code=start
using LeetCode

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


#=
  From a BLOG model in
  Yi Wu, Lei Li, Stuart Russell, Rastislav Bodik
  "Swift: Compiled Inference for Probabilistic Programming Languages"
  Page 2
  """
  [The model] demonstrates the open-universe urn-ball model. In this version
  the query asks for the color of the next random pick from an urn given the
  colors of balls drawn previously.
  """

  Summary Statistics
  parameters      mean       std   naive_se      mcse        ess       rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64   Float64?   Float64?      Float64?

    numBalls    2.6863    3.6429     0.0364    0.2942   107.7806     1.0000        4.7524
    ...
     drawnD5    1.7703    1.8844     0.0188    0.1436   120.3713     1.0076        5.3076
     colorD5    1.8228    0.3819     0.0038    0.0296   107.6107     0.9999        4.7450

   Distributions of variable numBalls (num:0)
   1.00000 =>    5969  (0.596900)
   2.00000 =>    1409  (0.140900)
   3.00000 =>     843  (0.084300)
   4.00000 =>     489  (0.048900)
   9.00000 =>     319  (0.031900)
   5.00000 =>     242  (0.024200)
   19.00000 =>     209  (0.020900)
   7.00000 =>     161  (0.016100)
   6.00000 =>     109  (0.010900)
   15.00000 =>      72  (0.007200)
   16.00000 =>      54  (0.005400)
   20.00000 =>      46  (0.004600)
   11.00000 =>      42  (0.004200)
   12.00000 =>      36  (0.003600)

   Distributions of variable drawnD5 (num:0)
   1.00000 =>    7070  (0.707000)
   2.00000 =>    1149  (0.114900)
   3.00000 =>     905  (0.090500)
   4.00000 =>     449  (0.044900)
   10.00000 =>     109  (0.010900)
   9.00000 =>      72  (0.007200)
   5.00000 =>      68  (0.006800)
   8.00000 =>      54  (0.005400)
   6.00000 =>      42  (0.004200)
   19.00000 =>      38  (0.003800)
   7.00000 =>      36  (0.003600)
   11.00000 =>       8  (0.000800)

   Distributions of variable colorD5
   blue       =>    8228  (0.822800)
   green      =>    1772  (0.177200)

   I.e. it's quite probable that there is just one ball (or perhaps two).

  Cf ~/blog/urn_ball.blog
     ~/psi/urn_ball.psi
     ~/webppl/urn_ball.wppl

=#
using Turing, StatsPlots, Distributions
include("jl_utils.jl")

# Blue: 1 Green: 2
# We observe 4 Green balls
@model function urn_ball(obs=[2,2,2,2])
    n = 5

    blue = 1
    green = 2

    numBalls ~ DiscreteUniform(1,20)

    # PG(5) throws
    # ERROR: LoadError: CTaskException:
    # DimensionMismatch("tried to assign 7 elements to 16 destinations")
    # for this filldist() expression
    # color ~ filldist(Categorical([0.9,0.1]),numBalls) # [Blue,Green]

    # PG(5) don't have any problems with this, though
    color = tzeros(numBalls)
    for i in 1:numBalls
        color[i] ~ Categorical([0.9,0.1]) # [Blue,Green]
    end

    drawn ~ filldist(DiscreteUniform(1,numBalls),n)

    for i in 1:length(obs)
        obs[i] ~ Dirac(color[drawn[i]])
    end

    drawnD5 ~ Dirac(drawn[5])
    colorD5 ~ Dirac(color[drawn[5]])

end

obs=[2,2,2,2] # Observed only green balls
model = urn_ball(obs)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:numBalls)
show_var_dist_pct(chns,:drawnD5)
show_var_dist_pct(chns,:colorD5,["green","blue"])


#***********************************************************************
#  Copyright 2006 John A. Trangenstein
#
#  This software is made available for research and instructional use
#  only.
#  You may copy and use this software without charge for these
#  non-commercial purposes, provided that the copyright notice and
#  associated text is reproduced on all copies.
#  For all other uses (including distribution of modified versions),
#  please contact the author at
#    John A. Trangenstein
#    Department of Mathematics
#    Duke University
#    Durham, NC 27708-0320
#    USA
#  or
#    johnt@math.duke.edu
#
#  This software is made available "as is" without any assurance that it
#  is completely correct, or that it will work for your purposes.
#  Use the software at your own risk.
#***********************************************************************

using OffsetArrays

function initsl(ncells::Int, fc::Int, lc::Int, fx::Int, lx::Int, ifirst::Int, ilast::Int,
                conservd::OffsetArray{Float64,1},
                x::OffsetArray{Float64,1})
#   ******************************************************************
#   interior values only; others defined by boundary conditions
#   ******************************************************************
    dx=(x_right-x_left)/ncells
    @inbounds for ie=ifirst:ilast+1
        x[ie]=x_left+ie*dx
    end
    ijump=max(ifirst-1,min(convert(Int,round(ncells*(jump-x_left)/(x_right-x_left))),ilast+1))
    @inbounds for ic=ifirst:ijump-1
        conservd[ic]=statelft
    end
    frac=(jump-x_left-ijump*dx)/(x_right-x_left)
    conservd[ijump]=statelft*frac+statergt*(1.0-frac)
    @inbounds for ic=ijump+1:ilast
        conservd[ic]=statergt
    end
end

function bcmesh(fim::Int,lam::Int,ncells::Int, x::OffsetArray{Float64,1})
#   ******************************************************************
#   right side
#   ******************************************************************
    if lam >= ncells
        it=ncells
        dx=x(it)-x(it-1)
        @inbounds for ie=it+1:lam+1
            x[ie]=x[it]+dx*(ie-it)
        end
    end
#   ******************************************************************
#   left side
#   ******************************************************************
    if fim < 0
        dx=x(1)-x(0)
        @inbounds for ie=fim:-1
            x[ie]=x[0]+dx*(ie)
        end
    end
end

function bccells(fic::Int,lac::Int,ncells::Int, conservd::OffsetArray{Float64,1})
#   ******************************************************************
#   right side
#   ******************************************************************
    if lac >= ncells
        # println("treating right side")
        @inbounds for ic=ncells:lac
            # outgoing wave:
            conservd[ic]=conservd[ncells-1]
        end
    end
#     ******************************************************************
#     left side
#     ******************************************************************
    if fic < 0
        # println("treating left side")
        @inbounds for ic=fic:-1
            # outgoing wave:
            # conservd[ic]=conservd[0]

            # specified value
            conservd[ic]=statelft
        end
    end
end


using ROC
using Statistics

auc_(scores,labels) = AUC(roc(scores, labels))

function R2_(scores, labels)
    dY = scores -labels
    Ysc = labels - mean(labels)*ones(size(labels))
    return 1 - (sum(dY.*dY) / sum(Ysc.*Ysc))
end

"""
#R2_(scores, labels) = 1 - sum((scores-labels).*scores-labels)

function explained_variance_score(y_true, y_pred)
    return(sum((y_pred .- mean(y_true)) .^ 2))
end

function total_variance_score(y_true, y_pred)
    return(sum((y_true .- mean(y_true)) .^ 2))
end

function r2_score(y_true, y_pred)
    #println(y_true)
    #println(y_pred)
    println((explained_variance_score(y_true, y_pred), total_variance_score(y_true, y_pred)))
    return(explained_variance_score(y_true, y_pred) / total_variance_score(y_true, y_pred))
end
"""


using Distributions, Statistics, Plots; pyplot()

dist1 = TriangularDist(0,1,1)
dist2 = TriangularDist(0,1,0)
N=10^6

data1, data2 = rand(dist1,N), rand(dist2,N)
dataSum = data1 + data2

mgf(s) = 4(1+(s-1)*MathConstants.e^s)*(MathConstants.e^s-1-s)/s^4

mgfPointEst(s) = mean([MathConstants.e^(s*z) for z in
					rand(dist1,20) + rand(dist2,20)])

p1 = histogram(dataSum, bins=80, normed=:true,
	ylims=(0,1.4), xlabel="z", ylabel="PDF")

sGrid = -1:0.01:1
p2 = plot(sGrid, mgfPointEst.(sGrid), c=:blue, ylims=(0,3.5))
p2 = plot!(sGrid, mgf.(sGrid), c=:red)
p2 = plot!( [minimum(sGrid),maximum(sGrid)],
	[minimum(sGrid),maximum(sGrid)].+1,
	c=:black, xlabel="s", ylabel="MGF")

plot(p1, p2, legend=:none, size=(800, 400))


using Plots
using ImplicitEquations
using LaTeXStrings

f(x, y) = x^4 - x^2 * y + y^3
g(u, w) = f(u, u * w)

p1 = plot(Eq(f, 0), title = L"x^4 - x^2y + y^3 = 0")
p2 = plot(Eq(g, 0), title = L"u^3(u - w + w^3) = 0")
plot(
    p1,
    p2,
    xlabel = "x",
    ylabel = "y",
    xlim = [-1, 1],
    ylim = [-1, 1],
    aspect_ratio = :equal
)


using LinearAlgebra
using Einsum
using ASE
using JuLIP
using JuLIP.Potentials

using JuLIPMaterials.CLE

println("=========================")
println(" Testing Elasticity      ")
println("=========================")

println("Testing Voigt conversions")
σ, ϵ = rand(6), rand(6)
σₘ = cauchy_stress(σ)
ϵₘ = cauchy_strain(ϵ)
@test voigt_stress(σₘ) ≈ σ
@test voigt_strain(ϵₘ) ≈ ϵ

# check invariance of scalar product σ⋅ϵ = sum_ij σ_ij⋅ϵ_ij =
@test  σ' * ϵ ≈ sum([σₘ[i,j] * ϵₘ[i,j] for i=1:3, j=1:3])

C11, C12, C44 = rand(3)
C = voigt_moduli(C11, C12, C44)
c = elastic_moduli(C)

# check invertability of c_ijkl <--> C_ij conversion
@test voigt_moduli(c) ≈ C

# check σ_ij = c_ijkl ϵ_kl <--> σ = C⋅ϵ
ϵ = rand(6)
ϵₘ = cauchy_strain(ϵ)
σ = zeros(3,3)
@einsum σ[i,j] = c[i,j,k,l] * ϵₘ[k,l]
@test cauchy_stress(C * ϵ) ≈ σ
@test voigt_stress(σ) ≈ C * ϵ

println("Testing `elastic_moduli`")

eam = JuLIP.Potentials.EAM(test_pots * "w_eam4.fs")
at = bulk(:W)
set_calculator!(at, eam)
variablecell!(at)
minimise!(at)
E0 = energy(at)
c = elastic_moduli(eam, at)
C = voigt_moduli(c)

# apply small random strain and check CLE prediction holds
ϵ = 1e-4*rand(6)
ϵₘ = cauchy_strain(ϵ)
apply_defm!(at, I(3) + ϵₘ)
σₘ = stress(at)
σ = voigt_stress(σₘ)

# check elastic energy matches potential energy
E = energy(at)
@test E - E0 ≈ 0.5 * ϵ' * σ * volume(at) atol=1e-7

# Hooke's law in Voigt form: σ = C ϵ
@test σₘ ≈ cauchy_stress(C * ϵ) atol=1e-5

# check Hooke's law in tensor form: σ_ij = c_ijkl ϵ_kl
σₜ = zeros(3,3)
@einsum σₜ[i,j] = c[i,j,k,l] * ϵₘ[k,l]
@test σₜ ≈ σₘ atol=1e-3

# check C has cubic symmetry and that C11, C12, C44 elastic constants match expected values for this potential
@show cubic_moduli(C, tol=1e-3)
@test all(isapprox.(cubic_moduli(C, tol=1e-3), [3.2644382942616716, 1.2618231295748539, 1.004134784461818], atol=1e-3))

println("Testing rotation of elastic moduli")

A = rotation_matrix(x=[1,1,1], y=[1,-1,0])
at = bulk(:W)
set_calculator!(at, eam)
c = elastic_moduli(at)
c_rot = rotate_moduli(c, A)
C = voigt_moduli(at)
C_rot = rotate_moduli(C, A)

# now we rotate the Atoms directly and compare
apply_defm!(at, A)
c_direct = elastic_moduli(at)
C_direct = voigt_moduli(at)

@test c_rot ≈ c_direct atol=1e-3
@test C_rot ≈ C_direct atol=1e-3


function init_pot_harm_3d( LF, ω, center )
    Npoints = LF.Nx * LF.Ny * LF.Nz
    Vpot = zeros( Float64,Npoints )
    for ip = 1:Npoints
        Vpot[ip] = 0.5 * ω^2 * norm(LF.lingrid[:,ip] - center[:])^2
    end
    return Vpot
end


# If want to use an empty function for normalizing/testing validity of parameters
noop(x::AbstractArray) = return true
noop(x::Real) = return true

# Copy population
import Base.copy
copy(x::abc_population_type) = abc_population_type(copy(x.theta),copy(x.weights),copy(x.dist), x.epsilon, copy(x.repeats))

# Compute weighted covariance of sample (from origin)
function cov_weighted(x::AbstractArray{Float64,2}, w::AbstractArray{Float64,1} )
  @assert(size(x,1)==length(w) )
  sumw = sum(w)
  @assert( sumw > 0. )
  if(sumw!= 1.0)
     w /= sum(w)
     sumw = 1.0
  end
  sumw2 = sum(w.*w)
  xbar = [ sum(x[:,i].*w) for i in 1:size(x,2) ]
  covar = zeros(size(x,2),size(x,2))
  for k in 1:size(x,2)
    for j in 1:size(x,2)
        for i in 1:size(x,1)
            @inbounds covar[j,k] += (x[i,j]-xbar[j])*(x[i,k]-xbar[k]) *w[i]
        end
    @inbounds covar[j,k] *= sumw/(sumw*sumw-sumw2)
    end
  end
  covar
end

# Compute weighted variances of sample (from origin)
#function var_weighted(x::AbstractArray{Float64,2}, w::AbstractArray{Float64,1} )
function var_weighted(x::Array{Float64,2}, w::Array{Float64,1} )
  @assert(size(x,1)==length(w) )
  sumw = sum(w)
  @assert( sumw > 0. )
  if(sumw!= 1.0)
     w /= sum(w)
     sumw = 1.0
  end
  sumw2 = sum(w.*w)
  xbar = [ sum(x[:,i].*w) for i in 1:size(x,2) ]
  covar = zeros(size(x,2))
    for j in 1:size(x,2)
        for i in 1:size(x,1)
            @inbounds covar[j] += (x[i,j]-xbar[j])*(x[i,j]-xbar[j]) *w[i]
        end
    @inbounds covar[j] *= sumw/(sumw*sumw-sumw2)
    end
  covar
end

function make_matrix_pd(A::AbstractArray{Float64,2}; epsabs::Float64 = 0.0, epsfac::Float64 = 1.0e-6)
  @assert(size(A,1)==size(A,2))
  println("make_matrix_pd input: ", A)
  flush(Base.stdout)
  B = 0.5*(A+A')
  itt = 1
  while !isposdef(B)
	eigvalB,eigvecB = eig(B)
        pos_eigval_idx = eigvalB.>0.0
	neweigval = (epsabs == 0.0) ? epsfac*minimum(abs.(eigvalB[pos_eigval_idx])) : epsabs
	eigvalB[.!pos_eigval_idx] = neweigval
	B = eigvecB *diagm(eigvalB)*eigvecB'
	println(itt,": ",B)
        flush(Base.stdout)
        #cholB = chol(B)
	itt +=1
	if itt>size(A,1)
	  error("There's a problem in make_matrix_pd.\n")
  	  break
	end
  end
  return B
end

# Compute Effective Sample Size for array of weights
function ess(w::AbstractArray{Float64,1})
  sumw = sum(w)
  sumw2 = sum(w.*w)
  return sumw*sumw/sumw2
end

# Compute Effective Sample Size for array of weights, ignorring any repeated elements
function ess(w::AbstractArray{Float64,1}, repeat::AbstractArray{Int64,1} )
  @assert(length(w)==length(repeat))
  sumw = sum(w[findall(x->x==0,repeat)])
  sumw2 = sum(w[findall(x->x==0,repeat)].^2)
  return sumw*sumw/sumw2
end

# Common summary stats and distances
function calc_summary_stats_mean_var(x::AbstractArray{T,2} )  where T <:Real
 @inbounds m = mean(x,dims=2);
 @inbounds v = var(x,dims=2,mean=vec(m));
 return vcat(m, v)
end

function calc_summary_stats_mean_var(x::AbstractArray{T,1} )  where T <:Real
 @inbounds m = mean(x);
 @inbounds v = var(x,mean=m);
 return vcat(m, v)
end

calc_dist_max(x::AbstractArray{Float64},y::AbstractArray{Float64}) = maximum(abs.(x.-y))
dist_scale = Array{Float64}(undef,0)
calc_scaled_dist_max(x::AbstractArray{Float64,1},y::AbstractArray{Float64,1}, scale::AbstractArray{Float64,1} = dist_scale) = maximum(abs.(x.-y)./scale)

function set_distance_scale(ds::Array{Float64,1})
  global dist_scale
  dist_scale = copy(ds)
end
function set_distance_scale( plan::abc_pmc_plan_type, theta::AbstractArray{Float64,1}; num_draw::Integer = 40,
                  ss::AbstractArray{T,1} = plan.calc_summary_stats(plan.gen_data(theta)) ) where T
  dist = Array{Float64}(undef,length(ss),num_draw)
  for i in 1:num_draw
    dist[:,i] = abs.(ss.-plan.calc_summary_stats(plan.gen_data(theta)))
  end
  set_distance_scale(vec(median(dist,dims=2)))
end

# Summary Stats & Distances based on Empirical CDF
# using Distributions
calc_summary_stats_ecdf(x::AbstractArray) = x # EmpiricalUnivariateDistribution(x)

#function calc_dist_ks(x::EmpiricalUnivariateDistribution, y::EmpiricalUnivariateDistribution)
function calc_dist_ks(x::AbstractArray{T1,1}, y::AbstractArray{T2,1}) where {T1<:Real, T2<:Real}
   maxd = 0.0
   cdf_x = ecdf(x)
   cdf_y = ecdf(y)
   #for v in x.values
   for v in x
     #d = abs(cdf(x,v)-cdf(y,v))
     d = abs(cdf_x(v)-cdf_y(v))
     if d>maxd maxd = d end
   end
   #for v in y.values
   for v in y
     #d = abs(cdf(x,v)-cdf(y,v))
     d = abs(cdf_x(v)-cdf_y(v))
     if d>maxd maxd = d end
   end
   maxd
end

function normalize_vector_of_log_values(x::Vector{Float64})
      logsum = logsumexp(x)
      param_vector = exp(param_vector-logsum)
end


"""
```
imgb = binarize(Intermodes(), img)
binarize!(Intermodes(), img)
```

Under the assumption that the histogram is bimodal the images histogram is
iterativly smoothed untill two modes remain. The threshold is then set to the
average value of the two thresholds. Falls back on
`HistogramThresholding.UnimodalRosin()` if only one mode is found.

# Arguments

The function arguments are described in more detail below.

##  `img`

An `AbstractGray` image which is to be binarized into background (`Gray(0)`)
and foreground (`Gray(1)`) pixels.

## Reference

C. A. Glasbey, “An Analysis of Histogram-Based Thresholding Algorithms,” CVGIP: Graphical Models and Image Processing, vol. 55, no. 6, pp. 532–537, Nov. 1993. [doi:10.1006/cgip.1993.1040](https://doi.org/10.1006/cgip.1993.1040)
"""
function binarize(algorithm::Intermodes,  img::AbstractArray{T,2}) where T <: AbstractGray
  binarize!(algorithm, copy(img))
end

function binarize!(algorithm::Intermodes,  img::AbstractArray{T,2}) where T <: AbstractGray
  edges, counts = build_histogram(img,  256)
  t = find_threshold(HistogramThresholding.Intermodes(), counts[1:end], edges)
  map!(img,img) do val
    val < t ? Gray(0) : Gray(1)
  end
end


module Day7
export day7

function parseLine(line)
    m = match(r"^([\w\s]*) bags contain ([\w\s,]*)\.$", line)
    if occursin("no other bags", m[2])
        return m[1], ()
    end
    m[1], map(split(m[2], ',')) do s
        let matches = match(r"(\d*) ([a-z\s]*) bags?", strip(s))
            parse(Int, matches[1]), matches[2]
        end
    end
end

function readInput(inputfile="inputs/input07.txt")
    data = Dict{String, Any}()
    for line in eachline(inputfile)
        newBag, contains = parseLine(line)
        data[newBag] = contains
    end
    data
end

function day7(inputfile="inputs/input07.txt")
    data = readInput(inputfile)
    (day7Part1(data), day7Part2(data))
end

function day7Part1(data)
    ## invert relationship data[bag] -> contains
    ## need map bag -> contained in
    containedIn = Dict{String, Any}()
    for (bag, contains) in data
        for (amount, containedBag) in contains
            if haskey(containedIn, containedBag)
                push!(containedIn[containedBag], bag)
            else
                containedIn[containedBag] = [bag]
            end
        end
    end

    ## find all connected nodes with "shiny gold"
    found = Set{String}()
    todo = containedIn["shiny gold"]
    while length(todo) > 0
        current = pop!(todo)
        if !(current in found) # have we looked at current already?
            push!(found, current)
            if haskey(containedIn, current) # current might be a root
                append!(todo, containedIn[current])
            end
        end
    end
    length(found)

end

function day7Part2(data)
    countBagsContained(data, "shiny gold")
end

## Little benchmark showed: memoization offers in fact NO speedup :(
countBagsContained(data, bag) = countBagsContained(data, bag, Dict{String, Int}())
function countBagsContained(data, bag, alreadyComputed)
    if !haskey(alreadyComputed, bag)
        amountContained = 0
        for (amount, subBag) in get(data, bag, ())
            amountContained += amount * (1 + countBagsContained(data, subBag, alreadyComputed))
        end
        alreadyComputed[bag] = amountContained
    end
    alreadyComputed[bag]
end

end #module


#=
   https://dtai.cs.kuleuven.be/problog/tutorial/basic/02_bayes.html
   """
   Since the random variables in the Bayesian network are all Boolean, we only need a single literal
   in the head of the rules. We can extend the Bayesian network to have a multi-valued variable by
   indicating the severity of the earthquake. The literal earthquake now has three possible values
   none, mild, heavy instead of previously two (no or yes).
   """

   Summary Statistics
   parameters      mean       std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64     Float64   Float64       Float64

    burglary    0.9798    0.1407     0.0014    0.0022   4381.1474    1.0005      499.2192
  earthquake    2.7649    0.4544     0.0045    0.0060   4707.8022    1.0000      536.4405
       alarm    0.9898    0.1005     0.0010    0.0015   4266.5933    0.9999      486.1661
    calls[1]    1.0000    0.0000     0.0000    0.0000         NaN       NaN           NaN
    calls[2]    0.8061    0.3954     0.0040    0.0052   4662.4821    1.0001      531.2764

   Distributions of variable earthquake
   none       =>    7782  (0.778200)
   mild       =>    2085  (0.208500)
   heavy      =>     133  (0.013300)

   Distributions of variable alarm (num:0)
   1.00000 =>    9898  (0.989800)
   0.00000 =>     102  (0.010200)

   Distributions of variable burglary (num:0)
   1.00000 =>    9798  (0.979800)
   0.00000 =>     202  (0.020200)

   Cf ~/blog/alarm_multi.blog
      ~/psi/alarm_multi.psi
      ~/webppl/alarm_multi.wppl
=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function alarm_multi()

    heavy = 1
    mild  = 2
    none  = 3
    earthquake_status = [heavy,mild,none]

    john = 1
    mary = 2
    people = [john,mary]
    num_people = length(people)

    burglary ~ flip(0.7)
    earthquake ~ Categorical([0.01,0.19,0.8])

    alarm ~ if     burglary && earthquake  == heavy flip(0.9)
            elseif burglary && earthquake  == mild  flip(0.85)
            elseif burglary && earthquake  == none  flip(0.80)
            elseif !burglary && earthquake == heavy flip(0.30)
            elseif !burglary && earthquake == mild  flip(0.10)
            else flip(0.0)
            end

    calls = tzeros(num_people)
    for p in people
        if alarm
            calls[p] ~ flip(0.8)
        else
            calls[p] ~ flip(0.01)
        end
    end

    true ~ Dirac(calls[john] == true)
    true ~ Dirac(calls[mary] == true)

    # true ~ Dirac(calls[john] == true)
    # true ~ Dirac(calls[mary] == false)

    return earthquake
end

model = alarm_multi()

num_chns = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 100_000)

# chns = sample(model, PG(15), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)
#

display(chns)
# display(plot(chns))

show_var_dist_pct(chns,:earthquake,["heavy","mild","none"])
show_var_dist_pct(chns,:alarm)
show_var_dist_pct(chns,:burglary)


# Julia program to implement Circle Sort

# Performs recursive circular swaps and returns true if atleast one swap occurs
function rec_sort(arr, high, low)
    isSwap = false

    # If concerned array is empty, Return False
    if (high == low)
        return false
    end

    # Storing the values of beg, end to later use in the recursive call
    highA = high
    lowA = low
    while( highA < lowA)
        if (arr[highA] > arr[lowA])
            arr[highA], arr[lowA] = arr[lowA], arr[highA]
            isSwap = true;
        end
        highA = highA + 1
        lowA = lowA - 1
    end

    # If the array has odd number of elements
    if (highA == lowA)
        if (arr[highA] > arr[lowA + 1])
            arr[high], arr[lowA+1] = arr[lowA+1], arr[high]
            isSwap = true
        end
    end

    mid = (low - high) ÷ 2
    isSwapA = rec_sort(arr, high, high+mid)
    isSwapB = rec_sort(arr, high+mid+1, low)

    return (isSwap || isSwapA || isSwapB)
end

function circle_sort(arr, n)
    while (rec_sort(arr, 1, n))
    end
end

print("How many numbers do you want to sort? ")
n = readline()
n = parse(Int64, n)
if (n <= 0)
    println("No numbers to sort!!!")
    exit()
end
arr = Int[]
print("Enter the numbers: ")
arr = [parse(Int, num) for num in split(readline())]
circle_sort(arr, n)
print("The numbers in sorted order is ")
print(arr)

"""
Time Complexity - O(n * log(n)), where 'n' is the size of the array
Space Complexity - O(n)

SAMPLE INPUT AND OUTPUT

SAMPLE I

How many numbers do you want to sort? 6
Enter the numbers: 76 243 89 25 8337 72
The numbers in sorted order is [25, 72, 76, 89, 243, 8337]

SAMPLE II

How many numbers do you want to sort? -3
No numbers to sort!!!
"""


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    BisectFractionPartition(normal, fraction=0.5, maxiter=10)

A method for partitioning spatial objects into two half spaces
defined by a `normal` direction and a `fraction` of points.
The partition is returned within `maxiter` bisection iterations.
"""
struct BisectFractionPartition{Dim,T} <: PartitionMethod
  normal::Vec{Dim,T}
  fraction::Float64
  maxiter::Int

  function BisectFractionPartition{Dim,T}(normal, fraction, maxiter) where {Dim,T}
    new(normalize(normal), fraction, maxiter)
  end
end

BisectFractionPartition(normal::Vec{Dim,T}, fraction=0.5, maxiter=10) where {Dim,T} =
  BisectFractionPartition{Dim,T}(normal, fraction, maxiter)

BisectFractionPartition(normal::NTuple{Dim,T}, fraction=0.5, maxiter=10) where {Dim,T} =
  BisectFractionPartition(Vec(normal), fraction, maxiter)

function partition(object, method::BisectFractionPartition)
  bbox = boundingbox(object)
  n = method.normal
  f = method.fraction
  c = coordinates(center(bbox))
  d = diagonal(bbox)

  # maximum number of bisections
  maxiter = method.maxiter

  iter = 0; p = 0
  a = c - d/2 * n
  b = c + d/2 * n
  while iter < maxiter
    m = (a + b) / 2

    p = partition(object, BisectPointPartition(n, Point(m)))
    g = nelements(p[1]) / nelements(object)

    g ≈ f && break
    g > f && (b = m)
    g < f && (a = m)

    iter += 1
  end

  p
end


@testset "$TEST $G" begin

g5 = DG(4)
add_edge!(g5,1,2); add_edge!(g5,2,3); add_edge!(g5,1,3); add_edge!(g5,3,4)
z = katz_centrality(g5, 0.4)
@test round.(z, digits=2) == [0.32, 0.44, 0.62, 0.56]

end # testset


export continuity

Space(d::IntervalDomain)=Chebyshev(d)

identity_fun(d::Interval)=Fun(eltype(d)[(d.b+d.a)/2,(d.b-d.a)/2],Chebyshev(d))

## Calculus

# the default domain space is higher to avoid negative ultraspherical spaces
Integral(d::IntervalDomain,n::Integer) = Integral(Ultraspherical(1,d),n)

for Func in (:DefiniteIntegral,:DefiniteLineIntegral)
    @eval begin
        $Func(d::IntervalDomain) = $Func(JacobiWeight(-.5,-.5,Chebyshev(d)))
        function $Func(α::Number,β::Number,d::IntervalDomain)
            @assert α == β
            @assert round(Int,α+.5) == α+.5
            @assert round(Int,α+.5) >= 0
            $Func(JacobiWeight(α,β,Ultraspherical(round(Int,α+.5),d)))
        end
        $Func(α::Number,β::Number) = $Func(α,β,Interval())
    end
end

## Evaluation

function continuity(d::Union{Vector,Tuple},order::Integer)

    m=length(d)
    B=zeros(Operator{mapreduce(eltype,promote_type,d)},m-1,m)

    for k=1:m-1
        B[k,k]=Evaluation(d[k],true,order)
        B[k,k+1]=-Evaluation(d[k+1],false,order)
    end
    B
end

function continuity(d::Union{Vector,Tuple},kr::UnitRange)
    @assert first(kr)==0
    m=length(d)
    B=zeros(Operator{mapreduce(eltype,promote_type,d)},length(kr)*(m-1),m)
    for r in kr
        B[(m-1)*r+1:(m-1)*(r+1),:]=continuity(d,r)
    end
    B
end

for OP in (:dirichlet,:neumann,:periodic)
    @eval $OP(d::Tuple)=$OP([d...])
end

for DT in (:IntervalDomain,:Space)
    @eval begin
        function dirichlet{T<:$DT}(d::Vector{T})
            m=length(d)
            B=zeros(Operator{mapreduce(eltype,promote_type,d)},2,m)
            B[1,1]=ldirichlet(d[1]);B[2,end]=rdirichlet(d[end])
            [B;
            continuity(d,0:1)]
        end

        function neumann{T<:$DT}(d::Vector{T})
            m=length(d)
            B=zeros(Operator{mapreduce(eltype,promote_type,d)},2,m)
            B[1,1]=ldirichlet(d[1]);B[2,end]=rdirichlet(d[end])
            [B;
            continuity(d,0:1)]
        end

        function periodic{T<:$DT}(d::Vector{T})
            m=length(d)
            B=zeros(Operator{mapreduce(eltype,promote_type,d)},2,m)
            B[1,1]=ldirichlet(d[1]);B[1,end]=-rdirichlet(d[end])
            B[2,1]=lneumann(d[1]);B[2,end]=-rneumann(d[end])
            [B;
            continuity(d,0:1)]
        end
    end
end


#=
   Same rank problem

   From
   Gunnar Blom, Lars Holst, Dennis Sandell:
   "Problems and Snapshots from the World of Probability"
   Page 22.

   100 students are seating (randomly) on seats numbered 0 to 99.
   The lectures says: assume that I rank you according to your age.
   How many will have the same rank number as the seat number?

   Here are three different approaches to this problem which all
   give the same expectation, i.e. E(same_rank) = 1.

   1) Implementing an "all_different" constraint, i.e.
      ensure that it's an permutation.

      n: 6 all_diff: true
      mean(numSameRank: 0.9653

      Distributions of variable numSameRank (num:0)
      1.00000 =>    3942  (0.394200)
      0.00000 =>    3638  (0.363800)
      2.00000 =>    1796  (0.179600)
      3.00000 =>     377  (0.037700)
      4.00000 =>     247  (0.024700)

   2) Generating randomInteger(n) _without_ the all_different
      constraint.

      n: 6 all_diff: false
      mean(numSameRank: 1.0145

      Distributions of variable numSameRank (num:0)
      1.00000 =>    4024  (0.402400)
      0.00000 =>    3312  (0.331200)
      2.00000 =>    1982  (0.198200)
      3.00000 =>     586  (0.058600)
      4.00000 =>      82  (0.008200)
      5.00000 =>      13  (0.001300)
      6.00000 =>       1  (0.000100)

      n: 100 all_diff: false
      mean(numSameRank: 1.0009

   3) Generate n 0/1 numbers with probability of 1 is 1/n.
      Summing this array gives E(same_rank).
      This is the method mentioned in the book, and is implemented
      as same_rank2.

      same_rank2 n: 6
      mean(numSameRank: 1.0064
      Distributions of variable numSameRank (num:0)
      1.00000 =>    4017  (0.401700)
      0.00000 =>    3321  (0.332100)
      2.00000 =>    2035  (0.203500)
      3.00000 =>     539  (0.053900)
      4.00000 =>      80  (0.008000)
      5.00000 =>       8  (0.000800)

      same_rank2 n: 100
      mean(numSameRank: 1.0219

      same_rank2 n: 1000
      mean(numSameRank: 0.9891

     I.e. all these has - also - the same expectation: E(same_rank) = (about) 1.

   Cf ~/webppl/same_rank.wppl

=#
using Turing, StatsPlots
include("jl_utils.jl")

@model function same_rank1(n=6,alldiff=true)
    x ~ filldist(DiscreteUniform(1,n),n)
    if (alldiff)
        # Ensure permutation (see all_different.jl)
        for i in 1:n
            for j in 1:n
                if i != j
                    true ~ Dirac(x[i] != x[j])
                end
            end
        end
    end
    numSameRank ~ Dirac(sum([x[i] == i ? 1 : 0 for i in 1:n]))
end

function run_same_rank1(n, all_diff)
    println("\nsame rank1 n: $n all_diff: $all_diff")
    model = same_rank1(n,all_diff)

    # chns = sample(model, Prior(), 10_000)
    # chns = sample(model, MH(), 10_000)
    # chns = sample(model, PG(5), 10_000)
    chns = sample(model, SMC(), 10_000)
    # chns = sample(model, IS(), 10_000)

    # display(chns)
    # display(plot(chns))
    println("mean(numSameRank: ", mean_val(chns,:numSameRank))

    if n < 10
        show_var_dist_pct(chns,:numSameRank)
    end
end

run_same_rank1(6,true)
run_same_rank1(6,false)
run_same_rank1(100,false)

# Generate n 0/1 values with probability 1/n of getting a 1
@model function same_rank2(n=6)
    x ~ filldist(flip(1/n),n)
    numSameRank ~ Dirac(sum(x))
end

function run_same_rank2(n)
    println("\nsame_rank2 n: $n")
    model = same_rank2(n)

    # chns = sample(model, Prior(), 10_000)
    # chns = sample(model, MH(), 10_000)
    # chns = sample(model, PG(5), 10_000)
    chns = sample(model, SMC(), 10_000)
    # chns = sample(model, IS(), 10_000)

    # display(chns)
    # display(plot(chns))
    println("mean(numSameRank: ", mean_val(chns,:numSameRank))
    if n < 10
        show_var_dist_pct(chns,:numSameRank)
    end
end

run_same_rank2(6)
run_same_rank2(100)
run_same_rank2(1000)


struct PriorityQueue{K, V}
    values::Array{Pair{K,V}}
    PriorityQueue{K,V}() where {K,V} = new(Array{Pair{K,V}}(undef, 0))
end

function push!(pq::PriorityQueue{K,V}, pair::Pair{K,V}) where {K,V}
    Base.push!(pq.values, pair)
    index = length(pq.values)
    while index > 1
        parent_index = div(index - 1, 2) + 1
        (pq.values[parent_index].first >= pair.first) && break
        pq.values[index] = pq.values[parent_index]
        index = parent_index
    end
    pq.values[index] = pair
    return nothing
end

function pop_highest!(pq::PriorityQueue)
    isempty(pq.values) && return nothing
    max_value = highest(pq)
    value = pop!(pq.values)
    index = 1
    while (index * 2) <= length(pq.values)
        cind1, cind2 = index * 2, index * 2 + 1
        flag = (cind2 >= length(pq.values) || pq.values[cind1] > pq.values[cind2])
        bigger_index = flag ? cind1 : cind2
        pq.values[bigger_index] <= value && break
        pq.values[index] = pq.values[bigger_index]
        index = bigger_index
    end
    pq.values[index] = value
    return max_value
end

pop_lowest!(pq::PriorityQueue) = !isempty(pq.values) ? pop!(pq.values) : nothing
highest(pq::PriorityQueue) = !isempty(pq.values) ? pq.values[1] : nothing
lowest(pq::PriorityQueue) = !isempty(pq.values) ? pq.values[end] : nothing


"""
Given the number of rows and columns, return an NTuple{4,Float64} containing
`(width, height, vspace, hspace)`, where `width` and `height` are the
width and height of each subplot and `vspace` and `hspace` are the vertical
and horizonal spacing between subplots, respectively.
"""
function sizes(nr::Int, nc::Int, subplot_titles::Bool=false)
    # NOTE: the logic of this function was mostly borrowed from plotly.py
    dx = 0.2 / nc
    dy = subplot_titles ? 0.55 / nr : 0.3 / nr
    width = (1. - dx * (nc - 1)) / nc
    height = (1. - dy * (nr - 1)) / nr
    vspace = nr == 1 ? 0.0 : (1 - height*nr)/(nr-1)
    hspace = nc == 1 ? 0.0 : (1 - width*nc)/(nc-1)
    width, height, vspace, hspace
end

function gen_layout(nr, nc, subplot_titles::Bool=false)
    w, h, dy, dx = sizes(nr, nc, subplot_titles)

    x = 0.0  # start from left
    y = 1.0  # start from top

    out = Layout()
    for col in 1:nc

        y = 1.0 # reset y as we start a new col
        for row in 1:nr
            subplot = sub2ind((nc, nr), col, row)

            out["xaxis$subplot"] = Dict{Any,Any}(:domain=>[x, x+w],
                                                 :anchor=> "y$subplot")
            out["yaxis$subplot"] = Dict{Any,Any}(:domain=>[y-h, y],
                                                 :anchor=> "x$subplot")

            y -= nr == 1 ? 0.0 : h + dy
         end

         x += nc == 1 ? 0.0 : w + dx
    end

    out

end

function handle_titles!(big_layout, sub_layout, ix::Int)
    # don't worry about it if the sub_layout doesn't have a title
    if !haskey(sub_layout.fields, "title") && !haskey(sub_layout.fields, :title)
        return big_layout
    end

    # check for symbol or string
    nm = haskey(sub_layout.fields, "title") ? "title" : :title

    ann = Dict{Any,Any}(:font => Dict{Any,Any}(:size => 16),
                        :showarrow => false,
                        :text => pop!(sub_layout.fields, nm),
                        :x => mean(big_layout["xaxis$(ix).domain"]),
                        :xanchor => "center",
                        :xref => "paper",
                        :y => big_layout["yaxis$(ix).domain"][2],
                        :yanchor => "bottom",
                        :yref => "paper")
    anns = get(big_layout.fields, :annotations, Dict{Any,Any}[])
    push!(anns, ann)
    big_layout[:annotations] = anns
    big_layout
end

function _cat(nr::Int, nc::Int, ps::Plot...)
    copied_plots = Plot[copy(p) for p in ps]
    subplot_titles = any(map(x -> haskey(x.layout.fields, :title) ||
                                  haskey(x.layout.fields, "title"), ps))
    layout = gen_layout(nr, nc, subplot_titles)

    for col in 1:nc, row in 1:nr
        ix = sub2ind((nc, nr), col, row)

        for trace in copied_plots[ix].data
            trace["xaxis"] = "x$ix"
            trace["yaxis"] = "y$ix"
        end

        handle_titles!(layout, copied_plots[ix].layout, ix)
        layout["xaxis$ix"] = merge(copied_plots[ix].layout["xaxis"], layout["xaxis$ix"])
        layout["yaxis$ix"] = merge(copied_plots[ix].layout["yaxis"], layout["yaxis$ix"])
    end

    Plot(vcat([p.data for p in copied_plots]...), layout)
end

Base.hcat(ps::Plot...) = _cat(1, length(ps), ps...)
Base.vcat(ps::Plot...) = _cat(length(ps), 1,  ps...)
Base.vect(ps::Plot...) = vcat(ps...)

function Base.hvcat(rows::Tuple{Vararg{Int}}, ps::Plot...)
    nr = length(rows)
    nc = rows[1]

    for (i, c) in enumerate(rows[2:end])
        c == nc || error("Saw $c columns in row $(i+1), expected $nc")
    end
    _cat(nr, nc, ps...)
end

# methods on syncplot
Base.hcat{TP<:SyncPlot}(sps::TP...) = TP(hcat([sp.plot for sp in sps]...))
Base.vcat{TP<:SyncPlot}(sps::TP...) = TP(vcat([sp.plot for sp in sps]...))
Base.vect{TP<:SyncPlot}(sps::TP...) = vcat(sps...)
Base.hvcat{TP<:SyncPlot}(rows::Tuple{Vararg{Int}}, sps::TP...) =
    TP(hvcat(rows, [sp.plot for sp in sps]...))


module Example
import Random

"""
    MyPareto(alpha::Float64,x0::Float64)

Pareto distribution for sampling. This does not depend on Distributions and
the RNG seed can be set. This may not be needed for future versions of Julia.
"""
struct MyPareto
    alpha::Float64
    x0::Float64
end
MyPareto(alpha=1.0) = MyPareto(alpha, 1.0)
Random.rand(p::MyPareto) = p.x0 * Random.rand()^(-one(p.alpha) / p.alpha)

"""
    makeparetodata()

Create a reproducible sorted array of Pareto distribution samples.
"""
function makeparetodata(alpha=0.5, seed=11)
    Random.seed!(seed)
    d = MyPareto(alpha)
    return [Random.rand(d) for _=1:10^6] |> sort!
end

end # module Example


# # Eady model of Baroclinic Instability
#
# A simulation of the growth of barolinic instability in the Eady n-layer model
# when we impose a vertical mean flow shear as a difference $\Delta U$ in the
# imposed, domain-averaged, zonal flow at each layer.

using FourierFlows, Plots, Printf

using FFTW: rfft, irfft
using Statistics: mean

import GeophysicalFlows.MultiLayerQG
import GeophysicalFlows.MultiLayerQG: energies

# ## Numerical parameters and time-stepping parameters
nx = 256                  # 2D resolution = nx^2
ny = nx

stepper = "FilteredRK4"   # timestepper
dt = 2e-3      # timestep
nsteps = 20000            # total number of time-steps
nsubs  = 100              # number of time-steps for plotting (nsteps must be multiple of nsubs)

# ## Physical parameters
Lx = 2π         # domain size
 μ = 1e-2       # bottom drag
 β = 0.1          # the y-gradient of planetary PV
 g = 1.0
f₀ = 2.0

# Vertical grid
nlayers = 5     # number of layers

total_depth = 1
Δh = total_depth / nlayers
z = [ (i - 1/2) * Δh for i = nlayers:-1:1 ] .- total_depth
H = [ Δh for i = 1:nlayers ]

# Density stratification
ρ₀ = 1
N² = 5

ρ_resting(z) = ρ₀ * (1 - N² * z / g)
ρ = ρ_resting.(z)

@info @sprintf("The largest Rossby radius of deformation is %.2f",
               sqrt(N²) * total_depth / f₀)

@info @sprintf("The smallest Rossby radius of deformation is %.2f",
               sqrt(N²) * Δh / f₀)

# Background shear
ΔU = 1 / (nlayers - 1)
U = [i * ΔU for i=nlayers-1:-1:0]

twod_grid = TwoDGrid(nx, Lx)

X, Y = gridpoints(twod_grid)
σx, σy = 0.7, 0.5
eta = f₀ / sum(H) * Δh * @. exp(-X^2/σx^2 -Y^2/σy^2)

# ## Problem setup
#
# We initialize a `Problem` by providing a set of keyword arguments,

prob = MultiLayerQG.Problem(nlayers,
                             nx = nx,
                             Lx = Lx,
                             f₀ = f₀,
                              g = g,
                              H = H,
                              ρ = ρ,
                              U = U,
                             dt = dt,
                        stepper = stepper,
                            eta = eta,
                              μ = μ,
                              β = β
                           )

# and define some shortcuts.
sol, clock, params, vars, grid = prob.sol, prob.clock, prob.params, prob.vars, prob.grid
x, y = grid.x, grid.y

# ## Setting initial conditions

# Our initial condition is some small amplitude random noise. We smooth our initial
# condidtion using the `timestepper`'s high-wavenumber `filter`.

qᵢ  = randn((nx, ny, nlayers))
for j in 1:nlayers
  qᵢ[:, :, j] = qᵢ[:, :, j] .- mean(qᵢ[:, :, j])
end
qhᵢ = prob.timestepper.filter .* rfft(qᵢ, (1, 2)) # only apply rfft in dims=1, 2
qᵢ  = irfft(qhᵢ, grid.nx, (1, 2)) # only apply irfft in dims=1, 2

MultiLayerQG.set_q!(prob, qᵢ)

KEᵢ, PEᵢ = MultiLayerQG.energies(prob)
KEᵢ

ψhᵢ = similar(qhᵢ)

MultiLayerQG.streamfunctionfrompv!(ψhᵢ, qhᵢ, prob.params, prob.grid)

for j in 1:nlayers
  ψhᵢ[:, :, j] = sqrt(5e-8 / KEᵢ[j]) * ψhᵢ[:, :, j]
end
ψᵢ  = irfft(ψhᵢ, grid.nx, (1, 2))

MultiLayerQG.set_ψ!(prob, ψᵢ)
nothing # hide

# ## Diagnostics

# Create Diagnostics -- `energies` function is imported at the top.
E = Diagnostic(energies, prob; nsteps=nsteps)
diags = [E] # A list of Diagnostics types passed to "stepforward!" will  be updated every timestep.

nothing # hide

# ## Output

# We choose folder for outputing `.jld2` files and snapshots (`.png` files).
filepath = "."
plotpath = "./plots_eady"
plotname = "snapshots"
filename = joinpath(filepath, "eady.jld2")
nothing # hide

# Do some basic file management
if isfile(filename); rm(filename); end
if !isdir(plotpath); mkdir(plotpath); end
nothing # hide

# And then create Output
get_sol(prob) = prob.sol # extracts the Fourier-transformed solution

out = Output(prob, filename, (:sol, get_sol))
nothing # hide

# ## Visualizing the simulation

# We define a function that plots the potential vorticity field and the evolution
# of energy and enstrophy.

function plot_output(prob)
  q = prob.vars.q
  x, y = prob.grid.x, prob.grid.y
  Lx, Ly = prob.grid.Lx, prob.grid.Ly

  layout = @layout Plots.grid(2, Int(round(nlayers/2)+1)+1)

  p = plot(layout=layout, size=(1000, 600), dpi=180)

  for m in 1:nlayers

    heatmap!(p[m], x, y, q[:, :, m],
         aspectratio = 1,
              legend = :none,
                   c = :deep,
               xlims = (-Lx/2, Lx/2),
               ylims = (-Ly/2, Ly/2),
               title = "q layer "*string(m),
              xticks = :none,
              yticks = :none,
            colorbar = false,
          framestyle = :box

         )
  end

  plot!(p[nlayers+1], nlayers,
             # label = ["KE1" "KE2"],
            legend = false,
         linewidth = 2,
             alpha = 0.7,
             xlims = (-0.05*μ*nsteps*dt, 1.01*μ*nsteps*dt),
             ylims = (1e-8, 1e2),
            yscale = :log10,
            yticks = 10.0.^(-8:2:2),
            xlabel = "μt",
             title = "KE at layers")

  plot!(p[nlayers+2], nlayers-1,
             # label = "PE",
            legend = false,
         linewidth = 2,
             alpha = 0.7,
             xlims = (-0.05*μ*nsteps*dt, 1.01*μ*nsteps*dt),
             ylims = (1e-9, 1e1),
            yscale = :log10,
            yticks = 10.0.^(-8:2:2),
            xlabel = "μt",
             title = "PE at interfaces")

  return p
end

nothing # hide

# ## Time-stepping the `Problem` forward

# Finally, we time-step the `Problem` forward in time.

p = plot_output(prob)

startwalltime = time()

anim = @animate for j = 0:Int(nsteps / nsubs)

    cfl = clock.dt * maximum([maximum(abs, vars.u) / grid.dx, maximum(abs, vars.v) / grid.dy])

    log = @sprintf("step: %04d, dt: %.3f, t: %0.2f, cfl: %.2f, walltime: %.2f min",
                    clock.step, clock.dt, clock.t, cfl, (time()-startwalltime)/60)

    log_KE = @sprintf("KE₁: %.3e     KE₂: %.3e     KE₃: %.3e     KE₄: %.3e     KE₅: %.3e",
                  E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][1][3], E.data[E.i][1][4], E.data[E.i][1][5])

    log_PE = @sprintf("      PE₁.₅: %.3e   PE₂.₅: %.3e   PE₃.₅: %.3e   PE₄.₅: %.3e",
                  E.data[E.i][2][1], E.data[E.i][2][2], E.data[E.i][2][3], E.data[E.i][2][4])

    # log_KE = @sprintf("KE₁: %.3e     KE₂: %.3e     KE₃: %.3e",
    #               E.data[E.i][1][1], E.data[E.i][1][2], E.data[E.i][1][3])
    #
    # log_PE = @sprintf("      PE₁.₅: %.3e   PE₂.₅: %.3e",
    #               E.data[E.i][2][1], E.data[E.i][2][2])
    #

    if j % (500 / nsubs) == 0
      println(log)
      println(log_KE)
      println(log_PE)
    end

    for m in 1:nlayers
      p[m][1][:z] = @. vars.q[:, :, m] + params.β * grid.y'
      if m==nlayers
        p[m][1][:z] = @. vars.q[:, :, m] + params.β * grid.y' + params.eta
      end
      push!(p[nlayers+1][m], μ*E.t[E.i], E.data[E.i][1][m])
    end

    for m in 1:nlayers-1
      push!(p[nlayers+2][m], μ*E.t[E.i], E.data[E.i][2][m])
    end

    stepforward!(prob, diags, nsubs)

    MultiLayerQG.updatevars!(prob)

end

mp4(anim, "multilayerqg_eady.mp4", fps=18)

# ## Save

#=
# Finally save the last snapshot.
savename = @sprintf("%s_%09d.png", joinpath(plotpath, plotname), clock.step)
savefig(savename)
=#


# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

"""
Either `Rational` or `Euclidean`, used for [`Spline`](@ref)s and [`SplineSurface`](@ref)s.
"""
abstract type CurveType end
abstract type Rational <: CurveType end
abstract type Euclidean <: CurveType end
export CurveType

"""
    Spline{P<:CurveType,S<:Number,N,M}

`M` is the curve order, i.e., the highest power of the parameterizing variable, u.
`P` determines the [`CurveType`](@ref).

All Spline types must implement:
```julia
point(curve,u)
```
and have field `controlpolygon`
"""
abstract type Spline{P<:CurveType,S<:Number,N,M} end

"""
    SplineSurface{P,S,N,M} <: ParametricSurface{S,N}

Curve order, `M`, is the same in the u and v direction and fixed over all spans.
`P` determines the [`CurveType`](@ref).
"""
abstract type SplineSurface{P<:CurveType,S,N,M} <: ParametricSurface{S,N} end
export Spline, SplineSurface

function toeuclidean(point)
    return SVector{length(point) - 1}([point[i] / point[end] for i in 1:(lastindex(point) - 1)])
end

function euclideancontrolpoints(curve::Spline{Rational,S,N,M}) where {S,N,M}
    return toeuclidean.(curve.controlpolygon)
end

function euclideancontrolpoints(curve::Spline{Euclidean,S,N,M}) where {S,N,M}
    return curve.controlpolygon
end

#Bernstein polynomial functions
#B(i,n,u) is broken into Bc,Bu because these two functions will be needed when constructing the matrics to compute moving lines or moving planes
Bc(i, n) = factorial(n) ÷ (factorial(i) * factorial(n - i)) #coefficient of the i,n Bernstein basis polynomial
Bu(i, n, u) = u^i * (1 - u)^(n - i) #polynomial part of Bernstein basis polynomial
B(i, n, u) = Bc(i, n) * Bu(i, n, u)


"""
    Models for random hypergraphs building:
        * random model
        * k-uniform model
        * d-uniform model
        * preferential-attachment
"""

"""
    random_model(nVertices::Int, nEdges::Int)

Generate a *random* hypergraph without any structural constraints.

# The Algorithm

Given two integer parameters *nVertices* and *nEdges* (the number of nodes and hyperedges, respectively),
the algorithm computes - for each hyperedge *he={1,...,m}* -
a random number *s ϵ [1, n]* (i.e. the hyperedge size).
Then, the algorithm selects uniformly at random *s* vertices from *V* to be added in *he*.
"""
function random_model(nVertices::Int, nEdges::Int)
    mx = Matrix{Union{Nothing,Bool}}(nothing, nVertices, nEdges)
    if nEdges == 0
        return Hypergraph(nVertices, nEdges)
    end
    for e in 1:size(mx,2)
        nv = rand(1:size(mx,1))
        mx[sample(1:size(mx,1), nv;replace=false), e] .= true
    end

    h = Hypergraph(mx)
    if all(length.(h.v2he) .> 0)
        return h
    else
        return random_model(nVertices, nEdges)
    end
end

"""
    random_kuniform_model(nVertices::Int, nEdges::Int, k::Int)

Generates a *k*-uniform hypergraph, i.e. an hypergraph where each hyperedge has size *k*.

# The Algorithm

The algorithm proceeds as the *randomH*, forcing the size of each hyperedge equal to *k*.
"""
function random_kuniform_model(nVertices::Int, nEdges::Int, k::Int)
    mx = Matrix{Union{Nothing,Bool}}(nothing, nVertices,nEdges)
    for e in 1:size(mx,2)
        mx[sample(1:size(mx,1), k;replace=false), e] .= true
    end
    Hypergraph(mx)
end

"""
    random_dregular_model(nVertices::Int, nEdges::Int, d::Int)

Generates a *d*-regular hypergraph, where each node has degree *d*.

# The Algorithm

The algorithm exploits the *k*-uniform approach described for the *random_kuniform_model* method
to build a *d*-regular hypergraph *H* having *nVertices* nodes and *nEdges* edges.
It returns the hypergraph H^* dual of *H*.
"""
function random_dregular_model(nVertices::Int, nEdges::Int, d::Int)
    mx = Matrix{Union{Nothing,Bool}}(nothing, nVertices,nEdges)
    for v in 1:size(mx,1)
        mx[v, sample(1:size(mx,2), d;replace=false)] .= true
    end
    Hypergraph(mx)
end

"""
    random_preferential_model(nVertices::Int, p::Real; H = random_model(5,5))

Generate a hypergraph with a preferential attachment rule between nodes, as presented in
*Avin, C., Lotker, Z., and Peleg, D.Random preferential attachment hyper-graphs. Computer Science 23 (2015).*

# The Algorithm

The algorithm starts with a random graph with 5 nodes and 5 edges. For this reason,
the generated random graph has at least 5 nodes and 5 edges.
It iteratively adds a node or a edge, according to a given parameter *p*,
which defines the probability of creating a new node or a new hyperedge.

More in detail, the connections with the new node/hyperedge are generated according to
a preferential attachment policy defined by _p_.

The so-built hypergraph will have `nVertices` vertices.

The starting hypergraph can be instantiated as preferred.
"""
function random_preferential_model(nVertices::Int, p::Real; H = random_model(5,5))
    while nhv(H) < nVertices
        r = rand()
        y = rand(1:nhv(H))
        Y = Dict{Int,Bool}()
        if r < p
            v = SimpleHypergraphs.add_vertex!(H)
            push!(Y,v=>true)
            for v in next_nodes(H,y-1)
                push!(Y,v)
            end
        else
            for v in next_nodes(H,y)
                push!(Y,v)
            end
        end
        SimpleHypergraphs.add_hyperedge!(H, vertices=Y)
    end
    H
end

"""
    next_nodes(h::Hypergraph, size::Int)

Selects nodes to add to a hyperedge.
"""
function next_nodes(h::Hypergraph, size::Int)
    nodes = Dict{Int, Bool}()

    ids = collect(1:nhv(h))
    degrees = length.(h.v2he)

    for s=1:size
        psum = Vector{Int}(undef, length(ids))

        psum[1] = degrees[ids[1]]
        for j=2:length(ids)
            psum[j] = psum[j-1] + degrees[ids[j]]
        end

        number = rand(1:psum[length(psum)])
        local bucket::Int
        index=1
        for i=1:length(psum)
            if number <= psum[i]
                bucket = ids[i]
                index = i
                break
            end
        end

        nodes[bucket] = true
        deleteat!(ids, index)
    end
    nodes
end


using ReadableNumbers
using Test

setprecision(BigFloat,128);

@test readable(BigFloat(pi)) == "3.14159_26535_89793_23846_26433_83279_50288_4195"
@test readable(factorial(BigInt(34))) == "295,232,799,039,604,140,847,618,609,643,520,000,000"
@test BigFloat(pi) == parse_readable(BigFloat, readable(BigFloat(pi)), '_')
@test factorial(BigInt(34)) == parse_readable(BigInt, readable(factorial(BigInt(34))), ',')


# general array for testing input to SpecialFunctions.jl
x32 = [-1.0; 0.0; 1e-6; 0.01:0.01:60.0]
x64 = [-1.0; 0.0; 1e-6; 0.01:0.01:600.0]

### Tests for besseli0
i032_SpecialFunctions = SpecialFunctions.besseli.(0, x32)
i064_SpecialFunctions = SpecialFunctions.besseli.(0, x64)

i0_64 = besseli0.(Float64.(x64))
i0_32 = besseli0.(Float32.(x32))

# make sure output types match input types
@test i0_64[1] isa Float64
@test i0_32[1] isa Float32

# test against SpecialFunctions.jl
@test i0_64 ≈ i064_SpecialFunctions
@test i0_32 ≈ i032_SpecialFunctions

### Tests for besseli0x
i0x32_SpecialFunctions = SpecialFunctions.besselix.(0, x32)
i0x64_SpecialFunctions = SpecialFunctions.besselix.(0, x64)

i0x_64 = besseli0x.(Float64.(x64))
i0x_32 = besseli0x.(Float32.(x32))

# make sure output types match input types
@test i0x_64[1] isa Float64
@test i0x_32[1] isa Float32

# test against SpecialFunctions.jl
@test i0x_64 ≈ i0x64_SpecialFunctions
@test i0x_32 ≈ i0x32_SpecialFunctions

### Tests for besseli1
i132_SpecialFunctions = SpecialFunctions.besseli.(1, x32)
i164_SpecialFunctions = SpecialFunctions.besseli.(1, x64)

i1_64 = besseli1.(Float64.(x64))
i1_32 = besseli1.(Float32.(x32))

# make sure output types match input types
@test i1_64[1] isa Float64
@test i1_32[1] isa Float32

# test against SpecialFunctions.jl
@test i1_64 ≈ i164_SpecialFunctions
@test i1_32 ≈ i132_SpecialFunctions

### Tests for besseli1x
i1x32_SpecialFunctions = SpecialFunctions.besselix.(1, x32)
i1x64_SpecialFunctions = SpecialFunctions.besselix.(1, x64)

i1x_64 = besseli1x.(Float64.(x64))
i1x_32 = besseli1x.(Float32.(x32))

# make sure output types match input types
@test i1x_64[1] isa Float64
@test i1x_32[1] isa Float32

# test against SpecialFunctions.jl
@test i1x_64 ≈ i1x64_SpecialFunctions
@test i1x_32 ≈ i1x32_SpecialFunctions

# test for besseli
# test small arguments and order
m = 0:1:200; x = 0.5f0:0.5f0:90.0f0
@test besseli(10, 1.0f0) isa Float32
@test besseli(2, 80.0f0) isa Float32
@test besseli(112, 80.0f0) isa Float32

for m in m, x in x
    @test besseli(m, x) ≈ Float32(SpecialFunctions.besseli(m, x))
end

#Float 64
m = 0:1:200; x = 0.1:0.5:150.0
@test besseli(10, 1.0) isa Float64
@test besseli(2, 80.0) isa Float64
@test besseli(112, 80.0) isa Float64
t = [besseli(m, x) for m in m, x in x]

@test t[10] isa Float64
@test t ≈ [SpecialFunctions.besseli(m, x) for m in m, x in x]

t = [besselix(m, x) for m in m, x in x]
@test t[10] isa Float64
@test t ≈ [SpecialFunctions.besselix(m, x) for m in m, x in x]


abstract type BandwidthSelector end

bandwidth(h::Number, kernel, rddata) = h
_string(::Number) = ""

struct ImbensKalyanaraman <: BandwidthSelector end

_string(::ImbensKalyanaraman) = "Imbens Kalyanaraman bandwidth"

function kernel_constant(::ImbensKalyanaraman, kernel)
    kernel = EquivalentKernel(kernel)
    νs = OffsetVector([kernel_moment(kernel, Val(j)) for j = 0:3], 0:3)
    πs = OffsetVector([squared_kernel_moment(kernel, Val(j)) for j = 0:3], 0:3)
    C1 = 1 / 4 * abs2((abs2(νs[2]) - νs[1] * νs[3]) / (νs[2] * νs[0] - abs2(νs[1])))
    C2_num = abs2(νs[2]) * πs[0] - 2 * νs[1] * νs[2] * πs[1] + abs2(νs[1]) * πs[2]
    C2_denom = abs2(νs[2] * νs[0] - abs2(νs[1]))
    C2 = C2_num / C2_denom
    (C2 / (4 * C1))^(1 / 5)
    #(πs[0]/abs2(νs[2]))^(1/5)
end

function bandwidth(ik::ImbensKalyanaraman, kernel::SupportedKernels, rddata::RDData)
    Z = rddata.ZsR.Zs
    Y = rddata.Ys

    rdd_df = DataFrame(Z = Z, Y = Y)

    N = length(Z)
    left_idx = findall(Z .< 0)
    right_idx = findall(Z .>= 0)
    N_left = length(left_idx)
    N_right = length(right_idx)

    # Step 1: Density and conditional variance at 0
    h₁ = 1.84 * std(Z) * N^(-1 / 5)

    h₁_left_idx = findall(-h₁ .<= Z .< 0)
    h₁_right_idx = findall(+h₁ .>= Z .>= 0)

    N_h₁_left = length(h₁_left_idx)
    N_h₁_right = length(h₁_right_idx)

    Ȳ_h₁_left = mean(Y[h₁_left_idx])
    sd_Y_h₁_left = std(Y[h₁_left_idx])
    Ȳ_h₁_right = mean(Y[h₁_right_idx])
    sd_Y_h₁_right = std(Y[h₁_right_idx])

    f̂₀ = (N_h₁_left + N_h₁_right) / 2 / N / h₁

    # Step 2: Estimation of second derivatives
    global_cubic_lm = lm(@formula(Y ~ 1 + (Z >= 0) + Z + Z^2 + Z^3), rdd_df)

    m̂₀_triple_prime = 6 * coef(global_cubic_lm)[5]

    h₂_left = 3.56 * (sd_Y_h₁_left^2 / f̂₀ / m̂₀_triple_prime^2)^(1 / 7) * N_left^(-1 / 7)
    h₂_right =
        3.56 * (sd_Y_h₁_right^2 / f̂₀ / m̂₀_triple_prime^2)^(1 / 7) * N_right^(-1 / 7)

    h₂_left_idx = findall(-h₂_left .<= Z .< 0)
    h₂_right_idx = findall(+h₂_right .>= Z .>= 0)

    N_h₂_left = length(h₂_left_idx)
    N_h₂_right = length(h₂_right_idx)

    quadratic_fit_left = lm(@formula(Y ~ 1 + Z + Z^2), rdd_df[h₂_left_idx, :])
    m̂₀_double_prime_left = 2 * last(coef(quadratic_fit_left))

    quadratic_fit_right = lm(@formula(Y ~ 1 + Z + Z^2), rdd_df[h₂_right_idx, :])
    m̂₀_double_prime_right = 2 * last(coef(quadratic_fit_right))

    # Step 3: Calculation of regularization terms

    r̂_left = 2160 * sd_Y_h₁_left^2 / N_h₂_left / h₂_left^4
    r̂_right = 2160 * sd_Y_h₁_right^2 / N_h₂_right / h₂_right^4

    Ckernel = kernel_constant(ik, kernel)# 3.4375

    regularized_squared_diff =
        abs2(m̂₀_double_prime_right - m̂₀_double_prime_left) + r̂_left + r̂_right
    ĥ_IK =
        Ckernel *
        ((sd_Y_h₁_left^2 + sd_Y_h₁_right^2) / f̂₀ / regularized_squared_diff)^(1 / 5) *
        N^(-1 / 5)
    ĥ_IK
end


## This file is outdated

export timeevolution, ETDRK4, BDF2, BDF4, BDF22

function RK(L,y,h)
    k1=L(y)
    k2=L(y+.5h*k1)
    k3=L(y+.5h*k2)
    k4=L(y+h*k3)

    y+h*(k1+2k2+2k3+k4)/6.
end

function BDF2(B,A::Operator,g::Function,bcs,u0,h,m,glp,tol=1000eps())
    SBDF2 = [B;I-2.0/3.0*h*A]

    u1=u0
    u2=chop(RK(g,u1,h),tol)
    u2,u1  = chop(SBDF2\[bcs;1/3.0*(4u2-u1)],tol),u2
    push!(glp,u2)
    yield()

    for k=1:m
        u2,u1 = chop(RK(g,u2,h),tol),u2
        u2,u1  = chop(SBDF2\[bcs;1/3.0*(4u2-u1)],tol),u2
        push!(glp,u2)
        yield()
    end

    u2
end

function BDF2(L::Operator,N::Function,u0,h,m,glp,tol=1000eps())
    SBDF2 = I-2.0/3.0*h*L

    u1=u0
    u2=chop(RK(N,u1,h),tol)
    u2,u1  = chop(SBDF2\(1/3.0*(4u2-u1)),tol),u2
    push!(glp,u2)
    yield()

    for k=1:m
        u2,u1 = chop(RK(N,u2,h),tol),u2
        u2,u1  = chop(SBDF2\(1/3.0*(4u2-u1)),tol),u2
        push!(glp,u2)
        yield()
    end

    u2
end

function ETDRK4(L::Operator,N::Function,u,t,h,m,glp,tol=10eps())
    z = L*h
    ez,ez2,h2ez2m1 = cache(exp(z)),cache(exp(z/2)),cache(h/2*expm1(z/2))
    hezα,h2ezβ,hezγ = cache(h*expα(z)),cache(2h*expβ(z)),cache(h*expγ(z))
    push!(glp,u)
    yield()

    for k=1:m
        ez2u = ez2*u
        Nut = chop(N(u,t),tol)

        a = ez2u + h2ez2m1*Nut
        Nath2 = chop(N(a,t+h/2),tol)
        b = ez2u + h2ez2m1*Nath2
        Nbth2 = chop(N(b,t+h/2),tol)
        c = ez2*a + h2ez2m1*(2Nbth2-Nut)

        u = chop(ez*u + hezα*Nut + h2ezβ*(Nath2+Nbth2) + hezγ*N(c,t+h),tol)
        push!(glp,u)
        yield()
        t+=h
    end

    u,t
end

## Multivariate

function BDF4(B::Vector,op::Operator,bcs::Vector,uin::MultivariateFun,h::Real,m::Integer,glp)
    nt=size(uin,2)
    d=domain(uin)
    SBE   = discretize([B;I-h*op],d,nt)            # backward euler for first 2 time steps
    SBDF2 = discretize([B;I-2.0/3.0*h*op],d,nt)    # BDF formula for subsequent itme steps
    SBDF3 = discretize([B;I-6.0/11.0*h*op],d,nt)    # BDF formula for subsequent itme steps
    SBDF4 = discretize([B;I-12.0/25.0*h*op],d,nt)    # BDF formula for subsequent itme steps

    u1=uin
    u2=SBE\[bcs;u1]
    push!(glp,pad(u2,80,80))
    yield()
    u3=SBDF2\[bcs;1/3.0*(4u2-u1)]
    push!(glp,pad(u3,80,80))#updates window
    yield()
    u4=SBDF3\[bcs;1/11.0*(18u3-9u2+2u1)]
    push!(glp,pad(u4,80,80))#updates window
    yield()

    for k=1:m
        u4,u3,u2,u1  = SBDF4\[bcs;1/25.0*(48u4-36u3+16u2-3u1)],u4,u3,u2

        push!(glp,pad(u4,80,80))#updates window
        yield()
    end
    u4
end

BDF4(B::Vector,op::Operator,uin::MultivariateFun,h::Real,m::Integer,glp)=BDF4(B,op,zeros(length(B)),uin,h,m,glp)

function BDF22(B::Vector,op::Operator,bcs::Vector,uin::Tuple{MultivariateFun,MultivariateFun},h::Real,m::Integer,glp)
    nt=size(uin[1],2)
    SBE  = discretize([B;I-h^2*op],domain(uin[1]),nt)            # backward euler for first 2 time steps
    SBDF = discretize([B;I-4.0/9.0*h^2*op],domain(uin[1]),nt)    # BDF formula for subsequent itme steps

    u1,u2=uin
    u3 =SBE\[bcs;2u2-u1]
    push!(glp,pad(u3,80,80))
    yield()
    u4 =SBE\[bcs;2u3-u2]
    push!(glp,pad(u4,80,80))
    yield()

    for k=1:m
        u4,u3,u2,u1  = SBDF\[bcs;1/9.0*(24u4-22u3+8u2-u1)],u4,u3,u2

        push!(glp,pad(u4,80,80)) #updates window
        yield()
    end
    u4
end

function BDF22(B::Vector,op::Operator,g::Function,bcs::Vector,uin::Tuple{MultivariateFun,MultivariateFun},h::Real,m::Integer,glp)
    nt=size(uin[1],2)
    SBE  = discretize([B;I-h^2*op],domain(uin[1]),nt)            # backward euler for first 2 time steps
    SBDF = discretize([B;I-4.0/9.0*h^2*op],domain(uin[1]),nt)    # BDF formula for subsequent itme steps

    u1,u2=uin
    u3 =SBE\[bcs;2u2-u1]
    push!(glp,pad(u3,80,80))
    yield()
    u4,u3,u2,u1=u3,u2,u1,u1
    for k=1:m
        u4,u3,u2,u1 =2u4 - u3 +h^2*g(u4),u4,u3,u2
        u4,u3,u2,u1  = SBDF\[bcs;1/9.0*(24u4-22u3+8u2-u1)],u4,u3,u2
        push!(glp,pad(u4,80,80))
        yield()
    end
    u4
end

BDF22(B::Vector,op::Operator,uin::MultivariateFun,h::Real,m::Integer,glp)=BDF22(B,op,zeros(length(B)),(uin,uin),h,m,glp)
BDF22(B::Vector,op::Operator,g::Function,uin::MultivariateFun,h::Real,m::Integer,glp)=BDF22(B,op,g,zeros(length(B)),(uin,uin),h,m,glp)

## GLPlot routines

#u_t = op*u
function timeevolution(B::Vector,op,bcs::Vector,uin::MultivariateFun,h::Real,m::Integer,glp)
    require("GLPlot")
    setplotter("GLPlot")
    nt=size(uin,2)
    d=domain(uin)
    SBE   = discretize([B;I-h*op],d,nt)            # backward euler for first 2 time steps
    SBDF2 = discretize([B;I-2.0/3.0*h*op],d,nt)    # BDF formula for subsequent itme steps
    SBDF3 = discretize([B;I-6.0/11.0*h*op],d,nt)    # BDF formula for subsequent itme steps
    SBDF4 = discretize([B;I-12.0/25.0*h*op],d,nt)    # BDF formula for subsequent itme steps

    u1=chop(uin,1000eps())
    u2=chop(SBE\[bcs;u1],1000eps())
    plot(pad(u2,80,80),glp...)#updates window
    u3=chop(SBDF2\[bcs;1/3.0*(4u2-u1)],1000eps())
    plot(pad(u3,80,80),glp...)#updates window
    u4=chop(SBDF3\[bcs;1/11.0*(18u3-9u2+2u1)],1000eps())
    plot(pad(u4,80,80),glp...)#updates window

    for k=1:m
        u4,u3,u2,u1  = chop(SBDF4\[bcs;1/25.0*(48u4-36u3+16u2-3u1)],1000eps()),u4,u3,u2

        plot(pad(u4,80,80),glp...)#updates window
    end
    u4
end

function timeevolution(B::Vector,op,bcs::Vector,uin::MultivariateFun,h::Real,m=5000)
    require("GLPlot")
    setplotter("GLPlot")
    timeevolution(B,op,bcs,uin,h,m,plot(pad(uin,80,80)))
end

timeevolution(B::Vector,op,uin::MultivariateFun,h::Real,dat...)=timeevolution(B,op,zeros(length(B)),uin,h,dat...)

timeevolution(B::Vector,op,bcs::Vector,uin::Fun,dat...)=timeevolution(B,op,bcs,ProductFun(uin),dat...)
timeevolution(B::Vector,op,uin::Fun,dat...)=timeevolution(B,op,ProductFun(uin),dat...)

timeevolution(B::Operator,dat...)=timeevolution([B],dat...)

#u_tt = op*u
function timeevolution2(B::Vector,op,bcs::Vector,uin::Tuple{MultivariateFun,MultivariateFun},h::Real,m::Integer,glp)
    require("GLPlot")
    setplotter("GLPlot")
    nt=size(uin[1],2)
    SBE  = discretize([B;I-h^2*op],domain(uin[1]),nt)            # backward euler for first 2 time steps
    SBDF = discretize([B;I-4.0/9.0*h^2*op],domain(uin[1]),nt)    # BDF formula for subsequent itme steps

    u1,u2=chop(uin[1],1000eps()),chop(uin[2],1000eps())
    u3 =chop(SBE\[bcs;2u2-u1],1000eps())
    u4 =chop(SBE\[bcs;2u3-u2],1000eps())

    for k=1:m
        u4,u3,u2,u1  = chop(SBDF\[bcs;1/9.0*(24u4-22u3+8u2-u1)],1000eps()),u4,u3,u2

        plot(pad(u4,80,80),glp...)#updates window
    end
    u4
end

#u_tt = op*u
function timeevolution2(B::Vector,op,g::Function,bcs::Vector,uin::Tuple{MultivariateFun,MultivariateFun},h::Real,m::Integer,glp)
    require("GLPlot")
    setplotter("GLPlot")
    nt=size(uin[1],2)
    SBE  = discretize([B;I-h^2*op],domain(uin[1]),nt)            # backward euler for first 2 time steps
    SBDF = discretize([B;I-4.0/9.0*h^2*op],domain(uin[1]),nt)    # BDF formula for subsequent itme steps

    u1,u2=chop(uin[1],1000eps()),chop(uin[2],1000eps())
    u3 =chop(SBE\[bcs;2u2-u1],1000eps())
    u4 =chop(2u3 - u2 +h^2*g(u3),1000eps())
    u4,u3,u2,u1  = chop(SBDF\[bcs;1/9.0*(24u4-22u3+8u2-u1)],1000eps()),u4,u3,u2
    plot(pad(u4,80,80),glp...)#updates window

    for k=1:m
        u4,u3,u2,u1  = chop(2u4 - u3 +h^2*g(u4),1000eps()),u4,u3,u2
        u4,u3,u2,u1  = chop(SBDF\[bcs;1/9.0*(24u4-22u3+8u2-u1)],1000eps()),u4,u3,u2
        plot(pad(u4,80,80),glp...)#updates window
    end
    u4
end

function timeevolution2(B::Vector,op,uin::Tuple{MultivariateFun,MultivariateFun},bcs::Vector,h::Real,m=5000)
    require("GLPlot")
    setplotter("GLPlot")
    timeevolution2(B,op,bcs,uin,h,m,plot(pad(uin[end],80,80)))
end

function timeevolution2(B::Vector,op,g::Function,uin::Tuple{MultivariateFun,MultivariateFun},bcs::Vector,h::Real,m=5000)
    require("GLPlot")
    setplotter("GLPlot")
    timeevolution2(B,op,g,bcs,uin,h,m,plot(pad(uin[end],80,80)))
end

timeevolution2(B::Vector,op,uin::Tuple{MultivariateFun,MultivariateFun},h::Real,dat...)=timeevolution2(B,op,uin,zeros(length(B)),h,dat...)
timeevolution2(B::Vector,op,uin::MultivariateFun,dat...)=timeevolution2(B,op,(uin,uin),dat...)
timeevolution2(B::Operator,dat...)=timeevolution2([B],dat...)

timeevolution2(B::Vector,op,g::Function,uin::Tuple{MultivariateFun,MultivariateFun},h::Real,dat...)=timeevolution2(B,op,g,uin,zeros(length(B)),h,dat...)
timeevolution2(B::Vector,op,g::Function,uin::MultivariateFun,dat...)=timeevolution2(B,op,g,(uin,uin),dat...)

timeevolution2(B::Vector,op,g::Function,uin::Tuple{Fun,Fun},dat...)=timeevolution2(B,op,g,(ProductFun(uin[1]),ProductFun(uin[2])),dat...)
timeevolution2(B::Vector,op,g::Function,uin::Fun,dat...)=timeevolution2(B,op,g,ProductFun(uin),dat...)
timeevolution2(B::Vector,op,uin::Fun,dat...)=timeevolution2(B,op,ProductFun(uin),dat...)

timeevolution(o::Integer,dat...)=o==2?timeevolution2(dat...):timeevolution(dat...)


function Polygon(
        points::AbstractVector;
        boundingbox = Rect2D(points),
        isconvex = isconvex(points),
        ishole = false
    )
    ps = convert(AbstractVector{Point}, points)
    NumType = eltype(eltype(ps))
    Polygon{NumType, typeof(ps)}(ps, boundingbox, isconvex, ishole)
end

iterate(polygons::MultiPolygon) = iterate(polygons.polygons)
iterate(polygons::MultiPolygon, state) = iterate(polygons.polygons, state)

vertices(poly::Polygon) = poly.points
boundingbox(poly::Polygon) = poly.boundingbox
ishole(poly::Polygon) = poly.ishole
isconvex(poly::Polygon) = poly.isconvex

"""
    in(point::AbstractPoint, poly::Polygon)

# An implementation of Hormann-Agathos (2001) Point in Polygon algorithm
# See: http://www.sciencedirect.com/science/article/pii/S0925772101000128
# Code segment adapted from PolygonClipping.jl
"""
function in(point::AbstractPoint, poly::Polygon)
    (point in boundingbox(poly)) || return false
    c = false
    detq(q1, q2, point) = (q1[1]-point[1])*(q2[2]-point[2])-(q2[1]-point[1])*(q1[2]-point[2])
    for (q1, q2) in Partition{2}(poly)
        if q1 == point
            @warn("point on polygon vertex - returning false")
            return false
        end
        if q2[2] == point[2]
            if q2[1] == point[1]
                @warn("point on polygon vertex - returning false")
                return false
            elseif (q1[2] == point[2]) && ((q2[1] > x) == (q1[1] < point[1]))
                @warn("point on edge - returning false")
                return false
            end
        end
        if (q1[2] < point[2]) != (q2[2] < point[2]) # crossing
            if q1[1] >= point[1]
                if q2[1] > point[1]
                    c = !c
                elseif ((detq(q1,q2,point) > 0) == (q2[2] > q1[2])) # right crossing
                    c = !c
                end
            elseif q2[1] > point[1]
                if ((detq(q1,q2,point) > 0) == (q2[2] > q1[2])) # right crossing
                    c = !c
                end
            end
        end
    end
    return c
end

function in(point::AbstractPoint, polygons::MultiPolygon)
    point in boundingbox(polygons) || return false
    for polygon in polygons
        if ishole(polygon)
            point in polygon && return false
        else
            point in polygon && return true
        end
    end
    return false
end

function isconvex(points::AbstractVector{<: AbstractPoint})
    length(points) < 4 && return true
    sign = false
    for (i, (a, b, c)) in enumerate(Partition{3}(points, true))
        Δ1 = c .- b
        Δ2 = a .- b
        det = cross(Δ1, Δ2)
        if i == 1
            sign = det > 0
        elseif (sign != (det > 0))
            return false
        end
    end
    return true
end


using Seismic
using Base.Test

# test of Conjugate Gradients
nd = 100
nm = 300
L = 10*randn(nd,nm)
m = 10.0*rand(nm)
d = zeros(nd)
d = L*m
d = d + 2.0*randn(nd)

println(norm(d))
println(norm(m))

m1,cost = ConjugateGradients(d,[MatrixMultiplyOp],[Dict(:matrix=>L)],Niter=20,mu=0.2)

ext = Seismic.Extent(size(d,1),size(d,2),1,1,1,
0,0,0,1,1,
1,1,1,1,1,
"","","","","",
"","","","","",
"")
h = Header[]
for ix = 1:size(d,2)
		push!(h,Seismic.InitSeisHeader())
		h[ix] = Seismic.InitSeisHeader()
		h[ix].tracenum = ix
		h[ix].n1 = size(d,1)
		h[ix].d1 = 1
end

SeisWrite("tmp_d.seis",d,h,ext)

ConjugateGradients("tmp_m.seis","tmp_d.seis",[MatrixMultiplyOp],[Dict(:matrix=>L)],"tmp_cost.txt",Niter=20,mu=0.2)
m2,ext = SeisRead("tmp_m.seis")

# test that quality factor between disk and memory based CG
# is greater than 50 Decibels
quality_factor = 10*log10(norm(m1[:],2)/norm(m2[:]-m1[:],2))
println("Quality factor = ",quality_factor)
@test quality_factor > 10.


function edit_distance{T}(a, b, submat::AbstractSubstitutionMatrix{T}, insertion_cost::T, deletion_cost::T)
    m = length(a)
    n = length(b)
    trace = Matrix{Trace}(m + 1, n + 1)
    D = Vector{T}(m + 1)
    D[1] = T(0)
    for i in 1:m
        D[i+1] = i * deletion_cost
        trace[i+1,1] = TRACE_INSERT
    end
    for j in 1:n
        d_diag = D[1]
        D[1] = j * insertion_cost
        trace[1,j+1] = TRACE_DELETE
        for i in 1:m
            del = D[i]   + deletion_cost
            ins = D[i+1] + insertion_cost
            mat = d_diag + submat[a[i],b[j]]
            # find the best score and its trace
            best = min(del, ins, mat)
            t = TRACE_NONE
            del == best && (t |= TRACE_DELETE)
            ins == best && (t |= TRACE_INSERT)
            mat == best && (t |= TRACE_MATCH)
            d_diag = D[i+1]
            D[i+1] = best
            trace[i+1,j+1] = t
        end
    end
    return D[m+1], trace
end

function edit_traceback(a, b, trace, endpos)
    # gap/character counts (reversed order)
    counts_a = [0, 0]
    counts_b = [0, 0]
    i, j = endpos
    while i ≥ 1 && j ≥ 2
        t = trace[i+1,j+1]
        if t & TRACE_MATCH > 0
            @match
        elseif t & TRACE_DELETE > 0
            @gap a
        elseif t & TRACE_INSERT > 0
            @gap b
        end
        # do not come here
        @assert false
    end
    while j ≥ 1 @gap a end
    while i ≥ 1 @gap b end
    reverse!(counts_a)
    reverse!(counts_b)
    return GappedSequence(a, counts_a), GappedSequence(b, counts_b)
end