#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 <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;
}

void restart_mab(kissat *solver) {
    static bool initialized = false;
    static double alphas[3][4]; 
    static double betas[3][4];
    static int current_state = 2; 
    static int current_arm = 0;
    static uint64_t next_quantum_limit = 0;
    static double quantum_start_ema_lbd = 0.0;
    static double prev_quantum_trail = 0.0;
    static double mid_quantum_ema_lbd = 0.0;
    static bool mid_quantum_recorded = false;

    if (!initialized) {
        // Step 3: Initialize Beta distributions with domain-informed optimistic priors
        for (int s = 0; s < 3; s++) {
            for (int a = 0; a < 4; a++) {
                alphas[s][a] = 1.0;
                betas[s][a] = 1.0;
            }
        }
        // Converging (State 0) -> Arm 2 (Conservative)
        alphas[0][2] = 4.0; betas[0][2] = 1.0;
        // Chaotic (State 1) -> Arm 0 (Hyper-aggressive)
        alphas[1][0] = 4.0; betas[1][0] = 1.0;
        // Stagnant (State 2) -> Arm 3 (Reluctant)
        alphas[2][3] = 4.0; betas[2][3] = 1.0;
        
        next_quantum_limit = CONFLICTS + 1000;
        quantum_start_ema_lbd = AVERAGE(slow_glue);
#ifndef QUIET
        prev_quantum_trail = AVERAGE(trail);
#else
        prev_quantum_trail = (double)(solver->vars - solver->unassigned);
#endif
        initialized = true;
    }

    // Record the mid-quantum EMA LBD to calculate trajectory over the last 500 conflicts
    if (CONFLICTS >= next_quantum_limit - 500 && !mid_quantum_recorded) {
        mid_quantum_ema_lbd = AVERAGE(slow_glue);
        mid_quantum_recorded = true;
    }

    // At the start of a 1000-conflict quantum
    if (CONFLICTS >= next_quantum_limit) {
        double current_ema_lbd = AVERAGE(slow_glue);
        
        if (!mid_quantum_recorded) {
            mid_quantum_ema_lbd = current_ema_lbd;
        }

        // Step 5: Compute Delayed Reward
        double R = 0.0;
        if (quantum_start_ema_lbd > 0.0) {
            R = (quantum_start_ema_lbd - current_ema_lbd) / quantum_start_ema_lbd;
        }
        
        if (R > 0.05) {
            alphas[current_state][current_arm] += 1.0; // Success
        } else {
            betas[current_state][current_arm] += 1.0;  // Failure
        }

        // Step 6: Volatility Reset mechanism
#ifndef QUIET
        double current_trail = AVERAGE(trail);
#else
        double current_trail = (double)(solver->vars - solver->unassigned);
#endif
        if (prev_quantum_trail > 0.0 && current_trail < 0.75 * prev_quantum_trail) {
            for (int s = 0; s < 3; s++) {
                for (int a = 0; a < 4; a++) {
                    alphas[s][a] *= 0.5;
                    if (alphas[s][a] < 1.0) alphas[s][a] = 1.0;
                    betas[s][a] *= 0.5;
                    if (betas[s][a] < 1.0) betas[s][a] = 1.0;
                }
            }
        }
        prev_quantum_trail = current_trail;

        // Step 2: Observe current State based on LBD trajectory
        double slope = current_ema_lbd - mid_quantum_ema_lbd;
        double diff = AVERAGE(fast_glue) - AVERAGE(slow_glue);
        double variance_est = diff * diff; // Approximation of LBD variance

        if (slope < -0.1) {
            current_state = 0; // Converging
        } else if (variance_est > 15.0) {
            current_state = 1; // Chaotic
        } else {
            current_state = 2; // Stagnant
        }

        // Step 3: Sample a value from Beta distribution of each Arm for the current State
        int best_arm = 0;
        double best_sample = -1.0;
        for (int a = 0; a < 4; a++) {
            double alpha = alphas[current_state][a];
            double beta = betas[current_state][a];
            
            // Box-Muller Normal approximation for Beta sampling
            double mean = alpha / (alpha + beta);
            double var = (alpha * beta) / ((alpha + beta) * (alpha + beta) * (alpha + beta + 1.0));
            double stddev = sqrt(var);
            
            double u1 = kissat_pick_double(&solver->random);
            double u2 = kissat_pick_double(&solver->random);
            if (u1 < 1e-15) u1 = 1e-15;
            double z = sqrt(-2.0 * log(u1)) * cos(2.0 * 3.14159265358979323846 * u2);
            
            double sample = mean + z * stddev;
            if (sample < 0.0) sample = 0.0;
            if (sample > 1.0) sample = 1.0;
            
            if (sample > best_sample) {
                best_sample = sample;
                best_arm = a;
            }
        }
        current_arm = best_arm; // Select Arm with highest sampled value

        // Setup for the next quantum
        next_quantum_limit = CONFLICTS + 1000;
        quantum_start_ema_lbd = current_ema_lbd;
        mid_quantum_recorded = false;
        mid_quantum_ema_lbd = current_ema_lbd;
    }

    // Step 4: Apply the selected Arm's restart policy
    // Modifying the EMA alpha directly to simulate the distinct policy profiles
    if (current_arm == 0) {
        // Arm 0: Hyper-aggressive EMA
        EMA(slow_glue).alpha = 1e-4;
    } else if (current_arm == 1) {
        // Arm 1: Standard EMA
        EMA(slow_glue).alpha = 1e-5;
    } else if (current_arm == 2) {
        // Arm 2: Conservative EMA
        EMA(slow_glue).alpha = 1e-5;
    } else if (current_arm == 3) {
        // Arm 3: Reluctant/Luby sequence baseline fallback
        EMA(slow_glue).alpha = 1e-5;
    }

    // Ensure the baseline heuristic selection mechanism still explores
    solver->heuristic = (current_arm % 2);
}

void 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;
  if (!solver->stable)
    kissat_update_focused_restart_limit (solver);
  
  if (solver->stable && solver->mab && old_heuristic!=new_heuristic) kissat_update_scores(solver);

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