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

#include <gmock/gmock.h>
#include <gtest/gtest.h>

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

using ::testing::AnyOf;
using ::testing::DoubleEq;
using ::testing::Eq;

namespace cobalt {

TEST(NumericEncodingTest, DoubleToIndexEqualsMaxValue) {
  double min_value = -4.0;
  double max_value = 6.0;
  uint64_t num_index_points = 6;
  double value = 6.0;
  uint64_t expected_index = 5u;
  RandomNumberGenerator gen;

  EXPECT_EQ(expected_index, DoubleToIndex(value, min_value, max_value, num_index_points, &gen));
}

TEST(NumericEncodingTest, DoubleToIndexEqualsMinValue) {
  double min_value = -4.0;
  double max_value = 6.0;
  uint64_t num_index_points = 6;
  double value = -4.0;
  uint64_t expected_index = 0u;
  RandomNumberGenerator gen;

  EXPECT_EQ(expected_index, DoubleToIndex(value, min_value, max_value, num_index_points, &gen));
}

TEST(NumericEncodingTest, DoubleToIndexEqualsInteriorIndexPoint) {
  double min_value = -4.0;
  double max_value = 6.0;
  uint64_t num_index_points = 6;
  double value = -2.0;
  uint64_t expected_index = 1u;
  RandomNumberGenerator gen;

  EXPECT_EQ(expected_index, DoubleToIndex(value, min_value, max_value, num_index_points, &gen));
}

TEST(NumericEncodingTest, DoubleToIndexBetweenTwoIndicesNoSeed) {
  double min_value = -4.0;
  double max_value = 6.0;
  uint64_t num_index_points = 6;
  double value = -1.0;
  uint64_t lower_index = 1u;
  uint64_t upper_index = 2u;
  RandomNumberGenerator gen;

  for (int i = 0; i < 1000; i++) {
    EXPECT_THAT(DoubleToIndex(value, min_value, max_value, num_index_points, &gen),
                AnyOf(Eq(lower_index), Eq(upper_index)));
  }
}

TEST(NumericEncodingTest, IntegerToIndexBetweenTwoIndicesNoSeed) {
  int64_t min_value = -4;
  int64_t max_value = 6;
  uint64_t num_index_points = 6;
  int64_t value = -1;
  uint64_t lower_index = 1u;
  uint64_t upper_index = 2u;
  RandomNumberGenerator gen;

  for (int i = 0; i < 1000; i++) {
    EXPECT_THAT(IntegerToIndex(value, min_value, max_value, num_index_points, &gen),
                AnyOf(Eq(lower_index), Eq(upper_index)));
  }
}

TEST(NumericEncodingTest, CountToIndex) {
  uint64_t max_count = 10u;
  uint64_t num_index_points = 6u;
  uint64_t count = 4u;
  RandomNumberGenerator gen;

  uint64_t expected = 8u;

  EXPECT_EQ(CountToIndex(count, max_count, num_index_points, &gen), expected);
}

TEST(NumericEncodingTest, HistogramBucketToIndex) {
  uint64_t max_count = 10u;
  uint64_t num_index_points = 6u;
  uint64_t bucket_count = 4u;
  uint32_t bucket_index = 7u;
  RandomNumberGenerator gen;

  uint64_t expected = 44u;

  EXPECT_EQ(HistogramBucketToIndex(bucket_count, bucket_index, max_count, num_index_points, &gen),
            expected);
}

TEST(NumericEncodingTest, DoubleFromIndex) {
  double min_value = -4.0;
  double max_value = 6.0;
  uint64_t num_index_points = 6u;

  double expected = min_value;
  for (uint64_t index = 0u; index <= num_index_points; index++) {
    EXPECT_THAT(DoubleFromIndex(index, min_value, max_value, num_index_points), DoubleEq(expected));
    expected += 2.0;
  }
}

TEST(NumericEncodingTest, IntegerFromIndex) {
  int64_t min_value = -4;
  int64_t max_value = 6;
  uint64_t num_index_points = 6u;

  int64_t expected = min_value;
  for (uint64_t index = 0u; index <= num_index_points; index++) {
    EXPECT_EQ(IntegerFromIndex(index, min_value, max_value, num_index_points), expected);
    expected += 2;
  }
}

TEST(NumericEncodingTest, CountFromIndex) {
  uint64_t max_count = 10u;
  uint64_t num_index_points = 6u;

  uint64_t expected = 0u;

  for (uint64_t index = num_index_points; index < 2 * num_index_points; index++) {
    EXPECT_EQ(CountFromIndex(index, max_count, num_index_points), expected);
    expected += 2;
  }
}

TEST(NumericEncodingTest, HistogramBucketFromIndex) {
  uint64_t max_count = 10u;
  uint64_t num_index_points = 6u;
  uint64_t num_buckets = 10u;

  uint64_t expected_bucket_count = 0u;
  uint64_t expected_bucket_index = 0u;

  for (uint64_t index = 0u; index < num_index_points * num_buckets; ++index) {
    uint64_t bucket_count;
    uint32_t bucket_index;

    HistogramBucketFromIndex(index, max_count, num_index_points, &bucket_count, &bucket_index);
    EXPECT_EQ(bucket_count, expected_bucket_count);
    EXPECT_EQ(bucket_index, expected_bucket_index);

    expected_bucket_count += 2;
    if (expected_bucket_count > max_count) {
      expected_bucket_count = 0;
      ++expected_bucket_index;
    }
  }
}

TEST(NumericEncodingTest, ValueAndEventVectorIndicesToIndex) {
  uint64_t num_index_points = 6;
  uint64_t max_event_vector_index = 9;

  uint64_t expected_index = 0;
  for (uint64_t value_index_in = 0; value_index_in < num_index_points; ++value_index_in) {
    for (uint64_t ev_index_in = 0; ev_index_in <= max_event_vector_index; ++ev_index_in) {
      uint64_t index =
          ValueAndEventVectorIndicesToIndex(value_index_in, ev_index_in, max_event_vector_index);
      EXPECT_EQ(index, expected_index);

      uint64_t value_index_out;
      uint64_t ev_index_out;
      ValueAndEventVectorIndicesFromIndex(index, max_event_vector_index, &value_index_out,
                                          &ev_index_out);
      EXPECT_EQ(value_index_in, value_index_out);
      EXPECT_EQ(ev_index_in, ev_index_out);
      ++expected_index;
    }
  }
}

TEST(NumericEncodingTest, IsCountIndex) {
  uint64_t num_index_points = 5;
  EXPECT_TRUE(IsCountIndex(7u, num_index_points));
  EXPECT_FALSE(IsCountIndex(3u, num_index_points));
}

}  // namespace cobalt