// tree/build-tree-utils.cc
  
  // Copyright 2009-2011  Microsoft Corporation;  Haihua Xu
  
  // 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 <set>
  #include <queue>
  #include "util/stl-utils.h"
  #include "tree/build-tree-utils.h"
  
  
  
  namespace kaldi {
  
  void WriteBuildTreeStats(std::ostream &os, bool binary, const BuildTreeStatsType &stats) {
    WriteToken(os, binary, "BTS");
    uint32 size = stats.size();
    WriteBasicType(os, binary, size);
    for (size_t i = 0; i < size; i++) {
      WriteEventType(os, binary, stats[i].first);
      bool nonNull = (stats[i].second != NULL);
      WriteBasicType(os, binary, nonNull);
      if (nonNull) stats[i].second->Write(os, binary);
    }
    if (os.fail()) {
      KALDI_ERR << "WriteBuildTreeStats: write failed.";
    }
    if (!binary) os << '
  ';
  }
  
  
  void ReadBuildTreeStats(std::istream &is, bool binary, const Clusterable &example, BuildTreeStatsType *stats) {
    KALDI_ASSERT(stats != NULL);
    KALDI_ASSERT(stats->empty());
    ExpectToken(is, binary, "BTS");
    uint32 size;
    ReadBasicType(is, binary, &size);
    stats->resize(size);
    for (size_t i = 0; i < size; i++) {
      ReadEventType(is, binary, &((*stats)[i].first));
      bool nonNull;
      ReadBasicType(is, binary, &nonNull);
      if (nonNull) (*stats)[i].second = example.ReadNew(is, binary);
      else (*stats)[i].second = NULL;
    }
  }
  
  
  bool PossibleValues(EventKeyType key,
                      const BuildTreeStatsType &stats,
                      std::vector<EventValueType> *ans) {
    bool all_present = true;
    std::set<EventValueType> values;
    BuildTreeStatsType::const_iterator iter = stats.begin(), end = stats.end();
    for (; iter != end; ++iter) {
      EventValueType val;
      if (EventMap::Lookup(iter->first, key, &val))
        values.insert(val);
      else
        all_present = false;
    }
    if (ans)
      CopySetToVector(values, ans);
    return all_present;
  }
  
  static void GetEventKeys(const EventType &vec, std::vector<EventKeyType> *keys) {
    keys->resize(vec.size());
    EventType::const_iterator iter = vec.begin(), end = vec.end();
    std::vector<EventKeyType>::iterator out_iter = keys->begin();
    for (; iter!= end; ++iter, ++out_iter)
      *out_iter = iter->first;
  }
  
  
  // recall:
  // typedef std::vector<std::pair<EventType, Clusterable*> > BuildTreeStatsType;
  void FindAllKeys(const BuildTreeStatsType &stats, AllKeysType keys_type, std::vector<EventKeyType> *keys_out) {
    KALDI_ASSERT(keys_out != NULL);
    BuildTreeStatsType::const_iterator iter = stats.begin(), end = stats.end();
    if (iter == end) return;  // empty set of keys.
    std::vector<EventKeyType> keys;
    GetEventKeys(iter->first, &keys);
    ++iter;
    for (; iter!= end; ++iter) {
      std::vector<EventKeyType> keys2;
      GetEventKeys(iter->first, &keys2);
      if (keys_type == kAllKeysInsistIdentical) {
        if (keys2 != keys)
          KALDI_ERR << "AllKeys: keys in events are not all the same [called with kAllKeysInsistIdentical and all are not identical.";
      } else if (keys_type == kAllKeysIntersection) {
        std::vector<EventKeyType> new_keys(std::max(keys.size(), keys2.size()));
        // following line relies on sorted order of event keys.
        std::vector<EventKeyType>::iterator end_iter =
            std::set_intersection(keys.begin(), keys.end(), keys2.begin(), keys2.end(), new_keys.begin());
        new_keys.erase(end_iter, new_keys.end());
        keys = new_keys;
      } else {  // union.
        KALDI_ASSERT(keys_type == kAllKeysUnion);
        std::vector<EventKeyType> new_keys(keys.size()+keys2.size());
        // following line relies on sorted order of event keys.
        std::vector<EventKeyType>::iterator end_iter =
            std::set_union(keys.begin(), keys.end(), keys2.begin(), keys2.end(), new_keys.begin());
        new_keys.erase(end_iter, new_keys.end());
        keys = new_keys;
      }
    }
    *keys_out = keys;
  }
  
  
  EventMap *DoTableSplit(const EventMap &orig, EventKeyType key,  const BuildTreeStatsType &stats,
                         int32 *num_leaves) {
    // First-- map the stats to each leaf in the EventMap.
    std::vector<BuildTreeStatsType> split_stats;
    SplitStatsByMap(stats, orig, &split_stats);
    // Now for each leaf that has stats in it, do the table split according to the given name.
    std::vector<EventMap*> splits(split_stats.size(), NULL);
    for (EventAnswerType leaf = 0; leaf < (EventAnswerType)split_stats.size(); leaf++) {
      if (!split_stats[leaf].empty()) {
        // first work out the possible values the name takes.
        std::vector<EventValueType> vals;  // vals are put here, sorted.
        bool all_present = PossibleValues(key, split_stats[leaf], &vals);
        KALDI_ASSERT(all_present);  // currently do not support mapping undefined values.
        KALDI_ASSERT(!vals.empty() && vals.front() >= 0);  // don't support mapping negative values
        // at present time-- would need different EventMap type, not TableEventMap.
        std::vector<EventMap*> table(vals.back()+1, (EventMap*)NULL);
        for (size_t idx = 0;idx < vals.size();idx++) {
          EventValueType val = vals[idx];
          if (idx == 0) table[val] = new ConstantEventMap(leaf);  // reuse current leaf.
          else table[val] = new ConstantEventMap( (*num_leaves)++ );  // else take new leaf id.
        }
        // takes ownershipof stats.
        splits[leaf] = new TableEventMap(key, table);
      }
    }
    EventMap *ans = orig.Copy(splits);
    DeletePointers(&splits);
    return ans;
  }
  
  EventMap *DoTableSplitMultiple(const EventMap &orig, const std::vector<EventKeyType> &keys,  const BuildTreeStatsType &stats, int32 *num_leaves) {
  
    if (keys.empty()) return orig.Copy();
    else {
      EventMap *cur = NULL;  // would make it &orig, except for const issues.
      for (size_t i = 0; i < keys.size(); i++) {
        EventMap *next = DoTableSplit( (cur ? *cur : orig), keys[i], stats, num_leaves);
        delete cur;  // delete intermediate maps.
        cur = next;
      }
      return cur;
    }
  }
  
  
  
  void SplitStatsByMap(const BuildTreeStatsType &stats, const EventMap &e, std::vector<BuildTreeStatsType> *stats_out) {
    BuildTreeStatsType::const_iterator iter, end = stats.end();
    KALDI_ASSERT(stats_out != NULL);
    stats_out->clear();
    size_t size = 0;
    for (iter = stats.begin(); iter != end; ++iter) {
      const EventType &evec = iter->first;
      EventAnswerType ans;
      if (!e.Map(evec, &ans)) // this is an error--could not map it.
        KALDI_ERR << "SplitStatsByMap: could not map event vector " << EventTypeToString(evec)
                  << "if error seen during tree-building, check that "
                  << "--context-width and --central-position match stats, "
                  << "and that phones that are context-independent (CI) during "
                  << "stats accumulation do not share roots with non-CI phones.";
      size = std::max(size, (size_t)(ans+1));
    }
    stats_out->resize(size);
    for (iter = stats.begin(); iter != end; ++iter) {
      const EventType &evec = iter->first;
      EventAnswerType ans;
      bool b = e.Map(evec, &ans);
      KALDI_ASSERT(b);
      (*stats_out)[ans].push_back(*iter);
    }
  }
  
  void SplitStatsByKey(const BuildTreeStatsType &stats_in, EventKeyType key, std::vector<BuildTreeStatsType> *stats_out) {
    BuildTreeStatsType::const_iterator iter, end = stats_in.end();
    KALDI_ASSERT(stats_out != NULL);
    stats_out->clear();
    size_t size = 0;
    // This loop works out size of output vector.
    for (iter = stats_in.begin(); iter != end; ++iter) {
      const EventType &evec = iter->first;
      EventValueType val;
      if (! EventMap::Lookup(evec, key, &val)) // no such key.
        KALDI_ERR << "SplitStats: key "<< key << " is not present in event vector " << EventTypeToString(evec);
      size = std::max(size, (size_t)(val+1));
    }
    stats_out->resize(size);
    // This loop splits up the stats.
    for (iter = stats_in.begin(); iter != end; ++iter) {
      const EventType &evec = iter->first;
      EventValueType val;
      EventMap::Lookup(evec, key, &val);  // will not fail.
      (*stats_out)[val].push_back(*iter);
    }
  }
  
  
  void FilterStatsByKey(const BuildTreeStatsType &stats_in,
                        EventKeyType key,
                        std::vector<EventValueType> &values,
                        bool include_if_present,  // true-> retain only in "values",
                        // false-> retain only not in "values".
                        BuildTreeStatsType *stats_out) {
    KALDI_ASSERT(IsSortedAndUniq(values));
    BuildTreeStatsType::const_iterator iter, end = stats_in.end();
    KALDI_ASSERT(stats_out != NULL);
    stats_out->clear();
  
    for (iter = stats_in.begin(); iter != end; ++iter) {
      const EventType &evec = iter->first;
      EventValueType val;
      if (! EventMap::Lookup(evec, key, &val)) // no such key. // HERE.
        KALDI_ERR << "SplitStats: key "<< key << " is not present in event vector " << EventTypeToString(evec);
      bool in_values = std::binary_search(values.begin(), values.end(), val);
      if (in_values == include_if_present)
        stats_out->push_back(*iter);
    }
  }
  
  
  Clusterable *SumStats(const BuildTreeStatsType &stats_in) {
    Clusterable *ans = NULL;
    BuildTreeStatsType::const_iterator iter = stats_in.begin(), end = stats_in.end();
    for (; iter != end; ++iter) {
      Clusterable *cl = iter->second;
      if (cl != NULL) {
        if (!ans)  ans = cl->Copy();
        else ans->Add(*cl);
      }
    }
    return ans;
  }
  
  BaseFloat SumNormalizer(const BuildTreeStatsType &stats_in) {
    BaseFloat ans = 0.0;
    BuildTreeStatsType::const_iterator iter = stats_in.begin(), end = stats_in.end();
    for (; iter != end; ++iter) {
      Clusterable *cl = iter->second;
      if (cl != NULL) ans += cl->Normalizer();
    }
    return ans;
  }
  
  BaseFloat SumObjf(const BuildTreeStatsType &stats_in) {
    BaseFloat ans = 0.0;
    BuildTreeStatsType::const_iterator iter = stats_in.begin(), end = stats_in.end();
    for (; iter != end; ++iter) {
      Clusterable *cl = iter->second;
      if (cl != NULL) ans += cl->Objf();
    }
    return ans;
  }
  
  
  void SumStatsVec(const std::vector<BuildTreeStatsType> &stats_in, std::vector<Clusterable*> *stats_out) {
    KALDI_ASSERT(stats_out != NULL && stats_out->empty());
    stats_out->resize(stats_in.size(), NULL);
    for (size_t i = 0;i < stats_in.size();i++) (*stats_out)[i] = SumStats(stats_in[i]);
  }
  
  BaseFloat ObjfGivenMap(const BuildTreeStatsType &stats_in, const EventMap &e) {
    std::vector<BuildTreeStatsType> split_stats;
    SplitStatsByMap(stats_in, e, &split_stats);
    std::vector<Clusterable*> summed_stats;
    SumStatsVec(split_stats, &summed_stats);
    BaseFloat ans = SumClusterableObjf(summed_stats);
    DeletePointers(&summed_stats);
    return ans;
  }
  
  // This function computes the best initial split of these stats [with this key].
  // Returns best objf change (>=0).
  BaseFloat ComputeInitialSplit(const std::vector<Clusterable*> &summed_stats,
                                const Questions &q_opts, EventKeyType key,
                                std::vector<EventValueType> *yes_set) {
    KALDI_ASSERT(yes_set != NULL);
    yes_set->clear();
    const QuestionsForKey &key_opts = q_opts.GetQuestionsOf(key);
  
    // "total" needed for optimization in AddToClustersOptimized,
    // and also used to work otu total objf.
    Clusterable *total = SumClusterable(summed_stats);
    if (total == NULL) return 0.0;  // because there were no stats or non-NULL stats.
    BaseFloat unsplit_objf = total->Objf();
  
    const std::vector<std::vector<EventValueType> > &questions_of_this_key = key_opts.initial_questions;
  
    int32 best_idx = -1;
    BaseFloat best_objf_change = 0;
  
    for (size_t i = 0; i < questions_of_this_key.size(); i++) {
      const std::vector<EventValueType> &yes_set = questions_of_this_key[i];
      std::vector<int32> assignments(summed_stats.size(), 0);  // 0 is index of "no".
      std::vector<Clusterable*> clusters(2);  // no and yes clusters.
      for (std::vector<EventValueType>::const_iterator iter = yes_set.begin(); iter != yes_set.end(); ++iter) {
        KALDI_ASSERT(*iter>=0);
        if (*iter < (EventValueType)assignments.size()) assignments[*iter] = 1;
      }
      kaldi::AddToClustersOptimized(summed_stats, assignments, *total, &clusters);
      BaseFloat this_objf = SumClusterableObjf(clusters);
  
      if (this_objf < unsplit_objf- 0.001*std::abs(unsplit_objf)) {  // got worse; should never happen.
        // of course small differences can be caused by roundoff.
        KALDI_WARN << "Objective function got worse when building tree: "<< this_objf << " < " << unsplit_objf;
        KALDI_ASSERT(!(this_objf < unsplit_objf - 0.01*(200 + std::abs(unsplit_objf))));  // do assert on more stringent check.
      }
  
      BaseFloat this_objf_change = this_objf - unsplit_objf;
      if (this_objf_change > best_objf_change) {
        best_objf_change = this_objf_change;
        best_idx = i;
      }
      DeletePointers(&clusters);
    }
    delete total;
    if (best_idx != -1)
      *yes_set = questions_of_this_key[best_idx];
    return best_objf_change;
  }
  
  
  // returns best delta-objf.
  // If key does not exist, returns 0 and sets yes_set_out to empty.
  BaseFloat FindBestSplitForKey(const BuildTreeStatsType &stats,
                                const Questions &q_opts,
                                EventKeyType key,
                                std::vector<EventValueType> *yes_set_out) {
    if (stats.size()<=1) return 0.0;  // cannot split if only zero or one instance of stats.
    if (!PossibleValues(key, stats, NULL)) {
      yes_set_out->clear();
      return 0.0;  // Can't split as key not always defined.
    }
    std::vector<Clusterable*> summed_stats;  // indexed by value corresponding to key. owned here.
    {  // compute summed_stats
      std::vector<BuildTreeStatsType> split_stats;
      SplitStatsByKey(stats, key, &split_stats);
      SumStatsVec(split_stats, &summed_stats);
    }
  
    std::vector<EventValueType> yes_set;
    BaseFloat improvement = ComputeInitialSplit(summed_stats,
                                                 q_opts, key, &yes_set);
    // find best basic question.
  
    std::vector<int32> assignments(summed_stats.size(), 0);  // assigns to "no" (0) by default.
    for (std::vector<EventValueType>::const_iterator iter = yes_set.begin(); iter != yes_set.end(); ++iter) {
      KALDI_ASSERT(*iter>=0);
      if (*iter < (EventValueType)assignments.size()) {
        // this guard necessary in case stats did not have all the
        // values present in "yes_set".
        assignments[*iter] = 1;  // assign to "yes" (1).
      }
    }
    std::vector<Clusterable*> clusters(2, (Clusterable*)NULL);  // no, yes.
    kaldi::AddToClusters(summed_stats, assignments, &clusters);
  
    EnsureClusterableVectorNotNull(&summed_stats);
    EnsureClusterableVectorNotNull(&clusters);
  
    // even if improvement == 0 we continue; if we do RefineClusters we may get further improvement.
    // now do the RefineClusters stuff.  Note that this is null-op if
    // q_opts.GetQuestionsOf(key).refine_opts.num_iters == 0.  We could check for this but don't bother;
    // it happens in RefineClusters anyway.
  
    if (q_opts.GetQuestionsOf(key).refine_opts.num_iters > 0) {
      // If we want to refine the questions... (a bit like k-means w/ 2 classes).
      // Note: the only reason we introduced the if-statement is so the yes_set
      // doesn't get modified (truncated, actually) if we do the refine stuff with
      // zero iters.
      BaseFloat refine_impr = RefineClusters(summed_stats, &clusters, &assignments,
                                             q_opts.GetQuestionsOf(key).refine_opts);
      KALDI_ASSERT(refine_impr > std::min(-1.0, -0.1*fabs(improvement)));
      // refine_impr should always be positive
      improvement += refine_impr;
      yes_set.clear();
      for (size_t i = 0;i < assignments.size();i++) if (assignments[i] == 1) yes_set.push_back(i);
    }
    *yes_set_out = yes_set;
  
    DeletePointers(&clusters);
  #ifdef KALDI_PARANOID
    {  // Check the "ans" is correct.
      KALDI_ASSERT(clusters.size() == 2 && clusters[0] == 0 && clusters[1] == 0);
      AddToClusters(summed_stats, assignments, &clusters);
      BaseFloat impr;
      if (clusters[0] == NULL || clusters[1] == NULL) impr = 0.0;
      else impr = clusters[0]->Distance(*(clusters[1]));
      if (!ApproxEqual(impr, improvement) && fabs(impr-improvement) > 0.01) {
        KALDI_WARN << "FindBestSplitForKey: improvements do not agree: "<< impr
                   << " vs. " << improvement;
      }
      DeletePointers(&clusters);
    }
  #endif
    DeletePointers(&summed_stats);
    return improvement; // objective-function improvement.
  }
  
  
  
  /*
    DecisionTreeBuilder is a class used in SplitDecisionTree
  */
  
  class DecisionTreeSplitter {
   public:
    EventMap *GetMap() {
      if (!yes_) {  // leaf.
        return new ConstantEventMap(leaf_);
      } else {
        return new SplitEventMap(key_, yes_set_, yes_->GetMap(), no_->GetMap());
      }
    }
    BaseFloat BestSplit() { return best_split_impr_; } // returns objf improvement (>=0) of best possible split.
    void DoSplit(int32 *next_leaf) {
      if (!yes_) {  // not already split; we are a leaf, so split.
        DoSplitInternal(next_leaf);
      } else {  // find which of our children is best to split, and split that.
        (yes_->BestSplit() >= no_->BestSplit() ? yes_ : no_)->DoSplit(next_leaf);
        best_split_impr_ = std::max(yes_->BestSplit(), no_->BestSplit());  // may have changed.
      }
    }
    DecisionTreeSplitter(EventAnswerType leaf, const BuildTreeStatsType &stats,
                        const Questions &q_opts): q_opts_(q_opts), yes_(NULL), no_(NULL), leaf_(leaf), stats_(stats) {
      // not, this must work when stats is empty too. [just gives zero improvement, non-splittable].
      FindBestSplit();
    }
    ~DecisionTreeSplitter() {
      delete yes_;
      delete no_;
    }
   private:
    void DoSplitInternal(int32 *next_leaf) {
      // Does the split; applicable only to leaf nodes.
      KALDI_ASSERT(!yes_);  // make sure children not already set up.
      KALDI_ASSERT(best_split_impr_ > 0);
      EventAnswerType yes_leaf = leaf_, no_leaf = (*next_leaf)++;
      leaf_ = -1;  // we now have no leaf.
      // Now split the stats.
      BuildTreeStatsType yes_stats, no_stats;
      yes_stats.reserve(stats_.size()); no_stats.reserve(stats_.size());  //  probably better than multiple resizings.
      for (BuildTreeStatsType::const_iterator iter = stats_.begin(); iter != stats_.end(); ++iter) {
        const EventType &vec = iter->first;
        EventValueType val;
        if (!EventMap::Lookup(vec, key_, &val)) KALDI_ERR << "DoSplitInternal: key has no value.";
        if (std::binary_search(yes_set_.begin(), yes_set_.end(), val)) yes_stats.push_back(*iter);
        else no_stats.push_back(*iter);
      }
  #ifdef KALDI_PARANOID
      {  // Check objf improvement.
        Clusterable *yes_clust = SumStats(yes_stats), *no_clust = SumStats(no_stats);
        BaseFloat impr_check = yes_clust->Distance(*no_clust);
        // this is a negated objf improvement from merging (== objf improvement from splitting).
        if (!ApproxEqual(impr_check, best_split_impr_, 0.01)) {
          KALDI_WARN << "DoSplitInternal: possible problem: "<< impr_check << " != " << best_split_impr_;
        }
        delete yes_clust; delete no_clust;
      }
  #endif
      yes_ = new DecisionTreeSplitter(yes_leaf, yes_stats, q_opts_);
      no_ = new DecisionTreeSplitter(no_leaf, no_stats, q_opts_);
      best_split_impr_ = std::max(yes_->BestSplit(), no_->BestSplit());
      stats_.clear();  // note: pointers in stats_ were not owned here.
    }
    void FindBestSplit() {
      // This sets best_split_impr_, key_ and yes_set_.
      // May just pick best question, or may iterate a bit (depends on
      // q_opts; see FindBestSplitForKey for details)
      std::vector<EventKeyType> all_keys;
      q_opts_.GetKeysWithQuestions(&all_keys);
      if (all_keys.size() == 0) {
        KALDI_WARN << "DecisionTreeSplitter::FindBestSplit(), no keys available to split on (maybe no key covered all of your events, or there was a problem with your questions configuration?)";
      }
      best_split_impr_ = 0;
      for (size_t i = 0; i < all_keys.size(); i++) {
        if (q_opts_.HasQuestionsForKey(all_keys[i])) {
          std::vector<EventValueType> temp_yes_set;
          BaseFloat split_improvement = FindBestSplitForKey(stats_, q_opts_, all_keys[i], &temp_yes_set);
          if (split_improvement > best_split_impr_) {
            best_split_impr_ = split_improvement;
            yes_set_ = temp_yes_set;
            key_ = all_keys[i];
          }
        }
      }
    }
  
  
  
    // Data members... Always used:
    const Questions &q_opts_;
    BaseFloat best_split_impr_;
  
    // If already split:
    DecisionTreeSplitter *yes_;
    DecisionTreeSplitter *no_;
  
    // Otherwise:
    EventAnswerType leaf_;
    BuildTreeStatsType stats_;  // vector of stats.  pointers inside there not owned here.
  
    // key and "yes set" of best split:
    EventKeyType key_;
    std::vector<EventValueType> yes_set_;
  
  };
  
  EventMap *SplitDecisionTree(const EventMap &input_map,
                              const BuildTreeStatsType &stats,
                              Questions &q_opts,
                              BaseFloat thresh,
                              int32 max_leaves,  // max_leaves<=0 -> no maximum.
                              int32 *num_leaves,
                              BaseFloat *obj_impr_out,
                              BaseFloat *smallest_split_change_out) {
    KALDI_ASSERT(num_leaves != NULL && *num_leaves > 0);  // can't be 0 or input_map would be empty.
    int32 num_empty_leaves = 0;
    BaseFloat like_impr = 0.0;
    BaseFloat smallest_split_change = 1.0e+20;
    std::vector<DecisionTreeSplitter*> builders;
    {  // set up "builders" [one for each current leaf].  This array is never extended.
      // the structures generated during splitting remain as trees at each array location.
      std::vector<BuildTreeStatsType> split_stats;
      SplitStatsByMap(stats, input_map, &split_stats);
      KALDI_ASSERT(split_stats.size() != 0);
      builders.resize(split_stats.size());  // size == #leaves.
      for (size_t i = 0;i < split_stats.size();i++) {
        EventAnswerType leaf = static_cast<EventAnswerType>(i);
        if (split_stats[i].size() == 0) num_empty_leaves++;
        builders[i] = new DecisionTreeSplitter(leaf, split_stats[i], q_opts);
      }
    }
  
    {  // Do the splitting.
      int32 count = 0;
      std::priority_queue<std::pair<BaseFloat, size_t> > queue;  // use size_t because logically these
      // are just indexes into the array, not leaf-ids (after splitting they are no longer leaf id's).
      // Initialize queue.
      for (size_t i = 0; i < builders.size(); i++)
        queue.push(std::make_pair(builders[i]->BestSplit(), i));
      // Note-- queue's size never changes from now.  All the alternatives leaves to split are
      // inside the "DecisionTreeSplitter*" objects, in a tree structure.
      while (queue.top().first > thresh
            && (max_leaves<=0 || *num_leaves < max_leaves)) {
        smallest_split_change = std::min(smallest_split_change, queue.top().first);
        size_t i = queue.top().second;
        like_impr += queue.top().first;
        builders[i]->DoSplit(num_leaves);
        queue.pop();
        queue.push(std::make_pair(builders[i]->BestSplit(), i));
        count++;
      }
      KALDI_LOG << "DoDecisionTreeSplit: split "<< count << " times, #leaves now " << (*num_leaves);
    }
  
    if (smallest_split_change_out)
      *smallest_split_change_out = smallest_split_change;
  
    EventMap *answer = NULL;
  
    {  // Create the output EventMap.
      std::vector<EventMap*> sub_trees(builders.size());
      for (size_t i = 0; i < sub_trees.size();i++) sub_trees[i] = builders[i]->GetMap();
      answer = input_map.Copy(sub_trees);
      for (size_t i = 0; i < sub_trees.size();i++) delete sub_trees[i];
    }
    // Free up memory.
    for (size_t i = 0;i < builders.size();i++) delete builders[i];
  
    if (obj_impr_out != NULL) *obj_impr_out = like_impr;
    return answer;
  }
  
  
  int ClusterEventMapGetMapping(const EventMap &e_in,
                                const BuildTreeStatsType &stats,
                                BaseFloat thresh,
                                std::vector<EventMap*> *mapping) {
    // First map stats
    KALDI_ASSERT(stats.size() != 0);
    std::vector<BuildTreeStatsType> split_stats;
    SplitStatsByMap(stats, e_in, &split_stats);
    std::vector<Clusterable*> summed_stats;
    SumStatsVec(split_stats, &summed_stats);
  
    std::vector<int32> indexes;
    std::vector<Clusterable*> summed_stats_contiguous;
    size_t max_index = 0;
    for (size_t i = 0;i < summed_stats.size();i++) {
      if (summed_stats[i] != NULL) {
        indexes.push_back(i);
        summed_stats_contiguous.push_back(summed_stats[i]);
        if (i > max_index) max_index = i;
      }
    }
    if (summed_stats_contiguous.empty()) {
      KALDI_WARN << "ClusterBottomUp: nothing to cluster.";
      return 0;  // nothing merged.
    }
  
    std::vector<int32> assignments;
    BaseFloat normalizer = SumClusterableNormalizer(summed_stats_contiguous),
        change;
    change = ClusterBottomUp(summed_stats_contiguous,
                             thresh,
                             0,  // no min-clust: use threshold for now.
                             NULL,  // don't need clusters out.
                             &assignments);  // this algorithm is quadratic, so might be quite slow.
  
  
    KALDI_ASSERT(assignments.size() == summed_stats_contiguous.size() && !assignments.empty());
    size_t num_clust = * std::max_element(assignments.begin(), assignments.end()) + 1;
    int32 num_combined = summed_stats_contiguous.size() - num_clust;
    KALDI_ASSERT(num_combined >= 0);
  
    KALDI_VLOG(2) <<  "ClusterBottomUp combined "<< num_combined
                  << " leaves and gave a likelihood change of " << change
                  << ", normalized = " << (change/normalizer)
                  << ", normalizer = " << normalizer;
    KALDI_ASSERT(change < 0.0001);  // should be negative or zero.
  
    KALDI_ASSERT(mapping != NULL);
    if (max_index >= mapping->size()) mapping->resize(max_index+1, NULL);
  
    for (size_t i = 0;i < summed_stats_contiguous.size();i++) {
      size_t index = indexes[i];
      size_t new_index = indexes[assignments[i]];  // index assigned by clusterig-- map to existing indices in the map,
      // that we clustered from, so we don't conflict with indices in other parts
      // of the tree.
      KALDI_ASSERT((*mapping)[index] == NULL || "Error: Cluster seems to have been "
                   "called for different parts of the tree with overlapping sets of "
                   "indices.");
      (*mapping)[index] = new ConstantEventMap(new_index);
    }
    DeletePointers(&summed_stats);
    return num_combined;
  }
  
  
  EventMap *RenumberEventMap(const EventMap &e_in, int32 *num_leaves) {
    EventType empty_vec;
    std::vector<EventAnswerType> initial_leaves;  // before renumbering.
    e_in.MultiMap(empty_vec, &initial_leaves);
    if (initial_leaves.empty()) {
      KALDI_ASSERT(num_leaves);
      if (num_leaves) *num_leaves = 0;
      return e_in.Copy();
    }
    SortAndUniq(&initial_leaves);
    EventAnswerType max_leaf_plus_one = initial_leaves.back() + 1;  // will typically, but not always, equal *num_leaves.
    std::vector<EventMap*> mapping(max_leaf_plus_one, (EventMap*)NULL);
    std::vector<EventAnswerType>::iterator iter = initial_leaves.begin(), end = initial_leaves.end();
    EventAnswerType cur_leaf = 0;
    for (; iter != end; ++iter) {
      KALDI_ASSERT(*iter >= 0 && *iter<max_leaf_plus_one);
      mapping[*iter] = new ConstantEventMap(cur_leaf++);
    }
    EventMap *ans = e_in.Copy(mapping);
    DeletePointers(&mapping);
    KALDI_ASSERT((size_t)cur_leaf == initial_leaves.size());
    if (num_leaves) *num_leaves = cur_leaf;
    return ans;
  }
  
  EventMap *MapEventMapLeaves(const EventMap &e_in,
                              const std::vector<int32> &mapping_in) {
    std::vector<EventMap*> mapping(mapping_in.size());
    for (size_t i = 0; i < mapping_in.size(); i++)
      mapping[i] = new ConstantEventMap(mapping_in[i]);
    EventMap *ans = e_in.Copy(mapping);
    DeletePointers(&mapping);
    return ans;
  }
  
  EventMap *ClusterEventMap(const EventMap &e_in, const BuildTreeStatsType &stats,
                            BaseFloat thresh, int32 *num_removed_ptr) {
    std::vector<EventMap*> mapping;
    int32 num_removed = ClusterEventMapGetMapping(e_in, stats, thresh, &mapping);
    EventMap *ans = e_in.Copy(mapping);
    DeletePointers(&mapping);
    if (num_removed_ptr != NULL) *num_removed_ptr = num_removed;
    return ans;
  
  }
  
  EventMap *ShareEventMapLeaves(const EventMap &e_in, EventKeyType key,
                                std::vector<std::vector<EventValueType> > &values,
                                int32 *num_leaves) {
    // the use of "pdfs" as the name of the next variable reflects the anticipated
    // use of this function.
    std::vector<std::vector<EventAnswerType> > pdfs(values.size());
    for (size_t i = 0; i < values.size(); i++) {
      EventType evec;
      for (size_t j = 0; j < values[i].size(); j++) {
        evec.push_back(MakeEventPair(key, values[i][j]));
        size_t size_at_start = pdfs[i].size();
        e_in.MultiMap(evec, &(pdfs[i]));  // append any corresponding pdfs to pdfs[i].
        if (pdfs[i].size() == size_at_start) {  // Nothing added... unexpected.
          KALDI_WARN << "ShareEventMapLeaves: had no leaves for key = " << key
                     << ", value = " << (values[i][j]);
        }
      }
      SortAndUniq(&(pdfs[i]));
    }
    std::vector<EventMap*> remapping;
    for (size_t i = 0; i < values.size(); i++) {
      if (pdfs[i].empty())
        KALDI_WARN << "ShareEventMapLeaves: no leaves in one bucket.";  // not expected.
      else {
        EventAnswerType map_to_this = pdfs[i][0];  // map all in this bucket
        // to this value.
        for (size_t j = 1; j < pdfs[i].size(); j++) {
          EventAnswerType leaf = pdfs[i][j];
          KALDI_ASSERT(leaf>=0);
          if (remapping.size() <= static_cast<size_t>(leaf))
            remapping.resize(leaf+1, NULL);
          KALDI_ASSERT(remapping[leaf] == NULL);
          remapping[leaf] = new ConstantEventMap(map_to_this);
        }
      }
    }
    EventMap *shared = e_in.Copy(remapping);
    DeletePointers(&remapping);
    EventMap *renumbered = RenumberEventMap(*shared, num_leaves);
    delete shared;
    return renumbered;
  }
  
  
  void DeleteBuildTreeStats(BuildTreeStatsType *stats) {
    KALDI_ASSERT(stats != NULL);
    BuildTreeStatsType::iterator iter = stats->begin(), end = stats->end();
    for (; iter!= end; ++iter) if (iter->second != NULL) { delete iter->second; iter->second = NULL; } // set to NULL for extra safety.
  }
  
  EventMap *GetToLengthMap(const BuildTreeStatsType &stats, int32 P,
                           const std::vector<EventValueType> *phones,
                           int32 default_length) {
    std::vector<BuildTreeStatsType> stats_by_phone;
    try {
      SplitStatsByKey(stats, P, &stats_by_phone);
    } catch(const KaldiFatalError &) {
      KALDI_ERR <<
          "You seem to have provided invalid stats [no central-phone key].";
    }
    std::map<EventValueType, EventAnswerType> phone_to_length;
    for (size_t p = 0; p < stats_by_phone.size(); p++) {
      if (! stats_by_phone[p].empty()) {
        std::vector<BuildTreeStatsType> stats_by_length;
        try {
          SplitStatsByKey(stats_by_phone[p], kPdfClass, &stats_by_length);
        } catch(const KaldiFatalError &) {
          KALDI_ERR <<
              "You seem to have provided invalid stats [no position key].";
        }
        size_t length = stats_by_length.size();
        for (size_t i = 0; i < length; i++) {
          if (stats_by_length[i].empty()) {
            KALDI_ERR << "There are no stats available for position " << i
                      << " of phone " << p;
          }
        }
        phone_to_length[p] = length;
      }
    }
    if (phones != NULL) {  // set default length for unseen phones.
      for (size_t i = 0; i < phones->size(); i++) {
        if (phone_to_length.count( (*phones)[i] ) == 0) {  // unseen.
          phone_to_length[(*phones)[i]] = default_length;
        }
      }
    }
    EventMap *ans = new TableEventMap(P, phone_to_length);
    return ans;
  }
  
  // Recursive routine that is a helper to ClusterEventMapRestricted.
  // returns number removed.
  static int32 ClusterEventMapRestrictedHelper(const EventMap &e_in,
                                               const BuildTreeStatsType &stats,
                                               BaseFloat thresh,
                                               std::vector<EventKeyType> keys,
                                               std::vector<EventMap*> *leaf_mapping) {
    if (keys.size() == 0) {
      return ClusterEventMapGetMapping(e_in, stats, thresh, leaf_mapping);
    } else {  // split on the key.
      int32 ans = 0;
      std::vector<BuildTreeStatsType> split_stats;
      SplitStatsByKey(stats, keys.back(), &split_stats);
      keys.pop_back();
      for (size_t i = 0; i< split_stats.size(); i++)
        if (split_stats[i].size() != 0)
          ans += ClusterEventMapRestrictedHelper(e_in, split_stats[i], thresh, keys, leaf_mapping);
      return ans;
    }
  }
  
  EventMap *ClusterEventMapRestrictedByKeys(const EventMap &e_in,
                                            const BuildTreeStatsType &stats,
                                            BaseFloat thresh,
                                            const std::vector<EventKeyType> &keys,
                                            int32 *num_removed) {
    std::vector<EventMap*> leaf_mapping;  // For output of ClusterEventMapGetMapping.
  
    int32 nr = ClusterEventMapRestrictedHelper(e_in, stats, thresh, keys, &leaf_mapping);
    if (num_removed != NULL) *num_removed = nr;
  
    EventMap *ans = e_in.Copy(leaf_mapping);
    DeletePointers(&leaf_mapping);
    return ans;
  }
  
  
  EventMap *ClusterEventMapRestrictedByMap(const EventMap &e_in,
                                           const BuildTreeStatsType &stats,
                                           BaseFloat thresh,
                                           const EventMap &e_restrict,
                                           int32 *num_removed_ptr) {
    std::vector<EventMap*> leaf_mapping;
  
    std::vector<BuildTreeStatsType> split_stats;
    int num_removed = 0;
    SplitStatsByMap(stats, e_restrict, &split_stats);
    for (size_t i = 0; i < split_stats.size(); i++) {
      if (!split_stats[i].empty())
        num_removed += ClusterEventMapGetMapping(e_in, split_stats[i], thresh,
                                                 &leaf_mapping);
    }
  
    if (num_removed_ptr != NULL) *num_removed_ptr = num_removed;
  
    EventMap *ans = e_in.Copy(leaf_mapping);
    DeletePointers(&leaf_mapping);
    return ans;
  }
  
  EventMap *ClusterEventMapToNClustersRestrictedByMap(
      const EventMap &e_in,
      const BuildTreeStatsType &stats,
      int32 num_clusters_required,
      const EventMap &e_restrict,
      int32 *num_removed_ptr) {
    std::vector<BuildTreeStatsType> split_stats;
    SplitStatsByMap(stats, e_restrict, &split_stats);
  
    if (num_clusters_required < split_stats.size()) {
      KALDI_WARN << "num-clusters-required is less than size of map. Not doing anything.";
      if (num_removed_ptr) *num_removed_ptr = 0;
      return e_in.Copy();
    }
  
    std::vector<std::vector<int32> > indexes(split_stats.size());
    std::vector<std::vector<Clusterable*> > summed_stats_contiguous(split_stats.size());
  
    BaseFloat normalizer = 0.0;
  
    size_t max_index = 0;
  
    int32 num_non_empty_clusters_required = num_clusters_required;
  
    int32 num_non_empty_clusters_in_map = 0;
    int32 num_non_empty_clusters = 0;
  
    for (size_t i = 0; i < split_stats.size(); i++) {
      if (!split_stats[i].empty()) {
        num_non_empty_clusters_in_map++;
  
        std::vector<BuildTreeStatsType> split_stats_i;
        SplitStatsByMap(split_stats[i], e_in, &split_stats_i);
        std::vector<Clusterable*> summed_stats_i;
        SumStatsVec(split_stats_i, &summed_stats_i);
  
        for (size_t j = 0; j < summed_stats_i.size(); j++) {
          if (summed_stats_i[j] != NULL) {
            num_non_empty_clusters++;
            indexes[i].push_back(j);
            summed_stats_contiguous[i].push_back(summed_stats_i[j]);
            if (j > max_index) max_index = j;
          }
        }
  
        normalizer += SumClusterableNormalizer(summed_stats_contiguous[i]);
      } else {
        // Even if split_stats[i] is empty, a cluster will be assigned to
        // that. To compensate, we decrease the num-clusters required.
        num_non_empty_clusters_required--;
      }
    }
  
    KALDI_VLOG(1) << "Number of non-empty clusters in map = " << num_non_empty_clusters_in_map;
    KALDI_VLOG(1) << "Number of non-empty clusters = " << num_non_empty_clusters;
  
    if (num_non_empty_clusters_required > num_non_empty_clusters) {
      KALDI_WARN << "Cannot get required num-clusters " << num_clusters_required
                 << " as number of non-empty clusters required is larger than "
                 << " number of non-empty clusters: " << num_non_empty_clusters_required
                 << " > " << num_non_empty_clusters;
      if (num_removed_ptr) *num_removed_ptr = 0;
      return e_in.Copy();
    }
  
    std::vector<std::vector<int32> > assignments;
    BaseFloat change = ClusterBottomUpCompartmentalized(
        summed_stats_contiguous,
        std::numeric_limits<BaseFloat>::infinity(),
        num_non_empty_clusters_required,
        NULL,  // don't need clusters out.
        &assignments);  // this algorithm is quadratic, so might be quite slow.
  
    KALDI_ASSERT(assignments.size() == split_stats.size());
  
    int32 num_combined = 0;
    for (size_t i = 0; i < split_stats.size(); i++) {
      KALDI_ASSERT(assignments[i].size() == summed_stats_contiguous[i].size());
      if (assignments[i].size() == 0) continue;
      size_t num_clust_i = *std::max_element(assignments[i].begin(),
                                             assignments[i].end()) + 1;
      num_combined += summed_stats_contiguous[i].size() - num_clust_i;
    }
  
    KALDI_VLOG(2) << "ClusterBottomUpCompartmentalized combined " << num_combined
                  << " leaves and gave a likelihood change of " << change
                  << ", normalized = " << (change / normalizer)
                  << ", normalizer = " << normalizer;
    KALDI_ASSERT(change < 0.0001);  // should be negative or zero.
  
    std::vector<EventMap*> leaf_mapping(max_index + 1, NULL);
  
    for (size_t i = 0; i < split_stats.size(); i++) {
      for (size_t j = 0; j < summed_stats_contiguous[i].size(); j++) {
        size_t index = indexes[i][j];
        size_t new_index = indexes[i][assignments[i][j]];
        // index assigned by clusterig-- map to existing indices in the map,
        // that we clustered from, so we don't conflict with indices in other parts
        // of the tree.
        KALDI_ASSERT(leaf_mapping[index] == NULL || "Error: Cluster seems to have been "
                     "called for different parts of the tree with overlapping sets of "
                     "indices.");
        leaf_mapping[index] = new ConstantEventMap(new_index);
      }
      DeletePointers(&summed_stats_contiguous[i]);
    }
  
    if (num_removed_ptr) *num_removed_ptr = num_combined;
  
    EventMap *ans = e_in.Copy(leaf_mapping);
    DeletePointers(&leaf_mapping);
    return ans;
  }
  
  EventMap *GetStubMap(int32 P,
                       const std::vector<std::vector<int32> > &phone_sets,
                       const std::vector<int32> &phone2num_pdf_classes,
                       const std::vector<bool> &share_roots,
                       int32 *num_leaves_out) {
  
    {  // Checking inputs.
      KALDI_ASSERT(!phone_sets.empty() && share_roots.size() == phone_sets.size());
      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]);
        }
      }
    }
  
    // Initially create a single leaf for each phone set.
  
    size_t max_set_size = 0;
    int32 highest_numbered_phone = 0;
    for (size_t i = 0; i < phone_sets.size(); i++) {
      max_set_size = std::max(max_set_size, phone_sets[i].size());
      highest_numbered_phone =
          std::max(highest_numbered_phone,
                   * std::max_element(phone_sets[i].begin(), phone_sets[i].end()));
    }
  
    if (phone_sets.size() == 1) {  // there is only one set so the recursion finishes.
      if (share_roots[0]) {  // if "shared roots" return a single leaf.
        return new ConstantEventMap( (*num_leaves_out)++ );
      } else {  // not sharing roots -> work out the length and return a
               // TableEventMap splitting on length.
        EventAnswerType max_len = 0;
        for (size_t i = 0; i < phone_sets[0].size(); i++) {
          EventAnswerType len;
          EventValueType phone = phone_sets[0][i];
          KALDI_ASSERT(static_cast<size_t>(phone) < phone2num_pdf_classes.size());
          len = phone2num_pdf_classes[phone];
          KALDI_ASSERT(len > 0);
          if (i == 0) max_len = len;
          else {
            if (len != max_len) {
              KALDI_WARN << "Mismatching lengths within a phone set: " << len
                         << " vs. " << max_len << " [unusual, but not necessarily fatal]. ";
              max_len = std::max(len, max_len);
            }
          }
        }
        std::map<EventValueType, EventAnswerType> m;
        for (EventAnswerType p = 0; p < max_len; p++)
          m[p] = (*num_leaves_out)++;
        return new TableEventMap(kPdfClass,  // split on hmm-position
                                 m);
      }
    } else if (max_set_size == 1
              && static_cast<int32>(phone_sets.size()) <= 2*highest_numbered_phone) {
      // create table map splitting on phone-- more efficient.
      // the part after the && checks that this would not contain a very sparse vector.
      std::map<EventValueType, EventMap*> m;
      for (size_t i = 0; i < phone_sets.size(); i++) {
        std::vector<std::vector<int32> > phone_sets_tmp;
        phone_sets_tmp.push_back(phone_sets[i]);
        std::vector<bool> share_roots_tmp;
        share_roots_tmp.push_back(share_roots[i]);
        EventMap *this_stub = GetStubMap(P, phone_sets_tmp, phone2num_pdf_classes,
                                         share_roots_tmp,
                                         num_leaves_out);
        KALDI_ASSERT(m.count(phone_sets_tmp[0][0]) == 0);
        m[phone_sets_tmp[0][0]] = this_stub;
      }
      return new TableEventMap(P, m);
    } else {
      // Do a split.  Recurse.
      size_t half_sz = phone_sets.size() / 2;
      std::vector<std::vector<int32> >::const_iterator half_phones =
          phone_sets.begin() + half_sz;
      std::vector<bool>::const_iterator half_share =
          share_roots.begin() + half_sz;
      std::vector<std::vector<int32> > phone_sets_1, phone_sets_2;
      std::vector<bool> share_roots_1, share_roots_2;
      phone_sets_1.insert(phone_sets_1.end(), phone_sets.begin(), half_phones);
      phone_sets_2.insert(phone_sets_2.end(), half_phones, phone_sets.end());
      share_roots_1.insert(share_roots_1.end(), share_roots.begin(), half_share);
      share_roots_2.insert(share_roots_2.end(), half_share, share_roots.end());
  
      EventMap *map1 = GetStubMap(P, phone_sets_1, phone2num_pdf_classes, share_roots_1, num_leaves_out);
      EventMap *map2 = GetStubMap(P, phone_sets_2, phone2num_pdf_classes, share_roots_2, num_leaves_out);
  
      std::vector<EventKeyType> all_in_first_set;
      for (size_t i = 0; i < half_sz; i++)
        for (size_t j = 0; j < phone_sets[i].size(); j++)
          all_in_first_set.push_back(phone_sets[i][j]);
      std::sort(all_in_first_set.begin(), all_in_first_set.end());
      KALDI_ASSERT(IsSortedAndUniq(all_in_first_set));
      return new SplitEventMap(P, all_in_first_set, map1, map2);
    }
  }
  
  // convert stats to different (possibly smaller) context-window size.
  bool ConvertStats(int32 oldN, int32 oldP, int32 newN, int32 newP,
                    BuildTreeStatsType *stats) {
    bool warned = false;
    KALDI_ASSERT(stats != NULL && oldN > 0 && newN > 0 && oldP >= 0
                 && newP >= 0 && newP < newN && oldP < oldN);
    if (newN > oldN) {  // can't add unseen context.
      KALDI_WARN << "Cannot convert stats to larger context: " << newN
                 << " > " << oldN;
      return false;
    }
    if (newP > oldP) {
      KALDI_WARN << "Cannot convert stats to have more left-context: " << newP
                 << " > " << oldP;
    }
    if (newN-newP-1 > oldN-oldP-1) {
      KALDI_WARN << "Cannot convert stats to have more right-context: " << (newN-newP-1)
                 << " > " << (oldN-oldP-1);
    }
    // if shift < 0, this is by how much we "shift down" key
    // values.
    int32 shift = newP - oldP;  // shift <= 0.
  
    for (size_t i = 0; i < stats->size(); i++) {
      EventType &evec = (*stats)[i].first;
      EventType evec_new;
      for (size_t j = 0; j < evec.size(); j++) {
        EventKeyType key = evec[j].first;
        if (key >= 0 && key < oldN) {  //
          key += shift;
          if (key >= 0 && key < newN) // within new context window...
            evec_new.push_back(std::make_pair(key, evec[j].second));
        } else {
          if (key != -1) {
            // don't understand this key value but assume for now
            // it's something that doesn't interact with the context window.
            if (!warned) {
              KALDI_WARN << "Stats had keys defined that we cannot interpret";
              warned = true;
            }
          }
          evec_new.push_back(evec[j]);
        }
      }
      evec = evec_new;  // Assign the modified EventVector with possibly
      // deleted keys.
    }
    return true;
  }
  
  
  } // end namespace kaldi