/*
 *  Copyright (c) 2014 The WebRTC project authors. All Rights Reserved.
 *
 *  Use of this source code is governed by a BSD-style license
 *  that can be found in the LICENSE file in the root of the source
 *  tree. An additional intellectual property rights grant can be found
 *  in the file PATENTS.  All contributing project authors may
 *  be found in the AUTHORS file in the root of the source tree.
 */

#define _USE_MATH_DEFINES

#include "webrtc/modules/audio_processing/beamformer/covariance_matrix_generator.h"

#include <cmath>

#include "testing/gtest/include/gtest/gtest.h"
#include "webrtc/modules/audio_processing/beamformer/matrix_test_helpers.h"

namespace webrtc {

using std::complex;

TEST(CovarianceMatrixGeneratorTest, TestBoxcar) {
  const float kWaveNumber = 10.3861f;
  const int kNumberMics = 3;
  const float kMicSpacing = 0.23f;
  const float kHalfWidth = 0.001f;
  const float kTolerance = 0.0001f;

  ComplexMatrix<float> actual_boxcar(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::Boxcar(
      kWaveNumber, kNumberMics, kMicSpacing, kHalfWidth, &actual_boxcar);

  complex<float>* const* actual_els = actual_boxcar.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].real(), 0.002f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].real(), 0.002f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].imag(), 0.f, kTolerance);
}

// Test Boxcar with large numbers to generate matrix with varying
// numbers.
TEST(CovarianceMatrixGeneratorTest, TestBoxcarLargeNumbers) {
  const float kWaveNumber = 10.3861f;
  const int kNumberMics = 3;
  const float kMicSpacing = 3.f;
  const float kHalfWidth = 0.1f;
  const float kTolerance = 0.0001f;

  ComplexMatrix<float> actual_boxcar(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::Boxcar(
      kWaveNumber, kNumberMics, kMicSpacing, kHalfWidth, &actual_boxcar);

  complex<float>* const* actual_els = actual_boxcar.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 0.2f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.0017f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].real(), -0.0017f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.0017f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 0.2f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].real(), 0.0017f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].real(), -0.0017f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].real(), 0.0017f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].real(), 0.2f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].imag(), 0.f, kTolerance);
}

TEST(CovarianceMatrixGeneratorTest, TestGappedUniformCovarianceMatrix2Mics) {
  const float kWaveNumber = 0.5775f;
  const int kNumberMics = 2;
  const float kMicSpacing = 0.05f;
  const float kHalfWidth = 0.001f;
  const float kTolerance = 0.0001f;

  ComplexMatrix<float> actual_covariance_matrix(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::GappedUniformCovarianceMatrix(
      kWaveNumber,
      kNumberMics,
      kMicSpacing,
      kHalfWidth,
      &actual_covariance_matrix);

  complex<float>* const* actual_els = actual_covariance_matrix.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 0.998f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.9978f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.9978f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 0.998f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
}

TEST(CovarianceMatrixGeneratorTest, TestGappedUniformCovarianceMatrix3Mics) {
  const float kWaveNumber = 10.3861f;
  const int kNumberMics = 3;
  const float kMicSpacing = 0.04f;
  const float kHalfWidth = 0.001f;
  const float kTolerance = 0.0001f;

  ComplexMatrix<float> actual_covariance_matrix(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::GappedUniformCovarianceMatrix(
      kWaveNumber,
      kNumberMics,
      kMicSpacing,
      kHalfWidth,
      &actual_covariance_matrix);

  complex<float>* const* actual_els = actual_covariance_matrix.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 0.998f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.9553f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].real(), 0.8327f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.9553f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 0.998f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].real(), 0.9553f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].real(), 0.8327f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].real(), 0.9553f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].real(), 0.998f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].imag(), 0.f, kTolerance);
}

TEST(CovarianceMatrixGeneratorTest, TestAngledCovarianceMatrix2Mics) {
  const float kSpeedOfSound = 340;
  const float kAngle = static_cast<float>(M_PI) / 4.f;
  const float kFrequencyBin = 6;
  const float kFftSize = 512;
  const int kNumberFrequencyBins = 257;
  const int kSampleRate = 16000;
  const int kNumberMics = 2;
  const float kMicSpacing = 0.04f;
  const float kTolerance = 0.0001f;
  ComplexMatrix<float> actual_covariance_matrix(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::AngledCovarianceMatrix(kSpeedOfSound,
                                                    kAngle,
                                                    kFrequencyBin,
                                                    kFftSize,
                                                    kNumberFrequencyBins,
                                                    kSampleRate,
                                                    kNumberMics,
                                                    kMicSpacing,
                                                    &actual_covariance_matrix);

  complex<float>* const* actual_els = actual_covariance_matrix.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 1.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.9952f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.9952f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 1.f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.0978f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), -0.0978f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
}

TEST(CovarianceMatrixGeneratorTest, TestAngledCovarianceMatrix3Mics) {
  const float kSpeedOfSound = 340;
  const float kAngle = static_cast<float>(M_PI) / 4.f;
  const float kFrequencyBin = 9;
  const float kFftSize = 512;
  const int kNumberFrequencyBins = 257;
  const int kSampleRate = 42000;
  const int kNumberMics = 3;
  const float kMicSpacing = 0.05f;
  const float kTolerance = 0.0001f;

  ComplexMatrix<float> actual_covariance_matrix(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::AngledCovarianceMatrix(kSpeedOfSound,
                                                    kAngle,
                                                    kFrequencyBin,
                                                    kFftSize,
                                                    kNumberFrequencyBins,
                                                    kSampleRate,
                                                    kNumberMics,
                                                    kMicSpacing,
                                                    &actual_covariance_matrix);

  complex<float>* const* actual_els = actual_covariance_matrix.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 1.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.8859f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].real(), 0.5696f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.8859f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 1.f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].real(), 0.8859f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].real(), 0.5696f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].real(), 0.8859f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].real(), 1.f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.4639f, kTolerance);
  EXPECT_NEAR(actual_els[0][2].imag(), 0.8219f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), -0.4639f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][2].imag(), 0.4639f, kTolerance);
  EXPECT_NEAR(actual_els[2][0].imag(), -0.8219f, kTolerance);
  EXPECT_NEAR(actual_els[2][1].imag(), -0.4639f, kTolerance);
  EXPECT_NEAR(actual_els[2][2].imag(), 0.f, kTolerance);
}

TEST(CovarianceMatrixGeneratorTest, TestDCCovarianceMatrix) {
  const int kNumberMics = 2;
  const float kHalfWidth = 0.01f;

  ComplexMatrix<float> actual_covariance_matrix(kNumberMics, kNumberMics);
  CovarianceMatrixGenerator::DCCovarianceMatrix(
      kNumberMics, kHalfWidth, &actual_covariance_matrix);

  complex<float>* const* actual_els = actual_covariance_matrix.elements();

  EXPECT_NEAR(actual_els[0][0].real(), 0.98f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].real(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].real(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].real(), 0.98f, kTolerance);

  EXPECT_NEAR(actual_els[0][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[0][1].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][0].imag(), 0.f, kTolerance);
  EXPECT_NEAR(actual_els[1][1].imag(), 0.f, kTolerance);
}

// PhaseAlignmentMasks is tested by AngledCovarianceMatrix and by
// InitBeamformerWeights in BeamformerUnittest.

}  // namespace webrtc
