// nnet3/nnet-discriminative-example.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 <cmath>
  #include "nnet3/nnet-discriminative-example.h"
  #include "nnet3/nnet-example-utils.h"
  
  namespace kaldi {
  namespace nnet3 {
  
  
  void NnetDiscriminativeSupervision::Write(std::ostream &os, bool binary) const {
    CheckDim();
    WriteToken(os, binary, "<NnetDiscriminativeSup>");
    WriteToken(os, binary, name);
    WriteIndexVector(os, binary, indexes);
    supervision.Write(os, binary);
    WriteToken(os, binary, "<DW>");  // for DerivWeights.  Want to save space.
    WriteVectorAsChar(os, binary, deriv_weights);
    WriteToken(os, binary, "</NnetDiscriminativeSup>");
  }
  
  bool NnetDiscriminativeSupervision::operator == (const NnetDiscriminativeSupervision &other) const {
    return name == other.name && indexes == other.indexes &&
        supervision == other.supervision &&
        deriv_weights.ApproxEqual(other.deriv_weights);
  }
  
  void NnetDiscriminativeSupervision::Read(std::istream &is, bool binary) {
    ExpectToken(is, binary, "<NnetDiscriminativeSup>");
    ReadToken(is, binary, &name);
    ReadIndexVector(is, binary, &indexes);
    supervision.Read(is, binary);
    ExpectToken(is, binary, "<DW>");
    ReadVectorAsChar(is, binary, &deriv_weights);
    ExpectToken(is, binary, "</NnetDiscriminativeSup>");
    CheckDim();
  }
  
  
  void NnetDiscriminativeSupervision::CheckDim() const {
    if (supervision.frames_per_sequence == -1) {
      // this object has not been set up.
      KALDI_ASSERT(indexes.empty());
      return;
    }
    KALDI_ASSERT(indexes.size() == supervision.num_sequences *
                 supervision.frames_per_sequence && !indexes.empty() &&
                 supervision.frames_per_sequence > 1);
    int32 first_frame = indexes[0].t,
        frame_skip = indexes[supervision.num_sequences].t - first_frame,
        num_sequences = supervision.num_sequences,
        frames_per_sequence = supervision.frames_per_sequence;
    int32 k = 0;
    for (int32 i = 0; i < frames_per_sequence; i++) {
      for (int32 j = 0; j < num_sequences; j++,k++) {
        int32 n = j, t = i * frame_skip + first_frame, x = 0;
        Index index(n, t, x);
        KALDI_ASSERT(indexes[k] == index);
      }
    }
    if (deriv_weights.Dim() != 0) {
      KALDI_ASSERT(deriv_weights.Dim() == indexes.size());
      KALDI_ASSERT(deriv_weights.Min() >= 0.0 &&
                   deriv_weights.Max() <= 1.0);
    }
  }
  
  NnetDiscriminativeSupervision::NnetDiscriminativeSupervision(const NnetDiscriminativeSupervision &other):
      name(other.name),
      indexes(other.indexes),
      supervision(other.supervision),
      deriv_weights(other.deriv_weights) { CheckDim(); }
  
  NnetDiscriminativeSupervision::NnetDiscriminativeSupervision(
      const std::string &name,
      const discriminative::DiscriminativeSupervision &supervision,
      const VectorBase<BaseFloat> &deriv_weights,
      int32 first_frame,
      int32 frame_skip):
      name(name),
      supervision(supervision),
      deriv_weights(deriv_weights) {
    // note: this will set the 'x' index to zero.
    indexes.resize(supervision.num_sequences *
                   supervision.frames_per_sequence);
    int32 k = 0, num_sequences = supervision.num_sequences,
        frames_per_sequence = supervision.frames_per_sequence;
    for (int32 i = 0; i < frames_per_sequence; i++) {
      for (int32 j = 0; j < num_sequences; j++,k++) {
        indexes[k].n = j;
        indexes[k].t = i * frame_skip + first_frame;
      }
    }
    KALDI_ASSERT(k == indexes.size());
    CheckDim();
  }
  
  void NnetDiscriminativeSupervision::Swap(NnetDiscriminativeSupervision *other) {
    name.swap(other->name);
    indexes.swap(other->indexes);
    supervision.Swap(&(other->supervision));
    deriv_weights.Swap(&(other->deriv_weights));
    if (RandInt(0, 5) == 0)
      CheckDim();
  }
  
  
  void NnetDiscriminativeExample::Write(std::ostream &os, bool binary) const {
    // Note: weight, label, input_frames and spk_info are members.  This is a
    // struct.
    WriteToken(os, binary, "<Nnet3DiscriminativeEg>");
    WriteToken(os, binary, "<NumInputs>");
    int32 size = inputs.size();
    WriteBasicType(os, binary, size);
    KALDI_ASSERT(size > 0 && "Attempting to write NnetDiscriminativeExample with no inputs");
    if (!binary) os << '
  ';
    for (int32 i = 0; i < size; i++) {
      inputs[i].Write(os, binary);
      if (!binary) os << '
  ';
    }
    WriteToken(os, binary, "<NumOutputs>");
    size = outputs.size();
    WriteBasicType(os, binary, size);
    KALDI_ASSERT(size > 0 && "Attempting to write NnetDiscriminativeExample with no outputs");
    if (!binary) os << '
  ';
    for (int32 i = 0; i < size; i++) {
      outputs[i].Write(os, binary);
      if (!binary) os << '
  ';
    }
    WriteToken(os, binary, "</Nnet3DiscriminativeEg>");
  }
  
  void NnetDiscriminativeExample::Read(std::istream &is, bool binary) {
    ExpectToken(is, binary, "<Nnet3DiscriminativeEg>");
    ExpectToken(is, binary, "<NumInputs>");
    int32 size;
    ReadBasicType(is, binary, &size);
    if (size < 1 || size > 1000000)
      KALDI_ERR << "Invalid size " << size;
    inputs.resize(size);
    for (int32 i = 0; i < size; i++)
      inputs[i].Read(is, binary);
    ExpectToken(is, binary, "<NumOutputs>");
    ReadBasicType(is, binary, &size);
    if (size < 1 || size > 1000000)
      KALDI_ERR << "Invalid size " << size;
    outputs.resize(size);
    for (int32 i = 0; i < size; i++)
      outputs[i].Read(is, binary);
    ExpectToken(is, binary, "</Nnet3DiscriminativeEg>");
  }
  
  void NnetDiscriminativeExample::Swap(NnetDiscriminativeExample *other) {
    inputs.swap(other->inputs);
    outputs.swap(other->outputs);
  }
  
  void NnetDiscriminativeExample::Compress() {
    std::vector<NnetIo>::iterator iter = inputs.begin(), end = inputs.end();
    // calling features.Compress() will do nothing if they are sparse or already
    // compressed.
    for (; iter != end; ++iter) iter->features.Compress();
  }
  
  NnetDiscriminativeExample::NnetDiscriminativeExample(const NnetDiscriminativeExample &other):
      inputs(other.inputs), outputs(other.outputs) { }
  
  void MergeSupervision(
      const std::vector<const NnetDiscriminativeSupervision*> &inputs,
      NnetDiscriminativeSupervision *output) {
    int32 num_inputs = inputs.size(),
        num_indexes = 0;
    for (int32 n = 0; n < num_inputs; n++) {
      KALDI_ASSERT(inputs[n]->name == inputs[0]->name);
      num_indexes += inputs[n]->indexes.size();
    }
    output->name = inputs[0]->name;
    std::vector<const discriminative::DiscriminativeSupervision*> input_supervision;
    input_supervision.reserve(inputs.size());
    for (int32 n = 0; n < num_inputs; n++)
      input_supervision.push_back(&(inputs[n]->supervision));
    discriminative::DiscriminativeSupervision output_supervision;
    discriminative::MergeSupervision(input_supervision,
                           &output_supervision);
    output->supervision.Swap(&(output_supervision));
  
    output->indexes.clear();
    output->indexes.reserve(num_indexes);
    for (int32 n = 0; n < num_inputs; n++) {
      const std::vector<Index> &src_indexes = inputs[n]->indexes;
      int32 cur_size = output->indexes.size();
      output->indexes.insert(output->indexes.end(),
                             src_indexes.begin(), src_indexes.end());
      std::vector<Index>::iterator iter = output->indexes.begin() + cur_size,
          end = output->indexes.end();
      // change the 'n' index to correspond to the index into 'input'.
      // Each example gets a different 'n' value, starting from 0.
      for (; iter != end; ++iter) {
        KALDI_ASSERT(iter->n == 0 && "Merging already-merged discriminative egs");
        iter->n = n;
      }
    }
    KALDI_ASSERT(output->indexes.size() == num_indexes);
    // OK, at this point the 'indexes' will be in the wrong order,
    // because they should be first sorted by 't' and next by 'n'.
    // 'sort' will fix this, due to the operator < on type Index.
    // TODO: Is this required?
    std::sort(output->indexes.begin(), output->indexes.end());
  
    // merge the deriv_weights.
    if (inputs[0]->deriv_weights.Dim() != 0) {
      int32 frames_per_sequence = inputs[0]->deriv_weights.Dim();
      output->deriv_weights.Resize(output->indexes.size(), kUndefined);
      KALDI_ASSERT(output->deriv_weights.Dim() ==
                   frames_per_sequence * num_inputs);
      for (int32 n = 0; n < num_inputs; n++) {
        const Vector<BaseFloat> &src_deriv_weights = inputs[n]->deriv_weights;
        KALDI_ASSERT(src_deriv_weights.Dim() == frames_per_sequence);
        // the ordering of the deriv_weights corresponds to the ordering of the
        // Indexes, where the time dimension has the greater stride.
        for (int32 t = 0; t < frames_per_sequence; t++) {
          output->deriv_weights(t * num_inputs + n) = src_deriv_weights(t);
        }
      }
    }
    output->CheckDim();
  }
  
  
  void MergeDiscriminativeExamples(
      bool compress,
      std::vector<NnetDiscriminativeExample> *input,
      NnetDiscriminativeExample *output) {
    int32 num_examples = input->size();
    KALDI_ASSERT(num_examples > 0);
    // we temporarily make the input-features in 'input' look like regular
    // NnetExamples, so that we can recycle the
    // MergeExamples() function.
    std::vector<NnetExample> eg_inputs(num_examples);
    for (int32 i = 0; i < num_examples; i++)
      eg_inputs[i].io.swap((*input)[i].inputs);
    NnetExample eg_output;
    MergeExamples(eg_inputs, compress, &eg_output);
    // swap the inputs back so that they are not really changed.
    for (int32 i = 0; i < num_examples; i++)
      eg_inputs[i].io.swap((*input)[i].inputs);
    // write to 'output->inputs'
    eg_output.io.swap(output->inputs);
  
    // Now deal with the discriminative-supervision 'outputs'.  There will
    // normally be just one of these, with name "output", but we
    // handle the more general case.
    int32 num_output_names = (*input)[0].outputs.size();
    output->outputs.resize(num_output_names);
    for (int32 i = 0; i < num_output_names; i++) {
      std::vector<const NnetDiscriminativeSupervision*> to_merge(num_examples);
      for (int32 j = 0; j < num_examples; j++) {
        KALDI_ASSERT((*input)[j].outputs.size() == num_output_names);
        to_merge[j] = &((*input)[j].outputs[i]);
      }
      MergeSupervision(to_merge,
                       &(output->outputs[i]));
    }
  }
  
  
  void GetDiscriminativeComputationRequest(const Nnet &nnet,
                                           const NnetDiscriminativeExample &eg,
                                           bool need_model_derivative,
                                           bool store_component_stats,
                                           bool use_xent_regularization,
                                           bool use_xent_derivative,
                                           ComputationRequest *request) {
    request->inputs.clear();
    request->inputs.reserve(eg.inputs.size());
    request->outputs.clear();
    request->outputs.reserve(eg.outputs.size());
    request->need_model_derivative = need_model_derivative;
    request->store_component_stats = store_component_stats;
    for (size_t i = 0; i < eg.inputs.size(); i++) {
      const NnetIo &io = eg.inputs[i];
      const std::string &name = io.name;
      int32 node_index = nnet.GetNodeIndex(name);
      if (node_index == -1 &&
          !nnet.IsInputNode(node_index))
        KALDI_ERR << "Nnet example has input named '" << name
                  << "', but no such input node is in the network.";
  
      request->inputs.resize(request->inputs.size() + 1);
      IoSpecification &io_spec = request->inputs.back();
      io_spec.name = name;
      io_spec.indexes = io.indexes;
      io_spec.has_deriv = false;
    }
    for (size_t i = 0; i < eg.outputs.size(); i++) {
      // there will normally be exactly one output , named "output"
      const NnetDiscriminativeSupervision &sup = eg.outputs[i];
      const std::string &name = sup.name;
      int32 node_index = nnet.GetNodeIndex(name);
      if (node_index == -1 &&
          !nnet.IsOutputNode(node_index))
        KALDI_ERR << "Nnet example has output named '" << name
                  << "', but no such output node is in the network.";
      request->outputs.resize(request->outputs.size() + 1);
      IoSpecification &io_spec = request->outputs.back();
      io_spec.name = name;
      io_spec.indexes = sup.indexes;
      io_spec.has_deriv = need_model_derivative;
  
      if (use_xent_regularization) {
        size_t cur_size = request->outputs.size();
        request->outputs.resize(cur_size + 1);
        IoSpecification &io_spec = request->outputs[cur_size - 1],
            &io_spec_xent = request->outputs[cur_size];
        // the IoSpecification for the -xent output is the same
        // as for the regular output, except for its name which has
        // the -xent suffix (and the has_deriv member may differ).
        io_spec_xent = io_spec;
        io_spec_xent.name = name + "-xent";
        io_spec_xent.has_deriv = use_xent_derivative;
      }
    }
    // check to see if something went wrong.
    if (request->inputs.empty())
      KALDI_ERR << "No inputs in computation request.";
    if (request->outputs.empty())
      KALDI_ERR << "No outputs in computation request.";
  }
  
  void ShiftDiscriminativeExampleTimes(int32 frame_shift,
                              const std::vector<std::string> &exclude_names,
                              NnetDiscriminativeExample *eg) {
    std::vector<NnetIo>::iterator input_iter = eg->inputs.begin(),
        input_end = eg->inputs.end();
    for (; input_iter != input_end; ++input_iter) {
      bool must_exclude = false;
      std::vector<string>::const_iterator exclude_iter = exclude_names.begin(),
          exclude_end = exclude_names.end();
      for (; exclude_iter != exclude_end; ++exclude_iter)
        if (input_iter->name == *exclude_iter)
          must_exclude = true;
      if (!must_exclude) {
        std::vector<Index>::iterator indexes_iter = input_iter->indexes.begin(),
            indexes_end = input_iter->indexes.end();
        for (; indexes_iter != indexes_end; ++indexes_iter)
          indexes_iter->t += frame_shift;
      }
    }
    // note: we'll normally choose a small enough shift that the output-data
    // shift will be zero after dividing by frame_subsampling_factor
    // (e.g. frame_subsampling_factor == 3 and shift = 0 or 1.
    std::vector<NnetDiscriminativeSupervision>::iterator
        sup_iter = eg->outputs.begin(),
        sup_end = eg->outputs.end();
    for (; sup_iter != sup_end; ++sup_iter) {
      std::vector<Index> &indexes = sup_iter->indexes;
      KALDI_ASSERT(indexes.size() >= 2 && indexes[0].n == indexes[1].n &&
                   indexes[0].x == indexes[1].x);
      int32 frame_subsampling_factor = indexes[1].t - indexes[0].t;
      KALDI_ASSERT(frame_subsampling_factor > 0);
  
      // We need to shift by a multiple of frame_subsampling_factor.
      // Round to the closest multiple.
      int32 supervision_frame_shift =
          frame_subsampling_factor *
          std::floor(0.5 + (frame_shift * 1.0 / frame_subsampling_factor));
      if (supervision_frame_shift == 0)
        continue;
      std::vector<Index>::iterator indexes_iter = indexes.begin(),
          indexes_end = indexes.end();
      for (; indexes_iter != indexes_end; ++indexes_iter)
        indexes_iter->t += supervision_frame_shift;
    }
  }
  
  size_t NnetDiscriminativeExampleStructureHasher::operator () (
      const NnetDiscriminativeExample &eg) const noexcept {
    // these numbers were chosen at random from a list of primes.
    NnetIoStructureHasher io_hasher;
    size_t size = eg.inputs.size(), ans = size * 35099;
    for (size_t i = 0; i < size; i++)
      ans = ans * 19157 + io_hasher(eg.inputs[i]);
    for (size_t i = 0; i < eg.outputs.size(); i++) {
      const NnetDiscriminativeSupervision &sup = eg.outputs[i];
      StringHasher string_hasher;
      IndexVectorHasher indexes_hasher;
      ans = ans * 17957 +
          string_hasher(sup.name) + indexes_hasher(sup.indexes);
    }
    return ans;
  }
  
  bool NnetDiscriminativeExampleStructureCompare::operator () (
      const NnetDiscriminativeExample &a,
      const NnetDiscriminativeExample &b) const {
    NnetIoStructureCompare io_compare;
    if (a.inputs.size() != b.inputs.size() ||
        a.outputs.size() != b.outputs.size())
      return false;
    size_t size = a.inputs.size();
    for (size_t i = 0; i < size; i++)
      if (!io_compare(a.inputs[i], b.inputs[i]))
        return false;
    size = a.outputs.size();
    for (size_t i = 0; i < size; i++)
      if (a.outputs[i].name != b.outputs[i].name ||
          a.outputs[i].indexes != b.outputs[i].indexes)
        return false;
    return true;
  }
  
  
  int32 GetNnetDiscriminativeExampleSize(const NnetDiscriminativeExample &a) {
    int32 ans = 0;
    for (size_t i = 0; i < a.inputs.size(); i++) {
      int32 s = a.inputs[i].indexes.size();
      if (s > ans)
        ans = s;
    }
    for (size_t i = 0; i < a.outputs.size(); i++) {
      int32 s = a.outputs[i].indexes.size();
      if (s > ans)
        ans = s;
    }
    return ans;
  }
  
  
  DiscriminativeExampleMerger::DiscriminativeExampleMerger(const ExampleMergingConfig &config,
                               NnetDiscriminativeExampleWriter *writer):
      finished_(false), num_egs_written_(0),
      config_(config), writer_(writer) { }
  
  
  void DiscriminativeExampleMerger::AcceptExample(NnetDiscriminativeExample *eg) {
    KALDI_ASSERT(!finished_);
    // If an eg with the same structure as 'eg' is already a key in the
    // map, it won't be replaced, but if it's new it will be made
    // the key.  Also we remove the key before making the vector empty.
    // This way we ensure that the eg in the key is always the first
    // element of the vector.
    std::vector<NnetDiscriminativeExample*> &vec = eg_to_egs_[eg];
    vec.push_back(eg);
    int32 eg_size = GetNnetDiscriminativeExampleSize(*eg),
        num_available = vec.size();
    bool input_ended = false;
    int32 minibatch_size = config_.MinibatchSize(eg_size, num_available,
                                                 input_ended);
    if (minibatch_size != 0) {  // we need to write out a merged eg.
      KALDI_ASSERT(minibatch_size == num_available);
  
      std::vector<NnetDiscriminativeExample*> vec_copy(vec);
      eg_to_egs_.erase(eg);
  
      // MergeDiscriminativeExamples() expects a vector of NnetDiscriminativeExample, not of pointers,
      // so use swap to create that without doing any real work.
      std::vector<NnetDiscriminativeExample> egs_to_merge(minibatch_size);
      for (int32 i = 0; i < minibatch_size; i++) {
        egs_to_merge[i].Swap(vec_copy[i]);
        delete vec_copy[i];  // we owned those pointers.
      }
      WriteMinibatch(&egs_to_merge);
    }
  }
  
  void DiscriminativeExampleMerger::WriteMinibatch(
      std::vector<NnetDiscriminativeExample> *egs) {
    KALDI_ASSERT(!egs->empty());
    int32 eg_size = GetNnetDiscriminativeExampleSize((*egs)[0]);
    NnetDiscriminativeExampleStructureHasher eg_hasher;
    size_t structure_hash = eg_hasher((*egs)[0]);
    int32 minibatch_size = egs->size();
    stats_.WroteExample(eg_size, structure_hash, minibatch_size);
    NnetDiscriminativeExample merged_eg;
    MergeDiscriminativeExamples(config_.compress, egs, &merged_eg);
    std::ostringstream key;
    key << "merged-" << (num_egs_written_++) << "-" << minibatch_size;
    writer_->Write(key.str(), merged_eg);
  }
  
  void DiscriminativeExampleMerger::Finish() {
    if (finished_) return;  // already finished.
    finished_ = true;
  
    // we'll convert the map eg_to_egs_ to a vector of vectors to avoid
    // iterator invalidation problems.
    std::vector<std::vector<NnetDiscriminativeExample*> > all_egs;
    all_egs.reserve(eg_to_egs_.size());
  
    MapType::iterator iter = eg_to_egs_.begin(), end = eg_to_egs_.end();
    for (; iter != end; ++iter)
      all_egs.push_back(iter->second);
    eg_to_egs_.clear();
  
    for (size_t i = 0; i < all_egs.size(); i++) {
      int32 minibatch_size;
      std::vector<NnetDiscriminativeExample*> &vec = all_egs[i];
      KALDI_ASSERT(!vec.empty());
      int32 eg_size = GetNnetDiscriminativeExampleSize(*(vec[0]));
      bool input_ended = true;
      while (!vec.empty() &&
             (minibatch_size = config_.MinibatchSize(eg_size, vec.size(),
                                                     input_ended)) != 0) {
        // MergeDiscriminativeExamples() expects a vector of
        // NnetDiscriminativeExample, not of pointers, so use swap to create that
        // without doing any real work.
        std::vector<NnetDiscriminativeExample> egs_to_merge(minibatch_size);
        for (int32 i = 0; i < minibatch_size; i++) {
          egs_to_merge[i].Swap(vec[i]);
          delete vec[i];  // we owned those pointers.
        }
        vec.erase(vec.begin(), vec.begin() + minibatch_size);
        WriteMinibatch(&egs_to_merge);
      }
      if (!vec.empty()) {
        int32 eg_size = GetNnetDiscriminativeExampleSize(*(vec[0]));
        NnetDiscriminativeExampleStructureHasher eg_hasher;
        size_t structure_hash = eg_hasher(*(vec[0]));
        int32 num_discarded = vec.size();
        stats_.DiscardedExamples(eg_size, structure_hash, num_discarded);
        for (int32 i = 0; i < num_discarded; i++)
          delete vec[i];
        vec.clear();
      }
    }
    stats_.PrintStats();
  }
  
  
  
  } // namespace nnet3
  } // namespace kaldi