//
// Created by bb on 2021/1/3.
//
#include <tvm/runtime/registry.h>
#include <tvm/tsl/te/operation.h>
#include <tvm/te/operation.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tsl/tir/expr.h>
#include <tvm/tir/stmt_functor.h>

#include <unordered_set>
namespace tvm {
namespace te {

Array<IterVar> TslBaseComputeOpNode::root_iter_vars() const {
  if (union_reduce_axis.size() == 0) return union_axis;
  Array<IterVar> ret = union_axis;
  for (IterVar iv : union_reduce_axis) {
    ret.push_back(iv);
  }
  return ret;
}

Array<PrimExpr> TslBaseComputeOpNode::output_unionshape(size_t idx) const {
  CHECK_LT(idx, num_outputs());
  // for now, all outputs of a BaseComputeOp have the same shape
  Array<PrimExpr> shape;
  for (const auto& ivar : this->union_axis) {
    const Range& r = ivar->dom;
    shape.push_back(r->extent);
  }
  return shape;
}

Array<TensorUnion> TslComputeOpNode::InputTensorUnions() const {
  Array<TensorUnion> ret;
  std::unordered_set<TensorUnion> visited;
  for (auto& e : body) {
    tir::PostOrderVisit(e, [&ret, &visited](const ObjectRef& n) {
      if (auto* pload = n.as<tir::TULoadNode>()) {
        TensorUnion t = Downcast<TensorUnion>(pload->producer);
        if (!visited.count(t)) {
          ret.push_back(t);
          visited.insert(t);
        }
      }
    });
  }
  return ret;
}

DataType TslComputeOpNode::output_dtype(size_t i) const {
  CHECK_LT(i, num_outputs());
  return body[i].dtype();
}
int TslComputeOpNode::num_outputs() const { return body.size(); }

TslComputeOp::TslComputeOp(std::string name, std::string tag, Map<String, ObjectRef> attrs,
                           Array<IterVar> union_axis, Array<PrimExpr> body) {
  if (!attrs.defined()) {
    attrs = Map<String, ObjectRef>();
  }
  auto n = make_object<TslComputeOpNode>();
  n->name = std::move(name);
  n->tag = std::move(tag);
  n->attrs = std::move(attrs);
  n->union_axis = std::move(union_axis);
  n->body = std::move(body);
  if (n->body[0]->IsInstance<tir::ReduceNode>()) {
    const tir::ReduceNode* reduce = n->body[0].as<tir::ReduceNode>();
    n->union_reduce_axis = reduce->axis;
  }
  // VerifyComputeOp(n.get()); //TODO: uncomment this after finishing verifier.
  data_ = std::move(n);
}

TensorUnion Tslcompute(Array<PrimExpr> union_shape, FCompute fcompute, std::string name,
                       std::string tag, Map<String, ObjectRef> attrs) {
  size_t ndim = union_shape.size();
  std::vector<IterVar> axis;
  std::vector<Var> args;
  for (size_t i = 0; i < ndim; ++i) {
    std::ostringstream os;
    os << "ax" << i;
    axis.emplace_back(
        IterVar(Range(0, union_shape[i]), Var(os.str(), union_shape[i].dtype()), kDataPar));
    args.push_back(axis.back()->var);
  }

  return TslComputeOp(name, tag, attrs, axis, {fcompute(args)}).output(0);
}

TVM_REGISTER_GLOBAL("te.TslComputeOp")
    .set_body_typed([](std::string name, std::string tag, Map<String, ObjectRef> attrs,
                       Array<IterVar> union_axis, Array<PrimExpr> body) {
      return TslComputeOp(name, tag, attrs, union_axis, body);
    });
/*class TslComputeVerifier final : protected tir::ExprVisitor {
 public:
  /// Special member functions
  //@{
  explicit ComputeVerifier(const ComputeOpNode* compute)
      : compute_(compute), reduce_(compute->body[0].as<tir::ReduceNode>()) {}
  virtual ~ComputeVerifier() = default;
  ComputeVerifier(const ComputeVerifier&) = delete;
  ComputeVerifier(ComputeVerifier&&) = delete;
  ComputeVerifier& operator=(const ComputeVerifier&) = delete;
  ComputeVerifier& operator=(ComputeVerifier&&) = delete;
  //@}

  /// Interface to perform compute verification
  void Run() {
    for (const PrimExpr e : compute_->body) {
      // Check for consistency of top level reductions
      const tir::ReduceNode* reduce = e.as<tir::ReduceNode>();
      CHECK((reduce && reduce_) || (!reduce && !reduce_)) << "All ComputeOp should be consistent "
                                                          << "with being Reduce operation or not.";

      if (reduce && reduce_) {
        CHECK(ReduceEqual(reduce, reduce_)) << "The Reduce inputs of ComputeOp should "
                                            << "have the same attribute except value_index";
      }

      level_ = 0;
      ExprVisitor::VisitExpr(e);
    }
  }

 protected:
  /// Visitor implementation
  //@{
  void VisitExpr(const PrimExpr& n) final {
    ++level_;
    ExprVisitor::VisitExpr(n);
    --level_;
  }

  void VisitExpr_(const tir::ReduceNode* op) final {
    // Check for non top level reductions
    CHECK(0 == level_) << "Reductions are only allowed at the top level of compute. "
                       << "Please create another tensor for further composition.";
  }
  //@}

 private:
  const ComputeOpNode* compute_{nullptr};   ///< ComputeOpNode to verify
  const tir::ReduceNode* reduce_{nullptr};  ///< Top level Reduce operation
  int level_{0};                            ///< Level of op being processed
};
}  // namespace*/ //TODO: complete this reduction verifier. it's not very significant in current setup though.
}  // namespace te
}  // namespace tvm
