text:
using Nextion
using Nextion.Draw
using Test

const LCD_WIDTH, LCD_HEIGHT = 480, 320
const DEFAULT_COLOUR = Colour.WHITE
const IMG_WIDTH, IMG_HEIGHT = 200, 135

@testset "draw" begin
    nexSerial = NexSerial("/dev/ttyUSB0")

    send(nexSerial, "page0")
    cls(nexSerial, DEFAULT_COLOUR)

    function print_several_lines(nexSerial; lcd_width=LCD_WIDTH, color1="WHITE", color2="BLACK")
        function swap(x, y)
            y, x
        end

        x, y = 0, 10
        w, h = lcd_width, 30
        font = Font(0)
        fontcolor = color1
        backcolor = color2
        xcenter = Alignment.Horizontal.CENTRE  # NONE/LEFT/CENTRE/RIGHT
        ycenter = Alignment.Vertical.CENTRE  # NONE/UP/CENTRE/DOWN
        sta = Background.SOLIDCOLOUR # NONE/CROPIMAGE/SOLIDCOLOUR/IMAGE/NOBACKCOLOUR

        lines = [
            "Hello Julia!",
            "",
            "This is Nextion",
            "",
            "I'm an intelligent display",
            "",
            "",
            "",
            "https://github.com/scls19fr",
            "/Nextion.jl"
        ]
        for line in lines
            xstr(nexSerial, x, y, w, h, font, fontcolor, backcolor, xcenter, ycenter, sta, line)
            y = y + h
            fontcolor, backcolor = swap(fontcolor, backcolor)
        end
    end

    function diagonal_line(nexSerial; lcd_width=LCD_WIDTH, lcd_height=LCD_HEIGHT)
        x1, y1 = 0, 0
        x2, y2 = lcd_width, lcd_height
        colour = Colour.BLUE
        line(nexSerial, x1, y1, x2, y2, colour)
    end

    function circle_frame_fill(nexSerial; lcd_width=LCD_WIDTH, lcd_height=LCD_HEIGHT, delay=1)
        x = Int(round(lcd_width / 2))
        y = Int(round(lcd_height / 2))
        r = 100
        colour = Colour.BLUE
        circle(nexSerial, x, y, r, colour)
        sleep(delay)
        circle(nexSerial, x, y, r, colour, Background.SOLIDCOLOUR)
    end

    function clear_screen_color_france(nexSerial, delay=1)
        cls(nexSerial, Colour.BLUE)

        sleep(delay)

        cls(nexSerial, Colour.WHITE)

        sleep(delay)

        cls(nexSerial, Colour.RED)
    end

    function france_flag(nexSerial; lcd_width=LCD_WIDTH, lcd_height=LCD_HEIGHT)
        w, h = Int(round(lcd_width/3)), lcd_height
        x1, y1 = 0, 0
        y2 = h
        x2 = x1 + w
        colour = Colour.BLUE
        rectangle(nexSerial, x1, y1, x2, y2, colour, Background.SOLIDCOLOUR)
        x1 = x1 + 2*w
        x2 = x1 + w
        colour = Colour.RED
        rectangle(nexSerial, x1, y1, x2, y2, colour, Background.SOLIDCOLOUR)
    end

    function rectangle_frame_fill(nexSerial; lcd_width=LCD_WIDTH, lcd_height=LCD_HEIGHT)
        w, h = 120, 80
        x1 = Int(round((lcd_width - w) / 2))
        y1 = Int(round((lcd_height - h) / 2))
        x2 = x1 + w
        y2 = y1 + h
        colour = Colour.RED
        rectangle(nexSerial, x1, y1, x2, y2, colour)
        sleep(1)
        rectangle(nexSerial, x1, y1, x2, y2, colour, Background.SOLIDCOLOUR)
    end

    function picture_julia(nexSerial; lcd_width=LCD_WIDTH, lcd_height=LCD_HEIGHT, img_width=IMG_WIDTH, img_height=IMG_HEIGHT, delay=0.5)
        x = Int(round((lcd_width - img_width) / 2))
        y = Int(round((lcd_height - img_height) / 2))
        for picid in 0:4
            picture(nexSerial, x, y, Picture(picid))
            sleep(delay)
        end
    end

    print_several_lines(nexSerial)

    sleep(2)
    cls(nexSerial, DEFAULT_COLOUR)

    diagonal_line(nexSerial)

    sleep(2)
    cls(nexSerial)

    circle_frame_fill(nexSerial)

    sleep(2)
    cls(nexSerial)

    clear_screen_color_france(nexSerial)

    sleep(2)
    cls(nexSerial)

    france_flag(nexSerial)

    sleep(2)
    cls(nexSerial)
    sleep(1)

    rectangle_frame_fill(nexSerial)

    sleep(2)
    cls(nexSerial)

    picture_julia(nexSerial)

    close(nexSerial)
    @test true

end


text:
using ASTInterpreter2
using Base.Meta
using REPL
using DebuggerFramework: execute_command, dummy_state

struct DummyState; end
REPL.LineEdit.transition(s::DummyState, _) = nothing

@assert step_through(ASTInterpreter2.enter_call_expr(:($(+)(1,2.5)))) == 3.5
@assert step_through(ASTInterpreter2.enter_call_expr(:($(sin)(1)))) == sin(1)
@assert step_through(ASTInterpreter2.enter_call_expr(:($(gcd)(10,20)))) == gcd(10, 20)

# Step into generated functions
@generated function generatedfoo(T)
    :(return $T)
end
callgenerated() = generatedfoo(1)
frame = ASTInterpreter2.enter_call_expr(:($(callgenerated)()))
state = dummy_state([frame])

# Step into the generated function itself
execute_command(state, state.stack[1], Val{:sg}(), "sg")

# Should now be in generated function
execute_command(state, state.stack[1], Val{:finish}(), "finish")

# Now finish the regular function
execute_command(state, state.stack[1], Val{:finish}(), "finish")

@assert isempty(state.stack)

# Optional arguments
function optional(n = sin(1))
    x = asin(n)
    cos(x)
end

frame = ASTInterpreter2.enter_call_expr(:($(optional)()))
state = dummy_state([frame])
# First call steps in
execute_command(state, state.stack[1], Val{:n}(), "n")
# cos(1.0)
execute_command(state, state.stack[1], Val{:n}(), "n")
# return
execute_command(state, state.stack[1], Val{:n}(), "n")

@assert isempty(state.stack)

# Macros
macro insert_some_calls()
    esc(quote
        x = sin(b)
        y = asin(x)
        z = sin(y)
    end)
end

# Work around the fact that we can't detect macro expansions if the macro
# is defined in the same file
include_string(Main, """
function test_macro()
    a = sin(5)
    b = asin(a)
    @insert_some_calls
    z
end
""","file.jl")

frame = ASTInterpreter2.enter_call_expr(:($(test_macro)()))
state = dummy_state([frame])
# a = sin(5)
execute_command(state, state.stack[1], Val{:n}(), "n")
# b = asin(5)
execute_command(state, state.stack[1], Val{:n}(), "n")
# @insert_some_calls
execute_command(state, state.stack[1], Val{:n}(), "n")
# TODO: Is this right?
execute_command(state, state.stack[1], Val{:n}(), "n")
# return z
execute_command(state, state.stack[1], Val{:n}(), "n")
execute_command(state, state.stack[1], Val{:n}(), "n")
@test isempty(state.stack)

# Test stepping into functions with keyword arguments
f(x; b = 1) = x+b
g() = f(1; b = 2)
frame = ASTInterpreter2.enter_call_expr(:($(g)()));
state = dummy_state([frame])
# Step to the actual call
execute_command(state, state.stack[1], Val{:nc}(), "nc")
execute_command(state, state.stack[1], Val{:nc}(), "nc")
execute_command(state, state.stack[1], Val{:nc}(), "nc")
# Step in
execute_command(state, state.stack[1], Val{:s}(), "s")
# Should get out in two steps
execute_command(state, state.stack[1], Val{:finish}(), "finish")
execute_command(state, state.stack[1], Val{:finish}(), "finish")
@assert isempty(state.stack)

# Test stepping into functions with exception frames
function f_exc()
    try
    catch err
    end
end

function g_exc()
    try
        error()
    catch err
        return err
    end
end

stack = @make_stack f_exc()
state = dummy_state(stack)

execute_command(state, state.stack[1], Val{:n}(), "n")
@assert isempty(state.stack)

stack = @make_stack g_exc()
state = dummy_state(stack)

execute_command(state, state.stack[1], Val{:n}(), "n")
execute_command(state, state.stack[1], Val{:n}(), "n")
@assert isempty(state.stack)
@assert state.overall_result isa ErrorException

# Test throwing exception across frames
function f_exc_inner()
    error()
end

function f_exc_outer()
    try
        f_exc_inner()
    catch err
        return err
    end
end

stack = @make_stack f_exc_outer()
state = dummy_state(stack)

execute_command(state, state.stack[1], Val{:s}(), "s")
execute_command(state, state.stack[1], Val{:n}(), "n")
execute_command(state, state.stack[1], Val{:n}(), "n")
@assert isempty(state.stack)
@assert state.overall_result isa ErrorException

# Test that symbols don't get an extra QuoteNode
f_symbol() = :limit => true

stack = @make_stack f_symbol()
state = dummy_state(stack)

execute_command(state, state.stack[1], Val{:s}(), "s")
execute_command(state, state.stack[1], Val{:finish}(), "finish")
execute_command(state, state.stack[1], Val{:finish}(), "finish")
@assert isempty(state.stack)
@assert state.overall_result == f_symbol()

# Test that we can step through varargs
f_va_inner(x) = x + 1
f_va_outer(args...) = f_va_inner(args...)

stack = @make_stack f_va_outer(1)
state = dummy_state(stack)

execute_command(state, state.stack[1], Val{:s}(), "s")
execute_command(state, state.stack[1], Val{:n}(), "n")
@assert !isempty(state.stack)
execute_command(state, state.stack[1], Val{:finish}(), "finish")
execute_command(state, state.stack[1], Val{:finish}(), "finish")
@assert isempty(state.stack)
@assert state.overall_result == 2

# Test that we step through kw wrappers
f(foo; bar=3) = foo+bar
stack = @make_stack f(2, bar=4)
@assert length(stack) > 1
state = dummy_state(stack)
execute_command(state, state.stack[1], Val{:n}(), "nc")
execute_command(state, state.stack[1], Val{:n}(), "nc")
@assert isempty(state.stack)
@assert state.overall_result == 6

# Test that we throw the right error when stepping through error functions
function foo_error(a,b)
    a > b && error()
    return a*b
end
stack = @make_stack foo_error(3,1)
state = dummy_state(stack)
try
    execute_command(state, state.stack[1], Val{:n}(), "n")
catch e
    @assert isa(e, ErrorException)
end

# Issue #17
struct B{T} end
function (::B)(y)
    x = 42*y
    return x + y
end

B_inst = B{Int}()
step_through(ASTInterpreter2.enter_call_expr(:($(B_inst)(10))))

ext: jl


text:
export upsample_nearest, ∇upsample_nearest,
    upsample_linear, ∇upsample_linear,
    upsample_bilinear, ∇upsample_bilinear,
    upsample_trilinear, ∇upsample_trilinear,
    pixel_shuffle

"""
    upsample_nearest(x, scale::NTuple{S,Int})
    upsample_nearest(x; size::NTuple{S,Int})

Upsamples the array `x` by integer multiples along the first `S` dimensions.
Subsequent dimensions of `x` are not altered.

Either the `scale` factors or the final output `size` can be specified.

See also [`upsample_bilinear`](@ref), for two dimensions of an `N=4` array.

# Example
```jldoctest
julia> upsample_nearest([1 2 3; 4 5 6], (2, 3))
4×9 Matrix{$Int}:
 1  1  1  2  2  2  3  3  3
 1  1  1  2  2  2  3  3  3
 4  4  4  5  5  5  6  6  6
 4  4  4  5  5  5  6  6  6

julia> ans == upsample_nearest([1 2 3; 4 5 6]; size=(4, 9))  # equivalent
true

julia> upsample_nearest([1 2 3; 4 5 6], (2,))
4×3 Matrix{$Int}:
 1  2  3
 1  2  3
 4  5  6
 4  5  6

julia> ans == upsample_nearest([1 2 3; 4 5 6], size=(4,))
true
```
"""
function upsample_nearest(x::AbstractArray; size::NTuple{S,Int}) where S
    xsize = Base.size(x)[1:S]
    all(size .% xsize .== 0) || throw(ArgumentError("expected output size divisible by input size"))
    return upsample_nearest(x, size .÷ xsize)
end

function upsample_nearest(x::AbstractArray{T,N}, scales::NTuple{S, <:Integer}) where {T,N,S}
    S in 1:N || throw(ArgumentError("can't upsample ndims(x)=$N with scale=$scales"))
    outsize = ntuple(d -> d<=S ? scales[d] * size(x,d) : size(x,d), N)
    out = similar(x, T, outsize)
    writesize = ntuple(N+S) do d
        d > 2S && return size(x, d-S)
        isodd(d) ? scales[cld(d,2)] : size(x, cld(d,2))
    end
    readsize = ntuple(N+S) do d
        d > 2S && return size(x, d-S)
        isodd(d) ? 1 : size(x, cld(d,2))
    end
    reshape(out, writesize) .= reshape(x, readsize)
    out
end

function ∇upsample_nearest(x::AbstractArray{T,N}, scales::NTuple{S, <:Integer}) where {T,N,S}
    outsize = ntuple(N) do d
        d > S && return size(x,d)
        rem(size(x,d), scales[d]) == 0 || throw(ArgumentError("expected input array evenly divisible by scale=$scales, got size(x)=$(size(x))"))
        div(size(x,d), scales[d])
    end
    tempsize = ntuple(N+S) do d
        d > 2S && return size(x, d-S)
        s = scales[cld(d,2)]
        isodd(d) ? s : div(size(x, cld(d,2)),s)
    end
    mid = sum(reshape(x, tempsize), dims=ntuple(d -> 2d-1, S))
    reshape(mid, outsize)
end

function rrule(::typeof(upsample_nearest), x::AbstractArray, s::Tuple)
    Ω = upsample_nearest(x, s)
    upsample_nearest_pullback(Δ) = (NoTangent(), ∇upsample_nearest(unthunk(Δ), s), NoTangent())
    return Ω, upsample_nearest_pullback
end

# utility function
@inline function compute_source_index_and_lambda(
    ratio, # 0 < ratio < 1
    output_index,
    input_size,
    output_size
)
    real_input_index = ratio*output_index
    input_index0 = floor(Int, real_input_index) # typecast to int was here in C++
    offset = (input_index0 < input_size - 1) ? 1 : 0
    input_index1 = input_index0 + offset
    lambda1 = real_input_index - input_index0
    lambda0 = 1 - lambda1
    return input_index0, input_index1, lambda0, lambda1
end

###########
# linear
###########
"""
    upsample_linear(x::AbstractArray{T,3}, scale::Real)
    upsample_linear(x::AbstractArray{T,3}; size::Integer)

Upsamples the first dimension of the array `x` by the upsample provided `scale`,
using linear interpolation. As an alternative to using `scale`, the resulting array `size`
can be directly specified with a keyword argument.

The size of the output is equal to
`(scale*S1, S2, S3)`, where `S1, S2, S3 = size(x)`.
"""
function upsample_linear(x::AbstractArray{<:Any,3}, scale::Real)
    outsize = floor(Int, scale * Base.size(x)[1])
    return upsample_linear(x; size=outsize)
end

function upsample_linear(x::AbstractArray{T,3}; size::Integer) where T
    w,c,n = Base.size(x)
    if w == size
        return x
    end
    y = similar(x, T, size, c, n)
    return upsample_linear_wcn!(y, x)
end

function upsample_linear(x::AbstractArray{T,3}; size::Integer) where T<:Integer
    y = float.(x)
    res = upsample_linear(y; size=size)
    return round.(T, res)
end

function upsample_linear_wcn!(output::AbstractArray{T,3}, input::AbstractArray{T,3}) where T
    size(input)[2:3] == size(output)[2:3] || error("Number of input and output channels and batches must match. Got input $(size(input)) and output $(size(output))")
    in_w, channels, batches = size(input)
    # treat batch and channel dimension as one for better parallelization granularity
    channels *= batches
    out_w, _, _ = size(output)
    output_slice_size = out_w

    # T() and // so that we can handle rationals (super slow)
    width_scale  = T((in_w - 1) // (out_w - 1))

    @inline idx(c, w) = c * in_w + w + 1

    @inbounds Threads.@threads for c in 0:channels-1
        for ow in 0:out_w-1
            iw0, iw1, w0lambda, w1lambda = compute_source_index_and_lambda(width_scale, ow, in_w, out_w)
            output_offset = c * output_slice_size + ow + 1
            output[output_offset] = (w0lambda * input[idx(c, iw0)] + # w0 * i00
                                     w1lambda * input[idx(c, iw1)])  # w1 * i01
        end
    end
    return output
end

"""
    ∇upsample_linear(Δ::AbstractArray{T,3}; size::Integer) where T

# Arguments
- `Δ`: Incoming gradient array, backpropagated from downstream layers
- `size`: Size of the image upsampled in the first place

# Outputs
- `dx`: Downsampled version of `Δ`
"""
function ∇upsample_linear(Δ::AbstractArray{T,3}; size::Integer) where T
    w, c, n = Base.size(Δ)
    out_w = size
    if w == out_w
        return Δ
    end
    dx = zero(similar(Δ, T, out_w, c, n))
    return ∇upsample_linear_wcn!(dx, Δ)
end

function ∇upsample_linear_wcn!(dx::AbstractArray{T,3}, Δ::AbstractArray{T,3}) where T
    size(dx)[2:3] == size(Δ)[2:3] || error("Number of input and output channels and batches must match. Got input $(size(input)) and output $(size(output))")
    in_w, channels, batches = size(dx)

    # treat batch and channel dimension as one for better parallelization granularity
    channels *= batches
    out_w, _, _ = size(Δ)
    output_slice_size = out_w

    width_scale  = T((in_w - 1) // (out_w - 1))

    @inline idx(c, w) = c * in_w + w + 1

    @inbounds Threads.@threads for c in 0:channels-1
        for ow in 0:out_w-1
            iw0, iw1, w0lambda, w1lambda = compute_source_index_and_lambda(width_scale, ow, in_w, out_w)
            output_offset = c * output_slice_size + ow + 1
            Δ_value = Δ[output_offset]
            dx[idx(c, iw0)] += w0lambda * Δ_value # i00
            dx[idx(c, iw1)] += w1lambda * Δ_value # i01
        end
    end
    return dx
end

function rrule(::typeof(upsample_linear), x; size)
    Ω = upsample_linear(x; size=size)
    function upsample_linear_pullback(Δ)
        (NoTangent(), ∇upsample_linear(unthunk(Δ); size=Base.size(x,1)))
    end
    return Ω, upsample_linear_pullback
end

###########
# bilinear
###########
"""
    upsample_bilinear(x::AbstractArray{T,4}, scale::NTuple{2,Real})
    upsample_bilinear(x::AbstractArray{T,4}; size::NTuple{2,Integer})

Upsamples the first 2 dimensions of the array `x` by the upsample factors stored in `scale`,
using bilinear interpolation. As an alternative to using `scale`, the resulting image `size`
can be directly specified with a keyword argument.

The size of the output is equal to
`(scale[1]*S1, scale[2]*S2, S3, S4)`, where `S1, S2, S3, S4 = size(x)`.

# Examples

```jldoctest
julia> x = reshape(Float32[1 2 3; 4 5 6], (2,3,1,1))
2×3×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
 1.0  2.0  3.0
 4.0  5.0  6.0

julia> upsample_bilinear(x, (2, 3))
4×9×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
 1.0  1.25  1.5  1.75  2.0  2.25  2.5  2.75  3.0
 2.0  2.25  2.5  2.75  3.0  3.25  3.5  3.75  4.0
 3.0  3.25  3.5  3.75  4.0  4.25  4.5  4.75  5.0
 4.0  4.25  4.5  4.75  5.0  5.25  5.5  5.75  6.0

julia> ans == upsample_bilinear(x; size=(4, 9))  # specify ouput size instead
true

julia> upsample_bilinear(x, (2.5, 3.5))  # non-integer scaling factors are allowed
5×10×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
 1.0   1.22222  1.44444  1.66667  1.88889  …  2.33333  2.55556  2.77778  3.0
 1.75  1.97222  2.19444  2.41667  2.63889     3.08333  3.30556  3.52778  3.75
 2.5   2.72222  2.94444  3.16667  3.38889     3.83333  4.05556  4.27778  4.5
 3.25  3.47222  3.69444  3.91667  4.13889     4.58333  4.80556  5.02778  5.25
 4.0   4.22222  4.44444  4.66667  4.88889     5.33333  5.55556  5.77778  6.0
```
"""
function upsample_bilinear(x::AbstractArray{<:Any,4}, scale::NTuple{2,Real})
    outsize = ntuple(i -> floor(Int, scale[i] * Base.size(x, i)), 2)
    return upsample_bilinear(x; size=outsize)
end

upsample_bilinear(x, scale::Real) = upsample_bilinear(x, (scale,scale))

function upsample_bilinear(x::AbstractArray{T,4}; size::NTuple{2,Integer}) where T
    w,h,c,n = Base.size(x)
    if (w,h) == size
        return x
    end
    y = similar(x, T, size..., c, n)
    return upsample_bilinear_whcn!(y, x)
end

function upsample_bilinear(x::AbstractArray{T,4}; size::NTuple{2,Integer}) where T<:Integer
    y = float.(x)
    res = upsample_bilinear(y; size=size)
    return round.(T, res)
end

# this is the core function which works on arrays of arbitrary size
# the implementation is a translation of https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/UpSampleMoreKernel.cpp
# which implements open-cv style linear interpolation / upsampling
# for simplicity, corners are aligned and all logic for other behaviour has been stripped
# - whcn because there is also a cwhn implementation
# - the function is parallelized using @threads
# - RGB types could be supported via reinterpreting
# - integer types need to be converted to Float and back
# - rationals work, but are slow
function upsample_bilinear_whcn!(output::AbstractArray{T,4}, input::AbstractArray{T,4}) where T
    size(input)[3:4] == size(output)[3:4] || error("Number of input and output channels and batches must match. Got input $(size(input)) and output $(size(output))")
    in_w, in_h, channels, batches = size(input)
    # treat batch and channel dimension as one for better parallelization granularity
    channels *= batches
    out_w, out_h, _, _ = size(output)
    output_slice_size = out_h * out_w

    # T() and // so that we can handle rationals (super slow)
    width_scale  = T((in_w - 1) // (out_w - 1))
    height_scale = T((in_h - 1) // (out_h - 1))

    @inline idx(c, h, w) = c * in_h * in_w + h * in_w + w + 1

    @inbounds Threads.@threads for c in 0:channels-1
        for oh in 0:out_h-1
            ih0, ih1, h0lambda, h1lambda = compute_source_index_and_lambda(height_scale, oh, in_h, out_h)
            for ow in 0:out_w-1
                iw0, iw1, w0lambda, w1lambda = compute_source_index_and_lambda(width_scale, ow, in_w, out_w)
                output_offset = c * output_slice_size + oh * out_w + ow + 1
                output[output_offset] =
                    (h0lambda * w0lambda * input[idx(c, ih0, iw0)] + # h0 * w0 * i00
                     h0lambda * w1lambda * input[idx(c, ih0, iw1)] + # h0 * w1 * i01
                     h1lambda * w0lambda * input[idx(c, ih1, iw0)] + # h1 * w0 * i10
                     h1lambda * w1lambda * input[idx(c, ih1, iw1)])  # h1 * w1 * i11
            end
        end
    end
    return output
end

"""
    ∇upsample_bilinear(Δ::AbstractArray{T,4}; size::NTuple{2,Integer}) where T

# Arguments
- `Δ`: Incoming gradient array, backpropagated from downstream layers
- `size`: Lateral (W,H) size of the image upsampled in the first place

# Outputs
- `dx`: Downsampled version of `Δ`
"""
function ∇upsample_bilinear(Δ::AbstractArray{T,4}; size::NTuple{2,Integer}) where T
    w, h, c, n = Base.size(Δ)
    out_w, out_h = size
    if (w,h) == (out_w, out_h)
        return Δ
    end
    dx = zero(similar(Δ, T, out_w, out_h, c, n))
    return ∇upsample_bilinear_whcn!(dx, Δ)
end

function ∇upsample_bilinear_whcn!(dx::AbstractArray{T,4}, Δ::AbstractArray{T,4}) where T
    size(dx)[3:4] == size(Δ)[3:4] || error("Number of input and output channels and batches must match. Got input $(size(input)) and output $(size(output))")
    in_w, in_h, channels, batches = size(dx)

    # treat batch and channel dimension as one for better parallelization granularity
    channels *= batches
    out_w, out_h, _, _ = size(Δ)
    output_slice_size = out_h * out_w

    width_scale  = T((in_w - 1) // (out_w - 1))
    height_scale = T((in_h - 1) // (out_h - 1))

    @inline idx(c, h, w) = c * in_h * in_w + h * in_w + w + 1

    @inbounds Threads.@threads for c in 0:channels-1
        for oh in 0:out_h-1
            ih0, ih1, h0lambda, h1lambda = compute_source_index_and_lambda(height_scale, oh, in_h, out_h)
            for ow in 0:out_w-1
                iw0, iw1, w0lambda, w1lambda = compute_source_index_and_lambda(width_scale, ow, in_w, out_w)
                output_offset = c * output_slice_size + oh * out_w + ow + 1
                Δ_value = Δ[output_offset]
                dx[idx(c, ih0, iw0)] += h0lambda * w0lambda * Δ_value # i00
                dx[idx(c, ih0, iw1)] += h0lambda * w1lambda * Δ_value # i01
                dx[idx(c, ih1, iw0)] += h1lambda * w0lambda * Δ_value # i10
                dx[idx(c, ih1, iw1)] += h1lambda * w1lambda * Δ_value # i11
            end
        end
    end
    return dx
end

function rrule(::typeof(upsample_bilinear), x; size)
    Ω = upsample_bilinear(x; size=size)
    function upsample_bilinear_pullback(Δ)
        (NoTangent(), ∇upsample_bilinear(unthunk(Δ); size=(Base.size(x,1),Base.size(x,2))))
    end
    return Ω, upsample_bilinear_pullback
end

###########
# trilinear
###########
"""
    upsample_trilinear(x::AbstractArray{T,5}, scale::NTuple{3,Real})
    upsample_trilinear(x::AbstractArray{T,5}; size::NTuple{3,Integer})

Upsamples the first 3 dimensions of the array `x` by the upsample factors stored in `scale`,
using trilinear interpolation. As an alternative to using `scale`, the resulting image `size`
can be directly specified with a keyword argument.

The size of the output is equal to
`(scale[1]*S1, scale[2]*S2, scale[3]*S3, S4, S5)`, where `S1, S2, S3, S4, S5 = size(x)`.

# Examples

```julia
upsample_trilinear(x, (2, 3, 4))
upsample_trilinear(x; size=(4, 9, 11))  # specify ouput size instead
upsample_trilinear(x, (2.5, 3.5, pi))  # non-integer scaling factors are allowed
```
"""
function upsample_trilinear(x::AbstractArray{<:Any,5}, scale::NTuple{3,Real})
    outsize = ntuple(i -> floor(Int, scale[i] * Base.size(x, i)), 3)
    return upsample_trilinear(x; size=outsize)
end

upsample_trilinear(x, scale::Real) = upsample_trilinear(x, (scale,scale,scale))

function upsample_trilinear(x::AbstractArray{T,5}; size::NTuple{3,Integer}) where T
    w,h,d,c,n = Base.size(x)
    if (w,h,d) == size
        return x
    end
    y = similar(x, T, size..., c, n)
    return upsample_trilinear_whdcn!(y, x)
end

function upsample_trilinear(x::AbstractArray{T,5}; size::NTuple{3,Integer}) where T<:Integer
    y = float.(x)
    res = upsample_trilinear(y; size=size)
    return round.(T, res)
end

function upsample_trilinear_whdcn!(output::AbstractArray{T,5}, input::AbstractArray{T,5}) where T
    size(input)[4:5] == size(output)[4:5] || error("Number of input and output channels and batches must match. Got input $(size(input)) and output $(size(output))")
    in_w, in_h, in_d, channels, batches = size(input)
    # treat batch and channel dimension as one for better parallelization granularity
    channels *= batches
    out_w, out_h, out_d, _, _ = size(output)
    output_slice_size = out_h * out_w * out_d

    # T() and // so that we can handle rationals (super slow)
    width_scale  = T((in_w - 1) // (out_w - 1))
    height_scale = T((in_h - 1) // (out_h - 1))
    depth_scale  = T((in_d - 1) // (out_d - 1))

    @inline idx(c, d, h, w) = c * in_d * in_h * in_w + d * in_h * in_w + h * in_w + w + 1

    @inbounds Threads.@threads for c in 0:channels-1
        for od in 0:out_d-1
            id0, id1, d0lambda, d1lambda = compute_source_index_and_lambda(depth_scale, od, in_d, out_d)
            for oh in 0:out_h-1
                ih0, ih1, h0lambda, h1lambda = compute_source_index_and_lambda(height_scale, oh, in_h, out_h)
                for ow in 0:out_w-1
                    iw0, iw1, w0lambda, w1lambda = compute_source_index_and_lambda(width_scale, ow, in_w, out_w)
                    output_offset = c * output_slice_size + od * out_w * out_h + oh * out_w + ow + 1
                    output[output_offset] =
                        d0lambda * h0lambda * w0lambda * input[idx(c, id0, ih0, iw0)] + # d0 * h0 * w0 * i000
                        d0lambda * h0lambda * w1lambda * input[idx(c, id0, ih0, iw1)] + # d0 * h0 * w1 * i001
                        d0lambda * h1lambda * w0lambda * input[idx(c, id0, ih1, iw0)] + # d0 * h1 * w0 * i010
                        d0lambda * h1lambda * w1lambda * input[idx(c, id0, ih1, iw1)] + # d0 * h1 * w1 * i011
                        d1lambda * h0lambda * w0lambda * input[idx(c, id1, ih0, iw0)] + # d1 * h0 * w0 * i100
                        d1lambda * h0lambda * w1lambda * input[idx(c, id1, ih0, iw1)] + # d1 * h0 * w1 * i101
                        d1lambda * h1lambda * w0lambda * input[idx(c, id1, ih1, iw0)] + # d1 * h1 * w0 * i110
                        d1lambda * h1lambda * w1lambda * input[idx(c, id1, ih1, iw1)]   # d1 * h1 * w1 * i111
                end
            end
        end
    end
    return output
end

"""
    ∇upsample_trilinear(Δ::AbstractArray{T,5}; size::NTuple{3,Integer}) where T

# Arguments
- `Δ`: Incoming gradient array, backpropagated from downstream layers
- `size`: Lateral size & depth (W,H,D) of the image upsampled in the first place

# Outputs
- `dx`: Downsampled version of `Δ`
"""
function ∇upsample_trilinear(Δ::AbstractArray{T,5}; size::NTuple{3,Integer}) where T
    w, h, d, c, n = Base.size(Δ)
    out_w, out_h, out_d = size
    if (w,h,d) == (out_w, out_h, out_d)
        return Δ
    end
    dx = zero(similar(Δ, T, size..., c, n))
    return ∇upsample_trilinear_whdcn!(dx, Δ)
end

function ∇upsample_trilinear_whdcn!(dx::AbstractArray{T,5}, Δ::AbstractArray{T,5}) where T
    size(dx)[4:5] == size(Δ)[4:5] || error("Number of input and output channels and batches must match. Got dx $(size(dx)) and Δ $(size(Δ))")
    in_w, in_h, in_d, channels, batches = size(dx)
    # treat batch and channel dimension as one for better parallelization granularity
    channels *= batches
    out_w, out_h, out_d, _, _ = size(Δ)
    output_slice_size = out_h * out_w * out_d

    # T() and // so that we can handle rationals (super slow)
    width_scale  = T((in_w - 1) // (out_w - 1))
    height_scale = T((in_h - 1) // (out_h - 1))
    depth_scale  = T((in_d - 1) // (out_d - 1))

    @inline idx(c, d, h, w) = c * in_d * in_h * in_w + d * in_h * in_w + h * in_w + w + 1

    @inbounds Threads.@threads for c in 0:channels-1
        for od in 0:out_d-1
            id0, id1, d0lambda, d1lambda = compute_source_index_and_lambda(depth_scale, od, in_d, out_d)
            for oh in 0:out_h-1
                ih0, ih1, h0lambda, h1lambda = compute_source_index_and_lambda(height_scale, oh, in_h, out_h)
                for ow in 0:out_w-1
                    iw0, iw1, w0lambda, w1lambda = compute_source_index_and_lambda(width_scale, ow, in_w, out_w)
                    output_offset = c * output_slice_size + od * out_w * out_h + oh * out_w + ow + 1
                    Δ_value = Δ[output_offset]
                    dx[idx(c, id0, ih0, iw0)] += d0lambda * h0lambda * w0lambda * Δ_value  # /* i000 */
                    dx[idx(c, id0, ih0, iw1)] += d0lambda * h0lambda * w1lambda * Δ_value  # /* i001 */
                    dx[idx(c, id0, ih1, iw0)] += d0lambda * h1lambda * w0lambda * Δ_value  # /* i010 */
                    dx[idx(c, id0, ih1, iw1)] += d0lambda * h1lambda * w1lambda * Δ_value  # /* i011 */
                    dx[idx(c, id1, ih0, iw0)] += d1lambda * h0lambda * w0lambda * Δ_value  # /* i100 */
                    dx[idx(c, id1, ih0, iw1)] += d1lambda * h0lambda * w1lambda * Δ_value  # /* i101 */
                    dx[idx(c, id1, ih1, iw0)] += d1lambda * h1lambda * w0lambda * Δ_value  # /* i110 */
                    dx[idx(c, id1, ih1, iw1)] += d1lambda * h1lambda * w1lambda * Δ_value  # /* i111 */
                end
            end
        end
    end
    return dx
end

function rrule(::typeof(upsample_trilinear), x; size)
    Ω = upsample_trilinear(x; size=size)
    function upsample_trilinear_pullback(Δ)
        (NoTangent(), ∇upsample_trilinear(unthunk(Δ); size=(Base.size(x,1), Base.size(x,2), Base.size(x,3))))
    end
    return Ω, upsample_trilinear_pullback
end

"""
    pixel_shuffle(x, r::Integer)

Pixel shuffling operation, upscaling by a factor `r`.

For 4-arrays representing `N` images, the operation converts input `size(x) == (W, H, r^2*C, N)`
to output of size `(r*W, r*H, C, N)`. For `D`-dimensional data, it expects `ndims(x) == D+2`
with channel and batch dimensions, and divides the number of channels by `r^D`.

Used in super-resolution networks to upsample towards high resolution features.
Reference: Shi et. al., "Real-Time Single Image and Video Super-Resolution ...", CVPR 2016, https://arxiv.org/abs/1609.05158

# Examples

```jldoctest
julia> x = [10i + j + channel/10 for i in 1:2, j in 1:3, channel in 1:4, batch in 1:1]
2×3×4×1 Array{Float64, 4}:
[:, :, 1, 1] =
 11.1  12.1  13.1
 21.1  22.1  23.1

[:, :, 2, 1] =
 11.2  12.2  13.2
 21.2  22.2  23.2

[:, :, 3, 1] =
 11.3  12.3  13.3
 21.3  22.3  23.3

[:, :, 4, 1] =
 11.4  12.4  13.4
 21.4  22.4  23.4

julia> pixel_shuffle(x, 2)  # 4 channels used up as 2x upscaling of image dimensions
4×6×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
 11.1  11.3  12.1  12.3  13.1  13.3
 11.2  11.4  12.2  12.4  13.2  13.4
 21.1  21.3  22.1  22.3  23.1  23.3
 21.2  21.4  22.2  22.4  23.2  23.4

julia> y = [i + channel/10 for i in 1:3, channel in 1:6, batch in 1:1]
3×6×1 Array{Float64, 3}:
[:, :, 1] =
 1.1  1.2  1.3  1.4  1.5  1.6
 2.1  2.2  2.3  2.4  2.5  2.6
 3.1  3.2  3.3  3.4  3.5  3.6

julia> pixel_shuffle(y, 2)  # 1D image, with 6 channels reduced to 3
6×3×1 Array{Float64, 3}:
[:, :, 1] =
 1.1  1.3  1.5
 1.2  1.4  1.6
 2.1  2.3  2.5
 2.2  2.4  2.6
 3.1  3.3  3.5
 3.2  3.4  3.6
```
"""
function pixel_shuffle(x::AbstractArray, r::Integer)
    ndims(x) > 2 || throw(ArgumentError("expected x with at least 3 dimensions"))
    d = ndims(x) - 2
    sizein = size(x)[1:d]
    cin, n = size(x, d+1), size(x, d+2)
    cin % r^d == 0 || throw(ArgumentError("expected channel dimension to be divisible by r^d = $(
        r^d), where d=$d is the number of spatial dimensions. Given r=$r, input size(x) = $(size(x))"))
    cout = cin ÷ r^d
    x = reshape(x, sizein..., ntuple(i->r, d)..., cout, n)
    perm = hcat(d+1:2d, 1:d) |> transpose |> vec  # = [d+1, 1, d+2, 2, ..., 2d, d]
    x = permutedims(x, (perm..., 2d+1, 2d+2))
    return reshape(x, map(s -> s*r, sizein)..., cout, n)
end


text:
abstract type AbstractAggregateCore end

"""
    AggregateCore{T,C1,C2}(molecules, coupling)
    AggregateCore{T,C1,C2}(molecules)

Mutable stuct which purpose is to store important information about the whole aggregate.

# Arguments
* `molecules`: Vector of molecules ([`Molecule`](@ref)).
* `coupling`: Matrix of couplings ``J_{nm}`` between molecules.
"""
struct AggregateCore{T<:Integer,C1<:ComputableType,C2<:ComputableType} <:
       AbstractAggregateCore
    molecules::Vector{Molecule{T,C1,C2}}
    coupling::Matrix{C1}
    molCount::Int64
    function AggregateCore{T,C1,C2}(
        molecules::Vector{Molecule{T,C1,C2}},
        coupling::Matrix{C1},
        molCount::Int64,
    ) where {T<:Integer,C1<:ComputableType,C2<:ComputableType}
        new(molecules, coupling, molCount)
    end
end

AggregateCore(molecules::Vector{Molecule{T,C1,C2}}, coupling::Matrix{C1}) where {T,C1,C2} =
    AggregateCore{T,C1,C2}(molecules, coupling, length(molecules))
AggregateCore(molecules::Vector{Molecule{T,C1,C2}}) where {T,C1,C2} =
    AggregateCore{T,C1,C2}(
        molecules,
        zeros(C1, (length(molecules) + 1, length(molecules) + 1)),
        length(molecules),
    )

Base.:(==)(x::AggregateCore, y::AggregateCore) =
    x.molecules == y.molecules && x.coupling == y.coupling && x.molCount == y.molCount

"""
    getNvib(aggCore)

Get maximum number of vibrational states of each molecule in the [`Aggregate`](@ref).

"""
function getNvib(aggCore::AggregateCore)
    # TODO: add typeofs
    NvibMols = Array{Array{Int64,1},1}(undef, 0)
    for mol in aggCore.molecules
        push!(NvibMols, fill(mol.Nvib, length(mol.modes)))
    end
    return NvibMols
end

"""
    getShifts(aggCore)

Get shifts for every mode on each molecule in the [`Aggregate`](@ref).

"""
function getShifts(aggCore::AggregateCore)
    # TODO: add typeofs
    shifts = Array{Array{Float64,1},1}(undef, 0)
    for mol in aggCore.molecules
        push!(shifts, getMolShifts(mol))
    end
    return shifts
end

"""
    getFrequencies(aggCore)

Get frequencies for every mode on each molecule in the [`Aggregate`](@ref).

"""
function getFrequencies(aggCore::AggregateCore)
    # TODO: add typeofs
    frequencies = Array{Array{Float64,1},1}(undef, 0)
    for mol in aggCore.molecules
        push!(frequencies, getMolFrequencies(mol))
    end
    return frequencies
end

"""
    getAggStateEnergy(agg, aggElState, aggVibState)

Get Hamiltonian of the [`Aggregate`](@ref) in a form of DenseOperator.

# Arguments
* `agg`: Instance of [`Aggregate`](@ref).
* `aggElState`: Aggregate electric state (e.g. [1, 1, 2]).
* `aggVibState`: Aggregate vibrational state (e.g. [[9, 1], [2, 5, 3], [10]]).
"""
function getAggStateEnergy(
    aggCore::AggregateCore,
    aggElState::Vector{U},
    aggVibState::Vector{Vector{U}},
)::Float64 where {U<:Integer}
    # TODO: typeof
    energy = 0.0::Float64
    for mol_i = 1:length(aggCore.molecules)
        molElState = aggElState[mol_i]
        molVibState = aggVibState[mol_i]
        energy += getMolStateEnergy(aggCore.molecules[mol_i], molElState, molVibState)
    end
    return energy
end

function elIndOrder(elInd::Vector{T})::T where {T<:Integer}
    len = length(elInd)
    ind = 1::T
    for i = 1:len
        ind += (elInd[i] - 1) * i
    end
    return ind
end

max_stars_repo_name: detrin/OpenQuantumSystems

max_issues_repo_name: detrin/OpenQuantumSystems

max_forks_repo_name: detrin/OpenQuantumSystems


text:
export check_errors, save, load_sol

######################################

using JLD2
using FileIO
using UUIDs
using Dates
using DataFrames

######################
# Helper methods for the solution of the solve method in stochasticDiffEe.jl library
######################

"""
   CHECK FOR ERRORS/WARNINGS DURING RUN

   can also remove warning trajectories by setting
                remove_warning_tr=true
"""
function check_errors(sol;remove_warning_tr=false)
    allok = true
    warnings_inxes = []
    for (i,s) in enumerate(sol)
        if (s.retcode != :Success)
            allok = false
            println("Trajectory ",i," has retcode: ", string(s.retcode))

            push!(warnings_inxes,i)

        end
    end

    if remove_warning_tr && !allok
        deleteat!(sol.u,warnings_inxes)
        println("Removing trajectories ", warnings_inxes)
    else
        println("All ok")
    end

end

function save(path, sol, args, model::AbstractModel, contour::AbstractContour, a, CC, tp, discretizationContour::AbstractContourDiscretization)
    ID = uuid1()

    if !isdir(path)
        mkdir(path)
    end

    save_sol(ID,path,sol,model,contour,a,CC,tp)
    make_config_file(ID,path,sol,args,contour,discretizationContour,model)
end

"""
   Load the sol object from filename

   returns
        sol: sol object from stochastifDiffEq.jl solver
        p: parameters used
        contour: contour
        a: lattice spacings (complex)
        CC: lattice points (complex)
        tp: t_param points (real)
"""
function save_sol(ID,path, sol, p, contour, a, CC, tp)

    path_sol = joinpath(path,"sol")
    if !isdir(path_sol)
        mkdir(path_sol)
    end

    filename = joinpath(path_sol,string(ID))

    @save filename :sol=sol :p=p :contour=contour :a=a  :CC=CC :tp=tp
end

"""
   Save the sol object with filename

   returns
        sol: sol object from stochastifDiffEq.jl solver
        p: parameters used
        contour: contour
        a: lattice spacings (complex)
        CC: lattice points (complex)
        tp: t_param points (real)
"""
function save_sol(filename, sol, p, contour, a, CC, tp)

    @save filename :sol=sol :p=p :contour=contour :a=a  :CC=CC :tp=tp
end

"""
    Make simulaiton configuration file
"""
function make_config_file(ID,path, sol, args, contour::AbstractContour, discretizationContour::AbstractContourDiscretization, model::AbstractModel)

    path_config = joinpath(path,"config")

    if !isdir(path_config)
        mkdir(path_config)
    end

    filename = joinpath(path_config,string(ID))

    line_shift = "\n"

    open(filename, "w") do io
        write(io, string("ID: ",        ID,                line_shift,
                  "Time: ",      Time(Dates.now()), line_shift,
                  "Date: ",      Date(Dates.now()), line_shift,
                  "Time_used: ", sol.elapsedTime,         line_shift)
            )
        write(io,"************ CONTOUR *************",line_shift)
        write(io, "contour: ", string(typeof(contour)), line_shift)
        write(io, string(contour),              line_shift)

        write(io,"************ DISCRETIZATION *************",line_shift)
        write(io, "discretizationContour: ", string(typeof(discretizationContour)), line_shift)
        write(io, string(discretizationContour), line_shift)

        write(io, "*********** MODEL ***********",line_shift)
        write(io, "model: ", string(typeof(model)), line_shift)
        write(io, string(model), line_shift)

        write(io, "*********** SOLVER ***********",line_shift)
        write(io, "algorithm: ", split(string(typeof(sol[1].alg)),"{")[1], line_shift)
        if hasproperty(sol[1].alg,:theta)
            write(io, "theta: ", string(sol[1].alg.theta), line_shift)
            write(io, "tol: ", string(args.tol), line_shift)
            write(io, "dtmax: ", string(args.dtmax), line_shift)
        end
        write(io, "tspan: ", string(sol[1].prob.tspan), line_shift)
        write(io, "dt: ", string(args.dt), line_shift)
        write(io, "N_Tr: ", string(args.N_Tr), line_shift)

    end
end

"""
   Load the sol object from filename

   returns
        sol: sol object from stochastifDiffEq.jl solver
        p: parameters used
        contour: contour
        a: lattice spacings (complex)
        CC: lattice points (complex)
        tp: t_param points (real)
"""
function load_sol(filename)
    lo = load(filename)
    sol = lo[":sol"]
    p = lo[":p"]
    contour = lo[":contour"]
    a = lo[":a"]
    CC = lo[":CC"]
    tp = lo[":tp"]

    return sol, p, contour, a, CC, tp
end

"""
   Read config file

   returns
        dictionary of all the config options
"""
function read_config(file_path)
    #df = DataFrame(ID = String[], Time = String[], Date = String[], Time_used = Float64[],
    #               contour = String[], tmax=Float64[],β=Float64[],A_β=Float64[],F_β=Float64[],
    #               discretizationContour=String[], NR_Points=Int[],NR_RT_FW_Points=Int[],
    #               NR_RT_BW_Points=Int[],NR_ET_Points=Int[],
    #               model=String[],g=Float64[],m=Float64[],λ=Float64[],
    #               algorithm=)
    D = Dict{String,Any}("filename"=>file_path)
    open(file_path) do f
        for l in eachline(f)

            # Check for int
            if !isnothing(try D[split(l,":")[1]] = parse(Int,split(l,":")[2]) catch ArgumentError end)
                continue
            end

            # Check for int
            if !isnothing(try D[split(l,":")[1]] = parse(Float64,split(l,":")[2]) catch ArgumentError end)
                continue
            end

            try
                D[split(l,":")[1]] = strip(join(split(l,":")[2:end],":"))
            catch BoundsError
            end

        end
    end

    return D
end

"""
   Get all configurations in path

   returns
        dataframe of all config options as columns and
"""
function get_configs(path)

    path_configs = joinpath(path,"configs")

    df = DataFrame()
    for file in readdir(path_configs,join=true)
        D = read_config(file)

        push!(df,D,cols=:union)
    end

    return df
end

"""
   Load from ID

   returns
        dataframe of all config options as columns and
"""
function load_sol_from_ID(ID,path)

    filename = joinpath(joinpath(path,"sol"),ID)

    return load_sol(filename)
end


text:
### 10 August 2016: currently this is not used. it has not proven to be faster than IHT itself
### left here just in case it proves useful later
export L0_reg_aiht

# DOUBLE OVERRELAXATION SCHEME TO ACCELERATE IHT
#
# This subroutine computes the (D)ouble (O)ver(R)elaxation for accelerating IHT steps.
function dor{T <: Float}(
    x       :: DenseArray{T,2},
    y       :: DenseArray{T,1},
    b       :: DenseArray{T,1},
    b0      :: DenseArray{T,1},
    b00     :: DenseArray{T,1},
    xb      :: DenseArray{T,1},
    xb0     :: DenseArray{T,1},
    xb00    :: DenseArray{T,1},
#    sortidx :: DenseArray{Int,1},
    IDX     :: BitArray{1},
    z1      :: DenseArray{T,1},
    z2      :: DenseArray{T,1},
    xz1     :: DenseArray{T,1},
    xz2     :: DenseArray{T,1},
    dif     :: DenseArray{T,1},
    r       :: DenseArray{T,1},
    r2      :: DenseArray{T,1},
    obj     :: Float,
    k       :: Int;
    n       :: Int = length(y),
    p       :: Int = length(b)
)

    # calculate first overrelaxation
    difference!(dif, xb, xb0, n=n)                  # dif = xb - xb0
    a1  = dot(dif,r) / sumabs2(dif)                 # a1  = dif'*r / || dif ||
    ypatzmw!(z1, b, a1, b, b0, n=p)                 # z1  = b + a1 * (b - b0)
    difference!(xz1, xb, xb0, n=n, a=1.0+a1, b=a1)  # xz1 = (1+a1)*xb - a1*xb0
    difference!(r2,y,xz1, n=n)                      # r2  = y - xz1

    # calculate second overrelaxation
    difference!(dif, xz1, xb00, n=n)                # dif = xz1 - xb00
    a2 = dot(dif,r2) / sumabs2(dif)                 # a2  = dif'*r2 / || dif ||
    ypatzmw!(z2, z1, a2, z1, b00)                   # z2 = z1 + a2 * (z1 - b00)

    # project z2 onto sparsity set
    project_k!(z2, k)

    # update residual information about z2
#    update_xb!(xz2, x, z2, sortidx, k, n=n, p=p)
    update_xb!(xz2, x, z2, IDX, k, n=n, p=p)
    difference!(r2,y,xz2)

    # if z2 is better than b0, then overwrite b with z2
    if sumabs(r2) < obj * 2
        println("Successful acceleration")
        copy!(b,z2)
        copy!(r,r2)
        copy!(xb,xz2)
    end

    return nothing
end

# ACCELERATED ITERATIVE HARD THRESHOLDING
#
# This function computes a hard threshold update
#
#    b+ = P_{S_k}(b + mu*X'(y - Xb))
#
# where mu is the step size (or learning rate) and P_{S_k} denotes the projection onto the set S_k defined by
#
#     S_k = { x in R^p : || x ||_0 <= k }.
#
# The projection in question preserves the largest k components of b in magnitude, and it sends the remaining
# p-k components to zero. This update is intimately related to a projected gradient step used in Landweber iteration.
# Unlike the Landweber method, this function performs a line search on mu whenever the step size exceeds a specified
# threshold omega given by
#
#     omega = || b+ - b ||_2^2 / || X(b+ - b) ||_2^2.
#
# By backtracking on mu, this function guarantees a stable estimation of a sparse b.
#
# This variant of IHT uses the double overrelaxation scheme described by Blumensath (2012) to accelerate convergence.
#
# Arguments:
#
# -- b is the iterate of p model components;
# -- x is the n x p design matrix;
# -- y is the vector of n responses;
# -- k is the model size;
# -- g is the negative gradient X'*(Y - Xbeta);
#
# Optional Arguments:
#
# -- p is the number of predictors. Defaults to length(b).
# -- n is the number of samples. Defaults to length(y).
# -- b0 is the previous iterate beta. Defaults to b.
# -- b00 is the previous b0. Defaults to b.
# -- xb = x*b.
# -- xb0 = x*b0.
# -- xb00 = x*b00.
# -- xk is a temporary array to store the k columns of x corresponding to the support of b.
# -- gk is a temporary array of k floats used to subset the k components of the gradient g with the support of b.
# -- xgk = x*gk.
# -- max_step is the maximum number of backtracking steps to take. Defaults to 50.
# -- sortidx is a vector to store the indices that would sort beta. Defaults to p zeros of type Int.
# -- betak is a vector to store the largest k values of beta. Defaults to k zeros of type T.
# -- IDX and IDX0 are BitArrays indicating the nonzero status of components of beta. They default to falses.
# -- r and r2 store the overall residual and partial accelerated residual, respectively. They default to zeroes.
# -- z1 and z2 store intermediate steps of the acceleration calculations. They default to zeroes.
# -- xz1 and xz2 store x*z1 and x*z2. They default to zeroes.
# -- dif is a temporary array to store the difference of two vectors
#
# coded by Kevin L. Keys (2015)
# klkeys@g.ucla.edu
# based on the HardLab demonstration code written in MATLAB by Thomas Blumensath
# http://www.personal.soton.ac.uk/tb1m08/sparsify/sparsify.html
function aiht{T <: Float}(
    x         :: DenseArray{T,2},
    y         :: DenseArray{T,1},
    b         :: DenseArray{T,1},
    g         :: DenseArray{T,1},
    obj       :: Float,
    k         :: Int,
    iter      :: Int;
    n         :: Int               = length(y),
    p         :: Int               = length(b),
    xk        :: DenseArray{T,2}   = zeros(T,n,k),
    b0        :: DenseArray{T,1}   = copy(b),
    b00       :: DenseArray{T,1}   = copy(b),
    xb        :: DenseArray{T,1}   = BLAS.gemv('N', one(T), x, b),
    xb0       :: DenseArray{T,1}   = copy(xb),
    xb00      :: DenseArray{T,1}   = copy(xb),
    gk        :: DenseArray{T,1}   = zeros(T,k),
    xgk       :: DenseArray{T,1}   = zeros(T,n),
    r         :: DenseArray{T,1}   = zeros(T,n),
    r2        :: DenseArray{T,1}   = zeros(T,n),
    z1        :: DenseArray{T,1}   = zeros(T,p),
    z2        :: DenseArray{T,1}   = zeros(T,p),
    dif       :: DenseArray{T,1}   = zeros(T,n),
    xz1       :: DenseArray{T,1}   = zeros(T,n),
    xz2       :: DenseArray{T,1}   = zeros(T,n),
#    sortidx   :: DenseArray{Int,1} = collect(1:p),
    IDX       :: BitArray{1}       = falses(p),
    IDX0      :: BitArray{1}       = copy(IDX),
    step_mult :: Float             = one(T),
    max_step  :: Int               = 50
)

    # which components of beta are nonzero?
    update_indices!(IDX, b, p=p)

    # if current vector is 0,
    # then take largest elements of d as nonzero components for b
    if sum(IDX) == 0
#        selectpermk!(sortidx,g,k, p=p)
#        IDX[sortidx[1:k]] = true
        a = select(g, k, by=abs, rev=true)
        threshold!(IDX, g, abs(a), n=p)
    end

    # store relevant columns of x
    # do so only if support has changed after first iteration
    if !isequal(IDX, IDX0) || iter < 2
        update_xk!(xk, x, IDX, k=k, p=p, n=n)   # xk = x[:,IDX]
    end

    # store relevant components of gradient
    fill_perm!(gk, g, IDX, k=k, p=p)    # gk = g[IDX]

    # now compute subset of x*g
    BLAS.gemv!('N', one(T), xk, gk, zero(T), xgk)

    # compute step size
    mu = step_mult * sumabs2(sdata(gk)) / sumabs2(sdata(xgk))
    isfinite(mu) || throw(error("Step size is not finite, is active set all zero?"))

    # take gradient step
    BLAS.axpy!(p, mu, sdata(g), 1, sdata(b), 1)

    # preserve top k components of b
    project_k!(b, k)

    # which indices of new beta are nonzero?
    copy!(IDX0, IDX)
    update_indices!(IDX, b, p=p)

    # must correct for equal entries at kth pivot of b
    # this is a total hack! but matching magnitudes are very rare
    # should not drastically affect performance, esp. with big data
    # hack randomly permutes indices of duplicates and retains one
    if sum(IDX) > k
        a = select(b, k, by=abs, rev=true)          # compute kth pivot
        duples = find(x -> abs(x) .== abs(a), b)    # find duplicates
        c = randperm(length(duples))                # shuffle
        d = duples[c[2:end]]                        # permute, clipping top
        b[d] = zero(T)                             # zero out duplicates
        IDX[d] = false                              # set corresponding indices to false
    end

    # update xb
#    update_xb!(xb, x, b, sortidx, k)
    update_xb!(xb, x, b, IDX, k)

    # update residuals
    difference!(r,y,xb, n=n)

    if iter > 2
#        dor(x, y, b, b0, b00, xb, xb0, xb00, sortidx, z1, z2, xz1, xz2, dif, r, r2, obj, k, n=n, p=p)
        dor(x, y, b, b0, b00, xb, xb0, xb00, IDX, z1, z2, xz1, xz2, dif, r, r2, obj, k, n=n, p=p)
    end

    # calculate omega
    omega_top = sqeuclidean(sdata(b),(b0))
    omega_bot = sqeuclidean(sdata(xb),sdata(xb0))

    # backtrack until mu sits below omega and support stabilizes
    mu_step = 0
    while mu*omega_bot > 0.99*omega_top && sum(IDX) != 0 && sum(IDX $ IDX0) != 0 && mu_step < max_step

        # stephalving
        mu /= 2

        # recompute gradient step
        copy!(b,b0)
        BLAS.axpy!(p, mu, sdata(g), 1, sdata(b), 1)

        # recompute projection onto top k components of b
        project_k!(b, k)

        # which indices of new beta are nonzero?
        update_indices!(IDX, b, p=p)

        # must correct for equal entries at kth pivot of b
        # this is a total hack! but matching magnitudes are very rare
        # should not drastically affect performance, esp. with big data
        # hack randomly permutes indices of duplicates and retains one
        if sum(IDX) > k
            a = select(b, k, by=abs, rev=true)          # compute kth pivot
            duples = find(x -> abs(x) .== abs(a), b)    # find duplicates
            c = randperm(length(duples))                # shuffle
            d = duples[c[2:end]]                        # permute, clipping top
            b[d] = zero(T)                             # zero out duplicates
            IDX[d] = false                              # set corresponding indices to false
        end

        # recompute xb
#        update_xb!(xb, x, b, sortidx, k)
        update_xb!(xb, x, b, IDX, k)

        # update residuals
        difference!(r,y,xb, n=n)

        if iter > 2
#            dor(x, y, b, b0, b00, xb, xb0, xb00, sortidx, z1, z2, xz1, xz2, dif, r, r2, obj, k, n=n, p=p)
            dor(x, y, b, b0, b00, xb, xb0, xb00, IDX, z1, z2, xz1, xz2, dif, r, r2, obj, k, n=n, p=p)
        end

        # calculate omega
        omega_top = sqeuclidean(sdata(b),sdata(b0))
        omega_bot = sqeuclidean(sdata(xb),sdata(xb0))

        # increment the counter
        mu_step += 1
    end

    return mu, mu_step
end

# ACCELERATED L0 PENALIZED LEAST SQUARES REGRESSION
#
# This routine solves the optimization problem
#
#     min 0.5*|| Y - XB ||_2^2
#
# subject to
#
#     B in S_k = { x in R^p : || x ||_0 <= k }.
#
# It uses Thomas Blumensath's accelerated iterative hard thresholding framework to keep B feasible.
# Note that the acceleration need not kick in, and this function could run slower than its IHT counterpart.
#
# Arguments:
# -- X is the n x p data matrix
# -- Y is the n x 1 continuous response vector
# -- k is the desired size of the support (active set)
#
# Optional Arguments:
# -- b is the p x 1 iterate. Warm starts should use this argument. Defaults to marginals:
#        b = cov(X,Y) / var(X)
# -- max_iter is the maximum number of iterations for the algorithm. Defaults to 1000.
# -- max_step is the maximum number of backtracking steps for the step size calculation. Defaults to 50.
# -- tol is the global tol. Defaults to 1e-4.
# -- quiet is a Boolean that controls algorithm output. Defaults to true (no output).
# -- several temporary arrays for intermediate steps of algorithm calculations:
#       Xk        = zeros(n,k)      # store k columns of X
#       r         = zeros(n)        # for || Y - XB ||_2^2
#       Xb        = zeros(n)        # X*beta
#       Xb0       = zeros(n)        # X*beta0
#       Xb00      = zeros(n)        # X*beta00
#       b0        = zeros(p)        # previous iterate beta0
#       b00       = zeros(p)        # previous beta0
#       df        = zeros(p)        # (negative) gradient
#       idx       = zeros(k)        # another temporary array of k floats
#       tempn     = zeros(n)        # temporary array of n floats
#       indices   = collect(1:p)    # indices that sort beta
#       support   = falses(p)       # indicates nonzero components of beta
#       support0  = copy(support)   # store previous nonzero indicators
#       r2 stores the overall residual and partial accelerated residual, respectively. They default to zeroes.
#       z1 and z2 store intermediate steps of the acceleration calculations. They default to zeroes.
#       xz1 and xz2 store x*z1 and x*z2. They default to zeroes.
#       dif is a temporary array to store the difference of two vectors
#
# Outputs are wrapped into a Dict with the following fields:
# -- time is the compute time for the algorithm. Note that this does not account for time spent initializing optional argument defaults
# -- iter is the number of iterations that the algorithm took
# -- loss is the optimal loss (residual sum of squares divided by sqrt of RSS with previous iterate)
# -- beta is the final iterate
#
# coded by Kevin L. Keys (2015)
# klkeys@g.ucla.edu
function L0_reg_aiht{T <: Float}(
    X        :: DenseArray{T,2},
    Y        :: DenseArray{T,1},
    k        :: Int;
    n        :: Int               = length(Y),
    p        :: Int               = size(X,2),
    Xk       :: DenseArray{T,2}   = zeros(T,n,k),
    b        :: DenseArray{T,1}   = zeros(T,p),
    b0       :: DenseArray{T,1}   = zeros(T,p),
    b00      :: DenseArray{T,1}   = zeros(T,p),
    df       :: DenseArray{T,1}   = zeros(T,p),
    z1       :: DenseArray{T,1}   = zeros(T,p),
    z2       :: DenseArray{T,1}   = zeros(T,p),
    Xb       :: DenseArray{T,1}   = zeros(T,n),
    Xb0      :: DenseArray{T,1}   = zeros(T,n),
    Xb00     :: DenseArray{T,1}   = zeros(T,n),
    tempn    :: DenseArray{T,1}   = zeros(T,n),
    r        :: DenseArray{T,1}   = zeros(T,n),
    r2       :: DenseArray{T,1}   = zeros(T,n),
    dif      :: DenseArray{T,1}   = zeros(T,n),
    xz1      :: DenseArray{T,1}   = zeros(T,n),
    xz2      :: DenseArray{T,1}   = zeros(T,n),
    idx      :: DenseArray{T,1}   = zeros(T,k),
    indices  :: DenseArray{Int,1} = collect(1:p),
    support  :: BitArray{1}       = falses(p),
    support0 :: BitArray{1}       = falses(p),
    tol      :: Float             = convert(T, 1e-4),
    max_iter :: Int               = 1000,
    max_step :: Int               = 50,
    quiet    :: Bool              = true
)

    # start timer
    tic()

    # first handle errors
    k            >= 0      || throw(ArgumentError("Value of k must be nonnegative!\n"))
    max_iter     >= 0      || throw(ArgumentError("Value of max_iter must be nonnegative!\n"))
    max_step     >= 0      || throw(ArgumentError("Value of max_step must be nonnegative!\n"))
    tol          >  eps(T) || throw(ArgumentError("Value of global tolerance must exceed machine precision!\n"))

    # initialize return values
    mm_iter   = 0                   # number of iterations of L0_reg
    mm_time   = zero(T)             # compute time *within* L0_reg
    next_obj  = oftype(tol, Inf)    # objective value
    next_loss = oftype(tol, Inf)    # loss function value

    # initialize floats
    current_obj = oftype(tol, Inf)   # tracks previous objective function value
    the_norm    = zero(T)            # norm(b - b0)
    scaled_norm = zero(T)            # the_norm / (norm(b0) + 1)
    mu          = zero(T)            # IHT step size

    # initialize integers
    i       = 0          # used for iterations in loops
    mu_step = 0          # counts number of backtracking steps for mu

    # initialize booleans
    converged = false    # scaled_norm < tol?

    # update X*beta
    update_xb!(Xb, X, b, indices, k, p=p, n=n)

    # update r and gradient
    difference!(r, Y, Xb, n=n)
    BLAS.gemv!('T', one(T), X, r, zero(T), df)

    # update loss and objective
    next_loss = Inf
    next_obj  = next_loss

    # guard against numerical instabilities
    isnan(next_loss) && throw(error("Loss function is NaN, something went wrong..."))

    # formatted output to monitor algorithm progress
    if !quiet
         println("\nBegin MM algorithm\n")
         println("Iter\tHalves\tMu\t\tNorm\t\tObjective")
         println("0\t0\tInf\t\tInf\t\tInf")
    end

    # main loop
    for mm_iter = 1:max_iter

        # notify and break if maximum iterations are reached.
        if mm_iter >= max_iter

            if !quiet
                print_with_color(:red, "MM algorithm has hit maximum iterations $(max_iter)!\n")
                print_with_color(:red, "Current Objective: $(current_obj)\n")
            end

            # send elements below tol to zero
            threshold!(b, tol, n=p)

            # calculate r piecemeal
            difference!(r, Y, Xb, n=n)

            # calculate loss and objective
            next_loss = sumabs2(r) / 2

            # stop timer
            mm_time = toq()

            # these are output variables for function
            # wrap them into a Dict and return
            output = Dict{String, Any}("time" => mm_time, "loss" => next_loss, "iter" => mm_iter, "beta" => b)

            return output
        end

        # save values from previous iterate
        copy!(b00,b0)           # b00  = b
        copy!(b0,b)             # b0   = b
        copy!(Xb00,Xb0)         # Xb00 = Xb0
        copy!(Xb0,Xb)           # Xb0  = Xb
        current_obj = next_obj

        # now perform AIHT step
#        (mu, mu_step) = aiht(X,Y,b,df,current_obj,k,mm_iter, n=n, p=p, max_step=max_step, IDX=support, IDX0=support0, b0=b0, xb=Xb, xb0=Xb0, xgk=tempn, xk=Xk, sortidx=indices, gk=idx, step_mult=one(T), b00=b00, r=r, r2=r2, z1=z1, z2=z2, dif=dif, xz1=xz1, xz2=xz2)
        (mu, mu_step) = aiht(X,Y,b,df,current_obj,k,mm_iter, n=n, p=p, max_step=max_step, IDX=support, IDX0=support0, b0=b0, xb=Xb, xb0=Xb0, xgk=tempn, xk=Xk, gk=idx, step_mult=one(T), b00=b00, r=r, r2=r2, z1=z1, z2=z2, dif=dif, xz1=xz1, xz2=xz2)

        # the IHT kernel gives us an updated x*b
        # use it to recompute residuals
        difference!(r, Y, Xb, n=n)

        # finally, recompute the gradient
        BLAS.gemv!('T', one(T), X, r, zero(T), df)

        # update loss, objective, and gradient
        next_loss = sumabs2(r) / 2
        next_obj  = next_loss

        # guard against numerical instabilities
        isnan(next_loss) && throw(error("Loss function is NaN, something went wrong..."))
        isinf(next_loss) && throw(error("Loss function is NaN, something went wrong..."))

        # track convergence
        the_norm    = chebyshev(b,b0)
        scaled_norm = the_norm / ( norm(b0,Inf) + one(T))
        converged   = scaled_norm < tol

        # output algorithm progress
        quiet || @printf("%d\t%d\t%3.7f\t%3.7f\t%3.7f\n", mm_iter, mu_step, mu, the_norm, next_obj)

        # check for convergence
        # if converged and in feasible set, then algorithm converged before maximum iteration
        # perform final computations and output return variables
        if converged

            # send elements below tol to zero
            threshold!(b, tol, n=p)

            # stop time
            mm_time = toq()

            if !quiet
                println("\nMM algorithm has converged successfully.")
                println("MM Results:\nIterations: $(mm_iter)")
                println("Final Loss: $(next_loss)")
                println("Total Compute Time: $(mm_time)")
            end

            # these are output variables for function
            # wrap them into a Dict and return
            output = Dict{String, Any}("time" => mm_time, "loss" => next_loss, "iter" => mm_iter, "beta" => b)

            return output
        end

        # algorithm is unconverged at this point.
        # if algorithm is in feasible set, then rho should not be changing
        # check descent property in that case
        # if rho is not changing but objective increases, then abort
        if next_obj > current_obj + tol
            if !quiet
                print_with_color(:red, "\nMM algorithm fails to descend!\n")
                print_with_color(:red, "MM Iteration: $(mm_iter)\n")
                print_with_color(:red, "Current Objective: $(current_obj)\n")
                print_with_color(:red, "Next Objective: $(next_obj)\n")
                print_with_color(:red, "Difference in objectives: $(abs(next_obj - current_obj))\n")
            end

            throw(error("Descent failure"))
            output = Dict{String, Any}("time" => -one(T), "loss" => -one(T), "iter" => -1, "beta" => fill!(b,Inf))

            return output
        end
    end # end main loop
end # end function


text:
# https://adventofcode.com/2020/day/19

example1 = """
0: 1 2
1: "a"
2: 1 3 | 3 1
3: "b"
"""

example2 = """
0: 4 1 5
1: 2 3 | 3 2
2: 4 4 | 5 5
3: 4 5 | 5 4
4: "a"
5: "b"
"""

function preprocess_rule(rule)
    if '"' in rule
        return strip(replace(rule, '"' => ""))
    end

    if '|' in rule
        rule = join(map(s -> "(?:$s)", split(rule, '|')), '|')
    end
    rule = replace(rule, r" +" => " ")
    rule = replace(rule, r"(\d+)" => s"(?:\1)")
    rule = "(?:$rule)"
    return rule
end

function parse_input(string)
    rules = Dict()
    for row in split(string, '\n', keepempty=false)
        key, rule = split(row, ":")
        rules[strip(key)] = preprocess_rule(rule)
    end
    return rules
end

function reduce_rule!(rules, name)
    replacement = pop!(rules, name)

    for (key, rule) in rules
        if key == name
            continue
        end
        rules[key] = replace(rule, "(?:$name)" => "(?:$replacement)")
        #rules[key] = replace(rules[key], r"\(\s+([a-z]+)\s+\)" => s"\1")
    end

    return rules
end

function reduce_rules(rules, ignore=["0"])
    select = map(s -> !isnothing(match(r"^[a-z ]+$", s)), values(rules))
    terminal_nodes = collect(keys(rules))[select]

    while length(terminal_nodes) > 0
        value = popfirst!(terminal_nodes)
        reduce_rule!(rules, value)

        select = map(s -> isnothing(match(r"[0-9]", s)), values(rules))
        terminal_nodes = collect(keys(rules))[select]
        terminal_nodes = [x for x in terminal_nodes if !(x in ignore)]
    end

    return rules
end

function remove_extra_characters(str)
    str = replace(str, " " => "")
    str = replace(str, r"\(?:([a-z]+)\)\(?:([a-z]+)\)" => s"\1\2")
    str = replace(str, r"(?<!\|)\(\?:([a-z]+)\)(?!\|)" => s"\1")
    return str
end

reduce_and_parse(input) = remove_extra_characters(reduce_rules(parse_input(input))["0"])

function isvalid(message, rule)
    return !isnothing(match(Regex("^$rule\$"), message))
end

@assert isvalid("aab", reduce_and_parse(example1))
@assert isvalid("aba", reduce_and_parse(example1))
@assert !isvalid("aaa", reduce_and_parse(example1))
@assert !isvalid("abaa", reduce_and_parse(example1))
@assert isvalid("aaaabb", reduce_and_parse(example2))
@assert isvalid("abbabb", reduce_and_parse(example2))
@assert isvalid("aabaab", reduce_and_parse(example2))
@assert isvalid("abaaab", reduce_and_parse(example2))

example3 = """
0: 4 1 5
1: 2 3 | 3 2
2: 4 4 | 5 5
3: 4 5 | 5 4
4: "a"
5: "b"

ababbb
bababa
abbbab
aaabbb
aaaabbb
"""

function count_valid_messages(rule, messages)
    return sum(isvalid.(split(messages, '\n'), rule))
end

function part1(input)
    rules, messages = split(input, "\n\n", keepempty=false)
    rule = reduce_and_parse(rules)
    return count_valid_messages(rule, messages)
end

@assert part1(example3) == 2

example4 = """
42: 9 14 | 10 1
9: 14 27 | 1 26
10: 23 14 | 28 1
1: "a"
11: 42 31
5: 1 14 | 15 1
19: 14 1 | 14 14
12: 24 14 | 19 1
16: 15 1 | 14 14
31: 14 17 | 1 13
6: 14 14 | 1 14
2: 1 24 | 14 4
0: 8 11
13: 14 3 | 1 12
15: 1 | 14
17: 14 2 | 1 7
23: 25 1 | 22 14
28: 16 1
4: 1 1
20: 14 14 | 1 15
3: 5 14 | 16 1
27: 1 6 | 14 18
14: "b"
21: 14 1 | 1 14
25: 1 1 | 1 14
22: 14 14
8: 42
26: 14 22 | 1 20
18: 15 15
7: 14 5 | 1 21
24: 14 1

abbbbbabbbaaaababbaabbbbabababbbabbbbbbabaaaa
bbabbbbaabaabba
babbbbaabbbbbabbbbbbaabaaabaaa
aaabbbbbbaaaabaababaabababbabaaabbababababaaa
bbbbbbbaaaabbbbaaabbabaaa
bbbababbbbaaaaaaaabbababaaababaabab
ababaaaaaabaaab
ababaaaaabbbaba
baabbaaaabbaaaababbaababb
abbbbabbbbaaaababbbbbbaaaababb
aaaaabbaabaaaaababaa
aaaabbaaaabbaaa
aaaabbaabbaaaaaaabbbabbbaaabbaabaaa
babaaabbbaaabaababbaabababaaab
aabbbbbaabbbaaaaaabbbbbababaaaaabbaaabba
"""

function part2(input)
    rules, messages = split(input, "\n\n", keepempty=false)
    rules = reduce_rules(parse_input(rules), ["0", "8", "11", "42", "31"])

    # 8: 42 | 42 8
    # 11: 42 31 | 42 11 31

    rules["8"] = "(?:(?:42)+)"
    rules["11"] = "((?:42)(?1)?(?:31))"

    for key in ["31", "42", "8", "11"]
        reduce_rule!(rules, key)
    end
    rule = remove_extra_characters(rules["0"])
    matched_messages = count_valid_messages(rule, messages)

    return matched_messages
end

@assert part1(example4) == 3
@assert part2(example4) == 12

test = read("data/day-19.txt", String)
println("Part 1: $(result1 = part1(test))")
println("Part 2: $(result2 = part2(test))")

@assert result1 == 168
@assert result2 == 277

max_stars_repo_name: twolodzko/advent-of-code

max_issues_repo_name: twolodzko/advent-of-code

max_forks_repo_name: twolodzko/advent-of-code


text:
# This file is a part of SimilaritySearch.jl

using Random
export BeamSearch

"""
    BeamSearch(bsize::Integer=16, beam=KnnResult(bsize))

BeamSearch is an iteratively improving local search algorithm that explores the graph using blocks of `bsize` elements and neighborhoods at the time.
Multithreading applications must have copies of this object due to shared cache objects.

- `bsize`: The size of the beam.
"""
@with_kw mutable struct BeamSearch <: LocalSearchAlgorithm
    bsize::Int32 = 8  # size of the search beam
end

Base.copy(bsearch::BeamSearch;
        bsize=bsearch.bsize
    ) = BeamSearch(; bsize)

function Base.copy!(dst::BeamSearch, src::BeamSearch)
    dst.bsize = src.bsize
end

Base.string(s::BeamSearch) = """{BeamSearch: bsize=$(s.bsize)}"""

const GlobalBeamKnnResult = [KnnResult(32)]  # see __init__ function
@inline function getbeam(bs::BeamSearch)
    @inbounds beam = GlobalBeamKnnResult[Threads.threadid()]
    empty!(beam, bs.bsize)
    beam
end

function __init__beamsearch()
    for i in 2:Threads.nthreads()
        push!(GlobalBeamKnnResult, KnnResult(32))
    end
end

### local search algorithm

function beamsearch_queue(index::SearchGraph, q, res::KnnResult, objID, vstate)
    @inbounds if !visited(vstate, objID)
        visit!(vstate, objID)
        d = evaluate(index.dist, q, index[objID])
        push!(res, objID, d)
    end
end

function beamsearch_init(bs::BeamSearch, index::SearchGraph, q, res::KnnResult, hints, vstate)
    for objID in hints
        beamsearch_queue(index, q, res, objID, vstate)
    end

    if length(res) == 0
        _range = 1:length(index)
        for i in 1:bs.bsize
           objID = rand(_range)
           beamsearch_queue(index, q, res, objID, vstate)
       end
    end
end

function beamsearch_inner(index::SearchGraph, q, res::KnnResult, beam::KnnResult, vstate)
    while length(beam) > 0
        prev_id, prev_dist = popfirst!(beam)
        @inbounds for childID in keys(index.links[prev_id])
            if !visited(vstate, childID)
                visit!(vstate, childID)
                d = evaluate(index.dist, q, index[childID])
                push!(res, childID, d) && push!(beam, childID, d)
                #d <= 0.9 * farthest(res).dist && push!(beam, childID, d)
            end
        end
    end
end

"""
    search(bs::BeamSearch, index::SearchGraph, q, res::KnnResult, hints, vstate)

Tries to reach the set of nearest neighbors specified in `res` for `q`.
- `bs`: the parameters of `BeamSearch`
- `index`: the local search index
- `q`: the query
- `res`: The result object, it stores the results and also specifies the kind of query
- `hints`: Starting points for searching, randomly selected when it is an empty collection
- `vstate`: A dictionary like object to store the visiting state of vertices
"""
function search(bs::BeamSearch, index::SearchGraph, q, res::KnnResult, hints, vstate)
    beamsearch_init(bs, index, q, res, hints, vstate)
    beam = getbeam(bs)
    push!(beam, first(res))
    beamsearch_inner(index, q, res, beam, vstate)
    res
end


text:
using GenerativeAD
using FileIO, BSON
using ValueHistories, DistributionsAD
using Flux
using ConditionalDists
using GenerativeModels
using DrWatson
using StatsBase
using EvalMetrics
using LinearAlgebra

# all the models will use the lowest anomaly score on normal validation data
master_path = datadir("experiments/images_leave-one-out")
outpath = datadir("experiments/images_leave-one-out_clean_val_score")
master_path = datadir("experiments/images_leave-one-in")
outpath = datadir("experiments/images_leave-one-in_clean_val_score")
mkpath(outpath)

models = readdir(master_path)
seeds = 1:1
acs = 1:10

function val_score(f)
	try
		data = load(f)
		val_score = Flux.mean(data[:val_scores][data[:val_labels].==0])
		return val_score
	catch e
		return NaN
	end
end
function collect_scores(path)
	try
		fs = readdir(path)
		fs = filter(f->(f[1:5] != "model"), fs)
		scores = map(f->val_score(joinpath(path,f)), fs)
		return scores
	catch e
		@warn "$path does not probably exist"
		return []
	end
end

function test_auc(f)
	try
		data = load(f)
		auc = EvalMetrics.auc_trapezoidal(EvalMetrics.roccurve(data[:tst_labels], vec(data[:tst_scores]))...)
		return auc
	catch e
		return NaN
	end
end
function collect_aucs(path)
	fs = readdir(path)
	fs = filter(x->!(occursin("model", x)), fs)
	scores = map(f->test_auc(joinpath(path,f)), fs)
end

function select_best_model(mp, dataset, model, acs)
	dp = joinpath(mp, dataset)
	sps = map(ac->joinpath(dp, "ac=$ac/seed=1"), acs) # paths of seed results
	sfs = map(x->filter(f->(f[1:5] != "model"), isdir(x) ? readdir(x) : []), sps) # all filenames
	scores = map(collect_scores, sps) # all scores

	best_models = []
	for (sf, score) in zip(sfs, scores)
		best_ind = argmin(score)
		push!(best_models, sf[best_ind])
	end

	if length(best_models) > 0
		return best_models, sps
	else
		return nothing, nothing
	end
end

function copy_models(acs, sps, best_models, dataset, outpath, model)
	for (ac, sp, bm) in zip(acs, sps, best_models)
		f = joinpath(sp, bm)
		if isfile(f)
			outp = joinpath(outpath, model, dataset, "ac=$ac/seed=1")
			mkpath(outp)
			outf = joinpath(outp, bm)
			cp(f, outf, force=true)
		end
	end
end

model = models[1]
dataset = "MNIST"

for model in models
	@info "Processing model $model..."
	mp = joinpath(master_path, model)
	datasets = readdir(mp)
	for dataset in datasets
		@info "   $dataset"
		best_models, sps = select_best_model(mp, dataset, model, acs)
		if !isnothing(best_models)
			copy_models(acs, sps, best_models, dataset, outpath, model)
		else
			@warn "No usable data found, skipping"
		end
	end
end


text:
@inline M{T}(i::T)           = (i == 1) ? T(1) : T(2)^(i - T(1)) + T(1)
@inline cc_iM(i::Int)           = (i == 1) ? 1 : Int(log2(i - 1) + 1)
@inline dM{T}(i::T)           = (i < T(3)) ? T(2^(i - 1)) : T(2^(i - 2))
@inline cc_g(i::Int, j::Int)     = (i == 1) ?  0.5  : (j - 1) / (M(i) - 1.0)
@inline cc_g(i::Int)            = Float64[cc_g(i, j) for j = 1:M(i)]
@inline cc_dg(i::Int, j::Int)    = i == 2 ? [0.0, 1.0][j] : cc_g(i, 2j)
@inline cc_dg(i::Int)           = Float64[cc_dg(i, j) for j = 1:dM(i)]

@inline function cc_bf_l(x::Float64, xij::Float64, mi)
    if (mi == 1)
        return 1.0
    end
    dx = (1.0 - (mi - 1.0) * abs(x - xij))
    return (dx > 0.0) ? dx : 0.0
end

cc_bf_l(x::Float64, xij::Float64) = cc_bf_l(x, xij, M(level(xij)))
cc_bf_l(x::Vector{Float64}, xij::Vector{Float64}) = prod(map(cc_bf_l, x, xij))

@inline function cc_bf_q(x::Float64, xij::Float64, mi)
	if (mi == 1)
		return 1.0
	end
	dx = 1.0 - ((mi - 1.0) * (x - xij))^2
	return (dx > 0.0) ?  dx : 0.0
end

cc_bf_q(x::Float64, xij::Float64) = cc_bf_q(x, xij, cc_M(level(xij)))
cc_bf_q(x::Vector{Float64}, xij::Vector{Float64}) = map(cc_bf_q, x, xij)

function cc_simpsonsw(i::Int, j::Int)
    @assert 1 ≤ j ≤ M(i)
    if i == 1
        return 1.0
    end
    w = 1 / (M(i) - 1) / 3
    if j == 1 || j == M(i)
        return w
    else
        w *= mod(j, 2) == 0 ? 4 : 2
    end
    return w
end

cc_simpsonsw(i::Int) = Float64[cc_simpsonsw(i, j) for j = 1:M(i)]

function cc_dsimpsonsw(i::Int)
    if i == 1
        return cc_simpsonsw(i)
    elseif i == 2
        return cc_simpsonsw(i) - [0;1.0;0]
    else
        return cc_simpsonsw(i) - vec(Float64[cc_simpsonsw(i - 1)'zeros(1, M(i - 1))])[1:M(i - 1) - 1]
    end
end


text:
# v0.6.0

for e in readdir() # Current directory
    println(e)
end

# Same for...
readdir("~") # Read home directory
readdir("~/documents")

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# ============================================================
# I/O Methods
# ============================================================

export load_data, load_model, save_model

"""
    load_data(file; mode = :csv, dlm = " ")

Load dataset from file
"""
function load_data(file; mode = :csv, dlm = nothing)
    if mode == :csv
        dataset = CSV.File(file; delim = dlm) |> DataFrame
    end

    return dataset
end

"""
    load_model(file; mode = :jld)

Load the trained machine learning model
"""
function load_model(file::T; kwargs...) where {T<:AbstractString}
    if file[end-3:end] == "jld2"
        nn = load_model(file, :jld)
    else
        nn = load_model(file, :tf)
    end

    return nn
end

function load_model(file::T, mode) where {T<:AbstractString}
    if mode == :jld
        JLD2.@load file nn
    elseif mode == :tf
        copy!(tf, pyimport("tensorflow"))
        nn = tf.keras.models.load_model(file)
    end

    return nn
end

"""
    save_model(nn; mode = :jld)

Save the trained machine learning model
"""
function save_model(nn; mode = :jld)
    if mode == :jld
        JLD2.@save "model.jld2" nn
    end
end

function save_model(nn::PyObject)
    nn.save("model.h5")
end


text:
baremodule Julio

import DefineSingletons

# I/O
function open end

# Task
function withtaskgroup end
function spawn! end

# Cancellation
function cancelscope end
function cancel! end
function oncancel end
function shield end
function withtimeout end
function timeout end
function sleep end

# Queue-like data structures
function channel end
function stack end
function queue end
function clone end

function maybetake! end
function tryput! end
function maybefetch end

# Low-level APIs
function select end
function sync end
function yield end
function checkpoint end
function iscancelled end

# DefineSingletons.@def_singleton var"try" isa Function
DefineSingletons.@def_singleton Events

macro SharedResource end

module Internal

import ..Julio: @SharedResource
using ..Julio: Julio

using ArgCheck: @argcheck
using ContextVariablesX
using ContextManagers: ContextManagers
using Reagents
using Reagents:
    Block,
    CAS,
    Computed,
    Identity,
    Map,
    PostCommit,
    Read,
    Reagents,
    Retry,
    Return,
    Update,
    WithNack

struct PleaseUseBaseDotReturns end
const Returns = PleaseUseBaseDotReturns()

include("utils.jl")
include("tracing.jl")

include("core.jl")
include("promises.jl")
include("base.jl")
include("stacks.jl")
include("queues.jl")
include("blocking.jl")
include("channels.jl")
include("taskgroups.jl")
include("locks.jl")

include("resizable.jl")
include("io.jl")

include("api.jl")
include("docs.jl")

function __init__()
    try
        Base.Experimental.register_error_hint(on_channel_handle_method_error, MethodError)
    catch err
        @warn "Failed to register an error hint" exception = (err, catch_backtrace())
    end
end

end  # module Internal

const Promise = Internal.Promise
const Lock = Internal.Lock

Internal.define_docstrings()

end  # baremodule Julio


text:
##########################
#     SOLVE METHODS      #
##########################

#
# function solve!(prob::Problem{T,Discrete}, solver::SequentialNewtonSolver) where T
#     V_ = newton_step!(prob, solver)
#     copyto!(prob.X, V_.X)
#     copyto!(prob.U, V_.U)
#     projection!(prob)
#     return solver
# end
#
# """
# Take a Projected Newton step using a purely sequential method (never forms block arrays)
# """
# function newton_step!(prob, solver::SequentialNewtonSolver)
#     V = solver.V
#     verbose = solver.opts.verbose
#
#     # Initial stats
#     update!(prob, solver)
#     J0 = cost(prob, V)
#     res0 = res2(solver, V)
#     viol0 = max_violation(solver)
#
#     # Projection
#     verbose ? println("\nProjection:") : nothing
#     projection!(prob, solver, V)
#     res1, = multiplier_projection!(solver, V)
#
#     # Solve KKT
#     J1 = cost(prob, V)
#     viol1 = max_violation(solver)
#     δx, δu, δλ = solveKKT(solver, V)
#
#     # Line Search
#     verbose ? println("\nLine Search") : nothing
#     V_ = line_search(prob, solver, δx, δu, δλ)
#     J_ = cost(prob, V_)
#     res_ = res2(solver, V_)
#     viol_ = max_violation(solver)
#
#     # Print Stats
#     if verbose
#         println("\nStats")
#         println("cost: $J0 → $J1 → $J_")
#         println("res: $res0 → $res1 → $res_")
#         println("viol: $viol0 → $viol1 → $viol_")
#     end
#
#     return V_
# end
#
# function line_search(prob::Problem, solver::SequentialNewtonSolver, δx, δu, δλ)
#     α = 1.0
#     s = 0.01
#     V = solver.V
#     J0 = cost(prob, V)
#     update!(prob, solver)
#     res0 = res2(solver, V)
#     count = 0
#     solver.opts.verbose ? println("res0: $res0") : nothing
#     while count < 10
#         V_ = V + (α.*(δx,δu,δλ),solver.active_set)
#
#         # Calculate residual
#         projection!(prob, solver, V_)
#         res, = multiplier_projection!(solver, V_)
#         J = cost(prob, V_)
#
#         # Calculate max violation
#         viol = max_violation(solver)
#
#         if solver.opts.verbose
#             println("cost: $J \t residual: $res \t feas: $viol")
#         end
#         if res < (1-α*s)*res0
#             solver.opts.verbose ? println("α: $α") : nothing
#             return V_
#         end
#         count += 1
#         α /= 2
#     end
#     return solver.V
# end
#
# """
# Calculates Hinv*Y'*((Y*Hinv*Y')\\y), which is the least norm solution to the problem
#     min xᵀHx
#      st Y*x = y
# where y is the vector of active constraints and Y is the constraint jacobian
# stores the result in vals.x and vals.u, with the intermediate lagrange multiplier stored in vals.λ
# """
# function _projection(solver::SequentialNewtonSolver)
#     y = active_constraints(solver)
#     calc_factors!(solver)
#     δλ = solve_cholesky(solver, y)
#
#     δx,δu = jac_T_mult(solver, δλ)
#     δx = -solver.Qinv .* δx
#     δu = -solver.Rinv .* δu
#     return δx, δu, δλ
# end
#
#
# function projection!(prob, solver::SequentialNewtonSolver, V::PrimalDualVars, active_set_update=true)
#     X,U = V.X, V.U
#     eps_feasible = solver.opts.feasibility_tolerance
#     count = 0
#     # cost_expansion!(prob, solver, V)
#     feas = Inf
#     while true
#         dynamics_constraints!(prob, solver, V)
#         update_constraints!(prob, solver, V)
#         dynamics_jacobian!(prob, solver, V)
#         constraint_jacobian!(prob, solver, V)
#         if active_set_update
#             active_set!(prob, solver)
#         end
#         y = active_constraints(solver)
#
#         viol = maximum(norm.(y,Inf))
#         if solver.opts.verbose
#             println("feas: ", viol)
#         end
#         if viol < eps_feasible || count > 10
#             break
#         else
#             δx, δu = _projection(solver)
#             _update_primals!(V, δx, δu)
#             count += 1
#         end
#     end
# end
#
# """
# Solve the least-squares problem to find the best multipliers,
#     given the current active multipliers λ
# """
# function _mult_projection(solver::SequentialNewtonSolver, λ)
#     N = length(solver.Q)
#
#     # Calculate Y'λ
#     δx,δu = jac_T_mult(solver, λ)
#
#     # Calculate g + Y'λ
#     for k = 1:N-1
#         δx[k] += solver.Q[k].x
#         δu[k] += solver.Q[k].u
#     end
#     δx[N] += solver.Q[N].x
#
#     # Calculate Y*(g + Y'λ)
#     r = jac_mult(solver, δx, δu)
#
#     # Solve (Y*Y')\r
#     eyes = [I for k = 1:N]
#     calc_factors!(solver, eyes, eyes)
#     δλ = solve_cholesky(solver, r)
#
#     return -δλ
# end
#
# function multiplier_projection!(solver::SequentialNewtonSolver, V)
#     λa = active_duals(V, solver.active_set)
#     δλ = _mult_projection(solver, λa)
#     _update_duals!(V, δλ, solver.active_set)
#     return res2(solver,V), δλ
# end
#
# """
# Solve the KKT system [H Y'][δz] = [g + Y'λ]
#                      [Y 0 ][δλ]   [   y   ]
#     using a sequential cholesky factorization of the Shur compliment S = (Y*Hinv*Y')
#     '''
#     r = g + Y'λ
#     δλ = Sinv*(y - Y*Hinv*r)
#     δz = -Hinv*(r + Y'δλ)
#     '''
# """
# function _solveKKT(solver::SequentialNewtonSolver, λ)
#     N = length(solver.Q)
#
#     # Calculate r = g + Y'λ
#     rx,ru = jac_T_mult(solver, λ)
#     for k = 1:N-1
#         rx[k] = solver.Q[k].x + rx[k]
#         ru[k] = solver.Q[k].u + ru[k]
#     end
#     rx[N] = solver.Q[N].x + rx[N]
#
#     # Calculate b0 = Y*Hinv*r
#     δx = solver.Qinv .* rx
#     δu = solver.Rinv .* ru
#     b0 = jac_mult(solver, δx, δu)
#
#     # Calculate y - Y*Hinv*r
#     y = active_constraints(solver)
#     b = y - b0
#
#     # Solve for δλ = (Y*Hinv*Y')\b
#     calc_factors!(solver)
#     δλ = solve_cholesky(solver, b)
#
#     # Solve for δz = -Hinv*(r + Y'δλ)
#     δx, δu = jac_T_mult(solver, δλ)
#     for k = 1:N-1
#         δx[k] = -solver.Qinv[k]*(δx[k] + rx[k])
#         δu[k] = -solver.Rinv[k]*(δu[k] + ru[k])
#     end
#     δx[N] = -solver.Qinv[N]*(δx[N] + rx[N])
#     return δx, δu, δλ
# end
#
# function solveKKT(solver::SequentialNewtonSolver, V)
#     λa = active_duals(V, solver.active_set)
#     δx, δu, δλ = _solveKKT(solver, λa)
#     return δx, δu, δλ
# end

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
module SquarenessTest

using Squareness
using Test

@testset "Bool" begin
    @test issquare(false)
    @test issquare(true)
end

@testset "Int8" begin
    @testset "is a square" begin
        dataset_sq_i8 = Int8[i^2 for i=0:11]
        @test all(issquare, dataset_sq_i8)
    end

    @testset "is not a square" begin
        dataset_nosq_i8 = setdiff(typemin(Int8):typemax(Int8), Int8[i^2 for i=0:11])
        @test all(!issquare, dataset_nosq_i8)
    end
end

@testset "UInt8" begin
    @testset "is a square" begin
        dataset_sq_ui8 = UInt8[i^2 for i=0:15]
        @test all(issquare, dataset_sq_ui8)
    end

    @testset "is not a square" begin
        dataset_nosq_ui8 = setdiff(typemin(UInt8):typemax(UInt8), UInt8[i^2 for i=0:15])
        @test all(!issquare, dataset_nosq_ui8)
    end
end

@testset "Int16" begin
    @testset "is a square" begin
        dataset_sq_i16 = Int16[i^2 for i=0:181]
        @test all(issquare, dataset_sq_i16)
    end

    @testset "is not a square" begin
        dataset_nosq_i16 = setdiff(typemin(Int16):typemax(Int16), Int16[i^2 for i=0:181])
        @test all(!issquare, dataset_nosq_i16)
    end
end

@testset "UInt16" begin
    @testset "is a square" begin
        dataset_sq_ui16 = UInt16[i^2 for i=0:255]
        @test all(issquare, dataset_sq_ui16)
    end

    @testset "is not a square" begin
        dataset_nosq_ui16 = setdiff(typemin(UInt16):typemax(UInt16), UInt16[i^2 for i=0:255])
        @test all(!issquare, dataset_nosq_ui16)
    end
end

@testset "Int32" begin
    @testset "is a square" begin
        dataset_sq_i32 = Int32[i^2 for i=0:46340]
        @test all(issquare, dataset_sq_i32)
    end

    @testset "is not a square" begin
        # dataset_nosq_i32 = setdiff(typemin(Int32):typemax(Int32), Int32[i^2 for i=0:46340])
        dataset_nosq_i32 = setdiff(rand(Int32, 1000000), Int32[i^2 for i=0:46340])
        @test all(!issquare, dataset_nosq_i32)
    end
end

@testset "UInt32" begin
    @testset "is a square" begin
        dataset_sq_ui32 = UInt32[i^2 for i=0:65535]
        @test all(issquare, dataset_sq_ui32)
    end

    @testset "is not a square" begin
        # dataset_nosq_ui32 = setdiff(typemin(UInt32):typemax(UInt32), UInt32[i^2 for i=0:65535])
        dataset_nosq_ui32 = setdiff(rand(UInt32, 1000000), UInt32[i^2 for i=0:65535])
        @test all(!issquare, dataset_nosq_ui32)
    end
end

@testset "Int64" begin
    @testset "is a square" begin
        # dataset_sq_i64 = Int64[i^2 for i=0:3037000499]
        dataset_sq_i64 = Int64[i^2 for i in rand(0:3037000499, 1000000)]
        @test all(issquare, dataset_sq_i64)
    end

    @testset "random (may not be a square)" begin
        dataset_random_i64 = rand(Int64, 1000000)
        @test all(@.(issquare(dataset_random_i64) == Squareness.issquare_with_isqrt(dataset_random_i64)))
    end
end

@testset "UInt64" begin
    @testset "is a square" begin
        # dataset_sq_ui64 = UInt64[i^2 for i=0:4294967295]
        dataset_sq_ui64 = UInt64[i^2 for i in rand(UInt64(0):UInt64(4294967295), 1000000)]
        @test all(issquare, dataset_sq_ui64)
    end

    @testset "random (may not be a square)" begin
        dataset_random_ui64 = rand(UInt64, 1000000)
        @test all(@.(issquare(dataset_random_ui64) == Squareness.issquare_with_isqrt(dataset_random_ui64)))
    end
end

@testset "Int128" begin
    @testset "is a square" begin
        # dataset_sq_i128 = Int128[i^2 for i=0:13043817825332782212]
        dataset_sq_i128 = Int128[i^2 for i in rand(int128"0":int128"13043817825332782212", 1000000)]
        @test all(issquare, dataset_sq_i128)
    end

    @testset "random (may not be a square)" begin
        dataset_random_i128 = rand(Int128, 1000000)
        @test all(@.(issquare(dataset_random_i128) == Squareness.issquare_with_isqrt(dataset_random_i128)))
    end
end

@testset "UInt128" begin
    @testset "is a square" begin
        # dataset_sq_ui128 = UInt128[i^2 for i=0:18446744073709551615]
        dataset_sq_ui128 = UInt128[i^2 for i in rand(uint128"0":uint128"18446744073709551615", 1000000)]
        @test all(issquare, dataset_sq_ui128)
    end

    @testset "random (may not be a square)" begin
        dataset_random_ui128 = rand(UInt128, 1000000)
        @test all(@.(issquare(dataset_random_ui128) == Squareness.issquare_with_isqrt(dataset_random_ui128)))
    end
end

@testset "BigInt" begin
    @testset "is a square" begin
        dataset_sq_mpz = BigInt[i^2 for i in (big.(rand(UInt128, 10000)) .<< 128 .+ rand(UInt128, 10000))]
        @test all(issquare, dataset_sq_mpz)
    end

    @testset "random (may not be a square)" begin
        dataset_random_mpz = big.(rand(Int128, 10000)) .<< 128 .+ rand(Int128, 10000)
        @test all(@.(issquare(dataset_random_mpz) == Squareness.issquare_with_isqrt(dataset_random_mpz)))
    end
end

end


text:
@testset "Test SplineSparse" begin

    # vals should be length 4
    @test_throws ErrorException SplineSparse{Float64,Int,1,2}(2, rand(3), [1, 1])

    for i in 1:10
        @test BasisMatrices.n_chunks(SplineSparse{Float64,Int64,i,i}) == i
        @test BasisMatrices.chunk_len(SplineSparse{Float64,Int64,i,i}) == i
    end

    @testset "Array" begin
        s = SplineSparse(1, 2, 3, 1:4, 1:2)
        want = [1 2 0; 0 3 4]
        @test Array(s) == want

        s = SplineSparse(2, 2, 6, 1:8, [1, 5, 2, 4])
        want = [1 2 0 0 3 4
                0 5 6 7 8 0]

        @test Array(s) == want
    end

    @testset "findnz" begin
        s = SplineSparse(1, 2, 3, 1:4, 1:2)
        want = ([1, 1, 2, 2], [1, 2, 2, 3], [1, 2, 3, 4])
        @test findnz(s) == want

        s = SplineSparse(2, 2, 6, 1:8, [1, 5, 2, 4])
        want = (
            [1, 1, 1, 1, 2, 2, 2, 2],
            [1, 2, 5, 6, 2, 3, 4, 5],
            [1, 2, 3, 4, 5, 6, 7, 8]
        )
        @test findnz(s) == want
    end

    @testset "row_kron" begin
        s1 = SplineSparse(1,2, 3, rand(1:10, 12), rand(1:2, 6))
        s2 = SplineSparse(1,2, 4, rand(1:10, 12), rand(1:3, 6))
        want = row_kron(Array(s1), Array(s2))
        s12 = row_kron(s1, s2)

        @test Array(s12) == want
        @test Array(row_kron(s12, s1)) == row_kron(Array(s12), Array(s1))
        @test Array(row_kron(s12, s12)) == row_kron(Array(s12), Array(s12))

        Base.zero(::Type{String}) = ""

        s1 = SplineSparse(1,2, 2, ["a", "b", "c", "d"], [1, 1])
        s2 = SplineSparse(1,3, 3, map(string, 1:6), [1, 1])

        want = ["a1" "a2" "a3" "b1" "b2" "b3"
                "c4" "c5" "c6" "d4" "d5" "d6"]
        have = row_kron(s1, s2)

        @test Array(have) == want

    end

    @testset "getindex" begin
        s = SplineSparse(1, 2, 10, rand(12), rand(1:9, 6))
        full_s = Array(s)

        for r in 1:size(s, 1)
            for c in 1:size(s, 1)
                @test s[r, c] == full_s[r, c]
            end
        end
    end

    @testset "*" begin
        s1 =  SplineSparse(1, 2, 3, 1:4, [1, 2])
        s2 =  SplineSparse(1, 3, 4, 1:6, [2, 1])
        s12 = SplineSparse(2, 3, 12, [1, 2, 3, 2, 4, 6, 12, 15, 18, 16, 20, 24],
                           [2, 6, 5, 9])

        x1 = rand(3)
        x2 = rand(4)
        x12 = rand(12)

        x1m = rand(3, 3)
        x2m = rand(4, 3)
        x12m = rand(12, 3)

        @test s1*x1 ≈ Array(s1)*x1
        @test s2*x2 ≈ Array(s2)*x2
        @test s12*x12 ≈ Array(s12)*x12

        @test s1*x1m ≈ Array(s1)*x1m
        @test s2*x2m ≈ Array(s2)*x2m
        @test s12*x12m ≈ Array(s12)*x12m

        # RowKron specialization
        rk = RowKron(s1, s2)

        @test rk * x12 ≈ Array(s12)*x12
        @test rk * x12m ≈ Array(s12)*x12m

        # mutating...
        out12 = zeros(2)
        out12m = zeros(2, 3)

        mul!(out12, rk, x12)
        @test out12 ≈ Array(s12) * x12

        mul!(out12m, rk, x12m)
        @test out12m ≈ Array(s12) * x12m

        # matrix evaluation of RowKron

        # errors with RowKron
        rk2 = RowKron(s1, s2, s12)
        @test_throws ArgumentError mul!(zeros(size(rk2,1)), rk2, rand(size(rk2, 2)))
        @test_throws DimensionMismatch mul!(zeros(size(rk,1)+1), rk, rand(size(rk, 2)))

    end

end


text:
@testset "mutable graph" begin

    g = complete_digraph(5)
    abm = ABM(Agent5, GraphSpace(g))
    for n in 1:5
        for _ in 1:5
            add_agent!(n, abm, rand(abm.rng))
        end
    end
    @test nv(abm) == 5
    @test abm.space.s[1] == 1:5
    @test abm.space.s[2] == 6:10

    rem_node!(abm, 2)
    @test nv(abm) == 4
    @test nagents(abm) == 20
# Last node became 2nd node (swapped places as per Graphs.jl)
    @test abm.space.s[2] == 21:25
    @test length(abm.space.s) == 4

    n = add_node!(abm)
    @test n == 5
    @test length(abm.space.s) == 5
    @test nv(abm) == 5
    add_edge!(abm, n, 2)

    a = add_agent!(n, abm, rand(abm.rng))
    ids = nearby_ids(a, abm)
    @test sort(ids) == 21:25

    model = ABM(Agent5, GraphSpace(SimpleGraph()))
    @test_throws ArgumentError add_agent!(model, rand(model.rng))
    add_node!(model)
    @test nv(model) == 1
    @test add_edge!(model, 1, 2) == false
    add_node!(model)
    @test add_edge!(model, 1, 2) == true

end


text:
abstract type EqualConstraint <: Constraint end

"""
    EqualConstant(x::SeaPearl.AbstractIntVar, v::Int)

Equality constraint, putting a constant value `v` for the variable `x` i.e. `x == v`.
"""
struct EqualConstant <: EqualConstraint
    x       ::SeaPearl.AbstractIntVar
    v       ::Int
    active  ::StateObject{Bool}
    function EqualConstant(x::SeaPearl.AbstractIntVar, v::Int, trailer)
        constraint = new(x, v, StateObject{Bool}(true, trailer))
        addOnDomainChange!(x, constraint)
        return constraint
    end
end

"""
    propagate!(constraint::EqualConstant, toPropagate::Set{Constraint}, prunedDomains::CPModification)

`EqualConstant` propagation function. Basically set the `x` domain to the constant value.
"""
function propagate!(constraint::EqualConstant, toPropagate::Set{Constraint}, prunedDomains::CPModification)
    # Reduce the domain to a singleton if possible
    if constraint.v in constraint.x.domain
        removed = assign!(constraint.x.domain, constraint.v)
        setValue!(constraint.active, false)
        triggerDomainChange!(toPropagate, constraint.x)
        addToPrunedDomains!(prunedDomains, constraint.x, removed)
        return true
    end

    # Reduce the domain to an empty set if value not in domain, not feasible => no propagation
    removed = removeAll!(constraint.x.domain)
    setValue!(constraint.active, false)
    addToPrunedDomains!(prunedDomains, constraint.x, removed)
    return false
end

variablesArray(constraint::EqualConstant) = [constraint.x]

function Base.show(io::IO, ::MIME"text/plain", con::EqualConstant)
    println(io, typeof(con), ": ", con.x.id, " == ", con.v, ", active = ", con.active)
    println(io, "   ", con.x)
end

function Base.show(io::IO, con::EqualConstant)
    print(io, typeof(con), ": ", con.x.id, " == ", con.v)
end

"""
    Equal(x::SeaPearl.AbstractIntVar, y::SeaPearl.AbstractIntVar)

Equality constraint between two variables, stating that `x == y`.
"""
struct Equal <: EqualConstraint
    x       ::SeaPearl.AbstractIntVar
    y       ::SeaPearl.AbstractIntVar
    active  ::StateObject{Bool}

    function Equal(x::SeaPearl.AbstractIntVar, y::SeaPearl.AbstractIntVar, trailer::Trailer)
        constraint = new(x, y, StateObject(true, trailer))
        addOnDomainChange!(x, constraint)
        addOnDomainChange!(y, constraint)
        return constraint
    end
end

"""
    propagate!(constraint::Equal, toPropagate::Set{Constraint}, prunedDomains::CPModification)

`Equal` propagation function.
"""
function propagate!(constraint::Equal, toPropagate::Set{Constraint}, prunedDomains::CPModification)
    if !constraint.active.value
        return true
    end
    xFormerLength = length(constraint.x.domain)
    yFormerLength = length(constraint.y.domain)

    prunedX = pruneEqual!(constraint.x, constraint.y)
    prunedY = pruneEqual!(constraint.y, constraint.x)

    if !isempty(prunedX)
        triggerDomainChange!(toPropagate, constraint.x)
        addToPrunedDomains!(prunedDomains, constraint.x, prunedX)
    end
    if !isempty(prunedY)
        triggerDomainChange!(toPropagate, constraint.y)
        addToPrunedDomains!(prunedDomains, constraint.y, prunedY)
    end

    if constraint in toPropagate
        pop!(toPropagate, constraint)
    end

    if length(constraint.x.domain) <= 1
        setValue!(constraint.active, false)
    end
    if isempty(constraint.x.domain) || isempty(constraint.y.domain)
        return false
    end
    return true
end

"""
    pruneEqual!(x::AbstractIntVar, y::AbstractIntVar)

Remove the values from the domain of `x` that are not in the domain of `y`.
"""
function pruneEqual!(x::AbstractIntVar, y::AbstractIntVar)
    toRemove = Int[]
    for val in x.domain
        if !(val in y.domain)
            push!(toRemove, val)
        end
    end

    for val in toRemove
        remove!(x.domain, val)
    end

    return toRemove
end

variablesArray(constraint::Equal) = [constraint.x, constraint.y]

function Base.show(io::IO, ::MIME"text/plain", con::Equal)
    println(io, typeof(con), ": ", con.x.id, " == ", con.y.id, ", active = ", con.active)
    println(io, "   ", con.x)
    print(io, "   ", con.y)
end

function Base.show(io::IO, con::Equal)
    print(io, typeof(con), ": ", con.x.id, " == ", con.y.id)
end