#include "src/algorithms/privacy/numeric_encoding.h"

#include <cmath>

#include "src/algorithms/random/distributions.h"

namespace cobalt {

uint64_t DoubleToIndex(double value, double min_value, double max_value, uint64_t num_index_points,
                       BitGeneratorInterface<uint32_t>* gen) {
  double interval_size = (max_value - min_value) / static_cast<double>(num_index_points - 1);
  double approx_index = (value - min_value) / interval_size;
  double lower_index = std::floor(approx_index);
  double distance_to_index = approx_index - lower_index;

  BernoulliDistribution dist(gen, distance_to_index);
  if (dist.Sample()) {
    return static_cast<uint64_t>(lower_index) + 1;
  }

  return static_cast<uint64_t>(lower_index);
}

uint64_t IntegerToIndex(int64_t value, int64_t min_value, int64_t max_value,
                        uint64_t num_index_points, BitGeneratorInterface<uint32_t>* gen) {
  return DoubleToIndex(static_cast<double>(value), static_cast<double>(min_value),
                       static_cast<double>(max_value), num_index_points, gen);
}

uint64_t CountToIndex(uint64_t count, uint64_t max_count, uint64_t num_index_points,
                      BitGeneratorInterface<uint32_t>* gen) {
  return IntegerToIndex(static_cast<int64_t>(count), 0, static_cast<int64_t>(max_count),
                        num_index_points, gen) +
         num_index_points;
}

uint64_t HistogramBucketAndCountToIndex(uint64_t bucket_count, uint32_t bucket_index,
                                        uint64_t max_count, uint64_t num_index_points,
                                        BitGeneratorInterface<uint32_t>* gen) {
  return IntegerToIndex(static_cast<int64_t>(bucket_count), 0, static_cast<int64_t>(max_count),
                        num_index_points, gen) +
         num_index_points * static_cast<uint64_t>(bucket_index);
}

double DoubleFromIndex(uint64_t index, double min_value, double max_value,
                       uint64_t num_index_points) {
  double interval_size = (max_value - min_value) / static_cast<double>(num_index_points - 1);
  return min_value + (interval_size * static_cast<double>(index));
}

int64_t IntegerFromIndex(uint64_t index, int64_t min_value, int64_t max_value,
                         uint64_t num_index_points) {
  return static_cast<int64_t>(DoubleFromIndex(index, static_cast<double>(min_value),
                                              static_cast<double>(max_value), num_index_points));
}

uint64_t CountFromIndex(uint64_t index, uint64_t max_count, uint64_t num_index_points) {
  return static_cast<uint64_t>(IntegerFromIndex(index - num_index_points, 0,
                                                static_cast<int64_t>(max_count), num_index_points));
}

void HistogramBucketAndCountFromIndex(uint64_t index, uint64_t max_count, uint64_t num_index_points,
                                      uint64_t* bucket_count, uint32_t* bucket_index) {
  *bucket_index = static_cast<uint32_t>(index / num_index_points);
  uint64_t numeric_index = index - (*bucket_index * num_index_points);
  *bucket_count = static_cast<uint64_t>(
      IntegerFromIndex(numeric_index, 0, static_cast<int64_t>(max_count), num_index_points));
}

uint64_t ValueAndEventVectorIndicesToIndex(uint64_t value_index, uint64_t event_vector_index,
                                           uint64_t max_event_vector_index) {
  return value_index * (max_event_vector_index + 1) + event_vector_index;
}

void ValueAndEventVectorIndicesFromIndex(uint64_t index, uint64_t max_event_vector_index,
                                         uint64_t* value_index, uint64_t* event_vector_index) {
  *event_vector_index = index % (max_event_vector_index + 1);
  *value_index = (index - *event_vector_index) / (max_event_vector_index + 1);
}

bool IsCountIndex(uint64_t index, uint64_t num_index_points) { return index >= num_index_points; }

}  // namespace cobalt