#include "allocate.h"
#include "clause.h"
#include "inline.h"
#include "internal.h"
#include "literal.h"
#include "predict.h"
#include "predictparam.h"
#include "print.h"
#include "sort.h"

#include <math.h>

struct sequences {
  double *features;
  double *vcg_p;
  double *vcg_n;
  double *lit;
  double *vg;
  double *cg;
};

typedef struct sequences sequences;

#define FEATURES 128

#define LESS_SEQ_VALUE(A, B) ((A) < (B))
#define ADD_FEATURE(X) (*((*next_feature)++) = (size) ? (X) : 0)

static void extract_sequence_features (kissat *solver, double *seq, int size,
                                       double **next_feature, bool derivate) {
  if (size < 0)
    size = 0;

  SORT (double, size, seq, LESS_SEQ_VALUE);
  double *begin = seq, *end = seq + size;

  while (begin < end && *begin == 0)
    begin++;

  ADD_FEATURE (1.0 * (begin - seq) / size);

  size = end - begin;

  for (double *p = begin; p < end; p++)
    *p /= *(end - 1);

  ADD_FEATURE (*begin);
  ADD_FEATURE (*(begin + (size / 4)));
  ADD_FEATURE (*(begin + (2 * size / 4)));
  ADD_FEATURE (*(begin + (3 * size / 4)));

  double sum = 0, sum2 = 0, mode_min = 0, mode_max = 0;
  unsigned mode_min_count = 0, mode_max_count = 0, count = 0;
  double rate = 0, entropy = 0;
  for (double *p = begin; p < end; p++) {
    sum += *p;
    sum2 += (*p) * (*p);

    if (p == begin || *p != *(p - 1)) {
      rate++;
      if (count)
        entropy += count * log (count);
      count = 0;
    }
    count++;

    if (count > mode_min_count) {
      mode_min_count = count;
      mode_min = *p;
    }
    if (count >= mode_max_count) {
      mode_max_count = count;
      mode_max = *p;
    }
  }
  entropy += count * log (count);

  ADD_FEATURE (sum / size);
  ADD_FEATURE (sqrt (fabs (sum2 / size - sum / size * sum / size)));
  ADD_FEATURE (mode_min);
  ADD_FEATURE (mode_max);

  ADD_FEATURE (rate / size);
  ADD_FEATURE (size == 1 ? 0 : 1 - entropy / size / log (size));

  if (!derivate)
    return;

  for (double *p = end - 1; p > begin; p--)
    *p -= *(p - 1);
  extract_sequence_features (solver, begin + 1, end - begin - 1,
                             next_feature, false);
}

static void extract_features (kissat *solver, sequences *seq) {
  double *next_feature = seq->features;

  unsigned idx = VARS;
  for (all_clauses (c)) {
    if (c->garbage)
      continue;
    double w = 1.0;
    for (all_literals_in_clause (lit, c)) {
      (NEGATED (lit) ? seq->vcg_n : seq->vcg_p)[idx]++;
      (NEGATED (lit) ? seq->vcg_n : seq->vcg_p)[IDX (lit)]++;
      w /= 2;
    }
    for (all_literals_in_clause (lit, c))
      seq->vg[IDX (lit)] += (c->size - 1) * w;
    idx++;
  }
  for (all_literals (lit)) {
    if (lit >= LITS)
      continue;
    for (all_binary_large_watches (watch, WATCHES (lit))) {
      if (!watch.type.binary)
        continue;
      unsigned other = watch.binary.lit;
      if (lit > other)
        continue;
      if (other >= LITS)
        continue;
      (NEGATED (lit) ? seq->vcg_n : seq->vcg_p)[idx]++;
      (NEGATED (lit) ? seq->vcg_n : seq->vcg_p)[IDX (lit)]++;
      (NEGATED (other) ? seq->vcg_n : seq->vcg_p)[idx]++;
      (NEGATED (other) ? seq->vcg_n : seq->vcg_p)[IDX (other)]++;
      seq->vg[IDX (lit)] += 1.0 / 4;
      seq->vg[IDX (other)] += 1.0 / 4;
      idx++;
    }
  }

  for (unsigned i = 0; i < VARS; i++) {
    double p = seq->vcg_p[i], n = seq->vcg_n[i];
    if (p + n == 0)
      seq->lit[i] = 0;
    else
      seq->lit[i] = (p > n ? p : n) / (p + n);
  }

  unsigned idx_c = 0;
  for (all_clauses (c)) {
    if (c->garbage)
      continue;
    for (all_literals_in_clause (lit, c)) {
      double p = seq->vcg_p[IDX (lit)];
      double n = seq->vcg_n[IDX (lit)];
      seq->cg[idx_c] += p + n - 1;
    }
    idx_c++;
  }
  for (all_literals (lit)) {
    if (lit >= LITS)
      continue;
    for (all_binary_large_watches (watch, WATCHES (lit))) {
      if (!watch.type.binary)
        continue;
      unsigned other = watch.binary.lit;
      if (lit > other)
        continue;
      if (other >= LITS)
        continue;
      seq->cg[idx_c] = seq->vcg_p[IDX (lit)] + seq->vcg_n[IDX (lit)] +
                       seq->vcg_p[IDX (other)] + seq->vcg_n[IDX (other)] -
                       2;
      idx_c++;
    }
  }

  extract_sequence_features (solver, seq->vcg_p, idx, &next_feature, true);
  extract_sequence_features (solver, seq->vcg_n, idx, &next_feature, true);
  extract_sequence_features (solver, seq->lit, VARS, &next_feature, true);
  extract_sequence_features (solver, seq->vg, VARS, &next_feature, true);
  extract_sequence_features (solver, seq->cg, idx_c, &next_feature, true);

  for (double *p = seq->features; p < next_feature; p++)
    kissat_verbose (solver, "feature %.10lf", *p);
}

static int predict (kissat *solver, sequences *seq) {
  extract_features (solver, seq);
  for (unsigned i = 0; i < TOPS; i++)
    seq->features[i] = seq->features[important_features[i]];
  double weight[VOTES];
  for (unsigned i = 0; i < VOTES; i++) {
    unsigned state = 0, coef = predict_coef[i][state];
    do {
      double combined = 1.0 * (coef >> 20) / 1023;
      int feature = (coef >> 16) & 0xf;
      if (feature == 0xf) {
        weight[i] = combined;
        break;
      }
      if (seq->features[feature] <= combined)
        state = (coef >> 8) & 0xff;
      else
        state = coef & 0xff;
      coef = predict_coef[i][state];
    } while (1);
  }
  SORT (double, VOTES, weight, LESS_SEQ_VALUE);
  double median = weight[VOTES >> 1];
  kissat_verbose (solver, "prediction score %.5lf", median);
  double threshold = 0.3;
  if (median < 0.5 - threshold)
    return 10;
  if (median > 0.5 + threshold)
    return 20;
  return 0;
}

static void init_sequences (kissat *solver, sequences *seq) {
  CALLOC (seq->features, FEATURES);
  CALLOC (seq->vcg_p, VARS + CLAUSES);
  CALLOC (seq->vcg_n, VARS + CLAUSES);
  CALLOC (seq->lit, VARS);
  CALLOC (seq->vg, VARS);
  CALLOC (seq->cg, CLAUSES);
}

static void release_sequences (kissat *solver, sequences *seq) {
  DEALLOC (seq->features, FEATURES);
  DEALLOC (seq->vcg_p, VARS + CLAUSES);
  DEALLOC (seq->vcg_n, VARS + CLAUSES);
  DEALLOC (seq->lit, VARS);
  DEALLOC (seq->vg, VARS);
  DEALLOC (seq->cg, CLAUSES);
}

void kissat_predict (kissat *solver) {
  sequences seq;
  init_sequences (solver, &seq);
  int res = predict (solver, &seq);
  kissat_verbose (solver, "instance predicted as %s",
                  res == 20 ? "unsat" : (res == 10 ? "sat" : "unknown"));
  if (res == 20) {
    kissat_set_option (solver, "mab", 0);
    solver->mab = false;
  }
  release_sequences (solver, &seq);
}
