text:
using MultivariatePolynomials
using TypedPolynomials

using Test
using Random
using Pkg
using LinearAlgebra
using NBInclude

"""
    @wrappedallocs(expr)

Given an expression, this macro wraps that expression inside a new function
which will evaluate that expression and measure the amount of memory allocated
by the expression. Wrapping the expression in a new function allows for more
accurate memory allocation detection when using global variables (e.g. when
at the REPL).

For example, `@wrappedallocs(x + y)` produces:

```julia
function g(x1, x2)
    @allocated x1 + x2
end
g(x, y)
```

You can use this macro in a unit test to verify that a function does not
allocate:

```
@test @wrappedallocs(x + y) == 0
```
"""
macro wrappedallocs(expr)
    argnames = [gensym() for a in expr.args]
    quote
        function g($(argnames...))
            @allocated $(Expr(expr.head, argnames...))
        end
        $(Expr(:call, :g, [esc(a) for a in expr.args]...))
    end
end

const mvp_test = joinpath(dirname(pathof(MultivariatePolynomials)), "..", "test")
const Mod = TypedPolynomials

@testset "TypedPolynomials" begin
    # include("jump.jl")
    include("generators.jl")
    include("macros.jl")
    include("calculus.jl")
    include("promotion.jl")
    include("polynomials.jl")
    include("substitution.jl")
    include("sequences.jl")
end

@testset "MultivariatePolynomials interface tests" begin
    include(joinpath(mvp_test, "utils.jl"))
    include(joinpath(mvp_test, "commutativetests.jl"))
end

@testset "Demo" begin
    @nbinclude(joinpath(@__DIR__, "..", "demo.ipynb"))
end


text:
# auxiliary function to perform tests for action
function do_action_tests!(
    r :: Reduced_lattice,
    m :: Mesh,
    a :: Action
    ) :: Nothing

    # init test points on mesh
    w_idx  = rand(1 : 31)
    v_idx  = rand(1 : 31)
    vp_idx = rand(1 : 31)

    # init test points for interpolations
    ws  = rand() * m.Ωs[end]
    vs  = rand() * m.νs[end]
    vsp = rand() * m.νs[end]

    wt  = rand() * m.Ωt[end]
    vt  = rand() * m.νt[end]
    vtp = rand() * m.νt[end]

    wu  = rand() * m.Ωs[end]
    vu  = rand() * m.νs[end]
    vup = rand() * m.νs[end]

    # fill self energy with random values
    a.Σ .= rand(31)

    # ensure antisymmetry
    a.Σ[1] = 0.0

    # generate coefficients for test function
    a111 = rand(3, length(a.Γ), length(r.sites))
    a110 = rand(3, length(a.Γ), length(r.sites))
    a101 = rand(3, length(a.Γ), length(r.sites))
    a011 = rand(3, length(a.Γ), length(r.sites))
    a100 = rand(3, length(a.Γ), length(r.sites))
    a010 = rand(3, length(a.Γ), length(r.sites))
    a001 = rand(3, length(a.Γ), length(r.sites))
    a000 = rand(3, length(a.Γ), length(r.sites))

    # define testfunction (linear in all frequencies, should interpolate exactly)
    f(ch, i, s, w, v, vp) = a111[ch, i, s] * w * v * vp +
                            a110[ch, i, s] * w * v      +
                            a101[ch, i, s] * w     * vp +
                            a011[ch, i, s]     * v * vp +
                            a100[ch, i, s] * w          +
                            a010[ch, i, s]     * v      +
                            a001[ch, i, s]         * vp +
                            a000[ch, i, s]

    # fill channels with values of testfunction
    for i in eachindex(a.Γ)
        for s in 1 : length(r.sites)
            for iw in 1 : m.num_Ω
                a.Γ[i].ch_s.q1[s, iw] = f(1, i, s, m.Ωs[iw], m.νs[end], m.νs[end])
                a.Γ[i].ch_t.q1[s, iw] = f(2, i, s, m.Ωt[iw], m.νt[end], m.νt[end])
                a.Γ[i].ch_u.q1[s, iw] = f(3, i, s, m.Ωs[iw], m.νs[end], m.νs[end])

                for iv in 1 : m.num_ν
                    a.Γ[i].ch_s.q2_1[s, iw, iv] = f(1, i, s, m.Ωs[iw], m.νs[iv], m.νs[end])
                    a.Γ[i].ch_t.q2_1[s, iw, iv] = f(2, i, s, m.Ωt[iw], m.νt[iv], m.νt[end])
                    a.Γ[i].ch_u.q2_1[s, iw, iv] = f(3, i, s, m.Ωs[iw], m.νs[iv], m.νs[end])

                    a.Γ[i].ch_s.q2_2[s, iw, iv] = f(1, i, s, m.Ωs[iw], m.νs[end], m.νs[iv])
                    a.Γ[i].ch_t.q2_2[s, iw, iv] = f(2, i, s, m.Ωt[iw], m.νt[end], m.νt[iv])
                    a.Γ[i].ch_u.q2_2[s, iw, iv] = f(3, i, s, m.Ωs[iw], m.νs[end], m.νs[iv])

                    for ivp in 1 : m.num_ν
                        a.Γ[i].ch_s.q3[s, iw, iv, ivp] = f(1, i, s, m.Ωs[iw], m.νs[iv], m.νs[ivp])
                        a.Γ[i].ch_t.q3[s, iw, iv, ivp] = f(2, i, s, m.Ωt[iw], m.νt[iv], m.νt[ivp])
                        a.Γ[i].ch_u.q3[s, iw, iv, ivp] = f(3, i, s, m.Ωs[iw], m.νs[iv], m.νs[ivp])
                    end
                end
            end
        end
    end

    # test self energy interpolation and extrapolation
    @testset "interpolation / extrapolation Σ" begin
        outer = m.σ[end] + rand()
        @test get_Σ(m.σ[w_idx], m, a) ≈ a.Σ[w_idx]
        @test get_Σ(outer, m, a)      ≈ a.Σ[end] * m.σ[end] / outer
    end

    # prepare buffer for vectorized routines
    temp = zeros(Float64, length(r.sites), length(a.Γ), 1)

    # test interpolations in s channel
    @testset "interpolations s channel" begin
        @testset "on mesh" begin

            # build buffers
            bs_q3   = get_buffer_s(m.Ωs[w_idx], m.νs[v_idx], m.νs[vp_idx], m)
            bs_q2_2 = get_buffer_s(m.Ωs[w_idx],         Inf, m.νs[vp_idx], m)
            bs_q2_1 = get_buffer_s(m.Ωs[w_idx], m.νs[v_idx],          Inf, m)
            bs_q1   = get_buffer_s(m.Ωs[w_idx],         Inf,          Inf, m)
            bs_bare = get_buffer_s(        Inf,         Inf,          Inf, m)

            # test sequential routine
            @testset "sequential" begin
                idx = rand(1 : length(r.sites))

                for i in eachindex(a.Γ)
                    @test get_vertex(idx, bs_q3,   a.Γ[i], 1) ≈ a.Γ[i].ch_s.q3[idx, w_idx, v_idx, vp_idx]
                    @test get_vertex(idx, bs_q2_2, a.Γ[i], 1) ≈ a.Γ[i].ch_s.q2_2[idx, w_idx, vp_idx]
                    @test get_vertex(idx, bs_q2_1, a.Γ[i], 1) ≈ a.Γ[i].ch_s.q2_1[idx, w_idx, v_idx]
                    @test get_vertex(idx, bs_q1,   a.Γ[i], 1) ≈ a.Γ[i].ch_s.q1[idx, w_idx]
                    @test get_vertex(idx, bs_bare, a.Γ[i], 1) ≈ 0.0
                end
            end

            # test vectorized routine
            @testset "vectorized" begin
                for i in eachindex(a.Γ)
                    temp .= 0.0; get_vertex_avx!(r, bs_q3,   a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_s.q3[:, w_idx, v_idx, vp_idx]
                    temp .= 0.0; get_vertex_avx!(r, bs_q2_2, a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_s.q2_2[:, w_idx, vp_idx]
                    temp .= 0.0; get_vertex_avx!(r, bs_q2_1, a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_s.q2_1[:, w_idx, v_idx]
                    temp .= 0.0; get_vertex_avx!(r, bs_q1,   a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_s.q1[:, w_idx]
                    temp .= 0.0; get_vertex_avx!(r, bs_bare, a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test norm(temp[:, i, 1]) ≈ 0.0
                end
            end
        end

        @testset "interpolated" begin

            # build buffers
            bs_q3   = get_buffer_s( ws,  vs, vsp, m)
            bs_q2_2 = get_buffer_s( ws, Inf, vsp, m)
            bs_q2_1 = get_buffer_s( ws,  vs, Inf, m)
            bs_q1   = get_buffer_s( ws, Inf, Inf, m)
            bs_bare = get_buffer_s(Inf, Inf, Inf, m)

            # test sequential routine
            @testset "sequential" begin
                idx = rand(1 : length(r.sites))

                for i in eachindex(a.Γ)
                    @test get_vertex(idx, bs_q3,   a.Γ[i], 1) ≈ f(1, i, idx, ws, vs, vsp)
                    @test get_vertex(idx, bs_q2_2, a.Γ[i], 1) ≈ f(1, i, idx, ws, m.νs[end], vsp)
                    @test get_vertex(idx, bs_q2_1, a.Γ[i], 1) ≈ f(1, i, idx, ws, vs, m.νs[end])
                    @test get_vertex(idx, bs_q1,   a.Γ[i], 1) ≈ f(1, i, idx, ws, m.νs[end], m.νs[end])
                    @test get_vertex(idx, bs_bare, a.Γ[i], 1) ≈ 0.0
                end
            end

            # test vectorized routine
            @testset "vectorized" begin
                for i in eachindex(a.Γ)
                    temp .= 0.0; get_vertex_avx!(r, bs_q3,   a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(1, i, j, ws, vs, vsp)              for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bs_q2_2, a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(1, i, j, ws, m.νs[end], vsp)       for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bs_q2_1, a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(1, i, j, ws, vs, m.νs[end])        for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bs_q1,   a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(1, i, j, ws, m.νs[end], m.νs[end]) for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bs_bare, a.Γ[i], 1, view(temp, :, i, 1), false, 1.0); @test norm(temp[:, i, 1]) ≈ 0.0
                end
            end
        end
    end

    # test interpolations in t channel
    @testset "interpolations t channel" begin
        @testset "on mesh" begin

            # build buffers
            bt_q3   = get_buffer_t(m.Ωt[w_idx], m.νt[v_idx], m.νt[vp_idx], m)
            bt_q2_2 = get_buffer_t(m.Ωt[w_idx],         Inf, m.νt[vp_idx], m)
            bt_q2_1 = get_buffer_t(m.Ωt[w_idx], m.νt[v_idx],          Inf, m)
            bt_q1   = get_buffer_t(m.Ωt[w_idx],         Inf,          Inf, m)
            bt_bare = get_buffer_t(        Inf,         Inf,          Inf, m)

            # test sequential routine
            @testset "sequential" begin
                idx = rand(1 : length(r.sites))

                for i in eachindex(a.Γ)
                    @test get_vertex(idx, bt_q3,   a.Γ[i], 2) ≈ a.Γ[i].ch_t.q3[idx, w_idx, v_idx, vp_idx]
                    @test get_vertex(idx, bt_q2_2, a.Γ[i], 2) ≈ a.Γ[i].ch_t.q2_2[idx, w_idx, vp_idx]
                    @test get_vertex(idx, bt_q2_1, a.Γ[i], 2) ≈ a.Γ[i].ch_t.q2_1[idx, w_idx, v_idx]
                    @test get_vertex(idx, bt_q1,   a.Γ[i], 2) ≈ a.Γ[i].ch_t.q1[idx, w_idx]
                    @test get_vertex(idx, bt_bare, a.Γ[i], 2) ≈ 0.0
                end
            end

            # test vectorized routine
            @testset "vectorized" begin
                for i in eachindex(a.Γ)
                    temp .= 0.0; get_vertex_avx!(r, bt_q3,   a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_t.q3[:, w_idx, v_idx, vp_idx]
                    temp .= 0.0; get_vertex_avx!(r, bt_q2_2, a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_t.q2_2[:, w_idx, vp_idx]
                    temp .= 0.0; get_vertex_avx!(r, bt_q2_1, a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_t.q2_1[:, w_idx, v_idx]
                    temp .= 0.0; get_vertex_avx!(r, bt_q1,   a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_t.q1[:, w_idx]
                    temp .= 0.0; get_vertex_avx!(r, bt_bare, a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test norm(temp[:, i, 1]) ≈ 0.0
                end
            end
        end

        @testset "interpolated" begin

            # build buffers
            bt_q3   = get_buffer_t( wt,  vt, vtp, m)
            bt_q2_2 = get_buffer_t( wt, Inf, vtp, m)
            bt_q2_1 = get_buffer_t( wt,  vt, Inf, m)
            bt_q1   = get_buffer_t( wt, Inf, Inf, m)
            bt_bare = get_buffer_t(Inf, Inf, Inf, m)

            # test sequential routine
            @testset "sequential" begin
                idx = rand(1 : length(r.sites))

                for i in eachindex(a.Γ)
                    @test get_vertex(idx, bt_q3,   a.Γ[i], 2) ≈ f(2, i, idx, wt, vt, vtp)
                    @test get_vertex(idx, bt_q2_2, a.Γ[i], 2) ≈ f(2, i, idx, wt, m.νt[end], vtp)
                    @test get_vertex(idx, bt_q2_1, a.Γ[i], 2) ≈ f(2, i, idx, wt, vt, m.νt[end])
                    @test get_vertex(idx, bt_q1,   a.Γ[i], 2) ≈ f(2, i, idx, wt, m.νt[end], m.νt[end])
                    @test get_vertex(idx, bt_bare, a.Γ[i], 2) ≈ 0.0
                end
            end

            # test vectorized routine
            @testset "vectorized" begin
                for i in eachindex(a.Γ)
                    temp .= 0.0; get_vertex_avx!(r, bt_q3,   a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(2, i, j, wt, vt, vtp)              for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bt_q2_2, a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(2, i, j, wt, m.νt[end], vtp)       for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bt_q2_1, a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(2, i, j, wt, vt, m.νt[end])        for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bt_q1,   a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(2, i, j, wt, m.νt[end], m.νt[end]) for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bt_bare, a.Γ[i], 2, view(temp, :, i, 1), false, 1.0); @test norm(temp[:, i, 1]) ≈ 0.0
                end
            end
        end
    end

    # test interpolations in u channel
    @testset "interpolations u channel" begin
        @testset "on mesh" begin

            # build buffers
            bu_q3   = get_buffer_u(m.Ωs[w_idx], m.νs[v_idx], m.νs[vp_idx], m)
            bu_q2_2 = get_buffer_u(m.Ωs[w_idx],         Inf, m.νs[vp_idx], m)
            bu_q2_1 = get_buffer_u(m.Ωs[w_idx], m.νs[v_idx],          Inf, m)
            bu_q1   = get_buffer_u(m.Ωs[w_idx],         Inf,          Inf, m)
            bu_bare = get_buffer_u(        Inf,         Inf,          Inf, m)

            # test sequential routine
            @testset "sequential" begin
                idx = rand(1 : length(r.sites))

                for i in eachindex(a.Γ)
                    @test get_vertex(idx, bu_q3,   a.Γ[i], 3) ≈ a.Γ[i].ch_u.q3[idx, w_idx, v_idx, vp_idx]
                    @test get_vertex(idx, bu_q2_2, a.Γ[i], 3) ≈ a.Γ[i].ch_u.q2_2[idx, w_idx, vp_idx]
                    @test get_vertex(idx, bu_q2_1, a.Γ[i], 3) ≈ a.Γ[i].ch_u.q2_1[idx, w_idx, v_idx]
                    @test get_vertex(idx, bu_q1,   a.Γ[i], 3) ≈ a.Γ[i].ch_u.q1[idx, w_idx]
                    @test get_vertex(idx, bu_bare, a.Γ[i], 3) ≈ 0.0
                end
            end

            # test vectorized routine
            @testset "vectorized" begin
                for i in eachindex(a.Γ)
                    temp .= 0.0; get_vertex_avx!(r, bu_q3,   a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_u.q3[:, w_idx, v_idx, vp_idx]
                    temp .= 0.0; get_vertex_avx!(r, bu_q2_2, a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_u.q2_2[:, w_idx, vp_idx]
                    temp .= 0.0; get_vertex_avx!(r, bu_q2_1, a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_u.q2_1[:, w_idx, v_idx]
                    temp .= 0.0; get_vertex_avx!(r, bu_q1,   a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ a.Γ[i].ch_u.q1[:, w_idx]
                    temp .= 0.0; get_vertex_avx!(r, bu_bare, a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test norm(temp[:, i, 1]) ≈ 0.0
                end
            end
        end

        @testset "interpolated" begin

            # build buffers
            bu_q3   = get_buffer_u( wu,  vu, vup, m)
            bu_q2_2 = get_buffer_u( wu, Inf, vup, m)
            bu_q2_1 = get_buffer_u( wu,  vu, Inf, m)
            bu_q1   = get_buffer_u( wu, Inf, Inf, m)
            bu_bare = get_buffer_u(Inf, Inf, Inf, m)

            # test sequential routine
            @testset "sequential" begin
                idx = rand(1 : length(r.sites))

                for i in eachindex(a.Γ)
                    @test get_vertex(idx, bu_q3,   a.Γ[i], 3) ≈ f(3, i, idx, wu, vu, vup)
                    @test get_vertex(idx, bu_q2_2, a.Γ[i], 3) ≈ f(3, i, idx, wu, m.νs[end], vup)
                    @test get_vertex(idx, bu_q2_1, a.Γ[i], 3) ≈ f(3, i, idx, wu, vu, m.νs[end])
                    @test get_vertex(idx, bu_q1,   a.Γ[i], 3) ≈ f(3, i, idx, wu, m.νs[end], m.νs[end])
                    @test get_vertex(idx, bu_bare, a.Γ[i], 3) ≈ 0.0
                end
            end

            # test vectorized routine
            @testset "vectorized" begin
                for i in eachindex(a.Γ)
                    temp .= 0.0; get_vertex_avx!(r, bu_q3,   a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(3, i, j, wu, vu, vup)              for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bu_q2_2, a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(3, i, j, wu, m.νs[end], vup)       for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bu_q2_1, a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(3, i, j, wu, vu, m.νs[end])        for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bu_q1,   a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test temp[:, i, 1]       ≈ [f(3, i, j, wu, m.νs[end], m.νs[end]) for j in 1 : length(r.sites)]
                    temp .= 0.0; get_vertex_avx!(r, bu_bare, a.Γ[i], 3, view(temp, :, i, 1), false, 1.0); @test norm(temp[:, i, 1]) ≈ 0.0
                end
            end
        end
    end

    return nothing
end

"""
    test_action() :: Nothing

Run consistency checks for available action implementations by benchmarking interpolation routines.
"""
function test_action() :: Nothing

    # init test dummys
    list = get_mesh(rand(), 1.0, 30, 0.4)
    m    = Mesh(31, 31, 31, 31, list, list, list, list, list, list)

    # run tests for action_su2
    @testset "action su2" begin
        # generate action dummy for square lattice Heisenberg model
        l = get_lattice("square", 5, verbose = false)
        r = get_reduced_lattice("heisenberg", [[1.0]], l, verbose = false)
        a = get_action_empty("su2", r, m)
        init_action!(l, r, a)

        # test if bare action is correctly initialized
        @testset "initialization" begin
            @test norm(a.Γ[1].bare) ≈ 0.25
            @test norm(a.Γ[2].bare) ≈ 0.0
        end

        # perform tests
        do_action_tests!(r, m, a)
    end

    # run tests for action_u1_dm
    @testset "action u1-dm" begin
        # generate action dummy for triangular lattice dm-c3 model
        l = get_lattice("triangular", 5, verbose = false)
        r = get_reduced_lattice("triangular-dm-c3", [[1.0, 1.0, 1.0]], l, verbose = false)
        a = get_action_empty("u1-dm", r, m)
        init_action!(l, r, a)

        # test if bare action is correctly initialized
        @testset "initialization" begin
            @test norm(a.Γ[1].bare) ≈ sqrt(0.125)
            @test norm(a.Γ[2].bare) ≈ sqrt(0.125)
            @test norm(a.Γ[3].bare) ≈ sqrt(0.125)
            @test norm(a.Γ[4].bare) ≈ 0.0
            @test norm(a.Γ[5].bare) ≈ 0.0
            @test norm(a.Γ[6].bare) ≈ 0.0
        end

        # perform tests
        do_action_tests!(r, m, a)
    end

    return nothing
end

ext: jl


text:
# Unit testing for scanning functions

using NumericExtensions
using Base.Test

# verification function

function safe_scan(x, f)
    y = similar(x)
    y[1] = x[1]
    for i = 2:length(x)
        y[i] = f(y[i-1], x[i])
    end
    y
end

function safe_scan(x, f, dim)
    @assert ndims(x) <= 3
    y = similar(x)
    siz = size(x)
    if dim == 1
        n = prod(siz[2:end])
        for j=1:n
            y[:,j] = safe_scan(x[:,j], f)
        end
    elseif dim == 2
        @assert ndims(x) >= 2
        if ndims(x) == 2
            m = siz[1]
            for i=1:m
                y[i,:] = safe_scan(x[i,:], f)
            end
        else
            for k=1:siz[3]
                y[:,:,k] = safe_scan(x[:,:,k], f, 2)
            end
        end

    elseif dim == 3
        @assert ndims(x) == 3
        y[:,:,1] = x[:,:,1]
        for k=2:siz[3]
            y[:,:,k] = f(y[:,:,k-1], x[:,:,k])
        end

    else
        error("dim > 3 is not supported in safe_scan.")
    end
    y
end

# data

x1 = randn(6)
y1 = randn(6)
z1 = randn(6)

x2 = randn(5, 6)
y2 = randn(5, 6)
z2 = randn(5, 6)

x3 = randn(3, 4, 5)
y3 = randn(3, 4, 5)
z3 = randn(3, 4, 5)

# vector scan

r = safe_scan(x1, +)

@test size(cumsum(x1)) == size(x1)
@test_approx_eq cumsum(x1) r

x1c = copy(x1)
@test cumsum!(x1c) === x1c
@test_approx_eq x1c r

dst = zeros(size(x1))
cumsum!(dst, x1)
@test_approx_eq dst r

@test_approx_eq cummax(x1) safe_scan(x1, max)
@test_approx_eq cummin(x1) safe_scan(x1, min)

@test_approx_eq cumsum(Abs2Fun(), x1) safe_scan(abs2(x1), +)
@test_approx_eq cumsum(Multiply(), x1, y1) safe_scan(x1 .* y1, +)
@test_approx_eq cumsum(FMA(), x1, y1, z1) safe_scan(x1 + y1 .* z1, +)

dst = zeros(size(x1))
cumsum!(dst, Abs2Fun(), x1)
@test_approx_eq dst safe_scan(abs2(x1), +)

# scan along dimensions

@test size(cumsum(x1, 1)) == size(x1)
@test size(cumsum(x2, 1)) == size(x2)
@test size(cumsum(x2, 2)) == size(x2)
@test size(cumsum(x3, 1)) == size(x3)
@test size(cumsum(x3, 2)) == size(x3)
@test size(cumsum(x3, 3)) == size(x3)

@test_approx_eq cumsum(x1, 1) safe_scan(x1, +, 1)
@test_approx_eq cumsum(x2, 1) safe_scan(x2, +, 1)
@test_approx_eq cumsum(x2, 2) safe_scan(x2, +, 2)
@test_approx_eq cumsum(x3, 1) safe_scan(x3, +, 1)
@test_approx_eq cumsum(x3, 2) safe_scan(x3, +, 2)
@test_approx_eq cumsum(x3, 3) safe_scan(x3, +, 3)

@test_approx_eq cummax(x1, 1) safe_scan(x1, max, 1)
@test_approx_eq cummax(x2, 1) safe_scan(x2, max, 1)
@test_approx_eq cummax(x2, 2) safe_scan(x2, max, 2)
@test_approx_eq cummax(x3, 1) safe_scan(x3, max, 1)
@test_approx_eq cummax(x3, 2) safe_scan(x3, max, 2)
@test_approx_eq cummax(x3, 3) safe_scan(x3, max, 3)

@test_approx_eq cummin(x1, 1) safe_scan(x1, min, 1)
@test_approx_eq cummin(x2, 1) safe_scan(x2, min, 1)
@test_approx_eq cummin(x2, 2) safe_scan(x2, min, 2)
@test_approx_eq cummin(x3, 1) safe_scan(x3, min, 1)
@test_approx_eq cummin(x3, 2) safe_scan(x3, min, 2)
@test_approx_eq cummin(x3, 3) safe_scan(x3, min, 3)

@test_approx_eq cumsum(Abs2Fun(), x1, 1) safe_scan(abs2(x1), +, 1)
@test_approx_eq cumsum(Abs2Fun(), x2, 1) safe_scan(abs2(x2), +, 1)
@test_approx_eq cumsum(Abs2Fun(), x2, 2) safe_scan(abs2(x2), +, 2)
@test_approx_eq cumsum(Abs2Fun(), x3, 1) safe_scan(abs2(x3), +, 1)
@test_approx_eq cumsum(Abs2Fun(), x3, 2) safe_scan(abs2(x3), +, 2)
@test_approx_eq cumsum(Abs2Fun(), x3, 3) safe_scan(abs2(x3), +, 3)

@test_approx_eq cumsum(Multiply(), x1, y1, 1) safe_scan(x1 .* y1, +, 1)
@test_approx_eq cumsum(Multiply(), x2, y2, 1) safe_scan(x2 .* y2, +, 1)
@test_approx_eq cumsum(Multiply(), x2, y2, 2) safe_scan(x2 .* y2, +, 2)
@test_approx_eq cumsum(Multiply(), x3, y3, 1) safe_scan(x3 .* y3, +, 1)
@test_approx_eq cumsum(Multiply(), x3, y3, 2) safe_scan(x3 .* y3, +, 2)
@test_approx_eq cumsum(Multiply(), x3, y3, 3) safe_scan(x3 .* y3, +, 3)

x1c = copy(x1); cumsum!(x1c, 1)
@test_approx_eq x1c cumsum(x1, 1)

x2c = copy(x2); cumsum!(x2c, 1)
@test_approx_eq x2c cumsum(x2, 1)

x2c = copy(x2); cumsum!(x2c, 2)
@test_approx_eq x2c cumsum(x2, 2)

x3c = copy(x3); cumsum!(x3c, 1)
@test_approx_eq x3c cumsum(x3, 1)

x3c = copy(x3); cumsum!(x3c, 2)
@test_approx_eq x3c cumsum(x3, 2)

x3c = copy(x3); cumsum!(x3c, 3)
@test_approx_eq x3c cumsum(x3, 3)

dst = similar(x2)
cumsum!(dst, Abs2Fun(), x2, 1)
@test_approx_eq dst cumsum(abs2(x2), 1)


text:
module jlSPICE
using DataFrames, Gtk, PyPlot

include("RunAnalysis.jl")
include("Load_LTSpice_Net.jl")
include("SPICE2Matrix.jl")

export RunACAnalysis,
DetermineTempCo,
SPICE2Matrix,
getElementCurrents,
plotACElCurrent,
ProcessSPICE_DF,
SPICE_DF2Matrix_ω,
LTSpiceLoad,
MakeNumericalVals

# Write your package code here.

end


text:
function distance(s1::AbstractString, s2::AbstractString)
    count = 0
    if length(s1) != length(s2)
        return throw(ArgumentError("LOL"))
    else
        for i in 1:length(s1)
            println("$(s1[i]) == $(s2[i])")
            if s1[i] != s2[i]
                count += 1
            end
        end
        return count
    end

end

max_stars_repo_name: Devlious/LearningJulia

max_issues_repo_name: Devlious/LearningJulia

max_forks_repo_name: Devlious/LearningJulia


text:
include("../aimajulia.jl");

using Test;

using Main.aimajulia;

#The following game tests are from the aima-python doctest

@test minimax_decision("A", Figure52Game()) == "A1";

@test alphabeta_full_search("A", Figure52Game()) == "A1";

@test alphabeta_search("A", Figure52Game()) == "A1";

@test play_game(Figure52Game(), alphabeta_player, alphabeta_player) == 3;

#=

    The following tests may fail sometimes because the tests run on random behavior.

    However, the results of tests that fail does not imply something is wrong.

=#

function colorize_testv_doctest_results(result::Bool)
    if (result)
        printstyled("Test Passed\n", color=:green, bold=true);
    else
        printstyled("Test Failed\n", color=:red, bold=true);
    end
end

randf52_result = play_game(Figure52Game(), random_player, random_player);
colorize_testv_doctest_results(randf52_result == 6);
println("Expression: play_game(Figure52Game(), random_player, random_player) == 6");
println("Evaluated: ", randf52_result, " == 6");

randttt_result = play_game(TicTacToeGame(), random_player, random_player);
colorize_testv_doctest_results(randttt_result == 0);
println("Expression: play_game(TicTacToeGame(), random_player, random_player) == 0");
println("Evaluated: ", randttt_result, " == 0");

max_stars_repo_name: mikhail-j/aima-julia-aimacode-fork

max_issues_repo_name: mikhail-j/aima-julia-aimacode-fork

max_forks_repo_name: mikhail-j/aima-julia-aimacode-fork


text:
"""
    ReplicaStrategy{N}

An abstract type that controles how [`lomc!`](@ref) uses replicas. A subtype of
`ReplicaStrategy{N}` operates on `N` replicas and must implement the following function:

* [`replica_stats(::ReplicaStrategy{N}, ::NTuple{N,ReplicaState})`](@ref) - return a tuple
  of `String`s or `Symbols` of replica statistic names and a tuple of the values.  These
  will be reported to the `DataFrame` returned by [`lomc!`](@ref)

Concrete implementations:

* [`NoStats`](@ref): run (possibly one) replica(s), but don't report any additional info.
* [`AllOverlaps`](@ref): report overlaps between all pairs of replica vectors.

"""
abstract type ReplicaStrategy{N} end

num_replicas(::ReplicaStrategy{N}) where {N} = N

"""
    replica_stats(::ReplicaStrategy{N}, replicas::NTuple{N,ReplicaState}) -> (names, values)

Return the names and values of statistics reported by `ReplicaStrategy`. `names` should be
a tuple of `Symbol`s or `String`s and `values` should be a tuple of the same length.
"""
replica_stats

"""
    NoStats(N=1) <: ReplicaStrategy{N}

The default [`ReplicaStrategy`](@ref). `N` replicas are run, but no statistics are collected.
"""
struct NoStats{N} <: ReplicaStrategy{N} end
NoStats(N=1) = NoStats{N}()

replica_stats(::NoStats, _) = (), ()

"""
    AllOverlaps(n=2, operator=nothing) <: ReplicaStrategy{n}

Run `n` replicas and report overlaps between all pairs of replica vectors. If operator is
not `nothing`, the overlap `dot(c1, operator, c2)` is reported as well. If operator is a tuple
of operators, the overlaps are computed for all operators.

Column names in the report are of the form c{i}_dot_c{j} for vector-vector overlaps, and
c{i}_Op{k}_c{j} for operator overlaps.

See [`ReplicaStrategy`](@ref) and [`AbstractHamiltonian`](@ref) (for an interface for
implementing operators).
"""
struct AllOverlaps{N,O} <: ReplicaStrategy{N}
    operators::O
end

function AllOverlaps(num_replicas=2, operator=nothing)
    if isnothing(operator)
        operators = ()
    elseif operator isa Tuple
        operators = operator
    else
        operators = (operator,)
    end
    return AllOverlaps{num_replicas, typeof(operators)}(operators)
end

function replica_stats(rs::AllOverlaps, replicas::NTuple{N}) where {N}
    # Not using broadcasting because it wasn't inferred properly.
    vecs = ntuple(i -> replicas[i].v, Val(N))
    return all_overlaps(rs.operators, vecs)
end

"""
    all_overlaps(operators, vectors)

Get all overlaps between vectors and operators. This function is overlpaded for `MPIData`.
"""
function all_overlaps(operators::Tuple, vecs::NTuple{N,AbstractDVec}) where {N}
    T = promote_type((valtype(v) for v in vecs)..., eltype.(operators)...)
    names = String[]
    values = T[]
    for i in 1:N, j in i+1:N
        push!(names, "c$(i)_dot_c$(j)")
        push!(values, dot(vecs[i], vecs[j]))
        for (k, op) in enumerate(operators)
            push!(names, "c$(i)_Op$(k)_c$(j)")
            push!(values, dot(vecs[i], op, vecs[j]))
        end
    end

    num_reports = (N * (N - 1) ÷ 2) * (length(operators) + 1)
    return SVector{num_reports,String}(names).data, SVector{num_reports,T}(values).data
end


text:
import PhysicsEngine

@testset "Renderer" begin
end

max_stars_repo_name: jonathanbuchanan/PhysicsEngine

max_issues_repo_name: jonathanbuchanan/PhysicsEngine

max_forks_repo_name: jonathanbuchanan/PhysicsEngine


text:
# keep up to date with misc/toc.jmd
table_of_contents = [
    :precalc		:calculator;
    :precalc		:variables;
    :precalc		:numbers_types;
    :precalc		:logical_expressions;
    :precalc		:vectors;
    :precalc		:ranges;
    :precalc		:functions;
    :precalc		:plotting;
    :precalc		:transformations;
    :precalc		:inversefunctions;
    :precalc		:polynomial;
    :precalc		:polynomial_roots;
    :precalc		:polynomials_package;
    :precalc		:rational_functions;
    :precalc		:exp_log_functions;
    :precalc		:trig_functions;
    :precalc            :julia_overview;

    :limits		:limits;
    :limits		:limits_extensions;
    :limits		:continuity;
    :limits		:intermediate_value_theorem;

    :derivatives		:derivatives;
    :derivatives		:numeric_derivatives;
    :derivatives		:mean_value_theorem;
    :derivatives		:optimization;
    :derivatives		:first_second_derivatives;
    :derivatives		:curve_sketching;
    :derivatives		:linearization;
    :derivatives		:newtons_method;
    :derivatives		:more_zeros;
    :derivatives		:lhospitals_rule;
    :derivatives		:implicit_differentiation;
    :derivatives		:related_rates;
    :derivatives		:taylor_series_polynomials;

    :integrals		:area;
    :integrals		:ftc;
    :integrals		:substitution;
    :integrals		:integration_by_parts;
    :integrals		:partial_fractions;
    :integrals		:improper_integrals;
    :integrals		:mean_value_theorem;
    :integrals		:area_between_curves;
    :integrals		:center_of_mass;
    :integrals		:volumes_slice;
    :integrals		:arc_length;
    :integrals		:surface_area;

    :ODEs		:odes;
    :ODEs		:euler;
    :ODEs		:solve;
    :ODEs		:differential_equations;

    :differentiable_vector_calculus		:polar_coordinates;
    :differentiable_vector_calculus		:vectors;
    :differentiable_vector_calculus		:vector_valued_functions;
    :differentiable_vector_calculus		:scalar_functions;
    :differentiable_vector_calculus		:scalar_functions_applications;
    :differentiable_vector_calculus		:vector_fields;
    :differentiable_vector_calculus		:plots_plotting;
    :alternatives		:makie_plotting;

    :integral_vector_calculus		:double_triple_integrals;
    :integral_vector_calculus		:line_inegrals;
    :integral_vector_calculus		:div_grad_curl;
    :integral_vector_calculus		:stokes_theorem;
    :integral_vector_calculus		:review;

    :alternatives   :symbolics;
    :misc           :getting_started_with_juila
    :misc           :bibliography;
    :misc           :quick_notes;
    :misc           :julia_interfaces;
    :misc           :calculus_with_juila;
    :misc           :unicode
]
# return :d,:f for previous and next
function prev_next(d,f)
    vals = [df for df ∈ eachrow(table_of_contents)]
    val = [d,f]
    i = findfirst(Ref(val,) .== vals)
    i == nothing && error(val)
    i == 1 && return (prev=nothing, next=vals[2])
    i == length(vals) && return (prev=vals[end-1], next=nothing)
    return (prev=vals[i-1], next=vals[i+1])
end


text:
push!(LOAD_PATH, abspath("src"))
push!(LOAD_PATH, abspath("."))


text:
function solution()
  # SOLUTION GOES HERE
end

if length(ARGS) == 1 && ARGS[1] == "-a"
  println(solution())
  exit(0)
end

for line in eachline(STDIN)
  iters = int(line)

  start = time_ns()
  for i in 1:iters
    solution()
  end
  end_ = time_ns()

  println(end_ - start)
end


text:
export GCV, GCVMetadata

import StatsFuns: log2π

const UniNormalOrExpLinQuad = Union{UnivariateGaussianDistributionsFamily, ExponentialLinearQuadratic}

struct GCVMetadata{ A <: AbstractApproximationMethod }
    approximation :: A
end

get_approximation(meta::GCVMetadata) = meta.approximation

struct GCV end

@node GCV Stochastic [ y, x, z, κ, ω ]

const DefaultGCVNodeMetadata = GCVMetadata(GaussHermiteCubature(20))

default_meta(::Type{ GCV }) = DefaultGCVNodeMetadata

@average_energy GCV (q_y_x::MultivariateNormalDistributionsFamily, q_z::NormalDistributionsFamily, q_κ::Any, q_ω::Any) = begin
    y_x_mean, y_x_cov = mean_cov(q_y_x)
    z_mean, z_var     = mean_var(q_z)
    κ_mean, κ_var     = mean_var(q_κ)
    ω_mean, ω_var     = mean_var(q_ω)

    ksi = (κ_mean ^ 2) * z_var + κ_var * ((z_mean ^ 2) + z_var)
    psi = @inbounds (y_x_mean[2] - y_x_mean[1]) ^ 2 + y_x_cov[1, 1] + y_x_cov[2, 2] - y_x_cov[1, 2] - y_x_cov[2, 1]
    A = exp(-ω_mean + ω_var / 2)
    B = exp(-κ_mean * z_mean + ksi / 2)

    0.5 * log2π + 0.5 * (z_mean * κ_mean + ω_mean) + 0.5 * (psi * A * B)
end

@average_energy GCV (q_y::NormalDistributionsFamily, q_x::NormalDistributionsFamily, q_z::NormalDistributionsFamily, q_κ::Any, q_ω::Any) = begin
    y_mean, y_var = mean_var(q_y)
    x_mean, x_var = mean_var(q_x)
    z_mean, z_var = mean_var(q_z)
    κ_mean, κ_var = mean_var(q_κ)
    ω_mean, ω_var = mean_var(q_ω)

    ksi = (κ_mean ^ 2) * z_var + (z_mean ^ 2) * κ_var + κ_var * z_var
    psi = (y_mean - x_mean) ^ 2 + y_var+x_var
    A = exp(-ω_mean + ω_var / 2)
    B = exp(-κ_mean * z_mean + ksi / 2)

    0.5 * log2π + 0.5 * (z_mean * κ_mean + ω_mean) + 0.5 * (psi * A * B)
end


text:
#=
We shall say that an n-digit number is pandigital if it makes use of all the digits 1 to n exactly once; for example, the 5-digit number, 15234, is 1 through 5 pandigital.

The product 7254 is unusual, as the identity, 39 × 186 = 7254, containing multiplicand, multiplier, and product is 1 through 9 pandigital.

Find the sum of all products whose multiplicand/multiplier/product identity can be written as a 1 through 9 pandigital.

HINT: Some products can be obtained in more than one way so be sure to only include it once in your sum.
=#
include("projecteulerutils.jl")
using Combinatorics
function calc()
  matches = Int64[]
  for c in permutations(collect(1:9))
    c1,c2,c3 = c[1:2],c[3:5],c[6:9] # 2x3
    d1,d2,d3 = c[1:1],c[2:5], c[6:9] # 1x4
    if nd(c1) * nd(c2) == nd(c3)
      push!(matches, nd(c3))
    end
    if nd(d1) * nd(d2) == nd(d3)
      push!(matches, nd(d3))
    end
  end
  sum(unique(matches))
end
@time println(calc())

max_stars_repo_name: daniel-beard/JuliaProjectEuler

max_issues_repo_name: daniel-beard/JuliaProjectEuler

max_forks_repo_name: daniel-beard/JuliaProjectEuler


text:
module FRSHouseholdGetter

    #
    # This module retrieves the main dataset. The data is retrieved from CSV files and assembled once, including generating sample weights,
    # but it can then be accessed multiple times in a session. Retrieval is currently by index (1,2,3...) only but retrieval by sernum/datayear,
    # or by some sort of query interface might be added later.
    #

    using CSV
    using DataFrames: DataFrame

    using ScottishTaxBenefitModel

    using .Definitions

    using .ModelHousehold:
        Household,
        OneIndex,
        uprate!

    using .HouseholdFromFrame:
        load_hhld_from_frame

    using .RunSettings: Settings

    using .Weighting:
        generate_weights

    using .Uprating: load_prices

    export
        initialise,
        get_household,
        num_households,
        get_household_of_person

    ## See scripts/performance/hhld_example.jl for the rationalle behind this wrapper
    #
    # If you're using a getter to get hhlds, wrap the array of hhlds in a struct,
    # so you can give the
    # array a type and declare a constant. This aviods type instability which can murder
    # performance of the getter.
    #
    struct HHWrapper
        hhlds      :: Vector{Household{Float64}}
        weight     :: Vector{Float64}
        dimensions :: Vector{Int}
        hh_map     :: Dict{OneIndex,Int}
        pers_map   :: Dict{OneIndex,Int}
    end

    const MODEL_HOUSEHOLDS =
        HHWrapper(
            Vector{Household{Float64}}(undef, 0 ),
            zeros(Float64,0),
            zeros(Int,3),
            Dict{OneIndex,Int}(),
            Dict{OneIndex,Int}())

    """
    Initialise the dataset. If this has already been done, do nothing unless
    `reset` is true.
    return (number of households available, num people loaded inc. kids, num hhls in dataset (should always = item[1]))
    """
    function initialise(
            settings :: Settings,
            ;
            reset :: Bool = false ) :: Tuple

        global MODEL_HOUSEHOLDS
        nhh = size( MODEL_HOUSEHOLDS.hhlds )[1]
        if( nhh > 0 ) && ( ! reset )
            # weird syntax to make tuple from array;
            # see e.g. https://discourse.julialang.org/t/array-to-tuple/9024
            return (MODEL_HOUSEHOLDS.dimensions...,)
        end
        load_prices( settings )
        hh_dataset = CSV.File("$(MODEL_DATA_DIR)/$(settings.household_name).tab" ) |> DataFrame
        people_dataset = CSV.File("$(MODEL_DATA_DIR)/$(settings.people_name).tab") |> DataFrame
        npeople = size( people_dataset)[1]
        nhhlds = size( hh_dataset )[1]
        resize!( MODEL_HOUSEHOLDS.hhlds, nhhlds )
        resize!( MODEL_HOUSEHOLDS.weight, nhhlds )
        MODEL_HOUSEHOLDS.weight .= 0

        for hseq in 1:nhhlds
            hh = load_hhld_from_frame( hseq, hh_dataset[hseq,:], people_dataset, FRS )
            MODEL_HOUSEHOLDS.hhlds[hseq] = hh
            uprate!( hh )
            MODEL_HOUSEHOLDS.hh_map[OneIndex( hh.hid, hh.data_year )] = hseq
            for pid in keys(hh.people)
                MODEL_HOUSEHOLDS.pers_map[OneIndex( pid, hh.data_year )] = hseq
            end
        end
        # println( "made pers_map as $(MODEL_HOUSEHOLDS.pers_map)")
        @time weight = generate_weights( nhhlds)
        for i in eachindex( weight )
            MODEL_HOUSEHOLDS.weight[i] = weight[i]
        end
        MODEL_HOUSEHOLDS.dimensions.=
            size(MODEL_HOUSEHOLDS.hhlds)[1],
            npeople,
            nhhlds
        return (MODEL_HOUSEHOLDS.dimensions...,)
    end

    function get_household( pos :: Integer ) :: Household
        hh = MODEL_HOUSEHOLDS.hhlds[pos]
        hh.weight = MODEL_HOUSEHOLDS.weight[pos]
        return hh
    end

    function get_household( hid :: BigInt, datayear :: Int ) :: Household
        pos :: Int = MODEL_HOUSEHOLDS.hh_map[ OneIndex( hid, datayear) ]
        return get_household( pos )
    end

    function get_household_of_person( pid :: BigInt, datayear :: Int ) :: Union{Nothing,Household}
        pos = get( MODEL_HOUSEHOLDS.pers_map, OneIndex( pid, datayear), nothing )
        if pos === nothing
            return nothing
        end
        return get_household( pos )
    end

    function get_num_households()::Integer
        return size( MODEL_HOUSEHOLDS.hhlds )[1]
    end

end # module


text:
abstract type Action end
abstract type ByPermutations <: Action end
abstract type ByLinearTransformation <: Action end

abstract type InducedActionHomomorphism{A,T} end
#=
implements:
* basis(action_hom)
* action(hom::InducedActionHomomorphism) → Action
* action(hom::InducedActionHomomorphism, g::GroupElement, feature) → the action of `g` on `feature`
=#

"""
Compute the vector g⋅v which is the result of acting g on v by permutation.
"""
function action(
    hom::InducedActionHomomorphism{<:ByPermutations},
    g::GroupElement,
    v::AbstractVector,
)
    return v^induce(hom, g)
end

"""
Compute the vector g⋅v which is the result of acting g on v by linear transformation.
"""
function action(
    hom::InducedActionHomomorphism{<:ByLinearTransformation},
    g::GroupElement,
    v::AbstractVector,
)
    return induce(hom, g) * v
end

function induce(ac::Action, hom::InducedActionHomomorphism, g::GroupElement)
    return throw(
        "No fallback is provided for $(typeof(ac)). You need to implement `induce(::$(typeof(ac)), ::$(typeof(hom)), ::$(typeof(g)))`.",
    )
end

induce(hom::InducedActionHomomorphism, g::GroupElement) =
    induce(action(hom), hom, g)

decompose(x::Any, hom::InducedActionHomomorphism) = throw(
    "No fallback is provided for $(typeof(x)). You need to implement `decompose(::$(typeof(x)), ::$(typeof(hom)))`.",
)
coeff_type(ac::ByLinearTransformation) = throw(
    "No fallback is provided for $(typeof(ac)). You need to implement `_coeff_type(::$(typeof(ac)))`.",
)
coeff_type(hom::InducedActionHomomorphism) = coeff_type(action(hom))
coeff_type(::ByPermutations) = Int

# Exceeding typemax(UInt16) here would mean e.g. that you're trying to block-diagonalize
# an SDP constraint of size 65535×65535, which is highly unlikely ;)
_int_type(::Type{<:InducedActionHomomorphism}) = UInt16
_int_type(hom::InducedActionHomomorphism) = _int_type(typeof(hom))

"""
Compute the (sparse) matrix corresponding to g acting by linear transformation.
"""
function induce(
    ac::ByLinearTransformation,
    hom::InducedActionHomomorphism,
    g::GroupElement,
)
    n = length(basis(hom))

    I = Int[]
    J = Int[]
    V = coeff_type(ac)[]

    for (i, f) in enumerate(basis(hom))
        k = action(action(hom), g, f)
        idcs, vals = decompose(k, hom)
        append!(I, fill(i, length(idcs)))
        append!(J, idcs)
        append!(V, vals)
    end
    return sparse(I, J, V)
end

# ala permutation action
function decompose(m::T, hom::InducedActionHomomorphism{A,T}) where {A,T}
    return [hom[m]], [one(coeff_type(action(hom)))]
end

struct ExtensionHomomorphism{A<:Action,T,B<:StarAlgebras.AbstractBasis{T}} <:
       InducedActionHomomorphism{A,T}
    action::A
    basis::B
end

# see the comment about UInt16 above
ExtensionHomomorphism(action::Action, basis) =
    ExtensionHomomorphism(_int_type(ExtensionHomomorphism), action, basis)

ExtensionHomomorphism(::Type{I}, action::Action, basis) where {I<:Integer} =
    ExtensionHomomorphism(action, StarAlgebras.Basis{I}(basis))

_int_type(::Type{<:ExtensionHomomorphism{A,T,B}}) where {A,T,B} = _int_type(B)
_int_type(::Type{<:StarAlgebras.AbstractBasis{T,I}}) where {T,I} = I

Base.getindex(ehom::ExtensionHomomorphism, i::Integer) = basis(ehom)[i]
Base.getindex(ehom::ExtensionHomomorphism{A,T}, f::T) where {A,T} =
    ehom.basis[f]

# interface:
StarAlgebras.basis(hom::ExtensionHomomorphism) = hom.basis
action(hom::ExtensionHomomorphism) = hom.action

function induce(::ByPermutations, hom::ExtensionHomomorphism, g::GroupElement)
    I = _int_type(hom)
    return Perm(vec(I[hom[action(action(hom), g, f)] for f in basis(hom)]))
end

struct CachedExtensionHomomorphism{A,T,G,H,E<:InducedActionHomomorphism{A,T}} <:
       InducedActionHomomorphism{A,T}
    ehom::E
    cache::Dict{G,H}
end

action(h::CachedExtensionHomomorphism) = action(h.ehom)
StarAlgebras.basis(h::CachedExtensionHomomorphism) = basis(h.ehom)
_int_type(chom::CachedExtensionHomomorphism) = _int_type(h.ehom)

Base.getindex(h::CachedExtensionHomomorphism, x::Any) = h.ehom[x]

function CachedExtensionHomomorphism(
    G::Group,
    action::Action,
    basis;
    precompute = false,
)
    hom = ExtensionHomomorphism(action, basis)
    S = typeof(induce(hom, one(G)))
    chom = CachedExtensionHomomorphism{eltype(G),S}(hom)
    if precompute
        # to make sure that the future access to chom is read-only, i.e. thread-safe
        # one may choose to precompute the images
        for g in G
            induce(chom, g)
        end
    end
    return chom
end

CachedExtensionHomomorphism{G,H}(hom::InducedActionHomomorphism) where {G,H} =
    CachedExtensionHomomorphism(hom, Dict{G,H}())

function _induce(ac::Action, hom::CachedExtensionHomomorphism, g::GroupElement)
    if !haskey(hom.cache, g)
        hom.cache[g] = induce(ac, hom.ehom, g)
    end
    return hom.cache[g]
end

induce(ac::Action, hom::CachedExtensionHomomorphism, g::GroupElement) =
    _induce(ac, hom, g)
# disabmiguation:
induce(
    ac::ByLinearTransformation,
    hom::CachedExtensionHomomorphism,
    g::GroupElement,
) = _induce(ac, hom, g)


text:
#X(u)_function.jl

"""
Implementation of FHIP's susceptibility formula χ(Ω) for the Optical Absorption of Polarons and uses parts of Devresse's, Sitter's & Goovaerts paper that expands upon calculations in the FHIP paper (particularly for ℜχ(Ω)).

See FHIP 1962, equations (35a), (44), (47a, b, c):
https://link.aps.org/doi/10.1103/PhysRev.127.1004

and Devreese's et al. 1971 paper, equations (15), (16), and Appendix A:
https://link.aps.org/doi/10.1103/PhysRevB.5.2367.
"""

using QuadGK # Have to include this as this script somehow doesn't pick it up in the main PolaronMakie.jl module.

"""
----------------------------------------------------------------------
Simple implementation for finite temperatures. Useful for most values of arguments.
----------------------------------------------------------------------
"""

"""
ℜχ(Ω::Float64, β::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the real part of χ(Ω) (equation (35a) in FHIP) for a given frequency Ω. v and w are the variational Polaron parameters that minimise the free energy, for the supplied α Frohlich coupling. β is the reduced thermodynamical beta = ħω/kT with k being Boltzman's constant. This attempts to just solve the integral directly using adaptive gaussian quadrature methods.
"""
function ℜχ(Ω, β, α, v, w)
    integrand(x) = (1 - cos(Ω * x)) * ℑS(x, 0.0, v, w, β, α)
    result = QuadGK.quadgk(x -> integrand(x), 0.0, Inf; atol = 1e-3)
    println(Ω, " ", result)
    return result
end

"""
ℑχ(Ω::Float64, β::Float64, α::Float64, v::Float64, w::Float64)

    Calculate the imaginary part of χ(Ω) (equations (35a) and (44) in FHIP) for a given frequency Ω. v and w are the variational Polaron parameters that minimise the free energy, for the supplied α Frohlich coupling. β is the reduced thermodynamical beta = ħω/kT with k being Boltzman's constant. This attempts to just solve the integral directly using adaptive gaussian quadrature methods.
"""
function ℑχ(Ω, β, α, v, w)
    integrand(x) = sin(Ω * x) * ℑS(x, 0.0, v, w, β, α)
    result = QuadGK.quadgk(x -> integrand(x), 0.0, Inf; atol = 1e-3)
    println(Ω, " ", result)
    return result
end

Ω_range = 0.01:0.1:10.01
ReX = [ℜχ(Ω, 0.4, 5, 4, 2.8) for Ω in Ω_range]
ImX = [ℑχ(Ω, 0.4, α, v, w) for Ω in Ω_range]
using Plots
plotly()
Plots.PlotlyBackend()
p = plot(Ω_range, [i[1] for i in ReX], ribbon = [i[2] for i in ReX], label = "ℜχ", xlabel = "Ω", ylabel = "χ")
plot!(p, Ω_range, [i[1] for i in ImX], ribbon = [i[2] for i in ImX], label = "ℑχ")
display(p)

# Create a complex number from the real and imaginary parts of χ(Ω).

function χ(Ω, β, α, v, w)

    R = (v^2 - w^2) / (w^2 * v)
    a_squared = β^2 / 4 + R * β * coth(β * v / 2)
    b = R * β / sinh(β * v / 2)

    D(x, y) = w^2 * (a_squared - β^2 / 4 - b * cos(v * x) * cosh(v * (y - β / 2)) + x^2 + y * (β - y)) / (β * v^2) + 1im * (w^2 * (b * sin(v * x) * sinh(v * (y - β / 2)) + 2 * x * (y - β / 2)) / (β * v^2))

    θ(x, y) = angle(D(x, y))
    r_squared(x, y) = abs2(D(x, y))

    S(x, y) = 2 * α * (cos(3 * θ(x, y) / 2) * cos(x) * cosh(y - β / 2) - sin(3 * θ(x, y) / 2) * sin(x) * sinh(y - β / 2)) / (3 * sqrt(π) * sinh(β / 2) * r_squared(x, y)^(3 / 4)) -2im * α * (cos(3 * θ(x, y) / 2) * sin(x) * sinh(y - β / 2) + sin(3 * θ(x, y) / 2) * cos(x) * cosh(y - β / 2)) / (3 * sqrt(π) * sinh(β / 2) * r_squared(x, y)^(3 / 4))

    integrand(x) = (1 - exp(1im * Ω * x)) * imag(S(x, 0.0))
    result = QuadGK.quadgk(x -> integrand(x), 0.0, Inf; atol = 1e-3)[1]
    println(Ω, " ", result)
    return result
end

function Γ(Ω, β, α, v, w)
    z = 1im * (Ω - χ(Ω, β, α, v, w) / Ω)
    abs(1 / z)
end


text:
import Base64.base64encode
import Colors.hex

export to_svg

"""
    to_path_string(path)

Returns SVG path syntax for the given path.
"""
function to_path_string(path::Path)
    string("M", join(["$(point[1]),$(point[2])" for point in path], "L"))
end

"""
	to_svg(plot)

Returns an SVG representation of the given plot as a string.
"""
function to_svg(plot::PenPlot)
    repr(MIME"image/svg+xml"(), plot)
end

function Base.show(io::IO, ::MIME"image/svg+xml", plot::PenPlot)
    ext = extent(plot)
    width = ext.lowerright[1] - ext.upperleft[1]
    height = ext.lowerright[2] - ext.upperleft[2]

    print(
        io,
        """<svg
xmlns="http://www.w3.org/2000/svg"
height="400"
width="100%"
viewBox="$(ext.upperleft[1]) $(ext.upperleft[2]) $(width) $(height)">""",
    )

    for layer in plot.layers
        print(
            io,
            """<g id="$(layer.name)"
stroke="#$(hex(layer.color))"
opacity="0.5"
fill="none">""",
        )

        for path in layer.paths
            print(
                io,
                """<path
  d="$(to_path_string(path))"
  vector-effect="non-scaling-stroke"
  />""",
            )
        end

        print(io, "</g>")
    end

    print(io, "</svg>")
end

function Base.show(io::IO, mime::MIME"image/svg+xml", path::Path)
    plot = PenPlot([path])
    Base.show(io, mime, plot)
end

function Base.show(io::IO, mime::MIME"image/svg+xml", paths::MultiPath)
    plot = PenPlot(paths)
    Base.show(io, mime, plot)
end


text:
using CFITSIO
using Documenter
DocMeta.setdocmeta!(CFITSIO, :DocTestSetup, :(using CFITSIO); recursive=true)

makedocs(;
    modules=[CFITSIO],
    authors="JuliaAstro",
    repo="https://github.com/JuliaAstro/CFITSIO.jl/blob/{commit}{path}#L{line}",
    sitename="CFITSIO.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://juliaastro.github.io/CFITSIO.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/JuliaAstro/CFITSIO.jl",
)


text:
"""
`Basics.determineEnergySharings(energy::Float64, NoEnergySharings::Int64)`
    ... to determine the NoEnergySharings sharings by using the Gauss-Legendre integration points in the interval
        [0., energy]; an Array{Tuple{Float64,Float64},1} is returned. This methods assumes that the package
        GaussQuadrature is 'used'.
"""
function Basics.determineEnergySharings(energy::Float64, NoEnergySharings::Int64)
    xx, ww = GaussQuadrature.legendre(NoEnergySharings)
    sharings = Tuple{Float64,Float64}[]
    for  x in xx
        xrel = (x+1.) / 2
        push!(sharings, (xrel*energy, energy - xrel*energy) )
    end
    return( sharings )
end

"""
`Basics.determineHoleShells(conf::Configuration)`
    ... to determine the hole-shells of a given non-relativistic configuration; a shellList::Array{Shell,1} is returned.
"""
function Basics.determineHoleShells(conf::Configuration)

    shellList = Shell[]
    for  (k,v) in conf.shells
       if   conf.shells[k] < 2*(2k.l+1)   push!(shellList, k)    end
    end

    return( shellList )
end

"""
`Basics.determineParity(conf::Configuration)`  ... to determine the parity of a given non-relativistic configuration.
"""
function Basics.determineParity(conf::Configuration)

    par = 1
    for  (k,v) in conf.shells
       if   iseven(k.l)    p = 1   else   p = -1    end
       par = par * (p^v)
    end

    if       par == 1    return( Basics.plus )
    elseif   par == -1   return( Basics.minus )
    else     error("stop b")
    end
end

"""
`Basics.determineParity(conf::ConfigurationR)`  ... to determine the parity of a given relativistic configuration.
"""
function Basics.determineParity(conf::ConfigurationR)

    par = 1
    for  (k,v) in conf.subshells
       l = Basics.subshell_l(k)
       if   iseven(l)    p = 1   else   p = -1    end
       par = par * (p^v)
    end

    if       par == 1    return( Basics.plus )
    elseif   par == -1   return( Basics.minus )
    else     error("stop b")
    end
end

"""
`Basics.determineSelectedLines(lineList::Array{Tuple{Int64,Int64},1}, initialMultiplet::Multiplet, finalMultiplet::Multiplet)`  ... to determine
                            the specified lines as tupels of initial- and final levels. A level index is 0 in lineList always refers to
                            all levels from the corresponding multiplet. A (unique) newList::Array{Tuple{Int64,Int64},1} is returned.
"""
function Basics.determineSelectedLines(lineList::Array{Tuple{Int64,Int64},1}, initialMultiplet::Multiplet, finalMultiplet::Multiplet)
    newList = Tuple{Int64,Int64}[]
    for tupl in lineList
        if      tupl[1] >  length(initialMultiplet.levels)   error("Initial level index $(tupl[1]) must be <= No. initial levels.")
        elseif  tupl[2] >  length(finalMultiplet.levels)     error("Final level index $(tupl[2]) must be <= No. final levels.")
        elseif  tupl[1] == 0  &&   tupl[2] == 0    error("Tupel (0,0) is not allowed.")
        elseif  tupl[1] == 0       for i = 1:length(initialMultiplet.levels)   push!(newList, (i, tupl[2]))   end
        elseif  tupl[2] == 0       for i = 1:length(finalMultiplet.levels)     push!(newList, (tupl[1], i))   end
        else                                                                   push!(newList, (tupl[1], tupl[2]))
        end
    end
    newList = unique(newList)
    return( newList)
end

"""
`Basics.determineSelectedPathways(pathwayList::Array{Tuple{Int64,Int64,Int64},1}, initialMultiplet::Multiplet, intermediateMultiplet::Multiplet,
                               finalMultiplet::Multiplet)`  ... to determine the specified pathways as tupels of initial-, intermediate-
                               and final levels. A level index 0 in pathwayList always refers to all levels from the corresponding multiplet.
                               A (unique) newList::Array{Tuple{Int64,Int64,Int64},1} is returned.
"""
function Basics.determineSelectedPathways(pathwayList::Array{Tuple{Int64,Int64,Int64},1}, initialMultiplet::Multiplet, intermediateMultiplet::Multiplet,
                                   finalMultiplet::Multiplet)
    newList = Tuple{Int64,Int64,Int64}[]
    for tupl in pathwayList
        if      tupl[1] >  length(initialMultiplet.levels)       error("Initial level index $(tupl[1]) must be <= No. initial levels.")
        elseif  tupl[2] >  length(intermediateMultiplet.levels)  error("Intermediate level index $(tupl[2]) must be <= No. intermediate levels.")
        elseif  tupl[3] >  length(finalMultiplet.levels)         error("Final level index $(tupl[3]) must be <= No. final levels.")
        elseif  (tupl[1] + tupl[2] == 0  || tupl[1] + tupl[3] == 0  ||  tupl[2] + tupl[3] == 0)   &&  error("Only tupels with one 0-index are allowed.")
        elseif  tupl[1] == 0       for i = 1:length(initialMultiplet.levels)          push!(newList, (i, tupl[2], tupl[3]))   end
        elseif  tupl[2] == 0       for i = 1:length(intermediateMultiplet.levels)     push!(newList, (tupl[1], i, tupl[3]))   end
        elseif  tupl[3] == 0       for i = 1:length(finalMultiplet.levels)            push!(newList, (tupl[1], tupl[2], i))   end
        else                                                                          push!(newList, (tupl[1], tupl[2], tupl[3]))
        end
    end
    newList = unique(newList)
    return( newList)
end


text:
using DifferentiableNAS
using Distributions
using Flux
using CUDA

@testset "Depthwise test" begin
    layer = DepthwiseConv((3,3),3=>3,stride=1,pad=1,bias=false) |> gpu
    test_image = rand(Float32, 32, 32, 3, 1) |> gpu
    params = Flux.params(layer)
    grads = gradient(params) do
       sum(layer(test_image))
    end
    @show [grads.grads[param] for param in params]
end
conv_answer_dict = Dict(
    # Known-good answers for 1d convolution operations
    1 => Dict(
        "y_pad"  => [1, 4,  7, 10, 13, 10.],
        "y_dil"  => [5, 8, 11.],
        "y_flip" => [5, 8, 11, 14.],

        "dx"        => [ 8, 18, 27, 36, 13.],
        "dx_stride" => [ 8,  4, 20, 10,  0.],
        "dx_pad"    => [ 9, 18, 27, 36, 33.],
        "dx_dil"    => [10, 16, 27,  8, 11.],
        "dx_flip"   => [ 5, 18, 27, 36, 28.],

        "dw"        => [134, 100.],
        "dw_stride" => [ 48,  34.],
        "dw_pad"    => [135, 150.],
        "dw_dil"    => [102,  54.],
        "dw_flip"   => [110, 148.],
    ),

    # Known-good answers for 2d convolution operations
    2 => Dict(
        "y_pad" => [
            1  9  29  49  48;
            4 29  79 129 115;
            7 39  89 139 122;
            10 49  99 149 129;
            13 59 109 159 136;
            10 40  70 100  80.
        ],
        "y_dil" => [
            48   98;
            58  108;
            68  118.
        ],
        "y_flip" => [
            51  101  151;
            61  111  161;
            71  121  171;
            81  131  181.
        ],

        "dx" => [
            116  374   674  258;
            243  700  1200  407;
            313  800  1300  437;
            383  900  1400  467;
            177  386   586  159.
        ],
        "dx_stride" => [
            116  58  516  258;
            87  29  387  129;
            196  98  596  298;
            147  49  447  149;
            0   0    0    0.
        ],
        "dx_pad" => [
            152  470   850   911;
            261  700  1200  1240;
            340  800  1300  1319;
            419  900  1400  1398;
            370  746  1126  1087.
        ],
        "dx_dil" => [
            192  392   96  196;
            232  432  116  216;
            416  766  184  334;
            174  324   58  108;
            204  354   68  118.
        ],
        "dx_flip" => [
            51  254   454   453;
            163  700  1200  1087;
            193  800  1300  1157;
            223  900  1400  1227;
            162  586   886   724.
        ],

        "dw" => [
            17378  11738;
            16250  10610.
        ],
        "dw_stride" => [
            5668  3888;
            5312  3532.
        ],
        "dw_pad" => [
            18670  22550;
            19850  23430.
        ],
        "dw_dil" => [
            8632  3652;
            7636  2656.
        ],
        "dw_flip" => [
            12590  19550;
            13982  20942.
        ],
    ),

    # Known-good answers for 3d convolution operations (these are getting rather large)
    3 => Dict(
        "y_pad"  => reshape([
            1, 4, 7, 10, 13, 10, 9, 29, 39, 49, 59, 40, 29, 79, 89, 99, 109, 70, 49, 129,
            139, 149, 159, 100, 48, 115, 122, 129, 136, 80, 26, 80, 94, 108, 122, 80, 126,
            322, 358, 394, 430, 260, 206, 502, 538, 574, 610, 360, 286, 682, 718, 754, 790,
            460, 220, 502, 524, 546, 568, 320, 146, 360, 374, 388, 402, 240, 446, 1042, 1078,
            1114, 1150, 660, 526, 1222, 1258, 1294, 1330, 760, 606, 1402, 1438, 1474, 1510,
            860, 420, 942, 964, 986, 1008, 560, 205, 456, 467, 478, 489, 270, 517, 1133, 1159,
            1185, 1211, 660, 577, 1263, 1289, 1315, 1341, 730, 637, 1393, 1419, 1445, 1471,
            800, 392, 847, 862, 877, 892, 480.
        ], (6,5,4)),
        "y_dil"  => reshape([608, 644, 680, 788, 824, 860.], (3,2,1)),
        "y_flip" => reshape([
            686, 722, 758, 794, 866, 902, 938, 974, 1046, 1082, 1118, 1154, 1406, 1442,
            1478, 1514, 1586, 1622, 1658, 1694, 1766, 1802, 1838, 1874.
        ], (4,3,2)),

        "dx"        => reshape([
            2576, 5118, 5658, 6198, 3010, 5948, 11576, 12512, 13448, 6420, 8468, 16256,
            17192, 18128, 8580, 4092, 7718, 8114, 8510, 3950, 9624, 18316, 19108, 19900,
            9340, 18680, 34992, 36288, 37584, 17320, 22280, 41472, 42768, 44064, 20200,
            9776, 17756, 18260, 18764, 8340, 4168, 7438, 7690, 7942, 3450, 6972, 11896,
            12256, 12616, 5140, 8052, 13696, 14056, 14416, 5860, 2804, 4278, 4386, 4494,
            1510.
        ], (5,4,3)),
        "dx_stride" => reshape([
            2576, 2254, 3152, 2758, 0, 1932, 1610, 2364, 1970, 0, 5456, 4774, 6032,
            5278, 0, 4092, 3410, 4524, 3770, 0, 1288, 966, 1576, 1182, 0, 644, 322,
            788, 394, 0, 2728, 2046, 3016, 2262, 0, 1364, 682, 1508, 754, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.
        ], (5,4,3)),
        "dx_pad"    => reshape([
            4220, 6343, 7116, 7889, 6550, 8490, 12276, 13312, 14348, 11606, 12350,
            17456, 18492, 19528, 15546, 11989, 16664, 17469, 18274, 14333, 16200,
            22628, 23616, 24604, 19392, 25336, 34992, 36288, 37584, 29320, 30216,
            41472, 42768, 44064, 34200, 26236, 35664, 36652, 37640, 28940, 22816,
            30831, 31636, 32441, 24794, 32522, 43668, 44704, 45740, 34742, 36462,
            48848, 49884, 50920, 38602, 29501, 39264, 40037, 40810, 30733.
        ], (5,4,3)),
        "dx_dil"    => reshape([
            4864, 5152, 9696, 4508, 4760, 6304, 6592, 12396, 5768, 6020, 3648,
            3864, 7120, 3220, 3400, 4728, 4944, 9100, 4120, 4300, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2432, 2576, 4544, 1932, 2040,
            3152, 3296, 5804, 2472, 2580, 1216, 1288, 1968, 644, 680, 1576, 1648,
            2508, 824, 860.
        ], (5,4,3)),
        "dx_flip"   => reshape([
            686, 2094, 2202, 2310, 1588, 2924, 7544, 7904, 8264, 5124, 3644, 9344,
            9704, 10064, 6204, 3138, 7430, 7682, 7934, 4616, 4836, 11980, 12484,
            12988, 7792, 14936, 34992, 36288, 37584, 21640, 17816, 41472, 42768,
            44064, 25240, 12620, 28412, 29204, 29996, 16728, 7030, 15646, 16042,
            16438, 9084, 17772, 38968, 39904, 40840, 22276, 19932, 43648, 44584,
            45520, 24796, 12362, 26742, 27282, 27822, 14992.
        ], (5,4,3)),

        "dw"        => reshape([1.058184e6, 1.0362e6,    948264,    926280,
                                    618504,   596520,    508584,    486600], (2,2,2)),
        "dw_stride" => reshape([    74760,     72608,     64000,     61848,
                                    31720,     29568,     20960,     18808.], (2,2,2)),
        "dw_pad"    => reshape([1.26055e6, 1.30805e6, 1.40327e6, 1.44923e6,
                                1.73731e6, 1.77589e6, 1.83259e6, 1.86731e6], (2,2,2)),
        "dw_dil"    => reshape([   250320,    241512,    206280,    197472,
                                    74160,     65352,     30120,     21312.], (2,2,2)),
        "dw_flip"   => reshape([    639480,   670200,    793080,    823800,
                                    1.25388e6, 1.2846e6, 1.40748e6,  1.4382e6], (2,2,2)),
    ),
)

@testset "Depthwise Convolution" begin
    # Start with some easy-to-debug cases that we have worked through and _know_ work
    for rank in (2,3)
        @testset "depthwiseconv$(rank)d" begin
            # Pull out known-good answers for y = depthwiseconv(x, w)
            y_pad = conv_answer_dict[rank]["y_pad"]  |> gpu
            y_dil = conv_answer_dict[rank]["y_dil"]  |> gpu
            y_flip = conv_answer_dict[rank]["y_flip"]  |> gpu

            # We can always derive y_plain and y_stride from the other answers.
            y_plain = y_pad[((2:(size(y_pad,idx)-1)) for idx in 1:rank)...]  |> gpu
            y_stride = y_pad[((2:2:(size(y_pad,idx)-1)) for idx in 1:rank)...]  |> gpu

            # Same for dx and dw:
            dx = conv_answer_dict[rank]["dx"]  |> gpu
            dx_stride = conv_answer_dict[rank]["dx_stride"]  |> gpu
            dx_pad = conv_answer_dict[rank]["dx_pad"]  |> gpu
            dx_dil = conv_answer_dict[rank]["dx_dil"]  |> gpu
            dx_flip = conv_answer_dict[rank]["dx_flip"]  |> gpu

            dw = conv_answer_dict[rank]["dw"] |> gpu
            dw_stride = conv_answer_dict[rank]["dw_stride"]  |> gpu
            dw_pad = conv_answer_dict[rank]["dw_pad"]  |> gpu
            dw_dil = conv_answer_dict[rank]["dw_dil"]  |> gpu
            dw_flip = conv_answer_dict[rank]["dw_flip"]  |> gpu

            # We generate x and w from the shapes we know they must be
            x = reshape(Float64[1:prod(size(dx));], size(dx)..., 1, 1)  |> gpu
            w = reshape(Float64[1:prod(size(dw));], size(dw)..., 1, 1)  |> gpu

            # A "drop channels and batch dimension" helper
            ddims(x) = dropdims(x, dims=(rank+1, rank+2))

            for conv in (depthwiseconv,)
                @testset "$(conv)" begin
                    # First, your basic convolution with no parameters
                    cdims = DepthwiseConvDims(x, w)
                    @test ddims(conv(x, w, cdims)) == y_plain

                    # Next, test convolution on views and alternate datatypes:
                    @test isapprox(ddims(conv(Float32.(x), Float32.(w), cdims)), Float32.(y_plain), rtol = 1.0e-7)

                    # Next, introduce stride:
                    cdims = DepthwiseConvDims(x, w; stride=2)
                    @test isapprox(ddims(conv(x, w, cdims)), y_stride, rtol = 1.0e-7)

                    # Next, introduce dilation:
                    cdims = DepthwiseConvDims(x, w; dilation=2)
                    @test isapprox(ddims(conv(x, w, cdims)), y_dil, rtol = 1.0e-7)

                    # Next, introduce padding:
                    cdims = DepthwiseConvDims(x, w; padding=1)
                    @test isapprox(ddims(conv(x, w, cdims)), y_pad, rtol = 1.0e-7)
                end
            end

            # Test all implementations/interfaces
            for (∇conv_filter, ∇conv_data) in (
                    (∇depthwiseconv_filter,        ∇depthwiseconv_data),
                )
                @testset "$(∇conv_filter)/$(∇conv_data)" begin
                    # First, your basic convolution with no parameters
                    cdims = DepthwiseConvDims(x, w)
                    dy = depthwiseconv(x, w, cdims)
                    @test ddims(∇conv_filter(x, dy, cdims)) == dw
                    @test ddims(∇conv_data(dy, w,  cdims)) == dx

                    # Next, test convolution on views and alternate datatypes:
                    @test ddims(∇conv_filter(x, view(dy, repeat([:], ndims(dy))...), cdims)) == dw
                    @test ddims(∇conv_data(view(dy, repeat([:], ndims(dy))...), w,   cdims)) == dx

                    @test ddims(∇conv_filter(Float32.(x), Float32.(dy), cdims)) == dw
                    @test ddims(∇conv_data(Float32.(dy),  Float32.(w),  cdims)) == dx

                    # Next, introduce stride:
                    cdims = DepthwiseConvDims(x, w; stride=2)
                    dy = depthwiseconv(x, w, cdims)
                    @test ddims(∇conv_filter(x, dy, cdims)) == dw_stride
                    @test ddims(∇conv_data(dy, w,  cdims)) == dx_stride

                    # Next, introduce dilation:
                    cdims = DepthwiseConvDims(x, w; dilation=2)
                    dy = depthwiseconv(x, w, cdims)
                    @test ddims(∇conv_filter(x, dy, cdims)) == dw_dil
                    @test ddims(∇conv_data(dy, w,  cdims)) == dx_dil

                    # Next, introduce padding:
                    cdims = DepthwiseConvDims(x, w; padding=1)
                    dy = depthwiseconv(x, w, cdims)
                    @test ddims(∇conv_filter(x, dy, cdims)) == dw_pad
                    @test ddims(∇conv_data(dy, w,  cdims)) == dx_pad
                end
            end
        end
    end
end

@testset "Flip test" begin
    batch = rand(Float32,32,32,3,8)
    orig = copy(batch)
    flips = flip_batch!(batch)
    for i=1:size(batch,4)
        if flips[i]
            @test orig[:,:,:,i]==batch[:,end:-1:1,:,i]
        else
            @test orig[:,:,:,i]==batch[:,:,:,i]
        end
    end
end

@testset "Shift test" begin
    batch = rand(Float32,32,32,3,8)
    orig = copy(batch)
    shifts = shift_batch!(batch)
    for i=1:size(batch,4)
        shiftx = shifts[i,1]
        shifty = shifts[i,2]
        @test batch[5-shiftx:28-shiftx,5-shifty:28-shifty,:,i]==orig[5:28,5:28,:,i]
        #@test batch[:,:,:,i]
    end
end

@testset "Cutout test" begin
    batch = rand(Float32,32,32,3,8)
    orig = copy(batch)
    cutouts = cutout_batch!(batch, 16)
end

max_stars_repo_name: huangyxi/DifferentiableNAS

max_issues_repo_name: huangyxi/DifferentiableNAS

max_forks_repo_name: huangyxi/DifferentiableNAS


text:
struct InputBuffer <: Group{:X}
    X::Vector{Float32}
end

# ==================================
# Leaky integrate-and-fire neuron
# ==================================

@with_kw struct LIF <: Group{:u}
    u::Vector{Float32}
    y::Vector{Float32}
    θ::Vector{Float32}
    I_ff::Vector{Float32}
    I_rc::Vector{Float32}
    n::Vector{Float32}
    dt::Float32 = 0.1
    p::Float32 = 0.05
    γ::Float32 = 0.01
end

LIF(N::Int; args...) = LIF(;
    u = fill(0.0f0, N),
    y = fill(0.0f0, N),
    θ = fill(5.0f0, N),
    I_ff = fill(0.0f0, N),
    I_rc = fill(0.0f0, N),
    n = fill(0.0f0, N),
    args...
)

@generated on_reset(lif::Elem{LIF}) = @resolve(quote
    u = zero(u)
    y = zero(y)
    I_ff = zero(I_ff)
    I_rc = zero(I_rc)
    n = zero(n)
end, lif)

@generated on_step(lif::Elem{LIF}) = @resolve(quote
    # Spiking and after-spike reset:
    if u > θ
        u = zero(u)
        y = one(y)
        n += one(n)
    else
        y = zero(y)
    end

    # Membrane potential:
    du = I_ff - I_rc - u
    u += du * dt

    # Prepare to integrate I_rc again:
    I_rc = zero(I_rc)
    # I_ff only changes when the stimulus changes

end, lif)

@generated has_spike(lif::Elem{LIF}) = @resolve(:(y > zero(y)), lif)

@generated on_learn(lif::Elem{LIF}) = @resolve(quote
    # Threshold adaptation:
    dθ = n - p
    θ += dθ * γ
end, lif)

# ===================
# Synaptic weights
# ===================

@with_kw struct Synapses{L} <: Group{:w}
    w::Matrix{Float32}
    rule::L
end

# Dispatch on learning rule:
@inline on_learn(pre, syn::Elem{<:Synapses}, post) =
    on_learn(syn.o.rule, pre, syn, post)

@generated on_ff_rate(pre, syn::Elem{<:Synapses}, post::Elem{LIF}) =
@resolve(quote
    I_ff_post += X_pre * w
end, syn, pre=pre, post=post)

@generated on_rc_spike(pre, syn::Elem{<:Synapses}, post::Elem{LIF}) =
@resolve(quote
    I_rc_post += w
end, syn, post=post)

function no_self_connections!(syns::Synapses)
    @assert size(syns.w,1) == size(syns.w,2)
    @inbounds @simd for i in 1:size(syns.w,1)
        syns.w[i,i] = zero(eltype(syns.w))
    end
end

# =======================
# Excitatory plasticity
# =======================

@with_kw struct HebbOja
    β::Float32 = 0.001
end

@generated on_learn(rule::HebbOja, pre, syns, post) = @resolve(quote
    dw = X_pre * n_post - w * n_post^2
    w += dw * β
end, syns, rule, pre=pre, post=post)

# =======================
# Inhibitory plasticity
# =======================

@with_kw struct Foldiak
    α::Float32 = 0.1
end

@generated on_learn(rule::Foldiak, pre, syns, post) = @resolve(quote
    dw = n_pre * n_post - p_post^2
    w += dw * α
    w = max(w, zero(w))
end, syns, rule, pre=pre, post=post)

# =======================
# Network and training
# =======================

struct SAILnet{IN, FF, RC}
    input::IN
    lifs::LIF
    ff_path::FF
    rc_path::RC
end

function train_one(t, net::SAILnet, substeps=50, spike_rec=nothing)
    on_reset.(net.lifs)
    on_ff_rate.(net.ff_path)
    for t in t+1:t+substeps
        on_step.(net.lifs)
        on_rc_spike.(net.rc_path, has_spike)
        record!(spike_rec, t, has_spike)
    end
    on_learn.(net.lifs)
    on_learn.(net.ff_path)
    on_learn.(net.rc_path)
    no_self_connections!(net.rc_path.syns)
    t
end


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder
using Pkg

name = "FastJet_Julia_Wrapper"
version = v"0.8.7"
julia_versions = [v"1.6", v"1.7", v"1.8"]

# Collection of sources required to build FastJet_Julia_Wrapper
sources = [
	GitSource("https://github.com/jstrube/FastJet_Julia_Wrapper.git", "dc12b746c4ac0ec03e506113d21b53ff02f8e1c0"),
]

# Bash recipe for building across all platforms
script = raw"""
cd ${WORKSPACE}/srcdir/FastJet_Julia_Wrapper
mkdir build && cd build
cmake -DJulia_PREFIX=${prefix} -DCMAKE_INSTALL_PREFIX=${prefix} -DCMAKE_FIND_ROOT_PATH=${prefix} -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} -DCMAKE_BUILD_TYPE=Release ..
make -j${nproc}
make install
install_license $WORKSPACE/srcdir/FastJet_Julia_Wrapper/LICENSE.md
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
include("../../L/libjulia/common.jl")
platforms = expand_cxxstring_abis(vcat(libjulia_platforms.(julia_versions)...))

# the plugins aren't found on win. Disable for now, but this is not a fundamental limitation.
filter!(!Sys.iswindows, platforms)

# The products that we will ensure are always built
products = [
    LibraryProduct("libfastjetwrap", :libfastjetwrap)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    Dependency(PackageSpec(name="libcxxwrap_julia_jll")),
    Dependency("FastJet_jll"),
    BuildDependency(PackageSpec(name="libjulia_jll"))
]

build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies;
    preferred_gcc_version=v"8", julia_compat="1.6")

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
module Sads

using Distributions
using KernelDensity: kde
using Polynomials: polyfit

abstract type CrossSection end

include("crosssections.jl")
include("kalman.jl")
include("gvf.jl")
include("ensemble.jl")

"""
Estimate discharge by assimilating observed water surface elevations along a river channel.

"""
function estimate(qwbm, H, W, x, rₚ, ri; nₚ=Uniform(0.02, 0.05), nsamples=1000, nens=100)
    zm, zs, dm, ds, qb = priors(qwbm, H, W, x, nₚ, rₚ, nsamples, nens)
    Qₚ = Truncated(LogNormal(log(qwbm/sqrt(1+dm^2)), log(1+dm^2)), qb[1], qb[2])
    Qa, A0, n = assimilate(H, W, x, maximum(W, dims=2), maximum(H, dims=2),
                           Qₚ, nₚ, rₚ, Normal(zm, 0.1), nens, ri)
    Qa, A0, n
end

"""
Estimate bed slope by assimilating observed water surface elevations.

"""
function bed_slope(S, H, W, x, hbf, wbf, Qₚ, nₚ, rₚ, zₚ, nens; ϵₒ=0.01)
    S0 = [mean(S[j, :]) > 0 ? mean(S[j, :]) : mean(S[j, :][S[j,:] .> 0]) for j in 1:size(S, 1)]
    S0[isnan.(S0)] .= minimum(S[S .> 0])
    Se = rand.(Normal(1., 0.2), length(x), nens)
    Se = [Se[i, j] .* S0[i] for i in 1:length(x), j in 1:nens]
    ze = zeros(length(x), nens)
    Qe, ne, re, ze[1, :] = lhs_ensemble(nens, Qₚ, nₚ, rₚ, zₚ)
    he = gvf_ensemble!(mean(H, dims=2), mean(W, dims=2), Se, x, hbf, wbf, Qe, ne, re, ze)
    i = findall(he[1, :] .> 0)
    h = ze .+ he .* ((re .+ 1) ./ re)'
    X = Se[:, i]
    XA = h[:, i]
    d = mean(H, dims=2)[:, 1]
    E = rand(Normal(ϵₒ, 1e-6), length(d), length(i)) .* rand([-1, 1], length(d), length(i))
    A = letkf(X, d, XA, E, [collect(1:length(x))], [collect(1:length(d))], diagR=true)
    mean(A, dims=2)[:, 1]
end

"""
Estimate parameters for discharge estimation, i.e. A₀ and n.

"""
function estimate_Q_params(H, W, S, ri, ze, re, ne, Qa)
    nr, nt = size(Qa)
    nens = length(ne)
    A0 = zeros(nr)
    n = zeros(nr)
    Wbf = maximum(W, dims=2)
    Hbf = maximum(H, dims=2)
    Hmin = H[:, findmin(H[1, :])[2]]
    Wmin = W[:, findmin(H[1, :])[2]]
    ybf = Hbf .- ze
    ymin = Hmin .- ze
    ybf[ybf .< 0] .= 0.0
    ymin[ymin .< 0] .= 0.0
    A0e = Wbf .* (ymin ./ ybf).^(1 ./ re') .* (re ./ (re .+ 1))' .* ymin
    # A0e = Wmin .* (re ./ (re .+ 1))' .* ymin
    A0e = [mean(A0e[ri[j]:ri[j+1], e]) for j in 1:nr, e in 1:nens]
    dA = (H .- Hmin) .* (W .+ Wmin) ./ 2
    dA = reshape([mean(dA[ri[j]:ri[j+1], t]) for j in 1:nr, t in 1:nt]', nr*nt, 1)
    Wr = reshape([mean(W[ri[j]:ri[j+1], t]) for j in 1:nr, t in 1:nt]', nr*nt, 1)
    Sr = reshape([mean(S[ri[j]:ri[j+1], t]) for j in 1:nr, t in 1:nt]', nr*nt, 1)
    A = repeat(A0e, outer=nt) .+ dA
    A[A .< 0] .= 0.0
    Sr[Sr .< 0] .= minimum(Sr[Sr .> 0])
    Qe = (1 ./ ne') .* A.^(5/3) .* Wr.^(-2/3) .* Sr.^(1/2)
    X = zeros(nr*2, nens)
    X[1:2:end, :] = A0e
    X[2:2:end, :] = repeat(ne', outer=nr)
    XA = Qe
    d = reshape(Qa, nr*nt)
    E = rand(Normal(0.1*mean(Qa), 1e-6), length(d), nens) .* rand([-1, 1], length(d), nens)
    A = letkf(X, d, XA, E, [[2*j-1;2*j] for j in 1:nr], [collect(nt*j-1:nt*j) for j in 1:nr], diagR=true)
    A0 = mean(A[1:2:end, :], dims=2)
    n = mean(A[2:2:end, :], dims=2)
    A0, n
end

"""
Assimilate SWOT observations for river reach.

- `H`: water surface elevation
- `W`: water surface width
- `x`: downstream distance for each cross section
- `wbf`: bankfull width
- `hbf`: bankfull depth
- `Qₚ`: prior probability distribution for discharge
- `nₚ`: prior probability distribution for roughness coefficient
- `rₚ`: prior probability distribution for channel shape parameter
- `zₚ`: prior distribution for downstream bed elevation
- `nens`: ensemble size
- `ri`: reach definition indices
- `ϵₒ`: observation error standard deviation
- `cv_thresh`: threshold for using dynamic or time-constant bed slope

"""
function assimilate(H, W, x, wbf, hbf, Qₚ, nₚ, rₚ, zₚ, nens, ri; ϵₒ=0.01, logQ=false)
    min_ensemble_size = 5
    Qa = zeros(length(ri)-1, size(H, 2))
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    So = bed_slope(S, H, W, x, hbf, wbf, Qₚ, nₚ, rₚ, zₚ, nens)
    ze = zeros(length(x), nens)
    Qe, ne, re, ze[1, :] = lhs_ensemble(nens, Qₚ, nₚ, rₚ, zₚ)
    for t in 1:size(H, 2)
        rri = [ri[1:end-1]; length(x) + 1]
        he = gvf_ensemble!(H[:, t], W[:, t], So, x, hbf, wbf, Qe, ne, re, ze)
        i = findall(he[1, :] .> 0)
        if length(i) < min_ensemble_size
            Sf = [S[j, t] > 0 ? S[j, t] : minimum(S[:, t][S[:, t] .> 0]) for j=1:length(x)]
            he = [(Qe[e] .* ne[e]) ./ (W[j, t] .* Sf[j].^0.5).^(3/5) for j=1:length(x), e=1:nens]
        end
        i = findall(he[1, :] .> 0)
        h = ze .+ he .* ((re .+ 1) ./ re)'
        X = repeat(Qe[i]', outer=length(ri)-1)
        XA = h[:, i]
        d = H[:, t]
        E = rand(Normal(ϵₒ, 1e-6), length(d), length(i)) .* rand([-1, 1], length(d), length(i))
        if logQ
            A = letkf(log.(X), d, XA, E, [[j] for j in 1:length(ri)-1],
                      [collect(rri[j]:rri[j+1]-1) for j in 1:length(ri)-1], diagR=true)
            Qa[:, t] = mean(exp.(A), dims=2)
        else
            A = letkf(X, d, XA, E, [[j] for j in 1:length(ri)-1],
                      [collect(rri[j]:rri[j+1]-1) for j in 1:length(ri)-1], diagR=true)
            A[A .< 0] .= 0.0
            Qa[:, t] = mean(A, dims=2)  # could also use the absolute of the mean
        end
    end
    # remove outliers
    outlier = findall(mean(Qa,dims=2)[:, 1] ./ mean(Qa) .> 1.5)
    Qa[outlier, :] .= mean(Qa[[i for i in setdiff(Set(1:size(Qa, 1)), Set(outlier))], :], dims=1)
    # estimate SWOT parameters
    # A0, n = estimate_Q_params(H, W, S, ri, ze, re, ne, Qa)
    Qa #, A0, n
end

"""
Assimilate SWOT observations and apply hydraulic geometry constraints to estimate discharge.

"""
function assimilateHG(H, W, x, wbf, hbf, Qₚ, nₚ, rₚ, zₚ, nens, ri, ϵₒ=0.01; loc=0)
    min_ensemble_size = 5
    Qa = zeros(length(ri)-1, size(H, 2))
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    So = bed_slope(S, H, W, x, hbf, wbf, Qₚ, nₚ, rₚ, zₚ, nens)
    ze = zeros(length(x), nens)
    Qe, ne, re, ze[1, :] = lhs_ensemble(nens, Qₚ, nₚ, rₚ, zₚ)
    for t in 1:size(H, 2)
        rri = [ri[1:end-1]; length(x) + 1]
        he = gvf_ensemble!(H[:, t], W[:, t], So, x, hbf, wbf, Qe, ne, re, ze)
        h = he .* ((re .+ 1) ./ re)' .+ ze
        i = findall([!((any(diff(h[:, e]) .< 0)) | (he[1, e] < 0))  for e in 1:size(h, 2)]) # i = findall(he[1, :] .> 0)
        if length(i) > 1
            X = repeat(Qe[i]', outer=length(ri)-1)
            if loc == 0
                XA = zeros(length(x), length(i))
                for j in 1:length(x)
                    py = polyfit(log.(Qe[i]), log.(he[j, i]), 1)
                    c = exp(py[0])
                    f = py[1]
                    XA[j, :] = c .* Qe[i]'.^f .* ((re[i] .+ 1) ./ re[i])'
                end
                XA = XA .+ ze[:, i]
                d = H[:, t]
            else
                py = polyfit(log.(Qe[i]), log.(he[loc, i]), 1)
                c=exp(py[0])
                f=py[1]
                XA = c .* X.^f .* ((re[i] .+ 1) ./ re[i])' .+ ze[loc, i]
                d = [H[loc, t]]
            end
            E = rand(Normal(0.1, 1e-6), length(d), length(i)) .* rand([-1, 1], length(d), length(i))
            A = letkf(X, d, XA, E, [[j] for j in 1:length(ri)-1], [collect(rri[j]:rri[j+1]-1) for j in 1:length(ri)-1], diagR=true)
            A[A .<= 0] .= 0.0
            A = reshape(mean(A, dims=1), 1, length(i))
            Qc = zeros(nens)
            Qc .= mean(A)
            Qc[i] = A
        else
            Qc = Qe
        end
        he = gvf_ensemble!(H[:, t], W[:, t], So, x, hbf, wbf, Qc, ne, re, ze)
        i = findall(he[1, :] .> 0)
        if length(i) < min_ensemble_size
            Sf = [S[j, t] > 0 ? S[j, t] : minimum(S[:, t][S[:, t] .> 0]) for j=1:length(x)]
            he = [(Qe[e] .* ne[e]) ./ (W[j, t] .* Sf[j].^0.5).^(3/5) for j=1:length(x), e=1:nens]
        end
        i = findall(he[1, :] .> 0)
        h = ze .+ he .* ((re .+ 1) ./ re)'
        X = repeat(Qc[i]', outer=length(ri)-1)
        XA = h[:, i]
        d = H[:, t]
        E = rand(Normal(ϵₒ, 1e-6), length(d), length(i)) .* rand([-1, 1], length(d), length(i))
        A = Sads.letkf(X, d, XA, E, [[j] for j in 1:length(ri)-1],
                       [collect(rri[j]:rri[j+1]-1) for j in 1:length(ri)-1], diagR=true)
        A[A .< 0] .= 0.0;
        Qa[:, t] = mean(A, dims=2)  # could also use the absolute of the mean
    end
    # remove outliers
    outlier = findall(mean(Qa, dims=2)[:, 1] ./ mean(Qa) .> 1.5)
    #Qa[outlier, :] .= mean(Qa[[i for i in setdiff(Set(1:size(Qa, 1)), Set(outlier))], :], dims=1)
    Qa
end

"""
Estimate the prior probability distributions of bed elevation and discharge using rejection sampling.

"""
function priors(qwbm, H, W, x, nₚ, rₚ, nsamples, nens)
    zbnds, qbnds = prior_bounds(qwbm, nsamples, H, W, x, nₚ, rₚ)
    dm, ds = discharge_prior(qwbm, nens, nsamples, H, W, x, nₚ, rₚ, Normal(minimum(H[1, :]), 0.1))
    zm, zs = bed_elevation_prior(qwbm, nens, nsamples, H, W, x, nₚ, rₚ, zbnds, [dm, ds], qbnds)
    zm, zs, dm, ds, zbnds, qbnds
end

"""
Estimate the mean of the discharge prior probability distribution.

"""
function discharge_prior_mean(qwbm, H, W, x, nₚ, rₚ, nens)
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]' S]
    wbf = maximum(W, dims=2)
    hbf = maximum(H, dims=2)
    Qₚ = LogNormal(log(qwbm/sqrt(1+1.5^2)),log(1+1.5^2)) #Uniform(qwbm/10, qwbm*10)
    zₚ = Uniform(minimum(H[1, :])-30, minimum(H[1, :]))
    So = bed_slope(S, H, W, x, hbf, wbf, Qₚ, nₚ, rₚ, zₚ, nens)
    ze = zeros(length(x), nens)
    Qe, ne, re, ze[1, :] = lhs_ensemble(nens, Qₚ, nₚ, rₚ, zₚ)
    he = gvf_ensemble!(mean(H, dims=2), mean(W, dims=2), So, x, hbf, wbf, Qe, ne, re, ze)
    i = findall(he[1, :] .> 0)
    X = zeros(1, length(i))
    X[1, :] = log.(Qe[i])
    r = repeat(re[i]', outer=length(x))
    XA = wbf .* (1 ./ (hbf .- ze[:, i])).^(1 ./ r) .* ((r .+ 1) ./ r).^(1 ./ r) .* he[:, i].^(1 ./ r)
    d = mean(W, dims=2)[:, 1]
    E = rand(Normal(0.25*mean(W), 1e-6), length(d), length(i)) .* rand([-1, 1], length(d), length(i))
    A = letkf(X, d, XA, E, [[1]], [collect(1:length(d))], diagR=true)
    qm = exp.(A)
    Qₚ = Normal(mean(qm), std(qm)) #Uniform(mean(qm)/10,10*mean(qm))
    Qa = lhs_ensemble(nens, Qₚ)[1]
    Qa = abs.(Qa)
    he = gvf_ensemble!(mean(H, dims=2), mean(W, dims=2), So, x, hbf, wbf, Qa, ne, re, ze)
    i = findall(he[1, :] .> 0)
    r = repeat(re[i]', outer=length(x))
    X = zeros(1,length(i))
    X[1,:] = log.(Qa[i])
    YA = wbf .* (1 ./ (hbf .- ze[:, i])).^(1 ./ r) .* ((r .+ 1) ./ r).^(1 ./ r) .* he[:, i].^(1 ./ r)
    nmin = minimum([size(XA, 2),size(YA, 2)])
    db = d .- XA[:,1:nmin]
    da = d.- YA[:,1:nmin]
    du = YA[:, 1:nmin] .- XA[:, 1:nmin]
    Qe = (tr(du * db') / tr((XA .- mean(XA, dims=2)) * (XA .- mean(XA, dims=2))')) .* Qe
    he = gvf_ensemble!(mean(H, dims=2), mean(W, dims=2), So, x, hbf, wbf, Qe, ne, re, ze)
    i = findall(he[1, :] .> 0)
    r = repeat(re[i]', outer=length(x))
    X = zeros(1,length(i))
    X[1, :] = log.(Qe[i])
    YA = wbf .* (1 ./ (hbf .- ze[:, i])).^(1 ./ r) .* ((r .+ 1) ./ r).^(1 ./ r) .* he[:, i].^(1 ./ r)
    nmin = minimum([size(XA, 2), size(YA, 2)])
    A = letkf(X, d, YA, (tr(da * db') / tr(E * E'))*E[:, 1:nmin], [[1]], [collect(1:length(d))], diagR=true)
    mean(exp.(A))
end

"""
Estimate the prior probability distribution of discharge.

"""
function discharge_prior(qwbm, nens, nsamples, H, W, x, nₚ, rₚ, zₚ)
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    S0 = [mean(S[j, :]) > 0 ? mean(S[j, :]) : maximum(S[j, :]) for j in 1:size(S, 1)]
    obs = [std(sample(H[end, :], nens)) for _ in 1:nsamples]
    Fobs = kde(obs)
    h = zeros(nsamples)
    δ = lhs_ensemble(nsamples, Uniform(0.5, 2.0))[1]
    t = sample(collect(1:size(H, 2)), nsamples)
    for s in 1:nsamples
        μ = log(qwbm / sqrt(1 + δ[s]^2))
        σ = log(1 + δ[s]^2)
        Qₚ = LogNormal(μ, σ)
        Qe, ne, re, ze = prior_ensemble(x, Qₚ, nₚ, rₚ, zₚ, nens)
        he = gvf_ensemble!(H[:, t[s]], W[:, t[s]], S0, x, maximum(H, dims=2), maximum(W, dims=2), Qe, [mean(ne) for i in 1:nens], [mean(re) for i in 1:nens], ze)
        i = findall(he[1, :] .> 0)
        h[s] = std(he[end, i] .* (mean(re) .+ 1) ./ mean(re) .+ ze[end, i])
    end
    δ = δ[.!isnan.(h)]
    h = h[.!isnan.(h)]
    Fmod = kde(h)
    L = 1
    accepted = [rand(Uniform(0, L)) * pdf(Fmod, s) <= pdf(Fobs, s) for s in  h]
    if length(findall(accepted)) < 2
        δₘ, δₛ = δ[findmin(abs.(h .- mean(obs)))[2]], 0.01
    else
        δₘ, δₛ = mean(δ[accepted]), std(δ[accepted])
    end
    δₘ, δₛ
end

"""
Estimate the prior probability distribution of bed elevation.

"""
function bed_elevation_prior(qwbm, nens, nsamples, H, W, x, nₚ, rₚ, zbnds, dpars, qbnds)
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    S0 = [mean(S[j, :]) > 0 ? mean(S[j, :]) : maximum(S[j, :]) for j in 1:size(S, 1)]
    # S0 = [mean(S[j, :]) > 0 ? mean(S[j, :]) : minimum(S[j, :][S[j, :] .> 0]) for j in 1:size(S, 1)]
    dm, ds = dpars
    Qₚ = Truncated(LogNormal(log(qwbm/sqrt(1+dm^2)), log(1+dm^2)), qbnds[1], qbnds[2])
    zₚ = Uniform(zbnds...)
    Qe, ne, re, ze = prior_ensemble(x, Qₚ, nₚ, rₚ, zₚ, nsamples)
    So = bed_slope(S, H, W, x, maximum(H, dims=2), maximum(W, dims=2), Qₚ, nₚ, rₚ, zₚ, nens)
    he = gvf_ensemble!(mean(H, dims=2), mean(W, dims=2), So, x, maximum(H, dims=2), maximum(W, dims=2), Qe, ne, re, ze)
    h  = he .* ((re .+ 1) ./ re)' .+ ze
    i = findall([!((any(diff(h[:, e]) .< 0)) | (he[1, e] < 0))  for e in 1:size(h, 2)])
    X = reshape(ze[1, i], 1, length(i))
    XA  = he[:,i] .* ((re[i] .+ 1) ./ re[i])' .+ ze[:,i]
    d = mean(H, dims=2)[:, 1]
    E = rand(Normal(0.1, 1e-6), length(d), length(i)) .* rand([-1, 1], length(d), length(i))
    A = letkf(X, d, XA, E, [[1]], [collect(1:length(d))], diagR=true)
    zm = mean(A)
    zs = std(A)
    zm, zs
end

"""
Estimate the prior probability distribution of the channel shape parameter.

"""
function r_prior(nens, H, W, rinit, loc)
    wi = unique(i -> W[loc, i], 1:length(W[loc, :]))
    re = Sads.lhs_ensemble(nens, rinit)[1]
    X = zeros(1, nens)
    XA = zeros(length(wi)-2, nens); YA = zeros(length(wi)-2,nens);
    X[1, :] = log.(re)
    for t=1:length(wi)-2
        ze =  (W[loc, wi[t]].^re .* H[loc, wi[t+1]] .- W[loc, wi[t+1]].^re .* H[loc, wi[t]]) ./ (W[loc, wi[t]].^re .- W[loc, wi[t+1]].^re)
        ze[findall(ze .>= minimum(H[loc, :]))] .= minimum(H[loc, :]) - 0.01
        XA[t, :] = ((H[loc, wi[t+2]] .- ze) ./ (H[loc, wi[t+1]] .- ze)).^(1 ./ re)
    end
    d = W[loc, wi][3:end] ./ W[loc, wi][2:end-1]
    ϵ = std(W[loc, wi[1:end-1]] ./ W[loc, wi[2:end]])
    E = rand(Normal(ϵ, 1e-6), length(d), nens) .* rand([-1, 1], length(d), nens)
    A = Sads.letkf(X, d, XA, E, [[1]], [collect(1:length(d))])
    ra = exp.(A)
    for t=1:length(wi)-2
        ze = (W[loc, wi[t]].^ra .* H[loc, wi[t+1]] .- W[loc, wi[t+1]].^ra .* H[loc, wi[t]]) ./ (W[loc, wi[t]].^ra .- W[loc, wi[t+1]].^ra)
        ze[findall(ze .>= minimum(H[loc, :]))] .= minimum(H[loc, :]) - 0.01
        YA[t, :] = ((H[loc, wi[t+2]] .- ze) ./ (H[loc, wi[t+1]] .- ze)).^(1 ./ ra)
    end
    db = d .- XA
    da = d .- YA
    A = Sads.letkf(X, d, XA, (tr(da * db') / tr(E * E')) * E, [[1]], [collect(1:length(d))]);
    mean(exp.(A)), std(exp.(A))
end

"""
Estimate the bounds of the prior distributions.

"""
function prior_bounds(qwbm, nsamples, H, W, x, nₚ, rₚ, zlim=20)
    S = diff(H, dims=1) ./ diff(x)
    S = [S[1, :]'; S]
    S0 = [mean(S[j, :]) > 0 ? mean(S[j, :]) : maximum(S[j, :]) for j in 1:size(S, 1)]
    # We will use arbitrary values that are large enough to represent the uninformative priors
    zₚ = Uniform(minimum(H[1, :]) - zlim, minimum(H[1, :]))
    Qₚ = Uniform(qwbm / 10, qwbm * 10)
    ze = zeros(length(x), nsamples)
    Qe, ze[1, :] = lhs_ensemble(nsamples, Qₚ, zₚ)
    re = rand(rₚ, nsamples)
    ne = rand(nₚ, nsamples)
    ze = zeros(length(x), nsamples)
    Qe, ze[1, :] = lhs_ensemble(nsamples, Qₚ, zₚ)
    he = gvf_ensemble!(mean(H, dims=2), mean(W, dims=2), S0, x, maximum(H, dims=2), maximum(W, dims=2), Qe, ne, re, ze)
    zbnds, qbnds = try
        i = findall(he[1, :] .> 0)
        h = he[end, i] .* (mean(re) .+ 1) ./ mean(re) .+ ze[end, i]
        j = i[(h .> minimum(H[end, :])) .& (h .< maximum(H[end, :]))]
        [minimum(ze[1, j]), maximum(ze[1, j])], [minimum(Qe[j]), maximum(Qe[j])]
    catch
        [minimum(ze[1, :]), maximum(ze[1, :])], [minimum(Qe[1, :]), maximum(Qe[1, :])]
    end
    zbnds, qbnds
end

"""
Estimate At-a-station Hydraulic Geometry coefficients.

"""
function estimate_ahg(r, n, wbf, ybf, S)
    p = 2 / 3
    q = 1 / 2
    κ = S^q / n
    δ = 1 + r + r * p
    R = (1 + r) / r
    α = wbf.^((r + r * p) / δ) .* (ybf ./ R).^(-(1 + p) / δ) .* κ.^(-1 / δ)
    c = wbf.^(-r / δ) .* (ybf / R).^(1 / δ) .* κ.^(-r / δ)
    k = wbf.^(-r * p / δ) .* (ybf / R).^(p / δ) .* κ.^((1 + r) / δ)
    b = 1 / δ
    f = r / δ
    m = p * r / δ
    return α, c, k, b, f, m
end

end


text:
import ContinuumArrays: apply, MulQuasiMatrix
import OffsetArrays: OffsetArray

@testset "$(rpad("Lagrange Basis Tests",80))" begin

    x = [0.0, 1.0]
    l = Lagrange(x)
    d = Derivative(axes(l,1))

    @test apply(*,d,l) isa LagrangeDerivative
    @test d*l isa LagrangeDerivative
    @test l' isa LagrangeDerivative

    @test grid(l) == x
    @test basis(l) == l.b
    @test nodes(l) == x
    @test nbasis(l) == 2
    @test nnodes(l) == 2
    @test eachindex(l) == 1:2
    @test order(l)  == 2
    @test degree(l) == 1

    l0 = Lagrange(OffsetArray(x, 0:1))
    l1 = Lagrange([0.0,  1.0 ])
    l2 = Lagrange([0,    1   ])
    l3 = Lagrange([0.25, 0.75])
    l4 = Lagrange{Float32}(x)

    @test hash(l) == hash(l0)
    @test hash(l) == hash(l1)
    @test hash(l) == hash(l2)
    @test hash(l) != hash(l3)
    @test hash(l) == hash(l4)

    @test l == l0
    @test l == l1
    @test l == l2
    @test l != l3
    @test l == l4

    @test  isequal(l, l0)
    @test  isequal(l, l1)
    @test !isequal(l, l2)
    @test !isequal(l, l3)
    @test !isequal(l, l4)

    @test  isapprox(l, l0)
    @test  isapprox(l, l1)
    @test  isapprox(l, l2)
    @test !isapprox(l, l3)
    @test  isapprox(l, l4)

    @test l != LagrangeGauß(2)
    @test l == LagrangeLobatto(2)

    y = rand(5)

    z1 =      [ l[y[i], j] for i in eachindex(y), j in eachindex(l)]
    z2 = hcat([ l[y[i], :] for i in eachindex(y)]...)'
    z3 = hcat([ l[y,    j] for j in eachindex(l)]...)
    z4 =        l[y,    :]

    @test z1 == z2 == z3 == z4

    z1 =      [ (d*l)[y[i], j] for i in eachindex(y), j in eachindex(l)]
    z2 = hcat([ (d*l)[y[i], :] for i in eachindex(y)]...)'
    z3 = hcat([ (d*l)[y,    j] for j in eachindex(l)]...)
    z4 =        (d*l)[y,    :]

    @test z1 == z2 == z3 == z4

    @test l(0.0, 1) == 1.0
    @test l(0.5, 1) == 0.5
    @test l(1.0, 1) == 0.0

    @test l(0.0, 2) == 0.0
    @test l(0.5, 2) == 0.5
    @test l(1.0, 2) == 1.0

    @test l[0.0, 1] == 1.0
    @test l[0.5, 1] == 0.5
    @test l[1.0, 1] == 0.0

    @test l[0.0, 2] == 0.0
    @test l[0.5, 2] == 0.5
    @test l[1.0, 2] == 1.0

    @test (d*l)[0.0, 1] == -1.0
    @test (d*l)[0.5, 1] == -1.0
    @test (d*l)[1.0, 1] == -1.0

    @test (d*l)[0.0, 2] == 1.0
    @test (d*l)[0.5, 2] == 1.0
    @test (d*l)[1.0, 2] == 1.0

    @test (d*l)[1, 1] == -1.0
    @test (d*l)[2, 1] == -1.0

    @test (d*l)[1, 2] == 1.0
    @test (d*l)[2, 2] == 1.0

end


text:
include("runtests_setup.jl")
println("testing types: ", GridapDistributedPETScWrappers.C.petsc_type)
for (i, ST) in enumerate(GridapDistributedPETScWrappers.C.petsc_type)
  if GridapDistributedPETScWrappers.have_petsc[i]
    println("testing datatype ", ST)
  # @testset "Scalar type $ST" begin # uncomment when nested test results can be printed
    include("error.jl")
    include("ksp.jl")
    include("vec.jl")
    include("is.jl")
    include("mat.jl")
    #include("ts.jl")
  end
  # end
end

#@test GridapDistributedPETScWrappers.petsc_sizeof(GridapDistributedPETScWrappers.C.PETSC_BOOL) == 4


text:
# Miscellaneous helper macros and functions
#

# Helper macro to iterate over the DG grid. Generates the needed loops
# and indices: `eh`, `ev`, `e`, `k,`, `j`, `i`, `ijk`.
macro visitQ(nhorzelem, nvertelem, Nqk, Nq, expr)
    return esc(
        quote
            for eh in 1:($nhorzelem)
                for ev in 1:($nvertelem)
                    e = ev + (eh - 1) * $nvertelem
                    for k in 1:($Nqk)
                        for j in 1:($Nq)
                            for i in 1:($Nq)
                                ijk = i + $Nq * ((j - 1) + $Nq * (k - 1))
                                $expr
                            end
                        end
                    end
                end
            end
        end,
    )
end

# Helper macro to iterate over a 3D array. Used for walking the
# interpolated grid.
macro visitI(nlong, nlat, nlevel, expr)
    return esc(quote
        for lo in 1:($nlong)
            for la in 1:($nlat)
                for le in 1:($nlevel)
                    $expr
                end
            end
        end
    end)
end

# Helpers to extract data from the various state arrays
function extract_state(dg, state, ijk, e, st::AbstractStateType)
    bl = dg.balance_law
    FT = eltype(state)
    num_state = number_states(bl, st)
    local_state = MArray{Tuple{num_state}, FT}(undef)
    for s in 1:num_state
        local_state[s] = state[ijk, s, e]
    end
    return Vars{vars_state(bl, st, FT)}(local_state)
end

max_stars_repo_name: mwarusz/CLIMA

max_issues_repo_name: mwarusz/CLIMA

max_forks_repo_name: mwarusz/CLIMA


text:
using BinaryProvider # requires BinaryProvider 0.3.0 or later

# Parse some basic command-line arguments
const verbose = "--verbose" in ARGS
const prefix = Prefix(get([a for a in ARGS if a != "--verbose"], 1, joinpath(@__DIR__, "usr")))
products = [
    LibraryProduct(prefix, String["libogg"], :libogg),
]

# Download binaries from hosted location
bin_prefix = "https://github.com/staticfloat/OggBuilder/releases/download/v1.3.3-3"

# Listing of files generated by BinaryBuilder:
download_info = Dict(
    BinaryProvider.Linux(:aarch64, :glibc, :blank_abi) => ("$bin_prefix/Ogg.aarch64-linux-gnu.tar.gz", "fb9ac4883b0a4baef13f3016b542744a1d786d186038fabe5f57518197695f5d"),
    BinaryProvider.Linux(:armv7l, :glibc, :eabihf) => ("$bin_prefix/Ogg.arm-linux-gnueabihf.tar.gz", "161f8cf6b8b4bd495c41f1d88899a6064a0444306c16d1dbad554642a3e89611"),
    BinaryProvider.Linux(:i686, :glibc, :blank_abi) => ("$bin_prefix/Ogg.i686-linux-gnu.tar.gz", "635708f9c87a796ad65c392c62846da3350c1d934a470ee8b86f95611a77699f"),
    BinaryProvider.Linux(:i686, :musl, :blank_abi) => ("$bin_prefix/Ogg.i686-linux-musl.tar.gz", "54e0f6d77f1db4afd2de69058c3d15403581aad36dda83d2ad61771113d0f410"),
    BinaryProvider.Windows(:i686, :blank_libc, :blank_abi) => ("$bin_prefix/Ogg.i686-w64-mingw32.tar.gz", "51cffd1229735d56a030ab78fb2968f3584d1e2c337c0e1aad357dc2b74fcfa1"),
    BinaryProvider.Linux(:powerpc64le, :glibc, :blank_abi) => ("$bin_prefix/Ogg.powerpc64le-linux-gnu.tar.gz", "b325fe0ce06e1d1588caa8d24da9dc750079cfc3c3e6d32408187dd639f141bb"),
    BinaryProvider.MacOS(:x86_64, :blank_libc, :blank_abi) => ("$bin_prefix/Ogg.x86_64-apple-darwin14.tar.gz", "b6d105845344223b959b788ab40c6143f085b7e1f35225b94c11c5b16612730f"),
    BinaryProvider.Linux(:x86_64, :glibc, :blank_abi) => ("$bin_prefix/Ogg.x86_64-linux-gnu.tar.gz", "5e09218b923643fd2a1e371b9288bc6dafa3adcb80e65761ac1343b8d3533d8d"),
    BinaryProvider.Linux(:x86_64, :musl, :blank_abi) => ("$bin_prefix/Ogg.x86_64-linux-musl.tar.gz", "52e0097cd1969f5594b939a8d50cb228cebb6a579061da8d70d14c8e930ca2c8"),
    BinaryProvider.FreeBSD(:x86_64, :blank_libc, :blank_abi) => ("$bin_prefix/Ogg.x86_64-unknown-freebsd11.1.tar.gz", "92245167b21af799d71abfcecba6885799e5e7740ace33c1623a9195a40a0348"),
    BinaryProvider.Windows(:x86_64, :blank_libc, :blank_abi) => ("$bin_prefix/Ogg.x86_64-w64-mingw32.tar.gz", "64e0d3412c6304e12156ffaa76765e048651d828dbb7d53ed64fbd7dd3b008c9"),
)

# Install unsatisfied or updated dependencies:
unsatisfied = any(!satisfied(p; verbose=verbose) for p in products)
if haskey(download_info, platform_key())
    url, tarball_hash = download_info[platform_key()]
    if unsatisfied || !isinstalled(url, tarball_hash; prefix=prefix)
        # Download and install binaries
        install(url, tarball_hash; prefix=prefix, force=true, verbose=verbose)
    end
elseif unsatisfied
    # If we don't have a BinaryProvider-compatible .tar.gz to download, complain.
    # Alternatively, you could attempt to install from a separate provider,
    # build from source or something even more ambitious here.
    error("Your platform $(triplet(platform_key())) is not supported by this package!")
end

# Write out a deps.jl file that will contain mappings for our products
write_deps_file(joinpath(@__DIR__, "deps.jl"), products)