// gmmbin/gmm-init-biphone.cc
  
  // Copyright 2017   Hossein Hadian
  
  // 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 "base/kaldi-common.h"
  #include "util/common-utils.h"
  #include "gmm/am-diag-gmm.h"
  #include "tree/event-map.h"
  #include "tree/context-dep.h"
  #include "hmm/hmm-topology.h"
  #include "hmm/transition-model.h"
  
  namespace kaldi {
  // This function reads a file like:
  // 1 2 3
  // 4 5
  // 6 7 8
  // where each line is a list of integer id's of phones (that should have their pdfs shared).
  void ReadSharedPhonesList(std::string rxfilename, std::vector<std::vector<int32> > *list_out) {
    list_out->clear();
    Input input(rxfilename);
    std::istream &is = input.Stream();
    std::string line;
    while (std::getline(is, line)) {
      list_out->push_back(std::vector<int32>());
      if (!SplitStringToIntegers(line, " \t\r", true, &(list_out->back())))
        KALDI_ERR << "Bad line in shared phones list: " << line << " (reading "
                  << PrintableRxfilename(rxfilename) << ")";
      std::sort(list_out->rbegin()->begin(), list_out->rbegin()->end());
      if (!IsSortedAndUniq(*(list_out->rbegin())))
        KALDI_ERR << "Bad line in shared phones list (repeated phone): " << line
                  << " (reading " << PrintableRxfilename(rxfilename) << ")";
    }
  }
  
  EventMap
  *GetFullBiphoneStubMap(const std::vector<std::vector<int32> > &phone_sets,
                         const std::vector<int32> &phone2num_pdf_classes,
                         const std::vector<int32> &ci_phones_list,
                         const std::vector<std::vector<int32> > &bi_counts,
                         int32 biphone_min_count,
                         const std::vector<int32> &mono_counts,
                         int32 mono_min_count) {
  
    {  // Check the inputs
      KALDI_ASSERT(!phone_sets.empty());
      std::set<int32> all_phones;
      for (size_t i = 0; i < phone_sets.size(); i++) {
        KALDI_ASSERT(IsSortedAndUniq(phone_sets[i]));
        KALDI_ASSERT(!phone_sets[i].empty());
        for (size_t j = 0; j < phone_sets[i].size(); j++) {
          KALDI_ASSERT(all_phones.count(phone_sets[i][j]) == 0);  // Check not present.
          all_phones.insert(phone_sets[i][j]);
        }
      }
    }
  
  
    int32 numpdfs_per_phone = phone2num_pdf_classes[1];
    int32 current_pdfid = 0;
    std::map<EventValueType, EventMap*> level1_map;  // key is 1
  
    for (size_t i = 0; i < ci_phones_list.size(); i++) {
      std::map<EventValueType, EventAnswerType> level2_map;
      level2_map[0] = current_pdfid++;
      if (numpdfs_per_phone == 2) level2_map[1] = current_pdfid++;
      level1_map[ci_phones_list[i]] = new TableEventMap(kPdfClass, level2_map);
    }
  
    // If there is not enough data for a biphone, we will revert to monophone
    // and if there is not enough data for the monophone either, we will revert
    // to zerophone (which is like a global garbage pdf) after initializing it.
    int32 zerophone_pdf = -1;
    // If a monophone state is created for a phone-set, the corresponding pdf will
    // be stored in this vector.
    std::vector<int32> monophone_pdf(phone_sets.size(), -1);
  
    for (size_t i = 0; i < phone_sets.size(); i++) {
  
      if (numpdfs_per_phone == 1) {
        // Create an event map for level2:
        std::map<EventValueType, EventAnswerType> level2_map;  // key is 0
        level2_map[0] = current_pdfid++;  // no-left-context case
        for (size_t j = 0; j < phone_sets.size(); j++) {
          int32 pdfid = current_pdfid++;
          std::vector<int32> pset = phone_sets[j];  // All these will have a
                                                    // shared pdf with id=pdfid
          for (size_t k = 0; k < pset.size(); k++)
            level2_map[pset[k]] = pdfid;
        }
        std::vector<int32> pset = phone_sets[i];  // All these will have a
                                                  // shared event-map child
        for (size_t k = 0; k < pset.size(); k++)
          level1_map[pset[k]] = new TableEventMap(0, level2_map);
      } else {
        KALDI_ASSERT(numpdfs_per_phone == 2);
        std::vector<int32> right_phoneset = phone_sets[i];  // All these will have a shared
                                                  // event-map child
        // Create an event map for level2:
        std::map<EventValueType, EventMap*> level2_map;  // key is 0
        {  // Handle CI phones
          std::map<EventValueType, EventAnswerType> level3_map;  // key is kPdfClass
          level3_map[0] = current_pdfid++;
          level3_map[1] = current_pdfid++;
          level2_map[0] = new TableEventMap(kPdfClass, level3_map);  // no-left-context case
          for (size_t i = 0; i < ci_phones_list.size(); i++)  // ci-phone left-context cases
            level2_map[ci_phones_list[i]] = new TableEventMap(kPdfClass, level3_map);
        }
        for (size_t j = 0; j < phone_sets.size(); j++) {
          std::vector<int32> left_phoneset = phone_sets[j];  // All these will have a
          // shared subtree with 2 pdfids
          std::map<EventValueType, EventAnswerType> level3_map;  // key is kPdfClass
          if (bi_counts.empty() ||
              bi_counts[left_phoneset[0]][right_phoneset[0]] >= biphone_min_count) {
            level3_map[0] = current_pdfid++;
            level3_map[1] = current_pdfid++;
          } else if (mono_counts.empty() ||
                     mono_counts[right_phoneset[0]] > mono_min_count) {
            //  Revert to mono.
            KALDI_VLOG(2) << "Reverting to mono for biphone (" << left_phoneset[0]
                          << "," << right_phoneset[0] << ")";
            if (monophone_pdf[i] == -1) {
              KALDI_VLOG(1) << "Reserving mono PDFs for phone-set " << i;
              monophone_pdf[i] = current_pdfid++;
              current_pdfid++; // num-pdfs-per-phone is 2
            }
            level3_map[0] = monophone_pdf[i];
            level3_map[1] = monophone_pdf[i] + 1;
          } else {
            KALDI_VLOG(2) << "Reverting to zerophone for biphone ("
                          << left_phoneset[0]
                          << "," << right_phoneset[0] << ")";
            // Revert to zerophone
            if (zerophone_pdf == -1) {
              KALDI_VLOG(1) << "Reserving zero PDFs.";
              zerophone_pdf = current_pdfid++;
              current_pdfid++; // num-pdfs-per-phone is 2
            }
            level3_map[0] = zerophone_pdf;
            level3_map[1] = zerophone_pdf + 1;
          }
  
          for (size_t k = 0; k < left_phoneset.size(); k++) {
            int32 left_phone = left_phoneset[k];
            level2_map[left_phone] = new TableEventMap(kPdfClass, level3_map);
          }
        }
        for (size_t k = 0; k < right_phoneset.size(); k++) {
          std::map<EventValueType, EventMap*> level2_copy;
          for (auto const& kv: level2_map)
            level2_copy[kv.first] = kv.second->Copy(std::vector<EventMap*>());
          int32 right_phone = right_phoneset[k];
          level1_map[right_phone] = new TableEventMap(0, level2_copy);
        }
      }
  
    }
    KALDI_LOG << "Num PDFs: " << current_pdfid;
    return new TableEventMap(1, level1_map);
  }
  
  
  ContextDependency*
  BiphoneContextDependencyFull(std::vector<std::vector<int32> > phone_sets,
                               const std::vector<int32> phone2num_pdf_classes,
                               const std::vector<int32> &ci_phones_list,
                               const std::vector<std::vector<int32> > &bi_counts,
                               int32 biphone_min_count,
                               const std::vector<int32> &mono_counts,
                               int32 mono_min_count) {
    // Remove all the CI phones from the phone sets
    std::set<int32> ci_phones;
    for (size_t i = 0; i < ci_phones_list.size(); i++)
      ci_phones.insert(ci_phones_list[i]);
    for (int32 i = phone_sets.size() - 1; i >= 0; i--) {
      for (int32 j = phone_sets[i].size() - 1; j >= 0; j--) {
        if (ci_phones.find(phone_sets[i][j]) != ci_phones.end()) {  // Delete it
          phone_sets[i].erase(phone_sets[i].begin() + j);
          if (phone_sets[i].empty())   // If empty, delete the whole entry
            phone_sets.erase(phone_sets.begin() + i);
        }
      }
    }
  
    std::vector<bool> share_roots(phone_sets.size(), false);  // Don't share roots
    // N is context size, P = position of central phone (must be 0).
    int32 P = 1, N = 2;
    EventMap *pdf_map = GetFullBiphoneStubMap(phone_sets,
                                              phone2num_pdf_classes,
                                              ci_phones_list, bi_counts,
                                              biphone_min_count, mono_counts,
                                              mono_min_count);
    return new ContextDependency(N, P, pdf_map);
  }
  
  
  } // end namespace kaldi
  
  /* This function reads the counts of biphones and monophones from a text file
     generated for chain flat-start training. On each line there is either a
     biphone count or a monophone count:
     <left-phone-id> <right-phone-id> <count>
     <monophone-id> <count>
     The phone-id's are according to phones.txt.
  
     It's more efficient to load the biphone counts into a map because
     most entries are zero, but since there are not many biphones, a 2-dim vector
     is OK. */
  static void ReadPhoneCounts(std::string &filename, int32 num_phones,
                              std::vector<int32> *mono_counts,
                              std::vector<std::vector<int32> > *bi_counts) {
    // The actual phones start from id = 1 (so the last phone has id = num_phones).
    mono_counts->resize(num_phones + 1, 0);
    bi_counts->resize(num_phones + 1, std::vector<int>(num_phones + 1, 0));
    std::ifstream infile(filename);
    std::string line;
    while (std::getline(infile, line)) {
      std::istringstream iss(line);
      int a, b;
      long c;
      if ((std::istringstream(line) >> a >> b >> c)) {
        // It's a biphone count.
        KALDI_ASSERT(a >= 0 && a <= num_phones);  // 0 means no-left-context
        KALDI_ASSERT(b > 0 && b <= num_phones);
        KALDI_ASSERT(c >= 0);
        (*bi_counts)[a][b] = c;
      } else if ((std::istringstream(line) >> b >> c)) {
        // It's a monophone count.
        KALDI_ASSERT(b > 0 && b <= num_phones);
        KALDI_ASSERT(c >= 0);
        (*mono_counts)[b] = c;
      } else {
        KALDI_ERR << "Bad line in phone stats file: " << line;
      }
    }
  }
  
  int main(int argc, char *argv[]) {
    try {
      using namespace kaldi;
      using kaldi::int32;
  
      const char *usage =
          "Initialize a biphone context-dependency tree with all the
  "
          "leaves (i.e. a full tree). Intended for end-to-end tree-free models.
  "
          "Usage:  gmm-init-biphone <topology-in> <dim> <model-out> <tree-out> 
  "
          "e.g.: 
  "
          " gmm-init-biphone topo 39 bi.mdl bi.tree
  ";
  
      bool binary = true;
      std::string shared_phones_rxfilename, phone_counts_rxfilename;
      int32 min_biphone_count = 100, min_mono_count = 20;
      std::string ci_phones_str;
      std::vector<int32> ci_phones;  // Sorted, uniqe vector of
      // context-independent phones.
  
      ParseOptions po(usage);
      po.Register("binary", &binary, "Write output in binary mode");
      po.Register("shared-phones", &shared_phones_rxfilename,
                  "rxfilename containing, on each line, a list of phones "
                  "whose pdfs should be shared.");
      po.Register("ci-phones", &ci_phones_str, "Colon-separated list of "
                  "integer indices of context-independent phones.");
      po.Register("phone-counts", &phone_counts_rxfilename,
                  "rxfilename containing, on each line, a biphone/phone and "
                  "its count in the training data.");
      po.Register("min-biphone-count", &min_biphone_count, "Minimum number of "
                  "occurences of a biphone in training data to reserve pdfs "
                  "for it.");
      po.Register("min-monophone-count", &min_mono_count, "Minimum number of "
                  "occurences of a monophone in training data to reserve pdfs "
                  "for it.");
      po.Read(argc, argv);
  
      if (po.NumArgs() != 4) {
        po.PrintUsage();
        exit(1);
      }
  
  
      std::string topo_filename = po.GetArg(1);
      int dim = 0;
      if (!ConvertStringToInteger(po.GetArg(2), &dim) || dim <= 0 || dim > 10000)
        KALDI_ERR << "Bad dimension:" << po.GetArg(2)
                  << ". It should be a positive integer.";
      std::string model_filename = po.GetArg(3);
      std::string tree_filename = po.GetArg(4);
  
      if (!ci_phones_str.empty()) {
        SplitStringToIntegers(ci_phones_str, ":", false, &ci_phones);
        std::sort(ci_phones.begin(), ci_phones.end());
        if (!IsSortedAndUniq(ci_phones) || ci_phones.empty() || ci_phones[0] == 0)
          KALDI_ERR << "Invalid --ci-phones option: " << ci_phones_str;
      }
  
      Vector<BaseFloat> glob_inv_var(dim);
      glob_inv_var.Set(1.0);
      Vector<BaseFloat> glob_mean(dim);
      glob_mean.Set(1.0);
  
      HmmTopology topo;
      bool binary_in;
      Input ki(topo_filename, &binary_in);
      topo.Read(ki.Stream(), binary_in);
  
      const std::vector<int32> &phones = topo.GetPhones();
  
      std::vector<int32> phone2num_pdf_classes(1 + phones.back());
      for (size_t i = 0; i < phones.size(); i++) {
        phone2num_pdf_classes[phones[i]] = topo.NumPdfClasses(phones[i]);
        // For now we only support 1 or 2 pdf's per phone
        KALDI_ASSERT(phone2num_pdf_classes[phones[i]] == 1 ||
                     phone2num_pdf_classes[phones[i]] == 2);
      }
  
      std::vector<int32> mono_counts;
      std::vector<std::vector<int32> > bi_counts;
      if (!phone_counts_rxfilename.empty()) {
        ReadPhoneCounts(phone_counts_rxfilename, phones.size(),
                        &mono_counts, &bi_counts);
        KALDI_LOG << "Loaded mono/bi phone counts.";
      }
  
  
      // Now the tree:
      ContextDependency *ctx_dep = NULL;
      std::vector<std::vector<int32> > shared_phones;
      if (shared_phones_rxfilename == "") {
        shared_phones.resize(phones.size());
        for (size_t i = 0; i < phones.size(); i++)
          shared_phones[i].push_back(phones[i]);
      } else {
        ReadSharedPhonesList(shared_phones_rxfilename, &shared_phones);
        // ReadSharedPhonesList crashes on error.
      }
      ctx_dep = BiphoneContextDependencyFull(shared_phones, phone2num_pdf_classes,
                                             ci_phones, bi_counts,
                                             min_biphone_count,
                                             mono_counts, min_mono_count);
  
      int32 num_pdfs = ctx_dep->NumPdfs();
  
      AmDiagGmm am_gmm;
      DiagGmm gmm;
      gmm.Resize(1, dim);
      {  // Initialize the gmm.
        Matrix<BaseFloat> inv_var(1, dim);
        inv_var.Row(0).CopyFromVec(glob_inv_var);
        Matrix<BaseFloat> mu(1, dim);
        mu.Row(0).CopyFromVec(glob_mean);
        Vector<BaseFloat> weights(1);
        weights.Set(1.0);
        gmm.SetInvVarsAndMeans(inv_var, mu);
        gmm.SetWeights(weights);
        gmm.ComputeGconsts();
      }
  
      for (int i = 0; i < num_pdfs; i++)
        am_gmm.AddPdf(gmm);
  
      // Now the transition model:
      TransitionModel trans_model(*ctx_dep, topo);
  
      {
        Output ko(model_filename, binary);
        trans_model.Write(ko.Stream(), binary);
        am_gmm.Write(ko.Stream(), binary);
      }
  
      // Now write the tree.
      ctx_dep->Write(Output(tree_filename, binary).Stream(),
                     binary);
  
      delete ctx_dep;
      return 0;
    } catch(const std::exception &e) {
      std::cerr << e.what();
      return -1;
    }
  }