#include "restart.h"
#include "backtrack.h"
#include "bump.h"
#include "decide.h"
#include "internal.h"
#include "kimits.h"
#include "logging.h"
#include "print.h"
#include "reluctant.h"
#include "report.h"
#include "factorl.h"

#include <inttypes.h>
#include <math.h>

bool kissat_restarting (kissat *solver) {
  assert (solver->unassigned);
  if (!GET_OPTION (restart))
    return false;
  if (!solver->level)
    return false;
  if (CONFLICTS < solver->limits.restart.conflicts)
    return false;
  if (solver->stable)
    return kissat_reluctant_triggered (&solver->reluctant);
  const double fast = AVERAGE (fast_glue);
  const double slow = AVERAGE (slow_glue);
  const double margin = (100.0 + GET_OPTION (restartmargin)) / 100.0;
  const double limit = margin * slow;
  kissat_extremely_verbose (solver,
                            "restart glue limit %g = "
                            "%.02f * %g (slow glue) %c %g (fast glue)",
                            limit, margin, slow,
                            (limit > fast    ? '>'
                             : limit == fast ? '='
                                             : '<'),
                            fast);
  return (limit <= fast);
}

void kissat_update_focused_restart_limit (kissat *solver) {
  assert (!solver->stable);
  limits *limits = &solver->limits;
  uint64_t restarts = solver->statistics.restarts;
  uint64_t delta = GET_OPTION (restartint);
  if (restarts)
    delta += kissat_logn (restarts) - 1;
  limits->restart.conflicts = CONFLICTS + delta;
  kissat_extremely_verbose (solver,
                            "focused restart limit at %" PRIu64
                            " after %" PRIu64 " conflicts ",
                            limits->restart.conflicts, delta);
}

static unsigned reuse_stable_trail (kissat *solver) {
  const heap *const scores = kissat_get_scores(solver);
  const unsigned next_idx = kissat_next_decision_variable (solver);
  const double limit = kissat_get_heap_score (scores, next_idx);
  unsigned level = solver->level, res = 0;
  while (res < level) {
    frame *f = &FRAME (res + 1);
    const unsigned idx = IDX (f->decision);
    const double score = kissat_get_heap_score (scores, idx);
    if (score <= limit)
      break;
    res++;
  }
  return res;
}

static unsigned reuse_focused_trail (kissat *solver) {
  const links *const links = solver->links;
  const unsigned next_idx = kissat_next_decision_variable (solver);
  const unsigned limit = links[next_idx].stamp;
  LOG ("next decision variable stamp %u", limit);
  unsigned level = solver->level, res = 0;
  while (res < level) {
    frame *f = &FRAME (res + 1);
    const unsigned idx = IDX (f->decision);
    const unsigned score = links[idx].stamp;
    if (score <= limit)
      break;
    res++;
  }
  return res;
}

static unsigned reuse_trail (kissat *solver) {
  assert (solver->level);
  assert (!EMPTY_STACK (solver->trail));

  if (!GET_OPTION (restartreusetrail))
    return 0;

  unsigned res;

  if (solver->stable)
    res = reuse_stable_trail (solver);
  else
    res = reuse_focused_trail (solver);

  LOG ("matching trail level %u", res);

  if (res) {
    INC (restarts_reused_trails);
    ADD (restarts_reused_levels, res);
    LOG ("restart reuses trail at decision level %u", res);
  } else
    LOG ("restarts does not reuse the trail");

  return res;
}

static double safe_mab_average (double reward, unsigned selected) {
  return selected ? reward / selected : 0.0;
}

static double safe_mab_reward (kissat *solver) {
  if (solver->mab_decisions <= 1.0)
    return 0.0;
  if (solver->mab_conflicts <= 1)
    return 0.0;
  const double decisions = log2 (solver->mab_decisions);
  const double conflicts = log2 ((double) solver->mab_conflicts);
  if (decisions <= 0.0 || conflicts <= 0.0)
    return 0.0;
  return decisions / conflicts;
}

static void restart_mab (kissat *solver) {
  unsigned stable_restarts = 0;
  solver->mab_reward[solver->heuristic] += safe_mab_reward (solver);

  for (all_variables (idx))
    solver->mab_chosen[idx] = 0;
  solver->mab_chosen_tot = 0;
  solver->mab_decisions = 0;
  solver->mab_conflicts = 0;

  for (unsigned i = 0; i < solver->mab_heuristics; i++)
    stable_restarts += solver->mab_select[i];

  static double recent_gains[10];
  static int gain_index;
  static double momentum = 1.0;

  const double current_gain = safe_mab_average (
      solver->mab_reward[solver->heuristic],
      solver->mab_select[solver->heuristic]);
  recent_gains[gain_index] = current_gain;
  gain_index = (gain_index + 1) % 10;

  double avg_gain = 0.0;
  for (int i = 0; i < 10; i++)
    avg_gain += recent_gains[i];
  avg_gain /= 10.0;

  momentum *= (current_gain > avg_gain) ? 1.1 : 0.9;
  if (momentum <= 0.0)
    momentum = 1.0;

  const double adaptive_c = solver->mabc / (momentum * (stable_restarts + 1));

  if (stable_restarts < solver->mab_heuristics) {
    solver->heuristic = solver->heuristic == 0 ? 1 : 0;
  } else {
    double best_ucb = -1.0;
    solver->heuristic = 0;
    for (unsigned i = 0; i < solver->mab_heuristics; i++) {
      const unsigned selected = solver->mab_select[i];
      if (!selected) {
        solver->heuristic = i;
        break;
      }
      const double avg = safe_mab_average (solver->mab_reward[i], selected);
      const double bonus =
          sqrt (adaptive_c * log ((double) stable_restarts + 1.0) / selected);
      const double ucb = avg + bonus;
      if (i && ucb <= best_ucb)
        continue;
      best_ucb = ucb;
      solver->heuristic = i;
    }
  }

  solver->mab_select[solver->heuristic]++;
}

int kissat_restart (kissat *solver) {
  START (restart);
  INC (restarts);
  ADD (restarts_levels, solver->level);
  if (solver->stable)
    INC (stable_restarts);
  else
    INC (focused_restarts);

  unsigned old_heuristic = solver->heuristic;
  if (solver->stable && solver->mab) 
      restart_mab(solver);
  unsigned new_heuristic = solver->heuristic;

  unsigned level = old_heuristic==new_heuristic?reuse_trail (solver):0;

  kissat_extremely_verbose (solver,
                            "restarting after %" PRIu64 " conflicts"
                            " (limit %" PRIu64 ")",
                            CONFLICTS, solver->limits.restart.conflicts);
  LOG ("restarting to level %u", level);
  if (solver->stable && solver->mab) solver->heuristic = old_heuristic;
  kissat_backtrack_in_consistent_state (solver, level);
  if (solver->stable && solver->mab) solver->heuristic = new_heuristic;
  kissat_factor_learnt (solver);

  if (!solver->inconsistent && !solver->stable)
    kissat_update_focused_restart_limit (solver);
  if (!solver->inconsistent && solver->stable && solver->mab &&
      old_heuristic != new_heuristic)
    kissat_update_scores (solver);

  REPORT (1, 'R');
  STOP (restart);

  return solver->inconsistent ? 20 : 0;
}
