/*
 *  Copyright (C) 2004-2021 Edward F. Valeev
 *
 *  This file is part of Libint.
 *
 *  Libint is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Lesser General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  Libint 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 Lesser General Public License for more details.
 *
 *  You should have received a copy of the GNU Lesser General Public License
 *  along with Libint.  If not, see <http://www.gnu.org/licenses/>.
 *
 */


#ifndef _libint2_include_tensor_h_
#define _libint2_include_tensor_h_

#include <libint2/util/cxxstd.h>
#if LIBINT2_CPLUSPLUS_STD < 2011
# error "The simple Libint API requires C++11 support"
#endif

#include <cstdlib>
#include <numeric>
#include <vector>
#include <algorithm>

namespace libint2 {
    template <typename T> 
    struct Tensor {
    public:
      Tensor() = default;
      Tensor(const Tensor&) = default;
      Tensor(Tensor&&) = default;
      ~Tensor() = default;

      template <class ... Dims> Tensor(Dims ... dims) : dims_{ std::forward<Dims>(dims)... } {
        strides_.resize(sizeof...(dims));
        // used in transform to compute strides
        struct stride {
            size_t value;
            stride() : value(1) {}
            size_t operator()(size_t dim) { auto old_value = value; value *= dim; return old_value; }
        };
        // row-major order of dimensions
        std::transform(dims_.rbegin(), dims_.rend(),
                       strides_.rbegin(), stride());
        size_t size = strides_.size() == 0 ? 0 : strides_[0] * dims_[0];
        data_.resize(size);
      }

      T* data() { return &data_[0]; }
      const T* data() const { return static_cast<const T*>(this->data()); }
      template <class ... MultiIndex> T* data(MultiIndex ... index) {
        assert(sizeof...(MultiIndex) <= dims_.size());
        auto index_list = { std::forward<MultiIndex>(index)... };
        size_t ordinal = std::inner_product(index_list.begin(), index_list.end(),
                                            strides_.begin(), 0);
        return &data_[ordinal];
      }
      template <class ... MultiIndex> const T* data(MultiIndex ... index) const {
          return static_cast<const T*>(this->data());
      }

      template <class ... MultiIndex> const T& operator()(MultiIndex ... index) {
        assert(sizeof...(MultiIndex) == dims_.size());
        auto index_list = { std::forward<MultiIndex>(index)... };
        size_t ordinal = std::inner_product(index_list.begin(), index_list.end(),
                                            strides_.begin(), 0);
        return data_[ordinal];
      }

    private:
      std::vector<size_t> dims_;
      std::vector<size_t> strides_;
      std::vector<T> data_;
  
    }; // class definition
} // namespace libint2

#endif  /* _libint2_include_tensor_h_*/
