// nnet2bin/nnet-relabel-egs.cc
  
  // Copyright 2014   Vimal Manohar
  
  // 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.
  
  /** @brief Relabels neural network egs with the read pdf-id alignments
  */
  
  #include <sstream>
  
  #include "base/kaldi-common.h"
  #include "util/common-utils.h"
  #include "nnet2/nnet-example.h"
  
  namespace kaldi {
    
    // this functions splits an egs key like <utt_id>-<frame_id> into 
    // separate utterance id and frame id on the last delimiter.
    // Returns false if the delimiter is not found in the key.
    bool SplitEgsKey(const std::string &key, 
                      std::string *utt_id, int32 *frame_id) {
      size_t start = 0, found = 0, end = key.size();
      utt_id->clear();
   
      found = key.find_last_of("-", end);
      // start != end condition is for when the delimiter is at the end
      
      if (found != start && start != end && found < end) {
        *utt_id = key.substr(start, found - start);
        std::istringstream tmp(key.substr(found + 1, end));
        tmp >> *frame_id;
        return true;
      }
  
      return false;
    }
  }
  
  int main(int argc, char *argv[]) {
    using namespace kaldi;
    using namespace kaldi::nnet2;
  
    typedef kaldi::int32 int32;
    typedef kaldi::int64 int64;
    try {
      const char *usage =
          "Relabel neural network egs with the read pdf-id alignments, "
          "zero-based..
  "
          "Usage: nnet-relabel-egs [options] <pdf-aligment-rspecifier> "
          "<egs_rspecifier1> ... <egs_rspecifierN> "
          "<egs_wspecifier1> ... <egs_wspecifierN>
  "
          "e.g.: 
  "
          " nnet-relabel-egs ark:1.ali egs_in/egs.1.ark egs_in/egs.2.ark "
          "egs_out/egs.1.ark egs_out/egs.2.ark
  "
          "See also: nnet-get-egs, nnet-copy-egs, steps/nnet2/relabel_egs.sh
  ";
  
      ParseOptions po(usage);
  
      po.Read(argc, argv);
  
      // Here we expect equal number of input egs archive and output egs archives. 
      // So the total number of arguments including the alignment specifier must be odd.
      if (po.NumArgs() < 3 || po.NumArgs() % 2 == 0) {
        po.PrintUsage();
        exit(1);
      }
  
      std::string alignments_rspecifier = po.GetArg(1);
      int32 num_archives = (po.NumArgs() - 1) / 2;
      
      SequentialInt32VectorReader ali_reader(alignments_rspecifier);
  
      unordered_map<std::string, std::vector<int32>* > utt_to_pdf_ali;
  
      // Keep statistics
      int32 num_ali = 0;
      int64 num_frames_ali = 0, num_frames_egs = 0, 
            num_frames_missing = 0, num_frames_relabelled = 0;
  
      // Read alignments and put the pointer in an unordered map
      // indexed by the key. This is so that we can efficiently find the 
      // alignment corresponding to the utterance to 
      // which a particular frame belongs
      for (; !ali_reader.Done(); ali_reader.Next(), num_ali++) {
        std::string key = ali_reader.Key();
        std::vector<int32> *alignment = new std::vector<int32>(ali_reader.Value());
        std::pair<std::string, std::vector<int32>* > map(key, alignment);
        utt_to_pdf_ali.insert(map);
        num_frames_ali += alignment->size();
      }
  
      // Read archives of egs sequentially
      for (int32 i = 0; i < num_archives; i++) {
        std::string egs_rspecifier(po.GetArg(i+2));
        std::string egs_wspecifier(po.GetArg(i+2+num_archives));
  
        SequentialNnetExampleReader egs_reader(egs_rspecifier);
        NnetExampleWriter egs_writer(egs_wspecifier);
  
        for (; !egs_reader.Done(); egs_reader.Next(), num_frames_egs++) {
        
          std::string key(egs_reader.Key());
  
          std::string utt_id;
          int32 frame_id;
  
          if (!SplitEgsKey(key, &utt_id, &frame_id)) {
            KALDI_ERR << "Unable to split key " << key << " on delimiter - " 
                      << " into utterance id and frame id";
          }
          NnetExample eg(egs_reader.Value());
  
          if (utt_to_pdf_ali.find(utt_id) == utt_to_pdf_ali.end()) {
            KALDI_WARN << "Unable to find utterance id " << utt_id;
            egs_writer.Write(key, eg);
            num_frames_missing++;
            continue;
          }
          const std::vector<int32> *alignment = utt_to_pdf_ali[utt_id];
  
          int32 num_frames_in_eg = eg.labels.size();
          for (int32 t_offset = 0; t_offset < num_frames_in_eg; t_offset++) {
            int32 t = frame_id + t_offset;
            if (t >= static_cast<int32>(alignment->size())) {
              KALDI_ERR << "Time index " << t << " out of range for alignment, "
                        << "should be < " << alignment->size();
            }
            if (eg.GetLabelSingle(t_offset) != (*alignment)[t])
              num_frames_relabelled++; 
            eg.SetLabelSingle(t_offset, (*alignment)[t]);
          }
          egs_writer.Write(key, eg);
        }
      }
  
      unordered_map<std::string, std::vector<int32>*>::iterator iter;
      
      for (iter = utt_to_pdf_ali.begin(); iter != utt_to_pdf_ali.end(); ++iter)
        delete iter->second;
      
      KALDI_LOG << "Read " << num_ali << " alignments containing a total of " 
                << num_frames_ali << " frames; labelled " 
                << num_frames_egs - num_frames_missing << " frames out of " 
                << num_frames_egs << " examples; labels changed for " 
                << num_frames_relabelled << " of those frames.
  .";
  
      return (num_frames_missing > 0.5  * num_frames_egs);
  
    } catch(const std::exception &e) {
      std::cerr << e.what();
      return -1;
    }
  }