#include <Tpetra_Core.hpp>
#include <Tpetra_Map.hpp>
#include <Teuchos_CommHelpers.hpp>
#include <Tpetra_Map_decl.hpp>

namespace test{

template<class Scalar = Tpetra::Details::DefaultTypes::scalar_type,
         class LocalOrdinal = Tpetra::Details::DefaultTypes::local_ordinal_type,
         class GlobalOrdinal =Tpetra::Details::DefaultTypes::global_ordinal_type,
         class Node = Tpetra::Details::DefaultTypes::node_type>
class BlockMultiVector;

template<class Scalar = Tpetra::Details::DefaultTypes::scalar_type,
         class LocalOrdinal = Tpetra::Details::DefaultTypes::local_ordinal_type,
         class GlobalOrdinal =Tpetra::Details::DefaultTypes::global_ordinal_type,
         class Node = Tpetra::Details::DefaultTypes::node_type>
class BlockVector;

template<class Scalar = Tpetra::Details::DefaultTypes::scalar_type,
         class LocalOrdinal = Tpetra::Details::DefaultTypes::local_ordinal_type,
         class GlobalOrdinal = Tpetra::Details::DefaultTypes::global_ordinal_type,
         class Node = Tpetra::Details::DefaultTypes::node_type>
class MultiVector;

template<class Scalar = Tpetra::Details::DefaultTypes::scalar_type,
         class LocalOrdinal = Tpetra::Details::DefaultTypes::local_ordinal_type,
         class GlobalOrdinal = Tpetra::Details::DefaultTypes::global_ordinal_type,
         class Node = Tpetra::Details::DefaultTypes::node_type>
class Vector;

template<class LocalOrdinal = Tpetra::Details::DefaultTypes::local_ordinal_type,
         class GlobalOrdinal = Tpetra::Details::DefaultTypes::global_ordinal_type,
         class Node = Tpetra::Details::DefaultTypes::node_type>
class DistObject;

template <class LocalOrdinal, class GlobalOrdinal, class Node>
class DistObject{
public:
  using local_ordinal_type = LocalOrdinal;
  using global_ordinal_type = GlobalOrdinal;
  using node_type = Node;
  using map_type = Tpetra::Map<local_ordinal_type, global_ordinal_type, node_type>;

  explicit DistObject (const Teuchos::RCP<const map_type>& map) : map_(*map){}

  DistObject (const DistObject<LocalOrdinal, GlobalOrdinal, Node>&) = default;
  DistObject& operator= (const DistObject< LocalOrdinal, GlobalOrdinal, Node>&) = default;
  DistObject (DistObject< LocalOrdinal, GlobalOrdinal, Node>&&) = default;
  DistObject& operator= (DistObject< LocalOrdinal, GlobalOrdinal, Node>&&) = default;
  virtual ~DistObject () = default;

  virtual Teuchos::RCP<const map_type> getMap () const {
    return Teuchos::rcp(new map_type(map_));
  }

public:
  map_type map_;
};

//
// TPETRA
//
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
class MultiVector : public DistObject<LocalOrdinal, GlobalOrdinal, Node>
{
public:
  using map_type = Tpetra::Map<LocalOrdinal, GlobalOrdinal, Node>;
  using local_ordinal_type = typename map_type::local_ordinal_type;
  using global_ordinal_type = typename map_type::global_ordinal_type;
  using node_type = typename map_type::node_type;
  using base_type = DistObject<local_ordinal_type,global_ordinal_type, node_type>;

  MultiVector() : base_type(){}
  MultiVector (const Teuchos::RCP<const map_type>& map, const size_t numVecs) : base_type (map){}

protected:
  MultiVector (const MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& X,
	       const size_t j)
    : base_type (X.getMap ()){}

public:
  size_t getLocalLength () const{
    if (this->getMap ().is_null ()) {
      return static_cast<size_t> (0);
    } else {
      return this->getMap()->getNodeNumElements();
    }
  }

  Teuchos::RCP<Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> >
  getVectorNonConst (const size_t j){
    typedef Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> V;
    return Teuchos::rcp (new V (*this, j));
  }
};

template <class Scalar, class LocalOrdinal, class GlobalOrdinal,class Node>
class Vector : public MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>
{
private:
  using base_type = MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>;

public:
  typedef LocalOrdinal local_ordinal_type;
  typedef GlobalOrdinal global_ordinal_type;
  typedef Node node_type;
  typedef typename base_type::map_type map_type;

  Vector () : base_type(){}
  Vector (const MultiVector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& X, const size_t j)
    : base_type (X, j){}

  Vector (const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>&) = default;
  Vector (Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>&&) = default;
  Vector& operator= (const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>&) = default;

  Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>&
  operator=(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>&&) = default;

  virtual ~Vector () = default;
};


//
// TPETRA-BLOCK
//
template<class Scalar, class LO, class GO, class Node>
class BlockMultiVector : public DistObject<LO, GO, Node>
{
public:
  using dist_object_type = DistObject<LO, GO, Node>;
  using map_type = Tpetra::Map<LO, GO, Node>;
  using mv_type = MultiVector<Scalar, LO, GO, Node>;
  using scalar_type = Scalar;

  BlockMultiVector (const BlockMultiVector<Scalar, LO, GO, Node>&) = default;
  BlockMultiVector (BlockMultiVector<Scalar, LO, GO, Node>&&) = default;

  BlockMultiVector<Scalar, LO, GO, Node>&
  operator= (const BlockMultiVector<Scalar, LO, GO, Node>&) = default;

  BlockMultiVector<Scalar, LO, GO, Node>&
  operator= (BlockMultiVector<Scalar, LO, GO, Node>&&) = default;

  BlockMultiVector(): dist_object_type(Teuchos::null), blockSize_(0){}

  BlockMultiVector (const map_type& meshMap,
                    const LO blockSize,
                    const LO numVecs)
    : dist_object_type (Teuchos::rcp (new map_type (meshMap))), // shallow copy
      meshMap_(meshMap),
      pointMap_(makePointMap (meshMap, blockSize)),
      mv_ (Teuchos::rcpFromRef (pointMap_), numVecs),
      blockSize_(blockSize){}

  map_type makePointMap (const map_type& meshMap, const LO blockSize)
  {
    typedef Tpetra::global_size_t GST;
    const GST gblNumMeshMapInds = static_cast<GST> (meshMap.getGlobalNumElements ());
    const size_t lclNumMeshMapIndices = static_cast<size_t> (meshMap.getLocalNumElements ());
    const GST gblNumPointMapInds = gblNumMeshMapInds * static_cast<GST> (blockSize);
    const size_t lclNumPointMapInds = lclNumMeshMapIndices * static_cast<size_t> (blockSize);
    const GO indexBase = meshMap.getIndexBase ();
    return map_type (gblNumPointMapInds, lclNumPointMapInds, indexBase, meshMap.getComm ());
  }

protected:
  map_type meshMap_;
  map_type pointMap_;
  mv_type mv_;
  LO blockSize_;
};

template<class Scalar, class LO, class GO, class Node>
class BlockVector : public BlockMultiVector<Scalar, LO, GO, Node> {
private:
  typedef BlockMultiVector<Scalar, LO, GO, Node> base_type;
  typedef Vector<Scalar, LO, GO, Node> vec_type;
  typedef Tpetra::Map<LO, GO, Node> map_type;

public:
  BlockVector() : base_type(){}

  BlockVector (const BlockVector<Scalar, LO, GO, Node>&) = default;
  BlockVector (BlockVector<Scalar, LO, GO, Node>&&) = default;
  BlockVector<Scalar, LO, GO, Node>& operator= (const BlockVector<Scalar, LO, GO, Node>&) = default;
  BlockVector<Scalar, LO, GO, Node>& operator= (BlockVector<Scalar, LO, GO, Node>&&) = default;

  BlockVector (const map_type& meshMap, const LO blockSize)
    : base_type (meshMap, blockSize, 1){}

  vec_type getVectorView () {
    Teuchos::RCP<vec_type> vPtr = this->mv_.getVectorNonConst(0);
    return *vPtr;
  }
};
}

using tcomm    = Teuchos::Comm<int>;
using map_t    = Tpetra::Map<>;
using vec_type = test::BlockVector<>;

class MyObject{
private:
  map_t map_;
public:
  using state_type = vec_type;

  MyObject(){
    Teuchos::RCP<const tcomm> comm_ = Teuchos::rcp (new Teuchos::MpiComm<int>(MPI_COMM_WORLD));
    const int numGlobalEntries = 15*comm_->getSize();
    map_ = map_t(numGlobalEntries, 0, comm_);
  }
  state_type createState() const{
    state_type v(map_, 5);
    return v;
  }
};

template<class DT1, class DT2>
struct Registry{
  DT1 d1_;
  DT2 d2_;
  DT1 & getA(){ return d1_; }
  DT2 & getB(){ return d2_; }
};

void print(int rank, const std::string & s1, const std::string & s2, vec_type r){
  std::cout << "\n";
  auto tpmv  = r.getVectorView();
  std::cout << s1 << "ext = " << r.getMap()->getGlobalNumElements() << "\n";
  std::cout << s2 << "ext = " << tpmv.getMap()->getGlobalNumElements() << "\n";
}

template<class RegistryType>
class Foo{
  int rank_;
  RegistryType reg_;
public:
  Foo(RegistryType && reg) : reg_(std::move(reg))
  {
    MPI_Comm_rank(MPI_COMM_WORLD, &rank_);
    auto & a = reg_.getA();
    auto & b = reg_.getB();
    print(rank_, "CONSTR-block-A ", "CONSTR-tpetr-A ", a);
    print(rank_, "CONSTR-block-B ", "CONSTR-tpetr-B ", b);
  }

  void dummy() {
    auto & a = reg_.getA();
    auto & b = reg_.getB();
    print(rank_, "DUMMY-block-A ", "DUMMY-tpetr-A ", a);
    print(rank_, "DUMMY-block-B ", "DUMMY-tpetr-B ", b);
  }
};

template<class T>
auto createFoo(const T& system){
  using state_t    = typename T::state_type;
  using registry_t = Registry<state_t, state_t>;
  registry_t reg{system.createState(), system.createState()};
  return Foo<registry_t>(std::move(reg));
}

template<class T>
auto createRegistry(const T & system){
  using state_t    = typename T::state_type;
  using registry_t = Registry<state_t, state_t>;
  registry_t reg{system.createState(), system.createState()};
  return reg;
}

int main(int argc, char **argv)
{
  Tpetra::ScopeGuard tpetraScope (&argc, &argv);
  {
    {
      std::cout << "starting case that works\n";
      MyObject object;
      auto reg = createRegistry(object);
      Foo<decltype(reg)> f(std::move(reg));
      f.dummy();
    }

    {
      std::cout << "\n";
      std::cout << "\n";
      std::cout << "starting case that does NOT work\n";
      MyObject object;
      auto f = createFoo(object);
      f.dummy();
    }
  }

  return 0;
}