// Copyright (c) 2024 CINN Authors. All Rights Reserved.
//
// 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.

#pragma once
#include <optional>
#include <unordered_map>
#include <vector>
#include "paddle/cinn/common/integer_set.h"
#include "paddle/cinn/common/ir_util.h"
#include "paddle/cinn/ir/ir.h"
#include "paddle/cinn/ir/ir_base.h"
#include "paddle/cinn/ir/ir_mutator.h"

namespace cinn {
namespace common {

class IterMapToExprNormalizer : public ir::IRMutator<> {
 public:
  explicit IterMapToExprNormalizer(const SymbolicExprAnalyzer& analyzer)
      : analyzer_(analyzer) {}

  void Convert(Expr* expr) { Visit(expr, expr); }

 private:
  void Visit(const Expr* expr, Expr* op) override;

  ir::IndexExpr ConvertIterSum(ir::IterSum* expr);

  ir::IndexExpr ConvertIterSplit(ir::IterSplit* expr);

 private:
  common::SymbolicExprAnalyzer analyzer_;
};

class IterMapRewriter : public ir::IRMutator<> {
 public:
  explicit IterMapRewriter(const std::vector<ir::Var>& input_iters,
                           const SymbolicExprAnalyzer& analyzer)
      : analyzer_(analyzer) {
    for (const auto& iter : input_iters) {
      // Remove unit extent `for` ?
      // if (IsOne(iter->upper_bound)) {
      //   var_map_[iter->name] = ir::IterSum::Make({}, iter->lower_bound);
      // }
      if (IsZero(iter->lower_bound)) {
        auto tmp =
            ir::IterMark::Make(ir::IndexExpr(iter.ptr()), iter->upper_bound);
        auto mark = tmp.As<ir::IterMark>();
        var_map_[iter->name] = ir::IterSplit::Make(tmp);
        input_marks_.push_back(*mark);
      } else {
        PADDLE_THROW(::common::errors::InvalidArgument(
            "iter should start from 0, but got %d", iter->lower_bound));
      }
    }
  }

  void Visit(Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }

  void Rewrite(Expr* expr) {
    IRMutator::Visit(expr, expr);
    *expr = ToIterSum(*expr);
  }

  void Visit(const ir::_Var_* op, Expr* expr) override;

  void Visit(const ir::Add* op, Expr* expr) override;

  void Visit(const ir::Sub* op, Expr* expr) override;

  void Visit(const ir::Mul* op, Expr* expr) override;

  void Visit(const ir::Div* op, Expr* expr) override;

  void Visit(const ir::Mod* op, Expr* expr) override;

 private:
  static Expr ToIterSum(const Expr& expr);

  static void AddToLhs(ir::IterSum* lhs, const ir::IterSplit& rhs, int sign);

  static void AddToLhs(ir::IterSum* lhs, const ir::IterSum& rhs, int sign);

  static void MulToLhs(ir::IterSum* lhs, const ir::IndexExpr& rhs);

  Expr PreprocessDividend(const Expr& dividend);

  Expr SplitDivConst(Expr lhs, ir::IndexExpr base, ir::IndexExpr rhs);

  Expr SplitModConst(Expr lhs, ir::IndexExpr base, ir::IndexExpr rhs);

  /*!
   * \brief This function will find the iter which has the expected scale.
   * For example:
   * expr:
   *  IterSum(IterSplit(IterMark(i), scale = 32),
   *          IterSplit(IterMark(j), scale = 8),
   *          base = 0)                    // i * 32 + j * 8
   * ret: `1` when expected_scale = 8
   * ret: `0` when expected_scale = 32
   *
   * \param expr the input IterSum to search.
   * \param skip_flag the flag to indicate whether a Iter should be skipped.
   * \param match_source Whether to only match the same source.
   * \param rbegin the last position in reverse searching. -1 means the last.
   * \param first_possible_unit_extent_pos the first possible position of the
   * unit extent.
   *  \return the index of the Iter with expected scale. return -1
   * if not found.
   */
  int32_t FindSplitWithExactScale(const ir::IterSum& expr,
                                  const std::vector<bool>& skip_flag,
                                  const ir::IndexExpr& expected_scale,
                                  const Expr& match_source,
                                  int32_t rbegin = -1,
                                  int32_t first_possible_unit_extent_pos = 0);

  /*!
   * \brief Find the first possible position where IterSplit->extent = 1.
   * \param expr the input IterSum to search.
   * \return the index of the first IterSplit with extent = 1,
   * return IterSum.args.size if not found.
   */
  int32_t FindFirstPossibleUnitExtentIndex(const ir::IterSum& expr);

  /*!
   * \brief This function will find the base Iter which has the smallest scale.
   *
   * For example:
   * expr:
   *  IterSum(IterSplit(IterMark(i), scale = 32),
   *          IterSplit(IterMark(j), scale = 8),
   *          base = 0)                    // i * 32 + j * 8
   *
   * ret: `1` when match_source = nullptr
   * ret: `0` when match_source = IterMark(i)
   *
   * \param expr the input IterSum to search.
   * \param skip_flag the flag to indicate whether a Iter should be skipped.
   * \param match_source Whether to only match the same source.
   * \param rbegin the last position in reverse searching. -1 means the last.
   * \return the index of the base Iter. return -1 if not found.
   */
  int32_t FindBaseSplit(const ir::IterSum& expr,
                        const std::vector<bool>& skip_flag,
                        const Expr& match_source,
                        int32_t rbegin = -1);

  /*!
   * \brief TryFuse will create new IterMark and returns an aggregated IterSum
   * that only has one IterSplit with the new IterMark.
   *
   * For example:
   * expr:
   *  IterSum(IterSplit(IterMark(i), scale = 32),
   *          IterSplit(IterMark(j), scale = 8),
   *          base = 0)                    // i * 32 + j * 8
   * ret:
   *  IterSum(IterSplit(IterMark(i * 4 + j), scale = 8),
   *          base = 0)                    // Treat `i * 4 + j` as a IterMark
   *
   * \param expr the input IterSum.
   * \return the IterSum after fused.
   */
  std::optional<Expr> TryFuse(const Expr& expr);

  /*!
   * \brief TryFuseSameSource will simplify the IterSum by fusing IterSplits
   * with same source. same source means the IterSplits have same IterMark.
   *
   * For example:
   * expr:
   *  IterSum(IterSplit(IterMark(f), lower = 4, ext = 8 scale = 4),
   *          IterSplit(IterMark(f), lower = 1, ext = 4 scale = 1),
   *          base = 0)                        // f /4 * 4 + f % 4
   * ret:
   *  IterSum(IterSplit(IterMark(f), scale = 1),
   *          base = 0)                        // f
   *
   * \param expr the input IterSum
   * \return the IterSum after fused.
   */
  std::optional<Expr> TryFuseSameSource(const Expr& expr);

  std::unordered_map<std::string, Expr> var_map_;
  std::vector<ir::IterMark> input_marks_;
  std::unordered_map<Expr, Expr> sum_fuse_map_;
  common::SymbolicExprAnalyzer analyzer_;
};

class SimplifyBlockBinding : public ir::IRMutator<> {
 public:
  explicit SimplifyBlockBinding(const std::vector<ir::Var>& loop_var,
                                const SymbolicExprAnalyzer& analyzer)
      : loop_var_(loop_var), analyzer_(analyzer) {}

  static void SimplifyBindings(ir::Expr expr,
                               const std::vector<ir::Expr>& loop_srefs,
                               const SymbolicExprAnalyzer& analyzer) {
    std::vector<ir::Var> loop_var;
    for (const ir::Expr& sref : loop_srefs) {
      const ir::For* loop = sref.As<ir::For>();
      loop_var.emplace_back(loop->loop_var);
    }
    SimplifyBlockBinding(loop_var, analyzer)(&expr);
  }
  void operator()(Expr* expr) { IRMutator::Visit(expr, expr); }

 private:
  void Visit(const ir::For* op, Expr* expr) override;

  void Visit(const ir::ScheduleBlockRealize* op, Expr* expr) override;

  std::vector<ir::Var> loop_var_;
  common::SymbolicExprAnalyzer analyzer_;
};

void IterMapSimplify(std::vector<Expr>& indices,  // NOLINT
                     const std::vector<cinn::ir::Var>& input_iters,
                     const SymbolicExprAnalyzer& analyzer);
}  // namespace common
}  // namespace cinn