formula.hpp

#pragma once

extern "C" {
#include "ipasir.h"
}

#include "problem.hpp"
#include "tools/Logger.h"
#include <limits>

// lol
#define inf_time std::numeric_limits<double>::max() * 0.5

#define UNSAT -3;

// this capsules the ipasir interface
// clauses added will be added to all past and future steps
class Formula {
private:
  void *solver;
  int num_base_variables     = 0;
  unsigned new_clauses_begin = 0;
  int current_step           = -1;
  // makestep x variable x value
  unsigned initial_clauses_index = 0;

  std::vector<int> clauses;
  std::vector<std::pair<unsigned, unsigned>> assumptions;

  void activate_assumptions() {
    for (auto [index, value] : assumptions) {
      ipasir_assume(solver, state.back()[index][value]);
    }
  }

  void add_clauses_for_new_step() {
    for (size_t i = 0; i < clauses.size(); ++i) {
      if (clauses[i] > 0) {
        clauses[i] += num_base_variables;
      } else if (clauses[i] < 0) {
        clauses[i] -= num_base_variables;
      }
      // 0 -> 0
      ipasir_add(solver, clauses[i]);
    }
  }

  // only new clauses
  // known clauses are added in increase makespan
  void add_clauses_for_all_steps() {
    if (new_clauses_begin == clauses.size()) {
      return;
    }
    std::vector<int> new_clauses(clauses.begin() + new_clauses_begin,
                                 clauses.end());
    new_clauses_begin = clauses.size();
    assert(new_clauses.back() == 0);
    for (int step = 0; step < current_step + 1; ++step) {
      for (unsigned i = 0; i < new_clauses.size(); ++i) {
        ipasir_add(solver, new_clauses[i]);
        if (new_clauses[i] > 0) {
          new_clauses[i] -= num_base_variables;
        } else if (new_clauses[i] < 0) {
          new_clauses[i] += num_base_variables;
        }
        // 0 -> 0
      }
    }
  }

public:
  std::vector<std::vector<std::vector<int>>> state;
  std::vector<std::vector<unsigned>> action;

  Formula(const Problem &problem,
          const std::vector<int> &additional_clauses = {}) {
    solver = ipasir_init();

    // glucose fix
    // ipasir_set_learn(solver, NULL, 0, NULL);

    std::cout << "Log: " << getTime() << " "
              << "Using the incremental SAT solver" << ipasir_signature()
              << std::endl;

    // set state variables
    state.emplace_back(problem.vars.size());
    for (size_t i = 0; i < state[0].size(); ++i) {
      state[0][i].resize(problem.vars[i].values.size());
      if (state[0][i].size() == 2) {
        // no dual-rail encoding for binary values
        ++num_base_variables;
        state[0][i] = {num_base_variables, -num_base_variables};
      } else {
        for (size_t v = 0; v < problem.vars[i].values.size(); ++v) {
          state[0][i][v] = ++num_base_variables;
        }
      }
    }

    // set action variables
    action.emplace_back(problem.actions.size());
    for (size_t i = 0; i < action[0].size(); ++i) {
      action[0][i] = ++num_base_variables;
    }

    // TODO the last set of action vars is not needed
    increase_makespan();

    // at least one value
    // is not necessary for the fixed initial sate
    for (size_t i = 0; i < state[0].size(); ++i) {
      if (state[0][i].size() > 2) {
        for (size_t v = 0; v < state[0][i].size(); ++v) {
          add_s(i, v, true, true);
        }
        close();
      }
    }

    // mutex which are generated by sas transform of fastdownward
    for (const auto &mutex : problem.mutexes) {
      for (size_t i = 0; i < mutex.size() - 1; ++i) {
        for (size_t j = i + 1; j < mutex.size(); ++j) {
          add_s(mutex[i].first, mutex[i].second, false, true);
          add_s(mutex[j].first, mutex[j].second, false, true);
          close();
        }
      }
    }
    initial_clauses_index = clauses.size();
    clauses.insert(clauses.begin(), additional_clauses.begin(),
                   additional_clauses.end());
  }

  ~Formula() { ipasir_release(solver); }

  unsigned increase_makespan(unsigned steps = 1) {
    for (unsigned i = 0; i < steps; ++i) {
      // set state variables
      state.push_back(state.back());
      for (size_t i = 0; i < state.back().size(); ++i) {
        for (size_t v = 0; v < state.back()[i].size(); ++v) {
          if (state.back()[i][v] > 0) {
            state.back()[i][v] += num_base_variables;
          } else {
            state.back()[i][v] -= num_base_variables;
          }
        }
      }

      // set action variables
      action.push_back(action.back());
      for (size_t i = 0; i < action.back().size(); ++i) {
        action.back()[i] += num_base_variables;
      }

      add_clauses_for_new_step();

      current_step++;
    }

    return current_step + 1;
  }

  // assuming the goal
  inline void assume_at_last(int variable, unsigned value) {
    assumptions.emplace_back(variable, value);
  }

  inline void add_initial_s_atom(int variable, unsigned value,
                                 bool polarity = true) {
    ipasir_add(solver, (polarity ? 1 : -1) * state[0][variable][value]);
    ipasir_add(solver, 0);
  }

  // clauses added will be added to all past and future steps
  inline void add_s(int variable, unsigned value, bool polarity = true,
                    bool next = false) {
    if (next) {
      clauses.push_back((polarity ? 1 : -1) *
                        state[current_step + 1][variable][value]);
    } else {
      clauses.push_back((polarity ? 1 : -1) *
                        state[current_step][variable][value]);
    }
  }

  inline void add_a(int index, bool polarity = true, bool next = false) {
    if (next) {
      clauses.push_back((polarity ? 1 : -1) * action[current_step + 1][index]);
    } else {
      clauses.push_back((polarity ? 1 : -1) * action[current_step][index]);
    }
  }

  inline void close() { clauses.push_back(0); }

  // assume goal and try to solve
  bool solve(double time_limit = inf_time) {
    add_clauses_for_all_steps();
    activate_assumptions();
    double end_time = getTime() + time_limit;
    ipasir_set_terminate(solver, &end_time, [](void *end_time) {
      return (int)(getTime() > *(double *)(end_time));
    });
    int satRes = ipasir_solve(solver);
    return satRes == 10;
  }

  int get_state_value(unsigned variable, unsigned t) const {
    for (unsigned i = 0; i < state[t][variable].size(); ++i) {
      int lit = state[t][variable][i];
      if (lit > 0) {
        // TODO not important values
        if (ipasir_val(solver, lit) >= 0) {
          return i;
        }
      } else {
        if (ipasir_val(solver, -lit) <= 0) {
          return i;
        }
      }
    }
    assert(false);
    return -1;
  }

  bool get_action_value(unsigned index, unsigned t) const {
    // not important variables should be false, a minimum of actions is
    // preferred
    return ipasir_val(solver, action[t][index]) > 0;
  }

  void get_clauses(std::vector<int> &c) {
    c = std::vector<int>(clauses.begin() + initial_clauses_index,
                         clauses.end());
  }
};