// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/variations/limited_layer_entropy_cost_tracker.h"

#include <cstdint>
#include <limits>
#include <optional>
#include <string_view>
#include <utility>
#include <vector>

#include "base/strings/strcat.h"
#include "base/strings/string_number_conversions.h"
#include "base/time/time.h"
#include "components/variations/entropy_provider.h"
#include "components/variations/proto/layer.pb.h"
#include "components/variations/proto/study.pb.h"
#include "components/variations/proto/variations_seed.pb.h"
#include "components/variations/variations_layers.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"

using ::testing::UnorderedElementsAreArray;

namespace variations {

namespace {

inline constexpr uint32_t kTestLayerId = 1001;
inline constexpr int kTestLayerMemberId = 2001;
inline constexpr std::string_view kTestClientId = "test_client_id";

constexpr int64_t kTimeMin = 0;

// The following value ensures slot 0 is selected  (among 100 slots) when the
// limited entropy provider is used.
inline constexpr std::string_view kTestLimitedEntropyRandomizationSource =
    "limited_entropy_randomization_source_964";

int64_t Timestamp(base::Time time) {
  return static_cast<int64_t>(time.InSecondsFSinceUnixEpoch());
}

Study::Experiment CreateExperiment(int weight) {
  Study::Experiment experiment;
  experiment.set_probability_weight(weight);
  return experiment;
}

Study::Experiment CreateGoogleWebExperiment(int weight,
                                            int google_web_experiment_id) {
  auto experiment = CreateExperiment(weight);
  experiment.set_google_web_experiment_id(google_web_experiment_id);
  return experiment;
}

Study::Experiment CreateTriggerExperiment(
    int weight,
    int google_web_trigger_experiment_id) {
  auto experiment = CreateExperiment(weight);
  experiment.set_google_web_trigger_experiment_id(
      google_web_trigger_experiment_id);
  return experiment;
}

std::vector<Study::Experiment> CreateExperimentsWithTwoBitsOfEntropy() {
  // Create 3 experiments with a total probability weight of 100. Only the 25%
  // group has a google_web_experiment_id so the entropy used should be
  // -log2(1/4) = 2.
  std::vector<Study::Experiment> experiments = {
      CreateExperiment(10), CreateGoogleWebExperiment(25, 100001),
      CreateExperiment(65)};
  return experiments;
}

LayerMemberReference CreateLayerMemberReference(
    uint32_t layer_id,
    const std::vector<uint32_t>& layer_member_ids) {
  LayerMemberReference layer_member_reference;
  layer_member_reference.set_layer_id(layer_id);
  for (uint32_t layer_member_id : layer_member_ids) {
    layer_member_reference.add_layer_member_ids(layer_member_id);
  }
  return layer_member_reference;
}

// Creates a test study.
Study CreateTestStudy(const std::vector<Study::Experiment>& experiments,
                      const std::optional<const LayerMemberReference>&
                          layer_member_reference = std::nullopt) {
  Study study;
  study.set_name("test_study");
  study.set_consistency(Study::PERMANENT);

  for (size_t i = 0; i < experiments.size(); ++i) {
    Study_Experiment* experiment_to_add = study.add_experiment();
    experiment_to_add->MergeFrom(experiments[i]);
    experiment_to_add->set_name(
        base::StrCat({"test_experiment_", base::NumberToString(i)}));
  }

  if (layer_member_reference.has_value()) {
    auto* layer_member_reference_to_add = study.mutable_layer();
    layer_member_reference_to_add->MergeFrom(*layer_member_reference);
  }
  return study;
}

// Creates a test layer member with slot ranges given as <start, end> pairs.
Layer::LayerMember CreateLayerMember(
    int layer_member_id,
    const std::vector<std::pair<int, int>>& slot_ranges) {
  Layer::LayerMember layer_member;
  layer_member.set_id(layer_member_id);

  for (const auto& slot_range : slot_ranges) {
    auto* slot = layer_member.add_slots();
    slot->set_start(slot_range.first);
    slot->set_end(slot_range.second);
  }
  return layer_member;
}

Layer CreateLayer(int layer_id,
                  int num_slots,
                  Layer::EntropyMode entropy_mode,
                  const std::vector<Layer::LayerMember>& layer_members) {
  Layer layer;
  layer.set_id(layer_id);
  layer.set_num_slots(num_slots);
  layer.set_entropy_mode(entropy_mode);

  for (const auto& layer_member : layer_members) {
    auto* member_to_add = layer.add_members();
    member_to_add->MergeFrom(layer_member);
  }

  return layer;
}

}  // namespace

class LimitedLayerEntropyCostTrackerTest : public ::testing::Test {
 public:
  LimitedLayerEntropyCostTrackerTest()
      : entropy_providers_(
            kTestClientId,
            // Using 100 as the test LES value (as opposed to the production
            // value of 8000) since this test suite is for using the limited
            // entropy randomization source.
            {0, 100},
            kTestLimitedEntropyRandomizationSource) {}

 protected:
  const EntropyProviders entropy_providers_;
};

TEST_F(LimitedLayerEntropyCostTrackerTest, TestConstructor_WithLimitedLayer) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 15);
  using EntropyEventList = LimitedLayerEntropyCostTracker::EntropyEventList;
  EXPECT_EQ(15, limited_entropy_tracker.entropy_limit_in_bits_);
  EXPECT_EQ(kTestLayerId, limited_entropy_tracker.limited_layer_id_);
  EXPECT_EQ(size_t(2),
            limited_entropy_tracker.entropy_events_by_member_id_.size());
  EXPECT_EQ(EntropyEventList({{0, 1.0}}),
            limited_entropy_tracker.entropy_events_by_member_id_[1]);
  EXPECT_EQ(EntropyEventList({{0, 1.0}}),
            limited_entropy_tracker.entropy_events_by_member_id_[2]);
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestConstructor_LayerMembersUsingEntropyAboveLimit) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 24}}), CreateLayerMember(2, {{25, 49}}),
       CreateLayerMember(3, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 1);
  using EntropyEventList = LimitedLayerEntropyCostTracker::EntropyEventList;
  EXPECT_EQ(1, limited_entropy_tracker.entropy_limit_in_bits_);
  EXPECT_EQ(kTestLayerId, limited_entropy_tracker.limited_layer_id_);
  EXPECT_EQ(size_t(3),
            limited_entropy_tracker.entropy_events_by_member_id_.size());
  // Note that the entropy used by layer members 1 and 2 is 2 bits, which is
  // above the total entropy limit of 1 bit.
  EXPECT_EQ(EntropyEventList({{0, 2.0}}),
            limited_entropy_tracker.entropy_events_by_member_id_[1]);
  EXPECT_EQ(EntropyEventList({{0, 2.0}}),
            limited_entropy_tracker.entropy_events_by_member_id_[2]);
  EXPECT_EQ(EntropyEventList({{0, 1.0}}),
            limited_entropy_tracker.entropy_events_by_member_id_[3]);

  // The total entropy used is zero because no study entropy has been added to
  // the limited_entropy_tracker at this stage.
  EXPECT_FALSE(limited_entropy_tracker.includes_study_entropy_);
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_StudyReferencingNoLayerMember) {
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(25, 100001),
      CreateGoogleWebExperiment(25, 200002),
      CreateExperiment(50),
  };
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments, CreateLayerMemberReference(kTestLayerId, {}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);

  EXPECT_FALSE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_StudyReferencingLimitedLayer) {
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(25, 100001),
      CreateGoogleWebExperiment(25, 200002),
      CreateExperiment(50),
  };
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 2);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(2, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
}

TEST_F(
    LimitedLayerEntropyCostTrackerTest,
    TestAddEntropyUsedByStudy_StudyReferencingLimitedLayerUsingFallbackField) {
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(25, 100001),
      CreateGoogleWebExperiment(25, 200002),
      CreateExperiment(50),
  };
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  LayerMemberReference fallback_layer_member_reference;
  fallback_layer_member_reference.set_layer_id(kTestLayerId);
  fallback_layer_member_reference.set_layer_member_id(kTestLayerMemberId);
  auto test_study =
      CreateTestStudy(experiments, fallback_layer_member_reference);
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 2);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(2, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_MultipleStudies) {
  using EntropyEventList = LimitedLayerEntropyCostTracker::EntropyEventList;
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(25, 100001),
      CreateGoogleWebExperiment(25, 200002),
      CreateExperiment(50),
  };
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 49}}),
       CreateLayerMember(kTestLayerMemberId + 1, {{50, 99}})});
  auto test_study_1 = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  auto test_study_2 = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId,
                                 {kTestLayerMemberId, kTestLayerMemberId + 1}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 5);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_1));
  EXPECT_EQ(3, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_2));
  EXPECT_EQ(5, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          // Entropy event for the layer member
          {kTimeMin, 1.0},
          // Entropy event for test_study_1.
          {kTimeMin, 2.0},
          // Entropy event for test_study_2.
          {kTimeMin, 2.0},
      })));
  EXPECT_THAT(limited_entropy_tracker
                .entropy_events_by_member_id_[kTestLayerMemberId + 1],
            UnorderedElementsAreArray(EntropyEventList({
                // Entropy event for the layer member
                {kTimeMin, 1.0},
                // Entropy event for test_study_2.
                {kTimeMin, 2.0},
            })));
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_EntropyLimitReachedByLayerMember) {
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(25, 100001),
      CreateGoogleWebExperiment(25, 200002),
      CreateExperiment(50),
  };
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 24}}),
       CreateLayerMember(kTestLayerMemberId + 1, {{25, 49}}),
       CreateLayerMember(kTestLayerMemberId + 2, {{50, 99}})});
  auto test_study_1 = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  auto test_study_2 = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 1);

  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_1));
  EXPECT_EQ(4, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_TRUE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_2));
  EXPECT_EQ(6, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_TRUE(limited_entropy_tracker.IsEntropyLimitExceeded());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_EntropyLimitReachedAfterAddingStudy) {
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(25, 100001),
      CreateGoogleWebExperiment(25, 200002),
      CreateExperiment(50),
  };
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study_1 = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  auto test_study_2 = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));

  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 3);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_1));
  EXPECT_EQ(2, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_2));
  EXPECT_EQ(4, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_TRUE(limited_entropy_tracker.IsEntropyLimitExceeded());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_IsTimeAware) {
          using EntropyEventList = LimitedLayerEntropyCostTracker::EntropyEventList;
  // Create two experiments arms. Studies created from these experiments will
  // consume 1 bit of entropy.
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(1, 100001),
      CreateGoogleWebExperiment(1, 200002),
  };

  // Use a fixed time to to ensure it doesn't overlap with the current time.
  base::Time now;
  ASSERT_TRUE(base::Time::FromString("2026-02-27 12:34:56", &now));
  constexpr auto kOneDay = base::Days(1);

  // Create a layer with 2 slots, and one layer member using one of the slots
  // (this consuming 1 bit of entropy). The other slot is unused.
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/2, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 0}})});

  // Create a study that is not active at the current time (ended in the
  // past).
  auto past_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  past_study.mutable_filter()->set_end_date(Timestamp(now - kOneDay));
  auto future_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  future_study.set_google_web_visibility_start_date(Timestamp(now + kOneDay));
  auto current_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  current_study.mutable_filter()->set_start_date(Timestamp(now - 3 * kOneDay));
  current_study.mutable_filter()->set_end_date(Timestamp(now + 3 * kOneDay));
  current_study.set_google_web_visibility_start_date(
      Timestamp(now - 2 * kOneDay));
  current_study.set_google_web_visibility_end_date(
      Timestamp(now + 2 * kOneDay));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);

  // The past study consumes entropy in the past.
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(past_study));
  ASSERT_EQ(2.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},  // past study doesn't specify a start date.
          {past_study.filter().end_date(), -1.0},
       })));

  // The future consumes entropy in the future. This doesn't overlap with the
  // past study, so the tracker should have still have the same maximum..
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(future_study));
  ASSERT_EQ(2.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},  // past_study doesn't specify a start date.
          {past_study.filter().end_date(), -1.0},
          {future_study.google_web_visibility_start_date(), 1.0},
     })));

  // The current study consumes entropy in the present and overlaps with both
  // the past and future studies.  Note that the event list tracks the latest
  // start and earliest end for current_study.
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(current_study));
  ASSERT_EQ(3.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          // Entropy event for the layer member
          {kTimeMin, 1.0},
          // past_study starts.
          {kTimeMin, 1.0},
          // current_study starts.
          {current_study.google_web_visibility_start_date(), 1.0},
          // past_study ends.
          {past_study.filter().end_date(), -1.0},
          // future_study starts.
          {future_study.google_web_visibility_start_date(), 1.0},
          // current_study ends.
          {current_study.google_web_visibility_end_date(), -1.0},
     })));

  // Adding another overlapping past study. Max entropy goes up by 1.
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(past_study));
  ASSERT_EQ(4.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          // Entropy event for the layer member
          {kTimeMin, 1.0},
          // past_study 1 and 2 start at the same time, kTimeMin.
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},
          // current_study starts.
          {current_study.google_web_visibility_start_date(), 1.0},
          // past_study 1 and 2 end at the same time.
          {past_study.filter().end_date(), -1.0},
          {past_study.filter().end_date(), -1.0},
          // future_study starts.
          {future_study.google_web_visibility_start_date(), 1.0},
          // current_study ends.
          {current_study.google_web_visibility_end_date(), -1.0},
     })));

  // Adding another overlapping future study. Max entropy stays the same, it
  // just occurs at multiple times.
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(future_study));
  ASSERT_EQ(4.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          // Entropy event for the layer member
          {kTimeMin, 1.0},
          // past_study 1 and 2 start at the same time, kTimeMin.
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},
          // current_study starts.
          {current_study.google_web_visibility_start_date(), 1.0},
          // past_study 1 and 2 end at the same time.
          {past_study.filter().end_date(), -1.0},
          {past_study.filter().end_date(), -1.0},
          // future_study 1 and 2 start at the same time.
          {future_study.google_web_visibility_start_date(), 1.0},
          {future_study.google_web_visibility_start_date(), 1.0},
          // current_study ends.
          {current_study.google_web_visibility_end_date(), -1.0},
     })));

  // Validate that negative entropy events are processed before positive ones
  // occurring at the same time.
  Study past_study_3 = past_study;
  past_study_3.mutable_filter()->set_end_date(
      current_study.google_web_visibility_start_date());
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(past_study_3));
  ASSERT_EQ(4.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_FALSE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          // Entropy event for the layer member
          {kTimeMin, 1.0},
          // past_study 1, 2 and 3 start at the same time, kTimeMin.
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},
          // past_study 3 ending is processed before current_studies 1 and 2
          // starting, though they occur at the same time.
          {current_study.google_web_visibility_start_date(), -1.0},
          // current_study starts.
          {current_study.google_web_visibility_start_date(), 1.0},
          // past_study 1 and 2 end at the same time.
          {past_study.filter().end_date(), -1.0},
          {past_study.filter().end_date(), -1.0},
          // future_study 1 and 2 start at the same time.
          {future_study.google_web_visibility_start_date(), 1.0},
          {future_study.google_web_visibility_start_date(), 1.0},
          // current_study ends.
          {current_study.google_web_visibility_end_date(), -1.0},
     })));

  // Adding another overlapping instance of current_study exceeds the entropy
  //limit.
  ASSERT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(current_study));
  ASSERT_EQ(5.0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  ASSERT_TRUE(limited_entropy_tracker.IsEntropyLimitExceeded());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(EntropyEventList({
          // Entropy event for the layer member
          {kTimeMin, 1.0},
          // past_study 1, 2 and 3 start at the same time, kTimeMin.
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},
          {kTimeMin, 1.0},
          // past_study 3 ending is processed before current_studies 1 and 2
          // starting, though they occur at the same time.
          {current_study.google_web_visibility_start_date(), -1.0},
          // current_study 1 and 2 start at the same time.
          {current_study.google_web_visibility_start_date(), 1.0},
          {current_study.google_web_visibility_start_date(), 1.0},
          // past_study 1 and 2 end at the same time.
          {past_study.filter().end_date(), -1.0},
          {past_study.filter().end_date(), -1.0},
          // future_study 1 and 2 start at the same time.
          {future_study.google_web_visibility_start_date(), 1.0},
          {future_study.google_web_visibility_start_date(), 1.0},
          // current_study 1 and 2 end at the same time.
          {current_study.google_web_visibility_end_date(), -1.0},
          {current_study.google_web_visibility_end_date(), -1.0},
     })));
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_NoGoogleWebExperimentID) {
  // Experiments without google_web_experiment_id are excluded from entropy
  // calculation.
  std::vector<Study::Experiment> experiments = {
      CreateExperiment(10), CreateExperiment(20), CreateExperiment(30),
      CreateExperiment(40)};
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_LaunchedStudy) {
  // Experiments without google_web_experiment_id are excluded from entropy
  // calculation.
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {
          CreateLayerMember(kTestLayerMemberId,
                            {{0, 49}}),  // 50% of the population.
          CreateLayerMember(kTestLayerMemberId + 1,
                            {{50, 99}}),  // 05% of the population.
      });
  auto launched_study = CreateTestStudy(
      {CreateGoogleWebExperiment(100, 100001)},  // 100% launched arm
      CreateLayerMemberReference(
          kTestLayerId,
          {kTestLayerMemberId, kTestLayerMemberId + 1}));  // 100% population.
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(launched_study));
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_LaunchedAndActiveStudies) {
  using EntropyEventList = LimitedLayerEntropyCostTracker::EntropyEventList;
  // Experiments without google_web_experiment_id are excluded from entropy
  // calculation.
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {
          CreateLayerMember(kTestLayerMemberId,
                            {{0, 49}}),  // 50% of the population.
          CreateLayerMember(kTestLayerMemberId + 1,
                            {{50, 99}}),  // 50% of the population.
      });
  auto launched_study = CreateTestStudy(
      {CreateGoogleWebExperiment(100, 100001)},  // 100% launched arm
      CreateLayerMemberReference(
          kTestLayerId,
          {kTestLayerMemberId, kTestLayerMemberId + 1}));  // 100% population.
  auto active_study = CreateTestStudy(
      {
          CreateGoogleWebExperiment(50, 100001),  // 50% active arm 1
          CreateGoogleWebExperiment(50, 200002),  // 50% active arm 2
      },
      CreateLayerMemberReference(kTestLayerId,
                                 {kTestLayerMemberId}));  // 25% population.
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(launched_study));
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(active_study));
  EXPECT_EQ(2,  // member: 1 bit; active study: 1 bit, launched study: 0 bits
            limited_entropy_tracker.GetMaxEntropyUsedForTesting());
  EXPECT_THAT(
      limited_entropy_tracker.entropy_events_by_member_id_[kTestLayerMemberId],
      UnorderedElementsAreArray(
          EntropyEventList({{kTimeMin, 1.0}, {kTimeMin, 1.0}})));
  EXPECT_THAT(limited_entropy_tracker
                  .entropy_events_by_member_id_[kTestLayerMemberId + 1],
              UnorderedElementsAreArray(EntropyEventList({{0, 1.0}})));
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_SkipZeroProbabilityWeightedStudies) {
  std::vector<Study::Experiment> experiments = {
      CreateGoogleWebExperiment(0, 100001), CreateExperiment(100)};
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_NoExperiments) {
  std::vector<Study::Experiment> experiments;
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_WithGoogleWebTriggerExpID) {
  std::vector<Study::Experiment> experiments = {
      CreateTriggerExperiment(25, 100001), CreateTriggerExperiment(25, 100002),
      CreateExperiment(50)};
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(2, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_SessionConsistency) {
  std::vector<Study::Experiment> experiments = {
      CreateTriggerExperiment(25, 100001), CreateTriggerExperiment(25, 100002),
      CreateExperiment(50)};
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 99}})});
  auto test_study = CreateTestStudy(
      experiments,
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  test_study.set_consistency(Study::SESSION);
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 13);

  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study));
  EXPECT_EQ(0, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_StudiesOfDifferentEntropyInLayerMember) {
  // TODO(siakabaro): The test name doesn't seems to match the code?
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(kTestLayerMemberId, {{0, 49}})});
  auto test_study_1 = CreateTestStudy(
      CreateExperimentsWithTwoBitsOfEntropy(),
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  auto test_study_2 = CreateTestStudy(
      CreateExperimentsWithTwoBitsOfEntropy(),
      CreateLayerMemberReference(kTestLayerId, {kTestLayerMemberId}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 5);

  // Expecting 5 bits of total usage with 4 bits from the two 2-bit studies,
  // and 1 bit from the layer member.
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_1));
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_2));
  EXPECT_EQ(5, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_MultipleLayerMembers) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 74}}),
       CreateLayerMember(3, {{75, 99}})});
  std::vector<Study> studies = {
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(kTestLayerId, {1})),
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(kTestLayerId, {2})),
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(kTestLayerId, {3})),
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(kTestLayerId, {3})),
  };
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 6);

  // Expecting 6 bits of total usage with 4 bits from the two 2-bit studies in
  // layer member #3, and 2 bit from the 25% layer member itself.
  for (const Study& study : studies) {
    EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(study));
  }
  EXPECT_EQ(6, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_ReferencingMultipleLayerMembers) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 74}}),
       CreateLayerMember(3, {{75, 99}})});
  auto test_study_1 =
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(kTestLayerId,
                                                 /*layer_member_ids=*/{1, 2}));
  auto test_study_2 =
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(kTestLayerId,
                                                 /*layer_member_ids=*/{3}));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);

  // Entropy usages:
  // - layer member 1:
  //   1 bit from the layer member + 2 bits from the first study = 3 bits.
  // - layer member 2:
  //   2 bit from the layer member + 2 bits from the first study = 4 bits.
  // - layer member 3:
  //   2 bit from the layer member + 2 bits from the second study = 4 bits.
  // - Therefore the layer uses a maximum of 4 bits.
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_1));
  EXPECT_TRUE(limited_entropy_tracker.AddEntropyUsedByStudy(test_study_2));
  EXPECT_EQ(4, limited_entropy_tracker.GetMaxEntropyUsedForTesting());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_InvalidLayerId) {
  auto test_layer = CreateLayer(
      0u, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_NoEntropy) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 0.0);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest, TestAddEntropyUsedByStudy_NoSlots) {
  auto test_layer = CreateLayer(kTestLayerId, /*num_slots=*/0u,
                                /*entropy_mode=*/Layer::LIMITED, {});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_WrongEntropyMode) {
  auto test_layer = CreateLayer(kTestLayerId, /*num_slots=*/100,
                                /*entropy_mode=*/Layer::DEFAULT, {});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_NoLayerMembers) {
  auto test_layer = CreateLayer(kTestLayerId, /*num_slots=*/100,
                                /*entropy_mode=*/Layer::LIMITED, {});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_OverlappingLayerMembers) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{40, 99}})});
  ASSERT_FALSE(VariationsLayers::AreSlotBoundsValid(test_layer));
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_InvalidLayerMemberId) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(0, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_DuplicatedLayerMemberId) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{1, 49}}), CreateLayerMember(1, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_Reference_InvalidLayerId) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_TRUE(limited_entropy_tracker.IsValid());
  auto test_study =
      CreateTestStudy(CreateExperimentsWithTwoBitsOfEntropy(),
                      CreateLayerMemberReference(0,  // Invalid layer id.
                                                 /*layer_member_ids=*/{1, 2}));
  limited_entropy_tracker.AddEntropyUsedByStudy(test_study);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_Reference_DanglingLayerId) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_TRUE(limited_entropy_tracker.IsValid());
  auto test_study = CreateTestStudy(
      CreateExperimentsWithTwoBitsOfEntropy(),
      CreateLayerMemberReference(kTestLayerId + 7,  // Undefined layer id.
                                 /*layer_member_ids=*/{1, 2}));
  limited_entropy_tracker.AddEntropyUsedByStudy(test_study);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_Reference_InvalidLayerMemberId) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_TRUE(limited_entropy_tracker.IsValid());
  auto test_study = CreateTestStudy(
      CreateExperimentsWithTwoBitsOfEntropy(),
      CreateLayerMemberReference(
          kTestLayerId,
          /*layer_member_ids=*/{0}));  // Invalid layer member id.
  limited_entropy_tracker.AddEntropyUsedByStudy(test_study);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

TEST_F(LimitedLayerEntropyCostTrackerTest,
       TestAddEntropyUsedByStudy_Reference_DanglingLayerMemberId) {
  auto test_layer = CreateLayer(
      kTestLayerId, /*num_slots=*/100, /*entropy_mode=*/Layer::LIMITED,
      {CreateLayerMember(1, {{0, 49}}), CreateLayerMember(2, {{50, 99}})});
  LimitedLayerEntropyCostTracker limited_entropy_tracker(test_layer, 4);
  EXPECT_TRUE(limited_entropy_tracker.IsValid());
  auto test_study = CreateTestStudy(
      CreateExperimentsWithTwoBitsOfEntropy(),
      CreateLayerMemberReference(
          kTestLayerId,
          /*layer_member_ids=*/{7}));  // Undefined layer member id.
  limited_entropy_tracker.AddEntropyUsedByStudy(test_study);
  EXPECT_FALSE(limited_entropy_tracker.IsValid());
}

}  // namespace variations