#include "pyTensor.h"

#include <type_traits>

#include "pybind11/operators.h"
#include "pybind11/stl.h"
#include "pybind11/numpy.h"

#include "taco/type.h"
#include "taco/tensor.h"

#if CUDA_BUILT
#include <cuda_runtime_api.h>
#endif

// Add Python dictionary initializer with {tuple(coordinate) : data} pairs


namespace taco{
namespace pythonBindings{

static void checkBounds(const std::vector<int>& dims, const std::vector<int>& indices){

  // Check for potential scalar access. Don't throw error if syntax valid
  if(dims.empty() && (indices.empty() || (indices[0] == 0 && indices.size() == 1))){
    return;
  }

  if(dims.size() != indices.size()){
    std::ostringstream o;
    o << "Incorrect number of dimensions when indexing. Tensor is order " << dims.size() << " but got index of "
                                                                                            "size " << indices.size();
    o << ". To index multiple dimensions only \"fancy\" notation is supported. For example to access the first "
         "element of a matrix, use A[0, 0] instead of A[0][0].";
    throw py::value_error(o.str());
  }

  for(size_t i = 0; i < dims.size(); ++i){
    if(indices[i] >= dims[i]){
      std::ostringstream o;
      o << "Index out of range for dimension " << i << ". Dimension shape is " << dims[i] << " but index value is "
           << indices[i];
      throw py::index_error(o.str());
    }
  }
}

template<typename T>
static Tensor<T> fromNpArr(py::buffer_info& array_buffer, Format& fmt, bool copy){

  std::vector<ssize_t> buf_shape = array_buffer.shape;
  std::vector<int> shape(buf_shape.begin(), buf_shape.end());
  const ssize_t size = array_buffer.size;

  // Creat row-major dense tensor
  Tensor<T> tensor(shape, fmt);
  TensorStorage& storage = tensor.getStorage();
  void *buf_data = array_buffer.ptr;
  Array::Policy policy = Array::Policy::UserOwns;
  if(should_use_CUDA_codegen()){
#if CUDA_BUILT
    taco_iassert(should_use_CUDA_unified_memory());
    buf_data = cuda_unified_alloc(size * array_buffer.itemsize);
    cudaMemcpy(buf_data, array_buffer.ptr, size * array_buffer.itemsize, cudaMemcpyDefault);
    policy = Array::Policy::Free;
#else
    taco_iassert(false);
#endif
  }
  else if(copy){
    buf_data = new T[size];
    memcpy(buf_data, array_buffer.ptr, size*array_buffer.itemsize);
    policy = Array::Policy::Delete;
  }

  storage.setValues(makeArray(static_cast<T*>(buf_data), size, policy));
  tensor.setStorage(storage);
  return tensor;
}

template<typename T>
static Tensor<T> fromNumpyF(py::array_t<T, py::array::f_style> &array, bool copy) {

  py::buffer_info array_buffer = array.request();
  const ssize_t dims = array_buffer.ndim;

  // Creat col-major dense tensor
  std::vector<int> ordering;
  for(int i = dims-1; i >= 0; --i){
    ordering.push_back(i);
  }

  Format fmt(std::vector<ModeFormatPack>(dims, dense), ordering);
  return fromNpArr<T>(array_buffer, fmt, copy);
}


template<typename T>
static Tensor<T> fromNumpyC(py::array_t<T, py::array::c_style>  &array, bool copy) {
  py::buffer_info array_buffer = array.request();
  const ssize_t dims = array_buffer.ndim;
  Format fmt(std::vector<ModeFormatPack>(dims, dense));
  return fromNpArr<T>(array_buffer, fmt, copy);
}

template<typename IdxType, typename T>
static Tensor<T> fromSpMatrix(py::array_t<IdxType> &ind_ptr, py::array_t<IdxType> &inds, py::array_t<T> &data,
                               const std::vector<int> &dims, bool copy, bool CSR){

  py::buffer_info ind_ptr_buf = ind_ptr.request();
  py::buffer_info inds_buf = inds.request();
  py::buffer_info data_buf = data.request();

  if(ind_ptr_buf.ndim != 1 || inds_buf.ndim != 1 || data_buf.ndim != 1) {
    throw py::value_error("Data arrays must be 1D.");
  }

  IdxType *mat_ptr  = static_cast<IdxType *>(ind_ptr_buf.ptr);
  IdxType *mat_ind  = static_cast<IdxType *>(inds_buf.ptr);
  T *mat_data = static_cast<T *>(data_buf.ptr);
  Array::Policy policy = Array::Policy::UserOwns;

  if(should_use_CUDA_codegen()){
    taco_iassert(should_use_CUDA_unified_memory());
    // TODO: Should copy arrays to unified memory
    taco_not_supported_yet;
  }
  else if(copy){
    mat_ptr = new IdxType[ind_ptr_buf.size];
    mat_ind = new IdxType[inds_buf.size];
    mat_data = new T[data_buf.size];
    memcpy(mat_ptr, ind_ptr_buf.ptr, ind_ptr_buf.size*ind_ptr_buf.itemsize);
    memcpy(mat_ind, inds_buf.ptr, inds_buf.size * inds_buf.itemsize);
    memcpy(mat_data, data_buf.ptr, data_buf.size * data_buf.itemsize);
    policy = Array::Policy::Delete;
  }

  // Create CSR Matrix
  Tensor<T> tensor;
  if(CSR){
    tensor = makeCSR(util::uniqueName("csr"), dims, mat_ptr, mat_ind, mat_data, policy);
  } else{
    tensor = makeCSC(util::uniqueName("csc"), dims, mat_ptr, mat_ind, mat_data, policy);
  }

  return tensor;
}

template<typename T>
static py::tuple toSpMatrix(Tensor<T> &tensor, bool tocsr) {
  if(tensor.getOrder() != 2) {
    throw py::value_error("Must be a matrix to convert to scipy");
  }

  // Force computation of the tensor
  tensor.pack();
  if(tensor.needsCompute()){
    tensor.evaluate();
  }

  int *ptr, *idx;
  T* vals;
  size_t ptr_arr_size, idx_arr_size, val_arr_size;

  // We may get a matrix in any format so we copy into a new tensor. Also we remove any explicit 0s before
  // moving to the scipy representation since the scipy contructor from dense arrays seems to do this as well.
  Tensor<T> t(tensor.getDimensions(), tocsr? CSR: CSC);

  for (auto& value : tensor) {
    if (value.second != 0) {
      t.insert(value.first.toVector(), value.second);
    }
  }
  t.pack();

  if(tocsr){
    getCSRArrays(t, &ptr, &idx, &vals);
  }else {
    getCSCArrays(t, &ptr, &idx, &vals);
  }

  // Could return these arrays without the memcpy. Would need to get the data pointers and change the
  // taco policies to UserOwn but would need to check the old policy to ensure that we free the right
  // way in general in the py capsules below. This code works so left with the double copy for now.
  auto index = t.getStorage().getIndex();
  ptr_arr_size = index.getModeIndex(1).getIndexArray(0).getSize();
  idx_arr_size = index.getModeIndex(1).getIndexArray(1).getSize();
  val_arr_size = t.getStorage().getValues().getSize();


  int *np_ptr = new int[ptr_arr_size];
  int *np_idx = new int[idx_arr_size];
  T   *np_vals   = new T[val_arr_size];

  memcpy(np_ptr, ptr, ptr_arr_size*sizeof(int));
  memcpy(np_idx, idx, idx_arr_size*sizeof(int));
  memcpy(np_vals, vals, val_arr_size*sizeof(T));

  py::capsule free_ptr(np_ptr, [](void *f) {
      int *p = static_cast<int *>(f);
      delete[] p;
  });

  py::capsule free_idx(np_idx, [](void *f) {
      int *p = static_cast<int *>(f);
      delete[] p;
  });

  py::capsule free_vals(np_vals, [](void *f) {
      T *p = static_cast<T *>(f);
      delete[] p;
  });

  py::array_t<int> ptr_arr({ptr_arr_size}, {sizeof(int)}, np_ptr, free_ptr);
  py::array_t<int> idx_arr({idx_arr_size}, {sizeof(int)}, np_idx, free_idx);
  py::array_t<T> val_arr({val_arr_size}, {sizeof(T)}, np_vals, free_vals);

  return py::make_tuple(ptr_arr, idx_arr, val_arr);
}

template<typename CType, typename idxVar>
static inline Access accessGetter(Tensor<CType>& tensor, idxVar& var) {
  return tensor(var);
}

template<typename CType>
static inline CType elementGetter(Tensor<CType>& tensor, std::vector<int> coords) {
  checkBounds(tensor.getDimensions(), coords);
  if(tensor.getOrder() == 0) {
    return tensor.at({});
  }
  return tensor.at(coords);
}

template<typename CType, typename pyType>
static inline void elementSetter(Tensor<CType> &tensor, std::vector<int> coords, pyType value) {
  checkBounds(tensor.getDimensions(), coords);
  if(tensor.getOrder() == 0) {
    tensor = static_cast<CType>(value);
  }
  tensor.insert(coords, static_cast<CType>(value));
}

template<typename CType>
static void insert(Tensor<CType> &tensor, std::vector<int> coords, double value) {
  checkBounds(tensor.getDimensions(), coords);
  if(tensor.getOrder() == 0) {
    tensor = static_cast<CType>(value);
  }
  tensor.insert(coords, static_cast<CType>(value));
}

template<typename CType, typename pyType>
static inline void singleElementSetter(Tensor<CType> &tensor, int coord, pyType value) {
  elementSetter<CType, pyType>(tensor, {coord}, value);
}

template<typename CType, typename VarType, typename ExprType>
static inline void exprSetter(Tensor<CType> &tensor, VarType idx, ExprType expr) {
  tensor(idx) = expr;
}

template<typename CType, typename VarType, typename SType>
static inline void exprScalarSetter(Tensor<CType> &tensor, VarType idx, SType scalar) {
  tensor(idx) = IndexExpr(scalar);
}

template<typename T>
static Tensor<T> makeTensor(std::string s, std::vector<int> shape, std::vector<ModeFormatPack> fmt) {
  return Tensor<T>(s, shape, Format(fmt));
}


template<typename T>
class PyTensorIter {

public:
  PyTensorIter(Tensor<T> &tensor) : end(tensor.end()), it(tensor.begin()) {
  }

  py::tuple advance() {
    // Ignore explicit zeros
    while (it != end && it->second == static_cast<T>(0)) {
      ++it;
    }

    if (it == end) {
      throw py::stop_iteration();
    }

    const auto coords = it->first.toVector();
    const auto val = it->second;
    ++it;

    return py::make_tuple(coords, val);
  }

private:
  const typename Tensor<T>::template const_iterator<int,T> end;
  typename Tensor<T>::template const_iterator<int,T> it;
};


template<typename CType>
static void declareTensor(py::module &m, const std::string typestr) {

  std::string pyIterName = std::string("py_tensor_iterator") + typestr;
  py::class_<PyTensorIter<CType>>(m, pyIterName.c_str())
          .def("__iter__", [](PyTensorIter<CType> &it) -> PyTensorIter<CType>&
                  { return it; })
          .def("__next__", &PyTensorIter<CType>::advance);

  using typedTensor = Tensor<CType>;

  m.def("to_sp_matrix", &toSpMatrix<CType>);

  m.def("fromNpF", &fromNumpyF<CType>);
  m.def("fromNpC", &fromNumpyC<CType>);

  m.def("fromSpMatrix", &fromSpMatrix<int, CType>);

  std::string pyClassName = std::string("Tensor") + typestr;
  py::class_<typedTensor, TensorBase>(m, pyClassName.c_str(), py::buffer_protocol())

          .def(py::init<>())

          .def(py::init<std::string>(), py::arg("name"))

          .def(py::init<CType>(), py::arg("value"))

          .def(py::init<std::string, std::vector<int>, ModeFormat>(), py::arg("name"), py::arg("shape"),
               py::arg("format") = ModeFormat::compressed)

          .def(py::init<std::string, std::vector<int>, Format>(), py::arg("name"), py::arg("shape"),
               py::arg("format"))

          .def(py::init(&makeTensor<CType>))

          .def(py::init<TensorBase>())

          .def_buffer([](typedTensor &t) -> py::buffer_info {

              if(!isDense(t.getFormat())){
                throw py::value_error("Cannot export a compressed tensor. Make sure all dimensions are dense "
                                      "using to_dense() before attempting this conversion.");
              }

              // Force computation of the tensor
              t.pack();
              if(t.needsCompute()){
                t.evaluate();
              }

              void *ptr = t.getStorage().getValues().getData();

              std::vector<ssize_t> shape (t.getDimensions().begin(), t.getDimensions().end());
              std::vector<ssize_t> row_major_strides;

              for(size_t i = 0; i < shape.size(); ++i) {
                ssize_t currentStride = sizeof(CType);
                for(size_t j = i + 1; j < shape.size(); ++j){
                  currentStride *= shape[j];
                }
                row_major_strides.push_back(currentStride);
              }

              std::vector<ssize_t> strides;
              for(const int &permutation : t.getFormat().getModeOrdering()){
                strides.push_back(row_major_strides[permutation]);
              }

              return py::buffer_info(
                      ptr,                                         /* Pointer to buffer */
                      sizeof(CType),                               /* Size of one scalar */
                      py::format_descriptor<CType>::format(),      /* Python struct-style format descriptor */
                      t.getOrder(),                                /* Number of dimensions */
                      shape,                                       /* Buffer dimensions */
                      strides                                      /* Strides (in bytes) for each index */
              );
          })

          .def("set_name", &TensorBase::setName)

          .def("get_name", &TensorBase::getName)

          .def("order", &TensorBase::getOrder)

          .def("get_shape", &TensorBase::getDimension, py::arg("axis"))

          .def("dtype", &TensorBase::getComponentType)

          .def("get_dimensions", &TensorBase::getDimensions)

          .def("format", &TensorBase::getFormat)

          .def("pack", &typedTensor::pack)

          // only bind .compile(), not .compile(IndexStmt, bool)
          .def("compile", [](typedTensor &self) { self.compile(); } )

          .def("assemble", &typedTensor::assemble)

          .def("evaluate", &typedTensor::evaluate)

          .def("compute", &typedTensor::compute)

          .def("insert", &insert<CType>)

          .def("remove_explicit_zeros", &typedTensor::removeExplicitZeros)

          .def("transpose", [](typedTensor &self, std::vector<int> dims, Format format, std::string name) -> typedTensor {
              return self.transpose(name, dims, format);
          }, py::is_operator())

          .def("__getitem__", [](typedTensor& self, const int &index) -> CType {
               return elementGetter<CType>(self, {index});
            }, py::is_operator())

          .def("__getitem__", [](typedTensor& self, const std::vector<int> &indices) -> CType {
               return elementGetter<CType>(self, indices);
            }, py::is_operator())

          .def("__getitem__", [](typedTensor& self, std::nullptr_t ptr) -> Access{
            if(self.getOrder() != 0) {
              throw py::index_error("Can only index scalar tensors with None.");
            }
            return self();
          }, py::is_operator())

          .def("__iter__", [](typedTensor &t) {return PyTensorIter<CType>(t); } )

          .def("__getitem__", &accessGetter<CType, IndexVar&>, py::is_operator())

          .def("__getitem__", &accessGetter<CType, std::vector<IndexVar>&>, py::is_operator())

          //  Set scalars to expression using none
          .def("__setitem__", [](typedTensor& self, std::nullptr_t ptr, const IndexExpr expr) -> void {
              self = expr;
          }, py::is_operator())

          .def("__setitem__", [](typedTensor& self, std::nullptr_t ptr, const Access access) -> void {
              self = access;
          }, py::is_operator())

          .def("__setitem__", [](typedTensor& self, std::nullptr_t ptr, const TensorVar tensorVar) -> void {
              self = tensorVar;
          }, py::is_operator())

          // Set expressions with varying types
          .def("__setitem__", &exprSetter<CType, IndexVar, IndexExpr>, py::is_operator())

          .def("__setitem__", &exprSetter<CType, IndexVar, Access>, py::is_operator())

          .def("__setitem__", &exprSetter<CType, IndexVar, TensorVar>, py::is_operator())

          .def("__setitem__", &exprSetter<CType, std::vector<IndexVar>, IndexExpr>, py::is_operator())

          .def("__setitem__", &exprSetter<CType, std::vector<IndexVar>, Access>, py::is_operator())

          .def("__setitem__", &exprSetter<CType, std::vector<IndexVar>, TensorVar>, py::is_operator())

          .def("__setitem__", &exprScalarSetter<CType, IndexVar, int64_t>, py::is_operator())

          .def("__setitem__", &exprScalarSetter<CType, IndexVar, double>, py::is_operator())

          .def("__setitem__", &exprScalarSetter<CType, std::vector<IndexVar>, int64_t>, py::is_operator())

          .def("__setitem__", &exprScalarSetter<CType, std::vector<IndexVar>, double>, py::is_operator())

          // This is a hack that exploits pybind11's resolution order. If we get here all other methods to resolve the
          // function failed and we throw an error. There may be better was to handle this in pybind.
          .def("__getitem__", [](typedTensor& self, const py::object &indices) -> void {
              std::ostringstream o;
              o << "Indices must be an iterable of integers or IndexVars but got " << indices;
              throw py::index_error(o.str());
          }, py::is_operator())

          .def("__setitem__", [](typedTensor& self, const py::object &indices, py::object value) -> void {
              std::ostringstream o;
              o << "Indices must be an iterable of IndexVars assigned to an index expression or a "
                   "value that can be transformed to an index expression (float or int) but got "
                   << indices << " and " << value << ". Note that element assignment is disabled in this release"
                   "and replace with .insert which increment the element at a given position (see the docs).";
              throw py::index_error(o.str());
          }, py::is_operator())

          .def("__repr__",   [](typedTensor& self) -> std::string{
              std::ostringstream o;
              o << self;
              return o.str();
          }, py::is_operator());

}



void defineTensor(py::module &m) {

  py::implicitly_convertible<ModeFormat, Format>();
  py::implicitly_convertible<std::vector<ModeFormat>, Format>();

  py::class_<TensorBase>(m, "TensorBase")
          .def("dtype", &TensorBase::getComponentType);

  declareTensor<bool>(m, "Bool");
  declareTensor<int8_t>(m, "Int8");
  declareTensor<int16_t>(m, "Int16");
  declareTensor<int32_t>(m, "Int32");
  declareTensor<int64_t>(m, "Int64");
  declareTensor<uint8_t>(m, "UInt8");
  declareTensor<uint16_t>(m, "UInt16");
  declareTensor<uint32_t>(m, "UInt32");
  declareTensor<uint64_t>(m, "UInt64");
  declareTensor<float>(m, "Float");
  declareTensor<double>(m, "Double");
}

}}