abstract type AbstractSampleBuf{T, N} <: AbstractArray{T, N} end

"""
Represents a multi-channel regularly-sampled buffer that stores its own sample
rate (in samples/second). The wrapped data is an N-dimensional array. A 1-channel
sample can be represented with a 1D array or an Mx1 matrix, and a C-channel
buffer will be an MxC matrix. So a 1-second stereo audio buffer sampled at
44100Hz with 32-bit floating-point samples in the time domain would have the
type SampleBuf{Float32, 2}.
"""
mutable struct SampleBuf{T, N} <: AbstractSampleBuf{T, N}
    data::Array{T, N}
    samplerate::Float64
end

# define constructor so conversion is applied to `sr`
SampleBuf(arr::Array{T, N}, sr::Real) where {T, N} = SampleBuf{T, N}(arr, sr)

"""
Represents a multi-channel regularly-sampled buffer representing the frequency-
domain spectrum of a `SampleBuf`. The wrapped data is an N-dimensional array. A
1-channel sample can be represented with a 1D array or an Mx1 matrix, and a
C-channel buffer will be an MxC matrix. So a 1-second stereo audio buffer
sampled at 44100Hz with 32-bit floating-point samples in the time domain would
have the type SampleBuf{Float32, 2}.
"""
mutable struct SpectrumBuf{T, N} <: AbstractSampleBuf{T, N}
    data::Array{T, N}
    samplerate::Float64
end

# define constructor so conversion is applied to `sr`
SpectrumBuf(arr::Array{T, N}, sr::Real) where {T, N} = SpectrumBuf{T, N}(arr, sr)

SampleBuf(T::Type, sr, dims...) = SampleBuf(Array{T}(undef, dims...), sr)
SpectrumBuf(T::Type, sr, dims...) = SpectrumBuf(Array{T}(undef, dims...), sr)
SampleBuf(T::Type, sr, len::Quantity) = SampleBuf(T, sr, inframes(Int,len,sr))
SampleBuf(T::Type, sr, len::Quantity, ch) =
    SampleBuf(T, sr, inframes(Int,len,sr), ch)
SpectrumBuf(T::Type, sr, len::Quantity) =
    SpectrumBuf(T, sr, inframes(Int,len, sr))
SpectrumBuf(T::Type, sr, len::Quantity, ch) =
    SpectrumBuf(T, sr, inframes(Int,len, sr))

# terminology:
# sample - a single value representing the amplitude of 1 channel at some point in time (or frequency)
# channel - a set of samples running in parallel
# frame - a collection of samples from each channel that were sampled simultaneously

# audio methods
samplerate(buf::AbstractSampleBuf) = buf.samplerate
nchannels(buf::AbstractSampleBuf{T, 2}) where {T} = size(buf.data, 2)
nchannels(buf::AbstractSampleBuf{T, 1}) where {T} = 1
nframes(buf::AbstractSampleBuf) = size(buf.data, 1)

function samplerate!(buf::AbstractSampleBuf, sr)
    buf.samplerate = sr

    buf
end

# define audio methods on raw buffers as well
nframes(arr::AbstractArray) = size(arr, 1)
nchannels(arr::AbstractArray) = size(arr, 2)

# it's important to define Base.similar so that range-indexing returns the
# right type, instead of just a bare array
Base.similar(buf::SampleBuf, ::Type{T}, dims::Dims) where {T} = SampleBuf(Array{T}(undef, dims), samplerate(buf))
Base.similar(buf::SpectrumBuf, ::Type{T}, dims::Dims) where {T} = SpectrumBuf(Array{T}(undef, dims), samplerate(buf))
domain(buf::AbstractSampleBuf) = range(0.0, stop=(nframes(buf)-1)/samplerate(buf), length=nframes(buf))

# There's got to be a better way to define these functions, but the dispatch
# and broadcast behavior for AbstractArrays is complex and has subtle differences
# between Julia versions, so we basically just override functions here as they
# come up as problems
import Base: +, -, *, /
import Base.broadcast

find_buf(args::Tuple) = find_buf(find_buf(args[1]), Base.tail(args))
find_buf(x) = x

for btype in (:SampleBuf, :SpectrumBuf)

        @eval find_buf(bc::Base.Broadcast.Broadcasted{Broadcast.ArrayStyle{$btype{T,N}}}) where {T,N} = find_buf(bc.args)
        @eval find_buf(::$btype, args::Tuple{$btype}) = args[1]
        @eval find_buf(::Any, args::Tuple{$btype}) = args[1]
        @eval find_buf(a::$btype, rest) = a

        @eval Base.BroadcastStyle(::Type{$btype{T,N}}) where {T,N} = Broadcast.ArrayStyle{$btype{T,N}}()
        @eval function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{$btype{T,N}}}, ::Type{ElType}) where {T,N,ElType}
            A = find_buf(bc)
            $btype(Array{ElType}(undef, length.(axes(bc))), samplerate(A))
        end

end # for btype

typename(::SampleBuf{T, N}) where {T, N} = "SampleBuf{$T, $N}"
unitname(::SampleBuf) = "s"
srname(::SampleBuf) = "Hz"
typename(::SpectrumBuf{T, N}) where {T, N} = "SpectrumBuf{$T, $N}"
unitname(::SpectrumBuf) = "Hz"
srname(::SpectrumBuf) = "s"

# from @mbauman's Sparklines.jl package
const ticks = ['▁','▂','▃','▄','▅','▆','▇','█']
# 3-arg version (with explicit mimetype) is needed because we subtype AbstractArray,
# and there's a 3-arg version defined in show.jl
function show(io::IO, ::MIME"text/plain", buf::AbstractSampleBuf)
    println(io, "$(nframes(buf))-frame, $(nchannels(buf))-channel $(typename(buf))")
    len = nframes(buf) / samplerate(buf)
    ustring = unitname(buf)
    srstring = srname(buf)
    print(io, "$(len)$ustring sampled at $(samplerate(buf))$srstring")
    nframes(buf) > 0 && showchannels(io, buf)
end

function showchannels(io::IO, buf::AbstractSampleBuf, widthchars=80)
    # number of samples per block
    blockwidth = round(Int, nframes(buf)/widthchars, RoundUp)
    nblocks = round(Int, nframes(buf)/blockwidth, RoundUp)
    blocks = Array{Char}(undef, nblocks, nchannels(buf))
    for blk in 1:nblocks
        i = (blk-1)*blockwidth + 1
        n = min(blockwidth, nframes(buf)-i+1)
        peaks = Compat.maximum(abs.(float.(buf[(1:n) .+ i .- 1, :])), dims=1)
        # clamp to -60dB, 0dB
        peaks = clamp.(20log10.(peaks), -60.0, 0.0)
        idxs = trunc.(Int, (peaks.+60)/60 * (length(ticks)-1)) .+ 1
        blocks[blk, :] = ticks[idxs]
    end
    for ch in 1:nchannels(buf)
        println(io)
        print(io, String(blocks[:, ch]))
    end
end

"""Get a pointer to the underlying data for the buffer. Will return a Ptr{T},
where T is the element type of the buffer. This is particularly useful for
passing to C libraries to fill the buffer"""
channelptr(buf::Array, channel, frameoffset=0) =
    pointer(buf) + ((channel-1)*nframes(buf)+frameoffset) * sizeof(eltype(buf))
channelptr(buf::AbstractSampleBuf, channel, frameoffset=0) =
    channelptr(buf.data, channel, frameoffset)

"""Mix the channels of the source array into the channels of the dest array,
using coefficients from the `mix` matrix. To mix an M-channel buffer to a
N-channel buffer, `mix` should be MxN. `src` and `dest` should not share
memory."""
function mix!(dest::AbstractMatrix, src::AbstractMatrix, mix::AbstractArray)
    inchans = nchannels(src)
    outchans = nchannels(dest)
    size(mix) == (inchans, outchans) || error("Mix Matrix should be $(inchans)x$(outchans)")
    mul!(dest, src, mix)
end

function mix!(dest::AbstractVector, src::AbstractVector, mix::AbstractArray)
    mix!(reshape(dest, (length(dest), 1)), reshape(src, (length(src), 1)), mix)
    dest
end

function mix!(dest::AbstractVector, src::AbstractMatrix, mix::AbstractArray)
    mix!(reshape(dest, (length(dest), 1)), src, mix)
    dest
end

function mix!(dest::AbstractMatrix, src::AbstractVector, mix::AbstractArray)
    mix!(dest, reshape(src, (length(src), 1)), mix)
end

"""Mix the channels of the source array into the channels of the dest array,
using coefficients from the `mix` matrix. To mix an M-channel buffer to a
N-channel buffer, `mix` should be MxN. `src` and `dest` should not share
memory."""
function mix(src::AbstractArray, mix::AbstractArray)
    dest = similar(src, (nframes(src), size(mix, 2)))
    mix!(dest, src, mix)
end

"""Mix the channels of the `src` array into the mono `dest` array."""
function mono!(dest::AbstractArray, src::AbstractArray)
    mix!(dest, src, ones(nchannels(src), 1) ./ nchannels(src))
end

"""Mix the channels of the `src` array into a mono array."""
function monoch(src::AbstractArray)
    dest = similar(src, (nframes(src), 1))
    mono!(dest, src)
end

# the index types that Base knows how to handle. Separate out those that index
# multiple results
const BuiltinMultiIdx = Union{Colon,
                        Vector{Int},
                        Vector{Bool},
                        AbstractRange{Int}}
const BuiltinIdx = Union{Int, BuiltinMultiIdx}
# the index types that will need conversion to built-in index types. Each of
# these needs a `toindex` method defined for it
const ConvertIdx{T1 <: Quantity, T2 <: Int} = Union{T1,
                                                # Vector{T1}, # not supporting vectors of Quantities (yet?)
                                                # Range{T1}, # not supporting ranges (yet?)
                                                ClosedInterval{T2},
                                                ClosedInterval{T1}}

"""
    toindex(buf::SampleBuf, I)

Convert the given index value to one that Base knows how to use natively for
indexing
"""
function toindex end

toindex(buf::SampleBuf, t::Number) = t
toindex(buf::SampleBuf, t::FrameQuant) = inframes(Int, t) + 1
toindex(buf::SampleBuf, t::Unitful.Time) = inframes(Int, t, samplerate(buf)) + 1
toindex(buf::SampleBuf, t::Unitful.AbstractQuantity) = throw(Unitful.DimensionError(t, s))
toindex(buf::SpectrumBuf, f::Number) = f
toindex(buf::SpectrumBuf, f::FrameQuant) = inframes(Int, f) + 1
toindex(buf::SpectrumBuf, f::Unitful.Frequency) = inframes(Int, f, samplerate(buf)) + 1
toindex(buf::SpectrumBuf, f::Unitful.AbstractQuantity) = throw(Unitful.DimensionError(f, Hz))

# indexing by vectors of Quantities not yet supported
toindex(buf::AbstractSampleBuf, I::ClosedInterval{Int}) =
    toindex(buf, minimum(I)*frames):toindex(buf, maximum(I)*frames)
toindex(buf::AbstractSampleBuf, I::ClosedInterval{T}) where {T <: Quantity} =
    toindex(buf, minimum(I)):toindex(buf, maximum(I))

# AbstractArray interface methods
Base.size(buf::AbstractSampleBuf) = size(buf.data)
Base.IndexStyle(::Type{T}) where {T <: AbstractSampleBuf} = Base.IndexLinear()
# this is the fundamental indexing operation needed for the AbstractArray interface
Base.getindex(buf::AbstractSampleBuf, i::Int) = buf.data[i];

Base.getindex(buf::AbstractSampleBuf, I::ConvertIdx) = buf[toindex(buf, I)]
Base.getindex(buf::AbstractSampleBuf, I1::ConvertIdx, I2::BuiltinIdx) =
    buf[toindex(buf, I1), I2]
# In Julia 0.5 scalar indices are now dropped, so by default indexing
# buf[5, 1:2] gives you a 2-frame single-channel buffer instead of a 1-frame
# two-channel buffer. The following getindex method defeats the index dropping
Base.getindex(buf::AbstractSampleBuf, I1::Int, I2::BuiltinMultiIdx) = buf[I1:I1, I2]

function Base.setindex!(buf::AbstractSampleBuf, val, i::Int)
    buf.data[i] = val
end

# equality
import Base.==
==(buf1::AbstractSampleBuf, buf2::AbstractSampleBuf) =
    samplerate(buf1) == samplerate(buf2) &&
    buf1.data == buf2.data

FFTW.fft(buf::SampleBuf) = SpectrumBuf(FFTW.fft(buf.data), nframes(buf)/samplerate(buf))
FFTW.ifft(buf::SpectrumBuf) = SampleBuf(FFTW.ifft(buf.data), nframes(buf)/samplerate(buf))

# does a per-channel convolution on SampleBufs
for buftype in (:SampleBuf, :SpectrumBuf)
    @eval function DSP.conv(b1::$buftype{T, 1}, b2::$buftype{T, 1}) where {T}
        if !isapprox(samplerate(b1), samplerate(b2))
            error("Resampling convolution not yet supported")
        end
        $buftype(conv(b1.data, b2.data), samplerate(b1))
    end

    @eval function DSP.conv(b1::$buftype{T, N1}, b2::$buftype{T, N2}) where {T, N1, N2}
        if !isapprox(samplerate(b1), samplerate(b2))
            error("Resampling convolution not yet supported")
        end
        if nchannels(b1) != nchannels(b2)
            error("Broadcasting convolution not yet supported")
        end
        out = $buftype(T, samplerate(b1), nframes(b1)+nframes(b2)-1, nchannels(b1))
        for ch in 1:nchannels(b1)
            out[:, ch] = conv(b1.data[:, ch], b2.data[:, ch])
        end

        out
    end

    @eval function DSP.conv(b1::$buftype{T, 1}, b2::StridedVector{T}) where {T}
        $buftype(conv(b1.data, b2), samplerate(b1))
    end

    @eval DSP.conv(b1::StridedVector{T}, b2::$buftype{T, 1}) where {T} = conv(b2, b1)

    @eval function DSP.conv(b1::$buftype{T, 2}, b2::StridedMatrix{T}) where {T}
        if nchannels(b1) != nchannels(b2)
            error("Broadcasting convolution not yet supported")
        end
        out = $buftype(T, samplerate(b1), nframes(b1)+nframes(b2)-1, nchannels(b1))
        for ch in 1:nchannels(b1)
            out[:, ch] = conv(b1.data[:, ch], b2[:, ch])
        end

        out
    end

    @eval DSP.conv(b1::StridedMatrix{T}, b2::$buftype{T, 2}) where {T} = conv(b2, b1)
end


import Oceananigans.Grids: required_halo_size

"""
    ScalarBiharmonicDiffusivity{F, N, K}

Holds viscosity and diffusivities for models with prescribed isotropic diffusivities.
"""
struct ScalarBiharmonicDiffusivity{F, N, K} <: AbstractScalarBiharmonicDiffusivity{F}
    ν :: N
    κ :: K

    function ScalarBiharmonicDiffusivity{F}(ν::N, κ::K) where {F, N, K}
        return new{F, N, K}(ν, κ)
    end
end

# Aliases that allow specify the floating type, assuming that the discretization is Explicit in time
          ScalarBiharmonicDiffusivity(FT::DataType;         kwargs...) = ScalarBiharmonicDiffusivity(ThreeDimensionalFormulation(), FT; kwargs...)
  VerticalScalarBiharmonicDiffusivity(FT::DataType=Float64; kwargs...) = ScalarBiharmonicDiffusivity(VerticalFormulation(), FT; kwargs...)
HorizontalScalarBiharmonicDiffusivity(FT::DataType=Float64; kwargs...) = ScalarBiharmonicDiffusivity(HorizontalFormulation(), FT; kwargs...)

required_halo_size(::ScalarBiharmonicDiffusivity) = 2

"""
    ScalarBiharmonicDiffusivity(FT=Float64; νh=0, κh=0, νz=nothing, κz=nothing)

Returns parameters for a scalar biharmonic diffusivity model.

Keyword arguments
=================

  - `ν`: Viscosity. `Number`, `AbstractArray`, or `Function(x, y, z, t)`.

  - `κ`: Diffusivity. `Number`, `AbstractArray`, or `Function(x, y, z, t)`, or
         `NamedTuple` of diffusivities with entries for each tracer.

  - `discrete_form`: `Boolean`.

  - `formulation`: formulation used for the discretization of the diffusivity operator.
                   Options are `VerticalFormulation`, `HorizontalFormulation` and
                   `ThreeDimensionalFormulation`.

"""
function ScalarBiharmonicDiffusivity(formulation=ThreeDimensionalFormulation(), FT=Float64;
                                     ν=0, κ=0,
                                     discrete_form = false)

    ν = convert_diffusivity(FT, ν, Val(discrete_form))
    κ = convert_diffusivity(FT, κ, Val(discrete_form))
    return ScalarBiharmonicDiffusivity{typeof(formulation)}(ν, κ)
end

function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F}) where {F}
    κ = tracer_diffusivities(tracers, closure.κ)
    return ScalarBiharmonicDiffusivity{F}(closure.ν, κ)
end

@inline viscosity(closure::ScalarBiharmonicDiffusivity, args...) = closure.ν
@inline diffusivity(closure::ScalarBiharmonicDiffusivity, ::Val{tracer_index}, args...) where tracer_index = closure.κ[tracer_index]

calculate_diffusivities!(diffusivities, closure::ScalarBiharmonicDiffusivity, args...) = nothing

Base.show(io::IO, closure::ScalarBiharmonicDiffusivity{F}) where {F} =
    print(io, "ScalarBiharmonicDiffusivity: " *
              "(ν=$(closure.ν), κ=$(closure.κ)), " *
              "formulation: $F")


module Simulacrum

using HyperGraphs, Symbolics
import SymbolicUtils: Term, Sym, Symbolic
using SimpleTraits
using ModelingToolkit
using Catalyst

const SymTypes = Union{Num, Symbolic}
export SymTypes

include("patterns.jl")
include("bridges/catalyst.jl")
include("connections.jl")
include("operations.jl")
include("reactions.jl")
include("utils.jl")
include("wip.jl")

export

# bridges: catalyst
catalyst_reactions,

# connections
connect, connect!, couple,

# operations
clone,

# patterns
AbstractPattern, indices,
AllToAll, Cycle, Path,

# reactions
isunimolecular, ismultimolecular,
reversible_reaction, @reversible,

# utils
make_var

end # module


using BSON: @save
using Flux.Tracker: data

function get_replay_batch(pos_buffer, π_buffer, res_buffer; batch_size = 32)
  idxs = sample(1:length(pos_buffer), batch_size, replace=false)

  pos_replay = pos_buffer[idxs]
  π_replay = hcat(π_buffer[idxs]...)
  res_replay = res_buffer[idxs]

  return pos_replay, π_replay, res_replay
end

function save_model(nn::NeuralNet)
  # checking for path to save
  path = joinpath(dirname(@__DIR__), "models")
  !isdir(path) && mkdir(path)

  @save path * "/agz_model.bson", nn

  # saving weights
  # checking for path to save weights
  path = path * "/weights"
  !isdir(path) && mkdir(path)

  base_net, value, policy = nn.base_net, nn.value, nn.policy

  bn_weights  = cpu.(data.(params(base_net)))
  val_weights = cpu.(data.(params(value)))
  pol_weights = cpu.(data.(params(policy)))

  @save path * "/agz_base.bson" bn_weights
  @save path * "/agz_value.bson" val_weights
  @save path * "/agz_policy.bson" pol_weights
end

function train(env::GameEnv; num_games::Int = 25000, memory_size::Int = 500000,
  batch_size::Int = 32, epochs = 1, ckp_freq::Int = 1000, readouts::Int = 800,
  tower_height::Int = 19, model = nothing, start_training_after = 50000)

  # @assert 0 ≤ tower_height ≤ 19
  cur_nn = model === nothing ? NeuralNet(env; tower_height = tower_height) : model

  # prev_nn = deepcopy(cur_nn)

  pos_buffer = Vector{Position}()
  π_buffer   = Vector{Vector{Float32}}()
  res_buffer = Vector{Int}()

  push_data(vec::Vector{T}, data::T) where T = vcat(vec, data)
  shrink(vec::Vector{T}) where T = vec[end-memory_size+1:end]

  opt = Momentum(2f-2)

  for i = 1:num_games
    player  = selfplay(env, cur_nn, readouts)
    p, π, v = extract_data(player)

    pos_buffer, π_buffer, res_buffer =
      push_data.((pos_buffer, π_buffer, res_buffer), (p, π, v))

    if length(pos_buffer) > memory_size
      pos_buffer, π_buffer, res_buffer = shrink.((pos_buffer, π_buffer, res_buffer))
    end

    if length(pos_buffer) >= start_training_after
      replay_pos, replay_π, replay_res = get_replay_batch(pos_buffer, π_buffer, res_buffer;
							  batch_size = batch_size)
      loss = _train(cur_nn, (replay_pos, replay_π, replay_res), opt; epochs = epochs)
      result = player.result_string
      num_moves = player.root.position.n
      println("Episode $i over. Loss: $loss. Winner: $result. Moves: $num_moves.")
    end

  # cur_is_winner = evaluate(cur_nn, prev_nn; num_games = eval_games, ro = readouts)
  # print("Evaluated. ")
  # if cur_is_winner
  #   prev_nn = deepcopy(cur_nn)
  #   print("Model updated. ")
  # else
  #   cur_nn = deepcopy(prev_nn)
  #   print("Model retained. ")
  # end

    if i % ckp_freq == 0
      save_model(cur_nn)
      print("Model saved. ")
    end
  end
  return cur_nn
end


module TrajOptPlots

using MeshCat
# using GeometryTypes
using GeometryBasics
using Rotations
using CoordinateTransformations
using TrajectoryOptimization
using RecipesBase
using StaticArrays
using FileIO
using MeshIO
using RobotDynamics
using RobotZoo
using Colors

import RobotDynamics: orientation

export
    visualize!,
    waypoints!,
    traj3!

include("meshes.jl")
include("rigid_bodies.jl")
include("2d_vis.jl")
include("models.jl")

end # module


using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()
cd(@__DIR__)
using Revise
using ZigZagBoomerang
const ZZB = ZigZagBoomerang

using ZigZagBoomerang: Trace, sevent, waiting_time_ref, freezing_time!, ab,
     smove_forward!, grad_correct!, λ, sλ̄, reflect_sticky!, freezing_time,
     refresh_sticky_vel!, sep
using DataStructures
using LinearAlgebra
#using Observables
using Colors
using Random
using SparseArrays
using Test
using FileIO
using Dates
using Statistics
@time using Makie, AbstractPlotting
println(now())
using ForwardDiff
const ρ0 = 0.0
const d = 2
const dist = 1.
const R = 3.0
 r = -R:0.05:R
#=
mutable struct Level
    story
    gradϕ!
    logdensity
    potential
    ys
    minpoints
    T
    Target
    User
    θ0
    x0
end
=#
mutable struct Game
    kill_on_boundary
    pressed
    released
    ys
    canvas
    score
    speed
    timeleft
    auto
    potential
end
ZigZagBoomerang.λ(∇ϕx, θ, F::Union{BouncyParticle, Boomerang}) = ZZB.pos(dot(∇ϕx, θ)) + 0.0

function bouncy_inner!(game, X, Ξ, ∇ϕ!, ∇ϕx, t, x, θ, c, b, t′, f, θf, tfrez, tref, told, (acc, num),
    Flow::Union{ZigZag, BouncyParticle, Boomerang}, κ, args...; strong_upperbounds = false, factor=1.5, adapt=false)
    canvas = game.canvas
    factorized = Flow isa ZigZag
    while true
        ispressed(canvas.scene, Keyboard.q) && error("end")
        for (key, pressed) in game.pressed
            if pressed && !ispressed(canvas.scene, key)
                game.pressed[key] = false
            end
        end
        for (key, released) in game.released
            if released && ispressed(canvas.scene, key)
                game.released[key] = false
            end
        end

        finished = remove(game, x)
        finished |= ispressed(canvas.scene, Keyboard.enter)
        if ispressed(canvas.scene, Keyboard.a) && !game.pressed[Keyboard.a]
            game.auto[] = !game.auto[]
            game.pressed[Keyboard.a] = true
        end
        X[] = [point(POTENTIAL[], x)]
        sleep(0.002)
        yield()
        tᶠ, i = findmin(tfrez) # could be implemented with a queue
        tt, j = findmin([tref, tᶠ, t′...])
        τ = tt - t
        t, x, θ = smove_forward!(τ, t, x, θ, f, Flow)
        # move forward
        if j == 1 # refreshments of velocities
            if (!game.auto[] && ispressed(canvas.scene, Keyboard.c)) || (game.auto[] && !ispressed(canvas.scene, Keyboard.c))
                θ, θf = refresh_sticky_vel!(θ, θf, f, Flow)
            elseif ispressed(canvas.scene, Keyboard.equal)
                θ *= 1.1
                θf *= 1.1
            elseif ispressed(canvas.scene, Keyboard.minus)
                θ *= 0.9
                θf *= 0.9
            end
            tref = t + waiting_time_ref(Flow) # regenerate refreshment time
            b = ab(x, θ, c, Flow) # regenerate reflection time
            told = t
            t′ = t + poisson_times(b)
            tfrez = freezing_time!(tfrez, t, x, θ, f, Flow)
            for i in eachindex(f) # make function later...
                if !f[i]
                    tfrez[i] = t - log(rand())/(κ[i]*abs(θf[i]))
                end
            end
            game.speed[] = norm(θ)
        elseif j == 2 # get frozen or unfrozen in i
            if f[i] # if free
                if abs(x[i]) > 1e-8
                    tfrez[i] = t + freezing_time(x[i], θ[i], Flow.μ[i], Flow) # wrong zero of curve
                    error("x[i] = $(x[i]) !≈ 0 at $(tfrez[i])")
                end
                x[i] = -0*θ[i]
                θf[i], θ[i] = θ[i], 0.0 # stop and save speed
                f[i] = false # change tag
                tfrez[i] = t - log(rand())/(κ[i]*abs(θf[i])) # sticky time
                # tfrez[i] = t - log(rand()) # option 2
                if !(strong_upperbounds) #not strong upperbounds, draw new waiting time
                    b = ab(x, θ, c, Flow) # regenerate reflection time
                    told = t
                    t′ = t .+ poisson_times(b)
                end
                game.released[Keyboard.s] = false
            elseif !ispressed(canvas.scene, Keyboard.s)  && !game.released[Keyboard.s] # is frozen ->  unfreeze
                @assert x[i] == 0 && θ[i] == 0
                θ[i], θf[i] = θf[i], 0.0 # restore speed
                f[i] = true # change tag
                tfrez[i] = t + freezing_time(x[i], θ[i], Flow.μ[i], Flow)
                b = ab(x, θ, c, Flow) # regenerate reflection time
                told = t
                t′ = t .+ poisson_times(b)
                game.released[Keyboard.s] = true
            else
                tfrez[i] = t - log(rand())/(κ[i]*abs(θf[i])) # sticky time
            end
        elseif factorized
            i = j - 2

            ∇ϕx = ∇ϕ!(∇ϕx, x, args...)
            l, lb = ZZB.sλ(∇ϕx[i], i, x, θ, Flow), ZZB.sλ̄(b[i], t - told)
            num += 1

            left = (θ[1]*θ[2] > 0) ⊻ (i == 2)

            if ((game.auto[] || !inbounds(game, x)) && rand()*lb <= l ) ||  ( left && ispressed(canvas.scene, Keyboard.left) && !game.pressed[Keyboard.left]) ||
                ( !left && ispressed(canvas.scene, Keyboard.right) && !game.pressed[Keyboard.right])
                if ispressed(canvas.scene, Keyboard.right)
                    game.pressed[Keyboard.right] = true
                end
                if ispressed(canvas.scene, Keyboard.left)
                    game.pressed[Keyboard.left] = true
                end
                acc += 1
                if l > lb
                    !adapt && error("Tuning parameter `c` too small. l/lb = $(l/lb)")
                    acc = num = 0
                    adapt!(c, i, factor)
                end
                θ = ZZB.reflect!(i, ∇ϕx[i], x, θ, Flow)
                b = ab(x, θ, c, Flow)
                told = t
                t′ = t .+ poisson_times(b)
                tfrez = freezing_time!(tfrez, t, x, θ, f, Flow)
   #             if !inbounds(game, x)
    #                game.score[] = max(0, game.score[]-1)
     #           end
            else
                b = ab(x, θ, c, Flow)
                told = t
                t′ = t .+ poisson_times(b)
                finished || continue
            end
        else #   t′ usual bouncy particle / boomerang step
            ∇ϕx = ∇ϕ!(∇ϕx, x, args...)
            ∇ϕx = grad_correct!(∇ϕx, x, Flow)
            l, lb = λ(∇ϕx, θ, Flow), sλ̄(b, t - told) # CHECK if depends on f
            num += 1
            if ((game.auto[] || !inbounds(game, x)) && rand()*lb <= l ) ||  (ispressed(canvas.scene, Keyboard.space) && !game.pressed[Keyboard.space]) # reflect!
                if ispressed(canvas.scene, Keyboard.space)
                    game.pressed[Keyboard.space] = true
                end
                acc += 1
                if l > lb
                    !adapt && error("Tuning parameter `c` too small.")
                    c *= factor
                end
                θ = reflect_sticky!(∇ϕx, x, θ, f, Flow)
                b = ab(x, θ, c, Flow) # regenerate reflection time
                told = t
                t′ = t + poisson_times(b)
                tfrez = freezing_time!(tfrez, t, x, θ, f, Flow)

#                if !inbounds(game, x)
 #                   game.score[] = max(0, game.score[]-1)
  #              end
            else # nothing happened
                b = ab(x, θ, c, Flow)
                told = t
                t′ = t + poisson_times(b)
                finished || continue
            end
        end

        push!(Ξ, sevent(t, x, θ, f, Flow))
        return finished, t, x, θ, t′, tref, tfrez, told, f, θf, (acc, num), c, b
    end

end
ZZB.Trace(t0, x0, θ0, f, Flow::ZigZag) = ZZB.Trace(t0, x0, θ0, f, BouncyParticle(Flow.Γ, Flow.μ, 0.0))
function ZZB.sevent(t, x, θ, f, Z::ZigZag)
    t, copy(x), copy(θ), copy(f)
end
ZZB.freezing_time!(tfrez, t0, x0, θ0, f, Flow::ZigZag) = ZZB.freezing_time!(tfrez, t0, x0, θ0, f,  BouncyParticle(Flow.Γ, Flow.μ, 0.0))
function ZZB.ab(x, θ, c, Z::ZigZag)
    a = [c + (ZZB.idot(Z.Γ, i, x)  - ZZB.idot(Z.Γ, i, Z.μ))'*θ[i] for i in eachindex(x)]
    b = [c + θ[i]'*ZZB.idot(Z.Γ, i, θ) for i in eachindex(x)]
    a, b
end
poisson_times((a, b)::Tuple{Array{Float64,1},Array{Float64,1}}) = ZZB.poisson_time.(a, b, rand(length(a)))
poisson_times((a, b)) = ZZB.poisson_time(a, b, rand())

function bouncy(game, X, ∇ϕ!, t0, x0, θ0, T, c, Flow::Union{ZigZag, BouncyParticle, Boomerang},
    κ, args...;  strong_upperbounds = false, adapt=false, factor=2.0)
    factorized = Flow isa ZigZag
    X[] .= [point(x0)]
    t, x, θ, ∇ϕx = t0, deepcopy(x0), deepcopy(θ0), deepcopy(θ0)
    told = t0
    θf = 0*θ # tags
    f = [true for _ in eachindex(x)]
    Ξ = Trace(t0, x0, θ0, f, Flow)
    push!(Ξ, sevent(t, x0, θ0, f, Flow))
    tref = waiting_time_ref(Flow) #refreshment times
    tfrez = zero(x)
    tfrez = freezing_time!(tfrez, t0, x0, θ0, f, Flow) #freexing times
    num = acc = 0
    b = ab(x, θ, c, Flow)
    t′ = t .+ poisson_times(b) # reflection time
    finished = false
    while !finished && t < T  && (!game.kill_on_boundary || inbounds(game, 0.9x))
        finished, t, x, θ, t′, tref, tfrez, told, f, θf, (acc, num), c, b = bouncy_inner!(game, X, Ξ, ∇ϕ!, ∇ϕx, t, x, θ, c, b, t′, f, θf, tfrez, tref, told, (acc, num),
                Flow, κ, args...; strong_upperbounds = strong_upperbounds, factor = factor, adapt = adapt)
        game.timeleft[] = T - t
    end
    return Ξ, (t, x, θ), (acc, num), c
end

phi(x, y, rho = 0.0) =  1/(2*pi*sqrt(1-rho^2))*exp(-0.5*(x^2 + y^2 - 2x*y*rho)/(1-rho^2))
sphi(x, y, σ1, σ2 = σ1, ρ=0.0) = phi(x/σ1, y/σ2, ρ)/(σ1*σ2)
logdensity1(x, y) = log(sphi(x, y+0.2x, 0.9))
logdensity2(x, y) = log(0.5sphi(x - dist, y - dist, 0.6) + 0.5sphi(x + dist, y + dist, 0.6))
#potential(xy) = -0.1exp(logdensity(xy[1], xy[2])-log(phi(xy[1], xy[2])))
getpotential(logdensity, boom=false) = function (xy)
     -(2-1.75boom)*exp(logdensity(xy[1], xy[2]) - boom*log(phi(xy[1], xy[2])))
end
getpotentialxy(logdensity, boom=false) = function (x,y)
    -(2-1.75boom)*exp(logdensity(x,y) - boom*log(phi(x, y)))
end

potentials = [getpotentialxy(logdensity1), getpotentialxy(logdensity2)]
potential(xy) = POTENTIAL[](xy[1], xy[2])
#potential(x,y) = POTENTIAL[](x, x)

gradϕ(logdensity) = function (x,i)
    -ForwardDiff.partials(logdensity(ForwardDiff.Dual{}(x[1], 1.0*(i==1)), ForwardDiff.Dual{}(x[2], 1.0*(i==2))))[]
end

gradϕ!(logdensity, g = gradϕ(logdensity)) = function (y, x)
    y[1] = g(x,1)
    y[2] = g(x,2)
    y
end
function remove(game, x)
    ys = game.ys[]
    for i in eachindex(ys)
        if norm(x - ys[i]) < 0.15
            deleteat!(ys, i)
            #Makie.notify(game.ys)
            game.ys[] = game.ys[]
            length(ys) == 0 && return true
            game.score[] = game.score[] + 1
            return false
        end
    end
    return false
end

point(x) = Point3f0(x[1], x[2], 0.03+potential(x))
inbounds(x) = -R < x[1] < R && -R < x[2] < R

point(potential, x) = Point3f0(x[1], x[2], 0.03+potential(x[1],x[2]))
point(i::Int, x) = Point3f0(x[1], x[2], 0.03+potentials[i](x[1],x[2]))
inbounds(game, x) = -R < x[1] < R && -R < x[2] < R

κ = [100.0, 100.0]
t0 = 0.0
x0 = [1.0, 0.]
θ0 = [1.0, 1.0]
c = 100.0
T = 1000.0

trace, (tT, xT, θT), (acc, num) = sspdmp(gradϕ!(logdensity2), t0, x0, θ0, T, c, BouncyParticle(sparse(I(d)), 0*x0, 0.1), 0.01*κ; adapt=false)
ts, xs = sep(collect(discretize(trace, T/500)))

score = Node(0)
SCORE = lift(x->"Score "*string(x), score)
speed = Node(norm(θ0))
timeleft = Node(T)
SPEED = lift(x->"Speed "*string(round(x, digits=2)), speed)
TIMELEFT = lift(y->"Time left "*string(round(y, digits=0)),  timeleft)
auto = Node(false)
AUTO = lift(x -> x ? "Autopilot on" : "Manual", auto)
MESSAGE = Node(" "^10*"Get ready - SPACE to reflect"*" "^10*"\n Q to quit")
ys = Node([rand(2)])
POTENTIAL = Node{Any}(getpotentialxy(logdensity1))
X = Node([point(x0)])
M = lift(p -> [p(x, y) for x in r, y in r], POTENTIAL)
YS = lift((i, ys)->point.(i, ys), POTENTIAL, ys)
ELEMENTS = [TIMELEFT, SPEED, SCORE, AUTO]
MARKER = Node(:circle)
TITLE = Node(" "^20*"Bouncy Particle Sampler Game"*" "^20)

canvas = Figure(resolution=(1500,1000))
#canvas[2:2,1] =  ax0 = Axis(canvas, aspect=DataAspect())
#hidedecorations!(ax0)
canvas[3:10,1] =  ax1 = Axis(canvas, aspect=DataAspect())
hidedecorations!(ax1)
hidespines!(ax1)
canvas[1:2,1:5] =  ax = Axis(canvas, title = "RETURN - continue to next lvl, SPACE - reflect, Q - quit, S* - stick, C* - Crank-Nicolson, A - auto pilot, +-* - speed (*hold)", aspect=DataAspect())
text!(ax, TITLE, align=(:center, :bottom), textsize=.85)
text!(ax, MESSAGE, color=RGB(0.1, 0.3, 0.7), align=(:center, :bottom), textsize=0.65, position=(-0, -2.2))
hidedecorations!(ax)
hidespines!(ax)
lscene = LScene(canvas[3:10, 2:5], scenekw = (camera = cam3d!, raw = false))

using Colors
#surface!(lscene, r, r, [potential((x1, x2)) for x1 in r, x2 in r], colormap=:atlantic, show_axis=false)
contour3d!(lscene, r, r, M, levels=30, linewidth=3.0)

scatter!(lscene, X, marker=MARKER, color=:red, markersize=200)
meshscatter!(lscene, YS, color=:orange, markersize=0.05)
#text!(lscene.scene, SCORE, position=(1.5R, 0R, 1), strokecolor=:black, strokewidth=3.0, color=RGB(0.6, 0.2, 0.0), textsize=1.0, rotation= .75pi)
for i in eachindex(ELEMENTS)
    text!(ax1, ELEMENTS[i], strokecolor=:black, strokewidth=.0, color=RGB(0.6, 0.2, 0.0), position=(0, -2*i), textsize=2.0)
end

display(canvas)

println("Find your Makie window and reflect with the SPACE key, quit with Q")

level = [
    (courage="Good start!", title="Bouncy particle sampler (BPS)", message="Stear the particle through the posterior landscape and collect the sample. \nChange direction with a 'contour reflection' at the countour lines with SPACE. ", logdensity=logdensity1, F=BouncyParticle(sparse(I(d)), 0*x0, 2.0, 0.995), x0=[0.1,1.0], θ0=[1.0,0.3], κ = [100.0, 100.0], T=100.0, ys=[[2.0,2.0]])
    (courage="Improving your ESS!", title="Bouncy particle sampler 2", message="Collect some more posterior points. Randomize your energy to reach\n the center with the Crank-Nicolson scheme (Hold C and reflect with SPACE).", logdensity=logdensity1, F=BouncyParticle(sparse(I(d)), 0*x0, 2.0, 0.995), x0=[0.1,1.0], θ0=[1.0,0.3], T=100.0, ys=0.9*[randn(2) for _ in 1:50])
    (courage="Great mixing!",title="Bouncy particle sampler 3", message="A double well potential! Collect the points.\nContinue with RETURN or try the `autopilot' (the original sampler) with A.", logdensity=logdensity2, F=BouncyParticle(sparse(I(d)), 0*x0, 2.0, 0.995), x0=[0.1,1.0], θ0=[1.0,0.3], κ = [100.0, 100.0], T=100.0, ys=[[0.6*(randn(2)-[dist,dist]) for _ in 1:50];[ 0.6*(randn(2) + [dist,dist]) for _ in 1:50]])

    (courage="Too easy!", title="Zig-Zag", message="Try to collect the point with the Zig-Zag.\n It always moves on the diagonals. Change direction with the arrow keys. \n Continue with RETURN.", logdensity=logdensity1, F=ZigZag(sparse(I(d)), 0*x0), x0=[0.1,1.0], θ0=[1.0,1.0], κ = [100.0, 100.0], T=100.0, ys=[[2.0,2.0]])
    (courage="Almost ergodic!", title="Zig-Zag 2", message="A double well potential! Collect the points. Try also A(utopilot) to see\n how the actual Zig-Zag sampler explores this posterior landscape.", logdensity=logdensity2, F=ZigZag(sparse(I(d)), 0*x0),x0=[0.1,1.0], θ0=[1.0,1.0], κ = [100.0, 100.0], T=1000.0, ys=[[0.6*(randn(2)-[dist,dist]) for _ in 1:50];[ 0.6*(randn(2) + [dist,dist]) for _ in 1:50]])

    (courage="So Hamiltonian!", title="Boomerang", message="The Boomerang moves on circles and does 'contour reflections' with SPACE.\n Change your energy level with +/- to reach the center and the outskirts.", logdensity=logdensity2, F=Boomerang(sparse(I(d)), 0*x0, 2.0, 0.995), x0=[0.1,1.0], θ0=[1.0,0.5], κ = [100.0, 100.0], T=100.0, ys=[[1.4,1.4], [0.1,-.1]])
    (courage="Done. You achieved mastership!", title="Sticky Bouncy particle sampler", message="Some of the points are on the coordinate axes.\nSTICK to the axis with S to collect them all (release to unstick).", logdensity=logdensity2, F=BouncyParticle(sparse(I(d)), 0*x0, 2.0, 0.995), x0=[0.1,1.0], θ0=[1.0,0.5], κ = [100.0, 100.0],   T=100.0, ys=[[0.6*rand(Bool,2).*(randn(2)-[dist,dist]) for _ in 1:50];[ 0.6*rand(Bool,2).*(randn(2) + [dist,dist]) for _ in 1:50]])

    ]
sleep(1.0)
for i in 1:length(level)
    lvl = level[i]
    ∇ϕ! = gradϕ!(lvl.logdensity)
    x0 = lvl.x0
    θ0 = lvl.θ0
    ys[] = lvl.ys
    auto[] = false
    MESSAGE[] = lvl.message
    sleep(2.)
    game = Game(false, Dict(Keyboard.space => false, Keyboard.s => false, Keyboard.left => false, Keyboard.right => false, Keyboard.a => false), Dict(Keyboard.s => false), ys, canvas, score, speed, timeleft, auto, getpotential(lvl.logdensity))
    POTENTIAL[] = getpotentialxy(lvl.logdensity, lvl.F isa Boomerang)
    MARKER[] = :circle
    #sleep(1.5)

    TITLE[] = lvl.title

    trace, (tT, xT, θT), (acc, num) = bouncy(game, X, ∇ϕ!, t0, x0, θ0, lvl.T, c, lvl.F, κ; adapt=false)

   # MARKER[] = :xcross

    MESSAGE[] = "$(lvl.courage) \nCurrent score: $(score[]) posterior\nsamples."
    sleep(1.5)
end


using AtomsBase
# Key functionality to integrate DFTK and AtomsBase

function construct_system(lattice::AbstractMatrix, atoms::Vector, positions::Vector, magnetic_moments::Vector)
    @assert length(atoms) == length(positions)
    @assert isempty(magnetic_moments) || length(magnetic_moments) == length(atoms)
    atomsbase_atoms = map(enumerate(atoms)) do (i, element)
        kwargs = Dict{Symbol, Any}(:potential => element, )
        if !isempty(magnetic_moments)
            kwargs[:magnetic_moment] = normalize_magnetic_moment(magnetic_moments[i])
        end
        if element isa ElementPsp
            kwargs[:pseudopotential] = element.psp.identifier
        end

        position = lattice * positions[i] * u"bohr"
        if atomic_symbol(element) == :X  # dummy element ... should solve this upstream
            Atom(:X, position; atomic_symbol=:X, atomic_number=0, atomic_mass=0u"u", kwargs...)
        else
            Atom(atomic_symbol(element), position; kwargs...)
        end
    end
    periodic_system(atomsbase_atoms, collect(eachcol(lattice)) * u"bohr")
end

function parse_system(system::AbstractSystem{D}) where {D}
    if !all(periodicity(system))
        error("DFTK only supports calculations with periodic boundary conditions.")
    end

    # Parse abstract system and return data required to construct model
    mtx = austrip.(reduce(hcat, bounding_box(system)))
    T = eltype(mtx)
    lattice = zeros(T, 3, 3)
    lattice[1:D, 1:D] .= mtx

    # Cache for instantiated pseudopotentials. This is done to ensure that identical
    # atoms are indistinguishable in memory, which is used in the Model constructor
    # to deduce the atom_groups.
    cached_pspelements = Dict{String, ElementPsp}()
    atoms = map(system) do atom
        if hasproperty(atom, :potential)
            atom.potential
        elseif hasproperty(atom, :pseudopotential)
            get!(cached_pspelements, atom.pseudopotential) do
                ElementPsp(atomic_symbol(atom); psp=load_psp(atom.pseudopotential))
            end
        else
            ElementCoulomb(atomic_symbol(atom))
        end
    end

    positions = map(system) do atom
        coordinate = zeros(T, 3)
        coordinate[1:D] = lattice[1:D, 1:D] \ T.(austrip.(position(atom)))
        Vec3{T}(coordinate)
    end

    magnetic_moments = map(system) do atom
        hasproperty(atom, :magnetic_moment) || return nothing
        getproperty(atom, :magnetic_moment)
    end
    if all(m -> isnothing(m) || iszero(m) || isempty(m), magnetic_moments)
        empty!(magnetic_moments)
    else
        magnetic_moments = normalize_magnetic_moment.(magnetic_moments)
    end

    # TODO Use system to determine n_electrons

    (; lattice, atoms, positions, magnetic_moments)
end

function _call_with_system(f, system::AbstractSystem, args...; kwargs...)
    @assert !(:magnetic_moments in keys(kwargs))
    parsed = parse_system(system)
    f(parsed.lattice, parsed.atoms, parsed.positions, args...;
      parsed.magnetic_moments, kwargs...)
end

"""
    Model(system::AbstractSystem; kwargs...)

AtomsBase-compatible Model constructor. Sets structural information (`atoms`, `positions`,
`lattice`, `n_electrons` etc.) from the passed `system`.
"""
Model(system::AbstractSystem; kwargs...) = _call_with_system(Model, system; kwargs...)

# Generate equivalent functions for AtomsBase
for fun in (:model_atomic, :model_DFT, :model_LDA, :model_PBE, :model_SCAN)
    @eval function $fun(system::AbstractSystem, args...; kwargs...)
        _call_with_system($fun, system, args...; kwargs...)
    end
end

# Extra methods to AtomsBase functions for DFTK data structures
"""
    atomic_system(model::DFTK.Model, magnetic_moments=[])

Construct an AtomsBase atomic system from a DFTK model and associated magnetic moments.
"""
function AtomsBase.atomic_system(model::Model, magnetic_moments=[])
    construct_system(model.lattice, model.atoms, model.positions, magnetic_moments)
end

"""
    periodic_system(model::DFTK.Model, magnetic_moments=[])

Construct an AtomsBase atomic system from a DFTK model and associated magnetic moments.
"""
function AtomsBase.periodic_system(model::Model, magnetic_moments=[])
    atomic_system(model, magnetic_moments)
end

AtomsBase.chemical_formula(model::Model) = chemical_formula(atomic_symbol.(model.atoms))


# high-level interface for CURAND
#
# the interface is split in two levels:
# - functions that extend the Random standard library, and take an RNG as first argument,
#   will only ever dispatch to CURAND and as a result are limited in the types they support.
# - functions that take an array will dispatch to either CURAND or GPUArrays
# - `cu`-prefixed functions are provided for constructing GPU arrays from only an eltype

## seeding

seed!(rng::RNG=generator()) = generate_seeds(rng)

## in-place

# uniform
Random.rand!(rng::RNG, A::CuArray{Float32}) = generate_uniform(rng, A)
Random.rand!(rng::RNG, A::CuArray{Float64}) = generate_uniform_double(rng, A)

# normal
Random.randn!(rng::RNG, A::CuArray{Float32}; mean=0, stddev=1) = generate_normal(rng, A, mean, stddev)
Random.randn!(rng::RNG, A::CuArray{Float64}; mean=0, stddev=1) = generate_normal_double(rng, A, mean, stddev)

# log-normal
rand_logn!(rng::RNG, A::CuArray{Float32}; mean=0, stddev=1) = generate_log_normal(rng, A, mean, stddev)
rand_logn!(rng::RNG, A::CuArray{Float64}; mean=0, stddev=1) = generate_log_normal_double(rng, A, mean, stddev)

# log-normal
rand_poisson!(rng::RNG, A::CuArray{Cuint}; lambda=1) = generate_poisson(rng, A, lambda)

## out of place

Random.rand(rng::RNG, ::Type{X}, dims::Dims) where {X} = rand!(rng, CuArray{X}(undef, dims))
Random.randn(rng::RNG, ::Type{X}, dims::Dims; kwargs...) where {X} = randn!(rng, CuArray{X}(undef, dims); kwargs...)
rand_logn(rng::RNG, ::Type{X}, dims::Dims; kwargs...) where {X} = rand_logn!(rng, CuArray{X}(undef, dims); kwargs...)
rand_poisson(rng::RNG, ::Type{X}, dims::Dims; kwargs...) where {X} = rand_poisson!(rng, CuArray{X}(undef, dims); kwargs...)

# specify default types
Random.rand(rng::RNG, dims::Integer...; kwargs...) = rand(rng, Float32, dims...; kwargs...)
Random.randn(rng::RNG, dims::Integer...; kwargs...) = randn(rng, Float32, dims...; kwargs...)
rand_logn(rng::RNG, dims::Integer...; kwargs...) = rand_logn(rng, Float32, dims...; kwargs...)
rand_poisson(rng::RNG, dims::Integer...; kwargs...) = rand_poisson(rng, Cuint, dims...; kwargs...)

# convenience
Random.randn(rng::RNG, ::Type{X}, dim1::Integer, dims::Integer...; kwargs...) where {X} =
    randn(rng, X, Dims((dim1, dims...)); kwargs...)
rand_logn(rng::RNG, ::Type{X}, dim1::Integer, dims::Integer...; kwargs...) where {X} =
    rand_logn(rng, X, Dims((dim1, dims...)); kwargs...)
rand_poisson(rng::RNG, ::Type{X}, dim1::Integer, dims::Integer...; kwargs...) where {X} =
    rand_poisson(rng, X, Dims((dim1, dims...)); kwargs...)

## functions that dispatch to either CURAND or GPUArrays

uniform_rng(::CuArray{<:Union{Float32,Float64}}) = generator()
uniform_rng(A::CuArray) = GPUArrays.global_rng(A)

normal_rng(::CuArray{<:Union{Float32,Float64}}) = generator()
normal_rng(::CuArray{T}) where {T} =
    error("CuArrays does not support generating normally distributed numbers of type $T")

logn_rng(::CuArray{<:Union{Float32,Float64}}) = generator()
logn_rng(::CuArray{T}) where {T} =
    error("CuArrays does not support generating lognormally distributed numbers of type $T")

poisson_rng(::CuArray{Cuint}) = generator()
poisson_rng(::CuArray{T}) where {T} =
    error("CuArrays does not support generating Poisson distributed numbers of type $T")

Random.rand!(A::CuArray; kwargs...) = rand!(uniform_rng(A), A; kwargs...)
Random.randn!(A::CuArray; kwargs...) = randn!(normal_rng(A), A; kwargs...)
rand_logn!(A::CuArray; kwargs...) = rand_logn!(logn_rng(A), A; kwargs...)
rand_poisson!(A::CuArray; kwargs...) = rand_poisson!(poisson_rng(A), A; kwargs...)
rand_logn(A::CuArray; kwargs...) = rand_logn!(logn_rng(A), A; kwargs...)
rand_poisson(A::CuArray; kwargs...) = rand_poisson!(poisson_rng(A), A; kwargs...)

rand(::Type{X}, args...; kwargs...) where {X} = rand!(CuArray{X}(undef, args...); kwargs...)
randn(::Type{X}, args...; kwargs...) where {X} = randn!(CuArray{X}(undef, args...); kwargs...)
rand_logn(::Type{X}, args...; kwargs...) where {X} = rand_logn!(CuArray{X}(undef, args...); kwargs...)
rand_poisson(::Type{X}, args...; kwargs...) where {X} = rand_poisson!(CuArray{X}(undef, args...); kwargs...)

# specify default types
rand(args...; kwargs...) where {X} = rand(Float32, args...; kwargs...)
randn(args...; kwargs...) where {X} = randn(Float32, args...; kwargs...)
rand_logn(args...; kwargs...) where {X} = rand_logn(Float32, args...; kwargs...)
rand_poisson(args...; kwargs...) where {X} = rand_poisson(Cuint, args...; kwargs...)


"""
Summary statistics.
"""

using CSV
using DataFrames

function sumstats(fh; alleles = false, dropna = true)
    # Parses .sumstats files. See docs/file_formats_sumstats.txt
    #dtype_dict = {"SNP": str,   "Z": float, "N": float, "A1": str, "A2": str}
    compression = get_compression(fh)
    usecols = ["SNP", "Z", "N"]
    if alleles usecols += ["A1", "A2"] end

    try
        x = read_csv(fh, delim_whitespace = True, na_values = ".")
    catch e
        println("Improperly formatted sumstats file:", str(e.args))
    end

    if dropna
        x = x.dropna(how = "any")
    end

    return x
end

""" Estimate h² and partitioned h² """
function estimate_h2()
    args = deepcopy(LDScore.args)

    # TODO: more checking

    M_annot, w_ld_cname, ref_ld_cnames, sumstats, novar_cols = _read_ld_sumstats()

    ref_ld = sumstats[!, "LD"]
    ref_ld = reshape(ref_ld, (size(ref_ld)[1], 1))
    n_snp = size(sumstats)[1]
    n_blocks = min(n_snp, args["n-blocks"])
    n_annot = size(ref_ld_cnames)[1]
    χ²_max = args["χ²-max"]
    old_weights = false

    if n_annot == 1
        if args["two-step"] == nothing && args["intercept-h²"] == nothing
            args["two-step"] = 30
        end
    else
        old_weights = true
        if args["χ²-max"] == nothing
            χ²_max = max(0.001 * max(sumstats[!, "N"]), 80)
        end
    end

    s(x) = reshape(x, (n_snp, 1))
    χ² = s(sumstats[!, "Z"] .^ 2)

    if χ²_max != nothing
        ii = vec(χ² < χ²_max)
        sumstats = sumstats.ix[ii, :]

        n_snp = sum(ii) # TODO double check that this works
        ref_ld = [sumstats[ref_ld_cnames]]
        χ² = reshape(χ²[ii], (n_snp, 1))
    end

    ĥ² = Hsq(
        χ², ref_ld, s(sumstats[!, w_ld_cname]), s(sumstats[!, "N"]), M_annot,
        n_blocks = n_blocks, intercept = args["intercept-h²"],
        slow = false, step1_ii = args["two-step"],
        old_weights = old_weights,
    )

    if args["overlap-annot"]
        overlap_matrix, M_tot = _read_annot(args)

        df_results = ĥ²._overlap_output(
            ref_ld_cnames, overlap_matrix, M_annot, M_tot,
            args["print-coefficients"],
        )
        df_results.to_csv(
            args["out"] * ".results",
            sep = "\t", index = false,
        )
    end

    return ĥ²
end # function estimate_h2

# TODO
function _read_M(args, n_annot) end
function _check_variance(M_annot, ref_ld) end
function _read_w_ld(args) end
function _merge_and_log(ld, sumstats, noun) end

function _read_sumstats(args, fh; alleles=false, dropna=false)
    # Parser summary statistics
    @info "Reading summary statistics from" fh
    sumstats = parse_sumstats(fh, alleles=alleles, dropna=dropna)
    @info "Read summary statistics for" size(sumstats) "SNPs."
    m = size(sumstats)
    sumstats = sumstats.drop_duplicates(subset="SNP")
    if m > size(sumstats)
        @info "Dropped SNPs with duplicated rs numbers:" m-size(sumstats)
    end
    return sumstats
end

function _read_ref_ld()
    # Read reference LD Scores
    ref_ld = _read_chr_split_files(
        LDScore.args["ref-ld-chr"], LDScore.args["ref-ld"],
        "reference panel LD Score", ldscore_fromlist,
    )
    @info "Num. of SNPs read for reference panel LD Scores:" size(ref_ld)
    return ref_ld
end

function _load_testset_1()
    df_ref_id = DataFrames.DataFrame(CSV.File(LDScore.args["ref-ld"] * ".l2.ldscore"))
    df_h2 = DataFrames.DataFrame(CSV.File(LDScore.args["h²"]))

    M_annot = reshape([[155881.2526]], (1, 1))
    w_ld_cname = "CHR"
    ref_ld_cnames = ["LD_0"]
    sumstats = DataFrames.innerjoin(df_ref_id, df_h2, on = "SNP")
    novar_cols = Dict([("LD_0", false)])
    return (M_annot, w_ld_cname, ref_ld_cnames, sumstats, novar_cols)
end

#= TODO: Just loading testing set now,
    need to actually implement to load generic datasets =#
function _read_ld_sumstats(alleles = false, dropna = true)
    return _load_testset_1()
end


using Turing
using Turing: Sampler
using Test

alg = SGHMC(1000, 0.01, 0.1)
sampler = Sampler(alg)
@test isa(alg, SGHMC)
@test isa(sampler, Sampler{<:Turing.SGHMC})

alg = SGHMC(200, 0.01, 0.1, :m)
sampler = Sampler(alg)
@test isa(alg, SGHMC)
@test isa(sampler, Sampler{<:Turing.SGHMC})

alg = SGHMC(1000, 0.01, 0.1, :s)
sampler = Sampler(alg)
@test isa(alg, SGHMC)
@test isa(sampler, Sampler{<:Turing.SGHMC})


module DEBmicroTrait

using FieldMetadata,
      Unitful,
      Roots,
      LinearAlgebra,
      Statistics,
      CSV,
      DataFrames,
      DataFramesMeta,
      Dierckx

import FieldMetadata: @default, @description, @units, @bounds, @logscaled, @flattenable, @plottable, @selectable, @chain,
                      default, description, units, bounds, logscaled, flattenable, plottable, selectable

# Field metadata columns
@chain columns @units @description

using Unitful: J, kJ, W, L, g, kg, cm, m, s, hr, d, mol, mmol, μmol
Unitful.register(@__MODULE__);
@unit bp "bp" Basepair 1u"m" true;

export AbstractMetabolism, AbstractMetabolismC, MetabolismC
export AbstractAssimilation, AbstractAssimilationC, AssimilationC
export AbstractTurnover, Turnover
export AbstractParams, Params
export AbstractIsolate, IsolateComposition
export AbstractSetup, Setup

include("components/metabolism.jl")
include("components/supeca.jl")
include("components/assimilation.jl")
include("components/turnover.jl")
include("components/thermostoichwizard.jl")
include("components/soil_properties.jl")
include("components/depolymerization.jl")
include("allometry/cell_composition.jl")
include("allometry/cell_assimilation.jl")
include("allometry/cell_metabolism.jl")
include("allometry/cell_turnover.jl")
include("allometry/cell_depolymerization.jl")
include("setup.jl")
include("model.jl")

end


module SeisProcessing
    using Interpolations
    using Requires,
    FFTW,
    LinearAlgebra,
    DSP
    import SeisMain

    include("Processing/Processing.jl")
    include("Modelling/Modelling.jl")
    include("Tools/Tools.jl")
end


function v_analytic(m::Model,p::Param,id,it)
    vf = m.v[id,it]
    c = m.c[id,it]
    if it == p.nT
        cons = m.avec
        cash = m.avec
        convert(Point,cash,vfun(id,it,cons,cash,vf,p))
    else
        cons = scaleGrid(p.cfloor_plot,gety(c.env)[2],p.k,logorder = 1)
        cash = scaleGrid(p.a_low,getx(vf.env)[2],p.k,logorder = 1)
        # deleteat!(cons,length(cons))
        # deleteat!(cash,length(cash))
        pts = convert(Point,cash,vfun(id,it,cons,cash,vf,p))
        vcat(pts, vf.env.v[2:end])  # connect at second point
    end
end

function v_analytic(m::BModel,p::Param,id,iy,it)
    vf = m.v[id,iy,it]
    c = m.c[id,iy,it]
    cons = scaleGrid(p.cfloor_plot,gety(c.env)[2],p.k,logorder = 1)
    cash = scaleGrid(p.nT == it ? p.a_lowT : p.a_low,getx(vf.env)[2],p.k,logorder = 1)
    # deleteat!(cons,length(cons))
    # deleteat!(cash,length(cash))
    pts = convert(Point,cash,vfun(x->u(x,id==2,p),it,cons,cash,vf,p))
    vcat(pts, vf.env.v[2:end])  # connect at second point
    # vf.env.v[2:end]  # connect at second point
end

function v_analytic(m::Model,p::Param,id,iy,it)
    vf = m.v[id,iy,it]
    c = m.c[id,iy,it]
    cons = scaleGrid(p.cfloor_plot,gety(c.env)[2],p.k,logorder = 1)
    cash = scaleGrid(p.a_low,getx(vf.env)[2],p.k,logorder = 1)
    # deleteat!(cons,length(cons))
    # deleteat!(cash,length(cash))
    pts = convert(Point,cash,vfun(id,it,cons,cash,vf,p))
    vcat(pts, vf.env.v[2:end])  # connect at second point
    # vf.env.v[2:end]  # connect at second point
end

function plot_s(s::Simulation)

    # inc, cons, w
    py = plot(s.inc',leg = false, title = "income",yticks = [0,5,10,15],ylims = (-1,19))
    if s.p.ν > 0
        pc = plot(s.cons[:,1:(s.p.nT-1)]',leg = false, title = "consumption",ylims = (0,7))
        pw1 = plot(s.w1[:,1:(s.p.nT-1)]',leg = false, title = "w1",ylims = (0,maximum(s.p.a_high) + 0.1*maximum(s.p.a_high)))

    else
        pc = plot(s.cons',leg = false, title = "consumption",ylims = (0,7))
        pw1 = plot(s.w1',leg = false, title = "w1",ylims = (0,maximum(s.p.a_high) + 0.1*maximum(s.p.a_high)))
    end
    i_retires = findall(s.ret_age .> 0)
    scatter!(pc, s.ret_age[i_retires], [s.cons[i,s.ret_age[i]] for i in i_retires], m = (:rect, 2, 0.6, :white))
    pw0 = plot(s.w0',leg = false, title = "w0",ylims = (0,maximum(s.p.a_high) + 0.1*maximum(s.p.a_high)))
    scatter!(pw0, s.ret_age[i_retires], [s.w0[i,s.ret_age[i]] for i in i_retires], m = (:rect, 2, 0.6, :white))

    ppr = plot(s.prob_work',leg = false, title = "p(work)")
    scatter!(ppr, s.ret_age[i_retires], [s.prob_work[i,s.ret_age[i]] for i in i_retires], m = (:rect, 2, 0.6, :white))

    pbeq = if s.p.ν > 0
        histogram(s.w0[:,s.p.nT], title = "Bequests", bins = range(0,stop = 30.0, length = 50), normalize = false, legend = false)
    else
        histogram(s.w1[:,s.p.nT], title = "Bequests", bins = range(0,stop = 30.0, length = 50), normalize =false, legend = false)
    end
    pretage = histogram(s.ret_age, title = "Ret Age", bins = range(0,stop = s.p.nT, length = s.p.nT), normalize = false, legend = false)
    mutil = map(x -> mean(s.util[:,x]), 1:s.p.nT)
    pmutil = plot(1:s.p.nT, mutil, title ="mean utility",ylims = (0,1),yticks = range(0,1,length = 6),legend = false)

    pempty = plot(legend=false,grid=false,foreground_color_subplot=:white)
    asize = 10
    # annotate!(pempty, [(0.3,1,Plots.text("gamma = $(s.p.gamma)",     :left, asize)),
    #                    (0.3,0.9,Plots.text("R = $(s.p.R)",     :left, asize)),
    #                    (0.3,0.8,Plots.text("beta = $(round(s.p.beta,digits=2))", :left, asize)),
    #                    (0.3,0.7,Plots.text("alpha = $(s.p.alpha)",   :left, asize)),
    #                    (0.3,0.6,Plots.text("sigma = $(s.p.sigma)",   :left, asize)),
    #                    (0.3,0.5,Plots.text("lambda = $(s.p.lambda)", :left, asize))#,
    #                    #(0.3,0.4,Plots.text("rho = $(s.p.ρ)",         :left, asize))
    #                    ])
    plot(py,pc,pw0,pretage,
        ppr,pw1,pbeq , pmutil,
          layout = (2,4), size = (800,400))
end

function plot_s(s::BSimulation)

    # inc, cons, w
    py = plot(s.inc',leg = false, title = "income")
    pc = plot(s.cons',leg = false, title = "consumption")
    i_retires = findall(s.file_age .> 0)
    scatter!(pc, s.file_age[i_retires], [s.cons[i,s.file_age[i]] for i in i_retires], m = (:rect, 2, 0.6, :white))
    pw0 = plot(s.w0',leg = false, title = "w0")
    scatter!(pw0, s.file_age[i_retires], [s.w0[i,s.file_age[i]] for i in i_retires], m = (:rect, 2, 0.6, :white))

    ppr = plot(s.prob_nofile',leg = false, title = "p(nofile)",ylims = (-0.05,1.05))
    scatter!(ppr, s.file_age[i_retires], [s.prob_nofile[i,s.file_age[i]] for i in i_retires], m = (:rect, 2, 0.6, :white))

    pw1 = plot(s.w1',leg = false, title = "w1")
    pempty = plot(legend=false,grid=false,foreground_color_subplot=:white)
    asize = 10
    annotate!(pempty, [(0.1,1,Plots.text("gamma = $(s.p.gamma)",     :left, asize)),
                       (0.1,0.9,Plots.text("R = $(round(s.p.R,digits=2))",     :left, asize)),
                       (0.1,0.8,Plots.text("beta = $(round(s.p.beta,digits=2))", :left, asize)),
                       (0.1,0.7,Plots.text("alpha = $(s.p.alpha)",   :left, asize)),
                       (0.1,0.6,Plots.text("sigma = $(s.p.sigma)",   :left, asize)),
                       (0.1,0.5,Plots.text("lambda = $(s.p.lambda)", :left, asize)),
                       (0.1,0.4,Plots.text("rho = $(s.p.ρ)",         :left, asize))])
    plot(py,pc,pw0,ppr,pw1,
          pempty ,
          layout = (2,3))
end

@recipe function f(m::FModel,p::Param;id=1)
    grid --> true
    xticks := true
    legend --> true
    cg = cgrad(:plasma)
    # c1 = colorant"red"
    # c2 = colorant"blue"
    # alow,ahi = extrema(m.avec)
    # aspect = (ahi-alow)/(ahi - 0.0)

    nT = size(m.v)[2]
    # cols = range(c1,stop=c2,length=nT)
    cols = cg[range(0.0,stop=1.0,length=nT)]
    xrange = zeros(2)
    xa = extrema(m.avec)
    xrange[1] = xa[1] .- diff(vcat(xa...))[1] .* 0.01
    xrange[2] = xa[2]

    layout := grid(1,2)
    for i in 1:nT
        vt = v_analytic(m,p,id,i)
        lab = ((i==1)|(i==nT)) ? "$i" : ""
        @series begin
            seriestype --> :path
            linewidth --> 1
            legend --> :bottomright
            seriescolor --> cols[i]
            label := lab
            subplot := 1  # value function
            yguide := "value"
            xguide := "Cash on Hand M"
            xlims := xrange
            ylims := (-15,15)
            getx(vt),gety(vt)
        end
        @series begin
            seriestype --> :path
            linewidth --> 1
            legend --> :bottomright
            seriescolor --> cols[i]
            # subplot := 2  #
            label := lab
            # xlims := (alow,ahi)
            # ylims := (0,ahi)
            ylims := extrema(m.avec)
            xlims := extrema(m.avec)
            yguide := "consumption"
            xguide := "Cash on Hand M"
            # aspect_ratio := aspect
            getx(m.c[id,i].env),gety(m.c[id,i].env)
        end
    end
end

@recipe function f(m::GModel,p::Param;id=1,iy=nothing,it=nothing)
    grid --> true
    xticks := true
    legend --> false
    cg = cgrad(:plasma)
    # c1 = colorant"red"
    # c2 = colorant"blue"

    xrange = zeros(2)
    xa = extrema(m.avec)
    xrange[1] = xa[1] .- diff(vcat(xa...))[1] .* 0.01
    xrange[2] = xa[2]

    nT = p.nT
    cols = cg[range(0.0,stop=1.0,length=nT)]

    layout := grid(1,2)

    if isnothing(iy) & !isnothing(it)
        legend := true
        title --> ["value period $it" "consumption period $it"]
        for jy in 1:p.ny
            vt = v_analytic(m,p,id,jy,it)
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend := :bottomright
                label := "y$jy"
                # seriescolor --> cols[i]
                subplot := 1  # value function
                yguide := "value"
                xguide := "Cash on Hand M"
                xlims := xrange
                # ylims --> (-15,15)
                getx(vt),gety(vt)
                # getx(m.v[id,iy,it].env),gety(m.v[id,iy,it].env)
            end
        end
        for jy in 1:p.ny
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> :bottomright
                # seriescolor --> cols[i]
                subplot := 2  #
                label := "y$jy"
                xlims := xa
                ylims := xa
                yguide := "consumption"
                xguide := "Cash on Hand M"
                # aspect_ratio := aspect
                getx(m.c[id,jy,it].env),gety(m.c[id,jy,it].env)
            end
        end
    elseif isnothing(it) & !isnothing(iy)
        for i in 1:nT
            vt = v_analytic(m,p,id,iy,i)
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> :bottomright
                seriescolor --> cols[i]
                label := ((i == 1) || (i == nT)) ? "$it" : ""
                # label := lab
                xlims := xrange
                # ylims --> (-15,15)
                subplot := 1  # value function
                yguide := "value"
                xguide := "Cash on Hand M"
                getx(vt),gety(vt)
            end
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> :bottomright
                seriescolor --> cols[i]
                subplot := 2  #
                label := ((i == 1) || (i == nT)) ? "$it" : ""
                # label := lab
                xlims := xa
                ylims := xa
                yguide := "consumption"
                xguide := "Cash on Hand M"
                # aspect_ratio := aspect
                getx(m.c[id,iy,i].env),gety(m.c[id,iy,i].env)
            end
        end
    else
        println("you need to either give it or iy. not both. not none.")
    end

end

@recipe function f(m::BModel,p::Param;id=1,iy=nothing,it=nothing)
    grid --> true
    xticks := true
    # legend --> false
    cg = cgrad(:plasma)
    cols = cg[range(0.0,stop=1.0,length=p.nT)]

    # cg = cgrad(:inferno)
    # c1 = colorant"red"
    # c2 = colorant"blue"
    xrange = zeros(2)
    xa = extrema(m.avec)
    xrange[1] = xa[1] .- diff(vcat(xa...))[1] .* 0.01
    xrange[2] = xa[2]
    nT = p.nT
    # cols = range(c1,stop=c2,length=nT)

    layout := grid(1,3)

    rates = map(z->mean(z),m.rgrid[1,:,1:(nT-1)])

    if isnothing(iy) & !isnothing(it)
        legend := true
        title --> ["value period $it" "consumption period $it"]
        for jy in 1:p.ny
            vt = v_analytic(m,p,id,jy,it)
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> :bottomright
                label := "y$jy"
                # seriescolor --> cols[i]
                subplot := 1  # value function
                yguide := "value"
                xguide := "Cash on Hand M"
                xlims := xrange
                # ylims --> (-15,15)
                getx(vt),gety(vt)
                # getx(m.v[id,iy,it].env),gety(m.v[id,iy,it].env)
            end
        end
        for jy in 1:p.ny
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> :bottomright
                # seriescolor --> cols[i]
                subplot := 2  #
                label := "y$jy"
                xlims := xa
                ylims := xa
                yguide := "consumption"
                xguide := "Cash on Hand M"
                # aspect_ratio := aspect
                getx(m.c[id,jy,it].env),gety(m.c[id,jy,it].env)
            end
        end
        for jy in 1:p.ny
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> :topleft
                label := "y$jy"

                # seriescolor --> cols[i]
                subplot := 3  #

                yguide := "interest rate"
                xguide := "age"
                ylims := (0,maximum(rates))

                # aspect_ratio := aspect
                # println(rates[jy,:])
                collect(1:(p.nT-1)),rates[jy,:]
            end
        end
    elseif isnothing(it) & !isnothing(iy)
        for i in 1:nT
            vt = v_analytic(m,p,id,iy,i)
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> false
                seriescolor --> cols[i]
                # label := lab
                label := ((i == 1) || (i == nT)) ? "$it" : ""

                xlims := xrange
                # ylims --> (-15,15)
                subplot := 1  # value function
                yguide := "value"
                xguide := "Cash on Hand M"
                getx(vt),gety(vt)
            end
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> false

                seriescolor --> cols[i]
                subplot := 2  #
                # label := lab
                label := ((i == 1) || (i == nT)) ? "$it" : ""

                xlims := xa
                ylims := xa
                yguide := "consumption"
                xguide := "Cash on Hand M"
                # aspect_ratio := aspect
                getx(m.c[id,iy,i].env),gety(m.c[id,iy,i].env)
            end
        end

        for jy in 1:p.ny
            @series begin
                seriestype --> :path
                linewidth --> 1
                legend --> true
                label := "y$jy"
                # seriescolor --> cols[i]
                subplot := 3  #
                ylims := (0,maximum(rates))

                yguide := "interest rate"
                xguide := "age"
                # aspect_ratio := aspect
                # println(rates[jy,:])
                collect(1:(p.nT-1)),rates[jy,:]
            end
        end
    else
        println("you need to either give it or iy. not both. not none.")
    end
end

struct tester
    m :: Matrix
end

@recipe function f(tt::tester)
    n,m = size(tt.m)
    for i in 1:n
        @series begin
            seriestype --> :path
            series_annotations := ["$i" for i in 1:m]
            (1:m,tt.m[i,:])
        end
    end
end

@recipe function f(l::Vector{Point{T}};numerate=false,marker=false) where T
    # defaults
    grid --> true
    xticks := true
    @series begin
        seriestype --> :path
        linecolor --> :black
        linewidth --> 1
        if marker
            markershape --> :circle
            markerstrokecolor --> :black
            markercolor --> :white
            markersize --> 1
        end
        if numerate
            series_annotations := ["$i" for i in 1:length(l.v)]
        end
        (getx(l),gety(l))
    end
end

@recipe function f(l::MLine;numerate=false,marker=false)
    # defaults
    grid --> true
    xticks := true
    legend := false
    @series begin
        seriestype --> :path
        linecolor --> :black
        linewidth --> 1
        if marker
            markershape --> :circle
            markerstrokecolor --> :black
            markercolor --> :white
            markersize --> 1
        end
        if numerate
            series_annotations := ["$i" for i in 1:length(l.v)]
        end
        (getx(l),gety(l))
    end
end

@recipe function f(L::Vector{MLine{T}};numerate=false,mrk=true) where T
    for l in L
        @series begin
            # subplot := 1
            seriestype := :path
            linewidth := 1
            if mrk
                # println("illegal")
                markershape --> :circle
                markerstrokecolor --> :black
                markercolor --> :white
                markersize := 3
            end
            if numerate
                series_annotations := ["$i" for i in sortperm(getx(l))]
            end
            (getx(l),gety(l))
        end
    end
end
#
@recipe function f(x::Envelope;numerate=false,mrk=true)

    # defaults
    grid --> true
    xticks := true

    any_isec = length(x.isects) > 0

    # println("marker = $mrk, title = $(get!(plotattributes,:title,""))")
    # println("marker = $mrk, numerate = $numerate, title = $(get!(plotattributes,:title,""))")

    # actual envelope
    @series begin
        # subplot := 1
        seriestype := :path
        linecolor --> :red
        linewidth --> 2
        if mrk
            markershape := :circle
            markercolor := :white
            # markeralpha := 0.5
            markerstrokecolor := :black
            markersize := 3
        end
        if numerate
            series_annotations := ["$i" for i in 1:length(getx(x.env))]
        end
        (getx(x.env),gety(x.env))
    end
    if any_isec && mrk
        @series begin
            seriestype := :scatter
            markershape := :star6
            markerstrokecolor := :black
            markercolor := :yellow
            markersize := 5
            (getx(x.isects),gety(x.isects))
        end
    end
end
#
@recipe function f(x::Envelope{T},L::Vector{MLine{T}};numerate=false,removed=false,mrk=true) where T

    # defaults
    grid --> true
    xticks := true
    legend --> false

    any_isec = length(x.isects) > 0

    # println("marker = $mrk, title = $(get!(plotattributes,:title,""))")
    # println("marker = $mrk, numerate = $numerate, title = $(get!(plotattributes,:title,""))")

    # lien arrya
    for l in L
        @series begin
            # subplot := 1
            seriestype := :path
            linewidth := 1
            if mrk
                # println("illegal")
                markershape --> :circle
                markerstrokecolor --> :black
                markercolor --> :white
                markersize := 3
            end
            if numerate
                series_annotations := ["$i" for i in sortperm(getx(l))]
            end
            (getx(l),gety(l))
        end
    end

    # actual envelope
    @series begin
        # subplot := 1
        seriestype := :path
        linecolor --> :red
        linewidth --> 2
        if mrk
            markershape := :circle
            markercolor := :white
            # markeralpha := 0.5
            markerstrokecolor := :black
            markersize := 3
        end
        if numerate
            series_annotations := ["$i" for i in 1:length(getx(x.env))]
        end
        (getx(x.env),gety(x.env))
    end
    if any_isec && mrk
        @series begin
            seriestype := :scatter
            markershape := :star6
            markerstrokecolor := :black
            markercolor := :yellow
            markersize := 5
            (getx(x.isects),gety(x.isects))
        end
    end
end

# plot env and initial mline next to each other, showing new points and removed ones
@recipe function f(x::Envelope{T},L::MLine{T};numerate=false) where T

    # defaults
    grid --> true
    xticks := true
    legend --> false
    layout := (1,2)
    title --> ["Backwards-Bending" "Envelope"]

    any_isec = length(x.isects) > 0
    any_rmv = length(x.removed) > 0

    # println("marker = $mrk, title = $(get!(plotattributes,:title,""))")
    # println("marker = $mrk, numerate = $numerate, title = $(get!(plotattributes,:title,""))")

    # dirty line
    @series begin
        subplot := 1
        seriestype := :path
        linewidth := 1
        markershape --> :circle
        markerstrokecolor --> :black
        markercolor --> :white
        markersize := 3
        if numerate
            series_annotations := ["$i" for i in 1:length(getx(L))]
        end
        (getx(L),gety(L))
    end
    if any_rmv
        @series begin
            subplot := 1
            seriestype = :scatter
            markershape := :rect
            markersize := 3
            markerstrokecolor := :black
            markercolor := :white
            markeralpha := 0.5
            (getx(L)[x.removed],gety(L)[x.removed])
        end
    end

    # actual envelope
    @series begin
        subplot := 2
        seriestype := :path
        linecolor --> :red
        linewidth --> 1
        markershape --> :circle
        markerstrokecolor --> :black
        markercolor --> :white
        markersize := 3
        (getx(x.env),gety(x.env))
    end
    if any_isec
        @series begin
            subplot := 2
            seriestype := :scatter
            markershape := :star6
            markerstrokecolor := :black
            markercolor := :yellow
            markersize := 5
            (getx(x.isects),gety(x.isects))
        end
    end

end

#
# # @recipe function f(x::Envelope; removed=false,num=false,marker=false)
# @recipe function f(x::Envelope;numerate=false,removed=false,mrk=true)
#
#     # defaults
#     grid --> true
#     xticks := true
#
#     if !x.env_clean
#         legend --> :bottomright
#     else
#         legend --> false
#     end
#     any_isec = length(x.isects) > 0
#
#     # println("marker = $mrk, title = $(get!(plotattributes,:title,""))")
#     # println("marker = $mrk, numerate = $numerate, title = $(get!(plotattributes,:title,""))")
#
#     # if line array exists, plot
#     if length(x.L) > 0
#         for l in x.L
#             @series begin
#                 # subplot := 1
#                 seriestype := :path
#                 linewidth := 1
#                 if mrk
#                     # println("illegal")
#                     markershape --> :circle
#                     markerstrokecolor --> :black
#                     markercolor --> :white
#                     markersize := 3
#                 end
#                 if numerate
#                     series_annotations := ["$i" for i in sortperm(getx(l))]
#                 end
#                 (getx(l),gety(l))
#             end
#         end
#     end
#     # plot envelope, if exists
#     if x.env_clean
#         if removed
#             for l in 1:length(x.L)
#                 ir = x.removed[l]
#                 if length(ir) > 0
#                     @series begin
#                         seriestype = :scatter
#                         markershape := :rect
#                         markersize := 3
#                         markerstrokecolor := :black
#                         markercolor := :white
#                         markeralpha := 0.5
#                         (getx(x.L[l])[ir],gety(x.L[l])[ir])
#                     end
#                 end
#             end
#         end
#         @series begin
#             # subplot := 1
#             seriestype := :path
#             linecolor --> :red
#             linewidth --> 2
#             if mrk
#                 markershape := :circle
#                 markercolor := :white
#                 # markeralpha := 0.5
#                 markerstrokecolor := :black
#                 markersize := 3
#             end
#             if numerate
#                 series_annotations := ["$i" for i in 1:length(getx(x.env))]
#             end
#             (getx(x.env),gety(x.env))
#         end
#         if any_isec && mrk
#             @series begin
#                 seriestype := :scatter
#                 markershape := :star6
#                 markerstrokecolor := :black
#                 markercolor := :yellow
#                 markersize := 5
#                 (getx(x.isects),gety(x.isects))
#             end
#         end
#     end
# end

function f3a()

    fs = Function[x->ones(length(x));x->0.5x;x->x .- 2;x->2x .- 8]
    xs = [
         collect(range(-1 , stop = 0.9 ,length= 6))   ,
         collect(range(1  , stop = 7   ,length= 19))  ,
         collect(range(2  , stop = 7   ,length= 15))  ,
         collect(range(4  , stop = 8   ,length= 25))]
    ls = [MLine(i[1],i[2](i[1])) for i in zip(xs,fs)]
    return ls
end
function f3b()

    fs = Function[x->ones(length(x));x->0.5x;x->x .- 2;x->2x .- 8]
    xs = [
         collect(range(-1 , stop = 3.1 , length = 6))   ,
         collect(range(1  , stop = 7   , length = 19))  ,
         collect(range(2  , stop = 7   , length = 15))  ,
         collect(range(4  , stop = 8   , length = 25))]
    ls = [MLine(i[1],i[2](i[1])) for i in zip(xs,fs)]
    # create an envelope
    return ls
end
function f3c()

    fs = Function[x->ones(length(x));x -> 0.5x;x->x .- 2;x -> 2x .- 8.5]
    xs = [
         collect(range(-1 , stop = 3.1 , length= 6))   ,
         collect(range(1  , stop = 7   , length= 4))  ,
         collect(range(2  , stop = 7   , length= 15))  ,
         collect(range(4  , stop = 8   , length= 25))]
    ls = [MLine(i[1],i[2](i[1])) for i in zip(xs,fs)]
    # create an envelope
    return ls
end

function tplot1()
    n = 15
    x1 = collect(range(0 , stop = 10,length = n))
    x2 = collect(range(-1, stop = 9 ,length = n))
    L1 = MLine(x1,x1)
    L2 = MLine(x2,ones(n)*5)
    e = [L1,L2]
    x1,x2,e #,plot(e)
end

function tplot2()
    n = 15
    x1 = collect(range(0  , stop=10 ,length= n))
    x2 = collect(range(-1 , stop=9  ,length= n))
    L1 = MLine(x1,x1)
    L2 = MLine(x2,ones(n)*5)
    LL = [L1,L2]

    # a = splitLine(LL)
    e = upper_env(LL)
    p2 = plot(e,LL)
    savefig(p2,joinpath(@__DIR__,"..","images","tplot2.png"))
    p2
end
function tplot3a()

    en = f3a()

    p1 = plot(en,title = "non-overlapping grids")

    e = upper_env(en)
    p2 = plot(e,en,title = "Envelope")
    plot(p1,p2)

end

function tplot3b()

    en = f3b()

    p1 = plot(en,title = "overlapping grids")

    e = upper_env(en)
    p2 = plot(en)

    plot(p1,p2)

end
function tplot3c()

    en = f3c()

    p1 = plot(en)

    e = upper_env(en)
    removed!(en)
    p2 = plot(en,removed=true)

    plot(p1,p2)
    gui()
    return en

end

function tplot4()

    fs = Function[x->ones(length(x));x->0.5x;x->x .- 2;x->2x .- 8]
    xs = [
         collect(range(0.1 , stop =  1.5 ,length = 5))   ,
         collect(range(1   , stop =  7   ,length = 19))  ,
         collect(range(2   , stop =  7   ,length = 15))  ,
         collect(range(4   , stop =  8   ,length = 25))]
    ls = [MLine(i[1],i[2](i[1])) for i in zip(xs,fs)]
    e = DCEGM.upper_env(ls)

    p1 = plot(ls)
    p2 = plot(e,ls,title="envelope")
    p = plot(p1,p2)
    savefig(p,joinpath(@__DIR__,"..","images","tplot4.png"))
    p
end

function tplot5()

    x1 = collect(-0.9:0.3:2.7)
    L1 = MLine(x1, x1)
    x2 = collect(0.0:0.1:1)
    L2 = MLine(x2, 2 .* x2)
    x3 = collect(1.0:0.45:2.9)
    L3 = MLine(x3, (0.1 .* x3) .+ 1.9 )
    ls = [L1,L2,L3]
    e = DCEGM.upper_env(ls)
    p1 = plot(ls)
    p2 = plot(e,ls,title="envelope")

    p = plot(p1,p2)
    savefig(p,joinpath(@__DIR__,"..","images","tplot5.png"))
    p

end

function splitf()
    x = [1,2,3,1.5,2.1,2.9]
    y = [1,1.5,1.7,1.2,1.8,2.1]
    L = MLine(x,y)
    p1 = plot(L,title="original",numerate=true)
    e = splitLine(L)
    p2 = plot(e,title="split MLine",numerate=true)
    p = plot(p1,p2)
    savefig(p,joinpath(@__DIR__,"..","images","split.png"))
    p
end

function splitf2()
    x = [1,2,3,2.9,2.5,1.9,1.8,1.5,2.1,2.9]
    y = [1,1.5,1.7,1.6,1.55,1.4,1.3,1.2,1.8,2.1]
    L = MLine(x,y)
    p1 = plot(L,title="original",numerate=true)
    e = splitLine(L)
    p2 = plot(e,title="split MLine",numerate=true)
    p = plot(p1,p2)
    savefig(p,joinpath(@__DIR__,"..","images","split2.png"))
    p
end

function splitf3()
    x = [1,2,3,2.9,2.5,1.9,1.8,1.5,2.1,2.9]
    y = [1,1.5,1.7,1.6,1.55,1.4,1.3,1.2,1.8,2.1]
    L = MLine(x,y)
    e = DCEGM.secondary_envelope(L)
    p2 = plot(e,L,numerate=true)
    savefig(p2,joinpath(@__DIR__,"..","images","split3.png"))
    p2
end

function splitf3a()
    x = [1,2,3,2.9,2.5,1.9,1.8,1.5,2.1,2.9]
    y = [1,1.5,1.7,1.6,1.55,1.4,1.3,1.2,1.8,2.1]
    L = MLine(x,y)
    e = secondary_envelope(L)
    plot(e,L)

end

function test_upper_env_dec()
    n = 15
    x1 = collect(range(0,stop = 10,length = n))
    x2 = collect(range(-1,stop = 9,length = n))
    x = vcat(x1,x2)
    y = vcat(x1[end:-1:1],ones(n)*5)
    L = MLine(x,y)
    e = splitLine(L)
    p1 = plot(L,numerate=true,title = "initial line")
    p2 = plot(e,marker=true,title = "split line")
    e = upper_env(e)
    p3 = plot(e,marker=true,title = "Envelope w new pts")
    p4 = plot(p1,p2,p3,layout = (1,3),size=(1200,500))
    savefig(p4,joinpath(@__DIR__,"..","images","descending_upper.png"))
    L,e,x1,x2,p4
end

function split_test()
    f1(x) = ones(length(x))
    f2(x) = 0.5x
    f3(x) = x .-2
    f4(x) = 2x .-8
    x1 = collect(range(0.9,stop = 2.1,length = 14))
    x2 = collect(range(1,stop = 7,length = 19))
    x3 = collect(range(1,stop = 7,length = 15))
    x4 = collect(range(1,stop = 8,length = 25))
    X = [x1...,x2...,x3...,x4...]
    L = MLine([x1...,x2...,x3...,x4...],vcat(f1(x1),f2(x2),f3(x3),f4(x4)))
    en = splitLine(L)
    en,L
end

function demo(;n = 10,k = 10)
    a = Array{MLine{Float64}}(undef,10)
    for i in 1:2:n
        a[i] = MLine(sort([0.0; 10*rand(k); 10]), [0.0; rand(k+1).*collect(range(0.5,stop=10,length=k+1))])
        a[i+1] = MLine([0.0; 10],[10.0-i ; -(10-i)*1.5/i+10-i] )
    end
    p1 = plot(a,legend=false,title = "$n lines")
    e = upper_env(a)
    p2 = plot(e,a,title = "envelope")
    p =plot(p1,p2)
    savefig(p,joinpath(@__DIR__,"..","images","demo.png"))
    p
end

function demo2(;n = 10,k = 10)
    a = Array{MLine{Float64}}(undef,10)
    slope = 1.5
    for i in 1:2:n
        a[i] = MLine(sort([0.0; 10*rand(k); 10]), [0.0; rand(k+1).*collect(range(0.5,stop=10,length=k+1))])
        a[i+1] = MLine([0.0; 10],[10.0-i ; -(10-i)*slope/i+10-i] )
    end
    # add a particularly tricky line
    tricky = [a[2].v[1]]
    push!(tricky,DCEGM.Point(a[2].v[2].x,a[2].v[2].y + slope))
    push!(a, MLine(tricky))

    p1 = plot(a,legend=false,title = "$n lines")
    e = upper_env(a)
    p2 = plot(e,a,title = "first intersection on grid")
    p = plot(p1,p2)
    savefig(p,joinpath(@__DIR__,"..","images","demo2.png"))
    p,e
end

function allplots()
    tplot1()
    tplot2()
    # savefig(joinpath(p,"tplot3a.png"))
    # tplot3b()
    # savefig(joinpath(p,"tplot3b.png"))
    # tplot3c()
    # savefig(joinpath(p,"tplot3c.png"))
    tplot4()
    tplot5()
    splitf()
    splitf2()
    splitf3()
    demo()
    demo2()
    test_upper_env_dec()
end


# Download latest binary shared library of CoolProp project
import JSON
const destpathbase = abspath(joinpath(@__FILE__,"..","..","lib"));
println(destpathbase);
const srctablepath = abspath(joinpath(@__FILE__,"..","..","src","Tables"));
println(srctablepath);
m1=0
m2=0
m3=0
function canconvert(ty::Type,matrix::Array,database,sizearray)
  global m1,m2,m3
  j=1
  si=Int(0)
  r::Tuple
  r=(NaN,)
  for(j in 1:size(matrix)[1])
    try
      r=ty(tuple(matrix[j,:]...))
    catch
      return matrix[j,:],false
    end
    if (sizearray!=nothing)
      m1>0 && (r2=rpad(r[2],m1,' '))
      m2>0 && (r3=rpad(r[3],m2,' '))
      m3>0 && (r4=rpad(r[4],m3,' '))
      m1>0 && (r=(r[1],r2,r3,r4,r[5],r[6]))
      si+=writetup(database,r)
    else
      b=[(r[i]==matrix[j,i] || isnan(r[i])) for i in 1:length(r)]
      if(sum(b)!=length(r))
        println(matrix[j,:])
        println(r)
        return false,matrix[j,:]
      end
    end
  end
  (sizearray!=nothing) && push!(sizearray,(si,j))
  return true
end
function writetup(database,t::Tuple)
  i=0
  i+=write(database,t...)
  return i;
end
function buildbinarydatabase(database,si)
  ret=true
  global m1,m2,m3
  propnames=["Criticals",
              "LiquidsDensities",
              "FormationEnergy",
              "Compounds",
              "Cp",
              "LiquidsCp",
              "VaporizHeat",
              "LiquidThermal",
              "VaporPressure",
              "LiquidViscos",
              "VaporThermal",
              "VaporViscos"]

  tablenames=["perryCriticals_Table2_141.table",
              "perryDensities_Table2_32.table",
              "perryEnergiesOfFormation_Table2_179.table",
              "perryFormulaCompounds.table",
              "perryHeatCapIdealGas_Tables156_155.table",
              "perryHeatCapLiquids_Table2_153.table",
              "perryHeatsofVaporization_Table2_150.table",
              "perryLiqidThermalConductivity_Table2_315.table",
              "perryLiquidsVaporPressure_Table2_8.table",
              "perryLiquidViscosity_Table2_313.table",
              "perryVaporThermalConductivity_Table2_314.table",
              "perryVaporViscosity_Table2_312.table"]

  rowtypes=[Tuple{[UInt16;[Float64 for i=1:6]]...},
            Tuple{[UInt16;[Float64 for i=1:10];UInt8]...},
            Tuple{[UInt16;[Float64 for i=1:5]]...},
            Tuple{[UInt16,ASCIIString,ASCIIString,ASCIIString,Float64,Float64]...},
            Tuple{[UInt16;[Float64 for i=1:10];UInt8]...},
            Tuple{[UInt16;[Float64 for i=1:10];UInt8]...},
            Tuple{[UInt16;[Float64 for i=1:10]]...},
            Tuple{[UInt16;[Float64 for i=1:10]]...},
            Tuple{[UInt16;[Float64 for i=1:10]]...},
            Tuple{[UInt16;[Float64 for i=1:10]]...},
            Tuple{[UInt16;[Float64 for i=1:9]]...},
            Tuple{[UInt16;[Float64 for i=1:9]]...}]

  for i=1:length(tablenames)
    tablename=tablenames[i]
    file=open(joinpath(srctablepath, tablename))
    matrix=readdlm(file,';',header=false)
    close(file)
    if tablename=="perryFormulaCompounds.table"
      m1=(maximum(Int[length(i) for i in matrix[:,2]]))
      m2=(maximum(Int[length(i) for i in matrix[:,3]]))
      m3=(maximum(Int[length(i) for i in matrix[:,4]]))
    else
      m1=0;m2=0;m3=0;
    end
    ty=rowtypes[i]
    if canconvert(ty,matrix,database,si)
      database==nothing && println("$tablename convertion simulated.")
    else
      println("$tablename convertion failed.")
      ret=false
    end
  end
  if si!=nothing
    databaseaddress=open(joinpath(destpathbase,"perryanalytic.addresses"),"w")
    println("Writing addresses .....")
    startaddress=0;
    for i=1:length(tablenames)
      s=si[i]
      println(databaseaddress,propnames[i],';',startaddress,';',s[2],';',round(Int,s[1]/s[2]),';',rowtypes[i])
      startaddress+=s[1]
    end
    close(databaseaddress)
  end
  return ret
end

try
  mkdir(destpathbase)
  const OS_ARCH_CoolProp = (WORD_SIZE == 64) ? "64bit" : "32bit__cdecl";
  const OS_ARCH_FreeSteam = (WORD_SIZE == 64) ? "win64" : "win32";
  const latestVersion_CoolProp = JSON.parse(readall(download("http://sourceforge.net/projects/coolprop/best_release.json")))["release"]["filename"][11:15];
  const latestVersion_FreeSteam = "2.1"
  println("CoolProp latestVersion = $latestVersion_CoolProp, by default I am going to install it...")

  @windows_only begin
      urlbase = "http://netix.dl.sourceforge.net/project/coolprop/CoolProp/$latestVersion_CoolProp/shared_library/Windows/$OS_ARCH_CoolProp/"
      download(joinpath(urlbase,"CoolProp.dll"),joinpath(destpathbase,"CoolProp.dll"))
      download(joinpath(urlbase,"CoolProp.lib"),joinpath(destpathbase,"CoolProp.lib"))
      download(joinpath(urlbase,"exports.txt"),joinpath(destpathbase,"exports.txt"))
      println("downloaded => lib/CoolProp.dll")
      # FreeSteam
      println("I'm Getting FreeSteam Binaries...")
      urlbase = "http://cdn.rawgit.com/DANA-Laboratory/FreeSteamBinary/master/$latestVersion_FreeSteam/$OS_ARCH_FreeSteam/"
      download(joinpath(urlbase,"freesteam.dll"),joinpath(destpathbase,"freesteam.dll"))
      println("downloaded => lib/freesteam.dll")
  end
  @linux_only begin
      # CoolProp
      urlbase = "http://netix.dl.sourceforge.net/project/coolprop/CoolProp/$latestVersion_CoolProp/shared_library/Linux/64bit/libCoolProp.so.$latestVersion_CoolProp"
      download(urlbase,joinpath(destpathbase,"libCoolProp.so"))
      println("downloaded => lib/libCoolProp.so")
      # FreeSteam
      println("I'm Getting FreeSteam Binaries...")
      urlbase = "http://cdn.rawgit.com/DANA-Laboratory/FreeSteamBinary/master/"
      download(joinpath(urlbase,"$latestVersion_FreeSteam","linux","libfreesteam.so.1.0"),joinpath(destpathbase,"libfreesteam.so"))
      println("downloaded => lib/libfreesteam.so")
      download(joinpath(urlbase,"libgsl.so.0"),joinpath(destpathbase,"libgsl.so"))
      println("downloaded => lib/libgsl.so")
      download(joinpath(urlbase,"libgslcblas.so.0"),joinpath(destpathbase,"libgslcblas.so"))
      println("downloaded => lib/libgslcblas.so")
  end
catch err
  println("$err\n => If lib folder exists, and you want to refresh existing files, remove lib folder and retry")
end
if buildbinarydatabase(nothing, nothing) # Simulate convertion
  database=open(joinpath(destpathbase,"perryanalytic.binary"),"w")
  si=Vector{Tuple{Int,Int}}()
  println("********** Simulation done, I'm writing perryanalytic.binary **********")
  buildbinarydatabase(database,si)
  close(database)
end


export GbmModel, GbmRegressor, GbmClassifier, GbmRanker

@with_kw mutable struct GbmModel <: BaseEstimator
    pyo::PyObject = PyNULL()
    name::String = "lightgbm"
    max_bin::Int = 16
    objective::String = "binary:logistic"
    max_depth::Int = 5
    min_child_weight::Float32 = 1
    gamma::Float32 = 0
    num_leaves::Int = 31
    min_child_samples::Int = 1000
    max_position::Int = 200
    label_gain::String = join(0:30, ',')
    n_estimators::Int = 50
    reg_lambda::Float32 = 0.1
    subsample::Float32 = 1
    colsample_bytree::Float32 = 1
    colsample_bylevel::Float32 = 1
    bagging_temperature::Float32 = 1
end

is_classifier(m::GbmModel) = occursin(r"binary|multi", m.objective)
is_ranker(m::GbmModel) = occursin("rank", m.objective)

function gridparams(m::GbmModel)
    @unpack objective, name = m
    rank = is_ranker(m)
    grid = [
        "n_estimators" => [20, 50, 100],
        "reg_lambda" => [0.1, 1, 10],
        [
            "name" => "lightgbm",
            "objective" => rank ? ["rank:ndcg"] : [objective],
            "num_leaves" => [15, 31],
            "max_position" => rank ? [100, 200, 500, 1000] : [1],
            "label_gain" => rank ? [join(0:30, ','), ""] : [""],
            "min_child_samples" => [1000, 100, 5000],
            "subsample" => [1, 0.8],
            "colsample_bytree" => [1, 0.8]
        ],
        [
            "name" => ["xgboost"],
            "objective" => rank ? ["rank:pairwise", "rank:map", "rank:ndcg"] : [objective],
            "max_depth" => [5, 3, 7],
            "min_child_weight" => [1, 10, 100],
            "gamma" => [0, 0.01, 0.1, 1],
            "subsample" => [1, 0.8],
            "colsample_bytree" => [1, 0.8],
        ],
        [
            "name" => ["catboost"],
            "objective" => !rank ? [objective] :
                ["rank:queryrmse", "rank:yetirank",
                "rank:yetirankpairwise", "rank:pairlogit",
                "rank:pairlogitpairwise"],
            "max_depth" => [5, 3, 7],
            "min_child_samples" => !usegpu() ? [1000] : [1000, 100, 5000],
            "bagging_temperature" => [0, 1, 2, 5, 10],
            "colsample_bylevel" => usegpu() ? [1] : [1, 0.8],
        ]
    ]
    return griparams(grid)
end

function fit!(m::GbmModel, x, y, w = nothing; group = nothing, columns = nothing)
    columns = something(columns, string.(1:size(x, 1)))
    if is_ranker(m) && isnothing(group) && ndims(x) == 3
        group = repeat([size(x, 2)], size(x, 3))
    end
    x, y, w, group = pymat(x), vec(y), vec(w), vec(group)
    x = DataFrame(x, columns = columns)
    @unpack name, max_bin, objective, max_depth, min_child_weight, gamma = m
    @unpack num_leaves, min_child_samples, max_position, label_gain = m
    @unpack max_depth, n_estimators, reg_lambda, subsample, colsample_bytree = m
    @unpack colsample_bylevel, bagging_temperature = m
    num_class = occursin("reg", objective) ? 1 : length(unique(y))
    objective = occursin("binary", objective) && num_class > 2 ? "multi:softmax" : objective
    num_class = occursin("multi", objective) ? num_class : 1
    y = occursin("binary", objective) && name == "xgboost" ? (y .+ 1) ./ 2 : y
    objective = denormalize_objective(objective, name)
    min_child_samples = min(length(y) ÷ 10, min_child_samples)
    if name == "xgboost"
        @from xgboost imports XGBModel, XGBRanker
        XGBModel = !is_ranker(m) ? XGBModel : XGBRanker
        tree_method = usegpu() && is_ranker(m) ? "gpu_hist" : "auto"
        pyo = XGBModel(n_jobs = 20, max_bin = max_bin, tree_method = tree_method, single_precision_histogram = true, n_gpus = countgpus(), num_class = num_class)
        pyo.set_params(;@NT(objective, max_depth, min_child_weight, n_estimators, reg_lambda, subsample, gamma)...)
        tree_method == "auto" && pyo.set_params(colsample_bytree = colsample_bytree)
        if isnothing(group)
            @time pyo.fit(x, y, sample_weight = w, eval_set = [(x, y)], sample_weight_eval_set = bra(w), verbose = true)
        else
            if length(group) != length(w)
                @imports numpy as np
                w = map(sum, np.split(w, cumsum(group[1:end - 1])))
                @show length(w), length(group)
            end
            @time pyo.fit(x, y, group, sample_weight = w, eval_set = [(x, y)], sample_weight_eval_set = bra(w), eval_group = bra(group), verbose = true)
        end
    elseif name == "lightgbm"
        @from lightgbm imports LGBMModel
        @from unidecode imports unidecode
        device = usegpu() ? "gpu" : "cpu"
        pyo = LGBMModel(n_jobs = 20, max_bin = max_bin - 1, device = device, label_gain = label_gain, num_class = num_class)
        pyo.set_params(;@NT(objective, num_leaves, min_child_samples, max_position, n_estimators, reg_lambda, subsample, colsample_bytree)...)
        x.columns = columns′ = [replace(filter(!isspace, unidecode(c)), ":" => "-") for c in columns]
        @time pyo.fit(x, y, w, group = group, eval_set = [(x, y)], eval_sample_weight = bra(w), eval_group = bra(group), feature_name = columns′)
    elseif name == "catboost"
        @from catboost imports CatBoost
        @imports numpy as np
        min_child_samples = usegpu() ? min_child_samples : 1
        colsample_bylevel = usegpu() ? 1 : colsample_bylevel
        task_type = usegpu() ? "GPU" : "CPU"
        devices = string("0:", countgpus() - 1)
        loss_function = objective
        pyo = CatBoost(pairs((thread_count = 20, max_bin = max_bin - 1, task_type = task_type, devices = devices, learning_rate = 0.1)))
        pyo.set_params(;@NT(loss_function, max_depth, n_estimators, reg_lambda, colsample_bylevel, min_child_samples, bagging_temperature)...)
        group = isnothing(group) ? nothing : np.repeat(axes(group, 1), group)
        @time pyo.fit(x, y, sample_weight = w, group_id = group)
    end
    @pack! m = pyo
    BSON.bson("model.bson", model = m)
    visualize(m, columns)
    return m
end

function predict_gbm(m::GbmModel, x)
    @unpack name, pyo = m
    if name == "xgboost"
        pyo.predict(pymat(x), output_margin = true, validate_features = false)
    elseif name == "lightgbm"
        pyo.predict(pymat(x), raw_score = true)
    else
        pyo.predict(pymat(x))
    end
end

function predict_proba(m::GbmModel, x)
    ŷ = predict_gbm(m, x)
    if ndims(ŷ) == 1
        ŷ = sigmoid.(ŷ)
        hcat(1 .- ŷ,  ŷ)
    else
        softmax(ŷ, dims = 2)
    end
end

function predict(m::GbmModel, x)
    ŷ = predict_gbm(m, x)
    !is_classifier(m) && return ŷ
    if ndims(ŷ) == 1
        signone.(ŷ)
    else
        [I[2] - 1 for I in argmax(ŷ, dims = 2)]
    end
end

istree(m::GbmModel) = true

function visualize(m::GbmModel, columns)
    @unpack pyo = m
    mname, tname = pymodulename(pyo), pytypename(pyo)
    plot_tree = parseenv("PLOT_TREE", false)
    graphs, w = [], []
    if occursin("lightgbm", mname)
        bst = tname == "Booster" ? pyo : pyo.booster_
        w = bst.feature_importance()
        @from lightgbm imports create_tree_digraph
        plot_tree && for n in 1:min(5, bst.num_trees())
            push!(graphs, create_tree_digraph(bst, tree_index = n - 1))
        end
    elseif occursin("xgboost", mname)
        bst = tname == "Booster" ? pyo : pyo.get_booster()
        score = bst.get_score()
        columns = collect(keys(score))
        w = collect(values(score))
        @from xgboost imports to_graphviz
        plot_tree && for n in 1:min(5, bst.best_iteration)
            push!(graphs, to_graphviz(bst, num_trees = n - 1))
        end
    else
        return
    end
    write_feaimpt(w, columns)
    plot_tree && for (n, g) in enumerate(graphs)
        try g.render("gbm/$n.pdf", cleanup = "true") catch e; println(e) end
    end
    plot_tree && try bash("pdfcat -o gbm.pdf gbm/*.pdf && rm -rf gbm") catch end
end

function save(m::GbmModel, dst = "model.gbm")
    @unpack pyo = m
    mname, tname = pymodulename(pyo), pytypename(pyo)
    if occursin("lightgbm", mname)
        bst = tname == "Booster" ? pyo : pyo.booster_
    elseif occursin("xgboost", mname)
        bst = tname == "Booster" ? pyo : pyo.get_booster()
    elseif occursin("catboost", mname)
        bst = pyo
    else
        return
    end
    bst.save_model(dst)
end

function denormalize_objective(objective, name)
    mapping = Dict(v => k for (k, v) in objective_mapping[name])
    get(mapping, objective, objective)
end

function normalize_objective(objective)
    mapping = merge(values(objective_mapping)...)
    get(mapping, objective, objective)
end

const objective_mapping = Dict(
    "xgboost" => Dict(
        "reg:squarederror" => "reg:mse",
        "reg:squaredlogerror" => "reg:msle",
        "reg:gamma" => "reg:gamma",
        "reg:tweedie" => "reg:tweedie",
        "reg:logistic" => "reg:logistic",
        "binary:logistic" => "binary:logistic",
        "binary:logitraw" => "binary:logitraw",
        "binary:hinge" => "binary:hinge",
        "count:poisson" => "reg:poisson",
        "survival:cox" => "reg:cox",
        "multi:softmax" => "multi:softmax",
        "multi:softprob" => "multi:softprob",
        "rank:pairwise" => "rank:pairwise",
        "rank:ndcg" => "rank:ndcg",
        "rank:map" => "rank:map"
    ),
    "lightgbm" => Dict(
        "regression" => "reg:mse",
        "regression_l1" => "reg:mae",
        "huber" => "reg:huber",
        "fair" => "reg:fair",
        "poisson" => "reg:poisson",
        "quantile" => "reg:quantile",
        "mape" => "reg:mape",
        "gamma" => "reg:gamma",
        "tweedie" => "reg:tweedie",
        "binary" => "binary:logistic",
        "cross_entropy" => "binary:cross_entropy",
        "cross_entropy_lambda" => "binary:cross_entropy_lambda",
        "multiclass" => "multi:softmax",
        "multiclassova" => "multi:ova",
        "lambdarank" => "rank:ndcg"
    ),
    "catboost" => Dict(
        "MAE" => "reg:mae",
        "MAPE" => "reg:mape",
        "Poisson" => "reg:poisson",
        "Quantile" => "reg:quantile",
        "RMSE" => "reg:mse",
        "LogLinQuantile" => "reg:loglinquantile",
        "Huber" => "reg:huber",
        "Expectile" => "reg:expectile",
        "Logloss" => "binary:logistic",
        "CrossEntropy" => "binary:crossentropy",
        "MultiClass" => "multi:softmax",
        "MultiClassOneVsAll" => "multi:ova",
        "PairLogit" => "rank:pairlogit",
        "PairLogitPairwise" => "rank:pairlogitpairwise",
        "YetiRank" => "rank:yetirank",
        "YetiRankPairwise" => "rank:yetirankpairwise",
        "QueryCrossEntropy" => "rank:querycrossentropy",
        "QueryRMSE" => "rank:queryrmse",
        "QuerySoftMax" => "rank:querysoftmax"
    )
)

GbmRegressor(;ka...) = GbmModel(;objective = "reg:mse", ka...)

GbmClassifier(;ka...) = GbmModel(;objective = "binary:logistic", ka...)

GbmRanker(;ka...) = GbmModel(;objective = "rank:ndcg", ka...)