/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 * 
 *   http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 *  Copyright (c) 2018 by Contributors
 * \file nnvm/layout.h
 * \brief Layout expression.
 *        The layout is composed of upper cases, lower cases and numbers,
 *        where upper case indicates a (super-)dimension and
 *        the corresponding lower case with factor size indicates the split (sub-)dimension.
 *        For example, NCHW16c can describe a 5-D tensor of
 *        [batch_size, channel, height, width, channel_block].
 *        Here sub-dimension channel_block=16 is the split of super-dimension C (channel).
 */
#ifndef NNVM_LAYOUT_H_
#define NNVM_LAYOUT_H_

#include <dmlc/parameter.h>
#include <string>
#include <sstream>
#include <vector>
#include <utility>
#include <algorithm>

namespace nnvm {

class Layout {
 public:
  using LayoutDim = char;

  /*! \brief default constructor */
  Layout() : name_("__undef__") {} // NOLINT(*)

  /*!
   * \brief construct from a string.
   * \param layout input in layout convention:
   *        upper case indicates a dimension and
   *        the corresponding lower case with factor size
   *        indicates the split dimension.
   *        return undefined layout if "__undef__" is passed.
   */
  inline Layout(const std::string& layout) { // NOLINT(*)
    parse(layout);
  }
  /*!
   * \brief copy constructor from another layout
   * \param s the source layout
   */
  inline Layout(const Layout& s) { // NOLINT(*)
    this->parse(s.name_);
  }
  /*!
   * \brief move constructor from Layout
   * \param src the source layout
   */
  inline Layout(Layout&& src) { // NOLINT(*)
    this->swap(src);
  }
  /*!
   * \brief assignment from another layout.
   * \param src source layout
   * \return reference of self
   */
  inline Layout& operator=(const Layout& src) {
    this->parse(src.name_);
    return *this;
  }
  /*!
   * \brief assignment from rvalue of another layout.
   * \param src source layout
   * \return reference of self
   */
  inline Layout& operator=(Layout&& src) {
    Layout(std::move(src)).swap(*this); // NOLINT(*)
    return *this;
  }
  /*!
   * \brief assignment from string.
   * \param src source layout
   * \return reference of self
   */
  inline Layout& operator=(const std::string& src) {
    this->parse(src);
    return *this;
  }
  /*!
   * \return whether two layout equals
   * \param s the layout to compare against
   */
  inline bool operator==(const Layout& s) const {
    return name_ == s.name_;
  }
  /*!
   * \return whether two layout not equal
   * \param s the layout to compare against
   */
  inline bool operator!=(const Layout& s) const {
    return !(*this == s);
  }

  /*!
   * \brief Append the current layout by another.
   * @param other the layout to be appended
   * @return a new layout
   */
  inline Layout operator+(const Layout& other) const {
    if (!this->defined() && !other.defined()) {
      return Layout::Undef();
    } else if (!this->defined()) {
      return other;
    } else if (!other.defined()) {
      return *this;
    }
    return Layout(this->name_ + other.name_);
  }

  /*!
   * \brief Check whether a given dimension is a super-dimension.
   * \param dim input dimension
   * \return Whether a given dimension is a super-dimension.
   */
  static inline bool is_superdim(LayoutDim dim) {
    return dim >= 'A' && dim <= 'Z';
  }

  /*!
   * \brief Check whether a given dimension is a sub-dimension.
   * \param dim input dimension
   * \return Whether a given dimension is a sub-dimension.
   */
  static inline bool is_subdim(LayoutDim dim) {
    return dim >= 'a' && dim <= 'z';
  }

  /*!
   * \brief Convert a given dimension to super-dimension.
   * \param dim input dimension
   * \return The converted description.
   */
  static inline LayoutDim to_superdim(LayoutDim dim) {
    if (is_subdim(dim)) {
      return dim - 'a' + 'A';
    }
    return dim;
  }

  /*!
   * \brief Convert a given dimension to sub-dimension.
   * \param dim input dimension
   * \return The converted description.
   */
  static inline LayoutDim to_subdim(LayoutDim dim) {
    if (is_superdim(dim)) {
      return dim - 'A' + 'a';
    }
    return dim;
  }

  /*!
   * \brief Return an undefined layout.
   * \return a (global) undefined layout.
   */
  static inline const Layout& Undef() {
    static Layout undef;
    return undef;
  }

  /*!
   * \brief Swap current object with other
   * \param other another object to be swapped.
   */
  inline void swap(Layout& other) {  // NOLINT(*)
    std::swap(name_, other.name_);
    std::swap(superdim_pos_, other.superdim_pos_);
    std::swap(subdim_pos_, other.subdim_pos_);
    std::swap(subdim_size_, other.subdim_size_);
    std::swap(layout_simplified_, other.layout_simplified_);
  }

  /*!
   * \brief Two layouts are convertible only if
   *        they have same set of super-dimensions.
   *        e.g., NCHW, NCHW16c, NHWC are convertible between each other,
   *        but NCHW, CHW, OIHW are not.
   * \param dst the target layout
   * \return Whether can be converted to dst layout.
   */
  inline bool convertible(const Layout &dst) const {
    if (!this->defined() || !dst.defined()) return false;
    for (size_t i = 0; i < kUniqueDim; ++i) {
      if ((superdim_pos_[i] >= 0 && dst.superdim_pos_[i] < 0) ||
          (superdim_pos_[i] < 0 && dst.superdim_pos_[i] >= 0)) {
        return false;
      }
    }
    return true;
  }

  /*!
   * \brief Returns a sublayout which is the portion of the object
   *        that starts at dimension \p pos and spans \p len dimensions
   *        (or until the end of the layout, whichever comes first).
   * \param pos The start position.
   * \param len The length of the sub-layout.
   * \return A newly constructed Layout object.
   */
  inline Layout sublayout(size_t pos, size_t len) const {
    if (pos > ndim()) return Layout::Undef();
    if (pos + len > ndim()) len = ndim() - pos;
    if (len == 0) return Layout::Undef();
    std::ostringstream new_layout;
    for (size_t i = pos; i < pos + len; ++i) {
      if (is_subdim(layout_simplified_[i])) {
        auto block_size = this->subsizeof(layout_simplified_[i]);
        CHECK_GT(block_size, 0);
        new_layout << block_size;
      }
      new_layout << layout_simplified_[i];
    }
    return Layout(new_layout.str());
  }

  /*! \return A newly constructed reversed Layout object. */
  inline Layout reverse() const {
    if (!this->defined()) return Layout::Undef();
    std::ostringstream new_layout;
    for (int64_t i = this->ndim() - 1; i >= 0; --i) {
      if (is_subdim(layout_simplified_[i])) {
        auto block_size = this->subsizeof(layout_simplified_[i]);
        CHECK_GT(block_size, 0);
        new_layout << block_size;
      }
      new_layout << layout_simplified_[i];
    }
    return Layout(new_layout.str());
  }

  /*!
   * \brief Split \p dim by \p size and put the sub-dimension to position \p target_pos.
   * \param dim The source dimension to be split. It must be a super-dimension.
   * \param target_pos The target position of the newly split sub-dimension.
   * \param size size of the sub-dimension.
   * \return A newly constructed Layout object.
   */
  inline Layout split(LayoutDim dim, size_t target_pos, uint32_t size) const {
    CHECK(target_pos <= this->ndim()) << "Invalid split position "
                                      << target_pos << " for layout " << name_;
    CHECK(is_superdim(dim)) << "Cannot split a sub-dimension " << dim;
    CHECK(this->contains(dim)) << "Axis " << dim << " does not exist in " << name_;
    CHECK(!this->contains(to_subdim(dim))) << "Dimension " << dim
                                           << " has already been split in "
                                           << name_;
    CHECK(size > 0) << "Invalid split size " << size;
    std::ostringstream new_layout;
    for (size_t i = 0; i <= this->ndim(); ++i) {
      if (i == target_pos) {
        new_layout << size << Layout::to_subdim(dim);
      }
      if (i == this->ndim()) break;
      new_layout << this->at(i);
    }
    Layout x(new_layout.str());
    return x;
  }

  using iterator = std::vector<LayoutDim>::const_iterator;
  using reverse_iterator = std::vector<LayoutDim>::const_reverse_iterator;

  /*! \return begin iterator */
  inline iterator begin() const {
    return layout_simplified_.begin();
  }
  /*! \return end iterator */
  inline iterator end() const {
    return layout_simplified_.end();
  }
  /*! \return rbegin iterator */
  inline reverse_iterator rbegin() const {
    return layout_simplified_.rbegin();
  }
  /*! \return rend iterator */
  inline reverse_iterator rend() const {
    return layout_simplified_.rend();
  }

  /*! \return number of dimensions */
  inline size_t ndim() const {
    return layout_simplified_.size();
  }

  /*!
   * \brief The description of the \p i-th dimension.
   *        If it is a sub-dimension, the size will be returned as well,
   *        e.g., 16c. Otherwise a single character is returned, e.g., C.
   * \param i The position
   * \return the description of the dimension.
   */
  inline std::string at(size_t i) const {
    CHECK_LT(i, this->ndim()) << "position " << i
                              << " exceeds ndim=" << this->ndim();
    std::ostringstream repr;
    if (is_subdim(layout_simplified_[i])) {
      auto factor = subsizeof(layout_simplified_[i]);
      CHECK_GT(factor, 0);
      repr << factor;
    }
    repr << layout_simplified_[i];
    return repr.str();
  }

  /*!
   * \brief return the index of the input dimension.
   *        If it is not found in the layout or the layout is undefined,
   *        return -1.
   * \param dim the input dimension.
   * \return the index or -1 if not found.
   */
  inline int32_t indexof(LayoutDim dim) const {
    if (!this->defined()) return -1;
    else if (is_superdim(dim)) return superdim_pos_[dim - 'A'];
    else if (is_subdim(dim)) return subdim_pos_[dim - 'a'];
    return -1;
  }

  /*!
   * \param dim the input super-dimension or sub-dimension.
   * \return the size of the sub-dimension of \p dim (if \p dim is a super-dimension),
   *         or the size of \p dim itself (if \p dim is a sub-dimension).
   *         Return -1 if \p dim is not in the layout or the layout is undefined.
   */
  inline int64_t subsizeof(LayoutDim dim) const {
    CHECK(is_superdim(dim) || is_subdim(dim)) << "Invalid dim " << dim;
    if (!this->defined() || !this->contains(to_subdim(dim))) {
      return -1;
    }
    int idx = to_subdim(dim) - 'a';
    return subdim_size_[idx];
  }

  /*!
   * \brief Whether the layout contains a dimension.
   * \param dim dimension to be checked.
   * \return Whether the layout contains the dimension.
   */
  inline bool contains(LayoutDim dim) const {
    if (is_superdim(dim)) {
      return superdim_pos_[dim-'A'] >= 0;
    } else if (is_subdim(dim)) {
      return subdim_pos_[dim-'a'] >= 0;
    }
    return false;
  }

  inline LayoutDim operator[](size_t i) const {
    return layout_simplified_[i];
  }

  /*! \return whether the layout is defined */
  inline bool defined() const {
    return name_ != "__undef__";
  }

  /*! \return the string description of the layout */
  inline const std::string& name() const {
    return name_;
  }

  /*!
   * \brief Write layout in JSON format.
   * \param writer JSONWriter
   */
  inline void Save(dmlc::JSONWriter* writer) const {
    writer->Write(name_);
  }

  /*!
   * \brief Load layout from JSON.
   * \param reader JSONReader
   */
  inline void Load(dmlc::JSONReader* reader) {
    std::string tmp;
    reader->Read(&tmp);
    this->parse(tmp);
  }

  /*!
   * \brief allow output string of layout to ostream
   * \param os the output stream
   * \param l the layout
   * \return the ostream
   */
  friend std::ostream& operator<<(std::ostream& os, const Layout& l) {
    os << l.name_;
    return os;
  }

 private:
  static const uint32_t kUniqueDim = 26;

  std::string name_;
  int32_t superdim_pos_[kUniqueDim];
  int32_t subdim_pos_[kUniqueDim];
  int64_t subdim_size_[kUniqueDim];
  std::vector<LayoutDim> layout_simplified_;

  void parse(const std::string& layout) {
    name_ = layout;
    std::fill_n(superdim_pos_, kUniqueDim, -1);
    std::fill_n(subdim_pos_, kUniqueDim, -1);
    std::fill_n(subdim_size_, kUniqueDim, -1);
    layout_simplified_.clear();

    if (layout == "__undef__") return;

    int32_t factor = 0;
    uint32_t curr = 0;
    for (size_t i = 0; i < layout.size(); ++i) {
      const LayoutDim c = layout.at(i);
      if (is_superdim(c)) {
        int pos = c - 'A';
        CHECK_EQ(factor, 0) << "Invalid layout " << layout
                            << ": invalid factor size " << factor
                            << " before dimension " << c;
        CHECK_EQ(superdim_pos_[pos], -1) << "Invalid layout " << layout
                                           << ": duplicate dimension " << c;
        superdim_pos_[pos] = curr++;
        layout_simplified_.push_back(c);
      } else if (is_subdim(c)) {
        int pos = c - 'a';
        CHECK_GT(factor, 0) << "Invalid layout " << layout << ": invalid factor size "
                            << factor << " for dimension " << c;
        CHECK_EQ(subdim_pos_[pos], -1) << "Invalid layout " << layout
                                           << ": duplicate dimension " << c;
        CHECK_EQ(subdim_size_[pos], -1) << "Invalid layout " << layout
                                         << ": duplicate dimension " << c;
        subdim_pos_[pos] = curr++;
        subdim_size_[pos] = factor;
        layout_simplified_.push_back(c);
        factor = 0;
      } else if (c >= '0' && c <= '9') {
        CHECK(factor >= 0) << "Invalid layout " << layout << ": _ is adjacent to a number.";
        factor = factor * 10 + c - '0';
      } else {
        LOG(FATAL) << "Invalid layout " << layout;
      }
    }
    CHECK(!layout_simplified_.empty()) << "Invalid layout " << layout;
    for (LayoutDim dim : layout_simplified_) {
      CHECK(is_superdim(dim) || superdim_pos_[dim-'a'] >= 0)
        << "Invalid layout " << layout << ": missing axis "
        << static_cast<char>(dim - 'a' + 'A');
    }
  }
};

}  // namespace nnvm

#endif  // NNVM_LAYOUT_H_