text:
using Test, ImageFeatures, Images, Colors, FixedPointNumbers

@testset "circular_offsets" begin
    img = [  0x4c  0x19  0xac  0x2e  0x8c  0xcc  0x96  0x4c  0xdb  0x4f
             0x01  0x75  0x65  0x57  0x59  0x24  0xe9  0x8e  0x1f  0xd3
             0xb3  0x92  0xdd  0xdb  0x2e  0xf2  0xcc  0xc1  0xac  0x97
             0xc8  0x75  0xa8  0x38  0x6d  0x77  0xc3  0xf4  0x63  0x15
             0x7c  0xdf  0x54  0x11  0x0f  0x03  0xea  0x01  0x5e  0xe6
             0xa5  0x87  0x40  0x22  0xeb  0x71  0xf7  0xba  0xb2  0x6f
             0x9d  0xa7  0x58  0x7e  0xf8  0x0f  0x91  0xd0  0x8e  0x04
             0x58  0xd3  0x8b  0x5d  0x6f  0x8b  0xaa  0x77  0x53  0x46
             0x9a  0x74  0xd8  0x2f  0x38  0x4e  0x4b  0x18  0x4a  0xc7
             0x1a  0xe0  0x52  0x5a  0x6a  0x03  0xe8  0xcb  0x95  0xfc
          ]

    global img_gray = map(i -> Images.Gray(reinterpret(N0f8, i)), img)
end

@test ImageFeatures.circular_offsets(8, 1) == [ (-0.0,1.0), (-0.70711,0.70711), (-1.0,0.0), (-0.70711,-0.70711), (-0.0,-1.0), (0.70711,-0.70711), (1.0,-0.0), (0.70711,0.70711)]

@testset "Original" begin
    uniform_params = ImageFeatures.PatternCache(1)
    @test lbp_original(BitArray([false]), uniform_params)[1] == 0
    @test lbp_original(BitArray([false, true]), uniform_params)[1] == 1
    @test lbp_original(BitArray([true, false, true]), uniform_params)[1] == 5
    @test lbp_original(BitArray([true, false, true, true, false]), uniform_params)[1] == 22

    img = zeros(Gray{N0f8}, 10, 10)

    lbp_image = lbp(img)
    @test all(lbp_image .== 255)

    lbp_image = lbp(img, 10, 2)
    @test all(lbp_image .== 1023)

    lbp_image = lbp(img_gray)
    expected_lbp = [ 247  226  255  224  239  247  250  232  255  250
                     131  209  176   72  149    0  253   80    0  127
                     251  228  255  254   32  255  182  120  220  190
                     247    8  157   44   79  135   43  255   22   56
                     163  255   60   16   16    0  219    0   33  255
                     215  184   48   96  251  196  255  122  220  188
                     167  251   96  205  255    0  139  127   62   56
                     131  247  122  140   47   79  223   30   54  120
                     215  138  253    0    9   21   18    0   33  251
                     143  255   46  111  255   14  255  254  206  255
                    ]
    @test all(lbp_image .== expected_lbp)

    lbp_image = lbp(img_gray, 8, 1)
    expected_lbp = [ 255   48  255  112  127  255  245  112  255  241
                       8  184  208   32  154    0  255  160    0  255
                     253  114  255  247    0  255  214  225  243  215
                     255    0  159    3   47   30   13  255  134  128
                      28  255  195  128  128    0  253    0   72  255
                     191  193  192   32  253   50  255  245  243  195
                      94  253   96   62  255    0   28  239  199  128
                      28  255  225    3   79  111  191  135  198  225
                     190    4  255    0    9  142  132    0    8  253
                      14  255    7  111  255    2  255  247   55  255
                    ]
    @test all(lbp_image .== expected_lbp)
end

@testset "Uniform" begin
    uniform_params = ImageFeatures.PatternCache(4)
    @test uniform_params.table[BitArray([true, false, true, false])] == 14
    ret, uniform_params = lbp_uniform(BitArray([false, false, false, false]), uniform_params)
    @test ret == 1
    ret, uniform_params = lbp_uniform(BitArray([false, true, false, false]), uniform_params)
    @test ret == 2
    ret, uniform_params = lbp_uniform(BitArray([false, false, true, false]), uniform_params)
    @test ret == 3

    uniform_params = ImageFeatures.PatternCache(8)
    @test uniform_params.table[[true, false, true, false, true, false, true, false]] == 58

    lbp_image = lbp(img_gray, lbp_uniform)
    expected_lbp = [  1  58   6  11  15   1  58  58   6  58
                      2  58  58  58  58  14  10  58  14  21
                      3  58   6  12  16   6  58  20  58  58
                      1   5  58  58  58  17  58   6  58  24
                     58   6   7  13  13  14  58  14  58   6
                     58  58   8   9   3  58   6  58  58  58
                     58   3   9  58   6  14  58  21  23  24
                      2   1  58  58  58  58  19  22  58  20
                     58  58  10  14  58  58  58  14  58   3
                      4   6  58  58   6  18   6  12  58   6
                    ]
    @test all(lbp_image .== expected_lbp)

    lbp_image = lbp(img_gray, 8, 1, lbp_uniform)
    expected_lbp = [    1   7   1  17  23   1  58  17   1  30
                        2  58  58  18  58   8   1  58   8   1
                        3  58   1  19   8   1  58  15  28  58
                        1   8  11  20  58  24  58   1  58  21
                        4   1  12  21  21   8   3   8  58   1
                        5   9  13  18   3  58   1  58  28  12
                       58   3  14  22   1   8   4  26  29  21
                        4   1  15  20  58  58   5  27  58  15
                       58  10   1   8  58  58  58   8   2   3
                        6   1  16  58   1  25   1  19  58   1
                    ]
    @test all(lbp_image .== expected_lbp)
end

@testset "Modified" begin
    img = zeros(Gray{N0f8}, 10, 10)

    lbp_image = modified_lbp(img)
    @test all(lbp_image .== 255)

    lbp_image = modified_lbp(img, 10, 2)
    @test all(lbp_image .== 1023)

    lbp_image = modified_lbp(img_gray)
    expected_lbp = [ 231  226  241  224  233  230  242  232  245  250
                     163  209  177   89  149  131  113  208  160  121
                     195  228  234  252  230  199  162  104   76  188
                     131   74   29   46   79  135   11   28   22   56
                     131   93   60   87  147  197  147   48  105  120
                     131  184  116  227  233  196  163  106  216   60
                     135   50  104  205  182   86  139   29   62   56
                     131   53  122  156   47   79  142   30   54  120
                     199  170   61   86   25   23   19  113  225  249
                     143   31   46  110   30  142   47  238  206  190
                    ]
    @test all(lbp_image .== expected_lbp)

    lbp_image = modified_lbp(img_gray, 8, 1)
    expected_lbp = [  126  116  248  112  120  118  240  113  250  241
                       92   56  216  233  154   28  232  225   88  233
                       60  114  125  227  118   62  214  225   51  195
                       28   37  139  199   47   30   13  131  198  193
                       28  171  195  166  156   58  156  193  104  225
                       28  193  226  124  121   50   94   96  177  195
                       30  228   97   62  223  167   29  143  195  193
                       28  206  225  147   79   39   23  135  198  225
                       30   29  203  227  137  142  204  200  120  249
                       31  143  199   39  135   23   63  103   55  199
                    ]
    @test all(lbp_image .== expected_lbp)
end

@testset "Rotation Invariant" begin
    uniform_params = ImageFeatures.PatternCache(1)
    @test lbp_rotation_invariant(BitArray([false]), uniform_params)[1] == 0
    @test lbp_rotation_invariant(BitArray([false, true]), uniform_params)[1] == 1
    @test lbp_rotation_invariant(BitArray([true, false, true]), uniform_params)[1] == 3
    @test lbp_rotation_invariant(BitArray([true, false, true, true, false]), uniform_params)[1] == 11
    @test lbp_rotation_invariant(BitArray([false, true, true, true, false, true, true, false]), uniform_params)[1] == 59

    img = zeros(Gray{N0f8}, 10, 10)
    lbp_image = lbp(img)
    @test all(lbp_image .== 255)

    lbp_image = lbp(img, 10, 2)
    @test all(lbp_image .== 1023)

    lbp_image = lbp(img_gray, lbp_rotation_invariant)
    expected_lbp = [ 127   23  255    7  127  127   95   29  255   95
                       7   29   11    9   43    0  127    5    0  127
                     127   39  255  127    1  255   91   15   55   95
                     127    1   59   11   61   15   43  255   11    7
                      29  255   15    1    1    0  111    0    9  255
                      95   23    3    3  127   19  255   61   55   47
                      61  127    3   55  255    0   23  127   31    7
                       7  127   61   25   47   61  127   15   27   15
                      95   21  127    0    9   21    9    0    9  127
                      31  255   23  111  255    7  255  127   59  255
                    ]
    @test all(lbp_image .== expected_lbp)

    lbp_image = lbp(img_gray, 8, 1, lbp_rotation_invariant)
    expected_lbp = [ 255    3  255    7  127  255   95    7  255   31
                       1   23   13    1   53    0  255    5    0  255
                     127   39  255  127    0  255   91   15   63   95
                     255    0   63    3   47   15   13  255   13    1
                       7  255   15    1    1    0  127    0    9  255
                     127    7    3    1  127   25  255   95   63   15
                      47  127    3   31  255    0    7  127   31    1
                       7  255   15    3   61  111  127   15   27   15
                      95    1  255    0    9   29    9    0    1  127
                       7  255    7  111  255    1  255  127   55  255
                    ]
    @test all(lbp_image .== expected_lbp)
end

@testset "Direction Coded" begin
    img = zeros(Gray{N0f8}, 10, 10)

    lbp_image = direction_coded_lbp(img)
    @test all(lbp_image .== 255)

    lbp_image = direction_coded_lbp(img, 10, 2)
    @test all(lbp_image .== 1023)

    lbp_image = direction_coded_lbp(img_gray)
    expected_lbp = [ 126   15  254   82  252  126  221  211  254  204
                      37   26   36   74   27  175  182  157  187   63
                     218  115  190  237  167  250   88   21  185  141
                     106  110  139   82  100   17   40  175  144   96
                     120  187   16  188  173  187  211  250  225  191
                     114  181  224  131  158  123  238   92  253  149
                      72  143  194  253  235  238  244   43   13  112
                     113  110    8  142   29  113  227    0  220   69
                      98  163  167  254   61  142  161  234  161  159
                     213  170   26  105  171   87  170  169  163  175
                    ]
    @test all(lbp_image .== expected_lbp)

    lbp_image = direction_coded_lbp(img_gray, 8, 1)
    expected_lbp = [ 175  224  191  132   63  191  119  197  191   18
                     122  165    8  163  244  250  191  103  238  254
                     167  205  190  122  234  175   37   84   78  114
                     171  187  242  181    8   84    8  250   23   43
                      13  254    5   62  122  254  231  175   75  254
                     143   76   11  226  183  237  171   54   78   84
                      33  242  131   93  235  187   29  232  112   47
                      92  187   16  176  101  108  202    0   55   81
                     137  218  250  175  125  146   90  170  106  246
                      86  170  133  104  234  247  170  106  202  234
                    ]
    @test all(lbp_image .== expected_lbp)
end

@testset "Multi Block" begin
    img = zeros(Gray{Float64}, 12, 12)
    img[1:4, 1:4] .= Gray(6.0)
    img[9:12, 1:4] .= Gray(3.0)
    img[5:8, 5:8] .= Gray(1.0)

    @test multi_block_lbp(img, 1, 1, 4, 4) == 130

    img = zeros(Gray{Float64}, 12, 12)
    img[9:12, 1:4] .= Gray(6.0)
    img[5:8, 9:12] .= Gray(3.0)
    img[5:8, 5:8] .= Gray(1.0)

    @test multi_block_lbp(img, 1, 1, 4, 4) == 18
end

@testset "Descriptor" begin
    img = Gray{N0f8}[   0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0
                        0 1 1 0 1 1 0 1 1 0 ]

    descriptor = create_descriptor(img, 1, 1)
    expected = [0,0,0,0,18,0,0,9,0,0,0,0,9,0,3,61]
    @test all(descriptor .== expected)

    img = Gray{N0f8}[   0 1 1 0 1 1 0 1 1 0 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        0 1 1 0 1 1 0 1 1 0 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        0 1 1 0 1 1 0 1 1 0 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        0 1 1 0 1 1 0 1 1 0 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1
                        1 1 1 1 1 1 1 1 1 1 1 1 ]

    descriptor = create_descriptor(img, 2, 2)
    expected = [1,0,0,0,0,0,0,0,1,0,0,0,0,0,2,32,1,0,0,0,0,0,0,0,1,0,0,0,0,0,2,32,1,0,0,0,0,0,0,0,1,0,0,0,0,0,2,32,1,0,0,0,0,0,0,0,1,0,0,0,0,0,2,32]
    @test all(descriptor .== expected)
end


text:
# This set of tests is specifically for the ordering
# of operations, when code blocks are eval-ed or re-eval-ed

@testset "EvalOrder" begin
    flush_td(); set_globals()

    Random.seed!(0) # seed for extra testing of when code is ran
    # Generate random numbers in a sequence then we'll check
    # things are done in the same sequence on the page
    a, b, c, d, e = randn(5)
    Random.seed!(0)

    # Create a page "foo.md" which we'll use and abuse
    foo = raw"""
        @def hascode = true
        ```julia:ex
        println(randn())
        ```
        \output{ex}
        """

    # XXX FIRST PASS --> EVAL

    h = foo |> fd2html_td

    @test isapproxstr(h, """

$(F.htmlesc("""println(randn())"""))

$a

                """)

    # XXX TEXT MODIFICATION + IN SCOPE --> NO REEVAL

    foo *= "etc"

    h = foo |> fd2html_td

    @test isapproxstr(h, """

$(F.htmlesc("""println(randn())"""))

$a
 etc

                """)

    # XXX CODE ADDITION --> NO REEVAL OF FIRST BLOCK

    foo *= raw"""
        ```julia:ex2
        println(randn())
        ```
        \output{ex2}
        ```julia:ex3
        println(randn())
        ```
        \output{ex3}
        """

    h = foo |> fd2html_td

    @test isapproxstr(h, """

            $(F.htmlesc(raw"""
            println(randn())
            """))

$a
 etc

            $(F.htmlesc(raw"""
            println(randn())
            """))

$b

            $(F.htmlesc(raw"""
            println(randn())"""))

$c

        """)

    # XXX CODE MODIFICATION --> REEVAL OF BLOCK AND AFTER

    foo = raw"""
        @def hascode = true
        ```julia:ex
        println(randn())
        ```
        \output{ex}
        ```julia:ex2
        # modif
        println(randn())
        ```
        \output{ex2}
        ```julia:ex3
        println(randn())
        ```
        \output{ex3}
        """

    h = foo |> fd2html_td

    @test isapproxstr(h, """

$(F.htmlesc(raw"""
            println(randn())
            """))

$a

$(F.htmlesc(raw"""
            # modif
            println(randn())
            """))

$d

$(F.htmlesc(raw"""
            println(randn())
            """))

$e

        """)
end

@testset "EvalOrder2" begin
    flush_td(); set_globals()

    # Inserting new code block
    h = raw"""
        @def hascode = true
        ```julia:ex1
        a = 5
        println(a)
        ```
        \output{ex1}
        ```julia:ex2
        a += 3
        println(a)
        ```
        \output{ex2}
        """ |> fd2html_td

    @test isapproxstr(h, """

$(F.htmlesc(raw"""
                a = 5
                println(a)
                """))

5

$(F.htmlesc(raw"""
                a += 3
                println(a)
                """))

8

            """)

    h = raw"""
        @def hascode = true
        @def reeval  = true
        ```julia:ex1
        a = 5
        println(a)
        ```
        \output{ex1}
        ```julia:ex1b
        a += 1
        println(a)
        ```
        \output{ex1b}
        ```julia:ex2
        a += 3
        println(a)
        ```
        \output{ex2}
        """ |> fd2html_td

    @test isapproxstr(h, """

$(F.htmlesc(raw"""
                a = 5
                println(a)
                """))

5

$(F.htmlesc(raw"""
                a += 1
                println(a)
                """))

6

$(F.htmlesc(raw"""
                a += 3
                println(a)
                """))

9

            """)
end

# https://github.com/tlienart/Franklin.jl/issues/832
@testset "#832" begin
    s = """
       ```!
       x=5 # a
       ```
       """ |> F.fd2html
    @test isapproxstr(s, """

x=5 # a

5

        """)
    s = """
       ```!
       x=5; # a
       ```
       """ |> F.fd2html
    @test isapproxstr(s, """

x=5; # a

        """)
end

ext: jl


text: include("singlehost.jl")


text:
import Base: length, eltype

abstract type GroupElement end

struct FiniteGroup{E <: GroupElement}
    e::Set{E}

    function (::Type{<:FiniteGroup})(gen)
        new{Base.eltype(gen)}(generate_(gen))
    end
end

function generate_(gen)
    e = Set(gen)
    last = e
    while true
        found = setdiff((x * y for x in gen, y in last), e)
        if isempty(found) break end
        union!(e, found)
        last = found
    end
    e
end

# Iteration
iterate(fg::FiniteGroup) = iterate(fg.e)
iterate(fg::FiniteGroup, i) = iterate(fg.e, i)
eltype(fg::FiniteGroup) = eltype(fg.e)
length(fg::FiniteGroup) = length(fg.e)


text:
function selectReaction(a::Vector{Float64}, suma::Float64, nr::Int)
    t = rand() * suma
    rxn = 0
    while t >= 0.0 && rxn < nr
        rxn += 1
        t -= a[rxn]
    end
    rxn
end

function selectReaction(a::SubArray, suma::Float64, nr::Int)
    t = rand() * suma
    rxn = 0
    while t >= 0.0 && rxn < nr
        rxn += 1
        t -= a[rxn]
    end
    rxn
end

function directMethod!(p::Path, F::Function, ν::Matrix{Int}, tf::Float64, nr::Int)
    #a = zeros(Float64, nr)
    fly = zeros(nr)
    while true
        # stop state-space exploration
        moleculeLimit!(p, tf) && break

        # Propensities
        F(p)
        suma = sum(p.a)

        if suma <= 0.0
            break
        end

        # Time step
        dt = randexp()/suma
        if (p.t + dt) > tf
            p.t = tf
            break
        end
        p.t += dt

        # calc on the fly
        @simd for i in eachindex(fly)
            @inbounds fly[i] += p.a[i]*dt
        end

        # Reaction selection & state update
        rxn = selectReaction(p.a, suma, nr)
        p.ra[rxn] += 1
        dx = view(ν, rxn, :)
        @simd for i in eachindex(p.x)
            @inbounds p.x[i] += dx[i]
        end

    end
    p.fly += fly./p.θ
end
directMethod = directMethod!

function nDirectMethod!(p::Path, F::Function, ν::Matrix{Int}, tf::Float64, nr::Int, nd::Int)
    nstep = 0
    fly = zeros(nr)
    while nstep < nd
        moleculeLimit!(p, tf) && break

        # Propensities
        F(p)

        suma = sum(p.a)
        if suma <= 0.0
            break
        end

        # Time step
        dt = randexp()/suma
        if (p.t + dt) > tf
            p.t = tf
            break
        end
        p.t += dt

        # calc on the fly
        @simd for i in eachindex(fly)
            @inbounds fly[i] += p.a[i]*dt
        end
        # Reaction selection & state update
        rxn = selectReaction(p.a, suma, nr)
        p.ra[rxn] += 1
        dx = view(ν, rxn, :)
        @simd for i in eachindex(p.x)
            @inbounds p.x[i] += dx[i]
        end

        nstep += 1
    end
    p.fly += fly./p.θ
end

function moleculeLimit!(p::Path, tf::Float64)
    if any(p.x .> 1500)
        p.t = tf
        p.dm .+= Inf
        return true
    end
    return false
end

max_stars_repo_name: pzuliani/SPICE

max_issues_repo_name: pzuliani/SPICE

max_forks_repo_name: pzuliani/SPICE


text:
""" Julia program to Check if a number is a Mystery Number or Not.

A mystery number is that number which can be expressed as a sum of two
numbers and those two numbers must be the reverse of each other
"""

function reverse(n)
    rev = 0
    while n > 0
        rev = rev * 10
        rev = rev + (n % 10);
        n = n ÷ 10
    end
    return rev
end

function check_mystery_num(n)
    # Let us iterate from 1 upto n/2, to check if there is any pair, such that thier
    # sum equals the given number and check if the pair are reverse of each other
    temp = n ÷ 2
    for i in 1:(temp)
        if (reverse(i) == (n - i))
            return true
        end
    end
    return false
end

print("Enter the number: ")
num = readline()
num = parse(Int, num)
res = check_mystery_num(num)
if res
    println("The given number $num is a Mystery Number.")
else
    println("The given number $num is not a Mystery Number.")
end

"""
Time Complexity: O(n*log(n)), where 'n' is the given number
Space Complexity: O(1)

SAMPLE INPUT AND OUTPUT

SAMPLE 1
Enter the number: 43234
The given number 43234 is a Mystery Number.

SAMPLE 2
Enter the number: 1234
The given number 1234 is not a Mystery Number.

"""


text:
include("includes.jl")

using TimeManager

function thread_printer()
	p=(
		"Printer",
		function()
			#init
      i=0
			while true
				global Messages
				pushMsg(p[1], "PRINT $i")
        showMessages()
				tsleep(1)
        i+=1
			end
			#cleanUp
		end,
	)
end

function thread_compute()
	p=(
		"Compute",
		function()
			timeRef = Ref(0.0)
      i=0
			while true
				if OnTime(0.25, timeRef)
					pushMsg(p[1], "CALC $i")
          i+=1
				end
				tsleep(0.001)
			end
		end,
	)
end

function thread_renderer()
	p=(
		"Renderer",
		function()
      timeRef = Ref(0.0)
      i=0
			while true
				if OnTime(0.5, timeRef)
					pushMsg(p[1], "RENDER $i")
          i+=1
				end
				tsleep(0.01)
			end
		end,
	)
end

function thread_sound()
	p=(
		"Sound",
		function()
      timeRef = Ref(0.0)
      i=0
			while true
				if OnTime(0.75, timeRef)
					pushMsg(p[1], "SOUND $i")
          i+=1
				end
				tsleep(0.1)
			end
		end,
	)
end

function main()
  #LoggerManager.log(()->begin
    println("Test Threads")
    list=Function[]
    push!(list, tinit(thread_printer()))
    push!(list, tinit(thread_compute()))
    push!(list, tinit(thread_renderer()))
    push!(list, tinit(thread_sound()))
    ThreadManager.run(list)
  #end)
end

main()

max_stars_repo_name: Gilga/JuliaOpenGL

max_issues_repo_name: Gilga/JuliaOpenGL

max_forks_repo_name: Gilga/JuliaOpenGL


text:
using MLDatasets: MNIST
using Flux: onehotbatch
using ProtoGrad: Conv, Linear, Dropout, Compose, maxpool, relu, softmax
using ProtoGrad: SupervisedObjective, cross_entropy, class_error
using ProtoGrad: Adam, forever, within_training_mode
using ProgressMeter: @showprogress
using StatsBase: sample
using Serialization: serialize, deserialize

train_x, train_y = MNIST.traindata(Float32, dir=joinpath(".", "datasets", "MNIST"))
test_x, test_y = MNIST.testdata(Float32, dir=joinpath(".", "datasets", "MNIST"))

train_x_with_channel = reshape(train_x, size(train_x)[1:2]..., 1, size(train_x)[3])
train_y_onehot = onehotbatch(train_y, 0:9)

height, width, input_channels, dataset_size = size(train_x_with_channel)
output_channels = 16
num_classes = 10

out_conv_size = (div(height, 4) - 3, div(width, 4) - 3, 16)

m = Compose(
    Dropout(0.1),
    Conv(input_channels=>6, (5, 5)),
    x -> relu.(x),                       # can also use ReLU()
    x -> maxpool(x, (2, 2)),             # can also use MaxPool((2, 2))
    Dropout(0.1),
    Conv(6=>output_channels, (5, 5)),
    x -> relu.(x),
    x -> maxpool(x, (2, 2)),
    x -> reshape(x, :, size(x, 4)),      # can also use Reshape((...))
    Dropout(0.1),
    Linear(prod(out_conv_size)=>120),
    x -> relu.(x),
    Dropout(0.1),
    Linear(120=>84),
    x -> relu.(x),
    Dropout(0.1),
    Linear(84=>num_classes),
    softmax,
)

batch_size = 128

training_batches = forever() do
    idx = sample(1:size(train_x_with_channel)[end], batch_size, replace = false)
    return (train_x_with_channel[:, :, :, idx], train_y_onehot[:, idx])
end

f = SupervisedObjective(cross_entropy, training_batches)

optimizer = Adam()
iterations = Iterators.Stateful(optimizer(m, f))

num_epochs = 20
num_batches_per_epoch = 500

within_training_mode() do
    @showprogress for epoch in 1:num_epochs
        for output in Iterators.take(iterations, num_batches_per_epoch)
            global m = output.solution
        end
    end
end

model_filename = "./serialized_conv_model"
println("saving model to $(model_filename)")
serialize(model_filename, m)

m_deserialized = deserialize(model_filename)

test_x_with_channel = reshape(test_x, size(test_x)[1:2]..., 1, size(test_x)[3])
test_y_onehot = onehotbatch(test_y, 0:9)

err = class_error(m_deserialized(test_x_with_channel), test_y_onehot)

println("test error = $(err * 100)%")


text:
using MahjongTiles

using Test

@testset "unit tests" begin

end


text:
"""Converts 100.000.000,00 to 100000000,00"""
@inline function remove_thousands_separator(value::T, ts::Nothing) :: T where {T<:AbstractString}
    return value
end

@inline function remove_thousands_separator(value::T, ts::Char) :: T where {T<:AbstractString}
    return replace(value, ts => "")
end

"""Converts 1000,00 to 1000.00"""
@inline function fix_decimal_separator(value::T, ds::Char) :: T where {T<:AbstractString}
    if ds != '.'
        return replace(value, ds => '.')
    else
        return value
    end
end

@inline function _parse_raw_value(raw_value::String, ::Type{T}, format::CSVFormat) where {T<:AbstractString}
    stripped = strip(raw_value)

    # strips '"' character if on both sides of string
    if startswith(stripped, '"') && endswith(stripped, '"')
        return strip(stripped, '"')
    else
        return stripped
    end
end

@inline function _parse_raw_value(raw_value::String, ::Type{T}, format::CSVFormat) where {T<:Number}
    value = remove_thousands_separator(raw_value, format.thousands_separator)
    value = fix_decimal_separator(value, format.decimal_separator)
    return parse(T, value)
end

@inline function _parse_raw_value(raw_value::String, ::Type{T}, format::CSVFormat) where {T<:Dates.TimeType}
    return Date(raw_value, format.date_format)
end

@inline function _parse_raw_value(raw_value::String, ::Type{Union{Missing, T}}, format::CSVFormat) where {T}
    if raw_value == format.missing_str
        return missing
    else
        return _parse_raw_value(raw_value, T, format)
    end
end

# doesn't know how to parse other types
function _parse_raw_value(raw_value::String, ::Type{T}, format::CSVFormat) where {T}
    error("Parsing $T not implemented.")
end

function _read_column(raw_column::Vector{String}, ::Type{T}, ROW_OFFSET::Int, FST_DATAROW_INDEX::Int,
    rows::Int, format::CSVFormat) where {T}

    col_data = _create_table_column(T, rows + ROW_OFFSET)
    @inbounds for r in FST_DATAROW_INDEX:rows
        r_ = r + ROW_OFFSET # r_ is the line index of the destination table. If raw contains a header, r_ = r - 1 . Otherwise, r_ = r
        col_data[r_] = _parse_raw_value(raw_column[r], T, format)
    end
    return col_data
end

const REGEX_WITH_QUOTES = r"^\s*\"\s*(?<str>.*)\s*\"\s*$"
const REGEX_WITHOUT_QUOTES = r"^\s*(?<str>.*)\s*$"

function _read_data!(table::Table, raw::Array{String,2}, format::CSVFormat; header::Bool=true)

    rows, cols = size(raw)

    # Check if header is consistent with schema
    @assert cols == length(table.schema.names) "CSV not consistent with given schema. ncols in file: $cols; ncols in schema: $(length(table.schema.names))."

    if header
        ROW_OFFSET = -1
    else
        ROW_OFFSET = 0
    end

    FST_DATAROW_INDEX = 1 - ROW_OFFSET # this is the index of the first data row in $raw

    col_array = Vector{Any}(undef, cols)
    @inbounds for col in 1:cols
        try
            col_type = table.schema.types[col]
            col_array[col] = _read_column(raw[:,col], col_type, ROW_OFFSET, FST_DATAROW_INDEX, rows, format)
        catch err
            @error("Error parsing column $col - $(table.schema.names[col]) <: $(table.schema.types[col])")
            rethrow()
        end
    end

    # Set table data without checks
    table.data = col_array
    return table
end

function _readraw(input, format::CSVFormat, use_mmap::Bool, skipstart, comments::Bool, comment_char::Char)
    return DelimitedFiles.readdlm(input, format.dlm, String; use_mmap=use_mmap, skipstart=skipstart, comments=comments, comment_char=comment_char)
end

"""
    readcsv(input, schema::Schema, format::CSVFormat=CSVFormat(); header::Bool=true, use_mmap::Bool=false, skipstart=0, comments::Bool=false, comment_char::Char='#')

`header` Tells if the input file has a header in the first line. Default is `true`.
"""
function readcsv(input, schema::Schema, format::CSVFormat=CSVFormat(); header::Bool=true, use_mmap::Bool=false, skipstart=0, comments::Bool=false, comment_char::Char='#')
    raw = _readraw(input, format, use_mmap, skipstart, comments, comment_char)
    tb = Table(schema)
    return _read_data!(tb, raw, format; header=header)
end

"""
    readcsv(input, format::CSVFormat=CSVFormat(); header::Bool=true, use_mmap::Bool=false, skipstart=0, comments::Bool=false, comment_char::Char='#')

Uses Schema inference.
"""
function readcsv(input, format::CSVFormat=CSVFormat(); header::Bool=true, use_mmap::Bool=false, skipstart=0, comments::Bool=false, comment_char::Char='#')
    raw = _readraw(input, format, use_mmap, skipstart, comments, comment_char)
    schema = infer_schema(raw, format, header)
    tb = Table(schema)
    return _read_data!(tb, raw, format; header=header)
end

function readcsv(input, types::Vector{T}, format::CSVFormat=CSVFormat(); use_mmap::Bool=false, skipstart=0, comments::Bool=false, comment_char::Char='#') where {T<:Type}
    raw = _readraw(input, format, use_mmap, skipstart, comments, comment_char)

    rows, cols = size(raw)
    @assert cols == length(types) "Number of cols in file is not $(length(types)): found $cols"

    # first row contains header
    header = [ Symbol(raw[1, col]) for col in 1:cols ]
    schema = Schema(header, types)
    table = Table(schema)

    return _read_data!(table, raw, format; header=true)
end

# Converts to string
function _write_string(io::IO, value::T, format::CSVFormat) where {T<:AbstractFloat}
    local result::String
    result = tostring(value)

    if format.decimal_separator != '.'
        result = replace(result, '.' => ',')
    end

    # TODO : support thousands_separator
    write(io, result)
end

_write_string(io::IO, value::T, format::CSVFormat) where {T<:Integer} = write(io, value)

function _write_string(io::IO, value::T, format::CSVFormat) where {T<:AbstractString}
    # Apply quotes if there´s a delimiter inside the string (replicate MS Excel behavior)
    if in(format.dlm, value)
        write(io, '"')
        write(io, value)
        write(io, '"')
    else
        write(io, value)
    end
end

_write_string(io::IO, value::Int, format::CSVFormat) = write(io, string(value))
_write_string(io::IO, value::T, format::CSVFormat) where {T<:Dates.TimeType} = write(io, Dates.format(value, format.date_format))
_write_string(io::IO, ::Missing, format::CSVFormat) = write(io, format.missing_str)

# Fallback
function _write_string(io::IO, value, format::CSVFormat)
    x = string(value)
    # Apply quotes if there´s a delimiter inside the string (replicate MS Excel behavior)
    if in(format.dlm, x)
        write(io, '"')
        write(io, x)
        write(io, '"')
    else
        write(io, x)
    end
end

function writecsv(filepath::String, tb::Union{AbstractDataFrame, Table}, format::CSVFormat=CSVFormat(); header::Bool=true)

    LB = '\n' # line break
    io = open(filepath, "w")
    rows, cols = size(tb)

    try
        # writes header to file
        if header
            h = names(tb)
            for i in eachindex(h)
                write(io, h[i])
                i != cols && write(io, format.dlm)
            end
            write(io, LB)
        end

        # data
        for r in 1:rows
            for c in 1:cols
                _write_string(io, tb[r,c], format)
                c != cols && write(io, format.dlm)
            end
            write(io, LB)
        end

    finally
        close(io)
    end

    nothing
end

# TODO: functions to append to existing Table
#=
function readcsv!(filepath::String, tb::Table;
    dlm::Char=';', null_str::String="", decimal_separator::Char=',', thousands_separator::Nullable{Char}=Nullable{Char}(), header=true)

end
=#

#function _read_column!(col_data::Vector{T}, raw_column::Vector{String}) where {T}

#end

#function _read_value!(col_data::NullableVector{T}, row_index::Int, value::String) where {T}

#end


text:
mutable struct Group
    memberCount::Int
    yeses::Array{Int}
    Group() = new(0, zeros(Int,1,26))
end

char2num(c::Char) = codepoint(c) - 96

num2char(n::Int) = Char(n + 96)

function addGroupMember(line::String, group::Group)
    for yes in line
        group.yeses[char2num(yes)] += 1
    end
    group.memberCount += 1
end

function addGroupMember2(line::String, group::Group)
    map(c -> (group.yeses[char2num(c)] += 1), collect(line))
    group.memberCount += 1
end

function countAnyYeses(group)
    result = 0
    for yes in group.yeses
        if yes >= 1
            result += 1
        end
    end
    return result
end

function countAnyYeses2(group)
    mapreduce(yes -> yes >= 1 ? 1 : 0, +, group.yeses)
end

function countAllYeses(group)
    result = 0
    for yes in group.yeses
        if yes == group.memberCount
            result += 1
        end
    end
    return result
end

function collectGroups(fileName)
    groups = []
    groupCounter = 1
    currentGroup = Group()
    open(fileName) do file
        for line in eachline(file)
            if length(line) == 0
                push!(groups, currentGroup)
                currentGroup = Group()
                groupCounter += 1
            else
                addGroupMember2(line, currentGroup)
            end
        end
    end
    return groups
end

groups = collectGroups("data.txt")

totalValidYeses = 0
for group in groups
    validGroupYeses = countAnyYeses2(group)
    global totalValidYeses += validGroupYeses
    println(group, " valid yeses =", validGroupYeses)
end
print("Total Valid Yeses = ", totalValidYeses)


text:
BrysonDenham_EXP=[:(x2[j]),:(u1[j])]; L=1/6;
```
using NLOptControl
BrysonDenham_EXP=[:(x2[j]),:(u1[j])]; L=1/6;
integrationConfig = :trapezoidal
```
@testset "BrysonDenham with (:integrationScheme=>$(integrationConfig)) using expressions)" for integrationConfig in integrationConfigs
  n=define(numStates=2,numControls=1,X0=[0.,1],XF=[0.,-1.],XL=[0.,NaN],XU=[L,NaN]);
  dynamics!(n,BrysonDenham_EXP)
  configure!(n;(:integrationScheme=>integrationConfig),(:finalTimeDV=>false),(:tf=>1.0));
  obj=integrate!(n,:(0.5*u1[j]^2));
  @NLobjective(n.ocp.mdl,Min,obj);
  optimize!(n);
  @show n.r.ocp.dfsOpt[:tSolve]
  @show 4/(9*L)
  @show n.r.ocp.objVal[1]
  @test isapprox(4/(9*L),n.r.ocp.objVal[1],atol=tol)
end

BrysonDenham_EXP=[:(x2[j]),:(u1[j])]; L=1/6;
@testset "BrysonDenham using slackVariables with (:integrationScheme=>$(integrationConfig)) using expressions)" for integrationConfig in integrationConfigs
  n=define(numStates=2,numControls=1,X0=[0.,1],XF=[0.,-1.],XL=[0.,NaN],XU=[L,NaN]);
  dynamics!(n,BrysonDenham_EXP)
  configure!(n;(:x0slackVariables=>true),(:xFslackVariables=>true),(:integrationScheme=>integrationConfig),(:finalTimeDV=>false),(:tf=>1.0));
  obj=integrate!(n,:(0.5*u1[j]^2));
  @NLobjective(n.ocp.mdl,Min,obj + 100*(n.ocp.x0s[1] + n.ocp.x0s[2] + n.ocp.xFs[1] + n.ocp.xFs[2]));
  optimize!(n);
  @show n.r.ocp.dfsOpt[:tSolve]
  @test isapprox(4/(9*L),n.r.ocp.objVal[1],atol=tol)
end

dx=[:(sin(x2[j])),:(u1[j])]
@testset "BeamProblem with (:integrationScheme=>$(integrationConfig)) using expressions and solver settings)" for integrationConfig in integrationConfigs
  n=define(numStates=2,numControls=1,XL=[-0.05,-1.0],XU=[-0.05,1.0]);
  dynamics!(n,dx)
  SS=[(:name=>:Ipopt),(:max_cpu_time=>5.0)]
  configure!(n;(:integrationScheme=>integrationConfig),(:finalTimeDV=>false),(:tf=>1.0),(:solverSettings=>SS));
  obj=integrate!(n,:( u1[j]^2 + 350*cos(x2[j]) ) )
  @NLobjective(n.ocp.mdl,Min,obj);
  optimize!(n);
  @show n.r.ocp.dfsOpt[:tSolve]
  @test isapprox(350,n.r.ocp.objVal[1],atol=tol)
end

dx=Array{Expr}(2);dx[1]=:(x[j]);dx[2]=:(u[j]-1.625);
@testset "MoonLander with (:integrationScheme=>$(integrationConfig)) using new names for states and controls; also naming them x and u to check for bugs)" for integrationConfig in integrationConfigs
  n=define(numStates=2,numControls=1,X0=[10.,-2],XF=[0.,0.],CL=[0.],CU=[3.]);
  states!(n,[:h,:x];descriptions=["h(t)","x(t)"]);
  controls!(n,[:u];descriptions=["u(t)"]);
  dynamics!(n,dx)
  configure!(n;(:integrationScheme=>integrationConfig),(:finalTimeDV=>true));
  obj=integrate!(n,:(u[j]));
  @NLobjective(n.ocp.mdl, Min, obj);
  optimize!(n);
  @show n.r.ocp.dfsOpt[:tSolve]
  @test isapprox(8.9253,n.r.ocp.objVal[1],atol=tol)
end

dx=Array{Expr}(2);dx[1]=:(x[j]);dx[2]=:(u[j]-1.625);
@testset "MoonLander with (:integrationScheme=>$(integrationConfig)) testing functionality to scale the states and controls)" for integrationConfig in integrationConfigs
  n=define(numStates=2,numControls=1,X0=[10.,-2],XF=[0.,0.],CL=[0.],CU=[3.],XS=[2.,.8],CS=[1.5]);
  states!(n,[:h,:x];descriptions=["h(t)","x(t)"]);
  controls!(n,[:u];descriptions=["u(t)"]);
  dynamics!(n,dx)
  configure!(n;(:integrationScheme=>integrationConfig),(:finalTimeDV=>true));
  obj=integrate!(n,:(u[j]));
  @NLobjective(n.ocp.mdl, Min, obj);
  optimize!(n);
  @show n.r.ocp.dfsOpt[:tSolve]
  @test isapprox(8.9253,n.r.ocp.objVal[1],atol=tol)
end

############################
# Benchmarking Test
v0 = -2
h0 = 10
pts = 10000
tf_star = 2/3*v0 +4/3*(1/2*v0^2+3/2*h0)^0.5
ts = tf_star/2+v0/3

function optimal_solution(t)
  if t < ts
    h = -3/4*t^2 + v0*t + h0
    v = -3/2*t + v0
    u = 0
  else
    h = 3/4*t^2 + (-3*ts + v0)*t + 3/2*ts^2 + h0
    v = 3/2*t + (-3*ts + v0)
    u = 3
  end
  return h, v, u
end

t_opt = linspace(0,tf_star,pts)
h_opt = zeros(pts); v_opt = zeros(pts); u_opt = zeros(pts);
for i in 1:pts
  h, v, u = optimal_solution(t_opt[i])
  h_opt[i] = h; v_opt[i] = v; u_opt[i] = u;
end

opt_runs = 2
col_pts = 10:2:12
Nck_vec = [[col_pts[i]] for i in 1:length(col_pts)]
opt_num = length(Nck_vec)

@testset "MoonLander test to make sure that a benchmark test will be able to be executed with (:integrationScheme=>$(integrationConfig)) " for integrationConfig in integrationConfigs

  # final average results
  t_opt_ave = NaN*zeros(opt_num)
  h_error_ave = NaN*zeros(opt_num)
  v_error_ave = NaN*zeros(opt_num)
  u_error_ave = NaN*zeros(opt_num)
  max_error_ave = NaN*zeros(opt_num)

  for num in 1:length(col_pts)
    n=define(numStates=2,numControls=1,X0=[h0,v0],XF=[0.,0.],XL=[-20,-20],XU=[20,20],CL=[0.],CU=[3.]);
    n.s.ocp.interpolationOn = true
    n.s.ocp.tfMax = 1000.0
    n.s.ocp.numInterpPts = pts  # must be the same as above to calculate error
    n.s.ocp.tfOptimal = tf_star # must be the same as above to calculate error
    states!(n,[:h,:v];descriptions=["h(t)","v(t)"]);
    controls!(n,[:T];descriptions=["T(t)"]);
    dx=[:(v[j]),:(T[j]-1.5)]
    dynamics!(n,dx)

    if (integrationConfig == :lgrImplicit || integrationConfig == :lgrExplicit)
      Nck = Nck_vec[num]
      configure!(n;(:integrationScheme=>integrationConfig),(:Nck=>Nck),(:finalTimeDV=>true));
    else
      configure!(n;(:integrationScheme=>integrationConfig),(:N=>col_pts[num]),(:finalTimeDV=>true));
    end
    x1=n.r.ocp.x[:,1];x2=n.r.ocp.x[:,2];
    obj=integrate!(n,:(T[j]));
    @NLobjective(n.ocp.mdl, Min, obj);
    setvalue(n.ocp.tf, 1.5)
#    for i in 1:length(x1); setvalue(x1[i], 0.0); setvalue(x2[i], 0.0);  end
    # cache functions; inital optimization
    optimize!(n);

    # RESET temp variables for averaging results
    t_solve = NaN*zeros(opt_runs)
    h_error = NaN*zeros(opt_runs)
    v_error = NaN*zeros(opt_runs)
    u_error = NaN*zeros(opt_runs)
    max_error = NaN*zeros(opt_runs)

    for j in 1:opt_runs
      optimize!(n);
      if n.r.ocp.status == :Optimal
        t_solve[j] = n.r.ocp.tSolve
        h_error[j] = maximum(abs.(n.r.ocp.Xpts[:,1] - h_opt))
        v_error[j] = maximum(abs.(n.r.ocp.Xpts[:,2] - v_opt))
        u_error[j] = maximum(abs.(n.r.ocp.Upts[:,1] - u_opt))
        max_error[j] = maximum([h_error[j], v_error[j]])
      end
    end

    t_opt_ave[num] = mean(t_solve[.!isnan.(t_solve)],1)[1]
    h_error_ave[num] = mean(h_error[.!isnan.(h_error)],1)[1]
    v_error_ave[num] = mean(v_error[.!isnan.(v_error)],1)[1]
    u_error_ave[num] = mean(u_error[.!isnan.(u_error)],1)[1]
    max_error_ave[num] = mean(max_error[.!isnan.(max_error)],1)[1]
  end
  @show maximum(max_error_ave)
  @test isapprox(0, maximum(max_error_ave),atol=big_tol)
end

# Benchmarking Test
############################


text:
using SparseArrays: SparseMatrixCSC, sparse
using LinearAlgebra: mul!
using Combinatorics: permutations
using LinearAlgebra: rank, svd, Diagonal

"""
`struct SymmetricBasis`

### Constructors

Option 1: pass a `PIBasis`
```julia
SymmetricBasis(pibasis, φ)
```
All possible permutation-invariant basis functions will be symmetrised and
then reduced to a basis (rather than spanning set)

Option 2: pass a `OneParticleBasis`
```julia
SymmetricBasis(φ, basis1p, maxν, maxdeg;
               Deg = NaiveTotalDegree())
```
will first construct a `PIBasis` from these inputs and then call the first
constructor.
"""
struct SymmetricBasis{BOP, PROP} <: ACEBasis
   pibasis::PIBasis{BOP}
   A2Bmap::SparseMatrixCSC{PROP, Int}
end

Base.length(basis::SymmetricBasis{BOP, PROP}) where {BOP, PROP} =
      size(basis.A2Bmap, 1)

fltype(basis::SymmetricBasis{BOP, PROP}) where {BOP, PROP} =  PROP
# rfltype(basis::SymmetricBasis) = rfltype(basis.pibasis)

SymmetricBasis(φ::AbstractProperty, args...; kwargs...) =
      SymmetricBasis(PIBasis(args...; kwargs..., property = φ), φ)

function SymmetricBasis(pibasis, φ::TP) where {TP}

   # AA index -> AA spec
   AAspec = get_spec(pibasis)

   # construct the reverse mapping AA spec -> iAA and A spec -> iA
   invAAspec = Dict{Any, Int}()
   for (iAA, AA) in enumerate(AAspec)
      invAAspec[AA] = iAA
   end
   invAspec = Dict{Any, Int}()
   for (iA, A) in enumerate(get_spec(pibasis.basis1p))
      invAspec[A] = iA
   end

   # FUTURE: here we could analyze the different symbols and choose
   # which ones are of the "m"-type. e.g. there might be m1, m2 symbols. We
   # could create the convention that any symbol starting with m will be
   # treated like an m from the Ylm basis...

   # allocate the datastructure that computes and caches the
   # coupling coefficients
   # TODO: should this be stored with the basis?
   #       or maybe written to a file on disk? and then flushed every time
   #       we finish with a basis construction???
   #rotc = Rot3DCoeffs(rfltype(pibasis))
   rotc = Rot3DCoeffs(φ, rfltype(pibasis))
   # allocate triplet format
   Irow, Jcol, vals = Int[], Int[], TP[]
   # count the number of PI basis functions = number of rows
   idxB = 0

   # loop through AA basis, but skip most of them ...
   for (iAA, AA) in enumerate(AAspec)
      # AA = [b1, b2, ...], each bi = (n = ..., l = .., m = ...)
      # skip it unless all m are zero, because we want to consider each
      # (nn, ll, ...) block only once.
      # the loop over all possible `mm` must be taken care of inside
      # the `coupling_coeffs` implementation
      if !all(b.m == 0 for b in AA)
         continue
      end
      # get the coupling coefficients
      # here, we could help out a bit and make it easier to find the
      # the relevant basis functions?
      # AAcols will be in spec format i.e. named tuples
      U, AAcols = coupling_coeffs(AA, rotc)
      # loop over the rows of U -> each specifies a basis function
      for irow = 1:size(U, 1)
         idxB += 1
         # loop over the columns of U / over brows
         for (icol, bcol) in enumerate(AAcols)
            # put bcol into the correct order
            bcol_ordered = _get_ordered(bcol, invAspec)
            # this is a subtle step: bcol and bcol_ordered are equivalent
            # permutation-invariant basis functions. This means we will
            # add the same PI basis function several times, but in the call to
            # `sparse` the values will just be added.
            if !haskey(invAAspec, bcol_ordered)
               @show bcol_ordered
               @show degree(bcol_ordered, NaiveTotalDegree(), pibasis.basis1p)
            end
            idxAA = invAAspec[bcol_ordered]
            push!(Irow, idxB)
            push!(Jcol, idxAA)
            push!(vals, U[irow, icol])
         end
      end
   end
   # TODO: filter and throw out everything that hasn't been used!!
   # create CSC: [   triplet    ]  nrows   ncols
   A2Bmap = sparse(Irow, Jcol, vals, idxB, length(AAspec))
   return SymmetricBasis(pibasis, A2Bmap)
end

"""
produce the ordered tuple defining the AA basis function uniquely.
"""
function _get_ordered(bb, invAspec)
   iAs = [invAspec[b] for b in bb]
   return bb[ sortperm(iAs, rev = true) ]
end

function coupling_coeffs(bb, rotc::Rot3DCoeffs)
   # bb = [ b1, b2, b3, ...)
   # bi = (μ = ..., n = ..., l = ..., m = ...)
   #    (μ, n) -> n; only the l and m are used in the angular basis
   if length(bb) == 0
      # return [1.0,], [bb,]
		error("correlation order 0 is currently not allowed")
   end
   # convert to a format that the Rotations3D implementation can understand
   # this utility function splits the bb = (b1, b2 ...) with each
   # b1 = (μ = ..., n = ..., l = ..., m = ...) into
   #    l, and a new n = (μ, n)
   ll, nn = _b2llnn(bb)
   # now we can call the coupling coefficient construiction!!
   U, Ms = rpe_basis(rotc, nn, ll)

   # but now we need to convert the m spec back to complete basis function
   # specifications
   rpebs = [ _nnllmm2b(bb[1], nn, ll, mm) for mm in Ms ]

   return U, rpebs
end

function rpe_basis(A::Rot3DCoeffs,
						 nn::SVector{N, TN},
						 ll::SVector{N, Int}) where {N, TN}
	Ure, Mre = Rotations3D.re_basis(A, ll)
	G = _gramian(nn, ll, Ure, Mre)
   S = svd(G)
   rk = rank(Diagonal(S.S); rtol =  1e-7)
	Urpe = S.U[:, 1:rk]'
	return Diagonal(sqrt.(S.S[1:rk])) * Urpe * Ure, Mre
end

function _gramian(nn, ll, Ure, Mre)
   N = length(nn)
   nre = size(Ure, 1)
   G = zeros(Complex{Float64}, nre, nre)
   for σ in permutations(1:N)
      if (nn[σ] != nn) || (ll[σ] != ll); continue; end
      for (iU1, mm1) in enumerate(Mre), (iU2, mm2) in enumerate(Mre)
         if mm1[σ] == mm2
            for i1 = 1:nre, i2 = 1:nre
               G[i1, i2] += coco_dot(Ure[i1, iU1], Ure[i2, iU2])
            end
         end
      end
   end
   return G
end

_nnllmm2b(b, nn, ll, mm) = [ _nlm2b(b, n, l, m) for (n, l, m) in zip(nn, ll, mm) ]

@generated function _nlm2b(b::NamedTuple{ALLKEYS}, n::NamedTuple{NKEYS}, l, m) where {ALLKEYS, NKEYS}
   code =
      ( "( _b = (" * prod("$(k) = n.$(k), " for k in NKEYS) * "l = l, m = m ); "
         *
        " b = (" * prod("$(k) = _b.$(k), " for k in ALLKEYS) * ") )" )
   :( $(Meta.parse(code)) )
end

function _b2llnn(bb)
   @assert all( iszero(b.m) for b in bb )
   ll = SVector( [b.l for b in bb]... )
   nn = SVector( [_all_but_lm(b) for b in bb]... )
   return ll, nn
end

"""
return a names tuple containing all values in b except those corresponding
to l and m keys
"""
@generated function _all_but_lm(b::NamedTuple{NAMES}) where {NAMES}
   code = "n = ("
   for k in NAMES
      if !(k in [:l, :m])
         code *= "$(k) = b.$(k), "
      end
   end
   code *= ")"
   quote
      $(Meta.parse(code))
      n
   end
end

# ---------------- A modified sparse matmul

using SparseArrays: AbstractSparseMatrixCSC, DenseInputVecOrMat,
				        nonzeros, rowvals, nzrange

function genmul!(C::StridedVecOrMat, A::AbstractSparseMatrixCSC, B::DenseInputVecOrMat, mulop)
    size(A, 2) == size(B, 1) || throw(DimensionMismatch())
    size(A, 1) == size(C, 1) || throw(DimensionMismatch())
    size(B, 2) == size(C, 2) || throw(DimensionMismatch())
    nzv = nonzeros(A)
    rv = rowvals(A)
    fill!(C, zero(eltype(C)))
    for k in 1:size(C, 2)
        @inbounds for col in 1:size(A, 2)
            αxj = B[col,k]
            for j in nzrange(A, col)
                C[rv[j], k] += mulop(nzv[j], αxj)
            end
        end
    end
    return C
end

# for (T, t) in ((Adjoint, adjoint), (Transpose, transpose))
#    @eval function mul!(C::StridedVecOrMat, xA::$T{<:Any,<:AbstractSparseMatrixCSC}, B::DenseInputVecOrMat, α::Number, β::Number)

using LinearAlgebra: Transpose

function genmul!(C::StridedVecOrMat, xA::Transpose{<:Any,<:AbstractSparseMatrixCSC}, B::DenseInputVecOrMat, mulop)
   A = xA.parent
   size(A, 2) == size(C, 1) || throw(DimensionMismatch())
   size(A, 1) == size(B, 1) || throw(DimensionMismatch())
   size(B, 2) == size(C, 2) || throw(DimensionMismatch())
   nzv = nonzeros(A)
   rv = rowvals(A)
   for k in 1:size(C, 2)
       @inbounds for col in 1:size(A, 2)
           tmp = zero(eltype(C))
           for j in nzrange(A, col)
               tmp += mulop(nzv[j], B[rv[j],k])
           end
           C[col,k] += tmp
       end
   end
   return C
end

# ---------------- Evaluation code

alloc_temp(basis::SymmetricBasis, args...) =
      (  AA = alloc_B(basis.pibasis),
         tmppi = alloc_temp(basis.pibasis) )

alloc_B(basis::SymmetricBasis, args...) =
      zeros(fltype(basis), length(basis))

function evaluate!(B, tmp, basis::SymmetricBasis,
                   cfg::AbstractConfiguration)
   # compute AA
   evaluate!(tmp.AA, tmp.tmppi, basis.pibasis, cfg)
   evaluate!(B, tmp, basis, tmp.AA)
   return B
end

# this function allows us to attach multiple symmetric bases to a single
#  𝑨 basis
#  TODO: this is extremely inefficient for multiple species But it is simple
#        and clean and will do for now...
function evaluate!(B, tmp, basis::SymmetricBasis,
                   AA::AbstractVector{<: Number})
   genmul!(B, basis.A2Bmap, AA, *)
end

# ---- gradients

function gradtype(basis::SymmetricBasis)
   φ = zero(eltype(basis.A2Bmap))
   dAA = zero(gradtype(basis.pibasis))
   return typeof(coco_o_daa(φ, dAA))
end

alloc_temp_d(basis::SymmetricBasis, nmax::Integer) =
      (  AA = alloc_B(basis.pibasis),
         dAA = alloc_dB(basis.pibasis, nmax),
         tmppi = alloc_temp(basis.pibasis),
         tmpdpi = alloc_temp_d(basis.pibasis, nmax) )

alloc_dB(basis::SymmetricBasis, nmax::Integer) =
      zeros(gradtype(basis), length(basis), nmax)

function evaluate_d!(dB, tmpd, basis::SymmetricBasis,
                     cfg::AbstractConfiguration)
   # compute AA
   evaluate_ed!(tmpd.AA, tmpd.dAA, tmpd.tmpdpi, basis.pibasis, cfg)
   evaluate_d!(dB, tmpd, basis, tmpd.AA, tmpd.dAA)
   return dB
end

function evaluate_d!(dB, tmpd, basis::SymmetricBasis,
                     AA::AbstractVector{<: Number}, dAA)
   genmul!(dB, basis.A2Bmap, dAA, ACE.coco_o_daa)
end


text:
@model function model_parametric_cases(
    num_impute,        # [Int] num. of days for which to impute infections
    num_total_days,    # [Int] days of observed data + num. of days to forecast
    cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
    deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
    π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    epidemic_start,    # [AbstractVector{<:Int}]
    population,        # [AbstractVector{<:Real}]
    serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
	num_iar_steps,         # [Int] number of weeks (21)
    iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
    lockdown,         # [Int] number of weeks (21)
    num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
	weekday,
	holiday,
	covariates,
	covariates_start,
	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    hospit,
	seroprev_mean,
    seroprev_std,
    seroprev_idx,
	π4,
    predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    ::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	invlink = log,#KLogistic(V(3))
	link    = exp,#KLogit(V(k))
) where {TV, V}
    ############ 0.) some definitions
	num_obs        = length(cases)
    num_Rt_steps   = covariates_start-lockdown-1
	num_sero       = length(seroprev_mean)
	num_covariates = size(covariates, 2)
    cases_start    = num_obs-num_case_obs
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
    num_time_steps = predict ? num_total_days : num_obs

    ############# 2.) time varying reproduction number
	R0        ~ truncated(Normal(3., 1.), 1., 5.)
	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
	σ_rt      ~ truncated(Normal(0.05, .03), 0, .3)
	latent_Rt ~ RandomWalk(num_Rt_steps, σ_rt, invlink(R1))
	α         ~ filldist(Exponential(0.2), num_covariates)

    Rt                                = TV(undef, num_time_steps)
    Rt[1:lockdown]                   .= R0
    Rt[lockdown+1:covariates_start-1] = link.(latent_Rt)
    Rt[covariates_start:num_obs]      = link.(
        latent_Rt[end] .+ covariates * α
    )
    if predict #if prediction
		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
	end

    ############ 3.) infection dynamics
	## 3.1) initialize serial interval
	num_si   = 15
	si       = serial_intervals
	# μ_si   ~ truncated(Normal(5.06, 0.33), 0, Inf)
	# σ_si   = 2.11
	# si     = pdf.( Ref(GammaMeanStd(μ_si, σ_si)), 1:num_si )
	# si    /= sum(si)

	## 3.2) seed epidemic
	τ ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
	T = typeof(τ)
    y ~ truncated(Exponential(τ),T(0),T(1000))

	newly_infected         = TV(undef, num_time_steps)
	cumulative_infected    = TV(undef, num_time_steps)
	newly_infected[1]      = y
	cumulative_infected[1] = zero(V)
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
	end

    ## 3.3) run disease dynamics
	for t = (num_impute + 1):num_time_steps
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * si[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		# number of new infections (unobserved)
		newly_infected[t] = St * Rt[t] * effectively_infectious
	end

	########### 4.) derive observables
	## 4.1) initialize delay distributions
	num_i2c = 40
	num_i2h = 40
	num_i2d = 60

	μ_i2c ~ truncated(Normal(4.5, 1.), 3, 7)
    μ_i2h ~ truncated(Normal(12., 3.), 9, 16)
    μ_i2d ~ truncated(Normal(20., 4.), 10, 30)

	ϕ_i2c = 5.41
	ϕ_i2h = 5.41
	ϕ_i2d = 14.26

	i2c = pdf.( Ref(NegativeBinomial2(μ_i2c, ϕ_i2c)), 1:num_i2c )
	i2h = pdf.( Ref(NegativeBinomial2(μ_i2h, ϕ_i2h)), 1:num_i2h )
	i2d = pdf.( Ref(NegativeBinomial2(μ_i2d, ϕ_i2d)), 1:num_i2d )

	i2c /= sum(i2c)
	i2h /= sum(i2h)
	i2d /= sum(i2d)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	num_days = 7
	weekdayeffect = TV(undef, num_days)
	for i in 1:num_days
		weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end

	### holiday effect
	holidayeffect ~ Beta(1,1)

    ### 4.2) derive observables from infections
	ifr       ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
    iar       ~ Beta(1,10)

    expected_daily_cases     = TV(undef, num_time_steps)
    expected_daily_hospit    = TV(undef, num_time_steps)
    expected_daily_deaths    = TV(undef, num_time_steps)
	expected_seropositive    = TV(undef, num_time_steps)

	expected_daily_cases[1]  = 1e-15 * newly_infected[1]
    expected_daily_hospit[1] = 1e-15 * newly_infected[1]
    expected_daily_deaths[1] = 1e-15 * newly_infected[1]
	for t = 2:num_time_steps
		periodiceffect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		expected_daily_cases[t]  = iar * sum(newly_infected[τ] * i2c[t - τ] for τ = (t - 1):-1:max(t-num_i2c,1)) * periodiceffect
        expected_daily_hospit[t] = ihr * sum(newly_infected[τ] * i2h[t - τ] for τ = (t - 1):-1:max(t-num_i2h,1))
        expected_daily_deaths[t] = ifr * sum(newly_infected[τ] * i2d[t - τ] for τ = (t - 1):-1:max(t-num_i2d,1))
	end

	serodelay = 16 # 14d self deferral + 16d seroconversion (see Erikstrup et al. https://doi.org/10.1093/cid/ciaa849)
	expected_seropositive[1:serodelay]    .= zero(V)
	expected_seropositive[serodelay+1:end] = cumulative_infected[1:end-serodelay]

    ########### 4.) compare model to observations
	## 4.1) observe cases
	ϕ_c  ~ truncated(Normal(25, 10), 0, Inf)
	ts_c = cases_start:num_obs
	μs_c = expected_daily_cases[ts_c]
	!all( 0 .< μs_c .< population ) && ( @error "expected_daily_cases < 0"; Turing.@addlogprob! -Inf; return)
	cases[ts_c] ~ arraydist(NegativeBinomial2.(μs_c, ϕ_c))

    ## 4.3) observe deaths
	ϕ_d   ~ truncated(Normal(20, 10), 0, Inf)
    ts_d  = epidemic_start:cases_start-1
	μs_d  = expected_daily_deaths[ts_d]
	!all( 0 .< μs_d .< population ) && (@error "expected_daily_deaths"; Turing.@addlogprob! -Inf; return)
	deaths[ts_d] ~ arraydist(NegativeBinomial2.(μs_d, ϕ_d))

    ## 4.3) seroprevalence study
    select_seroprev = 2:4
	seroprev_σ ~ arraydist( InverseGamma2.(seroprev_std[ select_seroprev ], Ref(.1)) )
    seroprev_μ = expected_seropositive[ seroprev_idx[ select_seroprev ] ]
	seroprev_mean[select_seroprev] ~ MvNormal(seroprev_μ, seroprev_σ) # select only early observations

    return (
		expected_daily_cases  = expected_daily_cases,
        expected_daily_deaths = expected_daily_deaths,
        expected_daily_hospit = expected_daily_hospit,
        Rt = Rt,
		expected_seropositive = expected_seropositive,
		iar = iar,
    )
end

## ============================================================================
#                           parametetric model with hospit
# =============================================================================
@model function model_parametric_hospit(
    num_impute,        # [Int] num. of days for which to impute infections
    num_total_days,    # [Int] days of observed data + num. of days to forecast
    cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
    deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
    π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    epidemic_start,    # [AbstractVector{<:Int}]
    population,        # [AbstractVector{<:Real}]
    serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
	num_iar_steps,         # [Int] number of weeks (21)
    iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
    lockdown,         # [Int] number of weeks (21)
    num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
	weekday,
	holiday,
	covariates,
	covariates_start,
	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    hospit,
	seroprev_mean,
    seroprev_std,
    seroprev_idx,
	π4,
    predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    ::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	invlink = log,#KLogistic(V(3))
	link    = exp,#KLogit(V(k))
) where {TV, V}
    ############ 0.) some definitions
	num_obs        = length(cases)
    num_Rt_steps   = covariates_start-lockdown-1
	num_sero       = length(seroprev_mean)
	num_covariates = size(covariates, 2)
    cases_start    = num_obs-num_case_obs
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
    num_time_steps = predict ? num_total_days : num_obs

    ############# 2.) time varying reproduction number
	R0        ~ truncated(Normal(3., 1.), 1., 5.)
	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
	σ_rt      ~ truncated(Normal(0.05, .03), 0, .3)
	latent_Rt ~ RandomWalk(num_Rt_steps, σ_rt, invlink(R1))
	α         ~ filldist(Exponential(0.2), num_covariates)

    Rt                                = TV(undef, num_time_steps)
    Rt[1:lockdown]                   .= R0
    Rt[lockdown+1:covariates_start-1] = link.(latent_Rt)
    Rt[covariates_start:num_obs]      = link.(
        latent_Rt[end] .+ covariates * α
    )
    if predict #if prediction
		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
	end

    ############ 3.) infection dynamics
	## 3.1) initialize serial interval
	num_si   = 15
	si       = serial_intervals
	# μ_si   ~ truncated(Normal(5.06, 0.33), 0, Inf)
	# σ_si   = 2.11
	# si     = pdf.( Ref(GammaMeanStd(μ_si, σ_si)), 1:num_si )
	# si    /= sum(si)

	## 3.2) seed epidemic
	τ ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
	T = typeof(τ)
    y ~ truncated(Exponential(τ),T(0),T(1000))

	newly_infected         = TV(undef, num_time_steps)
	cumulative_infected    = TV(undef, num_time_steps)
	newly_infected[1]      = y
	cumulative_infected[1] = zero(V)
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
	end

    ## 3.3) run disease dynamics
	for t = (num_impute + 1):num_time_steps
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * si[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		# number of new infections (unobserved)
		newly_infected[t] = St * Rt[t] * effectively_infectious
	end

	########### 4.) derive observables
	## 4.1) initialize delay distributions
	num_i2c = 40
	num_i2h = 40
	num_i2d = 60

	μ_i2c ~ truncated(Normal(4.5, 1.), 3, 7)
    μ_i2h ~ truncated(Normal(12., 3.), 9, 16)
    μ_i2d ~ truncated(Normal(20., 4.), 10, 30)

	ϕ_i2c = 5.41
	ϕ_i2h = 5.41
	ϕ_i2d = 14.26

	i2c = pdf.( Ref(NegativeBinomial2(μ_i2c, ϕ_i2c)), 1:num_i2c )
	i2h = pdf.( Ref(NegativeBinomial2(μ_i2h, ϕ_i2h)), 1:num_i2h )
	i2d = pdf.( Ref(NegativeBinomial2(μ_i2d, ϕ_i2d)), 1:num_i2d )

	i2c /= sum(i2c)
	i2h /= sum(i2h)
	i2d /= sum(i2d)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	num_days = 7
	weekdayeffect = TV(undef, num_days)
	for i in 1:num_days
		weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end

	### holiday effect
	holidayeffect ~ Beta(1,1)

    ### 4.2) derive observables from infections
	ifr       ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
    iar       ~ Beta(1,10)

    expected_daily_cases     = TV(undef, num_time_steps)
    expected_daily_hospit    = TV(undef, num_time_steps)
    expected_daily_deaths    = TV(undef, num_time_steps)
	expected_seropositive    = TV(undef, num_time_steps)

	expected_daily_cases[1]  = 1e-15 * newly_infected[1]
    expected_daily_hospit[1] = 1e-15 * newly_infected[1]
    expected_daily_deaths[1] = 1e-15 * newly_infected[1]
	for t = 2:num_time_steps
		periodiceffect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		expected_daily_cases[t]  = iar * sum(newly_infected[τ] * i2c[t - τ] for τ = (t - 1):-1:max(t-num_i2c,1)) * periodiceffect
        expected_daily_hospit[t] = ihr * sum(newly_infected[τ] * i2h[t - τ] for τ = (t - 1):-1:max(t-num_i2h,1))
        expected_daily_deaths[t] = ifr * sum(newly_infected[τ] * i2d[t - τ] for τ = (t - 1):-1:max(t-num_i2d,1))
	end

	serodelay = 16 # 14d self deferral + 16d seroconversion (see Erikstrup et al. https://doi.org/10.1093/cid/ciaa849)
	expected_seropositive[1:serodelay]    .= zero(V)
	expected_seropositive[serodelay+1:end] = cumulative_infected[1:end-serodelay]

    ########### 4.) compare model to observations
	# 4.2) observe hospitalizations
	ϕ_h   ~ truncated(Normal(25, 10), 0, Inf)
	ts_h  = epidemic_start:num_obs
	μs_h  = expected_daily_hospit[ts_h]
	!all( 0 .< μs_h .< population ) && (@error "expected_daily_hospit"; Turing.@addlogprob! -Inf; return)
	hospit[ts_h] ~ arraydist(NegativeBinomial2.(μs_h, ϕ_h))

    ## 4.3) seroprevalence study
    select_seroprev = 2:4
	seroprev_σ ~ arraydist( InverseGamma2.(seroprev_std[ select_seroprev ], Ref(.1)) )
    seroprev_μ = expected_seropositive[ seroprev_idx[ select_seroprev ] ]
	seroprev_mean[select_seroprev] ~ MvNormal(seroprev_μ, seroprev_σ) # select only early observations

    return (
		expected_daily_cases  = expected_daily_cases,
        expected_daily_deaths = expected_daily_deaths,
        expected_daily_hospit = expected_daily_hospit,
        Rt = Rt,
		expected_seropositive = expected_seropositive,
		iar = iar,
    )
end

## ==========================================================================
#                       parametric model with deaths
# ===========================================================================
@model function model_parametric_deaths(
    num_impute,        # [Int] num. of days for which to impute infections
    num_total_days,    # [Int] days of observed data + num. of days to forecast
    cases,            # [AbstractVector{<:AbstractVector{<:Int}}] reported infected
    deaths,            # [AbstractVector{<:AbstractVector{<:Int}}] reported deaths; rows indexed by i > N contain -1 and should be ignored
    π,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    π2,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    epidemic_start,    # [AbstractVector{<:Int}]
    population,        # [AbstractVector{<:Real}]
    serial_intervals,  # [AbstractVector{<:Real}] fixed pre-calculated serial interval (SI) using empirical data from Neil
	num_iar_steps,         # [Int] number of weeks (21)
    iar_index,        # [Vector{Array{Int64,1}}] Macro region index for each state
    lockdown,         # [Int] number of weeks (21)
    num_case_obs,       # [Int] 21 = max_date ("2020-05-11") - testing_date (2020-05-11) see publication
	weekday,
	holiday,
	covariates,
	covariates_start,
	π3,                 # [AbstractVector{<:AbstractVector{<:Real}}] h * s
    hospit,
	seroprev_mean,
    seroprev_std,
    seroprev_idx,
	π4,
    predict=false,     # [Bool] if `false`, will only compute what's needed to `observe` but not more
    ::Type{TV} = Vector{Float64},
	::Type{V}  = Float64;
	invlink = log,#KLogistic(V(3))
	link    = exp,#KLogit(V(k))
) where {TV, V}
    ############ 0.) some definitions
	num_obs        = length(cases)
    num_Rt_steps   = covariates_start-lockdown-1
	num_sero       = length(seroprev_mean)
	num_covariates = size(covariates, 2)
    cases_start    = num_obs-num_case_obs
	# If we don't want to predict the future, we only need to compute up-to time-step `num_obs_states[m]`
    num_time_steps = predict ? num_total_days : num_obs

    ############# 2.) time varying reproduction number
	R0        ~ truncated(Normal(3., 1.), 1., 5.)
	R1        ~ truncated(Normal(.8, .1), .5, 1.1)
	σ_rt      ~ truncated(Normal(0.05, .03), 0, .3)
	latent_Rt ~ RandomWalk(num_Rt_steps, σ_rt, invlink(R1))
	α         ~ filldist(Exponential(0.2), num_covariates)

    Rt                                = TV(undef, num_time_steps)
    Rt[1:lockdown]                   .= R0
    Rt[lockdown+1:covariates_start-1] = link.(latent_Rt)
    Rt[covariates_start:num_obs]      = link.(
        latent_Rt[end] .+ covariates * α
    )
    if predict #if prediction
		Rt[num_obs+1:num_total_days] .= Rt[num_obs]
	end

    ############ 3.) infection dynamics
	## 3.1) initialize serial interval
	num_si   = 15
	si       = serial_intervals
	# μ_si   ~ truncated(Normal(5.06, 0.33), 0, Inf)
	# σ_si   = 2.11
	# si     = pdf.( Ref(GammaMeanStd(μ_si, σ_si)), 1:num_si )
	# si    /= sum(si)

	## 3.2) seed epidemic
	τ ~ Exponential(1 / 0.03) # `Exponential` has inverse parameterization of the one in Stan
	T = typeof(τ)
    y ~ truncated(Exponential(τ),T(0),T(1000))

	newly_infected         = TV(undef, num_time_steps)
	cumulative_infected    = TV(undef, num_time_steps)
	newly_infected[1]      = y
	cumulative_infected[1] = zero(V)
	for t in 2:num_impute
		newly_infected[t] = y
		cumulative_infected[t] = cumulative_infected[t-1] + y
	end

    ## 3.3) run disease dynamics
	for t = (num_impute + 1):num_time_steps
		# Update cumulatively infected
		cumulative_infected[t] = cumulative_infected[t-1] + newly_infected[t - 1]
		# Adjusts for portion of population that is susceptible
		St = max(population - cumulative_infected[t], 0) / population
		# effective number of infectious individuals
		effectively_infectious = sum(newly_infected[τ] * si[t - τ] for τ = (t - 1):-1:max(t-num_si,1))
		# number of new infections (unobserved)
		newly_infected[t] = St * Rt[t] * effectively_infectious
	end

	########### 4.) derive observables
	## 4.1) initialize delay distributions
	num_i2c = 40
	num_i2h = 40
	num_i2d = 60

	μ_i2c ~ truncated(Normal(4.5, 1.), 3, 7)
    μ_i2h ~ truncated(Normal(12., 3.), 9, 16)
    μ_i2d ~ truncated(Normal(20., 4.), 10, 30)

	ϕ_i2c = 5.41
	ϕ_i2h = 5.41
	ϕ_i2d = 14.26

	i2c = pdf.( Ref(NegativeBinomial2(μ_i2c, ϕ_i2c)), 1:num_i2c )
	i2h = pdf.( Ref(NegativeBinomial2(μ_i2h, ϕ_i2h)), 1:num_i2h )
	i2d = pdf.( Ref(NegativeBinomial2(μ_i2d, ϕ_i2d)), 1:num_i2d )

	i2c /= sum(i2c)
	i2h /= sum(i2h)
	i2d /= sum(i2d)

	### day of week effect
	weekdayeffect_simplex ~ Dirichlet([7,7,7,7,7,3,1])
	num_days = 7
	weekdayeffect = TV(undef, num_days)
	for i in 1:num_days
		weekdayeffect[i] = weekdayeffect_simplex[i] * num_days
	end

	### holiday effect
	holidayeffect ~ Beta(1,1)

    ### 4.2) derive observables from infections
	ifr       ~ truncated(Normal(6/1000, 3/1000), 1/1000, 10/1000)
	ihr       ~ truncated(Normal(1/100,1/100),.1/100,5/100)
    iar       ~ Beta(1,10)

    expected_daily_cases     = TV(undef, num_time_steps)
    expected_daily_hospit    = TV(undef, num_time_steps)
    expected_daily_deaths    = TV(undef, num_time_steps)
	expected_seropositive    = TV(undef, num_time_steps)

	expected_daily_cases[1]  = 1e-15 * newly_infected[1]
    expected_daily_hospit[1] = 1e-15 * newly_infected[1]
    expected_daily_deaths[1] = 1e-15 * newly_infected[1]
	for t = 2:num_time_steps
		periodiceffect = (1 - holidayeffect * holiday[t]) * weekdayeffect[weekday[t]]
		expected_daily_cases[t]  = iar * sum(newly_infected[τ] * i2c[t - τ] for τ = (t - 1):-1:max(t-num_i2c,1)) * periodiceffect
        expected_daily_hospit[t] = ihr * sum(newly_infected[τ] * i2h[t - τ] for τ = (t - 1):-1:max(t-num_i2h,1))
        expected_daily_deaths[t] = ifr * sum(newly_infected[τ] * i2d[t - τ] for τ = (t - 1):-1:max(t-num_i2d,1))
	end

	serodelay = 16 # 14d self deferral + 16d seroconversion (see Erikstrup et al. https://doi.org/10.1093/cid/ciaa849)
	expected_seropositive[1:serodelay]    .= zero(V)
	expected_seropositive[serodelay+1:end] = cumulative_infected[1:end-serodelay]

    ########### 4.) compare model to observations

    ## 4.3) observe deaths
	ϕ_d   ~ truncated(Normal(20, 10), 0, Inf)
    ts_d  = epidemic_start:num_obs
	μs_d  = expected_daily_deaths[ts_d]
	!all( 0 .< μs_d .< population ) && (@error "expected_daily_deaths"; Turing.@addlogprob! -Inf; return)
	deaths[ts_d] ~ arraydist(NegativeBinomial2.(μs_d, ϕ_d))

    ## 4.3) seroprevalence study
    select_seroprev = 2:4
	seroprev_σ ~ arraydist( InverseGamma2.(seroprev_std[ select_seroprev ], Ref(.1)) )
    seroprev_μ = expected_seropositive[ seroprev_idx[ select_seroprev ] ]
	seroprev_mean[select_seroprev] ~ MvNormal(seroprev_μ, seroprev_σ) # select only early observations

    return (
		expected_daily_cases  = expected_daily_cases,
        expected_daily_deaths = expected_daily_deaths,
        expected_daily_hospit = expected_daily_hospit,
        Rt = Rt,
		expected_seropositive = expected_seropositive,
		iar = iar,
    )
end

max_stars_repo_name: andreaskoher/CovidSurvey

max_issues_repo_name: andreaskoher/CovidSurvey

max_forks_repo_name: andreaskoher/CovidSurvey


text:
export →,
    rescale,
    MappedDict,
    mapped_dict,
    forward_map

"""
A `MappedDict` has a dictionary and a map. The domain of the dictionary is
mapped to a different one. Evaluating the `MappedDict` in a point uses the
inverse map to evaluate the underlying dictionary element in the corresponding
point.
"""
abstract type MappedDict{S,T} <: DerivedDict{S,T} end

const MappedDict1d{S <: Number,T <: Number} = MappedDict{S,T}
const MappedDict2d{S <: SVector{2},T <: Number} = MappedDict{S,T}

mapped_dict(dict::Dictionary, map) = mapped_dict1(dict, map)
# Allow dispatch on first argument (dict)
mapped_dict1(dict, map) = mapped_dict2(dict, map)
# Allow dispatch on second argument (map)
mapped_dict2(dict, map) = MappedDict(dict, map)

mapped_dict1(dict::MappedDict, map) = mapped_dict(superdict(dict), map ∘ forward_map(dict))
mapped_dict2(dict, map::IdentityMap) = dict

# Convenience function, similar to apply_map for grids etcetera
@deprecate apply_map(dict::Dictionary, map) mapped_dict(dict, map)

# In the constructor we check the domain and codomain types.
# The domain of the MappedDict is defined by the range of the map, because the
# domain of the underlying dict is mapped to the domain of the MappedDict.
# Hence, the domain type of the map has to equal the domain type of the dictionary.
MappedDict(dict::Dictionary{T1,T}, map::Map{T1}) where {T1,T} =
    MappedDict{codomaintype(map,T1),T}(dict, map)
MappedDict(dict::Dictionary{S1,T1}, map::Map{S2}) where {S1,S2,T1} =
    MappedDict(dict, convert(Map{S1}, map))

MappedDict{S,T}(dict::Dictionary, map) where {S,T} = GenericMappedDict{S,T}(dict, map)

MappedDict{S,T}(dict::Dictionary, map::ScalarAffineMap{S}) where {S,T} =
    ScalarMappedDict{S,T}(dict, map)

"Concrete mapped dictionary for any dictionary and any map."
struct GenericMappedDict{S,T,D,M} <: MappedDict{S,T}
    superdict   ::  D
    map         ::  M
end

GenericMappedDict{S,T}(dict::Dictionary, map) where {S,T} =
    GenericMappedDict{S,T,typeof(dict),typeof(map)}(dict, map)

"Concrete mapped dictionary for any dictionary with a scalar map."
struct ScalarMappedDict{S,T,D} <: MappedDict{S,T}
    superdict   ::  D
    map         ::  ScalarAffineMap{S}
end

ScalarMappedDict{S,T}(dict::Dictionary, map) where {S,T} =
    GenericMappedDict{S,T,typeof(dict),typeof(map)}(dict, map)

forward_map(dict::MappedDict) = dict.map
inverse_map(dict::MappedDict) = inverse(forward_map(dict))

similardictionary(dict::MappedDict, dict2::Dictionary) = mapped_dict(dict2, forward_map(dict))

hasderivative(dict::MappedDict) = false
hasantiderivative(dict::MappedDict) = false

hasderivative(dict::MappedDict1d) =
    hasderivative(superdict(dict)) && isaffine(forward_map(dict))
hasantiderivative(dict::MappedDict1d) =
    hasantiderivative(superdict(dict)) && isaffine(forward_map(dict))

interpolation_grid(dict::MappedDict) = _grid(dict, superdict(dict), forward_map(dict))
_grid(dict::MappedDict, sdict, map) = map_grid(map, interpolation_grid(sdict))

function unmap_grid(dict::MappedDict, grid::MappedGrid)
    if iscompatible(forward_map(dict), forward_map(grid))
        supergrid(grid)
    else
        apply_map(grid, inverse_map(dict))
    end
end

unmap_grid(dict::MappedDict, grid::AbstractGrid) = apply_map(grid, inverse_map(dict))

isreal(s::MappedDict) = isreal(superdict(s)) && isreal(forward_map(s))

unsafe_eval_element(s::MappedDict, idx, y) =
    unsafe_eval_element(superdict(s), idx, leftinverse(forward_map(s), y))

function unsafe_eval_element_derivative(s::MappedDict1d, idx, y, order)
    @assert order == 1
    x = leftinverse(forward_map(s), y)
    d = unsafe_eval_element_derivative(superdict(s), idx, x, order)
    z = d / jacdet(forward_map(s), y)
end

function unsafe_eval_element_derivative(dict::MappedDict, idx, y, order)
    @assert maximum(order) <= 1
    m = forward_map(dict)
    x = leftinverse(m, y)
    J = jacobian(m, y)
    g = eval_gradient(superdict(dict), idx, x)
    sum(order .* (J' \ g))
end

function eval_expansion(s::MappedDict{S,T}, coef, y::S) where {S,T}
    if in_support(s, first(eachindex(s)), y)
        eval_expansion(superdict(s), coef, leftinverse(forward_map(s), y))
    else
        zero(codomaintype(s))
    end
end

support(dict::MappedDict) = DomainSets.map_domain(forward_map(dict), support(superdict(dict)))

support(dict::MappedDict, idx) = forward_map(dict).(support(superdict(dict), idx))

function dict_in_support(set::MappedDict, idx, y, threshold = default_threshold(y))
    x = leftinverse(forward_map(set), y)
    y1 = forward_map(set)(x)
    if norm(y-y1) < threshold
        in_support(superdict(set), idx, x)
    else
        false
    end
end

iscompatible(s1::MappedDict, s2::MappedDict) =
    iscompatible(forward_map(s1),forward_map(s2)) && iscompatible(superdict(s1),superdict(s2))

measure(dict::MappedDict) = apply_map(measure(superdict(dict)), forward_map(dict))

hasmeasure(dict::MappedDict) = hasmeasure(superdict(dict))

for f in (:isorthonormal, :isorthogonal, :isbiorthogonal)
    @eval $f(dict::MappedDict, measure::MappedWeight) =
        iscompatible(forward_map(dict),forward_map(measure)) && $f(superdict(dict), supermeasure(measure))
end

function innerproduct_native(d1::MappedDict, i, d2::MappedDict, j, measure::MappedWeight; options...)
    if iscompatible(d1,d2) && iscompatible(forward_map(d1),forward_map(measure))
        innerproduct(superdict(d1), i, superdict(d2), j, supermeasure(measure); options...)
    else
        innerproduct1(d1, i, d2, j, measure; options...)
    end
end

function gram(::Type{T}, dict::MappedDict, m::MappedWeight; options...) where {T}
    if iscompatible(forward_map(dict), forward_map(m))
        wrap_operator(dict, dict, gram(T, superdict(dict), supermeasure(m); options...))
    else
        default_gram(T, dict, m; options...)
    end
end

function gram(::Type{T}, dict::MappedDict, m::DiscreteWeight, grid::MappedGrid, weights; options...) where {T}
    if iscompatible(forward_map(grid), forward_map(dict))
        wrap_operator(dict, dict, gram(T, superdict(dict), supermeasure(m); options...))
    else
        default_gram(T, dict, m; options...)
    end
end

###############
# Transforms
###############

# We assume that maps do not affect a transform, as long as both the source and
# destination are mapped in the same way. If a source or destination grid is not
# mapped, we attempt to apply the inverse map to the grid and continue.
# For example, a mapped Fourier basis may have a PeriodicEquispacedGrid on a
# general interval. It is not necessarily a mapped grid.

transform_dict(s::MappedDict; options...) = mapped_dict(transform_dict(superdict(s); options...), forward_map(s))

function hasgrid_transform(dict::MappedDict, gb::GridBasis, grid::AbstractGrid)
    sgrid = unmap_grid(dict, grid)
    T = codomaintype(gb)
    hasgrid_transform(superdict(dict), GridBasis{T}(sgrid), sgrid)
end

function transform_from_grid(::Type{T}, s1::GridBasis, s2::MappedDict, grid; options...) where {T}
    sgrid = unmap_grid(s2, grid)
    op = transform_from_grid(T, GridBasis{T}(sgrid), superdict(s2), sgrid; options...)
    wrap_operator(s1, s2, op)
end

function transform_to_grid(::Type{T}, s1::MappedDict, s2::GridBasis, grid; options...) where {T}
    sgrid = unmap_grid(s1, grid)
    op = transform_to_grid(T, superdict(s1), GridBasis{T}(sgrid), sgrid; options...)
    wrap_operator(s1, s2, op)
end

###################
# Evaluation
###################

function evaluation(::Type{T}, dict::MappedDict, gb::GridBasis, grid; options...) where {T}
    sgrid = unmap_grid(dict, grid)
    op = evaluation(T, superdict(dict), GridBasis{T}(sgrid), sgrid; options...)
    wrap_operator(dict, gb, op)
end

###################
# Differentiation
###################

function derivative_dict(Φ::MappedDict, order; options...)
    map = forward_map(Φ)
    superdiff = similardictionary(Φ, derivative_dict(superdict(Φ), order; options...))
    if isaffine(map)
        superdiff
    elseif jacobian(map) isa ConstantMap
        superdiff
    else
        # TODO: this is only correct in 1D
        @assert order == 1
        mi = inverse(map)
        (t -> 1/jacdet(map, mi(t))) * superdiff
    end
end

function antiderivative_dict(Φ::MappedDict, order; options...)
    map = forward_map(Φ)
    if isaffine(map)
        similardictionary(Φ, antiderivative_dict(superdict(Φ), order; options...))
    elseif jacobian(map) isa ConstantMap
        similardictionary(Φ, antiderivative_dict(superdict(Φ), order; options...))
    else
        error("Don't know the antiderivative of a dictionary mapped by $(map)")
    end
end

# TODO: generalize to other orders
function differentiation(::Type{T}, dsrc::MappedDict1d, ddest::MappedDict1d, order::Int; options...) where {T}
    @assert isaffine(forward_map(dsrc))
    D = differentiation(T, superdict(dsrc), superdict(ddest), order; options...)
    S = ScalingOperator{T}(dest(D), jacobian(forward_map(dsrc),1)^(-order))
    wrap_operator(dsrc, ddest, S*D)
end

function differentiation(::Type{T}, dsrc::MappedDict1d, ddest::DerivedDict, order; options...) where {T}
    @assert order == 1
    @assert ddest isa WeightedDict1d
    @assert iscompatible(forward_map(dsrc), forward_map(superdict(ddest)))
    D = differentiation(T, superdict(dsrc), superdict(superdict(ddest)), order; options...)
    wrap_operator(dsrc, ddest, D)
end

function antidifferentiation(::Type{T}, dsrc::MappedDict1d, ddest::MappedDict1d, order::Int; options...) where {T}
    @assert isaffine(forward_map(dsrc))
    D = antidifferentiation(T, superdict(dsrc), superdict(ddest), order; options...)
    S = ScalingOperator{T}(dest(D), jacobian(forward_map(dsrc),1)^(order))
    wrap_operator(dsrc, ddest, S*D)
end

#################
# Special cases
#################

mapped_dict(dict::GridBasis, map) = GridBasis(map_grid(map, grid(s)), coefficienttype(s))

"Rescale a function set to an interval [a,b]."
function rescale(s::Dictionary1d, a::Number, b::Number)
    T = promote_type(domaintype(s), typeof(a), typeof(b))
    if abs(a-infimum(support(s))) < 10eps(T) && abs(b-supremum(support(s))) < 10eps(T)
        s
    else
        m = mapto(support(s), T(a)..T(b))
        mapped_dict(s, m)
    end
end

rescale(s::Dictionary, d::AbstractInterval) = rescale(s, infimum(d), supremum(d))

"Map a dictionary to a domain"
(→)(Φ::Dictionary, domain::Domain) = rescale(Φ, domain)
# The symbol is \to

(∘)(Φ::Dictionary, map::AbstractMap) = mapped_dict(Φ, map)

# "Preserve Tensor Product Structure"
function rescale(s::TensorProductDict, a::SVector{N}, b::SVector{N}) where {N}
    scaled_sets = [ rescale(component(s,i), a[i], b[i]) for i in 1:N]
    tensorproduct(scaled_sets...)
end

plotgrid(S::MappedDict, n) = apply_map(plotgrid(superdict(S),n), forward_map(S))

#################
# Arithmetic
#################

function (*)(dict1::MappedDict, dict2::MappedDict, coef_src1, coef_src2)
    @assert iscompatible(superdict(dict1),superdict(dict2))
    mset,mcoef = (*)(superdict(dict1),superdict(dict2),coef_src1,coef_src2)
    mapped_dict(mset, forward_map(dict1)), mcoef
end

## Printing

Display.object_parentheses(d::MappedDict) = true
Display.stencil_parentheses(d::MappedDict) = true
Display.displaystencil(d::MappedDict) = _stencil(d, superdict(d), forward_map(d))
_stencil(d::MappedDict, dict, map) = [dict, " ∘ ", map]
_stencil(d::MappedDict, dict, map::ScalarAffineMap) = [dict, " → ", map_domain(map, support(dict))]

show(io::IO, d::MappedDict) = _do_show(io, d, superdict(d), forward_map(d))
_do_show(io::IO, d::MappedDict, dict, map::ScalarAffineMap) =
    print(io, repr(dict), " → ", repr(map_domain(map, support(dict))))
_do_show(io::IO, d::MappedDict, dict, map) =
    print(io, "MappedDict(", repr(dict), ", ", repr(map), ")")


text:
"""
    cyclebits!(dgt::AbstractVector{<:Integer})

Right-shift digits (in place).
"""
function cyclebits!(dgt::AbstractVector{<:Integer})
    p = dgt[end]
    for i = length(dgt):-1:2
        dgt[i] = dgt[i-1]
    end
    dgt[1] = p
    dgt
end

"""
    translation_check(dgt::AbstractVector{<:Integer}, I0::Integer, base::Integer)

Check if the digits is of minimum index under cycling.

Inputs:
-------
- `dgt  : Digits.
- `I0`  : Index of the `dgt`.
- `base`: Base.

Outputs:
--------
- `Q`: Whether `dgt` is minimum.
- `R`: Periodicity of the `dgt`, return 0 if !Q.
"""
function translation_check(dgt::AbstractVector{<:Integer}, I0::Integer, base::Integer)
    Q, R = true, length(dgt)
    cyclebits!(dgt)
    for i=1:length(dgt)-1
        In = index(dgt, base=base)
        if In < I0
            Q = false
            R = 0
            change!(dgt, I0, base=base)
            break
        elseif In == I0
            R = i
            break
        end
        cyclebits!(dgt)
    end
    Q, R
end

"""
    translation_check_min!(dgt::AbstractVector{<:Integer}, Im::Integer, base::Integer)

Given an index `Im`, cycle the digits to see:

1. Whether `Im` is the smallest index.
2. Whether the shifted digits have the same index as `Im`.
3. If the shifted digits have the same index as `Im`, find the cycling number.

Check if the digits is of minimum index under cycling.
This function may change vector `dgt`.

Inputs:
-------
- `dgt  : Digits.
- `Im`  : Minimum index.
- `base`: Base.
- `R`   : (Optional) Periodicity of `dgt`, `R = length(dgt)` by default.

Outputs:
--------
- `Qm`: Whether `Im` is minimum index.
- `Qs`: Whether `dgt` can be shifted to have index `Im`.
- `M` : Cycling number by which `dgt` is shifted to have index `Im`.
"""
function translation_check_min!(dgt::AbstractVector{<:Integer}, Im::Integer, base::Integer, R::Integer=length(dgt))
    Qm, Qs, M = true, false, 0
    for i=0:R-1
        In = index(dgt, base=base)
        if In < Im
            Qm, Qs = false, false
            break
        elseif In == Im && !Qs
            Qs = true
            M = i
            break
        end
        cyclebits!(dgt)
    end
    Qm, Qs, M
end

"""
translation_index(dgt::AbstractVector{<:Integer}, base::Integer)

Given a digits, return the minimum index under shifting.

Inputs:
-------
- `dgt  : Digits.
- `base`: Base.

Outputs:
--------
- `Im`: Minimum index.
- `M` : Cycling number by which `dgt` is shifted to have index `Im`.
"""
function translation_index(dgt::AbstractVector{<:Integer}, base::Integer)
    I0 = index(dgt, base=base)
    Im, M = I0, 0
    cyclebits!(dgt)
    for i=1:length(dgt)-1
        In = index(dgt, base=base)
        if In == I0
            break
        elseif In < Im
            Im, M = In, i
        end
        cyclebits!(dgt)
    end
    Im, M
end

"""
    spinflip(v::AbstractVector{<:Integer}, base::Integer)

Flip spins Sz on each site.
"""
function spinflip(v::AbstractVector{<:Integer}, base::Integer)
    vf = Vector{eltype(v)}(undef, length(v))
    for i = 1:length(vf)
        vf[i] = base - v[i] - 1
    end
    vf
end

"""
    translation_parity_check(parity, dgt::AbstractVector{<:Integer}, I0::Integer, base::Integer)

Check whether a state is a valid representing state, and whether it is reflect-translation-invariant.

Inputs:
-------
- `parity`: Parity function, can be spatio-reflection or spin-reflection.
- `dgt`   : Digits.
- `I0`    : Index of the `dgt`.
- `base`  : Base.

Outputs:
--------
- `Qm`: Whether `dgt` is a representing vector.
- `Qp`: Whether `dgt` is reflect-translation-invariant.
- `R` : Minimum R that Tᴿ⋅|dgt⟩ = |dgt⟩.
- `T` : Minimum M that Tᴹ⋅P⋅|dgt⟩ = |dgt⟩, 0 if it is not reflect-translation-invariant.
"""
function translation_parity_check(parity, dgt::AbstractVector{<:Integer}, I0::Integer, base::Integer)
    Q, R = translation_check(dgt, I0, base)
    Qm, Qp, M = if Q
        translation_check_min!(parity(dgt), I0, base, R)
    else
        false, false, 0
    end
    Qm, Qp, R, M
end

"""
    translation_parity_index(parity, dgt::AbstractVector{<:Integer}, base::Integer)

Return the index of a digits, and also tell whether the minimum index is in the other parity sector, and return the number of shifts.

Inputs:
-------
- `parity`: Parity function, can be spatio-reflection or spin-reflection.
- `dgt`   : Digits.
- `base`  : Base.

Outputs:
--------
- `Qp`: Whether the minimum index is in the other parity sector.
- `Im`: Minimum index.
- `M` : Minimum M that Tᴹ⋅|dgt⟩ = |Im⟩ or Tᴹ⋅P⋅|dgt⟩ = |Im⟩.
"""
function translation_parity_index(parity, dgt::AbstractVector{<:Integer}, base::Integer)
    Im1, T1 = translation_index(dgt, base)
    Im2, T2 = translation_index(parity(dgt), base)
    Qp, Im, M = if Im2 < Im1
        true, Im2, T2
    else
        false, Im1, T1
    end
    Qp, Im, M
end


text:
Leaf = Void
mutable struct Node
    name::Symbol
    distribution # TODO find a way to type this
    value::Any
    loglikelihood::Float64
    israndom::Bool
    next::Union{Node,Leaf}
end
Trace = Union{Leaf, Node}

function copyupto(trace::Trace, upto::Node)
    if trace isa Leaf
        trace,trace
    elseif trace == upto
        newupto = Node(upto.name, upto.distribution, upto.value, upto.loglikelihood, upto.israndom, nothing)
        newupto,newupto
    else
        head = Node(trace.name, trace.distribution, trace.value, trace.loglikelihood, trace.israndom, nothing)
        rest,upto = copyupto(trace.next,upto)
        head.next = rest
        head,upto
    end
end

function extend!(previous::Trace, name::Symbol, distribution, value, israndom::Bool)
    node = Node(name, distribution, value, logpdf(distribution, value), israndom, nothing)
    if !(previous isa Leaf)
        previous.next = node
    end
    node
end

function getrandnode(n::Int, node::Trace)
    (n<=1 || node isa Leaf) ? node : getrandnode(n-(node.israndom ? 1 : 0), node.next)
end

function countrandom(node::Trace)
    node isa Leaf ? 0 : (node.israndom ? 1 : 0) + countrandom(node.next)
end

function Distributions.loglikelihood(trace::Trace)
    loglikelihood = 1.0
    current = trace
    while !(current isa Leaf)
        loglikelihood += current.loglikelihood
        current = current.next
    end
    loglikelihood
end

function propose(trace::Trace, perturb)
    nrandoms = countrandom(trace)
    nodeid = rand(1:nrandoms)
    node = getrandnode(nodeid, trace)

    proposal,node = copyupto(trace,node)

    node.value = perturb(node)
    node.loglikelihood = logpdf(node.distribution, node.value)
    node.next = nothing

    return proposal
end


text:
@testset "timestep.jl" begin
    println("             > timestep.jl")

    timestep = FixedTimestep(0.3)
    @test get_step(timestep) == 0.3

    timestep = RandomTimestep(Exponential(0.5))
    @test get_step(timestep) > 0
    @test get_step(timestep) != get_step(timestep)
end


text:
Base.@kwdef struct ConstantSmagorinsky{T} <: IsotropicDiffusivity{T}
               C :: T = 0.23
              Pr :: T = 1.0
    ν_background :: T = 0.0
    κ_background :: T = 0.0
end

"""
    ConstantSmagorinsky(T=Float64; C=0.23, Pr=1.0, ν_background=1e-6,
                            κ_background=1e-7)

Returns a `ConstantSmagorinsky` closure object of type `T` with

* `C`            : Smagorinsky constant
* `Pr`           : Prandtl number
* `ν_background` : background viscosity
* `κ_background` : background diffusivity
"""
ConstantSmagorinsky(T; kwargs...) =
      typed_keyword_constructor(T, ConstantSmagorinsky; kwargs...)

"Return the filter width for Constant Smagorinsky on a Regular Cartesian grid."
Δ(i, j, k, grid::RegularCartesianGrid, ::ConstantSmagorinsky) = geo_mean_Δ(grid)

# tr_Σ² : ccc
#   Σ₁₂ : ffc
#   Σ₁₃ : fcf
#   Σ₂₃ : cff

"Return the double dot product of strain at `ccc`."
function ΣᵢⱼΣᵢⱼ_ccc(i, j, k, grid, u, v, w)
    return (
                    tr_Σ²(i, j, k, grid, u, v, w)
            + 2 * ▶xy_cca(i, j, k, grid, Σ₁₂², u, v, w)
            + 2 * ▶xz_cac(i, j, k, grid, Σ₁₃², u, v, w)
            + 2 * ▶yz_acc(i, j, k, grid, Σ₂₃², u, v, w)
            )
end

"Return the double dot product of strain at `ffc`."
function ΣᵢⱼΣᵢⱼ_ffc(i, j, k, grid, u, v, w)
    return (
                  ▶xy_ffa(i, j, k, grid, tr_Σ², u, v, w)
            + 2 *    Σ₁₂²(i, j, k, grid, u, v, w)
            + 2 * ▶yz_afc(i, j, k, grid, Σ₁₃², u, v, w)
            + 2 * ▶xz_fac(i, j, k, grid, Σ₂₃², u, v, w)
            )
end

"Return the double dot product of strain at `fcf`."
function ΣᵢⱼΣᵢⱼ_fcf(i, j, k, grid, u, v, w)
    return (
                  ▶xz_faf(i, j, k, grid, tr_Σ², u, v, w)
            + 2 * ▶yz_acf(i, j, k, grid, Σ₁₂², u, v, w)
            + 2 *    Σ₁₃²(i, j, k, grid, u, v, w)
            + 2 * ▶xy_fca(i, j, k, grid, Σ₂₃², u, v, w)
            )
end

"Return the double dot product of strain at `cff`."
function ΣᵢⱼΣᵢⱼ_cff(i, j, k, grid, u, v, w)
    return (
                  ▶yz_aff(i, j, k, grid, tr_Σ², u, v, w)
            + 2 * ▶xz_caf(i, j, k, grid, Σ₁₂², u, v, w)
            + 2 * ▶xy_cfa(i, j, k, grid, Σ₁₃², u, v, w)
            + 2 *    Σ₂₃²(i, j, k, grid, u, v, w)
            )
end

# Temporarily set filter widths to cell-size (rather than distance between cell centers, etc.)
Δ_ccc = Δ
Δ_ffc = Δ
Δ_fcf = Δ
Δ_cff = Δ

ν_ccc(i, j, k, grid, clo::ConstantSmagorinsky, eos, g, u, v, w, T, S) =
    (clo.C * Δ_ccc(i, j, k, grid, clo))^2 * sqrt(2 * ΣᵢⱼΣᵢⱼ_ccc(i, j, k, grid, u, v, w)) + clo.ν_background

ν_ffc(i, j, k, grid, clo::ConstantSmagorinsky, eos, g, u, v, w, T, S) =
    (clo.C * Δ_ffc(i, j, k, grid, clo))^2 * sqrt(2 * ΣᵢⱼΣᵢⱼ_ffc(i, j, k, grid, u, v, w)) + clo.ν_background

ν_fcf(i, j, k, grid, clo::ConstantSmagorinsky, eos, g, u, v, w, T, S) =
    (clo.C * Δ_fcf(i, j, k, grid, clo))^2 * sqrt(2 * ΣᵢⱼΣᵢⱼ_fcf(i, j, k, grid, u, v, w)) + clo.ν_background

ν_cff(i, j, k, grid, clo::ConstantSmagorinsky, eos, g, u, v, w, T, S) =
    (clo.C * Δ_cff(i, j, k, grid, clo))^2 * sqrt(2 * ΣᵢⱼΣᵢⱼ_cff(i, j, k, grid, u, v, w)) + clo.ν_background

κ_ccc(i, j, k, grid, clo::ConstantSmagorinsky, eos, g, u, v, w, T, S) =
    (clo.C * Δ_ccc(i, j, k, grid, clo))^2 * sqrt(2 * ΣᵢⱼΣᵢⱼ_ccc(i, j, k, grid, u, v, w)) / clo.Pr + clo.κ_background


text:
using Crux
using Test

## elapsed
@test Crux.elapsed(1000, 100)
@test !Crux.elapsed(1001, 100)
@test Crux.elapsed(47001:47501, 500)
@test !Crux.elapsed(47001:47499, 500)


text:
module ProbLineSearch

using LinearAlgebra
using Distributions
using ArgCheck, Parameters, Formatting

# Original code (MATLAB) license:
#
# Copyright (c) 2015 (post NIPS 2015 release 4.0), Maren Mahsereci, Philipp Hennig
# mmahsereci@tue.mpg.de, phennig@tue.mpg.de
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""
Example usage:
```julia
x, gp, search = PLS.probLineSearch(plsfun, x₀, 0.01)
xs = [PLS.evaluate(plsfun, x₀), x]; gps = [[],gp];
for i in 1:20
    search.search_direction = -x.J
    PLS.project_1d!(x, search.search_direction)
    x, gp = PLS.probLineSearch(plsfun, x, search)
    push!(xs, copy(x))
    push!(gps, gp)
end
```

"""

"""
Statistics of accepted step sizes and number of func evals
example usage: PLSHistory(Float32, 10)
"""
# TODO: mv PLSHistory => PLSStatistics, make
# TODO: cat PLSHistory = {all_evals, all_EI, all_wolfe, all_messages}
@with_kw mutable struct PLSHistory{T <: AbstractFloat}
    evals::Vector{Int} = []       # number of functions evals
    recording::Bool = false            # number of outer iterations so far.
    αs::Vector{T} = []                 # accepted αs
    ewma_α::T = 1.0                    # (exp. wgt) moving average of α (across outer loops)
    EI::Vector{T} = []                 # Expected Improvement at acceptance
    Wolfe::Vector{T} = []              # Wolfe probability at acceptance
    msg::Vector{String} = []   # Description of acceptance reason
end
PLSHistory(α₀::T) where T <: AbstractFloat = PLSHistory{T}(ewma_α=α₀)
PLSHistory(T::Type; kwargs...) = PLSHistory{T}(;kwargs...)
PLSHistory() = PLSHistory{Float64}()

Base.length(x::PLSHistory) = length(x.evals)
increment_eval!(x::PLSHistory) = x.recording ? x.evals[end] += 1 : (push!(x.evals, 1); x.recording=true);

"""
PLS constants (see Mahsereci & Hennig, 2017), which shouldn't need to change,
but are available here just in case.
"""
@with_kw struct PLSConstantParams{T<:AbstractFloat}
    verbosity::Int = 1
    c1::T = 0.05  # > pct of ∇₀ decrease required
    c2::T = 0.5   # grad at accepted point must be above this pct of ∇₀
    niter::Int = 7
    wolfeThresh::T = 0.3   # ≥ this prob reqd, for Wolfe conditions to hold
    αGrowth::T = 1.3       # amount to grow α₀ at each outer iteration.
end
# PLSConstantParams(T::Type) = PLSConstantParams{T}()
PLSConstantParams(T::Type; kwargs...) = PLSConstantParams{T}(;kwargs...)

"""
Definition of the line search: the starting point, the direction etc.
"""
mutable struct PLSSearchDefn{T<:AbstractFloat}
    x₀::Vector{T}
    f₀::T
    α₀::T
    search_direction::Vector{T}
    σ²_f::T                # assumed var of f
    σ²_J::Vector{T}        # assumed var of Jf (Jacobian)
    extrap_amt::T          # extrapolation amount
    denom::T               # magnitude |y′(0)| => Divisor for normalisation
end
PLSSearchDefn(T::Type, args...) = PLSSearchDefn{T}(args..., T(1), T(1))
var∇1d(x::PLSSearchDefn) = (x.search_direction.^2)' * x.σ²_J

Base.copy(x::PLSSearchDefn) = PLSSearchDefn([deepcopy(getfield(x, nm)) for nm in fieldnames(PLSSearchDefn)]...)

include("./fun.jl")
include("./posteriorgp.jl")
include("./utils.jl")
# include("./plot.jl")
include("./main.jl")

export PLSHistory, PLSConstantParams, PLSSearchDefn, PLSPosterior, PLSEvaluation
export PLSBespokeFunction
export probLineSearch

end # module

max_stars_repo_name: ornithos/ProbLineSearch