// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

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

#include <gtest/gtest.h>

#include "src/lib/statusor/status_macros.h"

namespace cobalt {

TEST(CountMin, IncrementStringAndGetCount) {
  std::vector<std::pair<std::string, uint64_t>> test_cases = {
      {"Hello World", 100},
      {"Hallo Welt", 20},
      {"Hola Mundo", 73},
      {"Bonjour, Monde", 22},
  };
  auto count_min = CountMin<uint64_t>::MakeSketch(2000, 10);
  for (auto [data, count] : test_cases) {
    count_min.Increment(data, count);
  }

  for (auto [data, count] : test_cases) {
    EXPECT_EQ(count, count_min.GetCount(data));
  }
}

TEST(CountMin, IncrementCellAndGetValue) {
  size_t num_cells_per_hash = 3;
  size_t num_hashes = 2;
  auto count_min = CountMin<uint64_t>::MakeSketch(num_cells_per_hash, num_hashes);
  ASSERT_EQ(count_min.size(), num_cells_per_hash * num_hashes);

  std::vector<uint64_t> expected_sketch;
  for (size_t i = 0; i < count_min.size(); ++i) {
    auto count = static_cast<uint64_t>(i);
    count_min.IncrementCell(i, count);
    count_min.IncrementCell(i, count);
    expected_sketch.push_back(2 * count);
  }

  for (size_t i = 0; i < count_min.size(); ++i) {
    CB_ASSERT_OK_AND_ASSIGN(uint64_t cell_value, count_min.GetCellValue(i));
    EXPECT_EQ(expected_sketch[i], cell_value);
  }
}

TEST(CountMin, IncrementCellOutOfRange) {
  auto count_min = CountMin<uint64_t>::MakeSketch(3, 2);
  EXPECT_FALSE(count_min.IncrementCell(count_min.size(), 1).ok());
}

TEST(CountMin, GetCellValueOutOfRange) {
  auto count_min = CountMin<uint64_t>::MakeSketch(3, 2);
  EXPECT_FALSE(count_min.GetCellValue(count_min.size()).ok());
}

TEST(CountMin, MakeSketchFromCells) {
  size_t num_cells_per_hash = 3;
  size_t num_hashes = 2;
  uint64_t expected_cell_value = 10;
  std::vector<uint64_t> cells(num_cells_per_hash * num_hashes, expected_cell_value);
  CB_ASSERT_OK_AND_ASSIGN(auto count_min, CountMin<uint64_t>::MakeSketchFromCells(
                                              num_cells_per_hash, num_hashes, cells));

  for (size_t i = 0; i < count_min.size(); ++i) {
    CB_ASSERT_OK_AND_ASSIGN(uint64_t cell_value, count_min.GetCellValue(i));
    EXPECT_EQ(expected_cell_value, cell_value);
  }
}

TEST(CountMin, MakeSketchFromCellsWrongDimensions) {
  size_t num_cells_per_hash = 3;
  size_t num_hashes = 2;
  uint64_t expected_cell_value = 10;
  std::vector<uint64_t> cells(num_cells_per_hash * num_hashes + 1, expected_cell_value);
  EXPECT_FALSE(CountMin<uint64_t>::MakeSketchFromCells(num_cells_per_hash, num_hashes, cells).ok());
}

TEST(CountMin, GetCellIndices) {
  size_t num_cells_per_hash = 3;
  size_t num_hashes = 2;
  auto count_min = CountMin<uint64_t>::MakeSketch(num_cells_per_hash, num_hashes);

  std::string data = "Hello World";
  count_min.Increment(data, 1);

  std::vector<size_t> updated_indices;
  for (size_t i = 0; i < count_min.size(); ++i) {
    CB_ASSERT_OK_AND_ASSIGN(uint64_t count, count_min.GetCellValue(i));
    if (count != 0) {
      updated_indices.push_back(i);
    }
  }
  ASSERT_EQ(updated_indices.size(), num_hashes);
  EXPECT_EQ(updated_indices, count_min.GetCellIndices(data));
}

}  // namespace cobalt