/*
 * 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 operator.h
* \brief definition of operator
* \author Chuntao Hong, Zhang Chen
*/

#ifndef MXNET_CPP_OPERATOR_H_
#define MXNET_CPP_OPERATOR_H_

#include <map>
#include <string>
#include <vector>
#include "mxnet-cpp/base.h"
#include "mxnet-cpp/op_map.h"
#include "mxnet-cpp/symbol.h"

namespace mxnet {
namespace cpp {
class Mxnet;
/*!
* \brief Operator interface
*/
class Operator {
 public:
  /*!
  * \brief Operator constructor
  * \param operator_name type of the operator
  */
  explicit Operator(const std::string &operator_name);
  Operator &operator=(const Operator &rhs);
  /*!
  * \brief set config parameters
  * \param name name of the config parameter
  * \param value value of the config parameter
  * \return reference of self
  */
  template <typename T>
  Operator &SetParam(const std::string &name, const T &value) {
    std::string value_str;
    std::stringstream ss;
    ss << value;
    ss >> value_str;

    params_[name] = value_str;
    return *this;
  }
  /*!
  * \brief set config parameters from positional inputs
  * \param pos the position of parameter
  * \param value value of the config parameter
  * \return reference of self
  */
  template <typename T>
  Operator &SetParam(int pos, const T &value) {
    std::string value_str;
    std::stringstream ss;
    ss << value;
    ss >> value_str;

    params_[arg_names_[pos]] = value_str;
    return *this;
  }
  /*!
  * \brief add an input symbol
  * \param name name of the input symbol
  * \param symbol the input symbol
  * \return reference of self
  */
  Operator &SetInput(const std::string &name, const Symbol &symbol);
  /*!
  * \brief add an input symbol
  * \param symbol the input symbol
  */
  template<int N = 0>
  void PushInput(const Symbol &symbol) {
    input_symbols_.push_back(symbol.GetHandle());
  }
  /*!
  * \brief add input symbols
  * \return reference of self
  */
  Operator &operator()() { return *this; }
  /*!
  * \brief add input symbols
  * \param symbol the input symbol
  * \return reference of self
  */
  Operator &operator()(const Symbol &symbol) {
    input_symbols_.push_back(symbol.GetHandle());
    return *this;
  }
  /*!
  * \brief add a list of input symbols
  * \param symbols the vector of the input symbols
  * \return reference of self
  */
  Operator &operator()(const std::vector<Symbol> &symbols) {
    for (auto &s : symbols) {
      input_symbols_.push_back(s.GetHandle());
    }
    return *this;
  }
  /*!
  * \brief create a Symbol from the current operator
  * \param name the name of the operator
  * \return the operator Symbol
  */
  Symbol CreateSymbol(const std::string &name = "");

  /*!
  * \brief add an input ndarray
  * \param name name of the input ndarray
  * \param ndarray the input ndarray
  * \return reference of self
  */
  Operator &SetInput(const std::string &name, const NDArray &ndarray);
  /*!
  * \brief add an input ndarray
  * \param ndarray the input ndarray
  */
  template<int N = 0>
  Operator &PushInput(const NDArray &ndarray) {
    input_ndarrays_.push_back(ndarray.GetHandle());
    return *this;
  }
  /*!
  * \brief add positional inputs
  */
  template <class T, class... Args, int N = 0>
  Operator &PushInput(const T &t, Args... args) {
    SetParam(N, t);
    PushInput<Args..., N+1>(args...);
    return *this;
  }
  /*!
  * \brief add the last positional input
  */
  template <class T, int N = 0>
  Operator &PushInput(const T &t) {
    SetParam(N, t);
    return *this;
  }
  /*!
  * \brief add input ndarrays
  * \param ndarray the input ndarray
  * \return reference of self
  */
  Operator &operator()(const NDArray &ndarray) {
    input_ndarrays_.push_back(ndarray.GetHandle());
    return *this;
  }
  /*!
  * \brief add a list of input ndarrays
  * \param ndarrays the vector of the input ndarrays
  * \return reference of self
  */
  Operator &operator()(const std::vector<NDArray> &ndarrays) {
    for (auto &s : ndarrays) {
      input_ndarrays_.push_back(s.GetHandle());
    }
    return *this;
  }
  /*!
  * \brief add input ndarrays
  * \return reference of self
  */
  template <typename... Args>
  Operator &operator()(Args... args) {
    PushInput(args...);
    return *this;
  }
  std::vector<NDArray> Invoke();
  void Invoke(NDArray &output);
  void Invoke(std::vector<NDArray> &outputs);

 private:
  std::map<std::string, std::string> params_desc_;
  bool variable_params_ = false;
  std::map<std::string, std::string> params_;
  std::vector<SymbolHandle> input_symbols_;
  std::vector<NDArrayHandle> input_ndarrays_;
  std::vector<std::string> input_keys_;
  std::vector<std::string> arg_names_;
  AtomicSymbolCreator handle_;
  static OpMap*& op_map();
};
}  // namespace cpp
}  // namespace mxnet

#endif  // MXNET_CPP_OPERATOR_H_