// =====================================================================================================================
//  fennec, a free and open source game engine
//  Copyright © 2025  Medusa Slockbower
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  This program is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <https://www.gnu.org/licenses/>.
// =====================================================================================================================

#ifndef FENNEC_CONTAINERS_SET_H
#define FENNEC_CONTAINERS_SET_H

// https://programming.guide/robin-hood-hashing.html

#include <fennec/containers/optional.h>
#include <fennec/containers/set.h>
#include <fennec/lang/compare.h>
#include <fennec/math/ext/primes.h>
#include <fennec/memory/allocator.h>
#include <fennec/lang/hashing.h>

namespace fennec
{

///
///
/// \brief wrapper for sets of elements
/// \details
/// This data-structure behaves like a set, but does not use pointers, instead storing the table in-array
///
/// | Property  | Value      |
/// |:----------|:----------:|
/// | stable    | \emoji x   |
/// | access    | \f$O(1)\f$ |
/// | insertion | \f$O(1)\f$ |
/// | deletion  | \f$O(1)\f$ |
/// | space     | \f$O(1)\f$ |
///
/// \tparam TypeT The type to contain
template<typename TypeT, class Hash = hash<TypeT>, class Equals = equality<TypeT>, class Alloc = allocator<TypeT>>
struct set {

// Definitions =========================================================================================================
public:
	using alloc_t = typename allocator_traits<Alloc>::template rebind<TypeT>;
	using hash_t = Hash;
	using equal_t = Equals;
	using elem_t = TypeT;

	class iterator;
	class value_iterator;
	static constexpr size_t npos = -1;
	static constexpr double default_load = 0.8;

private:
	struct node {
		optional<elem_t> value;
		int psl;

		constexpr node() = default;
		constexpr ~node() = default;
	};

// Constructors ========================================================================================================
public:

	///
	/// \brief Default Constructor, initializes empty set
	constexpr set()
		: _alloc()
		, _hash()
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	};

	///
	/// \brief Hash Copy Constructor, initializes empty set with a hash
	constexpr set(const hash_t& hash)
		: _alloc()
		, _hash(hash)
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Hash Move Constructor, initializes empty set with a hash
	constexpr set(hash_t&& hash) noexcept
		: _alloc()
		, _hash(hash)
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Alloc Copy Constructor, initializes empty set with an allocator
	constexpr set(const alloc_t& alloc)
		: _alloc(alloc)
		, _hash()
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Alloc Move Constructor, initializes empty set with an allocator
	constexpr set(alloc_t&& alloc) noexcept
		: _alloc(alloc)
		, _hash()
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Hash Alloc Copy Constructor, initializes empty set with a hash and allocator
	constexpr set(const hash_t& hash, const alloc_t& alloc)
		: _alloc(alloc)
		, _hash(hash)
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Hash Copy Alloc Move Constructor, initializes empty set with a hash and allocator
	constexpr set(const hash_t& hash, alloc_t&& alloc) noexcept
		: _alloc(alloc)
		, _hash(hash)
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Hash Move Alloc Move Constructor, initializes empty set with a hash and allocator
	constexpr set(hash_t&& hash, alloc_t&& alloc) noexcept
		: _alloc(alloc)
		, _hash(hash)
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Hash Move Alloc Copy Constructor, initializes empty set with a hash and allocator
	constexpr set(hash_t&& hash, const alloc_t& alloc) noexcept
		: _alloc(alloc)
		, _hash(hash)
		, _size(0)
		, _sumpsl(0)
		, _load(default_load) {
	}

	///
	/// \brief Set Copy Constructor
	/// \param set Set to copy
	constexpr set(const set& set)
		: _alloc(set._alloc)
		, _hash(set._hash)
		, _size(set._size)
		, _sumpsl(set._sumpsl)
		, _load(set._load) {
	}

	///
	/// \brief Set Move Constructor
	/// \param set Set to move
	constexpr set(set&& set) noexcept
		: _alloc(fennec::move(set._alloc))
		, _hash(fennec::move(set._hash))
		, _size(fennec::move(set._size))
		, _sumpsl(set._sumpsl)
		, _load(set._load) {
	}

	///
	/// \brief Destructor, destructs all elements and releases the allocation
	constexpr ~set() {
		for (size_t i = 0; i < capacity(); ++i) {
			_alloc[i].value = nullopt;
		}
	}


// Properties ==========================================================================================================

	///
	/// \returns Size of the set in elements
	constexpr size_t size() const {
		return _size;
	}

	///
	/// \returns Capacity of the set in elements
	constexpr size_t capacity() const {
		return _alloc.capacity();
	}


// Access ==============================================================================================================

	///
	/// \brief Find an Element
	/// \param val Value to find
	/// \returns An iterator at the location of the first instance of `value`
	constexpr iterator find(const elem_t& val, size_t c = 0) const {
		if (capacity() == 0) {
			return end();
		}
		size_t s = _hash(val) % capacity();              // Initial search index
		int    psl = (_size != 0) ? _sumpsl / _size : 0; // Initial psl
		size_t i = (s + psl) % capacity();               // Median search
		size_t n = 0;

		// Check the first element;
		if (_alloc[i].psl >= psl && _alloc[i].value) {
			if (*_alloc[i].value == val) {
				return iterator(this, i);
			}
		}

		// Loop while there is a value and its psl is greater than our probe
		while (c > 0) {
			++n;
			size_t i0 = (i + capacity() - n) % capacity(); // Prevent index underflow
			size_t i1 = (i + n) % capacity();
			int    p0 = psl - n, p1 = psl + n;
			bool   c0  = p0 >= 0 && _alloc[i0].psl >= p0, c1 = _alloc[i1].psl >= p1; // Check that we are in range

			if (c0 && _alloc[i0].value) {
				if (*_alloc[i0].value == val) {
					if (c-- == 0) {
						return iterator(this, i0);
					}
				}
			}

			if (c1 && _alloc[i1].value) {
				if (*_alloc[i1].value == val) {
					if (c-- == 0) {
						return iterator(this, i1);
					}
				}
			}

			if (not(c0 or c1)) {
				break;
			}
		}

		return iterator(this, npos);
	}

	///
	/// \brief Check if a set contains a value
	/// \param val Value to check
	/// \returns `true` if `val` can be found, `false` otherwise
	constexpr bool contains(const elem_t& val) const {
		return this->find(val) != end();
	}

	///
	/// \brief Iterator Access
	/// \param it Location to access
	/// \returns A pointer to the element, `nullptr` if not found.
	///          The value should not be changed in a manner that will change the hash of the element.
	constexpr elem_t* at(const iterator& it) {
		if (it == end()) {
			return nullptr;
		}
		if (not _alloc[it._i].value) {
			return nullptr;
		}
		return &*_alloc[it._i].value;
	}

	///
	/// \brief Iterator Const Access
	/// \param it Location to access
	/// \returns A const-qualified pointer to the element, `nullptr` if not found.
	constexpr const elem_t* at(const iterator& it) const {
		if (not _alloc[it._i].value) return nullptr;
		return &*_alloc[it._i].value;
	}

// Insertion & Deletion ================================================================================================

	///
	/// \brief Move Insertion
	/// \param val Value to insert
	constexpr void insert(elem_t&& val) {
		this->_insert(fennec::forward<elem_t>(val));
	}

	///
	/// \brief Copy Insertion
	/// \param val Value to insert
	constexpr void insert(const elem_t& val) {
		this->_insert(val);
	}

	///
	/// \brief Emplace Insertion
	/// \tparam ArgsT Argument types
	/// \param args Arguments to construct with
	template<typename...ArgsT>
	constexpr void emplace(ArgsT&&...args) {
		this->_insert(fennec::forward<ArgsT>(args)...);
	}

	///
	/// \brief Element Erase
	/// \param it Location to erase
	constexpr void erase(iterator it) {
		size_t i = it._i;
		if (i >= capacity()) {
			return;
		} // These are separated due to compilers being inconsistent
		if (not _alloc[i].value) {
			return;
		}

		_alloc[i].value = nullopt;
		_sumpsl -= _alloc[i].psl;
		--_size;
		size_t p = i;
		while (_alloc[i = (i + 1) % capacity()].value) {
			if (_alloc[i].psl == 0) break;

			fennec::swap(_alloc[i - 1].value, _alloc[i].value);
			--_alloc[p].psl, --_sumpsl;
			p = i;
		}
	}

	///
	/// \brief Element Erase
	/// \param val Value to erase
	constexpr void erase(const elem_t& val) {
		this->erase(this->find(val));
	}


// ITERATOR ============================================================================================================

	///
	/// \brief Class for Iterating the Set
	class iterator {
	public:
		constexpr ~iterator() {
			_set = nullptr;
		}

		// prefix operator
		constexpr iterator& operator++() {
			while (++_i < _set->capacity()) {
				if (_set->_alloc[_i].value) {
					return *this;
				}
			}
			_i = npos;
			return *this;
		}

		constexpr iterator operator++(int) {
			iterator prev = *this;
			this->operator++();
			return prev;
		}

		constexpr const elem_t& operator*() const {
			return *_set->_alloc[_i].value;
		}

		constexpr const elem_t* operator->() const {
			if (not _set->_alloc[_i].value) return nullptr;
			return &*_set->_alloc[_i].value;
		}

		constexpr bool operator==(const iterator& it) const {
			return _set == it._set and _i == it._i;
		}

		constexpr bool operator!=(const iterator& it) const {
			return _set != it._set or _i != it._i;
		}

	private:
		const set* _set;
		size_t _i;
		friend set;

		constexpr iterator(const set* set, size_t i)
			: _set(set)
			, _i(i) {
		}
	};

	class value_iterator {
	public:
		constexpr ~value_iterator() {
			_set = nullptr;
		}

		// prefix operator
		constexpr value_iterator& operator++() {
			while (_psl <= _set->_alloc[_i].psl) {
				if (not _set->_alloc[_i].value) {
					break;
				}
			}
			_i = npos;
			return *this;
		}

		constexpr value_iterator operator++(int) {
			value_iterator prev = *this;
			this->operator++();
			return prev;
		}

		constexpr const elem_t& operator*() const {
			return *_set->_alloc[_i].value;
		}

		constexpr const elem_t* operator->() const {
			if (not _set->_alloc[_i].value) return nullptr;
			return &*_set->_alloc[_i].value;
		}

		constexpr bool operator==(const iterator& it) const {
			return _set == it._set and _i == it._i;
		}

		constexpr bool operator!=(const iterator& it) const {
			return _set != it._set or _i != it._i;
		}

	private:
		const set* _set;
		size_t _i;
		int _psl;
		elem_t _value;
		friend set;

		constexpr value_iterator(const set* set, size_t i, int psl, const elem_t& value)
			: _set(set)
			, _i(i)
			, _value(value) {
		}
	};

	constexpr iterator begin() const {
		iterator it(this, 0);
		if (not _alloc[it._i].value) {
			++it;
		}
		return it;
	}

	constexpr iterator end() const {
		return iterator(this, npos);
	}


// PRIVATE =============================================================================================================

private:
	constexpr void _expand() {
		set cpy; // Create a new set
		cpy._alloc.callocate(
			fennec::next_prime2(_alloc.capacity())
		);

		// rehash
		for (size_t i = 0; i < capacity(); ++i) {
			if (_alloc[i].value) {
				cpy.insert(fennec::move(*_alloc[i].value));
			}
		}

		// Swap buffers
		fennec::swap(_alloc, cpy._alloc);
	}

	template<typename...ArgsT>
	constexpr void _insert(ArgsT&&...args) {
		if (_size == 0 or double(_size) / capacity() >= _load) { // expand when full
			_expand();
		}

		elem_t value(fennec::forward<ArgsT>(args)...);
		size_t i = _hash(value) % capacity(); // Initial search index
		int psl = 0;
		while (_alloc[i].value) { // Search for empty cell
			if (_equal(*_alloc[i].value, value)) { // Check to see if this element is already inserted
				return;
			}
			if (psl > _alloc[i].psl) { // When psl is higher, swap
				_sumpsl += psl - _alloc[i].psl;
				fennec::swap(_alloc[i].psl, psl);
				fennec::swap(*_alloc[i].value, value);
			}
			i = (i + 1) % capacity(); ++psl;
		}
		_alloc[i].value = fennec::move(value);
		_sumpsl += (_alloc[i].psl = psl);
		++_size;
	}

	allocation<node, alloc_t> _alloc;
	hash_t  _hash;
	equal_t _equal;
	size_t  _size;
	size_t  _sumpsl;
	double  _load;
};

}

#endif // FENNEC_CONTAINERS_SET_H
