text:
using DifferentiableStateSpaceModels, Symbolics, LinearAlgebra, Zygote, Test
using DifferentiableStateSpaceModels.Examples
using DifferentiableStateSpaceModels: order_vector_by_symbols,
                                      fill_array_by_symbol_dispatch,
                                      all_fields_equal

# # # Use while testing internals
# m = @include_example_module(Examples.rbc_observables)
# # Basic Steady State
# p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
# p_d = (α=0.5, β=0.95)
# c = SolverCache(m, Val(1), p_d)
# sol = generate_perturbation(m, p_d, p_f, Val(1); cache = c) # manually passing in order
# generate_perturbation_derivatives!(m, p_d, p_f, c)
# ex = DifferentiableStateSpaceModels.exfiltrated

@testset "Construction" begin
    m = @include_example_module(Examples.rbc_observables)

    # bookkeeping tests
    @test m.n_y == 2
    @test m.n_x == 2
    @test m.n_p == 6
    @test m.n_ϵ == 1
    @test m.n_z == 2
    @test m.η == reshape([0; -1], 2, m.n_ϵ)
    @test m.Q == [1.0 0.0 0.0 0.0; 0.0 0.0 1.0 0.0]

    # function tests (steady state)
    # Basic Steady State
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)
    sol = generate_perturbation(m, p_d, p_f, Val(1)) # manually passing in first derivative, but not really required
    sol = generate_perturbation(m, p_d, p_f) # Default is first-order
    @inferred generate_perturbation(m, p_d, p_f)
    @inferred generate_perturbation(m, p_d, p_f, Val(1))
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test sol.retcode == :Success

    # Call all variables differentiated
    sol = generate_perturbation(m, merge(p_d, p_f), nothing)
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]

    # With a prebuilt cache
    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache=c)
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @inferred generate_perturbation(m, p_d, p_f; cache=c)
end

@testset "Function Evaluation" begin
    m = @include_example_module(Examples.rbc_observables)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)
    p_d_symbols = collect(Symbol.(keys(p_d)))
    c = SolverCache(m, Val(1), p_d)

    # Create parameter vector in the same ordering the internal algorithms would
    p = order_vector_by_symbols(merge(p_d, p_f), m.mod.p_symbols)

    y = zeros(m.n_y)
    x = zeros(m.n_x)

    m.mod.ȳ!(y, p)
    m.mod.x̄!(x, p)
    @test y ≈ [5.936252888048733, 6.884057971014498]
    @test x ≈ [47.39025414828825, 0.0]

    m.mod.H_yp!(c.H_yp, y, x, p)
    @test c.H_yp ≈ [0.028377570562199098 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]

    m.mod.H_y!(c.H_y, y, x, p)
    @test c.H_y ≈ [-0.0283775705621991 0.0; 1.0 -1.0; 0.0 1.0; 0.0 0.0]

    m.mod.H_xp!(c.H_xp, y, x, p)
    @test c.H_xp ≈ [0.00012263591151906127 -0.011623494029190608
                    1.0 0.0
                    0.0 0.0
                    0.0 1.0]

    m.mod.H_x!(c.H_x, y, x, p)
    @test c.H_x ≈ [0.0 0.0
                   -0.98 0.0
                   -0.07263157894736837 -6.884057971014498
                   0.0 -0.2]

    m.mod.Ψ!(c.Ψ, y, x, p)
    @test c.Ψ[1] ≈
          [-0.009560768753410337 0.0 0.0 0.0 -2.0658808482697935e-5 0.0019580523687917364 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.009560768753410338 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           -2.0658808482697935e-5 0.0 0.0 0.0 -3.881681383327978e-6 0.00012263591151906127 0.0 0.0
           0.0019580523687917364 0.0 0.0 0.0 0.00012263591151906127 -0.011623494029190608 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0]
    @test c.Ψ[2] ≈ zeros(8, 8)
    @test c.Ψ[3] ≈ [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0007663134567721225 -0.07263157894736837
           0.0 0.0 0.0 0.0 0.0 0.0 -0.07263157894736837 -6.884057971014498]
    @test c.Ψ[4] ≈ zeros(8, 8)

    m.mod.Γ!(c.Γ, p)
    @test c.Γ ≈ [0.01]

    m.mod.Ω!(c.Ω, p)
    @test c.Ω ≈ [0.01, 0.01]

    # The derivative ones dispatch by the derivative symbol
    fill_array_by_symbol_dispatch(m.mod.H_x_p!, c.H_x_p, p_d_symbols, y, x, p)
    @test c.H_x_p ≈ [[0.0 0.0
            0.0 0.0
            -0.4255060477077458 -26.561563542978472
            0.0 0.0], [0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]]
    fill_array_by_symbol_dispatch(m.mod.H_yp_p!, c.H_yp_p, p_d_symbols, y, x, p)
    @test c.H_yp_p ≈ [[0.011471086498795562 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0],
           [0.029871126907577997 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]]

    fill_array_by_symbol_dispatch(m.mod.H_y_p!, c.H_y_p, p_d_symbols, y, x, p)
    @test c.H_y_p ≈
          [[0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0], [0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]]

    fill_array_by_symbol_dispatch(m.mod.H_xp_p!, c.H_xp_p, p_d_symbols, y, x, p)
    @test c.H_xp_p ≈ [[0.000473180436623283 -0.06809527035753198
            0.0 0.0
            0.0 0.0
            0.0 0.0], [0.00012909043317795924 -0.01223525687283222
                       0.0 0.0
                       0.0 0.0
                       0.0 0.0]]

    fill_array_by_symbol_dispatch(m.mod.Γ_p!, c.Γ_p, p_d_symbols, p)

    @test c.Γ_p ≈ [[0.0], [0.0]]

    fill_array_by_symbol_dispatch(m.mod.H_p!, c.H_p, p_d_symbols, y, x, p)

    @test c.H_p ≈ [[-0.06809527035753199, 0.0, -26.561563542978472, 0.0],
           [-0.1773225633743801, 0.0, 0.0, 0.0]]

    fill_array_by_symbol_dispatch(m.mod.Ω_p!, c.Ω_p, p_d_symbols, p)
    @test c.Ω_p ≈ [[0.0, 0.0], [0.0, 0.0]]
end

@testset "Evaluation into cache and check solution" begin
    m = @include_example_module(Examples.rbc_observables)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)
    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)
    # Create parameter vector in the same ordering the internal algorithms would

    @test c.y ≈ [5.936252888048733, 6.884057971014498]
    @test c.x ≈ [47.39025414828825, 0.0]
    @test c.H_yp ≈ [0.028377570562199098 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
    @test c.H_y ≈ [-0.0283775705621991 0.0; 1.0 -1.0; 0.0 1.0; 0.0 0.0]
    @test c.H_xp ≈ [0.00012263591151906127 -0.011623494029190608
                    1.0 0.0
                    0.0 0.0
                    0.0 1.0]
    @test c.H_x ≈ [0.0 0.0
                   -0.98 0.0
                   -0.07263157894736837 -6.884057971014498
                   0.0 -0.2]
    @test c.Ψ[1] ≈
          [-0.009560768753410337 0.0 0.0 0.0 -2.0658808482697935e-5 0.0019580523687917364 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.009560768753410338 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           -2.0658808482697935e-5 0.0 0.0 0.0 -3.881681383327978e-6 0.00012263591151906127 0.0 0.0
           0.0019580523687917364 0.0 0.0 0.0 0.00012263591151906127 -0.011623494029190608 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0]
    @test c.Ψ[2] ≈ zeros(8, 8)
    @test c.Ψ[3] ≈ [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
           0.0 0.0 0.0 0.0 0.0 0.0 0.0007663134567721225 -0.07263157894736837
           0.0 0.0 0.0 0.0 0.0 0.0 -0.07263157894736837 -6.884057971014498]
    @test c.Ψ[4] ≈ zeros(8, 8)
    @test c.Γ ≈ [0.01]
    @test c.Ω ≈ [0.01, 0.01]

    # solution tests
    @test c.y ≈ [5.936252888048733, 6.884057971014498]
    @test c.x ≈ [47.39025414828824, 0.0]
    @test c.g_x ≈ [
        0.0957964300241661 0.6746869652586178
        0.07263157894736878 6.884057971014507
    ]
    @test c.h_x ≈ [0.9568351489232028 6.209371005755889; -1.5076865909646354e-18 0.2]
    @test c.Σ ≈ [1e-4]
    @test c.Ω ≈ [0.01, 0.01]
    @test c.η == reshape([0; -1], 2, m.n_ϵ)
    @test c.B ≈ [0.0; -0.01]
    @test c.Q ≈ [1.0 0.0 0.0 0.0; 0.0 0.0 1.0 0.0]
    @test c.C_1 ≈ [0.0957964300241661 0.6746869652585828; 1.0 0.0]
    @test Array(c.V) ≈ [
        0.07005411173180227 0.00015997603451513485
        0.00015997603451513485 0.00010416666666666667
    ]

    # Check the solution type matches these all
    fields_to_compare = (:y, :x, :g_x, :B, :Q, :η, :Γ)
    @test all_fields_equal(c, sol, fields_to_compare)
    @test c.h_x ≈ sol.A
    @test c.C_1 ≈ sol.C
    @test c.V.L ≈ sol.x_ergodic.C.L # choleskys
    @test c.V.U ≈ sol.x_ergodic.C.U
    @test c.Ω ≈ sol.D.σ
end

@testset "Evaluate Derivatives into cache" begin
    m = @include_example_module(Examples.rbc_observables)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)
    p_d_symbols = collect(Symbol.(keys(p_d)))
    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)
    @test :Success == generate_perturbation_derivatives!(m, p_d, p_f, c) # Solves and fills the cache
    @inferred generate_perturbation_derivatives!(m, p_d, p_f, c)

    @test c.H_x_p ≈ [[0.0 0.0
            0.0 0.0
            -0.4255060477077458 -26.561563542978472
            0.0 0.0], [0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]]
    @test c.H_yp_p ≈ [[0.011471086498795562 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0],
           [0.029871126907577997 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]]
    @test c.H_y_p ≈
          [[0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0], [0.0 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]]
    @test c.H_xp_p ≈ [[0.000473180436623283 -0.06809527035753198
            0.0 0.0
            0.0 0.0
            0.0 0.0], [0.00012909043317795924 -0.01223525687283222
                       0.0 0.0
                       0.0 0.0
                       0.0 0.0]]
    @test c.Γ_p ≈ [[0.0], [0.0]]
    @test c.H_p ≈ [[-0.06809527035753199, 0.0, -26.561563542978472, 0.0],
           [-0.1773225633743801, 0.0, 0.0, 0.0]]
    @test c.Ω_p ≈ [[0.0, 0.0], [0.0, 0.0]]

    # solution
    @test hcat(c.y_p...) ≈ [
        55.78596896689701 76.10141579073955
        66.89124302798608 105.01995379122064
    ]
    @test hcat(c.x_p...) ≈ [555.2637030544529 1445.9269000240533; 0.0 0.0]
    @test c.g_x_p ≈ [
        [
            -0.12465264193058262 5.596211904442805
            -1.2823781479976832e-15 66.89124302798608
        ],
        [
            -1.6946742377792863 -0.8343618226192915
            -1.1080332409972313 105.01995379122064
        ],
    ]
    @test c.h_x_p ≈ [
        [0.12465264193058134 61.29503112354326; 0.0 0.0],
        [0.586640996782055 105.85431561383992; 0.0 0.0],
    ]
    @test c.Σ_p ≈ [[0.0], [0.0]]
    @test hcat(c.Ω_p...) ≈ [0.0 0.0; 0.0 0.0]
    @test c.B_p ≈ [[0.0; 0.0], [0.0; 0.0]]
    @test c.C_1_p ≈ [
        [-0.12465264193057919 5.596211904442171; 0.0 0.0],
        [-1.6946742377792825 -0.8343618226202246; 0.0 0.0],
    ]
    @test c.V_p ≈ [
        [1.584528257999749 0.0015841155991973127; 0.0015841155991973127 0.0],
        [3.336643488330957 0.002750404724942799; 0.002750404724942799 0.0],
    ]
end

@testset "Steady State with Initial Conditions" begin
    m = @include_example_module(Examples.rbc_solve_steady_state)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)

    sol = generate_perturbation(m, p_d, p_f)
    @inferred generate_perturbation(m, p_d, p_f)
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test sol.retcode == :Success
end

@testset "Steady State with Initial Conditions" begin
    m = @include_example_module(Examples.rbc_solve_steady_state)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)

    # Starting at steady state in optimizer, so should be 0 or 1 iteration
    settings = PerturbationSolverSettings(; print_level=0)
    sol = generate_perturbation(m, p_d, p_f; settings)
    @inferred generate_perturbation(m, p_d, p_f)
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test sol.retcode == :Success

    # At an initial condition outside of the steady state
    m = @include_example_module(Examples.rbc_solve_steady_state_different_iv)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)

    sol = generate_perturbation(m, p_d, p_f; settings)
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test sol.retcode == :Success

    # Change the order of the parameters/etc.
    p_f = (σ=0.01, δ=0.02, Ω_1=0.01)
    p_d = (α=0.5, ρ=0.2, β=0.95)
    sol = generate_perturbation(m, p_d, p_f; settings)  # no p_f and rearranged
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test sol.retcode == :Success

    m = @include_example_module(Examples.rbc_solve_steady_state_different_iv)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)

    settings = PerturbationSolverSettings(; print_level=0, nlsolve_iterations=2)  # simulate failure by insufficient iterationrs
    sol = generate_perturbation(m, p_d, p_f; settings)
    @test sol.retcode == :SteadyStateFailure
end

@testset "Construction and solution with no Omega" begin
    m = @include_example_module(Examples.rbc)
    p_f = (ρ=0.2, δ=0.02, σ=0.01)
    p_d = (α=0.5, β=0.95)
    sol = generate_perturbation(m, p_d, p_f)
    @inferred generate_perturbation(m, p_d, p_f)
    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test sol.retcode == :Success

    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache =c)
    @test :Success == generate_perturbation_derivatives!(m, p_d, p_f, c)  # Solves and fills the cache
    @inferred generate_perturbation(m, p_d, p_f; cache =c)
    @inferred generate_perturbation_derivatives!(m, p_d, p_f, c)

    @inferred generate_perturbation(m, p_d, p_f; cache =c)
    @test sol.retcode == :Success

    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828824, 0.0]
    @test hcat(c.y_p...) ≈ [
        55.78596896689701 76.10141579073955
        66.89124302798608 105.01995379122064
    ]
    @test hcat(c.x_p...) ≈ [555.2637030544529 1445.9269000240533; 0.0 0.0]
    @test c.g_x ≈ [
        0.0957964300241661 0.6746869652586178
        0.07263157894736878 6.884057971014507
    ]
    @test c.h_x ≈ [0.9568351489232028 6.209371005755889; -1.5076865909646354e-18 0.2]
    @test c.g_x_p ≈ [
        [
            -0.12465264193058262 5.596211904442805
            -1.2823781479976832e-15 66.89124302798608
        ],
        [
            -1.6946742377792863 -0.8343618226192915
            -1.1080332409972313 105.01995379122064
        ],
    ]
    @test c.h_x_p ≈ [
        [0.12465264193058134 61.29503112354326; 0.0 0.0],
        [0.586640996782055 105.85431561383992; 0.0 0.0],
    ]
    @test c.Σ ≈ [1e-4]
    @test c.Σ_p ≈ [[0.0], [0.0]]
end

@testset "Evaluate rbc_observables_separate_variance derivatives into cache" begin
    m = @include_example_module(Examples.rbc_observables_separate_variance)

    p_f = (δ=0.02, σ=0.01, Ω_2=0.012)
    p_d = (α=0.5, β=0.95, ρ=0.2, Ω_1=0.011)
    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)
    @test :Success == generate_perturbation_derivatives!(m, p_d, p_f, c)

    @test c.y ≈ [5.936252888048733, 6.884057971014498]
    @test c.x ≈ [47.39025414828824, 0.0]
    @test hcat(c.y_p...) ≈ [
        55.78596896689701 76.10141579073955 0.0 0.0
        66.89124302798608 105.01995379122064 0.0 0.0
    ]
    @test hcat(c.x_p...) ≈ [555.2637030544529 1445.9269000240533 0.0 0.0; 0.0 0.0 0.0 0.0]
    @test c.g_x ≈ [
        0.0957964300241661 0.6746869652586178
        0.07263157894736878 6.884057971014507
    ]
    @test c.h_x ≈ [0.9568351489232028 6.209371005755889; -1.5076865909646354e-18 0.2]
    @test c.g_x_p ≈ [
        [
            -0.12465264193058262 5.596211904442805
            -1.2823781479976832e-15 66.89124302798608
        ],
        [
            -1.6946742377792863 -0.8343618226192915
            -1.1080332409972313 105.01995379122064
        ],
        [1.3921362894665956e-18 0.29450084693461975; 0.0 0.0],
        zeros(2, 2),
    ]
    @test c.h_x_p ≈ [
        [0.12465264193058134 61.29503112354326; 0.0 0.0],
        [0.586640996782055 105.85431561383992; 0.0 0.0],
        [-1.3921362894665956e-18 -0.29450084693461975; 4.7271045362244914e-18 1.0],
        zeros(2, 2),
    ]
    @test c.Σ ≈ [1e-4]
    @test c.Σ_p ≈ [[0.0], [0.0], [0.0], [0.0]]

    @test c.Ω ≈ [0.011, 0.012]
    @test hcat(c.Ω_p...) ≈ [0.0 0.0 0.0 1.0; 0.0 0.0 0.0 0.0]

    @test c.η == reshape([0; -1], 2, m.n_ϵ)
    @test c.Q ≈ [1.0 0.0 0.0 0.0; 0.0 0.0 1.0 0.0]

    # Should be identical to c
    @test sol.n_y == 2
    @test sol.n_x == 2
    @test sol.n_ϵ == 1
    @test sol.n_z == 2
    @test c.y ≈ sol.y
    @test c.x ≈ sol.x
    @test c.g_x ≈ sol.g_x
    @test c.h_x ≈ sol.A
    @test c.B ≈ sol.B
    @test c.Ω ≈ sol.D.σ
    @test c.Q ≈ sol.Q
    @test c.η ≈ sol.η
    @test sol.retcode == :Success
end

@testset "Construction and solution no p_f" begin
    m = @include_example_module(Examples.rbc)
    p_f = nothing
    p_d = (α=0.5, β=0.95, ρ=0.2, δ=0.02, σ=0.01)
    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)
    generate_perturbation_derivatives!(m, p_d, p_f, c)
    @inferred generate_perturbation(m, p_d, p_f; cache = c)
    @inferred generate_perturbation_derivatives!(m, p_d, p_f, c)

    @test c.y ≈ [5.936252888048733, 6.884057971014498]
    @test c.x ≈ [47.39025414828825, 0.0]
    @test c.H_yp ≈ [0.028377570562199098 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
    @test c.H_y ≈ [-0.0283775705621991 0.0; 1.0 -1.0; 0.0 1.0; 0.0 0.0]
    @test c.H_xp ≈ [
        0.00012263591151906127 -0.011623494029190608
        1.0 0.0
        0.0 0.0
        0.0 1.0
    ]
    @test c.H_x ≈ [
        0.0 0.0
        -0.98 0.0
        -0.07263157894736837 -6.884057971014498
        0.0 -0.2
    ]
    @test c.Ψ[1] ≈ [
        -0.009560768753410337 0.0 0.0 0.0 -2.0658808482697935e-5 0.0019580523687917364 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.009560768753410338 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        -2.0658808482697935e-5 0.0 0.0 0.0 -3.881681383327978e-6 0.00012263591151906127 0.0 0.0
        0.0019580523687917364 0.0 0.0 0.0 0.00012263591151906127 -0.011623494029190608 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    ]
    @test c.Ψ[2] ≈ zeros(8, 8)
    @test c.Ψ[3] ≈ [
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0007663134567721225 -0.07263157894736837
        0.0 0.0 0.0 0.0 0.0 0.0 -0.07263157894736837 -6.884057971014498
    ]
    @test c.Ψ[4] ≈ zeros(8, 8)
    @test c.Γ ≈ [0.01]
    @test c.Γ_p ≈ [[0.0], [0.0], [0.0], [0.0], [1.0]]
    @test hcat(c.H_p...) ≈ [
        -0.06809527035753199 -0.1773225633743801 0.0 0.16003361344537806 0.0
        0.0 0.0 0.0 47.39025414828824 0.0
        -26.561563542978472 0.0 0.0 0.0 0.0
        0.0 0.0 -0.0 0.0 0.0
    ]

    # solution tests
    @test c.y ≈ [5.936252888048733, 6.884057971014498]
    @test c.x ≈ [47.39025414828824, 0.0]
    @test hcat(c.y_p...) ≈ [
        55.78596896689701 76.10141579073955 0.0 -116.07178189943077 0.0
        66.89124302798608 105.01995379122064 0.0 -94.78050829657676 0.0
    ]
    @test hcat(c.x_p...) ≈ [
        555.2637030544529 1445.9269000240533 0.0 -1304.9490272717098 0.0
        0.0 0.0 0.0 0.0 0.0
    ]
    @test c.g_x ≈ [
        0.0957964300241661 0.6746869652586178
        0.07263157894736878 6.884057971014507
    ]
    @test c.h_x ≈ [0.9568351489232028 6.209371005755889; -1.5076865909646354e-18 0.2]
    @test c.g_x_p ≈ [
        [
            -0.12465264193058262 5.596211904442805
            -1.2823781479976832e-15 66.89124302798608
        ],
        [
            -1.6946742377792863 -0.8343618226192915
            -1.1080332409972313 105.01995379122064
        ],
        [4.640454298222595e-19 0.2945008469346586; 0.0 0.0],
        [
            0.6277268890968761 -5.0369653468355455
            1.0000000000000024 -94.78050829657676
        ],
        [0.0 0.0; 0.0 0.0],
    ]
    @test c.h_x_p ≈ [
        [0.12465264193058134 61.29503112354326; 0.0 0.0],
        [0.586640996782055 105.85431561383992; 0.0 0.0],
        [
            -4.640454298222595e-19 -0.2945008469346586
            1.5757015120748295e-18 1.0
        ],
        [-0.6277268890968737 -89.74354294974118; 0.0 0.0],
        [0.0 0.0; 0.0 0.0],
    ]
    @test c.Σ ≈ [1e-4]
    @test c.Σ_p ≈ [[0.0], [0.0], [0.0], [0.0], [0.02]]
end

@testset "Construction and solution, multiple shocks" begin
    m = @include_example_module(Examples.rbc_multiple_shocks)
    p_f = nothing
    p_d = (α=0.5, β=0.95, ρ=0.2, δ=0.02, σ=0.01)
    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)
    generate_perturbation_derivatives!(m, p_d, p_f, c)
    @inferred generate_perturbation(m, p_d, p_f; cache = c)
    @inferred generate_perturbation_derivatives!(m, p_d, p_f, c)

    @test sol.y ≈ [5.936252888048733, 6.884057971014498]
    @test sol.x ≈ [47.39025414828825, 0.0]
    @test c.H_yp ≈ [0.028377570562199098 0.0; 0.0 0.0; 0.0 0.0; 0.0 0.0]
    @test c.H_y ≈ [-0.0283775705621991 0.0; 1.0 -1.0; 0.0 1.0; 0.0 0.0]
    @test c.H_xp ≈ [
        0.00012263591151906127 -0.011623494029190608
        1.0 0.0
        0.0 0.0
        0.0 1.0
    ]
    @test c.H_x ≈ [
        0.0 0.0
        -0.98 0.0
        -0.07263157894736837 -6.884057971014498
        0.0 -0.2
    ]
    @test c.Ψ[1] ≈ [
        -0.009560768753410337 0.0 0.0 0.0 -2.0658808482697935e-5 0.0019580523687917364 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.009560768753410338 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        -2.0658808482697935e-5 0.0 0.0 0.0 -3.881681383327978e-6 0.00012263591151906127 0.0 0.0
        0.0019580523687917364 0.0 0.0 0.0 0.00012263591151906127 -0.011623494029190608 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
    ]
    @test c.Ψ[2] ≈ zeros(8, 8)
    @test c.Ψ[3] ≈ [
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
        0.0 0.0 0.0 0.0 0.0 0.0 0.0007663134567721225 -0.07263157894736837
        0.0 0.0 0.0 0.0 0.0 0.0 -0.07263157894736837 -6.884057971014498
    ]
    @test c.Ψ[4] ≈ zeros(8, 8)
    @test c.Γ ≈ [0.01 0.0; 0.01 0.01]
    @test c.Γ_p ≈ [
        [0.0 0.0; 0.0 0.0],
        [0.0 0.0; 0.0 0.0],
        [0.0 0.0; 0.0 0.0],
        [0.0 0.0; 0.0 0.0],
        [1.0 0.0; 1.0 1.0],
    ]

    # solution tests
    @test c.y ≈ [5.936252888048733, 6.884057971014498]
    @test c.x ≈ [47.39025414828824, 0.0]
    @test hcat(c.y_p...) ≈ [
        55.78596896689701 76.10141579073955 0.0 -116.07178189943077 0.0
        66.89124302798608 105.01995379122064 0.0 -94.78050829657676 0.0
    ]
    @test hcat(c.x_p...) ≈ [
        555.2637030544529 1445.9269000240533 0.0 -1304.9490272717098 0.0
        0.0 0.0 0.0 0.0 0.0
    ]
    @test c.g_x ≈ [
        0.0957964300241661 0.6746869652586178
        0.07263157894736878 6.884057971014507
    ]
    @test c.h_x ≈ [0.9568351489232028 6.209371005755889; -1.5076865909646354e-18 0.2]
    @test c.g_x_p ≈ [
        [
            -0.12465264193058262 5.596211904442805
            -1.2823781479976832e-15 66.89124302798608
        ],
        [
            -1.6946742377792863 -0.8343618226192915
            -1.1080332409972313 105.01995379122064
        ],
        [4.640454298222595e-19 0.2945008469346586; 0.0 0.0],
        [
            0.6277268890968761 -5.0369653468355455
            1.0000000000000024 -94.78050829657676
        ],
        [0.0 0.0; 0.0 0.0],
    ]
    @test c.h_x_p ≈ [
        [0.12465264193058134 61.29503112354326; 0.0 0.0],
        [0.586640996782055 105.85431561383992; 0.0 0.0],
        [
            -4.640454298222595e-19 -0.2945008469346586
            1.5757015120748295e-18 1.0
        ],
        [-0.6277268890968737 -89.74354294974118; 0.0 0.0],
        [0.0 0.0; 0.0 0.0],
    ]
    @test c.Σ ≈ [0.0001 0.0001; 0.0001 0.0002]
    @test c.Σ_p[1] ≈ zeros(2,2)
    @test c.Σ_p[2] ≈ zeros(2,2)
    @test c.Σ_p[3] ≈ zeros(2,2)
    @test c.Σ_p[4] ≈ zeros(2,2)
    @test c.Σ_p[5] ≈  [0.02 0.02; 0.02 0.04]
end

@testset "Schur Decomposition Failure" begin
    m = @include_example_module(Examples.rbc)
    p_f = (ρ=0.2, δ=0.02, σ=0.01)
    p_d = (α=100.0, β=0.0) # garbage parameters
    settings = PerturbationSolverSettings(; print_level = 0)
    sol = generate_perturbation(m, p_d, p_f;settings)
    @test sol.retcode == :Failure
end

@testset "BK Condition Failure" begin
    m = @include_example_module(Examples.rbc)
    p_f = nothing
    p_d = (α=0.5, β=0.95, ρ=1.01, δ=0.02, σ=0.01) # rho > 1
    settings = PerturbationSolverSettings(; print_level = 2)
    sol = generate_perturbation(m, p_d, p_f;settings)
    @test sol.retcode == :BlanchardKahnFailure
end

@testset "Callbacks" begin
    calculate_steady_state_callback_triggered = false
    function calculate_steady_state_callback(ret, m, cache, settings, p)
        calculate_steady_state_callback_triggered = true

        return nothing
    end
    evaluate_functions_callback_triggered = false
    function evaluate_functions_callback(ret, m, cache, settings, p)
        evaluate_functions_callback_triggered = true
        return nothing
    end
    solve_first_order_callback_triggered = false
    function solve_first_order_callback(ret, m, cache, settings)
        solve_first_order_callback_triggered = true
        return nothing
    end
    solve_first_order_p_callback_triggered = false
    function solve_first_order_p_callback(ret, m, cache, settings)
        solve_first_order_p_callback_triggered = true
        return nothing
    end

    settings = PerturbationSolverSettings(;
        print_level = 0,
        calculate_steady_state_callback,
        evaluate_functions_callback,
        solve_first_order_callback,
        solve_first_order_p_callback,
    )

    m = @include_example_module(Examples.rbc_observables)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)

    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f;settings, cache = c)
    generate_perturbation_derivatives!(m, p_d, p_f, c, ;settings)

    @test calculate_steady_state_callback_triggered == true
    @test evaluate_functions_callback_triggered == true
    @test solve_first_order_callback_triggered == true
    @test solve_first_order_p_callback_triggered == true
end

@testset "First Order Pullback inference" begin
    m = @include_example_module(Examples.rbc_observables)
    p_f = (ρ=0.2, δ=0.02, σ=0.01, Ω_1=0.01)
    p_d = (α=0.5, β=0.95)

    c = SolverCache(m, Val(1), p_d)
    sol = generate_perturbation(m, p_d, p_f; cache = c)

    _, pb = Zygote.pullback(generate_perturbation, m, p_d, p_f, Val(1))
    @inferred pb(sol)
end


text:
using Documenter
using DocumenterCitations
using Weave
using GeometricIntegrators

bib = CitationBibliography("GeometricIntegrators.bib")

makedocs(bib,
    sitename = "GeometricIntegrators.jl",
    format = Documenter.HTML(
               prettyurls = get(ENV, "CI", nothing) == "true",
               mathengine = MathJax3(Dict(
                 :tex => Dict(
                   :macros => Dict(
                     :bigtimes => "\\mathop{\\vcenter{\\huge\\times}}",
                     :ext => "\\mathsf{d}",
                     :id => "\\mathop{id}",
                     :rank => "\\mathop{rank}",
                     :spn => "\\mathop{span}",
                     :identity => "\\mathbb{I}",
                     :abs => ["\\left \\vert #1 \\right \\vert", 1],
                     :mf => ["\\mathcal{#1}", 1],
                     :tb => ["\\mathsf{T}_{#1} #2", 2, ""],
                     :cb => ["\\mathsf{T}_{#1}^{*} #2", 2, ""],
                   ),
                   :inlineMath => [["\$","\$"], ["\\(","\\)"]],
                   :tags => "ams",
                 ),
                 :options => Dict(
                   :ignoreHtmlClass => "tex2jax_ignore",
                   :processHtmlClass => "tex2jax_process",
                 ),
               ), true)
             ),
    pages = ["Home" => "index.md",
             "Tutorial" => "tutorial/tutorial.md",
             "Equations" => "equations.md",
             "Integrators" => [
                "Overview"    => "integrators.md",
                "Runge-Kutta" => "integrators/rk.md",
                "Splitting"   => "integrators/splitting.md",
                "Variational" => "integrators/variational.md",
                "VPRK"        => "integrators/vprk.md",
                "SPARK"       => "integrators/spark.md",
                "CGVI"        => "integrators/cgvi.md",
               #  "DGVI"        => "integrators/dgvi.md",
               #  "HPG"         => "integrators/hpg.md",
               ],
             "Modules" => [
               # "Discontinuities"     => "modules/discontinuities.md",
                "Equations"           => "modules/equations.md",
                "Integrators"         => "modules/integrators.md",
                "Simulations"         => "modules/simulations.md",
                "Solutions"           => "modules/solutions.md",
               ],
             "Tableaus" => [
                "Runge-Kutta Methods" => "tableaus/rungekutta.md",
                "Partitioned Runge-Kutta Methods" => "tableaus/rungekutta_partitioned.md",
                "Splitting Methods"   => "tableaus/splitting.md",
                "VPRK Methods"        => "tableaus/vprk.md",
                "SPARK Methods"       => "tableaus/spark.md",
               ],
             "Developer Docs" =>[
                "Code Integration"    => "developer/code_integration.md",
                "Custom Integrators"  => "developer/custom_integrators.md",
               ],
             "Release Notes" => "releasenotes.md",
             "Bibliography" => "bibliography.md",
             ]
)

deploydocs(
    repo   = "github.com/JuliaGNI/GeometricIntegrators.jl",
    push_preview = true,
)

ext: jl


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

struct TeamProfileApi <: SwaggerApi
    client::Swagger.Client
end

function _swaggerinternal_teamProfileGet(_api::TeamProfileApi; visibility=nothing, _mediaType=nothing)
    _ctx = Swagger.Ctx(_api.client, "GET", Dict{String, Any}, "/team.profile.get", ["slackAuth"])
    Swagger.set_param(_ctx.query, "visibility", visibility)  # type String
    Swagger.set_header_accept(_ctx, ["application/json"])
    Swagger.set_header_content_type(_ctx, (_mediaType === nothing) ? ["application/x-www-form-urlencoded"] : [_mediaType])
    return _ctx
end

"""

Retrieve a team's profile.
Param: visibility::String
Return: Dict{String, Any}
"""
function teamProfileGet(_api::TeamProfileApi; visibility=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_teamProfileGet(_api; visibility=visibility, _mediaType=_mediaType)
    Swagger.exec(_ctx)
end

function teamProfileGet(_api::TeamProfileApi, response_stream::Channel; visibility=nothing, _mediaType=nothing)
    _ctx = _swaggerinternal_teamProfileGet(_api; visibility=visibility, _mediaType=_mediaType)
    Swagger.exec(_ctx, response_stream)
end

export teamProfileGet


text:
"""
	mutable struct MomentumUpdater <: AbstractFieldUpdater

	MomentumUpdater(; method = :MOM, lam = 0.01, gamma = 0.9, beta1 = 0.9, beta2 = 0.999)

A field updater implementing various momentum-based iteration techniques.
These techniques track past gradient and step directions
to obtain a better current step direction.

Options for method include:
* MOM - Standard momentum update
* RMS - RMSprop update
* ADAM - Combination of ADAgrad and RMSprop, used commonly in NN training
* NGS - Nesterov accelerated gradient
"""
mutable struct MomentumUpdater <: AbstractFieldUpdater
	method :: Symbol
	lam    :: Float64
	gamma  :: Float64
	beta1  :: Float64
	beta2  :: Float64

	step   :: Dict{Int,FieldGrid{Float64}} # Final step for each method
	mt     :: Dict{Int,FieldGrid{Float64}} # First moment of gradient moving average
	vt     :: Dict{Int,FieldGrid{Float64}} # Second moment of gradient moving average
	temp   :: Dict{Int,FieldGrid{Float64}}

	system :: Option{FieldSystem}

	function MomentumUpdater(; method::Symbol = :MOM, lam::Real = 0.01,
		gamma::Real = 0.9, beta1::Real = 0.9, beta2::Real = 0.999)
		@assert 0.0 < gamma < 1.0
		@assert 0.0 < beta1 < 1.0
		@assert 0.0 < beta2 < 1.0

		if !(method in (:MOM, :RMS, :ADAM, :NAG))
			error("Invalid `method` parameter for MomentumUpdater. Options include: (:MOM, :RMS, :ADAM, :NAG)")
		end
		return new(method, lam, gamma, beta1, beta2, Dict(), Dict(), Dict(), Dict(), nothing)
	end
end

#==============================================================================#

Base.show(io::IO, updater::MomentumUpdater) = @printf(io, "MomentumUpdater(method = %s, lam = %.3g)", updater.method, updater.lam)

function setup!(updater::MomentumUpdater, sys::FieldSystem)
	updater.system = sys

	# Copy fields to step and temp
	copydict!(updater.step, sys.residuals)
	copydict!(updater.temp, sys.fields)

	# Initialize first and second moments of gradient
	for mid in keys(sys.fields)
		updater.mt[mid] = zeros(Float64, sys.dims)
		updater.vt[mid] = zeros(Float64, sys.dims)
	end

	return nothing
end

function step!(updater::MomentumUpdater)
	@assert !isnothing(updater.system)
	sys = updater.system

	# Update steps
	if updater.method == :MOM
		_step_mom!(updater)
	elseif updater.method == :RMS
		_step_rms!(updater)
	elseif updater.method == :ADAM
		_step_adam!(updater)
	elseif updater.method == :NAG
		_step_nag!(updater)
	end

	# Take step
	for (mid, omega) in sys.fields
		@. omega += updater.lam * updater.step[mid]
	end

	return nothing
end

function _step_mom!(updater::MomentumUpdater)
	sys = updater.system
	residuals!(sys)

	# Update the first moment history
	for (mid, m) in updater.mt
		@. m = updater.beta1 * m + (1 - updater.beta1) * sys.residuals[mid]
	end

	# Copy updated momentum to step
	copydict!(updater.step, updater.mt)

	return nothing
end

function _step_rms!(updater::MomentumUpdater)
	sys = updater.system
	residuals!(sys)

	# Update the second moment history
	for (mid, v) in updater.vt
		@. v = updater.beta2 * v + (1 - updater.beta2) * sys.residuals[mid]^2
	end

	# Update the step given new residuals and momentum
	for (mid, s) in updater.step
		@. s = sys.residuals[mid] / (sqrt(updater.vt[mid]) + 1e-8)
	end

	return nothing
end

function _step_adam!(updater::MomentumUpdater)
	sys = updater.system
	residuals!(sys)

	# Update the first moment history
	for (mid, m) in updater.mt
		@. m = updater.beta1 * m + (1 - updater.beta1) * sys.residuals[mid]
	end

	# Update the second moment history
	for (mid, v) in updater.vt
		@. v = updater.beta2 * v + (1 - updater.beta2) * sys.residuals[mid]^2
	end

	# Compute the new step
	for (mid, s) in updater.step
		 @. s = updater.mt[mid] / (sqrt(updater.vt[mid]) + 1e-8)
	end

	return nothing
end

function _step_nag!(updater::MomentumUpdater)
	sys = updater.system

	# Store current fields in temp
	copydict!(updater.temp, sys.fields)

	# Take a step with previous momentum term
	for (mid, omega) in sys.fields
		@. omega += updater.lam * updater.mt[mid]
	end

	# Compute gradient at update location
	residuals!(sys)

	# Update momentum w/ new residuals
	for (mid, m) in updater.mt
		@. m = updater.gamma * m + (1 - updater.gamma) * sys.residuals[mid]
	end

	# Copy updated momentum to step
	copydict!(updater.step, updater.mt)

	# Revert fields back to their previous values
	copydict!(sys.fields, updater.temp)

	return nothing
end


text:
using Images # requires Package Images.jl
using Colors # requires Package Colors.jl
using Mocha
using Compat

type ImageClassifier
  net           :: Net
  channel_order :: NTuple{3,Int} # The channel order of the trained net, (1,2,3) means RGB
  sp_order      :: NTuple{2,Int} # The spatial order (1,2) means width-height (row major)
  grayscale     :: Bool          # Whether we are working with gray (single-channel) images
  classes       :: Vector        # Names for each class

  data_layer :: MemoryDataLayer
  pred_blob  :: Blob
  pred       :: Array
  data_type  :: Type
  batch_size :: Int
  image_wh   :: NTuple{2, Int}
end

function ImageClassifier(net::Net, softmax_blob_name::Symbol;
    channel_order = (1,2,3), # default RGB
    sp_order  = (1,2),       # default row-major (width,height,channels)
    grayscale = false,       # default not grayscale
    classes = [],            # default no name information available
    )

  if length(net.data_layers) != 1 || !isa(net.layers[net.data_layers[1]], MemoryDataLayer)
    error("The network should contain exactly one MemoryDataLayer")
  end
  data_layer = net.layers[net.data_layers[1]]
  batch_size = data_layer.batch_size
  data_type  = eltype(data_layer.data[1])
  image_wh   = (size(data_layer.data[1],1), size(data_layer.data[1],2))

  @assert length(data_layer.data) == 1
  if size(data_layer.data[1], 4) != batch_size
    # to make our life easier
    error("Please make the num dimension ($(size(data_layer.data[1],4))) of MemoryDataLayer's data source the same to batch_size ($(batch_size))")
  end

  pred_blob = net.output_blobs[softmax_blob_name]
  pred = Array(data_type, size(pred_blob))

  return ImageClassifier(net, channel_order, sp_order, grayscale, classes,
      data_layer, pred_blob, pred, data_type, batch_size, image_wh)
end

################################################################################
# Interfaces
################################################################################
#-- Classify a single image via filename
function classify(classifier::ImageClassifier, filename::AbstractString)
  results = classify(classifier, AbstractString[filename])
  return results[1]
end

#-- Classify a bunch of images via filename
function classify{T<:AbstractString}(classifier::ImageClassifier, filenames::Vector{T})
  return classify(classifier, map(Images.imread, filenames))
end

#-- Classify a single image
function classify(classifier::ImageClassifier, image::Image)
  results = classify(classifier, Image[image])
  return results[1]
end

#-- Classify a bunch of images
function classify(classifier::ImageClassifier, images::Vector{Image})
  results = classify_batch(classifier, images, 1)
  idx = classifier.batch_size+1
  while idx <= length(images)
    ret_batch = classify_batch(classifier, images, idx)
    results = [results, ret_batch]
    idx += length(res_batch)
  end

  return results
end

################################################################################
# Implementations
################################################################################
#-- Propress images
function preprocess(classifier::ImageClassifier, images::Vector{Image})
  map(images) do image
    if (width(img), height(img)) != classifier.image_wh
      error("Image size ($(width(img)),$(height(img))) does not match the network input $(classifier.image_wh), and Julia does not have imresize yet........")
    end

    if classifier.grayscale
      if colorspace(image) != "Gray"
        image = convert(Image{Gray}, image)
      end
    else
      image = convert(Image{RGB}, image)
      image = separate(image) # separate color channels
    end

    data = convert(Array, image)
    data = reshape(data, size(data,1),size(data,2),size(data,3))
    if spatialorder(image) == ["x","y"]
      # row major
    elseif spatialorder(image) == ["y","x"]
      data = permutedims(data, (2,1,3)) # permute to row-major
    else
      error("Unknown spatialorder: $(spatialorder(image))")
    end

    # now data is in canonical row-major RGB/Gray format
    if classifier.sp_order != (1,2)
      @assert classifier.sp_order == (2,1)
      data = permutedims(data, (2,1,3))
    end

    if classifier.grayscale
      data2 = convert(Array{classifier.data_type}, data)
    else
      data2 = Array(classifier.data_type, tuple(classifier.image_wh..., 3))
      for i = 1:3
        data2[:,:,i] = convert(Array{classifier.data_type},
            data[:,:,classifier.channel_order[i]])
      end
    end

    data2 # pre-processed raw image data
  end
end

#-- Classify a batch or less
function classify_batch(classifier::ImageClassifier, images::Vector{Image}, idx::Int)
  #-- prepare the data
  idx_end = min(idx+classifier.batch_size-1, length(images))
  images = preprocess(classifier, images[idx:idx_end])
  for i = 1:length(images)
    classifier.data_layer.data[1][:,:,:,i] = images[i]
  end

  #-- run the network
  forward(classifier.net)
  copy!(classifier.pred, classifier.pred_blob)
  results = map(1:length(images)) do i
    pred = classifier.pred[:,i]
    i_class = indmax(pred)
    if !isempty(classifier.classes)
      ret = classifier.classes[i_class]
    else
      ret = string(i_class-1)
    end
    return (pred, ret)
  end

  return results
end


text:
@test true
push!(RUN, 0)

max_issues_repo_name: rfourquet/InlineTest

max_forks_repo_name: rfourquet/InlineTest


text:
include("hdfstests.jl")
include("yarntests.jl")

test_hdfs()
test_yarn()


text:
using aoc2019.computer: load_program, io, run

input = joinpath(@__DIR__, "input")
p = load_program(input)

function damage(stdout, display=false)
    for thing in stdout
        if thing > 256
            return thing
        elseif display
            c = convert(Char, thing)
            print("$c")
        end
    end
end

function program(stdin, input)
    for char in input
        put!(stdin, char)
    end
end

walkcode = """
NOT A T
OR T J
NOT B T
OR T J
NOT C T
OR T J
AND D J
WALK
"""

runcode = """
NOT A J
NOT B T
OR T J
NOT C T
OR T J
AND D J
NOT E T
AND H T
OR E T
AND T J
RUN
"""

stdin, stdout = io()
task = @async run(copy(p), stdin, stdout)

program(stdin, walkcode)

while !istaskdone(task); yield() end
close(stdout)

p1 = damage(stdout)
@assert p1 == 19353074

stdin, stdout = io()
task = @async run(copy(p), stdin, stdout)

program(stdin, runcode)

while !istaskdone(task); yield() end
close(stdout)

p2 = damage(stdout)
@assert p2 == 1147582556

print("-----------------------------------------------------------------------\n")
print("springdroid adventure -- part one\n    damage : $p1\n")
print("springdroid adventure -- part two\n    damage : $p2\n")
print("-----------------------------------------------------------------------\n")


text:
version = "0.8.0"
linux_build_id = "1"
macos_build_id = "1"
windows_build_id = "vc14_1"

function is_64_bits()
    return Int == Int64
end

if VERSION >= v"0.7.0-DEV.3073"
    const _HOME = Sys.BINDIR
else
    const _HOME = JULIA_HOME
end

function unpack(file, directory)
    if is_unix()
        run(`mkdir -p $directory`)
        run(`tar xjf $file --directory=$directory`)
    elseif is_windows()
        exe7z = joinpath(_HOME, "7z.exe")
        run(pipeline(`$exe7z x $file -y -so`, `$exe7z x -si -y -ttar -o$directory`))
    end
end

if is_windows()
    build_id = windows_build_id
    unpacked_file = joinpath("Library", "bin", "chemfiles.dll")
    if is_64_bits()
        platform = "win-64"
    else
        assert(Int == Int32)
        platform = "win-32"
    end
elseif is_linux()
    build_id = linux_build_id
    unpacked_file = joinpath("lib", "libchemfiles.so.$version")
    if !is_64_bits()
        error("There is no prebuilt chemfiles library for 32bit Linux")
    end
    platform = "linux-64"
elseif is_apple()
    build_id = macos_build_id
    unpacked_file = joinpath("lib", "libchemfiles.$version.dylib")
    if !is_64_bits()
        error("There is no prebuilt chemfiles library for 32bit macOS")
    end
    platform = "osx-64"
else
    error("Could not find a prebuilt chemfiles library for the current platform")
end

URL = "https://anaconda.org/conda-forge/chemfiles-lib/$version/download/$platform/chemfiles-lib-$version-$build_id.tar.bz2"
LOCAL_ARCHIVE = joinpath(@__DIR__, basename(URL))

info("Downloading pre-build library ...")
download(URL, LOCAL_ARCHIVE)
info("Unpacking library ...")
unpack(LOCAL_ARCHIVE, joinpath(@__DIR__, "usr"))

LIBPATH = joinpath(@__DIR__, "usr", unpacked_file)

write(joinpath(@__DIR__, "deps.jl"), """
# This is an auto-generated file; do not edit

# Macro to load a library
macro checked_lib(libname, path)
    Base.Libdl.dlopen_e(path) == C_NULL && error("Unable to load \\n\\n\$libname (\$path)\\n\\nPlease re-run Pkg.build(package), and restart Julia.")
    quote const \$(esc(libname)) = \$path end
end

# Load dependencies
@checked_lib libchemfiles "$(escape_string(LIBPATH))"
""")

info("Done")


text:
immutable ConcretePoint
x::Float64
y::Float64
end

max_stars_repo_path: Module 2/Chapter03/ConcrertePoint.jl

max_issues_repo_path: Module 2/Chapter03/ConcrertePoint.jl

max_forks_repo_path: Module 2/Chapter03/ConcrertePoint.jl


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""VolumeNodeResources is a set of resource limits for scheduling of volumes.

    IoK8sApiStorageV1beta1VolumeNodeResources(;
        count=nothing,
    )

    - count::Int32 : Maximum number of unique volumes managed by the CSI driver that can be used on a node. A volume that is both attached and mounted on a node is considered to be used once, not twice. The same rule applies for a unique volume that is shared among multiple pods on the same node. If this field is nil, then the supported number of volumes on this node is unbounded.
"""
mutable struct IoK8sApiStorageV1beta1VolumeNodeResources <: SwaggerModel
    count::Any # spec type: Union{ Nothing, Int32 } # spec name: count

    function IoK8sApiStorageV1beta1VolumeNodeResources(;count=nothing)
        o = new()
        validate_property(IoK8sApiStorageV1beta1VolumeNodeResources, Symbol("count"), count)
        setfield!(o, Symbol("count"), count)
        o
    end
end # type IoK8sApiStorageV1beta1VolumeNodeResources

const _property_map_IoK8sApiStorageV1beta1VolumeNodeResources = Dict{Symbol,Symbol}(Symbol("count")=>Symbol("count"))
const _property_types_IoK8sApiStorageV1beta1VolumeNodeResources = Dict{Symbol,String}(Symbol("count")=>"Int32")
Base.propertynames(::Type{ IoK8sApiStorageV1beta1VolumeNodeResources }) = collect(keys(_property_map_IoK8sApiStorageV1beta1VolumeNodeResources))
Swagger.property_type(::Type{ IoK8sApiStorageV1beta1VolumeNodeResources }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiStorageV1beta1VolumeNodeResources[name]))}
Swagger.field_name(::Type{ IoK8sApiStorageV1beta1VolumeNodeResources }, property_name::Symbol) =  _property_map_IoK8sApiStorageV1beta1VolumeNodeResources[property_name]

function check_required(o::IoK8sApiStorageV1beta1VolumeNodeResources)
    true
end

function validate_property(::Type{ IoK8sApiStorageV1beta1VolumeNodeResources }, name::Symbol, val)
end


text:
module LocalDescent

using LinearAlgebra

import ..Bracketing:
search,
GoldenSection

abstract type LocalDescentMethods end

# Line Search; Chapter 4.2

struct LineSearch <: LocalDescentMethods
    max_iter
    LineSearch(; max_iter = 100) = new(max_iter)
end

function search(params::LineSearch, f, x_k, d_k)
    """
    Line Search
    Algorithm 4.1
    """
    objective = α->f(x_k + α * d_k)::Real
    α_star = search(GoldenSection(max_iter = params.max_iter), objective, 0)
    x_k + α_star * d_k
end

# Strong Backtracking; Chapter 4.3

struct StrongBacktracking <: LocalDescentMethods
    α0
    β
    σ
    α_factor
    max_iter_zooming
    function StrongBacktracking(;
        α0 = 1, β = 1e-4, σ = 0.1, α_factor = 2,
        max_iter_zooming = 1_000)
        new(α0, β, σ, α_factor, max_iter_zooming)
    end
end

struct DescentPointK
    f  # f(x_k)
    ∇df  # ∇f(x_k) ⋅ d_k
    x  # x_k
    d  # d_k
    DescentPointK(f, ∇f, x_k, d_k) = new(f(x_k), ∇f(x_k) ⋅ d_k, x_k, d_k)
end

struct StrongWolfeConditions
    """
    Strong Wolfe conditions are termination conditions for local descent
        search of optimal optimization step α.
    """
    first
    second
    function StrongWolfeConditions(params, f, ∇f, point_k)
        new(α->f(point_k.x + α * point_k.d) <= point_k.f + params.β * α * point_k.∇df,
            α->abs(∇f(point_k.x + α * point_k.d) ⋅ point_k.d) <= -params.σ * point_k.∇df)
    end
end

struct BracketConditions
    """
    Bracket conditions are termination condition for a search of a bracket
        [α_left, α_right] such that the bracket contains a local minimum of
        the function f(x + αd).
    """
    first
    second
    third
    function BracketConditions(wolfe_cond, f, ∇f, point_k)
        new(α->f(point_k.x + α * point_k.d) >= point_k.f,
            α->!wolfe_cond.first(α),
            α->(∇f(point_k.x + α * point_k.d) ⋅ point_k.d) >= 0)
    end
end

mutable struct MutBracket
    left
    right
end

function search(params::StrongBacktracking, f, ∇f, x_k, d_k)
    point_k = DescentPointK(f, ∇f, x_k, d_k)
    wolfe_cond = StrongWolfeConditions(params, f, ∇f, point_k)
    bracket_cond = BracketConditions(wolfe_cond, f, ∇f, point_k)
    bracket = find_bracket(params, bracket_cond)
    α = zoom_bracket!(bracket, params, wolfe_cond, bracket_cond.third)
    x_k + α * d_k
end

function find_bracket(params, bracket_cond; max_iter = 1_000_000)::MutBracket
    """
    `max_iter` is a very high number since upon reaching this limit the algorithm
    will crash with error. This is because a failure to find a bracket means there's
    something wrong with either the algorithm, or the objective function.
    """
    bracket = MutBracket(0, params.α0)
    for _ = 1:max_iter
        if bracket_cond.first(bracket.right) ||
            bracket_cond.second(bracket.right) ||
            bracket_cond.third(bracket.right)
            return bracket
        end
        bracket.left, bracket.right = bracket.right, params.α_factor * bracket.right
    end
    error("LocalDescent.find_bracket: max number of iterations reached")
end

function zoom_bracket!(bracket::MutBracket, params, wolfe_cond, is_ascending)::Real
    for _ = 1:params.max_iter_zooming
        if (wolfe_cond.first(bracket.right) &&
            wolfe_cond.second(bracket.right)) ||
            abs(bracket.left - bracket.right) < eps()
            return bracket.right
        end
        α = bracket.left + (bracket.right - bracket.left) / 2
        if is_ascending(α)
            bracket.right = α
        else
            bracket.left = α
        end
    end
    @warn "LocalDescent.zoom_bracket: max number of iterations reached; bracket: $bracket"
    bracket.right
end

end # module


text:
using BCVUMPS
using Plots
using JLD2
using FileIO
using Random

# 2x2 Ising
begin
    model = Ising22(2)
    magplot = plot()
    xlabel!("β")
    ylabel!("mag")
    eneplot = plot()
    xlabel!("β")
    ylabel!("ene")
    for D = [2,4,8]
        mag = []
        ene = []
        for β = 0.26:0.005:0.36
            printstyled("BCVUMPS D = $(D), β = $(β) \n"; bold=true, color=:red)

            env = bcvumps_env(model, β, D)
            mag = [mag;magnetisation(env, model, β)]
            ene = [ene;energy(env, model, β)]
            chkp_file = "./data/$(model)_β$(round(β+0.005,digits=4))_D$(D).jld2"
            if isfile(chkp_file) == false
                save(chkp_file, "env", env)
            end
        end
        β = 0.26:0.005:0.36
        plot!(magplot,β,mag,seriestype = :scatter,title = "magnetisation", label = "BCVUMPS D = $(D)", lw = 3)
        plot!(eneplot,β,ene,seriestype = :scatter,title = "energy", label = "BCVUMPS D = $(D)", lw = 3)
    end
    for L = [8,12,16]
        mcmag = []
        mcene = []
        for β = 0.26:0.005:0.36
            printstyled("MCMC L = $(L) β = $(β) \n"; bold=true, color=:red)
            mag22,ene22 = MCMC(model,L,β,10000,100000)
            mcmag = [mcmag;mag22]
            mcene = [mcene;ene22]
        end
        β = 0.26:0.005:0.36
        plot!(magplot,β,mcmag,label = "MCMC L = $(L)", lw = 2)
        plot!(eneplot,β,mcene,label = "MCMC L = $(L)", lw = 2)
    end
    t = plot(magplot, eneplot, layout = (1, 2), legend=:bottomleft)
    savefig(t,"./plot/$(model).svg")
end


text:
using Pkg
# Pkg.add("Mollifiers")
# Pkg.add("WeightedNorms")
# Pkg.add("MISPTools")
Pkg.add("PyCall")
# using Mollifiers
# using WeightedNorms
# using MISPExamples
# using MISPTools
using PyCall

max_stars_repo_name: jgoldfar/misppy

max_issues_repo_name: jgoldfar/misppy

max_forks_repo_name: jgoldfar/misppy


text:
function autocovariance(z::AbstractVector{T}, k::Integer) where {T}
    N = length(z)
    m = mean(z)
    γ = zero(T)
    @inbounds @simd for i in 1:N - k
        γ += (z[i + k] - m) * (z[i] - m)
    end
    return γ / N
end

function autocorrelation(z::AbstractVector, k::Integer)
    return autocovariance(z, k) / autocovariance(z, 0)
end

function autocovariance_matrix(z::AbstractVector, k::Integer)
    N = length(z)
    T = typeof(zero(eltype(z)) / 1)
    γ = Vector{T}(undef, k)
    for i in 1:k
        γ[i] = autocovariance(z, i - 1)
    end
    Γ = Matrix{T}(undef, k, k)
    for j in 1:k
        for i in 1:k
            Γ[i,j] = γ[abs(i - j) + 1]
        end
    end
    return Γ
end

function autocorrelation_matrix(z::AbstractVector, k::Integer)
    return autocovariance_matrix(z, k) ./ autocovariance(z, 0)
end

function partial_autocorrelation(z::AbstractVector, k::Integer)
    Ρ = autocorrelation_matrix(z, k)
    Ρstar = copy(Ρ)
    Ρstar[1:k - 1,k] = Ρstar[2:k,1]
    Ρstar[k,k] = autocorrelation(z, k)
    return det(Ρstar) / det(cholesky(Symmetric(Ρ)))
end


text:
# compare-hyperparameters-multivariatetpseudohuber-approxwasserstein
#
# This script tries to compute an approximation of the wasserstein metric
# for the samples drawn from the Bayesian regession for the AIS dataset.
#
# NOTE: Make sure you have run the compare-hyperparameters-multivariatetpseudohuber
# script, as this script uses the samples created from that run in order to
# compute the approximation wasserstein measures.

include("experiment_utils.jl")

using SteinDiscrepancy: wassersteindiscrete, wasserstein1d

# this controls the seed used for the subsampling
@parseintcli seed "s" "seed" 7
# set the sampler
@parsestringcli sampler "q" "sampler" "runsgrld"
# set the delta param for the Huber prior
@parsestringcli nu "n" "nu" "10.0"
# set the nu param degrees of freedom for the d-distribution
@parsestringcli delta "d" "delta" "0.1"
# Select a solver for Stein discrepancy optimization problem
#solver = "clp"
solver = "gurobi"
# the n & d used
n = 2000; d = 4
# the thinby parameter
thinby = 100
# distname is multivariatetpseudohuber
distname = "multivariatetpseudohuber"
# the SGRLD epsilons to run the experiemnt
epsilons = [1e-7, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2]
# batchsize
batchsize = 30
# wasserstein n - the n to thin from
wasserstein_n = 5000
# the directory the data is stored
dir = "compare-hyperparameters-multivariatetpseudohuber"
dirout = @sprintf("%s-approxwasserstein", dir)

# prep the MARLA data
marla_jsonpath = @sprintf(
    "results/%s-MARLA/data/julia_distname=%s_dataset=ais_n=2000000_thinby=20_epsilon=0.01_nu=%s_delta=%s_d=4_seed=7.json",
     dir,
     distname,
     nu,
     delta)

marla_json = load_json(marla_jsonpath)
marla_betas0 = marla_json["betas"]
marla_betas = Array(Float64, wasserstein_n, d)

for epsilon in epsilons
    println("Computing approx wasserstein for epsilon=$(epsilon)")
    epsilon = string(epsilon)
    # load up the jsons
    sample_jsonpath = @sprintf(
        "results/%s/data/julia_distname=%s_n=%s_thinby=%s_sampler=%s_epsilon=%s_nu=%s_delta=%s_batchsize=%s_d=%s_seed=%s.json",
         dir,
         distname,
         n,
         thinby,
         sampler,
         epsilon,
         nu,
         delta,
         batchsize,
         d,
         seed)
    sample_json = load_json(sample_jsonpath)
    # prep the data for computing
    sample_betas0 = sample_json["betas"]
    sample_betas = Array(Float64, n, d)
    for j in 1:d
        sample_betas[:,j] = convert(Array{Float64,1}, sample_betas0[j])
        thinrate = round(Int, length(marla_betas0[1]) / wasserstein_n)
        marla_betas[:,j] = convert(Array{Float64,1},
                                   marla_betas0[j][thinrate:thinrate:(wasserstein_n*thinrate)])
    end
    # conmpute the marginal 1d and 2d wasserstein metrics
    marginal1d = Dict{Any,Any}[]
    marginal2d = Dict{Any,Any}[]
    for ii in 1:d
        for jj in 1:d
            if ii >= jj
                continue
            end
            tic()
            (wasserstein, numnodes, numedges, status) =
                wassersteindiscrete(xpoints=sample_betas[:,[ii,jj]],
                                    ypoints=marla_betas[:,[ii,jj]],
                                    solver=solver)
            solvetime = toc()
            result = Dict{Any, Any}(
                "wasserstein" => wasserstein,
                "numnodes" => numnodes,
                "numedges" => numedges,
                "wasserstein_n" => wasserstein_n,
                "betai" => ii,
                "betaj" => jj,
                "solvetime" => solvetime,
            )
            push!(marginal2d, result)
        end
        all_marla_betas = convert(Array{Float64,1}, marla_betas0[ii])
        (wasserstein, numnodes, numedges, status) =
            wasserstein1d(xpoints=sample_betas[:,ii],
                          ypoints=all_marla_betas)
        result = Dict{Any, Any}(
            "wasserstein" => wasserstein,
            "betai" => ii,
        )
        push!(marginal1d, result)
    end

    sample_json["marginal1d_wasserstein"] = marginal1d
    sample_json["marginal2d_wasserstein"] = marginal2d

    save_json(
        sample_json;
        dir=dirout,
        distname=distname,
        n=n,
        thinby=thinby,
        sampler=sampler,
        epsilon=epsilon,
        nu=nu,
        delta=delta,
        batchsize=batchsize,
        d=d,
        seed=seed,
        wassn=wasserstein_n
    )
end


text:
abstract type AbstractGridWorld <: RLBase.AbstractEnv end

@forward AbstractGridWorld.world get_grid, get_objects, get_num_objects, get_height, get_width
get_half_size(env::AbstractGridWorld) = (2 * (get_height(env) ÷ 4) + 1, 2 * (get_width(env) ÷ 4) + 1)

Random.rand(rng::Random.AbstractRNG, f::Function, env::AbstractGridWorld; max_try = 1000) = rand(rng, f, get_world(env), max_try = max_try)

#####
# utils
#####

macro generate_getters(type)
    T = getfield(__module__, type)::Union{Type,DataType}
    defs = Expr(:block)
    for field in fieldnames(T)
        get = Symbol(:get_, field)
        qn = QuoteNode(field)
        push!(defs.args, :($(esc(get))(instance::$type) = getfield(instance, $qn)))
    end
    return defs
end

macro generate_setters(type)
    T = getfield(__module__, type)::Union{Type,DataType}
    defs = Expr(:block)
    for field in fieldnames(T)
        set = Symbol(:set_, field, :!)
        qn = QuoteNode(field)
        push!(defs.args, :($(esc(set))(instance::$type, x) = setfield!(instance, $qn, x)))
    end
    return defs
end


text:
# Autogenerated wrapper script for OpenCVQt_jll for x86_64-apple-darwin14-cxx11
export libopencv_calib3d, libopencv_core, libopencv_dnn, libopencv_features2d, libopencv_flann, libopencv_gapi, libopencv_highgui, libopencv_imgcodecs, libopencv_imgproc, libopencv_ml, libopencv_objdetect, libopencv_photo, libopencv_stitching, libopencv_video, libopencv_videoio

using Qt_jll
using Libglvnd_jll
## Global variables
PATH = ""
LIBPATH = ""
LIBPATH_env = "DYLD_FALLBACK_LIBRARY_PATH"
LIBPATH_default = "~/lib:/usr/local/lib:/lib:/usr/lib"

# Relative path to `libopencv_calib3d`
const libopencv_calib3d_splitpath = ["lib", "libopencv_calib3d.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_calib3d_path = ""

# libopencv_calib3d-specific global declaration
# This will be filled out by __init__()
libopencv_calib3d_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_calib3d = "@rpath/libopencv_calib3d.4.5.dylib"

# Relative path to `libopencv_core`
const libopencv_core_splitpath = ["lib", "libopencv_core.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_core_path = ""

# libopencv_core-specific global declaration
# This will be filled out by __init__()
libopencv_core_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_core = "@rpath/libopencv_core.4.5.dylib"

# Relative path to `libopencv_dnn`
const libopencv_dnn_splitpath = ["lib", "libopencv_dnn.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_dnn_path = ""

# libopencv_dnn-specific global declaration
# This will be filled out by __init__()
libopencv_dnn_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_dnn = "@rpath/libopencv_dnn.4.5.dylib"

# Relative path to `libopencv_features2d`
const libopencv_features2d_splitpath = ["lib", "libopencv_features2d.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_features2d_path = ""

# libopencv_features2d-specific global declaration
# This will be filled out by __init__()
libopencv_features2d_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_features2d = "@rpath/libopencv_features2d.4.5.dylib"

# Relative path to `libopencv_flann`
const libopencv_flann_splitpath = ["lib", "libopencv_flann.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_flann_path = ""

# libopencv_flann-specific global declaration
# This will be filled out by __init__()
libopencv_flann_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_flann = "@rpath/libopencv_flann.4.5.dylib"

# Relative path to `libopencv_gapi`
const libopencv_gapi_splitpath = ["lib", "libopencv_gapi.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_gapi_path = ""

# libopencv_gapi-specific global declaration
# This will be filled out by __init__()
libopencv_gapi_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_gapi = "@rpath/libopencv_gapi.4.5.dylib"

# Relative path to `libopencv_highgui`
const libopencv_highgui_splitpath = ["lib", "libopencv_highgui.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_highgui_path = ""

# libopencv_highgui-specific global declaration
# This will be filled out by __init__()
libopencv_highgui_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_highgui = "@rpath/libopencv_highgui.4.5.dylib"

# Relative path to `libopencv_imgcodecs`
const libopencv_imgcodecs_splitpath = ["lib", "libopencv_imgcodecs.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_imgcodecs_path = ""

# libopencv_imgcodecs-specific global declaration
# This will be filled out by __init__()
libopencv_imgcodecs_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_imgcodecs = "@rpath/libopencv_imgcodecs.4.5.dylib"

# Relative path to `libopencv_imgproc`
const libopencv_imgproc_splitpath = ["lib", "libopencv_imgproc.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_imgproc_path = ""

# libopencv_imgproc-specific global declaration
# This will be filled out by __init__()
libopencv_imgproc_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_imgproc = "@rpath/libopencv_imgproc.4.5.dylib"

# Relative path to `libopencv_ml`
const libopencv_ml_splitpath = ["lib", "libopencv_ml.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_ml_path = ""

# libopencv_ml-specific global declaration
# This will be filled out by __init__()
libopencv_ml_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_ml = "@rpath/libopencv_ml.4.5.dylib"

# Relative path to `libopencv_objdetect`
const libopencv_objdetect_splitpath = ["lib", "libopencv_objdetect.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_objdetect_path = ""

# libopencv_objdetect-specific global declaration
# This will be filled out by __init__()
libopencv_objdetect_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_objdetect = "@rpath/libopencv_objdetect.4.5.dylib"

# Relative path to `libopencv_photo`
const libopencv_photo_splitpath = ["lib", "libopencv_photo.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_photo_path = ""

# libopencv_photo-specific global declaration
# This will be filled out by __init__()
libopencv_photo_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_photo = "@rpath/libopencv_photo.4.5.dylib"

# Relative path to `libopencv_stitching`
const libopencv_stitching_splitpath = ["lib", "libopencv_stitching.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_stitching_path = ""

# libopencv_stitching-specific global declaration
# This will be filled out by __init__()
libopencv_stitching_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_stitching = "@rpath/libopencv_stitching.4.5.dylib"

# Relative path to `libopencv_video`
const libopencv_video_splitpath = ["lib", "libopencv_video.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_video_path = ""

# libopencv_video-specific global declaration
# This will be filled out by __init__()
libopencv_video_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_video = "@rpath/libopencv_video.4.5.dylib"

# Relative path to `libopencv_videoio`
const libopencv_videoio_splitpath = ["lib", "libopencv_videoio.4.5.0.dylib"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libopencv_videoio_path = ""

# libopencv_videoio-specific global declaration
# This will be filled out by __init__()
libopencv_videoio_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libopencv_videoio = "@rpath/libopencv_videoio.4.5.dylib"

# Inform that the wrapper is available for this platform
wrapper_available = true

"""
Open all libraries
"""
function __init__()
    # This either calls `@artifact_str()`, or returns a constant string if we're overridden.
    global artifact_dir = find_artifact_dir()

    global PATH_list, LIBPATH_list
    # Initialize PATH and LIBPATH environment variable listings
    # From the list of our dependencies, generate a tuple of all the PATH and LIBPATH lists,
    # then append them to our own.
    foreach(p -> append!(PATH_list, p), (Qt_jll.PATH_list, Libglvnd_jll.PATH_list,))
    foreach(p -> append!(LIBPATH_list, p), (Qt_jll.LIBPATH_list, Libglvnd_jll.LIBPATH_list,))

    global libopencv_calib3d_path = normpath(joinpath(artifact_dir, libopencv_calib3d_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_calib3d_handle = dlopen(libopencv_calib3d_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_calib3d_path))

    global libopencv_core_path = normpath(joinpath(artifact_dir, libopencv_core_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_core_handle = dlopen(libopencv_core_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_core_path))

    global libopencv_dnn_path = normpath(joinpath(artifact_dir, libopencv_dnn_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_dnn_handle = dlopen(libopencv_dnn_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_dnn_path))

    global libopencv_features2d_path = normpath(joinpath(artifact_dir, libopencv_features2d_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_features2d_handle = dlopen(libopencv_features2d_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_features2d_path))

    global libopencv_flann_path = normpath(joinpath(artifact_dir, libopencv_flann_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_flann_handle = dlopen(libopencv_flann_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_flann_path))

    global libopencv_gapi_path = normpath(joinpath(artifact_dir, libopencv_gapi_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_gapi_handle = dlopen(libopencv_gapi_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_gapi_path))

    global libopencv_highgui_path = normpath(joinpath(artifact_dir, libopencv_highgui_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_highgui_handle = dlopen(libopencv_highgui_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_highgui_path))

    global libopencv_imgcodecs_path = normpath(joinpath(artifact_dir, libopencv_imgcodecs_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_imgcodecs_handle = dlopen(libopencv_imgcodecs_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_imgcodecs_path))

    global libopencv_imgproc_path = normpath(joinpath(artifact_dir, libopencv_imgproc_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_imgproc_handle = dlopen(libopencv_imgproc_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_imgproc_path))

    global libopencv_ml_path = normpath(joinpath(artifact_dir, libopencv_ml_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_ml_handle = dlopen(libopencv_ml_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_ml_path))

    global libopencv_objdetect_path = normpath(joinpath(artifact_dir, libopencv_objdetect_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_objdetect_handle = dlopen(libopencv_objdetect_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_objdetect_path))

    global libopencv_photo_path = normpath(joinpath(artifact_dir, libopencv_photo_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_photo_handle = dlopen(libopencv_photo_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_photo_path))

    global libopencv_stitching_path = normpath(joinpath(artifact_dir, libopencv_stitching_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_stitching_handle = dlopen(libopencv_stitching_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_stitching_path))

    global libopencv_video_path = normpath(joinpath(artifact_dir, libopencv_video_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_video_handle = dlopen(libopencv_video_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_video_path))

    global libopencv_videoio_path = normpath(joinpath(artifact_dir, libopencv_videoio_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libopencv_videoio_handle = dlopen(libopencv_videoio_path, RTLD_LAZY | RTLD_DEEPBIND)
    push!(LIBPATH_list, dirname(libopencv_videoio_path))

    # Filter out duplicate and empty entries in our PATH and LIBPATH entries
    filter!(!isempty, unique!(PATH_list))
    filter!(!isempty, unique!(LIBPATH_list))
    global PATH = join(PATH_list, ':')
    global LIBPATH = join(vcat(LIBPATH_list, [joinpath(Sys.BINDIR, Base.LIBDIR, "julia"), joinpath(Sys.BINDIR, Base.LIBDIR)]), ':')

end  # __init__()


text:
# calculate the spin correlation functions
# using the fixed point MPS obtained from the non-frustrated MPO

# run the code in the root directory

using TensorKit
using TensorOperations
using TensorKitAD
using Zygote
using KrylovKit
using LinearAlgebra
using Plots

using Revise
using statmech_tm_solver

# load prev data
A = quickload("data/triangular_AF_ising/checkpoints/ckpt_iTEBD_chi32_alternative")

# mps transfer matrix
@tensor M[-1, -2; -3, -4] := A[-2, 1, -4] * A'[-3, -1, 1]
W, V = eig(M)
Vinv = inv(V)
space(V)

V = permute(V, (2,), (1, 3))
Vinv = permute(Vinv, (1, 2), (3, ))

@show diag(W.data)

# observable
O = TensorMap(zeros, ComplexF64, ℂ^4, ℂ^4)
O[1, 1] = 1
O[2, 2] = 1
O[3, 3] = -1
O[4, 4] = -1
@show O

@tensor M1[-1, -2; -3, -4] := A[-2, 1, -4] * O[2, 1] * A'[-3, -1, 2]
M1 = permute(M1, (2, 3), (1, 4))

# environment
ER = convert_to_tensormap(convert_to_array(V)[:, :, end], 1)
EL = convert_to_tensormap(convert_to_array(Vinv)[end, :, :], 1)

@tensor ML[-1] := EL[3, 4] * M1[4, 1, 3, 2] * V[2, 1, -1]
ML = permute(ML, (), (1,))
@tensor MR[-1] := ER[4, 3] * M1[2, 3, 1, 4] * Vinv[-1, 1, 2]

# local observable <Z>
@tensor EL[1, 2] * M1[2, 4, 1, 3] * ER[3, 4]

# spin correlation, quasi-long-range order observed
corrs = zeros(300)
for n in 1:300
    corr_n = (ML * W^n * MR / tr(W^n))[1] |> real
    corrs[n] = corr_n
    @show n, corr_n
end

# log log plot of the spin correlation
gr()
plot(3:300, corrs[2:end])

plot(xs, ys)


text:
using FluxReconstruction, Test

ps = FRPSpace1D(0, 1, 100, 2)

#--- values ---#
V = vandermonde_matrix(ps.deg, ps.xpl)
ψf = vandermonde_matrix(ps.deg, [-1.0, 1.0])

lf = zeros(2, ps.deg + 1)
for i in axes(lf, 1)
    lf[i, :] .= V' \ ψf[i, :]
end

@test lf[1, :] ≈ ps.ll
@test lf[2, :] ≈ ps.lr

#--- derivatives ---#
Vr = ∂vandermonde_matrix(ps.deg, ps.xpl)

∂l = zeros(ps.deg + 1, ps.deg + 1)
for i = 1:ps.deg+1
    ∂l[i, :] .= V' \ Vr[i, :]
end

@test ∂l ≈ ps.dl

dVf = ∂vandermonde_matrix(ps.deg, [-1.0, 1.0])
∂lf = zeros(2, ps.deg + 1)
for i = 1:2
    ∂lf[i, :] .= V' \ dVf[i, :]
end


text:
"""
Contains some functions for calculating cross-correlations, autocorrelations,
and Pearson correlations among rate units in a rate network
"""

using Statistics
import DSP.xcorr

"""
	xcorr_unbiased(x::Vector, y::Vector, maxlag::Int64)

Returns the cross-correlation of two signals normalized by the length of the
available convolution window.

Written by Rainer Engelken (re2365@columbia.edu).
"""
function xcorr_unbiased(x, y, maxlag)
	l = length(x) # we assume same length for x and y
	lags = -maxlag:maxlag
	scale = l .- abs.(lags)
	scale[scale .<= 0] .= 1 # avoid zero division
    xcorr(x, y)[l .+ lags] ./ scale
end

"""
	avg_autocorrs(A::Matrix{T}, maxlag::Int64, wait::Int64=0) where T <: AbstractFloat

Returns the average autocorrelation for the activity of each unit described in
the `N x NT` matrix `A`.

Note the autocorrelations are not mean-subtracted. The `maxlag` parameter
gives the maximum shift of the signal, and the `wait` parameter gives the number
of timesteps in the beginning to disregard.
"""
function avg_autocorrs(A::Matrix{T}, maxlag::Int64, wait::Int64=0) where T <: AbstractFloat

	(N, NT) = size(A)

	unit_autocorrs = mapslices(view(A, :, (wait + 1):NT), dims=2) do x
		return xcorr_unbiased(x, x, maxlag)
	end

	return vec(mean(unit_autocorrs, dims=1))
end


text:
export TransformedSequence

struct TransformedSequence <: Sequence
  seq
  f
end

function sample(seq::TransformedSequence, t::AbstractVector)
  seq.f.(sample(seq.seq, t))
end

for f in (:exp, :log, :sin, :cos)
  @eval begin
    function Base.$f(s::Sequence)
      TransformedSequence(s, $f)
    end
  end
end


text:
module StaticCompiler

import GPUCompiler
import LLVM
import LLVM_full_jll
import Libdl

export generate_shlib, generate_shlib_fptr, compile

module TestRuntime
    # dummy methods
    signal_exception() = return
    malloc(sz) = C_NULL
    report_oom(sz) = return
    report_exception(ex) = return
    report_exception_name(ex) = return
    report_exception_frame(idx, func, file, line) = return

    # for validation
    sin(x) = Base.sin(x)
end

struct TestCompilerParams <: GPUCompiler.AbstractCompilerParams end
GPUCompiler.runtime_module(::GPUCompiler.CompilerJob{<:Any,TestCompilerParams}) = TestRuntime

const linker = Sys.isunix() ? "ld.lld" : Sys.isapple() ? "ld64.lld" : "lld-link"

function generate_shlib(f, tt, path::String = tempname(), name = GPUCompiler.safe_name(repr(f)))
    open(path, "w") do io
        target = GPUCompiler.NativeCompilerTarget(;reloc=LLVM.API.LLVMRelocPIC, extern=true)
        source = GPUCompiler.FunctionSpec(f, Base.to_tuple_type(tt), false, name)
        params = TestCompilerParams()
        job = GPUCompiler.CompilerJob(target, source, params)
        obj, _ = GPUCompiler.codegen(:obj, job; strip=true, only_entry=false, validate=false)
        write(io, obj)
        flush(io)
        run(`$(StaticCompiler.LLVM_full_jll.PATH)/$linker -shared -o $path.$(Libdl.dlext) $path`)
        rm(path)
    end
    path, name
end
function generate_shlib_fptr(f, tt, path::String=tempname(), name = GPUCompiler.safe_name(repr(f)); temp::Bool=true)
    generate_shlib(f, tt, path, name)
    ptr = Libdl.dlopen("$(abspath(path)).$(Libdl.dlext)", Libdl.RTLD_LOCAL)
    fptr = Libdl.dlsym(ptr, "julia_$name")
    @assert fptr != C_NULL
    if temp
        atexit(()->rm("$path.$(Libdl.dlext)"))
    end
    fptr
end
function generate_shlib_fptr(path::String, name)
    ptr = Libdl.dlopen("$(abspath(path)).$(Libdl.dlext)", Libdl.RTLD_LOCAL)
    fptr = Libdl.dlsym(ptr, "julia_$name")
    @assert fptr != C_NULL
    fptr
end

# Return an LLVM module
function compile(f, tt, name = GPUCompiler.safe_name(repr(f)))
    target = GPUCompiler.NativeCompilerTarget(;reloc=LLVM.API.LLVMRelocPIC, extern=true)
    source = GPUCompiler.FunctionSpec(f, Base.to_tuple_type(tt), false, name)
    params = TestCompilerParams()
    job = GPUCompiler.CompilerJob(target, source, params)
    m, _ = GPUCompiler.codegen(:llvm, job; strip=true, only_entry=false, validate=false)
    return m
end

end # module


text:
using Test
import LinearAlgebra.norm
@testset "Notebook: deep water" begin
    η(x)=exp.(-x.^2); # Gaussian initial data for the surface deformation.
    v(x)=zero(x)      # we set the initial velocity as zero to avoid inconsistencies among different models.
    init=Init(η,v);
    param_unstable = (
    # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
    ϵ  = 1/4,    # nonlinearity dimensionless parameter
    μ  = Inf,    # inifinite-depth case
    # Numerical parameters
    N  = 2^11,   # number of collocation points
    L  = 2*π,    # half-length of the numerical tank (-L,L)
    T  = 1.5,      # final time of computation
    dt = 0.01,  # timestep
                );
    WW2_unstable = Problem( WWn(param_unstable,dealias = 1), init, param_unstable )
    solve!(WW2_unstable);
    WW2_rectified = Problem( WWn(param_unstable,dealias = 1,δ=0.01), init, param_unstable,label="rWW2" ) # δ is the strength of the regularization
    solve!(WW2_rectified);
    WW2η,=solution(WW2_unstable,T=1);rWW2η,=solution(WW2_rectified,T=1);
    errWW2=norm(WW2η-rWW2η)/sqrt(length(WW2η))
    @test errWW2 < 1e-4 && errWW2 > 1e-5
    WW2η,=solution(WW2_unstable)
    @test any(isnan,WW2η) && !any(isnan,rWW2η)

    param_heap = (
    # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
    ϵ  = 1/4,    # nonlinearity dimensionless parameter
    μ  = Inf,    # infinite-depth case
    # Numerical parameters
    N  = 2^9,   # number of collocation points
    L  = 2*π,    # half-length of the numerical tank (-L,L)
    T  = 1,      # final time of computation
    dt = 0.01,  # timestep
    )
    init=Init(x->exp.(-x.^2),x->zero(x));
    WW_heap = Problem( WaterWaves(param_heap,dealias = 1, verbose=false), init, param_heap );
    rWW2_heap = Problem( WWn(param_heap,dealias = 1, δ=0.02), init, param_heap, label="rWW2" ) # δ is the strength of the regularization
    Matsuno_heap = Problem( Matsuno(param_heap,dealias = 1), init, param_heap )
    AkersNicholls_heap = Problem( AkersNicholls(param_heap,dealias = 1), init, param_heap )

    solve!(WW_heap);solve!(rWW2_heap);solve!(Matsuno_heap);solve!(AkersNicholls_heap);
    WWη,WWv,WWx,=solution(WW_heap)
    rWW2η,=solution(rWW2_heap,x=WWx);errWW2=norm(WWη-rWW2η)/sqrt(length(WWx))
    Matη,=solution(Matsuno_heap,x=WWx);errMat=norm(WWη-Matη)/sqrt(length(WWx))
    ANη,=solution(AkersNicholls_heap,x=WWx);errAN=norm(WWη-ANη)/sqrt(length(WWx))

    @test errWW2≈0.0008247817846683348
    @test errMat≈0.001381758315810658
    @test errAN ≈0.0007527210525751591

    param_rand = (
    # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
    ϵ  = 1/4,    # nonlinearity dimensionless parameter
    μ  = Inf,    # infinite-depth case
    # Numerical parameters
    N  = 2^10,   # number of collocation points
    L  = 2*π,    # half-length of the numerical tank (-L,L)
    T  = 1,      # final time of computation
    dt = 0.01,  # timestep
    )
    mesh=Mesh(param_rand);x=mesh.x;
    init= Random(x;L=2,a=(1,0)); # we set the initial velocity as zero to avoid possible inconsistencies among different models.
    WW_rand = Problem( WaterWaves(param_rand,dealias = 1, verbose=false), init, param_rand );
    rWW2_rand = Problem( WWn(param_rand,dealias = 1, δ=0.05), init, param_rand, label="rWW2" );
    Matsuno_rand = Problem( Matsuno(param_rand,dealias = 1), init, param_rand );
    AkersNicholls_rand = Problem( AkersNicholls(param_rand,dealias = 1), init, param_rand );

    solve!(WW_rand);solve!(rWW2_rand);solve!(Matsuno_rand);solve!(AkersNicholls_rand);

    WWη,WWv,WWx,=solution(WW_rand)
    rWW2η,=solution(rWW2_rand,x=WWx);errWW2=norm(WWη-rWW2η)/sqrt(length(WWx))
    Matη,=solution(Matsuno_rand,x=WWx);errMat=norm(WWη-Matη)/sqrt(length(WWx))
    ANη,=solution(AkersNicholls_rand,x=WWx);errAN=norm(WWη-ANη)/sqrt(length(WWx))

    @test 1e-4< errWW2 <2*1e-3
    @test 1e-4< errMat <2*1e-3
    @test 1e-4< errAN  <2*1e-3
end

@testset "Notebook: Example" begin
    param = (
    # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
    μ  = 1,     # shallow-water dimensionless parameter
    ϵ  = 1/4,   # nonlinearity dimensionless parameter
    # Numerical parameters
    N  = 2^9,  # number of collocation points
    L  = 10,    # half-length of the numerical tank (-L,L)
    T  = 5,     # final time of computation
    dt = 0.01,  # timestep
                );

    z(x) = exp.(-x.^4); # surface deformation
    v(x) = 0*z(x);      # zero initial velocity
    init = Init(z,v);   # generate the initial data with correct type

    model1=WaterWaves(param,verbose=false) # The water waves system
    model2=WWn(param;n=2,dealias=1,δ=1/10) # The quadratic model (WW2)
    # type `?WaterWaves` or `?WWn` to see details and signification of arguments
    problem1=Problem(model1, init, param, solver=RK4(model1));
    problem2=Problem(model2, init, param, solver=RK4(model2));

    solve!(problem1);solve!(problem2);

    s1,=solution(problem1);
    s2,=solution(problem2);
    err=norm(s1-s2)/sqrt(length(s1))

    @test err≈0.026407484058654316

end

@testset "Notebook: Full dispersion" begin

    param = (
    # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
    μ = 1,     # shallow-water dimensionless parameter
    ϵ = 1,     # nonlinearity dimensionless parameter
    c = 1.05,  # velocity of the solitary wave
    # Numerical parameters
    N = 2^7,  # number of collocation points
    L = 75    # half-length of the numerical tank (-L,L)
                    );
    # Compute the WGN solitary wave with velocity c (type `?SolitaryWaveWhithamGreenNaghdi` for more details)
    (ηWGN,uWGN,vWGN,mesh) = SolitaryWaveWhithamGreenNaghdi(param,verbose=false);
    # Compute the SGN solitary wave with velocity c (there is an exact formula)
    (ηSGN,uSGN,vSGN) = SolitaryWaveSerreGreenNaghdi(param);

    @test norm(ηWGN-ηSGN)/sqrt(param.N) ≈ 0.00018818525806149597
    @test norm(uWGN-uSGN)/sqrt(param.N) ≈ 0.00018061723014529588
    @test norm(vWGN-vSGN)/sqrt(param.N) ≈ 0.00018747696929027137

    T=2*param.L/param.c # final time of computation
    dt=T/10^3   # timestep
    init = Init(mesh,ηWGN,vWGN);           # use solitary wave as initial data
    model = WhithamGreenNaghdi(param)      # define the model
    pb = Problem(                          # set up the initial-value problem to be solved
        model,
        init,
        merge(param,(T=T,dt=dt)),
        solver=RK4(model)   )
    solve!(pb)

    ηfin,=solution(pb)
    err=norm(ηWGN-ηfin)/sqrt(param.N)

    @test err ≈ 2.5500799396961685e-6

    T=2*param.L/param.c # final time of computation
    dt=T/10^3   # timestep
    initWGN = Init(mesh,ηWGN,vWGN);       # WGN solitary wave as initial data
    initSGN = Init(mesh,ηSGN,vSGN);       # SGN solitary wave as initial data
    model = WaterWaves(param, dealias=0, verbose=false) # The water waves system
    pbSGN = Problem(                         # initial-value problem with SGN initial data
        model,
        initSGN,
        merge(param,(T=T,dt=dt)),
        solver=RK4(model)   )
    pbWGN = Problem(                         # initial-value problem with WGN initial data
        model,
        initWGN,
        merge(param,(T=T,dt=dt)),
        solver=RK4(model)   )
    solve!(pbSGN);solve!(pbWGN);

    ηSGNfin,vSGNfin,xSGNfin=solution(pbSGN)
    ηWGNfin,vWGNfin,xWGNfin=solution(pbWGN)
    ηSGNinit,vSGNinit,xSGNinit=solution(pbSGN,T=0)
    ηWGNinit,vWGNinit,xWGNinit=solution(pbWGN,T=0)

    errSGN = norm(ηSGNinit-ηSGNfin)/sqrt(param.N)
    errWGN = norm(ηWGNinit-ηWGNfin)/sqrt(param.N)

    @test errSGN ≈ 0.00043699621029233914
    @test errWGN ≈ 7.999694318031115e-5

end

@testset "Notebook: Shallow water" begin

    param = (
        # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
        μ  = 0.1,   # shallow-water dimensionless parameter
        ϵ  = 0.25,   # nonlinearity dimensionless parameter
        # Numerical parameters
        N  = 2^8,  # number of collocation points
        L  = π,    # half-length of the numerical tank (-L,L)
        T  = 1,    # final time of computation
        dt = 0.01, # timestep
                    );

    init= Init(x->exp.(-x.^4),zero); # we set the initial velocity as zero to avoid inconsistencies among different models.

    WW = Problem( WaterWaves(param,dealias = 1, verbose=false), init, param )
    GN = Problem( SerreGreenNaghdi(param; dealias = 0, iterate = true, precond = true), init, param )
    NH = Problem( NonHydrostatic(param; dealias = 0, iterate = true, precond = true), init, param )
    SRD= Problem( SquareRootDepth(param; dealias = 0, iterate = true, precond = true), init, param )
    WGN = Problem( WhithamGreenNaghdi(param; dealias = 0, iterate = true, precond = true), init, param )
    IK = Problem( IsobeKakinuma(param; dealias = 0, iterate = true, precond = true), init, param )

    solve!(WW);solve!(GN);solve!(NH);solve!(SRD);solve!(WGN);solve!(IK);

    WWη,WWv,WWx,=solution(WW)
    GNη,=solution(GN,x=WWx);errGN=norm(WWη-GNη)/norm(WWη)
    SRDη,=solution(SRD,x=WWx);errSRD=norm(WWη-SRDη)/norm(WWη)
    NHη,=solution(NH,x=WWx);errNH=norm(WWη-NHη)/norm(WWη)
    WGNη,=solution(WGN,x=WWx);errWGN=norm(WWη-WGNη)/norm(WWη)
    IKη,=solution(IK,x=WWx);errIK=norm(WWη-IKη)/norm(WWη)

    @test  errGN  ≈ 0.019492085082947582
    @test  errSRD ≈ 0.032762377673748344
    @test  errNH  ≈ 0.019380927144301712
    @test  errWGN ≈ 0.004015850164345053
    @test  errIK  ≈ 0.002501413432973788

end

@testset "Notebook: Hammack-Segur" begin
    d  = 0.1;  # depth of the basin (in m)
    g  = 9.81; # gravitational acceleration
    λ  = 0.1;  # typical horizontal length (=d to simplify)
    T  = 40;   # final time (in s)
    L  = 156.16;# half-length of the numerical tank (in m)
    param = (
        # Physical parameters. Variables are non-dimensionalized as in Lannes, The water waves problem, isbn:978-0-8218-9470-5
        μ  = 1,    # shallow-water dimensionless parameter
        ϵ  = 1,    # nonlinearity dimensionless parameter
        # Numerical parameters
        N  = 2^12, # number of collocation points
        L  = L/d,       # half-length of the numerical tank ( = 156.16 m)
        T  = T*sqrt(g*d)/λ, # final time of computation (= 50 s)
        dt = T*sqrt(g*d)/λ/10^3, # timestep
                    )
    using Elliptic
    a = 0.005 # half-amplitude of the initial data (in m)
    sn(x) = getindex.(ellipj.(9.25434*x*d,0.9999^2),1);  # regularized step function
    η0(x) = (-1/2*a.+1/2*a*sn(x)).*(x*d.<.61).*(x*d.>-1.83)/d;
    init = Init(x->2*η0(x),x->zero(x));         # generate the initial data with correct type

    model_WW2=WWn(param;n=2,dealias=1,δ=1/10) # The quadratic water waves model (WW2)
    #model_SGN=SerreGreenNaghdi(param) # The Serre-Green-Naghdi model (SGN)
    #model_WGN=WhithamGreenNaghdi(param) # The fully dispersive Whitham-Green-Naghdi model (WGN)
    # type `?WaterWaves` or `?WWn`, etc. to see details and signification of arguments
    WW2=Problem(model_WW2, init, param) ;
    #SGN=Problem(model_SGN, init, param) ;
    #WGN=Problem(model_WGN, init, param) ;

    solve!(WW2);#solve!(SGN);solve!(WGN);

    using HDF5
    data2a = h5read("./HammackSegur.h5","data2a")
    data2b = h5read("./HammackSegur.h5","data2b")
    data2c = h5read("./HammackSegur.h5","data2c")
    data2d = h5read("./HammackSegur.h5","data2d")
    data2e = h5read("./HammackSegur.h5","data2e")

    function gauge(p::Problem;x=0,T=nothing)
        if T==nothing
            times=p.times.ts
        elseif T[1]==T
            times=[T]
        else
            times=T
        end
        times/sqrt(g*d)*λ,[solution(p,T=ti,x=[x])[1][1] for ti in times]*d
    end
    nothing

    mesh_of_times = Mesh(WW2.times.ts /sqrt(g*d)*λ)

    ta=data2a[:,1]/sqrt(g*d)*λ
    ηa=data2a[:,2]*d*2/3
    ga=interpolate(mesh_of_times,gauge(WW2,x=1)[2],ta)

    @test norm(ηa-ga)/norm(ηa) ≈ 0.1672677795663852

    tb=(data2b[:,1].+50)/sqrt(g*d)*λ
    ηb=data2b[:,2]*d*2/3
    gb=interpolate(mesh_of_times,gauge(WW2,x=51)[2],tb)

    @test norm(ηb-gb)/norm(ηb) ≈ 0.13526329598410575

    tc=(data2c[:,1].+100)/sqrt(g*d)*λ
    ηc=data2c[:,2]*d*2/3
    gc=interpolate(mesh_of_times,gauge(WW2,x=101)[2],tc)

    @test norm(ηc-gc)/norm(ηc) ≈ 0.2599790311975911

    td=(data2d[:,1].+150)/sqrt(g*d)*λ
    ηd=data2d[:,2]*d*2/3
    gd=interpolate(mesh_of_times,gauge(WW2,x=151)[2],td)

    @test norm(ηd-gd)/norm(ηd) ≈ 0.4981462329721109

    te=(data2e[:,1].+200)/sqrt(g*d)*λ
    ηe=data2e[:,2]*d*2/3
    ge=interpolate(mesh_of_times,gauge(WW2,x=201)[2],te)

    @test norm(ηe-ge)/norm(ηe) ≈ 0.4921927105753341

end


text:
const ENDNOTES = """

#### Notes

* Scotland actually has 3 lower rates of income tax rather than the single 20% basic rate shown - currently 19%,20% and 21%. Changing the 20% 'basic rate'
causes all three to move in sync.
* Likewise, changing the *Universal Credit: single 25+ adult* field changes the rates for young people and couples.
* Scotland is [in the process of switching working-age families to Universal Credit from 'Legacy Benefits' (Income Support, Housing Benefit, etc.) ](https://commonslibrary.parliament.uk/constituency-data-universal-credit-roll-out/). I've written a [note on how this is modelled](https://stb-blog.virtual-worlds.scot/articles/2021/11/12/uc-legacy.html) - the code is [here](https://github.com/grahamstark/ScottishTaxBenefitModel.jl/blob/master/src/UCTransition.jl);
* the [£20pw 'uplift' to Universal Credit and Working Tax Credit](https://www.legislation.gov.uk/uksi/2021/313/pdfs/uksiem_20210313_en.pdf), now [scheduled for removal in April 2022](https://www.gov.uk/government/news/budget-2021-what-you-need-to-know) is modelled as already removed.

Poverty and Inequality:
* based on [equivalised before housing costs income](https://www.gov.scot/publications/poverty-in-scotland-methodology/pages/household-income-definition/).
* uses 60% of *Scottish* (not UK) median income as poverty line (official target uses the UK median).
* poverty line is [relative](https://www.gov.scot/publications/poverty-in-scotland-methodology/pages/poverty-definition/) - so may be different after your changes.

#### Key Assumptions

* *No behavioural changes*: increasing or decreasing taxes doesn't cause people to change how they work and earn;
* the model reports *entitlements to benefits and liability to taxes*, not receipts and payments - so we may overstate the costs of benefits since some eligible families may not claim the things they're entitled to. With taxes, some may be paid with a considerable delay, and some evaded or avoided.

See [the model blog](https://stb-blog.virtual-worlds.scot/) for more gory details (*content warning - very boring and rambling)*.

#### Known Problems

This is a new model. I'm now reasonably confident of its essential accuracy - it passes an [extensive test suite](https://github.com/grahamstark/ScottishTaxBenefitModel.jl/tree/master/test) but there are some aspects that
require investigation in the coming months. Notably:

* Income Tax revenues seem to be overstated by around £1bn pa compared to [official forecasts](https://www.fiscalcommission.scot/publications/scotlands-economic-and-fiscal-forecasts-august-2021/). Possibly much of this is due to how pension tax relief is treated;
* measures of inequality seem low compared to official statistics.

I'd very much welcome contributions and suggestions. If you spot anything odd or if you have any ideas for how this can be improved, you can:

* [Open an issue on GitHub](https://github.com/grahamstark/ScottishTaxBenefitModel.jl/issues); or
* [email me](mailto:graham.stark@virtual-worlds.biz).

### To Find Out More

You'll have to do some reading, I'm afraid. Some links:

* **Tax Benefit Models**: [A short introduction to microsimulation and tax benefit models](https://stb.virtual-worlds.scot/intro.html). Originally written for the Open University, it covers all the essential ideas. | [Blog Posts about the Model](https://stb-blog.virtual-worlds.scot/);
* **Poverty and Inequality**: [My Notes](https://stb.virtual-worlds.scot/poverty.html) | [World Bank Handbook](http://documents.worldbank.org/curated/en/488081468157174849/Handbook-on-poverty-and-inequality) | [Official Figures for Scotland](https://data.gov.scot/poverty/);
* **Scotland's Finances**: [Scottish Fiscal Commission](https://www.fiscalcommission.scot/publications/scotlands-economic-and-fiscal-forecasts-august-2021/) | [Scottish Government Budget Documents](https://www.gov.scot/budget/).

* Created with [Julia](https://julialang.org/) | [Dash](https://dash-julia.plotly.com/) | [Plotly](https://plotly.com/julia/) | [Poverty and Inequality Measures](https://github.com/grahamstark/PovertyAndInequalityMeasures.jl);
* Part of the [Scottish Tax Benefit Model](https://github.com/grahamstark/ScottishTaxBenefitModel.jl);
* Open Source software released under the [MIT Licence](https://github.com/grahamstark/Visualisations.jl/blob/main/LICENSE). [Source Code](https://github.com/grahamstark/Visualisations.jl).

"""


text:
module Figures

include("data.jl")
include("plotlib.jl")
include("plotutils.jl")

using PlotlyJS

export create_figure_1

plot_dir = "plots"

"""
	create_figure_1("medium")

Convenience function to create figure 1 of the notebook.
"""
function create_figure_1(plot_size::String)
	width, height = PlotUtils.get_plot_resolution(plot_size)
	loader_1 = Data.get_moon_data_loader(offset = 1.0, seed = 1337)
	loader_2 = Data.get_moon_data_loader(offset = 0.5, seed = 1337)
	p = PlotLib.plot_train_and_test_data(loader_1, loader_2; train_title = "Δ0.5", test_title = "Δ1.0")
	PlotlyJS.savefig(p, "$(plot_dir)/figure_1.png", width = width, height = height)
end

end


text:
#
# Unit tests for UnitfulAstrodynamics.jl
#

include("R2BP.jl")
include("CR3BP.jl")
include("Propagation.jl")
include("Visualizations.jl")