// nnet3/convolution-test.cc
  
  // Copyright 2017    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/convolution.h"
  #include "util/common-utils.h"
  
  namespace kaldi {
  namespace nnet3 {
  namespace time_height_convolution {
  
  // for testing purposes, create a random ConvolutionModel.
  static void GetRandomConvolutionModel(ConvolutionModel *model) {
  start:
    {
      model->num_filters_in = RandInt(1, 10);
      model->num_filters_out = RandInt(1, 10);
      model->height_in = RandInt(1, 10);
      int32 min_height_offset = RandInt(-2, 0),
          max_height_offset = RandInt(0, 2),
          min_time_offset = RandInt(-2, 0),
          max_time_offset = RandInt(0, 2);
  
      model->height_out = RandInt(1, model->height_in);
      model->height_subsample_out = 1;
      if (RandInt(0, 1) == 0) {
        if (model->height_out % 2 == 0) {
          model->height_out /= 2;
          model->height_subsample_out = 2;
        } else if (model->height_out % 3 == 0) {
          model->height_out /= 3;
          model->height_subsample_out = 3;
        }
      }
      std::vector<int32> all_time_offsets;
      int32 max_offsets = RandInt(1, 10);
      model->offsets.clear();
      model->required_time_offsets.clear();
      for (int32 i = 0; i < max_offsets; i++) {
        ConvolutionModel::Offset o;
        o.time_offset = RandInt(min_time_offset, max_time_offset);
        o.height_offset = RandInt(min_height_offset, max_height_offset);
        all_time_offsets.push_back(o.time_offset);
        model->offsets.push_back(o);
      }
      SortAndUniq(&(model->offsets));
      SortAndUniq(&all_time_offsets);
      std::random_shuffle(all_time_offsets.begin(), all_time_offsets.end());
      int32 num_required_offsets = RandInt(1, all_time_offsets.size());
      for (int32 i = 0; i < num_required_offsets; i++)
        model->required_time_offsets.insert(all_time_offsets[i]);
      model->ComputeDerived();
    }
    if (!model->Check()) {
      KALDI_WARN << "Regenerating model because it didn't pass the check: "
                 << model->Info();
      goto start;
    }
  }
  
  // for testing purposes, create a set of input and output indexes for
  // a convolution computation that are computable given this model.
  static void GetRandomConvolutionIndexes(const ConvolutionModel &model,
                                          std::vector<Index> *input_indexes,
                                          std::vector<Index> *output_indexes) {
    KALDI_ASSERT(model.Check());
  
    std::vector<std::pair<int32, int32> > n_x_pairs;
    int32 num_n_x_pairs = RandInt(1, 3);
    for (int32 i = 0; i < num_n_x_pairs; i++) {
      int32 n = RandInt(0, 3), x = RandInt(0, 1);
      n_x_pairs.push_back(std::pair<int32, int32>(n, x));
    }
    SortAndUniq(&n_x_pairs);
    num_n_x_pairs = n_x_pairs.size();
  
  
    // 'output_t_values' is the set of *possible* output
    // t values; we'll later sub-sample from these.
    std::vector<int32> output_t_values;
  
    {
      int32 out_t_start = RandInt(-5, 5), out_t_step = RandInt(1, 3),
          num_t_out = RandInt(1, 4);
      for (int32 i = 0; i < num_t_out; i++)
        output_t_values.push_back(out_t_start + i * out_t_step);
    }
  
    input_indexes->clear();
    output_indexes->clear();
    for (size_t i = 0; i < n_x_pairs.size(); i++) {
      std::vector<int32> chosen_output_t_values;
      while (chosen_output_t_values.empty()) {
        for (size_t j = 0; j < output_t_values.size(); j++)
          if (RandInt(0, 1) != 0)
            chosen_output_t_values.push_back(output_t_values[j]);
      }
      KALDI_ASSERT(IsSortedAndUniq(chosen_output_t_values));
  
      std::set<int32> required_input_t_values,
          usable_input_t_values;
      for (size_t j = 0; j < chosen_output_t_values.size(); j++) {
        std::set<int32>::const_iterator iter;
        int32 t_out = chosen_output_t_values[j];
        for (iter = model.required_time_offsets.begin();
             iter != model.required_time_offsets.end(); iter++) {
          int32 offset = *iter;
          required_input_t_values.insert(t_out + offset);
        }
        for (iter = model.all_time_offsets.begin();
             iter != model.all_time_offsets.end(); iter++) {
          int32 offset = *iter;
          usable_input_t_values.insert(t_out + offset);
        }
      }
  
      // add to output_indexes
      for (size_t j = 0; j < chosen_output_t_values.size(); j++) {
        int32 t_out = chosen_output_t_values[j];
        Index index;
        index.n = n_x_pairs[i].first;
        index.x = n_x_pairs[i].second;
        index.t = t_out;
        output_indexes->push_back(index);
      }
  
      std::vector<int32> chosen_input_t_values(required_input_t_values.begin(),
                                               required_input_t_values.end());
      for (std::set<int32>::const_iterator iter = usable_input_t_values.begin();
           iter != usable_input_t_values.end(); ++iter) {
        int32 t = *iter;
        if (RandInt(0, 1) == 0)
          chosen_input_t_values.push_back(t);
      }
      SortAndUniq(&chosen_input_t_values);
  
      // add to input_indexes
      for (size_t j = 0; j < chosen_input_t_values.size(); j++) {
        int32 t_in = chosen_input_t_values[j];
        Index index;
        index.n = n_x_pairs[i].first;
        index.x = n_x_pairs[i].second;
        index.t = t_in;
        input_indexes->push_back(index);
      }
    }
  }
  
  
  void UnitTestTimeHeightConvolutionIo() {
    for (int32 i = 0; i < 10; i++) {
      KALDI_LOG << "iter = " << i;
      // Create a ConvolutionModel and test its I/O.
      ConvolutionModel conv_model;
      GetRandomConvolutionModel(&conv_model);
      std::ostringstream os1, os2;
      bool binary = (RandInt(0, 1) == 0);
      conv_model.Write(os1, binary);
      std::istringstream is(os1.str());
      ConvolutionModel conv_model2;
      conv_model2.Read(is, binary);
      conv_model2.Write(os2, binary);
      KALDI_ASSERT(os1.str() == os2.str() && conv_model2.Check());
    }
  }
  
  void TestComputationIo(const ConvolutionComputation &computation) {
    std::ostringstream os1, os2;
    bool binary = (RandInt(0, 1) == 0);
    computation.Write(os1, binary);
    std::istringstream is(os1.str());
    ConvolutionComputation computation2;
    computation2.Read(is, binary);
    computation2.Write(os2, binary);
    KALDI_ASSERT(os1.str() == os2.str());
    computation2.Check();
  }
  
  
  // This function exects indexes.size() == matrix->NumRows();
  // it sets to zero any row i of the matrix for which
  // indexes[i].t == kNoTime.
  void ZeroBlankRows(const std::vector<Index> &indexes,
                     CuMatrix<BaseFloat> *matrix) {
    KALDI_ASSERT(static_cast<int32>(indexes.size()) == matrix->NumRows());
    int32 num_rows = matrix->NumRows();
    if (num_rows == 0) return;
    Vector<BaseFloat> mask(num_rows, kUndefined);
    mask.Set(1.0);
    const Index *indexes_ptr = &(indexes[0]);
    BaseFloat *mask_ptr = mask.Data();
    for (int32 r = 0; r < num_rows; r++) {
      if (indexes_ptr[r].t == kNoTime)
        mask_ptr[r] = 0.0;
    }
    CuVector<BaseFloat> cu_mask;
    cu_mask.Swap(&mask);
    matrix->MulRowsVec(cu_mask);
  }
  
  // This is a 'dumb' implementation of convolution, created to compare
  // with ConvolveForward.
  void ConvolveForwardSimple(
      const ConvolutionModel &model,
      const std::vector<Index> &input_indexes,
      const std::vector<Index> &output_indexes,
      const CuMatrixBase<BaseFloat> &input_cu,
      const CuMatrixBase<BaseFloat> &params_cu,
      CuMatrixBase<BaseFloat> *output_cu) {
    // these loops will be very slow on GPU, so do it all on CPU.
    Matrix<BaseFloat> input(input_cu), params(params_cu),
        output(*output_cu);
    std::unordered_map<Index, int32, IndexHasher> index_to_row;
    int32 input_rows = input.NumRows(),
        output_rows = output.NumRows();
    for (int32 r_in = 0; r_in < input_rows; r_in++) {
      if (input_indexes[r_in].t != kNoTime) {
        index_to_row[input_indexes[r_in]] = r_in;
      }
    }
    int32 num_offsets = model.offsets.size(),
        num_filters_in = model.num_filters_in,
        num_filters_out = model.num_filters_out,
        height_in = model.height_in,
        height_out = model.height_out,
        height_subsample_out = model.height_subsample_out;
    for (int32 r_out = 0; r_out < output_rows; r_out++) {
      Index index_out = output_indexes[r_out];
      if (index_out.t == kNoTime)
        continue;
      SubVector<BaseFloat> output_row(output, r_out);
      for (int32 o = 0; o < num_offsets; o++) {
        int32 time_offset = model.offsets[o].time_offset,
            height_offset = model.offsets[o].height_offset;
        Index index_in(index_out);
        index_in.t += time_offset;
        std::unordered_map<Index, int32, IndexHasher>::const_iterator iter =
            index_to_row.find(index_in);
        if (iter != index_to_row.end()) {
          SubMatrix<BaseFloat> params_part(params, 0, params.NumRows(),
                                           o * num_filters_in, num_filters_in);
          int32 r_in = iter->second;
          SubVector<BaseFloat> input_row(input, r_in);
          for (int32 h_out_subsampled = 0;
               h_out_subsampled < height_out;
               h_out_subsampled++) {
            int32 h_out = h_out_subsampled * height_subsample_out,
                h_in = h_out + height_offset;
            if (h_in < 0 || h_in >= height_in)
              continue;
            SubVector<BaseFloat> output_part(output_row,
                                             h_out_subsampled * num_filters_out,
                                             num_filters_out),
                input_part(input_row, h_in * num_filters_in, num_filters_in);
            output_part.AddMatVec(1.0, params_part, kNoTrans, input_part, 1.0);
          }
        }
      }
    }
    output_cu->CopyFromMat(output);
  }
  
  
  
  void TestRunningComputation(const ConvolutionModel &conv_model,
                              const std::vector<Index> &input_indexes,
                              const std::vector<Index> &output_indexes,
                              const ConvolutionComputation &computation) {
    CuMatrix<BaseFloat> input(input_indexes.size(), conv_model.InputDim(),
                              kSetZero, kStrideEqualNumCols),
        output(output_indexes.size(), conv_model.OutputDim(),
               kSetZero, kStrideEqualNumCols),
        output2(output),
        params(conv_model.ParamRows(), conv_model.ParamCols());
    input.SetRandn();
    params.SetRandn();
    ZeroBlankRows(input_indexes, &input);
    ConvolveForward(computation, input, params, &output);
    ZeroBlankRows(output_indexes, &output);
  
    ConvolveForwardSimple(conv_model, input_indexes, output_indexes,
                          input, params, &output2);
    KALDI_LOG << "Tested convolution for model: "
              << conv_model.Info();
    if (!output.ApproxEqual(output2, 0.001)) {
      KALDI_LOG << "Output is: " << output;
      KALDI_LOG << "Output2 is: " << output2;
      KALDI_ERR << "Convolution test failure.";
    }
  }
  
  
  void TestDataBackprop(const ConvolutionModel &conv_model,
                        const std::vector<Index> &input_indexes,
                        const std::vector<Index> &output_indexes,
                        const ConvolutionComputation &computation) {
    CuMatrix<BaseFloat>
        input_deriv(input_indexes.size(), conv_model.InputDim(),
                    kSetZero, kStrideEqualNumCols),
        input(input_indexes.size(), conv_model.InputDim(),
              kSetZero, kStrideEqualNumCols),
        output(output_indexes.size(), conv_model.OutputDim(),
               kSetZero, kStrideEqualNumCols),
        output_deriv(output_indexes.size(), conv_model.OutputDim(),
                     kSetZero, kStrideEqualNumCols),
        params(conv_model.ParamRows(), conv_model.ParamCols());
  
    input.SetRandn();
    params.SetRandn();
    output_deriv.SetRandn();
  
    ZeroBlankRows(output_indexes, &output_deriv);
    ConvolveBackwardData(computation, params, output_deriv, &input_deriv);
    ZeroBlankRows(input_indexes, &input_deriv);
    ZeroBlankRows(input_indexes, &input);
  
    // define the objf as TraceMatMat(output_deriv, output, kTrans).
    // we can work it out from the backpropagated data-derivative.
    BaseFloat expected_objf = TraceMatMat(input_deriv, input, kTrans);
  
    ConvolveForward(computation, input, params, &output);
    ZeroBlankRows(output_indexes, &output);
  
    BaseFloat observed_objf = TraceMatMat(output, output_deriv, kTrans);
  
    KALDI_LOG << "Expected objf = " << expected_objf
              << ", observed objf = " << observed_objf;
    if (!ApproxEqual(expected_objf, observed_objf, 0.1) &&
        fabs(expected_objf) < 1.0) {
      KALDI_ERR << "Difference in objf too large.";
    }
  }
  
  
  void TestParamsBackprop(const ConvolutionModel &conv_model,
                          const std::vector<Index> &input_indexes,
                          const std::vector<Index> &output_indexes,
                          const ConvolutionComputation &computation) {
    CuMatrix<BaseFloat>
        input(input_indexes.size(), conv_model.InputDim(),
              kSetZero, kStrideEqualNumCols),
        output(output_indexes.size(), conv_model.OutputDim(),
               kSetZero, kStrideEqualNumCols),
        output_deriv(output_indexes.size(), conv_model.OutputDim(),
                     kSetZero, kStrideEqualNumCols),
        params(conv_model.ParamRows(), conv_model.ParamCols()),
        params_deriv(conv_model.ParamRows(), conv_model.ParamCols());
  
    input.SetRandn();
    params.SetRandn();
    output_deriv.SetRandn();
  
    BaseFloat alpha = 0.5 * RandInt(1, 3);
  
    ZeroBlankRows(output_indexes, &output_deriv);
    ZeroBlankRows(input_indexes, &input);
  
    ConvolveBackwardParams(computation, input, output_deriv, alpha,
                           &params_deriv);
  
    BaseFloat expected_objf = TraceMatMat(params_deriv, params, kTrans) / alpha;
  
    ConvolveForward(computation, input, params, &output);
  
    ZeroBlankRows(output_indexes, &output);
  
    BaseFloat observed_objf = TraceMatMat(output, output_deriv, kTrans);
  
    KALDI_LOG << "Expected objf = " << expected_objf
              << ", observed objf = " << observed_objf;
    if (!ApproxEqual(expected_objf, observed_objf, 0.1) &&
        fabs(expected_objf) < 1.0) {
      KALDI_ERR << "Difference in objf too large.";
    }
  }
  
  
  
  void UnitTestTimeHeightConvolutionCompile() {
    for (int32 i = 0; i < 10; i++) {
      KALDI_LOG << "iter = " << i;
      // Create a ConvolutionModel
      ConvolutionModel conv_model;
      GetRandomConvolutionModel(&conv_model);
      std::vector<Index> input_indexes, output_indexes;
      GetRandomConvolutionIndexes(conv_model, &input_indexes, &output_indexes);
  
      ConvolutionComputationOptions opts;
      ConvolutionComputation computation;
      std::vector<Index> input_indexes_modified, output_indexes_modified;
      CompileConvolutionComputation(conv_model, input_indexes, output_indexes,
                                    opts, &computation,
                                    &input_indexes_modified,
                                    &output_indexes_modified);
      TestComputationIo(computation);
      TestRunningComputation(conv_model,
                             input_indexes_modified,
                             output_indexes_modified,
                             computation);
      TestDataBackprop(conv_model,
                       input_indexes_modified,
                       output_indexes_modified,
                       computation);
      TestParamsBackprop(conv_model,
                         input_indexes_modified,
                         output_indexes_modified,
                         computation);
      std::ostringstream os;
      os << "
  Input-indexes: ";
      WriteIndexVector(os, false, input_indexes);
      os << "
  Input-indexes-modified: ";
      WriteIndexVector(os, false, input_indexes_modified);
      os << "
  Output-indexes: ";
      WriteIndexVector(os, false, output_indexes);
      os << "
  Output-indexes-modified: ";
      WriteIndexVector(os, false, output_indexes_modified);
      KALDI_LOG << os.str();
    }
  }
  
  
  void UnitTestTimeHeightConvolution() {
    UnitTestTimeHeightConvolutionIo();
    UnitTestTimeHeightConvolutionCompile();
  }
  
  
  
  } // namespace time_height_convolution
  } // namespace nnet3
  } // namespace kaldi
  
  
  int main() {
    using namespace kaldi;
    using namespace kaldi::nnet3;
    using namespace kaldi::nnet3::time_height_convolution;
    for (int32 loop = 0; loop < 2; loop++) {
  #if HAVE_CUDA == 1
      CuDevice::Instantiate().SetDebugStrideMode(true);
      if (loop == 0)
        CuDevice::Instantiate().SelectGpuId("no"); // -1 means no GPU
      else
        CuDevice::Instantiate().SelectGpuId("optional"); // -2 .. automatic selection
  #endif
      for (int32 i = 0; i < 5; i++) {
        UnitTestTimeHeightConvolution();
      }
    }
  }