#Generate an ouput ImagineSignal for commanding the piezo to repeat the sample sequence `one_cyc` `ncycles` times.
#Appends those samples to `pos` and also returns dummy camera and laser signals so that Imagine doesn't complain
#`ncycles` default to 20 about seconds-worth of cycles based on the sample rate of `pos`
function pos_commands(pos::ImagineSignal, one_cyc, ncycles=ceil(Int, 20.0 / ustrip(inv(samprate(pos))*length(one_cyc))))
    @assert isempty(pos)
    rig_sigs = rigtemplate(rig_name(pos); sample_rate = samprate(pos))
    add_sequence!(pos, "bidi_cycle", one_cyc)
    cycle_nsamps = length(one_cyc)
    cam = first(getcameras(rig_sigs)) #not used
    las = first(getlasers(rig_sigs)) #not used
    add_sequence!(cam, "quiet", falses(cycle_nsamps)) #just to pass validation
    for i = 1:ncycles
        append!(pos, "bidi_cycle")
        append!(cam, "quiet")
        append!(las, "quiet")
    end
    pos_mon = getname(rig_sigs, monitor_name(pos))
    return [pos;cam;las;pos_mon]
end

function pos_commands(rig::AbstractString, pos_name::AbstractString, pstart, pstop, stack_rate, ncycles=ceil(Int, 20.0/ustrip(inv(stack_rate)));
                      sample_rate = 100000Hz, lp_cutoff=3.5*stack_rate)
    mod_cyc = vcat(gen_bidi_pos(pstart, pstop, 1/stack_rate, sample_rate; lp_cutoff=lp_cutoff)...)
    return pos_commands(rig, pos_name, mod_cyc, ncycles; sample_rate = sample_rate)
end

function pos_commands(rig::AbstractString, pos_name::AbstractString, mod_cyc, ncycles=ceil(Int, 20.0/ustrip(inv(stack_rate))); sample_rate = 100000Hz)
    rig_sigs = rigtemplate(rig; sample_rate = sample_rate)
    pos = getname(rig_sigs, pos_name)
    return pos_commands(pos, mod_cyc, ncycles)
end

function append_warmup!(pos, cam, las, high_las, pos_cycle, ncycs_warmup)
    nsamps_cyc = Base.length(pos_cycle) #should always be even
    add_sequence!(cam, "cam_warmup", falses(nsamps_cyc))
    add_sequence!(las, "las_warmup", falses(nsamps_cyc))
    add_sequence!(pos, "pos_warmup", pos_cycle)
    if rig_name(pos) == "ocpi-2"
        add_sequence!(high_las, "high_las_warmup", trues(nsamps_cyc))
    end
    for i = 1:ncycs_warmup
        append!(pos, "pos_warmup")
        append!(cam, "cam_warmup")
        append!(las, "las_warmup")
        append!(high_las, "high_las_warmup")
    end
	return pos, cam, las, high_las
end

function get_trial_intervals(ctr_i, lag, sample_rate, nsamps_exp, nsamps_flash)
	nsamps_from_ctr = div(nsamps_flash,2)
	nsamps_lag = ImagineInterface.calc_num_samps(lag, sample_rate)
    flash_start = ctr_i-nsamps_lag-nsamps_from_ctr
    flash_stop = ctr_i-nsamps_lag+nsamps_from_ctr
    exp_stop = flash_stop
    exp_start = flash_stop-nsamps_exp+1
    flash_itv = ClosedInterval(flash_start, flash_stop)
    exp_itv = ClosedInterval(exp_start, exp_stop)
    return exp_itv, flash_itv
end

function check_overlap(itvs::Vector{T}) where {T<:ClosedInterval}
    for i = 2:length(itvs)
        if !isempty(intersect(itvs[i-1], itvs[i]))
            error("Pulse intervals overlap")
        end
    end
    return itvs
end

#given a measured piezo cycle, a set of desired slice locations `slice_zs`, and a set of forward and reverse lags,
#return a tuple (cam_samps, las_samps, shift) with sample vectors for controlling camera exposure and laser pulsing
function get_cyc_pulses(mon_cyc, slice_zs, fwd_lags, back_lags, exp_time, flash_time, sample_rate::HasInverseTimeUnits)
	nsamps_exp = ImagineInterface.calc_num_samps(exp_time, sample_rate)
	nsamps_flash = ImagineInterface.calc_num_samps(flash_time, sample_rate)
    @assert iseven(length(mon_cyc))
    @assert issorted(slice_zs) || issorted(slice_zs; rev=true)
    nsamps_sweep = div(length(mon_cyc),2)
    cams_fwd = ClosedInterval{Int}[]
    cams_back = ClosedInterval{Int}[]
    lasers_fwd = ClosedInterval{Int}[]
    lasers_back = ClosedInterval{Int}[]
    nslices = length(slice_zs)
    #get_trial_intervals can return intervals lying outside the cycle bounds (negative first indices and last indices greater than 2*nsamps_sweep)
    #check extrema of intervals (smin and smax below)
    #if both of these extrema lie outside of 2*nsamps_sweep, throw an error
    #if only one of these does, it's okay.  return cam_samps, las_samps, and the shift value
    for s = 1:nslices
        ctr_fwd, ctr_back = find_bidi_locations(mon_cyc, sample_rate, slice_zs[s])
        cam_fwd, las_fwd = get_trial_intervals(ctr_fwd, fwd_lags[s], sample_rate, nsamps_exp, nsamps_flash)
        cam_back, las_back = get_trial_intervals(ctr_back, back_lags[s], sample_rate, nsamps_exp, nsamps_flash)
        #cam_back+=nsamps_sweep
        #las_back+=nsamps_sweep
        push!(cams_fwd, cam_fwd)
        push!(cams_back, cam_back)
        push!(lasers_fwd, las_fwd)
        push!(lasers_back, las_back)
    end
    las_itvs = vcat(lasers_fwd, lasers_back)
    cam_itvs = vcat(cams_fwd, cams_back)
    shift = 0
    all_itvs = vcat(las_itvs, cam_itvs)
    smin = minimum(minimum.(all_itvs))
    smax = maximum(maximum.(all_itvs))
    if smin < 1 && smax > 2*nsamps_sweep
        error("Cannot create pulse sequence without overlapping pulses across cycles")
    elseif !(smin >0 && (smax <= 2*nsamps_sweep))
        if abs(smin-1) > abs(smax - 2*nsamps_sweep)
            @show shift = smin-1
        else
            @show shift = smax - 2*nsamps_sweep
        end
    end
    las_samps = ImagineInterface.gen_pulses(2*nsamps_sweep, check_overlap(las_itvs.-shift))
    cam_samps = ImagineInterface.gen_pulses(2*nsamps_sweep, check_overlap(cam_itvs.-shift))
    return cam_samps, las_samps, shift #shift is the shift that should be applied to sample vectors to get them into proper alignment
end

function append_template!(pos, cam, las, high_las, slice_zs, exp_time, flash_time; freq=0.1Hz) #appends a stack sequence to each signal input. Images only on the forward sweep.
    #pad by 5μm if we have room
    z_min = minimum(slice_zs)
    z_max = maximum(slice_zs)
    z_pad = minimum([abs(z_min - 0.0μm), abs(800μm - z_max), 5.0μm])
    z_min-=z_pad
    z_max+=z_pad
    mod_cyc = vcat(gen_bidi_pos(z_min, z_max, 1/freq, samprate(pos))...)
    z_range = z_max - z_min #should also match the command (not necessarily the monitored signal)
    pos_speed = freq * z_range #in microns per second
    pad_nsecs = z_pad / pos_speed #nsecs to cover the z_pad distance
    pad_nsamps = ImagineInterface.calc_num_samps(pad_nsecs, samprate(pos))
    flash_cyc, cam_cyc, monshft = flash_cam_cycs(mod_cyc, mod_cyc, slice_zs, pad_nsamps, flash_time, exp_time, samprate(pos))
    #only take pictures on the forward sweep
    midpoint = div(length(flash_cyc),2) + 1
    flash_cyc[midpoint:end] .= false
    cam_cyc[midpoint:end] .= false
    append!(pos, "pos_template", mod_cyc)
    append!(cam, "cam_template", cam_cyc)
    append!(las, "las_template", flash_cyc)
    append!(high_las, "highlas_template", trues(length(mod_cyc)))
end

function multistack_bidi(cam_name, las_name, pos_name, pos_cyc, mon_cyc, nstacks_fwd, slice_zs, fwd_toffsets, back_toffsets, exp_time, flash_time, sample_rate, ncycs_warmup)
	ocpi2 = rigtemplate("ocpi-2"; sample_rate = sample_rate)
	pos_new = getname(ocpi2, pos_name)
	cam_new = getname(ocpi2, cam_name)
	las_new = getname(ocpi2, "all lasers")
	high_las_new = getname(ocpi2, las_name)

	FastScanning.append_warmup!(pos_new, cam_new, las_new, high_las_new, pos_cyc, ncycs_warmup)

    @show nsamps_offset = ImagineAnalyses.mon_delay(pos_cyc, mon_cyc)
	#temporal shift to make calculations easier (will shift things back later)
    mon_cyc = circshift(mon_cyc, -nsamps_offset)
	if nsamps_offset < 0
		error("can't handle negative offset")
	elseif nsamps_offset > 0
		append!(las_new, "delay", fill(false, nsamps_offset))
		append!(cam_new, "delay")
		append!(high_las_new, "delayhigh", fill(true, nsamps_offset))
	end
	cam_samps, las_samps, shift = get_cyc_pulses(mon_cyc, slice_zs, fwd_toffsets, back_toffsets, exp_time, flash_time, sample_rate)
	@assert shift == 0 #TODO: handle this

	nsamps_cyc = length(pos_cyc)
	pos_cyc_name = "pos_cyc"
	add_sequence!(pos_new, pos_cyc_name, pos_cyc)
	add_sequence!(cam_new, "cam_cyc", cam_samps)
	add_sequence!(las_new, "las_cyc", las_samps)
	add_sequence!(high_las_new, "high_cyc", trues(nsamps_cyc))

	for i = 1:nstacks_fwd
		append!(pos_new, "pos_cyc")
		append!(cam_new, "cam_cyc")
		append!(las_new, "las_cyc")
		append!(high_las_new, "high_cyc")
	end
    if nsamps_offset > 0 #add another positioner cycle because we have too many cam and las samples
        append!(pos_new, pos_cyc_name)
        append!(cam_new, "extra_cyc", fill(false, nsamps_cyc-nsamps_offset))
        append!(las_new, "extra_cyc")
        append!(high_las_new, "extra_cychigh", fill(true, nsamps_cyc-nsamps_offset))
    end
	return [pos_new; cam_new; las_new; high_las_new]
end

function slicetiming_commands!(pos::T1, mod_cyc, mon_cyc, las, high_las, cam, lags, slice_zs) where {T1<:ImagineSignal}
    nsamps_cyc = Base.length(mod_cyc) #should always be even
    nsamps_sweep = div(nsamps_cyc,2)
    ncycs_ignore = ceil(Int, (ustrip(samprate(pos))*5) / nsamps_cyc) #5 seconds of warmup
    #@show nsamps_offset = mon_lag_nsamps(mod_cyc, mon_cyc)
    @show nsamps_offset = ImagineAnalyses.mon_delay(mod_cyc, mon_cyc)
	#temporal shift to make calculations easier (will shift things back later)
    mon_cyc = circshift(mon_cyc, -nsamps_offset)
    ntrials=1 #number of trials per lag value per z slice
    #We want to deliver the laser pulse during the global shutter time (tglobal in PCO camera manual)
    #If the user images with full chip then the global exposure segment doesn't begin until _10ms_ after the exposure starts.
    #Therefore with a 1ms laser pulse we need an 11ms exposure time, and the pulse is placed in the last 1ms
    exp_time = 0.002s
    flash_time = 0.001s
    empty!(pos)
	if nsamps_offset < 0
		error("Cannot handle negative offset (but could with a bit more coding work)")
	elseif nsamps_offset > 0
		append!(las, "delay", fill(false, nsamps_offset))
		append!(cam, "delay")
		append!(high_las, "delayhigh", fill(true, nsamps_offset))
	end
    nsamps_flash = ImagineInterface.calc_num_samps(flash_time, samprate(pos))
    if iseven(nsamps_flash)
	    nsamps_flash -= 1 #force odd so that fwd and reverse flashes are easier to align
    end
	nsamps_exp = ImagineInterface.calc_num_samps(exp_time, samprate(pos))

	#Sequence from here:
	# one slow image stack
	# several fast "warmup" piezo cycles with no camera / laser activity (ncycs_ignore)
	# ntrials * nlags piezo cycles in which we take a pair of slices for each cycle
	# repeat above step for each slice, so a total of ntrials * nlags * nslices cycles (and twice that many images)
    ncycs_ignore = ceil(Int, (ustrip(samprate(pos))*5) / nsamps_cyc) #5 seconds of warmup
	append_template!(pos, cam, las, high_las, slice_zs, exp_time, flash_time) #appends slow stack sequence to each signal input
	append_warmup!(pos, cam, las, high_las, mod_cyc, ncycs_ignore) #note: the warmup cycles are named differently but sequence is same as other fast pos cycles

    pos_cyc_name = "fast_pos_cyc"
	add_sequence!(pos, pos_cyc_name, mod_cyc)
	add_sequence!(high_las, "laser_high", trues(nsamps_cyc)) #will append one of these for every cycle as we go
	#Now iterate through slices, lags, and trials
	for s = 1:length(slice_zs)
		fwd_ctr, back_ctr = find_bidi_locations(mon_cyc, samprate(pos), slice_zs[s])
		for l = 1:Base.length(lags)
            cyc_exp_name = "cyc_exp_$(l)_slice_$s"
            cyc_flash_name = "cyc_flash_$(l)_slice_$s"
            has_partial = false #do we need to add a partial cycle before and after each cycle due to exposure pulse timing?
            cam_samps, las_samps, timeshift = get_cyc_pulses(mon_cyc, [slice_zs[s]], [lags[l];], [lags[l];], exp_time, flash_time, samprate(pos))
			add_sequence!(cam, cyc_exp_name, cam_samps)
			add_sequence!(las, cyc_flash_name, las_samps)
            #if timeshift is nonzero then the imaging pulses lie partially outside the cycle and we need an extra cycle
            if timeshift < 0
                warn("Adding an extra piezo cycle in slice $s with lag $l for exposures or flashes that begin before the cycle")
                has_partial = true
                add_sequence!(cam, "cam_partial_before_$(s)_$(l)", falses(nsamps_cyc+timeshift))
                add_sequence!(las, "las_partial_before_$(s)_$(l)", falses(nsamps_cyc+timeshift))
                add_sequence!(cam, "cam_partial_after_$(s)_$(l)",  falses(-timeshift))
                add_sequence!(las, "las_partial_after_$(s)_$(l)",  falses(-timeshift))
                @show s
                @show l
            elseif timeshift > 0
                has_partial = true
                warn("Adding an extra piezo cycle in slice $s with lag $l for exposures or flashes that end after the cycle")
                add_sequence!(cam, "cam_partial_before_$(s)_$(l)", falses(timeshift))
                add_sequence!(las, "las_partial_before_$(s)_$(l)", falses(timeshift))
                add_sequence!(cam, "cam_partial_after_$(s)_$(l)", falses(nsamps_cyc-timeshift))
                add_sequence!(las, "las_partial_after_$(s)_$(l)", falses(nsamps_cyc-timeshift))
                @show s
                @show l
            end
            if has_partial
                append!(pos, pos_cyc_name)
                append!(cam, "cam_partial_before_$(s)_$(l)")
                append!(las, "las_partial_before_$(s)_$(l)")
                append!(high_las, "laser_high")
            end
			for t = 1:ntrials
				append!(pos, pos_cyc_name)
				append!(cam, cyc_exp_name)
				append!(las, cyc_flash_name)
				append!(high_las, "laser_high")
			end
            if has_partial
                append!(cam, "cam_partial_after_$(s)_$(l)")
                append!(las, "las_partial_after_$(s)_$(l)")
            end
		end
	end
    if nsamps_offset > 0 #add another positioner cycle because we have too many cam and las samples
        append!(pos, pos_cyc_name)
        append!(cam, "extra_cyc", fill(false, nsamps_cyc-nsamps_offset))
        append!(las, "extra_cyc")
        append!(high_las, "extra_cychigh", fill(true, nsamps_cyc-nsamps_offset))
    end
    @assert length(pos) == length(cam) == length(las) == length(high_las)
    return [pos; cam; las; high_las]
end

function slicetiming_commands(pos::T1, mon_cyc, las_name::AbstractString, cam_name::AbstractString, lags, slice_zs) where {T1<:ImagineSignal}
    mod_cyc = get_samples(pos, first(sequence_names(pos))) #relies on the cycle being the first (and only) subsequence
    @assert all(sequence_names(pos).==first(sequence_names(pos)))
    rig_sigs = rigtemplate(rig_name(pos); sample_rate = samprate(pos))
    cam = getname(rig_sigs, cam_name)
    las = high_las = all_las = -1
    @assert rig_name(pos) == "ocpi-2"
    if las_name != "all lasers"
        las = getname(rig_sigs, "all lasers")
        high_las = getname(rig_sigs, las_name)
        all_las = [las;high_las]
    else
        error("Please specify a specific laser, not all lasers")
    end
    if !ImagineInterface.ispos(pos) error("Only positioner signals are currently supported.") end
    if !ImagineInterface.iscam(cam) error("$cam_name is not a camera trigger signal.") end
    if !ImagineInterface.islas(las) error("$las_name is not a laser trigger signal.") end

    slicetiming_commands!(deepcopy(pos), mod_cyc, mon_cyc, las, high_las, cam, lags, slice_zs)
end


"""
    insert!(tree::RTree, pt::Point, [id], val)
    insert!(tree::RTree, br::Rect, [id], val)
    insert!(tree::RTree, elem::SpatialElem)

Inserts `val` value identified by `br` bounding box (or point `pt`) and
`id` (if tree elememnts support `HasID` trait) into the `tree`.
"""
function Base.insert!(tree::RTree, el::Any)
    check_eltype_rtree(el, tree)
    _insert!(tree, el, 1)
    tree.nelems += 1
    tree.nelem_insertions += 1
    return tree
end

Base.insert!(tree::RTree{T,N,SpatialElem{T,N,K,V}}, br::Rect{T,N}, id::Any, val::Any) where {T,N,K,V} =
    insert!(tree, SpatialElem{T,N,K,V}(br, id, val))

Base.insert!(tree::RTree{T,N,SpatialElem{T,N,K,V}},
             pt::Point{T,N}, id::Any, val::Any) where {T,N,K,V} =
    insert!(tree, SpatialElem{T,N,K,V}(Rect(pt), id, val))

Base.insert!(tree::RTree{T,N,SpatialElem{T,N,Nothing,V}},
             br::Union{Rect{T,N}, Point{T,N}}, val::Any) where {T,N,V} =
    insert!(tree, br, nothing, val)

# inserts the child into the node rebalancing R-tree as needed
# returns true if MBR was changed and upstream R-tree readjustments are required
function _insert!(node::Node, child::Any, con::SubtreeContext)
    if length(node) < capacity(node, con.tree) # underfull node
        oldmbr = mbr(node) # this has to happen before _attach!() modifies node's MBR
        _attach!(node, child, con.tree)
        if !in(mbr(child), oldmbr) && hasparent(node)
            #@debug "_insert!(): mbr change"
            _updatembr!(parent(node), pos_in_parent(node), oldmbr, con)
            return true
        else
            return false
        end
    elseif variant(con.tree) == RTreeStar && (node !== con.tree.root) && !isoverflow(con, level(node))
        return _insert!_fullnode_rstar(node, child, con)
    else
        return _insert!_fullnode(node, child, con)
    end
end

function _insert!(con::SubtreeContext, node::Any, lev::Integer)
    #@debug "_insert!(con): len=$(length(con.tree)) node_lev=$(node isa Node ? level(node) : 0) node_len=$(node isa Node ? length(node) : 0)"
    par = choose_subtree(con.tree.root, lev, mbr(node), con.tree)
    @assert level(par) == lev
    _insert!(par, node, con)
end

_insert!(tree::RTree, node::Any, lev::Int) =
    _insert!(SubtreeContext(tree), node, lev)

# index of the child that would have the least MBR enlargement by br
function find_least_enlargement(node::Node, br::Rect)
    min_enl = Inf
    best_ix = 0
    best_area = Inf
    for i in eachindex(children(node))
        child_mbr = mbr(@inbounds node[i])
        child_area = area(child_mbr)
        enl = combined_area(child_mbr, br) - child_area
        if enl < min_enl || (min_enl == enl && child_area < best_area)
            min_enl = enl
            best_area = child_area
            best_ix = i
        end
    end
    return best_ix
end

# index of the child that would have the least MBR overlap with br
function find_least_overlap(node::Node, br::Rect, tree::RTree)
    min_enl = Inf
    best_area = Inf
    best_ix = 0
    enls = Vector{Float64}(undef, length(node)) # FIXME use pool?
    # find combined region and enlargement of every entry and store it.
    for (i, child) in enumerate(children(node))
        child_area = area(mbr(child))
        enls[i] = enl = combined_area(br, mbr(child)) - child_area
        if enl < min_enl || (enl == min_enl && child_area < best_area)
            min_enl = enl
            best_area = child_area
            best_ix = i
        end
    end

    abs(min_enl) <= eps(min_enl) && return best_ix # no enlargement

    if length(node) > tree.nearmin_overlap
        # sort entries in increasing order of enlargement
        enl_order = sortperm(enls) # FIXME reuse from pool
        @assert min_enl == enls[enl_order[1]] <= enls[enl_order[end]]
        niter = tree.nearmin_overlap
    else
        niter = length(node)
        enl_order = collect(1:niter) # FIXME reuse from pool
    end

    # calculate overlap of most important original entries (near minimum overlap cost).
    min_delta_olap = Inf
    best_ix = 0
    best_area = Inf
    best_enl = Inf
    for i in 1:niter
        node_ix = enl_order[i]
        node_mbr = mbr(node[node_ix])
        comb_mbr = combine(node_mbr, br)
        delta_olap = 0.0

        for (j, child) in enumerate(children(node))
            if node_ix != j
                olap = overlap_area(comb_mbr, mbr(child))
                if olap > 0.0
                    delta_olap += olap - overlap_area(node_mbr, mbr(child))
                end
            end
        end

        node_area = area(node_mbr)
        # keep the one with the least, in the order:
        #   delta overlap, enlargement, area
        if (delta_olap < min_delta_olap) || (delta_olap == min_delta_olap &&
           (enls[node_ix] < best_enl || (enls[node_ix] == best_enl &&
            node_area < best_area)))
            min_delta_olap = delta_olap
            best_ix = node_ix
            best_area = node_area
            best_enl = enls[node_ix]
        end
    end
    return best_ix
end

# choose the node (for inserting the given mbr at specified level)
function choose_subtree(node::Node, lev::Int, br::Rect, tree::RTree)
    while level(node) > lev
        if variant(tree) == RTreeLinear || variant(tree) == RTreeQuadratic
            child_ix = find_least_enlargement(node, br)
        elseif variant(tree) == RTreeStar
            child_ix = level(node) == 1 ?
                    find_least_overlap(node, br, tree) : # for the node pointing to leaves
                    find_least_enlargement(node, br)
        else
            throw(SpatialIndexException("RTree variant not supported"))
        end
        node = node[child_ix]
    end
    level(node) == lev ||
        throw(SpatialIndexException("Couldn't find the node of level $lev in RTree of height $(height(tree))"))
    return node
end

# R*-tree-specific insertion into the full node
# splits the node and rebalances the R-tree
function _insert!_fullnode_rstar(node::Node, child::Any, con::SubtreeContext)
    #@debug "_insert!_fullnode_rstar(): lev=$(level(node)) len=$(length(node))"
    @assert isa(con.level_overflow, BitVector)
    @assert hasparent(node)

    setoverflow!(con, level(node)) # the R* insert could only be done once per level
    oldmbr = mbr(node)
    _attach!(node, child, con.tree, force=true) # force-push the child exceeding the capacity

    node_center = center(mbr(node))
    # calculate relative distance of every entry from the node MBR (ignore square root.)
    children_dist = sqrdistance.(Ref(node_center), center.(mbr.(children(node)))) # FIXME reuse array?
    # return the indices sorted by the increasing order of distances.
    # Since children_dist is sorted in ascending order
    # it would match the reinsertion order suggested in the paper.
    children_order = sortperm(children_dist) # FIXME reuse array?
    nkeep = length(node) - floor(Int, length(node) * con.tree.reinsert_factor)

    newnode = acquire(con.tree, typeof(node), level(node))   # the replacement node
    for i in 1:nkeep
        _attach!(newnode, node[children_order[i]], con.tree)
    end

    # Divertion from R*-Tree algorithm proposed in libspatialindex:
    # First update the path to the root, then start reinserts, to avoid complicated handling
    # of changes to the same node from multiple insertions.
    #@debug "_insert!_fullnode_rstar(): replace the node"
    _replace!(node, newnode, oldmbr, con) # replace the node (so (nkeep+1):... children are detached from the tree)
    #@debug "_insert!_fullnode_rstar(): reinsert $(length((nkeep+1):length(children_order))) children of lev=$(level(node)) node"
    # don't use par during/after reinsertion, it might have changed
    for i in (nkeep+1):length(children_order)
        child = node[children_order[i]]
        #@debug "_insert!_fullnode_rstar(): reinsert #$i child (type=$(typeof(child)) lev=$(child isa Node ? level(child) : "data") parent=$(isa(child, Node) && hasparent(child)))"
        _insert!(con, child, level(node))
        con.tree.nnode_reinsertions += 1
    end

    release(con.tree, node) # return node to the pool
    return true
end

# R-tree-specific insertion into the full node
# splits the node and rebalances the R-tree
function _insert!_fullnode(node::Node, child::Any, con::SubtreeContext)
    #@debug "_insert!_fullnode(): lev=$(level(node)) len=$(length(node)) parent=$(hasparent(node))"
    oldmbr = mbr(node)
    _attach!(node, child, con.tree, force=true) # temporary overflow the node capacity
    n1, n2 = _split!(node, con.tree)

    if node === con.tree.root
        #@debug "_insert!_fullnode(): newroot tree_height=$(height(con.tree)) tree_len=$(length(con.tree))"
        # the node must be the root, make the new root and attach n1, n2 to it
        newroot = acquire(con.tree, Branch, level(node) + 1)
        _attach!(newroot, n1, con.tree)
        _attach!(newroot, n2, con.tree)
        # tree grows
        con.tree.root = newroot
        con.tree.nnodes_perlevel[level(node)] = 2
        push!(con.tree.nnodes_perlevel, 1)
    else # non-root, n1 and n2 replace the node in its parent
        _replace!(node, n1, n2, oldmbr, con)
    end
    return true
end


using aoc2019.sif: image, resolve, checksum

input = joinpath(@__DIR__, "input")
bits = [parse(Int, char) for char in readline(input)]
img = image(bits, 25, 6)

p1 = checksum(img)
@assert p1 == 2806

p2 = resolve(img)

print("-----------------------------------------------------------------------\n")
print("space image format -- part one\n    corruption check : $p1\n")
print("space image format -- part two\n    password         :\n")
for y in 1:6
    print("        ")
    for x in 1:25
        glyph = p2[x, y] == 1 ? "#" : " "
        print("$glyph")
    end
    println()
end
print("-----------------------------------------------------------------------\n")


# This file is a part of Julia. License is MIT: https://julialang.org/license

# Support for @simd for

module SimdLoop

export @simd, simd_outer_range, simd_inner_length, simd_index

# Error thrown from ill-formed uses of @simd
struct SimdError <: Exception
    msg::AbstractString
end

# Parse iteration space expression
#       symbol '=' range
#       symbol 'in' range
function parse_iteration_space(x)
    (isa(x, Expr) && (x.head == :(=) || x.head == :in)) || throw(SimdError("= or in expected"))
    length(x.args) == 2 || throw(SimdError("simd range syntax is wrong"))
    isa(x.args[1], Symbol) || throw(SimdError("simd loop index must be a symbol"))
    x.args # symbol, range
end

# reject invalid control flow statements in @simd loop body
function check_body!(x::Expr)
    if x.head === :break || x.head == :continue
        throw(SimdError("$(x.head) is not allowed inside a @simd loop body"))
    elseif x.head === :macrocall && x.args[1] === Symbol("@goto")
        throw(SimdError("$(x.args[1]) is not allowed inside a @simd loop body"))
    end
    for arg in x.args
        check_body!(arg)
    end
    return true
end
check_body!(x::QuoteNode) = check_body!(x.value)
check_body!(x) = true

# @simd splits a for loop into two loops: an outer scalar loop and
# an inner loop marked with :loopinfo. The simd_... functions define
# the splitting.
# Custom iterators that do not support random access cannot support
# vectorization. In order to be compatible with `@simd` annotated loops,
#they should override `simd_inner_length(v::MyIter, j) = 1`,
#`simd_outer_range(v::MyIter) = v`, and `simd_index(v::MyIter, j, i) = j`.

# Get range for outer loop.
simd_outer_range(r) = 0:0

# Get trip count for inner loop.
@inline simd_inner_length(r, j) = Base.length(r)

# Construct user-level element from original range, outer loop index j, and inner loop index i.
@inline simd_index(r, j, i) = (@inbounds ret = r[i+firstindex(r)]; ret)

# Compile Expr x in context of @simd.
function compile(x, ivdep)
    (isa(x, Expr) && x.head == :for) || throw(SimdError("for loop expected"))
    length(x.args) == 2 || throw(SimdError("1D for loop expected"))
    check_body!(x)

    var,range = parse_iteration_space(x.args[1])
    r = gensym("r") # Range value
    j = gensym("i") # Iteration variable for outer loop
    n = gensym("n") # Trip count for inner loop
    i = gensym("i") # Trip index for inner loop
    quote
        # Evaluate range value once, to enhance type and data flow analysis by optimizers.
        let $r = $range
            for $j in Base.simd_outer_range($r)
                let $n = Base.simd_inner_length($r,$j)
                    if zero($n) < $n
                        # Lower loop in way that seems to work best for LLVM 3.3 vectorizer.
                        let $i = zero($n)
                            while $i < $n
                                local $var = Base.simd_index($r,$j,$i)
                                $(x.args[2])        # Body of loop
                                $i += 1
                                $(Expr(:loopinfo, Symbol("julia.simdloop"), ivdep))  # Mark loop as SIMD loop
                            end
                        end
                    end
                end
            end
        end
        nothing
    end
end

"""
    @simd

Annotate a `for` loop to allow the compiler to take extra liberties to allow loop re-ordering

!!! warning
    This feature is experimental and could change or disappear in future versions of Julia.
    Incorrect use of the `@simd` macro may cause unexpected results.

The object iterated over in a `@simd for` loop should be a one-dimensional range.
By using `@simd`, you are asserting several properties of the loop:

* It is safe to execute iterations in arbitrary or overlapping order, with special consideration for reduction variables.
* Floating-point operations on reduction variables can be reordered, possibly causing different results than without `@simd`.

In many cases, Julia is able to automatically vectorize inner for loops without the use of `@simd`.
Using `@simd` gives the compiler a little extra leeway to make it possible in more situations. In
either case, your inner loop should have the following properties to allow vectorization:

* The loop must be an innermost loop
* The loop body must be straight-line code. Therefore, [`@inbounds`](@ref) is
    currently needed for all array accesses. The compiler can sometimes turn
    short `&&`, `||`, and `?:` expressions into straight-line code if it is safe
    to evaluate all operands unconditionally. Consider using the [`ifelse`](@ref)
    function instead of `?:` in the loop if it is safe to do so.
* Accesses must have a stride pattern and cannot be "gathers" (random-index
    reads) or "scatters" (random-index writes).
* The stride should be unit stride.

!!! note
    The `@simd` does not assert by default that the loop is completely free of loop-carried
    memory dependencies, which is an assumption that can easily be violated in generic code.
    If you are writing non-generic code, you can use `@simd ivdep for ... end` to also assert that:

* There exists no loop-carried memory dependencies
* No iteration ever waits on a previous iteration to make forward progress.
"""
macro simd(forloop)
    esc(compile(forloop, nothing))
end

macro simd(ivdep, forloop)
    if ivdep == :ivdep
        esc(compile(forloop, Symbol("julia.ivdep")))
    else
        throw(SimdError("Only ivdep is valid as the first argument to @simd"))
    end
end

end # module SimdLoop


# SPD-License-Identifier: MIT

using ArgParse
using BayesianModels
using CUDA
using HDF5
using JLD2
using LinearAlgebra
using Random
using TOML

CUDA.allowscalar(false)

include("utils.jl")

const DEFAULTS_MODEL = Dict(
    "ncomponents" => 1,
    "covtype" => "full",
    "priortype" => "categorical",
    "noise_init" => 0.1,
    "pstrength" => 1
)

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--checkpoint-rate", "-c"
            default = 100
            arg_type = Int
            help = "save a checkpoint of the model every C updates"
        "--learning-rate", "-l"
            default = 0.1
            arg_type = Float64
            help = "constant learning rate of the training"
        "--mini-batch-size", "-m"
            default = 1
            arg_type = Int
            help = "number of utterances per mini-batch"
        "--single-precision", "-S"
            action = :store_true
            help = "use float32 data type for computations"
        "--start-from", "-s"
            default = ""
            help = "start training from the given checkpoint model"
        "--start-update", "-k"
            default = 1
            arg_type = Int
            help = "start training the given update number"
        "--steps", "-n"
            default = 1
            arg_type = Int
            help = "number of training steps"
        "--use-gpu", "-g"
            action = :store_true
            help = "use a GPU for training the model"
        "config"
            required = true
            help = "GMM configuration file"
        "stats"
            required = true
            help = "statistics of the data to initialize the model"
        "features"
            required = true
            help = "HDF5 archive containing the features"
        "outdir"
            required = true
            help = "directory"
    end
    s.description = """
    Train a Gaussian Mixture Model on a set of utterances. The list
    of utterances to use is read from the standard input. Upon
    completion, The final model will be stored in outdir/final.jld2.
    """
    parse_args(s)
end

function main(args)
    config = TOML.parsefile(args["config"])
    outdir = args["outdir"]
    T = args["single-precision"] ? Float32 : Float64

    utts = collect(eachline(stdin))

    model = nothing
    modelconf = nothing
    if isempty(args["start-from"])
        modelconf = merge(DEFAULTS_MODEL, config["model"])
        f = h5open(args["stats"], "r")
        μ = read(f["mean"])
        Σ = read(f["cov"])
        close(f)
        model = makemodel(T, modelconf, μ, Σ)
    else
        dict = load(args["start-from"])
        model, modelconf = loadmodel(args["start-from"])
    end

    if args["use-gpu"]
        init_gpu()
        model |> gpu!
    end

    checkrate = args["checkpoint-rate"]
    lrate = args["learning-rate"]
    steps = args["steps"]
    mb_size = args["mini-batch-size"]
    inferenceconf = Dict(
        "learning-rate" => lrate,
        "steps" => steps,
        "mb_size" => mb_size
    )
    params = filter(isbayesianparam, getparams(model))
    cache = Dict()
    h5open(args["features"], "r") do ffea
        for step in args["start-update"]:steps
            # Clear cache for gradient computation
            empty!(cache)

            # Randomize the utterance list
            shuffle!(utts)[1:mb_size]

            # Load the minibatch data
            X = hcat([read(ffea[utts[n]]) for n in 1:mb_size]...)
            X = convert(Array{T}, X)
            if args["use-gpu"] X = X |> CuArray end

            # One gradient step
            scale = length(utts) / mb_size
            𝓛, llh, KL = elbo(model, X; cache = cache, stats_scale = scale,
                              detailed = true)
            ∇ξ = ∇elbo(model, cache, params)
            gradstep(∇ξ, lrate = lrate)

            # Log the status of the training in JSON format
            norm = scale * size(X,2)
            norm𝓛 = 𝓛 / norm
            normllh = llh / norm
            normKL = KL / norm
            println("{\"step\": $step, \"nsteps\":$steps, \"elbo\": $norm𝓛, \"llh\":$normllh, \"kl\":$normKL}")

            # Checkpoint
            if step % checkrate == 0
                model |> cpu!
                savemodel(outdir, model, modelconf, inferenceconf, step)
                if args["use-gpu"] model |> gpu!  end
                params = filter(isbayesianparam, getparams(model))
            end
        end
    end

    model |> cpu!
    savemodel(outdir, model, modelconf, inferenceconf)
end

args = parse_commandline()
main(args)


"""
Describes a single water reservoir, using volumes (m^3) and discharges (m^3/s).
It is linked to a series of input rivers that provied water, and has to
fulfil some requirements with the output (tap or industry water, environmental
flow, with or without uncertainty).

TODO:

  * what about topology? (where inputs and diversions can be other reservoirs!).
    Store the rivers inside the basin, and only use names here? Could also get
    rid of different arrays for rivers. Could also store information about the links
    to improve the modelling (e.g., add the wave propagation along this edge).

  * consider uncertainty about purposes (like drinking water consumption)?
"""
immutable Reservoir
  name::AbstractString
  variant::AbstractString
  purposes::Array{Purpose, 1}
  hydropower::Array{HydropowerUnit, 1}
  outputs::Array{DamOutput, 1}

  # Reservoir bathymetry.
  capacity::Tuple{Float64, Float64} # 10^6 m^3; minimum is enough to ensure drinking water intake, by hypothesis.
  bathymetry::ReservoirBathymetry

  # Topography.
  rivers_in::Array{NaturalRiver, 1}
  rivers_diverted::Array{DivertedRiver, 1}

  # Private constructor, without units nor defaut values.
  function Reservoir(name::AbstractString, variant::AbstractString, purposes::Array{Purpose, 1},
                     hydropower::Array{HydropowerUnit, 1}, outputs::Array{DamOutput, 1},
                     capacity::Tuple{Float64, Float64}, bathymetry::ReservoirBathymetry,
                     rivers_in::Array{NaturalRiver, 1}, rivers_diverted::Array{DivertedRiver, 1}
                    )
    ## Some consistency check for the inputs.
    if length(rivers_in) + length(rivers_diverted) == 0
      error("Must have at least one input river!")
    end

    ## If there is hydropower, complement the purposes and outputs.
    # First check if these hydropower units are notalready in the given arrays!
    purposeHPs = HydropowerUnit[getHydropowerUnit(p) for p in filter(p -> isHydropower(p), purposes)]
    outputHPs = HydropowerUnit[getHydropowerUnit(o) for o in filter(o -> isHydropower(o), outputs)]
    for hp in hydropower
      if ! in(hp, purposeHPs)
        push!(purposes, DeterministicHydropower(hp))
      end
      if ! in(hp, outputHPs)
        push!(outputs, HydropowerDamOutput(hp))
      end
    end

    ## Build the object.
    obj = new(name, variant, purposes, hydropower, outputs,
              capacity, bathymetry,
              rivers_in, rivers_diverted)

    ## Check consistency.
    # Check consistency of the bathymetry.
    if capacity[1] > capacity[2]
      error("Inconsistent capacities: the minimum is higher than the maximum!")
    end

    # Check consistency of scenarios.
    countScenariosWithConsistency(obj)
    countTimeStepsWithConsistency(obj)
    getPeriodWithConsistency(obj)
    getTimePointsWithConsistency(obj)

    # Check consistency of the ids.
    ids = vec(vcat([getId(p) for p in purposes], [getId(hp) for hp in hydropower], [getId(o) for o in outputs]))
    ids_cnt = filter((k,v) -> v > 1, countmap(ids))
    if length(ids_cnt) > 0
      error("IDs are not consistent! Already met: " * string(ids_cnt) * ".")
    end

    ## Done!
    return obj
  end
end

# Public constructor, with units and keyword arguments.
Reservoir(;name::AbstractString=error("Missing reservoir name"),
           variant::AbstractString="vanilla",
           purposes::Array{Purpose, 1}=Purpose[],
           hydropower::Array{HydropowerUnit, 1}=HydropowerUnit[],
           outputs::Array{DamOutput, 1}=DamOutput[],

           capacity::Tuple{Volume, Volume}=error("Missing reservoir capacities"),
           bathymetry::ReservoirBathymetry=NullReservoirBathymetry(),

           rivers_in::Array{NaturalRiver, 1}=NaturalRiver[],
           rivers_diverted::Array{DivertedRiver, 1}=DivertedRiver[]
          ) = Reservoir(name, variant, purposes, hydropower, outputs,
                        _from_unitful(capacity), bathymetry,
                        rivers_in, rivers_diverted)

copy(r::Reservoir; name::AbstractString=r.name,
                   variant::AbstractString=r.variant,
                   purposes::Array{Purpose, 1}=r.purposes,
                   hydropower::Array{HydropowerUnit, 1}=r.hydropower,
                   outputs::Array{DamOutput, 1}=r.outputs,
                   capacity::Tuple{Union{Float64, Volume}, Union{Float64, Volume}}=r.capacity,
                   bathymetry::ReservoirBathymetry=r.bathymetry,
                   rivers_in::Array{NaturalRiver, 1}=r.rivers_in,
                   rivers_diverted::Array{DivertedRiver, 1}=r.rivers_diverted
    ) =
  Reservoir(name, variant, purposes, hydropower, outputs,
            _from_unitful(capacity), bathymetry,
            rivers_in, rivers_diverted)

Base.isempty(r::Reservoir) = false

### Generic getters.

getName(r::Reservoir) = r.name
getVariant(r::Reservoir) = r.variant
getMinimumCapacity(r::Reservoir) = r.capacity[1]
getMaximumCapacity(r::Reservoir) = r.capacity[2]

getHydropowerUnits(r::Reservoir) = r.hydropower
hasHydropower(r::Reservoir) = length(getHydropowerUnits(r)) > 0
getHydropowerIds(r::Reservoir) = [getId(hp) for hp in getHydropowerUnits(r)]
getPenstockHydropower(r::Reservoir) = sum([hp.max_discharge for hp in getHydropowerUnits(r)])

getPurposes(r::Reservoir) = r.purposes
getOutputs(r::Reservoir) = r.outputs
getBathymetry(r::Reservoir) = r.bathymetry

function _check_purpose_subtype(purposeType::DataType)
  if ! (purposeType <: Purpose)
    error("Wrong type argument: must be a subtype of Purpose; got " * string(purposeType) * ".")
  end
end
function hasPurpose(r::Reservoir, purposeType::DataType)
  _check_purpose_subtype(purposeType)
  return length(filter(t -> isa(t, purposeType), getPurposes(r))) > 0
end
function getPurpose(r::Reservoir, purposeType::DataType)
  _check_purpose_subtype(purposeType)
  # Return an array of purposes, as the type has no unicity guarantee (as opposed to IDs).
  return filter(t -> isa(t, purposeType), getPurposes(r))
end
getPurposeIds(r::Reservoir) = Symbol[getId(o) for o in getPurposes(r)]
hasPurpose(r::Reservoir, s::Symbol) = length(filter(o -> getId(o) == s, getPurposes(r))) > 0
function getPurpose(r::Reservoir, s::Symbol; nullable::Bool=false)
  elt = filter(o -> getId(o) == s, getPurposes(r))
  if nullable
    if length(elt) == 0
      return Nullable{Symbol}()
    else
      return Nullable(elt[1])
    end
  else
    if length(elt) == 0
      error("Purpose " * string(s) * " not found.")
    else
      return elt[1]
    end
  end
end

getOutputIds(r::Reservoir) = Symbol[getId(o) for o in getOutputs(r)]
hasOutput(r::Reservoir, s::Symbol) = length(filter(o -> getId(o) == s, getOutputs(r))) > 0
function getOutput(r::Reservoir, s::Symbol; nullable::Bool=false)
  elt = filter(o -> getId(o) == s, getOutputs(r))
  if nullable
    if length(elt) == 0
      return Nullable{Symbol}()
    else
      return Nullable(elt[1])
    end
  else
    if length(elt) == 0
      error("Purpose " * string(s) * " not found.")
    else
      return elt[1]
    end
  end
end

name(r::Reservoir) = getName(r)
variant(r::Reservoir) = getVariant(r)
minCapacity(r::Reservoir) = getMinimumCapacity(r)
maxCapacity(r::Reservoir) = getMaximumCapacity(r)
penstockHydropower(r::Reservoir) = getPenstockHydropower(r)
hydropowerUnits(r::Reservoir) = getHydropowerUnits(r)
outputs(r::Reservoir) = getOutputs(r)
outputsId(r::Reservoir) = getOutputIds(r)
output(r::Reservoir, s::Symbol; kwargs...) = getOutput(r, s; kwargs...)
bathymetry(r::Reservoir) = getBathymetry(r)
purpose(r::Reservoir, s::Symbol; kwargs...) = getPurpose(r, s; kwargs...)
purpose(r::Reservoir, purposeType::DataType) = getPurpose(r, purposeType)

# Counting things.

"""
Counts the number of scenarios within this reservoir.
"""
countScenarios(r::Reservoir) = countScenarios(first(rivers(r)))

"""
Counts the number of scenarios within this reservoir and ensures the whole data structure is consistent.
"""
function countScenariosWithConsistency(r::Reservoir)
  cnt = -1

  for river in rivers(r)
    if cnt < 0
      cnt = countScenariosWithConsistency(river)
    elseif cnt != countScenariosWithConsistency(river)
      error("Inconsistent data in reservoir " * getName(r) * ": different number of scenarios within the reservoir (while checking river " * getName(river) * ").")
    end
  end

  if cnt < 0
    error("Inconsistent data in reservoir " * getName(r) * ": no time series found in the rivers.")
  end

  return cnt
end

"""
Counts the number of time steps within this reservoir.
"""
countTimeSteps(r::Reservoir) = countTimeSteps(first(rivers(r)))

"""
Counts the number of time steps within this reservoir and ensures the whole data structure is consistent.
"""
function countTimeStepsWithConsistency(r::Reservoir)
  cnt = -1

  for river in rivers(r)
    if cnt < 0
      cnt = countTimeStepsWithConsistency(river)
    elseif cnt != countTimeStepsWithConsistency(river)
      error("Inconsistent data in reservoir " * getName(reservoir) * ": different number of time steps within the reservoir " *
            "(while checking river " * getName(river) * ").")
    end
  end

  if cnt < 0
    error("Inconsistent data in reservoir " * getName(reservoir) * ": no time series found in the rivers.")
  end

  return cnt
end

"""
Retrieves the periodicity of this reservoir data.
"""
getPeriod(r::Reservoir) = getPeriod(getRivers(r)[1])

"""
Retrieves the periodicity of this reservoir data and ensures the whole data structure is consistent.
"""
function getPeriodWithConsistency(r::Reservoir)
  period = :Unknown

  for river in rivers(r)
    if period == :Unknown
      period = getPeriodWithConsistency(river)
    elseif period != getPeriodWithConsistency(river)
      error("Inconsistent data in reservoir " * getName(reservoir) * ": different period within the reservoir " *
            "(while checking river " * getName(river) * ").")
    end
  end

  if period == :Unknown
    error("Inconsistent data in reservoir " * getName(reservoir) * ": no time series found in the rivers.")
  end

  return period::Period
end

"""
Retrieves the time series of this reservoir data, i.e. the time points where historical information is known.
"""
getTimePoints(r::Reservoir) = getTimePoints(getRivers(r)[1])

"""
Retrieves the time instants of this reservoir data, i.e. the time points where historical information is known.
This function also ensures the whole data structure is consistent.
"""
function getTimePointsWithConsistency(r::Reservoir)
  ts = nothing

  for river in rivers(r)
    if ts == nothing
      ts = getTimePointsWithConsistency(river)
    elseif ts != getTimePointsWithConsistency(river)
      error("Inconsistent data in reservoir " * getName(reservoir) * ": different time series within the reservoir " *
            "(while checking river " * getName(river) * ").")
    end
  end

  if ts == nothing
    error("Inconsistent data in reservoir " * getName(r) * ": no time series found in the rivers.")
  end

  return ts::Array{DateTime, 1}
end

### Getters for rivers.

# Iterate through all rivers.
rivers(r::Reservoir) = RiverIterator(r)
naturalRivers(r::Reservoir) = NaturalRiverIterator(r)
divertedRivers(r::Reservoir) = DivertedRiverIterator(r)

getNaturalRivers(r::Reservoir) = r.rivers_in
getDivertedRivers(r::Reservoir) = r.rivers_diverted
getRivers(r::Reservoir) = vcat(getNaturalRivers(r), getDivertedRivers(r))

getNaturalRiver(r::Reservoir, naturalRiver::Int) = r.rivers_in[naturalRiver]
getDivertedRiver(r::Reservoir, divertedRiver::Int) = r.rivers_diverted[divertedRiver]

hasRiver(r::Reservoir, river::AbstractString) = length(filter(nr -> getName(nr) == river, getRivers(r))) >= 1
function getRiver(r::Reservoir, river::AbstractString)
  elt = filter(nr -> getName(nr) == river, getRivers(r))
  if length(elt) < 1
    return Nullable{River}()
  else
    return Nullable(elt[1])
  end
end

# Get river name.
getDivertedRiverNames(reservoir::Reservoir) = UTF8String[r.name for r in divertedRivers(reservoir)]
getNaturalRiverNames(reservoir::Reservoir) = UTF8String[r.name for r in naturalRivers(reservoir)]

function getRiverNames(reservoir::Reservoir)
  names = getNaturalRiverNames(reservoir)
  append!(names, getDivertedRiverNames(reservoir))
  return names
end

### Agregators.

# Total inflow.

"""
Compute the total inflow for this reservoir for the given scenario,
i.e. this function agregates all tributaries to the reservoir.
The output is given along time, either as an array (`asTimeSeries=false`, default),
or as a time series object  (`asTimeSeries=true`).
"""
function totalInflow(r::Reservoir, scenario::Int; asTimeSeries::Bool=false)
  if asTimeSeries
    anyRiver = getRivers(r)[1].scenarios[1]
    return TimeArray(timestamp(anyRiver), totalInflow(r, scenario, asTimeSeries=false), colnames(anyRiver), meta(anyRiver))
  else
    return sum([values(riv.scenarios[scenario]) for riv in eachriver(r)])
  end
end

"""
Compute the total inflow for this reservoir,
i.e. this function agregates all tributaries to the reservoir.
The output is given along time, either as an array (`asTimeSeries=false`, default),
or as a time series object  (`asTimeSeries=true`).
"""
function totalInflow(r::Reservoir; kwargs...)
  if countScenarios(r) > 1
    error("This reservoir has multiple scenarios, please indicate which one to use.")
  end
  return totalInflow(r, 1; kwargs...)
end

"""
Compute the total inflow for this reservoir for the given range of scenarios,
i.e. this function agregates all tributaries to the reservoir.
The output is given along time, either as an array (`asTimeSeries=false`, default),
or as a time series object  (`asTimeSeries=true`).
"""
function totalInflow{T}(r::Reservoir, scenarios::Range{T}; kwargs...)
  kwargs = Dict(kwargs)
  if haskey(kwargs, :asTimeSeries) && kwargs[:asTimeSeries]
    return TimeSeries.TimeArray{Float64, 1, DateTime, Array{Float64, 1}}[totalInflow(r, scenario; kwargs...) for scenario in scenarios]
  else
    return Float64[totalInflow(r, scenario; kwargs...) for scenario in scenarios]
  end
end

# No agregator for outputs or purposes: they have a high variability in the way to take them into account.
# Such a method would only make sense for a subset of all purposes.

### Altering the set of scenarios.

"""
Removes the given set of scenarios from the reservoir, of which a modified copy is given.
"""
function removeScenarios(reservoir::Reservoir, toRemove::IntSet)
  nScenarios = countScenarios(reservoir)
  indices = trues(nScenarios)
  for i in toRemove
    indices[i] = false
  end

  newRivers = NaturalRiver[copy(r, scenarios=r.scenarios[indices]) for r in eachnaturalriver(reservoir)]
  newDivertedRivers = DivertedRiver[copy(r, scenarios=r.scenarios[indices]) for r in eachdivertedriver(reservoir)]

  return copy(reservoir, rivers_in=newRivers, rivers_diverted=newDivertedRivers)
end

removeScenarios(r::Reservoir, toRemove::Int) = removeScenarios(r, IntSet(toRemove))
removeScenarios(r::Reservoir, toRemove::Array{Int, 1}) = removeScenarios(r, IntSet(toRemove))

"""
Limits the scenarios in the reservoir to the given time steps.
"""
function keepTimeSteps(reservoir::Reservoir, timeSteps)
  newRivers = NaturalRiver[keepTimeSteps(r, timeSteps) for r in eachnaturalriver(reservoir)]
  newDivertedRivers = DivertedRiver[keepTimeSteps(r, timeSteps) for r in eachdivertedriver(reservoir)]

  return copy(reservoir, rivers_in=newRivers, rivers_diverted=newDivertedRivers)
end

"""
Limits the scenarios in the river to the given time steps.
"""
keepTimeSteps(r::River, timeSteps) = copy(r, scenarios=[s[timeSteps] for s in r.scenarios])

### River iterators.

# Natural rivers.
immutable NaturalRiverIterator
  r::Reservoir
end

eachnaturalriver(r::Reservoir) = NaturalRiverIterator(r)
Base.start(::NaturalRiverIterator) = 1
Base.next(i::NaturalRiverIterator, s) = (i.r.rivers_in[s], s + 1)
Base.done(i::NaturalRiverIterator, s) = s > length(i.r.rivers_in)
Base.eltype(i::NaturalRiverIterator) = NaturalRiver
Base.length(i::NaturalRiverIterator) = length(i.r.rivers_in)

# Diverted rivers.
immutable DivertedRiverIterator
  r::Reservoir
end

eachdivertedriver(r::Reservoir) = DivertedRiverIterator(r)
Base.start(::DivertedRiverIterator) = 1
Base.next(i::DivertedRiverIterator, s) = (i.r.rivers_diverted[s], s + 1)
Base.done(i::DivertedRiverIterator, s) = s > length(i.r.rivers_diverted)
Base.eltype(i::DivertedRiverIterator) = DivertedRiver
Base.length(i::DivertedRiverIterator) = length(i.r.rivers_diverted)

# Finally for both at once.
# Implementation detail: first natural rivers, then diverted.
# TODO: redo with Iterators.chain()?
immutable RiverIterator
  r::Reservoir
  nIn::Int

  function RiverIterator(reservoir::Reservoir)
    return new(reservoir, length(reservoir.rivers_in))
  end
end

eachriver(r::Reservoir) = RiverIterator(r)
Base.start(::RiverIterator) = 1
Base.next(i::RiverIterator, s) = ((s <= i.nIn) ? i.r.rivers_in[s] : i.r.rivers_diverted[s - i.nIn], s + 1)
Base.done(i::RiverIterator, s) = s > i.nIn + length(i.r.rivers_diverted)
Base.eltype(i::RiverIterator) = River
Base.length(i::RiverIterator) = i.nIn + length(i.r.rivers_diverted)

### Iterate over scenarios in a reservoir.
# Provides a copy of the reservoir with a single scenario.
immutable ScenarioReservoirIterator
  r::Reservoir
  max::Int

  function ScenarioReservoirIterator(reservoir::Reservoir)
    return new(reservoir, countScenarios(reservoir))
  end
end

eachscenario(r::Reservoir) = ScenarioReservoirIterator(r)
Base.start(::ScenarioReservoirIterator) = 1
Base.next(i::ScenarioReservoirIterator, s) = (
  copy(i.r,
       rivers_in=NaturalRiver[copy(riv, scenarios=[riv.scenarios[s]]) for riv in eachnaturalriver(i.r)],
       rivers_diverted=DivertedRiver[copy(riv, scenarios=[riv.scenarios[s]]) for riv in eachdivertedriver(i.r)]),
  s + 1
)
Base.done(i::ScenarioReservoirIterator, s) = s > i.max
Base.eltype(i::ScenarioReservoirIterator) = Reservoir
Base.length(i::ScenarioReservoirIterator) = i.max


# dummycamera

dummycamTimer = nothing
exposurefactor = camSettings.exposureTime/12.0 #dummy setting to simulate ideal 12.0 time

# Typically imported functions, made to act like t dummy camera
function getimage!(cam::String,image::Array{UInt8};normalize::Bool=false)
    image .= round.(UInt8,clamp.(rand(UInt8,size(image,1),size(image,2)).*exposurefactor,UInt8(0),UInt8(255)))
    wait(dummycamTimer)
end
function exposure!(cam::String)
    exposure = 12.0 #dummy value
    updateExposureFactor(exposure,camSettings.gain) #dummy brightness change
    return exposure
end
function exposure!(cam::String,exposure::Real)
    updateExposureFactor(exposure,camSettings.gain) #dummy brightness change
end
function gain!(cam::String)
    gain = 0.0 #dummy value
    updateExposureFactor(camSettings.exposureTime,gain) #dummy brightness change
    return gain
end
function gain!(cam::String,gain::Real)
    updateExposureFactor(camSettings.exposureTime,gain) #dummy brightness change
end
function start!(cam::String)
    global dummycamTimer
    dummycamTimer = Timer(0.0,interval=1/camSettings.acquisitionFramerate)
end
function stop!(cam::String)
    global dummycamTimer
    close(dummycamTimer)
    dummycamTimer = nothing
end
function framerate!(cam::String,framerate::Real)
    global dummycamTimer
    dummycamTimer = Timer(0.0,interval=1/camSettings.acquisitionFramerate)
end
function imagedims!(cam::String,dims)
    global camImage
    camImage = Array{UInt8}(undef,dims...)
end
function offsetdims!(cam::String, dims)
    # do something with dims
end
function sensordims(cam::String)
    return (2048,1536)
end
function isrunning(cam::String)
    return camRunning
end

# Helper functions for dummy behaviour
function updateExposureFactor(exposure,gain)
    global exposurefactor
    exposurefactor = (exposure/12.0)*(1.0 + gain)
end

# Main functions
function runCamera()
    global dummycamTimer
    global perfGrabFramerate
    global cam, camImage, camImageFrameBuffer
    global camSettingsLimits

    # Initialize dummy camera
    cam = "dummycam"

    camSettingsLimitsUpdater!(cam,camSettingsLimits)

    start!(cam)

    camImage = Array{UInt8}(undef,camSettings.width,camSettings.height)

    perfGrabTime = time()
    while gui_open
        if isrunning(cam)
            #cim_id, cim_timestamp, cim_exposure = getimage!(camera, previewimage)
            try
                getimage!(cam,camImage,normalize=false)
            catch err
                if err isa EOFError
                    # Do nothing. This happens if the camera is stopped before the camRunning
                    # bool is set due to async
                else
                    rethrow()
                end
            end
            # Loop timing
            perfGrabFramerate = 1/(time() - perfGrabTime)
            perfGrabTime = time()
        end
        yield()
    end
    stop!(cam)
end


"""
Abstract type to represent Electroencephalography (EEG) data.

The following types inherit from the EEG type and can be used to process your data:

- `GeneralEEG`: Used to store data without assumption of any experimental paradigm.
- `SSR`: Used to store data acquired with a steady state response experiment paradigm.
- `TR`: Used to store data acquired with a transient response experiment paradigm.

# Examples
```julia
data = # load your EEG data using for example read_EEG()

samplingrate(data)  # Returns the sampling rate
channelnames(data)  # Returns the channel names
```
"""
abstract type EEG <: NeuroimagingMeasurement end

import Base.show
function Base.show(io::IO, a::EEG)
    time_length = round.(size(a.data, 1) / samplingrate(Float64, a) / 60)
    println(
        io,
        "EEG measurement of $time_length mins with $(size(a.data,2)) channels sampled at $(a.samplingrate)",
    )
end

"""
Type for storing general EEG data without assumption of any experimental paradigm.

# Examples
```julia
s = read_EEG(filename)
s = rereference(s, "Cz")
s = remove_channel!(s, "Cz")
```
"""
mutable struct GeneralEEG <: EEG
    data::Array
    sensors::Array{Sensor}
    triggers::Dict
    system_codes::Dict
    samplingrate::typeof(1.0u"Hz")
    reference_channel::Array{AbstractString,1}
    file_path::AbstractString
    file_name::AbstractString
    processing::Dict
    header::Dict
end
GeneralEEG(args...; kwargs...) = read_EEG(args...; kwargs...)

"""
    samplingrate(t::Type, s::EEG)
    samplingrate(s::EEG)

Return the sampling rate of an EEG type in Hz as the requested type.
If no type is provided, the sampling rate is returned as a floating point number.

# Examples
```julia
s = read_EEG(filename)
samplingrate(s)
```
"""
samplingrate(s::EEG) = s.samplingrate |> u"Hz"
samplingrate(t, s::EEG) = convert(t, s.samplingrate |> u"Hz" |> ustrip)

"""
    times(t::Type, s::EEG)
    times(s::EEG)

Return the times associated with each sample in seconds.

# Examples
```julia
s = read_EEG(filename)
times(s)
```
"""
times(s::EEG) = collect(1:size(data(s), 1)) ./ samplingrate(Float64, s) .* u"s"
times(t, s::EEG) = convert.(t, times(s) |> ustrip)

"""
    channelnames(s::EEG)

Return the names of sensors in EEG measurement.

# Examples
```julia
s = read_EEG(filename)
channelnames(s)
```
"""
channelnames(s::EEG) = labels(s.sensors)

"""
    channelnames(s::EEG, i::Int, l::AbstractString)
    channelnames(s::EEG, l::AbstractVector{AbstractString})

Change the names of `i`th sensors in an EEG measurement `s` to `l`.
Or change the name of all sensors by pass a vector of strings.

# Examples
```julia
s = read_EEG(filename)
channelnames(s, 1, "Fp1")
```
"""
function channelnames(s::EEG, i::Int, l::S) where {S<:AbstractString}

    s.sensors[i].label = l
    return s
end
function channelnames(s::EEG, l::AbstractVector{S}) where {S<:AbstractString}
    @assert length(l) == length(channelnames(s))
    for li = 1:length(l)
        s = channelnames(s, li, l[li])
    end
    return s
end

"""
    data(s::EEG)
    data(s::EEG, channel::AbstractString)
    data(s::EEG, channels::AbstractVector{AbstractString})

Return the data stored in the type object.
This may be useful for integration with custom processing.

# Examples
```julia
s = read_EEG(filename)
data(s)
data(s, "Cz")
data(s, ["Cz", "P2"])
```
"""
data(s::EEG) = s.data
data(s::EEG, channel::AbstractString) = data(s, [channel])
function data(s::EEG, channels::AbstractVector{S}) where {S<:AbstractString}
    return data(s, _channel_indices(s, channels))
end
function data(s::EEG, channels::AbstractVector{S}) where {S<:Int}
    return deepcopy(s.data)[:, channels]
end

"""
    sensors(s::EEG)
    channelnames(s::EEG, l::AbstractVector{AbstractString})

Returns the sensors for an EEG recording.

# Examples
```julia
s = read_EEG(filename)
sensors(s)
```
"""
sensors(s::EEG) = s.sensors
electrodes(s::EEG) = s.sensors

#######################################
#
# EEG type operations
#
#######################################

import Base.hcat
"""
    hcat(a::EEG, b::EEG)

Concatenate two EEG measurements together, effectively creating a single long measurement.

# Examples
```julia
hcat(a, b)
```
"""
function hcat(a::EEG, b::EEG)

    if channelnames(a) != channelnames(b)
        throw(
            ArgumentError(
                string("Channels do not match $(channelnames(a)) != $(channelnames(b))"),
            ),
        )
    end

    if haskey(a.processing, "epochs")
        @warn("Epochs have already been extracted and will no longer be valid")
    end
    if haskey(a.processing, "statistics")
        @warn("Statistics have already been calculated and will no longer be valid")
    end

    @debug(
        "Appending two EEGs with $(size(a.data, 2)) .& $(size(b.data, 2)) channels and lengths $(size(a.data, 1)) $(size(b.data, 1))"
    )

    join_triggers(a, b)
    a.data = [a.data; b.data]

    return a
end

"""
Append the trigger information of one EEG type to another.
Places the trigger information at the end of first file

#### Example

```julia
join_triggers(a, b)
```
"""
function join_triggers(a, b; offset = size(a.data, 1))

    a.triggers["Index"] = [a.triggers["Index"]; (b.triggers["Index"] .+ offset)]
    a.triggers["Code"] = [a.triggers["Code"]; b.triggers["Code"]]
    a.triggers["Duration"] = [a.triggers["Duration"]'; b.triggers["Duration"]']'

    a
end

#######################################
#
# Change reference channels
#
#######################################

"""
    rereference(a::EEG, refChan::Union{AbstractString, Array{AbstractString}}; kwargs...)

Reference data to specified channel(s).

#### Example

```julia
a = rereference(a, "Cz")
# or
a = rereference(a, ["P9", "P10"])
```

"""
function rereference(
    a::EEG,
    refChan::Union{S,Array{S}};
    kwargs...,
) where {S<:AbstractString}

    a.data = rereference(a.data, refChan, channelnames(a))

    a.reference_channel = [refChan]

    return a
end

#######################################
#
# Manipulate channels
#
#######################################

"""
    add_channel(a::EEG, data::Vector, chanLabel::AbstractString)

Add a channel to the EEG type with specified channel names.

# Examples
```julia
s = read_EEG(filename)
new_channel = mean(s.data, 2)
s = add_channel(s, new_channel, "MeanChannelData")
```
"""
function add_channel(a::EEG, data::Vector, chanLabel::AbstractString; kwargs...)

    @info("Adding channel $chanLabel")

    a.data = hcat(a.data, data)
    push!(a.sensors, Electrode(chanLabel, Talairach(NaN, NaN, NaN), Dict()))

    return a
end

"""
    remove_channel!(a::EEG, channelname::AbstractString)
    remove_channel!(a::EEG, channelnames::Array{AbstractString})
    remove_channel!(a::EEG, channelidx::Int)
    remove_channel!(a::EEG, channelidxs::Array{Int})

Remove channel(s) from EEG as specifed by `channelname` or `channelidx`.

# Examples
```julia
a = read_EEG(filename)
remove_channel!(a, ["TP8", "Cz"])
```
"""
function remove_channel!(a::EEG, channel_name::S; kwargs...) where {S<:AbstractString}
    @debug("Removing channel $(channel_name)")
    remove_channel!(a, [channel_name])
end

function remove_channel!(
    a::EEG,
    channel_names::Array{S};
    kwargs...,
) where {S<:AbstractString}
    @debug("Removing channels $(join(channel_names, " "))")
    remove_channel!(a, _channel_indices(a, channel_names))
end

remove_channel!(a::EEG, channel_names::Int; kwargs...) =
    remove_channel!(a, [channel_names]; kwargs...)

function remove_channel!(a::EEG, channel_idx::Array{Int}; kwargs...)

    channel_idx = channel_idx[channel_idx.!=0]

    @debug("Removing channel(s) $channel_idx")
    if any(channel_idx .== 0)
        @warn("Failed to remove a channel")
    end

    keep_idx = [1:size(a.data)[end];]
    for c in sort(channel_idx, rev = true)
        try
            splice!(keep_idx, c)
        catch
            # Nothing
        end
    end

    if haskey(a.processing, "epochs")
        if size(a.processing["epochs"], 3) == size(a.data, 2)
            @debug("Removing channel(s) from epoch data")
            a.processing["epochs"] = a.processing["epochs"][:, :, keep_idx]
        end
    end
    if haskey(a.processing, "sweeps")
        if size(a.processing["sweeps"], 3) == size(a.data, 2)
            @debug("Removing channel(s) from sweep data")
            a.processing["sweeps"] = a.processing["sweeps"][:, :, keep_idx]
        end
    end

    a.data = a.data[:, keep_idx]
    a.sensors = a.sensors[keep_idx]

    return a
end

"""
    keep_channel!(a::EEG, channelname::AbstractString)
    keep_channel!(a::EEG, channelnames::Array{AbstractString})
    keep_channel!(a::EEG, channelidxs::Array{Int})

Remove all channels except those requested from EEG.

# Examples
```julia
a = read_EEG(filename)
keep_channel!(a, ["P8", "Cz"])
```
"""
function keep_channel!(a::EEG, channel_name::AbstractString; kwargs...)
    keep_channel!(a, [channel_name]; kwargs...)
end

function keep_channel!(
    a::EEG,
    channel_names::AbstractVector{S};
    kwargs...,
) where {S<:AbstractString}
    @info("Keeping channel(s) $(join(channel_names, " "))")
    keep_channel!(a, _channel_indices(a, channel_names))
end

function keep_channel!(a::EEG, channel_idx::AbstractVector{Int}; kwargs...)

    remove_channels = [1:size(a.data, 2);]

    channel_idx = sort(channel_idx, rev = true)
    for c in channel_idx
        splice!(remove_channels, c)
    end

    remove_channel!(a, remove_channels; kwargs...)
end

#######################################
#
# Trim channels
#
#######################################

"""
    trim_channel(a::EEG, stop::Int; start::Int=1)

Trim EEG recording by removing data after `stop` specifed samples
and optionally before `start` samples.

# Examples
```julia
s = trim_channel(s, 8192*300, start=8192)
```
"""
function trim_channel(a::EEG, stop::Int; start::Int = 1, kwargs...)

    @info("Trimming $(size(a.data)[end]) channels between $start and $stop")

    a.data = a.data[start:stop, :]

    to_keep = findall((a.triggers["Index"] .>= start) .& (a.triggers["Index"] .<= stop))
    a.triggers["Index"] = a.triggers["Index"][to_keep]
    a.triggers["Duration"] = a.triggers["Duration"][to_keep]
    a.triggers["Code"] = a.triggers["Code"][to_keep]
    a.triggers["Index"] .-= (start - 1)

    to_keep =
        findall((a.system_codes["Index"] .>= start) .& (a.system_codes["Index"] .<= stop))
    a.system_codes["Index"] = a.system_codes["Index"][to_keep]
    a.system_codes["Duration"] = a.system_codes["Duration"][to_keep]
    a.system_codes["Code"] = a.system_codes["Code"][to_keep]
    a.system_codes["Index"] .-= (start - 1)

    return a
end

#######################################
#
# Merge channels
#
#######################################

"""
    merge_channels(a::EEG, merge_Chans::Array{S}, new_name::S) where {S<:AbstractString}
    merge_channels(a::EEG, merge_Chans::S, new_name::S) where {S<:AbstractString}

Average `EEG` channels listed in `merge_Chans` and label the averaged channel as `new_name`.

# Examples
```julia
s = merge_channels(s, ["P6", "P8"], "P68")
```
"""
function merge_channels(
    a::EEG,
    merge_Chans::Array{S},
    new_name::S;
    kwargs...,
) where {S<:AbstractString}

    @debug("Number of original channels: $(length(channelnames(a)))")

    keep_idxs = _channel_indices(a, merge_Chans)

    if sum(keep_idxs .== 0) > 0
        @warn(
            "Could not merge as these channels don't exist: $(join(vec(merge_Chans[keep_idxs .== 0]), " "))"
        )
        keep_idxs = keep_idxs[keep_idxs.>0]
    end

    @info("Merging channels $(join(vec(channelnames(a)[keep_idxs,:]), " "))")
    @debug("Merging channels $keep_idxs")

    a = add_channel(
        a,
        vec(Statistics.mean(a.data[:, keep_idxs], dims = 2)),
        new_name;
        kwargs...,
    )
end

function merge_channels(
    a::EEG,
    merge_Chans::S,
    new_name::S;
    kwargs...,
) where {S<:AbstractString}
    a = merge_channels(a, [merge_Chans], new_name; kwargs...)
end

"""
    read_EEG(fname::AbstractString)
    read_EEG(args...)

Read a file or IO stream and store the data in an `GeneralEEG` type.

# Arguments

- `fname`: Name of the file to be read
- `min_epoch_length`: Minimum epoch length in samples. Shorter epochs will be removed (0)
- `max_epoch_length`: Maximum epoch length in samples. Longer epochs will be removed (0 = all)
- `valid_triggers`: Triggers that are considered valid, others are removed ([1,2])
- `stimulation_amplitude`: Amplitude of stimulation (NaN)
- `remove_first`: Number of epochs to be removed from start of recording (0)
- `max_epochs`: Maximum number of epochs to retain (0 = all)

# Supported file formats

- BIOSEMI (.bdf)
"""
function read_EEG(
    fname::AbstractString;
    valid_triggers::Array{Int} = [1, 2],
    min_epoch_length::Int = 0,
    max_epoch_length::Int = 0,
    remove_first::Int = 0,
    max_epochs::Int = 0,
    kwargs...,
)

    @info("Importing EEG from file: $fname")
    file_path, file_name, ext = fileparts(fname)

    #
    # Read file data
    #

    # Import raw data
    if ext == "bdf"
        data, triggers, system_codes, samplingrate, reference_channel, header =
            import_biosemi(fname; kwargs...)
    else
        warn("File type $ext is unknown")
    end

    # Create electrodes
    elecs = Electrode[]
    for e in header["chanLabels"]
        push!(elecs, Electrode(e, Talairach(NaN * u"m", NaN * u"m", NaN * u"m"), Dict()))
    end

    # Create EEG type
    if unit(samplingrate) == unit(1.0u"Hz")
        #nothing
    else
        samplingrate = samplingrate * 1.0u"Hz"
    end

    a = GeneralEEG(
        data,
        elecs,
        triggers,
        system_codes,
        samplingrate,
        [reference_channel],
        file_path,
        file_name,
        Dict(),
        header,
    )

    #
    # Clean up
    #

    # Remove status channel information
    remove_channel!(a, "Status")

    # Clean epoch index
    a.triggers = clean_triggers(
        a.triggers,
        valid_triggers,
        min_epoch_length,
        max_epoch_length,
        remove_first,
        max_epochs,
    )

    # Try and match sensor names to known locations
    locs = read_elp(joinpath(datadep"BioSemi64Locations", "biosemi64.elp"))
    new_sens, idx = match_sensors(locs, labels(sensors(a)))
    if length(new_sens) == length(sensors(a))
        @debug("Sucsessfully matches location of all sensors, using new locations")
        a.sensors = new_sens
    end

    return a
end

function trigger_channel(a::EEG; kwargs...)

    create_channel(a.triggers, a.data, samplingrate(Float64, a))
end

function system_code_channel(a::EEG; kwargs...)

    create_channel(a.system_codes, a.data, samplingrate(Float64, a))
end

"""
    epoch_rejection(a::EEG; retain_percentage::Number = 0.95, kwargs...)

Reject epochs such that `retain_percentage` is retained.
"""
function epoch_rejection(a::EEG; retain_percentage::Number = 0.95, kwargs...)

    a.processing["epochs"] =
        epoch_rejection(a.processing["epochs"], retain_percentage; kwargs...)

    return a
end


module REFPROPwrap
"""
Initializing by calling the SETUPdll with required fluid name

currently support only pure fluids with one component
"""
function refprop_init(fluid::String)
    ncmax = Ref{Cint}(2)
    refpropcharlength  = Ref{Cint}(255)
    filepathlength  = Ref{Cint}(255)
    lengthofreference = Ref{Cint}(3)
    errormessagelength =  Ref{Cint}(255)
    hfld_length = Ref{Cint}(refpropcharlength[]*ncmax[])
    hf = Array{UInt8}(undef,refpropcharlength[]*ncmax[])
    hrf = Array{UInt8}(undef, lengthofreference[])
    herr = Array{UInt8}(undef, errormessagelength[])
    hfmix = Array{UInt8}(undef,refpropcharlength[])
    hf_name = "/home/aaditya/Desktop/julia/REFPROP/FLUIDS/"*fluid*".FLD\0"
    hfmix_name = "/home/aaditya/Desktop/julia/REFPROP/FLUIDS/HMX.BNC"
    hrf_name = "DEF"
    herr_name = "Ok"
    hf[1:length(hf_name)] = collect(hf_name)
    hfmix[1:length(hfmix_name)] = collect(hfmix_name)
    hrf[1:length(hrf_name)] = collect(hrf_name)
    herr[1:length(herr_name)] = collect(herr_name)
    i = Ref{Cint}(1)
    ierr = Ref{Cint}(0)
    ccall((:SETUPdll,"/home/aaditya/Desktop/julia/librefprop.so"),Cvoid,(Ref{Cint},Ptr{UInt8},Ptr{UInt8},Ptr{UInt8},Ref{Cint},Ptr{UInt8},Ref{Cint},Ref{Cint},Ref{Cint},Ref{Cint}), i, hf, hfmix, hrf, ierr, herr, hfld_length, refpropcharlength, lengthofreference, errormessagelength)
    return String(herr)
end #refprop_init
"""
sample function to call the properties using pressure and enthalpy as the inputs
"""
function refprop_call(pre::Float64, enthal::Float64)
    #initializing the arrays and variables
    ncmax = Ref{Cint}(2)
    ierr = Ref{Cint}(0)
    herr = Array{UInt8}(undef, 255)

    x = Array{Cdouble}(zeros(ncmax[])); xliq = Array{Cdouble}(zeros(ncmax[])); xvap = Array{Cdouble}(zeros(ncmax[])); f = Array{Cdouble}(zeros(ncmax[]));

    wm = Ref{Cdouble}(0)
    vs = Ref{Cdouble}(0)
    ccall((:WMOLdll,"/home/aaditya/Desktop/julia/librefprop.so"), Cvoid, (Ptr{Cdouble}, Ref{Cdouble}), x, wm)
    p = Ref{Cdouble}(pre)
    e = Ref{Cdouble}(0)
    h = Ref{Cdouble}(enthal*wm[]/1000.0)
    s = Ref{Cdouble}(0)
    cv = Ref{Cdouble}(0)
    cp = Ref{Cdouble}(0)
    hjt = Ref{Cdouble}(0)
    temp = Ref{Cdouble}(0)
    rhop = Ref{Cdouble}(0) #(rho/wm[])
    rhol = Ref{Cdouble}(0)
    rhov = Ref{Cdouble}(0)
    qual = Ref{Cdouble}(0)

    #ccall((:THERMdll,"/home/aaditya/Desktop/julia/librefprop.so"),Cvoid,(Ref{Cdouble}, Ref{Cdouble},Ptr{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble},Ref{Cdouble}), temp, rhop, x, p, e, h, s, cv, cp, vs, hjt);
    ccall((:PHFLSHdll,"/home/aaditya/Desktop/julia/librefprop.so"),Cvoid,(Ref{Cdouble}, Ref{Cdouble}, Ptr{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ptr{Cdouble}, Ptr{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cint}, Ptr{UInt8}), p,h,x,temp,rhop,rhol,rhov,xliq,xvap,qual,e,s,cv,cp,vs,ierr,herr);

    return p[], h[]/wm[]*1000.0, temp[], rhop[]*wm[], rhol[]*wm[], rhov[]*wm[], qual[],  e[]/wm[]*1000.0, s[]/wm[]*1000.0, cv[]/wm[]*1000.0, cp[]/wm[]*1000.0, vs[]
end #refprop_call

export refprop_inti, refprop_call

end # module


module jLRS

# standard library
using Printf

# third party libraries
using Distributions
using HDF5
using Proj4
using Glob
using AstroLib
using DataFrames
using Dates
using LinearAlgebra
using Random
using SatelliteToolbox
using SQLite
using SparseArrays
using StaticArrays

jLRS_seed = 123456
# Seed the generator at the beginning. So the *same* code with the *same* inputs
# will always produce the same numbers. If even one parameter changes, we will
# get different results!
Random.seed!(jLRS_seed)

# Type definitions required by the rest of the module
include("types.jl")

# Common math-y functions
include("misc.jl")

# Solar data required for many calculations
include("solar.jl")

# Surface samplings (BRDF, land cover, etc.)
include("surface.jl")

# SIF-related functions
include("SIF.jl")

# Functions for OCO-type
include("OCO.jl")

# Functions for Geostationary-type
include("Geostationary.jl")

# Functions for SwathNadir-type
include("SwathNadirOrbiter.jl")

# Main calculations
include("light_response.jl")

# Plotting tools
include("plotting.jl")

@info "You have loaded jLRS - thanks!"

end # module


# To do:
#
# * Test that CoPs are in region
# * Test that friction cone constraints are satisfied
# * Move utilities to appropriate places
# * better CoM kinematic constraints
# * Initial condition modification

## Packages
using CentroidalTrajOpt
using LinearAlgebra

using MultilinearOpt
using Polyhedra: hrep, polyhedron
using CentroidalTrajOpt: extrude

using RigidBodyDynamics
using RigidBodyDynamics.PDControl
using RigidBodyDynamics.Contact
using MeshCatMechanisms
using GeometryTypes: Point, GLNormalMesh
using QPControl
using QPWalkingControl
using AtlasRobot
using RigidBodySim
using Rotations
using OSQP
using MathOptInterface

using QPControl.Trajectories: PointTrajectory
using QPWalkingControl: PosePlan, set_pose_plan!

import Polyhedra
import Blink

using MeshCat: RGBA

const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities

include("optimizers.jl")
include("util.jl")

function create_atlas()
    urdf = AtlasRobot.urdfpath()
    mechanism = parse_urdf(urdf, floating=true, remove_fixed_tree_joints=false)
    link_colors = Dict(map(body -> string(body) => RGBA(0.7f0, 0.7f0, 0.7f0, 0.3f0), bodies(mechanism)))
    visuals = URDFVisuals(AtlasRobot.urdfpath(); package_path=[AtlasRobot.packagepath()], link_colors=link_colors)
    remove_fixed_tree_joints!(mechanism)
    foot_points = AtlasRobot.foot_contact_points(mechanism)
    sole_frames = AtlasRobot.add_sole_frames!(mechanism)
    nominal_state = MechanismState(mechanism)
    AtlasRobot.setnominal!(nominal_state; kneebend=1.6)
    floating_joint = first(joints(mechanism))
    Δz = translation(transform_to_root(nominal_state, first(values(sole_frames))))[3]
    configuration(nominal_state, floating_joint)[end] -= Δz
    setdirty!(nominal_state)
    pelvis = findbody(mechanism, "pelvis");
    mechanism, nominal_state, foot_points, sole_frames, floating_joint, pelvis, visuals
end

function create_environment()
    region_data = ContactRegion{Float64}[]
    push!(region_data, ContactRegion(
            AffineMap(one(RotMatrix{3}), zero(SVector{3})),
            0.7,
            0.0,
            Float64[1 0; 0 1; -1 0; 0 -1],
            0.2 * ones(4)
    ))
    push!(region_data, ContactRegion(
            AffineMap(one(RotMatrix{3}) * RotXYZ(0.1, -0.2, 0.3), SVector(0.75, 0.3, 0.2)),
            0.7,
            0.0,
            Float64[1 0; 0 1; -1 0; 0 -1],
            0.2 * ones(4)
    ))
    push!(region_data, ContactRegion(
            AffineMap(one(RotMatrix{3}) * RotXYZ(-0.1, 0.2, 0.3), SVector(0.85, 1.15, 0.1)),
            0.7,
            0.0,
            Float64[1 0; 0 1; -1 0; 0 -1],
            0.2 * ones(4)
    ))
    push!(region_data, ContactRegion(
        AffineMap(one(RotMatrix{3}) * RotXYZ(0.0, 0.0, 0.0), SVector(1.7, 1.2, 0.2)),
        0.7,
        0.0,
        Float64[1 0; 0 1; -1 0; 0 -1],
        0.2 * ones(4)
    ))
    region_data
end

function create_contact_model(
        mechanism::Mechanism,
        foot_points::AbstractDict{BodyID, <:AbstractVector{<:Point3D}},
        region_data::Vector{<:ContactRegion}; region_offset) # TODO
    contact_model = ContactModel()
    normal_model = hunt_crossley_hertz(; k=500e3)
    k_tangential = 7e3
    b_tangential = 250.#2 * sqrt(k_tangential * mass(mechanism) / 10)
    world_frame = root_frame(mechanism)
    foot_collision_elements = CollisionElement[]
    for (bodyid, points) in foot_points
        body = findbody(mechanism, bodyid)
        for point in points
            push!(foot_collision_elements, CollisionElement(body, point.frame, Point(point.v)))
        end
    end
    push!(contact_model, foot_collision_elements)
    debug_with_flat_ground = false
    if debug_with_flat_ground
        geometry = HalfSpace(SVector(0., 0., 1.), 0.0)
        frame = world_frame
        group = CollisionElement[CollisionElement(root_body(mechanism), frame, geometry)]
        push!(contact_model, group)
        tangential_model = ViscoelasticCoulombModel(10.0, k_tangential, b_tangential)
        contact_force_model = SplitContactForceModel(normal_model, tangential_model)
        set_contact_force_model!(contact_model, foot_collision_elements, group, contact_force_model)
    else
        region_thickness = 0.3
        for (i, region) in enumerate(region_data)
            b̄ = region.b .+ flipsign.(region_offset, region.b) # FIXME
            geometry = Contact.HRep(extrude(hrep(region.A, b̄), region_thickness))
            frame = CartesianFrame3D("region_$i")
            add_frame!(root_body(mechanism), Transform3D(frame, world_frame, region.transform))
            element = CollisionElement(root_body(mechanism), frame, geometry)
            group = CollisionElement[element]
            push!(contact_model, group)
            tangential_model = ViscoelasticCoulombModel(region.μ, k_tangential, b_tangential)
            contact_force_model = SplitContactForceModel(normal_model, tangential_model)
            set_contact_force_model!(contact_model, foot_collision_elements, group, contact_force_model)
        end
    end
    return contact_model
end

function create_controller(
        mechanism::Mechanism,
        contact_body_ids::AbstractVector{BodyID},
        floating_joint::Joint,
        foot_points::AbstractDict{BodyID, <:AbstractVector{<:Point3D}},
        μ::Number,
        pelvis::RigidBody,
        state0::MechanismState,
        nominal_state::MechanismState,
        com_trajectory,
        pose_plans::AbstractDict{BodyID, <:PosePlan}
    )
    # Low level controller
    optimizer = OSQP.Optimizer(verbose=false, eps_abs=1e-5, eps_rel=1e-5, max_iter=5000, adaptive_rho_interval=25)
    pelvis_gains_xy = QPWalkingControl.critically_damped_gains(100.0)
    pelvis_gains_z = QPWalkingControl.critically_damped_gains(25.0)
    lowlevel = MomentumBasedController{4}(mechanism, optimizer, floatingjoint = floating_joint);
    for bodyid in contact_body_ids
        points = foot_points[bodyid]
        body = findbody(mechanism, bodyid)
        for point in points
            normal = FreeVector3D(default_frame(body), 0.0, 0.0, 1.0)
            contact = addcontact!(lowlevel, body, point, normal, μ)
            contact.maxnormalforce[] = 1e6 # TODO
            contact.weight[] = 1e-3
        end
    end

    # Linear momentum controller
    world_frame = root_frame(mechanism)
    com_gains = QPWalkingControl.critically_damped_gains(10.0)
    linear_momentum_controller = PDCoMController(com_gains, PointTrajectory(world_frame, com_trajectory), mass(mechanism))

    # State machine
    contacts = Dict(BodyID(body) => contact for (body, contact) in lowlevel.contacts)
    state_machine = CoMTrackingStateMachine(mechanism, contacts)
    for (bodyid, plan) in pose_plans
        set_pose_plan!(state_machine, bodyid, plan)
    end

    # High level controller
    pelvis_gains = PDGains(
        Diagonal(SVector(pelvis_gains_xy.k, pelvis_gains_xy.k, pelvis_gains_z.k)),
        Diagonal(SVector(pelvis_gains_xy.d, pelvis_gains_xy.d, pelvis_gains_z.d))
    )
    HumanoidQPController(lowlevel, pelvis, nominal_state,
        state_machine, collect(values(state_machine.end_effector_controllers)), linear_momentum_controller, pelvisgains = pelvis_gains)
end

## Robot setup
mechanism, state0, foot_points, sole_frames, floating_joint, pelvis, visuals = create_atlas()
contact_body_ids = sort(collect(keys(sole_frames)), by=x -> x.value) # establishes order once and for all

## Environment
region_data = create_environment()

## Parameters
g = mechanism.gravitational_acceleration.v
inside_foot_radius = Inf
outside_foot_radius = 0.0
for (bodyid, points) in foot_points
    body = findbody(mechanism, bodyid)
    sole_frame = sole_frames[bodyid]
    for point in points
        point_sole_frame = fixed_transform(body, point.frame, sole_frame) * point
        dist = norm(point_sole_frame.v)
        global inside_foot_radius = min(inside_foot_radius, dist) # TODO: not technically correct
        global outside_foot_radius = max(outside_foot_radius, dist)
    end
end
max_cop_distance = inside_foot_radius# - 0.01
min_com_to_contact_distance = 0.8
max_com_to_contact_distance = 1.05
min_inter_contact_distance = 2 * outside_foot_radius
region_offset = outside_foot_radius + 0.04

## Collision setup
contact_model = create_contact_model(mechanism, foot_points, region_data, region_offset=region_offset)

## Load results
load = false
if load
    if @isdefined result
        @error "result already defined."
    else
        import StaticArrays, StaticUnivariatePolynomials, QPControl.Trajectories # for field types of CentroidalTrajectoryResult
        result = load_result(joinpath(ENV["CENTROIDAL_TRAJ_OPT_RESULT_DIR"], "results", "four_regions_kinematically_infeasible.jld2"));
    end
end

## Initial conditions
c0 = center_of_mass(state0).v
ċ0 = center_of_mass_velocity(state0).v
p0s = map(contact_body_ids) do bodyid
    sole_frame = sole_frames[bodyid]
    translation(transform_to_root(state0, sole_frame))
end
contacts0 = map(p0s) do p0 # TODO
    region_data[1] => p0
end

## Final conditions
cf = nothing
# cf = c0# + SVector(0.8, 0.5, 0.05)
# cf = c0 + SVector(0.05, 0.0, -0.05)

# check_kinematic_constraints_satisfied(contacts0, c0, max_com_to_contact_distance, min_inter_contact_distance)

## Create visualizer
using MeshCat
newvis = false
if newvis || (!@isdefined vis) || isempty(vis.core.scope.pool.connections)
    vis = Visualizer()
    # open(vis)
    # wait(vis)
end
delete!(vis)

## Centroidal trajectory visualization
cvis = CentroidalTrajectoryVisualizer(vis, region_data, norm(g), length(contact_body_ids))

## Robot visualization
mvis = MechanismVisualizer(mechanism, visuals, vis)
copyto!(mvis, state0)
# for sole_frame in values(sole_frames)
#     setelement!(mvis, sole_frame, 0.1)
# end

## Environment visualization
@time setelement!(mvis, contact_model)#, MeshLambertMaterial(color=RGBA(0.9, 0.9, 0.5, 0.95)))

## GUI
gui = GUI(mvis)
if isempty(vis.core.scope.pool.connections)
    open(gui)
    wait(gui)
end
sleep(1)

## Optimizer
# optimizer_factory = baron_optimizer_factory()
optimizer_factory = scip_optimizer_factory()
# optimizer_factory = gurobi_optimizer_factory()
# optimizer_factory = cplex_optimizer_factory()
# optimizer_factory = ipopt_optimizer_factory()

## Problem
problem = CentroidalTrajectoryProblem(optimizer_factory, region_data, c0, ċ0, contacts0;
    cf=cf, g=g,
    min_Δt=0.6, max_Δt=1.5,
    # min_Δt=0.6, max_Δt=0.6,
    max_cop_distance=max_cop_distance, min_com_to_contact_distance=min_com_to_contact_distance,
    max_com_to_contact_distance=max_com_to_contact_distance, min_inter_contact_distance=min_inter_contact_distance,
    num_pieces=10, com_degree=3,
    objective_type=ObjectiveTypes.FEASIBILITY);
    # objective_type=ObjectiveTypes.MAX_HEIGHT);

disallow_jumping!(problem)

## Final region constraint
fix.(problem.z_vars[problem.pieces(problem.pieces[end]), problem.regions(problem.regions[end])], 1.0)

## Additional kinematic constraints
# TODO: use interval arithmetic to get rough boxes

if optimizer_factory.constructor == SCIP.Optimizer
    set_optimize_hook(problem.model, scip_optimize_hook)
end
if optimizer_factory.constructor == Ipopt.Optimizer
    @info "Reformulating binary variables -> complementarity constraints"
    complementarity_reformulate!(problem.model)
end

relax = optimizer_factory.constructor == Gurobi.Optimizer || optimizer_factory.constructor == CPLEX.Optimizer
if relax
    @info "Relaxing bilinearities."
    relaxbilinear!(problem.model, method=:Logarithmic1D, disc_level=50, constraints_to_skip=problem.friction_cone_quadratic_constraints)
end

## Feasibility
result = CentroidalTrajOpt.solve!(problem);

if backend(problem.model).optimizer.model isa SCIP.Optimizer
    mscip = backend(problem.model).optimizer.model.mscip
    # SCIP.print_statistics(mscip)
    SCIP.print_heuristic_statistics(mscip)
    SCIP.SCIPprintVersion(mscip, Libc.FILE(0))
end

reoptimize = false
if reoptimize
    if backend(problem.model).optimizer.model isa SCIP.Optimizer
        mscip = backend(problem.model).optimizer.model.mscip
        SCIP.SCIPsetEmphasis(mscip, SCIP.SCIP_PARAMEMPHASIS_OPTIMALITY, true)
    end
    @objective problem.model Max sum(problem.c_vars[problem.c_coeffs(3)])
    result = CentroidalTrajOpt.solve!(problem; ignore_optimize_hook=true)
end

## Result visualization
set_com_trajectory!(cvis, result)
set_state!(cvis, result, 0.0)
plan_animation = Animation(cvis, result, show_cones=true)
setanimation!(vis, plan_animation)

## Optimality
# Setting objective function straight away somehow prevents SCIP from finding a feasible point, so:
# SCIP.SCIPsetEmphasis(problem.model.moi_backend.optimizer.model.mscip, SCIP.SCIP_PARAMEMPHASIS_OPTIMALITY, true);
# set_objective!(problem.model, MOI.MIN_SENSE,
# set_objective(problem.model, MOI.MAX_SENSE, sum(problem.z_vars))
# result = solve!(problem);
# set_objective(problem.model, MOI.MIN_SENSE, sum(problem.Δts))
# result = CentroidalTrajOpt.solve!(problem);

if optimizer_factory.constructor == BARON.Optimizer
    @info "BARON problem file: $(backend(problem.model).optimizer.model.optimizer.inner.problem_file_name)"
    @info "BARON result file: $(backend(problem.model).optimizer.model.optimizer.inner.result_file_name)"
end

## Tests
using Test
c = result.center_of_mass
ċ = map_subfunctions(derivative, c)
c̈ = map_subfunctions(derivative, ċ)
rs = result.centers_of_pressure
fs = result.contact_forces
# τns = result.contact_normal_torques
ps = result.contact_positions
ns = normals(problem)

# Times
@test issorted(result.break_times)
@test all(x -> x >= 0, result.break_times)

T = last(result.break_times)

# Initial and final conditions
@test c(0) ≈ c0 atol=1e-12
@test ċ(0) ≈ ċ0 atol=1e-12
# @test ċ(T) ≈ zeros(3) atol=1e-12
# @test c̈(T) ≈ zeros(3) atol=1e-12

max_constraint_violation = 0.0

for t in range(0, T, length=100)
    ftot = sum(fs[contact](t) for contact in problem.contacts.val)

    # Dynamics
    @test c̈(t) ≈ g + ftot atol=1e-3

    # Torque about CoM
    # τ = sum((rs[j](t) - c(t)) × fs[j](t) + n * τns[j](t) for j in problem.contacts.val)
    τ = sum((rs[j](t) - c(t)) × fs[j](t) for j in problem.contacts.val)

    constraint_violation = norm(τ) / norm(g)
    global max_constraint_violation = max(max_constraint_violation, constraint_violation)

    if !relax
        @test τ ≈ zeros(3) atol=1e-4
    end

    # Friction cones
    for j in problem.contacts.val
        zs = result.contact_indicators[j](t)
        @test sum(zs) <= 1 + 1e-6
        region_idx = findfirst(isequal(1), zs)

        f = fs[j]
        # τn = τns[j]
        p = ps[j]
        r = rs[j]

        f_t = f(t)
        @test norm(p(t) - r(t)) < max_cop_distance + 1e-5

        if region_idx == nothing
            @test f_t ≈ zero(f_t) atol=1e-4
            # @test τn_t ≈ zero(τn_t) atol=1e-4
        else
            region = region_data[region_idx]
            n = ns[problem.regions(region_idx)]
            fn_t = f_t ⋅ n
            μ = region.μ
            @test fn_t >= -1e-7
            @test norm(f_t - n * fn_t) <= μ * fn_t + 1e-6
            # @test -μrot * fn_t <= τn_t
            # @test τn_t <= μrot * fn_s
        end
    end
end
@info "Max constraint violation: $max_constraint_violation"

# Continuity around breaks
for t in result.break_times
    @test c(t - 1e-8) ≈ c(t + 1e-8) atol=1e-3
    @test ċ(t - 1e-8) ≈ ċ(t + 1e-8) atol=1e-3
end

## Simulation
simulate = true
if simulate
    # Pose plans
    pose_plans = Dict{BodyID, PosePlan{Float64}}()
    for (i, bodyid) in enumerate(contact_body_ids)
        pose0 = transform_to_root(state0, sole_frames[bodyid])
        pose_plans[bodyid] = PosePlan(pose0, result.contact_positions[i], result.contact_indicators[i], region_data)
    end

    # Controller
    μ_control = 0.7
    nominal_state = state0 # TODO
    controller = create_controller(mechanism, contact_body_ids, floating_joint, foot_points, μ_control, pelvis,
        state0, nominal_state, result.center_of_mass, pose_plans);

    state = MechanismState(mechanism)
    copyto!(state, state0)
    Δt = 1 / 500
    pcontroller = PeriodicController(similar(velocity(state)), Δt, controller)
    damping = JointDamping{Float64}(mechanism, AtlasRobot.urdfpath())
    dynamics = Dynamics(
        mechanism,
        pcontroller,#SumController(similar(velocity(state)), (pcontroller, damping));
        contact_model=contact_model)
    # callback = CallbackSet(RealtimeRateLimiter(poll_interval=pi / 100), )
    callback = CallbackSet(gui; max_fps=30)
    T = last(result.break_times)
    tspan = (0., T + 2)
    contact_state = SoftContactState(contact_model)
    odeproblem = ODEProblem(dynamics, (state, contact_state), tspan)#, callback=callback)

    # simulate
    @time sol = RigidBodySim.solve(odeproblem, Tsit5(), abs_tol = 1e-8, dt = 1e-6, dtmax=1e-3);
    sim_animation = Animation(mvis, sol)

    # setanimation!(vis, plan_animation);
    # setanimation!(vis, sim_animation);
    setanimation!(vis, merge(plan_animation, sim_animation));
end

## Set up visualizer for creating videos
video = false
if video
    open(vis, Blink.Window(Dict(:width => 1280, :height => 960, :useContentSize => true)))
    sleep(1)
    setanimation!(vis, merge(plan_animation, sim_animation));
end

## Save results
save = false
if save
    CentroidalTrajOpt.Serialization.save_result(result)
end

set_state!(cvis, result, 0.0)
copyto!(mvis, state0)