/*
Copyright 2018 Google Inc. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS-IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

#include "geometrical_acoustics/acoustic_source.h"

#include <cmath>
#include <random>
#include <vector>

#include "third_party/googletest/googletest/include/gtest/gtest.h"
#include "base/constants_and_types.h"
#include "geometrical_acoustics/test_util.h"

namespace vraudio {

namespace {

TEST(AcousticSource, GeneratedRayNonDirectionalDataTest) {
  std::default_random_engine engine(0);
  std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
  const std::array<float, kNumReverbOctaveBands> energies{
      {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f}};
  AcousticSource source({0.1f, 0.2f, 0.3f}, energies, [&engine, &distribution] {
    return distribution(engine);
  });
  AcousticRay ray = source.GenerateRay();

  EXPECT_EQ(ray.type(), AcousticRay::RayType::kSpecular);
  for (size_t i = 0; i < kNumReverbOctaveBands; ++i) {
    EXPECT_FLOAT_EQ(ray.energies().at(i), energies.at(i));
  }
  EXPECT_FLOAT_EQ(ray.t_near(), 0.0f);
  const float expected_origin[3] = {0.1f, 0.2f, 0.3f};
  ExpectFloat3Close(ray.origin(), expected_origin);
}

TEST(AcousticSource, GeneratedRayDirectionDistributionTest) {
  std::default_random_engine engine(0);
  std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
  const std::array<float, kNumReverbOctaveBands> energies{
      {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}};
  AcousticSource source({0.0f, 0.0f, 0.0f}, energies, [&engine, &distribution] {
    return distribution(engine);
  });

  // Check that the direction vectors on generated rays are all normalized,
  // i.e., of unit lengths.
  for (size_t i = 0; i < 1000; ++i) {
    AcousticRay ray = source.GenerateRay();
    const float direction_squared_norm =
        ray.direction()[0] * ray.direction()[0] +
        ray.direction()[1] * ray.direction()[1] +
        ray.direction()[2] * ray.direction()[2];
    EXPECT_FLOAT_EQ(direction_squared_norm, 1.0f);
  }

  // For a ray whose direction is uniformly distributed over a sphere:
  // - The PDF of theta is sin(theta), and the CDF is (1 - cos(theta)) / 2.
  // - The PDF of phi is 1 / 2 pi, and the CDF is 0.5 + phi / 2 pi.
  ValidateDistribution(100000, 100, [&source]() {
    AcousticRay ray = source.GenerateRay();
    const float cos_theta = ray.direction()[2];
    return 0.5f * (1.0f - cos_theta);
  });

  ValidateDistribution(100000, 100, [&source]() {
    AcousticRay ray = source.GenerateRay();
    const float* direction = ray.direction();
    const float phi = std::atan2(direction[1], direction[0]);
    return 0.5f + phi / 2.0f / static_cast<float>(M_PI);
  });
}

// Similar to GeneratedRayDirectionDistributionTest, but tests on the whole
// set of rays generated by AcousticSource::GenerateStratifiedRays() that they
// are uniformly distributed.
TEST(AcousticSource, GeneratedStratifiedRaysDirectionDistributionTest) {
  std::default_random_engine engine(0);
  std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
  const std::array<float, kNumReverbOctaveBands> energies{
      {1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f, 1.0f}};
  AcousticSource source({0.0f, 0.0f, 0.0f}, energies, [&engine, &distribution] {
    return distribution(engine);
  });

  // Check that the direction vectors on generated rays are all normalized,
  // i.e., of unit lengths.
  const size_t num_rays = 10000;
  const size_t sqrt_num_rays = 100;
  std::vector<AcousticRay> rays =
      source.GenerateStratifiedRays(num_rays, sqrt_num_rays);
  for (const AcousticRay& ray : rays) {
    const float direction_squared_norm =
        ray.direction()[0] * ray.direction()[0] +
        ray.direction()[1] * ray.direction()[1] +
        ray.direction()[2] * ray.direction()[2];
    EXPECT_FLOAT_EQ(direction_squared_norm, 1.0f);
  }

  // For a ray whose direction is uniformly distributed over a sphere:
  // - The PDF of theta is sin(theta), and the CDF is (1 - cos(theta)) / 2.
  // - The PDF of phi is 1 / 2 pi, and the CDF is 0.5 + phi / 2 pi.
  size_t ray_index = 0;
  ValidateDistribution(num_rays, 100, [&rays, &ray_index]() {
    const AcousticRay& ray = rays[ray_index];
    ++ray_index;
    const float cos_theta = ray.direction()[2];
    return 0.5f * (1.0f - cos_theta);
  });

  ray_index = 0;
  ValidateDistribution(num_rays, 100, [&rays, &ray_index]() {
    const AcousticRay& ray = rays[ray_index];
    ++ray_index;
    const float* direction = ray.direction();
    const float phi = std::atan2(direction[1], direction[0]);
    return 0.5f + phi / kTwoPi;
  });
}

}  // namespace

}  // namespace vraudio