// cudamatrix/cu-block-matrix-test.cc

// Copyright 2013  Johns Hopkins University (author: Daniel Povey)

// 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.


#include <iostream>
#include <vector>
#include <cstdlib>

#include "base/kaldi-common.h"
#include "util/common-utils.h"
#include "cudamatrix/cu-matrix-lib.h"

using namespace kaldi;


namespace kaldi {

template<typename Real>
static bool ApproxEqual(const CuBlockMatrix<Real> &A,
                        const CuBlockMatrix<Real> &B,
                        float tol = 0.001) {
  CuMatrix<Real> Acopy(A), Bcopy(B);
  return Acopy.ApproxEqual(Bcopy, tol);
}




template<class Real>
static void UnitTestCuBlockMatrixIO() {
  for (int32 i = 0; i < 10; i++) {
    int32 num_blocks = Rand() % 5;
    std::vector<CuMatrix<Real> > data(num_blocks);
    for (int32 b = 0; b < num_blocks; b++) {
      int32 dimM = 100 + Rand() % 255, dimN = 10 + Rand() % 20;
      if (b % 2 == 0) std::swap(dimM, dimN);
      data[b].Resize(dimM, dimN);
      data[b].SetRandn();
    }
    CuBlockMatrix<Real> B(data);

    std::ostringstream os;
    bool binary = (i % 4 < 2);
    B.Write(os, binary);

    CuBlockMatrix<Real> B2;
    std::istringstream is(os.str());
    B2.Read(is, binary);

    CuMatrix<Real> mat(B), mat2(B2);
    AssertEqual(mat, mat2);
    if (!data.empty())
      KALDI_ASSERT(mat.Sum() != 0.0);
  }
}



template<class Real>
static void UnitTestCuBlockMatrixAddMatBlock() {
  for (int32 i = 0; i < 20; i++) {
    int32 num_blocks = Rand() % 5;
    std::vector<CuMatrix<Real> > data(num_blocks);
    for (int32 b = 0; b < num_blocks; b++) {
      int32 dimM = 100 + Rand() % 255, dimN = 10 + Rand() % 20;
      // early failures will have small dim for easier eyeballing.
      if (b % 2 == 0) std::swap(dimM, dimN);
      data[b].Resize(dimM, dimN);
      data[b].SetRandn();
    }
    CuBlockMatrix<Real> B(data);
    int32 B_num_rows = B.NumRows(), B_num_cols = B.NumCols();
    // will do X += A B

    MatrixTransposeType transB = (i % 2 == 1 ? kTrans : kNoTrans),
        transA = (i % 3 == 1 ? kTrans : kNoTrans);
    if (transB == kTrans) std::swap(B_num_rows, B_num_cols);

    int32 X_num_rows = 100 + Rand() % 255, X_num_cols = B_num_cols,
        A_num_rows = X_num_rows, A_num_cols = B_num_rows;
    if (data.size() == 0) { X_num_rows = 0; A_num_rows = 0; }
    if (transA == kTrans) std::swap(A_num_rows, A_num_cols);

    Real alpha = 2.0, beta = -1.0;
    CuMatrix<Real> X(X_num_rows, X_num_cols);
    X.SetRandn();
    CuMatrix<Real> A(A_num_rows, A_num_cols);
    A.SetRandn();

    CuMatrix<Real> Xcopy(X), Bcopy(B), Xorig(X), Aorig(A);
    Xcopy.AddMatMat(alpha, A, transA, Bcopy, transB, beta);
    X.AddMatBlock(alpha, A, transA, B, transB, beta);

    AssertEqual(X, Xcopy);
  }
}


template<class Real>
static void UnitTestCuBlockMatrixAddMatMat() {
  for (int32 i = 0; i < 20; i++) {
    int32 num_blocks = Rand() % 5 + 1;
    std::vector<CuMatrix<Real> > data(num_blocks);
    for (int32 b = 0; b < num_blocks; b++) {
      int32 dimM = 100 + Rand() % 255, dimN = 10 + Rand() % 20;
      if (i == 0) { dimM = 1; dimN = 1; }
      // early failures will have small dim for easier eyeballing.
      if (b % 2 == 0) std::swap(dimM, dimN);
      data[b].Resize(dimM, dimN);
      KALDI_LOG << "dimM " << dimM << ", dimN " << dimN << ", stride " << data[b].Stride();
      data[b].SetRandn();
    }

    CuBlockMatrix<Real> B(data);
    int32 B_num_rows = B.NumRows(), B_num_cols = B.NumCols();
    // will do B += C D

    int32 C_num_rows = B_num_rows, C_num_cols = 100 + Rand() % 255;
    if (C_num_rows == 0) C_num_cols = 0;
    int32 D_num_rows = C_num_cols, D_num_cols = B_num_cols;

    MatrixTransposeType transC = (i % 2 == 1 ? kTrans : kNoTrans),
        transD = (i % 3 == 1 ? kTrans : kNoTrans);
    if (transC == kTrans) std::swap(C_num_rows, C_num_cols);
    if (transD == kTrans) std::swap(D_num_rows, D_num_cols);

    CuMatrix<Real> C(C_num_rows, C_num_cols), D(D_num_rows, D_num_cols);
    C.SetRandn();
    D.SetRandn();

    CuMatrix<Real> Bmat(B);

    Real alpha = 2.0, beta = -1.0;

    CuBlockMatrix<Real> Bcopy(B);

    B.AddMatMat(alpha, C, transC, D, transD, beta);

    Bmat.AddMatMat(alpha, C, transC, D, transD, beta);

    // Now check that the block-structured part of Bmat is the
    // same as B.
    Bcopy.CopyFromMat(Bmat); // copy block-structured part from Bmat to Bcopy.

    if (!ApproxEqual(B, Bcopy)) {
      KALDI_WARN << "CuBlockMatrixTest failure, please report to maintainers: Bcopy = "
                 << Bcopy << ", B = " << B << ", C = " << C << ", D = " << D
                 << ", Bmat = " << B << " transD = " << transD << ", transC = "
                 << transC;
      KALDI_ERR << "Please give this log to the maintainers.";
    }
    KALDI_ASSERT(Bmat.Sum() != 0 || B_num_rows == 0);
  }
}


template<typename Real> void CuBlockMatrixUnitTest() {
  UnitTestCuBlockMatrixIO<Real>();
  UnitTestCuBlockMatrixAddMatBlock<Real>();
  UnitTestCuBlockMatrixAddMatMat<Real>();
}


} // namespace kaldi


int main() {
  SetVerboseLevel(1);
  int32 loop = 0;
#if HAVE_CUDA == 1
  for (; loop < 2; loop++) {
    CuDevice::Instantiate().SetDebugStrideMode(true);
    if (loop == 0)
      CuDevice::Instantiate().SelectGpuId("no"); // -1 means no GPU
    else
      CuDevice::Instantiate().SelectGpuId("yes"); // -2 .. automatic selection
#endif

    kaldi::CuBlockMatrixUnitTest<float>();
#if HAVE_CUDA == 1
    if (CuDevice::Instantiate().DoublePrecisionSupported()) {
      kaldi::CuBlockMatrixUnitTest<double>();
    } else {
      KALDI_WARN << "Double precision not supported";
    }
#else
    kaldi::CuBlockMatrixUnitTest<double>();
#endif

    if (loop == 0)
      KALDI_LOG << "Tests without GPU use succeeded.";
    else
      KALDI_LOG << "Tests with GPU use (if available) succeeded.";
#if HAVE_CUDA == 1
  }
  CuDevice::Instantiate().PrintProfile();
#endif
  return 0;
}