/**
 * @file   numpy_tools.hh
 *
 * @author Lars Pastewka <lars.pastewka@imtek.uni-freiburg.de>
 * @author Till Junge <till.junge@epfl.ch>
 *
 * @date   02 Dec 2019
 *
 * @brief  Convenience functionality for working with (py::'s) numpy arrays.
 *         These are implemented header-only, in order to avoid an explicit
 *         dependency on py::
 *
 * Copyright © 2018 Lars Pastewka, Till Junge
 *
 * µGrid 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, or (at
 * your option) any later version.
 *
 * µGrid 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 µGrid; see the file COPYING. If not, write to the
 * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
 * * Boston, MA 02111-1307, USA.
 *
 * Additional permission under GNU GPL version 3 section 7
 *
 * If you modify this Program, or any covered work, by linking or combining it
 * with proprietary FFT implementations or numerical libraries, containing parts
 * covered by the terms of those libraries' licenses, the licensors of this
 * Program grant you additional permission to convey the resulting work.
 *
 */

#ifndef SRC_LIBMUGRID_NUMPY_TOOLS_HH_
#define SRC_LIBMUGRID_NUMPY_TOOLS_HH_

#include "field_typed.hh"
#include "field_collection_global.hh"
#include "raw_memory_operations.hh"

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

#include <algorithm>
#include <limits>

namespace py = pybind11;

namespace muGrid {

  /**
   * base class for numpy related exceptions
   */
  class NumpyError : public RuntimeError {
   public:
    //! constructor
    explicit NumpyError(const std::string & what) : RuntimeError(what) {}
    //! constructor
    explicit NumpyError(const char * what) : RuntimeError(what) {}
  };

  namespace internal {

    //! convert strides into the storage order description used by muGrid
    template <typename T, int flags = py::array::forcecast>
    inline std::tuple<StorageOrder, DynCcoord_t, Shape_t>
    detect_storage_order(const DynCcoord_t & nb_subdomain_grid_pts,
                         const Shape_t & components_shape, Index_t nb_sub_pts,
                         py::array_t<T, flags> & array) {
      auto & dim{nb_subdomain_grid_pts.get_dim()};
      if (static_cast<py::ssize_t>(dim + components_shape.size()) !=
          array.ndim()) {
        std::stringstream s;
        s << "Internal error: Array dimension (= " << array.ndim()
          << ") do not match expected number of dimensions (" << dim
          << " spatial dimensions and shape of components " << components_shape
          << ").";
        throw NumpyError(s.str());
      }

      // Compute strides for SubPt and Pixel iterations passed on to directly
      // to the WrappedField if the storage order of this buffer does not match
      // the storage order intrinsically supported by muGrid.
      Shape_t strides{};

      auto nb_components{get_nb_from_shape(components_shape)};

      Shape_t components_strides{};
      DynCcoord_t pixels_strides{};
      for (size_t i = 0; i < components_shape.size(); ++i) {
        components_strides.push_back(array.strides(i));
        if (nb_components > 1) {
          strides.push_back(array.strides(i) / array.itemsize());
        }
      }
      if (nb_components > 1 and !components_shape.empty() and nb_sub_pts == 1) {
        // Insert a token stride for the single sub-point
        strides.push_back(0);
      }
      for (int i = 0; i < dim; ++i) {
        pixels_strides.push_back(array.strides(array.ndim() - dim + i));
        strides.push_back(array.strides(array.ndim() - dim + i) /
                          array.itemsize());
      }

      // Determine the storage order of the whole buffer. We first check the
      // order of pixels vs. sub-points storage. This is our first guess for
      // the storage order of the whole buffer.
      StorageOrder fc_storage_order{StorageOrder::ColMajor};
      if (components_strides.size() > 0) {
        fc_storage_order = StorageOrder::Unknown;
        const auto c{std::minmax_element(components_strides.begin(),
                                         components_strides.end())};
        const auto p{
            std::minmax_element(pixels_strides.begin(), pixels_strides.end())};
        if (*c.first > *p.second) {
          // this may be row-major, but we need to check that buffer is
          // contiguous
          if (*c.first ==
              nb_subdomain_grid_pts[p.second - pixels_strides.begin()] *
                  (*p.second)) {
            fc_storage_order = StorageOrder::RowMajor;
          }
        } else {
          if (*c.second <= *p.first) {
            // this may be column-major, but we need to check that buffer is
            // contiguous
            if (*p.first ==
                components_shape[c.second - components_strides.begin()] *
                    (*c.second)) {
              fc_storage_order = StorageOrder::ColMajor;
            }
          }
        }

        // the components support compact strides, but we must normalize these
        // smallest entry
        Index_t smallest_stride{*std::min_element(components_strides.begin(),
                                                  components_strides.end())};
        for (auto && s : components_strides) {
          s /= smallest_stride;
        }

        // determine order of the sub-points
        if (!CcoordOps::is_buffer_contiguous(components_shape,
                                             components_strides)) {
          fc_storage_order = StorageOrder::Unknown;
        }
      }

      // determine if the storage order of the components matches the storage
      // order determined above for the pixels vs. sub-point portion
      if (components_strides.size() > 1) {
        if (fc_storage_order == StorageOrder::ColMajor) {
          // check if the whole thing is column-major
          for (size_t i = 1; i < components_strides.size() - 1; ++i) {
            if (components_strides[i + 1] < components_strides[i]) {
              fc_storage_order = StorageOrder::Unknown;
            }
          }
        } else {
          // check if the whole thing is row-major
          for (size_t i = 1; i < components_strides.size() - 1; ++i) {
            if (components_strides[i + 1] > components_strides[i]) {
              fc_storage_order = StorageOrder::Unknown;
            }
          }
        }
      }

      // the pixels generally support arbitrary strides, but we must normalize
      // to the smallest entry
      Index_t smallest_stride{
          *std::min_element(pixels_strides.begin(), pixels_strides.end())};
      if (smallest_stride > 0) {
        for (auto && s : pixels_strides) {
          s /= smallest_stride;
        }
      }

      // If we have more than one sub-pt, the field to be wrapped must match
      // the native storage order of muGrid, because otherwise there is no way
      // to switch between SubPt and Pixel iteration.
      if (nb_sub_pts > 1) {
        if (fc_storage_order == StorageOrder::Unknown) {
          throw NumpyError(
              "Cannot wrap a field with more than one sub-point that does not "
              "match the intrinsic muGrid storage order.");
        }
      }

      if (fc_storage_order == StorageOrder::ColMajor or
          fc_storage_order == StorageOrder::RowMajor) {
        // This is a buffer that matches the intrinsic storage order of muGrid.
        // We clear the SubPt and Pixel shapes because they are correctly
        // computed by the respective Field.
        strides.clear();
      }

      return std::make_tuple(fc_storage_order, pixels_strides, strides);
    }

  }  // namespace internal

  /**
   * Construct a NumpyProxy given that we only know the number of components
   * of the field. The constructor will complain if the grid dimension differs
   * but will wrap any field whose number of components match. For example,
   * a 3x3 grid with 8 components could look like this:
   *    1. (8, 3, 3)
   *    2. (2, 4, 3, 3)
   *    3. (2, 2, 2, 3, 3)
   * The method `get_components_shape` returns the shape of the component part
   * of the field in this case. For the above examples, it would return:
   *    1. (8,)
   *    2. (2, 4)
   *    3. (2, 2, 2)
   * Note that a field with a single component can be passed either with a
   * shape having leading dimension of one or without any leading dimension.
   * In the latter case, `get_component_shape` will return a vector of size 0.
   * The same applies for fields with a single quadrature point, whose
   * dimension can be omitted. In general, the shape of the field needs to
   * look like this:
   *    (component_1, component_2, sub_pt, grid_x, grid_y, grid_z)
   * where the number of components and grid indices can be arbitrary.
   */
  template <typename T, int flags = py::array::forcecast,
            class Collection_t = GlobalFieldCollection>
  class NumpyProxy {
   public:
    /**
     * Construct a NumpyProxy given that we only know the number of components
     * of the field. The constructor will complain if the grid dimension differs
     * but will wrap any field whose number of components match. For example,
     * a 3x3 grid with 8 components could look like this:
     *    1. (8, 3, 3)
     *    2. (2, 4, 3, 3)
     *    3. (2, 2, 2, 3, 3)
     * The method `get_components_shape` returns the shape of the component part
     * of the field in this case. For the above examples, it would return:
     *    1. (8,)
     *    2. (2, 4)
     *    3. (2, 2, 2)
     * Note that a field with a single component can be passed either with a
     * shape having leading dimension of one or without any leading dimension.
     * In the latter case, `get_component_shape` will return a vector of size 0.
     * The same applies for fields with a single quadrature point, whose
     * dimension can be omitted. In general, the shape of the field needs to
     * look like this:
     *    (component_1, component_2, sub_pt, grid_x, grid_y, grid_z)
     * where the number of components and grid indices can be arbitrary.
     */
    NumpyProxy(DynCcoord_t nb_domain_grid_pts,
               DynCcoord_t nb_subdomain_grid_pts,
               DynCcoord_t subdomain_locations, Index_t nb_dof_per_pixel,
               py::array_t<T, flags> & array,
               const Unit & unit = Unit::unitless())
        : collection{}, field{}, iter_type{IterUnit::Pixel} {
      // Sanity check 1: Are the sizes of array and field equal?
      Index_t size{get_nb_from_shape(nb_subdomain_grid_pts) * nb_dof_per_pixel};
      if (size != array.size()) {
        std::stringstream s;
        s << "The numpy array has a size of " << array.size() << ", but the "
          << "muGrid field reports a size of " << size << " entries.";
        throw NumpyError(s.str());
      }

      // Sanity check 2: Has the array a dimension of least the spatial dim?
      Index_t dim{nb_subdomain_grid_pts.get_dim()};
      if (array.ndim() < dim) {
        std::stringstream s;
        s << "The numpy array has a dimension of " << array.ndim()
          << ", but the muGrid field reports a pixels dimension of " << dim
          << ".";
        throw NumpyError(s.str());
      }

      // Sanity check 3: Do the last three dimensions of the array agree with
      // the size of the grid?
      if (!std::equal(nb_subdomain_grid_pts.begin(),
                      nb_subdomain_grid_pts.end(),
                      array.shape() + array.ndim() - dim)) {
        std::stringstream s;
        s << "The numpy array has shape " << array.request().shape << ", but "
          << "the muGrid field reports a grid of " << nb_subdomain_grid_pts
          << " pixels. The numpy array must equal the grid size in its last "
          << "dimensions.";
        throw NumpyError(s.str());
      }

      Index_t nb_array_components{1};
      Shape_t components_shape{};
      for (Dim_t i = 0; i < array.ndim() - dim; ++i) {
        components_shape.push_back(array.shape(i));
        nb_array_components *= array.shape(i);
      }

      // Sanity check 4: Are the number of array components identical to the
      // number of sub-point degrees of freedom of the field? (This should not
      // fail if the above tests have not failed.)
      if (nb_array_components != nb_dof_per_pixel) {
        std::stringstream s;
        s << "The numpy array has shape " << array.request().shape << ", but "
          << "the muGrid field reports " << nb_dof_per_pixel
          << " components. The numpy array "
          << "must equal the number of components in its first dimensions.";
        throw NumpyError(s.str());
      }

      auto storage_order{internal::detect_storage_order(
          nb_subdomain_grid_pts, components_shape, 1, array)};
      auto & fc_storage_order{std::get<0>(storage_order)};
      auto & pixels_strides{std::get<1>(storage_order)};
      auto & strides{std::get<2>(storage_order)};

      FieldCollection::SubPtMap_t map{};
      map["proxy_subpt"] = 1;
      this->collection = std::make_unique<Collection_t>(
          nb_domain_grid_pts, nb_subdomain_grid_pts, subdomain_locations,
          pixels_strides, map, fc_storage_order);
      // Note: The pointer should be obtained from array.template mutable_data()
      // but this has a guard if the array is not writeable. We need a constant
      // WrappedField to clean up this code.
      this->field = std::make_unique<WrappedField<T>>(
          "proxy_field", *collection, components_shape, array.size(),
          static_cast<T *>(array.request().ptr), "proxy_subpt", unit, strides);
    }

    /**
     * Construct a NumpyProxy given that we know the shape of the leading
     * component indices. The constructor will complain if both the grid
     * dimensions and the component dimensions differ. `get_component_shape`
     * returns exactly the shape passed to this constructor.
     *
     * In general, the shape of the field needs to look like this:
     *    (component_1, component_2, sub_pt, grid_x, grid_y, grid_z)
     * where the number of components and grid indices can be arbitrary. The
     * sub_pt dimension can be omitted if there is only a single sub-point.
     */
    NumpyProxy(DynCcoord_t nb_domain_grid_pts,
               DynCcoord_t nb_subdomain_grid_pts,
               DynCcoord_t subdomain_locations, Index_t nb_sub_pts,
               Shape_t components_shape, py::array_t<T, flags> & array,
               const Unit & unit = Unit::unitless())
        : collection{}, field{}, iter_type{IterUnit::SubPt} {
      // Sanity check: Do the array dimensions agree and shapes agree?
      Index_t dim{nb_subdomain_grid_pts.get_dim()};
      bool shape_matches{false};
      Shape_t sub_pt_shape{components_shape};
      if (static_cast<py::ssize_t>(dim + components_shape.size() + 1) ==
          array.ndim()) {
        shape_matches = std::equal(nb_subdomain_grid_pts.begin(),
                                   nb_subdomain_grid_pts.end(),
                                   array.shape() + array.ndim() - dim) &&
                        nb_sub_pts == array.shape(components_shape.size()) &&
                        std::equal(components_shape.begin(),
                                   components_shape.end(), array.shape());
        sub_pt_shape.push_back(nb_sub_pts);
      } else if (static_cast<py::ssize_t>(dim + components_shape.size()) ==
                 array.ndim()) {
        // For a field with a single quad point, we can omit that dimension.
        shape_matches = std::equal(nb_subdomain_grid_pts.begin(),
                                   nb_subdomain_grid_pts.end(),
                                   array.shape() + array.ndim() - dim) &&
                        nb_sub_pts == 1 &&
                        std::equal(components_shape.begin(),
                                   components_shape.end(), array.shape());
        // Make sure we return a numpy field with identical shape
        this->iter_type = IterUnit::Pixel;
      }
      if (!shape_matches) {
        std::stringstream s;
        s << "The numpy array has shape " << array.request().shape << ", but "
          << "the muGrid field reports a grid of " << nb_subdomain_grid_pts
          << " pixels with " << nb_sub_pts << " sub-"
          << (nb_sub_pts == 1 ? "point" : "points") << " holding a quantity of "
          << "shape " << components_shape << ".";
        throw NumpyError(s.str());
      }

      auto storage_order{internal::detect_storage_order(
          nb_subdomain_grid_pts, sub_pt_shape, nb_sub_pts, array)};
      auto & fc_storage_order{std::get<0>(storage_order)};
      auto & pixels_strides{std::get<1>(storage_order)};
      auto & sub_pt_iter_strides{std::get<2>(storage_order)};

      FieldCollection::SubPtMap_t map{};
      map["proxy_subpt"] = nb_sub_pts;
      this->collection = std::make_unique<Collection_t>(
          nb_domain_grid_pts, nb_subdomain_grid_pts, subdomain_locations,
          pixels_strides, map, fc_storage_order);
      // Note: The pointer should be obtained from array.template mutable_data()
      // but this has a guard if the array is not writeable. We need a constant
      // WrappedField to clean up this code.
      this->field = std::make_unique<WrappedField<T>>(
          "proxy_field", *collection, components_shape, array.size(),
          static_cast<T *>(array.request().ptr), "proxy_subpt", unit,
          sub_pt_iter_strides);
    }

    /**
     * move constructor
     */
    NumpyProxy(NumpyProxy && other) = default;

    WrappedField<T> & get_field() { return *this->field; }

    Collection_t & get_collection() { return *this->collection; }

    const Shape_t get_components_shape() const {
      return this->field->get_components_shape();
    }

    Shape_t get_sub_pt_shape() const { return this->field->get_sub_pt_shape(); }

    IterUnit get_iter_type() const { return this->iter_type; }

   protected:
    std::unique_ptr<Collection_t> collection;
    std::unique_ptr<WrappedField<T>> field;
    IterUnit iter_type;  // IterUnit that will give input shape
  };

  /* Wrap a column-major field into a numpy array, without copying the data */
  template <typename T>
  py::array_t<T, py::array::f_style>
  numpy_wrap(const TypedFieldBase<T> & field,
             IterUnit iter_type = IterUnit::SubPt) {
    Shape_t shape{field.get_shape(iter_type)};
    Shape_t strides{field.get_strides(iter_type, sizeof(T))};
    return py::array_t<T, py::array::f_style>(shape, strides, field.data(),
                                              py::capsule([]() {}));
  }

  /* Copy a column-major field into a numpy array */
  template <typename T>
  py::array_t<T, py::array::f_style>
  numpy_copy(const TypedFieldBase<T> & field,
             IterUnit iter_type = IterUnit::SubPt) {
    const Shape_t shape{field.get_shape(iter_type)};
    py::array_t<T> array(shape);
    Shape_t array_strides(array.strides(), array.strides() + array.ndim());
    // numpy arrays have stride in bytes
    for (auto && s : array_strides)
      s /= sizeof(T);
    muGrid::raw_mem_ops::strided_copy(shape, field.get_strides(iter_type),
                                      array_strides, field.data(),
                                      array.mutable_data());
    return std::move(array);
  }

  /* Turn any type that can be enumerated into a tuple */
  template <typename T>
  py::tuple to_tuple(T a) {
    py::tuple t(a.get_dim());
    ssize_t i = 0;
    for (auto && v : a) {
      t[i] = v;
      i++;
    }
    return t;
  }

}  // namespace muGrid

#endif  // SRC_LIBMUGRID_NUMPY_TOOLS_HH_
