text:
"""
Internal abstract type.
"""
abstract type AbstractSimulation end

"""
Stores sample value and corresponding log probability.
"""
mutable struct Sample
	value::Any
	logprob::Float64
end

"""
Stores distributions of environment variables.
"""
const Environment = Dict{Symbol, Sampleable}

"""
Stores samples of environment distributions.
"""
const EnvironmentSample = Dict{Symbol, Sample}

"""
Stores values of environment samples.
"""
const EnvironmentValue = Dict{Symbol, Any}

"""
Samples environment, returning EnvironmentValue or array (default).
"""
function rand(env::Environment; flat::Bool=true)
	value = EnvironmentValue(k => rand(dist) for (k, dist) in env)
	return flat ? flatten(env, value) : value
end

"""
Computes log probability of environment sample.
"""
logprob(sample::EnvironmentSample) = sum(s.logprob for (_, s) in sample)

"""
Stores inferred properties of environment variable.
"""
struct EnvironmentInfo
	n::Int64	# environment variable dimensionality
	t::Type		# environment variable type
end

"""
MDP object for AST. Wraps simulation and contains auxiliary information and parameters.
"""
mutable struct ASTMDP <: CommonRLInterface.AbstractEnv
	sim::AbstractSimulation						# simulation wrapping system under test
	reward_bonus::Float64						# bonus for reaching event (default zero)
	env_info::Dict{Symbol, EnvironmentInfo}		# inferred environment properties
end

"""
Infers dimension of action space.
"""
act_dim(mdp::ASTMDP) = sum(info.n for info in values(mdp.env_info))

"""
Infers dimension of state space.
"""
obs_dim(mdp::ASTMDP) = length(observe(mdp.sim))

"""
Flattens EnvironmentValue into single array.
"""
function flatten(env::Environment, value::EnvironmentValue)
	action = Float32[]
	for k in sort(collect(keys(env)))
		array = flatten(env[k], value[k])
		append!(action, array)
	end
	return action
end

"""
Reconstructs EnvironmentValue from single array.
"""
function unflatten(mdp::ASTMDP, action::Vector{<:Real})
	value = EnvironmentValue()
	env = environment(mdp.sim)

	i = 0
	for k in sort(collect(keys(env)))
		n = mdp.env_info[k].n
		array = action[i+1:i+n]
		value[k] = unflatten(env[k], array)
		i += n
	end
	return value
end

"""
Converts EnvironmentValue to EnvironmentSample, calculating associated log probabilities.
"""
function create_sample(env::Environment, value::EnvironmentValue)
	sample = EnvironmentSample()
	for k in keys(env)
		dist, val = env[k], value[k]
		sample[k] = Sample(val, logpdf(dist, val))
	end
	return sample
end

"""
Calculates AST reward from MDP, sample, and before/after distances.
"""
function reward(mdp::ASTMDP, sample::EnvironmentSample, ds::NTuple{2,<:Real})
	r = logprob(sample)
	r += isevent(mdp.sim) ? mdp.reward_bonus : 0.0
	r += ds[1] - ds[2]
	return r
end

"""
Constructor for ASTMDP object. Infers various properties of MDP.
"""
function ASTMDP(sim::AbstractSimulation; reward_bonus=0.0)
	reset!(sim)
	env = environment(sim)
	env_info = Dict{Symbol, EnvironmentInfo}()
	for (k, dist) in env
		sample = rand(dist)
		array = flatten(dist, sample)
		env_info[k] = EnvironmentInfo(length(array), typeof(sample))
	end
	return ASTMDP(sim, reward_bonus, env_info)
end


text:
# Script for running the MCMC algorithm
# set number of threads for BLAS

# load case models
include("rickermodel.jl")

# set up problem
problem = set_up_problem(ploton = false)

problem.alg_param.N = 1000
problem.alg_param.R = 52000
problem.alg_param.burn_in = 2000
problem.data.y = Array(readtable("y_data_set_2.csv"))[:,1]
#problem.data.y = Array(readtable("y_data_200_obs_3.csv"))[:,1]
#problem.data.y = Array(readtable("y_10000_obs.csv"))[:,1]

#Array(readtable("y.csv"))[:,1]
#problem.data.y = Array(readtable("y_data_set_abc.csv"))[:,1] #Array(readtable("y.csv"))[:,1]

problem.alg_param.print_interval = 10000 #problem.alg_param.R
# test starting at true parameters
#problem.model_param.theta_0 = problem.model_param.theta_true

#MCWM
#jobname = "mcwm"
#problem.alg_param.alg = "MCWM"
#problem.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.4, 1., 0.8, 25)

# old setting of mcwm
#problem.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.4, 1., 0.8, 25)

# PMCMC
jobname = "pmcmc"
problem.alg_param.alg = "PMCMC"
problem.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.4, 1., 0.8, 25)

# use AM alg for adaptive updating
#problem.adaptive_update = AMUpdate(eye(3), 2.4/sqrt(3), 1., 0.7, 25)

#problem.adaptive_update = noAdaptation(2.4/sqrt(3)*eye(3))

# or, use AM gen alg for adaptive updating
#problem.adaptive_update = AMUpdate_gen(eye(3), 2.4/sqrt(3), 0.2, 1., 0.8, 25)

# run adaptive PMCMC
tic()
res_MCMC = mcmc(problem)
time_MCMC = toc()
@printf "Run time (s): %.4f \n" time_MCMC

# profiling
#using ProfileView
#Profile.clear()
#res_MCMC = @profile MCMC(problem)
#ProfileView.view()

# write outputs
res = res_MCMC[1]

Theta = res.Theta_est
loglik = res.loglik_est
accept_vec = res.accept_vec
prior_vec = res.prior_vec

loglik_avec_priorvec = zeros(3, length(loglik))
loglik_avec_priorvec[1,:] = loglik
loglik_avec_priorvec[2,:] = accept_vec
loglik_avec_priorvec[3,:] = prior_vec

algorithm_parameters = zeros(10, 2)

algorithm_parameters[1,1] = problem.alg_param.burn_in
algorithm_parameters[2:4,1] = problem.model_param.theta_true
algorithm_parameters[5:7,1] = problem.model_param.theta_0
algorithm_parameters[8:end,:] = problem.prior_dist.prior_parameters

writetable("Results/Theta"*jobname*".csv", convert(DataFrame, Theta))
writetable("Results/loglik_avec_priorvec"*jobname*".csv", convert(DataFrame, loglik_avec_priorvec))
writetable("Results/algorithm_parameters"*jobname*".csv", convert(DataFrame, algorithm_parameters))

ext: jl

max_stars_repo_path: Ricker model/run_mcmc.jl

max_stars_repo_name: SamuelWiqvist/adamcmcpaper

max_issues_repo_path: Ricker model/run_mcmc.jl

max_issues_repo_name: SamuelWiqvist/adamcmcpaper

max_forks_repo_path: Ricker model/run_mcmc.jl

max_forks_repo_name: SamuelWiqvist/adamcmcpaper


text:
@testset "Dicts can be diffed" begin
    d1 = Dict(
        :foo => "foo",
        :bar => "bar",
        :baz => Dict(
            :fizz => "fizz",
            :buzz => "buzz"
        )
    )

    @testset "One Changed" begin
        d = deepdiff(d1, Dict(
            :foo => "foo",
            :bar => "biz",
            :baz => Dict(
                :fizz => "fizz",
                :buzz => "buzz"
            )
        ))
        @test added(d) == Set()
        @test removed(d) == Set()
        @test changed(d) == Dict{Symbol, DeepDiffs.DeepDiff}(:bar => deepdiff("bar", "biz"))
    end

    @testset "One Removed" begin
        d = deepdiff(d1, Dict(
            :foo => "foo",
            :baz => Dict(
                :fizz => "fizz",
                :buzz => "buzz"
            )
        ))
        @test added(d) == Set()
        @test removed(d) == Set([:bar])
        @test changed(d) == Dict()
    end

    @testset "One Added" begin
        d = deepdiff(d1, Dict(
            :foo => "foo",
            :bar => "bar",
            :biz => "biz",
            :baz => Dict(
                :fizz => "fizz",
                :buzz => "buzz"
            )
        ))
        @test added(d) == Set([:biz])
        @test removed(d) == Set()
        @test changed(d) == Dict()
    end

    @testset "Inner Dict Modified" begin
        d = deepdiff(d1, Dict(
            :foo => "foo",
            :bar => "bar",
            :baz => Dict(
                :fizz => "fizz",
                :buzz => "bizzle"
            )
        ))
        @test added(d) == Set()
        @test removed(d) == Set()
        @test changed(d) == Dict{Symbol, DeepDiffs.DeepDiff}(:baz => deepdiff(
            Dict(
                :fizz => "fizz",
                :buzz => "buzz"
            ),
            Dict(
                :fizz => "fizz",
                :buzz => "bizzle"
            )
        ))
    end

    @testset "Totally Removed" begin
        d = deepdiff(d1, Dict())
        @test added(d) == Set()
        @test removed(d) == Set([:foo, :bar, :baz])
        @test changed(d) == Dict()
    end

    @testset "Totally added" begin
        d = deepdiff(Dict(), d1)
        @test added(d) == Set([:foo, :bar, :baz])
        @test removed(d) == Set()
        @test changed(d) == Dict()
    end
end


text:
##########
# TArray #
##########

"""
    TArray{T}(dims, ...)

Implementation of data structures that automatically perform copy-on-write after task copying.

For each `TArray` object, we have a task-local array object, i.e. if
you access (read or write) the same `TArray` object in different
tasks, you may be dealing with different array objects each belonging
to a different task.  These task-specific array objects, however,
share the same parent `TArray` object.

More specifically, we store the underlying arrays in the field
`TArray.data`, a dictionary whose keys are tasks, instead of storing
them in the task local storage, for performance considerations.

Usage:

```julia
TArray(dim)
```

Example:

```julia
ta = TArray(4)              # init
for i in 1:4 ta[i] = i end  # assign
Array(ta)                   # convert to 4-element Array{Int64,1}: [1, 2, 3, 4]
```
"""
mutable struct TArray{T, N, A <: AbstractArray{T, N}} <: AbstractArray{T, N}
    orig_task :: Task
    data::Dict{Task, Tuple{Int, A}}
    function TArray{T, N, A}() where {T, N, A <: AbstractArray{T, N}}
        d = Dict{Task, Tuple{Int, A}}()
        res = new(current_task(), d)
        register_to_keeper(res)
        return res
    end
end

TArray{T}(d::Integer...) where T = TArray(T, d)
TArray{T}(::UndefInitializer, d::Integer...) where T = TArray(T, d)
TArray{T}(::UndefInitializer, dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)
TArray{T,N}(d::Vararg{<:Integer,N}) where {T,N} = TArray(T, d)
TArray{T,N}(::UndefInitializer, d::Vararg{<:Integer,N}) where {T,N} = TArray{T,N}(d)
TArray{T,N}(dim::NTuple{N,Int}) where {T,N} = TArray(T, dim)

function TArray(T::Type, dim)
    N_dim = length(dim)
    res = TArray{T, N_dim, Array{T, N_dim}}()
    n = n_copies()
    d = Array{T}(undef, dim)
    _set_local_storage(res, n, d)
    res
end

TArray(x::AbstractArray) = convert(TArray, x)

# TArray House-Keeper
const TArrayKeeper = Vector{WeakRef}()
register_to_keeper(x::TArray) = push!(TArrayKeeper, WeakRef(x))
function copy_tarrays(task1::Task, task2::Task)
    filter!(x -> x.value !== nothing, TArrayKeeper)
    for wref in TArrayKeeper
        ta = wref.value
        if ta !== nothing && haskey(ta.data, task1) && !haskey(ta.data, task2)
            ta.data[task2] = ta.data[task1]
        end
    end
end

# _local_storage

_get_local_storage(x) = x
function _get_local_storage(x::TArray; copydata=false)
    n, d = x.data[current_task()]
    copydata || return d
    cn   = n_copies()
    newd = d
    if cn > n
        newd = deepcopy(d)
        _set_local_storage(x, cn, newd)
    end
    return newd
end

function _set_local_storage(x::TArray{T, N}, n_copies::Int, d::AbstractArray{T, N}) where {T, N}
    x.data[current_task()] = (n_copies, d)
end

localize(x) = x
localize(x::AbstractArray) = TArray(x)

# Constructors
"""
     tzeros(dims, ...)

Construct a distributed array of zeros.
Trailing arguments are the same as those accepted by `TArray`.

```julia
tzeros(dim)
```

Example:

```julia
tz = tzeros(4)              # construct
Array(tz)                   # convert to 4-element Array{Int64,1}: [0, 0, 0, 0]
```
"""
function tzeros(T::Type, dim)
    res = TArray{T,length(dim), Array{T, length(dim)}}();
    n = n_copies()
    d = zeros(T,dim)
    _set_local_storage(res, n, d)
    return res
end

tzeros(::Type{T}, d1::Integer, drest::Integer...) where T = tzeros(T, convert(Dims, tuple(d1, drest...)))
tzeros(d1::Integer, drest::Integer...) = tzeros(Float64, convert(Dims, tuple(d1, drest...)))
tzeros(d::Dims) = tzeros(Float64, d)

"""
     tfill(val, dim, ...)

Construct a TArray of a specified value.

```julia
tfill(val, dim)
```

Example:

```julia
tz = tfill(9.0, 4)            # construct
Array(tz)                     # convert to 4-element Array{Float64,1}:  [9.0  9.0  9.0  9.0]
```
"""
function tfill(val::Real, dim)
    res = TArray{typeof(val), length(dim), Array{typeof(val), length(dim)}}();
    n = n_copies()
    d = fill(val,dim)
    _set_local_storage(res, n, d)
    return res
end

#
# Conversion between TArray and Array
#

function Base.convert(::Type{Array}, x::TArray)
    return convert(Array{eltype(x), ndims(x)}, x)
end
function Base.convert(::Type{Array{T, N}}, x::TArray{T, N}) where {T, N}
    c = convert(Array{T, N}, deepcopy(_get_local_storage(x)))
    return c
end

function Base.convert(::Type{TArray}, x::AbstractArray)
    return convert(TArray{eltype(x), ndims(x)}, x)
end
function Base.convert(::Type{TArray{T, N}}, x::AbstractArray{T, N}) where {T, N}
    res = TArray{T, N, typeof(x)}()
    n   = n_copies()
    _set_local_storage(res, n, x)
    return res
end

#
# Representation
#
function Base.show(io::IO, ::MIME"text/plain", x::TArray)
    arr = x.data[x.orig_task][2]
    @warn "Here shows the originating task's storage, " *
        "not the current task's storage. " *
        "Please explicitly call show(::TArray) to display the current task's version of a TArray."
    show(io,  MIME("text/plain"), arr)
end

Base.show(io::IO, x::TArray) = show(io, _get_local_storage(x))

function Base.summary(io::IO, x::TArray)
  print(io, "Task Local Array: ")
  summary(io, _get_local_storage(x))
end

#
# Forward many methods to the underlying array
#
for F in (:size,
          :iterate,
          :firstindex, :lastindex, :axes)
    @eval Base.$F(a::TArray, args...) = $F(_get_local_storage(a), args...)
end

#
# Similarity implementation
#

Base.similar(x::TArray, ::Type{T}, dims::Dims) where T = TArray(similar(_get_local_storage(x), T, dims))

for op in [:(==), :≈]
    @eval Base.$op(x::TArray, y::AbstractArray) = Base.$op(_get_local_storage(x), y)
    @eval Base.$op(x::AbstractArray, y::TArray) = Base.$op(x, _get_local_storage(y))
    @eval Base.$op(x::TArray, y::TArray) = Base.$op(_get_local_storage(x), _get_local_storage(y))
end

#
# Array Stdlib
#

# Indexing Interface
Base.@propagate_inbounds function Base.getindex(x::TArray{T, N}, I::Vararg{Int,N}) where {T, N}
    return _get_local_storage(x)[I...]
end

Base.@propagate_inbounds function Base.setindex!(x::TArray{T, N}, e, I::Vararg{Int,N}) where {T, N}
    d = _get_local_storage(x; copydata=true)
    d[I...] = e
end

function Base.push!(x::TArray, e)
    d = _get_local_storage(x; copydata=true)
    push!(d, e)
end

function Base.pop!(x::TArray)
    d = _get_local_storage(x; copydata=true)
    pop!(d)
end

# Other methods from stdlib

Base.view(x::TArray, inds...; kwargs...) =
    Base.view(_get_local_storage(x), inds...; kwargs...) |> localize
Base.:-(x::TArray) = (- _get_local_storage(x)) |> localize
Base.transpose(x::TArray) = transpose(_get_local_storage(x)) |> localize
Base.adjoint(x::TArray) = adjoint(_get_local_storage(x)) |> localize
Base.repeat(x::TArray; kw...) = repeat(_get_local_storage(x); kw...) |> localize

Base.hcat(xs::Union{TArray{T,1}, TArray{T,2}}...) where T =
    hcat(_get_local_storage.(xs)...) |> localize
Base.vcat(xs::Union{TArray{T,1}, TArray{T,2}}...) where T =
    vcat(_get_local_storage.(xs)...) |> localize
Base.cat(xs::Union{TArray{T,1}, TArray{T,2}}...; dims) where T =
    cat(_get_local_storage.(xs)...; dims = dims) |> localize

Base.reshape(x::TArray, dims::Union{Colon,Int}...) = reshape(_get_local_storage(x), dims) |> localize
Base.reshape(x::TArray, dims::Tuple{Vararg{Union{Int,Colon}}}) =
    reshape(_get_local_storage(x), Base._reshape_uncolon(_get_local_storage(x), dims)) |> localize
Base.reshape(x::TArray, dims::Tuple{Vararg{Int}}) = reshape(_get_local_storage(x), dims) |> localize

Base.permutedims(x::TArray, perm) = permutedims(_get_local_storage(x), perm) |> localize
Base.PermutedDimsArray(x::TArray, perm) = PermutedDimsArray(_get_local_storage(x), perm) |> localize
Base.reverse(x::TArray; dims) = reverse(_get_local_storage(x), dims = dims) |> localize

Base.sum(x::TArray; dims = :) = sum(_get_local_storage(x), dims = dims) |> localize
Base.sum(f::Union{Function,Type},x::TArray) = sum(f.(_get_local_storage(x))) |> localize
Base.prod(x::TArray; dims=:) = prod(_get_local_storage(x); dims=dims) |> localize
Base.prod(f::Union{Function, Type}, x::TArray) = prod(f.(_get_local_storage(x))) |> localize

Base.findfirst(x::TArray, args...) = findfirst(_get_local_storage(x), args...) |> localize
Base.maximum(x::TArray; dims = :) = maximum(_get_local_storage(x), dims = dims) |> localize
Base.minimum(x::TArray; dims = :) = minimum(_get_local_storage(x), dims = dims) |> localize

Base.:/(x::TArray, y::TArray) = _get_local_storage(x) / _get_local_storage(y) |> localize
Base.:/(x::AbstractArray, y::TArray) = x / _get_local_storage(y) |> localize
Base.:/(x::TArray, y::AbstractArray) = _get_local_storage(x) / y |> localize
Base.:\(x::TArray, y::TArray) = _get_local_storage(x) \ _get_local_storage(y) |> localize
Base.:\(x::AbstractArray, y::TArray) = x \ _get_local_storage(y) |> localize
Base.:\(x::TArray, y::AbstractArray) = _get_local_storage(x) \ y |> localize
Base.:*(x::TArray, y::TArray) = _get_local_storage(x) * _get_local_storage(y) |> localize
Base.:*(x::AbstractArray, y::TArray) = x * _get_local_storage(y) |> localize
Base.:*(x::TArray, y::AbstractArray) = _get_local_storage(x) * y |> localize

# broadcast
Base.BroadcastStyle(::Type{<:TArray}) = Broadcast.ArrayStyle{TArray}()
Broadcast.broadcasted(::Broadcast.ArrayStyle{TArray}, f, args...) = f.(_get_local_storage.(args)...) |> localize

import LinearAlgebra
import LinearAlgebra:  \, /, inv, det, logdet, logabsdet, norm

LinearAlgebra.inv(x::TArray) = inv(_get_local_storage(x)) |> localize
LinearAlgebra.det(x::TArray) = det(_get_local_storage(x)) |> localize
LinearAlgebra.logdet(x::TArray) = logdet(_get_local_storage(x)) |> localize
LinearAlgebra.logabsdet(x::TArray) = logabsdet(_get_local_storage(x)) |> localize
LinearAlgebra.norm(x::TArray, p::Real = 2) =
    LinearAlgebra.norm(_get_local_storage(x), p) |> localize

import LinearAlgebra: dot
dot(x::TArray, ys::TArray) = dot(_get_local_storage(x), _get_local_storage(ys)) |> localize
dot(x::AbstractArray, ys::TArray) = dot(x, _get_local_storage(ys)) |> localize
dot(x::TArray, ys::AbstractArray) = dot(_get_local_storage(x), ys) |> localize

using Statistics
Statistics.mean(x::TArray; dims = :) = mean(_get_local_storage(x), dims = dims) |> localize
Statistics.std(x::TArray; kw...) = std(_get_local_storage(x), kw...) |> localize

# TODO
# * NNlib


text:
module Partitions

using Random

export partition, ballanced_partition, Partitioner, IndexPartitioner,
    training_set, testing_set

abstract type AbstractPartitioner end

struct Partitioner{T<:AbstractMatrix{Float64},L<:AbstractVector{Bool}} <: AbstractPartitioner
    x::T
    y::L
    rows::Vector{Vector{Int}}
end

struct Partition
    xtrain::Matrix{Float64}
    xtest::Matrix{Float64}
    ytrain::Vector{Bool}
    ytest::Vector{Bool}
end

struct IndexPartitioner <: AbstractPartitioner
    rows::Vector{Vector{Int}}
end

IndexPartitioner(x, y, krow::Vector{Vector{Int}}) = IndexPartitioner(krow)

struct IndexPartition
    idxtrain::Vector{Int}
    idxtest::Vector{Int}
end

@inline training_set(p::Partition) = p.xtrain, p.ytrain
@inline testing_set(p::Partition) = p.xtest, p.ytest
@inline training_set(p::IndexPartition) = p.idxtrain
@inline testing_set(p::IndexPartition) = p.idxtest

Base.IteratorSize(p::AbstractPartitioner) = HasLength()
Base.IteratorEltype(p::AbstractPartitioner) = HasEltype()

@inline Base.length(p::AbstractPartitioner) = length(p.rows)
@inline Base.eltype(p::Partitioner) = Partition
@inline Base.eltype(p::IndexPartitioner) = IndexPartition

Base.iterate(p::AbstractPartitioner) = iterate(p, 1)

function Base.iterate(p::Partitioner, k::Integer)
    k > length(p.rows) && return nothing
    test = p.rows[k]
    ktrain = setdiff(1:length(p.rows), k)
    train = vcat(p.rows[ktrain]...)
    return Partition(p.x[train,:], p.x[test,:], p.y[train], p.y[test]), k + 1
end

function Base.iterate(p::IndexPartitioner, k::Integer)
    k > length(p.rows) && return nothing
    test = p.rows[k]
    ktrain = setdiff(1:length(p.rows), k)
    train = vcat(p.rows[ktrain]...)
    return IndexPartition(train, test), k + 1
end

function partition(::Type{T}, x::AbstractArray, y::AbstractArray, nfold::Integer, shfl::Bool=false) where T<:AbstractPartitioner

    size(x, 1) != length(y) && error("Sizes of inputs do not match!")

    krow = Vector{Vector{Int}}(undef, nfold)
    sz = floor(Int, length(y) / nfold)
    res = length(y) - (sz * nfold)

    idx = collect(1:length(y))
    if shfl
        shuffle!(idx)
    end

    ks = 1
    for k = 1:nfold
        ke = ks + sz
        if k > res
            ke -= 1
        end
        krow[k] = idx[ks:ke]
        ks = ke + 1
    end

    return T(x, y, krow)
end

# attempts to ballance the number of relayed spikes per test-set per fold
# such that each test set has an approximatley equal number of relayed spikes
function ballanced_partition(::Type{T}, x::AbstractArray, y::AbstractArray, nfold::Integer, shfl::Bool=false) where T<:AbstractPartitioner

    size(x, 1) != length(y) && error("Sizes of inputs do not match!")

    krow = Vector{Vector{Int}}(undef, nfold)

    kr = findall(>(0), y)
    kf = setdiff(1:length(y), kr)
    if shfl
        shuffle!(kr)
        shuffle!(kf)
    end

    nrf = floor(Int, length(kr) / nfold)
    nff = floor(Int, length(kf) / nfold)
    rres = length(kr) - (nrf * nfold)
    fres = length(kf) - (nff * nfold)

    ksr = 1
    ksf = 1
    for k = 1:nfold
        ker = ksr + nrf
        kef = ksf + nff
        if k > rres
            ker -= 1
        end
        if k > fres
            kef -= 1
        end
        krow[k] = shuffle!(vcat(kr[ksr:ker], kf[ksf:kef]))
        ksr = ker + 1
        ksf = kef + 1
    end

    return T(x, y, krow)
end

end


text:
module PointBasedValueIteration

using POMDPs
using POMDPPolicies
using POMDPModelTools
using POMDPLinter
using BeliefUpdaters
using LinearAlgebra
using Distributions
using FiniteHorizonPOMDPs

import POMDPs: Solver, solve
import Base: ==, hash, convert
import FiniteHorizonPOMDPs: InStageDistribution, FixedHorizonPOMDPWrapper

export
    PBVISolver,
    solve

include("pbvi.jl")

end # module


text:
using UnPack: @unpack
using MuladdMacro: @muladd
using DiffEqBase: @..
using LinearAlgebra: Diagonal
using NQCDynamics: RingPolymers, ndofs

export BCOCB

struct BCOCB <: StochasticDiffEq.StochasticDiffEqAlgorithm end

StochasticDiffEq.alg_compatible(prob::DiffEqBase.AbstractSDEProblem, alg::BCOCB) = true

struct BCOCBCache{U,A,K,uEltypeNoUnits,F} <: StochasticDiffEq.StochasticDiffEqMutableCache
    tmp::U
    rtmp::A
    vtmp::A
    k::K
    cayley::Vector{Matrix{uEltypeNoUnits}}
    half::uEltypeNoUnits
    friction::F
end

function StochasticDiffEq.alg_cache(::BCOCB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{true}})
    tmp = zero(u)
    rtmp = zero(u.x[1])
    vtmp = zero(rtmp)
    k = zeros(size(rtmp))

    cayley = RingPolymers.cayley_propagator(p.beads, dt; half=true)
    half = uEltypeNoUnits(1//2)

    friction = FrictionCache(p, dt)

    BCOCBCache(tmp, rtmp, vtmp, k, cayley, half, friction)
end

struct MDEFCache{flatV,G,N,C,M,S}
    flatvtmp::flatV
    tmp1::flatV
    tmp2::flatV
    gtmp::G
    noise::N
    c1::C
    c2::C
    sqrtmass::M
    σ::S
end

function FrictionCache(sim::RingPolymerSimulation{<:DynamicsMethods.ClassicalMethods.MDEF}, dt)
    DoFs = ndofs(sim)
    atoms = length(sim.atoms)
    beads = length(sim.beads)

    flatvtmp = zeros(DoFs*atoms)
    tmp1 = zero(flatvtmp)
    tmp2 = zero(flatvtmp)

    gtmp = zeros(DoFs*atoms, DoFs*atoms, beads)
    noise = zero(flatvtmp)

    c1 = Diagonal(zeros(DoFs*atoms, DoFs*atoms))
    c2 = zero(c1)

    sqrtmass = 1 ./ sqrt.(repeat(sim.atoms.masses; inner=DoFs))
    σ = zero(sqrtmass)

    MDEFCache(flatvtmp, tmp1, tmp2, gtmp, noise, c1, c2, sqrtmass, σ)
end

struct LangevinCache{C,M,S}
    c1::C
    c2::C
    sqrtmass::M
    σ::S
end

function FrictionCache(sim::RingPolymerSimulation{<:DynamicsMethods.ClassicalMethods.ThermalLangevin}, dt)
    γ0 = sim.method.γ
    γ = [γ0, 2sqrt.(2sim.beads.normal_mode_springs[2:end])...]
    c1 = exp.(-dt.*γ)
    c2 = sqrt.(1 .- c1.^2)
    sqrtmass = 1 ./ repeat(sqrt.(sim.atoms.masses'), ndofs(sim))
    σ = zero(sqrtmass)

    LangevinCache(c1, c2, sqrtmass, σ)
end

function StochasticDiffEq.initialize!(integrator, cache::BCOCBCache)
    @unpack t,uprev,p = integrator
    v = integrator.uprev.x[1]
    r = integrator.uprev.x[2]

    integrator.f.f1(cache.k,v,r,p,t)
end

@muladd function StochasticDiffEq.perform_step!(integrator,cache::BCOCBCache,f=integrator.f)
    @unpack t,dt,sqdt,uprev,u,p,W = integrator
    @unpack rtmp, vtmp, k, cayley, half, friction = cache

    v1 = uprev.x[1]
    r1 = uprev.x[2]

    step_B!(vtmp, v1, half*dt, k)
    @.. rtmp = r1

    RingPolymers.transform_to_normal_modes!(p.beads, rtmp)
    RingPolymers.transform_to_normal_modes!(p.beads, vtmp)

    step_C!(vtmp, rtmp, cayley)

    step_O!(friction, integrator, vtmp, rtmp, t+half*dt)

    step_C!(vtmp, rtmp, cayley)

    RingPolymers.transform_from_normal_modes!(p.beads, rtmp)
    RingPolymers.transform_from_normal_modes!(p.beads, vtmp)

    @.. u.x[2] = rtmp

    integrator.f.f1(k,vtmp,u.x[2],p,t+dt)
    step_B!(u.x[1], vtmp, half*dt, k)
end

function step_C!(v::R, r::R, cayley::Vector{<:Matrix}) where {R<:RingPolymers.RingPolymerArray}
    for i in axes(r, 3)
        for j in v.quantum_atoms
            for k in axes(r, 1)
                rtmp = cayley[i][1,1] * r[k,j,i] + cayley[i][1,2] * v[k,j,i]
                vtmp = cayley[i][2,1] * r[k,j,i] + cayley[i][2,2] * v[k,j,i]
                r[k,j,i] = rtmp
                v[k,j,i] = vtmp
            end
        end
    end

    for j in v.classical_atoms
        for k in axes(r, 1)
            rtmp = cayley[1][1,1] * r[k,j,1] + cayley[1][1,2] * v[k,j,1]
            vtmp = cayley[1][2,1] * r[k,j,1] + cayley[1][2,2] * v[k,j,1]
            r[k,j,1] = rtmp
            v[k,j,1] = vtmp
        end
    end
end

function step_O!(friction::MDEFCache, integrator, v, r, t)
    @unpack W, p, dt, sqdt = integrator
    @unpack flatvtmp, tmp1, tmp2, gtmp, noise, c1, c2, sqrtmass, σ = friction

    integrator.g(gtmp,r,p,t)
    Λ = gtmp
    @.. σ = sqrt(get_temperature(p, t)) * sqrtmass

    @views for i in axes(r, 3)
        @.. noise = σ * W.dW[:,:,i][:] / sqdt

        γ, c = LAPACK.syev!('V', 'U', Λ[:,:,i]) # symmetric eigen
        clamp!(γ, 0, Inf)
        for (j,i) in enumerate(diagind(c1))
            c1[i] = exp(-γ[j]*dt)
            c2[i] = sqrt(1 - c1[i]^2)
        end

        copyto!(flatvtmp, v[:,:,i])
        mul!(tmp1, transpose(c), flatvtmp)
        mul!(tmp2, c1, tmp1)
        mul!(flatvtmp, c, tmp2)

        mul!(tmp1, transpose(c), noise)
        mul!(tmp2, c2, tmp1)
        mul!(tmp1, c, tmp2)

        @.. flatvtmp += tmp1
        copyto!(v[:,:,i], flatvtmp)
    end

end

function step_O!(friction::LangevinCache, integrator, v::R, r::R, t) where {R<:RingPolymers.RingPolymerArray}
    @unpack W, p, dt, sqdt = integrator
    @unpack c1, c2, sqrtmass, σ = friction

    @.. σ = sqrt(get_temperature(p, t)) * sqrtmass

    for i in axes(r, 3)
        for j in v.quantum_atoms
            @. v[:,j,i] = c1[i] * v[:,j,i] + c2[i] * σ[:,j] * W.dW[:,j,i] / sqdt
        end
    end

    for j in v.classical_atoms
        @. v[:,j,1] = c1[1] * v[:,j,1] + c2[1] * σ[:,j] * W.dW[:,j,1] / sqdt
    end
end

DynamicsMethods.select_algorithm(::RingPolymerSimulation{<:DynamicsMethods.ClassicalMethods.AbstractMDEF}) = BCOCB()
DynamicsMethods.select_algorithm(::RingPolymerSimulation{<:DynamicsMethods.ClassicalMethods.ThermalLangevin}) = BCOCB()


text:
@testset "SymbolicCSC0" begin
    Q = @expr [-x x 0;
            1 -2 1;
            0 1 -1]
    csc = SparseCSC(Q)
    y = convert(AbstractMatrixSymbolic{Float64}, csc)
    println(typeof(y))
    convert(AbstractMatrixSymbolic{Float64}, y)
end

@testset "SymbolicCSR1" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = SparseCSR(3, 3, v, [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
    println(typeof(m))
    println(m)
    @test seval(m) == SparseCSR(3, 3, [seval(x)^i + i for i = 1:10], [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicCSR2" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = convert(AbstractMatrixSymbolic{Float64}, SparseCSR(3, 3, v, [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3]))
    println(typeof(m))
    println(m)
    @test seval(m) == SparseCSR(3, 3, [seval(x)^i + i for i = 1:10], [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicCSC1" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = SparseCSC(3, 3, v, [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
    @test seval(m) == SparseCSC(3, 3, [seval(x)^i + i for i = 1:10], [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicCOO1" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = SparseCOO(3, 3, v, [1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 1, 2, 3, 1, 2, 3])
    @test seval(m) == SparseCOO(3, 3, [seval(x)^i + i for i = 1:10], [1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicCSR2" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = SparseCSR(3, 3, v, [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
    @test seval(m, :x) == SparseCSR(3, 3, [i * seval(x)^(i-1) for i = 1:10], [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicCSC2" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = SparseCSC(3, 3, v, [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
    @test seval(m, :x) == SparseCSC(3, 3, [i * seval(x)^(i-1) for i = 1:10], [1, 4, 7, 10], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicCOO2" begin
    @bind x = 10
    v = [x^i + i for i = 1:9]
    m = SparseCOO(3, 3, v, [1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 1, 2, 3, 1, 2, 3])
    @test seval(m, :x) == SparseCOO(3, 3, [i * seval(x)^(i-1) for i = 1:10], [1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 1, 2, 3, 1, 2, 3])
end

@testset "SymbolicMat1" begin
    @bind x = 10
    m = [x^i + i for i = 1:3, j = 1:3]
    result = seval(m)
    @test result == [seval(x)^i + i for i = 1:3, j = 1:3]
end

@testset "SymbolicMat2" begin
    @bind x = 10
    m = [x^i + i for i = 1:3, j = 1:3]
    @test seval(m, :x) == [i*seval(x)^(i-1) for i = 1:3, j = 1:3]
end

@testset "SymbolicMat3" begin
    m = @expr [1 x 3; 3 4 y]
    @test m[1,1].val == 1
    @test m[1,2].var == :x
    @test m[1,3].val == 3
    @test m[2,1].val == 3
    @test m[2,2].val == 4
    @test m[2,3].var == :y
end

@testset "SymbolicMat4" begin
    m = SparseCSR(@expr [0 x 0; 0 z y])
    println(m)
    x0 = 10.0
    y0 = 0.8
    z0 = 4
    test = SymbolicEnv()
    @bind test begin
        x = x0
        y = y0
        z = z0
    end
    test2 = SymbolicEnv()
    @bind test2 begin
        x = 1.0
        y = 0.0
        z = 0.0
    end
    result = seval(m, :x, test)
    expected = seval(m, test2)
    println(typeof(result))
    println(typeof(expected))
    @test isapprox(expected.val, result.val)
end

@testset "SymbolicMat5" begin
    m = SparseCSC(@expr [0 x 0; 0 z y])
    println(m)
    x0 = 10.0
    y0 = 0.8
    z0 = 4
    test = SymbolicEnv()
    @bind test begin
        x = x0
        y = y0
        z = z0
    end
    test2 = SymbolicEnv()
    @bind test2 begin
        x = 1.0
        y = 0.0
        z = 0.0
    end
    result = seval(m, :x, test)
    expected = seval(m, test2)
    @test isapprox(expected.val, result.val)
end

@testset "SymbolicMat6" begin
    m = SparseCOO(@expr [0 x 0; 0 z y])
    println(m)
    x0 = 10.0
    y0 = 0.8
    z0 = 4
    test = SymbolicEnv()
    @bind test begin
        x = x0
        y = y0
        z = z0
    end
    test2 = SymbolicEnv()
    @bind test2 begin
        x = 1.0
        y = 0.0
        z = 0.0
    end
    result = seval(m, :x, test)
    expected = seval(m, test2)
    @test isapprox(expected.val, result.val)
end

@testset "SymbolicMat7" begin
    m = sparse(@expr [0 x 0; 0 z y])
    println(m)
    x0 = 10.0
    y0 = 0.8
    z0 = 4
    test = SymbolicEnv()
    @bind test begin
        x = x0
        y = y0
        z = z0
    end
    test2 = SymbolicEnv()
    @bind test2 begin
        x = 1.0
        y = 0.0
        z = 0.0
    end
    result = seval(m, :x, test)
    expected = seval(m, test2)
    @test isapprox(expected, result)
end


text:
"""
ksintegrateNaive: integrate kuramoto-sivashinsky equation (Julia)
       u_t = -u*u_x - u_xx - u_xxxx, domain x in [0,Lx], periodic BCs

 inputs
          u = initial condition (vector of u(x) values on uniform gridpoints))
         Lx = domain length
         dt = time step
         Nt = number of integration timesteps
      nsave = save every nsave-th time step

 outputs
          u = final state, vector of u(x, Nt*dt) at uniform x gridpoints

This a line-by-line translation of a Matlab code into Julia. It uses out-of-place
FFTs and doesn't pay any attention to the allocation of temporary vectors within
the time-stepping loop. Hence the name "naive".
"""
function ksintegrateNaive(u, Lx, dt, Nt)
    Nx = length(u)                  # number of gridpoints
    x = collect(0:(Nx-1)/Nx)*Lx
    kx = vcat(0:Nx/2-1, 0, -Nx/2+1:-1)  # integer wavenumbers: exp(2*pi*kx*x/L)
    alpha = 2*pi*kx/Lx              # real wavenumbers:    exp(alpha*x)
    D = 1im*alpha                   # D = d/dx operator in Fourier space
    L = alpha.^2 - alpha.^4         # linear operator -D^2 - D^4 in Fourier space
    G = -0.5*D                      # -1/2 D operator in Fourier space

    # Express PDE as u_t = Lu + N(u), L is linear part, N nonlinear part.
    # Then Crank-Nicolson Adams-Bashforth discretization is
    #
    # (I - dt/2 L) u^{n+1} = (I + dt/2 L) u^n + 3dt/2 N^n - dt/2 N^{n-1}
    #
    # let A = (I - dt/2 L)
    #     B = (I + dt/2 L), then the CNAB timestep formula
    #
    # u^{n+1} = A^{-1} (B u^n + 3dt/2 N^n - dt/2 N^{n-1})

    # convenience variables
    dt2  = dt/2
    dt32 = 3*dt/2
    A =  ones(Nx) + dt2*L
    B = (ones(Nx) - dt2*L).^(-1)

    Nn  = G.*fft(u.*u) # -u u_x (spectral), notation Nn = N^n     = N(u(n dt))
    Nn1 = copy(Nn)     #                   notation Nn1 = N^{n-1} = N(u((n-1) dt))
    u  = fft(u)        # transform u to spectral

    # timestepping loop
    for n = 0:Nt
        Nn1 = copy(Nn)                 # shift nonlinear term in time: N^{n-1} <- N^n
        Nn  = G.*fft(real(ifft(u)).^2) # compute Nn = -u u_x

        u = B .* (A .* u + dt32*Nn - dt2*Nn1)
    end

    real(ifft(u))
end

max_stars_repo_name: raphbacher/julia-intro

max_issues_repo_name: raphbacher/julia-intro

max_forks_repo_name: raphbacher/julia-intro


text:
### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ f299bb23-aa3a-4a89-abc8-91459c649a43
begin
	using Revise
	using Pkg
	Pkg.activate("../.")
end

# ╔═╡ b88ba8d6-33c8-4d5b-a673-7030fc768596
using JavaCall, BioformatsLoader, View5D

# ╔═╡ 930855c0-2439-4889-a399-85b23e780b5e
using TestImages, Noise, Napari, ImageShow, Images, ImageFiltering, Statistics

# ╔═╡ 6f9b960e-6742-46d5-b1a5-88c69e24348f
md"## Package loading and some initialization"

# ╔═╡ ecbe0fe3-33b5-4529-b00d-76bf28b66eb4
JavaCall.addClassPath(BioformatsLoader.get_bf_path())

# ╔═╡ a41f9c9d-7a84-4dd4-9b46-62271f4318cb
JavaCall.addOpts("-Xmx2048M"); # Set maximum memory to 2 Gigabyte

# ╔═╡ 77515a38-b65b-11eb-388d-bb59872c6ca1
BioformatsLoader.init()

# ╔═╡ cb5ccedb-2b29-4518-a106-a2579ac82c6e
md"## Testimage Demo"

# ╔═╡ d8ae01cd-7205-4155-80d4-922de74a0464
img = testimage("resolution_test_512")

# ╔═╡ 6e59a98b-a8f6-42e1-972a-85108f5f898e
arr = Float32.(img);

# ╔═╡ d5e27923-fd85-4354-9778-6d14f830bb46
img_color = testimage("mandrill")

# ╔═╡ 7f7a302a-4c5a-4eb2-a41c-9ab0247532c6
arr_color = reshape(permutedims(Float32.(channelview(img_color)), (3,2,1)), (512, 512, 1, 3));

# ╔═╡ 7a6b5f57-bc33-4632-9eb0-259df20198a5
md"## The type of `img` is not a pure array but rather a type which encodes color information like `Gray`"

# ╔═╡ 896db842-feab-4d29-811c-53968122b67f
typeof(img)

# ╔═╡ 08afca09-245c-44ed-b29d-f9fe383b4a06
typeof(img_color)

# ╔═╡ 2f4de0ae-0f22-4426-8db1-35b3047b4376
md"# Different Image Viewer
* View5D
* Napari
"

# ╔═╡ 997308fd-6609-4677-9303-f7d1cea83f17
@vt arr_color

# ╔═╡ 89800d5e-c1ba-4d42-83f5-76b7d1561bee
v = napari.view_image(reshape(arr_color, (512, 512, 3)))

# ╔═╡ dadbf7b7-28fb-41ab-86a0-75dd93a222ea

# ╔═╡ a3d556e5-246d-46a9-b117-76ccb9da6da2
md"## BioFormatsLoader
"

# ╔═╡ a35825e6-7c1a-4772-bd83-49a70171baa0
data = BioformatsLoader.bf_import("/home/fxw/Downloads/150707_WTstack.lsm")

# ╔═╡ cbc4344f-840f-46e5-84ab-6b13b4347f9e
data

# ╔═╡ 1ab43110-151b-42cf-a9c6-6fd0149de564
view5d(arraydata(data[1]))

# ╔═╡ 88beed13-2cdf-4886-b5ba-b613d418ff80
md"## Some Filtering to reduce Noise"

# ╔═╡ 04869acd-5f59-40c3-a236-7f414773ddfe
data_raw = permutedims(Float32.(arraydata(data[1])), (2, 3, 4, 1, 5));

# ╔═╡ 81aeb9f1-3cb8-4364-b492-957b0da70151
@vt data_raw

# ╔═╡ cef373b6-e19d-48cb-b287-1b24c096be95
data_med = mapwindow(median, data_raw, (5, 5, 1, 1, 1));

# ╔═╡ 176f0d6b-690d-48c8-8ac2-dada844a2747
@vt data_med

# ╔═╡ 30db4158-baef-49e3-a57c-d357fe28d4f3
@vt data_raw

# ╔═╡ Cell order:
# ╟─6f9b960e-6742-46d5-b1a5-88c69e24348f
# ╠═f299bb23-aa3a-4a89-abc8-91459c649a43
# ╠═b88ba8d6-33c8-4d5b-a673-7030fc768596
# ╠═ecbe0fe3-33b5-4529-b00d-76bf28b66eb4
# ╠═a41f9c9d-7a84-4dd4-9b46-62271f4318cb
# ╠═930855c0-2439-4889-a399-85b23e780b5e
# ╠═77515a38-b65b-11eb-388d-bb59872c6ca1
# ╟─cb5ccedb-2b29-4518-a106-a2579ac82c6e
# ╠═d8ae01cd-7205-4155-80d4-922de74a0464
# ╠═6e59a98b-a8f6-42e1-972a-85108f5f898e
# ╠═d5e27923-fd85-4354-9778-6d14f830bb46
# ╠═7f7a302a-4c5a-4eb2-a41c-9ab0247532c6
# ╟─7a6b5f57-bc33-4632-9eb0-259df20198a5
# ╠═896db842-feab-4d29-811c-53968122b67f
# ╠═08afca09-245c-44ed-b29d-f9fe383b4a06
# ╠═2f4de0ae-0f22-4426-8db1-35b3047b4376
# ╠═997308fd-6609-4677-9303-f7d1cea83f17
# ╠═89800d5e-c1ba-4d42-83f5-76b7d1561bee
# ╠═dadbf7b7-28fb-41ab-86a0-75dd93a222ea
# ╠═a3d556e5-246d-46a9-b117-76ccb9da6da2
# ╠═a35825e6-7c1a-4772-bd83-49a70171baa0
# ╠═cbc4344f-840f-46e5-84ab-6b13b4347f9e
# ╠═1ab43110-151b-42cf-a9c6-6fd0149de564
# ╠═88beed13-2cdf-4886-b5ba-b613d418ff80
# ╠═04869acd-5f59-40c3-a236-7f414773ddfe
# ╠═81aeb9f1-3cb8-4364-b492-957b0da70151
# ╠═cef373b6-e19d-48cb-b287-1b24c096be95
# ╠═176f0d6b-690d-48c8-8ac2-dada844a2747
# ╠═30db4158-baef-49e3-a57c-d357fe28d4f3


text:
using ShaderAbstractions, LinearAlgebra
using ShaderAbstractions: VertexArray
using Test

struct Bla <: ShaderAbstractions.AbstractContext end

import GeometryTypes, AbstractPlotting, GeometryBasics

m = GeometryTypes.GLNormalMesh(GeometryTypes.Sphere(GeometryTypes.Point3f0(0), 1f0))

mvao = VertexArray(m)
instances = VertexArray(positions = rand(GeometryBasics.Point{3, Float32}, 100))

x = ShaderAbstractions.InstancedProgram(
    Bla(), "hi",
    mvao,
    instances,

    model = GeometryTypes.Matf4f0(I),
    view = GeometryTypes.Matf4f0(I),
    projection = GeometryTypes.Matf4f0(I),

)

x.program.source |> println


text:
using Revise
using POMDPStressTesting
using Distributions
using Parameters
using MATLAB

# Automatic transmission Simulink model used
# in the 2019 ARCH-COMP falsification category.

@with_kw mutable struct ATParams
    model_name::String = "Autotrans_shift" # Simulink model name
    model::String = joinpath(@__DIR__, "models", model_name * ".mdl") # Simulink model full path
    time_range::Vector = [0, 30] # Simulation time
    throttle_range::Vector = [0, 100] # Throttle open percentage
    brake_range::Vector = [0, 325] # Brake torque
    max_actions::Int = 4 # Maximum number of actions
    speed_thresh::Real = 120
    stl::Function = yout->any(yout[:,1] .>= speed_thresh)
end

# Implement abstract GrayBox.Simulation
@with_kw mutable struct ATSim <: GrayBox.Simulation
    params::ATParams = ATParams() # Parameters
    action_matrix::Matrix{Float64} = Matrix{Float64}(undef, 0, 3) # Set of actions (note, Float64 not Real to allow MATLAB to convert)
    t::Real = 0 # Current time
    tout::Array = []
    yout::Array = []
end

# Override from GrayBox
function GrayBox.transition!(sim::ATSim)
    # New time from "remaining" time (i.e. starting at sim.t)
    dist_time::Uniform = Uniform(sim.t, sim.params.time_range[end])
    dist_throttle::Uniform = Uniform(sim.params.throttle_range...)
    dist_brake::Uniform = Uniform(sim.params.brake_range...)

    sampled_time = rand(dist_time)
    sampled_throttle = rand(dist_throttle)
    sampled_brake = rand(dist_brake)
    sim.t = sampled_time # Keep track of latest time.

    input = [sampled_time sampled_throttle sampled_brake] # New input action.
    sim.action_matrix = [sim.action_matrix; input] # Append to action set.

    logprob::Real = logpdf(dist_time, sampled_time) +
                    logpdf(dist_throttle, sampled_throttle) +
                    logpdf(dist_brake, sampled_brake)
    return logprob::Real
end

# Override from BlackBox
function BlackBox.initialize!(sim::ATSim)
    sim.t = 0 # Reset current simulation time
    sim.action_matrix = Matrix{Float64}(undef, 0, 3) # Reset action set
    sim.tout = []
    sim.yout = []

    startime, endtime = sim.params.time_range
    times = startime:endtime
    @mput endtime # pass to MATLAB
    @mput times # pass to MATLAB

    # Load Simulink model in MATLAB
    eval_string("load_system('$(sim.params.model)')")
end

# Override from BlackBox
BlackBox.distance(sim::ATSim) = max(minimum(sim.params.speed_thresh .- sim.yout[:,1]), 0) # Non-negative

# Override from BlackBox
BlackBox.isevent(sim::ATSim) = sim.params.stl(sim.yout)

# Override from BlackBox
BlackBox.isterminal(sim::ATSim) = BlackBox.isevent(sim) || size(sim.action_matrix, 1) >= sim.params.max_actions

# Override from BlackBox
function BlackBox.evaluate!(sim::ATSim)
    logprob::Real = GrayBox.transition!(sim) # Step simulation

    x = sim.action_matrix # Matrix (not array of arrays)
    @mput x # Pass to MATLAB

    # `evalc` to suppress "data type override" print-out.
    eval_string("[~,results] = evalc('sim(''$(sim.params.model_name)'', ''StopTime'', num2str(endtime), ''LoadExternalInput'', ''on'', ''ExternalInput'', mat2str(x), ''SaveTime'', ''on'', ''TimeSaveName'', ''tout'', ''SaveOutput'', ''on'', ''OutputSaveName'', ''yout'', ''SaveFormat'', ''Array'', ''SolverType'', ''variable-step'', ''OutputOption'', ''SpecifiedOutputTimes'', ''ZeroCrossAlgorithm'', ''Adaptive'', ''OutputTimes'', mat2str(times))');")

    # MATLAB.jl cannot directly handle nested structs
    sim.tout = mat"results.tout"
    sim.yout = mat"results.yout"

    d::Real = BlackBox.distance(sim) # Calculate miss distance
    event::Bool = BlackBox.isevent(sim) # Check event indication
    if event
        @info "Failure event found mid-simulation!"
        header = [:time :speed :rpms :gears]
        display(vcat(header, [sim.tout sim.yout])) # TODO. add stl (:failure) category
    end

    return (logprob::Real, d::Real, event::Bool)
end

function setup_ast()
    # Create gray-box simulation object
    sim::GrayBox.Simulation = ATSim()

    # AST MDP formulation object
    mdp::ASTMDP = ASTMDP(sim)
    mdp.params.debug = true # record metrics
    mdp.params.seed = 1
    mdp.params.top_k = 10

    # Hyperparameters for MCTS-PW as the solver
    solver = MCTSPWSolver(depth=sim.params.time_range[end],
                          exploration_constant=10.0,
                          k_action=0.1,
                          alpha_action=0.85,
                          n_iterations=20)

    planner = solve(solver, mdp)

    return planner
end

planner = setup_ast()
action_trace = search!(planner)
final_state = playback(planner, action_trace)

print_metrics(planner)

nothing # Suppress REPL


text:
using Test
import SearchBall
const sb = SearchBall

@testset "SearchBall" begin

include("configuration.jl")
include("game_state.jl")
include("game_view_gtk.jl")
include("geometry_utils.jl")
include("graphics_2d.jl")
include("orders.jl")
include("strategy_utils.jl")

end


text:
module EndowmentsLH

using ArgCheck, DocStringExtensions, Random
using DataFrames, Distributions, Statistics
using CommonLH

include("types.jl");
include("marginals.jl");
include("endowment.jl")
include("endowment_draws.jl")

# Endowment
export AbstractMarginal, UniformMarginal, LogNormalMarginal, NormalMarginal, BetaMarginal, BoundedMarginal, UnboundedMarginal, PercentileMarginal, UnknownMarginal
export Endowment
export label, name, isbounded, lb, ub, marginal, marginal_quantile, validate_draws

# EndowmentDraws
export EndowmentDraws
export validate_draws, add_draws!, replace_draws!,
    select_rows, draw_fixed_percentiles, get_draws, endow_quantiles,
    get_label, get_meta, has_endowment, type_endowments,
    corr_matrix, formatted_corr_matrix

end # module


text:
#  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
#  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
#  Released under the modified BSD license. See COPYING.md for more details.

module MIPLearn

using PyCall

global DynamicLazyConstraintsComponent = PyNULL()
global JuMPSolver = PyNULL()
global MinPrecisionThreshold = PyNULL()
global miplearn = PyNULL()
global ObjectiveValueComponent = PyNULL()
global PrimalSolutionComponent = PyNULL()
global PyFileInstance = PyNULL()
global PyJuMPInstance = PyNULL()
global StaticLazyConstraintsComponent = PyNULL()
global traceback = PyNULL()
global UserCutsComponent = PyNULL()
global MemorySample = PyNULL()
global Hdf5Sample = PyNULL()

include("solvers/structs.jl")

include("utils/log.jl")
include("utils/exceptions.jl")
include("instance/abstract_instance.jl")
include("instance/jump_instance.jl")
include("instance/file_instance.jl")
include("solvers/jump_solver.jl")
include("solvers/learning_solver.jl")
include("solvers/macros.jl")
include("utils/benchmark.jl")
include("utils/parse.jl")

include("bb/BB.jl")

function __init__()
    copy!(miplearn, pyimport("miplearn"))
    copy!(traceback, pyimport("traceback"))
    copy!(DynamicLazyConstraintsComponent, miplearn.DynamicLazyConstraintsComponent)
    copy!(UserCutsComponent, miplearn.UserCutsComponent)
    copy!(ObjectiveValueComponent, miplearn.ObjectiveValueComponent)
    copy!(PrimalSolutionComponent, miplearn.PrimalSolutionComponent)
    copy!(StaticLazyConstraintsComponent, miplearn.StaticLazyConstraintsComponent)
    copy!(MinPrecisionThreshold, miplearn.MinPrecisionThreshold)
    copy!(MemorySample, miplearn.features.sample.MemorySample)
    copy!(Hdf5Sample, miplearn.features.sample.Hdf5Sample)
    __init_PyFileInstance__()
    __init_PyJuMPInstance__()
    __init_JuMPSolver__()

    py"""
    import numpy as np

    def to_str_array(values):
        if values is None:
            return None
        return np.array(values, dtype="S")

    def from_str_array(values):
        return [v.decode() for v in values]
    """
end

to_str_array(values) = py"to_str_array"(values)
from_str_array(values) = py"from_str_array"(values)

function convert(::Type{SparseMatrixCSC}, o::PyObject)
    I, J, V = pyimport("scipy.sparse").find(o)
    return sparse(I .+ 1, J .+ 1, V, o.shape...)
end

function PyObject(m::SparseMatrixCSC)
    pyimport("scipy.sparse").csc_matrix(
        (m.nzval, m.rowval .- 1, m.colptr .- 1),
        shape = size(m),
    ).tocoo()
end

export DynamicLazyConstraintsComponent,
    UserCutsComponent,
    ObjectiveValueComponent,
    PrimalSolutionComponent,
    StaticLazyConstraintsComponent,
    MinPrecisionThreshold,
    Hdf5Sample

end # module


text:
using Scalpels
using Test

@testset "Scalpels.jl" begin

    @testset "Load code to generate simulated data" begin
        @test_nowarn include("../examples/generate_sim_data.jl")
    end
    include("../examples/generate_sim_data.jl")

    @testset "Check Code runs" begin
        @test_nowarn make_simulated_rvs_and_ccfs_to_test_scalpels()
        (rvs, ccfs) = make_simulated_rvs_and_ccfs_to_test_scalpels()
        @test_nowarn clean_rvs_scalpels(rvs,ccfs)
        @test_nowarn rms_clean_rvs_vs_num_basis_scalpels(rvs,ccfs)
    end

    @testset "Test accuracy" begin
        @test_nowarn make_simulated_rvs_and_ccfs_to_test_scalpels()
        (rvs, ccfs) = make_simulated_rvs_and_ccfs_to_test_scalpels()
        rms_vs_num_basis =  rms_clean_rvs_vs_num_basis_scalpels(rvs,ccfs)
        @test all(rms_vs_num_basis[5:end] .< 3.0)
    end

end


text:
# This file contains code that was formerly a part of Julia. License is MIT: http://julialang.org/license

using Base: @deprecate, depwarn
using Compat: @dep_vectorize_1arg, @dep_vectorize_2arg

for f in (:digamma, :trigamma, :zeta, :eta, :erfcx, :erfi, :dawson, :airyai, :airyaiprime,
          :airybi, :airybiprime, :besselj0, :besselj1, :bessely0, :bessely1, :erf, :erfc)
    @eval @dep_vectorize_1arg Number $f
end

for f in (:invdigamma, :erfinc, :erfcinv)
    @eval @dep_vectorize_1arg Real $f
end

for f in (:polygamma, :zeta, :besseli, :besselix, :besselj, :besseljx, :besselk, :besselkx,
          :bessely, :besselyx, :besselh, :besselhx, :hankelh1, :hankelh2, :hankelh1x, :hankelh2x)
    @eval @dep_vectorize_2arg Number $f
end

@deprecate airy(z::Number) airyai(z)
@deprecate airyx(z::Number) airyaix(z)
@deprecate airyprime(z::Number) airyaiprime(z)
@deprecate airy{T<:Number}(x::AbstractArray{T}) airyai.(x)
@deprecate airyx{T<:Number}(x::AbstractArray{T}) airyaix.(x)
@deprecate airyprime{T<:Number}(x::AbstractArray{T}) airyprime.(x)

function _airy(k::Integer, z::Complex128)
    depwarn("`airy(k,x)` is deprecated, use `airyai(x)`, `airyaiprime(x)`, `airybi(x)` or `airybiprime(x)` instead.",:airy)
    id = Int32(k==1 || k==3)
    if k == 0 || k == 1
        return _airy(z, id, Int32(1))
    elseif k == 2 || k == 3
        return _biry(z, id, Int32(1))
    else
        throw(ArgumentError("k must be between 0 and 3"))
    end
end
function _airyx(k::Integer, z::Complex128)
    depwarn("`airyx(k,x)` is deprecated, use `airyaix(x)`, `airyaiprimex(x)`, `airybix(x)` or `airybiprimex(x)` instead.",:airyx)
    id = Int32(k==1 || k==3)
    if k == 0 || k == 1
        return _airy(z, id, Int32(2))
    elseif k == 2 || k == 3
        return _biry(z, id, Int32(2))
    else
        throw(ArgumentError("k must be between 0 and 3"))
    end
end

for afn in (:airy,:airyx)
    _afn = Symbol("_"*string(afn))
    suf  = string(afn)[5:end]
    @eval begin
        function $afn(k::Integer, z::Complex128)
            afn = $(QuoteNode(afn))
            suf = $(QuoteNode(suf))
            depwarn("`$afn(k,x)` is deprecated, use `airyai$suf(x)`, `airyaiprime$suf(x)`, `airybi$suf(x)` or `airybiprime$suf(x)` instead.",$(QuoteNode(afn)))
            $_afn(k,z)
        end

        $afn(k::Integer, z::Complex) = $afn(k, float(z))
        $afn{T<:AbstractFloat}(k::Integer, z::Complex{T}) = throw(MethodError($afn,(k,z)))
        $afn(k::Integer, z::Complex64) = Complex64($afn(k, Complex128(z)))
        $afn(k::Integer, x::Real) = $afn(k, float(x))
        $afn(k::Integer, x::AbstractFloat) = real($afn(k, complex(x)))

        function $afn{T<:Number}(k::Number, x::AbstractArray{T})
            $afn.(k,x)
        end
        function $afn{S<:Number}(k::AbstractArray{S}, x::Number)
            $afn.(k,x)
        end
        function $afn{S<:Number,T<:Number}(k::AbstractArray{S}, x::AbstractArray{T})
            $afn.(k,x)
        end
    end
end


text:
using StatsBase
using Distances

module kNN
    export knn, kernelregression

    using StatsBase
    using Distances
    using NearestNeighbors
	using SmoothingKernels

    include("majority_vote.jl")
    include("classifier.jl")
    include("regress.jl")
end


text:
export ExplicitSolverStep, GeneralizedAlphaSolverStep, ImplicitStaticSolver, LinearStaticSolver,
LinearGeneralizedAlphaSolverStep, LinearGeneralizedAlphaSolverInit!
@doc raw"""
    ExplicitSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64)

Central Difference explicit solver for `M a + fint(u) = fext(u)`. `a`, `v`, `u` are acceleration, velocity and displacement.

```math
\begin{aligned}
u_{n+1} =& u_n + dtv_n + dt^2/2 a_n \\
v_{n+1} =& v_n + dt/2(a_n + a_{n+1}) \\
M a_{n+1} + f^{int}(u_{n+1}) =& f^{ext}_{n+1} \\
M a_{n+1} =& f^{ext}_{n+1} - f^{int}(u_{n+1}) \\
\end{aligned}
```

!!! info
    You need to call SolverInitial! before the first time step, if $f^{ext}_0 \neq 0$.
    Otherwise we assume the initial acceleration `globdat.acce[:] = 0`.
"""
function ExplicitSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64)
    assembleMassMatrix!(globdat, domain)

    u = globdat.state[:]
    ∂u  = globdat.velo[:]
    ∂∂u = globdat.acce[:]

    globdat.time  += Δt
    updateTimeDependentEssentialBoundaryCondition!(domain, globdat)
    fext = getExternalForce(domain, globdat)

    u += Δt*∂u + 0.5*Δt*Δt*∂∂u
    ∂u += 0.5*Δt * ∂∂u

    domain.state[domain.eq_to_dof] = u[:]
    fint  = assembleInternalForce( globdat, domain, Δt)
    ∂∂up = globdat.M\(fext - fint)

    ∂u += 0.5 * Δt * ∂∂up

    globdat.Dstate = globdat.state[:]
    globdat.state = u[:]
    globdat.velo = ∂u[:]
    globdat.acce = ∂∂up[:]

    commitHistory(domain)
    updateStates!(domain, globdat)

    return globdat, domain
end

@doc raw"""
    ImplicitStaticSolver(globdat::GlobalData, domain::Domain;
        N::Int64 = 10, ε::Float64 = 1.e-6, maxiterstep::Int64=100)

Solves
$$K(u) = F$$
using the incremental load method. Specifically, at step $i$, we solve
```math
f^{int}(u_i) = \frac{i}{N} f^{ext}
```

- `globdat`, GlobalData

- `domain`, Domain

- `N`, an integer, load stepping steps

- `ε`, Float64, absolute error for Newton convergence

-  `maxiterstep`, Int64, maximum iteration number for Newton convergence
"""
function ImplicitStaticSolver(globdat::GlobalData, domain::Domain;
        N::Int64 = 10, ε::Float64 = 1.e-6, maxiterstep::Int64=100)
    assembleMassMatrix!(globdat, domain)
    fext = getExternalForce(domain, globdat)

    for iterstep = 1:N

        # Newton's method
        Newtoniterstep, Newtonconverge = 0, false

        while  !Newtonconverge

            Newtoniterstep += 1

            fint, stiff = assembleStiffAndForce( globdat, domain )

            res = fint - iterstep/N * fext

            Δstate = stiff\res

            globdat.state -= Δstate
            # @show Newtoniterstep, norm(res)
            if (norm(res) < ε  || Newtoniterstep > maxiterstep)
                if Newtoniterstep > maxiterstep
                    @error "$Newtoniterstep Newton iteration does not converge"
                end
                Newtonconverge = true
            end
            updateStates!(domain, globdat)
        end
        commitHistory(domain)
        globdat.Dstate = copy(globdat.state)
    end
    globdat, domain
end

@doc raw"""
    LinearStaticSolver(globdat::GlobalData, domain::Domain)

Solves the linear static problem
$$\begin{aligned}\div \sigma &= f\\ \sigma = H \epsilon \end{aligned}$$

"""
function LinearStaticSolver(globaldata::GlobalData, domain::Domain)
    @assert globaldata.time ≈ 0.0
    fext = getExternalForce(domain, globaldata)
    fint, stiff = assembleStiffAndForce( globaldata, domain, 0.0)
    res = fext - fint
    globaldata.state = stiff\res
    updateStates!(domain, globaldata)
    return globaldata, domain
end

@doc raw"""
    GeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64,
    ρ::Float64 = 0.0, ε::Float64 = 1e-8, ε0::Float64 = 1e-8, maxiterstep::Int64=100,
    η::Float64 = 1.0, failsafe::Bool = false, verbose::Bool = false)

Implicit solver for
$$Ma  + f_{int}(u) = fext$$
Here ``a``, ``v``, ``u`` are acceleration, velocity and displacement respectively.

- `ρ`: controls the damping effect of the α-scheme, ρ∈[0,1], ρ=1 corresponds to the maximum damping
- `ε`: Float64, absolute error for Newton convergence
- `ε0`: Float64, relative error for Newton convergence
- `max_iter`: Int64, maximum iteration number for Newton convergence
- `η`: Float64, Newton step size at the first iteration
- `failsafe`: Bool, if failsafe is true, when the Newton fails to converge,
              revert back, and return false

The nonlinear $\alpha$
```math
u_{n+1} = u_n + dtv_n + dt^2/2 ((1 - 2\beta)a_n + 2\beta a_{n+1})
v_{n+1} = v_n + dt((1 - \gamma)a_n + \gamma a_{n+1})
2\beta = 0.5*(1 - αm + αf)^2
\gamma = 0.5 - \alpha_m + \alpha_f
```

```math
a_{n+1-\alpha_m} = (1-\alpha_m)a_{n+1} + \alpha_m a_{n}
v_{n+1-\alpha_f} = (1-\alpha_f)v_{n+1} + \alpha_f v_{n}
u_{n+1-\alpha_f} = (1-\alpha_f)u_{n+1} + \alpha_f u_{n}
M a_{n+1-\alpha_m} + f^{int}(u_{n+1-\alpha_f}) = f^{ext}_{n+1-\alpha_f}
```

'a_{n+1}' is solved by

```math
M ((1-\alpha_m)a_{n+1} + \alpha_m a_{n})
+ f^{int}((1-\alpha_f)(u_n + dtv_n + dt^2/2 ((1 - 2\beta)a_n + 2\beta a_{n+1}))) + \alpha_f u_{n})
= f^{ext}_{n+1-\alpha_f}
```

As for ```\alpha_m``` and ``\alpha_f``
```math
\alpha_m = (2\rho_{\infty} - 1)/(\rho_{\infty} + 1)
\alpha_f = \rho_{\infty}/(\rho_{\infty} + 1)
```

use the current states `a`, `v`, `u`, `time` in globdat, and update these stetes to next time step
update domain history, when failsafe is true, and Newton's solver fails, nothing will be changed.

You need to call SolverInitial! before the first time step, if f^{ext}_0 != 0.
SolverInitial! updates a_0 in the globdat.acce
a_0 = M^{-1}(- f^{int}(u_0) + f^{ext}_0)

We assume globdat.acce[:] = a_0 and so far initialized to 0
We also assume the external force is conservative (it does not depend on the current deformation)
"""
function GeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64;
        ρ::Float64 = 0.0, ε::Float64 = 1e-8, ε0::Float64 = 1e-8, maxiterstep::Int64=100,
        η::Float64 = 1.0, failsafe::Bool = false, verbose::Bool = false)
    assembleMassMatrix!(globdat, domain)
    @assert 0<=ρ<=1
    αm = (2ρ-1)/(1+ρ)
    αf = ρ/(1+ρ)
    local res0

    β2 = 0.5*(1 - αm + αf)^2
    γ = 0.5 - αm + αf

    # compute solution at uⁿ⁺¹
    failsafe_time = copy(globdat.time)
    globdat.time  += (1 - αf)*Δt

    # domain.Dstate = uⁿ
    failsafe_Dstate = copy(domain.Dstate)
    failsafe_state = copy(domain.state)
    domain.Dstate = domain.state[:]

    updateTimeDependentEssentialBoundaryCondition!(domain, globdat)
    M = globdat.M

    ∂∂u = globdat.acce[:] #∂∂uⁿ
    u = globdat.state[:]  #uⁿ
    ∂u  = globdat.velo[:] #∂uⁿ

    fext = getExternalForce(domain, globdat)
    # fext = zeros(domain.neqs)
    ∂∂up = ∂∂u[:]

    Newtoniterstep, Newtonconverge = 0, false

    norm0 = Inf

    while !Newtonconverge

        Newtoniterstep += 1

        domain.state[domain.eq_to_dof] = (1 - αf)*(u + Δt*∂u + 0.5 * Δt * Δt * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u
        fint, stiff = assembleStiffAndForce( globdat, domain, Δt)
        res = M * (∂∂up *(1 - αm) + αm*∂∂u)  + fint - fext
        if Newtoniterstep==1
            res0 = res
        end
        A = M*(1 - αm) + (1 - αf) * 0.5 * β2 * Δt^2 * stiff
        Δ∂∂u = A\res

        while η * norm(Δ∂∂u) > norm0
            η /= 2.0
            verbose && (@info "η", η)
        end
        ∂∂up -= η*Δ∂∂u

        if maxiterstep==1
            break
        end
        verbose && println("$Newtoniterstep/$maxiterstep, abs = $(norm(res)), rel = $(norm(res)/norm(res0))")
        if (norm(res)< ε || norm(res)< ε0*norm(res0) ||Newtoniterstep > maxiterstep)
            if Newtoniterstep > maxiterstep
                # Newton method does not converge
                if failsafe
                    globdat.time = failsafe_time
                    domain.state = failsafe_state[:]
                    domain.Dstate = failsafe_Dstate[:]
                    return false
                end
                # When failsafe is not on, test the gradient
                function f(∂∂up)
                    domain.state[domain.eq_to_dof] = (1 - αf)*(u + Δt*∂u + 0.5 * Δt * Δt * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u
                    fint, stiff = assembleStiffAndForce( globdat, domain )
                    fint, (1 - αf) * 0.5 * β2 * Δt^2 * stiff
                end
                gradtest(f, ∂∂up)
                @warn("Newton iteration cannot converge $(norm(res))"); Newtonconverge = true
            else
                Newtonconverge = true
                verbose && printstyled("[Newmark] Newton converged $Newtoniterstep\n", color=:green)
            end
        end

        η = min(1.0, 2η)
        norm0 = norm(Δ∂∂u)
    end

    #update globdat to the next time step
    globdat.Dstate = globdat.state[:]
    globdat.state += Δt * ∂u + Δt^2/2 * ((1 - β2) * ∂∂u + β2 * ∂∂up)
    globdat.velo += Δt * ((1 - γ) * ∂∂u + γ * ∂∂up)
    globdat.acce = ∂∂up[:]
    globdat.time  += αf*Δt

    #commit history in domain
    commitHistory(domain)
    updateStates!(domain, globdat)
    if options.save_history>=2
        fint, stiff = assembleStiffAndForce( globdat, domain, Δt)
        push!(domain.history["fint"], fint)
        push!(domain.history["fext"], fext)
    end
    if options.save_history>=1
        push!(domain.history["time"], [globdat.time])
    end
    return globdat, domain

end

function LinearGeneralizedAlphaSolverInit!(globdat::GlobalData, domain::Domain, Δt::Float64;
    ρ::Float64 = 0.0)
    assembleMassMatrix!(globdat, domain)
    @assert 0<=ρ<=1
    αm = (2ρ-1)/(1+ρ)
    αf = ρ/(1+ρ)
    β2 = 0.5*(1 - αm + αf)^2
    γ = 0.5 - αm + αf
    M = globdat.M

    _, stiff = assembleStiffAndForce(domain, Δt)
    A = M*(1 - αm) + (1 - αf) * 0.5 * β2 * Δt^2 * stiff
    STORAGE["LinearGeneralizedAlphaSolverInit!"] = lu(A)
    nothing
end

"""
    LinearGeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64;
    ρ::Float64 = 0.0)

Solves the dynamic structural equation with time-independent coefficient matrix. This is useful for linear elasticity simulation.
"""
function LinearGeneralizedAlphaSolverStep(globdat::GlobalData, domain::Domain, Δt::Float64;
    ρ::Float64 = 0.0)
    if haskey(STORAGE, "LinearGeneralizedAlphaSolverInit")
        error("You must call LinearGeneralizedAlphaSolverInit! before LinearGeneralizedAlphaSolverStep")
    end
    A = STORAGE["LinearGeneralizedAlphaSolverInit!"]
    M = globdat.M
    @assert 0<=ρ<=1
    αm = (2ρ-1)/(1+ρ)
    αf = ρ/(1+ρ)
    β2 = 0.5*(1 - αm + αf)^2
    γ = 0.5 - αm + αf

    # compute solution at uⁿ⁺¹
    globdat.time  += (1 - αf)*Δt

    # domain.Dstate = uⁿ
    domain.Dstate = domain.state[:]

    updateTimeDependentEssentialBoundaryCondition!(domain, globdat)

    ∂∂u = globdat.acce[:] #∂∂uⁿ
    u = globdat.state[:]  #uⁿ
    ∂u  = globdat.velo[:] #∂uⁿ

    fext = getExternalForce(domain, globdat)
    ∂∂up = ∂∂u[:]

    domain.state[domain.eq_to_dof] = (1 - αf)*(u + Δt*∂u + 0.5 * Δt * Δt * ((1 - β2)*∂∂u + β2*∂∂up)) + αf*u

    fint = assembleInternalForce(domain, Δt)
    res = M * (∂∂up *(1 - αm) + αm*∂∂u)  + fint - fext
    Δ∂∂u = A\res
    ∂∂up -= Δ∂∂u

    #update globdat to the next time step
    globdat.Dstate = globdat.state[:]
    globdat.state += Δt * ∂u + Δt^2/2 * ((1 - β2) * ∂∂u + β2 * ∂∂up)
    globdat.velo += Δt * ((1 - γ) * ∂∂u + γ * ∂∂up)
    globdat.acce = ∂∂up[:]
    globdat.time  += αf*Δt

    #commit history in domain
    commitHistory(domain)
    updateStates!(domain, globdat)
    if options.save_history>=2
        fint, stiff = assembleStiffAndForce( globdat, domain, Δt)
        push!(domain.history["fint"], fint)
        push!(domain.history["fext"], fext)
    end
    if options.save_history>=1
        push!(domain.history["time"], [globdat.time])
    end
    return globdat, domain

end


text:
"`x' = contract(d, x)` implies x ∈ d`"
function contract end

"`codomain(f, t)` -- codomain of method `f(::t...)`"
function codomain end

# zt - fix me
contract(i::Interval, x) = softmax()
contract(::Type{Any}, x) = x

# zt - whatabout closed vs open interval
"Interval [l, u]"
struct Interval{T}
  l::T
  u::T
end
l..u = Interval(l, u)

codomain(::typeof(sin), t) =  -1..1
codomain(::typeof(cos), t) =  -1..1
codomain(::typeof(sqrt), t) = -1..1

# Unless otherwise specified, assume a function is unconstrained
codomain(f, t) = Any

# zt - something a bit odd about using Any for codmain of arbitray function#
# zt - maybe its better to think of it as something else

# zt - we only want to bother contracting if the function is contractable
invertinvoke(traits::trait(Contract), f, t, y, ϕ) =
  invertinvoke(f, t, contract(traits, codomain(f, t), y), ϕ)

# move this into
invertinvoke(traits, f, t, y, ϕ) =
  invertinvoke(t, f, t, y, ϕ)

# Handle this case, ...
function invertinvoke(::typeof(dupl), t, y, ϕ::trait(contract))
end

# TODO
# How to specify different contractions
# How to have a better solution than tagtraits here
# -- Could have Trait{}


text:
@compat abstract type Link end     # Link types define linkfun!, linkinv!, and mueta!

type CauchitLink <: Link end
type CloglogLink  <: Link end
type IdentityLink <: Link end
type InverseLink  <: Link end
type LogitLink <: Link end
type LogLink <: Link end
type ProbitLink <: Link end
type SqrtLink <: Link end

linkfun(::CauchitLink, μ) = tan(pi * (μ - oftype(μ, 0.5)))
linkinv(::CauchitLink, η) = oftype(η, 0.5) + atan(η) / pi
mueta(::CauchitLink, η) = one(η) / (pi * (one(η) + abs2(η)))

linkfun(::CloglogLink, μ) = log(-log1p(-μ))
linkinv(::CloglogLink, η) = -expm1(-exp(η))
mueta(::CloglogLink, η) = exp(η) * exp(-exp(η))

linkfun(::IdentityLink, μ) = μ
linkinv(::IdentityLink, η) = η
mueta(::IdentityLink, η) = 1

linkfun(::InverseLink, μ) = inv(μ)
linkinv(::InverseLink, η) = inv(η)
mueta(::InverseLink, η) = -inv(abs2(η))

linkfun(::LogitLink, μ) = logit(μ)
linkinv(::LogitLink, η) = logistic(η)
function mueta(::LogitLink, η)
    e = exp(-abs(η))
    f = one(η) + e
    return e / (f * f)
end

linkfun(::LogLink, μ) = log(μ)
linkinv(::LogLink, η) = exp(η)
mueta(::LogLink, η) = exp(η)

linkfun(::ProbitLink, μ) = -sqrt2 * erfcinv(2μ)
linkinv(::ProbitLink, η) = erfc(-η / sqrt2) / 2
mueta(::ProbitLink, η) = exp(-abs2(η) / 2) / sqrt2π

linkfun(::SqrtLink, μ) = sqrt(μ)
linkinv(::SqrtLink, η) = abs2(η)
mueta(::SqrtLink, η) = 2η

canonicallink(::Bernoulli) = LogitLink()
canonicallink(::Binomial) = LogitLink()
canonicallink(::Gamma) = InverseLink()
canonicallink(::Normal) = IdentityLink()
canonicallink(::Poisson) = LogLink()

# For the "odd" link functions we evaluate the linear predictor such that mu is closest to zero where the precision is higher
function glmvar(::Union{Bernoulli,Binomial}, link::Union{CauchitLink,InverseLink,LogitLink,ProbitLink}, μ, η)
    μ = linkinv(link, ifelse(η < 0, η, -η))
    μ * (1 - μ)
end
glmvar(::Union{Bernoulli,Binomial}, ::Link, μ, η) = μ * (1 - μ)
glmvar(::Gamma, ::Link, μ, η) = abs2(μ)
glmvar(::Normal, ::Link, μ, η) = 1
glmvar(::Poisson, ::Link, μ, η) = μ

mustart(::Bernoulli, y, wt) = (y + oftype(y, 0.5)) / 2
mustart(::Binomial, y, wt) = (wt * y + oftype(y, 0.5)) / (wt + one(y))
mustart(::Gamma, y, wt) = y == 0 ? oftype(y, 0.1) : y
mustart(::Normal, y, wt) = y
mustart(::Poisson, y, wt) = y + oftype(y, 0.1)

function devresid(::Bernoulli, y, μ)
    if y == 1
        return -2 * log(μ)
    elseif y == 0
        return -2 * log1p(-μ)
    end
    throw(ArgumentError("y should be 0 or 1 (got $y)"))
end
function devresid(::Binomial, y, μ)
    if y == 1
        return -2 * log(μ)
    elseif y == 0
        return -2 * log1p(-μ)
    else
        return 2 * (y * (log(y) - log(μ)) + (1 - y)*(log1p(-y) - log1p(-μ)))
    end
end
devresid(::Gamma, y, μ) = -2 * (log(y / μ) - (y - μ) / μ)
devresid(::Normal, y, μ) = abs2(y - μ)
devresid(::Poisson, y, μ) = 2 * (xlogy(y, y / μ) - (y - μ))

# Whether a dispersion parameter has to be estimated for a distribution
dispersion_parameter(::Union{Bernoulli, Binomial, Poisson}) = false
dispersion_parameter(::Union{Gamma, Normal, InverseGaussian}) = true

# Log-likelihood for an observation
loglik_obs(::Bernoulli, y, μ, wt, ϕ) = wt*logpdf(Bernoulli(μ), y)
loglik_obs(::Binomial, y, μ, wt, ϕ) = logpdf(Binomial(Int(wt), μ), Int(y*wt))
loglik_obs(::Gamma, y, μ, wt, ϕ) = wt*logpdf(Gamma(1/ϕ, μ*ϕ), y)
loglik_obs(::Normal, y, μ, wt, ϕ) = wt*logpdf(Normal(μ, sqrt(ϕ)), y)
loglik_obs(::Poisson, y, μ, wt, ϕ) = wt*logpdf(Poisson(μ), y)


text:
#Broadcasting machinery
#######################

#This is directly from the Broadcasting interface docs
struct GBVectorStyle <: Broadcast.AbstractArrayStyle{1} end
struct GBMatrixStyle <: Broadcast.AbstractArrayStyle{2} end
Base.BroadcastStyle(::Type{<:GBVector}) = GBVectorStyle()
Base.BroadcastStyle(::Type{<:GBMatrix}) = GBMatrixStyle()
Base.BroadcastStyle(::Type{<:Transpose{T, <:GBMatrix} where T}) = GBMatrixStyle()
#
GBVectorStyle(::Val{0}) = GBVectorStyle()
GBVectorStyle(::Val{1}) = GBVectorStyle()
GBVectorStyle(::Val{2}) = GBMatrixStyle()
GBVectorStyle(::Val{N}) where N = Broadcast.DefaultArrayStyle{N}()
GBMatrixStyle(::Val{0}) = GBMatrixStyle()
GBMatrixStyle(::Val{1}) = GBMatrixStyle()
GBMatrixStyle(::Val{2}) = GBMatrixStyle()
GBMatrixStyle(::Val{N}) where N = Broadcast.DefaultArrayStyle{N}()
Broadcast.BroadcastStyle(::GBMatrixStyle, ::GBVectorStyle) = GBMatrixStyle()

function Base.similar(
    bc::Broadcast.Broadcasted{GBMatrixStyle},
    ::Type{ElType}
) where {ElType}
    return GBMatrix{ElType}(axes(bc))
end

function Base.similar(
    bc::Broadcast.Broadcasted{GBVectorStyle},
    ::Type{ElType}
) where {ElType}
    return GBVector{ElType}(axes(bc))
end

#Find the modifying version of a function.
modifying(::typeof(mul)) = mul!
modifying(::typeof(eadd)) = eadd!
modifying(::typeof(emul)) = emul!

@inline function Base.copy(bc::Broadcast.Broadcasted{GBMatrixStyle})
    f = bc.f
    l = length(bc.args)
    if l == 1
        return map(f, first(bc.args))
    else
        left = first(bc.args)
        right = last(bc.args)
        if left isa Broadcast.Broadcasted
            left = copy(left)
        end
        if right isa Broadcast.Broadcasted
            right = copy(right)
        end
        if left isa GBArray && right isa GBArray

            add = defaultadd(f)
            return add(left, right, f)
        else
            return map(f, left, right)
        end
    end
end

@inline function Base.copy(bc::Broadcast.Broadcasted{GBVectorStyle})
    f = bc.f
    l = length(bc.args)
    if l == 1
        return map(f, first(bc.args))
    else
        left = first(bc.args)
        right = last(bc.args)
        if left isa Broadcast.Broadcasted
            left = copy(left)
        end
        if right isa Broadcast.Broadcasted
            right = copy(right)
        end
        if left isa GBArray && right isa GBArray
            add = defaultadd(f)
            return add(left, right, f)
        else
            return map(f, left, right)
        end
    end
end

#Elementwise Broadcasts
#####################
#function Base.broadcasted(::typeof(*), A::GBArray, B::GBArray)
#    emul(A, B, BinaryOps.TIMES)
#end
#
#function Base.broadcasted(::typeof(+), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.PLUS)
#end
#
#function Base.broadcasted(::typeof(-), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.MINUS)
#end
#
##TODO: fix tricky gotchas, this will do type-specific (ie sometimes integer) division.
#function Base.broadcasted(::typeof(/), A::GBArray, B::GBArray)
#    emul(A, B, BinaryOps.DIV)
#end
#
##TODO: fix tricky gotchas, this will do type-specific (ie sometimes integer) division.
#function Base.broadcasted(::typeof(\), A::GBArray, B::GBArray)
#    emul(A, B, BinaryOps.RDIV)
#end
#
#function Base.broadcasted(::typeof(==), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.EQ)
#end
#
#function Base.broadcasted(::typeof(!=), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.NE)
#end
#
#function Base.broadcasted(::typeof(<), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.LT)
#end
#
#function Base.broadcasted(::typeof(>), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.GT)
#end
#
#function Base.broadcasted(::typeof(<=), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.LE)
#end
#
#function Base.broadcasted(::typeof(>=), A::GBArray, B::GBArray)
#    eadd(A, B, BinaryOps.GE)
#end
#
#function Base.broadcasted(::typeof(^), A::GBArray, B::GBArray)
#    emul(A, B, BinaryOps.POW)
#end
#
##Map Broadcasts
########################
#function Base.broadcasted(::typeof(*), u::GBArray, x::valid_union)
#    map(BinaryOps.TIMES, u, x)
#end
#function Base.broadcasted(::typeof(*), x::valid_union, u::GBArray)
#    map(BinaryOps.TIMES, x, u)
#end
#
#function Base.broadcasted(::typeof(/), u::GBArray, x::valid_union)
#    map(BinaryOps.DIV, u, x)
#end
#function Base.broadcasted(::typeof(/), x::valid_union, u::GBArray)
#    map(BinaryOps.DIV, x, u;)
#end
#
#function Base.broadcasted(::typeof(^), u::GBArray, x::valid_union)
#    map(BinaryOps.POW, u, x)
#end
#function Base.broadcasted(::typeof(^), x::valid_union, u::GBArray)
#    map(BinaryOps.POW, x, u)
#end
#
#
#function Base.broadcasted(::typeof(-), u::GBArray, x::valid_union)
#    map(BinaryOps.MINUS, u, x)
#end
#function Base.broadcasted(::typeof(-), x::valid_union, u::GBArray)
#    map(BinaryOps.MINUS, x, u)
#end
#function Base.broadcasted(::typeof(+), u::GBArray, x::valid_union)
#    map(BinaryOps.PLUS, u, x)
#end
#function Base.broadcasted(::typeof(+), x::valid_union, u::GBArray)
#    map(BinaryOps.PLUS, x, u)
#end


text:
# This file was generated, do not modify it. # hide
mach = machine(mdl, X, y)
fit!(mach)

fp = fitted_params(mach)
coefs = fp.coefs
intercept = fp.intercept
for (name, val) in coefs
    println("$(rpad(name, 8)):  $(round(val, sigdigits=3))")
end
println("Intercept: $(round(intercept, sigdigits=3))")


text:
using Base.Meta

# # the "transformable trait"... I and O are the shapes (num dims) of the input and output
# abstract TransformationType{I, O}
#     immutable Transformable{I,O} <: TransformationType{I,O} end
#     immutable NotTransformable{I,O} <: TransformationType{I,O} end

# transformation_type(t) = error()

# TODO:
# - add input_size/output_size, tranformation/Transformation to LearnBase
# - add traits to LearnBase?

export
    transformation,
    is_learnable,
    Center,
    Affine

# sigmoid(x) = 1 ./ (1 .+ exp(-x))

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

abstract AbstractTransformation{I,O}

# we can wrap any object which implements the interface, and it becomes a transformation
immutable Transformation{T,I,O} <: AbstractTransformation{I,O}
    t::T
end

transformation(x) = Transformation{typeof(x), length(input_size(x)), length(output_size(x))}(x)
Base.show{T,I,O}(io::IO, t::Transformation{T,I,O}) = print(io, "T{$I-->$O}{$(t.t)}")

# just pass it through.  we only need to define `transform!` now
transform(t::Transformation, x)     = (y = zeros(output_size(t.t)); transform!(y, t.t, x))
transform!(y, t::Transformation, x) = transform!(y, t.t, x)

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

# # the "learnable trait"... are there parameters to fit/learn?
# abstract LearnableTrait
#     immutable Learnable <: LearnableTrait end
#     immutable NotLearnable <: LearnableTrait end

# # by default, it can't be learned
# learnable(x) = NotLearnable()

# are there parameters to be learned in this tranformation?
is_learnable(x) = false

# can we take a partial derivative with respect to those parameters?
is_differentiable(x) = false

# the non-bang version must have an empty constructor
# learn{T}(::Type{T}, args...)        = learn!(transformation(T()), args...)
# learn!(t::Transformation, args...)     = learn!(learnable(t.t), t.t, args...)
# learn!(::NotLearnable, t, args...)  = error() # TODO: maybe just do nothing?
# learn!(::Learnable, t, args...)     = learn!(t, args...)

# build a new object and then learn
learn{T}(::Type{T}, args...)        = learn!(transformation(T()), args...)
learn!(t::Transformation, args...)  = learn!(t.t, args...)

# learn!(t::Transformation, args...)  = learn!(Val{is_learnable(t.t)}, t.t, args...)
# learn!(::Type{Val{false}}, args...)    = error("Called learn but is_learnable($(typeof(t.t))) == false") # TODO: maybe just do nothing?
# learn!(::Type{Val{true}}, args...)     = learn!(args...)

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

is_stochastic(x) = false

# rand(t::Transformation, args...) = rand(Val{is_stochastic(t.t)}, t.t, args...)
# rand(::Type{Val{false}}, t, args...) = error("Called rand but is_stochastic($(typeof(t.t))) == false")

# # the "randomness trait"... is it static or stochastic?
# abstract RandomnessTrait
#     immutable Static <: RandomnessTrait end
#     immutable Stochastic <: RandomnessTrait end

# # by default, everything is static
# randomness(x) = Static()

# # pass through to RandomnessTrait definition
# rand(t::Transformation, args...)    = rand(randomness(t.t), t.t, args...)
# rand(::Static, t, args...)          = error()
# rand(::Stochastic, t, args...)      = rand(t, args...)

# -------------------------------------------------------
# some sample transformations
# notice how we don't need to inherit from a common type tree?

abstract DemoTransform

# we can set traits for a bunch of types like this.  could also use a Union
learnable(::DemoTransform) = Learnable()

immutable Center{A<:AbstractArray} <: DemoTransform
    mu::A
end

input_size(t::Center)               = size(t.mu)
output_size(t::Center)              = size(t.mu)

transform!(y, t::Center, x)         = (y[:] = x - t.mu)
learn!(t::Center, x)                = (mean!(t.mu, x))  # recomputes for full dataset

immutable Affine{W<:AbstractMatrix, B<:AbstractVector} <: DemoTransform
    w::W
    b::B
end

input_size(t::Affine)               = (size(t.w,2),)
output_size(t::Affine)              = size(t.b)

transform!(y, t::Affine, x)         = (y[:] = t.w * x + t.b)

# -------------------------------------------------------
# Distributions are generating transformations

import Distributions:
    Distribution,
    Normal,
    Bernoulli,
    NegativeBinomial,
    Poisson

is_stochastic(::Distribution) = true

## Exported types ##

# export
#     # static transformations
#     IdentityTransformation,
#     LogTransformation,
#     ExpTransformation,
#     LogisticTransformation,
#     LogitTransformation,
#     ScaleTransformation,
#     ShiftTransformation,

#     # stochastic transformations
#     GaussianTransformation,
#     PoissonTransformation,
#     BernoulliTransformation,
#     NegativeBinomialTransformation,
#     GeneralizedLinearTransformation


text:
"""
newton(f,dfdx,x1)

Use Newton's method to find a root of `f` starting from `x1`, where
`dfdx` is the derivative of `f`. Returns a vector of root estimates.
"""
function newton(f,dfdx,x1)
    # Operating parameters.
    funtol = 100*eps();  xtol = 100*eps();  maxiter = 40;

    x = [x1]
    y = f(x1)
    dx = Inf   # for initial pass below
    k = 1

    while (abs(dx) > xtol) && (abs(y) > funtol) && (k < maxiter)
        dydx = dfdx(x[k])
        dx = -y/dydx            # Newton step
        push!(x,x[k]+dx)        # append new estimate

        k = k+1
        y = f(x[k])
    end

    if k==maxiter
        @warn "Maximum number of iterations reached."
    end

    return x
end

"""
secant(f,x1,x2)

Use the secant method to find a root of `f` starting from `x1` and
`x2`. Returns a vector of root estimates.
"""
function secant(f,x1,x2)
    # Operating parameters.
    funtol = 100*eps();  xtol = 100*eps();  maxiter = 40;

    x = [x1,x2]
    y1 = f(x1); y2 = 100;
    dx = Inf   # for initial pass below
    k = 2

    while (abs(dx) > xtol) && (abs(y2) > funtol) && (k < maxiter)
        y2 = f(x[k])
        dx = -y2 * (x[k]-x[k-1]) / (y2-y1)   # secant step
        push!(x,x[k]+dx)        # append new estimate

        k = k+1
        y1 = y2    # current f-value becomes the old one next time
    end

    if k==maxiter
        @warn "Maximum number of iterations reached."
    end

    return x
end

"""
newtonsys(f,jac,x1)

Use Newton's method to find a root of a system of equations,
starting from `x1`. The functions `f` and `jac should return the
residual vector and the Jacobian matrix, respectively. Returns
history of root estimates as a vector of vectors.
"""
function newtonsys(f,jac,x1)
    # Operating parameters.
    funtol = 1000*eps();  xtol = 1000*eps();  maxiter = 40;

    x = [float(x1)]
    y,J = f(x1),jac(x1)
    dx = Inf   # for initial pass below
    k = 1

    while (norm(dx) > xtol) && (norm(y) > funtol) && (k < maxiter)
        dx = -(J\y)             # Newton step
        push!(x,x[k] + dx)    # append to history
        k += 1
        y,J = f(x[k]),jac(x[k])
    end

    if k==maxiter
        @warn "Maximum number of iterations reached."
    end

    return x
end

"""
fdjac(f,x0,y0)

Compute a finite-difference approximation of the Jacobian matrix for
`f` at `x0`, where `y0`=`f(x0)` is given.
"""
function fdjac(f,x0,y0)

delta = sqrt(eps())   # FD step size
m,n = length(y0),length(x0)
if n==1
    J = (f(x0+delta) - y0) / delta
else
    J = zeros(m,n)
    In = I(n)
    for j = 1:n
        J[:,j] = (f(x0 + delta*In[:,j]) - y0) / delta
    end
end

return J
end

"""
levenberg(f,x1,tol)

Use Levenberg's quasi-Newton iteration to find a root of the system
`f`, starting from `x1`, with `tol` as the stopping tolerance in
both step size and residual norm. Returns root estimates as a
matrix, one estimate per column.
"""
function levenberg(f,x1,tol=1e-12)

# Operating parameters.
ftol = tol;  xtol = tol;  maxiter = 40;

x = zeros(length(x1),maxiter)
x = [float(x1)]
fk = f(x1)
k = 1;  s = Inf;
Ak = fdjac(f,x1,fk)   # start with FD Jacobian
jac_is_new = true

lambda = 10;
while (norm(s) > xtol) && (norm(fk) > ftol) && (k < maxiter)
    # Compute the proposed step.
    B = Ak'*Ak + lambda*I
    z = Ak'*fk
    s = -(B\z)

    xnew = x[k] + s
    fnew = f(xnew)

    # Do we accept the result?
    if norm(fnew) < norm(fk)    # accept
        y = fnew - fk
        push!(x,xnew)
        fk = fnew
        k += 1

        lambda = lambda/10   # get closer to Newton
        # Broyden update of the Jacobian.
        Ak = Ak + (y-Ak*s)*(s'/(s'*s))
        jac_is_new = false
    else                       # don't accept
        # Get closer to steepest descent.
        lambda = 4lambda
        # Re-initialize the Jacobian if it's out of date.
        if !jac_is_new
            Ak = fdjac(f,x[k],fk)
            jac_is_new = true
        end
    end
end

if (norm(fk) > 1e-3)
    @warn "Iteration did not find a root."
end

return x
end


text:
module NorvigsTrieSpellchecker
using DataStructures

include("spellcheker.jl")

export correction, load_text

end # module


text:
using FIFOStreams
using Test

@static if VERSION >= v"1.3"
    using diffutils_jll
else
    _diff(f) = Sys.isunix() ? f("diff") : @test_broken false
end

a = read(joinpath(@__DIR__, "a"), String)
b = read(joinpath(@__DIR__, "b"), String)
diff_ab = read(joinpath(@__DIR__, "diff_ab"), String)

for T in [(Sys.isunix() ? [UnixFIFOStream] : [])..., FallbackFIFOStream]

@testset "diffutils - $T" begin
    _diff() do diff
        s = FIFOStreamCollection(T, 2)
        io = IOBuffer()
        attach(s, pipeline(ignorestatus(`$diff $(path(s, 1)) $(path(s, 2))`); stdout=io))
        s1, s2 = s
        print(s1, a)
        print(s2, b)
        close(s)
        out = String(take!(io))
        @test out == diff_ab
    end
end

s = FIFOStream()
# just test that this doesn't error
@test close(s) === nothing

# doctests are broken on Windows, let's not bother
Sys.iswindows() && continue

using Documenter
DocMeta.setdocmeta!(
    FIFOStreams,
    :DocTestSetup,
    :(using FIFOStreams; const FIFOStream = $T);
    recursive=true,
)
@testset "doctests - $T" begin
    doctest(FIFOStreams)
end

end


text:
export tap

import Base: show

"""
    tap(tapFn::F) where { F <: Function }

Creates a tap operator, which performs a side effect
for every emission on the source Observable, but return an Observable that is identical to the source.

# Arguments
- `tapFn::Function`: side-effect tap function with `(data) -> Nothing` signature

# Producing

Stream of type `<: Subscribable{L}` where `L` refers to type of source stream

# Examples
```jldoctest
using Rocket

source = from([ 1, 2, 3 ])
subscribe!(source |> tap((d) -> println("In tap: \$d")), logger())
;

# output

In tap: 1
[LogActor] Data: 1
In tap: 2
[LogActor] Data: 2
In tap: 3
[LogActor] Data: 3
[LogActor] Completed

```

See also: [`tap_on_subscribe`](@ref), [`tap_on_complete`](@ref), [`logger`](@ref)
"""
tap(tapFn::F) where { F <: Function } = TapOperator{F}(tapFn)

struct TapOperator{F} <: InferableOperator
    tapFn :: F
end

function on_call!(::Type{L}, ::Type{L}, operator::TapOperator{F}, source) where { L, F }
    return proxy(L, source, TapProxy{F}(operator.tapFn))
end

operator_right(operator::TapOperator, ::Type{L}) where L = L

struct TapProxy{F} <: ActorProxy
    tapFn :: F
end

actor_proxy!(::Type{L}, proxy::TapProxy{F}, actor::A) where { L, A, F } = TapActor{L, A, F}(proxy.tapFn, actor)

struct TapActor{L, A, F} <: Actor{L}
    tapFn :: F
    actor :: A
end

on_next!(actor::TapActor{L}, data::L) where L = begin actor.tapFn(data); next!(actor.actor, data) end
on_error!(actor::TapActor, err)               = error!(actor.actor, err)
on_complete!(actor::TapActor)                 = complete!(actor.actor)

Base.show(io::IO, ::TapOperator)         = print(io, "TapOperator()")
Base.show(io::IO, ::TapProxy)            = print(io, "TapProxy()")
Base.show(io::IO, ::TapActor{L}) where L = print(io, "TapActor($L)")


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))
Pkg.instantiate()

# ## Setup
using Dojo
using Random
using LinearAlgebra
using Plots

include("methods/utils.jl")
include("methods/data_processing.jl")
include("methods/quasi_newton.jl")
include("hardware_examples/methods.jl")
include("hardware_examples/visualization.jl")

# ## Open visualizer
vis= Visualizer()
open(vis)

# ## Mechanism
S = 7
timestep= 1/148 * S
gscaled = -9.81*20
mech = get_mechanism(:block,
	timestep=timestep,
	gravity=gravityscaled,
	friction_coefficient=0.2,
	radius=0.00,
	side=0.50,
	mode=:box);

# ## Simualate
initialize!(mech, :block,
	x=[0.0, -1.0, 1.0],
	v=[0.0, 2.0, 1.0],
	ω=[2.0, 5.0, 10.0])
storage = simulate!(mech, 5.0,
	record=true,
    opts=SolverOptions(btol=1e-6, rtol=1e-6, verbose=false))

# ## Visualize
visualize(mech, storage, vis=vis, show_contact=true)

################################################################################
## Generate & Save Dataset
################################################################################
generate_hardware_dataset(N=400, sleep_ratio=0.0002, S=S)
generate_hardware_dataset(N=400, sleep_ratio=0.0002, S=1)

################################################################################
## Load Dataset
################################################################################
params0, trajs0, pairs0 = open_dataset(:hardwarebox; N=400, S=S)
params1, trajs1, pairs1 = open_dataset(:hardwarebox; N=400, S=1)

pairs0
params0
data0 = params0[:data]

################################################################################
## Optimization Objective: Evaluation & Gradient
################################################################################
clean_loss(:block, pairs0, data0,
	n_sample=250,
	opts=SolverOptions(btol=3e-4, rtol=3e-4),
	timestep=timestep,
	gravity=gravityscaled)

[clean_loss(:block, pairs0, data0 + [i;zeros(55)],
	opts=SolverOptions(btol=3e-4, rtol=3e-4),
	timestep=timestep,
	gravity=gravityscaled)
	for i in Vector(-0.10:0.01:0.1)]
[clean_loss(:block, pairs0, data0 + [0;i;zeros(54)],
	opts=SolverOptions(btol=3e-4, rtol=3e-4),
	timestep=timestep,
	gravity=gravityscaled)
	for i in Vector(-0.10:0.01:0.1)]

plot(hcat([p[1][1:3] for p in pairs0]...)')
plot(hcat([p[1][4:6] for p in pairs0]...)')
plot(hcat([p[1][7:10] for p in pairs0]...)')
plot(hcat([p[1][11:13] for p in pairs0]...)')

################################################################################
## Optimization Algorithm: Quasi Newton:
## We learn a single coefficient of friction and a side length = 2 params
################################################################################
function d2data(d)
	friction_coefficient = d[1]
	side = d[2]
	data = [friction_coefficient; 0.0; 0.0; 0.0; +side; +side; -side;
			friction_coefficient; 0.0; 0.0; 0.0; +side; -side; -side;
			friction_coefficient; 0.0; 0.0; 0.0; -side; +side; -side;
			friction_coefficient; 0.0; 0.0; 0.0; -side; -side; -side;
			friction_coefficient; 0.0; 0.0; 0.0; +side; +side; +side;
			friction_coefficient; 0.0; 0.0; 0.0; +side; -side; +side;
			friction_coefficient; 0.0; 0.0; 0.0; -side; +side; +side;
			friction_coefficient; 0.0; 0.0; 0.0; -side; -side; +side;
			]
	return data
end
∇d2data = ForwardDiff.jacobian(d -> d2data(d), zeros(2))

d0 = [0.40, 0.50]
lower = [0.00, 0.05]
upper = [0.80, 1.50]

function f0(d; rot=0)
	return clean_loss(:block, pairs0, d2data(d),
		n_sample=200,
		timestep=timestep,
		gravity=gravityscaled,
		rot=rot,
		opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
end

function fgH0(d;
	rot=0)
	f, g, H = clean_loss(:block, pairs0, d2data(d),
		n_sample=200,
		timestep=timestep,
		gravity=gravityscaled,
		rot=rot,
		opts=SolverOptions(btol=3e-4, rtol=3e-4))
	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
end
dsol, Dsol = quasi_newton_solve(f0, fgH0, d0,
	iter=100,
	gtol=1e-8,
	ftol=1e-6,
	lower=lower,
	upper=upper,
	reg=1e-6,
	Δrot=20)

f0(d0)
f0(dsol)
d0
dsol #[0.171, 0.933]
## We can learn the coefficient of friction and the side dimenson of the cube
## form 35*0.40 seconds of recording. We use the simulator to evaluate the loss
## and its gradients by differentiating through the simulator. With gradient
## information we can use L-BFGS

################################################################################
## Visualization
################################################################################

using Plots; pyplot()
x=range(0.00,stop=0.80,length=10)
y=range(0.10,stop=1.50,length=10)
f(x,y) = log(10, clean_loss(:block, pairs0, d2data([x,y]), n_sample=50, timestep=timestep, gravity=gravityscaled,
	opts=SolverOptions(btol=3e-4, rtol=3e-4))[1])
plot(x,y,f,st=:surface,camera=(25,75))

################################################################################
## Optimization Algorithm: Quasi Newton:
## We learn a single coefficient of friction and a 8 contact locations [x,y,z] -> 25 params in total
################################################################################
function d2data(d)
	friction_coefficient = d[1]
	data = [friction_coefficient; 0;0;0; +d[2:4];
			friction_coefficient; 0;0;0; +d[5:7];
			friction_coefficient; 0;0;0; +d[8:10];
			friction_coefficient; 0;0;0; +d[11:13];
			friction_coefficient; 0;0;0; +d[14:16];
			friction_coefficient; 0;0;0; +d[17:19];
			friction_coefficient; 0;0;0; +d[20:22];
			friction_coefficient; 0;0;0; +d[23:25];
			]
	return data
end
∇d2data = ForwardDiff.jacobian(d -> d2data(d), zeros(25))
d0 = [0.40,
	+1.30, +1.40, -1.50,
	+1.50, -1.30, -1.10,
	-1.50, +1.30, -1.10,
	-1.50, -1.40, -1.30,
	+1.10, +1.20, +1.50,
	+1.20, -1.40, +1.50,
	-1.50, +1.40, +1.30,
	-1.10, -1.40, +1.50]
## d0 = [0.40,
## 	+0.50, +0.50, -0.50,
## 	+0.50, -0.50, -0.50,
## 	-0.50, +0.50, -0.50,
## 	-0.50, -0.50, -0.50,
## 	+0.50, +0.50, +0.50,
## 	+0.50, -0.50, +0.50,
## 	-0.50, +0.50, +0.50,
## 	-0.50, -0.50, +0.50]
lower = [0.00,
	+0.05, +0.05, -1.50,
	+0.05, -1.50, -1.50,
	-1.50, +0.05, -1.50,
	-1.50, -1.50, -1.50,
	+0.05, +0.05, +0.05,
	+0.05, -1.50, +0.05,
	-1.50, +0.05, +0.05,
	-1.50, -1.50, +0.05]
upper = [0.80,
	+1.50, +1.50, -0.05,
	+1.50, -0.05, -0.05,
	-0.05, +1.50, -0.05,
	-0.05, -0.05, -0.05,
	+1.50, +1.50, +1.50,
	+1.50, -0.05, +1.50,
	-0.05, +1.50, +1.50,
	-0.05, -0.05, +1.50]

function f0(d; rot=0, n_sample=50)
	return clean_loss(:block, pairs0, d2data(d), n_sample=n_sample, timestep=timestep, gravity=gravityscaled, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))[1]
end

function fgH0(d; rot=0, n_sample=50)
	f, g, H = clean_loss(:block, pairs0, d2data(d), n_sample=n_sample, timestep=timestep, gravity=gravityscaled, rot=rot, opts=SolverOptions(btol=3e-4, rtol=3e-4))
	return f, ∇d2data' * g, ∇d2data' * H * ∇d2data
end
dsol, Dsol = quasi_newton_solve(f0, fgH0, d0, iter=50, gtol=1e-8, ftol=1e-6,
	lower=lower, upper=upper, reg=1e+1, Δrot=500, n_sample0=20, Δn_sample=5, n_sample_max=1000)

f0(d0)
f0(dsol)
d0
dsol
## We can learn the coefficient of friction and the side dimenson of the cube
## form 35*0.40 seconds of recording. We use the simulator to evaluate the loss
## and its gradients by differentiating through the simulator. With gradient
## information we can use L-BFGS

jldsave(joinpath("..", "results", "sol.jld2"), dsol=dsol, Dsol=Dsol)

################################################################################
## Visualization
################################################################################

# ## Visualize Dataset with solution
datasol = d2data(dsol)
for i = 1:8
	datasol[(i-1)*7 + 4] = 0.05
end
set_simulator_data!(mech, datasol)
for traj in trajs0[1:10]
	visualize(mech, traj, vis=vis, show_contact = true)
	sleep(0.5)
end

# ## Visualize Progress made during 'learning'
for (i,dsol) in enumerate(Dsol[1:20])
	datasol = d2data(dsol)
	for i = 1:8
		datasol[(i-1)*7 + 4] = 0.05
	end
	set_simulator_data!(mech, datasol)
	visualize(mech, trajs0[i], vis=vis, show_contact = true)
	sleep(1.5)
end


text:
#################################################################################################
# AIM 1: X
#################################################################################################
## Deps
using CodecZlib
using StatsPlots
using Distributions
using DelimitedFiles
using Plots.PlotMeasures

## Constants
const noise_levels = ["$(s)" for s in collect(2.0:0.5:3.5)]
const data_dir = "/Users/jordiabante/OneDrive - Johns Hopkins/CpelNano/Data/Simulations/Aim-1"

# Colors
const blnd_col = ["#999999","#E69F00","#56B4E9","#009E73","#F0E442","#0072B2","#D55E00","#CC79A7"]
const cllr_color_code = Dict("Nanopolish" => "#D55E00", "DeepSignal" => "#009E73", "Megalodon" => "#56B4E9")
const s_color_cde = Dict(0.5 => blnd_col[1], 1.0 => blnd_col[2], 1.5 => blnd_col[3], 2.0 => blnd_col[4],
                         2.5 => blnd_col[5], 3.0 => blnd_col[6], 3.5 => blnd_col[7], 4.0 => blnd_col[8])

## Default attributes
default(titlefont=(14, "arial"),guidefont=(16, "arial"),tickfont=(12, "arial"))

# Function to compute accuracy
function comp_accu(truth, pred)

    # Return accuracy
    return sum(truth .== pred) / length(truth)

end

# Function to compute precision
function comp_prec(truth, pred)

    # Get positives
    pos_ind = pred .== 1

    # Get true positives
    true_pos = sum(truth[pos_ind] .== pred[pos_ind])

    # Return precision
    return true_pos / sum(pos_ind)

end

# Function to compute sensitivity (recall, or true positive rate)
function comp_sens(truth, pred)

    # Get positives
    pos_ind = pred .== 1

    # Get true positives
    if length(truth[pos_ind]) != 0
        true_pos = sum(truth[pos_ind] .== pred[pos_ind])
    else
        return 0.0
    end

    # Return sensitivity
    return true_pos / sum(truth .== 1)

end

# Function to compute specificity (selectivity or true negative rate)
function comp_spec(truth, pred)

    # Get negatives
    neg_ind = pred .== -1

    # Get prediction & truth
    if length(truth[neg_ind]) != 0
        true_neg = sum(truth[neg_ind] .== pred[neg_ind])
    else
        return 0.0
    end

    # Return specificity
    return true_neg / sum(truth .== -1)

end

function map_noise_ex(caller, s)

    # Print noise level
    println("Working on sigma=$(s)")

    # Read in
    stream = GzipDecompressorStream(open("$(data_dir)/$(caller)/gm12878_chr22_sigma_$(s)_$(caller)_tuples_sample.tsv.gz"))
    in_data = readdlm(stream)
    close(stream)

    # Remove NaN
    kp_ind = .! map(isequal("n/a"), in_data[:,2])
    in_data = in_data[kp_ind,:]

    # Get data
    true_x = in_data[:,1]
    pred_x = in_data[:,3]

    # Return tuple
    return comp_accu(true_x, pred_x), comp_prec(true_x, pred_x), comp_sens(true_x, pred_x), comp_spec(true_x, pred_x)

end

#################################################################################################
# Performance of methylation callers with X_{n}
#################################################################################################

# Init plots
p1 = plot(xlabel="signal noise level (sd)", ylabel="accuracy", ylim=(0.5, 1.0));
p2 = plot(xlabel="signal noise level (sd)", ylabel="precision", ylim=(0.5, 1.0));
p3 = plot(xlabel="signal noise level (sd)", ylabel="sensitivity", ylim=(0.5, 1.0));
p4 = plot(xlabel="signal noise level (sd)", ylabel="specificity", ylim=(0.5, 1.0));

for caller in ["Nanopolish"]

    # Print caller
    println("Working on $(caller)")

    # Get error in call
    map_out = map(s -> map_noise_ex(caller, s), noise_levels)

    # Unravel map out
    accu_vec = [x[1] for x in map_out]
    prec_vec = [x[2] for x in map_out]
    sens_vec = [x[3] for x in map_out]
    spec_vec = [x[4] for x in map_out]
    println("Acc.= $(accu_vec)\nPrec.= $(prec_vec)\nSens.=$(sens_vec)\nSpec.=$(spec_vec)")

    # Update plot
    col = cllr_color_code[caller]
    plot!(p1, noise_levels, accu_vec, seriestype=:scatter, markershape=:circle, color=col, label="")
    plot!(p1, noise_levels, accu_vec, seriestype=:line, alpha=0.5, color=col, label="")
    plot!(p2, noise_levels, prec_vec, seriestype=:scatter, markershape=:circle, color=col, label="")
    plot!(p2, noise_levels, prec_vec, seriestype=:line, alpha=0.5, color=col, label="")
    plot!(p3, noise_levels, sens_vec, seriestype=:scatter, markershape=:circle, color=col, label="")
    plot!(p3, noise_levels, sens_vec, seriestype=:line, alpha=0.5, color=col, label="")
    plot!(p4, noise_levels, spec_vec, seriestype=:scatter, markershape=:circle, color=col, label="")
    plot!(p4, noise_levels, spec_vec, seriestype=:line, alpha=0.5, color=col, label="")

end

# Generate plot & store
plot(p1,p2,p3,p4,layout=(4, 1),size=(600, 1000),top_margin=10px,bottom_margin=10px,left_margin=20px,right_margin=20px,legend=(0.2, 0.3))
savefig("$(data_dir)/Benchmark-Callers-X.pdf")


text:
module IterationControl

import Base.*
import EarlyStopping: done, StoppingCriterion, StoppingCriterion
using InteractiveUtils
using EarlyStopping

const ES = EarlyStopping

# make a list of all controls:
const CONTROLS = [:Step,
                  :Info,
                  :Warn,
                  :Error,
                  :Callback,
                  :WithLossDo,
                  :WithTrainingLossesDo,
                  :WithNumberDo,
                  :Data]
for criterion in subtypes(StoppingCriterion)
    control = split(string(criterion), ".") |> last |> Symbol
    push!(CONTROLS, control)
end

# export controls:
for control in CONTROLS
    eval(:(export $control))
end

# re-export stopping criterion from EarlyStopping.jl:
for criterion in subtypes(StoppingCriterion)
    ex = split(string(criterion), ".") |> last |> Symbol
    eval(:(export $ex))
end

include("utilities.jl")
include("api.jl")
include("stopping_controls.jl")
include("composite_controls.jl")
include("wrapped_controls.jl")
include("controls.jl")
include("train.jl")
include("square_rooter.jl")
include("deprecated.jl")

end # module