// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#include "main.h"

#include <Eigen/Tensor>

template <typename Scalar>
static void test_default() {
  Tensor<Scalar, 1> vec(6);

  // Fixme: we should check that the generated numbers follow a uniform
  // distribution instead.
  // For low-precision types (half, bfloat16), the RNG has limited distinct
  // values (e.g. 128 for bfloat16), so adjacent collisions are possible.
  // Retry a few times to avoid spurious failures.
  bool all_distinct = false;
  for (int attempt = 0; attempt < 10 && !all_distinct; ++attempt) {
    vec.setRandom();
    all_distinct = true;
    for (int i = 1; i < 6; ++i) {
      if (vec(i) == vec(i - 1)) {
        all_distinct = false;
        break;
      }
    }
  }
  VERIFY(all_distinct);
}

template <typename Scalar>
static void test_normal() {
  Tensor<Scalar, 1> vec(6);

  // Fixme: we should check that the generated numbers follow a gaussian
  // distribution instead.
  bool all_distinct = false;
  for (int attempt = 0; attempt < 10 && !all_distinct; ++attempt) {
    vec.template setRandom<Eigen::internal::NormalRandomGenerator<Scalar>>();
    all_distinct = true;
    for (int i = 1; i < 6; ++i) {
      if (vec(i) == vec(i - 1)) {
        all_distinct = false;
        break;
      }
    }
  }
  VERIFY(all_distinct);
}

struct MyGenerator {
  MyGenerator() {}
  MyGenerator(const MyGenerator&) {}

  // Return a random value to be used.  "element_location" is the
  // location of the entry to set in the tensor, it can typically
  // be ignored.
  int operator()(Eigen::DenseIndex element_location, Eigen::DenseIndex /*unused*/ = 0) const {
    return static_cast<int>(3 * element_location);
  }

  // Same as above but generates several numbers at a time.
  internal::packet_traits<int>::type packetOp(Eigen::DenseIndex packet_location,
                                              Eigen::DenseIndex /*unused*/ = 0) const {
    const int packetSize = internal::packet_traits<int>::size;
    EIGEN_ALIGN_MAX int values[packetSize];
    for (int i = 0; i < packetSize; ++i) {
      values[i] = static_cast<int>(3 * (packet_location + i));
    }
    return internal::pload<typename internal::packet_traits<int>::type>(values);
  }
};

static void test_custom() {
  Tensor<int, 1> vec(6);
  vec.setRandom<MyGenerator>();

  for (int i = 0; i < 6; ++i) {
    VERIFY_IS_EQUAL(vec(i), 3 * i);
  }
}

EIGEN_DECLARE_TEST(tensor_random) {
  CALL_SUBTEST((test_default<float>()));
  CALL_SUBTEST((test_normal<float>()));
  CALL_SUBTEST((test_default<double>()));
  CALL_SUBTEST((test_normal<double>()));
  CALL_SUBTEST((test_default<Eigen::half>()));
  CALL_SUBTEST((test_normal<Eigen::half>()));
  CALL_SUBTEST((test_default<Eigen::bfloat16>()));
  CALL_SUBTEST((test_normal<Eigen::bfloat16>()));
  CALL_SUBTEST(test_custom());
}