// cudamatrix/cu-packed-matrix-test.cc
  //
  // Copyright 2013  Ehsan Variani
  //                 Lucas Ondel
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // 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
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  // UnitTests for testing cu-sp-matrix.h methods.
  //
  
  #include <iostream>
  #include <vector>
  #include <cstdlib>
  
  #include "base/kaldi-common.h"
  #include "cudamatrix/cu-device.h"
  #include "cudamatrix/cu-sp-matrix.h"
  #include "cudamatrix/cu-vector.h"
  #include "cudamatrix/cu-math.h"
  
  using namespace kaldi;
  
  namespace kaldi {
  
  /*
   * INITIALIZERS
   */
  
  /*
   * ASSERTS
   */
  template<typename Real>
  static void AssertEqual(const CuPackedMatrix<Real> &A,
                          const CuPackedMatrix<Real> &B,
                          float tol = 0.001) {
    KALDI_ASSERT(A.NumRows() == B.NumRows());
    for (MatrixIndexT i = 0; i < A.NumRows(); i++)
      for (MatrixIndexT j = 0; j <= i; j++)
        KALDI_ASSERT(std::abs(A(i, j) - B(i, j))
                     < tol * std::max(1.0, (double) (std::abs(A(i, j)) + std::abs(B(i, j)))));
  }
  
  template<typename Real>
  static void AssertEqual(const PackedMatrix<Real> &A,
                          const PackedMatrix<Real> &B,
                          float tol = 0.001) {
    KALDI_ASSERT(A.NumRows() == B.NumRows());
    for (MatrixIndexT i = 0; i < A.NumRows(); i++)
      for (MatrixIndexT j = 0; j <= i; j++)
        KALDI_ASSERT(std::abs(A(i, j) - B(i, j))
                     < tol * std::max(1.0, (double) (std::abs(A(i, j)) + std::abs(B(i, j)))));
  }
  
  template<typename Real>
  static void AssertDiagEqual(const PackedMatrix<Real> &A,
                          const CuPackedMatrix<Real> &B,
                          float value,
                          float tol = 0.001) {
    for (MatrixIndexT i = 0; i < A.NumRows(); i++) {
      KALDI_ASSERT(std::abs((A(i, i)+value) - B(i, i))  
                   < tol * std::max(1.0, (double) (std::abs(A(i, i)) + std::abs(B(i, i) + value))));
    }
  }
  template<typename Real>
  static void AssertDiagEqual(const PackedMatrix<Real> &A,
                          const PackedMatrix<Real> &B,
                          float value,
                          float tol = 0.001) {
    for (MatrixIndexT i = 0; i < A.NumRows(); i++) {
      KALDI_ASSERT(std::abs((A(i, i)+value) - B(i, i))  
                   < tol * std::max(1.0, (double) (std::abs(A(i, i)) + std::abs(B(i, i) + value))));
    }
  }
  
  template<typename Real>
  static void AssertEqual(const PackedMatrix<Real> &A,
                          const CuPackedMatrix<Real> &B,
                          float tol = 0.001) {
    KALDI_ASSERT(A.NumRows() == B.NumRows());
    for (MatrixIndexT i = 0; i < A.NumRows(); i++)
      for (MatrixIndexT j = 0; j <= i; j++)
        KALDI_ASSERT(std::abs(A(i, j) - B(i, j))
                     < tol * std::max(1.0, (double) (std::abs(A(i, j)) + std::abs(B(i, j)))));
  }
  
  template<typename Real>
  static bool ApproxEqual(const PackedMatrix<Real> &A,
                          const PackedMatrix<Real> &B, Real tol = 0.001) {
    KALDI_ASSERT(A.NumRows() == B.NumRows());
    PackedMatrix<Real> diff(A);
    diff.AddPacked(1.0, B);
    Real a = std::max(A.Max(), -A.Min()), b = std::max(B.Max(), -B.Min()),
        d = std::max(diff.Max(), -diff.Min());
    return (d <= tol * std::max(a, b));
  }
  
  /*
   * Unit Tests
   */
  template<typename Real>
  static void UnitTestCuPackedMatrixConstructor() { 
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 10 * i;
  
      PackedMatrix<Real> A(dim);
      A.SetRandn();
      CuPackedMatrix<Real> B(A);
      CuPackedMatrix<Real> C(B);
      AssertEqual(B, C);
    }
  }
  
  template<typename Real>
  static void UnitTestCuPackedMatrixCopy() { 
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 10 * i;
      
      PackedMatrix<Real> A(dim);
      A.SetRandn();
      CuPackedMatrix<Real> B(A);
  
      CuPackedMatrix<Real> C(dim);
      C.CopyFromPacked(A);
      CuPackedMatrix<Real> D(dim);
      D.CopyFromPacked(B);
      AssertEqual(C, D);
  
      PackedMatrix<Real> E(dim);
      D.CopyToPacked(&E);
      AssertEqual(A, E);
    }
  }
  
  template<typename Real>
  static void UnitTestCuPackedMatrixTrace() {
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 5 * i + Rand() % 10;
      
      PackedMatrix<Real> A(dim);
      A.SetRandn();
      CuPackedMatrix<Real> B(A);
      
      AssertEqual(A.Trace(), B.Trace());
    }
  }
  
  template<typename Real>
  static void UnitTestCuPackedMatrixScale() {
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 5 * i + Rand() % 10;
      
      PackedMatrix<Real> A(dim);
      A.SetRandn();
      CuPackedMatrix<Real> B(A);
  
      Real scale_factor = 23.5896223;
      A.Scale(scale_factor); 
      B.Scale(scale_factor);
      AssertEqual(A, B);
    }
  }
  
  template<typename Real>
  static void UnitTestCuPackedMatrixScaleDiag() {
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 5 * i + Rand() % 10;
      
      PackedMatrix<Real> A(dim);
      A.SetRandn();
      CuPackedMatrix<Real> B(A);
  
      Real scale_factor = 23.5896223;
      A.ScaleDiag(scale_factor); 
      B.ScaleDiag(scale_factor);
      AssertEqual(A, B);
    }
  }
  
  
  
  template<typename Real>
  static void UnitTestCuPackedMatrixAddToDiag() {
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 5 * i + Rand() % 10;
      
      PackedMatrix<Real> A(dim);
      A.SetRandn();
      CuPackedMatrix<Real> B(A);
  
      Real value = Rand() % 50;
      B.AddToDiag(value); 
      
      AssertDiagEqual(A, B, value);
    }
  }
  
  template<typename Real>
  static void UnitTestCuPackedMatrixSetUnit() {
    for (MatrixIndexT i = 1; i < 10; i++) {
      MatrixIndexT dim = 5 * i + Rand() % 10;
      
      CuPackedMatrix<Real> A(dim);
      A.SetUnit();
      
      for (MatrixIndexT i = 0; i < A.NumRows(); i++) {
        for (MatrixIndexT j = 0; j < A.NumRows(); j++) {
          if (i != j) { 
            KALDI_ASSERT(A(i, j) == 0);
          } else {
            KALDI_ASSERT(A(i, j) == 1.0);
          }
        }
      } 
    }
  }
  
  
  template<typename Real> void CudaPackedMatrixUnitTest() {
    UnitTestCuPackedMatrixConstructor<Real>();
    //UnitTestCuPackedMatrixCopy<Real>();
    UnitTestCuPackedMatrixTrace<Real>();
    UnitTestCuPackedMatrixScale<Real>();
    UnitTestCuPackedMatrixAddToDiag<Real>();
    UnitTestCuPackedMatrixSetUnit<Real>();
  }
  
  } // namespace kaldi
  
  
  int main() {
    using namespace kaldi;
  #if HAVE_CUDA == 1
    CuDevice::Instantiate().SetDebugStrideMode(true);
    // Select the GPU
    CuDevice::Instantiate().SelectGpuId("yes");
  #endif
    kaldi::CudaPackedMatrixUnitTest<float>();
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().DoublePrecisionSupported()) {
      kaldi::CudaPackedMatrixUnitTest<double>();
    } else {
      KALDI_WARN << "Double precision not supported";
    }
  #else
    kaldi::CudaPackedMatrixUnitTest<double>();
  #endif
    
    KALDI_LOG << "Tests succeeded";
  #if HAVE_CUDA == 1
    CuDevice::Instantiate().PrintProfile();
  #endif
    return 0;
  }