##########################################################################
# Copyright 2017 Samuel Ridler.
#
# 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.
##########################################################################

# initialise travel data for fGraph
# arcTravelTimes[i,j] gives travel time along arc j (in fGraph) for travel mode i
function initFNetTravels!(net::Network, arcTravelTimes::Array{Float,2})
	numModes = size(arcTravelTimes, 1) # shorthand
	@assert(length(net.fGraph.arcs) == size(arcTravelTimes, 2)) # check number of arcs

	# create fNetTravels
	net.fNetTravels = [NetTravel(false) for i = 1:numModes]
	for i = 1:numModes
		fNetTravel = net.fNetTravels[i]
		fNetTravel.modeIndex = i
		fNetTravel.arcTimes = arcTravelTimes[i,:]
	end
end

# return true if fNode i is in rGraph, false otherwise
function isFNodeInRGraph(net::Network, i::Int)
	return (net.fNodeRNode[i] != nullIndex)
end

# create the reduced graph from the full graph, for the given network
function createRGraphFromFGraph!(net::Network)
	# assumptions:
	# - no nodes are connected to themselves
	# - fGraph has no duplicate arcs

	net.rGraph = Graph(true)

	# shorthand names:
	fGraph = net.fGraph
	fNodes = fGraph.nodes
	fArcs = fGraph.arcs
	numFNodes = length(fNodes)
	rGraph = net.rGraph

	##############################################################################

	# some useful fGraph connectivity data

	# shorthand:
	fNodesOut = fGraph.light.fadjlist # fNodesOut[i] gives indices of fNodes incident to arcs outgoing from fNode i
	fNodesIn = fGraph.light.badjlist # fNodesIn[i] gives indices of fNodes incident to arcs incoming to fNode i

	fNodeNumAdjacent = [length(union(fNodesIn[i], fNodesOut[i])) for i = 1:numFNodes] # number of fNodes adjacent to each fNode

	# check that no node is connected to itself
	for i = 1:numFNodes
		@assert(!in(i, fNodesIn[i]))
	end

	##############################################################################

	# find intersection/decision nodes in fGraph, add to rGraph

	# only keep nodes with >= 3 adjacent nodes, or 2 adjacent nodes if the incoming and outgoing arcs cannot be replaced
	# Also, for easier programming, put leaf nodes in fGraph into rGraph, hopefully there are not too many...
	keepNode = Vector{Bool}(undef, numFNodes) # find which fNodes belong in rGraph
	keepNode[:] .= false
	for i = 1:numFNodes
		if fNodeNumAdjacent[i] >= 3
			keepNode[i] = true
		elseif fNodeNumAdjacent[i] == 2 && (length(fNodesIn[i]) + length(fNodesOut[i]) == 3)
			keepNode[i] = true
		elseif fNodeNumAdjacent[i] == 1
			keepNode[i] = true
		end
	end
	if !any(keepNode) keepNode[1] = true end # if there are still no rNodes, make one (can happen if network is just a loop)
	rGraph.nodes = deepcopy(fNodes[keepNode]) # need to renumber rNode indices later
	net.rNodeFNode = findall(keepNode)
	net.fNodeRNode = [nullIndex for i = 1:numFNodes]
	net.fNodeRNode[keepNode] = [1:sum(keepNode);]

	# special consideration: need to add one decision node to any loop with currently only one rNode
	# Otherwise, fNodes in the loop still have a decision as to how to reach the rNode
	for ri = 1:length(rGraph.nodes)
		fi = net.rNodeFNode[ri]
		for fj = fNodesOut[fi] # fNodes outgoing from fNode[fi]
			firstFNode = fj
			prevFNode = fi
			fNode = fj
			# traverse fGraph from rNode[ri], until another rNode or dead end (leaf node?) reached
			while !isFNodeInRGraph(net, fNode) # !(isFNodeInRGraph(net, fNode) || fNodeNumAdjacent[fNode] == 1)
				# traverse along an arc to next node
				nextFNode = first(setdiff(fNodesOut[fNode], prevFNode)) # fNodesOut[fNode] = indices of fNodes outgoing from fNode[fNode]
				prevFNode = fNode
				fNode = nextFNode
			end
			# if we have looped back to same rNode, add firstFNode as decision nodes
			if ri == net.fNodeRNode[fNode]
				keepNode[firstFNode] = true
				net.fNodeRNode[firstFNode] = numFNodes + 1 # a temporary value that will make isFNodeInRGraph(net, firstFNode) return true
				@assert(isFNodeInRGraph(net, firstFNode))
			end
		end
	end

	rGraph.nodes = deepcopy(fNodes[keepNode])
	# renumber rNode indices
	for i = 1:length(rGraph.nodes)
		rGraph.nodes[i].index = i
	end
	net.rNodeFNode = findall(keepNode)
	net.fNodeRNode = [nullIndex for i = 1:numFNodes]
	net.fNodeRNode[keepNode] = [1:sum(keepNode);]

	# shorthand names:
	rNodes = rGraph.nodes
	fNodeRNode = net.fNodeRNode
	rNodeFNode = net.rNodeFNode
	numRNodes = length(rNodes)

	@assert(numRNodes > 0) # should not be possible at this point, and would cause problems below if false

	# check index conversion between fNodes and rNodes
	@assert(all([i == fNodeRNode[rNodeFNode[rNodes[i].index]] for i = 1:numRNodes]))
	# @assert(all([i == fNodeRNode[rNodeFNode[i]] for i = 1:numRNodes]))

	##############################################################################

	# calculate rArcs, rArcFNodes, fNodeRArcs
	# find outgoing arcs from each rNode to other rNodes
	net.fNodeFromRNodes = [Vector{Int}() for i = 1:numFNodes]
	net.rArcFNodes = []
	net.fNodeRArcs = [Vector{Int}() for i = 1:numFNodes]
	rArcFNodes = net.rArcFNodes # shorthand
	fNodeRArcs = net.fNodeRArcs # shorthand
	rArcs = rGraph.arcs # shorthand
	numRArcs = 0
	for ri = 1:numRNodes
		fi = rNodeFNode[ri]
		for fj = fNodesOut[fi] # fNodes outgoing from fNode[fi]
			numRArcs += 1 # will create another rArc
			# traverse fGraph from rNode[ri], until another rNode or dead end (leaf node?) reached
			fNodesOnRArc = [fi] # indices of fNodes on rArc
			prevFNode = fi
			fNode = fj
			while !isFNodeInRGraph(net, fNode)
				push!(fNodeRArcs[fNode], numRArcs)
				push!(fNodesOnRArc, fNode)
				# traverse along an arc to next node
				nextFNode = first(setdiff(fNodesOut[fNode], prevFNode)) # fNodesOut[fNode] = indices of fNodes outgoing from fNode
				prevFNode = fNode
				fNode = nextFNode
			end
			push!(fNodesOnRArc, fNode)
			rj = fNodeRNode[fNode] # rNode at end of rArc, fNodeRNode[fNode] = nullIndex if fNode[fNode] is not in rGraph
			@assert(rj != nullIndex) # should not happen; earlier we put all leaf nodes in fGraph in rGraph
			@assert(ri != rj) # ri == rj would mean that arc is a loop, should have added enough rNodes to remove this problem earlier

			# add rArc
			arc = Arc()
			arc.index = numRArcs
			arc.fromNodeIndex = ri
			arc.toNodeIndex = rj
			push!(rArcs, arc)
			push!(rArcFNodes, fNodesOnRArc)
		end
	end

	@assert(numRArcs > 0)

	# calculate rArcFNodeIndex from rArcFNodes
	net.rArcFNodeIndex = [Dict{Int,Int}() for i = 1:numRArcs]
	for ri = 1:numRArcs
		for (i, fi) in enumerate(rArcFNodes[ri])
			net.rArcFNodeIndex[ri][fi] = i
		end
	end

	# set fNodeRArcs for each rNode (for corresponding fNode)
	for i = 1:numRArcs
		fromRNode = rArcs[i].fromNodeIndex
		toRNode = rArcs[i].toNodeIndex
		push!(fNodeRArcs[rNodeFNode[fromRNode]], i)
		push!(fNodeRArcs[rNodeFNode[toRNode]], i)
	end

	# check that fNodeRArcs is not empty for rNodes, and other nodes have 1-2 entries
	for fi = 1:numFNodes
		if isFNodeInRGraph(net, fi)
			@assert(!isempty(fNodeRArcs[fi]))
		else
			@assert(length(fNodeRArcs[fi]) >= 1 && length(fNodeRArcs[fi]) <= 2)
		end
	end

	# check that start and end fNodes on each rArc are also in rGraph, and arcs are numbered 1 to n
	for i = 1:numRArcs
		@assert(isFNodeInRGraph(net, rArcFNodes[i][1]))
		@assert(isFNodeInRGraph(net, rArcFNodes[i][end]))
		@assert(rArcs[i].index == i)
	end

	##############################################################################

	# calculate fNodeToRNodes and fNodeFromRNodes

	net.fNodeToRNodes = [Vector{Int}() for i = 1:numFNodes]
	net.fNodeFromRNodes = [Vector{Int}() for i = 1:numFNodes]
	net.fNodeToRNodeNextFNode = [Dict{Int,Int}() for i = 1:numFNodes]
	@unpack fNodeToRNodes, fNodeFromRNodes, fNodeToRNodeNextFNode = net

	# traverse each rArc
	for rArc = rArcs
		fNodesOnRArc = rArcFNodes[rArc.index]
		@assert(isFNodeInRGraph(net, fNodesOnRArc[1]))
		@assert(isFNodeInRGraph(net, fNodesOnRArc[end]))
		startRNode = fNodeRNode[fNodesOnRArc[1]]
		endRNode = fNodeRNode[fNodesOnRArc[end]]
		# for fNodes on arc (excluding rNodes):
		for i = 2:length(fNodesOnRArc)-1
			fNode = fNodesOnRArc[i]
			push!(fNodeToRNodes[fNode], endRNode)
			push!(fNodeFromRNodes[fNode], startRNode)
			fNodeToRNodeNextFNode[fNode][endRNode] = fNodesOnRArc[i+1]
		end
	end
	# for rNodes, set corresponding fNode to have fNodeToRNodes and fNodeFromRNodes return the rNode
	for ri = 1:numRNodes
		fi = rNodeFNode[ri]
		fNodeFromRNodes[fi] = [ri]
		fNodeToRNodes[fi] = [ri]
	end

	# check that fNodeFromRNodes and fNodeToRNodes have been filled
	for i = 1:numFNodes
		@assert(!isempty(fNodeFromRNodes[i]))
		@assert(!isempty(fNodeToRNodes[i]))
	end

	##############################################################################

	# arc distances

	# set rArcFArcs
	rArcFArcs = net.rArcFArcs = [Vector{Int}() for i = 1:numRArcs]
	for rArc in rArcs
		fNodesOnRArc = rArcFNodes[rArc.index]
		for i = 1:length(fNodesOnRArc)-1
			fNode = fNodesOnRArc[i]
			nextFNode = fNodesOnRArc[i+1]
			push!(rArcFArcs[rArc.index], fGraph.nodePairArcIndex[fNode, nextFNode])
		end
		@assert(length(rArcFArcs[rArc.index]) >= 1) # at least one fArc per rArc
	end
	# @assert(sort(vcat(rArcFArcs...)) == 1:length(fArcs)) # rArcFArcs contains each fArc exactly once

	# set arc distances
	@assert(all(arc -> !isnan(arc.distance) && arc.distance >= 0, fArcs))
	for rArc in rArcs
		rArc.distance = 0.0
		for fArcIndex in rArcFArcs[rArc.index]
			rArc.distance += fArcs[fArcIndex].distance
		end
	end

	##############################################################################

	# set fNodeToRNodeDist, fNodeFromRNodeDist
	fNodeToRNodeDist = net.fNodeToRNodeDist = [Dict{Int, Float}() for i = 1:numFNodes]
	fNodeFromRNodeDist = net.fNodeFromRNodeDist = [Dict{Int, Float}() for i = 1:numFNodes]

	for rArc in rArcs
		# traverse arc, calculate fNodeFromRNodeDist and fNodeToRNodeDist
		startRNode = rArc.fromNodeIndex
		endRNode = rArc.toNodeIndex
		dist = 0.0
		for fArcIndex in rArcFArcs[rArc.index]
			fArc = fArcs[fArcIndex]
			fNodeFromRNodeDist[fArc.fromNodeIndex][startRNode] = dist
			dist += fArc.distance
			fNodeToRNodeDist[fArc.toNodeIndex][endRNode] = rArc.distance - dist
		end
	end

	# for rNodes, check that corresponding fNode has fNodeFromRNodeDist and fNodeToRNodeDist return 0
	for ri = 1:numRNodes
		fi = rNodeFNode[ri]
		@assert(fNodeFromRNodeDist[fi][ri] == 0.0)
		@assert(fNodeToRNodeDist[fi][ri] == 0.0)
	end

	# for fNodes along each rArc, fNodeFromRNodeDist and fNodeToRNodeDist should add to rArc distance
	for rArc in rArcs
		fNodesOnRArc = rArcFNodes[rArc.index]
		startRNode = fNodeRNode[fNodesOnRArc[1]]
		endRNode = fNodeRNode[fNodesOnRArc[end]]
		for i = 2:length(fNodesOnRArc)-1
			fNode = fNodesOnRArc[i]
			@assert(isapprox(fNodeFromRNodeDist[fNode][startRNode] + fNodeToRNodeDist[fNode][endRNode], rArc.distance; rtol = eps(Float)))
		end
	end

	##############################################################################

	# calculate fields for rGraph: arcDists, light
	initGraph!(rGraph)
end

# calculate arcTimes, spTimes, and spSuccs, for rNetTravels
# also, calculate fNodeToRNodeTime, and fNodeFromRNodeTime, for fNetTravels...
# can only call after rGraph is created
function createRNetTravelsFromFNetTravels!(net::Network;
	rNetTravelsLoaded::Vector{NetTravel} = NetTravel[])

	# shorthand names:
	@unpack fGraph, rGraph, fNetTravels, rArcFNodes, fNodeRNode, rNodeFNode, fNodeToRNodes = net
	numTravelModes = length(fNetTravels)
	numFNodes = length(fGraph.nodes)
	numRNodes = length(rGraph.nodes)
	numRArcs = length(rGraph.arcs)

	rNetTravels = net.rNetTravels = [NetTravel(true) for i = 1:numTravelModes]

	@assert(length(fGraph.arcs) == length(fNetTravels[1].arcTimes))

	for travelModeIndex = 1:numTravelModes
		# shorthand:
		fNetTravel = fNetTravels[travelModeIndex]
		rNetTravel = rNetTravels[travelModeIndex]

		rNetTravel.arcTimes = Vector{Float}(undef, numRArcs)
		fNetTravel.fNodeToRNodeTime = [Dict{Int, Float}() for i = 1:numFNodes]
		fNetTravel.fNodeFromRNodeTime = [Dict{Int, Float}() for i = 1:numFNodes]
		fNetTravel.rArcFNodesTimes = [Vector{Float}() for i = 1:numRArcs]
		fNodeToRNodeTime = fNetTravel.fNodeToRNodeTime # shorthand
		fNodeFromRNodeTime = fNetTravel.fNodeFromRNodeTime # shorthand
		rArcFNodesTimes = fNetTravel.rArcFNodesTimes

		rNetTravel.modeIndex = fNetTravel.modeIndex
		rNetTravel.arcDists = rGraph.arcDists
		rNetTravel.fadjList = rGraph.fadjList

		for rArcIndex = 1:numRArcs
			fNodesOnRArc = rArcFNodes[rArcIndex]
			startRNode = fNodeRNode[fNodesOnRArc[1]]
			endRNode = fNodeRNode[fNodesOnRArc[end]]

			# traverse arc, calculate rArcTime and fNodeFromRNodeTime
			rArcTime = 0.0
			for i = 1:length(fNodesOnRArc)-1
				fNode = fNodesOnRArc[i]
				nextFNode = fNodesOnRArc[i+1]
				fNodeFromRNodeTime[fNode][startRNode] = rArcTime
				push!(rArcFNodesTimes[rArcIndex], rArcTime)
				fArcIndex = fGraph.nodePairArcIndex[fNode, nextFNode]
				rArcTime += fNetTravel.arcTimes[fArcIndex]
			end
			push!(rArcFNodesTimes[rArcIndex], rArcTime)
			rNetTravel.arcTimes[rArcIndex] = rArcTime

			# calculate fNodeToRNodeTime
			for i = 2:length(fNodesOnRArc)-1
				fNode = fNodesOnRArc[i]
				fNodeToRNodeTime[fNode][endRNode] = rArcTime - fNodeFromRNodeTime[fNode][startRNode]
			end
			fNodeToRNodeTime[fNodesOnRArc[end]][endRNode] = 0.0 # time to self is 0

			# check length of rArcFNodesTimes is same as length of rArcFNodes for current arc
			@assert(length(fNodesOnRArc) == length(rArcFNodesTimes[rArcIndex]))
		end

		# # for rNodes, set corresponding fNode to have fNodeFromRNodeTime and fNodeToRNodeTime return 0
		# for ri = 1:numRNodes
			# fi = rNodeFNode[ri]
			# # fNodeToRNodeTime[fi][ri] = 0.0 # should already be true?
			# # fNodeFromRNodeTime[fi][ri] = 0.0 # should already be true?
		# end

		# for rNodes, check that corresponding fNode has fNodeFromRNodeTime and fNodeToRNodeTime return 0
		for ri = 1:numRNodes
			fi = rNodeFNode[ri]
			@assert(fNodeToRNodeTime[fi][ri] == 0.0)
			@assert(fNodeFromRNodeTime[fi][ri] == 0.0)
		end

		# for fNodes along each rArc, fNodeFromRNodeTime and fNodeToRNodeTime should add to rArc time
		for rArcIndex = 1:numRArcs
			rArcTime = rNetTravel.arcTimes[rArcIndex]
			fNodesOnRArc = rArcFNodes[rArcIndex]
			startRNode = fNodeRNode[fNodesOnRArc[1]]
			endRNode = fNodeRNode[fNodesOnRArc[end]]
			for i = 2:length(fNodesOnRArc)-1
				fNode = fNodesOnRArc[i]
				@assert(isapprox(fNodeFromRNodeTime[fNode][startRNode] + fNodeToRNodeTime[fNode][endRNode], rArcTime; rtol = eps(Float)))
			end
		end

		if rNetTravelsLoaded == []
			# calculate shortest path data for rNetTravel
			calcRNetTravelShortestPaths!(net, rNetTravel)
		else
			rNetTravelLoaded = rNetTravelsLoaded[travelModeIndex]
			@assert(rNetTravel.arcTimes == rNetTravelLoaded.arcTimes)
			@assert(rNetTravel.arcDists == rNetTravelLoaded.arcDists)
			@assert(rNetTravel.fadjList == rNetTravelLoaded.fadjList)

			# set rNetTravel values by rNetTravelLoaded
			for fname in [:spFadjIndex, :spNodePairArcIndex, :spFadjArcList]
				setfield!(rNetTravel, fname, getfield(rNetTravelLoaded, fname))
			end

			rNetTravel.spTimes = calcRNetTravelShortestPathTimes(net, rNetTravel)
			rNetTravel.spDists = calcRNetTravelShortestPathDists(net, rNetTravel)
		end

	end
end

# For a given rNetTravel, calculate and store shortest path data
# Uses repeated Dijkstra's shortest path algorithm,
# for sparse graphs, this is faster than Floyd-Warshall algorithm
function calcRNetTravelShortestPaths!(net::Network, rNetTravel::NetTravel)
	@assert(rNetTravel.isReduced)

	# shorthand
	rGraph = net.rGraph
	rArcs = rGraph.arcs
	rArcTimes = rNetTravel.arcTimes
	fadjList = rGraph.fadjList
	numRNodes = length(rGraph.nodes)

	# create reverse graph
	revRGraph = LightGraphs.reverse(rGraph.light)

	# create arc time matrix for reverse graph, for use by LightGraphs Dijkstra SP function
	revRArcsTimes = spzeros(FloatSpTime, numRNodes, numRNodes)
	for rArc in rArcs
		# add time of this arc to revRArcsTimes if first/shortest arc connecting the two nodes
		# (there can be multiple arcs connecting a pair of nodes)
		t1 = revRArcsTimes[rArc.toNodeIndex, rArc.fromNodeIndex] # 0 if not yet filled
		t2 = rArcTimes[rArc.index]
		revRArcsTimes[rArc.toNodeIndex, rArc.fromNodeIndex] = (t1 == 0 ? t2 : min(t1, t2))
	end

	# calculate shortest paths for each possible destination node
	spData = LightGraphs.DijkstraState{Float} # init
	spTimes = rNetTravel.spTimes = Array{FloatSpTime,2}(undef, numRNodes, numRNodes)
	spSuccs = Array{IntRNode,2}(undef, numRNodes, numRNodes)
	for i = 1:numRNodes
		spData = LightGraphs.dijkstra_shortest_paths(revRGraph, i, revRArcsTimes)
		spTimes[:,i] = spData.dists
		spSuccs[:,i] = spData.parents
	end

	# change spSuccs where needed
	# where a path i -> l has spSuccs[i,l] = k, but spSuccs[i,k] = j,
	# need to change all cases where spSuccs[i,-] = k to instead have spSuccs[i,-] = j
	# this case can happen when an arc (i,k) has the same travel time as a path (i,j,k)
	for rArc in rArcs
		i = rArc.fromNodeIndex
		j = k = rArc.toNodeIndex
		while j != spSuccs[i,j]
			j = spSuccs[i,j]
		end
		if j != k
			# for any paths from i that have successor node k, change successor to be j
			spSuccs[i, spSuccs[i,:] .== k] .= j
		end
	end

	# set spNodePairArcIndex and spFadjArcList
	spNodePairArcIndex = rNetTravel.spNodePairArcIndex = spzeros(Int, numRNodes, numRNodes)
	spFadjArcList = rNetTravel.spFadjArcList = deepcopy(fadjList)
	for rArc in rArcs
		i = rArc.fromNodeIndex
		j = rArc.toNodeIndex
		t = FloatSpTime(rArcTimes[rArc.index])
		@assert(spTimes[i,j] <= t)
		if spSuccs[i,j] == j && spTimes[i,j] == t
			spNodePairArcIndex[i,j] = rArc.index
			spFadjArcList[i][findfirst(isequal(j), fadjList[i])] = rArc.index
		end
	end

	# set spFadjIndex
	spFadjIndex = rNetTravel.spFadjIndex = Array{IntFadj,2}(undef, numRNodes, numRNodes)
	for i = 1:numRNodes, j = 1:numRNodes
		spFadjIndex[i,j] = (i == j ? nullIndex : findfirst(isequal(spSuccs[i,j]), fadjList[i]))
	end

	# check stored shortest path times are same as those from traversing stored shortest paths
	# this can be slow, and should not be necessary unless shortest path code above has changed
	if checkMode && false
		checkRNetTravelShortestPathTimes(net, rNetTravel)
	end

	# check assumptions:
	for i = 1:numRNodes, j = [1:i-1;i+1:numRNodes]
		@assert(0 < spTimes[i,j] < Inf)
	end

	rNetTravel.spDists = calcRNetTravelShortestPathDists(net, rNetTravel)
end

# For a given travel mode, check that the stored shortest paths data
# gives paths with the same travel time as stored in rNetTravel.spTimes
function checkRNetTravelShortestPathTimes(net::Network, rNetTravel::NetTravel)
	# shorthand:
	numRNodes = length(net.rGraph.nodes)
	travelModeIndex = rNetTravel.modeIndex
	@unpack spTimes, spNodePairArcIndex = rNetTravel

	@assert(size(spTimes) == (numRNodes, numRNodes))

	# check each pair of start and end nodes
	for startRNode = 1:numRNodes, endRNode = 1:numRNodes
		if startRNode == endRNode
			@assert(spTimes[startRNode, endRNode] == 0.0)
			@assert(spNodePairArcIndex[startRNode, endRNode] == 0)
		else
			# follow path from start to end node, calculate and compare travel time
			i = startRNode
			t = 0.0
			while i != endRNode
				j = shortestPathNextRNode(net, travelModeIndex, i, endRNode)
				rArcIndex = shortestPathNextRArc(net, travelModeIndex, i, endRNode)
				@assert(rArcIndex == rNetTravel.spNodePairArcIndex[i,j])
				# @assert(j == net.rGraph.arcs[rArcIndex].toNodeIndex)
				t += rNetTravel.arcTimes[rArcIndex]
				i = j # go to next node
			end
			@assert(isapprox(spTimes[startRNode, endRNode], t; rtol = eps(FloatSpTime)))
		end
	end
end

# Calculate time of shortest path between each pair of nodes for rNetTravel.
# Requires some shortest path data (spFadjIndex, spNodePairArcIndex) to be calculated already.
# See also: calcRNetTravelShortestPathDists
function calcRNetTravelShortestPathTimes(net::Network, rNetTravel::NetTravel)
	@assert(rNetTravel.isReduced)

	# shorthand
	rGraph = net.rGraph
	fadjList = rGraph.fadjList
	n = length(rGraph.nodes)
	@unpack spFadjIndex, spFadjArcList, spNodePairArcIndex, arcTimes = rNetTravel

	spTimes = fill(FloatSpTime(Inf), n, n) # spTimes[i,j] = shortest path time from node i to j
	for i = 1:n spTimes[i,i] = 0 end
	for arc in rGraph.arcs
		i = arc.fromNodeIndex
		j = arc.toNodeIndex
		if spNodePairArcIndex[i,j] == arc.index
			spTimes[i,j] = arcTimes[arc.index]
		end
	end

	spNextNode = zeros(Int, n)
	spNextArcTime = zeros(Float, n)
	visited = fill(false, n) # visited[i] = true if node i has been visited and so had time populated
	path = zeros(Int, n) # path node indices
	for k = 1:n # root node of shortest path tree, find time from each node to this node
		# get data for paths towards node k
		spNextNode[k] = 0
		for j = 1:n
			if j == k continue end
			fadjIndex = spFadjIndex[j,k]
			spNextNode[j] = fadjList[j][fadjIndex]
			spNextArcTime[j] = arcTimes[spFadjArcList[j][fadjIndex]]
		end

		fill!(visited, false)
		visited[k] = true
		for i = 1:n
			if i == k || visited[i] continue end

			# trace path from i to k, or to node that already has data for path to k
			j = i
			p = 0
			while j != k && !visited[j]
				p += 1
				path[p] = j
				j = spNextNode[j]
			end

			# backtrack path, setting time from node j to k
			t = spTimes[j,k]
			while p > 0
				j = path[p]
				t += spNextArcTime[j]
				spTimes[j,k] = t
				visited[j] = true
				p -= 1
			end
		end
	end

	@assert(all(t -> t >= 0 && t != FloatSpTime(Inf), spTimes))

	return spTimes
end

# Calculate distance of shortest path between each pair of nodes for rNetTravel.
# Requires some shortest path data (spFadjIndex, spNodePairArcIndex) to be calculated already.
# See also: calcRNetTravelShortestPathTimes
function calcRNetTravelShortestPathDists(net::Network, rNetTravel::NetTravel)
	@assert(rNetTravel.isReduced)

	# shorthand
	rGraph = net.rGraph
	fadjList = rGraph.fadjList
	n = length(rGraph.nodes)
	@unpack spFadjIndex, spFadjArcList, spNodePairArcIndex, arcDists = rNetTravel

	spDists = fill(FloatSpDist(Inf), n, n) # spDists[i,j] = shortest path distance from node i to j
	for i = 1:n spDists[i,i] = 0 end
	for arc in rGraph.arcs
		i = arc.fromNodeIndex
		j = arc.toNodeIndex
		if spNodePairArcIndex[i,j] == arc.index
			spDists[i,j] = arc.distance
		end
	end

	spNextNode = zeros(Int, n)
	spNextArcDist = zeros(Float, n)
	visited = fill(false, n) # visited[i] = true if node i has been visited and so had distance populated
	path = zeros(Int, n) # path node indices
	for k = 1:n # root node of shortest path tree, find distance from each node to this node
		# get data for paths towards node k
		spNextNode[k] = 0
		for j = 1:n
			if j == k continue end
			fadjIndex = spFadjIndex[j,k]
			spNextNode[j] = fadjList[j][fadjIndex]
			spNextArcDist[j] = arcDists[spFadjArcList[j][fadjIndex]]
		end

		fill!(visited, false)
		visited[k] = true
		for i = 1:n
			if i == k || visited[i] continue end

			# trace path from i to k, or to node that already has data for path to k
			j = i
			p = 0
			while j != k && !visited[j]
				p += 1
				path[p] = j
				j = spNextNode[j]
			end

			# backtrack path, setting distance from node j to k
			d = spDists[j,k]
			while p > 0
				j = path[p]
				d += spNextArcDist[j]
				spDists[j,k] = d
				visited[j] = true
				p -= 1
			end
		end
	end

	@assert(all(d -> d >= 0 && d != FloatSpDist(Inf), spDists))

	return spDists
end

# for the shortest path from startRNode to endRNode,
# travelling with the given travel mode,
# return the index of the successor rNode of startRNode
function shortestPathNextRNode(net::Network, travelModeIndex::Int, startRNode::Int, endRNode::Int)
	@assert(startRNode != endRNode)

	rNetTravel = net.rNetTravels[travelModeIndex] # shorthand
	fadjIndex = rNetTravel.spFadjIndex[startRNode, endRNode]
	return net.rGraph.fadjList[startRNode][fadjIndex]
end

# for the shortest path from startRNode to endRNode,
# travelling with the given travel mode,
# return the index of the rArc to travel along from startRNode
function shortestPathNextRArc(net::Network, travelModeIndex::Int, startRNode::Int, endRNode::Int)
	@assert(startRNode != endRNode)

	rNetTravel = net.rNetTravels[travelModeIndex] # shorthand
	fadjIndex = rNetTravel.spFadjIndex[startRNode, endRNode]
	return rNetTravel.spFadjArcList[startRNode][fadjIndex]
end

# Given the travel mode, and indices of start and end node (in full graph),
# returns travel time for shortest path, and the first and last rNodes in the path (if any)
function shortestPathData(net::Network, travelModeIndex::Int, startFNode::Int, endFNode::Int)

	# shorthand:
	fNetTravel = net.fNetTravels[travelModeIndex]
	rNetTravel = net.rNetTravels[travelModeIndex]
	@unpack fNodeFromRNodeTime, fNodeToRNodeTime = fNetTravel
	@unpack fNodeRArcs, fNodeToRNodes, fNodeFromRNodes, isFNodeCommon, fNodeCommonFNodeIndex = net

	# return values:
	shortestTravelTime = Inf
	rNodes = (nullIndex, nullIndex) # first and last rNodes in path

	if startFNode == endFNode
		shortestTravelTime = 0.0
		rNodes = (nullIndex, nullIndex)
		return shortestTravelTime, rNodes
	end

	# if startFNode or endFNode are common, can look up stored shortest path data
	if isFNodeCommon[startFNode]
		i = fNodeCommonFNodeIndex[startFNode]
		shortestTravelTime = fNetTravel.commonFNodeToFNodeTime[i, endFNode]
		rNodes = fNetTravel.commonFNodeToFNodeRNodes[i, endFNode]
		return shortestTravelTime, rNodes
	elseif isFNodeCommon[endFNode]
		i = fNodeCommonFNodeIndex[endFNode]
		shortestTravelTime = fNetTravel.fNodeToCommonFNodeTime[startFNode, i]
		rNodes = fNetTravel.fNodeToCommonFNodeRNodes[startFNode, i]
		return shortestTravelTime, rNodes
	end

	# see if travel is possible without using any nodes in rGraph
	if !isFNodeInRGraph(net, startFNode) && !isFNodeInRGraph(net, endFNode) && fNodeRArcs[startFNode] == fNodeRArcs[endFNode]
		rArc = findRArcFromFNodeToFNode(net, startFNode, endFNode)
		if rArc != nullIndex
			# calculate travel time on rArc
			fromRNode = net.rGraph.arcs[rArc].fromNodeIndex
			toRNode = net.rGraph.arcs[rArc].toNodeIndex
			travelTime = fNodeFromRNodeTime[endFNode][fromRNode] - fNodeFromRNodeTime[startFNode][fromRNode]
			@assert(travelTime > 0)
			@assert(isapprox(travelTime + fNodeToRNodeTime[endFNode][toRNode], fNodeToRNodeTime[startFNode][toRNode]; rtol = eps(Float)))

			shortestTravelTime = travelTime # not necessarily shortest travel time, still need to check paths that use rNodes
			rNodes = (nullIndex, nullIndex)
		end
	end

	# for each pair of starting and ending rNodes (near given fNodes), calculate travel time
	for startRNode = fNodeToRNodes[startFNode], endRNode = fNodeFromRNodes[endFNode]
		travelTime = rNetTravel.spTimes[startRNode, endRNode]
		travelTime += fNodeToRNodeTime[startFNode][startRNode]
		travelTime += fNodeFromRNodeTime[endFNode][endRNode]
		if travelTime < shortestTravelTime
			shortestTravelTime = travelTime
			rNodes = (startRNode, endRNode)
		end
	end

	return shortestTravelTime, rNodes
end

# Returns the travel time of the shortest path between two nodes; see shortestPathData function.
function shortestPathTravelTime(net::Network, travelModeIndex::Int, startFNode::Int, endFNode::Int)
	return shortestPathData(net, travelModeIndex, startFNode, endFNode)[1]
end

# Returns distance of shortest path (shortest by time) between two nodes.
# The startRNode and endRNode of shortest path can be given as kwargs if known.
function shortestPathDistance(net::Network, travelModeIndex::Int, startFNode::Int, endFNode::Int;
	startRNode::Int = nullIndex - 1, endRNode::Int = nullIndex - 1)
	# note that the default value of startRNode and endRNode are nullIndex - 1 and not just nullIndex;
	# this is to allow them to be set to nullIndex which indicates that the shortest path has no rNodes.

	# shorthand:
	rArcs = net.rGraph.arcs
	@unpack fNodeToRNodeDist, fNodeFromRNodeDist, isFNodeCommon, fNodeCommonFNodeIndex = net
	fNetTravel = net.fNetTravels[travelModeIndex]
	rNetTravel = net.rNetTravels[travelModeIndex]

	dist = 0.0
	if startFNode == endFNode return dist end

	# if startFNode or endFNode are common, can look up stored shortest path distance
	if isFNodeCommon[startFNode]
		i = fNodeCommonFNodeIndex[startFNode]
		return fNetTravel.commonFNodeToFNodeDist[i, endFNode]
	elseif isFNodeCommon[endFNode]
		i = fNodeCommonFNodeIndex[endFNode]
		return fNetTravel.fNodeToCommonFNodeDist[startFNode, i]
	end

	if startRNode == nullIndex - 1 || endRNode == nullIndex - 1
		@assert(startRNode == nullIndex - 1 && endRNode == nullIndex - 1)
		# get start and end rNodes
		(shortestTravelTime, rNodes) = shortestPathData(net, travelModeIndex, startFNode, endFNode)
		(startRNode, endRNode) = rNodes
	elseif startRNode == nullIndex || endRNode == nullIndex
		@assert(startRNode == nullIndex && endRNode == nullIndex)
	else
		@assert(in(startRNode, net.fNodeToRNodes[startFNode]))
		@assert(in(endRNode, net.fNodeFromRNodes[endFNode]))
	end

	if startRNode == nullIndex
		# shortest path does not use any rNodes
		@assert(endRNode == nullIndex)
		rArcIndex = findRArcFromFNodeToFNode(net, startFNode, endFNode)
		@assert(rArcIndex != nullIndex) # otherwise would need to use an rNode
		rArc = rArcs[rArcIndex]
		fromRNode = rArc.fromNodeIndex
		toRNode = rArc.toNodeIndex
		dist = fNodeFromRNodeDist[endFNode][fromRNode] - fNodeFromRNodeDist[startFNode][fromRNode]
		@assert(dist >= 0)
		@assert(isapprox(dist + fNodeToRNodeDist[endFNode][toRNode], fNodeToRNodeDist[startFNode][toRNode]; rtol = eps(Float)))
	else
		# shortest path uses at least one rNode
		@assert(startRNode != nullIndex && endRNode != nullIndex)
		dist += net.fNodeToRNodeDist[startFNode][startRNode]
		dist += rNetTravel.spDists[startRNode, endRNode]
		dist += net.fNodeFromRNodeDist[endFNode][endRNode]
	end

	return dist
end

# Find and return shortest path (as represented by list of nodes) from startFNode to endFNode,
# starting at startFNode at startTime, with a given travel mode
# Also return time at which each node in path will be reached
function shortestPath(net::Network, travelModeIndex::Int, startFNode::Int, endFNode::Int, startTime::Float)

	# shorthand:
	fNetTravel = net.fNetTravels[travelModeIndex]
	rNetTravel = net.rNetTravels[travelModeIndex]
	@unpack fNodeToRNodeTime, fNodeFromRNodeTime, rArcFNodesTimes = fNetTravel
	@unpack fGraph, rGraph, rArcFNodes, rNodeFNode, fNodeRArcs, fNodeToRNodeNextFNode = net
	rArcs = rGraph.arcs

	# return values:
	fNodeList = Vector{Int}() # list of indices of nodes to travel through, in traversal order
	fNodeTimeList = Vector{Float}() # time at which each node in fNodeList will be reached

	if startFNode == endFNode
		fNodeList = [startFNode]
		fNodeTimeList = [startTime]
		return fNodeList, fNodeTimeList
	end

	# find which rNodes are used in the shortest path
	(travelTime, rNodes) = shortestPathData(net, travelModeIndex, startFNode, endFNode)
	startRNode = rNodes[1]
	endRNode = rNodes[2]

	if startRNode == nullIndex
		@assert(endRNode == nullIndex)
		# no nodes from rGraph are used, travel along a single rArc from startFNode to endFNode
		rArcIndices = fNodeRArcs[startFNode]
		@assert(length(rArcIndices) <= 2)
		toRNode = rArcs[rArcIndices[1]].toNodeIndex
		if fNodeToRNodeTime[startFNode][toRNode] - fNodeToRNodeTime[endFNode][toRNode] < 0 # negative travel time
			toRNode = rArcs[rArcIndices[2]].toNodeIndex
		end
		# create fNodeList, fNodeTimeList
		fNodeList = [startFNode]
		fNodeTimeList = [startTime]
		fNode = startFNode
		currentTime = startTime
		while fNode != endFNode
			# find next arc, add next node and arc time
			nextFNode = fNodeToRNodeNextFNode[fNode][toRNode]
			arcTime = fNetTravel.arcTimes[fGraph.nodePairArcIndex[fNode, nextFNode]]
			currentTime += arcTime
			push!(fNodeList, nextFNode)
			push!(fNodeTimeList, currentTime)
			fNode = nextFNode # go to next node
		end
		return fNodeList, fNodeTimeList
	end

	# travel using rGraph
	if startRNode != nullIndex
		@assert(endRNode != nullIndex)
		currentTime = startTime

		# find path from startFNode to startRNode
		fNodeList = [startFNode]
		fNodeTimeList = [startTime]
		fNode = startFNode
		startRNodeTime = startTime + fNodeToRNodeTime[startFNode][startRNode] # time to reach first rNode
		while fNode != rNodeFNode[startRNode]
			# find next arc, add next node and arc time
			nextFNode = fNodeToRNodeNextFNode[fNode][startRNode]
			# arcTime = fNodeToRNodeTime[fNode][startRNode] - fNodeToRNodeTime[nextFNode][startRNode]
			# arcTime = fNetTravel.arcTimes[fGraph.nodePairArcIndex[fNode, nextFNode]]
			# currentTime += arcTime
			currentTime = startRNodeTime - fNodeToRNodeTime[nextFNode][startRNode]
			push!(fNodeList, nextFNode)
			push!(fNodeTimeList, currentTime)
			fNode = nextFNode # go to next node
		end

		# find path from after startRNode to endRNode (rNodes[1] to rNodes[2])
		rNode = startRNode
		while rNode != endRNode
			# find next rArc, add fNodes
			rArc = shortestPathNextRArc(net, travelModeIndex, rNode, endRNode)
			nextRNode = rArcs[rArc].toNodeIndex
			fNodesTimesForRArc = rArcFNodesTimes[rArc]
			rNodeTime = currentTime
			fNodesOnRArc = rArcFNodes[rArc]
			i = 2
			while i <= length(fNodesOnRArc)
				fNode = fNodesOnRArc[i]
				currentTime = rNodeTime + fNodesTimesForRArc[i]
				push!(fNodeList, fNode)
				push!(fNodeTimeList, currentTime)
				i += 1 # go to next node
			end
			rNode = nextRNode # go to next node
		end

		# find path from after endRNode to endFNode
		fNode = rNodeFNode[endRNode]
		if fNode != endFNode
			rArcIndices = fNodeRArcs[endFNode]
			@assert(length(rArcIndices) <= 2) # endFNode should only be on at most two rArcs, as it is not in rGraph
			rArc = nullIndex
			for rArcIndex in rArcIndices
				if rArcFNodes[rArcIndex][1] == fNode
					rArc = rArcIndex
				end
			end
			fNodesOnRArc = rArcFNodes[rArc]
			fNodeTime = currentTime
			i = 1
			while fNodesOnRArc[i] != endFNode
				i += 1 # go to next node on rArc
				fNode = fNodesOnRArc[i]
				currentTime = fNodeTime + fNodeFromRNodeTime[fNode][endRNode]
				push!(fNodeList, fNode)
				push!(fNodeTimeList, currentTime)
			end
		end
	end

	return fNodeList, fNodeTimeList
end

# set the common fNodes of the network,
# and the shortest path travel data between each fNode and commonFNodes
function setCommonFNodes!(net::Network, commonFNodes::Vector{Int})
	numFNodes = length(net.fGraph.nodes) # shorthand

	# add common fNodes to net
	commonFNodes = sort(unique(commonFNodes))
	@assert(all(commonFNode -> 1 <= commonFNode <= numFNodes, commonFNodes))
	numCommonFNodes = length(commonFNodes) # shorthand
	net.isFNodeCommon = [false for i = 1:numFNodes]
	fNodeCommonFNodeIndex = [nullIndex for i = 1:numFNodes]
	for (i, commonFNode) in enumerate(commonFNodes)
		fNodeCommonFNodeIndex[commonFNode] = i
	end

	# calculate and store shortest path travel data between all fNodes and commonFNodes
	for fNetTravel in net.fNetTravels
		fNetTravel.commonFNodeToFNodeTime = Array{Float,2}(undef, numCommonFNodes, numFNodes)
		fNetTravel.commonFNodeToFNodeDist = Array{Float,2}(undef, numCommonFNodes, numFNodes)
		fNetTravel.commonFNodeToFNodeRNodes = Array{Tuple{Int,Int},2}(undef, numCommonFNodes, numFNodes)

		fNetTravel.fNodeToCommonFNodeTime = Array{Float,2}(undef, numFNodes, numCommonFNodes)
		fNetTravel.fNodeToCommonFNodeDist = Array{Float,2}(undef, numFNodes, numCommonFNodes)
		fNetTravel.fNodeToCommonFNodeRNodes = Array{Tuple{Int,Int},2}(undef, numFNodes, numCommonFNodes)

		for commonFNode in commonFNodes, fNode = 1:numFNodes
			i = fNodeCommonFNodeIndex[commonFNode]

			(travelTime, rNodes) = shortestPathData(net, fNetTravel.modeIndex, commonFNode, fNode)
			fNetTravel.commonFNodeToFNodeTime[i,fNode] = travelTime
			fNetTravel.commonFNodeToFNodeRNodes[i,fNode] = rNodes
			fNetTravel.commonFNodeToFNodeDist[i,fNode] = shortestPathDistance(net, fNetTravel.modeIndex, commonFNode, fNode; startRNode = rNodes[1], endRNode = rNodes[2])

			(travelTime, rNodes) = shortestPathData(net, fNetTravel.modeIndex, fNode, commonFNode)
			fNetTravel.fNodeToCommonFNodeTime[fNode,i] = travelTime
			fNetTravel.fNodeToCommonFNodeRNodes[fNode,i] = rNodes
			fNetTravel.fNodeToCommonFNodeDist[fNode,i] = shortestPathDistance(net, fNetTravel.modeIndex, fNode, commonFNode; startRNode = rNodes[1], endRNode = rNodes[2])
		end
	end

	# add common fNodes to net
	# need to do this after storing results from shortest path calculations above,
	# otherwise the shortestPathData function may try to use values before they are initialised.
	net.commonFNodes = commonFNodes
	net.fNodeCommonFNodeIndex = fNodeCommonFNodeIndex
	for commonFNode in commonFNodes
		net.isFNodeCommon[commonFNode] = true
	end
end

# Return index of the unique rArc that two fNodes have in common,
# where the rArc contains fromFNode before toFNode
# return nullIndex if no such arc found
function findRArcFromFNodeToFNode(net::Network, fromFNode::Int, toFNode::Int)
	@assert(fromFNode != toFNode)
	@assert(isFNodeInRGraph(net, fromFNode) + isFNodeInRGraph(net, toFNode) < 2) # otherwise result may not be unique, as multiple rArcs can connect two rNodes

	rArc = nullIndex
	numRArcsFound = 0
	for rArcIndex in net.fNodeRArcs[fromFNode], rArcIndex2 in net.fNodeRArcs[toFNode]
		if rArcIndex == rArcIndex2 && net.rArcFNodeIndex[rArcIndex][fromFNode] < net.rArcFNodeIndex[rArcIndex][toFNode]
			rArc = rArcIndex
			numRArcsFound += 1
		end
	end
	@assert(numRArcsFound <= 1) # result should be unique

	return rArc
end


module MergingModifiers

println("\nDemonstrating merging modifiers")

using Modia

# Desired:
#   using ModiaMath: plot
#
# In order that these packages need not to be defined in the user environment, they are included via Modia:
using Modia.ModiaMath: plot

@model M begin
  x = Float(start=1.0, info="state")
@equations begin
    der(x) + 2x = 0
  end
end

result = simulate(M, 1, logTranslation=true, removeSingularities=false)
plot(result, ("x"), heading="M", figure=1)

# ------------------------------

@model MInstance begin
  @extends M(x=Var(start=2.0))  # info forgotten
end

result = simulate(MInstance, 1, logTranslation=true, removeSingularities=false)
plot(result, ("x"), heading="MInstance", figure=2)

# ------------------------------

@model MInstance2 begin
  @extends M(y=Var(start=2.0)) # Should give error message for missing y
end

result = simulate(MInstance2, 1, logTranslation=true, removeSingularities=false)
plot(result, ("x"), heading="MInstance2", figure=3)

# ------------------------------

#=
merge(c, e, v) = begin @show c; dump(c); @show e; dump(e.initializers[1].name); dump(e.initializers[1].fdef.args[2].args[2].args[2]); @show v return c end

@model MInstance3 begin
# @extends M(x(start=2.0))  # UndefVarError: x not defined
# @extends M(x=Var(start=2.0))
  @extends M(x=merge(Var(start=2.0), M, :x))
end

result = simulate(MInstance3, 1, logTranslation=true, removeSingularities=false)
plot(result, ("x"), heading="MInstance3", figure=4)
=#

end


#####
##### PushVector
#####

mutable struct PushVector{T} <: AbstractVector{T}
    data::Vector{T}
    len::Int
end

PushVector{T}() where {T} = PushVector{T}(T[], 0)
Base.size(v::PushVector) = (v.len,)
maxlength(v::PushVector) = length(v.data)

Base.IndexStyle(::PushVector) = Base.IndexLinear()
Base.@propagate_inbounds function Base.getindex(A::PushVector, i::Int)
    @boundscheck checkbounds(A, i)
    return @inbounds(A.data[i])
end

Base.@propagate_inbounds function Base.setindex!(A::PushVector, v, i::Int)
    @boundscheck checkbounds(A, i)
    return @inbounds(A.data[i] = v)
end

function Base.push!(A::PushVector, v)
    len = length(A)
    if len < maxlength(A)
        @inbounds(A.data[len + 1] = v)
    else
        push!(A.data, v)
    end
    A.len += 1
    return A
end

Base.empty!(A::PushVector) = (A.len = 0)

#####
##### Node
#####

@enum Direction::UInt8 Left Right
Base.reverse(x::Direction) = Base.bitcast(Direction, xor(UInt8(x), one(UInt8)))

const SEED = UInt(0x167ef00a9)
mutable struct Node{T}
    key::T
    priority::UInt
    nchildren::Int
    left::Union{Nothing,Node{T}}
    right::Union{Nothing,Node{T}}

    function Node{T}(key::T) where {T}
        priority = hash(key, SEED)
        return new{T}(key, priority, 1, nothing, nothing)
    end
end

Node(key::T) where {T} = Node{T}(key)

function Base.show(io::IO, node::Node{T}) where {T}
    hasleft = lchild(node) !== nothing
    hasright = rchild(node) !== nothing
    print(
        io,
        "Node{$T}($(node.key), $(node.priority), $(node.nchildren), $hasleft, $hasright)",
    )
    return nothing
end

mutable struct Treap{T}
    root::Union{Nothing,Node{T}}
    path::PushVector{Tuple{Node{T},Direction}}
    nodebuffer::Vector{Node{T}}
end
Treap{T}() where {T} = Treap{T}(nothing, PushVector{Tuple{Node{T},Direction}}(), Node{T}[])

lchild(n::Node) = n.left
rchild(n::Node) = n.right
hasleft(n::Node) = lchild(n) !== nothing
hasright(n::Node) = rchild(n) !== nothing

@inline setleft!(n::Node, v::Node) = (n.left = v)
@inline setleft!(n::Node) = (n.left = nothing)
@inline setright!(n::Node, v::Node) = (n.right = v)
@inline setright!(n::Node) = (n.right = nothing)

function reassign!(node::Node{T}, key::T) where {T}
    node.key = key
    node.priority = hash(key, SEED)
    node.nchildren = 1
    node.left = nothing
    node.right = nothing
    return nothing
end

# Compute number of children.
function update!(node::Node)
    subtree = 1
    left = lchild(node)
    left === nothing || (subtree += left.nchildren)
    right = rchild(node)
    right === nothing || (subtree += right.nchildren)
    node.nchildren = subtree
    return subtree
end

function rotate_right!(root::Node)
    newroot = lchild(root)
    newroot === nothing && return root
    right = rchild(newroot)
    right === nothing ? setleft!(root) : setleft!(root, right)
    setright!(newroot, root)
    update!(root)
    update!(newroot)
    return newroot
end

function rotate_left!(root::Node)
    newroot = rchild(root)
    newroot === nothing && return root
    left = lchild(newroot)
    left === nothing ? setright!(root) : setright!(root, left)
    setleft!(newroot, root)
    update!(root)
    update!(newroot)
    return newroot
end

function search(root::Node, key)
    root.key == key && return root
    next = (key < root.key) ? lchild(root) : rchild(root)
    return next === nothing ? nothing : search(next, key)
end
Base.in(key, root::Node) = (search(root, key) !== nothing)
Base.haskey(root::Node, key) = (search(root, key) !== nothing)

function nchildren(root::Node, key)
    count = 0
    while true
        if root.key == key
            right = rchild(root)
            if right !== nothing
                count += right.nchildren
            end
            return count
        elseif key < root.key
            count += 1
            right = rchild(root)
            if right !== nothing
                count += right.nchildren
            end
            next = lchild(root)
        else
            next = rchild(root)
        end
        next === nothing && error()
        root = next
    end
end

#####
##### insertion
#####

# Node only method
function Base.insert!(root::Node, node)
    if node.key < root.key
        _left = lchild(root)
        left = setleft!(root, _left === nothing ? node : insert!(_left, node))
        if left.priority > root.priority
            root = rotate_right!(root)
        end
    else
        _right = rchild(root)
        right = setright!(root, _right === nothing ? node : insert!(_right, node))
        if right.priority > root.priority
            root = rotate_left!(root)
        end
    end
    update!(root)
    return root
end

# Using treap auxiliary data structures
function insertdown!(treap::Treap, node)
    _current = root(treap)
    _current === nothing && return nothing
    current = _current
    path = treap.path
    empty!(path)
    while true
        if node.key < current.key
            push!(path, (current, Left))
            next = lchild(current)
        else
            push!(path, (current, Right))
            next = rchild(current)
        end
        next === nothing && return nothing
        current = next
    end
end

function insertup!(treap::Treap, node::Node)
    path = treap.path
    i = length(path)
    @inbounds while i > 0
        root, direction = path[i]
        if direction == Left
            left = setleft!(root, node)
            left.priority > root.priority ? rotate_right!(root) : break
        else
            right = setright!(root, node)
            right.priority > root.priority ? rotate_left!(root) : break
        end
        update!(node)
        i -= 1
    end

    # If the path is empty, then we percolated updates all the way to the root and
    # need to update the "treap's" root.
    # Otherwise, we need to bookkeep the children numbers on the path we visited.
    if iszero(i)
        treap.root = node
    else
        @inbounds for j in Base.OneTo(i)
            root, _ = path[j]
            root.nchildren += 1
        end
    end
    return treap
end

#####
##### deletion
#####

macro _split(sym, f)
    return esc(quote
        if $sym !== nothing
            _tmp = delete!($sym, key)
            _tmp === nothing ? $f(root) : $f(root, _tmp)
        end
    end)
end

function Base.delete!(root::Node{T}, key::T) where {T}
    left = lchild(root)
    right = rchild(root)
    # Branch left if possible
    if (key < root.key)
        @_split left setleft!
        # Branch right if possible
    elseif (key > root.key)
        @_split right setright!
        # Key is at node, try removal of left or right.
    elseif left === nothing
        root = right
    elseif right === nothing
        root = left
        # Node has two children, rotate and move down.
    elseif left.priority < right.priority
        old = root
        root = rotate_left!(root)
        _tmp = delete!(old, key)
        _tmp === nothing ? setleft!(root) : setleft!(root, _tmp)
    else
        old = root
        root = rotate_right!(root)
        _tmp = delete!(old, key)
        _tmp === nothing ? setright!(root) : setright!(root, _tmp)
    end
    root === nothing || update!(root)
    return root
end

function deletefind!(treap::Treap, current::Node, key)
    path = treap.path
    empty!(path)
    while true
        current.key == key && return current
        if key < current.key
            push!(path, (current, Left))
            next = lchild(current)
        else
            push!(path, (current, Right))
            next = rchild(current)
        end
        next === nothing && return nothing
        current = next
    end
end

function handleroot!(treap::Treap, node)
    path = treap.path
    left = lchild(node)
    right = rchild(node)

    if left === nothing
        if right === nothing
            # Delete node - return true to indicate that upper level function
            # should simply return.
            treap.root = nothing
            push!(treap.nodebuffer, node)
            return true
        else
            newroot = rotate_left!(node)
            push!(path, (newroot, Left))
            treap.root = newroot
        end
    elseif right === nothing
        newroot = rotate_right!(node)
        push!(path, (newroot, Right))
        treap.root = newroot
    elseif left.priority < right.priority
        newroot = rotate_left!(node)
        push!(path, (right, Left))
        treap.root = newroot
    else
        newroot = rotate_right!(node)
        push!(path, (left, Right))
        treap.root = newroot
    end
    return false
end

function deletedown!(treap::Treap, root, node)
    updateroot = (node == root)
    if updateroot && handleroot!(treap, node)
        return nothing
    end

    # Shift the node down until it falls off the bottom of the tree.
    path = treap.path
    while true
        parent, direction = @inbounds(path[end])
        left = lchild(node)
        right = rchild(node)

        # Easy path, both children are nothing, we just need to
        if left === nothing
            if right === nothing
                direction === Left ? setleft!(parent) : setright!(parent)
                push!(treap.nodebuffer, node)
                break
            else
                rotate_left!(node)
                direction === Left ? setleft!(parent, right) : setright!(parent, right)
                push!(path, (right, Left))
            end
        elseif right === nothing
            rotate_right!(node)
            direction === Left ? setleft!(parent, left) : setright!(parent, left)
            push!(path, (left, Right))
        elseif left.priority < right.priority
            # Right will become the successor
            rotate_left!(node)
            direction === Left ? setleft!(parent, right) : setright!(parent, right)
            push!(path, (right, Left))
        else
            rotate_right!(node)
            direction === Left ? setleft!(parent, left) : setright!(parent, left)
            push!(path, (left, Right))
        end
    end

    for (n, _) in path
        n.nchildren -= 1
    end
    return nothing
end

#####
##### misc
#####

function traverse(f::F, root::Node) where {F}
    left = lchild(root)
    if left !== nothing
        traverse(f, left)
    end
    f(root)
    right = rchild(root)
    if right !== nothing
        traverse(f, right)
    end
    return nothing
end

function list(root::Node{T}) where {T}
    v = T[]
    traverse(x -> push!(v, x.key), root)
    return v
end

#####
##### Treap
#####

function Base.length(treap::Treap)
    base = root(treap)
    return base === nothing ? 0 : base.nchildren
end
root(treap::Treap) = treap.root

function nchildren(treap::Treap, key)
    base = root(treap)
    base === nothing && return -1
    return nchildren(base, key)
end

function Base.in(key0, treap::Treap{T}) where {T}
    node = root(treap)
    node === nothing && return false
    return in(convert(T, key0), node)
end

function unsafe_push!(treap::Treap{T}, key0) where {T}
    buffer = treap.nodebuffer
    key = convert(T, key0)
    if !isempty(buffer)
        node = pop!(buffer)
        reassign!(node, key)
    else
        node = Node(key)
    end

    base = root(treap)
    if base === nothing
        treap.root = node
        return treap
    end

    insertdown!(treap, node)
    insertup!(treap, node)
    return treap
end

function Base.push!(treap::Treap{T}, key0) where {T}
    key = convert(T, key0)
    in(key, treap) || unsafe_push!(treap, key)
    return treap
end

function Base.delete!(treap::Treap{T}, key0) where {T}
    key = convert(T, key0)
    # If the root doesn't exist, then obviously we can't delete the key.
    base = root(treap)
    base === nothing && return false

    # If `deletefind!` doesn't find the node - then we've done a bit of extra work by
    # keeping track of our downward path.
    #
    # But, if we DO find the node, then we can save redoing the work we just did, which
    # is potentially much more useful.
    node = deletefind!(treap, base, key)
    node === nothing && return false
    deletedown!(treap, base, node)
    return true
end

#####
##### Treap verification
#####

function istree(treap::Treap)
    base = root(treap)
    return base === nothing ? true : istree(base)
end
function istree(node::Node{T}) where {T}
    v = T[]
    traverse(node) do n
        push!(v, n.key)
    end
    return issorted(v)
end

function isheap(treap::Treap)
    base = root(treap)
    return base === nothing ? true : isheap(base)
end
function isheap(node::Node)
    passed = true
    left = lchild(node)
    if left !== nothing
        passed &= isheap(left)
        passed &= (node.priority > left.priority)
    end
    right = rchild(node)
    if right !== nothing
        passed &= isheap(right)
        passed &= (node.priority > right.priority)
    end

    return passed
end

#####
##### AbstractTrees
#####

function AbstractTrees.children(node::Node)
    if hasleft(node)
        if hasright(node)
            return (lchild(node), rchild(node))
        else
            return (lchild(node),)
        end
    else
        if hasright(node)
            return (rchild(node),)
        else
            return ()
        end
    end
end

AbstractTrees.printnode(io::IO, node::Node) = print(io, "($(node.key), $(node.priority))")


# module for representing robot model
# modeling as differential wheeled robot

using Plots, Random, Distributions
pyplot()

include(joinpath(split(@__FILE__, "src")[1], "src/slam/graph_based_slam/logger_agent.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/slam/graph_based_slam/psi_sensor.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/slam/fast_slam_2/fast_slam_2_agent.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/model/agent/agent.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/model/agent/puddle_ignore_agent.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/model/agent/dp_policy_agent.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/model/agent/q_agent.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/model/sensor/sensor.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/common/error_calculation/error_calculation.jl"))
include(joinpath(split(@__FILE__, "src")[1], "src/common/state_transition/state_transition.jl"))

mutable struct DifferentialWheeledRobot
  pose
  radius
  color
  agent
  delta_time
  noise_expon
  noise_norm
  dist_until_noise
  bias_rate_spd
  bias_rate_yr
  stuck_exporn
  escape_exporn
  time_until_stuck
  time_until_escape
  kidnap_exporn
  kidnap_x_uniform
  kidnap_y_uniform
  kidnap_theta_uniform
  time_until_kidnap
  sensor
  is_stuck
  traj_x
  traj_y

  # init
  function DifferentialWheeledRobot(pose::Array, radius::Float64,
                                    color::String, agent,
                                    delta_time::Float64; sensor=nothing,
                                    noise_per_meter::Int64=0,
                                    noise_std::Float64=0.0,
                                    bias_rate_stds::Array=[0.0, 0.0],
                                    exp_stuck_time::Float64=1e100,
                                    exp_escape_time::Float64=1e-100,
                                    exp_kidnap_time::Float64=1e100,
                                    kidnap_rx::Array=[-5.0, 5.0],
                                    kidnap_ry::Array=[-5.0, 5.0])
      self = new()
      self.pose = pose
      self.radius = radius
      self.color = color
      self.agent = agent
      self.delta_time = delta_time
      self.noise_expon = Exponential(1.0/(1e-100 + noise_per_meter))
      self.noise_norm = Normal(0.0, noise_std)
      self.dist_until_noise = rand(self.noise_expon)
      self.bias_rate_spd = rand(Normal(1.0, bias_rate_stds[1]))
      self.bias_rate_yr = rand(Normal(1.0, bias_rate_stds[2]))
      self.stuck_exporn = Exponential(exp_stuck_time)
      self.escape_exporn = Exponential(exp_escape_time)
      self.time_until_stuck = rand(self.stuck_exporn)
      self.time_until_escape = rand(self.escape_exporn)
      self.kidnap_exporn = Exponential(exp_kidnap_time)
      self.kidnap_x_uniform = Uniform(kidnap_rx[1], kidnap_rx[2])
      self.kidnap_y_uniform = Uniform(kidnap_ry[1], kidnap_ry[2])
      self.kidnap_theta_uniform = Uniform(0.0, 2 * pi)
      self.time_until_kidnap = rand(self.kidnap_exporn)
      self.sensor = sensor
      self.is_stuck = false
      self.traj_x = [pose[1]]
      self.traj_y = [pose[2]]
      return self
  end
end

function noise(self::DifferentialWheeledRobot, pose::Array,
               speed::Float64, yaw_rate::Float64, delta_time::Float64)
  self.dist_until_noise -= abs(speed) * delta_time + self.radius * abs(yaw_rate) * delta_time
  if self.dist_until_noise <= 0.0
    self.dist_until_noise += rand(self.noise_expon)
    pose[3] += rand(self.noise_norm)
  end
  return pose
end

function bias(self::DifferentialWheeledRobot, speed::Float64,
              yaw_rate::Float64)
  return speed * self.bias_rate_spd, yaw_rate * self.bias_rate_yr
end

function stuck(self::DifferentialWheeledRobot, speed::Float64,
             yaw_rate::Float64, delta_time::Float64)
  if self.is_stuck
    self.time_until_escape -= delta_time
    if self.time_until_escape <= 0.0
      self.time_until_escape += rand(self.escape_exporn)
      self.is_stuck = false
    end
  else
    self.time_until_stuck -= delta_time
    if self.time_until_stuck <= 0.0
      self.time_until_stuck += rand(self.stuck_exporn)
      self.is_stuck = true
    end
  end
  return speed * (!self.is_stuck), yaw_rate * (!self.is_stuck)
end

function kidnap(self::DifferentialWheeledRobot, pose::Array,
                delta_time::Float64)
  self.time_until_kidnap -= delta_time
  if self.time_until_kidnap <= 0.0
    self.time_until_kidnap += rand(self.kidnap_exporn)
    return [rand(self.kidnap_x_uniform), rand(self.kidnap_y_uniform),
            rand(self.kidnap_theta_uniform)]
  else
    return pose
  end
end

function circle(x, y, r)
  theta = LinRange(0, 2*pi, 500)
  x .+ r * sin.(theta), y .+ r * cos.(theta)
end

function draw!(self::DifferentialWheeledRobot)
  # current pose
  x, y, theta = self.pose[1], self.pose[2], self.pose[3]
  xn = x + self.radius * cos(theta)
  yn = y + self.radius * sin(theta)

  # draw robot and current pose
  plot!([x, xn], [y, yn], color=self.color,
        legend=false, aspect_ratio=true)
  plot!(circle(x, y, self.radius), linecolor=self.color,
        aspect_ratio=true)

  # draw trajectory
  plot!(self.traj_x, self.traj_y, color="black",
        legend=false, aspect_ratio=true)

  # draw observation
  observation = []
  if self.sensor !== nothing
    data(self.sensor, self.pose)
    draw!(self.sensor, self.pose)
    observation = self.sensor.last_data
  end

  # draw estimation
  spd, yr = draw_decision!(self.agent, observation)

  if self.agent.estimator !== nothing
    draw_lon_lat_error!(self.agent.estimator.estimated_pose, self.pose)
  end

  # next pose
  spd, yr = bias(self, spd, yr)
  spd, yr = stuck(self, spd, yr, self.delta_time)
  self.pose = state_transition(spd, yr, self.delta_time, self.pose)
  self.pose = noise(self, self.pose, spd, yr, self.delta_time)
  self.pose = kidnap(self, self.pose, self.delta_time)
  push!(self.traj_x, self.pose[1]), push!(self.traj_y, self.pose[2])
end


export add_walsh, add_residue

function add_walsh(mgr::Ptr{Manager}, f_pp::Ptr{Ptr{Node}}, g_pp::Ptr{Ptr{Node}}, n::Int)
    result = Cudd_addWalsh(mgr, f_pp, g_pp, n)
    if result == C_NULL
        throw(OutOfMemoryError())
    end
    return result
end

function add_residue(mgr::Ptr{Manager}, n::Int, m::Int, options::Int, top::Int)
    result = Cudd_addResidue(mgr, n, m, options, top)
    if result == C_NULL
        throw(OutOfMemoryError())
    end
    return result
end


export NoGate

using YaoBase, YaoArrayRegister, YaoBlocks
using LuxurySparse

struct NoGate <: PrimitiveBlock{1} end

YaoBlocks.apply!(r::AbstractRegister, x::NoGate) = r
YaoBase.instruct!(r::ArrayReg, x::NoGate, locs) = r
YaoBlocks.mat_matchreg(r::ArrayReg, x::NoGate) = x

YaoBlocks.mat(::Type{T}, x::NoGate) where T = IMatrix{2, T}()


using Test, Omniscape

d = run_omniscape("input/config2.ini")
a, b, c = run_omniscape("input/config.ini")

e = run_omniscape("input/config3.ini")

@test typeof(a) == Array{Float64,2}
@test typeof(b) == Array{Float64,2}
@test typeof(c) == Array{Float64,2}
@test typeof(d) == Array{Float64,2}
@test typeof(e) == Array{Float64,2}
@test b ≈ d

rm("test1_output", recursive = true)
rm("test2_output", recursive = true)
rm("test3_output", recursive = true)


using DSGE, ModelConstructors, Distributed
using DataFrames, OrderedCollections, Dates, Plots
GR.inline("pdf")

####################################
# To use:
# Just run in the Julia REPL
# include("decompose_forecast.jl")
####################################

use_parallel = false # set to true if you want to use parallel workers
n_workers    = 10
make_plots   = true  # set to true if you want to produce forecast decomposition plots

############
# Estimate
############

# Initialize a new instance of an An Schorfheide model,
# a canonical New Keynesian DSGE model augmented for
# Bayesian estimation.
m = AnSchorfheide()

# See "test_smc.jl" for more information on setting up SMC
# We choose settings so the estimation finishes faster,
# so the particles will not approximate the actual posterior well.
m <= Setting(:data_vintage, "200101", true, "vint", "Data vintage")
m <= Setting(:date_forecast_start, DSGE.quartertodate("2020-Q1"))
m <= Setting(:test_decomp_fcast, "true", true, "test_decomp_fcast",
             "Tag for estimation output file name indicating this file is for testing a forecast decomposition")
m <= Setting(:use_parallel_workers, use_parallel)
m <= Setting(:n_particles, 800)
m <= Setting(:n_Φ, 50) # number of stages in SMC
m <= Setting(:λ, 2.0) # the curvature of the tempering schedule
m <= Setting(:adaptive_tempering_target_smc, false) # Setting a fixed tempering schedule
m <= Setting(:reoptimize, false)
m <= Setting(:sampling_method, :SMC)
forecast_string = "new" # string to distinguish the forecast for the new model

# Load data and create fake old data
df = load_data(m)
df_old = deepcopy(df[1:end - 4, :]) # subset data into a pseudo "old" and "new" data so old data is 1 year old

# Set up our "fake" old model
m_old = deepcopy(m)
m_old <= Setting(:data_vintage, "190101", true, "vint", "Data vintage") # fake data vintage, 1 year old
m_old <= Setting(:date_forecast_start, DSGE.quartertodate("2019-Q1"))
forecast_string_old = "old" # forecast string for the old forecast

# Estimate
if use_parallel
    myprocs = addprocs(n_workers)
    @everywhere using DSGE, OrderedCollections
end

estimate(m, df_to_matrix(m, df), verbose = :low, run_csminwel = false)
estimate(m_old, df_to_matrix(m_old, df_old), verbose = :low, run_csminwel = false)

if use_parallel
    rmprocs(myprocs)
end

#####################
# Decompose forecast
#####################
# We now produce a forecast decomposition, which decomposes the difference
# in the forecasts of the old model/data and the new model/data into the effect from
# data revisions, news (e.g. new data), and parameter re-estimation.
# To save on time, we only compute the modal forecast decomposition,
# and we will not use the DSGE's indivdual shocks for the decomposition.

# Run forecast
DSGE.decompose_forecast(m, m_old, df, df_old, :mode, :none, :none, [:obs, :pseudo],
                        check = true, forecast_string_new = forecast_string,
                        forecast_string_old = forecast_string_old)
DSGE.decomposition_means(m, m_old, :mode, :none, :none, [:obs, :pseudo],
                         verbose = :none, forecast_string_new = forecast_string,
                         forecast_string_old = forecast_string_old)

if make_plots
    # Set up plots
    kwargs = Dict{Symbol, Any}()
    kwargs[:tick_size]      = 1
    kwargs[:hist_label]     = ""
    kwargs[:forecast_label] = ""
    kwargs[:start_date]     = DSGE.iterate_quarters(date_mainsample_end(m), -4)
    kwargs[:end_date]       = DSGE.iterate_quarters(date_mainsample_end(m), 8)  # 2 years out

    # Make plots
    DSGE.plot_forecast_decomposition(m, m_old, collect(keys(m.observables)),
                                     :obs, :mode, :none, :none;
                                     individual_shocks = false,
                                     forecast_string_new = forecast_string,
                                     forecast_string_old = forecast_string_old,
                                     kwargs...)
    DSGE.plot_forecast_decomposition(m, m_old, collect(keys(m.pseudo_observables)),
                                     :pseudo, :mode, :none, :none;
                                     individual_shocks = false,
                                     forecast_string_new = forecast_string,
                                     forecast_string_old = forecast_string_old,
                                     kwargs...)
end

nothing


using CSV
using DataFrames
using DelimitedFiles
using YAML
using Query

using SLiDE

UNITS = "billions of us dollars (USD)"

# ******************************************************************************************
#   READ BLUENOTE DATA -- For benchmarking!
# ******************************************************************************************
BLUE2 = joinpath("data", "windc_output", "2z_windcdatabase")
b2_lst = [x for x in readdir(joinpath(SLIDE_DIR, BLUE2)) if occursin(".csv", x)]
# b2 = Dict(Symbol(k[1:end-4]) => sort(edit_with(
#     read_file(joinpath(BLUE2, k))), Rename(:Val, :value))) for k in b2_lst)

BLUE_DIR = joinpath("data", "windc_output", "2a_io_national_cgeparm_raw")
bluenote_lst = [x for x in readdir(joinpath(SLIDE_DIR, BLUE_DIR)) if occursin(".csv", x)]
bluenote = Dict(Symbol(k[1:end-4]) => sort(edit_with(
    read_file(joinpath(BLUE_DIR, k)), Rename(:Val, :value))) for k in bluenote_lst)

# Add supply/use info for checking.
BLUE_DIR_IN = joinpath("data", "windc_output", "1b_stream_windc_base")
[bluenote[k] = sort(edit_with(read_file(joinpath(BLUE_DIR_IN, string(k, "_units.csv"))), [
    Rename.([:Dim1,:Dim2,:Dim3,:Dim4,:Val], [:yr,:i,:j,:units,:value]);
    Replace.([:i,:j], "upper", "lower")])) for k in [:supply, :use]]

[bluenote[k][!,:value] .= round.(bluenote[k][:,:value]*1E-3, digits=3) # convert millions -> billions USD
    for k in [:supply,:use]]

# ******************************************************************************************
#   READ SETS AND SLiDE SUPPLY/USE DATA.
# ******************************************************************************************
y = read_file(joinpath("data","readfiles","list_sets.yml"));
set = Dict(Symbol(k) => sort(read_file(joinpath(y["Path"]..., ensurearray(v)...)))[:,1]
    for (k,v) in y["Input"])

# Read supply/use data.
DATA_DIR = joinpath("data", "input")
io_lst = convert_type.(Symbol, ["supply", "use"])
jo = Dict(k => read_file(joinpath(DATA_DIR, string(k, ".csv"))) for k in io_lst)

# Perform some minor edits that will apply to all parameters in order to maintain
# consistency with the data in build_national_cgeparm_raw.gdx for easier benchmarking.

for k in keys(jo)
    global jo[k] = edit_with(jo[k], [Rename(:value, k); Drop.([:value,:units], [0.,"all"], "==")])
    global jo[k] = jo[k] |> @filter(_.yr in set[:yr]) |> DataFrame

    # global jo[k][!,:value] .= round.(jo[k][:,:value]*1E-3, digits=3)
    # global jo[k][!,:units] .= UNITS
end

jo[:supply], jo[:use] = fill_zero(jo[:supply], jo[:use]; permute_keys = true);

df = jo[:supply][:,find_oftype(jo[:supply], Not(AbstractFloat))]
df = hcat(df, [edit_with(jo[k],Rename(:value,k))[:,[k]] for k in sort(collect(keys(jo)))]...)
df = fill_zero((yr = set[:yr], i = set[:i], j = set[:j]), df)

# df = copy(df)
# x = (i = set[:i], j = set[:j])
# col = [:a, :b]

function add_filter(df::DataFrame, col::Array{Symbol,1}, set)
    df = copy(df)
    cols = intersect(find_oftype(df, Not(AbstractFloat)), find_oftype(df, Not(Bool)))
    inset = [col in collect(keys(set)) for col in cols]

    vals_sets = [[set[k] for k in cols[inset]]; unique.(eachcol(df[:,cols[.!inset]]))]
    cols_sets = [cols[inset]; cols[.!inset]]

    df_sets = edit_with(DataFrame(permute(NamedTuple{Tuple(cols_sets,)}(vals_sets,))),
        Add.(col, true))
    df = leftjoin(df, df_sets, on = cols_sets)

    df = edit_with(df, Replace.(col, missing, false))
    return df
end

s = Dict()

mutable struct Filter
    set::NamedTuple
    col::Array{Symbol,1}
end

filters = []

push!(filters, Filter((yr = set[:yr], i = set[:i], j = set[:j]),     [:ys0, :id0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = set[:fd]),    [:fd0]))
push!(filters, Filter((yr = set[:yr], i = set[:va], j = set[:j]),    [:va0]))
push!(filters, Filter((yr = set[:yr], i = set[:ts], j = set[:j]),    [:ts0]))

push!(filters, Filter((yr = set[:yr], i = set[:i], j = "exports"),   [:x0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "imports"),   [:m0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "margins"),   [:mrg0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "trncost"),   [:trn0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "ciffob"),    [:cif0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "duties"),    [:duty0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "tax"),       [:tax0]))
push!(filters, Filter((yr = set[:yr], i = set[:i], j = "subsidies"), [:sbd0]))

# Create DataFrame from filters.
df_filtered = copy(df[:,[:yr,:i,:j]])
[global df_filtered = add_filter(df_filtered, f.col, f.set) for f in filters]

# Create final DataFrames
ij = fill_zero((yr = set[:yr], i = set[:i], j = set[:j]))[:,1:end-1]
i  = fill_zero((yr = set[:yr], i = set[:i]))[:,1:end-1]
fd = fill_zero((yr = set[:yr], i = set[:i], fd = set[:fd]))[:,1:end-1]
ts = fill_zero((yr = set[:yr], ts = set[:ts], j = set[:j]))[:,1:end-1]
va = fill_zero((yr = set[:yr], va = set[:va], j = set[:j]))[:,1:end-1]

kij = [:ys0,:id0]
vij = [:supply,:use]

ki = [:x0,:m0,:mrg0,:trn0,:cif0,:duty0,:tax0,:sbd0]
vi = [:use; fill(:supply,7)]

[global i[!,k] .= df[df_filtered[:,k],v] for (k,v) in zip(ki,vi)]
[global ij[!,k] .= df[df_filtered[:,k],v] for (k,v) in zip(kij,vij)]
fd[!,:fd0] = df[df_filtered[:,:fd0],:use]
ts[!,:ts0] = df[df_filtered[:,:ts0],:use]
va[!,:va0] = df[df_filtered[:,:va0],:use]

# Treat negative inputs as outputs.
ij[!,:ys0] = ij[:,:ys0] - min.(0, ij[:,:id0])
ij[!,:id0] = max.(0, ij[:,:id0])


export CDS_MVp_MT_subfunc
"""
wrapper around Julia threads to compute a multi-threaded matrix vector product in CDS formulated
This is a subfunction of CDS_MVp_MT.jl
"""
function CDS_MVp_MT_subfunc(
        R  ::Array{TF,2},
        x  ::Vector{TF},
        y  ::Vector{TF},
        r0 ::Int,
        c0 ::Int,
        r1 ::Int,
        i  ::Int) where {TF<:Real}

  @Threads.threads for r = r0 : r1
      c = r - r0 + c0 #original
      @inbounds y[r] = y[r] + R[r,i] * x[c]#original
    end
  return y
end


global drag_get=Vector{Function}(),
drag_set=Vector{Function}(),
drag_type=Vector{Symbol}()
#= c for common
r for read-only
i for infinite =#
function drag_work(from::Int,to::Int;shift::Bool=false,ctrl::Bool=false)::Bool
	if (drag_type[to]!=:c)||(from==to)
		return false
	end
	f1=deepcopy(drag_get[from](from))::Pair
	f2=deepcopy(drag_get[to](to))::Pair
	if isa(f2.first,EI) # =>空格
		if drag_type[from]==:i
			drag_set[to](to,Pair(f1.first,ctrl ? stack(f1.first) : 0x1))
		else
			if ctrl # 分半
				mi=f1.second>>1
				drag_set[from](from,Pair(mi==0 ? EI() : f1.first,mi))
				drag_set[to](to,Pair(deepcopy(f1.first),f1.second-mi))
			elseif shift # 分一
				if f1.second==0x1
					return false
				end
				drag_set[from](from,Pair(f1.first,0x1))
				drag_set[to](to,Pair(deepcopy(f1.first),f1.second-0x1))
			else # 全移
				drag_set[from](from,Pair(EI(),0x0))
				drag_set[to](to,f1)
			end
		end
	elseif id(f2.first)==id(f2.first) # 同类
		st=stack(f2.first)
		if st==f2.second
			return false
		end
		if drag_type[from]==:i # 填满
			drag_set[to](to,Pair(f2.first,st))
		else
			sum=f1.second+f2.second
			if sum>st
				drag_set[from](from,Pair(f1.first,sum-st))
				drag_set[to](to,Pair(f2.first,st))
			else
				drag_set[from](from,Pair(EI(),0x0))
				drag_set[to](to,Pair(f2.first,sum))
			end
		end
	elseif drag_type[from]==:c
		drag_set[from](from,f2)
		drag_set[to](to,f1)
	else
		return false
	end
	return true
end


################################################################################
# MOI wrapper
################################################################################
# This file is largely inspired from MOI wrappers existing in other
# JuliaOpt packages:
# https://www.juliaopt.org/
# The authors are indebted to the developers of JuliaOpt for
# the current MOI wrapper.
#
#--------------------------------------------------
# Specifications
#--------------------------------------------------
# The MOI wrapper works only for KNITRO version >= 11.0.
#
# - linear coefs are handled directly by KNITRO, via
#  `KN_add_con_linear_struct , KN_add_obj_linear_struct
#
# - quadratic coefs are also handled directly by KNITRO, via
#  `KN_add_con_quad_struct , KN_add_obj_quad_struct
#
# - NLP data are handled via KNITRO's callbacks
#
# NB: If `model.nlp_data` is empty, KNITRO would not use any
#     callbacks during solve.
#
#--------------------------------------------------

import MathOptInterface
const MOI  = MathOptInterface
const MOIU = MathOptInterface.Utilities

# TODO
const SF = Union{MOI.ScalarAffineFunction{Float64},
                 MOI.ScalarQuadraticFunction{Float64}}
const VAF = MOI.VectorAffineFunction{Float64}
const VOV = MOI.VectorOfVariables

const SS = Union{MOI.EqualTo{Float64},
                 MOI.GreaterThan{Float64},
                 MOI.LessThan{Float64},
                 MOI.Interval{Float64}}
# LinSets
const LS = Union{MOI.EqualTo{Float64},
                 MOI.GreaterThan{Float64},
                 MOI.LessThan{Float64}}
# VecLinSets
const VLS = Union{MOI.Nonnegatives,
                  MOI.Nonpositives,
                  MOI.Zeros}

##################################################
# Define custom error for MOI wrapper.
struct UpdateObjectiveError <: MOI.NotAllowedError end
struct AddVariableError <: MOI.NotAllowedError end
struct AddConstraintError <: MOI.NotAllowedError end

# Import some utils.
include(joinpath("MOI_wrapper", "utils.jl"))

##################################################
mutable struct VariableInfo
    has_lower_bound::Bool # Implies lower_bound == Inf
    has_upper_bound::Bool # Implies upper_bound == Inf
    is_fixed::Bool        # Implies lower_bound == upper_bound and !has_lower_bound and !has_upper_bound.
    name::String
end
VariableInfo() = VariableInfo(false, false, false, "")

##################################################
# EmptyNLPEvaluator for non-NLP problems.
struct EmptyNLPEvaluator <: MOI.AbstractNLPEvaluator end
MOI.features_available(::EmptyNLPEvaluator) = [:Grad, :Jac, :Hess]
MOI.initialize(::EmptyNLPEvaluator, features) = nothing
MOI.eval_objective(::EmptyNLPEvaluator, x) = NaN
function MOI.eval_constraint(::EmptyNLPEvaluator, g, x)
    return
end
MOI.eval_objective_gradient(::EmptyNLPEvaluator, g, x) = nothing
MOI.jacobian_structure(::EmptyNLPEvaluator) = Tuple{Int64,Int64}[]
MOI.hessian_lagrangian_structure(::EmptyNLPEvaluator) = Tuple{Int64,Int64}[]
function MOI.eval_constraint_jacobian(::EmptyNLPEvaluator, J, x)
    return
end
function MOI.eval_hessian_lagrangian(::EmptyNLPEvaluator, H, x, σ, μ)
    return
end

empty_nlp_data() = MOI.NLPBlockData([], EmptyNLPEvaluator(), false)

##################################################
# MOI Optimizer
mutable struct Optimizer <: MOI.AbstractOptimizer
    inner::Union{Model, Nothing}
    # We only keep in memory some information about variables
    # as we cannot delete variables, we do not have to store an index.
    variable_info::Vector{VariableInfo}
    # Get number of solve for restart.
    number_solved::Int
    # Specify if NLP is loaded inside KNITRO to avoid double definition.
    nlp_loaded::Bool
    nlp_data::MOI.NLPBlockData
    # Store index of nlp constraints.
    nlp_index_cons::Vector{Cint}
    # Store optimization sense.
    sense::MOI.OptimizationSense
    # Store the structure of the objective.
    objective::Union{MOI.SingleVariable,MOI.ScalarAffineFunction{Float64},MOI.ScalarQuadraticFunction{Float64},Nothing}
    # Constraint counters.
    number_zeroone_constraints::Int
    number_integer_constraints::Int
    # Constraint mappings.
    constraint_mapping::Dict{MOI.ConstraintIndex, Union{Cint, Vector{Cint}}}
    license_manager::Union{LMcontext, Nothing}
    options::Dict{String, Any}
end

function set_options(model::Optimizer, options)
    # Set KNITRO option.
    for (name, value) in options
        sname = string(name)
        MOI.set(model, MOI.RawParameter(sname), value)
    end
    return
end

function Optimizer(;license_manager=nothing, options...)
    # Create KNITRO context.
    if isa(license_manager, LMcontext)
        kc = Model(license_manager)
    else
        kc = Model()
    end
    # Convert Symbol to String in options dictionnary.
    options_dict = Dict{String, Any}()
    for (name, value) in options
        options_dict[string(name)] = value
    end
    model = Optimizer(kc, [], 0, false, empty_nlp_data(),
                      Cint[], MOI.FEASIBILITY_SENSE, nothing, 0, 0,
                      Dict{MOI.ConstraintIndex, Int}(), license_manager, Dict())

    set_options(model, options)
    return model
end

# Print Optimizer.
function Base.show(io::IO, model::Optimizer)
    println(io, "A MathOptInterface model with backend:")
    println(io, model.inner)
    return
end

# copy
MOIU.supports_default_copy_to(model::Optimizer, copy_names::Bool) = true
function MOI.copy_to(model::Optimizer, src::MOI.ModelLike; kws...)
    return MOI.Utilities.automatic_copy_to(model, src; kws...)
end

function MOI.empty!(model::Optimizer)
    # Free KNITRO model properly.
    if model.inner != nothing
        KN_free(model.inner)
    end
    # Handle properly license manager
    if isa(model.license_manager, LMcontext)
        model.inner = Model(model.license_manager)
    else
        model.inner = Model()
    end
    empty!(model.variable_info)
    model.number_solved = 0
    model.nlp_data = empty_nlp_data()
    model.nlp_loaded = false
    model.nlp_index_cons = Cint[]
    model.sense = MOI.FEASIBILITY_SENSE
    model.objective = nothing
    model.number_zeroone_constraints = 0
    model.number_integer_constraints = 0
    model.constraint_mapping = Dict()
    model.license_manager = model.license_manager
    set_options(model, model.options)
    return
end

function MOI.is_empty(model::Optimizer)
    return isempty(model.variable_info) &&
           model.nlp_data.evaluator isa EmptyNLPEvaluator &&
           model.sense == MOI.FEASIBILITY_SENSE &&
           model.number_solved == 0 &&
           isa(model.objective, Nothing) &&
           model.number_zeroone_constraints == 0 &&
           model.number_integer_constraints == 0 &&
           !model.nlp_loaded
end

# Some utilities.
number_variables(model::Optimizer) = length(model.variable_info)
number_constraints(model::Optimizer) = KN_get_number_cons(model.inner)

# Getter for solver's name.
MOI.get(model::Optimizer, ::MOI.SolverName) = "Knitro"

# MOI.Silent.
MOI.supports(model::Optimizer, ::MOI.Silent) = true
function MOI.get(model::Optimizer, ::MOI.Silent)
    return KN_get_int_param(model.inner, "outlev") == 0
end

function MOI.set(model::Optimizer, ::MOI.Silent, value)
    outlev = value ? 0 : 2
    KN_set_param(model.inner, "outlev", outlev)
    return
end

# MOI.TimeLimitSec.
MOI.supports(model::Optimizer, ::MOI.TimeLimitSec) = true
function MOI.get(model::Optimizer, ::MOI.TimeLimitSec)
    return KN_get_double_param(model.inner, "maxtime_cpu")
end

function MOI.set(model::Optimizer, ::MOI.TimeLimitSec, value)
    # By default, maxtime is set to 1e8 in Knitro.
    maxtime = isnothing(value) ? 1e8 : value
    KN_set_param(model.inner, "maxtime_cpu", value)
    return
end

# MOI.RawParameters
function MOI.supports(model::Optimizer, param::MOI.RawParameter)
    name = param.name
    if name in KNITRO_OPTIONS || haskey(KN_paramName2Indx, name)
        return true
    end
    return false
end

function MOI.set(model::Optimizer, p::MOI.RawParameter, value)
    if !MOI.supports(model, p)
        throw(MOI.UnsupportedAttribute)
    end
    model.options[p.name] = value
    if p.name == "option_file"
        KN_load_param_file(model.inner, value)
    elseif p.name == "tuner_file"
        KN_load_tuner_file(model.inner, value)
    else
        KN_set_param(model.inner, p.name, value)
    end
    return
end

function MOI.get(model::Optimizer, p::MOI.RawParameter)
    if haskey(model.options, p.name)
        return model.options[p.name]
    end
    error("RawParameter with name $(p.name) is not set.")
end

##################################################
# Optimize
##################################################
function MOI.optimize!(model::Optimizer)
    KN_set_param(model.inner, "datacheck", 0)
    KN_set_param(model.inner, "hessian_no_f", 1)

    # Add NLP structure if specified.
    if !isa(model.nlp_data.evaluator, EmptyNLPEvaluator) && !model.nlp_loaded
        # Instantiate NLPEvaluator once and for all.
        features = MOI.features_available(model.nlp_data.evaluator)
        has_hessian = (:Hess in features)
        # Build initial features for solver.
        init_feat = Symbol[]
        model.nlp_data.has_objective && push!(init_feat, :Grad)
        has_hessian && push!(init_feat, :Hess)
        num_nlp_constraints = length(model.nlp_data.constraint_bounds)
        num_nlp_constraints > 0 && push!(init_feat, :Jac)
        # initialize!
        MOI.initialize(model.nlp_data.evaluator, init_feat)

        # Load NLP structure inside Knitro.
        offset = number_constraints(model)
        load_nlp_constraints(model)
        num_cons = KN_get_number_cons(model.inner)

        # The callbacks must match the signature of the callbacks
        # defined in knitro.h.
        # Objective callback (set both objective and constraint evaluation).
        function eval_f_cb(kc, cb, evalRequest, evalResult, userParams)
            # Evaluate objective if specified in nlp_data.
            if model.nlp_data.has_objective
                evalResult.obj[1] = MOI.eval_objective(model.nlp_data.evaluator,
                                                       evalRequest.x)
            end
            # Evaluate nonlinear term in constraint.
            MOI.eval_constraint(model.nlp_data.evaluator,
                                evalResult.c, evalRequest.x)
            return 0
        end

        # Objective gradient callback.
        function eval_grad_cb(kc, cb, evalRequest, evalResult, userParams)
            # Evaluate non-linear term in objective gradient.
            if model.nlp_data.has_objective
                MOI.eval_objective_gradient(model.nlp_data.evaluator,
                                            evalResult.objGrad,
                                            evalRequest.x)
            end
            # Evaluate non linear part of jacobian.
            MOI.eval_constraint_jacobian(model.nlp_data.evaluator,
                                         evalResult.jac,
                                         evalRequest.x)
            return 0
        end

        if has_hessian
            # Hessian callback.
            function eval_h_cb(kc, cb, evalRequest, evalResult, userParams)
                MOI.eval_hessian_lagrangian(model.nlp_data.evaluator,
                                            evalResult.hess,
                                            evalRequest.x,
                                            evalRequest.sigma,
                                            view(evalRequest.lambda, offset+1:num_cons))
                return 0
            end
        else
            eval_h_cb = nothing
        end

        # Be careful that sometimes objective is not evaluated here.
        # In any case, NLP objective has precedence over model.objective.
        if model.nlp_data.has_objective
            if num_nlp_constraints == 0
                # Add only a callback for objective if no NLP constraint
                cb = KN_add_objective_callback(model.inner, eval_f_cb)
            else
                cb = KN_add_eval_callback(model.inner, true, model.nlp_index_cons, eval_f_cb)
            end
        elseif num_nlp_constraints > 0
            cb = KN_add_eval_callback(model.inner, false, model.nlp_index_cons, eval_f_cb)

            # If a objective is specified in model.objective, load it.
            !isa(model.objective, Nothing) && add_objective!(model, model.objective)
        end

        # Gradient structure.
        nV = (model.nlp_data.has_objective) ? KN_DENSE : 0
        # Get jacobian structure.
        jacob_structure = MOI.jacobian_structure(model.nlp_data.evaluator)
        nnzJ = length(jacob_structure)
        if nnzJ == 0
            KN_set_cb_grad(model.inner, cb, eval_grad_cb, nV=nV)
        else
            # Take care to convert 1-indexing to 0-indexing!
            # KNITRO supports only Int32 array for integer.
            jacIndexVars = Int32[j - 1 for (_, j) in jacob_structure]
            # NLP constraints are set after all other constraints
            # inside Knitro.
            jacIndexCons = Int32[i - 1 + offset for (i, _) in jacob_structure]
            KN_set_cb_grad(model.inner, cb, eval_grad_cb, nV=nV,
                           jacIndexCons=jacIndexCons, jacIndexVars=jacIndexVars)
        end

        if has_hessian
            # Get hessian structure.
            hessian_structure = MOI.hessian_lagrangian_structure(model.nlp_data.evaluator)
            nnzH = length(hessian_structure)
            # Take care to convert 1-indexing to 0-indexing!
            # Knitro supports only Int32 array for integer.
            hessIndexVars1 = Int32[i - 1 for (i, _) in hessian_structure]
            hessIndexVars2 = Int32[j - 1 for (_, j) in hessian_structure]

            KN_set_cb_hess(model.inner, cb, nnzH, eval_h_cb,
                           hessIndexVars1=hessIndexVars1,
                           hessIndexVars2=hessIndexVars2)
        end

        model.nlp_loaded = true
    elseif !isa(model.objective, Nothing)
        add_objective!(model, model.objective)
    end

    KN_solve(model.inner)
    model.number_solved += 1
    return
end

include(joinpath("MOI_wrapper", "variables.jl"))
include(joinpath("MOI_wrapper", "constraints.jl"))
include(joinpath("MOI_wrapper", "objective.jl"))
include(joinpath("MOI_wrapper", "results.jl"))
include(joinpath("MOI_wrapper", "nlp.jl"))


__precompile__(true)

# TODO x/0 = 0
# TODO ϕ * 2, ϕ1 + 5, etc...
# TODO ϕ .< 3, ϕ1 == ϕ2, etc... (how would this work?)
# TODO assignment tuple indexing, e.g. ϕ["a", 1]
# TODO push(ϕ, dims...), push(ϕ, Pair{Symbol}...)
# TODO better / faster broadcast
# TODO custom error type for NegativeError
# TODO also, better error type of others (e.g. Bounds error, DimensionMismatch, ...)
# TODO broadcast_reduce tag-team
# TODO map_reduce as well
# TODO add @boundscheck and @inbounds where applicable

module Factors

import Base: .==, .!=, .<, .<=, .>, .>=, in, ==, *, /, +, -, hash
import DataFrames

export
    # Dimension stuff
    Dimension,
    RangeDimension,

    name,
    eltype,
    support,
    spttype,
    update,

    Assignment,

    Factor,
    # basic access
    scope,
    potential,
    lengths,
    # interact with dimensions
    getdim,
    push,
    pattern,
    pattern_states,
    getdim,
    reducedim!,
    Z,
    marginalize,
    # assignment conversions
    at2sub,
    sub2at,
    at2a,
    a2at,
    a2sub,
    sub2a,

    # negatives stuff
    set_negative_mode,
    NegativeMode,
    NegativeIgnore,
    NegativeWarn,
    NegativeError

include("dimensions.jl")   # define Dimension
include("factors_main.jl") # define Factors
include("auxiliary.jl")    # define Assignment

include("factors_broadcast.jl")
include("factors_index.jl")
include("factors_iter.jl")
include("factors_join.jl")
include("factors_map.jl")
include("factors_reduce.jl")

include("dataframes.jl")
include("errors.jl")
include("io.jl")
include("negatives.jl")

end # module