-- A graph which is guaranteed to always have edges with the same weight in both directions ("undirected edges");
-- under the hood this is still stored as a directed graph.
-- Thus, for consistency, edges (except for trivial cycles) will be iterated twice when using `:edges()`,
-- once as `x, y, weight` and a second time as `y, x, weight`.

local heap = require("data_structures.heap")
local table_heap = require("data_structures.table_heap")
local union_find = require("data_structures.union_find")
local graph = require("data_structures.graph")

local undirected_graph = {}

-- Create a new graph from undirected edges.
-- An undirected edge between nodes `x` and `y` can be specified as
-- `nodes[x][y] = weight` or `nodes[y][x] = weight`.
-- If both are specified, the weights must match.
function undirected_graph.new(nodes)
	nodes = nodes or {}
	for from, tos in pairs(nodes) do
		for to, weight in pairs(tos) do
			assert(nodes[to], "destination node missing")
			if nodes[to][from] == nil then
				nodes[to][from] = weight
			else
				assert(nodes[to][from] == weight, "weights don't match")
			end
		end
	end
	return graph.new(nodes)
end

-- Set the weight of an edge in both directions. Overrides a previous weight.
-- `node` and `other_node` must already exist in the graph.
function undirected_graph:set_weight(
	node,
	other_node,
	weight -- weight of the edge, `nil` to remove; use `true` if there is no `weight`
)
	graph.set_weight(self, node, other_node, weight)
	graph.set_weight(self, other_node, node, weight)
end

-- Partitions the graph into connected components
--> Iterator over connected components (undirected subgraphs)
function undirected_graph:connected_components()
	local seen = {}
	local iterator, state, root = self:nodes()
	return function()
		repeat
			root = iterator(state, root)
		until not seen[root] -- note: terminates if `root == nil`
		if root == nil then
			return nil
		end
		local connected_component = undirected_graph.new()
		connected_component:add_node(root)
		local to_visit = { root } -- stack of nodes to visit (depth-first traversal)
		seen[root] = true
		repeat
			local node = table.remove(to_visit)
			for neighbor, weight in self:neighbors(node) do
				if not connected_component:has_node(neighbor) then
					connected_component:add_node(neighbor)
				end
				connected_component:set_weight(node, neighbor, weight)
				if not seen[neighbor] then
					seen[neighbor] = true
					table.insert(to_visit, neighbor)
				end
			end
		until #to_visit == 0
		return connected_component
	end
end

-- Finds a Minimum Spanning Forest using Prim's algorithm
function undirected_graph:msf_prim()
	local spanning_forest = undirected_graph.new()

	local function grow_spanning_tree(root)
		local min_dist = {} -- [node] = minimum distance to reach from any node of spanning tree
		local predec = {} -- [node] = other node such that the weight of the edge is minimal
		-- Immediate neighbors by their distance to the spanning tree (always a single edge!)
		local neighbors = table_heap.new({}, function(a, b)
			return min_dist[a] < min_dist[b]
		end)
		-- Update the neighbors of a node which has already been added to the spanning tree
		local function update_neighbors(node)
			for neighbor, weight in self:neighbors(node) do
				if not spanning_forest:has_node(neighbor) then -- neighbor not in spanning tree?
					if min_dist[neighbor] == nil then -- add neighbor
						min_dist[neighbor], predec[neighbor] = weight, node
						neighbors:push(neighbor)
					elseif weight < min_dist[neighbor] then -- update neighbor: cheaper edge found
						min_dist[neighbor], predec[neighbor] = weight, node
						neighbors:decrease(neighbor)
					end
				end
			end
			-- These entries aren't needed anymore, node is now part of the spanning tree
			min_dist[node], predec[node] = nil, nil
		end

		spanning_forest:add_node(root)
		update_neighbors(root) -- update nodes reachable from this "root"
		while #neighbors > 0 do -- grow the spanning tree while there are still reachable nodes
			-- Pick the closest neighbor of the current spanning tree...
			local node = neighbors:pop()
			spanning_forest:add_node(node) -- ... add the node
			spanning_forest:set_weight(predec[node], node, min_dist[node]) -- ... and connect it using the cheapest edge
			update_neighbors(node) -- now update the neighbors of the spanning tree
		end
	end

	-- Grow a spanning tree for each root that isn't yet part of one
	local n_conn_comps = 0
	for root in self:nodes() do
		if not spanning_forest:has_node(root) then
			n_conn_comps = n_conn_comps + 1
			grow_spanning_tree(root)
		end
	end

	return spanning_forest, n_conn_comps
end

-- Finds a Minimum Spanning Forest using Kruskal's algorithm
function undirected_graph:msf_kruskal()
	local spanning_forest = undirected_graph.new()

	-- Build a heap of edges by weight (we could also sort, but using a heap is more efficient in the best case)
	local edges = {}
	for node, other_node, weight in self:edges() do
		table.insert(edges, { node = node, other_node = other_node, weight = weight })
	end
	edges = heap.new(edges, function(a, b)
		return a.weight < b.weight
	end)

	local connected_components = union_find.new()
	local n_conn_comps = 0
	for node in self:nodes() do
		spanning_forest:add_node(node)
		connected_components:make_set(node)
		n_conn_comps = n_conn_comps + 1
	end

	while n_conn_comps > 1 and not edges:empty() do
		local min_edge = edges:pop()
		local node, other_node = min_edge.node, min_edge.other_node
		-- Nodes are in two distinct connected components currently
		-- <=> adding the edge does not introduce a cycle
		if connected_components:find(node) ~= connected_components:find(other_node) then
			-- Add edge, connecting the two components.
			connected_components:union(node, other_node)
			spanning_forest:set_weight(node, other_node, min_edge.weight)
			n_conn_comps = n_conn_comps - 1
		end
	end

	return spanning_forest, n_conn_comps
end

-- Default choice for finding Minimum Spanning Forests is Prim's algorithm:
-- Prim theoretically runs faster than Kruskal for graphs with n > m,
-- requiring O(m log n) vs. O(m log m) time.
--> spanning forest (undirected graph), number of connected components
undirected_graph.msf = undirected_graph.msf_prim -- TODO (...) benchmark Kruskal vs Prim for small & sparse graphs

return require("class")(undirected_graph, graph)