/*
 * 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/tirx/op.h>

#include "../utils.h"

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

/*!
 * \brief Check if buffer indices are all Vars and expr
 * \param buffer_access The BufferLoad or BufferStore
 * \return The indices if the indices are all Vars, otherwise std::nullopt
 */
ffi::Optional<ffi::Array<Var>> CheckTrivialBufferIndices(
    const ffi::Array<PrimExpr>& buffer_access) {
  ffi::Array<Var> indices;
  for (const PrimExpr& index : buffer_access) {
    if (index->IsInstance<IntImmNode>()) {
      continue;
    }
    auto var = index.as<PrimVar>();
    if (!var.has_value()) {
      return std::nullopt;
    }
    indices.push_back(var.value());
  }
  return indices;
}

ffi::Optional<ffi::Array<Var>> CheckTrivialBufferAccess(const BufferRegion& buffer_region) {
  ffi::Array<Var> indices;
  indices.reserve(buffer_region->region.size());
  for (const Range& range : buffer_region->region) {
    if (!tirx::is_one(range->extent)) {
      return std::nullopt;
    }
    if (range->min->IsInstance<IntImmNode>()) {
      continue;
    }
    if (auto var = range->min.as<PrimVar>()) {
      indices.push_back(var.value());
    } else {
      return std::nullopt;
    }
  }
  return indices;
}

/*! \brief The schedule error class when the padding size is invalid. */
class InvalidPaddingError : public ScheduleError {
 public:
  InvalidPaddingError(IRModule mod, SBlock block, ffi::Array<int64_t> padding)
      : mod_(std::move(mod)), block_(std::move(block)), padding_(std::move(padding)) {}
  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {block_}; }
  ffi::String FastErrorString() const final {
    return "ScheduleError: The padding size for the block is invalid.";
  }
  ffi::String DetailRenderTemplate() const final {
    std::ostringstream os;
    os << "The padding for the block {0} are invalid. It should be a list of "
       << block_->iter_vars.size() << " positive integers. Got " << padding_;
    return os.str();
  }

  static void Check(const ScheduleState& self, const SBlock& block, ffi::Array<int64_t> padding) {
    if (padding.size() != block->iter_vars.size()) {
      throw InvalidPaddingError(self->mod, block, padding);
    }
    for (int64_t pad : padding) {
      if (pad <= 0) {
        throw InvalidPaddingError(self->mod, block, padding);
      }
    }
  }

 private:
  IRModule mod_;
  SBlock block_;
  ffi::Array<int64_t> padding_;
};

/*! \brief The schedule error class when the block body is not an Einsum pattern. */
class NonEinsumError : public ScheduleError {
 public:
  explicit NonEinsumError(IRModule mod, SBlock block)
      : mod_(std::move(mod)), block_(std::move(block)) {}

  IRModule mod() const final { return mod_; }
  ffi::Array<ffi::ObjectRef> LocationsOfInterest() const final { return {block_}; }
  ffi::String FastErrorString() const final {
    return "ScheduleError: The block is not a computation of Einsum pattern.";
  }
  ffi::String DetailRenderTemplate() const final {
    return "The block {0} not a computation of Einsum pattern.";
  }

 private:
  IRModule mod_;
  SBlock block_;
};

/*! \brief Data structure that represents a Einsum computation. */
struct Einsum {
  // The output buffer
  ffi::Array<BufferVar> output_buffers;
  // The indices of the output buffer
  ffi::Map<BufferVar, ffi::Array<Var>> output_indices;
  // The input buffers
  ffi::Array<BufferVar> input_buffers;
  // The indices of the input buffers
  ffi::Map<BufferVar, ffi::Array<Var>> input_indices;
};

struct BufferPadding {
  BufferVar buffer;
  BufferVar padded_buffer;

  static BufferPadding FromBufferRegion(const BufferRegion& buffer_region,
                                        const ffi::Map<Var, PrimExpr>& iter_extents) {
    BufferPadding result;
    result.buffer = buffer_region->buffer;
    ffi::Array<PrimExpr> shape;
    shape.reserve(buffer_region->region.size());
    int ndim = buffer_region->region.size();
    for (int i = 0; i < ndim; ++i) {
      PrimExpr pos = buffer_region->region[i]->min;
      TVM_FFI_ICHECK(pos->IsInstance<IntImmNode>() || pos.as<PrimVar>());
      if (pos->IsInstance<IntImmNode>()) {
        shape.push_back(IntImm(pos.ty(), 1));
      } else if (ffi::Optional<PrimExpr> extent = iter_extents.Get(pos.as_or_throw<Var>())) {
        shape.push_back(extent.value());
      } else {
        shape.push_back(buffer_region->buffer->shape[i]);
      }
    }
    result.padded_buffer = decl_buffer(shape, result.buffer->dtype, result.buffer.name() + "_pad",
                                       result.buffer.scope());
    return result;
  }

  Stmt MakeCopyBlock(bool is_read, ffi::Array<SBlock>* blocks, arith::AnalyzerObj* analyzer) {
    ffi::Array<Var> loop_vars;
    ffi::Array<Range> loop_doms;
    ffi::Array<IterVar> iter_vars;
    ffi::Array<Range> instance_dom;
    ffi::Array<PrimExpr> indices;
    int ndim = buffer->shape.size();
    for (int i = 0; i < ndim; ++i) {
      PrimExpr dim{nullptr};
      if (is_read) {
        dim = padded_buffer->shape[i];
      } else {
        dim = buffer->shape[i];
      }
      Range dom = Range::FromMinExtent(IntImm(dim.ty(), 0), dim);
      loop_vars.push_back(Var("i" + std::to_string(i), dim.ty()));
      loop_doms.push_back(dom);
      IterVar iter_var(dom, PrimVar("v" + std::to_string(i), dim.ty()), kDataPar);
      instance_dom.push_back(Range::FromMinExtent(iter_var->var, IntImm(dim.ty(), 1)));
      iter_vars.push_back(iter_var);
      indices.push_back(iter_var->var);
    }
    Stmt body{nullptr};
    if (is_read) {
      PrimExpr predicate = IntImm::Bool(true);
      for (int i = 0; i < ndim; ++i) {
        if (!analyzer->CanProveEqual(buffer->shape[i], padded_buffer->shape[i])) {
          predicate = predicate && (indices[i] < buffer->shape[i]);
        }
      }
      PrimExpr rhs = BufferLoad(buffer, indices);
      body =
          BufferStore(padded_buffer, if_then_else(predicate, rhs, MakeConst(rhs.ty(), 0)), indices);
    } else {
      body = BufferStore(buffer, BufferLoad(padded_buffer, indices), indices);
    }
    BufferRegion read_region(buffer, instance_dom);
    BufferRegion write_region(padded_buffer, instance_dom);
    if (!is_read) {
      std::swap(read_region, write_region);
    }
    SBlock new_block(iter_vars, {read_region}, {write_region}, padded_buffer.name(),
                     std::move(body));
    blocks->push_back(new_block);
    ffi::Array<PrimExpr> prim_loop_vars;
    prim_loop_vars.reserve(loop_vars.size());
    for (const Var& var : loop_vars) prim_loop_vars.push_back(var.as_or_throw<PrimExpr>());
    body = SBlockRealize(prim_loop_vars, IntImm::Bool(true), new_block);
    for (int i = ndim - 1; i >= 0; --i) {
      body = For(loop_vars[i].as_or_throw<PrimVar>(), loop_doms[i]->min, loop_doms[i]->extent,
                 ForKind::kSerial, std::move(body));
    }
    return body;
  }
};

Einsum ExtractEinsum(const ScheduleState& self, const SBlock& block) {
  Einsum result;
  std::unordered_set<const VarNode*> buffer_used;
  int n_reads = block->reads.size();
  for (int i = 0; i < n_reads; ++i) {
    const BufferVar& buffer = block->reads[i]->buffer;
    if (buffer_used.count(buffer.get()) != 0) {
      throw NonEinsumError(self->mod, block);
    }
    buffer_used.insert(buffer.get());
    if (ffi::Optional<ffi::Array<Var>> opt_indices = CheckTrivialBufferAccess(block->reads[i])) {
      result.input_buffers.push_back(buffer);
      result.input_indices.Set(buffer, opt_indices.value());
    } else {
      throw NonEinsumError(self->mod, block);
    }
  }
  int n_writes = block->writes.size();
  for (int i = 0; i < n_writes; ++i) {
    const BufferVar& buffer = block->writes[i]->buffer;
    if (buffer_used.count(buffer.get()) != 0) {
      throw NonEinsumError(self->mod, block);
    }
    buffer_used.insert(buffer.get());
    if (ffi::Optional<ffi::Array<Var>> opt_indices = CheckTrivialBufferAccess(block->writes[i])) {
      result.output_buffers.push_back(buffer);
      result.output_indices.Set(buffer, opt_indices.value());
    } else {
      throw NonEinsumError(self->mod, block);
    }
  }
  return result;
}

class BufferNotAllocatedInScopeError : public ScheduleError {
 public:
  explicit BufferNotAllocatedInScopeError(IRModule mod, BufferVar buffer)
      : mod_(std::move(mod)), buffer_(std::move(buffer)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: The buffer is not allocated as an intermediate buffer in current "
           "PrimFunc.";
  }

  ffi::String DetailRenderTemplate() const final {
    std::ostringstream os;
    os << "The buffer " << buffer_.name()
       << " is not allocated as an intermediate buffer in current PrimFunc.";
    return os.str();
  }

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

 private:
  IRModule mod_;
  BufferVar buffer_;
};

/*! \brief The schedule error class when the producer block cannot be padded. */
class InvalidProducerError : public ScheduleError {
 public:
  explicit InvalidProducerError(IRModule mod, SBlock producer)
      : mod_(std::move(mod)), producer_(std::move(producer)) {}

  ffi::String FastErrorString() const final {
    return "ScheduleError: The producer block cannot be padded.";
  }

  ffi::String DetailRenderTemplate() const final {
    std::ostringstream os;
    os << "The producer block {0} cannot be padded. It should write to a single buffer and the "
          "body should be a BufferStore.";
    return os.str();
  }

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

 private:
  IRModule mod_;
  BufferVar buffer_;
  SBlock producer_;
};

class PadEinsumBufferReplacer : public StmtExprMutator {
 public:
  Stmt VisitStmt_(const SBlockNode* old_block_ptr) final {
    SBlock old_block = ffi::GetRef<SBlock>(old_block_ptr);
    SBlock block = StmtMutator::VisitStmt_(old_block_ptr).as_or_throw<SBlock>();
    ffi::Array<IterVar> iter_vars;
    iter_vars.reserve(block->iter_vars.size());
    for (const IterVar& iter_var : block->iter_vars) {
      if (ffi::Optional<PrimExpr> new_dom = iter2padded_extents.Get(iter_var->var)) {
        ffi::ObjectPtr<IterVarNode> new_iter_var = ffi::make_object<IterVarNode>(*iter_var.get());
        new_iter_var->dom = Range::FromMinExtent(iter_var->dom->min, new_dom.value());
        iter_vars.push_back(IterVar(new_iter_var));
      } else {
        iter_vars.push_back(iter_var);
      }
    }
    ffi::Array<BufferRegion> reads;
    reads.reserve(block->reads.size());
    for (const BufferRegion& read : block->reads) {
      if (ffi::Optional<BufferVar> buffer = buffer_map_.Get(read->buffer)) {
        reads.push_back(BufferRegion(buffer.value(), read->region));
      } else {
        reads.push_back(read);
      }
    }
    ffi::Array<BufferRegion> writes;
    writes.reserve(block->writes.size());
    for (const BufferRegion& write : block->writes) {
      if (ffi::Optional<BufferVar> buffer = buffer_map_.Get(write->buffer)) {
        writes.push_back(BufferRegion(buffer.value(), write->region));
      } else {
        writes.push_back(write);
      }
    }
    SBlock new_block =
        SBlock(iter_vars, reads, writes, block->name_hint, block->body, block->init,
               /*alloc_buffers=*/{}, /*match_buffers=*/{}, /*annotations=*/block->annotations);
    block_sref_reuse_.Set(old_block, new_block);
    return new_block;
  }

  Stmt VisitStmt_(const ForNode* old_for_ptr) final {
    For old_for = ffi::GetRef<For>(old_for_ptr);
    For new_for = StmtMutator::VisitStmt_(old_for_ptr).as_or_throw<For>();
    if (ffi::Optional<PrimExpr> new_extent = loop_var2padded_extent.Get(new_for->loop_var)) {
      ffi::ObjectPtr<ForNode> new_for_ptr = ffi::make_object<ForNode>(*new_for.get());
      new_for_ptr->extent = new_extent.value();
      new_for = For(new_for_ptr);
    }
    return new_for;
  }

  Stmt VisitStmt_(const BufferStoreNode* old_store_ptr) final {
    BufferStore store = StmtMutator::VisitStmt_(old_store_ptr).as_or_throw<BufferStore>();
    if (ffi::Optional<BufferVar> buffer = buffer_map_.Get(store->buffer)) {
      return BufferStore(buffer.value(), store->value, store->indices);
    } else {
      return store;
    }
  }

  Expr VisitExpr_(const BufferLoadNode* old_load_ptr) final {
    BufferLoad load = ExprMutator::VisitExpr_(old_load_ptr).as_or_throw<BufferLoad>();
    if (ffi::Optional<BufferVar> buffer = buffer_map_.Get(load->buffer)) {
      return BufferLoad(buffer.value(), load->indices);
    } else {
      return load;
    }
  }

  ffi::Map<Var, PrimExpr> iter2padded_extents;
  ffi::Map<Var, PrimExpr> loop_var2padded_extent;
  ffi::Map<BufferVar, BufferVar> buffer_map_;
  ffi::Map<SBlock, SBlock> block_sref_reuse_;
};

void PadEinsum(ScheduleState self, const StmtSRef& block_sref, const ffi::Array<int64_t>& padding) {
  arith::Analyzer analyzer;
  // Step 1: Input checking and error handling
  const SBlockNode* block = TVM_SREF_TO_SBLOCK(block_sref);
  SBlockRealize realize = GetSBlockRealize(self, block_sref);
  StmtSRef scope_sref = GetScopeRoot(self, block_sref, /*require_stage_pipeline=*/true);
  const SBlockNode* scope_block = TVM_SREF_TO_SBLOCK(scope_sref);
  InvalidPaddingError::Check(self, ffi::GetRef<SBlock>(block), padding);
  // Step 2. Extract the Einsum pattern
  ExtractEinsum(self, ffi::GetRef<SBlock>(block));
  // Step 3. Figure out the padding needed
  PadEinsumBufferReplacer replacer;
  for (int i = 0, n = padding.size(); i < n; ++i) {
    const IterVar& iter = block->iter_vars[i];
    PrimExpr dom = iter->dom->extent;
    PrimExpr pad_imm = IntImm(dom.ty(), padding[i]);
    PrimExpr new_dom = analyzer->Simplify(ceildiv(dom, pad_imm) * pad_imm);
    if (!analyzer->CanProveEqual(new_dom, dom)) {
      replacer.iter2padded_extents.Set(iter->var, new_dom);
      if (auto loop_var = realize->iter_values[i].as<PrimVar>()) {
        replacer.iter2padded_extents.Set(loop_var.value(), new_dom);
        replacer.loop_var2padded_extent.Set(loop_var.value(), new_dom);
      }
    }
  }
  auto f_needs_padding = [&replacer](const ffi::Array<Range>& region) {
    for (const Range& range : region) {
      if (auto var = range->min.as<PrimVar>()) {
        if (replacer.iter2padded_extents.count(var.value())) {
          return true;
        }
      }
    }
    return false;
  };
  // Step 3. Convert the subtree under the scope root
  ffi::Array<Stmt> scope_body;
  if (const auto* seq_stmt = scope_block->body.as<SeqStmtNode>()) {
    scope_body = seq_stmt->seq;
  } else {
    scope_body.push_back(scope_block->body);
  }
  // Step 4. Find out the block of our interest
  int pos = -1;
  for (int i = 0; i < static_cast<int>(scope_body.size()); ++i) {
    bool found = false;
    PostOrderVisit(scope_body[i], [&found, &block](const ffi::ObjectRef& node) {
      if (node.get() == block) {
        found = true;
      }
    });
    if (found) {
      pos = i;
      break;
    }
  }
  TVM_FFI_ICHECK_NE(pos, -1);
  // Step 5. For each buffer, if it needs padding, create a new buffer and a new block
  ffi::Array<Stmt> read_blocks;
  ffi::Array<Stmt> write_blocks;
  ffi::Array<SBlock> new_copy_blocks;
  ffi::Array<BufferVar> alloc_buffers;
  for (const BufferRegion& buffer_region : block->reads) {
    if (f_needs_padding(buffer_region->region)) {
      BufferPadding bp =
          BufferPadding::FromBufferRegion(buffer_region, replacer.iter2padded_extents);
      replacer.buffer_map_.Set(bp.buffer, bp.padded_buffer);
      read_blocks.push_back(bp.MakeCopyBlock(true, &new_copy_blocks, analyzer.get()));
      alloc_buffers.push_back(bp.padded_buffer);
    }
  }
  for (const BufferRegion& buffer_region : block->writes) {
    if (f_needs_padding(buffer_region->region)) {
      BufferPadding bp =
          BufferPadding::FromBufferRegion(buffer_region, replacer.iter2padded_extents);
      replacer.buffer_map_.Set(bp.buffer, bp.padded_buffer);
      write_blocks.push_back(bp.MakeCopyBlock(false, &new_copy_blocks, analyzer.get()));
      alloc_buffers.push_back(bp.padded_buffer);
    }
  }
  // Step 6. Create new scope body
  ffi::Array<Stmt> new_scope_body;
  for (int i = 0; i < static_cast<int>(scope_body.size()); ++i) {
    if (i != pos) {
      new_scope_body.push_back(scope_body[i]);
      continue;
    }
    new_scope_body.insert(new_scope_body.end(), read_blocks.begin(), read_blocks.end());
    new_scope_body.push_back(replacer(scope_body[i]));
    new_scope_body.insert(new_scope_body.end(), write_blocks.begin(), write_blocks.end());
  }
  // Step 7. Create new scope
  SBlock new_scope_block{nullptr};
  {
    ffi::ObjectPtr<SBlockNode> n = ffi::make_object<SBlockNode>(*scope_block);
    n->body = SeqStmt::Flatten(new_scope_body);
    n->alloc_buffers.insert(n->alloc_buffers.end(), alloc_buffers.begin(), alloc_buffers.end());
    new_scope_block = SBlock(n);
  }
  replacer.block_sref_reuse_.Set(ffi::GetRef<SBlock>(scope_block), new_scope_block);
  // Step 8. Do replacement and update flags
  self->Replace(scope_sref, new_scope_block, replacer.block_sref_reuse_);
  for (const SBlock& block : new_copy_blocks) {
    StmtSRef block_sref = self->stmt2ref.at(block.get());
    SBlockInfo& block_info = self->block_info[block_sref];
    block_info.affine_binding = true;
    block_info.region_cover = true;
    block_info.stage_pipeline = true;
  }
}

/******** Instruction Registration ********/

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

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

  static void UnpackedApplyToSchedule(Schedule sch, SBlockRV block, ffi::Array<int64_t> padding) {
    sch->PadEinsum(block, padding);
  }

  static ffi::String UnpackedAsPython(ffi::Array<ffi::String> outputs, ffi::String block,
                                      ffi::Array<int64_t> padding) {
    PythonAPICall py("pad_einsum");
    py.Input("block", block);
    py.Input("padding", padding);
    return py.Str();
  }

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

TVM_REGISTER_INST_KIND_TRAITS(PadEinsumTraits);

}  // namespace s_tir
}  // namespace tvm
