text:
using Gen

@testset "optional positional args (calling + GFI)" begin

    @gen function foo(x, y=2, (grad)(z::Float64=3.0))
        @param theta::Float64
        a = @trace(normal(x+y+z, 1), :a)
        b = @trace(normal(a+theta, 1), :b)
        return (x, y, z, x+y+z)
    end

    # initialize theta to zero for non-gradient tests
    init_param!(foo, :theta, 0.)

    # test directly calling with varying args
    @test foo(1) == (1, 2, 3, 6)
    @test foo(5, 4) == (5, 4, 3, 12)
    @test foo(2, 3, 4) == (2, 3, 4, 9)

    # test simulate using default args
    tr = simulate(foo, (1,))
    @test get_args(tr) == (1, 2, 3)
    @test get_retval(tr) == (1, 2, 3, 6)

    # test generate using default args
    tr, w = generate(foo, (1,), choicemap(:a => 5))
    @test get_args(tr) == (1, 2, 3)
    @test isapprox(w, logpdf(normal, 5, 1+2+3, 1))

    # test update with varying args and argdiffs
    new_tr, w_diff, _, _ = update(tr, (2, 3, 4),
        (UnknownChange(),  UnknownChange(), UnknownChange()), choicemap())
    @test get_args(new_tr) == (2, 3, 4)
    @test new_tr[:a] == tr[:a]
    @test new_tr[:b] == tr[:b]
    @test isapprox(w_diff, logpdf(normal, new_tr[:a], 2+3+4, 1) -
                           logpdf(normal, tr[:a], 1+2+3, 1))

    new_tr, w_diff, _, _ = update(tr, (2,),
        (UnknownChange(),), choicemap(:a => 6))
    @test get_args(new_tr) == (2, 2, 3)
    @test new_tr[:a] == 6
    @test new_tr[:b] == tr[:b]
    @test isapprox(w_diff, (logpdf(normal, new_tr[:a], 2+2+3, 1) +
                            logpdf(normal, new_tr[:b], new_tr[:a], 1)) -
                           (logpdf(normal, tr[:a], 1+2+3, 1) +
                            logpdf(normal, tr[:b], tr[:a], 1)))

    # test regenerate with varying args and argdiffs
    new_tr, w_diff, _ = regenerate(tr, (2, 3, 4),
        (UnknownChange(),  UnknownChange(), UnknownChange()), select(:b))
    @test get_args(new_tr) == (2, 3, 4)
    @test new_tr[:a] == tr[:a]
    expected_score = (logpdf(normal, new_tr[:a], 2+3+4, 1) +
                      logpdf(normal, new_tr[:b], new_tr[:a], 1))
    @test isapprox(get_score(new_tr), expected_score)
    expected_diff = (logpdf(normal, new_tr[:a], 2+3+4, 1) -
                     logpdf(normal, tr[:a], 1+2+3, 1))
    @test isapprox(w_diff, expected_diff)

    new_tr, w_diff, _ = regenerate(tr, (2, 3),
        (UnknownChange(), UnknownChange()), select(:b))
    @test get_args(new_tr) == (2, 3, 3)
    @test new_tr[:a] == tr[:a]
    expected_score = (logpdf(normal, new_tr[:a], 2+3+3, 1) +
                      logpdf(normal, new_tr[:b], new_tr[:a], 1))
    @test isapprox(get_score(new_tr), expected_score)
    expected_diff = (logpdf(normal, new_tr[:a], 2+3+3, 1) -
                     logpdf(normal, tr[:a], 1+2+3, 1))
    @test isapprox(w_diff, expected_diff)

    # test choice_gradients and accumulate_param_gradients!
    function foo_lpdf(x, y, z, a, b)
        lpdf = 0.
        lpdf += logpdf(normal, a, x+y+z, 1)
        lpdf += logpdf(normal, b, a, 1)
        return lpdf
    end

    tr, _ = generate(foo, (2,), choicemap(:a => 5, :b => 6))
    arg_grads = accumulate_param_gradients!(tr, nothing)
    @test length(arg_grads) == 3 # return gradients for all args, incl. optional
    _, _, z_grad = arg_grads
    @test isapprox(z_grad, finite_diff(foo_lpdf, (2, 2, 3, 5, 6), 3, dx))

    arg_grads, _, _ = choice_gradients(tr, select(:a), nothing)
    @test length(arg_grads) == 3 # return gradients for all args, incl. optional
    _, _, z_grad = arg_grads
    @test isapprox(z_grad, finite_diff(foo_lpdf, (2, 2, 3, 5, 6), 3, dx))

    # test nested calls
    @gen function bar()
        return @trace(foo(1), :foo)
    end
    @test bar() == (1, 2, 3, 6)

    tr, w = generate(bar, (), choicemap((:foo => :a) => 5))
    @test get_args(tr) == ()
    sub_tr = Gen.get_call(tr, :foo).subtrace
    @test get_args(sub_tr) == (1, 2, 3) # args of subtrace should have defaults
    @test get_score(tr) == get_score(sub_tr)
    @test isapprox(w, logpdf(normal, 5, 1+2+3, 1))

    # test optional GenerativeFunction arguments
    @gen function outer(inner::GenerativeFunction=foo)
        @trace(inner(1), :inner)
    end

    @test outer() == (1, 2, 3, 6)
    tr, w = generate(outer, (), choicemap((:inner => :a) => 5))
    @test get_args(tr) == (foo,)
    sub_tr = Gen.get_call(tr, :inner).subtrace
    @test get_args(sub_tr) == (1, 2, 3) # args of subtrace should have defaults
    @test get_score(tr) == get_score(sub_tr)
    @test isapprox(w, logpdf(normal, 5, 1+2+3, 1))

    @gen function foo2(x)
        @trace(normal(x, 1), :a)
        return x
    end
    @test outer(foo2) == 1
    tr, w = generate(outer, (foo2,), choicemap((:inner => :a) => 5))
    @test get_args(tr) == (foo2,)
    sub_tr = Gen.get_call(tr, :inner).subtrace
    @test get_args(sub_tr) == (1,)
    @test get_score(tr) == get_score(sub_tr)
    @test isapprox(w, logpdf(normal, 5, 1, 1))

end

@testset "optional positional args (combinators)" begin

    @gen function foo(t::Int, y_prev::Bool, z1::Float64=0.25, z2::Float64=0.75)
        y = @trace(bernoulli(y_prev ? z1 : z2), :y)
    end

    # test call_at
    at_foo = call_at(foo, Int)
    addr = 3
    tr, w = generate(at_foo, (0, true, addr), choicemap((addr => :y) => true))
    @test tr[addr => :y] == true
    @test isapprox(w, logpdf(bernoulli, tr[addr => :y], 0.25))
    tr, w = generate(at_foo, (0, true, 0.5, addr), choicemap((addr => :y) => true))
    @test tr[addr => :y] == true
    @test isapprox(w, logpdf(bernoulli, tr[addr => :y], 0.5))

    # test map
    map_foo = Map(foo)
    addr = 1
    tr, w = generate(map_foo,
        ([1, 2, 3], [false, true, false]),
        choicemap((addr => :y) => true))
    @test tr[addr => :y] == true
    @test isapprox(w, logpdf(bernoulli, tr[addr => :y], 0.75))
    tr, w = generate(map_foo,
        ([1, 2, 3], [true, false, true], [0.5, 0.25, 0.0]),
        choicemap((addr => :y) => true))
    @test tr[addr => :y] == true
    @test isapprox(w, logpdf(bernoulli, tr[addr => :y], 0.5))

    # test unfold
    unfold_foo = Unfold(foo)
    addr = 1
    tr, w = generate(unfold_foo, (3, false),
        choicemap((addr => :y) => true))
    @test tr[addr => :y] == true
    @test isapprox(w, logpdf(bernoulli, tr[addr => :y], 0.75))
    tr, w = generate(unfold_foo, (3, false, 0.05, 0.95),
        choicemap((addr => :y) => true))
    @test tr[addr => :y] == true
    @test isapprox(w, logpdf(bernoulli, tr[addr => :y], 0.95))

    # test recurse
    @gen function production(depth::Int, branches::Int=1)
        n_children = @trace(bernoulli(0.5), :branch) ? branches : 0
        return Production(depth, [depth+1 for i=1:n_children])
    end

    @gen function aggregation(depth::Int, child_outputs::Vector, combiner=sum)
        if length(child_outputs) == 0
            out = depth
        else
            out = combiner(child_outputs)
        end
    end

    bar1 = Recurse(production, aggregation, 1, Int, Int, Any)
    constraints = choicemap()
    constraints[(1, Val(:production)) => :branch] = true
    constraints[(2, Val(:production)) => :branch] = false
    tr, w = generate(bar1, (1, 1), constraints)
    @test get_retval(tr) == 2
    constraints[(2, Val(:production)) => :branch] = true
    tr, w = zip([generate(bar1, (1, 1), constraints) for i in 1:100]...)
    @test all(get_retval.(tr) .>= 2)

    @gen new_prod(d::Int) =
        @trace(production(d, 2))
    @gen new_aggr(d::Int, cvals::Vector) =
        @trace(aggregation(d, cvals, xs -> join(string.(xs))))

    bar2 = Recurse(new_prod, new_aggr, 2, Int, Int, Any)
    constraints = choicemap()
    constraints[(1, Val(:production)) => :branch] = true
    constraints[(Gen.get_child(1, 1, 2), Val(:production)) => :branch] = false
    constraints[(Gen.get_child(1, 2, 2), Val(:production)) => :branch] = false
    tr, w = generate(bar2, (1, 1), constraints)
    @test get_retval(tr) == "22"
    constraints[(Gen.get_child(1, 1, 2), Val(:production)) => :branch] = true
    constraints[(Gen.get_child(1, 2, 2), Val(:production)) => :branch] = true
    constraints[(Gen.get_child(2, 1, 2), Val(:production)) => :branch] = false
    constraints[(Gen.get_child(2, 2, 2), Val(:production)) => :branch] = false
    constraints[(Gen.get_child(3, 1, 2), Val(:production)) => :branch] = false
    constraints[(Gen.get_child(3, 2, 2), Val(:production)) => :branch] = false
    tr, w = generate(bar2, (1, 1), constraints)
    @test get_retval(tr) == "3333"
end


text:
#!/usr/bin/env julia

function main()
    f = readlines("input")

    # `hold` will contain the previous depth
    # `counter' the number of times the depth increases
    hold = 0
    counter = 0

    for i in range(1, length(f))
        val = parse(Int, f[i])

        # There's no depth before the first one
        if i > 1
            if val > hold
                counter += 1
            end
        end
        hold = val
    end

    print(counter)
end

main()

ext: jl


text:
function get_GLRK_ω_matrix(s)
    ω = zeros(s, s+1)
    g = TableauGauss(s)
    ω[s,s+1] = 1
    for i in 1:s-1
        for j in 1:s
            ω[i,j] = g.b[j] * g.c[j]^(i-1)
        end
    end
    return ω
end


text:
frameattrs = (
    width = 1,
    color = :black,
    visible = true,
    style = nothing
)

yaxisattrs = (
    tick = (
    # tick marks
        ticks   = WilkinsonTicks(; k_min = 5, k_ideal = 7, k_max = 15),
        autolimitmargin = 0.05f0,
        size    = 10f0,
        visible = true,
        color   = RGBA(Colors.colorant"lightgrey", 0.7),
        align   = 0f0,
        width   = 1f0,
        style   = nothing,

        # tick labels
        label = (
            size      = 20f0,
            formatter = Formatting.format,
            visible   = true,
            font      = "DejaVu Sans",
            color     = RGBf0(0, 0, 0),
            spacing   = 20f0,
            padding   = 5f0,
            rotation  = 0f0,
            align     = (:center, :top),
            position  = Left()::MakieLayout.Side

        ),
    ),
)

xaxisattrs = (
    tick = (
    # tick marks
        ticks   = WilkinsonTicks(; k_min = 5, k_ideal = 7, k_max = 12),
        autolimitmargin = 0.05f0,
        size    = 10f0,
        visible = true,
        color   = RGBA(Colors.colorant"lightgrey", 0.7),
        align   = 0f0,
        width   = 1f0,
        style   = nothing,

        # tick labels
        label = (
            size      = 20f0,
            formatter = Formatting.format,
            visible   = true,
            font      = "DejaVu Sans",
            color     = RGBf0(0, 0, 0),
            spacing   = 20f0,
            padding   = 5f0,
            rotation  = 0f0,
            align     = (:center, :top),
            position  = Bottom()

        ),
    ),
)

@recipe(GeoAxis, limits) do scene
    merge(scene.attributes,
    Theme(
        samples = 100,
        show_axis = false,

        frames = (
            top = frameattrs,
            left = frameattrs,
            right = frameattrs,
            bottom = frameattrs
        ),

        grid = (
            visible = true,
            color   = RGBf0(0, 0, 0),
            width   = 1f0,
            style   = nothing
        ),

        x = xaxisattrs,
        y = yaxisattrs,

        crs = (
            source = LonLat(),
            dest   = LonLat()
        )
    )
    )
end

convert_arguments(::Type{<: GeoAxis}, xmin::Real, xmax::Real, ymin::Real, ymax::Real) = (FRect2D(xmin, xmax, ymin, ymax),)

function convert_arguments(::Type{<: GeoAxis}, xs::Tuple, ys::Tuple)
    xmin, xmax = xs
    ymin, ymax = ys
    return (FRect2D(xmin, xmax, ymin, ymax),)
end

function convert_arguments(::Type{<: GeoAxis}, xs::AbstractVector{<: Number}, ys::AbstractVector{<: Number})
    xmin, xmax = extrema(xs)
    ymin, ymax = extrema(ys)
    return (FRect2D(xmin, xmax, ymin, ymax),)
end

# function AbstractPlotting.calculated_attributes!(plot::GeoAxis)
#     @extract plot (x, y, crs, limits)
#
#     lift(limits, crs.source, crs.dest, x.tick.label.size, y.tick.label.size) do limits, source, dest, xticklabelsize, yticklabelsize
#         newrect = transform(source, dest, [limits[LEFT] limits[TOP]; limits[RIGHT] limits[BOTTOM]])
#         textscale = maximum(diff.(eachcol(newrect)))
#
#         x.tick.label.textsize[] = xticklabelsize * textscale/100
#         y.tick.label.textsize[] = yticklabelsize * textscale/100
#
#     end
#
# end

function AbstractPlotting.plot!(plot::GeoAxis{T}) where T

    @extract plot (x, y)

    draw_frames!(plot)

    draw_ticks!(plot)

end

function draw_frames!(plot::GeoAxis{T}) where T

    @extract plot (frames, crs, samples)

    @extract frames (top, bottom, left, right)

    @extract crs (source, dest)

    # initialize frames
    topline = Node(Vector{Point2f0}())
    bottomline = Node(Vector{Point2f0}())
    leftline = Node(Vector{Point2f0}())
    rightline = Node(Vector{Point2f0}())

    # initialize the line vectors
    lift(plot.limits, source, dest, samples) do lims, source, dest, samples

        lonrange = LinRange(lims[LEFT], lims[RIGHT], samples)
        latrange = LinRange(lims[BOTTOM], lims[TOP], samples)

        topline[] = Point2f0.(transform.(source, dest, [Point2f0(lon, lims[TOP]) for lon in lonrange]))
        leftline[] = Point2f0.(transform.(source, dest, [Point2f0(lims[LEFT], lat) for lat in latrange]))
        rightline[] = Point2f0.(transform.(source, dest, [Point2f0(lims[RIGHT], lat) for lat in latrange]))
        bottomline[] = Point2f0.(transform.(source, dest, [Point2f0(lon, lims[BOTTOM]) for lon in lonrange]))

    end

    # plot the frames
    lines!.(plot, (top, bottom, left, right), (topline, bottomline, leftline, rightline))

end

function draw_ticks!(plot::GeoAxis)

    @extract plot (x, y)

    xtickvalues = Node{Vector{<: AbstractFloat}}(Vector{Float64}())
    ytickvalues = Node{Vector{<: AbstractFloat}}(Vector{Float64}())

    xlinevec = Node(Vector{Point2f0}())
    ylinevec = Node(Vector{Point2f0}())

    xtickannotations = Node(Vector{Tuple{String, Point2f0}}())
    ytickannotations = Node(Vector{Tuple{String, Point2f0}}())

    lift(x.tick.ticks, y.tick.ticks, x.tick.label.position, y.tick.label.position, plot.limits, plot.samples, plot.crs.source, plot.crs.dest, x.tick.label.size, y.tick.label.size) do xticks_struct, yticks_struct, xtickp, ytickp, limits, samples, source, dest, xticklabelsize, yticklabelsize

        xtickvalues[] = MakieLayout.compute_tick_values(xticks_struct, limits[LEFT], limits[RIGHT], 100f0)
        ytickvalues[] = MakieLayout.compute_tick_values(yticks_struct, limits[BOTTOM], limits[TOP], 100f0)

        xticklabels = MakieLayout.get_tick_labels(xticks_struct, xtickvalues[])
        yticklabels = MakieLayout.get_tick_labels(yticks_struct, ytickvalues[])

        # silently update the backend value without calling
        # these Observables' listener functions.
        xlinevec.val = Vector{Point2f0}()
        ylinevec.val = Vector{Point2f0}()

        for xtick in xtickvalues[]
            append!(xlinevec.val, transform.(source, dest, Point2f0.(xtick, LinRange(limits[BOTTOM], limits[TOP], samples))))
            push!(xlinevec.val, Point2f0(NaN))
        end
        for ytick in ytickvalues[]
            append!(ylinevec.val, transform.(source, dest, Point2f0.(LinRange(limits[LEFT], limits[RIGHT], samples), ytick)))
            push!(ylinevec.val, Point2f0(NaN))
        end

        # notify the observables that they have changed
        AbstractPlotting.notify!(xlinevec)
        AbstractPlotting.notify!(ylinevec)

        # now for the tick placement

        # first, we do the y ticks (latitude)
        ytickpositions, ytickstrings = (nothing, nothing)

        if typeof(ytickp) <: MakieLayout.Side

            xpos = limits[ytickp]

            ytickpositions = Point2f0.(xpos, ytickvalues[])

            ytickstrings = yticklabels

        elseif ytickp isa NTuple{2, <: MakieLayout.Side}

            xpos1 = limits[ytickp[1]]
            xpos2 = limits[ytickp[2]]

            ytickpositions = Point2f0.([(x, y) for x in (xpos1, xpos2), y in ytickvalues[]])

            ytickstrings = repeat(yticklabels, 2)

        else
            @warn "Unsupported tick position format given!"
        end

        (isnothing(ytickstrings) || isnothing(ytickpositions)) || (ytickannotations[] = to2tuple.(ytickstrings, transform.(source, dest, ytickpositions)))

        xtickpositions, xtickstrings = (nothing, nothing)

        if typeof(xtickp) <: MakieLayout.Side

            ypos = limits[xtickp]

            xtickpositions = Point2f0.(xtickvalues[], ypos)

            xtickstrings = xticklabels

        elseif xtickp isa NTuple{2, <: MakieLayout.Side}

            ypos1 = limits[ytickp[1]]
            ypos2 = limits[ytickp[2]]

            xtickpositions = Point2f0.([(x, y) for y in (ypos1, ypos2), x in xtickvalues[]])

            xtickstrings = repeat(xticklabels, 2)

        else
            @warn "Unsupported tick position format given!" xtickp
        end

        (isnothing(xtickstrings) || isnothing(xtickpositions)) || (xtickannotations[] = to2tuple.(xtickstrings, transform.(source, dest, xtickpositions)))

    end

    # plot the damn thing

    # x ticks
    lines!(
        plot,
        xlinevec;
        visible = x.tick.visible,
        color   = x.tick.color,
        align   = x.tick.align,
        linewidth = x.tick.width,
        linestyle = x.tick.style,
    )
    # y ticks
    lines!(
        plot,
        ylinevec;
        visible = y.tick.visible,
        color   = y.tick.color,
        align   = y.tick.align,
        linewidth = y.tick.width,
        linestyle = y.tick.style,
    )

    # annotations

    annotations!(
        plot,
        x.tick.label,
        xtickannotations;
    )

    annotations!(
        plot,
        y.tick.label,
        ytickannotations;
    )

end


text:
# Local ZigZag
using SparseArrays
d = 8
S = 1.3I + 0.5sprandn(d, d, 0.1)
const Γ = S*S'

∇ϕ(x, i, Γ) = ZigZagBoomerang.idot(Γ, i, x) # sparse computation
∇ϕmoving(t, x, θ, i, t′, F, Γ) = ZigZagBoomerang.idot_moving!(Γ, i, t, x, θ, t′, F) # sparse computation

@testset "ZigZag" begin

    t0 = 0.0
    x0 = rand(d)
    θ0 = rand([-1.0, 1.0], d)

    c = .6*[norm(Γ[:, i], 2) for i in 1:d]

    Z = ZigZag(0.9Γ, x0*0)
    T = 1000.0

    trace, _, acc = @time pdmp(∇ϕ, t0, x0, θ0, T, c, Z, Γ)
    @show acc[1]/acc[2]
    dt = 0.5
    ts, xs = sep(collect(discretize(trace, dt)))

    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 2.5/sqrt(T)
end

@testset "SZigZag" begin

    t0 = 0.0
    x0 = rand(d)
    θ0 = rand([-1.0,-0.5,0.5,1.0], d)

    c = .7*[norm(Γ[:, i], 2) for i in 1:d]

    Z = ZigZag(0.9Γ, x0*0)
    T = 1000.0

    trace, _, acc = @time spdmp(∇ϕ, t0, x0, θ0, T, c, Z, Γ)
    @show acc[1]/acc[2]
    dt = 0.5
    ts, xs = sep(collect(discretize(trace, dt)))
    @testset "subtrace" begin
        J = 1:2:d
        ts2, xs2 = sep(collect(discretize(subtrace(trace, J), dt)))
        @test ts2 ≈ ts[eachindex(ts2)]
        @test (xs2) ≈ (getindex.(xs[eachindex(ts2)], Ref(J)))
    end
    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 2.5/sqrt(T)
end

@testset "SZigZagSelfMoving" begin

    t0 = 0.0
    x0 = rand(d)
    θ0 = rand([-1.0,-0.5,0.5,1.0], d)

    c = .7*[norm(Γ[:, i], 2) for i in 1:d]

    Z = ZigZag(0.9Γ, x0*0)
    T = 1000.0

    trace, _, acc = @time spdmp(∇ϕmoving, t0, x0, θ0, T, c, Z, SelfMoving(), Γ)
    @show acc[1]/acc[2]
    dt = 0.5
    ts, xs = sep(collect(discretize(trace, dt)))

    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 2.5/sqrt(T)
end

@testset "FactBoomerang" begin

    t0 = 0.0
    x0 = 0.2rand(d)

    #Γ0 = sparse(I, d, d)
    c = [norm(Γ[:, i], 2) for i in 1:d]
    Γ0 = copy(Γ)
    for i in 1:d
        #Γ0[d,d] = 1
    end
    Z = FactBoomerang(0.85Γ0, x0*0, 0.3)
    θ0 = sqrt(Diagonal(Z.Γ))\randn(d)

    T = 3000.0

    trace, _, acc = @time pdmp(∇ϕ, t0, x0, θ0, T, c, Z, Z.Γ)
    @show acc[1]/acc[2]
    dt = 0.5
    ts, xs = sep(collect(discretize(trace, dt)))

    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 4/sqrt(T)
end

@testset "SFactBoomerang" begin

    t0 = 0.0
    x0 = rand(d)

    #Γ0 = sparse(I, d, d)
    c = [norm(Γ[:, i], 2) for i in 1:d]
    Γ0 = copy(Γ)
    for i in 1:d
        #Γ0[d,d] = 1
    end
    Z = FactBoomerang(1.2Γ0, x0*0, 0.3)
    θ0 = sqrt(Diagonal(Z.Γ))\randn(d)
    T = 3000.0

    trace, _, acc = @time spdmp(∇ϕ, t0, x0, θ0, T, c, Z, Z.Γ)
    @show acc[1]/acc[2]
    dt = 0.5
    ts, xs = sep(collect(discretize(trace, dt)))

    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 4/sqrt(T)
end

@testset "Boomerang" begin
    t0 = 0.0
    θ0 = randn(d)
    x0 = randn(d)
    c = 6.0
    Γ0 = copy(Γ)
    B = Boomerang(Γ0, x0*0, 0.5)
    ∇ϕ!(y, x) = mul!(y, Γ, x)
    T = 3000.0
    trace, _, acc = @time pdmp(∇ϕ!, t0, x0, θ0, T, c, B)
    @show acc[1]/acc[2]
    dt = 0.1
    ts, xs = sep(collect(discretize(trace, dt)))
    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ0)))) < 2.5/sqrt(T)
end

@testset "Bouncy Particle Sampler" begin
    t0 = 0.0
    θ0 = randn(d)
    x0 = randn(d)
    c = 1.1
    Γ0 = copy(Γ)
    B = BouncyParticle(Γ0, x0*0, 0.5)
    ∇ϕ!(y, x) = mul!(y, Γ, x)
    T = 300.0
    trace, _, acc, more = @time pdmp(∇ϕ!, t0, x0, θ0, T, c, B, progress=true)
    @show more
    @show acc[1]/acc[2]
    dt = 0.1
    ts, xs = sep(collect(discretize(trace, dt)))
    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ0)))) < 2/sqrt(T)
end

@testset "ZigZag (independent)" begin

    t0 = 0.0
    x0 = rand(d)
    θ0 = rand([-1.0, 1.0], d)

    c = 10.0*[norm(Γ[:, i], 2) for i in 1:d]
    Γ0 = sparse(I, d, d)
    Z = ZigZag(Γ0, x0*0)

    T = 1000.0

    trace, _, acc = @time pdmp(∇ϕ, t0, x0, θ0, T, c, Z, Γ)
    @show acc[1]/acc[2]
    dt = 0.5
    ts, xs = sep(collect(discretize(trace, dt)))

    @test mean(abs.(mean(xs))) < 2/sqrt(T)
    @test mean(abs.(cov(xs) - inv(Matrix(Γ)))) < 2.5/sqrt(T)
end

@testset "FactBoomerang1" begin
    ϕ(x) = [cos(π*x[1]) + x[1]^2/2] # not needed
    # gradient of ϕ(x)
    ∇ϕ(x) = [-π*sin(π*x[1]) + x[1]]
    ∇ϕ(x, i) = -π*sin(π*x[1]) + x[1] # (REPLACE IT WITH AUTOMATIC DIFFERENTIATION)
    c = [3.5π]
    λref = 1.5
    n = 1
    x0 = randn(n)
    θ0 = randn(n)
    t0 = 0.0
    T = 10100.0
    Γ = sparse(Matrix(1.0I, n, n))
    B = FactBoomerang(Γ, x0*0, λref)
    trace, _,  acc = pdmp(∇ϕ, t0, x0, θ0, T, c, B)
    @show acc[1]/acc[2]
    m = mean(last.(collect(trace)))
    dt = 0.1
    ts, xs = sep(collect(discretize(trace, dt)))
    @test mean(xs)[1] < 2.5/sqrt(T)
end


text:
using LinearAlgebra

"""
    initbox(x::AbstractVector{T}) where {T}

Create a bounding box and initialize with a single point.
"""
function initbox(x::AbstractVector{T}) where {T}
    sdim = length(x)
    box = fill(zero(T), 2*sdim)
    for i in 1:sdim
        box[2*i-1] = box[2*i] = x[i];
    end
    return box
end

function initbox(x::AbstractMatrix{T}) where {T}
	box = initbox(x[1, :])
	return updatebox!(box, x)
end

"""
    updatebox!(box::AbstractVector{T}, x::AbstractVector{T}) where {T}

Update a box with another location, or create a new box.

If the  `box` does not have  the correct dimensions,  it is correctly sized.

`box` = bounding box
    for 1-D `box=[minx,maxx]`, or
    for 2-D `box=[minx,maxx,miny,maxy]`, or
    for 3-D `box=[minx,maxx,miny,maxy,minz,maxz]`
`x` = vector defining a point.  The `box` is expanded to include the
    supplied location `x`.
"""
function updatebox!(box::AbstractVector{T}, x::AbstractVector{T}) where {T}
    sdim = length(x)
    for i in 1:sdim
    	box[2*i-1] = min(box[2*i-1],x[i]);
    	box[2*i]   = max(box[2*i],x[i]);
    end
    return box
end

function updatebox!(box::AbstractVector{T}, x::AbstractMatrix{T}) where {T}
    for i in 1:size(x, 1)
    	updatebox!(box, x[i, :])
    end
    return box
end

"""
    boundingbox(x::AbstractArray{T}) where {T}

Compute the bounding box of the points in `x`.

`x` = holds points, one per row.

Returns `box` = bounding box
    for 1-D `box=[minx,maxx]`, or
    for 2-D `box=[minx,maxx,miny,maxy]`, or
    for 3-D `box=[minx,maxx,miny,maxy,minz,maxz]`
"""
function boundingbox(x::AbstractArray{T}) where {T}
	box = initbox(x[1, :])
	for i in 2:size(x, 1)
		updatebox!(box, x)
	end
    return box
end

"""
    inflatebox!(box::AbstractVector{T}, inflatevalue::T) where {T}

Inflate the box by the value supplied.
"""
function inflatebox!(box::AbstractVector{T}, inflatevalue::T) where {T}
    abox = deepcopy(box)
    sdim = Int(length(box)/2);
    for i=1:sdim
        box[2*i-1] = min(abox[2*i-1],abox[2*i]) - inflatevalue;
        box[2*i]   = max(abox[2*i-1],abox[2*i]) + inflatevalue;
    end
    return box
end

"""
    inbox(box::AbstractVector{T}, x::AbstractVector{T}) where {T}

Is the given location inside the box?

- `box` = vector entries arranged as [minx,maxx,miny,maxy,minz,maxz] (or
  adjusted in an obvious way for lower space dimension).

Note: point on the boundary of the box is counted as being inside.
"""
function inbox(box::AbstractVector{T}, x::AbstractVector{T}) where {T}
    inrange(rangelo,rangehi,r) = ((r>=rangelo) && (r<=rangehi));
    sdim=length(x);
    @assert 2*sdim == length(box)
    if !inrange(box[1], box[2], x[1])
        return false # short-circuit
    end
    for i=2:sdim
        if !inrange(box[2*i-1], box[2*i], x[i])
            return false # short-circuit
        end
    end
    return true
end

function inbox(box::AbstractVector{T}, x::AbstractArray{T}) where {T}
    return inbox(box, vec(x))
end

"""
    boxesoverlap(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}

Do the given boxes overlap?
"""
function boxesoverlap(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}
    dim=Int(length(box1)/2);
    @assert 2*dim == length(box2) "Mismatched boxes"
    for i=1:dim
        if box1[2*i-1]>box2[2*i]
            return false;
        end
        if box1[2*i]<box2[2*i-1]
            return false;
        end
    end
    return  true;
end

"""
    intersectboxes(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}

Compute the intersection of two boxes.

The function returns an empty box (length(b) == 0) if the intersection is
empty; otherwise a box is returned.
"""
function intersectboxes(box1::AbstractVector{T}, box2::AbstractVector{T}) where {T}
    @assert length(box1) == length(box2) "Mismatched boxes"
    b = copy(box1)
    dim=Int(length(box1)/2);
    @assert 2*dim == length(box2) "Wrong box data"
    for i=1:dim
        lb = max(box1[2*i-1], box2[2*i-1])
        ub = min(box1[2*i], box2[2*i])
        if (ub <= lb) # intersection is empty
            return eltype(box1)[] # box of length zero signifies empty intersection
        end
        b[2*i-1] = lb
        b[2*i]   = ub
    end
    return  b;
end


text:
"""
    decode_single(input::AbstractVector{UInt8}; offset=zero(UInt))
Decode a single value of ULEB128 from `input`, starting with `offset + 1`-th byte.
"""
@inline function decode_single(input::Ptr{UInt8};
    offset::UInt=zero(UInt)
    )
    output = zero(UInt32)
    shift = 0x00
    for j in 1:5
        byte = unsafe_load(input, offset + j)#input[offset + j]
        output |= (UInt32(byte & 0x7f) << shift)
        if (byte & 0x80 == 0)
            offset += j
            break
        end
        @assert j < 5 "incorrect encoding of a ULEB128 number"
        shift += 0x07
    end
    return output, offset
end

"""
    decode_multiple!(output::AbstractVector{UInt32}, input::AbstractVector{UInt8};
    offset=zero(UInt))

Decode `count` values of ULEB128 from `input`, starting with `offset + 1`-th byte.
"""
function decode_multiple!(output::Ptr{UInt32},
    input::Ptr{UInt8};
    count::Integer = length(output),
    offset::UInt = zero(UInt)
    )
    for i in 1:count
        o = zero(UInt32)
        shift = 0x00
        for j in 1:5
            byte = unsafe_load(input, offset + j) #input[offset + j]
            o |= (UInt32(byte & 0x7f) << shift)
            if (byte & 0x80 == 0)
                offset += j
                break
            end
            @assert j < 5 "incorrect encoding of a ULEB128 number"
            shift += 0x07
        end
        unsafe_store!(output, o, i) # output[i] = o
    end
    return output, offset
end

function size_n(input::AbstractVector{UInt8}, n::Integer, offset::UInt)
    r = zero(UInt32)
    for i in 1:n
        for j in 1:5
            byte = input[offset + j]
            if (byte & 0x80 == 0)
                r += j
                offset += j
                break
            end
            @assert j < 5 "incorrect encoding of a ULEB128 number"
        end
    end
    r
end


text:
# Autogenerated wrapper script for Calcium_jll for aarch64-apple-darwin
export libcalcium

using FLINT_jll
using Arb_jll
using Antic_jll
using GMP_jll
using MPFR_jll
JLLWrappers.@generate_wrapper_header("Calcium")
JLLWrappers.@declare_library_product(libcalcium, "@rpath/libcalcium-0.4.1.dylib")
function __init__()
    JLLWrappers.@generate_init_header(FLINT_jll, Arb_jll, Antic_jll, GMP_jll, MPFR_jll)
    JLLWrappers.@init_library_product(
        libcalcium,
        "lib/libcalcium.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# Groups of functions with a shared signature defined for vectors
# https://stat.ethz.ch/R-manual/R-patched/library/methods/html/S4groupGeneric.html
# Translated to julia, adding scalar versions of operators
# Arith
# "+", "-", "*", "^", "%%", "%/%", "/"
#const arith_group = [:(+), :(-), :(*), :(/), :(^), :(.+), :(.-), :(.*), :(./), :(.^), :(div), :(mod), :(fld), :(rem)]
const arith_group = [:+, :-, :.+, :.-, :.*, :./, :.^, :div, :mod, :fld, :rem] # Just scalar arith for vectors

# Compare
# "=="), ">"), "<"), "!="), "<="), ">="
const compare_group = [:.==, :.>, :.<, :.!=, :.<=, :.>=]

# Logic
# "&"), "|".
const logic_group = [:&, :|]

# Ops
# "Arith", "Compare", "Logic"
const ops_group = vcat( arith_group, compare_group, logic_group )

# Math
# "abs", "sign", "sqrt", "ceiling", "floor", "trunc", "cummax", "cummin", "cumprod", "cumsum", "log", "log10", "log2", "log1p", "acos", "acosh", "asin", "asinh", "atan", "atanh", "exp", "expm1", "cos", "cosh", "sin", "sinh", "tan", "tanh", "gamma", "lgamma", "digamma", "trigamma"
const math_group = [:abs, :sign, :sqrt, :ceil, :floor, :trunc, :cummax, :cummin, :cumprod, :cumsum, :log, :log10, :log2, :log1p, :acos, :acosh, :asin, :asinh, :atan, :atanh, :exp, :expm1, :cos, :cosh, :sin, :sinh, :tan, :tanh, :gamma, :lgamma, :digamma, :trigamma]

# Math2
# "round", "signif"
# leaving out for now

# Summary
# "max", "min", "range", "prod", "sum", "any", "all"
const summary_group =[:maximum, :minimum, :range, :prod, :sum, :any, :all, :eltype, :unique, :minimum, :maximum, :extrema, :first, :last, :maxabs, :minabs, :any, :all] # sum and prod special

# Complex
# "Arg", "Conj", "Im", "Mod", "Re"
# leaving out for now

### Operators, methods that take two arguments and return a modified RLEVector
function ^(x::RLEVector,y::Integer) # Necessary to prevent an ambiguity warning
  rv = ^(x.runvalues,y)
  RLEVector(rv,x.runends)
end

.^(x::Base.Irrational{:e}, y::RLEVectors.RLEVector) = .^(x,y) # Ambig fix
for op in ops_group
    @eval begin
        # Rle, Rle
        function ($op)(x::RLEVector, y::RLEVector)
            (runends, runvalues_x, runvalues_y) = disjoin(x, y)
            runvalues = $(op)(runvalues_x, runvalues_y)
            RLEVector{eltype(runvalues), eltype(runends)}(runvalues, runends)
        end
        # Rle, Number
        ($op){T<:Integer}(x::RLEVector{Bool,T},y::Bool) = RLEVector( ($op)(x.runvalues,y), x.runends ) # Ambig fix
        ($op)(x::RLEVector,y::Number) = RLEVector( ($op)(x.runvalues,y), x.runends )
        # Number, Rle
        ($op){T<:Integer}(y::Bool, x::RLEVector{Bool,T}) = RLEVector( ($op)(y,x.runvalues), x.runends ) # Ambig fix
        ($op)(y::Number, x::RLEVector) = RLEVector( ($op)(y,x.runvalues), x.runends )
    end
end

## Methods that delegate to the runvalues and return an RLEVector
## Methods that take one argument, an RLEVector, and delegate to rle.runvalues and return an RLEVector
for op in math_group
  @eval ($op)(x::RLEVector) = RLEVector( ($op)(x.runvalues), x.runends )
end

## Methods that take one argument, an RLEVector, and delegate to rle.runvalues and return something other than an RLEVector
for op in setdiff(summary_group,[:sum,:prod])
  @eval ($op)(x::RLEVector) = ($op)(x.runvalues)
end

## Methods that take two arguments, delegate to rle.runvalues and return something other than an RLEVector
in{T1,T2<:Integer}(y::T1, x::RLEVector{T1,T2}) = in(y, x.runvalues)

# Defaulting to fun(itr) for some things
for op in [:findmin, :findmax]
  @eval begin
    function ($op)(x::RLEVector)
      m = ($op)(x.runvalues)
      (m[1], rfirst(x,m[2]))
    end
  end
end

indexin(x::RLEVector,y::RLEVector) = RLEVector( indexin(x.runvalues,y), x.runends)
indexin(x::RLEVector,y::AbstractVector) = RLEVector( indexin(x.runvalues,y), x.runends )
indexin(x::AbstractVector,y::RLEVector) = Int[ i == 0 ? 0 : y.runends[i] for i in indexin(x,y.runvalues) ]

findin(x,y::RLEVector) = findin(x,y.runvalues)
function findin(x::RLEVector,y::RLEVector)
  runs = findin(x.runvalues,y.runvalues)
  re = x.runends
  vcat( [ rfirst(x,i):re[i] for i in runs ]... )  # hashing in above findin takes the vast majority of the time, don't sweat the time here
end

function findin(x::RLEVector,y::UnitRange) # ambig fix
  runs = findin(x.runvalues,y)
  re = x.runends
  vcat( [ rfirst(x,i):re[i] for i in runs ]... ) # hashing in above findin takes the vast majority of the time, don't sweat the time here
end

function findin(x::RLEVector,y)
  runs = findin(x.runvalues,y)
  re = x.runends
  vcat( [ rfirst(x,i):re[i] for i in runs ]... )  # hashing in above findin takes the vast majority of the time, don't sweat the time here
end

function median(x::RLEVector)
  len = length(x)
  len < 2 && return(middle(x.runvalues[1]))
  sorted = sort(x)
  mid = fld(len,2)
  mid_run = ind2run(sorted,mid)
  if mod(len,2) == 0 && mid == sorted.runends[mid_run] # even numbered and at end of run, avg with next value
    median = middle(sorted.runvalues[mid_run], sorted.runvalues[mid_run+1])
  else
    median = middle(sorted.runvalues[mid_run])
  end
  return(median)
end

function sum{T1,T2}(x::RLEVector{T1,T2})
  rval = zero(T1)
   @simd for i in 1:nrun(x)
    @inbounds rval = rval + (x.runvalues[i] * rwidth(x, i))
  end
  return(rval)
end

function mean{T1,T2}(x::RLEVector{T1,T2})
  rval = sum(x) / length(x)
end


text:
export qr, qr_back, copyltu!, lq, lq_back

"""
		qr(A) -> Tuple{AbstractMatrix, AbstractMatrix}

private QR method, call LinearAlgebra.qr to achieve forward calculation, while
return Tuple type
"""
function qr(A)
	res = LinearAlgebra.qr(A)
	Matrix(res.Q), res.R
end

"""
		qr(A, pivot) -> Tuple{AbstractMatrix, AbstractMatrix, AbstractVector}
"""
function qr(A::AbstractMatrix, pivot::Val{true})
    res = LinearAlgebra.qr(A, pivot)
    Q, R, P = Matrix(res.Q), res.R, res.P
end

"""
		lq(A) -> Tuple{AbstractMatrix, AbstractMatrix}
"""
function lq(A)
		res = LinearAlgebra.lq(A)
    res.L, Matrix(res.Q)
end

trtrs!(c1::Char, c2::Char, c3::Char, r::AbstractMatrix, b::AbstractVecOrMat) = LinearAlgebra.LAPACK.trtrs!(c1, c2, c3, r, b)

"""
    copyltu!(A::AbstractMatrix) -> AbstractMatrix

copy the lower triangular to upper triangular.
"""
function copyltu!(A::AbstractMatrix)
    m, n = size(A)
    for i=1:m
        A[i,i] = real(A[i,i])
        for j=i+1:n
            @inbounds A[i,j] = conj(A[j,i])
        end
    end
    A
end

"""
    qr_back_fullrank(q, r, dq, dr) -> Matrix

backward for QR decomposition, for input matrix (in forward pass) with M > N.

References:
    Seeger, M., Hetzel, A., Dai, Z., Meissner, E., & Lawrence, N. D. (2018). Auto-Differentiating Linear Algebra.
"""
@generated function qr_back_fullrank(q, r, dq, dr)
    dqnot0 = !(dq <: Nothing)
    drnot0 = !(dr <: Nothing)
    (!dqnot0 && !drnot0) && return :(nothing)
    ex = drnot0 && dqnot0 ? :(r*dr' - dq'*q) : (dqnot0 ? :(-dq'*q) : :(r*dr'))
    :(b = $(dqnot0 ? :(dq) : :(ZeroAdder())) + q*copyltu!($ex);
    trtrs!('U', 'N', 'N', r, do_adjoint(b))')
end

do_adjoint(A::Matrix) = Matrix(A')

"""
    qr_back(A, q, r, dq, dr) -> Matrix

backward for QR decomposition, for an arbituary shaped input matrix.

References:
    Seeger, M., Hetzel, A., Dai, Z., Meissner, E., & Lawrence, N. D. (2018). Auto-Differentiating Linear Algebra.
    Differentiable Programming Tensor Networks, Hai-Jun Liao, Jin-Guo Liu, Lei Wang, Tao Xiang
"""
@generated function qr_back(A, q, r, dq, dr)
    dqnot0 = !(dq <: Nothing)
    drnot0 = !(dr <: Nothing)
    (!dqnot0 && !drnot0) && return :(nothing)
    ex = quote
        size(r, 1) == size(r, 2) && return qr_back_fullrank(q, r, dq ,dr)
        M, N = size(r)
        B = view(A,:,M+1:N)
        U = view(r,:,1:M)
        D = view(r,:,M+1:N)
        $(if drnot0
            :(dD = view(dr,:,M+1:N);
            da = qr_back_fullrank(q, U, $(dqnot0 ? :(dq+B*dD') : :(B*dD')), view(dr,:,1:M));
            db = q*dD)
        else
            :(da = qr_back_fullrank(q, U, dq, nothing);
            db = zero(B))
        end)
        hcat(da, db)
    end
end

"""
    lq_back(A, l, q, dl, dq) -> Matrix

backward for LQ decomposition, for an arbituary shaped input matrix.

References:
    Seeger, M., Hetzel, A., Dai, Z., Meissner, E., & Lawrence, N. D. (2018). Auto-Differentiating Linear Algebra.
    Differentiable Programming Tensor Networks, Hai-Jun Liao, Jin-Guo Liu, Lei Wang, Tao Xiang
"""
function lq_back_fullrank(L, Q, dL, dQ)
    M = ZeroAdder()
    dL === nothing || (M += L'*dL)
    dQ === nothing || (M -= dQ*Q')
    C = copyltu!(M)*Q
    if dQ !== nothing
        C += dQ
    end
    #inv(L)' * C
    trtrs!('L', 'C', 'N', L, C)
end

"""
    lq_back(A, L, Q, dL, dQ) -> Matrix

backward for QR decomposition, for an arbituary shaped input matrix.

References:
    Seeger, M., Hetzel, A., Dai, Z., Meissner, E., & Lawrence, N. D. (2018). Auto-Differentiating Linear Algebra.
    HaiJun's paper.
"""
@generated function lq_back(A, L, Q, dL, dQ)
    dunot0 = !(dQ <: Nothing)
    dlnot0 = !(dL <: Nothing)
    (!dunot0 && !dlnot0) && return :(nothing)
    ex = quote
        N, M = size(L)
        M == N && return lq_back_fullrank(L, Q, dL, dQ)
        B = view(A,M+1:N,:)
        U = view(L,1:M,:)
        D = view(L,M+1:N,:)
        $(if dlnot0
            :(dD = view(dL,M+1:N,:);
            da = lq_back_fullrank(U, Q, view(dL,1:M,:), $(dunot0 ? :(dQ+dD'*B) : :(dD'*B)));
            db = dD*Q)
        else
            :(da = lq_back_fullrank(U, Q, nothing, dQ);
            db = zero(B))
        end)
        vcat(da, db)
    end
end


text:
using Plots,LaTeXStrings
using LinearAlgebra,FundamentalsNumericalComputation

z = [-1,3];
v = z - norm(z)*[1,0];
P = I - 2*(v*v')/dot(v,v)
Pz = P*z;

gr()
default(legend=nothing,linewidth=1)
plt = plot(aspect_ratio=1,showaxis=false,grid=false,xaxis=:none,yaxis=false)

plot!([0,z[1]],[0,z[2]],aspect_ratio=1)
plot!([Pz[1],z[1]],[Pz[2],z[2]])
plot!([0,Pz[1]],[0,Pz[2]])

plot!([-2,2]*v[2],[2 ,-2]*v[1],l=(:dash,:black))

scatter!([0],[0],m=(5,:black,:black))
xlims!(-1,4)
ylims!(-1, 3)

annotate!(-0.68,1.5,text(L"\mathbf{z}"))
annotate!(2,1,text(L"\mathbf{v}"))
annotate!(1.5,-0.3,text(L"\|\mathbf{z}\|\,\mathbf{e}_1"))
annotate!(-0.2,0,text(L"0"))
annotate!(1.6,1.9,text(L"S"))

savefig(plt,"hhreflect.svg")

max_stars_repo_name: tobydriscoll/fnc

max_issues_repo_name: tobydriscoll/fnc

max_forks_repo_name: tobydriscoll/fnc


text:
using Base.Test
using RandomFerns

features = reshape([0:0.01:10],1001,1);
labels = vec(sin(features) + (randn(size(features))/2));

println("\n ### k fold CV - Random Fern Regression ###")
r2 = kfoldCV_fern(labels, features, 10, 5)
@test mean(r2) > 0.3

println("\n ### k fold CV - Random Fern Ensemble Regression ###")
r2 = kfoldCV_fernEnsemble(labels, features, 10, 5, 100)
@test mean(r2) > 0.5


text:
##
using Shapefile
path = "/p/mnt/data/raw-data/megaregions/us_ttw_v3_US_only_epsg5070v2.shp"
table = Shapefile.Table(path)

## get the nodes and edges
tracts = unique(vcat(table.Ofips, table.Dfips))
tract2ind = Dict(tracts .=> 1:length(tracts))

src = map(i->tract2ind[i],table.Ofips)
dst = map(i->tract2ind[i],table.Dfips)

## commutes.net has the edge
# 71785 4122 1.000000
# 71785 8371 0.285714
# which map to
# 4122 "48323950500"
# 8371 "48507950301"
# 71785 "48507950200"
# we have
#=
julia> tract2ind["48507950200"]
69521

julia> tract2ind["48507950301"]
41675

julia> tract2ind["48323950500"]
6193

Which maps to the edge weight ESTDIVMOE
(exept, not... it's Flow./Moe)
=#

## get the coordinates
# lots of issues in getting this stuff, but they were all worked out
# and the following code s
function getxy(table)
  tract2xy = Dict{String,Tuple{Float64,Float64}}()
  Ofips = table.Ofips
  Dfips = table.Dfips
  for (i,geom) in enumerate(Shapefile.shapes(table))
    fsrc::String = Ofips[i]
    fdst::String = Dfips[i]
    @assert(length(geom.points) == 2)
    srcxy = (geom.points[1].x,geom.points[1].y)
    dstxy = (geom.points[2].x,geom.points[2].y)

    if haskey(tract2xy, fsrc)
      @assert srcxy == tract2xy[fsrc]
    else
      tract2xy[fsrc] = srcxy
    end

    if haskey(tract2xy, fdst)
      @assert dstxy == tract2xy[fdst]
    else
      tract2xy[fdst] = dstxy
    end
  end
  return tract2xy
end
tract2xy= getxy(table)

##
tracts_xy = map(i->tract2xy[i],tracts)
xy = [map(first, tracts_xy) map(i->i[2], tracts_xy)]
##
using SparseArrays
A = sparse(src, dst, table.Flow./table.Moe, length(tracts), length(tracts))
##
using MatrixNetworks
is_connected(A) # it isn't... we'll output largest cc below
## From https://github.com/dgleich/DGFun.jl/blob/master/src/io.jl
using DelimitedFiles
using SparseArrays
function writeSMAT(filename::AbstractString, A::SparseMatrixCSC{T,Int}; values::Bool=true) where T
    open(filename, "w") do outfile
        write(outfile, join((size(A,1), size(A,2), nnz(A)), " "), "\n")

        rows = rowvals(A)
        vals = nonzeros(A)
        m, n = size(A)
        for j = 1:n
           for nzi in nzrange(A, j)
              row = rows[nzi]
              val = vals[nzi]
              if values
                write(outfile, join((row-1, j-1, val), " "), "\n")
              else
                write(outfile, join((row-1, j-1, 1), " "), "\n")
              end
           end
        end
    end
end
##
A32 = Float32.(A)
writeSMAT("commutes-all.smat", A32)
writedlm("commutes-all.xy.epsg5070", xy)
writedlm("commutes-all.labels", tracts)
##

## Let's also get the lat-long
# curl https://www2.census.gov/geo/docs/reference/cenpop2010/tract/CenPop2010_Mean_TR.txt -O
using CSV
ctracts = CSV.read("CenPop2010_Mean_TR.txt")
##
using Printf
cids = map(i->@sprintf("%02i",i), ctracts.STATEFP) .* map(i->@sprintf("%03i",i), ctracts.COUNTYFP) .* map(i->@sprintf("%06i",i), ctracts.TRACTCE)
##
cids2row = Dict(cids .=> 1:length(cids))
all(map(i->haskey(cids2row,i), tracts))
##
xyll = [ctracts.LONGITUDE[map(i->cids2row[i], tracts)] ctracts.LATITUDE[map(i->cids2row[i], tracts)]]
writedlm("commutes-all.xy", xyll)

## This is no longer used...
## Let's also output the connected network

#=
Acc,p = largest_component(A)
writeSMAT("commutes-scc.smat", Acc)
writedlm("commutes-scc.xy.epsg5070", xy[p,:])
writedlm("commutes-scc.labels", tracts[p])
writedlm("commutes-scc.xy", xyll[p,:])
=#

max_stars_repo_name: empet/more-graphs

max_issues_repo_name: empet/more-graphs

max_forks_repo_name: empet/more-graphs


text:
pomdp = BabyPOMDP()

bu = DiscreteUpdater(pomdp)
b0 = initialize_belief(bu, initial_state_distribution(pomdp))

# these values were gotten from FIB.jl
alphas = [-29.4557 -36.5093; -19.4557 -16.0629]
policy = AlphaVectorPolicy(pomdp, alphas)

# initial belief is 100% confidence in baby not being hungry
@test isapprox(value(policy, b0), -16.0629)
@test isapprox(value(policy, [0.0,1.0]), -16.0629)

# because baby isn't hungry, policy should not feed (return false)
@test action(policy, b0) == false

# try pushing new vector
push!(policy, [0.0,0.0], true)

@test value(policy, b0) == 0.0
@test action(policy, b0) == true


text:
module utils
if VERSION >= v"0.7.0"
    using Statistics
else
    nothing
end

include("base.jl")
export flatten, unzip
export slicedimview, nextpow2, allequal, offdiagind, offdiag
export findnmax, findnmin
export rotatexy

include("old.jl")

include("misc.jl")
export from_choose
export nanmaximum, nanminimum, nan2zero!, nanmap, nanmapslices
export twodresample
export mmtoin, intomm
export wavedec2d

include("plots.jl")
export significance
export sslice

include("rsync.jl")

include("statsbase.jl")
export zscore_base, zscore_base!
export _bootstrap, bootstrap_confint

include("sorting.jl")
export bubble_sort!, bubble_sort_perm

include("images.jl")
export nanfilt, nanfilt!, imfilter_gaussian
end


text:
using ElectronDisplay
ElectronDisplay.CONFIG.showable = showable
ElectronDisplay.CONFIG.single_window = true
ElectronDisplay.CONFIG.focus = false
using ImageMagick, FileIO
using WGLMakie, Makie, Test
using ReferenceTests
using ReferenceTests: database_filtered

excludes = Set([
    "Streamplot animation",
    "Transforming lines",
    "image scatter",
    "Line GIF",
    "surface + contour3d",
    # Hm weird, looks like some internal JSServe error missing an Observable:
    "Errorbars x y low high",
    "Rangebars x y low high",
    # These are a bit sad, since it's just missing interpolations
    "FEM mesh 2D",
    "FEM polygon 2D",
    # missing transparency & image
    "Wireframe of a Surface",
    "Image on Surface Sphere",
    "Surface with image",
    # Marker size seems wrong in some occasions:
    "Hbox",
    "UnicodeMarker",
    # Not sure, looks pretty similar to me! Maybe blend mode?
    "Test heatmap + image overlap",
    "Stars",
    "heatmaps & surface",
    "OldAxis + Surface"
])
excludes2 = Set(["short_tests_83", "short_tests_78", "short_tests_40", "short_tests_13", "short_tests_5", "short_tests_41"])
database = database_filtered(excludes, excludes2)

recorded = joinpath(@__DIR__, "recorded")
rm(recorded; force=true, recursive=true); mkdir(recorded)
ReferenceTests.record_tests(database; recording_dir=recorded)
ReferenceTests.reference_tests(recorded; difference=0.06)


text:
# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>
#
# This script can be used to check the
# output of the acoustic model and its decoding
# using two random utterances taken from the test set

using HMMGradients, Flux, Zygote
using Random, Statistics, LinearAlgebra
using FiniteStateTransducers
using DSP, MFCC
using BSON, JLD2, LibSndFile, FileIO
include("WFSTs.jl")
include("Models.jl")
include("Utils.jl")

setup="2a"

include("conf/$(setup)/feat_conf.jl")
include("conf/$(setup)/model_conf.jl")

# get transition matrix
lexicon, ilexicon = get_lexicon()
H, L = get_HL(lexicon)
a, A = get_aA(H)
ippsym = get_iisym(H)

BSON.@load "models/$(setup)/current_modely.bson" best_modely
Flux.testmode!(best_modely)

feat_dir = get_feat_dir(setup)
data = load(joinpath(feat_dir,"test.jld2"))
uttID2feats, uttID2text = data["uttID2feats"], data["uttID2text"]

uttIDs = [keys(uttID2text)...]
i,j = rand(uttIDs),rand(uttIDs)

xi,xj = feats_post(uttID2feats[i]), feats_post(uttID2feats[j])
yi,yj = best_modely(Flux.unsqueeze(xi,3)),best_modely(Flux.unsqueeze(xj,3))

gammai = logposterior(size(yi,1),a,A,yi[:,:])
gammaj = logposterior(size(yj,1),a,A,yj[:,:])

outi = posterior2phones(ippsym,gammai)
outj = posterior2phones(ippsym,gammaj)
outi[outi .== "<SIL>"] .= " "
outj[outj .== "<SIL>"] .= " "

using Plots
pyplot()
psi = prod([prod(lexicon[w]) for w in split(uttID2text[i])].*" ")
p1i = heatmap(xi', title=uttID2text[i])
p2i = heatmap(yi[:,:]', clims = (maximum(yj)-20,maximum(yj)), title=psi)
p3i = heatmap(gammai',  clims = (-20,0), title=prod(outi))

psj = prod([prod(lexicon[w]) for w in split(uttID2text[j])].*" ")
p1j = heatmap(xj', title=uttID2text[j])
p2j = heatmap(yj[:,:]', clims = (maximum(yj)-20,maximum(yj)), title=psj)
p3j = heatmap(gammaj',  clims = (-20,0), title=prod(outj))

plot(p1i,p1j,p2i,p2j,p3i,p3j,layout=(3,2))


text:
schedList = DataFrame(Schedule = [	"OD",
	"STD",
	"XS",
	"XXS",
	"SCH10",
	"SCH20",
	"SCH30",
	"SCH40",
	"SCH60",
	"SCH80",
	"SCH100",
	"SCH120",
	"SCH140",
	"SCH160"] )


text:
include("..\\libs\\FunctionsRobot.jl")

function putmarker!(robot::CoordsRobot)
    x, y = get_xy(get_coords(robot))
    if x - y <= 0
        putmarker!(get_robot(robot))
    end
end

function lestnizaWithPartitions!(r::Robot)
    back_path = BackPath(r)
    R = CoordsRobot(r)
    marks!(R)
    BackPath(R)
    back!(R, back_path)
end


text:
#test1
F=readdlm("F.txt")

M=readdlm("M.txt")

W=readdlm("W.txt")

x=readdlm("x.txt")

y=readdlm("y.txt")

Fishmetica.testgmx(x,M,F)
Fishmetica.testgmy(x,y,M,F,W,1,1)

max_stars_repo_name: picestcode/Fishmetica

max_issues_repo_name: picestcode/Fishmetica

max_forks_repo_name: picestcode/Fishmetica


text:
using PkgBenchmark

#--track-allocation=all
config =
    BenchmarkConfig(id=nothing, juliacmd=`julia -O3`, env=Dict("JULIA_NUM_THREADS" => 4))

results = benchmarkpkg("Manifolds", config)
export_markdown("benchmark/results.md", results)


text:
using UltraDark
using Test
using MPI
using NPZ
using PencilFFTs

resol = 64

include("../../examples/init_soliton.jl")

output_dir = mktempdir()

for grid_type in [Grids, PencilGrids]

    this_output_dir = joinpath(output_dir, "$grid_type")
    output_times = [0, 1]

    output_config = OutputConfig(this_output_dir, output_times; box=true, slice=false)
    options = Config.SimulationConfig()

    grids = grid_type(10.0, resol)

    mass = 10
    position = [0, 0, 0]
    velocity = [0, 0, 0]
    phase = 0
    t0 = 0

    add_soliton(grids, mass, position, velocity, phase, t0)

    @test simulate(grids, options, output_config) == nothing
end

# for grid_type in [Grids, PencilGrids]
#     output_initial = npzread(joinpath(output_dir, "$grid_type", "psi_1.npy"))
#     output_final = npzread(joinpath(output_dir, "$grid_type", "psi_2.npy"))

#     @test all(output_initial .≈ output_final)
# end

for i in [1, 2]
    grids_output = npzread(joinpath(output_dir, "Grids", "psi_$i.npy"))
    pencil_grids_output = npzread(joinpath(output_dir, "PencilGrids", "psi_$i.npy"))

    # Fail for multiprocess - there are order 1e-12 differences
    @test all(grids_output .≈ pencil_grids_output)
end


text:
# In this visualization we wish to Generate an image of the Mandelbrot
# set, and interactive images of Julia sets. To do this, we will need
# to create our own painter function and colorscheme.
using Cairo, Colors, ColorTypes
using ComplexVisual; @ComplexVisual.import_huge
using ComplexVisualGtk; @ComplexVisualGtk.import_huge
using ComplexPortraits; @ComplexPortraits.import_huge

# We first need a painter function that redraws a portrait
# every time the action position changes. There are simply
# too many pixel options to cache every drawn image
struct CV_PortraitPainterWithoutCache{trafoT, CS} <: CV_2DCanvasPainter  # {{{
    trafo       :: trafoT
    colorscheme :: CS
end
function CV_PortraitPainterWithoutCache(trafo::trafoT=identity,
        colorscheme::CS=ComplexPortraits.cs_j()) where {trafoT, CS}
    return CV_PortraitPainterWithoutCache{trafoT, CS}(trafo, colorscheme)
end

# CV_PortraitPainterWithoutCache performs the same painting
# function as CV_PortraitPainter. See `painter.jl` for the original
import ComplexVisual: cv_paint
function cv_paint(cc::CV_2DCanvasContext{canvasT},
    portrait_painter::CV_PortraitPainterWithoutCache{CS}
    )  where {canvasT <: CV_Math2DCanvas, CS}   # {{{
    canvas = cc.canvas
    color_matrix = ComplexPortraits.portrait(
    canvas.corner_ul, canvas.corner_lr,
    portrait_painter.trafo,
    no_pixels=(canvas.pixel_height, canvas.pixel_width),
    point_color=portrait_painter.colorscheme)
    color_argb32 = map(color_matrix) do x
        convert(Colors.ARGB32, x).color
    end :: Matrix{UInt32}
    surface = cc.canvas.surface
    Cairo.flush(surface)
    surface.data .= color_argb32'
    Cairo.mark_dirty(surface)
    return nothing
end

# This is our approximation function for the Mandelbrot and Julia
# sets. Our goal is to see if the function `f(z) = z^2 + c` stays
# bounded after infinitely many iterations `f ∘ f ∘ ... ∘ f (z)`.
# Since we can't do infinite iterations, we use a finite number of
# iterations and a finite bounding circle. If
# `|fⁿᵈⁱᵗᵉʳ(z)| < maxabs`, we assume it will stay bounded.
function trafo(z; c=0.0 + 0.0im, maxabs=4, nditer=16)
    k = 1
    while abs(z) < maxabs && k < nditer
        z = z^2 + c
        k += 1
    end
    k == nditer && return 0
    return k % 360
end

# Now we need a custom colorscheme to paint each number
# of iterations a different color
function cs_circ(; colormap=hsv_colors())
    stepfct = generate_stepfct(length(colormap))
    let cm = RGB.(colormap)
      return function(z, fz)
        fz < 1 && return RGB(0, 0, 0)
        color = cm[stepfct(1/fz)]
        return RGB(color.r, color.g, color.b)
      end
    end
end

# We set up the standard left-right layout now
domain = CV_Math2DCanvas(-2.0+1.5im, 1.0-1.5im, 250)
codomain = CV_Math2DCanvas(-2.0+2.0im, 2.0-2.0im, 200)
action_pos = CV_TranslateByOffset(ComplexF64)
style = cv_color(1, 1, 1, 1) → cv_linewidth(3) → cv_antialias_best
cross = [[0.0+0.05im, 0.0-0.05im, 0.0+0.0im, 0.05+0.0im, -0.05+0.0im]]

# `:portrait_painter_domain` draws a custom phase portrait once
# at the beginning of the visualization on the left (domain) side.
#
# `:img_painter_domain`, `:img_painter_codomain` are options to
# change the interactive set that gets drawn on mouseclick. On the
# left side, we wish to show a small reticle to indicate where the
# user clicks. On the right side we draw the corresponding Julia set.
lr_painters_kwargs = Dict(
    :portrait_painter_domain => CV_PortraitPainter(
        w -> trafo(0.0+0.0im; c=w, maxabs=4, nditer=32), cs_circ()
    ),
    :img_painter_domain => style ↦ CV_LinePainter(
        action_pos, cross
    ),
    :img_painter_codomain => CV_PortraitPainterWithoutCache(
        w -> trafo(w; c=action_pos(0.0+0.0im), maxabs=4, nditer=16),
        cs_circ()
    ),
    :action_pos => action_pos
)

# Finally, we can see our work!
scene = cv_scene_lr_std(identity, domain, codomain; lr_painters_kwargs)
handler = cvg_visualize(scene); cvg_wait_for_destroy(handler.window)
cvg_close(handler); cv_destroy(scene)


text:
##############################################################################
##
## compare
## compare always return a value between 0 and 1.
##
##############################################################################

function compare(dist::PreMetric, s1::AbstractString, s2::AbstractString)
    1.0 - evaluate(dist, s1, s2)
end

function compare(dist::Union{Hamming, Levenshtein, DamerauLevenshtein},
    s1::AbstractString, s2::AbstractString)
    len = max(length(s1), length(s2))
    len == 0 ? 1.0 : 1.0 - evaluate(dist, s1, s2) / len
end

function compare(dist::AbstractQGram, s1::AbstractString, s2::AbstractString)
    # When string length < q for qgram distance, returns s1 == s2
    len1 = length(s1) ; len2 = length(s2)
    min(len1, len2) <= (dist.q - 1) && return convert(Float64, s1 == s2)
    if typeof(dist) <: QGram
        1 - evaluate(dist, s1, s2) / (len1 + len2 - 2 * dist.q + 2)
    else
        1 - evaluate(dist, s1, s2)
    end
end

##############################################################################
##
## Winkler
##
##############################################################################

struct Winkler{T1 <: PreMetric, T2 <: Real, T3 <: Real} <: PreMetric
    dist::T1
    scaling_factor::T2      # scaling factor. Default to 0.1
    boosting_limit::T3      # boost threshold. Default to 0.7
end

# restrict to distance between 0 and 1
Winkler(x) = Winkler(x, 0.1, 0.7)

function compare(dist::Winkler, s1::AbstractString, s2::AbstractString)
    score = compare(dist.dist, s1, s2)
    l = common_prefix(s1, s2, 4)[1]
    # common prefix adjustment
    if score >= dist.boosting_limit
        score += l * dist.scaling_factor * (1 - score)
    end
    return score
end

##############################################################################
##
## Partial
## http://chairnerd.seatgeek.com/fuzzywuzzy-fuzzy-string-matching-in-python/
##
##############################################################################
struct Partial{T <: PreMetric} <: PreMetric
    dist::T
end

# general
function compare(dist::Partial, s1::AbstractString, s2::AbstractString)
    s2, len2, s1, len1 = reorder(s1, s2)
    len1 == len2 && return compare(dist.dist, s1, s2)
    len1 == 0 && return compare(dist.dist, "", "")
    iter = QGramIterator(s2, len2, len1)
    out = 0.0
    x = iterate(iter)
    while x !== nothing
        s, state = x
        curr = compare(dist.dist, s1, s)
        out = max(out, curr)
        x = iterate(iter, state)
    end
    return out
end

# Specialization for RatcliffObershelp distance
# Code follows https://github.com/seatgeek/fuzzywuzzy/blob/master/fuzzywuzzy/fuzz.py
function compare(dist::Partial{RatcliffObershelp}, s1::AbstractString, s2::AbstractString)
    s2, len2, s1, len1 = reorder(s1, s2)
    len1 == len2 && return compare(dist.dist, s1, s2)
    out = 0.0
    result = matching_blocks(s1, s2)
    for r in result
        # here I difffer from fuzz.py by making sure the substring of s2 has length len1
        s2_start = r[2] - r[1] + 1
        s2_end = s2_start + len1 - 1
        if s2_start <= 0
            s2_end += 1 - s2_start
            s2_start += 1 - s2_start
        elseif s2_end > len2
            s2_start += len2 - s2_end
            s2_end += len2 - s2_end
        end
        i2_start =  nextind(s2, 0, s2_start)
        i2_end = s2_end == len2 ? lastindex(s2) : (nextind(s2, 0, s2_end + 1) - 1)
        curr = compare(RatcliffObershelp(), s1, SubString(s2, i2_start, i2_end))
        out = max(out, curr)
    end
    return out
end

##############################################################################
##
## TokenSort
## http://chairnerd.seatgeek.com/fuzzywuzzy-fuzzy-string-matching-in-python/
##
##############################################################################
struct TokenSort{T <: PreMetric} <: PreMetric
    dist::T
end

function compare(dist::TokenSort, s1::AbstractString, s2::AbstractString)
    s1 = join(sort!(split(s1)), " ")
    s2 = join(sort!(split(s2)), " ")
    compare(dist.dist, s1, s2)
end

##############################################################################
##
## TokenSet
## http://chairnerd.seatgeek.com/fuzzywuzzy-fuzzy-string-matching-in-python/
##
##############################################################################
struct TokenSet{T <: PreMetric} <: PreMetric
    dist::T
end

function compare(dist::TokenSet, s1::AbstractString, s2::AbstractString)
    v0, v1, v2 = _separate!(split(s1), split(s2))
    s0 = join(v0, " ")
    s1 = join(Iterators.flatten((v0, v1)), " ")
    s2 = join(Iterators.flatten((v0, v2)), " ")
    if isempty(s0)
        # otherwise compare(dist, "", "a")== 1.0
        compare(dist.dist, s1, s2)
    else
        max(compare(dist.dist, s0, s1),
            compare(dist.dist, s1, s2),
            compare(dist.dist, s0, s2))
    end
end

# separate 2 vectors in intersection, setdiff1, setdiff2 (all sorted)
function _separate!(v1::Vector, v2::Vector)
    sort!(v1)
    sort!(v2)
    out = eltype(v1)[]
    start = 1
    i1 = 0
    while i1 < length(v1)
        i1 += 1
        x = v1[i1]
        i2 = searchsortedfirst(v2, x, start, length(v2), Base.Forward)
        i2 > length(v2) && break
        if i2 > 0 && v2[i2] == x
            deleteat!(v1, i1)
            deleteat!(v2, i2)
            push!(out, x)
            i1 -= 1
            start = i2
        end
    end
    return out, v1, v2
end

##############################################################################
##
## TokenMax
##
##############################################################################
struct TokenMax{T <: PreMetric} <: PreMetric
    dist::T
end

function compare(dist::TokenMax, s1::AbstractString, s2::AbstractString)
    dist0 = compare(dist.dist, s1, s2)
    s2, len2, s1, len1 = reorder(s1, s2)
    unbase_scale = 0.95
    # if one string is much much shorter than the other
    if len2 >= 1.5 * len1
        # if strings are of dissimilar length, use partials
        partial = compare(Partial(dist.dist), s1, s2)
        ptsor = compare(TokenSort(Partial(dist.dist)), s1, s2)
        ptser = compare(TokenSet(Partial(dist.dist)), s1, s2)
        partial_scale = len2 > (8 * len1) ? 0.6 : 0.9
        return max(dist0,
                partial * partial_scale,
                ptsor * unbase_scale * partial_scale,
                ptser * unbase_scale * partial_scale)
    else
        ptsor = compare(TokenSort(dist.dist), s1, s2)
        ptser = compare(TokenSet(dist.dist), s1, s2)
        return max(dist0,
                ptsor * unbase_scale,
                ptser * unbase_scale)
    end
end


text:
"""
    with_tracers(tracers, initial_tuple, tracer_default)

Create a tuple corresponding to the solution variables `u`, `v`, `w`,
and `tracers`. `initial_tuple` is a `NamedTuple` that at least has
fields `u`, `v`, and `w`, and may have some fields corresponding to
the names in `tracers`. `tracer_default` is a function that produces
a default tuple value for each tracer if not included in `initial_tuple`.
"""
function with_tracers(tracers, initial_tuple::NamedTuple, tracer_default; with_velocities=false)
    solution_values = [] # Array{Any, 1}
    solution_names = []

    if with_velocities
        push!(solution_values, initial_tuple.u)
        push!(solution_values, initial_tuple.v)
        push!(solution_values, initial_tuple.w)

        append!(solution_names, [:u, :v, :w])
    end

    for name in tracers
        tracer_elem = name ∈ propertynames(initial_tuple) ?
                        getproperty(initial_tuple, name) :
                        tracer_default(tracers, initial_tuple)

        push!(solution_values, tracer_elem)
    end

    append!(solution_names, tracers)

    return NamedTuple{Tuple(solution_names)}(Tuple(solution_values))
end


text:
module test_jive_If

using Jive # @If

@If true module A
f() = 1
end

@If false module B
f() = 1
end

using Test
@test isdefined(@__MODULE__, :A)
@test !isdefined(@__MODULE__, :B)
@test A.f() == 1

end # module test_jive_If


text:
# PST 3201 power suuply driver

mutable struct PST3201
    handle
	initialized
    instr_dict
end
PST3201(handle) = PST3201(
                handle,
				false,
                Dict(
                    "C1"=>"1",
                    "C2"=>"2",
                    "C3"=>"3",
                    "on"=>"1",
                    "off"=>"0",
                    "1"=>"on",
                    "0"=>"off",
                    "voltage"=>"VOLT",
                    "current"=>"CURR",
                    )
                    )

# MISC
function connect!(obj::PST3201, handle)
	obj.handle = handle
    obj.connected = 1
	return nothing
end
function disconnect!(obj::PST3201)
	obj.handle = 0
    obj.connected = 0
	return nothing
end
reset_instr(obj::PST3201) = write(obj.handle, "*RST")
clear_instr(obj::PST3201) = write(obj.handle, "*CLS")
get_idn(obj::PST3201) = query(obj.handle, "*IDN?")

# SOURCE VOLTAGE
function set_source_lev(obj::PST3201, ch, lev)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:VOLT $lev"
    write(obj.handle, cmd)
end
function get_source_lev(obj::PST3201, ch)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:VOLT?"
    value = parse(Float64, strip(query(obj.handle, cmd)))
end

# PROTECTION
function set_volt_protection(obj::PST3201, ch, lev)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:PROT:VOLT $lev"
    write(obj.handle, cmd)
end
function get_volt_protection(obj::PST3201, ch)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:PROT:VOLT?"
    value = parse(Float64, strip(query(obj.handle, cmd)))
end
function set_curr_protection(obj::PST3201, ch, state)
    ch = obj.instr_dict[ch]
    state = obj.instr_dict[state]
    cmd = ":CHAN$ch:PROT:CURR $state"
    write(obj.handle, cmd)
end
function get_curr_protection(obj::PST3201, ch)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:PROT:CURR?"
    value = strip(query(obj.handle, cmd))
    value = obj.instr_dict[value]
end
function set_max_curr(obj::PST3201, ch, lev)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:CURR $lev"
    write(obj.handle, cmd)
end
function get_max_curr(obj::PST3201, ch)
    ch = obj.instr_dict[ch]
    cmd = ":CHAN$ch:CURR?"
    value = parse(Float64, strip(query(obj.handle, cmd)))
end
function reset_protections(obj::PST3201)
    cmd = ":OUTPut:PROTection:CLEar"
    write(obj.handle, cmd)
end

# MEASUREMENTS
function get_meas(obj::PST3201, ch, fct)
    ch = obj.instr_dict[ch]
    fct = obj.instr_dict[fct]
    cmd = ":CHAN$ch:MEAS:$fct?"
    value = parse(Float64, strip(query(obj.handle, cmd)))
end

# OUTPUT ON/OFF
function set_outp(obj::PST3201, ch, state)
    ch = obj.instr_dict[ch]
    state = obj.instr_dict[state]
    cmd = ":OUTPut:STATe $state"
    write(obj.handle, cmd)
end
function get_outp(obj::PST3201, ch)
    ch = obj.instr_dict[ch]
    cmd = ":OUTPut:STATe?"
    value = strip(query(obj.handle, cmd))
    value = obj.instr_dict[value]
end

max_stars_repo_name: iuliancioarca/TIVM

max_issues_repo_name: iuliancioarca/TIVM

max_forks_repo_name: iuliancioarca/TIVM


text:
#!/usr/bin/env julia
using ArgParse

s = ArgParseSettings(
        description = "Calculate time evolution for given spin geometry."
    )

@add_arg_table s begin
    "--geometry"
    "--N"
        arg_type = Int
    "--d"
        arg_type = Float64
    "--gamma"
        default = "1.0"
    "--edipole"
    "--phi"
        arg_type = Float64
    "--theta"
        arg_type = Float64
    "--o"
    "--T"
    "--method"
        range_tester = x->(x in ["master", "meanfield", "mpc"])
end

parameters = parse_args(s)

# Output
const odir = parameters["o"]
@assert isdir(odir)

# Integration
const T = float(eval(Meta.parse(parameters["T"])))
@assert T[1]<T[end]

# System parameters
const γ = float(parameters["gamma"])
@assert 0<γ

# Initial state (Bloch state)
const phi = float(parameters["phi"])
const theta = float(parameters["theta"])

# System geometry
using quantumoptics, CollectiveSpins
const cs = CollectiveSpins

const edipole = float(eval(Meta.parse(parameters["edipole"])))
const geomN = int(parameters["N"])
const d = float(parameters["d"])
const geomstring = parameters["geometry"]
if geomstring=="chain"
    const systemgeometry = cs.geometry.chain(d, 2*geomN+1)
elseif geomstring=="square"
    const systemgeometry = cs.geometry.square(d; Nx=2*geomN+1, Ny=2*geomN+1)
elseif geomstring=="hexagonal"
    const systemgeometry = cs.geometry.hexagonal(d; Nr=geomN)
elseif geomstring=="cube"
    const systemgeometry = cs.geometry.cube(d; Nx=2*geomN+1, Ny=2*geomN+1, Nz=2*geomN+1)
else
    error("Unknown geometry: $(geomstring)")
end

const system = SpinCollection(systemgeometry, edipole, γ)
const N = length(system.spins)
const index_center = int((N+1)/2)

const method = parameters["method"]

const sx = Float64[]
const sy = Float64[]
const sz = Float64[]

tic()
if method=="meanfield"
    state0 = cs.meanfield.blochstate(phi,theta,N)
    function fout(t, state)
        push!(sx, cs.meanfield.sx(state)[index_center])
        push!(sy, cs.meanfield.sy(state)[index_center])
        push!(sz, cs.meanfield.sz(state)[index_center])
    end
    cs.meanfield.timeevolution(T, system, state0; fout=fout)
elseif method=="mpc"
    state0 = cs.mpc.blochstate(phi,theta,N)
    function fout(t, state)
        push!(sx, cs.mpc.sx(state)[index_center])
        push!(sy, cs.mpc.sy(state)[index_center])
        push!(sz, cs.mpc.sz(state)[index_center])
    end
    cs.mpc.timeevolution(T, system, state0; fout=fout)
elseif method=="master"
    embed(op::Operator) = QuantumOptics.embed(cs.quantum.basis(system), [index_center], [op])
    function fout(t, rho::Operator)
        push!(sx, real(expect(embed(sigmax), rho)))
        push!(sy, real(expect(embed(sigmay), rho)))
        push!(sz, real(expect(embed(sigmaz), rho)))
    end
    Ψ₀ = cs.quantum.blochstate(phi,theta,N)
    ρ₀ = Ψ₀⊗dagger(Ψ₀)
    cs.quantum.timeevolution(T, system, ρ₀; fout=fout)
end
t = toc()

keyparameters = Dict(
    "phi"=>parameters["phi"],
    "theta"=>parameters["theta"],
    "edipole"=>parameters["edipole"],
    "d"=>parameters["d"],
    "N"=>parameters["N"],
    "T"=>parameters["T"],
    "method"=>parameters["method"]
    )

name = QuantumOptics.io.dict2filename(keyparameters)
f = open(joinpath(odir, name), "w")
write(f, "# Time sx sy sz\n")
write(f, "#calctime=$(t)\n")
for i=1:length(T)
    write(f, "$(T[i]);$(sx[i]);$(sy[i]);$(sz[i])\n")
end

close(f)


text:
x = [ complex(cos(θ),sin(θ)) for θ in 2π*(0:10)/10 ]

sum(x.^2)

y = [ a//(a+1) for a = 1:10 ]

prod(y.^2)

convert(Array{Float64}, y)


text:
@testset "smc.jl" begin
    @testset "SMC constructor" begin
        sampler = AdvancedPS.SMC(10)
        @test sampler.nparticles == 10
        @test sampler.resampler == AdvancedPS.ResampleWithESSThreshold()

        sampler = AdvancedPS.SMC(15, 0.6)
        @test sampler.nparticles == 15
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)

        sampler = AdvancedPS.SMC(20, AdvancedPS.resample_multinomial, 0.6)
        @test sampler.nparticles == 20
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)

        sampler = AdvancedPS.SMC(25, AdvancedPS.resample_systematic)
        @test sampler.nparticles == 25
        @test sampler.resampler === AdvancedPS.resample_systematic
    end

    # Smoke tests
    @testset "models" begin
        mutable struct NormalModel <: AbstractMCMC.AbstractModel
            a::Float64
            b::Float64

            NormalModel() = new()
        end

        function (m::NormalModel)(rng::Random.AbstractRNG)
            # First latent variable.
            m.a = a = rand(rng, Normal(4, 5))

            # First observation.
            AdvancedPS.observe(Normal(a, 2), 3)

            # Second latent variable.
            m.b = b = rand(rng, Normal(a, 1))

            # Second observation.
            return AdvancedPS.observe(Normal(b, 2), 1.5)
        end

        sample(NormalModel(), AdvancedPS.SMC(100))

        # failing test
        mutable struct FailSMCModel <: AbstractMCMC.AbstractModel
            a::Float64
            b::Float64

            FailSMCModel() = new()
        end

        function (m::FailSMCModel)(rng::Random.AbstractRNG)
            m.a = a = rand(rng, Normal(4, 5))
            m.b = b = rand(rng, Normal(a, 1))
            if a >= 4
                AdvancedPS.observe(Normal(b, 2), 1.5)
            end
        end

        @test_throws ErrorException sample(FailSMCModel(), AdvancedPS.SMC(100))
    end

    @testset "logevidence" begin
        Random.seed!(100)

        mutable struct TestModel <: AbstractMCMC.AbstractModel
            a::Float64
            x::Bool
            b::Float64
            c::Float64

            TestModel() = new()
        end

        function (m::TestModel)(rng::Random.AbstractRNG)
            # First hidden variables.
            m.a = rand(rng, Normal(0, 1))
            m.x = x = rand(rng, Bernoulli(1))
            m.b = rand(rng, Gamma(2, 3))

            # First observation.
            AdvancedPS.observe(Bernoulli(x / 2), 1)

            # Second hidden variable.
            m.c = rand(rng, Beta())

            # Second observation.
            return AdvancedPS.observe(Bernoulli(x / 2), 0)
        end

        chains_smc = sample(TestModel(), AdvancedPS.SMC(100))

        @test all(isone(p.f.x) for p in chains_smc.trajectories)
        @test chains_smc.logevidence ≈ -2 * log(2)
    end

    @testset "PG constructor" begin
        sampler = AdvancedPS.PG(10)
        @test sampler.nparticles == 10
        @test sampler.resampler == AdvancedPS.ResampleWithESSThreshold()

        sampler = AdvancedPS.PG(60, 0.6)
        @test sampler.nparticles == 60
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_systematic, 0.6)

        sampler = AdvancedPS.PG(80, AdvancedPS.resample_multinomial, 0.6)
        @test sampler.nparticles == 80
        @test sampler.resampler ===
              AdvancedPS.ResampleWithESSThreshold(AdvancedPS.resample_multinomial, 0.6)

        sampler = AdvancedPS.PG(100, AdvancedPS.resample_systematic)
        @test sampler.nparticles == 100
        @test sampler.resampler === AdvancedPS.resample_systematic
    end

    @testset "logevidence" begin
        Random.seed!(100)

        mutable struct TestModel <: AbstractMCMC.AbstractModel
            a::Float64
            x::Bool
            b::Float64
            c::Float64

            TestModel() = new()
        end

        function (m::TestModel)(rng::Random.AbstractRNG)
            # First hidden variables.
            m.a = rand(rng, Normal(0, 1))
            m.x = x = rand(rng, Bernoulli(1))
            m.b = rand(rng, Gamma(2, 3))

            # First observation.
            AdvancedPS.observe(Bernoulli(x / 2), 1)

            # Second hidden variable.
            m.c = rand(rng, Beta())

            # Second observation.
            return AdvancedPS.observe(Bernoulli(x / 2), 0)
        end

        chains_pg = sample(TestModel(), AdvancedPS.PG(10), 100)

        @test all(isone(p.trajectory.f.x) for p in chains_pg)
        @test mean(x.logevidence for x in chains_pg) ≈ -2 * log(2) atol = 0.01
    end

    @testset "Replay reference" begin
        mutable struct Model <: AbstractMCMC.AbstractModel
            a::Float64
            b::Float64

            Model() = new()
        end

        function (m::Model)(rng)
            m.a = rand(rng, Normal())
            AdvancedPS.observe(Normal(), m.a)

            m.b = rand(rng, Normal())
            return AdvancedPS.observe(Normal(), m.b)
        end

        pg = AdvancedPS.PG(1)
        first, second = sample(Model(), pg, 2)

        first_model = first.trajectory.f
        second_model = second.trajectory.f

        # Single Particle - must be replaying
        @test first_model.a ≈ second_model.a
        @test first_model.b ≈ second_model.b
        @test first.logevidence ≈ second.logevidence
    end
end

# @testset "pmmh.jl" begin
#     @turing_testset "pmmh constructor" begin
#         N = 2000
#         s1 = PMMH(N, SMC(10, :s), MH(1,(:m, s -> Normal(s, sqrt(1)))))
#         s2 = PMMH(N, SMC(10, :s), MH(1, :m))
#         s3 = PIMH(N, SMC())
#
#         c1 = sample(gdemo_default, s1)
#         c2 = sample(gdemo_default, s2)
#         c3 = sample(gdemo_default, s3)
#     end
#     @numerical_testset "pmmh inference" begin
#         alg = PMMH(2000, SMC(20, :m), MH(1, (:s, GKernel(1))))
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain, atol = 0.1)
#
#         # PMMH with prior as proposal
#         alg = PMMH(2000, SMC(20, :m), MH(1, :s))
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain, atol = 0.1)
#
#         # PIMH
#         alg = PIMH(2000, SMC())
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain)
#
#         # MoGtest
#         pmmh = PMMH(2000,
#             SMC(10, :z1, :z2, :z3, :z4),
#             MH(1, :mu1, :mu2))
#         chain = sample(MoGtest_default, pmmh)
#
#         check_MoGtest_default(chain, atol = 0.1)
#     end
# end

# @testset "ipmcmc.jl" begin
#     @turing_testset "ipmcmc constructor" begin
#         Random.seed!(125)
#
#         N = 50
#         s1 = IPMCMC(10, N, 4, 2)
#         s2 = IPMCMC(10, N, 4)
#
#         c1 = sample(gdemo_default, s1)
#         c2 = sample(gdemo_default, s2)
#     end
#     @numerical_testset "ipmcmc inference" begin
#         alg = IPMCMC(30, 500, 4)
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain)
#
#         alg2 = IPMCMC(15, 100, 10)
#         chain2 = sample(MoGtest_default, alg2)
#         check_MoGtest_default(chain2, atol = 0.2)
#     end
# end


text:
plan_clenshawcurtis(μ) = length(μ) > 1 ? FFTW.plan_r2r!(μ, FFTW.REDFT00) : ones(μ)'

"""
Compute nodes of the Clenshaw—Curtis quadrature rule.
"""
clenshawcurtisnodes(::Type{T}, N::Int) where T = T[cospi(k/(N-one(T))) for k=0:N-1]

"""
Compute weights of the Clenshaw—Curtis quadrature rule with modified Chebyshev moments of the first kind ``\\mu``.
"""
clenshawcurtisweights(μ::Vector) = clenshawcurtisweights!(copy(μ))
clenshawcurtisweights!(μ::Vector) = clenshawcurtisweights!(μ, plan_clenshawcurtis(μ))
function clenshawcurtisweights!(μ::Vector{T}, plan) where T
    N = length(μ)
    rmul!(μ, inv(N-one(T)))
    plan*μ
    μ[1] *= half(T); μ[N] *= half(T)
    return μ
end

# Chebyshev-T coefficients to values at Clenshaw-Curtis nodes

applyTN_plan(x) = length(x) > 1 ? FFTW.plan_r2r!(x, FFTW.REDFT00) : fill!(similar(x),1)'

applyTN!(x::Vector{T}) where {T<:AbstractFloat} = applyTN!(x,applyTN_plan(x))

function applyTN!(x::Vector{T},plan) where T<:AbstractFloat
    x[1] *= 2; x[end] *=2
    plan*x
    rmul!(x,half(T))
end
applyTN(x::Vector{T},plan) where {T<:AbstractFloat} = applyTN!(copy(x),plan)
applyTN(x::Vector{T}) where {T<:AbstractFloat} = applyTN!(copy(x))

# Values at Clenshaw-Curtis nodes to Chebyshev-T coefficients

applyTNinv_plan(x) = length(x) > 1 ? FFTW.plan_r2r!(x, FFTW.REDFT00) : ones(x)'

applyTNinv!(x::Vector{T}) where {T<:AbstractFloat} = applyTNinv!(x,applyTNinv_plan(x))

function applyTNinv!(x::Vector{T},plan) where T<:AbstractFloat
    plan*x
    x[1] /= 2;x[end] /= 2
    rmul!(x,inv(length(x)-one(T)))
end
applyTNinv(x::Vector{T},plan) where {T<:AbstractFloat} = applyTNinv!(copy(x),plan)
applyTNinv(x::Vector{T}) where {T<:AbstractFloat} = applyTNinv!(copy(x))

# sin(nθ) coefficients to values at Clenshaw-Curtis nodes except ±1

applyUN_plan(x) = length(x) > 0 ? FFTW.plan_r2r!(x, FFTW.RODFT00) : fill!(similar(x),1)'

applyUN!(x::AbstractVector{T}) where {T<:AbstractFloat} = applyUN!(x,applyUN_plan(x))

function applyUN!(x::AbstractVector{T},plan) where T<:AbstractFloat
    plan*x
    rmul!(x,half(T))
end
applyUN(x::AbstractVector{T},plan) where {T<:AbstractFloat} = applyUN!(copy(x),plan)
applyUN(x::AbstractVector{T}) where {T<:AbstractFloat} = applyUN!(copy(x))


text:
# Copyright (c) 2013: Joey Huchette and contributors
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

using CPLEX
using Test

@testset "MathOptInterface Tests" begin
    for file in readdir("MathOptInterface")
        include(joinpath("MathOptInterface", file))
    end
end


text:
VERSION >= v"0.4.0-dev+6641" && __precompile__()
module StructIO

    using Base: @pure
    using Base.Meta
    export @struct, unpack, pack, fix_endian

    needs_bswap(endianness) = (ENDIAN_BOM == 0x01020304) ?
        endianness == :LittleEndian : endianness == :BigEndian
    fix_endian(x, endianness) = needs_bswap(x) ? bswap(x) : x

    # Alignment traits
    abstract PackingStrategy
    immutable Packed <: PackingStrategy; end
    immutable Default <: PackingStrategy; end
    function strategy
    end

    # Sizeof computation
    round_up(offset, alignment) = offset +
        mod(alignment - mod(offset, alignment), alignment)

    @pure function sizeof(T::DataType, ::Type{Default})
        Core.sizeof(T)
    end

    @pure function sizeof(T::DataType, ::Type{Packed})
        @assert nfields(T) != 0 && isbits(T)
        sum(sizeof, T.types)
    end
    sizeof(T::DataType) = sizeof(T, nfields(T) == 0 ? Default : strategy(T))

    # Generates methods for unpack!, pack!, and sizeof
    macro struct(typ, annotations...)
        alignment = :align_default
        if length(annotations) == 1
            ann = annotations[1]
            if isa(ann, Symbol) || haskey(alignments, ann)
                alignment = ann
            end
        end
        typname = typ.args[2]
        isexpr(typname,:(<:)) && (typname = typname.args[1])
        isexpr(typname,:curly) && (typname = typname.args[1])
        ret = Expr(:toplevel, typ)
        if alignment == :align_default
            push!(ret.args, :(StructIO.strategy(::Type{$typname}) = StructIO.Default))
        else
            @assert alignment == :align_packed
            push!(ret.args, :(StructIO.strategy(::Type{$typname}) = StructIO.Packed))
        end
        push!(ret.args, :(Base.sizeof(::Type{$typname}) = StructIO.sizeof($typname)))
        push!(ret.args, :(Base.sizeof(::$typname) = StructIO.sizeof($typname)))
        esc(ret)
    end

    function unsafe_unpack(io::IO, T::Type, target, endianness, ::Type{Default})
        if nfields(T) == 0
            sz = Core.sizeof(T)
            unsafe_read(io, target, sz)
            if needs_bswap(endianness)
                # Special case small sizes, LLVM should turn this into a jump
                # table
                if sz == 1
                elseif sz == 2
                    ptr = Base.unsafe_convert(Ptr{UInt16}, target)
                    unsafe_store!(ptr, bswap(unsafe_load(ptr)))
                elseif sz == 4
                    ptr = Base.unsafe_convert(Ptr{UInt32}, target)
                    unsafe_store!(ptr, bswap(unsafe_load(ptr)))
                elseif sz == 8
                    ptr = Base.unsafe_convert(Ptr{UInt64}, target)
                    unsafe_store!(ptr, bswap(unsafe_load(ptr)))
                else
                    for i = 0:div(sz,2)
                        ptrhigh = Base.unsafe_convert(Ptr{UInt8}, target) + 8*(sz-i)
                        ptrlow = Base.unsafe_convert(Ptr{UInt8}, target) + 8*i
                        high = unsafe_load(ptrhigh)
                        low = unsafe_load(ptrlow)
                        unsafe_store(ptrhigh, low)
                        unsafe_store(ptrlow, high)
                    end
                end
            end
        elseif !needs_bswap(endianness)
            sz = Core.sizeof(T)
            unsafe_read(io, target, sz)
        else
            reached = 0
            for i = 1:nfields(T)
                fT = fieldtype(T, i)
                foffs = fieldoffset(T, i)
                skip(io, reached - foffs)
                reached = foffs
                unsafe_unpack(io, fT,
                    Base.unsafe_convert(Ptr{Void}, target) + foffs, endianness, Default)
                reached += Core.sizeof(fT)
            end
        end
    end

    function unsafe_unpack(io::IO, T::Type, target, endianness, ::Type{Packed})
        nfields(T) == 0 && return unsafe_unpack(io, T, target, endianness, Default)
        for i = 1:nfields(T)
            fT = fieldtype(T, i)
            unsafe_unpack(io, fT,
                Base.unsafe_convert(Ptr{Void}, target) + fieldoffset(T, i), endianness, Packed)
        end
    end

    function unpack(io::IO, T::Type, endianness = :NativeEndian)
        r = Ref{T}()
        unsafe_unpack(io, T, r, endianness, nfields(T) == 0 ? Default : strategy(T))
        r[]
    end

end # module


text:
using MAT
using CuArrays
using MatrixProductStates
BLAS.set_num_threads(nworkers())

#Spin model parameters
const na = 60 # number of atoms
const od = 33 # optical depth of cloud
const ad = 2.5*60/na #atomic distance in micro m
const cloud_sig = 23.6344/ad # cloud length
const cloud_coup = exp.(-((1:na)-(na+1)/2).^2/2/cloud_sig^2) # atomic cloud distribution
const gam_1d = cloud_coup/sum(cloud_coup)/2*od # coupling of atoms
const gam_eg = 1.0 # eg spontaneous decay rate
const del_p = 0.0 # detuning of pump beam
const k_wg = 0.5*pi # waveguide wavevector
const rj = collect(1.0:na) #sort!(rand(na)*na) # atom positions
const k_in = k_wg # pump beam wavevector
const gam_exp = 2*pi*6.065e6 # spontaneous emmision rate
const f_amp_exp = sqrt.([0.3657, 2.9549, 7.7137, 10.0961]/gam_exp/1e-6)
const f_amp = f_amp_exp[2]

#Rydberg specific parameters
const om = 10.0/6.065/2.0 # control beam Rabi frequency
const k_c = 0.0 # control wavevector
const del_s = 0.0 # two photon detuning from s level
const gam_sp = 0.100/6.065
const gam_pg = 0.100/6.065

# rmax = 20 na = 70
#uu = [524.328+89.9671im, 524.328-89.9671im, 649.737+219.429im, 649.737-219.429im, 488.928]
#lam = [-0.228929 + 0.189648im, -0.228929 - 0.189648im, 0.146279+0.297638im, 0.146279-0.297638im, 0.480032]

# rmax = 30 na = 60
#uu = [1055.32, 1156.39 + 264.599im, 1156.39 - 264.599im, 781.661]
#lam = [-0.220397, 0.0354409 + 0.244048im, 0.0354409 - 0.244048im, 0.396262]
#uu = [15158.7, 11895.1 - 1929.93im, 11895.1 + 1929.93im, 2399.43, 16.4076]
#lam = [-0.103573, 0.015448 + 0.126927im, 0.015448 - 0.126927im, 0.30542, 0.606657]
#uu = [90866.8, 59173.0 - 18223.1im, 59173.0 + 18223.1im, 5257.17, 107.47, 1.26558]
#lam = [-0.0606795, 0.0112804 + 0.0784944im, 0.0112804 - 0.0784944im, 0.246432, 0.48336, 0.720283]

# rmax = 20 na = 60
const uu = [0.704528 + 135.691im, 0.704528 - 135.691im, 225.917]
const lam = [0.226209 + 0.32374im, 0.226209 - 0.32374im, 0.475997]

#uu = [ 2255.07, 686.082 + 91.2384im, 686.082 - 91.2384im, 477.07]
#lam = [-0.0998312, 0.0646318 + 0.25377im, 0.0646318-0.25377im, 0.424938]
#uu = [-508.38 - 1018.61im, -508.38 + 1018.61im, -1546.1 - 1942.74im, -1546.1 + 1942.74im,
#        28.5203]
#lam = [ -0.0374226 + 0.193987im, -0.0374226 - 0.193987im, 0.250214 + 0.0887921im,
#       0.250214 - 0.0887921im, 0.58104]
#uu = [-658.963 + 91.3732im, -658.963 - 91.3732im, 309.325 + 1085.49im, 309.325 - 1085.49im,
#       709.204, 4.33617]
#lam = [-0.0921073 + 0.171392im, -0.0921073 - 0.171392im, 0.124872 + 0.194013im,
#       0.124872 - 0.194013im, 0.382177, 0.670268]

# rmax = 20 na = 50
#uu = [1740.29, 2277.16 + 454.662im, 2277.16 - 454.662im, 945.959, 6.06471]
#lam = [-0.126487, 0.0183872 + 0.140738im, 0.0183872 - 0.140738im, 0.296793, 0.601841]
#uu = [ 57.1125, 170.91 + 211.667im, 170.91 - 211.667im, 336.281]
#lam = [ -0.298989, 0.0548962 + 0.260436im, 0.0548962 - 0.260436im, 0.382366]

# rmax = 20 na = 40
#uu = [19014.9 - 3038.05im, 19014.9 + 3038.05im, 2635.02, 47.4863, 0.531659]
#lam = [-0.0167917 + 0.0377784im, -0.0167917 - 0.0377784im, 0.155726, 0.386972, 0.653397]

# rmax = 20 na = 30
#uu = [-2652.81-1864.93im, -2652.81 + 1864.93im, 274.344, 5.91125, 0.0704647]
#lam = [0.0313486 + 0.0360902im, 0.0313486 - 0.0360902im, 0.179561, 0.407673, 0.666632]

# rmax = 20 na = 20
#uu = [1350.25, 707.234, 15.7581, 0.426333, 0.0053846]
#lam = [-0.0138335, 0.0500349, 0.19712, 0.418572, 0.672633]

# pure dephasing
const gam_gg = 0.0
const gam_ss = 0.0
const gam_ee = 0.0

# simulation parameters
const dt = 0.01
const t_fin = 60
const d = 4
const d_max = 200
const measure_int = 5
const path_data = string(homedir(), "/data/")
const base_filename = string(path_data,"Ryd_Dens_N",na,"_D",d_max,
    "_Tf",t_fin,"_f",round(f_amp,3),"_dt",dt,"_three_exp_float32_gsp",
    round(gam_sp,3),"_gpg",round(gam_pg,3))
const g2_filename = string(base_filename, "_g2")

# input pulse envelope
function f(t)

    f_amp

end

# local spin operators
const hgg = [1 0 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 0]
const hee = [0 0 0 0; 0 1 0 0; 0 0 0 0; 0 0 0 0]
const hss = [0 0 0 0; 0 0 0 0; 0 0 1 0; 0 0 0 0]
const hpp = [0 0 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 1]
const hge = [0 1 0 0; 0 0 0 0; 0 0 0 0; 0 0 0 0]
const hes = [0 0 0 0; 0 0 1 0; 0 0 0 0; 0 0 0 0]
const hsp = [0 0 0 0; 0 0 0 0; 0 0 0 1; 0 0 0 0]
const hgp = [0 0 0 1; 0 0 0 0; 0 0 0 0; 0 0 0 0]
const heg = hge'
const hse = hes'
const hps = hsp'
const hpg = hgp'
const id = eye(4)

# types used in mps representation
const TN = Complex{Float32}
const TA = CuArray

function time_evolve()

    #Create measurement operators
    IDlmps = makemps(TN, TA, [jj->id[:]], na, d^2)
    ERmpo = makempo(TN, TA, [jj->id jj->im*sqrt(gam_1d[jj]/2)*exp(-im*k_wg*rj[jj])*hge;
                        jj->0 jj->id], na, d)
    ERmpo[1][1, :, 2, :] = f(0.0)*id + im*sqrt(gam_1d[1]/2)*exp(-im*k_wg*rj[1])*hge
    IRmpo = applyMPOtoMPO(ERmpo, conj_mpo(ERmpo))
    IRlmps = mpo_to_mps(TN, TA, IRmpo)

    #step times and measurement times
    t = 0.0:dt:t_fin
    t_m = t[1:measure_int:end]
    tstep = length(t) - 1
    times = zeros(tstep, 4)

    file = matopen(string(g2_filename, "_temp.mat"), "r")
    rho_temp = read(file, "rho")
    tr_rho = read(file, "tr_rho")
    I_r = read(file, "I_r")
    times[1, :] = read(file, "times")
    i_last = read(file, "i_last")
    close(file)
    rho = Array{TA{TN, 3}, 1}(na)
    copy!(rho, rho_temp)
    dims = mpsdims(rho_temp)
    rho_temp = 0

    # temporary arrays for time evolution
    rho1 = similar(rho)
    rho2 = similar(rho)
    rho3 = similar(rho)
    rho1 .= copy.(rho)
    rho2 .= copy.(rho)
    rho3 .= copy.(rho)
    mpo_size = 6 + 2*length(lam)
    env1 = build_env(TN, TA, dims, dims)
    env2 = build_env(TN, TA, dims, dims)
    env3 = build_env(TN, TA, dims, dims)
    envop = build_env(TN, TA, dims, dims, ones(na + 1)*mpo_size)

    # time evolution operators for Runge Kutta algorithm
    L1 = construct_L_Ryd(TN, TA, 1, dt/2, t[i_last])
    L2 = construct_L_Ryd(TN, TA, 0, 1, t[i_last] + dt/2)
    L3 = construct_L_Ryd(TN, TA, 1, dt/2, t[i_last + 1])
    L = [L1, L2, L3]

    println("Initialisation complete")
    println(typeof(L))
    println(typeof(rho))
    println(typeof(env1))
    println(typeof(envop))

    # second time evolution
    for i = (i_last + 1):tstep
        time_in = time()

        # Runge-Kutta 4th order time step
        RK4stp_apply_H_half!(dt, L, rho, rho1, rho2, rho3, env1, env2, env3, envop)
        times[i,2] = time() - times[i,1] - time_in

        # measurements (every measure_int time steps)
        if rem(i, measure_int) == 0
            mind = div(i, measure_int) + 1
            tr_rho[mind] = scal_prod_no_conj(IDlmps, rho)
            I_r[mind] = scal_prod_no_conj(IRlmps,rho)
            times[i, 3] = time() - times[i, 2] - times[i, 1] - time_in
        end

        # saving temporary file
        if rem(i, measure_int*20) == 0
            write_data_file(string(g2_filename, "_temp.mat"), i, t_m, tr_rho,
                            I_r, rho, times)
        end

    end

    # saving final data to file
    write_data_file(string(g2_filename, ".mat"), tstep, t_m, tr_rho,
                            I_r, rho, times)

end

function write_data_file(filename, i, t_m, tr_rho, I_r, rho, times)

    file = matopen(filename, "w")
    write(file, "TN", string(TN))
    write(file, "TA", string(TA))
    write(file, "na", na)
    write(file, "gam_1d", gam_1d)
    write(file, "gam_eg", gam_eg)
    write(file, "del_p", del_p)
    write(file, "k_wg", k_wg)
    write(file, "rj", rj)
    write(file, "k_in", k_in)
    write(file, "f_amp", f_amp)
    write(file, "f", f.(t_m))
    write(file, "om", om)
    write(file, "k_c", k_c)
    write(file, "del_s", del_s)
    write(file, "uu", uu)
    write(file, "lam", lam)
    write(file, "gam_gg", gam_gg)
    write(file, "gam_ss", gam_ss)
    write(file, "gam_ee", gam_ee)
    write(file, "gam_sp", gam_sp)
    write(file, "gam_pg", gam_pg)
    write(file, "dt", dt)
    write(file, "t_fin", t_fin)
    write(file, "d_max", d_max)
    write(file, "measure_int", measure_int)

    write(file, "i_last", i)
    write(file, "t_m", collect(t_m))
    write(file, "tr_rho", tr_rho)
    write(file, "I_r", I_r)
    write(file, "rho", collect.(rho))
    write(file, "times", sum(times,1))
    close(file)

end

function construct_L_Ryd(::Type{TN}, ::Type{TA}, con, delt, time) where {TN, TA}

    LMPO = Array{TA{TN, 4}, 1}(na)
	drj = diff(rj)
	ph = exp.(im*k_wg*drj)
	cp = sqrt.(delt*gam_1d/2)
    no_exp = length(lam)
    max_ind = 6 + 2*no_exp

	LMPO[1] = zeros(TN, 1, d^2, max_ind, d^2)
	LMPO[1][1, :, 1, :] = kron(id, id)
	LMPO[1][1, :, 2, :] = cp[1]*ph[1]*(kron(id, hge) - kron(heg, id))
	LMPO[1][1, :, 3, :] = cp[1]*ph[1]*(kron(hge, id))
	LMPO[1][1, :, 4, :] = cp[1]'*ph[1]'*(kron(id, hge))
	LMPO[1][1, :, 5, :] = cp[1]'*ph[1]'*(kron(hge, id) - kron(id, heg))
    for exp_ind = 1:no_exp
        LMPO[1][1, :, 5 + 2*exp_ind - 1, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*
            (kron(hss, id) + kron(hpp, id))
        LMPO[1][1, :, 5 + 2*exp_ind, :] = sqrt(im*delt*uu[exp_ind]*lam[exp_ind])*
            (kron(id, hss) + kron(id, hpp))
    end
	LMPO[1][1, :, max_ind, :] = local_l(1, con, delt, time)

	LMPO[na] = zeros(TN, max_ind, d^2, 1, d^2)
	LMPO[na][1, :, 1, :] = local_l(na, con, delt, time)
	LMPO[na][2, :, 1, :] = cp[na]*kron(hge, id)
	LMPO[na][3, :, 1, :] = -cp[na]*(kron(heg, id) - kron(id, hge))
	LMPO[na][4, :, 1, :] = -cp[na]'*(kron(id, heg) - kron(hge, id))
	LMPO[na][5, :, 1, :] = cp[na]'*kron(id, hge)
    for exp_ind = 1:no_exp
        LMPO[na][5 + 2*exp_ind - 1, :, 1, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*
            (kron(hss, id) + kron(hpp, id))
        LMPO[na][5 + 2*exp_ind, :, 1, :] = sqrt(im*delt*uu[exp_ind]*lam[exp_ind])*
            (kron(id, hss) + kron(id, hpp))
    end
	LMPO[na][max_ind, :, 1, :] = kron(id, id)

    for jj = 2:(na - 1)
        LMPO[jj] = zeros(TN, max_ind, d^2, max_ind, d^2)
        LMPO[jj][1, :, 1, :] = kron(id, id)
        LMPO[jj][1, :, 2, :] = cp[jj]*ph[jj]*(kron(id, hge) - kron(heg, id))
        LMPO[jj][1, :, 3, :] = cp[jj]*ph[jj]*(kron(hge, id))
        LMPO[jj][1, :, 4, :] = cp[jj]'*ph[jj]'*(kron(id, hge))
        LMPO[jj][1, :, 5, :] = cp[jj]'*ph[jj]'*(kron(hge, id) - kron(id, heg))
        LMPO[jj][1, :, max_ind, :] = local_l(jj, con, delt, time)
        LMPO[jj][2, :, max_ind, :] = cp[jj]*kron(hge,id)
        LMPO[jj][3, :, max_ind, :] = -cp[jj]*(kron(heg, id) - kron(id, hge))
        LMPO[jj][4, :, max_ind, :] = -cp[jj]'*(kron(id, heg) - kron(hge, id))
        LMPO[jj][5, :, max_ind, :] = cp[jj]'*kron(id, hge)
        LMPO[jj][max_ind, :, max_ind, :] = kron(id, id)
        LMPO[jj][2, :, 2, :] = ph[jj]*kron(id, id)
        LMPO[jj][3, :, 3, :] = ph[jj]*kron(id, id)
        LMPO[jj][4, :, 4, :] = ph[jj]'*kron(id, id)
        LMPO[jj][5, :, 5, :] = ph[jj]'*kron(id, id)
        for exp_ind = 1:no_exp
            LMPO[jj][1, :, 5 + 2*exp_ind - 1, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*
                (kron(hss, id) + kron(hpp, id))
            LMPO[jj][1, :, 5 + 2*exp_ind, :] = sqrt(im*delt*uu[exp_ind]*lam[exp_ind])*
                (kron(id, hss) + kron(id, hpp))
            LMPO[jj][5 + 2*exp_ind - 1, :, max_ind, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*
                (kron(hss, id) + kron(hpp, id))
            LMPO[jj][5 + 2*exp_ind, :, max_ind, :] = sqrt(im*delt*uu[exp_ind]*lam[exp_ind])*
                (kron(id, hss) + kron(id, hpp))
            LMPO[jj][5 + 2*exp_ind - 1, :, 5 + 2*exp_ind - 1, :] = lam[exp_ind]*kron(id, id)
            LMPO[jj][5 + 2*exp_ind, :, 5 + 2*exp_ind, :] = lam[exp_ind]*kron(id, id)
        end
    end

    return LMPO

end

function update_L_Ryd!(LMPO, con, delt, time)

    no_exp = length(lam)
    max_ind = 6 + 2*no_exp
    LMPO[1][1,:,max_ind,:] = local_l(1, con, delt, time)
    LMPO[na][1,:,1,:] = local_l(na, con, delt, time)
    for jj = 2:na-1
        LMPO[jj][1,:,max_ind,:] = local_l(jj, con, delt, time)
    end

end

function local_l(jj, con, delt, time)

    con/na*kron(id, id) + delt*(
        (im*del_p - (gam_eg + gam_1d[jj] + gam_ee)/2)*kron(hee, id) -
        (im*del_p + (gam_eg + gam_1d[jj] + gam_ee)/2)*kron(id, hee) +
        (im*del_s - gam_ss/2 - gam_sp/2)*kron(hss, id) -
        (im*del_s + gam_ss/2 + gam_sp/2)*kron(id, hss) +
        (-gam_gg/2)*kron(hgg, id) -
        (+gam_gg/2)*kron(id, hgg) +
        (-gam_pg/2)*kron(hpp, id) -
        (+gam_pg/2)*kron(id, hpp) +
        gam_gg*kron(hgg, hgg) +
        gam_ss*kron(hss, hss) +
        gam_ee*kron(hee, hee) +
        gam_sp*kron(hps, hps) +
        gam_pg*kron(hgp, hgp) +
        (gam_eg + gam_1d[jj])*kron(hge, hge) +
        im*f(time)*sqrt(gam_1d[jj]/2)*exp(im*k_in*rj[jj])*(kron(heg, id) - kron(id, hge)) +
        im*f(time)*sqrt(gam_1d[jj]/2)*exp(-im*k_in*rj[jj])*(kron(hge, id) - kron(id, heg)) +
        im*om*exp(im*k_c*rj[jj])*(kron(hse, id) - kron(id, hes)) +
        im*om*exp(-im*k_c*rj[jj])*(kron(hes, id) - kron(id, hse)))

end

function measure_excitations!(rho, popj, IDlmps, locmat)

    imat = IDlmps[1][1, :, 1]
    na = length(rho)
    for jj = 1:na
        IDlmps[jj][1, :, 1] = locmat
        popj[jj] = scal_prod_no_conj(IDlmps, rho)
        IDlmps[jj][1, :, 1] = imat
    end
end

time_evolve()