/**
 * @file   field_map.hh
 *
 * @author Till Junge <till.junge@epfl.ch>
 *
 * @date   15 Aug 2019
 *
 * @brief  Implementation of the base class of all field maps
 *
 * Copyright © 2019 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_FIELD_MAP_HH_
#define SRC_LIBMUGRID_FIELD_MAP_HH_

#include "core/types.hh"
#include "core/enums.hh"
#include "grid/iterators.hh"
#include "collection/field_collection.hh"

#include "Eigen/Dense"

#include <cassert>
#include <type_traits>
#include <memory>
#include <functional>

namespace muGrid {
    /**
     * base class for field map-related exceptions
     */
    class FieldMapError : public muGrid::RuntimeError {
       public:
        //! constructor
        explicit FieldMapError(const std::string & what)
            : muGrid::RuntimeError(what) {}
        //! constructor
        explicit FieldMapError(const char * what)
            : muGrid::RuntimeError(what) {}
    };

    // forward declaration
    template <typename T, typename MemorySpace>
    class TypedFieldBase;

    /**
     * Dynamically sized field map. Field maps allow iterating over the pixels
     * or quadrature points of a field and to select the shape (in a matrix
     * sense) of the iterate. For example, it allows to iterate in 2×2 matrices
     * over the quadrature points of a strain field for a two-dimensional
     * problem.
     */
    template <typename T, Mapping Mutability>
    class FieldMap {
       public:
        //! stored scalar type
        using Scalar = T;

        //! const-correct field depending on mapping mutability (HostSpace only)
        using Field_t =
            std::conditional_t<Mutability == Mapping::Const,
                               const TypedFieldBase<T, HostSpace>,
                               TypedFieldBase<T, HostSpace>>;

        //! dynamically mapped eigen type
        using PlainType = Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>;

        //! return type for iterators over this- map
        template <Mapping MutVal>
        using Return_t = std::conditional_t<MutVal == Mapping::Const,
                                            Eigen::Map<const PlainType>,
                                            Eigen::Map<PlainType>>;

        //! Input type for matrix-like values (used for setting uniform values)
        using EigenRef = Eigen::Ref<const PlainType>;

        /**
         * zip-container for iterating over pixel index and stored value
         * simultaneously
         */
        using PixelEnumeration_t = akantu::containers::ZipContainer<
            FieldCollection::PixelIndexIterable, FieldMap &>;
        /**
         * zip-container for iterating over pixel or quadrature point index and
         * stored value simultaneously
         */
        using Enumeration_t =
            akantu::containers::ZipContainer<FieldCollection::IndexIterable,
                                             FieldMap &>;

        //! determine whether a field is mutably mapped at compile time
        constexpr static Mapping FieldMutability() { return Mutability; }

        //! determine whether a field map is statically sized at compile time
        constexpr static bool IsStatic() { return false; }

        //! Default constructor
        FieldMap() = delete;

        /**
         * Constructor from a field. The iter_type can be the natural
         * sub-division of the field, or `muGrid::IterUnit::Pixel. The default
         * case is a map iterating over sub-division points with a matrix of
         * shape (nb_components × 1) per field entry
         */
        explicit FieldMap(Field_t & field,
                          const IterUnit & iter_type = IterUnit::SubPt);

        /**
         * Constructor from a field with explicitly chosen shape of iterate.
         * (the number of columns is inferred). The iter_type can be the natural
         * sub-division of the field, or `muGrid::IterUnit::Pixel
         */
        FieldMap(Field_t & field, Index_t nb_rows,
                 const IterUnit & iter_type = IterUnit::SubPt);

        //! Copy constructor
        FieldMap(const FieldMap & other) = delete;

        //! Move constructor
        FieldMap(FieldMap && other);

        //! Destructor
        virtual ~FieldMap() = default;

        //! Copy assignment operator (delete because of reference member)
        FieldMap & operator=(const FieldMap & other) = delete;

        //! Move assignment operator (delete because of reference member)
        FieldMap & operator=(FieldMap && other) = delete;

        //! Assign a matrix-like value to every entry
        template <bool IsMutableField = Mutability == Mapping::Mut>
        std::enable_if_t<IsMutableField, FieldMap> &
        operator=(const EigenRef & val) {
            if (not((val.rows() == this->nb_rows) and
                    (val.cols() == this->nb_cols))) {
                std::stringstream error_str{};
                error_str << "Expected an array/matrix with shape ("
                          << this->nb_rows << " × " << this->nb_cols
                          << "), but received a value of shape (" << val.rows()
                          << " × " << val.cols() << ").";
                throw FieldMapError(error_str.str());
            }
            for (auto && entry : *this) {
                entry = val;
            }
            return *this;
        }

        //! Addition-assign a matrix-like value to every entry
        template <bool IsMutableField = Mutability == Mapping::Mut>
        std::enable_if_t<IsMutableField, FieldMap> &
        operator+=(const EigenRef & val) {
            if (not((val.rows() == this->nb_rows) and
                    (val.cols() == this->nb_cols))) {
                std::stringstream error_str{};
                error_str << "Expected an array/matrix with shape ("
                          << this->nb_rows << " × " << this->nb_cols
                          << "), but received a value of shape (" << val.rows()
                          << " × " << val.cols() << ").";
                throw FieldMapError(error_str.str());
            }
            for (auto && entry : *this) {
                entry += val;
            }
            return *this;
        }

        //! Subtraction-assign a matrix-like value to every entry
        template <bool IsMutableField = Mutability == Mapping::Mut>
        std::enable_if_t<IsMutableField, FieldMap> &
        operator-=(const EigenRef & val) {
            if (not((val.rows() == this->nb_rows) and
                    (val.cols() == this->nb_cols))) {
                std::stringstream error_str{};
                error_str << "Expected an array/matrix with shape ("
                          << this->nb_rows << " × " << this->nb_cols
                          << "), but received a value of shape (" << val.rows()
                          << " × " << val.cols() << ").";
                throw FieldMapError(error_str.str());
            }
            for (auto && entry : *this) {
                entry -= val;
            }
            return *this;
        }

        //! Assign a scalar value to every entry
        template <bool IsMutableField = Mutability == Mapping::Mut>
        std::enable_if_t<IsMutableField, FieldMap> &
        operator=(const Scalar & val) {
            if (not(this->nb_rows == 1 && this->nb_cols == 1)) {
                std::stringstream error_str{};
                error_str << "Expected an array/matrix with shape ("
                          << this->nb_rows << " × " << this->nb_cols
                          << "), but received a scalar value.";
                throw FieldMapError(error_str.str());
            }
            for (auto && entry : *this) {
                entry(0, 0) = val;
            }
            return *this;
        }

        //! Addition-assign a scalar value to every entry
        template <bool IsMutableField = Mutability == Mapping::Mut>
        std::enable_if_t<IsMutableField, FieldMap> &
        operator+=(const Scalar & val) {
            if (not(this->nb_rows == 1 && this->nb_cols == 1)) {
                std::stringstream error_str{};
                error_str << "Expected an array/matrix with shape ("
                          << this->nb_rows << " × " << this->nb_cols
                          << "), but received a scalar value.";
                throw FieldMapError(error_str.str());
            }
            for (auto && entry : *this) {
                entry(0, 0) += val;
            }
            return *this;
        }

        //! Subtraction-assign a scalar value to every entry
        template <bool IsMutableField = Mutability == Mapping::Mut>
        std::enable_if_t<IsMutableField, FieldMap> &
        operator-=(const Scalar & val) {
            if (not(this->nb_rows == 1 && this->nb_cols == 1)) {
                std::stringstream error_str{};
                error_str << "Expected an array/matrix with shape ("
                          << this->nb_rows << " × " << this->nb_cols
                          << "), but received a scalar value.";
                throw FieldMapError(error_str.str());
            }
            for (auto && entry : *this) {
                entry(0, 0) -= val;
            }
            return *this;
        }

        /**
         * forward-declaration for `mugrid::FieldMap`'s iterator
         */
        template <Mapping MutIter>
        class Iterator;
        //! stl
        using iterator =
            Iterator<(Mutability == Mapping::Mut) ? Mapping::Mut
                                                  : Mapping::Const>;

        //! stl
        using const_iterator = Iterator<Mapping::Const>;

        //! stl
        iterator begin();

        //! stl
        iterator end();

        //! stl
        const_iterator cbegin();

        //! stl
        const_iterator cend();

        //! stl
        const_iterator begin() const;

        //! stl
        const_iterator end() const;

        /**
         * returns the number of iterates produced by this map (corresponds to
         * the number of field entries if Iteration::Subpt, or the number of
         * pixels/voxels if Iteration::Pixel); In parallel, if the processor
         * does not hold a part of the field (empty processor) this empty
         * processor returns size=0.
         */
        size_t size() const;

        //! random access operator
        Return_t<Mutability> operator[](size_t index) {
            assert(this->is_initialised);
            assert(index >= 0);
            assert(index <=
                   static_cast<size_t>(this->field.get_nb_buffer_entries()));
            return Return_t<Mutability>{this->data_ptr + index * this->stride,
                                        this->nb_rows, this->nb_cols};
        }

        //! random const access operator
        Return_t<Mapping::Const> operator[](size_t index) const {
            assert(this->is_initialised);
            assert(index >= 0);
            assert(index <=
                   static_cast<size_t>(this->field.get_nb_buffer_entries()));
            return Return_t<Mapping::Const>{this->data_ptr +
                                                index * this->stride,
                                            this->nb_rows, this->nb_cols};
        }

        //! query the size from the field's collection and set data_ptr
        void set_data_ptr();

        /**
         * return an iterable proxy over pixel indices and stored values
         * simultaneously. Throws a `muGrid::FieldMapError` if the iteration
         * type is over quadrature points
         */
        PixelEnumeration_t enumerate_pixel_indices_fast();

        /**
         * return an iterable proxy over pixel/quadrature indices and stored
         * values simultaneously
         */
        Enumeration_t enumerate_indices();

        //! evaluate and return the sum of the map. In parallel, if the
        //! processor does not hold a part of the field (empty processor) this
        //! empty processor returns a zero matrix of the proper size and type.
        PlainType sum() const;

        //! evaluate and return the mean value of the map. In parallel, if the
        //! processor does not hold a part of the field (empty processor) this
        //! empty processor returns a zero matrix of the proper size and type.
        PlainType mean() const;

        //! return const reference to the mapped field
        const Field_t & get_field() const;
        Index_t get_stride() const;

       protected:
        //! mapped field. Needed for query at initialisations
        const Field_t & field;
        const IterUnit iteration;  //!< type of map iteration
        const Index_t stride;      //!< precomputed stride
        const Index_t nb_rows;     //!< number of rows of the iterate
        const Index_t nb_cols;     //!< number of columns fo the iterate

        /**
         * Pointer to mapped data; is also unknown at construction and set in
         * the map's begin function
         */
        T * data_ptr{nullptr};

        //! keeps track of whether the map has been initialised.
        bool is_initialised{false};
        //! shared_ptr used for latent initialisation
        std::shared_ptr<std::function<void()>> callback{nullptr};
    };

    template <typename T, Mapping Mutability>
    template <Mapping MutIter>
    class FieldMap<T, Mutability>::Iterator {
       public:
        //! convenience alias
        using FieldMap_t = std::conditional_t<MutIter == Mapping::Const,
                                              const FieldMap, FieldMap>;
        //! stl
        using value_type =
            typename FieldMap<T, Mutability>::template Return_t<MutIter>;

        //! stl
        using cvalue_type =
            typename FieldMap<T, Mutability>::template Return_t<Mapping::Const>;

        //! Default constructor
        Iterator() = delete;

        //! Constructor to beginning, or to end
        Iterator(FieldMap_t & map, bool end)
            : map{map}, index{end ? map.size() : 0} {}

        //! Copy constructor
        Iterator(const Iterator & other) = delete;

        //! Move constructor
        Iterator(Iterator && other) = default;

        //! Destructor
        virtual ~Iterator() = default;

        //! Copy assignment operator
        Iterator & operator=(const Iterator & other) = default;

        //! Move assignment operator
        Iterator & operator=(Iterator && other) = default;

        //! pre-increment
        Iterator & operator++() {
            this->index++;
            return *this;
        }
        //! dereference
        inline value_type operator*() { return this->map[this->index]; }
        //! dereference (force const operator[] to get const Eigen::Map)
        inline cvalue_type operator*() const {
            return static_cast<const FieldMap<T, Mutability> &>(this->map)[this->index];
        }
        //! equality
        inline bool operator==(const Iterator & other) const {
            return this->index == other.index;
        }
        //! inequality
        inline bool operator!=(const Iterator & other) const {
            return not(*this == other);
        }

       protected:
        //! FieldMap being iterated over
        FieldMap_t & map;
        //! current iteration index
        size_t index;
    };

}  // namespace muGrid

#endif  // SRC_LIBMUGRID_FIELD_MAP_HH_
