//
// Created by Dan Crosby on 11/28/22.
//

#ifndef EAGLE0_ITERABLEBITSET_HPP
#define EAGLE0_ITERABLEBITSET_HPP

#include <algorithm>
#include <bit>
#include <cstdlib>

#define ITERABLE_BITSET_INDEX_CHECKS false

namespace common {

template<size_t N>
class IterableBitset {
private:
    using STORE_TYPE = uint64_t;
    constexpr static int STORE_BITS = sizeof(STORE_TYPE) * 8;
    constexpr static int BIT_SHIFT = std::countr_zero(sizeof(STORE_TYPE) * 8);
    constexpr static int BIT_MASK = STORE_BITS - 1;
    constexpr static int STORE_COUNT = (N + STORE_BITS - 1) / STORE_BITS;

    STORE_TYPE store[STORE_COUNT] = {};

public:
    IterableBitset() = default;

    IterableBitset(const IterableBitset& rhs) = default;
    IterableBitset(IterableBitset&& rhs) noexcept = default;

    ~IterableBitset() = default;

    void clear() { std::fill(std::begin(store), std::end(store), 0); }

    auto operator=(const IterableBitset& rhs) -> IterableBitset& = default;
    auto operator=(IterableBitset&& rhs) noexcept -> IterableBitset& = default;

    auto operator==(const IterableBitset& rhs) const -> bool {
        for (int i = 0; i < STORE_COUNT; i++) {
            if (store[i] != rhs.store[i]) return false;
        }
        return true;
    };

    void Set(size_t idx) {
#if ITERABLE_BITSET_INDEX_CHECKS
        assert(idx < N);
#endif
        store[idx >> BIT_SHIFT] |= 1UL << (idx & BIT_MASK);
    }

    void Unset(size_t idx) {
#if ITERABLE_BITSET_INDEX_CHECKS
        assert(idx < N);
#endif
        store[idx >> BIT_SHIFT] &= ~(1UL << (idx & BIT_MASK));
    }

    [[nodiscard]] auto IsSet(size_t idx) const -> bool {
#if ITERABLE_BITSET_INDEX_CHECKS
        assert(idx < N);
#endif
        return (store[idx >> BIT_SHIFT] >> (idx & BIT_MASK)) & 1U;
    }

    [[nodiscard]] auto size() const -> size_t {
        size_t sz = 0;
        for (size_t i = 0; i < STORE_COUNT; i++) { sz += std::popcount(store[i]); }

        return sz;
    }

    [[nodiscard]] auto empty() const -> bool {
        for (size_t i = 0; i < STORE_COUNT; i++) {
            if (store[i]) return false;
        }
        return true;
    }

    auto operator+=(const IterableBitset<N>& rhs) -> IterableBitset<N>& {
        for (int i = 0; i < STORE_COUNT; i++) { store[i] |= rhs.store[i]; }
        return *this;
    }

    auto operator+(const IterableBitset<N>& rhs) const -> IterableBitset<N> {
        IterableBitset<N> ib = *this;
        ib += rhs;
        return ib;
    }

    auto operator-=(const IterableBitset<N>& rhs) -> IterableBitset<N>& {
        for (int i = 0; i < STORE_COUNT; i++) { store[i] &= ~(rhs.store[i]); }
        return *this;
    }

    auto operator-(const IterableBitset<N>& rhs) const -> IterableBitset<N> {
        IterableBitset<N> ib = *this;
        ib -= rhs;
        return ib;
    }

    class Iter {
    public:
        using value_type = size_t;
        using pointer = size_t*;
        using reference = size_t&;

    private:
        size_t storeIndex;
        size_t bitIndex;

        const IterableBitset& bs;

        void ToNext() {
            do {
                const STORE_TYPE current = bs.store[storeIndex] >> bitIndex;
                if (current != 0) {
                    // Skip ahead by the number of trailing zeroes
                    const auto ctz = std::countr_zero(current);
                    bitIndex += ctz;
                    return;
                }
                bitIndex = 0;
                storeIndex++;
            } while (storeIndex < STORE_COUNT);
        }

    public:
        Iter(const IterableBitset& bs, const int startStore, const int startBit)
            : storeIndex(startStore),
              bitIndex(startBit),
              bs(bs) {
            ToNext();
        }

        auto operator==(const Iter& rhs) const -> bool {
            return storeIndex == rhs.storeIndex && bitIndex == rhs.bitIndex;
        }
        auto operator!=(const Iter& rhs) const -> bool {
            return storeIndex != rhs.storeIndex || bitIndex != rhs.bitIndex;
        }
        void operator++() {
            bitIndex++;
            if (bitIndex >= STORE_BITS) {
                storeIndex++;
                bitIndex = 0;
            }
            ToNext();
        }

        auto operator*() const -> size_t { return (storeIndex << BIT_SHIFT) + bitIndex; }
    };

    [[nodiscard]] constexpr auto begin() const -> Iter { return Iter(*this, 0, 0); }
    [[nodiscard]] constexpr auto end() const -> Iter {
        return Iter(*this, N >> BIT_SHIFT, N & BIT_MASK);
    }
};

}  // namespace common

#endif  // EAGLE0_ITERABLEBITSET_HPP
