/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2013  Virginia Tech
 *
 *  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 3 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 <http://www.gnu.org/licenses/>.
 *
 */

#ifndef TILEDARRAY_BITSET_H__INCLUDED
#define TILEDARRAY_BITSET_H__INCLUDED

#include <TiledArray/error.h>
#include <TiledArray/transform_iterator.h>
#include <climits>
#include <iosfwd>
#include <iomanip>

namespace TiledArray {
  namespace detail {

    /// Fixed size bitset

    /// Bitset is similar to \c std::bitset except the size is set at runtime.
    /// This bit set has very limited functionality, because it is only intended
    /// to fit the needs of \c SparseShape. The block type may be an integral
    /// or char type.
    /// \tparam Block The type used to store the data [ default = \c unsigned \c long ]
    template <typename Block = unsigned long>
    class Bitset {
    private:
      static_assert( (std::is_integral<Block>::value || std::is_same<Block, char>::value),
          "Bitset template type Block must be an integral or char type");

      static const std::size_t block_bits; ///< The number of bits in a block
      static const Block zero;
      static const Block one;
      static const Block xffff;

    public:

      class reference {

        friend class Bitset<Block>;

        reference(Block& block, Block mask) : block_(block), mask_(mask) { }

        // Not allowed
        void operator&();

      public:

        reference& operator=(bool value) {
          assign(value);
          return *this;
        }

        reference& operator=(const reference& other) {
          assign(other);
          return *this;
        }

        reference& operator|=(bool value) {
          if(value) set();
          return *this;
        }

        reference& operator&=(bool value) {
          if(!value) reset();
          return *this;
        }

        reference& operator^=(bool value) {
          if(value) flip();
          return *this;
        }

        reference& operator-=(bool value) {
          if(value) reset();
          return *this;
        }

        operator bool() const { return block_ & mask_; }

        bool operator~() const { return !(block_ & mask_); }

        reference& flip() {
          block_ ^= mask_;
          return *this;
        }

     private:

        void assign(bool value) {
          if(value)
            set();
          else
            reset();
        }

        void set() { block_ |= mask_; }

        void reset() { block_ &= ~mask_; }

        Block& block_;
        const Block mask_;
      }; // class reference

    private:

      /// Operation to provide iterator access to bits
      class ConstTransformOp {
      public:
        typedef std::size_t argument_type;
        typedef Block result_type;

        ConstTransformOp(const Bitset<Block>& bitset) : bitset_(&bitset) { }

        ConstTransformOp(const ConstTransformOp& other) : bitset_(other.bitset_) { }


        ConstTransformOp& operator=(const ConstTransformOp& other) {
          bitset_ = other.bitset_;

          return *this;
        }

        Block operator()(std::size_t i) const { return (*bitset_)[i]; }

      private:
        const Bitset<Block>* bitset_;
      }; // class ConstTransformOp

      /// Operation to provide const iterator access to bits
      class TransformOp {
      public:
        typedef std::size_t argument_type;
        typedef reference result_type;

        TransformOp(const Bitset<Block>& bitset) : bitset_(bitset) { }
        reference operator()(std::size_t i) const { return const_cast<Bitset<Block>&>(bitset_)[i]; }
      private:
        const Bitset<Block>& bitset_;
      }; // class TransformOp

    public:
      typedef Bitset<Block> Bitset_; ///< This object type
      typedef Block block_type; ///< The type used to store the data
      typedef Block value_type; ///< The value type
      typedef Block const_reference; ///< Constant reference to a bit
      typedef std::size_t size_type; ///< indexing size type
      typedef UnaryTransformIterator<Block, ConstTransformOp> const_iterator; ///< Const iterator type
      typedef UnaryTransformIterator<Block, TransformOp> iterator; ///< Iterator type

      /// Construct a bitset that contains \c s bits.

      /// \param s The number of bits
      /// \throw std::bad_alloc If bitset allocation fails.
      Bitset(size_type s) :
          size_(s),
          blocks_((size_ / block_bits) + (size_ % block_bits ? 1 : 0)),
          set_((size_ ? new block_type[blocks_] : NULL))
      {
        std::fill_n(set_, blocks_, zero);
      }

      /// Construct a bitset that contains \c s bits.

      /// \tparam InIter The input iterator type
      /// \param first The first element of a set of bits to be set
      /// \param last The last element of a set of bits to be set
      /// \throw std::bad_alloc If bitset allocation fails.
      template <typename InIter>
      Bitset(InIter first, InIter last) :
          size_(std::distance(first, last)),
          blocks_((size_ / block_bits) + (size_ % block_bits ? 1 : 0)),
          set_((size_ ? new block_type[blocks_] : NULL))
      {
        // Initialize to zero
        std::fill_n(set_, blocks_, zero);

        for(size_type i = 0; first != last; ++i, ++first)
          if(*first)
            set(i);
      }

      /// Copy constructor for bitset

      /// \param other The bitset to copy
      /// \throw std::bad_alloc If bitset allocation fails.
      Bitset(const Bitset<Block>& other) :
          size_(other.size_),
          blocks_(other.blocks_),
          set_((size_ ? new block_type[blocks_] : NULL))
      {
        std::copy(other.set_, other.set_ + blocks_, set_);
      }

      /// Destructor
      ~Bitset() { delete [] set_; }

      /// Assignment operator

      /// This will only copy the data from \c other. It will not change the size
      /// of the bitset.
      /// \param other The bitset to copy
      /// \throw std::runtime_error If the bitset sizes are not equal.
      Bitset<Block>& operator=(const Bitset<Block>& other) {
        if(blocks_ == other.blocks_) {
          if(this != &other) {
            size_ = other.size_;
            std::copy(other.set_, other.set_ + blocks_, set_);
          }
        } else {
          Bitset<Block>(other).swap(*this);
        }

        return *this;
      }

      /// Or-assignment operator

      /// Or-assign all bits from the two ranges
      /// \param other The bitset to be or-assigned to this bitset
      /// \throw std::range_error If the bitset sizes are not equal.
      Bitset<Block>& operator|=(const Bitset<Block>& other) {
        TA_ASSERT(size_ == other.size_);
        for(size_type i = 0; i < blocks_; ++i)
          set_[i] |= other.set_[i];

        return *this;
      }

      /// And-assignment operator

      /// And-assign all bits from the two ranges
      /// \param other The bitset to be and-assigned to this bitset
      /// \throw std::range_error If the bitset sizes are not equal.
      Bitset<Block>& operator&=(const Bitset<Block>& other) {
        TA_ASSERT(size_ == other.size_);
        for(size_type i = 0; i < blocks_; ++i)
          set_[i] &= other.set_[i];

        return *this;
      }


      /// And-assignment operator

      /// And-assign all bits from the two ranges
      /// \param other The bitset to be and-assigned to this bitset
      /// \throw std::range_error If the bitset sizes are not equal.
      Bitset<Block>& operator^=(const Bitset<Block>& other) {
        TA_ASSERT(size_ == other.size_);
        for(size_type i = 0; i < blocks_; ++i)
          set_[i] ^= other.set_[i];

        return *this;
      }

    private:

      static void left_shift(Bitset<Block>& dest, const Bitset<Block>& source, size_type n) {
        // Compute shifts
        const size_type block_shift = dest.block_index(n);
        const size_type bit_shift = dest.bit_index(n);

        // Compute iteration ranges
        block_type* last = dest.set_ + dest.blocks_ - 1;
        const block_type* const first = dest.set_ + block_shift;
        const block_type* base = source.set_ + source.blocks_ - 1 - block_shift;

        if(bit_shift == 0) {
          // Shift by unit strides
          while(last >= first)
            *last-- = *base--;
        } else {
          // Shift by non-unit strides
          const block_type* base1 = base - 1;
          const size_type reverse_bit_shift = block_bits - bit_shift;
          while(last > first) {
            *last-- = (*base << bit_shift) | (*base1 >> reverse_bit_shift);
            base = base1--;
          }
          *last-- = (*base << bit_shift);
        }

        // Zero the head
        while(last >= dest.set_)
          *last-- = zero;

        // Zero the tail
        const size_type extra_bits = dest.size_ % block_bits;
        if (extra_bits != 0)
            dest.set_[dest.blocks_ - 1] &= ~(xffff << extra_bits);
      }

      static void right_shift(Bitset<Block>& dest, const Bitset<Block>& source, size_type n) {
        // Compute shifts
        size_type  const block_shift = block_index(n);
        size_type const bit_shift = bit_index(n);

        // Compute iterator ranges
        const block_type* base = source.set_ + block_shift;
        block_type* first = dest.set_;
        const block_type* const end = dest.set_ + dest.blocks_;
        const block_type* const last = end - 1 - block_shift;

        if (bit_shift == 0) {
          // Shift by unit strides
          while(first <= last)
              *first++ = *base++;
        } else {
          // Shift by non-unit strides
            size_type const reverse_bit_shift = block_bits - bit_shift;
            const block_type* base1 = base + 1;

            while(first < last) {
                *first++ = (*base >> bit_shift) | (*base1  << reverse_bit_shift);
                base = base1++;
            }
            *first++ = *base >> bit_shift;
        }

        // Zero the tail
        while(first < end)
          *first++ = zero;
      }


    public:

      Bitset<Block>& operator<<=(size_type n) {
        if(n >= size_)
          reset();
        else if(n > 0ul)
          left_shift(*this, *this, n);
        return *this;
      }

      Bitset<Block> operator<<(size_type n) {
        Bitset<Block> temp = Bitset<Block>(size_);
        if(n < size_)
          left_shift(temp, *this, n);
        return temp;
      }

      Bitset<Block>& operator>>=(size_type n) {
        if(n >= size_)
          reset();
        else if(n > 0ul)
          right_shift(*this, *this, n);
        return *this;
      }

      Bitset<Block> operator>>(size_type n) {
        Bitset<Block> temp = Bitset<Block>(size_);
        if(n < size_)
          right_shift(temp, *this, n);
        return temp;
      }

      /// Bit accessor operator

      /// \param i The bit to access
      /// \return The value of i-th bit of the bit set
      /// \throw std::out_of_range If \c i is greater than or equal to the size
      const_reference operator[](size_type i) const {
        TA_ASSERT(i < size_);
        return mask(i) & set_[block_index(i)];
      }

      reference operator[](size_type i) {
        TA_ASSERT(i < size_);
        return reference(set_[block_index(i)], mask(i));
      }

      operator bool() const {
        const block_type* const end = set_ + blocks_;
        for(const block_type* first = set_; first != end; ++first)
          if(*first)
            return true;

        return false;
      }

      bool operator!() const {
        const block_type* const end = set_ + blocks_;
        for(const block_type* first = set_; first != end; ++first)
          if(*first)
            return false;

        return true;
      }

      const_iterator begin() const {
        return const_iterator(0, ConstTransformOp(*this));
      }

      iterator begin() {
        return iterator(0, TransformOp(*this));
      }

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

      iterator end() {
        return iterator(size_, TransformOp(*this));
      }

      /// Set a bit value

      /// \param i The bit to be set
      /// \param value The new value of the bit
      /// \throw std::out_of_range When \c i is >= size.
      void set(size_type i, bool value = true) {
        TA_ASSERT(i < size_);
        if(value)
          set_[block_index(i)] |= mask(i);
        else
          reset(i);
      }

      /// Set all bits

      /// \throw nothing
      void set() {
        std::fill_n(set_, blocks_, xffff);

        // Zero the tail
        const size_type extra_bits = bit_index(size_);
        if (extra_bits != 0)
            set_[blocks_ - 1] &= ~(xffff << extra_bits);
      }

      /// Set all bits from first to last

      /// \param first The first bit in the range to set
      /// \param last The last bit in the range to set
      void set_range(size_type first, size_type last) {
        if(last >= size_)
          last = size_ - 1;
        TA_ASSERT(first < last);

        // Get iterator and shift values
        block_type* first_block = set_ + block_index(first);
        const size_type first_shift = bit_index(first);
        block_type* const last_block = set_ + block_index(last);
        const size_type last_shift = block_bits - bit_index(last) - 1;

        // Set the first and last bits
        if(first_block == last_block) {
          *first_block |= (xffff << first_shift) & (xffff >> last_shift);
        } else {
          *first_block++ |= (xffff << first_shift);
          *last_block |= (xffff >> last_shift);
        }

        // Set all blocks between the first and last blocks.
        while(first_block < last_block)
          *first_block++ = xffff;
      }

      /// Set elements separated by \c stride

      /// \param first The first bit to set
      /// \param stride The distance between each set bit
      void set_stride(size_type first, size_type stride) {
        for(; first < size_; first += stride)
          set_[block_index(first)] |= mask(first);
      }

      /// Reset a bit

      /// \param i The bit to be reset
      /// \throw std::out_of_range When \c i is >= size.
      void reset(size_type i) {
        TA_ASSERT(i < size_);
        set_[block_index(i)] &= ~mask(i);
      }

      /// Set all bits

      /// \throw nothing
      void reset() {
        std::fill_n(set_, blocks_, zero);
      }

      /// Flip a bit

      /// \param i The bit to be flipped
      /// \throw std::out_of_range When \c i is >= size.
      void flip(size_type i) {
        TA_ASSERT(i < size_);
        set_[block_index(i)] ^= mask(i);
      }

      /// Flip all bits

      /// \throw nothing
      void flip() {
        for(size_type i = 0; i < blocks_; ++i)
          set_[i] = ~set_[i];
      }

      /// Count the number of non-zero bits

      /// \return The number of non-zero bits
      size_type count() const {
        size_type c = 0ul;
        for(size_type i = 0ul; i < blocks_; ++i) {
          block_type v = set_[i]; // temp
          v = v - ((v >> 1) & xffff / 3);
          v = (v & xffff / 15 * 3) + ((v >> 2) & xffff / 15 * 3);
          v = (v + (v >> 4)) & xffff / 255 * 15;
          c += block_type(v * (xffff / 255)) >> (sizeof(block_type) - 1) * CHAR_BIT; // count
        }
        return c;
      }

      /// Data pointer accessor

      /// The pointer to the data points to a contiguous block of memory of type
      /// \c block_type that contains \c num_blocks() elements.
      /// \return A pointer to the first element of the bitset data
      /// \throw nothing
      const block_type* get() const { return set_; }


      /// Data pointer accessor

      /// The pointer to the data points to a contiguous block of memory of type
      /// \c block_type that contains \c num_blocks() elements.
      /// \return A pointer to the first element of the bitset data
      /// \throw nothing
      block_type* get() { return set_; }

      /// Bitset size

      /// \return Number of bits in the bitset
      /// \throw nothing
      size_type size() const { return size_; }

      /// Bitset block size

      /// \return Number of block_type elements used to store the bitset array.
      /// \throw nothing
      size_type num_blocks() const { return blocks_; }

      void swap(Bitset_& other) {
        std::swap(size_, other.size_);
        std::swap(blocks_, other.blocks_);
        std::swap(set_, other.set_);
      }

    private:

      /// Calculate block index

      /// \return The block index that contains the i-th bit
      /// \throw nothing
      static size_type block_index(size_type i) { return i / block_bits; }

      /// Calculate the bit index

      /// \return The bit index that contains the i-th bit in the block
      /// \throw nothing
      static size_type bit_index(size_type i) { return i % block_bits; }

      /// Construct a mask

      /// \return A \c block_type that contains a single bit "on" bit at the i-th
      /// bit index.
      static block_type mask(size_type i) { return one << bit_index(i); }

      size_type size_;    ///< The number of bits in the set
      size_type blocks_;  ///< The number of blocks used to store the bits
      block_type* set_;   ///< An array that store the bits
    }; // class Bitset

    template <typename B>
    inline void swap(Bitset<B>& b0, Bitset<B>& b1) {
      b0.swap(b1);
    }

    // Bitset static constant data
    template <typename Block>
    const std::size_t Bitset<Block>::block_bits =
        8ul * sizeof(typename Bitset<Block>::block_type);
    template <typename Block>
    const Block Bitset<Block>::zero = Block(0);
    template <typename Block>
    const Block Bitset<Block>::one = Block(1);
    template <typename Block>
    const Block Bitset<Block>::xffff = ~Block(0);

    /// Bitwise and operator of bitset.

    /// \tparam Block The bitset block type
    /// \param left The left-hand bitset
    /// \param right The right-hand bitset
    /// \return The a intersection of the \c left and \c right bitsets
    template <typename Block>
    Bitset<Block> operator&(Bitset<Block> left, const Bitset<Block>& right) {
      left &= right;
      return left;
    }

    /// Bitwise or operator of bitset.

    /// \tparam Block The bitset block type
    /// \param left The left-hand bitset
    /// \param right The right-hand bitset
    /// \return The union of the \c left and \c right bitsets
    template <typename Block>
    Bitset<Block> operator|(Bitset<Block> left, const Bitset<Block>& right) {
      left |= right;
      return left;
    }


    /// Bitwise xor operator of bitset.

    /// \tparam Block The bitset block type
    /// \param left The left-hand bitset
    /// \param right The right-hand bitset
    /// \return The union of the \c left and \c right bitsets
    template <typename Block>
    Bitset<Block> operator^(Bitset<Block> left, const Bitset<Block>& right) {
      left ^= right;
      return left;
    }

    template <typename Block>
    std::ostream& operator<<(std::ostream& os, const Bitset<Block>& bitset) {
      os << std::hex;
      for(long i = bitset.num_blocks() - 1l; i >= 0l; --i)
        os << std::setfill('0') << std::setw(sizeof(Block)*2) << bitset.get()[i] << " ";

      os << std::setbase(10) << std::setw(0);
      return os;
    }


  } // namespace detail
} // namespace TiledArray

#endif // TILEDARRAY_BITSET_H__INCLUDED
