text:
# Model modified from
# https://discourse.julialang.org/t/how-to-drop-the-dropout-layers-in-flux-jl-when-assessing-model-performance/19924

using Flux, Statistics, Random, Test
# Uncomment to use MNIST data
# using MLDatasets: MNIST
using StableRNGs
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated, partition
using Legolas, LegolasFlux

# This should store all the information needed
# to construct the model.
Base.@kwdef struct DigitsConfig
    seed::Int = 5
    dropout_rate::Float32 = 0f1
end

# Here's our model object itself, just a `DigitsConfig` and
# a `chain`. We keep the config around so it's easy to save out
# later.
struct DigitsModel
    chain::Chain
    config::DigitsConfig
end

# Ensure Flux can recurse into our model to find params etc
Flux.@functor DigitsModel (chain,)

# Construct the actual model from a config object. This is the only
# constructor that should be used, to ensure the model is created just
# from the config object alone.
function DigitsModel(config::DigitsConfig=DigitsConfig())
    dropout_rate = config.dropout_rate
    Random.seed!(config.seed)
    chain = Chain(Dropout(dropout_rate),
                  Conv((3, 3), 1 => 32, relu),
                  BatchNorm(32, relu),
                  MaxPool((2, 2)),
                  Dropout(dropout_rate),
                  Conv((3, 3), 32 => 16, relu),
                  Dropout(dropout_rate),
                  MaxPool((2, 2)),
                  Dropout(dropout_rate),
                  Conv((3, 3), 16 => 10, relu),
                  Dropout(dropout_rate),
                  x -> reshape(x, :, size(x, 4)),
                  Dropout(dropout_rate),
                  Dense(90, 10),
                  softmax)
    return DigitsModel(chain, config)
end

# Our model acts on input just by applying the chain.
(m::DigitsModel)(x) = m.chain(x)

# Here, we define a schema extension of the `legolas-flux.model` schema.
# We add our `DigitsConfig` object, as well as the epoch and accuracy.
const DigitsRow = Legolas.@row("digits.model@1" > "legolas-flux.model@1",
                               config::DigitsConfig,
                               epoch::Union{Missing, Int},
                               accuracy::Union{Missing, Float32})

# Construct a `DigitsRow` from a model by collecting the weights.
# This can then be saved with e.g. `LegolasFlux.write_model_row`.
function DigitsRow(model::DigitsModel; epoch=missing, accuracy=missing)
    return DigitsRow(; weights=fetch_weights(model), model.config, epoch, accuracy)
end

# Construct a `DigitsModel` from a row satisfying the `DigitsRow` schema,
# i.e. one with a `weights` and `config::DigitsConfig`.
# This could be the result of `LegolasFlux.read_model_row`.
function DigitsModel(row)
    m = DigitsModel(row.config)
    load_weights!(m, row.weights)
    return m
end

# Increase to get more training/test data
N_train = 1_000
N_test = 50

##
# to use MNIST data, uncomment these
# train_x, train_y = MNIST.traindata(Float32, 1:N_train)
# test_x,  test_y  = MNIST.testdata(Float32, 1:N_test)

# Random data:
rng = StableRNG(735)
train_x = rand(rng, Float32, 28, 28, N_train)
train_y = rand(rng, 0:9, N_train)
test_x = rand(rng, Float32, 28, 28, N_test)
test_y = rand(rng, 0:9, N_test)
##

# Partition into batches of size 32
batch_size = 32
train = [(reshape(train_x[:, :, I], 28, 28, 1, :), onehotbatch(train_y[I], 0:9))
         for I in partition(1:N_train, batch_size)]

tX = reshape(test_x, 28, 28, 1, :)
tY = onehotbatch(test_y, 0:9)

function accuracy(m, x, y)
    testmode!(m)
    val = mean(onecold(m(x)) .== onecold(y))
    trainmode!(m)
    return val
end

function train_model!(m; N = N_train)
    loss = (x, y) -> crossentropy(m(x), y)
    opt = ADAM()
    evalcb = throttle(() -> @show(accuracy(m, tX, tY)), 5)
    Flux.@epochs 1 Flux.train!(loss, params(m), Iterators.take(train, N), opt; cb=evalcb)
    return accuracy(m, tX, tY)
end

m = DigitsModel()

# increase N to actually train more than a tiny amount
acc = train_model!(m; N=10)

# Let's serialize out the weights into a `DigitsRow`.
# We could save this here with `write_model_row`.
row = DigitsRow(m; epoch=1, accuracy=acc)

testmode!(m)
input = tX[:, :, :, 1:1]
output = m(input)
label = tY[:, 1]

# Let's now reconstruct the model from the `row` and check that we get
# the same outputs.
m2 = DigitsModel(row)
testmode!(m2)
output2 = m2(input)

@test output ≈ output2


text:
using LightXML
using Codecs
using Compat

struct VTKNode
  coords::Vector{Float64}
end

struct VTKElement
  vertices::Vector{Int}
  vtknum::Int
end

abstract type AbstractVTKXML end
#abstract type AbstractVTKXML end

abstract type AbstractVTKXMLBinary <: AbstractVTKXML end
#abstract type AbstractVTKXMLBinary <: AbstractVTKXML end

function add_data!(vtkxml::T, data) where T <: AbstractVTKXMLBinary
    write(vtkxml.buffer, data)
end

struct VTKXMLBinaryCompressed <: AbstractVTKXMLBinary
    buffer::IOBuffer
end
VTKXMLBinaryCompressed() = VTKXMLBinaryCompressed(IOBuffer())

function write_data!(vtkxml::VTKXMLBinaryCompressed, xmlele::XMLElement)
    uncompressed_size = vtkxml.buffer.size
    compressed_data = encode(Zlib, takebuf_array(vtkxml.buffer))
    compressed_size = length(compressed_data)
    header = UInt32[1, uncompressed_size, uncompressed_size, compressed_size]
    header_binary = bytestring(encode(Base64, reinterpret(UInt8, header)))
    data_binary = bytestring(encode(Base64, compressed_data))
    add_text(xmlele, header_binary)
    add_text(xmlele, data_binary)
end

struct VTKXMLBinaryUncompressed <: AbstractVTKXMLBinary
    buffer::IOBuffer
end
VTKXMLBinaryUncompressed() = VTKXMLBinaryUncompressed(IOBuffer())

function write_data!(vtkxml::VTKXMLBinaryUncompressed, xmlele::XMLElement)
    uncompressed_size = vtkxml.buffer.size
    uncompressed_data = takebuf_array(vtkxml.buffer)
    header = UInt32[uncompressed_size]
    header_binary = bytestring(encode(Base64, reinterpret(UInt8, header)))
    data_binary = bytestring(encode(Base64, uncompressed_data))
    add_text(xmlele, header_binary)
    add_text(xmlele, data_binary)
end

struct VTKXMLASCII <: AbstractVTKXML
    buffer::IOBuffer
end
VTKXMLASCII() = VTKXMLASCII(IOBuffer())

function add_data!(vtkxml::VTKXMLASCII, data)
    print(vtkxml.buffer, data, " ")
end

function add_data!(vtkxml::VTKXMLASCII, data::T) where T <: AbstractArray
    for comp in data
        print(vtkxml.buffer, comp, " ")
    end
end

function write_data!(vtkxml::VTKXMLASCII, xmlele::XMLElement)
    add_text(xmlele, takebuf_string(vtkxml.buffer))
end

function write_VTKXML(filename::String, nodes::Vector{VTKNode},
                             fin_els::Vector{VTKElement},
                             binary::Bool=true, compress::Bool=false)
     if (!binary) && compress
        error("Can only compress when using Binary format")
    end

    if binary
        if compress
            write_VTKXML1(filename, nodes, fin_els, binary, compress, VTKXMLBinaryCompressed())
        else
            write_VTKXML1(filename, nodes, fin_els, binary, compress, VTKXMLBinaryUncompressed())
        end
    else
        write_VTKXML1(filename, nodes, fin_els, binary, compress, VTKXMLASCII())
    end
end

function write_VTKXML1(filename::String, nodes::Vector{VTKNode},
                       fin_els::Vector{VTKElement}, binary::Bool,
                       compress::Bool, vtkxml::P) where P <: AbstractVTKXML

    if binary
        VTK_FORMAT = "binary"
    else
        VTK_FORMAT = "ascii"
    end

    xdoc = XMLDocument()
    xroot = create_root(xdoc, "VTKFile")
    set_attribute(xroot, "type", "UnstructuredGrid")
    set_attribute(xroot, "version", "0.1")
    set_attribute(xroot, "byte_order", "LittleEndian")

    if compress
        set_attribute(xroot, "compressor", "vtkZLibDataCompressor")
    end

    xgrid = new_child(xroot, "UnstructuredGrid")

    xpiece = new_child(xgrid, "Piece")
    set_attribute(xpiece, "NumberOfPoints", length(nodes))

    ncells = length(fin_els)
    set_attribute(xpiece, "NumberOfCells", length(fin_els))

    # Points
    xpoints = new_child(xpiece, "Points")

    # Coordinates for points
    xcoords = new_child(xpoints, "DataArray")
    set_attribute(xcoords, "type", "Float64")
    set_attribute(xcoords, "name", "Points")
    set_attribute(xcoords, "format", VTK_FORMAT)
    set_attribute(xcoords, "NumberOfComponents", "3")

    for node in nodes
        for coord in node.coords
            add_data!(vtkxml, coord)
        end
    end
    write_data!(vtkxml, xcoords)

    # Cells
    xcells = new_child(xpiece, "Cells")

    # Cell location
    xcellconn = new_child(xcells, "DataArray")
    set_attribute(xcellconn, "type", "Int")
    set_attribute(xcellconn, "Name", "connectivity")
    set_attribute(xcellconn, "format", VTK_FORMAT)

    for fin_el in fin_els
        for vertex in fin_el.vertices
            add_data!(vtkxml, vertex-1)
        end
    end
    write_data!(vtkxml, xcellconn)

    xcell_offsets = new_child(xcells, "DataArray")
    set_attribute(xcell_offsets, "type", "Int")
    set_attribute(xcell_offsets, "Name", "offsets")
    set_attribute(xcell_offsets, "format", VTK_FORMAT)
    offset = 0
    for fin_el in fin_els
        offset += length(fin_el.vertices)
        add_data!(vtkxml, offset)
    end
    write_data!(vtkxml, xcell_offsets)

    xcell_types = new_child(xcells, "DataArray")

    set_attribute(xcell_types, "type", "UInt8")
    set_attribute(xcell_types, "Name", "types")
    set_attribute(xcell_types, "format", VTK_FORMAT)
    for fin_el in fin_els
        add_data!(vtkxml, UInt8(fin_el.vtknum))
    end
    write_data!(vtkxml, xcell_types)

    save_file(xdoc, filename)
end

ext: jl


text:
function bipartite_matching_test()
    W = sprand(10,8,0.5)
    bipartite_matching(W)
    bipartite_matching([10;12;13],[1;2;3],[3;2;4])

    A = [5 1 1 1 1;
     1 5 1 1 1;
     1 1 5 1 1;
     1 1 1 5 1;
     1 1 1 1 5;
     1 1 1 1 1]
    A = sparse(A)
    M1 = bipartite_matching(A)
    (m1,m2) = edge_list(M1)

    if M1.weight != 25 || !isequal(m1,m2)
        error("bipartite matching failed")
    end

    A = ones(Int64,6,5) - eye(Int64,6,5)
    A = sparse(A')
    (ei,ej,ev) = csr_to_sparse(A.colptr,A.rowval,A.nzval)
    ai = [1;2;3;4;5;ei]
    aj = [1;2;3;4;5;ej]
    av = [5;5;5;5;5;ev]

    M2 = bipartite_matching(av,ai,aj)
    mi = MatrixNetworks.edge_indicator(M2,ai,aj)
    mitrue = zeros(Int64,length(av))
    mitrue[1:5] = 1

    if M2.weight != 25 || !isequal(m1,m2) || sum(mi) !=5 || !isequal(mi,mitrue)
        error("bipartite matching failed")
    end

    return true
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""PriorityLevelConfigurationStatus represents the current state of a \"request-priority\".

    IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus(;
        conditions=nothing,
    )

    - conditions::Vector{IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationCondition} : `conditions` is the current state of \"request-priority\".
"""
mutable struct IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus <: SwaggerModel
    conditions::Any # spec type: Union{ Nothing, Vector{IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationCondition} } # spec name: conditions

    function IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus(;conditions=nothing)
        o = new()
        validate_property(IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus, Symbol("conditions"), conditions)
        setfield!(o, Symbol("conditions"), conditions)
        o
    end
end # type IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus

const _property_map_IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus = Dict{Symbol,Symbol}(Symbol("conditions")=>Symbol("conditions"))
const _property_types_IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus = Dict{Symbol,String}(Symbol("conditions")=>"Vector{IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationCondition}")
Base.propertynames(::Type{ IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus }) = collect(keys(_property_map_IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus))
Swagger.property_type(::Type{ IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus[name]))}
Swagger.field_name(::Type{ IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus }, property_name::Symbol) =  _property_map_IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus[property_name]

function check_required(o::IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus)
    true
end

function validate_property(::Type{ IoK8sApiFlowcontrolV1alpha1PriorityLevelConfigurationStatus }, name::Symbol, val)
end


text:
# This file was generated, do not modify it. # hide
xgb = fitted_params(mtm).best_model
@show xgb.subsample
@show xgb.colsample_bytree


text:
export Mem

module Mem

import ..CLArrays: context
using OpenCL
import Base: pointer, eltype, cconvert, convert, unsafe_copy!
using GPUArrays: supports_double, device, global_memory
using GPUArrays

immutable OwnedPtr{T}
    ptr::cl.CL_mem
    mapped::Bool
    context::cl.Context
end
pointer(p::OwnedPtr) = p.ptr
eltype{T}(::Type{OwnedPtr{T}}) = T
cconvert(::Type{cl.CL_mem}, p::OwnedPtr) = pointer(p)
context(p::OwnedPtr) = p.context

function convert{T}(::Type{OwnedPtr{T}}, p::OwnedPtr)
    OwnedPtr{T}(pointer(p), p.mapped, context(p))
end

function retain(p::OwnedPtr)
    err = cl.api.clRetainMemObject(p)
    return err
end

const current_allocated_mem = Ref(0)
function free(p::OwnedPtr)
    cl.api.clReleaseMemObject(p)
end

function pressure_gc!(device, bytes)
    # If we used up 80% of our device. lets free see if we can free stuff with a gc swipe.
    # try three times, first only with a an unforced swipe
    mem = current_allocated_mem[] + bytes
    if (global_memory(device) * 0.8) < mem
        println("calling unforced gc")
        gc(false) # non force
    end
    if (global_memory(device) * 0.8) < mem
        println("force swipe")
        gc() # non forced
    end
    current_allocated_mem[] += bytes
end

function alloc(T, elems::Integer, ctx::cl.Context, flags = cl.CL_MEM_READ_WRITE)
    dev = GPUArrays.device(ctx)
    if T == Float64 && !supports_double(dev)
        error("Float64 is not supported by your device: $dev. Make sure to convert all types for the GPU to Float32")
    end
    nbytes = cl.cl_uint(elems * sizeof(T))
    pressure_gc!(dev, nbytes)
    err_code = Ref{cl.CL_int}()
    mem_id = cl.api.clCreateBuffer(
        ctx.id, flags, nbytes,
        C_NULL,
        err_code
    )
    if err_code[] != cl.CL_SUCCESS
        throw(cl.CLError(err_code[]))
    end
    OwnedPtr{T}(mem_id, false, ctx)
end

function cl.set_arg!(k::cl.Kernel, idx::Integer, arg::Mem.OwnedPtr)
    arg_boxed = Base.RefValue(pointer(arg))
    cl.@check cl.api.clSetKernelArg(k.id, cl.cl_uint(idx-1), sizeof(cl.CL_mem), arg_boxed)
    return k
end

function unsafe_copy!(
        queue, hostref::Ref, dev_offset::Integer, ptr::OwnedPtr, nbytes::Integer
    )
    n_evts  = UInt(0)
    evt_ids = C_NULL
    ret_evt = Ref{cl.CL_event}()
    cl.@check cl.api.clEnqueueReadBuffer(
        queue.id, ptr, cl.cl_bool(true),
        dev_offset, nbytes, hostref,
        n_evts, evt_ids, ret_evt
    )
    return
end

function unsafe_copy!(
        queue, ptr::OwnedPtr, dev_offset::Integer, hostref::Ref, nbytes::Integer
    )
    n_evts  = UInt(0)
    evt_ids = C_NULL
    ret_evt = Ref{cl.CL_event}()
    cl.@check cl.api.clEnqueueWriteBuffer(
        queue.id, ptr, cl.cl_bool(true),
        dev_offset, nbytes, hostref,
        n_evts, evt_ids, ret_evt
    )
    return
end

export OwnedPtr

end

using .Mem


text:
using Pkg
Pkg.activate(@__DIR__)
using Soss
using ZigZagBoomerang

include("zigzag.jl")

m = @model x begin
    α ~ Normal()
    β ~ Normal()
    yhat = α .+ β .* x
    y ~ For(eachindex(x)) do j
        Normal(yhat[j], 2.0)
    end
end

x = randn(20);
obs = -0.1 .+ 2x + 1randn(20);
posterior = m(x=x) | (y=obs,)

T = 100.0
trace, final, (num, acc) = @time zigzag(posterior, T)

# trace is a continous object, discretize to obtain samples
ts, xs = ZigZagBoomerang.sep(discretize(trace, 0.1))

xs = xform(posterior).(xs)

using Plots
p = plot(ts, getindex.(xs, 1))
plot!(ts, getindex.(xs, 2), color=:red)
p2 = plot(first.(xs), last.(xs))
p


text:
const braille_signs = ['⠁' '⠂' '⠄' '⡀';
                              '⠈' '⠐' '⠠' '⢀']

function tobraille(a::BitArray{2})
    res = UInt16(0x2800)
    for x in 1:2, y in 1:4
        a[y, x] && (res |= UInt16(braille_signs[x,y]))
    end
    return Char(res)
end

tobraille(a::AbstractArray) = tobraille(a .> 0)

function arr2braille(arr)
    dyarr, dxarr = size(arr)
    dx = ceil(Int, dxarr/2)
    dy = ceil(Int, dyarr/4)
    a = zeros(Int, dy*4, dx*2)
    a[1:dyarr, 1:dxarr] = arr
    res = fill(Char(0x2800), (dy, dx))
    for y in 1:dy, x in 1:dx
        y1 = 1+(y-1)*4
        x1 = 1+(x-1)*2
        res[y,x] = tobraille(a[y1:y1+3, x1:x1+1])
    end
    return res
end

@compat function update_board!(old::Array{Char}, arr::Array{Char})
    buf = IOBuffer()
    for i in findall(old .!= arr)
        if VERSION < v"0.7.0-DEV.3025"
            y,x = ind2sub(size(arr), i)
        else
            y,x = Tuple(i)
        end
        put(buf, [x, y], string(arr[y,x]))
    end
    print(String(take!(buf)))
end


text:
"""
    transformer_encoder(planes, depth, heads, headplanes, mlppanes; dropout = 0.)

Transformer as used in the base ViT architecture.
([reference](https://arxiv.org/abs/2010.11929)).

# Arguments
- `planes`: number of input channels
- `depth`: number of attention blocks
- `heads`: number of attention heads
- `headplanes`: number of hidden channels per head
- `mlppanes`: number of hidden channels in the MLP block
- `dropout`: dropout rate
"""
function transformer_encoder(planes, depth, heads, headplanes, mlpplanes; dropout = 0.)
  layers = [Chain(SkipConnection(prenorm(planes, MHAttention(planes, headplanes, heads; dropout)), +),
                  SkipConnection(prenorm(planes, mlp_block(planes, mlpplanes; dropout)), +))
            for _ in 1:depth]

  Chain(layers...)
end

"""
    vit(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = (16, 16), planes = 1024,
        depth = 6, heads = 16, mlppanes = 2048, headplanes = 64, dropout = 0.1, emb_dropout = 0.1,
        pool = :class, nclasses = 1000)

Creates a Vision Transformer (ViT) model.
([reference](https://arxiv.org/abs/2010.11929)).

# Arguments
- `imsize`: image size
- `inchannels`: number of input channels
- `patch_size`: size of the patches
- `planes`: the number of channels fed into the main model
- `depth`: number of blocks in the transformer
- `heads`: number of attention heads in the transformer
- `mlpplanes`: number of hidden channels in the MLP block in the transformer
- `headplanes`: number of hidden channels per head in the transformer
- `dropout`: dropout rate
- `emb_dropout`: dropout rate for the positional embedding layer
- `pool`: pooling type, either :class or :mean
- `nclasses`: number of classes in the output
"""
function vit(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = (16, 16), planes = 1024,
             depth = 6, heads = 16, mlppanes = 2048, headplanes = 64, dropout = 0.1, emb_dropout = 0.1,
             pool = :class, nclasses = 1000)

  im_height, im_width = imsize
  patch_height, patch_width = patch_size

  @assert (im_height % patch_height == 0) && (im_width % patch_width == 0)
  "Image dimensions must be divisible by the patch size."
  @assert pool in [:class, :mean]
  "Pool type must be either :class (class token) or :mean (mean pooling)"

  npatches = (im_height ÷ patch_height) * (im_width ÷ patch_width)
  patchplanes = inchannels * patch_height * patch_width

  return Chain(Chain(PatchEmbedding(patch_height, patch_width),
                     Dense(patchplanes, planes),
                     ClassTokens(planes),
                     ViPosEmbedding(planes, npatches + 1),
                     Dropout(emb_dropout),
                     transformer_encoder(planes, depth, heads, headplanes, mlppanes; dropout),
                     (pool == :class) ? x -> x[:, 1, :] : _seconddimmean),
               Chain(LayerNorm(planes), Dense(planes, nclasses)))
end

struct ViT
  layers
end

"""
    ViT(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = (16, 16), planes = 1024,
        depth = 6, heads = 16, mlppanes = 2048, headplanes = 64, dropout = 0.1, emb_dropout = 0.1,
        pool = :class, nclasses = 1000)

Creates a Vision Transformer (ViT) model.
([reference](https://arxiv.org/abs/2010.11929)).

# Arguments
- `imsize`: image size
- `inchannels`: number of input channels
- `patch_size`: size of the patches
- `planes`: the number of channels fed into the main model
- `depth`: number of blocks in the transformer
- `heads`: number of attention heads in the transformer
- `mlpplanes`: number of hidden channels in the MLP block in the transformer
- `headplanes`: number of hidden channels per head in the transformer
- `dropout`: dropout rate
- `emb_dropout`: dropout rate for the positional embedding layer
- `pool`: pooling type, either :class or :mean
- `nclasses`: number of classes in the output
"""
function ViT(imsize::NTuple{2} = (256, 256); inchannels = 3, patch_size = (16, 16), planes = 1024,
            depth = 6, heads = 16, mlppanes = 2048, headplanes = 64,
            dropout = 0.1, emb_dropout = 0.1, pool = :class, nclasses = 1000)

  layers = vit(imsize; inchannels, patch_size, planes, depth, heads, mlppanes, headplanes,
               dropout, emb_dropout, pool, nclasses)

  ViT(layers)
end

(m::ViT)(x) = m.layers(x)

backbone(m::ViT) = m.layers[1]
classifier(m::ViT) = m.layers[2]

@functor ViT


text:
"""
    rts(
        f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)},
        mask::AbstractArray{Bool, 3},
        vsz::NTuple{3, Real};
        pad::NTuple{3, Integer} = (0, 0, 0),
        Dkernel::Symbol = :k,
        bdir::NTuple{3, Real} = (0, 0, 1),
        lstol::Integer = 4,
        delta::Real = 0.15,
        mu::Real = 1e5,
        rho::Real = 10,
        tol::Real = 1e-2,
        maxit::Integer = 20,
        verbose::Bool = false,
    ) -> typeof(similar(f))

Rapid two-step dipole inversion with sparsity priors [1].

### Arguments
- `f::AbstractArray{<:AbstractFloat, N ∈ (3, 4)}`: unwrapped (multi-echo) local
    field/phase
- `mask::AbstractArray{Bool, 3}`: binary mask of region of interest
- `vsz::NTuple{3, Real}`: voxel size

### Keywords
- `pad::NTuple{3, Integer} = (0, 0, 0)`: zero padding array
    - `< 0`: no padding
    - `≥ 0`: minimum padding to fast fft size
- `bdir::NTuple{3, Real} = (0, 0, 1)`: unit vector of B field direction
- `Dkernel::Symbol = :k`: dipole kernel method
- `lstol::Integer = 4`: stopping tolerance (# of iterations) for lsmr solver
- `delta::Real = 0.15`: threshold for ill-conditioned k-space region
- `mu::Real = 1e5`: regularization parameter for tv minimization
- `rho::Real = 10`: Lagrange multiplier penalty parameter
- `tol::Real = 1e-2`: stopping tolerance
- `maxit::Integer = 20`: maximum number of iterations
- `verbose::Bool = false`: print convergence information

### Returns
- `typeof(similar(f))`: susceptibility map

### References
[1] Kames C, Wiggermann V, Rauscher A. Rapid two-step dipole inversion for
    susceptibility mapping with sparsity priors.
    Neuroimage. 2018 Feb 15;167:276-83.
"""
function rts(
    f::AbstractArray{<:AbstractFloat, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real};
    pad::NTuple{3, Integer} = (0, 0, 0),
    Dkernel::Symbol = :k,
    bdir::NTuple{3, Real} = (0, 0, 1),
    lstol::Integer = 4,
    delta::Real = 0.15,
    mu::Real = 1e5,
    rho::Real = 10,
    tol::Real = 1e-2,
    maxit::Integer = 20,
    verbose::Bool = false,
    tau::Real = 0.7,
    gamma::Real = 5,
) where {N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))
    return _rts!(
        tzero(f), f, mask, vsz, pad, Dkernel, bdir, lstol, delta, mu, rho,
        tol, maxit, verbose, tau, gamma
    )
end

function _rts!(
    x::AbstractArray{<:AbstractFloat, N},
    f::AbstractArray{T, N},
    mask::AbstractArray{Bool, 3},
    vsz::NTuple{3, Real},
    pad::NTuple{3, Integer},
    Dkernel::Symbol,
    bdir::NTuple{3, Real},
    lstol::Integer,
    delta::Real,
    mu::Real,
    rho::Real,
    tol::Real,
    maxit::Integer,
    verbose::Bool,
    tau::Real,
    gamma::Real,
) where {T, N}
    N ∈ (3, 4) || throw(ArgumentError("arrays must be 3d or 4d, got $(N)d"))

    size(x) == size(f) || throw(DimensionMismatch())
    size(mask) == size(f)[1:3] || throw(DimensionMismatch())

    Dkernel ∈ (:k, :kspace, :i, :ispace) ||
        throw(ArgumentError("Dkernel must be one of :k, :kspace, :i, :ispace, got :$(Dkernel)"))

    # convert scalars
    _zero = zero(T)

    ρ = convert(T, rho)
    δ = convert(T, delta)
    μ = convert(T, mu)
    γ = convert(T, gamma)
    τ = convert(T, tau)
    ϵ = convert(T, tol)

    iγ = inv(γ)
    iρ = inv(ρ)

    # pad to fast fft size
    xp = padfastfft(@view(f[:,:,:,1]), pad, rfft=true)
    m = padfastfft(mask, pad, rfft=true)

    # initialize variables and fft
    sz0 = size(mask)
    sz  = size(m)
    sz_ = (sz[1]>>1 + 1, sz[2], sz[3])

    x0 = similar(xp)

    D  = Array{T}(undef, sz_)   # dipole kernel
    L  = Array{T}(undef, sz_)   # laplace kernel
    M  = Array{T}(undef, sz_)   # abs(D) > δ
    iA = Array{T}(undef, sz_)   # 1 / (mu*M - rho*Δ)

    FFTW.set_num_threads(num_cores())
    P = plan_rfft(xp)
    iP = inv(P)

    # lsmr
    A = LinearMap{complex(T)}(
        (Dv, v) -> _A_rts!(Dv, v, D),
        length(D),
        ishermitian = true,
        ismutating = true
    )

    WS = LSMRWorkspace(A)
    B̂ = reshape(WS.b, sz_)      # lsmr rhs
    X̂ = reshape(WS.x, sz_)      # lsmr solution, re-use as in-place rfft var
    F̂ = reshape(WS.v, sz_)      # re-use as pre-computed rhs for admm

    # admm
    px = similar(xp)
    py = similar(xp)
    pz = similar(xp)

    dx = similar(xp)
    dy = similar(xp)
    dz = similar(xp)

    vx = similar(xp)
    vy = similar(xp)
    vz = similar(xp)

    # get kernels
    D = _dipole_kernel!(D, F̂, xp, sz0, vsz, bdir, P, Dkernel, :rfft)
    L = _laplace_kernel!(L, F̂, xp, vsz, P, negative=true)

    # mask of well-conditioned frequencies
    @inbounds @batch for I in eachindex(M)
        M[I] = ifelse(abs(D[I]) > δ, μ, _zero)
    end

    @inbounds for t in axes(f, 4)
        if verbose && size(f, 4) > 1
            @printf("Echo: %d/%d\n", t, size(f, 4))
        end

        xp = padarray!(xp, @view(f[:, :, :, t]))

        ######################################################################
        # Step 1: Well-conditioned
        ######################################################################
        B̂ = mul!(B̂, P, xp)
        lsmr!(WS; lambda=_zero, atol=_zero, btol=_zero, maxit=lstol)

        xp = mul!(xp, iP, X̂)

        @batch for I in eachindex(xp)
            xp[I] *= m[I]
        end

        ######################################################################
        # Step 2: Ill-conditioned
        ######################################################################
        X̂ = mul!(X̂, P, xp)

        @batch for I in eachindex(F̂)
            a = muladd(ρ, L[I], M[I])
            if iszero(a)
                F̂[I] = _zero
                iA[I] = _zero
            else
                ia = inv(a)
                F̂[I] = ia * M[I] * X̂[I]
                iA[I] = ρ*ia
            end
        end

        nr = typemax(T)
        px = tfill!(px, _zero)
        py = tfill!(py, _zero)
        pz = tfill!(pz, _zero)
        vx, vy, vz = _gradfp!(vx, vy, vz, xp, vsz)

        if verbose
            @printf("\n iter\t  ||x-xprev||/||x||\n")
        end

        for i in 1:maxit
            x0, xp = xp, x0

            ##################################################################
            # x - subproblem
            ##################################################################
            # v = y - p
            xp = _gradfp_adj!(xp, vx, vy, vz, vsz)

            X̂ = mul!(X̂, P, xp)
            @batch for I in eachindex(X̂)
                X̂[I] *= iA[I]
                X̂[I] += F̂[I]
            end

            xp = mul!(xp, iP, X̂)

            ##################################################################
            # convergence check
            ##################################################################
            # [1] ndx = norm(x - xprev)
            # [2] nx  = norm(x)
            @batch threadlocal=zeros(T, 2)::Vector{T} for I in eachindex(xp)
                a, b = xp[I], x0[I]
                threadlocal[1] = muladd(a-b, a-b, threadlocal[1])
                threadlocal[2] = muladd(a, a, threadlocal[2])
            end
            ndx, nx = sqrt.(sum(threadlocal::Vector{Vector{T}}))

            if verbose
                @printf("%3d/%d\t    %.4e\n", i, maxit, ndx/nx)
            end

            if ndx < ϵ*nx || i == maxit
                break
            end

            ##################################################################
            # y - subproblem and Lagrange multiplier update
            ##################################################################
            dx, dy, dz = _gradfp!(dx, dy, dz, xp, vsz)

            # y = shrink(∇x + p, 1/ρ)   y-subproblem
            # p = p + ∇x - y            Lagrange multiplier update
            # v = y - p                 pre-compute for x-problem
            # d = ∇x - y                pre-compute for primal residual
            @batch for I in eachindex(px)
                d = dx[I]
                p = px[I] + d
                y = ifelse(abs(p) ≤ iρ, _zero, copysign(abs(p)-iρ, p))
                px[I] = p - y
                vx[I] = y - p + y
                dx[I] = d - y
            end

            @batch for I in eachindex(py)
                d = dy[I]
                p = py[I] + d
                y = ifelse(abs(p) ≤ iρ, _zero, copysign(abs(p)-iρ, p))
                py[I] = p - y
                vy[I] = y - p + y
                dy[I] = d - y
            end

            @batch for I in eachindex(pz)
                d = dz[I]
                p = pz[I] + d
                y = ifelse(abs(p) ≤ iρ, _zero, copysign(abs(p)-iρ, p))
                pz[I] = p - y
                vz[I] = y - p + y
                dz[I] = d - y
            end

            ##################################################################
            # parameter update
            ##################################################################
            nr0 = nr

            # nr = norm(∇x - y), primal residual
            @batch threadlocal=zero(T)::T for I in eachindex(dx)
                threadlocal += sqrt(dx[I]*dx[I] + dy[I]*dy[I] + dz[I]*dz[I])
            end
            nr = sum(threadlocal::Vector{T})

            if nr > τ*nr0
                ρ0 = ρ
                ρ *= γ
                iρ = inv(ρ)

                # update constant rhs and ilhs
                @batch for I in eachindex(iA)
                    iaprev = iA[I]
                    if !iszero(iaprev)
                        ia = inv(muladd(ρ, L[I], M[I]))
                        F̂[I] *= ia*ρ0*inv(iaprev)
                        iA[I] = ρ*ia
                    end
                end

                # update scaled dual variables
                @batch for I in eachindex(px)
                    px[I] *= iγ
                end

                @batch for I in eachindex(py)
                    py[I] *= iγ
                end

                @batch for I in eachindex(pz)
                    pz[I] *= iγ
                end
            end
        end

        unpadarray!(@view(x[:,:,:,t]), xp)

        if verbose
            println()
        end

    end

    return x
end

function _A_rts!(Dv, v, D)
    @inbounds @batch for I in eachindex(Dv)
        Dv[I] = D[I]*v[I]
    end
    return Dv
end


text:
using Random

"""
    ContGridWorld

A Continuous grid world domain with generic dynamics to allow for different types of mazes.
    - state: [y, x]

# Args

"""
mutable struct ContGridWorld{F} <: AbstractEnvironment
    # State is stored as (y,x)
    state::Array{Float64, 1}
    walls::Array{Bool, 2}

    # goals stored as (y,x)
    goals::Array{Int, 2}
    reward_funcs::Dict{Int, F}

    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}
    ContGridWorld(walls, goals, rew_funcs::Dict{Int, F}, max_action_noise, drift_noise, normalized) where {F} =
        new{F}([0.0, 0.0], walls, goals, rew_funcs, max_action_noise, drift_noise, normalized, false, [[0,0] for i in 1:4], fill(false, 4), zeros(2))
end

module ContGridWorldParams
UP = 1
RIGHT = 2
DOWN = 3
LEFT = 4

ACTIONS = [UP, RIGHT, DOWN, LEFT]

const BODY_RADIUS = 0.2
const STEP = 0.5

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

MinimalRLCore.is_terminal(env::ContGridWorld) = begin
    goal_id = which_goal(env)
    if goal_id ∈ keys(env.reward_funcs)
        is_terminal(env.reward_funcs[goal_id])
    else
        false
    end
end

MinimalRLCore.get_reward(env::ContGridWorld) = begin
    goal_id = which_goal(env)
    if goal_id ∈ keys(env.reward_funcs)
        get_reward(env.reward_funcs[goal_id])
    else
        0.0
    end
end
MinimalRLCore.get_state(env::ContGridWorld) = (env.normalized ? env.state./size(env) : env.state, env.collision)

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

valid_projected_state(env::ContGridWorld, prj) = begin
    gw_size = size(env.walls)
    prj[1] >= 1 && prj[2] >= 1 && prj[1] <= gw_size[1]  && prj[2] <= gw_size[2]
end

project(env::ContGridWorld, state) = [Int64(floor(state[1])) + 1, Int64(floor(state[2])) + 1]
project(env::ContGridWorld, state, loc) = begin; loc[1] = Int64(floor(state[1]) + 1); loc[2] = Int64(floor(state[2]) + 1); end;

function is_wall(env::ContGridWorld, state::Array{Float64, 1})
    prj = project(env, state)
    return is_wall(env, prj)
end

function is_wall(env::ContGridWorld, prj::Array{Int64, 1})
    gw_size = size(env.walls)
    if !valid_projected_state(env, prj)
        return false
    end
    env.walls[prj[1], prj[2]]
end

which_goal(env::ContGridWorld) = which_goal(env, env.state)

function which_goal(env::ContGridWorld, state::Array{Float64, 1})
    prj = project(env, state)
    if !valid_projected_state(env, prj)
        return -1
    end
    env.goals[prj[1], prj[2]]
end

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

    #Approximate by the square...
    CRP = ContGridWorldParams
    new_state = copy(state)
    # env.new_state .= state
    # new_state = copy(env.new_state)

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

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

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

    for i in 1:4
        project(env::ContGridWorld, new_state .+ CRP.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 == CRP.DOWN
            new_state[1] = (wall_piece[1]) + CRP.BODY_RADIUS
        elseif action == CRP.UP
            new_state[1] = (wall_piece[1] - 1) - CRP.BODY_RADIUS
        end
    elseif env.collision_check[1]
        new_state[1] = (env.edge_locs[1][1] - 1) - CRP.BODY_RADIUS
    elseif env.collision_check[2]
        new_state[1] = (env.edge_locs[2][1]) + CRP.BODY_RADIUS
    end

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

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

    return new_state, collided
end

function handle_collision!(env::ContGridWorld, action)
    env.state, collided = handle_collision(env, env.state, action)
    return collided
end

function MinimalRLCore.reset!(env::ContGridWorld, rng::Random.AbstractRNG=Random.GLOBAL_RNG)
    state = random_state(env, rng)

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

function MinimalRLCore.reset!(env::ContGridWorld, state::AbstractArray)
    if is_wall(env, state)
        throw("Cannot reset environment to invalid state: $(state)")
    end

    env.state = state
    return state
end

function mini_step!(env::ContGridWorld, step, action)
    env.state .+= step
    collision = handle_collision!(env, action)
    return collision
end

function MinimalRLCore.environment_step!(env::ContGridWorld, action, rng=Random.GLOBAL_RNG)

    CRP = ContGridWorldParams
    next_step = zeros(2)

    if action == CRP.UP
        next_step[1] = (CRP.STEP - rand(rng)*env.max_action_noise)
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == CRP.DOWN
        next_step[1] = -(CRP.STEP - rand(rng)*env.max_action_noise)
        next_step[2] = randn(rng)*env.drift_noise
    elseif action == CRP.RIGHT
        next_step[2] = CRP.STEP - rand(rng)*env.max_action_noise
        next_step[1] = randn(rng)*env.drift_noise
    elseif action == CRP.LEFT
        next_step[2] = -(CRP.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 mini_step(env::ContGridWorld, state, step, action)

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

# function _step(env::ContGridWorld, 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