#pragma once

#include <stdexcept>

#include "sequences/pef_global_parameters.hpp"
#include "utils/util.hpp"
#include "vectors/bit_vector.hpp"

namespace pef {

struct compact_elias_fano {
    struct offsets {
        offsets() {}

        offsets(uint64_t base_offset, uint64_t universe, uint64_t n,
                pef_global_parameters const& params)
            : universe(universe)
            , n(n)
            , log_sampling0(params.ef_log_sampling0)
            , log_sampling1(params.ef_log_sampling1)

            , lower_bits(universe > n ? tongrams::util::msb(universe / n) : 0)
            , mask((uint64_t(1) << lower_bits) - 1)
            // pad with a zero on both sides as sentinels
            , higher_bits_length(n + (universe >> lower_bits) + 2)
            , pointer_size(tongrams::util::ceil_log2(higher_bits_length))
            , pointers0((higher_bits_length - n) >> log_sampling0)  // XXX
            , pointers1(n >> log_sampling1)

            , pointers0_offset(base_offset)
            , pointers1_offset(pointers0_offset + pointers0 * pointer_size)
            , higher_bits_offset(pointers1_offset + pointers1 * pointer_size)
            , lower_bits_offset(higher_bits_offset + higher_bits_length)
            , end(lower_bits_offset + n * lower_bits) {
            assert(n > 0);
        }

        uint64_t universe;
        uint64_t n;
        uint64_t log_sampling0;
        uint64_t log_sampling1;

        uint64_t lower_bits;
        uint64_t mask;
        uint64_t higher_bits_length;
        uint64_t pointer_size;
        uint64_t pointers0;
        uint64_t pointers1;

        uint64_t pointers0_offset;
        uint64_t pointers1_offset;
        uint64_t higher_bits_offset;
        uint64_t lower_bits_offset;
        uint64_t end;
    };

    static uint64_t bitsize(pef_global_parameters const& params,
                            uint64_t universe, uint64_t n) {
        return offsets(0, universe, n, params).end;
    }

    template <typename Iterator>
    static void write(tongrams::bit_vector_builder& bvb, Iterator begin,
                      uint64_t universe, uint64_t n,
                      pef_global_parameters const& params) {
        using tongrams::util::ceil_div;
        uint64_t base_offset = bvb.size();
        offsets of(base_offset, universe, n, params);
        // initialize all the bits to 0
        bvb.zero_extend(of.end - base_offset);

        uint64_t sample1_mask = (uint64_t(1) << of.log_sampling1) - 1;
        uint64_t offset;

        // utility function to set 0 pointers
        auto set_ptr0s = [&](uint64_t begin, uint64_t end, uint64_t rank_end) {
            uint64_t begin_zeros = begin - rank_end;
            uint64_t end_zeros = end - rank_end;

            for (uint64_t ptr0 =
                     ceil_div(begin_zeros, uint64_t(1) << of.log_sampling0);
                 (ptr0 << of.log_sampling0) < end_zeros; ++ptr0) {
                if (!ptr0)
                    continue;
                offset = of.pointers0_offset + (ptr0 - 1) * of.pointer_size;
                assert(offset + of.pointer_size <= of.pointers1_offset);
                bvb.set_bits(offset, (ptr0 << of.log_sampling0) + rank_end,
                             of.pointer_size);
            }
        };

        uint64_t last = 0;
        uint64_t last_high = 0;
        Iterator it = begin;
        for (size_t i = 0; i < n; ++i) {
            uint64_t v = *it++;
            if (i && v < last) {
                std::cout << "at position " << i << "/" << n << std::endl;
                std::cout << v << " < " << last << std::endl;
                throw std::runtime_error("sequence is not sorted");
            }
            assert(v <= universe);

            uint64_t high = (v >> of.lower_bits) + i + 1;
            uint64_t low = v & of.mask;

            bvb.set(of.higher_bits_offset + high, 1);

            offset = of.lower_bits_offset + i * of.lower_bits;
            assert(offset + of.lower_bits <= of.end);
            bvb.set_bits(offset, low, of.lower_bits);

            if (i && (i & sample1_mask) == 0) {
                uint64_t ptr1 = i >> of.log_sampling1;
                assert(ptr1 > 0);
                offset = of.pointers1_offset + (ptr1 - 1) * of.pointer_size;
                assert(offset + of.pointer_size <= of.higher_bits_offset);
                bvb.set_bits(offset, high, of.pointer_size);
            }

            // write pointers for the run of zeros in [last_high, high)
            set_ptr0s(last_high + 1, high, i);
            last_high = high;
            last = v;
        }

        // pointers to zeros after the last 1
        set_ptr0s(last_high + 1, of.higher_bits_length, n);  // XXX
    }

    struct enumerator {
        typedef std::pair<uint64_t, uint64_t> value_type;  // (position, value)

        enumerator() {}

        enumerator(tongrams::bit_vector const& bv, uint64_t offset,
                   uint64_t universe, uint64_t n,
                   pef_global_parameters const& params)
            : m_bv(&bv)
            , m_of(offset, universe, n, params)
            , m_position(size())
            , m_value(m_of.universe) {}

        value_type move(uint64_t position) {
            assert(position <= m_of.n);
            if (position == m_position) {
                return value();
            }

            uint64_t skip = position - m_position;
            // optimize small forward skips
            if (LIKELY(position > m_position &&
                       skip <= linear_scan_threshold)) {
                m_position = position;
                if (UNLIKELY(m_position == size())) {
                    m_value = m_of.universe;
                } else {
                    tongrams::bit_vector::unary_iterator he = m_high_enumerator;
                    for (size_t i = 0; i < skip; ++i) {
                        he.next();
                    }
                    m_value = ((he.position() - m_of.higher_bits_offset -
                                m_position - 1)
                               << m_of.lower_bits) |
                              read_low();
                    m_high_enumerator = he;
                }
                return value();
            }

            return slow_move(position);
        }

        value_type next_geq(uint64_t lower_bound) {
            if (lower_bound == m_value) {
                return value();
            }

            uint64_t high_lower_bound = lower_bound >> m_of.lower_bits;
            uint64_t cur_high = m_value >> m_of.lower_bits;
            uint64_t high_diff = high_lower_bound - cur_high;

            if (LIKELY(lower_bound > m_value &&
                       high_diff <= linear_scan_threshold)) {
                // optimize small skips
                next_reader next_value(*this, m_position + 1);
                uint64_t val;
                do {
                    m_position += 1;
                    if (LIKELY(m_position < size())) {
                        val = next_value();
                    } else {
                        val = m_of.universe;
                        break;
                    }
                } while (val < lower_bound);

                m_value = val;
                return value();
            } else {
                return slow_next_geq(lower_bound);
            }
        }

        uint64_t size() const {
            return m_of.n;
        }

        value_type next() {
            m_position += 1;
            assert(m_position <= size());
            if (LIKELY(m_position < size())) {
                m_value = read_next();
            } else {
                m_value = m_of.universe;
            }
            return value();
        }

        uint64_t prev_value() const {
            if (m_position == 0) {
                return 0;
            }

            uint64_t prev_high = 0;
            if (LIKELY(m_position < size())) {
                prev_high =
                    m_bv->predecessor1(m_high_enumerator.position() - 1);
            } else {
                prev_high = m_bv->predecessor1(m_of.lower_bits_offset - 1);
            }
            prev_high -= m_of.higher_bits_offset;

            uint64_t prev_pos = m_position - 1;
            uint64_t prev_low = m_bv->get_word56(m_of.lower_bits_offset +
                                                 prev_pos * m_of.lower_bits) &
                                m_of.mask;
            return ((prev_high - prev_pos - 1) << m_of.lower_bits) | prev_low;
        }

        uint64_t position() const {
            return m_position;
        }

    private:
        value_type NOINLINE slow_move(uint64_t position) {
            if (UNLIKELY(position == size())) {
                m_position = position;
                m_value = m_of.universe;
                return value();
            }

            uint64_t skip = position - m_position;
            uint64_t to_skip;
            if (position > m_position && (skip >> m_of.log_sampling1) == 0) {
                to_skip = skip - 1;
            } else {
                uint64_t ptr = position >> m_of.log_sampling1;
                uint64_t high_pos = pointer1(ptr);
                uint64_t high_rank = ptr << m_of.log_sampling1;
                m_high_enumerator = tongrams::bit_vector::unary_iterator(
                    *m_bv, m_of.higher_bits_offset + high_pos);
                to_skip = position - high_rank;
            }

            m_high_enumerator.skip(to_skip);
            m_position = position;
            m_value = read_next();
            return value();
        }

        value_type NOINLINE slow_next_geq(uint64_t lower_bound) {
            if (UNLIKELY(lower_bound >= m_of.universe)) {
                return move(size());
            }

            uint64_t high_lower_bound = lower_bound >> m_of.lower_bits;
            uint64_t cur_high = m_value >> m_of.lower_bits;
            uint64_t high_diff = high_lower_bound - cur_high;

            // XXX bounds checking!
            uint64_t to_skip;
            if (lower_bound > m_value &&
                (high_diff >> m_of.log_sampling0) == 0) {
                // note: at the current position in the bitvector there
                // should be a 1, but since we already consumed it, it
                // is 0 in the enumerator, so we need to skip it
                to_skip = high_diff;
            } else {
                uint64_t ptr = high_lower_bound >> m_of.log_sampling0;
                uint64_t high_pos = pointer0(ptr);
                uint64_t high_rank0 = ptr << m_of.log_sampling0;

                m_high_enumerator = tongrams::bit_vector::unary_iterator(
                    *m_bv, m_of.higher_bits_offset + high_pos);
                to_skip = high_lower_bound - high_rank0;
            }

            m_high_enumerator.skip0(to_skip);
            m_position = m_high_enumerator.position() -
                         m_of.higher_bits_offset - high_lower_bound;

            next_reader read_value(*this, m_position);
            while (true) {
                if (UNLIKELY(m_position == size())) {
                    m_value = m_of.universe;
                    return value();
                }
                auto val = read_value();
                if (val >= lower_bound) {
                    m_value = val;
                    return value();
                }
                m_position++;
            }
        }

        static const uint64_t linear_scan_threshold = 8;

        inline value_type value() const {
            return value_type(m_position, m_value);
        }

        inline uint64_t read_low() {
            return m_bv->get_word56(m_of.lower_bits_offset +
                                    m_position * m_of.lower_bits) &
                   m_of.mask;
        }

        inline uint64_t read_next() {
            assert(m_position < size());
            uint64_t high = m_high_enumerator.next() - m_of.higher_bits_offset;
            return ((high - m_position - 1) << m_of.lower_bits) | read_low();
        }

        struct next_reader {
            next_reader(enumerator& e, uint64_t position)
                : e(e)
                , high_enumerator(e.m_high_enumerator)
                , high_base(e.m_of.higher_bits_offset + position + 1)
                , lower_bits(e.m_of.lower_bits)
                , lower_base(e.m_of.lower_bits_offset + position * lower_bits)
                , mask(e.m_of.mask)
                , bv(*e.m_bv) {}

            ~next_reader() {
                e.m_high_enumerator = high_enumerator;
            }

            uint64_t operator()() {
                uint64_t high = high_enumerator.next() - high_base;
                uint64_t low = bv.get_word56(lower_base) & mask;
                high_base += 1;
                lower_base += lower_bits;
                return (high << lower_bits) | low;
            }

            enumerator& e;
            tongrams::bit_vector::unary_iterator high_enumerator;
            uint64_t high_base, lower_bits, lower_base, mask;
            tongrams::bit_vector const& bv;
        };

        inline uint64_t pointer(uint64_t offset, uint64_t i) const {
            if (i == 0) {
                return 0;
            } else {
                return m_bv->get_word56(offset + (i - 1) * m_of.pointer_size) &
                       ((uint64_t(1) << m_of.pointer_size) - 1);
            }
        }

        inline uint64_t pointer0(uint64_t i) const {
            return pointer(m_of.pointers0_offset, i);
        }

        inline uint64_t pointer1(uint64_t i) const {
            return pointer(m_of.pointers1_offset, i);
        }

        tongrams::bit_vector const* m_bv;
        offsets m_of;

        uint64_t m_position;
        uint64_t m_value;
        tongrams::bit_vector::unary_iterator m_high_enumerator;
    };
};
}  // namespace pef
