text:
using FactCheck, FileIO, QuartzImageIO, Images, Colors, FixedPointNumbers, TestImages

# Saving notes:
# autumn_leaves and toucan fail as of November 2015. The "edges" of the
# leaves are visibly different after a save+load cycle. Not sure if the
# reader or the writer is to blame. Probably an alpha channel issue.
# Mri-stack and multichannel timeseries OME are both image stacks,
# but the save code only saves the first frame at the moment.

facts("FileIO default") do
    imagedir = joinpath(dirname(@__FILE__), "images")
    images = readdir(imagedir)
    for image in images
    	context(image) do
	    	img = load(joinpath(imagedir, image))
	    	@fact isa(img, Image) --> true
    	end
    end
end

facts("OS X reader") do
    context("Autumn leaves") do
        name = "autumn_leaves"
        img = testimage(name)
        @fact colorspace(img) --> "RGBA"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGBA{UFixed16}
        out_name = joinpath(tempdir(), name * ".png")
        # Calling `save` relies on FileIO dispatching to us, so we
        # make the call explicit.
        QuartzImageIO.save_(out_name, img, "public.png")
        # Ideally, the `convert` would not be necessary, but the
        # saving step goes through a conversion, so we need to do it
        # in this test
        @pending load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Camerman") do
        name = "cameraman"
        img = testimage(name)
        @fact colorspace(img) --> "Gray"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> Gray{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Earth Apollo") do
        name = "earth_apollo17"
        img = testimage(name)
        @fact colorspace(img) --> "RGB4"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGB4{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Fabio") do
        name = "fabio"
        img = testimage(name)
        @fact colorspace(img) --> "Gray"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> Gray{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("House") do
        name = "house"
        img = testimage(name)
        @fact colorspace(img) --> "GrayA"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> GrayA{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Jetplane") do
        name = "jetplane"
        img = testimage(name)
        @fact colorspace(img) --> "GrayA"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> GrayA{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Lighthouse") do
        name = "lighthouse"
        img = testimage(name)
        @fact colorspace(img) --> "RGB4"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGB4{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Mandrill") do
        name = "mandrill"
        img = testimage(name)
        @fact colorspace(img) --> "RGB"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGB{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Moonsurface") do
        name = "moonsurface"
        img = testimage(name)
        @fact colorspace(img) --> "Gray"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> Gray{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Mountainstream") do
        name = "mountainstream"
        img = testimage(name)
        @fact colorspace(img) --> "RGB4"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGB4{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("MRI Stack") do
        name = "mri-stack"
        img = testimage(name)
        @fact colorspace(img) --> "Gray"
        @fact ndims(img) --> 3
        @fact colordim(img) --> 0
        @fact eltype(img) --> Gray{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        # Stack saving isn't implemented yet
        @pending load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("M51") do
        name = "m51"
        img = testimage(name)
        @fact colorspace(img) --> "Gray"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> Gray{UFixed16}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("HeLa cells") do
        name = "hela-cells"
        img = testimage(name)
        @fact colorspace(img) --> "RGB"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGB{UFixed16}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Blobs GIF") do
        name = "blobs"
        img = testimage(name)
        @fact colorspace(img) --> "RGB4"
        @fact ndims(img) --> 2
        @fact colordim(img) --> 0
        @fact eltype(img) --> RGB4{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        @fact load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
    context("Multichannel timeseries OME") do
        name = "multi-channel-time-series.ome"
        img = testimage(name)
        @fact colorspace(img) --> "Gray"
        @fact ndims(img) --> 3
        @fact colordim(img) --> 0
        @fact eltype(img) --> Gray{UFixed8}
        out_name = joinpath(tempdir(), name * ".png")
        QuartzImageIO.save_(out_name, img, "public.png")
        # Stack saving isn't implemented yet
        @pending load(out_name) --> convert(Image{RGBA{UFixed8}}, img)
    end
end

facts("Streams") do
    name = "lighthouse"
    img = testimage(name)
    out_name = joinpath(tempdir(), name * ".png")
    context("saving") do
        open(out_name, "w") do io
            QuartzImageIO.save(Stream(format"PNG", io), img)
        end
        imgcmp = load(out_name)
        @fact convert(Image{RGB4}, imgcmp) --> img
    end
    context("loading") do
        imgcmp = open(out_name) do io
            QuartzImageIO.load(Stream(format"PNG", io))
        end
        @fact convert(Image{RGB4}, imgcmp) --> img
    end
end

FactCheck.exitstatus()


text:
"Stochastic Gradient Hamiltonian Monte Carlo Sampling"
abstract type SGHMC <: Algorithm end

isapproximate(::Type{<:SamplingAlgorithm}) = true

defΩ(::Type{SGHMC}) = Omega.SimpleΩ{Int, Float64}

"Stochastic Gradient Hamiltonian Monte Carlo with Langevin Dynamics Friction: https://arxiv.org/pdf/1402.4102.pdf"
function sghmc(ygen, nsteps, stepsize, current_q::AbstractVector, ω, state)
  d = length(current_q)
  q = transform(current_q)
  p = randn(d)
  current_p = p

  # Rejects proposals outside domain TODO: Something smarter
  # https://stats.stackexchange.com/questions/73885/mcmc-on-a-bounded-parameter-space
  # any(notunit, q) && return (current_q, false, state)

  # construct friction and noise estimate matrices
  Bhat = 0.1
  C = 1

  for i = 1:nsteps
    #@show i
    # generate a predicate and get the gradient of its ϵ
    predicate, state = ygen(state)
    invq = inv_transform(q)
    @assert !(any(isnan(invq)))
    ∇U(q) = gradient(predicate, unlinearize(invq, ω), invq) .* jacobian(invq)
    # ∇U(q) = gradient(predicate, unlinearize(q, ω), q)

    # update the location and momentum parameters
    q = q .+ stepsize .* p
    @assert !(any(isnan(q)))
    #@show q[1], p[1]
    # any(notunit, q) && return (current_q, false, state)
    #@show state
    term = rand(MvNormal(d, 2 * stepsize .* (C .- Bhat)))
    @assert !(any(isnan(term)))
    p = p - stepsize .* ∇U(q) - stepsize .* C * p + term
    @assert !(any(isnan(p)))
    #@show ∇U(q)
  end

  # no MH step necessary
  q = inv_transform(q)
  @assert !(any(isnan(q)))
  return (q, true, state)
end

"Sample from `ω | y == true` with Stochastic Gradient Hamiltonian Monte Carlo"
function Base.rand(ΩT::Type{OT}, ygen, state, alg::Type{SGHMC};
                   n=1000,
                   nsteps = 100,
                   stepsize = 0.001,
                   cb = default_cbs(n)) where {OT <: Ω}
  ω = ΩT()
  predicate, state = ygen(state)
  predicate(ω) # Initialize omega
  ωvec = linearize(ω) # UNNEC

  # ωsamples = ΩT[] # FIXME: preallocate (and use inbounds)
  ωsamples = [] # FIXME: preallocate (and use inbounds)
  accepted = 0
  m = div(n, 10)
  @show n
  @showprogress 1 "Running SGHMC Chain" for i = 1:n
    ωvec, wasaccepted, state = sghmc(ygen, nsteps, stepsize, ωvec, ω, state)
    @show mean(ωvec)
    @show std(ωvec)
    push!(ωsamples, unlinearize(ωvec, ω)) # UNNEC
    if wasaccepted
      accepted += 1
    end
    i % m == 0 && print_with_color(:light_blue,  "acceptance ratio: $(accepted/i)\n")
  end
  ωsamples
end

"Sample from `x | y == true` with Metropolis Hasting"
function Base.rand(x::Union{RandVar, UTuple{<:RandVar}}, ygen, state, alg::Type{SGHMC};
                   n::Integer = 1000,
                   nsteps = 100,
                   stepsize = 0.001,
                   ΩT::OT = Omega.SimpleΩ{Int, Float64},
                   cb = default_cbs(n)) where {OT}
  map(x, rand(ΩT, ygen, state, alg;
              n = n, nsteps = nsteps, stepsize = stepsize, cb = cb))
end

ext: jl


text:
module BuildMiBench

export build_automotive,
    build_consumer,
    build_network,
    build_office,
    build_security,
    build_telecomm

function build_automotive(wd)
    println("Build: automotive")
    cwd=pwd()

    apps = ["basicmath","bitcount","qsort","susan"]
    println("Entering $wd ...")
    cd(wd)
    for app in apps
        cd(app)
        println("Building $app")
        run(`make`)
        cd("..")
    end
    cd(cwd)
end

function build_consumer(wd)
    println("Build: consumer")
    cwd=pwd()

    apps = ["jpeg","lame","mad","tiff-v3.5.4","typeset"]
    subdirs = ["jpeg-6a","lame3.70","mad-0.14.2b",".","lout-3.24"]
    println("Entering $wd ...")
    cd(wd)
    for (app,sub) in zip(apps,subdirs)
        cd(app)
        println("Building $app")
        cd(sub)
        if app == "tiff-v3.5.4"
            run(`./configure --noninteractive`)
        elseif app == "mad"
            run(`./configure --without-id3tag`)
        end
        run(`make`)
        cd("..")
        if app != "tiff-v3.5.4"
            cd("..")
        end
    end
    cd(cwd)
end

function build_network(wd)
    println("Build: network")
    cwd=pwd()

    apps = ["dijkstra","patricia"]
    println("Entering $wd ...")
    cd(wd)
    for app in apps
        cd(app)
        println("Building $app")
        run(`make`)
        cd("..")
    end
    cd(cwd)
end

function build_office(wd)
    println("Build: office")
    cwd=pwd()

    apps = ["ghostscript","ispell","rsynth","sphinx","stringsearch"]
    println("Entering $wd ...")
    cd(wd)
    for app in apps
        cd(app)
        println("Building $app")
        if app == "ghostscript"
            cd("src")
            # use XCFLAGS to fix compilation error bug https://bugs.ghostscript.com/show_bug.cgi?id=692443
            run(`make XCFLAGS=-DHAVE_SYS_TIME_H=1`)
            cd("..")
        elseif app == "ispell"
            run(`make ispell`)
        elseif app == "rsynth"
            run(`./configure --host='x86_64-pc-linux-gnu'`)
            run(`make`)
        elseif app == "sphinx"
            run(`./configure`)
            run(`make`)
        elseif app == "stringsearch"
            run(`make`)
        end
        cd("..")
    end
    cd(cwd)
end

function build_security(wd)
    println("Build: security")
    cwd=pwd()

    apps = ["blowfish","pgp","rijndael","sha"]
    println("Entering $wd ...")
    cd(wd)
    # handle pgp separately
    app = "pgp"
    deleteat!(apps,apps .== app)
    cd(app)
    println("Building $app")
    cd("src")
    run(`make clean`)
    run(`make linux-portable`)
    cd("..")
    cd("..")

    # handle the rest of the apps
    for app in apps
        cd(app)
        println("Building $app")
        run(`make`)
        cd("..")
    end
    cd(cwd)
end

function build_telecomm(wd)
    println("Build: telecomm")
    cwd=pwd()

    apps = ["adpcm","CRC32","FFT","gsm"]
    println("Entering $wd ...")
    cd(wd)
    # handle adpcm separately
    app = "adpcm"
    deleteat!(apps,apps .== app)
    cd(app)
    println("Building $app")
    cd("src")
    run(`make`)
    cd("..")
    cd("..")

    # handle the rest of the apps
    for app in apps
        cd(app)
        println("Building $app")
        run(`make`)
        cd("..")
    end
    cd(cwd)
end

end

max_stars_repo_name: emcores/MiBench

max_issues_repo_name: emcores/MiBench

max_forks_repo_name: emcores/MiBench


text:
"""
TODO: Document mutation module.
"""
module mutation
import StatsBase.sample
importall common, variation
using core, utility, representation, individual, species
export Mutation, MutationDefinition, mutate!

"""
The base type used by all mutation operation definitions.
"""
abstract MutationDefinition <: VariationDefinition

"""
The base type for all mutation operations.
"""
abstract Mutation <: Variation

"""
Subjects the first `n` individuals from a provided collection to mutation,
performed using a supplied mutation operator.
"""
operate!(m::Mutation, srcs::IndividualCollection, n::Int) =
  operate!(m, srcs.stages[m.stage][1:n])

"""
Subjects a list of chromosomes to mutation, performed using a supplied mutation
operator.
"""
operate!{T}(m::Mutation, srcs::Vector{IndividualStage{T}}) =
  for src in srcs; operate!(m, src); end

"""
Performs the mutation operation on the provided individual stage provided that
it is in a valid state.
"""
operate!{T}(m::Mutation, src::IndividualStage{T}) =
  if valid(src); set(src, mutate!(m, get(src))); end

"""
All mutation operators accept a single individual.
"""
num_inputs(m::Mutation) = 1

"""
All mutation operators produce a single individual.
"""
num_outputs(m::Mutation) = 1

# Load all mutation operations.
include("mutation/null.jl")
include("mutation/bit_flip.jl")
include("mutation/gaussian.jl")
include("mutation/single_swap.jl")
include("mutation/reverse.jl")
end


text:
#=
functions to calculate summary statistics on the performance of a backtest
=#

function cum_pnl(results::Results)::Float64
    result = 0.0
    @inbounds for val in values(results.backtest)
        pnl::Vector = val[:PNL].values[:]
        result += sum(pnl)
    end
    return result
end


text:
using Swagger
using Test

# set the current julia executable to be used henceforth
("JULIA" in keys(ENV)) || (ENV["JULIA"] = joinpath(Sys.BINDIR, "julia"))

const gencmd = joinpath(dirname(@__FILE__()), "petstore", "generate.sh")
@info("Generating petstore", gencmd)
run(`$gencmd`)

if ENV["RUNNER_OS"] == "Linux"
    @info("Running petstore tests")
    include("petstore/runtests.jl")
else
    @info("Skipping petstore tests in OSX (can not run petstore docker on travis OSX)")
end


text:
function parseInput(input)
    chars = split(input, ",")
    chars = parse.(Int, chars)
    return chars
end

function part1(codes)
    i = 1
    while codes[i] != 99
        operation = codes[i]
        pos1 = codes[i+1] + 1
        pos2 = codes[i+2] + 1
        dest = codes[i+3] + 1
        if operation == 1
            codes[dest] = codes[pos1] + codes[pos2]
        elseif operation == 2
            codes[dest] = codes[pos1] * codes[pos2]
        end
        i += 4
    end
    return codes
end

function run(s)
    input = parseInput(split(s, "\n")[1])
    input[2] = 12
    input[3] = 2
    return part1(input)[1]
end

function main()
    res, time, memory = @timed run(ARGS[1])
    println("_duration:$(time*1000)")
    println(res)
end

main()


text:
struct Dynamic{V <: Real}
    τ::V
end

mutable struct Potential{V <: Real}
    a::V
    b::V
    temp::V
    λ::Integer
end

mutable struct System{V <: Real}
    N::Int
    L::V
    ρ::V
    vol::V
    press::AbstractArray
    rc::V
    rng::Random.AbstractRNG
    ener::AbstractArray
end

"""
    The system created should always allocate the pressure
and energy arrays for the user.
"""
function System(N, Lbox, ρ, volumen, rc, rng)
    P = Vector{Float32}(undef, N)
    E = Vector{Float32}(undef, N)

    return System{Float32}(N, Lbox, ρ, volumen, P, rc, rng, E)
end


text:
using PyCall
cdsapi = pyimport("cdsapi")

c = cdsapi.Client()

c.retrieve("reanalysis-era5-pressure-levels",
    Dict("variable" => "temperature",
        "pressure_level" => "1000",
        "product_type" => "reanalysis",
        "year" => "2008",
        "month" => "01",
        "day" => "01",
        "time" => "12:00",
        "format" => "netcdf"),
    "era5output.netcdf"
)


text:
using GNSSTools
using JLD
using PyPlot
pygui(true)

GM = 3.986004418e14  # m³s⁻²
eop = get_eop();

# N_rand = 25
# lats = rand(-90:0.001:90, N_rand)  # degrees
# longs = rand(-180:0.001:180, N_rand)  # degrees
# heights = fill(100, N_rand)  # meters
# llas = (lats, longs, heights)
# user_lla = llas
freq = L1_freq
# user_lla = (40.01, -105.2437, 1655);

file_name = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/ch1_doppler_curves_user_lla.jld"
# file_name = "ch1_doppler_curves_user_lla.jld"
user_lla = load(file_name, "user_lla")

Δt = 1  # seconds
total_time = 1*24*60*60  # seconds
t_range = Array(0:Δt:total_time);
min_elevation = 10  # degrees

# GPS
a = 26_560e3;  # m
plane_num = 6;
sat_per_plane = 4;
incl = 56;
ΔΩ = 360/plane_num;
Δf_per_plane = 360/sat_per_plane/2;  # degrees
orbital_period = 2π*sqrt(a^3 / GM)  # seconds
doppler_gps, doppler_rate_gps, doppler_bounds_gps,
doppler_rate_bounds_gps, raw_dopplers_gps = doppler_distribution(a, plane_num,
                                           sat_per_plane, incl,
                                           t_range, user_lla,
                                           freq, show_plot=false,
                                           print_steps=true,
                                           eop=eop,
                                           min_elevation=min_elevation,
                                           ΔΩ=ΔΩ,
                                           Δf_per_plane=Δf_per_plane,
										   return_constellation=false,
                                           return_raw_dopplers=true);
doppler_raw_gps, doppler_rate_raw_gps,
elevations_gps, max_elevations_gps = raw_dopplers_gps

# Iridium
a = Rₑ+780*1000;  # meters
plane_num = 6;
sat_per_plane = 11;
incl = 86;
ΔΩ = 360/plane_num/2;
Δf_per_plane = 360/sat_per_plane/2;  # degrees
orbital_period = 2π*sqrt(a^3 / GM)  # seconds
doppler_iridium, doppler_rate_iridium, doppler_bounds_iridium,
doppler_rate_bounds_iridium, raw_dopplers_iridium = doppler_distribution(a, plane_num,
                                            sat_per_plane, incl,
                                            t_range, user_lla,
                                            freq, show_plot=false,
                                            print_steps=true,
                                            eop=eop,
                                            ΔΩ=ΔΩ,
                                            Δf_per_plane=Δf_per_plane,
                                            return_raw_dopplers=true);
doppler_raw_iridium, doppler_rate_raw_iridium,
elevations_iridium, max_elevations_iridium = raw_dopplers_iridium

# Starlink
a = Rₑ+550*1000;  # meters
plane_num = 24;
sat_per_plane = 66;
incl = 53;
ΔΩ = 360/plane_num;
Δf_per_plane = 360/sat_per_plane/2;  # degrees
orbital_period = 2π*sqrt(a^3 / GM)  # seconds
doppler_starlink, doppler_rate_starlink, doppler_bounds_starlink,
doppler_rate_bounds_starlink, raw_dopplers_starlink = doppler_distribution(a, plane_num,
                                            sat_per_plane, incl,
                                            t_range, user_lla,
                                            freq, show_plot=false,
                                            print_steps=true,
                                            eop=eop,
                                            ΔΩ=ΔΩ,
                                            Δf_per_plane=Δf_per_plane,
                                            return_raw_dopplers=true);
doppler_raw_starlink, doppler_rate_raw_starlink,
elevations_starlink, max_elevations_starlink = raw_dopplers_starlink

# OneWeb
a = Rₑ+1200*1000;  # meters
plane_num = 18;
sat_per_plane = 36;
incl = 87.9;  # degrees
ΔΩ = 360/plane_num/2;
Δf_per_plane = 360/sat_per_plane/2;  # degrees
orbital_period = 2π*sqrt(a^3 / GM)  # seconds
doppler_oneweb, doppler_rate_oneweb, doppler_bounds_oneweb,
doppler_rate_bounds_oneweb, raw_dopplers_oneweb = doppler_distribution(a, plane_num,
                                            sat_per_plane, incl,
                                            t_range, user_lla,
                                            freq, show_plot=false,
                                            print_steps=true,
                                            eop=eop,
                                            ΔΩ=ΔΩ,
                                            Δf_per_plane=Δf_per_plane,
                                            return_raw_dopplers=true);
doppler_raw_oneweb, doppler_rate_raw_oneweb,
elevations_oneweb, max_elevations_oneweb = raw_dopplers_oneweb

save("ch1_doppler_curves.jld",
     "doppler_gps", doppler_gps,
	 "doppler_iridium", doppler_iridium,
	 "doppler_starlink", doppler_starlink,
	 "doppler_oneweb", doppler_oneweb,
	 "doppler_rate_gps", doppler_rate_gps,
	 "doppler_rate_iridium", doppler_rate_iridium,
	 "doppler_rate_starlink", doppler_rate_starlink,
	 "doppler_rate_oneweb", doppler_rate_oneweb,
	 "doppler_bounds_gps", doppler_bounds_gps,
	 "doppler_bounds_iridium", doppler_bounds_iridium,
	 "doppler_bounds_starlink", doppler_bounds_starlink,
	 "doppler_bounds_oneweb", doppler_bounds_oneweb,
	 "doppler_raw_gps", doppler_raw_gps,
	 "doppler_raw_iridium", doppler_raw_iridium,
	 "doppler_raw_starlink", doppler_raw_starlink,
	 "doppler_raw_oneweb", doppler_raw_oneweb,
	 "doppler_rate_raw_gps", doppler_rate_raw_gps,
	 "doppler_rate_raw_iridium", doppler_rate_raw_iridium,
	 "doppler_rate_raw_starlink", doppler_rate_raw_starlink,
	 "doppler_rate_raw_oneweb", doppler_rate_raw_oneweb,
	 "freq", freq,
	 "user_lla", user_lla,
	 "t_range", t_range,
	 "min_elevation", min_elevation,
	 "elevations_gps", elevations_gps,
	 "elevations_iridium", elevations_iridium,
	 "elevations_starlink", elevations_starlink,
	 "elevations_oneweb", elevations_oneweb,
	 "max_elevations_gps", max_elevations_gps,
	 "max_elevations_iridium", max_elevations_iridium,
	 "max_elevations_starlink", max_elevations_starlink,
	 "max_elevations_oneweb", max_elevations_oneweb)

# file_name = "/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/ch1_doppler_curves.jld"
# data = load(file_name,
#             "doppler_gps",
# 	       "doppler_iridium",
# 	       "doppler_starlink",
# 	       "doppler_oneweb",
# 	       "doppler_rate_gps",
# 	       "doppler_rate_iridium",
# 	       "doppler_rate_starlink",
# 	       "doppler_rate_oneweb",
# 	       "doppler_bounds_gps",
# 	       "doppler_bounds_iridium",
# 	       "doppler_bounds_starlink",
# 	       "doppler_bounds_oneweb",
# 	       "doppler_raw_gps",
# 	       "doppler_raw_iridium",
# 	       "doppler_raw_starlink",
# 	       "doppler_raw_oneweb",
# 	       "doppler_rate_raw_gps",
# 	       "doppler_rate_raw_iridium",
# 	       "doppler_rate_raw_starlink",
# 	       "doppler_rate_raw_oneweb",
# 	       "freq",
# 	       "user_lla",
# 	       "t_range",
# 	       "min_elevation")

# doppler_gps, doppler_iridium, doppler_starlink, doppler_oneweb,
# doppler_rate_gps, doppler_rate_iridium, doppler_rate_starlink, doppler_rate_oneweb,
# doppler_bounds_gps, doppler_bounds_iridium, doppler_bounds_starlink, doppler_bounds_oneweb,
# doppler_raw_gps, doppler_raw_iridium, doppler_raw_starlink, doppler_raw_oneweb,
# doppler_rate_raw_gps, doppler_rate_raw_iridium, doppler_rate_raw_starlink, doppler_rate_raw_oneweb,
# freq, user_lla, t_range, min_elevation = data

# # Make plots
# fig = figure(figsize=(8.5,9))
# # GPS
# ax1 = fig.add_subplot(4,2,1)
# ax1.hist(Array(doppler_gps.edges[1])[2:end]./1000,
#          Array(doppler_gps.edges[1])./1000, weights=doppler_gps.weights,
#          color="k");
# xlabel("Doppler (kHz)")
# ylabel("Probability")
# title("GPS")
# ax2 = fig.add_subplot(4,2,2)
# ax2.hist(Array(doppler_rate_gps.edges[1])[2:end],
#          Array(doppler_rate_gps.edges[1]), weights=doppler_rate_gps.weights,
#          color="k");
# xlabel("Doppler Rate (Hz/s)")
# ylabel("Probability")
# title("GPS")
# # Iridium
# ax3 = fig.add_subplot(4,2,3)
# ax3.hist(Array(doppler_iridium.edges[1])[2:end]./1000,
#          Array(doppler_iridium.edges[1])./1000, weights=doppler_iridium.weights,
#          color="k");
# xlabel("Doppler (kHz)")
# ylabel("Probability")
# title("Iridium")
# ax4 = fig.add_subplot(4,2,4)
# ax4.hist(Array(doppler_rate_iridium.edges[1])[2:end],
#          Array(doppler_rate_iridium.edges[1]), weights=doppler_rate_iridium.weights,
#          color="k");
# xlabel("Doppler Rate (Hz/s)")
# ylabel("Probability")
# title("Iridium")
# # Starlink
# ax5 = fig.add_subplot(4,2,5)
# ax5.hist(Array(doppler_starlink.edges[1])[2:end]./1000,
#          Array(doppler_starlink.edges[1])./1000, weights=doppler_starlink.weights,
#          color="k");
# xlabel("Doppler (kHz)")
# ylabel("Probability")
# title("Starlink")
# ax6 = fig.add_subplot(4,2,6)
# ax6.hist(Array(doppler_rate_starlink.edges[1])[2:end],
#          Array(doppler_rate_starlink.edges[1]), weights=doppler_rate_starlink.weights,
#          color="k");
# xlabel("Doppler Rate (Hz/s)")
# ylabel("Probability")
# title("Starlink")
# # OneWeb
# ax7 = fig.add_subplot(4,2,7)
# ax7.hist(Array(doppler_oneweb.edges[1])[2:end]./1000,
#          Array(doppler_oneweb.edges[1])./1000, weights=doppler_oneweb.weights,
#          color="k");
# xlabel("Doppler (kHz)")
# ylabel("Probability")
# title("OneWeb")
# ax8 = fig.add_subplot(4,2,8)
# ax8.hist(Array(doppler_rate_oneweb.edges[1])[2:end],
#          Array(doppler_rate_oneweb.edges[1]), weights=doppler_rate_oneweb.weights,
#          color="k");
# xlabel("Doppler Rate (Hz/s)")
# ylabel("Probability")
# title("OneWeb")

# subplots_adjust(hspace=0.7, wspace=0.35, top=0.93, left=0.1, right=0.9)
# savefig("/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/figures/ch1_doppler_distributions.svg", dpi=300)

# using PyCall
# basemap = pyimport("mpl_toolkits.basemap")
# # lon_0 is central longitude of projection.
# # resolution = 'c' means use crude resolution coastlines.
# fig = figure(figsize=(7.5, 2.5))
# ax1 = fig.add_subplot(1,1,1)
# m = basemap.Basemap(projection="robin", lon_0=0, resolution="c")
# m.drawcoastlines(color="black", linewidth=0.25)
# m.fillcontinents(color="#69b2a2",lake_color="#A6CAE0")
# # draw parallels and meridians.
# m.drawparallels(range(-90., 120., step=30.), linewidth=0.5, color="dimgrey")
# m.drawmeridians(range(0., 360., step=60.), linewidth=0.5, color="dimgrey")
# m.drawmapboundary(fill_color="#A6CAE0")
# # m.fillcontinents(color="grey", alpha=0.3)
# color = "tomato"
# marker = "."
# markersize = 5
# N = min(length(user_lla[1]), length(user_lla[2]), length(user_lla[3]))
# for i in 1:N
# 	lat = user_lla[1][i]
# 	lon	= user_lla[2][i]
# 	xpt, ypt = m(lon,lat)
# 	m.plot(xpt, ypt, color=color, marker=marker, markersize=markersize)
# end
# # title("$N Random Observation Locations")
# subplots_adjust(bottom=0.01, left=0.1, right=0.9, top=0.99)
# savefig("/home/sjbilardi/Dropbox/Apps/Overleaf/MastersThesis/figures/ch1_doppler_curves_random_location_map.pdf",
#         dpi=300)


text:
module ThermodynamicIntegration

using AdvancedHMC
using ForwardDiff
using Trapz
using ProgressMeter
using Random
using Requires
using Statistics

export ThermInt
export TIParallelThreads

const GLOBAL_RNG = Random.MersenneTwister(42)

const ADBACKEND = Ref(:ForwardDiff)

set_adbackend(ad::String) = set_adbackend(Symbol(ad))
set_adbackend(ad::Symbol) = set_adbackend(Val(ad))
set_adbackend(::Val{:ForwardDiff}) = ADBACKEND[] = :ForwardDiff
function set_adbackend(::Any)
    return error(
        "ad should be :ForwardDiff, :Zygote or :ReverseDiff\n" *
        "For Zygote and ReverseDiff, make sure to have" *
        " `using Zygote/ReverseDiff` in your script",
    )
end

include("thermint.jl")

function __init__()
    @require Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0" include("turing.jl")
    @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
        set_adbackend(::Val{:Zygote}) = ADBACKEND[] = :Zygote
    end
    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
        set_adbackend(::Val{:ReverseDiff}) = ADBACKEND[] = :ReverseDiff
    end
end

end


text:
# growth_witness()
# smoothen()

"""
    growth_witness(shift::AbstractArray, norm::AbstractArray, dt, [b]; skip=0) -> g
    growth_witness(df::DataFrame, [b]; skip=0) -> g
Compute the growth witness
```math
G^{(n)} = S^{(n)} - \\frac{\\vert\\mathbf{c}^{(n+1)}\\vert -
          \\vert\\mathbf{c}^{(n)}\\vert}{\\vert\\mathbf{c}^{(n)}\\vert d\\tau},
```
where `S` is the `shift` and \$\\vert\\mathbf{c}^{(n)}\\vert ==\$ `norm[n, 1]`.
Setting `b ≥ 1` a sliding average over `b` time steps is computed using
[`smoothen()`](@ref). The first `skip` time steps are skipped.
`mean(growth_witness)` is approximately the same as [`growth_estimator`](@ref) with `h=0`.

See also [`growth_estimator`](@ref).
"""
function growth_witness(shift::AbstractArray, norm::AbstractArray, dt; skip = 0)
    l = length(norm)
    @assert length(shift) == l "`shift` and `norm` arrays need to have the same length."
    l =  l-skip
    @assert l > 0 "`skip` must be larger than length of `shift` and `norm`"
    n = l - 1
    g = Vector{eltype(shift)}(undef, l)
    for i in 1:n
        g[i] = shift[skip+i] - (norm[skip+i+1] - norm[skip+i])/(dt*norm[skip+i])
    end
    # pad the vector g at the end
    g[n+1] = @views mean(g[1:n])
    return g
end
function growth_witness(shift::AbstractArray, norm::AbstractArray, dt, b; kwargs...)
    g_raw = growth_witness(shift, norm, dt; kwargs...)
    return smoothen(g_raw, b)
end
"""
    growth_witness(df::DataFrame, [b]; shift=:shift, norm=:norm, dτ=df.dτ[end], skip=0)
Calculate the growth witness directly from a `DataFrame` returned by
[`lomc!`](@ref). The keyword arguments `shift` and `norm`
can be used to change the names of the relevant columns.
"""
function growth_witness(
    df::DataFrame, b=Val(0);
    shift=:shift, norm=:norm, dτ=df.dτ[end], kwargs...
)
    shift_vec = getproperty(df, Symbol(shift))
    norm_vec = getproperty(df, Symbol(norm))
    return growth_witness(shift_vec, norm_vec, dτ, b; kwargs...)
end

"""
    smoothen(noisy::AbstractVector, b)
Smoothen the array `noisy` by averaging over a sliding window of length `b` and
wrapping `noisy` periodically. The [`mean(noisy)`](@ref) is preserved.
"""
function smoothen(noisy::AbstractVector, b::Integer)
    l = length(noisy)
    @assert 1 ≤ b ≤ l
    smooth = zeros(float(eltype(noisy)), l)
    offset = b÷2 + 1
    for i in 1:l
        for j = 1:b # average over `b` elements of `noisy`
            smooth[i] += noisy[mod1(i-offset+j, l)]/b
        end
    end
    return smooth
end
smoothen(noisy::AbstractVector, ::Val{0}) = noisy


text:
GraphPlot._src_index(e::LightGraphs.AbstractEdge, g::LightGraphs.AbstractGraph) = LightGraphs.src(e)
GraphPlot._dst_index(e::LightGraphs.AbstractEdge, g::LightGraphs.AbstractGraph) = LightGraphs.dst(e)

function cgplot(g::AbstractColoredGraph; kw...)
    if :edgelabel ∉ keys(kw)
        clist = []
        for e=edges(g)
            push!(clist, color(e))
        end
        kw = (edgelabel=clist, kw...)
    end

    GraphPlot.gplot(g; kw...)
end


text:
struct NotEqualSet <: MOI.AbstractVectorSet end

Base.copy(::NotEqualSet) = NotEqualSet()

MOI.dimension(s::NotEqualSet) = 2


text:
export SymMatrix
export square_getindex, symmetric_setindex!

"""
    struct SymMatrix{T} <: AbstractMatrix{T}
        Q::Vector{T}
        n::Int
    end

Symmetric ``n \\times n`` matrix storing the vectorized upper triangular
part of the matrix in the `Q` vector (this corresponds to the vectorized
form of `MathOptInterface.PositiveSemidefiniteConeTriangle`).
It implement the `AbstractMatrix` interface except for `setindex!` as it might
break its symmetry. The [`symmetric_setindex!`](@ref) function should be used
instead.
"""
struct SymMatrix{T} <: AbstractMatrix{T}
    Q::Vector{T}
    n::Int
end

_undef_sym(T::Type, n) = Vector{T}(undef, trimap(n, n))
function SymMatrix{T}(::UndefInitializer, dims::Dims{2}) where T
    dims[1] != dims[2] && error("Expected same dimension for `SymMatrix`, got `$(dims)`.")
    n = dims[1]
    return SymMatrix(_undef_sym(T, n), n)
end
Base.similar(Q::SymMatrix{T}, dims::Tuple{Base.OneTo, Vararg{Base.OneTo}}) where {T} = similar(Matrix{T}, dims)
Base.similar(Q::SymMatrix, T::Type, dims::Tuple{Base.OneTo, Vararg{Base.OneTo}}) = similar(Matrix{T}, dims)
Base.similar(Q::SymMatrix, T::Type, dims::Dims{2}) = similar(SymMatrix{T}, dims)
Base.copy(Q::SymMatrix) = SymMatrix(copy(Q.Q), Q.n)
Base.map(f::Function, Q::SymMatrix) = SymMatrix(map(f, Q.Q), Q.n)

# i <= j
#function trimap(i, j, n) # It was used when Q was the lower triangular part
#    div(n*(n+1), 2) - div((n-i+1)*(n-i+2), 2) + j-i+1
#end

# i <= j
trimap(i, j) = div(j * (j - 1), 2) + i

function trimat(::Type{T}, f, n, σ) where {T}
    Q = _undef_sym(T, n)
    for j in 1:n
        for i in 1:j
            Q[trimap(i, j)] = f(σ[i], σ[j])
        end
    end
    return SymMatrix(Q, n)
end

Base.size(Q::SymMatrix) = (Q.n, Q.n)

"""
    square_getindex!(Q::SymMatrix, I)

Return the `SymMatrix` corresponding to `Q[I, I]`.
"""
function square_getindex(Q::SymMatrix{T}, I) where T
    n = length(I)
    q = _undef_sym(T, n)
    k = 0
    for (j, Ij) in enumerate(I)
        for (i, Ii) in enumerate(I)
            i > j && break
            k += 1
            q[k] = Q[Ii, Ij]
        end
    end
    return SymMatrix(q, n)
end

"""
    symmetric_setindex!(Q::SymMatrix, value, i::Integer, j::Integer)

Set `Q[i, j]` and `Q[j, i]` to the value `value`.
"""
function symmetric_setindex!(Q::SymMatrix, value, i::Integer, j::Integer)
    Q.Q[trimap(min(i, j), max(i, j))] = value
end

function Base.getindex(Q::SymMatrix, i::Integer, j::Integer)
    return Q.Q[trimap(min(i, j), max(i, j))]
end
Base.getindex(Q::SymMatrix, I::Tuple) = Q[I...]
Base.getindex(Q::SymMatrix, I::CartesianIndex) = Q[I.I]


text:
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using LinearAlgebra, InvertibleNetworks, Test

###############################################################################
# Leaky ReLU

nx = 28
ny = 28
n_in = 4
batchsize = 2

X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Invertibility
err = norm(X - LeakyReLUinv(LeakyReLU(X)))
@test isapprox(err, 0f0, atol=1f-6)

err = norm(X - LeakyReLU(LeakyReLUinv(X)))
@test isapprox(err, 0f0, atol=1f-6)

# Gradient test
function objective(X, Y)
    Y0 = LeakyReLU(X)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = LeakyReLUgrad(ΔY, Y0)  # Use Y0 to recompute forward state
    return f, ΔX
end

# Observed data
Y = ReLU(X)

# Gradient test for X
maxiter = 5
print("\nGradient test leaky ReLU\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

###############################################################################
# ReLU

nx = 28
ny = 28
n_in = 4
batchsize = 2

X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Gradient test
function objective(X, Y)
    Y0 = ReLU(X)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = ReLUgrad(ΔY, X)    # Pass X, as ReLU not invertible
    return f, ΔX
end

# Observed data
Y = ReLU(X)

# Gradient test for X
maxiter = 5
print("\nGradient test ReLU\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

###############################################################################
# Sigmoid

nx = 12
ny = 12
n_in = 4
batchsize = 2

X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Invertibility
err = norm(X - SigmoidInv(Sigmoid(X))) / norm(X)
@test isapprox(err, 0f0, atol=1f-5)

# Gradient test sigmoid
function objective(X, Y)
    Y0 = Sigmoid(X)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = SigmoidGrad(ΔY, Y0)
    return f, ΔX
end

# Observed data
Y = Sigmoid(X)

# Gradient test for X
maxiter = 5
print("\nGradient test Sigmoid\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Shifted and Scaled Sigmoid
low = 0.5f0
high = 1f0

nx = 12
ny = 12
n_in = 4
batchsize = 2

X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Invertibility
err = norm(X - SigmoidInv(Sigmoid(X; low=low, high=high); low=low, high=high)) / norm(X)
@test isapprox(err, 0f0, atol=1f-5)

# Gradient test Shifted and Scaled Sigmoid
function objective(X, Y)
    Y0 = Sigmoid(X; low=low, high=high)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = SigmoidGrad(ΔY, Y0; low=low, high=high)
    return f, ΔX
end

# Observed data
Y = Sigmoid(X; low=low, high=high)

# Gradient test for X
maxiter = 5
print("\nGradient test scaled and shifted Sigmoid\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

###############################################################################
# Gated linear unit (GaLU)

nx = 12
ny = 12
n_in = 4
batchsize = 2

X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Gradient test GaLU
function objective(X, Y)
    Y0 = GaLU(X)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = GaLUgrad(ΔY, X)
    return f, ΔX
end

# Observed data
Y = GaLU(X)

# Gradient test for X
maxiter = 5
print("\nGradient test GaLU\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

###############################################################################
# Sigmoid2 (scaled version of Sigmoid)

nx = 12
ny = 12
n_in = 4
batchsize = 2

X = randn(Float32, nx, ny, n_in, batchsize)
X0 = randn(Float32, nx, ny, n_in, batchsize)
dX = X - X0

# Gradient test Sigmoid2
L = Sigmoid2Layer()
function objective(X, Y)
    Y0 = L.forward(X)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = L.backward(ΔY, Y0)
    return f, ΔX
end

# Observed data
Y = L.forward(X)

# Gradient test for X
maxiter = 5
print("\nGradient test Sigmoid2\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

###############################################################################
# ExpClamp

nx = 12
ny = 12
n_in = 4
batchsize = 3

X = glorot_uniform(nx, ny, n_in, batchsize)
X0 = glorot_uniform(nx, ny, n_in, batchsize)
dX = X - X0

# Invertibility
err = norm(X - ExpClampInv(ExpClamp(X))) / norm(X)
@test isapprox(err, 0f0, atol=1f-5)

# Gradient test sigmoid
function objective(X, Y)
    Y0 = ExpClamp(X)
    ΔY = Y0 - Y
    f = .5f0*norm(ΔY)^2
    ΔX = ExpClampGrad(ΔY, Y0)
    return f, ΔX
end

# Observed data
Y = ExpClamp(X)

# Gradient test for X
maxiter = 5
print("\nGradient test ExpClamp\n")
f0, ΔX = objective(X0, Y)
h = .1f0
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)
for j=1:maxiter
    f = objective(X0 + h*dX, Y)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)


text:
using Stheno, Plots
plotly();

############################## Define model and data ################################

# Define additive model.
σ_noise, gpc = 1e-1, GPC();
f₁ = GP(pick_dims(EQ(), 1), gpc);
f₂ = GP(pick_dims(EQ(), 2), gpc);
f₃ = f₁ + f₂;
y = f₃ + GP(Noise(σ_noise^2), gpc);

# Sample from marginal process to generate toy data.
rng, S = MersenneTwister(123456), 25;
X = ColsAreObs(randn(2, 100));
ŷ = rand(rng, y(X));

############################### Set up plotting ##################################

# Define plotting locations.
P, Q, pos, b = 50, 50, 1, 4.0;
Ps, Qs, Xp = linspace(-b, b, P), linspace(-b, b, Q), Matrix{Float64}(2, P * Q);
for p in 1:P, q in 1:Q
    Xp[1, pos] = Ps[p]
    Xp[2, pos] = Qs[q]
    pos += 1
end
Xp = ColsAreObs(Xp);

exact_plot, approx_plot = plot(), plot();

################################ Do exact inference ################################

@show logpdf(y(X), ŷ);
f₁′, f₂′, f₃′ = (f₁, f₂, f₃) | (y(X) ← ŷ);
f₁′Xp, f₂′Xp, f₃′Xp = rand(rng, [f₁′(Xp), f₂′(Xp), f₃′(Xp)], S);
μ₁′, σ₁′ = marginals(f₁′(Xp));
μ₂′, σ₂′ = marginals(f₂′(Xp));
μ₃′, σ₃′ = marginals(f₃′(Xp));

################################ Do approximate inference ###############################

M = 15;
Z₁ = ColsAreObs(hcat(linspace(-b, b, M), zeros(M))');
Z₂ = ColsAreObs(hcat(zeros(M), linspace(-b, b, M))');
u₁, u₂ = f₁(Z₁), f₂(Z₂);
u = BlockGP([u₁, u₂]);
@show elbo(f₃(X), ŷ, u, σ_noise);
μ′u, Σ′u = Stheno.optimal_q(f₃(X), ŷ, u, σ_noise);
conditioner = Stheno.Titsias(u, μ′u, Σ′u, gpc);
f₁′u, f₂′u, f₃′u = (f₁, f₂, f₃) | conditioner;


text:
"""
```
eqcond(m::Schorf)
```

Expresses the equilibrium conditions in canonical form using Γ0, Γ1, C, Ψ, and Π matrices.
Using the mappings of states/equations to integers defined in an_schorfheide.jl, coefficients are
specified in their proper positions.

### Outputs

* `Γ0` (`n_states` x `n_states`) holds coefficients of current time states.
* `Γ1` (`n_states` x `n_states`) holds coefficients of lagged states.
* `C`  (`n_states` x `1`) is a vector of constants
* `Ψ`  (`n_states` x `n_shocks_exogenous`) holds coefficients of iid shocks.
* `Π`  (`n_states` x `n_states_expectational`) holds coefficients of expectational states.
"""
function eqcond(m::AnSchorfheide) # do not edit this function definition
    return eqcond(m, 1)
end

function eqcond(m::AnSchorfheide, reg::Int) # do not edit these inputs
    endo = m.endogenous_states # for convenience, make references to these fields
    exo  = m.exogenous_shocks  # which hold indices, e.g. endo[:y_t] will return an Int
    ex   = m.expected_shocks   # for the index of the `y_t` variable.
    eq   = m.equilibrium_conditions # you can rename these abbreviations if you want

    Γ0 = zeros(n_states(m), n_states(m)) # initialize containers (do not edit)
    Γ1 = zeros(n_states(m), n_states(m))
    C  = zeros(n_states(m))
    Ψ  = zeros(n_states(m), n_shocks_exogenous(m))
    Π  = zeros(n_states(m), n_shocks_expectational(m))

    # Regime-switching parameters (you should probably not edit this call)
    for para in m.parameters      # if you're new to DSGE.jl, then revisit this block of code later.
        if !isempty(para.regimes) # if you don't need regime-switching, then you can comment out this block of code
            if (haskey(get_settings(m), :model2para_regime) ? haskey(get_setting(m, :model2para_regime), para.key) : false)
                toggle_regime!(para, reg, get_setting(m, :model2para_regime)[para.key])
            else
                toggle_regime!(para, reg)
            end
        end
    end

    # You should definitely make edits to the following code!

    ### ENDOGENOUS STATES ###

    ### 1. Consumption Euler Equation

    Γ0[eq[:eq_euler], endo[:y_t]]  = 1.
    Γ0[eq[:eq_euler], endo[:R_t]] = 1/m[:τ]
    Γ0[eq[:eq_euler], endo[:g_t]] = -(1-m[:ρ_g])
    Γ0[eq[:eq_euler], endo[:z_t]] = -m[:ρ_z]/m[:τ]
    Γ0[eq[:eq_euler], endo[:Ey_t]] = -1
    Γ0[eq[:eq_euler], endo[:Eπ_t]] = -1/m[:τ]

    ### 2. NK Phillips Curve

    Γ0[eq[:eq_phillips], endo[:y_t]] = -m[:κ]
    Γ0[eq[:eq_phillips], endo[:π_t]] = 1
    Γ0[eq[:eq_phillips], endo[:g_t]] = m[:κ]
    Γ0[eq[:eq_phillips], endo[:Eπ_t]] = -1/(1+m[:rA]/400)

    ### 3. Monetary Policy Rule

    Γ0[eq[:eq_mp], endo[:y_t]] = -(1-m[:ρ_R])*m[:ψ_2]
    Γ0[eq[:eq_mp], endo[:π_t]] = -(1-m[:ρ_R])*m[:ψ_1]
    Γ0[eq[:eq_mp], endo[:R_t]] = 1
    Γ0[eq[:eq_mp], endo[:g_t]] = (1-m[:ρ_R])*m[:ψ_2]
    Γ1[eq[:eq_mp], endo[:R_t]] = m[:ρ_R]
    Ψ[eq[:eq_mp], exo[:rm_sh]] = 1

    ### 4. Output lag

    Γ0[eq[:eq_y_t1], endo[:y_t1]] = 1
    Γ1[eq[:eq_y_t1], endo[:y_t]] = 1

    ### 5. Government spending

    Γ0[eq[:eq_g], endo[:g_t]] = 1
    Γ1[eq[:eq_g], endo[:g_t]] = m[:ρ_g]
    Ψ[eq[:eq_g], exo[:g_sh]] = 1

    ### 6. Technology

    Γ0[eq[:eq_z], endo[:z_t]] = 1
    Γ1[eq[:eq_z], endo[:z_t]] = m[:ρ_z]
    Ψ[eq[:eq_z], exo[:z_sh]] = 1

    ### 7. Expected output

    Γ0[eq[:eq_Ey], endo[:y_t]] = 1
    Γ1[eq[:eq_Ey], endo[:Ey_t]] = 1
    Π[eq[:eq_Ey], ex[:Ey_sh]] = 1

    ### 8. Expected inflation

    Γ0[eq[:eq_Eπ], endo[:π_t]] = 1
    Γ1[eq[:eq_Eπ], endo[:Eπ_t]] = 1
    Π[eq[:eq_Eπ], ex[:Eπ_sh]] = 1

    # Ensure parameter regimes are in 1 at the end (you should probably not edit)
    for para in m.parameters      # if you're new to DSGE.jl, then revisit this block of code later.
        if !isempty(para.regimes) # if you don't need regime-switching, then you can comment out this block of code
            toggle_regime!(para, 1)
        end
    end

    return Γ0, Γ1, C, Ψ, Π # do not edit this return
end


text:
function ticks_and_labels end

module Formatters
using Showoff

function scientific(ticks::AbstractVector)
    return Showoff.showoff(ticks, :scientific)
end

function plain(ticks::AbstractVector)
    try
        Showoff.showoff(ticks, :plain)
    catch e
        Base.showerror(stderr, e)
        println("with ticks: ", ticks)
        String["-Inf", "Inf"]
    end
end

end
using .Formatters

"""
    $(SIGNATURES)

Plots a 2-dimensional axis.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Axis2D) do scene
    return Attributes(visible=true, showgrid=true, showticks=true, showtickmarks=true, padding=0.0,
                      ticks=Attributes(ranges_labels=(automatic, automatic), formatter=Formatters.plain,
                                       gap=1, title_gap=3, linewidth=(1, 1),
                                       linecolor=((:black, 0.4), (:black, 0.4)), linestyle=(nothing, nothing),
                                       textcolor=(:black, :black), textsize=(5, 5), rotation=(0.0, 0.0),
                                       align=((:center, :top), (:right, :center)),
                                       font=lift(dim2, theme(scene, :font))),
                      tickmarks=Attributes(length=(1.0, 1.0), linewidth=(1, 1),
                                           linecolor=((:black, 0.4), (:black, 0.2)),
                                           linestyle=(nothing, nothing)),
                      grid=Attributes(linewidth=(0.5, 0.5), linecolor=((:black, 0.3), (:black, 0.3)),
                                      linestyle=(nothing, nothing)),
                      frame=Attributes(linewidth=1.0, linecolor=:black, linestyle=nothing,
                                       axis_position=nothing, axis_arrow=false, arrow_size=2.5,
                                       frames=((false, false), (false, false))),
                      names=Attributes(axisnames=("x", "y"), title=nothing, textcolor=(:black, :black),
                                       textsize=(6, 6), rotation=(0.0, -1.5pi),
                                       align=((:center, :top), (:center, :bottom)),
                                       font=lift(dim2, theme(scene, :font))))
end

"""
    $(SIGNATURES)

Plots a 3-dimensional Axis.

## Attributes
$(ATTRIBUTES)
"""
@recipe(Axis3D) do scene
    q1 = qrotation(Vec3f0(1, 0, 0), -0.5f0 * pi)
    q2 = qrotation(Vec3f0(0, 0, 1), 1.0f0 * pi)
    tickrotations3d = (qrotation(Vec3f0(0, 0, 1), -1.5pi), q2, qrotation(Vec3f0(1, 0, 0), -0.5pi) * q2)
    axisnames_rotation3d = tickrotations3d
    tickalign3d = ((:left, :center), # x axis
                   (:right, :center), # y axis
                   (:right, :center))
    axisnames_align3d = tickalign3d
    tick_color = RGBAf0(0.5, 0.5, 0.5, 0.6)
    grid_color = RGBAf0(0.5, 0.5, 0.5, 0.4)
    grid_thickness = 1
    gridthickness = ntuple(x -> 1.0f0, Val(3))
    tsize = 5 # in percent
    return Attributes(visible=true, showticks=(true, true, true), showaxis=(true, true, true),
                      showgrid=(true, true, true), scale=Vec3f0(1), padding=0.1,
                      names=Attributes(axisnames=("x", "y", "z"), textcolor=(:black, :black, :black),
                                       rotation=axisnames_rotation3d, textsize=(6.0, 6.0, 6.0),
                                       align=axisnames_align3d, font=lift(dim3, theme(scene, :font)), gap=3),
                      ticks=Attributes(ranges_labels=(automatic, automatic), formatter=Formatters.plain,
                                       textcolor=(tick_color, tick_color, tick_color),
                                       rotation=tickrotations3d, textsize=(tsize, tsize, tsize),
                                       align=tickalign3d, gap=3, font=lift(dim3, theme(scene, :font))),
                      frame=Attributes(linecolor=(grid_color, grid_color, grid_color),
                                       linewidth=(grid_thickness, grid_thickness, grid_thickness),
                                       axiscolor=(:black, :black, :black)))
end

isaxis(x) = false
isaxis(x::Union{Axis2D,Axis3D}) = true

const Limits{N} = NTuple{N,Tuple{Number,Number}}

function default_ticks(limits::Limits, ticks, scale_func::Function)
    return default_ticks.(limits, (ticks,), scale_func)
end

# function default_ticks(limits::Limits, ticks::Tuple, scale_func::Function)
#     default_ticks.(limits, ticks, scale_func)
# end
function default_ticks(limits::Tuple{Number,Number}, ticks, scale_func::Function)
    return default_ticks(limits..., ticks, scale_func)
end

function default_ticks(lmin::Number, lmax::Number, ticks::AbstractVector{<:Number}, scale_func::Function)
    return scale_func.((filter(t -> lmin <= t <= lmax, ticks)))
end
function default_ticks(lmin::Number, lmax::Number, ::Automatic, scale_func::Function)
    # scale the limits
    scaled_ticks, mini, maxi = optimize_ticks(scale_func(lmin), scale_func(lmax); k_min=4, # minimum number of ticks # minimum number of ticks
                                              k_max=8)
    length(scaled_ticks) == 1 && isnan(scaled_ticks[1]) && return [-Inf, Inf]
    return scaled_ticks
end

function default_ticks(lmin::Number, lmax::Number, ticks::Integer, scale_func=identity)
    scaled_ticks, mini, maxi = optimize_ticks(scale_func(lmin), scale_func(lmax); k_min=ticks, # minimum number of ticks # minimum number of ticks
                                              k_max=ticks, # maximum number of ticks # maximum number of ticks
                                              k_ideal=ticks,
                                              # `strict_span = false` rewards cases where the span of the
                                              # chosen  ticks is not too much bigger than amin - amax:
                                              strict_span=false)
    return scaled_ticks
end

function default_ticks(x::Automatic, limits::Tuple, n)
    return default_ticks(limits, n, identity)
end

function default_ticks(ticks::Tuple, limits::Tuple, n::Tuple)
    return default_ticks.(ticks, (limits,), n)
end

default_ticks(ticks::Tuple, limits::Limits, n) = default_ticks.(ticks, limits, (n,))

default_ticks(ticks::Tuple, limits::Limits, n::Tuple) = default_ticks.(ticks, limits, n)

default_ticks(ticks::AbstractVector{<:Number}, limits, n) = ticks

function default_labels(x::NTuple{N,Any}, formatter::Function) where {N}
    return default_labels.(x, formatter)
end

function default_labels(x::AbstractVector, y::AbstractVector, formatter::Function=Formatters.plain)
    return default_labels.((x, y), formatter)
end

function default_labels(ticks::AbstractVector, formatter::Function=Formatters.plain)
    if applicable(formatter, ticks)
        return formatter(ticks) # takes the whole array
    elseif applicable(formatter, first(ticks))
        return formatter.(ticks)
    else
        error("Formatting function $(formatter) is neither applicable to $(typeof(ticks)) nor $(eltype(ticks)).")
    end
end

default_labels(x::Automatic, ranges, formatter) = default_labels(ranges, formatter)
default_labels(x::Tuple, ranges::Tuple, formatter) = default_labels.(x, ranges, (formatter,))
default_labels(x::Tuple, ranges, formatter) = default_labels.(x, (ranges,), (formatter,))
default_labels(x::AbstractVector{<:AbstractString}, ranges, formatter::Function) = x
default_labels(x::AbstractVector{<:AbstractString}, ranges::AbstractVector, formatter::Function) = x

function convert_arguments(::Type{<:Axis2D}, limits::Rect)
    e = extrema(limits)
    return (((e[1][1], e[2][1]), (e[1][2], e[2][2])),)
end

function convert_arguments(::Type{<:Axis3D}, limits::Rect)
    e = (minimum(limits), maximum(limits))
    return (((e[1][1], e[2][1]), (e[1][2], e[2][2]), (e[1][3], e[2][3])),)
end

a_length(x::AbstractVector) = length(x)
a_length(x::Automatic) = x

function calculated_attributes!(::Type{<:Union{Axis2D,Axis3D}}, plot)
    ticks = plot.ticks
    args = (plot.padding, plot[1], ticks.ranges, ticks.labels, ticks.formatter)
    ticks[:ranges_labels] = lift(args...) do pad, lims, ranges, labels, formatter
        limit_widths = map(x -> x[2] - x[1], lims)
        pad = (limit_widths .* pad)
        # pad the drawn limits and use them as the ranges
        lim_pad = map((lim, p) -> (lim[1] - p, lim[2] + p), lims, pad)
        num_ticks = labels === automatic ? automatic : a_length.(labels)
        ranges = default_ticks(ranges, lim_pad, num_ticks)
        labels = default_labels(labels, ranges, formatter)
        return (ranges, labels)
    end
    return
end

function draw_ticks(textbuffer, dim, origin, ticks, linewidth, linecolor, linestyle, textcolor, textsize,
                    rotation, align, font)
    for (tick, str) in ticks
        pos = ntuple(i -> i != dim ? origin[i] : tick, Val(2))
        push!(textbuffer, str, Point(pos), rotation=rotation[dim], textsize=textsize[dim], align=align[dim],
              color=textcolor[dim], font=font[dim])
    end
end

function draw_tickmarks(linebuffer, dim, origin, ticks, dir::NTuple{N}, linewidth, linecolor, linestyle,
                        textcolor, textsize, rotation, align, font) where {N}
    for (tick, str) in ticks
        pos = ntuple(i -> i != dim ? origin[i] : tick, Val(2))
        posf0 = Point2f0(pos)
        dirf0 = Pointf0{N}(dir)
        append!(linebuffer, [posf0, posf0 .+ dirf0], color=linecolor[dim], linewidth=linewidth[dim])
    end
end

function draw_grid(linebuffer, dim, origin, ticks, dir::NTuple{N}, linewidth, linecolor, linestyle) where {N}
    dirf0 = Pointf0{N}(dir)
    for (tick, str) in ticks
        tup = ntuple(i -> i != dim ? origin[i] : tick, Val(N))
        posf0 = Pointf0{N}(tup)
        append!(linebuffer, [posf0, posf0 .+ dirf0], color=linecolor[dim], linewidth=linewidth[dim])
    end
end

function draw_frame(linebuffer, limits::NTuple{N,Any}, linewidth, linecolor, linestyle, axis_position,
                    axis_arrow, arrow_size, frames) where {N}
    mini = minimum.(limits)
    maxi = maximum.(limits)
    rect = Rect(Vec(mini), Vec(maxi .- mini))
    origin = Vec{N}(0.0)
    if (origin in rect) && axis_position == :origin
        for i in 1:N
            start = unit(Point{N,Float32}, i) * Float32(mini[i])
            to = unit(Point{N,Float32}, i) * Float32(maxi[i])
            if false #axis_arrow
                arrows(scene, [start, to], linewidth=linewidth, linecolor=linecolor, linestyle=linestyle,
                       arrowsize=arrow_size)
            else
                append!(linebuffer, [start, to], linewidth=linewidth, color=linecolor)
            end
        end
    end
    limit_widths = maxi .- mini
    frames = convert_attribute(frames, key"frames"())
    for side in 1:2
        from = Point{N,Float32}(getindex.(limits, side))
        # if axis is drawn at origin, and we draw frame from origin,
        # we already did this
        if !(from == origin && axis_position == :origin)
            for otherside in 1:2
                for dim in 1:N
                    if frames[N - dim + 1][3 - otherside]
                        p = ntuple(i -> i == dim ? limits[i][otherside] : limits[i][side], Val(N))
                        to = Point{N,Float32}(p)
                        append!(linebuffer, [from, to], linewidth=linewidth, color=linecolor)
                    end
                end
            end
        end
    end
end

function draw_titles(textbuffer, xticks, yticks, origin, limit_widths, tickfont, tick_size, tick_gap,
                     tick_title_gap, axis_labels, textcolor, textsize, rotation, align, font, title)
    tickspace_x = maximum(map(yticks) do tick
                              str = last(tick)
                              tick_bb = text_bb(str, to_font(tickfont[2]), tick_size[2])
                              return widths(tick_bb)[1]
                          end)

    tickspace_y = widths(text_bb(last(first(xticks)), to_font(tickfont[1]), tick_size[1]))[2]

    model_inv = inv(transformationmatrix(textbuffer)[])

    tickspace = transform(model_inv, (tickspace_x, tickspace_y))
    title_start = origin .- (tick_gap .+ tickspace .+ tick_title_gap)
    half_width = origin .+ (limit_widths ./ 2.0)

    posx = (half_width[1], title_start[2])
    posy = (title_start[1], half_width[2])
    positions = (posx, posy)
    for i in 1:2
        if !isempty(axis_labels[i])
            push!(textbuffer, axis_labels[i], positions[i], textsize=textsize[i], align=align[i],
                  rotation=rotation[i], color=textcolor[i], font=font[i])
        end
    end
    if title !== nothing
        # TODO give title own text attributes
        push!(textbuffer, title, (half_width[1], (origin .+ (limit_widths))[2]), textsize=textsize[1],
              align=(:center, :top), rotation=rotation[1], color=textcolor[1], font=font[1])
    end
end

function ticks_and_labels(x)
    st, s = extrema(x)
    r = LinRange(st, s, 4)
    return zip(r, string.(round.(r, 4)))
end

function transform(model::Mat4, x::T) where {T}
    x4d = to_ndim(Vec4f0, x, 0.0)
    return to_ndim(T, model * x4d, 0.0)
end
un_transform(model::Mat4, x) = transform(inv(model), x)

to2tuple(x) = ntuple(i -> x, Val(2))
to2tuple(x::Tuple{<:Any,<:Any}) = x

function draw_axis2d(textbuffer, frame_linebuffer, grid_linebuffer, tickmarks_linebuffer, m, padding, limits,
                     xyrange_labels, showgrid, showticks, showtickmarks,
                     # grid attributes
                     g_linewidth, g_linecolor, g_linestyle,

                     # tick attributes
                     t_linewidth, t_linecolor, t_linestyle, t_textcolor, t_textsize, t_rotation, t_align,
                     t_font, t_gap, t_title_gap,

                     #tickmark attribures
                     tm_linewidth, tm_linecolor, tm_linestyle, tm_length,

                     # frame attributes
                     f_linewidth, f_linecolor, f_linestyle, f_axis_position, f_axis_arrow, f_arrow_size,
                     f_frames,

                     # title / axis name attributes
                     ti_labels, ti_textcolor, ti_textsize, ti_rotation, ti_align, ti_font, ti_title)
    xyrange, labels = xyrange_labels
    start!(textbuffer)
    start!(frame_linebuffer)
    foreach(start!, grid_linebuffer)
    foreach(start!, tickmarks_linebuffer)
    # limits = limits Vec2f0(padding)
    # limits ((xmin, xmax), (ymin, ymax))
    limit_widths = map(x -> x[2] - x[1], limits)
    pad = (limit_widths .* to2tuple(padding))
    # pad the drawn limits
    limits = map((lim, p) -> (lim[1] - p, lim[2] + p), limits, pad)

    # recalculate widths
    limit_widths = map(x -> x[2] - x[1], limits)

    % = mean(limit_widths) / 100 # percentage
    xyticks = zip.(xyrange, labels)
    model_inv = inv(transformationmatrix(textbuffer)[])

    ti_textsize = ti_textsize .* %
    t_textsize = t_textsize .* %
    t_gap = transform(model_inv, to2tuple(t_gap .* %))
    tm_length = transform(model_inv, to2tuple(tm_length .* %))
    t_title_gap = transform(model_inv, to2tuple(t_title_gap .* %))

    origin = first.(limits)
    dirs = ((0.0, Float64(limit_widths[2])), (Float64(limit_widths[1]), 0.0))
    foreach(1:2, dirs, xyticks) do dim, dir, ticks
        if showgrid[dim]
            draw_grid(grid_linebuffer[dim], dim, origin, ticks, dir, g_linewidth, g_linecolor, g_linestyle)
        end
    end
    tm_offsets = ((0.0, Float64(tm_length[2])), (Float64(tm_length[1]), Float64(0.0)))
    foreach(1:2, tm_offsets, xyticks) do dim, offset, ticks
        if showtickmarks[dim]
            draw_tickmarks(tickmarks_linebuffer[dim], dim, origin .- offset, ticks, offset, t_linewidth,
                           t_linecolor, t_linestyle, t_textcolor, t_textsize, t_rotation, t_align, t_font)
        end
    end
    t_offsets = ((0.0, Float64(t_gap[2] + tm_length[2])), (Float64(t_gap[1] + tm_length[1]), Float64(0.0)))
    foreach(1:2, t_offsets, xyticks) do dim, offset, ticks
        if showticks[dim]
            draw_ticks(textbuffer, dim, origin .- offset, ticks, t_linewidth, t_linecolor, t_linestyle,
                       t_textcolor, t_textsize, t_rotation, t_align, t_font)
        end
    end

    draw_frame(frame_linebuffer, limits, f_linewidth, f_linecolor, f_linestyle, f_axis_position, f_axis_arrow,
               f_arrow_size, f_frames)

    draw_titles(textbuffer, xyticks..., origin, limit_widths, t_font, t_textsize, t_gap, t_title_gap,
                ti_labels, ti_textcolor, ti_textsize, ti_rotation, ti_align, ti_font, ti_title)
    finish!(textbuffer)
    finish!(frame_linebuffer)
    foreach(finish!, grid_linebuffer)
    foreach(finish!, tickmarks_linebuffer)
    return
end

# for axis, we don't want to have plot!(scene, args) called on it, so we need to overload it directly
function plot!(scene::SceneLike, ::Type{<:Axis2D}, attributes::Attributes, args...)
    # create "empty" plot type - empty meaning containing no plots, just attributes + arguments
    cplot = Axis2D(scene, attributes, args)
    # Disable any non linear transform for the axis plot!
    cplot.transformation.transform_func[] = identity
    g_keys = (:linewidth, :linecolor, :linestyle)
    f_keys = (:linewidth, :linecolor, :linestyle, :axis_position, :axis_arrow, :arrow_size, :frames)
    t_keys = (:linewidth, :linecolor, :linestyle, :textcolor, :textsize, :rotation, :align, :font, :gap,
              :title_gap)
    tm_keys = (:linewidth, :linecolor, :linestyle, :length)
    ti_keys = (:axisnames, :textcolor, :textsize, :rotation, :align, :font, :title)

    g_args = getindex.(cplot.grid, g_keys)
    f_args = getindex.(cplot.frame, f_keys)
    t_args = getindex.(cplot.ticks, t_keys)
    tm_args = getindex.(cplot.tickmarks, tm_keys)
    ti_args = getindex.(cplot.names, ti_keys)

    textbuffer = TextBuffer(cplot, Point{2})

    grid_linebuffer = Node((LinesegmentBuffer(cplot, Point{2}, transparency=true,
                                              linestyle=lift(first, cplot[:grid, :linestyle])),
                            LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                              linestyle=lift(last, cplot[:grid, :linestyle]))))

    frame_linebuffer = Node(LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                              linestyle=cplot.frame.linestyle))
    tickmarks_linebuffer = Node((LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                                   linestyle=lift(first, cplot[:tickmarks, :linestyle])),
                                 LinesegmentBuffer(cplot, Point{2}; transparency=true,
                                                   linestyle=lift(last, cplot[:tickmarks, :linestyle]))))
    map_once(draw_axis2d, Node(textbuffer), frame_linebuffer, grid_linebuffer, tickmarks_linebuffer,
             transformationmatrix(scene), cplot.padding, cplot[1], cplot.ticks.ranges_labels,
             lift.((dim2,), (cplot.showgrid, cplot.showticks, cplot.showtickmarks))..., g_args..., t_args...,
             tm_args..., f_args..., ti_args...)
    push!(scene, cplot)
    return cplot
end

function labelposition(ranges, dim, dir, tgap, origin::StaticVector{N}) where {N}
    a, b = extrema(ranges[dim])
    whalf = Float32(((b - a) / 2))
    halfaxis = unit(Point{N,Float32}, dim) .* whalf

    return origin .+ (halfaxis .+ (normalize(dir) * tgap))
end

_widths(x::Tuple{<:Number,<:Number}) = x[2] - x[1]
_widths(x) = Float32(maximum(x) - minimum(x))

to3tuple(x::Tuple{Any}) = (x[1], x[1], x[1])
to3tuple(x::Tuple{Any,Any}) = (x[1], x[2], x[2])
to3tuple(x::Tuple{Any,Any,Any}) = x
to3tuple(x) = ntuple(i -> x, Val(3))

function draw_axis3d(textbuffer, linebuffer, limits, ranges_labels, args...)
    # make sure we extend all args to 3D
    ranges, ticklabels = ranges_labels
    args3d = to3tuple.(args)
    (showaxis, showticks, showgrid, axisnames, axisnames_color, axisnames_size, axisrotation, axisalign, axisnames_font, titlegap, gridcolors, gridthickness, axiscolors, ttextcolor, trotation, ttextsize, talign, tfont, tgap) = args3d # splat to names

    N = 3
    start!(textbuffer)
    start!(linebuffer)
    mini, maxi = first.(limits), last.(limits)

    origin = Point{N,Float32}(min.(mini, first.(ranges)))
    limit_widths = max.(last.(ranges), maxi) .- origin
    % = minimum(limit_widths) / 100 # percentage
    ttextsize = (%) .* ttextsize
    axisnames_size = (%) .* axisnames_size

    titlegap = (%) .* titlegap
    tgap = (%) .* tgap
    for i in 1:N
        axis_vec = unit(Point{N,Float32}, i)
        width = Float32(limit_widths[i])
        stop = origin .+ (width .* axis_vec)
        if showaxis[i]
            append!(linebuffer, [origin, stop], color=axiscolors[i], linewidth=1.5f0)
        end
        if showticks[i]
            range = ranges[i]
            j = mod1(i + 1, N)
            tickdir = unit(Point{N,Float32}, j)
            tickdir, offset2 = if i != 2
                tickdir = unit(Vec{N,Float32}, j)
                tickdir, Float32(limit_widths[j] + tgap[i]) * tickdir
            else
                tickdir = unit(Vec{N,Float32}, 1)
                tickdir, Float32(limit_widths[1] + tgap[i]) * tickdir
            end
            for (j, tick) in enumerate(range)
                labels = ticklabels[i]
                if length(labels) >= j
                    str = labels[j]
                    if !isempty(str)
                        startpos = (origin .+ ((Float32(tick - origin[i]) * axis_vec)) .+ offset2)
                        push!(textbuffer, str, startpos, color=ttextcolor[i], rotation=trotation[i],
                              textsize=ttextsize[i], align=talign[i], font=tfont[i])
                    end
                end
            end
            if !isempty(axisnames[i])
                tick_widths = if length(ticklabels[i]) >= 3
                    widths(text_bb(ticklabels[i][end - 1], to_font(tfont[i]), ttextsize[i]))[1]
                else
                    0.0f0
                end
                pos = (labelposition(ranges, i, tickdir, titlegap[i] + tick_widths, origin) .+ offset2)
                push!(textbuffer, to_latex(axisnames[i]), pos, textsize=axisnames_size[i],
                      color=axisnames_color[i], rotation=axisrotation[i], align=axisalign[i],
                      font=axisnames_font[i])
            end
        end
        if showgrid[i]
            c = gridcolors[i]
            thickness = gridthickness[i]
            for _j in (i + 1):(i + N - 1)
                j = mod1(_j, N)
                dir = unit(Point{N,Float32}, j)
                range = ranges[j]
                for tick in range
                    offset = Float32(tick - origin[j]) * dir
                    append!(linebuffer, [origin .+ offset, stop .+ offset], color=c, linewidth=thickness)
                end
            end
        end
        finish!(textbuffer)
        finish!(linebuffer)
    end
    return
end

function plot!(scene::SceneLike, ::Type{<:Axis3D}, attributes::Attributes, args...)
    axis = Axis3D(scene, attributes, args)
    # Disable any non linear transform for the axis plot!
    axis.transformation.transform_func[] = identity
    textbuffer = TextBuffer(axis, Point{3}, transparency=true)
    linebuffer = LinesegmentBuffer(axis, Point{3}, transparency=true)

    tstyle, ticks, frame = to_value.(getindex.(axis, (:names, :ticks, :frame)))
    titlevals = getindex.(tstyle, (:axisnames, :textcolor, :textsize, :rotation, :align, :font, :gap))
    framevals = getindex.(frame, (:linecolor, :linewidth, :axiscolor))
    tvals = getindex.(ticks, (:textcolor, :rotation, :textsize, :align, :font, :gap))
    args = (getindex.(axis, (:showaxis, :showticks, :showgrid))..., titlevals..., framevals..., tvals...)
    map_once(draw_axis3d, Node(textbuffer), Node(linebuffer), axis[1], axis.ticks.ranges_labels, args...)
    push!(scene, axis)
    return axis
end


text:
# =========================================================================== #
# Types
# =========================================================================== #

abstract type DistributionType
end

struct UniformDistribution{T} <: DistributionType
    a::T
    b::T
end
function UniformDistribution(a::T, b::T) where {T}
    a = min(a, b)
    b = max(a, b)
    UniformDistribution{T}(a, b)
end
function UniformDistribution() where {T}
    UniformDistribution{T}(0.0, 1.0)
end

struct NormalDistribution{T} <: DistributionType
    mu::T
    sigma::T
end
function NormalDistribution(mu::T, sigma::T) where {T}
    if sigma < 0
        error("sigma must be positive")
    end
    NormalDistribution{T}(mu, sigma)
end
function NormalDistribution(;mu = 0.0, sigma = 1.0)
    NormalDistribution(mu, sigma)
end

struct CauchyDistribution{T} <: DistributionType
    mu::T
    c::T
end
function CauchyDistribution(mu::T, c::T) where {T}
    CauchyDistribution{T}(mu, c)
end
function CauchyDistribution(;mu = 0.0, c = 1.0)
    CauchyDistribution(mu, c)
end

struct MultiNormalDistribution{T} <: DistributionType
    mu::Vector{T}
    sigma::Matrix{T}
    L::LowerTriangular{T, Matrix{T}}
    inv_sigma::Matrix{T}
end
function MultiNormalDistribution(mu::Vector{T}, sigma::Matrix{T}) where {T}
    n_dim = length(mu)
    sz = size(sigma)
    if sz[1] != sz[2]
        error("Non-square sigma matrix")
    end
    if sz[1] != n_dim
        error("dimension mismatch between mu and sigma")
    end

    fact = cholesky(sigma)

    id = Matrix{T}( one(T) * I, n_dim, n_dim )
    inv_sigma = fact.L \ ( fact.U \ id )
    MultiNormalDistribution{T}(mu, sigma, fact.L, inv_sigma)
end

max_stars_repo_name: HomoModelicus/julia

max_issues_repo_name: HomoModelicus/julia

max_forks_repo_name: HomoModelicus/julia


text:
####
# Batch sampling methods
# function to sample b new data points
# optimise for  b << b_max with b_max all unseen test datapoints
###

"""
    thompson_sampling(f_pre::GPpredict,b)

Brute force Thompson sampling of all unseen data obtain b new samples.
Use GPpredict object, which contains the GP_prediction model and the unseen design points

"""

function ts_sampler(f_pre::GPpredict,b)
    # some checks
    @assert length(f_pre.x_test) ≥ b
    # index selected
    selected = Int[]
    #estimate distributions
    Ndist = marginals(f_pre.f̂_pred)
    #b randomsample form every distribution and set to matrix
    f̂_matrix = map(x->rand(x,b),Ndist) |> x-> hcat(x...)  |> permutedims
    # get max form every column
    # avoid retake, previous selected = -inf
    for n in eachcol(f̂_matrix)
        n[selected]  .= - Inf64
       i = argmax(n)
       push!(selected, i)
   end
   return f_pre.x_test[selected]
end

"""
    ei_sampler(f_pre::GPpredict,fmax,b;ϵ = 0)

Brute force Expected improvement (EI) to obtain b samples
from the non-evaluated
design space.
These means that for all unseen combination the  EI is calculated and the best b
constructs are returned
Uses a GPpredict object,
which contains the GP_prediction model and the unseen design points

ϵ is a paramters to balance between exploring and exploiting.
higher ϵ values will explore the search space more.
"""

function ei_sampler(f_pre::GPpredict,b,fmax;ϵ = 0)
    @assert length(f_pre.x_test) ≥ b
    ei_values = ei(f_pre,fmax,;ϵ = 0) |> x -> sort( x,:EI,rev=true)
    return ei_values.x_test[1:b]
end

"""
    pi_sampler(f_pre::GPpredict,fmax,b;ϵ = 0)

Brute force Probability of improvement (PI) to obtain b samples
from the non-evaluated
design space.
The  PI is calculated for all unseen construct, and the b constructs
with the highest value are returned.
Uses a GPpredict object,
which contains the GP_prediction model and the unseen design points

ϵ is a paramters to balance between exploring and exploiting.
higher ϵ values will explore the search space more.
"""
function pi_sampler(f_pre::GPpredict,b,fmax;ϵ = 0)
    @assert length(f_pre.x_test) ≥ b
    pi_values = pi(f_pre,fmax;ϵ = 0) |> x -> sort( x ,:PI,rev=true)
    return pi_values.x_test[1:b]
end

"""
    gpucb_sampler(f_pre::GPpredict,b,β)

Brute force calculation of the Gaussian Process Upper Confidence Bound (GP-UCB)
to obtain b samples from the non-evaluated design space.
The  (GP-UCB) is calculated for all unseen construct,
and the b constructs with the highest value are returned.
Uses a GPpredict object,
which contains the GP_prediction model and the unseen design points

β is a paramters to balance between exploring and exploiting.
higher β values will explore the search space more.
"""
function gpucb_sampler(f_pre::GPpredict,b;β = 1)
    @assert length(f_pre.x_test) ≥ b
    gpubc_values = gp_ubc(f_pre::GPpredict,β) |> x -> sort( x ,:UCB,rev=true)
    return gpubc_values.x_test[1:b]
end

function ucb_sampler(f_pre::GPpredict,λ,b)
    @assert length(f_pre.x_test) ≥ b
    ubc_values = ubc(f_pre::GPpredict,λ) |> x -> sort( x ,:UCB,rev=true)
    return ubc_values.x_test[1:b]
end

"""
    optimial_β(optimial_β(n,t;δ= 0.1)

  Optimal β for gp_ucb for sequential samler
  with n = number of training datapoints
  and t = iteration number
  δ ∈ [0,1]

  Gaussian Process Optimization in the Bandit Setting: No Regret and Experimental Design
  N. Srinivas, A. Krause, S.M. Kakade, M. Seeger.
  arXiv e-prints, pp. arXiv:0912.3995. 2009.
"""
optimial_β(n,t;δ = 0.2) = 2*log(n*t^2*π^2/(6δ))

######
# auxilary function for the acquisition function
# all formals can be found on
######

"""
    ei(f_pre::GPpredict,fmax,;ϵ = 0)
Return the expect impovent value for the give predicted GP model
ϵ is a paramters to balance between exploring and exploiting.
higher ϵ values will explore the search space more.
"""

function ei(f_pre::GPpredict,fmax;ϵ = 0)
    # estimate distribution
    Ndist = marginals(f_pre.f̂_pred)
    σ_t = std.(Ndist)
    μ_t = mean.(Ndist)
    # calculate z
    z =  (μ_t .- fmax.-ϵ) ./ σ_t
    EI = ((μ_t .- fmax .-ϵ) .*cdf.(Normal(),z)) .+ σ_t .* pdf.(Normal(),z)
    return DataFrame(EI = EI, x_test = f_pre.x_test)
end

"""
    pi(f_pre::GPpredict,fmax,;ϵ = 0)
Return the probability of improvent value for the give predicted GP model
ϵ is a paramters to balance between exploring and exploiting.
higher ϵ values will explore the search space more.
"""

function pi(f_pre,fmax;ϵ = 0)
    # estimate distribution
    Ndist = marginals(f_pre.f̂_pred)
    σ_t = std.(Ndist)
    μ_t = mean.(Ndist)
    # calculate z
    z =  (μ_t .- fmax.-ϵ) ./ σ_t
    PI = cdf.(Normal(),z)
    return  DataFrame(PI = PI, x_test = f_pre.x_test)
end

"""
    ubc(f_pre::GPpredict,λ)

Return the upper bound confience intravel for given prediction

"""

function ubc(f_pre::GPpredict,λ)
        Ndist = marginals(f_pre.f̂_pred)
        σ_t = std.(Ndist)
        μ_t = mean.(Ndist)
        UCB = μ_t .+ λ.*σ_t
        return DataFrame(UCB = UCB, x_test = f_pre.x_test)
end
"""
    gp_ubc(f_pre::GPpredict,λ)

Return the gp- upper bound confience intravel for given predictions.

"""
gp_ubc(f_pre::GPpredict,β) =  ubc(f_pre::GPpredict,sqrt(β))

#####
# need?
######

"""
    thompson_sampling(f_pre::GPpredict,b)

Brute force Thompson sampling of all unseen data obtain b new samples.
Use GPpredict object,
which contains the GP_prediction model and the unseen design points

"""
function ts_sampler_stheno(f_pre::GPpredict,b)
    @assert length(f_pre.x_test) ≥ b
    # indices selected
    selected = Int[]
    f̂_matrix = rand(f_pre.f̂_pred,b)
    for j in eachcol(f̂_matrix)
        j[selected]  .= - Inf64
       i = argmax(j)
       push!(selected, i)
   end
   return f_pre.x_test[selected]
end

"""
    save_thompson_sampling(f, x::AbstractVector,n_samples, σ²)

A brute force Thompson sampling algorithm that uses the initer undiscovered search space.
Save version that uses the eigenvalue decomposition of the posterior covariance matrix to assure that the covariance matrix is positive definite
"""
function save_thompson_sampling(f_post, constructs_test::AbstractVector,
                                            n_samples, σ²_opt)
    myeigen = (-(eigvals!(cov(f_post(constructs_test)))[1])) + σ²_opt
    if  myeigen > 0
        println("usetrick")
        sigma_add = myeigen+σ²_opt
        nextsample = thompson_sampling(f_post, constructs_test, n_samples,
                                            sigma_add)
    else
        nextsample = thompson_sampling(f_post, constructs_test, n_samples,
                                                σ²_opt)
    end
    return nextsample
end

####

max_stars_repo_name: SamDier/BOMoD

max_issues_repo_name: SamDier/BOMoD

max_forks_repo_name: SamDier/BOMoD


text:
mutable struct Args
    η::Float32 # learning rate
    ρ::Float32 # regularization paramater (for data augmentation)
    batch_size::Int64 # batch size
    train_prop::Float32 # % of train data to be used for training (validation set is 1 - train_prop)
    epochs::Int64 # number of epochs
    seed::Int64 # random seed
    cuda::Bool  # attempt to use GPU
    input_dims::Array{Int64}  # input size
    nclasses::Int64  # classes
    lr_patience::Int64  # non-improving iterations before learning rate drop
    γ::Float32  # amount to drop lr (1/γ)
    convergence::Int64 # non-improving iterations to quit
    shuffle::Bool #whether to shuffle data before splitting
    scale::Bool #whether to apply Z score scaling
    tblogging::Bool  # use tensorboard
    save_model::Bool
    save_path::String  # results path
    train_dir::String
    testdir::String
    Args() = new(
        1e-3, 1e-1, 32, 0.8, 10, 0, true, [9, 128, 1], 6, 5, 10.0, 10, true, true, false, true, "output_" * Dates.format(now(), "Y-mm-dd-HMS"), DIR*"/../data/train", DIR*"/../data/test"
    )
end


text:
#
# This file is part of the DiscreteEvents.jl Julia package, MIT license
#
# Paul Bayer, 2020
#
# This is a Julia package for discrete event simulation
#
using DiscreteEvents, Distributions, Test

println("... timed events ...")

a = [0.0]
b = [0.0]
c = [0.0]
d = [0.0]
incr!(x) = x[1] += 1

clk = Clock()

foreach(i -> event!(clk, fun(incr!, a), i), 1:5) # 1st case
@test length(clk.sc.events) == 5
foreach(i -> event!(clk, fun(incr!, a), Normal(i, 0)), 6:10) # 2nd case
@test length(clk.sc.events) == 10
event!(clk, fun(incr!, b), 1, 1) # 3rd case
@test length(clk.sc.events) == 11
event!(clk, fun(incr!, c), 1, Normal(1, 0)) # 4th case
@test length(clk.sc.events) == 12
event!(clk, fun(incr!, d), Normal(1, 0), Normal(1, 0)) # 5th case
@test length(clk.sc.events) == 13

run!(clk, 10)
@test a[1] == 10
@test b[1] == 10
@test c[1] == 10
@test d[1] == 10

# test events with Timing
resetClock!(clk)
a[1] = b[1] = c[1] = d[1] = 0
foreach(i -> event!(clk, fun(incr!, a), at, i), 1:5)
foreach(i -> event!(clk, fun(incr!, a), at, Normal(i,0)), 6:10)
foreach(i -> event!(clk, fun(incr!, b), after, i), 1:5)
foreach(i -> event!(clk, fun(incr!, b), after, Normal(i,0)), 6:10)
event!(clk, fun(incr!, c), every, 1)
event!(clk, fun(incr!, d), every, Normal(1,0))
@test length(clk.sc.events) == 22

run!(clk, 10)
@test a[1] == 10
@test b[1] == 10
@test c[1] == 10
@test d[1] == 10

# the same with default clock
resetClock!(𝐶)
a[1] = b[1] = c[1] = d[1] = 0
foreach(i -> event!(fun(incr!, a), i), 1:5) # 1st case
foreach(i -> event!(fun(incr!, a), Normal(i, 0)), 6:10) # 2nd case
event!(fun(incr!, b), 1, 1) # 3rd case
event!(fun(incr!, c), 1, Normal(1, 0)) # 4th case
event!(fun(incr!, d), Normal(1, 0), Normal(1, 0)) # 5th case
@test length(𝐶.sc.events) == 13

run!(𝐶, 10)
@test a[1] == 10
@test b[1] == 10
@test c[1] == 10
@test d[1] == 10

# test events with Timing and default clock
resetClock!(𝐶)
a[1] = b[1] = c[1] = d[1] = 0
foreach(i -> event!(fun(incr!, a), at, i), 1:5)
foreach(i -> event!(fun(incr!, a), at, Normal(i,0)), 6:10)
foreach(i -> event!(fun(incr!, b), after, i), 1:5)
foreach(i -> event!(fun(incr!, b), after, Normal(i,0)), 6:10)
event!(fun(incr!, c), every, 1)
event!(fun(incr!, d), every, Normal(1,0))
@test length(𝐶.sc.events) == 22

run!(𝐶, 10)
@test a[1] == 10
@test b[1] == 10
@test c[1] == 10
@test d[1] == 10

println("... conditional events ...")
is(x, y) = x[1] == y
tim(x, clk) = x[1] = clk.time
resetClock!(clk)
a[1] = b[1] = c[1] = d[1] = 0
event!(clk, fun(incr!, a), every, 1)
event!(clk, fun(tim, b, clk), (fun(is, d, 3), fun(x->x[1]≥5, a)))

# mock clk.state = Busy
state = clk.state
clk.state = DiscreteEvents.Busy()
event!(clk, fun(incr!, c), (fun(is, a, 0), fun(is, b, 0))) # execute immediately
@test c[1] == 1
clk.state = state

event!(clk, fun(event!, clk, fun(tim, d, clk), fun(is, c, 1)), 3) # execute immediately at 3

@test length(clk.sc.cevents) == 1  # only the 1st cevent is scheduled
@test DiscreteEvents._evaluate(clk.sc.cevents[1].cond) == (false, false)
@test DiscreteEvents._nextevent(clk).t == 1
clk.Δt = 0
event!(clk, fun(incr!, c), fun(≥, ()->tau(clk), 3))
@test clk.Δt == 0.01

run!(clk, 6)
@test tau(clk) == 6
@test a[1] == 6
@test b[1] ≈ 5 + clk.Δt            # 1st cevent must have fired
@test c[1] == 2                    # 2nd cevent must have fired
@test d[1] == 3
@test length(clk.sc.cevents) == 0  # and has disappeared
@test length(clk.sc.events) == 1   # but repeat event is still scheduled

resetClock!(clk)
a[1] = 0
event!(clk, ()->a[1] += 1, every, 1, n=100) # test repeat event count
run!(clk, 200)
@test clk.time == 200
@test clk.evcount == 100
@test a[1] == 100

println("... sampling ...")
clk = Clock(1)  # clock with sample rate 1
periodic!(clk, fun(incr!, a), 0)
@test clk.Δt == 0.01
clk.time = 10
clk.evcount = 1000
periodic!(clk, fun(incr!, a), 0)
@test clk.Δt == 0.0001

clk = Clock(1)  # clock with sample rate 1
@test clk.time == 0
@test clk.tn == 1
@test clk.Δt == 1

b[1] = 0
periodic!(clk, fun(incr!, b))
@test length(clk.sc.samples) == 1
run!(clk, 10)
@test clk.time == 10
@test b[1] == 10
sample_time!(clk, 0.5)
run!(clk, 10)
@test clk.time == 20
@test b[1] == 30

println("... everything events and sampling ...")
a[1] = b[1] = 0
clk = Clock(0.5)
event!(clk, fun(incr!, a), every, 0.5)
event!(clk, fun(incr!, a), every, TruncatedNormal(1, 0.3, 0.5, 2))
event!(clk, fun(event!, clk, fun(incr!, a), fun((c, x)->c.time ≥ x, clk, clk.time+2rand())), every, Exponential(5))
periodic!(clk, fun(incr!, b))
run!(clk, 10000)
@test a[1] == clk.evcount
@test b[1] == 20000

println("... everything events and sampling with 𝐶 ...")
a[1] = b[1] = 0
resetClock!(𝐶)
sample_time!(0.5)
event!(fun(incr!, a), every, 0.5)
event!(fun(incr!, a), every, TruncatedNormal(1, 0.3, 0.5, 2))
event!(fun(event!, fun(incr!, a), fun(x->tau() ≥ x, tau()+2rand())), every, Exponential(5))
periodic!(fun(incr!, b))
run!(clk, 10000)
@test a[1] == clk.evcount
@test b[1] == 20000


text:
#
println("Dep) Test of the AutoIonization and PlasmaShift modules with ASF from an internally generated initial- and final-state multiplet.")
#
setDefaults("print summary: open", "zzz-AutoIonization-Plasma.sum")
setDefaults("method: continuum, Galerkin")
setDefaults("method: normalization, pure sine")   ## setDefaults("method: normalization, pure Coulomb")    setDefaults("method: normalization, pure sine")

wa = Atomic.Computation("xx",  Nuclear.Model(10.); grid=JAC.Radial.Grid("grid: by given parameters"; rnt = 2.0e-6, h = 5.0e-2, hp = 1.0e-2, NoPoints = 600),
                        initialConfigs=[Configuration("1s 2s^2 2p^6")],
                        finalConfigs  =[Configuration("1s^2 2s^2 2p^4"), Configuration("1s^2 2s 2p^5"), Configuration("1s^2 2p^6")],
                        process = JAC.AugerInPlasma,
                        processSettings = PlasmaShift.AutoIonizationSettings(PlasmaShift.DebyeHueckel, 0.25, 0., 0, true, true, Tuple{Int64,Int64}[(1,0)]) )

wb = perform(wa)
setDefaults("print summary: close", "")


text:
import Brillouin
import Brillouin.KPaths: Bravais

@doc raw"""
Extract the high-symmetry ``k``-Point path corresponding to the passed model
using `Brillouin.jl`. Uses the conventions described in the reference work by
Cracknell, Davies, Miller, and Love (CDML). Of note, this has minor differences to
the ``k``-path reference
([Y. Himuma et. al. Comput. Mater. Sci. **128**, 140 (2017)](https://doi.org/10.1016/j.commatsci.2016.10.015))
underlying the path-choices of `Brillouin.jl`, specifically for oA and mC Bravais types.

Issues a warning in case the passed lattice does not match the expected primitive.
"""
function high_symmetry_kpath(model; kline_density=20)
    # Change units from Number of kpoints per inverse Angström (i.e. unit of length)
    # to number of kpoints per inverse bohrs.
    # TODO Change interface to atomic units ...
    kline_density = kline_density * austrip(1u"Å")

    if model.n_dim == 1  # Return fast for 1D model
        # TODO Is this special-casing of 1D is not needed for Brillouin.jl any more
        #
        # Just use irrfbz_path(1, DirectBasis{1}([1.0]))
        # (see https://github.com/JuliaMolSim/DFTK.jl/pull/496/files#r725205860)
        #
        # Length of the kpath is recip_lattice[1, 1] in 1D
        n_points = max(2, 1 + ceil(Int, kline_density * model.recip_lattice[1, 1]))
        kcoords  = [@SVector[coord, 0, 0] for coord in range(-1//2, 1//2, length=n_points)]
        klabels  = Dict("Γ" => zeros(3), "-½" => [-0.5, 0.0, 0.0], "½" => [0.5, 0, 0])
        return (kcoords=kcoords, klabels=klabels,
                kpath=[["-½", "½"]], kbranches=[1:length(kcoords)])
    end

    # - Brillouin.jl expects the input direct lattice to be in the conventional lattice
    #   in the convention of the International Table of Crystallography Vol A (ITA).
    # - spglib uses this convention for the returned conventional lattice,
    #   so it can be directly used as input to Brillouin.jl
    # - The output k-Points and reciprocal lattices will be in the CDML convention.
    conv_latt = get_spglib_lattice(model; to_primitive=false)
    sgnum     = spglib_spacegroup_number(model)  # Get ITA space-group number
    direct_basis   = Bravais.DirectBasis(collect(eachcol(conv_latt)))
    primitive_latt = Bravais.primitivize(direct_basis, Bravais.centering(sgnum, 3))

    primitive_latt ≈ collect(eachcol(model.lattice)) || @warn(
        "DFTK's model.lattice and Brillouin's primitive lattice do not agree. " *
        "The kpath selected to plot the band structure might not be most appropriate."
    )

    kp     = Brillouin.irrfbz_path(sgnum, direct_basis)
    kinter = Brillouin.interpolate(kp, density=kline_density)

    # TODO Need to take care of time-reversal symmetry here!
    #      See https://github.com/JuliaMolSim/DFTK.jl/pull/496/files#r725203554

    # Need to double the points whenever a new path starts
    # (for temporary compatibility with pymatgen)
    # TODO Remove this later
    kcoords = empty(first(kinter.kpaths))
    for kbranch in kinter.kpaths
        idcs = findall(k -> any(sum(abs2, k - kcomp) < 1e-5
                                for kcomp in values(kp.points)), kbranch)
        @assert length(idcs) ≥ 2
        idcs = idcs[2:end-1]  # Don't duplicate first and last
        idcs = sort(append!(idcs, 1:length(kbranch)))
        append!(kcoords, kbranch[idcs])
    end

    T = eltype(kcoords[1])
    klabels = Dict{String,Vector{T}}(string(key) => val for (key, val) in kp.points)
    kpath   = [[string(el) for el in path] for path in kp.paths]
    (; kcoords, klabels, kpath)
end

@timing function compute_bands(basis, ρ, kcoords, n_bands;
                               eigensolver=lobpcg_hyper,
                               tol=1e-3,
                               show_progress=true,
                               kwargs...)
    # Create basis with new kpoints, where we cheat by not using any symmetry operations.
    ksymops     = [[identity_symop()] for _ in 1:length(kcoords)]
    bs_basis    = PlaneWaveBasis(basis, kcoords, ksymops)
    ham = Hamiltonian(bs_basis; ρ=ρ)

    band_data = diagonalize_all_kblocks(eigensolver, ham, n_bands + 3;
                                        n_conv_check=n_bands,
                                        tol=tol, show_progress=show_progress, kwargs...)
    if !band_data.converged
        @warn "Eigensolver not converged" iterations=band_data.iterations
    end
    merge((basis=bs_basis, ), select_eigenpairs_all_kblocks(band_data, 1:n_bands))
end

function split_into_branches(kcoords, data::Dict, klabels::Dict)
    # kcoords in cartesian coordinates, klabels uses cartesian coordinates
    function getlabel(kcoord; tol=1e-4)
        findfirst(c -> norm(c - kcoord) < tol, klabels)
    end

    branches = Any[(kindices = [0], kdistances=[0.0], ), ]
    for (ik, kcoord) in enumerate(kcoords)
        previous_kcoord = ik == 1 ? kcoords[1] : kcoords[ik - 1]
        if !isnothing(getlabel(kcoord)) && !isnothing(getlabel(previous_kcoord))
            # New branch encountered
            previous_distance = branches[end].kdistances[end]
            push!(branches, (kindices=[ik], kdistances=[previous_distance]))
        else
            # Keep adding to current branch
            distance = branches[end].kdistances[end] + norm(kcoord - kcoords[ik - 1])
            push!(branches[end].kdistances, distance)
            push!(branches[end].kindices, ik)
        end
    end

    map(branches[2:end]) do branch
        branch_data = Dict(key => data[key][branch.kindices, :, :] for key in keys(data))
        ret = (klabels=(getlabel(kcoords[branch.kindices[1]]),
                        getlabel(kcoords[branch.kindices[end]])),
               kdistances=branch.kdistances,
               kindices=branch.kindices)
        merge(ret, (; branch_data...))
    end
end

function prepare_band_data(band_data; datakeys=[:λ, :λerror],
                           klabels=Dict{String, Vector{Float64}}())
    basis = band_data.basis
    n_spin   = basis.model.n_spin_components
    n_kcoord = length(basis.kpoints) ÷ n_spin
    n_bands  = nothing

    # Convert coordinates to Cartesian
    kcoords_cart = [basis.kpoints[ik].coordinate_cart for ik in krange_spin(basis, 1)]
    klabels_cart = Dict(lal => basis.model.recip_lattice * vec for (lal, vec) in klabels)

    # Split data into branches
    data = Dict{Symbol, Any}()
    for key in datakeys
        hasproperty(band_data, key) || continue
        n_bands = length(band_data[key][1])
        data_per_kσ = similar(band_data[key][1], n_kcoord, n_bands, n_spin)
        for σ in 1:n_spin, (ito, ik) in enumerate(krange_spin(basis, σ))
            data_per_kσ[ito, :, σ] = band_data[key][ik]
        end
        data[key] = data_per_kσ
    end
    @assert !isnothing(n_bands)
    branches = split_into_branches(kcoords_cart, data, klabels_cart)

    tick_labels    = String[branches[1].klabels[1]]
    tick_distances = Float64[branches[1].kdistances[1]]
    for (i, br) in enumerate(branches)
        # Ignore branches with a single k-point
        branches[i].klabels[1] == branches[i].klabels[2] && continue

        label    = branches[i].klabels[2]
        distance = branches[i].kdistances[end]
        if i != length(branches) && branches[i+1].klabels[1] != label
            # Next branch is not continuous from the current
            label = label * " | " * branches[i+1].klabels[1]
        end
        push!(tick_labels, label)
        push!(tick_distances, distance)
    end

    (branches=branches, ticks=(distances=tick_distances, labels=tick_labels),
     n_bands=n_bands, n_kcoord=n_kcoord, n_spin=n_spin, basis=basis)
end

"""
    is_metal(band_data, εF, tol)

Determine whether the provided bands indicate the material is a metal,
i.e. where bands are cut by the Fermi level.
"""
function is_metal(band_data, εF, tol=1e-4)
    n_bands = length(band_data.λ[1])
    n_kpoints = length(band_data.λ)
    for ib in 1:n_bands
        some_larger = any(band_data.λ[ik][ib] - εF < -tol for ik in 1:n_kpoints)
        some_smaller = any(εF - band_data.λ[ik][ib] < -tol for ik in 1:n_kpoints)
        some_larger && some_smaller && return true
    end
    false
end

"""
Compute and plot the band structure. `n_bands` selects the number of bands to compute.
If this value is absent and an `scfres` is used to start the calculation a default of
`n_bands_scf + 5sqrt(n_bands_scf)` is used. Unlike the rest of DFTK bands energies
are plotted in eV unless a different `unit` (any Unitful unit) is selected.
"""
function plot_bandstructure(basis, ρ, n_bands;
                            εF=nothing, kline_density=20, unit=u"eV", kwargs...)
    mpi_nprocs() > 1 && error("Band structures with MPI not supported yet")
    if !isdefined(DFTK, :PLOTS_LOADED)
        error("Plots not loaded. Run 'using Plots' before calling plot_bandstructure.")
    end

    # Band structure calculation along high-symmetry path
    kcoords, klabels, kpath = high_symmetry_kpath(basis.model; kline_density)
    println("Computing bands along kpath:")
    println("       ", join(join.(kpath, " -> "), "  and  "))
    band_data = compute_bands(basis, ρ, kcoords, n_bands; kwargs...)

    plotargs = ()
    if kline_density ≤ 10
        plotargs = (markersize=2, markershape=:circle)
    end

    plot_band_data(band_data; εF=εF, klabels=klabels, unit=unit, plotargs...)
end
function plot_bandstructure(scfres; n_bands=nothing, kwargs...)
    # Convenience wrapper for scfres named tuples
    n_bands_scf = length(scfres.occupation[1])
    isnothing(n_bands) && (n_bands = ceil(Int, n_bands_scf + 5sqrt(n_bands_scf)))
    plot_bandstructure(scfres.basis, scfres.ρ, n_bands; εF=scfres.εF, kwargs...)
end


text:
using Flora.RNN
using MLDatasets # FashionMNIST

# load full test set
test_x,  test_y  = FashionMNIST.testdata() # 10_000

picks = [
    [0, 3, 9], # 0 T-Shirt, 3 Dress, 9 Ankle boot
    [1, 2, 5],
]

for pick in picks
    p = load(:fashion, pick)
    tX, tY = prepare(test_x, test_y, pick)
    show_accuracy(p.m, p.pick, tX, tY)
end

#=
pick = [0, 3, 9]
mean(argmax(m(x)) .== argmax(y)) = 0.929
pick = [1, 2, 5]
mean(argmax(m(x)) .== argmax(y)) = 0.9866666666666667
=#