# BigWig Statistics
# =================

"""
    coverage(reader, chrom, chromstart, chromend; usezoom=false)::Int

Compute the coverage of values in `[chromstart, chromend]` of `chrom`.

If `usezoom` is `true`, this function tries to use precomputed statistics (zoom) in the file.
This is often faster but not exact in most cases.
"""
function coverage(reader::Reader, chrom::AbstractString, chromstart::Integer, chromend::Integer; usezoom=false)::Int
    chromid = reader.chroms[chrom][1]
    chromstart -= 1
    if usezoom
        zoom = BBI.find_best_zoom(reader.zooms, UInt32(chromend - chromstart))
        if zoom !== nothing
            return BBI.coverage(zoom, chromid, UInt32(chromstart), UInt32(chromend))
        end
    end
    return exact_coverage(reader, chromid, UInt32(chromstart), UInt32(chromend))
end

function exact_coverage(reader::Reader, chromid::UInt32, chromstart::UInt32, chromend::UInt32)
    cov::Int = 0
    for record in OverlapIterator(reader, chromid, chromstart, chromend)
        cov += BBI.coverage2((record.chromstart, record.chromend), (chromstart, chromend))
    end
    return cov
end

"""
    mean(reader, chrom, chromstart, chromend; usezoom=false)::Float32

Compute the mean of values in `[chromstart, chromend]` of `chrom`.

This function returns `NaN32` if there are no data in that range.

See `coverage` for the `usezoom` keyword argument.
"""
function mean(reader::Reader, chrom::AbstractString, chromstart::Integer, chromend::Integer; usezoom=false)::Float32
    chromid = reader.chroms[chrom][1]
    chromstart -= 1
    if usezoom
        zoom = BBI.find_best_zoom(reader.zooms, UInt32(chromend - chromstart))
        if zoom !== nothing
            return BBI.mean(zoom, chromid, UInt32(chromstart), UInt32(chromend))
        end
    end
    return exact_mean(reader, chromid, UInt32(chromstart), UInt32(chromend))
end

function exact_mean(reader::Reader, chromid::UInt32, chromstart::UInt32, chromend::UInt32)
    # compute size-weighted mean
    sum = 0.0f0
    size = 0
    for record in OverlapIterator(reader, chromid, chromstart, chromend)
        cov = BBI.coverage2((record.chromstart, record.chromend), (chromstart, chromend))
        sum += record.value * cov
        size += cov
    end
    return sum / size
end

"""
    minimum(reader, chrom, chromstart, chromend; usezoom=false)::Float32

Compute the minimum of values in `[chromstart, chromend]` of `chrom`.

This function returns `NaN32` if there are no data in that range.

See `coverage` for the `usezoom` keyword argument.
"""
function minimum(reader::Reader, chrom::AbstractString, chromstart::Integer, chromend::Integer; usezoom=false)::Float32
    chromid = reader.chroms[chrom][1]
    chromstart -= 1
    if usezoom
        zoom = BBI.find_best_zoom(reader.zooms, UInt32(chromend - chromstart))
        if zoom !== nothing
            return BBI.minimum(zoom, chromid, UInt32(chromstart), UInt32(chromend))
        end
    end
    return exact_extrema(reader, chromid, UInt32(chromstart), UInt32(chromend))[1]
end

"""
    maximum(reader, chrom, chromstart, chromend; usezoom=false)::Float32

Compute the maximum of values in `[chromstart, chromend]` of `chrom`.

This function returns `NaN32` if there are no data in that range.

See `coverage` for the `usezoom` keyword argument.
"""
function maximum(reader::Reader, chrom::AbstractString, chromstart::Integer, chromend::Integer; usezoom=false)::Float32
    chromid = reader.chroms[chrom][1]
    chromstart -= 1
    if usezoom
        zoom = BBI.find_best_zoom(reader.zooms, UInt32(chromend - chromstart))
        if zoom !== nothing
            return BBI.maximum(zoom, chromid, UInt32(chromstart), UInt32(chromend))
        end
    end
    return exact_extrema(reader, chromid, UInt32(chromstart), UInt32(chromend))[2]
end

function exact_extrema(reader::Reader, chromid::UInt32, chromstart::UInt32, chromend::UInt32)
    min = +Inf32
    max = -Inf32
    defined = false
    for record in OverlapIterator(reader, chromid, chromstart, chromend)
        min = Base.min(min, record.value)
        max = Base.max(max, record.value)
        defined = true
    end
    return defined ? (min, max) : (NaN32, NaN32)
end

"""
    std(reader, chrom, chromstart, chromend; usezoom=false)::Float32

Compute the standard deviation of values in `[chromstart, chromend]` of `chrom`.

See `coverage` for the `usezoom` keyword argument.
"""
function std(reader::Reader, chrom::AbstractString, chromstart::Integer, chromend::Integer; usezoom=false)::Float32
    chromid = reader.chroms[chrom][1]
    chromstart -= 1
    if usezoom
        zoom = BBI.find_best_zoom(reader.zooms, UInt32(chromend - chromstart))
        if zoom !== nothing
            return BBI.std(zoom, chromid, UInt32(chromstart), UInt32(chromend))
        end
    end
    return exact_std(reader, chromid, UInt32(chromstart), UInt32(chromend))
end

function exact_std(reader::Reader, chromid::UInt32, chromstart::UInt32, chromend::UInt32)
    sum = 0.0f0
    ssq = 0.0f0
    size = 0
    for record in OverlapIterator(reader, chromid, chromstart, chromend)
        cov = BBI.coverage2((record.chromstart, record.chromend), (chromstart, chromend))
        sum += record.value * cov
        ssq += record.value^2 * cov
        size += cov
    end
    return sqrt((ssq - sum^2 / size) / (size - 1))
end


# series of deconvolution tools

export selectFactorType
export approxDeconv, deconvSolveKey
export approxDeconvBelief

## Initial version of selecting the dimension of a factor -- will be consolidated with existing infrastructure later

"""
    $SIGNATURES

Inverse solve of predicted noise value and returns tuple of (newly calculated-predicted, and known measurements) values.

Notes
- Only works for first value in measurement::Tuple at this stage.
- "measured" is used as starting point for the "calculated-predicted" values solve.
- Not all factor evaluation cases are support yet.
- NOTE only works on `.threadid()==1` at present, see #1094
- This function is still part of the initial implementation and needs a lot of generalization improvements.

DevNotes
- TODO Test for various cases with multiple variables.
- TODO make multithread-safe, and able, see #1094
- TODO Test for cases with `nullhypo`
- FIXME FactorMetadata object for all use-cases, not just empty object.
- TODO resolve #1096 (multihypo)
  - TODO Test cases for `multihypo`.
- TODO figure out if there is a way to consolidate with evalFactor and approxConv?
  - basically how to do deconv for just one sample with unique values (wrt TAF)
- TODO N should not be hardcoded to 100

Related

[`approxDeconv`](@ref), [`_solveCCWNumeric!`](@ref)
"""
function approxDeconv(fcto::DFGFactor,
                      ccw::CommonConvWrapper = _getCCW(fcto);
                      N::Int=100,
                      measurement::AbstractVector=sampleFactor(ccw, N),
                      retries::Int=3  )
  #
  # but what if this is a partial factor -- is that important for general cases in deconv?
  _setCCWDecisionDimsConv!(ccw)

  # FIXME This does not incorporate multihypo??
  varsyms = getVariableOrder(fcto)
  # vars = getPoints.(getBelief.(dfg, varsyms, solveKey) )
  fcttype = getFactorType(fcto)

  # get measurement dimension
  zDim = _getZDim(fcto)
  # TODO consider using ccw.cpt[thrid].res # likely needs resizing
  res_ = zeros(zDim)
  # TODO, consolidate fmd with getSample/sampleFactor and _buildLambda
  fctSmpls = deepcopy(measurement)
  fmd = _getFMdThread(ccw)

  # TODO assuming vector on only first container in measurement::Tuple
  makeTarget = (i) -> measurement[i] # TODO does not support copy-primitive types like Float64, only Ref()
  # makeTarget = (i) -> view(measurement[1][i],:)
  # makeTarget = (i) -> view(measurement[1], :, i)

  # NOTE
  # build a lambda that incorporates the multihypo selections
  # set these first
  # ccw.cpt[].activehypo / .p / .params  # params should already be set from construction
  hyporecipe = _prepareHypoRecipe!(nothing, N, 0, length(varsyms))
  # Juila 1.7 allows destructure assign `(;a,b) = namedtype`
  # certainidx, allelements, activehypo, mhidx =
  # only doing the current active hypo
  @assert hyporecipe.activehypo[2][1] == 1 "deconv was expecting hypothesis nr == (1, 1:d)"

  islen1 = zDim == 1

  for idx in 1:N
    # towards each particle in their own thread (not 100% ready yet, factors should be separate memory)
    thrid = Threads.threadid()
    cpt_ = ccw.cpt[thrid]
    targeti_ = makeTarget(idx)

    # TODO must first resolve hypothesis selection before unrolling them -- deferred #1096
    cpt_.activehypo = hyporecipe.activehypo[2][2]

    onehypo!, _ = _buildCalcFactorLambdaSample( ccw,
                                                idx,
                                                cpt_,
                                                targeti_,
                                                measurement  )
    #

    # lambda with which to find best measurement values
    hypoObj = (tgt) -> (targeti_.=tgt; onehypo!() )

    # find solution via SubArray view pointing to original memory location
    if fcttype isa AbstractManifoldMinimize
      sfidx = ccw.varidx
      targeti_ .= _solveLambdaNumericMeas(fcttype, hypoObj, res_, measurement[idx], ccw.vartypes[sfidx](), islen1)
    else
      targeti_ .= _solveLambdaNumeric(fcttype, hypoObj, res_, measurement[idx], islen1)
    end

  end

  # return (deconv-prediction-result, independent-measurement)
  # r_meas = map(m->m[1], measurement)
  # r_fctSmpls = map(m->m[1], fctSmpls)
  return measurement, fctSmpls
end

"""
    $SIGNATURES

Generalized deconvolution to find the predicted measurement values of the factor `fctsym` in `dfg`.
Inverse solve of predicted noise value and returns tuple of (newly predicted, and known "measured" noise) values.

Notes
- Opposite operation contained in `approxConv`.
- For more notes see [`solveFactorMeasurements`](@ref).

Related

[`approxConv`](@ref), `deconvSolveKey`
"""
function approxDeconv(dfg::AbstractDFG,
                      fctsym::Symbol,
                      solveKey::Symbol=:default;
                      retries::Int=3  )
  #

  # which factor
  fct = getFactor(dfg, fctsym)
  pts = getPoints(getBelief(dfg,getVariableOrder(fct)[1], solveKey))
  N = length(pts)
  pts = approxDeconv(fct, N=N, retries=retries )
  return pts
end

function approxDeconv(dfg::AbstractDFG,
                      fctlbl::Symbol,
                      factorType::AbstractRelative,
                      solveKey::Symbol=:default;
                      tfg::AbstractDFG=initfg(),
                      retries::Int=3 )
  #

  # build a local temporary graph copy containing the same values but user requested factor type.
  fct = getFactor(dfg, fctlbl)
  fctT = getFactorType(fct)
  lbls = getVariableOrder(fct)
  for lb in lbls
    exists(tfg, lb) ? nothing : addVariable!(tfg, lb, getVariableType(dfg, lb))
    initManual!(tfg, lb, getBelief(dfg, lb, solveKey))
  end

  # add factor type requested by user
  f_ = addFactor!(tfg, lbls, factorType, graphinit=false)

  # peform the deconvolution operation on the temporary graph with user desired factor instead.
  approxDeconv(tfg, getLabel(f_); retries=retries)
end

# try default constructor
approxDeconv( dfg::AbstractDFG,
              fctlbl::Symbol,
              factorType::Type{<:AbstractRelative},
              w...;
              kw... ) = approxDeconv(dfg, fctlbl, factorType(), w...; kw...)
#

approxDeconvBelief(dfg::AbstractDFG, lb::Symbol, w...;kw...) = manikde!(approxDeconv(dfg, lb, w...;kw...)[1], getManifold(getFactorType(dfg, lb)))

"""
    $SIGNATURES

Calculate the relative difference between two variables and across solveKeys.

Example

```julia
fg = generateCanonicalFG_Kaess()
solveTree!(fg, storeOld=true)
# calculate the relative motion induced by the solver from init to solve.
pts = deconvSolveKey(fg, :x1, :default, :x1, :graphinit)
```

Notes
- Can pass user `tfg::AbstractDFG` for better in-place operations.

DevNotes
- TODO use dfg, rather than building new tfg internally.

Related

[`approxDeconv`](@ref), [`mmd`](@ref)
"""
function deconvSolveKey(dfg::AbstractDFG,
                        refSym::Symbol,
                        refKey::Symbol,
                        tstSym::Symbol,
                        tstKey::Symbol;
                        tfg = initfg()  )
  #
  # create a new temporary factor graph for calculations

  # add the first "reference" variable
  Xref = getBelief(dfg, refSym, refKey)
  refSym_ = Symbol(refSym, "_ref")
  refVarType = getVariableType(dfg, refSym)
  if !exists(tfg, refSym_)
    addVariable!(tfg, refSym_, refVarType)
  end
  initManual!(tfg, refSym_, Xref)

  # add the second "test" variable
  tstVarType = getVariableType(dfg, tstSym)
  Xtst = getBelief(dfg, tstSym, tstKey)
  tstSym_ = Symbol(tstSym, "_tst")
  if !exists(tfg, tstSym_)
    addVariable!(tfg, tstSym_, tstVarType)
  end
  initManual!(tfg, tstSym_, Xtst)

  # add the new dummy factor with default manifold for computations
  fctType = selectFactorType(refVarType, tstVarType)
  nf = addFactor!(tfg, [refSym_; tstSym_], fctType())

  # TODO connect from dfg all other data that might form part of FactorMetadata in tfg
  pts = approxDeconv(tfg, nf.label)

  # assuming tfg was passed in by the user
  deleteFactor!(tfg, nf.label)

  # return result
  return pts, fctType
end

#


# In this module, all parameters for the cat state preparation are defined:

module parameters
export parameters_qb

parameters_para=Dict(
    "N" =>2,
    "w" => 20.0f0,
    "force_mag" => 5.0f0,
    "max_episode_steps" => 150,
    "dt" => 0.00025f0,
    "n_substeps" => 80,
    "gamma" => 1.0f0)
end


using HTTP, JSON, JuMP
include("REopt.jl")
include("GHPGHX.jl")

using .REopt
using .GHPGHX

function job(req::HTTP.Request)
    d = JSON.parse(String(req.body))
    timeout = pop!(d, "timeout_seconds")
    tol = pop!(d, "tolerance")
    m = xpress_model(timeout, tol)
    @info "Starting REopt with timeout of $(timeout) seconds..."
	error_response = Dict()
	results = Dict()
	try
    	results = reopt(m, d)
	catch e
		@error "Something went wrong in the Julia code!" exception=(e, catch_backtrace())
		error_response["error"] = sprint(showerror, e)
	end
	finalize(backend(m))
	GC.gc()
	if isempty(error_response)
    	@info "REopt model solved with status $(results["status"])."
    	return HTTP.Response(200, JSON.json(results))
	else
		@info "An error occured in the Julia code."
		return HTTP.Response(500, JSON.json(error_response))
	end
end

function ghpghx(req::HTTP.Request)
    inputs_dict = JSON.parse(String(req.body))
    @info "Starting GHPGHX" #with timeout of $(timeout) seconds..."
    results, inputs_params = GHPGHX.ghp_model(inputs_dict)
    # Create a dictionary of the results data needed for REopt
    ghpghx_results = GHPGHX.get_ghpghx_results_for_reopt(results, inputs_params)
    @info "GHPGHX model solved" #with status $(results["status"])."
    return HTTP.Response(200, JSON.json(ghpghx_results))
end

function health(req::HTTP.Request)
    return HTTP.Response(200, JSON.json(Dict("Julia-api"=>"healthy!")))
end

# define REST endpoints to dispatch to "service" functions
const ROUTER = HTTP.Router()

HTTP.@register(ROUTER, "POST", "/job", job)
HTTP.@register(ROUTER, "POST", "/ghpghx", ghpghx)
HTTP.@register(ROUTER, "GET", "/health", health)
HTTP.serve(ROUTER, "0.0.0.0", 8081, reuseaddr=true)


#!/usr/bin/env julia

using Plots; plotlyjs()

h1=histogram(randn(1000))

gui(h1)

readline()


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

using Flux, DiffEqFlux
using OrdinaryDiffEq
using DiffEqCallbacks
using Interpolations: interpolate, LinearInterpolation

using DiffEqFlux: ODEFunction, basic_tgrad, ODEProblem, ZygoteVJP, InterpolatingAdjoint, solve

import ForwardDiff
import Optim
import ProgressMeter

import SciMLBase: RightRootFind

using FMIImport: fmi2ComponentState, fmi2ComponentStateInstantiated, fmi2ComponentStateInitializationMode, fmi2ComponentStateEventMode, fmi2ComponentStateContinuousTimeMode, fmi2ComponentStateTerminated, fmi2ComponentStateError, fmi2ComponentStateFatal
import ChainRulesCore: ignore_derivatives

using FMIImport: fmi2StatusOK, FMU2Solution, FMU2Event

"""
The mutable struct representing an abstract (simulation mode unknown) NeuralFMU.
"""
abstract type NeuralFMU end

"""
Structure definition for a NeuralFMU, that runs in mode `Model Exchange` (ME).
"""
mutable struct ME_NeuralFMU <: NeuralFMU
    neuralODE::NeuralODE
    fmu::FMU

    currentComponent

    tspan
    saveat
    saved_values
    recordValues

    valueStack

    customCallbacks::Array
    callbacks::Array

    x0::Array{Float64}
    firstRun::Bool
    solution::FMU2Solution

    tolerance::Union{Real, Nothing}
    parameters::Union{Dict{<:Any, <:Any}, Nothing}
    setup::Union{Bool, Nothing}
    reset::Union{Bool, Nothing}
    instantiate::Union{Bool, Nothing}
    freeInstance::Union{Bool, Nothing}
    terminate::Union{Bool, Nothing}

    progressMeter

    #componentShadow::Union{FMU2ComponentShadow, Nothing}

    solveCycle::UInt

    function ME_NeuralFMU()
        inst = new()
        inst.currentComponent = nothing
        inst.progressMeter = nothing

        return inst
    end
end

"""
Structure definition for a NeuralFMU, that runs in mode `Co-Simulation` (CS).
"""
mutable struct CS_NeuralFMU{F, C} <: NeuralFMU
    model
    fmu::F
    currentComponent # ::Union{C, Nothing}

    tspan
    saveat
    solution::FMU2Solution

    function CS_NeuralFMU{F, C}() where {F, C}
        inst = new{F, C}()

        inst.currentComponent = nothing

        return inst
    end
end

##### EVENT HANDLING START

function startCallback(nfmu, t)
    ignore_derivatives() do
        nfmu.solveCycle += 1
        @debug "[$(nfmu.solveCycle)][FIRST STEP]"

        @assert ForwardDiff.value(t) == nfmu.tspan[1] "startCallback(...): Called for non-start-point t=$(ForwardDiff.value(t))"

        # allocate component shadow for second run
        if nfmu.fmu.executionConfig.useComponentShadow

            if nfmu.currentComponent == nothing
                nfmu.currentComponent = FMU2ComponentShadow()
                @debug "[$(nfmu.solveCycle)][INIT SHADOW]"
                push!(nfmu.fmu.components, nfmu.currentComponent)
            end

            # real comp renew
            realComponent = nfmu.currentComponent.realComponent
            realComponent, nfmu.x0 = prepareFMU(nfmu.fmu, realComponent, nfmu.instantiate, nfmu.freeInstance, nfmu.terminate, nfmu.reset, nfmu.setup, nfmu.parameters, nfmu.tspan[1], nfmu.tspan[end], nfmu.tolerance; x0=nfmu.x0, pushComponents=false)
            setComponent!(nfmu.currentComponent, realComponent)

            handleEvents(nfmu.currentComponent)
            #@assert fmi2EnterContinuousTimeMode(nfmu.currentComponent) == fmi2StatusOK

            nfmu.currentComponent.log = (nfmu.solveCycle == 2)

        else
            nfmu.currentComponent, nfmu.x0 = prepareFMU(nfmu.fmu, nfmu.currentComponent, nfmu.instantiate, nfmu.freeInstance, nfmu.terminate, nfmu.reset, nfmu.setup, nfmu.parameters, nfmu.tspan[1], nfmu.tspan[end], nfmu.tolerance; x0=nfmu.x0)

            handleEvents(nfmu.currentComponent)
            #@assert fmi2EnterContinuousTimeMode(nfmu.currentComponent) == fmi2StatusOK
        end

        # re-init every callback except #1 (the startCallback itself)
        # for i in 2:length(nfmu.callbacks)
        #     callback = nfmu.callbacks[i]
        #     callback.initialize(callback, u, t, integrator)
        # end
    end
end

function stopCallback(nfmu, t)
    ignore_derivatives() do
        @debug "[$(nfmu.solveCycle)][LAST STEP]"

        @assert ForwardDiff.value(t) == nfmu.tspan[end] "stopCallback(...): Called for non-start-point t=$(ForwardDiff.value(t))"

        if nfmu.fmu.executionConfig.useComponentShadow

            if nfmu.solveCycle == 2
                prepare!(nfmu.currentComponent)
                @debug "[$(nfmu.solveCycle)][PREPARED SHADOW]"
            end

            #comp = finishFMU(nfmu.fmu, nfmu.currentComponent.realComponent, nfmu.freeInstance; popComponent=false)
            #setComponent!(nfmu.currentComponent, comp)

        else
            nfmu.currentComponent = finishFMU(nfmu.fmu, nfmu.currentComponent, nfmu.freeInstance)
        end

    end
end

# Read next time event from fmu and provide it to the integrator
function time_choice(nfmu, integrator, tStart, tStop)

    # last call may be after simulation end
    if nfmu.currentComponent == nothing
        return nothing
    end

    if !nfmu.fmu.executionConfig.handleTimeEvents
        return nothing
    end

    if nfmu.currentComponent.eventInfo.nextEventTimeDefined == fmi2True
        #@debug "time_choice(...): $(nfmu.currentComponent.eventInfo.nextEventTime) at t=$(ForwardDiff.value(integrator.t))"

        if nfmu.currentComponent.eventInfo.nextEventTime >= tStart && nfmu.currentComponent.eventInfo.nextEventTime <= tStop
            return nfmu.currentComponent.eventInfo.nextEventTime
        else
            # the time event is outside the simulation range!
            @debug "Next time event @$(nfmu.currentComponent.eventInfo.nextEventTime)s is outside simulation time range ($(tStart), $(tStop)), skipping."
            return nothing
        end
    else
        #@debug "time_choice(...): nothing at t=$(ForwardDiff.value(integrator.t))"
        return nothing
    end
end

# Handles events and returns the values and nominals of the changed continuous states.
function handleEvents(c::Union{FMU2Component, FMU2ComponentShadow})

    @assert c.state == fmi2ComponentStateEventMode "handleEvents(...): Must be in event mode!"

    #@debug "Handle Events..."

    # trigger the loop
    c.eventInfo.newDiscreteStatesNeeded = fmi2True

    valuesOfContinuousStatesChanged = fmi2False
    nominalsOfContinuousStatesChanged = fmi2False
    nextEventTimeDefined = fmi2False
    nextEventTime = 0.0

    numCalls = 0
    while c.eventInfo.newDiscreteStatesNeeded == fmi2True
        numCalls += 1
        fmi2NewDiscreteStates!(c, c.eventInfo)

        if c.eventInfo.valuesOfContinuousStatesChanged == fmi2True
            valuesOfContinuousStatesChanged = fmi2True
        end

        if c.eventInfo.nominalsOfContinuousStatesChanged == fmi2True
            nominalsOfContinuousStatesChanged = fmi2True
        end

        if c.eventInfo.nextEventTimeDefined == fmi2True
            nextEventTimeDefined = fmi2True
            nextEventTime = c.eventInfo.nextEventTime
        end

        if c.eventInfo.terminateSimulation == fmi2True
            @error "handleEvents(...): FMU throws `terminateSimulation`!"
        end

        @assert numCalls <= c.fmu.executionConfig.maxNewDiscreteStateCalls "handleEvents(...): `fmi2NewDiscreteStates!` exceeded $(c.fmu.executionConfig.maxNewDiscreteStateCalls) calls, this may be an error in the FMU. If not, you can change the max value for this FMU in `fmu.executionConfig.maxNewDiscreteStateCalls`."
    end

    c.eventInfo.valuesOfContinuousStatesChanged = valuesOfContinuousStatesChanged
    c.eventInfo.nominalsOfContinuousStatesChanged = nominalsOfContinuousStatesChanged
    c.eventInfo.nextEventTimeDefined = nextEventTimeDefined
    c.eventInfo.nextEventTime = nextEventTime

    @assert fmi2EnterContinuousTimeMode(c) == fmi2StatusOK "FMU is not in state continuous time after event handling."

    return nothing
end

# Returns the event indicators for an FMU.
function condition(nfmu::ME_NeuralFMU, out::SubArray{<:ForwardDiff.Dual{T, V, N}, A, B, C, D}, x, t, integrator) where {T, V, N, A, B, C, D} # Event when event_f(u,t) == 0

    @assert nfmu.fmu.components[end].state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time."

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

    if isa(t, ForwardDiff.Dual)
        t = ForwardDiff.value(t)
    end

    if all(isa.(x, ForwardDiff.Dual))
        x = collect(ForwardDiff.value(e) for e in x)
    end

    # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN
    nfmu.neuralODE.model(x) # evaluate NeuralFMU (set new states)
    fmi2SetTime(nfmu.fmu, t)

    buf = zeros(nfmu.fmu.modelDescription.numberOfEventIndicators)
    fmi2GetEventIndicators!(nfmu.fmu.components[end], buf)

    out[:] = collect(ForwardDiff.Dual{T, V, N}(V(buf[i]), ForwardDiff.partials(out[i])) for i in 1:length(out))

    return nothing
end
function condition(nfmu::ME_NeuralFMU, out, x, t, integrator) # Event when event_f(u,t) == 0

    # last call may be after simulation end
    # if nfmu.currentComponent == nothing
    #     return nothing
    # end

    @assert nfmu.fmu.components[end].state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time."

    #@debug "State condition..."

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

    if isa(t, ForwardDiff.Dual)
        t = ForwardDiff.value(t)
    end

    if all(isa.(x, ForwardDiff.Dual))
        x = collect(ForwardDiff.value(e) for e in x)
    end

    # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN
    nfmu.neuralODE.model(x) # evaluate NeuralFMU (set new states)
    fmi2SetTime(nfmu.fmu, t)

    # ToDo: Input Function

    #@debug "Condition..."

    fmi2GetEventIndicators!(nfmu.fmu.components[end], out)

    return nothing
end
function conditionSingle(nfmu::ME_NeuralFMU, index, x, t, integrator)

    # last call may be after simulation end
    # if nfmu.currentComponent == nothing
    #     return 1.0
    # end

    @assert nfmu.fmu.components[end].state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time."

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

    if isa(t, ForwardDiff.Dual)
        t = ForwardDiff.value(t)
    end

    if all(isa.(x, ForwardDiff.Dual))
        x = collect(ForwardDiff.value(e) for e in x)
    end

    #@debug "ConditionSingle..."

    # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN
    nfmu.neuralODE.model(x) # evaluate NeuralFMU (set new states)
    fmi2SetTime(nfmu.fmu, t)

    # ToDo: Input Function

    indicators = zeros(nfmu.fmu.modelDescription.numberOfEventIndicators)
    fmi2GetEventIndicators!(nfmu.fmu, indicators)

    return indicators[index]
end

function f_optim(x, nfmu, right_x_fmu) # , idx, direction::Real)
    # propagete the new state-guess `x` through the NeuralFMU
    nfmu.neuralODE.model(x)
    #indicators = fmi2GetEventIndicators(nfmu.fmu)
    return Flux.Losses.mse(right_x_fmu, fmi2GetContinuousStates(nfmu.fmu)) # - min(-direction*indicators[idx], 0.0)
end

# Handles the upcoming events.
function affectFMU!(nfmu::ME_NeuralFMU, integrator, idx)

    c = nfmu.fmu.components[end]

    @assert c.state == fmi2ComponentStateContinuousTimeMode "affectFMU!(...): Must be in continuous time mode!"

    t = integrator.t
    if isa(t, ForwardDiff.Dual)
        t = ForwardDiff.value(t)
    end

    x = integrator.u
    if all(isa.(x, ForwardDiff.Dual))
        x = collect(ForwardDiff.value(e) for e in x)
    end

    left_x_fmu = fmi2GetContinuousStates(c)
    fmi2SetTime(c, t)
    # Todo set inputs

    # Event found - handle it
    @assert fmi2EnterEventMode(c) == fmi2StatusOK
    handleEvents(c)

    # ignore_derivatives() do
    #     if idx == 0
    #         @debug "affectFMU!(...): Handle time event at t=$t"
    #     end

    #     if idx > 0
    #         @debug "affectFMU!(...): Handle state event at t=$t"
    #     end
    # end

    left_x = nothing
    right_x = nothing

    if c.eventInfo.valuesOfContinuousStatesChanged == fmi2True

        left_x = x

        right_x_fmu = fmi2GetContinuousStates(c) # the new FMU state after handled events

        ignore_derivatives() do
            #@debug "affectFMU!(_, _, $idx): NeuralFMU state event from $(left_x) (fmu: $(left_x_fmu)). Indicator [$idx]: $(indicators[idx]). Optimizing new state ..."
        end

        # ToDo: Problem-related parameterization of optimize-call
        #result = optimize(x_seek -> f_optim(x_seek, nfmu, right_x_fmu), left_x, LBFGS(); autodiff = :forward)
        #result = Optim.optimize(x_seek -> f_optim(x_seek, nfmu, right_x_fmu, idx, sign(indicators[idx])), left_x, NelderMead())
        result = Optim.optimize(x_seek -> f_optim(x_seek, nfmu, right_x_fmu), left_x, NelderMead())

        #display(result)

        right_x = Optim.minimizer(result)
        integrator.u = right_x

    end

    if c.eventInfo.nominalsOfContinuousStatesChanged == fmi2True
        x_nom = fmi2GetNominalsOfContinuousStates(nfmu.fmu)
    end

    ignore_derivatives() do
        if idx != -1
            e = FMU2Event(t, UInt64(idx), left_x, right_x)
            push!(nfmu.solution.events, e)
        end
    end

    #@assert fmi2EnterContinuousTimeMode(c) == fmi2StatusOK
end

# Does one step in the simulation.
function stepCompleted(nfmu::ME_NeuralFMU, x, t, integrator, tStart, tStop)

    # there might be no component!
    # @assert nfmu.currentComponent.state == fmi2ComponentStateContinuousTimeMode "stepCompleted(...): Must be in continuous time mode."

    if nfmu.progressMeter !== nothing
        ProgressMeter.update!(nfmu.progressMeter, floor(Integer, 1000.0*(t-tStart)/(tStop-tStart)) )
    end

    if nfmu.currentComponent != nothing && (!nfmu.fmu.executionConfig.useComponentShadow || nfmu.currentComponent.realComponent != nothing)
        (status, enterEventMode, terminateSimulation) = fmi2CompletedIntegratorStep(nfmu.currentComponent, fmi2True)

        if enterEventMode == fmi2True
            affectFMU!(nfmu, integrator, -1)
        else
            # ToDo: set inputs here
            #fmiSetReal(myFMU, InputRef, Value)
        end
    end

end

# save FMU values
function saveValues(nfmu, c::FMU2Component, recordValues, x, t, integrator)

    if isa(t, ForwardDiff.Dual)
        t = ForwardDiff.value(t)
    end

    if all(isa.(x, ForwardDiff.Dual))
        x = collect(ForwardDiff.value(e) for e in x)
    end

    # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN
    nfmu.neuralODE.model(x) # evaluate NeuralFMU (set new states)

    fmi2SetTime(c, t)

    (fmi2GetReal(c, recordValues)...,)
end

function fd_eltypes(e::SubArray{<:ForwardDiff.Dual{T, V, N}}) where {T, V, N}
    return (T, V, N)
end

function fd_eltypes(e::Vector{<:ForwardDiff.Dual{T, V, N}}) where {T, V, N}
    return (T, V, N)
end

function fx(nfmu,
    dx,#::Array{<:Real},
    x,#::Array{<:Real},
    p::Array,
    t::Real)

    dx_tmp = fx(nfmu,x,p,t)

    if all(isa.(dx_tmp, ForwardDiff.Dual))
        if all(isa.(dx, ForwardDiff.Dual))
            dx[:] = dx_tmp
        else
            dx[:] = collect(ForwardDiff.value(e) for e in dx_tmp)
        end
        #dx[:] = collect(ForwardDiff.value(e) for e in dx_tmp)
    else
        if all(isa.(dx, ForwardDiff.Dual))
            #dx_tmp = collect(ForwardDiff.value(e) for e in dx)
            #fmi2GetDerivatives!(nfmu.currentComponent, dx_tmp)
            T, V, N = fd_eltypes(dx)
            dx[:] = collect(ForwardDiff.Dual{T, V, N}(V(dx_tmp[i]), ForwardDiff.partials(dx[i])    ) for i in 1:length(dx))
        else
            dx[:] = dx_tmp
        end
    end

    #return dx_tmp
    return dx
end

function fx(nfmu,
    x,#::Array{<:Real},
    p::Array,
    t::Real)

    # if all(isa.(x, ForwardDiff.Dual))
    #     x = collect(ForwardDiff.value(e) for e in x)
    # end

    c = nfmu.currentComponent

    if c === nothing
        # this should never happen!
        return zeros(length(x))
    end

    dx = nfmu.neuralODE.re(p)(x)

    ignore_derivatives() do

        if isa(t, ForwardDiff.Dual)
            t = ForwardDiff.value(t)
        end

        # if t == nfmu.tspan[1]
        #     @debug ["fx (start)"]
        # end

        # if t == nfmu.tspan[end]
        #     @debug ["fx (stop)"]
        # end

        fmi2SetTime(c, t)

    end

    # build up ẋ interpolation polynominal
    # ignore_derivatives() do
    #     if nfmu.fmu.executionConfig.useComponentShadow
    #         if nfmu.fmu.ẋy_interp == nothing
    #             if t <= nfmu.tspan[end]
    #                 if length(nfmu.fmu.t_cache) == 0 || nfmu.fmu.t_cache[end] < t
    #                     push!(nfmu.fmu.ẋy_cache, collect(ForwardDiff.value(e) for e in dx) )
    #                     push!(nfmu.fmu.t_cache, t)
    #                 end
    #             end

    #             if t >= nfmu.tspan[end] # endpoint

    #             end
    #         end
    #     end
    # end

    return dx
end

##### EVENT HANDLING END

"""
Constructs a ME-NeuralFMU where the FMU is at an arbitrary location inside of the NN.

# Arguents
    - `fmu` the considered FMU inside the NN
    - `model` the NN topology (e.g. Flux.chain)
    - `tspan` simulation time span
    - `alg` a numerical ODE solver

# Keyword arguments
    - `saveat` time points to save the NeuralFMU output, if empty, solver step size is used (may be non-equidistant)
    - `fixstep` forces fixed step integration
    - `recordFMUValues` additionally records internal FMU variables (currently not supported because of open issues)
"""
function ME_NeuralFMU(fmu::FMU2,
                      model,
                      tspan,
                      alg=nothing;
                      saveat=[],
                      recordFMUValues = nothing,
                      callbacks = [],
                      kwargs...)

    nfmu = ME_NeuralFMU()
    nfmu.fmu = fmu
    nfmu.customCallbacks = callbacks

    nfmu.saved_values = nothing

    nfmu.recordValues = prepareValueReference(fmu, recordFMUValues)

    nfmu.neuralODE = NeuralODE(model, tspan, alg; saveat=saveat, kwargs...)

    nfmu.tspan = tspan
    nfmu.saveat = saveat

    nfmu
end

"""
Constructs a CS-NeuralFMU where the FMU is at an arbitrary location inside of the NN.

# Arguents
    - `fmu` the considered FMU inside the NN
    - `model` the NN topology (e.g. Flux.chain)
    - `tspan` simulation time span

# Keyword arguments
    - `saveat` time points to save the NeuralFMU output, if empty, solver step size is used (may be non-equidistant)
"""
function CS_NeuralFMU(fmu::Union{FMU2, Vector{<:FMU2}},
                      model,
                      tspan;
                      saveat=[],
                      recordValues = [])

    nfmu = nothing
    if typeof(fmu) == FMU2
        nfmu = CS_NeuralFMU{FMU2, FMU2Component}()
        nfmu.currentComponent = nothing
    else
        nfmu = CS_NeuralFMU{Vector{FMU2}, Vector{FMU2Component} }()
        nfmu.currentComponent = Vector{Union{FMU2Component, Nothing}}(nothing, length(fmu))
    end

    nfmu.fmu = fmu

    nfmu.model = model # Chain(model.layers...)

    nfmu.tspan = tspan
    nfmu.saveat = saveat

    nfmu
end

function prepareFMU(fmu::FMU2, c::Union{Nothing, FMU2Component}, instantiate::Union{Nothing, Bool}, freeInstance::Union{Nothing, Bool}, terminate::Union{Nothing, Bool}, reset::Union{Nothing, Bool}, setup::Union{Nothing, Bool}, parameters::Union{Dict{<:Any, <:Any}, Nothing}, t_start, t_stop, tolerance;
    x0::Union{Array{<:Real}, Nothing}=nothing, pushComponents::Bool=true)

    c = nothing

    ignore_derivatives() do
        if instantiate === nothing
            instantiate = fmu.executionConfig.instantiate
        end

        if freeInstance === nothing
            freeInstance = fmu.executionConfig.freeInstance
        end

        if terminate === nothing
            terminate = fmu.executionConfig.terminate
        end

        if reset === nothing
            reset = fmu.executionConfig.reset
        end

        if setup === nothing
            setup = fmu.executionConfig.setup
        end

        # instantiate (hard)
        if instantiate

            # remove old one if we missed it (callback)
            if c != nothing
                if freeInstance
                    fmi2FreeInstance!(c)
                    @debug "[AUTO-RELEASE INST]"
                end
            end

            c = fmi2Instantiate!(fmu; pushComponents=pushComponents)
            @debug "[NEW INST]"
        else
            if c === nothing
                c = fmu.components[end]
            end
        end

        # soft terminate (if necessary)
        if terminate
            retcode = fmi2Terminate(c; soft=true)
            @assert retcode == fmi2StatusOK "fmi2Simulate(...): Termination failed with return code $(retcode)."
        end

        # soft reset (if necessary)
        if reset
            retcode = fmi2Reset(c; soft=true)
            @assert retcode == fmi2StatusOK "fmi2Simulate(...): Reset failed with return code $(retcode)."
        end

        # enter setup (hard)
        if setup
            retcode = fmi2SetupExperiment(c, t_start, t_stop; tolerance=tolerance)
            @assert retcode == fmi2StatusOK "fmi2Simulate(...): Setting up experiment failed with return code $(retcode)."

            retcode = fmi2EnterInitializationMode(c)
            @assert retcode == fmi2StatusOK "fmi2Simulate(...): Entering initialization mode failed with return code $(retcode)."
        end

        if x0 === nothing
            x0 = fmi2GetContinuousStates(c)
        else
            retcode = fmi2SetContinuousStates(c, x0)
            @assert retcode == fmi2StatusOK "fmi2Simulate(...): Setting initial state failed with return code $(retcode)."
        end

        if parameters !== nothing
            retcodes = fmi2Set(c, collect(keys(parameters)), collect(values(parameters)) )
            @assert all(retcodes .== fmi2StatusOK) "fmi2Simulate(...): Setting initial parameters failed with return code $(retcode)."
        end

        # exit setup (hard)
        if setup
            retcode = fmi2ExitInitializationMode(c)
            @assert retcode == fmi2StatusOK "fmi2Simulate(...): Exiting initialization mode failed with return code $(retcode)."
        end
    end # ignore_derivatives

    @assert c.state == fmi2ComponentStateEventMode "After prepareFMU(...), FMU must be in event mode!"

    return c, x0
end

function prepareFMU(fmu::Vector{FMU2}, c::Vector{Union{Nothing, FMU2Component}}, instantiate::Union{Nothing, Bool}, freeInstance::Union{Nothing, Bool}, terminate::Union{Nothing, Bool}, reset::Union{Nothing, Bool}, setup::Union{Nothing, Bool}, parameters::Union{Vector{Union{Dict{<:Any, <:Any}, Nothing}}, Nothing}, t_start, t_stop, tolerance;
    x0::Union{Vector{Union{Array{<:Real}, Nothing}}, Nothing}=nothing)

    ignore_derivatives() do
        for i in 1:length(fmu)

            if instantiate === nothing
                instantiate = fmu[i].executionConfig.instantiate
            end

            if freeInstance === nothing
                freeInstance = fmu[i].executionConfig.freeInstance
            end

            if terminate === nothing
                terminate = fmu[i].executionConfig.terminate
            end

            if reset === nothing
                reset = fmu[i].executionConfig.reset
            end

            if setup === nothing
                setup = fmu[i].executionConfig.setup
            end

            # instantiate (hard)
            if instantiate
                # remove old one if we missed it (callback)
                if c[i] != nothing
                    if freeInstance
                        fmi2FreeInstance!(c[i])
                        @debug "[AUTO-RELEASE INST]"
                    end
                end

                c[i] = fmi2Instantiate!(fmu[i])
                @debug "[NEW INST]"
            else
                if c[i] === nothing
                    c[i] = fmu[i].components[end]
                end
            end

            # soft terminate (if necessary)
            if terminate
                retcode = fmi2Terminate(c[i]; soft=true)
                @assert retcode == fmi2StatusOK "fmi2Simulate(...): Termination failed with return code $(retcode)."
            end

            # soft reset (if necessary)
            if reset
                retcode = fmi2Reset(c[i]; soft=true)
                @assert retcode == fmi2StatusOK "fmi2Simulate(...): Reset failed with return code $(retcode)."
            end

            # enter setup (hard)
            if setup
                retcode = fmi2SetupExperiment(c[i], t_start, t_stop; tolerance=tolerance)
                @assert retcode == fmi2StatusOK "fmi2Simulate(...): Setting up experiment failed with return code $(retcode)."

                retcode = fmi2EnterInitializationMode(c[i])
                @assert retcode == fmi2StatusOK "fmi2Simulate(...): Entering initialization mode failed with return code $(retcode)."
            end

            if x0 !== nothing
                if x0[i] === nothing
                    x0[i] = fmi2GetContinuousStates(c[i])
                else
                    retcode = fmi2SetContinuousStates(c[i], x0[i])
                    @assert retcode == fmi2StatusOK "fmi2Simulate(...): Setting initial state failed with return code $(retcode)."
                end
            end

            if parameters !== nothing
                if parameters[i] !== nothing
                    retcodes = fmi2Set(c[i], collect(keys(parameters[i])), collect(values(parameters[i])) )
                    @assert all(retcodes .== fmi2StatusOK) "fmi2Simulate(...): Setting initial parameters failed with return code $(retcode)."
                end
            end

            # exit setup (hard)
            if setup
                retcode = fmi2ExitInitializationMode(c[i])
                @assert retcode == fmi2StatusOK "fmi2Simulate(...): Exiting initialization mode failed with return code $(retcode)."
            end
        end

    end # ignore_derivatives

    return c, x0
end

function finishFMU(fmu::FMU2, c::Union{FMU2Component, Nothing}, freeInstance::Union{Nothing, Bool}; popComponent::Bool=true)

    ignore_derivatives() do
        if freeInstance === nothing
            freeInstance = fmu.executionConfig.freeInstance
        end

        if c != nothing
            if freeInstance
                fmi2FreeInstance!(c; popComponent=popComponent)
                @debug "[RELEASED INST]"
            end
            c = nothing
        end

    end # ignore_derivatives

    return c
end

function finishFMU(fmu::Vector{FMU2}, c::Vector{Union{FMU2Component, Nothing}}, freeInstance::Union{Nothing, Bool})

    ignore_derivatives() do
        for i in 1:length(fmu)
            if freeInstance === nothing
                freeInstance = fmu[i].executionConfig.freeInstance
            end

            if c[i] != nothing
                if freeInstance
                    fmi2FreeInstance!(c[i])
                    @debug "[RELEASED INST]"
                end
                c[i] = nothing
            end
        end

    end # ignore_derivatives

    return c
end

"""
Evaluates the ME_NeuralFMU in the timespan given during construction or in a custom timespan from `t_start` to `t_stop` for a given start state `x_start`.

# Keyword arguments
    - `reset`, the FMU is reset every time evaluation is started (default=`true`).
    - `setup`, the FMU is set up every time evaluation is started (default=`true`).
"""
function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nothing,
                              t_start::Union{Real, Nothing} = nothing,
                              t_stop::Union{Real, Nothing} = nothing;
                              showProgress::Bool = false,
                              tolerance::Union{Real, Nothing} = nothing,
                              parameters::Union{Dict{<:Any, <:Any}, Nothing} = nothing,
                              setup::Union{Bool, Nothing} = nothing,
                              reset::Union{Bool, Nothing} = nothing,
                              instantiate::Union{Bool, Nothing} = nothing,
                              freeInstance::Union{Bool, Nothing} = nothing,
                              terminate::Union{Bool, Nothing} = nothing,
                              kwargs...)

    saving = (length(nfmu.recordValues) > 0)
    sense = nfmu.fmu.executionConfig.sensealg
    inPlace = nfmu.fmu.executionConfig.inPlace
    tspan = getfield(nfmu.neuralODE, :tspan)

    ignore_derivatives() do
        @debug "ME_NeuralFMU..."

        nfmu.firstRun = true

        nfmu.solution = FMU2Solution(nfmu.fmu)

        nfmu.tolerance = tolerance
        nfmu.parameters = parameters
        nfmu.setup = setup
        nfmu.reset = reset
        nfmu.instantiate = instantiate
        nfmu.freeInstance = freeInstance
        nfmu.terminate = terminate
        nfmu.currentComponent = nothing
        nfmu.solveCycle = 0
        nfmu.progressMeter = nothing
        nfmu.callbacks = []

        nfmu.x0 = x_start

        # simulation tspan
        if t_start === nothing
            t_start = tspan[1]
        end
        if t_stop === nothing
            t_stop = tspan[end]
        end
        nfmu.tspan = (t_start, t_stop)

        # remove last run's shadow
        if nfmu.fmu.executionConfig.useComponentShadow
            if length(nfmu.fmu.components) > 0 && nfmu.fmu.components[end] == nfmu.currentComponent
                ind = findall(x -> x==nfmu.currentComponent, nfmu.fmu.components)

                if length(ind) == 1
                    @debug ["SHADOW CLEANUP"]
                    deleteat!(nfmu.fmu.components, ind)
                end
            end

            if nfmu.progressMeter != nothing
                ProgressMeter.finish!(nfmu.progressMeter)
            end
        end

        # from here on, we are in event mode, if `setup=false` this is the job of the user
        # @assert nfmu.currentComponent.state == fmi2ComponentStateEventMode "FMU needs to be in event mode after setup (end)."

        # this is not necessary for NeuralFMUs:
        #x0 = fmi2GetContinuousStates(nfmu.currentComponent)
        #x0_nom = fmi2GetNominalsOfContinuousStates(nfmu.currentComponent)

        # initial event handling
        # handleEvents(nfmu.currentComponent)
        # fmi2EnterContinuousTimeMode(nfmu.currentComponent)

        nfmu.fmu.hasStateEvents = (nfmu.fmu.modelDescription.numberOfEventIndicators > 0)
        # nfmu.fmu.hasTimeEvents = (nfmu.currentComponent.eventInfo.nextEventTimeDefined == fmi2True)

        startcb = FunctionCallingCallback((u, t, integrator) -> startCallback(nfmu, t);
                 funcat=[nfmu.tspan[1]], func_start=true, func_everystep=false)
        push!(nfmu.callbacks, startcb)

        # integrator step callback

        if showProgress

            nfmu.progressMeter = ProgressMeter.Progress(1000; desc="Simulating ME-NeuralFMU ...", color=:blue, dt=1.0) #, barglyphs=ProgressMeter.BarGlyphs("[=> ]"))
            ProgressMeter.update!(nfmu.progressMeter, 0) # show it!

            stepCb = FunctionCallingCallback((x, t, integrator) -> stepCompleted(nfmu, x, t, integrator, t_start, t_stop);
                                                func_everystep=true,
                                                func_start=true)
            push!(nfmu.callbacks, stepCb)
        end

        # state event callback

        if nfmu.fmu.hasStateEvents && nfmu.fmu.executionConfig.handleStateEvents
            if nfmu.fmu.executionConfig.useVectorCallbacks

                eventCb = VectorContinuousCallback((out, x, t, integrator) -> condition(nfmu, out, x, t, integrator),
                                                (integrator, idx) -> affectFMU!(nfmu, integrator, idx),
                                                Int64(nfmu.fmu.modelDescription.numberOfEventIndicators);
                                                rootfind=RightRootFind,
                                                save_positions=(false, false),
                                                interp_points=nfmu.fmu.executionConfig.rootSearchInterpolationPoints)
                push!(nfmu.callbacks, eventCb)
            else

                for idx in 1:nfmu.fmu.modelDescription.numberOfEventIndicators
                    eventCb = ContinuousCallback((x, t, integrator) -> conditionSingle(nfmu, idx, x, t, integrator),
                                                    (integrator) -> affectFMU!(nfmu, integrator, idx);
                                                    rootfind=RightRootFind,
                                                    save_positions=(false, false),
                                                    interp_points=nfmu.fmu.executionConfig.rootSearchInterpolationPoints)
                    push!(nfmu.callbacks, eventCb)
                end
            end
        end

        # time event handling

        #if nfmu.fmu.executionConfig.handleTimeEvents # && nfmu.fmu.hasTimeEvents
            timeEventCb = IterativeCallback((integrator) -> time_choice(nfmu, integrator, t_start, t_stop),
                                            (integrator) -> affectFMU!(nfmu, integrator, 0),
                                            Float64;
                                            initial_affect=false,
                                            save_positions=(false,false))

            push!(nfmu.callbacks, timeEventCb)
        #end

        if saving
            nfmu.solution.values = SavedValues(Float64, Tuple{collect(Float64 for i in 1:length(nfmu.recordValues))...})

            savingCB = SavingCallback((x, t, integrator) -> saveValues(nfmu, nfmu.fmu.components[end], nfmu.recordValues, x, t, integrator),
                            nfmu.solution.values,
                              saveat=nfmu.saveat)
            push!(nfmu.callbacks, savingCB)
        end

        # custom callbacks
        for cb in nfmu.customCallbacks
            push!(nfmu.callbacks, cb)
        end

        # stoping cb
        stopcb = FunctionCallingCallback((u, t, integrator) -> stopCallback(nfmu, t);
                                    funcat=[nfmu.tspan[end]])
        push!(nfmu.callbacks, stopcb)

        # auto pick sensealg

        if sense === nothing

            # currently, Callbacks only working with ForwardDiff
            if length(nfmu.callbacks) > 2 # only start and stop callback
                p_len = 0
                fp = Flux.params(nfmu)
                if length(fp) > 0
                    p_len = length(fp[1])
                end
                fds_chunk_size = p_len
                # limit to 256 because of RAM usage
                if fds_chunk_size > 256
                    fds_chunk_size = 256
                end
                if fds_chunk_size < 1
                    fds_chunk_size = 1
                end

                sense = ForwardDiffSensitivity(;chunk_size=fds_chunk_size, convert_tspan=true)
                #sense = QuadratureAdjoint(autojacvec=ZygoteVJP())

            else
                sense = InterpolatingAdjoint(autojacvec=ZygoteVJP()) # EnzymeVJP()

                if inPlace === true
                    #@info "NeuralFMU: The configuration orders to use `inPlace=true`, but this is currently not supported by the (automatically) determined sensealg `Zygote`. Switching to `inPlace=false`. If you don't want this behaviour, explicitely choose a sensealg instead of `nothing`."
                    inPlace = false # Zygote only works for out-of-place (currently)
                end
            end
        end

    end # ignore_derivatives

    # setup problem and parameters
    p = nfmu.neuralODE.p
    prob = nothing

    if inPlace
        ff = ODEFunction{true}((dx, x, p, t) -> fx(nfmu, dx, x, p, t),
                               tgrad=nothing)
        prob = ODEProblem{true}(ff, nfmu.x0, nfmu.tspan, p)
    else
        ff = ODEFunction{false}((x, p, t) -> fx(nfmu, x, p, t),
                                tgrad=basic_tgrad)
        prob = ODEProblem{false}(ff, nfmu.x0, nfmu.tspan, p)
    end

    if (length(nfmu.callbacks) == 2) # only start and stop callback, so the system is pure continuous
        startCallback(nfmu, nfmu.tspan[1])
        nfmu.solution.states = solve(prob, nfmu.neuralODE.args...; sensealg=sense, saveat=nfmu.saveat, nfmu.neuralODE.kwargs...)
        stopCallback(nfmu, nfmu.tspan[end])
    else
        nfmu.solution.states = solve(prob, nfmu.neuralODE.args...; sensealg=sense, saveat=nfmu.saveat, callback = CallbackSet(nfmu.callbacks...), nfmu.neuralODE.kwargs...)
    end

    # this code is executed after the initial solve-evaluation (further ForwardDiff-solve-evaluations will happen after this!)
    ignore_derivatives() do

        nfmu.solution.success = (nfmu.solution.states.retcode == :Success)

    end # ignore_derivatives

    return nfmu.solution
end

"""
Evaluates the CS_NeuralFMU in the timespan given during construction or in a custum timespan from `t_start` to `t_stop` with a given time step size `t_step`.

Via optional argument `reset`, the FMU is reset every time evaluation is started (default=`true`).
"""
function (nfmu::CS_NeuralFMU{F, C})(inputFct,
                                 t_step::Real,
                                 t_start::Real = nfmu.tspan[1],
                                 t_stop::Real = nfmu.tspan[end];
                                 tolerance::Union{Real, Nothing} = nothing,
                                 parameters::Union{Dict{<:Any, <:Any}, Nothing} = nothing,
                                 setup::Union{Bool, Nothing} = nothing,
                                 reset::Union{Bool, Nothing} = nothing,
                                 instantiate::Union{Bool, Nothing} = nothing,
                                 freeInstance::Union{Bool, Nothing} = nothing,
                                 terminate::Union{Bool, Nothing} = nothing) where {F, C}

    ignore_derivatives() do
        nfmu.solution = FMU2Solution(nfmu.fmu)
    end

    nfmu.currentComponent, _ = prepareFMU(nfmu.fmu, nfmu.currentComponent, instantiate, freeInstance, terminate, reset, setup, parameters, t_start, t_stop, tolerance)

    ts = collect(t_start:t_step:t_stop)
    nfmu.currentComponent.skipNextDoStep = true # skip first fim2DoStep-call
    model_input = inputFct.(ts)
    valueStack = nfmu.model.(model_input)

    ignore_derivatives() do
        nfmu.solution.success = true
    end

    nfmu.solution.values = SavedValues{Float64, Vector{Float64}}(ts, valueStack )

    # this is not possible in CS, clean-up happens at the next call
    #nfmu.currentComponent = finishFMU(nfmu.fmu, nfmu.currentComponent, freeInstance)

    return nfmu.solution
end
function (nfmu::CS_NeuralFMU{Vector{F}, Vector{C}})(inputFct,
                                         t_step::Real,
                                         t_start::Real = nfmu.tspan[1],
                                         t_stop::Real = nfmu.tspan[end];
                                         tolerance::Union{Real, Nothing} = nothing,
                                         parameters::Union{Vector{Union{Dict{<:Any, <:Any}, Nothing}}, Nothing} = nothing,
                                         setup::Union{Bool, Nothing} = nothing,
                                         reset::Union{Bool, Nothing} = nothing,
                                         instantiate::Union{Bool, Nothing} = nothing,
                                         freeInstance::Union{Bool, Nothing} = nothing,
                                         terminate::Union{Bool, Nothing} = nothing) where {F, C}

    ignore_derivatives() do
        nfmu.solution = FMU2Solution(nfmu.fmu)
    end

    nfmu.currentComponent, _ = prepareFMU(nfmu.fmu, nfmu.currentComponent, instantiate, freeInstance, terminate, reset, setup, parameters, t_start, t_stop, tolerance)

    ts = collect(t_start:t_step:t_stop)
    model_input = inputFct.(ts)
    valueStack = nfmu.model.(model_input)

    ignore_derivatives() do
        nfmu.solution.success = true
    end

    nfmu.solution.values = SavedValues{Float64, Vector{Float64}}(ts, valueStack )

    # this is not possible in CS, clean-up happens at the next call
    #nfmu.currentComponent = finishFMU(nfmu.fmu, nfmu.currentComponent, freeInstance)

    return nfmu.solution
end

# adapting the Flux functions
function Flux.params(neuralfmu::ME_NeuralFMU)
    Flux.params(neuralfmu.neuralODE)
end

function Flux.params(neuralfmu::CS_NeuralFMU)
    Flux.params(neuralfmu.model)
end

# FMI version independent dosteps

"""
Wrapper. Call ```fmi2EvaluateME``` for more information.
"""
function fmiEvaluateME(str::fmi2Struct,
                     x::Array{<:Real},
                     t::Real = (typeof(str) == FMU2 ? str.components[end].t : str.t),
                     setValueReferences::Array{fmi2ValueReference} = zeros(fmi2ValueReference, 0),
                     setValues::Array{<:Real} = zeros(Real, 0),
                     getValueReferences::Array{fmi2ValueReference} = zeros(fmi2ValueReference, 0))
    fmi2EvaluateME(str, x, t,
                setValueReferences,
                setValues,
                getValueReferences)
end

"""
Wrapper. Call ```fmi2DoStepCS``` for more information.
"""
function fmiDoStepCS(str::fmi2Struct,
                     dt::Real,
                     setValueReferences::Array{fmi2ValueReference} = zeros(fmi2ValueReference, 0),
                     setValues::Array{<:Real} = zeros(Real, 0),
                     getValueReferences::Array{fmi2ValueReference} = zeros(fmi2ValueReference, 0))
    fmi2DoStepCS(str, dt, setValueReferences, setValues, getValueReferences)
end

"""
Wrapper. Call ```fmi2InputDoStepCSOutput``` for more information.
"""
function fmiInputDoStepCSOutput(str::fmi2Struct,
                                dt::Real,
                                u::Array{<:Real})
    fmi2InputDoStepCSOutput(str, dt, u)
end

# function ChainRulesCore.rrule(f::Union{typeof(fmi2SetupExperiment),
#                                        typeof(fmi2EnterInitializationMode),
#                                        typeof(fmi2ExitInitializationMode),
#                                        typeof(fmi2Reset),
#                                        typeof(fmi2Terminate)}, args...)

#     y = f(args...)

#     function pullback(ȳ)
#         return collect(ZeroTangent() for arg in args)
#     end

#     return y, fmi2EvaluateME_pullback
# end


## ExoplanetsSysSim/examples/setup_abc_christiansen.jl
## (c) 2016 Eric B. Ford & Danley C. Hsu

module EvalSysSimModel
  export setup, get_param_vector, get_ss_obs #, evaluate_model
  export gen_data, calc_summary_stats, calc_distance, is_valid
  using ExoplanetsSysSim
  include(joinpath(Pkg.dir(),"ExoplanetsSysSim","examples","hsu_etal_2017", "christiansen_func.jl"))
  include(joinpath(pwd(), "param_file.jl"))

  sim_param_closure = SimParam()
  summary_stat_ref_closure =  CatalogSummaryStatistics()

    function is_valid(param_vector::Vector{Float64})
      global sim_param_closure
      update_sim_param_from_vector!(param_vector,sim_param_closure)
      const sigma_log_r = get_real(sim_param_closure,"sigma_log_planet_radius")
      const sigma_log_P = get_real(sim_param_closure,"sigma_log_planet_period")
      const rate_tab::Array{Float64,2} = get_any(sim_param_closure, "obs_par", Array{Float64,2})
      const lambda = sum_kbn(rate_tab)
    # println("Lambda = ", lambda)
      if sigma_log_r <= 0. || sigma_log_P<=0. || lambda > 10. || any(x -> x < 0., rate_tab)
         return false
      end
      return true
    end

#=
    function norm_christiansen(param_vector::Vector{Float64})
      log_z_sum = logsumexp(param_vector)
      param_vector = exp(param_vector-z_sum)
    end
=#

    function gen_data(param_vector::Vector{Float64})
      global sim_param_closure
      update_sim_param_from_vector!(param_vector,sim_param_closure)
      cat_phys = generate_kepler_physical_catalog(sim_param_closure)
      cat_phys_cut = ExoplanetsSysSim.generate_obs_targets(cat_phys, sim_param_closure)
      cat_obs = ExoplanetsSysSim.observe_kepler_targets_single_obs(cat_phys_cut, sim_param_closure)
      return cat_obs
    end

    function calc_summary_stats(cat::KeplerObsCatalog )
      global sim_param_closure
      #sum_stat = calc_summary_stats_obs_binned_rates(cat[1], sim_param_closure)
      sum_stat = calc_summary_stats_obs_binned_rates(cat, sim_param_closure)
      return sum_stat
    end

    function calc_distance(sum_stat_obs::CatalogSummaryStatistics,sum_stat_sim::CatalogSummaryStatistics, n::Integer = 0)
      global sim_param_closure
      dist1 = calc_distance_vector_binned(sum_stat_obs,sum_stat_sim, 1, sim_param_closure)
      num_available = length(dist1)
      num_to_use = n>0 ? min(n,num_available) : num_available
      return calc_scalar_distance(dist1[1:num_to_use])
    end

  function setup()
    global sim_param_closure = setup_sim_param_christiansen()

    sim_param_closure = set_test_param(sim_param_closure)
    add_param_fixed(sim_param_closure,"stellar_catalog","q1_q12_christiansen.jld")

    ### Use simulated planet candidate catalog data
    add_param_fixed(sim_param_closure,"num_kepler_targets",150969)  # For "observed" catalog
    cat_obs = simulated_read_kepler_observations(sim_param_closure)
    #inv_det_prob(cat_obs, sim_param_closure)
    #println("------------------------")
    global summary_stat_ref_closure = calc_summary_stats_obs_binned_rates(cat_obs,sim_param_closure, trueobs_cat = true)
    ###

    ### Use real planet candidate catalog data
    #add_param_fixed(sim_param_closure,"koi_catalog","q1_q16_koi_cand.csv")
    #cat_obs = setup_actual_planet_candidate_catalog(setup_star_table_christiansen(sim_param_closure), sim_param_closure)
    #global summary_stat_ref_closure = calc_summary_stats_obs_binned_rates(cat_obs,sim_param_closure, trueobs_cat = true)
    ###

    num_targ = get_int(sim_param_closure,"num_kepler_targets")
    add_param_fixed(sim_param_closure,"num_targets_sim_pass_one",num_targ)  # Set universal simulated catalog size

  end

  get_param_vector() = make_vector_of_sim_param(sim_param_closure)
  get_ss_obs() = summary_stat_ref_closure

  function set_simparam_ss(sim_param::ExoplanetsSysSim.SimParam, ss_true::ExoplanetsSysSim.CatalogSummaryStatistics)
    global sim_param_closure = sim_param
    global summary_stat_ref_closure = ss_true
  end

end  # module EvalSysSimModel

include(joinpath(Pkg.dir("ABC"),"src/composite.jl"))

module SysSimABC
  export setup_abc, run_abc
  import ABC
  import Distributions
  using CompositeDistributions
  import ExoplanetsSysSim
  import EvalSysSimModel
  include(joinpath(Pkg.dir(),"ExoplanetsSysSim","examples","hsu_etal_2017", "christiansen_func.jl"))
  include(joinpath(pwd(), "param_file.jl"))

  function setup_abc(num_dist::Integer = 0)
    EvalSysSimModel.setup()
    #println("# setup_abc: ",EvalSysSimModel.sim_param_closure)
    theta_true = EvalSysSimModel.get_param_vector()
    param_prior = CompositeDist( Distributions.ContinuousDistribution[Distributions.Uniform(0., 1.) for x in 1:length(theta_true)] )
    in_parallel = nworkers() > 1 ? true : false

    calc_distance_ltd(sum_stat_obs::ExoplanetsSysSim.CatalogSummaryStatistics,sum_stat_sim::ExoplanetsSysSim.CatalogSummaryStatistics) = EvalSysSimModel.calc_distance(sum_stat_obs,sum_stat_sim,num_dist)

    global abc_plan = ABC.abc_pmc_plan_type(EvalSysSimModel.gen_data,EvalSysSimModel.calc_summary_stats,calc_distance_ltd, param_prior, make_proposal_dist=ABC.make_proposal_dist_gaussian_diag_covar, is_valid=EvalSysSimModel.is_valid,
                                     num_part=40, num_max_attempt=50, num_max_times=200, epsilon_init=9.9e99, target_epsilon=1.0e-100, in_parallel=in_parallel, adaptive_quantiles = false, epsilon_reduction_factor=0.9, tau_factor=2.0);
  end

  function run_abc_largegen(pop::ABC.abc_population_type, ss_true::ExoplanetsSysSim.CatalogSummaryStatistics, epshist_targ::Float64, npart::Integer = 1000, num_dist::Integer = 0)
    sim_param_closure = setup_sim_param_christiansen()
    sim_param_closure = set_test_param(sim_param_closure)
    #add_param_fixed(sim_param_closure,"num_kepler_targets",150000)  # For "observed" catalog
    EvalSysSimModel.set_simparam_ss(sim_param_closure, ss_true)

    theta_true = EvalSysSimModel.get_param_vector()
    param_prior = CompositeDist( Distributions.ContinuousDistribution[Distributions.Uniform(0., 1.) for x in 1:length(theta_true)] )
    in_parallel = nworkers() > 1 ? true : false

    calc_distance_ltd(sum_stat_obs::ExoplanetsSysSim.CatalogSummaryStatistics,sum_stat_sim::ExoplanetsSysSim.CatalogSummaryStatistics) = EvalSysSimModel.calc_distance(sum_stat_obs,sum_stat_sim,num_dist)

    global abc_plan = ABC.abc_pmc_plan_type(EvalSysSimModel.gen_data,EvalSysSimModel.calc_summary_stats,calc_distance_ltd, param_prior, is_valid=EvalSysSimModel.is_valid,
                                     num_part=npart, num_max_attempt=50, num_max_times=1, epsilon_init=9.9e99, target_epsilon=1.0e-100, in_parallel=in_parallel);

    println("# run_abc_largegen: ",EvalSysSimModel.sim_param_closure)
    sampler_plot = abc_plan.make_proposal_dist(pop, abc_plan.tau_factor)
    theta_plot = Array{Float64}(size(pop.theta, 1), npart)
    for i in 1:npart
      theta_plot[:,i], dist_plot, attempts_plot = ABC.generate_theta(abc_plan, sampler_plot, ss_true, epshist_targ)
    end
    return theta_plot
  end

  function run_abc(abc_plan::ABC.abc_pmc_plan_type)
    println("# run_abc: ",EvalSysSimModel.sim_param_closure)
    ss_true = EvalSysSimModel.get_ss_obs()
    #println("True catalog SS: ", ss_true)
    pop_out = ABC.run_abc(abc_plan,ss_true;verbose=true)
  end

  function run_abc(abc_plan::ABC.abc_pmc_plan_type, pop::ABC.abc_population_type)
    dist_threshold = maximum(pop.dist)
    EvalSysSimModel.add_param_fixed(EvalSysSimModel.sim_param_closure,"minimum ABC dist skip pass 2",dist_threshold)
    println("# run_abc: ",EvalSysSimModel.sim_param_closure)
    ss_true = EvalSysSimModel.get_ss_obs()
    pop_out = ABC.run_abc(abc_plan,ss_true,pop;verbose=true)
  end

end # module SysSimABC


# the core loop for several CRC calculations.  it is used by CRC32,
# whose zlib implementation in C runs ~2.5 times faster than this
# code.

function generic_loop(::Type{D}, remainder::A, data,
        table::Vector{A}, word_size) where {D<:Unsigned, A<:Unsigned}
    for word::D in data
        word ⊻= remainder%D
        remainder = (remainder >>> word_size) ⊻ table[1 + word]
    end
    remainder
end

# try accessing the lowest byte directly to avoid the shift (slow)
struct Bytes
    a::UInt8
    b::UInt8
    c::UInt8
    d::UInt8
end

function struct_loop(::Type{UInt8}, remainder::UInt64, data::Vector{UInt8}, table::Vector{UInt64}, ignored)
    r::UInt64 = remainder
    b = reinterpret(Bytes, r)
    for i in 1:length(data)
        @inbounds word::UInt8 = data[i]
        word ⊻= b.a
        r = (r >>> 8) ⊻ table[1 + word]
        b = reinterpret(Bytes, r)
    end
    return r
end

# try accessing via aliased arrays

function array_loop(::Type{UInt8}, remainder::UInt64, data::Vector{UInt8}, table::Vector{UInt64}, ignored)
    r = Vector{UInt64}(undef, 1)
    @inbounds r[1] = remainder
    b = reinterpret(UInt8, r)
    for i in 1:length(data)
        @inbounds word::UInt8 = data[i]
        @inbounds word ⊻= b[1]
        @inbounds r[1] = (r[1] >>> 8) ⊻ table[1 + word]  # @inbounds no help here
    end
    r[1]
end

# according to my notes this was slightly (~15%) faster than
# generic_loop.  repeated here, however, it seems to be identical to
# the generic case.

function byte_loop(::Type{UInt8}, remainder::A, data::Vector{UInt8}, table::Vector{A}, ignored) where {A<:Unsigned}
    for i in 1:length(data)
        word = data[i]
        word ⊻= remainder%UInt8
        remainder = (remainder >>> 8) ⊻ table[1 + word]
    end
    remainder
end

# timing with a standard data set (the table is "meaningless")

Random.seed!(0)  # repeatable results
TABLE = UInt64[rand(UInt64) for _ in 1:256]
DATA = UInt8[rand(UInt8) for _ in 1:10_000_000]
#KNOWN = 17880998858342525629  # 1 million
KNOWN = 0x65280cfa186b00d5  # 10 million
# KNOWN = 10790464939549839568  # 100 million
using Printf
function run_std(loop)
    loop(UInt8, zero(UInt64), DATA[1:10], TABLE, 8)  # warmup
    @printf("%15s: ", loop)
    @time (result = loop(UInt8, zero(UInt64), DATA, TABLE, 8))
    @assert result == KNOWN
end

for loop in (generic_loop, byte_loop, array_loop)
    run_std(loop)
end
run_std(array_loop)

#  generic_loop: elapsed time: 0.023491408 seconds (16 bytes allocated)
#   struct_loop: elapsed time: 2.027252231 seconds (1680000168 bytes allocated)
#     byte_loop: elapsed time: 0.02304929 seconds (16 bytes allocated)
#    array_loop: elapsed time: 0.045178428 seconds (168 bytes allocated)

# benchmarks against Libz

SIZE = 300_000_000
# SIZE = 300_000

function time_libz()
    println("libz")
    data = rand(UInt8, SIZE)
    check = crc32(data)
    @time crc32(data)
    for tables in (Single, Multiple)
        ours = crc(CRC_32, tables=tables)
        @assert ours(data) == check
        println(tables)
        @time ours(data)
    end
end

function time_no_tables()
    println("no_tables")
    ours = crc(CRC_15, tables=NoTables)
    data = rand(UInt8, round(Int, SIZE//10))
    @assert ours(CHECK) == CRC_15.check
    @time ours(data)
end

function time_padded()
    println("padded")
    ours = crc(CRC_64)
    data = rand(UInt8, SIZE)
    @assert ours(CHECK) == CRC_64.check
    @time ours(data)
end

Random.seed!(0)  # repeatable results

time_libz()
time_no_tables()
time_padded()


module PIPplot

export PIPfullplot

using JuLIP, NBodyIPs, DataFrames, Plots, JLD, StaticArrays
using NBodyIPs.Data: Dat
using Plots; gr()

using PyCall
unshift!(PyVector(pyimport("sys")["path"]), Pkg.dir()*"/PIPplot")

@pyimport plottools
@pyimport numpy
@pyimport matplotlib.pyplot as plt

function unfold(A)
    V = []
    for x in A
        if x === A
            push!(V, x)
        else
            append!(V, unfold(x))
        end
    end
    V
end

function PIPpotplot(potname, IP)
    e0 = string(IP.orders[1].c)
    cutoff2b = string(IP.orders[3].D.rcut)
    cutoff3b = string(IP.orders[4].D.rcut)
    cutoff4b = string(IP.orders[5].D.rcut)
	r0 = 2.89
	rr = linspace(2, 9.0, 200)
	P = plot(rr, IP.(rr); yaxis = ([-0.1, 0.2],), label="2b");
	plot!(P, rr, IP.(rr, rr, rr), label="3b")
	W4 = rr -> IP(SVector(2.9, rr, 2.9, 2.9, 2.9, 2.9))
	title = "Potname=" * potname[3:end-4] * " rc=(" * cutoff2b * "/" * cutoff3b * "/" * cutoff4b * ")"*" e0=" * e0
	plot!(P, rr, W4.(rr), title=title, xlabel="Interatomic distance (Å)", titlefont=8, label="4b", legend=true)
    vline!([2.89], label="r0", color="black")
end

function PIPenergyplot(test_data, IP, s; e_data = 1, returnplot = true)
    d=Dict([(i,test_data[i].config_type) for i in 1:length(test_data)])
    configselection = [[i,d[i]] for i in 1:s:length(test_data)]
    configtypes = unique([configselection[i][2] for i in 1:length(configselection)])
    if e_data == 1
        e_data = hcat([[test_data[configselection[i][1]].E, energy(IP, Atoms(test_data[configselection[i][1]])), configselection[i][2]] for i in 1:length(configselection)]...)
    end
    RMSE_E = sqrt(mean((e_data[1,:]-e_data[2,:]).^2))/54
    if returnplot == true
        xmin = findmin(convert(Array{Float64,1},unfold(e_data[1,:])))[1]/54
        xmax = findmax(convert(Array{Float64,1},unfold(e_data[1,:])))[1]/54
        x = linspace(xmin, xmax, 1000)
        f(x) = x#(abs.(ymin - ymax)/abs.(xmin-xmax))x
        p = plot(x, f(x), color="black", legend=false)
        colors = ["red", "green", "blue", "yellow", "purple", "brown", "cyan", "violet"]
        for i in 1:length(configtypes)
            config = configtypes[i]
            indices = find(e_data[3, 1:end] .== config)
            target = convert(Array{Float64,1}, unfold(e_data[1,findmin(indices)[1]:findmax(indices)[1]]))/54
            test = convert(Array{Float64,1}, unfold(e_data[2,findmin(indices)[1]:findmax(indices)[1]]))/54
            data = transpose(numpy.array([target, test], numpy.float64))
            thin_data, thin_weights = plottools.thin_points(data, r=0.001)
            p = scatter!(p, thin_data[:,1], thin_data[:,2], markersize=4*(thin_weights).^(0.5), markeralpha=0.1, label=config, legend=:bottomright, xlabel="Target Energy per atom (eV) RMSE = "*string(RMSE_E)[1:8], ylabel="Predicted Energy per atom (eV)", legendfont=7, color=colors[i])  ##yaxis = (:log10, (1,100000)), markerstrokealpha=1, markerstrokecolor="black",
        end
        return p
    else
        return e_data
    end
end

function PIPforceplot(test_data, IP, s; f_data = 1, returnplot=true)
    d=Dict([(i,test_data[i].config_type) for i in 1:length(test_data)])
    configselection = [[i,d[i]] for i in 1:s:length(test_data)] ###which different ones are there
    configtypes = unique([configselection[i][2] for i in 1:length(configselection)])
    if f_data == 1
        f_data = hcat([[unfold(test_data[configselection[i][1]].F), unfold(forces(IP, Atoms(test_data[configselection[i][1]]))), configselection[i][2]] for i in 1:length(configselection)]...)
    end
    RMSE_F = sqrt(mean((unfold(f_data[1,:]) - unfold(f_data[2,:])).^2))
    if returnplot == true
        xmin = findmin(convert(Array{Float64,1},unfold(f_data[1,:])))[1]
        xmax = findmax(convert(Array{Float64,1},unfold(f_data[1,:])))[1]
        x = linspace(xmin, xmax, 1000)
        f(x) = x
        p = plot(x, f(x), color="black", legend=false)
        colors = ["red", "green", "blue", "yellow", "purple", "brown", "cyan", "violet"]
        for i in 1:length(configtypes)
            config = configtypes[i]
            indices = find(f_data[3, 1:end] .== config)
            target = convert(Array{Float64,1}, unfold(f_data[1,findmin(indices)[1]:findmax(indices)[1]]))
            test = convert(Array{Float64,1}, unfold(f_data[2,findmin(indices)[1]:findmax(indices)[1]]))
            data = transpose(numpy.array([test, target], numpy.float64))
            thin_data, thin_weights = plottools.thin_points(data, r=0.02)
            p = scatter!(p, thin_data[:,1], thin_data[:,2], markersize=4*(thin_weights).^(0.5), markeralpha=0.1, label=config, legend=:bottomright, xlabel="Target Force (ev/Å) RMSE = "*string(RMSE_F)[1:8], ylabel="Predicted Force (ev/Å)", legendfont=7, color=colors[i], legend=:false)  ##markerstrokealpha=1, markerstrokecolor="black",
        end
        return p
    else
        return f_data
    end
end

function PIP3Body3D(potname, IP; p2b = 2.25, p2e = 4.5, p2s = 50, theta = 60)
	V3 = IP
	W3 = (r1, r2, r3) -> V3(SVector(r1, r2, r3))
	elist = []
	R = []
	for r in linspace(p2b,p2e,p2s)
	    for r12 in linspace(p2b,p2e,p2s)
	        x = [r12*cosd(theta), r12*sind(theta)]
	        y = [r, 0]
	        z = x - y
	        push!(R, [norm(x),norm(y),norm(z)])
	    end
	end
	for i in 1:length(R)
		e = W3.(R[i][1], R[i][2], R[i][3])
		push!(elist, e)
	end
	Z = Matrix(0,p2s)
	for i in 1:p2s:(p2s*p2s)
	    v = transpose(vec(elist[i:i+(50-1)]))
	    Z = cat(1,Z,v)
	end
	X, Y = linspace(p2b,p2e,p2s), linspace(p2b,p2e,p2s)
	Z = convert(Array{Float64,2},Z)
	heatmap(X,Y,Z, xlabel="Interatomic distance r12 (Å)", ylabel="Interatomic distance r13 (Å)")
    plot!([X[1], X[end]], [Y[1], Y[end]], color="white", legend=false)
end

function PIPforceplot2(test_data, IP, s; f_data = 1, returnplot=true)
    d=Dict([(i,test_data[i].config_type) for i in 1:length(test_data)])
    configselection = [[i,d[i]] for i in 1:s:length(test_data)] ###which different ones are there
    configtypes = unique([configselection[i][2] for i in 1:length(configselection)])
    if f_data == 1
        f_data = hcat([[unfold(test_data[configselection[i][1]].F), unfold(forces(IP, Atoms(test_data[configselection[i][1]]))), configselection[i][2]] for i in 1:length(configselection)]...)
    end
    RMSE_F = sqrt(mean((unfold(f_data[1,:]) - unfold(f_data[2,:])).^2))
    if returnplot == true
        xmin = findmin(convert(Array{Float64,1},unfold(f_data[1,:])))[1]
        xmax = findmax(convert(Array{Float64,1},unfold(f_data[1,:])))[1]
        ymin = findmin(convert(Array{Float64,1},unfold(f_data[2,:])))[1]
        ymax = findmax(convert(Array{Float64,1},unfold(f_data[2,:])))[1]
        p = plot()
        colors = ["red", "green", "blue", "yellow", "purple", "brown", "cyan", "violet"]
        for i in 1:length(configtypes)
            config = configtypes[i]
            indices = find(f_data[3, 1:end] .== config)
            target = convert(Array{Float64,1}, unfold(f_data[1,findmin(indices)[1]:findmax(indices)[1]]))
            test = convert(Array{Float64,1}, unfold(f_data[2,findmin(indices)[1]:findmax(indices)[1]]))
            error = abs.(test-target)
            data = transpose(numpy.array([target, error], numpy.float64))
            thin_data, thin_weights = plottools.thin_points(data, r=0.02)
            p = scatter!(p, thin_data[:,1], thin_data[:,2]+1e-10, markersize=4*(thin_weights).^(0.5), markeralpha=0.1, label=config, legend=:bottomright, xlabel="Target Force (ev/Å) RMSE = "*string(RMSE_F)[1:8], ylabel="| Predicted Force Error | (ev/Å)", yaxis = (:log10, (0.00005,2)), legendfont=7, color=colors[i], legend=:false)  ##markerstrokealpha=1, markerstrokecolor="black",
        end
        return p
    else
        return f_data
    end
end

function PIPenergyplot2(test_data, IP, s; e_data = 1, returnplot = true)
    d=Dict([(i,test_data[i].config_type) for i in 1:length(test_data)])
    configselection = [[i,d[i]] for i in 1:s:length(test_data)]
    configtypes = unique([configselection[i][2] for i in 1:length(configselection)])
    if e_data == 1
        e_data = hcat([[test_data[configselection[i][1]].E, energy(IP, Atoms(test_data[configselection[i][1]])), configselection[i][2]] for i in 1:length(configselection)]...)
    end
    RMSE_E = sqrt(mean((e_data[1,:]-e_data[2,:]).^2))/54
    if returnplot == true
        p = plot()
        colors = ["red", "green", "blue", "yellow", "purple", "brown", "cyan", "violet"]
        for i in 1:length(configtypes)
            config = configtypes[i]
            indices = find(e_data[3, 1:end] .== config)
            target = convert(Array{Float64,1}, unfold(e_data[1,findmin(indices)[1]:findmax(indices)[1]]))/54
            test = convert(Array{Float64,1}, unfold(e_data[2,findmin(indices)[1]:findmax(indices)[1]]))/54
            error = abs.(test-target)
            data = transpose(numpy.array([target, error], numpy.float64))
            thin_data, thin_weights = plottools.thin_points(data, r=0.01)
            p = scatter!(p, thin_data[:,1], thin_data[:,2]+1e-10, markersize=4*(thin_weights).^(0.5), markeralpha=0.1, label=config, legend=:bottomright, xlabel="Target Energy per atom (eV) RMSE = "*string(RMSE_E)[1:8], yaxis = (:log10, (0.00005,0.04)), ylabel="Predicted Energy per atom (eV)", legendfont=7, color=colors[i], legend=:false)  ##yaxis = (:log10, (1,100000)), markerstrokealpha=1, markerstrokecolor="black",
        end
        return p
    else
        return e_data
    end
end

function PIPfullplot(potname, test_data, IP; s=10)
    if isfile("./data_"*potname[3:end-4]*".jld") == true
        file = load("./data_"*potname[3:end-4]".jld")
        e_data = file["e_data"]
        f_data = file["f_data"]
    else
        e_data = PIPenergyplot(test_data, IP, s, returnplot=false)
        f_data = PIPforceplot(test_data, IP, s, returnplot=false)
        file = jldopen("./data_"*potname[3:end-4]".jld", "w")
        write(file, "e_data", e_data)
        write(file, "f_data", f_data)
        close(file)
    end
    p1 = PIPpotplot(potname, IP)
    p2 = PIP3Body3D(potname, IP)
    p3 = PIPenergyplot(test_data, IP, e_data=e_data, s)
    p4 = PIPforceplot(test_data, IP, f_data=f_data, s)
    p5 = PIPenergyplot2(test_data, IP, e_data=e_data, s)
    p6 = PIPforceplot2(test_data, IP, f_data=f_data, s)
    p7 = plot(p1,p2,p3,p4,p5,p6, layout=(3,2), size=(900,1500))
    #png(p7, "plot")
end

end # module


struct MLMPCSolver <: Solver
    solver
    updater::Nullable{Any}
end
function set_rng!(s::MLMPCSolver, rng::AbstractRNG)
    set_rng!(s.solver, rng)
    set_rng!(get(s.updater), rng)
end

mutable struct MLMPCAgent <: Policy
    updater::Updater
    previous_belief::Nullable{Any}
    policy::POMDPs.Policy
    previous_action::Nullable{MLAction}
end

function solve(solver::MLMPCSolver, problem::MLMDP, up=get(solver.updater, nothing))
    internal_problem = MLPOMDP{MLState, MLAction, MLPhysicalState, typeof(problem.dmodel), typeof(problem.rmodel)}(problem.dmodel, problem.rmodel, problem.discount)
    policy = solve(solver.solver, problem)
#=
    if up == nothing
        up = POMDPs.updater(policy)
    else
        set_problem!(up, internal_problem)
    end
=#

    set_problem!(up, internal_problem)
    #up.problem = internal_problem
    return MLMPCAgent(up, nothing, policy, nothing)
    #return up, policy
    #return MLMPCAgent(policy)
end

function action(agent::MLMPCAgent, state::MLState)
    o = MLPhysicalState(state)
    if isnull(agent.previous_belief)
        belief = initialize_belief(agent.updater, ParticleGenerator(agent.policy.mdp, state))
    else
        belief = update(agent.updater, get(agent.previous_belief), get(agent.previous_action), o)
    end
    ml = most_likely_state(belief)
    a = action(agent.policy, ml)
    agent.previous_action = a
    agent.previous_belief = belief
    return a
end

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

################################heuristics

struct Simple <: Policy #
  mdp
  A::NoCrashActionSpace
  sweeping_up::Bool
end
type SimpleSolver <: Solver end

Simple(mdp) = Simple(mdp,actions(mdp),true)
solve(solver::SimpleSolver, problem::MDP) = Simple(problem)
solve(solver::SimpleSolver, problem::EmbeddedBehaviorMDP) = Simple(problem.base)
solve(solver::SimpleSolver, problem::POMDP) = Simple(problem)
POMDPs.updater(::Simple) = POMDPToolbox.FastPreviousObservationUpdater{MLObs}()
create_policy(s::SimpleSolver, problem::MDP) = Simple(problem)
create_policy(s::SimpleSolver, problem::POMDP) = Simple(problem)

set_rng!(solver::SimpleSolver, rng::AbstractRNG) = nothing

function action(p::Simple,s::Union{MLState,MLObs},a::MLAction=create_action(p.mdp))
# lane changes if there is an opportunity
    goal_lane = p.mdp.rmodel.target_lane
    y_desired = goal_lane
    dmodel = p.mdp.dmodel
    lc = sign(y_desired-s.cars[1].y) * dmodel.lane_change_rate
    acc = dmodel.adjustment_acceleration

    #if can't move towards desired lane sweep through accelerating and decelerating

    # TODO need an equivalent of is_safe that can operate on observations
    if is_safe(p.mdp, s, MLAction(0.,lc))
        return MLAction(0.,lc)
    end
    # maintain distance from other cars

    # maintain distance
    nbhd = get_neighborhood(dmodel.phys_param,s,1)

    if nbhd[2] == 0 && nbhd[5] == 0
        return MLAction(0.,0.)
    end

    dist_ahead = nbhd[2] != 0 ? s.cars[nbhd[2]].x - s.cars[1].x : Inf
    dist_behind = nbhd[5] != 0 ? s.cars[nbhd[5]].x - s.cars[1].x : Inf

    sgn = abs(dist_ahead) <= abs(dist_behind) ? -1 : 1

    accel = sgn * acc

    max_accel = max_safe_acc(p.mdp, s, 0.0)

    return MLAction(min(accel, max_accel),0.)
end
#action(p::Simple, b::BehaviorBelief, a::MLAction=create_action(p.mdp)) = action(p, b.physical, a)

struct BehaviorSolver <: Solver
    b::BehaviorModel
    keep_lane::Bool
    rng::AbstractRNG
end
struct BehaviorPolicy <: Policy
    problem::NoCrashProblem
    b::BehaviorModel
    keep_lane::Bool
    rng::AbstractRNG
end
solve(s::BehaviorSolver, p::NoCrashProblem) = BehaviorPolicy(p, s.b, s.keep_lane, s.rng)

function action(p::BehaviorPolicy, s::MLState, a::MLAction=MLAction(0.0,0.0))
    #println("LC")
    nbhd = get_neighborhood(p.problem.dmodel.phys_param, s, 1)
    acc = gen_accel(p.b, p.problem.dmodel, s, nbhd, 1, p.rng)
    if p.keep_lane
        lc = 0.0
    else
        #lc = 1.0

        lc = gen_lane_change(p.b, p.problem.dmodel, s, nbhd, 1, p.rng)
    end
    return MLAction(acc, lc)
end
#action(p::BehaviorPolicy, b::AggressivenessBelief, a::MLAction=MLAction(0.0,0.0)) = action(p, most_likely_state(b))
#action(p::BehaviorPolicy, b::BehaviorParticleBelief, a::MLAction=MLAction(0.0,0.0)) = action(p, most_likely_state(b))

type IDMLaneSeekingSolver <: Solver
    b::BehaviorModel
    rng::AbstractRNG
end

type IDMLaneSeekingPolicy <: Policy
    problem::NoCrashProblem
    b::BehaviorModel
    rng::AbstractRNG
end
solve(s::IDMLaneSeekingSolver, p::NoCrashProblem) = IDMLaneSeekingPolicy(p, s.b, s.rng)

function action(p::IDMLaneSeekingPolicy, s::MLState, a::MLAction=MLAction(0.0,0.0))
    nbhd = get_neighborhood(p.problem.dmodel.phys_param, s, 1)
    acc = gen_accel(p.b, p.problem.dmodel, s, nbhd, 1, p.rng)
    # try to positive lanechange
    # lc = problem.dmodel.lane_change_rate * !is_lanechange_dangerous(pp,s,nbhd,1,1)
    lc = p.problem.dmodel.lane_change_rate
    if is_safe(p.problem, s, MLAction(acc, lc))
        return MLAction(acc, lc)
    end
    return MLAction(acc, 0.0)
end


module AutomotiveDrivingModels

using Printf
using LinearAlgebra
using Parameters
using StaticArrays
using Distributions
using Reexport
using Random
using DataFrames
using Tricks: static_hasmethod

include(joinpath(@__DIR__, "Vec", "Vec.jl"))
@reexport using .Vec

# Roadways

export CurvePt,
       Curve,
       CurveIndex,
       CurveProjection,
       CURVEINDEX_START,
       get_lerp_time,
       index_closest_to_point,
       get_curve_index,
       curveindex_end

include("roadways/curves.jl")

export
    LaneTag,
    LaneBoundary,
    Lane,
    LaneConnection,
    SpeedLimit,
    RoadSegment,
    Roadway,
    RoadProjection,
    RoadIndex,

    NULL_BOUNDARY,
    NULL_LANETAG,
    NULL_ROADINDEX,

    DEFAULT_SPEED_LIMIT,
    DEFAULT_LANE_WIDTH,

    is_in_exits,
    is_in_entrances,
    is_at_curve_end,
    is_between_segments_lo,
    is_between_segments_hi,
    is_between_segments,
    has_segment,
    has_lanetag,
    has_next,
    has_prev,
    next_lane,
    prev_lane,
    next_lane_point,
    prev_lane_point,
    connect!,
    move_along,
    n_lanes_left,
    leftlane,
    n_lanes_right,
    rightlane,
    lanes,
    lanetags,
    read

include("roadways/roadways.jl")

export
    Frenet,
    NULL_FRENET

include("roadways/frenet.jl")

export
        gen_straight_curve,
        gen_straight_segment,
        gen_straight_roadway,
        gen_stadium_roadway,
        gen_bezier_curve

include("roadways/roadway_generation.jl")

## agent definitions

export
    AbstractAgentDefinition,
    AgentClass,
    VehicleDef,
    BicycleModel,
    length,
    width,
    class

include("agent-definitions/agent_definitions.jl")

export
    Entity,
    Scene,
    EntityScene,
    capacity,
    id2index,
    get_by_id,
    get_first_available_id,
    posf,
    posg,
    vel,
    velf,
    velg,
    VehicleState,
    CustomVehicleState,
    get_center,
    get_footpoint,
    get_front,
    get_rear,
    get_lane

include("states/entities.jl")
include("states/scenes.jl")
include("states/interface.jl")
include("states/vehicle_state.jl")

## Collision Checkers

export
    ConvexPolygon,
    CPAMemory,
    CollisionCheckResult,
    to_oriented_bounding_box!,
    get_oriented_bounding_box,
    is_colliding,
    is_potentially_colliding,
    get_collision_time,
    get_first_collision,
    is_collision_free,
    get_distance,
    get_edge,
    collision_checker,
    polygon

include("collision-checkers/minkowski.jl")
include("collision-checkers/parallel_axis.jl")

## Feature Extraction

export AbstractFeature,
       EntityFeature,
       SceneFeature,
       TemporalFeature,
       extract_features,
       extract_feature,
       featuretype,

       # provided feature functions
       posgx,
       posgy,
       posgθ,
       posfs,
       posft,
       posfϕ,
       vel,
       velfs,
       velft,
       velgx,
       velgy,
       time_to_crossing_right,
       time_to_crossing_left,
       estimated_time_to_lane_crossing,
       iswaiting,
       iscolliding,
       distance_to,
       acc,
       accfs,
       accft,
       jerk,
       jerkft,
       turn_rate_g,
       turn_rate_f,
       isbraking,
       isaccelerating

include("feature-extraction/features.jl")

export lane_width,
       markerdist_left,
       markerdist_right,
       road_edge_dist_left,
       road_edge_dist_right,
       lane_offset_left,
       lane_offset_right,
       has_lane_left,
       has_lane_right,
       lane_curvature

include("feature-extraction/lane_features.jl")

export
    VehicleTargetPoint,
    VehicleTargetPointFront,
    VehicleTargetPointCenter,
    VehicleTargetPointRear,
    targetpoint_delta,
    find_neighbor,
    NeighborLongitudinalResult,
    FrenetRelativePosition,
    get_frenet_relative_position,
    dist_to_front_neighbor,
    front_neighbor_speed,
    time_to_collision

include("feature-extraction/neighbors_features.jl")

export
    LidarSensor,
    nbeams,
    observe!,
    RoadlineLidarSensor,
    nlanes,
    LanePortion,
    RoadwayLidarCulling,
    ensure_leaf_in_rlc!,
    get_lane_portions

include("feature-extraction/lidar_sensor.jl")

## Actions

export
    propagate,
    EntityAction,
    LaneFollowingAccel,
    AccelTurnrate,
    AccelDesang,
    LatLonAccel,
    AccelSteeringAngle,
    PedestrianLatLonAccel

include("actions/interface.jl")
include("actions/lane_following_accel.jl")
include("actions/accel_turnrate.jl")
include("actions/accel_desang.jl")
include("actions/lat_lon_accel.jl")
include("actions/accel_steering.jl")
include("actions/pedestrian_lat_lon_accel.jl")

export
    DriverModel,
    StaticDriver,
    action_type,
    set_desired_speed!,
    observe!,
    reset_hidden_state!,
    reset_hidden_states!

include("behaviors/interface.jl")

export
    LaneFollowingDriver,
    StaticLaneFollowingDriver,
    PrincetonDriver,
    IntelligentDriverModel,
    ProportionalSpeedTracker,
    track_longitudinal!,
    LateralDriverModel,
    ProportionalLaneTracker,
    LatLonSeparableDriver,
    Tim2DDriver,
    SidewalkPedestrianModel,
    LaneChangeChoice,
    LaneChangeModel,
    get_lane_offset,
    DIR_RIGHT,
    DIR_MIDDLE,
    DIR_LEFT,
    MOBIL,
    TimLaneChanger

include("behaviors/lane_following_drivers.jl")
include("behaviors/princeton_driver.jl")
include("behaviors/speed_trackers.jl")
include("behaviors/intelligent_driver_model.jl")
include("behaviors/lateral_driver_models.jl")
include("behaviors/lane_change_models.jl")
include("behaviors/MOBIL.jl")
include("behaviors/tim_lane_changer.jl")
include("behaviors/lat_lon_separable_driver.jl")
include("behaviors/tim_2d_driver.jl")
include("behaviors/sidewalk_pedestrian_model.jl")

export
    simulate,
    simulate!,
    run_callback,
    CollisionCallback,
    observe_from_history!,
    simulate_from_history!,
    simulate_from_history,
    maximum_entities

include("simulation/simulation.jl")
include("simulation/callbacks.jl")
include("simulation/simulation_from_history.jl")

end # AutomotiveDrivingModels