#include "solver/abstract/abstraction_module.h"

#include <algorithm>

#include "bv/bitvector.h"
#include "node/kind_info.h"
#include "node/node_manager.h"
#include "node/node_ref_vector.h"
#include "node/node_utils.h"
#include "solver/abstract/abstraction_lemmas.h"
#include "solver/bv/aig_bitblaster.h"
#include "util/hash.h"

#ifndef NDEBUG
#include "solving_context.h"
#endif

namespace bzla::abstract {

using namespace node;

/* --- AbstractionModule public --------------------------------------------- */

AbstractionModule::AbstractionModule(Env& env, SolverState& state)
    : d_env(env),
      d_logger(env.logger()),
      d_solver_state(state),
      d_rewriter(env.rewriter()),
      d_active_term_abstractions(state.backtrack_mgr()),
      d_active_assertion_abstractions(state.backtrack_mgr()),
      d_assertion_abstractions_cache(state.backtrack_mgr()),
      d_lemma_cache(state.backtrack_mgr()),
      // Some lemmas are not valid for bit-vectors of size 1 or 2. Hence, we
      // only start abstracting from size 3+ instead of guarding each lemma
      // separately.
      d_opt_minimum_size(std::max(env.options().abstraction_bv_size(),
                                  static_cast<uint64_t>(3))),
      d_opt_eager_refine(env.options().abstraction_eager_refine()),
      d_opt_value_inst_limit(env.options().abstraction_value_limit()),
      d_opt_value_inst_only(env.options().abstraction_value_only()),
      d_opt_abstract_assertions(env.options().abstraction_assert()),
      d_opt_assertion_refinements(env.options().abstraction_assert_refs()),
      d_opt_inc_bitblast(env.options().abstraction_inc_bitblast()),
      d_stats(env.statistics(), "solver::abstract::")
{
  bool opt_initial_lemmas = env.options().abstraction_initial_lemmas();

  NodeManager& nm = d_env.nm();
  if (env.options().abstraction_bv_mul())
  {
    auto& mul_abstr_lemmas = d_abstr_lemmas[Kind::BV_MUL];
    mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_POW2>(nm));
    mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_NEG_POW2>(nm));
    mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_IC>(nm));
    mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_ODD>(nm));
    if (!opt_initial_lemmas)
    {
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF1>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF3>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN3>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN4>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN5>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN6>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF14>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF15>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN9>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF18>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN11>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN12>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REFN13>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF13>(nm));
      mul_abstr_lemmas.emplace_back(new Lemma<LemmaKind::MUL_REF12>(nm));
    }
  }

  if (env.options().abstraction_bv_udiv())
  {
    auto& udiv_abstr_lemmas = d_abstr_lemmas[Kind::BV_UDIV];
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_POW2>(nm));
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF1>(nm));
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF2>(nm));
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF3>(nm));
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF4>(nm));
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF5>(nm));
    udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF6>(nm));
    if (!opt_initial_lemmas)
    {
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF7>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF8>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF9>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF10>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF11>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF12>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF13>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF14>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF15>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF16>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF17>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF18>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF19>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF20>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF21>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF23>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF24>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF25>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF26>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF27>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF28>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF29>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF30>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF31>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF32>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF33>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF34>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF36>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF37>(nm));
      udiv_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UDIV_REF38>(nm));
    }
  }

  if (env.options().abstraction_bv_urem())
  {
    auto& urem_abstr_lemmas = d_abstr_lemmas[Kind::BV_UREM];
    urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_POW2>(nm));
    urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF1>(nm));
    urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF2>(nm));
    urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF3>(nm));
    urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF4>(nm));
    urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF5>(nm));
    // urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF6>(nm));
    if (!opt_initial_lemmas)
    {
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF7>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF8>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF9>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF10>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF11>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF12>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF13>(nm));
      urem_abstr_lemmas.emplace_back(new Lemma<LemmaKind::UREM_REF14>(nm));
    }
  }

  if (env.options().abstraction_ite())
  {
    d_abstr_lemmas.try_emplace(Kind::ITE);
  }

  if (env.options().abstraction_bv_add())
  {
    auto& add_abstr_lemmas = d_abstr_lemmas[Kind::BV_ADD];
    add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_ZERO>(nm));
    add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_SAME>(nm));
    add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_INV>(nm));
    add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_OVFL>(nm));
    add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_NOOVFL>(nm));
    add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_OR>(nm));
    if (!opt_initial_lemmas)
    {
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF1>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF2>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF3>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF4>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF5>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF6>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF7>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF8>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF9>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF10>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF11>(nm));
      add_abstr_lemmas.emplace_back(new Lemma<LemmaKind::ADD_REF12>(nm));
    }
  }
}

AbstractionModule::~AbstractionModule() {}

void
AbstractionModule::register_abstraction(const Node& node)
{
  assert(is_abstraction(node));
  if (d_opt_abstract_assertions
      && d_abstraction_cache_assertions.find(node)
             != d_abstraction_cache_assertions.end())
  {
    d_active_assertion_abstractions.push_back(node);
  }
  else
  {
    d_active_term_abstractions.push_back(node);
  }
}

bool
AbstractionModule::is_abstraction(const Node& node)
{
  return node.kind() == Kind::AM_ABSTRACT;
}

void
AbstractionModule::check()
{
  Log(1);
  Log(1) << "*** check abstractions";
  util::Timer timer(d_stats.time_check);
  ++d_stats.num_checks;

#ifndef NDEBUG
  // verify_lemmas();
  // std::unordered_map<LemmaKind, uint64_t> map;
  // score_lemmas(Kind::BV_MUL, 4, map);
#endif

  d_added_lemma = false;

  // New abstraction may be added while checking
  for (size_t i = 0; i < d_active_term_abstractions.size(); ++i)
  {
    // Do not use reference here, since d_active_abstractions may change when
    // calling check_term_abstraction().
    const Node abstr_term = d_active_term_abstractions[i];
    check_term_abstraction(abstr_term);
  }

  // Check abstracted assertions.
  // Assertion abstractions must be checked after we are done checking that
  // all term abstractions are consistent.
  if (!d_added_lemma && d_opt_abstract_assertions)
  {
    check_assertion_abstractions();
  }

  // Abstraction refinements that are added as "last resort" are only added if
  // we expanded all violated assertions.
  // This includes:
  //   - value instantiation lemmas
  //   - bit-blasting lemmas
  if (!d_added_lemma)
  {
    for (const auto& [node, lemma, lk] : d_lemma_buffer)
    {
      if (lemma_no_abstract(lemma, lk))
      {
        // Increment counter for value instantiations
        if (is_lemma_kind_value(lk))
        {
          ++d_value_insts[node];
        }
      }
    }
  }
  d_lemma_buffer.clear();
}

const Node&
AbstractionModule::process(const Node& term)
{
  NodeManager& nm = d_env.nm();

  node_ref_vector visit{term};
  do
  {
    const Node& cur     = visit.back();
    auto [it, inserted] = d_abstraction_cache.try_emplace(cur);
    if (inserted)
    {
      // No need to go below quantifiers since abstraction will happen when
      // quantifier instantiation is added.
      if (cur.kind() == Kind::FORALL || cur.kind() == Kind::AM_ABSTRACT)
      {
        it->second = cur;
      }
      else
      {
        visit.insert(visit.end(), cur.begin(), cur.end());
      }
      continue;
    }
    else if (it->second.is_null())
    {
      Node rebuilt = d_env.rewriter().rewrite(
          utils::rebuild_node(nm, cur, d_abstraction_cache));

      if (abstract(rebuilt))
      {
        it->second = nm.mk_node(Kind::AM_ABSTRACT, {rebuilt});
        ++d_stats.num_terms;
        d_stats.terms << rebuilt.kind();
      }
      else
      {
        it->second = rebuilt;
      }
    }
    visit.pop_back();
  } while (!visit.empty());

  return d_abstraction_cache.at(term);
}

const Node&
AbstractionModule::process_assertion(const Node& assertion, bool is_lemma)
{
  Log(2) << "process assertion: " << assertion << " (" << is_lemma << ")";
  const Node& processed = process(assertion);

  // Do not abstract assertions that are values, consts or lemmas.
  if (d_opt_abstract_assertions && !processed.is_value()
      && !processed.is_const() && !is_lemma)
  {
    if (processed.kind() != Kind::AM_ABSTRACT)
    {
      Node abstr = d_env.nm().mk_node(Kind::AM_ABSTRACT, {processed});
      Log(2) << "abstract assertion: " << processed << " (abstr: " << abstr
             << ", orig: " << assertion << ")";
      d_abstraction_cache[abstr] = abstr;
      d_abstraction_cache_assertions.emplace(abstr, assertion);
      d_active_assertion_abstractions.push_back(abstr);
      return d_abstraction_cache.at(abstr);
    }
    return processed;
  }

  // Map original assertion to processed assertion for unsat cores.
  if (processed != assertion)
  {
    d_abstraction_cache_assertions.emplace(processed, assertion);
  }

  return processed;
}

bool
AbstractionModule::is_processed_assertion(const Node& assertion)
{
  auto it = d_abstraction_cache_assertions.find(assertion);
  return it != d_abstraction_cache_assertions.end() && it->second != assertion;
}

const Node&
AbstractionModule::get_original_assertion(const Node& processed_assertion)
{
  auto it = d_abstraction_cache_assertions.find(processed_assertion);
  assert(it != d_abstraction_cache_assertions.end());
  return it->second;
}

Node
AbstractionModule::remove_abstractions(const Node& node) const
{
  NodeManager& nm = d_env.nm();

  std::unordered_map<Node, Node> cache;
  node_ref_vector visit{node};
  do
  {
    const Node& cur     = visit.back();
    auto [it, inserted] = cache.try_emplace(cur);
    if (inserted)
    {
      visit.insert(visit.end(), cur.begin(), cur.end());
      continue;
    }
    else if (it->second.is_null())
    {
      if (cur.kind() == Kind::AM_ABSTRACT)
      {
        it->second = cache.at(cur[0]);
      }
      else
      {
        it->second = utils::rebuild_node(nm, cur, cache);
      }
    }
    visit.pop_back();
  } while (!visit.empty());

  return cache.at(node);
}

/* --- AbstractionModule private -------------------------------------------- */

bool
AbstractionModule::abstract(const Node& node) const
{
  Kind k = node.kind();
  assert(d_abstr_lemmas.find(k) == d_abstr_lemmas.end()
         || KindInfo::num_children(k) > 1);
  return d_abstr_lemmas.find(k) != d_abstr_lemmas.end()
         && d_opt_minimum_size > 0 && node[1].type().is_bv()
         && node[1].type().bv_size() >= d_opt_minimum_size;
}

bool
AbstractionModule::check_lemma(const AbstractionLemma* lem,
                               const Node& val_x,
                               const Node& val_s,
                               const Node& val_t,
                               const Node& x,
                               const Node& s,
                               const Node& t)
{
  Node inst = lem->instance(val_x, val_s, val_t);
  Node lemma;
  if (!inst.is_null())
  {
    inst = d_rewriter.eval(inst);
    assert(inst.is_value());
    if (!inst.value<bool>())
    {
      lemma = lem->instance(x, s, t);
    }
  }
  else
  {
    inst = lem->instance(val_x, val_s, val_t, val_x, val_s, val_t);
    if (!inst.is_null())
    {
      inst = d_rewriter.eval(inst);
      assert(inst.is_value());
      if (!inst.value<bool>())
      {
        lemma = lem->instance(val_x, val_s, val_t, x, s, t);
      }
    }
  }

  if (!lemma.is_null())
  {
    Log(2) << lem->kind() << " inconsistent";
    return lemma_no_abstract(lemma, lem->kind());
  }

  return false;
}

void
AbstractionModule::check_term_abstraction(const Node& abstr)
{
  assert(abstr.kind() == Kind::AM_ABSTRACT);

  const Node& node = abstr[0];
  Log(2) << "check abstraction: " << abstr;
  Log(2) << "abstracted node: " << abstr[0];

  Kind kind = node.kind();
  assert(d_abstr_lemmas.find(kind) != d_abstr_lemmas.end());

  if (kind == Kind::ITE)
  {
    check_term_abstraction_ite(abstr, node);
    return;
  }

  NodeManager& nm   = d_env.nm();
  const Node& x     = node[0];
  const Node& s     = node[1];
  const Node& t     = abstr;
  Node val_x        = d_solver_state.value(x);
  Node val_s        = d_solver_state.value(s);
  Node val_t        = d_solver_state.value(t);
  Node val_expected = d_rewriter.eval(nm.mk_node(kind, {val_x, val_s}));

  if (val_t == val_expected)
  {
    Log(2) << "skip: assignment correct";
    return;
  }

  Log(2) << "x: " << x;
  Log(2) << "s: " << s;
  Log(2) << "t: " << t;
  Log(2) << "val_x: " << val_x;
  Log(2) << "val_s: " << val_s;
  Log(2) << "val_t: " << val_t;

  bool added_lemma = false;
  if (!d_opt_value_inst_only)
  {
    auto it = d_abstr_lemmas.find(kind);
    assert(it != d_abstr_lemmas.end());
    const auto& to_check = it->second;
    for (const auto& lem : to_check)
    {
      added_lemma = check_lemma(lem.get(), val_x, val_s, val_t, x, s, t);
      if (!added_lemma && KindInfo::is_commutative(kind))
      {
        added_lemma = check_lemma(lem.get(), val_s, val_x, val_t, s, x, t);
      }
      if (added_lemma && !d_opt_eager_refine)
      {
        break;
      }
    }
  }

  // Inconsistent value, but no lemma violated, add value-based lemma.
  if (!added_lemma)
  {
    uint64_t limit = d_opt_value_inst_limit > 0
                         ? val_x.type().bv_size() / d_opt_value_inst_limit
                         : 0;
    if (kind != Kind::EQUAL && d_value_insts[node] < limit)
    {
      LemmaKind lk = lemma_kind_value(kind);
      Log(2) << lk << " inconsistent";
      Node lemma =
          nm.mk_node(Kind::IMPLIES,
                     {nm.mk_node(Kind::AND,
                                 {
                                     nm.mk_node(Kind::EQUAL, {x, val_x}),
                                     nm.mk_node(Kind::EQUAL, {s, val_s}),
                                 }),
                      nm.mk_node(Kind::EQUAL, {t, val_expected})});
      d_lemma_buffer.emplace_back(node, lemma, lk);
      if (kind == Kind::BV_MUL && val_x == val_s)
      {
        ++d_value_insts_square[node];
      }
    }
    // Incrementally bit-blast abstracted term starting from LSB
    else if (d_opt_inc_bitblast
             && (kind == Kind::BV_MUL || kind == Kind::BV_ADD))
    {
      constexpr uint64_t increment = 32;
      const auto& bv_t = val_t.value<BitVector>();
      const auto& bv_e = val_expected.value<BitVector>();
      auto bv_xor      = bv_t.bvxor(bv_e);
      uint64_t upper               = bv_xor.count_trailing_zeros() + increment;
      uint64_t size                = bv_t.size();
      upper = std::min(upper, size) - 1;

      Node extr_x = nm.mk_node(Kind::BV_EXTRACT, {x}, {upper, 0});
      Node extr_s = nm.mk_node(Kind::BV_EXTRACT, {s}, {upper, 0});
      Node extr_t = nm.mk_node(Kind::BV_EXTRACT, {t}, {upper, 0});
      Node term   = nm.mk_node(kind, {extr_x, extr_s});
      Node lemma  = nm.mk_node(Kind::EQUAL, {extr_t, term});
      d_lemma_buffer.emplace_back(node,
                                  lemma,
                                  upper + 1 == size ? LemmaKind::BITBLAST_FULL
                                                    : LemmaKind::BITBLAST_INC);
    }
    // Use special square encoding for BV_MUL if value instantiations were all
    // square instantiations.
    else if (kind == Kind::BV_MUL && d_value_insts_square[node] > 0
             && d_value_insts_square[node] >= d_value_insts[node]
             && val_x == val_s)
    {
      d_lemma_buffer.emplace_back(
          node,
          nm.mk_node(
              Kind::IMPLIES,
              {nm.mk_node(Kind::EQUAL, {x, s}),
               nm.mk_node(Kind::EQUAL, {t, nm.mk_node(Kind::BV_MUL, {x, x})})}),
          LemmaKind::BITBLAST_BV_MUL_SQUARE);
      d_value_insts_square[node] = 0;  // Reset, next time we fully bit-blast
    }
    else
    {
      // Fully bit-blast abstracted term
      Node term  = nm.mk_node(kind, {x, s});
      Node lemma = nm.mk_node(Kind::EQUAL, {t, term});
      LemmaKind lk;
      switch (kind)
      {
        case Kind::BV_MUL: lk = LemmaKind::BITBLAST_BV_MUL; break;
        case Kind::BV_UDIV: lk = LemmaKind::BITBLAST_BV_UDIV; break;
        case Kind::BV_UREM: lk = LemmaKind::BITBLAST_BV_UREM; break;
        default: lk = LemmaKind::BITBLAST_FULL;
      }
      d_lemma_buffer.emplace_back(node, lemma, lk);
    }
  }
}

void
AbstractionModule::check_term_abstraction_ite(const Node& abstr,
                                              const Node& node)
{
  const Node& c  = node[0];
  const Node& bt = node[1];
  const Node& bf = node[2];
  Node val_c     = d_solver_state.value(c);
  Node val_t     = d_solver_state.value(abstr);

  bool cond = val_c.value<bool>();
  if (cond)
  {
    Node val_bt = d_solver_state.value(bt);
    if (val_t == val_bt)
    {
      Log(2) << "skip: assignment correct";
      return;
    }
  }
  else
  {
    Node val_bf = d_solver_state.value(bf);
    if (val_t == val_bf)
    {
      Log(2) << "skip: assignment correct";
      return;
    }
  }

  NodeManager& nm = d_env.nm();
  if (cond)
  {
    lemma_no_abstract(
        nm.mk_node(Kind::IMPLIES, {c, nm.mk_node(Kind::EQUAL, {abstr, bt})}),
        LemmaKind::ITE_EXPAND);
  }
  else
  {
    lemma_no_abstract(nm.mk_node(Kind::IMPLIES,
                                 {nm.mk_node(Kind::NOT, {c}),
                                  nm.mk_node(Kind::EQUAL, {abstr, bf})}),
                      LemmaKind::ITE_EXPAND);
  }
}

bool
AbstractionModule::check_assertion_abstractions()
{
  uint64_t nadded = 0;
  NodeManager& nm = d_env.nm();
  for (size_t i = 0, size = d_active_assertion_abstractions.size(); i < size;
       ++i)
  {
    const Node& abstr = d_active_assertion_abstractions[i];
    auto it           = d_assertion_abstractions_cache.find(abstr);
    if (it != d_assertion_abstractions_cache.end())
    {
      continue;
    }
    assert(abstr.kind() == Kind::AM_ABSTRACT);
    const Node& assertion = abstr[0];
    Node val              = d_solver_state.value(assertion);
    if (!val.value<bool>())
    {
      Log(2) << "violated assertion: " << abstr;
      Log(2) << "abstr assertion:    " << assertion;
      Node lemma = nm.mk_node(Kind::EQUAL, {assertion, abstr});
      lemma_no_abstract(lemma, LemmaKind::ASSERTION);
      d_assertion_abstractions_cache.insert(abstr);
      ++nadded;
      if (nadded >= d_opt_assertion_refinements)
      {
        break;
      }
    }
  }
  return nadded > 0;
}

bool
AbstractionModule::lemma_no_abstract(const Node& lemma, LemmaKind lk)
{
  // Make sure that lemma is rewritten before adding to the cache.
  Node lem = d_rewriter.rewrite(lemma);
  // Cache lemma so that we won't consider it for abstraction.
  d_abstraction_cache.emplace(lem, lem);
  auto [it, inserted] = d_lemma_cache.insert(lem);
  assert(inserted || d_opt_eager_refine);
  if (inserted && d_solver_state.lemma(lem))
  {
    d_added_lemma = true;
    d_stats.lemmas << lk;
    return true;
  }
  return false;
}

bool
AbstractionModule::lemma_abstract(const Node& lemma, LemmaKind lk)
{
  if (d_solver_state.lemma(lemma))
  {
    d_added_lemma = true;
    d_stats.lemmas << lk;
    return true;
  }
  return false;
}

void
AbstractionModule::score_lemmas(
    Kind kind,
    uint64_t bv_size,
    std::unordered_map<LemmaKind, uint64_t>& rank_map) const
{
  NodeManager& nm = d_env.nm();
  uint64_t max    = 1u << bv_size;
  std::vector<Node> values;
  std::vector<std::vector<std::vector<bool>>> results_lemmas(
      max, std::vector<std::vector<bool>>(max, std::vector<bool>(max, true)));

  // Create all possible values [0, max[
  for (uint64_t i = 0; i < max; ++i)
  {
    values.push_back(nm.mk_value(BitVector::from_ui(bv_size, i)));
  }

  // Compute all results for kind
  uint64_t optimal_score = 0;
  for (uint64_t i = 0; i < values.size(); ++i)
  {
    for (uint64_t j = 0; j < values.size(); ++j)
    {
      for (uint64_t k = 0; k < values.size(); ++k)
      {
        Node val = d_rewriter.eval(
            nm.mk_node(Kind::EQUAL,
                       {values[k], nm.mk_node(kind, {values[i], values[j]})}));
        assert(val.is_value());
        if (val.value<bool>())
        {
          ++optimal_score;
        }
      }
    }
  }

  std::cout << std::fixed;
  uint64_t max_score   = max * max * max;
  uint64_t final_score = max_score;
  std::cout << "lemma score (worst: " << final_score << ", best: " << max * max
            << ")" << std::endl;

  for (const auto& lem : d_abstr_lemmas.at(kind))
  {
    uint64_t score            = 0;
    uint64_t prev_final_score = final_score;
    // Compute result for each triplet (x, s, t)
    for (uint64_t i = 0; i < values.size(); ++i)
    {
      for (uint64_t j = 0; j < values.size(); ++j)
      {
        // const Node& expected = results[i][j];
        for (uint64_t k = 0; k < values.size(); ++k)
        {
          Node inst = lem->instance(values[i], values[j], values[k]);
          if (inst.is_null())
          {
            inst = lem->instance(values[i],
                                 values[j],
                                 values[k],
                                 values[i],
                                 values[j],
                                 values[k]);
          }
          bool res = true;
          if (!inst.is_null())
          {
            inst = d_rewriter.rewrite(inst);
            assert(inst.is_value());
            res = inst.value<bool>();
          }

          // check commutative case
          if (kind == Kind::BV_MUL)
          {
            Node instc = lem->instance(values[j], values[i], values[k]);
            if (instc.is_null())
            {
              instc = lem->instance(values[j],
                                    values[i],
                                    values[k],
                                    values[j],
                                    values[i],
                                    values[k]);
            }
            if (!instc.is_null())
            {
              instc = d_rewriter.rewrite(instc);
              res   = res & instc.value<bool>();
            }
          }

          auto overall_res = results_lemmas[i][j][k];
          // Count cases when lemma is true (including false positives)
          if (res)
          {
            ++score;
          }
          // Count number of ruled out triplets
          else if (overall_res)
          {
            --final_score;
          }
          results_lemmas[i][j][k] = overall_res & res;
        }
      }
    }
    rank_map[lem->kind()] = score;
    int64_t diff          = final_score - prev_final_score;
    std::cout << lem->kind() << ": " << score << "/" << max_score
              << " (final: " << final_score << ", diff: " << diff << ", "
              << static_cast<double>(diff) / max_score * 100 << "%)"
              << std::endl;
  }
  std::cout << "final score:   " << final_score << " "
            << static_cast<double>(final_score) / max_score * 100
            << "% (wrong results: " << final_score - (max * max) << ")"
            << std::endl;
  std::cout << "optimal score: " << optimal_score << " "
            << static_cast<double>(optimal_score) / max_score * 100 << "%"
            << std::endl;
}

void
AbstractionModule::rank_lemmas_by_circuit_size()
{
  Env env(d_env.nm(), d_env.sat_factory());
  bv::AigBitblaster bb;
  NodeManager& nm = d_env.nm();
  Type bv32       = nm.mk_bv_type(32);
  std::unordered_map<Kind, std::vector<std::pair<LemmaKind, uint64_t>>>
      lemma_sizes;

  std::unordered_map<Kind, uint64_t> circuit_size;
  for (const auto& [kind, lemmas] : d_abstr_lemmas)
  {
    if (KindInfo::num_children(kind) == 2)
    {
      Node x                       = nm.mk_const(bv32);
      Node s                       = nm.mk_const(bv32);
      Node t                       = nm.mk_const(bv32);
      uint64_t size_overall_before = bb.num_aig_ands();
      for (const auto& lem : lemmas)
      {
        Node inst = lem->instance(x, s, t);
        if (inst.is_null())
        {
          // Conditional on x == s, hence manual computation needed
          if (lem->kind() == LemmaKind::BITBLAST_BV_MUL_SQUARE)
          {
            inst =
                nm.mk_node(Kind::IMPLIES,
                           {nm.mk_node(Kind::EQUAL, {x, s}),
                            nm.mk_node(Kind::EQUAL,
                                       {t, nm.mk_node(Kind::BV_MUL, {x, x})})});
          }
          else if (lem->kind() == LemmaKind::MUL_POW2)
          {
            Node val_pow2 =
                nm.mk_value(BitVector::mk_one(bv32.bv_size()).ibvshl(2));
            inst = lem->instance(val_pow2, s, t, x, s, t);
          }
          else if (lem->kind() == LemmaKind::MUL_NEG_POW2)
          {
            Node val_pow2 = nm.mk_value(
                BitVector::mk_one(bv32.bv_size()).ibvshl(2).ibvneg());
            inst = lem->instance(val_pow2, s, t, x, s, t);
          }
          else
          {
            inst = lem->instance(x, s, t, x, s, t);
          }
        }
        if (!inst.is_null())
        {
          inst                 = env.rewriter().rewrite(inst);
          uint64_t size_before = bb.num_aig_ands();
          bb.bitblast(inst);
          uint64_t circuit_size = bb.num_aig_ands() - size_before;
          lemma_sizes[kind].emplace_back(lem->kind(), circuit_size);
        }
      }
      std::cout << kind << " total lemma size: "
                << bb.num_aig_ands() - size_overall_before << std::endl;
      {
        Node x               = nm.mk_const(bv32);
        Node s               = nm.mk_const(bv32);
        Node t               = nm.mk_node(kind, {x, s});
        uint64_t size_before = bb.num_aig_ands();
        bb.bitblast(t);
        std::cout << kind
                  << " circuit size: " << bb.num_aig_ands() - size_before
                  << std::endl;
        circuit_size[kind] = bb.num_aig_ands() - size_before;
      }
    }
  }

  std::unordered_map<LemmaKind, uint64_t> rank_map;
  for (auto& [k, lemmas] : lemma_sizes)
  {
    std::sort(lemmas.begin(), lemmas.end(), [](const auto& p1, const auto& p2) {
      return p1.second < p2.second;
    });
    uint64_t sum = 0;
    bool reached = false;
    for (const auto& [lk, size] : lemmas)
    {
      rank_map.emplace(lk, size);
      sum += size;
      std::cout << size << " " << lk << " (sum: " << sum << "/"
                << circuit_size[k] << ")" << std::endl;
      if (!reached && sum >= circuit_size[k])
      {
        std::cout << "--- circuit size reached (" << circuit_size[k] << ") ---"
                  << std::endl;
        reached = true;
      }
    }
    std::sort(d_abstr_lemmas.at(k).begin(),
              d_abstr_lemmas.at(k).end(),
              [&rank_map](const auto& l1, const auto& l2) {
                return rank_map[l1->kind()] < rank_map[l2->kind()];
              });
    std::cout << "score: " << k << std::endl;
    std::unordered_map<LemmaKind, uint64_t> rm;
    score_lemmas(k, 6, rm);
  }

  std::cout << "final ranking:" << std::endl;
  std::cout << "std::unordered_map<LemmaKind, uint64_t> rank_map = {";
  for (const auto& [lk, size] : rank_map)
  {
    std::cout << "{LemmaKind::" << lk << "," << size << "}," << std::endl;
  }
  std::cout << "};" << std::endl;
  abort();
}

void
AbstractionModule::rank_lemmas_by_score()
{
  std::unordered_map<LemmaKind, uint64_t> rank_map;
  score_lemmas(Kind::BV_MUL, 6, rank_map);
  score_lemmas(Kind::BV_UDIV, 6, rank_map);
  score_lemmas(Kind::BV_UREM, 6, rank_map);

  std::cout << "std::unordered_map<LemmaKind, uint64_t> rank_map = {";
  for (const auto& [lk, score] : rank_map)
  {
    std::cout << "{LemmaKind::" << lk << "," << score << "}," << std::endl;
  }
  std::cout << "};" << std::endl;
  abort();
}

void
AbstractionModule::print_initial_lemmas() const
{
  Node x, s, t;

  NodeManager& nm                = d_env.nm();
  Type bv4                       = nm.mk_bv_type(4);
  x                              = nm.mk_const(bv4, "x");
  s                              = nm.mk_const(bv4, "s");
  t                              = nm.mk_const(bv4, "t");
  std::vector<BitVector> pow2    = {BitVector::from_ui(4, 1),
                                    BitVector::from_ui(4, 2),
                                    BitVector::from_ui(4, 4),
                                    BitVector::from_ui(4, 8)};
  std::vector<BitVector> negpow2 = {BitVector::from_ui(4, 1).bvneg(),
                                    BitVector::from_ui(4, 2).bvneg(),
                                    BitVector::from_ui(4, 4).bvneg()};
  for (const auto& lemma : d_abstr_lemmas.at(Kind::BV_MUL))
  {
    const auto lk = lemma->kind();
    if (lk == LemmaKind::MUL_REF1)
    {
      break;
    }

    if (lk == LemmaKind::MUL_POW2)
    {
      uint64_t i = 1;
      for (const auto& p2 : pow2)
      {
        Node npow2 = nm.mk_value(p2);
        std::cout
            << "(define-fun lemma_" << lemma->kind() << "_" << i
            << " ((x (_ BitVec 4)) (s (_ BitVec 4)) (t (_ BitVec 4))) Bool ";
        std::cout << lemma->instance(npow2, s, t, x, s, t);
        std::cout << ")" << std::endl;
        std::cout
            << "(define-fun lemma_" << lemma->kind() << "_c" << i++
            << " ((x (_ BitVec 4)) (s (_ BitVec 4)) (t (_ BitVec 4))) Bool ";
        std::cout << lemma->instance(npow2, x, t, s, x, t);
        std::cout << ")" << std::endl;
      }
    }
    else if (lk == LemmaKind::MUL_NEG_POW2)
    {
      uint64_t i = 1;
      for (const auto& p2 : negpow2)
      {
        Node npow2 = nm.mk_value(p2);
        std::cout
            << "(define-fun lemma_" << lemma->kind() << "_" << i
            << " ((x (_ BitVec 4)) (s (_ BitVec 4)) (t (_ BitVec 4))) Bool ";
        std::cout << lemma->instance(npow2, s, t, x, s, t);
        std::cout << ")" << std::endl;
        std::cout
            << "(define-fun lemma_" << lemma->kind() << "_c" << i++
            << " ((x (_ BitVec 4)) (s (_ BitVec 4)) (t (_ BitVec 4))) Bool ";
        std::cout << lemma->instance(npow2, x, t, s, x, t);
        std::cout << ")" << std::endl;
      }
    }
    else
    {
      std::cout
          << "(define-fun lemma_" << lemma->kind()
          << " ((x (_ BitVec 4)) (s (_ BitVec 4)) (t (_ BitVec 4))) Bool ";
      std::cout << lemma->instance(x, s, t);
      std::cout << ")" << std::endl;
      std::cout
          << "(define-fun lemma_" << lemma->kind() << "_c"
          << " ((x (_ BitVec 4)) (s (_ BitVec 4)) (t (_ BitVec 4))) Bool ";
      std::cout << lemma->instance(s, x, t);
      std::cout << ")" << std::endl;
    }
  }
}

#ifndef NDEBUG
void
AbstractionModule::verify_lemmas() const
{
  option::Options opts;
  NodeManager& nm = d_env.nm();
  SolvingContext ctx(nm, opts, d_env.sat_factory());

  for (uint64_t size = 4; size < 32; ++size)
  {
    std::cout << std::endl;
    std::cout << "check size=" << size << std::endl;
    Node x = nm.mk_const(nm.mk_bv_type(size), "x");
    Node s = nm.mk_const(nm.mk_bv_type(size), "s");
    Node t = nm.mk_const(nm.mk_bv_type(size), "t");
    for (const auto& [k, lemmas] : d_abstr_lemmas)
    {
      Node term = nm.mk_node(k, {x, s});
      ctx.push();
      std::cout << "check: " << k << std::endl;
      Node eq = nm.mk_node(Kind::EQUAL, {term, t});
      ctx.assert_formula(eq);
      size_t i = 0;
      for (const auto& lemma : lemmas)
      {
        std::cout << "\r" << ++i << "/" << lemmas.size() << std::flush;
        ctx.push();
        Node inst = lemma->instance(x, s, t);
        // may be null if lemma cannot be instantiated (if not applicable, e.g.,
        // for pow2 lemmas)
        if (inst.is_null())
        {
          continue;
        }
        inst = nm.mk_node(Kind::NOT, {inst});
        ctx.assert_formula(inst);
        Result res = ctx.solve();
        if (res != Result::UNSAT)
        {
          std::cout << std::endl;
          std::cout << lemma->kind() << " failed" << std::endl;
          std::cout << "(assert " << eq << ")" << std::endl;
          std::cout << "(assert " << inst << ")" << std::endl;
          std::cout << "x: " << ctx.get_value(x) << std::endl;
          std::cout << "s: " << ctx.get_value(s) << std::endl;
          std::cout << "t: " << ctx.get_value(t) << std::endl;
        }
        ctx.pop();
      }
      std::cout << std::endl;
      ctx.pop();
    }
  }
}
#endif

AbstractionModule::Statistics::Statistics(util::Statistics& stats,
                                          const std::string& prefix)
    : num_terms(stats.new_stat<uint64_t>(prefix + "terms::total")),
      num_checks(stats.new_stat<uint64_t>(prefix + "num_checks")),
      terms(stats.new_stat<util::HistogramStatistic>(prefix + "terms")),
      lemmas(stats.new_stat<util::HistogramStatistic>(prefix + "lemmas")),
      time_check(stats.new_stat<util::TimerStatistic>(prefix + "time_check"))
{
}

}  // namespace bzla::abstract
