#include "sorted_uint_vec.hpp"
#include <terark/bitmap.hpp>
#include <terark/num_to_str.hpp>
#include <terark/io/FileStream.hpp>
#include <terark/io/StreamBuffer.hpp>
#include <terark/util/function.hpp>
#include <boost/intrusive_ptr.hpp>

namespace terark {

static inline size_t safe_bsr_u64(uint64_t val) {
    return val ? terark_bsr_u64(val) : 0;
}

static inline uint64_t encode_signed(int64_t x) {
	if (x < 0)
		return ((-x-1) << 1) | 1;
	else
		return (+x << 1);
}

static inline int64_t decode_signed(uint64_t x) {
	if (x & 1)
		return -int64_t(x >> 1) - 1;
	else
		return +int64_t(x >> 1);
}

template<class Uint>
static inline int64_t
varbits_get_signed(const Uint* base, size_t bitpos, size_t width) {
	Uint val = febitvec::s_get_uint(base, bitpos, width);
	return decode_signed(val);
}

/*
template<class Uint>
static inline void
varbits_set_signed(Uint* base, size_t bitpos, size_t width, int64_t sval) {
	BOOST_STATIC_ASSERT(sizeof(Uint) >= 8);
	Uint val = encode_signed(sval);
	febitvec::s_set_uint(base, bitpos, width, val);
}
*/

// Lagrange polynomial interpolation:
static inline
int64_t lagrange(int64_t x, int64_t a, int64_t f_a
                          , int64_t m, int64_t f_m
                          , int64_t b, int64_t f_b
                          , size_t log2_range) {
    assert(0 == a);
    assert(m == b/2);
    assert(b == int64_t(1) << log2_range);
    assert(x != 0);
    assert(x != m);
    typedef long double RealNum;
#if 0
    int64_t f_x = int64_t(0
          + f_a*RealNum((x-m)*(x-b))/((a-m)*(a-b))
          + f_m*RealNum((x-a)*(x-b))/((m-a)*(m-b))
          + f_b*RealNum((x-a)*(x-m))/((b-a)*(b-m))
        );
#else
    int64_t one = 1;
//  int64_t denominator;
    int64_t deno_a = +(one << (log2_range*2 - 1));
//  int64_t deno_m = -(one << (log2_range*2 - 2));
    int64_t deno_m = -(deno_a >> 1);
//  int64_t deno_b = +(one << (log2_range*2 - 1));
    int64_t deno_b = deno_a;
    int64_t f_x = int64_t(0
          + f_a*RealNum((x-m)*(x-b))/deno_a
          + f_m*RealNum((x-a)*(x-b))/deno_m
          + f_b*RealNum((x-a)*(x-m))/deno_b
        );
#endif
    return f_x;
};

static inline uint32_t GetDDWidthType(const byte_t* ba) {
	return ba[0] & 0x0F;
}

// for other diffdiffWidthType
static inline uint64_t GetLoWater_x(const byte_t* ba, size_t* headerLen) {
	byte_t type = (ba[0] >> 4) & 3;
	uint64_t val = ba[0] >> 6; // get low 2 bits
	switch (type) {
	case 0: // total 10 bits, max 1023
		*headerLen = 2;
		val |= uint64_t(ba[1]) << 2;
		break;
	case 1: // total 18 bits, max 256K - 1
		*headerLen = 3;
		val |= uint64_t(unaligned_load<uint16_t>(ba+1)) << 2;
		break;
	case 2: // total 26 bits, max  64M -1
		*headerLen = 4;
		val |= uint64_t(unaligned_load<uint32_t>(ba+1) & 0x00FFFFFF) << 2;
		break;
	case 3: // total 50 bits, max   1P -1
		*headerLen = 7;
		val |= uint64_t(unaligned_load<uint64_t>(ba+1) & ~(uint64_t(-1LL)<<48)) << 2;
		break;
	}
	return val;
}

static const byte_t g_used_small_width[65] = {0,1,2,3,4,5,6,7,8,9,10
	,  12 // width = 11, type = NA
	,  12 // width = 12, type = 12 +++
	,  16 // width = 13, type = NA
	,  16 // width = 14, type = NA
	,  16 // width = 15, type = NA
	,  16 // width = 16, type = 13 +++
	,  20 // width = 17, type = NA
	,  20 // width = 18, type = NA
	,  20 // width = 19, type = NA
	,  20 // width = 20, type = 14 +++
	      // 21 numbers above --------------
	,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
	,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
	,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
};

static size_t getWidthType(size_t smallWidth, size_t largeWidth) {
	// type  width
	// 11 ->  1 bit, no large units
	// 13 -> 16
	// 14 -> 20
	// 15 ->  2 bits, no large units
	static const byte_t widthToType[21] = {0,1,2,3,4,5,6,7,8,9,10
		, 255 // width = 11, type = NA
		,  12 // width = 12, type = 12 +++
		, 255 // width = 13, type = NA
		, 255 // width = 14, type = NA
		, 255 // width = 15, type = NA
		,  13 // width = 16, type = 13 +++
		, 255 // width = 17, type = NA
		, 255 // width = 18, type = NA
		, 255 // width = 19, type = NA
		,  14 // width = 20, type = 14 +++
	};
	assert(smallWidth <= 20);
	assert(smallWidth <= largeWidth);
	uint32_t widthType;
	if (1 == smallWidth && 1 == largeWidth) {
		widthType = 11;
	}
	else if (2 == smallWidth && 2 == largeWidth) {
		widthType = 15;
	}
	else {
		widthType = widthToType[smallWidth];
		assert(widthType < 16);
	}
	return widthType;
}

static inline
void WriteHeader(valvec<byte_t>* data, uint64_t loWater,
				 size_t smallWidth, size_t largeWidth) {
	const size_t widthType = getWidthType(smallWidth, largeWidth);

	// +-------------+-------------+---------------------------+
	// |  loWaterLo  |  loWaterHi  |        widthType          |
	// |    2 bits   | Length Type |   Small Unit Width Type   |
	// +-------------+-------------+---------------------------+
	// |   7  |   6  |   5  |   4  |   3  |   2  |   1  |   0  |
	// +-------------+-------------+---------------------------+

	const uint08_t loWaterLo = uint08_t(loWater &  3) << 6;
	const uint64_t loWaterHi = uint64_t(loWater >> 2);
	const uint64_t one = 1;

// LoWaterType use 2 bits, it map to how many bytes of 'loWaterHi' requires
#define WriteHeaderData(LoWaterType, HiBytes) \
	static_assert(LoWaterType >= 0 && LoWaterType < 4, "Bad LoWaterType"); \
	data->append(uint08_t(loWaterLo | (LoWaterType << 4) | widthType)); \
	data->append((const byte*)(&loWaterHi), HiBytes)

	if (false) {}
	else if (loWaterHi < one <<  8) { WriteHeaderData(0, 1); }
	else if (loWaterHi < one << 16) { WriteHeaderData(1, 2); }
	else if (loWaterHi < one << 24) { WriteHeaderData(2, 3); }
	else if (loWaterHi < one << 48) { WriteHeaderData(3, 6); }
	else {
		abort(); // should not happen
	}
}

static const uint64_t
// 21*3 = 63, high 1 bit into overflowEx, seperate next high 3 bits
Zc01_21_03 = 0b0001001001001001001001001001001001001001001001001001001001001001,
Bits_20_03 = 0b0000000111000111000111000111000111000111000111000111000111000111,
Bits_10_06 = 0b0000000000111111000000111111000000111111000000111111000000111111,
Bits_05_01 = 0b0000000000000001000000000001000000000001000000000001000000000001,

Zc03_10_06 = 0b0000000111000111000111000111000111000111000111000111000111000111,
Zc01_10_06 = 0b0000000001000001000001000001000001000001000001000001000001000001,

Zc02_12_05 = 0b0000000110001100011000110001100011000110001100011000110001100011,
Zc01_12_05 = 0b0000000010000100001000010000100001000010000100001000010000100001,
Bits_12_05 = 0b0000000001111100000111110000011111000001111100000111110000011111,
Bits_06_01 = 0b0000000000000100000000010000000001000000000100000000010000000001,

// Zc03 -- 3 bits of 1 in one 7 bits group
Zc03_09_07 = 0b0000011100001110000111000011100001110000111000011100001110000111,
Zc01_09_07 = 0b0000000100000010000001000000100000010000001000000100000010000001,
Bits_08_07 = 0b0000000000000001111111000000011111110000000111111100000001111111,
Bits_04_14 = 0b0000000000000000000000111111111111110000000000000011111111111111,

// 7*9 = 63, high 1 bit into overflowEx, seperate next high 7 bits
Zc03_07_09 = 0b0000000111000000111000000111000000111000000111000000111000000111,
Zc01_07_09 = 0b0000000001000000001000000001000000001000000001000000001000000001,
Bits_06_09 = 0b0000000000000000000111111111000000000111111111000000000111111111,
Bits_03_18 = 0b0000000000000000000000000001000000000000000001000000000000000001,

Bits_06_10 = 0b0000000000000011111111110000000000111111111100000000001111111111;

static inline size_t s_get_mask(size_t bits) {
	if (sizeof(size_t)*8 == bits)
		return size_t(-1);
	else
		return ~(size_t(-1) << bits);
}

#define WordUnits      (64/Width)
#define UnitMask       ( (size_t(+1) <<    Width) - 1 )
#define HiExMaskToLow  ( (size_t(+1) << 64%Width) - 1 )
#define HiExMaskAtHig   ~(size_t(-1) >> 64%Width)

struct SortedUintVec::ObjectHeader {
	uint64_t  units : 48; // SortedUintVec::m_size
	uint64_t  log2_blockUnits : 4; // 6 or 7
	uint64_t  offsetWidth_1   : 6; // 1 ~ 64, = real_offset_width - 1
	uint64_t  sampleWidth_1   : 6; // 1 ~ 64, = real_sample_width - 1
	uint64_t  indexOffset     :48;
	uint64_t  is_overall_full_sorted: 1;
	uint64_t  is_samples_full_sorted: 1;
	uint64_t  reserved : 14;
};

void SortedUintVec::get2(size_t idx, size_t aVal[2]) const {
    static_assert(sizeof(ObjectHeader)==16, "sizeof(ObjectHeader) must be 16");
    assert(m_is_sorted_uint_vec);
    assert(idx < m_size);
    size_t blockIdx = idx >> m_log2_blockUnits;
    size_t blockUnits = size_t(1) << m_log2_blockUnits;
    size_t indexWidth = m_offsetWidth + m_sampleWidth;
    size_t indexBitPos = indexWidth * blockIdx;
    auto   indexBase = (const size_t*)(m_index);
    size_t offset0 = febitvec::s_get_uint(indexBase, indexBitPos, m_offsetWidth); indexBitPos += m_offsetWidth;
    size_t sample0 = febitvec::s_get_uint(indexBase, indexBitPos, m_sampleWidth);
    indexBitPos += m_sampleWidth;
    size_t offset1 = febitvec::s_get_uint(indexBase, indexBitPos, m_offsetWidth); indexBitPos += m_offsetWidth;
    size_t sample1 = febitvec::s_get_uint(indexBase, indexBitPos, m_sampleWidth);
    size_t subIdx = idx & (blockUnits-1);

    // pLargeBase is always aligned at size_t, faster for getting large units.
    // 'pData' may not aligned at size_t,
    // 'pData' is used for read multiple small units by size_t
    auto   pLargeBase = (size_t const*)(m_data.data() + sizeof(ObjectHeader));

    auto   header = m_data.data() + sizeof(ObjectHeader) + offset0;
    assert(offset0 + 2 <= offset1);
    assert(sample0 <= sample1 || !m_is_samples_full_sorted);

    switch (GetDDWidthType(header)) {
    case 0: { // bits = 0, all diffdiff are zero
        assert(offset0 + 2 <= offset1);
        size_t headerLen;
        size_t loWater = GetLoWater_x(header, &headerLen);
        aVal[0] = sample0 + loWater * subIdx;
        if (subIdx + 1 < blockUnits) {
            aVal[1] = aVal[0] + loWater;
        } else {
            aVal[1] = sample1;
        }
        break; }
    case 1:
#define Width 1
#include "sorted_uint_vec_case.hpp"
    case 2:
#define Width 2
#include "sorted_uint_vec_case.hpp"
    case 3:
#define Width 3
#include "sorted_uint_vec_case.hpp"
    case 4:
#define Width 4
#include "sorted_uint_vec_case.hpp"
    case 5:
#define Width 5
#include "sorted_uint_vec_case.hpp"
    case 6:
#define Width 6
#include "sorted_uint_vec_case.hpp"
    case 7:
#define Width 7
#include "sorted_uint_vec_case.hpp"
    case 8:
#define Width 8
#include "sorted_uint_vec_case.hpp"
    case 9:
#define Width 9
#include "sorted_uint_vec_case.hpp"
    case 10:
#define Width 10
#include "sorted_uint_vec_case.hpp"
    case 11: {
        // smallWidth and largeWidth are both 1, this is an optimization
        size_t headerLen;
        size_t const loWater = GetLoWater_x(header, &headerLen);
        size_t const* pData = (const size_t*)(uintptr_t(header) + headerLen);
        size_t i = 0, ddsum = 0;
        for (; i < subIdx / 64; ++i) {
            ddsum += fast_popcount64(pData[i]);
        }
        uint64_t w = pData[i];
        ddsum += fast_popcount_trail(w, subIdx % 64);
        aVal[0] = sample0 + ddsum + loWater * subIdx;
        if (subIdx + 1 < blockUnits)
            aVal[1] = aVal[0] + loWater + ((w >> subIdx % 64) & 1);
        else
            aVal[1] = sample1;
        break; }
    case 12:
#define Width 12
#include "sorted_uint_vec_case.hpp"
    case 13:
#define Width 16
#include "sorted_uint_vec_case.hpp"
    case 14:
#define Width 20
#include "sorted_uint_vec_case.hpp"
    case 15: {
        // smallWidth and largeWidth are both 2, this is an optimization
        size_t headerLen;
        size_t const loWater = GetLoWater_x(header, &headerLen);
        byte_t const* pData = (const byte_t*)(uintptr_t(header) + headerLen);
        size_t const old_style_block_size = headerLen + blockUnits * 2 / 8;
        if (offset1 - offset0 > old_style_block_size) {
            size_t fixWidth = (pData[0] & 63) + 1;
            size_t baseBitPos = (headerLen + offset0) * 8 + 6;
            size_t a = 0;
            size_t m = blockUnits/2;
            size_t b = blockUnits;
            size_t log2_range = m_log2_blockUnits;
            int64_t headerVal = loWater;
            int64_t f_a = 0;
            int64_t f_b = sample1 - sample0;
            int64_t f_m = (f_a+f_b)/2 + decode_signed(headerVal);
            auto computeVal = [=](size_t x) {
                assert(x > 0);
                uint64_t val;
                if (terark_likely(blockUnits/2 != x)) {
                    size_t x_idx = (x < blockUnits/2) ? (x-1) : (x-2);
                    size_t bitpos = baseBitPos + fixWidth*x_idx;
                    int64_t diff = varbits_get_signed(pLargeBase, bitpos, fixWidth);
                    int64_t f_x_lag = lagrange(x, a, f_a, m, f_m, b, f_b, log2_range);
                    val = f_x_lag + diff;
                //  printf("read: x = %3zd , f_x_lag = %3zd , diff = %3zd\n", x, f_x_lag, diff);
                } else {
                    val = f_m;
                }
                return sample0 + val;
            };
            if (terark_likely(subIdx)) {
                aVal[0] = computeVal(subIdx);
            } else {
                aVal[0] = sample0;
            }
            if (terark_likely(subIdx + 1 < blockUnits)) {
                aVal[1] = computeVal(subIdx + 1);
            } else {
                aVal[1] = sample1;
            }
            break; // break the case
        }
        size_t i = 0, ddsum = 0;
        for (; i < subIdx / 32; ++i) {
            uint64_t w = unaligned_load<uint64_t>(pData, i);
            w = (w & 0x3333333333333333ull) + ((w >> 2) & 0x3333333333333333ull);
            w = (w & 0x0F0F0F0F0F0F0F0Full) + ((w >> 4) & 0x0F0F0F0F0F0F0F0Full);
            w = (w * 0x0101010101010101ull) >> 56;
            ddsum += w;
        }
        uint64_t w = unaligned_load<uint64_t>(pData, i);
        size_t shiftBits = (subIdx % 32) * 2;
        if (terark_likely(subIdx % 32)) {
            uint64_t x = w & ~(size_t(-1) << shiftBits);
            x = (x & 0x3333333333333333ull) + ((x >> 2) & 0x3333333333333333ull);
            x = (x & 0x0F0F0F0F0F0F0F0Full) + ((x >> 4) & 0x0F0F0F0F0F0F0F0Full);
            x = (x * 0x0101010101010101ull) >> 56;
            ddsum += x;
        }
        aVal[0] = sample0 + ddsum + loWater * subIdx;
        if (subIdx + 1 < blockUnits)
            aVal[1] = aVal[0] + loWater + ((w >> shiftBits) & 3);
        else
            aVal[1] = sample1;
        break; }
    } // switch
}

/// aVal capacity must be at least blockUnits(64 or 128)
void SortedUintVec::get_block(size_t blockIdx, size_t* aVals) const {
    assert(blockIdx << m_log2_blockUnits < m_size + (size_t(1) << m_log2_blockUnits) - 1);
    assert(m_is_sorted_uint_vec);
    size_t blockUnits = size_t(1) << m_log2_blockUnits;
    size_t indexWidth = m_offsetWidth + m_sampleWidth;
    size_t indexBitPos = indexWidth * blockIdx;
    auto   indexBase = (const size_t*)(m_index);
    size_t offset0 = febitvec::s_get_uint(indexBase, indexBitPos, m_offsetWidth); indexBitPos += m_offsetWidth;
    size_t sample0 = febitvec::s_get_uint(indexBase, indexBitPos, m_sampleWidth);
    indexBitPos += m_sampleWidth;
    size_t offset1 = febitvec::s_get_uint(indexBase, indexBitPos, m_offsetWidth); indexBitPos += m_offsetWidth;
    size_t sample1 = febitvec::s_get_uint(indexBase, indexBitPos, m_sampleWidth);
    // pLargeBase is always aligned at size_t, faster for getting large units.
    // 'pData' may not aligned at size_t,
    // 'pData' is used for read multiple small units by size_t
    auto   pLargeBase = (size_t const*)(m_data.data() + sizeof(ObjectHeader));

    auto   header = m_data.data() + sizeof(ObjectHeader) + offset0;
    assert(offset0 + 2 <= offset1);
    assert(sample0 <= sample1 || !m_is_samples_full_sorted);

    switch (GetDDWidthType(header)) {
    case 0: { // bits = 0, all diffdiff are zero
        assert(offset0 + 2 <= offset1);
        size_t headerLen;
        size_t loWater = GetLoWater_x(header, &headerLen);
        size_t val = sample0;
        for (size_t i = 0; i < blockUnits; ++i) {
            aVals[i] = val;
            val += loWater;
        }
        break; }
    case 1:
#define Width 1
#include "sorted_uint_vec_get_block_case.hpp"
    case 2:
#define Width 2
#include "sorted_uint_vec_get_block_case.hpp"
    case 3:
#define Width 3
#include "sorted_uint_vec_get_block_case.hpp"
    case 4:
#define Width 4
#include "sorted_uint_vec_get_block_case.hpp"
    case 5:
#define Width 5
#include "sorted_uint_vec_get_block_case.hpp"
    case 6:
#define Width 6
#include "sorted_uint_vec_get_block_case.hpp"
    case 7:
#define Width 7
#include "sorted_uint_vec_get_block_case.hpp"
    case 8:
#define Width 8
#include "sorted_uint_vec_get_block_case.hpp"
    case 9:
#define Width 9
#include "sorted_uint_vec_get_block_case.hpp"
    case 10:
#define Width 10
#include "sorted_uint_vec_get_block_case.hpp"
    case 11: {
        // smallWidth and largeWidth are both 1, this is an optimization
        size_t headerLen;
        size_t const loWater = GetLoWater_x(header, &headerLen);
        size_t const* pData = (const size_t*)(uintptr_t(header) + headerLen);
        size_t val = sample0;
        for(size_t i = 0; i < blockUnits / TERARK_WORD_BITS; ++i) {
            auto   w = unaligned_load<size_t>(&pData[i]);
            for (size_t j = 0; j < TERARK_WORD_BITS; ++j) {
                aVals[i*TERARK_WORD_BITS + j] = val;
                val += loWater + (w & 1);
                w >>= 1;
            }
        }
        break; }
    case 12:
#define Width 12
#include "sorted_uint_vec_get_block_case.hpp"
    case 13:
#define Width 16
#include "sorted_uint_vec_get_block_case.hpp"
    case 14:
#define Width 20
#include "sorted_uint_vec_get_block_case.hpp"
    case 15: {
        // smallWidth and largeWidth are both 2, this is an optimization
        size_t headerLen;
        size_t const loWater = GetLoWater_x(header, &headerLen);
        byte_t const* pData = (const byte_t*)(uintptr_t(header) + headerLen);
        size_t const old_style_block_size = headerLen + blockUnits * 2 / 8;
        if (offset1 - offset0 > old_style_block_size) {
            size_t log2_range = m_log2_blockUnits;
            size_t headerVal = loWater;
            int64_t f_a = 0;
            int64_t f_b = sample1 - sample0;
            int64_t f_m = (f_a+f_b)/2 + decode_signed(headerVal);
            size_t plainWidth = (pData[0] & 63) + 1;
            assert(plainWidth <= 64);
            size_t bitpos = (headerLen + offset0) * 8 + 6;
            size_t a = 0;
            size_t b = blockUnits;
            size_t m = (b - a) / 2;
            aVals[0] = sample0;
            size_t i = 1;
            for (; i < blockUnits/2; ++i) {
                int64_t diff = varbits_get_signed(pLargeBase, bitpos, plainWidth);
                int64_t f_x_lag = lagrange(i, a, f_a, m, f_m, b, f_b, log2_range);
                aVals[i] = sample0 + f_x_lag + diff;
                bitpos += plainWidth;
            }
            aVals[i++] = sample0 + f_m; // do not inc bitpos
            for (; i < blockUnits; ++i) {
                int64_t diff = varbits_get_signed(pLargeBase, bitpos, plainWidth);
                int64_t f_x_lag = lagrange(i, a, f_a, m, f_m, b, f_b, log2_range);
                aVals[i] = sample0 + f_x_lag + diff;
                bitpos += plainWidth;
            }
            break; // break the case
        }
        size_t val = sample0;
        for(size_t i = 0; i < blockUnits / (TERARK_WORD_BITS/2); ++i) {
            auto   w = unaligned_load<size_t>(pData, i);
            for (size_t j = 0; j < TERARK_WORD_BITS/2; ++j) {
                aVals[i*(TERARK_WORD_BITS/2) + j] = val;
                val += loWater + (w & 3);
                w >>= 2;
            }
        }
        break; }
    } // switch
}

size_t SortedUintVec::get(size_t idx) const {
	assert(idx < m_size);
	assert(m_is_sorted_uint_vec);
	size_t aVal[2];
	get2(idx, aVal);
	return aVal[0];
}

SortedUintVec::SortedUintVec() {
	m_index = NULL;
	m_log2_blockUnits = 0;
	m_offsetWidth = 0;
	m_sampleWidth = 0;
	m_is_sorted_uint_vec = true;
	m_is_overall_full_sorted = true;
	m_is_samples_full_sorted = true;
#if TERARK_WORD_BITS == 64
	m_padding = 0xCCCCCCCC;
#endif
	m_size = 0;
}

SortedUintVec::~SortedUintVec() {
}

SortedUintVec::SortedUintVec(SortedUintVec&& o) noexcept { swap(o); }

SortedUintVec& SortedUintVec::operator=(SortedUintVec&& o) noexcept {
  if (this != &o) {
    swap(o);
  }
  return *this;
}

size_t SortedUintVec::lower_bound(size_t lo, size_t hi, size_t key) const {
    assert(lo <= hi);
    if (terark_unlikely(lo == hi)) {
        return lo;
    }
    size_t blockUnits = block_units(), mask = blockUnits - 1;
    size_t hiBlk = (hi + mask) >> m_log2_blockUnits;
    size_t loBlk = lo >> m_log2_blockUnits;
    size_t i = loBlk;
    if (loBlk + 1 < hiBlk) {
        auto indexBase = (const size_t* )(m_index);
        auto sampleWidth = m_sampleWidth;
        auto indexWidth = m_offsetWidth + sampleWidth;
        if (!m_is_samples_full_sorted) {
            size_t nextSample = s_get_block_min_val(
                indexBase, sampleWidth, indexWidth, loBlk+1);
            if (nextSample < key) {
                i = ++loBlk;
                lo = loBlk << m_log2_blockUnits;
            }
            else
                goto SearchBlock;
        }
        size_t j = hiBlk;
        while (i < j) {
            size_t mid_idx = (i + j) / 2;
            size_t mid_pos = indexWidth * (mid_idx + 1) - sampleWidth;
            size_t mid_val = febitvec::s_get_uint(indexBase, mid_pos, sampleWidth);
            if (mid_val < key)
                i = mid_idx + 1;
            else
                j = mid_idx;
        }
        if (terark_unlikely(loBlk == i)) {
            size_t hit_pos = indexWidth * (i + 1) - sampleWidth;
            size_t hit_val = febitvec::s_get_uint(indexBase, hit_pos, sampleWidth);
            if (key < hit_val) {
                return lo;
            }
        }
        else {
            i--;
        }
    }
SearchBlock:
    assert(blockUnits <= 128);
    size_t block[128]; get_block(i, block);
    size_t l = (i+0 == loBlk) ? lo & mask : 0;
    size_t u = (i+1 == hiBlk) ? ((hi-1) & mask) + 1 : blockUnits;
    size_t k = lower_bound_n(block, l, u, key);
    size_t p = k + (i << m_log2_blockUnits);
    return p;
}

size_t SortedUintVec::upper_bound(size_t lo, size_t hi, size_t key) const {
    assert(lo <= hi);
    if (terark_unlikely(lo == hi)) {
        return lo;
    }
    size_t blockUnits = block_units(), mask = blockUnits - 1;
    size_t hiBlk = (hi + mask) >> m_log2_blockUnits;
    size_t loBlk = lo >> m_log2_blockUnits;
    size_t i = loBlk;
    if (loBlk + 1 < hiBlk) {
        auto indexBase = (const size_t*)(m_index);
        auto sampleWidth = m_sampleWidth;
        auto indexWidth = m_offsetWidth + sampleWidth;
        if (!m_is_samples_full_sorted) {
            size_t nextSample = s_get_block_min_val(
                indexBase, sampleWidth, indexWidth, loBlk+1);
            if (nextSample <= key) {
                i = ++loBlk;
                lo = loBlk << m_log2_blockUnits;
            }
            else
                goto SearchBlock;
        }
        size_t j = hiBlk;
        while (i < j) {
            size_t mid_idx = (i + j) / 2;
            size_t mid_pos = indexWidth * (mid_idx + 1) - sampleWidth;
            size_t mid_val = febitvec::s_get_uint(indexBase, mid_pos, sampleWidth);
            if (mid_val <= key)
                i = mid_idx + 1;
            else
                j = mid_idx;
        }
        if (terark_unlikely(loBlk == i)) {
            return lo;
        }
        i--;
    }
SearchBlock:
    assert(blockUnits <= 128);
    size_t block[128]; get_block(i, block);
    size_t l = (i+0 == loBlk) ? lo & mask : 0;
    size_t u = (i+1 == hiBlk) ? ((hi-1) & mask) + 1 : blockUnits;
    size_t k = upper_bound_n(block, l, u, key);
    size_t p = k + (i << m_log2_blockUnits);
    return p;
}

std::pair<size_t, size_t>
SortedUintVec::equal_range(size_t lo, size_t hi, size_t key) const {
    assert(lo <= hi);
    if (lo == hi) {
        return std::make_pair(lo, lo);
    }
    size_t log2 = m_log2_blockUnits;
    size_t blockUnits = size_t(1) << log2, mask = blockUnits - 1;
    size_t hiBlk = (hi + mask) >> log2;
    size_t loBlk = lo >> log2;
    size_t i = loBlk, j = hiBlk;
    if (loBlk + 1 < hiBlk) {
        auto indexBase = (const size_t* )(m_index);
        auto sampleWidth = m_sampleWidth;
        auto indexWidth = m_offsetWidth + sampleWidth;
        if (!m_is_samples_full_sorted) {
            size_t nextSample = s_get_block_min_val(
                indexBase, sampleWidth, indexWidth, loBlk+1);
            if (nextSample < key) {
                i = ++loBlk;
                lo = loBlk << log2;
            }
            else { // key <= nextSample
                auto loBase = lo & ~mask;
                if (nextSample == key) {
                    hi = upper_bound(loBase + blockUnits, hi, key);
                    // after calling upper_bound, to reduce stack space,
                    // because upper_bound also needs a block[128]
                    size_t block[128]; get_block(i, block);
                    size_t lo_k = lower_bound_n(block, lo & mask, blockUnits, key);
                    lo = loBase + lo_k;
                }
                else { // key < nextSample
                    size_t block[128]; get_block(i, block);
                    auto r = equal_range_n(block, lo & mask, blockUnits, key);
                    lo = loBase + r.first;
                    hi = loBase + r.second;
                }
                return std::make_pair(lo, hi);
            }
        }
        while (i < j) {
            size_t mid_idx = (i + j) / 2;
            size_t mid_pos = indexWidth * (mid_idx + 1) - sampleWidth;
            size_t mid_val = febitvec::s_get_uint(indexBase, mid_pos, sampleWidth);
            if (mid_val < key)
                i = mid_idx + 1;
            else if ( mid_val > key)
                j = mid_idx;
            else {
                size_t ii = i;
                {
                    size_t k = j;
                    while (i < k) {
                        mid_idx = (i + k) / 2;
                        mid_pos = indexWidth * (mid_idx + 1) - sampleWidth;
                        mid_val = febitvec::s_get_uint(indexBase, mid_pos, sampleWidth);
                        if (mid_val < key) // lower_bound
                            i = mid_idx + 1;
                        else
                            k = mid_idx;
                    }
                }
                {
                    size_t k = ii;
                    while (k < j) {
                        mid_idx = (k + j) / 2;
                        mid_pos = indexWidth * (mid_idx + 1) - sampleWidth;
                        mid_val = febitvec::s_get_uint(indexBase, mid_pos, sampleWidth);
                        if (mid_val <= key) // upper_bound
                            k = mid_idx + 1;
                        else
                            j = mid_idx;
                    }
                }
                break;
            }
        }
        if (terark_unlikely(j == loBlk)) {
            return std::make_pair(lo, lo);
        }
        if (i > loBlk) {
            i--;
        }
    }
    assert(i < j);
    assert(blockUnits <= 128);
    size_t block[128]; get_block(i, block);
    size_t hi_u = (j == hiBlk) ? ((hi-1) & (blockUnits-1)) + 1 : blockUnits;
    size_t lo_l = (i == loBlk) ? lo & (blockUnits-1) : 0;
    size_t lo_k, hi_k;
    j--;
    if (i < j) { // different blocks
        lo_k = lower_bound_n(block, lo_l, blockUnits, key);
        get_block(j, block);
        hi_k = upper_bound_n(block, 0, hi_u, key);
    }
    else {
        auto lohi = equal_range_n(block, lo_l, hi_u, key);
        lo_k = lohi.first;
        hi_k = lohi.second;
    }
    lo = lo_k + (i << log2);
    hi = hi_k + (j << log2);
    return std::make_pair(lo, hi);
}

void SortedUintVec::risk_set_data(const void* base, size_t bytes) {
	assert(m_is_sorted_uint_vec);
	auto oheader = (const ObjectHeader*)(base);
	size_t numUnits = oheader->units;
	size_t log2_blockUnits = oheader->log2_blockUnits;
	size_t offsetWidth = oheader->offsetWidth_1 + 1;
	size_t sampleWidth = oheader->sampleWidth_1 + 1;
	size_t blockUnits = size_t(1) << log2_blockUnits;
	if (64 == blockUnits || 128 == blockUnits) {
		// OK
	} else {
		THROW_STD(invalid_argument
			, "invalid log2_blockUnits = %zd, must be 6 or 7"
			  " (blockUnits must be 64 or 128)"
			, log2_blockUnits
		);
	}
	if (0 != m_size || !m_data.empty()) {
		THROW_STD(invalid_argument
			, "this function must be called on empty object"
		);
	}
	size_t indexEntries = ceiled_div(numUnits, blockUnits) + 1;
	size_t indexWidth   = offsetWidth + sampleWidth;
	size_t indexBits    = febitvec::align_bits(indexWidth * indexEntries);
	size_t indexBytes   = indexBits / 8;
	size_t indexOffset  = oheader->indexOffset;
	if (bytes < indexOffset + indexBytes) {
		THROW_STD(invalid_argument
			, "bytes = %zd is too small, indexOffset = %zd, indexBytes = %zd"
			, bytes, indexOffset, indexBytes
		);
	}
	m_data.risk_set_data((byte_t*)base, bytes);
	m_index = (const byte_t*)(base) + indexOffset;
	m_size = numUnits;
	m_log2_blockUnits = byte_t(log2_blockUnits);
	m_offsetWidth = byte_t(offsetWidth);
	m_sampleWidth = byte_t(sampleWidth);
	m_is_overall_full_sorted = oheader->is_overall_full_sorted;
	m_is_samples_full_sorted = oheader->is_samples_full_sorted;
}

void SortedUintVec::risk_release_ownership() {
	m_data.risk_release_ownership();
	clear();
}

void SortedUintVec::clear() {
	SortedUintVec().swap(*this);
}

void SortedUintVec::swap(SortedUintVec& y) {
	m_data.swap(y.m_data);
	std::swap(m_index          , y.m_index);
	std::swap(m_size           , y.m_size);
	std::swap(m_log2_blockUnits, y.m_log2_blockUnits);
	std::swap(m_offsetWidth    , y.m_offsetWidth    );
	std::swap(m_sampleWidth    , y.m_sampleWidth    );
#define DO_SWAP_BIT(x, y) {bool t = x; x = y; y = t;}
	DO_SWAP_BIT(m_is_sorted_uint_vec   , y.m_is_sorted_uint_vec);
	DO_SWAP_BIT(m_is_overall_full_sorted, y.m_is_overall_full_sorted);
	DO_SWAP_BIT(m_is_samples_full_sorted, y.m_is_samples_full_sorted);
#if TERARK_WORD_BITS == 64
	std::swap(m_padding        , y.m_padding);
#endif
}

/////////////////////////////////////////////////////////////////////////////
class SortedUintVec::Builder::Impl : public SortedUintVec::Builder {
	struct IndexItem {
		size_t  offset;
		size_t  sample;
	};
	valvec<byte_t>  m_data;
	valvec<int64_t> m_block;
	AutoFree<int64_t> m_diffvec;
	valvec<IndexItem> m_index;
	uint64_t m_lastValue;
	uint64_t m_maxValue;
	size_t   m_size;
	byte_t   m_log2_blockUnits;
	bool     m_is_sorted;
	bool     m_is_real_sorted;

	// m_smallToLarge[largeCount][smallWidth][largeWidth]
	AutoFree<unsigned[16][64]> m_smallToLarge;

    boost::intrusive_ptr<FileStream> m_fp;
    boost::intrusive_ptr<OutputBuffer> m_writer;
    size_t m_startPos;
	size_t m_indexOffset;

	size_t getBlockUnits() const {
		return m_block.capacity();
	}
	void inc_histogram(size_t largeCount, size_t smallWidth, size_t largeWidth) {
		if (m_smallToLarge) {
			size_t largeWidth_1 = largeWidth ? largeWidth-1 : 0;
			size_t smallWidthType = getWidthType(smallWidth, largeWidth);
			m_smallToLarge[largeCount][smallWidthType][largeWidth_1]++;
		}
	}
	void print_histogram() const;
	void append_block(uint64_t nextBlockFirstValue);
	void append_block_impl();
	void append_block_impl_lagrange(uint64_t nextBlockFirstValue);

    void init(size_t blockUnits);

public:
    Impl(size_t blockUnits, OutputBuffer* buffer, bool is_sorted);
    Impl(size_t blockUnits, fstring fpath, bool is_sorted);
	void push_back(uint64_t val) override;
	void append(const uint64_t* values, size_t count);
	void append(const valvec<uint64_t>& v) { append(v.data(), v.size()); }
    BuildResult finish(SortedUintVec* vec) override;
};

template<size_t BlockUnits>
static inline
size_t bits(size_t lo, size_t hi, const int64_t* ds) {
	size_t realsmallWidth = 1 + terark_bsr_u64(uint64_t(ds[hi-1] - ds[lo] + 1));
	size_t largeWidth = 1 + terark_bsr_u64(uint64_t(ds[BlockUnits-2]));
	size_t smallBits;
	size_t smallWidth = g_used_small_width[realsmallWidth];
	if (0 == smallWidth) {
		// this is very rare case!
		// zero means realSmallWidth == 1 and all values use largeWidth
		// effective Width is largeWidth+1
		return (largeWidth+1) * (BlockUnits-1);
	}
	if (64 % smallWidth == 0) {
		// last smallWidth bits can be used by largeUnits
		smallBits = smallWidth * (BlockUnits-1);
	} else {
		// tricky: for every 64 units, use uint64 words = smallWidth
		// and wasted bits == smallWidth
		smallBits = smallWidth * 64 * (BlockUnits/64);
	}
	size_t largeBits = largeWidth * (BlockUnits-1 - (hi-lo));
	return smallBits + largeBits;
}

template<size_t BlockUnits, size_t Dim>
static inline
void compute_lo_hi_index(const int64_t* ds, size_t* pLo, size_t* pHi) {
	uint32_t aa[Dim][Dim];
	for (size_t i = 0; i < Dim; ++i) {
		for (size_t j = 0; j < Dim; ++j)
			aa[i][j] = (uint32_t)bits<BlockUnits>(i, BlockUnits-1-j, ds);
	}
	size_t vMin = size_t(-1);
	size_t lo = size_t(-1), hi = size_t(-1);
	for (size_t i = 0; i < Dim; ++i) {
		for (size_t j = 0; j < Dim; ++j)
			if (vMin > aa[i][j])
				vMin = aa[i][j], lo = i, hi = j;
	}
	hi = BlockUnits-1-hi;
	*pLo = lo;
	*pHi = hi;
}

static inline
void compute_lo_hi_index(size_t blockUnits, const int64_t* ds,
						 size_t* pLo, size_t* pHi) {
	if (64 == blockUnits) {
		compute_lo_hi_index< 64, 32>(ds, pLo, pHi);
	} else if (128 == blockUnits) {
		compute_lo_hi_index<128, 64>(ds, pLo, pHi);
	} else {
		assert(0); // should not happen
		THROW_STD(invalid_argument,
			"bug: invalid blockUnits = %zd, should not happen", blockUnits);
	}
}

void SortedUintVec::Builder::Impl::append_block(uint64_t nextBlockFirstValue) {
	size_t const old_data_size = m_data.size();
	size_t const blockUnits = getBlockUnits();
	auto vals = m_block.data();
	bool isIncreasing = true;
	for (size_t i = 1; i < blockUnits; ++i) {
		if (vals[i-1] > vals[i]) {
			isIncreasing = false;
			break;
		}
	}
	if (isIncreasing) {
		append_block_impl();
	} else {
		append_block_impl_lagrange(nextBlockFirstValue);
	}
	m_indexOffset += m_data.size() - old_data_size;
	if (m_writer) {
		assert(sizeof(ObjectHeader) == old_data_size);
		auto blockData = m_data.data() + sizeof(ObjectHeader);
		auto blockSize = m_data.size() - sizeof(ObjectHeader);
		m_writer->ensureWrite(blockData, blockSize);
		m_data.risk_set_size(sizeof(ObjectHeader));
	}
	else {
		assert(m_data.size() == m_indexOffset);
	}
}

void SortedUintVec::Builder::Impl::append_block_impl() {
	const size_t blockUnits = getBlockUnits();
	auto vals = m_block.data();
	auto ds = m_diffvec.p;
	m_index.push_back({m_indexOffset - sizeof(ObjectHeader), (size_t)vals[0]});
	for (size_t i = 1; i < blockUnits; ++i) {
		ds[i-1] = vals[i] - vals[i-1];
	}
	std::sort(ds, ds+blockUnits-1); // sort the diff sequence
	if (ds[0] == ds[blockUnits-2]) {
		WriteHeader(&m_data, ds[0], 0, 0);
		inc_histogram(0, 0, 0);
		return;
	}
	if (ds[0] + 1 == ds[blockUnits-2]) {
		int64_t loWater = ds[0];
		WriteHeader(&m_data, loWater, 1, 1);
		inc_histogram(0, 1, 1);
		// re-generate the diff sequence
		for (size_t i = 1; i < blockUnits; ++i) {
			ds[i-1] = vals[i] - vals[i-1];
		}
		ds[blockUnits-1] = loWater; // to zero last unused bit
		// do not need to store large units
		for (size_t i = 0; i < blockUnits / 64; ++i) {
			uint64_t w = 0;
			auto pds = ds + i * 64;
			for (size_t j = 0; j < 64; j++) {
				w |= (pds[j] - loWater) << j; // don't need +1
			}
			m_data.append((byte_t*)&w, 8);
		}
		return;
	}
	if (ds[0] + 3 >= ds[blockUnits-2]) {
		int64_t loWater = ds[0];
		WriteHeader(&m_data, loWater, 2, 2);
		inc_histogram(0, 2, 2);
		// re-generate the diff sequence
		for (size_t i = 1; i < blockUnits; ++i) {
			ds[i-1] = vals[i] - vals[i-1];
		}
		ds[blockUnits-1] = loWater; // to zero last unused unit
		// do not need to store large units
		for (size_t i = 0; i < blockUnits / 32; ++i) {
			uint64_t w = 0;
			auto pds = ds + i * 32;
			for (size_t j = 0; j < 32; j++) {
				w |= (pds[j] - loWater) << (2*j); // don't need +1
			}
			m_data.append((byte_t*)&w, 8);
		}
		return;
	}
	size_t lo, hi;
	compute_lo_hi_index(blockUnits, ds, &lo, &hi);
	int64_t loWater = ds[lo], hiWater = ds[hi-1];
	size_t realsmallWidth = 1 + terark_bsr_u64(uint64_t(hiWater - loWater + 1));
	size_t largeWidth = 1 + terark_bsr_u64(uint64_t(ds[blockUnits-2]));
	assert(realsmallWidth <= largeWidth); // may be equal
	assert(largeWidth >= 2);
	size_t smallWidth = g_used_small_width[realsmallWidth];
	inc_histogram(blockUnits-1-(hi-lo), smallWidth, largeWidth);
	// re-generate the diff sequence
	for (size_t i = 1; i < blockUnits; ++i) {
		ds[i-1] = vals[i] - vals[i-1];
	}
	WriteHeader(&m_data, loWater, smallWidth, largeWidth);
	size_t wordUnits = 64/smallWidth;
	size_t smallOffset = m_data.size();
	size_t largeBitPos, largeBytes, largeUnits = 0;
	if (64 % smallWidth == 0) {

		// to set last small unit
		// last small unit will always be overwritten
		//
		// set it to loWater will mark it as 'smallUnit'
		// and 'largeUnits' will not be over increased
		ds[blockUnits-1] = loWater;

		for(size_t i = 0; i < blockUnits / wordUnits; ++i) {
			auto pds = ds + i * (64 / smallWidth);
			uint64_t w = 0;
			for (size_t j = 0; j < wordUnits; ++j) {
				if (pds[j] >= loWater && pds[j] <= hiWater)
					w |= (pds[j] - loWater + 1) << (smallWidth*j);
				else
					largeUnits++;
			}
			m_data.append((byte_t*)&w, 8);
		}
		largeBitPos = smallWidth * (blockUnits - 1);
		// largeWidth use 6 bits:
		largeBytes = (largeWidth * largeUnits - smallWidth + 6 + 7) / 8;
	}
	else {
		uint64_t encoded[44];
		size_t baseWords = smallWidth * (blockUnits/64); // tricky!
		assert(baseWords <= 40);
		size_t virtualWords = (blockUnits-1 + wordUnits-1) / wordUnits;
		assert(virtualWords <= 44);
		for(size_t i = 0; i < virtualWords; ++i) {
			auto pds = ds + i * wordUnits;
			size_t n;
			if (i < virtualWords-1 || (blockUnits-1)%wordUnits == 0)
				n = wordUnits;
			else
				n = (blockUnits-1) % wordUnits;
			uint64_t w = 0;
			for (size_t j = 0; j < n; ++j) {
				if (pds[j] >= loWater && pds[j] <= hiWater)
					w |= (pds[j] - loWater + 1) << (smallWidth * j);
				else
					largeUnits++;
			}
			encoded[i] = w;
		}
		size_t baseBitWidth = smallWidth * wordUnits;
		size_t extrBitWidth = 64 - baseBitWidth;
		size_t extrBitPos = 0;
		uint64_t extrBitsPacked[8];
		// high extrBitWidth bits are hole bits in each word
		// pack such extra bits to extrBitsPacked makes no holes
		for(size_t i = baseWords; i < virtualWords; ++i) {
			uint64_t w = encoded[i];
			assert(w <= uint64_t(-1) >> extrBitWidth);
			assert(extrBitPos + baseBitWidth <= sizeof(extrBitsPacked)*8);
			febitvec::s_set_uint(extrBitsPacked, extrBitPos, baseBitWidth, w);
			extrBitPos += baseBitWidth;
			// extrBitPos may exceeds real bitpos, but this is ok
			// because exceeded bits are all zero
		}
		extrBitPos = 0;
		for(size_t i = 0; i < baseWords; ++i) {
			uint64_t extrBits = febitvec::s_get_uint(extrBitsPacked, extrBitPos, extrBitWidth);
			encoded[i] |= extrBits << baseBitWidth;
			extrBitPos += extrBitWidth;
		}
		m_data.append((byte_t*)encoded, 8 * baseWords);
		largeBitPos = 64 * baseWords;
		// largeWidth use 6 bits:
		largeBytes = (largeWidth * largeUnits + 6 + 7) / 8;
	}
	if (largeUnits) {
		m_data.grow(largeBytes + 7); // 7 bytes extra for safe uint64
		m_data.pop_n(7); // return back 7 bytes
		largeBitPos += 8 * smallOffset;
		uint64_t* pData = (uint64_t*)(m_data.data());
		febitvec::s_set_uint(pData, largeBitPos, 6, largeWidth - smallWidth);
		largeBitPos += 6;
		for (size_t i = 0; i < blockUnits-1; ++i) {
			if (ds[i] < loWater || ds[i] > hiWater) {
				febitvec::s_set_uint(pData, largeBitPos, largeWidth, ds[i]);
				largeBitPos += largeWidth;
			}
		}
	} else {
		TERARK_IF_DEBUG(int xxx = 0; TERARK_UNUSED_VAR(xxx);,;);
	}
}

void SortedUintVec::Builder::Impl::append_block_impl_lagrange(uint64_t nextBlockFirstValue) {
    const size_t log2_range = m_log2_blockUnits;
    const size_t blockUnits = getBlockUnits();
    auto vals = m_block.data();
    m_index.push_back({ m_indexOffset - sizeof(ObjectHeader), (size_t)vals[0] });
    int64_t a = 0;
    int64_t b = blockUnits;
    int64_t m = (b-a)/2;
    int64_t f_a = 0;
    int64_t f_b = nextBlockFirstValue - vals[0];
    int64_t f_m = vals[m] - vals[0];
    int64_t f_m_encoded = encode_signed(f_m - (f_a + f_b) / 2);
    auto fix = (uint64_t*)m_diffvec.p;
    auto set_fix = [=](size_t idx, size_t x) {
        int64_t f_x_lag = lagrange(x, a, f_a, m, f_m, b, f_b, log2_range);
        int64_t f_x_real = vals[x] - vals[0];
        int64_t diff = f_x_real - f_x_lag;
    //  printf("build: x = %3zd , f_x_lag = %3zd , diff = %3zd\n", x, f_x_lag, diff);
        fix[idx] = encode_signed(diff);
    };
    for (size_t i = 1; i < blockUnits / 2; ++i) {
        set_fix(i - 1, i);
    }
    for (size_t i = blockUnits / 2 + 1; i < blockUnits; ++i) {
        set_fix(i - 2, i);
    }
//  uint64_t minFix = fix[0];
    uint64_t maxFix = fix[0];
    for (size_t i = 1; i < blockUnits - 2; ++i) {
//      if (minFix > fix[i]) minFix = fix[i];
        if (maxFix < fix[i]) maxFix = fix[i];
    }
    WriteHeader(&m_data, f_m_encoded, 2, 2); // special (2,2)
//  size_t f_width = (maxFix - minFix) < 4
//                 ? 2 : terark_bsr_u64(maxFix - minFix) + 1;
    size_t f_width = maxFix < 4 ? 2 : terark_bsr_u64(maxFix) + 1;
    size_t bitpos = m_data.size() * 8;
    m_data.grow(ceiled_div(6 + f_width*(blockUnits - 2), 8));
    uint64_t* pData = (uint64_t*)(m_data.data());
    febitvec::s_set_uint(pData, bitpos, 6, f_width - 1);
    bitpos += 6;
    for (size_t i = 0; i < blockUnits - 2; ++i) {
    //  uint64_t val = fix[i] - minFix; // !!don't use minFix
        uint64_t val = fix[i];
        febitvec::s_set_uint(pData, bitpos, f_width, val);
        bitpos += f_width;
    }
    assert(bitpos <= m_data.size() * 8);
}

void SortedUintVec::Builder::Impl::init(size_t blockUnits) {
    if (64 != blockUnits && 128 != blockUnits) {
        THROW_STD(invalid_argument
            , "invalid blockUnits = %zd, must be 64 or 128"
            , blockUnits);
    }
    //	if (fast_popcount(blockUnits) != 1) {
    //		THROW_STD(invalid_argument,
    //			"invalid blockUnits = %zd", blockUnits);
    //	}
    m_data.reserve(4096);
    m_data.grow(sizeof(ObjectHeader));
    m_index.reserve(256);
    m_lastValue = 0;
    m_maxValue = 0;
    m_size = 0;
    m_block.reserve(blockUnits);
    m_diffvec.alloc(blockUnits); // has 1 extra unit
    m_log2_blockUnits = terark_bsr_u64(blockUnits);

    if (getEnvBool("SortedUintVec_Builder_EnableHistogram")) {
        m_smallToLarge.alloc(blockUnits);
        memset(m_smallToLarge, 0, sizeof(*m_smallToLarge)*blockUnits);
    }

    m_indexOffset = sizeof(ObjectHeader);
}

SortedUintVec::Builder::Impl::Impl(size_t blockUnits, OutputBuffer* buffer, bool is_sorted) {
    init(blockUnits);
    m_is_sorted = is_sorted;
    m_is_real_sorted = true;
    // m_fp should not assign
    // for buffer->getOutputStream() can be stack object
    if (buffer) {
        auto fpSeekable = dynamic_cast<ISeekable*>(buffer->getOutputStream());
        assert(fpSeekable != nullptr);
        m_writer.reset(buffer);
        m_startPos = fpSeekable->tell() + m_writer->bufpos();
        m_writer->ensureWrite(m_data.data(), sizeof(ObjectHeader)); // all zero
    }
}

SortedUintVec::Builder::Impl::Impl(size_t blockUnits, fstring fpath, bool is_sorted) {
    init(blockUnits);
    m_is_sorted = is_sorted;
    m_is_real_sorted = true;
    m_fp.reset(new FileStream(fpath.c_str(), "wb+"));
    m_writer.reset(new OutputBuffer(m_fp.get()));
    m_startPos = 0;
    m_writer->ensureWrite(m_data.data(), sizeof(ObjectHeader)); // all zero
}

void SortedUintVec::Builder::Impl::push_back(uint64_t val) {
	const size_t blockUnits = getBlockUnits();
	if (terark_unlikely(m_block.size() >= blockUnits)) {
		if (terark_unlikely(m_block.size() > blockUnits)) {
			THROW_STD(invalid_argument
				, "blockUnits = %zd, but m_block.size = %zd"
				,  blockUnits, m_block.size()
			);
		}
		append_block(val);
		m_block.risk_set_size(0);
	}
	if (terark_unlikely(m_block.capacity() < blockUnits)) {
		m_block.reserve(blockUnits);
	}
	if (val < m_lastValue) {
		if (m_is_sorted) {
			THROW_STD(invalid_argument
				, "values must be in ascending order, m_lastValue = %llu, cur = %llu"
				,  ullong(m_lastValue), ullong(val)
			);
		}
		m_is_real_sorted = false;
	}
	else if (terark_unlikely(m_size > 0 && val - m_lastValue >= uint64_t(1)<<(20+31))) {
		THROW_STD(invalid_argument
			, "diff = 0x%llX is larger than 51 bits, m_lastValue = %llu, cur = %llu"
			,  ullong(val - m_lastValue), ullong(m_lastValue), ullong(val)
		);
	}
	m_maxValue = std::max(m_maxValue, val);
	m_block.unchecked_push_back(val);
	m_size++;
	m_lastValue = val;
}

void SortedUintVec::Builder::Impl::append(const uint64_t* values, size_t count) {
	const size_t blockUnits = getBlockUnits();
	if (terark_unlikely(count != blockUnits)) {
		THROW_STD(invalid_argument
			, "blockUnits = %zd, but argument block.size = %zd"
			,  blockUnits, count
		);
	}
	for (size_t i = 0; i < count; ++i) {
		push_back(values[i]);
	}
}

//#ifdef __GNUC__
//#   pragma GCC push_options
//#   pragma GCC optimize("-fno-tree-partial-pre")
//    // -ftree-partial-pre makes bad_alloc exception
//#endif

SortedUintVec::Builder::BuildResult
SortedUintVec::Builder::Impl::finish(SortedUintVec* vec) {
    const size_t blockUnits = getBlockUnits();
    if (m_block.size()) {
        while (m_block.size() < blockUnits)
            m_block.unchecked_push_back(m_lastValue);
        append_block(m_lastValue);
    }
    m_index.push_back({ m_indexOffset - sizeof(ObjectHeader), (size_t)m_lastValue });
    m_data.append("\0\0\0\0\0\0\0", 7); // zero padding
    if (m_writer) {
        while ((m_indexOffset + m_data.size()) % 8 != 0) { // align padding
            m_data.push_back(0);
        }
        auto pad_data = m_data.data() + sizeof(ObjectHeader);
        auto pad_size = m_data.size() - sizeof(ObjectHeader);
        m_writer->ensureWrite(pad_data, pad_size);
        m_indexOffset += pad_size;
#if !defined(NDEBUG)
        auto fpSeekable = dynamic_cast<ISeekable*>(m_writer->getOutputStream());
        assert(!m_fp || dynamic_cast<FileStream*>(fpSeekable) == m_fp.get());
        assert(size_t(fpSeekable->tell()) + m_writer->bufpos() == m_startPos + m_indexOffset);
#endif
        m_data.risk_set_size(sizeof(ObjectHeader));
    }
    else {
        while (m_data.size() % 8 != 0) { // align padding
            m_data.push_back(0);
        }
        m_indexOffset = m_data.size();
    }
    bool is_samples_full_sorted = true;
    if (!m_is_real_sorted) {
        for (size_t i = 1; i < m_index.size(); ++i) {
            if (m_index[i-1].sample > m_index[i].sample) {
                is_samples_full_sorted = false;
                break;
            }
        }
    }
    else {
        assert(m_lastValue == m_maxValue);
    }
    size_t numBlocks = (m_size + blockUnits - 1) / blockUnits;
    assert(numBlocks + 1 == m_index.size());

    // m_maxValue==0 cause terark_bsr_u64 yields undefined value
    // so use safe_bsr_u64
    size_t offsetWidth = 1 + safe_bsr_u64(m_indexOffset - sizeof(ObjectHeader));
    size_t sampleWidth = 1 + safe_bsr_u64(m_maxValue);

    size_t indexWidth = offsetWidth + sampleWidth;
    m_data.grow(ceiled_div(indexWidth*(numBlocks + 1), 8));
    BuildResult result;
    if (m_writer) {
        while ((m_indexOffset + m_data.size()) % 16 != 0) { // align padding
            m_data.push_back(0);
        }
        result.mem_size = m_indexOffset + m_data.size();
    }
    else {
        while (m_data.size() % 16 != 0) { // align padding
            m_data.push_back(0);
        }
        result.mem_size = m_data.size();
    }
    m_data.shrink_to_fit();
  {
    IndexItem* rawIndex = m_index.data();
    uint64_t * zipIndex;
    if (m_writer) {
        zipIndex = (uint64_t*)(m_data.data() + sizeof(ObjectHeader));
    }
    else {
        zipIndex = (uint64_t*)(m_data.data() + m_indexOffset);
    }
    for(size_t bitpos = 0, i = 0; i < numBlocks + 1; ++i) {
        size_t offset = rawIndex[i].offset;
        size_t sample = rawIndex[i].sample;
        febitvec::s_set_uint(zipIndex, bitpos, offsetWidth, offset); bitpos += offsetWidth;
        febitvec::s_set_uint(zipIndex, bitpos, sampleWidth, sample); bitpos += sampleWidth;
    }
  }
#if !defined(NDEBUG) && 0
    for (size_t bitpos = 0, i = 0; i < numBlocks + 1; ++i) {
        size_t offset = rawIndex[i].offset;
        size_t sample = rawIndex[i].sample;
        size_t offset2 = febitvec::s_get_uint(zipIndex, bitpos, offsetWidth); bitpos += offsetWidth;
        size_t sample2 = febitvec::s_get_uint(zipIndex, bitpos, sampleWidth); bitpos += sampleWidth;
        assert(offset2 == offset);
        assert(sample2 == sample);
    }
#endif
    auto oheader = (ObjectHeader*)m_data.data();
    oheader->units = m_size;
    oheader->log2_blockUnits = m_log2_blockUnits;
    oheader->offsetWidth_1 = offsetWidth - 1;
    oheader->sampleWidth_1 = sampleWidth - 1;
    oheader->indexOffset = m_indexOffset;
    oheader->is_samples_full_sorted = is_samples_full_sorted;
    oheader->is_overall_full_sorted = m_is_real_sorted;
    if (m_writer) {
        // m_fp   = header + data + index
        // m_data = header        + index
        //                  ^^^^^^
        //                  data is not in memory
        m_writer->ensureWrite(oheader + 1, m_data.size() - sizeof(ObjectHeader));
        m_writer->flush_buffer();
        auto fpOutput = m_writer->getOutputStream();
        auto fpInput = dynamic_cast<IInputStream*>(fpOutput);
        auto fpSeekable = dynamic_cast<ISeekable*>(fpOutput);
        fpOutput->flush(); // must flush, VC bug?
        auto endPos = fpSeekable->tell();
        fpSeekable->seek(m_startPos);
        m_writer->ensureWrite(oheader, sizeof(ObjectHeader));
        m_writer->flush_buffer();
        fpOutput->flush(); // must flush, VC bug?
        if (vec) {
            size_t indexSize = m_data.size() - sizeof(ObjectHeader);
            size_t  dataSize = m_indexOffset - sizeof(ObjectHeader);
            m_data.grow_no_init(dataSize);
            auto base = m_data.data();
            memmove(base + m_indexOffset, base + sizeof(ObjectHeader), indexSize);
            auto readSize = fpInput->read(base + sizeof(ObjectHeader), dataSize);
            TERARK_UNUSED_VAR(readSize);
            assert(readSize == dataSize);
            assert(dataSize % 8 == 0);
        }
        fpSeekable->seek(endPos);
        m_writer.reset();
        m_fp.reset();
    }
    else {
        if (NULL == vec) {
            THROW_STD(invalid_argument, "If without file, vec must not be NULL");
        }
    }
    if (vec) {
        vec->m_size = m_size;
        vec->m_data.swap(m_data);
        vec->m_index = vec->m_data.data() + m_indexOffset;
        vec->m_log2_blockUnits = m_log2_blockUnits;
        vec->m_offsetWidth = offsetWidth;
        vec->m_sampleWidth = sampleWidth;
        vec->m_is_overall_full_sorted = m_is_real_sorted;
        vec->m_is_samples_full_sorted = is_samples_full_sorted;
    }
    if (m_smallToLarge) {
        print_histogram();
    }
    result.size = m_size;
    return result;
}

//#ifdef __GNUC__
//#   pragma GCC pop_options
//#endif

void SortedUintVec::Builder::Impl::print_histogram() const {
	if (!m_smallToLarge) {
		fprintf(stderr, "ERROR: SortedUintVec::Builder::Impl: histogram is not enabled\n");
		return;
	}
	// wt: widthType
	// lw: largeWidth-1
	// lc: largeCount
	printf("wt lw\t(lc   freq)+\n");
	printf(" 0  0\t  0 %6u\n", m_smallToLarge[0][0][0]);
	const size_t blockUnits = getBlockUnits();
	for(size_t wt = 1; wt < 16; ++wt) {
		for (size_t lw = 0; lw < 64; ++lw) {
			size_t cn = 0;
			for (size_t lc = 0; lc < blockUnits; ++lc) {
				cn += m_smallToLarge[lc][wt][lw];
			}
			if (cn) {
				printf("%2zd %2zd", wt, lw+1);
				for (size_t lc = 0; lc < blockUnits; ++lc) {
					if (unsigned freq = m_smallToLarge[lc][wt][lw]) {
						printf("\t%3zd %6u", lc, freq);
					}
				}
				printf("\n");
			}
		}
	}
}

///////////////////////////////////////////////////////////////////////

SortedUintVec::Builder::~Builder() {}

SortedUintVec::Builder*
SortedUintVec::createBuilder(size_t blockUnits) {
	return createBuilder(true, blockUnits);
}

SortedUintVec::Builder*
SortedUintVec::createBuilder(size_t blockUnits, const char* fname) {
	return createBuilder(true, blockUnits, fname);
}

SortedUintVec::Builder*
SortedUintVec::createBuilder(size_t blockUnits, OutputBuffer* buffer) {
	return createBuilder(true, blockUnits, buffer);
}

SortedUintVec::Builder*
SortedUintVec::createBuilder(bool inputSorted, size_t blockUnits) {
	return new Builder::Impl(blockUnits, nullptr, inputSorted);
}

SortedUintVec::Builder*
SortedUintVec::createBuilder(bool inputSorted, size_t blockUnits, const char* fname) {
	if (NULL == fname)
		return new Builder::Impl(blockUnits, nullptr, inputSorted);
	return new Builder::Impl(blockUnits, fname, inputSorted);
}

SortedUintVec::Builder*
SortedUintVec::createBuilder(bool inputSorted, size_t blockUnits, OutputBuffer* buffer) {
	return new Builder::Impl(blockUnits, buffer, inputSorted);
}

} // namespace terark


/***********************************************************************

Copyright (c) 2014, Carnegie Mellon University
All rights reserved.

Authors: Jennifer King <jeking04@gmail.com>

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are
met:

  Redistributions of source code must retain the above copyright
  notice, this list of conditions and the following disclaimer.

  Redistributions in binary form must reproduce the above copyright
  notice, this list of conditions and the following disclaimer in the
  documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

*************************************************************************/

#include <vector>
#include <Eigen/Geometry>
#include <ompl/util/RandomNumbers.h>
#include <or_ompl/TSR.h>

using namespace or_ompl;

TSR::TSR()
    : _initialized(false)
{
}

TSR::TSR(
        const Eigen::Affine3d &T0_w,
        const Eigen::Affine3d &Tw_e,
        const Eigen::Matrix<double, 6, 2> &Bw)
    : _T0_w(T0_w)
    , _Tw_e(Tw_e)
    , _Bw(Bw)
    , _manipulator_index(-1)
    , _relative_body_name("NULL")
    , _relative_link_name("")
    , _initialized(true)
{
    _T0_w_inv = _T0_w.inverse();
    _Tw_e_inv = _Tw_e.inverse();
}

bool TSR::deserialize(std::stringstream &ss)
{
    return deserialize(static_cast<std::istream &>(ss));
}

bool TSR::deserialize(std::istream &ss)
{
    // Set _initialized to false in case an error occurs.
    _initialized = false;

    ss >> _manipulator_index
       >> _relative_body_name;

    if(_relative_body_name != "NULL")
        ss >> _relative_link_name;

    // Read in the T0_w matrix
    for(unsigned int c=0; c < 3; c++)
    for(unsigned int r=0; r < 3; r++)
        ss >> _T0_w.matrix()(r,c);

    for(unsigned int idx=0; idx < 3; idx++)
        ss >> _T0_w.translation()(idx);

    // Read in the Tw_e matrix
    for(unsigned int c=0; c < 3; c++)
    for(unsigned int r=0; r < 3; r++)
        ss >> _Tw_e.matrix()(r,c);

    for(unsigned int idx=0; idx < 3; idx++)
        ss >> _Tw_e.translation()(idx);

    // Read in the Bw matrix
    for(unsigned int r=0; r < 6; r++)
    for(unsigned int c=0; c < 2; c++)
        ss >> _Bw(r,c);

    // Check for an error.
    if (!ss)
        return false;

    _T0_w_inv = _T0_w.inverse();
    _Tw_e_inv = _Tw_e.inverse();
    _initialized = true;

    return true;
}

void TSR::serialize(std::ostream& ss)
{
    if (!_initialized)
        throw std::runtime_error("TSR is not initialized.");

    ss << _manipulator_index
       << ' ' << _relative_body_name;

    if(_relative_body_name != "NULL")
        ss << ' ' << _relative_link_name;

    // T0_w matrix
    for(unsigned int c=0; c < 3; c++)
    for(unsigned int r=0; r < 3; r++)
        ss << ' ' << _T0_w.matrix()(r, c);

    for(unsigned int idx=0; idx < 3; idx++)
        ss << ' ' << _T0_w.translation()(idx);

    // Tw_e matrix
    for(unsigned int c=0; c < 3; c++)
    for(unsigned int r=0; r < 3; r++)
        ss << ' ' << _Tw_e.matrix()(r, c);

    for(unsigned int idx=0; idx < 3; idx++)
        ss << ' ' << _Tw_e.translation()(idx);

    // Read in the Bw matrix
    for(unsigned int r=0; r < 6; r++)
    for(unsigned int c=0; c < 2; c++)
        ss << ' ' << _Bw(r, c);
}

Eigen::Matrix<double, 6, 1> TSR::distance(const Eigen::Affine3d &ee_pose) const {
    Eigen::Matrix<double, 6, 1> dist = Eigen::Matrix<double, 6, 1>::Zero();

    // First compute the pose of the w frame in world coordinates, given the ee_pose
    Eigen::Affine3d w_in_world = ee_pose * _Tw_e_inv;

    // Next compute the pose of the w frame relative to its original pose (as specified by T0_w)
    Eigen::Affine3d w_offset = _T0_w_inv * w_in_world;

    // Now compute the elements of the distance matrix
    dist(0,0) = w_offset.translation()(0);
    dist(1,0) = w_offset.translation()(1);
    dist(2,0) = w_offset.translation()(2);
    dist(3,0) = atan2(w_offset.rotation()(2,1), w_offset.rotation()(2,2));
    dist(4,0) = -asin(w_offset.rotation()(2,0));
    dist(5,0) = atan2(w_offset.rotation()(1,0), w_offset.rotation()(0,0));

    return dist;
}

Eigen::Matrix<double, 6, 1> TSR::displacement(const Eigen::Affine3d &ee_pose) const {

    Eigen::Matrix<double, 6, 1> dist = distance(ee_pose);
    Eigen::Matrix<double, 6, 1> disp = Eigen::Matrix<double, 6, 1>::Zero();

    for(unsigned int idx=0; idx < 6; idx++){
        if(dist(idx,0) < _Bw(idx,0)){
            disp(idx,0) = dist(idx,0) - _Bw(idx,0);
        }else if(dist(idx,0) > _Bw(idx,1)){
            disp(idx,0) = dist(idx,0) - _Bw(idx,1);
        }
    }

    return disp;
}

Eigen::Affine3d TSR::sampleDisplacementTransform(void) const {

    // First sample uniformly betwee each of the bounds of Bw
    std::vector<double> d_sample(6);

    ompl::RNG rng;
    for(unsigned int idx=0; idx < d_sample.size(); idx++){
        if(_Bw(idx,1) > _Bw(idx,0)){
            d_sample[idx] = rng.uniformReal(_Bw(idx,0), _Bw(idx,1));
        }
    }

    Eigen::Affine3d return_tf;
    return_tf.translation() << d_sample[0], d_sample[1], d_sample[2];

    // Convert to a transform matrix
    double roll = d_sample[3];
    double pitch = d_sample[4];
    double yaw = d_sample[5];

    double A = cos(yaw);
    double B = sin(yaw);
    double C = cos(pitch);
    double D = sin(pitch);
    double E = cos(roll);
    double F = sin(roll);
    return_tf.linear() << A*C, A*D*F - B*E, B*F + A*D*E,
        B*C, A*E + B*D*F, B*D*E - A*F,
        -D, C*F, C*E;

    return return_tf;
}

Eigen::Affine3d TSR::sample() const {

    Eigen::Affine3d tf = sampleDisplacementTransform();

    return _T0_w * tf * _Tw_e;
}


/*
 * Copyright 2014 RWTH Aachen University. All rights reserved.
 *
 * Licensed under the RWTH LM License (the "License");
 * you may not use this file except in compliance with the License.
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <boost/functional/hash.hpp>
#include <iomanip>
#include <iostream>
//#include <numeric>  // only for checking normalization
#include "output.h"

Output::Output(const int input_dimension,
               const int max_batch_size,
               const int max_sequence_length,
               const int num_oovs,
               const bool use_bias,
               ConstVocabularyPointer vocabulary,
               ActivationFunctionPointer activation_function)
    : Function(
          input_dimension,
          vocabulary->GetVocabularySize() + vocabulary->GetNumClasses() -
              vocabulary->ComputeShortlistSize(),
          max_batch_size,
          max_sequence_length),
      activation_function_(std::move(activation_function)),
      num_classes_(vocabulary->GetNumClasses()),
      max_class_size_(vocabulary->GetMaxClassSize()),
      num_oovs_(num_oovs),
      num_out_of_shortlist_words_(vocabulary->GetVocabularySize() -
                                  vocabulary->ComputeShortlistSize()),
      shortlist_size_(vocabulary->ComputeShortlistSize()),
      vocabulary_(vocabulary) {
  class_b_ = FastMalloc((num_classes_ + max_class_size_) * max_batch_size *
                        max_sequence_length);
  class_delta_ = FastMalloc((num_classes_ + max_class_size_) * max_batch_size *
                            max_sequence_length);
  class_weights_ = FastMalloc(num_classes_ * input_dimension);
  class_bias_ = use_bias ? FastMalloc(num_classes_) : nullptr;
  momentum_class_weights_ = FastMalloc(num_classes_ * input_dimension);
  momentum_class_bias_ = use_bias ? FastMalloc(num_classes_) : nullptr;

  word_b_ = class_b_ + num_classes_ * max_batch_size;
  word_delta_ = class_delta_ + num_classes_ * max_batch_size;
  word_weights_ = FastMalloc(num_out_of_shortlist_words_ * input_dimension);
  word_bias_ = use_bias ? FastMalloc(num_out_of_shortlist_words_) : nullptr;

  int sum = 0;
  for (int i = 0; i < num_classes_; ++i) {
    word_offset_.push_back(sum);
    const int class_size = vocabulary_->GetClassSize(i);
    // we do not need weights for shortlist words
    if (class_size != 1)
      sum += class_size;
  }
}

Output::~Output() {
  FastFree(class_b_);
  FastFree(class_delta_);
  FastFree(class_weights_);
  FastFree(class_bias_);
  FastFree(word_weights_);
  FastFree(word_bias_);
  FastFree(momentum_class_weights_);
  FastFree(momentum_class_bias_);
}

const Real *Output::Evaluate(const Slice &slice, const Real x[]) {
  const Real *result = class_b_t_;
  // class part
  if (class_bias_) {
    for (size_t i = 0; i < slice.size(); ++i)
      FastCopy(class_bias_, num_classes_, class_b_t_ + i * num_classes_);
  }
  FastMatrixMatrixMultiply(1.0,
                           class_weights_,
                           false,
                           num_classes_,
                           input_dimension(),
                           x,
                           false,
                           slice.size(),
                           class_b_t_);
  activation_function_->Evaluate(num_classes_, slice.size(), class_b_t_);
  class_b_t_ += GetOffset();

  // word part
#pragma omp parallel for
  for (int i = 0; i < static_cast<int>(slice.size()); ++i) {
    const int clazz = vocabulary_->GetClass(slice[i]),
              class_size = vocabulary_->GetClassSize(clazz);
    // shortlist class?
    if (class_size > 1) {
      if (class_bias_) {
        FastCopy(word_bias_ + word_offset_[clazz],
                 class_size,
                 word_b_t_ + i * max_class_size_);
      }
      FastMatrixVectorMultiply(
          word_weights_ + word_offset_[clazz] * input_dimension(),
          false,
          class_size,
          input_dimension(),
          x + i * input_dimension(),
          word_b_t_ + i * max_class_size_);
      activation_function_->Evaluate(class_size, 1,
                                     word_b_t_ + i * max_class_size_);
    }
  }
  word_b_t_ += GetOffset();
  return result;
}

void Output::ComputeDelta(const Slice &slice, FunctionPointer f) {
  class_b_t_ -= GetOffset();
  word_b_t_ -= GetOffset();

#pragma omp parallel for
  for (int i = 0; i < static_cast<int>(slice.size()); ++i) {
    const int clazz = vocabulary_->GetClass(slice[i]),
              class_size = vocabulary_->GetClassSize(clazz);
    class_delta_t_[i * num_classes_ + clazz] = 1.;
    if (class_size == 1)
      continue;
    word_delta_t_[i * max_class_size_ + slice[i] - word_offset_[clazz] -
                  shortlist_size_] = 1.;
    activation_function_->MultiplyDerivative(
        class_size,
        1,
        word_b_t_ + i * max_class_size_,
        word_delta_t_ + i * max_class_size_);
  }
  activation_function_->MultiplyDerivative(num_classes_, slice.size(),
                                           class_b_t_, class_delta_t_);
}

void Output::AddDelta(const Slice &slice, Real delta_t[]) {
  // class part
  FastMatrixMatrixMultiply(1.0,
                           class_weights_,
                           true,
                           input_dimension(),
                           num_classes_,
                           class_delta_t_,
                           false,
                           slice.size(),
                           delta_t);
  class_delta_t_ += GetOffset();

  // word part
#pragma omp parallel for
  for (int i = 0; i < static_cast<int>(slice.size()); ++i) {
    const int clazz = vocabulary_->GetClass(slice[i]),
              class_size = vocabulary_->GetClassSize(clazz);
    // shortlist class?
    if (class_size == 1)
      continue;
    FastMatrixVectorMultiply(
        word_weights_ + word_offset_[clazz] * input_dimension(),
        true,
        class_size,  // rows of A, not op(A)!
        input_dimension(),
        word_delta_t_ + i * max_class_size_,
        delta_t + i * input_dimension());
  }
  word_delta_t_ += GetOffset();
}

const Real *Output::UpdateWeights(const Slice &slice,
                                  const Real learning_rate,
                                  const Real x[]) {
  const Real *result = class_b_t_;
  // class part
  class_delta_t_ -= GetOffset();
  if (class_bias_) {
    for (size_t i = 0; i < slice.size(); ++i) {
      FastMultiplyByConstantAdd(-learning_rate,
                                class_delta_t_ + i * num_classes_,
                                num_classes_,
                                momentum_class_bias_);
    }
  }
  FastMatrixMatrixMultiply(-learning_rate,
                           class_delta_t_,
                           false,
                           num_classes_,
                           slice.size(),
                           x,
                           true,
                           input_dimension(),
                           momentum_class_weights_);
  class_b_t_ += GetOffset();

  // word part
  word_delta_t_ -= GetOffset();
  for (size_t i = 0; i < slice.size(); ++i) {
    const int clazz = vocabulary_->GetClass(slice[i]),
              class_size = vocabulary_->GetClassSize(clazz);
    if (class_size == 1)
      continue;
    if (class_bias_) {
      FastMultiplyByConstantAdd(-learning_rate,
                                word_delta_t_ + i * max_class_size_,
                                class_size,
                                word_bias_ + word_offset_[clazz]);
    }
    FastOuterProduct(
        -learning_rate,
        word_delta_t_ + i * max_class_size_,
        class_size,
        x + i * input_dimension(),
        input_dimension(),
        word_weights_ + word_offset_[clazz] * input_dimension());
  }
  word_b_t_ += GetOffset();
  return result;
}

void Output::UpdateMomentumWeights(const Real momentum) {
  FastAdd(momentum_class_weights_,
          num_classes_ * input_dimension(),
          class_weights_,
          class_weights_);
  FastMultiplyByConstant(momentum_class_weights_,
                         num_classes_ * input_dimension(),
                         momentum,
                         momentum_class_weights_);
  if (class_bias_) {
    FastAdd(momentum_class_bias_,
            num_classes_,
            class_bias_,
            class_bias_);
    FastMultiplyByConstant(momentum_class_bias_,
                           num_classes_,
                           momentum,
                           momentum_class_bias_);
  }
}

void Output::ResetMomentum() {
  FastZero(num_classes_ * input_dimension(), momentum_class_weights_);
  FastZero(num_classes_, momentum_class_bias_);
}

void Output::Reset(const bool is_dependent) {
  class_b_t_ = class_b_;
  class_delta_t_ = class_delta_;
  word_b_t_ = word_b_;
  word_delta_t_ = word_delta_;

  FastZero((num_classes_ + max_class_size_) * max_batch_size() *
           max_sequence_length(), class_b_);
  FastZero((num_classes_ + max_class_size_) * max_batch_size() *
           max_sequence_length(), class_delta_);
}

void Output::RandomizeWeights(Random *random) {
//  const Real sigma = 1. / sqrt(input_dimension());
  const Real sigma = 0.1;
  random->ComputeGaussianRandomNumbers(num_classes_ * input_dimension(),
                                       0.,
                                       sigma,
                                       class_weights_);
  random->ComputeGaussianRandomNumbers(
      num_out_of_shortlist_words_ * input_dimension(),
      0.,
      sigma,
      word_weights_);
  FastZero(num_classes_ * input_dimension(), momentum_class_weights_);
  if (class_bias_) {
    random->ComputeGaussianRandomNumbers(num_classes_,
                                         0.,
                                         sigma,
                                         class_bias_);
    random->ComputeGaussianRandomNumbers(num_out_of_shortlist_words_,
                                         0.,
                                         sigma,
                                         word_bias_);
    FastZero(num_classes_, momentum_class_bias_);
  }
}

void Output::Read(std::ifstream *input_stream) {
  input_stream->read(reinterpret_cast<char *>(class_weights_),
                     num_classes_ * input_dimension() * sizeof(Real));
  input_stream->read(reinterpret_cast<char *>(momentum_class_weights_),
                     num_classes_ * input_dimension() * sizeof(Real));
  if (num_out_of_shortlist_words_ > 0) {
    input_stream->read(
        reinterpret_cast<char *>(word_weights_),
        num_out_of_shortlist_words_ * input_dimension() * sizeof(Real));
  }
  if (class_bias_) {
    input_stream->read(reinterpret_cast<char *>(class_bias_),
                       num_classes_ * sizeof(Real));
    input_stream->read(reinterpret_cast<char *>(momentum_class_bias_),
                       num_classes_ * sizeof(Real));
    if (num_out_of_shortlist_words_ > 0) {
      input_stream->read(reinterpret_cast<char *>(word_bias_),
                         num_out_of_shortlist_words_ * sizeof(Real));
    }
  }
}

void Output::Write(std::ofstream *output_stream) {
  output_stream->write(reinterpret_cast<char *>(class_weights_), num_classes_ *
                       input_dimension() * sizeof(Real));
  output_stream->write(reinterpret_cast<char *>(momentum_class_weights_),
                       num_classes_ * input_dimension() * sizeof(Real));
  if (num_out_of_shortlist_words_ > 0) {
    output_stream->write(reinterpret_cast<char *>(
        word_weights_),
        num_out_of_shortlist_words_ * input_dimension() * sizeof(Real));
  }
  if (class_bias_) {
    output_stream->write(reinterpret_cast<char *>(class_bias_),
                         num_classes_ * sizeof(Real));
    output_stream->write(reinterpret_cast<char *>(momentum_class_bias_),
                         num_classes_ * sizeof(Real));
    if (num_out_of_shortlist_words_ > 0) {
      output_stream->write(reinterpret_cast<char *>(word_bias_),
                           num_out_of_shortlist_words_ * sizeof(Real));
    }
  }
}

Real Output::ComputeLogProbability(const Slice &slice,
                                   const Real x[],
                                   const bool verbose,
                                   ProbabilitySequenceVector *probabilities) {
  Real log_probability = 0.;
  if (!slice.empty()) {
    const int offset = num_classes_ * max_batch_size();
    for (size_t i = 0; i < slice.size(); ++i) {
      const int clazz = vocabulary_->GetClass(slice[i]),
                class_size = vocabulary_->GetClassSize(clazz);
      Real probability = x[num_classes_ * i + clazz];
/*
      assert(clazz < num_classes_);
      assert(clazz >= 0);
      const double z = std::accumulate(x + num_classes_ * i, x + num_classes_ * (i + 1), 0.);
      assert(z > 0.99999 && z < 1.00001);
*/
      if (vocabulary_->HasUnk() &&
          slice[i] == vocabulary_->GetIndex(vocabulary_->unk()))
        probability /= num_oovs_ + 1.;  // <unk> may represent multiple words
      log_probability += log(probability);
      if (class_size > 1) {
        const Real word_probability = x[offset + max_class_size_ * i +
                                        slice[i] - word_offset_[clazz] -
                                        shortlist_size_];
/*
        const int start = offset + max_class_size_ * i,
                  word = slice[i] - word_offset_[clazz] - shortlist_size_;
        assert(word >= 0);
        assert(word < max_class_size_);
        const double z = std::accumulate(x + start, x + start + max_class_size_, 0.);
        assert(z > 0.99999 && z < 1.00001);
*/
        probability *= word_probability;
        log_probability += log(word_probability);
      }
      if (probabilities) {
        probabilities->at(i).push_back(probability);
      }
      if (verbose) {
        std::cout << "\tp( " << vocabulary_->GetWord(slice[i]) <<
                     " | ... ) \t = [1gram] " << std::setprecision(8) <<
                     probability << " [ " << std::setprecision(5) <<
                     log(probability) / log(10.) << " ]\n";
      }
    }
  }
  return log_probability;
}


#ifndef BENCHMARK_EXECUTOR_HPP_
#define BENCHMARK_EXECUTOR_HPP_

#include "application.hpp"
#include "benchmark_data.hpp"
#include "types.hpp"

#include <boost/test/unit_test.hpp>

#include <type_traits>

#include <cmath>

namespace gearshifft {
  /**
   * Benchmark body with NR_RUNS repetitions of iFFT(FFT()) implementation.
   * Implementation is given with TFunctor.
   * Depending on TFunctor::InputIsReal it uses RealType or ComplexType test data.
   * FFT output will be [normalized and] compared to original input.
   */
  template<typename T_Context,
           typename T_FFT_Wrapper,
           typename T_FFT_Normalized,
           typename T_Precision,
           typename T_Extents
           >
  struct BenchmarkExecutor {
    using ApplicationT = Application<T_Context>;
    using ResultT = typename ApplicationT::ResultT;
    static constexpr int NR_RUNS = ApplicationT::NR_RUNS; // includes warmups
    static constexpr double ERROR_BOUND = ApplicationT::ERROR_BOUND;
    static constexpr size_t NDim = std::tuple_size<T_Extents>::value;
    using VectorT = typename std::conditional<T_FFT_Wrapper::IsComplex,
                                              typename BenchmarkData<T_Precision,NDim>::ComplexVector,
                                              typename BenchmarkData<T_Precision,NDim>::RealVector>::type;
    static_assert(NDim<=3,"NDim<=3");

    void operator()(const T_Extents& extents) {
      const auto& dataset = BenchmarkData<T_Precision,NDim>::data(extents);

      VectorT data_buffer;
      dataset.copyTo(data_buffer);
      assert(data_buffer.data());

      auto fft = T_FFT_Wrapper();
      ResultT result;
      result.template init<T_FFT_Wrapper::IsComplex,
                           T_FFT_Wrapper::IsInplace,
                           T_Precision >
                       (extents);

      int r;
      if(result.getID()==0) {
        // hidden initial warmup, which is not recorded (i.e. run=0 will be overwritten)
        for(r=0; r<2*NR_RUNS; ++r) {
          dataset.copyTo(data_buffer);
          fft(result, data_buffer, extents);
        }
      }

      try {
        const double error_bound = ERROR_BOUND<0.0 ? ErrorBound<T_Precision>()() : ERROR_BOUND;
        for(r=0; r<NR_RUNS; ++r)
        {
          result.setRun(r);
          dataset.copyTo(data_buffer);
          fft(result, data_buffer, extents);

          double deviation = 0.0; // sample standard deviation
          size_t mismatches = 0; // nr of mismatches
          // compute deviation and mismatches
          dataset.template check_deviation<!T_FFT_Normalized::value>
            (deviation, mismatches, data_buffer, error_bound);

          result.setValue(RecordType::Deviation, deviation);
          result.setValue(RecordType::Mismatches, static_cast<double>(mismatches));

          if(std::isnan(deviation) || deviation>error_bound) {
            std::stringstream msg;
            msg << "mismatches=" << mismatches
                << " deviation=" << deviation
                << " errorbound=" << error_bound;
            result.setError(r, msg.str());
            throw std::runtime_error(msg.str());
          }
        }
      } catch(const std::runtime_error& e) {
        result.setError(r, e.what());
        ApplicationT::getInstance().addRecord(result);
        BOOST_FAIL( e.what() );
      }

      ApplicationT::getInstance().addRecord(result);
    }
  };

} // gearshifft

#endif /* FIXTURE_BENCHMARK_HPP_ */


#pragma once

#include <libdevcore/Address.h>
#include <libdevcore/FixedHash.h>
#include <libdevcrypto/Common.h>

#include <boost/asio.hpp>
#include <boost/multiprecision/cpp_int.hpp>
#include <chrono>
#include <filesystem>
#include <string>
#include <type_traits>

namespace taraxa {

namespace fs = std::filesystem;

using dev::Address;
using dev::AddressSet;
using dev::bytes;
using dev::bytesConstRef;
using dev::h256;
using dev::h256Hash;
using dev::h256s;
using dev::h64;
using dev::Secret;
using dev::u256;

using EthBlockNumber = uint64_t;

// time related
using time_point_t = std::chrono::system_clock::time_point;
using time_stamp_t = unsigned long;

using uint128_t = boost::multiprecision::uint128_t;
using uint256_t = boost::multiprecision::uint256_t;
using uint512_t = boost::multiprecision::uint512_t;
using uint1024_t = boost::multiprecision::uint1024_t;

// newtork related
using end_point_udp_t = boost::asio::ip::udp::endpoint;
using socket_udp_t = boost::asio::ip::udp::socket;
using resolver_udp_t = boost::asio::ip::udp::resolver;

using uint256_hash_t = dev::FixedHash<32>;
using uint512_hash_t = dev::FixedHash<64>;
using uint520_hash_t = dev::FixedHash<65>;
using uint160_hash_t = dev::FixedHash<20>;

using secret_t = dev::SecureFixedHash<32>;
using public_t = uint512_hash_t;
using addr_t = uint160_hash_t;
using sig_t = uint520_hash_t;
using vote_hash_t = uint256_hash_t;
using blk_hash_t = uint256_hash_t;
using trx_hash_t = uint256_hash_t;
using sig_hash_t = uint256_hash_t;

using gas_t = uint64_t;
using key_t = std::string;
using level_t = uint64_t;
using val_t = dev::u256;
using root_t = dev::h256;
using dag_blk_num_t = uint64_t;

using vec_blk_t = std::vector<blk_hash_t>;
using vec_trx_t = std::vector<trx_hash_t>;
using trx_num_t = vec_trx_t::size_type;
using byte = uint8_t;
using bytes = std::vector<byte>;
using node_id_t = uint512_hash_t;
using trx_nonce_t = val_t;

// val_t type related helper functions
inline val_t operator+=(val_t const &val, val_t const &other) { return val + other; }
inline std::string toString(val_t const &val) {
  std::stringstream strm;
  strm << val;
  return strm.str();
}

inline bytes str2bytes(std::string const &str) {
  assert(str.size() % 2 == 0);
  bytes data;
  // convert it to byte stream
  for (size_t i = 0; i < str.size(); i += 2) {
    std::string s = str.substr(i, 2);
    auto t = std::stoul(s, nullptr, 16);
    assert(t < 256);
    data.push_back(static_cast<uint8_t>(t));
  }
  return data;
}

inline std::string bytes2str(bytes const &data) {
  // convert it to str
  std::stringstream ss;
  ss << std::hex << std::noshowbase << std::setfill('0');
  for (auto const &d : data) {
    ss << std::setfill('0') << std::setw(2) << unsigned(d);
  }
  return ss.str();
}

}  // namespace taraxa


#define BOOST_TEST_MODULE "test_bond_length_interaction"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/forcefield/MultipleBasin/MultipleBasinForceField.hpp>
#include <mjolnir/forcefield/MultipleBasin/MultipleBasin2BasinUnit.hpp>
#include <mjolnir/forcefield/MultipleBasin/MultipleBasin3BasinUnit.hpp>
#include <mjolnir/forcefield/local/BondLengthInteraction.hpp>
#include <mjolnir/forcefield/local/HarmonicPotential.hpp>
#include <mjolnir/util/make_unique.hpp>

#include <random>

BOOST_AUTO_TEST_CASE(MultipleBasin_2Basin_numerical_difference)
{
    mjolnir::LoggerManager::set_default_logger("test_multiple_basin_forcefield.log");
    using traits_type      = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type        = traits_type::real_type;
    using coord_type       = traits_type::coordinate_type;
    using boundary_type    = traits_type::boundary_type;
    using system_type      = mjolnir::System<traits_type>;
    using potential_type   = mjolnir::HarmonicPotential<real_type>;
    using interaction_type = mjolnir::BondLengthInteraction<traits_type, potential_type>;

    const real_type k(10.0);
    const real_type native1(1.5);
    const real_type native2(2.0);

    mjolnir::LocalForceField<traits_type>    loc1, loc2, locc;
    mjolnir::GlobalForceField<traits_type>   glo1, glo2, gloc;
    mjolnir::ExternalForceField<traits_type> ext1, ext2, extc;

    std::vector<std::pair<std::array<std::size_t, 2>, potential_type>> param1, param2;
    param1.emplace_back(std::array<std::size_t, 2>{{0,1}}, potential_type(k, native1));
    param2.emplace_back(std::array<std::size_t, 2>{{0,1}}, potential_type(k, native2));

    loc1.emplace(mjolnir::make_unique<interaction_type>("none", std::move(param1)));
    loc2.emplace(mjolnir::make_unique<interaction_type>("none", std::move(param2)));

    std::unique_ptr<mjolnir::MultipleBasinUnitBase<traits_type>> unit1 =
        mjolnir::make_unique<mjolnir::MultipleBasin2BasinUnit<traits_type>>(
            -10.0, "short", "long", 0.0, 0.0,
            std::make_tuple(std::move(loc1), std::move(glo1), std::move(ext1)),
            std::make_tuple(std::move(loc2), std::move(glo2), std::move(ext2)));

    std::vector<std::unique_ptr<mjolnir::MultipleBasinUnitBase<traits_type>>> units;
    units.push_back(std::move(unit1));

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<real_type> uni(-1.0, 1.0);

    mjolnir::MultipleBasinForceField<traits_type> forcefield(
            std::make_tuple(std::move(locc), std::move(gloc), std::move(extc)),
            std::move(units));

    for(std::size_t i = 0; i < 1000; ++i)
    {
        system_type sys(2, boundary_type{});

        sys.at(0).mass  = 1.0;
        sys.at(1).mass  = 1.0;
        sys.at(0).rmass = 1.0;
        sys.at(1).rmass = 1.0;

        sys.at(0).position = coord_type( 0.0 + 0.01 * uni(mt), 0.0 + 0.01 * uni(mt), 0.0 + 0.01 * uni(mt));
        sys.at(1).position = coord_type( 1.0 + 0.01 * uni(mt), 1.0 + 0.01 * uni(mt), 1.0 + 0.01 * uni(mt));
        sys.at(0).velocity = coord_type( 0.0, 0.0, 0.0);
        sys.at(1).velocity = coord_type( 0.0, 0.0, 0.0);
        sys.at(0).force    = coord_type( 0.0, 0.0, 0.0);
        sys.at(1).force    = coord_type( 0.0, 0.0, 0.0);

        sys.at(0).name  = "X";
        sys.at(1).name  = "X";
        sys.at(0).group = "TEST";
        sys.at(1).group = "TEST";

        const auto init = sys;

        forcefield.initialize(init); // don't forget this

        constexpr real_type tol = 1e-4;
        constexpr real_type dr  = 1e-5;
        for(std::size_t idx=0; idx<2; ++idx)
        {
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = forcefield.calc_energy(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc F(x)
                forcefield.calc_force(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = forcefield.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::X(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = forcefield.calc_energy(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc F(x)
                forcefield.calc_force(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = forcefield.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Y(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = forcefield.calc_energy(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc F(x)
                forcefield.calc_force(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = forcefield.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Z(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
        }
    }
}

BOOST_AUTO_TEST_CASE(MultipleBasin_3Basin_numerical_difference)
{
    mjolnir::LoggerManager::set_default_logger("test_multiple_basin_forcefield.log");
    using traits_type      = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type        = traits_type::real_type;
    using coord_type       = traits_type::coordinate_type;
    using boundary_type    = traits_type::boundary_type;
    using system_type      = mjolnir::System<traits_type>;
    using potential_type   = mjolnir::HarmonicPotential<real_type>;
    using interaction_type = mjolnir::BondLengthInteraction<traits_type, potential_type>;

    const real_type k(10.0);
    const real_type native1(1.6);
    const real_type native2(1.8);
    const real_type native3(1.9);

    mjolnir::LocalForceField<traits_type>    loc1, loc2, loc3, locc;
    mjolnir::GlobalForceField<traits_type>   glo1, glo2, glo3, gloc;
    mjolnir::ExternalForceField<traits_type> ext1, ext2, ext3, extc;

    std::vector<std::pair<std::array<std::size_t, 2>, potential_type>> param1, param2, param3;
    param1.emplace_back(std::array<std::size_t, 2>{{0,1}}, potential_type(k, native1));
    param2.emplace_back(std::array<std::size_t, 2>{{0,1}}, potential_type(k, native2));
    param3.emplace_back(std::array<std::size_t, 2>{{0,1}}, potential_type(k, native3));

    loc1.emplace(mjolnir::make_unique<interaction_type>("none", std::move(param1)));
    loc2.emplace(mjolnir::make_unique<interaction_type>("none", std::move(param2)));
    loc3.emplace(mjolnir::make_unique<interaction_type>("none", std::move(param3)));

    std::unique_ptr<mjolnir::MultipleBasinUnitBase<traits_type>> unit1 =
        mjolnir::make_unique<mjolnir::MultipleBasin3BasinUnit<traits_type>>(
            "short", "middle", "long", -10.0, -10.0, -10.0, 0.0, 0.0, 0.0,
            std::make_tuple(std::move(loc1), std::move(glo1), std::move(ext1)),
            std::make_tuple(std::move(loc2), std::move(glo2), std::move(ext2)),
            std::make_tuple(std::move(loc3), std::move(glo3), std::move(ext3)));

    std::vector<std::unique_ptr<mjolnir::MultipleBasinUnitBase<traits_type>>> units;
    units.push_back(std::move(unit1));

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<real_type> uni(-1.0, 1.0);

    mjolnir::MultipleBasinForceField<traits_type> forcefield(
            std::make_tuple(std::move(locc), std::move(gloc), std::move(extc)),
            std::move(units));

    for(std::size_t i = 0; i < 1000; ++i)
    {
        system_type sys(2, boundary_type{});

        sys.at(0).mass  = 1.0;
        sys.at(1).mass  = 1.0;
        sys.at(0).rmass = 1.0;
        sys.at(1).rmass = 1.0;

        sys.at(0).position = coord_type( 0.0 + 0.01 * uni(mt), 0.0 + 0.01 * uni(mt), 0.0 + 0.01 * uni(mt));
        sys.at(1).position = coord_type( 1.0 + 0.01 * uni(mt), 1.0 + 0.01 * uni(mt), 1.0 + 0.01 * uni(mt));
        sys.at(0).velocity = coord_type( 0.0, 0.0, 0.0);
        sys.at(1).velocity = coord_type( 0.0, 0.0, 0.0);
        sys.at(0).force    = coord_type( 0.0, 0.0, 0.0);
        sys.at(1).force    = coord_type( 0.0, 0.0, 0.0);

        sys.at(0).name  = "X";
        sys.at(1).name  = "X";
        sys.at(0).group = "TEST";
        sys.at(1).group = "TEST";

        const auto init = sys;

        forcefield.initialize(init); // don't forget this

        constexpr real_type tol = 1e-4;
        constexpr real_type dr  = 1e-5;
        for(std::size_t idx=0; idx<2; ++idx)
        {
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = forcefield.calc_energy(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc F(x)
                forcefield.calc_force(sys);

                mjolnir::math::X(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = forcefield.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::X(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = forcefield.calc_energy(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc F(x)
                forcefield.calc_force(sys);

                mjolnir::math::Y(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = forcefield.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Y(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
            {
                // ----------------------------------------------------------------
                // reset positions
                sys = init;

                // calc U(x-dx)
                const auto E0 = forcefield.calc_energy(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc F(x)
                forcefield.calc_force(sys);

                mjolnir::math::Z(sys.position(idx)) += dr;

                // calc U(x+dx)
                const auto E1 = forcefield.calc_energy(sys);

                // central difference
                const auto dE = (E1 - E0) * 0.5;

                BOOST_TEST(-dE == dr * mjolnir::math::Z(sys.force(idx)),
                           boost::test_tools::tolerance(tol));
            }
        }
    }
}


//
// Copyright (c) 2018 CNRS
//

#include "pinocchio/multibody/model.hpp"
#include "pinocchio/multibody/data.hpp"
#include "pinocchio/algorithm/kinematics.hpp"
#include "pinocchio/algorithm/crba.hpp"
#include "pinocchio/algorithm/joint-configuration.hpp"
#include "pinocchio/parsers/sample-models.hpp"

#include <iostream>

#include <boost/test/unit_test.hpp>
#include <boost/utility/binary.hpp>

BOOST_AUTO_TEST_SUITE(BOOST_TEST_MODULE)

BOOST_AUTO_TEST_CASE(test_kinematics_zero)
{
  using namespace Eigen;
  using namespace pinocchio;

  Model model;
  buildModels::humanoidRandom(model);

  Data data(model), data_ref(model);

  model.lowerPositionLimit.head<3>().fill(-1.);
  model.upperPositionLimit.head<3>().fill(1.);
  VectorXd q = randomConfiguration(model);

  forwardKinematics(model,data_ref,q);
  crba(model,data,q);
  updateGlobalPlacements(model,data);

  for(Model::JointIndex i = 1; i < (Model::JointIndex)model.njoints; ++i)
  {
    BOOST_CHECK(data.oMi[i] == data_ref.oMi[i]);
    BOOST_CHECK(data.liMi[i] == data_ref.liMi[i]);
  }
}

BOOST_AUTO_TEST_CASE(test_kinematics_first)
{
  using namespace Eigen;
  using namespace pinocchio;

  Model model;
  buildModels::humanoidRandom(model);

  Data data(model);

  model.lowerPositionLimit.head<3>().fill(-1.);
  model.upperPositionLimit.head<3>().fill(1.);
  VectorXd q = randomConfiguration(model);
  VectorXd v(VectorXd::Zero(model.nv));

  forwardKinematics(model,data,q,v);

  for(Model::JointIndex i = 1; i < (Model::JointIndex)model.njoints; ++i)
  {
    BOOST_CHECK(data.v[i] == Motion::Zero());
  }
}

BOOST_AUTO_TEST_CASE(test_kinematics_second)
{
  using namespace Eigen;
  using namespace pinocchio;

  Model model;
  buildModels::humanoidRandom(model);

  Data data(model);

  model.lowerPositionLimit.head<3>().fill(-1.);
  model.upperPositionLimit.head<3>().fill(1.);
  VectorXd q = randomConfiguration(model);
  VectorXd v(VectorXd::Zero(model.nv));
  VectorXd a(VectorXd::Zero(model.nv));

  forwardKinematics(model,data,q,v,a);

  for(Model::JointIndex i = 1; i < (Model::JointIndex)model.njoints; ++i)
  {
    BOOST_CHECK(data.v[i] == Motion::Zero());
    BOOST_CHECK(data.a[i] == Motion::Zero());
  }
}

BOOST_AUTO_TEST_SUITE_END()


#ifndef RHF_HPP
#define RHF_HPP

#include "EigenSolver.hpp"
#include "Integrals.hpp"
#include <armadillo>

namespace willow {
namespace qcmol {

class RHF {
public:
  RHF(const std::vector<libint2::Atom> &atoms, const libint2::BasisSet &bs,
      const Integrals &ints, const bool l_print = true,
      const bool l_grad = false, const int qm_chg = 0,
      const std::vector<QAtom> &atoms_Q = std::vector<QAtom>());

  arma::mat densityMatrix();
  arma::mat energyWeightDensityMatrix();
  arma::mat g_matrix(const arma::vec &TEI, arma::mat &Dm);

  template <libint2::Operator obtype>
  arma::mat compute_1body_ints_deriv(const libint2::BasisSet &bs,
                                     const std::vector<libint2::Atom> &atoms,
                                     const arma::mat &Dm);

  arma::mat compute_1body_nuclear_deriv(const libint2::BasisSet &bs,
                                        const std::vector<libint2::Atom> &atoms,
                                        const arma::mat &Dm,
                                        const std::vector<QAtom> &atoms_Q);

  arma::mat compute_nuclear_gradient(const std::vector<libint2::Atom> &atoms,
                                     const std::vector<QAtom> &atoms_Q);

  arma::mat compute_2body_ints_deriv(const libint2::BasisSet &bs,
                                     const std::vector<libint2::Atom> &atoms,
                                     const arma::mat &Dm,
                                     const arma::mat &Schwartz);

  double rhf_energy() { return (E_nuc + E_hf); };
  double get_nocc() { return m_nocc; };

  arma::mat get_gradient() { return m_grad; };
  virtual double get_energy() { return rhf_energy(); };

  EigenSolver eig_solver;

protected:
  arma::mat m_grad;

private:
  size_t m_nocc;
  double E_nuc;
  double E_hf;
};

} // namespace qcmol
} // namespace willow

#endif


#ifndef NETKET_IMAGINARY_TIME_PROPAGATION_HPP
#define NETKET_IMAGINARY_TIME_PROPAGATION_HPP

#include <Eigen/Dense>

#include "Dynamics/TimeStepper/abstract_time_stepper.hpp"
#include "Operator/MatrixWrapper/matrix_wrapper.hpp"
#include "Operator/operator.hpp"
#include "Output/json_output_writer.hpp"
#include "Stats/stats.hpp"

namespace netket {

class ImagTimePropagation {
 public:
  using StateVector = Eigen::VectorXcd;
  using Stepper = ode::AbstractTimeStepper<StateVector>;
  using Matrix = AbstractMatrixWrapper<>;

  using ObsEntry = std::pair<std::string, std::unique_ptr<Matrix>>;

  ImagTimePropagation(const AbstractOperator& hamiltonian, Stepper& stepper,
                      double t0, StateVector initial_state,
                      const std::string& matrix_type = "sparse")
      : matrix_(CreateMatrixWrapper<>(hamiltonian, matrix_type)),
        stepper_(stepper),
        t_(t0),
        state_(std::move(initial_state)) {
    ode_system_ = [this](const StateVector& x, StateVector& dxdt,
                         double /*t*/) { dxdt.noalias() = -matrix_->Apply(x); };
  }

  void AddObservable(const AbstractOperator& observable,
                     const std::string& name,
                     const std::string& matrix_type = "sparse") {
    auto wrapper = CreateMatrixWrapper(observable, matrix_type);
    observables_.emplace_back(name, std::move(wrapper));
  }

  void Advance(double dt) {
    // Propagate the state
    stepper_.Propagate(ode_system_, state_, t_, dt);
    // renormalize the state to prevent unbounded growth of the norm
    state_.normalize();
    ComputeObservables(state_);
    t_ += dt;
  }

  void ComputeObservables(const StateVector& state) {
    const auto mean_variance = matrix_->MeanVariance(state);
    obsmanager_.Reset("Energy");
    obsmanager_.Push("Energy", mean_variance[0].real());
    obsmanager_.Reset("EnergyVariance");
    obsmanager_.Push("EnergyVariance", mean_variance[1].real());

    for (const auto& entry : observables_) {
      const auto& name = entry.first;
      const auto& obs = entry.second;
      obsmanager_.Reset(name);

      const auto value = obs->Mean(state).real();
      obsmanager_.Push(name, value);
    }
  }

  const ObsManager& GetObsManager() const { return obsmanager_; }

  double GetTime() const { return t_; }
  void SetTime(double t) { t_ = t; }

 private:
  std::unique_ptr<Matrix> matrix_;
  Stepper& stepper_;
  ode::OdeSystemFunction<StateVector> ode_system_;

  double t_;
  StateVector state_;

  std::vector<ObsEntry> observables_;
  ObsManager obsmanager_;
};

}  // namespace netket

#endif  // NETKET_IMAGINARY_TIME_PROPAGATION_HPP