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

#ifndef TILEDARRAY_SYMM_PERMUTATION_H__INCLUDED
#define TILEDARRAY_SYMM_PERMUTATION_H__INCLUDED

#include <array>
#include <vector>
#include <map>
#include <set>

#include <algorithm>

#include <TiledArray/type_traits.h>
#include <TiledArray/error.h>

namespace TiledArray {

  /**
   * \addtogroup symmetry
   * @{
   */

  namespace symmetry {

    namespace detail {

      /// Create a permuted copy of an array

      /// \note a more efficient version of detail::permute_array specialized for TiledArray::symmetry::Permutation
      /// \tparam Perm The permutation type
      /// \tparam Arg The input array type
      /// \tparam Result The output array type
      /// \param[in] perm The permutation
      /// \param[in] arg The input array to be permuted
      /// \param[out] result The output array that will hold the permuted array
      template <typename Perm, typename Arg, typename Result>
      inline void permute_array(const Perm& perm, const Arg& arg, Result& result) {
        TA_ASSERT(result.size() == arg.size());
        if (perm.domain_size() < arg.size()) {
          std::copy(arg.begin(), arg.end(), result.begin()); // if perm does not map every element of arg, copy arg to result first
        }
        for(const auto& p: perm.data()) {
          TA_ASSERT(result.size() > p.second);
          TA_ASSERT(arg.size() > p.second);
          result[p.second] = arg[p.first];
        }
      }
    } // namespace detail

    /// Permutation of a sequence of objects indexed by base-0 indices.

    /// \warning Unlike TiledArray::Permutation, this does not fix domain size.
    ///
    /// Permutation class is used as an argument in all permutation operations on
    /// other objects. Permutations can be applied to sequences of objects:
    /// \code
    ///   b = p * a; // apply permutation p to sequence a and assign the result to sequence b.
    ///   a *= p;    // apply permutation p (in-place) to sequence a.
    /// \endcode
    /// Permutations can also be composed, e.g. multiplied and inverted:
    /// \code
    ///   p3 = p1 * p2;      // computes product of permutations of p1 and p2
    ///   p1_inv = p1.inv(); // computes inverse of p1
    /// \endcode
    ///
    /// \note
    ///
    /// \par
    /// Unlike TiledArray::Permutation, which is internally represented in one-line form,
    /// TiledArray::symmetry::Permutation is internally
    /// represented in compressed two-line form.
    /// E.g. the following permutation in Cauchy's two-line form,
    /// \f$
    ///   \left(
    ///   \begin{tabular}{ccccc}
    ///     0 & 1 & 2 & 3 & 4 \\
    ///     0 & 2 & 3 & 1 & 4
    ///   \end{tabular}
    ///   \right)
    /// \f$
    /// , is represented in compressed form as \f$ \{ 1 \to 2, 2 \to 3, 3 \to 1 \} \f$ . This means
    /// that 0th element of a sequence is mapped by this permutation into the 0th element of the permuted
    /// sequence (hence 0 is referred to as a <em>fixed point</em> of this permutation; so is 4);
    /// similarly, 1st element of a sequence is mapped by this permutation into the 2nd element of
    /// the permuted sequence (hence 2 is referred as the \em image of 1 under the action of this Permutation;
    /// similarly, 1 is the image of 3, etc.). Set \f$ \{1, 2, 3\} \f$ is referred to
    /// as \em domain  (or \em support) of this Permutation. Note that (by definition) Permutation
    /// maps its domain into itself (i.e. it's a bijection).
    ///
    /// \par
    /// As a reminder, permutation
    /// \f$
    ///   \left(
    ///   \begin{tabular}{ccccc}
    ///     0 & 1 & 2 & 3 & 4 \\
    ///     0 & 2 & 3 & 1 & 4
    ///   \end{tabular}
    ///   \right)
    /// \f$
    /// is represented in one-line form as \f$ \{0, 2, 3, 1, 4\} \f$. Note that the one-line representation
    /// is redundant as multiple distinct one-line representations correspond to the same
    /// compressed form, e.g. \f$ \{0, 2, 3, 1, 4\} \f$ and \f$ \{0, 2, 3, 1\} \f$ correspond to the
    /// same \f$ \{ 1 \to 2, 2 \to 3, 3 \to 1 \} \f$ compressed form.
    ///
    /// \par
    /// Another non-redundant representation of Permutation is as a set of cycles. For example,
    /// permutation \f$ \{0 \to 3, 1 \to 2, 2 \to 1, 0 \to 3 \} \f$ is represented uniquely as the
    /// following set of cycles: (0,3)(1,2).
    /// The canonical format for the cycle decomposition used by Permutation class is defined as follows:
    /// <ul>
    ///  <li> Cycles of length 1 are skipped.
    ///  <li> Each cycle is in order of increasing elements.
    ///  <li> Cycles are in the order of increasing first elements.
    /// </ul>
    /// Cycle representation is convenient for some operations, but is less efficient for others.
    /// Thus cycle representation can be computed on request, but internally the compressed form is used.
    ///
    class Permutation {
    public:
      typedef Permutation Permutation_;
      typedef unsigned int index_type;
      typedef std::map<index_type,index_type> Map;
      typedef Map::const_iterator const_iterator;

    private:

      /// Two-line representation of permutation
      Map p_;

      static std::ostream& print_map(std::ostream& output, const Map& p) {
        for (auto i=p.cbegin(); i!=p.cend();) {
          output << i->first << "->" << i->second;
          if (++i != p.cend())
            output << ", ";
        }
        return output;
      }
      friend inline std::ostream& operator<<(std::ostream& output, const Permutation& p);

      /// Validate permutation specified in one-line form as an iterator range
      /// \return \c true if each element of \c [first,last) is non-negative and unique
      template <typename InIter>
      static bool valid_permutation(InIter first, InIter last) {
        bool result = true;
        for(; first != last; ++first) {
          const auto value = *first;
          result = result && value >= 0 && (std::count(first, last, *first) == 1ul);
        }
        return result;
      }

      /// Validate permutation specified in compressed two-line form as an index->index associative container
      /// \note Map can be std::map<index,index>, std::unordered_map<index,index>, or any similar container.
      template <typename Map>
      static bool valid_permutation(const Map& input) {
        std::set<index_type> keys;
        std::set<index_type> values;
        for(const auto& e: input) {
          const auto& key = e.first;
          const auto& value = e.second;
          if (keys.find(key) == keys.end())
            keys.insert(key);
          else {
            //print_map(std::cout, input);
            return false; // key is found more than once
          }
          if (values.find(value) == values.end())
            values.insert(value);
          else {
            //print_map(std::cout, input);
            return false; // value is found more than once
          }
        }
        return keys == values; // must map domain into itself
      }

      // Used to select the correct constructor based on input template types
      struct Enabler { };

    public:

      Permutation() = default; // makes an identity permutation
      Permutation(const Permutation&) = default;
      Permutation(Permutation&&) = default;
      ~Permutation() = default;
      Permutation& operator=(const Permutation&) = default;
      Permutation& operator=(Permutation&& other) = default;

      /// Construct permutation using its 1-line form given by range [first,last)

      /// \tparam InIter An input iterator type
      /// \param first The beginning of the iterator range
      /// \param last The end of the iterator range
      /// \throw TiledArray::Exception if invalid input is given. \sa valid_permutation(first,last)
      template <typename InIter,
          typename std::enable_if<TiledArray::detail::is_input_iterator<InIter>::value>::type* = nullptr>
      Permutation(InIter first, InIter last)
      {
        TA_ASSERT( valid_permutation(first, last) );
        size_t i=0;
        for(auto e=first; e!=last; ++e, ++i) {
          auto p_i = *e;
          if (i != p_i) p_[i] = p_i;
        }
      }

      /// std::vector constructor

      /// Construct permutation using 1-line form input as a vector
      /// \param a vector that specifies permutation in 1-line form
      explicit Permutation(const std::vector<index_type>& a) :
          Permutation(a.begin(), a.end())
      {
      }

      /// Construct permutation with an initializer list

      /// \param list An initializer list of integers
      explicit Permutation(std::initializer_list<index_type> list) :
          Permutation(list.begin(), list.end())
      {
      }

      /// Construct permutation using its compressed 2-line form given by std::map

      /// \param p the map
      Permutation(Map p) : p_(std::move(p))
      {
        TA_ASSERT(valid_permutation(p_));
      }

      /// Permutation domain size

      /// \return The number of elements in the domain of permutation
      unsigned int domain_size() const { return p_.size(); }

      /// @name Iterator accessors
      ///
      /// Permutation iterators dereference into \c std::pair<index_type,index_type>, where \c first is the domain index, \c second is its image
      /// E.g. \c Permutation::begin() of \f$ \{1->3, 2->1, 3->2\} \f$ dereferences to \c std::pair {1,3} .
      ///
      /// @{

      /// Begin element iterator factory function

      /// \return An iterator that points to the beginning of the range
      const_iterator begin() const { return p_.begin(); }

      /// Begin element iterator factory function

      /// \return An iterator that points to the beginning of the range
      const_iterator cbegin() const { return p_.cbegin(); }

      /// End element iterator factory function

      /// \return An iterator that points to the end of the range
      const_iterator end() const { return p_.end(); }

      /// End element iterator factory function

      /// \return An iterator that points to the end of the range
      const_iterator cend() const { return p_.cend(); }

      /// @}

      /// Computes image of an element under this permutation

      /// \param e input index
      /// \return the image of element \c e; if \c e is in the domain of this permutation, returns its image, otherwise returns \c e
      index_type operator[](index_type e) const {
        auto e_image_iter = p_.find(e);
        if (e_image_iter != p_.end())
          return e_image_iter->second;
        else
          return e;
      }

      /// Computes the domain of this permutation

      /// \tparam Set a container type in which the result will be returned
      /// \return the domain of this permutation, as a Set
      template <typename Set>
      Set domain() const {
        Set result;
        for(const auto& e: this->p_) {
          result.insert(result.begin(), e.first);
        }
        //std::sort(result.begin(), result.end());
        return result;
      }

      /// Test if an index is in the domain of this permutation

      /// \tparam Integer an integer type
      /// \param i an index whose presence in domain is tested
      /// \return \c true , if \c i is in the domain, \c false otherwise
      template <typename Integer,
                typename std::enable_if<std::is_integral<Integer>::value>::type* = nullptr>
      bool is_in_domain(Integer i) const {
        return p_.find(i) != p_.end();
      }

      /// Cycles decomposition

      /// Certain algorithms are more efficient with permutations represented as a set of cyclic transpositions.
      /// This function returns the set of cycles that represent this permutation. For example,
      /// permutation \f$ \{3, 2, 1, 0 \} \f$ is represented as the following set of cycles: (0,3)(1,2).
      /// The canonical format for the cycles is:
      /// <ul>
      ///  <li> Cycles of length 1 are skipped.
      ///  <li> Each cycle is in order of increasing elements.
      ///  <li> Cycles are in the order of increasing first elements.
      /// </ul>
      /// \return the set of cycles (in canonical format) that represent this permutation
      std::vector<std::vector<index_type> > cycles() const {

        std::vector<std::vector<index_type>> result;

        std::set<index_type> placed_in_cycle;

        // safe to use non-const operator[] due to properties of Permutation (domain==image)
        auto& p_nonconst_ = const_cast<Map&>(p_);

        // 1. for each i compute its orbit
        // 2. if the orbit is longer than 1, sort and add to the list of cycles
        for(const auto& e: p_) {
          auto i = e.first;
          if (placed_in_cycle.find(i) == placed_in_cycle.end()) { // not in a cycle yet?
            std::vector<index_type> cycle(1,i);
            placed_in_cycle.insert(i);

            index_type next_i = p_nonconst_[i];
            while (next_i != i) {
              cycle.push_back(next_i);
              placed_in_cycle.insert(next_i);
              next_i = p_nonconst_[next_i];
            }

            if (cycle.size() != 1) {
              std::sort(cycle.begin(), cycle.end());
              result.emplace_back(cycle);
            }

          } // this i already in a cycle
        } // loop over i

        return result;
      }

      /// Product of this permutation by \c other

      /// \param other a Permutation
      /// \return \c other * \c this, i.e. \c this applied first, then \c other
      Permutation mult(const Permutation& other) const {
        // 1. domain of product = domain of this U domain of other
        using iset = std::set<index_type>;
        auto product_domain = this->domain<iset>();
        auto other_domain = other.domain<iset>();
        product_domain.insert(other_domain.begin(), other_domain.end());

        Map result;
        for(const auto& d: product_domain) {
          const auto d_image = other[(*this)[d]];
          if (d_image != d)
            result[d] = d_image;
        }

        return Permutation(result);
      }

      /// Construct the inverse of this permutation

      /// The inverse of permutation \f$ P \f$ is defined as \f$ P \times P^{-1} = P^{-1} \times P = I \f$,
      /// where \f$ I \f$ is the identity permutation.
      /// \return The inverse of this permutation
      Permutation inv() const {
        Map result;
        for(const auto& iter: p_) {
          const auto i = iter.first;
          const auto pi = iter.second;
          result.insert(std::make_pair(pi,i));
        }
        return Permutation(std::move(result));
      }

      /// Raise this permutation to the n-th power

      /// Constructs the permutation \f$ P^n \f$, where \f$ P \f$ is this
      /// permutation.
      /// \param n Exponent value
      /// \return This permutation raised to the n-th power
      Permutation pow(int n) const {
        // Initialize the algorithm inputs
        int power;
        Permutation value;
        if(n < 0) {
          value = inv();
          power = -n;
        } else {
          value = *this;
          power = n;
        }

        Permutation result;

        while(power) {
            if(power & 1)
              result = result.mult(value);
            value = value.mult(value);
            power >>= 1;
        }

        return result;
      }

      /// Data accessor

      /// gives direct access to \c std::map that encodes the Permutation
      /// \return \c std::map<index_type,index_type> object encoding the permutation in compressed two-line form
      const Map& data() const { return p_; }

      /// Idenity permutation factory method

      /// \return the identity permutation
      Permutation identity() const {
        return Permutation();
      }

      /// Serialize permutation

      /// MADNESS compatible serialization function
      /// \tparam Archive The serialization archive type
      /// \param[in,out] ar The serialization archive
      template <typename Archive>
      void serialize(Archive& ar) {
        ar & p_;
      }

    }; // class Permutation

    /// Permutation equality operator

    /// \param p1 The left-hand permutation to be compared
    /// \param p2 The right-hand permutation to be compared
    /// \return \c true if all elements of \c p1 and \c p2 are equal and in the
    /// same order, otherwise \c false.
    inline bool operator==(const Permutation& p1, const Permutation& p2) {
      return (p1.domain_size() == p2.domain_size())
             && p1.data() == p2.data();
    }

    /// Permutation inequality operator

    /// \param p1 The left-hand permutation to be compared
    /// \param p2 The right-hand permutation to be compared
    /// \return \c true if any element of \c p1 is not equal to that of \c p2,
    /// otherwise \c false.
    inline bool operator!=(const Permutation& p1, const Permutation& p2) {
      return ! operator==(p1, p2);
    }

    /// Permutation less-than operator

    /// \param p1 The left-hand permutation to be compared
    /// \param p2 The right-hand permutation to be compared
    /// \return \c true if the elements of \c p1 are lexicographically less than
    /// that of \c p2, otherwise \c false.
    inline bool operator<(const Permutation& p1, const Permutation& p2) {
      if (&p1 == &p2) return false;
      return std::lexicographical_compare(p1.data().begin(), p1.data().end(),
          p2.data().begin(), p2.data().end());
    }

    /// Add permutation to an output stream

    /// \param[out] output The output stream
    /// \param[in] p The permutation to be added to the output stream
    /// \return The output stream
    inline std::ostream& operator<<(std::ostream& output, const Permutation& p) {
      output << "{";
      Permutation::print_map(output, p.data());
      output << "}";
      return output;
    }

    /// Inverse permutation operator

    /// \param perm The permutation to be inverted
    /// \return \c perm.inverse()
    inline Permutation operator -(const Permutation& perm) {
      return perm.inv();
    }

    /// Permutation multiplication operator

    /// \param p1 The left-hand permutation
    /// \param p2 The right-hand permutation
    /// \return The product of p1 and p2 (which is the permutation of \c p2
    /// by \c p1).
    inline Permutation operator*(const Permutation& p1, const Permutation& p2) {
      return p1.mult(p2);
    }


    /// return *this ^ other
    inline Permutation& operator*=(Permutation& p1, const Permutation& p2) {
      return (p1 = p1 * p2);
    }

    /// Raise \c perm to the n-th power

    /// Constructs the permutation \f$ P^n \f$, where \f$ P \f$ is the
    /// permutation \c perm.
    /// \param perm The base permutation
    /// \param n Exponent value
    /// \return This permutation raised to the n-th power
    inline Permutation operator^(const Permutation& perm, int n) {
      return perm.pow(n);
    }

    /** @}*/


    /// Permute a \c std::array

    /// \tparam T The element type of the array
    /// \tparam N The size of the array
    /// \param perm The permutation
    /// \param a The array to be permuted
    /// \return A permuted copy of \c a
    /// \throw TiledArray::Exception When the dimension of the permutation is not
    /// equal to the size of \c a.
    template <typename T, std::size_t N>
    inline std::array<T,N> operator*(const Permutation& perm, const std::array<T, N>& a) {
      std::array<T,N> result;
      symmetry::detail::permute_array(perm, a, result);
      return result;
    }

    /// In-place permute a \c std::array

    /// \tparam T The element type of the array
    /// \tparam N The size of the array
    /// \param[out] a The array to be permuted
    /// \param[in] perm The permutation
    /// \return A reference to \c a
    /// \throw TiledArray::Exception When the dimension of the permutation is not
    /// equal to the size of \c a.
    template <typename T, std::size_t N>
    inline std::array<T,N>& operator*=(std::array<T,N>& a, const Permutation& perm) {
      const std::array<T,N> temp = a;
      symmetry::detail::permute_array(perm, temp, a);
      return a;
    }

    /// permute a \c std::vector<T>

    /// \tparam T The element type of the vector
    /// \tparam A The allocator type of the vector
    /// \param perm The permutation
    /// \param v The vector to be permuted
    /// \return A permuted copy of \c v
    /// \throw TiledArray::Exception When the dimension of the permutation is not
    /// equal to the size of \c v.
    template <typename T, typename A>
    inline std::vector<T> operator*(const Permutation& perm, const std::vector<T, A>& v) {
      std::vector<T> result(v.size());
      symmetry::detail::permute_array(perm, v, result);
      return result;
    }

    /// In-place permute a \c std::array

    /// \tparam T The element type of the vector
    /// \tparam A The allocator type of the vector
    /// \param[out] v The vector to be permuted
    /// \param[in] perm The permutation
    /// \return A reference to \c v
    /// \throw TiledArray::Exception When the dimension of the permutation is not
    /// equal to the size of \c v.
    template <typename T, typename A>
    inline std::vector<T, A>& operator*=(std::vector<T, A>& v, const Permutation& perm) {
      const std::vector<T, A> temp = v;
      symmetry::detail::permute_array(perm, temp, v);
      return v;
    }

  } // namespace symmetry

  /** @}*/

} // namespace TiledArray

#endif // TILEDARRAY_SYMM_PERMUTATION_H__INCLUDED
