#ifndef GDWG_GRAPH_H
#define GDWG_GRAPH_H

#include <initializer_list>
#include <algorithm>
#include <exception>
#include <iterator>
#include <memory>
#include <numeric>
#include <optional>
#include <set>
#include <sstream>
#include <string>
#include <vector>

namespace gdwg {
	template<typename N, typename E>
	class edge {
	 protected:
		N src;
		N dst;

	 public:
		edge(const N& s, const N& d)
		: src(s)
		, dst(d) {}
		virtual ~edge() = default;

		virtual auto print_edge() const -> std::string = 0;
		virtual auto is_weighted() const noexcept -> bool = 0;
		virtual auto get_weight() const -> std::optional<E> = 0;

		auto get_nodes() const -> std::pair<N, N> {
			return std::pair<N, N>{this->src, this->dst};
		}

		auto operator==(const edge& other) const noexcept -> bool {
			if (this->src != other.src) {
				return false;
			}
			return this->dst == other.dst;
		};
	};

	template<typename N, typename E>
	class weighted_edge : virtual public edge<N, E> {
	 protected:
		E weight;

	 public:
		weighted_edge(const N& s, const N& d, const E& w)
		: edge<N, E>(s, d)
		, weight(w) {}
		virtual ~weighted_edge() = default;

		virtual auto print_edge() const -> std::string override {
			std::stringstream ss{};
			ss << this->src << " -> " << this->dst << " | W | " << this->weight;
			return ss.str();
		}
		virtual auto is_weighted() const noexcept -> bool override {
			return true;
		}
		virtual auto get_weight() const -> std::optional<E> override {
			return this->weight;
		}
		auto operator==(const weighted_edge<N, E>& other) const noexcept -> bool {
			if (edge<N, E>::operator!=(*this, other)) {
				return false;
			}
			return this->weight == other.weight;
		};
	};

	template<typename N, typename E>
	class unweighted_edge : virtual public edge<N, E> {
	 private:
	 public:
		unweighted_edge(const N& s, const N& d)
		: edge<N, E>(s, d) {}
		virtual ~unweighted_edge() = default;

		virtual auto print_edge() const -> std::string override {
			std::stringstream ss{};
			ss << this->src << " -> " << this->dst << " | U";
			return ss.str();
		}
		virtual auto is_weighted() const noexcept -> bool override {
			return false;
		}
		virtual auto get_weight() const -> std::optional<E> override {
			return std::nullopt;
		}
	};

	template<typename N, typename E>
	class graph {
	 private:
		// We are not allowed to store redundant copies of nodes, so we use
		// an unordered_set of N to avoid unnecessary copies. As the edge class
		// copies the strings internally we cannot use them for our internal
		// edge representation.
		// We define two custom hashing functions here to force our shared_ptrs
		// to use their internally stored strings for comparison
		// such that we can avoid creating a map which inevitably stores a
		// redundant copy of the string as the key.
		struct shared_edge {
			std::shared_ptr<N> from;
			std::shared_ptr<N> to;
			std::optional<E> weight;

			// Required for multiset comparison.
			auto operator==(const shared_edge& other) const noexcept -> bool {
				if (*this->from != *other.from) {
					return false;
				}
				if (*this->to != *other.to) {
					return false;
				}
				return this->weight == other.weight;
			}
		};

		using nodes_comp =
		    decltype([](const std::shared_ptr<N>& a, const std::shared_ptr<N>& b) noexcept -> bool { return *a < *b; });
		using edges_comp = decltype([](const shared_edge& a, const shared_edge& b) -> bool {
			if (*a.from != *b.from) {
				return *a.from < *b.from;
			}
			if (a.to == nullptr) { // necessary for our std::lower_bound
				return true;
			}
			if (b.to == nullptr) {
				return false;
			}
			if (*a.to != *b.to) {
				return *a.to < *b.to;
			}
			return a.weight < b.weight;
		});

	 public:
		using nodes_t = std::set<std::shared_ptr<N>, nodes_comp>;
		using edges_t = std::set<shared_edge, edges_comp>;

	 private:
		nodes_t node_ptrs{};
		edges_t shared_edges{};

	 private:
		// Efficiently find iterators for some N.
		auto get_edges_from_n(const N& n) const noexcept -> auto {
			const auto lower = this->shared_edges.lower_bound(
			    shared_edge{.from = std::make_shared<N>(n), .to = nullptr, .weight = std::nullopt});

			auto upper = lower;
			for (; upper != std::end(this->shared_edges); ++upper) {
				if (*upper->from != n) {
					break;
				}
			}
			return std::make_pair(lower, upper);
		}
		// Transforms a shared_edge to a unique_ptr of the appropriate edge.
		static auto make_edge(const shared_edge& e) -> std::unique_ptr<edge<N, E>> {
			if (not e.weight.has_value()) {
				return std::make_unique<unweighted_edge<N, E>>(*e.from, *e.to);
			}
			return std::make_unique<weighted_edge<N, E>>(*e.from, *e.to, *e.weight);
		};

	 public:
		class iterator {
			friend graph;

		 public:
			using value_type = struct {
				N from;
				N to;
				std::optional<E> weight;
			};
			using reference = value_type;
			using pointer = void;
			using difference_type = std::ptrdiff_t;
			using iterator_category = std::bidirectional_iterator_tag;

		 private:
			edges_t::iterator iter;
			value_type contents;

		 private:
			static auto iter_to_value(const edges_t::iterator i) -> value_type {
				return value_type{
				    .from = *i->from,
				    .to = *i->to,
				    .weight = i->weight,
				};
			}

		 public:
			// Iterator constructor
			iterator() noexcept = default;
			explicit iterator(edges_t::iterator it) noexcept
			: iter(it) {}

			// Iterator source
			auto operator*() -> reference {
				this->contents = iter_to_value(this->iter);
				return contents;
			}

			// Iterator traversal
			auto operator++() -> iterator& {
				++this->iter;
				return *this;
			}
			auto operator++(int) -> iterator {
				iterator ret = *this;
				++(*this);
				return ret;
			}
			auto operator--() -> iterator& {
				--this->iter;
				return *this;
			}
			auto operator--(int) -> iterator {
				iterator ret = *this;
				--(*this);
				return ret;
			}

			// Iterator comparison
			friend auto operator==(const iterator& a, const iterator& b) noexcept -> bool {
				return a.iter == b.iter;
			}
		};
		[[nodiscard]] auto begin() const -> iterator {
			return iterator{std::begin(this->shared_edges)};
		}
		[[nodiscard]] auto end() const -> iterator {
			return iterator{std::end(this->shared_edges)};
		}

	 public:
		graph() noexcept {}
		graph(std::initializer_list<N> il) noexcept
		: graph(std::begin(il), std::end(il)) {}
		template<typename InputIt>
		graph(InputIt first, InputIt last) {
			std::transform(first, last, std::inserter(this->node_ptrs, std::end(this->node_ptrs)), [](const auto& n) {
				return std::make_shared<N>(n);
			});
		}
		graph(graph&& other) noexcept
		: node_ptrs(std::move(other.node_ptrs))
		, shared_edges(std::move(other.shared_edges)) {
			other.clear();
		}
		graph(const graph& other)
		: node_ptrs(other.node_ptrs)
		, shared_edges(other.shared_edges) {}

		auto operator=(graph&& other) noexcept {
			if (&other != this) {
				this->node_ptrs = std::move(other.node_ptrs);
				this->shared_edges = std::move(other.shared_edges);
				other.clear();
			}
			return *this;
		}
		auto operator=(const graph& other) -> graph& {
			if (&other != this) {
				this->node_ptrs = other.node_ptrs;
				this->shared_edges = other.shared_edges;
			}
			return *this;
		}
		friend auto operator==(const graph& a, const graph& b) noexcept -> bool {
			if (not std::equal(std::begin(a.node_ptrs),
			                   std::end(a.node_ptrs),
			                   std::begin(b.node_ptrs),
			                   std::end(b.node_ptrs),
			                   [](const auto& a_v, const auto& b_v) { return *a_v == *b_v; }))
			{
				return false;
			}
			return a.shared_edges == b.shared_edges;
		}
		friend auto operator<<(std::ostream& os, const graph& g) -> std::ostream& {
			os << '\n';
			for (const auto& n : g.node_ptrs) {
				os << *n << " (\n";

				const auto& [lower, upper] = g.get_edges_from_n(*n);
				std::for_each(lower, upper, [&](const auto& e) {
					const auto print = make_edge(e);
					os << "  " << print->print_edge() << '\n';
				});

				os << ")\n";
			}
			return os;
		}

		auto insert_node(const N& value) -> bool {
			return this->node_ptrs.insert(std::make_shared<N>(value)).second;
		}
		auto insert_edge(const N& s, const N& d, std::optional<E> w = std::nullopt) -> bool {
			const auto s_it = this->node_ptrs.find(std::make_shared<N>(s));
			const auto d_it = this->node_ptrs.find(std::make_shared<N>(d));
			if (s_it == std::end(this->node_ptrs) || d_it == std::end(this->node_ptrs)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::insert_edge when either src or dst node "
				                         "does "
				                         "not exist");
			}

			const auto insert = shared_edge{.from = *s_it, .to = *d_it, .weight = w};
			if (this->shared_edges.contains(insert)) {
				return false;
			}

			this->shared_edges.insert(std::move(insert));
			return true;
		}
		auto replace_node(const N& old_node, const N& new_node) -> bool {
			const auto old_it = this->node_ptrs.find(std::make_shared<N>(old_node));
			if (old_it == std::end(this->node_ptrs)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::replace_node on a node that doesn't exist");
			}

			const auto insert = std::make_shared<N>(new_node);
			if (this->node_ptrs.contains(insert)) {
				return false;
			}

			this->node_ptrs.erase(old_it);
			this->node_ptrs.insert(insert);
			return true;
		}
		auto merge_replace_node(const N& old_node, const N& new_node) -> void {
			const auto old_it = this->node_ptrs.find(std::make_shared<N>(old_node));
			const auto new_it = this->node_ptrs.find(std::make_shared<N>(new_node));
			if (old_it == std::end(this->node_ptrs) or new_it == std::end(this->node_ptrs)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::merge_replace_node on old or new data if they "
				                         "don't exist in the graph");
			}

			// This operation completely transforms all nodes and replaces all
			// existences of old_node with new_node (using the existing shared_ptr).
			auto result = edges_t{};
			std::transform(std::begin(this->shared_edges),
			               std::end(this->shared_edges),
			               std::inserter(result, std::end(result)),
			               [&](const auto& e) {
				               return shared_edge{.from = (e.from == *old_it ? *new_it : e.from),
				                                  .to = (e.to == *old_it ? *new_it : e.to),
				                                  .weight = e.weight};
			               });

			this->shared_edges = std::move(result);
			this->node_ptrs.erase(old_it);
		}
		auto erase_node(const N& n) -> bool {
			const auto node_it = this->node_ptrs.find(std::make_shared<N>(n));
			if (node_it == std::end(this->node_ptrs)) {
				return false;
			}

			std::erase_if(this->shared_edges, [&](const auto& e) {
				if (*e.from == **node_it) {
					return true;
				}
				if (*e.to == **node_it) {
					return true;
				}
				return false;
			});
			this->node_ptrs.erase(node_it);
			return true;
		}
		auto erase_edge(const N& s, const N& d, std::optional<E> w = std::nullopt) -> bool {
			const auto s_it = this->node_ptrs.find(std::make_shared<N>(s));
			const auto d_it = this->node_ptrs.find(std::make_shared<N>(d));
			if (s_it == std::end(this->node_ptrs) || d_it == std::end(this->node_ptrs)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::erase_edge on src or dst if they don't exist "
				                         "in the graph");
			}

			const auto& [lower, upper] = this->get_edges_from_n(s);
			if (const auto it = std::find_if(lower,
			                                 upper,
			                                 [&](const auto& e) {
				                                 if (*e.from != s) {
					                                 return false;
				                                 }
				                                 if (*e.to != d) {
					                                 return false;
				                                 }
				                                 return e.weight == w;
			                                 });
			    it != upper)
			{
				this->shared_edges.erase(it);
				return true;
			}

			return false;
		}
		auto erase_edge(const iterator& i) -> iterator {
			return iterator{this->shared_edges.erase(i.iter)};
		}
		auto erase_edge(const iterator& i, const iterator& s) -> iterator {
			return iterator{this->shared_edges.erase(i.iter, s.iter)};
		}

		auto clear() noexcept -> void {
			this->node_ptrs.clear();
			this->shared_edges.clear();
		}

		[[nodiscard]] auto is_node(const N& value) const noexcept -> bool {
			return this->node_ptrs.contains(std::make_shared<N>(value));
		}
		[[nodiscard]] auto empty() const noexcept -> bool {
			return this->node_ptrs.empty();
		}
		[[nodiscard]] auto nodes() const -> std::vector<N> {
			auto result = std::vector<N>{};
			std::transform(std::begin(this->node_ptrs),
			               std::end(this->node_ptrs),
			               std::back_inserter(result),
			               [](const auto& n) { return *n; });
			return result;
		}
		[[nodiscard]] auto is_connected(const N& s, const N& d) const -> bool {
			if (not this->is_node(s) or not this->is_node(d)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::is_connected if src or dst node don't exist "
				                         "in the graph");
			}

			const auto& [lower, upper] = this->get_edges_from_n(s);
			return std::any_of(lower, upper, [&](const auto& e) { return *e.to == d; });
		}

		[[nodiscard]] auto edges(const N& s, const N& d) const -> std::vector<std::unique_ptr<edge<N, E>>> {
			if (not this->is_node(s) or not this->is_node(d)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::edges if src or dst node don't exist in the "
				                         "graph");
			}

			const auto& [lower, upper] = this->get_edges_from_n(s);

			auto result = std::vector<std::unique_ptr<edge<N, E>>>{};
			std::for_each(lower, upper, [&](const auto& e) {
				if (*e.to != d) {
					return;
				}
				result.push_back(make_edge(e));
			});
			return result;
		}
		[[nodiscard]] auto find(const N& s, const N& d, std::optional<E> w = std::nullopt) const -> iterator {
			const auto& [lower, upper] = this->get_edges_from_n(s);

			if (const auto it = std::find_if(lower,
			                                 upper,
			                                 [&](const auto& e) {
				                                 if (*e.to != d) {
					                                 return false;
				                                 }
				                                 return e.weight == w;
			                                 });
			    it != upper)
			{
				return iterator{it};
			}
			return std::end(*this);
		}
		[[nodiscard]] auto connections(const N& s) const -> std::vector<N> {
			if (not this->is_node(s)) {
				throw std::runtime_error("Cannot call gdwg::graph<N, E>::connections if src doesn't exist in the "
				                         "graph");
			}

			const auto& [lower, upper] = this->get_edges_from_n(s);

			auto result = std::set<N>{};
			std::transform(lower, upper, std::inserter(result, std::end(result)), [](const auto& e) {
				const auto n = *e.to;
				return n;
			});

			return std::vector<N>{std::begin(result), std::end(result)};
		}
	};

} // namespace gdwg

#endif // GDWG_GRAPH_H