/*
 * 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.
 */

/*!
 * \file shape.h
 * \brief definition of shape
 * \author Chuntao Hong, Zhang Chen
 */

#ifndef MXNET_CPP_SHAPE_H_
#define MXNET_CPP_SHAPE_H_

#include <istream>
#include <ostream>
#include <algorithm>
#include <vector>
#include "mxnet-cpp/base.h"

namespace mxnet {
namespace cpp {

/*!
 * \brief dynamic shape class that can hold shape
 *   of arbirary dimension
 */
struct Shape {
 public:
  /*! \brief constructor */
  Shape() : ndim_(0), num_heap_allocated_(0), data_heap_(nullptr) {}
  /*!
   * \brief constructor from a vector of index_t
   * \param v the vector
   */
  explicit Shape(const std::vector<index_t>& v) : ndim_(v.size()) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      std::copy(v.begin(), v.end(), data_stack_);
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      std::copy(v.begin(), v.end(), data_heap_);
    }
  }
  /*!
   * \brief constructor one dimmension shape
   * \param s1 size of the first dimmension
   */
  explicit Shape(index_t s1) : ndim_(1) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
    }
  }
  /*!
   * \brief constructor two dimmension shape
   * \param s1 size of the first dimmension
   * \param s2 size of the second dimmension
   */
  Shape(index_t s1, index_t s2) : ndim_(2) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
    }
  }
  /*!
   * \brief constructor three dimmension shape
   * \param s1 size of the first dimmension
   * \param s2 size of the second dimmension
   * \param s3 size of the third dimmension
   */
  Shape(index_t s1, index_t s2, index_t s3) : ndim_(3) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
      data_stack_[2]      = s3;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
      data_heap_[2]       = s3;
    }
  }
  /*!
   * \brief constructor four dimmension shape
   * \param s1 size of the first dimmension
   * \param s2 size of the second dimmension
   * \param s3 size of the third dimmension
   * \param s4 size of the fourth dimmension
   */
  Shape(index_t s1, index_t s2, index_t s3, index_t s4) : ndim_(4) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
      data_stack_[2]      = s3;
      data_stack_[3]      = s4;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
      data_heap_[2]       = s3;
      data_heap_[3]       = s4;
    }
  }
  /*!
   * \brief constructor five dimmension shape
   * \param s1 size of the first dimmension
   * \param s2 size of the second dimmension
   * \param s3 size of the third dimmension
   * \param s4 size of the fourth dimmension
   * \param s5 size of the fifth dimmension
   */
  Shape(index_t s1, index_t s2, index_t s3, index_t s4, index_t s5) : ndim_(5) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      data_stack_[0]      = s1;
      data_stack_[1]      = s2;
      data_stack_[2]      = s3;
      data_stack_[3]      = s4;
      data_stack_[4]      = s5;
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      data_heap_[0]       = s1;
      data_heap_[1]       = s2;
      data_heap_[2]       = s3;
      data_heap_[3]       = s4;
      data_heap_[4]       = s5;
    }
  }
  /*!
   * \brief constructor from Shape
   * \param s the source shape
   */
  Shape(const Shape& s) : ndim_(s.ndim_) {
    if (ndim_ <= kStackCache) {
      data_heap_          = nullptr;
      num_heap_allocated_ = 0;
      std::copy(s.data_stack_, s.data_stack_ + ndim_, data_stack_);
    } else {
      data_heap_          = new index_t[ndim_];
      num_heap_allocated_ = ndim_;
      std::copy(s.data_heap_, s.data_heap_ + ndim_, data_heap_);
    }
  }
#if MSHADOW_IN_CXX11
  /*!
   * \brief move constructor from Shape
   * \param s the source shape
   */
  Shape(Shape&& s)
      : ndim_(s.ndim_), num_heap_allocated_(s.num_heap_allocated_), data_heap_(s.data_heap_) {
    if (ndim_ <= kStackCache) {
      std::copy(s.data_stack_, s.data_stack_ + ndim_, data_stack_);
    }
    // remove data heap space from s
    s.data_heap_ = nullptr;
  }
#endif
  /*! \brief destructor */
  ~Shape() {
    // data_heap_ can be nullptr
    delete[] data_heap_;
  }
  /*!
   * \brief copy shape from content betwen two iterators
   * \param begin the beginning of iterator
   * \param end the end of the iterator
   * \tparam RandomAccessIterator iterator type
   */
  template <typename RandomAccessIterator>
  inline void CopyFrom(RandomAccessIterator begin, RandomAccessIterator end) {
    this->SetDim(end - begin);
    std::copy(begin, end, data());
  }
  /*!
   * \brief assignment from shape
   * \param shape source shape
   * \return reference of self
   */
  inline Shape& operator=(const Shape& shape) {
    this->SetDim(shape.ndim_);
    const index_t* src = shape.data();
    std::copy(src, src + ndim_, data());
    return *this;
  }
  /*!
   * \brief assignment from vector
   * \param shape source shape
   * \return reference of self
   */
  inline Shape& operator=(const std::vector<index_t>& shape) {
    this->CopyFrom(shape.begin(), shape.end());
    return *this;
  }
  /*! \return the data content of the shape */
  inline const index_t* data() const {
    return ndim_ <= kStackCache ? data_stack_ : data_heap_;
  }
  /*! \return the data content of the shape */
  inline index_t* data() {
    return ndim_ <= kStackCache ? data_stack_ : data_heap_;
  }
  /*! \brief return number of dimension of the tensor inside */
  inline index_t ndim(void) const {
    return ndim_;
  }
  /*!
   * \brief get corresponding index
   * \param i dimension index
   * \return the corresponding dimension size
   */
  inline index_t& operator[](index_t i) {
    return data()[i];
  }
  /*!
   * \brief get corresponding index
   * \param i dimension index
   * \return the corresponding dimension size
   */
  inline const index_t& operator[](index_t i) const {
    return data()[i];
  }
  /*! \brief total number of elements in the tensor */
  inline size_t Size(void) const {
    size_t size      = 1;
    const index_t* d = this->data();
    for (index_t i = 0; i < ndim_; ++i) {
      size *= d[i];
    }
    return size;
  }
  /*!
   * \return whether two shape equals
   * \param s the shape to compare against
   */
  inline bool operator==(const Shape& s) const {
    if (ndim_ != s.ndim_)
      return false;
    if (ndim_ <= kStackCache) {
      for (index_t i = 0; i < ndim_; ++i) {
        if (data_stack_[i] != s.data_stack_[i])
          return false;
      }
    } else {
      for (index_t i = 0; i < ndim_; ++i) {
        if (data_heap_[i] != s.data_heap_[i])
          return false;
      }
    }
    return true;
  }
  /*!
   * \return whether two shape not equals
   * \param s the shape to compare against
   */
  inline bool operator!=(const Shape& s) const {
    return !(*this == s);
  }

  friend std::ostream& operator<<(std::ostream& os, const Shape& shape);
  friend std::istream& operator>>(std::istream& is, Shape& shape);

 private:
  // the shape will be stored in data_stack_
  // when dimension is smaller than kStackCache
  // when it is bigger, it will be stored in data_heap_;
  /*! \brief size of in stack space */
  static const index_t kStackCache = 5;
  /*! \brief number of dimnsion of the shape */
  index_t ndim_;
  /*! \brief number of cells allocated in data_heap_ */
  index_t num_heap_allocated_;
  /*! \brief in stack space used to store shape when it is small */
  index_t data_stack_[kStackCache];
  /*! \brief space to store shape when dimension is big*/
  index_t* data_heap_;
  /*!
   * \brief internal function to set the dimension
   * \param dim the dimension of the shape
   */
  inline void SetDim(index_t dim) {
    if (dim > kStackCache && dim > num_heap_allocated_) {
      // data_heap_ can be nullptr
      delete[] data_heap_;
      data_heap_          = new index_t[dim];
      num_heap_allocated_ = dim;
    }
    ndim_ = dim;
  }
};

/*!
 * \brief allow string printing of the shape
 * \param os the output stream
 * \param shape the shape
 * \return the ostream
 */
inline std::ostream& operator<<(std::ostream& os, const Shape& shape) {
  os << '(';
  for (index_t i = 0; i < shape.ndim(); ++i) {
    if (i != 0)
      os << ',';
    os << static_cast<int>(shape[i]);  // Supports negative Shape 'special codes' for inferring
  }
  // python style tuple
  if (shape.ndim() == 1)
    os << ',';
  os << ')';
  return os;
}

/*!
 * \brief read shape from the istream
 * \param is the input stream
 * \param shape the shape
 * \return the istream
 */
inline std::istream& operator>>(std::istream& is, Shape& shape) {
  // get (
  while (true) {
    char ch = is.get();
    if (ch == '(')
      break;
    if (!isspace(ch)) {
      is.setstate(std::ios::failbit);
      return is;
    }
  }
  index_t idx;
  std::vector<index_t> tmp;
  while (is >> idx) {
    tmp.push_back(idx);
    char ch;
    do {
      ch = is.get();
    } while (isspace(ch));
    if (ch == ',') {
      while (true) {
        ch = is.peek();
        if (isspace(ch)) {
          is.get();
          continue;
        }
        if (ch == ')') {
          is.get();
          break;
        }
        break;
      }
      if (ch == ')')
        break;
    } else if (ch == ')') {
      break;
    } else {
      is.setstate(std::ios::failbit);
      return is;
    }
  }
  shape.CopyFrom(tmp.begin(), tmp.end());
  return is;
}

}  // namespace cpp
}  // namespace mxnet

#endif  // MXNET_CPP_SHAPE_H_