// nnet3/nnet-compute-test.cc

// Copyright 2015  Johns Hopkins University (author: Daniel Povey)
//           2015  Xiaohui Zhang

// 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 "nnet3/nnet-nnet.h"
#include "nnet3/nnet-compile.h"
#include "nnet3/nnet-analyze.h"
#include "nnet3/nnet-test-utils.h"
#include "nnet3/nnet-utils.h"
#include "nnet3/nnet-optimize.h"
#include "nnet3/nnet-compute.h"
#include "nnet3/nnet-am-decodable-simple.h"
#include "nnet3/decodable-simple-looped.h"

namespace kaldi {
namespace nnet3 {

void UnitTestNnetComputationIo(NnetComputation *computation) {
  bool binary = (Rand() % 2 == 0);
  std::ostringstream os;
  computation->Write(os, binary);
  const std::string &original_output = os.str();
  std::istringstream computation_is(original_output);
  computation->Read(computation_is, binary);
  std::istringstream computation_is2(original_output);
  NnetComputation computation2;
  computation2.Read(computation_is2, binary);

  std::ostringstream os2, os3;
  computation->Write(os2, binary);
  computation2.Write(os3, binary);

  if (binary) {
    if (!(os2.str() == original_output)) {
      KALDI_ERR << "Outputs differ for computation";
    }
  }
}

void UnitTestComputationRequestIo(ComputationRequest *request) {
  bool binary = (Rand() % 2 == 0);
  std::ostringstream os;
  request->Write(os, binary);
  const std::string &original_output = os.str();
  std::istringstream request_is(original_output);
  request->Read(request_is, binary);
  std::istringstream request_is2(original_output);
  ComputationRequest request2;
  request2.Read(request_is2, binary);

  std::ostringstream os2, os3;
  request->Write(os2, binary);
  request2.Write(os3, binary);
  KALDI_ASSERT(*request == request2);

  if (binary) {
    KALDI_ASSERT(os2.str() == original_output);
    KALDI_ASSERT(os3.str() == original_output);
  }
}

// this checks that a couple of different decodable objects give the same
// answer.
void TestNnetDecodable(Nnet *nnet) {
  int32 num_frames = 5 + RandInt(1, 100),
      input_dim = nnet->InputDim("input"),
      output_dim = nnet->OutputDim("output"),
      ivector_dim = std::max<int32>(0, nnet->InputDim("ivector"));
  Matrix<BaseFloat> input(num_frames, input_dim);

  SetBatchnormTestMode(true, nnet);
  SetDropoutTestMode(true, nnet);

  input.SetRandn();
  Vector<BaseFloat> ivector(ivector_dim);
  ivector.SetRandn();

  Vector<BaseFloat> priors(RandInt(0, 1) == 0 ? output_dim : 0);
  if (priors.Dim() != 0) {
    priors.SetRandn();
    priors.ApplyExp();
  }

  Matrix<BaseFloat> output1(num_frames, output_dim),
      output2(num_frames, output_dim);

  {
    NnetSimpleComputationOptions opts;
    opts.frames_per_chunk = RandInt(5, 25);
    CachingOptimizingCompiler compiler(*nnet);
    DecodableNnetSimple decodable(opts, *nnet, priors, input, &compiler,
                                  (ivector_dim != 0 ? &ivector : NULL));
    for (int32 t = 0; t < num_frames; t++) {
      SubVector<BaseFloat> row(output1, t);
      decodable.GetOutputForFrame(t, &row);
    }
  }

  {
    NnetSimpleLoopedComputationOptions opts;
    // caution: this may modify nnet, by changing how it consumes iVectors.
    DecodableNnetSimpleLoopedInfo info(opts, priors, nnet);
    DecodableNnetSimpleLooped decodable(info, input,
                                        (ivector_dim != 0 ? &ivector : NULL));
    for (int32 t = 0; t < num_frames; t++) {
      SubVector<BaseFloat> row(output2, t);
      decodable.GetOutputForFrame(t, &row);
    }
  }


  // the components that we exclude from this test, are excluded because they
  // all take "optional" right context, and this destroys the equivalence that
  // we are testing.
  if (!NnetIsRecurrent(*nnet) &&
      nnet->Info().find("statistics-extraction") == std::string::npos &&
      nnet->Info().find("TimeHeightConvolutionComponent") == std::string::npos &&
      nnet->Info().find("RestrictedAttentionComponent") == std::string::npos) {
    // this equivalence will not hold for recurrent nnets, or those that
    // have the statistics-extraction/statistics-pooling layers,
    // or in general for nnets with convolution components (because these
    // might have 'optional' context if required-time-offsets != time-offsets.
    for (int32 t = 0; t < num_frames; t++) {
      SubVector<BaseFloat> row1(output1, t),
          row2(output2, t);
      KALDI_ASSERT(row1.ApproxEqual(row2));
    }
  }
}

void UnitTestNnetCompute() {
  for (int32 n = 0; n < 20; n++) {
    struct NnetGenerationOptions gen_config;
    bool test_collapse_model = (RandInt(0, 1) == 0);

    std::vector<std::string> configs;
    GenerateConfigSequence(gen_config, &configs);
    Nnet nnet;
    for (size_t j = 0; j < configs.size(); j++) {
      KALDI_LOG << "Input config[" << j << "] is: " << configs[j];
      std::istringstream is(configs[j]);
      nnet.ReadConfig(is);
    }

    ComputationRequest request;
    std::vector<Matrix<BaseFloat> > inputs;
    ComputeExampleComputationRequestSimple(nnet, &request, &inputs);

    // Test CollapseModel().  Note: lines with 'collapse' in some part of them
    // are not necessary for the rest of the test to run; they only test
    // CollapseModel().
    if (test_collapse_model) {
      // this model collapsing code requires that test mode is set for batchnorm
      // and dropout components.
      SetBatchnormTestMode(true, &nnet);
      SetDropoutTestMode(true, &nnet);
    }

    NnetComputation computation;
    Compiler compiler(request, nnet);
    CompilerOptions opts;
    compiler.CreateComputation(opts, &computation);

    Nnet nnet_collapsed(nnet);
    CollapseModelConfig collapse_config;
    NnetComputation computation_collapsed;

    if (test_collapse_model) {
      CollapseModel(collapse_config, &nnet_collapsed);
      Compiler compiler_collapsed(request, nnet_collapsed);
      compiler_collapsed.CreateComputation(opts, &computation_collapsed);
      computation_collapsed.ComputeCudaIndexes();
    }


    {
      std::ostringstream os;
      computation.Print(os, nnet);
      KALDI_LOG << "Generated computation is: " << os.str();
      UnitTestNnetComputationIo(&computation);
      UnitTestComputationRequestIo(&request);
    }
    CheckComputationOptions check_config;
    // we can do the rewrite check since it's before optimization.
    check_config.check_rewrite = true;
    ComputationChecker checker(check_config, nnet, computation);
    checker.Check();

    if (RandInt(0, 1) == 0) {
      NnetOptimizeOptions opt_config;

      Optimize(opt_config, nnet,
               MaxOutputTimeInRequest(request),
               &computation);
      {
        std::ostringstream os;
        computation.Print(os, nnet);
        KALDI_LOG << "Optimized computation is: " << os.str();
      }
    }

    NnetComputeOptions compute_opts;
    if (RandInt(0, 1) == 0)
      compute_opts.debug = true;

    computation.ComputeCudaIndexes();
    NnetComputer computer(compute_opts,
                          computation,
                          nnet,
                          &nnet);
    // provide the input to the computation.
    for (size_t i = 0; i < request.inputs.size(); i++) {
      CuMatrix<BaseFloat> temp(inputs[i]);
      KALDI_LOG << "Input sum is " << temp.Sum();
      computer.AcceptInput(request.inputs[i].name, &temp);

    }
    computer.Run();


    const CuMatrixBase<BaseFloat> &output(computer.GetOutput("output"));

    KALDI_LOG << "Output sum is " << output.Sum();

    if (test_collapse_model) {
      NnetComputer computer_collapsed(compute_opts,
                                      computation_collapsed,
                                      nnet_collapsed,
                                      &nnet_collapsed);
      for (size_t i = 0; i < request.inputs.size(); i++) {
        CuMatrix<BaseFloat> temp(inputs[i]);
        KALDI_LOG << "Input sum is " << temp.Sum();
        computer_collapsed.AcceptInput(request.inputs[i].name, &temp);
      }
      computer_collapsed.Run();
      const CuMatrixBase<BaseFloat> &output_collapsed(
          computer_collapsed.GetOutput("output"));
      KALDI_LOG << "Output sum [collapsed] is " << output_collapsed.Sum();
      if (!ApproxEqual(output, output_collapsed)) {
        KALDI_ERR << "Regular and collapsed computations' outputs differ";
      }
    }

    CuMatrix<BaseFloat> output_deriv(output.NumRows(), output.NumCols());
    output_deriv.SetRandn();
    // output_deriv sum won't be informative so don't print it.
    if (request.outputs[0].has_deriv) {
      computer.AcceptInput("output", &output_deriv);
      computer.Run();
      for (size_t i = 0; i < request.inputs.size(); i++) {
        if (request.inputs[i].has_deriv) {
          const CuMatrixBase<BaseFloat> &in_deriv =
              computer.GetOutput(request.inputs[i].name);
          KALDI_LOG << "Input-deriv sum for input '"
                    << request.inputs[i].name << "' is " << in_deriv.Sum();
        }
      }
    }
    TestNnetDecodable(&nnet);
  }
}

} // namespace nnet3
} // namespace kaldi

int main() {
  using namespace kaldi;
  using namespace kaldi::nnet3;
  // uncommenting the following activates extra checks during optimization, that
  // can help narrow down the source of problems.
  //SetVerboseLevel(4);


  for (kaldi::int32 loop = 0; loop < 2; loop++) {
#if HAVE_CUDA == 1
    CuDevice::Instantiate().SetDebugStrideMode(true);
    if (loop == 0)
      CuDevice::Instantiate().SelectGpuId("no");
    else
      CuDevice::Instantiate().SelectGpuId("yes");
#endif
    UnitTestNnetCompute();
  }

  KALDI_LOG << "Nnet tests succeeded.";

  return 0;
}