using Compat
import Base: Forward
import SortingAlgorithms: RADIX_MASK, RADIX_SIZE, uint_mapping
using FastGroupBy
using BenchmarkTools

using SortingAlgorithms, Base.Order, Compat, IndexedTables, DataFrames
import Base: ht_keyindex, rehash!, _setindex!
import SortingAlgorithms: uint_mapping, RADIX_SIZE, RADIX_MASK
import DataFrames: DataFrame, AbstractDataFrame
import IndexedTables: IndexedTable, column
import PooledArrays.PooledArray
import CategoricalArrays.CategoricalArray
import SplitApplyCombine.groupreduce

using StatsBase

T = Int
S = Int

N  = 1_000_000;
K = 100;
by = nothing;
val = nothing;
gc();
srand(1);
by = rand(Int32(1):Int32(round(N/K)), N);
val = rand(Int32(1):Int32(5), N)

#
# @time sumby_lessmem_chain(by, val);
#
# by = nothing;
# val = nothing;
# gc();
# srand(1);
# @time val = rand(Int32(1):Int32(5), N);
# @time by = rand(Int32(1):Int32(round(N/K)), N);
# @time sumby(by, val);
#
# function abc()
#   by1, val1 = copy(by), copy(val)
#   @elapsed sumby_lessmem_chain(by1, val1)
# end
#
# function def()
#   by1, val1 = copy(by), copy(val)
#   @elapsed sumby(by1, val1)
# end

# a1 = mean([abc() for i = 1:5][2:end])
# b1 = mean([def() for i = 1:5][2:end])
# 1 - a1/b1
2+2


const AXIOM_NODES = ["001", "000", "01"]
const DELETED = "del"

const NOISE_LEVEL_NODES = 0.6
const NOISE_LEVEL_MISSING_OUTPUT = 0.5
const EXCITATORY_MIN = 0
const EXCITATORY_MAX = 2
const INHIBITORY_THRESHOLD = 0.5

const MAX_TURN = 60

const FIGHT_THRESHOLD = 0.5
const CROSS_OVER_THRESHOLD = 0.5
const INFLUENCE_RADIUS = 2

const DISTANCE_SCALE_FACTOR = 4

const MUTATION_PROBABILITY = 0

const MAX_NEURONS = 20
const MIN_ORGANISMS = 0 # 20
const MAX_DEVELOPMENT_STEPS = 99


import AmplNLWriter

@testset "[nl_convert] check special conversion cases" begin
    special_cases = [:cbrt, :abs2, :inv, :log2, :log1p, :exp2, :expm1, :sec,
                     :csc, :cot, :sind, :cosd, :tand, :asind, :acosd, :atand,
                     :secd, :cscd, :cotd, :sech, :csch, :coth, :asech, :acsch]
    for func in special_cases
        x = rand()
        expr = Expr(:call, func, x)
        @test isapprox(eval(AmplNLWriter.convert_formula(expr)), eval(expr), atol=1e-6)
    end
    # These functions need input >1
    for func in [:acoth, :asec, :acsc, :acot, :asecd, :acscd, :acotd]
        x = rand() + 1
        expr = Expr(:call, func, x)
        @test isapprox(eval(AmplNLWriter.convert_formula(expr)),
               eval(expr), atol=1e-6)
    end
end

@testset "[nl_convert] check numeric values" begin
    x = rand()
    @test AmplNLWriter.convert_formula(:($x)) == :($x)
    x = -rand()
    @test AmplNLWriter.convert_formula(:($x)) == :($x)
end

@testset "[nl_convert] check unary, binary and n-ary plus" begin
    expr = :(+(1))
    @test AmplNLWriter.convert_formula(expr) == :(1)
    expr = :(1 + 2)
    @test AmplNLWriter.convert_formula(expr) == :(1 + 2)
    expr = :(1 + 2 + 3)
    @test AmplNLWriter.convert_formula(expr) == :(sum(1, 2, 3))
end

@testset "[nl_convert] check unary, binary and n-ary minus" begin
    expr = :(- x)
    @test AmplNLWriter.convert_formula(expr) == :(neg(x))
    expr = :(x - y)
    @test AmplNLWriter.convert_formula(expr) == :(x - y)
    expr = :(x - y - z)
    @test AmplNLWriter.convert_formula(expr) == :((x - y) - z)
end

@testset "[nl_convert] check n-ary multiplication" begin
    expr = :(x * y * z)
    @test AmplNLWriter.convert_formula(expr) == :(x * (y * z))
end

@testset "[nl_convert] check comparison expansion" begin
    expr = :(1 < 2 < 3)
    @test AmplNLWriter.convert_formula(expr) == :(1 < 2 && 2 < 3)
    expr = :(1 < 2 < 3 < 4)
    @test AmplNLWriter.convert_formula(expr) == :((1 < 2 && 2 < 3) && 3 < 4)
end

@testset "[nl_convert] check user defined functions error" begin
    m = Model(solver=AmplNLSolver("any_solver"))

    myf(x,y) = (x-1)^2+(y-2)^2
    JuMP.register(m, :myf, 2, myf, autodiff=true)

    @variable(m, x[1:2] >= 0.5)
    @NLobjective(m, Min, myf(x[1], x[2]))

    @test_throws Exception solve(m)

    # clean up temp file
    for file in readdir(AmplNLWriter.solverdata_dir)
        if file != ".gitignore"
            rm(joinpath(AmplNLWriter.solverdata_dir, file))
        end
    end
end


@with_kw mutable struct PedestrianAvoidancePOMDPFrenet{B <: Updater} <: DriverModel{LatLonAccel}

    timestep::Float64 = 0.2
    t_current::Float64 = 0
    tick::Int64 = 0

    ego_vehicle::Vehicle = Vehicle(VehicleState(VecSE2(0., 0., 0.), 0.), VehicleDef(), 1)
    a::LatLonAccel = LatLonAccel(0.0, 0)
    desired_velocity::Float64 = 50.0 / 3.6

    obstacles::Vector{ConvexPolygon} = []
    sensor_observations::Vector{Vehicle} = []
    objects_tracked = Dict{Int64, Float64}()

    update_tick_high_level_planner::Int64 = 1

    env::CrosswalkEnv = CrosswalkEnv(CrosswalkParams())
    sensor::AutomotiveSensors.GaussianSensor = AutomotiveSensors.GaussianSensor(AutomotiveSensors.LinearNoise(10, 0., 0.),
                                                               AutomotiveSensors.LinearNoise(10, 0., 0.), 0, 0, MersenneTwister(1))

    pomdp::SingleOCFPOMDP = SingleOCFPOMDP()
    policy::AlphaVectorPolicy{SingleOCFPOMDP,SingleOCFAction} = AlphaVectorPolicy(pomdp, Vector{Vector{Float64}}())
    updater::B = SingleOCFUpdater(pomdp)

    policy_dec::DecPolicy = DecPolicy(policy, pomdp, policy.action_map, (x,y) -> min.(x,y))

    b::SingleOCFBelief = SingleOCFBelief([],[])
    b_dict::Dict{Int64, SingleOCFBelief} = Dict{Int64, SingleOCFBelief}()

    prediction_obstacle::Array{Float64} = []

    brake_request::Bool = false
    collision_rate::Float64 = 0.0
    risk::Float64 = 0.0
    ttc::Float64 = 0.0

end

function AutomotiveDrivingModels.observe!(model::PedestrianAvoidancePOMDPFrenet, scene::Scene, roadway::Roadway, egoid::Int)

    ego = scene[findfirst(egoid, scene)]
    model.pomdp.ego_vehicle = ego
    model.ego_vehicle = ego
    #model.sensor_observations = measure(PerfectSensor(), ego, scene, roadway, model.obstacles)

    # get observations from sensor, only visible objects with sensor noise
    model.sensor_observations = measure(model.sensor, ego, scene, roadway, model.obstacles)

    # add sensor delay 0.2s
    model.sensor_observations = objects_time_delay(model)

    # initialization of the belief for the absent state
    if (model.t_current == 0 )
        model.b_dict[PEDESTRIAN_OFF_KEY] = initBeliefAbsentPedestrian(model.pomdp, ego.state.posF.t, ego.state.v)
    end

    ################ High Level Planner ###################################################
    if (  true ) #model.tick % model.update_tick_high_level_planner == 0 )

        #println("--------------------------POMDP high level planner----------------------- t: ", model.t_current)
        #println("EGO: x/y:", ego.state.posG.x, " / ",  ego.state.posG.y, " v: ", ego.state.v)

        # prepare observation dictionary / remove observations outside state space and add absent state
        observations = get_observations_state_space(model, ego, model.sensor_observations)

        # update belief dictionary
        b_new = update(model.pomdp, model.updater, model.b_dict, SingleOCFAction(model.a.a_lon, model.a.a_lat), observations)
        model.b_dict = deepcopy(b_new)
        #println("dict: ", model.b_dict)

#=
b_dict_tmp = Dict{Int64, SingleOCFBelief}()
for oid in keys(model.b_dict)
    if ( oid == PEDESTRIAN_OFF_KEY )
        b_states = []
        b_prob = []
        for (s, prob) in weighted_iterator(model.b_dict[PEDESTRIAN_OFF_KEY] )
            if ( is_state_absent(model.pomdp, s) )
                (ego_y_state_space, ego_v_state_space) = getEgoDataInStateSpace(model.pomdp, ego.state.posF.t, ego.state.v)
                s = SingleOCFState(ego_y_state_space, ego_v_state_space, -10.0, -10.0, model.pomdp.PED_THETA_RANGE[1] ,0.0)
            end
            push!(b_states, s)
            push!(b_prob, prob)
        end
        normalize!(b_prob, 1)
        b_dict_tmp[PEDESTRIAN_OFF_KEY] =  SingleOCFBelief(b_states, b_prob)
    else
        b_dict_tmp[oid] = model.b_dict[oid]
    end
end
=#
        # use policy and belief dictionary to calculate next action
        act = action(model.policy_dec, model.b_dict)
        model.a = LatLonAccel(act.lateral_movement, act.acc)
        #println("action combined: ", model.a)

        # dummy functionality to test transition function / belief update
        #=
        if (model.tick > 2 )
           model.a = LatLonAccel(1.0, -2.0)
            println("manual intervention")
        end
        =#

    end
    model.tick += 1
    model.t_current = model.t_current + model.timestep
end

function AutomotiveDrivingModels.propagate(veh::Vehicle, action::LatLonAccel, roadway::Roadway, Δt::Float64)

    # new velocity
    v_ = veh.state.v + action.a_lon*Δt
    v_ = clamp(v_, 0., v_)

    # lateral offset
    delta_y = action.a_lat * Δt   # a_lat corresponds to lateral velocity --> a_lat == v_lat
    if v_ <= 0.
        delta_y = 0.
    end
    y_ = veh.state.posG.y + delta_y
    y_ = clamp(y_, -1.0, 1.0)

    s_new = v_ * Δt

    # longitudional distance based on required velocity and lateral offset
    if (abs(delta_y) < abs(s_new) )
        delta_x = sqrt(s_new^2 - delta_y^2 )
    else
        delta_x = 0
    end

    delta_x = clamp(delta_x, 0., delta_x)
    x_ = veh.state.posG.x + delta_x

    return VehicleState(VecSE2(x_, y_, veh.state.posG.θ), roadway, v_)
end

function AutomotiveDrivingModels.get_name(model::PedestrianAvoidancePOMDPFrenet)
    return "Pedestrian Avoidance System (POMDP)"
end

AutomotiveDrivingModels.rand(rng::AbstractRNG, model::PedestrianAvoidancePOMDPFrenet) = model.a

@with_kw mutable struct ObservationCallback

    sensor_observations::Vector{Vector{Vehicle}}
    ego_vehicle::Vector{Vehicle}
    ego_a::Vector{Float64}
    collision::Vector{Bool}

    b_dict::Vector{Dict{Int64, SingleOCFBelief}}
    action::Vector{SingleOCFAction}

    collision_rate::Vector{Float64}
    ttc::Vector{Float64}
    risk::Vector{Float64}
    brake_request::Vector{Bool}
    prediction_obstacle::Vector{Array{Float64}}

end

function AutomotiveDrivingModels.run_callback(
        callback::ObservationCallback,
        scenes::Vector{Scene},
        actions::Union{Nothing, Vector{Frame{A}}},
        roadway::R,
        models::Dict{I,M},
        tick::Int) where {A<:EntityAction, S,D,I,R,M<:DriverModel}

    push!(callback.sensor_observations, models[1].sensor_observations)
    push!(callback.ego_vehicle, models[1].ego_vehicle)
    push!(callback.ego_a, models[1].a.a_lon)

    collision = is_crash(scenes[tick])
    push!(callback.collision, collision)

    push!(callback.b_dict, deepcopy(models[1].b_dict))
    push!(callback.action, SingleOCFAction(models[1].a.a_lon, models[1].a.a_lat))

    push!(callback.collision_rate, models[1].collision_rate)
    push!(callback.ttc, models[1].ttc)
    push!(callback.risk, models[1].risk)
    push!(callback.brake_request, models[1].brake_request)
    push!(callback.prediction_obstacle, models[1].prediction_obstacle)

    return collision
end


module DICE

using JuMP;
#Ipopt is distributed under the EPL.
#We don't package it here though, and assume you have this set up on your system already.
using Ipopt;

include("Abstractions.jl")

#DICE Versions
abstract type Version end

#Configuration options & parameters, model settings and output
include("BaseTypes.jl")

#Scenarios
abstract type Scenario end
include("Scenarios.jl")

struct DICENarrative
    constants::Options
    parameters::Parameters
    model::JuMP.Model
    scenario::Scenario
    version::Version
    variables::Variables
    equations::Equations
    results::Results
end

function Base.show(io::IO, ::MIME"text/plain", dice::DICENarrative)
    println(io, "$(dice.scenario) scenario using $(dice.version).")
    show(io, dice.model)
end

function solve end
function options end

export solve, options

# A list of all linear solvers that Ipopt will accept.
# Note that DICE.jl is only tested against ma97 and mumps.
# We use this externally so users cannot accidentally send in nonsense strings,
# although the result is string based and used throughout as such.
@exported_enum ipoptLinearSolver ma27 ma57 ma77 ma86 ma97 pardiso wsmp mumps custom

# Test if system has HSL MA97 installed, fall back to MUMPS if not.
function selectLinearSolver(solver_name::ipoptLinearSolver = ma97)
    prob = Ipopt.CreateIpoptProblem(1,[1.],[1.],1,[1.],[1.],1,1,sum,sum,sum,sum,sum);
    Ipopt.AddIpoptStrOption(prob, "sb", "yes");
    Ipopt.AddIpoptIntOption(prob, "print_level", 0);
    solver_name_string = string(solver_name);
    # Initially, we must check that coinhsl is installed at all if we want a HSL solver.
    # If not, Ipopt will default to try and find any dynamically linked libhsl. If it
    # cannot find one it will hard panic and we can't capture that failure.
    if occursin("ma", solver_name_string)
        Ipopt.AddIpoptStrOption(prob, "linear_solver", "ma27");
        try
            # No HSL was found on the system, we return with fallback.
            return runLinearSolverCheck(prob, solver_name_string)
        catch
            #Continue, we have access to at least some version of HSL
        end
    end
    try
        # Outer try will fail if solver string is not in the list of
        # possible Ipopt solvers.
        # For now that's ma27, ma57, ma77, ma86, ma97, pardiso, wsmp, mumps, custom
        Ipopt.AddIpoptStrOption(prob, "linear_solver", solver_name_string);
        try
            # Inner try attempts to run the dummy program and will crash because
            # the dummy is malformed.
            # No error code will be returned if the solver is extant, and an Invalid_Option
            # if the library is not found.
            runLinearSolverCheck(prob, solver_name_string)
        catch
            # We can use the requested solver.
            solver_name_string
        end
    catch
        "mumps"
    end
end

function runLinearSolverCheck(prob::IpoptProblem, solver_name_string::String)
    result_code = Ipopt.IpoptSolve(prob);
    if Ipopt._STATUS_CODES[result_code] == :Invalid_Option
        @info "Unable to set linear_solver = $(solver_name_string), defaulting to MUMPS."
        return "mumps"
    else
        error("Attempts to identify linear solvers on system returned unexpected results.");
    end
end

# Include all version implementations
include("2013R.jl")
include("2016.jl")

end


"""
    get_buff_modifier(buff)

Compute the mulitplier associated with stat buffs (multiplied by 12 to return
an integer). As a result, the multiplier for no buff effect is 0. Inputs should
be between -4 and 4.

# Examples
```jldoctest
julia> get_buff_modifier(Int8(0))
12
"""
function get_buff_modifier(i::UInt32, agent::UInt8)
    a, d = agent == 0x01 ?
        (UInt8((i ÷ UInt32(13230))   % 9),
        UInt8((i ÷ UInt32(119070))   % 9)) :
        (UInt8((i ÷ UInt32(1071630)) % 9),
        UInt8((i ÷ UInt32(9644670))  % 9))
    a1, a2 = a > UInt8(3) ? (a, UInt8(4)) :
        (UInt8(4), UInt8(8) - a)
    d1, d2 = d > UInt8(3) ? (d, UInt8(4)) :
        (UInt8(4), UInt8(8) - d)
    return a1*d2, a2*d1
end

evaluate_fast_moves(state::DynamicIndividualState,
    static_state::StaticIndividualState, using_fm::Tuple{Bool, Bool}) =
    DynamicIndividualState(
        add_energy(damage(state[0x01],
            using_fm[0x02] ? calculate_damage(
                static_state[0x02].stats.attack,
                state.data,
                static_state[0x01],
                static_state[0x02].fast_move,
            ) : 0x0000), (using_fm[0x01] ?
            get_energy(static_state[0x01].fast_move) : Int8(0))),
        add_energy(damage(state[0x02],
            using_fm[0x01] ? calculate_damage(
                static_state[0x01].stats.attack,
                state.data,
                static_state[0x02],
                static_state[0x01].fast_move,
            ) : 0x0000), (using_fm[0x02] ?
            get_energy(static_state[0x02].fast_move) : Int8(0))),
        state.data)

"""
    evaluate_charged_move(state, static_state, cmp, move_id, charge,
        shielding, buffs_applied)

Takes in the dynamic state, the static state, the attacking agent, the move,
the charge, whether or not the opponent shields, and whether or not buffs are
applied (say in the case of a random buff move) and returns
the dynamic state after the charged move has occurred, with precisely one copy
"""
function evaluate_charged_move(state::DynamicIndividualState,
    static_state::StaticIndividualState, cmp::UInt32, move_id::UInt8,
    charge::UInt8, shielding::Bool)
    next_state = state
    agent = isodd(cmp) ? 0x01 : 0x02
    d_agent = get_other_agent(agent)
    data = next_state.data
    move = move_id == 0x01 ? static_state[agent].charged_move_1 :
        static_state[agent].charged_move_2

    buff_chance = move.buffChance
    if buff_chance == Int8(100)
        data = apply_buff(data, move, agent)
    elseif buff_chance != Int8(0)
        data += agent == 0x01 ? move_id == 0x01 ? UInt32(2205) : UInt32(4410) :
                             move_id == 0x01 ? UInt32(6615) : UInt32(8820)
    end

    attacking_team = subtract_energy(next_state[agent], get_energy(move))

    if shielding
        defending_team = damage(next_state[d_agent], 0x0001)
        data -= d_agent == 0x01 ? 1 : 3
    else
        defending_team = damage(next_state[d_agent], calculate_damage(
                static_state[agent].stats.attack,
                state.data,
                static_state[d_agent],
                move,
                charge = charge
            ))
    end
    return DynamicIndividualState(
        agent == 1 ? attacking_team : defending_team,
        agent == 2 ? attacking_team : defending_team,
        # go from cmp 4 (2 then 1) to cmp 1
        # or go from cmp 3 (1 then 2) to cmp 2
        # or go from cmp 2 to 0
        # or go from cmp 1 to 0
        data  - (cmp == 4 ? (get_hp(defending_team) != 0x0000 ?
                    1323 : 1764) :
                (cmp == 3 ? (get_hp(defending_team) != 0x0000 ?
                    441 : 1323) :
                (cmp == 2 ? 882 : 441)))
    )
end

function apply_buff(data::UInt32, move::ChargedMove, agent::UInt8)
    buffs1 = get_buffs(data, 1)
    buffs2 = get_buffs(data, 2)
     return agent == 1 ? data + Int32(13230) * clamp(
        get_atk(move.self_buffs), -Int8(buffs1[1]), 9 - Int8(buffs1[1])) +
        Int32(119070) * clamp(get_def(move.opp_buffs),
        -Int8(buffs1[2]), 9 - Int8(buffs1[2]))  +
        Int32(1071630) * clamp(get_atk(move.opp_buffs),
        -Int8(buffs2[1]), 9 - Int8(buffs2[1]))  +
        Int32(9644670) * clamp(get_def(move.self_buffs),
        -Int8(buffs2[2]), 9 - Int8(buffs2[2])) :
        data + Int32(13230) * clamp(get_atk(move.opp_buffs),
        -Int8(buffs1[1]), 9 - Int8(buffs1[1])) +
        Int32(119070) * clamp(get_def(move.self_buffs),
        -Int8(buffs1[2]), 9 - Int8(buffs1[2]))  +
        Int32(1071630) * clamp(get_atk(move.self_buffs),
        -Int8(buffs2[1]), 9 - Int8(buffs2[1]))  +
        Int32(9644670) * clamp(get_def(move.opp_buffs),
        -Int8(buffs2[2]), 9 - Int8(buffs2[2]))
end

function step_timers(state::DynamicIndividualState, fmCooldown1::Int8,
    fmCooldown2::Int8)
    fmPending = get_fast_moves_pending(state)
    data = state.data
    if fmCooldown1 != Int8(0)
        data += (UInt32(fmCooldown1) - fmPending[1]) * 9
    elseif fmPending[1] != UInt32(0)
        data -= 9
    end
    if fmCooldown2 != Int8(0)
        data += (UInt32(fmCooldown2) - fmPending[2]) * 63
    elseif fmPending[2] != UInt32(0)
        data -= 63
    end

    return DynamicIndividualState(state[0x01], state[0x02], data)
end

"""
    min_score(state, static_state)

Given the state and the static state (here just for starting hp values), compute
the PvPoke-like score that would occur if the first agent stopped attacking
altogether. This is currently only used in computing the final score, but it
could be used as strict bounds for α/β pruning, for example.
"""
function min_score(state::DynamicIndividualState,
    static_state::StaticIndividualState)
     return 0.5 * (static_state[0x02].stats.hitpoints - get_hp(state[0x02])) /
        static_state[0x02].stats.hitpoints
end

"""
    max_score(state, static_state)

Given the state and the static state (here just for starting hp values), compute
the PvPoke-like score that would occur if the second agent stopped attacking
altogether. This is currently only used in computing the final score, but it
could be used as strict bounds for α/β pruning, for example.
"""
function max_score(state::DynamicIndividualState,
    static_state::StaticIndividualState)
     return 0.5 + 0.5 * get_hp(state[0x01]) / static_state[0x01].stats.hitpoints
end

"""
    battle_score(state, static_state)

Given the state and the static state (here just for starting hp values), compute
the PvPoke-like score for the battle. Note that this can also be computed for
battles in progress, and thus differs from PvPoke's use cases
"""
function battle_score(state::DynamicIndividualState,
    static_state::StaticIndividualState)
    return min_score(state, static_state) + max_score(state, static_state) - 0.5
end


# Julia wrapper for the sgp4 Python library:
# https://pypi.python.org/pypi/sgp4/

module SGP4

using PyCall, Dates

import Base.getindex

const sgp4io = PyNULL()
const earth_gravity = PyNULL()
const sgp4_propagation = PyNULL()

function __init__()
    copy!(sgp4io, pyimport_conda("sgp4.io", "sgp4", "conda-forge"))
    copy!(earth_gravity, pyimport("sgp4.earth_gravity"))
    copy!(sgp4_propagation, pyimport("sgp4.propagation"))
end

export GravityModel,
       twoline2rv,
       propagate

struct GravityModel
    model::PyObject # can be any of {wgs72old, wgs72, wgs84}
end

mutable struct SGP4Sat
    s::PyObject
end
getindex(sat::SGP4Sat, sym::Symbol) = sat.s[sym]

GravityModel(ref::AbstractString) = GravityModel(earth_gravity[ref])

# sgp4.io convenience functions
function twoline2rv(line1::String, line2::String, grav::GravityModel)
    return SGP4Sat(sgp4io["twoline2rv"](line1,line2,grav.model))
end

function sgp4( sat::SGP4Sat,
               dtmin::Real )
    r,v = sgp4_propagation["sgp4"](sat.s, dtmin)
    return vcat(r...), vcat(v...)
end

"""
Propagate the satellite by `dtmin` minutes.

Returns (position, velocity)
"""
function propagate( sat::SGP4Sat,
                    dtmin::Real )
    sgp4(sat,dtmin)
end

function propagate(sats::Vector{SGP4Sat},
                   dtmin::Real)
    f = x->propagate(x, dtmin)
    f.(sats)
end

"""
Propagate the satellite from its epoch to the date/time specified

Returns (position, velocity) at the specified time
"""
function propagate( sat::SGP4Sat,
                    year::Real,
                    month::Real,
                    day::Real,
                    hour::Real,
                    min::Real,
                    sec::Real )
    (pos, vel) = sat.s[:propagate](year,month,day,hour,min,sec)

    # check for errors
    if sat.s[:error] != 0
        println(sat.s[:error_message])
    end

    return ([pos...],[vel...])
end

function propagate( sat::SGP4Sat,
                    t::DateTime )
    propagate(sat,
              Dates.year(t),
              Dates.month(t),
              Dates.day(t),
              Dates.hour(t),
              Dates.minute(t),
              Dates.second(t) + Dates.millisecond(t)/1000)
end

function propagate(sat::SGP4Sat,
                   t::AbstractVector{DateTime})
    pos = zeros(3, length(t))
    vel = zeros(3, length(t))

    for (idx, ti) in enumerate(t)
        pos[:,idx],vel[:,idx] = propagate(sat, ti)
    end
    return (pos,vel)
end

"Generate a satellite ephemeris"
function propagate( sat::SGP4Sat,
                    tstart::DateTime,
                    tstop::DateTime,
                    tstep::Dates.TimePeriod )
    propagate(sat, tstart:tstep:tstop)
end

"tstep specified in seconds"
function propagate( sat::SGP4Sat,
                    tstart::DateTime,
                    tstop::DateTime,
                    tstep::Real )
    propagate(sat,tstart,tstop,Dates.Second(tstep))
end

"Propagate many satellites to a common time"
function propagate( sats::Vector{SGP4Sat},
                    year::Real,
                    month::Real,
                    day::Real,
                    hour::Real,
                    min::Real,
                    sec::Real )
    f = x->propagate(x, year, month, day, hour, min, sec)
    f.(sats)
end

function propagate( sats::Vector{SGP4Sat},
                    t::DateTime )
    f = x->propagate(x, Dates.year(t), Dates.month(t), Dates.day(t), Dates.hour(t), Dates.minute(t), Dates.second(t) + Dates.millisecond(t)/1000)
    f.(sats)
end

end #module


using MeasureBase: logdensityof, Likelihood
export RadioLikelihood, logdensityof, MultiRadioLikelihood
using LinearAlgebra

struct RadioLikelihood{T,A} <: MB.AbstractMeasure
    lklhds::T
    ac::A
end

struct MultiRadioLikelihood{L} <: MB.AbstractMeasure
    lklhds::L
end

"""
    `MultiRadioLikelihood(lklhd1, lklhd2, ...)`
Combines multiple likelihoods into one object that is useful for fitting multiple days/frequencies.

```julia
lklhd1 = RadioLikelihood(dcphase1, dlcamp1)
lklhd2 = RadioLikelihood(dcphase2, dlcamp2)

lklhd = MultiRadioLikelihood(lklhd1, lklhd2)
```

"""
MultiRadioLikelihood(lklhds::RadioLikelihood...) = MultiRadioLikelihood(lklhds)

function MB.logdensityof(lklhds::MultiRadioLikelihood, m)
    sum(Base.Fix2(logdensityof, m), lklhds.lklhds)
end

"""
    `RadioLikelihood(data1, data2, ...)`
Forms a radio likelihood from a set of data products. These data products must share
the same array data/configuration. If you want to form a likelihood from multiple arrays
such as when fitting different wavelengths or days, you can combine them using
`MultiRadioLikelihood`

```julia
lklhd1 = RadioLikelihood(dcphase1, dlcamp1)
lklhd2 = RadioLikelihood(dcphase2, dlcamp2)

lklhd = MultiRadioLikelihood(lklhd1, lklhd2)
```
"""
function RadioLikelihood(data::EHTObservation...)
    ls = Tuple(map(makelikelihood, data))
    acs = arrayconfig.(data)
    #@argcheck acs[1] == acs[2]
    RadioLikelihood{typeof(ls), typeof(acs[1])}(ls, acs[1])
end

function RadioLikelihood(data::MT.Likelihood...)
    return RadioLikelihood{typeof(data), Nothing}(data, nothing)
end

function Base.show(io::IO, d::RadioLikelihood{T}) where {T}
    println(io, "RadioLikelihood{$T}")
    println(io, "\tNumber of data products: ", length(d.lklhds))
end

"""
    `logclosure_amplitudes(vis, ac::ArrayConfiguration)`

Compute the log-closure amplitudes for a set of visibilities and an array configuration

# Notes
This uses a closure design matrix for the computation.
"""
function logclosure_amplitudes(vis::AbstractArray{<:Complex}, ac::ArrayConfiguration)
    lva = log.(abs.(vis))
    return ac.designmat*lva
end

"""
    `closure_phases(vis, ac::ArrayConfiguration)`

Compute the closure phases for a set of visibilities and an array configuration

# Notes
This uses a closure design matrix for the computation.
"""
function closure_phases(vis::AbstractArray{<:Complex}, ac::ArrayConfiguration)
    ph = angle.(vis)
    return ac.designmat*ph
end

amplitudes(vis::AbstractArray{<:Complex}, ac::ArrayConfiguration) = abs.(vis)
phase(vis::AbstractArray{<:Complex}) = angle.(vis)

function MB.logdensityof(d::RadioLikelihood, m::ComradeBase.AbstractModel)
    ac = d.ac
    vis = visibilities(m, ac)
    return logdensityof(d, vis)
end

function MB.logdensityof(d::RadioLikelihood, vis::AbstractArray)
    # We use a for loop here since Zygote plays nice with this
    acc = logdensityof(first(d.lklhds), vis)
    @inbounds for l in d.lklhds[begin+1:end]
        acc += logdensityof(l, vis)
    end
    return acc
end

#function MB.logdensityof(d::RadioLikelihood, m::ComradeBase.AbstractModel)
#    acc = logdensityof(first(d.lklhds), m)
#    @inbounds for i in 2:length(d.lklhds)
#        acc += logdensityof(d.lklhds[i], m)
#    end
#    return acc
#end

function makelikelihood(data::Comrade.EHTObservation{<:Real, <:Comrade.EHTVisibilityDatum})
    errinv = inv.(data[:error])
    vis = StructArray{Complex{eltype(data[:visr])}}((data[:visr],data[:visi]))
    ℓ = Likelihood(vis) do (μ, )
        ComplexNormal{(:μ, :τ)}(μ, errinv)
    end
    return ℓ
end

function makelikelihood(data::Comrade.EHTObservation{<:Real, <:Comrade.EHTVisibilityAmplitudeDatum})
    τ = inv.(data[:error])
    amp = getdata(data, :amp)
    ℓ = Likelihood(amp) do (μ, )
        AmpNormal{(:μ, :τ)}(abs.(μ), τ)
    end
    return ℓ
end

function makelikelihood(data::Comrade.EHTObservation{<:Real, <:Comrade.EHTLogClosureAmplitudeDatum})
    dmat = data.config.designmat
    τ  = inv.(data[:error])
    #covm = σvis*dmat*transpose(σvis)

    #C = Cholesky(Symmetric(covm))
    #σ = C.L

    f = Base.Fix2(logclosure_amplitudes, data.config)
    amp = data[:amp]
    ℓ = Likelihood(amp) do (μ,)
        AmpNormal{(:μ, :τ)}(f(μ), τ)
    end
    #return AmpNormal{(:μ, :τ)}(amp, τ)
    return ℓ
end

function makelikelihood(data::Comrade.EHTObservation{<:Real, <:Comrade.EHTClosurePhaseDatum})
    τ = inv.(getdata(data, :error)).^2
    f = Base.Fix2(closure_phases, data.config)
    phase = data[:phase]
    ℓ = Likelihood(phase) do (μ,)
        CPVonMises{(:μ, :κ)}(f(μ), τ)
    end

    return ℓ
end


#######################################################################################################################################################################################################
# Changes to these constants
# General
#     2020-May-30: put the 2017 mat file in an artifact
#     2020-Aug-30: add the updated 2021 mat file along with that of 2017 into a new artifact
#
#######################################################################################################################################################################################################const FILE_SUN  = artifact"land_model_spectrum_V1" * "/sun.mat";
const OPTI_2017 = artifact"land_model_spectrum_V1" * "/Optipar2017_ProspectD.mat";
const OPTI_2021 = artifact"land_model_spectrum_V1" * "/Optipar2021_ProspectPRO_CX.mat";
WAVELENGTHS = [collect(400:10:650.1); collect(655:5:770.1); collect(780:25:2400.1)];

#######################################################################################################################################################################################################
# Changes to this structure
# General
#     2021-Aug-10: refactor the structure with renamed fields
#     2021-Aug-10: add a constructor within the structure to avoid external initialization
#     2021-Oct-19: sort variable to prognostic and dignostic catergories
#
#######################################################################################################################################################################################################
"""
$(TYPEDEF)

Structure that stores wave length information.

$(TYPEDFIELDS)

---
# Examples
```julia
wls = WaveLengthSet{FT}();
wls = WaveLengthSet{FT}(collect(FT,400:5:2500));
wls = WaveLengthSet{FT}(collect(FT,400:5:2500); opti=Emerald.OPTI_2017);
```
"""
mutable struct WaveLengthSet{FT<:AbstractFloat}
    # parameters that do not change with time
    "Indicies of Λ_NIR in Λ"
    IΛ_NIR::Vector{Int}
    "Indicies of Λ_PAR in Λ"
    IΛ_PAR::Vector{Int}
    "Indicies of Λ_SIF in Λ"
    IΛ_SIF::Vector{Int}
    "Indicies of Λ_SIFE in Λ"
    IΛ_SIFE::Vector{Int}
    "Number of wavelength bins"
    NΛ::Int
    "Number of wavelength bins for PAR"
    NΛ_PAR::Int
    "Number of wavelength bins for SIF"
    NΛ_SIF::Int
    "Number of wavelength bins for SIF excitation"
    NΛ_SIFE::Int
    "Wavelength limits for NIR `[nm]`"
    WL_NIR::Vector{FT}
    "Wavelength limits for PAR `[nm]`"
    WL_PAR::Vector{FT}
    "Wavelength limits for SIF emission `[nm]`"
    WL_SIF::Vector{FT}
    "Wavelength limits for SIF excitation `[nm]`"
    WL_SIFE::Vector{FT}
    "Standard wavelength (boundaries) `[nm]`"
    SΛ::Vector{FT}
    "Differential wavelength `[nm]`"
    ΔΛ::Vector{FT}
    "Differential wavelength for PAR `[nm]`"
    ΔΛ_PAR::Vector{FT}
    "Differential wavelength for SIF excitation `[nm]`"
    ΔΛ_SIFE::Vector{FT}
    "Wavelength (bins) `[nm]`"
    Λ::Vector{FT}
    "Wavelength bins for PAR `[nm]`"
    Λ_PAR::Vector{FT}
    "Wavelength bins for SIF `[nm]`"
    Λ_SIF::Vector{FT}
    "Wavelength bins for SIF excitation `[nm]`"
    Λ_SIFE::Vector{FT}

    # prognostic variables that change with time

    # dignostic variables that change with time

    # caches to speed up calculations

    # constructors
    WaveLengthSet{FT}(swl::Vector{FT}=FT.(WAVELENGTHS); opti::String=OPTI_2021) where {FT<:AbstractFloat} = (
        _dwl    = diff(swl);
        _λ      = zeros(FT, length(swl)-1);
        _opti   = matread(opti)["optipar"];
        _λ_opti = _opti["wl"];
        @inbounds for _i in 1:length(swl)-1
            _wo    = findall( swl[_i] .<= _λ_opti .< swl[_i+1] );
            _λ[_i] = nanmean(_λ_opti[_wo]);
        end;
        _iλ_nir  = findall( 700 .<= _λ .<= 2500 );
        _iλ_par  = findall( 400 .<= _λ .<= 700  );
        _iλ_sif  = findall( 640 .<= _λ .<= 850  );
        _iλ_sife = findall( 400 .<= _λ .<= 750  );

        return new(_iλ_nir, _iλ_par, _iλ_sif, _iλ_sife, length(_λ), length(_iλ_par), length(_iλ_sif), length(_iλ_sife), FT[700,2500], FT[400,700], FT[640,850], FT[400,750], swl, _dwl, _dwl[_iλ_par],
                   _dwl[_iλ_sife], _λ, _λ[_iλ_par], _λ[_iλ_sif], _λ[_iλ_sife])
    )
end


export TabularLearner

"""
    TabularLearner{S, T}

Use a `Dict{S,Vector{T}}` to store action probabilities.
"""
struct TabularLearner{S,T} <: AbstractLearner
    table::Dict{S,Vector{T}}
end

TabularLearner() = TabularLearner{Int,Float32}()
TabularLearner{S}() where {S} = TabularLearner{S,Float32}()
TabularLearner{S,T}() where {S,T} = TabularLearner(Dict{S,Vector{T}}())

(p::TabularLearner)(env::AbstractEnv) = p(ChanceStyle(env), env)

function (p::TabularLearner)(::ExplicitStochastic, env::AbstractEnv)
    if get_current_player(env) == get_chance_player(env)
        [a.prob for a::ActionProbPair in get_actions(env)]
    else
        p(DETERMINISTIC, env)  # treat it just like a normal one
    end
end

function (p::TabularLearner)(::RLBase.AbstractChanceStyle, env::AbstractEnv)
    get!(p.table, get_state(env)) do
        n = length(get_legal_actions(env))
        fill(1 / n, n)
    end
end

RLBase.update!(p::TabularLearner, experience::Pair) =
    p.table[first(experience)] = last(experience)


# This file is a part of MGVI.jl, licensed under the MIT License (MIT).

using MGVI

using Random

if :ModelPolyfit ∉ names(Main)
    include("test_models/model_polyfit.jl")
    import .ModelPolyfit
end

Test.@testset "test_mgvi_optimize_step" begin

    model = ModelPolyfit.model
    true_params = ModelPolyfit.true_params
    starting_point = ModelPolyfit.starting_point

    rng = MersenneTwister(145)
    data = rand(rng, model(true_params), 1)[1]

    first_iteration = mgvi_kl_optimize_step(rng,
                                            model, data, starting_point;
                                            jacobian_func=FwdRevADJacobianFunc,
                                            residual_sampler=ImplicitResidualSampler)

    next_iteration = first_iteration
    next_iteration = mgvi_kl_optimize_step(rng,
                                           model, data, next_iteration.result;
                                           jacobian_func=FwdRevADJacobianFunc,
                                           residual_sampler=ImplicitResidualSampler)

end


PIDNames = Dict(
    1 => "d",
    -1 => "d̄",
    2 => "u",
    -2 => "ū",
    3 => "s",
    -3 => "s̄",
    4 => "c",
    -4 => "c̄",
    5 => "b",
    -5 => "b̄",
    6 => "t",
    -6 => "t̄",
    11 => "e⁻",
    -11 => "e⁺",
    12 => "νₑ",
    -12 => "ν̄ₑ",
    13 => "μ⁻",
    -13 => "μ⁺",
    14 => "ν(μ)",
    -14 => "ν̄(μ)",
    15 => "τ⁻",
    -15 => "τ⁺",
    16 => "ν(τ)",
    -16 => "ν̄(τ)",
    17 => "τ'",
    -17 => "τ'",
    18 => "ν(τ')",
    -18 => "ν̄(τ')",
    21 => "g",
    22 => "γ",
    23 => "Z⁰",
    24 => "W⁺",
    -24 => "W⁻",
    25 => "H⁰",
    111 => "π⁰",
    113 => "ρ(770)⁰",
    115 => "a₂(1320)⁰",
    117 => "ρ₃(1690)⁰",
    119 => "a₄(1970)⁰",
    130 => "K(L)",
    211 => "π⁺",
    -211 => "π⁻",
    213 => "ρ(770)⁺",
    -213 => "ρ(770)⁻",
    215 => "a₂(1320)⁺",
    -215 => "a₂(1320)⁻",
    217 => "ρ₃(1690)⁺",
    -217 => "ρ₃(1690)⁻",
    219 => "a₄(1970)⁺",
    -219 => "a₄(1970)⁻",
    221 => "η",
    223 => "ω(782)",
    225 => "f₂(1270)",
    227 => "ω₃(1670)",
    229 => "f₄(2050)",
    310 => "K(S)",
    311 => "K⁰",
    -311 => "K̄⁰",
    313 => "K*(892)",
    -313 => "K̄*(892)",
    315 => "K(2)*(1430)",
    -315 => "K̄(2)*(1430)",
    317 => "K(3)*(1780)",
    -317 => "K̄(3)*(1780)",
    319 => "K(4)*(2045)",
    -319 => "K̄(4)*(2045)",
    321 => "K⁺",
    -321 => "K⁻",
    323 => "K*(892)",
    -323 => "K̄*(892)",
    325 => "K(2)*(1430)",
    -325 => "K̄(2)*(1430)",
    327 => "K(3)*(1780)",
    -327 => "K̄(3)*(1780)",
    329 => "K(4)*(2045)",
    -329 => "K̄(4)*(2045)",
    331 => "eta'(958)",
    333 => "ϕ(1020)",
    335 => "f(2)'(1525)",
    337 => "ϕ₃(1850)",
    411 => "D⁺",
    -411 => "D⁻",
    413 => "D*(2010)",
    -413 => "D̄*(2010)",
    415 => "D(2)*(2460)",
    -415 => "D̄(2)*(2460)",
    421 => "D⁰",
    -421 => "D̄⁰",
    423 => "D*(2007)",
    -423 => "D̄*(2007)",
    425 => "D(2)*(2460)",
    -425 => "D̄(2)*(2460)",
    431 => "Dₛ⁺",
    -431 => "Dₛ⁻",
    433 => "D(s)*",
    -433 => "D̄(s)*",
    435 => "D(s2)*(2573)",
    -435 => "D̄(s2)*(2573)",
    441 => "eta(c)(1S)",
    443 => "J/ψ(1S)",
    445 => "χ(c2)(1P)",
    511 => "B⁰",
    -511 => "B̄⁰",
    513 => "B*",
    -513 => "B*",
    515 => "B(2)*(5747)",
    -515 => "B̄(2)*(5747)",
    521 => "B⁺",
    -521 => "B⁻",
    523 => "B*",
    -523 => "B̄*",
    525 => "B(2)*(5747)",
    -525 => "B̄(2)*(5747)",
    531 => "Bₛ⁰",
    -531 => "B̄ₛ⁰",
    533 => "B(s)*",
    -533 => "B̄(s)*",
    535 => "B(s2)*(5840)",
    -535 => "B̄(s2)*(5840)",
    541 => "B(c)",
    -541 => "B̄(c)",
    553 => "Υ(1S)",
    555 => "χ(b2)(1P)",
    1103 => "(dd)₁",
    -1103 => "(dd)₁",
    1112 => "Δ(1620)⁻",
    -1112 => "Δ̄(1620)⁺",
    1114 => "Δ(1232)⁻",
    -1114 => "Δ̄(1232)⁺",
    1116 => "Δ(1905)⁻",
    -1116 => "Δ̄(1905)⁺",
    1118 => "Δ(1950)⁻",
    -1118 => "Δ̄(1950)⁺",
    1212 => "Δ(1620)⁰",
    -1212 => "Δ̄(1620)⁰",
    1214 => "N(1520)⁰",
    -1214 => "N̄(1520)⁰",
    1216 => "Δ(1905)⁰",
    -1216 => "Δ̄(1905)⁰",
    1218 => "N(2190)⁰",
    -1218 => "N̄(2190)⁰",
    2101 => "(ud)₀",
    -2101 => "(ud)₀",
    2103 => "(ud)₁",
    -2103 => "(ud)₁",
    2112 => "n",
    -2112 => "n̄",
    2114 => "Δ(1232)⁰",
    -2114 => "Δ̄(1232)⁰",
    2116 => "N(1675)⁰",
    -2116 => "N̄(1675)⁰",
    2118 => "Δ(1950)⁰",
    -2118 => "Δ̄(1950)⁰",
    2122 => "Δ(1620)⁺",
    -2122 => "Δ̄(1620)⁻",
    2124 => "N(1520)⁺",
    -2124 => "N̄(1520)⁻",
    2126 => "Δ(1905)⁺",
    -2126 => "Δ̄(1905)⁻",
    2128 => "N(2190)⁺",
    -2128 => "N̄(2190)⁻",
    2203 => "(uu)₁",
    -2203 => "(uu)₁",
    2212 => "p",
    -2212 => "p̄",
    2214 => "Δ(1232)⁺",
    -2214 => "Δ̄(1232)⁻",
    2216 => "N(1675)⁺",
    -2216 => "N̄(1675)⁻",
    2218 => "Δ(1950)⁺",
    -2218 => "Δ̄(1950)⁻",
    2222 => "Δ(1620)⁺⁺",
    -2222 => "Δ̄(1620)⁻⁻",
    2224 => "Δ(1232)⁺⁺",
    -2224 => "Δ̄(1232)⁻⁻",
    2226 => "Δ(1905)⁺⁺",
    -2226 => "Δ̄(1905)⁻⁻",
    2228 => "Δ(1950)⁺⁺",
    -2228 => "Δ̄(1950)⁻⁻",
    3101 => "(sd)₀",
    -3101 => "(sd)₀",
    3103 => "(sd)₁",
    -3103 => "(sd)₁",
    3112 => "Σ⁻",
    -3112 => "Σ̄⁺",
    3114 => "Σ(1385)⁻",
    -3114 => "Σ̄(1385)⁺",
    3116 => "Σ(1775)⁻",
    -3116 => "Σ̄(1775)⁺",
    3118 => "Σ(2030)⁻",
    -3118 => "Σ̄(2030)⁺",
    3122 => "Λ",
    -3122 => "Λ̄",
    3124 => "Λ(1520)",
    -3124 => "Λ̄(1520)",
    3126 => "Λ(1820)",
    -3126 => "Λ̄(1820)",
    3128 => "Λ(2100)",
    -3128 => "Λ̄(2100)",
    3201 => "(su)₀",
    -3201 => "(su)₀",
    3203 => "(su)₁",
    -3203 => "(su)₁",
    3212 => "Σ⁰",
    -3212 => "Σ̄⁰",
    3214 => "Σ(1385)⁰",
    -3214 => "Σ̄(1385)⁰",
    3216 => "Σ(1775)⁰",
    -3216 => "Σ̄(1775)⁰",
    3218 => "Σ(2030)⁰",
    -3218 => "Σ̄(2030)⁰",
    3222 => "Σ⁺",
    -3222 => "Σ̄⁻",
    3224 => "Σ(1385)⁺",
    -3224 => "Σ̄(1385)⁻",
    3226 => "Σ(1775)⁺",
    -3226 => "Σ̄(1775)⁻",
    3228 => "Σ(2030)⁺",
    -3228 => "Σ̄(2030)⁻",
    3303 => "(ss)₁",
    -3303 => "(ss)₁",
    3312 => "Ξ⁻",
    -3312 => "Ξ̄⁺",
    3314 => "Ξ(1530)⁻",
    -3314 => "Ξ̄(1530)⁺",
    3322 => "Ξ⁰",
    -3322 => "Ξ̄⁰",
    3324 => "Ξ(1530)⁰",
    -3324 => "Ξ̄(1530)⁰",
    3334 => "Ω⁻",
    -3334 => "Ω̄⁺",
    4101 => "(cd)₀",
    -4101 => "(cd)₀",
    4103 => "(cd)₁",
    -4103 => "(cd)₁",
    4112 => "Σ(c)(2455)",
    -4112 => "Σ̄(c)(2455)",
    4114 => "Σ(c)(2520)",
    -4114 => "Σ̄(c)(2520)",
    4122 => "Λ(c)",
    -4122 => "Λ̄(c)",
    4132 => "Ξ(c)",
    -4132 => "Ξ̄(c)",
    4201 => "(cu)₀",
    -4201 => "(cu)₀",
    4203 => "(cu)₁",
    -4203 => "(cu)₁",
    4212 => "Σ(c)(2455)",
    -4212 => "Σ̄(c)(2455)",
    4214 => "Σ(c)(2520)",
    -4214 => "Σ̄(c)(2520)",
    4222 => "Σ(c)(2455)",
    -4222 => "Σ̄(c)(2455)",
    4224 => "Σ(c)(2520)",
    -4224 => "Σ̄(c)(2520)",
    4232 => "Ξ(c)",
    -4232 => "Ξ̄(c)",
    4301 => "(cs)₀",
    -4301 => "(cs)₀",
    4303 => "(cs)₁",
    -4303 => "(cs)₁",
    4312 => "Ξ(c)'",
    -4312 => "Ξ̄(c)'",
    4314 => "Ξ(c)(2645)",
    -4314 => "Ξ̄(c)(2645)",
    4322 => "Ξ(c)'",
    -4322 => "Ξ̄(c)'",
    4324 => "Ξ(c)(2645)",
    -4324 => "Ξ̄(c)(2645)",
    4332 => "Ω(c)",
    -4332 => "Ω̄(c)",
    4334 => "Ω(c)(2770)",
    -4334 => "Ω̄(c)(2770)",
    4403 => "(cc)₁",
    -4403 => "(cc)₁",
    5101 => "(bd)₀",
    -5101 => "(bd)₀",
    5103 => "(bd)₁",
    -5103 => "(bd)₁",
    5112 => "Σ(b)",
    -5112 => "Σ̄(b)",
    5114 => "Σ(b)*",
    -5114 => "Σ̄(b)*",
    5122 => "Λ(b)",
    -5122 => "Λ̄(b)",
    5132 => "Ξ(b)",
    -5132 => "Ξ̄(b)",
    5201 => "(bu)₀",
    -5201 => "(bu)₀",
    5203 => "(bu)₁",
    -5203 => "(bu)₁",
    5222 => "Σ(b)",
    -5222 => "Σ̄(b)",
    5224 => "Σ(b)*",
    -5224 => "Σ̄(b)*",
    5232 => "Ξ(b)",
    -5232 => "Ξ̄(b)",
    5301 => "(bs)₀",
    -5301 => "(bs)₀",
    5303 => "(bs)₁",
    -5303 => "(bs)₁",
    5332 => "Ω(b)",
    -5332 => "Ω̄(b)",
    5401 => "(bc)₀",
    -5401 => "(bc)₀",
    5403 => "(bc)₁",
    -5403 => "(bc)₁",
    5503 => "(bb)₁",
    -5503 => "(bb)₁",
    10111 => "a₀(1450)⁰",
    10113 => "b₁(1235)⁰",
    10115 => "π₂(1670)⁰",
    10211 => "a₀(1450)⁺",
    -10211 => "a₀(1450)⁻",
    10213 => "b₁(1235)⁺",
    -10213 => "b₁(1235)⁻",
    10215 => "π₂(1670)⁺",
    -10215 => "π₂(1670)⁻",
    10221 => "f₀(1370)",
    10223 => "h₁(1170)",
    10225 => "η₂(1645)",
    10311 => "K(0)*(1430)",
    -10311 => "K̄(0)*(1430)",
    10313 => "K₁(1270)⁰",
    -10313 => "K̄₁(1270)⁰",
    10315 => "K₂(1770)⁰",
    -10315 => "K̄₂(1770)⁰",
    10321 => "K(0)*(1430)",
    -10321 => "K̄(0)*(1430)",
    10323 => "K₁(1270)⁺",
    -10323 => "K₁(1270)⁻",
    10325 => "K₂(1770)⁺",
    -10325 => "K₂(1770)⁻",
    10331 => "f₀(1710)",
    10333 => "h₁(1415)",
    10411 => "D(0)*(2300)",
    -10411 => "D̄(0)*(2300)",
    10421 => "D(0)*(2300)",
    -10421 => "D̄(0)*(2300)",
    10423 => "D₁(2420)⁰",
    -10423 => "D̄₁(2420)⁰",
    10431 => "D(s0)*(2317)",
    -10431 => "D̄(s0)*(2317)",
    10433 => "Dₛ₁(2536)⁺",
    -10433 => "Dₛ₁(2536)⁻",
    10441 => "χ(c0)(1P)",
    10443 => "h(c)(1P)",
    10551 => "χ(b0)(1P)",
    10553 => "h(b)(1P)",
    11112 => "Δ(1900)⁻",
    -11112 => "Δ̄(1900)⁺",
    11114 => "Δ(1700)⁻",
    -11114 => "Δ̄(1700)⁺",
    11116 => "Δ(1930)⁻",
    -11116 => "Δ̄(1930)⁺",
    11212 => "Δ(1900)⁰",
    -11212 => "Δ̄(1900)⁰",
    11216 => "Δ(1930)⁰",
    -11216 => "Δ̄(1930)⁰",
    12112 => "N(1440)⁰",
    -12112 => "N̄(1440)⁰",
    12114 => "Δ(1700)⁰",
    -12114 => "Δ̄(1700)⁰",
    12116 => "N(1680)⁰",
    -12116 => "N̄(1680)⁰",
    12122 => "Δ(1900)⁺",
    -12122 => "Δ̄(1900)⁻",
    12126 => "Δ(1930)⁺",
    -12126 => "Δ̄(1930)⁻",
    12212 => "N(1440)⁺",
    -12212 => "N̄(1440)⁻",
    12214 => "Δ(1700)⁺",
    -12214 => "Δ̄(1700)⁻",
    12216 => "N(1680)⁺",
    -12216 => "N̄(1680)⁻",
    12222 => "Δ(1900)⁺⁺",
    -12222 => "Δ̄(1900)⁻⁻",
    12224 => "Δ(1700)⁺⁺",
    -12224 => "Δ̄(1700)⁻⁻",
    12226 => "Δ(1930)⁺⁺",
    -12226 => "Δ̄(1930)⁻⁻",
    13112 => "Σ(1660)⁻",
    -13112 => "Σ̄(1660)⁺",
    13114 => "Σ(1670)⁻",
    -13114 => "Σ̄(1670)⁺",
    13116 => "Σ(1915)⁻",
    -13116 => "Σ̄(1915)⁺",
    13122 => "Λ(1404)",
    -13122 => "Λ̄(1404)",
    13124 => "Λ(1690)",
    -13124 => "Λ̄(1690)",
    13126 => "Λ(1830)",
    -13126 => "Λ̄(1830)",
    13212 => "Σ(1660)⁰",
    -13212 => "Σ̄(1660)⁰",
    13214 => "Σ(1670)⁰",
    -13214 => "Σ̄(1670)⁰",
    13216 => "Σ(1915)⁰",
    -13216 => "Σ̄(1915)⁰",
    13222 => "Σ(1660)⁺",
    -13222 => "Σ̄(1660)⁻",
    13224 => "Σ(1670)⁺",
    -13224 => "Σ̄(1670)⁻",
    13226 => "Σ(1915)⁺",
    -13226 => "Σ̄(1915)⁻",
    13314 => "Ξ(1820)⁻",
    -13314 => "Ξ̄(1820)⁺",
    13324 => "Ξ(1820)⁰",
    -13324 => "Ξ̄(1820)⁰",
    14122 => "Λ(c)(2595)",
    -14122 => "Λ(c)(2595)",
    20113 => "a₁(1260)⁰",
    20213 => "a₁(1260)⁺",
    -20213 => "a₁(1260)⁻",
    20223 => "f₁(1285)",
    20313 => "K₁(1400)⁰",
    -20313 => "K̄₁(1400)⁰",
    20315 => "K₂(1820)⁰",
    -20315 => "K̄₂(1820)⁰",
    20323 => "K₁(1400)⁺",
    -20323 => "K₁(1400)⁻",
    20325 => "K₂(1820)⁺",
    -20325 => "K₂(1820)⁻",
    20333 => "f₁(1420)",
    20433 => "Dₛ₁(2460)⁺",
    -20433 => "Dₛ₁(2460)⁻",
    20443 => "χ(c1)(1P)",
    20553 => "χ(b1)(1P)",
    20555 => "Υ₂(1D)",
    21112 => "Δ(1910)⁻",
    -21112 => "Δ̄(1910)⁺",
    21114 => "Δ(1920)⁻",
    -21114 => "Δ̄(1920)⁺",
    21212 => "Δ(1910)⁰",
    -21212 => "Δ̄(1910)⁰",
    21214 => "N(1700)⁰",
    -21214 => "N̄(1700)⁰",
    22112 => "Δ(1910)⁰",
    -22112 => "Δ̄(1910)⁰",
    22114 => "Δ(1920)⁰",
    -22114 => "Δ̄(1920)⁰",
    22122 => "Δ(1910)⁺",
    -22122 => "Δ̄(1910)⁻",
    22124 => "N(1700)⁺",
    -22124 => "N̄(1700)⁻",
    22212 => "N(1535)⁺",
    -22212 => "N̄(1535)⁻",
    22214 => "Δ(1920)⁺",
    -22214 => "Δ̄(1920)⁻",
    22222 => "Δ(1910)⁺⁺",
    -22222 => "Δ̄(1910)⁻⁻",
    22224 => "Δ(1920)⁺⁺",
    -22224 => "Δ̄(1920)⁻⁻",
    23112 => "Σ(1750)⁻",
    -23112 => "Σ̄(1750)⁺",
    23114 => "Σ(1940)⁻",
    -23114 => "Σ̄(1940)⁺",
    23122 => "Λ(1600)",
    -23122 => "Λ̄(1600)",
    23124 => "Λ(1890)",
    -23124 => "Λ̄(1890)",
    23126 => "Λ(2110)",
    -23126 => "Λ̄(2110)",
    23212 => "Σ(1750)⁰",
    -23212 => "Σ̄(1750)⁰",
    23214 => "Σ(1940)⁰",
    -23214 => "Σ̄(1940)⁰",
    23222 => "Σ(1750)⁺",
    -23222 => "Σ̄(1750)⁻",
    23224 => "Σ(1940)⁺",
    -23224 => "Σ̄(1940)⁻",
    30113 => "ρ(1700)⁰",
    30213 => "ρ(1700)⁺",
    -30213 => "ρ(1700)⁻",
    30223 => "ω(1650)",
    30313 => "K*(1680)",
    -30313 => "K*(1680)",
    30323 => "K*(1680)",
    -30323 => "K*(1680)",
    30443 => "ψ(3770)",
    31114 => "Δ(1600)⁻",
    -31114 => "Δ̄(1600)⁺",
    31214 => "N(1720)⁰",
    -31214 => "N̄(1720)⁰",
    32112 => "N(1650)⁰",
    -32112 => "N̄(1650)⁰",
    32114 => "Δ(1600)⁰",
    -32114 => "Δ̄(1600)⁰",
    32124 => "N(1720)⁺",
    -32124 => "N̄(1720)⁻",
    32212 => "N(1650)⁺",
    -32212 => "N̄(1650)⁻",
    32214 => "Δ(1600)⁺",
    -32214 => "Δ̄(1600)⁻",
    32224 => "Δ(1600)⁺⁺",
    -32224 => "Δ̄(1600)⁻⁻",
    33122 => "Λ(1670)",
    -33122 => "Λ̄(1670)",
    42112 => "N(1710)⁰",
    -42112 => "N̄(1710)⁰",
    42212 => "N(1710)⁺",
    -42212 => "N̄(1710)⁻",
    43122 => "Λ(1800)",
    -43122 => "Λ̄(1800)",
    53122 => "Λ(1810)",
    -53122 => "Λ̄(1810)",
    100111 => "π(1300)⁰",
    100113 => "ρ(1450)⁰",
    100211 => "π(1300)⁺",
    -100211 => "π(1300)⁻",
    100213 => "ρ(1450)⁺",
    -100213 => "ρ(1450)⁻",
    100221 => "η(1295)",
    100313 => "K*(1410)",
    -100313 => "K̄*(1410)",
    100321 => "K(1460)⁺",
    -100321 => "K(1460)⁻",
    100323 => "K*(1410)",
    -100323 => "K̄*(1410)",
    100331 => "η(1475)",
    100333 => "ϕ(1680)",
    100441 => "eta(c)(2S)",
    100443 => "ψ(2S)",
    100445 => "χ(c2)(3930)",
    100553 => "Υ(2S)",
    100555 => "χ(b2)(2P)",
    103316 => "Ξ(1950)⁻",
    -103316 => "Ξ̄(1950)⁺",
    103326 => "Ξ(1950)⁰",
    -103326 => "Ξ̄(1950)⁰",
    104122 => "Λ(c)(2625)",
    -104122 => "Λ̄(c)(2625)",
    104312 => "Ξ(c)(2815)",
    -104312 => "Ξ̄(c)(2815)",
    104314 => "Ξ(c)(2790)",
    -104314 => "Ξ̄(c)(2790)",
    104322 => "Ξ(c)(2815)",
    -104322 => "Ξ̄(c)(2815)",
    104324 => "Ξ(c)(2790)",
    -104324 => "Ξ̄(c)(2790)",
    110551 => "χ(b0)(2P)",
    120553 => "χ(b1)(2P)",
    200553 => "Υ(3S)",
    203312 => "Ξ(1690)⁻",
    -203312 => "Ξ̄(1690)⁺",
    203316 => "Ξ(2030)⁻",
    -203316 => "Ξ̄(2030)⁺",
    203322 => "Ξ(1690)⁰",
    -203322 => "Ξ̄(1690)⁰",
    203326 => "Ξ(2030)⁰",
    -203326 => "Ξ̄(2030)⁰",
    203338 => "Ω(2250)⁻",
    -203338 => "Ω̄(2250)⁺",
    204126 => "Λ(c)(2880)",
    -204126 => "Λ̄(c)(2880)",
    300553 => "Υ(4S)",
    1000223 => "ω(1420)",
    9000111 => "a₀(980)⁰",
    9000113 => "π₁(1400)⁰",
    9000115 => "a₂(1700)⁰",
    9000211 => "a₀(980)⁺",
    -9000211 => "a₀(980)⁻",
    9000213 => "π₁(1400)⁺",
    -9000213 => "π₁(1400)⁻",
    9000215 => "a₂(1700)⁺",
    -9000215 => "a₂(1700)⁻",
    9000221 => "f₀(500)",
    9000311 => "K(0)*(700)",
    -9000311 => "K(0)*(700)",
    9000321 => "K(0)*(700)",
    -9000321 => "K(0)*(700)",
    9000323 => "K₁(1650)⁺",
    -9000323 => "K₁(1650)⁻",
    9000325 => "K₂(1580)⁺",
    -9000325 => "K₂(1580)⁻",
    9000329 => "K₄(2500)⁺",
    -9000329 => "K₄(2500)⁻",
    9000443 => "ψ(4040)",
    9000553 => "Υ(10860)",
    9010111 => "π(1800)⁰",
    9010113 => "π₁(1600)⁰",
    9010211 => "π(1800)⁺",
    -9010211 => "π(1800)⁻",
    9010213 => "π₁(1600)⁺",
    -9010213 => "π₁(1600)⁻",
    9010221 => "f₀(980)",
    9010315 => "K(2)*(1980)",
    -9010315 => "K̄(2)*(1980)",
    9010321 => "K(1830)⁺",
    -9010321 => "K(1830)⁻",
    9010325 => "K(2)*(1980)",
    -9010325 => "K̄(2)*(1980)",
    9010327 => "K₃(2320)⁺",
    -9010327 => "K₃(2320)⁻",
    9010443 => "ψ(4160)",
    9010553 => "Υ(11020)",
    9020113 => "a₁(1640)⁰",
    9020213 => "a₁(1640)⁺",
    -9020213 => "a₁(1640)⁻",
    9020221 => "η(1405)",
    9020321 => "K(0)*(1950)",
    -9020321 => "K̄(0)*(1950)",
    9020325 => "K₂(2250)⁺",
    -9020325 => "K₂(2250)⁻",
    9020443 => "ψ(4415)",
    9030221 => "f₀(1500)",
    9050225 => "f₂(1950)",
    9060225 => "f₂(2010)",
    9080225 => "f₂(2300)",
    9090225 => "f₂(2340)"
)


struct DiffRuleVariable
    diffrule::DiffRule
    vars::OffsetVector{Variable,Vector{Variable}}
    funcold::Func
end
function DiffRuleVariable(∂m::∂Model, dr::DiffRule, func::Func)
    nparents = num_parents(dr)
    vars = OffsetVector{Variable}(undef, -nparents:last(last(dr.sections)))
    drv = DiffRuleVariable( dr, vars, func )
    vparents = parents(func)
    moldvars = ∂m.mold.vars
    for i ∈ -nparents:-1
        # drv.vars[i] = getvar!(∂m, moldvars[vparents[i + nparents + 1]])
        drv.vars[i] = getvar!(∂m, moldvars[vparents[i + nparents + 1]])
    end
    drv.vars[0] = var0 = getvar!(∂m.m, diffsym(moldvars[func.output[]]))
    drv
end

function add_section!(∂m::∂Model, drv::DiffRuleVariable, sectionid::Int)
    @unpack vars, funcold, diffrule = drv
    @unpack sections, returns, instructions, dependencies = diffrule
    section = sections[sectionid]
    if isone(sectionid)
        ret = returns[1]
    else
        str = ReverseDiffExpressionsBase.section_two_returns(diffrule)
        if sectionid == 2
            str || return
            ret = returns[2]
        else
            if (vars[sectionid - num_parents(diffrule) - 3]).tracked
                ret = returns[sectionid + str - 1]
            else
                return
            end
        end
    end
    for i ∈ section
        func = Func(instructions[i], false)
        # (instructions[i].instr === :constrain_pullback!) && uses!(func, ∂m.gradvar)
        for j ∈ dependencies[i]
            uses!(func, vars[j])
        end
        vars[i] = if i == ret && sectionid != 2
            if isone(sectionid)
                retv = getvar!(∂m, ∂m.mold.vars[funcold.output[]])
            else
                parent = vars[sectionid - num_parents(diffrule) - 3]
                retv = getvar!(∂m.m, diffsym(name(parent)))
            end
        else
            retv = getvar!(∂m.m, gensym(:temp))
        end
        retv.tracked = true
        returns!(func, retv)
        retold = ∂m.mold.vars[funcold.output[]]
        addfunc!(∂m.m, func)
    end
end

function add_fallback_diffrule!(∂m::∂Model, func::Func)
    @unpack m, mold = ∂m
    # add func as var, return DiffRuleVariable
    v = addvarref!(∂m.m, asexpr(func.instr))
    drv = DiffRuleVariable(∂m, ReverseDiffExpressionsBase.FALLBACK_RULES[length(parents(func))], func)

    retvsym = gensym(:rrule_LHS)
    rrulefunc = Func(Instruction(:ChainRules,:rrule), false)
    drv.vars[1] = rruleretv = addvar!(∂m.m, retvsym)
    rruleretv.tracked = true
    returns!(rrulefunc, rruleretv)
    uses!(rrulefunc, v)
    foreach(arg -> uses!(rrulefunc, arg), @view(drv.vars[begin:-1]))
    addfunc!(m, rrulefunc)
    funcfirst = Func(Instruction(:Base,:first), false)
    uses!(funcfirst, rruleretv)
    drv.vars[2] = retvar = getvar!(∂m, mold.vars[func.output[]])
    returns!(funcfirst, retvar)
    retvar.tracked = true
    addfunc!(m, funcfirst)
    drv
end
function add_diffrule!(∂m::∂Model, dr::DiffRule, func::Func)
    drv = DiffRuleVariable(∂m, dr, func)
    add_section!(∂m, drv, 1)
    drv
end

function add_view_diff!(∂m::∂Model, func::Func)
    gi_forward = Func(func.instr, false)
    gi_reverse = Func(func.instr, false)
    vparents = parents(func)
    moldvars = ∂m.mold.vars
    p = moldvars[first(vparents)]
    inputvf = getvar!(∂m, p); inputvf.tracked = true
    inputvr = getvar!(∂m.m, diffsym(name(p))); inputvr.tracked = true
    uses!(gi_forward, inputvf)
    uses!(gi_reverse, inputvr)
    for i ∈ @view(vparents[2:end])
        indold = moldvars[i]
        if isref(indold)
            ind = addvar!(∂m.m, Symbol(""))
            ind.ref = indold.ref
        else
            ind = addvar!(∂m.m, name(indold))
        end
        uses!(gi_forward, ind)
        uses!(gi_reverse, ind)
    end
    retv = moldvars[func.output[]]
    retvf = getvar!(∂m, retv); retvf.tracked = true
    returns!(gi_forward, retvf)
    retvr = getvar!(∂m.m, diffsym(name(retv))); retvr.tracked = true
    returns!(gi_reverse, retvr)
    push!(inputvr.paireddeps, identifier(retvr))
    addfunc!(∂m.m, gi_forward)
    addfunc!(∂m.m, gi_reverse)
    nothing
end

function add_probdist_diff!(∂m::∂Model, func::Func)
    ∂loglik = Func(func.instr, true)
    moldvars = ∂m.mold.vars
    for i ∈ parents(func)
        uses!(∂loglik, getvar!(∂m, moldvars[i]))
    end
    tup = addvar!(∂m.m, gensym(:tup))
    tup.tracked = true
    returns!(∂loglik, tup)
    addfunc!(∂m.m, ∂loglik)
    getconstindex!(∂m.m, ∂m.m.vars[0], tup, 1)
    ∂tup = addvar!(∂m.m, gensym(:∂tup))
    getconstindex!(∂m.m, ∂tup, tup, 2)
    for (ind,i) ∈ enumerate(parents(func))
        p = moldvars[i]
        if istracked(p)
            retv = getvar!(∂m.m, diffsym(name(p))); retv.tracked = true
            getconstindex!(∂m.m, retv, ∂tup, ind)
        end
    end
end
function add_constrain_diff!(∂m::∂Model, funcold::Func)
    drvc = DiffRuleVariable(∂m, DERIVATIVERULES[InstructionArgs(Instruction(:DistributionParameters,:constrain),3)], funcold)
    constrain_pullback = Func(Instruction(:constrain_pullback!), false)
    uses!(constrain_pullback, ∂m.gradvar)
    for i ∈ -3:-1
        uses!(constrain_pullback, drvc.vars[i])
    end
    tempv = getvar!(∂m.m, gensym(:constrainpullbacktup)); tempv.tracked = true
    returns!(constrain_pullback, tempv)
    addfunc!(∂m.m, constrain_pullback)
    retv = getvar!(∂m, ∂m.mold.vars[funcold.output[]]); retv.tracked = true
    getconstindex!(∂m.m, retv, tempv, 1)
    ∂retv = getvar!(∂m.m, diffsym(name(retv))); ∂retv.tracked = true
    getconstindex!(∂m.m, ∂retv, tempv, 2)

    constrain_reverse = Func(Instruction(:constrain_reverse!), false)
    uses!(constrain_reverse, ∂retv)
    uses!(constrain_reverse, drvc.vars[-1])
    nothingvar = addvar!(∂m.m, gensym(:nothing))
    returns!(constrain_reverse, nothingvar)
    push!(∂m.m.targets, identifier(nothingvar))
    addfunc!(∂m.m, constrain_reverse)
    nothing
end
function add_no_diff!(∂m::∂Model, funcold::Func)
    func = Func(funcold.instr, funcold.probdistapi)
    for i ∈ parents(funcold)
        uses!(func, getvar!(∂m, ∂m.mold.vars[i]))
    end
    returns!(func, getvar!(∂m, ∂m.mold.vars[funcold.output[]]))
    addfunc!(∂m.m, func)
    nothing
end
function has_tracked_parents(∂m::∂Model, func::Func)
    any(i -> istracked(∂m.mold.vars[i]), parents(func))
end

function add_func_diff!(∂m::∂Model, func::Func)
    has_tracked_parents(∂m, func) || return add_no_diff!(∂m, func)
    if func.probdistapi
        return add_probdist_diff!(∂m, func)
    elseif func.instr.instr === :constrain
        return add_constrain_diff!(∂m, func)
    elseif isindexfunc(func)
        if iszero(∂m.mold.vars[func.output[]].varid)
            return add_no_diff!(∂m, func)
        else
            return add_view_diff!(∂m, func)
        end
    end
    instrargs = InstructionArgs(func)
    dr = get(DERIVATIVERULES, instrargs, nothing)
    drv = isnothing(dr) ? add_fallback_diffrule!(∂m, func) : add_diffrule!(∂m, dr, func)
    for i ∈ 2:length(drv.diffrule.sections)
        add_section!(∂m, drv, i)
    end
end


# FALLBACK function
function getLatticeByBondDistance(
            :: Type{L},
            unitcell        :: U,
            bonddistance    :: Integer,
            origin          :: Integer = 1
        ) :: L where {
            D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N},
            U<:AbstractUnitcell{S,B},
            DL,LLS,LLB,SL<:AbstractSite{LLS,DL},BL<:AbstractBond{LLB,0},
            L<:AbstractLattice{SL,BL,U}
        }

    # throw an error as this is not implemented yet
    error("Either function 'getLatticeByBondDistance' is not yet implemented for unitcells of type " * string(U) * ", i.e. N=" *
        string(N) * " / D=" * string(D) * " or you passed the wrong lattice type L = " * string(L))
end

# generate a lattice in 2D by bond distance (open boundary conditions)
function getLatticeByBondDistance(
            :: Type{L},
            unitcell        :: U,
            bonddistance    :: Integer,
            origin          :: Integer = 1
        ) :: L where {
            D,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,2},
            U<:AbstractUnitcell{S,B},
            DL,LLS,LLB,SL<:AbstractSite{LLS,DL},BL<:AbstractBond{LLB,0},
            L<:AbstractLattice{SL,BL,U}
        }

    # Format of checklist item
    #   Tuple{S, NTuple, Int64, Int64}
    # denoting
    #   Site - UC Indices (i,j,k,...) - Site Index (in UC) - Bonddistance to origin

    # checklist for all sites that are checked if added etc
    checklist = Tuple{S, NTuple{2,Int64}, Int64, Int64}[]
    # list for all accepted sites that are checked
    accepted  = Tuple{S, NTuple{2,Int64}, Int64, Int64}[]

    # List of all accepted bonds
    bond_list = BL[]

    # push the origin site to the checklist
    push!(
        checklist, (deepcopy(site(unitcell,origin)), (0,0), origin, 0)
    )

    # iterate while the checklist is not empty
    while size(checklist,1) > 0

        # get the item that was on the checklist for longest
        item_to_handle = popfirst!(checklist)

        # check if the item is already in the accepted list
        found = false
        for acc in accepted
            # check if a site with index acc[3] in unitcell acc[2] is already found
            if item_to_handle[2] == acc[2] && item_to_handle[3] == acc[3]
                found = true
                break
            end
        end
        # if the element is found, continue with the next checklist item
        if found
            continue
        end

        # Item is NOT in the accepted list, push it into the list
        push!(accepted, item_to_handle)
        # set the current item index to the length of the list
        index_from = length(accepted)

        # search through all connections of the unitcell to find the ones relevant persuing
        for b in bonds(unitcell)

            # check if the bond actually describes a bond emerging from the site
            if from(b) != item_to_handle[3]
                continue
            end

            # build up the element that the bond connects to

            # find out in which copy of unitcell the site is located
            uc_to = (
                item_to_handle[2][1] + wrap(b)[1],
                item_to_handle[2][2] + wrap(b)[2]
            )
            # find out which site index it is
            alpha_to = to(b)

            # find out if the item already exists within the accepted list
            index_to = -1
            for (index,item_handled) in enumerate(accepted)
                # check if the item is the particular item
                if item_handled[2] == uc_to && item_handled[3] == alpha_to
                    # save the list index
                    index_to = index
                    # break the loop
                    break
                end
            end
            # determine whether the element is already inside the list (index was set to > 0)
            if index_to > 0
                # create a new bonds
                bond_new_1 = newBond(
                    BL,
                    index_from,
                    index_to,
                    label(b),
                    NTuple{0,Int64}()
                )
                bond_new_2 = newBond(
                    BL,
                    index_to,
                    index_from,
                    label(b),
                    NTuple{0,Int64}()
                )
                # check if the bonds are already added
                if !(bond_new_1 in bond_list)
                    push!(bond_list, bond_new_1)
                end
                if !(bond_new_2 in bond_list)
                    push!(bond_list, bond_new_2)
                end
            else
                # the respective site is not in list, maybe add it to the checklist?
                if item_to_handle[4] < bonddistance
                    # it can be added, create a new site object
                    site_new = deepcopy(site(unitcell, to(b)))
                    # set the new position
                    point!(site_new, (uc_to[1] .* a1(unitcell)) .+ (uc_to[2] .* a2(unitcell)) .+ point(site_new))
                    # format for sites on checklist: ([pos_x, pos_y, pos_z], [uc_i, uc_j, uc_k], index_in_UC, bd_current)
                    push!(checklist, (
                        site_new,
                        uc_to,
                        to(b),
                        item_to_handle[4]+1
                    ))
                end
            end
        end
    end

    # insert missing connections (if (i to j) is present, insert (j to i))
    # TODO

    # save everything to a Lattice object and return it
    return newLattice(
            L,
        	Vector{Float64}[],
        	SL[newSite(SL, point(s[1]), label(s[1])) for s in accepted],
        	bond_list,
            unitcell
        )
end

# generate a lattice in 3D by bond distance (open boundary conditions)
function getLatticeByBondDistance(
            :: Type{L},
            unitcell        :: U,
            bonddistance    :: Integer,
            origin          :: Integer = 1
        ) :: L where {
            D,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,3},
            U<:AbstractUnitcell{S,B},
            DL,LLS,LLB,SL<:AbstractSite{LLS,DL},BL<:AbstractBond{LLB,0},
            L<:AbstractLattice{SL,BL,U}
        }

    # Format of checklist item
    #   Tuple{S, NTuple, Int64, Int64}
    # denoting
    #   Site - UC Indices (i,j,k,...) - Site Index (in UC) - Bonddistance to origin

    # checklist for all sites that are checked if added etc
    checklist = Tuple{S, NTuple{3,Int64}, Int64, Int64}[]
    # list for all accepted sites that are checked
    accepted  = Tuple{S, NTuple{3,Int64}, Int64, Int64}[]

    # List of all accepted bonds
    bond_list = BL[]

    # push the origin site to the checklist
    push!(
        checklist, (deepcopy(site(unitcell,origin)), (0,0,0), origin, 0)
    )

    # iterate while the checklist is not empty
    while size(checklist,1) > 0

        # get the item that was on the checklist for longest
        item_to_handle = popfirst!(checklist)

        # check if the item is already in the accepted list
        found = false
        for acc in accepted
            # check if a site with index acc[3] in unitcell acc[2] is already found
            if item_to_handle[2] == acc[2] && item_to_handle[3] == acc[3]
                found = true
                break
            end
        end
        # if the element is found, continue with the next checklist item
        if found
            continue
        end

        # Item is NOT in the accepted list, push it into the list
        push!(accepted, item_to_handle)
        # set the current item index to the length of the list
        index_from = length(accepted)

        # search through all connections of the unitcell to find the ones relevant persuing
        for b in bonds(unitcell)

            # check if the bond actually describes a bond emerging from the site
            if from(b) != item_to_handle[3]
                continue
            end

            # build up the element that the bond connects to

            # find out in which copy of unitcell the site is located
            uc_to = (
                item_to_handle[2][1] + wrap(b)[1],
                item_to_handle[2][2] + wrap(b)[2],
                item_to_handle[2][3] + wrap(b)[3]
            )
            # find out which site index it is
            alpha_to = to(b)

            # find out if the item already exists within the accepted list
            index_to = -1
            for (index,item_handled) in enumerate(accepted)
                # check if the item is the particular item
                if item_handled[2] == uc_to && item_handled[3] == alpha_to
                    # save the list index
                    index_to = index
                    # break the loop
                    break
                end
            end
            # determine whether the element is already inside the list (index was set to > 0)
            if index_to > 0
                # create a new bonds
                bond_new_1 = newBond(
                    BL,
                    index_from,
                    index_to,
                    label(b),
                    NTuple{0,Int64}()
                )
                bond_new_2 = newBond(
                    BL,
                    index_to,
                    index_from,
                    label(b),
                    NTuple{0,Int64}()
                )
                # check if the bonds are already added
                if !(bond_new_1 in bond_list)
                    push!(bond_list, bond_new_1)
                end
                if !(bond_new_2 in bond_list)
                    push!(bond_list, bond_new_2)
                end
            else
                # the respective site is not in list, maybe add it to the checklist?
                if item_to_handle[4] < bonddistance
                    # it can be added, create a new site object
                    site_new = deepcopy(site(unitcell, to(b)))
                    # set the new position
                    point!(site_new, (uc_to[1] .* a1(unitcell)) .+ (uc_to[2] .* a2(unitcell)) .+ (uc_to[3] .* a3(unitcell)) .+ point(site_new))
                    # format for sites on checklist: ([pos_x, pos_y, pos_z], [uc_i, uc_j, uc_k], index_in_UC, bd_current)
                    push!(checklist, (
                        site_new,
                        uc_to,
                        to(b),
                        item_to_handle[4]+1
                    ))
                end
            end
        end
    end

    # insert missing connections (if (i to j) is present, insert (j to i))
    # TODO

    # save everything to a Lattice object and return it
    return newLattice(
            L,
        	Vector{Float64}[],
        	SL[newSite(SL, point(s[1]), label(s[1])) for s in accepted],
        	bond_list,
            unitcell
        )
end

# wrapper
function getLatticeByBondDistance(
            unitcell        :: U,
            bonddistance    :: Integer,
            origin          :: Integer = 1
        ) :: Lattice{S,Bond{LB,0},U} where {
            D,N,LS,LB,S<:AbstractSite{LS,D},B<:AbstractBond{LB,N},
            U<:AbstractUnitcell{S,B}
        }

    # call the general function
    return getLatticeByBondDistance(Lattice{S,Bond{LB,0},U}, unitcell, bonddistance, origin)
end

# export the function
export getLatticeByBondDistance


using DifferentiableStateSpaceModels, SparseArrays, LinearAlgebra, Parameters, Test,
      TimerOutputs, BenchmarkTools

@testset "Dense SW First Order" begin
    m = @include_example_module(Examples.SW07)

    p = [10, 0.51, 10, 0, 0.7, 0.742, 0, 0, 0.24, 0.2696, 6.0144, 0.025, 1.5, 0.6361, 1.5,
         0.3243, 0.8087, 0.47, 0.6, 1.9423, 1.5, 1.488, 0.2347, 0.0593, 0.8762, 0.9977,
         0.5799, 0.9957, 0.7165, 0, 0, 0, 0.3982, 0.18]
    generate_perturbation(m, p)
    reset_timer!()
    sol = generate_perturbation(m, p)
    print_timer()
    @test sol.retcode == :Success
end