#pragma once
#include <functional>
#include <memory>
#include <utility>
#include <iostream>
#include <memory>

#include <useful/base.hpp>

using namespace uf::sized;

namespace hm
{
    template<typename Tp, typename Allocator = std::allocator<Tp>>
    class storage
    {
    public:
        using value_type = Tp;
        using pointer = value_type*;
        using const_pointer = const value_type*;
        using allocator_type = Allocator;
        using size_type = u64;

    private:
        allocator_type m_allocator;
        std::vector<u8, allocator_type> m_constructed;
        pointer m_data = nullptr;
        size_type m_constructed_count = 0;

    public:
        storage() = default;

        template<typename A>
        storage(const storage& other, A&& allocator) : m_allocator(std::forward<A>(allocator)), m_constructed(other.m_constructed), m_data(m_allocator.allocate(other.size()))
        {
            for (size_type i = 0; i < size(); ++i)
                if (m_constructed[i])
                    construct(i, other.m_data[i]);
        }

        storage(const storage& other) : storage(other, std::allocator_traits<allocator_type>::select_on_container_copy_construction(other.get_allocator())) { }

        template<typename A>
        storage(storage&& other, A&& allocator) noexcept : m_allocator(std::forward<A>(allocator)), m_constructed(std::move(other.m_constructed)), m_data(std::move(other.m_data))
        {
            other.self_destruction();
        }

        storage(storage&& other) noexcept : storage(std::move(other), std::move(other.m_allocator)) { }

        template<typename A>
        storage(A&& allocator, size_type size) : m_allocator(std::forward<A>(allocator)),  m_constructed(size, 0, m_allocator), m_data(m_allocator.allocate(size)) { }

        template<typename A>
        storage(A&& allocator) : storage(std::forward<A>(allocator), 0) { }

        storage(size_type size) : storage(Allocator(), size) { }

        storage& operator=(const storage& other)
        {
            if (std::allocator_traits<allocator_type>::propagate_on_container_copy_assignment::value || m_allocator == other.m_allocator)
            {
                m_allocator = other.m_allocator;
                self_destruction();
            }
            else
            {
                self_destruction();
                m_allocator = other.m_allocator;
            }
            m_data = m_allocator.allocate(other.m_constructed.size());
            m_constructed = other.m_constructed;
            for (size_type i = 0; i < m_constructed.size(); ++i)
                if (m_constructed[i])
                    new(m_data + i) value_type(other.m_data[i]);
            return *this;
        }

        storage& operator=(storage&& other) noexcept
        {
            if constexpr (std::allocator_traits<allocator_type>::propagate_on_container_move_assignment::value)
            {
                m_allocator = other.m_allocator;
                self_destruction();
            }
            else
            {
                if (m_allocator == other.m_allocator)
                {
                    m_allocator = std::move(other.m_allocator);
                    self_destruction();
                }
                else
                {
                    self_destruction();
                    m_allocator = std::move(other.m_allocator);
                }
            }
            m_data = std::move(other.m_data);
            m_constructed = std::move(other.m_constructed);
            other.m_data = nullptr;
            return *this;
        }

        ~storage() noexcept
        {
            self_destruction();
        }

        void resize(size_type size)
        {
            if (size <= m_constructed.size())
                return; // TODO
            auto new_data = m_allocator.allocate(size);
            for (size_type i = 0; i < m_constructed.size(); ++i)
            {
                if (m_constructed[i])
                {
                    new(new_data + i) value_type(m_data[i]);
                    destruct(i);
                }
            }
            m_constructed.resize(size, 0);
            m_allocator.deallocate(m_data);
            m_data = new_data;
        }

        template<typename... Args>
        void construct(size_type index, Args&&... args)
        {
            new(m_data + index) value_type(std::forward<Args>(args)...);
            m_constructed[index] = 1;
            ++m_constructed_count;
        }

        template<typename... Args>
        void assign(size_type index, Args&&... args)
        {
            m_data[index] = value_type(std::forward<Args>(args)...);
        }

        template<typename... Args>
        void construct_or_assign(size_type index, Args&&... args)
        {
            if (is_constructed(index))
                assign(index, std::forward<Args>(args)...);
            else
                construct(index, std::forward<Args>(args)...);
        }

        void destruct(size_type index) noexcept
        {
            if (!m_constructed[index])
                return;
            std::destroy_at(m_data + index);
            m_constructed[index] = 0;
            --m_constructed_count;
        }

        bool is_constructed(size_type index) const noexcept
        {
            return m_constructed[index];
        }

        allocator_type get_allocator() const noexcept
        {
            return m_allocator;
        }

        size_type size() const noexcept
        {
            return m_constructed.size();
        }

        size_type constructed_count() const noexcept
        {
            return m_constructed_count;
        }

        const Tp* data() const noexcept
        {
            return m_data;
        }

        Tp* data() noexcept
        {
            return m_data;
        }

        const Tp& operator[](u64 index) const noexcept
        {
            return m_data[index];
        }

        Tp& operator[](u64 index) noexcept
        {
            return m_data[index];
        }

    private:
        void self_destruction() noexcept
        {
            if (!m_data)
                return;
            for (size_type i = 0; i < size(); ++i)
                if (m_constructed[i])
                    destruct(i);
            m_allocator.deallocate(m_data, size());
            m_data = nullptr;
        }
    };

    template<typename T>
    class allocator
    {
    public:
        using size_type = std::size_t;
        using difference_type = std::ptrdiff_t;
        using pointer = T*;
        using const_pointer = const T*;
        using reference = typename std::add_lvalue_reference<T>::type;
        using const_reference = typename std::add_lvalue_reference<const T>::type;
        using value_type = T;

    public:
        allocator() noexcept = default;

        allocator(const allocator&) noexcept = default;

        template <class U>
        allocator(const allocator<U>&) noexcept { }

        ~allocator() = default;

        pointer allocate(size_type n)
        {
            return static_cast<pointer>(::operator new(n * sizeof(value_type)));
        }

        void deallocate(pointer p, size_type) noexcept
        {
            ::operator delete(p);
        }

        bool operator==(const allocator&)
        {
            return true;
        }

        bool operator!=(const allocator&)
        {
            return true;
        }
    };

    enum class cell_state : u8
    {
        presented, erased, non_init
    };

    template<typename Tp>
    class hash_map_iterator
    {
    public:
        using iterator_category = std::forward_iterator_tag;
        using value_type = Tp;
        using difference_type = std::ptrdiff_t;
        using reference = value_type&;
        using pointer = value_type*;
        using size_type = u64;

    private:
        const std::vector<cell_state>& m_states;
        value_type* m_base;
        size_type m_pos;

    public:
        reference operator*() const
        {
            return m_base[m_pos];
        }

        pointer operator->() const
        {
            return m_base + m_pos;
        }

        hash_map_iterator& operator++()
        {
            ++m_pos;
            for (; m_pos < m_states.size(); ++m_pos)
                if (m_states[m_pos] == cell_state::presented)
                    break;
            return *this;
        }

        hash_map_iterator operator++(int)
        {
            hash_map_iterator result(*this);
            this->operator++();
            return result;
        }

        friend bool operator==(const hash_map_iterator<value_type>& a, const hash_map_iterator<value_type>& b)
        {
            return a.m_base == b.m_base && a.m_pos == b.m_pos;
        }

        friend bool operator!=(const hash_map_iterator<value_type>& a, const hash_map_iterator<value_type>& b)
        {
            return !operator==(a, b);
        }

    private:
        // TODO check variadic
        template<typename, typename, typename, typename, typename>
        friend class hash_map;

        template<typename>
        friend class hash_map_const_iterator;

        hash_map_iterator(const std::vector<cell_state>& states, value_type* base, size_type pos) : m_states(states), m_base(base), m_pos(pos) { }
    };

    template<typename Tp>
    class hash_map_const_iterator
    {
    public:
        using iterator_category = std::forward_iterator_tag;
        using value_type = Tp;
        using difference_type = std::ptrdiff_t;
        using reference = const value_type&;
        using pointer = const value_type*;
        using size_type = u64;

    private:
        const std::vector<cell_state>& m_states;
        value_type* m_base;
        size_type m_pos;

    public:
        hash_map_const_iterator() = default;

        hash_map_const_iterator(const hash_map_iterator<value_type>& other) : m_states(other.m_states), m_base(other.m_base), m_pos(other.m_pos) { }

        reference operator*() const
        {
            return *m_base;
        }

        pointer operator->() const
        {
            return m_base;
        }

        hash_map_const_iterator& operator++()
        {
            ++m_pos;
            for (; m_pos < m_states.size(); ++m_pos)
                if (m_states == cell_state::presented)
                    break;
            return *this;
        }

        hash_map_const_iterator operator++(int)
        {
            hash_map_iterator result(*this);
            this->operator++();
            return result;
        }

        friend bool operator==(const hash_map_const_iterator<value_type>& a, const hash_map_const_iterator<value_type>& b)
        {
            return a.m_states.data() == b.m_states.data() && a.m_pos == b.m_pos;
        }

        friend bool operator!=(const hash_map_const_iterator<value_type>& a, const hash_map_const_iterator<value_type>& b)
        {
            return !operator==(a, b);
        }

    private:
        // TODO check variadic
        template<typename, typename, typename, typename, typename>
        friend class hash_map;

        hash_map_const_iterator(const std::vector<cell_state>& states, value_type* base, size_type pos) : m_states(states), m_base(base), m_pos(pos) { }
    };

    template<typename K, typename T, typename Hash = std::hash<K>, typename Pred = std::equal_to<K>, typename Alloc = allocator<std::pair<const K, T>>>
    class hash_map
    {
    public:
        using key_type = K;
        using mapped_type = T;
        using hasher = Hash;
        using key_equal = Pred;
        using allocator_type = Alloc;
        using value_type = std::pair<const key_type, mapped_type>;
        using pointer = typename std::allocator_traits<allocator_type>::pointer;
        using const_pointer = typename std::allocator_traits<allocator_type>::const_pointer;
        using reference = value_type&;
        using const_reference = const value_type&;
        using iterator = hash_map_iterator<value_type>;
        using const_iterator = hash_map_const_iterator<value_type>;
        using size_type = std::size_t;

    private:
        float m_max_load_factor = 0.5f;
        u64 m_size = 0;
        hasher m_hash;
        key_equal m_key_equal;
        storage<value_type, allocator_type> m_storage;
        std::vector<cell_state> m_states;

    public:
        hash_map() = default;

        explicit hash_map(size_type n) : m_storage(n), m_states(n, cell_state::non_init) { }

        explicit hash_map(const allocator_type& allocator) : m_storage(allocator) { }

        template<typename InputIterator>
        hash_map(InputIterator first, InputIterator last, size_type n = 0) : hash_map(n)
        {
            std::for_each(first, last, [this](const value_type& v){ insert(v); });
        }

        hash_map(std::initializer_list<value_type> l, size_type n = 0) : hash_map(l.begin(), l.end(), n) { }

        hash_map(const hash_map& other, const allocator_type& allocator) : m_max_load_factor(other.m_max_load_factor), m_size(other.m_size),
                                                                           m_hash(other.m_hash), m_key_equal(other.m_key_equal),
                                                                           m_storage(other.m_storage, allocator) { }

        hash_map(hash_map&& other, const allocator_type& allocator) : m_max_load_factor(other.m_max_load_factor), m_size(other.m_size),
                                                                      m_hash(std::move(other.m_hash)), m_key_equal(std::move(other.m_key_equal)),
                                                                      m_storage(std::move(other.m_storage), allocator) { }

        hash_map& operator=(std::initializer_list<value_type> l)
        {
            *this = hash_map<K, T, Hash, Pred, Alloc>(l);
        }

        allocator_type get_allocator() const noexcept
        {
            return m_storage.m_allocator;
        }

        bool empty() const noexcept
        {
            return !m_size;
        }

        size_type size() const noexcept
        {
            return m_size;
        }

        size_type max_size() const noexcept
        {
            return std::numeric_limits<size_type>::max();
        }

        iterator begin() noexcept
        {
            size_type index = 0;
            for (; index < m_storage.size(); ++index)
                if (m_states[index] == cell_state::presented)
                    break;
            return hash_map_iterator<value_type>(m_states, m_storage.data(), index);
        }

        const_iterator begin() const noexcept;

        const_iterator cbegin() const noexcept;

        iterator end() noexcept
        {
            return hash_map_iterator<value_type>(m_states, m_storage.data(), m_storage.size());
        }

        const_iterator end() const noexcept;

        const_iterator cend() const noexcept;

        template<typename... Args>
        std::pair<iterator, bool> emplace(Args&&... args)
        {
            static_assert (sizeof...(Args) == 2);
            const auto cell = get_key_cell(std::get<0>(std::tuple<Args&...>(args...)));
            const auto prev_state = m_states[cell];
            raw_insert_cell(cell, std::forward<Args>(args)...);
            return {create_iterator(cell), prev_state != cell_state::presented};
        }

        template <typename... Args>
        std::pair<iterator, bool> try_emplace(const key_type& k, Args&&... args)
        {
            return perform_try_emplace(k, std::forward<Args>(args)...);
        }

        template <typename... Args>
        std::pair<iterator, bool> try_emplace(key_type&& k, Args&&... args)
        {
            return perform_try_emplace(std::move(k), std::forward<Args>(args)...);
        }

        std::pair<iterator, bool> insert(const value_type& x)
        {
            const auto cell = raw_insert(x.first, x.second);
            return create_iterator(cell);
        }

        std::pair<iterator, bool> insert(value_type&& x)
        {
            const auto cell = raw_insert(std::move(x.first), std::move(x.second));
            return create_iterator(cell);
        }

        template<typename InputIterator>
        void insert(InputIterator first, InputIterator last)
        {
            for (; first != last; ++first)
                raw_insert(first->first, first->second);
        }

        void insert(std::initializer_list<value_type> l)
        {
            for (const auto& [k, v] : l)
                raw_insert(k, v);
        }

        template <typename Value>
        std::pair<iterator, bool> insert_or_assign(const key_type& k, Value&& v)
        {
            return perform_insert_or_assign(k, std::forward<Value>(v));
        }

        template <typename Value>
        std::pair<iterator, bool> insert_or_assign(key_type&& k, Value&& v)
        {
            return perform_insert_or_assign(std::move(k), std::forward<Value>(v));
        }

        template<typename K1, typename T1, typename Hash1, typename Pred1, typename Alloc1>
        friend bool operator==(const hash_map<K1, T1, Hash1, Pred1, Alloc1>&, const hash_map<K1, T1, Hash1, Pred1, Alloc1>&);

        Hash hash_function() const
        {
            return m_hash;
        }

        Pred key_eq() const
        {
            return m_key_equal;
        }

        size_type count(const key_type& k) const
        {
            return contains(k);
        }

        bool contains(const key_type& k) const
        {
            return m_states[get_key_cell(k)] == cell_state::presented;
        }

        mapped_type& operator[](const key_type& k)
        {
            return perform_brackets_operator(k);
        }

        mapped_type& operator[](key_type&& k)
        {
            return perform_brackets_operator(std::move(k));
        }

        const mapped_type& at(const key_type& k) const
        {
            auto cell = get_key_cell(k);
            if (m_states[cell] != cell_state::presented)
                throw std::out_of_range("Key is not presented in a hash_map");
            return get_value(cell);
        }

        mapped_type& at(const key_type& k)
        {
            return const_cast<mapped_type&>(static_cast<const hash_map&>(*this).at(k));
        }

        size_type bucket_count() const noexcept
        {
            return m_storage.size();
        }

        size_type bucket(const key_type& k) const
        {
            const auto cell = get_key_cell(k);
            if (m_states[cell] == cell_state::presented)
                return cell;
            return npos;
        }

        iterator erase(const_iterator position);

        iterator erase(iterator position);

        size_type erase(const key_type& x);

        iterator erase(const_iterator first, const_iterator last);

        void clear() noexcept;

        void swap(hash_map& x);

        template<typename _H2, typename _P2>
        void merge(hash_map<K, T, _H2, _P2, Alloc>& source);

        template<typename _H2, typename _P2>
        void merge(hash_map<K, T, _H2, _P2, Alloc>&& source);

        iterator find(const key_type& x);

        const_iterator find(const key_type& x) const;

        void rehash(size_type n)
        {
            if (n <= m_storage.size())
                return;
            storage<value_type, allocator_type> new_storage(m_storage.get_allocator(), n);
            std::vector<cell_state> new_states(n, cell_state::non_init);
            for (size_type i = 0; i < m_storage.size(); ++i)
            {
                if (m_states[i] != cell_state::presented)
                    continue;
                const auto new_cell = m_hash(get_key(i)) % n;
                new_storage.construct(new_cell, m_storage[i]);
                new_states[new_cell] = cell_state::presented;
            }
            m_states = std::move(new_states);
            m_storage = std::move(new_storage);
        }

        void reserve(size_type n)
        {
            // TODO problem with m_storage.size()
        }

        float load_factor() const noexcept
        {
            return static_cast<float>(m_storage.constructed_count()) / m_storage.size();
        }

        float max_load_factor() const noexcept
        {
            return m_max_load_factor;
        }

        void max_load_factor(float z)
        {
            m_max_load_factor = z;
            if (load_factor() <= m_max_load_factor)
                return;
            rehash(load_factor() / m_max_load_factor * 1.66);
        }

    private:
        /* std functions fwd analogs */

        template<typename Key, typename Value>
        std::pair<iterator, bool> perform_insert_or_assign(Key&& k, Value&& v)
        {
            const auto cell = get_key_cell(k);
            const auto prev_state = m_states[cell];
            raw_insert_cell(cell, std::forward<Key>(k), std::forward<Value>(v), true);
            return {create_iterator(cell), prev_state != cell_state::presented};
        }

        template <typename Key, typename... Args>
        std::pair<iterator, bool> perform_try_emplace(Key&& k, Args&&... args)
        {
            const auto cell = get_key_cell(k);
            if (m_states[cell] == cell_state::presented)
                return {create_iterator(cell), false};
            m_storage.destruct(cell);
            m_storage.construct(cell, std::piecewise_construct, std::forward_as_tuple(std::forward<Key>(k)), std::forward_as_tuple(std::forward<Args...>(args)...));
            m_states[cell] = cell_state::presented;
            return {create_iterator(cell), true};
        }

        template<typename Key>
        mapped_type& perform_brackets_operator(Key&& key)
        {
            const auto cell = get_key_cell(key);
            if (m_states[cell] == cell_state::non_init)
                raw_insert(std::forward<Key>(key), mapped_type());
            else if (m_states[cell] == cell_state::erased)
                raw_insert(std::forward<Key>(key), mapped_type(), true);
            return get_value(cell);
        }

        /* end */

        static constexpr size_type npos = std::numeric_limits<size_type>::max();

        size_type get_key_cell(const key_type& k) const
        {
            const auto index = m_hash(k) % m_storage.size();
            auto i = index;
            do
            {
                if (m_states[i] == cell_state::non_init)
                    return i;
                if (m_key_equal(k, get_key(i)))
                    return i;
                if (++i == m_storage.size())
                    i = 0;
            } while (i != index);
            return npos;
        }

        template<typename Key, typename Value>
        void raw_insert_cell(size_type cell, Key&& key, Value&& value, bool replace_if_exist = false)
        {
            if (m_states[cell] == cell_state::non_init)
                create_pair(cell, std::forward<Key>(key), std::forward<Value>(value));
            else if (m_states[cell] == cell_state::erased || replace_if_exist)
                replace_value(cell, std::forward<Value>(value));
        }

        template<typename Key, typename Value>
        size_type raw_insert(Key&& key, Value&& value, bool replace_if_exist = false)
        {
            const auto cell = get_key_cell(key);
            raw_insert_cell(cell, std::forward<Key>(key), std::forward<Value>(value), replace_if_exist);
            return cell;
        }

        template<typename Key, typename Value>
        void create_pair(size_type index, Key&& key, Value&& value)
        {
            m_states[index] = cell_state::presented;
            m_storage.construct(index, std::pair(std::forward<Key>(key), std::forward<Value>(value)));
        }

        template<typename Value>
        void replace_value(size_type index, Value&& value)
        {
            m_states[index] = cell_state::presented;
            m_storage[index].second = std::forward<Value>(value);
        }

        void erase_index(size_type index)
        {
            m_states[index] = cell_state::erased;
        }

        const key_type& get_key(size_type index) const
        {
            return m_storage[index].first;
        }

        mapped_type& get_value(size_type index)
        {
            return m_storage[index].second;
        }

        iterator create_iterator(size_type index)
        {
            return hash_map_iterator(m_states, m_storage.data(), index);
        }
    };
}
// namespace hm