#include "geep.h"
#include "analyze.h"
#include "backtrack.h"
#include "decide.h"
#include "frames.h"
#include "inline.h"
#include "inlinevector.h"
#include "internal.h"
#include "print.h"
#include "propbeyond.h"
#include "proprobe.h"
#include "report.h"
#include "terminate.h"
#include "watch.h"

#include <string.h>

static void init_geep (kissat *solver) {
  const value *walked_values = solver->phases.walked;
  unsigned assigned = 0, unassigned = 0;
  CLEAR_STACK (solver->backbone);
  CLEAR_STACK (solver->partition);
  struct flags *flags = solver->flags;
  const bool equivalences = GET_OPTION (geepequiv);
  for (all_variables (idx)) {
    struct flags *f = flags + idx;
    if (!f->active)
      continue;
    unsigned lit = LIT (idx);
    const value value = walked_values[idx];
    if (value)
      assigned++;
    else
      unassigned++;
    if (value < 0)
      lit = NOT (lit);
    PUSH_STACK (solver->backbone, lit);
    if (equivalences)
      PUSH_STACK (solver->partition, lit);
    f->gackbone = value;
    assert (f->gackbone == value);
  }
  kissat_message (solver,
                  "[geep] "
                  "initially assigned %u %.0f%%, %u unassigned %.0f%%",
                  assigned, kissat_percent (assigned, solver->active),
                  unassigned, kissat_percent (unassigned, solver->active));
  (void) assigned, (void) unassigned;
  if (equivalences)
    PUSH_STACK (solver->partition, INVALID_LIT);
  kissat_phase (solver, "geep", GET (geeps),
                "initialized %zu backbone %scandidates",
                SIZE_STACK (solver->backbone),
                equivalences ? "and equivalences " : "");
}

static void refine_backbone (kissat *solver) {
  if (EMPTY_STACK (solver->backbone))
    return;
  LOG ("refining backbone");
  struct flags *flags = solver->flags;
  assert (SIZE_STACK (solver->frames) > 1);
  struct frame *frame = &FRAME (1);
  unsigned *begin = solver->trail.begin + frame->trail;
  unsigned *end = solver->propagate;
  assert (begin != end);
  assert (begin[0] == frame->decision);
#ifndef QUIET
  unsigned removed = 0;
#endif
  for (const unsigned *p = begin + 1; p != end; p++) {
    unsigned lit = *p, idx = IDX (lit);
    struct flags *f = flags + idx;
    const int gackbone = f->gackbone;
    if (!gackbone)
      continue;
    const int value = NEGATED (lit) ? -1 : 1;
    if (gackbone == value)
      continue;
    f->gackbone = 0;
#ifndef QUIET
    removed++;
#endif
  }
#ifndef QUIET
  size_t old_size = SIZE_STACK (solver->backbone);
  assert (old_size >= removed);
  size_t new_size = old_size - removed;
  if (old_size == new_size)
    LOG ("no refinement necessary");
  else
    kissat_very_verbose (
        solver,
        "refining from %zu to %zu global sweeping backbone candidates",
        old_size, new_size);
#endif
}

static void refine_partition (kissat *solver) {
  if (EMPTY_STACK (solver->partition))
    return;
  LOG ("refining partition");
  size_t old_candidates = 0, new_candidates = 0;
  size_t old_classes = 0, new_classes = 0;
  unsigneds new_partition, class0, class1, classx;
  INIT_STACK (new_partition);
  INIT_STACK (class0);
  INIT_STACK (class1);
  INIT_STACK (classx);
  const value *values = solver->values;
  size_t size_partition = SIZE_STACK (solver->partition);
  uint64_t ticks =
      1 + kissat_cache_lines (size_partition, sizeof (unsigned));
  ADD (geep_ticks, ticks);
  for (all_stack (unsigned, lit, solver->partition)) {
    if (lit == INVALID_LIT) {
      old_classes++;
      size_t size_class0 = SIZE_STACK (class0);
      size_t size_class1 = SIZE_STACK (class1);
      size_t size_classx = SIZE_STACK (classx);
      if (size_class0 > 1) {
        new_classes++;
        new_candidates += size_class0;
        for (all_stack (unsigned, lit, class0))
          PUSH_STACK (new_partition, lit);
        PUSH_STACK (new_partition, INVALID_LIT);
      }
      if (size_class1 > 1) {
        new_classes++;
        new_candidates += size_class1;
        for (all_stack (unsigned, lit, class1))
          PUSH_STACK (new_partition, lit);
        PUSH_STACK (new_partition, INVALID_LIT);
      }
      if (size_classx > 1) {
        new_classes++;
        new_candidates += size_classx;
        for (all_stack (unsigned, lit, classx))
          PUSH_STACK (new_partition, lit);
        PUSH_STACK (new_partition, INVALID_LIT);
      }
      CLEAR_STACK (class0);
      CLEAR_STACK (class1);
      CLEAR_STACK (classx);
    } else {
      old_candidates++;
      const unsigned idx = IDX (lit);
      if (ACTIVE (idx)) {
        int value = values[lit];
        if (value < 0)
          PUSH_STACK (class0, lit);
        else if (value > 0)
          PUSH_STACK (class1, lit);
        else
          PUSH_STACK (classx, lit);
      }
    }
  }
  RELEASE_STACK (class0);
  RELEASE_STACK (class1);
  RELEASE_STACK (classx);
  RELEASE_STACK (solver->partition);
  solver->partition = new_partition;
#ifndef QUIET
  if (old_candidates == new_candidates)
    LOG ("no refinement of equivalence candidates necessary");
  else
    kissat_very_verbose (
        solver,
        "refining from %zu to %zu global sweeping equivalence candidates",
        old_candidates, new_candidates);
  if (old_classes == new_classes)
    LOG ("no refinement of equivalences classes necessary");
  else
    kissat_very_verbose (
        solver,
        "refining from %zu to %zu global sweeping equivalence classes",
        old_classes, new_classes);
#endif
}

static unsigned pick_backbone_candidate (kissat *solver) {
  return POP_STACK (solver->backbone);
}

static bool pick_equivalence_candidate (kissat *solver, unsigned *first_ptr,
                                        unsigned *second_ptr) {
  unsigned first, second, end_of_class;
RESTART:
  end_of_class = POP_STACK (solver->partition);
  assert (end_of_class == INVALID_LIT), (void) end_of_class;
  first = second = INVALID_LIT;
  for (;;) {
    if (EMPTY_STACK (solver->partition))
      return false;
    first = TOP_STACK (solver->partition);
    if (first == INVALID_LIT)
      goto RESTART;
    (void) POP_STACK (solver->partition);
    if (ACTIVE (IDX (first)))
      break;
  }
  for (;;) {
    if (EMPTY_STACK (solver->partition))
      return false;
    second = TOP_STACK (solver->partition);
    if (second == INVALID_LIT)
      goto RESTART;
    (void) POP_STACK (solver->partition);
    if (first != second && ACTIVE (IDX (second)))
      break;
  }
  if (!EMPTY_STACK (solver->partition))
    PUSH_STACK (solver->partition, INVALID_LIT);
  *first_ptr = first;
  *second_ptr = second;
  return true;
}

static bool sweep_backbone_candidate (kissat *solver, unsigned candidate) {
  const unsigned candidate_idx = IDX (candidate);
  struct flags *f = solver->flags + candidate_idx;
  if (!f->active)
    return false;
  if (!f->gackbone)
    return false;
  f->gackbone = 0;
  assert (!VALUE (candidate));
  LOG ("global sweeping backbone candidate %s", LOGLIT (candidate));
  assert (!solver->level);
  assert (!solver->inconsistent);
  const unsigned not_candidate = NOT (candidate);
  INC (geep_decisions);
  kissat_internal_assume (solver, not_candidate);
  clause *conflict = kissat_probing_propagate (solver, 0, true);
  bool res;
  if (conflict) {
    LOG ("learning global sweeping backbone %s", LOGLIT (candidate));
    INC (geeped);
    INC (geep_backbones);
    kissat_backtrack_without_updating_phases (solver, 0);
    kissat_learned_unit (solver, candidate);
    conflict = kissat_probing_propagate (solver, 0, true);
    assert (!conflict || solver->inconsistent);
    (void) conflict;
    res = true;
  } else {
    if (GET_OPTION (geeprefine)) {
      refine_backbone (solver);
      if (GET_OPTION (geeprefine) > 1)
        refine_partition (solver);
    }
    kissat_backtrack_without_updating_phases (solver, 0);
    res = false;
  }
  LOG ("globally swept backbone candidate %s %s", LOGLIT (candidate),
       res ? "successfully" : "unsuccessfully");
  return res;
}

static bool sweep_equivalence_candidate (kissat *solver, unsigned first,
                                         unsigned second) {
  assert (ACTIVE (IDX (first)));
  assert (ACTIVE (IDX (second)));
  assert (!VALUE (first));
  LOG ("checking global sweeping equivalence of %s and %s", LOGLIT (first),
       LOGLIT (second));
  assert (!solver->level);
  assert (!solver->inconsistent);
  INC (geep_decisions);
  LOG ("assuming first literal %s", LOGLIT (first));
  kissat_internal_assume (solver, first);
  clause *conflict = kissat_probing_propagate (solver, 0, true);
  bool res;
  if (conflict) {
    unsigned unit = NOT (first);
    LOG ("learning global sweeping backbone %s "
         "(negation of first literal)",
         LOGLIT (unit));
    INC (geep_backbones);
    kissat_backtrack_without_updating_phases (solver, 0);
    kissat_learned_unit (solver, unit);
    conflict = kissat_probing_propagate (solver, 0, true);
    assert (!conflict || solver->inconsistent);
    (void) conflict;
    res = false;
    goto DONE;
  }
  value second_value = VALUE (second);
  if (second_value < 0) {
    LOG ("first literal %s implies second %s negated", LOGLIT (first),
         LOGLIT (second));
  FAILED:
    if (GET_OPTION (geeprefine) > 2)
      refine_partition (solver);
    kissat_backtrack_without_updating_phases (solver, 0);
    res = false;
    goto DONE;
  }
  if (second_value > 0)
    LOG ("first literal %s already implies second %s through propagation",
         LOGLIT (first), LOGLIT (second));
  else {
    assert (!second_value);
    unsigned not_second = NOT (second);
    assert (solver->level == 1);
    INC (geep_decisions);
    LOG ("assuming negation of second literal %s", LOGLIT (not_second));
    kissat_internal_assume (solver, not_second);
    conflict = kissat_probing_propagate (solver, 0, true);
    if (!conflict) {
      LOG ("could not determine that first literal %s implies second %s",
           LOGLIT (first), LOGLIT (second));
      goto FAILED;
    }
    LOG ("learned that first literal %s implies second %s", LOGLIT (first),
         LOGLIT (second));
  }
  kissat_backtrack_without_updating_phases (solver, 0);
  assert (solver->level == 0);
  INC (geep_decisions);
  LOG ("assuming second literal %s", LOGLIT (second));
  kissat_internal_assume (solver, second);
  conflict = kissat_probing_propagate (solver, 0, true);
  if (conflict) {
    unsigned unit = NOT (second);
    LOG ("learning global sweeping backbone %s"
         "(negation of second literal)",
         LOGLIT (unit));
    INC (geep_backbones);
    kissat_backtrack_without_updating_phases (solver, 0);
    kissat_learned_unit (solver, unit);
    conflict = kissat_probing_propagate (solver, 0, true);
    assert (!conflict || solver->inconsistent);
    (void) conflict;
    res = false;
    goto DONE;
  }
  value first_value = VALUE (first);
  if (first_value < 0) {
    LOG ("second literal %s implies first %s negated", LOGLIT (second),
         LOGLIT (first));
    goto FAILED;
  }
  if (first_value > 0)
    LOG ("second literal %s already implies first %s through propagation",
         LOGLIT (second), LOGLIT (first));
  else {
    assert (!first_value);
    unsigned not_first = NOT (first);
    assert (solver->level == 1);
    INC (geep_decisions);
    LOG ("assuming negation of first literal %s", LOGLIT (not_first));
    kissat_internal_assume (solver, not_first);
    conflict = kissat_probing_propagate (solver, 0, true);
    if (!conflict) {
      LOG ("could not determine that second literal %s implies first %s",
           LOGLIT (second), LOGLIT (first));
      goto FAILED;
    }
    LOG ("learned that second literal %s implies first %s", LOGLIT (second),
         LOGLIT (first));
  }
  LOG ("determined equivalence of %s and %s", LOGLIT (first),
       LOGLIT (second));
  kissat_backtrack_without_updating_phases (solver, 0);
  kissat_new_binary_clause (solver, NOT (first), second);
  kissat_new_binary_clause (solver, first, NOT (second));
  INC (geep_equivalences);
  INC (geeped);
  res = true;
DONE:
  assert (!solver->level);
  return res;
}

static void flush_backbone (kissat *solver) {
  LOG ("flushing backbone");
  unsigned *begin = BEGIN_STACK (solver->backbone);
  unsigned *end = END_STACK (solver->backbone);
  unsigned *p = begin, *q = p;
  struct flags *flags = solver->flags;
  while (p != end) {
    unsigned lit = *p++;
    unsigned idx = IDX (lit);
    struct flags *f = flags + idx;
    if (!f->active)
      continue;
    if (!f->gackbone)
      continue;
    *q++ = lit;
  }
  SET_END_OF_STACK (solver->backbone, q);
#ifndef QUIET
  size_t flushed = end - q, size = end - begin;
  kissat_phase (solver, "geep", GET (geeps),
                "flushed %zu backbone candidates %.0f%%", flushed,
                kissat_percent (flushed, size));
#endif
}

static void reset_backbone (kissat *solver, unsigned minimum) {
  LOG ("reset backbone");
  unsigned *begin = BEGIN_STACK (solver->backbone);
  unsigned *end = END_STACK (solver->backbone);
  unsigned *p = begin, *q = p;
  struct flags *flags = solver->flags;
  const value *walked_values = solver->phases.walked;
  while (p != end) {
    unsigned lit = *p++;
    unsigned idx = IDX (lit);
    struct flags *f = flags + idx;
    if (!f->active)
      continue;
    if (!f->gackbone)
      continue;
    const value old_value = NEGATED (lit) ? -1 : 1;
    const value new_value = walked_values[idx];
    assert (new_value);
    if (old_value != new_value) {
      f->gackbone = 0;
      continue;
    }
    *q++ = lit;
  }
  SET_END_OF_STACK (solver->backbone, q);
#ifndef QUIET
  size_t new_size = q - begin, old_size = end - begin;
  kissat_message (solver,
                  "[geep] reset from %zu to %zu backbone candidates "
                  "(due to new minimum %u)",
                  old_size, new_size, minimum);
#else
  (void) minimum;
#endif
}

static void reset_partition (kissat *solver, unsigned minimum) {
  LOG ("reset partition");
  size_t old_candidates = 0, new_candidates = 0;
  size_t old_classes = 0, new_classes = 0;
  unsigneds new_partition, class0, class1;
  INIT_STACK (new_partition);
  INIT_STACK (class0);
  INIT_STACK (class1);
  const value *walked_values = solver->phases.walked;
  for (all_stack (unsigned, lit, solver->partition)) {
    if (lit == INVALID_LIT) {
      old_classes++;
      size_t size_class0 = SIZE_STACK (class0);
      size_t size_class1 = SIZE_STACK (class1);
      if (size_class0 > 1) {
        new_classes++;
        new_candidates += size_class0;
        for (all_stack (unsigned, lit, class0))
          PUSH_STACK (new_partition, lit);
        PUSH_STACK (new_partition, INVALID_LIT);
      }
      if (size_class1 > 1) {
        new_classes++;
        new_candidates += size_class1;
        for (all_stack (unsigned, lit, class1))
          PUSH_STACK (new_partition, lit);
        PUSH_STACK (new_partition, INVALID_LIT);
      }
      CLEAR_STACK (class0);
      CLEAR_STACK (class1);
    } else {
      old_candidates++;
      const unsigned idx = IDX (lit);
      if (ACTIVE (idx)) {
        int walked_value = walked_values[idx];
        if (NEGATED (lit))
          walked_value = -walked_value;
        assert (walked_value);
        if (walked_value < 0)
          PUSH_STACK (class0, lit);
        else
          PUSH_STACK (class1, lit);
      }
    }
  }
  assert (EMPTY_STACK (class0));
  assert (EMPTY_STACK (class1));
  RELEASE_STACK (class0);
  RELEASE_STACK (class1);
  RELEASE_STACK (solver->partition);
  solver->partition = new_partition;
#ifndef QUIET
  kissat_message (solver,
                  "[geep] reset from %zu to %zu equivalence candidates "
                  "(due to new minimum %u)",
                  old_candidates, new_candidates, minimum);
  kissat_message (solver,
                  "[geep] reset from %zu to %zu equivalence classes "
                  "(due to new minimum %u)",
                  old_classes, new_classes, minimum);
#else
  (void) minimum;
#endif
}

void kissat_reset_global_sweeping (kissat *solver, unsigned minimum) {
  if (EMPTY_STACK (solver->backbone))
    LOG ("no backbone candidates left to reset");
  else
    reset_backbone (solver, minimum);
  if (GET_OPTION (geepequiv)) {
    if (EMPTY_STACK (solver->partition))
      LOG ("no equivalence left to reset");
    else
      reset_partition (solver, minimum);
  }
}

static void geep_backbones (kissat *solver, uint64_t limit,
                            size_t *backbones_ptr) {

  assert (!EMPTY_STACK (solver->backbone));
  size_t size_candidates = SIZE_STACK (solver->backbone);
  kissat_phase (solver, "geep", GET (geeps),
                "sweeping %zu backbone candidates %.0f%%", size_candidates,
                kissat_percent (size_candidates, solver->active));
  (void) size_candidates;

  size_t backbones = 0;

  while (!solver->inconsistent) {
    if (solver->statistics.geep_ticks > limit)
      break;
    if (TERMINATED (geep_terminated_1))
      break;
    if (EMPTY_STACK (solver->backbone))
      break;
    unsigned lit = pick_backbone_candidate (solver);
    if (sweep_backbone_candidate (solver, lit))
      backbones++;
  }
  if (!EMPTY_STACK (solver->backbone))
    flush_backbone (solver);

#ifndef QUIET
  kissat_phase (solver, "geep", GET (geeps), "derived %zu backbones",
                backbones);
  size_t remaining_backbone_candidates = SIZE_STACK (solver->backbone);
  if (remaining_backbone_candidates)
    kissat_phase (
        solver, "geep", GET (geeps),
        "%zu backbone candidates remain "
        "(%.0f%% out of %zu)",
        remaining_backbone_candidates,
        kissat_percent (remaining_backbone_candidates, size_candidates),
        size_candidates);
  else
    kissat_phase (solver, "geep", GET (geeps),
                  "all %zu backbone candidates tried", size_candidates);
#endif

  *backbones_ptr = backbones;
}

static void geep_equivalences (kissat *solver, uint64_t limit,
                               size_t *equivalences_ptr) {
  size_t equivalences = 0;
  while (!solver->inconsistent) {
    if (solver->statistics.geep_ticks > limit)
      break;
    if (TERMINATED (geep_terminated_2))
      break;
    if (EMPTY_STACK (solver->partition))
      break;
    unsigned first, second;
    if (pick_equivalence_candidate (solver, &first, &second) &&
        sweep_equivalence_candidate (solver, first, second))
      equivalences++;
  }
  kissat_phase (solver, "geep", GET (geeps), "derived %zu equivalences",
                equivalences);
  if (EMPTY_STACK (solver->partition))
    kissat_phase (solver, "geep", GET (geeps),
                  "all equivalence candidates tried");

  *equivalences_ptr = equivalences;
}

bool kissat_global_sweep (kissat *solver) {
  if (solver->inconsistent)
    return false;
  assert (solver->probing);
  assert (solver->watching);
  assert (!solver->level);
  if (!GET_OPTION (geep))
    return false;
  START (geep);
  INC (geeps);
  assert (!solver->geeping);
  solver->geeping = true;
  if (EMPTY_STACK (solver->backbone) && EMPTY_STACK (solver->partition))
    init_geep (solver);
  SET_EFFORT_LIMIT (limit, geep, geep_ticks);
#ifndef QUIET
  uint64_t budget = limit - solver->statistics.geep_ticks;
#endif
  size_t backbones = 0, equivalences = 0;
  if (EMPTY_STACK (solver->backbone))
    kissat_phase (solver, "geep", GET (geeps), "no backbone candidates");
  else
    geep_backbones (solver, limit, &backbones);
  if (!solver->inconsistent && solver->statistics.geep_ticks <= limit &&
      !TERMINATED (geep_terminated_3) && !EMPTY_STACK (solver->partition))
    geep_equivalences (solver, limit, &equivalences);
#ifndef QUIET
  if (solver->statistics.geep_ticks < limit) {
    uint64_t delta = limit - solver->statistics.geep_ticks;
    kissat_phase (solver, "geep", GET (geeps),
                  "sweeping has %" PRIu64 " ticks left %.2f%%", delta,
                  kissat_percent (delta, budget));
  } else {
    uint64_t delta = solver->statistics.geep_ticks - limit;
    kissat_phase (solver, "geep", GET (geeps),
                  "limit exceeded by %" PRIu64 " ticks %.2f%%", delta,
                  kissat_percent (delta, budget));
  }
#endif
  assert (solver->geeping);
  solver->geeping = false;
  bool success = backbones || equivalences;
  REPORT (!success, 'g');
  STOP (geep);
  return success;
}
