text:
module EdgeType
    @enum EType SIM HAD
end

"""
    ZXGraph{T, P}

This is the type for representing the graph-like ZX-diagrams.
"""
struct ZXGraph{T<:Integer, P} <: AbstractZXDiagram{T, P}
    mg::Multigraph{T}
    ps::Dict{T, P}
    st::Dict{T, SpiderType.SType}
    et::Dict{Tuple{T, T}, EdgeType.EType}
    layout::ZXLayout{T}
    phase_ids::Dict{T,Tuple{T, Int}}
    scalar::Scalar{P}
    master::ZXDiagram{T, P}
    inputs::Vector{T}
    outputs::Vector{T}
end

function Base.copy(zxg::ZXGraph{T, P}) where {T, P}
    ZXGraph{T, P}(
        copy(zxg.mg), copy(zxg.ps),
        copy(zxg.st), copy(zxg.et), copy(zxg.layout),
        deepcopy(zxg.phase_ids), copy(zxg.scalar),
        copy(zxg.master), copy(zxg.inputs), copy(zxg.outputs)
    )
end

"""
    ZXGraph(zxd::ZXDiagram)

Convert a ZX-diagram to graph-like ZX-diagram.

```jldoctest
julia> using ZXCalculus

julia> zxd = ZXDiagram(2); push_gate!(zxd, Val{:CNOT}(), 2, 1);

julia> zxg = ZXGraph(zxd)
ZX-graph with 6 vertices and 5 edges:
(S_1{input} <-> S_5{phase = 0//1⋅π})
(S_2{output} <-> S_5{phase = 0//1⋅π})
(S_3{input} <-> S_6{phase = 0//1⋅π})
(S_4{output} <-> S_6{phase = 0//1⋅π})
(S_5{phase = 0//1⋅π} <-> S_6{phase = 0//1⋅π})

```
"""
function ZXGraph(zxd::ZXDiagram{T, P}) where {T, P}
    zxd = copy(zxd)
    nzxd = copy(zxd)

    simplify!(Rule{:i1}(), nzxd)
    simplify!(Rule{:h}(), nzxd)
    simplify!(Rule{:i2}(), nzxd)
    match_f = match(Rule{:f}(), nzxd)
    while length(match_f) > 0
        for m in match_f
            vs = m.vertices
            if check_rule(Rule{:f}(), nzxd, vs)
                rewrite!(Rule{:f}(), nzxd, vs)
                v1, v2 = vs
                set_phase!(zxd, v1, phase(zxd, v1) + phase(zxd, v2))
                set_phase!(zxd, v2, zero(P))
            end
        end
        match_f = match(Rule{:f}(), nzxd)
    end

    vs = spiders(nzxd)
    vH = T[]
    vZ = T[]
    vB = T[]
    for v in vs
        if spider_type(nzxd, v) == SpiderType.H
            push!(vH, v)
        elseif spider_type(nzxd, v) == SpiderType.Z
            push!(vZ, v)
        else
            push!(vB, v)
        end
    end
    eH = [(neighbors(nzxd, v, count_mul = true)[1], neighbors(nzxd, v, count_mul = true)[2]) for v in vH]

    rem_spiders!(nzxd, vH)
    et = Dict{Tuple{T, T}, EdgeType.EType}()
    for e in edges(nzxd.mg)
        et[(src(e), dst(e))] = EdgeType.SIM
    end
    zxg = ZXGraph{T, P}(nzxd.mg, nzxd.ps, nzxd.st, et, nzxd.layout, nzxd.phase_ids, nzxd.scalar, zxd, nzxd.inputs, nzxd.outputs)

    for e in eH
        v1, v2 = e
        add_edge!(zxg, v1, v2)
    end

    return zxg
end

Graphs.has_edge(zxg::ZXGraph, vs...) = has_edge(zxg.mg, vs...)
Graphs.nv(zxg::ZXGraph) = nv(zxg.mg)
Graphs.ne(zxg::ZXGraph) = ne(zxg.mg)
Graphs.outneighbors(zxg::ZXGraph, v::Integer) = outneighbors(zxg.mg, v)
Graphs.inneighbors(zxg::ZXGraph, v::Integer) = inneighbors(zxg.mg, v)
Graphs.neighbors(zxg::ZXGraph, v::Integer) = neighbors(zxg.mg, v)
Graphs.degree(zxg::ZXGraph, v::Integer) = degree(zxg.mg, v)
Graphs.indegree(zxg::ZXGraph, v::Integer) = degree(zxg, v)
Graphs.outdegree(zxg::ZXGraph, v::Integer) = degree(zxg, v)
function Graphs.rem_edge!(zxg::ZXGraph, v1::Integer, v2::Integer)
    if rem_edge!(zxg.mg, v1, v2)
        delete!(zxg.et, (min(v1, v2), max(v1, v2)))
        return true
    end
    return false
end

function Graphs.add_edge!(zxg::ZXGraph, v1::Integer, v2::Integer, edge_type::EdgeType.EType = EdgeType.HAD)
    if has_vertex(zxg.mg, v1) && has_vertex(zxg.mg, v2)
        if v1 == v2
            if edge_type == EdgeType.HAD
                set_phase!(zxg, v1, phase(zxg, v1)+1)
                add_power!(zxg, -1)
            end
            return true
        else
            if has_edge(zxg, v1, v2)
                if is_hadamard(zxg, v1, v2)
                    add_power!(zxg, -2)
                    return rem_edge!(zxg, v1, v2)
                else
                    return false
                end
            elseif add_edge!(zxg.mg, v1, v2)
                zxg.et[(min(v1, v2), max(v1, v2))] = edge_type
                return true
            end
        end
    end
    return false
end

spider_type(zxg::ZXGraph, v::Integer) = zxg.st[v]
phase(zxg::ZXGraph, v::Integer) = zxg.ps[v]
function set_phase!(zxg::ZXGraph{T, P}, v::T, p::P) where {T, P}
    if has_vertex(zxg.mg, v)
        while p < 0
            p += 2
        end
        p = rem(p, 2)
        zxg.ps[v] = p
        return true
    end
    return false
end
nqubits(zxg::ZXGraph) = zxg.layout.nbits

qubit_loc(zxg::ZXGraph{T, P}, v::T) where {T, P} = qubit_loc(zxg.layout, v)
function column_loc(zxg::ZXGraph{T, P}, v::T) where {T, P}
    c_loc = column_loc(zxg.layout, v)
    if c_loc !== nothing
        if spider_type(zxg, v) == SpiderType.Out
            nb = neighbors(zxg, v)
            if length(nb) == 1
                nb = nb[1]
                spider_type(zxg, nb) == SpiderType.In && return 3//1
                c_loc = floor(column_loc(zxg, nb) + 2)
            else
                c_loc = 1000
            end
        end
        if spider_type(zxg, v) == SpiderType.In
            nb = neighbors(zxg, v)[1]
            spider_type(zxg, nb) == SpiderType.Out && return 1//1
            c_loc = ceil(column_loc(zxg, nb) - 2)
        end
    end
    return c_loc
end

function is_hadamard(zxg::ZXGraph, v1::Integer, v2::Integer)
    if has_edge(zxg, v1, v2)
        src = min(v1, v2)
        dst = max(v1, v2)
        return zxg.et[(src, dst)] == EdgeType.HAD
    end
    return false
end
spiders(zxg::ZXGraph) = vertices(zxg.mg)

function rem_spiders!(zxg::ZXGraph{T, P}, vs::Vector{T}) where {T, P}
    if rem_vertices!(zxg.mg, vs)
        for v in vs
            delete!(zxg.ps, v)
            delete!(zxg.st, v)
            delete!(zxg.phase_ids, v)
            rem_vertex!(zxg.layout, v)
        end
        return true
    end
    return false
end
rem_spider!(zxg::ZXGraph{T, P}, v::T) where {T, P} = rem_spiders!(zxg, [v])

function add_spider!(zxg::ZXGraph{T, P}, st::SpiderType.SType, phase::P = zero(P), connect::Vector{T}=T[]) where {T<:Integer, P}
    v = add_vertex!(zxg.mg)[1]
    set_phase!(zxg, v, phase)
    zxg.st[v] = st
    if st in (SpiderType.Z, SpiderType.X)
        zxg.phase_ids[v] = (v, 1)
    end
    if all(has_vertex(zxg.mg, c) for c in connect)
        for c in connect
            add_edge!(zxg, v, c)
        end
    end
    return v
end
function insert_spider!(zxg::ZXGraph{T, P}, v1::T, v2::T, phase::P = zero(P)) where {T<:Integer, P}
    v = add_spider!(zxg, SpiderType.Z, phase, [v1, v2])
    rem_edge!(zxg, v1, v2)
    return v
end

tcount(cir::ZXGraph) = sum([phase(cir, v) % 1//2 != 0 for v in spiders(cir)])

function print_spider(io::IO, zxg::ZXGraph{T}, v::T) where {T<:Integer}
    st_v = spider_type(zxg, v)
    if st_v == SpiderType.Z
        printstyled(io, "S_$(v){phase = $(phase(zxg, v))"*(zxg.ps[v] isa Phase ? "}" : "⋅π}"); color = :green)
    elseif st_v == SpiderType.In
        print(io, "S_$(v){input}")
    elseif st_v == SpiderType.Out
        print(io, "S_$(v){output}")
    end
end

function Base.show(io::IO, zxg::ZXGraph{T}) where {T<:Integer}
    println(io, "ZX-graph with $(nv(zxg)) vertices and $(ne(zxg)) edges:")
    vs = sort!(spiders(zxg))
    for i = 1:length(vs)
        for j = i+1:length(vs)
            if has_edge(zxg, vs[i], vs[j])
                print(io, "(")
                print_spider(io, zxg, vs[i])
                if is_hadamard(zxg, vs[i], vs[j])
                    printstyled(io, " <-> "; color = :blue)
                else
                    print(io, " <-> ")
                end
                print_spider(io, zxg, vs[j])
                print(io, ")\n")
            end
        end
    end
end

function round_phases!(zxg::ZXGraph{T, P}) where {T<:Integer, P}
    ps = zxg.ps
    for v in keys(ps)
        while ps[v] < 0
            ps[v] += 2
        end
        ps[v] = rem(ps[v], 2)
    end
end

"""
    is_interior(zxg::ZXGraph, v)

Return `true` if `v` is a interior spider of `zxg`.
"""
function is_interior(zxg::ZXGraph{T, P}, v::T) where {T, P}
    if has_vertex(zxg.mg, v)
        (spider_type(zxg, v) == SpiderType.In || spider_type(zxg, v) == SpiderType.Out) && return false
        for u in neighbors(zxg, v)
            if spider_type(zxg, u) == SpiderType.In || spider_type(zxg, u) == SpiderType.Out
                return false
            end
        end
        return true
    end
    return false
end

get_outputs(zxg::ZXGraph) = get_outputs(zxg.master)
get_inputs(zxg::ZXGraph) = get_inputs(zxg.master)

# TODO: remove it?
function spider_sequence(zxg::ZXGraph{T, P}) where {T, P}
    nbits = nqubits(zxg)
    if nbits > 0
        vs = spiders(zxg)
        spider_seq = Vector{Vector{T}}(undef, nbits)
        for q = 1:nbits
            spider_seq[q] = Vector{T}()
        end
        for v in vs
            if qubit_loc(zxg, v) !== nothing
                q_loc = Int(qubit_loc(zxg, v))
                q_loc > 0 && push!(spider_seq[q_loc], v)
            end
        end
        for q = 1:nbits
            sort!(spider_seq[q], by = (v -> column_loc(zxg, v)))
        end
        return spider_seq
    end
end

function generate_layout!(zxg::ZXGraph{T, P}) where {T, P}
    layout = zxg.layout
    nbits = length(zxg.inputs)
    vs_frontier = copy(zxg.inputs)
    vs_generated = Set(vs_frontier)
    for i = 1:nbits
        set_qubit!(layout, vs_frontier[i], i)
        set_column!(layout, vs_frontier[i], 1//1)
    end

    curr_col = 1//1

    while !(isempty(vs_frontier))
        vs_after = Set{Int}()
        for v in vs_frontier
            nb_v = neighbors(zxg, v)
            for v1 in nb_v
                if !(v1 in vs_generated) && !(v1 in vs_frontier)
                    push!(vs_after, v1)
                end
            end
        end
        for i = 1:length(vs_frontier)
            v = vs_frontier[i]
            set_loc!(layout, v, i, curr_col)
            push!(vs_generated, v)
        end
        vs_frontier = collect(vs_after)
        curr_col += 1
    end
    gad_col = 2//1
    for v in spiders(zxg)
        if degree(zxg, v) == 1 && spider_type(zxg, v) == SpiderType.Z
            v1 = neighbors(zxg, v)[1]
            set_loc!(layout, v, -1//1, gad_col)
            set_loc!(layout, v1, 0//1, gad_col)
            push!(vs_generated, v, v1)
            gad_col += 1
        elseif degree(zxg, v) == 0
            set_loc!(layout, v, 0//1, gad_col)
            gad_col += 1
            push!(vs_generated, v)
        end
    end
    for q = 1:length(zxg.outputs)
        set_loc!(layout, zxg.outputs[q], q, curr_col + 1)
        set_loc!(layout, neighbors(zxg, zxg.outputs[q])[1], q, curr_col)
    end
    for q = 1:length(zxg.inputs)
        set_qubit!(layout, neighbors(zxg, zxg.inputs[q])[1], q)
    end
    return layout
end

scalar(zxg::ZXGraph) = zxg.scalar

function add_global_phase!(zxg::ZXGraph{T, P}, p::P) where {T, P}
    add_phase!(zxg.scalar, p)
    return zxg
end

function add_power!(zxg::ZXGraph, n)
    add_power!(zxg.scalar, n)
    return zxg
end


text:
module JNI

export init_new_vm, destroy_vm, is_jni_loaded, is_env_loaded,
    # Types.jl
    jint, jlong, jbyte, jboolean, jchar, jshort, jfloat, jdouble, jsize,
    jvoid, jobject, jclass, jthrowable, jweak, jmethodID, jfieldID, jstring, jarray,
    JNINativeMethod, jobjectArray, jbooleanArray, jbyteArray, jshortArray, jintArray,
    jlongArray, jfloatArray, jdoubleArray, jcharArray, jvalue, jobjectRefType,
    # Constants.jl
    JNI_FALSE, JNI_TRUE

include("Threads.jl")
include("Constants.jl")
include("Types.jl")
include("Interfaces.jl")

import Libdl

using .Threads
using .Constants
using .Types
using .Interfaces

# Global variables

# JavaVM **
const jvmptrref = Ref(Ptr{JavaVM}(C_NULL))
# JNIEnv ** (array of JNIEnv used by different threads)
const jnienvptrs = [Ptr{JNIEnv}(C_NULL)]

const jninativeinterfaceref = Ref(JNINativeInterface())
const jniinvokeinterfaceref = Ref(JNIInvokeInterface())

# API Structs
const Char = UInt8

struct JavaVMOption
    optionString::Ptr{Char}
    extraInfo::Ptr{Nothing}
end

JavaVMOption(optionString::Ptr{Char}) = JavaVMOption(optionString, C_NULL)

struct JavaVMInitArgs
    version::jint
    nOptions::jint
    options::Ptr{JavaVMOption}
    ignoreUnrecognized::jboolean
end

JavaVMInitArgs(version::jint, jopts::Vector{JavaVMOption}, ignoreUnrecognized::jboolean) =
    JavaVMInitArgs(version, convert(jint, length(jopts)), convert(Ptr{JavaVMOption}, pointer(jopts)), ignoreUnrecognized)

struct JNIError <: Exception
    msg::String
end

# JNI API
"""
    init_new_vm(opts)

Initialize a new Java virtual machine.
"""
function init_new_vm(libpath,opts)
    libjvm = Libdl.dlopen(libpath)
    create = Libdl.dlsym(libjvm, :JNI_CreateJavaVM)
	jopts = [JavaVMOption(pointer(x)) for x in opts]
    Threads.resize_nthreads!(jnienvptrs)
    GC.@preserve jvmptrref jnienvptrs jopts begin
        vm_args = JavaVMInitArgs(JNI_VERSION_1_8, jopts, JNI_TRUE)
        res = @ccall $create(jvmptrref::Ref{Ptr{JavaVM}}, jnienvptrs::Ref{Ptr{JNIEnv}}, Ref(vm_args)::Ptr{JavaVMInitArgs})::Cint
        res < 0 && throw(JNIError("Unable to initialise Java VM: $(res)"))
    end
    jvm = unsafe_load(jvmptrref[])
    jniinvokeinterfaceref[] = unsafe_load(jvm)
    jnienv = unsafe_load(jnienvptrs[1])
    jninativeinterfaceref[] = unsafe_load(jnienv)
    _attachthreads()
    return
end

function destroy_vm()
    if !is_env_loaded()
        throw(JNIError("Called destroy without initialising Java VM"))
    end
    _detachthreads()
    destroy_fn = jniinvokeinterfaceref[].DestroyJavaVM
    res = @ccall $destroy_fn(jvmptrref[]::Ptr{JavaVM})::Cint
    res < 0 && throw(JavaCallError("Unable to destroy Java VM"))
    _resetglobalvars()
    nothing
end

is_jni_loaded() = jninativeinterfaceref[].GetVersion != C_NULL
is_env_loaded() = jnienvptrs[1] != C_NULL

# Include file with generated jni interface
include("generated_jni_interface.jl")

# Private functions

function _resetglobalvars()
    jniinvokeinterfaceref[] = JNIInvokeInterface()
    empty!(jnienvptrs)
    push!(jnienvptrs, Ptr{JNIEnv}(C_NULL))
    jninativeinterfaceref[] = JNINativeInterface()
    jvmptrref[] = Ptr{JavaVM}(C_NULL)
end

function _attachthreads()
    Threads.@threads for i=1:Threads.nthreads()
        _attachthread(Ref(jnienvptrs, Threads.threadid()))
    end
end

function _attachthread(ppenv_thread = Ref{Ptr{JNIEnv}}(C_NULL))
    res = ccall(jniinvokeinterfaceref[].AttachCurrentThread, Cint, (Ptr{Nothing}, Ptr{Ptr{JNIEnv}}, Ptr{Nothing}), jvmptrref[], ppenv_thread, C_NULL)
    res < 0 && throw(JNIError("Unable to attach thread id: $(Threads.threadid())"))
    return ppenv_thread[]
end

function _detachthreads()
    Threads.@threads for i=1:Threads.nthreads()
        _detachthread()
    end
    nothing
end

function _detachthread()
    res = ccall(jniinvokeinterfaceref[].DetachCurrentThread, Cint, (Ptr{JavaVM},), jvmptrref[])
    res < 0 && throw(JNIError("Unable to detach thread id: $(Threads.threadid())"))
    nothing
end

end

ext: jl


text:
module Reproduction

function CAWS(current, other...; brood_size=3)
    averages = mean.(eachrow(hcat(current, other...)))
    make_sol(averages) = [rand() < chance for chance in averages]
    unique([make_sol(averages) for _ in 1:brood_size])
end

function GA(a, b)
    i = rand(1:length(a))
    vcat(a[1:i], b[i+1:end])
end

function jaya(current, better, worse)

end

function jaya_simp(current::BitArray, better, worse)

end

function jaya_trad(current::BitArray, better, worse)

end

function TBO(current::BitArray, best, mean)

end

function CBO(current, worst, mean)

end

function LBO()

end

function Rao1()

end

function Rao2()

end

function mutate()

end

end

max_stars_repo_name: Dibillilia/JuliaOptHeuristics

max_issues_repo_name: Dibillilia/JuliaOptHeuristics

max_forks_repo_name: Dibillilia/JuliaOptHeuristics


text:
"""
    This file generates an array saved as a plain text that contains the integrated trajectory together with each marginal theoretical distribution. The normalization factor associated with the distribution in `q` is passed as an argument.
    Example:
    ```
    julia> results = marginaldistributions(filename, Potential("quartic", x->x^4/4.), Thermostat("logistic", A, x -> exp(x/A)/(A*(1.+ exp(x/A)).^2)), normalizationfactor)
    ```
    """

function marginaldistributions(data::Matrix{Float64},p::Parameters, normalizationfactor::Float64)

    beta = 1./p.T

    rhoq = 1./normalizationfactor*exp(-beta*Float64[p.potential.f(i) for i in data[:,2]])
    rhop = sqrt(beta/(2.*pi))*exp(-beta*data[:,3].^2/2.)
    rhoz = [p.thermo.distribution(i) for i in data[:,4]]
    q1 = hcat(data[:,2], rhoq)
    p1 = hcat(data[:,3], rhop)
    z1 = hcat(data[:,4], rhoz)
    results = hcat(q1, p1)
    results = hcat(results,z1)

end

max_stars_repo_name: dapias/STODs

max_issues_repo_name: dapias/STODs

max_forks_repo_name: dapias/STODs


text:
#=

# Wrapper types

The Vulkan API possesses a few data structures that exhibit a different behavior. Each structure type has been wrapped carefully to automate the underlying API patterns. We list all of these here along with their properties and features that we hope will free the developer of some of the heaviest patterns and boilerplate of the Vulkan API.

## Handles

Handles are opaque pointers to internal Vulkan objects. Almost every handle must be created and destroyed with API commands. Some handles have a parent handle (see [Parent handle access](@ref Parent-handle-access) for navigating through the resulting handle hierarchy), which *must not* be destroyed before its children. For this we provide wrappers around creation functions with an automatic finalization feature (see [Automatic finalization](@ref Automatic-finalization)) that uses a simple reference couting system. This alleviates the burden of tracking when a handle can be freed and freeing it, in conformance with the Vulkan specification.

Most handles are typically created with a `*CreateInfo` or `*AllocateInfo` structure, that packs creation parameters to be provided to the API creation function. To allow for nice one-liners that don't involve long create info names, [these create info parameters are exposed](@ref expose-create-info-args) in the creation function, automatically building the create info structure.

!!! tip
    Most handle types have constructors defined that wrap around the creation function and automatically unwrap the result (see [Error handling](@ref Error-handling)).

### Automatic finalization

In the Vulkan API, handles are created with the functions `vkCreate*` and `vkAllocate*`, and most of them must be destroyed after use with a call to `vkDestroy*` or `vkFree*`. More importantly, they must be destroyed with the same allocator and parent handle that created them.

To automate this, new mutable handle types were defined to allow for the registration of a [finalizer](https://docs.julialang.org/en/v1/base/base/#Base.finalizer). The `create_*` and `allocate_*` wrappers automatically register the corresponding destructor in a finalizer, so that you don't need to worry about destructors (except for `CommandBuffer`s and `DescriptorSet`s, see below). The finalizer of a handle, and therefore its API destructor, will execute when there are no program-accessible references to this handle. Because finalizers may run in arbitrary order in Julia, and some handle types such as `VkDevice` require to be destroyed only after all their children, a simple thread-safe reference counting system is used to make sure that a handle is destroyed **only after all its children are destroyed**.

As an exception, because they are meant to be freed in batches, `CommandBuffer`s and `DescriptorSet`s do not register any destructor and are not automatically freed. Those handles will have to explicitly freed with [`free_command_buffers`](@ref) and [`free_descriptor_sets`](@ref) respectively.

Finalizers can be run eagerly with [`finalize`](https://docs.julialang.org/en/v1/base/base/#Base.finalize), which allows one to reclaim resources early. The finalizers won't run twice if triggered manually.

!!! danger
    You should **never** explicitly call a destructor, except for `CommandBuffer` and `DescriptorSet`. Otherwise, the object will be destroyed twice and will lead to a segmentation fault.

!!! note
    If you need to construct a handle from an opaque pointer (obtained, for example, via an external library such as a `VkSurfaceKHR` from GLFW), you can use the constructor `(::Type{<:Handle})(ptr::Ptr{Cvoid}, destructor[, parent])` as in

    ```julia
    surface_ptr = GLFW.CreateWindowSurface(instance, window)
    SurfaceKHR(surface_ptr, x -> destroy_surface_khr(instance, x), instance)
    ```

    If the surface doesn't need to be destroyed (for example, if the external library does it automatically), the `identity` function should be passed in as destructor.

### Handle constructors

Handles that can only be created with a single API constructor have constructors defined which wrap the relevant create/allocate\* function and `unwrap` the result.

For example, `Instance(layers, extensions)` is equivalent to `unwrap(create_instance(layers, extensions))`.

If the API constructor returns an error, an exception will be raised (see [Error handling](@ref)).

### Parent handle access

Handles store their parent handle if they have one. For example, `Pipeline`s have a `device` field as a [`Device`](@ref), which itself contains a `physical_device` field and so on until the instance that has no parent. This reduces the number of objects that must be carried around in user programs.

`Base.parent` was extended to navigate this hierarchy, where for example `parent(device) == device.physical_device` and `parent(physical_device) == physical_device.instance`.

## Structures

Vulkan structures, such as `Extent2D`, `InstanceCreateInfo` and `PhysicalDeviceFeatures` were wrapped into two different structures each: a high-level structure, which should be used most of the time, and an intermediate structure used for maximal performance whenever required.

### High-level structures

High-level structures were defined to ressemble idiomatic Julia structures, replacing C types by idiomatic Julia types. They abstract most pointers away, using Julia arrays and strings, and use `VersionNumbers` instead of integers. Equality and hashing are implemented with [AutoHashEquals.jl](https://github.com/andrewcooke/AutoHashEquals.jl) to facilitate their use in dictionaries.

### Intermediate structures

Intermediate structures wrap C-compatible structures and embed pointer data as dependencies. Therefore, as long as the intermediate structure lives, all pointer data contained within the C-compatible structure will be valid. These structures are mostly used internally by Vulkan.jl, but they can be used with [intermediate functions](@ref Intermediate-functions) for maximum performance, avoiding the overhead incurred by high-level structures which require back and forth conversions with C-compatible structures for API calls.

These intermediate structures share the name of the high-level structures, starting with an underscore. For example, the high-level structure [`InstanceCreateInfo`](@ref) has an intermediate counterpart [`_InstanceCreateInfo`](@ref).

Note that intermediate structures can only be used with other intermediate structures. `convert` methods allow the conversion between arbitrary high-level and intermediate structures, if required.

!!! tip
    Outside performance-critical sections such as tight loops, high-level structures are much more convenient to manipulate and should be used instead.

## Bitmask flags

In the Vulkan API, certain flags use a bitmask structure. A bitmask is a logical `or` combination of several bit values, whose meaning is defined by the bitmask type. In Vulkan, the associated flag type is defined as a `UInt32`, which allows any value to be passed in as a flag. This opens up the door to incorrect usage that may be hard to debug. To circumvent that, most bitmask flags were wrapped with an associated type which prevents combinations with flags of other bitmask types.

For example, consider the core `VkSampleCountFlags` type (alias for `UInt32`) with bits defined via the enumerated type `VkSampleCountFlagBits`:

```@repl
using Vulkan.VkCore
VK_SAMPLE_COUNT_1_BIT isa VkSampleCountFlagBits
VK_SAMPLE_COUNT_1_BIT === VkSampleCountFlagBits(1)
VK_SAMPLE_COUNT_1_BIT === VkSampleCountFlags(1)
VK_SAMPLE_COUNT_1_BIT | VK_SAMPLE_COUNT_2_BIT === VkSampleCountFlags(3)
VK_SAMPLE_COUNT_1_BIT & VK_SAMPLE_COUNT_2_BIT === VkSampleCountFlags(0)
VK_SAMPLE_COUNT_1_BIT & VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR === VkSampleCountFlags(1)
```

Those two types are combined into one `SampleCountFlag`:

```@repl
using Vulkan
SampleCountFlag <: BitMask
SurfaceTransformFlagKHR <: BitMask # another bitmask flag
SAMPLE_COUNT_1_BIT | SAMPLE_COUNT_2_BIT === SampleCountFlag(3)
SAMPLE_COUNT_1_BIT & SAMPLE_COUNT_2_BIT === SampleCountFlag(0)
SAMPLE_COUNT_1_BIT & SURFACE_TRANSFORM_IDENTITY_BIT_KHR
UInt32(typemax(SampleCountFlag)) === UInt32(VkCore.VK_SAMPLE_COUNT_FLAG_BITS_MAX_ENUM)
```

All functions that were expecting a `VkSampleCountFlags` (`UInt32`) value will have their wrapped versions expect a value of type `SampleCountFlag`. Furthermore, the `*FLAG_BITS_MAX_ENUM` values are removed. This value is the same for all enums and can be accessed via `typemax(T)` where `T` is a `BitMask` (e.g. `SampleCountFlag`).

=#

max_stars_repo_name: JuliaGPU/VulkanAbstraction

max_issues_repo_name: JuliaGPU/VulkanAbstraction

max_forks_repo_name: JuliaGPU/VulkanAbstraction


text:
using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.Model{Float64}())
config = MOIT.TestConfig()
config_with_basis = MOIT.TestConfig(basis = true)

@testset "Interval" begin
    bridged_mock = MOIB.Constraint.SplitInterval{Float64}(mock)
    MOIT.basic_constraint_tests(bridged_mock, config,
                                include=[(MOI.SingleVariable,
                                          MOI.Interval{Float64}),
                                         (MOI.ScalarAffineFunction{Float64},
                                          MOI.Interval{Float64}),
                                         (MOI.ScalarQuadraticFunction{Float64},
                                          MOI.Interval{Float64})])
    MOIU.set_mock_optimize!(mock,
         (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [5.0, 5.0], con_basis =
             [(MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [MOI.BASIC],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [MOI.NONBASIC],
              (MOI.SingleVariable, MOI.GreaterThan{Float64})                => [MOI.BASIC, MOI.BASIC]],
              (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [0],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [-1]),
         (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [2.5, 2.5], con_basis =
             [(MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [MOI.NONBASIC],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [MOI.BASIC],
              (MOI.SingleVariable, MOI.GreaterThan{Float64})                => [MOI.BASIC, MOI.BASIC]],
              (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [1],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [0]),
         (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0, 1.0],
                  (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [1],
                  (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [0],
              con_basis =
             [(MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [MOI.NONBASIC],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [MOI.BASIC],
              (MOI.SingleVariable, MOI.GreaterThan{Float64})                => [MOI.BASIC, MOI.BASIC]]),
         (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [6.0, 6.0], con_basis =
             [(MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [MOI.BASIC],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [MOI.NONBASIC],
              (MOI.SingleVariable, MOI.GreaterThan{Float64})                => [MOI.BASIC, MOI.BASIC]]))
    MOIT.linear10test(bridged_mock, config_with_basis)
    MOIU.set_mock_optimize!(mock,
         (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [0.0, 0.0],
               (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [0],
               (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [-1],
              con_basis =
             [(MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [MOI.BASIC],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [MOI.BASIC],
              (MOI.SingleVariable, MOI.GreaterThan{Float64})                => [MOI.NONBASIC, MOI.NONBASIC]],
              (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [0],
              (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64})    => [0]))
    MOIT.linear10btest(bridged_mock, config_with_basis)

    ci = first(MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.Interval{Float64}}()))
    vis = MOI.get(bridged_mock, MOI.ListOfVariableIndices())
    newf = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([1.0, -1.0], vis), 0.0)
    MOI.set(bridged_mock, MOI.ConstraintFunction(), ci, newf)
    @test MOI.get(bridged_mock, MOI.ConstraintFunction(), ci) ≈ newf

    MOI.modify(bridged_mock, ci, MOI.ScalarCoefficientChange(vis[2], 1.0))
    modified_f = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(ones(2), vis), 0.0)
    @test MOI.get(bridged_mock, MOI.ConstraintFunction(), ci) ≈ modified_f

    test_delete_bridge(bridged_mock, ci, 2, ((MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}, 0),
                                            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64},    0)))
end


text:
module NewtonsMethod

using LinearAlgebra, ForwardDiff

function newtonroot(f, f_prime; x_0,  tol = 1E-7, maxiter = 1000)
    x_old=x_0
    normdiff=Inf
    iter=1

    if(abs(f_prime(x_old))<=tol)
        return (value=x_old,normdiff=f(x_old),iter=0)
    end

    while normdiff>tol && iter<=maxiter
        x_new= x_old - f(x_old)/f_prime(x_old)
        if isnan(x_new)
            return (nothing, nothing, nothing)
        end
        normdiff=norm(x_new-x_old)
        x_old=x_new
        iter = iter + 1

    end
    if(normdiff>tol)
        return (nothing, nothing, nothing)
    end

    return (value=x_old,normdiff=normdiff,iter=iter)
end

function newtonroot(f; x_0, tol = 1E-7, maxiter = 1000)
    #find derivate of f

    x_old=x_0
    normdiff=Inf
    iter=1

    D(f) = x -> ForwardDiff.derivative(f, x)
    f_prime=D(f)

    if(abs(f_prime(x_old))<=tol)
        return (value=x_old,normdiff=f(x_old),iter=0)
    end

    while normdiff>tol && iter<=maxiter
        x_new= x_old - f(x_old)/f_prime(x_old)
        if isnan(x_new)
            return (nothing, nothing, nothing)
        end
        normdiff=norm(x_new-x_old)
        x_old=x_new

        iter += 1
    end
    if(normdiff>tol)
        return (nothing, nothing, nothing)
    end
    return (value=x_old,normdiff=normdiff,iter=iter)
end

end


text:
using MLJBase

docstring = (Base.Docs.doc)((Base.Docs.Binding)(Main, :multiclass_recall))

@test string(docstring) == "An instance of type "*
    "[`MulticlassTruePositiveRate`](@ref). Query the "*
    "[`MulticlassTruePositiveRate`](@ref) doc-string for details. \n"

true


text:
using LittleManComputer
using Test

@testset "All Tests" begin

include("assem_tests.jl")
include("disassem_tests.jl")
include("simulator_tests.jl")

end


text:
using Documenter, SemiparametricAUC

makedocs(
  modules  = [SemiparametricAUC],
  clean    = false,
  format   = :html,
  sitename = :"Perform semiparametric AUC model in Julia",
  authors  = "Som B. Bohora",
  # assets   = ["assets/favicon.ico"],
  pages    = [
              "Introduction"    => "index.md",
              "Installation"    => "install.md",
              "Article"         => "paper.md",
              "Example"         => "example-julia.md",
              "References"      => "command.md",
              "Code"            => "code.md",
              "sAUC in R"       => "example.md",
              "sAUC in R Shiny" => "r-shiny.md",
              "saucpy in Python"=> "example-python.md",
              "Developer"       => "bohora.md",
              "Report Bugs"     => "bug-report.md"
  ]
)

deploydocs(
    deps   = Deps.pip("mkdocs", "python-markdown-math"),
    target = "build",
    repo   = "github.com/sbohora/SemiparametricAUC.jl.git",
    make   = nothing,
    julia  = "release",
    osname = "osx"
)


text:
function render!(
    rendermodel::RenderModel,
    scene::EntityFrame{S,D,I};
    car_color::Colorant=COLOR_CAR_OTHER, # default color
    car_colors::Dict{I,Colorant}=Dict{I,Colorant}(), #  id -> color
    ) where {S,D,I}

    for veh in scene
        render!(rendermodel, veh, get(car_colors, veh.id, car_color))
    end

    rendermodel
end

function render(roadway::R;
    canvas_width::Int=DEFAULT_CANVAS_WIDTH,
    canvas_height::Int=DEFAULT_CANVAS_HEIGHT,
    rendermodel = RenderModel(),
    cam::Camera = FitToContentCamera(),
    ) where {R<:Roadway}

    s = CairoRGBSurface(canvas_width, canvas_height)
    ctx = creategc(s)
    clear_setup!(rendermodel)

    render!(rendermodel, roadway)

    camera_set!(rendermodel, cam, canvas_width, canvas_height)
    render(rendermodel, ctx, canvas_width, canvas_height)
    return s
end

function render(ctx::CairoContext, scene::EntityFrame{S,D,I}, roadway::R;
    rendermodel::RenderModel=RenderModel(),
    cam::Camera=SceneFollowCamera(),
    car_colors::Dict{I,Colorant}=Dict{I,Colorant}(),
    ) where {S,D,I,R}

    canvas_width = floor(Int, Cairo.width(ctx))
    canvas_height = floor(Int, Cairo.height(ctx))

    clear_setup!(rendermodel)

    render!(rendermodel, roadway)
    render!(rendermodel, scene, car_colors=car_colors)

    camera_set!(rendermodel, cam, scene, roadway, canvas_width, canvas_height)

    render(rendermodel, ctx, canvas_width, canvas_height)
    ctx
end
function render(scene::EntityFrame{S,D,I}, roadway::R;
    canvas_width::Int=DEFAULT_CANVAS_WIDTH,
    canvas_height::Int=DEFAULT_CANVAS_HEIGHT,
    rendermodel::RenderModel=RenderModel(),
    cam::Camera=SceneFollowCamera(),
    car_colors::Dict{I,Colorant}=Dict{I,Colorant}(), # id
    ) where {S,D,I,R}

    s, ctx = get_surface_and_context(canvas_width, canvas_height)
    render(ctx, scene, roadway, rendermodel=rendermodel, cam=cam, car_colors=car_colors)
    s
end

function get_pastel_car_colors(scene::EntityFrame{S,D,I}; saturation::Float64=0.85, value::Float64=0.85) where {S,D,I}
    retval = Dict{I,Colorant}()
    n = length(scene)
    for (i,veh) in enumerate(scene)
        retval[veh.id] = convert(RGB, HSV(180*(i-1)/max(n-1,1), saturation, value))
    end
    return retval
end


text:
import LinearAlgebra.mul!

function add!(zd1::VTUDataField{Float64}, zd2::VTUDataField{Float64})
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] += zd2.dat[i]
	end
	return nothing
end

function minus!(zd1::VTUDataField{Float64}, zd2::VTUDataField{Float64})
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] -= zd2.dat[i]
	end
	return nothing
end

function add!(zd1::VTUDataField{Float64}, a::Number)
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] += a
	end
	return nothing
end

function mul!(zd1::VTUDataField{Float64}, c::Number)
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] *= c
	end
	return nothing
end

function mul!(zd1::VTUDataField{Float64}, zd2::VTUDataField{Float64}, zd3::VTUDataField{Float64})
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] = zd2.dat[i] * zd3.dat[i]
	end
	return nothing
end

function mul!(zd1::VTUDataField{Float64}, zd2::VTUDataField{Float64}, fac::Float64)
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] = zd2.dat[i] * fac
	end
	return nothing
end

function pow!(zd1::VTUDataField{Float64}, a::Number)
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] ^= a
	end
	return nothing
end

function div!(zd1::VTUDataField{Float64}, a::Number)
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] /= a
	end
	return nothing
end

function max!(zd1::VTUDataField{Float64}, zd2::VTUDataField{Float64})
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] = max(zd1.dat[i],zd2.dat[i])
	end
	return nothing
end

function min!(zd1::VTUDataField{Float64}, zd2::VTUDataField{Float64})
	@inbounds for i in 1:length(zd1.dat)
		zd1.dat[i] = min(zd1.dat[i],zd2.dat[i])
	end
	return nothing
end

function add!(zd1::VTUData, zd2::VTUData)
	for (dat1,dat2) in zip(zd1.interp_data,zd2.interp_data)
		add!(dat1,dat2)
	end
	return nothing
end

function minus!(zd1::VTUData, zd2::VTUData)
	for (dat1,dat2) in zip(zd1.interp_data,zd2.interp_data)
		minus!(dat1,dat2)
	end
	return nothing
end

function add!(zd1::VTUData, a::Number)
	for dat in zd1.interp_data
		add!(dat,a)
	end
	return nothing
end

function mul!(zd1::VTUData, c::Number)
	for dat in zd1.interp_data
		mul!(dat,c)
	end
	return nothing
end

function mul!(zd1::VTUData, zd2::VTUData, zd3::VTUData)
	for (dat1,dat2,dat3) in zip(zd1.interp_data,zd2.interp_data,zd3.interp_data)
		mul!(dat1,dat2,dat3)
	end
	return nothing
end

function mul!(zd1::VTUData, zd2::VTUData, fac::Float64)
	for (dat1,dat2) in zip(zd1.interp_data,zd2.interp_data)
		mul!(dat1,dat2,fac)
	end
	return nothing
end

function pow!(zd1::VTUData, a::Number)
	for dat in zd1.interp_data
		pow!(dat,a)
	end
	return nothing
end

function div!(zd1::VTUData, a::Number)
	for dat in zd1.interp_data
		div!(dat,a)
	end
	return nothing
end

function div!(zd1::VTUData, zd2::VTUData, zd3::VTUData)
	for (dat1,dat2,dat3) in zip(zd1.interp_data,zd2.interp_data,zd3.interp_data)
		div!(dat1,dat2,dat3)
	end
	return nothing
end

function max!(zd1::VTUData, zd2::VTUData)
	for (dat1,dat2) in zip(zd1.interp_data,zd2.interp_data)
		max!(dat1,dat2)
	end
	return nothing
end

function min!(zd1::VTUData, zd2::VTUData)
	for (dat1,dat2) in zip(zd1.interp_data,zd2.interp_data)
		min!(dat1,dat2)
	end
	return nothing
end

import Base: +,-,*,/,^

function +(tpf1::VTUData, tpf2::VTUData)
	ret = similar(tpf1)
	add!(ret,tpf1)
	add!(ret,tpf2)
	return ret
end

function -(tpf1::VTUData, tpf2::VTUData)
	ret = similar(tpf1)
	add!(ret,tpf1)
	minus!(ret,tpf2)
	return ret
end

function *(tpf1::VTUData, tpf2::VTUData)
	ret = similar(tpf1)
	mul!(ret,tpf1,tpf2)
	return ret
end

function /(tpf1::VTUData, tpf2::VTUData)
	ret = similar(tpf1)
	div!(ret,tpf1,tpf2)
	return ret
end

function +(tpf::VTUData, a::Number)
	ret = similar(tpf)
	add!(ret,tpf)
	add!(ret,a)
	return ret
end
+(a::Number,tpf::VTUData) = tpf+a

function -(tpf::VTUData, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	add!(ret,-a)
	return ret
end

function -(a::T,tpf::VTUData) where T<:Number
	ret = similar(tpf)
	add!(ret,a)
	minus!(ret,tpf)
	return ret
end

function *(tpf::VTUData, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	mul!(ret,a)
	return ret
end
*(a::T,tpf::VTUData) where T<:Number = tpf*a

function /(tpf::VTUData, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	div!(ret,a)
	return ret
end

function ^(tpf::VTUData, a::T) where T<:Number
	ret = similar(tpf)
	add!(ret,tpf)
	pow!(ret,a)
	return ret
end

max_stars_repo_name: nagelt/VTUFileHandler

max_issues_repo_name: nagelt/VTUFileHandler

max_forks_repo_name: nagelt/VTUFileHandler


text:
using Distributions
using Random

using ImportanceWeightedRiskMetrics
using POMDPGym
using POMDPModelTools
using POMDPPolicies
using POMDPs
using POMDPSimulators

using ParallelTreeSampling

# Runs Monte Carlo baseline.
function run_baseline(amdp::MDP, s, distrib_fn; N=1000)
    function_policy = (s) -> rand(distrib_fn(amdp, s))
    costs = [sum(collect(simulate(HistoryRecorder(), amdp, function_policy, s)[:r]))
             for _ in 1:N]
    output = (costs, [])
    return output
end

function run_baseline(rmdp::RMDP, s, px; N=1000)
    noise_fn = (s) -> rand(px)  # Random noise.
    function_policy = FunctionPolicy(noise_fn)
    costs = [sum(collect(simulate(HistoryRecorder(), rmdp, function_policy, s)[:r]))
             for _ in 1:N]
    output = (costs, [])
    return output
end

# Runs MCTS-based sampling.
function run_mcts(mdp::MDP, s, nominal_distrib_fn;
                  N=1000, c=0.3, vloss=0.0, α=0.1, β=1.0, γ=0.0)
    solver = PISSolver(; depth=100,
                       exploration_constant=c,
                       n_iterations=N,
                       enable_action_pw=false,  # Needed for discrete cases.
                       k_state=Inf,             # Needed for discrete cases (to always transition).
                       virtual_loss=vloss,
                       nominal_distrib_fn=nominal_distrib_fn,
                       show_progress=true,
                       tree_in_info=true,
                       α=α);
    planner = solve(solver, mdp);
    a, info = action_info(planner, s; tree_in_info=true, β=β, γ=γ)
    tree = info[:tree]
    output = (tree.costs, [], tree.weights, info)
    return output, planner
end

# Returns risk metrics.
function eval_metrics(costs; weights=nothing, alpha=1e-4)
    costs = [Float64(cost) for cost in costs]
    if weights == nothing
        weights = ones(length(costs))
    else
        weights = [Float64(weight) for weight in weights]
    end

    metrics = IWRiskMetrics(costs, weights, alpha);
    return metrics
end


text:
"""
    intensity_sppa(pressures::Array{Unitful.Pressure, N}, medium::Medium, excitation::Excitation)

Spatial Peak, Pulse Averaged intensity [W/cm²]

`pressures` is expected to be an array with up to 4 dimensions. The first dimension is time,
and the remaining dimensions are space
"""
function intensity_sppa(pressures::Array{T, N}, medium, excitation) where {T <: Unitful.Pressure, N}
    # account for sampling frequency by multiplying pressure values with dt
    dt = 1/excitation.fs
    intensities = map(p -> intensity(p, medium), pressures)
    i_sppa = sum(intensities * dt; dims=1) / excitation.pulse_duration
    return format_intensity.(squeeze(i_sppa))
end

"""
    intensity(pressure::Unitful.Pressure, medium::Medium)

Instantaneous intensity [W/cm²]

`pressure` is expected to be a single number
"""
function intensity(pressure::Unitful.Pressure, medium::Medium)
    return format_intensity(pressure^2 / (medium.density * medium.c))
end

format_intensity(input) = uconvert(u"W/cm^2", input)

"""
    intensity_spta(pressures::Array{Unitful.Pressure, N}, medium::Medium, excitation::Excitation)

Spatial Peak, Time Averaged intensity [W/cm²]

`pressures` is expected to be an array with up to 4 dimensions. The first dimension is time,
and the remaining dimensions are space
"""
function intensity_spta(pressures::Array{T, N}, medium::Medium, excitation::Excitation) where {T <: Unitful.Pressure, N}
    return intensity_sppa(pressures, medium, excitation) * excitation.duty_cycle
end

"""
    energy_spta(pressures::Array{Unitful.Pressure, N}, medium::Medium, excitation::Excitation)

Spatial Peak, Time Averaged energy [J/cm²]
"""
function energy_spta(pressures::Array{T, N}, medium::Medium, excitation::Excitation) where {T <: Unitful.Pressure, N}
    return intensity_spta(pressures, medium, excitation) * excitation.total_duration
end

"""
    mechanical_index(pressures::Array{Unitful.Pressure, N}, excitation::Excitation)

Mechanical index [dimensionless]

`pressures` is expected to be an array with up to 4 dimensions. The first dimension is time,
and the remaining dimensions are space
"""
function mechanical_index(pressures::Array{T, N}, excitation::Excitation) where {T <: Unitful.Pressure, N}
    peak_rarefaction_pressure = pn_pressure(uconvert.(u"MPa", pressures))
    center_frequency = sqrt(uconvert(u"MHz", excitation.f0))
    return ustrip(peak_rarefaction_pressure / center_frequency)
end

"""
    pn_pressure(pressures::Array{Unitful.Pressure, N})

Peak Negative pressure [units unchanged]

`pressures` is expected to be an array with up to 4 dimensions. The first dimension is time,
and the remaining dimensions are space
"""
function pn_pressure(pressures::Array{T, N}) where {T <: Unitful.Pressure, N}
    peak_neg_pressure = abs.(minimum(pressures; dims=1))
    return squeeze(peak_neg_pressure)
end

"""
    pp_pressure(pressures::Array{Unitful.Pressure, N})

Peak Positive pressure [units unchanged]

`pressures` is expected to be an array with up to 4 dimensions. The first dimension is time,
and the remaining dimensions are space
"""
function pp_pressure(pressures::Array{T, N}) where {T <: Unitful.Pressure, N}
    peak_pos_pressure = abs.(maximum(pressures; dims=1))
    return squeeze(peak_pos_pressure)
end

"""
    ptp_pressure(pressures::Array{Unitful.Pressure, N})

Peak To Peak pressure [units unchanged]

`pressures` is expected to be an array with up to 4 dimensions. The first dimension is time,
and the remaining dimensions are space
"""
function ptp_pressure(pressures::Array{T, N}) where {T <: Unitful.Pressure, N}
    ptp_pressure = pp_pressure(pressures) .+ pn_pressure(pressures)
    return ptp_pressure
end


text:
"""
    convertRateToPercent
converts a rate (e.g. 0.081) to percentage (e.g. 8,1)
"""
function convertRateToPercent(rate::Float64)::Float64
    return rate*100
end

max_stars_repo_name: RainerNeu/tradeStatistics

max_issues_repo_name: RainerNeu/tradeStatistics

max_forks_repo_name: RainerNeu/tradeStatistics


text:
module ColadaSeismic

using Logging
using JUDI
using PyCall
using LinearAlgebra
using SlimOptim
using Random
using Statistics
using NLopt

# h5geo includes
include("h5geo/h5geo.jl")

# Colada includes
include("TimeModeling/inner/util.jl")
include("TimeModeling/inner/h5judi_physical_parameter.jl")
include("TimeModeling/h5judi_forward_modeling.jl")
include("TimeModeling/h5judi_rtm.jl")
include("TimeModeling/h5judi_fwi.jl")
include("TimeModeling/h5judi_twri.jl")
include("TimeModeling/h5judi_modeling.jl")

# ONLY WHEN COLADA IS NOT AVAILABLE
# h5geo = pyimport("h5geopy._h5geo")
# h5geo.sr.setSpatialReferenceFromUserInput("EPSG:32056")

end # module


text:
export visualize_wavefield, plot_result, add_initial_model,
    Ricker, Gauss, compute_lame_parameters,
    gradtest, compute_loss_and_grads_GPU, compute_forward_GPU, variable_source,
    sampling_compute_loss_and_grads_GPU,
    sampling_compute_loss_and_grads_GPU_v2,
    sampling_compute_loss_and_grads_GPU_v3,
    visualize_model, aggregate_wavefield

function visualize_model(vp::Array{Float64, 2}, params::Union{ElasticPropagatorParams,AcousticPropagatorParams, MPIAcousticPropagatorParams})

    clf()
    if isa(params, MPIAcousticPropagatorParams)
        pcolormesh([0:params.NX-1;]*params.DELTAX/1e3,[0:params.NY-1;]*params.DELTAY/1e3,  Array(vp'))
    else
        pcolormesh([0:params.NX+1;]*params.DELTAX/1e3,[0:params.NY+1;]*params.DELTAY/1e3,  Array(vp'))
    end
    axis("scaled")
    colorbar(shrink=0.4)
    xlabel("x (km)")
    ylabel("z (km)")
    gca().invert_yaxis()
    tight_layout()

end

function aggregate_wavefield(param::Union{MPIAcousticPropagatorParams, MPIElasticPropagatorParams}, vals::Array{Array{Float64, 3}, 1})
    M, N = param.M, param.N
    nt, m, n = size(vals[1])
    Vals = zeros(nt, m*M, n*N)
    k = 0
    for i = 1:M
        for j = 1:N
            k += 1
            Vals[:, (i-1)*m+1:i*m, (j-1)*n+1:j*n] = vals[k]
        end
    end
    Vals
end

function visualize_wavefield(val::Array{Array{Float64, 3}, 1},
    param::Union{MPIAcousticPropagatorParams, MPIElasticPropagatorParams})
    val = aggregate_wavefield(param, val)
    visualize_wavefield(val, param)
end

"""
    visualize_wavefield(val::Array{Float64, 3},
        param::Union{ElasticPropagatorParams,AcousticPropagatorParams})
    visualize_wavefield(val::Array{Array{Float64, 3}, 1},
        param::Union{MPIAcousticPropagatorParams, MPIElasticPropagatorParams})

Visualizes the wavefield and returns the handle of the animation. You can save the animation to gif via

```julia
p = visualize_wavefield(...)
saveanim(p, "myfigure.gif")
```
"""
function visualize_wavefield(val::Array{Float64, 3},
    param::Union{ElasticPropagatorParams,AcousticPropagatorParams, MPIAcousticPropagatorParams, MPIElasticPropagatorParams})

    figure()
    vmin = -3std(val)
    vmax = 3std(val)
    dt = max(1, div(param.NSTEP, 20))

    pl = matshow(val[1,:,:]', cmap="jet", vmin=vmin, vmax=vmax)
    t = title("Time = $(round(0.0, digits=2))")
    colorbar(shrink=0.8)
    gca().xaxis.set_ticks_position("bottom")
    xlabel("X")
    ylabel("Y")

    function update(i)
        pl[:set_data](val[i,:,:]')
        t.set_text("Time = $(round((i-1)*param.DELTAT, digits=3))")
    end
    p = animate(update, [2:dt:size(val, 1);])

end

function plot_result(sess, var, feed_dict, iter; figure_dir::String="figures", result_dir::String="results", var_name=nothing)

    if !isdir(figure_dir)
        figure_dir="./"
    end
    if !isdir(result_dir)
        result_dir=figure_dir
    end

    x = run(sess, var, feed_dict=feed_dict)

    std_images = zeros(size(x,2), size(x,3)) * 0
    mean_images = zeros(size(x,2), size(x,3)) * 0
    for i = 1:size(x,2)
        for j = 1:size(x, 3)
            std_images[i,j] = std(x[:, i, j, 1])
            mean_images[i,j] = mean(x[:, i, j, 1])
        end
    end

    figure()
    imshow(std_images')
    colorbar(shrink=0.5)
    title("Iteration = $iter")
    if isnothing(var_name)
        savefig(joinpath(figure_dir, "std_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
        writedlm(joinpath(result_dir, "std_$(lpad(iter,5,"0")).txt"), std_images')
    else
        savefig(joinpath(figure_dir, "std_$(var_name)_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
        writedlm(joinpath(result_dir, "std_$(var_name)_$(lpad(iter,5,"0")).txt"), std_images')
    end

    figure()
    fig, ax = subplots(div(size(x)[1],2), 2)
    for i = 1:div(size(x)[1],2)
        for j = 1:2
            ax[i, j].get_xaxis().set_visible(false)
            ax[i, j].get_yaxis().set_visible(false)
            k = (i-1) * 2 + j
            pcm= ax[i, j].imshow(x[k,:,:]')
            # fig.colorbar(pcm, ax=ax[i,j])
        end
    end
    subplots_adjust(wspace=0, hspace=0)
    suptitle("Iteration = $iter")
    if isnothing(var_name)
        savefig(joinpath(figure_dir, "inv_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
        writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), mean_images')
    else
        savefig(joinpath(figure_dir, "inv_$(var_name)_$(lpad(iter,5,"0")).png"), bbox_inches="tight")
        writedlm(joinpath(result_dir, "inv_$(var_name)_$(lpad(iter,5,"0")).txt"), mean_images')
    end
    # writedlm(joinpath(result_dir, "inv_$(lpad(iter,5,"0")).txt"), vp*scale)
    close("all")
end

function add_initial_model(x, v)
    if ndims(x) == 3
        size_x = size(x)[2:3]
    else
        size_x = size(x)
    end
    if size_x <= size(v)
        i = (size(v)[1]-size_x[1])÷2 +1 :(size(v)[1]-size_x[1])÷2 + size_x[1]
        j = (size(v)[2]-size_x[2])÷2 +1 : (size(v)[2]-size_x[2])÷2 + size_x[2]
        if ndims(x) == 3
            v = stack([scatter_add(v, i, j, x[k]) for k in 1:size(x)[1]])
        else
            v = scatter_add(v, i, j, x)
        end
        @warn "size(x) <= size(vp)", size_x, size(v)
    elseif  size_x > size(v)
        if ndims(x) == 3
            v = v + tf.slice(x, (size(x).-(size(x)[1], size(v)...)).÷2, (size(x)[1], size(v)...))
        else
            v = v + tf.slice(x, (size(x).-size(v)).÷2, size(v))
        end
        @warn "size(x) > size(vp)", size_x, size(v)
    else
        error("Size error: ", size(v), size_x)
    end
    return v
end

"""
    Ricker(epp::Union{ElasticPropagatorParams, AcousticPropagatorParams},
    a::Union{PyObject, <:Real},
    shift::Union{PyObject, <:Real},
    amp::Union{PyObject, <:Real}=1.0)

Returns a Ricker wavelet (a tensor).
- `epp`: a `ElasticPropagatorParams` or an `AcousticPropagatorParams`
- `a`: Width parameter
- `shift`: Center of the Ricker wavelet
- `amp`: Amplitude of the Ricker wavelet

```math
f(x) = \\mathrm{amp}A (1 - x^2/a^2) exp(-x^2/2 a^2)
```
where
```math
A = 2/sqrt(3a)pi^1/4
```
"""
function Ricker(epp::Union{ElasticPropagatorParams, AcousticPropagatorParams, MPIAcousticPropagatorParams, MPIElasticPropagatorParams},
        a::Union{PyObject, <:Real},
        shift::Union{PyObject, <:Real},
        amp::Union{PyObject, <:Real}=1.0)
    a = convert_to_tensor(a)
    shift = convert_to_tensor(shift)
    amp = convert_to_tensor(amp)
    NT, T = epp.NSTEP, epp.NSTEP*epp.DELTAT
    A = 2 / (sqrt(3 * a) * (pi^0.25))
    wsq = a^2
    vec = collect(1:NT)-shift
    xsq = vec^2
    mod = (1 - xsq / wsq)
    gauss = exp(-xsq / (2 * wsq))
    total = amp * A * mod * gauss
    return total
end

function Gauss(epp::Union{ElasticPropagatorParams, AcousticPropagatorParams, MPIAcousticPropagatorParams, MPIElasticPropagatorParams},
        a::Union{<:Real, PyObject},
        shift::Union{<:Real, PyObject, Missing} = missing,
        amp::Union{<:Real, PyObject} = 1.0)
    if ismissing(shift)
        shift = 1.2/a
    end
    amp = convert_to_tensor(amp)
    a = convert_to_tensor(a)
    shift = convert_to_tensor(shift, dtype=Float64)
    t = collect(0:epp.NSTEP-1)*epp.DELTAT
    A = π^2*a^2
    # -a*2.0*a*(t - t0) .* exp(-a*(t - t0)^2)
    amp * 2.0 * A * exp(-A*(t - shift)^2)
end

function compute_lame_parameters(NX::Int, NY::Int, vp::Float64, vs::Float64, rho::Float64)
    λ = zeros(NX+2, NY+2)
    μ = zeros(NX+2, NY+2)
    ρ = zeros(NX+2, NY+2)
    for j = 1:NY+2
        for i = 1:NX+2
            ρ[i,j] = rho
            μ[i,j] = rho*vs*vs
            λ[i,j] = rho*(vp*vp - 2. * vs*vs)
        end
    end
    λ, μ, ρ
end

function compute_lame_parameters(vp::Union{PyObject, Array{Float64,2}}, vs::Union{PyObject, Array{Float64,2}},
                                 rho::Union{PyObject, Array{Float64,2}})
    ρ = rho
    μ = rho .* vs .* vs
    λ = rho .* (vp .* vp - 2. * vs .* vs)
    return λ, μ, ρ
end

function compute_lame_parameters(param::MPIElasticPropagatorParams, vp::Float64, vs::Float64, rho::Float64)
    compute_lame_parameters(param.n+2, param.n+2, vp, vs, rho)
end

function compute_properties(vp::Union{PyObject, Array{Float64}},
                            vs::Union{PyObject, Array{Float64}},
                            ρ::Union{PyObject, Array{Float64}})
    vs = convert_to_tensor(vs)
    vp = convert_to_tensor(vp)
    ρ = convert_to_tensor(ρ)
    μ = ρ .* vs .* vs
    λ = ρ .* vp .* vp - 2 * μ
    λ, μ, ρ
end

function gradtest(sess::PyObject, o::PyObject, g::PyObject, loss::PyObject)
    sess = Session(); init(sess)
    X, G, L = run(sess, [o, g, loss])
    v_ = rand(size(X)...)*median(X)

    gs_ =  @. 1 / 10^(1:5)
    s = zeros(5)
    w = zeros(5)
    for i = 1:5
        @info i
        g_ = gs_[i]
        run(sess, assign(o, X + g_*v_))
        Li = run(sess, loss)
        s[i] = Li - L
        w[i] = Li - L - g_ * sum(v_ .* G)
    end

    close("all")
    loglog(gs_, abs.(s), "*-", label="finite difference")
    loglog(gs_, abs.(w), "+-", label="automatic differentiation")
    loglog(gs_, gs_.^2 * 0.5*abs(w[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
    loglog(gs_, gs_ * 0.5*abs(s[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

    plt.gca().invert_xaxis()
    legend()
    xlabel("\$\\gamma\$")
    ylabel("Error")
end

"""
    compute_loss_and_grads_GPU(model::Function, src::Union{Array{AcousticSource},Array{ElasticSource}},
    rcv_sim::Union{Array{AcousticReceiver}, Array{ElasticReceiver}}, Rs::Array{Array{Float64}}, vs::Union{PyObject, Array{PyObject}})

Computes the loss and the gradients of the model on all available GPUs.
- `model`, `src`, `rcv_sim`: Model, Source and Receiver returned by [`load_acoustic_source`](@ref), [`load_acoustic_model`](@ref),
    [`load_acoustic_receiver`](@ref) or [`load_elastic_source`](@ref), [`load_elastic_model`](@ref),
    [`load_elastic_receiver`](@ref)
- `Rs`: True receiver data
- `vs`: Variables to be optimized. It or its elements must have `ref` type.
"""
function compute_loss_and_grads_GPU(model::Function, src::Union{Array{AcousticSource},Array{ElasticSource}},
         rcv_sim::Union{Array{AcousticReceiver}, Array{ElasticReceiver}}, Rs::Array{Array{Float64,2},1}, vs::Union{PyObject, Array{PyObject}}; aux_loss::Union{PyObject, Missing}=missing)

    flag = isa(vs, PyObject)
    function run_on_gpu_device(gpu_id, jobs)
        local loss, g
        @pywith tf.device("/gpu:$(gpu_id)") begin
            [SimulatedObservation!(model(src[i]), rcv_sim[i]) for i = jobs]
            loss = sum([sum((rcv_sim[i].rcvv-Rs[i])^2) for i = jobs]) #+reg * tf.cast(sum(sum.(vs.^2)), tf.float64)
            if !ismissing(aux_loss)
                loss += aux_loss
            end
            g = gradients(loss, vs)
        end
        return loss, g
    end

    n = length(Rs)
    n_gpu = min(length(use_gpu()), n)
    if n_gpu==0
        error("No GPU resources available")
    end

    losses = Array{PyObject}(undef, n_gpu)
    gs = Array{flag ? PyObject : Array{PyObject}}(undef, n_gpu)
    for i = 1:n_gpu
        jobs = [k for k=1:n if k%n_gpu==i-1]
        println("GPU $i --> sources $jobs")
        losses[i], gs[i] = run_on_gpu_device(i-1, jobs)
    end
    gs = [gradients_GPU(x) for x in losses]
    return losses, gs
end

# function sampling_compute_loss_and_grads_GPU(models, src::AcousticSource,
#     rcv_sim::Array{AcousticReceiver},
#     Rs::Array{PyObject}, reg::Float64=1.0)
#     local loss, d
#     @show length(models), length(rcv_sim)
#     @assert length(models)==length(rcv_sim)
#     n_models = length(models)
#     for model in models
#         if model.param.IT_DISPLAY!=0
#             @warn "To use GPU, IT_DISPLAY must be 0. Setting IT_DISPLAY to 0..."
#             model.param.IT_DISPLAY = 0
#         end
#     end

#     function run_on_gpu_device!(gpu_id, jobs)
#         local d
#         @pywith tf.device("/gpu:$(gpu_id)") begin
#             [SimulatedObservation!(models[i](src), rcv_sim[i]) for i = jobs]
#             x = vcat([reshape(rcv_sim[i].rcvv,1,-1) for i in jobs]...)
#             y = vcat([reshape(r, 1, -1) for r in Rs]...)
#             d = dist(x,y,1)
#             @info size(x), size(y), size(d)
#         end
#         return d
#     end

#     n_gpu = min(length(use_gpu()), n_models)

#     if n_gpu==0
#         @warn "not using GPU"
#         [SimulatedObservation!(models[i](src), rcv_sim[i]) for i = 1:n_models]
#     else
#         ds = Array{PyObject}(undef, n_gpu)
#         for i = 1:n_gpu
#             jobs = [k for k=1:n_models if k%n_gpu==i-1]
#             println("GPU $i --> sources $jobs")
#             ds[i] = run_on_gpu_device!(i-1, jobs)
#         end
#     end

#     @pywith tf.device("/cpu:0") begin
#         # x = vcat([reshape(rcv_sim[i].rcvv,1,-1) for i = 1:n_models]...)
#         # y = vcat([reshape(r, 1, -1) for r in Rs]...)
#         # @info size(x), size(y)
#         # d = dist(x,y,1)
#         d = vcat(ds...)
#         a = tf.ones((size(d,1),), dtype=tf.float64)/Float64(size(d,1))
#         b = tf.ones((size(d,2),), dtype=tf.float64)/Float64(size(d,2))
#         loss = sinkhorn(a, b, d, reg=reg)
#     end
#     return loss, d
# end

function sampling_compute_loss_and_grads_GPU(models, src::AcousticSource, rcv_sim::Array{AcousticReceiver}, Rs::Array{PyObject}, reg::Float64=1.0)

    local loss, d
    @show length(models), length(rcv_sim)
    @assert length(models)==length(rcv_sim)
    n_models = length(models)

    function run_on_gpu_device!(gpu_id, jobs)
        local d
        @pywith tf.device("/gpu:$(gpu_id)") begin
            [SimulatedObservation!(models[i](src), rcv_sim[i]) for i = jobs]
            x = vcat([reshape(rcv_sim[i].rcvv,1,-1) for i in jobs]...)
            y = vcat([reshape(r, 1, -1) for r in Rs]...)
            d = dist(x,y,1)
            # @info size(x), size(y), size(d)
        end
        return d
    end

    n_gpu = min(length(use_gpu()), n_models)

    if n_gpu==0
        @warn "not using GPU"
        [SimulatedObservation!(models[i](src), rcv_sim[i]) for i = 1:n_models]
    else
        ds = Array{PyObject}(undef, n_gpu)
        for i = 1:n_gpu
            jobs = [k for k=1:n_models if k%n_gpu==i-1]
            println("GPU $i --> sources $jobs")
            ds[i] = run_on_gpu_device!(i-1, jobs)
        end
    end

    @pywith tf.device("/cpu:0") begin
        d = vcat(ds...)
        a = tf.ones((size(d,1),), dtype=tf.float64)/Float64(size(d,1))
        b = tf.ones((size(d,2),), dtype=tf.float64)/Float64(size(d,2))
        loss = sinkhorn(a, b, d, reg=reg, method="lp")
    end
    return loss, d
end

"""
- `models` : nsrc * batch_size
- `src` : nsrc
- `rcv_sim` : nsrc * batch_size
- `Rs` : nsrc * sample_size
"""
function sampling_compute_loss_and_grads_GPU_v2(models, src, rcv_sim,  Rs; reg::Float64=1.0, method::String="sinkhorn")

    local loss
    nsrc, batch_size = size(models)
    @assert nsrc == length(src)
    @assert size(rcv_sim,1) == nsrc
    @assert size(rcv_sim,2) == batch_size
    @assert size(Rs,1)==nsrc
    sample_size = size(Rs,2)

    function run_on_gpu_device!(gpu_id, src_ids)
        local ds
        @pywith tf.device("/gpu:$(gpu_id)") begin
            x = Any[[] for i=1:length(src_ids)]
            y = Any[[] for i=1:length(src_ids)]
            for (k,i) in enumerate(src_ids)
                for j = 1:batch_size
                    SimulatedObservation!(models[i, j](src[i]), rcv_sim[i, j])
                    push!(x[k], reshape(rcv_sim[i, j].rcvv, 1, -1))
                end
                x[k] = vcat(x[k]...)
            end
            for (k,i) in enumerate(src_ids)
                for j = 1:sample_size
                    push!(y[k], reshape(Rs[i, j], 1, -1))
                end
                y[k] = vcat(y[k]...)
            end

            ds = Array{PyObject}(undef, length(src_ids))
            for i = 1:length(src_ids)
                ds[i] = dist(x[i], y[i], 1)
                # @info ds[i]
            end
        end
        return ds
    end

    n_gpu = min(length(use_gpu()), nsrc)

    if n_gpu==0
        @error "You must use GPU for sampling_compute_loss_and_grads_GPU_v2"
    else
        ds = PyObject[]
        for i = 1:n_gpu
            jobs = [k for k=1:nsrc if k%n_gpu==i-1]
            println("GPU $i --> sources $jobs with $(batch_size) vp each")
            [push!(ds, d) for d in run_on_gpu_device!(i-1, jobs)]
        end
        @assert length(ds)==nsrc
    end

    @pywith tf.device("/cpu:0") begin
        a = tf.ones((size(ds[1],1),), dtype=tf.float64)/Float64(size(ds[1],1))
        b = tf.ones((size(ds[1],2),), dtype=tf.float64)/Float64(size(ds[1],2))
        loss = constant(0.0)
        for i = 1:nsrc
            loss += sinkhorn(a, b, ds[i], reg=reg, method=method)
        end
    end
    return loss

end

"""

Mathcing statistics.
- `models` : nsrc * batch_size
- `src` : nsrc
- `rcv_sim` : nsrc * batch_size
- `Rs` : nsrc * sample_size
"""
function sampling_compute_loss_and_grads_GPU_v3(models, src, rcv_sim, Rs; reg::Float64=1.0, method::String="sinkhorn")

    local loss
    nsrc, batch_size = size(models)
    @assert nsrc == length(src)
    @assert size(rcv_sim,1) == nsrc
    @assert size(rcv_sim,2) == batch_size
    @assert size(Rs,1)==nsrc
    sample_size = size(Rs,2)

    function run_on_gpu_device!(gpu_id, src_ids)
        local ds
        @pywith tf.device("/gpu:$(gpu_id)") begin
            x = Any[[] for i=1:length(src_ids)]
            y = Any[[] for i=1:length(src_ids)]
            for (k,i) in enumerate(src_ids)
                for j = 1:batch_size
                    SimulatedObservation!(models[i, j](src[i]), rcv_sim[i, j])
                    push!(x[k], reshape(rcv_sim[i, j].rcvv, 1, -1))
                end
                x[k] = vcat(x[k]...)
            end
            for (k,i) in enumerate(src_ids)
                for j = 1:sample_size
                    push!(y[k], reshape(Rs[i, j], 1, -1))
                end
                y[k] = vcat(y[k]...)
            end

            ds = Array{PyObject}(undef, length(src_ids))
            for i = 1:length(src_ids)
                ds[i] = sum((mean(x[i], dims=1) - mean(y[i], dims=1))^2 +
                                 (std(x[i], dims=1) - std(y[i], dims=1))^2)
                @info ds[i]
            end
        end
        return ds
    end

    n_gpu = min(length(use_gpu()), nsrc)

    if n_gpu==0
        @error "You must use GPU for sampling_compute_loss_and_grads_GPU_v2"
    else
        ds = PyObject[]
        for i = 1:n_gpu
            jobs = [k for k=1:nsrc if k%n_gpu==i-1]
            println("GPU $i --> sources $jobs with $(batch_size) vp each")
            [push!(ds, d) for d in run_on_gpu_device!(i-1, jobs)]
        end
        @assert length(ds)==nsrc
    end

    @pywith tf.device("/cpu:0") begin
        # a = tf.ones((size(ds[1],1),), dtype=tf.float64)/Float64(size(ds[1],1))
        # b = tf.ones((size(ds[1],2),), dtype=tf.float64)/Float64(size(ds[1],2))
        loss = constant(0.0)
        for i = 1:nsrc
            loss += ds[i]
        end
    end
    return loss
end

function compute_forward_GPU(model::Function, src::Union{Array{AcousticSource},Array{ElasticSource}},
    rcv::Union{Array{AcousticReceiver}, Array{ElasticReceiver}})

    function run_on_gpu_device(gpu_id, jobs)
        local rcvv
        @pywith tf.device("/gpu:$(gpu_id)") begin
            [SimulatedObservation!(model(src[i]), rcv[i]) for i = jobs]
            rcvv = [rcv[i].rcvv for i = jobs]
        end
        return rcvv
    end

    n = length(src)
    n_gpu = min(length(use_gpu()), n)
    if n_gpu==0
        error("No GPU resources available")
    end

    Rs = Array{PyObject}(undef, n)
    for i = 1:n_gpu
        jobs = [k for k=1:n if k%n_gpu==i-1]
        println("GPU $i --> sources $jobs")
        Rs[jobs] = run_on_gpu_device(i-1, jobs)
    end

    return Rs
end

function variable_source(param::Union{ElasticPropagatorParams,AcousticPropagatorParams},
        x::Union{PyObject, Float64}, y::Union{PyObject, Float64}, v::Union{Array{Float64},PyObject};
        sigma::Union{Float64, Nothing} = nothing)
    x = convert_to_tensor(x)
    y = convert_to_tensor(y)
    v = convert_to_tensor(v)
    if isnothing(sigma)
        sigma = 1.0
    end
    srci = zeros(Int64, (param.NX+2)*(param.NY+2))
    srcj = zeros(Int64, (param.NX+2)*(param.NY+2))
    xs = zeros((param.NX+2)*(param.NY+2))
    ys = zeros((param.NX+2)*(param.NY+2))
    k = 1
    for i = 1:param.NX+2
        for j = 1:param.NY+2
            srci[k] = i
            srcj[k] = j
            xs[k] = (i-1)*param.DELTAX
            ys[k] = (j-1)*param.DELTAY
            k += 1
        end
    end
    # magn = 1/(2π*sigma) * exp( - ((xs-x)^2 + (ys-y)^2)/2sigma^2)
    h = param.DELTAX
    # srcv = reshape(v, length(v), 1) * exp(-((xs-x)^2+(ys-y)^2)/(h^2))'
    magn = 1/(2π*sigma) * exp( - (((xs-x)/(sqrt(2sigma)*param.DELTAX))^2 +
                    ((ys-y)/(sqrt(2sigma)*param.DELTAY))^2))
    srcv = reshape(v, length(v), 1) * magn'
    return srci, srcj, srcv
end

function variable_source(param::AcousticPropagatorParams,
    ix::Array{Int64,1}, iy::Array{Int64,1}; name::String = "src", source_scale::Float64 = 1.0)
    @assert length(ix)==length(iy)
    v = Variable(zeros(param.NSTEP+1, length(ix)), name=name) * source_scale
    AcousticSource(ix, iy, v)
end


text:
# This file uses metaprogramming techniques to implement vector-scalar and scalar-vector
# operations (eg: scaling)

vsops = [:*, :+, :-, :/, :./, :.*]
vsopnames_l = [:ET_vsltimes, :ET_vslplus, :ET_lminus, :ET_ldiv, :ET_lddiv, :ET_ldtimes]
vsopnames_r = [:ET_vsrtimes, :ET_vsrplus, :ET_rminus, :ET_rdiv, :ET_rddiv, :ET_rdtimes]

# generate rhs ops
for i = 1:length(vsops)
    op = vsops[i]
    opname = vsopnames_l[i]

    eval(quote
         immutable $opname{T,R,ET1<:VectorizedExpression, ET2<:Number} <: VectorizedExpression{T,R}
         lhs::ET1
         rhs::ET2
         end

         @inline function getindex(A::$opname, i...)
         return $op(A.lhs[i...], A.rhs)
         end

         @inline function length(A::$opname)
         return length(A.lhs)
         end

         @inline function size(A::$opname)
             return size(A.lhs)
         end

        @inline function ($op){T,R}(lhs::VectorizedExpression{T,R}, rhs::Number)
            return $opname{typeof($op(T(0),rhs)),R,typeof(lhs), typeof(rhs)}(lhs, rhs)
        end

        export $op, $opname
    end)
end

#generate lhs ops
for i = 1:length(vsops)
    op = vsops[i]
    opname = vsopnames_r[i]

    eval(quote
         immutable $opname{T,R,ET1<:Number, ET2<:VectorizedExpression} <: VectorizedExpression{T,R}
         lhs::ET1
         rhs::ET2
         end

         @inline function getindex(A::$opname, i...)
         return $op(A.lhs, A.rhs[i...])
         end

         @inline function length(A::$opname)
         return length(A.rhs)
         end

         @inline function size(A::$opname)
             return size(A.rhs)
         end

        @inline function ($op){T,R}(lhs::Number, rhs::VectorizedExpression{T,R})
            return $opname{typeof($op(T(0),lhs)) ,R,typeof(lhs), typeof(rhs)}(lhs, rhs)
        end

        export $op, $opname
    end
    )
end


text:
"""

Read summary output file created by stansummary.

$(SIGNATURES)

# Extended help

### Required arguments
```julia
* `m`                                  : A Stan model object, e.g. SampleModel
```

### Optional arguments
```julia
* `printsummary=false`                 : Print cmdstan summary
```

### Returns
```julia
* `df`                                 : Dataframe containing the cmdstan summary
```

"""
function read_summary(m::T, printsummary=false) where {T <: CmdStanModels}

  fname = "$(m.output_base)_summary.csv"
  !isfile(fname) && stan_summary(m, printsummary)

  df = CSV.read(fname, DataFrame; delim=",", comment="#")

  cnames = lowercase.(convert.(String, String.(names(df))))
  cnames[1] = "parameters"
  cnames[4] = "std"
  cnames[8] = "ess"
  rename!(df, Symbol.(cnames), makeunique=true)
  df[!, :parameters] = Symbol.(df[!, :parameters])

  df

end   # end of read_samples


text:
if basename(pwd()) == "aoc"
    cd("2015/14")
end

function loadreindeer(filename::AbstractString)
    map(eachmatch(r"(.*) can fly (\d+) km/s for (\d+) seconds, but then must rest for (\d+) seconds\.", readchomp(filename))) do m
        m[1], (parse(Int, m[i]) for i in 2:4)...
    end
end

function position(reindeer, time::Integer)
    _, speed, on, off = reindeer
    fulls, final = divrem(time, on + off)
    (fulls * on + min(final, on)) * speed
end
position(time::Integer) = Base.Fix2(position, time)

part1(filename::AbstractString, n) = part1(loadreindeer(filename), n)
part1(reindeer, n = 2503) = maximum(position(n), reindeer)

part2(filename::AbstractString, n = 2503) = part2(loadreindeer(filename), n)

function part2(reindeer, n = 2503)
    scores = Dict(r[1] => 0 for r in reindeer)
    for i in 1:n
        top = maximum(position(i), reindeer)
        for r in reindeer
            if position(r, i) == top
                scores[r[1]] += 1
            end
        end
    end
    maximum(values(scores))
end

max_stars_repo_name: stellartux/advent-of-code

max_issues_repo_name: stellartux/advent-of-code

max_forks_repo_name: stellartux/advent-of-code


text:
using Kafka
import Kafka: writeobj, readobj
using Test

io = IOBuffer()

writeobj(io, "test")
writeobj(io, 42)
writeobj(io, UInt8[1, 2, 3])
writeobj(io, ["hello", "world"])

seek(io, 0)

@test readobj(io, String) == "test"
@test readobj(io, Int64) == 42
@test readobj(io, Vector{UInt8}) == UInt8[1, 2, 3]
@test readobj(io, Vector{String}) == ["hello", "world"]

include("integration.jl")

println("Ok.")


text:
function Sch_solve_Emin_cg( LF::LF3dGrid, Vpot, Ncol::Int64;
                            v0 = nothing,
                            α_t=3e-5, Niter=1000,
                            verbose=false )
  #
  Npoints = LF.Nx * LF.Ny * LF.Nz
  ΔV = LF.LFx.h * LF.LFy.h * LF.LFz.h
  #
  # Setup trial solution
  #
  if v0 == nothing
    # random trial solution
    srand(1234)
    v = rand( Npoints, Ncol )
    v = orthonormalize( LF, v )
  else
    v = copy(v0)
    # v0 is assumed to be ortonormalized
  end
  #
  Energies = calc_Energies( LF, Vpot, v )
  #
  Etot_old = Energies.Total
  #
  g     = zeros(Float64,Npoints,Ncol)
  g_old = zeros(Float64,Npoints,Ncol)
  d     = zeros(Float64,Npoints,Ncol)
  d_old = zeros(Float64,Npoints,Ncol)
  #
  β = 0
  α = 0
  for iter = 1:Niter
    g = calc_grad( LF, Vpot, v )
    #
    if iter != 1
      β = trace( g' * g)/trace( g_old'*g_old )
      #β = trace( (g-g_old)'*g ) / trace( g_old'*g_old )
      #β = trace( (g-g_old)'*g )/ trace( (g-g_old)'*d )
    end
    d = -g + β*d_old
    #
    # compute gradient at trial step
    #
    v2 = orthonormalize(LF, v + α_t*d)
    g_t = calc_grad( LF, Vpot, v2 )
    #
    # compute estimate of best step and update current trial vectors
    #
    denum = trace( (g - g_t)' * d )
    if denum != 0.0
      α = abs( α_t * trace( g'*d )/denum )
    else
      α = 0.0
    end
    v = v + α*d
    #
    v = orthonormalize( LF, v )
    Energies = calc_Energies( LF, Vpot, v )
    Etot = Energies.Total
    #
    if verbose
      @printf("%8d %16.9e %16.9e\n", iter, Etot, abs(Etot-Etot_old))
    end
    if abs(Etot-Etot_old) < 1.e-7
      if verbose
        @printf("Emin CG converges in %8d iterations\n", iter)
      end
      break
    end
    Etot_old = Etot
    #
    g_old = g[:,:]
    d_old = d[:,:]
  end
  return Energies, v
end


text:
using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

function istractionboundary(x)
    if x[1] ≈ 1.0
        return true
    else
        return false
    end
end

function tractionfunc(x)
    return [x[1], 0.0]
end

lambda = 1.0
mu = 2.0
stiffness = CutCell.plane_strain_voigt_hooke_matrix(lambda, mu)
polyorder = 1
numqp = 2
basis = TensorProductBasis(2, polyorder)
nf = CutCell.number_of_basis_functions(basis)
vquad = tensor_product_quadrature(2, numqp)
facequads = CutCell.face_quadratures(numqp)

mesh = CutCell.Mesh([0.0, 0.0], [1.0, 1.0], [1, 1], nf)
cellmaps = CutCell.cell_maps(mesh)
nodalconnectivity = CutCell.nodal_connectivity(mesh)
cellconnectivity = CutCell.cell_connectivity(mesh)
cellmatrix = CutCell.bilinear_form(basis, vquad, stiffness, cellmaps[1])

sysmatrix = CutCell.SystemMatrix()
sysrhs = CutCell.SystemRHS()

CutCell.assemble_bilinear_form!(sysmatrix, cellmatrix, nodalconnectivity, 2)
CutCell.assemble_traction_force_linear_form!(
    sysrhs,
    tractionfunc,
    basis,
    facequads,
    cellmaps,
    nodalconnectivity,
    cellconnectivity,
    istractionboundary,
)

matrix = CutCell.make_sparse(sysmatrix,mesh)
rhs = CutCell.rhs(sysrhs,mesh)

bottomnodeids = CutCell.bottom_boundary_node_ids(mesh)
leftnodeids = CutCell.left_boundary_node_ids(mesh)

CutCell.apply_dirichlet_bc!(matrix,rhs,bottomnodeids,2,0.0,2)
CutCell.apply_dirichlet_bc!(matrix,rhs,leftnodeids,1,0.0,2)

sol = matrix\rhs
disp = reshape(sol,2,:)

m = (lambda+2mu)
e11 = 1/(m - lambda^2/m)
e22 = - lambda/m*e11
u1 = e11
u2 = e22

testdisp = [0. 0. u1  u1
            0. u2 0.  u2]
@test allapprox(disp,testdisp)

#########################################################
# Test 4 element quadratic basis uniform tension
lambda = 1.0
mu = 2.0
stiffness = CutCell.plane_strain_voigt_hooke_matrix(lambda, mu)
polyorder = 2
numqp = 3
basis = TensorProductBasis(2, polyorder)
nf = CutCell.number_of_basis_functions(basis)
vquad = tensor_product_quadrature(2, numqp)
facequads = CutCell.face_quadratures(numqp)

mesh = CutCell.Mesh([0.0, 0.0], [1.0, 1.0], [2, 2], nf)
cellmaps = CutCell.cell_maps(mesh)
nodalconnectivity = CutCell.nodal_connectivity(mesh)
cellconnectivity = CutCell.cell_connectivity(mesh)
cellmatrix = CutCell.bilinear_form(basis, vquad, stiffness, cellmaps[1])

sysmatrix = CutCell.SystemMatrix()
sysrhs = CutCell.SystemRHS()

CutCell.assemble_bilinear_form!(sysmatrix, cellmatrix, nodalconnectivity, 2)
CutCell.assemble_traction_force_linear_form!(
    sysrhs,
    tractionfunc,
    basis,
    facequads,
    cellmaps,
    nodalconnectivity,
    cellconnectivity,
    istractionboundary,
)

matrix = CutCell.make_sparse(sysmatrix,mesh)
rhs = CutCell.rhs(sysrhs,mesh)

bottomnodeids = CutCell.bottom_boundary_node_ids(mesh)
leftnodeids = CutCell.left_boundary_node_ids(mesh)

CutCell.apply_dirichlet_bc!(matrix,rhs,bottomnodeids,2,0.0,2)
CutCell.apply_dirichlet_bc!(matrix,rhs,leftnodeids,1,0.0,2)

sol = matrix\rhs
disp = reshape(sol,2,:)

nodalcoordinates = CutCell.nodal_coordinates(mesh)

function displacement_of_mesh(nodalcoordinates,u1,u2)
    disp = copy(nodalcoordinates)
    disp[1,:] .*= u1
    disp[2,:] .*= u2
    return disp
end

testdisp = displacement_of_mesh(nodalcoordinates,u1,u2)

@test allapprox(testdisp,disp)

#########################################################
# Test stress calculation
lambda = 1.0
mu = 2.0
stiffness = CutCell.plane_strain_voigt_hooke_matrix(lambda, mu)
polyorder = 1
numqp = 2
basis = TensorProductBasis(2, polyorder)
nf = CutCell.number_of_basis_functions(basis)
vquad = tensor_product_quadrature(2, numqp)
facequads = CutCell.face_quadratures(numqp)

mesh = CutCell.Mesh([0.0, 0.0], [1.0, 1.0], [1, 1], nf)
cellmaps = CutCell.cell_maps(mesh)
nodalconnectivity = CutCell.nodal_connectivity(mesh)
cellconnectivity = CutCell.cell_connectivity(mesh)
cellmatrix = CutCell.bilinear_form(basis, vquad, stiffness, cellmaps[1])

sysmatrix = CutCell.SystemMatrix()
sysrhs = CutCell.SystemRHS()

CutCell.assemble_bilinear_form!(sysmatrix, cellmatrix, nodalconnectivity, 2)
CutCell.assemble_traction_force_linear_form!(
    sysrhs,
    tractionfunc,
    basis,
    facequads,
    cellmaps,
    nodalconnectivity,
    cellconnectivity,
    istractionboundary,
)

matrix = CutCell.make_sparse(sysmatrix,mesh)
rhs = CutCell.rhs(sysrhs,mesh)

bottomnodeids = CutCell.bottom_boundary_node_ids(mesh)
leftnodeids = CutCell.left_boundary_node_ids(mesh)

CutCell.apply_dirichlet_bc!(matrix,rhs,bottomnodeids,2,0.0,2)
CutCell.apply_dirichlet_bc!(matrix,rhs,leftnodeids,1,0.0,2)

sol = matrix\rhs

detjac = CutCell.determinant_jacobian(mesh)
jac = CutCell.jacobian(mesh)
stresscellmatrix = CutCell.stress_cell_mass_matrix(basis,vquad,detjac)

sysmatrix = CutCell.SystemMatrix()
sysrhs = CutCell.SystemRHS()

CutCell.assemble_bilinear_form!(sysmatrix,stresscellmatrix,nodalconnectivity,3)
CutCell.assemble_stress_linear_form!(sysrhs,basis,vquad,stiffness,sol,nodalconnectivity,jac)

matrix = CutCell.make_sparse_stress_operator(sysmatrix,mesh)
rhs = CutCell.stress_rhs(sysrhs,mesh)

stress = matrix\rhs
stress = reshape(stress,3,:)

teststress = zeros(3,4)
teststress[1,:] .= 1.0

@test allapprox(stress,teststress,1e2eps())


text:
abstract type SteadyStateDiffEqAlgorithm <: DiffEqBase.AbstractSteadyStateAlgorithm end

struct SSRootfind{F} <: SteadyStateDiffEqAlgorithm
  nlsolve::F
end
SSRootfind(;nlsolve=(f,u0,abstol) -> (res=NLsolve.nlsolve(f,u0,ftol = abstol);res.zero)) = SSRootfind(nlsolve)

struct DynamicSS{A,AT,RT} <: SteadyStateDiffEqAlgorithm
  alg::A
  abstol::AT
  reltol::RT
end
DynamicSS(alg;abstol = 1e-8, reltol = 1e-6) = DynamicSS(alg,abstol,reltol)


text:
type LmResp{V<:FPVector} <: ModResp  # response in a linear model
    mu::V                                  # mean response
    offset::V                              # offset added to linear predictor (may have length 0)
    wts::V                                 # prior weights (may have length 0)
    y::V                                   # response
    function LmResp(mu::V, off::V, wts::V, y::V)
        n = length(y); length(mu) == n || error("mismatched lengths of mu and y")
        ll = length(off); ll == 0 || ll == n || error("length of offset is $ll, must be $n or 0")
        ll = length(wts); ll == 0 || ll == n || error("length of wts is $ll, must be $n or 0")
        new(mu,off,wts,y)
    end
end
LmResp{V<:FPVector}(y::V) = LmResp{V}(fill!(similar(y), zero(eltype(V))), similar(y, 0), similar(y, 0), y)

function updatemu!{V<:FPVector}(r::LmResp{V}, linPr::V)
    n = length(linPr); length(r.y) == n || error("length(linPr) is $n, should be $(length(r.y))")
    length(r.offset) == 0 ? copy!(r.mu, linPr) : map!(Add(), r.mu, linPr, r.offset)
    deviance(r)
end
updatemu!{V<:FPVector}(r::LmResp{V}, linPr) = updatemu!(r, convert(V,vec(linPr)))

type WtResid <: Functor{3} end
evaluate{T<:FP}(::WtResid,wt::T,y::T,mu::T) = (y - mu)*sqrt(wt)
result_type{T<:FP}(::WtResid,wt::T,y::T,mu::T) = T

deviance(r::LmResp) = length(r.wts) == 0 ? sumsqdiff(r.y, r.mu) : wsumsqdiff(r.wts,r.y,r.mu)
residuals(r::LmResp)= length(r.wts) == 0 ? r.y - r.mu : map(WtResid(),r.wts,r.y,r.mu)

type LinearModel{T<:LinPred} <: LinPredModel
    rr::LmResp
    pp::T
end

cholfact(x::LinearModel) = cholfact(x.pp)

function StatsBase.fit{LinPredT<:LinPred}(::Type{LinearModel{LinPredT}}, X::Matrix, y::Vector)
    rr = LmResp(float(y)); pp = LinPredT(X)
    installbeta!(delbeta!(pp, rr.y)); updatemu!(rr, linpred(pp,0.))
    LinearModel(rr, pp)
end
StatsBase.fit(::Type{LinearModel}, X::Matrix, y::Vector) = StatsBase.fit(LinearModel{DensePredQR}, X, y)

lm(X, y) = fit(LinearModel, X, y)
lmc(X, y) = fit(LinearModel{DensePredChol}, X, y)

## scale(m) -> estimate, s, of the scale parameter
## scale(m,true) -> estimate, s^2, of the squared scale parameter
function scale(x::LinearModel, sqr::Bool=false)
    ssqr = deviance(x.rr)/df_residual(x)
    sqr ? ssqr : sqrt(ssqr)
end

function coeftable(mm::LinearModel)
    cc = coef(mm)
    se = stderr(mm)
    tt = cc ./ se
    CoefTable(hcat(cc,se,tt,ccdf(FDist(1, df_residual(mm)), abs2(tt))),
              ["Estimate","Std.Error","t value", "Pr(>|t|)"],
              ["x$i" for i = 1:size(mm.pp.X, 2)], 4)
end

predict(mm::LinearModel, newx::Matrix) =  newx * coef(mm)

function confint(obj::LinearModel, level::Real)
    hcat(coef(obj),coef(obj)) + stderr(obj) *
    quantile(TDist(df_residual(obj)), (1. - level)/2.) * [1. -1.]
end
confint(obj::LinearModel) = confint(obj, 0.95)


text:
@model function bernoulli_mixture(x)
    # Mixture prior.
    w ~ Dirichlet(2, 1.0)

    # Latent probability.
    p ~ DiscreteNonParametric([0.3, 0.7], w)

    # Observation.
    x ~ Bernoulli(p)
end

bernoulli_example(x = false) = bernoulli_mixture(x)

@model function gmm(x, K)
    N = length(x)

    # Cluster centers.
    μ ~ filldist(Normal(), K)

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments.
    z ~ filldist(DiscreteNonParametric(1:K, w), N)

    # Observations.
    for n = 1:N
        x[n] ~ Normal(μ[z[n]], 1.0)
    end
end

gmm_example(x = [0.1, -0.05, 1.0], K = 2) = gmm(x, K)

@model function gmm_loopy(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # Cluster centers.
    μ = Vector{T}(undef, K)
    for k = 1:K
        μ[k] ~ Normal()
    end

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments & observations.
    z = Vector{Int}(undef, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
        x[n] ~ Normal(μ[z[n]], 1.0)
    end
end

# same as gmm_loopy, but with an affine transformation on μ.
@model function gmm_shifted(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # Cluster centers.
    μ = Vector{T}(undef, K)
    for k = 1:K
        μ[k] ~ Normal()
    end

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments & observations.
    z = Vector{Int}(undef, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
        x[n] ~ Normal(4μ[z[n]] - 1, 1.0)
    end
end

# K clusters, each one around i for i = 1:K with variance 0.5
@model function hmm(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # State sequence.
    s = zeros(Int, N)

    # Emission matrix.
    m = Vector{T}(undef, K)

    # Transition matrix.
    T = Vector{Vector{T}}(undef, K)

    # Assign distributions to each element
    # of the transition matrix and the
    # emission matrix.
    for i = 1:K
        T[i] ~ Dirichlet(K, 1.0)
        m[i] ~ Normal(i, 0.5)
    end

    # Observe each point of the input.
    s[1] ~ Categorical(K)
    x[1] ~ Normal(m[s[1]], 0.1)

    for i = 2:N
        s[i] ~ Categorical(T[s[i-1]])
        x[i] ~ Normal(m[s[i]], 0.1)
    end
end

hmm_example(x = [1.1, 0.95, 2.2], K = 2) = hmm(x, K)

@model function changepoint(y)
    N = length(y)
    α = 1 / mean(y)
    λ₁ ~ Exponential(α)
    λ₂ ~ Exponential(α)
    τ ~ DiscreteUniform(1, N)

    for n in 1:N
        y[n] ~ Poisson(τ > n ? λ₁ : λ₂)
    end
end

# @model function reverse_deps(x)
#     m = Vector{Float64}(undef, 2)
#     m[1] ~ Normal()
#     m[2] ~ Normal()
#     x ~ MvNormal(m)
# end

###########################################################################
# data from R. Neal paper:
const data_neal = [-1.48, -1.40, -1.16, -1.08, -1.02, 0.14, 0.51, 0.53, 0.78]
const α_neal = 10.0

function stickbreak(v)
    K = length(v) + 1
    cumprod_one_minus_v = cumprod(1 .- v)

    return map(1:K) do k
        if k == 1
            v[1]
        elseif k == K
            cumprod_one_minus_v[K - 1]
        else
            v[k] * cumprod_one_minus_v[k - 1]
        end
    end
end

# from https://luiarthur.github.io/TuringBnpBenchmarks/dpsbgmm
@model function imm_stick(y, α, K)
    N = length(y)
    crm = DirichletProcess(α)

    v ~ filldist(StickBreakingProcess(crm), K - 1)
    w = stickbreak(v)

    # Cluster assignments
    z = zeros(Int, N)
    for n = 1:N
        z[n] ~ Categorical(w)
    end

    # Cluster centers
    L = identity(K)
    μ ~ filldist(Normal(), L)

    # Observations
    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

imm_stick_example(y = data_neal, α = α_neal, K = 10) = imm_stick(y, α, K)