/**
 * Copyright 2020-2026, XGBoost Contributors
 */
#include <gtest/gtest.h>
#include <xgboost/context.h>  // for Context

#include <memory>  // for unique_ptr
#include <tuple>   // for tuple
#include <vector>  // for vector

#include "../../../../src/tree/gpu_hist/expand_entry.cuh"  // for GPUExpandEntry
#include "../../../../src/tree/gpu_hist/histogram.cuh"
#include "../../../../src/tree/gpu_hist/row_partitioner.cuh"  // for RowPartitioner
#include "../../../../src/tree/hist/hist_param.h"             // for HistMakerTrainParam
#include "../../../../src/tree/param.h"                       // for TrainParam
#include "../../categorical_helpers.h"                        // for OneHotEncodeFeature
#include "../../helpers.h"
#include "../../histogram_helpers.h"  // for BuildEllpackPage
#include "dummy_quantizer.cuh"

namespace xgboost::tree {
TEST(Histogram, DeviceHistogramStorage) {
  // Ensures that node allocates correctly after reaching `kStopGrowingSize`.
  auto ctx = MakeCUDACtx(0);
  constexpr size_t kNBins = 128;
  constexpr int kNNodes = 4;
  constexpr size_t kStopGrowing = kNNodes * kNBins * 2u;
  DeviceHistogramStorage histogram{};
  histogram.Reset(&ctx, kNBins, kNNodes);
  for (int i = 0; i < kNNodes; ++i) {
    histogram.AllocateHistograms(&ctx, {i});
  }
  ASSERT_EQ(histogram.Data().size(), kStopGrowing);
  histogram.Reset(&ctx, kNBins, kNNodes);

  // Use allocated memory but do not erase nidx_map.
  for (int i = 0; i < kNNodes; ++i) {
    histogram.AllocateHistograms(&ctx, {i});
  }
  for (int i = 0; i < kNNodes; ++i) {
    ASSERT_TRUE(histogram.HistogramExists(i));
  }

  // Add two new nodes
  histogram.AllocateHistograms(&ctx, {kNNodes});
  histogram.AllocateHistograms(&ctx, {kNNodes + 1});

  // Old cached nodes should still exist
  for (int i = 0; i < kNNodes; ++i) {
    ASSERT_TRUE(histogram.HistogramExists(i));
  }

  // Should be deleted
  ASSERT_FALSE(histogram.HistogramExists(kNNodes));
  // Most recent node should exist
  ASSERT_TRUE(histogram.HistogramExists(kNNodes + 1));

  // Add same node again - should fail
  EXPECT_ANY_THROW(histogram.AllocateHistograms(&ctx, {kNNodes + 1}););
}

TEST(Histogram, SubtractionTrack) {
  auto ctx = MakeCUDACtx(0);

  auto page = BuildEllpackPage(&ctx, 64, 4);
  auto cuts = page->CutsShared();
  FeatureGroups fg{*cuts, true, std::numeric_limits<std::size_t>::max()};
  auto n_total_bins = cuts->TotalBins();

  // 2 nodes
  auto max_cached_hist_nodes = 2ull;
  DeviceHistogramBuilder histogram;
  histogram.Reset(&ctx, max_cached_hist_nodes, n_total_bins, false);
  histogram.AllocateHistograms(&ctx, {0, 1, 2});
  GPUExpandEntry root;
  root.nidx = 0;
  auto need_build = histogram.SubtractHist<GPUExpandEntry>(&ctx, {root}, {0}, {1});

  std::vector<GPUExpandEntry> candidates(2);
  candidates[0].nidx = 1;
  candidates[1].nidx = 2;

  need_build = histogram.SubtractHist(&ctx, candidates, {3, 5}, {4, 6});
  ASSERT_EQ(need_build.size(), 2);
  ASSERT_EQ(need_build[0], 4);
  ASSERT_EQ(need_build[1], 6);
}

namespace {
void TestDeterministicHistogram(bool is_dense, std::size_t shm_size, bool force_global) {
  Context ctx = MakeCUDACtx(0);
  size_t constexpr kBins = 256, kCols = 120, kRows = 16384, kRounds = 16;
  float constexpr kLower = -1e-2, kUpper = 1e2;

  float sparsity = is_dense ? 0.0f : 0.5f;
  auto matrix = RandomDataGenerator(kRows, kCols, sparsity).GenerateDMatrix();
  auto batch_param = BatchParam{kBins, tree::TrainParam::DftSparseThreshold()};

  for (auto const& batch : matrix->GetBatches<EllpackPage>(&ctx, batch_param)) {
    auto* page = batch.Impl();

    tree::RowPartitioner row_partitioner;
    row_partitioner.Reset(&ctx, kRows, page->base_rowid);
    auto ridx = row_partitioner.GetRows(0);

    bst_bin_t num_bins = kBins * kCols;
    dh::device_vector<GradientPairInt64> histogram(num_bins);
    auto d_histogram = dh::ToSpan(histogram);
    auto gpair = GenerateGradientsFixedPoint(&ctx, kRows, 1, kLower, kUpper).gpair;

    FeatureGroups feature_groups{page->Cuts(), page->IsDenseCompressed(), shm_size};

    DeviceHistogramBuilder builder;
    builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), num_bins, force_global);
    page->Visit(&ctx, {}, [&](auto&& acc) {
      builder.BuildHistogram(&ctx, acc, feature_groups.DeviceAccessor(ctx.Device()),
                             gpair.View(ctx.Device()).Values(), ridx, d_histogram);
    });

    std::vector<GradientPairInt64> histogram_h(num_bins);
    dh::safe_cuda(cudaMemcpy(histogram_h.data(), d_histogram.data(),
                             num_bins * sizeof(GradientPairInt64), cudaMemcpyDeviceToHost));

    for (std::size_t i = 0; i < kRounds; ++i) {
      dh::device_vector<GradientPairInt64> new_histogram(num_bins);
      auto d_new_histogram = dh::ToSpan(new_histogram);

      DeviceHistogramBuilder builder;
      builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), num_bins, force_global);
      page->Visit(&ctx, {}, [&](auto&& acc) {
        builder.BuildHistogram(&ctx, acc, feature_groups.DeviceAccessor(ctx.Device()),
                               gpair.View(ctx.Device()).Values(), ridx, d_new_histogram);
      });

      std::vector<GradientPairInt64> new_histogram_h(num_bins);
      dh::safe_cuda(cudaMemcpy(new_histogram_h.data(), d_new_histogram.data(),
                               num_bins * sizeof(GradientPairInt64), cudaMemcpyDeviceToHost));
      for (size_t j = 0; j < new_histogram_h.size(); ++j) {
        ASSERT_EQ(new_histogram_h[j].GetQuantisedGrad(), histogram_h[j].GetQuantisedGrad());
        ASSERT_EQ(new_histogram_h[j].GetQuantisedHess(), histogram_h[j].GetQuantisedHess());
      }
    }

    {
      auto gpair = GenerateGradientsFixedPoint(&ctx, kRows, 1, kLower, kUpper).gpair;

      // Use a single feature group to compute the baseline.
      FeatureGroups single_group(page->Cuts());

      dh::device_vector<GradientPairInt64> baseline(num_bins);
      DeviceHistogramBuilder builder;
      // Single group must use global memory.
      builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), num_bins, /*force_global=*/true);
      page->Visit(&ctx, {}, [&](auto&& acc) {
        builder.BuildHistogram(&ctx, acc, single_group.DeviceAccessor(ctx.Device()),
                               gpair.View(ctx.Device()).Values(), ridx, dh::ToSpan(baseline));
      });

      std::vector<GradientPairInt64> baseline_h(num_bins);
      dh::safe_cuda(cudaMemcpy(baseline_h.data(), baseline.data().get(),
                               num_bins * sizeof(GradientPairInt64), cudaMemcpyDeviceToHost));

      for (size_t i = 0; i < baseline.size(); ++i) {
        ASSERT_NEAR(baseline_h[i].GetQuantisedGrad(), histogram_h[i].GetQuantisedGrad(),
                    baseline_h[i].GetQuantisedGrad() * 1e-3);
      }
    }
  }
}

class TestGPUDeterministic : public ::testing::TestWithParam<std::tuple<bool, std::size_t, bool>> {
 protected:
  void Run() {
    auto [is_dense, shm_size, force_global] = this->GetParam();
    if (shm_size > dh::MaxSharedMemoryOptin(0) && !force_global) {
      force_global = true;  // We will have to skip this test to avoid false check in the builder.
    }
    TestDeterministicHistogram(is_dense, shm_size, force_global);
  }
};
}  // anonymous namespace

TEST_P(TestGPUDeterministic, Histogram) { this->Run(); }

INSTANTIATE_TEST_SUITE_P(Histogram, TestGPUDeterministic,
                         ::testing::Combine(::testing::Bool(),
                                            ::testing::Values(48 * 1024, 64 * 1024, 160 * 1024),
                                            ::testing::Bool()));

void ValidateCategoricalHistogram(size_t n_categories, common::Span<GradientPairInt64> onehot,
                                  common::Span<GradientPairInt64> cat) {
  auto cat_sum = std::accumulate(cat.cbegin(), cat.cend(), GradientPairInt64{});
  for (size_t c = 0; c < n_categories; ++c) {
    auto zero = onehot[c * 2];
    auto one = onehot[c * 2 + 1];

    auto chosen = cat[c];
    auto not_chosen = cat_sum - chosen;
    ASSERT_EQ(zero, not_chosen);
    ASSERT_EQ(one, chosen);
  }
}

// Test 1 vs rest categorical histogram is equivalent to one hot encoded data.
void TestGPUHistogramCategorical(size_t num_categories) {
  auto ctx = MakeCUDACtx(0);
  size_t kRows = std::max(static_cast<decltype(num_categories)>(340), num_categories);
  size_t constexpr kBins = 256;
  auto x = GenerateRandomCategoricalSingleColumn(kRows, num_categories);
  auto cat_m = GetDMatrixFromData(x, kRows, 1);
  cat_m->Info().feature_types.HostVector().push_back(FeatureType::kCategorical);
  auto batch_param = BatchParam{kBins, tree::TrainParam::DftSparseThreshold()};
  tree::RowPartitioner row_partitioner;
  row_partitioner.Reset(&ctx, kRows, 0);
  auto ridx = row_partitioner.GetRows(0);
  dh::device_vector<GradientPairInt64> cat_hist(num_categories);

  auto gpair = GenerateRandomGradients(kRows, 0, 2);
  gpair.SetDevice(DeviceOrd::CUDA(0));
  GradientQuantiserGroup quantiser_group{
      &ctx, linalg::MakeVec(ctx.Device(), gpair.ConstDeviceSpan()), MetaInfo()};
  linalg::Matrix<GradientPairInt64> gpairs_i64;
  CalcQuantizedGpairs(&ctx, linalg::MakeTensorView(&ctx, gpair.ConstDeviceSpan(), gpair.Size(), 1),
                      quantiser_group.DeviceSpan(), &gpairs_i64);
  /**
   * Generate hist with cat data.
   */
  for (auto const& batch : cat_m->GetBatches<EllpackPage>(&ctx, batch_param)) {
    auto* page = batch.Impl();
    FeatureGroups single_group(page->Cuts());
    DeviceHistogramBuilder builder;
    builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), num_categories, false);
    page->Visit(&ctx, {}, [&](auto&& acc) {
      builder.BuildHistogram(&ctx, acc, single_group.DeviceAccessor(ctx.Device()),
                             gpairs_i64.View(ctx.Device()).Values(), ridx, dh::ToSpan(cat_hist));
    });
  }

  /**
   * Generate hist with one hot encoded data.
   */
  auto x_encoded = OneHotEncodeFeature(x, num_categories);
  auto encode_m = GetDMatrixFromData(x_encoded, kRows, num_categories);
  dh::device_vector<GradientPairInt64> encode_hist(2 * num_categories);
  for (auto const& batch : encode_m->GetBatches<EllpackPage>(&ctx, batch_param)) {
    auto* page = batch.Impl();
    FeatureGroups single_group(page->Cuts());
    DeviceHistogramBuilder builder;
    builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), encode_hist.size(), false);
    page->Visit(&ctx, {}, [&](auto&& acc) {
      builder.BuildHistogram(&ctx, acc, single_group.DeviceAccessor(ctx.Device()),
                             gpairs_i64.View(ctx.Device()).Values(), ridx, dh::ToSpan(encode_hist));
    });
  }

  std::vector<GradientPairInt64> h_cat_hist(cat_hist.size());
  thrust::copy(cat_hist.begin(), cat_hist.end(), h_cat_hist.begin());

  std::vector<GradientPairInt64> h_encode_hist(encode_hist.size());
  thrust::copy(encode_hist.begin(), encode_hist.end(), h_encode_hist.begin());
  ValidateCategoricalHistogram(num_categories, common::Span<GradientPairInt64>{h_encode_hist},
                               common::Span<GradientPairInt64>{h_cat_hist});
}

TEST(Histogram, GPUHistCategorical) {
  for (size_t num_categories = 2; num_categories < 8; ++num_categories) {
    TestGPUHistogramCategorical(num_categories);
  }
  // Larger than the shared memory size, must use global memory since there's no feature
  // group with a single feature.
  auto max_shmem = dh::MaxSharedMemoryOptin(0);
  auto n_categories = common::DivRoundUp(max_shmem, sizeof(GradientPairInt64)) * 2;
  TestGPUHistogramCategorical(n_categories);
}

namespace {
// Atomic add as type cast for test.
XGBOOST_DEV_INLINE int64_t atomicAdd(int64_t* dst, int64_t src) {  // NOLINT
  uint64_t* u_dst = reinterpret_cast<uint64_t*>(dst);
  uint64_t u_src = *reinterpret_cast<uint64_t*>(&src);
  uint64_t ret = ::atomicAdd(u_dst, u_src);
  return *reinterpret_cast<int64_t*>(&ret);
}
}  // namespace

void TestAtomicAdd() {
  size_t n_elements = 1024;
  dh::device_vector<int64_t> result_a(1, 0);
  auto d_result_a = result_a.data().get();

  dh::device_vector<int64_t> result_b(1, 0);
  auto d_result_b = result_b.data().get();

  /**
   * Test for simple inputs
   */
  std::vector<int64_t> h_inputs(n_elements);
  for (size_t i = 0; i < h_inputs.size(); ++i) {
    h_inputs[i] = (i % 2 == 0) ? i : -i;
  }
  dh::device_vector<int64_t> inputs(h_inputs);
  auto d_inputs = inputs.data().get();

  dh::LaunchN(n_elements, [=] __device__(size_t i) {
    AtomicAdd64As32(d_result_a, d_inputs[i]);
    atomicAdd(d_result_b, d_inputs[i]);
  });
  ASSERT_EQ(result_a[0], result_b[0]);

  /**
   * Test for positive values that don't fit into 32 bit integer.
   */
  thrust::fill(inputs.begin(), inputs.end(), (std::numeric_limits<uint32_t>::max() / 2));
  thrust::fill(result_a.begin(), result_a.end(), 0);
  thrust::fill(result_b.begin(), result_b.end(), 0);
  dh::LaunchN(n_elements, [=] __device__(size_t i) {
    AtomicAdd64As32(d_result_a, d_inputs[i]);
    atomicAdd(d_result_b, d_inputs[i]);
  });
  ASSERT_EQ(result_a[0], result_b[0]);
  ASSERT_GT(result_a[0], std::numeric_limits<uint32_t>::max());
  CHECK_EQ(thrust::reduce(inputs.begin(), inputs.end(), int64_t(0)), result_a[0]);

  /**
   * Test for negative values that don't fit into 32 bit integer.
   */
  thrust::fill(inputs.begin(), inputs.end(), (std::numeric_limits<int32_t>::min() / 2));
  thrust::fill(result_a.begin(), result_a.end(), 0);
  thrust::fill(result_b.begin(), result_b.end(), 0);
  dh::LaunchN(n_elements, [=] __device__(size_t i) {
    AtomicAdd64As32(d_result_a, d_inputs[i]);
    atomicAdd(d_result_b, d_inputs[i]);
  });
  ASSERT_EQ(result_a[0], result_b[0]);
  ASSERT_LT(result_a[0], std::numeric_limits<int32_t>::min());
  CHECK_EQ(thrust::reduce(inputs.begin(), inputs.end(), int64_t(0)), result_a[0]);
}

TEST(Histogram, AtomicAddInt64) { TestAtomicAdd(); }

TEST(Histogram, Quantiser) {
  auto ctx = MakeCUDACtx(0);
  std::size_t n_samples{16};
  HostDeviceVector<GradientPair> gpair(n_samples, GradientPair{1.0, 1.0});
  gpair.SetDevice(ctx.Device());

  GradientQuantiserGroup quantiser_group{
      &ctx, linalg::MakeVec(ctx.Device(), gpair.ConstDeviceSpan()), MetaInfo()};
  auto quantiser = quantiser_group[0];
  for (auto v : gpair.ConstHostVector()) {
    auto gh = quantiser.ToFloatingPoint(quantiser.ToFixedPoint(v));
    ASSERT_EQ(gh.GetGrad(), 1.0);
    ASSERT_EQ(gh.GetHess(), 1.0);
  }
}
namespace {
enum CacheMode {
  kNoCache = 0,
  kCopy = 1,
  kDirect = 2,
};

class HistogramExternalMemoryTest
    : public ::testing::TestWithParam<std::tuple<float, bool, CacheMode>> {
 public:
  void Run(float sparsity, bool force_global, CacheMode cache_mode) {
    auto ctx = MakeCUDACtx(0);
    bst_idx_t n_samples{512}, n_features{12}, n_batches{3};
    std::vector<std::unique_ptr<RowPartitioner>> partitioners;
    auto rng = RandomDataGenerator{n_samples, n_features, sparsity}.Batches(n_batches);
    bst_bin_t n_bins = 16;
    std::shared_ptr<DMatrix> p_fmat;
    switch (cache_mode) {
      case kCopy:
      case kDirect: {
        p_fmat = rng.CacheHostRatio(0.5)
                     .Device(ctx.Device())
                     .Bins(n_bins)
                     .OnHost(true)
                     .MinPageCacheBytes(n_bins * n_features)
                     .GenerateExtMemQuantileDMatrix("cache", true);
        break;
      }
      case kNoCache: {
        p_fmat = rng.GenerateSparsePageDMatrix("cache", true);
        break;
      }
    }

    BatchParam p{n_bins, TrainParam::DftSparseThreshold()};
    if (cache_mode == kDirect) {
      p.prefetch_copy = false;
    } else if (cache_mode == kCopy) {
      p.prefetch_copy = true;
    }

    std::unique_ptr<FeatureGroups> fg;
    dh::device_vector<GradientPairInt64> single_hist;
    dh::device_vector<GradientPairInt64> multi_hist;

    auto gpair = GenerateGradientsFixedPoint(&ctx, n_samples).gpair;
    std::shared_ptr<common::HistogramCuts> cuts;

    std::size_t row_stride = 0;
    {
      /**
       * Multi page.
       */
      std::int32_t k{0};
      for (auto const& page : p_fmat->GetBatches<EllpackPage>(&ctx, p)) {
        auto impl = page.Impl();
        row_stride = impl->info.row_stride;
        if (k == 0) {
          // Initialization
          fg = std::make_unique<FeatureGroups>(impl->Cuts());
          auto init = GradientPairInt64{0, 0};
          multi_hist = decltype(multi_hist)(impl->Cuts().TotalBins(), init);
          single_hist = decltype(single_hist)(impl->Cuts().TotalBins(), init);
          cuts = std::make_shared<common::HistogramCuts>(impl->Cuts());
        }

        partitioners.emplace_back(std::make_unique<RowPartitioner>());
        partitioners.back()->Reset(&ctx, impl->Size(), impl->base_rowid);

        auto ridx = partitioners.at(k)->GetRows(0);
        auto d_histogram = dh::ToSpan(multi_hist);
        DeviceHistogramBuilder builder;
        builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), d_histogram.size(),
                      force_global);
        impl->Visit(&ctx, {}, [&](auto&& acc) {
          builder.BuildHistogram(&ctx, acc, fg->DeviceAccessor(ctx.Device()),
                                 gpair.View(ctx.Device()).Values(), ridx, d_histogram);
        });
        ++k;
      }
      ASSERT_EQ(k, n_batches);
    }

    {
      /**
       * Single page.
       */
      RowPartitioner partitioner;
      partitioner.Reset(&ctx, p_fmat->Info().num_row_, 0);

      auto concat = EllpackPageImpl(&ctx, cuts, sparsity == 0.0, row_stride, n_samples);
      std::vector<float> hess(p_fmat->Info().num_row_, 1.0f);
      std::size_t offset = 0;
      for (auto const& page : p_fmat->GetBatches<EllpackPage>(&ctx, p)) {
        bst_idx_t num_elements = concat.Copy(&ctx, page.Impl(), offset);
        offset += num_elements;
      }
      auto ridx = partitioner.GetRows(0);
      auto d_histogram = dh::ToSpan(single_hist);
      DeviceHistogramBuilder builder;
      builder.Reset(&ctx, HistMakerTrainParam::CudaDefaultNodes(), d_histogram.size(),
                    force_global);
      concat.Visit(&ctx, {}, [&](auto&& acc) {
        builder.BuildHistogram(&ctx, acc, fg->DeviceAccessor(ctx.Device()),
                               gpair.View(ctx.Device()).Values(), ridx, d_histogram);
      });
    }

    std::vector<GradientPairInt64> h_single(single_hist.size());
    thrust::copy(single_hist.begin(), single_hist.end(), h_single.begin());
    std::vector<GradientPairInt64> h_multi(multi_hist.size());
    thrust::copy(multi_hist.begin(), multi_hist.end(), h_multi.begin());

    for (std::size_t i = 0; i < single_hist.size(); ++i) {
      ASSERT_EQ(h_single[i].GetQuantisedGrad(), h_multi[i].GetQuantisedGrad()) << i;
      ASSERT_EQ(h_single[i].GetQuantisedHess(), h_multi[i].GetQuantisedHess());
    }
  }
};
}  // namespace

TEST_P(HistogramExternalMemoryTest, ExternalMemory) {
  std::apply(&HistogramExternalMemoryTest::Run, std::tuple_cat(std::make_tuple(this), GetParam()));
}

INSTANTIATE_TEST_SUITE_P(
    Histogram, HistogramExternalMemoryTest,
    ::testing::Combine(::testing::Values(0.0f, 0.2f, 0.8f), ::testing::Bool(),
                       ::testing::Values(kNoCache, kDirect, kCopy)),
    [](::testing::TestParamInfo<HistogramExternalMemoryTest::ParamType> const& info) {
      std::stringstream ss;
      auto const& p = info.param;
      ss << "sparsity_0" << (std::get<0>(p) * 10) << "_global_" << std::get<1>(p) << "_dcache_";
      switch (std::get<2>(p)) {
        case kNoCache:
          ss << "nocache";
          break;
        case kDirect:
          ss << "direct";
          break;
        case kCopy:
          ss << "copy";
          break;
      }
      return ss.str();
    });
}  // namespace xgboost::tree
