/*!
 * Copyright (c) 2015 by Contributors
 * \file upsampling-inl.h
 * \brief
 * \author Bing Xu
*/
#ifndef MXNET_OPERATOR_UPSAMPLING_INL_H_
#define MXNET_OPERATOR_UPSAMPLING_INL_H_

#include <dmlc/logging.h>
#include <dmlc/parameter.h>
#include <mxnet/operator.h>
#include <algorithm>
#include <map>
#include <vector>
#include <string>
#include <utility>
#include "./operator_common.h"

namespace mxnet {
namespace op {

namespace up_enum {
enum UpSamplingOpInputs {kData, kWeight};
enum UpSamplingOpOutputs {kOut};
enum UpSamplingType {kNearest, kBilinear};
enum UpSamplingMultiInputMode {kConcat, kSum};
}  // namespace up_enum

struct UpSamplingParam : public dmlc::Parameter<UpSamplingParam> {
  index_t scale;
  index_t num_filter;
  int sample_type;
  int num_args;
  int multi_input_mode;
  uint64_t workspace;
  DMLC_DECLARE_PARAMETER(UpSamplingParam) {
    DMLC_DECLARE_FIELD(scale)
    .set_range(1, 1000)
    .describe("Up sampling scale");
    DMLC_DECLARE_FIELD(num_filter)
    .describe("Input filter. Only used by bilinear sample_type.")
    .set_default(0);
    DMLC_DECLARE_FIELD(sample_type)
    .add_enum("nearest", up_enum::kNearest)
    .add_enum("bilinear", up_enum::kBilinear)
    .describe("upsampling method");
    DMLC_DECLARE_FIELD(multi_input_mode)
    .add_enum("concat", up_enum::kConcat)
    .add_enum("sum", up_enum::kSum)
    .set_default(up_enum::kConcat)
    .describe("How to handle multiple input. concat means concatenate upsampled "
    "images along the channel dimension. sum means add all images together, "
    "only available for nearest neighbor upsampling.");
    DMLC_DECLARE_FIELD(num_args).set_lower_bound(1)
    .describe("Number of inputs to be upsampled. For nearest neighbor "
    "upsampling, this can be 1-N; the size of output will be"
    "(scale*h_0,scale*w_0) and all other inputs will be upsampled to the"
    "same size. For bilinear upsampling this must be 2; 1 input and 1 weight.");
    DMLC_DECLARE_FIELD(workspace).set_default(512).set_range(0, 8192)
    .describe("Tmp workspace for deconvolution (MB)");
  }
};  // struct UpSamplingParam

template<typename xpu, typename DType>
class UpSamplingNearestOp : public Operator {
 public:
  explicit UpSamplingNearestOp(UpSamplingParam p) {
    this->param_ = p;
  }

  virtual void Forward(const OpContext &ctx,
                       const std::vector<TBlob> &in_data,
                       const std::vector<OpReqType> &req,
                       const std::vector<TBlob> &out_data,
                       const std::vector<TBlob> &aux_args) {
    using namespace mshadow;
    using namespace mshadow::expr;
    CHECK_EQ(in_data.size(), static_cast<size_t>(param_.num_args));
    CHECK_EQ(out_data.size(), 1U);
    if (req[up_enum::kOut] == kNullOp) {
      return;
    }
    Stream<xpu> *s = ctx.get_stream<xpu>();
    Tensor<xpu, 4, DType> out = out_data[up_enum::kOut].get<xpu, 4, DType>(s);
    if (param_.num_args > 1) {
      int begin = 0;
      for (int i = 0; i < param_.num_args; ++i) {
        Tensor<xpu, 4, DType> data = in_data[i].get<xpu, 4, DType>(s);
        int end = begin + data.size(1);
        int scale = out_data[up_enum::kOut].size(2)/in_data[i].size(2);
        if (param_.multi_input_mode == up_enum::kSum) {
          if (i == 0) {
            Assign(out, req[up_enum::kOut], upsampling_nearest(data, scale));
          } else {
            out += upsampling_nearest(data, scale);
          }
        } else {
          Assign(slice<1>(out, begin, end), req[up_enum::kOut], upsampling_nearest(data, scale));
        }
        begin = end;
      }
    } else {
      Tensor<xpu, 4, DType> data = in_data[up_enum::kData].get<xpu, 4, DType>(s);
      Assign(out, req[up_enum::kOut], upsampling_nearest(data, param_.scale));
    }
  }

  virtual void Backward(const OpContext &ctx,
                        const std::vector<TBlob> &out_grad,
                        const std::vector<TBlob> &in_data,
                        const std::vector<TBlob> &out_data,
                        const std::vector<OpReqType> &req,
                        const std::vector<TBlob> &in_grad,
                        const std::vector<TBlob> &aux_args) {
    using namespace mshadow;
    using namespace mshadow::expr;
    CHECK_EQ(out_grad.size(), 1U);
    CHECK_EQ(in_grad.size(), static_cast<size_t>(param_.num_args));
    Stream<xpu> *s = ctx.get_stream<xpu>();
    Tensor<xpu, 4, DType> grad = out_grad[up_enum::kOut].get<xpu, 4, DType>(s);
    if (param_.num_args > 1) {
      int begin = 0;
      for (int i = 0; i < param_.num_args; ++i) {
        Tensor<xpu, 4, DType> input_grad = in_grad[i].get<xpu, 4, DType>(s);
        mshadow::Shape<2> in_shape = Shape2(input_grad.shape_[2], input_grad.shape_[3]);
        int end = begin + input_grad.size(1);
        int scale = grad.size(2)/in_shape[0];
        if (param_.multi_input_mode == up_enum::kSum) {
          Assign(input_grad, req[i],
                 pool<mshadow::red::sum>(grad,
                                         in_shape,
                                         scale,
                                         scale,
                                         scale,
                                         scale));
        } else {
          Assign(input_grad, req[i],
                 pool<mshadow::red::sum>(slice<1>(grad, begin, end),
                                         in_shape,
                                         scale,
                                         scale,
                                         scale,
                                         scale));
        }
        begin = end;
      }
    } else {
      Tensor<xpu, 4, DType> input_grad = in_grad[up_enum::kData].get<xpu, 4, DType>(s);
      mshadow::Shape<2> in_shape = Shape2(input_grad.shape_[2], input_grad.shape_[3]);
      Assign(input_grad, req[up_enum::kData],
             pool<mshadow::red::sum>(grad,
                                     in_shape,
                                     param_.scale,
                                     param_.scale,
                                     param_.scale,
                                     param_.scale));
    }
  }

 private:
  UpSamplingParam param_;
};  // class UpSamplingNearestOp

template<typename xpu>
Operator *CreateOp(UpSamplingParam param, int dtype);


#if DMLC_USE_CXX11
class UpSamplingProp : public OperatorProperty {
 public:
  void Init(const std::vector<std::pair<std::string, std::string> >& kwargs) override {
    param_.Init(kwargs);
  }

  std::map<std::string, std::string> GetParams() const override {
    return param_.__DICT__();
  }

  std::vector<std::string> ListArguments() const override {
    if (param_.sample_type == up_enum::kNearest) {
      std::vector<std::string> ret;
      for (int i = 0; i < param_.num_args; ++i) {
        ret.push_back(std::string("arg") + std::to_string(i));
      }
      return ret;
    } else {
      return {"data", "weight"};
    }
  }

  bool InferShape(std::vector<TShape> *in_shape,
                  std::vector<TShape> *out_shape,
                  std::vector<TShape> *aux_shape) const override {
    CHECK_GE(in_shape->size(), 1U);
    const TShape &dshape = (*in_shape)[0];
    TShape oshape = dshape;
    if (param_.sample_type == up_enum::kNearest) {
      CHECK_EQ(in_shape->size(), static_cast<size_t>(param_.num_args));
      oshape[1] = 0;
      for (auto& shape : *in_shape) {
        CHECK_EQ(shape.ndim(), 4U) << \
          "UpSamplingNearest: Input data should be 4D in (batch, channel, y, x)";
        int oh = dshape[2]*param_.scale, ow = dshape[3]*param_.scale;
        CHECK_EQ(oh%shape[2], 0U) << "UpSamplingNearest: input height of " << shape[2] << \
          "does not divide output height of " << oh;
        CHECK_EQ(ow%shape[3], 0U) << "UpSamplingNearest: input width of " << shape[3] << \
          "does not divide output width of " << ow;
        if (param_.multi_input_mode == up_enum::kSum) {
          CHECK(oshape[1] == 0 || oshape[1] == shape[1]) << \
            "Number of channels must be the same when multi_input_mode==sum";
          oshape[1] = shape[1];
        } else {
          oshape[1] += shape[1];
        }
      }
    } else {
      CHECK_EQ(in_shape->size(), 2U) << "Input:[data, weight]";
      CHECK_EQ(dshape.ndim(), 4U) << \
        "UpSamplingBilinear: Input data should be 4D in (batch, channel, y, x)";
      if (dshape.ndim() ==  0) return false;
      int kernel = 2 * param_.scale - param_.scale % 2;
      SHAPE_ASSIGN_CHECK(*in_shape,
                         up_enum::kWeight,
                         mshadow::Shape4(dshape[1], 1, kernel, kernel));
      oshape = dshape;
    }
    oshape[2] = dshape[2] * param_.scale;
    oshape[3] = dshape[3] * param_.scale;
    out_shape->clear();
    out_shape->push_back(oshape);
    return true;
  }

  bool InferType(std::vector<int> *in_type,
                 std::vector<int> *out_type,
                 std::vector<int> *aux_type) const override {
    CHECK_GE(in_type->size(), 1U);
    int dtype = (*in_type)[0];
    CHECK_NE(dtype, -1) << "First input must have specified type";
    for (index_t i = 0; i < in_type->size(); ++i) {
      if ((*in_type)[i] == -1) {
        (*in_type)[i] = dtype;
      } else {
        CHECK_EQ((*in_type)[i], dtype) << "This layer requires uniform type. "
                                       << "Expected " << dtype << " v.s. given "
                                       << (*in_type)[i] << " at " << ListArguments()[i];
      }
    }
    out_type->clear();
    out_type->push_back(dtype);
    return true;
  }

  OperatorProperty* Copy() const override {
    auto ptr = new UpSamplingProp();
    ptr->param_ = this->param_;
    return ptr;
  }

  std::string TypeString() const override {
    return "UpSampling";
  }

  std::vector<int> DeclareBackwardDependency(
    const std::vector<int> &out_grad,
    const std::vector<int> &in_data,
    const std::vector<int> &out_data) const override {
    if (param_.sample_type == up_enum::kNearest) {
      return {out_grad[up_enum::kOut]};
    } else {
      return {out_grad[up_enum::kOut], in_data[up_enum::kData], in_data[up_enum::kWeight]};
    }
  }

  std::vector<std::pair<int, void*> > BackwardInplaceOption(
    const std::vector<int> &out_grad,
    const std::vector<int> &in_data,
    const std::vector<int> &out_data,
    const std::vector<void*> &in_grad) const override {
    return {};
  }

  std::vector<ResourceRequest> ForwardResource(
      const std::vector<TShape> &in_shape) const override {
    if (param_.sample_type == up_enum::kNearest) {
      return {};
    } else {
      return {ResourceRequest::kTempSpace};
    }
  }

  std::vector<ResourceRequest> BackwardResource(
      const std::vector<TShape> &in_shape) const override {
    if (param_.sample_type == up_enum::kNearest) {
      return {};
    } else {
      return {ResourceRequest::kTempSpace};
    }
  }

  Operator* CreateOperator(Context ctx) const override {
    LOG(FATAL) << "Not Implemented";
    return NULL;
  }

  Operator* CreateOperatorEx(Context ctx, std::vector<TShape> *in_shape,
                             std::vector<int> *in_type) const override;


 private:
  UpSamplingParam param_;
};  // class UpSamplingProp
#endif  // DMLC_USE_CXX11
}  // namespace op
}  // namespace mxnet

#endif  // MXNET_OPERATOR_UPSAMPLING_INL_H_