/*
All modification made by Intel Corporation: © 2016 Intel Corporation

All contributions by the University of California:
Copyright (c) 2014, 2015, The Regents of the University of California (Regents)
All rights reserved.

All other contributions:
Copyright (c) 2014, 2015, the respective contributors
All rights reserved.
For the list of contributors go to https://github.com/BVLC/caffe/blob/master/CONTRIBUTORS.md


Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of Intel Corporation nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <cmath>

#include "gtest/gtest.h"

#include "caffe/common.hpp"
#include "caffe/syncedmem.hpp"
#include "caffe/util/math_functions.hpp"

#include "caffe/test/test_caffe_main.hpp"

namespace caffe {

template <typename Dtype>
class RandomNumberGeneratorTest : public ::testing::Test {
 protected:
  RandomNumberGeneratorTest()
     : mean_bound_multiplier_(3.8),  // ~99.99% confidence for test failure.
       sample_size_(10000),
       seed_(1701),
       data_(new SyncedMemory(sample_size_ * sizeof(Dtype))),
       data_2_(new SyncedMemory(sample_size_ * sizeof(Dtype))),
       int_data_(new SyncedMemory(sample_size_ * sizeof(int))),
       int_data_2_(new SyncedMemory(sample_size_ * sizeof(int))) {}

  virtual void SetUp() {
    Caffe::set_random_seed(this->seed_);
  }

  Dtype sample_mean(const Dtype* const seqs, const int sample_size) {
    Dtype sum = 0;
    for (int i = 0; i < sample_size; ++i) {
      sum += seqs[i];
    }
    return sum / sample_size;
  }

  Dtype sample_mean(const Dtype* const seqs) {
    return sample_mean(seqs, sample_size_);
  }

  Dtype sample_mean(const int* const seqs, const int sample_size) {
    Dtype sum = 0;
    for (int i = 0; i < sample_size; ++i) {
      sum += Dtype(seqs[i]);
    }
    return sum / sample_size;
  }

  Dtype sample_mean(const int* const seqs) {
    return sample_mean(seqs, sample_size_);
  }

  Dtype mean_bound(const Dtype std, const int sample_size) {
    return mean_bound_multiplier_ * std / sqrt(static_cast<Dtype>(sample_size));
  }

  Dtype mean_bound(const Dtype std) {
    return mean_bound(std, sample_size_);
  }

  void RngGaussianFill(const Dtype mu, const Dtype sigma, void* cpu_data) {
    Dtype* rng_data = static_cast<Dtype*>(cpu_data);
    caffe_rng_gaussian(sample_size_, mu, sigma, rng_data);
  }

  void RngGaussianChecks(const Dtype mu, const Dtype sigma,
                         const void* cpu_data, const Dtype sparse_p = 0) {
    const Dtype* rng_data = static_cast<const Dtype*>(cpu_data);
    const Dtype true_mean = mu;
    const Dtype true_std = sigma;
    // Check that sample mean roughly matches true mean.
    const Dtype bound = this->mean_bound(true_std);
    const Dtype sample_mean = this->sample_mean(
        static_cast<const Dtype*>(cpu_data));
    EXPECT_NEAR(sample_mean, true_mean, bound);
    // Check that roughly half the samples are above the true mean.
    int num_above_mean = 0;
    int num_below_mean = 0;
    int num_mean = 0;
    int num_nan = 0;
    for (int i = 0; i < sample_size_; ++i) {
      if (rng_data[i] > true_mean) {
        ++num_above_mean;
      } else if (rng_data[i] < true_mean) {
        ++num_below_mean;
      } else if (rng_data[i] == true_mean) {
        ++num_mean;
      } else {
        ++num_nan;
      }
    }
    EXPECT_EQ(0, num_nan);
    if (sparse_p == Dtype(0)) {
      EXPECT_EQ(0, num_mean);
    }
    const Dtype sample_p_above_mean =
        static_cast<Dtype>(num_above_mean) / sample_size_;
    const Dtype bernoulli_p = (1 - sparse_p) * 0.5;
    const Dtype bernoulli_std = sqrt(bernoulli_p * (1 - bernoulli_p));
    const Dtype bernoulli_bound = this->mean_bound(bernoulli_std);
    EXPECT_NEAR(bernoulli_p, sample_p_above_mean, bernoulli_bound);
  }

  void RngUniformFill(const Dtype lower, const Dtype upper, void* cpu_data) {
    CHECK_GE(upper, lower);
    Dtype* rng_data = static_cast<Dtype*>(cpu_data);
    caffe_rng_uniform(sample_size_, lower, upper, rng_data);
  }

  void RngUniformChecks(const Dtype lower, const Dtype upper,
                        const void* cpu_data, const Dtype sparse_p = 0) {
    const Dtype* rng_data = static_cast<const Dtype*>(cpu_data);
    const Dtype true_mean = (lower + upper) / 2;
    const Dtype true_std = (upper - lower) / sqrt(12);
    // Check that sample mean roughly matches true mean.
    const Dtype bound = this->mean_bound(true_std);
    const Dtype sample_mean = this->sample_mean(rng_data);
    EXPECT_NEAR(sample_mean, true_mean, bound);
    // Check that roughly half the samples are above the true mean, and none are
    // above upper or below lower.
    int num_above_mean = 0;
    int num_below_mean = 0;
    int num_mean = 0;
    int num_nan = 0;
    int num_above_upper = 0;
    int num_below_lower = 0;
    for (int i = 0; i < sample_size_; ++i) {
      if (rng_data[i] > true_mean) {
        ++num_above_mean;
      } else if (rng_data[i] < true_mean) {
        ++num_below_mean;
      } else if (rng_data[i] == true_mean) {
        ++num_mean;
      } else {
        ++num_nan;
      }
      if (rng_data[i] > upper) {
        ++num_above_upper;
      } else if (rng_data[i] < lower) {
        ++num_below_lower;
      }
    }
    EXPECT_EQ(0, num_nan);
    EXPECT_EQ(0, num_above_upper);
    EXPECT_EQ(0, num_below_lower);
    if (sparse_p == Dtype(0)) {
      EXPECT_EQ(0, num_mean);
    }
    const Dtype sample_p_above_mean =
        static_cast<Dtype>(num_above_mean) / sample_size_;
    const Dtype bernoulli_p = (1 - sparse_p) * 0.5;
    const Dtype bernoulli_std = sqrt(bernoulli_p * (1 - bernoulli_p));
    const Dtype bernoulli_bound = this->mean_bound(bernoulli_std);
    EXPECT_NEAR(bernoulli_p, sample_p_above_mean, bernoulli_bound);
  }

  void RngBernoulliFill(const Dtype p, void* cpu_data) {
    int* rng_data = static_cast<int*>(cpu_data);
    caffe_rng_bernoulli(sample_size_, p, rng_data);
  }

  void RngBernoulliChecks(const Dtype p, const void* cpu_data) {
    const int* rng_data = static_cast<const int*>(cpu_data);
    const Dtype true_mean = p;
    const Dtype true_std = sqrt(p * (1 - p));
    const Dtype bound = this->mean_bound(true_std);
    const Dtype sample_mean = this->sample_mean(rng_data);
    EXPECT_NEAR(sample_mean, true_mean, bound);
  }

#ifndef CPU_ONLY

  void RngGaussianFillGPU(const Dtype mu, const Dtype sigma, void* gpu_data) {
    Dtype* rng_data = static_cast<Dtype*>(gpu_data);
    caffe_gpu_rng_gaussian(sample_size_, mu, sigma, rng_data);
  }

  void RngUniformFillGPU(const Dtype lower, const Dtype upper, void* gpu_data) {
    CHECK_GE(upper, lower);
    Dtype* rng_data = static_cast<Dtype*>(gpu_data);
    caffe_gpu_rng_uniform(sample_size_, lower, upper, rng_data);
  }

  // Fills with uniform integers in [0, UINT_MAX] using 2 argument form of
  // caffe_gpu_rng_uniform.
  void RngUniformIntFillGPU(void* gpu_data) {
    unsigned int* rng_data = static_cast<unsigned int*>(gpu_data);
    caffe_gpu_rng_uniform(sample_size_, rng_data);
  }

#endif

  int num_above_mean;
  int num_below_mean;

  Dtype mean_bound_multiplier_;

  size_t sample_size_;
  uint32_t seed_;

  shared_ptr<SyncedMemory> data_;
  shared_ptr<SyncedMemory> data_2_;
  shared_ptr<SyncedMemory> int_data_;
  shared_ptr<SyncedMemory> int_data_2_;
};

TYPED_TEST_CASE(RandomNumberGeneratorTest, TestDtypes);

TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussian) {
  const TypeParam mu = 0;
  const TypeParam sigma = 1;
  void* gaussian_data = this->data_->mutable_cpu_data();
  this->RngGaussianFill(mu, sigma, gaussian_data);
  this->RngGaussianChecks(mu, sigma, gaussian_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussian2) {
  const TypeParam mu = -2;
  const TypeParam sigma = 3;
  void* gaussian_data = this->data_->mutable_cpu_data();
  this->RngGaussianFill(mu, sigma, gaussian_data);
  this->RngGaussianChecks(mu, sigma, gaussian_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniform) {
  const TypeParam lower = 0;
  const TypeParam upper = 1;
  void* uniform_data = this->data_->mutable_cpu_data();
  this->RngUniformFill(lower, upper, uniform_data);
  this->RngUniformChecks(lower, upper, uniform_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniform2) {
  const TypeParam lower = -7.3;
  const TypeParam upper = -2.3;
  void* uniform_data = this->data_->mutable_cpu_data();
  this->RngUniformFill(lower, upper, uniform_data);
  this->RngUniformChecks(lower, upper, uniform_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngBernoulli) {
  const TypeParam p = 0.3;
  void* bernoulli_data = this->int_data_->mutable_cpu_data();
  this->RngBernoulliFill(p, bernoulli_data);
  this->RngBernoulliChecks(p, bernoulli_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngBernoulli2) {
  const TypeParam p = 0.9;
  void* bernoulli_data = this->int_data_->mutable_cpu_data();
  this->RngBernoulliFill(p, bernoulli_data);
  this->RngBernoulliChecks(p, bernoulli_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussianTimesGaussian) {
  const TypeParam mu = 0;
  const TypeParam sigma = 1;

  // Sample from 0 mean Gaussian.
  TypeParam* gaussian_data_1 =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  this->RngGaussianFill(mu, sigma, gaussian_data_1);

  // Sample from 0 mean Gaussian again.
  TypeParam* gaussian_data_2 =
      static_cast<TypeParam*>(this->data_2_->mutable_cpu_data());
  this->RngGaussianFill(mu, sigma, gaussian_data_2);

  // Multiply Gaussians.
  for (int i = 0; i < this->sample_size_; ++i) {
    gaussian_data_1[i] *= gaussian_data_2[i];
  }

  // Check that result has mean 0.
  TypeParam mu_product = pow(mu, 2);
  TypeParam sigma_product = sqrt(pow(sigma, 2) / 2);
  this->RngGaussianChecks(mu_product, sigma_product, gaussian_data_1);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniformTimesUniform) {
  // Sample from Uniform on [-2, 2].
  const TypeParam lower_1 = -2;
  const TypeParam upper_1 = -lower_1;
  TypeParam* uniform_data_1 =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  this->RngUniformFill(lower_1, upper_1, uniform_data_1);

  // Sample from Uniform on [-3, 3].
  const TypeParam lower_2 = -3;
  const TypeParam upper_2 = -lower_2;
  TypeParam* uniform_data_2 =
      static_cast<TypeParam*>(this->data_2_->mutable_cpu_data());
  this->RngUniformFill(lower_2, upper_2, uniform_data_2);

  // Multiply Uniforms.
  for (int i = 0; i < this->sample_size_; ++i) {
    uniform_data_1[i] *= uniform_data_2[i];
  }

  // Check that result does not violate checked properties of Uniform on [-6, 6]
  // (though it is not actually uniformly distributed).
  const TypeParam lower_prod = lower_1 * upper_2;
  const TypeParam upper_prod = -lower_prod;
  this->RngUniformChecks(lower_prod, upper_prod, uniform_data_1);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussianTimesBernoulli) {
  // Sample from 0 mean Gaussian.
  const TypeParam mu = 0;
  const TypeParam sigma = 1;
  TypeParam* gaussian_data =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  this->RngGaussianFill(mu, sigma, gaussian_data);

  // Sample from Bernoulli with p = 0.3.
  const TypeParam bernoulli_p = 0.3;
  int* bernoulli_data =
      static_cast<int*>(this->int_data_->mutable_cpu_data());
  this->RngBernoulliFill(bernoulli_p, bernoulli_data);

  // Multiply Gaussian by Bernoulli.
  for (int i = 0; i < this->sample_size_; ++i) {
    gaussian_data[i] *= bernoulli_data[i];
  }

  // Check that result does not violate checked properties of sparsified
  // Gaussian (though it is not actually a Gaussian).
  this->RngGaussianChecks(mu, sigma, gaussian_data, 1 - bernoulli_p);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniformTimesBernoulli) {
  // Sample from Uniform on [-1, 1].
  const TypeParam lower = -1;
  const TypeParam upper = 1;
  TypeParam* uniform_data =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  this->RngUniformFill(lower, upper, uniform_data);

  // Sample from Bernoulli with p = 0.3.
  const TypeParam bernoulli_p = 0.3;
  int* bernoulli_data =
      static_cast<int*>(this->int_data_->mutable_cpu_data());
  this->RngBernoulliFill(bernoulli_p, bernoulli_data);

  // Multiply Uniform by Bernoulli.
  for (int i = 0; i < this->sample_size_; ++i) {
    uniform_data[i] *= bernoulli_data[i];
  }

  // Check that result does not violate checked properties of sparsified
  // Uniform on [-1, 1] (though it is not actually uniformly distributed).
  this->RngUniformChecks(lower, upper, uniform_data, 1 - bernoulli_p);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngBernoulliTimesBernoulli) {
  // Sample from Bernoulli with p = 0.5.
  const TypeParam p_a = 0.5;
  int* bernoulli_data_a =
      static_cast<int*>(this->int_data_->mutable_cpu_data());
  this->RngBernoulliFill(p_a, bernoulli_data_a);

  // Sample from Bernoulli with p = 0.3.
  const TypeParam p_b = 0.3;
  int* bernoulli_data_b =
      static_cast<int*>(this->int_data_2_->mutable_cpu_data());
  this->RngBernoulliFill(p_b, bernoulli_data_b);

  // Multiply Bernoullis.
  for (int i = 0; i < this->sample_size_; ++i) {
    bernoulli_data_a[i] *= bernoulli_data_b[i];
  }
  int num_ones = 0;
  for (int i = 0; i < this->sample_size_; ++i) {
    if (bernoulli_data_a[i] != TypeParam(0)) {
      EXPECT_EQ(TypeParam(1), bernoulli_data_a[i]);
      ++num_ones;
    }
  }

  // Check that resulting product has roughly p_a * p_b ones.
  const TypeParam sample_p = this->sample_mean(bernoulli_data_a);
  const TypeParam true_mean = p_a * p_b;
  const TypeParam true_std = sqrt(true_mean * (1 - true_mean));
  const TypeParam bound = this->mean_bound(true_std);
  EXPECT_NEAR(true_mean, sample_p, bound);
}

#ifndef CPU_ONLY

TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussianGPU) {
  const TypeParam mu = 0;
  const TypeParam sigma = 1;
  void* gaussian_gpu_data = this->data_->mutable_gpu_data();
  this->RngGaussianFillGPU(mu, sigma, gaussian_gpu_data);
  const void* gaussian_data = this->data_->cpu_data();
  this->RngGaussianChecks(mu, sigma, gaussian_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussian2GPU) {
  const TypeParam mu = -2;
  const TypeParam sigma = 3;
  void* gaussian_gpu_data = this->data_->mutable_gpu_data();
  this->RngGaussianFillGPU(mu, sigma, gaussian_gpu_data);
  const void* gaussian_data = this->data_->cpu_data();
  this->RngGaussianChecks(mu, sigma, gaussian_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniformGPU) {
  const TypeParam lower = 0;
  const TypeParam upper = 1;
  void* uniform_gpu_data = this->data_->mutable_gpu_data();
  this->RngUniformFillGPU(lower, upper, uniform_gpu_data);
  const void* uniform_data = this->data_->cpu_data();
  this->RngUniformChecks(lower, upper, uniform_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniform2GPU) {
  const TypeParam lower = -7.3;
  const TypeParam upper = -2.3;
  void* uniform_gpu_data = this->data_->mutable_gpu_data();
  this->RngUniformFillGPU(lower, upper, uniform_gpu_data);
  const void* uniform_data = this->data_->cpu_data();
  this->RngUniformChecks(lower, upper, uniform_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniformIntGPU) {
  unsigned int* uniform_uint_gpu_data =
      static_cast<unsigned int*>(this->int_data_->mutable_gpu_data());
  this->RngUniformIntFillGPU(uniform_uint_gpu_data);
  const unsigned int* uniform_uint_data =
      static_cast<const unsigned int*>(this->int_data_->cpu_data());
  TypeParam* uniform_data =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  for (int i = 0; i < this->sample_size_; ++i) {
    uniform_data[i] = static_cast<const TypeParam>(uniform_uint_data[i]);
  }
  const TypeParam lower = 0;
  const TypeParam upper = UINT_MAX;
  this->RngUniformChecks(lower, upper, uniform_data);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngGaussianTimesGaussianGPU) {
  const TypeParam mu = 0;
  const TypeParam sigma = 1;

  // Sample from 0 mean Gaussian.
  TypeParam* gaussian_gpu_data_1 =
      static_cast<TypeParam*>(this->data_->mutable_gpu_data());
  this->RngGaussianFillGPU(mu, sigma, gaussian_gpu_data_1);

  // Sample from 0 mean Gaussian again.
  TypeParam* gaussian_gpu_data_2 =
      static_cast<TypeParam*>(this->data_2_->mutable_gpu_data());
  this->RngGaussianFillGPU(mu, sigma, gaussian_gpu_data_2);

  // Multiply Gaussians.
  TypeParam* gaussian_data_1 =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  const TypeParam* gaussian_data_2 =
      static_cast<const TypeParam*>(this->data_2_->cpu_data());
  for (int i = 0; i < this->sample_size_; ++i) {
    gaussian_data_1[i] *= gaussian_data_2[i];
  }

  // Check that result does not violate checked properties of Gaussian
  // (though it is not actually a Gaussian).
  TypeParam mu_product = pow(mu, 2);
  TypeParam sigma_product = sqrt(pow(sigma, 2) / 2);
  this->RngGaussianChecks(mu_product, sigma_product, gaussian_data_1);
}


TYPED_TEST(RandomNumberGeneratorTest, TestRngUniformTimesUniformGPU) {
  // Sample from Uniform on [-2, 2].
  const TypeParam lower_1 = -2;
  const TypeParam upper_1 = -lower_1;
  TypeParam* uniform_gpu_data_1 =
      static_cast<TypeParam*>(this->data_->mutable_gpu_data());
  this->RngUniformFillGPU(lower_1, upper_1, uniform_gpu_data_1);

  // Sample from Uniform on [-3, 3].
  const TypeParam lower_2 = -3;
  const TypeParam upper_2 = -lower_2;
  TypeParam* uniform_gpu_data_2 =
      static_cast<TypeParam*>(this->data_2_->mutable_gpu_data());
  this->RngUniformFillGPU(lower_2, upper_2, uniform_gpu_data_2);

  // Multiply Uniforms.
  TypeParam* uniform_data_1 =
      static_cast<TypeParam*>(this->data_->mutable_cpu_data());
  const TypeParam* uniform_data_2 =
      static_cast<const TypeParam*>(this->data_2_->cpu_data());
  for (int i = 0; i < this->sample_size_; ++i) {
    uniform_data_1[i] *= uniform_data_2[i];
  }

  // Check that result does not violate properties of Uniform on [-7, -3].
  const TypeParam lower_prod = lower_1 * upper_2;
  const TypeParam upper_prod = -lower_prod;
  this->RngUniformChecks(lower_prod, upper_prod, uniform_data_1);
}

#endif

}  // namespace caffe
