/******************************************************************************
 * Top contributors (to current version):
 *   Andrew Reynolds, Morgan Deters, 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.
 * ****************************************************************************
 *
 * Model for Theory UF.
 */

#include "cvc5_private.h"

#ifndef CVC5__THEORY_UF_MODEL_H
#define CVC5__THEORY_UF_MODEL_H

#include <vector>

#include "expr/node.h"

namespace cvc5::internal {
namespace theory {

class TheoryModel;
class Rewriter;

namespace uf {

class UfModelTreeNode
{
public:
  UfModelTreeNode(){}
  /** the data */
  std::map< Node, UfModelTreeNode > d_data;
  /** the value of this tree node (if all paths lead to same value) */
  Node d_value;
public:
  //is this model tree empty?
  bool isEmpty() { return d_data.empty() && d_value.isNull(); }
  //clear
  void clear();
  /** setValue function */
  void setValue(const TheoryModel* m,
                Node n,
                Node v,
                std::vector<int>& indexOrder,
                bool ground,
                int argIndex);
  /** getFunctionValue */
  Node getFunctionValue(const std::vector<Node>& args,
                        int index,
                        Node argDefaultValue);
  /** update function */
  void update(const TheoryModel* m);
  /** simplify function */
  void simplify( Node op, Node defaultVal, int argIndex );
  /** is total ? */
  bool isTotal( Node op, int argIndex );
public:
 void debugPrint(std::ostream& out,
                 const TheoryModel* m,
                 std::vector<int>& indexOrder,
                 int ind = 0,
                 int arg = 0);
};

class UfModelTree
{
private:
  //the op this model is for
  Node d_op;
  //the order we will treat the arguments
  std::vector< int > d_index_order;
  //the data
  UfModelTreeNode d_tree;
public:
  //constructors
  UfModelTree(){}
  UfModelTree( Node op ) : d_op( op ){
    TypeNode tn = d_op.getType();
    for( int i=0; i<(int)(tn.getNumChildren()-1); i++ ){
      d_index_order.push_back( i );
    }
  }
  UfModelTree( Node op, std::vector< int >& indexOrder ) : d_op( op ){
    d_index_order.insert( d_index_order.end(), indexOrder.begin(), indexOrder.end() );
  }
  /** clear/reset the function */
  void clear() { d_tree.clear(); }
  /** setValue function
    *
    * For each argument of n with ModelBasisAttribute() set to true will be considered default arguments if ground=false
    *
    */
  void setValue(const TheoryModel* m, Node n, Node v, bool ground = true)
  {
    d_tree.setValue( m, n, v, d_index_order, ground, 0 );
  }
  /** setDefaultValue function */
  void setDefaultValue(const TheoryModel* m, Node v)
  {
    d_tree.setValue( m, Node::null(), v, d_index_order, false, 0 );
  }
  /** getFunctionValue
   * Returns a representation of this function. The body of the function is
   * rewritten if r is non-null.
   */
  Node getFunctionValue(const std::vector<Node>& args, Rewriter* r);
  /** getFunctionValue for args with set prefix */
  Node getFunctionValue(const std::string& argPrefix, Rewriter* r);
  /** update
    *   This will update all values in the tree to be representatives in m.
    */
  void update(const TheoryModel* m) { d_tree.update(m); }
  /** simplify the tree */
  void simplify() { d_tree.simplify( d_op, Node::null(), 0 ); }
  /** is this tree total? */
  bool isTotal() { return d_tree.isTotal( d_op, 0 ); }
  /** is this tree empty? */
  bool isEmpty() { return d_tree.isEmpty(); }
public:
 void debugPrint(std::ostream& out, const TheoryModel* m, int ind = 0)
 {
   d_tree.debugPrint(out, m, d_index_order, ind);
 }
};

}
}
}  // namespace cvc5::internal

#endif
