/*
 * 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.
 */
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>

#include "../../../tirx/transform/ir_utils.h"
#include "../utils.h"

namespace tvm {
namespace s_tir {
using namespace tvm::tirx;

/*! \brief Information used to create new padding block */
struct PaddingSBlockInfo {
  /*! \brief In-bound block iter regions, wrt loop vars. */
  ffi::Array<Range> in_bound_region;
  /*! \brief In-bound value, wrt block iter vars. */
  PrimExpr in_bound_value;
  /*! \brief Condition of in-bound write, wrt loop vars. */
  PrimExpr in_bound_predicate;
  /*! \brief Padding value, should be a constant. */
  PrimExpr pad_value;
};

class PaddingPatternMatchError : public ScheduleError {
 public:
  PaddingPatternMatchError(IRModule mod, SBlock block, const std::string& error_msg)
      : mod_(std::move(mod)), block_(std::move(block)), error_msg_(error_msg) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: decompose_padding expect the block to match padding pattern\n  " +
           error_msg_;
  }

  ffi::String DetailRenderTemplate() const final {
    std::ostringstream os;
    os << "ScheduleError: decompose_padding expect the block {0} to match padding pattern\n  "
       << error_msg_;
    return os.str();
  }

  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {block_}; }

  IRModule mod_;
  SBlock block_;
  std::string error_msg_;
};

/*!
 * \brief Helper class to analyze and check the padding pattern of the block,
 * then return the padding information.
 */
class PaddingInfoAnalyzer {
 public:
  static PaddingSBlockInfo CheckAndGetPaddingInfo(IRModule mod, const SBlockRealizeNode* realize,
                                                  const ffi::Map<Var, Range>& dom_map,
                                                  arith::AnalyzerObj* analyzer) {
    PaddingInfoAnalyzer padding_analyzer(analyzer);
    if (!padding_analyzer.MatchPadding(realize, dom_map)) {
      throw PaddingPatternMatchError(mod, realize->block, padding_analyzer.error_msg_);
    }
    return padding_analyzer.info_;
  }

 private:
  explicit PaddingInfoAnalyzer(arith::AnalyzerObj* analyzer) : analyzer_(analyzer) {}

  /*! \brief Detect padding pattern and update result. */
  bool MatchPadding(const SBlockRealizeNode* realize, const ffi::Map<Var, Range>& dom_map) {
    // Step 1. Check match padding computation pattern.
    // A[...] = T.if_then_else(predicate, B[...], imm)
    SBlock block = realize->block;
    std::unordered_map<const VarNode*, PrimExpr> iter_values;
    for (size_t i = 0; i < realize->iter_values.size(); ++i) {
      Var block_var = block->iter_vars[i]->var;
      iter_values[block_var.get()] = realize->iter_values[i];
    }
    const BufferStoreNode* store = block->body.as<BufferStoreNode>();
    if (!store) {
      SetError("Block body expect a BufferStore to the write buffer");
      return false;
    }
    const CallNode* if_then_else = store->value.as<CallNode>();
    if (!if_then_else || !if_then_else->op.same_as(tirx::builtin::if_then_else())) {
      SetError("Value of BufferStore expect to be constrained by a padding predicate");
      return false;
    }
    PrimExpr pad_predicate = Substitute(if_then_else->args[0].as_or_throw<PrimExpr>(), iter_values);
    PrimExpr in_bound_value = if_then_else->args[1].as_or_throw<PrimExpr>();
    PrimExpr pad_value = if_then_else->args[2].as_or_throw<PrimExpr>();
    if (!is_const_number(pad_value)) {
      SetError("Pad value should be constant");
      return false;
    }

    // Step 2. Check in-bound computation to be effectiveless.
    if (SideEffect(if_then_else->args[1].as_or_throw<PrimExpr>()) > CallEffectKind::kReadState) {
      SetError("Inbound computation should not have side-effect");
      return false;
    }

    // Step 3. Analyze in-bound write region.
    PrimExpr in_bound_predicate = RewritePredicate(pad_predicate && realize->predicate);
    if (analyzer_->CanProveEqual(in_bound_predicate, 1)) {
      SetError("The in-bound predicate is trivial");
      return false;
    }
    ffi::Array<Range> in_bound_region = this->EstimateInBoundRegion(
        /*iter_values=*/realize->iter_values, /*dom_map=*/dom_map,
        /*in_bound_predicate=*/in_bound_predicate);
    if (in_bound_region.empty()) {
      return false;
    }

    // Step 4. Update result information.
    info_.in_bound_value = if_then_else->args[1].as_or_throw<PrimExpr>();
    info_.in_bound_region = in_bound_region;
    info_.in_bound_predicate = in_bound_predicate;
    info_.pad_value = pad_value;
    return true;
  }

  /*! \brief Rewrite predicate to left recursive conjunction, drop likely annotation. */
  PrimExpr RewritePredicate(const PrimExpr& predicate) {
    PrimExpr res = IntImm::Bool(true);
    std::function<void(PrimExpr)> update = [&res, &update](PrimExpr e) {
      arith::PVar<PrimExpr> a, b;
      if ((a && b).Match(e)) {
        update(a.Eval());
        update(b.Eval());
      } else {
        if (const CallNode* call = e.as<CallNode>()) {
          if (call->op.same_as(builtin::likely())) {
            e = call->args[0].as_or_throw<PrimExpr>();
          }
        }
        res = res && e;
      }
    };
    update(predicate);
    return analyzer_->Simplify(res);
  }

  /*! \brief Return iteration region of block vars where the padding predicate evals to true. */
  ffi::Array<Range> EstimateInBoundRegion(const ffi::Array<PrimExpr>& iter_values,
                                          const ffi::Map<Var, Range>& dom_map,
                                          const PrimExpr& in_bound_predicate) {
    ffi::Array<Range> region;

    arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(analyzer_);
    auto res = arith::DetectIterMap(iter_values, dom_map, in_bound_predicate,
                                    arith::IterMapLevel::Surjective, analyzer_ref);
    if (res->indices.empty()) {
      SetError("Block iters are not independent wrt padding condition");
      return {};
    }
    for (const arith::IterSumExpr& sum : res->indices) {
      if (sum->args.empty()) {
        region.push_back(Range::FromMinExtent(sum->base, IntImm(sum->base.ty(), /* value */ 1)));
      } else {
        TVM_FFI_ICHECK_EQ(sum->args.size(), 1U);
        if (!analyzer_->CanProveEqual(sum->args[0]->scale, 1)) {
          SetError("Strided iteration is not supported");
          return {};
        }
        region.push_back(Range::FromMinExtent(sum->base, sum->args[0]->extent));
      }
    }
    return region;
  }

  void SetError(const std::string& msg) { error_msg_ = msg; }

  /*! \brief padding info analyse result. */
  PaddingSBlockInfo info_;
  /*! \brief current error message. */
  std::string error_msg_;
  /*! \brief arithmetic analyzer. */
  arith::AnalyzerObj* analyzer_;
};

/*! \brief Create block to fill constant pad values into full region */
static std::pair<Stmt, SBlockRealize> CreateConstBlock(const SBlockRealizeNode* realize,
                                                       const PaddingSBlockInfo& info,
                                                       const ffi::Array<For>& loops,
                                                       const Stmt& highest_pos_inclusive,
                                                       arith::AnalyzerObj* analyzer) {
  const SBlock& block = realize->block;
  ffi::Array<IterVar> new_iter_vars;
  ffi::Map<Var, PrimExpr> repl_dict;

  // create new block itervars
  for (size_t i = 0; i < block->iter_vars.size(); ++i) {
    const IterVar& origin_iter = block->iter_vars[i];
    PrimVar new_var = origin_iter->var.CopyWithSuffix("");
    new_iter_vars.push_back(IterVar(origin_iter->dom, new_var, IterVarType::kDataPar));
    repl_dict.Set(origin_iter->var, new_var);
  }

  // rewrite expr helper
  auto rewrite_expr = [&repl_dict, analyzer](const PrimExpr& e) {
    return analyzer->Simplify(Substitute(e, repl_dict));
  };

  // create new write region
  TVM_FFI_ICHECK_EQ(block->writes.size(), 1U);
  BufferRegion write_region = BufferRegion(
      block->writes[0]->buffer, block->writes[0]->region.Map([rewrite_expr](const Range& r) {
        return Range::FromMinExtent(rewrite_expr(r->min), rewrite_expr(r->extent));
      }));

  // create block to fill const pad values
  BufferStore store = block->body.as_or_throw<BufferStore>();
  store.CopyOnWrite()->value = info.pad_value;
  store.CopyOnWrite()->indices = store->indices.Map(rewrite_expr);
  SBlock new_block(/*iter_vars=*/new_iter_vars, /*reads=*/{}, /*writes=*/{write_region},
                   /*name_hint=*/block->name_hint + "_pad_const", /*body=*/std::move(store));

  // create new loop vars
  ffi::Array<Var> new_loop_vars;
  for (const For& loop : loops) {
    Var new_var = loop->loop_var.CopyWithSuffix("");
    new_loop_vars.push_back(new_var);
    repl_dict.Set(loop->loop_var, new_var.as_or_throw<PrimExpr>());
    if (loop.same_as(highest_pos_inclusive)) {
      break;
    }
  }

  // create new block realize node
  ffi::Array<PrimExpr> new_iter_values;
  for (size_t i = 0; i < realize->iter_values.size(); ++i) {
    new_iter_values.push_back(rewrite_expr(realize->iter_values[i]));
  }
  SBlockRealize new_realize(/*iter_values=*/new_iter_values,
                            /*predicate=*/rewrite_expr(realize->predicate),
                            /*block=*/new_block);

  // create new loops
  Stmt nest_stmt_root = new_realize;
  for (size_t i = 0; i < new_loop_vars.size(); ++i) {
    For loop = loops[i];
    nest_stmt_root = For(new_loop_vars[i].as_or_throw<PrimVar>(), loop->min, loop->extent,
                         ForKind::kSerial, nest_stmt_root);
  }

  return {nest_stmt_root, new_realize};
}

/*! \brief Create block to fill in-bound region values. */
static std::pair<Stmt, SBlockRealize> CreateInBoundBlock(const SBlockRealizeNode* realize,
                                                         const PaddingSBlockInfo& info,

                                                         const ffi::Array<For>& loops,
                                                         const Stmt& highest_pos_inclusive,
                                                         arith::AnalyzerObj* analyzer) {
  const SBlock& block = realize->block;
  ffi::Array<IterVar> new_iter_vars;
  ffi::Map<Var, PrimExpr> repl_dict;

  // record loop ranges to be mutated
  ffi::Map<Var, Range> new_loop_ranges;
  for (const For& loop : loops) {
    new_loop_ranges.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
    if (loop.same_as(highest_pos_inclusive)) {
      break;
    }
  }

  // create new block iter vars and iter bindings
  ffi::Array<PrimExpr> new_iter_binding;
  for (size_t i = 0; i < info.in_bound_region.size(); ++i) {
    // add new block itervar
    const IterVar& origin_itervar = block->iter_vars[i];
    PrimVar new_var = origin_itervar->var.CopyWithSuffix("");
    Range new_range =
        Range::FromMinExtent(IntImm(new_var.ty(), 0), info.in_bound_region[i]->extent);
    new_iter_vars.push_back(IterVar(new_range, new_var, IterVarType::kDataPar));
    repl_dict.Set(origin_itervar->var, new_var + info.in_bound_region[i]->min);

    // update new loop range
    if (auto opt = realize->iter_values[i].as<Var>(); opt && new_loop_ranges.count(opt.value())) {
      // if the block binding is the loop var with single child, mutate loop range
      // instead of insert extra block predicate
      auto loop_var = opt.value();
      new_loop_ranges.Set(loop_var, new_range);
      new_iter_binding.push_back(realize->iter_values[i]);
      repl_dict.Set(loop_var, loop_var.as_or_throw<PrimExpr>() + info.in_bound_region[i]->min);
      analyzer->Bind(loop_var, new_range, /*allow_override=*/true);
    } else {
      new_iter_binding.push_back(
          analyzer->Simplify(realize->iter_values[i] - info.in_bound_region[i]->min));
    }
  }

  // rewrite helpers
  auto rewrite_expr = [&repl_dict, analyzer](const PrimExpr& e) {
    return analyzer->Simplify(Substitute(e, repl_dict));
  };
  auto rewrite_region = [rewrite_expr](const Region& region) {
    return region.Map([rewrite_expr](const Range& r) {
      return Range::FromMinExtent(rewrite_expr(r->min), rewrite_expr(r->extent));
    });
  };

  // create new read/write region for in-bound accesses
  ffi::Array<BufferRegion> reads, writes;
  for (const BufferRegion& read : block->reads) {
    reads.push_back(BufferRegion(read->buffer, rewrite_region(read->region)));
  }
  for (const BufferRegion& write : block->writes) {
    writes.push_back(BufferRegion(write->buffer, rewrite_region(write->region)));
  }

  // create new block realize node
  BufferStore store = block->body.as_or_throw<BufferStore>();
  store.CopyOnWrite()->value = rewrite_expr(info.in_bound_value);
  store.CopyOnWrite()->indices = store->indices.Map(rewrite_expr);
  SBlock new_block(/*iter_vars=*/new_iter_vars, /*reads=*/reads, /*writes=*/writes,
                   /*name_hint=*/block->name_hint, /*body=*/std::move(store));
  PrimExpr new_predicate = rewrite_expr(info.in_bound_predicate);
  SBlockRealize new_realize(/*iter_values=*/new_iter_binding, /*predicate=*/new_predicate,
                            /*block=*/new_block);

  // create new loops
  Stmt nest_stmt_root = new_realize;
  for (const For& loop : loops) {
    auto it = new_loop_ranges.find(loop->loop_var);
    PrimExpr min = it == new_loop_ranges.end() ? loop->min : (*it).second->min;
    PrimExpr extent = it == new_loop_ranges.end() ? loop->extent : (*it).second->extent;
    nest_stmt_root = For(loop->loop_var, min, extent, loop->kind, nest_stmt_root,
                         loop->thread_binding, loop->annotations, loop->step, loop->span);
    if (loop.same_as(highest_pos_inclusive)) {
      break;
    }
  }
  return {nest_stmt_root, new_realize};
}

/*!
 * \brief A helper class to create a new scope that contains decomposed padding blocks.
 */
class DecomposePaddingBlockReplacer : public StmtMutator {
 public:
  /*! \brief Replacement information */
  struct ReplaceDesc {
    /*! \brief loop above which to insert const pad value filling code. */
    For const_filling_pos;
    /*! \brief loop under which to insert in bound value filling code. */
    For in_bound_filling_pos;
    /*! \brief const pad value filling loop. */
    Stmt const_filling_loop;
    /*! \brief highest in bound value filling loop with single child. */
    Stmt in_bound_filling_loop;
    /*! \brief const pad value filling block. */
    SBlockRealize const_filling_block;
    /*! \brief in bound value filling block. */
    SBlockRealize in_bound_filling_block;
  };

  static SBlock Replace(SBlock scope_root, const ReplaceDesc& desc) {
    DecomposePaddingBlockReplacer replacer(desc);
    return replacer(std::move(scope_root)).as_or_throw<SBlock>();
  }

 private:
  explicit DecomposePaddingBlockReplacer(const ReplaceDesc& desc) : desc_(desc) {}

  Stmt VisitStmt_(const ForNode* op) final {
    Stmt new_loop;
    if (op == desc_.in_bound_filling_pos.get()) {
      // position to rewrite inbound filling code
      new_loop = desc_.in_bound_filling_loop;
    } else {
      new_loop = StmtMutator::VisitStmt_(op);
    }
    if (op == desc_.const_filling_pos.get()) {
      // position to insert pad value filling code
      return std::move(SeqStmt({desc_.const_filling_loop, new_loop}));
    }
    return new_loop;
  }

 private:
  const ReplaceDesc& desc_;
};

StmtSRef DecomposePaddingImpl(ScheduleState self, const StmtSRef& block_sref,
                              const StmtSRef& loop_sref, bool check_only) {
  /*!
   *  Check
   *    - the block is a compact block
   *    - the loop is an ancester of the block
   *    - the block match padding pattern
   *  Mutate
   *    - generate new block to fill padding values
   *    - trim original block to write non-padding part only
   */
  // Condition Checks and Information Collection
  const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
  const SBlockRealizeNode* realize = GetSBlockRealize(self, block_sref).get();
  ffi::Map<Var, Range> dom_map;
  arith::Analyzer analyzer;

  // Check 1. check the block is complete.
  StmtSRef scope_root_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/false);
  CheckCompleteBlock(self, block_sref, scope_root_sref);

  // Check 2. Check loop_sref is an ancestor of block_sref. Also collect
  //   - the highest loop position (inclusive) to insert const pad value filling code above.
  //   - the highest loop position (inclusive) to replace with in-bound value filling code.
  ffi::Array<StmtSRef> loop_srefs = GetLoops(block_sref);
  ffi::Array<For> loops;
  bool found_const_filling_pos = false;
  bool found_in_bound_filling_pos = false;
  For const_filling_pos = ffi::GetRef<For>(loop_sref->StmtAs<ForNode>());
  For in_bound_filling_pos{nullptr};
  for (auto it = loop_srefs.rbegin(); it != loop_srefs.rend(); ++it) {
    For cur_loop = ffi::GetRef<For>((*it)->StmtAs<ForNode>());
    Range range = Range::FromMinExtent(cur_loop->min, cur_loop->extent);
    dom_map.Set(cur_loop->loop_var, range);
    analyzer->Bind(cur_loop->loop_var, range);
    loops.push_back(cur_loop);

    if (cur_loop.same_as(const_filling_pos)) {
      TVM_FFI_ICHECK(!found_const_filling_pos);
      found_const_filling_pos = true;
      if (!found_in_bound_filling_pos) {
        found_in_bound_filling_pos = true;
        in_bound_filling_pos = cur_loop;
      }
    } else if (!found_in_bound_filling_pos) {
      if (!cur_loop->body->IsInstance<ForNode>() &&
          !cur_loop->body->IsInstance<SBlockRealizeNode>()) {
        found_in_bound_filling_pos = true;
      } else {
        in_bound_filling_pos = cur_loop;
      }
    }
  }
  TVM_FFI_ICHECK(in_bound_filling_pos.defined());
  if (!found_const_filling_pos) {
    throw LoopPositionError(self->mod, const_filling_pos, ffi::GetRef<SBlock>(block),
                            "decompose_padding");
  }

  // Check 3. match padding pattern and return padding operation info.
  PaddingSBlockInfo info =
      PaddingInfoAnalyzer::CheckAndGetPaddingInfo(self->mod, realize, dom_map, analyzer.get());

  // IR Manipulation
  // Step 1. Create const pad value filling part and in-bound value filling part.
  DecomposePaddingBlockReplacer::ReplaceDesc replace_desc;
  replace_desc.const_filling_pos = const_filling_pos;
  replace_desc.in_bound_filling_pos = in_bound_filling_pos;
  std::tie(replace_desc.const_filling_loop, replace_desc.const_filling_block) =
      CreateConstBlock(realize, info, loops, const_filling_pos, analyzer.get());
  std::tie(replace_desc.in_bound_filling_loop, replace_desc.in_bound_filling_block) =
      CreateInBoundBlock(realize, info, loops, in_bound_filling_pos, analyzer.get());

  // Step 2. Execute IR replacement.
  SBlock old_scope_root_block = ffi::GetRef<SBlock>(scope_root_sref->StmtAs<SBlockNode>());
  SBlock new_scope_root =
      DecomposePaddingBlockReplacer::Replace(old_scope_root_block, replace_desc);
  if (check_only) {
    return block_sref;
  }

  // Step 3. Update schedule states.
  self->Replace(scope_root_sref, new_scope_root,
                {{old_scope_root_block, new_scope_root},
                 {ffi::GetRef<SBlock>(block), replace_desc.in_bound_filling_block->block}});
  auto new_block_sref = self->stmt2ref.at(replace_desc.const_filling_block->block.get());

  // Set block info of created const pad value filling block
  SBlockInfo& block_info = self->block_info[new_block_sref];
  block_info.affine_binding = true;
  block_info.region_cover = true;
  block_info.stage_pipeline = true;

  // If the const pad value filling block is lifted out of the original subtree,
  // set the region_cover flag as false since region_cover is the property under the subtree.
  bool preserve_stage_pipeline = true;
  for (const StmtSRef& consumer_sref : GetConsumers(self, block_sref)) {
    StmtSRef lca = GetSRefLowestCommonAncestor({consumer_sref, block_sref});
    const StmtSRefNode* parent = new_block_sref->parent;
    bool is_under_lca = false;
    while (parent) {
      if (parent == lca.get()) {
        is_under_lca = true;
        break;
      }
      parent = parent->parent;
    }
    if (!is_under_lca) {
      preserve_stage_pipeline = false;
      self->block_info[consumer_sref].region_cover = false;
    }
  }
  if (!preserve_stage_pipeline) {
    self->block_info[scope_root_sref].stage_pipeline = false;
  }
  return new_block_sref;
}

StmtSRef DecomposePadding(ScheduleState self, const StmtSRef& block_sref,
                          const StmtSRef& loop_sref) {
  return DecomposePaddingImpl(self, block_sref, loop_sref, false);
}

bool CanDecomposePadding(ScheduleState self, const StmtSRef& block_sref,
                         const StmtSRef& loop_sref) {
  try {
    DecomposePaddingImpl(self, block_sref, loop_sref, true);
  } catch (const tvm::ffi::Error& e) {
    return false;
  }
  return true;
}

/******** FFI ********/

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def(
      "s_tir.schedule.CanDecomposePadding", [](Schedule self, SBlockRV block_rv, LoopRV loop_rv) {
        return CanDecomposePadding(self->state(), self->GetSRef(block_rv), self->GetSRef(loop_rv));
      });
}

/******** InstructionKind Registration ********/

struct DecomposPaddingTraits : public UnpackedInstTraits<DecomposPaddingTraits> {
  static constexpr const char* kName = "DecomposePadding";
  static constexpr bool kIsPure = false;

 private:
  static constexpr size_t kNumInputs = 2;
  static constexpr size_t kNumAttrs = 0;
  static constexpr size_t kNumDecisions = 0;

  static SBlockRV UnpackedApplyToSchedule(Schedule sch, SBlockRV block_rv, LoopRV loop_rv) {
    return sch->DecomposePadding(block_rv, loop_rv);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs, ffi::String block_rv,
                                      LoopRV loop_rv) {
    PythonAPICall py("decompose_padding");
    py.Input("block", block_rv);
    py.Input("loop", loop_rv);
    py.SingleOutput(outputs);
    return py.Str();
  }

  template <typename>
  friend struct ::tvm::s_tir::UnpackedInstTraits;
};

TVM_REGISTER_INST_KIND_TRAITS(DecomposPaddingTraits);

}  // namespace s_tir
}  // namespace tvm
