text:
#
# 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


text:
"""
     foreach_setup(z)
Return `Channel` for `foreach` threaded computation from iterable `z`.
"""
function foreach_setup(z)
    return Channel{Int}(length(z)) do ch
        foreach(i -> put!(ch, i), z)
    end
end

ext: jl


text:
#  BLOSUM Clustered Target Frequencies=qij
#  Blocks Database = /data/blocks_5.0/blocks.dat
#  Cluster Percentage: >= 62

"""
BLOSUM62 probabilities *P(aa)* for each residue on the `UngappedAlphabet`.
SUM:  0.9987
"""
const BLOSUM62_Pi = ContingencyTable(Float64[ 0.0741,
    0.0516,
    0.0445,
    0.0536,
    0.0246,
    0.034,
    0.0543,
    0.0741,
    0.0262,
    0.0679,
    0.0988,
    0.0579,
    0.0249,
    0.0473,
    0.0387,
    0.0573,
    0.0507,
    0.0132,
    0.0321,
    0.0729 ], UngappedAlphabet())

"""
Table with conditional probabilities of residues based on BLOSUM62. The normalization is
done row based. The firts row contains the *P(aa|A)* and so one.
"""
const BLOSUM62_Pij = ContingencyTable(Float64[
#    A      R      N      D      C      Q      E      G      H      I      L      K      M      F      P      S      T      W      Y      V
0.2901 0.0310 0.0256 0.0297 0.0216 0.0256 0.0405 0.0783 0.0148 0.0432 0.0594 0.0445 0.0175 0.0216 0.0297 0.0850 0.0499 0.0054 0.0175 0.0688 ;
0.0446 0.3450 0.0388 0.0310 0.0078 0.0484 0.0523 0.0329 0.0233 0.0233 0.0465 0.1202 0.0155 0.0174 0.0194 0.0446 0.0349 0.0058 0.0174 0.0310 ;
0.0427 0.0449 0.3169 0.0831 0.0090 0.0337 0.0494 0.0652 0.0315 0.0225 0.0315 0.0539 0.0112 0.0180 0.0202 0.0697 0.0494 0.0045 0.0157 0.0270 ;
0.0410 0.0299 0.0690 0.3974 0.0075 0.0299 0.0914 0.0466 0.0187 0.0224 0.0280 0.0448 0.0093 0.0149 0.0224 0.0522 0.0354 0.0037 0.0112 0.0243 ;
0.0650 0.0163 0.0163 0.0163 0.4837 0.0122 0.0163 0.0325 0.0081 0.0447 0.0650 0.0203 0.0163 0.0203 0.0163 0.0407 0.0366 0.0041 0.0122 0.0569 ;
0.0559 0.0735 0.0441 0.0471 0.0088 0.2147 0.1029 0.0412 0.0294 0.0265 0.0471 0.0912 0.0206 0.0147 0.0235 0.0559 0.0412 0.0059 0.0206 0.0353 ;
0.0552 0.0497 0.0405 0.0902 0.0074 0.0645 0.2965 0.0350 0.0258 0.0221 0.0368 0.0755 0.0129 0.0166 0.0258 0.0552 0.0368 0.0055 0.0166 0.0313 ;
0.0783 0.0229 0.0391 0.0337 0.0108 0.0189 0.0256 0.5101 0.0135 0.0189 0.0283 0.0337 0.0094 0.0162 0.0189 0.0513 0.0297 0.0054 0.0108 0.0243 ;
0.0420 0.0458 0.0534 0.0382 0.0076 0.0382 0.0534 0.0382 0.3550 0.0229 0.0382 0.0458 0.0153 0.0305 0.0191 0.0420 0.0267 0.0076 0.0573 0.0229 ;
0.0471 0.0177 0.0147 0.0177 0.0162 0.0133 0.0177 0.0206 0.0088 0.2710 0.1679 0.0236 0.0368 0.0442 0.0147 0.0250 0.0398 0.0059 0.0206 0.1767 ;
0.0445 0.0243 0.0142 0.0152 0.0162 0.0162 0.0202 0.0213 0.0101 0.1154 0.3755 0.0253 0.0496 0.0547 0.0142 0.0243 0.0334 0.0071 0.0223 0.0962 ;
0.0570 0.1071 0.0415 0.0415 0.0086 0.0535 0.0708 0.0432 0.0207 0.0276 0.0432 0.2781 0.0155 0.0155 0.0276 0.0535 0.0397 0.0052 0.0173 0.0328 ;
0.0522 0.0321 0.0201 0.0201 0.0161 0.0281 0.0281 0.0281 0.0161 0.1004 0.1968 0.0361 0.1606 0.0482 0.0161 0.0361 0.0402 0.0080 0.0241 0.0924 ;
0.0338 0.0190 0.0169 0.0169 0.0106 0.0106 0.0190 0.0254 0.0169 0.0634 0.1142 0.0190 0.0254 0.3869 0.0106 0.0254 0.0254 0.0169 0.0888 0.0550 ;
0.0568 0.0258 0.0233 0.0310 0.0103 0.0207 0.0362 0.0362 0.0129 0.0258 0.0362 0.0413 0.0103 0.0129 0.4935 0.0439 0.0362 0.0026 0.0129 0.0310 ;
0.1099 0.0401 0.0541 0.0489 0.0175 0.0332 0.0524 0.0663 0.0192 0.0297 0.0419 0.0541 0.0157 0.0209 0.0297 0.2199 0.0820 0.0052 0.0175 0.0419 ;
0.0730 0.0355 0.0434 0.0375 0.0178 0.0276 0.0394 0.0434 0.0138 0.0533 0.0651 0.0454 0.0197 0.0237 0.0276 0.0927 0.2465 0.0059 0.0178 0.0710 ;
0.0303 0.0227 0.0152 0.0152 0.0076 0.0152 0.0227 0.0303 0.0152 0.0303 0.0530 0.0227 0.0152 0.0606 0.0076 0.0227 0.0227 0.4924 0.0682 0.0303 ;
0.0405 0.0280 0.0218 0.0187 0.0093 0.0218 0.0280 0.0249 0.0467 0.0436 0.0685 0.0312 0.0187 0.1308 0.0156 0.0312 0.0280 0.0280 0.3178 0.0467 ;
0.0700 0.0219 0.0165 0.0178 0.0192 0.0165 0.0233 0.0247 0.0082 0.1646 0.1303 0.0261 0.0316 0.0357 0.0165 0.0329 0.0494 0.0055 0.0206 0.2689 ],
UngappedAlphabet())


text:
using GeometricBase
using GeometricBase.Config
using GeometricIntegrators.Integrators
using GeometricIntegrators.Tableaus
using GeometricProblems.LotkaVolterra2d
using Test

using GeometricEquations: _get_v, _get_f, _get_v̄, _get_f̄

int = get_config(:ig_extrapolation)

Δt = 0.01
q₀ = [1.0, 1.0]
parameters = (a₁=1.0, a₂=1.0, b₁=-1.0, b₂=-2.0)

ode  = lotka_volterra_2d_ode(q₀; params=parameters)
pode = lotka_volterra_2d_pode(q₀; params=parameters)
hode = lotka_volterra_2d_hode(q₀; params=parameters)
iode = lotka_volterra_2d_iode(q₀; params=parameters)
lode = lotka_volterra_2d_lode(q₀; params=parameters)

dae  = lotka_volterra_2d_dae(q₀; params=parameters)
pdae = lotka_volterra_2d_pdae(q₀; params=parameters)
hdae = lotka_volterra_2d_hdae(q₀; params=parameters)
idae = lotka_volterra_2d_idae(q₀; params=parameters)
ldae = lotka_volterra_2d_ldae(q₀; params=parameters)

@test InitialGuessODE(int,  ode, Δt) == InitialGuessODE(int(0.0, Δt), _get_v̄( ode), Δt)
@test InitialGuessODE(int, iode, Δt) == InitialGuessODE(int(0.0, Δt), _get_v̄(iode), Δt)
@test InitialGuessODE(int, lode, Δt) == InitialGuessODE(int(0.0, Δt), _get_v̄(lode), Δt)

@test InitialGuess(int, ode, Δt) == InitialGuessODE(int(0.0, Δt), _get_v̄(ode), Δt)
@test InitialGuess(int, dae, Δt) == InitialGuessODE(int(0.0, Δt), _get_v̄(dae), Δt)

@test InitialGuess(int, hode, Δt) == InitialGuessPODE(int(0.0, Δt), _get_v̄(hode), _get_f̄(hode), Δt)
@test InitialGuess(int, hdae, Δt) == InitialGuessPODE(int(0.0, Δt), _get_v̄(hdae), _get_f̄(hdae), Δt)

@test InitialGuess(int, iode, Δt) == InitialGuessIODE(int(0.0, Δt), _get_v̄(iode), _get_f̄(iode), Δt)
@test InitialGuess(int, idae, Δt) == InitialGuessIODE(int(0.0, Δt), _get_v̄(idae), _get_f̄(idae), Δt)

@test InitialGuess(int, pode, Δt) == InitialGuessPODE(int(0.0, Δt), _get_v̄(pode), _get_f̄(pode), Δt)
@test InitialGuess(int, pdae, Δt) == InitialGuessPODE(int(0.0, Δt), _get_v̄(pdae), _get_f̄(pdae), Δt)

@test InitialGuess(int, lode, Δt) == InitialGuessIODE(int(0.0, Δt), _get_v̄(lode), _get_f̄(lode), Δt)
@test InitialGuess(int, ldae, Δt) == InitialGuessIODE(int(0.0, Δt), _get_v̄(ldae), _get_f̄(ldae), Δt)

# Reference Solution

ref_prev = integrate(ode, TableauGauss(8), -Δt, 1)
ref_next = integrate(ode, TableauGauss(8), +Δt, 1)

tₚ = ref_prev.t[end]
qₚ = ref_prev.q[end]
vₚ = zero(qₚ)

tₙ = ref_next.t[end]
qₙ = ref_next.q[end]
vₙ = zero(qₙ)

ode.v(tₚ, qₚ, vₚ, ode.parameters)
ode.v(tₙ, qₙ, vₙ, ode.parameters)

# InitialGuessODE

igode = InitialGuessODE(int, _get_v(ode), Δt)

t₀ = ode.t₀
q₀ = ode.q₀[begin]
v₀ = zero(q₀)

t₁ = ode.t₀ - Δt
q₁ = zero(q₀)
v₁ = zero(v₀)

t₂ = one(Δt)
q₂ = zero(q₀)
v₂ = zero(v₀)

ode.v(t₀, q₀, v₀, ode.parameters)

initialize!(igode, t₀, q₀, v₀, t₁, q₁, v₁)
evaluate!(igode, q₁, v₁, q₀, v₀, q₂, v₂, t₂)

# println("IG-ODE")
# println(q₁ .- qₚ)
# println(v₁ .- vₚ)
# println(q₂ .- qₙ)
# println(v₂ .- vₙ)
# println()

@test q₁ ≈ qₚ atol=1E-14
@test v₁ ≈ vₚ atol=1E-14

@test q₂ ≈ qₙ atol=1E-8
@test v₂ ≈ vₙ atol=1E-5

# InitialGuessIODE

igiode = InitialGuessIODE(int, _get_v̄(iode), _get_f̄(iode), Δt)

t₀ = iode.t₀
q₀ = iode.q₀[begin]
p₀ = iode.p₀[begin]
v₀ = zero(q₀)
f₀ = zero(p₀)

t₁ = iode.t₀ - Δt
q₁ = zero(q₀)
p₁ = zero(p₀)
v₁ = zero(v₀)
f₁ = zero(f₀)

t₂ = one(Δt)
q₂ = zero(q₀)
p₂ = zero(p₀)
v₂ = zero(v₀)
f₂ = zero(f₀)

iode.v̄(t₀, q₀, v₀, iode.parameters)

initialize!(igiode, t₀, q₀, p₀, v₀, f₀, t₁, q₁, p₁, v₁, f₁)
evaluate!(igiode, q₁, p₁, v₁, f₁, q₀, p₀, v₀, f₀, q₂, v₂, t₂)

# println("IG-IODE")
# println(q₁ .- qₚ)
# println(v₁ .- vₚ)
# println(q₂ .- qₙ)
# println(v₂ .- vₙ)
# println()

@test q₁ ≈ qₚ atol=1E-14
@test v₁ ≈ vₚ atol=1E-14

@test q₂ ≈ qₙ atol=1E-8
@test v₂ ≈ vₙ atol=1E-5

# InitialGuessPODE

igpode = InitialGuessPODE(int, _get_v(pode), _get_f(pode), Δt)

t₀ = pode.t₀
q₀ = pode.q₀[begin]
p₀ = pode.p₀[begin]
v₀ = zero(q₀)
f₀ = zero(p₀)

t₁ = pode.t₀ - Δt
q₁ = zero(q₀)
p₁ = zero(p₀)
v₁ = zero(v₀)
f₁ = zero(f₀)

t₂ = one(Δt)
q₂ = zero(q₀)
p₂ = zero(p₀)
v₂ = zero(v₀)
f₂ = zero(f₀)

pode.v(t₀, q₀, p₀, v₀, pode.parameters)
pode.f(t₀, q₀, p₀, f₀, pode.parameters)

initialize!(igpode, t₀, q₀, p₀, v₀, f₀, t₁, q₁, p₁, v₁, f₁)
evaluate!(igpode, q₁, p₁, v₁, f₁, q₀, p₀, v₀, f₀, q₂, v₂, t₂)

# println("IG-PODE")
# println(q₁ .- qₚ)
# println(v₁ .- vₚ)
# println(q₂ .- qₙ)
# println(v₂ .- vₙ)
# println()

@test q₁ ≈ qₚ atol=1E-14
@test v₁ ≈ vₚ atol=1E-14

@test q₂ ≈ qₙ atol=1E-8
@test v₂ ≈ vₙ atol=1E-5