# Copyright 2017 Stephen L. Smith and Frank Imeson
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#########  Input Parsing ############################

"""
parse the problem instance file given as a text file with the following format:
N: int
M: int
Symmetric: true/false
Triangle: true/false
1 5
5 1 8 9 10 11
8 2 3 4 6 7 12 13 14
DISTANCE_MATRIX

Note: the line 5 1 8 9 10 11 shows that the second set contains 5 vertices:
1,8,9,10,11.

TSPLIB Parser defined by:
  http://comopt.ifi.uni-heidelberg.de/software/TSPLIB95/TSPFAQ.html
"""
function read_file(filename)
	  if !isfile(filename)
		    error("the problem instance  ", filename, " does not exist")
	  end

    # Setup
    INF = 9999
    RE_INT = r"-?\d+"
    RE_DEC = r"-?\d+\.\d+"
    RE_NUMBER = r"-?\d+\.?\d*"

    parse_state = "UNKNOWN_FORMAT"
    data_type = ""
    data_format = ""
    dist = zeros(Int64, 0, 0)
    sets = Any[]
    vid00 = vid01 = 1
    coords = Any[]
    set_data = Int64[]
    num_vertices = -1
    num_sets = -1

    # parse
    s = open(filename, "r")
    for line in readlines(s)
        line = strip(line)
        # debug
#        println(line)
#        println(parse_state)

        # auto format select
        if parse_state == "UNKNOWN_FORMAT"
            if occursin(r"^\s*NAME\s*:\s*\w+", uppercase(line))
                parse_state = "TSPLIB_HEADER"
            elseif occursin(r"^\s*N\s*:\s*\w+", uppercase(line))
                parse_state = "SIMPLE_HEADER"
            end
        end

        # Parse Setup
        if parse_state == "TSPLIB_HEADER"
            value = strip(split(line,":")[end])
            if occursin(r"^\s*NAME\s*:\s*\w+", uppercase(line))
            elseif occursin(r"^\s*TYPE\s*:\s*\w+\s*$", uppercase(line))
            elseif occursin(r"^\s*DIMENSION\s*:\s*\d+\s*$", uppercase(line))
                num_vertices = parse(Int64, value)
                dist = zeros(Int64, num_vertices, num_vertices)
            elseif occursin(r"^\s*GTSP_SETS\s*:\s*\d+\s*$", uppercase(line))
                num_sets = parse(Int64, value)
            elseif occursin(r"^\s*EDGE_WEIGHT_TYPE\s*:\s*\w+\s*$", uppercase(line))
                data_type = value
                data_format = value
            elseif occursin(r"^\s*EDGE_WEIGHT_FORMAT\s*:\s*\w+\s*$", uppercase(line))
                if data_type == "EXPLICIT"
                    data_format = value
                end
            elseif occursin(r"^\s*EDGE_WEIGHT_SECTION\s*:?\s*$", uppercase(line))
                parse_state = "TSPLIB_MATRIX_DATA"
            elseif occursin(r"^\s*NODE_COORD_SECTION\s*:?\s*$", uppercase(line))
                parse_state = "TSPLIB_COORD_DATA"
            end

        # Parse matrix data
        elseif parse_state == "TSPLIB_MATRIX_DATA"
            if occursin(r"^[\d\se+-\.]+$", line)
                for x in split(line)
                    cost = parse(Int64, x)
                    # tested
                    if data_format == "FULL_MATRIX"
                        dist[vid00, vid01] = cost
                        vid01 += 1
                        if vid01 > num_vertices
                            vid00 += 1
                            vid01  = 1
                        end
                    # tested
                    elseif data_format == "LOWER_DIAG_ROW"
                        dist[vid00, vid01] = cost
                        dist[vid01, vid00] = cost
                        vid01 += 1
                        if vid01 > vid00
                            vid00 += 1
                            vid01  = 1
                        end
                    # not tested
                    elseif data_format == "LOWER_ROW"
                        println("WARNING: Not tested")

                        if vid00 == 0 && vid01 == 0
                            vid00 = 1
                        end
                        dist[vid00, vid01] = cost
                        dist[vid01, vid00] = cost
                        vid01 += 1
                        if vid01 >= vid00
                            vid00 += 1
                            vid01  = 0
                        end
                    # tested
                    elseif data_format == "UPPER_DIAG_ROW"
                        dist[vid00, vid01] = cost
                        dist[vid01, vid00] = cost
                        vid01 += 1
                        if vid01 > num_vertices
                            vid00 += 1
                            vid01  = vid00
                        end
                    # tested
                    elseif data_format == "UPPER_ROW"
                        if vid00 == 1 && vid01 == 1
                            vid01 = 2
                        end
                        dist[vid00, vid01] = cost
                        dist[vid01, vid00] = cost
                        vid01 += 1
                        if vid01 > num_vertices
                            vid00 += 1
                            vid01  = vid00+1
                        end
                    end
                end
            elseif occursin(r"^\s*DISPLAY_DATA_SECTION\s*:?\s*$", uppercase(line))
                parse_state = "TSPLIB_DISPLAY_DATA"
            elseif occursin(r"^\s*GTSP_SET_SECTION\s*:?\s*$", uppercase(line))
                parse_state = "TSPLIB_SET_DATA"
            end

        # Parse display data
        elseif parse_state == "TSPLIB_DISPLAY_DATA"
            if occursin(r"^\s*GTSP_SET_SECTION\s*:?\s*$", uppercase(line))
                parse_state = "TSPLIB_SET_DATA"
            end

        # Parse coord data
        elseif parse_state == "TSPLIB_COORD_DATA"
            if occursin(r"\s*\d+\s*", uppercase(line))
                coord = [parse(Float64, x) for x = split(line)[2:end]]
                push!(coords, coord)
            elseif occursin(r"^\s*GTSP_SET_SECTION\s*:?\s*$", uppercase(line))
                parse_state = "TSPLIB_SET_DATA"
            end

        # Parse set data
        elseif parse_state == "TSPLIB_SET_DATA"
            if occursin(r"\d+", uppercase(line))
                for x = split(line)
                    push!(set_data, parse(Int64,x))
                end
            elseif occursin(r"^\s*EOF\s*$", uppercase(line))
                parse_state = "TSPLIB"
            end

        # Parse header (simple)
        elseif parse_state == "SIMPLE_HEADER"
            if occursin(r"^\s*N\s*:\s*\w+", uppercase(line))
                value = strip(split(strip(line),":")[end])
                num_vertices = parse(Int64, value)
                dist = zeros(Int64, num_vertices, num_vertices)
            elseif occursin(r"^\s*M\s*:\s*\d+\s*$", uppercase(line))
                value = strip(split(strip(line),":")[end])
                num_sets = parse(Int64, value)
                parse_state = "SIMPLE_SETS"
            end

        # Parse set data (simple)
        elseif parse_state == "SIMPLE_SETS"
            if occursin(r"^[\d\se+-\.]+$", line)
                sid = parse(Int64, split(line)[1])
                set = [parse(Int64, x) for x in split(line)[2:end]]
                push!(sets, set)
                if sid == num_sets
                    parse_state = "SIMPLE_MATRIX"
                end
            end

        # Parse set data (simple)
        elseif parse_state == "SIMPLE_MATRIX"
            if occursin(r"^[\d\se+-\.]+$", line)
                for x in split(line)
                    cost = parse(Int64, x)
                    dist[vid00, vid01] = cost
                    vid01 += 1
                    if vid01 > num_vertices
                        vid00 += 1
                        vid01  = 1
                    end
                end
            else
                parse_state = "SIMPLE"
            end
        end

    end
    close(s)
    if occursin(r"TSPLIB", parse_state)
        parse_state = "TSPLIB"
    end

    # Convert coordinate data to matrix data
    if parse_state == "TSPLIB" && data_type != "EXPLICIT"
        # tested
        if data_format == "EUC_2D"
            for vid00 in 1:num_vertices
                for vid01 in 1:num_vertices
                    if vid00 == vid01
                        dist[vid00, vid01] = INF
                    else
                        dx = coords[vid00][1] - coords[vid01][1]
                        dy = coords[vid00][2] - coords[vid01][2]
                        cost = sqrt(dx^2 + dy^2)
                        dist[vid00, vid01] = nint(cost)
                    end
                end
            end
        # not tested
        elseif data_format == "MAN_2D"
            println("Warning: MAN_2D not tested")
            for vid00 in 1:num_vertices
                for vid01 in 1:num_vertices
                    if vid00 == vid01
                        dist[vid00, vid01] = INF
                    else
                        dx = coords[vid00][1] - coords[vid01][1]
                        dy = coords[vid00][2] - coords[vid01][2]
                        cost = abs(dx) + abs(dy)
                        dist[vid00, vid01] = nint(cost)
                    end
                end
            end
        # not working...
        elseif data_format == "GEO"
            RRR = 6378.388
            PI = 3.141592
            DEBUG01 = false
            TSPLIB_GEO = true
            for vid00 in 1:num_vertices
                d, m    = degree_minutes(coords[vid00][1])
                lat00   = PI * (d + 5.0 * m / 3.0) / 180.0
                d, m    = degree_minutes(coords[vid00][2])
                long00  = PI * (d + 5.0 * m / 3.0) / 180.0
                for vid01 in 1:num_vertices
                    d, m    = degree_minutes(coords[vid01][1])
                    lat01   = PI * (d + 5.0 * m / 3.0) / 180.0
                    d, m    = degree_minutes(coords[vid01][2])
                    long01  = PI * (d + 5.0 * m / 3.0) / 180.0
                    if vid00 == vid01
                        dist[vid00, vid01] = INF
                    else
                        if TSPLIB_GEO
                            q1      = cos(long00 - long01)
                            q2      = cos(lat00 - lat01)
                            q3      = cos(lat00 + lat01)
                            cost    = RRR * acos(0.5 * ((1.0 + q1) * q2 - (1.0 - q1) * q3)) + 1.0
                            dist[vid00, vid01] = floor(Int64, cost)
                        else
                            # http://andrew.hedges.name/experiments/haversine/
                            R       = 6373
                            dlat    = abs(coords[vid01][1] - coords[vid00][1])
                            dlong   = abs(coords[vid01][2] - coords[vid00][2])
                            a       = (sind(dlat/2))^2 + cosd(coords[vid00][1])*cosd(coords[vid01][1])*(sind(dlong/2))^2
                            c       = 2*atan2(sqrt(a), sqrt(1-a))
                            cost    = R * c
                            dist[vid00, vid01] = floor(Int64, cost)
                        end
                        if DEBUG01
                            println("lat00 = ", coords[vid00][1], ", long00 = ", coords[vid00][2], ", lat01 = ", coords[vid01][1], ", long01 = ", coords[vid01][2], ", dist = ", floor(Int64, cost))
                            DEBUG01 = false
                        end
                    end
                end
            end
        # tested
        elseif data_format == "ATT"
            for vid00 in 1:num_vertices
                for vid01 in 1:num_vertices
                    if vid00 == vid01
                        dist[vid00, vid01] = INF
                    else
                        dx = coords[vid00][1] - coords[vid01][1]
                        dy = coords[vid00][2] - coords[vid01][2]
                        r  = sqrt((dx^2 + dy^2)/10.0)
                        cost = ceil(r)
                        dist[vid00, vid01] = nint(cost)
                    end
                end
            end
        # tested: not sure if working or not
        elseif data_format == "CEIL_2D"
            for vid00 in 1:num_vertices
                for vid01 in 1:num_vertices
                    if vid00 == vid01
                        dist[vid00, vid01] = INF
                    else
                        dx = coords[vid00][1] - coords[vid01][1]
                        dy = coords[vid00][2] - coords[vid01][2]
                        cost = sqrt(dx^2 + dy^2)
                        dist[vid00, vid01] = ceil(cost)
                    end
                end
            end
        else
            error("coordinate type $data_format not supported")
        end
    end

    # construct sets
    if parse_state == "TSPLIB"
        i = 1
        sid00 = 1
        set = Int64[]
        set_data = set_data[2:end]
        while i <= length(set_data)
            x = set_data[i]
            if x == -1
                push!(sets, set)
                set = Int64[]
                i += 1 # skip set id
            else
                push!(set, x)
            end
            i += 1
        end
        if num_sets != length(sets)
            error("number of sets doesn't match set size")
        end
    end

    if length(sets) <= 1
        error("must have more than 1 set")
    end

	  membership = findmember(num_vertices, sets)

    return num_vertices, num_sets, sets, dist, membership
end

""" Computing degrees and minutes for GEO instances """
function degree_minutes(num)
	if num > 0
		deg = floor(Int64, num)
		return deg, num - deg
	else
		deg = ceil(Int64, num)
		return deg, num - deg
	end
end

#####################################################
#####################################################
######### Output Printing ###########################

""" print the main parameter settings """
function print_params(param::Dict{Symbol,Any})
	if param[:print_output] > 0
		println("\n", "--------- Problem Data ------------")
		println("Instance Name      : ", param[:problem_instance])
	    println("Number of Vertices : ", param[:num_vertices])
	    println("Number of Sets     : ", param[:num_sets])
		println("Initial Tour       : ", (param[:init_tour] == "rand" ?
				"Random" : "Random Insertion"))
		println("Maximum Removals   : ", param[:max_removals])
		println("Trials             : ", param[:cold_trials])
		println("Restart Attempts   : ", param[:warm_trials])
		println("Rate of Adaptation : ", param[:epsilon])
		println("Prob of Reopt      : ", param[:prob_reopt])
		println("Maximum Time       : ", param[:max_time])
		println("Tour Budget        : ", (param[:budget] == typemin(Int64) ?
				"None" : param[:budget]))
		println("-----------------------------------\n")
	end
end

"""
Print the powers in an easy-to-read format for debugging the adaptive weights
"""
function print_powers(powers)
	println("--  Printing powers -- ")
	println("Insertions:")
	for power in powers["insertions"]
		println(power.name, " ", power.value, ": ", power.weight)
	end

	println("\n Removals:")
	for power in powers["removals"]
		println(power.name, " ", power.value, ": ", power.weight)
	end
	print("\n")

	println("\n Noises:")
	for power in powers["noise"]
		println(power.name, " ", power.value, ": ", power.weight)
	end
	print("\n")
end

"""print statement at the beginning of a cold trial"""
function print_cold_trial(count::Dict{Symbol,Real}, param::Dict{Symbol,Any}, best::Tour)
	if param[:print_output] == 2
		println("\n||--- trial ", count[:cold_trial],
		" --- initial cost ", best.cost, " ---||")
	end
end

"""print details at the end of each warm trial"""
function print_warm_trial(count::Dict{Symbol,Real}, param::Dict{Symbol,Any},
							best::Tour, iter_count::Int)
	if param[:print_output] == 2
		println("-- ", count[:cold_trial], ".", count[:warm_trial],
		" - iterations ", iter_count, ":  ", "cost ", best.cost)
	end
end

""" print best cost so far """
function print_best(count::Dict{Symbol,Real}, param::Dict{Symbol,Any},
							best::Tour, lowest::Tour, init_time::Float64)
	if param[:print_output] == 1 && time() - count[:print_time] > param[:print_time]
		count[:print_time] = time()
		println("-- trial ", count[:cold_trial], ".", count[:warm_trial], ":",
				"  Cost = ", min(best.cost, lowest.cost),
				"  Time = ", round(count[:print_time] - init_time, digits=1), " sec")

	elseif (param[:print_output] == 3 && time() - count[:print_time] > 0.5) ||
		param[:budget_met] || param[:timeout]

		count[:print_time] = time()
		if param[:warm_trials] > 0
			progress = (count[:cold_trial] - 1)/param[:cold_trials] +
	 	   			 (count[:warm_trial])/param[:warm_trials]/param[:cold_trials]
		else
			progress = (count[:cold_trial] - 1)/param[:cold_trials]
		end
		tcurr = round(count[:print_time] - init_time, digits=1)
		cost = min(best.cost, lowest.cost)
		progress_bar(param[:cold_trials], progress, cost, tcurr)
	end
end

""" a string representing the progress bar """
function progress_bar(trials, progress, cost, time)
	ticks, trials_per_bar, total_length = 6, 5, 31
	progress == 1.0 && (progress -= 0.0001)
	n = floor(Int64, progress * trials/trials_per_bar)
	start_number = n * trials_per_bar
	trials_in_bar = min(trials_per_bar, trials - start_number)

	progress_in_bar = (progress * trials - start_number)/trials_in_bar
	bar_length = min(total_length - 1, (trials - start_number) * ticks)

	progress_bar = "|"
	for i=1:total_length
		if i == bar_length + 1
			progress_bar *= "|"
		elseif i > bar_length + 1
			progress_bar *= " "
		elseif i % ticks == 1
			progress_bar *= string(start_number + ceil(Int64, i / ticks))
		elseif i <= ceil(Int64, bar_length * progress_in_bar)
			progress_bar *= "="
		else
			progress_bar *= " "
		end
	end
	print(" ", progress_bar, "  Cost = ", cost, "  Time = ", time, " sec      \r")
end

"""print tour summary at end of execution"""
function print_summary(lowest::Tour, timer::Float64, member::Array{Int64,1},
						param::Dict{Symbol,Any})
	if param[:print_output] == 3 && !param[:timeout] && !param[:budget_met]
		progress_bar(param[:cold_trials], 1.0, lowest.cost, round(timer, digits=1))
	end
	if param[:print_output] > -1
		if (param[:print_output] > 0 || param[:output_file] == "None")
			println("\n\n", "--------- Tour Summary ------------")
			println("Cost              : ", lowest.cost)
			println("Total Time        : ", round(timer, digits=2), " sec")
			println("Solver Timeout?   : ", param[:timeout])
			println("Tour is Feasible? : ", tour_feasibility(lowest.tour, member,
																	param[:num_sets]))
			order_to_print = (param[:output_file] == "None" ?
					lowest.tour : "printed to " * param[:output_file])
			println("Output File       : ",  param[:output_file])
			println("Tour Ordering     : ",  order_to_print)
			println("-----------------------------------")
		end
		if param[:output_file] != "None"
			s = open(param[:output_file], "w")
			write(s, "Problem Instance : ", param[:problem_instance], "\n")
			write(s, "Vertices         : ", string(param[:num_vertices]), "\n")
			write(s, "Sets             : ", string(param[:num_sets]), "\n")
			write(s, "Comment          : To avoid ~0.5sec startup time, use the Julia REPL\n")
			write(s, "Host Computer    : ", gethostname(), "\n")
			write(s, "Solver Time      : ", string(round(timer, digits=3)), " sec\n")
			write(s, "Tour Cost        : ", string(lowest.cost), "\n")
			write(s, "Tour             : ", string(lowest.tour))
			close(s)
		end
	end
end

"""
init function defined by TSPLIB
  http://comopt.ifi.uni-heidelberg.de/software/TSPLIB95/TSPFAQ.html
"""
function nint(x::Float64)
    return floor(Int64, x + 0.5)
end


using Mill: ArrayNode, BagNode, TreeNode, WeightedBagNode
using SparseArrays, DataFrames

import Mill: sparsify, mapdata
import LearnBase: nobs

@testset "creating bags" begin
    k = [2, 2, 2, 1, 1, 3]
    @test Mill.bags(k).bags == [1:3,4:5,6:6]
end

metadata = fill("metadata", 4)
a = BagNode(ArrayNode(rand(3,4)),[1:4], metadata)
b = BagNode(ArrayNode(rand(3,4)),[1:2,3:4], metadata)
c = BagNode(ArrayNode(rand(3,4)),[1:1,2:2,3:4], metadata)
d = BagNode(ArrayNode(rand(3,4)),[1:4,0:-1], metadata)
wa = WeightedBagNode(ArrayNode(rand(3,4)),[1:4], rand(1:4, 4), metadata)
wb = WeightedBagNode(ArrayNode(rand(3,4)),[1:2,3:4], rand(1:4, 4), metadata)
wc = WeightedBagNode(ArrayNode(rand(3,4)),[1:1,2:2,3:4], rand(1:4, 4), metadata)
wd = WeightedBagNode(ArrayNode(rand(3,4)),[1:4,0:-1], rand(1:4, 4), metadata)
e = ArrayNode(rand(2, 2))

f = TreeNode((wb,b))
g = TreeNode((c, wc))
h = TreeNode((wc,c))
i = TreeNode((
              b,
              TreeNode((
                        b,
                        BagNode(
                                BagNode(
                                        ArrayNode(rand(2,4)),
                                        [1:1, 2:2, 3:3, 4:4]
                                       ),
                                [1:3, 4:4]
                               )
                       ))
             ))
k = TreeNode((a = wb, b = b))
l = TreeNode((a = wc, b = c))

@testset "testing nobs" begin
    @test nobs(a) == nobs(wa) == 1
    @test nobs(b) == nobs(wb) == 2
    @test nobs(c) == nobs(wc) == 3
    @test nobs(d) == nobs(wd) == 2
    @test nobs(e) == 2
    @test nobs(f) == nobs(wb) == nobs(b) == nobs(k)
    @test nobs(g) == nobs(c) == nobs(wc) == nobs(l)
    @test nobs(h) == nobs(wc) == nobs(c)
    @test nobs(i) == nobs(b)
end

@testset "testing ArrayNode hcat and vcat" begin
    @test all(catobs(e, e).data .== hcat(e.data, e.data) .== reduce(catobs, [e,e]).data)
    @test all(hcat(e, e).data .== hcat(e.data, e.data))
    @test all(vcat(e, e).data .== vcat(e.data, e.data))
    x = ArrayNode(randn(2,3),rand(3))
    @test typeof(catobs(x,x[0:-1])) .== ArrayNode{Array{Float64,2},Array{Float64,1}}
    @test typeof(reduce(catobs, [x, x[0:-1]])) .== ArrayNode{Array{Float64,2},Array{Float64,1}}
    @test all(cat(e, e, dims = ndims(e)).data .== hcat(e.data, e.data))
end

@testset "testing BagNode hcat" begin
    @test all(catobs(a, b, c).data.data .== hcat(a.data.data, b.data.data, c.data.data))
    @test all(reduce(catobs, [a, b, c]).data.data .== hcat(a.data.data, b.data.data, c.data.data))
    @test catobs(a, b, c).bags.bags == vcat(a.bags, b.bags, c.bags).bags
    @test reduce(catobs, [a, b, c]).bags.bags == vcat(a.bags, b.bags, c.bags).bags
    @test all(catobs(c, a).data.data .== hcat(c.data.data, a.data.data))
    @test all(reduce(catobs, [c, a]).data.data .== hcat(c.data.data, a.data.data))
    @test catobs(c, a).bags.bags == vcat(c.bags, a.bags).bags
    @test reduce(catobs, [c, a]).bags.bags == vcat(c.bags, a.bags).bags
    @test all(catobs(a, c).data.data .== hcat(a.data.data, c.data.data))
    @test all(reduce(catobs, [a, c]).data.data .== hcat(a.data.data, c.data.data))
    @test catobs(a, c).bags.bags == vcat(a.bags, c.bags).bags
    @test reduce(catobs, [a, c]).bags.bags == vcat(a.bags, c.bags).bags
    @test all(catobs(a, d).data.data .== hcat(a.data.data, d.data.data))
    @test all(reduce(catobs, [a, d]).data.data .== hcat(a.data.data, d.data.data))
    @test catobs(a, d).bags.bags == vcat(a.bags, d.bags).bags
    @test reduce(catobs, [a, d]).bags.bags == vcat(a.bags, d.bags).bags
    @test all(catobs(d, a).data.data .== hcat(d.data.data, a.data.data))
    @test all(reduce(catobs, [d, a]).data.data .== hcat(d.data.data, a.data.data))
    @test catobs(d, a).bags.bags == vcat(d.bags, a.bags).bags
    @test reduce(catobs, [d, a]).bags.bags == vcat(d.bags, a.bags).bags
    @test all(catobs(d).data.data .== hcat(d.data.data))
    @test all(reduce(catobs, [d]).data.data .== hcat(d.data.data))
    @test catobs(d).bags.bags == vcat(d.bags).bags
    @test reduce(catobs, [d]).bags.bags == vcat(d.bags).bags
    @test all(cat(a, b, dims = ndims(a)).data.data .== hcat(a.data.data, b.data.data))
end

@testset "testing WeightedBagNode hcat" begin
    @test all(catobs(wa, wb, wc).data.data .== hcat(wa.data.data, wb.data.data, wc.data.data))
    @test all(reduce(catobs, [wa, wb, wc]).data.data .== hcat(wa.data.data, wb.data.data, wc.data.data))
    @test catobs(wa, wb, wc).bags.bags == vcat(wa.bags, wb.bags, wc.bags).bags
    @test reduce(catobs, [wa, wb, wc]).bags.bags == vcat(wa.bags, wb.bags, wc.bags).bags
    @test all(catobs(wc, wa).data.data .== hcat(wc.data.data, wa.data.data))
    @test all(reduce(catobs, [wc, wa]).data.data .== hcat(wc.data.data, wa.data.data))
    @test catobs(wc, wa).bags.bags == vcat(wc.bags, wa.bags).bags
    @test reduce(catobs, [wc, wa]).bags.bags == vcat(wc.bags, wa.bags).bags
    @test all(catobs(wa, wc).data.data .== hcat(wa.data.data, wc.data.data))
    @test all(reduce(catobs, [wa, wc]).data.data .== hcat(wa.data.data, wc.data.data))
    @test catobs(wa, wc).bags.bags == vcat(wa.bags, wc.bags).bags
    @test reduce(catobs, [wa, wc]).bags.bags == vcat(wa.bags, wc.bags).bags
    @test all(catobs(wa, wd).data.data .== hcat(wa.data.data, wd.data.data))
    @test all(reduce(catobs, [wa, wd]).data.data .== hcat(wa.data.data, wd.data.data))
    @test catobs(wa, wd).bags.bags == vcat(wa.bags, wd.bags).bags
    @test reduce(catobs, [wa, wd]).bags.bags == vcat(wa.bags, wd.bags).bags
    @test all(catobs(wd, wa).data.data .== hcat(wd.data.data, wa.data.data))
    @test all(reduce(catobs, [wd, wa]).data.data .== hcat(wd.data.data, wa.data.data))
    @test catobs(wd, wa).bags.bags == vcat(wd.bags, wa.bags).bags
    @test reduce(catobs, [wd, wa]).bags.bags == vcat(wd.bags, wa.bags).bags
    @test all(catobs(wd).data.data .== hcat(wd.data.data))
    @test all(reduce(catobs, [wd]).data.data .== hcat(wd.data.data))
    @test catobs(wd).bags.bags == vcat(wd.bags).bags
    @test reduce(catobs, [wd]).bags.bags == vcat(wd.bags).bags
end

@testset "testing hierarchical hcat on tree nodes" begin
    @test all(catobs(f, h).data[1].data.data .== hcat(wb.data.data, wc.data.data))
    @test all(reduce(catobs, [f, h]).data[1].data.data .== hcat(wb.data.data, wc.data.data))
    @test all(catobs(f, h).data[2].data.data .== hcat(b.data.data, c.data.data))
    @test all(reduce(catobs, [f, h]).data[2].data.data .== hcat(b.data.data, c.data.data))
    @test all(catobs(f, h, f).data[1].data.data .== hcat(wb.data.data, wc.data.data, wb.data.data))
    @test all(reduce(catobs, [f, h, f]).data[1].data.data .== hcat(wb.data.data, wc.data.data, wb.data.data))

    @test all(catobs(k, l).data[1].data.data .== hcat(wb.data.data, wc.data.data))
    @test all(catobs(k, l).data[2].data.data .== hcat(b.data.data, c.data.data))
    @test nobs(catobs(k,l)) == nobs(k) + nobs(l)

    # different tuple length
    @test_throws Exception catobs(f, i)
    @test_throws Exception reduce(catobs, [f, i])
    @test_throws Exception catobs(g, i)
    @test_throws Exception reduce(catobs, [g, i])
    # different content
    @test_throws Exception catobs(f, g)
    @test_throws Exception reduce(catobs, [f, g])
end

@testset "testing BagNode indexing" begin
    @test all(a[1].data.data .== a.data.data)
    @test a[1].bags.bags == [1:4]
    @test all(b[1:2].data.data .== b.data.data)
    @test b[1:2].bags.bags == [1:2,3:4]
    @test all(b[2].data.data .== b.data.data[:,3:4])
    @test b[2].bags.bags == [1:2]
    @test all(b[1].data.data .== b.data.data[:,1:2])
    @test b[1].bags.bags == [1:2]
    @test all(c[1].data.data .== c.data.data[:,1:1])
    @test c[1].bags.bags == [1:1]
    @test all(c[[1,2]].data.data .== c.data.data[:,1:2])
    @test c[[1,2]].bags.bags == [1:1,2:2]
    @test all(c[[2,1]].data.data .== c.data.data[:,[2,1]])
    @test c[[2,1]].bags.bags == [1:1,2:2]
    @test all(d[[2,1]].data.data .== d.data.data)
    @test d[[2,1]].bags.bags == [0:-1,1:4]
    @test all(d[1:2].data.data .== d.data.data)
    @test d[1:2].bags.bags == [1:4,0:-1]
    @test ismissing(d[2].data)
    @test d[2].bags.bags == [0:-1]
    @test isempty(a[2:1].bags.bags)
    @test ismissing(a[2:1].data)
end

@testset "testing WeightedBagNode indexing" begin
    @test all(wa[1].data.data .== wa.data.data)
    @test wa[1].bags.bags == [1:4]
    @test all(wb[1:2].data.data .== wb.data.data)
    @test wb[1:2].bags.bags == [1:2,3:4]
    @test all(wb[2].data.data .== wb.data.data[:,3:4])
    @test wb[2].bags.bags == [1:2]
    @test all(wb[1].data.data .== wb.data.data[:,1:2])
    @test wb[1].bags.bags == [1:2]
    @test all(wc[1].data.data .== wc.data.data[:,1:1])
    @test wc[1].bags.bags == [1:1]
    @test all(wc[[1,2]].data.data .== wc.data.data[:,1:2])
    @test wc[[1,2]].bags.bags == [1:1,2:2]
    @test all(wc[[2,1]].data.data .== wc.data.data[:,[2,1]])
    @test wc[[2,1]].bags.bags == [1:1,2:2]
    @test all(wd[[2,1]].data.data .== wd.data.data)
    @test wd[[2,1]].bags.bags == [0:-1,1:4]
    @test all(wd[1:2].data.data .== wd.data.data)
    @test wd[1:2].bags.bags == [1:4,0:-1]
    @test ismissing(wd[2].data)
    @test wd[2].bags.bags == [0:-1]
end

@testset "testing nested ragged array" begin
    x = BagNode(ArrayNode(rand(3,10)),[1:2,3:3,0:-1,4:5,6:6,7:10])
    y = BagNode(x,[1:2,3:3,4:5,6:6])
    @test all(y[1].data.data.data .== x.data.data[:,1:3])
    @test y[1].data.bags.bags == [1:2,3:3]
    @test all(y[1:2].data.data.data .== x.data.data[:,1:3])
    @test y[1:2].data.bags.bags == [1:2,3:3,0:-1]
    @test all(y[2:3].data.data.data .== x.data.data[:,4:6])
    @test y[2:3].data.bags.bags == [0:-1,1:2,3:3]
end

@testset "testing TreeNode" begin
    x = TreeNode((ArrayNode(rand(3,2)),ArrayNode(rand(3,2)),ArrayNode(randn(3,2))))
    y = TreeNode((ArrayNode(rand(3,2)),ArrayNode(rand(3,2)),ArrayNode(randn(3,2))))
    @test all(catobs(x,y).data[1].data .== hcat(x.data[1].data,y.data[1].data))
    @test all(reduce(catobs, [x,y]).data[1].data .== hcat(x.data[1].data,y.data[1].data))
    @test all(catobs(x,y).data[2].data .== hcat(x.data[2].data,y.data[2].data))
    @test all(reduce(catobs, [x,y]).data[2].data .== hcat(x.data[2].data,y.data[2].data))
    @test all(catobs(x,y).data[3].data .== hcat(x.data[3].data,y.data[3].data))
    @test all(reduce(catobs, [x,y]).data[3].data .== hcat(x.data[3].data,y.data[3].data))
    @test all(cat(x,y, dims = ndims(x)).data[3].data .== hcat(x.data[3].data,y.data[3].data))

    @test all(k[1].data[1].data.data .==  wb[1].data.data)
    @test all(k[2].data[1].data.data .==  wb[2].data.data)
    @test all(l[2:3].data[1].data.data .==  wc[2:3].data.data)
end

@testset "testing sparsify" begin
    @test typeof(sparsify(zeros(10,10),0.05)) <: SparseMatrixCSC
    @test typeof(sparsify(randn(10,10),0.05)) <: Matrix
    @test typeof(sparsify(randn(10),0.05)) <: Vector
end

@testset "testing sparsify and mapdata" begin
    x = TreeNode((TreeNode((ArrayNode(randn(5,5)), ArrayNode(zeros(5,5)))), ArrayNode(zeros(5,5))))
    xs = mapdata(i -> sparsify(i,0.05), x)
    @test typeof(xs.data[2].data) <: SparseMatrixCSC
    @test typeof(xs.data[1].data[2].data) <: SparseMatrixCSC
    @test all(xs.data[1].data[1].data .== x.data[1].data[1].data)
end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

using ACE.SphericalHarmonics: SHBasis, index_y
using StaticArrays
using JuLIP: AbstractCalculator, Atoms, JVec
using JuLIP.Potentials: SitePotential, SZList, ZList
using NeighbourLists: neigs

import JuLIP, JuLIP.MLIPs

export PIPotential

# from pibasis:
# StandardEvaluator, DAGEvaluator, graphevaluator, standardevaluator

"""
`struct PIPotential` : specifies a PIPotential, which is basically defined
through a PIBasis and its coefficients
"""
mutable struct PIPotential{T, NZ, TPI, TEV} <: SitePotential
   pibasis::TPI
   coeffs::NTuple{NZ, Vector{T}}
   dags::NTuple{NZ, CorrEvalGraph{T, Int}}
   evaluator::TEV
end

cutoff(V::PIPotential) = cutoff(V.pibasis)

==(V1::PIPotential, V2::PIPotential) =
      (V1.pibasis == V2.pibasis) && (V1.coeffs == V2.coeffs)

# TODO: this doesn't feel right ... should be real(T)?
fltype(::PIPotential{T}) where {T} = real(T)

zlist(V::PIPotential) = zlist(V.pibasis)

graphevaluator(V::PIPotential) =
   PIPotential(V.pibasis, V.coeffs, V.dags, DAGEvaluator())

standardevaluator(V::PIPotential) =
   PIPotential(V.pibasis, V.coeffs, V.dags, StandardEvaluator())

maxorder(V::PIPotential) = maxorder(V.pibasis)

# ------------------------------------------------------------
#   Initialisation code
# ------------------------------------------------------------

combine(basis::PIBasis, coeffs::AbstractVector) =
      PIPotential(basis, identity.(collect(coeffs)))

# assemble from basis with global coeff vector
function PIPotential(basis::PIBasis, coeffs::Vector{<: Number})
   coeffs_t = ntuple(iz0 -> coeffs[basis.inner[iz0].AAindices], numz(basis))
   return PIPotential(basis, coeffs_t)
end

# assemble from basis with coeff vectors separated into individual species
function PIPotential(basis::PIBasis, coeffs_t::Tuple)
   dags = ntuple(iz0 -> _getdagfrombasis(basis.inner[iz0], coeffs_t[iz0]), numz(basis))
   return PIPotential(basis, coeffs_t, dags, DAGEvaluator())
end

function _getdagfrombasis(inner, c)
   nodes = inner.dag.nodes
   vals = inner.dag.vals
   dag = DAG.CorrEvalGraph(nodes, zeros(eltype(c), length(nodes)),
                           inner.dag.num1, inner.dag.numstore)
   for (i, idx) in enumerate(vals)
      if idx != 0  # idx = 0 means this is an auxiliary basis with 0 coefficient!
         dag.vals[i] = c[idx]
      end
   end
   return dag
end

# ------------------------------------------------------------
#   experimental sparsification code

function deletezeros(V::PIPotential)
   idx0 = 0
   innernew = []
   coeffs = []
   for iz0 = 1:length(V.coeffs)
      # zeros
      Iz = findall(V.coeffs[iz0] .== 0)
      # non-zeros
      Inz = findall(V.coeffs[iz0] .!= 0)
      # inverse mapping from old to new (local) indices
      invInz = Dict{Int, Int}()
      for (i, val) in enumerate(Inz)
         invInz[val] = i
      end
      # extract the inner basis to start rebuilding it
      inner = V.pibasis.inner[iz0]     # old
      iAA2iA = inner.iAA2iA[Inz, :]    # new
      # construct the new b2iAA mapping
      b2iAA = Dict{ACE.PIBasisFcn, Int}()   # new
      for (b, iAA) in inner.b2iAA
         if haskey(invInz, iAA)
            b2iAA[b] = invInz[iAA]
         end
      end
      # create the new inner basis
      push!(innernew,
            ACE.InnerPIBasis( inner.orders[Inz],
                              iAA2iA,
                              b2iAA,
                              inner.b2iA,
                              (idx0+1):(idx0+length(Inz)),
                              inner.z0,
                              ACE.DAG.CorrEvalGraph{Int, Int}() ) )
      # and the new coefficients
      push!(coeffs, V.coeffs[iz0][Inz])
   end
   pibasis = PIBasis( V.pibasis.basis1p, V.pibasis.zlist,
                      tuple( innernew... ), V.pibasis.evaluator )
   return standardevaluator(PIPotential( pibasis, tuple(coeffs...) ))
end

# ------------------------------------------------------------
#   FIO code
# ------------------------------------------------------------

write_dict(V::PIPotential) = Dict(
      "__id__" => "ACE_PIPotential",
     "pibasis" => write_dict(V.pibasis),
      "coeffs" => [ write_dict.(V.coeffs)... ] )

#    if ntests > 0
#       tests = ACE.Random.rand_nhd(Nat, J::ScalarBasis, species = :X)
#       Pr = V.pibasis.basis1p.J
#       for ntest = 1:5
#
#          r = ACE.rand_radial(V.pibasis.basis1p.J)
#          Pr = evaluate(B3.J, r)
#          push!(rtests, Dict("r" => r, "Pr" => Pr))
#       end
#
#    end
#    return D
# end

read_dict(::Val{:SHIPs_PIPotential}, D::Dict; tests = true) =
   read_dict(Val{:ACE_PIPotential}(), D)

read_dict(::Val{:ACE_PIPotential}, D::Dict; tests = true) =
   PIPotential( read_dict(D["pibasis"]),
                tuple( read_dict.( D["coeffs"] )... ) )

# function test_compat(V::PIPotential, rtests, tests)
#
# end

# ------------------------------------------------------------
#   Dispatching the Evaluation codes
# ------------------------------------------------------------

alloc_temp(V::PIPotential, maxN::Integer) =  alloc_temp(V.evaluator, V, maxN)
alloc_temp_d(V::PIPotential, maxN::Integer) =  alloc_temp_d(V.evaluator, V, maxN)
evaluate!(tmp, V::PIPotential, args...) = evaluate!(tmp, V, V.evaluator, args...)
evaluate_d!(dEs, tmp, V::PIPotential, args...) = evaluate_d!(dEs, tmp, V, V.evaluator, args...)

# ------------------------------------------------------------
#   Standard Evaluation code
# ------------------------------------------------------------

# TODO: generalise the R, Z, allocation
alloc_temp(::StandardEvaluator, V::PIPotential{T}, maxN::Integer) where {T} =
   (
      R = zeros(JVec{real(T)}, maxN),
      Z = zeros(AtomicNumber, maxN),
      tmp_pibasis = alloc_temp(V.pibasis, maxN),
   )

# compute one site energy
function evaluate!(tmp, V::PIPotential, ::StandardEvaluator,
                   Rs::AbstractVector{JVec{T}},
                   Zs::AbstractVector{<:AtomicNumber},
                   z0::AtomicNumber) where {T}
   iz0 = z2i(V, z0)
   A = evaluate!(tmp.tmp_pibasis.A, tmp.tmp_pibasis.tmp_basis1p,
                 V.pibasis.basis1p, Rs, Zs, z0)
   inner = V.pibasis.inner[iz0]
   c = V.coeffs[iz0]
   Es = zero(T)
   @inbounds for iAA = 1:length(inner)
      Esi = one(Complex{T}) * c[iAA]    # TODO: OW - NASTY!!!
      for α = 1:inner.orders[iAA]
         Esi *= A[inner.iAA2iA[iAA, α]]
      end
      Es += real(Esi)
   end
   return Es
end

# TODO: generalise the R, Z, allocation
alloc_temp_d(::StandardEvaluator, V::PIPotential{T}, N::Integer) where {T} =
      (
      dAco = zeros(fltype(V.pibasis),
                   maximum(length(V.pibasis.basis1p, iz) for iz=1:numz(V))),
       tmpd_pibasis = alloc_temp_d(V.pibasis, N),
       dV = zeros(JVec{real(T)}, N),
        R = zeros(JVec{real(T)}, N),
        Z = zeros(AtomicNumber, N),
        dAAt = zero(MVector{maxorder(V), fltype(V.pibasis.basis1p)})
      )

# compute one site energy
function evaluate_d!(dEs, tmpd, V::PIPotential, ::StandardEvaluator,
                     Rs::AbstractVector{<: JVec{T}},
                     Zs::AbstractVector{AtomicNumber},
                     z0::AtomicNumber
                     ) where {T}
   iz0 = z2i(V, z0)
   basis1p = V.pibasis.basis1p
   tmpd_1p = tmpd.tmpd_pibasis.tmpd_basis1p
   Araw = tmpd.tmpd_pibasis.A

   # stage 1: precompute all the A values
   A = evaluate!(Araw, tmpd_1p, basis1p, Rs, Zs, z0)

   # stage 2: compute the coefficients for the ∇A_{klm} = ∇ϕ_{klm}
   dAco = tmpd.dAco
   c = V.coeffs[iz0]
   inner = V.pibasis.inner[iz0]
   dAAt = tmpd.dAAt
   fill!(dAco, 0)
   for iAA = 1:length(inner)
      c_ = c[iAA]
      ord = inner.orders[iAA]
      Afwd = one(eltype(A))
      dAAt[1] = 1
      for α = 1:ord-1
         Afwd *= A[inner.iAA2iA[iAA, α]]
         dAAt[α+1] = Afwd
      end
      Abwd = one(eltype(A))
      for α = ord:-1:2
         Abwd *= A[inner.iAA2iA[iAA, α]]
         dAAt[α-1] *= Abwd
      end
      for α = 1:ord
         iAα = inner.iAA2iA[iAA, α]
         dAco[iAα] += c_ * dAAt[α]
      end

      # for α = 1:inner.orders[iAA]
      #    CxA_α = c[iAA]
      #    for β = 1:inner.orders[iAA]
      #       if β != α
      #          CxA_α *= A[inner.iAA2iA[iAA, β]]
      #       end
      #    end
      #    iAα = inner.iAA2iA[iAA, α]
      #    dAco[iAα] += CxA_α
      # end
   end

   # stage 3: get the gradients
   fill!(dEs, zero(JVec{T}))
   dAraw = tmpd.tmpd_pibasis.dA
   for (iR, (R, Z)) in enumerate(zip(Rs, Zs))
      evaluate_d!(Araw, dAraw, tmpd_1p, basis1p, R, Z, z0)
      iz = z2i(basis1p, Z)
      zinds = basis1p.Aindices[iz, iz0]
      for iA = 1:length(basis1p, iz, iz0)
         dEs[iR] += real(dAco[zinds[iA]] * dAraw[zinds[iA]])
      end
   end

   return dEs
end

# ------- Evaluation code

_maxstore(V::PIPotential) = maximum( dag.numstore for dag in V.dags )

alloc_temp(::DAGEvaluator, V::PIPotential{T}, maxN::Integer) where {T} =
   (
   R = zeros(JVec{real(T)}, maxN),
   Z = zeros(AtomicNumber, maxN),
   tmp_basis1p = alloc_temp(V.pibasis.basis1p),
   AA = zeros(fltype(V.pibasis.basis1p), _maxstore(V)),
   A = alloc_B(V.pibasis.basis1p)
   )

function evaluate!(tmp, V::PIPotential, ::DAGEvaluator, Rs, Zs, z0)
   AAdag = tmp.AA
   A = tmp.A
   iz0 = z2i(V, z0)
   dag = V.dags[iz0]
   nodes = dag.nodes
   vals = dag.vals
   @assert length(A) >= dag.num1
   @assert length(AAdag) >= dag.numstore

   evaluate!(A, tmp.tmp_basis1p, V.pibasis.basis1p, Rs, Zs, z0)

   Es = zero(fltype(V))
   @inbounds for i = 1:dag.num1
      AAdag[i] = a = A[i]
      Es = muladd(vals[i], real(a), Es)
   end

   @inbounds for i = (dag.num1+1):dag.numstore
      n1, n2 = nodes[i]
      AAdag[i] = a = AAdag[n1] * AAdag[n2]
      Es = muladd(vals[i], real(a), Es)
   end

   @inbounds for i = (dag.numstore+1):length(dag)
      n1, n2 = nodes[i]
      a = AAdag[n1] * AAdag[n2]
      Es = muladd(vals[i], real(a), Es)
   end

   return Es
end

alloc_temp_d(::DAGEvaluator, V::PIPotential{T}, maxN::Integer) where {T} =
   (
   R = zeros(JVec{real(T)}, maxN),
   Z = zeros(AtomicNumber, maxN),
    dV = zeros(JVec{real(T)}, maxN),
   tmpd_basis1p = alloc_temp_d(V.pibasis.basis1p),
   AA = zeros(fltype(V.pibasis.basis1p), _maxstore(V)),
   B = zeros(fltype(V.pibasis.basis1p), _maxstore(V)),
   A = alloc_B(V.pibasis.basis1p),
   dA = alloc_dB(V.pibasis.basis1p)
    )

function evaluate_d!(dEs, tmpd, V::PIPotential{T}, ::DAGEvaluator, Rs, Zs, z0) where {T}
   iz0 = z2i(V, z0)
   AA = tmpd.AA
   B = tmpd.B
   tmpd_basis1p = tmpd.tmpd_basis1p
   basis1p = V.pibasis.basis1p
   dag = V.dags[iz0]
   nodes = dag.nodes
   coeffs = dag.vals

   # we start from the representation
   #     V = sum B[i] AA[i]
   # i.e. this vector represents the constributions c[i] ∂AA[i]
   copy!(B, coeffs)

   # FORWARD PASS
   # ------------
   # evaluate the 1-particle basis
   # this puts the first `dag.num1` 1-b (trivial) correlations into the
   # storage array, and from these we can build the rest
   evaluate!(AA, tmpd_basis1p, basis1p, Rs, Zs, z0)

   # Stage 2 of evaluate!
   # go through the dag and store the intermediate results we need
   @inbounds @fastmath for i = (dag.num1+1):dag.numstore
      n1, n2 = nodes[i]
      AA[i] = muladd(AA[n1], AA[n2], AA[i])
   end

   # BACKWARD PASS
   # --------------
   # fill the B array -> coefficients of the derivatives
   #  AA_i = AA_{n1} * AA_{n2}
   #  ∂AA_i = AA_{n1} * ∂AA_{n2} + AA_{n1} * AA_{n2}
   #  c_{n1} * ∂AA_{n1} <- (c_{n1} + c_i AA_{n2}) ∂AA_{n1}
   @inbounds @fastmath for i = length(dag):-1:(dag.numstore+1)
      c = coeffs[i]
      n1, n2 = nodes[i]
      B[n1] = muladd(c, AA[n2], B[n1])
      B[n2] = muladd(c, AA[n1], B[n2])
   end
   # in stage 2 c = C[i] is replaced with b = B[i]
   @inbounds @fastmath for i = dag.numstore:-1:(dag.num1+1)
      n1, n2 = nodes[i]
      b = B[i]
      B[n1] = muladd(b, AA[n2], B[n1])
      B[n2] = muladd(b, AA[n1], B[n2])
   end

   # stage 3: get the gradients
   fill!(dEs, zero(JVec{T}))
   A = tmpd.A
   dA = tmpd.dA
   for (iR, (R, Z)) in enumerate(zip(Rs, Zs))
      evaluate_d!(A, dA, tmpd_basis1p, basis1p, R, Z, z0)
      iz = z2i(basis1p, Z)
      inds = basis1p.Aindices[iz, iz0]
      for iA = 1:length(basis1p, iz, iz0)
         dEs[iR] += real(B[inds[iA]] * dA[inds[iA]])
      end
   end

   return dEs
end


export BasicDQNLearner

using Random
using Flux

"""
    BasicDQNLearner(;kwargs...)

See paper: [Playing Atari with Deep Reinforcement Learning](https://www.cs.toronto.edu/~vmnih/docs/dqn.pdf)

This is the very basic implementation of DQN. Compared to the traditional Q learning, the only difference is that,
in the updating step it uses a batch of transitions sampled from an experience buffer instead of current transition.
And the `approximator` is usually a [`NeuralNetworkApproximator`](@ref).
You can start from this implementation to understand how everything is organized and how to write your own customized algorithm.
# Keywords
- `approximator`::[`AbstractApproximator`](@ref): used to get Q-values of a state.
- `loss_func`: the loss function to use. TODO: provide a default [`huber_loss`](@ref)?
- `γ::Float32=0.99f0`: discount rate.
- `batch_size::Int=32`
- `min_replay_history::Int=32`: number of transitions that should be experienced before updating the `approximator`.
- `seed=nothing`.
"""
mutable struct BasicDQNLearner{Q,F,R} <: AbstractLearner
    approximator::Q
    loss_func::F
    γ::Float32
    batch_size::Int
    min_replay_history::Int
    rng::R
    loss::Float32
end

Flux.functor(x::BasicDQNLearner) = (Q = x.approximator,), y -> begin
    x = @set x.approximator = y.Q
    x
end

(learner::BasicDQNLearner)(obs) =
    obs |>
    get_state |>
    x ->
        send_to_device(device(learner.approximator), x) |>
        learner.approximator |>
        send_to_host

function BasicDQNLearner(;
    approximator::Q,
    loss_func::F,
    γ = 0.99f0,
    batch_size = 32,
    min_replay_history = 32,
    seed = nothing,
) where {Q,F}
    rng = MersenneTwister(seed)
    BasicDQNLearner{Q,F,typeof(rng)}(
        approximator,
        loss_func,
        γ,
        batch_size,
        min_replay_history,
        rng,
        0.0,
    )
end

function RLBase.update!(learner::BasicDQNLearner, t::AbstractTrajectory)
    length(t) < learner.min_replay_history && return

    inds = rand(learner.rng, 1:length(t), learner.batch_size)
    batch = map(get_trace(t, :state, :action, :reward, :terminal, :next_state)) do x
        consecutive_view(x, inds)
    end

    Q, γ, loss_func, batch_size =
        learner.approximator, learner.γ, learner.loss_func, learner.batch_size
    s, r, t, s′ = map(
        x -> send_to_device(device(Q), x),
        (batch.state, batch.reward, batch.terminal, batch.next_state),
    )
    a = CartesianIndex.(batch.action, 1:batch_size)

    gs = gradient(params(Q)) do
        q = Q(s)[a]
        q′ = vec(maximum(Q(s′); dims = 1))
        G = r .+ γ .* (1 .- t) .* q′
        loss = loss_func(G, q)
        ignore() do
            learner.loss = loss
        end
        loss
    end

    update!(Q, gs)
end


export GLFont, TextField, StyledTextSegment, TextCursor, FontProperties

immutable StyledTextSegment
	segment::UnitRange{Int}
	style::Dict{ASCIIString, Any}
end
type TextCursor
	x::Float32
	y::Float32
	intend::Float32
end

type TextField
	id::Symbol
	text::UTF8String
	newLineIndexes::Array{UnitRange{Int}, 1}
	#words::Array{SubString, 1}
	styles::Array{StyledTextSegment, 1}
	selection::UnitRange{Int}
	x::Float32
	y::Float32
	area::Shape
	hasFocus::Bool
	function TextField(id::Symbol, text::String, x::Real, y::Real, area::Shape)
		defaultStyle 	= @compat Dict{ASCIIString, Any}("textColor" => Float32[1,1,1,1], "backgroundColor" => Float32[0,0,0,0])
		styles 			= [StyledTextSegment(1:length(text), defaultStyle)]
		new(id, utf8(text), build_line_indexes(text), styles, length(text) : length(text)-1, float32(x), float32(y), area, false)
	end

	function TextField(id::Symbol,text::String, styles::Array{StyledTextSegment, 1}, x::Real, y::Real, area::Shape)
		new(id, utf8(text), build_line_indexes(text), styles, length(text):length(text)-1, float32(x), float32(y), area, false)
	end

	function TextField(id::Symbol,text::String,
			styles::Array{StyledTextSegment, 1},
			selection::Range,
			x::Real, y::Real, area::Shape)
		new(id, utf8(text), build_line_indexes(text), styles, selection, float32(x), float32(y), area, false)
	end
end

type FontProperties
	lineHeight::GLfloat
    advance::GLfloat
end

immutable GLFont
    data::Dict{Symbol, Any}
	props::Vector{Vec2}
	function GLFont(name::String)
	    flStream            = open("$(name).txt")
	    width               = int(readline(flStream))
	    height              = int(readline(flStream))
	    lineHeight::GLfloat = int(readline(flStream))
	    advance::GLfloat 	= 0

	    values                   = split(readline(flStream))
        charCodet                = char(int(values[1]))
        advancet                 = int(values[2])
        xt::GLfloat              = int(values[3]) / width
        x2t::GLfloat             = (int(values[3]) + advance) / width
        yt::GLfloat              =  ((height - int(values[4])  - lineHeight) / height)
        texLineHeightt::GLfloat  = lineHeight / height

        tmp_uv 	  = Array(Vec2, 1, 5)
        tmp_uv[1] = Vec2(advancet, lineHeight)
        tmp_uv[2] = Vec2(xt,  yt)
        tmp_uv[3] = Vec2(xt,  yt + texLineHeightt)
        tmp_uv[4] = Vec2(x2t, yt + texLineHeightt)
        tmp_uv[5] = Vec2(x2t, yt)

	    uv = vcat(tmp_uv)
	    i = 1
	    for line in eachline(flStream)
	        values                  = split(line)
	        charCode                = char(int(values[1]))
	        advance                 = int(values[2])
	        x::GLfloat              = int(values[3]) / width
	        x2::GLfloat             = (int(values[3]) + advance) / width
	        y::GLfloat              = ((height - int(values[4]) - lineHeight) / height)
	        texLineHeight::GLfloat  = lineHeight / height
	        tmp_uv[1] = Vec2(advance, lineHeight)
	        tmp_uv[2] = Vec2(x,  y)
	        tmp_uv[3] = Vec2(x,  y + texLineHeight)
	        tmp_uv[4] = Vec2(x2, y + texLineHeight)
	        tmp_uv[5] = Vec2(x2, y)
	        uv = vcat(uv, tmp_uv)
	        i += 1
	    end
	    close(flStream)
		data = @compat Dict{Symbol, Any}(
			:dontdelete_uv_index 		=> GLBuffer(GLint[1:6;], 1),
			:dontdelete_indexes 		=> indexbuffer(GLuint[0:5;]),
			:dontdelete_uv 				=> Texture(uv),
			:dontdelete_font_texture 	=> Texture("$(name).bmp")
		)
	    new(data, vec(uv[1,:]))
	end
end


function add_re!(sys)
    re = RenewableDispatch(
        "WindBusA",
        true,
        get_component(Bus, sys, "bus5"),
        0.0,
        0.0,
        1.200,
        PrimeMovers.WT,
        (min = 0.0, max = 0.0),
        1.0,
        TwoPartCost(0.220, 0.0),
        100.0,
    )
    add_component!(sys, re)
    copy_time_series!(re, get_component(PowerLoad, sys, "bus2"))

    for g in get_components(HydroEnergyReservoir, sys)
        tpc = get_operation_cost(g)
        smc = StorageManagementCost(
            variable = get_variable(tpc),
            fixed = get_fixed(tpc),
            start_up = 0.0,
            shut_down = 0.0,
            energy_shortage_cost = 10.0,
            energy_surplus_cost = 10.0,
        )
        set_operation_cost!(g, smc)
    end
end

function run_test_sim(result_dir::String)
    sim_name = "results_sim"
    sim_path = joinpath(result_dir, sim_name)

    if ispath(sim_path)
        @info "Reading UC system from" sim_path
        c_sys5_hy_uc = System(
            joinpath(sim_path, "..", "c_sys5_hy_uc.json"),
            time_series_read_only = true,
        )
        @info "Reading ED system from" sim_path
        c_sys5_hy_ed = System(
            joinpath(sim_path, "..", "c_sys5_hy_ed.json"),
            time_series_read_only = true,
        )
        executions = tryparse.(Int, readdir(sim_path))
        sim = joinpath(sim_path, string(maximum(executions)))
        @info "Reading results from last execution" sim
    else
        @info "Building UC system from"
        c_sys5_hy_uc = PSB.build_system(PSB.SIIPExampleSystems, "5_bus_hydro_uc_sys")
        @info "Building ED system from"
        c_sys5_hy_ed = PSB.build_system(PSB.SIIPExampleSystems, "5_bus_hydro_ed_sys")

        @info "Adding extra RE"
        add_re!(c_sys5_hy_uc)
        add_re!(c_sys5_hy_ed)
        to_json(c_sys5_hy_uc, joinpath(sim_path, "..", "c_sys5_hy_uc.json"))
        to_json(c_sys5_hy_ed, joinpath(sim_path, "..", "c_sys5_hy_ed.json"))

        mkpath(result_dir)
        GLPK_optimizer =
            optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => GLPK.GLP_MSG_OFF)

        template_hydro_st_uc = template_unit_commitment()
        set_device_model!(template_hydro_st_uc, HydroDispatch, FixedOutput)
        set_device_model!(
            template_hydro_st_uc,
            HydroEnergyReservoir,
            HydroDispatchReservoirStorage,
        )

        template_hydro_st_ed = template_economic_dispatch()
        set_device_model!(template_hydro_st_ed, HydroDispatch, FixedOutput)
        set_device_model!(
            template_hydro_st_ed,
            HydroEnergyReservoir,
            HydroDispatchReservoirStorage,
        )
        template_hydro_st_ed.services = Dict() #remove ed services

        problems = SimulationProblems(
            UC = OperationsProblem(
                template_hydro_st_uc,
                c_sys5_hy_uc,
                optimizer = GLPK_optimizer,
                system_to_file = false,
            ),
            ED = OperationsProblem(
                template_hydro_st_ed,
                c_sys5_hy_ed,
                optimizer = GLPK_optimizer,
                constraint_duals = [:CopperPlateBalance],
                system_to_file = false,
                balance_slack_variables = true,
            ),
        )

        sequence = SimulationSequence(
            problems = problems,
            feedforward_chronologies = Dict(("UC" => "ED") => Synchronize(periods = 24)),
            intervals = Dict(
                "UC" => (Hour(24), Consecutive()),
                "ED" => (Hour(1), Consecutive()),
            ),
            feedforward = Dict(
                ("ED", :devices, :ThermalStandard) => SemiContinuousFF(
                    binary_source_problem = PSI.ON,
                    affected_variables = [PSI.ACTIVE_POWER],
                ),
                ("ED", :devices, :HydroEnergyReservoir) => IntegralLimitFF(
                    variable_source_problem = PSI.ACTIVE_POWER,
                    affected_variables = [PSI.ACTIVE_POWER],
                ),
            ),
            cache = Dict(
                ("UC",) => TimeStatusChange(PSY.ThermalStandard, PSI.ON),
                ("UC", "ED") => StoredEnergy(PSY.HydroEnergyReservoir, PSI.ENERGY),
            ),
            ini_cond_chronology = InterProblemChronology(),
        )
        sim = Simulation(
            name = "results_sim",
            steps = 2,
            problems = problems,
            sequence = sequence,
            simulation_folder = result_dir,
        )
        build!(sim)
        execute!(sim)
    end

    results = SimulationResults(sim)
    results_uc = get_problem_results(results, "UC")
    set_system!(results_uc, c_sys5_hy_uc)
    results_ed = get_problem_results(results, "ED")
    set_system!(results_ed, c_sys5_hy_ed)

    return results_uc, results_ed
end

function run_test_prob()
    c_sys5_hy_uc = PSB.build_system(PSB.SIIPExampleSystems, "5_bus_hydro_uc_sys")
    add_re!(c_sys5_hy_uc)
    GLPK_optimizer =
        optimizer_with_attributes(GLPK.Optimizer, "msg_lev" => GLPK.GLP_MSG_OFF)

    template_hydro_st_uc = template_unit_commitment()
    set_device_model!(template_hydro_st_uc, HydroDispatch, FixedOutput)
    set_device_model!(
        template_hydro_st_uc,
        HydroEnergyReservoir,
        HydroDispatchReservoirStorage,
    )

    prob = OperationsProblem(
        template_hydro_st_uc,
        c_sys5_hy_uc,
        optimizer = GLPK_optimizer,
        horizon = 12,
        use_parameters = true,
    )
    build!(prob, output_dir = mktempdir())
    solve!(prob)
    res = ProblemResults(prob)
    return res
end


export WallBoundaries, FreeSlip, NoSlip

abstract type SlipComponent{N} end

struct FreeSlip{N} <: SlipComponent{N}
    FreeSlip(::Dim{N}, ::Integer) where {N} = new{N}()
end

Base.show(io::IO, ::Type{FreeSlip{N}}) where {N} = print(io, "FreeSlip")

mutable struct NoSlip{N} <: SlipComponent{N}
    virtual_velocities::Vector{SVectorF{N}}

    NoSlip(::Dim{N}, number_of_particles::Integer) where {N} =
        new{N}([zeros(SVectorF{N}) for _ = 1:number_of_particles])
end

Base.show(io::IO, ::Type{NoSlip{N}}) where {N} = print(io, "NoSlip")

mutable struct WallBoundaries{N,S} <: BoundaryComponent{N}
    positions::Vector{SVectorF{N}}
    velocities::Vector{SVectorF{N}}
    acceleration::SVectorF{N}
    virtual_pressures::Vector{Float}
    slip_component::S

    function WallBoundaries(
        positions::Vector{SVectorF{N}},
        velocities::Vector{SVectorF{N}},
        acceleration::SVectorF{N},
        slip_component::Type{S} = FreeSlip,
    ) where {N,S<:SlipComponent{N}}
        number_of_particles = length(positions)
        @checkeq(length(velocities), number_of_particles)
        new{N,S}(
            positions,
            velocities,
            acceleration,
            fill(0.0, number_of_particles),
            slip_component(Dim{N}(), number_of_particles),
        )
    end
end

Base.show(io::IO, ::Type{WallBoundaries{N,S}}) where {N,S} =
    print(io, "WallBoundaries{$(S)}")

get_positions(wall::WallBoundaries) = wall.positions
get_velocities(wall::WallBoundaries) = wall.velocities
get_pressures(wall::WallBoundaries) = wall.virtual_pressures

function compute_mass_density(
    wall::WallBoundaries,
    energy_component::EnergyComponent,
    eos::EquationOfState,
    idx_fluid::Unsigned,
    idx_boundary::Unsigned,
)
    mass_density_from(
        eos,
        wall.virtual_pressures[idx_boundary],
        get_specific_energy(energy_component, idx_fluid),
    )
end

function updateboundaries!(
    wall::WallBoundaries{N},
    positions::AbstractVector{SVectorF{N}},
    velocities::Velocities{N},
    standard_mass::StandardMass{ConstantWidthKernel{N}},
    pressures::Pressures,
) where {N}
    update_wall_properties!(
        wall,
        positions,
        get_velocities(velocities),
        get_mass_densities(standard_mass),
        get_pressures(pressures),
        get_kernel(standard_mass),
    )
end

function evolveboundaries!(wall::WallBoundaries, Δt::Number)
    wall.positions .+= Δt * wall.velocities
end

function update_wall_properties!(
    wall::WallBoundaries{N},
    fluid_positions::AbstractVector{SVectorF{N}},
    fluid_velocities::AbstractVector{SVectorF{N}},
    fluid_mass_densities::AbstractVector{Float},
    fluid_pressures::AbstractVector{Float},
    W::ConstantWidthKernel{N},
) where {N}
    h = get_kernel_width(W)
    h⁻¹ = get_inverse_kernel_width(W)
    for (i, (rᵢ, vᵢ)) in enumerate(zip(wall.positions, wall.velocities))
        ΣPⱼWᵢⱼ::Float = 0.0
        ΣρⱼrᵢⱼWᵢⱼ = zeros(SVectorF{N})
        ΣvⱼWᵢⱼ = zeros(SVectorF{N})
        ΣWᵢⱼ::Float = 0.0
        for (rⱼ, vⱼ, ρⱼ, Pⱼ) in zip(
            fluid_positions,
            fluid_velocities,
            fluid_mass_densities,
            fluid_pressures,
        )
            rᵢⱼ = rⱼ - rᵢ
            qᵢⱼ = norm(rᵢⱼ) * h⁻¹
            if nonzeroat(W, qᵢⱼ)
                Wᵢⱼ = W(qᵢⱼ, h)
                ΣPⱼWᵢⱼ += Pⱼ * Wᵢⱼ
                ΣρⱼrᵢⱼWᵢⱼ .+= (ρⱼ * Wᵢⱼ) * rᵢⱼ
                accumulate_virtual_velocity(
                    ΣvⱼWᵢⱼ,
                    vⱼ,
                    Wᵢⱼ,
                    wall.slip_component,
                )
                ΣWᵢⱼ += W_rᵢⱼ_h
            end
        end
        ΣWᵢⱼ⁻¹ = 1 / ΣWᵢⱼ
        wall.virtual_pressures[i] =
            (ΣPⱼWᵢⱼ - wall.acceleration ⋅ ΣρⱼrᵢⱼWᵢⱼ) * ΣWᵢⱼ⁻¹
        set_virtual_velocity(wall.slip_component, vᵢ, ΣvⱼWᵢⱼ, ΣWᵢⱼ⁻¹, i)
    end
end

function accumulate_virtual_velocity(
    ::SVectorF{N},
    ::SVectorF{N},
    ::Number,
    ::FreeSlip{N},
) where {N} end

function set_virtual_velocity(
    ::FreeSlip{N},
    ::SVectorF{N},
    ::SVectorF{N},
    ::Number,
    ::Integer,
) where {N} end

function accumulate_virtual_velocity(
    ΣvⱼWᵢⱼ::SVectorF{N},
    vⱼ::SVectorF{N},
    Wᵢⱼ::Number,
    no_slip::NoSlip{N},
) where {N}
    ΣvⱼWᵢⱼ .+= vⱼ * Wᵢⱼ
end

function set_virtual_velocity(
    no_slip::NoSlip{N},
    vᵢ::SVectorF{N},
    ΣvⱼWᵢⱼ::SVectorF{N},
    ΣWᵢⱼ⁻¹::Number,
    i::Integer,
) where {N}
    ṽᵢ = ΣvⱼWᵢⱼ * ΣWᵢⱼ⁻¹
    no_slip.virtual_velocities[i] = 2 * vᵢ - ṽᵢ
end


using ModernGL, GeometryTypes, GLAbstraction, GLWindow, Images

# Load our texture. See "downloads.jl" to get the images.
img = load(GLAbstraction.dir("tutorials", "images", "kitten.png"))

# Create the window. This sets all the hints and makes the context current.
window = create_glcontext("Textures exercise 2", resolution=(800,600))

vao = glGenVertexArrays()
glBindVertexArray(vao)

# The positions of the vertices in our rectangles
# This is a two-rectangle solution; alternatively, one could modify
# the fragment shader to compute the reflection.
vertex_positions = Point{2,Float32}[(-0.5,  0.5),     # top-left
                                    ( 0.5,  0.5),     # top-right
                                    ( 0.5,  0.0),     # middle-right
                                    (-0.5,  0.0),     # middle-left
                                    ( 0.5, -0.5),     # bottom-right
                                    (-0.5, -0.5)]     # bottom-left
# Duplicate the middle ones
vertex_positions = vertex_positions[[1:4;4;3;5:6]]

# The texture coordinates of each vertex
vertex_texcoords = Vec2f0[(0, 0),
                          (1, 0),
                          (1, 1),
                          (0, 1)]
vertex_texcoords = vertex_texcoords[[1:4;4;3;2;1]]

# Specify how vertices are arranged into faces
elements = Face{3,UInt32,-1}[(0,1,2),
                             (2,3,0),
                             (4,5,6),
                             (6,7,4)]

# The vertex shader---note the `vert` in front of """
vertex_shader = vert"""
#version 150

in vec2 position;
in vec2 texcoord;

out vec2 Texcoord;

void main()
{
    Texcoord = texcoord;
    gl_Position = vec4(position, 0.0, 1.0);
}
"""

# The fragment shader
fragment_shader = frag"""
# version 150

in vec2 Texcoord;

out vec4 outColor;

uniform sampler2D tex;

void main()
{
    outColor = texture(tex, Texcoord);
}
"""

# Link everything together, using the corresponding shader variable as
# the Dict key
bufferdict = Dict(:position=>Buffer(vertex_positions),
                  :texcoord=>Buffer(vertex_texcoords),
                  :tex=>Texture(data(img)),
                  :indexes=>indexbuffer(elements)) # special for element buffers

ro = std_renderobject(bufferdict,
                      LazyShader(vertex_shader, fragment_shader))

# Draw until we receive a close event
while !GLFW.WindowShouldClose(window)
    glClear(GL_COLOR_BUFFER_BIT)
    GLAbstraction.render(ro)
    GLFW.SwapBuffers(window)
    GLFW.PollEvents()
    if GLFW.GetKey(window, GLFW.KEY_ESCAPE) == GLFW.PRESS
        GLFW.SetWindowShouldClose(window, true)
    end
end

GLFW.DestroyWindow(window)  # needed if you're running this from the REPL


module ParameterJuMP

using SparseArrays

using JuMP
export index
using MathOptInterface
const MOI = MathOptInterface
const MOIU = MOI.Utilities

export
ModelWithParams, ParameterRef, add_parameter, add_parameters, all_parameters, Param, parametrized_dual_objective_value

# types
# ------------------------------------------------------------------------------
struct ParameterRef <: JuMP.AbstractVariableRef
    ind::Int64 # local reference
    model::JuMP.Model
end

# Reference to a constraint in which the parameter has coefficient coef
struct ParametrizedConstraintRef{C}
    cref::JuMP.ConstraintRef{JuMP.Model, C}
    coef::Float64
end

const CtrRef{F, S} = ConstraintRef{JuMP.Model,MOI.ConstraintIndex{F,S}, JuMP.ScalarShape}
const SAF = MOI.ScalarAffineFunction{Float64}
const EQ = MOI.EqualTo{Float64}
const LE = MOI.LessThan{Float64}
const GE = MOI.GreaterThan{Float64}

mutable struct ParameterData
    inds::Vector{Int64}
    next_ind::Int64

    current_values::Vector{Float64}
    future_values::Vector{Float64}

    # Whether the JuMP model is in sync with the value of the parameters
    sync::Bool

    # constraints where it is a RHS
    constraints_map::Dict{Int64, Vector{ParametrizedConstraintRef}}

    parameters_map_saf_in_eq::Dict{CtrRef{SAF, EQ}, JuMP.GenericAffExpr{Float64,ParameterRef}}
    parameters_map_saf_in_le::Dict{CtrRef{SAF, LE}, JuMP.GenericAffExpr{Float64,ParameterRef}}
    parameters_map_saf_in_ge::Dict{CtrRef{SAF, GE}, JuMP.GenericAffExpr{Float64,ParameterRef}}

    names::Dict{ParameterRef, String}

    # overload addvariable
    # variables_map::Dict{Int64, Vector{Int64}}
    # variables_map_lb::Dict{Int64, Vector{Int64}}
    # variables_map_ub::Dict{Int64, Vector{Int64}}

    # solved::Bool

    has_deletion::Bool
    index_map::Dict{Int64, Int64}

    lazy::Bool

    no_duals::Bool
    dual_values::Vector{Float64}
    function ParameterData()
        new(Int64[],
            1,
            Float64[],
            Float64[],
            true,
            Dict{Int64, Vector{JuMP.ConstraintRef}}(),
            Dict{CtrRef{SAF, EQ}, JuMP.GenericAffExpr{Float64,ParameterRef}}(),
            Dict{CtrRef{SAF, LE}, JuMP.GenericAffExpr{Float64,ParameterRef}}(),
            Dict{CtrRef{SAF, GE}, JuMP.GenericAffExpr{Float64,ParameterRef}}(),
            Dict{ParameterRef, String}(),
            false,
            Dict{Int64, Int64}(),
            false,
            false,
            Float64[]
            )
    end
end

no_duals(data::ParameterData) = data.no_duals
no_duals(model::JuMP.Model) = no_duals(_getparamdata(model))

set_no_duals(model::JuMP.Model) = set_no_duals(_getparamdata(model))
function set_no_duals(data::ParameterData)
    if isempty(data.current_values)
        data.no_duals = true
    elseif no_duals(data)
        @warn "No duals mode is already activated"
    else
        error("Parameter JuMP's no duals mode can only be activated in empty models.")
    end
end

"""
    JuMP.index(p::ParameterRef)::Int64

Return the internal index of the parameter `p`.
"""
JuMP.index(p::ParameterRef) = index(_getparamdata(p.model), p)::Int64
JuMP.index(model::JuMP.Model, p::ParameterRef) = index(_getparamdata(model), p)::Int64
function JuMP.index(data::ParameterData, p::ParameterRef)::Int64
    if data.has_deletion
        return data.index_map[p.ind]
    else
        return p.ind
    end
end

lazy_duals(data::ParameterData) = data.lazy
lazy_duals(model::JuMP.Model) = lazy_duals(_getparamdata(model))

set_lazy_duals(model::JuMP.Model) = set_lazy_duals(_getparamdata(model))
function set_lazy_duals(data::ParameterData)
    if isempty(data.current_values)
        data.lazy = true
    elseif lazy_duals(data)
        @warn "Lazy mode is already activated"
    else
        error("Parameter JuMP's lazy mode can only be activated in empty models.")
    end
end

set_not_lazy_duals(model::JuMP.Model) = set_not_lazy_duals(_getparamdata(model))
function set_not_lazy_duals(data::ParameterData)
    if isempty(data.current_values)
        data.lazy = false
    elseif !lazy_duals(data)
        @warn "Lazy mode is already de-activated"
    else
        error("Parameter JuMP's lazy mode can only be de-activated in empty models.")
    end
end

_get_param_dict(data::ParameterData, ::Type{EQ}) = data.parameters_map_saf_in_eq
_get_param_dict(data::ParameterData, ::Type{LE}) = data.parameters_map_saf_in_le
_get_param_dict(data::ParameterData, ::Type{GE}) = data.parameters_map_saf_in_ge

_getmodel(p::ParameterRef) = p.model
_getparamdata(p::ParameterRef)::ParameterData = _getparamdata(_getmodel(p))::ParameterData
# _getparamdata(model::JuMP.Model) = model.ext[:params]::ParameterData
function _getparamdata(model::JuMP.Model)::ParameterData
    # TODO checks dict twice
    !haskey(model.ext, :params) && error("In order to use Parameters the model must be created with the ModelWithParams constructor")
    return model.ext[:params]
end

"""
    is_sync(model::JuMP.Model)

Test if the JuMP Model is updated to the latest value of all parameters.
"""
is_sync(data::ParameterData) = data.sync
is_sync(model::JuMP.Model) = is_sync(_getparamdata(model))

"""
    add_parameter(model::JuMP.Model, val::Real)::ParameterRef

Adds one parameter fixed at `val` to the model `model`.

    add_parameter(model::JuMP.Model)::ParameterRef

Adds one parameter fixed at zero to the model `model`.
"""
function add_parameter(model::JuMP.Model, val::Real)
    params = _getparamdata(model)::ParameterData

    ind = params.next_ind
    params.next_ind += 1

    push!(params.inds, ind)
    push!(params.current_values, 0.0)
    push!(params.future_values, val)
    if !iszero(val)
        params.sync = false
    end
    push!(params.dual_values, NaN)

    params.constraints_map[ind] = ParametrizedConstraintRef[]

    return ParameterRef(ind, model)
end
add_parameter(model) = add_parameter(model, 0.0)

_addsizehint!(vec::Vector, len::Integer) = sizehint!(vec, len+length(vec))

"""
    add_parameters(model::JuMP.Model, val::Vector{R})::Vector{ParameterRef}

Adds one parameter for each element of the vector `val`.

    add_parameters(model::JuMP.Model, N::Integer)::Vector{ParameterRef}

Adds `N` parameters to the model, all of them fixed in zero.
"""
function add_parameters(model::JuMP.Model, N::Integer)
    return add_parameters(model::JuMP.Model, zeros(N))
end
function add_parameters(model::JuMP.Model, val::AbstractArray{R,N}) where {R,N}
    params = _getparamdata(model)::ParameterData

    nparam = length(val)
    out = ParameterRef[]
    out = similar(val, ParameterRef)
    sizehint!(out, nparam)
    _addsizehint!(params.inds, nparam)
    _addsizehint!(params.current_values, nparam)
    _addsizehint!(params.future_values, nparam)
    _addsizehint!(params.dual_values, nparam)

    for i in eachindex(val)
        ind = params.next_ind
        params.next_ind += 1

        push!(params.inds, ind)
        push!(params.current_values, 0.0)
        push!(params.future_values, val[i])
        if !iszero(val[i])
            params.sync = false
        end
        push!(params.dual_values, NaN)

        params.constraints_map[ind] = ParametrizedConstraintRef[]

        out[i] = ParameterRef(ind, model)
    end

    return out
end

"""
    all_parameters(model::JuMP.AbstractModel)::Vector{ParameterRef}

Returns a list of all parameters currently in the model. The parameters are
ordered by creation time.
"""
function all_parameters(model::JuMP.AbstractModel)
    data = _getparamdata(model)
    return ParameterRef[ParameterRef(ind, model) for ind in data.inds]
end

# solve
# ------------------------------------------------------------------------------

"""
    ModelWithParams(args...; kwargs...)::JuMP.Model

Returns a JuMP.Model able to handle `Parameters`.

`args` and `kwargs` are the same parameters that would be passed
to the regular `Model` constructor.

Example using GLPK solver:

```julia
    model = ModelWithParams(with_optimizer(GLPK.Optimizer))
```
"""
function ModelWithParams(args...; kwargs...)
    m = JuMP.Model(args...; kwargs...)
    enable_parameters(m)
    return m
end

"""
    enable_parameters(m::JuMP.Model)::Nothing

Enables a `JuMP.Model` to handle `Parameters`.
"""
function enable_parameters(m::JuMP.Model)
    if haskey(m.ext, :params)
        error("Model already has parameter enabled")
    end
    # test and compare hook
    initialize_parameter_data(m)
    JuMP.set_optimize_hook(m, parameter_optimizehook)
    return nothing
end

function initialize_parameter_data(m::JuMP.Model)
    m.ext[:params] = ParameterData()
end

# The constraint has been deleted. We cannot keep `dict` in sync with deletions
# as we return a `JuMP.ConstraintRef`, not a custom type when the user create a
# parametrized constraint.
function _update_dicts_with_deletion(data::ParameterData, S::Type)
    dict = _get_param_dict(data, S)
    to_delete = eltype(keys(dict))[]
    for (cref, gaep) in dict
        if !JuMP.is_valid(cref.model, cref)
            push!(to_delete, cref)
        end
    end
    for cref in to_delete
        delete!(dict, cref)
    end
end
function _update_dicts_with_deletion(data::ParameterData)
    _update_dicts_with_deletion(data, EQ)
    _update_dicts_with_deletion(data, LE)
    _update_dicts_with_deletion(data, GE)
end
function parameter_optimizehook(m::JuMP.Model; kwargs...)
    data = _getparamdata(m)::ParameterData

    _update_dicts_with_deletion(data)

    # sync model rhs to newest parameter values
    sync(data)

    ret = JuMP.optimize!(m::JuMP.Model, ignore_optimize_hook = true, kwargs...)

    # update duals for later query
    if !no_duals(data) && JuMP.has_duals(m)
        _update_duals(data)
    end
    return ret
end

"""
    sync(model::JuMP.Model)::Nothing

Forces the model to update its constraints to the new values of the
Parameters.
"""
function sync(data::ParameterData)
    if !is_sync(data)
        _update_constraints(data)
        data.current_values .= data.future_values
        data.sync = true
    end
    return nothing
end
function sync(model::JuMP.Model)
    sync(_getparamdata(model))
end

function parametrized_dual_objective_value(model::JuMP.AbstractModel)
    params = all_parameters(model)
    linear = [param => dual(param) for param in params]
    obj = dual_objective_value(model)
    constant = obj - sum(value(term.first) * term.second for term in linear)
    return GenericAffExpr(constant, linear)
end

# constraints
# ------------------------------------------------------------------------------
include("constraints.jl")

# JuMP variable interface
# ------------------------------------------------------------------------------
include("variable_interface.jl")

# operators and mutable arithmetics
# ------------------------------------------------------------------------------
include("operators.jl")
include("mutable_arithmetics.jl")

# other
# ------------------------------------------------------------------------------
include("macros.jl")
include("print.jl")

end


# Tests for Decision trees/ Random forest
function t_decisiontree()

Ac = j4pr.DataGenerator.fish(20)  # classification dataset
Ar = j4pr.DataGenerator.fishr(20) # regression dataset

Wclass = [j4pr.tree(), j4pr.randomforest(), j4pr.adaboostump()] # only default arguments

Wreg = [j4pr.treer(), j4pr.randomforestr()]			# only default arguments

# Test classification
for w in Wclass
	Test.@test try
		wt1 = Ac |> w
		wt2 = j4pr.strip(Ac) |> w

		result = Ac |> wt1
		result2 = Ac |> wt2

		# Do not check decisions, they may vary based on sampling etc.
		true
	catch
		false
	end
end

# Test regression
for w in Wreg
	Test.@test try
		wt1 = Ar |> w
		wt2 = j4pr.strip(Ar) |> w

		result = Ar |> wt1
		result2 = Ar |> wt2

		# Do not check decisions, they may vary based on sampling etc.
		true
	catch
		false
	end
end

end


@testset "Forecast data" begin
    ps_dict = PowerSystems.parsestandardfiles(joinpath(MATPOWER_DIR, "case5_re.m"))
    da_time_series = PowerSystems.read_data_files(joinpath(FORECASTS_DIR, "5bus_ts");
                                                  REGEX_FILE=r"da_(.*?)\.csv")
    rt_time_series = PowerSystems.read_data_files(joinpath(FORECASTS_DIR, "5bus_ts");
                                                  REGEX_FILE=r"rt_(.*?)\.csv")
    ps_dict = PowerSystems.assign_ts_data(ps_dict,rt_time_series)

    buses, generators, storage, branches, loads, loadZones, shunts, services =
        PowerSystems.ps_dict2ps_struct(ps_dict)
    sys_5 = PowerSystems.System(buses, generators, loads, branches, storage,
                                     ps_dict["baseMVA"])

    forecast_gen = PowerSystems.make_forecast_dict(da_time_series["gen"],
                                                   Dates.Day(1), 24, generators)
    forecast_load = PowerSystems.make_forecast_dict(da_time_series, Dates.Day(1), 24, loads)
    forecast_gen = PowerSystems.make_forecast_dict(da_time_series["gen"], Dates.Day(1), 24,
                                                   generators)
    forecast_load = PowerSystems.make_forecast_dict(da_time_series, Dates.Day(1), 24, loads)
    forecast_gen = PowerSystems.make_forecast_dict(da_time_series["gen"], Dates.Day(1), 24,
                                                   generators);
    forecast_load = PowerSystems.make_forecast_dict(da_time_series, Dates.Day(1), 24,
                                                    loads);

    ps_dict = PowerSystems.parsestandardfiles(joinpath(MATPOWER_DIR, "RTS_GMLC.m"))
    da_time_series = PowerSystems.read_data_files(joinpath(FORECASTS_DIR,
                                                           "RTS_GMLC_forecasts");
                                                  REGEX_FILE=r"DAY_AHEAD(.*?)\.csv")
    rt_time_series = PowerSystems.read_data_files(joinpath(FORECASTS_DIR,
                                                           "RTS_GMLC_forecasts");
                                                  REGEX_FILE=r"REAL_TIME(.*?)\.csv")
    ps_dict = PowerSystems.assign_ts_data(ps_dict,rt_time_series)

    buses, generators, storage, branches, loads, loadZones, shunts, services =
        PowerSystems.ps_dict2ps_struct(ps_dict)
    sys_5 = PowerSystems.System(buses, generators, loads, branches, storage,
                                     ps_dict["baseMVA"])

    forecast_gen = PowerSystems.make_forecast_dict(da_time_series["gen"],
                                                   Dates.Day(1), 24, generators);
    forecast_load = PowerSystems.make_forecast_dict(rt_time_series, Dates.Day(1),
                                                    288, loads, loadZones)

    forecast_struct = PowerSystems.make_forecast_array(forecast_load)
end


module ReinforcementLearningZoo

const RLZoo = ReinforcementLearningZoo
export RLZoo

export GaussianNetwork

using CircularArrayBuffers
using ReinforcementLearningBase
using ReinforcementLearningCore
using Setfield: @set
using Logging
using Flux.Losses
using Dates
using IntervalSets
using Random
using Random: shuffle
using CUDA
using Zygote
using Zygote: ignore, @ignore
using Flux
using Flux: onehot, normalise
using StatsBase
using StatsBase: sample, Weights, mean
using LinearAlgebra: dot
using MacroTools
using Distributions: Categorical, Normal, logpdf
using StructArrays

include("patch.jl")
include("algorithms/algorithms.jl")

end # module


include("filtering.jl");
gr(markerstrokecolor = :black,markersize = 8,)
##
list = ["PreVehicle",
    "Altanserin",
    "SB242084",
    "Way_100135",
    "Methysergide",
    "Citalopram",
    "SB242084_opt",
    "Saline"]
s = filter(r->r.Treatment in list &&
    r.Trial < 51 &&
    # r.Trial_Travel_to < 40 &&
    r.MouseID != "pc7",
    streaks)
b = filter(r->r.Treatment in list &&
    r.Trial < 51 &&
    # r.Trial_Travel_to < 40 &&
    r.MouseID != "pc7",
    bouts)
p = copy(pokes)
for df in [s,b, p]
    df[df.Treatment .== "PreVehicle",:Treatment] .= "Control"
    df[df.Treatment .== "Saline",:Treatment] = [o ? "Optogenetic" : "Control" for o in df[df.Treatment .== "Saline",:Stim]]
    df[df.Treatment .== "SB242084_opt",:Phase] .=  "SB242084_opt"
    df[df.Treatment .== "SB242084_opt",:Treatment] = [o ? "SB242084_opt" : "Control" for o in df[df.Treatment .== "SB242084_opt",:Stim]]
end
##
#look at survival function without censor data. That is looking only at last bout in a trial, not bout interrupted by a reward
f_bouts = filter(r-> r.Omissions_plus_one >= 2 &&
    r.Leave &&
    r.Treatment in ["Control", "Citalopram", "Optogenetic", "SB242084", "Altanserin"],
    b)
f_bouts.BoutOut
f_bouts.Bout_duration
res, plt = mediansurvival_analysis(f_bouts,:Bout_duration,:Treatment)
Drug_colors!(res)
plt = @df res scatter(string.(:Treatment), :Mean, yerror = :Sem, color = :color,
    xlabel = "Group", ylabel = "Median survival time", label = "")
savefig(joinpath(figs_loc,"ForagingMeeting/mediansurvival.pdf"))

f_res = function_analysis(f_bouts,:Bout_duration, survivalrate_algorythm; grouping = :Treatment, step =0.05)
Drug_colors!(f_res)
@df f_res plot(:Bout_duration,:Mean, ribbon = :Sem, group = :Treatment,
    color = :color, linecolor = :color, xlims = (0,15),
    xlabel = "Consecutive unrewarded poking time", ylabel = "Survival rate")
savefig(joinpath(figs_loc,"ForagingMeeting/SurvivalRate.pdf"))
##
b2 = filter(r -> r.Treatment in ["Control", "Citalopram", "Optogenetic", "SB242084", "Altanserin"], b)
KMdf = combine(groupby(b2,[:MouseID,:Treatment]), [:Bout_duration, :Leave] => ((d,l) -> fit(KaplanMeier,d,l)) => :KMfit)
transform!(KMdf, :KMfit => ByRow(SwitchingAnalysis.KM_median) => :KM_median)
dropmissing!(KMdf)
res2 = combine(groupby(KMdf, :Treatment), :KM_median .=> [mean, sem])
Drug_colors!(res2)
plt = @df res2 scatter(string.(:Treatment), :KM_median_mean, yerror = :KM_median_sem, color = :color,
    xlabel = "Group", ylabel = "Median Kaplan Meier survival time", label = "")
savefig(joinpath(figs_loc,"ForagingMeeting/KMmediansurvival.pdf"))

survdf = survival_analysis(b2,:Bout_duration, :Leave; grouping = :Treatment, step =0.05)
Drug_colors!(survdf)
@df survdf plot(:Bout_duration, :Mean, ribbon = :Sem, group = :Treatment,
    color = :color, linecolor = :color, xlims = (0,10),
    xlabel = "Consecutive unrewarded poking time", ylabel = "Survival rate")
savefig(joinpath(figs_loc,"ForagingMeeting/KMSurvivalRate.pdf"))
##
f_pokes = filter(r -> r.Treatment in ["Control", "Citalopram", "Optogenetic", "SB242084", "Altanserin"], p)
f_pokes.Treatment = categorical(f_pokes.Treatment)
levels!(f_pokes.Treatment,["Control", "Citalopram", "Optogenetic", "SB242084", "Altanserin"])
gd = groupby(f_pokes, [:MouseID, :Day, :Treatment, :Bout])
res3 = combine(gd, [:PokeIn, :PokeOut] => ((i,o) -> (In_bout = i .- first(i), Out_bout = o .- first(i))) => AsTable,
    :LastPoke => :Leave,
    :Protocol)
transform!(res3, :Out_bout => zscore)
verbFullBasic = @formula(Leave ~ 1 + Out_bout_zscore * Protocol +  (1+Out_bout_zscore+Protocol|MouseID));
FullBasic = fit(MixedModel,verbFullBasic, res3, Bernoulli(), fast = true)

verbSimpleBasic = @formula(Leave ~ 1 + Out_bout_zscore + Protocol +  (1+Out_bout_zscore+Protocol|MouseID));
SimpleBasic = fit(MixedModel,verbSimpleBasic, res3, Bernoulli(), fast = true)

SimpleFullBasicTest = MixedModels.likelihoodratiotest(LeaveSimpleBasic,LeaveBasic)

verbTreatment = @formula(Leave ~ 1 + Out_bout_zscore * Protocol * Treatment +  (1+Out_bout_zscore+Protocol|MouseID));
LeaveTreatment = fit(MixedModel,verbTreatment, res3, Bernoulli(), fast = true)

TreatmentTest = MixedModels.likelihoodratiotest(LeaveBasic,LeaveTreatment)


struct GeoPlot end
struct GeoZPlot end

# We only look at arrays with X, Y, Z dims here.
# Otherwise they fall back to DimensionalData.jl recipes
@recipe function f(A::AbstractGeoArray)
    ddplot(A) = DimArray(A; dims=_maybe_mapped(dims(A)))

    A = GeoArray(A)
    if !(get(plotattributes, :seriestype, :none) in (:none, :heatmap))
        DD.DimensionalPlot(), ddplot(A)
    elseif all(hasdim(A, (SpatialDim, SpatialDim)))
        # Heatmap or multiple heatmaps. Use GD recipes.
        GeoPlot(), _prepare(A)
    elseif hasdim(A, ZDim) && ndims(A) == 1
        # Z dim plot, but for spatial data we want Z on the Y axis
        GeoZPlot(), _prepare(A)
    else
        DD.DimensionalPlot(), ddplot(A)
    end
end

# Plot 3d arrays as multiple tiled plots
@recipe function f(::GeoPlot, A::GeoArray{T,3,<:Tuple{<:SpatialDim,<:SpatialDim,D}}) where {T,D}
    nplots = size(A, 3)
    if nplots > 1
        :layout --> nplots
        :title --> permutedims(string.(val(dims(A, D))))
        for i in 1:nplots
            @series begin
                seriestype := :heatmap
                aspect_ratio := 1
                subplot := i
                GeoPlot(), A[:, :, i]
                slice = A[:, :, i]
                ys, xs = map(_prepare, dims(A))
                xs, ys, parent(slice)
            end
        end
    else
        GeoPlot(), A[:, :, 1]
    end
end
# Plot a sinlge 2d map
@recipe function f(::GeoPlot, A::GeoArray{T,2,<:Tuple{<:SpatialDim,<:SpatialDim}}) where T
    # If colorbar is close to symmetric (< 25% difference) use a symmetric
    # colormap and set symmetric limits so zero shows up as a neutral color.
    A_min, A_max = extrema(skipmissing(A))
    if (A_min + A_max) / abs(A_max - A_min) < 0.25
        A_limit = max(abs(A_min), abs(A_max))
        clims = (-A_limit, A_limit)
        :seriescolor --> :balance
    else
        clims = A_min, A_max
    end

    yguide, xguide = label(dims(A))

    :seriestype --> :heatmap
    :title --> "$(_maybename(A)) $(DD.refdims_title(A))"
    :xguide --> xguide
    :yguide --> yguide
    :clims --> clims
    :colorbar_title --> name(A)

    ys, xs = map(_prepare, dims(A))

    if get(plotattributes, :seriestype, :none) == :contourf
        :linewidth --> 0
        :levels --> range(clims[1], clims[2], length=20)
        xs, ys, clamp.(A, clims[1], clims[2])
    else
        xs, ys, parent(A)
    end
end

# # Plot a vertical 1d line
@recipe function f(::GeoZPlot, A::GeoArray{T,1}) where T
    z_dim = dims(A, ZDim)
    yguide = label(z_dim)
    xguide = label(A)
    :title --> "$(_maybename(A))$(DD.refdims_title(A))"
    :xguide --> xguide
    :yguide --> yguide
    :label --> ""
    z = map(_prepare, dims(A))
    parent(A), z
end

# Plots heatmaps pixels are centered.
# So we should center, and use the projected value.
_prepare(d::Dimension) = d |> _maybe_shift |> _maybe_mapped |> index
# Convert arrays to a consistent missing value and Forward array order
_prepare(A::AbstractGeoArray) =
    _maybe_replace_missing(A) |>
    A -> reorder(A, ForwardIndex) |>
    A -> reorder(A, ForwardRelation) |>
    A -> permutedims(A, DD.commondims(>:, (ZDim, YDim, XDim, TimeDim, Dimension), dims(A)))

_maybename(A) = _maybename(name(A))
_maybename(n::Symbol) = n == Symbol("") ? "" : string(n, ": ")
_maybename(n::Name{N}) where N = _maybename(N)
_maybename(n::NoName) = ""

_maybe_replace_missing(A::AbstractArray{<:AbstractFloat}) = replace_missing(A, eltype(A)(NaN))
_maybe_replace_missing(A) = A

_maybe_shift(d) = _maybe_shift(sampling(d), d)
_maybe_shift(::Intervals, d) = DD.maybeshiftlocus(Center(), d)
_maybe_shift(sampling, d) = d

_maybe_mapped(dims::Tuple) = map(_maybe_mapped, dims)
_maybe_mapped(dim::Dimension) = _maybe_mapped(mode(dim), dim)
_maybe_mapped(mode::IndexMode, dim::Dimension) = dim
_maybe_mapped(mode::Projected, dim::Dimension) = _maybe_mapped(mappedcrs(mode), dim)
_maybe_mapped(::Nothing, dim::Dimension) = dim
_maybe_mapped(::GeoFormat, dim::Dimension) = convertmode(Mapped, dim)


module ParticleFilter

using RayCast.Bresenham: cast_heading
using Serialization
using RoboLib.Util: rangebearing2point, binarize, img2grid, grid2img, binarize, euclidean
using mushr_pf.SensorModels: LaserScanModel, DiscBeamPDF
using mushr_pf.SensorModels
using mushr_pf.MotionModels: AckerParams, AckerData, step!, BasicStochasticPredictModel, predict!, reset!
using mushr_pf.MotionModels
using RoboLib.Geom: Pose2D, Scale2D, T2D, R2D
using CoordinateTransformations, Rotations
using StaticArrays

include("particlefilter.jl")

struct OccMap{M<:AbstractMatrix{Bool}, TWM, TMW, P, O, S}
    map::M
    T_WM::TWM
    T_MW::TMW
    pos_WM::P
    ori_WM::O
    scale_WM::S
end
function OccMap(map::M, pos_WM::P, ori_WM::O, scale_WM::S) where {M<:AbstractMatrix{Bool},P<:T2D,O<:R2D,S<:Real}
    T_WM = Translation(pos_WM) ∘ LinearMap(ori_WM) ∘ LinearMap(Scale2D(scale_WM))
    T_MW = inv(T_WM)
    OccMap{M, typeof(T_WM), typeof(T_MW), P, O, S}(map, T_WM, T_MW, pos_WM, ori_WM, scale_WM)
end
OccMap(testfn, map, pos_WM, ori_WM, scale_WM) = OccMap(binarize(map, testfn), pos_WM, ori_WM, scale_WM)

# Dangerous indeed!
@inline (m::OccMap)(x::Integer, y::Integer) = @inbounds m.map[x,y]

function setup_predict_model(c::Dict{Symbol, Any})
    mc = c[:predictconf]

    if mc[:type] == :ackermann
        data = AckerData{c[:dtype]}()
        params = AckerParams{c[:dtype]}(mc[:car_length])
    else
        error("\"$mc[:type]\" not a known model")
    end

    if mc[:stochastic]
        process_dist = MvNormal([0,0,0], [mc[:sigma_x], mc[:sigma_y], mc[:sigma_theta]])
        ctrl_dist = MvNormal([0,0], [mc[:vel_sigma], mc[:delta_sigma]])
        init_dist = MvNormal([0,0,0], [mc[:init_sigma_x], mc[:init_sigma_y], mc[:init_sigma_theta]])
        model = BasicStochasticPredictModel(data, params, ctrl_dist, process_dist, init_dist)
    else
        error("not yet implemented")
    end
    model
end

function setup_reweight_model(om::OccMap, c::Dict{Symbol, Any})
    sc = c[:reweightconf]

    if sc[:type] == :beammodel
        #TODO make configurable
        raymethodW = RangeMethodW(om, sc[:max_range_meters])
        beammodel = DiscBeamPDF(om.scale_WM, sc[:max_range_meters]; a_short_lin=sc[:a_short_lin], a_short_exp=sc[:a_short_exp], a_max=sc[:a_max], a_rand=sc[:a_rand], a_hit=sc[:a_hit], sigma_hit=sc[:sigma_hit], lambda_short=sc[:lambda_short])
        reweightmodel = LaserScanModel{c[:dtype]}(beammodel, raymethodW, sc[:inv_squash_factor], sc[:nrays], sc[:method])
    else
        error("not yet implemented")
    end
    reweightmodel
end

struct RangeMethodW{D, M<:OccMap, MR<:Real}
    occmap::M
    mr_M::MR
    function RangeMethodW{D}(map::M, mr_W::MR) where {D,M,MR}
        new{D,M,MR}(map, mr_W/map.scale_WM)
    end
    RangeMethodW(map, mr_W) = RangeMethodW{Int}(map, mr_W)
end

function (rm::RangeMethodW{D})(T_WP::Pose2D{T}, angle) where {D,T}
    T_MP = rm.occmap.T_MW ∘ T_WP
    x_M, y_M, theta_M = T_MP.statev
    x_G, y_G, theta_G = img2grid(x_M, y_M, theta_M)
    xhit_G, yhit_G = cast_heading(D, x_G, y_G, theta_G+angle, rm.mr_M, rm.occmap)
    T(rm.occmap.scale_WM * euclidean(x_G, y_G, xhit_G, yhit_G))::T
end

function setup_resampler(c)
    rs = c[:resampler]
    if rs === :naive
        return resample_naive!
    elseif rs === :lowvar
        error("fix me") #TODO
        return resample_lowvar!
    else
        error("Resampler \"$rs\" unknown!")
    end
end

function setup_pf(occmap::OccMap, c::Dict{Symbol, Any})
    predictmodel = setup_predict_model(c)
    reweightmodel = setup_reweight_model(occmap, c)
    resampler = setup_resampler(c)
    particles = BufferedWeightedParticleBelief{Pose2D{c[:dtype]}}(c[:max_particles])

    StatefulParticleFilter(particles, predictmodel, reweightmodel, resampler, c[:rng])
end

end # module


@testset "cache" begin
    @everywhere GC.gc()
    #= set available memory to 80KB on each worker
    @everywhere empty!(MemPool.datastore)
    @everywhere empty!(MemPool.lru_order)
    @everywhere MemPool.max_memsize[] = 8*10^4
    =#

    thunks1 = map(delayed(_ -> rand(10^3), cache=true), workers())
    sum1 = delayed((x...)->sum([x...]))(map(delayed(sum), thunks1)...)
    thunks2 = map(delayed(-), thunks1)
    sum2 = delayed((x...)->sum([x...]))(map(delayed(sum), thunks2)...)
    s1 = collect(sum1)
    @test -s1 == collect(sum2)
    @test s1 == collect(sum1)
    @test -collect(sum1) == collect(sum2)

    # Issue #246
    z = delayed(identity)(10)
    zz = delayed(identity; cache=true)(10)
    @test collect(z) == collect(z)
    @test collect(zz) == collect(zz)
end


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

const DiagMatLike{T} = Union{Diagonal{T}, PDiagMat{T}}

# get_diagonal is only defined for diagonal matrices, unlike LinearAlgebra.diag:
get_diagonal(A::DiagMatLike) = diag(A)

without_chol(A::AbstractMatrix) = A
without_chol(A::PDMat) = A.mat
without_chol(A::PDiagMat) = Diagonal(diag(A))
without_chol(A::PDSparseMat) = A.mat
without_chol(A::ScalMat) = Diagonal(Fill(A.value, A.dim))

without_chol_pullback(thunked_Δy) = (NoTangent(), thunked_Δy)
function ChainRulesCore.rrule(::typeof(without_chol), A::AbstractMatrix)
    return without_chol(A), without_chol_pullback
end

cholesky_L(A::AbstractMatrix) = cholesky(A).L
cholesky_L(A::AbstractSparseMatrix) = sparse(cholesky(A).L)
cholesky_L(A::DiagMatLike) = Diagonal(sqrt.(get_diagonal(A)))

# Workaround for missing SVector() pullback in Zygote:
_svector(x::NTuple{N,T}) where {N,T<:Real} = SVector(x)

function ChainRulesCore.rrule(::typeof(_svector), x::NTuple{N,T}) where {N,T<:Real}
    function _svector_pullback(thunked_Δy)
        Δy = unthunk(thunked_Δy)
        return ChainRulesCore.NoTangent(), (Δy...,)
    end

    return _svector(x), _svector_pullback
end

# ToDo: Check Zygote.pullback performance of `reduce(vcat, A)`
_flatten_vec_of_vec(A::AbstractVector{<:AbstractVector{<:Real}}) = reduce(vcat, A)

_flatten_vec_of_vec(A::AbstractVector{U}) where {N,T<:Real,U<:Union{NTuple{N,T},StaticVector{N,T}}} = copy(reinterpret(T, A))

function ChainRulesCore.rrule(::typeof(_flatten_vec_of_vec), A::AbstractVector{U}) where {N,T<:Real,U<:Union{NTuple{N,T},StaticVector{N,T}}}
    function _flatten_vec_of_vec_pullback(thunked_Δy)
        Δy = unthunk(thunked_Δy)
        return ChainRulesCore.NoTangent(), copy(reinterpret(U, Δy))
    end

    return _flatten_vec_of_vec(A), _flatten_vec_of_vec_pullback
end

# ToDo (if ArraysOfArrays gets added to deps): _flatten_vec_of_vec(A::VectorOfSimilarVectors{<:Real}) = vec(flatview(C))