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

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

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

using testing::DoubleNear;
using testing::UnorderedElementsAre;
using testing::UnorderedElementsAreArray;

namespace cobalt {
class RapporNoiseTest : public ::testing::Test {
 protected:
  void SetUp() override { gen_ = std::make_unique<TestSecureRandomNumberGenerator>(0); }

  TestSecureRandomNumberGenerator* GetGenerator() { return gen_.get(); }

 private:
  std::unique_ptr<TestSecureRandomNumberGenerator> gen_;
};

TEST_F(RapporNoiseTest, PIsZero) {
  std::vector<uint64_t> indices = {2, 3, 5, 8, 10};
  uint64_t max_index = 10;
  double p = 0.0;
  auto noised = ApplyRapporNoise(indices, max_index, p, GetGenerator());
  EXPECT_THAT(noised, UnorderedElementsAreArray(indices));
}

TEST_F(RapporNoiseTest, PIsOne) {
  std::vector<uint64_t> indices = {2, 3, 5, 8, 10};
  uint64_t max_index = 10;
  double p = 1.0;
  std::vector<uint64_t> noised = ApplyRapporNoise(indices, max_index, p, GetGenerator());
  EXPECT_THAT(noised, UnorderedElementsAre(0ul, 1ul, 4ul, 6ul, 7ul, 9ul));
}

// Check that the total number of indices changed over 1000 trials is within 3
// standard deviations of the expected number.
TEST_F(RapporNoiseTest, NumIndicesChanged) {
  std::vector<uint64_t> indices;
  uint64_t num_trials = 1000;
  int total_indices_changed = 0;

  uint64_t max_index = 99;
  double p = 0.1;

  for (uint64_t i = 0; i < num_trials; i++) {
    total_indices_changed += ApplyRapporNoise(indices, max_index, p, GetGenerator()).size();
  }

  double stddev = std::sqrt(static_cast<double>(num_trials * (1 + max_index)) * p * (1 - p));
  double interval_radius = 3 * stddev;
  double expected_indices_changed = static_cast<double>(num_trials * (1 + max_index)) * p;
  EXPECT_THAT(static_cast<double>(total_indices_changed),
              DoubleNear(expected_indices_changed, interval_radius));
}

TEST_F(RapporNoiseTest, DebiasRapporCountPIsZero) {
  double p = 0.0;
  uint64_t raw_count = 500;
  uint64_t num_contributions = 1000;

  EXPECT_THAT(DebiasRapporCount(raw_count, num_contributions, p),
              DoubleNear(static_cast<double>(raw_count), 0.001));
}

TEST_F(RapporNoiseTest, DebiasRapporCountNegativeEstimate) {
  double p = 0.25;
  uint64_t raw_count = 500;
  uint64_t num_contributions = 10000;

  EXPECT_THAT(DebiasRapporCount(raw_count, num_contributions, p), DoubleNear(-4000, 0.001));
}

TEST_F(RapporNoiseTest, DebiasRapporCount) {
  double p = 0.1;
  uint64_t raw_count = 2000;
  uint64_t num_contributions = 10000;

  EXPECT_THAT(DebiasRapporCount(raw_count, num_contributions, p), DoubleNear(1250, 0.001));
}

}  // namespace cobalt