/*
 * 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.
 */

/*!
 * \file tvm/arith/analyzer.cc
 */
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>

#include "const_fold.h"
#include "product_normal_form.h"

namespace tvm {
namespace arith {

AnalyzerObj::AnalyzerObj()
    : const_int_bound(this),
      modular_set(this),
      rewrite_simplify(this),
      canonical_simplify(this),
      int_set(this),
      z3_prover(this) {}

void AnalyzerObj::Bind(const Var& var, const PrimExpr& expr, bool allow_override) {
  PrimExpr new_expr = expr;
  new_expr = this->canonical_simplify(new_expr);
  new_expr = this->rewrite_simplify(new_expr);

  this->const_int_bound.Update(var, this->const_int_bound(new_expr), allow_override);
  this->modular_set.Update(var, this->modular_set(new_expr), allow_override);
  this->rewrite_simplify.Update(var, new_expr, allow_override);
  this->canonical_simplify.Update(var, new_expr, allow_override);
  this->int_set.Update(var, this->int_set(new_expr), allow_override);
  this->transitive_comparisons.Bind(var, expr, allow_override);
  this->z3_prover.Bind(var, expr, allow_override);
}

void AnalyzerObj::Bind(const Var& var, const Range& range, bool allow_override) {
  TVM_FFI_ICHECK(range.defined());
  if (tirx::is_one(range->extent)) {
    this->Bind(var, range->min, allow_override);
  } else {
    this->const_int_bound.Bind(var, range, allow_override);
    this->int_set.Bind(var, range, allow_override);
    this->transitive_comparisons.Bind(var, range, allow_override);
    this->z3_prover.Bind(var, range, allow_override);
  }
  // skip modular_set
  // skip rewrite simplify
}

void AnalyzerObj::MarkGlobalNonNegValue(const PrimExpr& value) {
  // decompose value as symbol * scale + offset
  int64_t offset = 0;
  PrimType value_ty = value.ty();
  PrimExpr symbol_scale = tirx::MakeConst(value_ty, 0);

  auto fcollect_sum = [&](PrimExpr val, int sign) {
    if (const auto* intimm = val.as<IntImmNode>()) {
      offset += intimm->value * sign;
    } else {
      if (sign > 0) {
        symbol_scale = symbol_scale + val;
      } else {
        symbol_scale = symbol_scale - val;
      }
    }
  };
  UnpackSum(value, fcollect_sum);

  // split out the symbol and non-symbolic part
  int64_t cscale = 1;
  PrimExpr symbol = tirx::MakeConst(value_ty, 1);
  auto fcollect_prod = [&](PrimExpr val) {
    if (const auto* intimm = val.as<IntImmNode>()) {
      cscale *= intimm->value;
    } else {
      symbol = symbol * val;
    }
  };
  UnpackReduction<tirx::MulNode>(symbol_scale, fcollect_prod);
  if (cscale <= 0) return;
  // override the constant int bound by marking it as non-negative
  // NOTE: there might be future opportunities of more bound hint
  // this is a simple step and covers all the current needs
  //
  // We may consider enhance the sub analyzer to directly take
  // MarkPositiveVar so their bounds do not overlap
  if (auto prim_var = symbol.as<tirx::PrimVar>()) {
    Var var = *prim_var;
    // skip non-index type, keep it to be compatible
    // with any_dim that do not represent any value
    if (!IsIndexTypedExpr(var.as_or_throw<PrimExpr>())) return;
    bool allow_override = true;
    // mark the constant bound is sufficient
    // we cannot mark interval set as that will cause relaxation of the var
    // during bound proof which is not our intention
    this->const_int_bound.Update(var, ConstIntBound(-offset, ConstIntBound::kPosInf),
                                 allow_override);
  }
}

void AnalyzerObj::Bind(const ffi::Map<Var, Range>& variables, bool allow_override) {
  for (const auto& iter : variables) {
    this->Bind(iter.first, iter.second, allow_override);
  }
}

void ConstraintContext::EnterWithScope() {
  TVM_FFI_ICHECK(recovery_functions_.size() == 0);
  // entering the scope.
  recovery_functions_.push_back(analyzer_->const_int_bound.EnterConstraint(constraint_));
  recovery_functions_.push_back(analyzer_->modular_set.EnterConstraint(constraint_));
  recovery_functions_.push_back(analyzer_->rewrite_simplify.EnterConstraint(constraint_));
  recovery_functions_.push_back(analyzer_->int_set.EnterConstraint(constraint_));
  recovery_functions_.push_back(analyzer_->transitive_comparisons.EnterConstraint(constraint_));
  recovery_functions_.push_back(analyzer_->z3_prover.EnterConstraint(constraint_));
}

void ConstraintContext::ExitWithScope() {
  while (recovery_functions_.size()) {
    auto& func = recovery_functions_.back();
    if (func) {
      func();
    }
    recovery_functions_.pop_back();
  }
}

bool AnalyzerObj::CanProveGreaterEqual(const PrimExpr& expr, int64_t lower_bound) {
  if (const auto* ptr = expr.as<tirx::IntImmNode>()) {
    return ptr->value >= lower_bound;
  }
  auto bd = this->const_int_bound(this->rewrite_simplify(expr));
  if (bd->min_value >= lower_bound) return true;
  return false;
}

bool AnalyzerObj::CanProveLess(const PrimExpr& expr, int64_t upper_bound) {
  if (const auto* ptr = expr.as<tirx::IntImmNode>()) {
    return ptr->value < upper_bound;
  }
  auto bd = this->const_int_bound(this->rewrite_simplify(expr));
  if (bd->max_value < upper_bound) return true;
  return false;
}

bool AnalyzerObj::CanProveEqual(const PrimExpr& lhs, const PrimExpr& rhs) {
  const auto* clhs = lhs.as<IntImmNode>();
  const auto* crhs = rhs.as<IntImmNode>();
  if (clhs && crhs) return clhs->value == crhs->value;
  if (is_pos_inf(lhs) || is_neg_inf(lhs) || is_pos_inf(rhs) || is_neg_inf(rhs)) {
    return lhs.same_as(rhs);
  }
  return CanProve(lhs - rhs == 0);
}

bool AnalyzerObj::CanProveLessEqualThanSymbolicShapeValue(const PrimExpr& lhs,
                                                          const PrimExpr& shape) {
  if (this->CanProve(lhs <= shape, ProofStrength::kSymbolicBound)) return true;
  // no need to do further attempt if shape is already a constant.
  if (tirx::is_const_int(shape)) return false;
  // collect constant scale and ignore symbolic part
  // so 32 * n => cscale = 32
  int64_t cscale = 1;
  auto fcollect = [&](const PrimExpr& expr) {
    if (auto* ptr = expr.as<IntImmNode>()) {
      cscale *= ptr->value;
    }
  };
  UnpackReduction<tirx::MulNode>(shape, fcollect);
  PrimExpr const_shape_bound = IntImm(shape.ty(), std::abs(cscale));
  if (this->CanProve(lhs <= const_shape_bound, ProofStrength::kSymbolicBound)) return true;
  return false;
}

bool AnalyzerObj::CanProve(const PrimExpr& expr, ProofStrength strength) {
  // Avoid potentially expensive simplification unless required.
  if (const auto* ptr = expr.as<IntImmNode>()) {
    return ptr->value != 0;
  }
  PrimExpr simplified = Simplify(expr);
  const int64_t* as_int = tirx::as_const_int(simplified);
  if (as_int && *as_int) return true;
  if (strength >= ProofStrength::kSymbolicBound) {
    // NOTE: we intentionally only pattern match common bound predicate i < bound
    // and put this implementation at the top-level.
    // This is to avoid repeatitive calling of this function
    // that causes speed issues.
    // This strategy can only be called from top-level and not from sub-analyzers.
    ffi::Optional<PrimExpr> pos_diff;
    int lower_bound = 0;
    if (const auto* ptr_lt = expr.as<tirx::LTNode>()) {
      pos_diff = ptr_lt->b - ptr_lt->a;
      lower_bound = 1;
    }
    if (const auto* ptr_le = expr.as<tirx::LENode>()) {
      pos_diff = ptr_le->b - ptr_le->a;
      lower_bound = 0;
    }
    if (const auto* ptr_gt = expr.as<tirx::GTNode>()) {
      pos_diff = ptr_gt->a - ptr_gt->b;
      lower_bound = 1;
    }
    if (const auto* ptr_ge = expr.as<tirx::GENode>()) {
      pos_diff = ptr_ge->a - ptr_ge->b;
      lower_bound = 0;
    }
    if (pos_diff) {
      IntSet iset = this->int_set(this->Simplify(pos_diff.value()));
      if (iset.HasLowerBound()) {
        ConstIntBound relaxed_lower_bound = this->const_int_bound(this->Simplify(iset.min()));
        if (relaxed_lower_bound->min_value >= lower_bound) return true;
      }
    }
  }

  // Z3 is an expensive best-effort fallback. Gate it behind the higher
  // kSymbolicBound strength so the common kDefault path (including deeply
  // recursive internal CanProve calls) never pays the prover cost.
  if (strength >= ProofStrength::kSymbolicBound && z3_prover.CanProve(simplified)) {
    return true;
  }
  return false;
}

PrimExpr AnalyzerObj::Simplify(const PrimExpr& expr, int steps) {
  PrimExpr res = expr;

  // Always starts with a canonical simplification, as some structural property
  // of an expression might be destroyed by rewrite simplification.
  res = this->canonical_simplify(res);

  for (int i = 0; i < steps; ++i) {
    if (tirx::is_const_int(res)) {
      return res;
    }
    if (i % 2 == 0) {
      res = this->rewrite_simplify(res);
    } else {
      res = this->canonical_simplify(res);
    }
  }

  return res;
}

Analyzer AnalyzerObj::Clone() const {
  Analyzer cloned;
  cloned->const_int_bound.CopyFrom(this->const_int_bound);
  cloned->modular_set.CopyFrom(this->modular_set);
  cloned->rewrite_simplify.CopyFrom(this->rewrite_simplify);
  cloned->canonical_simplify.CopyFrom(this->canonical_simplify);
  cloned->int_set.CopyFrom(this->int_set);
  cloned->transitive_comparisons.CopyFrom(this->transitive_comparisons);
  cloned->z3_prover.CopyFrom(this->z3_prover);
  return cloned;
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::ObjectDef<AnalyzerObj>();
  refl::GlobalDef()
      .def("arith.Analyzer", []() { return Analyzer(); })
      .def("arith.AnalyzerClone", [](Analyzer analyzer) { return analyzer->Clone(); })
      .def("arith.EnterZ3ContextScope", []() { EnterZ3ContextScope(); })
      .def("arith.ExitZ3ContextScope", []() { ExitZ3ContextScope(); })
      .def("arith.AnalyzerConstIntBound",
           [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->const_int_bound(expr); })
      .def("arith.AnalyzerConstIntBoundUpdate",
           [](Analyzer analyzer, const Var& var, const ConstIntBound& info, bool allow_override) {
             analyzer->const_int_bound.Update(var, info, allow_override);
           })
      .def("arith.AnalyzerConstIntBoundIsBound",
           [](Analyzer analyzer, const Var& var) { return analyzer->const_int_bound.IsBound(var); })
      .def("arith.AnalyzerModularSetUpdate",
           [](Analyzer analyzer, const Var& var, const ModularSet& info, bool allow_override) {
             analyzer->modular_set.Update(var, info, allow_override);
           })
      .def("arith.AnalyzerIntSetUpdate",
           [](Analyzer analyzer, const Var& var, const IntSet& info, bool allow_override) {
             analyzer->int_set.Update(var, info, allow_override);
           })
      .def("arith.AnalyzerModularSet",
           [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->modular_set(expr); })
      .def("arith.AnalyzerSimplify", [](Analyzer analyzer, const PrimExpr& expr,
                                        int steps) { return analyzer->Simplify(expr, steps); })
      .def("arith.AnalyzerRewriteSimplify",
           [](Analyzer analyzer, const PrimExpr& expr) { return analyzer->rewrite_simplify(expr); })
      .def("arith.AnalyzerGetRewriteSimplifyStats",
           [](Analyzer analyzer) { return analyzer->rewrite_simplify.GetStatsCounters(); })
      .def("arith.AnalyzerResetRewriteSimplifyStats",
           [](Analyzer analyzer) { analyzer->rewrite_simplify.ResetStatsCounters(); })
      .def("arith.AnalyzerCanonicalSimplify",
           [](Analyzer analyzer, const PrimExpr& expr) {
             return analyzer->canonical_simplify(expr);
           })
      .def("arith.AnalyzerIntSet",
           [](Analyzer analyzer, const PrimExpr& expr,
              ffi::Optional<ffi::Map<Var, IntSet>> opt_dom_map) {
             if (opt_dom_map.has_value()) {
               return analyzer->int_set(expr, opt_dom_map.value());
             }
             return analyzer->int_set(expr);
           })
      .def_packed("arith.AnalyzerBind",
                  [](ffi::PackedArgs args, ffi::Any* ret) {
                    TVM_FFI_ICHECK(args.size() == 3 || args.size() == 4)
                        << "AnalyzerBind expects 3 or 4 arguments, but got " << args.size();
                    Analyzer analyzer = args[0].cast<Analyzer>();
                    bool allow_override = args.size() >= 4 && args[3].cast<bool>();
                    if (auto opt_range = args[2].try_cast<Range>()) {
                      analyzer->Bind(args[1].cast<Var>(), opt_range.value(), allow_override);
                    } else {
                      analyzer->Bind(args[1].cast<Var>(), args[2].cast<PrimExpr>(), allow_override);
                    }
                  })
      .def("arith.AnalyzerCanProve",
           [](Analyzer analyzer, const PrimExpr& expr, int strength) {
             return analyzer->CanProve(expr, static_cast<ProofStrength>(strength));
           })
      .def("arith.EnterAllowUintAsIndex", []() { ++arith::uint_as_index::g_depth; })
      .def("arith.ExitAllowUintAsIndex",
           []() {
             TVM_FFI_ICHECK(arith::uint_as_index::g_depth > 0)
                 << "ExitAllowUintAsIndex without a matching Enter";
             --arith::uint_as_index::g_depth;
           })
      .def("arith.GetAllowUintAsIndex", []() { return arith::uint_as_index::Enabled(); })
      .def("arith.AnalyzerSetMaximumRewriteSteps",
           [](Analyzer analyzer, int64_t maximum) {
             analyzer->rewrite_simplify.SetMaximumRewriteSteps(maximum);
           })
      .def("arith.AnalyzerEnterConstraintContext",
           [](Analyzer analyzer, const PrimExpr& constraint) {
             // can't use make_shared due to noexcept(false) decl in destructor,
             // see https://stackoverflow.com/a/43907314
             auto ctx = std::shared_ptr<With<ConstraintContext>>(
                 new With<ConstraintContext>(analyzer, constraint));
             auto fexit = [ctx](ffi::PackedArgs, ffi::Any*) mutable { ctx.reset(); };
             return ffi::Function::FromPacked(fexit);
           })
      .def_method("arith.AnalyzerCanProveEqual", &AnalyzerObj::CanProveEqual)
      .def("arith.AnalyzerTryCompare",
           [](Analyzer analyzer, const PrimExpr& lhs, const PrimExpr& rhs,
              bool propagate_inequalities) {
             return static_cast<int64_t>(
                 analyzer->transitive_comparisons.TryCompare(lhs, rhs, propagate_inequalities));
           })
      .def("arith.AnalyzerIsZ3Enabled",
           [](Analyzer analyzer) { return analyzer->z3_prover.IsEnabled(); })
      .def("arith.AnalyzerGetSMTLIB2",
           [](Analyzer analyzer, ffi::Optional<PrimExpr> expr) {
             return analyzer->z3_prover.GetSMTLIB2(expr);
           })
      .def("arith.AnalyzerSetZ3TimeoutMs",
           [](Analyzer analyzer, int64_t timeout_ms) {
             analyzer->z3_prover.SetTimeoutMs(static_cast<unsigned>(timeout_ms));
           })
      .def("arith.AnalyzerSetZ3RLimit",
           [](Analyzer analyzer, int64_t rlimit) {
             analyzer->z3_prover.SetRLimit(static_cast<unsigned>(rlimit));
           })
      .def("arith.AnalyzerGetZ3Stats",
           [](Analyzer analyzer) { return analyzer->z3_prover.GetStats(); })
      .def("arith.AnalyzerGetEnabledExtensions",
           [](Analyzer analyzer) {
             return static_cast<std::int64_t>(analyzer->rewrite_simplify.GetEnabledExtensions());
           })
      .def("arith.AnalyzerSetEnabledExtensions", [](Analyzer analyzer, int64_t flags) {
        analyzer->rewrite_simplify.SetEnabledExtensions(
            static_cast<RewriteSimplifier::Extension>(flags));
      });
}

}  // namespace arith
}  // namespace tvm
