/******************************************************************************
 * Top contributors (to current version):
 *   Andrew Reynolds, Clark Barrett, Aina Niemetz
 *
 * 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 model class.
 */
#include "theory/theory_model.h"

#include "expr/attribute.h"
#include "expr/cardinality_constraint.h"
#include "expr/node_algorithm.h"
#include "expr/skolem_manager.h"
#include "expr/sort_to_term.h"
#include "expr/subs.h"
#include "options/quantifiers_options.h"
#include "options/smt_options.h"
#include "options/theory_options.h"
#include "options/uf_options.h"
#include "smt/env.h"
#include "theory/trust_substitutions.h"
#include "theory/uf/function_const.h"
#include "theory/uf/theory_uf_model.h"
#include "util/rational.h"
#include "util/uninterpreted_sort_value.h"

using namespace std;
using namespace cvc5::internal::kind;
using namespace cvc5::context;

namespace cvc5::internal {
namespace theory {

TheoryModel::TheoryModel(Env& env, std::string name, bool enableFuncModels)
    : EnvObj(env),
      d_name(name),
      d_equalityEngine(nullptr),
      d_using_model_core(false),
      d_enableFuncModels(enableFuncModels)
{
  // must use function models when ufHo is enabled
  Assert(d_enableFuncModels || !logicInfo().isHigherOrder());
  d_true = nodeManager()->mkConst(true);
  d_false = nodeManager()->mkConst(false);
}

TheoryModel::~TheoryModel() {}

void TheoryModel::finishInit(eq::EqualityEngine* ee)
{
  Assert(ee != nullptr);
  d_equalityEngine = ee;
  // we do not do congruence on any kind in the model equality engine, with
  // the exception of HO_APPLY for the sake of higher-order.
  d_equalityEngine->addFunctionKind(Kind::HO_APPLY);
  // do not interpret APPLY_UF if we are not assigning function values
  if (!d_enableFuncModels)
  {
    setSemiEvaluatedKind(Kind::APPLY_UF);
  }
  // Equal and not terms are not relevant terms. In other words, asserted
  // equalities and negations of predicates (as terms) do not need to be sent
  // to the model. Regardless, theories should send information to the model
  // that ensures that all assertions are satisfied.
  setIrrelevantKind(Kind::EQUAL);
  setIrrelevantKind(Kind::NOT);
}

void TheoryModel::reset(){
  d_modelCache.clear();
  d_semiEvalCacheSet = false;
  d_semiEvalCache.clear();
  d_sep_heap = Node::null();
  d_sep_nil_eq = Node::null();
  d_reps.clear();
  d_assignExcSet.clear();
  d_aesMaster.clear();
  d_aesSlaves.clear();
  d_rep_set.clear();
  d_uf_terms.clear();
  d_ho_uf_terms.clear();
  d_uf_models.clear();
  d_using_model_core = false;
  d_model_core.clear();
}

void TheoryModel::setHeapModel( Node h, Node neq ) { 
  d_sep_heap = h;
  d_sep_nil_eq = neq;
}

bool TheoryModel::getHeapModel(Node& h, Node& neq) const
{
  if (d_sep_heap.isNull() || d_sep_nil_eq.isNull())
  {
    return false;
  }
  h = d_sep_heap;
  neq = d_sep_nil_eq;
  return true;
}

std::vector<Node> TheoryModel::getDomainElements(TypeNode tn) const
{
  // must be an uninterpreted sort
  Assert(tn.isUninterpretedSort());
  std::vector<Node> elements;
  const std::vector<Node>* type_reps = d_rep_set.getTypeRepsOrNull(tn);
  if (type_reps == nullptr || type_reps->empty())
  {
    // This is called when t is a sort that does not occur in this model.
    // Sorts are always interpreted as non-empty, thus we add a single element.
    // We use mkGroundValue here, since domain elements must all be
    // of UNINTERPRETED_SORT_VALUE kind.
    elements.push_back(NodeManager::mkGroundValue(tn));
    return elements;
  }
  // The representatives are skolems. We convert them to uninterpreted sort
  // values here for printing purposes.
  NodeManager* nm = nodeManager();
  for (size_t i = 0, size = type_reps->size(); i < size; i++)
  {
    UninterpretedSortValue usv = UninterpretedSortValue(tn, Integer(i));
    elements.push_back(nm->mkConst<UninterpretedSortValue>(usv));
  }
  return elements;
}

Node TheoryModel::getValue(TNode n) const
{
  //apply substitutions
  Node nn = d_env.getTopLevelSubstitutions().apply(n);
  nn = rewrite(nn);
  Trace("model-getvalue-debug") << "[model-getvalue] getValue : substitute " << n << " to " << nn << std::endl;
  //get value in model
  nn = getModelValue(nn);
  if (nn.isNull())
  {
    return nn;
  }
  if (nn.getKind() == Kind::FUNCTION_ARRAY_CONST)
  {
    // return the lambda instead
    nn = uf::FunctionConst::toLambda(nn);
  }
  if (nn.getKind() == Kind::LAMBDA)
  {
    if (options().theory.condenseFunctionValues)
    {
      // normalize the body. Do not normalize the entire node, which
      // involves array normalization.
      NodeManager* nm = nodeManager();
      nn = nm->mkNode(Kind::LAMBDA, nn[0], rewrite(nn[1]));
    }
  }
  else
  {
    //normalize
    nn = rewrite(nn);
  }
  Trace("model-getvalue") << "[model-getvalue] getValue( " << n << " ): " << std::endl
                          << "[model-getvalue] returning " << nn << std::endl;
  Assert(nn.getType() == n.getType());
  return nn;
}

Node TheoryModel::simplify(TNode n) const
{
  std::unordered_set<Node> syms;
  expr::getSymbols(n, syms);
  // if there are free symbols
  if (!syms.empty())
  {
    // apply a substitution mapping those symbols to their model values
    Subs subs;
    for (const Node& s : syms)
    {
      subs.add(s, getValue(s));
    }
    // rewrite the result
    return rewrite(subs.apply(n));
  }
  return n;
}

bool TheoryModel::isModelCoreSymbol(Node s) const
{
  if (!d_using_model_core)
  {
    return true;
  }
  Assert(s.isVar() && s.getKind() != Kind::BOUND_VARIABLE);
  return d_model_core.find(s) != d_model_core.end();
}

size_t TheoryModel::getCardinality(const TypeNode& tn) const
{
  //for now, we only handle cardinalities for uninterpreted sorts
  Assert(tn.isUninterpretedSort());
  if (d_rep_set.hasType(tn))
  {
    Trace("model-getvalue-debug")
        << "Get cardinality sort, #rep : "
        << d_rep_set.getNumRepresentatives(tn) << std::endl;
    return d_rep_set.getNumRepresentatives(tn);
  }
  Trace("model-getvalue-debug")
      << "Get cardinality sort, unconstrained, return 1." << std::endl;
  return 1;
}

Node TheoryModel::getModelValue(TNode n) const
{
  std::unordered_map<Node, Node>::iterator it = d_modelCache.find(n);
  if (it != d_modelCache.end()) {
    return (*it).second;
  }
  Trace("model-getvalue-debug") << "Get model value " << n << " ... ";
  Trace("model-getvalue-debug") << d_equalityEngine->hasTerm(n) << std::endl;
  Kind nk = n.getKind();
  if (n.isConst() || nk == Kind::BOUND_VARIABLE)
  {
    d_modelCache[n] = n;
    return n;
  }

  Node ret = n;
  NodeManager* nm = nodeManager();

  // If it has children, evaluate them. We do this even if n is an unevaluatable
  // or semi-unevaluatable operator. Notice this includes quantified
  // formulas. It may not be possible in general to evaluate bodies of
  // quantified formulas, because they have free variables. Regardless, we
  // may often be able to evaluate the body of a quantified formula to true,
  // e.g. forall x. P(x) where P = lambda x. true.
  if (n.getNumChildren() > 0)
  {
    Trace("model-getvalue-debug")
        << "Get model value children " << n << std::endl;
    std::vector<Node> children;
    if (n.getKind() == Kind::APPLY_UF)
    {
      Node op = getModelValue(n.getOperator());
      Trace("model-getvalue-debug") << "  operator : " << op << std::endl;
      children.push_back(op);
    }
    else if (n.getMetaKind() == kind::metakind::PARAMETERIZED)
    {
      children.push_back(n.getOperator());
    }
    // evaluate the children
    for (size_t i = 0, nchild = n.getNumChildren(); i < nchild; ++i)
    {
      if (n.isClosure() && i==0)
      {
        // do not evaluate bound variable lists
        ret = n[0];
      }
      else
      {
        ret = getModelValue(n[i]);
      }
      Trace("model-getvalue-debug")
          << "  " << n << "[" << i << "] is " << ret << std::endl;
      children.push_back(ret);
    }
    ret = nm->mkNode(nk, children);
    Trace("model-getvalue-debug") << "ret (pre-rewrite): " << ret << std::endl;
    ret = rewrite(ret);
    Trace("model-getvalue-debug") << "ret (post-rewrite): " << ret << std::endl;
    // special cases
    if (ret.getKind() == Kind::CARDINALITY_CONSTRAINT)
    {
      const CardinalityConstraint& cc =
          ret.getOperator().getConst<CardinalityConstraint>();
      Trace("model-getvalue-debug")
          << "get cardinality constraint " << cc.getType() << std::endl;
      size_t cval = getCardinality(cc.getType());
      Assert(cval > 0);
      ret = nm->mkConst(Integer(cval) <= cc.getUpperBound());
    }
    // if the value was constant, we return it. If it was non-constant,
    // we only return it if we are an evaluated kind. This can occur if the
    // children of n failed to evaluate.
    if (ret.isConst() || (
     d_unevaluated_kinds.find(nk) == d_unevaluated_kinds.end()
      && d_semi_evaluated_kinds.find(nk) == d_semi_evaluated_kinds.end()))
    {
      d_modelCache[n] = ret;
      return ret;
    }
    // Note that we discard the evaluation of the arguments here
    Trace("model-getvalue-debug") << "Failed to evaluate " << ret << std::endl;
  }
  // We revert to the original rewritten form of n here.
  ret = rewrite(n);
  Trace("model-getvalue-debug")
      << "Look up " << ret << " in equality engine" << std::endl;
  // return the representative of the term in the equality engine, if it exists
  TypeNode t = ret.getType();
  if (!logicInfo().isHigherOrder() && t.isFunction())
  {
    // functions are in the equality engine, but *not* as first-class members
    // when higher-order is disabled. In this case, we cannot query
    // representatives for functions since they are "internal" nodes according
    // to the equality engine despite hasTerm returning true. However, they are
    // first class members when higher-order is enabled. Hence, the special
    // case here.
  }
  else if (d_equalityEngine->hasTerm(ret))
  {
    Trace("model-getvalue-debug")
        << "get value from representative " << ret << "..." << std::endl;
    ret = getRepresentative(ret);
    d_modelCache[n] = ret;
    Trace("model-getvalue-debug") << "...set rep is " << ret << std::endl;
    return ret;
  }

  Kind rk = ret.getKind();
  // If we are a ground term that is *not* unevaluated, we assign an arbitrary
  // value.
  if (d_unevaluated_kinds.find(rk) == d_unevaluated_kinds.end()
      && !expr::hasBoundVar(ret))
  {
    // If we are a semi-evaluated kind, then we need to check whether we are
    // entailed equal to an existing term. For example, if we are a datatype
    // selector S(x), x is equal to y, and S(y) is a term in the equality
    // engine of this model, then the value of S(x) must be equal to the value
    // of S(y).
    if (d_semi_evaluated_kinds.find(rk) != d_semi_evaluated_kinds.end())
    {
      Node retSev = evaluateSemiEvalTerm(ret);
      if (retSev.isNull())
      {
        // If null, if an argument is not constant, then it *might* be the
        // same as a constrained application of rk. In this case, we cannot
        // evaluate to an arbitrary value, and instead we use ret itself
        // (unevaluated) as the model value.
        for (const Node& rc : ret)
        {
          if (!rc.isConst())
          {
            retSev = ret;
            break;
          }
        }
        
      }
      // if the result was entailed, return it
      if (!retSev.isNull())
      {
        d_modelCache[n] = retSev;
        Trace("model-getvalue-debug")
            << "...semi-evaluated entailed is " << retSev << std::endl;
        return retSev;
      }
      // otherwise we return an arbtirary value below.
      Trace("model-getvalue-debug")
          << "...not semi-evaluated entailed" << std::endl;
    }
    if (t.isFunction())
    {
      if (d_enableFuncModels)
      {
        std::map<Node, Node>::const_iterator entry = d_uf_models.find(n);
        // If we have not assigned the function, assign it now
        if (entry == d_uf_models.end())
        {
          assignFunctionDefault(n);
          entry = d_uf_models.find(n);
          Assert(entry != d_uf_models.end());
        }
        ret = entry->second;
      }
      else
      {
        // if func models not enabled, throw an error
        Unreachable();
      }
    }
    else if (!t.isFirstClass() || t.isRegExp())
    {
      // this is the class for regular expressions
      // we simply invoke the rewriter on them
      ret = rewrite(ret);
    }
    else
    {
      // Unknown term - return first enumerated value for this type
      TypeEnumerator te(n.getType());
      ret = *te;
    }
    d_modelCache[n] = ret;
    return ret;
  }

  d_modelCache[n] = n;
  return n;
}

/** add term */
void TheoryModel::addTermInternal(TNode n)
{
  Assert(d_equalityEngine->hasTerm(n));
  Trace("model-builder-debug2") << "TheoryModel::addTerm : " << n << std::endl;
  //must collect UF terms
  if (n.getKind() == Kind::APPLY_UF)
  {
    Node op = n.getOperator();
    if( std::find( d_uf_terms[ op ].begin(), d_uf_terms[ op ].end(), n )==d_uf_terms[ op ].end() ){
      d_uf_terms[ op ].push_back( n );
      Trace("model-builder-fun") << "Add apply term " << n << std::endl;
    }
  }
  else if (n.getKind() == Kind::HO_APPLY)
  {
    Node op = n[0];
    if( std::find( d_ho_uf_terms[ op ].begin(), d_ho_uf_terms[ op ].end(), n )==d_ho_uf_terms[ op ].end() ){
      d_ho_uf_terms[ op ].push_back( n );
      Trace("model-builder-fun") << "Add ho apply term " << n << std::endl;
    }
  }
  // all functions must be included, marked as higher-order
  if( n.getType().isFunction() ){
    Trace("model-builder-fun") << "Add function variable (without term) " << n << std::endl;
    if( d_uf_terms.find( n )==d_uf_terms.end() ){
      d_uf_terms[n].clear();
    }
    if( d_ho_uf_terms.find( n )==d_ho_uf_terms.end() ){
      d_ho_uf_terms[n].clear();
    }
  }
}

/** assert equality */
bool TheoryModel::assertEquality(TNode a, TNode b, bool polarity)
{
  Assert(d_equalityEngine->consistent());
  if (a == b && polarity) {
    return true;
  }
  Trace("model-builder-assertions")
      << "(assert " << (polarity ? "(= " : "(not (= ") << a << " " << b
      << (polarity ? "));" : ")));");
  d_equalityEngine->assertEquality( a.eqNode(b), polarity, Node::null() );
  Trace("model-builder-assertions")
      << (d_equalityEngine->consistent() ? "" : "...inconsistent!")
      << std::endl;
  return d_equalityEngine->consistent();
}

/** assert predicate */
bool TheoryModel::assertPredicate(TNode a, bool polarity)
{
  Assert(d_equalityEngine->consistent());
  if ((a == d_true && polarity) ||
      (a == d_false && (!polarity))) {
    return true;
  }
  if (a.getKind() == Kind::EQUAL)
  {
    Trace("model-builder-assertions") << "(assert " << (polarity ? " " : "(not ") << a << (polarity ? ");" : "));") << endl;
    d_equalityEngine->assertEquality( a, polarity, Node::null() );
  }
  else
  {
    Trace("model-builder-assertions") << "(assert " << (polarity ? "" : "(not ") << a << (polarity ? ");" : "));") << endl;
    d_equalityEngine->assertPredicate( a, polarity, Node::null() );
  }
  return d_equalityEngine->consistent();
}

/** assert equality engine */
bool TheoryModel::assertEqualityEngine(const eq::EqualityEngine* ee,
                                       const std::set<Node>* termSet)
{
  Assert(d_equalityEngine->consistent());
  eq::EqClassesIterator eqcs_i = eq::EqClassesIterator( ee );
  for (; !eqcs_i.isFinished(); ++eqcs_i) {
    Node eqc = (*eqcs_i);
    bool predicate = false;
    bool predTrue = false;
    bool predFalse = false;
    Trace("model-builder-debug") << "...asserting terms in equivalence class " << eqc;
    if (eqc.getType().isBoolean()) {
      predicate = true;
      predTrue = ee->areEqual(eqc, d_true);
      predFalse = ee->areEqual(eqc, d_false);
      Trace("model-builder-debug") << ", pred = " << predTrue << "/" << predFalse;
    }
    Trace("model-builder-debug") << std::endl;
    eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, ee);
    bool first = true;
    Node rep;
    for (; !eqc_i.isFinished(); ++eqc_i) {
      Node n = *eqc_i;
      // notice that constants are always relevant
      if (termSet != nullptr && termSet->find(n) == termSet->end()
          && !n.isConst())
      {
        Trace("model-builder-debug")
            << "...skip node " << n << " in eqc " << eqc << std::endl;
        continue;
      }
      if (predicate) {
        if (predTrue || predFalse)
        {
          if (!assertPredicate(n, predTrue))
          {
            return false;
          }
        }
        else {
          if (first) {
            rep = n;
            first = false;
          }
          else {
            Node eq = (n).eqNode(rep);
            Trace("model-builder-assertions")
                << "(assert " << eq << ");" << std::endl;
            d_equalityEngine->assertEquality(eq, true, Node::null());
            if (!d_equalityEngine->consistent())
            {
              return false;
            }
          }
        }
      } else {
        if (first) {
          rep = n;
          //add the term (this is specifically for the case of singleton equivalence classes)
          if (!rep.getType().isRegExp())
          {
            d_equalityEngine->addTerm( rep );
            Trace("model-builder-debug") << "Add term to ee within assertEqualityEngine: " << rep << std::endl;
          }
          first = false;
        }
        else {
          if (!assertEquality(n, rep, true))
          {
            return false;
          }
        }
      }
    }
  }
  return true;
}

void TheoryModel::assertSkeleton(TNode n)
{
  Trace("model-builder-reps") << "Assert skeleton : " << n << std::endl;
  Trace("model-builder-reps") << "...rep eqc is : " << getRepresentative(n)
                              << std::endl;
  d_reps[n] = n;
}

void TheoryModel::assignRepresentative(const Node& r,
                                       const Node& n,
                                       bool isFinal)
{
  Trace("model-builder-reps") << "Assign rep : " << r << " " << n << std::endl;
  Assert(r.getType() == n.getType());
  TypeNode tn = r.getType();
  if (isFinal && logicInfo().isHigherOrder() && tn.isFunction())
  {
    assignFunctionDefinition(r, n);
  }
  else
  {
    d_reps[r] = n;
  }
  d_rep_set.add(tn, n);
}

void TheoryModel::setAssignmentExclusionSet(TNode n,
                                            const std::vector<Node>& eset)
{
  // should not be assigned yet
  Assert(d_assignExcSet.find(n) == d_assignExcSet.end());
  Trace("model-builder-debug")
      << "Exclude values of " << n << " : " << eset << std::endl;
  std::vector<Node>& aes = d_assignExcSet[n];
  aes.insert(aes.end(), eset.begin(), eset.end());
}

void TheoryModel::setAssignmentExclusionSetGroup(
    const std::vector<TNode>& group, const std::vector<Node>& eset)
{
  if (group.empty())
  {
    return;
  }
  // for efficiency, we store a single copy of eset and set a slave/master
  // relationship
  setAssignmentExclusionSet(group[0], eset);
  std::vector<Node>& gslaves = d_aesSlaves[group[0]];
  for (unsigned i = 1, gsize = group.size(); i < gsize; ++i)
  {
    Node gs = group[i];
    // set master
    d_aesMaster[gs] = group[0];
    // add to slaves
    gslaves.push_back(gs);
  }
}

bool TheoryModel::getAssignmentExclusionSet(TNode n,
                                            std::vector<Node>& group,
                                            std::vector<Node>& eset)
{
  // does it have a master?
  std::map<Node, Node>::iterator itm = d_aesMaster.find(n);
  if (itm != d_aesMaster.end())
  {
    return getAssignmentExclusionSet(itm->second, group, eset);
  }
  std::map<Node, std::vector<Node> >::iterator ita = d_assignExcSet.find(n);
  if (ita == d_assignExcSet.end())
  {
    return false;
  }
  eset.insert(eset.end(), ita->second.begin(), ita->second.end());
  group.push_back(n);
  // does it have slaves?
  ita = d_aesSlaves.find(n);
  if (ita != d_aesSlaves.end())
  {
    group.insert(group.end(), ita->second.begin(), ita->second.end());
  }
  return true;
}

bool TheoryModel::hasAssignmentExclusionSets() const
{
  return !d_assignExcSet.empty();
}

bool TheoryModel::isUsingModelCore() const { return d_using_model_core; }
void TheoryModel::setUsingModelCore()
{
  d_using_model_core = true;
  d_model_core.clear();
}

void TheoryModel::recordModelCoreSymbol(Node sym) { d_model_core.insert(sym); }

void TheoryModel::setUnevaluatedKind(Kind k) { d_unevaluated_kinds.insert(k); }

void TheoryModel::setSemiEvaluatedKind(Kind k)
{
  d_semi_evaluated_kinds.insert(k);
}

void TheoryModel::setIrrelevantKind(Kind k) { d_irrKinds.insert(k); }

const std::set<Kind>& TheoryModel::getIrrelevantKinds() const
{
  return d_irrKinds;
}

bool TheoryModel::isLegalElimination(TNode x, TNode val)
{
  return !expr::hasSubtermKinds(d_unevaluated_kinds, val);
}

bool TheoryModel::hasTerm(TNode a)
{
  return d_equalityEngine->hasTerm( a );
}

Node TheoryModel::getRepresentative(TNode a) const
{
  if (d_equalityEngine->hasTerm(a))
  {
    Node r = d_equalityEngine->getRepresentative( a );
    std::map<Node, Node>::const_iterator itr = d_reps.find(r);
    // note that d_reps[r]=r for equivalence classes that haven't been assigned
    if (itr != d_reps.end() && itr->second != r)
    {
      return itr->second;
    }
    // special case: functions are constructed lazily so if we are higher-order
    // and are looking for the representative of a function eqc, compute it now
    if (logicInfo().isHigherOrder() && a.getType().isFunction())
    {
      assignFunctionDefault(r);
      itr = d_reps.find(r);
      Assert(itr != d_reps.end());
      return itr->second;
    }
    return r;
  }
  return a;
}

bool TheoryModel::areEqual(TNode a, TNode b) const
{
  if( a==b ){
    return true;
  }else if( d_equalityEngine->hasTerm( a ) && d_equalityEngine->hasTerm( b ) ){
    return d_equalityEngine->areEqual( a, b );
  }else{
    return false;
  }
}

bool TheoryModel::areDisequal(TNode a, TNode b)
{
  if( d_equalityEngine->hasTerm( a ) && d_equalityEngine->hasTerm( b ) ){
    return d_equalityEngine->areDisequal( a, b, false );
  }else{
    return false;
  }
}

bool TheoryModel::hasUfTerms(Node f) const
{
  return d_uf_terms.find(f) != d_uf_terms.end();
}

const std::vector<Node>& TheoryModel::getUfTerms(Node f) const
{
  const auto it = d_uf_terms.find(f);
  Assert(it != d_uf_terms.end());
  return it->second;
}

bool TheoryModel::areFunctionValuesEnabled() const
{
  return d_enableFuncModels;
}

void TheoryModel::assignFunctionDefault(Node f) const
{
  if (d_uf_models.find(f) != d_uf_models.end())
  {
    return;
  }
  if (logicInfo().isHigherOrder())
  {
    Trace("model-builder") << "  Assign function value for " << f
                           << " based on curried HO_APPLY" << std::endl;
    assignFunctionDefaultHo(f);
    return;
  }
  Trace("model-builder") << "  Assign function value for " << f
                         << " based on APPLY_UF" << std::endl;
  uf::UfModelTree ufmt(f);
  options::DefaultFunctionValueMode dfvm =
      options().theory.defaultFunctionValueMode;
  Node default_v;
  std::map<Node, std::vector<Node>>::const_iterator itu = d_uf_terms.find(f);
  if (itu != d_uf_terms.end())
  {
    for (const Node& un : itu->second)
    {
      std::vector<TNode> children;
      children.push_back(f);
      Trace("model-builder-debug") << "  process term : " << un << std::endl;
      for (size_t j = 0; j < un.getNumChildren(); ++j)
      {
        Node rc = getRepresentative(un[j]);
        Trace("model-builder-debug2")
            << "    get rep : " << un[j] << " returned " << rc << std::endl;
        Assert(rewrite(rc) == rc);
        children.push_back(rc);
      }
      Node simp = nodeManager()->mkNode(un.getKind(), children);
      Node v = getRepresentative(un);
      Trace("model-builder")
          << "  Setting (" << simp << ") to (" << v << ")" << endl;
      ufmt.setValue(this, simp, v);
      if (dfvm == options::DefaultFunctionValueMode::FIRST)
      {
        default_v = v;
      }
    }
  }
  TypeNode rangeType = f.getType().getRangeType();
  if (dfvm == options::DefaultFunctionValueMode::HOLE)
  {
    NodeManager* nm = nodeManager();
    SkolemManager* sm = nm->getSkolemManager();
    std::vector<Node> cacheVals;
    cacheVals.push_back(nm->mkConst(SortToTerm(rangeType)));
    default_v = sm->mkSkolemFunction(SkolemId::GROUND_TERM, cacheVals);
  }
  else if (default_v.isNull())
  {
    // choose default value from model if none exists
    TypeEnumerator te(rangeType);
    default_v = (*te);
  }
  ufmt.setDefaultValue(this, default_v);
  bool condenseFuncValues = options().theory.condenseFunctionValues;
  if (condenseFuncValues)
  {
    ufmt.simplify();
  }
  std::stringstream ss;
  ss << "_arg_";
  Rewriter* r = condenseFuncValues ? d_env.getRewriter() : nullptr;
  Node val = ufmt.getFunctionValue(ss.str(), r);
  Trace("model-builder-debug") << "...assign via function" << std::endl;
  assignFunctionDefinition(f, val);
}

void TheoryModel::assignFunctionDefaultHo(Node f) const
{
  Assert(logicInfo().isHigherOrder());
  // collect all HO_APPLY terms, modulo equality of the function, which will
  // determine the relevant points of the function value we construct below.
  std::vector<Node> hoTerms;
  std::map<Node, std::vector<Node>>::const_iterator itht;
  if (d_equalityEngine->hasTerm(f))
  {
    // if we are in the equality, look up each function in the equivalence class
    Node r = d_equalityEngine->getRepresentative(f);
    eq::EqClassIterator eqc_i = eq::EqClassIterator(r, d_equalityEngine);
    while (!eqc_i.isFinished())
    {
      Node n = *eqc_i;
      ++eqc_i;
      itht = d_ho_uf_terms.find(n);
      if (itht != d_ho_uf_terms.end())
      {
        hoTerms.insert(hoTerms.end(), itht->second.begin(), itht->second.end());
      }
    }
  }
  else
  {
    // otherwise just take the function itself
    itht = d_ho_uf_terms.find(f);
    if (itht != d_ho_uf_terms.end())
    {
      hoTerms.insert(hoTerms.end(), itht->second.begin(), itht->second.end());
    }
  }
  Trace("model-builder-debug") << "Assign HO function " << f << std::endl;
  TypeNode type = f.getType();
  std::vector<TypeNode> argTypes = type.getArgTypes();
  std::vector<Node> args;
  std::vector<TNode> apply_args;
  options::DefaultFunctionValueMode dfvm =
      options().theory.defaultFunctionValueMode;
  for (size_t i = 0; i < argTypes.size(); i++)
  {
    Node v = nodeManager()->mkBoundVar(argTypes[i]);
    args.push_back(v);
    if (i > 0)
    {
      apply_args.push_back(v);
    }
  }
  // Depending on the default value mode, maybe set the current value (curr).
  // We also remember a default value (currPre) in case there are no terms
  // to assign below.
  TypeNode rangeType = type.getRangeType();
  Node curr, currPre;
  if (dfvm == options::DefaultFunctionValueMode::HOLE)
  {
    NodeManager* nm = nodeManager();
    SkolemManager* sm = nm->getSkolemManager();
    std::vector<Node> cacheVals;
    cacheVals.push_back(nm->mkConst(SortToTerm(rangeType)));
    currPre = sm->mkSkolemFunction(SkolemId::GROUND_TERM, cacheVals);
    curr = currPre;
  }
  else
  {
    TypeEnumerator te(rangeType);
    currPre = (*te);
    if (dfvm == options::DefaultFunctionValueMode::FIRST_ENUM)
    {
      curr = currPre;
    }
  }
  curr = currPre;
  for (const Node& hn : hoTerms)
  {
    Trace("model-builder-debug") << "    process : " << hn << std::endl;
    Assert(hn.getKind() == Kind::HO_APPLY);
    Assert(areEqual(hn[0], f));
    // get representative of the argument, which note may recursively compute
    // more function values.
    Node hni = getRepresentative(hn[1]);
    Trace("model-builder-debug2")
        << "      get rep : " << hn[1] << " returned " << hni << std::endl;
    Assert(hni.getType() == args[0].getType());
    // rewrite to ensure the equality is properly oriented, as required by
    // function constants
    hni = rewrite(args[0].eqNode(hni));
    // get representative of the returned term, which note may recursively
    // compute more function values.
    Node hnv = getRepresentative(hn);
    Trace("model-builder-debug2")
        << "      get rep val : " << hn << " returned " << hnv << std::endl;
    // hnv is expected to be constant but may not be the case if e.g. a
    // non-trivial lambda is given as argument to this function.
    if (!apply_args.empty())
    {
      // Convert to lambda, which is necessary if hnv is a function array
      // constant.
      hnv = uf::FunctionConst::toLambda(hnv);
      Assert(!hnv.isNull() && hnv.getKind() == Kind::LAMBDA
             && hnv[0].getNumChildren() + 1 == args.size());
      std::vector<TNode> largs(hnv[0].begin(), hnv[0].end());
      Assert(largs.size() == apply_args.size());
      hnv = hnv[1].substitute(
          largs.begin(), largs.end(), apply_args.begin(), apply_args.end());
      hnv = rewrite(hnv);
    }
    Assert(hnv.getType() == curr.getType());
    if (curr.isNull())
    {
      curr = hnv;
    }
    else
    {
      curr = nodeManager()->mkNode(Kind::ITE, hni, hnv, curr);
    }
  }
  // if curr was not set, we set it to currPre.
  if (curr.isNull())
  {
    curr = currPre;
  }
  Node val = nodeManager()->mkNode(
      Kind::LAMBDA, nodeManager()->mkNode(Kind::BOUND_VAR_LIST, args), curr);
  Trace("model-builder-debug")
      << "...assign via ho function to " << val << std::endl;
  assignFunctionDefinition(f, val);
}

void TheoryModel::assignFunctionDefinition(Node f, Node f_def) const
{
  Trace("model-builder") << "  Assigning function (" << f << ") to (" << f_def << ")" << endl;
  Assert(d_uf_models.find(f) == d_uf_models.end());

  if (logicInfo().isHigherOrder())
  {
    // we must rewrite the function value since the definition needs to be a
    // constant value. This does not need to be the case if we are assigning a
    // lambda to the equivalence class in isolation, so we do not assert that
    // f_def is constant here.
    f_def = rewrite(f_def);
    Trace("model-builder-debug")
        << "Model value (post-rewrite) : " << f_def << std::endl;
  }

  d_uf_models[f] = f_def;

  if (logicInfo().isHigherOrder() && d_equalityEngine->hasTerm(f))
  {
    Trace("model-builder-debug") << "  ...function is first-class member of equality engine" << std::endl;
    // assign to representative if higher-order
    Node r = d_equalityEngine->getRepresentative( f );
    //always replace the representative, since it is initially assigned to itself
    Trace("model-builder") << "    Assign: Setting function rep " << r << " to " << f_def << endl;
    // should not have assigned it yet, unless it was set to itself
    Assert(d_reps.find(r) == d_reps.end() || d_reps[r] == r)
        << "Function already assigned " << d_reps[r];
    d_reps[r] = f_def;  
    // also assign to other assignable functions in the same equivalence class
    eq::EqClassIterator eqc_i = eq::EqClassIterator(r,d_equalityEngine);
    while( !eqc_i.isFinished() ) {
      Node n = *eqc_i;
      // if an unassigned variable function
      if (isAssignableUf(n) && !hasAssignedFunctionDefinition(n))
      {
        d_uf_models[n] = f_def;
        Trace("model-builder") << "  Assigning function (" << n << ") to function definition of " << f << std::endl;
      }
      ++eqc_i;
    }
    Trace("model-builder-debug") << "  ...finished." << std::endl;
  }
}

bool TheoryModel::hasAssignedFunctionDefinition(Node f) const
{
  return d_uf_models.find(f) != d_uf_models.end();
}

const std::string& TheoryModel::getName() const { return d_name; }

std::string TheoryModel::debugPrintModelEqc() const
{
  std::stringstream ss;
  ss << "--- Equivalence classes:" << std::endl;
  ss << d_equalityEngine->debugPrintEqc() << std::endl;
  ss << "--- Representative map: " << std::endl;
  for (const std::pair<const Node, Node>& r : d_reps)
  {
    ss << r.first << " -> " << r.second << std::endl;
  }
  ss << "---" << std::endl;
  return ss.str();
}

struct IsModelValueTag
{
};
struct IsModelValueComputedTag
{
};
/** Attribute true for expressions that are quasi-values */
using IsModelValueAttr = expr::Attribute<IsModelValueTag, bool>;
using IsModelValueComputedAttr = expr::Attribute<IsModelValueComputedTag, bool>;

bool TheoryModel::isBaseModelValue(TNode n) const
{
  if (n.isConst())
  {
    return true;
  }
  Kind k = n.getKind();
  if (k == Kind::REAL_ALGEBRAIC_NUMBER || k == Kind::LAMBDA
      || k == Kind::WITNESS)
  {
    // we are a value if we are one of the above kinds
    return true;
  }
  return false;
}

bool TheoryModel::isAssignableUf(const Node& n) const
{
  return n.getKind() != Kind::LAMBDA;
}

Node TheoryModel::evaluateSemiEvalTerm(TNode n) const
{
  Assert(d_semi_evaluated_kinds.find(n.getKind())
         != d_semi_evaluated_kinds.end());
  if (!d_semiEvalCacheSet)
  {
    d_semiEvalCacheSet = true;
    // traverse
    eq::EqClassesIterator eqcs_i = eq::EqClassesIterator(d_equalityEngine);
    for (; !eqcs_i.isFinished(); ++eqcs_i)
    {
      Node eqc = (*eqcs_i);
      eq::EqClassIterator eqc_i = eq::EqClassIterator(eqc, d_equalityEngine);
      for (; !eqc_i.isFinished(); ++eqc_i)
      {
        Node t = (*eqc_i);
        Kind k = t.getKind();
        if (d_semi_evaluated_kinds.find(k) == d_semi_evaluated_kinds.end())
        {
          continue;
        }
        Node eqcv = getModelValue(eqc);
        Assert(t.hasOperator());
        Node op = t.getOperator();
        std::vector<Node> targs = getModelValueArgs(t);
        NodeTrie& nt = d_semiEvalCache[op];
        nt.addOrGetTerm(eqcv, targs);
        Trace("semi-eval") << "Semi-eval: SET " << targs << " = " << eqcv
                           << std::endl;
      }
    }
  }
  Trace("semi-eval") << "Semi-eval: EVALUATE " << n << "..." << std::endl;
  Node op = n.getOperator();
  std::unordered_map<Node, NodeTrie>::iterator it = d_semiEvalCache.find(op);
  if (it != d_semiEvalCache.end())
  {
    std::vector<Node> nargs = getModelValueArgs(n);
    Trace("semi-eval") << "Semi-eval: lookup " << nargs << std::endl;
    Node result = it->second.existsTerm(nargs);
    Trace("semi-eval") << "...result is " << result << std::endl;    
    return result;
  }
  return Node::null();
}

std::vector<Node> TheoryModel::getModelValueArgs(TNode n) const
{
  std::vector<Node> args;
  for (const Node& nc : n)
  {
    args.emplace_back(getModelValue(nc));
  }
  return args;
}

bool TheoryModel::isValue(TNode n) const
{
  IsModelValueAttr imva;
  IsModelValueComputedAttr imvca;
  // The list of nodes we are processing, and the current child index of that
  // node we are inspecting. This vector always specifies a single path in the
  // original term n. Notice this index accounts for operators, where the
  // operator of a term is treated as the first child, and subsequently all
  // other children are shifted up by one.
  std::vector<std::pair<TNode, size_t>> visit;
  // the last computed value of whether a node was a value
  bool currentReturn = false;
  visit.emplace_back(n, 0);
  std::pair<TNode, size_t> v;
  while (!visit.empty())
  {
    v = visit.back();
    TNode cur = v.first;
    bool finishedComputing = false;
    // if we just pushed to the stack, do initial checks
    if (v.second == 0)
    {
      if (cur.getAttribute(imvca))
      {
        // already cached
        visit.pop_back();
        currentReturn = cur.getAttribute(imva);
        continue;
      }
      if (isBaseModelValue(cur))
      {
        finishedComputing = true;
        currentReturn = true;
      }
      else if (cur.getNumChildren() == 0)
      {
        // PI is a (non-base) value. We require it as a special case here
        // since nullary operators are represented internal as variables.
        // All other non-constant terms with zero children are not values.
        finishedComputing = true;
        currentReturn = (cur.getKind() == Kind::PI);
      }
      else if (rewrite(cur) != cur)
      {
        // non-rewritten terms are never model values
        finishedComputing = true;
        currentReturn = false;
      }
    }
    else if (!currentReturn)
    {
      // if the last child was not a value, we are not a value
      finishedComputing = true;
    }
    if (!finishedComputing)
    {
      // The only non-constant operators we consider are APPLY_UF and
      // APPLY_SELECTOR. All other operators are either builtin, or should be
      // considered constants, e.g. constructors.
      Kind k = cur.getKind();
      bool hasOperator = k == Kind::APPLY_UF || k == Kind::APPLY_SELECTOR;
      size_t nextChildIndex = v.second;
      if (hasOperator && nextChildIndex > 0)
      {
        // if have an operator, we shift the child index we are looking at
        nextChildIndex--;
      }
      if (nextChildIndex == cur.getNumChildren())
      {
        // finished, we are a value
        currentReturn = true;
      }
      else
      {
        visit.back().second++;
        if (hasOperator && v.second == 0)
        {
          // if we have an operator, process it as the first child
          visit.emplace_back(cur.getOperator(), 0);
        }
        else
        {
          Assert(nextChildIndex < cur.getNumChildren());
          // process the next child, which may be shifted from v.second to
          // account for the operator
          visit.emplace_back(cur[nextChildIndex], 0);
        }
        continue;
      }
    }
    visit.pop_back();
    cur.setAttribute(imva, currentReturn);
    cur.setAttribute(imvca, true);
  }
  Assert(n.getAttribute(imvca));
  return currentReturn;
}

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