/**
 * Copyright 2022-2024, XGBoost Contributors
 *
 * @brief Utilities for estimating initial score.
 */
#include <thrust/execution_policy.h>            // cuda::par
#include <thrust/iterator/counting_iterator.h>  // thrust::make_counting_iterator

#include <cstddef>  // std::size_t

#include "../collective/aggregator.cuh"  // for GlobalSum
#include "../common/device_helpers.cuh"  // dh::MakeTransformIterator
#include "fit_stump.h"
#include "xgboost/base.h"     // GradientPairPrecise, GradientPair, XGBOOST_DEVICE
#include "xgboost/context.h"  // Context
#include "xgboost/linalg.h"   // TensorView, Tensor, Constant
#include "xgboost/logging.h"  // CHECK_EQ
#include "xgboost/span.h"     // span

namespace xgboost::tree::cuda_impl {
void FitStump(Context const* ctx, MetaInfo const& info,
              linalg::TensorView<GradientPair const, 2> gpair, linalg::VectorView<float> out) {
  auto n_targets = out.Size();
  CHECK_EQ(n_targets, gpair.Shape(1));
  linalg::Vector<GradientPairPrecise> sum = linalg::Constant(ctx, GradientPairPrecise{}, n_targets);
  CHECK(out.Contiguous());

  // Reduce by column
  auto key_it = dh::MakeTransformIterator<bst_target_t>(
      thrust::make_counting_iterator(0ul),
      [=] XGBOOST_DEVICE(std::size_t i) -> bst_target_t { return i / gpair.Shape(0); });
  auto grad_it = dh::MakeTransformIterator<GradientPairPrecise>(
      thrust::make_counting_iterator(0ul),
      [=] XGBOOST_DEVICE(std::size_t i) -> GradientPairPrecise {
        auto target = i / gpair.Shape(0);
        auto sample = i % gpair.Shape(0);
        return GradientPairPrecise{gpair(sample, target)};
      });
  auto d_sum = sum.View(ctx->Device());
  CHECK(d_sum.CContiguous());

  dh::XGBCachingDeviceAllocator<char> alloc;
  auto policy = thrust::cuda::par(alloc);
  thrust::reduce_by_key(policy, key_it, key_it + gpair.Size(), grad_it,
                        thrust::make_discard_iterator(), dh::tbegin(d_sum.Values()));

  auto rc = collective::GlobalSum(ctx, info,
                                  linalg::MakeVec(reinterpret_cast<double*>(d_sum.Values().data()),
                                                  d_sum.Size() * 2, ctx->Device()));
  SafeColl(rc);

  thrust::for_each_n(policy, thrust::make_counting_iterator(0ul), n_targets,
                     [=] XGBOOST_DEVICE(std::size_t i) mutable {
                       out(i) = static_cast<float>(
                           CalcUnregularizedWeight(d_sum(i).GetGrad(), d_sum(i).GetHess()));
                     });
}
}  // namespace xgboost::tree::cuda_impl
