/*
 *  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_DIST_EVAL_DIST_EVAL_BASE_H__INCLUDED
#define TILEDARRAY_DIST_EVAL_DIST_EVAL_BASE_H__INCLUDED

#include <TiledArray/tensor_impl.h>
#include <TiledArray/permutation.h>
#include <TiledArray/perm_index.h>
#include <TiledArray/type_traits.h>

namespace TiledArray {
  namespace detail {

    /// Distributed evaluator implementation object

    /// This class is used as the base class for other distributed evaluation
    /// implementation classes. It has several pure virtual function that are
    /// used by derived classes to implement the distributed evaluate. This
    /// class can also handles permutation of result tiles if necessary.
    /// \tparam Tile The output tile type
    /// \tparam Policy The tensor policy class
    template <typename Tile, typename Policy>
    class DistEvalImpl : public TensorImpl<Policy>, public madness::CallbackInterface {
    public:
      typedef DistEvalImpl<Tile, Policy> DistEvalImpl_; ///< This object type
      typedef TiledArray::detail::TensorImpl<Policy> TensorImpl_;
                                          ///< Tensor implementation base class

      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 type
      typedef typename eval_trait<value_type>::type eval_type; ///< Tile evaluation type

    private:
      madness::uniqueidT id_; ///< Globally unique object identifier.
      PermIndex source_to_target_; ///< Functor used to permute a source index to a target index.
      PermIndex target_to_source_; ///< Functor used to permute a target index to a source index.

      // The following variables are used to track the total number of tasks run
      // on the local node, task_count_, and the number of tiles set on this
      // node, set_counter_. They are used to track the progress of work done by
      // this node, which allows us to wait for the completion of these tasks
      // without waiting for all tasks.

      volatile int task_count_; ///< Total number of local tasks
      madness::AtomicInt set_counter_; ///< The number of tiles set by this node

    protected:


      /// Permute \c index from a source index to a target index

      /// \param index An ordinal index in the source index space
      /// \return The ordinal index in the target index space
      size_type perm_index_to_target(size_type index) const {
        TA_ASSERT(index < TensorImpl_::trange().tiles_range().volume());
        return (source_to_target_ ? source_to_target_(index) : index);
      }

      /// Permute \c index from a target index to a source index

      /// \param index An ordinal index in the target index space
      /// \return The ordinal index in the source index space
      size_type perm_index_to_source(size_type index) const {
        TA_ASSERT(index < TensorImpl_::trange().tiles_range().volume());
        return (target_to_source_ ? target_to_source_(index) : index);
      }

    public:
      /// Constructor

      /// \param world The world where the tensor lives
      /// \param trange The tiled range object
      /// \param shape The tensor shape object
      /// \param pmap The tile-process map
      /// \param perm The permutation that is applied to tile indices
      /// \note \c trange and \c shape will be permuted by \c perm before
      /// storing the data.
      DistEvalImpl(World& world, const trange_type& trange,
          const shape_type& shape, const std::shared_ptr<pmap_interface>& pmap,
          const Permutation& perm) :
        TensorImpl_(world, trange, shape, pmap),
        id_(world.unique_obj_id()),
        source_to_target_(),
        target_to_source_(),
        task_count_(-1),
        set_counter_()
      {
        set_counter_ = 0;

        if(perm) {
          Permutation inv_perm(-perm);
          range_type source_range = inv_perm * trange.tiles_range();
          source_to_target_ = PermIndex(source_range, perm);
          target_to_source_ = PermIndex(trange.tiles_range(), inv_perm);
        }
      }

      virtual ~DistEvalImpl() { }


      /// Unique object id accessor

      /// \return This object's unique identifier
      const madness::uniqueidT& id() const { return id_; }

      /// Get tile at index \c i

      /// \param i The index of the tile
      /// \return Tile at index i
      virtual Future<value_type> get_tile(size_type i) const = 0;

      /// Discard a tile that is not needed

      /// This function handles the cleanup for tiles that are not needed in
      /// subsequent computation.
      /// \param i The index of the tile
      virtual void discard_tile(size_type i) const = 0;

      /// Set tensor value

      /// This will store \c value at ordinal index \c i . Typically, this
      /// function should be called by a task function.
      /// \param i The index in the result space where value will be stored
      /// \param value The value to be stored at index \c i
      void set_tile(size_type i, const value_type& value) {
        // Store value
        madness::DistributedID id(id_, i);
        TensorImpl_::world().gop.send(TensorImpl_::owner(i), id, value);

        // Record the assignment of a tile
        DistEvalImpl_::notify();
      }

      /// Set tensor value with a future

      /// This will store \c value at ordinal index \c i . Typically, this
      /// function should be called by a task function.
      /// \param i The index in the result space where value will be stored
      /// \param f The future value to be stored at index \c i
      void set_tile(size_type i, Future<value_type> f) {
        // Store value
        madness::DistributedID id(id_, i);
        TensorImpl_::world().gop.send(TensorImpl_::owner(i), id, f);

        // Record the assignment of a tile
        f.register_callback(this);
      }

      /// Tile set notification
      virtual void notify() { set_counter_++; }

      /// Wait for all tiles to be assigned
      void wait() const {
        const int task_count = task_count_;
        if(task_count > 0) {
          try {
            TensorImpl_::world().await([this,task_count] ()
                { return this->set_counter_ == task_count; });
          } catch(...) {
            std::stringstream ss;
            ss << "!! ERROR TiledArray: Aborting due to exception.\n"
               << "!! ERROR TiledArray: rank=" << TensorImpl_::world().rank()
               << " id=" << id_ << " " << set_counter_ << " of " << task_count << " tiles set\n";
            std::cout << ss.str().c_str();
            throw;
          }
        }
      }

    private:

      /// Evaluate the tiles of this tensor

      /// This function will evaluate the children of this distributed evaluator
      /// and evaluate the tiles for this distributed evaluator. It will block
      /// until the tasks for the children are evaluated (not for the tasks of
      /// this object).
      /// \return The number of tiles that will be set by this process
      virtual int internal_eval() = 0;

    public:

      /// Evaluate this tensor expression object

      /// This function will evaluate the children of this distributed evaluator
      /// and evaluate the tiles for this distributed evaluator. It will block
      /// until the tasks for the children are evaluated (not for the tasks of
      /// this object).
      void eval() {
        TA_ASSERT(task_count_ == -1);
        task_count_ = this->internal_eval();
        TA_ASSERT(task_count_ >= 0);
      }

    }; // class DistEvalImpl


    /// Tensor expression object

    /// This object holds a tensor expression. It is used to store various type
    /// of tensor expressions that depend on the pimpl used to construct the
    /// expression.
    /// \tparam Tile The output tile type
    /// \tparam Policy The tensor policy class
    template <typename Tile, typename Policy>
    class DistEval {
    public:
      typedef DistEval<Tile, Policy> DistEval_; ///< This class type
      typedef DistEvalImpl<Tile, Policy> impl_type; ///< Implementation base class type
      typedef typename impl_type::size_type size_type; ///< Size type
      typedef typename impl_type::trange_type trange_type; ///< Tiled range type for this object
      typedef typename impl_type::range_type range_type; ///< Range type this tensor
      typedef typename impl_type::shape_type shape_type; ///< Tensor shape type
      typedef typename impl_type::pmap_interface pmap_interface; ///< Process map interface type
      typedef typename impl_type::value_type value_type; ///< Tile type
      typedef typename impl_type::eval_type eval_type; ///< Tile evaluation type
      typedef Future<value_type> future; ///< Future of tile type

    private:
      std::shared_ptr<impl_type> pimpl_; ///< pointer to the implementation object

    public:
      /// Constructor

      /// \param pimpl A pointer to the expression implementation object
      template <typename Impl>
      DistEval(const std::shared_ptr<Impl>& pimpl) :
        pimpl_(std::static_pointer_cast<impl_type>(pimpl))
      {
        TA_ASSERT(pimpl_);
      }

      /// Copy constructor

      /// Create a shallow copy of \c other .
      /// \param other The object to be copied.
      DistEval(const DistEval_& other) : pimpl_(other.pimpl_) { }

      /// Assignment operator

      /// Create a shallow copy of \c other
      /// \param other The object to be copied
      /// \return A reference to this object
      DistEval_& operator=(const DistEval_& other) {
        pimpl_ = other.pimpl_;
        return *this;
      }

      /// Evaluate this tensor object with the given result variable list

      /// \c v is the dimension ordering that the parent expression expects.
      /// The returned future will be evaluated once the tensor has been evaluated.
      void eval() { return pimpl_->eval(); }


      /// Tensor tile size array accessor

      /// \return The size array of the tensor tiles
      const range_type& range() const { return pimpl_->range(); }

      /// Tensor tile volume accessor

      /// \return The number of tiles in the tensor
      size_type size() const { return pimpl_->size(); }

      /// Query a tile owner

      /// \param i The tile index to query
      /// \return The process ID of the node that owns tile \c i
      ProcessID owner(size_type i) const { return pimpl_->owner(i); }

      /// Query for a locally owned tile

      /// \param i The tile index to query
      /// \return \c true if the tile is owned by this node, otherwise \c false
      bool is_local(size_type i) const { return pimpl_->is_local(i); }

      /// Query for a zero tile

      /// \param i The tile index to query
      /// \return \c true if the tile is zero, otherwise \c false
      bool is_zero(size_type i) const { return pimpl_->is_zero(i); }

      /// Tensor process map accessor

      /// \return A shared pointer to the process map of this tensor
      const std::shared_ptr<pmap_interface>& pmap() const { return pimpl_->pmap(); }

      /// Query the density of the tensor

      /// \return \c true if the tensor is dense, otherwise false
      bool is_dense() const { return pimpl_->is_dense(); }

      /// Tensor shape accessor

      /// \return A reference to the tensor shape map
      const shape_type& shape() const { return pimpl_->shape(); }

      /// Tiled range accessor

      /// \return The tiled range of the tensor
      const trange_type& trange() const { return pimpl_->trange(); }

      /// Tile move

      /// Tile is removed after it is set.
      /// \param i The tile index
      /// \return Tile \c i
      future get(size_type i) const { return pimpl_->get_tile(i); }

      /// Discard a tile that is not needed

      /// This function handles the cleanup for tiles that are not needed in
      /// subsequent computation.
      /// \param i The index of the tile
      virtual void discard(size_type i) const { pimpl_->discard_tile(i); }

      /// World object accessor

      /// \return A reference to the world object
      World& world() const { return pimpl_->world(); }

      /// Unique object id

      /// \return The unique id for this object
      madness::uniqueidT id() const { return pimpl_->id(); }

      /// Wait for all local tiles to be evaluated
      void wait() const { pimpl_->wait(); }

    }; // class DistEval

  }  // namespace detail
} // namespace TiledArray

#endif // TILEDARRAY_DIST_EVAL_DIST_EVAL_BASE_H__INCLUDED
