/*
  Copyright (C) 2018 David Robillard <d@drobilla.net>
  Copyright (C) 2006-2007 Chris Hamilton <chamilton@cs.dal.ca>

  This program is free software: you can redistribute it and/or modify it under
  the terms of the GNU General Public License as published by the Free Software
  Foundation, either version 2 of the License, or (at your option) any later
  version.

  This program is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
  details.

  You should have received a copy of the GNU General Public License along with
  this program.  If not, see <https://www.gnu.org/licenses/>.
*/

#ifndef CHILBERT_BIGBITVEC_HPP
#define CHILBERT_BIGBITVEC_HPP

#include "chilbert/Operations.hpp"

#include <algorithm>
#include <cassert>
#include <climits>
#include <cstddef>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <memory>

namespace chilbert {

class CBigBitVec
{
public:
	using Rack = uintptr_t;

	static constexpr size_t bits_per_rack = sizeof(Rack) * CHAR_BIT;

	/** Mask for a bit that can be incremented like an index.
	 *
	 * This enables fast iteration while setting or resetting bits since it is
	 * internally stored as a mask which avoids repeated index math.
	 */
	class Mask
	{
	public:
		void operator++()
		{
			if ((m_mask <<= 1) == 0) {
				m_mask = 1;
				++m_rack;
			}
		}

		void operator--()
		{
			if ((m_mask >>= 1) == 0) {
				m_mask = Rack{1} << (bits_per_rack - 1);
				--m_rack;
			}
		}

		bool operator==(const Mask& mask) const
		{
			return m_rack == mask.m_rack && m_mask == mask.m_mask;
		}

		bool operator!=(const Mask& mask) const { return !operator==(mask); }

	private:
		friend class CBigBitVec;

		Mask(const size_t index)
		  : m_rack{index / bits_per_rack}
		  , m_mask{Rack{1} << (index - m_rack * bits_per_rack)}
		{
		}

		size_t m_rack;
		Rack   m_mask;
	};

	explicit CBigBitVec(const size_t bits)
	  : m_pcRacks{make_racks(num_racks(bits))}
	  , m_iRacks{bits == 0 ? 0 : num_racks(bits)}
	{
	}

	CBigBitVec(const size_t bits, const Rack value)
	  : CBigBitVec{bits}
	{
		m_pcRacks[0] = value;
	}

	CBigBitVec(const CBigBitVec& vec)
	  : m_pcRacks{make_racks(vec.m_iRacks)}
	  , m_iRacks{vec.m_iRacks}
	{
		if (vec.m_pcRacks) {
			memcpy(
			  m_pcRacks.get(), vec.m_pcRacks.get(), sizeof(Rack) * m_iRacks);
		}
	}

	CBigBitVec(CBigBitVec&& vec) = default;

	/// Return the size in bits
	size_t size() const { return m_iRacks * bits_per_rack; }

	/// Set all bits to zero
	CBigBitVec& reset()
	{
		memset(m_pcRacks.get(), 0, sizeof(Rack) * m_iRacks);
		return *this;
	}

	/// Set all bits to one
	CBigBitVec& set()
	{
		memset(m_pcRacks.get(), 0xFF, sizeof(Rack) * m_iRacks);
		return *this;
	}

	/// Truncate to a given precision in bits (zero MSBs)
	CBigBitVec& truncate(const size_t bits)
	{
		const Mask m = mask(bits);

		m_pcRacks[m.m_rack] &= (m.m_mask - 1);

		for (size_t i = m.m_rack + 1; i < m_iRacks; ++i) {
			m_pcRacks[i] = 0;
		}

		return *this;
	}

	bool operator==(const CBigBitVec& vec) const
	{
		return (rackCount() == vec.rackCount() &&
		        (rackCount() == 0 ||
		         !memcmp(racks(), vec.racks(), rackCount() * sizeof(Rack))));
	}

	bool operator!=(const CBigBitVec& vec) const { return !(*this == vec); }

	bool operator<(const CBigBitVec& vec) const
	{
		assert(size() == vec.size());

		for (size_t ri = 0; ri < m_iRacks; ++ri) {
			size_t i = m_iRacks - ri - 1;
			if (m_pcRacks[i] < vec.m_pcRacks[i]) {
				return true;
			} else if (m_pcRacks[i] > vec.m_pcRacks[i]) {
				return false;
			}
		}
		return false;
	}

	/// Non-resizing assignment
	CBigBitVec& operator=(const CBigBitVec& vec)
	{
		if (m_iRacks < vec.m_iRacks) {
			m_iRacks  = vec.m_iRacks;
			m_pcRacks = make_racks(m_iRacks);
			memcpy(
			  m_pcRacks.get(), vec.m_pcRacks.get(), sizeof(Rack) * m_iRacks);
		} else if (vec.m_iRacks > 0) {
			m_iRacks = vec.m_iRacks;
			memcpy(
			  m_pcRacks.get(), vec.m_pcRacks.get(), sizeof(Rack) * m_iRacks);
		} else {
			m_iRacks = 0;
			m_pcRacks.reset();
		}

		return *this;
	}

	CBigBitVec& operator=(CBigBitVec&& vec) = default;

	/// Return the value of the bit covered by `mask`
	bool test(const Mask mask) const
	{
		return m_pcRacks[mask.m_rack] & mask.m_mask;
	}

	/// Return the value of the `index`th bit
	bool test(const size_t index) const { return test(mask(index)); }

	/// Set the bit covered by `mask` to 1
	CBigBitVec& set(const Mask mask)
	{
		m_pcRacks[mask.m_rack] |= mask.m_mask;
		return *this;
	}

	/// Set the `index`th bit to 1
	CBigBitVec& set(const size_t index) { return set(mask(index)); }

	/// Reset the bit covered by `mask` to 0
	CBigBitVec& reset(const Mask mask)
	{
		m_pcRacks[mask.m_rack] &= ~mask.m_mask;
		return *this;
	}

	/// Reset the `index`th bit to 0
	CBigBitVec& reset(const size_t index) { return reset(mask(index)); }

	/// Set the bit covered by `mask` to `value`
	CBigBitVec& set(const Mask mask, const bool value)
	{
		auto& rack = m_pcRacks[mask.m_rack];
		rack ^= (-Rack{value} ^ rack) & mask.m_mask;
		return *this;
	}

	/// Set the `index`th bit to `value`
	CBigBitVec& set(const size_t index, const bool value)
	{
		return set(mask(index), value);
	}

	/// Flip the value of the bit covered by `mask`
	CBigBitVec& flip(const Mask mask)
	{
		m_pcRacks[mask.m_rack] ^= mask.m_mask;
		return *this;
	}

	/// Flip the value of the `index`th bit
	CBigBitVec& flip(const size_t index) { return flip(mask(index)); }

	CBigBitVec& operator&=(const CBigBitVec& vec)
	{
		for (size_t i = 0; i < std::min(m_iRacks, vec.m_iRacks); ++i) {
			m_pcRacks[i] &= vec.m_pcRacks[i];
		}

		return *this;
	}

	CBigBitVec operator&(const CBigBitVec& vec) const
	{
		CBigBitVec t(*this);
		t &= vec;

		return t;
	}

	CBigBitVec& operator|=(const CBigBitVec& vec)
	{
		for (size_t i = 0; i < std::min(m_iRacks, vec.m_iRacks); ++i) {
			m_pcRacks[i] |= vec.m_pcRacks[i];
		}

		return *this;
	}

	CBigBitVec operator|(const CBigBitVec& vec) const
	{
		CBigBitVec t(*this);
		t |= vec;

		return t;
	}

	CBigBitVec& operator^=(const CBigBitVec& vec)
	{
		for (size_t i = 0; i < std::min(m_iRacks, vec.m_iRacks); ++i) {
			m_pcRacks[i] ^= vec.m_pcRacks[i];
		}

		return *this;
	}

	CBigBitVec operator^(const CBigBitVec& vec) const
	{
		CBigBitVec t(*this);
		t ^= vec;

		return t;
	}

	CBigBitVec operator~() const
	{
		CBigBitVec t(*this);
		t.flip();
		return t;
	}

	CBigBitVec& operator<<=(const size_t bits)
	{
		if (bits == 0) {
			return *this;
		} else if (bits >= size()) {
			reset();
			return *this;
		}

		const Index index{bits};

		if (index.rack > 0) {
			// Shift entire racks
			for (size_t i = m_iRacks - 1; i >= index.rack; --i) {
				m_pcRacks[i] = m_pcRacks[i - index.rack];
			}
			for (size_t i = 0; i < index.rack; ++i) {
				m_pcRacks[i] = 0;
			}
		}

		if (index.bit > 0) {
			// Shift bits within racks
			size_t bi = bits_per_rack - index.bit;
			size_t i;
			for (i = m_iRacks - 1; i >= index.rack + 1; --i) {
				m_pcRacks[i] <<= index.bit;
				m_pcRacks[i] |= m_pcRacks[i - 1] >> bi;
			}
			m_pcRacks[i] <<= index.bit;
		}

		return *this;
	}

	CBigBitVec operator<<(const size_t bits) const
	{
		CBigBitVec t(*this);
		t <<= bits;
		return t;
	}

	CBigBitVec& operator>>=(const size_t bits)
	{
		if (bits == 0) {
			return *this;
		} else if (bits >= size()) {
			reset();
			return *this;
		}

		const Index index{bits};

		if (index.rack > 0) {
			// Shift entire racks
			size_t i;
			for (i = 0; i < m_iRacks - index.rack; ++i) {
				m_pcRacks[i] = m_pcRacks[i + index.rack];
			}
			for (; i < m_iRacks; ++i) {
				m_pcRacks[i] = 0;
			}
		}

		if (index.bit > 0) {
			// Shift bits within racks
			size_t bi = bits_per_rack - index.bit;
			size_t i;
			for (i = 0; i < m_iRacks - index.rack - 1; ++i) {
				m_pcRacks[i] >>= index.bit;
				m_pcRacks[i] |= m_pcRacks[i + 1] << bi;
			}
			m_pcRacks[i] >>= index.bit;
		}

		return *this;
	}

	CBigBitVec operator>>(const size_t bits) const
	{
		CBigBitVec t(*this);
		t >>= bits;
		return t;
	}

	/// Right-rotate the least significant `width` bits by `bits` positions
	CBigBitVec& rotr(const size_t bits, size_t width)
	{
		// Fill in the width, if necessary.
		if (width <= 0) {
			width = size();
		}

		// Modulo the number of bits.
		assert(bits < width);
		if (bits == 0) {
			return *this;
		}

		// Ensure we are truncated appropriately.
		truncate(width);

		CBigBitVec t1(*this);
		(*this) >>= bits;
		t1 <<= (width - bits);
		(*this) |= t1;

		truncate(width);

		return *this;
	}

	/// Left-rotate the least significant `width` bits by `bits` positions
	CBigBitVec& rotl(const size_t bits, size_t width)
	{
		// Fill in the width, if necessary.
		if (width <= 0) {
			width = size();
		}

		// Modulo the number of bits.
		assert(bits < width);
		if (bits == 0) {
			return *this;
		}

		// Ensure we are truncated appropriately.
		truncate(width);

		CBigBitVec t1(*this);
		(*this) <<= bits;
		t1 >>= (width - bits);
		(*this) |= t1;

		truncate(width);

		return *this;
	}

	/// Return true iff all bits are zero
	bool none() const
	{
		for (size_t i = 0; i < m_iRacks; ++i) {
			if (m_pcRacks[i]) {
				return false;
			}
		}
		return true;
	}

	/// Return 1 + the index of the first set bit, or 0 if there are none
	size_t find_first() const
	{
		for (size_t i = 0; i < m_iRacks; ++i) {
			const int j = ffs(m_pcRacks[i]);
			if (j) {
				return (i * bits_per_rack) + static_cast<size_t>(j);
			}
		}
		return 0;
	}

	/// Return the number of set bits
	size_t count() const
	{
		size_t c = 0;
		for (size_t i = 0; i < m_iRacks; ++i) {
			c += static_cast<size_t>(pop_count(m_pcRacks[i]));
		}
		return c;
	}

	/// Flip all bits (one's complement)
	CBigBitVec& flip()
	{
		for (size_t i = 0; i < m_iRacks; ++i) {
			m_pcRacks[i] = ~m_pcRacks[i];
		}
		return *this;
	}

	/// Return the first rack
	Rack& rack() { return m_pcRacks[0]; }
	Rack  rack() const { return m_pcRacks[0]; }

	/// Return a pointer to the racks
	Rack*       racks() { return m_pcRacks.get(); }
	const Rack* racks() const { return m_pcRacks.get(); }

	/// Return the number of racks
	size_t rackCount() const { return m_iRacks; }

	template <class BitVec>
	class iterator_base : public Mask
	{
	public:
		iterator_base& operator++()
		{
			Mask::operator++();
			return *this;
		}

		iterator_base& operator--()
		{
			Mask::operator--();
			return *this;
		}

		bool operator==(const iterator_base& iterator_base) const
		{
			return m_vec == iterator_base.m_vec &&
			       Mask::operator==(iterator_base);
		}

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

		bool operator*() const { return m_vec->test(*this); }

	protected:
		iterator_base(BitVec& vec, const size_t index)
		  : Mask{index}
		  , m_vec{&vec}
		{
		}

		BitVec* m_vec;
	};

	class iterator : public iterator_base<CBigBitVec>
	{
	public:
		void set() { m_vec->set(*this); }
		void reset() { m_vec->reset(*this); }

	private:
		friend class CBigBitVec;

		iterator(CBigBitVec& vec, const size_t index)
		  : iterator_base{vec, index}
		{
		}
	};

	class const_iterator : public iterator_base<const CBigBitVec>
	{
	private:
		friend class CBigBitVec;

		const_iterator(const CBigBitVec& vec, const size_t index)
		  : iterator_base{vec, index}
		{
		}
	};

	Mask mask(const size_t i = 0) const
	{
		assert(i <= size());
		return Mask{i};
	}

	iterator begin(const size_t i = 0) { return iterator(*this, i); }

	iterator end() { return iterator(*this, size()); }

	const_iterator begin(const size_t i = 0) const
	{
		return const_iterator(*this, i);
	}

	const_iterator end() const { return const_iterator(*this, size()); }

private:
	struct Index
	{
		Index(const size_t bits)
		  : rack{bits / bits_per_rack}
		  , bit{bits - rack * bits_per_rack}
		{
			assert(bit < bits_per_rack);
		}

		size_t rack;
		size_t bit;
	};

	struct RacksDeleter
	{
		void operator()(Rack* const racks) { free(racks); }
	};

	using RacksPtr = std::unique_ptr<Rack[], RacksDeleter>;

	static size_t num_racks(const size_t bits)
	{
		return (std::max(bits, size_t(1)) + bits_per_rack - 1) / bits_per_rack;
	}

	static RacksPtr make_racks(const size_t n)
	{
		return RacksPtr{static_cast<Rack*>(calloc(n, sizeof(Rack)))};
	}

	// Right rotates entire racks (in place).
	void rackRotr(const size_t k)
	{
		assert(k < m_iRacks);

		if (k == 0) {
			return;
		}

		size_t c = 0;
		for (size_t v = 0; c < m_iRacks; ++v) {
			size_t     t   = v;
			size_t     tp  = v + k;
			const Rack tmp = m_pcRacks[v];
			c++;
			while (tp != v) {
				m_pcRacks[t] = m_pcRacks[tp];
				t            = tp;
				tp += k;
				if (tp >= m_iRacks) {
					tp -= m_iRacks;
				}
				c++;
			}
			m_pcRacks[t] = tmp;
		}
	}

	RacksPtr m_pcRacks;
	size_t   m_iRacks;
};

template <>
void
grayCode(CBigBitVec& value)
{
	CBigBitVec::Rack s = 0;

	for (intptr_t i = static_cast<intptr_t>(value.rackCount() - 1); i >= 0;
	     --i) {
		const auto t = value.racks()[i] & 1;
		grayCode(value.racks()[i]);
		value.racks()[i] ^= (s << (CBigBitVec::bits_per_rack - 1));
		s = t;
	}
}

template <>
void
grayCodeInv(CBigBitVec& value)
{
	CBigBitVec::Rack s = 0;

	for (intptr_t i = static_cast<intptr_t>(value.rackCount() - 1); i >= 0;
	     --i) {
		auto& rack = value.racks()[i];
		grayCodeInv(rack);
		if (s) {
			rack = ~rack;
		}
		s = rack & 1;
	}
}

inline std::ostream&
operator<<(std::ostream& os, const CBigBitVec& vec)
{
	for (size_t i = 0; i < vec.size(); ++i) {
		os << vec.test(vec.size() - i - 1);
	}
	return os;
}

} // namespace chilbert

#endif