using Test
using ViennaRNA
using ViennaRNA: LibRNA
using Unitful

@testset "FoldCompound" begin
    for s in ["GGGAAACCC",
              ["GGG-A-CC-", "G-GGC-G-G"],
              ["GGG-A-CC-", "G-GGC-G-G", "CGC----CG"]]
        fc = FoldCompound(s)
        @test length(fc) == 9
        @test sum(size(fc)) == length(fc)

        # test FoldCompound properties
        @test fc.circular         == LibRNA.VRNA_MODEL_DEFAULT_CIRC
        @test fc.dangles          == LibRNA.VRNA_MODEL_DEFAULT_DANGLES
        @test fc.gquadruplex      == LibRNA.VRNA_MODEL_DEFAULT_GQUAD
        @test fc.has_exp_matrices == false
        @test fc.log_ML           == LibRNA.VRNA_MODEL_DEFAULT_LOG_ML
        @test fc.min_loop_size    == LibRNA.TURN
        @test fc.no_GU_basepairs  == LibRNA.VRNA_MODEL_DEFAULT_NO_GU
        @test fc.no_GU_closure    == LibRNA.VRNA_MODEL_DEFAULT_NO_GU_CLOSURE
        @test fc.no_lonely_pairs  == LibRNA.VRNA_MODEL_DEFAULT_NO_LP
        @test fc.nstrands         == 1
        @test fc.params_name      == "RNA - Turner 2004"
        @test fc.special_hairpins == LibRNA.VRNA_MODEL_DEFAULT_SPECIAL_HP
        @test fc.temperature      == 37u"°C"
        @test fc.uniq_ML          == LibRNA.VRNA_MODEL_DEFAULT_UNIQ_ML

        # options
        fc = FoldCompound(s; options=ViennaRNA.LibRNA.VRNA_OPTION_MFE)
        @test length(fc) == 9
        # use different energy parameter sets
        fc = FoldCompound(s; params=:RNA_Turner1999)
        @test length(fc) == 9
        @test_throws ArgumentError FoldCompound(s; params=:UNKNOWN_ENERGY_PARAMS)
        # use different temperatures to evaluate energies
        fc = FoldCompound(s; temperature=55u"°C")
        @test length(fc) == 9
        @test fc.temperature == 55u"°C"
        fc = FoldCompound(s; temperature=310u"K")
        @test length(fc) == 9
        @test fc.temperature == 310u"K"
        @test_throws Unitful.DimensionError FoldCompound(s; temperature=100u"m")

        # circular
        fc = FoldCompound(s; circular=true)
        @test length(fc) == 9
        @test fc.circular == true
        # dangles
        fc = FoldCompound(s; dangles=1)
        @test length(fc) == 9
        @test fc.dangles == 1
        @test_throws ArgumentError FoldCompound(s; dangles=-1)
        @test_throws ArgumentError FoldCompound(s; dangles=4)
        # gquadruplex
        fc = FoldCompound(s; gquadruplex=true)
        @test length(fc) == 9
        @test fc.gquadruplex == true
        # log_ML
        fc = FoldCompound(s; log_ML=true)
        @test length(fc) == 9
        @test fc.log_ML == true
        # min_loop_size
        fc = FoldCompound(s; min_loop_size=2)
        @test length(fc) == 9
        @test fc.dangles == 2
        @test_throws ArgumentError FoldCompound(s; min_loop_size=-1)
        # no_GU_basepairs
        fc = FoldCompound(s; no_GU_basepairs=true)
        @test length(fc) == 9
        @test fc.no_GU_basepairs == true
        # no_GU_closure
        fc = FoldCompound(s; no_GU_closure=true)
        @test length(fc) == 9
        @test fc.no_GU_closure == true
        # no_lonely_pairs
        fc = FoldCompound(s; no_lonely_pairs=true)
        @test length(fc) == 9
        @test fc.no_lonely_pairs == true
        # special_hairpins
        fc = FoldCompound(s; special_hairpins=false)
        @test length(fc) == 9
        @test fc.special_hairpins == false
        # uniq_ML
        fc = FoldCompound(s; uniq_ML=true)
        @test length(fc) == 9
        @test fc.uniq_ML == true
        # show
        buf = IOBuffer()
        show(buf, MIME("text/plain"), fc)
        @test length(String(take!(buf))) > 0
    end

    # multiple strands
    fc = FoldCompound("GGG&AAA&CCC")
    @test length(fc) == 9
    @test size(fc) == (3, 3, 3)
    @test fc.nstrands == 3

    # multistrand alifold (not supported yet)
    @test_throws ArgumentError fc = FoldCompound(["GGG&AAA&CCC", "G-G&AA-&CCA"])
    # @test length(fc) == 9
    # @test size(fc) == (3, 3, 3)
    # @test fc.nstrands == 3

    # multiple strands don't work in comparative (alifold) mode
    @test_throws ArgumentError FoldCompound(["GGGA", "GCGG", "G"])
    @test_throws ArgumentError FoldCompound(["GG&C", "AA&U"])

    # test setting of min_loop_size with single and multiple strands
    # - this didn't work in ViennaRNA-2.5.0 with multiple strands
    #   (e.g. "GGG&CCC") and h=3 (min_loop_size was always set to 0 by
    #   ViennaRNA). Works as of ViennaRNA-2.5.1.
    for h = 0:5
        fc = FoldCompound("GGG"; min_loop_size=h)
        @test fc.min_loop_size == h
        fc = FoldCompound("GGG&CCC"; min_loop_size=h)
        @test fc.min_loop_size == h
        fc = FoldCompound("GGG&CCC&AAAA"; min_loop_size=h)
        @test fc.min_loop_size == h
        fc = FoldCompound(["GGG", "G-C"]; min_loop_size=h)
        @test fc.min_loop_size == h
        fc = FoldCompound(["GGG", "G-C", "-AG"]; min_loop_size=h)
        @test fc.min_loop_size == h
    end
end

@testset "Pairtable" begin
    pt = Pairtable("(((...)))")
    @test length(pt) == 9
    # getindex
    @test pt[1] == 9
    @test pt[2] == 8
    @test pt[3] == 7
    @test pt[4] == 0
    @test pt[5] == 0
    @test pt[6] == 0
    @test pt[7] == 3
    @test pt[8] == 2
    @test pt[9] == 1
    @test_throws BoundsError pt[0]
    @test_throws BoundsError pt[length(pt) + 1]
    # setindex!
    pt[3] = 4
    @test pt[3] == 4
    @test pt[4] == 3
    @test pt[7] == 0
    @test_throws ArgumentError pt[1] = 1
    @test_throws ArgumentError pt[0] = 0
    @test_throws ArgumentError pt[length(pt) + 1] = 0
    @test_throws ArgumentError pt[1] = -1
    @test_throws ArgumentError pt[1] = length(pt) + 1
    # show
    buf = IOBuffer()
    show(buf, MIME("text/plain"), pt)
    @test length(String(take!(buf))) > 0
    # basepairs
    @test Set(basepairs(Pairtable("(((...)))"))) == Set([(1,9), (2,8), (3,7)])
end

@testset "bp_distance" begin
    @test bp_distance("((..))", "(....)") == 1
    @test_throws ArgumentError bp_distance("()", ".")
end

@testset "tree_edit_dist" begin
    @test tree_edit_dist("...", "()") == 3.0f0
end

@testset "mean_bp_distance" begin
    seq = "GGGAAACCC"
    fc = FoldCompound(seq)
    @test_throws ArgumentError mean_bp_distance(fc)
    partfn(fc)
    @test mean_bp_distance(fc) isa AbstractFloat
    @test mean_bp_distance(seq) isa AbstractFloat
end

@testset "ensemble_defect" begin
    seq = "GGGAAACCCC"
    str = "(((....)))"
    fc = FoldCompound(seq)
    partfn(fc)
    @test ensemble_defect(fc, str) isa AbstractFloat
    @test ensemble_defect(fc, Pairtable(str)) isa AbstractFloat
    @test_throws ArgumentError ensemble_defect(fc, ".")
    @test_throws ArgumentError ensemble_defect(fc, Pairtable("."))
    @test ensemble_defect(seq, Pairtable(str)) isa AbstractFloat
    @test ensemble_defect(seq, str) isa AbstractFloat
end

@testset "prob_of_structure" begin
    seq = "GGGAAACCCC"
    str = "(((....)))"
    fc = FoldCompound(seq)
    partfn(fc)
    @test prob_of_structure(fc, str) isa AbstractFloat
    @test_throws ArgumentError prob_of_structure(fc, ".")
    @test prob_of_structure(seq, str) isa AbstractFloat
end

@testset "energy" begin
    seq = "GGGAAAACCCC"
    str = "(((....)))."
    fc = FoldCompound(seq)
    @test energy(fc, str) isa Unitful.Quantity
    @test energy(seq, str) isa Unitful.Quantity
    redirect_stdout(devnull) do
        # suppress stdout output from ViennaRNA for verbose option
        @test energy(fc, str; verbose=true) isa Unitful.Quantity
        @test energy(seq, str; verbose=true) isa Unitful.Quantity
    end
    @test_throws ArgumentError energy(fc, ".")
    @test_throws ArgumentError energy("(...)", ".")
    # test different energy parameter sets
    @test energy(FoldCompound(seq; params=:RNA_Turner1999), str) ≈ -2.90u"kcal/mol" atol=1e-3u"kcal/mol"
    @test energy(FoldCompound(seq; params=:RNA_Andronescu2007), str) ≈ -2.15u"kcal/mol" atol=1e-3u"kcal/mol"
    # test different temperatures
    @test energy(FoldCompound(seq; temperature=30u"°C"), str) ≈ -3.52u"kcal/mol" atol=1e-3u"kcal/mol"
    @test energy(FoldCompound(seq; temperature=37u"°C"), str) ≈ -2.90u"kcal/mol" atol=1e-3u"kcal/mol"
    @test energy(FoldCompound(seq; temperature=55u"°C"), str) ≈ -1.24u"kcal/mol" atol=1e-3u"kcal/mol"
    @test energy(FoldCompound(seq; temperature=323.15u"K"), str) ≈ -1.69u"kcal/mol" atol=1e-3u"kcal/mol" # 323.15u"K" is 50u"°C"
end

@testset "mfe" begin
    seq = "GGGAAACCCC"
    fc = FoldCompound(seq)
    @test mfe(fc) isa Tuple{String,Unitful.Quantity}
    @test mfe(seq) isa Tuple{String,Unitful.Quantity}
end

@testset "partfn" begin
    seq = "GGGAAACCCC"
    fc = FoldCompound(seq)
    @test partfn(fc) isa Tuple{String,Unitful.Quantity}
    @test partfn(seq) isa Tuple{String,Unitful.Quantity}
end

@testset "bpp" begin
    seq = "GGGAAACCCC"
    n = length(seq)

    fc = FoldCompound(seq)
    @test_throws ArgumentError bpp(fc)
    partfn(fc)
    p = bpp(fc)
    @test eltype(p) <: AbstractFloat
    @test size(p) == (n,n)

    p = bpp(seq)
    @test eltype(p) <: AbstractFloat
    @test size(p) == (n,n)
end

@testset "pbacktrack" begin
    seq = "GGGGGAAAAACCCCCCCCAUUCA"
    fc = FoldCompound(seq; uniq_ML=true)
    partfn(fc)
    @test length(pbacktrack(fc)) == 1
    @test length(pbacktrack(seq)) == 1
    s = pbacktrack(fc; num_samples=10)
    @test length(s) == 10
    s = pbacktrack(fc; num_samples=5,
                   options=ViennaRNA.LibRNA.VRNA_PBACKTRACK_NON_REDUNDANT)
    @test length(s) == 5
end

@testset "mea" begin
    seq = "GGGGGAAAAACCCCCCCCAUUCA"
    fc = FoldCompound(seq)
    partfn(fc)
    @test mea(fc) isa Tuple{String,AbstractFloat}
    @test mea(fc; gamma=1.0) isa Tuple{String,AbstractFloat}
    @test mea(fc; gamma=0.5) isa Tuple{String,AbstractFloat}
    @test mea(seq) isa Tuple{String,AbstractFloat}
    @test mea(seq; gamma=1.0) isa Tuple{String,AbstractFloat}
    @test mea(seq; gamma=0.5) isa Tuple{String,AbstractFloat}
    fc = FoldCompound(seq)
    @test_throws ArgumentError mea(fc)
end

@testset "centroid" begin
    seq = "GGGGGAAAAACCCCCCCCAUUCA"
    fc = FoldCompound(seq)
    partfn(fc)
    @test centroid(fc) isa Tuple{String,AbstractFloat}
    @test centroid(seq) isa Tuple{String,AbstractFloat}
    fc = FoldCompound(seq)
    @test_throws ArgumentError centroid(fc)
end

@testset "subopt" begin
    seq = "GGGGGAAAAACCCCCCCCAUUCA"
    fc = FoldCompound(seq; uniq_ML=true)
    s = subopt(fc; delta=5u"kcal/mol")
    @test s isa Vector{Tuple{String,Unitful.Quantity}}
    s = subopt(fc; delta=5u"kcal/mol", sorted=true)
    @test s isa Vector{Tuple{String,Unitful.Quantity}}
    s = subopt(seq; delta=5u"kcal/mol")
    @test s isa Vector{Tuple{String,Unitful.Quantity}}
    s = subopt(seq; delta=5u"kcal/mol", sorted=true)
    @test s isa Vector{Tuple{String,Unitful.Quantity}}
end

@testset "subopt_zuker" begin
    seq = "GGGGAAAACCCC"
    # TODO: the following sequences generate warnings
    # seq = "GGGUAAACCCAUUCAC"
    # seq = "GGGGGAAAAACCCCCCCCAUUCA"
    fc = FoldCompound(seq)
    s = subopt_zuker(fc)
    @test s isa Vector{Tuple{String,Unitful.Quantity}}
    s = subopt_zuker(seq)
    @test s isa Vector{Tuple{String,Unitful.Quantity}}
end

@testset "inverse_fold" begin
    s = "((((....))))"
    @test inverse_fold("A"^length(s), s) isa Tuple{String,AbstractFloat}
    @test_throws ArgumentError inverse_fold("A", "()")
    @test inverse_pf_fold("A"^length(s), s) isa Tuple{String,Unitful.Quantity}
    @test_throws ArgumentError inverse_pf_fold("A", "()")
end

@testset "neighbors" begin
    fc = FoldCompound("GGGAAACCC")
    pt = Pairtable(".((...)).")
    @test neighbors(fc, pt) == [[(-2, -8)], [(-3, -7)], [(1, 9)]]
end

@testset "plot_coords" begin
    function test_plot_xy(s, x, y)
        @test x isa Vector{Float32}
        @test length(x) == length(s)
        @test y isa Vector{Float32}
        @test length(y) == length(s)
    end
    s = "(((...)))"
    pt = Pairtable(s)

    # test: plot_coords(::String), plot_coords(::Pairtable)
    x, y = plot_coords(s)
    test_plot_xy(s, x, y)
    x, y = plot_coords(pt)
    test_plot_xy(pt, x, y)
    for plot_type in (:simple, :naview, :circular, :turtle, :puzzler)
        x, y = plot_coords(s; plot_type)
        test_plot_xy(s, x, y)
        x, y = plot_coords(pt; plot_type)
        test_plot_xy(pt, x, y)
    end
    @test_throws ArgumentError x, y = plot_coords(s; plot_type = :unknown)
    @test_throws ArgumentError x, y = plot_coords(pt; plot_type = :unknown)
    # zero-sized inputs
    x, y = plot_coords("")
    test_plot_xy("", x, y)
    x, y = plot_coords(Pairtable(""))
    test_plot_xy(Pairtable(""), x, y)
end

@testset "heat_capacity" begin
    s = "GGGAAACCC"
    fc = FoldCompound(s)
    Tmin = 10u"°C"
    Tmax = 60u"°C"
    Tincrement = 1u"°C"
    mpoints = 2
    # TODO: this range call doesn't work without ustrip
    n = length(range(start=ustrip(Tmin), stop=ustrip(Tmax), step=ustrip(Tincrement)))

    hcs = heat_capacity(fc, Tmin, Tmax, Tincrement; mpoints)
    @test hcs isa Vector{Tuple{typeof(1.0f0u"°C"),typeof(1.0f0u"kcal/mol/K")}}
    @test length(hcs) == n
    hcs = heat_capacity(fc, Tmin, Tmax; mpoints)
    @test hcs isa Vector{Tuple{typeof(1.0f0u"°C"),typeof(1.0f0u"kcal/mol/K")}}
    @test length(hcs) == n
    hcs = heat_capacity(s, Tmin, Tmax, Tincrement; mpoints)
    @test hcs isa Vector{Tuple{typeof(1.0f0u"°C"),typeof(1.0f0u"kcal/mol/K")}}
    @test length(hcs) == n
    hcs = heat_capacity(s, Tmin, Tmax; mpoints)
    @test hcs isa Vector{Tuple{typeof(1.0f0u"°C"),typeof(1.0f0u"kcal/mol/K")}}
    @test length(hcs) == n
end

include("utils.jl")
include("plot_structure.jl")


struct ContourLike end

contourlevels(spacing=1.7, n=12) = (spacing^i for i=0:(n-1))

# 1D plot
@recipe function f(A::NMRData{T,1}; normalize=true) where T
    Afwd = DimensionalData.forwardorder(A) # make sure data axes are in forwards order
    dim = dims(Afwd, 1)

    # force 1D to be a line plot
    seriestype := :path
    markershape --> :none

    # set default title
    title --> DimensionalData.refdims_title(Afwd)
    legend --> false

    xguide --> "$(Afwd[dim,:label]) chemical shift / ppm"
    xflip --> true
    xgrid --> false
    xtick_direction --> :out

    yguide --> ""
    yshowaxis --> false
    yticks --> nothing

    delete!(plotattributes, :normalize)
    val(dim), parent(Afwd)./(normalize ? scale(A) : 1)
end

# multiple 1D plots
@recipe function f(v::Vector{<:NMRData{T,1} where T}; normalize=true, vstack=false)
    # force 1D to be a line plot
    seriestype := :path
    markershape --> :none

    # get the first entry to determine axis label
    Afwd = DimensionalData.forwardorder(v[1])
    dim = dims(Afwd, 1)

    xguide --> "$(Afwd[dim,:label]) chemical shift / ppm"
    xflip --> true
    xgrid --> false
    xtick_direction --> :out

    yguide --> ""
    yshowaxis --> false
    yticks --> nothing

    delete!(plotattributes, :vstack)
    delete!(plotattributes, :normalize)

    voffset = 0
    vdelta = maximum([
            maximum(abs.(A)) / (normalize ? A[:ns]*A[:rg] : 1)
            for A in v]) / length(v)

    # TODO add guide lines
    # if vstack
    #     yticks --> voffset .+ (0:length(v)-1)*vdelta
    # else
    #     yticks --> [0,]
    # end

    for A in v
        @series begin
            seriestype := :path
            Afwd = DimensionalData.forwardorder(A) # make sure data axes are in forwards order
            dim = dims(Afwd, 1)
            label --> label(A)
            val(dim), parent(Afwd)./(normalize ? scale(A) : 1) .+ voffset
        end
        vstack && (voffset += vdelta)
    end
end

# 2D plot
@recipe f(d::D) where {D<:NMRData{T,2} where T} = SimpleTraits.trait(HasPseudoDimension{D}), d

@recipe function f(::Type{Not{HasPseudoDimension{D}}}, d::D) where {D<:NMRData{T, 2} where T}
    dfwd = DimensionalData.forwardorder(d) # make sure data axes are in forwards order
    dfwd = DimensionalData.maybe_permute(dfwd, (YDim, XDim))
    y, x = dims(dfwd)

    # set default title
    title --> DimensionalData.refdims_title(dfwd)
    legend --> false
    framestyle --> :box

    xguide --> "$(dfwd[x,:label]) chemical shift / ppm"
    xflip --> true
    xgrid --> false
    xtick_direction --> :out

    yguide --> "$(dfwd[y,:label]) chemical shift / ppm"
    yflip --> true
    ygrid --> false
    ytick_direction --> :out

    # generate light and dark colours for plot contours, based on supplied colour
    # - create a 5-tone palette with the same hue as the passed colour, and select the
    # fourth and second entries to provide dark and light shades
    basecolor = get(plotattributes, :linecolor, :blue)
    colors = sequential_palette(hue(convert(HSV,parse(Colorant, basecolor))),5)[[4,2]]

    #delete!(plotattributes, :normalize)
    #scale = normalize ? A[:ns]*A[:rg] : 1
    #val(dim), parent(A) ./ scale
    @series begin
        levels --> 5*dfwd[:noise].*contourlevels()
        linecolor := colors[1]
        primary := true
        val(x), val(y), data(dfwd)
    end
    @series begin
        levels --> -5*dfwd[:noise].*contourlevels()
        linecolor := colors[2]
        primary := false
        val(x), val(y), data(dfwd)
    end
end

@recipe function f(::Type{HasPseudoDimension{D}}, d::D) where {D<:NMRData{T, 2} where T}
    # TODO define recipe for pseudo2D data
    dfwd = DimensionalData.forwardorder(d) # make sure data axes are in forwards order
    dfwd = DimensionalData.maybe_permute(dfwd, (YDim, XDim))
    y, x = dims(dfwd)
    @info "plot recipe for pseudo2D data not yet defined"
    val(x), val(y), data(dfwd)
end

# multiple 2D plots
@recipe f(v::Vector{D}; normalize=true) where {D<:NMRData{T,2}} where {T} = SimpleTraits.trait(HasPseudoDimension{D}), v

@recipe function f(::Type{Not{HasPseudoDimension{D}}}, v::Vector{D}; normalize=true) where {D<:NMRData{T,2}} where {T}
    @info "plotting vector of 2D NMR data"
    n = length(v)
    hues = map(h->HSV(h,0.5,0.5), (0:n-1) .* (360/n))

    # force 1D to be a line plot
    seriestype := :contour

    # get the first entry to determine axis label
    dfwd = DimensionalData.forwardorder(v[1]) # make sure data axes are in forwards order
    dfwd = DimensionalData.maybe_permute(dfwd, (YDim, XDim))
    y, x = dims(dfwd)

    # set default title
    title --> DimensionalData.refdims_title(dfwd)
    legend --> false
    framestyle --> :box

    xguide --> "$(dfwd[x,:label]) chemical shift / ppm"
    xflip --> true
    xgrid --> false
    xtick_direction --> :out

    yguide --> "$(dfwd[y,:label]) chemical shift / ppm"
    yflip --> true
    ygrid --> false
    ytick_direction --> :out

    delete!(plotattributes, :normalize)

    h = 0.0
    for d in v
        dfwd = DimensionalData.forwardorder(d) # make sure data axes are in forwards order
        dfwd = DimensionalData.maybe_permute(dfwd, (YDim, XDim))
        y, x = dims(dfwd)
        label --> label(d)
        # generate light and dark colours for plot contours, based on supplied colour
        # - create a 5-tone palette with the same hue as the passed colour, and select the
        # fourth and second entries to provide dark and light shades
        colors = sequential_palette(h, 5)[[4,2]]

        @series begin
            levels --> 5*dfwd[:noise].*contourlevels()
            linecolor := colors[1]
            primary := true
            val(x), val(y), data(dfwd)
        end
        @series begin
            levels --> -5*dfwd[:noise].*contourlevels()
            linecolor := colors[2]
            primary := false
            val(x), val(y), data(dfwd)
        end
        h += 360.0/n
    end
end

@recipe function f(::Type{HasPseudoDimension{D}}, v::Vector{D}; normalize=true) where {D<:NMRData{T,2}} where {T}
    @info "plot recipe for vector of pseudo2D NMR data not yet implemented"
    delete!(plotattributes, :normalize)
    # TODO just make repeat calls to single plot recipe
    for d in v
        @series begin
            HasPseudoDimension{D}, d
        end
    end
end


abstract type AbstractPassive{T} <: AbstractObject{T} end

mutable struct Passive{T} <: AbstractPassive{T}
    name::String
    id::Int
    potential::Union{Symbol,T}
    temperature::Union{T,Missing}
    material::NamedTuple
    charge_density_model::AbstractChargeDensity{T}
    geometry::AbstractGeometry{T}
    geometry_positive::Vector{AbstractGeometry{T}}
    geometry_negative::Vector{AbstractGeometry{T}}

    Passive{T}() where T <: SSDFloat = new{T}()
end

function Passive{T}(dict::Dict, inputunit_dict::Dict{String,Unitful.Units}) where T <: SSDFloat
    pass = Passive{T}()
    haskey(dict, "name") ? pass.name = dict["name"] : pass.name = "external part"
    haskey(dict, "id") ? pass.id = dict["id"] : pass.id = -1
    haskey(dict,"potential") ? pass.potential = T(dict["potential"]) : pass.potential = :floating
    haskey(dict, "temperature") ? pass.temperature = T(dict["temperature"]) : pass.temperature = missing
    pass.material = material_properties[materials[dict["material"]]]
    pass.charge_density_model = if haskey(dict, "charge_density_model")
        ChargeDensity(T, dict["charge_density_model"], inputunit_dict)
    else
        ZeroChargeDensity{T}()
    end
    pass.geometry = Geometry(T, dict["geometry"], inputunit_dict)
    pass.geometry_positive, pass.geometry_negative = get_decomposed_volumes(pass.geometry)
    return pass
end

function println(io::IO, d::Passive{T}) where {T}
    println("\t________"*"Passive{$T} $(d.name) $(d.id)"*"________\n")
    println("\t---General Properties---")
    println("\t-Potential: \t\t $(d.potential) V")
    println("\t-Material: \t $(d.material.name)")
    println()
end
print(io::IO, d::Passive{T}) where {T} = print(io, "Passive $(d.name) - id $(d.id) - $(d.potential) V")
show(io::IO, d::Passive) = print(io, d)
show(io::IO,::MIME"text/plain", d::Passive) = show(io, d)


"""
Data structure
--------------
mutable struct: SELF.recording

Functions
---------
SELF.recording.filter()

SELF.recording.resample()

SELF.recording.detect_spikes()
"""
module Neurons

export recording, normalize_trace, filter_trace, resample, detect_spikes, detect_spikes1, get_PSD, extract_spikes, extract_features, cluster_spikes, find_redundant, extract_isih,
		plot_detection, plot_waveforms

using DSP
using Images
using Statistics
using MultivariateStats
using GaussianMixtures
using Clustering
using LinearAlgebra
using Distances
using PyPlot

using ..MyUtils

# main type
"""
Fields
------
+ file::String
+ sf::Real
+ sf_preprocessed::Real
+ n_channels::Real
+ trace::AbstractArray # supposed to contain sample X channel
+ trial_marker::AbstractArray
+ trace_preprocessed::AbstractArray
+ preprocessing_steps::Vector{String}
+ spike_idx::AbstractArray
+ detection_threshold::AbstractArray
+ isih::AbstractArray
+ spikeforms::AbstractArray # channel (val_source) X channel (t_source) X time of spike/spikeform X sample
+ features::AbstractArray # channel X feature space representation of spikes
+ spike_class::AbstractArray # channel X clusters (as spike indices)
+ BIC::AbstractArray
"""
mutable struct recording
    ### FIELDS
    file::String
    sf::Real
    sf_preprocessed::Real
    n_channels::Real
    trace::Array{Array{Float64,2}}#AbstractArray # supposed to contain sample X channel
    trial_marker::AbstractArray
    trace_preprocessed::Array{Array{Float64,2}}#AbstractArray
    preprocessing_steps::Vector{String}
    spike_idx::AbstractArray
	detection_threshold::AbstractArray
    isih::AbstractArray
    spikeforms::AbstractArray # channel (val_source) X channel (t_source) X time of spike/spikeform X sample
    features::AbstractArray # channel X feature space representation of spikes
    spike_class::AbstractArray # channel X clusters (as spike indices)
    BIC::AbstractArray

    ### CONSTRUCTOR
    # recording(file, sf, trace, trial_marker) =
    #    new(file, sf, 0, trace, trial_marker, Array{Float64,2}(undef,0,0), Vector{String}(undef,0), Array{Float64,3}Array{Float64,2}(undef,0,0,0))
    # recording() = new("", 0, 0, Array{Float64,2}(undef,0,0), Array{Float64,2}(undef,0,0), Array{Float64,2}(undef,0,0), Vector{String}(undef,0), Array{Float64,3}(undef,0,0,0))
    recording() = new("", 0, 0, 0, [], [], [], Vector{String}(undef,0), [], [], [], [], [], [])
end

# functions

function normalize_trace(n::Neurons.recording, what="trace")

	x = deepcopy( getfield( n, Symbol(what) ) )
	sd = std.(x)
	mn = mean.(x)
	[x[ii]=(x[ii] .- mn[ii]) ./ sd[ii] for ii in 1:n.n_channels] ### normalization
	n.trace_preprocessed = x
	n.preprocessing_steps = ["normalized"];

end

"""
Documentation ...
"""
function filter_trace(n::Neurons.recording, filter_type::String, freqs::AbstractArray, filter_design::AbstractArray, filter_order::AbstractArray, bandwidth::AbstractArray, from_chnnls::AbstractArray, what::String="trace")
"""
    filter(n::recording, band::Tuple{Real}, filter_type)
    filter_type: lowpass, highpass, bandpass, notch
    freqs = highpass: freq | lowpass: freq | passband (Wn1,Wn2) frequency | Notch: array of frequencies | CommonAverage: channels used for averaging
    what: spcifies whether filter works on trace or trace_preprocessed. If the latter trace_preprocessed is replaced, respectively.
"""

    ######################
    ### raw trace
    ######################
    if what == "trace"
        nyq = n.sf/2;
        freqs_n = freqs./nyq; # normalized to nquist frequncy
        n.trace_preprocessed = []
        n.sf_preprocessed = n.sf # restore original sampling frequency

        if filter_type == "Bandpass"
            rt = Bandpass( freqs_n[1], freqs_n[2] );
            dm = getfield( Main, Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            ### loop over channels and filter trace
            for trc in n.trace
                #trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                #push!(n.trace_preprocessed, filtfilt(f, trc_padded)[pads+1:end-pads]);
                push!(n.trace_preprocessed, filtfilt(f, trc));
            end
        elseif filter_type == "Lowpass"
            rt = Lowpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for trc in n.trace
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # push!(n.trace_preprocessed, filtfilt(f, trc_padded)[pads+1:end-pads]);
                push!(n.trace_preprocessed, filtfilt(f, trc));
            end
        elseif filter_type == "Highpass"
            rt = Highpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for trc in n.trace
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # push!(n.trace_preprocessed, filt(f, trc_padded)[pads+1:end-pads]);
                push!(n.trace_preprocessed, filtfilt(f, trc));
            end
        elseif filter_type == "Notch"
            for trc in n.trace
                trc = flatten(trc)
                for ff in freqs_n.*nyq
                    nn = iirnotch(ff, bandwidth[1]; fs=n.sf)
                    trc = filtfilt(nn, trc);
                end
                push!(n.trace_preprocessed, trc);
            end
        elseif filter_type == "CommonAverage"
            chnnls = freqs
            average_trace = [median([n.trace[ch][idx] for ch in chnnls]) for idx in 1:length(n.trace[1])]
            # n.trace_preprocessed = [channel .- average_trace for channel in n.trace[from_chnnls]]
			[n.trace_preprocessed[ch] = n.trace_preprocessed[ch] .- average_trace for ch in from_chnnls]
            #n.trace_preprocessed = [ch .- mean(n.trace) for ch in n.trace]
        end
        n.preprocessing_steps = [string( filter_type, ": ",freqs )];

    ######################
    ### preprocessed trace
    ######################
    elseif what == "trace_preprocessed" && !isempty(n.preprocessing_steps)
        nyq = n.sf_preprocessed/2;
        freqs_n = freqs./nyq; # normalized to nquist frequncy

        if filter_type == "Bandpass"
            rt = Bandpass( freqs_n[1], freqs_n[2] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # n.trace_preprocessed[itrc] = filt(f, trc_padded)[pads+1:end-pads];
                n.trace_preprocessed[itrc] = filtfilt(f, trc);
            end
        elseif  filter_type == "Lowpass"
            rt = Lowpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # n.trace_preprocessed[itrc] = filt(f, trc_padded)[pads+1:end-pads];
                n.trace_preprocessed[itrc] = filtfilt(f, trc);
            end
        elseif filter_type == "Highpass"
            rt = Highpass( freqs_n[1] );
            dm = getfield( Main,Symbol(filter_design[1]) )(filter_order[1]);
            f = digitalfilter(rt, dm);
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                # trc_padded = flatten([[.0 for ii in 1:pads], trc, [.0 for ii in 1:pads]])
                # n.trace_preprocessed[itrc] = filt(f, trc_padded)[pads+1:end-pads];
                n.trace_preprocessed[itrc] = filtfilt(f, trc);
            end
        elseif filter_type == "Notch"
            for (itrc, trc) in enumerate(n.trace_preprocessed)
                trc = flatten(trc)
                for ff in freqs_n.*nyq
                    nn = iirnotch(ff, bandwidth[1]; fs=n.sf_preprocessed)
                    trc = filtfilt(nn, trc);
                end
                n.trace_preprocessed[itrc] = trc;
            end
        elseif filter_type == "CommonAverage"
            chnnls = freqs
            average_trace = [median([n.trace_preprocessed[ch][idx] for ch in chnnls]) for idx in 1:length(n.trace_preprocessed[1])]
            # n.trace_preprocessed = [channel .- average_trace for channel in n.trace_preprocessed]
			[n.trace_preprocessed[ch] = n.trace_preprocessed[ch] .- average_trace for ch in from_chnnls]
			#n.trace_preprocessed = [ch .- mean(n.trace_preprocessed) for ch in n.trace_preprocessed]
        end
        push!(n.preprocessing_steps, string( filter_type, ": ", freqs ));
    else
        print(n.preprocessing_steps)
        @warn "SELF.trace_preprocessed is empty. Process SELF.trace first"
    end
end

"""
Arguments:

+ n::Neurons.recording
+ chans::Array{Int64,1}
+ SNR::Array{Float64,1}
+ lockout::Int64=30;
+ polarity::String="negative",
+ ww::Array{Int64,1}=[5,15],
+ stp::Int64=1,
+ w_lockout::Array{Int64,1}=[10],
+ method::String="petrantonakis"

"""
function detect_spikes(
	n::Neurons.recording,
	chans::Array{Int64,1},
	SNR::Array{Float64,1},
	lockout::Int64=30;
	#
	polarity::String="negative",
	ww::Array{Int64,1}=[5,15],
	stp::Int64=1,
	w_lockout::Array{Int64,1}=[10],
	method::String="petrantonakis")
#
#   spikes are stored as:s
#   n.spike_idx: channel X w X Dw/indices of peaks (Dw only if method is "petrantonakis")
#
#   w: window length if method is "petrantonakis"
#   (for consistency if method is "threshold" the same structure is stored )

### WAVEFORM/THRESHOLD BASED SPIKE DETECITON
    if method == "petrantonakis" ### ref: https://www.frontiersin.org/articles/10.3389/fnins.2015.00452/full
        # Dw = [ [[] for ii in 1:n.n_channels] for jj in eachindex(ww) ] # ww X channel
        # Spks = [ [CartesianIndex{1}[] for ii in 1:n.n_channels] for jj in eachindex(ww) ] # ww X channel
        Dw_spk = [ [[[],[]] for ii in eachindex(ww)] for jj in 1:n.n_channels ] ### ch X w X Dw,spk_idx
		n.detection_threshold = []

        # for (ii,ch) in enumerate([n.trace_preprocessed[2]]) ### loop over channels
		for ii in chans ### loop over channels

			ch = n.trace_preprocessed[ii]

			for (jj,w) in enumerate(ww) ### loop over w

                ### compute the Dw signal (euklidean distance signal)
                f = ch[collect(0:w-1).+collect(1:stp:length(ch)-w+1)'] ### cut signal in overlapping chunks
				Dw_spk[ii][jj][1] = flatten(sqrt.(sum(diff(f, dims=2).^2, dims=1))); ### compute euclidean distance between consecutive chunks

                ### normalization of Dw
				sig = median( Dw_spk[ii][jj][1] ./ 0.6745 )
				thr = SNR[jj]*sig
				append!(n.detection_threshold, thr)
                # Dw_spk[ii][jj][1] = Dw_spk[ii][jj][1] ./ th ### store the standardized Dw signal
                ### find all local mixima in Dw
                loc_maxima = findlocalmaxima(Float64.(Dw_spk[ii][jj][1])) ### find all local maxima
                ### remove local mixima < thr
                Dw_spk[ii][jj][2] = loc_maxima[findall(Dw_spk[ii][jj][1][loc_maxima] .> thr)] ### extract all local maxima > SNR, SNR depends on w

                ### remove redundant triggered spikes and keep the maximum of each
                if length(Dw_spk[ii][jj][2]) > 1

					Dw_spk[ii][jj][2] = find_redundant( Dw_spk[ii][jj][1], Dw_spk[ii][jj][2], lockout );

				end

			end #for w: jj

            ### remove spikes that are detected only by one value of w
			loop_cntr = 0
			while length(Dw_spk[ii][1][2]) != length(Dw_spk[ii][2][2]) ### necessary if there a spikes in value of w with less spikes that are not in w with more spikes

				loop_cntr += 1

				if loop_cntr == 10 ### 2 should be enough

					break;

				end

				if !any( [(length(Dw_spk[ii][1][2])==0), (length(Dw_spk[ii][2][2])==0)] )

					if length(Dw_spk[ii][1][2]) > length(Dw_spk[ii][2][2])

						dd = pairwise(Euclidean(), cartesian2matrix(Dw_spk[ii][2][2])', cartesian2matrix(Dw_spk[ii][1][2])', dims=2)
					    deleteat!(Dw_spk[ii][1][2], findall(flatten(.!any(dd .< w_lockout, dims=1)))) ### delete all spikes in Dw_spk[ii][1][2] that do not have a partner spike in Dw_spk[ii][2][2]

					elseif length(Dw_spk[ii][1][2]) < length(Dw_spk[ii][2][2])

						dd = pairwise(Euclidean(), cartesian2matrix(Dw_spk[ii][1][2])', cartesian2matrix(Dw_spk[ii][2][2])', dims=2)
					    deleteat!(Dw_spk[ii][2][2], findall(flatten(.!any(dd .< w_lockout, dims=1))))

					end

				else ### if one has no spikes at all

					Dw_spk[ii][1][2] = [];
					Dw_spk[ii][2][2] = [];
					# Dw_spk[ii][1][1] = []
					# Dw_spk[ii][2][1] = []

				end

			end

			if length(Dw_spk[ii][1][2]) != length(Dw_spk[ii][2][2])

				@info "@Channel ", ii, ": different spike numbers for values of w"

			end

			### scatter plot of w1 X w2
	        if (length(spk_idx_w1)-length(spk_idx_w1) == 0) && length(recording.spike_idx[channel][1][2]) != 0

	            ax1 = subplot(1,3,1)
	            w1 = recording.spike_idx[channel][1][1][spk_idx_w1]
	            w2 = recording.spike_idx[channel][2][1][spk_idx_w2]
	            pols = cartesian2polar(w1, w2)
	            r = pols[:,1]
	            phi = pols[:,2]
	            ax1.scatter( r, phi, color=pcols[2])

	        end

        end #for channels: ii

    ### THRESHOLD BASED SPIKE DETECTION
    elseif method == "threshold"

		Dw_spk = [ [[[],[]] for ii in [1]] for jj in 1:n.n_channels ] ### ch X w X DW,spk_idx
		n.detection_threshold = ones(n.n_channels) .* NaN
        x = n.trace_preprocessed[chans]

        sig = [median(abs.(ch) ./ 0.6745 ) for ch in x]
		thr = SNR[1] .* sig
		n.detection_threshold[chans] = thr

        # x = .-x ./ th

        if polarity == "both"

			ind_max = findlocalmaxima.(x)
            ind_max = [[lm[1] for lm in ch] for ch in ind_max] # cartesian coordinates to indices
            ind_min = findlocalminima.(x)
            ind_min = [[lm[1] for lm in ch] for ch in ind_min]
            ind_max = [ind_max[ii][findall(x[ii][ind_max[ii]] .> thr[ii])] for ii in eachindex(chans)]    # local minima with minimal SNR
            ind_min = [ind_min[ii][findall(x[ii][ind_min[ii]] .< -thr[ii])] for ii in eachindex(chans)]    # local maxima with minimal SNR
            ind = [sort(flatten([ind_max[ii], ind_min[ii]])) for ii in 1:length(ind_max)] # channel X index_of_spikes

		elseif polarity == "positive"

			ind = findlocalmaxima.(x)
            ind = [[lm[1] for lm in ch] for ch in ind] # cartesian coordinates to indices
            ind = [ind[ii][findall(x[ii][ind[ii]] .> thr[ii])] for ii in eachindex(chans)]    # local minima with minimal SNR

		elseif polarity == "negative"

			ind = findlocalminima.(x)
            ind = [[lm[1] for lm in ch] for ch in ind] # cartesian coordinates to indices
            ind = [ind[ii][findall(x[ii][ind[ii]] .< -thr[ii])] for ii in eachindex(chans)]

		end

        ### remove redundantly triggered spikes and store the maximum respectively
        for (ii,ch) in enumerate(chans)

			chans_spike_idxs = ind[ii]
    		Dw_spk[ch][1][2] = CartesianIndex.( find_redundant( x[ii], chans_spike_idxs, lockout ) )

		end

    end # method

    ### store clean indices in struct
    # n.spike_idx = ind_clean#[ind_clean[ii][d[ii]] for ii in 1:length(ind_clean)]
    n.spike_idx = Dw_spk; # channel X w X Dw/index of peak (Dw only if method is "petrantonakis")
end

"""
Documentation ...
"""
function extract_isih(n::Neurons.recording)
### Extracts the inter-spike intervals
    dt = 1/n.sf_preprocessed
    isih = []
    for ch in  eachindex(n.spike_idx)
        push!(isih, [])
        for w in eachindex(n.spike_idx[1])
            push!(isih[ch], [])
            if length(n.spike_idx[ch][w][2]) > 1
                isih[ch][w] = cartesian2matrix(diff(n.spike_idx[ch][w][2])) .* dt
            end
        end
    end
    n.isih = isih # value_source X w X time_source X time/values
end

"""
Documentation ...
"""
function extract_spikes(n::Neurons.recording, mp, w::Int64=1)
### Extract spike spikeforms and time stamps in ms.
#   extracts the spikeforms at indices n.spike_idx (normalized around zero)
#   from the filtered signal x using a fixed window around the
#   window +-mp of the spikes.
#	ONLY FOR ONE VALUE OF w
#   structure: n.spikeforms[val_source][time_source][spike_time/value]
#

    dt = 1/n.sf_preprocessed
	window = collect(mp[1]:1:mp[2]);
	spkfrms = []

	for ii in 1:n.n_channels # loop over channels as value source

		push!(spkfrms, [])
	    # for w in 1:size(n.spike_idx[1],1)
	    #     push!(spkfrms[ii], [])

			for (jj,spk_idxs) in enumerate(n.spike_idx) # loop over channels for time source

				push!(spkfrms[ii], [[],[]]) ### t, s

				for spk in spk_idxs[w][2] # loop over individual spikes

				    # skip spikes that have negative indices
	                spike_window_idxs = spk.I .+ window # generate window for each spike
	                spike_window_times = spike_window_idxs .* dt .* 1000 # spike times in ms

					if spike_window_idxs[1] > 0 && spike_window_idxs[end] <= length(n.trace_preprocessed[ii])# skip spikes that are not complete in window

						push!( spkfrms[ii][jj][1], spike_window_times ) # store spike times
	                    push!( spkfrms[ii][jj][2], n.trace_preprocessed[ii][spike_window_idxs] .- mean( n.trace_preprocessed[ii][spike_window_idxs]) ) #store spike values

					end
	            end
	        end
	    # end
	end
	n.spikeforms = spkfrms # value_source X w X time_source X time/values

end

"""
Documentation ...
"""
function extract_features(n::Neurons.recording)
    n.features = []
    for ch in 1:n.n_channels
        push!(n.features, [])
        if !isempty(n.spikeforms[ch][ch][2])
            ### make matrix for PCA
            m = zeros(length(n.spikeforms[ch][ch][2]), length(n.spikeforms[ch][ch][2][1])); # #spikes X #samples per spike
            [m[ii,:] = n.spikeforms[ch][ch][2][ii] for ii = 1:length(n.spikeforms[ch][ch][2])]; # spike X sample
            ### fit PCA
            M = fit(PCA, m)
            n.features[ch] = projection(M) # projection(M): spike X PC
        end
    end
end

"""
Documentation ...
"""
function cluster_spikes(n::Neurons.recording, algorithm::AbstractArray, init_method::AbstractArray, feats::AbstractArray, n_clusters::AbstractArray, nIter_EM::AbstractArray, nIter_kmeans_init::AbstractArray, cov_kind::AbstractArray)

############
# cluster given features using _algorithm_
# BIC optimization for cluster number
#
########
# algorithm: kmeans
#
########
# algorithm: gmm
# init_method: :split or :kmeans

    n.spike_class = []
    ###
    if algorithm[1] == "kmeans"
        for ch in 1:n.n_channels
            #push!(n.spike_class, [])
            if size(n.features[ch], 2) > feats[end] # if there are enough features
                result = kmeans(n.features[ch][:,feats]', n_clusters[1]);
                clusters = []
                for cl = 1:n_clusters[1]
                    push!(clusters, findall(result.assignments .== cl))
                end
                push!(n.spike_class, clusters)
            else ### push empty array for that channel
                push!(n.spike_class, [])
            end
        end
    ###
    elseif algorithm[1] == "gmm"
        n.BIC = []
        ### for each channel fit gmm with n components specified in
        ### n_clusters, compute BIC fear each value, find minimum BIC and use
        ### for clustering
        for ch in 1:n.n_channels
            if size(n.features[ch],2) > length(feats) ### if enough features
                N = size(n.features[ch][:,feats], 1)
                BICS = Array{Float64,2}(undef, 0, 2)
                for cmpnts in n_clusters
                    if size(n.features[ch], 1) > cmpnts
                        M = GMM(cmpnts, n.features[ch][:,feats];
                            method=init_method[1], kind=cov_kind[1], nInit=nIter_kmeans_init[1], nIter=nIter_EM[1], nFinal=nIter_EM[1])
                        LL = sum(llpg(M, n.features[ch][:,feats]))
                        par = cmpnts + cmpnts*length(feats) + cmpnts* length(feats)*(length(feats)+1)/2 # number of parameters
                        bic = -2*LL + log.(N)*par
                        BICS = vcat(BICS, [cmpnts bic])
                    end
                end
                push!(n.BIC, BICS)
                ### fit model with optimal number of components and cluster spikes
                opt_cmpnts = Int(BICS[findmin(BICS[:,2])[2]])
                push!(n.spike_class, [])
                M = GMM(opt_cmpnts, n.features[ch][:,feats];
                    method=init_method[1], kind=cov_kind[1], nInit=nIter_kmeans_init[1], nIter=nIter_EM[1], nFinal=nIter_EM[1])
                posteriors = gmmposterior( M, n.features[ch][:,feats] )
                classes = [findmax(posteriors[1], dims=2)[2][spk][2] for spk in 1:length(findmax(posteriors[1], dims=2)[2])] ### extract cluster index for each spike from CartesianCoordinates
                n.spike_class[ch] = [[] for ii in 1:opt_cmpnts]
                for (spk_idx,cl) in enumerate(classes)
                    push!(n.spike_class[ch][cl], spk_idx)
                end
            else
                push!(n.spike_class, [])
                push!(n.BIC, [])
            end
        end
        # for ch in 1:n.n_channels
        #     print(ch)
        #     push!(n.spike_class, [])
        #     if size(n.features[ch],2) > length(feats) ### if enough features
        #         M = GMM(n_clusters[1], n.features[ch][:,feats];
        #             method=:kmeans, kind=cov_kind[1], nInit=nIter_kmeans_init[1], nIter=nIter_EM[1], nFinal=nIter_EM[1])
        #         posteriors = gmmposterior( M, n.features[ch][:,feats] )
        #         classes = [findmax(posteriors[1], dims=2)[2][spk][2] for spk in 1:length(findmax(posteriors[1], dims=2)[2])] ### extract cluster index for each spike from CartesianCoordinates
        #         n.spike_class[ch] = [[] for ii in 1:n_clusters[1]]
        #         for (spk_idx,cl) in enumerate(classes)
        #             push!(n.spike_class[ch][cl], spk_idx)
        #         end
        #     else
        #         n.spike_class[ch] = [[] for ii in 1:n_clusters[1]]
        #     end
        # end
    end
end

"""
Up- or downsampling
Acts on either n.trace or n.trace_preprocessed
"""
function resample(n::Neurons.recording, frequency::Real)
    #idx = 1:factor:size(n.trace)[1]
    #n.trace_preprocessed = n.trace[idx];
    #n.sf_decimated = n.sf/factor;
    if isempty(n.preprocessing_steps)
        n.trace_preprocessed = DSP.resample( n.trace,frequency//convert(Int64, n.sf) );
    else
        n.trace_preprocessed = DSP.resample( n.trace_preprocessed,frequency//convert(Int64, n.sf_preprocessed) );
    end
    n.sf_preprocessed = frequency;
    push!(n.preprocessing_steps, string( "Decimated to: ",frequency ))
end

"""
Up- or downsampling
Acts on either n.trace or n.trace_preprocessed
"""
function get_PSD(n::Neurons.recording, which, psd_type)
    if which == "trace"
        return [getfield( Main, Symbol(psd_type) )( flatten(ch), fs=n.sf ) for ch in n.trace]
    elseif which == "trace_preprocessed"
        return [getfield( Main, Symbol(psd_type) )( flatten(ch), fs=n.sf_preprocessed ) for ch in n.trace_preprocessed]
    end
end

function find_redundant(ref::AbstractArray, ind::AbstractArray, lockout::Int)
    ### ref: 1D array containing amplitude like values
    ### ind: 1D array containing indices of identified events, e.g. spikes
    ### lockout: Integer defining minimum distance of evvents in samples that are treated as separate events

    ### make sure it's cartesian!
    ind = CartesianIndex.(ind)
    ### matrix with all pairs of spike indices
    # mm = Iterators.product(flatten(cartesian2matrix(ind)), flatten(cartesian2matrix(ind)))
	### differences between all pairs of spike indices
	# dd = [ float(abs(diff([x for x in m])[1][1])) for m in mm ]
	dd = pairwise(Euclidean(), flatten(cartesian2matrix(ind))', flatten(cartesian2matrix(ind))', dims=2)
	dd[diagind(dd)] .= NaN # set diagonal to NaN
	sup_diag = dd[diagind(dd, 1)]
	gg = findall(x -> x > lockout, sup_diag )

	if all(sup_diag .> lockout) ### no double triggered spikes

		groups = [[spk_idx] for spk_idx in ind]

	elseif all(sup_diag .< lockout) ### only one double triggered spike

		groups = [ind]

	else ### single and double triggered spikes

		groups = []
	    current_group = []

		for (ij,spk_id) in enumerate(ind)

			push!(current_group, spk_id)

			if in(ij, gg) || ij == length(ind) ### ij == length(ind) required for last element

				push!(groups, current_group)
	            current_group = []

			end
		end
	end
	### find maxima of groups and store indices
    ind_clean = []
    for gr in groups

		max_id = findmax( abs.(ref[CartesianIndex.( gr )]) )
		# max_id = findmax(ref[gr])
    	push!(ind_clean, gr[max_id[2]])

	end

    #################
    ### DELETE ME ###a
    #################
    # ind_clean = []
    # ind_ = deepcopy(ind)
    # if length(ind_) > 1
    #     groups = []
    #     current_group = []
    #     push!(current_group, ind_[1])
    #     deleteat!(ind_, 1)
    #
    #     ### extract groups of detected event that are too close to each other to represent seperate spikes
    #     while !isempty(ind_)
    #         ### while next data point is to close to preceeding one, add it to current group
    #         while ind_[1][1]-current_group[end][1] < lockout ### double index because CartesianIndex
    #             push!(current_group, ind_[1])
    #             deleteat!(ind_, 1)
    #             if isempty(ind_) ### if already at the last spike index
    #                 break;
    #             end
    #         end
    #
    #         ### add current_group to group
    #         push!(groups, current_group)
    #         current_group = []
    #         if !isempty(ind_) ### deals with if last index was already deleted in while loop
    #             push!(current_group, ind_[1])
    #             deleteat!(ind_, 1)
    #         end
    #     end
    #
    #     ### for each group extract the maximum as representation
    #     ind_c = []
    #     for gr in groups
    #         if length(gr) > 1 ### if the group is a "group"
    #             gg = [gr[ii] for ii in eachindex(gr)]
    #             push!( ind_c, gg[findmax(ref[gg])[2]] )
    #         else ### if the group has only one entry
    #             push!(ind_c, gr[1])
    #         end
    #     end
    #     push!(ind_clean, ind_c)
    # else ### if no spikes in channel, add empty list to kepp size consistent
    #     push!(ind_clean, [])
    # end

    return ind_clean
end

################
### PLOTTING ###
################

function plot_detection(
        recording::Neurons.recording;
        channel::Int64,
        method::String="petrantonakis",
        polarity::String="negative")

	colors = myColors() # from MyHelpers.jl
	pcols = map(col -> (red(col), green(col), blue(col)), colors)

    if method == "petrantonakis"

        fig = figure(figsize=(13,5))
        rc("xtick",labelsize=8)
        rc("ytick",labelsize=8)

		### time course of Dw signal
        Dw_w1 = recording.spike_idx[channel][1][1]
        Dw_w2 = recording.spike_idx[channel][2][1]
		### indices of spikes
		spk_idx_w1 = recording.spike_idx[channel][1][2]
        spk_idx_w2 = recording.spike_idx[channel][2][2]
		### spike indices from cartesian to matrix
		if !isempty(spk_idx_w1) && !isempty(spk_idx_w2)
			spk_idx_w1 = cartesian2matrix(spk_idx_w1)
			spk_idx_w2 = cartesian2matrix(spk_idx_w2)
		end
        w1_threshold = recording.detection_threshold[1]
        w2_threshold = recording.detection_threshold[2]

        ### scatter plot of w1 X w2
        if (length(spk_idx_w1)-length(spk_idx_w1) == 0) && length(recording.spike_idx[channel][1][2]) != 0

            ax1 = subplot(1,3,1)
            w1 = recording.spike_idx[channel][1][1][spk_idx_w1]
            w2 = recording.spike_idx[channel][2][1][spk_idx_w2]
            pols = cartesian2polar(w1, w2)
            r = pols[:,1]
            phi = pols[:,2]
            ax1.scatter( r, phi, color=pcols[2])

        end

        ax2 = subplot(1,3,2)
        ax2.plot( Dw_w1, color=pcols[1], zorder=1 )
        ax2.plot( 1:length(Dw_w1), ones(length(Dw_w1))*w1_threshold, color=:red, linestyle="--" )
        if length(recording.spike_idx[channel][1][2]) > 0

            spk_idx_w1 = cartesian2matrix(recording.spike_idx[channel][1][2])
            ax2.scatter(spk_idx_w1, Dw_w1[spk_idx_w1], color=pcols[2], zorder=2)

        end

        ax3 = subplot(1,3,3)
        ax3.plot(Dw_w2, color=pcols[1], zorder=1)
        ax3.plot( 1:length(Dw_w2), ones(length(Dw_w2))*w2_threshold, color=:red, linestyle="--" )
		if length(recording.spike_idx[channel][2][2]) > 0

            spk_idx_w2 = cartesian2matrix(recording.spike_idx[channel][2][2])
            ax3.scatter(spk_idx_w2, Dw_w2[spk_idx_w2], color=pcols[2], zorder=2)

        end

    elseif method == "threshold"

        fig = figure(figsize=(13,5))
        rc("xtick",labelsize=8)
        rc("ytick",labelsize=8)

        x = recording.trace_preprocessed[channel]
        spk_idx = cartesian2matrix(recording.spike_idx[channel][1][2])

		if polarity == "negative"

			thr = -recording.detection_threshold[channel]

		elseif polarity == "positive"

			thr = recording.detection_threshold[channel]

		elseif polarity == "both"

			thr = [-recording.detection_threshold[channel], recording.detection_threshold[channel]]
        end

        plot( x )
        scatter( spk_idx, x[spk_idx], color=pcols[2], zorder=2 )

        if length(thr) != 2

			plot( 1:length(x), ones(length(x))*thr, color=:red, linestyle="--" )

		else

			plot( 1:length(x), ones(length(x))*thr[1], color=:red, linestyle="--" )
            plot( 1:length(x), ones(length(x))*thr[2], color=:red, linestyle="--" )
        end
    end
end

function plot_waveforms(
	recording::Neurons.recording,
	t::Array{Float64,1},
	ylims::Array{Int64,2},
	w::Int64=1;
	#
	t_source )
	##########################
	##########################
	##########################

	colors = myColors() # from MyHelpers.jl
	pcols = map(col -> (red(col), green(col), blue(col)), colors)
	half_window = abs(minimum(t))/1000

	fig = figure(figsize=(15,7), dpi=300)
	rc("xtick",labelsize=4)
	rc("ytick",labelsize=4)

	# ylims = [-1000 1000]

	if t_source == "val_source"
		corresponding = true
	else
		corresponding = false
	end

	for val_source = 1:recording.n_channels
	    #### t_cource determines whose channels spike times are used to cut out the the waveforms: if t_cource = val_source then channel and spiketimes are matched, i.e. 1<->1, 2<->2 etc
		if corresponding
			t_source = val_source
		end
	#     t_source = val_source

	    ### sort spikes according to amplitude
	    sorted_idx = sortperm([maximum(spk) for spk in recording.spikeforms[val_source][t_source][2]] .- [minimum(spk) for spk in recording.spikeforms[val_source][t_source][2]], rev=true )
	    ### how many spikes to plot
	    # n = length(sorted_idx)
	    n = 10

	    if !isempty(sorted_idx) ### if there are spikes in sorted_idx, i.e. if channel has spikes

	        spks_to_plot = []

	        ### loop over clusters of current channel
	        if !isempty(recording.spike_class)
	            for cl in recording.spike_class[t_source] ### plot the clusters based on the t_source channel
	                push!(spks_to_plot, intersect(cl, sorted_idx))
	            end
	        else
	            push!(spks_to_plot, sorted_idx) ### i.e. one "cluster" if not yet clustered
	        end

	        ### check if there are enough spikes if number of spikes to plot is specified
	        if length(sorted_idx) < n
	            n = length(sorted_idx)
	        end

	        ### plot spikes belonging to clusters
			offset = 0
	        for (ii,spk_cl) in enumerate(spks_to_plot)
	            if !isempty(spk_cl) ### cluster is not empty
	                subplot(2,5,val_source)
				   ### plot individual traces
				   plot(t, reduce(hcat, recording.spikeforms[val_source][t_source][2][spk_cl]),
					  color=pcols[ii], alpha=0.3, linewidth=0.1, zorder=1) ### plot individual spikes
				   ### plot mean waveform +- std
				   clusters_mean = mean(reduce(hcat, recording.spikeforms[val_source][t_source][2][spk_cl]), dims=2) .+ offset
				   clusters_std = std(reduce(hcat, recording.spikeforms[val_source][t_source][2][spk_cl]), dims=2)
				   plot(t, clusters_mean, color=pcols[ii], alpha=0.5, linewidth=0.5, zorder=2) ### plot mean spike waveform of cluster
				   fill_between(t, flatten(clusters_mean+clusters_std), flatten(clusters_mean-clusters_std), color=pcols[ii], alpha=0.2, linewidth= 0.1)
				   ###
				   title(string("Channel ", val_source, " (", n, ")"), fontsize=5)
				   ylim(ylims[1], ylims[2])
	# 			   legend()
	            end
				offset = offset +0
	        end

	        ### add x- and ylabel
	        if val_source in collect(6:10)
	            subplot(2,5,val_source)
	            xlabel("ms", fontsize=5)
	        elseif val_source in [1,6]
	            subplot(2,5,val_source)
	#             ylabel("amplitude")
	        end
			if val_source in collect(1:5)
			   xticks(-half_window*1000:1:half_window*1000, []) ### half_window comes from waveform extraction step
			end
			if val_source in [2:5 7:10]
			   yticks(ylims[1]:100:ylims[2], [])
			end

	    else ### plot empty axes of channel has no spikes
	        subplot(2,5,val_source)
	        plot(t, t,
	            linewidth=0)
	        title(string("Channel ", val_source, " (", n, ")"))
	        ### add x- and ylabel
	        ### add x- and ylabel
	        if val_source in collect(6:10)
	            subplot(2,5,val_source)
	            xlabel("ms", fontsize=5)
	        elseif val_source in [1,6]
	            subplot(2,5,val_source)
	#             ylabel("amplitude")
	        end
			if val_source in collect(1:5)
			   xticks(-4:2:4, [])
			end
			if val_source in [2:5 7:10]
			   yticks(ylims[1]:100:ylims[2], [])
			end

	    end

	end

end

###
end


@test @inferred(Octavian.cld_fast(Octavian.StaticInt(6), Octavian.StaticInt(2))) == @inferred(Octavian.cld_fast(11, Octavian.StaticInt(4))) == @inferred(Octavian.cld_fast(Octavian.StaticInt(7), 3)) == @inferred(Octavian.cld_fast(9, 4)) == 3
@test @inferred(Octavian.cld_fast(Octavian.StaticInt(7), Octavian.StaticInt(2))) == @inferred(Octavian.cld_fast(13, Octavian.StaticInt(4))) == @inferred(Octavian.cld_fast(Octavian.StaticInt(8), 2)) == @inferred(Octavian.cld_fast(15, 4)) == 4


"""
    restoreoptsum!(m::LinearMixedModel, io::IO)
    restoreoptsum!(m::LinearMixedModel, fnm::AbstractString)

Read, check, and restore the `optsum` field from a JSON stream or filename.
"""
function restoreoptsum!(m::LinearMixedModel{T}, io::IO) where {T}
    dict = JSON3.read(io)
    ops = m.optsum
    okay =
        (setdiff(propertynames(ops), keys(dict)) == [:lowerbd]) &&
        all(ops.lowerbd .≤ dict.initial) &&
        all(ops.lowerbd .≤ dict.final)
    if !okay
        throw(ArgumentError("initial or final parameters in io do not satify lowerbd"))
    end
    for fld in (:feval, :finitial, :fmin, :ftol_rel, :ftol_abs, :maxfeval, :nAGQ, :REML)
        setproperty!(ops, fld, getproperty(dict, fld))
    end
    ops.initial_step = copy(dict.initial_step)
    ops.xtol_rel = copy(dict.xtol_rel)
    copyto!(ops.initial, dict.initial)
    copyto!(ops.final, dict.final)
    for (v, f) in (:initial => :finitial, :final => :fmin)
        if !isapprox(objective(updateL!(setθ!(m, getfield(ops, v)))), getfield(ops, f))
            throw(ArgumentError("model m at $v does not give stored $f"))
        end
    end
    ops.optimizer = Symbol(dict.optimizer)
    ops.returnvalue = Symbol(dict.returnvalue)
    # provides compatibility with fits saved before the introduction of fixed sigma
    ops.sigma = get(dict, :sigma, nothing)
    fitlog = get(dict, :fitlog, nothing)
    ops.fitlog = if isnothing(fitlog)
        # compat with fits saved before fitlog
        [(ops.initial, ops.finitial, ops.final, ops.fmin)]
    else
        [(convert(Vector{T}, first(entry)), T(last(entry))) for entry in fitlog]
    end
    return m
end

function restoreoptsum!(m::LinearMixedModel, fnm::AbstractString)
    open(fnm, "r") do io
        restoreoptsum!(m, io)
    end
end

"""
    saveoptsum(io::IO, m::LinearMixedModel)
    saveoptsum(fnm::AbstractString, m::LinearMixedModel)

Save `m.optsum` (w/o the `lowerbd` field) in JSON format to an IO stream or a file

The reason for omitting the `lowerbd` field is because it often contains `-Inf`
values that are not allowed in JSON.
"""
saveoptsum(io::IO, m::LinearMixedModel) = JSON3.write(io, m.optsum)
function saveoptsum(fnm::AbstractString, m::LinearMixedModel)
    open(fnm, "w") do io
        saveoptsum(io, m)
    end
end

# TODO: write methods for GLMM
# TODO, maybe: something nice for the MixedModelBootstrap


abstract Any_Area_is_an_Area

type AnArea_Regular <: Any_Area_is_an_Area
  data_based_on :: ASCIIString
  lon  :: Array{Float64,1}
  lat  :: Array{Float64,1}
  plon :: Array{Float64,2}
  plat :: Array{Float64,2}
  plon_selected_on_by :: Array{Bool,2}
  data_selected_on_plon_by :: Array{Bool,2}
end

# Select time series
function Selected_within{N}(p_area :: AnArea_Regular,
                            sst :: Array{Float64,N}; map_to :: Symbol = :default)
  plon_s = p_area.plon_selected_on_by
  data_s = p_area.data_selected_on_plon_by
  sst1 = reshape(sst ,length(plon_s), reduce(*,size(sst,[3:ndims(sst);]...)))

  if (map_to == :default )
    sst1 = sst1[find(plon_s[:]),:];
    sst1[find(!data_s[:]),:] = NaN
    sst1 = reshape(sst1,size(data_s)...,size(sst,(3:ndims(sst))...)...)
  elseif(map_to == :whole )
    Filter1= find(plon_s[:])
    Filter2 = Filter1[find(data_s[:])]
    pFilter = fill(false,plon_s); pFilter[Filter2] = true;
    sst1[find(!pFilter[:]),:] = NaN
    sst1 = reshape(sst1, size(sst))
  else
    error("No such 'map_to' stratagy yet")
  end
  return sst1
end

function Selected_within(p_area :: AnArea_Regular,
                         sst :: Array{Float64,2}; map_to :: Symbol = :default)
  plon_selected_on_by = p_area.plon_selected_on_by
  data_selected_on_plon_by = p_area.data_selected_on_plon_by
  if (map_to == :default )
    sst1 = reshape(sst[plon_selected_on_by],size(data_selected_on_plon_by))
    sst1[!data_selected_on_plon_by] = NaN
  elseif(map_to == :whole )
    sst1 = copy(sst)
    dreal = find(plon_selected_on_by)
    dreal = dreal[find(data_selected_on_plon_by)]
    mask = fill(true,size(plon_selected_on_by) )
    mask[dreal] = false
    sst1[mask] = NaN
  else
    error("No such 'map_to' stratagy yet")
  end
  return sst1
end

export Any_Area_is_an_Area, AnArea_Regular
export Selected_within
include("deal_with_models/Regular_Area_Select.jl")
include("special_areas.jl")


# --------------- Base type for model:
module Models

using Flux, LinearAlgebra

include("base.jl")

export AbstractFittedModel, AbstractDifferentiableModel,
    FluxModel, LogisticModel, BayesianLogisticModel,
    RTorchModel, PyTorchModel,
    probs, logits

end


# Necessary to have initialize set u_modified to false if all don't do anything
# otherwise unnecessary save
INITIALIZE_DEFAULT(cb,u,t,integrator) = u_modified!(integrator, false)

struct ContinuousCallback{F1,F2,F3,F4,T,T2,I} <: AbstractContinuousCallback
  condition::F1
  affect!::F2
  affect_neg!::F3
  initialize::F4
  idxs::I
  rootfind::Bool
  interp_points::Int
  save_positions::BitArray{1}
  abstol::T
  reltol::T2
  ContinuousCallback(condition::F1,affect!::F2,affect_neg!::F3,
                     initialize::F4,idxs::I,rootfind,
                     interp_points,save_positions,abstol::T,reltol::T2) where {F1,F2,F3,F4,T,T2,I} =
                       new{F1,F2,F3,F4,T,T2,I}(condition,
                                               affect!,affect_neg!,
                                               initialize,idxs,rootfind,interp_points,
                                               BitArray(collect(save_positions)),
                                               abstol,reltol)
end

ContinuousCallback(condition,affect!,affect_neg!;
                   initialize = INITIALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=true,
                   save_positions=(true,true),
                   interp_points=10,
                   abstol=10eps(),reltol=0) = ContinuousCallback(
                              condition,affect!,affect_neg!,initialize,
                              idxs,
                              rootfind,interp_points,
                              save_positions,abstol,reltol)

function ContinuousCallback(condition,affect!;
                   initialize = INITIALIZE_DEFAULT,
                   idxs = nothing,
                   rootfind=true,
                   save_positions=(true,true),
                   affect_neg! = affect!,
                   interp_points=10,
                   abstol=10eps(),reltol=0)

 ContinuousCallback(
            condition,affect!,affect_neg!,initialize,idxs,
            rootfind,interp_points,
            collect(save_positions),abstol,reltol)

end

struct DiscreteCallback{F1,F2,F3} <: AbstractDiscreteCallback
  condition::F1
  affect!::F2
  initialize::F3
  save_positions::BitArray{1}
  DiscreteCallback(condition::F1,affect!::F2,
                   initialize::F3,save_positions) where {F1,F2,F3} = new{F1,F2,F3}(condition,
                                                                                affect!,initialize,
                                                                                BitArray(collect(save_positions)))
end
DiscreteCallback(condition,affect!;
        initialize = INITIALIZE_DEFAULT,save_positions=(true,true)) = DiscreteCallback(condition,affect!,initialize,save_positions)

# DiscreteCallback(condition,affect!,save_positions) = DiscreteCallback(condition,affect!,save_positions)

struct CallbackSet{T1<:Tuple,T2<:Tuple} <: DECallback
  continuous_callbacks::T1
  discrete_callbacks::T2
end

CallbackSet(callback::AbstractDiscreteCallback) = CallbackSet((),(callback,))
CallbackSet(callback::AbstractContinuousCallback) = CallbackSet((callback,),())
CallbackSet() = CallbackSet((),())
CallbackSet(cb::Nothing) = CallbackSet()

# For Varargs, use recursion to make it type-stable
CallbackSet(callbacks::Union{DECallback,Nothing}...) = CallbackSet(split_callbacks((), (), callbacks...)...)

@inline split_callbacks(cs, ds) = cs, ds
@inline split_callbacks(cs, ds, c::Nothing, args...) = split_callbacks(cs, ds, args...)
@inline split_callbacks(cs, ds, c::AbstractContinuousCallback, args...) = split_callbacks((cs..., c), ds, args...)
@inline split_callbacks(cs, ds, d::AbstractDiscreteCallback, args...) = split_callbacks(cs, (ds..., d), args...)
@inline function split_callbacks(cs, ds, d::CallbackSet, args...)
  split_callbacks((cs...,d.continuous_callbacks...), (ds..., d.discrete_callbacks...), args...)
end

# Recursively apply initialize! and return whether any modified u
function initialize!(cb::CallbackSet,u,t,integrator::DEIntegrator)
  initialize!(u,t,integrator,false,cb.continuous_callbacks...,cb.discrete_callbacks...)
end
initialize!(cb::CallbackSet{Tuple{},Tuple{}},u,t,integrator::DEIntegrator) = false
function initialize!(u,t,integrator::DEIntegrator,any_modified::Bool,
                     c::DECallback,cs::DECallback...)
  c.initialize(c,u,t,integrator)
  initialize!(u,t,integrator,any_modified || integrator.u_modified,cs...)
end
function initialize!(u,t,integrator::DEIntegrator,any_modified::Bool,
                     c::DECallback)
  c.initialize(c,u,t,integrator)
  any_modified || integrator.u_modified
end

# Helpers
Base.isempty(cb::CallbackSet) = isempty(cb.continuous_callbacks) && isempty(cb.discrete_callbacks)
Base.isempty(cb::AbstractContinuousCallback) = false
Base.isempty(cb::AbstractDiscreteCallback) = false

#======================================================#
# Callback handling
#======================================================#

function get_tmp(integrator::DEIntegrator, callback)
  _tmp = get_tmp_cache(integrator)
  _tmp === nothing && return nothing
  _cache = first(_tmp)
  if callback.idxs isa Nothing
    tmp = _cache
  elseif !(callback.idxs isa Number)
    tmp = @view _cache[callback.idxs]
  else
    tmp = nothing
  end
  return tmp
end

function get_condition(integrator::DEIntegrator, callback, abst)
  tmp = get_tmp(integrator, callback)
  ismutable = !(tmp === nothing)
  ismutable && !(typeof(callback.idxs) isa Number) ? integrator(tmp,abst,Val{0},idxs=callback.idxs) :
                                                     tmp = integrator(abst,Val{0},idxs=callback.idxs)
  return callback.condition(tmp,abst,integrator)
end

# Use Recursion to find the first callback for type-stability

# Base Case: Only one callback
function find_first_continuous_callback(integrator, callback::AbstractContinuousCallback)
  (find_callback_time(integrator,callback,1)...,1,1)
end

# Starting Case: Compute on the first callback
function find_first_continuous_callback(integrator, callback::AbstractContinuousCallback, args...)
  find_first_continuous_callback(integrator,find_callback_time(integrator,callback,1)...,1,1,args...)
end

function find_first_continuous_callback(integrator,tmin::Number,upcrossing::Number,
                                        event_occured::Bool,idx::Int,counter::Int,
                                        callback2)
  counter += 1 # counter is idx for callback2.
  tmin2,upcrossing2,event_occurred2 = find_callback_time(integrator,callback2,counter)

  if event_occurred2 && (tmin2 < tmin || !event_occured)
    return tmin2,upcrossing2,true,counter,counter
  else
    return tmin,upcrossing,event_occured,idx,counter
  end
end

function find_first_continuous_callback(integrator,tmin::Number,upcrossing::Number,event_occured::Bool,idx::Int,counter::Int,callback2,args...)
  find_first_continuous_callback(integrator,find_first_continuous_callback(integrator,tmin,upcrossing,event_occured,idx,counter,callback2)...,args...)
end

@inline function determine_event_occurance(integrator,callback,counter)
  event_occurred = false
  if callback.interp_points!=0
    addsteps!(integrator)
  end
  Θs = range(typeof(integrator.t)(0), stop=typeof(integrator.t)(1), length=callback.interp_points)
  interp_index = 0
  # Check if the event occured
  if typeof(callback.idxs) <: Nothing
    previous_condition = callback.condition(integrator.uprev,integrator.tprev,integrator)
  else
    @views previous_condition = callback.condition(integrator.uprev[callback.idxs],integrator.tprev,integrator)
  end

  if integrator.event_last_time == counter && ODE_DEFAULT_NORM(previous_condition,integrator.t) < 100ODE_DEFAULT_NORM(integrator.last_event_error,integrator.t)

    # If there was a previous event, utilize the derivative at the start to
    # chose the previous sign. If the derivative is positive at tprev, then
    # we treat the value as positive, and derivative is negative then we
    # treat the value as negative, reguardless of the postiivity/negativity
    # of the true value due to it being =0 sans floating point issues.

    # Only due this if the discontinuity did not move it far away from an event
    # Since near even we use direction instead of location to reset

    if callback.interp_points==0
      addsteps!(integrator)
    end

    abst = integrator.tprev+sign(integrator.dt)*100*eps(integrator.tprev) # Evaluate condition slightly in future
    tmp_condition = get_condition(integrator, callback, abst)
    prev_sign = tmp_condition > previous_condition ? 1.0 : -1.0
  else
    prev_sign = sign(previous_condition)
  end

  prev_sign_index = 1
  if typeof(callback.idxs) <: Nothing
    next_sign = sign(callback.condition(integrator.u,integrator.t,integrator))
  else
    @views next_sign = sign(callback.condition(integrator.u[callback.idxs],integrator.t,integrator))
  end

  if ((prev_sign<0 && !(typeof(callback.affect!)<:Nothing)) || (prev_sign>0 && !(typeof(callback.affect_neg!)<:Nothing))) && prev_sign*next_sign<=0
    event_occurred = true
    interp_index = callback.interp_points
  elseif callback.interp_points!=0 && !isdiscrete(integrator.alg) # Use the interpolants for safety checking
    for i in 2:length(Θs)
      abst = integrator.tprev+integrator.dt*Θs[i]
      new_sign = get_condition(integrator, callback, abst)
      if ((prev_sign<0 && !(typeof(callback.affect!)<:Nothing)) || (prev_sign>0 && !(typeof(callback.affect_neg!)<:Nothing))) && prev_sign*new_sign<0
        event_occurred = true
        interp_index = i
        break
      else
        prev_sign_index = i
      end
    end
  end

  event_occurred,interp_index,Θs,prev_sign,prev_sign_index
end

function find_callback_time(integrator,callback,counter)
  event_occurred,interp_index,Θs,prev_sign,prev_sign_index = determine_event_occurance(integrator,callback,counter)
  if event_occurred
    if typeof(callback.condition) <: Nothing
      new_t = zero(typeof(integrator.t))
    else
      if callback.interp_points!=0
        top_Θ = Θs[interp_index] # Top at the smallest
        bottom_θ = Θs[prev_sign_index]
      else
        top_Θ = typeof(integrator.t)(1)
        bottom_θ = typeof(integrator.t)(0)
      end
      if callback.rootfind && !isdiscrete(integrator.alg)
        zero_func = (Θ) -> begin
          abst = integrator.tprev+integrator.dt*Θ
          return get_condition(integrator, callback, abst)
        end
        if zero_func(top_Θ) == 0
          Θ = top_Θ
        else
          if integrator.event_last_time == counter &&
            abs(zero_func(bottom_θ)) < 100abs(integrator.last_event_error) &&
            prev_sign_index == 1

            # Determined that there is an event by derivative
            # But floating point error may make the end point negative

            sign_top = sign(zero_func(top_Θ))
            bottom_θ += 2eps(typeof(bottom_θ))
            iter = 1
            while sign(zero_func(bottom_θ)) == sign_top && iter < 12
              bottom_θ *= 5
              iter += 1
            end
            iter == 12 && error("Double callback crossing floating pointer reducer errored. Report this issue.")
          end
          Θ = prevfloat(find_zero(zero_func, (bottom_θ,top_Θ), Roots.AlefeldPotraShi(), atol = callback.abstol/100))
          integrator.last_event_error = ODE_DEFAULT_NORM(zero_func(Θ),integrator.t+integrator.dt*Θ)
        end
        #Θ = prevfloat(...)
        # prevfloat guerentees that the new time is either 1 floating point
        # numbers just before the event or directly at zero, but not after.
        # If there's a barrier which is never supposed to be crossed,
        # then this will ensure that
        # The item never leaves the domain. Otherwise Roots.jl can return
        # a float which is slightly after, making it out of the domain, causing
        # havoc.
        new_t = integrator.dt*Θ
      elseif interp_index != callback.interp_points && !isdiscrete(integrator.alg)
        new_t = integrator.dt*Θs[interp_index]
      else
        # If no solve and no interpolants, just use endpoint
        new_t = integrator.dt
      end
    end
  else
    new_t = zero(typeof(integrator.t))
  end

  new_t,prev_sign,event_occurred
end

function apply_callback!(integrator,callback::ContinuousCallback,cb_time,prev_sign)
  if cb_time == zero(typeof(integrator.t))
    error("Event repeated at the same time. Please report this error")
  end
  change_t_via_interpolation!(integrator,integrator.tprev+cb_time)

  # handle saveat
  _, savedexactly = savevalues!(integrator)
  saved_in_cb = true

  @inbounds if callback.save_positions[1]
    # if already saved then skip saving
    savedexactly || savevalues!(integrator,true)
  end

  integrator.u_modified = true

  if prev_sign < 0
    if typeof(callback.affect!) <: Nothing
      integrator.u_modified = false
    else
      callback.affect!(integrator)
    end
  elseif prev_sign > 0
    if typeof(callback.affect_neg!) <: Nothing
      integrator.u_modified = false
    else
      callback.affect_neg!(integrator)
    end
  end

  if integrator.u_modified
    reeval_internals_due_to_modification!(integrator)
    @inbounds if callback.save_positions[2]
      savevalues!(integrator,true)
      saved_in_cb = true
    end
    return true,saved_in_cb
  end
  false,saved_in_cb
end

#Base Case: Just one
@inline function apply_discrete_callback!(integrator,callback::DiscreteCallback)
  saved_in_cb = false
  if callback.condition(integrator.u,integrator.t,integrator)
    # handle saveat
    _, savedexactly = savevalues!(integrator)
    saved_in_cb = true
    @inbounds if callback.save_positions[1]
      # if already saved then skip saving
      savedexactly || savevalues!(integrator,true)
    end
    integrator.u_modified = true
    callback.affect!(integrator)
    @inbounds if callback.save_positions[2]
      savevalues!(integrator,true)
      saved_in_cb = true
    end
  end
  integrator.u_modified,saved_in_cb
end

#Starting: Get bool from first and do next
@inline function apply_discrete_callback!(integrator,callback::DiscreteCallback,args...)
  apply_discrete_callback!(integrator,apply_discrete_callback!(integrator,callback)...,args...)
end

@inline function apply_discrete_callback!(integrator,discrete_modified::Bool,saved_in_cb::Bool,callback::DiscreteCallback,args...)
  bool,saved_in_cb2 = apply_discrete_callback!(integrator,apply_discrete_callback!(integrator,callback)...,args...)
  discrete_modified || bool, saved_in_cb || saved_in_cb2
end

@inline function apply_discrete_callback!(integrator,discrete_modified::Bool,saved_in_cb::Bool,callback::DiscreteCallback)
  bool,saved_in_cb2 = apply_discrete_callback!(integrator,callback)
  discrete_modified || bool, saved_in_cb || saved_in_cb2
end


"""
    submit_limit_order!(
        ob::OrderBook{Sz,Px,Oid,Aid},
        orderid::Oid,
        side::OrderSide,
        limit_price::Real,
        limit_size::Real,
        [, acct_id::Aid, fill_mode::OrderTraits ]
    )

Enter limit order with size `limit_size`, price `limit_price` with `side::OrderSide` into `ob::OrderBook`.

If an account if `acct_id` is provided, account holdings are tracked in `ob.acct_map`.

Order execution logic can be modified according to the argument `fill_mode::`[`OrderTraits`](@ref) which
defaults to `fill_mode=VANILLA_FILLTYPE`, representing the default order matching mode.

`submit_limit_order!` returns tuple of
 - `new_open_order::Order` representing the order left in the book after matching. Is `nothing` if no order was inserted
 - `order_match_lst::Vector{Order}` representing the matched orders if the order crosses the book.
 - `left_to_trade::Sz` representing the size of the portion of the order which could neither inserted nor matched.

"""
function submit_limit_order!(
    ob::OrderBook{Sz,Px,Oid,Aid},
    orderid::Oid,
    side::OrderSide,
    limit_price::Real,
    limit_size::Real,
    acct_id::Union{Nothing,Aid}=nothing,
    fill_mode::OrderTraits=VANILLA_FILLTYPE,
) where {Sz,Px,Oid,Aid}
    # Part 0 - Check Arguments
    if !((limit_price > zero(limit_price)) && (limit_size > zero(limit_size)))
        error("Both limit_price and limit_size must be positive")
    end
    ## Part 1 - Cross the Order if appropriate
    best_bid, best_ask = best_bid_ask(ob) # check that price in right range
    if allows_cross(fill_mode) &&
       isbuy(side) &&
       !isnothing(best_ask) &&
       (limit_price >= best_ask) # order is bid (buy) and can cross
        cross_match_lst, remaining_size = _walk_order_book_bysize!(
            ob.ask_orders, Sz(limit_size), Px(limit_price), fill_mode
        )
    elseif allows_cross(fill_mode) &&
           !isbuy(side) &&
           !isnothing(best_ask) &&
           (limit_price >= best_ask) # order is ask (sell) and can cross
        cross_match_lst, remaining_size = _walk_order_book_bysize!(
            ob.bid_orders, Sz(limit_size), Px(limit_price), fill_mode
        )
    else # order can or does not cross, return empty matches
        cross_match_lst, remaining_size = Vector{Order{Sz,Px,Oid,Aid}}(), Sz(limit_size)
    end

    ## Part 2 - Rest the remaining order in the book if possible
    if allows_book_insert(fill_mode) && !iszero(remaining_size) # if there are remaining shares, try to add remaining to the LOB
        best_bid, best_ask = best_bid_ask(ob) # new best bid and ask
        if isbuy(side) && (isnothing(best_ask) || (limit_price < best_ask)) # if order is a buy and limit price is valid for resting order
            # create and insert order object into BID book
            new_open_order = Order{Sz,Px,Oid,Aid}(
                side, remaining_size, limit_price, orderid, acct_id
            )
            insert_order!(ob.bid_orders, new_open_order)
            # if account_id present, add order account map
            isnothing(acct_id) || _add_order_acct_map!(ob.acct_map, acct_id, new_open_order)
            # set remaining size to zero
            remaining_size = zero(Sz)
        elseif !isbuy(side) && (isnothing(best_bid) || (limit_price > best_bid)) # if order is a sell and limit price is valid for resting order
            # create and insert order object into ASK book
            new_open_order = Order{Sz,Px,Oid,Aid}(
                side, remaining_size, limit_price, orderid, acct_id
            )
            insert_order!(ob.ask_orders, new_open_order)
            # if account_id present, add order account map
            !isnothing(acct_id) &&
                _add_order_acct_map!(ob.acct_map, acct_id, new_open_order)
            # set remaining size to zero
            remaining_size = zero(Sz)
        else # if not appropriate to insert new order, register new order as nothing
            new_open_order = nothing
        end
    else
        new_open_order = nothing
    end

    ## Part 3 - Return information
    return (
        new_open_order, cross_match_lst, remaining_size
    )::Tuple{Union{Order{Sz,Px,Oid,Aid},Nothing},Vector{Order{Sz,Px,Oid,Aid}},Sz}
end

@inline _is_best_price_inside_limit(::OneSidedBook, ::Nothing) = true

@inline function _is_best_price_inside_limit(
    sb::OneSidedBook{Sz,Px,Oid,Aid}, limit_price::Px
) where {Sz,Px,Oid,Aid}
    if isbidbook(sb)
        return sb.best_price >= limit_price
    else
        return sb.best_price <= limit_price
    end
end

"""
    _walk_order_book_bysize!(
        sb::OneSidedBook{Sz,Px,Oid,Aid},
        order_size::Sz,
        limit_price::Union{Px,Nothing},
        order_mode::OrderTraits,
    )

Cross (limit or market) order with opposite single side of book.
Order size is specified in number of shares.
Return matches and outstanding size.

__Notes__
 - If `limit_price::Nothing`, order is treated as Market Order
 - Function expects `order_size>0` and `limit_price>0`

"""
function _walk_order_book_bysize!(
    sb::OneSidedBook{Sz,Px,Oid,Aid},
    order_size::Sz,
    limit_price::Union{Px,Nothing},
    order_mode::OrderTraits,
)::Tuple{Vector{Order{Sz,Px,Oid,Aid}},Sz} where {Sz,Px,Oid,Aid}
    # Allocate memory for order output
    order_match_lst = Vector{Order{Sz,Px,Oid,Aid}}()
    shares_left = order_size # remaining quantity to trade
    # Perform initial available liquidity check
    if isallornone(order_mode) && (_size_available(sb, limit_price) < order_size)
        return order_match_lst, shares_left
    end
    # Perform matching logic
    limit_price_check = Base.Fix2(_is_best_price_inside_limit, limit_price)
    while !isempty(sb.book) && !iszero(shares_left) && limit_price_check(sb) # while book not empty, order not done and best price within limit price
        price_queue::OrderQueue = _popfirst_queue!(sb)
        if price_queue.total_volume[] <= shares_left # If entire queue is to be wiped out
            append!(order_match_lst, price_queue.queue) # Add all of the orders to the match list
            shares_left -= price_queue.total_volume[] # decrement what's left to trade
        else
            while !isempty(price_queue) && !iszero(shares_left) # while not done and queue not empty
                best_ord::Order = popfirst!(price_queue) # pop out best order
                if shares_left >= best_ord.size # Case 1: Limit order gets wiped out
                    # Add best_order to match list & decrement outstanding MO
                    push!(order_match_lst, best_ord)
                    shares_left -= best_ord.size
                else
                    # shares_left < best_ord.size # Case 2: Market Order gets wiped out
                    # Return the difference: LO.size-MO.size back to the order book
                    return_ord = copy_modify_size(best_ord, best_ord.size - shares_left)
                    pushfirst!(price_queue, return_ord)
                    # Add remainder to match list & decrement outstanding MO
                    best_ord = copy_modify_size(best_ord, shares_left)
                    push!(order_match_lst, best_ord)
                    shares_left -= best_ord.size
                end
            end
            if !isempty(price_queue) # If price queue wasn't killed, put it back into the OneSidedBook
                _insert_queue!(sb, price_queue)
            end
        end
    end
    # Return results
    return order_match_lst, shares_left
end

"""
Fill market order by walking the book using Price/Arrival priority.
__This version walks the book with funds rather than trade size.__
"""
function _walk_order_book_byfunds!(
    sb::OneSidedBook{Sz,Px,Oid,Aid},
    order_funds::Real,
    limit_price::Union{Px,Nothing},
    order_mode::OrderTraits,
)::Tuple{Vector{Order{Sz,Px,Oid,Aid}},Real} where {Sz,Px,Oid,Aid}
    # Allocate memory for order output
    order_match_lst = Vector{Order{Sz,Px,Oid,Aid}}()
    funds_left = order_funds # remaining quantity to trade
    # Perform initial available liquidity check
    if isallornone(order_mode) && (_size_available(sb, limit_price) < order_size)
        return order_match_lst, funds_left
    end
    # Perform matching logic
    limit_price_check = Base.Fix2(_is_best_price_inside_limit, limit_price)
    while !isempty(sb.book) && !iszero(funds_left) && limit_price_check(sb) # while book not empty, order not done and best price within limit price
        price_queue::OrderQueue = popfirst_queue!(sb)
        if (price_queue.total_volume[] * price_queue.price) <= funds_left # If entire queue is to be wiped out
            append!(order_match_lst, price_queue.queue) # Add all of the orders to the match list
            funds_left -= price_queue.total_volume[] * price_queue.price # decrement what's left to trade
        else
            while !isempty(price_queue) && !iszero(funds_left) # while not done and queue not empty
                best_ord::Order = popfirst!(price_queue) # pop out best order
                if funds_left >= (best_ord.size * best_ord.price) # Case 1: Limit order gets wiped out
                    # Add best_order to match list & decrement outstanding MO
                    push!(order_match_lst, best_ord)
                    funds_left -= (best_ord.size * best_ord.price)
                else # funds_left < best_ord.size # Case 2: Market Order gets wiped out
                    # Return the difference: LO.size-MO.size back to the order book
                    rem_match_size = floor(Sz, funds_left / best_ord.price)
                    return_ord = copy_modify_size(best_ord, best_ord.size - rem_match_size)
                    pushfirst!(price_queue, return_ord)
                    # Add remainder to match list & decrement outstanding MO
                    best_ord = copy_modify_size(best_ord, rem_match_size)
                    push!(order_match_lst, best_ord)
                    funds_left -= (best_ord.size * best_ord.price)
                end
            end
            if !isempty(price_queue) # If price queue wasn't killed, put it back into the OneSidedBook
                _insert_queue!(sb, price_queue)
            end
        end
    end
    # Return results
    return order_match_lst, funds_left
end

"""
    submit_market_order!(ob::OrderBook,side::OrderSide,mo_size[,fill_mode::OrderTraits])

Submit market order to `ob::OrderBook` with `side::OrderSide` and size `mo_size`.
Optionally `fill_mode::OrderTraits` may be provided to modify fill logic.
Market orders are filled by price-time priority.

Returns tuple `( ord_lst::Vector{Order}, left_to_trade::Sz )`
where
 - `ord_lst` is a list of _limit orders_ that _market order_ matched with
 - `left_to_trade` is the remaining size of un-filled order ( `==0` if order is complete, `>0` if incomplete)

__Note:__ Only `fill_mode.allornone` will be used from `fill_mode::OrderTraits`.
All other entries will be ignored.

"""
function submit_market_order!(
    ob::OrderBook{Sz,Px,Oid,Aid},
    side::OrderSide,
    mo_size::Real;
    fill_mode::OrderTraits=VANILLA_FILLTYPE,
) where {Sz,Px,Oid,Aid}
    mo_size > zero(mo_size) || error("market order argument mo_size must be positive")
    if isbuy(side)
        return _walk_order_book_bysize!(ob.ask_orders, Sz(mo_size), nothing, fill_mode)
    else
        return _walk_order_book_bysize!(ob.bid_orders, Sz(mo_size), nothing, fill_mode)
    end
end

"""
    submit_market_order_byfunds!(ob::OrderBook,side::Symbol,funds[,mode::OrderTraits])

Submit market order to `ob::OrderBook` `side::OrderSide` and available funds `funds::Real`.
Optionally `fill_mode::OrderTraits` may be provided to modify fill logic.
Market orders are filled by price-time priority.

Functionality is exactly the same as `submit_market_order!` except _available funds_ (max total price paid on order)
is provided, rather than _number of shares_ (order size).

Returns tuple `( ord_lst::Vector{Order}, funds_leftover )`
where
 - `ord_lst` is a list of _limit orders_ that _market order_ matched with
 - `funds_leftover` is the amount of remaining funds if not enough liquidity was available ( `==0` if order is complete, `>0` if incomplete)

__Note:__ Only `fill_mode.allornone` will be considered from `fill_mode::OrderTraits`.
All other entries will be ignored.

"""
function submit_market_order_byfunds!(
    ob::OrderBook, side::OrderSide, funds::Real, fill_mode::OrderTraits=VANILLA_FILLTYPE
)
    funds > zero(funds) || error("market order argument funds must be positive")
    if isbuy(side)
        return _walk_order_book_byfunds!(ob.ask_orders, funds, nothing, fill_mode)
    else
        return _walk_order_book_byfunds!(ob.bid_orders, funds, nothing, fill_mode)
    end
end

# Order Cancellation functions

"""
    cancel_order!(ob::OrderBook, o::Order)
    cancel_order!(ob::OrderBook, orderid, side, price [, acct_id=nothing])

Cancels Order `o`, or order with matching information from OrderBook.

Provide `acct_id` if known to guarantee correct account tracking.
"""
function cancel_order!(
    ob::OrderBook{Sz,Px,Oid,Aid}, orderid::Oid, side::OrderSide, price
) where {Sz,Px,Oid,Aid}
    # Delete order from bid (buy) or ask (sell) book
    if isbuy(side)
        popped_ord = pop_order!(ob.bid_orders, Px(price), orderid)
    else
        popped_ord = pop_order!(ob.ask_orders, Px(price), orderid)
    end
    # Delete order from account maps
    if !isnothing(popped_ord) && !isnothing(popped_ord.acctid)
        _delete_order_acct_map!(ob.acct_map, popped_ord.acctid, popped_ord.orderid)
    end
    return popped_ord
end

cancel_order!(ob::OrderBook, o::Order) = cancel_order!(ob, o.orderid, o.side, o.price)


using Serialization
using Parameters
using Setfield
using Dates
using Suppressor
using LinearAlgebra

""" Configuration struct for running experiments """
@with_kw struct TrainConfig
    # metadata
    name::String
    result_path::String = "results"
    data_path::String
    log_path::String = ""
    dtype::Type = Float64
    # assume by default that the system has two threads per core
    blas_threads::Int = Integer(0.5 * length(Sys.cpu_info()))
    # train configuration
    num_epochs::Int
    p_init_fn = initial_params
    patience::Int
    verbose::Bool = true
    logevery::Int = 1
    # model configuration
    ode::O where {O<:Ode}
    link::L where {L<:Link}
    mod::M where {M<:Modifier} = Nominal()
    feed::F where {F<:Feedthrough} = NoFeedthrough()
    # DynamicsTrainer arguments
    solver::Any
    adaptive::Bool = true
    sensealg::Any
    abstol = 1e-6
    reltol = 1e-3
    dt = 1.0
    initial_condition = solve_ss
    loss = Flux.Losses.mse
    regularization = no_regularization
    regularization_coef = 0.0
    optim = ADAMW(dtype(1e-3))
    wdecay = WeightDecay(dtype(0.0))
    sched = Step(λ = dtype(1e-3), γ = dtype(1.0), 1)
    validation_split::Float64 = 0.0
    test_split::Float64 = 0.0
    train_split::Float64 = 1.0
    batch_size::Int
    time_dropout::Float64
    time_horizon::Int
    time_horizon_increment::Int
    time_horizon_max::Int
    callback_tasks = []
end

""" Unpack and run model train configuration """
function run_config(c::TrainConfig; runs = 1)
    @unpack_TrainConfig c

    # set number of threads to be used by BLAS
    LinearAlgebra.BLAS.set_num_threads(blas_threads)

    # display multithreading information
    println(
        "Number of BLAS threads: $(ccall((:openblas_get_num_threads64_, Base.libblas_name), Cint, ())
), Number of Julia threads: $(Threads.nthreads())",
    )

    for _ = 1:runs
        # create experiment result path
        savepath = joinpath(result_path, name, Dates.format(now(), "HH_MM_SS_MS"))
        if ~ispath(savepath)
            mkpath(savepath)
        end

        # write config
        serialize(joinpath(savepath, "config.jls"), c)

        # write human readable config as txt file
        config_string = @capture_out show(c)
        open(joinpath(savepath, "config.txt"), "w") do io
            write(io, config_string)
        end

        # read data
        dataset = deserialize(data_path)
        t, u, y = dataset["t"], dataset["u"], dataset["y"]

        # model training
        trainer = DynamicsTrainer(
            ode,
            link,
            mod,
            feed,
            solver,
            sensealg,
            Array{dtype,1}(t),
            Array{dtype,3}(u),
            Array{dtype,3}(y);
            savepath = savepath,
            dtype = dtype,
            validation_split = validation_split,
            test_split = test_split,
            abstol = abstol,
            reltol = reltol,
            dt = dt,
            adaptive = adaptive,
            initial_condition = initial_condition,
            loss = loss,
            regularization = regularization,
            regularization_coef = regularization_coef,
            optim = optim,
            wdecay = wdecay,
            sched = sched,
            batch_size = batch_size,
            callback_tasks = callback_tasks,
        )

        train(
            trainer,
            num_epochs;
            p_init = p_init_fn(trainer; log_path = log_path),
            time_dropout = time_dropout,
            time_horizon = time_horizon,
            time_horizon_increment = time_horizon_increment,
            time_horizon_max = time_horizon_max,
            patience = patience,
            logevery = logevery,
            verbose = verbose,
        )
    end
end