// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project

#ifndef KOKKOS_ALGORITHMS_UNITTESTS_TEST_RANDOM_HPP
#define KOKKOS_ALGORITHMS_UNITTESTS_TEST_RANDOM_HPP

#include <gtest/gtest.h>
#include <iostream>
#include <cstdlib>
#include <cstdio>
#include <Kokkos_Timer.hpp>
#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
import kokkos.dyn_rank_view;
import kokkos.random;
#else
#include <Kokkos_Core.hpp>
#include <Kokkos_DynRankView.hpp>
#include <Kokkos_Random.hpp>
#endif
#include <impl/Kokkos_Half_FloatingPointWrapper.hpp>
#include <cmath>
#include <chrono>
#include <vector>
#include <algorithm>
#include <numeric>

namespace Test {
namespace AlgoRandomImpl {

// This test runs the random number generators and uses some statistic tests to
// check the 'goodness' of the random numbers:
//    (i)   mean:         the mean is expected to be 0.5*RAND_MAX
//    (ii)  variance:     the variance is 1/3*mean*mean
//    (iii) covariance:   the covariance is 0
//    (iv)  1-tupledistr: the mean, variance and covariance of a 1D Histrogram
//    of random numbers (v)   3-tupledistr: the mean, variance and covariance of
//    a 3D Histrogram of random numbers

#define HIST_DIM3D 24
#define HIST_DIM1D (HIST_DIM3D * HIST_DIM3D * HIST_DIM3D)

struct RandomProperties {
  uint64_t count;
  double mean;
  double variance;
  double covariance;
  double min;
  double max;

  KOKKOS_INLINE_FUNCTION
  RandomProperties() {
    count      = 0;
    mean       = 0.0;
    variance   = 0.0;
    covariance = 0.0;
    min        = 1e64;
    max        = -1e64;
  }

  KOKKOS_INLINE_FUNCTION
  RandomProperties& operator+=(const RandomProperties& add) {
    count += add.count;
    mean += add.mean;
    variance += add.variance;
    covariance += add.covariance;
    min = add.min < min ? add.min : min;
    max = add.max > max ? add.max : max;
    return *this;
  }
};

template <class GeneratorPool, class Scalar>
struct test_random_functor {
  using rnd_type = typename GeneratorPool::generator_type;

  using value_type  = RandomProperties;
  using device_type = typename GeneratorPool::device_type;

  GeneratorPool rand_pool;
  const double mean;

  // NOTE (mfh 03 Nov 2014): Kokkos::rand::max() is supposed to define
  // an exclusive upper bound on the range of random numbers that
  // draw() can generate.  However, for the float specialization, some
  // implementations might violate this upper bound, due to rounding
  // error.  Just in case, we leave an extra space at the end of each
  // dimension, in the View types below.
  using type_1d =
      Kokkos::View<int[HIST_DIM1D + 1], typename GeneratorPool::device_type>;
  type_1d density_1d;
  using type_3d =
      Kokkos::View<int[HIST_DIM3D + 1][HIST_DIM3D + 1][HIST_DIM3D + 1],
                   typename GeneratorPool::device_type>;
  type_3d density_3d;

  test_random_functor(GeneratorPool rand_pool_, type_1d d1d, type_3d d3d)
      : rand_pool(rand_pool_),
        mean(0.5 * Kokkos::rand<rnd_type, Scalar>::max()),
        density_1d(d1d),
        density_3d(d3d) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(int /*i*/, RandomProperties& prop) const {
    using Kokkos::atomic_fetch_add;

    rnd_type rand_gen = rand_pool.get_state();
    for (int k = 0; k < 1024; ++k) {
      const Scalar tmp = Kokkos::rand<rnd_type, Scalar>::draw(rand_gen);
      prop.count++;
      prop.mean += tmp;
      prop.variance += (tmp - mean) * (tmp - mean);
      const Scalar tmp2 = Kokkos::rand<rnd_type, Scalar>::draw(rand_gen);
      prop.count++;
      prop.mean += tmp2;
      prop.variance += (tmp2 - mean) * (tmp2 - mean);
      prop.covariance += (tmp - mean) * (tmp2 - mean);
      const Scalar tmp3 = Kokkos::rand<rnd_type, Scalar>::draw(rand_gen);
      prop.count++;
      prop.mean += tmp3;
      prop.variance += (tmp3 - mean) * (tmp3 - mean);
      prop.covariance += (tmp2 - mean) * (tmp3 - mean);

      // NOTE (mfh 03 Nov 2014): Kokkos::rand::max() is supposed to
      // define an exclusive upper bound on the range of random
      // numbers that draw() can generate.  However, for the float
      // specialization, some implementations might violate this upper
      // bound, due to rounding error.  Just in case, we have left an
      // extra space at the end of each dimension of density_1d and
      // density_3d.
      //
      // Please note that those extra entries might not get counted in
      // the histograms.  However, if Kokkos::rand is broken and only
      // returns values of max(), the histograms will still catch this
      // indirectly, since none of the other values will be filled in.

      const Scalar theMax = Kokkos::rand<rnd_type, Scalar>::max();

      const uint64_t ind1_1d =
          static_cast<uint64_t>(1.0 * HIST_DIM1D * tmp / theMax);
      const uint64_t ind2_1d =
          static_cast<uint64_t>(1.0 * HIST_DIM1D * tmp2 / theMax);
      const uint64_t ind3_1d =
          static_cast<uint64_t>(1.0 * HIST_DIM1D * tmp3 / theMax);

      const uint64_t ind1_3d =
          static_cast<uint64_t>(1.0 * HIST_DIM3D * tmp / theMax);
      const uint64_t ind2_3d =
          static_cast<uint64_t>(1.0 * HIST_DIM3D * tmp2 / theMax);
      const uint64_t ind3_3d =
          static_cast<uint64_t>(1.0 * HIST_DIM3D * tmp3 / theMax);
      atomic_fetch_add(&density_1d(ind1_1d), 1);
      atomic_fetch_add(&density_1d(ind2_1d), 1);
      atomic_fetch_add(&density_1d(ind3_1d), 1);
      atomic_fetch_add(&density_3d(ind1_3d, ind2_3d, ind3_3d), 1);
    }
    rand_pool.free_state(rand_gen);
  }
};

template <class DeviceType>
struct test_histogram1d_functor {
  using value_type      = RandomProperties;
  using execution_space = typename DeviceType::execution_space;
  using memory_space    = typename DeviceType::memory_space;

  // NOTE (mfh 03 Nov 2014): Kokkos::rand::max() is supposed to define
  // an exclusive upper bound on the range of random numbers that
  // draw() can generate.  However, for the float specialization, some
  // implementations might violate this upper bound, due to rounding
  // error.  Just in case, we leave an extra space at the end of each
  // dimension, in the View type below.
  using type_1d   = Kokkos::View<int[HIST_DIM1D + 1], memory_space>;
  using size_type = typename type_1d::size_type;
  type_1d density_1d;
  double mean;

  test_histogram1d_functor(type_1d d1d, int num_draws)
      : density_1d(d1d), mean(1.0 * num_draws / HIST_DIM1D * 3) {}

  KOKKOS_INLINE_FUNCTION void operator()(const size_type i,
                                         RandomProperties& prop) const {
    const double count = density_1d(i);
    prop.mean += count;
    prop.variance += 1.0 * (count - mean) * (count - mean);
    // prop.covariance += 1.0*count*count;
    prop.min = count < prop.min ? count : prop.min;
    prop.max = count > prop.max ? count : prop.max;
    if (i < static_cast<size_type>(HIST_DIM1D - 1)) {
      prop.covariance += (count - mean) * (density_1d(i + 1) - mean);
    }
  }
};

template <class DeviceType>
struct test_histogram3d_functor {
  using value_type      = RandomProperties;
  using execution_space = typename DeviceType::execution_space;
  using memory_space    = typename DeviceType::memory_space;

  // NOTE (mfh 03 Nov 2014): Kokkos::rand::max() is supposed to define
  // an exclusive upper bound on the range of random numbers that
  // draw() can generate.  However, for the float specialization, some
  // implementations might violate this upper bound, due to rounding
  // error.  Just in case, we leave an extra space at the end of each
  // dimension, in the View type below.
  using type_3d =
      Kokkos::View<int[HIST_DIM3D + 1][HIST_DIM3D + 1][HIST_DIM3D + 1],
                   memory_space>;
  using size_type = typename type_3d::size_type;
  type_3d density_3d;
  double mean;

  test_histogram3d_functor(type_3d d3d, int num_draws)
      : density_3d(d3d), mean(1.0 * num_draws / HIST_DIM1D) {}

  KOKKOS_INLINE_FUNCTION void operator()(const size_type i,
                                         RandomProperties& prop) const {
    const double count = density_3d(
        i / (HIST_DIM3D * HIST_DIM3D),
        (i % (HIST_DIM3D * HIST_DIM3D)) / HIST_DIM3D, i % HIST_DIM3D);
    prop.mean += count;
    prop.variance += (count - mean) * (count - mean);
    if (i < static_cast<size_type>(HIST_DIM1D - 1)) {
      const double count_next =
          density_3d((i + 1) / (HIST_DIM3D * HIST_DIM3D),
                     ((i + 1) % (HIST_DIM3D * HIST_DIM3D)) / HIST_DIM3D,
                     (i + 1) % HIST_DIM3D);
      prop.covariance += (count - mean) * (count_next - mean);
    }
  }
};

//
// Templated test that uses the above functors.
//
template <class RandomGenerator, class Scalar>
struct test_random_scalar {
  using rnd_type = typename RandomGenerator::generator_type;

  test_random_scalar(
      typename test_random_functor<RandomGenerator, int>::type_1d& density_1d,
      typename test_random_functor<RandomGenerator, int>::type_3d& density_3d,
      RandomGenerator& pool, unsigned int num_draws) {
    using Kokkos::parallel_reduce;
    using std::cout;
    using std::endl;

    {
      cout << " -- Testing randomness properties" << endl;

      RandomProperties result;
      using functor_type = test_random_functor<RandomGenerator, Scalar>;
      parallel_reduce(num_draws / 1024,
                      functor_type(pool, density_1d, density_3d), result);

      // printf("Result: %lf %lf
      // %lf\n",result.mean/num_draws/3,result.variance/num_draws/3,result.covariance/num_draws/2);
      double tolerance       = 1.6 * std::sqrt(1.0 / num_draws);
      double mean_expect     = 0.5 * Kokkos::rand<rnd_type, Scalar>::max();
      double variance_expect = 1.0 / 3.0 * mean_expect * mean_expect;
      double mean_eps = mean_expect / (result.mean / num_draws / 3) - 1.0;
      double variance_eps =
          variance_expect / (result.variance / num_draws / 3) - 1.0;
      double covariance_eps =
          result.covariance / num_draws / 2 / variance_expect;
#if !KOKKOS_BHALF_T_IS_FLOAT
      if (!std::is_same_v<Scalar, Kokkos::Experimental::bhalf_t>) {
#endif
        EXPECT_LT(std::abs(mean_eps), tolerance);
        EXPECT_LT(std::abs(variance_eps), 1.5 * tolerance);
        EXPECT_LT(std::abs(covariance_eps), 2.0 * tolerance);
#if !KOKKOS_BHALF_T_IS_FLOAT
      }
#endif
    }
    {
      cout << " -- Testing 1-D histogram" << endl;

      RandomProperties result;
      using functor_type =
          test_histogram1d_functor<typename RandomGenerator::device_type>;
      parallel_reduce(HIST_DIM1D, functor_type(density_1d, num_draws), result);
      double mean_eps_expect       = 0.0001;
      double variance_eps_expect   = 0.07;
      double covariance_eps_expect = 0.06;
      double tolerance             = 6 * std::sqrt(1.0 / HIST_DIM1D);
      double mean_expect           = 1.0 * num_draws * 3 / HIST_DIM1D;
      double variance_expect =
          1.0 * num_draws * 3 / HIST_DIM1D * (1.0 - 1.0 / HIST_DIM1D);
      double covariance_expect = -1.0 * num_draws * 3 / HIST_DIM1D / HIST_DIM1D;
      double mean_eps          = mean_expect / (result.mean / HIST_DIM1D) - 1.0;
      double variance_eps =
          variance_expect / (result.variance / HIST_DIM1D) - 1.0;
      double covariance_eps =
          (result.covariance / HIST_DIM1D - covariance_expect) / mean_expect;

#if !KOKKOS_HALF_T_IS_FLOAT
      if (std::is_same_v<Scalar, Kokkos::Experimental::half_t>) {
        mean_eps_expect       = 0.0003;
        variance_eps_expect   = 1.0;
        covariance_eps_expect = 5.0e4;
      }
#endif

#if !KOKKOS_BHALF_T_IS_FLOAT
      if (!std::is_same_v<Scalar, Kokkos::Experimental::bhalf_t>) {
#endif
        EXPECT_LT(std::abs(mean_eps), mean_eps_expect);
        EXPECT_LT(std::abs(variance_eps), variance_eps_expect);
        EXPECT_LT(std::abs(covariance_eps), covariance_eps_expect);
#if !KOKKOS_BHALF_T_IS_FLOAT
      }
#endif

      cout << "Density 1D: " << mean_eps << " " << variance_eps << " "
           << (result.covariance / HIST_DIM1D / HIST_DIM1D) << " || "
           << tolerance << " " << result.min << " " << result.max << " || "
           << result.variance / HIST_DIM1D << " "
           << 1.0 * num_draws * 3 / HIST_DIM1D * (1.0 - 1.0 / HIST_DIM1D)
           << " || " << result.covariance / HIST_DIM1D << " "
           << -1.0 * num_draws * 3 / HIST_DIM1D / HIST_DIM1D << endl;
    }
    {
      cout << " -- Testing 3-D histogram" << endl;

      RandomProperties result;
      using functor_type =
          test_histogram3d_functor<typename RandomGenerator::device_type>;
      parallel_reduce(HIST_DIM1D, functor_type(density_3d, num_draws), result);

      double variance_factor = 1.2;
      double tolerance       = 6 * std::sqrt(1.0 / HIST_DIM1D);
      double mean_expect     = 1.0 * num_draws / HIST_DIM1D;
      double variance_expect =
          1.0 * num_draws / HIST_DIM1D * (1.0 - 1.0 / HIST_DIM1D);
      double covariance_expect = -1.0 * num_draws / HIST_DIM1D / HIST_DIM1D;
      double mean_eps          = mean_expect / (result.mean / HIST_DIM1D) - 1.0;
      double variance_eps =
          variance_expect / (result.variance / HIST_DIM1D) - 1.0;
      double covariance_eps =
          (result.covariance / HIST_DIM1D - covariance_expect) / mean_expect;

#if !KOKKOS_HALF_T_IS_FLOAT
      if (std::is_same_v<Scalar, Kokkos::Experimental::half_t>) {
        variance_factor = 7;
      }
#endif

#if !KOKKOS_BHALF_T_IS_FLOAT
      if (!std::is_same_v<Scalar, Kokkos::Experimental::bhalf_t>) {
#endif
        EXPECT_LT(std::abs(mean_eps), tolerance);
        EXPECT_LT(std::abs(variance_eps), variance_factor);
        EXPECT_LT(std::abs(covariance_eps), variance_factor);
#if !KOKKOS_BHALF_T_IS_FLOAT
      }
#endif

      cout << "Density 3D: " << mean_eps << " " << variance_eps << " "
           << result.covariance / HIST_DIM1D / HIST_DIM1D << " || " << tolerance
           << " " << result.min << " " << result.max << endl;
    }
  }
};

template <class RandomGenerator>
void test_random(unsigned int num_draws) {
  using std::cout;
  using std::endl;
  typename test_random_functor<RandomGenerator, int>::type_1d density_1d("D1d");
  typename test_random_functor<RandomGenerator, int>::type_3d density_3d("D3d");

  uint64_t ticks =
      std::chrono::high_resolution_clock::now().time_since_epoch().count();
  cout << "Test Seed:" << ticks << endl;

  RandomGenerator pool(ticks);

  cout << "Test Scalar=int" << endl;
  test_random_scalar<RandomGenerator, int> test_int(density_1d, density_3d,
                                                    pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=unsigned int" << endl;
  test_random_scalar<RandomGenerator, unsigned int> test_uint(
      density_1d, density_3d, pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=int64_t" << endl;
  test_random_scalar<RandomGenerator, int64_t> test_int64(
      density_1d, density_3d, pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=uint64_t" << endl;
  test_random_scalar<RandomGenerator, uint64_t> test_uint64(
      density_1d, density_3d, pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=half" << endl;
  test_random_scalar<RandomGenerator, Kokkos::Experimental::half_t> test_half(
      density_1d, density_3d, pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=bhalf" << endl;
  test_random_scalar<RandomGenerator, Kokkos::Experimental::bhalf_t> test_bhalf(
      density_1d, density_3d, pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=float" << endl;
  test_random_scalar<RandomGenerator, float> test_float(density_1d, density_3d,
                                                        pool, num_draws);
  deep_copy(density_1d, 0);
  deep_copy(density_3d, 0);

  cout << "Test Scalar=double" << endl;
  test_random_scalar<RandomGenerator, double> test_double(
      density_1d, density_3d, pool, num_draws);
}

template <class ExecutionSpace, class Pool>
struct TestDynRankView {
  using ReducerType      = Kokkos::MinMax<double, Kokkos::HostSpace>;
  using ReducerValueType = typename ReducerType::value_type;

  Kokkos::DynRankView<double, ExecutionSpace> A;

  TestDynRankView(int n) : A("a", n) {}

  KOKKOS_FUNCTION void operator()(int i, ReducerValueType& update) const {
    if (A(i) < update.min_val) update.min_val = A(i);
    if (A(i) > update.max_val) update.max_val = A(i);
  }

  void run() {
    Pool random(13);
    double min = 10.;
    double max = 100.;
    ExecutionSpace exec;
    Kokkos::fill_random(exec, A, random, min, max);

    ReducerValueType val;
    Kokkos::parallel_reduce(
        Kokkos::RangePolicy<ExecutionSpace>(exec, 0, A.size()), *this,
        ReducerType(val));

    exec.fence();
    ASSERT_GE(val.min_val, min);
    ASSERT_LE(val.max_val, max);
  }
};

template <class ExecutionSpace, class GeneratorPool>
struct generate_random_stream {
  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>;

  ViewType vals;
  GeneratorPool rand_pool;
  int samples;
  bool use_specific_gen;

  generate_random_stream(ViewType vals_, GeneratorPool rand_pool_, int samples_,
                         bool use_specific_gen_)
      : vals(vals_),
        rand_pool(rand_pool_),
        samples(samples_),
        use_specific_gen(use_specific_gen_) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(int i) const {
    typename GeneratorPool::generator_type rand_gen =
        use_specific_gen ? rand_pool.get_state(i) : rand_pool.get_state();

    for (int k = 0; k < samples; k++) vals(i, k) = rand_gen.urand64();

    rand_pool.free_state(rand_gen);
  }
};

// NOTE: this doesn't test the statistical independence of multiple streams
// generated by a Random pool, it only tests for complete duplicates.
template <class ExecutionSpace, class Pool>
void test_duplicate_stream() {
  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>;

  // Heuristic to create a "large enough" number of streams.
  int n_streams = ExecutionSpace{}.concurrency() * 4;
  int samples   = 8;

  Pool rand_pool(42);
  ViewType vals_d("Vals", n_streams, samples);

  Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(0, n_streams),
                       generate_random_stream<ExecutionSpace, Pool>(
                           vals_d, rand_pool, samples, false));

  auto vals_h =
      Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, vals_d);

  /*
  To quickly find streams that are identical, we sort them by the first number,
  if that's equal then the second and so on. We then test each neighbor pair
  for duplicates.
  */
  std::vector<size_t> indices(n_streams);
  std::iota(indices.begin(), indices.end(), 0);

  auto comparator = [&](size_t i, size_t j) {
    for (int k = 0; k < samples; k++) {
      if (vals_h(i, k) != vals_h(j, k)) return vals_h(i, k) < vals_h(j, k);
    }
    return false;
  };
  std::sort(indices.begin(), indices.end(), comparator);

  for (int i = 0; i < n_streams - 1; i++) {
    int idx1 = indices[i];
    int idx2 = indices[i + 1];

    int k = 0;
    while (k < samples && vals_h(idx1, k) == vals_h(idx2, k)) k++;
    ASSERT_LT(k, samples) << "Duplicate streams found";
  }
}

template <class ExecutionSpace, class GeneratorPool>
struct compare_random_streams {
  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>;

  ViewType vals;
  GeneratorPool rand_pool;
  int samples;

  compare_random_streams(ViewType vals_, GeneratorPool rand_pool_, int samples_)
      : vals(vals_), rand_pool(rand_pool_), samples(samples_) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(int i, std::size_t& mismatches) const {
    // this is problematic: on a GPU when launching with more than a single
    // thread the generator returned is in principle random, using atomic
    // locks to acquire a state.
    // but we only launch it with a single thread so its ok.
    typename GeneratorPool::generator_type rand_gen = rand_pool.get_state(i);

    for (int k = 0; k < samples; k++)
      if (vals(i, k) != rand_gen.urand64()) mismatches++;

    rand_pool.free_state(rand_gen);
  }
};

template <class ExecutionSpace, class Pool, class... Args>
void test_async_initialization(Args... args) {
  // using 2D View here to reuse functions from other test
  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>;

  int samples = 123456;

  // use default execution space instance to generate reference values
  Pool rand_pool_A(args...);

  ViewType vals_d("Vals", 1, samples);
  // create two, distinct ExecutionSpace instances
  auto instances =
      Kokkos::Experimental::partition_space(ExecutionSpace{}, 1, 1);
  // use first instance to initialize values of vals_d
  Kokkos::parallel_for(
      Kokkos::RangePolicy<ExecutionSpace>(instances.at(0), 0, 1),
      generate_random_stream<ExecutionSpace, Pool>(vals_d, rand_pool_A, samples,
                                                   true));
  instances.at(0).fence();

  // use second instance to initialize another Pool using the same seed
  Pool rand_pool_B(instances.at(1), args...);

  std::size_t mismatches;
  // compare values in stream of rand_pool_B with vals_d
  Kokkos::parallel_reduce(
      Kokkos::RangePolicy<ExecutionSpace>(instances.at(1), 0, 1),
      compare_random_streams<ExecutionSpace, Pool>(vals_d, rand_pool_B,
                                                   samples),
      mismatches);

  EXPECT_EQ(mismatches, 0lu) << "Stream from async constructed pool does not "
                                "match stream from default constructed pool";
}

template <class ExecutionSpace, class Pool>
void test_offset_stream() {
  // using 2D View here to reuse functions from other test
  using ViewType = Kokkos::View<uint64_t**, ExecutionSpace>;

  // arbitrary values for the test
  const int seed = 8432348;

  const int num_streams         = 123456;
  const int samples_per_streams = 100;

  // use default execution space instance to generate reference values
  Pool ref_pool(seed, num_streams);

  ViewType ref_vals("Vals", num_streams, samples_per_streams);
  Kokkos::parallel_for(Kokkos::RangePolicy<ExecutionSpace>(0, num_streams),
                       generate_random_stream<ExecutionSpace, Pool>(
                           ref_vals, ref_pool, samples_per_streams, true));
  Kokkos::fence();

  // create two, distinct ExecutionSpace instances
  auto instances =
      Kokkos::Experimental::partition_space(ExecutionSpace{}, 1, 1);

  // divide number of streams in halves
  const int streams_half = num_streams / 2;

  // give a half to a pool on first instance
  Pool rand_pool_A(instances.at(0), seed, streams_half);

  // and the other half to a pool on second instance with an offset
  const int offset = num_streams - streams_half;
  Pool rand_pool_B(instances.at(1), seed, offset, streams_half);

  std::size_t mismatches;
  // compare values in stream of rand_pool_A with ref_vals(0, streams_half)
  Kokkos::parallel_reduce(
      Kokkos::RangePolicy<ExecutionSpace>(instances.at(0), 0, streams_half),
      compare_random_streams<ExecutionSpace, Pool>(ref_vals, rand_pool_A,
                                                   samples_per_streams),
      mismatches);

  EXPECT_EQ(mismatches, 0lu)
      << "First half of reference streams doesn't match streams of pool A";

  // compare values in stream of rand_pool_B with ref_vals(streams_half,
  // num_streams)
  auto last_half_range = std::make_pair(offset, num_streams);
  auto sub_refs = Kokkos::subview(ref_vals, last_half_range, Kokkos::ALL);

  Kokkos::parallel_reduce(
      Kokkos::RangePolicy<ExecutionSpace>(instances.at(1), 0, streams_half),
      compare_random_streams<ExecutionSpace, Pool>(sub_refs, rand_pool_B,
                                                   samples_per_streams),
      mismatches);

  EXPECT_EQ(mismatches, 0lu)
      << "Second half of reference streams doesn't match streams of pool B";
}

// Test rand/rand64 with ranges that previously caused signed-overflow UB.
template <class GeneratorPool>
void test_rand_range_overflow() {
  using exec_space = typename GeneratorPool::device_type::execution_space;
  using gen_type   = typename GeneratorPool::generator_type;

  constexpr int32_t start32 = -gen_type::MAX_RAND - 1;
  constexpr int32_t end32   = gen_type::MAX_RAND;
  constexpr int64_t start64 = -gen_type::MAX_RAND64 - 1;
  constexpr int64_t end64   = gen_type::MAX_RAND64;

  GeneratorPool pool(42);

  int64_t n_nonmin32 = 0;
  Kokkos::parallel_reduce(
      "test_rand_range_overflow_32", Kokkos::RangePolicy<exec_space>(0, 1000),
      KOKKOS_LAMBDA(int /*i*/, int64_t& n) {
        auto gen = pool.get_state();
        for (int k = 0; k < 64; ++k)
          if (gen.rand(start32, end32) != start32) ++n;
        pool.free_state(gen);
      },
      n_nonmin32);
  EXPECT_GT(n_nonmin32, 0) << "rand(INT_MIN,INT_MAX) always returned INT_MIN";

  int64_t n_nonmin64 = 0;
  Kokkos::parallel_reduce(
      "test_rand_range_overflow_64", Kokkos::RangePolicy<exec_space>(0, 1000),
      KOKKOS_LAMBDA(int /*i*/, int64_t& n) {
        auto gen = pool.get_state();
        for (int k = 0; k < 64; ++k)
          if (gen.rand64(start64, end64) != start64) ++n;
        pool.free_state(gen);
      },
      n_nonmin64);
  EXPECT_GT(n_nonmin64, 0)
      << "rand64(INT64_MIN,INT64_MAX) always returned INT64_MIN";
}

}  // namespace AlgoRandomImpl

TEST(TEST_CATEGORY, Random_XorShift64) {
  using ExecutionSpace = TEST_EXECSPACE;

#if defined(KOKKOS_ENABLE_SYCL) || defined(KOKKOS_ENABLE_CUDA) || \
    defined(KOKKOS_ENABLE_HIP)
  const int num_draws = 132141141;
#else  // SERIAL, HPX, OPENMP
  const int num_draws = 10240000;
#endif
  AlgoRandomImpl::test_random<Kokkos::Random_XorShift64_Pool<ExecutionSpace>>(
      num_draws);
  AlgoRandomImpl::test_random<Kokkos::Random_XorShift64_Pool<
      Kokkos::Device<ExecutionSpace, typename ExecutionSpace::memory_space>>>(
      num_draws);
  AlgoRandomImpl::TestDynRankView<
      ExecutionSpace, Kokkos::Random_XorShift64_Pool<ExecutionSpace>>(10000)
      .run();
}

TEST(TEST_CATEGORY, Random_XorShift1024_0) {
  using ExecutionSpace = TEST_EXECSPACE;

#if defined(KOKKOS_ENABLE_SYCL) || defined(KOKKOS_ENABLE_CUDA) || \
    defined(KOKKOS_ENABLE_HIP)
  const int num_draws = 52428813;
#else  // SERIAL, HPX, OPENMP
  const int num_draws = 10130144;
#endif
  AlgoRandomImpl::test_random<Kokkos::Random_XorShift1024_Pool<ExecutionSpace>>(
      num_draws);
  AlgoRandomImpl::test_random<Kokkos::Random_XorShift1024_Pool<
      Kokkos::Device<ExecutionSpace, typename ExecutionSpace::memory_space>>>(
      num_draws);
  AlgoRandomImpl::TestDynRankView<
      ExecutionSpace, Kokkos::Random_XorShift1024_Pool<ExecutionSpace>>(10000)
      .run();
}

TEST(TEST_CATEGORY, Random_SFC64) {
  using ExecutionSpace = TEST_EXECSPACE;

#if defined(KOKKOS_ENABLE_SYCL)
  GTEST_SKIP() << "Failing on Intel GPUs";  // FIXME_SYCL
  const int num_draws = 14203454;
#elif defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP)
  const int num_draws = 132141141;
#else  // SERIAL, HPX, OPENMP
  const int num_draws = 10240000;
#endif
  AlgoRandomImpl::test_random<Kokkos::Random_SFC64_Pool<ExecutionSpace>>(
      num_draws);
  AlgoRandomImpl::test_random<Kokkos::Random_SFC64_Pool<
      Kokkos::Device<ExecutionSpace, typename ExecutionSpace::memory_space>>>(
      num_draws);
  AlgoRandomImpl::TestDynRankView<ExecutionSpace,
                                  Kokkos::Random_SFC64_Pool<ExecutionSpace>>(
      10000)
      .run();
}

TEST(TEST_CATEGORY, Random_XorShift64_rand_range_overflow) {
  AlgoRandomImpl::test_rand_range_overflow<
      Kokkos::Random_XorShift64_Pool<TEST_EXECSPACE>>();
}

TEST(TEST_CATEGORY, Random_XorShift1024_rand_range_overflow) {
  AlgoRandomImpl::test_rand_range_overflow<
      Kokkos::Random_XorShift1024_Pool<TEST_EXECSPACE>>();
}

TEST(TEST_CATEGORY, Random_SFC64_rand_range_overflow) {
#if defined(KOKKOS_ENABLE_SYCL)
  GTEST_SKIP() << "Failing on Intel GPUs";  // FIXME_SYCL
#endif
  AlgoRandomImpl::test_rand_range_overflow<
      Kokkos::Random_SFC64_Pool<TEST_EXECSPACE>>();
}

TEST(TEST_CATEGORY, Multi_streams) {
  using ExecutionSpace = TEST_EXECSPACE;
#if defined(KOKKOS_ENABLE_SYCL) && defined(KOKKOS_IMPL_ARCH_NVIDIA_GPU)
  if constexpr (std::is_same_v<ExecutionSpace, Kokkos::SYCL>) {
    GTEST_SKIP() << "Failing on NVIDIA GPUs";  // FIXME_SYCL
  }
#endif

  using Pool64   = Kokkos::Random_XorShift64_Pool<ExecutionSpace>;
  using Pool1024 = Kokkos::Random_XorShift1024_Pool<ExecutionSpace>;

  using SFC64Pool = Kokkos::Random_SFC64_Pool<ExecutionSpace>;

  AlgoRandomImpl::test_duplicate_stream<ExecutionSpace, Pool64>();
  AlgoRandomImpl::test_duplicate_stream<ExecutionSpace, Pool1024>();

  AlgoRandomImpl::test_duplicate_stream<ExecutionSpace, SFC64Pool>();

  // Test with construction from seed
  AlgoRandomImpl::test_async_initialization<ExecutionSpace, Pool64>(42);
  AlgoRandomImpl::test_async_initialization<ExecutionSpace, Pool1024>(42);

  // Test with construction from seed and num_states
  AlgoRandomImpl::test_async_initialization<ExecutionSpace, Pool64>(42, 1);
  AlgoRandomImpl::test_async_initialization<ExecutionSpace, Pool1024>(42, 1);

  // FIXME : The tests fail because the get_state(state_idx) method does not
  // acquire state locks.
  // AlgoRandomImpl::test_async_initialization<ExecutionSpace, SFC64Pool>(42);
  // AlgoRandomImpl::test_async_initialization<ExecutionSpace, SFC64Pool>(42,
  // 1);
}

TEST(TEST_CATEGORY, Offset_streams) {
  using ExecutionSpace = TEST_EXECSPACE;
  using SFC64Pool      = Kokkos::Random_SFC64_Pool<ExecutionSpace>;

  AlgoRandomImpl::test_offset_stream<ExecutionSpace, SFC64Pool>();
}

}  // namespace Test
#endif