text:
using Gen: generate_generative_function

@testset "static IR" begin

counter = 0

#@gen (static, nojuliacache) function bar(y, b)
    #u = begin counter += 1; y + b end
    #v = @trace(normal(0, 1), :v)
    #z = u + v
    #return z
#end

builder = StaticIRBuilder()
y = add_argument_node!(builder, name=:y)
b = add_argument_node!(builder, name=:b)
zero = add_constant_node!(builder, 0.)
one = add_constant_node!(builder, 1.)
u = add_julia_node!(builder, (y, b) -> begin counter += 1; y + b end, inputs=[y, b], name=:u)
v = add_addr_node!(builder, normal, inputs=[zero, one], addr=:v, name=:v)
z = add_julia_node!(builder, (u, v) -> u + v, inputs=[u, v], name=:z)
set_return_node!(builder, z)
ir = build_ir(builder)
bar = eval(generate_generative_function(ir, :bar, track_diffs=false, cache_julia_nodes=false))

#@gen (static, nojuliacache) function foo(a, b)
    #@param theta::Float64
    #x = @trace(normal(0, 1), :x)
    #y = @trace(normal(x, a), :y)
    #z = @trace(bar(y, b), :z)
    #w = z + 1 + a + theta
    #return w
#end

builder = StaticIRBuilder()
a = add_argument_node!(builder, name=:a)
b = add_argument_node!(builder, name=:b)
theta = add_trainable_param_node!(builder, :theta, typ=QuoteNode(Float64))
zero = add_constant_node!(builder, 0.)
one = add_constant_node!(builder, 1.)
x = add_addr_node!(builder, normal, inputs=[zero, one], addr=:x, name=:x)
y = add_addr_node!(builder, normal, inputs=[x, a], addr=:y, name=:y)
z = add_addr_node!(builder, bar, inputs=[y, b], addr=:z, name=:z)
w = add_julia_node!(builder, (z, a, theta) -> z + 1 + a + theta, inputs=[z, a, theta], name=:w)
set_return_node!(builder, w)
ir = build_ir(builder)
foo = eval(generate_generative_function(ir, :foo, track_diffs=false, cache_julia_nodes=false))

theta_val = rand()
set_param!(foo, :theta, theta_val)

#@gen (static, nojuliacache) function const_fn()
    #return 1
#end

builder = StaticIRBuilder()
one = add_constant_node!(builder, 2)
set_return_node!(builder, one)
ir = build_ir(builder)
const_fn = eval(generate_generative_function(ir, :const_fn, track_diffs=false, cache_julia_nodes=false))

Gen.load_generated_functions()

@testset "Julia call" begin
    @test const_fn() == 2
end

function expected_score(a, b, x, y, v)
    (logpdf(normal, x, 0, 1) +
        logpdf(normal, y, x, a) +
        logpdf(normal, v, 0, 1))
end

@testset "simulate" begin
    a, b = 1, 2
    trace = simulate(foo, (a, b))
    x = get_choices(trace)[:x]
    y = get_choices(trace)[:y]
    v = get_choices(trace)[:z => :v]
    z = y + b + v

    @test isapprox(expected_score(a, b, x, y, v), get_score(trace))

    expected_retval = z + 1 + a + theta_val
    @test isapprox(expected_retval, get_retval(trace))
end

@testset "generate" begin
    a, b = 1, 2
    y = 1.123
    v = 2.345
    constraints = choicemap((:y, y), (:z => :v, v))
    (trace, weight) = generate(foo, (a, b), constraints)
    x = trace[:x] # filled in using ancestral sampling
    z = y + b + v

    @test get_choices(trace)[:y] == y
    @test get_choices(trace)[:z => :v] == v

    @test isapprox(expected_score(a, b, x, y, v), get_score(trace))

    proposal_score = logpdf(normal, x, 0, 1)
    expected_weight = get_score(trace)  - proposal_score
    @test isapprox(weight, expected_weight)
end

@testset "project" begin
    a, b = 1, 2
    y = 1.123
    v = 2.345
    x = 3.456
    constraints = choicemap((:y, y), (:z => :v, v), (:x, x))
    (trace, weight) = generate(foo, (a, b), constraints)
    weight = project(trace, select(:y, :z => :v))
    expected_weight = logpdf(normal, y, x, a) + logpdf(normal, v, 0, 1)
    @test isapprox(weight, expected_weight)
end

@testset "update without tracked diffs" begin

    # generate initial trace
    a = 1
    b = 2
    trace = simulate(foo, (a, b))
    x_init = get_choices(trace)[:x]
    y_init = get_choices(trace)[:y]
    v_init = get_choices(trace)[:z => :v]
    w_init = get_retval(trace)

    # case 1: diff(a) is NoChange(), :y is not constrained, diff(b) is NoChange()
    counter = 0
    (new_trace, weight, retdiff, discard) = update(trace, (a, b), (NoChange(), NoChange()), choicemap())
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test get_retval(new_trace) == w_init
    @test isapprox(weight, 0.)
    @test retdiff == NoChange()
    @test isempty(discard)
    @test counter == 0
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_init, v_init))

    # case 2: diff(a) is NoChange(), no constraints, diff(b) is UnknownChange()
    a_new = 3
    counter = 0
    (new_trace, weight, retdiff, discard) = update(trace, (a_new, b), (UnknownChange(), NoChange()), choicemap())
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test isapprox(get_retval(new_trace), w_init + a_new - a)
    @test isapprox(weight, logpdf(normal, y_init, x_init, a_new) - logpdf(normal, y_init, x_init, a))
    @test retdiff == UnknownChange() # because w depends on a
    @test isempty(discard)
    @test counter == 0 # y blocks the dependnece on b on a
    @test isapprox(get_score(new_trace), expected_score(a_new, b, x_init, y_init, v_init))

    # case 3: diff(a) is UnknownChange(), no constraints, diff(b) is NoChange()
    counter = 0
    (new_trace, weight, retdiff, discard) = update(trace, (a, b), (NoChange(), UnknownChange()), choicemap())
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test get_retval(new_trace) == w_init
    @test isapprox(weight, 0.)
    @test retdiff == UnknownChange() # because w depends on z which depends on b
    @test isempty(discard)
    @test counter == 1 # because the call to bar depends on b
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_init, v_init))

    # case 4: diff(a) is NoChange(), :y is constrained, diff(b) is NoChange()
    counter = 0
    y_new = 1.123
    (new_trace, weight, retdiff, discard) = update(trace, (a, b), (NoChange(), NoChange()), choicemap((:y, y_new)))
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_new
    @test get_choices(new_trace)[:z => :v] == v_init
    @test isapprox(get_retval(new_trace), w_init + y_new - y_init)
    @test isapprox(weight, logpdf(normal, y_new, x_init, a) - logpdf(normal, y_init, x_init, a))
    @test retdiff == UnknownChange() # because the call to bar()
    @test discard[:y] == y_init
    @test counter == 1 # because the call to bar depends on y, which was constrained
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_new, v_init))

    # case 5: diff(a) is NoChange(), diff(b) is NoChange(), constrain :z => :v
    counter = 0
    v_new = 2.345
    (new_trace, weight, retdiff, discard) = update(trace, (a, b), (NoChange(), NoChange()), choicemap((:z => :v, v_new)))
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_new
    @test isapprox(get_retval(new_trace), w_init + v_new - v_init)
    @test isapprox(weight, logpdf(normal, v_new, 0, 1) - logpdf(normal, v_init, 0, 1))
    @test retdiff == UnknownChange() # because the call to bar()
    @test discard[:z => :v] == v_init
    @test counter == 1 # because the call to bar depends on y, which was constrained
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_init, v_new))
end

@testset "regenerate without tracked diffs" begin

    # generate initial trace
    a = 1
    b = 2
    trace = simulate(foo, (a, b))
    x_init = get_choices(trace)[:x]
    y_init = get_choices(trace)[:y]
    v_init = get_choices(trace)[:z => :v]
    w_init = get_retval(trace)

    # case 1: diff(a) is NoChange(), no selected choices, diff(b) is NoChange()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a, b), (NoChange(), NoChange()), select())
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test get_retval(new_trace) == w_init
    @test isapprox(weight, 0.)
    @test retdiff == NoChange()
    @test counter == 0
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_init, v_init))

    # case 2: diff(a) is NoChange(), no selected choices, diff(b) is UnknownChange()
    a_new = 3
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a_new, b), (UnknownChange(), NoChange()), select())
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test isapprox(get_retval(new_trace), w_init + a_new - a)
    @test isapprox(weight, logpdf(normal, y_init, x_init, a_new) - logpdf(normal, y_init, x_init, a))
    @test retdiff == UnknownChange() # because w depends on a
    @test counter == 0 # y blocks the dependnece on b on a
    @test isapprox(get_score(new_trace), expected_score(a_new, b, x_init, y_init, v_init))

    # case 3: diff(a) is UnknownChange(), no selected choices, diff(b) is NoChange()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a, b), (NoChange(), UnknownChange()), select())
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test get_retval(new_trace) == w_init
    @test isapprox(weight, 0.)
    @test retdiff == UnknownChange() # because w depends on z which depends on b
    @test counter == 1 # because the call to bar depends on b
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_init, v_init))

    # case 4: diff(a) is NoChange(), :y is selected, diff(b) is NoChange()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a, b), (NoChange(), NoChange()), select(:y))
    y_new = get_choices(new_trace)[:y]
    @test y_new != y_init
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:z => :v] == v_init
    @test isapprox(get_retval(new_trace), w_init + y_new - y_init)
    @test isapprox(weight, 0.)
    @test retdiff == UnknownChange() # because the call to bar()
    @test counter == 1 # because the call to bar depends on y, which was constrained
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_new, v_init))

    # case 5: diff(a) is NoChange(), diff(b) is NoChange(), select :z => :v
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a, b), (NoChange(), NoChange()), select(:z => :v))
    v_new = get_choices(new_trace)[:z => :v]
    @test v_new != v_init
    @test get_choices(new_trace)[:x] == x_init
    @test get_choices(new_trace)[:y] == y_init
    @test isapprox(get_retval(new_trace), w_init + v_new - v_init)
    @test isapprox(weight, 0.)
    @test retdiff == UnknownChange() # because the call to bar()
    @test counter == 1 # because the call to bar depends on y, which was constrained
    @test isapprox(get_score(new_trace), expected_score(a, b, x_init, y_init, v_new))

end

@testset "backprop" begin

    #@gen (static) function bar(mu_z::Float64)
        #z = @trace(normal(mu_z, 1), :z)
        #return z + mu_z
    #end

    # bar
    builder = StaticIRBuilder()
    mu_z = add_argument_node!(builder, name=:mu_z, typ=:Float64, compute_grad=true)
    one = add_constant_node!(builder, 1.)
    z = add_addr_node!(builder, normal, inputs=[mu_z, one], addr=:z, name=:z)
    retval = add_julia_node!(builder, (z, mu_z) -> z + mu_z, inputs=[z, mu_z], name=:retval)
    set_return_node!(builder, retval)
    ir = build_ir(builder)
    bar = eval(generate_generative_function(ir, :bar, track_diffs=false, cache_julia_nodes=false))

    #@gen (static) function foo(mu_a::Float64)
        #param theta::Float64
        #a = @trace(normal(mu_a, 1), :a)
        #b = @trace(normal(a, 1), :b)
        #bar = @trace(bar(a), :bar)
        #c = a * b * bar * theta
        #out = @trace(normal(c, 1), :out)
        #return out
    #end

    # foo
    builder = StaticIRBuilder()
    mu_a = add_argument_node!(builder, name=:mu_a, typ=:Float64, compute_grad=true)
    theta = add_trainable_param_node!(builder, :theta, typ=QuoteNode(Float64))
    one = add_constant_node!(builder, 1.)
    a = add_addr_node!(builder, normal, inputs=[mu_a, one], addr=:a, name=:a)
    b = add_addr_node!(builder, normal, inputs=[a, one], addr=:b, name=:b)
    bar_val = add_addr_node!(builder, bar, inputs=[a], addr=:bar, name=:bar_val)
    c = add_julia_node!(builder, (a, b, bar, theta) -> (a * b * bar * theta),
            inputs=[a, b, bar_val, theta], name=:c)
    retval = add_addr_node!(builder, normal, inputs=[c, one], addr=:out, name=:out)
    set_return_node!(builder, retval)
    ir = build_ir(builder)
    foo = eval(generate_generative_function(ir, :foo, track_diffs=false, cache_julia_nodes=false))

    Gen.load_generated_functions()

    function f(mu_a, theta, a, b, z, out)
        lpdf = 0.
        mu_z = a
        lpdf += logpdf(normal, z, mu_z, 1)
        lpdf += logpdf(normal, a, mu_a, 1)
        lpdf += logpdf(normal, b, a, 1)
        c = a * b * (z + mu_z) * theta
        lpdf += logpdf(normal, out, c, 1)
        return lpdf + 2 * out
    end

    mu_a = 1.
    theta = -0.5
    a = 2.
    b = 3.
    z = 4.
    out = 5.

    init_param!(foo, :theta, theta)

    @test get_param(foo, :theta) == theta
    @test get_param_grad(foo, :theta) == 0.

    # get the initial trace
    constraints = choicemap()
    constraints[:a] = a
    constraints[:b] = b
    constraints[:out] = out
    constraints[:bar => :z] = z
    (trace, _) = generate(foo, (mu_a,), constraints)

    # compute gradients with choice_gradients
    selection = select(:bar => :z, :a, :out)
    selection = StaticSelection(selection)
    retval_grad = 2.
    ((mu_a_grad,), value_trie, gradient_trie) = choice_gradients(trace, selection, retval_grad)

    # check input gradient
    @test isapprox(mu_a_grad, finite_diff(f, (mu_a, theta, a, b, z, out), 1, dx))

    # check value trie
    @test get_value(value_trie, :a) == a
    @test get_value(value_trie, :out) == out
    @test get_value(value_trie, :bar => :z) == z
    @test !has_value(value_trie, :b) # was not selected
    @test length(get_submaps_shallow(value_trie)) == 1
    @test length(get_values_shallow(value_trie)) == 2

    # check gradient trie
    @test length(get_submaps_shallow(gradient_trie)) == 1
    @test length(get_values_shallow(gradient_trie)) == 2
    @test !has_value(gradient_trie, :b) # was not selected
    @test isapprox(get_value(gradient_trie, :a), finite_diff(f, (mu_a, theta, a, b, z, out), 3, dx))
    @test isapprox(get_value(gradient_trie, :out), finite_diff(f, (mu_a, theta, a, b, z, out), 6, dx))
    @test isapprox(get_value(gradient_trie, :bar => :z), finite_diff(f, (mu_a, theta, a, b, z, out), 5, dx))

    # reset the trainable parameter gradient
    zero_param_grad!(foo, :theta)
    @test get_param(foo, :theta) == theta
    @test get_param_grad(foo, :theta) == 0.

    # compute gradients with accumulate_param_gradients!
    retval_grad = 2.
    (mu_a_grad,) = accumulate_param_gradients!(trace, retval_grad)

    # check input gradient
    @test isapprox(mu_a_grad, finite_diff(f, (mu_a, theta, a, b, z, out), 1, dx))

    # check trainable parameter gradient
    @test isapprox(get_param_grad(foo, :theta), finite_diff(f, (mu_a, theta, a, b, z, out), 2, dx))

end

# functions to test tracked diffs

# NOTE: this function will always return NoChange() for its retdiff

counter = 0

#@gen (static) function bar(a)
    #x = begin counter += 1; 0 end
    #return x
#end

builder = StaticIRBuilder()
a = add_argument_node!(builder, name=:a)
x = add_julia_node!(builder, () -> begin counter += 1; 0 end, inputs=[], name=:x)
set_return_node!(builder, x)
ir = build_ir(builder)
bar = eval(generate_generative_function(ir, :bar, track_diffs=false, cache_julia_nodes=false))

# this function will dynamically obtain NoChange() from the retdiff of bar(),
# which will cause its own retdiff to be NoChange():

#@gen (static) function foo(a)
    #x = @trace(bar(a), :x)
    #y = x + 1
    #return y
#end

builder = StaticIRBuilder()
a = add_argument_node!(builder, name=:a)
x = add_addr_node!(builder, bar, inputs=[a], addr=:x, name=:x)
y = add_julia_node!(builder, (x) -> x + 1, inputs=[x], name=:y)
set_return_node!(builder, y)
ir = build_ir(builder)

# generate a version of the function with tracked diffs
foo = eval(generate_generative_function(ir, :foo, track_diffs=true, cache_julia_nodes=false))

# generate a version of the function without tracked diffs
foo_without_tracked_diffs = eval(generate_generative_function(ir, :foo, track_diffs=false, cache_julia_nodes=false))

Gen.load_generated_functions()

@testset "update with tracked diffs" begin

    # generate initial trace from function with tracked diffs
    a = 1
    trace = simulate(foo, (a,))

    # case 1: diff(a) is NoChange, it will statically avoid running bar()
    counter = 0
    (new_trace, weight, retdiff, _) = update(trace, (a,), (NoChange(),), choicemap())
    @test counter == 0
    @test retdiff == NoChange()

    # case 2: diff(a) is UnknownChange, it will run bar(), but still return retdiff of NoChange()
    counter = 0
    (new_trace, weight, retdiff, _) = update(trace, (a,), (UnknownChange(),), choicemap())
    @test counter == 1
    @test retdiff == NoChange()

    # generate initial trace from function without tracked diffs
    a = 1
    trace = simulate(foo_without_tracked_diffs, (a,))

    # case 1: diff(a) is NoChange, it will statically avoid running bar()
    counter = 0
    (new_trace, weight, retdiff, _) = update(trace, (a,), (NoChange(),), choicemap())
    @test counter == 0
    @test retdiff == NoChange()

    # case 2: diff(a) is UnknownChange, it will run bar()
    counter = 0
    (new_trace, weight, retdiff, _) = update(trace, (a,), (UnknownChange(),), choicemap())
    @test counter == 1
    @test retdiff == UnknownChange()

end

@testset "regenerate with tracked diffs" begin

    # generate initial trace from function with tracked diffs
    a = 1
    trace = simulate(foo, (a,))

    # case 1: diff(a) is NoChange, it will statically avoid running bar()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a,), (NoChange(),), select())
    @test counter == 0
    @test retdiff == NoChange()

    # case 2: diff(a) is UnknownChange, it will run bar(), but still return retdiff of NoChange()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a,), (UnknownChange(),), select())
    @test counter == 1
    @test retdiff == NoChange()

    # generate initial trace from function without tracked diffs
    a = 1
    trace = simulate(foo_without_tracked_diffs, (a,))

    # case 1: diff(a) is NoChange, it will statically avoid running bar()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a,), (NoChange(),), select())
    @test counter == 0
    @test retdiff == NoChange()

    # case 2: diff(a) is UnknownChange, it will run bar()
    counter = 0
    (new_trace, weight, retdiff) = regenerate(trace, (a,), (UnknownChange(),), select())
    @test counter == 1
    @test retdiff == UnknownChange()
end

counter = 0

#@gen (static) function foo(a, b)
    #u = begin a + b; counter += 1 end
    #x = @trace(normal(u, 1), :x)
    #return x
#end

builder = StaticIRBuilder()
a = add_argument_node!(builder, name=:a)
b = add_argument_node!(builder, name=:b)
u = add_julia_node!(builder, (a, b) -> begin counter += 1; a + b end, inputs=[a, b], name=:u)
one = add_constant_node!(builder, 1.)
x = add_addr_node!(builder, normal, inputs=[u, one], addr=:x, name=:x)
set_return_node!(builder, x)
ir = build_ir(builder)
foo = eval(generate_generative_function(ir, :foo, track_diffs=false, cache_julia_nodes=true))

Gen.load_generated_functions()

@testset "cached julia nodes" begin

    counter = 0
    trace = simulate(foo, (1, 2))
    @test counter == 1

    counter = 0
    update(trace, (1, 2), (NoChange(), NoChange()), choicemap((:x, 1.0)))
    @test counter == 0 # would be 1 if nojuliacache was used

    counter = 0
    update(trace, (1, 2), (NoChange(), NoChange()), choicemap((:x, 1.0)))
    @test counter == 0 # would be 1 if nojuliacache was used

    counter = 0
    regenerate(trace, (1, 2), (NoChange(), NoChange()), select(:x))
    @test counter == 0 # would be 1 if nojuliacache was used

    counter = 0
    update(trace, (1, 3), (NoChange(), UnknownChange()), choicemap((:x, 1.0)))
    @test counter == 1

    counter = 0
    update(trace, (1, 3), (NoChange(), UnknownChange()), choicemap((:x, 1.0)))
    @test counter == 1

    counter = 0
    regenerate(trace, (1, 3), (NoChange(), UnknownChange()), select(:x))
    @test counter == 1
end

@testset "regression test for https://github.com/probcomp/Gen/issues/168" begin
    @gen (static) function model(var)
        mean = @trace(normal(0, 1), :mean)
        T = @trace(normal(mean, var), :T)
        return T
    end
    load_generated_functions()
    selection = StaticSelection(select(:mean))
    (tr, _) = generate(model, (1,))
    # At the time the issue was filed, this line produced a crash
    (tr, ) = mh(tr, selection)
end

end # @testset "static IR"


text:
#=
println("hither : ",Threads.threadid()," : ",pwd()) # hither : 1 : /tmp
Threads.@spawn cd( () -> ( println("tither : ",Threads.threadid(), " : ", pwd()); sleep(3)), "..") # tither : 2 : /
sleep(1);println("hither : ",Threads.threadid(), " : ", pwd());sleep(3) # hither : 1 : /
println("hither : ",Threads.threadid()," : ",pwd()) # hither : 1 : /tmp
=#

cd(@__DIR__)

println("hither : ",Threads.threadid()," : ",pwd()) # hither : 1 : ./CmdStan/test

Threads.@spawn cd( () -> ( println("tither : ",Threads.threadid(), " : ",
  pwd()); sleep(3)), "..") # tither : 2 : ./CmdStan

sleep(1)

println("hither : ",Threads.threadid()," : ", pwd())  # hither : 1 : ./CmdStan

sleep(3)

println("hither : ",Threads.threadid()," : ", pwd()) # hither : 1 : ./CmdStan/test

ext: jl


text:
struct KeyEntry
    keyid::UInt16
    tiff_tag_location::UInt16
    count::UInt16
    value_offset::UInt16
end

struct GeoKeys
    key_directory_version::UInt16
    key_reversion::UInt16
    minor_revision::UInt16
    number_of_keys::UInt16
    keys::Vector{KeyEntry}
end

struct SRID
    authority::Symbol
    code::Int64
end

# version numbers are fixed in the LAS specification
GeoKeys(keys::Vector{KeyEntry}) = GeoKeys(0x0001, 0x0001, 0x0000, UInt16(length(keys)), keys)

struct GeoDoubleParamsTag
    double_params::Vector{Float64}
end

struct GeoAsciiParamsTag
    ascii_params::String
    nb::Int  # number of bytes
    GeoAsciiParamsTag(s::AbstractString, nb::Integer) = new(ascii(s), Int(nb))
end

const id_geokeydirectorytag = UInt16(34735)
const id_geodoubleparamstag = UInt16(34736)
const id_geoasciiparamstag = UInt16(34737)

"test whether a vlr is a GeoKeyDirectoryTag, GeoDoubleParamsTag or GeoAsciiParamsTag"
is_srs(vlr::LasVariableLengthRecord) = vlr.record_id in (
    id_geokeydirectorytag,
    id_geodoubleparamstag,
    id_geoasciiparamstag)

# number of bytes
Base.sizeof(data::GeoKeys) = 8 * Int(data.number_of_keys) + 8
Base.sizeof(data::GeoDoubleParamsTag) = sizeof(data.double_params)
Base.sizeof(data::GeoAsciiParamsTag) = data.nb

"Construct a projection VLR based on an EPSG code"
function LasVariableLengthRecord(header::LasHeader, srid::SRID)
    if srid.authority != :epsg
        throw(ArgumentError("No other code than EPSG-code is implemented"))
    elseif srid.code < 0
        throw(ArgumentError("EPSG is not valid: negative integer input given"))
    end

    reserved = 0xAABB
    user_id = "LASF_Projection"
    description = "GeoTIFF GeoKeyDirectoryTag"
    record_id = id_geokeydirectorytag

    data = GeoKeys(srid.code)

    return LasVariableLengthRecord(
        reserved,
        user_id,
        record_id,
        description,
        data
    )
end

function Base.write(io::IO, data::GeoKeys)
    write(io, data.key_directory_version)
    write(io, data.key_reversion)
    write(io, data.minor_revision)
    write(io, data.number_of_keys)
    for keyEntry in data.keys
        write_key_entry(io, keyEntry)
    end
end

Base.write(io::IO, data::GeoDoubleParamsTag) = write(io, data.double_params)
Base.write(io::IO, data::GeoAsciiParamsTag) = writestring(io, data.ascii_params, data.nb)

"Create GeoKeys from EPSG code. Assumes CRS is projected and in meters."
function GeoKeys(epsg::Integer)
    #Standard types
    is_projected = KeyEntry(UInt16(1024), UInt16(0), UInt16(1), UInt16(1))         # Projected
    proj_linear_units = KeyEntry(UInt16(1025), UInt16(0), UInt16(1), UInt16(1))    # Units in meter
    projected_cs_type = KeyEntry(UInt16(3072), UInt16(0), UInt16(1), UInt16(epsg)) # EPSG code
    vertical_units = KeyEntry(UInt16(3076), UInt16(0), UInt16(1), UInt16(9001))    # Units in meter
    keys = [is_projected, proj_linear_units, projected_cs_type, vertical_units]
    GeoKeys(keys)
end

function write_key_entry(io, entry::KeyEntry)
    write(io, entry.keyid)
    write(io, entry.tiff_tag_location)
    write(io, entry.count)
    write(io, entry.value_offset)
end

"Get the EPSG code of the projection in the header"
function epsg_code(header::LasHeader)
    if is_wkt(header)
        throw(ArgumentError("WKT format projection information not implemented"))
    end
    vlrs = header.variable_length_records
    ind = findfirst(x -> x.record_id == id_geokeydirectorytag, vlrs)
    if ind === nothing
        nothing
    else
        vlrs[ind].data.keys[3].value_offset
    end
end

"Set the projection in the header, without altering the points"
function epsg_code!(header::LasHeader, epsg::Integer)
    # small check if epsg is valid
    if epsg < 0
        throw(ArgumentError("EPSG is not valid: negative integer input given"))
    elseif is_wkt(header)
        throw(ArgumentError("Setting the SRS in WKT format not implemented"))
    end

    # read old header metadata
    old_vlrlength = header.n_vlr == 0 ? 0 : sum(sizeof, header.variable_length_records)
    old_offset = header.data_offset

    # reconstruct VLRs
    vlrs = LasVariableLengthRecord[]
    srid = SRID(:epsg,epsg)
    push!(vlrs, LasVariableLengthRecord(header, srid))
    # keep existing non-SRS VLRs intact
    append!(vlrs, filter(!is_srs, header.variable_length_records))

    # update header
    header.variable_length_records = vlrs
    header.n_vlr = length(header.variable_length_records)
    new_vlrlength = header.n_vlr == 0 ? 0 : sum(sizeof, header.variable_length_records)
    # update offset to point data, assuming the VLRs come before the data, i.e. not extended VLR
    header.data_offset = old_offset - old_vlrlength + new_vlrlength
    header
end

function read_vlr_data(io::IO, record_id::Integer, nb::Integer)
    if record_id == id_geokeydirectorytag
        return read(io, GeoKeys)
    elseif record_id == id_geodoubleparamstag
        double_params = zeros(nb ÷ 8)
        read!(io, double_params)
        return GeoDoubleParamsTag(double_params)
    elseif record_id == id_geoasciiparamstag
        ascii_params = readstring(io, nb)
        return GeoAsciiParamsTag(ascii_params, nb)
    else
        return read(io, nb)
    end
end

function Base.read(io::IO, ::Type{GeoKeys})
    key_directory_version = read(io, UInt16)
    key_reversion = read(io, UInt16)
    minor_revision = read(io, UInt16)
    number_of_keys = read(io, UInt16)
    keys = KeyEntry[]
    for i in 1:number_of_keys
        keyid = read(io, UInt16)
        tiff_tag_location = read(io, UInt16)
        count = read(io, UInt16)
        value_offset = read(io, UInt16)
        push!(keys, KeyEntry(
            keyid,
            tiff_tag_location,
            count,
            value_offset
        ))
    end
    return GeoKeys(
        key_directory_version,
        key_reversion,
        minor_revision,
        number_of_keys,
        keys
    )
end


text:
"""
    struct ModelGen{G, defaultnames, Tdefaults}
        generator::G
        defaults::Tdefaults
    end

A `ModelGen` struct with model generator function of type `G`, and default arguments `defaultnames`
with values `Tdefaults`.
"""
struct ModelGen{G, argnames, defaultnames, Tdefaults}
    generator::G
    defaults::NamedTuple{defaultnames, Tdefaults}

    function ModelGen{argnames}(
        generator::G,
        defaults::NamedTuple{defaultnames, Tdefaults}
    ) where {G, argnames, defaultnames, Tdefaults}
        return new{G, argnames, defaultnames, Tdefaults}(generator, defaults)
    end
end

(m::ModelGen)(args...; kwargs...) = m.generator(args...; kwargs...)

"""
    getdefaults(modelgen::ModelGen)

Get a named tuple of the default argument values defined for a model defined by a generating function.
"""
getdefaults(modelgen::ModelGen) = modelgen.defaults

"""
    getargnames(modelgen::ModelGen)

Get a tuple of the argument names of the `modelgen`.
"""
getargnames(model::ModelGen{_G, argnames}) where {argnames, _G} = argnames

"""
    struct Model{F, argnames, Targs, missings}
        f::F
        args::NamedTuple{argnames, Targs}
        modelgen::Tgen
    end

A `Model` struct with model evaluation function of type `F`, arguments names `argnames`, arguments
types `Targs`, missing arguments `missings`, and corresponding model generator. `argnames` and
`missings` are tuples of symbols, e.g. `(:a, :b)`.  An argument with a type of `Missing` will be in
`missings` by default.  However, in non-traditional use-cases `missings` can be defined differently.
All variables in `missings` are treated as random variables rather than observations.

# Example

```julia
julia> Model(f, (x = 1.0, y = 2.0))
Model{typeof(f),(),(:x, :y),Tuple{Float64,Float64}}((x = 1.0, y = 2.0))

julia> Model{(:y,)}(f, (x = 1.0, y = 2.0))
Model{typeof(f),(:y,),(:x, :y),Tuple{Float64,Float64}}((x = 1.0, y = 2.0))
```
"""
struct Model{F, argnames, Targs, missings, Tgen} <: AbstractModel
    f::F
    args::NamedTuple{argnames, Targs}
    modelgen::Tgen

    """
        Model{missings}(f, args::NamedTuple, modelgen::ModelGen)

    Create a model with evalutation function `f` and missing arguments overwritten by `missings`.
    """
    function Model{missings}(
        f::F,
        args::NamedTuple{argnames, Targs},
        modelgen::Tgen
    ) where {missings, F, argnames, Targs, Tgen<:ModelGen}
        return new{F, argnames, Targs, missings, Tgen}(f, args, modelgen)
    end
end

"""
    Model(f, args::NamedTuple, modelgen::ModelGen)

    Create a model with evalutation function `f` and missing arguments deduced from `args`.
"""
@generated function Model(
    f::F,
    args::NamedTuple{argnames, Targs},
    modelgen::ModelGen{_G, argnames}
) where {F, argnames, Targs, _G}
    missings = Tuple(name for (name, typ) in zip(argnames, Targs.types) if typ <: Missing)
    return :(Model{$missings}(f, args, modelgen))
end

"""
    Model{missings}(modelgen::ModelGen, args::NamedTuple)

Create a copy of the model described by `modelgen(args...)`, with missing arguments
overwritten by `missings`.
"""
function Model{missings}(
    modelgen::ModelGen,
    args::NamedTuple{argnames, Targs}
) where {missings, argnames, Targs}
    model = modelgen(args...)
    return Model{missings}(model.f, args, modelgen)
end

"""
    (model::Model)([spl = SampleFromPrior(), ctx = DefaultContext()])

Sample from `model` using the sampler `spl`.
"""
function (model::Model)(
    spl::AbstractSampler=SampleFromPrior(),
    ctx::AbstractContext=DefaultContext()
)
    return model(VarInfo(), spl, ctx)
end

"""
    (model::Model)(vi::AbstractVarInfo[, spl = SampleFromPrior(), ctx = DefaultContext()])

Sample from `model` using the sampler `spl` storing the sample and log joint probability in `vi`.
Resets the `vi` and increases `spl`s `state.eval_num`.
"""
function (model::Model)(
    vi::AbstractVarInfo,
    spl::AbstractSampler=SampleFromPrior(),
    ctx::AbstractContext=DefaultContext()
)
    if Threads.nthreads() == 1
        return evaluate_singlethreaded(model, vi, spl, ctx)
    else
        return evaluate_multithreaded(model, vi, spl, ctx)
    end
end

function evaluate_singlethreaded(model, varinfo, sampler, context)
    resetlogp!(varinfo)
    if has_eval_num(sampler)
        sampler.state.eval_num += 1
    end
    return model.f(model, varinfo, sampler, context)
end

function evaluate_multithreaded(model, varinfo, sampler, context)
    resetlogp!(varinfo)
    if has_eval_num(sampler)
        sampler.state.eval_num += 1
    end
    wrapper = ThreadSafeVarInfo(varinfo)
    result = model.f(model, wrapper, sampler, context)
    setlogp!(varinfo, getlogp(wrapper))
    return result
end

"""
    getargnames(model::Model)

Get a tuple of the argument names of the `model`.
"""
getargnames(model::Model{_F, argnames}) where {argnames, _F} = argnames

"""
    getmissings(model::Model)

Get a tuple of the names of the missing arguments of the `model`.
"""
getmissings(model::Model{_F, _a, _T, missings}) where {missings, _F, _a, _T} = missings

getmissing(model::Model) = getmissings(model)
@deprecate getmissing(model) getmissings(model)

"""
    getgenerator(model::Model)

Get the model generator associated with `model`.
"""
getgenerator(model::Model) = model.modelgen


text:
abstract type AbstractAffineMap <: JalerkinType end

# """
# Data type for 1-d affine mappings, $F:[0,1] -> [x0,x1]$,
# defined as: $F(t)=a\cdot t + b$, $a, b \in \mathbb R$.
# """
struct AffineMap1D <: AbstractAffineMap
    "Multiplicative coefficient"
    a :: fp_precision()
    "Translation coefficient"
    b :: fp_precision()
    "Determinant of the jacobian (=a, in 1D)"
    det_jacobian :: fp_precision()
end

#--- Finite Element

abstract type FE_Family <: JalerkinType end
abstract type Abstract_FE <: JalerkinType end

"""
Lagrange family of finite elements
"""
struct Lagrange <: FE_Family end;

"""
Finite element (of a given type) on a mesh
"""
struct FiniteElement{Mesh <: AbstractMesh, FE_Fam <: FE_Family, order}
    mesh :: Mesh
    quad_rule :: AbstractQuadrature
end

"""
Constructor of a FiniteElement from a mesh, a FE family and a polinomial order"

A Gauss Quadrature rule is chosen such that the mass matrix is computed exactly
"""
function FiniteElement(mesh::Mesh, ::Type{FE_Fam}, order;
    quad_rule=GaussianQuadrature(2*order)) where {
        Mesh <: AbstractMesh,
        FE_Fam <: FE_Family }

    FiniteElement{Mesh, FE_Fam, order}(mesh, quad_rule)
end

"""
Return the mesh of a Finite Element
"""
@inline get_mesh(fe::FiniteElement) = fe.mesh

"""
Return the quadrature rule of a Finite Element
"""
@inline get_quad_rule(fe::FiniteElement) = fe.quad_rule

"""
Return the polynomial order of a Finite Element
"""
@inline get_order(fe::FiniteElement{Mesh, FE_T, order}) where {
    Mesh, FE_T, order} = order

max_stars_repo_name: rrgalvan/Jalerkin

max_issues_repo_name: rrgalvan/Jalerkin

max_forks_repo_name: rrgalvan/Jalerkin


text:
# Tue 13 Nov 2012 04:13:36 PM EST
#
# NeHe Tut 10 - Move around in a 3D world
#
# Q - quit
# B - turn texture alpha-blending on/off
# L - turn lights on/off
# F - change texture filter (linear, nearest, mipmap)
# PageUp/Down - look up/down
# Up/Down - move forward/backward
# Left/Right - turn left/right

# load necessary GLUT/OpenGL routines

using OpenGL
@OpenGL.version "1.0"
@OpenGL.load
using GLUT

### auxiliary functions

function SetupWorld(world_map::String)

    global numtriangles

    filein       = open(world_map)
    world_data   = readlines(filein)

    numtriangles = parseint(chomp(split(world_data[1],' ')[2]))

    sector       = zeros(numtriangles,3,5)

    loop = 1
    vert = 1
    line = 1

    while line <= length(world_data)-2
        if world_data[2+line][1] != '/' && world_data[2+line][1] != '\n'
            while vert <= 3
                (x, y, z, u, v)      = split(chomp(world_data[2+line]),' ')
                x                    = parsefloat(x)
                y                    = parsefloat(y)
                z                    = parsefloat(z)
                u                    = parsefloat(u)
                v                    = parsefloat(v)
                sector[loop,vert,:]  = [x,y,z,u,v]
                vert                 += 1
                line                 += 1
            end
            vert = 1
            loop += 1
        else
            line += 1
        end
    end

    return sector

end

### end of auxiliary functions

# initialize variables

global window

global numtriangles  = 0

global walkbias      = 0.0
global walkbiasangle = 0.0

global lookupdown    = 0.0

global xpos          = 0.0
global zpos          = 0.0

global yrot          = 0.0

global LightAmbient  = [0.5f0, 0.5f0, 0.5f0, 1.0f0]
global LightDiffuse  = [1.0f0, 1.0f0, 1.0f0, 1.0f0]
global LightPosition = [0.0f0, 0.0f0, 2.0f0, 1.0f0]

global filter        = 3
global light         = true
global blend         = false

global xtrans        = 0.0
global ytrans        = 0.0
global ztrans        = 0.0
global sceneroty     = 0.0

global tex           = Array(Uint32,3) # generating 3 textures

width                = 640
height               = 480

# initialize sector1 with SetupWorld

sector1 = SetupWorld(expanduser("~/.julia/GLUT/Examples/NeHe/tut10/world.txt"))

# load textures from images

function LoadGLTextures()
    global tex

    img, w, h = glimread(expanduser("~/.julia/GLUT/Examples/NeHe/tut10/mud.bmp"))

    glGenTextures(3,tex)
    glBindTexture(GL_TEXTURE_2D,tex[1])
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
    glTexImage2D(GL_TEXTURE_2D, 0, 3, w, h, 0, GL_RGB, GL_UNSIGNED_BYTE, img)

    glBindTexture(GL_TEXTURE_2D,tex[2])
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexImage2D(GL_TEXTURE_2D, 0, 3, w, h, 0, GL_RGB, GL_UNSIGNED_BYTE, img)

    glBindTexture(GL_TEXTURE_2D,tex[3])
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_NEAREST)
    glTexImage2D(GL_TEXTURE_2D, 0, 3, w, h, 0, GL_RGB, GL_UNSIGNED_BYTE, img)

    gluBuild2DMipmaps(GL_TEXTURE_2D, 3, w, h, GL_RGB, GL_UNSIGNED_BYTE, img)
end

# function to init OpenGL context

function initGL(w::Integer,h::Integer)
    glViewport(0,0,w,h)
    LoadGLTextures()
    glClearColor(0.0, 0.0, 0.0, 0.0)
    glClearDepth(1.0)
    glDepthFunc(GL_LESS)
    glEnable(GL_DEPTH_TEST)
    glShadeModel(GL_SMOOTH)

    # initialize lights
    glLightfv(GL_LIGHT1, GL_AMBIENT, LightAmbient)
    glLightfv(GL_LIGHT1, GL_DIFFUSE, LightDiffuse)
    glLightfv(GL_LIGHT1, GL_POSITION, LightPosition)

    glEnable(GL_LIGHT1)
    glEnable(GL_LIGHTING)

    # enable texture mapping and alpha blending
    glEnable(GL_TEXTURE_2D)
    glBlendFunc(GL_SRC_ALPHA, GL_ONE)

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()

    gluPerspective(45.0,w/h,0.1,100.0)

    glMatrixMode(GL_MODELVIEW)
end

# prepare Julia equivalents of C callbacks that are typically used in GLUT code

function ReSizeGLScene(w::Int32,h::Int32)
    if h == 0
        h = 1
    end

    glViewport(0,0,w,h)

    glMatrixMode(GL_PROJECTION)
    glLoadIdentity()

    gluPerspective(45.0,w/h,0.1,100.0)

    glMatrixMode(GL_MODELVIEW)

    return nothing
end

_ReSizeGLScene = cfunction(ReSizeGLScene, Void, (Int32, Int32))

function DrawGLScene()
    global xtrans
    global ytrans
    global ztrans
    global xpos
    global ypos
    global walkbias
    global sceneroty
    global yrot
    global lookupdown
    global tex
    global numtriangles
    global sector1

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
    glLoadIdentity()

    xtrans = -xpos
    ztrans = -zpos
    ytrans = -walkbias-0.25
    sceneroty = 360.0-yrot

    glRotate(lookupdown, 1.0, 0.0, 0.0)
    glRotate(sceneroty, 0.0, 1.0, 0.0)
    glTranslate(xtrans, ytrans, ztrans)

    glBindTexture(GL_TEXTURE_2D,tex[filter])

    for face = 1:numtriangles
        glBegin(GL_TRIANGLES)
            glNormal(0.0, 0.0, 1.0)
            x_m = sector1[face,1,1]
            y_m = sector1[face,1,2]
            z_m = sector1[face,1,3]
            u_m = sector1[face,1,4]
            v_m = sector1[face,1,5]
            glTexCoord(u_m,v_m)
            glVertex(x_m,y_m,z_m)

            x_m = sector1[face,2,1]
            y_m = sector1[face,2,2]
            z_m = sector1[face,2,3]
            u_m = sector1[face,2,4]
            v_m = sector1[face,2,5]
            glTexCoord(u_m,v_m)
            glVertex(x_m,y_m,z_m)

            x_m = sector1[face,3,1]
            y_m = sector1[face,3,2]
            z_m = sector1[face,3,3]
            u_m = sector1[face,3,4]
            v_m = sector1[face,3,5]
            glTexCoord(u_m,v_m)
            glVertex(x_m,y_m,z_m)
        glEnd()
    end

    glutSwapBuffers()

    return nothing
end

_DrawGLScene = cfunction(DrawGLScene, Void, ())

function keyPressed(the_key::Char,x::Int32,y::Int32)
    global blend
    global light
    global filter

    if the_key == int('q')
        glutDestroyWindow(window)
    elseif the_key == int('b')
        println("Blend was: $blend")
        blend = (blend ? false : true)
        if blend
            glEnable(GL_BLEND)
            glDisable(GL_DEPTH_TEST)
        else
            glDisable(GL_BLEND)
            glEnable(GL_DEPTH_TEST)
        end
        println("Blend is now: $blend")
    elseif the_key == int('l')
        println("Light was: $light")
        light = (light ? false : true)
        println("Light is now: $light")
        if light
            glEnable(GL_LIGHTING)
        else
            glDisable(GL_LIGHTING)
        end
    elseif the_key == int('f')
        println("Filter was: $filter")
        filter += 1
        if filter > 3
            filter = 1
        end
        println("Filter is now: $filter")
    end

    return nothing # keyPressed returns "void" in C. this is a workaround for Julia's "automatically return the value of the last expression in a function" behavior.
end

_keyPressed = cfunction(keyPressed, Void, (Char, Int32, Int32))

function specialKeyPressed(the_key::Int32,x::Int32,y::Int32)
    global lookupdown
    global xpos
    global zpos
    global walkbias
    global walkbiasangle
    global yrot

    if the_key == GLUT_KEY_PAGE_UP
        lookupdown -=0.2
    elseif the_key == GLUT_KEY_PAGE_DOWN
        lookupdown +=1.0
    elseif the_key == GLUT_KEY_UP
        xpos -=sin(degrees2radians(yrot))*0.05
        zpos -=cos(degrees2radians(yrot))*0.05
        walkbias +=10
        if walkbiasangle <= 359.0
            walkbiasangle = 0.0
        else
            walkbiasangle +=10
        end
        walkbias = sin(degrees2radians(walkbiasangle))/20.0
    elseif the_key == GLUT_KEY_DOWN
        xpos +=sin(degrees2radians(yrot))*0.05
        zpos +=cos(degrees2radians(yrot))*0.05
        walkbias -=10
        if walkbiasangle <= 1.0
            walkbiasangle = 359.0
        else
            walkbiasangle -=10
        end
        walkbias = sin(degrees2radians(walkbiasangle))/20.0
    elseif the_key == GLUT_KEY_LEFT
        yrot +=1.5
    elseif the_key == GLUT_KEY_RIGHT
        yrot -=1.5
    end

    return nothing # specialKeyPressed returns "void" in C-GLUT. this is a workaround for Julia's "automatically return the value of the last expression in a function" behavior.
end

_specialKeyPressed = cfunction(specialKeyPressed, Void, (Int32, Int32, Int32))

# run GLUT routines

glutInit()
glutInitDisplayMode(GLUT_RGBA | GLUT_DOUBLE | GLUT_ALPHA | GLUT_DEPTH)
glutInitWindowSize(width, height)
glutInitWindowPosition(0, 0)

window = glutCreateWindow("NeHe Tut 10")

glutDisplayFunc(_DrawGLScene)
glutFullScreen()

glutIdleFunc(_DrawGLScene)
glutReshapeFunc(_ReSizeGLScene)
glutKeyboardFunc(_keyPressed)
glutSpecialFunc(_specialKeyPressed)

initGL(width, height)

glutMainLoop()


text:
using BenchmarkTools
using UpperEnvelope
using Plots
import Random

Random.seed!(42);

tmp1 = sortslices(reshape(rand(15), (5, 3)), dims=1)
k1=tmp1[:, 1]
v1=tmp1[:, 2]
w1=1 .+ tmp1[:, 3]

tmp2 = sortslices(reshape(rand(15), (5, 3)), dims=1)
k2=tmp2[:, 1]
v2=tmp2[:, 2]
w2=1 .+ tmp2[:, 3]

envelope=compute_envelope((k1, v1, w1), (k2, v2, w2), true)

# Plot example (too lazy to make code prettier)
p0=plot(k1, v1, marker=(:black), line=(:black), label="Function 1")
p1=plot(k1, v1, marker=(:black), line=(:black), label="Function 1")
plot!(p0, k2, v2, marker=(:x, :darkblue), line=(:dash, :darkblue), label="Function 2")
plot!(p1, k2, v2, marker=(:x, :darkblue), line=(:dash, :darkblue), label="Function 2")
plot!(p1,envelope[1], [envelope[2] envelope[3]], line=(:purple, 3), marker=(:diamond, :purple, 3), label="Upper envelope")
plot(p0, p1, size=0.8.*(980, 440), dpi=190)

savefig("example.png")


text:
include("../../utils.jl")
using TimeZones

zdts = [now(tz"UTC") for _ in  1:100_000_000];

@gctime sum(hash, ["trashfire"^min(1000, i) for i in 1:100_000])

max_stars_repo_name: vtjnash/GCBenchmarks

max_issues_repo_name: vtjnash/GCBenchmarks

max_forks_repo_name: vtjnash/GCBenchmarks


text:
using Test

for testscen in 1:2
    valdir, scenario, use_permafrost, use_seaice = get_scenario(testscen)
    println(scenario)

    m = test_page_model()
    include("../src/components/RCPSSPScenario.jl")
    include("../src/components/TotalForcing.jl")

    scenario = addrcpsspscenario(m, scenario)
    totalforcing = add_comp!(m, TotalForcing)

    totalforcing[:f_CO2forcing] = readpagedata(m, "test/validationdata/$valdir/f_co2forcing.csv")
    totalforcing[:f_CH4forcing] = readpagedata(m, "test/validationdata/$valdir/f_ch4forcing.csv")
    totalforcing[:f_N2Oforcing] = readpagedata(m, "test/validationdata/$valdir/f_n2oforcing.csv")
    totalforcing[:f_lineargasforcing] = readpagedata(m, "test/validationdata/$valdir/f_LGforcing.csv")
    totalforcing[:exf_excessforcing] = scenario[:exf_excessforcing]
    totalforcing[:fs_sulfateforcing] = readpagedata(m, "test/validationdata/$valdir/fs_sulfateforcing.csv")

    p = load_parameters(m)
    p["y_year_0"] = 2015.
    p["y_year"] = Mimi.dim_keys(m.md, :time)
    set_leftover_params!(m, p)

    run(m)

    forcing = m[:TotalForcing, :ft_totalforcing]
    forcing_compare = readpagedata(m, "test/validationdata/$valdir/ft_totalforcing.csv")

    @test forcing ≈ forcing_compare rtol = 1e-3
end


text:
using EffectiveWaves
using LaTeXStrings
using Plots
height=500
# unicodeplots()
 pyplot(linewidth=3, size=(2.6*height,height))

 Plots.scalefontsizes(1.7)

## choose material

  # concrete
  medium = LimeStone
  inclusion1 = AirDry
  inclusion2 = Brick
  ωs = LinRange(real(medium.c/10000),real(medium.c),100) # k from 0 to 1
  # ωs = LinRange(real(medium.c/10000),real(medium.c/1000),100) # k from 0 to 1

  volfrac = 0.16
  r1 = 0.1; vol1 = 0.06
  r2 = 1.0; vol2 = volfrac - vol1

  sp1 = Specie(ρ=inclusion1.ρ ,r=r1, c=inclusion1.c, volfrac = vol1)
  sp2 = Specie(ρ=inclusion2.ρ ,r=r2, c=inclusion2.c, volfrac = volfrac-vol1)

  # True wavenumber
  kTs  = wavenumber_low_volumefraction(ωs, medium, [sp1,sp2]);
  # Approximate wavenumber
  kTLSs = two_species_approx_wavenumber(ωs, medium, [sp1,sp2]);

  speed_arr = [ ωs./real(kTs), ωs./real(kTLSs), 0 .*ωs + real(medium.c)]
  atten_arr = imag([kTs,kTLSs])

  styles = reshape([:solid,:solid,:dot],1,3)
  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$", "Lime Stone"] ,1, 3);
  ys_arr = speed_arr;
  xs = r1.*(ωs./real(medium.c));
  m =5;
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  plot(xs, ys_arr, xlabel=L"a_S k", ylabel="sound speed", labels=labs, line = styles);
  p1 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  ys_arr = atten_arr;
  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$"] ,1, 2);
  plot(xs, ys_arr, labels=labs, xlabel=L"a_S k", ylabel="attenuation");
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  p2 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  plot(p1,p2)
  savefig("media/compare_concrete.png")
  savefig("media/compare_concrete.pdf")
  gui()

## Zoomed in version
  ωs = LinRange(real(medium.c/10000),ωs[m],100) # k from 0 to 1
  m = length(ωs);
  kTs  = wavenumber_low_volumefraction(ωs, medium, [sp1,sp2]);
  kTLSs = two_species_approx_wavenumber(ωs, medium, [sp1,sp2]);

  speed_arr = [ ωs./real(kTs), ωs./real(kTLSs)]
  atten_arr = imag([kTs,kTLSs])

  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$"] ,1, 2);
  ys_arr = speed_arr;
  xs = r1.*(ωs./real(medium.c));
  plot(xs, ys_arr, xlabel=L"a_S k", ylabel="sound speed", labels=labs
                 , border = false);
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  p1 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  ys_arr = atten_arr;
  labs = reshape([L"Effective $k_{*}$",L"Approximate $k_{*LS}$"] ,1, 2);
  plot(xs, ys_arr, xlabel=L"a_S k", ylabel="attenuation", labels=labs
                 , border=false);
  y1 = min(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  y2 = max(ys_arr[1][1:m]..., ys_arr[2][1:m]...);
  dy = abs(y2- y1)/120.0;
  p2 = gray_square([xs[1],xs[m]],[y1,y2], plot!)

  plot(p1,p2)
  mkdir("media")
  savefig("media/compare_concrete_zoom.png")
  savefig("media/compare_concrete_zoom.pdf")

Plots.scalefontsizes(1/1.7)

max_stars_repo_name: arturgower/EffectiveWaves


text:
int(x) = (Int ∘ round)(x)
fint(x) = (Int ∘ floor)(x)
cint(x) = (Int ∘ ceil)(x)

"""
  loop_last(v)

  Returns an iterable yielding tuples (is_last, value).
"""
function loop_last(v)
    is_last = 1:length(v) .== length(v)
    return zip(is_last, v)
end

function loop_firstlast(v)
    is_first = 1:length(v) .== 1
    is_last = 1:length(v) .== length(v)
    return zip(is_first, is_last, v)
end

"""
  get_lr_widths(width::Int)

To split something with `width` in 2, get the lengths
of the left/right widths.

When width is even that's easy, when it's odd we need to
be careful.
"""
function get_lr_widths(width::Int)::Tuple{Int,Int}
    iseven(width) && return (int(width / 2), int(width / 2))
    return (fint(width / 2), cint(width / 2))
end

"""
Get a clean string representation of an expression
"""
expr2string(e::Expr) = replace_multi(
    string(e),
    '\n' => "",
    ' ' => "",
    r"#=.*=#" => "",
    "begin" => "",
    "end" => "",
)

"""
  get_file_format(nbytes; suffix="B")

Return a string with formatted file size.
"""
function get_file_format(nbytes; suffix = "B")
    for unit in ["", "K", "M", "G", "T", "P", "E", "Z", "Y"]
        if nbytes < 1024.0
            _nd = round(nbytes; digits = 2)
            return string(_nd) * " $unit$suffix"
        end
        nbytes = nbytes / 1024.0
    end
end

"""
    calc_nrows_ncols(n, aspect=(16, 9))

Computes the number of rows and columns to fit a given number `n` of subplots in a figure with aspect `aspect`.
Adapted from: stackoverflow.com/a/43366784
"""
function calc_nrows_ncols(n, aspect::Union{Number,NTuple} = (16, 9))
    w, h = if aspect isa Number
        (aspect, one(aspect))
    else
        aspect
    end
    factor = √(n / (w * h))
    cols = floor(Int, w * factor)
    rows = floor(Int, h * factor)
    rowFirst = w < h
    while rows * cols < n
        if rowFirst
            rows += 1
        else
            cols += 1
        end
        rowFirst = !rowFirst
    end
    return rows, cols
end


text:
using Test
using TestSetExtensions
using LinearAlgebra
using Qaintessent
using Random
using Qaintessent.MaxKColSubgraphQAOA

graphs = [
    Graph(3, [(1,2), (2,3)]),
    Graph(5, [(1,2), (2,3), (3,4), (4,5), (5,1), (1,3)]),
    Graph(4, [(1,2), (1,3), (1,4), (2,3), (2,4), (3,4)])
]

# Utility function to count the properly colored edges in a graph coloring (i.e. endpoints have different colors)
function properly_colored_edges(graph::Graph, coloring::Vector{Int})
    length(coloring) == graph.n || throw(ArgumentError("Length of coloring must equal the number of vertices."))
    return count(coloring[a] != coloring[b] for (a, b) ∈ graph.edges)
end

# Create a state ψ that corresponds to a given coloring
function ψ_from_coloring(n::Int, κ::Int, colors::Vector{Int})::Vector{ComplexF64}
    (n > 0 && κ > 0) || throw(DomainError("Parameters n and κ must be positive integers."))
    colors ⊆ 1:κ || throw(ArgumentError("Parameter `colors` may only contain colors in the range 1:$(κ)."))

    # Create ψ with |1> entries in the indices corresponding to the given colors, |0> elsewhere
    ψ = kron((color == colors[vertex] ? [0.0im, 1] : [1, 0.0im] for vertex ∈ 1:n for color ∈ 1:κ)...)
    ψ
end

# Utility function to compute the probabilities of the outcomes represented by a wavefunction ψ, sorted by descending probability
function wavefunction_distribution(ψ::Vector{ComplexF64}; as_bitstrings::Bool = true,
        include_zero = false)::Union{Vector{Tuple{Int, Float64}}, Vector{Tuple{Vector{Int}, Float64}}}
    distribution = [(i-1, abs(amplitude)^2) for (i, amplitude) ∈ enumerate(ψ) if abs(amplitude) > 0 || include_zero]
    if as_bitstrings
        N = Int(log2(length(ψ)))
        distribution = [(digits(i, base=2, pad=N) |> reverse, p) for (i, p) ∈ distribution]
    end

    return sort(distribution, by=(t -> -t[2]))
end

@testset ExtendedTestSet "max-k-col. subgraph QAOA Hamiltonians and gate matrices" begin
    @testset ExtendedTestSet "max-k-col. subgraph QAOA Hamiltonians" begin
        # Test that the phase separation Hamiltonian used in MaxKColSubgraphPhaseSeparationGate
        # implements the objective function correctly.
        @testset "Hamiltonian MaxKColSubgraphPhaseSeparationGate" begin
            κs = [4, 2, 3]

            hamiltonians = max_k_col_subgraph_phase_separation_hamiltonian.(graphs, κs)
            colorings = [rand(1:κ, graph.n) for (κ, graph) ∈ zip(κs, graphs)]

            for (H, graph, κ, coloring) ∈ zip(hamiltonians, graphs, κs, colorings)
                f_val = properly_colored_edges(graph, coloring)
                ψ_coloring = ψ_from_coloring(graph.n, κ, coloring)
                scaling = κ * length(graph.edges) - 4 * f_val # ψ_coloring should be scaled by this factor
                @test H * ψ_coloring ≈ ψ_coloring * scaling
            end
        end

        # Utility functions to perform circular shift on Vectors (that represent bitstrings)
        left_circ_shift(l::Vector) = [l[2:length(l)]; l[1:1]]
        right_circ_shift(l::Vector) = [l[length(l):length(l)]; l[1:length(l)-1]]

        # Shorthand for repeated function application
        # Taken from https://stackoverflow.com/questions/39895672/apply-function-repeatedly-a-specific-number-of-times
        (repeatApply)(f::Function, i::Int) = i==1 ? f : x->(repeatApply(f, i-1))(f(x))

        # Test that the Hamiltonian which the RNearbyValuesMixerGate is based on
        # maps coloring states to a superposition of their neighboring states
        @testset "Hamiltonian RNearbyValuesMixerGate" begin
            ds = [2, 5, 9, 10] # different numbers of colors
            rs = [1, 2, 5]
            nums_samples = [2, 5, 5, 5]

            for (d, num_samples) ∈ zip(ds, nums_samples)
                for r ∈ rs[rs .< d] # only test valid rs
                    # compute the Hamiltonian
                    gate = RNearbyValuesMixerGate(0., r, d)
                    hamiltonian = r_nearby_values_hamiltonian_onehot(gate)

                    # for each color (i.e. |1000..00>, |0100..00>, ..., |0000..01>):
                    for color ∈ randperm(d)[1:num_samples]
                        # compute superposition state bitstrings after applying Hamiltonian
                        ψ = ψ_from_coloring(1, d, [color])
                        ψ_bits = wavefunction_distribution(ψ)[1][1]
                        out_vec = hamiltonian * ψ # not a wavefunction (we apply a Hamiltonian, not a gate)
                        superposition_states = (wavefunction_distribution(out_vec)
                            |> (distr -> filter(d -> !(d[2] ≈ 0), distr))
                            .|> (d -> d[1]))

                        # compute bitstrings of nearby colors by shifting
                        near_values = [
                            [repeatApply(left_circ_shift, k)(ψ_bits) for k ∈ 1:r];
                            [repeatApply(right_circ_shift, k)(ψ_bits) for k ∈ 1:r]
                        ] # repeatApply is repeated function application defined above

                        # r-NV Mixer should map state to superposition of its neighbor colors
                        @test Set(superposition_states) == Set(near_values)
                    end
                end
            end
        end
    end

    # Utility function that returns the Hamming weights of all basis states of ψ
    hamming_weights(ψ) = (wavefunction_distribution(ψ)
    |> (distr -> filter(d -> !(d[2] ≈ 0), distr))
    .|> (d -> sum(d[1])))

    @testset ExtendedTestSet "max-k-col. subgraph QAOA mixer Hamming weights" begin
        # Test that the RNearbyValuesMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one)
        @testset "Hamming weight RNearbyValuesMixerGate" begin
            ds = [2, 5, 9, 10] # different numbers of colors
            rs = [1, 2, 5]
            nums_samples = [2, 5, 5, 5]

            for (d, num_samples) ∈ zip(ds, nums_samples)
                for r ∈ rs[rs .< d] # only test valid rs
                    gate = RNearbyValuesMixerGate(rand() * 10 + 0.1, r, d)
                    U = Qaintessent.matrix(gate)

                    for color ∈ randperm(d)[1:num_samples]
                        ψ = ψ_from_coloring(1, d, [color])
                        ψ_out = U * ψ

                        # All basis states should have Hamming weight 1
                        out_hamming_weights = hamming_weights(ψ_out)
                        @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                    end
                end
            end
        end

        # Test that the ParityRingMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one), if d == 2
        @testset "Hamming weight ParityRingMixerGate (d = 2)" begin
            ds = [2] # different numbers of colors

            for d ∈ ds
                gate = ParityRingMixerGate(rand() * 10 + 0.1, d)
                U = Qaintessent.matrix(gate)

                for color ∈ 1:d
                    ψ = ψ_from_coloring(1, d, [color])
                    ψ_out = U * ψ

                    # All basis states should have Hamming weight 1
                    out_hamming_weights = hamming_weights(ψ_out)
                    @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                end
            end
        end

        # Test that the ParityRingMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one), if d ≥ 3
        @testset "Hamming weight ParityRingMixerGate (d ≥ 3)" begin
            ds = [5, 8, 9] # different numbers of colors
            num_samples = 5

            for d ∈ ds
                gate = ParityRingMixerGate(rand() * 10 + 0.1, d)
                U = Qaintessent.matrix(gate)

                for color ∈ randperm(d)[1:num_samples]
                    ψ = ψ_from_coloring(1, d, [color])
                    ψ_out = U * ψ

                    # All basis states should have Hamming weight 1
                    out_hamming_weights = hamming_weights(ψ_out)
                    @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                end
            end
        end

        # Test that the PartitionMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one)
        @testset "Hamming weight PartitionMixerGate" begin
            ds = [5, 8, 9] # different numbers of colors
            num_samples = 5
            nums_partition_layers = [3, 2, 7]

            for (d, num_partition_layers) ∈ zip(ds, nums_partition_layers)
                # create num_partition_layers partition parts, each of
                # random length and with random entries
                partition = [
                    # create h x 2 matrix with random height 2 from a random permutation of the indices...
                    (reshape(randperm(d)[1:(2*rand(1:div(d, 2)))], (:, 2))
                    # ... then map the rows to tuples
                        |> eachrow .|> Tuple)
                    for _ ∈ 1:num_partition_layers
                ]

                gate = PartitionMixerGate(rand() * 10 + 0.1, d, partition)
                U = Qaintessent.matrix(gate)

                for color ∈ randperm(d)[1:num_samples]
                    ψ = ψ_from_coloring(1, d, [color])
                    ψ_out = U * ψ

                    # All basis states should have Hamming weight 1
                    out_hamming_weights = hamming_weights(ψ_out)
                    @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                end
            end
        end
    end

    # Test that the parity ring mixer is a special case of the partition mixer
    @testset "PartitionMixerGate generalizes ParityRingMixerGate" begin
        ds = [2, 5, 9, 10]

        for d ∈ ds
            β = rand() * 10 + 0.1

            # construct the parity mixer
            parity_mixer = ParityRingMixerGate(β, d)

            # construct the equivalent partition mixer
            odd_part = [(a, a+1) for a ∈ 1:2:(d-1)]
            even_part = [(a, (a % d) + 1) for a ∈ 2:2:d]
            last_part = isodd(d) ? [(d, 1)] : nothing
            partition = isodd(d) ? [odd_part, even_part, last_part] : [odd_part, even_part]
            partition_mixer = PartitionMixerGate(β, d, partition)

            @test Qaintessent.matrix(parity_mixer) ≈ Qaintessent.matrix(partition_mixer)
        end
    end

    @testset ExtendedTestSet "max-k-col. subgraph QAOA adjoints" begin
        # Test that `MaxKColSubgraphPhaseSeparationGate` has the correct adjoint.
        @testset "adjoint MaxKColSubgraphPhaseSeparationGate" begin
            γs = rand(length(graphs)) * 2π
            κs = [3, 2, 2]

            gates = MaxKColSubgraphPhaseSeparationGate.(γs, κs, graphs)
            for i ∈ 1:length(gates)
                @test Qaintessent.matrix(gates[i]) * Qaintessent.matrix(Base.adjoint(gates[i])) ≈ I
            end
        end

        # Test that `ParityRingMixerGate` has the correct adjoint.
        @testset "adjoint ParityRingMixerGate" begin
            N = 3 # number of test cases
            βs = rand(N) * 2π
            ds = rand(2:9, N)

            gates = ParityRingMixerGate.(βs, ds)
            for i ∈ 1:length(gates)
                @test Qaintessent.matrix(gates[i]) * Qaintessent.matrix(Base.adjoint(gates[i])) ≈ I
            end
        end

        # Test that `RNearbyValuesMixerGate` has the correct adjoint.
        @testset "adjoint RNearbyValuesMixerGate" begin
            βs = rand(4) * 2π
            ds = [2, 3, 6, 9]
            rs = [1, 2, 1, 5]

            gates = RNearbyValuesMixerGate.(βs, rs, ds)
            for i ∈ 1:length(gates)
                @test Qaintessent.matrix(gates[i]) * Qaintessent.matrix(Base.adjoint(gates[i])) ≈ I
            end
        end

        # Test that `PartitionMixerGate` has the correct adjoint.
        @testset "adjoint PartitionMixerGate" begin
            ds = [2, 3, 6, 9]
            nums_partition_layers = [2, 3, 4, 5]

            for (d, num_partition_layers) ∈ zip(ds, nums_partition_layers)
                # create num_partition_layers partition parts, each of
                # random length and with random entries
                partition = [
                    # create h x 2 matrix with random height 2 from a random permutation of the indices...
                    (reshape(randperm(d)[1:(2*rand(1:div(d, 2)))], (:, 2))
                    # ... then map the rows to tuples
                        |> eachrow .|> Tuple)
                    for _ ∈ 1:num_partition_layers
                ]

                β = rand() * 10 + 0.1
                gate = PartitionMixerGate(β, d, partition)
                @test Qaintessent.matrix(gate) * Qaintessent.matrix(Base.adjoint(gate)) ≈ I
            end
        end
    end

    @testset ExtendedTestSet "max-k-col. subgraph QAOA num wires" begin
        @testset "num wires MaxKColSubgraphPhaseSeparationGate" begin
            γs = rand(length(graphs)) * 2π
            κs = [3, 2, 2]

            gates = MaxKColSubgraphPhaseSeparationGate.(γs, κs, graphs)
            for (gate, graph, κ) ∈ zip(gates, graphs, κs)
                @test Qaintessent.num_wires(gate) == graph.n * κ
            end
        end

        @testset "num wires ParityRingMixerGate" begin
            N = 3 # number of test cases
            βs = rand(N) * 2π
            ds = rand(2:9, N)

            gates = ParityRingMixerGate.(βs, ds)
            for (gate, d) ∈ zip(gates, ds)
                @test Qaintessent.num_wires(gate) == d
            end
        end

        @testset "num wires RNearbyValuesMixerGate" begin
            βs = rand(4) * 2π
            ds = [2, 3, 6, 9]
            rs = [1, 2, 1, 5]

            gates = RNearbyValuesMixerGate.(βs, rs, ds)
            for (gate, d) ∈ zip(gates, ds)
                @test Qaintessent.num_wires(gate) == d
            end
        end

        @testset "num wires PartitionMixerGate" begin
            ds = [2, 3, 6, 9]
            nums_partition_layers = [2, 3, 4, 5]

            for (d, num_partition_layers) ∈ zip(ds, nums_partition_layers)
                # create num_partition_layers partition parts, each of
                # random length and with random entries
                partition = [
                    # create h x 2 matrix with random height 2 from a random permutation of the indices...
                    (reshape(randperm(d)[1:(2*rand(1:div(d, 2)))], (:, 2))
                    # ... then map the rows to tuples
                        |> eachrow .|> Tuple)
                    for _ ∈ 1:num_partition_layers
                ]

                β = rand() * 10 + 0.1
                gate = PartitionMixerGate(β, d, partition)
                @test Qaintessent.num_wires(gate) == d
            end
        end
    end
end


text:
using Revise

for file in ["common.jl", "measure.jl", "analyze.jl"]
    Revise.track(file)
    include(file)
end


text:
# Binary dithering and color dithering can be distinguished by the extra argument `arg`,
# which is either
#   - a color scheme (array of colors)
#   - a ColorSchemes.jl symbol
#   - the number of colors specified for clustering
#
# All functions in this file end up calling `colordither` and return IndirectArrays.

struct ColorNotImplementedError <: Exception
    algname::String
    ColorNotImplementedError(alg::AbstractDither) = new("$alg")
end
function Base.showerror(io::IO, e::ColorNotImplementedError)
    return print(
        io, e.algname, " algorithm currently doesn't support custom color palettes."
    )
end
colordither(alg, img, cs, metric) = throw(ColorNotImplementedError(alg))

##############
# Public API #
##############

# If `out` is specified, it will be changed in place...
function dither!(out::GenericImage, img::GenericImage, alg::AbstractDither, arg; kwargs...)
    return out .= _colordither(eltype(out), img, alg, arg; kwargs...)
end

# ...otherwise `img` will be changed in place.
function dither!(img::GenericImage, alg::AbstractDither, arg; kwargs...)
    return img .= _colordither(eltype(img), img, alg, arg; kwargs...)
end

# The return type can be chosen...
function dither(::Type{T}, img::GenericImage, alg::AbstractDither, arg; kwargs...) where {T}
    return _colordither(T, img, alg, arg; kwargs...)
end

# ...and defaults to the type of the input image.
function dither(
    img::GenericImage{T,N}, alg::AbstractDither, arg; kwargs...
) where {T<:Pixel,N}
    return _colordither(T, img, alg, arg; kwargs...)
end

#############################
# Low-level algorithm calls #
#############################

# Dispatch to dithering with custom color palettes on any image type
# when color palette is provided
function _colordither(
    ::Type{T},
    img::GenericImage,
    alg::AbstractDither,
    cs::AbstractVector{<:Pixel};
    metric::DifferenceMetric=DE_2000(),
    to_linear=false,
) where {T}
    to_linear && (@warn "Skipping transformation `to_linear` when dithering in color.")
    length(cs) >= 2 ||
        throw(DomainError(steps, "Color scheme for dither needs >= 2 colors."))

    index = colordither(alg, img, cs, metric)
    return IndirectArray(index, T.(cs))
end

# A special case occurs when a grayscale output image is to be dithered in colors.
# Since this is not possible, instead the return image will be of type of the color scheme.
function _colordither(
    ::Type{T},
    img::GenericImage,
    alg::AbstractDither,
    cs::AbstractVector{<:Color{<:Any,3}};
    metric::DifferenceMetric=DE_2000(),
    to_linear=false,
) where {T<:NumberLike}
    return _colordither(eltype(cs), img, alg, cs; metric=metric, to_linear=to_linear)
end


text:
using SmartShockFinder

# Determine the type of Network
n_dims = 2
network_type = CNN(n_dims)

# add mesh and polynomdegrees to then generate Traindata
meshes = SmartShockFinder.get_mesh_2d(6)
polydeg = collect(2:9)
train_setup = Traindata_Settings(polydeg, meshes)

# Parameters to train the Net
η = 0.001               # learning rate
β = 0.001                # regularization paramater
number_epochs = 100
Sb = 500                 # batch size
L = 10                   # early stopping parameter
hidden_units = []

model_setup = Model_Settings(η, β, number_epochs, Sb, L, hidden_units)

# generate Traindata
generate_traindataset2d(network_type, train_setup)

# generate validationdata
generate_validdataset2d(network_type, train_setup)

# Build network
train_network2dcnn(network_type, model_setup)


text:
using DSAGAnalysis
using Test

@testset "DSAGAnalysis.jl" begin
    # Write your tests here.
end


text:
function success!(ata_model::AbstractModel, message::String)
    push!(ata_model.output.infos, ["success", string("- ", message, "\n")])
    return nothing
end

function error!(ata_model::AbstractModel, message::String)
    push!(ata_model.output.infos, ["danger", string("- ", message, "\n")])
    return nothing
end


text:
export acceptance_prob

abstract type McmcMove end

acceptance_prob(move::McmcMove) = move.counts[1]/move.counts[2]

function reset_counts!(mcmc::McmcMove)
    mcmc.counts .= 0
end


text:
abstract type AbstractPlotList{T<:Tuple} end

plottype(::Type{<:AbstractPlotList{T}}) where {T} = T.parameters

"""
`Plotlist(plots...)`

Experimental feature. Create an object that can encode multiple series.
"""
struct PlotList{T<:Tuple} <: AbstractPlotList{T}
    plots::Tuple
    function PlotList(plots...)
        T = Tuple{unique(plottype.(plots))...}
        new{T}(plots)
    end
end

Base.getindex(m::PlotList, i) = getindex(m.plots, i)
Base.length(m::PlotList) = length(m.plots)
Base.iterate(m::PlotList, args...) = iterate(m.plots, args...)

@recipe(MultiplePlot) do scene
    default_theme(scene)
end

function default_theme(scene, ::Type{<:Combined{multipleplot, Tuple{P}}}) where {P<:AbstractPlotList}
    merge((default_theme(scene, pt) for pt in plottype(P))...)
end
# Allow MultiplePlot to prevail on user input: the plot type of each series will be defined in convert_arguments
plottype(::Type{<: Combined{Any}}, A::Type{<:MultiplePlot}, argvalues...) = A
plottype(::Type{<: Combined{T}}, A::Type{<:MultiplePlot}, argvalues...) where T = A

function convert_arguments(P::PlotFunc, m::PlotList)
    function convert_series(plot::PlotSpec)
        ptype = plottype(P, plottype(plot))
        to_plotspec(ptype, convert_arguments(ptype, plot.args...); plot.kwargs...)
    end
    pl = PlotList(convert_series.(m.plots)...)
    PlotSpec{MultiplePlot}(pl)
end

combine(val1, val2; palette = nothing) = val2

function combine!(theme1::Theme, theme2::Theme)
    palette = get(theme1, :palette, current_default_theme()[:palette])
    for (key, val) in theme2
        tv = convert(Node, get(theme1, key, nothing))
        pv = convert(Node, get(palette, key, nothing))
        theme1[key] = lift((t, p, v) -> combine(t, v, palette = p), tv, pv, val)
    end
    theme1
end
combine(theme1::Theme, theme2) = combine!(copy(theme1), theme2)

# This allows plotting an arbitrary combination of series form one argument
# The recipe framework can be constructed using this as a building block and computing
# PlotList with convert_arguments
function plot!(p::Combined{multipleplot, <:Tuple{PlotList}})
    mp = to_value(p[1]) # TODO how to preserve interactivity here, as number of series may change?
    theme = Theme(p)
    for s in mp.plots
        attr = combine(theme, Theme(; s.kwargs...))
        plot!(p, plottype(s), attr, s.args...)
    end
end

function default_theme(scene, ::Type{<:Combined{S, T}}) where {S<:Tuple, T}
    merge((default_theme(scene, Combined{pt}) for pt in S.parameters)...)
end

function plot!(p::Combined{S, <:Tuple{Vararg{Any, N}}}) where {S <: Tuple, N}
    for pt in S.parameters
        plot!(p, Combined{pt}, Theme(p), p[1:N]...)
    end
end

function Base.:*(::Type{<:Combined{S}}, ::Type{<:Combined{T}}) where {S, T}
    params1 = S isa Type{<:Tuple} ? S.parameters : [S]
    params2 = T isa Type{<:Tuple} ? T.parameters : [T]
    params = union(params1, params2)
    Combined{Tuple{params...}}
end


text:
# Julia wrapper for header: /usr/include/scip/scip_benders.h
# Automatically generated using Clang.jl wrap_c

function SCIPincludeBenders(scip, name, desc, priority, cutlp, cutpseudo, cutrelax, shareauxvars, benderscopy, bendersfree, bendersinit, bendersexit, bendersinitpre, bendersexitpre, bendersinitsol, bendersexitsol, bendersgetvar, benderscreatesub, benderspresubsolve, benderssolvesubconvex, benderssolvesub, benderspostsolve, bendersfreesub, bendersdata)
    ccall((:SCIPincludeBenders, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Cstring, Cstring, Cint, UInt32, UInt32, UInt32, UInt32, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{SCIP_BENDERSDATA}), scip, name, desc, priority, cutlp, cutpseudo, cutrelax, shareauxvars, benderscopy, bendersfree, bendersinit, bendersexit, bendersinitpre, bendersexitpre, bendersinitsol, bendersexitsol, bendersgetvar, benderscreatesub, benderspresubsolve, benderssolvesubconvex, benderssolvesub, benderspostsolve, bendersfreesub, bendersdata)
end

function SCIPincludeBendersBasic(scip, bendersptr, name, desc, priority, cutlp, cutpseudo, cutrelax, shareauxvars, bendersgetvar, benderscreatesub, bendersdata)
    ccall((:SCIPincludeBendersBasic, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{Ptr{SCIP_BENDERS}}, Cstring, Cstring, Cint, UInt32, UInt32, UInt32, UInt32, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{SCIP_BENDERSDATA}), scip, bendersptr, name, desc, priority, cutlp, cutpseudo, cutrelax, shareauxvars, bendersgetvar, benderscreatesub, bendersdata)
end

function SCIPsetBendersCopy(scip, benders, benderscopy)
    ccall((:SCIPsetBendersCopy, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, benderscopy)
end

function SCIPsetBendersFree(scip, benders, bendersfree)
    ccall((:SCIPsetBendersFree, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersfree)
end

function SCIPsetBendersInit(scip, benders, bendersinit)
    ccall((:SCIPsetBendersInit, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersinit)
end

function SCIPsetBendersExit(scip, benders, bendersexit)
    ccall((:SCIPsetBendersExit, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersexit)
end

function SCIPsetBendersInitpre(scip, benders, bendersinitpre)
    ccall((:SCIPsetBendersInitpre, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersinitpre)
end

function SCIPsetBendersExitpre(scip, benders, bendersexitpre)
    ccall((:SCIPsetBendersExitpre, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersexitpre)
end

function SCIPsetBendersInitsol(scip, benders, bendersinitsol)
    ccall((:SCIPsetBendersInitsol, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersinitsol)
end

function SCIPsetBendersExitsol(scip, benders, bendersexitsol)
    ccall((:SCIPsetBendersExitsol, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, bendersexitsol)
end

function SCIPsetBendersPresubsolve(scip, benders, benderspresubsolve)
    ccall((:SCIPsetBendersPresubsolve, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, benderspresubsolve)
end

function SCIPsetBendersSolveAndFreesub(scip, benders, benderssolvesubconvex, benderssolvesub, bendersfreesub)
    ccall((:SCIPsetBendersSolveAndFreesub, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}), scip, benders, benderssolvesubconvex, benderssolvesub, bendersfreesub)
end

function SCIPsetBendersPostsolve(scip, benders, benderspostsolve)
    ccall((:SCIPsetBendersPostsolve, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Cvoid}), scip, benders, benderspostsolve)
end

function SCIPfindBenders(scip, name)
    ccall((:SCIPfindBenders, libscip), Ptr{SCIP_BENDERS}, (Ptr{SCIP_}, Cstring), scip, name)
end

function SCIPgetBenders(scip)
    ccall((:SCIPgetBenders, libscip), Ptr{Ptr{SCIP_BENDERS}}, (Ptr{SCIP_},), scip)
end

function SCIPgetNBenders(scip)
    ccall((:SCIPgetNBenders, libscip), Cint, (Ptr{SCIP_},), scip)
end

function SCIPgetNActiveBenders(scip)
    ccall((:SCIPgetNActiveBenders, libscip), Cint, (Ptr{SCIP_},), scip)
end

function SCIPactivateBenders(scip, benders, nsubproblems)
    ccall((:SCIPactivateBenders, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Cint), scip, benders, nsubproblems)
end

function SCIPdeactivateBenders(scip, benders)
    ccall((:SCIPdeactivateBenders, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}), scip, benders)
end

function SCIPsetBendersPriority(scip, benders, priority)
    ccall((:SCIPsetBendersPriority, libscip), Cvoid, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Cint), scip, benders, priority)
end

function SCIPsolveBendersSubproblems(scip, benders, sol, result, infeasible, auxviol, type, checkint)
    ccall((:SCIPsolveBendersSubproblems, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_SOL}, Ptr{SCIP_RESULT}, Ptr{UInt32}, Ptr{UInt32}, SCIP_BENDERSENFOTYPE, UInt32), scip, benders, sol, result, infeasible, auxviol, type, checkint)
end

function SCIPgetBendersMasterVar(scip, benders, var, mappedvar)
    ccall((:SCIPgetBendersMasterVar, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_VAR}, Ptr{Ptr{SCIP_VAR}}), scip, benders, var, mappedvar)
end

function SCIPgetBendersSubproblemVar(scip, benders, var, mappedvar, probnumber)
    ccall((:SCIPgetBendersSubproblemVar, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_VAR}, Ptr{Ptr{SCIP_VAR}}, Cint), scip, benders, var, mappedvar, probnumber)
end

function SCIPgetBendersNSubproblems(scip, benders)
    ccall((:SCIPgetBendersNSubproblems, libscip), Cint, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}), scip, benders)
end

function SCIPaddBendersSubproblem(scip, benders, subproblem)
    ccall((:SCIPaddBendersSubproblem, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_}), scip, benders, subproblem)
end

function SCIPsetupBendersSubproblem(scip, benders, sol, probnumber)
    ccall((:SCIPsetupBendersSubproblem, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_SOL}, Cint), scip, benders, sol, probnumber)
end

function SCIPsolveBendersSubproblem(scip, benders, sol, probnumber, infeasible, type, solvecip, objective)
    ccall((:SCIPsolveBendersSubproblem, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_SOL}, Cint, Ptr{UInt32}, SCIP_BENDERSENFOTYPE, UInt32, Ptr{Cdouble}), scip, benders, sol, probnumber, infeasible, type, solvecip, objective)
end

function SCIPfreeBendersSubproblem(scip, benders, probnumber)
    ccall((:SCIPfreeBendersSubproblem, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Cint), scip, benders, probnumber)
end

function SCIPcheckBendersSubproblemOptimality(scip, benders, sol, probnumber, optimal)
    ccall((:SCIPcheckBendersSubproblemOptimality, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_SOL}, Cint, Ptr{UInt32}), scip, benders, sol, probnumber, optimal)
end

function SCIPgetBendersAuxiliaryVarVal(scip, benders, sol, probnumber)
    ccall((:SCIPgetBendersAuxiliaryVarVal, libscip), Cdouble, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_SOL}, Cint), scip, benders, sol, probnumber)
end

function SCIPcomputeBendersSubproblemLowerbound(scip, benders, probnumber, lowerbound, infeasible)
    ccall((:SCIPcomputeBendersSubproblemLowerbound, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Cint, Ptr{Cdouble}, Ptr{UInt32}), scip, benders, probnumber, lowerbound, infeasible)
end

function SCIPmergeBendersSubproblemIntoMaster(scip, benders, varmap, consmap, probnumber)
    ccall((:SCIPmergeBendersSubproblemIntoMaster, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{SCIP_HASHMAP}, Ptr{SCIP_HASHMAP}, Cint), scip, benders, varmap, consmap, probnumber)
end

function SCIPincludeBenderscut(scip, benders, name, desc, priority, islpcut, benderscutcopy, benderscutfree, benderscutinit, benderscutexit, benderscutinitsol, benderscutexitsol, benderscutexec, benderscutdata)
    ccall((:SCIPincludeBenderscut, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Cstring, Cstring, Cint, UInt32, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cvoid}, Ptr{SCIP_BENDERSCUTDATA}), scip, benders, name, desc, priority, islpcut, benderscutcopy, benderscutfree, benderscutinit, benderscutexit, benderscutinitsol, benderscutexitsol, benderscutexec, benderscutdata)
end

function SCIPincludeBenderscutBasic(scip, benders, benderscutptr, name, desc, priority, islpcut, benderscutexec, benderscutdata)
    ccall((:SCIPincludeBenderscutBasic, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERS}, Ptr{Ptr{SCIP_BENDERSCUT}}, Cstring, Cstring, Cint, UInt32, Ptr{Cvoid}, Ptr{SCIP_BENDERSCUTDATA}), scip, benders, benderscutptr, name, desc, priority, islpcut, benderscutexec, benderscutdata)
end

function SCIPsetBenderscutCopy(scip, benderscut, benderscutcopy)
    ccall((:SCIPsetBenderscutCopy, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{Cvoid}), scip, benderscut, benderscutcopy)
end

function SCIPsetBenderscutFree(scip, benderscut, benderscutfree)
    ccall((:SCIPsetBenderscutFree, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{Cvoid}), scip, benderscut, benderscutfree)
end

function SCIPsetBenderscutInit(scip, benderscut, benderscutinit)
    ccall((:SCIPsetBenderscutInit, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{Cvoid}), scip, benderscut, benderscutinit)
end

function SCIPsetBenderscutExit(scip, benderscut, benderscutexit)
    ccall((:SCIPsetBenderscutExit, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{Cvoid}), scip, benderscut, benderscutexit)
end

function SCIPsetBenderscutInitsol(scip, benderscut, benderscutinitsol)
    ccall((:SCIPsetBenderscutInitsol, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{Cvoid}), scip, benderscut, benderscutinitsol)
end

function SCIPsetBenderscutExitsol(scip, benderscut, benderscutexitsol)
    ccall((:SCIPsetBenderscutExitsol, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{Cvoid}), scip, benderscut, benderscutexitsol)
end

function SCIPsetBenderscutPriority(scip, benderscut, priority)
    ccall((:SCIPsetBenderscutPriority, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Cint), scip, benderscut, priority)
end

function SCIPstoreBenderscutCons(scip, benderscut, cons)
    ccall((:SCIPstoreBenderscutCons, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{SCIP_CONS}), scip, benderscut, cons)
end

function SCIPstoreBenderscutCut(scip, benderscut, cut)
    ccall((:SCIPstoreBenderscutCut, libscip), SCIP_RETCODE, (Ptr{SCIP_}, Ptr{SCIP_BENDERSCUT}, Ptr{SCIP_ROW}), scip, benderscut, cut)
end


text:
foo(x) = x^5
bar(x) = x^10
apply(f, x1, x2::Val{T}) where {T} = f(x1, T)

function f(x)
    g = x < 5 ? foo : bar
    return g(2)
end

struct DynamicMix <: CompilationContext end

allow(ctx::DynamicMix, m::Module) = m == TestMixtape

swap(e) = e
function swap(e::Expr)
    new = MacroTools.postwalk(e) do s
        isexpr(s, :call) || return s
        s.args[1] == Base.literal_pow || return s
        return Expr(:call, apply, Base.:(*), s.args[3:end]...)
    end
    return new
end

function transform(::DynamicMix, b)
    for (v, st) in b
        replace!(b, v, swap(st))
    end
    return b
end

# JIT compile an entry and call.
@testset "Dynamic overlay" begin
    fn = Mixtape.jit(DynamicMix(), f, Tuple{Int64})
    @test fn(3) == 10
    @test fn(6) == 20

    # Mixtape cached call.
    Mixtape.@load_call_interface()
    @test call(DynamicMix(), f, 3) == 10
    @test call(DynamicMix(), f, 6) == 20

    # Native.
    @test f(3) == 2^5
    @test f(6) == 2^10
end


text:
# Load Turing and MCMCChains.
using Turing
using CosmoCentral
using Distributed
using Dates
using ClusterManagers
using LinearAlgebra

n_process = 100
n_step = 400
#add  n_processes on the long queue
for i in 1:n_process
    time = string(Dates.now())
    ClusterManagers.addprocs_lsf(1; bsub_flags = `-q long -o Turing/outfiles/$time.out -e Turing/outfiles/$time.err -M 3000 -R 'hname!=teo26.ge.infn.it && hname!=teo25.ge.infn.it && hname!=teo27.ge.infn.it && hname!=teo28.ge.infn.it && hname!=g2farm7.ge.infn.it && hname!=totem04.ge.infn.it && hname!=totem07.ge.infn.it && hname!=totem08.ge.infn.it &&  hname!=g2farm8.ge.infn.it && hname!=teo22.ge.infn.it && hname!=teo24.ge.infn.it && hname!=teo28.ge.infn.it'`)
end

@everywhere using Distributed, Turing, LinearAlgebra
@everywhere using CosmoCentral

@everywhere w0waCDMCosmology = CosmoCentral.Flatw0waCDMCosmology()
println("Loaded central cosmology")

@everywhere @model function gdemo(vecpCℓData, Cov, ConvolvedDensity, EuclidBias, EuclidIA,
    CosmologicalGrid)
    w₀ ~ Uniform(-1.32, -0.72)
    wₐ ~ Uniform(-0.68, 0.72)
    ns ~ Uniform(0.92, 1.0)

    w0waCDMCosmology = CosmoCentral.Flatw0waCDMCosmology()
    w0waCDMCosmology.w0 = w₀
    w0waCDMCosmology.wa = wₐ
    w0waCDMCosmology.ns = ns
    CℓMCMC = CosmoCentral.EvaluateCℓMCMCStep(w0waCDMCosmology, ConvolvedDensity, EuclidBias,
    EuclidIA, CosmologicalGrid)

    ℓnumber = length(CℓMCMC.CℓArray[:,1,1])
    inumber = length(CℓMCMC.CℓArray[1,:,1])
    L = CosmoCentral.EliminationMatrix(inumber)

    vecpCℓMCMC = zeros(1,floor(Int,inumber*0.5*(inumber+1)))
    for ℓ in 1:ℓnumber
        vecCℓMCMC = vec(CℓMCMC.CℓArray[ℓ,:,:])
        vecpCℓMCMC = zeros(floor(Int,inumber*0.5*(inumber+1)))
        LinearAlgebra.mul!(vecpCℓMCMC, L, vecCℓMCMC)

        vecpCℓData[ℓ,:] ~ MvNormal(vecpCℓMCMC, Cov.Covariance[ℓ,:,:])
    end
end

@everywhere MultipolesArrayTemp = CosmoCentral.LogSpaced(10.,3000., 101)
@everywhere MultipolesArray = zeros(100)
@everywhere MultipolesWidths = CosmoCentral.Difference(MultipolesArrayTemp)
@everywhere for i in 1:100
    MultipolesArray[i] = (MultipolesArrayTemp[i+1]+MultipolesArrayTemp[i])/2
end
@everywhere CosmologicalGrid = CosmoCentral.CosmologicalGrid(ZArray = LinRange(0.001, 4., 500),
KArray = CosmoCentral.LogSpaced(1e-5, 50., 1000), ℓBinCenters = MultipolesArray,
ℓBinWidths = MultipolesWidths);

@everywhere AnalitycalDensity = CosmoCentral.AnalitycalDensity()
@everywhere CosmoCentral.NormalizeAnalitycalDensity!(AnalitycalDensity)

@everywhere InstrumentResponse = CosmoCentral.InstrumentResponse()

@everywhere ConvolvedDensity = CosmoCentral.ConvolvedDensity(DensityGridArray = ones(10, length(CosmologicalGrid.ZArray)))
@everywhere CosmoCentral.NormalizeConvolvedDensity!(ConvolvedDensity, AnalitycalDensity, InstrumentResponse, CosmologicalGrid)

@everywhere CosmoCentral.ComputeConvolvedDensityGrid!(CosmologicalGrid, ConvolvedDensity, AnalitycalDensity, InstrumentResponse)

@everywhere EuclidBias = CosmoCentral.EuclidBias()
@everywhere EuclidIA = CosmoCentral.ExtendedNLIA()

@everywhere CℓData = CosmoCentral.EvaluateCℓMCMCStep(w0waCDMCosmology, ConvolvedDensity, EuclidBias,
EuclidIA, CosmologicalGrid)
#TODO now only LL, but this need to be more flexible...maybe list with probes?
@everywhere Covaₗₘ = CosmoCentral.InstantiateEvaluateCovariance(CℓData, ConvolvedDensity, CosmologicalGrid, "Lensing",
"Lensing")
@everywhere CovCℓ = CosmoCentral.InstantiateEvaluateCovariance(Covaₗₘ)
@everywhere inumber = 10
@everywhere L = CosmoCentral.EliminationMatrix(inumber)
@everywhere vecpCℓData = zeros(100,floor(Int,inumber*0.5*(inumber+1)))
@everywhere for ℓ in 1:length(CℓData.CℓArray[:,1,1])
    vecCℓData = vec(CℓData.CℓArray[ℓ,:,:])
    vecpCℓDataT = zeros(floor(Int,inumber*0.5*(inumber+1)))
    LinearAlgebra.mul!(vecpCℓDataT, L, vecCℓData)
    vecpCℓData[ℓ,:] = vecpCℓDataT
    Matrix = (CovCℓ.Covariance[ℓ,:,:] .+ transpose(CovCℓ.Covariance[ℓ,:,:])) ./2
    CovCℓ.Covariance[ℓ,:,:] = Matrix
end
println(size(CovCℓ.Covariance[1,:,:]))

@everywhere model = gdemo(vecpCℓData, CovCℓ, ConvolvedDensity, EuclidBias, EuclidIA,
CosmologicalGrid)

chains = sample(model, MH(), MCMCDistributed(), n_step, n_process; save_state = true)
write("chain_one.jls", chains)