// nnet3/nnet-compile-looped.cc

// Copyright      2016  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-compile-looped.h"
#include "nnet3/nnet-optimize-utils.h"
#include "nnet3/nnet-utils.h"

namespace kaldi {
namespace nnet3 {


void ModifyNnetIvectorPeriod(int32 ivector_period,
                             Nnet *nnet) {
  KALDI_ASSERT(ivector_period > 0);
  std::vector<std::string> config_lines;
  nnet->GetConfigLines(false, &config_lines);
  std::ostringstream config_to_read;
  for (size_t i = 0; i < config_lines.size(); i++) {
    std::string s = config_lines[i];
    ConfigLine config_line;
    bool b = config_line.ParseLine(config_lines[i]);
    KALDI_ASSERT(b && "Could not parse config line.");
    if (config_line.FirstToken() == "component-node") {
      // What we're trying to do here is: find a line like:
      //  component-node name=foo component=foo input=Append(bar, ReplaceIndex(ivector, t, 0))
      // we want to replace it with something like:
      // component-node name=foo component=foo input=Append(bar, ReplaceIndex(ivector, t, 0))
      // .. and we want this to also work if instead of 'ivector' it has something like
      // Scale(0.5, ivector).  We assume that ReplaceIndex() expressions only occur in this
      // type of context.
      std::string whole_line = config_lines[i];
      std::string to_search_for = "ReplaceIndex(";
      std::string::size_type to_search_for_size = to_search_for.size();
      std::string::size_type pos = whole_line.find(to_search_for);
      if (pos != std::string::npos) {
        std::string::size_type comma_pos = whole_line.find(", t, 0)", pos);
        if (comma_pos != std::string::npos) {
          // if the line contained ReplaceIndex(ivector, t, 0),
          // descriptor_name would now be 'ivector'.
          std::string descriptor_name =
              whole_line.substr(pos + to_search_for_size,
                                comma_pos - (pos + to_search_for_size));
          // Note: 7, below, is the size of: ", t, 0)".
          std::string::size_type end_pos = comma_pos + 7;
          std::string::size_type expr_size = end_pos - pos;
          // e.g. expr_size would be strlen("ReplaceIndex(ivector, t, 0)").
          std::ostringstream to_replace_with;
          to_replace_with << "Round(" << descriptor_name << ", " << ivector_period << ")";
          whole_line.replace(pos, expr_size, to_replace_with.str());
          config_to_read << whole_line << "\n";
        } else {
          KALDI_ERR << "Could not process the ReplaceIndex expression in: "
                    << whole_line;
        }
      }
    }
  }
  if (!config_to_read.str().empty()) {
    std::istringstream is(config_to_read.str());
    nnet->ReadConfig(is);
  }
}


int32 GetChunkSize(const Nnet &nnet,
                   int32 frame_subsampling_factor,
                   int32 advised_chunk_size) {
  int32 modulus = nnet.Modulus();
  KALDI_ASSERT(modulus > 0 && frame_subsampling_factor > 0 &&
               advised_chunk_size > 0);
  int32 chunk_size = advised_chunk_size;
  while (1) {
    if (chunk_size % modulus == 0 &&
        chunk_size % frame_subsampling_factor == 0)
      return chunk_size;
    chunk_size++;
  }
}


/// Mod(m, n), defined for integers m and n where n > 0, returns
/// the modulus m % n, defined as the integer 0 <= i < n
/// such that i and m are congruent modulo n; for instance,
/// Mod(13, 10) = 3.
/// This is like the % operation in C/C++, except that it always returns a
/// positive value even for negative m; in 99% of cases where it makes a
/// difference, this is what you want.  In the C/C++ standard, the sign of a % b
/// for negative a is not specified (except by relation with the division '/'
/// operator), but in practice it would be <= 0 for almost all implementations.
template<class I> I  Mod(I m, I n) {
  I ans = m % n;
  if (ans < 0) ans += n;
  return ans;
}


static void CreateComputationRequestInternal(
    int32 begin_input_t, int32 end_input_t,
    int32 begin_output_t, int32 end_output_t,
    int32 num_sequences,
    int32 frame_subsampling_factor,
    const std::set<int32> &ivector_times,
    ComputationRequest *request) {
  request->inputs.reserve(2);
  request->inputs.clear();
  request->inputs.resize(1 + (ivector_times.empty() ? 0 : 1));
  request->inputs[0].name = "input";
  request->inputs[0].has_deriv = false;
  request->outputs.clear();
  request->outputs.resize(1);
  request->outputs[0].name = "output";
  request->outputs[0].has_deriv = false;
  if (!ivector_times.empty()) {
    request->inputs[1].name = "ivector";
    request->inputs[1].has_deriv = false;
  }

  // in the computation request the 'n' indexes (the sequence/utterance indexes)
  // have the larger stride than 't', although this is opposite to the way it's
  // done inside the computation.  This is for user convenience where it may be
  // easier to deal with submatrixes per sequence.
  for (int32 n = 0; n < num_sequences; n++) {
    int32 x = 0;
    for (int32 t = begin_input_t; t < end_input_t; t++) {
      request->inputs[0].indexes.push_back(Index(n, t, x));
    }
    for (int32 t = begin_output_t;
         t < end_output_t;
         t += frame_subsampling_factor)
      request->outputs[0].indexes.push_back(Index(n, t, x));
  }
  if (!ivector_times.empty()) {
    request->inputs.resize(2);
    request->inputs[1].name = "ivector";
    request->inputs[1].has_deriv = false;
    for (int32 n = 0; n < num_sequences; n++) {
      // note: std::sets store things in sorted order.
      for (std::set<int32>::const_iterator iter = ivector_times.begin();
           iter != ivector_times.end(); ++iter) {
        int32 t = *iter, x = 0;
        request->inputs[1].indexes.push_back(Index(n, t, x));
      }
    }
  }
}


void CreateLoopedComputationRequest(const Nnet &nnet,
                                    int32 chunk_size,
                                    int32 frame_subsampling_factor,
                                    int32 ivector_period,
                                    int32 left_context_begin,
                                    int32 right_context,
                                    int32 num_sequences,
                                    ComputationRequest *request1,
                                    ComputationRequest *request2,
                                    ComputationRequest *request3) {
  bool has_ivector = (nnet.InputDim("ivector") > 0);
  KALDI_ASSERT(chunk_size % frame_subsampling_factor == 0 &&
               chunk_size % nnet.Modulus() == 0 &&
               chunk_size % ivector_period == 0);
  KALDI_ASSERT(left_context_begin >= 0 && right_context >= 0);
  // note, 'end' is one past the last one.
  int32 chunk1_input_begin_t = - left_context_begin,
      chunk1_input_end_t = chunk_size + right_context,
      chunk2_input_begin_t = chunk1_input_end_t,
      chunk2_input_end_t = chunk2_input_begin_t + chunk_size,
      chunk3_input_begin_t = chunk2_input_end_t,
      chunk3_input_end_t = chunk3_input_begin_t + chunk_size;


  // work out the times at which i-vectors are required.
  std::set<int32> ivector_times1, ivector_times2, ivector_times3;
  if (has_ivector) {
    for (int32 t = chunk1_input_begin_t; t < chunk1_input_end_t; t++) {
      int32 ivector_t = t - Mod(t, ivector_period);
      ivector_times1.insert(ivector_t);
    }
    for (int32 t = chunk2_input_begin_t; t < chunk2_input_end_t; t++) {
      int32 ivector_t = t - Mod(t, ivector_period);
      if (ivector_times2.count(ivector_t) == 0 &&
	  ivector_times1.count(ivector_t) == 0)
        ivector_times2.insert(ivector_t);
    }
    for (int32 t = chunk3_input_begin_t; t < chunk3_input_end_t; t++) {
      int32 ivector_t = t - Mod(t, ivector_period);
      if (ivector_times3.count(ivector_t) == 0 &&
          ivector_times2.count(ivector_t) == 0 &&
	  ivector_times1.count(ivector_t) == 0)
        ivector_times3.insert(ivector_t);
    }
  }

  CreateComputationRequestInternal(
      chunk1_input_begin_t, chunk1_input_end_t,
      0, chunk_size,
      num_sequences, frame_subsampling_factor,
      ivector_times1,
      request1);

  CreateComputationRequestInternal(
      chunk2_input_begin_t, chunk2_input_end_t,
      chunk_size, chunk_size * 2,
      num_sequences, frame_subsampling_factor,
      ivector_times2,
      request2);

  CreateComputationRequestInternal(
      chunk3_input_begin_t, chunk3_input_end_t,
      chunk_size * 2, chunk_size * 3,
      num_sequences, frame_subsampling_factor,
      ivector_times3,
      request3);

}



void AddTimeOffsetToComputationRequest(int32 t_offset,
                                       ComputationRequest *request) {
  for (size_t i = 0; i < request->inputs.size(); i++) {
    size_t size = request->inputs[i].indexes.size();
    for (size_t j = 0; j < size; j++)
      request->inputs[i].indexes[j].t += t_offset;
  }
  for (size_t i = 0; i < request->outputs.size(); i++) {
    size_t size = request->outputs[i].indexes.size();
    for (size_t j = 0; j < size; j++)
      request->outputs[i].indexes[j].t += t_offset;
  }
}



static bool ExtrapolateComputationRequest(
    const ComputationRequest &request1,
    const ComputationRequest &request2,
    ComputationRequest *request3) {
  // accepts two computation requests 'request1' and 'request2' that
  // must be identical except for a time offset, and creates 'request3'
  // that is the extrapolation of the next term in sequence.
  *request3 = request2;
  KALDI_ASSERT(!request1.inputs.empty() && !request1.inputs[0].indexes.empty() &&
               !request2.inputs.empty() && !request2.inputs[0].indexes.empty());
  int32 t_offset = request2.inputs[0].indexes[0].t -
      request1.inputs[0].indexes[0].t;
  // the following is just to make sure that the inputs are structurally
  // equivalent.
  AddTimeOffsetToComputationRequest(-t_offset, request3);
  if (!(*request3 == request1))
    return false;  // there is somse structural difference, or
                   // the time offset is not consistent.
  // the following reverses the last call to AddTimeOffsetToComputationRequest,
  // then adds the offset we want.
  AddTimeOffsetToComputationRequest(2 * t_offset, request3);
  return true;
}


/* Internal version of CompileLooped where
   you specify the the number of computation requests (must be >= 3).
   Returns true on success.
   It's possible for the optimization to fail if you give too small
   a value of 'num_requests' (this depends on the network topology),
   and in that case this function will return false and you should re-try
   with a higher value of num_requests.
 */
static bool CompileLoopedInternal(
    const Nnet &nnet,
    NnetOptimizeOptions optimize_opts,
    const ComputationRequest &request1,
    const ComputationRequest &request2,
    const ComputationRequest &request3,
    int32 num_requests,
    NnetComputation *computation) {

  KALDI_ASSERT(num_requests >= 3);
  std::vector<ComputationRequest> extra_requests(num_requests - 3);
  const ComputationRequest *prev_request = &request2;
  const ComputationRequest *cur_request = &request3;
  for (int32 i = 0; i < num_requests - 3; i++) {
    if (!ExtrapolateComputationRequest(*prev_request, *cur_request,
                                       &(extra_requests[i]))) {
      KALDI_LOG << "prev_request is:";
      prev_request->Print(std::cerr);
      KALDI_LOG << "cur_request is:";
      cur_request->Print(std::cerr);
      KALDI_ERR << "Computation requests do not have the right relationship";
    }
    prev_request = cur_request;
    cur_request = &(extra_requests[i]);
  }

  std::vector<const ComputationRequest*> requests;
  requests.push_back(&request1);
  requests.push_back(&request2);
  requests.push_back(&request3);
  for (int32 i = 0; i < num_requests - 3; i++)
    requests.push_back(&(extra_requests[i]));
  Compiler compiler(requests, nnet);
  CompilerOptions compiler_opts;
  compiler.CreateComputation(compiler_opts, computation);
  optimize_opts.optimize_looped_computation = true;

  int32 dont_really_care = MaxOutputTimeInRequest(request3);
  Optimize(optimize_opts, nnet,
           dont_really_care, computation);

  return computation->commands.size() != 0 &&
      computation->commands.back().command_type == kGotoLabel;
}

void CompileLooped(const Nnet &nnet,
                   const NnetOptimizeOptions &optimize_opts,
                   const ComputationRequest &request1,
                   const ComputationRequest &request2,
                   const ComputationRequest &request3,
                   NnetComputation *computation) {
  int32 num_requests1 = 5, factor = 2, max_requests = 100,
      num_requests;

  Timer timer;

  for (num_requests = num_requests1; num_requests <= max_requests;
       num_requests *= factor) {
    if (CompileLoopedInternal(nnet, optimize_opts,
                             request1, request2, request3,
                             num_requests, computation)) {
      KALDI_LOG << "Spent " << timer.Elapsed()
                << " seconds in looped compilation.";
      return;
    } else {
      KALDI_VLOG(2) << "Looped compilation failed with "
                    << num_requests << " requests, trying "
                    << (num_requests * factor);
    }
  }
  KALDI_ERR << "Looped compilation failed with "
            << (num_requests/factor) << " requests, which "
            << "we expect should be enough... something "
            << "went wrong.";
}


void CreateLoopedComputationRequestSimple(const Nnet &nnet,
                                          int32 chunk_size,
                                          int32 frame_subsampling_factor,
                                          int32 ivector_period,
                                          int32 extra_left_context_begin,
                                          int32 extra_right_context,
                                          int32 num_sequences,
                                          ComputationRequest *request1,
                                          ComputationRequest *request2,
                                          ComputationRequest *request3) {
  int32 left_context, right_context;
  ComputeSimpleNnetContext(nnet, &left_context, &right_context);

  CreateLoopedComputationRequest(nnet, chunk_size, frame_subsampling_factor,
                                 ivector_period,
                                 extra_left_context_begin + left_context,
                                 extra_right_context + right_context,
                                 num_sequences, request1, request2, request3);
}

} // namespace nnet3
} // namespace kaldi