text:
"""
    WebsocketClient([; options...])

Constructs a new WebsocketClient, overriding [`clientConfig`](@ref) with the passed options.

# Example
```julia
using Websocket
client = WebsocketClient([; options...])
```
"""
struct WebsocketClient
    config::NamedTuple
    callbacks::Dict{Symbol, Union{Bool, Function}}
    flags::Dict{Symbol, Bool}

    function WebsocketClient(; config...)
        @debug "WebsocketClient"
        config = merge(clientConfig, (; config...))
        config = merge(config, (; maskOutgoingPackets = true, type = "client") )
        self = new(
            config,
            Dict{Symbol, Union{Bool, Function}}(
                :connect => false,
                :connectError => false,
            ),
            Dict{Symbol, Bool}(
                :isopen => false
            )
        )
    end
end
"""
    listen(callback::Function, client::Websocket.WebsocketClient, event::Symbol)
Register event callbacks onto a client. The callback must be a function with exactly one argument.

Valid events are:
- :connect
- :connectError

# Example
```julia
listen(client, :connect) do ws
    #...
end
```
"""
function listen(
    cb::Function,
    self::WebsocketClient,
    key::Symbol
)
    if !haskey(self.callbacks, key)
        return @warn "WebsocketClient has no listener for :$key."
    end
    if haskey(self.callbacks, key) && !(self.callbacks[key] isa Function)
        self.callbacks[key] = data -> (
            @async try
                cb(data)
            catch err
                err = CallbackError(err, catch_backtrace())
                err.log()
                exit()
            end
        )
    end
end

function makeConnection(
    self::WebsocketClient,
    urlString::String,
    headers::Dict{String, String};
        options...
)
    @debug "WebsocketClient.connect"
    if isopen(self)
        @warn """called "connect" on a WebsocketClient that is open or opening."""
        return
    end
    options = merge((; options...), clientOptions)
    connected = Condition()
    self.flags[:isopen] = true
    @async try
        connection = wait(connected)
        if self.callbacks[:connect] isa Function
            self.callbacks[:connect](connection)
            wait(connection.closed)
            self.flags[:isopen] = false
        else
            throw(error("""called "open" before registering ":connect" event."""))
        end
    catch err
        err = ConnectError(err, catch_backtrace())
        self.flags[:isopen] = false
        if self.callbacks[:connectError] isa Function
            self.callbacks[:connectError](err)
        else
            err.log()
            exit()
        end
    end

    try
        headers = makeHeaders(headers)
        if !(headers["Sec-WebSocket-Version"] in ["8", "13"])
            throw(error("only version 8 and 13 of websocket protocol supported."))
        end
        if haskey(headers, "Sec-WebSocket-Extensions")
            throw(error("websocket extensions not supported in client"))
        end
        connect(
            self.config,
            urlString,
            connected,
            headers;
                options...
        )
    catch err
        @async notify(connected, err; error = true)
    end
end

function validateHandshake(headers::Dict{String, String}, request::HTTP.Messages.Response)

    if request.status != 101
        throw(error("connection error with status: $(request.status)"))
    end
    if !HTTP.hasheader(request, "Connection", "Upgrade")
        throw(error("""did not receive "Connection: Upgrade" """))
    end
    if !HTTP.hasheader(request, "Upgrade", "websocket")
        throw(error("""did not receive "Upgrade: websocket" """))
    end
    if !HTTP.hasheader(request, "Sec-WebSocket-Accept", acceptHash(headers["Sec-WebSocket-Key"]))
        throw(error("""invalid "Sec-WebSocket-Accept" response from server"""))
    end
    if HTTP.hasheader(request, "Sec-WebSocket-Extensions")
        @warn "Server uses websocket extensions" (;
            value = HTTP.header(request, "Sec-WebSocket-Extensions"),
            caution = "Websocket extensions are not supported in the client and may cause connection closure."
        )...
    end
end

function connect(
    config::NamedTuple,
    url::String,
    connected::Condition,
    headers::Dict{String, String};
        options...
)
    @debug "WebsocketClient.connect"
    let self
        HTTP.open("GET", url, headers;
            options...
        ) do io
            tcp = io.stream.c.io isa TCPSocket ? io.stream.c.io : io.stream.c.io.bio
            VERSION >= v"1.3" && Sockets.nagle(tcp, config.useNagleAlgorithm) #Sockets.nagle needs Julia >= 1.3
            try
                request = startread(io)
                validateHandshake(headers, request)
                self = WebsocketConnection(io.stream, config)
                notify(connected, self)
            catch err
                notify(connected, err; error = true)
                return
            end
            startConnection(self, io)
        end
    end
end
"""
    open(client::WebsocketClient, url::String [, headers::Dict{String, String}; options...])
Open a new websocket client connection at the given `url`. Blocks until the connection is closed.

Optionally provide custom `headers` for the http request.

`; options...` are passed to the underlying [HTTP.request](https://juliaweb.github.io/HTTP.jl/stable/public_interface/#Requests-1)
"""
function Base.open(client::WebsocketClient, url::String, headers::Dict{String, String} = Dict{String, String}(); options...)
    makeConnection(client, url, headers; options...)
end
"""
    isopen(client::Websocket.WebsocketClient)::Bool
Returns a Bool indication if the client TCP connection is open.
"""
function Base.isopen(client::WebsocketClient)
    client.flags[:isopen]
end

max_stars_repo_name: citkane/Websocket

max_issues_repo_name: citkane/Websocket

max_forks_repo_name: citkane/Websocket


text:
macro _threads(ex)
    return quote
        if (Threads.nthreads() > 1) && (length(st) > 4096)
            $(Expr(:macrocall, Expr(:(.), :Threads, QuoteNode(Symbol("@threads"))), __source__, ex))
        else
            $ex
        end
    end |> esc
end

for N in [8, 16, 32, 64, 128]
    T = Symbol(:Int, N)
    UT = Symbol(:UInt, N)
    @eval begin
        log2i(x::$T) =
            !signbit(x) ? ($(N - 1) - leading_zeros(x)) :
            throw(ErrorException("nonnegative expected ($x)"))
        log2i(x::$UT) = $(N - 1) - leading_zeros(x)
    end
end

bmask(args...) = bmask(Int, args...)
bmask(::Type{T}) where {T<:Integer} = zero(T)
bmask(::Type{T}, positions::Int...) where {T<:Integer} = bmask(T, positions)

function bmask(::Type{T}, itr) where {T<:Integer}
    isempty(itr) && return 0
    ret = zero(T)
    for b in itr
        ret += one(T) << (b - 1)
    end
    return ret
end

@inline function bmask(::Type{T}, range::UnitRange{Int})::T where {T<:Integer}
    ((one(T) << (range.stop - range.start + 1)) - one(T)) << (range.start - 1)
end

@inline lmove(b::Int, mask::Int, k::Int)::Int = (b & ~mask) << k + (b & mask)

@inline function _group_shift(masks::Vector{Int}, shift_len::Vector{Int}, k::Int, k_prv::Int)
    # if current position in the contiguous region
    # since these bits will be moved together with
    # the first one, we don't need to generate a
    # new mask
    if k == k_prv + 1
        shift_len[end] += 1
    else
        # we generate a bit mask where the 1st to k-th bits are 1
        push!(masks, bmask(0:k-1))
        push!(shift_len, 1)
    end
end

@inline function group_shift(locations)
    masks = Int[]
    shift_len = Int[]
    k_prv = -1
    for k in locations
        _group_shift(masks, shift_len, k, k_prv)
        k_prv = k
    end
    return masks, shift_len
end

@inline function complement_group_shift(n::Int, locations)
    masks = Int[]
    shift_len = Int[]
    k_prv = -1
    for k in 1:n
        k in locations && continue
        _group_shift(masks, shift_len, k, k_prv)
        k_prv = k
    end
    return masks, shift_len
end

struct BitSubspace
    n::Int # number of bits in fullspace
    sz_subspace::Int # size of the subspace
    n_shifts::Int # number of shifts
    masks::Vector{Int} # shift masks
    shift_len::Vector{Int} # length of each shift
end

function Base.getindex(s::BitSubspace, i::Int)
    index = i - 1
    @inbounds for k in 1:s.n_shifts
        index = lmove(index, s.masks[k], s.shift_len[k])
    end
    return index
end

Base.firstindex(s::BitSubspace) = 1
Base.lastindex(s::BitSubspace) = s.sz_subspace
Base.length(s::BitSubspace) = s.sz_subspace
Base.eltype(::BitSubspace) = Int

function Base.iterate(s::BitSubspace, st::Int = 1)
    st <= length(s) || return
    return s[st], st + 1
end

function bsubspace(n::Int, locs)
    @assert issorted(locs)
    masks, shift_len = group_shift(locs)
    BitSubspace(n, 1 << (n - length(locs)), length(masks), masks, shift_len)
end

function bcomspace(n::Int, locs)
    @assert issorted(locs)
    masks, shift_len = complement_group_shift(n, locs)
    BitSubspace(n, 1 << length(locs), length(masks), masks, shift_len)
end

log2dim1(x) = log2i(size(x, 1))

function broutine!(st::AbstractVector, U::AbstractMatrix, locs::NTuple{N, Int}) where N
    n = log2dim1(st)
    subspace = bsubspace(n, locs)
    comspace = bcomspace(n, locs)
    indices = [idx + 1 for idx in comspace]
    for k in subspace
        idx = indices .+ k
        st[idx] = U * st[idx]
    end
    return st
end

function Base.show(io::IO, ::MIME"text/plain", s::BitSubspace)
    indent = get(io, :indent, 0)
    println(io, " "^indent, s.sz_subspace, "-element BitSubspace:")
    if s.sz_subspace < 5
        for k in 1:s.sz_subspace
            print(io, " "^(indent+1), string(s[k]; base=2, pad=s.n))
            if k != s.sz_subspace
                println(io)
            end
        end
    else # never print more than 4 elements
        println(io, " "^(indent+1), string(s[1]; base=2, pad=s.n))
        println(io, " "^(indent+1), string(s[2]; base=2, pad=s.n))
        println(io, " "^(indent+1), "⋮")
        println(io, " "^(indent+1), string(s[end-1]; base=2, pad=s.n))
        print(io, " "^(indent+1), string(s[end]; base=2, pad=s.n))
    end
end

function subspace_mul!(st::AbstractVector{T}, comspace, U, subspace, offset=0) where T
    # pretend this is deleted
    # if size(U, 1) == 2
    #     subspace_mul2x2!(st, comspace, U, subspace, offset)
    # else
    if size(U, 1) == 4
        subspace_mul4x4!(st, comspace, U, subspace, offset)
    elseif size(U, 1) == 8
        subspace_mul8x8!(st, comspace, U, subspace, offset)
    else
        subspace_mul_generic!(st, comspace, U, subspace, offset)
    end
    return st
end

# pretend this is deleted
# function subspace_mul2x2!(st::AbstractVector{T}, comspace, U, subspace, offset=0) where T
#     indices_1 = comspace[1] + 1
#     indices_2 = comspace[2] + 1
#     @inbounds for k in subspace
#         idx_1 = indices_1 + k + offset
#         idx_2 = indices_2 + k + offset

#         T1 = U[1, 1] * st[idx_1] + U[1, 2] * st[idx_2]
#         T2 = U[2, 1] * st[idx_1] + U[2, 2] * st[idx_2]

#         st[idx_1] = T1
#         st[idx_2] = T2
#     end
#     return st
# end

function subspace_mul4x4!(st::AbstractVector{T}, comspace, U, subspace, offset=0) where T
    Base.Cartesian.@nextract 4 indices i -> comspace[i] + 1

    Base.Cartesian.@nextract 4 U i->begin
        Base.Cartesian.@nextract 4 U_i j->U[i, j]
    end

    for k in subspace
        Base.Cartesian.@nextract 4 idx i-> k + indices_i + offset

        Base.Cartesian.@nexprs 4 i -> begin
            y_i = zero(T)
            Base.Cartesian.@nexprs 4 j -> begin
                y_i += U_i_j * st[idx_j]
            end
        end

        Base.Cartesian.@nexprs 4 i -> begin
            st[idx_i] = y_i
        end
    end
    return st
end

function subspace_mul8x8!(st::AbstractVector{T}, comspace, U, subspace, offset=0) where T
    Base.Cartesian.@nextract 8 indices i -> comspace[i] + 1

    Base.Cartesian.@nextract 8 U i->begin
        Base.Cartesian.@nextract 8 U_i j->U[i, j]
    end

    @inbounds for k in subspace
        Base.Cartesian.@nextract 8 idx i-> k + indices_i + offset

        Base.Cartesian.@nexprs 8 i -> begin
            y_i = zero(T)
            Base.Cartesian.@nexprs 8 j -> begin
                y_i += U_i_j * st[idx_j]
            end
        end

        Base.Cartesian.@nexprs 8 i -> begin
            st[idx_i] = y_i
        end
    end
    return st
end

function subspace_mul_generic!(st::AbstractVector{T}, indices, U, subspace, offset=0) where T
    y = similar(st, (size(U, 1), ))
    idx = similar(indices)

    @inbounds for k in subspace
        for i in 1:size(U, 1)
            idx[i] = indices[i] + k + offset
        end

        for i in 1:size(U, 1)
            y[i] = zero(T)
            for j in 1:size(U, 2)
                y[i] += U[i, j] * st[idx[j]]
            end
        end

        for i in 1:size(U, 1)
            st[idx[i]] = y[i]
        end
    end
    return st
end

function broutine2x2!(st::AbstractVector{T}, U::AbstractMatrix, locs::Tuple{Int}) where T
    U11 = U[1, 1]; U12 = U[1, 2];
    U21 = U[2, 1]; U22 = U[2, 2];
    step_1 = 1 << (first(locs) - 1)
    step_2 = 1 << first(locs)

    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 1)-step_1
            ST1 = U11 * st[j + 1] + U12 * st[j + 1 + step_1]
            ST2 = U21 * st[j + 1] + U22 * st[j + 1 + step_1]

            st[j + 1] = ST1
            st[j + 1 + step_1] = ST2
        end
    elseif step_1 == 2
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 2 i->begin
                ST1 = U11 * st[j + i] + U12 * st[j + i + step_1]
                ST2 = U21 * st[j + i] + U22 * st[j + i + step_1]
                st[j + i] = ST1
                st[j + i + step_1] = ST2
            end
        end
    elseif step_1 == 4
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 4 i->begin
                ST1 = U11 * st[j + i] + U12 * st[j + i + step_1]
                ST2 = U21 * st[j + i] + U22 * st[j + i + step_1]
                st[j + i] = ST1
                st[j + i + step_1] = ST2
            end
        end
    elseif step_1 == 8
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 8 i->begin
                ST1 = U11 * st[j + i] + U12 * st[j + i + step_1]
                ST2 = U21 * st[j + i] + U22 * st[j + i + step_1]
                st[j + i] = ST1
                st[j + i + step_1] = ST2
            end
        end
    else
        for j in 0:step_2:size(st, 1)-step_1
            for i in j:8:j+step_1-1
                Base.Cartesian.@nexprs 8 k->begin
                    ST1 = U11 * st[i + k] + U12 * st[i + step_1 + k]
                    ST2 = U21 * st[i + k] + U22 * st[i + step_1 + k]
                    st[i + k] = ST1
                    st[i + step_1 + k] = ST2
                end
            end
        end
    end
    return st
end

function broutine!(st::AbstractVector, U::AbstractMatrix, locs::NTuple{N, Int}) where N
    size(U, 1) == 2 && return broutine2x2!(st, U, locs)
    n = log2dim1(st)
    subspace = bsubspace(n, locs)
    comspace = bcomspace(n, locs)
    subspace_mul!(st, comspace, U, subspace)
    return st
end

function threaded_broutine!(st::AbstractVector, U::AbstractMatrix, locs::NTuple{N, Int}) where N
    n = log2dim1(st)
    subspace = bsubspace(n, locs)
    comspace = bcomspace(n, locs)
    indices = [idx + 1 for idx in comspace]
    @inbounds @views Threads.@threads for k in subspace
        idx = indices .+ k
        st[idx] = U * st[idx]
    end
    return st
end

ctrl_offset(locs, configs) = bmask(locs[i] for (i, u) in enumerate(configs) if u != 0)
ismatch(index::T, mask::T, target::T) where {T<:Integer} = (index & mask) == target

function broutine2x2!(st::AbstractVector, U::AbstractMatrix, locs::Tuple{Int}, ctrl_locs::NTuple{M, Int}, ctrl_configs::NTuple{M, Int}) where {N, M}
    step_1 = 1 << (first(locs) - 1)
    step_2 = 1 << first(locs)
    ctrl_mask = bmask(ctrl_locs)
    flag_mask = reduce(+, 1 << (ctrl_locs[i] - 1) for i in 1:length(ctrl_locs) if ctrl_configs[i] == 1)
    U11 = U[1, 1]; U12 = U[1, 2];
    U21 = U[2, 1]; U22 = U[2, 2];

    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 1)-step_1
            if ismatch(j, ctrl_mask, flag_mask)
                ST1 = U11 * st[j+1] + U12 * st[j + step_1 + 1]
                ST2 = U21 * st[j+1] + U22 * st[j + step_1 + 1]

                st[j + 1] = ST1
                st[j + step_1 + 1] = ST2
            end
        end
    elseif step_1 == 2
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 2 i->begin
                if ismatch(j+i-1, ctrl_mask, flag_mask)
                    ST1 = U11 * st[j+i] + U12 * st[j + step_1 + i]
                    ST2 = U21 * st[j+i] + U22 * st[j + step_1 + i]

                    st[j + i] = ST1
                    st[j + step_1 + i] = ST2
                end
            end
        end
    elseif step_1 == 4
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 4 i->begin
                if ismatch(j+i-1, ctrl_mask, flag_mask)
                    ST1 = U11 * st[j+i] + U12 * st[j + step_1 + i]
                    ST2 = U21 * st[j+i] + U22 * st[j + step_1 + i]

                    st[j + i] = ST1
                    st[j + step_1 + i] = ST2
                end
            end
        end
    elseif step_1 == 8
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 8 i->begin
                if ismatch(j+i-1, ctrl_mask, flag_mask)
                    ST1 = U11 * st[j+i] + U12 * st[j + step_1 + i]
                    ST2 = U21 * st[j+i] + U22 * st[j + step_1 + i]

                    st[j + i] = ST1
                    st[j + step_1 + i] = ST2
                end
            end
        end
    else
        for j in 0:step_2:size(st, 1)-step_1
            for i in j:8:j+step_1-1
                Base.Cartesian.@nexprs 8 k->begin
                    if ismatch(i+k-1, ctrl_mask, flag_mask)
                        ST1 = U11 * st[i+k] + U12 * st[i + step_1 + k]
                        ST2 = U21 * st[i+k] + U22 * st[i + step_1 + k]
                        st[i + k] = ST1
                        st[i + step_1 + k] = ST2
                    end
                end
            end
        end
    end
    return st
end

function broutine!(st::AbstractVector, U::AbstractMatrix, locs::NTuple{N, Int}, ctrl_locs::NTuple{M, Int}, ctrl_configs::NTuple{M, Int}) where {N, M}
    size(U, 1) == 2 && return broutine2x2!(st, U, locs, ctrl_locs, ctrl_configs)
    n = log2dim1(st)
    subspace = bsubspace(n, sort([locs..., ctrl_locs...]))
    comspace = bcomspace(n, locs)
    offset = ctrl_offset(ctrl_locs, ctrl_configs)
    subspace_mul!(st, comspace, U, subspace, offset)
    return st
end

function broutine!(st::AbstractVector{T}, ::Val{:X}, locs::Tuple{Int}) where T
    step_1 = 1 << (first(locs) - 1)
    step_2 = 1 << first(locs)

    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 1)-step_1
            tmp = st[j + 1]
            st[j + 1] = st[j + 1 + step_1]
            st[j + 1 + step_1] = tmp
        end
    elseif step_1 == 2
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 2 i->begin
                tmp = st[j + i]
                st[j + i] = st[j + i + step_1]
                st[j + i + step_1] = tmp
            end
        end
    elseif step_1 == 4
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 4 i->begin
                tmp = st[j + i]
                st[j + i] = st[j + i + step_1]
                st[j + i + step_1] = tmp
            end
        end
    elseif step_1 == 8
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 8 i->begin
                tmp = st[j + i]
                st[j + i] = st[j + i + step_1]
                st[j + i + step_1] = tmp
            end
        end
    else
        for j in 0:step_2:size(st, 1)-step_1
            for i in j:8:j+step_1-1
                Base.Cartesian.@nexprs 8 k->begin
                    tmp = st[i+k]
                    st[i+k] = st[i + step_1 + k]
                    st[i + step_1 + k] = tmp
                end
            end
        end
    end
    return st
end

function broutine!(st::AbstractVector{T}, ::Val{:Rz}, locs::Tuple{Int}, theta::Real) where T
    U11 = exp(-im * theta/2)
    U22 = exp(im * theta/2)
    step_1 = 1 << (first(locs) - 1)
    step_2 = 1 << first(locs)

    @inbounds if step_1 == 1
        for j in 0:step_2:size(st, 1)-step_1
            ST1 = U11 * st[j + 1]
            ST2 = U22 * st[j + 1 + step_1]

            st[j + 1] = ST1
            st[j + 1 + step_1] = ST2
        end
    elseif step_1 == 2
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 2 i->begin
                ST1 = U11 * st[j + i]
                ST2 = U22 * st[j + i + step_1]
                st[j + i] = ST1
                st[j + i + step_1] = ST2
            end
        end
    elseif step_1 == 4
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 4 i->begin
                ST1 = U11 * st[j + i]
                ST2 = U22 * st[j + i + step_1]
                st[j + i] = ST1
                st[j + i + step_1] = ST2
            end
        end
    elseif step_1 == 8
        for j in 0:step_2:size(st, 1)-step_1
            Base.Cartesian.@nexprs 8 i->begin
                ST1 = U11 * st[j + i]
                ST2 = U22 * st[j + i + step_1]
                st[j + i] = ST1
                st[j + i + step_1] = ST2
            end
        end
    else
        for j in 0:step_2:size(st, 1)-step_1
            for i in j:8:j+step_1-1
                Base.Cartesian.@nexprs 8 k->begin
                    ST1 = U11 * st[i + k]
                    ST2 = U22 * st[i + step_1 + k]
                    st[i + k] = ST1
                    st[i + step_1 + k] = ST2
                end
            end
        end
    end
    return st
end

using CUDA

function broutine!(st::CuVector{T}, U::AbstractMatrix, locs::Tuple{Int}) where T
    U11 = U[1, 1]; U12 = U[1, 2];
    U21 = U[2, 1]; U22 = U[2, 2];
    step_1 = 1 << (first(locs) - 1)
    step_2 = 1 << first(locs)

    function kernel(st)
        idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
        j = step_2 * idx - step_2
        for i in j+1:j+step_1
            ST1 = U11 * st[i] + U12 * st[i + step_1]
            ST2 = U21 * st[i] + U22 * st[i + step_1]

            st[i] = ST1
            st[i + step_1] = ST2
        end
        return
    end

    N = length(0:step_2:size(st, 1)-step_1)
    nblocks = ceil(Int, N/256)
    @cuda threads=256 blocks=nblocks kernel(st)
    return st
end

function broutine!(st::CuVector{T}, U::AbstractMatrix, locs::Tuple{Int}, ctrl_locs::NTuple{M, Int}, ctrl_configs::NTuple{M, Int}) where {T, M}
    step_1 = 1 << (first(locs) - 1)
    step_2 = 1 << first(locs)
    ctrl_mask = bmask(ctrl_locs)
    flag_mask = reduce(+, 1 << (ctrl_locs[i] - 1) for i in 1:length(ctrl_locs) if ctrl_configs[i] == 1)
    U11 = U[1, 1]; U12 = U[1, 2];
    U21 = U[2, 1]; U22 = U[2, 2];

    function kernel(st)
        idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
        j = step_2 * idx - step_2
        for i in j:j+step_1-1
            if ismatch(i, ctrl_mask, flag_mask)
                ST1 = U11 * st[i+1] + U12 * st[i + step_1 + 1]
                ST2 = U21 * st[i+1] + U22 * st[i + step_1 + 1]

                st[i + 1] = ST1
                st[i + step_1 + 1] = ST2
            end
        end
        return
    end

    N = length(0:step_2:size(st, 1)-step_1)
    nblocks = ceil(Int, N/256)
    @cuda threads=256 blocks=nblocks kernel(st)
    return st
end

using SymEngine

@vars α θ
st = Basic[1, α, 0, 0]
broutine!(st, [exp(-im * θ) 0; 0 exp(im * θ)], (1, ))

using BenchmarkTools

data = Dict(
    "X" => [],
    "T" => [],
    "H" => [],
    "CNOT" => [],
)

for n in 4:25
    st = rand(ComplexF64, 1<<n)
    t = @benchmark broutine!(r, $([0 1;1 0]), (3, )) setup=(r=copy($st))
    push!(data["X"], minimum(t).time)
end

for n in 4:25
    st = rand(ComplexF64, 1<<n)
    t = @benchmark broutine!(r, $([1 0;0 exp(im * π / 4)]), (3, )) setup=(r=copy($st))
    push!(data["T"], minimum(t).time)
end

for n in 4:25
    st = rand(ComplexF64, 1<<n)
    t = @benchmark broutine!(r, $([1/sqrt(2) 1/sqrt(2); 1/sqrt(2) -1/sqrt(2)]), (3, )) setup=(r=copy($st))
    push!(data["H"], minimum(t).time)
end

for n in 4:25
    st = rand(ComplexF64, 1<<n)
    t = @benchmark broutine!(r, $([0 1;1 0]), (2, ), (3, ), (1, )) setup=(r=copy($st))
    push!(data["CNOT"], minimum(t).time)
end

using JSON
d = JSON.parse(read("../../Documents/quantum-benchmarks/qulacs/data/Linux-CPython-3.8-64bit/0002_data.json", String))

qulacs = Dict(
    "X" => [each["stats"]["min"] * 1e9 for each in d["benchmarks"] if each["group"] == "X"],
    "T" => [each["stats"]["min"] * 1e9 for each in d["benchmarks"] if each["group"] == "T"],
    "H" => [each["stats"]["min"] * 1e9 for each in d["benchmarks"] if each["group"] == "H"],
    "CNOT" => [each["stats"]["min"] * 1e9 for each in d["benchmarks"] if each["group"] == "CNOT"],
)

write("data.json", JSON.json(data))

using Plots
p1 = plot(4:25, [qulacs["X"], data["X"]]; title="X gate", yaxis=:log10, xlabel="nqubits", ylabel="ns")
p2 = plot(4:25, [qulacs["T"], data["T"]]; title="T gate", yaxis=:log10, xlabel="nqubits", ylabel="ns")
p3 = plot(4:25, [qulacs["H"], data["H"]]; title="H gate", yaxis=:log10, xlabel="nqubits", ylabel="ns")
p4 = plot(4:25, qulacs["CNOT"]./data["CNOT"]; title="CNOT gate", yaxis=:log10, xlabel="nqubits", ylabel="ns")
plt = plot(p1, p2, p3, p4; layout=4)
savefig("benchmark.png")

# sanity check
# using YaoArrayRegister

# st = rand(ComplexF64, 1 << 10)
# broutine!(copy(st), [0 1;1 0], (3, )) ≈ instruct!(copy(st), Val(:X), (3, ))
# broutine!(copy(st), [1 0;0 exp(im * π / 4)], (3, )) ≈ instruct!(copy(st), [1 0;0 exp(im * π / 4)], (3, ))
# H = [1/sqrt(2) 1/sqrt(2); 1/sqrt(2) -1/sqrt(2)]
# broutine!(copy(st), H, (3, )) ≈ instruct!(copy(st), H, (3, ))
# broutine!(copy(st), [0 1;1 0], (2, ), (3, ), (1, )) ≈ instruct!(copy(st), Val(:X), (2, ), (3, ), (1, ))

ext: jl


text:
data=vec(readdlm("../kepler62ef_planets.txt",',',Float64) )
include("../test_ttv.jl")
@time ttv1,ttv2=test_ttv(5,40,20,data); # This call includes time to compile
Profile.clear_malloc_data()
@time ttv1,ttv2=test_ttv(5,40,20,data,WriteOutput=false,num_evals=100000);

println("# Ignore this: ",hash(ttv1)+hash(ttv2)) # This just makes sure optimizer doesn't optimize away important calculations.

# Test that results haven't changed much
data_inner_out=readdlm("inner_ttv.txt",Float64)
data_outer_out=readdlm("outer_ttv.txt",Float64)
data_inner_ref=readdlm("inner_ttv.txt.ref",Float64)
data_outer_ref=readdlm("outer_ttv.txt.ref",Float64)
max_abs_diff = max(maximum(abs(data_inner_out.-data_inner_ref)),maximum(abs(data_outer_out.-data_outer_ref)))
println("# Maximum difference: ",max_abs_diff)

max_stars_repo_name: nwtuchow/TTVFaster

max_issues_repo_name: nwtuchow/TTVFaster

max_forks_repo_name: nwtuchow/TTVFaster


text:
Base.@kwdef struct Model{T}
    name::String
    mesh::FEMesh
    ρ::Vector{T}
    D::Vector{SMatrix{3,3,T,9}}
    T₂::Vector{T}
    κ::Vector{T}
end


text:
using .MPI
import Libdl

# Low-level MPI handles.
const MPIHandle = Union{MPI.MPI_Comm, MPI.MPI_Info}

# MPI.jl wrapper types.
const MPIHandleWrapper = Union{MPI.Comm, MPI.Info}

const H5MPIHandle = let csize = sizeof(MPI.MPI_Comm)
    @assert csize in (4, 8)
    csize == 4 ? API.Hmpih32 : API.Hmpih64
end

h5_to_mpi_comm(handle::H5MPIHandle) = reinterpret(MPI.MPI_Comm, handle)
h5_to_mpi_info(handle::H5MPIHandle) = reinterpret(MPI.MPI_Info, handle)

mpi_to_h5(handle::MPIHandle) = reinterpret(H5MPIHandle, handle)
mpi_to_h5(mpiobj::MPIHandleWrapper) = mpi_to_h5(mpiobj.val)

"""
    MPIO(comm::MPI.Comm, info::MPI.Info)
    MPIO(comm::MPI.Comm; kwargs....)

The parallel MPI file driver. This requires the use of
[MPI.jl](https://github.com/JuliaParallel/MPI.jl), and a custom HDF5 binary that has been
built with MPI support.

- `comm` is the communicator over which the file will be opened.
- `info`/`kwargs` are MPI-IO options, and are passed to `MPI_FILE_OPEN`.

# External links

- $(h5doc("H5P_SET_FAPL_MPIO"))
- [Parallel HDF5](https://portal.hdfgroup.org/display/HDF5/Parallel+HDF5)
"""
struct MPIO <: Driver
    comm::MPI.Comm
    info::MPI.Info
end
 MPIO(comm::MPI.Comm; kwargs...) =
    MPIO(comm, MPI.Info(;kwargs...))

# Check whether the HDF5 libraries were compiled with parallel support.
try
    DRIVERS[API.h5fd_mpio_init()] = MPIO
    HDF5.HAS_PARALLEL[] = true
catch e
end

function set_driver!(fapl::Properties, mpio::MPIO)
    HDF5.has_parallel() || error(
        "HDF5.jl has no parallel support." *
        " Make sure that you're using MPI-enabled HDF5 libraries, and that" *
        " MPI was loaded before HDF5." *
        " See HDF5.jl docs for details."
    )
    # Note: HDF5 creates a COPY of the comm and info objects.
    GC.@preserve mpio begin
        API.h5p_set_fapl_mpio(fapl, mpi_to_h5(mpio.comm), mpi_to_h5(mpio.info))
    end
end

function get_driver(fapl::Properties, ::Type{MPIO})
    h5comm, h5info = API.h5p_get_fapl_mpio(fapl, H5MPIHandle)
    comm = MPI.Comm(h5_to_mpi_comm(h5comm))
    info = MPI.Info(h5_to_mpi_info(h5info))
    return MPIO(comm, info)
end

"""
    h5open(filename, [mode="r"], comm::MPI.Comm, [info::MPI.Info]; pv...)

Open or create a parallel HDF5 file using the MPI-IO driver.

Equivalent to `h5open(filename, mode; fapl_mpio=(comm, info), pv...)`.
Throws an informative error if the loaded HDF5 libraries do not include parallel
support.

See the [HDF5 docs](https://portal.hdfgroup.org/display/HDF5/H5P_SET_FAPL_MPIO)
for details on the `comm` and `info` arguments.
"""
function HDF5.h5open(
        filename::AbstractString, mode::AbstractString,
        comm::MPI.Comm, info::MPI.Info = MPI.Info(); pv...
    )
    HDF5.h5open(filename, mode; driver=MPIO(comm, info), pv...)
end

HDF5.h5open(filename::AbstractString, comm::MPI.Comm, args...; pv...) =
    HDF5.h5open(filename, "r", comm, args...; pv...)


text:
function __rendergeodesics(
    m::AbstractMetricParams{T},
    init_pos;
    image_width,
    image_height,
    kwargs...,
) where {T}
    image = zeros(T, (image_height, image_width))
    render_into_image!(
        m,
        init_pos,
        image;
        image_width = image_width,
        image_height = image_height,
        kwargs...,
    )
end

function __prerendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    cache::AbstractRenderCache;
    kwargs...,
) where {T}
    error("Not implemented for render cache strategy '$typeof(cache)'.")
end

function __prerendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    cache::DefaultCache;
    kwargs...,
) where {T}
    # default is endpoints for now
    __prerendergeodesics(m, init_pos, EndpointCache(); kwargs...)
end

function __prerendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    cache::SolutionCache;
    kwargs...,
) where {T}
    prerender_solutions(m, init_pos; kwargs...)
end

function __prerendergeodesics(
    m::AbstractMetricParams{T},
    init_pos,
    cache::EndpointCache;
    kwargs...,
) where {T}
    prerender_endpoints(m, init_pos; kwargs...)
end

#get_parameter_type(m, u, v, max_time) =
#    typeof(GeodesicTracer.integrator_problem(m, u, v, (0.0, max_time)).p)

function prerender_endpoints(
    m::AbstractMetricParams{T},
    init_pos::V;
    image_width,
    image_height,
    fov_factor,
    max_time,
    solver_opts...,
) where {T,V}
    y_mid = image_height ÷ 2
    x_mid = image_width ÷ 2

    #u_cache = Matrix{V}(undef, (image_height, image_width))
    #v_cache = Matrix{V}(undef, (image_height, image_width))
    #t_cache = Matrix{T}(undef, (image_height, image_width))
    #p_cache = Matrix{get_parameter_type(m, init_pos, init_pos, max_time)}(undef, (image_height, image_width))
    #retcode_cache = Matrix{Symbol}(undef, (image_height, image_width))

    point_cache =
        Matrix{GeodesicBase.geodesic_point_type(m)}(undef, (image_height, image_width))

    vs = fill(init_pos, image_width)
    us = fill(init_pos, image_width)

    for Y = 1:image_height

        β = T(y_to_β(Y, y_mid, fov_factor))
        α_generator_row = (T(x_to_α(X, x_mid, fov_factor)) for X = 1:image_width)

        calculate_velocities!(vs, m, init_pos, α_generator_row, β)
        simsols =
            tracegeodesics(m, us, vs, (T(0.0), max_time); save_on = false, solver_opts...)
        println("+ $Y / $image_height ...")

        point_cache[Y, :] = get_endpoint(m, simsols)
    end

    EndpointRenderCache(m, max_time, image_height, image_width, point_cache)
end

function prerender_solutions(
    m::AbstractMetricParams{T},
    init_pos;
    image_width,
    image_height,
    fov_factor,
    max_time,
    solver_opts...,
) where {T}
    y_mid = image_height ÷ 2
    x_mid = image_width ÷ 2

    # this hits the garbage collector like a truck
    # need a better pre-cache for the velocity and position vectors
    # maybe using a row by row renderer as before?
    # or can we somehow repeat us without copying, as with FillArrays.jl?
    # -- tried that, doesn't make enough of a difference

    vs = fill(init_pos, image_width)
    us = fill(init_pos, image_width)

    simsol_array = map(1:image_height) do Y

        β = T(y_to_β(Y, y_mid, fov_factor))
        α_generator_row = (T(x_to_α(X, x_mid, fov_factor)) for X = 1:image_width)

        calculate_velocities!(vs, m, init_pos, α_generator_row, β)
        simsols = tracegeodesics(
            m,
            us,
            vs,
            (T(0.0), max_time);
            save_on = false,
            solver_opts...,
        )
        println("+ $Y / $image_height ...")
        simsols
    end

    SolutionRenderCache(m, max_time, image_height, image_width, simsol_array)
end

function render_into_image!(
    m::AbstractMetricParams{T},
    init_pos,
    image;
    image_width,
    image_height,
    fov_factor,
    max_time,
    vf,
    solver_opts...,
) where {T}
    y_mid = image_height ÷ 2
    x_mid = image_width ÷ 2

    vs = fill(init_pos, image_width)
    us = fill(init_pos, image_width)
    for Y = 1:image_height

        β = T(y_to_β(Y, y_mid, fov_factor))
        α_generator_row = (T(x_to_α(X, x_mid, fov_factor)) for X = 1:image_width)

        calculate_velocities!(vs, m, init_pos, α_generator_row, β)

        # do the render
        simsols =
            tracegeodesics(m, us, vs, (T(0.0), max_time); save_on = false, solver_opts...)
        apply_to_location!(m, @view(image[Y, :]), simsols, vf, max_time)

        println("+ $Y / $image_height ...")
    end

    image
end

function apply_to_location!(m::AbstractMetricParams{T}, loc, sols, vf, max_time) where {T}
    @threads for i = 1:length(sols)
        loc[i] = vf(m, get_endpoint(m, sols[i]), max_time)
    end
end


text:
using JuMP, LinearAlgebra, Random, Gurobi, DataFrames, CSV, Printf;
root=("Please specify the root directory here")
include("Utilities.jl")
include("FullSpace.jl")
#------------------------------------------------------------------------------
#------------------------------------------------------------------------------
function SolveModel(A, obj, b, n)
    # Solves the forward optimization problem

# Step 1: Create the model
    m1 = Model(solver=GurobiSolver(OutputFlag = 0))
    @variable(m1, x[1:n])
    @objective(m1, Max, dot(obj, x))
    for i = 1:length(A[:, 1])
        @constraint(m1, dot(A[i, :], x) <= b[i])
    end

# Step 2: Solve the model
    status = solve(m1; suppress_warnings = true)
# Step 3: return optimal value;
    val = getvalue(x);
end

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

function SolveFP(cone, iteration)
    # Feasibility Problem to be solved at every iteration of the decomposition algorithm
    epsilon = 1
    m2 = Model(solver=GurobiSolver(OutputFlag = 1))
    @variable(m2, y[1:n]>=1e-4)
    @variable(m2, gamma[1:iteration, 1:c] >= 0)
    @variable(m2, w, Bin)

    for i = 1:iteration
        for j = 1:n
            @constraint(m2, y[j] == sum(gamma[i,k]*cone[i][k,j] for k = 1:length(cone[i][:,1])))
        end
    end

    @constraint(m2, sum(y) == epsilon)

    @objective(m2, Min, 0)
    status = solve(m2; suppress_warnings = true)
    return status
end

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

function SolveP2(C_bar, cone, iteration)

    m3 = Model(solver=GurobiSolver(OutputFlag = 0))
    @variable(m3, y[1:n])
    @variable(m3, gamma[1:iteration, 1:c] >= 0)

    for i = 1:iteration
        for j = 1:n
            @constraint(m3, y[j] == sum(gamma[i,k]*cone[i][k,j] for k = 1:length(cone[i][:,1])))
        end
    end

    @objective(m3, Min, sum((y[j]-C_bar[j])^2 for j = 1:n))
    status = solve(m3; suppress_warnings = true)
    return getvalue(y), getobjectivevalue(m3)
end

#------------------------------------------------------------------------------
#------------------------------------------------------------------------------

function SolveP1(A, b, x, NoS, J)
    # NoS (number of samples, i.e., I is always 1 here.)
    n = length(A[1, :, 1])
    m4 = Model(solver = GurobiSolver(OutputFlag = 0, IntFeasTol = 1e-9, MIPGapAbs = 1e-5, MIPGap = 0))
    M = 50
    epsilon = 1

    @variable(m4, xbar[1:NoS, 1:n], start = 0)
    @variable(m4, lambda[1:NoS, 1:c] >= 0)
    @variable(m4, s[1:NoS, 1:c] >= 0)
    @variable(m4, z[1:NoS, 1:c], Bin)
    @variable(m4, C[1:n]>=1e-4)
    @variable(m4, t[1:J, 1:n, 1:NoS] >= 0)

    @constraint(m4, S[k = 1:NoS, j = 1:n], C[j] - sum(lambda[k, i]*A[i, j, k] for i = 1:c) == 0)
    @constraint(m4, PF[k = 1:NoS, i = 1:c], sum(A[i, j, k]*xbar[k, j] for j = 1:n) + s[k, i] == b[i])
    @constraint(m4, CS1[k = 1:NoS, i = 1:c], s[k, i] <= M*(1-z[k, i]))
    @constraint(m4, CS2[k = 1:NoS, i = 1:c], lambda[k, i] <= M*z[k, i])
    @constraint(m4, VC[k = 1:NoS] , sum(z[k,:]) >= n)
    @constraint(m4, linearization1[j = 1:J, i = 1:n, k = 1:NoS], t[j,i,k] >= xbar[k,i] - x[j,i,k])
    @constraint(m4, linearization2[j = 1:J, i = 1:n, k = 1:NoS], t[j,i,k] >= -(xbar[k,i] - x[j,i,k]))
    @constraint(m4, sum(C) == epsilon)

    @objective(m4, Min, sum(t))
    status = solve(m4; suppress_warnings = true)
    loss1 = getobjectivevalue(m4)

    return getvalue(xbar), loss1, getvalue(z)
end

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

function SolveP1m(A, b, x, NoS, J, instance, initx, initz)
    # Solves P1 with multiple (NoS > 1) experiments
    n = length(A[1, :, 1])
    initz = round.(initz)
    m5 = Model(solver = GurobiSolver(OutputFlag = 1, TimeLimit = 7200, ConcurrentMIP = 3,
    IntFeasTol = 1e-9, MIPGapAbs = 1e-5, MIPGap = 0))
    M = 20
    epsilon = 1

    @variable(m5, xbar[j = 1:NoS, k = 1:n])
    @variable(m5, lambda[1:NoS, 1:c] >= 0)
    @variable(m5, s[1:NoS, 1:c] >= 0)
    @variable(m5, z[j = 1:NoS, k = 1:c], Bin)
    for j = 1:NoS-1
        for k = 1:c
            setvalue(z[j,k], initz[j,k])
        end
    end
    @variable(m5, C[1:n]>=1e-4)
    @variable(m5, t[1:J, 1:n, 1:NoS] >= 0)

    @constraint(m5, S[k = 1:NoS, j = 1:n], C[j] - sum(lambda[k, i]*A[i, j, k] for i = 1:c) == 0)
    @constraint(m5, PF[k = 1:NoS, i = 1:c], sum(A[i, j, k]*xbar[k, j] for j = 1:n) + s[k, i] == b[k, i])
    @constraint(m5, CS1[k = 1:NoS, i = 1:c], s[k, i] <= M*(1-z[k, i]))
    @constraint(m5, CS2[k = 1:NoS, i = 1:c], lambda[k, i] <= M*z[k, i])
    @constraint(m5, VC[k = 1:NoS] , sum(z[k,:]) >= n)
    @constraint(m5, epsilon == sum(C))
    @constraint(m5, linearization1[j = 1:J, i = 1:n, k = 1:NoS], t[j,i,k] >= xbar[k,i] - x[j,i,k])
    @constraint(m5, linearization2[j = 1:J, i = 1:n, k = 1:NoS], t[j,i,k] >= -(xbar[k,i] - x[j,i,k]))

    @objective(m5, Min, sum(t))
    status = solve(m5; suppress_warnings = true)
    loss = getobjectivevalue(m5)
    return getvalue(xbar), loss, getvalue(z), status
end

#------------------------------------------------------------------------------

function accuracy(x, x_hat)
NoS = length(x[:, 1])
n = length(x[1,:])
error = 0
    for i = 1:NoS
        if !isapprox(norm(x[i,:] - x_hat[i,:]), 0, atol = 1e-3)
            error += 1
        end
    end
    return error/NoS
end

#------------------------------------------------------------------------------
#   Main body of the program
#------------------------------------------------------------------------------
rng = MersenneTwister(1234)

# Edit the following parameters
n = 25 # Dimension of the forward problem
J_scheme = 1

# Dummy A and b
A = zeros(Float64, (2*dim)+1, dim) # Initialize A matrix
b = zeros(Float64, (2*dim)+1) # Initialize b vector
obj = rand(rng,1:1000,dim)
obj = obj/ norm(obj, 1) # Generate random objective coefficients

for i = 1:dim
    A[i, i] = 1
    b[i, 1] = 1
end

for i = 1:dim
    A[dim+1, i] = obj[i] + 100 + rand(rng, -10:10)
end
b[dim+1, 1] = rand(rng, sum(A[dim+1, :])/2:sum(A[dim+1, :]) - 1)

for i = 1:dim
    A[dim+1+i, i] = -1
    b[dim+1+i, 1] = 0
end

# Generate a random reference vector
noisy_p = rand(rng,1:1000,dim)
noisy_p = noisy_p/ norm(noisy_p, 1)

I = 100 # Number of experiments
TS = 100 # Size of test dataset

dim = n # Some part of this code uses dim for the dimension of the forward problem
c = length(A[:,1]) # Total number of constraints in the problem

for w in [10 20 100] # w decides the level of noise in data, (sigma = 1/w)
    global J # Initialize Number of samples per experiment
    if J_scheme == 1
        if w == 10
            J = 5
        elseif w == 20
            J = 10
        else
            J = 20
        end
    elseif J_scheme == 2
        J = 250
    end

    Total_Time = []
    Resolves = []

    Pred_Error_Evolution = zeros(Float64, I, 10)

    # Make directories to store results
    mkdir(string(root,"Results_dim_",dim,"_noise_", w))

    for instance = 1:10 # For every dim and w, solve 10 random instances

        t_instance_s = time() # # To store the time taken to solve one instance

        # Arrays to store training data
        b_train = zeros(Float64, I, c)
        A_train = zeros(Float64, c, n, I)
        true_x  = zeros(Float32, I, n)
        noisy_x = zeros(Float64, J, n, I)

        # Arrays to store test data
        test_x  = zeros(Float32, TS, n)
        b_test = zeros(Float64, TS, c)
        A_test = zeros(Float64, c, n, TS)

        # We also store some performance data at every step of the algorithm, initialize required arrays
        runtime = zeros(Float64,I)
        error =  zeros(Float64,I)
        Is_resolve = zeros(Float64,I)

        # Generate training data
        for T = 1:I
            b_train[T, :] = b
            A_train[:, :, T] = A
            for i = 1:dim
                A_train[dim+1, i, T] = 100 + rand(rng, -5000:5000)/100
            end
        	A_train[dim+1, :, T] = A_train[dim+1, :, T]/ 1000
            b_train[T, dim+1] = 0.6*sum(A_train[dim+1, :,1]) # Keep b the same across different experiments
            true_x[T, 1:n] = SolveModel(A_train[:, :, T], obj, b_train[T, :], n)
        end

        # Generate test data
        for test = 1:TS
            b_test[test, :] = b
            A_test[:, :, test] = A
            for i = 1:dim
                A_test[dim+1, i, test] = 100 + rand(rng, -5000:5000)/100
            end
            b_test[test, dim+1] = deepcopy(b_train[1, dim+1])*1000 # Same b as the training dataset
            test_x[test, 1:n] = SolveModel(A_test[:, :, test], obj, b_test[test, :], n)
        end

        # Add noise to the training data
        for T = 1:I
            for j = 1:J
                for k = 1:n
                    noisy_x[j, k, T] = true_x[T,k] + randn(rng, Float64)/w
                end
            end
        end

        # First solve the full-space (P1)
        batchlearning(obj, noisy_p, n, c, I, J, A_train, b_train, true_x, noisy_x, A_test, b_test, test_x)

        # Next solve the same problem with adaptive sampling
        resolve_counter = 0

        global cone = fill(zeros(n,n), 1) # Store the active constraints that form the cone for every experiment

        # Store estimates for warm-start (though we only use z for warm-starting)
        hat_x = zeros(Float64, I, n)
        hat_z = zeros(Float64, I, c)

        # Start the solution process
        for T = 1:I
            flag = 0 # Just to check if a resolve was required for this experiment or not
            # push!(sr, T)
            t_experiment_s = time() # To store the time taken to solve for one experiment
            istight = [] # To Store which constraints are active for this experiment

            # Solve (P1)_{\ell}, obtain \hat{x}_{\ell}^{*}
            hat_x[T, :], r, z = SolveP1(A_train[:,:,T], b_train[T,:], noisy_x[:,:,T], 1, J)
            # init_x[T, :] = temp_x
            hat_z[T, :] = z

            # Find the cone for current experiment
            for i = 1:c
                if isapprox(z[1, i], 1, atol = 1e-3) == 1
                    push!(istight, i)
                end
            end

            # To-Do: Check why this does not work in one step.
            if T == 1
                for j = 1:length(istight)
                    cone[T][j,:] = A_train[istight[j], :, T]
                end
            else
                B = zeros(Float64, length(istight), n) # Temporary matrix
                for j = 1:length(istight)
                    B[j,:] = A_train[istight[j], :, T]
                end
                push!(cone, B)
            end

            status = SolveFP(cone, T) # Check if (FP)_{\ell} is feasible

            if status == :Infeasible # Solve full problem if infeasible
                flag = 1 # Set flag to one because a resolve is required
                resolve_counter += 1

                z = []
                x, r, z, status = SolveP1m(A_train[:,:,1:T], b_train[1:T,:], noisy_x[:,:,1:T], T, J, instance, hat_x, hat_z)

                if status != :Optimal # If problem does not get solved in the cutoff time, stop this instance
                    break
                end

                # Update the estimated solutions with what we find after the resolve
                hat_x[1:T, :] = x
                hat_z[1:T, :] = z

                # Also update the cone
                global cone = []
                global cone = fill(zeros(n,n), 1)
                for k = 1:T
                    istight = []
                    for i = 1:c
                        if isapprox(z[k, i], 1, atol = 1e-3)
                            push!(istight, i)
                        end
                    end
                    if k == 1
                        for j = 1:length(istight)
                            cone[k][j,:] = A_train[istight[j], :, k]
                        end
                    else
                        B = zeros(Float64, length(istight), n)
                        for j = 1:length(istight)
                            B[j,:] = A_train[istight[j], :, k]
                        end
                        push!(cone, B)
                    end
                end

                status = SolveFP(cone, T) # This is not required but still do it.
                if status == :Infeasible
                    println("Potential problem with instance: ", instance)
                end
            end

            # Solve (P2)
            hat_c, loss = SolveP2(noisy_p, cone, T)

            # Finally compute the prediction error
            x_pred  = zeros(Float32, TS, n)

            for test = 1:TS
                x_pred[test, 1:n] = SolveModel(A_test[:, :, test], hat_c, b_test[test, :], n)
            end
            runtime[T] = time() - t_experiment_s
            error[T] = accuracy(test_x, x_pred)
            Is_resolve[T] = flag
        end
        ##----------------------------------------------

        # Store experiment level data
        cd(string(root,"Results_dim_",dim,"_noise_",w))
        df1 = DataFrame(Computation_Time = runtime, Prediction_Error = error, Resolve_Flag = Is_resolve)
        # Uncomment if required
        # s1 = @sprintf "InnerLoopResultsSummary_dim_%03.d_Noise_%03.d_Instance_%03.d.csv" dim w instance
        # CSV.write(s1, df1)

        push!(Total_Time, time() - t_instance_s)
        push!(Resolves, resolve_counter)
        Pred_Error_Evolution[:,instance] = error
    end

    # Store instance level data
    df2 = DataFrame(Instance = 1:1:I, Computation_Time = Total_Time, Resolves_reqd = Resolves)
    s2 = @sprintf "OuterLoopResultsSummary_dim_%03.d_Noise_%03.d.csv" dim w;
    CSV.write(s2, df2)
    df3 = DataFrame(Pred_Error_Evolution)
    s3 = @sprintf "ErrEvolSummary_dim_%03.d_Noise_%03.d.csv" dim w;
    CSV.write(s3, df3)
end

max_stars_repo_name: ddolab/InvLP

max_issues_repo_name: ddolab/InvLP

max_forks_repo_name: ddolab/InvLP


text:
#Test --- if ϵ (infectiousness of asymptomatics is zero) and external rates are zero then nothing should happen despite 1000 initial asymptomatics

function test_no_infecteds()
    #Load data completely susceptible Population
    u0,P,transport_matrix = KenyaCoV.model_ingredients_from_data("data/combined_population_estimates.csv",
                                                                 "data/optimal_transition_matrix.jld2",
                                                                "data/optimal_movement_matrix.jld2",
                                                                "data/flight_numbers.csv",
                                                                "data/projected_global_prevelance.csv")
    #You can modify other parameters directly
    P.τ = 0. #e.g. no treatment rate
    P.ϵ = 0.
    P.into_mom .= 0
    P.into_nai .= 0
    #Define initial conditions by modifying the completely susceptible population
    u0[30,3,1] += 1000#One asymptomatic in Nairobi

    #Go straight to solution using solver compiled in the KenyaCoV module
    sol_tl = KenyaCoV.solve_KenyaCoV_prob(u0,(0.,365.),P,1.)

    return ~any([sum(reshape(sol_tl[end],KenyaCoV.n,KenyaCoV.n_s,2)[i,7:8,1:2]) for i = 1:47] .!= 0)
end

function stays_nonnegative()
    u0,P,transport_matrix = KenyaCoV.model_ingredients_from_data("data/combined_population_estimates.csv",
                                                                 "data/optimal_transition_matrix.jld2",
                                                                 "data/optimal_movement_matrix.jld2",
                                                                 "data/flight_numbers.csv",
                                                                 "data/projected_global_prevelance.csv")
    P.τ = 0. #e.g. no treatment rate
    for (i,p) in enumerate(P.global_prev)
        P.global_prev[i] = 0.
    end
    # P.global_prev
    P.μ₁ = 10.
    P.ϵ = 1.
    P.β = 2.2*P.γ
    #Define initial conditions by modifying the completely susceptible population
    P.dt = 1.
    u0[30,4,1] += 10#One diseased in Nairobi
    prob = KenyaCoV.create_KenyaCoV_non_neg_prob(u0,(0.,365.),P)
    sol = solve(prob,FunctionMap(),dt = P.dt)
    ~any([any(u .< 0) for u in sol.u])
end
# function run_epidemic_infs_everywhere()
#     u0,P,transport_matrix = KenyaCoV.model_ingredients_from_data("data/combined_population_estimates.csv",
#                                                                  "data/optimal_transition_matrix.jld2",
#                                                                 "data/optimal_movement_matrix.jld2",
#                                                                 "data/flight_numbers.csv",
#                                                                 "data/projected_global_prevelance.csv")
#
#     u0[30,4,1] += 100#One diseased in Nairobi
#
#     #Create a JumpProblem which you can solve --- needs DifferentialEquations module for the solvers
#     jump_prob_tl = create_KenyaCoV_prob(u0,(0.,365.),P)
#     #Go straight to solution using solver compiled in the KenyaCoV module
#     sol_tl = solve_KenyaCoV_prob(u0,(0.,365.),P,1.)
# end


text:
# A tool for testing lossless video encoding

using VideoIO, ColorTypes, FixedPointNumbers, DataFrames

function collectexecoutput(exec::Cmd)
    out = Pipe(); err = Pipe()
    p = Base.open(pipeline(ignorestatus(exec), stdout=out, stderr=err))
    close(out.in); close(err.in)
    err_s = readlines(err); out_s = readlines(out)
    return (length(out_s) > length(err_s)) ? out_s : err_s
end

function createtestvideo(;filename::String="$(tempname()).mp4",duration::Real=5,
    width::Int64=1280,height::Int64=720,framerate::Real=30,testtype::String="testsrc2",
    encoder::String="libx264rgb")
    withenv(VideoIO.execenv) do
        collectexecoutput(`$(VideoIO.ffmpeg) -y -f lavfi -i
            $testtype=duration=$duration:size=$(width)x$(height):rate=$framerate
            -c:v $encoder -preset slow -crf 0 -c:a copy $filename`)
    end
    return filename
end

function testvideocomp!(df,preset,imgstack_gray)
    t = @elapsed VideoIO.save("video.mp4",imgstack_gray,framerate=30,codec_name = "libx264",
                                    encoder_options=(color_range=2, crf=0, "preset"=preset))
    fs = filesize("video.mp4")
    f = openvideo("video.mp4",target_format=VideoIO.AV_PIX_FMT_GRAY8)
    imgstack_gray_copy = []
    while !eof(f)
        push!(imgstack_gray_copy,read(f))
    end
    identical = !any(.!(imgstack_gray .== imgstack_gray_copy))
    push!(df,[preset,fs,t,identical])
end

imgstack_gray_noise = map(x->rand(Gray{N0f8},1280,720),1:1000)

f = openvideo(createtestvideo())
imgstack = []
while !eof(f)
    push!(imgstack,read(f))
end
imgstack_gray_testvid = map(x->convert.(Gray{N0f8},x),imgstack)

f = openvideo("videos/ladybird.mp4")
imgstack = []
while !eof(f)
    push!(imgstack,read(f))
end
imgstack_gray_ladybird = map(x->convert.(Gray{N0f8},x),imgstack)

df_noise = DataFrame(preset=[],filesize=[],time=[],identical=[])
df_testvid = DataFrame(preset=[],filesize=[],time=[],identical=[])
df_ladybird = DataFrame(preset=[],filesize=[],time=[],identical=[])
for preset in ["ultrafast","superfast","veryfast","faster","fast","medium","slow",
"slower","veryslow"]
    @show preset
    for rep in 1:3
        @show rep
        testvideocomp!(df_noise,preset,imgstack_gray_noise)
        testvideocomp!(df_testvid,preset,imgstack_gray_testvid)
        testvideocomp!(df_ladybird,preset,imgstack_gray_ladybird)
    end
end

noise_raw_size = size(imgstack_gray_noise[1],1) * size(imgstack_gray_noise[1],2) * length(imgstack_gray_noise)
testvid_raw_size = size(imgstack_gray_testvid[1],1) * size(imgstack_gray_testvid[1],2) * length(imgstack_gray_testvid)
ladybird_raw_size = size(imgstack_gray_ladybird[1],1) * size(imgstack_gray_ladybird[1],2) * length(imgstack_gray_ladybird)

df_noise[:filesize_perc] = 100*(df_noise[:filesize]./noise_raw_size)
df_testvid[:filesize_perc] = 100*(df_testvid[:filesize]./testvid_raw_size)
df_ladybird[:filesize_perc] = 100*(df_ladybird[:filesize]./ladybird_raw_size)
df_noise[:fps] = length(imgstack_gray_noise)./df_noise[:time]
df_testvid[:fps] = length(imgstack_gray_testvid)./df_testvid[:time]
df_ladybird[:fps] = length(imgstack_gray_ladybird)./df_ladybird[:time]

using Statistics
df_noise_summary = by(df_noise, :preset, identical = :identical => minimum, fps_mean = :fps => mean, fps_std = :fps => std, filesize_perc_mean = :filesize_perc => mean,filesize_perc_std = :filesize_perc => std)
df_testvid_summary = by(df_testvid, :preset, identical = :identical => minimum, fps_mean = :fps => mean, fps_std = :fps => std, filesize_perc_mean = :filesize_perc => mean,filesize_perc_std = :filesize_perc => std)
df_ladybird_summary = by(df_ladybird, :preset, identical = :identical => minimum, fps_mean = :fps => mean, fps_std = :fps => std, filesize_perc_mean = :filesize_perc => mean,filesize_perc_std = :filesize_perc => std)

@show df_noise_summary
@show df_testvid_summary
@show df_ladybird_summary

### Results (generated 2019-05-29 on a 2019 Macbook Pro)
### OUTDATED. Generated before change to VideoIO.save
#=
df_noise_summary = 9×6 DataFrame
│ Row │ preset    │ identical │ fps_mean │ fps_std │ filesize_perc_mean │ filesize_perc_std │
│     │ Any       │ Bool      │ Float64  │ Float64 │ Float64            │ Float64           │
├─────┼───────────┼───────────┼──────────┼─────────┼────────────────────┼───────────────────┤
│ 1   │ ultrafast │ true      │ 92.5769  │ 8.40224 │ 156.444            │ 0.0               │
│ 2   │ superfast │ true      │ 62.3509  │ 1.19652 │ 144.019            │ 0.0               │
│ 3   │ veryfast  │ true      │ 59.9182  │ 1.77294 │ 144.019            │ 0.0               │
│ 4   │ faster    │ true      │ 60.3482  │ 2.32679 │ 144.02             │ 0.0               │
│ 5   │ fast      │ true      │ 149.169  │ 1.56068 │ 100.784            │ 0.0               │
│ 6   │ medium    │ true      │ 146.141  │ 3.41282 │ 100.784            │ 0.0               │
│ 7   │ slow      │ true      │ 147.214  │ 1.23929 │ 100.784            │ 0.0               │
│ 8   │ slower    │ true      │ 138.808  │ 2.553   │ 100.784            │ 0.0               │
│ 9   │ veryslow  │ true      │ 132.505  │ 3.28558 │ 100.784            │ 0.0               │

df_testvid_summary = 9×6 DataFrame
│ Row │ preset    │ identical │ fps_mean │ fps_std │ filesize_perc_mean │ filesize_perc_std │
│     │ Any       │ Bool      │ Float64  │ Float64 │ Float64            │ Float64           │
├─────┼───────────┼───────────┼──────────┼─────────┼────────────────────┼───────────────────┤
│ 1   │ ultrafast │ true      │ 228.166  │ 75.1439 │ 4.80392            │ 0.0               │
│ 2   │ superfast │ true      │ 239.73   │ 54.2033 │ 3.62199            │ 0.0               │
│ 3   │ veryfast  │ true      │ 197.506  │ 13.1121 │ 3.59901            │ 0.0               │
│ 4   │ faster    │ true      │ 174.174  │ 18.0316 │ 3.60282            │ 0.0               │
│ 5   │ fast      │ true      │ 235.181  │ 7.40358 │ 3.44104            │ 0.0               │
│ 6   │ medium    │ true      │ 219.654  │ 3.27445 │ 3.40832            │ 0.0               │
│ 7   │ slow      │ true      │ 171.337  │ 3.92415 │ 3.33917            │ 0.0               │
│ 8   │ slower    │ true      │ 105.24   │ 6.59151 │ 3.25774            │ 5.43896e-16       │
│ 9   │ veryslow  │ true      │ 63.1136  │ 2.47291 │ 3.2219             │ 0.0               │

df_ladybird_summary = 9×6 DataFrame
│ Row │ preset    │ identical │ fps_mean │ fps_std  │ filesize_perc_mean │ filesize_perc_std │
│     │ Any       │ Bool      │ Float64  │ Float64  │ Float64            │ Float64           │
├─────┼───────────┼───────────┼──────────┼──────────┼────────────────────┼───────────────────┤
│ 1   │ ultrafast │ true      │ 176.787  │ 36.5227  │ 12.2293            │ 0.0               │
│ 2   │ superfast │ true      │ 135.925  │ 7.04431  │ 10.3532            │ 0.0               │
│ 3   │ veryfast  │ true      │ 117.115  │ 1.28102  │ 10.1954            │ 0.0               │
│ 4   │ faster    │ true      │ 94.39    │ 3.48494  │ 9.85604            │ 0.0               │
│ 5   │ fast      │ true      │ 69.657   │ 1.61004  │ 9.62724            │ 0.0               │
│ 6   │ medium    │ true      │ 54.9621  │ 0.568074 │ 9.51032            │ 0.0               │
│ 7   │ slow      │ true      │ 37.8888  │ 1.27484  │ 9.33622            │ 0.0               │
│ 8   │ slower    │ true      │ 20.1112  │ 1.04282  │ 9.25529            │ 0.0               │
│ 9   │ veryslow  │ true      │ 10.0016  │ 0.473213 │ 9.24999            │ 0.0               │
=#

# HISTOGRAM COMPARISON - useful for diagnosing range compression
# using PyPlot, ImageCore
# figure()
# hist(rawview(channelview(imgstack_gray_copy[1]))[:],0:256,label="copy")
# hist(rawview(channelview(imgstack_gray[1]))[:],0:256,label="original")
# legend()


text:
# module FourRooms

import Random, Base.size

# import IWER.GeneralValueFunction

module FourRoomsContParams

const BASE_WALLS = [0 0 0 0 0 1 0 0 0 0 0;
                    0 0 0 0 0 1 0 0 0 0 0;
                    0 0 0 0 0 0 0 0 0 0 0;
                    0 0 0 0 0 1 0 0 0 0 0;
                    0 0 0 0 0 1 0 0 0 0 0;
                    1 0 1 1 1 1 0 0 0 0 0;
                    0 0 0 0 0 1 1 1 0 1 1;
                    0 0 0 0 0 1 0 0 0 0 0;
                    0 0 0 0 0 1 0 0 0 0 0;
                    0 0 0 0 0 0 0 0 0 0 0;
                    0 0 0 0 0 1 0 0 0 0 0;]

const UP = 1
const RIGHT = 2
const DOWN = 3
const LEFT = 4

const ACTIONS = [UP, RIGHT, DOWN, LEFT]

const ROOM_TOP_LEFT = 1
const ROOM_TOP_RIGHT = 2
const ROOM_BOTTOM_LEFT = 3
const ROOM_BOTTOM_RIGHT = 4

const BODY_RADIUS = 0.1
const STEP = 0.5

const AGENT_BOUNDRIES = [
    [BODY_RADIUS, 0.0],
    [-BODY_RADIUS, 0.0],
    [0.0, BODY_RADIUS],
    [0.0, -BODY_RADIUS]
]

end

"""
    FourRoomsCont

    Four Rooms environment using RLCore abstract environment.
    - state: [y, x]

"""
mutable struct FourRoomsCont <: AbstractEnvironment
    state::Array{Float64, 1}
    walls::Array{Bool, 2}
    max_action_noise::Float64
    drift_noise::Float64
    normalized::Bool
    collision::Bool

    edge_locs::Array{Array{Int64, 1}, 1}
    collision_check::Array{Bool, 1}
    new_state::Array{Float64, 1}
    FourRoomsCont(size, walls, max_action_noise, drift_noise, normalized) =
        new(size, walls, max_action_noise, drift_noise, normalized, false, [[0,0] for i in 1:4], fill(false, 4), zeros(2))
end

FourRoomsCont(max_action_noise = 0.1, drift_noise = 0.001) =
    FourRoomsCont([0.0,0.0], convert(Array{Bool}, FourRoomsContParams.BASE_WALLS), max_action_noise, drift_noise, false)
FourRoomsCont(walls::Array{Int64, 2}) = FourRoomsCont([0.0,0.0], convert(Array{Bool}, walls), 0.0, 0.0, false)
function FourRoomsCont(size::Int, wall_list::Array{CartesianIndex{2}}, max_action_noise::Float64, drift_noise::Float64)
    walls = fill(false, size[1], size[2])
    for wall in wall_list
        walls[wall] = true
    end
    FourRoomsCont([0.0,0.0], walls, max_noise, drift_noise, false)
end

RLCore.is_terminal(env::FourRoomsCont) = false
RLCore.get_reward(env::FourRoomsCont) = 0
RLCore.get_state(env::FourRoomsCont) = (env.normalized ? env.state./size(env) : env.state, env.collision)
project(env::FourRoomsCont, state) = [Int64(floor(state[1])) + 1, Int64(floor(state[2])) + 1]
project(env::FourRoomsCont, state, loc) = begin; loc[1] = Int64(floor(state[1]) + 1); loc[2] = Int64(floor(state[2]) + 1); end;
function is_wall(env::FourRoomsCont, state::Array{Float64, 1})
    prj = project(env, state)
    return is_wall(env, prj)
end

function is_wall(env::FourRoomsCont, prj::Array{Int64, 1})
    if prj[1] < 1 || prj[2] < 1 || prj[1] > 11 || prj[2] > 11
        return false
    end
    env.walls[prj[1], prj[2]]
end

random_state(env::FourRoomsCont, rng) = [rand(rng)*size(env.walls)[1], rand(rng)*size(env.walls)[2]]
function random_start_state(env::FourRoomsCont, rng)
    state = random_state(env, rng)
    while is_wall(env, state)
        state = random_state(env, rng)
    end
    return state
end
Base.size(env::FourRoomsCont) = size(env.walls)
num_actions(env::FourRoomsCont) = 4
get_states(env::FourRoomsCont) = findall(x->x==false, env.walls)
RLCore.get_actions(env::FourRoomsCont) = FourRoomsContParams.ACTIONS

function handle_collision(env::FourRoomsCont, state, action)

    #Approximate by the square...
    frp = FourRoomsContParams
    new_state = copy(state)
    # env.new_state .= state
    # new_state = copy(env.new_state)

    new_state[1] = clamp(new_state[1], frp.BODY_RADIUS, size(env.walls)[1] - frp.BODY_RADIUS)
    new_state[2] = clamp(new_state[2], frp.BODY_RADIUS, size(env.walls)[2] - frp.BODY_RADIUS)

    # Really basic collision detection for 2-d plane worlds.

    collided = new_state[1] != state[1] || new_state[2] != state[2]

    # edge_locs =
    #     [project(env::FourRoomsCont, new_state .+ [frp.BODY_RADIUS, 0.0]),
    #      project(env::FourRoomsCont, new_state .+ [-frp.BODY_RADIUS, 0.0]),
    #      project(env::FourRoomsCont, new_state .+ [0.0, frp.BODY_RADIUS]),
    #      project(env::FourRoomsCont, new_state .+ [0.0, -frp.BODY_RADIUS])]
    # env.collision_check .= [is_wall(env, edge_locs[1]),
    #                        is_wall(env, edge_locs[2]),
    #                        is_wall(env, edge_locs[3]),
    #                        is_wall(env, edge_locs[4])]
    for i in 1:4
        project(env::FourRoomsCont, new_state .+ frp.AGENT_BOUNDRIES[i], env.edge_locs[i])
        env.collision_check[i] = is_wall(env, env.edge_locs[i])
    end

    collided = collided || any(env.collision_check)

    if env.collision_check[1] && env.collision_check[2]
        wall_piece = project(env, new_state)
        if action == frp.UP
            new_state[1] = (wall_piece[1]) + frp.BODY_RADIUS
        elseif action == frp.DOWN
            new_state[1] = (wall_piece[1] - 1) - frp.BODY_RADIUS
        end
    elseif env.collision_check[1]
        new_state[1] = (env.edge_locs[1][1] - 1) - frp.BODY_RADIUS
    elseif env.collision_check[2]
        new_state[1] = (env.edge_locs[2][1]) + frp.BODY_RADIUS
    end

    if env.collision_check[3] && env.collision_check[4]
        wall_piece = project(env, new_state)
        if action == frp.LEFT
            new_state[2] = (wall_piece[2]) + frp.BODY_RADIUS
        elseif action == frp.RIGHT
            new_state[2] = (wall_piece[2] - 1) - frp.BODY_RADIUS
        end
    elseif env.collision_check[3]
        new_state[2] = (env.edge_locs[3][2] - 1) - frp.BODY_RADIUS
    elseif env.collision_check[4]
        new_state[2] = (env.edge_locs[4][2]) + frp.BODY_RADIUS
    end

    new_state[1] = clamp(new_state[1], frp.BODY_RADIUS, size(env.walls)[1] - frp.BODY_RADIUS)
    new_state[2] = clamp(new_state[2], frp.BODY_RADIUS, size(env.walls)[2] - frp.BODY_RADIUS)

    return new_state, collided
end

function handle_collision!(env::FourRoomsCont, action)
    frp = FourRoomsContParams
    env.state, collided = handle_collision(env, env.state, action)
    return collided
end

function which_room(env::FourRoomsCont, state)
    frp = FourRoomsContParams
    room = -1
    if state[1] < 6
        # LEFT
        if state[2] < 6
            # TOP
            room = frp.ROOM_TOP_LEFT
        else
            # Bottom
            room = frp.ROOM_BOTTOM_LEFT
        end
    else
        # RIGHT
        if state[2] < 7
            # TOP
            room = frp.ROOM_TOP_RIGHT
        else
            # Bottom
            room = frp.ROOM_BOTTOM_RIGHT
        end
    end
    return room
end

function RLCore.reset!(env::FourRoomsCont; rng=Random.GLOBAL_RNG, kwargs...)
    state = random_state(env, rng)

    while is_wall(env, state)
        state = random_state(env, rng)
    end
    env.state = state
    return state
end

function mini_step!(env::FourRoomsCont, step, action)

    env.state .+= step
    collision = handle_collision!(env, action)
    return collision
end

function mini_step(env::FourRoomsCont, state, step, action)

    new_state = state .+ step
    new_state, collision = handle_collision(env, new_state, action)
    return new_state, collision
end

function RLCore.environment_step!(env::FourRoomsCont, action; rng=Random.GLOBAL_RNG, kwargs...)

    frp = FourRoomsContParams
    next_step = zeros(2)

    if action == frp.UP
        next_step[1] = -(frp.STEP - rand(rng)*env.max_action_noise)
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == frp.DOWN
        next_step[1] = frp.STEP - rand(rng)*env.max_action_noise
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == frp.RIGHT
        next_step[2] = frp.STEP - rand(rng)*env.max_action_noise
        next_step[1] = randn(rng)*env.drift_noise
    elseif action == frp.LEFT
        next_step[2] = -(frp.STEP - rand(rng)*env.max_action_noise)
        next_step[1] = randn(rng)*env.drift_noise
    else
        throw("Wrong action")
    end

    # mini_physics simulation for 1 second (== 10 steps of 0.1 seconds)
    Δt=1.0
    τ=10
    next_step .*= Δt/τ
    collision = false
    for t in 1:τ
        collision = mini_step!(env, next_step, action)
        if collision
            break
        end
    end
    env.collision = collision
end

function _step(env::FourRoomsCont, state, action; rng=Random.GLOBAL_RNG, kwargs...)

    frp = FourRoomsContParams
    next_step = zeros(2)

    if action == frp.UP
        next_step[1] = -(frp.STEP - rand(rng)*env.max_action_noise)
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == frp.DOWN
        next_step[1] = frp.STEP - rand(rng)*env.max_action_noise
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == frp.RIGHT
        next_step[2] = frp.STEP - rand(rng)*env.max_action_noise
        next_step[1] = randn(rng)*env.drift_noise
    elseif action == frp.LEFT
        next_step[2] = -(frp.STEP - rand(rng)*env.max_action_noise)
        next_step[1] = randn(rng)*env.drift_noise
    else
        throw("Wrong action")
    end

    Δt=1.0
    τ=10
    next_step .*= Δt/τ
    collision = false
    next_state = state[1]
    for t in 1:τ
        next_state, collision = mini_step(env, next_state, next_step, action)
        if collision
            break
        end
    end
    # next_state = handle_collision(env, next_state, action)
    return (next_state, collision), 0.0, false
end

function MonteCarloReturn(env::FourRoomsCont, gvf::GVF, start_state::Array{Float64, 1},
                          num_returns::Int64, γ_thresh::Float64=1e-6,
                          max_steps::Int64=Int64(1e7);
                          rng=Random.GLOBAL_RNG)

    returns = zeros(num_returns)
    for ret in 1:num_returns
        step = 0
        cur_state = start_state
        cumulative_gamma = 1.0
        while cumulative_gamma > γ_thresh && step < max_steps
            action = StatsBase.sample(rng, gvf.policy, cur_state)
            next_state, _, _ = _step(env, cur_state, action; rng=rng)
            c, γ, pi_prob = get(gvf, cur_state, action, next_state, nothing, nothing)
            returns[ret] += cumulative_gamma*c
            cumulative_gamma *= γ
            cur_state = next_state
            step += 1
        end
    end

    return returns
end

import ProgressMeter

function get_sequence(env::FourRoomsCont, num_steps, policy; seed=0, normalized=false)
    rng = Random.MersenneTwister(seed)
    env.normalized = normalized
    states = Array{Array{Float64, 1}}(undef, num_steps+1)
    _, state = start!(env; rng=rng)
    states[1] = copy(state[1])
    ProgressMeter.@showprogress 0.1 "Step:" for step in 1:num_steps
        action = StatsBase.sample(rng, policy, state)
        _, state, _, _ = step!(env, action; rng=rng)
        states[step+1] = copy(state[1])
    end
    return states
end


text:
"""
    lis(arr::Array{Int}, ::Val{:bs})

# Arguments:
- `arr`: sequence of integers

# Examples/Tests
```julia
julia> lis([3, 10, 2, 1, 20], Val(:bs))
[3, 10, 20]

julia> lis([2, 2, 3], Val(:bs))
[2, 3]
```

# Brief:
Given a sequence of integers, the function finds the longest, strictly increasing, subsequence in that sequence.

# For more information:
https://cp-algorithms.com/sequences/longest_increasing_subsequence.html

# Contributed by:
- [Igor Malheiros](https://github.com/igormalheiros)
"""

function lis(arr::Array{Int}, ::Val{:bs})
    len = length(arr)
    memo = ones(Int, len)
    p = ones(Int, len)

    lis_arr = Int[]

    len == 0 && return lis_arr # if `arr` is empty

    lis_len = 1
    memo[1] = arr[1]

    for i = 2:len
        # binary search in current `memo`, it uses @view and slicing for in-place verification
        p[i] = searchsortedfirst((@view memo[1:lis_len]), arr[i])
        # if arr[i] is greater than the greatest element in `memo`, it increases `lis_len`
        lis_len = max(lis_len, p[i])
        # updates `memo`
        memo[p[i]] = arr[i]
    end

    # Restoring
    last_pos = lis_len
    for i = len:-1:1
        if p[i] == last_pos
            push!(lis_arr, arr[i])
            last_pos -= 1
        end
    end

    reverse!(lis_arr)

    return lis_arr
end

max_issues_repo_name: AugustoCL/Julia

max_forks_repo_name: AugustoCL/Julia


text:
# Practical Machine learning
# Association rule based learning - Apriori TEST
# Chapter 7

using Base.Test
include("apriori.jl")

function _gen_dummy_data!(transactions)
    range = [1:10]
    for i in 1:length(transactions)
        transactions[i] = sample(range, sample(range, 1)[1], replace = false)
    end
    transactions
end

minsupp = 0.5
minconf = 0.2
T = Array(Array{Int64,1},10)
T[1] = [1,2]
T[2] = [1,3,4,5]
T[3] = [2,3,4,6]
T[4] = [1,2,3,4]
T[5] = [1,2,3,6]
T[6] = [1,2,3,5]
T[7] = [1,2,3,5,6]
T[8] = [1,3,4,5,6]
T[9] = [1,2,3,4,5,6]
T[10] = [1,2,3,4,5]

@test_approx_eq σ([2,3,4],T[1:5]) 2
@test_approx_eq supp([2,3],4,T[1:5]) 0.4
@test_approx_eq_eps conf([2,3],4,T[1:5]) 0.67 1e-2

F₂ = Array(Array{Int64,1},4)
F₂[1] = [3,4]
F₂[2] = [1,3]
F₂[3] = [1,2]
F₂[4] = [2,3]
@test length(gen_candidate(F₂)) == 1
@test gen_candidate(F₂)[1] == [1,2,3]

F = find_freq_itemset(T[1:5], minsupp)
@test length(F) == 2
@test Set(F[1]) == Set([1],[2],[3],[4])
@test Set(F[2]) == Set([3,4],[1,3],[1,2],[2,3])


text:
typealias Point AbstractVector
typealias Vec AbstractVector
typealias Mat AbstractMatrix

## 2D inverse graphics using Approximate Bayesian Computation
## ==========================================================

# This example constructs a generative model which creates image from polygon
# The inference task is to infer a distribution over the polygons conditioned on
# onbserving an image

# http://stackoverflow.com/questions/217578/point-in-polygon-aka-hit-test
function point_in_poly(poly::Mat, testx::Float64, testy::Float64)
  println("Checking if point is in poly")
  print(poly)
  nvert = size(poly,2)
  vertx = poly[1,:]
  verty = poly[2,:]
  c = flip(0,0.0) # HACK
  j = nvert
  print(vertx)
  print(verty)
  print(testy)

  println("Setting up Lens")
  println("Going into Loop")
  for i = 1:nvert
    firsttest = !((verty[i]>testy) == (verty[j]>testy))
    secondtest = testx < ((vertx[j]-vertx[i]) * (testy-verty[i]) / (verty[j]-verty[i]) + vertx[i])
    conjunction = firsttest & secondtest
    c = ifelse(conjunction, !c, c)
    j = i
  end
  return c
end

# Render a polygon to a monochrome image
function render(poly::Mat, width::Int, height::Int)
  # If the point is in the polygon render it 1.0, otherwie as 0.0
  image = RandVar{Float64}[ifelse(point_in_poly(poly, float(x), float(y)),1.0,0.0)
           for x = 1:width, y = 1:height]
  RandArray(image)
end

# ABC comparison.
function img_compare(image::Mat, data::Mat)
  abs(image - data)
end

function abc(poly::Mat, observation::Mat)
  @show typeof(poly)
  @show poly
  width = size(observation)[1]
  height = size(observation)[2]
  rendering = render(poly,width,height)
  condition = sum(img_compare(rendering,observation)) < 5
  return poly, condition
end

## Example ABC
function test_abc()
  testimage = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
               0.0 1.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
               0.0 1.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0 0.0
               0.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 0.0
               0.0 1.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0
               0.0 1.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 0.0
               0.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0
               0.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0
               0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
               0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0]
  abc(mvuniform(-1,0,10,2,3),testimage)
end

function main_polyfrompixels()
  poly,condition = test_abc()
  ndims(condition)
  to_dimacs(condition)
  # rand(poly,condition,1)
end


text:
using NLS_Fit
using Documenter

DocMeta.setdocmeta!(NLS_Fit, :DocTestSetup, :(using NLS_Fit); recursive=true)
doctest(NLS_Fit, fix=true)


text:
using WordTokenizers

export NaiveBayesClassifier

simpleTokenise(s) = WordTokenizers.tokenize(lowercase(replace(s, "."=>"")))

"""
Create a dict that maps elements in input array to their frequencies.
"""
function frequencies(xs)
    frequencies = Dict{eltype(xs),Int}()
    for x in xs
        frequencies[x] = get(frequencies, x, 0) + 1
    end
    return frequencies
end

"""
    features(::AbstractDict, dict)

Compute an Array, mapping the value corresponding to elements of `dict` to the input `AbstractDict`.
"""
function features(fs::AbstractDict, dict)
    bag = zeros(Int, size(dict))
    for i = 1:length(dict)
        bag[i] = get(fs, dict[i], 0)
    end
    return bag
end

features(s::AbstractString, dict) = features(frequencies(simpleTokenise(s)), dict)

Features{T<:Integer} = AbstractVector{T}

mutable struct NaiveBayesClassifier{T}
    dict::Vector{String}
    classes::Vector{T}
    weights::Matrix{Int}
end

"""
    NaiveBayesClassifier([dict, ]classes)

A Naive Bayes Classifier for classifying documents.

# Example
```julia-repl
julia> using TextAnalysis: NaiveBayesClassifier, fit!, predict
julia> m = NaiveBayesClassifier([:spam, :non_spam])
NaiveBayesClassifier{Symbol}(String[], Symbol[:spam, :non_spam], Array{Int64}(0,2))

julia> fit!(m, "this is spam", :spam)
NaiveBayesClassifier{Symbol}(["this", "is", "spam"], Symbol[:spam, :non_spam], [2 1; 2 1; 2 1])

julia> fit!(m, "this is not spam", :non_spam)
NaiveBayesClassifier{Symbol}(["this", "is", "spam", "not"], Symbol[:spam, :non_spam], [2 2; 2 2; 2 2; 1 2])

julia> predict(m, "is this a spam")
Dict{Symbol,Float64} with 2 entries:
  :spam     => 0.59883
  :non_spam => 0.40117
```
"""
NaiveBayesClassifier(dict, classes) =
    NaiveBayesClassifier(dict, classes,
             ones(Int, length(dict), length(classes)))

NaiveBayesClassifier(classes) = NaiveBayesClassifier(String[], classes)

probabilities(c::NaiveBayesClassifier) = c.weights ./ sum(c.weights, dims = 1)

"""
    extend!(model::NaiveBayesClassifier, dictElement)

Add the dictElement to dictionary of the Classifier `model`.
"""
function extend!(c::NaiveBayesClassifier, dictElement)
    push!(c.dict, dictElement)
    c.weights = vcat(c.weights, ones(Int, length(c.classes))')
    return c
end

"""
    fit!(model::NaiveBayesClassifier, str, class)
    fit!(model::NaiveBayesClassifier, ::Features, class)
    fit!(model::NaiveBayesClassifier, ::StringDocument, class)

Fit the weights for the model on the input data.
"""
function fit!(c::NaiveBayesClassifier, x::Features, class)
    n = findfirst(==(class), c.classes)
    c.weights[:, n] .+= x
    return c
end

function fit!(c::NaiveBayesClassifier, sd::AbstractDocument, class)
    fs = frequencies(tokens(sd))
    for k in keys(fs)
        k in c.dict || extend!(c, k)
    end
    fit!(c, features(fs, c.dict), class)
end

fit!(c::NaiveBayesClassifier, s::String, class) = fit!(c, StringDocument(s), class)

"""
    predict(::NaiveBayesClassifier, str)
    predict(::NaiveBayesClassifier, ::Features)
    predict(::NaiveBayesClassifier, ::StringDocument)

Predict probabilities for each class on the input Features or String.
"""
function predict(c::NaiveBayesClassifier, x::Features)
    ps = prod(probabilities(c) .^ x, dims = 1)
    ps ./= sum(ps)
    Dict(c.classes[i] => ps[i] for i = 1:length(c.classes))
end

predict(c::NaiveBayesClassifier, s::String) =
    predict(c, features(s, c.dict))

predict(c::NaiveBayesClassifier, sd::AbstractDocument) = predict(c, features(frequencies(tokens(sd)), c.dict))


text:
"""
Wrappers for the family of sf_readf_* functions, which read the given number
of frames into the given array. Returns the number of frames read.
"""
function sf_readf end

sf_readf(filePtr, dest::Array{T}, nframes) where T <: Union{Int16, PCM16Sample} =
ccall((:sf_readf_short, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{T}, Int64),
      filePtr, dest, nframes)

sf_readf(filePtr, dest::Array{T}, nframes) where T <: Union{Int32, PCM32Sample} =
ccall((:sf_readf_int, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{T}, Int64),
      filePtr, dest, nframes)

sf_readf(filePtr, dest::Array{Float32}, nframes) =
ccall((:sf_readf_float, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{Float32}, Int64),
      filePtr, dest, nframes)

sf_readf(filePtr, dest::Array{Float64}, nframes) =
ccall((:sf_readf_double, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{Float64}, Int64),
      filePtr, dest, nframes)

function SampledSignals.unsafe_read!(source::SndFileSource, buf::Array, frameoffset, framecount)
    total = min(framecount, nframes(source) - source.pos + 1)
    nread = 0
    readbuf = source.readbuf
    while nread < total
        n = min(size(readbuf, 2), total - nread)
        # transpose! needs the ranges to all use Ints, which on 32-bit systems
        # is an Int32, but sf_writef returns Int64 on both platforms, so we
        # convert to a platform-native Int. This also avoids a
        # type-inferrability problem where `nw` would otherwise change type.
        nr::Int = sf_readf(source.filePtr, readbuf, n)
        # the data comes in interleaved, so we need to transpose
        transpose!(view(buf, (1:nr) .+ frameoffset .+ nread, :),
                   view(readbuf, :, 1:nr))
        source.pos += nr
        nread += nr
        nr == n || break
    end

    nread
end

"""
Wrappers for the family of sf_writef_* functions, which write the given number
of frames in the source array to the file. Returns the number of frames written.
"""
function sf_writef end

sf_writef(filePtr, src::Array{T}, nframes) where T <: Union{Int16, PCM16Sample} =
ccall((:sf_writef_short, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{T}, Int64),
      filePtr, src, nframes)

sf_writef(filePtr, src::Array{T}, nframes) where T <: Union{Int32, PCM32Sample} =
ccall((:sf_writef_int, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{T}, Int64),
      filePtr, src, nframes)

sf_writef(filePtr, src::Array{Float32}, nframes) =
ccall((:sf_writef_float, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{Float32}, Int64),
      filePtr, src, nframes)

sf_writef(filePtr, src::Array{Float64}, nframes) =
ccall((:sf_writef_double, libsndfile), Int64,
      (Ptr{Cvoid}, Ptr{Float64}, Int64),
      filePtr, src, nframes)

# returns the number of samples written
function SampledSignals.unsafe_write(sink::SndFileSink, buf::Array, frameoffset, framecount)
    nwritten = 0
    writebuf = sink.writebuf
    while nwritten < framecount
        n = min(size(writebuf, 2), framecount - nwritten)
        # the data needs to be interleaved, so we need to transpose
        transpose!(view(writebuf, :, 1:n),
                   view(buf, (1:n) .+ frameoffset .+ nwritten, :))
        # transpose! needs the ranges to all use Ints, which on 32-bit systems
        # is an Int32, but sf_writef returns Int64 on both platforms, so we
        # convert to a platform-native Int. This also avoids a
        # type-inferrability problem where `nw` would otherwise change type.
        nw::Int = sf_writef(sink.filePtr, writebuf, n)
        sink.nframes += nw
        nwritten += nw
        nw == n || break
    end

    nwritten
end


text:
@inline _alloc_temp_array(size_1::Int, x::AbstractVector{T}) where T = zeros(T, size_1)
@inline _alloc_temp_array(size_1::Int, x::AbstractMatrix{T}) where T = zeros(T, size_1, size(x, 2))

const MatrixOrFactorization{T} = Union{AbstractMatrix{T}, Factorization{T}}

for N in (1, 2)
    for op in [:mul, :ldiv]
        square_kernel! = Symbol("_kron_", op, "_kernel_square!")

        op! = Symbol(op, "!")
        if N == 1
            op_expr_square = :(@views q[slc] = $op!(temp[1:n], m, q[slc]))
        else
            op_expr_square = :(@views q[slc, :] = $op!(temp[1:n, :], m, q[slc, :]))
        end

        @eval function $square_kernel!(temp::AbstractArray{T1, $N}, q::AbstractArray{T2, $N}, n::Int, i_left::Int, m::MatrixOrFactorization{T3}, i_right::Int) where {T1,T2,T3}
            # apply kron(I(l), m, I(r)) where m is square to the given vector x, overwriting x in the process

            if m isa Factorization || m isa Adjoint{T3, <:Factorization} || m isa Transpose{T3, <:Factorization} || m != I
                # if matrix is the identity, skip matmul/div, unless it's some sort factorization, then proceed anyway
                irc = i_right * n

                base_i, top_i = 0, (irc - i_right)
                @inbounds for i_l in 1:i_left
                    for i_r in 1:i_right
                        slc = base_i + i_r : i_right : top_i + i_r

                        $op_expr_square
                    end

                    base_i += irc
                    top_i += irc
                end
            end
            return q
        end

        square_func! = Symbol("_kron_", op, "_fast_square!")
        @eval function $square_func!(out::AbstractArray{T1, $N}, x::AbstractArray{T2, $N}, matrices) where {T1,T2}
            ns::Vector{Int} = [size(m, 1) for m in matrices]
            i_left::Int = 1
            i_right::Int = prod(ns)

            out = copyto!(out, x)
            temp = _alloc_temp_array(maximum(ns), x)

            for s in 1:length(ns)
                n = ns[s]
                i_right ÷= n
                $square_kernel!(temp, out, n, i_left, matrices[s], i_right)
                i_left *= n
            end
            return out
        end

        # NOTE: ldiv! seems to give erroneous results if m is not square,
        #       using allocating ldiv for now
        if N == 1
            if op == :mul
                op_expr_rect = :(@views q′[slc_out] = mul!(temp[1:r_h], m, q[slc_in]))
            else
                op_expr_rect = :(@views q′[slc_out] = m \ q[slc_in])
            end
        else
            if op == :mul
                op_expr_rect = :(@views q′[slc_out, :] = mul!(temp[1:r_h, :], m, q[slc_in, :]))
            else
                op_expr_rect = :(@views q′[slc_out, :] = m \ q[slc_in, :])
            end
        end

        rect_kernel! = Symbol("_kron_", op, "_kernel_rect!")
        @eval function $rect_kernel!(temp::AbstractArray{T1, $N}, q::AbstractArray{T2, $N}, r_h::Int, c_h::Int, i_left::Int, m::MatrixOrFactorization{T3}, i_right::Int) where {T1,T2,T3}
            # apply kron(I(i_left), m, I(i_right)) to the given vector q

            # don't bother checking for identity, since we know the matrix
            #  is rectangular here
            irc = i_right * c_h
            irr = i_right * r_h

            size_ = i_left * irr
            q′ = _alloc_temp_array(size_, q)

            base_i, base_j = 0, 0
            top_i, top_j = (irc - i_right), (irr - i_right)
            @inbounds for i_l in 1:i_left
                for i_r in 1:i_right
                    slc_in  = base_i + i_r : i_right : i_r + top_i
                    slc_out = base_j + i_r : i_right : i_r + top_j

                    $op_expr_rect
                end

                base_i += irc
                top_i += irc

                base_j += irr
                top_j += irr
            end

            return q′
        end

        rect_func! = Symbol("_kron_", op, "_fast_rect!")
        ri = (op == :mul) ? 1 : 2
        ci = (op == :mul) ? 2 : 1
        @eval function $rect_func!(out::AbstractArray{T1, $N}, x::AbstractArray{T2, $N}, matrices) where {T1,T2}
            r::Vector{Int} = [size(m, $ri) for m in matrices]
            c::Vector{Int} = [size(m, $ci) for m in matrices]
            i_left::Int = 1
            i_right::Int = prod(c)

            out_ = copy(x)
            temp = _alloc_temp_array(maximum(r), x)

            for h in 1:length(matrices)
                r_h, c_h = r[h], c[h]
                i_right ÷= c_h
                if r_h == c_h
                    $square_kernel!(temp, out_, r_h, i_left, matrices[h], i_right)
                else
                    out_ = $rect_kernel!(temp, out_, r_h, c_h, i_left, matrices[h], i_right)
                end
                i_left *= r_h
            end

            out = copyto!(out, out_)
            return out
        end
    end

    @eval function _kronsum_mul_fast!(out::AbstractArray{T1, $N}, x::AbstractArray{T2, $N}, matrices) where {T1,T2}
        ns::Vector{Int} = [size(m, 1) for m in matrices]
        i_left::Int = 1
        i_right::Int = prod(ns)

        out = fill!(out, zero(T1))
        temp = copy(x)
        # should use similar instead, but there seems to be a bug when using copy! with empty SparseArrays
        small_temp = _alloc_temp_array(maximum(ns), x)

        # this loop is technically parallelizable,
        #  though that'd end up using more memory
        for s in 1:length(ns)
            n = ns[s]
            i_right ÷= n
            copyto!(temp, x)
            out += _kron_mul_kernel_square!(small_temp, temp, n, i_left, matrices[s], i_right)
            i_left *= n
        end
        return out
    end
end