#include "mab_dc.h"
#include "inlineheap.h"
#include "internal.h"
#include "logging.h"

#include <math.h>

void kissat_mab_dc_init (kissat *solver) {
  solver->step_chb = 0.1 * GET_OPTION (stepchb);
  solver->heuristic = GET_OPTION (heuristic);
  for (unsigned i = 0; i < solver->mab_heuristics; i++) {
    solver->mab_reward[i] = 0;
    solver->mab_select[i] = 0;
  }
  solver->mabc = GET_OPTION (mabcint) + 0.1 * GET_OPTION (mabcdecimal);
  solver->mab_select[solver->heuristic]++;
}

void kissat_bump_chb (kissat *solver, unsigned v, double multiplier) {
  int64_t age =
      solver->statistics.conflicts - solver->conflicted_chb[v] + 1;
  double reward_chb = multiplier / age;
  double old_score = kissat_get_heap_score (&solver->scores_chb, v);
  double new_score =
      solver->step_chb * reward_chb + (1 - solver->step_chb) * old_score;
  LOG ("new CHB score[%u] = %g vs %g", v, new_score, old_score);
  kissat_update_heap (solver, &solver->scores_chb, v, new_score);
}

void kissat_decay_chb (kissat *solver) {
  if (solver->step_chb > solver->step_min_chb)
    solver->step_chb -= solver->step_dec_chb;
}

void kissat_update_conflicted_chb (kissat *solver) {
  flags *flags = solver->flags;
  for (all_stack (unsigned, idx, solver->analyzed))
    if (flags[idx].active)
      solver->conflicted_chb[idx] = solver->statistics.conflicts;
}

void kissat_restart_mab (kissat *solver) {
  unsigned stable_restarts = 0;
  if (solver->mab_decisions > 0 && solver->mab_conflicts > 0)
    solver->mab_reward[solver->heuristic] +=
        log2 (solver->mab_decisions) / log2 (solver->mab_conflicts);

  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];

  if (stable_restarts < solver->mab_heuristics) {
    solver->heuristic = solver->heuristic == 0 ? 1 : 0;
  } else {
    double ucb[2];
    solver->heuristic = 0;
    for (unsigned i = 0; i < solver->mab_heuristics; i++) {
      ucb[i] = solver->mab_reward[i] / solver->mab_select[i] +
               sqrt (solver->mabc * log (stable_restarts + 1) /
                     solver->mab_select[i]);
      if (i != 0 && ucb[i] > ucb[solver->heuristic])
        solver->heuristic = i;
    }
  }
  solver->mab_select[solver->heuristic]++;
}
