/*
 * 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/ir_mutator_with_analyzer.cc
 */
#include "ir_mutator_with_analyzer.h"

#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>

#include <vector>

namespace tvm {
namespace arith {

using namespace tirx;

namespace {

enum class CompareKind { kEQ, kLT, kLE, kGT, kGE };

bool TryGetIntImm(const PrimExpr& expr, int64_t* value) {
  if (const auto* imm = expr.as<IntImmNode>()) {
    *value = imm->value;
    return true;
  }
  return false;
}

void AppendFloorDivConstraints(const FloorDivNode* div, int64_t value, CompareKind kind,
                               std::vector<PrimExpr>* out) {
  int64_t divisor_value = 0;
  if (!TryGetIntImm(div->b, &divisor_value) || divisor_value <= 0) return;

  PrimType dtype = div->a.ty();
  PrimExpr divisor = IntImm(dtype, divisor_value);
  PrimExpr k = IntImm(dtype, value);
  PrimExpr lo = k * divisor;
  PrimExpr hi = (k + IntImm(dtype, 1)) * divisor;

  switch (kind) {
    case CompareKind::kEQ:
      out->push_back(div->a >= lo);
      out->push_back(div->a < hi);
      break;
    case CompareKind::kLT:
      out->push_back(div->a < lo);
      break;
    case CompareKind::kLE:
      out->push_back(div->a < hi);
      break;
    case CompareKind::kGT:
      out->push_back(div->a >= hi);
      break;
    case CompareKind::kGE:
      out->push_back(div->a >= lo);
      break;
  }
}

CompareKind InvertCompare(CompareKind kind) {
  switch (kind) {
    case CompareKind::kEQ:
      return CompareKind::kEQ;
    case CompareKind::kLT:
      return CompareKind::kGT;
    case CompareKind::kLE:
      return CompareKind::kGE;
    case CompareKind::kGT:
      return CompareKind::kLT;
    case CompareKind::kGE:
      return CompareKind::kLE;
  }
  return CompareKind::kEQ;
}

void CollectFloorDivConstraintsFromCompare(const PrimExpr& lhs, const PrimExpr& rhs,
                                           CompareKind kind, std::vector<PrimExpr>* out) {
  int64_t value = 0;
  if (const auto* div = lhs.as<FloorDivNode>()) {
    if (TryGetIntImm(rhs, &value)) AppendFloorDivConstraints(div, value, kind, out);
  }
  if (const auto* div = rhs.as<FloorDivNode>()) {
    if (TryGetIntImm(lhs, &value)) {
      AppendFloorDivConstraints(div, value, InvertCompare(kind), out);
    }
  }
}

void CollectDerivedConstraintFacts(const PrimExpr& condition, std::vector<PrimExpr>* out) {
  if (const auto* and_node = condition.as<AndNode>()) {
    CollectDerivedConstraintFacts(and_node->a, out);
    CollectDerivedConstraintFacts(and_node->b, out);
    return;
  }
  if (const auto* call = condition.as<CallNode>()) {
    if (call->op.same_as(tirx::builtin::bitwise_and()) && call->args.size() == 2) {
      PrimExpr lhs = call->args[0].as_or_throw<PrimExpr>();
      PrimExpr rhs = call->args[1].as_or_throw<PrimExpr>();
      if (lhs.ty().MatchesElementType(DLDataTypeCode::kDLBool, 8) &&
          rhs.ty().MatchesElementType(DLDataTypeCode::kDLBool, 8)) {
        CollectDerivedConstraintFacts(lhs, out);
        CollectDerivedConstraintFacts(rhs, out);
        return;
      }
    }
  }
  if (const auto* eq = condition.as<EQNode>()) {
    CollectFloorDivConstraintsFromCompare(eq->a, eq->b, CompareKind::kEQ, out);
  } else if (const auto* lt = condition.as<LTNode>()) {
    CollectFloorDivConstraintsFromCompare(lt->a, lt->b, CompareKind::kLT, out);
  } else if (const auto* le = condition.as<LENode>()) {
    CollectFloorDivConstraintsFromCompare(le->a, le->b, CompareKind::kLE, out);
  } else if (const auto* gt = condition.as<GTNode>()) {
    CollectFloorDivConstraintsFromCompare(gt->a, gt->b, CompareKind::kGT, out);
  } else if (const auto* ge = condition.as<GENode>()) {
    CollectFloorDivConstraintsFromCompare(ge->a, ge->b, CompareKind::kGE, out);
  }
}

void EnterConstraintFacts(WithGroup<ConstraintContext>* constraints, AnalyzerObj* analyzer,
                          const PrimExpr& condition) {
  arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(analyzer);
  constraints->Emplace(analyzer_ref, condition);
  std::vector<PrimExpr> derived;
  CollectDerivedConstraintFacts(condition, &derived);
  for (const PrimExpr& fact : derived) {
    constraints->Emplace(analyzer_ref, fact);
  }
}

}  // namespace

void IRMutatorWithAnalyzer::MarkBufferParamShapes(const tirx::PrimFunc& func) {
  // Mark all symbolic buffer-parameter shape values as positive.
  for (const tirx::Var& param : func->params) {
    if (!param->ty.as<tirx::BufferTypeNode>()) {
      continue;
    }
    tirx::BufferVar buffer(param);
    for (PrimExpr shape : buffer->shape) {
      analyzer_->MarkGlobalNonNegValue(shape);
    }
  }
}

ffi::Array<PrimExpr> IRMutatorWithAnalyzer::IterMapSimplifyWithContext(
    const ffi::Array<PrimExpr>& indices, bool non_trivial_only) {
  PrimExpr pred = IntImm::Bool(true);
  for (PrimExpr val : iter_predicates_) {
    pred = pred && val;
  }
  int n = indices.size();
  arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(this->analyzer_);
  ffi::Array<PrimExpr> simplified = arith::IterMapSimplify(
      indices, this->iter_vars_, pred, arith::IterMapLevel::Surjective, analyzer_ref);
  if (non_trivial_only) {
    for (int i = 0; i < n; ++i) {
      if (simplified[i]->IsInstance<IntImmNode>() && indices[i].as<PrimVar>()) {
        simplified.Set(i, indices[i]);
      }
    }
  }
  return simplified;
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const ForNode* op) {
  return constraint_scope_.WithNewScope([&]() -> Stmt {
    // record the loop variable as iterators
    Range dom = Range::FromMinExtent(op->min, op->extent);
    analyzer_->Bind(op->loop_var, dom);
    iter_vars_.Set(op->loop_var, dom);

    PrimExpr min = this->VisitPrimExpr(op->min);
    PrimExpr extent = this->VisitPrimExpr(op->extent);
    ffi::Optional<PrimExpr> step{std::nullopt};
    if (op->step.has_value()) {
      step = this->VisitPrimExpr(*op->step);
    }
    Stmt body = constraint_scope_.WithNewScope([&]() -> Stmt {
      EnterConstraintFacts(&constraint_scope_.Current(), analyzer_,
                           extent > IntImm(extent.ty(), 0));
      return this->VisitStmt(op->body);
    });
    if (min.same_as(op->min) && extent.same_as(op->extent) && body.same_as(op->body) &&
        step.same_as(op->step)) {
      return ffi::GetRef<Stmt>(op);
    } else {
      auto n = this->CopyOnWrite(op);
      n->min = std::move(min);
      n->extent = std::move(extent);
      n->step = std::move(step);
      n->body = std::move(body);
      return Stmt(n);
    }
  });
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const SBlockNode* op) {
  return constraint_scope_.WithNewScope([&]() -> Stmt {
    for (const auto& iter_var : op->iter_vars) {
      analyzer_->Bind(iter_var->var, iter_var->dom);
      iter_vars_.Set(iter_var->var, iter_var->dom);
    }
    return StmtExprMutator::VisitStmt_(op);
  });
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const BindNode* op) {
  Expr value = this->VisitExpr(op->value);
  if (auto prim_value = value.as<PrimExpr>()) {
    if (SideEffect(prim_value.value()) <= CallEffectKind::kPure) {
      analyzer_->Bind(op->var, prim_value.value());
    }
  }
  if (value.same_as(op->value)) return ffi::GetRef<Stmt>(op);
  auto n = this->CopyOnWrite(op);
  n->value = std::move(value);
  return Stmt(n);
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const IfThenElseNode* op) {
  return constraint_scope_.WithNewScope([&]() -> Stmt {
    PrimExpr condition = this->VisitPrimExpr(op->condition);
    PrimExpr real_condition = condition;

    if (auto call = condition.as<CallNode>()) {
      static const Op& likely_op = Op::Get("tirx.likely");
      if (call->op.same_as(likely_op)) {
        real_condition = call->args[0].as_or_throw<PrimExpr>();
      }
    }

    Stmt then_case;
    ffi::Optional<Stmt> else_case;
    constraint_scope_.WithNewScope([&]() {
      EnterConstraintFacts(&constraint_scope_.Current(), analyzer_, real_condition);
      WithRecordIterPredicate(real_condition, [&] { then_case = this->VisitStmt(op->then_case); });
    });
    if (op->else_case) {
      PrimExpr neg_condition = analyzer_->rewrite_simplify(Not(real_condition));
      constraint_scope_.WithNewScope([&]() {
        constraint_scope_.Current().Emplace(analyzer_, neg_condition);
        else_case = this->VisitStmt(op->else_case.value());
      });
    }
    if (is_one(real_condition)) return then_case;
    if (is_zero(real_condition)) {
      return else_case.value_or(Evaluate(0));
    }

    if (condition.same_as(op->condition) && then_case.same_as(op->then_case) &&
        else_case.same_as(op->else_case)) {
      return ffi::GetRef<Stmt>(op);
    } else {
      auto n = this->CopyOnWrite(op);
      n->condition = std::move(condition);
      n->then_case = std::move(then_case);
      n->else_case = std::move(else_case);
      return Stmt(n);
    }
  });
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const AttrStmtNode* op) {
  return constraint_scope_.WithNewScope([&]() -> Stmt {
    if (op->attr_key == tirx::attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) {
      IterVar iv = op->node.as_or_throw<IterVar>();
      TVM_FFI_ICHECK_NE(iv->thread_tag.length(), 0U);
      Range dom = Range::FromMinExtent(IntImm(op->value.ty(), 0), op->value);
      analyzer_->Bind(iv->var, dom);
      iter_vars_.Set(iv->var, dom);
    }
    return StmtExprMutator::VisitStmt_(op);
  });
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const AssertStmtNode* op) {
  PrimExpr condition = this->VisitPrimExpr(op->condition);
  constraint_scope_.Current().Emplace(analyzer_, condition);

  if (condition.same_as(op->condition)) {
    return ffi::GetRef<Stmt>(op);
  } else {
    auto n = this->CopyOnWrite(op);
    n->condition = std::move(condition);
    return Stmt(n);
  }
}

Stmt IRMutatorWithAnalyzer::VisitStmt_(const SeqStmtNode* op) {
  // SeqStmt does NOT get WithNewScope — constraints accumulate across siblings.
  return StmtExprMutator::VisitStmt_(op);
}

Expr IRMutatorWithAnalyzer::VisitExpr_(const CallNode* op) {
  // add condition context to if_then_else
  static const Op& if_then_else_op = Op::Get("tirx.if_then_else");
  if (op->op.same_as(if_then_else_op)) {
    PrimExpr cond = this->VisitPrimExpr(op->args[0].as_or_throw<PrimExpr>());
    Expr true_value, false_value;
    constraint_scope_.WithNewScope([&]() {
      EnterConstraintFacts(&constraint_scope_.Current(), analyzer_, cond);
      WithRecordIterPredicate(cond, [&] { true_value = this->VisitExpr(op->args[1]); });
    });
    {
      PrimExpr not_cond = Not(cond);
      constraint_scope_.WithNewScope([&]() {
        constraint_scope_.Current().Emplace(analyzer_, not_cond);
        WithRecordIterPredicate(not_cond, [&] { false_value = this->VisitExpr(op->args[2]); });
      });
    }
    if (is_zero(cond)) {
      return false_value;
    }
    if (is_one(cond)) {
      return true_value;
    }
    if (cond.same_as(op->args[0]) && true_value.same_as(op->args[1]) &&
        false_value.same_as(op->args[2])) {
      return ffi::GetRef<Expr>(op);
    } else {
      return Call(op->ty, op->op, {cond, true_value, false_value}, op->attrs, {}, op->span);
    }
  }
  return StmtExprMutator::VisitExpr_(op);
}

Expr IRMutatorWithAnalyzer::VisitExpr_(const LetNode* op) {
  PrimExpr value = this->VisitPrimExpr(op->value);
  if (SideEffect(value) <= CallEffectKind::kPure) {
    analyzer_->Bind(op->var, value);
  }
  // We keep the let-binding here
  // as sub-class may or maynot choose to replace it.
  PrimExpr body = this->VisitPrimExpr(op->body);
  if (value.same_as(op->value) && body.same_as(op->body)) {
    return ffi::GetRef<PrimExpr>(op);
  } else {
    return Let(op->var, value, body);
  }
}

Expr IRMutatorWithAnalyzer::VisitExpr_(const SelectNode* op) {
  PrimExpr cond = this->VisitPrimExpr(op->condition);
  PrimExpr true_value, false_value;
  constraint_scope_.WithNewScope([&]() {
    EnterConstraintFacts(&constraint_scope_.Current(), analyzer_, cond);
    true_value = VisitPrimExpr(op->true_value);
  });
  {
    PrimExpr neg_cond = analyzer_->rewrite_simplify(Not(cond));
    constraint_scope_.WithNewScope([&]() {
      constraint_scope_.Current().Emplace(analyzer_, neg_cond);
      false_value = VisitPrimExpr(op->false_value);
    });
  }
  if (is_zero(cond)) {
    return false_value;
  }
  if (is_one(cond)) {
    return true_value;
  }
  // normal path
  if (cond.same_as(op->condition) && true_value.same_as(op->true_value) &&
      false_value.same_as(op->false_value)) {
    return ffi::GetRef<PrimExpr>(op);
  } else {
    return Select(cond, true_value, false_value);
  }
}

Expr IRMutatorWithAnalyzer::VisitExpr_(const ReduceNode* op) {
  // Setup the domain information before simplification.
  for (const IterVar& iv : op->axis) {
    analyzer_->Bind(iv->var, iv->dom);
  }
  // Recursively call simplification when necessary.
  return StmtExprMutator::VisitExpr_(op);
}

}  // namespace arith
}  // namespace tvm
