# Julia wrapper for header: /Users/dls/jl/openssl-1.0.1e/include/openssl/aes.h
# Automatically generated using Clang.jl wrap_c, version 0.0.0

@c Ptr{Uint8} AES_options () libcrypto
@c Int32 AES_set_encrypt_key (Ptr{Uint8}, Int32, Ptr{AES_KEY}) libcrypto
@c Int32 AES_set_decrypt_key (Ptr{Uint8}, Int32, Ptr{AES_KEY}) libcrypto
@c Int32 private_AES_set_encrypt_key (Ptr{Uint8}, Int32, Ptr{AES_KEY}) libcrypto
@c Int32 private_AES_set_decrypt_key (Ptr{Uint8}, Int32, Ptr{AES_KEY}) libcrypto
@c None AES_encrypt (Ptr{Uint8}, Ptr{Uint8}, Ptr{AES_KEY}) libcrypto
@c None AES_decrypt (Ptr{Uint8}, Ptr{Uint8}, Ptr{AES_KEY}) libcrypto
@c None AES_ecb_encrypt (Ptr{Uint8}, Ptr{Uint8}, Ptr{AES_KEY}, Int32) libcrypto
@c None AES_cbc_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Int32) libcrypto
@c None AES_cfb128_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Int32}, Int32) libcrypto
@c None AES_cfb1_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Int32}, Int32) libcrypto
@c None AES_cfb8_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Int32}, Int32) libcrypto
@c None AES_ofb128_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Int32}) libcrypto
@c None AES_ctr128_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Uint8}, Ptr{Uint32}) libcrypto
@c None AES_ige_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{Uint8}, Int32) libcrypto
@c None AES_bi_ige_encrypt (Ptr{Uint8}, Ptr{Uint8}, Csize_t, Ptr{AES_KEY}, Ptr{AES_KEY}, Ptr{Uint8}, Int32) libcrypto
@c Int32 AES_wrap_key (Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Uint8}, Ptr{Uint8}, Uint32) libcrypto
@c Int32 AES_unwrap_key (Ptr{AES_KEY}, Ptr{Uint8}, Ptr{Uint8}, Ptr{Uint8}, Uint32) libcrypto


import JLD, FileIO
import JSON
import PyPlot
using LightGraphs, MetaGraphs

include("calc_xq.jl")
include("make_roadnet_figs.jl")
include("utilities.jl")
include("plot_rwd_dists.jl")

# make experiment dict
experiment_name = "mturk_fast"
f_loc = "logs"
include("experiment_utilities.jl")

i_loc = "imgs"
myval = []

run = Dict(:figs=>true,:xq=>true,:json=>true)

if run[:figs]
    info("Making Figures")
    make_figs(exp_dict,f_loc,i_loc)
end
if run[:xq]
    info("Calculating xQ")
    # run this second, because we create a new file based on solver to which reference is compared
    calculate_sq(exp_dict)
end

function get_exit_node(g::MetaGraphs.MetaGraph)
    e_nodes = g.gprops[:exit_nodes]
    if length(e_nodes) > 1
        error("something is wrong, should only be length 1 in this experiment")
    else
        return e_nodes[1]
    end
end

if run[:json]
    info("## Making .json file for experiment data")
    srand(5555) #for consistent results
    for expr in keys(exp_dict[:conditions])
        net_type = exp_dict[:name]
        inpts = exp_dict[:xQ][:inpts]

        fname =  net_type*"_$(expr)_"*make_label_from_keys(inpts)*"_solver"
        f = joinpath("logs",fname)

        info("processing $f")
        d = JLD.load(f*".jld")["problem_dict"]

        exp_json = Dict()

        for net in d
            # need to use a HACK because I accidentally placed 13 in the :exit_loc key, when the networks
            # actually have 8 as the exit in the :graph properties. the graph properties are used for
            # the mdp solution, so unless I want to re-create the data (which is expensive in time and money
            # I'm just going to write some hack code to get the exit node from the graph
            exit_loc = get_exit_node(net[2][:graph])
            dist_to_exit = length(a_star(net[2][:graph],net[2][:evader_start],exit_loc))
            dist_to_pursuer = length(a_star(net[2][:graph],net[2][:evader_start],net[2][:pursuer_start]))
            edge_to_node_ratio = ne(net[2][:graph])/nv(net[2][:graph])
            etnr = 2.5
            dte = 2.0
            dtp = 2.0
            if dist_to_exit < dte || dist_to_pursuer < dtp || edge_to_node_ratio > etnr
                # we don't want networks where exit or pursuer is too close at the beginning
                # or where the edge/node ratio is too high (i.e. too many edges)
                if dist_to_exit < dte
                    println("Exit too close (d=$dist_to_exit): skipping $(net[1])")
                end
                if dist_to_pursuer < dtp
                    println("Pursuer too close (d=$dist_to_pursuer): skipping $(net[1])")
                end
                if edge_to_node_ratio > etnr
                    println("Too many edges ($edge_to_node_ratio): skipping $(net[1])")
                end
                continue
            end

            n = exp_json[net[1]] = Dict()
            net_vals = net[2]

            n[:xQ] = net_vals[:solver_quality][:X_Q]
            n[:xP] = net_vals[:training_data][:X4]
            n[:dist_to_exit] = dist_to_exit
            n[:dist_to_pursuer] = dist_to_pursuer
            n[:edge_to_node_ratio] = edge_to_node_ratio

            # calculate whether experiment succeded or not, take reward distribution and draw at random
            # this will have a higher probability of drawing most repeated numbers, so it is like sampling
            # a parameterized distribution, with the limitaion that it will only draw an outcome that has
            # already been seen
            # If that value is greater than the mean expected reward of the trusted distribution (:X3_1) then the experiment is a success
            # otherwise it is a failure
            r_dist = net_vals[:training_data][:r_dist]
            myval = r_dist
            o = rand(r_dist)

            #  r_star_mean = net_vals[:training_data][:X3_1] # Old, WRONG way, leave hear for illustrative purposes
            r_star_mean = net_vals[:solver_quality][:R_star][:X3_1]
            #  if all(r_star_mean .> o) #old WRONG way
            #  if all(r_star_mean .> o) && all(o .< 0.0)
            if all(o .< 0.0)
                if n[:xQ] < 0.5 && n[:xP] > 0.5
                    # this is an unexpected outcome
                    plot_rwd_dists(net[1],n[:xQ],n[:xP],net_vals[:training_data][:r_dist],net_vals[:solver_quality][:R_star],o,fldr="figs/exp_data/"*net_type*"_$(expr)",fname="$(fname)_$(net[1])_bigfail",outcome=:fail)
                end
                if n[:xQ] > 1.15 && n[:xP] > 0.25
                    # this is an unexpected outcome
                    plot_rwd_dists(net[1],n[:xQ],n[:xP],net_vals[:training_data][:r_dist],net_vals[:solver_quality][:R_star],o,fldr="figs/exp_data/"*net_type*"_$(expr)",fname="$(fname)_$(net[1])_surprisefail",outcome=:fail)
                end
                outcome = "fail"
            else
                if n[:xQ] < 0.75 && n[:xP] < -0.1
                    # this is an unexpected outcome
                    plot_rwd_dists(net[1],n[:xQ],n[:xP],net_vals[:training_data][:r_dist],net_vals[:solver_quality][:R_star],o,fldr="figs/exp_data/"*net_type*"_$(expr)",fname="$(fname)_$(net[1])_surprisesuccess",outcome=:success)
                end
                outcome = "success"
            end
            n[:outcome] = outcome
            n[:num_nodes] = net_vals[:num_nodes]
            n[:trans_prob] = net_vals[:trans_prob]
            n[:rwd_draw] = o
            n[:trusted_mean] = r_star_mean
            outcome == "success" ? oc_bool = 1 : oc_bool = 0
            #  println("rand draw: $o")
            #  println("trusted mean: $r_star_mean")
            #  println("oucome: $outcome")
            #  readline()
            #  info(n[:outcome])

            n[:image_file] = joinpath(net_vals[:img_name])
            n[:mcts_depth] = net_vals[:mcts_depth]
        end

        json_fname = "logs/experiment_data_$(fname).json"
        info("Writing to $json_fname")
        open(json_fname,"w") do f
            dat = JSON.json(exp_json,4)
            write(f,dat)
        end
    end
end


#JuliaRL/src/display.jl

using ReinforcementLearning
using PyCall
using Random
using ReinforcementLearningEnvironments
using Flux
using Logging

include("./shared.jl")

"""
    display_agent(save_dir, duration::Int = 1000)

Runs the agent saved at `save_dir`, rendering the environment,
up to `duration` steps.
"""
function display_agent(save_dir, duration::Int = 5000)
    env = LunarLander()
    agent = RLCore.load(save_dir, Agent)
    Flux.testmode!(agent)

    stop_condition = StopAfterStep(duration)
    disp_hook = DoEveryNStep(1) do t, agent, env
        env.env.pyenv.render()
    end

    disp_hook = DoEveryNStep(1) do t, agent, env
        env.env.pyenv.render()
    end

    print_hook = DoEveryNEpisode(1) do t, agent, env
        println("- Ep N $t")
    end

    hook = ComposedHook(disp_hook, print_hook)

    run(agent, env, stop_condition, hook)
end


"""
    mutable struct SolidStateDetector{T <: SSDFloat, CS} <: AbstractConfig{T}

CS: Coordinate System: -> :cartesian / :cylindrical
"""
mutable struct SolidStateDetector{T <: SSDFloat, CS} <: AbstractConfig{T}
    name::String  # optional
    inputunits::Dict{String, Unitful.Units}
    world::World{T, 3}

    config_dict::Dict

    medium::NamedTuple # this should become a struct at some point

    semiconductors::Vector{Semiconductor{T}}
    contacts::Vector{Contact{T}}
    passives::Vector{Passive{T}}

    SolidStateDetector{T, CS}() where {T <: SSDFloat, CS} = new{T, CS}()
end

get_precision_type(d::SolidStateDetector{T}) where {T} = T
get_coordinate_system(d::SolidStateDetector{T, CS}) where {T, CS} = CS

function construct_units(dict::Dict)::Dict{String,Unitful.Units}
    result_dict::Dict{String,Unitful.Units} = Dict()
    haskey(dict,"length") ? result_dict["length"] = unit_conversion[dict["length"]] : result_dict["length"] = u"mm"
    haskey(dict,"angle") ? result_dict["angle"] = unit_conversion[dict["angle"]] : result_dict["angle"] = u"rad"
    haskey(dict,"potential") ? result_dict["potential"] = unit_conversion[dict["potential"]] : result_dict["potential"] = u"V"
    haskey(dict,"temperature") ? result_dict["temperature"] = unit_conversion[dict["temperature"]] : result_dict["temperature"] = u"K"
    result_dict
end

function construct_semiconductor(T, sc::Dict, inputunit_dict::Dict{String, Unitful.Units})
    Semiconductor{T}(sc, inputunit_dict)
end

function construct_passive(T, pass::Dict, inputunit_dict::Dict{String, Unitful.Units})
    Passive{T}(pass, inputunit_dict)
end

function construct_contact(T, contact::Dict, inputunit_dict::Dict{String, Unitful.Units})
    Contact{T}(contact, inputunit_dict)
end

function construct_objects(T, objects::Vector, semiconductors, contacts, passives, inputunit_dict)::Nothing
    for obj in objects
        if obj["type"] == "semiconductor"
            push!(semiconductors, construct_semiconductor(T, obj, inputunit_dict))
        elseif obj["type"] == "contact"
            push!(contacts, construct_contact(T, obj, inputunit_dict))
        elseif obj["type"] == "passive"
            push!(passives, construct_passive(T, obj, inputunit_dict))
        else
            @warn "please spcify the calss to bei either a \"semiconductor\", a \"contact\", or \"passive\""
        end
    end
    nothing
end

function SolidStateDetector{T}(config_file::Dict)::SolidStateDetector{T} where{T <: SSDFloat}
    grid_type = Symbol(config_file["setup"]["grid"]["coordinates"])
    c = SolidStateDetector{T, grid_type}()
    c.name = config_file["name"]
    c.config_dict = config_file
    c.inputunits = construct_units(config_file["setup"]["units"])
    c.world = World(T, config_file["setup"]["grid"], c.inputunits)
    c.medium = material_properties[materials[config_file["setup"]["medium"]]]
    c.semiconductors, c.contacts, c.passives = [], [], []
    construct_objects(T, config_file["setup"]["objects"], c.semiconductors, c.contacts, c.passives, c.inputunits)
    return c
end

function SolidStateDetector(parsed_dict::Dict)
    SolidStateDetector{Float32}(parsed_dict)
end

function Base.sort!(v::AbstractVector{<:AbstractGeometry})
    hierarchies::Vector{Int} = map(x->x.hierarchy,v)
    v_result::typeof(v) = []
    for idx in sort!(unique!(hierarchies))
        push!(v_result,filter(x->x.hierarchy == hierarchies[idx],v)...)
    end
    return v_result
end

function SolidStateDetector{T}(parsed_dict::Dict) where T
    SolidStateDetector{T}(parsed_dict)
end

function contains(c::SolidStateDetector, point::AbstractCoordinatePoint{T,3})::Bool where T
    for contact in c.contacts
        if point in contact
            return true
        end
    end
    for sc in c.semiconductors
        if point in sc
            return true
        end
    end
    return false
end

function println(io::IO, d::SolidStateDetector{T, CS}) where {T <: SSDFloat, CS}
    println("________"*d.name*"________\n")
    # println("Class: ",d.class)
    println("---General Properties---")
    println("-Environment Material: \t $(d.medium.name)")
    println("-Grid Type: \t $(CS)")
    println()
    println("# Semiconductors: $(length(d.semiconductors))")
    for (isc, sc)  in enumerate(d.semiconductors)
        println("\t_____Semiconductor $(isc)_____\n")
        println(sc)
    end
    println()
    println("# Contacts: $(length(d.contacts))")
    if length(d.contacts)<=5
        for c in d.contacts
            println(c)
        end
    end
    println()
    println("# Passives: $(length(d.passives))")
    if length(d.passives)<=5
        for p in d.passives
            # println(c)
        end
    end
end

function show(io::IO, d::SolidStateDetector{T}) where {T <: SSDFloat} println(d) end
function print(io::IO, d::SolidStateDetector{T}) where {T <: SSDFloat} println(d) end
function display(io::IO, d::SolidStateDetector{T} ) where {T <: SSDFloat} println(d) end
function show(io::IO,::MIME"text/plain", d::SolidStateDetector) where {T <: SSDFloat}
    show(io, d)
end

# ToDo: Test it
function generate_random_startpositions(d::SolidStateDetector{T}, n::Int, Volume::NamedTuple=bounding_box(d), rng::AbstractRNG = MersenneTwister(), min_dist_from_boundary = 0.0001) where T
    delta = T(min_dist_from_boundary)
    n_filled::Int = 0
    positions = Vector{CartesianPoint{T}}(undef,n)
    while n_filled < n
        sample=CylindricalPoint{T}(rand(rng,Volume[:r_range].left:0.00001:Volume[:r_range].right),rand(rng,Volume[:φ_range].left:0.00001:Volume[:φ_range].right),rand(rng,Volume[:z_range].left:0.00001:Volume[:z_range].right))
        if !(sample in d.contacts) && contains(d,sample) && contains(d,CylindricalPoint{T}(sample.r+delta,sample.φ,sample.z))&& contains(d,CylindricalPoint{T}(sample.r-delta,sample.φ,sample.z))&& contains(d,CylindricalPoint{T}(sample.r,sample.φ,sample.z+delta))&& contains(d,CylindricalPoint{T}(sample.r,sample.φ,sample.z-delta))
            n_filled += 1
            positions[n_filled]=CartesianPoint(sample)
        end
    end
    positions
end


average(a::T, b::T) where {T<:AbstractFloat} = (a + b) / 2

cpad(s::String, n::Integer) = rpad(lpad(s, (n + textwidth(s)) >> 1), n)

"""
densify(S::SparseMatrix, threshold=0.25)

Convert sparse `S` to `Diagonal` if `S` is diagonal or to `Array(S)` if
the proportion of nonzeros exceeds `threshold`.
"""
function densify(A::SparseMatrixCSC, threshold::Real = 0.25)
    m, n = size(A)
    if m == n && isdiag(A)  # convert diagonal sparse to Diagonal
        # the diagonal is always dense (otherwise rank deficit)
        # so make sure it's stored as such
        Diagonal(Vector(diag(A)))
    elseif nnz(A)/(m * n) ≤ threshold
        A
    else
        Array(A)
    end
end
densify(A::AbstractMatrix, threshold::Real = 0.3) = A

densify(A::SparseVector, threshold::Real = 0.3)  = Vector(A)
densify(A::Diagonal{T, SparseVector}, threshold::Real = 0.3) where T = Diagonal(Vector(A.diag))

"""
    RaggedArray{T,I}

A "ragged" array structure consisting of values and indices

# Fields
- `vals`: a `Vector{T}` containing the values
- `inds`: a `Vector{I}` containing the indices

For this application a `RaggedArray` is used only in its `sum!` method.
"""
struct RaggedArray{T,I}
    vals::Vector{T}
    inds::Vector{I}
end
function Base.sum!(s::AbstractVector{T}, a::RaggedArray{T}) where T
    for (v, i) in zip(a.vals, a.inds)
        s[i] += v
    end
    s
end

"""
    normalized_variance_cumsum(A::AbstractMatrix)

Return the cumulative sum of the squared singular values of `A` normalized to sum to 1
"""
function normalized_variance_cumsum(A::AbstractMatrix)
    vars = cumsum(abs2.(svdvals(A)))
    vars ./ last(vars)
end

"""
    ltriindprs

A row-major order `Vector{NTuple{2,Int}}` of indices in the strict lower triangle.
"""
const ltriindprs = NTuple{2,Int}[]

function checkindprsk(k::Integer)
    kchoose2 = (k * (k - 1)) >> 1
    if length(ltriindprs) < kchoose2
        sizehint!(empty!(ltriindprs), kchoose2)
        for i in 1:k, j in 1:(i-1)
            push!(ltriindprs, (i,j))
        end
    end
    ltriindprs
end

"""
    subscriptednames(nm, len)

Return a `Vector{String}` of `nm` with subscripts from `₁` to `len`
"""
function subscriptednames(nm, len)
    nd = ndigits(len)
    nd == 1 ?
        [string(nm, '₀' + j) for j in 1:len] :
        [string(nm, lpad(string(j), nd, '0')) for j in 1:len]
end
#=
updatedevresid!(r::GLM.GlmResp, η::AbstractVector) = updateμ!(r, η)

fastlogitdevres(η, y) = 2log1p(exp(iszero(y) ? η : -η))

function updatedevresid!(r::GLM.GlmResp{V,<:Bernoulli,LogitLink}, η::V) where V<:AbstractVector{<:AbstractFloat}
    map!(fastlogitdevres, r.devresid, η, r.y)
    r
end
=#


"""
like bloom_best2, but with continue in main loop instead of elseif
no significant change (as expected)

"""
function bloom_best end

function bloom_bestc(t::SearchSequence)
    n = sizeof(t)
    skip = n
    tlast = _nthbyte(t,n)
    bloom_mask = UInt64(_search_bloom_mask(tlast))
    bloom_bits = 1 # no. of bits set in bloom filter if <= 32
    bloom_skip = 1 # no. of bytes to skip if byte not in filter
    j = n
    while (j-=1)>=1  # j=n is already processed
        tj = _nthbyte(t,j)
        if tj == tlast && skip == n
            skip = n - j
        end
        if bloom_bits <= 32 # argument: is near max(p(bloom_skip)*bloom_skip)
            hash = _search_bloom_mask(tj)
            if hash&bloom_mask == 0
                if bloom_bits < 32
                    # put in bloom filter up to 32 bits
                    bloom_mask |= hash
                end
                bloom_bits += 1 # gets >32 to terminate bloom filter filling
            end
            if bloom_bits <= 32
                bloom_skip += 1
            end
        else
            if skip<n
                # we have finished adding hashes to bloom filter
                # and we have determined the skip distance if matching last byte
                # nothimg remains to be done in preprocessing - stop work.
                break
            end
        end
    end
    return t,bloom_mask,bloom_skip,bloom_bits,skip,tlast
end

function bloom_bestc(s::SearchSequence, t::SearchSequence, i::Integer)
    bloom_bestc(s,bloom_bestc(t),i)
end

function bloom_bestc(s::SearchSequence,p::Tuple,i::Integer,sv::MaybeVector=nothing)
    (t,bloom_mask,bloom_skip,bloom_bits,skip,tlast) = p
    DOSTATS = !(sv isa Nothing)
    n = sizeof(t)
    m = sizeof(s)
    if n == 0
        return 1 <= i <= m+1 ? max(1, i) : 0
    elseif m == 0
        return 0
    elseif n == 1
        return something(findnext(isequal(_nthbyte(t,1)), s, i), 0)
    end
    w = m - n
    if w < 0 || i - 1 > w
        return 0
    end

    tlast = _nthbyte(t,n)
    if DOSTATS loops = 0 end
    if DOSTATS bloomtests = 0 end
    if DOSTATS bloomskips = 0 end
    i +=n-1
    if bloom_bits <= 12 # best guess from benchmarks
        # do bloom test first in loop
        while i < m
            if DOSTATS loops += 1 end
            if DOSTATS bloomtests += 1 end
            if bloom_mask & _search_bloom_mask(_nthbyte(s,i)) == 0
                if DOSTATS bloomskips += 1 end
                i += bloom_skip
                continue
            end
            if _nthbyte(s,i) == tlast
                # check candidate
                j = 1
                while j < n
                    if _nthbyte(s,i-n+j) != _nthbyte(t,j)
                        break
                    end
                    j += 1
                    # match found?
                    if j == n
                        if DOSTATS sv[Int(SFloops)] = loops; sv[Int(SFtests)] = bloomtests; sv[Int(SFskips)] = bloomskips; sv[Int(SFbits)] = bitcount(bloom_mask) end
                        return i-n+1
                    end
                end
                # no match: skip and test bloom
                i += skip
                continue
            end
            i +=1
        end
    else
        # do byte test first
        while i < m
            if DOSTATS loops += 1 end
            if _nthbyte(s,i) == tlast
                # check candidate
                j = 1
                while j < n
                    if _nthbyte(s,i-n+j) != _nthbyte(t,j)
                        break
                    end
                    j += 1
                    # match found?
                    if j == n
                        if DOSTATS recordcase(sv, loops, bloomtests, bloomskips, bitcount(bloom_mask), skip) end
                        return i-n+1
                    end
                end
                # no match: skip and test bloom
                i += skip
                if DOSTATS bloomtests += 1 end
                if i>=w
                    break
                end
                continue
            end
            i += 1
            if DOSTATS bloomtests += 1 end
            if bloom_mask & _search_bloom_mask(_nthbyte(s,i)) == 0
                if DOSTATS bloomskips += 1 end
                i += bloom_skip
            end
        end
    end
    if i==m
        # test end match
        j = 1
        while j <= n
            if _nthbyte(s,i-n+j) != _nthbyte(t,j)
                break # not found
            end
            if j == n
                if DOSTATS recordcase(sv, loops, bloomtests, bloomskips, bitcount(bloom_mask), skip) end
                return i-n+1
            end # match at the very end
            j += 1
        end
    end
    if DOSTATS recordcase(sv, loops, bloomtests, bloomskips, bitcount(bloom_mask), skip) end
    0
end


#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons, Josef Kircher
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

using DifferentialEquations, DiffEqCallbacks

############ Model-Exchange ############

# Read next time event from fmu and provide it to the integrator
function time_choice(c::fmi2Component, integrator)
    eventInfo = fmi2NewDiscreteStates(c)
    fmi2EnterContinuousTimeMode(c)
    if Bool(eventInfo.nextEventTimeDefined)
        eventInfo.nextEventTime
    else
        Inf
    end
end

# Handles events and returns the values and nominals of the changed continuous states.
function handleEvents(c::fmi2Component, enterEventMode::Bool, exitInContinuousMode::Bool)

    if enterEventMode
        fmi2EnterEventMode(c)
    end

    eventInfo = fmi2NewDiscreteStates(c)

    valuesOfContinuousStatesChanged = eventInfo.valuesOfContinuousStatesChanged
    nominalsOfContinuousStatesChanged = eventInfo.nominalsOfContinuousStatesChanged

    #set inputs here
    #fmiSetReal(myFMU, InputRef, Value)

    while eventInfo.newDiscreteStatesNeeded == fmi2True
        # update discrete states
        eventInfo = fmi2NewDiscreteStates(c)
        valuesOfContinuousStatesChanged = eventInfo.valuesOfContinuousStatesChanged
        nominalsOfContinuousStatesChanged = eventInfo.nominalsOfContinuousStatesChanged

        if eventInfo.terminateSimulation == fmi2True
            @error "Event info returned error!"
        end
    end

    if exitInContinuousMode
        fmi2EnterContinuousTimeMode(c)
    end

    return valuesOfContinuousStatesChanged, nominalsOfContinuousStatesChanged

end

# Returns the event indicators for an FMU.
function condition(c::fmi2Component, out, x, t, integrator) # Event when event_f(u,t) == 0

    fmi2SetTime(c, t)
    fmi2SetContinuousStates(c, x)
    indicators = fmi2GetEventIndicators(c)

    copy!(out, indicators)

end

# Handles the upcoming events.
function affectFMU!(c::fmi2Component, integrator, idx)
    # Event found - handle it
    continuousStatesChanged, nominalsChanged = handleEvents(c, true, Bool(sign(idx)))

    if continuousStatesChanged == fmi2True
        integrator.u = fmi2GetContinuousStates(c)
    end

    if nominalsChanged == fmi2True
        x_nom = fmi2GetNominalsOfContinuousStates(c)
    end
    timeEventCb = PresetTimeCallback(2.0, (integrator) -> affect!(c, integrator, 0))
end

# Does one step in the simulation.
function stepCompleted(c::fmi2Component, x, t, integrator)

     (status, enterEventMode, terminateSimulation) = fmi2CompletedIntegratorStep(c, fmi2True)
     if enterEventMode == fmi2True
        affect!(c, integrator, 0)
     end

end

# Returns the state derivatives of the FMU.
function fx(c::fmi2Component, x, p, t)
    fmi2SetTime(c, t)
    fmi2SetContinuousStates(c, x)
    dx = fmi2GetDerivatives(c)
end

# save FMU values
function saveValues(c::fmi2Component, recordValues, u, t, integrator)
    Tuple(fmiGetReal(c, recordValues)...,)
end

"""
Source: FMISpec2.0.2[p.90 ff]: 3.2.4 Pseudocode Example

Simulates a FMU instance for the given simulation time interval.
State- and Time-Events are handled correctly.

Returns a tuple of type (ODESolution, DiffEqCallbacks.SavedValues).
If keyword `recordValues` is not set, a tuple of type (ODESolution, nothing) is returned for consitency.
"""
function fmi2SimulateME(c::fmi2Component, t_start::Real = 0.0, t_stop::Real = 1.0;
    solver = nothing,
    customFx = nothing,
    recordValues::fmi2ValueReferenceFormat = nothing,
    saveat = [],
    setup = true)

    recordValues = prepareValueReference(c, recordValues)
    solution = nothing
    callbacks = []
    savedValues = nothing

    if length(recordValues) > 0
        savedValues = SavedValues(Float64, Tuple{collect(Float64 for i in 1:length(recordValues))...})

        savingCB = SavingCallback((u,t,integrator) -> saveValues(c, recordValues, u, t, integrator),
                                  savedValues,
                                  saveat=saveat)
        push!(callbacks, savingCB)
    end

    if solver == nothing
        solver = Tsit5()
    end

    if customFx == nothing
        customFx = (x, p, t) -> fx(c, x, p, t)
    end

    if setup
        fmi2Reset(c)
        fmi2SetupExperiment(c, t_start, t_stop)
        fmi2EnterInitializationMode(c)
        fmi2ExitInitializationMode(c)
    end

    eventHandling = c.fmu.modelDescription.numberOfEventIndicators > 0

    # First evaluation of the FMU
    x0 = fmi2GetContinuousStates(c)
    x0_nom = fmi2GetNominalsOfContinuousStates(c)

    fmi2SetContinuousStates(c, x0)

    handleEvents(c, false, false)

    # Get states of handling initial Events
    x0 = fmi2GetContinuousStates(c)
    x0_nom = fmi2GetNominalsOfContinuousStates(c)

    p = []
    problem = ODEProblem(customFx, x0, (t_start, t_stop), p,)

    if eventHandling

        eventInfo = fmi2NewDiscreteStates(c)
        fmi2EnterContinuousTimeMode(c)

        timeEvents = (eventInfo.nextEventTimeDefined == fmi2True)

        eventCb = VectorContinuousCallback((out, x, t, integrator) -> condition(c, out, x, t, integrator),
                                           (integrator, idx) -> affectFMU!(c, integrator, idx),
                                           Int64(c.fmu.modelDescription.numberOfEventIndicators);
                                           rootfind = DiffEqBase.RightRootFind)
        push!(callbacks, eventCb)

        stepCb = FunctionCallingCallback((x, t, integrator) -> stepCompleted(c, x, t, integrator);
                                         func_everystep = true,
                                         func_start = true)
        push!(callbacks, stepCb)

        if timeEvents
            timeEventCb = IterativeCallback((integrator) -> time_choice(c, integrator),
                                            (integrator) -> affectFMU!(c, integrator, 0), Float64; initial_affect = true)

            push!(callbacks, timeEventCb)
        end
    end

    if length(callbacks) > 0
        solution = solve(problem, solver, callback = CallbackSet(callbacks...), saveat = saveat)
    else
        solution = solve(problem, solver, saveat = saveat)
    end

    return solution, savedValues
end

############ Co-Simulation ############

"""
Starts a simulation of the Co-Simulation FMU instance.

Returns a tuple of (success::Bool, DiffEqCallbacks.SavedValues) with success = `true` or `false`.
If keyword `recordValues` is not set, a tuple of type (success::Bool, nothing) is returned for consitency.

ToDo: Improove Documentation.
"""
function fmi2SimulateCS(c::fmi2Component, t_start::Real, t_stop::Real;
                        recordValues::fmi2ValueReferenceFormat = nothing,
                        saveat = [],
                        setup = true)

    recordValues = prepareValueReference(c, recordValues)
    variableSteps = c.fmu.modelDescription.CScanHandleVariableCommunicationStepSize

    success = false
    savedValues = nothing

    # default setup
    if length(saveat) == 0
        saveat = LinRange(t_start, t_stop, 100)
    end
    dt = (t_stop - t_start) / length(saveat)

    # setup if no variable steps
    if variableSteps == false
        if length(saveat) >= 2
            dt = saveat[2] - saveat[1]
        end
    end

    if setup
        fmi2Reset(c)
        fmi2SetupExperiment(c, t_start, t_stop)
        fmi2EnterInitializationMode(c)
        fmi2ExitInitializationMode(c)
    end

    t = t_start

    record = length(recordValues) > 0

    #numDigits = length(string(round(Integer, 1/dt)))

    if record
        savedValues = SavedValues(Float64, Tuple{collect(Float64 for i in 1:length(recordValues))...})

        i = 1

        values = (fmi2GetReal(c, recordValues)...,)
        DiffEqCallbacks.copyat_or_push!(savedValues.t, i, t)
        DiffEqCallbacks.copyat_or_push!(savedValues.saveval, i, values, Val{false})

        while t < t_stop
            if variableSteps
                if length(saveat) > i
                    dt = saveat[i+1] - saveat[i]
                else
                    dt = t_stop - saveat[i]
                end
            end

            fmi2DoStep(c, t, dt)
            t = t + dt #round(t + dt, digits=numDigits)
            i += 1

            values = (fmi2GetReal(c, recordValues)...,)
            DiffEqCallbacks.copyat_or_push!(savedValues.t, i, t)
            DiffEqCallbacks.copyat_or_push!(savedValues.saveval, i, values, Val{false})
        end

        success = true
    else
        i = 1
        while t < t_stop
            if variableSteps
                if length(saveat) > i
                    dt = saveat[i+1] - saveat[i]
                else
                    dt = t_stop - saveat[i]
                end
            end

            fmi2DoStep(c, t, dt)
            t = t + dt #round(t + dt, digits=numDigits)
            i += 1
        end

        success = true
    end

    success, savedValues
end

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

"""
Starts a simulation of the fmu instance for the matching fmu type, if both types are available, CS is preferred.

Returns:
    - a tuple of (success::Bool, DiffEqCallbacks.SavedValues) with success = `true` or `false` for CS-FMUs
    - a tuple of (ODESolution, DiffEqCallbacks.SavedValues) for ME-FMUs
    - if keyword `recordValues` is not set, a tuple of type (..., nothing)

ToDo: Improove Documentation.
"""
function fmi2Simulate(c::fmi2Component, t_start::Real = 0.0, t_stop::Real = 1.0;
                      recordValues::fmi2ValueReferenceFormat = nothing,
                      saveat = [],
                      setup = true)

    if fmi2IsCoSimulation(c.fmu)
        return fmi2SimulateCS(c, t_start, t_stop; recordValues=recordValues, saveat=saveat, setup=setup)
    elseif fmi2IsModelExchange(c.fmu)
        return fmi2SimulateME(c, t_start, t_stop; recordValues=recordValues, saveat=saveat, setup=setup)
    else
        error(unknownFMUType)
    end
end


# # Rabbit, Fox, Hawk
# ```@raw html
# <video width="auto" controls autoplay loop>
# <source src="../rabbit_fox_hawk.mp4" type="video/mp4">
# </video>
# ```
# This model is a variation on the [Predator-prey dynamics](@ref) example. It uses a 3-dimensional
# [`ContinuousSpace`](@ref), a realistic terrain for the agents, and pathfinding (with multiple
# pathfinders).
#
# Agents in this model are one of three species of animals: rabbits, foxes and hawks. Rabbits
# eat grass, and are hunted by foxes and hawks. While rabbits and foxes are restricted to
# walk on suitable portions of the map, hawks are capable of flight and can fly over a much
# larger region of the map.
#
# Similar to the [Predator-prey dynamics](@ref) example, agent types are distinguished using a
# `type` field. Agents also have an additional `energy` field, which is consumed to move around
# and reproduce. Eating food (grass or rabbits) replenishes `energy` by a fixed amount.
using Agents, Agents.Pathfinding
using Random
using FileIO # To load images you also need ImageMagick available to your project

mutable struct Animal <: AbstractAgent
    id::Int
    pos::NTuple{3,Float64}
    type::Symbol ## one of :rabbit, :fox or :hawk
    energy::Float64
end

## Some utility functions to create specific types of agents, and find the norm of a vector
Rabbit(id, pos, energy) = Animal(id, pos, :rabbit, energy)
Fox(id, pos, energy) = Animal(id, pos, :fox, energy)
Hawk(id, pos, energy) = Animal(id, pos, :hawk, energy)
norm(vec) = √sum(vec .^ 2)

# The environment is generated from a heightmap: a 2D matrix, where each value denotes the
# height of the terrain at that point. We segregate the model into 4 regions based on the
# height:
# - Anything below `water_level` is waster and cannot be walked on
# - The region between `water_level` and `grass_level` is flatland, that can be walked on
# - The part of the map between `grass_level` and `mountain_level` is too high for animals to
#   walk over, but it can be flown over
# - The terrain above `mountain_level` is completely inaccessible
#
# Grass is the food source for rabbits. It can grow anywhere from `water_level` to `grass_level`.
# The spread of grass across the terrain is specified using a BitArray. A value of
# 1 at a location indicates the presence of grass there, which can be consumed when it is eaten
# by a rabbit. The probability of grass growing is proportional to how close it is to the water.
#
# The `initialize_model` function takes in the URL to our heightmap, the thresholds for the 4
# regions, and some additional parameters for the model. It then creates and returns a model
# with the specified heightmap and containing the specified number of rabbits, foxes and hawks.

function initialize_model(
    heightmap_url,
    water_level = 8,
    grass_level = 20,
    mountain_level = 35;
    n_rabbits = 160,  ## initial number of rabbits
    n_foxes = 30,  ## initial number of foxes
    n_hawks = 30,  ## initial number of hawks
    Δe_grass = 25,  ## energy gained from eating grass
    Δe_rabbit = 30,  ## energy gained from eating one rabbit
    rabbit_repr = 0.06,  ## probability for a rabbit to (asexually) reproduce at any step
    fox_repr = 0.03,  ## probability for a fox to (asexually) reproduce at any step
    hawk_repr = 0.02, ## probability for a hawk to (asexually) reproduce at any step
    rabbit_vision = 6,  ## how far rabbits can see grass and spot predators
    fox_vision = 10,  ## how far foxes can see rabbits to hunt
    hawk_vision = 15,  ## how far hawks can see rabbits to hunt
    rabbit_speed = 1.3, ## movement speed of rabbits
    fox_speed = 1.1,  ## movement speed of foxes
    hawk_speed = 1.2, ## movement speed of hawks
    regrowth_chance = 0.03,  ## probability that a patch of grass regrows at any step
    dt = 0.1,   ## discrete timestep each iteration of the model
    seed = 42,  ## seed for random number generator
)

    ## Download and load the heightmap. The grayscale value is converted to `Float64` and
    ## scaled from 1 to 40
    heightmap = floor.(Int, convert.(Float64, load(download(heightmap_url))) * 39) .+ 1
    ## The x and y dimensions of the pathfinder are that of the heightmap
    dims = (size(heightmap)..., 50)
    ## The region of the map that is accessible to each type of animal (land-based or flying)
    ## is defined using `BitArrays`
    land_walkmap = BitArray(falses(dims...))
    air_walkmap = BitArray(falses(dims...))
    for i in 1:dims[1], j in 1:dims[2]
        ## land animals can only walk on top of the terrain between water_level and grass_level
        if water_level < heightmap[i, j] < grass_level
            land_walkmap[i, j, heightmap[i, j]+1] = true
        end
        ## air animals can fly at any height upto mountain_level
        if heightmap[i, j] < mountain_level
            air_walkmap[i, j, (heightmap[i, j]+1):mountain_level] .= true
        end
    end

    ## Generate the RNG for the model
    rng = MersenneTwister(seed)

    ## Note that the dimensions of the space do not have to correspond to the dimensions
    ## of the pathfinder. Discretisation is handled by the pathfinding methods
    space = ContinuousSpace((100., 100., 50.); periodic = false)

    ## Generate an array of random numbers, and threshold it by the probability of grass growing
    ## at that location. Although this causes grass to grow below `water_level`, it is
    ## effectively ignored by `land_walkmap`
    grass = BitArray(
        rand(rng, dims[1:2]...) .< ((grass_level .- heightmap) ./ (grass_level - water_level)),
    )
    properties = (
        ## The pathfinder for rabbits and foxes
        landfinder = AStar(space; walkmap = land_walkmap),
        ## The pathfinder for hawks
        airfinder = AStar(space; walkmap = air_walkmap, cost_metric = MaxDistance{3}()),
        Δe_grass = Δe_grass,
        Δe_rabbit = Δe_rabbit,
        rabbit_repr = rabbit_repr,
        fox_repr = fox_repr,
        hawk_repr = hawk_repr,
        rabbit_vision = rabbit_vision,
        fox_vision = fox_vision,
        hawk_vision = hawk_vision,
        rabbit_speed = rabbit_speed,
        fox_speed = fox_speed,
        hawk_speed = hawk_speed,
        heightmap = heightmap,
        grass = grass,
        regrowth_chance = regrowth_chance,
        water_level = water_level,
        grass_level = grass_level,
        dt = dt,
    )

    model = ABM(Animal, space; rng, properties)

    ## spawn each animal at a random walkable position according to its pathfinder
    for _ in 1:n_rabbits
        add_agent_pos!(
            Rabbit(
                nextid(model), ## Using `nextid` prevents us from having to manually keep track
                               ## of animal IDs
                random_walkable(model, model.landfinder),
                rand(model.rng, Δe_grass:2Δe_grass),
            ),
            model,
        )
    end
    for _ in 1:n_foxes
        add_agent_pos!(
            Fox(
                nextid(model),
                random_walkable(model, model.landfinder),
                rand(model.rng, Δe_rabbit:2Δe_rabbit),
            ),
            model,
        )
    end
    for _ in 1:n_hawks
        add_agent_pos!(
            Hawk(
                nextid(model),
                random_walkable(model, model.airfinder),
                rand(model.rng, Δe_rabbit:2Δe_rabbit),
            ),
            model,
        )
    end

    return model
end

# The `animal_step!` function dispatches to the proper function depending on the type of agent.
# The stepping functions for each type of agent are similar: They lose energy per step, and
# die if their energy ever reaches 0. They also have a random probability to reproduce at an
# iteration. Agents all move towards their food. In the case of rabbits, they also move away
# from any nearby predators.

function animal_step!(animal, model)
    if animal.type == :rabbit
        rabbit_step!(animal, model)
    elseif animal.type == :fox
        fox_step!(animal, model)
    else
        hawk_step!(animal, model)
    end
end

# Rabbits eat grass at their position, if it exists. If they see a predator, they run away.
# The direction in which they flee is dependent on all predators in their vision, with closer
# ones contributing more to the chosen direction. If there are no predators to flee from,
# rabbits walk around randomly.

function rabbit_step!(rabbit, model)
    ## Eat grass at this position, if any
    if get_spatial_property(rabbit.pos, model.grass, model) == 1
        model.grass[get_spatial_index(rabbit.pos, model.grass, model)] = 0
        rabbit.energy += model.Δe_grass
    end

    ## The energy cost at each step corresponds to the amount of time that has passed
    ## since the last step
    rabbit.energy -= model.dt
    ## All animals die if their energy reaches 0
    if rabbit.energy <= 0
        kill_agent!(rabbit, model, model.landfinder)
        return
    end

    ## Get a list of positions of all nearby predators
    predators = [
        x.pos for x in nearby_agents(rabbit, model, model.rabbit_vision) if
            x.type == :fox || x.type == :hawk
            ]
    ## If the rabbit sees a predator and isn't already moving somewhere
    if !isempty(predators) && is_stationary(rabbit, model.landfinder)
        ## Try and get an ideal direction away from predators
        direction = (0., 0., 0.)
        for predator in predators
            ## Get the direction away from the predator
            away_direction = (rabbit.pos .- predator)
            ## In case there is already a predator at our location, moving anywhere is
            ## moving away from it, so it doesn't contribute to `direction`
            all(away_direction .≈ 0.) && continue
            ## Add this to the overall direction, scaling inversely with distance.
            ## As a result, closer predators contribute more to the direction to move in
            direction = direction .+ away_direction ./ norm(away_direction) ^ 2
        end
        ## If the only predator is right on top of the rabbit
        if all(direction .≈ 0.)
            ## Move anywhere
            chosen_position = random_walkable(rabbit.pos, model, model.landfinder, model.rabbit_vision)
        else
            ## Normalize the resultant direction, and get the ideal position to move it
            direction = direction ./ norm(direction)
            ## Move to a random position in the general direction of away from predators
            position = rabbit.pos .+ direction .* (model.rabbit_vision / 2.)
            chosen_position = random_walkable(position, model, model.landfinder, model.rabbit_vision / 2.)
        end
        plan_route!(rabbit, chosen_position, model.landfinder)
    end

    ## Reproduce with a random probability, scaling according to the time passed each
    ## step
    rand(model.rng) <= model.rabbit_repr * model.dt && reproduce!(rabbit, model)

    ## If the rabbit isn't already moving somewhere, move to a random spot
    if is_stationary(rabbit, model.landfinder)
        plan_route!(
            rabbit,
            random_walkable(rabbit.pos, model, model.landfinder, model.rabbit_vision),
            model.landfinder
        )
    end

    ## Move along the route planned above
    move_along_route!(rabbit, model, model.landfinder, model.rabbit_speed, model.dt)
end

# Foxes hunt for rabbits, and eat rabbits within a unit radius of its position.

function fox_step!(fox, model)
    ## Look for nearby rabbits that can be eaten
    food = [x for x in nearby_agents(fox, model) if x.type == :rabbit]
    if !isempty(food)
        kill_agent!(rand(model.rng, food), model, model.landfinder)
        fox.energy += model.Δe_rabbit
    end

    ## The energy cost at each step corresponds to the amount of time that has passed
    ## since the last step
    fox.energy -= model.dt
    ## All animals die once their energy reaches 0
    if fox.energy <= 0
        kill_agent!(fox, model, model.landfinder)
        return
    end

    ## Random chance to reproduce every step
    rand(model.rng) <= model.fox_repr * model.dt && reproduce!(fox, model)

    ## If the fox isn't alreadu moving somewhere
    if is_stationary(fox, model.landfinder)
        ## Look for any nearby rabbits
        prey = [x for x in nearby_agents(fox, model, model.fox_vision) if x.type == :rabbit]
        if isempty(prey)
            ## Move anywhere if no rabbits were found
            plan_route!(
                fox,
                random_walkable(fox.pos, model, model.landfinder, model.fox_vision),
                model.landfinder,
            )
            return
        end
        ## Move toward a random rabbit
        plan_route!(fox, rand(model.rng, map(x -> x.pos, prey)), model.landfinder)
    end

    move_along_route!(fox, model, model.landfinder, model.fox_speed, model.dt)
end

# Hawks function similarly to foxes, except they can also fly. They dive down for prey and
# fly back up after eating it.

function hawk_step!(hawk, model)
    ## Look for rabbits nearby
    food = [x for x in nearby_agents(hawk, model) if x.type == :rabbit]
    if !isempty(food)
        ## Eat (kill) the rabbit
        kill_agent!(rand(model.rng, food), model, model.airfinder)
        hawk.energy += model.Δe_rabbit
        ## Fly back up
        plan_route!(hawk, hawk.pos .+ (0., 0., 7.), model.airfinder)
    end

    ## The rest of the stepping function is similar to that of foxes, except hawks use a
    ## different pathfinder
    hawk.energy -= model.dt
    if hawk.energy <= 0
        kill_agent!(hawk, model, model.airfinder)
        return
    end

    rand(model.rng) <= model.hawk_repr * model.dt && reproduce!(hawk, model)

    if is_stationary(hawk, model.airfinder)
        prey = [x for x in nearby_agents(hawk, model, model.hawk_vision) if x.type == :rabbit]
        if isempty(prey)
            plan_route!(
                hawk,
                random_walkable(hawk.pos, model, model.airfinder, model.hawk_vision),
                model.airfinder,
            )
        else
            plan_route!(hawk, rand(model.rng, map(x -> x.pos, prey)), model.airfinder)
        end
    end

    move_along_route!(hawk, model, model.airfinder, model.hawk_speed, model.dt)
end

# This function is called when an animal reproduces. The animal loses half its energy, and
# a copy of it is created and added to the model.

function reproduce!(animal, model)
    animal.energy = ceil(Int, animal.energy / 2)
    add_agent_pos!(Animal(nextid(model), animal.pos, animal.type, animal.energy), model)
end

# The model stepping function simulates the growth of grass

function model_step!(model)
    ## To prevent copying of data, obtain a view of the part of the grass matrix that
    ## doesn't have any grass, and grass can grow there
    growable = view(
        model.grass,
        model.grass .== 0 .& model.water_level .< model.heightmap .<= model.grass_level,
    )
    ## Grass regrows with a random probability, scaling with the amount of time passing
    ## each step of the model
    growable .= rand(model.rng, length(growable)) .< model.regrowth_chance * model.dt
end

# ## Visualization
# Now we use `InteractiveDynamics` to create a visualization of the model running in 3D space
#
# The agents are color-coded according to their `type`, to make them easily identifiable in
# the visualization.

using InteractiveDynamics
using GLMakie
GLMakie.activate!() # hide

animalcolor(a) =
    if a.type == :rabbit
        :brown
    elseif a.type == :fox
        :orange
    else
        :blue
    end

# We use `surface!` to plot the terrain as a mesh, and colour it using the `:terrain`
# colormap. Since the heightmap dimensions don't correspond to the dimensions of the space,
# we explicitly provide ranges to specify where the heightmap should be plotted.
function static_preplot!(ax, model)
    surface!(
        ax,
        (100/205):(100/205):100,
        (100/205):(100/205):100,
        model.heightmap;
        colormap = :terrain
    )
end

# Passing in a sample heightmap to the `initialize_model` function we created returns the generated
# model.
heightmap_url =
    "https://raw.githubusercontent.com/JuliaDynamics/" *
    "JuliaDynamics/master/videos/agents/rabbit_fox_hawk_heightmap.png"
model = initialize_model(heightmap_url)

abm_video(
    "rabbit_fox_hawk.mp4",
    model,
    animal_step!,
    model_step!;
    resolution = (700, 700),
    frames = 300,
    framerate = 20,
    ac = animalcolor,
    as = 1.0,
    static_preplot!
)

nothing # hide
# ```@raw html
# <video width="auto" controls autoplay loop>
# <source src="../rabbit_fox_hawk.mp4" type="video/mp4">
# </video>
# ```


#################### Missing Values Sampler ####################

#################### Types and Constructors ####################

struct MISSTune
  dims::Tuple
  valueinds::Array
  distrinds::Array
end

function MISSTune(s::AbstractStochastic)
  MISSTune(s.distr, s.value)
end

function MISSTune(d::Distribution, v)
  MISSTune((), findall(isnan.(v)), Int[])
end

function MISSTune(D::Array{UnivariateDistribution}, v::Array)
  inds = findall(isnan.(v))
  MISSTune(dims(D), inds, inds)
end

function MISSTune(D::Array{MultivariateDistribution}, v::Array)
  isvalueinds = falses(size(v))
  isdistrinds = falses(size(D))
  for sub in CartesianIndices(size(D))
    n = length(D[sub])
    for i in 1:n
      if isnan(v[sub, i])
        isvalueinds[sub, i] = isdistrinds[sub] = true
      end
    end
  end
  MISSTune(dims(D), findall(isvalueinds), findall(isdistrinds))
end

#################### Sampler Constructor ####################
"""
    MISS(params::ElementOrVector{Symbol})

Construct a `Sampler` object to sampling missing output values. The constructor
should only be used to sample stochastic nodes upon which no other stochastic
node depends. So-called ‘output nodes’ can be identified with the `keys()`
function. Moreover, when the `MISS` constructor is included in a vector of
`Sampler` objects to define a sampling scheme, it should be positioned at the
beginning of the vector. This ensures that missing output values are updated
before any other samplers are executed.

Returns a `Sampler{Dict{Symbol, MISSTune}}` type object.

* `params`: stochastic node(s) that contain missing values (`NaN`) to be updated with the sampler.
"""
function MISS(params::ElementOrVector{Symbol})
  params = asvec(params)
  samplerfx = function(model::Model, block::Integer)
    tune = gettune(model, block)
    if model.iter == 1
      for key in params
        miss = MISSTune(model[key])
        if !isempty(miss.valueinds)
          tune[key] = miss
        end
      end
      params = intersect(keys(model, :dependent), keys(tune))
    end
    for key in params
      node = model[key]
      miss = tune[key]
      node[miss.valueinds] = rand(node, miss)
      update!(model, node.targets)
    end
    nothing
  end
  Sampler(params, samplerfx, Dict{Symbol, MISSTune}())
end

#################### Sampling Functions ####################

rand(s::AbstractStochastic, miss::MISSTune) = rand_sub(s.distr, miss)

function rand_sub(d::Distribution, miss::MISSTune)
  x = rand(d)
  Float64[x[i] for i in miss.valueinds]
end

function rand_sub(D::Array{UnivariateDistribution}, miss::MISSTune)
  Float64[rand(d) for d in D[miss.distrinds]]
end

function rand_sub(D::Array{MultivariateDistribution}, miss::MISSTune)
  X = Array{Float64}(undef, miss.dims)
  for i in miss.distrinds
    d = D[i]
    X[ind2sub(D, i)..., 1:length(d)] = rand(d)
  end
  X[miss.valueinds]
end


using Test
using TightlyBound
using Suppressor

#include("../includes_laguerre.jl")
#include("../Ewald.jl")

TESTDIR=TightlyBound.TESTDIR

function loaddata(dirs; scf=true)
    tbc_list  = []
    dft_list = []

    for t in dirs
        #                println(t*"/qe.save")
        tfull = "$TESTDIR/"*t
        dft = TightlyBound.QE.loadXML(tfull*"/qe.save")
        tbc = []
        tbc_scf = []
        try
            if scf
                tbc_scf = TightlyBound.TB.read_tb_crys("projham_scf.xml.gz", directory=tfull)
            else
                tbc_scf = TightlyBound.TB.read_tb_crys("projham.xml.gz", directory=tfull)
            end
        catch
            tbc = TightlyBound.AtomicProj.projwfx_workf(dft, directory=tfull, writefile="projham.xml", skip_og=true, skip_proj=true, freeze=true, localized_factor = 0.15)
            if scf
                tbc_scf = TightlyBound.SCF.remove_scf_from_tbc(tbc)
                TightlyBound.TB.write_tb_crys(t*"/projham_scf.xml.gz", tbc_scf)
            else
                tbc_scf = tbc
            end

        end

        push!(dft_list, dft)
        push!(tbc_list, tbc_scf)
    end
    return     tbc_list, dft_list
end

function test_force()

    @testset "testing force dimer" begin

        if true
#        @suppress begin
            ft = open("$TESTDIR/data_forces/fil_MgS_dimer", "r");
            dirst = readlines(ft);
            close(ft);

            #        println(dirst)

#            for scf = [false true]
            f_cart = zeros(2,3)
            f_cartFD = 0.0

            for scf = [false, true]
                @suppress begin
                tbc_list, dft_list = loaddata(dirst, scf=scf);
                     database_rec = TightlyBound.FitTB.do_fitting_recursive(tbc_list,dft_list = dft_list,  fit_threebody=false, fit_threebody_onsite=false);

                x = 4;
                smearing = 0.01;

                    #en, tbc_x, flag = scf_energy(tbc_list[x].crys, database = database_rec)
                    #en, f_cart,stress = TightlyBound.Force_Stress.get_energy_force_stress(tbc_x, database_rec,   smearing = smearing);

                    en, f_cart,stress = TightlyBound.Force_Stress.get_energy_force_stress(tbc_list[x].crys, database_rec,   smearing = smearing);

                enFD, f_cartFD = TightlyBound.Force_Stress.finite_diff(tbc_list[x].crys, database_rec,1, 3,   smearing = smearing);
                end

#                println(scf, " xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx  ", f_cartFD, "  ", f_cart[1,3])

                @test abs(f_cartFD - f_cart[1,3]) < 1e-3
#                @test abs(f_cartFD - f_cart_fft[1,3]) < 1e-3
                #            println("SCF $scf TEST1 finite diff: ", f_cartFD , " autodiff:   ", f_cart[1,3], " xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
                #            println("TEST1 dft ref ", dft_list[x].forces[1,3])

            end
            #        x = 3;
            #        smearing = 0.01;
            #        en, f_cart,stress = Force_Stress.get_energy_force_stress(tbc_list[x].crys, database_rec,   smearing = smearing);
            #        enFD, f_cartFD = Force_Stress.finite_diff(tbc_list[x].crys, database_rec,1, 3,   smearing = smearing);

            #        @test abs(f_cartFD - f_cart[1,3]) < 1e-3
            #        println("TEST3 finite diff: ", f_cartFD , " autodiff:   ", f_cart[1,3], " xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
            #        println("TEST3 dft ref ", dft_list[x].forces[1,3])

        end
    end
end

function test_stress()

    @testset "testing force znse" begin
        @suppress begin

            ft = open("$TESTDIR/data_forces/fil_MgS_znse", "r");
            dirst = readlines(ft);
            close(ft);

            #        println(dirst)

            tbc_list, dft_list = loaddata(dirst, scf=false);

            database = TightlyBound.FitTB.do_fitting(tbc_list, fit_threebody=false, fit_threebody_onsite=false, do_plot=false);
            #        database = FitTB.do_fitting_recursive(tbc_list,dft_list = dft_list,  fit_threebody=true, fit_threebody_onsite=false);

            x = 1;
            smearing = 0.01;
            en, f_cart, stress = TightlyBound.Force_Stress.get_energy_force_stress(tbc_list[x].crys, database,   smearing = smearing);

            enFD, f_cartFD = TightlyBound.Force_Stress.finite_diff(tbc_list[x].crys, database,1, 3,   smearing = smearing);

            #        println("TEST force finite diff: ", f_cartFD , " autodiff:   ", f_cart[1,3], " xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
            #        println("TEST dft ref ", dft_list[x].forces[1,3])
            @test abs(f_cartFD - f_cart[1,3]) < 1e-4

            x=1
            enFD, stressFD = TightlyBound.Force_Stress.finite_diff(tbc_list[x].crys, database,1, 1, stress_mode=true,  smearing = smearing);

            #        println("TEST stress11 finite diff: ", stressFD , " autodiff:   ", stress[1,1], " xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
            #        println("TEST dft ref ", dft_list[x].stress[1,1])
            @test abs(stressFD - stress[1,1]) < 1e-5

            x=1
            enFD, stressFD = TightlyBound.Force_Stress.finite_diff(tbc_list[x].crys, database,1, 2, stress_mode=true,  smearing = smearing);

            #        println("TEST stress12 finite diff: ", stressFD , " autodiff:   ", stress[1,2], " xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
            #        println("TEST dft ref ", dft_list[x].stress[1,2])
            @test abs(stressFD - stress[1,2]) < 1e-5

            #        println("done")
        end
    end
end

test_force()

#println("sleep ")
#sleep(3)

test_stress()


using Raylib
using Raylib: RayCamera3D, rayvector,
    RAYWHITE, MAROON, RED, DARKGRAY

function main()
    screenWidth = 800
    screenHeight = 450

    Raylib.InitWindow(screenWidth, screenHeight,
                      "raylib [core] example - 3d camera mode")

    camera = RayCamera3D(
        rayvector(0, 10, 10),
        rayvector(0, 0, 0),
        rayvector(0, 1, 0),
        45,
        0,
    )

    cubePosition = rayvector(0, 0, 0)

    Raylib.SetTargetFPS(60)

    while !Raylib.WindowShouldClose()
        Raylib.BeginDrawing()
        Raylib.ClearBackground(RAYWHITE)

        Raylib.BeginMode3D(camera)
        Raylib.DrawCube(cubePosition, 2f0, 2f0, 2f0, RED)

        Raylib.DrawCubeWires(cubePosition, 2f0, 2f0, 2f0, MAROON)

        Raylib.DrawGrid(10, 1f0)

        Raylib.EndMode3D()
        Raylib.DrawText("Welcome to the third dimension!",
                        10, 40, 20, DARKGRAY)
        Raylib.EndDrawing()
    end

    Raylib.CloseWindow()

end


using Revise
using AOM
using BrowseTables
##
plot_results(RequiemFusion;remove = ["Black Dahlia", "uncleb"])
##
res = update_results(RequiemFusion)
lw = lastweek(res)
open_html_table(lw)
pub = adjust_to_show(lw)
open_html_table(pub)
##
using DataFrames
filt=filter(row -> row.name != "Mazsola",lw)
##
range(colorant"red", stop=colorant"green", length=15)
p = AOM.plot_results(filt)
savefig(p,"/Users/dariosarra/Documents/House/AOM/Stats/03-May-2020.png")
##
plot_results(BoethiasChampions)
res = update_results(BoethiasChampions)
lw = lastweek(res)
pub = adjust_to_show(lw)
open_html_table(pub)
p = AOM.plot_results(lw)
##
using Plots
pow = dropmissing(res)
gd = groupby(pow,:date)
p = combine(gd,:total_power=>sum=>:clan_power)
scatter(p[:,:date],p[:,:clan_power])


#shared image alignment computational procedure
#1.  Compute transform to align camera2's flipped padded half-image (step 8) with its flipped full image (step 7) (should be almost aligned already)
#2.  Compute transform to align camera2's flipped full image (step 7) with camera1's padded half-image (step 8)
#3.  Compose transform from 2 with transform from 1 and call recenter() on the composed transform.
#4.  For each camera2 half-image:
    # -flip y
    # -pad vertically to full chip size
    # -apply transform from #3
    # -nansum with camera1's corresponding padded half-image
    # -find minimal bounding box that contains all non-nan pixels and return image cropped to that region

function stitch(img_top_fixed::AbstractArray{T,2}, img_bottom_moving::AbstractArray{T,2}, tfm, out_type::DataType; kwargs...) where {T}
    if T!=out_type
        img_top_fixed = out_type.(img_top_fixed)
        img_bottom_moving = out_type.(img_bottom_moving)
    end
    return stitch(img_top_fixed, img_bottom_moving, tfm; kwargs...)
end

#Stitches img_top_fixed and img_bottom_moving using a transformation tfm that, when applied to img_bottom_moving, aligns it with img_top_fixed.
#Since img_top_fixed and img_bottom_moving are expected to have little-to-no overlap they are padded before applying the transformation.
#The transformed img_bottom_moving is added to img_top_fixed and the interesting (non-NaN) region of the summed image is returned.
#For most accurate stitching you must provide ysz_full, the size in the y dimension of full images used to compute tfm with the stitch_tfm function
function stitch(img_top_fixed::AbstractArray{T,2}, img_bottom_moving::AbstractArray{T,2}, tfm; flip_y_bottom = true, ysz_full=2048) where {T}
    yroi = size(img_top_fixed,2)
    @assert size(img_bottom_moving,2) == yroi
    yextra = ceil(Int, abs(tfm.translation[2]))
    #fillvals should be NaN, but currently interpolations doesn't handle it well (imputes NaNs at gridpoints)
    img_bottom_moving_pre = ypad(img_bottom_moving, ysz_full; pad_side=:both, flip_y=flip_y_bottom, fillval=0.0)
    img_top_fixed = ypad(img_top_fixed, ysz_full; pad_side=:both, flip_y=false, fillval=0.0)
    img_bottom_moving = warp_and_resample(img_bottom_moving_pre, tfm)
    imgsummed = nanplus(img_top_fixed,img_bottom_moving)
    ystart = div(ysz_full,2) - div(yroi,2) + 1
    return view(imgsummed, :, ystart:ystart+2*yroi-1)
end

#returns a transform that stitches a pair of corresponding shared images into a single image when passed to the stitch function
#Assumes that moving_roi contains the bottom half of the image and that it needs to be flipped vertically to align with fixed_roi.
function stitch_tfm(moving_full, fixed_roi, moving_roi; trunc_frac=0.1, kwargs...)
    @assert size(fixed_roi) == size(moving_roi)

    #bright bead clumps can bias calculation, so truncate them at this value:
    trunc_thr = trunc_frac * maximum_finite(fixed_roi)
    moving_roi = trunc_above.(moving_roi, trunc_thr)
    fixed_roi = trunc_above.(fixed_roi, trunc_thr)
    moving_full = trunc_above.(moving_full, trunc_thr)

    #flip and pad
    moving_full = flipy(moving_full)
    fixed_roi_pad = ypad(fixed_roi, size(moving_full,2); flip_y = false, fillval=NaN)
    moving_roi_pad = ypad(moving_roi, size(moving_full,2); flip_y = true, fillval=NaN)

    #################### Moving strip to moving full
    print("Registering the moving strip to the moving full image\n")
    mxshift_moving = (50, size(moving_roi,2)+50)
    m_roi_pwr = sum(abs2.(moving_roi_pad[.!(isnan.(moving_roi_pad))]))
    moving_pre_rigid, pre_rigid_mm = qd_rigid(moving_full, moving_roi_pad, mxshift_moving, [pi/60]; thresh=0.85*m_roi_pwr, minwidth_rot=[pi/10000], kwargs...)
    @show moving_pre_rigid

    #refine it further, allowing full affine
    moving_pre_tfm, pre_mm = qd_affine(moving_full, moving_roi_pad, (50,50); thresh=0.5*m_roi_pwr, initial_tfm=moving_pre_rigid, kwargs...)
    @show moving_pre_tfm

    #################### Moving strip to moving full
    print("Registering the moving full image to the fixed strip\n")
    mxshift_xcam = (50, round(Int, size(moving_roi, 2)+50))
    f_roi_pwr = sum(abs2.(fixed_roi_pad[.!(isnan.(fixed_roi_pad))]))
    moving_post_rigid, post_rigid_mm = qd_rigid(fixed_roi_pad, moving_full, mxshift_xcam, [pi/60]; thresh=0.5*f_roi_pwr, minwidth_rot=[pi/10000], kwargs...)
    @show moving_post_rigid

    #refine it further, allowing full affine
    print("Found rigid transform, now optimizing allowing general affine transforms\n")
    moving_post_tfm, post_mm = qd_affine(fixed_roi_pad, moving_full, (50,50); thresh=0.5*f_roi_pwr, initial_tfm=moving_post_rigid, kwargs...)
    @show moving_post_tfm

    return recenter(moving_post_tfm ∘ moving_pre_tfm, RegisterQD.center(moving_full)), post_mm
end

#green and magenta overlay
mg_overlay(img1, img2) = colorview(RGB, nanz.(img1), nanz.(img2), nanz.(img1))
#green and magenta overlay, one color per half-image
function stitched_mg_overlay(strip1::AbstractArray{T,2}, strip2::AbstractArray{T,2}, tfm, ysz_full; flip_y2=true) where {T}
    strip0 = zeros(T, size(strip1))
    simg1 = stitch(strip1, strip0, tfm; flip_y_bottom = false, ysz_full=ysz_full)
    simg2 = stitch(strip0, strip2, tfm; flip_y_bottom = flip_y2, ysz_full=ysz_full)
    return mg_overlay(simg1,simg2)
end

#By convention, we align images to camera1, which has the top portion of the split image on OCPI2
#This returns an initial guess translation mapping camera 2's image to camera1's image after camera2's image has been flipped (flip required on OCPI2)
#ysz_roi is the pixel size of the image in the vertical direction (the dimension that gets split between cameras)
#Note: currently not using this
function initial_guess_tfm(ysz_chip::Int, ysz_roi::Int)
    @assert iseven(ysz_chip)
    @assert iseven(ysz_roi)
    if ysz_roi > ysz_chip
        error("ROI size cannot be larger than chip size")
    end
    return Translation(0, -ysz_roi)
end

#Makes an initial guess based on stripsz
#function full2full(fixed, moving, stripsz::Int, mxshift, mxrot, rotstep; thresh = 0.1*sum(abs2.(moving[.!(isnan.(moving))])), tfm0 = IdentityTransformation(), kwargs...)
#    @assert size(fixed) == size(moving)
#    @show tfm_guess = initial_guess_tfm(size(fixed,2), stripsz) ∘ tfm0
#    #tfm, mm = register_padmatched(fixed, moving, mxshift, mxrot, rotstep; thresh=thresh, tfm0=tfm_guess, kwargs...)
#    tfm, mm = register_padmatched(fixed, moving, mxshift, mxrot, rotstep; thresh=thresh, kwargs...)
#end
#
#function register_padmatched(fixed::AbstractMatrix, moving::AbstractMatrix, mxshift, mxrot, rotstep; thresh = 0.1*sum(abs2.(moving[.!(isnan.(moving))])), tfm0 = IdentityTransformation(), kwargs...)
#    if size(fixed,1) != size(moving, 1)
#        error("Images must be equally sized in the first dimension")
#    end
#    fixed, moving = match_ypadding(fixed, moving)
#    tfm, mm = qd_rigid(copy(fixed), copy(moving), mxshift, [mxrot...], [rotstep...]; thresh=thresh, tfm0=tfm0, kwargs...)
#end
#
#function match_ypadding(img1, img2; padval=NaN)
#    if size(img1,2) > size(img2, 2)
#        img2 = ypad(img2, size(img1,2); flip_y = false, fillval=padval)
#    end
#    if size(img1,2) < size(img2, 2)
#        img1 = ypad(img1, size(img2,2); flip_y = false, fillval=padval)
#    end
#    return img1, img2
#end


"""
    Parameter(value::Float64, vary::Bool, prior::Distrbutions.UnivariateDistribution)

A parameter that can be optimised in the analysis procress.

### Parameters
- `value::Float64` or `Distributions.UnivariateDistribution` : the value for the given parameter.
"""
mutable struct Parameter
    value
end

"""
    make_layer
"""
function make_layer(thick, sld, isld, rough)
    return Layer(Parameter(thick), Parameter(sld), Parameter(isld), Parameter(rough))
end

"""
    Layer(thick::Jeff.Parameter, sld::Jeff.Parameter, isld::Jeff.Parameter, rough::Jeff.Parameter)

A description of a layer in a system.

### Parameters
- `thick::Jeff.Parameter` : the layer thickness.
- `sld::Jeff.Parameter` : layer real scattering length density.
- `isld::Jeff.Parameter` : layer imaginary scattering length density.
- `rough::Jeff.Parameter` : roughness with layer above.
"""
mutable struct Layer
    thick::Jeff.Parameter
    sld::Jeff.Parameter
    isld::Jeff.Parameter
    rough::Jeff.Parameter
end

"""
    Model(scale::Jeff.Parmeter, bkg::Jeff.Parameter, layers::Array{Jeff.Layer})

The model from which the reflectometry should be calculated.

### Parameters
- `scale::Jeff.Parameter` : scale factor for reflectometry.
- `bkg::Jeff.Parameter` : uniform background.
- `layers::Array{Jeff.Layer}` : the array of [`Jeff.Layer`] objects that describe the model.
"""
mutable struct Model
    scale::Jeff.Parameter
    bkg::Jeff.Parameter
    layers::Array{Jeff.Layer}
end

"""
    layers_to_array(layers::Array{Jeff.Layer})

Convert from an array of N Jeff.Layer objects to an Nx4 array.

### Parameters
- `layers::Array{Jeff.Layer}`: layers to be converted.

### Returns
- `::Array{Any, 2}` : an array describing the layers.
"""
function layers_to_array(layers::Array{Jeff.Layer})
    nlayers = size(layers, 1)
    array = Array{Any}(undef, nlayers, 4)
    for i in 1:nlayers
        array[i, 1] = layers[i].thick.value
        array[i, 2] = layers[i].sld.value
        array[i, 3] = layers[i].isld.value
        array[i, 4] = layers[i].rough.value
    end
    return array
end


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control

"""
    module ModiaMath.SimulationEngine

Simulation engine for implicit index 1 DAE models with events.

# Main developer

[Martin Otter](https://rmc.dlr.de/sr/de/staff/martin.otter/),
[DLR - Institute of System Dynamics and Control](https://www.dlr.de/sr/en)
"""
module SimulationEngine

export simulate!

import Sundials
import ModiaMath
import ModiaMath.DAE
using  StaticArrays

@static if VERSION >= v"0.7.0-DEV.2005"
    @eval using Printf
    using LinearAlgebra
end

# include code
# include("IDA_UserData.jl")  # is included within simulate.jl
include("simulate.jl")

end


### measurements ###

abstract type OutputMeasurementType end

"""
    FockDetection(s::ModeOccupation)

Measuring the probability of getting the [`ModeOccupation`](@ref) `s` at the output.

    Fields:
        - s::ModeOccupation
"""
struct FockDetection <: OutputMeasurementType
    s::ModeOccupation
    FockDetection(s::ModeOccupation) = at_most_one_photon_per_bin(s) ? new(s) : error("more than one detector per more")
end

"""
    PartitionCount(part_occupancy::PartitionOccupancy)

Measuring the probability of getting a specific count for a given partition `part_occupancy`.

    Fields:
        - part_occupancy::PartitionOccupancy
"""
struct PartitionCount <: OutputMeasurementType
    part_occupancy::PartitionOccupancy
    PartitionCount(part_occupancy::PartitionOccupancy) = new(part_occupancy)
end

"""
    PartitionCountsAll(part::Partition)

Measuring all possible counts probabilities in the partition `part`.

    Fields:
        - part::Partition
"""
struct PartitionCountsAll <: OutputMeasurementType
    part::Partition
    PartitionCountsAll(part::Partition) = new(part)
end

struct OutputMeasurement{T<:OutputMeasurementType}

    # as in most of boson sampling literature, this is for detectors blind to
    # internal degrees of freedom

    s::Union{ModeOccupation,Nothing} # position of the detectors for Fock measurement

    # function OutputMeasurement{T}(s::ModeOccupation) where {T<:OutputMeasurementType}
    #     if T == FockDetection
    #         return at_most_one_photon_per_bin(s) ? new(s, nothing) : error("more than one detector per more")
    #     else
    #         return error(T, " not implemented")
    #     end
    # end

    function OutputMeasurement{FockDetection}(s::ModeOccupation)

        @warn "OutputMeasurement{FockDetection} obsolete, replace with FockDetection"

        at_most_one_photon_per_bin(s) ? new(s) : error("more than one detector per more")

    end
    OutputMeasurement(s::ModeOccupation) = OutputMeasurement{FockDetection}(s::ModeOccupation)

end


using Pkg
pkg"activate ."
using Revise
using Layered
using Colors

function example()
    c, l = canvas(4.5, 3)
    n = 5

    defaultfont("Helvetica Neue Light")

    sls = layer!.((r, i) -> begin
        margin = 20
        avail_w = (r.width - 2margin) / n
        avail_h = (r.height - 2margin) / n
        Transform(translation = topleft(r) + P(margin + (i-0.5) * avail_w, margin + (i-0.5) * avail_h))
    end, l, c.rect, 1:n)

    rs = rect!.((r, i) -> begin
        Rect(P(0, 0), 80, 50, deg(0))
    end, sls, c.rect, 0:n-1) .+ Fill(Gray(0.5)) .+ Visible(true)

    crosses = polygon!.(r -> begin
        ncross(r.center, 4, 3, 0.3)
    end, sls, rs) .+ Fill("black") .+ Stroke("transparent")

    circlesleft = circle!.(r -> begin
        Circle(r.center - X(r.width/4), 10)
    end, sls, rs) .+ Linestyle(:dashed) .+
        Stroke.([i == 2 ? "lime" : "black" for i in 1:n])

    circlesright = circle!.(r -> begin
        Circle(r.center + X(r.width/4), 10)
    end, sls, rs) .+ Linestyle(:dashed) .+
        Stroke.([i == 4 ? "tomato" : "black" for i in 1:n])

    focuscircle = circle!(r -> begin
        Circle(topright(r) + P(-35, 35), 30)
    end, l, c.rect) + Linestyle(:dashed)

    textl = layer!(l, Transform())

    phases = txt!.((r, phase) -> begin
        Txt(bottomleft(r) + Y(0), "phase $phase", 10, :l, :t)
    end, textl, rs, 1:5)

    path!((c1, c2) -> begin
        Path(false, outertangents(c1, c2)...)
    end, l, focuscircle, circlesright[4]) + Opacity(0.3)#+ Linestyle(:dashed)

    arr = path!(l, rs[1], rs[end]) do r1, r2
        a = arrow(topright(r1), topright(r2), 9, 9, 1, 1, 0)
        a + normal(a.segments[1], -5)
    end + Fill("tomato") + Stroke("transparent")

    te = txt!(textl, arr) do arr
        pos = normfrom(arr.segments[1], 0.5, -4)
        Txt(pos, "time", 10, :c, :b, angle(arr.segments[1]))
    end

    r = rect_pre!(textl, te, te) do t
        Rect(t, 4)
    end + Fill("tomato") + Stroke("transparent")

    txt!(textl, focuscircle) do fc
        Txt(fc.center, "focus circle", 10, :c, :c)
    end

    txt!(textl, c.rect) do r
        Txt(bottomleft(r) + P(5, -5), "experimental paradigm", 12, :l, :b)
    end

    c

    # draw_svg(c, "example.svg")

end; example()


@testset "Differential evolution optimizer" begin

ss = symmetric_search_space(1, (0.0, 10.0))
fake_problem = FunctionBasedProblem(x -> 0.0, "test_problem", MinimizingFitnessScheme, ss)
DE = de_rand_1_bin(fake_problem, ParamsDict(
    :Population => reshape(collect(1.0:10.0), (1, 10)),
    :f => 0.4, :cr => 0.5, :NumParents => 3))

@testset "ask()/tell!()" begin
    for i in 1:NumTestRepetitions
        res = BlackBoxOptim.ask(DE)
        @test BlackBoxOptim.candi_pool_size(BlackBoxOptim.population(DE)) == 0 # no candidates in the pool as we just exhausted it
        @test length(res) == 2

        trial, target = res

        @test ndims(trial.params) == 1
        @test (1 <= trial.index <= popsize(DE.population))
        @test in(trial.params, DE.embed.searchSpace)

        @test ndims(target.params) == 1
        @test (1 <= target.index <= popsize(DE.population))
        @test in(target.params, DE.embed.searchSpace)

        @test trial.index == target.index

        BlackBoxOptim.tell!(DE, res)
        @test BlackBoxOptim.candi_pool_size(BlackBoxOptim.population(DE)) == 2 # test that all candidates returned to the pool
    end
end

end


using .Zygote: Zygote, @adjoint, @nograd, pullback

@adjoint istraining() = true, _ -> nothing
@nograd Bijectors._debug

@adjoint function mapvcat(f, args...)
    g(f, args...) = map(f, args...)
    return pullback(g, f, args...)
end
@adjoint function eachcolmaphcat(f, x1, x2)
    function g(f, x1, x2)
        init = reshape(f(view(x1, :, 1), x2[1]), :, 1)
        return reduce(hcat, [f(view(x1, :, i), x2[i]) for i in 2:size(x1, 2)]; init = init)
    end
    return pullback(g, f, x1, x2)
end
@adjoint function eachcolmaphcat(f, x)
    function g(f, x)
        init = reshape(f(view(x, :, 1)), :, 1)
        return reduce(hcat, [f(view(x, :, i)) for i in 2:size(x, 2)]; init = init)
    end
    return pullback(g, f, x)
end
@adjoint function _sum(f, args...)
    g(f, args...) = sum(map(f, args...))
    return pullback(g, f, args...)
end
@adjoint function _sumeachcol(f, x1, x2)
    g(f, x1, x2) = sum([f(view(x1, :, i), x2[i]) for i in 1:size(x1, 2)])
    return pullback(g, f, x1, x2)
end

@adjoint function logabsdetjac(b::Log{1}, x::AbstractVector)
    return -sum(log, x), Δ -> (nothing, -Δ ./ x)
end
@adjoint function logabsdetjac(b::Log{1}, x::AbstractMatrix)
    return -vec(sum(log, x; dims = 1)), Δ -> (nothing, .- Δ' ./ x)
end

# AD implementations
function jacobian(
    b::Union{<:ADBijector{<:ZygoteAD}, Inverse{<:ADBijector{<:ZygoteAD}}},
    x::Real
)
    return Zygote.gradient(b, x)[1]
end
function jacobian(
    b::Union{<:ADBijector{<:ZygoteAD}, Inverse{<:ADBijector{<:ZygoteAD}}},
    x::AbstractVector{<:Real}
)
    return Zygote.jacobian(b, x)
end
@adjoint function _logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return _logabsdetjac_scale(a, x, Val(0)), Δ -> (inv(a) .* Δ, nothing, nothing)
end
@adjoint function _logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    J = fill(inv.(a), length(x))
    return _logabsdetjac_scale(a, x, Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
@adjoint function _logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    J = fill(size(x, 1) / a, size(x, 2))
    return _logabsdetjac_scale(a, x, Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
@adjoint function _logabsdetjac_scale(a::AbstractVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    J = inv.(a)
    return _logabsdetjac_scale(a, x, Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
@adjoint function _logabsdetjac_scale(a::AbstractVector, x::AbstractMatrix, ::Val{1})
    Jᵀ = repeat(inv.(a), 1, size(x, 2))
    return _logabsdetjac_scale(a, x, Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end

## Positive definite matrices

@adjoint function replace_diag(::typeof(log), X)
    f(i, j) = i == j ? log(X[i, j]) : X[i, j]
    out = f.(1:size(X, 1), (1:size(X, 2))')
    out, ∇ -> begin
        g(i, j) = i == j ? ∇[i, j] / X[i, j] : ∇[i, j]
        (nothing, g.(1:size(X, 1), (1:size(X, 2))'))
    end
end
@adjoint function replace_diag(::typeof(exp), X)
    f(i, j) = ifelse(i == j, exp(X[i, j]), X[i, j])
    out = f.(1:size(X, 1), (1:size(X, 2))')
    out, ∇ -> begin
        g(i, j) = ifelse(i == j, ∇[i, j] * exp(X[i, j]), ∇[i, j])
        (nothing, g.(1:size(X, 1), (1:size(X, 2))'))
    end
end

@adjoint function _logpdf_with_trans_pd(
    d,
    X::AbstractMatrix{<:Real},
    transform::Bool,
)
    return pullback(_logpdf_with_trans_pd_zygote, d, X, transform)
end
function _logpdf_with_trans_pd_zygote(
    d,
    X::AbstractMatrix{<:Real},
    transform::Bool,
)
    T = eltype(X)
    Xcf = cholesky(X, check = false)
    if !issuccess(Xcf)
        Xcf = cholesky(X + max(eps(T), eps(T) * norm(X)) * I, check = true)
    end
    lp = getlogp(d, Xcf, X)
    if transform && isfinite(lp)
        U = Xcf.U
        @inbounds for i in 1:dim(d)
            lp += (dim(d) - i + 2) * log(U[i, i])
        end
        lp += dim(d) * log(T(2))
    end
    return lp
end

# Simplex adjoints

@adjoint function _simplex_bijector(X::AbstractVector, b::SimplexBijector)
    return _simplex_bijector(X, b), Δ -> (simplex_link_jacobian(X)' * Δ, nothing)
end
@adjoint function _simplex_inv_bijector(Y::AbstractVector, b::SimplexBijector)
    return _simplex_inv_bijector(Y, b), Δ -> (simplex_invlink_jacobian(Y)' * Δ, nothing)
end

@adjoint function _simplex_bijector(X::AbstractMatrix, b::SimplexBijector)
    return _simplex_bijector(X, b), Δ -> begin
        maphcat(eachcol(X), eachcol(Δ)) do c1, c2
            simplex_link_jacobian(c1)' * c2
        end, nothing
    end
end
@adjoint function _simplex_inv_bijector(Y::AbstractMatrix, b::SimplexBijector)
    return _simplex_inv_bijector(Y, b), Δ -> begin
        maphcat(eachcol(Y), eachcol(Δ)) do c1, c2
            simplex_invlink_jacobian(c1)' * c2
        end, nothing
    end
end

@adjoint function logabsdetjac(b::SimplexBijector, x::AbstractVector)
    return logabsdetjac(b, x), Δ -> begin
        (nothing, simplex_logabsdetjac_gradient(x) * Δ)
    end
end
@adjoint function logabsdetjac(b::SimplexBijector, x::AbstractMatrix)
    return logabsdetjac(b, x), Δ -> begin
        (nothing, maphcat(eachcol(x), Δ) do c, g
            simplex_logabsdetjac_gradient(c) * g
        end)
    end
end

# LocationScale fix

@adjoint function _clamp(x::T, dist::LocationScale) where {T <: Real}
    function f(x, dist)
        ϵ = _eps(T)
        bounds = (minimum(dist) + ϵ, maximum(dist) - ϵ)
        if x < bounds[1]
            clamped_x = bounds[1]
        elseif x > bounds[2]
            clamped_x = bounds[2]
        else
            clamped_x = x
        end
        DEBUG && _debug("x = $x, bounds = $bounds, clamped_x = $clamped_x")
        return clamped_x
    end
    return pullback(f, x, dist)
end