/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2014  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/>.
 *
 *  Justus Calvin
 *  Department of Chemistry, Virginia Tech
 *
 *  array_impl.h
 *  Oct 24, 2014
 *
 */

#ifndef TILEDARRAY_ARRAY_IMPL_H__INCLUDED
#define TILEDARRAY_ARRAY_IMPL_H__INCLUDED

#include <TiledArray/tensor_impl.h>
#include <TiledArray/distributed_storage.h>
#include <TiledArray/transform_iterator.h>
#include <TiledArray/type_traits.h>

namespace TiledArray {
  namespace detail {

    // Forward declaration
    template <typename> class TensorReference;
    template <typename> class TensorConstReference;
    template <typename, typename> class ArrayIiterator;

    /// Tensor tile reference

    /// \tparam Impl The TensorImpl type
    template <typename Impl>
    class TileReference {
    private:

      template <typename, typename>
      friend class ArrayIiterator;

      template <typename>
      friend class TileConstReference;

      typedef typename Impl::range_type range_type;
      typedef typename Impl::range_type::index index_type;
      typedef typename Impl::size_type ordinal_type;

      Impl* tensor_; ///< The tensor that owns the referenced tile
      ordinal_type index_; ///< The ordinal index of the tile

      // Not allowed
      TileReference<Impl>& operator=(const TileReference<Impl>&);
    public:

      TileReference(Impl* tensor, const typename Impl::size_type index) :
        tensor_(tensor), index_(index)
      { }

      TileReference(const TileReference<Impl>& other) :
        tensor_(other.tensor_), index_(other.index_)
      { }

      template <typename Value>
      TileReference<Impl>& operator=(const Value& value) {
        tensor_->set(index_, value);
        return *this;
      }

      typename Impl::future future() const {
        TA_ASSERT(tensor_);
        return tensor_->get(index_);
      }

      typename Impl::value_type get() const {
        // NOTE: return by value to avoid lifetime issues.
        TA_ASSERT(tensor_);
        return future().get();
      }

      operator typename Impl::future() const { return this->future(); }

      operator typename Impl::value_type() const { return get(); }

      /// Tile coordinate index accessor

      /// \return The coordinate index of the current tile
      index_type index() const {
        TA_ASSERT(tensor_);
        return tensor_->range().idx(index_);
      }

      /// Tile ordinal index accessor

      /// \return The ordinal index of the current tile
      ordinal_type ordinal() const {
        return index_;
      }

      /// Tile range factory function

      /// Construct a range object for the current tile
      range_type make_range() const {
        TA_ASSERT(tensor_);
        return tensor_->trange().make_tile_range(index_);
      }

    }; // class TileReference

    /// Tensor tile reference

    /// \tparam Impl The TensorImpl type
    template <typename Impl>
    class TileConstReference {
    private:

      template <typename, typename>
      friend class ArrayIiterator;

      const Impl* tensor_; ///< The tensor that owns the referenced tile
      typename Impl::size_type index_; ///< The index of the tensor

      // Not allowed
      TileConstReference<Impl>& operator=(const TileConstReference<Impl>&);
    public:

      TileConstReference(const Impl* tensor, const typename Impl::size_type index) :
        tensor_(tensor), index_(index)
      { }

      TileConstReference(const TileConstReference<Impl>& other) :
        tensor_(other.tensor_), index_(other.index_)
      { }

      TileConstReference(const TileReference<Impl>& other) :
        tensor_(other.tensor_), index_(other.index_)
      { }

      typename Impl::future future() const {
        TA_ASSERT(tensor_);
        return tensor_->get(index_);
      }

      typename Impl::value_type get() const {
        // NOTE: return by value to avoid lifetime issues.
        TA_ASSERT(tensor_);
        return future().get();
      }

      operator typename Impl::future() const { return tensor_->get(index_); }

      operator typename Impl::value_type() const { return get(); }
    }; // class TileConstReference

  } // namespace detail
} // namespace TiledArray


namespace madness {
  namespace detail {

    // The following class specializations are required so MADNESS will do the
    // right thing when given a TileReference or TileConstReference object
    // as an input for task functions.

    template <typename Impl>
    struct task_arg<TiledArray::detail::TileReference<Impl> > {
        typedef typename Impl::value_type type;
        typedef typename Impl::future holderT;
    }; // struct task_arg<TiledArray::detail::TileReference<Impl> >

    template <typename Impl>
    struct task_arg<TiledArray::detail::TileConstReference<Impl> > {
        typedef typename Impl::value_type type;
        typedef typename Impl::future holderT;
    }; // struct task_arg<TiledArray::detail::TileConstReference<Impl> >

  }  // namespace detail
}  // namespace madness


namespace TiledArray {
  namespace detail {

    /// Distributed tensor iterator

    /// This iterator will reference local tiles for a TensorImpl object. It can
    /// be used to get or set futures to a tile, or access the coordinate and
    /// ordinal index of the tile.
    /// \tparam Impl The TensorImpl type
    /// \tparam Reference The iterator reference type
    template <typename Impl, typename Reference>
    class ArrayIiterator {
    private:
      // Give access to other iterator types.
      template <typename, typename>
      friend class ArrayIiterator;

      Impl* array_;
      typename Impl::pmap_interface::const_iterator it_;

    public:
      typedef ptrdiff_t difference_type; ///< Difference type
      typedef typename Impl::future value_type; ///< Iterator dereference value type
      typedef PointerProxy<value_type> pointer; ///< Pointer type to iterator value
      typedef Reference reference; ///< Reference type to iterator value
      typedef std::forward_iterator_tag iterator_category; ///< Iterator category type
      typedef ArrayIiterator<Impl, Reference> ArrayIterator_; ///< This object type
      typedef typename Impl::range_type::index index_type;
      typedef typename Impl::size_type ordinal_type;
      typedef typename Impl::range_type range_type;
      typedef typename Impl::value_type tile_type;

    private:

      void advance() {
        TA_ASSERT(array_);
        const typename Impl::pmap_interface::const_iterator end =
            array_->pmap()->end();
        do {
          ++it_;
        } while((it_ != end) && array_->is_zero(*it_));
      }

    public:

      /// Default constructor
      ArrayIiterator() : array_(NULL), it_() { }

      /// Constructor
      ArrayIiterator(Impl* tensor, typename Impl::pmap_interface::const_iterator it) :
        array_(tensor), it_(it)
      { }

      /// Copy constructor

      /// \param other The transform iterator to copy
      ArrayIiterator(const ArrayIterator_& other) :
        array_(other.array_), it_(other.it_)
      { }

      /// Copy const iterator constructor

      /// \tparam R Iterator reference type
      /// \param other The transform iterator to copy
      template <typename I, typename R,
          typename std::enable_if<
              !((! std::is_const<Impl>::value) &&
              std::is_const<I>::value)
          >::type* = nullptr>
      ArrayIiterator(const ArrayIiterator<I, R>& other) :
        array_(other.array_), it_(other.it_)
      { }

      /// Copy operator

      /// \param other The transform iterator to copy
      /// \return A reference to this object
      ArrayIterator_& operator=(const ArrayIterator_& other) {
        array_ = other.array_;
        it_ = other.it_;

        return *this;
      }

      /// Copy operator

      /// \tparam R Iterator reference type
      /// \param other The transform iterator to copy
      /// \return A reference to this object
      template <typename R>
      ArrayIterator_& operator=(const ArrayIiterator<Impl, R>& other) {
        array_ = other.array_;
        it_ = other.it_;

        return *this;
      }

      /// Prefix increment operator

      /// \return A reference to this object after it has been incremented.
      ArrayIterator_& operator++() {
        advance();
        return *this;
      }

      /// Post-fix increment operator

      /// \return A copy of this object before it is incremented.
      ArrayIterator_ operator++(int) {
        ArrayIterator_ tmp(*this);
        advance();
        return tmp;
      }

      /// Equality operator

      /// \tparam R Iterator reference type
      /// \param other The iterator to compare to this iterator.
      /// \return \c true when the iterators are equal to each other, otherwise
      /// \c false.
      template <typename I, typename R>
      bool operator==(const ArrayIiterator<I, R>& other) const {
        return (array_ == other.array_) && (it_ == other.it_);
      }

      /// Inequality operator

      /// \tparam R Iterator reference type
      /// \param other The iterator to compare to this iterator.
      /// \return \c true when the iterators are not equal to each other,
      /// otherwise \c false.
      template <typename I, typename R>
      bool operator!=(const ArrayIiterator<I, R>& other) const {
        return (array_ != other.array_) || (it_ != other.it_);
      }

      /// Dereference operator

      /// \return A referenc to the current tile future.
      reference operator*() const {
        TA_ASSERT(array_);
        return reference(array_, *it_);
      }

      /// Arrow dereference operator

      /// \return A pointer-proxy to the current tile
      pointer operator->() const {
        TA_ASSERT(array_);
        return pointer(array_->get(*it_));
      }

      /// Tile coordinate index accessor

      /// \return The coordinate index of the current tile
      index_type index() const {
        TA_ASSERT(array_);
        return array_->range().idx(*it_);
      }

      /// Tile ordinal index accessor

      /// \return The ordinal index of the current tile
      ordinal_type ordinal() const {
        TA_ASSERT(array_);
        return *it_;
      }

      /// Tile range factory function

      /// Construct a range object for the current tile
      range_type make_range() const {
        TA_ASSERT(array_);
        TA_ASSERT(it_ != array_->pmap()->end());
        return array_->trange().make_tile_range(*it_);
      }

    }; // class TensorIterator

    /// Tensor implementation and base for other tensor implementation objects

    /// This implementation object holds the data for tensor object, which
    /// includes tiled range, shape, and tiles. The tiles are held in a
    /// distributed container, stored according to a given process map.
    /// \tparam Tile The tile or value_type of this tensor
    /// \note The process map must be set before data elements can be set.
    /// \note It is the users responsibility to ensure the process maps on all
    /// nodes are identical.
    template <typename Tile, typename Policy>
    class ArrayImpl : public TensorImpl<Policy> {
    public:
      typedef ArrayImpl<Tile, Policy> ArrayImpl_; ///< This object type
      typedef TensorImpl<Policy> TensorImpl_; ///< The base class of this object
      typedef typename TensorImpl_::size_type size_type; ///< Size type
      typedef typename TensorImpl_::trange_type trange_type; ///< Tiled range type for this object
      typedef typename TensorImpl_::range_type range_type; ///< Range type this tensor
      typedef typename TensorImpl_::shape_type shape_type; ///< Shape type
      typedef typename TensorImpl_::pmap_interface pmap_interface; ///< process map interface type
      typedef Tile value_type; ///< Tile or data type
      typedef typename eval_trait<Tile>::type eval_type; ///< The tile evaluation type
      typedef typename numeric_type<value_type>::type
          numeric_type; ///< the numeric type that supports Tile
      typedef DistributedStorage<value_type> storage_type; ///< The data container type
      typedef typename storage_type::future future; ///< Future tile type
      typedef TileReference<ArrayImpl_> reference; ///< Tile reference type
      typedef TileConstReference<ArrayImpl_> const_reference; ///< Tile constant reference type
      typedef ArrayIiterator<ArrayImpl_, reference> iterator; ///< Iterator type
      typedef ArrayIiterator<const ArrayImpl_, const_reference> const_iterator; ///< Constant iterator type

    private:

      storage_type data_; ///< Tile container

    public:

      /// Constructor

      /// The size of shape must be equal to the volume of the tiled range tiles.
      /// \param world The world where this tensor will live
      /// \param trange The tiled range for this tensor
      /// \param shape The shape of this tensor
      /// \param pmap The tile-process map
      /// \throw TiledArray::Exception When the size of shape is not equal to
      /// zero
      ArrayImpl(World& world, const trange_type& trange, const shape_type& shape,
          const std::shared_ptr<pmap_interface>& pmap) :
        TensorImpl_(world, trange, shape, pmap),
        data_(world, trange.tiles_range().volume(), pmap)
      { }

      /// Virtual destructor
      virtual ~ArrayImpl() { }

      /// Tile future accessor

      /// \tparam Index The index type
      /// \param i The tile index
      /// \return A \c future to tile \c i
      /// \throw TiledArray::Exception When tile \c i is zero
      template <typename Index>
      future get(const Index& i) const {
        TA_ASSERT(! TensorImpl_::is_zero(i));
        return data_.get(TensorImpl_::trange().tiles_range().ordinal(i));
      }

      /// Tile future accessor

      /// \tparam Integer An integer type
      /// \param i The tile index, as an \c std::initializer_list<Integer>
      /// \return A \c future to tile \c i
      /// \throw TiledArray::Exception When tile \c i is zero
      template <typename Integer>
      future get(const std::initializer_list<Integer>& i) const {
        return get<std::initializer_list<Integer>>(i);
      }

      /// Set tile

      /// Set the tile at \c i with \c value . \c Value type may be \c value_type ,
      /// \c Future<value_type> , or
      /// \c madness::detail::MoveWrapper<value_type> .
      /// \tparam Index The index type
      /// \tparam Value The value type
      /// \param i The index of the tile to be set
      /// \param value The object tat contains the tile value
      template <typename Index, typename Value>
      void set(const Index& i, const Value& value) {
        TA_ASSERT(! TensorImpl_::is_zero(i));
        data_.set(TensorImpl_::trange().tiles_range().ordinal(i), value);
      }

      /// Array begin iterator

      /// \return A const iterator to the first element of the array.
      iterator begin() {
        // Get the pmap iterator
        typename pmap_interface::const_iterator it = TensorImpl_::pmap()->begin();

        // Find the fist non-zero iterator
        const typename pmap_interface::const_iterator end = TensorImpl_::pmap()->end();
        while((it != end) && TensorImpl_::is_zero(*it)) ++it;

        // Construct and return the iterator
        return iterator(this, it);
      }

      /// Array begin iterator

      /// \return A const iterator to the first element of the array.
      const_iterator cbegin() const {
        // Get the pmap iterator
        typename pmap_interface::const_iterator it = TensorImpl_::pmap()->begin();

        // Find the fist non-zero iterator
        const typename pmap_interface::const_iterator end = TensorImpl_::pmap()->end();
        while((it != end) && TensorImpl_::is_zero(*it)) ++it;

        // Construct and return the iterator
        return const_iterator(this, it);
      }

      /// Array end iterator

      /// \return A const iterator to one past the last element of the array.
      iterator end() {
        return iterator(this, TensorImpl_::pmap()->end());
      }

      /// Array end iterator

      /// \return A const iterator to one past the last element of the array.
      const_iterator cend() const {
        return const_iterator(this, TensorImpl_::pmap()->end());
      }

      /// Unique object id accessor

      /// \return A const reference to this object unique id
      const madness::uniqueidT& id() const { return data_.id(); }

    }; // class ArrayImpl


#ifndef TILEDARRAY_HEADER_ONLY

    extern template
    class ArrayImpl<Tensor<double, Eigen::aligned_allocator<double> >, DensePolicy>;
    extern template
    class ArrayImpl<Tensor<float, Eigen::aligned_allocator<float> >, DensePolicy>;
    extern template
    class ArrayImpl<Tensor<int, Eigen::aligned_allocator<int> >, DensePolicy>;
    extern template
    class ArrayImpl<Tensor<long, Eigen::aligned_allocator<long> >, DensePolicy>;
//    extern template
//    class ArrayImpl<Tensor<std::complex<double>, Eigen::aligned_allocator<std::complex<double> > >, DensePolicy>;
//    extern template
//    class ArrayImpl<Tensor<std::complex<float>, Eigen::aligned_allocator<std::complex<float> > >, DensePolicy>;

    extern template
    class ArrayImpl<Tensor<double, Eigen::aligned_allocator<double> >, SparsePolicy>;
    extern template
    class ArrayImpl<Tensor<float, Eigen::aligned_allocator<float> >, SparsePolicy>;
    extern template
    class ArrayImpl<Tensor<int, Eigen::aligned_allocator<int> >, SparsePolicy>;
    extern template
    class ArrayImpl<Tensor<long, Eigen::aligned_allocator<long> >, SparsePolicy>;
//    extern template
//    class ArrayImpl<Tensor<std::complex<double>, Eigen::aligned_allocator<std::complex<double> > >, SparsePolicy>;
//    extern template
//    class ArrayImpl<Tensor<std::complex<float>, Eigen::aligned_allocator<std::complex<float> > >, SparsePolicy>;

#endif // TILEDARRAY_HEADER_ONLY

  }  // namespace detail
}  // namespace TiledArray

#endif // TILEDARRAY_ARRAY_IMPL_H__INCLUDED
