using Test, Onda, Dates, MsgPack

@testset "round trip" begin
    mktempdir() do root
        # generate a test dataset
        dataset = Dataset(joinpath(root, "test"); create=true)
        @test dataset isa Dataset
        @test isdir(dataset.path)
        @test isdir(joinpath(dataset.path, "samples"))
        duration_in_seconds = Second(10)
        duration_in_nanoseconds = Nanosecond(duration_in_seconds)
        uuid, recording = create_recording!(dataset)
        Ts = (UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64)
        sample_rate = 50.5
        signals = Dict(Symbol(:x, i) => Signal(Symbol.([:a, :b, :c], i), Nanosecond(0),
                                               duration_in_nanoseconds, Symbol(:unit, i),
                                               0.25, i, T, sample_rate, Symbol("lpcm.zst"),
                                               nothing) for (i, T) in enumerate(Ts))
        samples = Dict(k => Samples(v, true, rand(v.sample_type, 3, sample_count(v)))
                       for (k, v) in signals)
        for (name, s) in samples
            @test channel_count(s) == length(s.signal.channel_names)
            @test channel_count(s.signal) == length(s.signal.channel_names)
            @test sample_count(s) == size(s.data, 2)
            @test sizeof_samples(s.signal) == sizeof(s.data)
            @test encode(s) === s
            tmp = similar(s.data)
            encode!(tmp, s)
            @test tmp == s.data
            d = decode(s)
            @test decode(d) === d
            tmp = similar(d.data)
            decode!(tmp, d)
            @test tmp == d.data
            tmp = similar(d.data)
            decode!(tmp, s)
            @test tmp == d.data
            tmp = similar(d.data)
            decode!(tmp, s.signal.sample_resolution_in_unit, s.signal.sample_offset_in_unit,
                    s.data)
            @test tmp == d.data
            @test d.data == (s.data .* s.signal.sample_resolution_in_unit .+
                   s.signal.sample_offset_in_unit)
            if sizeof(s.signal.sample_type) >= 8
                # decoding from 64-bit to floating point is fairly lossy
                tmp = similar(s.data)
                @test isapprox(encode(d, nothing).data, s.data; rtol=10)
                encode!(tmp, d, nothing)
                @test isapprox(tmp, s.data; rtol=10)
                @test isapprox(encode(d, missing).data, s.data; rtol=10)
                encode!(tmp, d, missing)
                @test isapprox(tmp, s.data; rtol=10)
            else
                tmp = similar(s.data)
                encode!(tmp, d, nothing)
                @test tmp == s.data
                @test encode(d, nothing).data == s.data
                encode!(tmp, d, missing)
                @test isapprox(tmp, s.data; rtol=1)
                @test isapprox(encode(d, missing).data, s.data; rtol=1)
            end
            chs = s.signal.channel_names
            i = 93
            @test s[:, i].data == s.data[:, i:i]
            t = Onda.time_from_index(s.signal.sample_rate, i)
            t2 = t + Nanosecond(Second(3))
            j = Onda.index_from_time(s.signal.sample_rate, t2) - 1
            for ch_inds in (:, 1:2, 2:3, 1:3, [3, 1], [2, 3, 1], [1, 2, 3])
                @test s[chs[ch_inds], t].data == s[ch_inds, i].data
                @test s[chs[ch_inds], TimeSpan(t, t2)].data == s.data[ch_inds, i:j]
                @test s[chs[ch_inds], i:j].data == s.data[ch_inds, i:j]
                @test s[ch_inds, t].data == s[ch_inds, i].data
                @test s[ch_inds, TimeSpan(t, t2)].data == s.data[ch_inds, i:j]
                @test s[ch_inds, i:j].data == s.data[ch_inds, i:j]
            end
            @test size(s[:, TimeSpan(0, Second(1))].data, 2) == floor(s.signal.sample_rate)
            for i in 1:length(chs)
                @test channel(s, chs[i]) == i
                @test channel(s, i) == chs[i]
                @test channel(s.signal, chs[i]) == i
                @test channel(s.signal, i) == chs[i]
            end
            @test duration(s) == duration_in_seconds
            @test s[:, TimeSpan(0, duration(s))].data == s.data
            store!(dataset, uuid, name, s)
        end
        save_recordings_file(dataset)

        # read back in the test dataset, add some annotations
        old_dataset = dataset
        dataset = Dataset(joinpath(root, "test"))
        @test length(dataset.recordings) == 1
        uuid, recording = first(dataset.recordings)
        x1 = load(dataset, uuid, :x1)
        @test x1.signal == signals[:x1]
        xs = load(dataset, uuid, (:x3, :x2))
        @test xs[:x3].signal == signals[:x3]
        @test xs[:x2].signal == signals[:x2]
        xs = load(dataset, uuid)
        span = TimeSpan(Second(1), Second(2))
        xs_span = load(dataset, uuid, span)
        for (name, s) in samples
            xi = xs[name]
            @test xi.signal == signals[name]
            @test xi.encoded
            @test xi.data == s.data
            @test xi[:, span].data == xs_span[name].data
        end
        for i in 1:3
            annotate!(recording,
                      Annotation("value_$i", Nanosecond(i),
                                 Nanosecond(i + rand(1:1000000))))
        end
        save_recordings_file(dataset)

        # read back in annotations
        old_uuid = uuid
        old_recording = recording
        old_dataset = dataset
        dataset = Dataset(joinpath(root, "test"))
        uuid, recording = first(dataset.recordings)
        @test old_recording == recording
        delete!(dataset.recordings, uuid)
        uuid, recording = create_recording!(dataset)
        foreach(x -> annotate!(recording, x), old_recording.annotations)
        foreach(x -> store!(dataset, uuid, x, load(old_dataset, old_uuid, x)),
                keys(old_recording.signals))
        merge!(dataset, old_dataset; only_recordings=true)
        @test length(dataset.recordings) == 2
        r1 = dataset.recordings[old_uuid]
        @test r1 == old_recording
        r2 = dataset.recordings[uuid]
        @test r2 == recording
        @test old_uuid != uuid
        @test r1.signals == r2.signals
        @test r1.annotations == r2.annotations

        old_duration = duration(r2)
        new_duration = old_duration + Nanosecond(1)
        r2signals = set_span!(r2, TimeSpan(Nanosecond(0), new_duration))
        @test keys(r2signals) == keys(r2.signals)
        @test all(duration.(values(r2signals)) .== new_duration)
        set_span!(r2, TimeSpan(Nanosecond(0), old_duration))

        r = dataset.recordings[uuid]
        original_signals_length = length(r.signals)
        signal_name, signal = first(r.signals)
        signal_samples = load(dataset, uuid, signal_name)
        signal_samples_path = samples_path(dataset, uuid, signal_name)
        delete!(dataset, uuid, signal_name)
        @test r === dataset.recordings[uuid]
        @test length(r.signals) == (original_signals_length - 1)
        @test !haskey(r.signals, signal_name)
        @test !isfile(signal_samples_path)
        store!(dataset, uuid, signal_name, signal_samples)

        # read back everything, but without assuming an order on the metadata
        dataset = Dataset(joinpath(root, "test"))
        Onda.write_recordings_file(dataset.path,
                                   Onda.Header(dataset.header.onda_format_version, false),
                                   dataset.recordings)
        dataset = Dataset(joinpath(root, "test"))
        @test Dict(old_uuid => old_recording) == dataset.recordings
        delete!(dataset, old_uuid)
        save_recordings_file(dataset)

        # read back the dataset that should now be empty
        dataset = Dataset(joinpath(root, "test"))
        @test isempty(dataset.recordings)
        @test !isdir(joinpath(dataset.path, "samples", string(old_uuid)))

        # make sure samples directory is appropriately created if not present
        no_samples_path = joinpath(root, "no_samples_dir.onda")
        mkdir(no_samples_path)
        cp(joinpath(dataset.path, "recordings.msgpack.zst"),
           joinpath(no_samples_path, "recordings.msgpack.zst"))
        Dataset(no_samples_path; create=false)
        return @test isdir(joinpath(no_samples_path, "samples"))
    end
end

@testset "Error conditions" begin
    mktempdir() do root
        mkdir(joinpath(root, "i_exist.onda"))
        touch(joinpath(root, "i_exist.onda", "memes"))
        @test_throws ArgumentError Dataset(joinpath(root, "i_exist.onda"); create=true)

        dataset = Dataset(joinpath(root, "okay.onda"); create=true)
        uuid, recording = create_recording!(dataset)
        signal = Signal([:a], Nanosecond(0), Nanosecond(Second(10)), :mv, 0.25, 0.0, Int8,
                        100, Symbol("lpcm.zst"), nothing)
        @test_throws DimensionMismatch Samples(signal, true, rand(Int8, 2, 10))
        @test_throws ArgumentError Samples(signal, true, rand(Float32, 1, 10))
        samples = Samples(signal, true, rand(Int8, 1, 10 * 100))
        @test_throws ArgumentError store!(dataset, uuid, Symbol("***HI***"), samples)
        store!(dataset, uuid, :name_okay, samples)
        @test_throws ArgumentError store!(dataset, uuid, :name_okay, samples;
                                          overwrite=false)

        @test_throws ArgumentError Annotation("hi", Nanosecond(20), Nanosecond(4))

        mkdir(joinpath(root, "other.onda"))
        other = Dataset(joinpath(root, "other.onda"); create=true)  # Using existing empty directory
        create_recording!(other, uuid)
        @test_throws ArgumentError create_recording!(other, uuid)
        store!(other, uuid, :cool_stuff, samples)
        @test_throws ErrorException merge!(dataset, other; only_recordings=false)
        return @test_throws ArgumentError merge!(dataset, other; only_recordings=true)
    end
end

@testset "upgrade_onda_format_from_v0_2_to_v0_3!" begin
    mktempdir() do new_path
        old_path = joinpath(@__DIR__, "old_test_v0_2.onda")
        cp(old_path, new_path; force=true)
        dataset = Onda.upgrade_onda_format_from_v0_2_to_v0_3!(new_path,
                                                              (k, v) -> string(k, '.', v))
        @test dataset.path == new_path
        @test dataset.header.onda_format_version == v"0.3.0"
        @test dataset.header.ordered_keys
        old_recordings = MsgPack.unpack(Onda.zstd_decompress(read(joinpath(old_path,
                                                                           "recordings.msgpack.zst"))))[2]
        new_customs = MsgPack.unpack(Onda.zstd_decompress(read(joinpath(new_path,
                                                                        "recordings_custom.msgpack.zst"))))
        @test length(dataset.recordings) == 1
        @test length(new_customs) == 1
        uuid = first(keys(dataset.recordings))
        recording = first(values(dataset.recordings))
        old_recording = first(values(old_recordings))
        @test string(uuid) == first(keys(new_customs)) == first(keys(old_recordings))
        @test first(values(new_customs)) == old_recording["custom"]
        sorted_annotations = sort(collect(recording.annotations); by=first)
        sorted_old_annotations = sort(old_recording["annotations"];
                                      by=(x -> x["start_nanosecond"]))
        @test length(sorted_annotations) == length(sorted_old_annotations)
        for (ann, old_ann) in zip(sorted_annotations, sorted_old_annotations)
            @test ann.value == string(old_ann["key"], '.', old_ann["value"])
            @test ann.start_nanosecond.value == old_ann["start_nanosecond"]
            @test ann.stop_nanosecond.value == old_ann["stop_nanosecond"]
        end
        old_signals = old_recording["signals"]
        @test keys(recording.signals) == Set(Symbol.(keys(old_signals)))
        return for (signal_name, signal) in recording.signals
            old_signal = old_signals[string(signal_name)]
            @test signal.channel_names == Symbol.(old_signal["channel_names"])
            @test signal.start_nanosecond == Nanosecond(0)
            @test signal.stop_nanosecond ==
                  Nanosecond(old_recording["duration_in_nanoseconds"])
            @test signal.sample_unit == Symbol(old_signal["sample_unit"])
            @test signal.sample_resolution_in_unit ==
                  old_signal["sample_resolution_in_unit"]
            @test signal.sample_offset_in_unit == 0.0
            @test signal.sample_type ==
                  Onda.julia_type_from_onda_sample_type(old_signal["sample_type"])
            @test signal.sample_rate == old_signal["sample_rate"]
            @test signal.file_extension == Symbol(old_signal["file_extension"])
            @test signal.file_options == old_signal["file_options"]
        end
    end
end


include("MetaDict.jl")

export Image, ImageHeader, ImageFlags

@MRD.enm ImageType::UInt16 magnitude = 1 phase = 2 real = 3 imag = 4 complex = 5
@MRD.enm TypeIndex::UInt16 ushort = 1 short = 2 uint = 3 int = 4 float = 5 double = 6 cxfloat =
    7 cxdouble = 8

datatype_to_type = Base.Dict{TypeIndex.Enm,Type}([
    (TypeIndex.ushort, UInt16),
    (TypeIndex.short, Int16),
    (TypeIndex.uint, UInt32),
    (TypeIndex.int, Int32),
    (TypeIndex.float, Float32),
    (TypeIndex.double, Float64),
    (TypeIndex.cxfloat, ComplexF32),
    (TypeIndex.cxdouble, ComplexF64),
])
type_to_datatype = Base.Dict(value => key for (key, value) in datatype_to_type)
is_mrd_datatype(X) = X ∈ keys(type_to_datatype)

module ImageFlags
import FlagSets
FlagSets.@flagset Flags {_, UInt64} begin
    :is_navigation_data
    72057594037927936 --> :user1
    :user2
    :user3
    :user4
    :user5
    :user6
    :user7
    :user8
end
end

@with_kw struct ImageHeader
    version::UInt16 = 1
    flags::ImageFlags.Flags = ImageFlags.Flags(0)
    measurement_uid::UInt32 = 0
    field_of_view::NTuple{3,Float32} = (0, 0, 0)
    position::NTuple{3,Float32} = (0, 0, 0)
    read_dir::NTuple{3,Float32} = (1, 0, 0)
    phase_dir::NTuple{3,Float32} = (0, 1, 0)
    slice_dir::NTuple{3,Float32} = (0, 0, 1)
    patient_table_position::NTuple{3,Float32} = (0, 0, 0)
    average::UInt16 = 0
    slice::UInt16 = 0
    contrast::UInt16 = 0
    phase::UInt16 = 0
    repetition::UInt16 = 0
    set::UInt16 = 0
    acquisition_time_stamp::UInt32 = 0
    physiology_time_stamp::NTuple{3,UInt32} = (0, 0, 0)
    image_type::ImageType.Enm = ImageType.magnitude
    image_index::UInt16 = 0
    image_series_index::UInt16 = 0
    user_int::NTuple{8,Int32} = (0, 0, 0, 0, 0, 0, 0, 0)
    user_float::NTuple{8,Float32} = (0, 0, 0, 0, 0, 0, 0, 0)
end

struct Image
    header::ImageHeader
    data::Array{T,4} where {T<:Number}
    meta::MetaDict
    function Image(header::ImageHeader, data::Array, meta = MetaDict())
        is_mrd_datatype(eltype(data)) || error("Only MRD standard types supported")
        ndims(data) <= 4 || error("Images have a maximum of 4 dimensions")
        data = reshape(data, Val(4))
        new(header, data, meta)
    end

end

function Base.:(==)(img1::Image,img2::Image)
    return img1.header == img2.header && img1.data == img2.data && img1.meta == img2.meta
end

struct RawImageHeader

    version::UInt16
    data_type::TypeIndex.Enm
    flags::ImageFlags.Flags
    measurement_uid::UInt32
    matrix_size::NTuple{3,UInt16}
    field_of_view::NTuple{3,Float32}
    channels::UInt16
    position::NTuple{3,Float32}
    read_dir::NTuple{3,Float32}
    phase_dir::NTuple{3,Float32}
    slice_dir::NTuple{3,Float32}
    patient_table_position::NTuple{3,Float32}
    average::UInt16
    slice::UInt16
    contrast::UInt16
    phase::UInt16
    repetition::UInt16
    set::UInt16
    acquisition_time_stamp::UInt32
    physiology_time_stamp::NTuple{3,UInt32}
    image_type::ImageType.Enm
    image_index::UInt16
    image_series_index::UInt16
    user_int::NTuple{8,Int32}
    user_float::NTuple{8,Float32}
    attribute_string_len::UInt32
end

function RawImageHeader(image::Image)
    channels = UInt16(size(image.data)[4])
    matrix_size = Tuple{UInt16,UInt16,UInt16}(size(image.data)[1:3])
    data_type = type_to_datatype[eltype(image.data)]
    hdr = image.header
    attribute_string_len = ncodeunits(to_xml_string(image.meta))

    RawImageHeader(
        hdr.version,
        data_type,
        hdr.flags,
        hdr.measurement_uid,
        matrix_size,
        hdr.field_of_view,
        channels,
        hdr.position,
        hdr.read_dir,
        hdr.phase_dir,
        hdr.slice_dir,
        hdr.patient_table_position,
        hdr.average,
        hdr.slice,
        hdr.contrast,
        hdr.phase,
        hdr.repetition,
        hdr.set,
        hdr.acquisition_time_stamp,
        hdr.physiology_time_stamp,
        hdr.image_type,
        hdr.image_index,
        hdr.image_series_index,
        hdr.user_int,
        hdr.user_float,
        attribute_string_len
    )

end

@generated function ImageHeader(raw::RawImageHeader)
    names = fieldnames(ImageHeader)
    args = names .|> s -> Expr(:call,:getfield,:raw,QuoteNode(s))
    :(ImageHeader($(args...)))
end

function ImageHeader(acq::AcquisitionHeader; kwargs...)
    names = (:version, :measurement_uid, :position, :read_dir, :phase_dir, :slice_dir, :patient_table_position, :acquisition_time_stamp, :physiology_time_stamp, :user_int, :user_float )

    fields_acq = LittleDict(n => getfield(acq,n) for n in names)

    names_idx = (:average,:slice, :contrast,:phase,:repetition, :set)
    fields_idx = LittleDict(n => getfield(acq.idx,n) for n in names_idx)

    fields = merge(fields_acq,fields_idx,kwargs)

    return ImageHeader(;fields...)
end

write(io::IO, header::RawImageHeader) = fieldnames(RawImageHeader) .|> getfield $ header .|> MRD.unsafe_write $ io

function write(io::IO, img::Image)
    write(io, RawImageHeader(img))
    write(io,img.meta)
    Base.write(io, img.data)
end

read(io::IO, ::Type{RawImageHeader}) = RawImageHeader((fieldnames(RawImageHeader) .|> fieldtype $ RawImageHeader .|> MRD.unsafe_read $ io)...)

function read(io::IO, ::Type{Image})
    header = MRD.read(io, RawImageHeader)
    meta = MRD.read(io, MetaDict)
    data_type = datatype_to_type[header.data_type]

    data = Array{data_type,4}(undef, header.matrix_size[1], header.matrix_size[2], header.matrix_size[3], header.channels)
    Base.read!(io, data)
    return Image(ImageHeader(header), data, meta)

end

Base.precompile(Tuple{Type{RawImageHeader},Image})   # time: 0.089653455


using MusicXML

@MX begin
## Part List
### Piano
midiinstrument_piano = MidiInstrument(channel= 1, program =1, volume = 100, pan =0, id = "P1-I1")
scorepart_piano = ScorePart(name = "Piano", midiinstrument = midiinstrument_piano, id = "P1")

### Acoustic Guitar
midiinstrument_guitar = MidiInstrument(channel= 1, program =26, volume = 100, pan =0, id = "P2-I1")
scorepart_guitar = ScorePart(name = "Acoustic Guitar", midiinstrument = midiinstrument_guitar, id = "P2")

###
partlist = PartList(scoreparts = [scorepart_piano, scorepart_guitar])

## Part
### Piano

attributes1_piano = Attributes(
   time = Time(beats = 4, beattype = 4), # 4/4
   divisions = 4, # we want to use 16th notes at minimum
   clef = [Clef(number = 1, sign = "G", line = 2), Clef(number = 2, sign = "F", line = 4)], # Piano clefs
   staves = 2, # Piano staves
   key = Key(fifths = 0, mode = "major"), # no accidentals, major key
)

measure1_notes_piano = [
Note(pitch = Pitch(step = "C", alter = 0, octave = 4), duration =  4),
Note(pitch = Pitch(step = "D", alter = 0, octave = 4), duration =  4),
Note(pitch = Pitch(step = "E", alter = 0, octave = 4), duration =  4),
Note(pitch = Pitch(step = "F", alter = +1, octave = 4), duration =  4),
]

measure2_notes_piano = [
Note(pitch = Pitch(step = "G", alter = 0, octave = 5), duration =  1),
Note(pitch = Pitch(step = "G", alter = +1, octave = 5), duration =  1),
Note(pitch = Pitch(step = "B", alter = 0, octave = 5), duration =  1),
Note(pitch = Pitch(step = "A", alter = +1, octave = 5), duration =  1),
Note(rest = Rest(), duration =  4), # Rest
Note(pitch = Pitch(step = "A", alter = 0, octave = 5), duration =  4),
Note(pitch = Pitch(step = "B", alter = 0, octave = 5), duration =  4),
]

measures_piano = [
Measure(attributes = attributes1_piano, notes = measure1_notes_piano) # measure 1 has attributes
# Measure(notes = measure2_notes_piano)
]

part_piano = Part(measures = measures_piano, id = "P1")

### Guitar

attributes1_guitar = Attributes(
   time = Time(beats = 4, beattype = 4), # 4/4
   divisions = 4, # we want to use 16th notes at minimum
   clef = [Clef(number = 1, sign = "G", line = 2), Clef(number = 2, sign = "TAB", line = 6)], # Guitar Clefs
   staves = 2, # Guitar staves
   key = Key(fifths = 0, mode = "major"), # no accidentals, major key
)

measure1_notes_guitar = [
# G Major chord for a bar
Note(pitch = Pitch(step = "G", alter = 0, octave = 2), duration =  16),
Note(pitch = Pitch(step = "B", alter = 0, octave = 2), duration =  16, chord = Chord()),
Note(pitch = Pitch(step = "D", alter = 0, octave = 3), duration =  16, chord = Chord()),
Note(pitch = Pitch(step = "G", alter = 0, octave = 3), duration =  16, chord = Chord()),
Note(pitch = Pitch(step = "D", alter = 0, octave = 4), duration =  16, chord = Chord()),
Note(pitch = Pitch(step = "G", alter = 0, octave = 4), duration =  16, chord = Chord()),
]

measure2_notes_guitar = [
# G Major chord for half a bar
Note(pitch = Pitch(step = "G", alter = 0, octave = 2), duration =  8),
Note(pitch = Pitch(step = "B", alter = 0, octave = 2), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "D", alter = 0, octave = 3), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "G", alter = 0, octave = 3), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "D", alter = 0, octave = 4), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "G", alter = 0, octave = 4), duration =  8, chord = Chord()),

# G Major chord for half a bar
Note(pitch = Pitch(step = "G", alter = 0, octave = 2), duration =  8),
Note(pitch = Pitch(step = "B", alter = 0, octave = 2), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "D", alter = 0, octave = 3), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "G", alter = 0, octave = 3), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "D", alter = 0, octave = 4), duration =  8, chord = Chord()),
Note(pitch = Pitch(step = "G", alter = 0, octave = 4), duration =  8, chord = Chord()),
]

measures_guitar = [
Measure(attributes = attributes1_guitar, notes = measure1_notes_guitar) # measure 1 has attributes
Measure(notes = measure2_notes_guitar)
]

part_guitar = Part(measures = measures_guitar, id = "P2")

##
score = ScorePartwise(
partlist = partlist,
parts =  [part_piano, part_guitar],
)

writemusicxml("myscore.musicxml", score)

end


############################################################################################
"""
$(TYPEDEF)

MCMC Tuning container.

# Fields
$(TYPEDFIELDS)
"""
struct MCMCTune{
    A<:BaytesCore.UpdateBool,
    B<:BaytesCore.UpdateBool,
    F<:AbstractFloat,
    T<:Tagged,
    E<:Tuple,
    P<:Proposal} <: AbstractTune
    "Tagged Parameter."
    tagged::T
    "Current Phase in MCMC Cycle"
    phase::PhaseTune{E}
    "Stepsize container"
    stepsize::StepSizeTune{A,F}
    "Information for posterior covariance estimate"
    proposal::P
    "Boolean if generated quantities should be generated while sampling"
    generated::B
    "Current iteration number"
    iter::Iterator
    function MCMCTune(
        objective::Objective,
        phase::PhaseTune{E},
        stepsize::StepSizeTune{A,F},
        proposal::P,
        generate::B,
    ) where {A<:BaytesCore.UpdateBool,B<:BaytesCore.UpdateBool,F<:AbstractFloat,E<:Tuple,P<:Proposal}
        #!NOTE: Start with 0, so first proposal step will update iter to 1
        iter = Iterator(0)
        return new{A,B,F,typeof(objective.tagged),E,P}(
            objective.tagged, phase, stepsize, proposal, generate, iter
        )
    end
end

############################################################################################
"""
$(SIGNATURES)
Update MCMC tuning fields at current iteration.

# Examples
```julia
```

"""
function update!(
    tune::MCMCTune, result::S, acceptrate::T
) where {S<:ℓObjectiveResult,T<:Real}
    @unpack θᵤ = result
    ##  Update Current iteration counter
    update!(tune.iter)
    ## Update current sampling phase
    #!NOTE: Done AFTER tune.iter update; both iterations start from 0, so to check if phase needs to be updated, tune.iter has to be updated before
    update!(tune.phase, tune.iter.current)
    ## Update stepsize and proposal parameter given current sampling phase
    update!(tune.stepsize, acceptrate, tune.phase)
    update!(tune.proposal, θᵤ, tune.phase)
    ## Pack container
    return nothing
end

############################################################################################
#=
"""
$(SIGNATURES)
Some basic tuning output figures. Useful for debugging.

# Examples
```julia
```

"""
function printtune(tune::MCMCTune, diagparam::Integer=size(tune.proposal.chain, 2))
    ArgCheck.@argcheck tune.iter.current > 0 "Need to run a proposal step before accessing tuning configuration."
    println("#################################################")
    println("Tune iter: ", tune.iter.current)
    println("Phase iter: ", tune.phase.iter.current)
    println("Phase name: ", tune.phase.name[tune.phase.counter.current])
    println("Phase update: ", tune.phase.update.current)
    println("Proposal update: ", tune.proposal.adaption)
    println("Phase counter: ", tune.phase.counter.current)
    #!NOTE: This will be Dualaveraging Parameter for first update! in new samplingphase, hence in final sampling phase only the smoothed ϵ is returned, and the dualaverage has no impact cause it wont compute anymore when Exploration+update=false
    #!NOTE: Hence, the stepsize multiplier will only come into effect at NEXT iteration with un-smoothed epsilon, in case adaption is still performed
    println("Phase stepsize: ", tune.stepsize.ϵ)
    println(
        "Phase (avg) stepsize: ",
        exp(tune.stepsize.dualaverage.logϵ),
        " (",
        exp(tune.stepsize.dualaverage.logϵ̄),
        ")",
    )
    println(
        "Proposal chain current row: ", tune.proposal.chain[end, tune.phase.iter.current]
    )
    println("Proposal chain last element: ", tune.proposal.chain[end])
    println("Proposal covariance: ", tune.proposal.Σ.diag[1:diagparam])
    return nothing
end
=#
############################################################################################
# Export
export MCMCTune


#Finding PROPERTIES WITH NUMERICAL VALUES in Trivial inputs
tpropwithnumval = Set();
@info "Finding trivials with numerical value"
for p in coolproptrivialparameters
    missed = Set();
    for fluid in coolpropfluids
        try
            res = ("$(PropsSI(p, Compat.String(fluid)))");
            push!(tpropwithnumval, p);
        catch err
            push!(missed, fluid);
        end
    end
    if (length(missed) == length(coolpropfluids))
    # if (findfirst(fails_any_props_trivals, p) == 0)
        if findfirst(x->x==p, fails_any_props_trivals) == nothing
            @warn "None of the defined fluids return a number for $p (expected to have numerical value.)";
        else
            println("$p is trivial without any numericalvalue");
        end
    else
        length(missed) > 0 && println("Only $(collect(setdiff(coolpropfluids, missed))) return number for $p");
    end
end
if ((tpropwithnumval == Set(trivalwithnumval)) && (Set(setdiff(coolproptrivialparameters, tpropwithnumval)) == Set(fails_any_props_trivals)))
    @info "Trivial inputs with numerical values are as expected.";
else
    @warn "Trivial inputs with numerical values are not as expected.";
end


module Samplers

export GibbsSampler, SequentialMonteCarlo, SequentialImportanceSampler, MetropolisHastingsSampler
export MetropolisHastingsSampler, MetropolisWithinGibbsSampler

abstract type Sampler end

abstract type MetropolisWithinGibbsSampler <: Sampler  end

struct GibbsSampler <: MetropolisWithinGibbsSampler
    num_iterations::Int64
    num_burn_in::Int64
    skip::Int64
    assignment_types::Vector{String}
    function GibbsSampler(;num_iterations::Int64=1000, num_burn_in::Int64=0, skip::Int64=1,
                          assignment_types::Vector{String}=["random"])
        if num_iterations < 1
            error("Number of iterations should be positive.")
        end
        if num_burn_in < 0
            error("Number of burn in iterations should be positive.")
        end
        if !(all(map(x -> x in ["random", "all same cluster", "all different clusters"], assignment_types)))
            error("Initial assignment types can only be one of 'random', 'all same cluster' or 'all different clusters'")
        end
        return new(num_iterations, num_burn_in, skip, assignment_types)
    end
end

mutable struct MetropolisHastingsSampler <: MetropolisWithinGibbsSampler
    num_iterations::Int64
    num_burn_in::Int64
    proposal_radius::Int64
    skip::Int64
    adaptive::Bool
    assignment_types::Vector{String}
    function MetropolisHastingsSampler(;num_iterations=1000, num_burn_in=1000, proposal_radius=5, skip=1,
                                       adaptive=false, assignment_types::Vector{String}=["random"])
        if num_iterations < 1
            error("Number of iterations should be positive.")
        end
        if num_burn_in < 0
            error("Number of burn in iterations should be positive.")
        end
        if !(all(map(x -> x in ["random", "all same cluster", "all different clusters"], assignment_types)))
            error("Initial assignment types can only be one of 'random', 'all same cluster' or 'all different clusters'")
        end
        return new(num_iterations, num_burn_in, proposal_radius, skip, adaptive, assignment_types)
    end
end

struct SequentialMonteCarlo <: Sampler
    num_particles::Int64
    ess_threshold::Float64
    function SequentialMonteCarlo(;num_particles::Int64=1000, ess_threshold::T=0.5) where
                                  {T <: Real}
        if num_particles < 1
            error("Number of particles should be positive.")
        end
        if ess_threshold < 0
            error("ESS Threshold should be non-negative.")
        end
        if typeof(ess_threshold) <: AbstractFloat && ess_threshold > 1
            error("ESS Threshold should be either a non-negative integer or a fraction")
        end
        if ess_threshold < 1
            ess_threshold = num_particles*ess_threshold
        end
        return new(num_particles, ess_threshold)
    end
end

struct SequentialImportanceSampler <: Sampler
    num_particles::Int64
    function SequentialImportanceSampler(;num_particles=1000)
        return new(num_particles)
    end
end

end


module SeisCore

# packages
using Base, Core, SeisIO, SeisNoise, Dates, CSV, DataFrames, SCEDC, AWS, Distributed, Parallelism, AWSS3, Glob,
        Statistics, JLD2, HDF5, AbstractFFTs, Plots

# functions
include("download.jl")
include("process.jl")
include("correlate.jl")
include("big_correlate.jl")
end


# We document types/constructors/functors in one docstring until
# https://github.com/JuliaDocs/Documenter.jl/issues/558 is resolved
"""
    AbstractAggregation

Supertype for any aggregation operator.

See also: [`AggregationStack`](@ref), [`SegmentedSum`](@ref), [`SegmentedMax`](@ref),
    [`SegmentedMean`](@ref), [`SegmentedPNorm`](@ref), [`SegmentedLSE`](@ref).
"""
abstract type AbstractAggregation end

@inline _bagnorm(w::Nothing, b) = length(b)
@inline _bagnorm(w::AbstractVector, b) = @views sum(w[b])
@inline _bagnorm(w::AbstractMatrix, b) = @views vec(sum(w[:, b], dims=2))

@inline _weight(w::Nothing, _, _, ::Type{T}) where T = one(T)
@inline _weight(w::AbstractVector, _, j, _) = w[j]
@inline _weight(w::AbstractMatrix, i, j, _) = w[i, j]

@inline _weightsum(ws::Real, _) = ws
@inline _weightsum(ws::AbstractVector, i) = ws[i]

# more stable definitions for r_map and p_map
ChainRulesCore.rrule(::typeof(softplus), x) = softplus.(x), Δ -> (NoTangent(), Δ .* σ.(x))

# our definition of type min for Maybe{...} types
_typemin(t::Type) = typemin(t)
_typemin(::Type{Missing}) = missing
_typemin(::Type{Maybe{T}}) where T = typemin(T)

function _check_agg(a::AbstractAggregation, X::Missing) end
function _check_agg(a::AbstractAggregation, X::AbstractMatrix)
    if size(X, 1) != length(a.ψ)
        DimensionMismatch(
            "Different number of rows in input ($(size(X, 2))) and aggregation ($(length(a.ψ)))"
        ) |> throw
    end
end

include("segmented_sum.jl")
include("segmented_mean.jl")
include("segmented_max.jl")
include("segmented_pnorm.jl")
include("segmented_lse.jl")
# include("transformer.jl")

include("aggregation_stack.jl")

Base.vcat(as::AbstractAggregation...) = reduce(vcat, as |> collect)
function Base.reduce(::typeof(vcat), as::Vector{<:AbstractAggregation})
    AggregationStack(tuple(as...))
end

include("bagcount.jl")

(a::Union{AbstractAggregation, BagCount})(x::ArrayNode, args...) = ArrayNode(a(x.data, args...))

# definitions for mixed aggregations
const names = ["Sum", "Mean", "Max", "PNorm", "LSE"]
for p in filter(p -> length(p) > 1, collect(powerset(collect(1:length(names)))))
    s = Symbol("Segmented", names[p]...)
    @eval begin
        """
            $($(s))([t::Type, ]d::Int)

        Construct [`AggregationStack`](@ref) consisting of $($(
            join("[`Segmented" .* names[p] .* "`](@ref)", ", ", " and ")
       )) operator$($(length(p) > 1 ? "s" : "")).

        $($(
            all(in(["Sum", "Mean", "Max"]), names[p]) ? """
            # Examples
            ```jldoctest
            julia> $(s)(4)
            AggregationStack:
            $(join(" Segmented" .* names[p] .* "(ψ = Float32[0.0, 0.0, 0.0, 0.0])", "\n"))

            julia> $(s)(Float64, 2)
            AggregationStack:
            $(join(" Segmented" .* names[p] .* "(ψ = [0.0, 0.0])", "\n"))
            ```
            """ : ""
        ))

        See also: [`AbstractAggregation`](@ref), [`AggregationStack`](@ref), [`SegmentedSum`](@ref),
            [`SegmentedMax`](@ref), [`SegmentedMean`](@ref), [`SegmentedPNorm`](@ref), [`SegmentedLSE`](@ref).
        """
        function $s(d::Int)
            AggregationStack($((Expr(:call, Symbol("Segmented", n), :d) for n in names[p])...))
        end
    end
    @eval function $s(::Type{T}, d::Int) where T
        AggregationStack($((Expr(:call, Symbol("Segmented", n), :T, :d) for n in names[p])...))
    end
    @eval export $s
end


# add more julia processes
nprocs() < 4 ? addprocs(4-nprocs()) : nothing

using IncrementalInference
using HDF5, JLD, Gadfly, Colors, Cairo
using KernelDensityEstimate, Distributions
using Caesar, RoME
using RoMEPlotting
using DocStringExtensions
using ArgParse

## load all the model data
# d = odometry information
# f = laser scanner detections
# MM = multi-modal individual id references
# MMr = reworked to map to only one previous feature
# examplefolder, datafolder

include(joinpath(Pkg.dir("Caesar"),"examples","wheeled","loadVicPrkData.jl"))

# include(joinpath(ENV["HOME"],"Documents","blandauthlocal.jl"))
# backend_config, user_config = standardcloudgraphsetup(addrdict=addrdict)
# # Start new session
# addrdict["sessionId"] = "VICPRK_VID"
# addrdict["robotId"] = "Ute"
# fg = initfg(sessionname=user_config["sessionId"], robotname=user_config["robotId"], cloudgraph=backend_config)
# deleteServerSession!(fg.cg, user_config["sessionId"])

include(joinpath(Pkg.dir("Caesar"),"examples","wheeled","VicPrk_FrontEndUtils.jl"))

# run using local inference
function mainrun(dd,ff; frames=3, coord=Coord.Cartesian(xmin=-100,xmax=100,ymin=-150,ymax=50))

    currdirtime = now()
    imgdir = joinpath(ENV["HOME"], "Pictures", "vicprkimgs", "$(currdirtime)")
    mkdir(imgdir)

    ddd = deepcopy(dd)
    fff = deepcopy(ff)

    PL = Dict{Int, Any}()

    Podo=diagm([0.5;0.5;0.005])
    lsrNoise=diagm([0.1;1.0])

    # start factor graph
    N=75
    fg = initfg()
    lmoccur = Dict{Symbol, Int}()

    # init pose
    prevn = initFactorGraph!(fg, init=ddd[1][1:3])[1]

    findAddNewLandmarks!(fg, lmoccur, fff[1])
    # there might not be any landmarks yet
    findAddLandmFactorsByPose!(fg, :x0, fff[1], N=N)
    tree = wipeBuildNewTree!(fg, drawpdf=true)
    inferOverTree!(fg, tree, N=N)
    PL[1] = drawPoses(fg)
    Gadfly.draw(PDF(imgdir*"/x1.pdf", 30cm,20cm),PL[1])
    Gadfly.draw(PNG(imgdir*"/x1.png", 30cm,20cm),PL[1])

    for i in 2:frames
      PL[i] = donextframe!(fg, i, ddd, fff, lmoccur, Podo, N=N)
      PL[i].coord = coord
      Gadfly.draw(PDF(imgdir*"/x$(i).pdf", 30cm,20cm),PL[i])
      Gadfly.draw(PNG(imgdir*"/x$(i).png", 30cm,20cm),PL[i])
    end
    PL
end

# @async run('evince bt.pdf')
PL = mainrun(d,f; frames=40)

#


module HistoryLogger

using Printf
using RecipesBase
using Statistics
using POLO
using POLO: AbstractLogger
using POLO.Utility: LogLevel, Value, Decision, Gradient, Full, SimpleLogger
import POLO.log

struct History{L <: POLO.AbstractLogger} <: AbstractLogger
    fhist::Vector{Float64}
    inner::L

    function (::Type{History})()
        return new{SimpleLogger.Logger{None}()}(Vector{Float64}(), SimpleLogger.None())
    end

    function (::Type{History})(::Type{L}) where L <: LogLevel
        return new{SimpleLogger.Logger{L}}(Vector{Float64}(), SimpleLogger.Logger(L))
    end

    function (::Type{History})(logger::AbstractLogger)
        return new{typeof(logger)}(Vector{Float64}(), logger)
    end
end

function log(logger::History,k::Integer,fval::AbstractFloat,x::AbstractVector,g::AbstractVector)
    push!(logger.fhist,fval)
    log(logger.inner,k,fval,x,g)
end

@recipe function f(logger::History)
    length(logger.fhist) > 0 || error("No solution data. Has solver been run?")
    fmin = minimum(logger.fhist)
    fmax = maximum(logger.fhist)
    increment = std(logger.fhist)

    linewidth --> 4
    linecolor --> :black
    tickfontsize := 14
    tickfontfamily := "sans-serif"
    guidefontsize := 16
    guidefontfamily := "sans-serif"
    titlefontsize := 22
    titlefontfamily := "sans-serif"
    xlabel := "k"
    ylabel := "f"
    ylims --> (fmin-increment,fmax+increment)
    xlims --> (1,length(logger.fhist)+1)
    xticks --> 1:5:length(logger.fhist)
    yticks --> fmin:increment:fmax
    xformatter := (d) -> @sprintf("%.1f",d)
    yformatter := (d) -> begin
        if abs(d) <= sqrt(eps())
            "0.0"
        elseif (log10(abs(d)) < -2.0 || log10(abs(d)) > 3.0)
            @sprintf("%.4e",d)
        elseif log10(abs(d)) > 2.0
            @sprintf("%.1f",d)
        else
            @sprintf("%.2f",d)
        end
    end

    @series begin
        seriescolor --> :black
        1:1:length(logger.fhist),logger.fhist
    end
end

end


#             funcname: the name of the function to become wrapped
#             argType:  the name of the type of argument
#             cvtType:  the name of the type of the conversion target

macro wrapUnaryFunction(fname, argType, cvtType)
    @eval begin
      function ($fname)(x::($argType))
          ($fname)(convert(($cvtType), x))
      end
    end
end

macro wrapUnaryFunctionFloat65(fname, argType)
    @eval begin
      function ($fname)(a::($argType))
          value = (argType)( ($fname)(reflect(a.fp)) )
          if getstate(a)
             value = setstate(value)
          end
          value
      end
    end
end

macro wrapBinaryFunctionFloat65(fname, argType, arg2Type)
    @eval begin
      function ($fname)(a::($argType), b::($arg2Type))
          value = (argType)( ($fname)(reflect(a.fp),reflect(b.fp)) )
          if getstate(a)|getstate(b)
             value = setstate(value)
          end
          value
      end
    end
end


export PolytopeVertexMesh
export name, vertices, polytopes, materials, material_names, groups, vtk_type

struct PolytopeVertexMesh{Dim,T,P<:Polytope} <: AbstractMesh
    vertices::Vector{Point{Dim,T}}
    polytopes::Vector{P}
    materials::Vector{UInt8}
    material_names::Vector{String}
    name::String
    groups::Dict{String,BitSet}
end

name(mesh::PolytopeVertexMesh) = mesh.name
points(mesh::PolytopeVertexMesh) = mesh.vertices
vertices(mesh::PolytopeVertexMesh) = mesh.vertices
polytopes(mesh::PolytopeVertexMesh) = mesh.polytopes
materials(mesh::PolytopeVertexMesh) = mesh.materials
material_names(mesh::PolytopeVertexMesh) = mesh.material_names
groups(mesh::PolytopeVertexMesh) = mesh.groups

# constructors
function PolytopeVertexMesh(vertices::Vector{Point{Dim,T}},
                            polytopes::Vector{P},
                            name::String
                           ) where {Dim,T,P}
    return PolytopeVertexMesh(vertices, polytopes, name, Dict{String,BitSet}())
end

function PolytopeVertexMesh(vertices::Vector{Point{Dim,T}},
                            polytopes::Vector{P}
                           ) where {Dim,T,P}
    return PolytopeVertexMesh(vertices, polytopes, "", Dict{String,BitSet}())
end

function PolytopeVertexMesh(mesh::VolumeMesh{Dim,T,U}) where {Dim,T,U}
    return PolytopeVertexMesh(
             mesh.points,
             map(i->_materialize_face_connectivity(i, mesh), eachindex(mesh.types)),
             mesh.materials,
             mesh.material_names,
             mesh.name,
             mesh.groups
            )
end

vtk_type(::Type{<:Triangle})               = VTK_TRIANGLE
vtk_type(::Type{<:Quadrilateral})          = VTK_QUAD
vtk_type(::Type{<:QuadraticTriangle})      = VTK_QUADRATIC_TRIANGLE
vtk_type(::Type{<:QuadraticQuadrilateral}) = VTK_QUADRATIC_QUAD
vtk_type(::Type{<:Tetrahedron})            = VTK_TETRA
vtk_type(::Type{<:Hexahedron})             = VTK_HEXAHEDRON
vtk_type(::Type{<:QuadraticTetrahedron})   = VTK_QUADRATIC_TETRA
vtk_type(::Type{<:QuadraticHexahedron})    = VTK_QUADRATIC_HEXAHEDRON

function Base.show(io::IO, mesh::PolytopeVertexMesh{Dim,T,P}) where {Dim,T,P}
    print(io, "PolytopeVertexMesh{",Dim,", ",T,", ")
    if isconcretetype(P)
        println(io, alias_string(P),"{",vertex_type(P),"}}")
    else
        println(io, P,"}")
    end
    println(io, "  ├─ Name      : ", mesh.name)
    size_B = Base.summarysize(mesh)
    if size_B < 1e6
        println(io, "  ├─ Size (KB) : ", string(@sprintf("%.3f", size_B/1000)))
    else
        println(io, "  ├─ Size (MB) : ", string(@sprintf("%.3f",size_B/1e6)))
    end
    println(io, "  ├─ Vertices  : ", length(mesh.vertices))
    println(io, "  ├─ Polytopes : ", length(mesh.polytopes))
    poly_types = unique!(map(typeof, mesh.polytopes))
    npoly_types = length(poly_types)
    for i = 1:npoly_types
        poly_type = poly_types[i]
        npoly = count(x->x isa poly_type,  mesh.polytopes)
        if i === npoly_types
            println(io, "  │  └─ ", rpad(alias_string(poly_type), 22), ": ", npoly)
        else
            println(io, "  │  ├─ ", rpad(alias_string(poly_type), 22), ": ", npoly)
        end
    end
    println(io, "  ├─ Materials : ", length(mesh.material_names))
    ngroups = length(mesh.groups)
    println(io, "  └─ Groups    : ", ngroups)
    if 0 < ngroups ≤ 5
        group_keys = sort(collect(keys(mesh.groups)))
        for i = 1:ngroups
            if i === ngroups
                println(io, "     └─ ", group_keys[i])
            else
                println(io, "     ├─ ", group_keys[i])
            end
        end
    end
end


@testset "Alphabet" begin

    @testset "Creation, Iteration and getindex" begin

        # Creation & Iteration
        @test [ i for i in UngappedAlphabet() ] == Int[ i for i in 1:20 ]
        @test [ i for i in GappedAlphabet() ]   == Int[ i for i in 1:21 ]
        @test [ i for i in ReducedAlphabet("(AILMV)(NQST)(RHK)(DE)(FWY)CGP") ]  ==
           Int[ i for i in 1:8 ]
    end

    @testset "getindex and names" begin

        reduced = ReducedAlphabet("(AILMV)(NQST)(RHK)(DE)(FWY)CGP")
        strings = ["AILMV", "NQST", "RHK", "DE", "FWY", "C", "G", "P"]

        @test length(reduced) == length(strings)
        @test names(reduced) == strings

        for i in 1:length(reduced)
            @test reduced[strings[i]] == i
        end

        strings_gapped = ["A","R","N","D","C","Q","E","G","H","I","L","K","M","F","P","S",
            "T","W","Y","V","-"]

        @test names(GappedAlphabet()) == strings_gapped
        @test names(UngappedAlphabet()) == strings_gapped[1:end-1]

        for i in 1:20
            @test UngappedAlphabet()[strings_gapped[i]] == i
            @test GappedAlphabet()[strings_gapped[i]] == i
        end

        @test GappedAlphabet()["-"] == 21
        @test UngappedAlphabet()["-"] == 22
    end

    @testset "in" begin

        # Creation & Iteration
        @test in(Residue('A'), UngappedAlphabet())
        @test in(Residue('A'), GappedAlphabet())
        @test in(Residue('A'), ReducedAlphabet("(AILMV)(NQST)(RHK)(DE)(FWY)CGP"))

        @test !in(GAP, UngappedAlphabet())
        @test  in(GAP, GappedAlphabet())
        @test !in(GAP, ReducedAlphabet("(AILMV)(NQST)(RHK)(DE)(FWY)CGP"))

        @test !in(XAA, UngappedAlphabet())
        @test !in(XAA, GappedAlphabet())
        @test !in(XAA, ReducedAlphabet("(AILMV)(NQST)(RHK)(DE)(FWY)CGP"))
    end

    @testset "Show" begin

        tmp = IOBuffer()
        show(tmp, UngappedAlphabet())
        @test String(take!(tmp)) ==
        "UngappedAlphabet of length 20. Residues : res\"ARNDCQEGHILKMFPSTWYV\""

        show(tmp, GappedAlphabet())
        @test String(take!(tmp)) ==
        "GappedAlphabet of length 21. Residues : res\"ARNDCQEGHILKMFPSTWYV-\""

        show(tmp, ReducedAlphabet("(AILMV)(NQST)(RHK)(DE)(FWY)CGP"))
        @test String(take!(tmp)) ==
        "ReducedAlphabet of length 8 : \"(AILMV)(NQST)(RHK)(DE)(FWY)CGP\""
    end
end


# ## Ghost
set_camera!(vis, cam_pos=[-1,1,0], zoom=1)

z_sim = get_maximal_state(storage)
timesteps = [5, 10, 15]# .+ 150

for t in timesteps
    name = Symbol("robot_$t")
    build_robot(mech, vis=vis, name=name, color= magenta_light)
    z = z_sim[t]
    set_robot(vis, mech, z, name=name)
end

z_sim[220][1] -= 0.1
z_sim[220][13 + 1] -= 0.05
set_robot(vis, mech, z_sim[220])

z_sim[104][1] -= 0.025
z_sim[104][13 + 1] -= 0.05
set_robot(vis, mech, z_sim[104])


module Networks

using Printf, ..Lif, ..LifConfig

include("./Synapses.jl")
using .Synapses

export LIFNetwork, simulate, reset!, connect!

# by exporting the Synapses module, where 'using' is used to import
# this module, Synapses will be imported as if 'import' had been used
export Synapses
# ============================================================================ #
abstract type NeuralNetwork end

# ============================================================================ #
# LIFNetwork
# ============================================================================ #
mutable struct LIFNetwork{S<:BaseSynapse, N<:BaseNeuron} <: NeuralNetwork
    neurons::Array{N,1}
    synapses::Array{S,2}
end
# ---------------------------------------------------------------------------- #
function LIFNetwork(x::Array{T,1}) where {T<:BaseNeuron}
    n = length(x)
    LIFNetwork(x, Array{BaseSynapse,2}(undef, n, n))
end
# ---------------------------------------------------------------------------- #
function LIFNetwork(n::Integer, xi::Float64=0.2, typ::Type{T}=LIFNeuron) where {T<:BaseNeuron}
    x = Array{T,1}(undef, n)
    for k in eachindex(x)
        x[k] = typ()
        Lif.set!(x[k], :xi, xi)
    end
    LIFNetwork(x, Array{BaseSynapse,2}(undef, n, n))
end
# ---------------------------------------------------------------------------- #
function connect!(net::LIFNetwork)
    for k in 1:length(net)^2
        net.synapses[k] = Synapse()
    end
end
# # ---------------------------------------------------------------------------- #
# function connect!(net::LIFNetwork, args...)
#     d = Dict(args)
#     idx = keys(d)
#     n = length(net)
#
#     for r = 1:n
#         for c = 1:n
#             if (r,c) in idx
#                 net.synapses[r,c] = Synapse(d[(r,c)]...)
#             elseif !isassigned(net.synapses, r, c)
#                 net.synapses[r,c] = Synapse()
#             else
#                 @show(net.synapses[r,c])
#             end
#         end
#     end
# end
# ---------------------------------------------------------------------------- #
function connect!(net::LIFNetwork, k::Tuple{Integer, Integer}, p::Tuple)
    net.synapses[k...] = Synapse(p...)
end
# ---------------------------------------------------------------------------- #
Base.iterate(net::LIFNetwork) = (net[1], 1)
function Base.iterate(net::LIFNetwork, state::Integer)
    if state <= length(net)
        return (net[state], state+1)
    else
        return nothing
    end
end
Base.IteratorSize(net::LIFNetwork) = HasLength()
Base.IteratorEltype(net::LIFNetwork) = HasEltype()
Base.eltype(net::LIFNetwork) = eltype(net.neurons)
Base.length(net::LIFNetwork) = length(net.neurons)

Base.getindex(net::LIFNetwork, k::Integer) = net.neurons[k]
Base.getindex(net::LIFNetwork, k::UnitRange) = net.neurons[k]
Base.getindex(net::LIFNetwork, k::AbstractArray) = net.neurons[k]
# ---------------------------------------------------------------------------- #
function Base.show(io::IO, net::LIFNetwork)
    str = @sprintf("%dx1 Array{LIFNeuron,1}:\n", length(net.neurons))
    print(io, str)
    print(io, net.neurons)
    print(io, "\n\n")
    print(io, net.synapses)
end
# ---------------------------------------------------------------------------- #
function reset!(net::LIFNetwork)
    for x in net.neurons
        #NOTE: not sure why we have to specify the module here...
        Lif.reset!(x)
    end
    for s in net.synapses
        #NOTE: not sure why we have to specify the module here...
        Synapses.reset!(s)
    end
end
# ---------------------------------------------------------------------------- #
function collect_input(net::LIFNetwork, k::Integer, dt::Time, tnow::Time)
    totl = 0.0
    #input is received column-wise
    @inbounds for s in net.synapses[:,k]
        if isconnected(s)
            totl += receive!(net[k], s, dt, tnow)
        end
    end
    return totl
end
# ---------------------------------------------------------------------------- #
function distribute_output(net::LIFNetwork, k::Integer, spk::Bool, tnow::Time)
    #output is distributed row-wise
    @inbounds for s in net.synapses[k,:]
        send!(net[k], s, spk, tnow)
    end
end
# ---------------------------------------------------------------------------- #
# function update!(net::LIFNetwork, stim::Vector{Float64}, dt::Time, tnow::Time)
function update!(net::LIFNetwork, fstim::Function, dt::Time, tnow::Time)
    vm = zeros(length(net),1)
    spk = falses(length(net),1)

    @inbounds for k = 1:length(net)

        #get synaptic and stimulus input
        inp = collect_input(net, k, dt, tnow) + fstim(k, tnow)

        # update Vm
        vm[k], spk[k] = step!(net[k], inp, dt, tnow)

        #propagate output to synapses
        distribute_output(net, k, spk[k], tnow)

        if spk[k]
            vm[k] = net[k].p.vspike
        end
    end

    return vm, spk
end
# ============================================================================ #
# helper functions
# ============================================================================ #
# include("NetworksExt.jl")
# ============================================================================ #
function default_stim(id::Integer, t::Time)
    return 0.0
end
# ============================================================================ #
function simulate(net::T, fstim::Function, duration::Real=500.0,
    record::Vector{Int64}=Int64[]) where {T<:NeuralNetwork}

    tnow = 0.0
    dt = 0.043
    record_vm = false

    reset!(net)

    d = Dict{Symbol, AbstractArray}()

    times = 0.0:dt:duration
    if !isa(record, Array)
        record = [record]
    end
    if !isempty(record) && eltype(record) <: Integer
        filter!(x -> (x > 0 && x <= length(net)), record)
        d[:vm] = zeros(length(times),length(record))
        record_vm = true
    end
    d[:ts] = [Vector{Time}() for k in 1:length(net)]

    for k in eachindex(times)

        vm, spk = update!(net, fstim, dt, times[k])
        if record_vm
            d[:vm][k,:] = vm[record]
        end

        for j in 1:length(net)
            if spk[j]
                #NOTE: timestamps are stored in *SECONDS*
                push!(d[:ts][j], times[k]*1e-3)
            end
        end
    end

    d[:time] = times

    return d
end
# ============================================================================ #
# function test(n=2, xi=0.5)
#
#     net = LIFNetwork(n,xi)
#     for x in net
#         Lif.set!(x, :xi, 0.0)
#     end
#     connect!(net, (1,2)=>(Synapses.Static, .54, 3.0, 2.0))
#     stim(id, t) = t > 100.0 && id == 1 ? 0.2 : 0.0
#
#     d = simulate(net, fstim=stim, duration=500, record=[1])
#     return d
# end
# # ============================================================================ #

end #END MODULE


@testset "training.jl inside the solver (with backtracking)" begin

    println("Training inside the solver -------- ")

    generator = SeaPearl.HomogenousGraphColoringGenerator(10, 0.1)

    numInFeatures = 3
    # Model definition
    approximator_GNN = SeaPearl.GraphConv(64 => 64, Flux.leakyrelu)
    target_approximator_GNN = SeaPearl.GraphConv(64 => 64, Flux.leakyrelu)
    gnnlayers = 10

    approximator_model = SeaPearl.CPNN(
        graphChain = Flux.Chain(
            SeaPearl.GraphConv(numInFeatures => 64, Flux.leakyrelu),
            [approximator_GNN for i = 1:gnnlayers]...
        ),
        nodeChain = Flux.Chain(
            Flux.Dense(64, 32, Flux.leakyrelu),
            Flux.Dense(32, 32, Flux.leakyrelu),
            Flux.Dense(32, 16, Flux.leakyrelu),
        ),
        outputChain = Flux.Dense(16, generator.nb_nodes),
    )
    target_approximator_model = SeaPearl.CPNN(
        graphChain = Flux.Chain(
            SeaPearl.GraphConv(numInFeatures => 64, Flux.leakyrelu),
            [target_approximator_GNN for i = 1:gnnlayers]...
        ),
        nodeChain = Flux.Chain(
            Flux.Dense(64, 32, Flux.leakyrelu),
            Flux.Dense(32, 32, Flux.leakyrelu),
            Flux.Dense(32, 16, Flux.leakyrelu),
        ),
        outputChain = Flux.Dense(16, generator.nb_nodes),
    )

    # Agent definition
    agent = RL.Agent(
        policy = RL.QBasedPolicy(
            learner = RL.DQNLearner(
                approximator = RL.NeuralNetworkApproximator(
                    model = approximator_model,
                    optimizer = ADAM(0.0005f0)
                ),
                target_approximator = RL.NeuralNetworkApproximator(
                    model = target_approximator_model,
                    optimizer = ADAM(0.0005f0)
                ),
                loss_func = Flux.Losses.huber_loss,
                stack_size = nothing,
                γ = 0.99f0,
                batch_size = 32,
                update_horizon = 4,
                min_replay_history = 32,
                update_freq = 4,
                target_update_freq = 20,
            ),
            explorer = RL.EpsilonGreedyExplorer(
                ϵ_stable = 0.001,
                kind = :exp,
                ϵ_init = 1.0,
                warmup_steps = 0,
                decay_steps = 5000,
                step = 1,
                is_break_tie = false,
                #is_training = true,
                rng = MersenneTwister(33)
            )
        ),
        trajectory = RL.CircularArraySARTTrajectory(
            capacity = 800,
            state = SeaPearl.DefaultTrajectoryState[] => (),
        )
    )

    learnedHeuristic = SeaPearl.LearnedHeuristic(agent)

    initial_params = deepcopy(params(learnedHeuristic.agent.policy.learner.approximator.model))

    metrics,evalmetrics = SeaPearl.train!(
        valueSelectionArray=learnedHeuristic,
        generator=generator,
        nbEpisodes= 10, #3,
        strategy=SeaPearl.DFSearch(),
        variableHeuristic=SeaPearl.MinDomainVariableSelection{true}(),
        out_solver = false,
        verbose = false,
        evaluator=SeaPearl.SameInstancesEvaluator([learnedHeuristic],generator)  #need to be a vector of Heuristic
    )

    final_params = params(learnedHeuristic.agent.policy.learner.approximator.model)

    @test final_params != initial_params

end

@testset "training.jl outside the solver (without backtracking)" begin

    println("Let's try without backtracking   :   ------------")

    generator = SeaPearl.HomogenousGraphColoringGenerator(10, 0.1)

    numInFeatures = 3

    # Model definition
    approximator_GNN = SeaPearl.GraphConv(64 => 64, Flux.leakyrelu)
    target_approximator_GNN = SeaPearl.GraphConv(64 => 64, Flux.leakyrelu)
    gnnlayers = 10

    approximator_model = SeaPearl.CPNN(
        graphChain = Flux.Chain(
            SeaPearl.GraphConv(numInFeatures => 64, Flux.leakyrelu),
            [approximator_GNN for i = 1:gnnlayers]...
        ),
        nodeChain = Flux.Chain(
            Flux.Dense(64, 32, Flux.leakyrelu),
            Flux.Dense(32, 32, Flux.leakyrelu),
            Flux.Dense(32, 16, Flux.leakyrelu),
        ),
        outputChain = Flux.Dense(16, generator.nb_nodes),
    )
    target_approximator_model = SeaPearl.CPNN(
        graphChain = Flux.Chain(
            SeaPearl.GraphConv(numInFeatures => 64, Flux.leakyrelu),
            [target_approximator_GNN for i = 1:gnnlayers]...
        ),
        nodeChain = Flux.Chain(
            Flux.Dense(64, 32, Flux.leakyrelu),
            Flux.Dense(32, 32, Flux.leakyrelu),
            Flux.Dense(32, 16, Flux.leakyrelu),
        ),
        outputChain = Flux.Dense(16, generator.nb_nodes),
    )

    # Agent definition
    agent = RL.Agent(
        policy = RL.QBasedPolicy(
            learner = RL.DQNLearner(
                approximator = RL.NeuralNetworkApproximator(
                    model = approximator_model,
                    optimizer = ADAM(0.0005f0)
                ),
                target_approximator = RL.NeuralNetworkApproximator(
                    model = target_approximator_model,
                    optimizer = ADAM(0.0005f0)
                ),
                loss_func = Flux.Losses.huber_loss,
                stack_size = nothing,
                γ = 0.99f0,
                batch_size = 32,
                update_horizon = 4,
                min_replay_history = 32,
                update_freq = 4,
                target_update_freq = 20,
            ),
            explorer = RL.EpsilonGreedyExplorer(
                ϵ_stable = 0.001,
                kind = :exp,
                ϵ_init = 1.0,
                warmup_steps = 0,
                decay_steps = 5000,
                step = 1,
                is_break_tie = false,
                #is_training = true,
                rng = MersenneTwister(33)
            )
        ),
        trajectory = RL.CircularArraySLARTTrajectory(
            capacity = 800,
            state = SeaPearl.DefaultTrajectoryState[] => (),
            legal_actions_mask = Vector{Bool} => (generator.nb_nodes, ),

        )
    )

    learnedHeuristic = SeaPearl.LearnedHeuristic(agent)
    initial_params = deepcopy(params(learnedHeuristic.agent.policy.learner.approximator.model))

    metrics,evalmetrics = SeaPearl.train!(
        valueSelectionArray=learnedHeuristic,
        generator=generator,
        nbEpisodes=10,
        strategy=SeaPearl.DFSearch(),
        variableHeuristic=SeaPearl.MinDomainVariableSelection{true}(),
        out_solver = true,
        verbose = false,
        evaluator=nothing
    )

    final_params = params(learnedHeuristic.agent.policy.learner.approximator.model)

    @test final_params != initial_params

end


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Bus
        zone::Int
        name::String
        status::Int
        area::Int
        vm_ub::Float64
        vm_lb::Float64
        index::Int
        va::Float64
        vm::Float64
        base_kv::Float64
        type::Int
    end

# Arguments
- `zone::Int`: assigned control zone
- `name::String`: a flexible name for components, non required to be unique
- `status::Int`: a 0/1 value indicating if the component should be omitted or not (0 => omitted)
- `area::Int`: assigned control area
- `vm_ub::Float64`: an upper limit on voltage magnitude
- `vm_lb::Float64`: a lower limit on voltage magnitude
- `index::Int`: a unique identifier for components
- `va::Float64`: voltage angle
- `vm::Float64`: voltage magnitude
- `base_kv::Float64`: base voltage
- `type::Int`: bus type for power flow (1=PQ, 2=PV, 3=Slack)
"""
mutable struct Bus
    "assigned control zone"
    zone::Int
    "a flexible name for components, non required to be unique"
    name::String
    "a 0/1 value indicating if the component should be omitted or not (0 => omitted)"
    status::Int
    "assigned control area"
    area::Int
    "an upper limit on voltage magnitude"
    vm_ub::Float64
    "a lower limit on voltage magnitude"
    vm_lb::Float64
    "a unique identifier for components"
    index::Int
    "voltage angle"
    va::Float64
    "voltage magnitude"
    vm::Float64
    "base voltage"
    base_kv::Float64
    "bus type for power flow (1=PQ, 2=PV, 3=Slack)"
    type::Int
end

function Bus(; zone, name, status, area, vm_ub, vm_lb, index, va, vm, base_kv, type, )
    Bus(zone, name, status, area, vm_ub, vm_lb, index, va, vm, base_kv, type, )
end

# Constructor for demo purposes; non-functional.
function Bus(::Nothing)
    Bus(;
        zone=1,
        name="1",
        status=1,
        area=1,
        vm_ub=1.2,
        vm_lb=1.2,
        index=1,
        va=1.3,
        vm=1.3,
        base_kv=1.2,
        type=1,
    )
end

"""Get Bus zone."""
get_zone(value::Bus) = value.zone
"""Get Bus name."""
get_name(value::Bus) = value.name
"""Get Bus status."""
get_status(value::Bus) = value.status
"""Get Bus area."""
get_area(value::Bus) = value.area
"""Get Bus vm_ub."""
get_vm_ub(value::Bus) = value.vm_ub
"""Get Bus vm_lb."""
get_vm_lb(value::Bus) = value.vm_lb
"""Get Bus index."""
get_index(value::Bus) = value.index
"""Get Bus va."""
get_va(value::Bus) = value.va
"""Get Bus vm."""
get_vm(value::Bus) = value.vm
"""Get Bus base_kv."""
get_base_kv(value::Bus) = value.base_kv
"""Get Bus type."""
get_type(value::Bus) = value.type

"""Set Bus zone."""
set_zone!(value::Bus, val::Int) = value.zone = val
"""Set Bus name."""
set_name!(value::Bus, val::String) = value.name = val
"""Set Bus status."""
set_status!(value::Bus, val::Int) = value.status = val
"""Set Bus area."""
set_area!(value::Bus, val::Int) = value.area = val
"""Set Bus vm_ub."""
set_vm_ub!(value::Bus, val::Float64) = value.vm_ub = val
"""Set Bus vm_lb."""
set_vm_lb!(value::Bus, val::Float64) = value.vm_lb = val
"""Set Bus index."""
set_index!(value::Bus, val::Int) = value.index = val
"""Set Bus va."""
set_va!(value::Bus, val::Float64) = value.va = val
"""Set Bus vm."""
set_vm!(value::Bus, val::Float64) = value.vm = val
"""Set Bus base_kv."""
set_base_kv!(value::Bus, val::Float64) = value.base_kv = val
"""Set Bus type."""
set_type!(value::Bus, val::Int) = value.type = val


mutable struct Simulation{T <: SimulationModel}
    status::BUILD_STATUS
    problem::Union{Nothing, SciMLBase.DEProblem}
    tspan::NTuple{2, Float64}
    sys::PSY.System
    perturbations::Vector{<:Perturbation}
    x0_init::Vector{Float64}
    initialized::Bool
    tstops::Vector{Float64}
    callbacks::DiffEqBase.CallbackSet
    simulation_folder::String
    inputs::Union{Nothing, SimulationInputs}
    results::Union{Nothing, SimulationResults}
    console_level::Base.CoreLogging.LogLevel
    file_level::Base.CoreLogging.LogLevel
    multimachine::Bool
end

get_system(sim::Simulation) = sim.sys
get_simulation_inputs(sim::Simulation) = sim.inputs
get_initial_conditions(sim::Simulation) = sim.x0_init
get_tspan(sim::Simulation) = sim.tspan

function Simulation(
    ::Type{T},
    sys::PSY.System;
    tspan,
    initial_conditions,
    initialize_simulation,
    perturbations,
    simulation_folder,
    console_level,
    file_level,
) where {T <: SimulationModel}
    PSY.set_units_base_system!(sys, "DEVICE_BASE")

    return Simulation{T}(
        BUILD_INCOMPLETE,
        nothing,
        tspan,
        sys,
        perturbations,
        initial_conditions,
        !initialize_simulation,
        Vector{Float64}(),
        DiffEqBase.CallbackSet(),
        simulation_folder,
        nothing,
        nothing,
        console_level,
        file_level,
        false,
    )
end

function Simulation!(
    ::Type{T},
    system::PSY.System,
    simulation_folder::String,
    tspan::NTuple{2, Float64},
    perturbation::Perturbation;
    kwargs...,
) where {T <: SimulationModel}
    return Simulation!(T, system, simulation_folder, tspan, [perturbation]; kwargs...)
end

function Simulation(
    ::Type{T},
    system::PSY.System,
    simulation_folder::String,
    tspan::NTuple{2, Float64},
    perturbation::Perturbation;
    kwargs...,
) where {T <: SimulationModel}
    return Simulation(T, system, simulation_folder, tspan, [perturbation]; kwargs...)
end

"""
Builds the simulation object and conducts the indexing process. The initial conditions are stored in the system.

# Accepted Key Words
- `initialize_simulation::Bool : Runs the initialization routine. If false, simulation runs based on the operation point stored in System`
- `initial_conditions::Vector{Float64} : Allows the user to pass a vector with the initial condition values desired in the simulation. If initialize_simulation = true, these values are used as a first guess and overwritten.
- `system_to_file::Bool`: Serializes the initialized system
- `console_level::Logging`: Sets the level of logging output to the console. Can be set to Logging.Error, Logging.Warn, Logging.Info or Logging.Debug
- `file_level::Logging`: Sets the level of logging output to file. Can be set to Logging.Error, Logging.Warn, Logging.Info or Logging.Debug
"""
function Simulation!(
    ::Type{T},
    system::PSY.System,
    simulation_folder::String,
    tspan::NTuple{2, Float64},
    perturbations::Vector{<:Perturbation} = Vector{Perturbation}();
    kwargs...,
) where {T <: SimulationModel}
    check_folder(simulation_folder)
    # Instantiates the Simulation object
    sim = Simulation(
        T,
        system;
        tspan = tspan,
        initial_conditions = get(kwargs, :initial_conditions, Vector{Float64}()),
        initialize_simulation = get(kwargs, :initialize_simulation, true),
        simulation_folder = simulation_folder,
        perturbations = perturbations,
        console_level = get(kwargs, :console_level, Logging.Warn),
        file_level = get(kwargs, :file_level, Logging.Debug),
    )

    build!(sim; kwargs...)
    if get(kwargs, :system_to_file, false)
        PSY.to_json(system, joinpath(simulation_folder, "initialized_system.json"))
    end
    return sim
end

"""
Initializes the simulations and builds the indexing. The input system is not modified during the initialization

# Accepted Key Words
- `initialize_simulation::Bool : Runs the initialization routine. If false, simulation runs based on the operation point stored in System`
- `initial_conditions::Vector{Float64} : Initial conditions provided by the user. Must be used with initialize_simulation = false to not be override.`
- `system_to_file::Bool`: Serializes the original input system
- `console_level::Logging`: Sets the level of logging output to the console. Can be set to Logging.Error, Logging.Warn, Logging.Info or Logging.Debugg
- `file_level::Logging`: Sets the level of logging output to file. Can be set to Logging.Error, Logging.Warn, Logging.Info or Logging.Debug
"""
function Simulation(
    ::Type{T},
    system::PSY.System,
    simulation_folder::String,
    tspan::NTuple{2, Float64},
    perturbations::Vector{<:Perturbation} = Vector{Perturbation}();
    kwargs...,
) where {T <: SimulationModel}
    check_folder(simulation_folder)
    simulation_system = deepcopy(system)
    # Instantiates the Simulation object
    sim = Simulation(
        T,
        simulation_system;
        tspan = tspan,
        initial_conditions = get(kwargs, :initial_conditions, Vector{Float64}()),
        initialize_simulation = get(kwargs, :initialize_simulation, true),
        simulation_folder = simulation_folder,
        perturbations = perturbations,
        console_level = get(kwargs, :console_level, Logging.Warn),
        file_level = get(kwargs, :file_level, Logging.Debug),
    )
    build!(sim; kwargs...)
    if get(kwargs, :system_to_file, false)
        PSY.to_json(system, joinpath(simulation_folder, "input_system.json"))
    end
    return sim
end

function reset!(sim::Simulation{T}) where {T <: SimulationModel}
    @info "Rebuilding the simulation after reset"
    sim.inputs = SimulationInputs(T(), get_system(sim), sim.inputs.tspan)
    build!(sim)
    @info "Simulation reset to status $(sim.status)"
    return
end

function configure_logging(sim::Simulation, file_mode; kwargs...)
    return IS.configure_logging(
        console = true,
        console_stream = stderr,
        console_level = get(kwargs, :console_level, sim.console_level),
        file = true,
        filename = joinpath(sim.simulation_folder, SIMULATION_LOG_FILENAME),
        file_level = get(kwargs, :file_level, sim.file_level),
        file_mode = file_mode,
        tracker = nothing,
        set_global = false,
    )
end

function _build_inputs!(
    sim::Simulation{T},
    frequency_reference,
) where {T <: SimulationModel}
    simulation_system = get_system(sim)
    sim.inputs = SimulationInputs(T, simulation_system, frequency_reference)
    @debug "Simulation Inputs Created"
    return
end

function _get_flat_start(inputs::SimulationInputs)
    bus_count = get_bus_count(inputs)
    var_count = get_variable_count(inputs)
    initial_conditions = zeros(var_count)
    initial_conditions[1:bus_count] .= 1.0
    return initial_conditions
end

function _initialize_state_space(sim::Simulation{T}) where {T <: SimulationModel}
    simulation_inputs = get_simulation_inputs(sim)
    x0_init = sim.x0_init
    if isempty(x0_init) && sim.initialized
        @warn "Initial Conditions set to flat start"
        sim.x0_init = _get_flat_start(simulation_inputs)
    elseif isempty(x0_init) && !sim.initialized
        sim.x0_init = _get_flat_start(simulation_inputs)
    elseif !isempty(x0_init) && sim.initialized
        if length(sim.x0_init) != get_variable_count(simulation_inputs)
            throw(
                IS.ConflictingInputsError(
                    "The size of the provided initial state space does not match the model's state space.",
                ),
            )
        end
    elseif !isempty(x0_init) && !sim.initialized
        @warn "initial_conditions were provided with initialize_simulation. User's initial_conditions will be overwritten."
        if length(sim.x0_init) != get_variable_count(simulation_inputs)
            sim.x0_init = _get_flat_start(simulation_inputs)
        end
    else
        @assert false
    end
    return
end

function _pre_initialize_simulation!(sim::Simulation)
    _initialize_state_space(sim)
    if sim.initialized != true
        @info("Pre-Initializing Simulation States")
        sim.initialized = precalculate_initial_conditions!(sim)
        if !sim.initialized
            error(
                "The simulation failed to find an adequate initial guess for the initialization. Check the intialization routine.",
            )
        end
    else
        @warn(
            "No Pre-initialization conducted. If this is unexpected, check the initialization keywords"
        )
        sim.status = SIMULATION_INITIALIZED
    end
    return
end

function _get_jacobian(sim::Simulation{T}) where {T <: SimulationModel}
    inputs = get_simulation_inputs(sim)
    x0_init = get_initial_conditions(sim)
    return JacobianFunctionWrapper(T(inputs, x0_init, JacobianCache), x0_init)
end

function _build_perturbations!(sim::Simulation)
    @info "Attaching Perturbations"
    if isempty(sim.perturbations)
        @debug "The simulation has no perturbations"
        return DiffEqBase.CallbackSet(), [0.0]
    end
    inputs = get_simulation_inputs(sim)
    perturbations = sim.perturbations
    perturbations_count = length(perturbations)
    callback_vector = Vector{DiffEqBase.DiscreteCallback}(undef, perturbations_count)
    tstops = Vector{Float64}(undef, perturbations_count)
    for (ix, pert) in enumerate(perturbations)
        @debug pert
        condition = (x, t, integrator) -> t in [pert.time]
        affect = get_affect(inputs, get_system(sim), pert)
        callback_vector[ix] = DiffEqBase.DiscreteCallback(condition, affect)
        tstops[ix] = pert.time
    end
    sim.tstops = tstops
    sim.callbacks = DiffEqBase.CallbackSet((), tuple(callback_vector...))
    return
end

function _get_diffeq_problem(
    sim::Simulation,
    model::SystemModel{ResidualModel},
    jacobian::JacobianFunctionWrapper,
)
    x0 = get_initial_conditions(sim)
    dx0 = zeros(length(x0))
    simulation_inputs = get_simulation_inputs(sim)
    sim.problem = SciMLBase.DAEProblem(
        SciMLBase.DAEFunction{true}(
            model;
            # Currently commented for Sundials compatibility
            #jac = jacobian,
            tgrad = (dT, u, p, t) -> dT .= false,
            #jac_prototype = jacobian.Jv,
        ),
        dx0,
        x0,
        get_tspan(sim),
        simulation_inputs,
        differential_vars = get_DAE_vector(simulation_inputs),
    )
    sim.status = BUILT
    return
end

function _get_diffeq_problem(
    sim::Simulation,
    model::SystemModel{MassMatrixModel},
    jacobian::JacobianFunctionWrapper,
)
    simulation_inputs = get_simulation_inputs(sim)
    sim.problem = SciMLBase.ODEProblem(
        SciMLBase.ODEFunction{true}(
            model,
            mass_matrix = get_mass_matrix(simulation_inputs),
            jac = jacobian,
            jac_prototype = jacobian.Jv,
            # Necessary to avoid unnecessary calculations in Rosenbrock methods
            tgrad = (dT, u, p, t) -> dT .= false,
        ),
        sim.x0_init,
        get_tspan(sim),
        simulation_inputs,
    )
    sim.status = BUILT
    return
end

function _build!(sim::Simulation{T}; kwargs...) where {T <: SimulationModel}
    check_kwargs(kwargs, SIMULATION_ACCEPTED_KWARGS, "Simulation")
    # Branches are a super set of Lines. Passing both kwargs will
    # be redundant.
    TimerOutputs.reset_timer!(BUILD_TIMER)
    if get(kwargs, :disable_timer_outputs, false)
        TimerOutputs.disable_timer!(BUILD_PROBLEMS_TIMER)
    end

    TimerOutputs.@timeit BUILD_TIMER "Build Simulation" begin
        if get(kwargs, :all_branches_dynamic, false)
            TimerOutputs.@timeit BUILD_TIMER "AC Branch Transform to Dynamic" begin
                sys = get_system(sim)
                transform_branches_to_dynamic(sys, PSY.ACBranch)
            end
        elseif get(kwargs, :all_lines_dynamic, false)
            TimerOutputs.@timeit BUILD_TIMER "Line Transform to Dynamic" begin
                sys = get_system(sim)
                transform_branches_to_dynamic(sys, PSY.Line)
            end
        end
        TimerOutputs.@timeit BUILD_TIMER "Build Simulation Inputs" begin
            f_ref = get(kwargs, :frequency_reference, ReferenceBus)
            _build_inputs!(sim, f_ref)
            sim.multimachine = f_ref != ConstantFrequency
        end
        TimerOutputs.@timeit BUILD_TIMER "Pre-initialization" begin
            _pre_initialize_simulation!(sim)
        end
        if sim.status != BUILD_FAILED
            simulation_inputs = get_simulation_inputs(sim)
            try
                TimerOutputs.@timeit BUILD_TIMER "Calculate Jacobian" begin
                    jacobian = _get_jacobian(sim)
                end
                TimerOutputs.@timeit BUILD_TIMER "Make Model Function" begin
                    model = T(simulation_inputs, get_initial_conditions(sim), SimCache)
                end
                TimerOutputs.@timeit BUILD_TIMER "Initial Condition NLsolve refinement" begin
                    refine_initial_condition!(sim, model, jacobian)
                end
                TimerOutputs.@timeit BUILD_TIMER "Build Perturbations" begin
                    _build_perturbations!(sim)
                end
                TimerOutputs.@timeit BUILD_TIMER "Make DiffEq Problem" begin
                    _get_diffeq_problem(sim, model, jacobian)
                end
                @info "Simulations status = $(sim.status)"
            catch e
                bt = catch_backtrace()
                @error "$T failed to build" exception = e, bt
                sim.status = BUILD_FAILED
            end
        else
            @error "The simulation couldn't be initialized correctly. Simulations status = $(sim.status)"
        end
    end
    return
end

function build!(sim; kwargs...)
    logger = configure_logging(sim, "w")
    try
        Logging.with_logger(logger) do
            _build!(sim; kwargs...)
        end
    catch e
        @error "Build failed" exception = (e, catch_backtrace())
        return sim.status
    finally
        string_buffer = IOBuffer()
        TimerOutputs.print_timer(
            string_buffer,
            BUILD_TIMER,
            sortby = :firstexec,
            compact = true,
        )
        @info "\n$(String(take!(string_buffer)))\n"
        close(logger)
    end
    return
end

function simulation_pre_step!(sim::Simulation, reset_sim::Bool)
    reset_sim = sim.status == CONVERTED_FOR_SMALL_SIGNAL || reset_sim
    if sim.status == BUILD_FAILED
        error(
            "The Simulation status is $(sim.status). Can not continue, correct your inputs and build the simulation again.",
        )
    elseif sim.status != BUILT && !reset_sim
        error(
            "The Simulation status is $(sim.status). Use keyword argument reset_simulation = true",
        )
    end

    reset_sim && reset!(sim)
    return
end

function _prog_meter_enabled()
    return isa(stderr, Base.TTY) || (get(ENV, "CI", nothing) == "true")
end

function _execute!(sim::Simulation, solver; kwargs...)
    @debug "status before execute" sim.status
    simulation_pre_step!(sim, get(kwargs, :reset_simulation, false))
    sim.status = SIMULATION_STARTED
    time_log = Dict{Symbol, Any}()
    solution,
    time_log[:timed_solve_time],
    time_log[:solve_bytes_alloc],
    time_log[:sec_in_gc] = @timed SciMLBase.solve(
        sim.problem,
        solver;
        callback = sim.callbacks,
        tstops = sim.tstops,
        progress = get(kwargs, :enable_progress_bar, _prog_meter_enabled()),
        progress_steps = 1,
        kwargs...,
    )
    if solution.retcode == :Success
        sim.status = SIMULATION_FINALIZED
        sim.results = SimulationResults(
            get_simulation_inputs(sim),
            get_system(sim),
            time_log,
            solution,
        )
    else
        @error("The simulation failed with return code $(solution.retcode)")
        sim.status = SIMULATION_FAILED
    end
end

function execute!(sim::Simulation, solver; kwargs...)
    logger = configure_logging(sim, "a"; kwargs...)
    try
        Logging.with_logger(logger) do
            _execute!(sim, solver; kwargs...)
        end
    catch e
        @error "Execution failed" exception = (e, catch_backtrace())
        return sim.status = SIMULATION_FAILED
    finally
        close(logger)
        return sim.status
    end
end

function read_results(sim::Simulation)
    return sim.results
end

function get_dynamic_wrapper(sim::Simulation, name::AbstractString)
    return get_dynamic_wrapper(get_simulation_inputs(sim), name)
end

function get_setpoints(sim::Simulation)
    return get_setpoints(get_simulation_inputs(sim))
end


module Injection

using PilePenetration
using PoingrSimulator.Poingr
using PoingrSimulator.GeometricObjects

function main_output_initialize(args)
    INPUT = args.INPUT

    history_file = joinpath(INPUT.Output.directory, "history.csv")
    PilePenetration.outputhistory_head(history_file)
end

function main_output(args)
    grid         = args.grid
    pointstate   = args.pointstate
    pile         = args.rigidbody
    pile0        = args.rigidbody0
    INPUT        = args.INPUT

    history_file = joinpath(INPUT.Output.directory, "history.csv")
    force_inside_directory = joinpath(INPUT.Output.directory, "force inside directory")
    force_outside_directory = joinpath(INPUT.Output.directory, "force outside directory")

    tip_height = 2 * pile[1][1]              # 1D = 2R
    tapered_height = pile[2][2] - pile[3][2] # tapered length
    if haskey(INPUT, :OutputHistory)
        if haskey(INPUT.OutputHistory, :tip_height)
            tip_height = oftype(tip_height, INPUT.OutputHistory.tip_height)
        end
        if haskey(INPUT.OutputHistory, :tapered_height)
            tapered_height = oftype(tapered_height, INPUT.OutputHistory.tapered_height)
        end
    end

    PilePenetration.outputhistory_append(
        history_file,
        grid,
        pointstate,
        pile,
        tip_height,
        tapered_height,
        sum(layer -> layer.thickness, INPUT.SoilLayer),
        centroid(pile0),
    )

    if args.INPUT.Pile.vacuum
        vacuum_height = INPUT.Pile.vacuum_height
        pile_inside = Polygon(Vec(0.0, pile[1][2]), Vec(0.0, pile[3][2]), pile[3], pile[2], pile[1])
        inds = findall(xₚ -> (xₚ in pile_inside) && xₚ[2] > (PilePenetration.tip_height(pile) + vacuum_height), pointstate.x)
        deleteat!(pointstate, inds)
    end
end

end


module TestMOIWrapper

using Ipopt
using MathOptInterface
using Test

const MOI = MathOptInterface

const OPTIMIZER = Ipopt.Optimizer()
MOI.set(OPTIMIZER, MOI.Silent(), true)

# Without fixed_variable_treatment set, duals are not computed for variables
# that have lower_bound == upper_bound.
MOI.set(
    OPTIMIZER,
    MOI.RawParameter("fixed_variable_treatment"),
    "make_constraint",
)

# TODO(odow): add features to Ipopt so we can remove some of this caching.
const BRIDGED_OPTIMIZER = MOI.Bridges.full_bridge_optimizer(
    MOI.Utilities.CachingOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
        OPTIMIZER,
    ),
    Float64,
)

const CONFIG = MOI.Test.TestConfig(
    atol = 1e-4,
    rtol = 1e-4,
    optimal_status = MOI.LOCALLY_SOLVED,
    infeas_certificates = false,
)

const CONFIG_NO_DUAL = MOI.Test.TestConfig(
    atol = 1e-4,
    rtol = 1e-4,
    optimal_status = MOI.LOCALLY_SOLVED,
    infeas_certificates = false,
    duals = false,  # Don't check dual result!
)

function test_solvername()
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "Ipopt"
end

function test_supports_default_copy_to()
    @test MOI.Utilities.supports_default_copy_to(OPTIMIZER, false)
    @test !MOI.Utilities.supports_default_copy_to(OPTIMIZER, true)
end

function test_basicconstraint()
    # TODO(odow): these basic constraint tests are broken.
    # MOI.Test.basic_constraint_tests(OPTIMIZER, CONFIG)
end

function test_unittest()
    return MOI.Test.unittest(
        BRIDGED_OPTIMIZER,
        CONFIG,
        String[
            # VectorOfVariables-in-SecondOrderCone not supported
            "delete_soc_variables",
            # NumberOfThreads not supported
            "number_threads",
            # MOI.Integer not supported.
            "solve_integer_edge_cases",
            # ObjectiveBound not supported.
            "solve_objbound_edge_cases",
            # DualObjectiveValue not supported.
            "solve_result_index",
            # Returns NORM_LIMIT instead of DUAL_INFEASIBLE
            "solve_unbounded_model",
            # MOI.ZeroOne not supported.
            "solve_zero_one_with_bounds_1",
            "solve_zero_one_with_bounds_2",
            "solve_zero_one_with_bounds_3",
        ],
    )
end

function test_ConstraintDualStart()
    model = Ipopt.Optimizer()
    x = MOI.add_variables(model, 2)
    l = MOI.add_constraint(model, x[1], MOI.GreaterThan(1.0))
    u = MOI.add_constraint(model, x[1], MOI.LessThan(1.0))
    e = MOI.add_constraint(model, x[2], MOI.EqualTo(1.0))
    c = MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.0),
        MOI.LessThan(1.5),
    )
    @test MOI.get(model, MOI.ConstraintDualStart(), l) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), u) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), e) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), c) === nothing
    @test MOI.get(model, MOI.NLPBlockDualStart()) === nothing
    MOI.set(model, MOI.ConstraintDualStart(), l, 1.0)
    MOI.set(model, MOI.ConstraintDualStart(), u, -1.0)
    MOI.set(model, MOI.ConstraintDualStart(), e, -1.5)
    MOI.set(model, MOI.ConstraintDualStart(), c, 2.0)
    MOI.set(model, MOI.NLPBlockDualStart(), [1.0, 2.0])
    @test MOI.get(model, MOI.ConstraintDualStart(), l) == 1.0
    @test MOI.get(model, MOI.ConstraintDualStart(), u) == -1.0
    @test MOI.get(model, MOI.ConstraintDualStart(), e) == -1.5
    @test MOI.get(model, MOI.ConstraintDualStart(), c) == 2.0
    @test MOI.get(model, MOI.NLPBlockDualStart()) == [1.0, 2.0]
    MOI.set(model, MOI.ConstraintDualStart(), l, nothing)
    MOI.set(model, MOI.ConstraintDualStart(), u, nothing)
    MOI.set(model, MOI.ConstraintDualStart(), e, nothing)
    MOI.set(model, MOI.ConstraintDualStart(), c, nothing)
    MOI.set(model, MOI.NLPBlockDualStart(), nothing)
    @test MOI.get(model, MOI.ConstraintDualStart(), l) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), u) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), e) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), c) === nothing
    @test MOI.get(model, MOI.NLPBlockDualStart()) === nothing
end

function test_contlinear()
    MOI.Test.contlineartest(
        BRIDGED_OPTIMIZER,
        CONFIG,
        String[
            # Tests requiring DualObjectiveValue. Tested below.
            "linear1",
            "linear2",
            "linear10",
            "linear14",
            # Tests requiring infeasibility certificates
            "linear8a",
            "linear8b",
            "linear8c",
            "linear12",
            # An INVALID_MODEL because it contains an empty 0 == 0 row.
            "linear15",
        ],
    )
    MOI.Test.linear1test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
    MOI.Test.linear2test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
    MOI.Test.linear10test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
    return MOI.Test.linear14test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
end

function test_qp()
    return MOI.Test.qptest(BRIDGED_OPTIMIZER, CONFIG)
end

function test_qcp()
    MOI.empty!(BRIDGED_OPTIMIZER)
    return MOI.Test.qcptest(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
end

function test_nlptest()
    return MOI.Test.nlptest(OPTIMIZER, CONFIG)
end

function test_getters()
    return MOI.Test.copytest(
        MOI.instantiate(Ipopt.Optimizer, with_bridge_type = Float64),
        MOI.Utilities.Model{Float64}(),
    )
end

function test_boundsettwice()
    MOI.Test.set_lower_bound_twice(OPTIMIZER, Float64)
    return MOI.Test.set_upper_bound_twice(OPTIMIZER, Float64)
end

function test_nametest()
    return MOI.Test.nametest(BRIDGED_OPTIMIZER)
end

function test_validtest()
    return MOI.Test.validtest(BRIDGED_OPTIMIZER)
end

function test_emptytest()
    return MOI.Test.emptytest(BRIDGED_OPTIMIZER)
end

function test_solve_time()
    model = Ipopt.Optimizer()
    x = MOI.add_variable(model)
    @test isnan(MOI.get(model, MOI.SolveTime()))
    MOI.optimize!(model)
    @test MOI.get(model, MOI.SolveTime()) > 0.0
end

# Model structure for test_check_derivatives_for_naninf()
struct Issue136 <: MOI.AbstractNLPEvaluator end
MOI.initialize(::Issue136, ::Vector{Symbol}) = nothing
MOI.features_available(d::Issue136) = [:Grad, :Jac]
MOI.eval_objective(::Issue136, x) = x[1]
MOI.eval_constraint(::Issue136, g, x) = (g[1] = x[1]^(1 / 3))
MOI.eval_objective_gradient(::Issue136, grad_f, x) = (grad_f[1] = 1.0)
MOI.jacobian_structure(::Issue136) = Tuple{Int64,Int64}[(1, 1)]
function MOI.eval_constraint_jacobian(::Issue136, J, x)
    J[1] = (1 / 3) * x[1]^(1 / 3 - 1)
    return
end

function test_check_derivatives_for_naninf()
    model = Ipopt.Optimizer()
    x = MOI.add_variable(model)
    MOI.set(
        model,
        MOI.NLPBlock(),
        MOI.NLPBlockData(MOI.NLPBoundsPair.([-Inf], [0.0]), Issue136(), false),
    )
    # Failure to set check_derivatives_for_naninf="yes" may cause Ipopt to
    # segfault or return a NUMERICAL_ERROR status. Check that it is set to "yes"
    # by obtaining an INVALID_MODEL status.
    # MOI.set(model, MOI.RawParameter("check_derivatives_for_naninf"), "no")
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.INVALID_MODEL
end

function test_deprecation()
    model = Ipopt.Optimizer(print_level = 0)
    @test MOI.get(model, MOI.RawParameter("print_level")) == 0
end

function test_callback()
    model = Ipopt.Optimizer()
    MOI.set(model, MOI.RawParameter("print_level"), 0)
    x = MOI.add_variable(model)
    MOI.add_constraint(model, MOI.SingleVariable(x), MOI.GreaterThan(1.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    f = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.5)
    MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
    x_vals = Float64[]
    function my_callback(
        prob::IpoptProblem,
        alg_mod::Cint,
        iter_count::Cint,
        obj_value::Float64,
        inf_pr::Float64,
        inf_du::Float64,
        mu::Float64,
        d_norm::Float64,
        regularization_size::Float64,
        alpha_du::Float64,
        alpha_pr::Float64,
        ls_trials::Cint,
    )
        c = Ipopt.column(x)
        push!(x_vals, prob.x[c])
        @test isapprox(obj_value, 1.0 * x_vals[end] + 0.5, atol = 1e-1)
        return iter_count < 1
    end
    MOI.set(model, Ipopt.CallbackFunction(), my_callback)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.INTERRUPTED
    @test length(x_vals) == 2
end

function test_empty_optimize()
    model = Ipopt.Optimizer()
    err = ErrorException(
        "IPOPT: Failed to construct problem because there are 0 variables. " *
        "If you intended to construct an empty problem, one work-around is " *
        "to add a variable fixed to 0.",
    )
    @test_throws err MOI.optimize!(model)
end

end  # module TestMOIWrapper

runtests(TestMOIWrapper)


"""
    TurbulenceConvection

Module for solving the sub-grid-scale equations
using the Eddy-Diffusivity-Mass-Flux (EDMF) model.
"""
module TurbulenceConvection

include("Grids/FiniteDifferenceGrids.jl")
include("StateVecs/StateVecs.jl")
include("LinearSolvers/TriDiagSolvers.jl")
# include("EDMF/EDMF.jl")

end # module