/**
 * Copyright 2019 Huawei Technologies Co., Ltd
 *
 * Licensed 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.
 */

#ifndef PASS_CONVOLUTION_MODEL_H_
#define PASS_CONVOLUTION_MODEL_H_

#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/operation.h>

#include "ir_pass.h"
#include "poly/poly_util.h"

namespace akg {
namespace ir {
#define IS_ATTR_EXIST(attrs_, key_) ((attrs_).count(key_) > 0)

#define GET_INTIMM_ATTR(attrs_, key_) (attrs_)[key_].as<IntImm>()->value
#define GET_STRINGIMM_ATTR(attrs_, key_) (attrs_)[key_].as<StringImm>()->value
#define IS_STATIC true

#define GET_EXPR_ATTR(attrs_, key_, default_) \
  (IS_ATTR_EXIST(attrs_, key_) ? (Downcast<Expr>((attrs_)[key_])) : (default_))

#define GET_INTIMM_ATTR_DEFAULT(attrs_, key_, default_)                                                               \
  ((IS_ATTR_EXIST(attrs_, key_) && ((attrs_)[key_].as<IntImm>())) ? (static_cast<int>(GET_INTIMM_ATTR(attrs_, key_))) \
                                                                  : (default_))

#define GET_STRINGIMM_ATTR_DEFAULT(attrs_, key_, default_) \
  ((IS_ATTR_EXIST(attrs_, key_) && ((attrs_)[key_].as<StringImm>())) ? (GET_STRINGIMM_ATTR(attrs_, key_)) : (default_))

class Convolution : public IRVisitor {
 public:
  void Visit_(const AttrStmt *op) final {
    if (op->attr_key == "pragma_attrs") {
      attrs_ = Downcast<Map<std::string, NodeRef>>(op->node);
      if (attrs_.count(ATTR_CONV_BACKPROP_INPUT) > 0) {
        CHECK(attrs_[ATTR_CONV_BACKPROP_INPUT].as<IntImm>());
        isConvBackpropInput_ = GET_INTIMM_ATTR(attrs_, ATTR_CONV_BACKPROP_INPUT);
      }
      if (attrs_.count(ATTR_CONV_BACKPROP_FILTER) > 0) {
        CHECK(attrs_[ATTR_CONV_BACKPROP_FILTER].as<IntImm>());
        isConvBackpropFilter_ = GET_INTIMM_ATTR(attrs_, ATTR_CONV_BACKPROP_FILTER);
      }
    } else if (op->attr_key == "pragma_im2col") {
      isConv_ = true;
    }

    IRVisitor::Visit_(op);
  }

  Map<std::string, NodeRef> attrs_;
  bool isConv_{false};
  bool isConvBackpropInput_{false};
  bool isConvBackpropFilter_{false};
};

struct NCHW {
  Expr n{0};
  Expr c{0};
  Expr h{0};
  Expr w{0};

  NCHW() : n(0), c(0), h(0), w(0) {}

  NCHW(Expr n, Expr c, Expr h, Expr w) : n(n), c(floordiv(c + 15, 16) * 16), h(h), w(w) {}
};

struct Weight {
  Expr cout{0};
  Expr cin{0};
  Expr kh{0};
  Expr kw{0};
  Expr d_kh{0};
  Expr d_kw{0};

  Weight() : cout(0), cin(0), kh(0), kw(0), d_kh(0), d_kw(0) {}

  Weight(Expr cout, Expr cin, Expr kh, Expr kw)
      : cout(floordiv(cout + 15, 16) * 16), cin(floordiv(cin + 15, 16) * 16), kh(kh), kw(kw), d_kh(0), d_kw(0) {}

  Weight(Expr cout, Expr cin, Expr kh, Expr kw, Expr d_kh, Expr d_kw)
      : cout(floordiv(cout + 15, 16) * 16), cin(floordiv(cin + 15, 16) * 16), kh(kh), kw(kw), d_kh(d_kh), d_kw(d_kw) {}
};

struct Pad {
  Expr pad_top{0};
  Expr pad_bottom{0};
  Expr pad_left{0};
  Expr pad_right{0};

  Pad() : pad_top(0), pad_bottom(0), pad_left(0), pad_right(0) {}

  explicit Pad(Expr pad) : pad_top(pad), pad_bottom(pad), pad_left(pad), pad_right(pad) {}

  Pad(Expr pad_top, Expr pad_bottom, Expr pad_left, Expr pad_right)
      : pad_top(pad_top), pad_bottom(pad_bottom), pad_left(pad_left), pad_right(pad_right) {}
};

struct Stride {
  Expr stride_h{0};
  Expr stride_w{0};

  Stride() : stride_h(0), stride_w(0) {}

  explicit Stride(Expr stride) : stride_h(stride), stride_w(stride) {}

  Stride(Expr stride_h, Expr stride_w) : stride_h(stride_h), stride_w(stride_w) {}
};

struct Dilation {
  Expr dilation_h{0};
  Expr dilation_w{0};

  Dilation() : dilation_h(0), dilation_w(0) {}

  explicit Dilation(Expr dilation) : dilation_h(dilation), dilation_w(dilation) {}

  Dilation(Expr dilation_h, Expr dilation_w) : dilation_h(dilation_h), dilation_w(dilation_w) {}
};

struct Conv {
  NCHW input;
  Weight filter;
  NCHW output;
  Pad pad;
  Stride stride;
  Dilation dilation;
  Expr block_size{16};

  void infer_output_shape() {
    Expr out_n = this->input.n;
    Expr out_c = this->filter.cout;

    Expr in_h = this->input.h;
    Expr pad_top = this->pad.pad_top;
    Expr pad_bottom = this->pad.pad_bottom;
    Expr d_kh = this->filter.d_kh;
    Expr stride_h = this->stride.stride_h;
    Expr out_h = (in_h + pad_top + pad_bottom - d_kh) / stride_h + 1;

    Expr in_w = this->input.w;
    Expr pad_left = this->pad.pad_left;
    Expr pad_right = this->pad.pad_right;
    Expr d_kw = this->filter.d_kw;
    Expr stride_w = this->stride.stride_w;
    Expr out_w = (in_w + pad_left + pad_right - d_kw) / stride_w + 1;

    this->output = NCHW(out_n, out_c, out_h, out_w);
  }

  Conv() {}

  Conv(Expr n, Expr c, Expr h, Expr w, Expr cout, Expr kh, Expr kw, Expr pad, Expr stride, Expr dilation)
      : input(n, c, h, w),
        filter(cout, c, kh, kw, (kh - 1) * dilation + 1, (kw - 1) * dilation + 1),
        pad(pad),
        stride(stride),
        dilation(dilation) {
    infer_output_shape();
  }

  Conv(Expr n, Expr c, Expr h, Expr w, Expr cout, Expr kh, Expr kw, Expr pad_top, Expr pad_bottom, Expr pad_left,
       Expr pad_right, Expr stride_h, Expr stride_w, Expr dilation_h, Expr dilation_w)
      : input(n, c, h, w),
        filter(cout, c, kh, kw, (kh - 1) * dilation_h + 1, (kw - 1) * dilation_w + 1),
        pad(pad_top, pad_bottom, pad_left, pad_right),
        stride(stride_h, stride_w),
        dilation(dilation_h, dilation_w) {
    infer_output_shape();
  }
};

struct Tiling {
  Expr cut_b{0};
  Expr cut_ci{0};
  Expr cut_co{0};
  Expr cut_h{0};
  Expr cut_w{0};

  Expr cut_kh{0};
  Expr cut_kw{0};

  Expr cut_m{0};
  Expr cut_k{0};
  Expr cut_n{0};

  Tiling() : cut_b(0), cut_ci(0), cut_co(0), cut_h(0), cut_w(0), cut_kh(0), cut_kw(0), cut_m(0), cut_k(0), cut_n(0) {}

  Tiling(Expr cut_b, Expr cut_ci, Expr cut_co, Expr cut_h, Expr cut_w, Expr cut_kh, Expr cut_kw, Expr cut_m, Expr cut_k,
         Expr cut_n)
      : cut_b(cut_b),
        cut_ci(cut_ci),
        cut_co(cut_co),
        cut_h(cut_h),
        cut_w(cut_w),
        cut_kh(cut_kh),
        cut_kw(cut_kw),
        cut_m(cut_m),
        cut_k(cut_k),
        cut_n(cut_n) {}
};

struct IsolateInfo {
  Expr outer{0};
  Expr inner{0};

  IsolateInfo() : outer(0), inner(0) {}

  IsolateInfo(Expr outer, Expr inner) : outer(outer), inner(inner) {}
};

class ConvolutionModel {
 public:
  ConvolutionModel() {}

  ConvolutionModel(const Map<std::string, NodeRef> &attrs, bool is_dynamic);

  virtual ~ConvolutionModel() {}

  int infer_isolate(std::vector<IsolateInfo> *info, Expr len, Expr cut);

  int infer_isolate_overlap(std::vector<IsolateInfo> *info, std::vector<IsolateInfo> *win_info, Expr len, Expr cut,
                            Expr stride, Expr kernel_dilation, Expr head, Expr tail);

  /* L1 isolate Info */
  virtual int infer_L1_tile() = 0;

  virtual int get_ci_idx(int isolate_idx) const;
  virtual int get_co_idx(int isolate_idx) const;
  virtual int get_b_idx(int isolate_idx) const;
  virtual int get_h_idx(int isolate_idx) const;
  virtual int get_w_idx(int isolate_idx) const;
  virtual int get_kh_idx(int isolate_idx) const;
  virtual int get_kw_idx(int isolate_idx) const;

  IsolateInfo get_ci_isolate_info(int isolate_idx);
  IsolateInfo get_co_isolate_info(int isolate_idx);
  IsolateInfo get_b_isolate_info(int isolate_idx);
  IsolateInfo get_h_isolate_info(int isolate_idx);
  IsolateInfo get_h_win_isolate_info(int isolate_idx);
  IsolateInfo get_w_isolate_info(int isolate_idx);
  IsolateInfo get_w_win_isolate_info(int isolate_idx);
  IsolateInfo get_kh_isolate_info(int isolate_idx);
  IsolateInfo get_kw_isolate_info(int isolate_idx);

  /* L0 isolate Info */
  virtual int infer_L0_tile(int isolate_idx) = 0;

  int get_n_idx(int gemm_idx) const;
  int get_m_idx(int gemm_idx) const;
  int get_k_idx(int gemm_idx) const;

  IsolateInfo get_n_isolate_info(int gemm_idx);
  IsolateInfo get_m_isolate_info(int gemm_idx);
  IsolateInfo get_k_isolate_info(int gemm_idx);

  Expr calc_till_idx(std::vector<IsolateInfo> *info, int idx) const;

 protected:
  Map<std::string, NodeRef> attrs_;

 public:
  bool is_dynamic_{false};
  bool reduce_at_l1{false};
  int l1_reduce_base{1};
  int l0_reduce_base{1};

  Conv conv_;
  Tiling tile_;

  /* l1 info */
  std::vector<IsolateInfo> b_info;
  int b_base{0};
  std::vector<IsolateInfo> ci_info;
  int ci_base{0};
  std::vector<IsolateInfo> co_info;
  int co_base{0};
  std::vector<IsolateInfo> h_info;
  std::vector<IsolateInfo> h_win_info;
  int h_base{0};
  std::vector<IsolateInfo> w_info;
  std::vector<IsolateInfo> w_win_info;
  int w_base{0};
  std::vector<IsolateInfo> kh_info;
  int kh_base{0};
  std::vector<IsolateInfo> kw_info;
  int kw_base{0};

  /* l0 info */
  std::vector<IsolateInfo> m_info;
  int m_base{0};
  std::vector<IsolateInfo> k_info;
  int k_base{0};
  std::vector<IsolateInfo> n_info;
  int n_base{0};
};

class ConvolutionForwardModel : public ConvolutionModel {
 public:
  ConvolutionForwardModel() {}

  ConvolutionForwardModel(const Map<std::string, NodeRef> &attrs, bool is_dynamic);

  ~ConvolutionForwardModel() override = default;

  /* L1 isolate Info */
  int infer_L1_tile() override;

  int get_co_idx(int isolate_idx) const override;
  int get_h_idx(int isolate_idx) const override;
  int get_w_idx(int isolate_idx) const override;

  /* L0 isolate Info */
  int infer_L0_tile(int isolate_idx) override;
};

class ConvolutionBackpropInputModel : public ConvolutionModel {
 public:
  ConvolutionBackpropInputModel() {}

  ConvolutionBackpropInputModel(const Map<std::string, NodeRef> &attrs, bool is_dynamic);

  ~ConvolutionBackpropInputModel() override = default;

  /* L1 isolate Info */
  int infer_L1_tile() override;

  int get_co_idx(int isolate_idx) const override;
  int get_h_idx(int isolate_idx) const override;
  int get_w_idx(int isolate_idx) const override;

  /* L0 isolate Info */
  int infer_L0_tile(int isolate_idx) override;
};

class ConvolutionBackpropFilterModel : public ConvolutionModel {
 public:
  ConvolutionBackpropFilterModel() {}

  explicit ConvolutionBackpropFilterModel(const Map<std::string, NodeRef> &attrs, bool is_dynamic);

  ~ConvolutionBackpropFilterModel() override = default;

  /* L1 isolate Info */
  int infer_L1_tile() override;

  int get_ci_idx(int isolate_idx) const override;
  int get_co_idx(int isolate_idx) const override;
  int get_b_idx(int isolate_idx) const override;
  int get_h_idx(int isolate_idx) const override;
  int get_w_idx(int isolate_idx) const override;
  int get_kh_idx(int isolate_idx) const override;
  int get_kw_idx(int isolate_idx) const override;

  /* L0 isolate Info */
  int infer_L0_tile(int isolate_idx) override;
};
}  // namespace ir
}  // namespace akg

#endif  // PASS_CONVOLUTION_MODEL_H_
