/******************************************************************************
 * Top contributors (to current version):
 *   Andrew Reynolds, Haniel Barbosa, Mathias Preiner
 *
 * This file is part of the cvc5 project.
 *
 * Copyright (c) 2009-2025 by the authors listed in the file AUTHORS
 * in the top-level source directory and their institutional affiliations.
 * All rights reserved.  See the file COPYING in the top-level source
 * directory for licensing information.
 * ****************************************************************************
 *
 * Implementation of techniques for evaluating terms with recursively
 * defined functions.
 */

#include "theory/quantifiers/fun_def_evaluator.h"

#include "options/quantifiers_options.h"
#include "theory/quantifiers/quantifiers_attributes.h"
#include "theory/rewriter.h"

using namespace cvc5::internal::kind;

namespace cvc5::internal {
namespace theory {
namespace quantifiers {

FunDefEvaluator::FunDefEvaluator(Env& env) : EnvObj(env) {}

bool FunDefEvaluator::assertDefinition(Node q)
{
  Trace("fd-eval") << "FunDefEvaluator: assertDefinition " << q << std::endl;
  Node head = QuantAttributes::getFunDefHead(q);
  if (head.isNull())
  {
    size_t index;
    if (getDefinitionIndex(q, index))
    {
      Assert(q[1].getKind() == Kind::EQUAL);
      addDefinition(q[1][index], q[1][1 - index], q);
      return true;
    }
    Trace("fd-eval") << "...not a definition" << std::endl;
    // not a function definition
    return false;
  }
  Node body = QuantAttributes::getFunDefBody(q);
  Assert(!body.isNull());
  addDefinition(head, body, q);
  return true;
}

bool FunDefEvaluator::isDefinition(const Node& q) const
{
  size_t index;
  return getDefinitionIndex(q, index);
}

bool FunDefEvaluator::getDefinitionIndex(const Node& q, size_t& index) const
{
  Assert(q.getKind() == Kind::FORALL);
  if (q[1].getKind() == Kind::EQUAL)
  {
    size_t nvars = q[0].getNumChildren();
    // check if we are (f x) = t or t = (f x).
    for (size_t i = 0; i < 2; i++)
    {
      size_t nchild = q[1][i].getNumChildren();
      if (q[1][i].getKind() != Kind::APPLY_UF || nchild != nvars)
      {
        continue;
      }
      bool isMacro = true;
      // if this side of the equality is (f x1 ... xn) where the quantified
      // formula is (forall ((x1 T1) ... (xn Tn)) ...).
      for (size_t j = 0; j < nvars; j++)
      {
        if (q[1][i][j] != q[0][j])
        {
          isMacro = false;
          break;
        }
      }
      if (isMacro)
      {
        index = i;
        return true;
      }
    }
  }
  return false;
}

void FunDefEvaluator::addDefinition(const Node& head,
                                    const Node& body,
                                    const Node& q)
{
  // h possibly with zero arguments?
  Node f = head.hasOperator() ? head.getOperator() : head;
  Assert(d_funDefMap.find(f) == d_funDefMap.end())
      << "FunDefEvaluator::assertDefinition: function already defined";
  d_funDefs.push_back(q);
  FunDefInfo& fdi = d_funDefMap[f];
  fdi.d_quant = q;
  fdi.d_body = body;
  fdi.d_args.insert(fdi.d_args.end(), q[0].begin(), q[0].end());
  Trace("fd-eval") << "FunDefEvaluator: function " << f << " is defined with "
                   << fdi.d_args << " / " << fdi.d_body << std::endl;
}

Node FunDefEvaluator::evaluateDefinitions(Node n) const
{
  // should do standard rewrite before this call
  Assert(rewrite(n) == n);
  Trace("fd-eval") << "FunDefEvaluator: evaluateDefinitions " << n << std::endl;
  NodeManager* nm = nodeManager();
  std::unordered_map<TNode, unsigned> funDefCount;
  std::unordered_map<TNode, unsigned>::iterator itCount;
  std::unordered_map<TNode, Node> visited;
  std::unordered_map<TNode, Node>::iterator it;
  // to ensure all nodes are ref counted
  std::unordered_set<Node> keep;
  std::map<Node, FunDefInfo>::const_iterator itf;
  std::vector<TNode> visit;
  TNode cur;
  TNode curEval;
  Node f;
  visit.push_back(n);
  do
  {
    cur = visit.back();
    visit.pop_back();
    it = visited.find(cur);
    Trace("fd-eval-debug") << "evaluate subterm " << cur << std::endl;

    if (it == visited.end())
    {
      if (cur.isConst())
      {
        Trace("fd-eval-debug") << "constant " << cur << std::endl;
        visited[cur] = cur;
      }
      else if (cur.getKind() == Kind::ITE)
      {
        Trace("fd-eval-debug") << "ITE " << cur << std::endl;
        visited[cur] = Node::null();
        visit.push_back(cur);
        visit.push_back(cur[0]);
      }
      else
      {
        Trace("fd-eval-debug") << "recurse " << cur << std::endl;
        visited[cur] = Node::null();
        visit.push_back(cur);
        for (const Node& cn : cur)
        {
          visit.push_back(cn);
        }
      }
    }
    else
    {
      curEval = it->second;
      if (curEval.isNull())
      {
        Trace("fd-eval-debug") << "from arguments " << cur << std::endl;
        Node ret = cur;
        bool childChanged = false;
        std::vector<Node> children;
        Kind ck = cur.getKind();
        // If a parameterized node that is not APPLY_UF (which is handled below,
        // we add it to the children vector.
        if (ck != Kind::APPLY_UF
            && cur.getMetaKind() == metakind::PARAMETERIZED)
        {
          children.push_back(cur.getOperator());
        }
        else if (ck == Kind::ITE)
        {
          // get evaluation of condition
          it = visited.find(cur[0]);
          Assert(it != visited.end());
          Assert(!it->second.isNull());
          if (!it->second.isConst())
          {
            Trace("fd-eval") << "FunDefEvaluator: couldn't reduce condition of "
                                "ITE to const, FAIL\n";

            Trace("fd-eval")
                << "...failing eval was " << it->second << std::endl;
            return Node::null();
          }
          // pick child to evaluate depending on condition eval
          unsigned childIdxToEval = it->second.getConst<bool>() ? 1 : 2;
          Trace("fd-eval-debug2")
              << "FunDefEvaluator: result of ITE condition : "
              << it->second.getConst<bool>() << "\n";
          // the result will be the result of evaluation the child
          visited[cur] = cur[childIdxToEval];
          // push back self and child. The child will be evaluated first and
          // result will be the result of evaluation child
          visit.push_back(cur);
          visit.push_back(cur[childIdxToEval]);
          Trace("fd-eval-debug2") << "FunDefEvaluator: result will be from : "
                                  << cur[childIdxToEval] << "\n";
          continue;
        }
        unsigned child CVC5_UNUSED = 0;
        for (const Node& cn : cur)
        {
          it = visited.find(cn);
          Assert(it != visited.end());
          Assert(!it->second.isNull());
          childChanged = childChanged || cn != it->second;
          children.push_back(it->second);
          Trace("fd-eval-debug2") << "argument " << child++
                                  << " eval : " << it->second << std::endl;
        }
        if (cur.getKind() == Kind::APPLY_UF)
        {
          // need to evaluate it
          f = cur.getOperator();
          Trace("fd-eval-debug2")
              << "FunDefEvaluator: need to eval " << f << "\n";
          itf = d_funDefMap.find(f);
          itCount = funDefCount.find(f);
          if (itCount == funDefCount.end())
          {
            funDefCount[f] = 0;
            itCount = funDefCount.find(f);
          }
          if (itf == d_funDefMap.end()
              || itCount->second > options().quantifiers.sygusRecFunEvalLimit)
          {
            Trace("fd-eval")
                << "FunDefEvaluator: "
                << (itf == d_funDefMap.end() ? "no definition for "
                                             : "too many evals for ")
                << f << ", FAIL" << std::endl;
            return Node::null();
          }
          ++funDefCount[f];
          // get the function definition
          Node sbody = itf->second.d_body;
          Trace("fd-eval-debug2")
              << "FunDefEvaluator: definition: " << sbody << "\n";
          const std::vector<Node>& args = itf->second.d_args;
          if (!args.empty())
          {
            // invoke it on arguments using the evaluator
            sbody = evaluate(sbody, args, children);
            if (TraceIsOn("fd-eval-debug2"))
            {
              Trace("fd-eval-debug2")
                  << "FunDefEvaluator: evaluation with args:\n";
              for (const Node& ch : children)
              {
                Trace("fd-eval-debug2") << "..." << ch << "\n";
              }
              Trace("fd-eval-debug2")
                  << "FunDefEvaluator: results in " << sbody << "\n";
            }
            Assert(!sbody.isNull());
          }
          keep.insert(sbody);
          // our result is the result of the body
          visited[cur] = sbody;
          // If its not constant, we push back self and the substituted body.
          // Thus, we evaluate the body first; our result will be the result of
          // evaluating the body.
          if (!sbody.isConst())
          {
            Trace("fd-eval-debug2") << "FunDefEvaluator: will map " << cur
                                    << " from body " << sbody << "\n";
            visit.push_back(cur);
            visit.push_back(sbody);
          }
        }
        else
        {
          if (childChanged)
          {
            ret = nm->mkNode(cur.getKind(), children);
            ret = rewrite(ret);
            keep.insert(ret);
          }
          Trace("fd-eval-debug2") << "built from arguments " << ret << "\n";
          visited[cur] = ret;
        }
      }
      else if (cur != curEval && !curEval.isConst())
      {
        Trace("fd-eval-debug") << "from body " << cur << std::endl;
        Trace("fd-eval-debug") << "and eval  " << curEval << std::endl;
        // we had to evaluate our body, which should have a definition now
        it = visited.find(curEval);
        if (it == visited.end())
        {
          Trace("fd-eval-debug2") << "eval without definition\n";
          // this is the case where curEval was not a constant but it was
          // irreducible, for example (DT_SYGUS_EVAL e args)
          visited[cur] = curEval;
        }
        else
        {
          Trace("fd-eval-debug2")
              << "eval with definition " << it->second << "\n";
          visited[cur] = it->second;
        }
      }
    }
  } while (!visit.empty());
  Trace("fd-eval") << "FunDefEvaluator: return " << visited[n] << ", SUCCESS\n";
  Assert(visited.find(n) != visited.end());
  Assert(!visited.find(n)->second.isNull());
  return visited[n];
}

bool FunDefEvaluator::hasDefinitions() const { return !d_funDefMap.empty(); }

const std::vector<Node>& FunDefEvaluator::getDefinitions() const
{
  return d_funDefs;
}
Node FunDefEvaluator::getDefinitionFor(Node f) const
{
  std::map<Node, FunDefInfo>::const_iterator it = d_funDefMap.find(f);
  if (it != d_funDefMap.end())
  {
    return it->second.d_quant;
  }
  return Node::null();
}
Node FunDefEvaluator::getLambdaFor(Node f) const
{
  std::map<Node, FunDefInfo>::const_iterator it = d_funDefMap.find(f);
  if (it != d_funDefMap.end())
  {
    NodeManager* nm = nodeManager();
    return nm->mkNode(Kind::LAMBDA,
                      nm->mkNode(Kind::BOUND_VAR_LIST, it->second.d_args),
                      it->second.d_body);
  }
  return Node::null();
}

}  // namespace quantifiers
}  // namespace theory
}  // namespace cvc5::internal
