// nnet3/nnet-component-test.cc
  
  // Copyright 2015  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 "nnet3/nnet-nnet.h"
  #include "nnet3/nnet-simple-component.h"
  #include "nnet3/nnet-test-utils.h"
  
  namespace kaldi {
  namespace nnet3 {
  // Reset seeds for test time for RandomComponent
  static void ResetSeed(int32 rand_seed, const Component &c) {
    RandomComponent *rand_component =
      const_cast<RandomComponent*>(dynamic_cast<const RandomComponent*>(&c));
  
    if (rand_component != NULL) {
      srand(rand_seed);
      rand_component->ResetGenerator();
    }
  }
  
  // this is the same as calling StringsApproxEqual(), except it prints
  // a warning if it fails.
  bool CheckStringsApproxEqual(const std::string &a,
                               const std::string &b,
                               int32 tolerance = 3) {
    if (!StringsApproxEqual(a, b, tolerance)) {
      KALDI_WARN << "Strings differ: " << a
                 << "
  vs.
  " << b;
      return false;
    } else {
      return true;
    }
  }
  
  
  void TestNnetComponentIo(Component *c) {
    bool binary = (Rand() % 2 == 0);
    std::ostringstream os1;
    c->Write(os1, binary);
    std::istringstream is(os1.str());
    Component *c2 = Component::ReadNew(is, binary);
    std::ostringstream os2;
    c2->Write(os2, binary);
    if (!binary) {
      std::string s1 = os1.str(), s2 = os2.str();
      KALDI_ASSERT(CheckStringsApproxEqual(s1, s2));
    }
    delete c2;
  }
  
  void TestNnetComponentCopy(Component *c) {
    Component *c2 = c->Copy();
    if (!StringsApproxEqual(c->Info(), c2->Info())) {
      KALDI_ERR << "Expected info strings to be equal: '"
                << c->Info() << "' vs. '" << c2->Info() << "'";
    }
    delete c2;
  }
  
  void TestNnetComponentAddScale(Component *c) {
    Component *c2 = c->Copy();
    Component *c3 = c2->Copy();
    c3->Add(0.5, *c2);
    c2->Scale(1.5);
    KALDI_ASSERT(CheckStringsApproxEqual(c2->Info(), c3->Info()));
    delete c2;
    delete c3;
  }
  
  void TestNnetComponentVectorizeUnVectorize(Component *c) {
    if (!(c->Properties() & kUpdatableComponent))
      return;
    UpdatableComponent *uc = dynamic_cast<UpdatableComponent*>(c);
    KALDI_ASSERT(uc != NULL);
    UpdatableComponent *uc2 = dynamic_cast<UpdatableComponent*>(uc->Copy());
    uc2->Scale(0.0);
    Vector<BaseFloat> params(uc2->NumParameters());
    uc2->Vectorize(&params);
    KALDI_ASSERT(params.Min()==0.0 && params.Sum()==0.0);
    uc->Vectorize(&params);
    uc2->UnVectorize(params);
    KALDI_ASSERT(CheckStringsApproxEqual(uc2->Info(), uc->Info()));
    BaseFloat x = uc2->DotProduct(*uc2), y = uc->DotProduct(*uc),
        z = uc2->DotProduct(*uc);
    KALDI_ASSERT(ApproxEqual(x, y) && ApproxEqual(y, z));
    Vector<BaseFloat> params2(uc2->NumParameters());
    uc2->Vectorize(&params2);
    for(int i = 0; i < params.Dim(); i++)
      KALDI_ASSERT(params(i) == params2(i));
    delete uc2;
  }
  
  void TestNnetComponentUpdatable(Component *c) {
    if (!(c->Properties() & kUpdatableComponent))
      return;
    UpdatableComponent *uc = dynamic_cast<UpdatableComponent*>(c);
    if (uc == NULL) {
      KALDI_ASSERT(!(c->Properties() & kUpdatableComponent) &&
                   "Component returns updatable flag but does not inherit "
                   "from UpdatableComponent");
      return;
    }
    if(!(uc->Properties() & kUpdatableComponent)){
      // testing that if it declares itself as non-updatable,
      // Scale() and Add() have no effect.
      KALDI_ASSERT(uc->NumParameters() == 0);
      KALDI_ASSERT(uc->DotProduct(*uc) == 0);
      UpdatableComponent *uc2 = dynamic_cast<UpdatableComponent*>(uc->Copy());
      uc2->Scale(7.0);
      uc2->Add(3.0, *uc);
      KALDI_ASSERT(CheckStringsApproxEqual(uc2->Info(), uc->Info()));
      uc->Scale(0.0);
      KALDI_ASSERT(CheckStringsApproxEqual(uc2->Info(), uc->Info()));
      delete uc2;
    } else {
      KALDI_ASSERT(uc->NumParameters() != 0);
      UpdatableComponent *uc2 = dynamic_cast<UpdatableComponent*>(uc->Copy()),
          *uc3 = dynamic_cast<UpdatableComponent*>(uc->Copy());
  
      // testing some expected invariances of scale and add.
      uc2->Scale(5.0);
      uc2->Add(3.0, *uc3);
      uc3->Scale(8.0);
      // now they should both be scaled to 8 times the original component.
      if (!StringsApproxEqual(uc2->Info(), uc3->Info())) {
        KALDI_ERR << "Expected info strings to be equal: '"
                  << uc2->Info() << "' vs. '" << uc3->Info() << "'";
      }
      // testing that scaling by 0.5 works the same whether
      // done on the vectorized paramters or via Scale().
      Vector<BaseFloat> vec2(uc->NumParameters());
      uc2->Vectorize(&vec2);
      vec2.Scale(0.5);
      uc2->UnVectorize(vec2);
      uc3->Scale(0.5);
      KALDI_ASSERT(CheckStringsApproxEqual(uc2->Info(), uc3->Info()));
  
      // testing that Scale(0.0) works the same whether done on the vectorized
      // paramters or via SetZero(), and that unvectorizing something that's been
      // zeroed gives us zero parameters.
      uc2->Vectorize(&vec2);
      vec2.SetZero();
      uc2->UnVectorize(vec2);
      uc3->Scale(0.0);
      uc3->Vectorize(&vec2);
      KALDI_ASSERT(uc2->Info() == uc3->Info() && VecVec(vec2, vec2) == 0.0);
  
      delete uc2;
      delete uc3;
    }
  }
  
  
  /*
    This function gets the 'ComponentPrecomputedIndexes*' pointer from
    a component, given the num-rows in the matrix of inputs we're testing it
    with.  It uses a plausible arrangement of indexes.
  
    Note: in this file we primarily test simple components, and simple
    components don't return precomputed indexes; but we also test a
    few non-simple components that operate with the same set of indexes
    on the input and the output.  Simple components would return NULL.
   */
  ComponentPrecomputedIndexes *GetPrecomputedIndexes(const Component &c,
                                                     int32 num_rows) {
    std::vector<Index> input_indexes(num_rows);
    int32 num_t_values;
    if (num_rows % 3 == 0) { num_t_values = 3; }
    else if (num_rows % 2 == 0) { num_t_values = 2; }
    else { num_t_values = 1; }
  
    for (int32 i = 0; i < num_rows; i++) {
      input_indexes[i].n = i % num_t_values;
      input_indexes[i].x = 0;
      input_indexes[i].t = i / num_t_values;
    }
    std::vector<Index> output_indexes(input_indexes);
  
    if (c.Properties()&kReordersIndexes) {
      c.ReorderIndexes(&input_indexes, &output_indexes);
    }
    MiscComputationInfo misc_info;
    bool need_backprop = true;  // just in case.
    ComponentPrecomputedIndexes *ans = c.PrecomputeIndexes(misc_info,
                                                           input_indexes,
                                                           output_indexes,
                                                           need_backprop);
    // ans will be NULL in most cases.
    return ans;
  }
  
  // tests the properties kPropagateAdds, kBackpropAdds,
  // kBackpropNeedsInput, kBackpropNeedsOutput.
  void TestSimpleComponentPropagateProperties(const Component &c) {
    int32 properties = c.Properties();
    Component *c_copy = NULL;
    int32 rand_seed = Rand();
  
    if (RandInt(0, 1) == 0)
      c_copy = c.Copy();  // This will test backprop with an updatable component.
    MatrixStrideType input_stride_type = (c.Properties()&kInputContiguous) ?
        kStrideEqualNumCols : kDefaultStride;
    MatrixStrideType output_stride_type = (c.Properties()&kOutputContiguous) ?
        kStrideEqualNumCols : kDefaultStride;
    MatrixStrideType both_stride_type =
        (c.Properties()&(kInputContiguous|kOutputContiguous)) ?
        kStrideEqualNumCols : kDefaultStride;
  
    int32 input_dim = c.InputDim(),
        output_dim = c.OutputDim(),
        num_rows = RandInt(1, 100);
    CuMatrix<BaseFloat> input_data(num_rows, input_dim, kUndefined,
                                   input_stride_type);
    input_data.SetRandn();
    CuMatrix<BaseFloat> output_data3(num_rows, input_dim, kSetZero,
                                     output_stride_type);
    output_data3.CopyFromMat(input_data);
    CuMatrix<BaseFloat>
        output_data1(num_rows, output_dim, kSetZero, output_stride_type),
        output_data2(num_rows, output_dim, kSetZero, output_stride_type);
    output_data2.Add(1.0);
  
    if ((properties & kPropagateAdds) && (properties & kPropagateInPlace)) {
      KALDI_ERR << "kPropagateAdds and kPropagateInPlace flags are incompatible.";
    }
  
    ResetSeed(rand_seed, c);
    ComponentPrecomputedIndexes *indexes = GetPrecomputedIndexes(c, num_rows);
    void *memo = c.Propagate(indexes, input_data, &output_data1);
  
    ResetSeed(rand_seed, c);
    c.DeleteMemo(c.Propagate(indexes, input_data, &output_data2));
    if (properties & kPropagateInPlace) {
      ResetSeed(rand_seed, c);
      c.DeleteMemo(c.Propagate(indexes, output_data3, &output_data3));
      if (!output_data1.ApproxEqual(output_data3)) {
        KALDI_ERR << "Test of kPropagateInPlace flag for component of type "
                  << c.Type() << " failed.";
      }
    }
    if (properties & kPropagateAdds)
      output_data2.Add(-1.0); // remove the offset
    AssertEqual(output_data1, output_data2);
  
  
    CuMatrix<BaseFloat> output_deriv(num_rows, output_dim, kSetZero, output_stride_type);
    output_deriv.SetRandn();
    CuMatrix<BaseFloat> input_deriv1(num_rows, input_dim, kSetZero, input_stride_type),
        input_deriv2(num_rows, input_dim, kSetZero, input_stride_type);
    CuMatrix<BaseFloat> input_deriv3(num_rows, output_dim, kSetZero, both_stride_type);
    input_deriv3.CopyFromMat(output_deriv);
  
    input_deriv2.Add(1.0);
    CuMatrix<BaseFloat> empty_mat;
  
    // test with input_deriv1 that's zero
    c.Backprop("foobar", indexes,
               ((properties & kBackpropNeedsInput) ? input_data : empty_mat),
               ((properties & kBackpropNeedsOutput) ? output_data1 : empty_mat),
               output_deriv,
               memo,
               c_copy,
               &input_deriv1);
    // test with input_deriv2 that's all ones.
    c.Backprop("foobar", indexes,
               ((properties & kBackpropNeedsInput) ? input_data : empty_mat),
               ((properties & kBackpropNeedsOutput) ? output_data1 : empty_mat),
               output_deriv,
               memo,
               c_copy,
               &input_deriv2);
    // test backprop in place, if supported.
    if (properties & kBackpropInPlace) {
      c.Backprop("foobar", indexes,
                 ((properties & kBackpropNeedsInput) ? input_data : empty_mat),
                 ((properties & kBackpropNeedsOutput) ? output_data1 : empty_mat),
                 input_deriv3,
                 memo,
                 c_copy,
                 &input_deriv3);
    }
    c.DeleteMemo(memo);
  
    if (properties & kBackpropAdds)
      input_deriv2.Add(-1.0);  // subtract the offset.
    AssertEqual(input_deriv1, input_deriv2);
    if (properties & kBackpropInPlace)
      AssertEqual(input_deriv1, input_deriv3);
    delete c_copy;
    delete indexes;
  }
  
  bool TestSimpleComponentDataDerivative(const Component &c,
                                         BaseFloat perturb_delta) {
    MatrixStrideType input_stride_type = (c.Properties()&kInputContiguous) ?
        kStrideEqualNumCols : kDefaultStride;
    MatrixStrideType output_stride_type = (c.Properties()&kOutputContiguous) ?
        kStrideEqualNumCols : kDefaultStride;
  
    int32 input_dim = c.InputDim(),
        output_dim = c.OutputDim(),
        num_rows = RandInt(1, 100),
        rand_seed = Rand();
    int32 properties = c.Properties();
    CuMatrix<BaseFloat> input_data(num_rows, input_dim, kSetZero, input_stride_type),
        output_data(num_rows, output_dim, kSetZero, output_stride_type),
        output_deriv(num_rows, output_dim, kSetZero, output_stride_type);
    input_data.SetRandn();
    output_deriv.SetRandn();
  
    ResetSeed(rand_seed, c);
    ComponentPrecomputedIndexes *indexes = GetPrecomputedIndexes(c, num_rows);
    void *memo = c.Propagate(indexes, input_data, &output_data);
  
    CuMatrix<BaseFloat> input_deriv(num_rows, input_dim, kSetZero, input_stride_type),
        empty_mat;
    c.Backprop("foobar", indexes,
               ((properties & kBackpropNeedsInput) ? input_data : empty_mat),
               ((properties & kBackpropNeedsOutput) ? output_data : empty_mat),
               output_deriv, memo, NULL, &input_deriv);
    c.DeleteMemo(memo);
  
    int32 test_dim = 3;
    BaseFloat original_objf = TraceMatMat(output_deriv, output_data, kTrans);
    Vector<BaseFloat> measured_objf_change(test_dim),
        predicted_objf_change(test_dim);
    for (int32 i = 0; i < test_dim; i++) {
      CuMatrix<BaseFloat> perturbed_input_data(num_rows, input_dim,
                                               kSetZero, input_stride_type),
          perturbed_output_data(num_rows, output_dim,
                                kSetZero, output_stride_type);
      perturbed_input_data.SetRandn();
      perturbed_input_data.Scale(perturb_delta);
      // at this point, perturbed_input_data contains the offset at the input data.
      predicted_objf_change(i) = TraceMatMat(perturbed_input_data, input_deriv,
                                             kTrans);
      perturbed_input_data.AddMat(1.0, input_data);
  
      ResetSeed(rand_seed, c);
      c.DeleteMemo(c.Propagate(indexes, perturbed_input_data, &perturbed_output_data));
      measured_objf_change(i) = TraceMatMat(output_deriv, perturbed_output_data,
                                            kTrans) - original_objf;
    }
    KALDI_LOG << "Predicted objf-change = " << predicted_objf_change;
    KALDI_LOG << "Measured objf-change = " << measured_objf_change;
    BaseFloat threshold = 0.1;
    bool ans = ApproxEqual(predicted_objf_change, measured_objf_change, threshold);
    if (!ans)
      KALDI_WARN << "Data-derivative test failed, component-type="
                 << c.Type() << ", input-dim=" << input_dim
                 << ", output-dim=" << output_dim;
    if (c.Type() == "NormalizeComponent" && input_dim == 1) {
      // derivatives are mathematically zero, but the measured and predicted
      // objf have different roundoff and the relative differences are large.
      // this is not unexpected.
      KALDI_LOG << "Accepting deriv differences since it is NormalizeComponent "
                << "with dim=1.";
      return true;
    }
    else if (c.Type() == "ClipGradientComponent") {
      KALDI_LOG << "Accepting deriv differences since "
                << "it is ClipGradientComponent.";
      return true;
    }
    delete indexes;
    return ans;
  }
  
  
  // if test_derivative == false then the test only tests that the update
  // direction is downhill.  if true, then we measure the actual model-derivative
  // and check that it's accurate.
  // returns true on success, false on test failure.
  bool TestSimpleComponentModelDerivative(const Component &c,
                                          BaseFloat perturb_delta,
                                          bool test_derivative) {
    int32 input_dim = c.InputDim(),
        output_dim = c.OutputDim(),
        num_rows = RandInt(1, 100);
    int32 properties = c.Properties();
    if ((properties & kUpdatableComponent) == 0) {
      // nothing to test.
      return true;
    }
    MatrixStrideType input_stride_type = (c.Properties()&kInputContiguous) ?
        kStrideEqualNumCols : kDefaultStride;
    MatrixStrideType output_stride_type = (c.Properties()&kOutputContiguous) ?
        kStrideEqualNumCols : kDefaultStride;
  
    CuMatrix<BaseFloat> input_data(num_rows, input_dim, kSetZero, input_stride_type),
        output_data(num_rows, output_dim, kSetZero, output_stride_type),
        output_deriv(num_rows, output_dim, kSetZero, output_stride_type);
    input_data.SetRandn();
    output_deriv.SetRandn();
  
    ComponentPrecomputedIndexes *indexes = GetPrecomputedIndexes(c, num_rows);
    void *memo = c.Propagate(indexes, input_data, &output_data);
  
    BaseFloat original_objf = TraceMatMat(output_deriv, output_data, kTrans);
  
    Component *c_copy = c.Copy();
  
    const UpdatableComponent *uc = dynamic_cast<const UpdatableComponent*>(&c);
    UpdatableComponent *uc_copy = dynamic_cast<UpdatableComponent*>(c_copy);
    KALDI_ASSERT(uc != NULL && uc_copy != NULL);
    if (test_derivative) {
      uc_copy->Scale(0.0);
      uc_copy->SetAsGradient();
    }
  
    CuMatrix<BaseFloat> input_deriv(num_rows, input_dim,
                                    kSetZero, input_stride_type),
        empty_mat;
    c.Backprop("foobar", indexes,
               ((properties & kBackpropNeedsInput) ? input_data : empty_mat),
               ((properties & kBackpropNeedsOutput) ? output_data : empty_mat),
               output_deriv, memo, c_copy,
               (RandInt(0, 1) == 0 ? &input_deriv : NULL));
    c.DeleteMemo(memo);
  
    if (!test_derivative) { // Just testing that the model update is downhill.
      CuMatrix<BaseFloat> new_output_data(num_rows, output_dim,
                                          kSetZero, output_stride_type);
      c.DeleteMemo(c_copy->Propagate(indexes, input_data, &new_output_data));
  
      BaseFloat new_objf = TraceMatMat(output_deriv, new_output_data, kTrans);
      bool ans = (new_objf > original_objf);
      if (!ans) {
        KALDI_WARN << "After update, new objf is not better than the original objf: "
                   << new_objf << " <= " << original_objf;
      }
      delete c_copy;
      delete indexes;
      return ans;
    } else {
      // check that the model derivative is accurate.
      int32 test_dim = 3;
  
      Vector<BaseFloat> measured_objf_change(test_dim),
          predicted_objf_change(test_dim);
      for (int32 i = 0; i < test_dim; i++) {
        CuMatrix<BaseFloat> perturbed_output_data(num_rows, output_dim,
                                                  kSetZero, output_stride_type);
        Component *c_perturbed = c.Copy();
        UpdatableComponent *uc_perturbed =
            dynamic_cast<UpdatableComponent*>(c_perturbed);
        KALDI_ASSERT(uc_perturbed != NULL);
        uc_perturbed->PerturbParams(perturb_delta);
  
        predicted_objf_change(i) = uc_copy->DotProduct(*uc_perturbed) -
            uc_copy->DotProduct(*uc);
        c_perturbed->Propagate(indexes, input_data, &perturbed_output_data);
        measured_objf_change(i) = TraceMatMat(output_deriv, perturbed_output_data,
                                              kTrans) - original_objf;
        delete c_perturbed;
      }
      KALDI_LOG << "Predicted objf-change = " << predicted_objf_change;
      KALDI_LOG << "Measured objf-change = " << measured_objf_change;
      BaseFloat threshold = 0.1;
  
      bool ans = ApproxEqual(predicted_objf_change, measured_objf_change,
                             threshold);
      if (!ans)
        KALDI_WARN << "Model-derivative test failed, component-type="
                   << c.Type() << ", input-dim=" << input_dim
                   << ", output-dim=" << output_dim;
      delete c_copy;
      delete indexes;
      return ans;
    }
  }
  
  
  void UnitTestNnetComponent() {
    for (int32 n = 0; n < 200; n++)  {
      Component *c = GenerateRandomSimpleComponent();
      KALDI_LOG << c->Info();
      TestNnetComponentIo(c);
      TestNnetComponentCopy(c);
      TestNnetComponentAddScale(c);
      TestNnetComponentVectorizeUnVectorize(c);
      TestNnetComponentUpdatable(c);
      TestSimpleComponentPropagateProperties(*c);
      if (!TestSimpleComponentDataDerivative(*c, 1.0e-04) &&
          !TestSimpleComponentDataDerivative(*c, 1.0e-03) &&
          !TestSimpleComponentDataDerivative(*c, 1.0e-05) &&
          !TestSimpleComponentDataDerivative(*c, 1.0e-06))
        KALDI_ERR << "Component data-derivative test failed";
  
      if (!TestSimpleComponentModelDerivative(*c, 1.0e-04, false) &&
          !TestSimpleComponentModelDerivative(*c, 1.0e-03, false) &&
          !TestSimpleComponentModelDerivative(*c, 1.0e-06, false))
        KALDI_ERR << "Component downhill-update test failed";
  
      if (!TestSimpleComponentModelDerivative(*c, 1.0e-04, true) &&
          !TestSimpleComponentModelDerivative(*c, 1.0e-03, true) &&
          !TestSimpleComponentModelDerivative(*c, 1.0e-05, true) &&
          !TestSimpleComponentModelDerivative(*c, 1.0e-06, true))
        KALDI_ERR << "Component model-derivative test failed";
  
      delete c;
    }
  }
  
  } // namespace nnet3
  } // namespace kaldi
  
  int main() {
    using namespace kaldi;
    using namespace kaldi::nnet3;
  #if HAVE_CUDA == 1
    kaldi::int32 loop = 0;
    for (loop = 0; loop < 2; loop++) {
      //CuDevice::Instantiate().SetDebugStrideMode(true);
      if (loop == 0)
        CuDevice::Instantiate().SelectGpuId("no");
      else
        CuDevice::Instantiate().SelectGpuId("yes");
  #endif
      UnitTestNnetComponent();
  #if HAVE_CUDA == 1
    } // No for loop if 'HAVE_CUDA != 1',
    CuDevice::Instantiate().PrintProfile();
  #endif
    KALDI_LOG << "Nnet component tests succeeded.";
  
    return 0;
  }