// nnet3/nnet-compile-utils.cc
  
  // Copyright      2015-2017  Johns Hopkins University (author: Daniel Povey)
  //                2015                           (author: Vijayaditya Peddinti)
  
  // 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 <iterator>
  #include <sstream>
  #include "util/common-utils.h"
  #include "nnet3/nnet-compile-utils.h"
  
  namespace kaldi {
  namespace nnet3 {
  
  /**
     Gets counts of submatrices (the 1st members of pairs) in submat_lists.
     Also outputs, to 'submats_with_large_counts', a list of submatrix indexes
     that have counts over half of submat_lists.size().  (These will be separated
     out into their own AddRows() commands).
   */
  void GetSubmatCounts(
      const std::vector<std::vector<std::pair<int32, int32> > > &submat_lists,
      std::unordered_map<int32,int32> *submat_counts,
      std::vector<int32> *submats_with_large_counts) {
    auto iter = submat_lists.begin(), end = submat_lists.end();
    for (; iter != end; ++iter) {
      std::vector<std::pair<int32, int32> >::const_iterator
          iter2 = iter->begin(), end2 = iter->end();
      for (; iter2 != end2; ++iter2) {
        int32 submat_index = iter2->first;
        KALDI_ASSERT(submat_index >= 0);  // We don't expect -1's in submat_lists.
        std::unordered_map<int32,int32>::iterator
            iter = submat_counts->find(submat_index);
        if (iter == submat_counts->end())
          (*submat_counts)[submat_index] = 1;
        else
          iter->second++;
      }
    }
    auto counts_iter = submat_counts->begin(),
        counts_end = submat_counts->end();
    size_t cutoff = submat_lists.size() / 2;
    for (; counts_iter != counts_end; ++counts_iter)
      if (counts_iter->second > cutoff)
        submats_with_large_counts->push_back(counts_iter->first);
  }
  
  /**
     This function, used in SplitLocations(), is used to make separate
     'split lists' for certain high-count submatrix indexes, specified by
     the user in 'submats_to_separate'.  These split
     lists will be lists of pairs that are all either (-1, 1) or (submatrix_index, x)
     for a particular submatrix index (constant within the split list).
     These high-count lists will be written to 'split_lists'; they
     will eventually compile to AddRows() commands.  We write the remaining
     members of the lists in 'submat_lists' (the ones that did not make it
     into 'split_lists') to 'reduced_submat_lists'.
   */
  void SeparateSubmatsWithLargeCounts(
      const std::vector<int32> &submats_to_separate,
      const std::vector<std::vector<std::pair<int32, int32> > > &submat_lists,
      std::vector<std::vector<std::pair<int32, int32> > > *reduced_submat_lists,
      std::vector<std::vector<std::pair<int32, int32> > > *split_lists) {
    KALDI_ASSERT(split_lists->empty() && !submats_to_separate.empty());
    size_t num_to_separate = submats_to_separate.size(),
        num_rows = submat_lists.size();
    std::unordered_map<int32, size_t> submat_to_index;
    reduced_submat_lists->clear();
    reduced_submat_lists->resize(num_rows);
    split_lists->resize(num_to_separate);
    for (size_t i = 0; i < num_to_separate; i++) {
      (*split_lists)[i].resize(num_rows, std::pair<int32, int32>(-1, -1));
      int32 submat = submats_to_separate[i];
      submat_to_index[submat] = i;
    }
    for (size_t row = 0; row < submat_lists.size(); row++) {
      std::vector<std::pair<int32, int32> >::const_iterator
          iter = submat_lists[row].begin(), end = submat_lists[row].end();
      std::vector<std::pair<int32, int32> >
          &reduced_list = (*reduced_submat_lists)[row];
      // 'reduced_lists' will contain the pairs that don't make it into
      // 'split_lists'.
      for (; iter != end; ++iter) {
        int32 submat_index = iter->first;
        std::unordered_map<int32, size_t>::const_iterator map_iter =
            submat_to_index.find(submat_index);
        if (map_iter == submat_to_index.end()) { // not a large-count submatrix.
          reduced_list.push_back(*iter);
          continue;
        }
        size_t index = map_iter->second;
        std::pair<int32,int32> &p = (*split_lists)[index][row];
        if (p.first >= 0) {
          // we'd only reach here if the same submat index repeated in the same
          // row, which is possible but rare.
          reduced_list.push_back(*iter);
          continue;
        }
        p.first = submat_index;
        int32 src_row_index = iter->second;
        p.second = src_row_index;
      }
    }
  }
  
  void SplitLocations(
      const std::vector<std::vector<std::pair<int32, int32> > > &submat_lists,
      std::vector<std::vector<std::pair<int32, int32> > > *split_lists) {
    size_t num_rows = submat_lists.size(),
        num_output_lists = 0;
    auto iter = submat_lists.begin(), end = submat_lists.end();
    for (; iter != end; ++iter)
      if (iter->size() > num_output_lists)
        num_output_lists = iter->size();
    split_lists->clear();
    if (num_output_lists == 0)  // Odd, but could happen, maybe
      return;
    else if (num_output_lists == 1) {
      split_lists->resize(1);
      std::vector<std::pair<int32, int32> > &list = (*split_lists)[0];
      list.resize(num_rows, std::pair<int32, int32>(-1, -1));
      for (size_t i = 0; i < num_rows; i++) {
        if (!submat_lists[i].empty())
          list[i] = submat_lists[i][0];
      }
      return;
    }
  
    // counts for each submatrix index, of how many times it occurs.
    std::unordered_map<int32,int32> submat_counts;
    std::vector<int32> submats_with_large_counts;
    GetSubmatCounts(submat_lists, &submat_counts, &submats_with_large_counts);
    if (!submats_with_large_counts.empty()) {
      // There are submatrices with counts over half the num-rows.  We assign these
      // their own output lists.
  
      std::vector<std::vector<std::pair<int32, int32> > > reduced_submat_lists;
      SeparateSubmatsWithLargeCounts(submats_with_large_counts,
                                     submat_lists,
                                     &reduced_submat_lists,
                                     split_lists);
      // 'reduced_split_lists' is the result of recursing with input 'reduced_submat_lists';
      // we'll append its result to 'split_lists'.
      std::vector<std::vector<std::pair<int32, int32> > > reduced_split_lists;
      SplitLocations(reduced_submat_lists, &reduced_split_lists);
      size_t cur_num_lists = split_lists->size(),
          num_extra_lists = reduced_split_lists.size(),
          new_num_lists = cur_num_lists + num_extra_lists;
      split_lists->resize(new_num_lists);
      for (size_t i = 0; i < num_extra_lists; i++)
        (*split_lists)[cur_num_lists + i].swap(reduced_split_lists[i]);
      return;
      // and we're done.
    } else {
      // All the counts of submatrix indexes seem to be small so we are resigned to
      // only using AddRowsMulti commands.
      split_lists->resize(num_output_lists);
      for (size_t i = 0; i < num_output_lists; i++)
        (*split_lists)[i].resize(num_rows, std::pair<int32, int32>(-1, -1));
      for (size_t row = 0; row < num_rows; row++) {
        const std::vector<std::pair<int32, int32> > &this_list =
            submat_lists[row];
        size_t this_list_size = submat_lists[row].size();
        for (size_t i = 0; i < this_list_size; i++) {
          (*split_lists)[i][row] = this_list[i];
        }
      }
    }
  }
  
  
  /* If it is the case for some i >= 0 that all the .first elements of
     "location_vector" are either i or -1, then output i to first_value and the
     .second elements into "second_values", and return true.  Otherwise return
     false and the outputs are don't-cares. */
  bool ConvertToIndexes(
      const std::vector<std::pair<int32, int32> > &location_vector,
      int32 *first_value,
      std::vector<int32> *second_values)  {
    *first_value = -1;
    second_values->clear();
    second_values->reserve(location_vector.size());
    std::vector<std::pair<int32, int32> >::const_iterator iter;
    for (iter = location_vector.begin(); iter < location_vector.end(); ++iter)  {
      if (iter->first != -1) {
        if (*first_value == -1)
          *first_value = iter->first;
        if (iter->first != *first_value)
          return false;
        second_values->push_back(iter->second);
      } else  {
        second_values->push_back(-1);
      }
    }
    return true;
  }
  
  
  // see declaration in header for documentation
  void EnsureContiguousProperty(
      const std::vector<int32> &indexes,
      std::vector<std::vector<int32> > *indexes_out) {
    indexes_out->clear();
    indexes_out->reserve(3);
    if (indexes.empty()) return;
    int32 max_value = *std::max_element(indexes.begin(), indexes.end());
    if (max_value == -1) return;
    std::vector<int32> num_segments_seen(max_value + 1, 0);
    int32 dim = indexes.size(), num_output_vectors = 0;
    for (int32 i = 0; i < dim;) {
      // note, we increment i within the loop.
      if (indexes[i] == -1) {
        i++;
        continue;
      }
      int32 value = indexes[i], start_index = i;
      for (; i < dim && indexes[i] == value; i++);
      int32 end_index = i;  // one past the end.
      // the input 'indexes' contains a sequence of possibly-repeated instances of
      // the value 'value', starting at index 'start_index', with 'end_index' as
      // one past the end.
      int32 this_num_segments_seen = num_segments_seen[value]++;
      if (this_num_segments_seen >= num_output_vectors) {  // we have nowhere to
                                                           // put it.
        indexes_out->resize(++num_output_vectors);
        indexes_out->back().resize(dim, -1);  // fill newly added vector with -1's.
      }
      std::vector<int32> &this_out_vec((*indexes_out)[this_num_segments_seen]);
      std::vector<int32>::iterator iter = this_out_vec.begin() + start_index,
          end = this_out_vec.begin() + end_index;
      // Fill the appropriate range of the output vector with 'value'
      for (; iter != end; ++iter) *iter = value;
    }
  }
  
  
  
  /**
     This function splits a vector of pairs into a list of vectors of pairs.
     [note: by 'vector' we mean something that has a meaningful index that we care
     about; by 'list' we mean a collection of elements to be iterated over, without
     (in this case) meaningful indexes or even order.
  
     @param [in] list   A vector of pairs; these pairs should be either (-1,-1)
                        or (a,b) for a >= 0, b >= 0.  At least one element of 'list'
                        must be different from (-1,-1).
     @param [out] split_lists   A list, in arbitrary order, of vectors of pairs.
                       It has the following relationship with 'list':
                       - Size: for each j, split_lists[j].size() == list.size().
                       - Contents must match input: For each i:
                         -  If list[i] == (-1, -1), then
                            split_lists[j][i] == (-1, -1) for all j.
                         -  If list[i] != (-1, -1), then
                           split_lists[j][i] == (-1, -1) for *all but one* j, and
                           for the remaining j, split_lists[j][i] == list[i].
                       - Uniqueness: for no j should split_lists[j] contain
                         any duplicate elements (except the pair (-1,-1), which  is
                         allowed to exist in duplicate form).
                       To satisfy the above conditions, this function will create
                       as many lists in split_lists (i.e. as many j values) as the
                       number of times that the most frequent pair in 'list'
                       repeats other than the pair (-1,-1), e.g. if the pair
                       (10,11) appears 4 times in 'list' and that is the most,
                       split_lists->size() == 4.
  */
  void SplitPairList(std::vector<std::pair<int32, int32> >& list,
                     std::vector<std::vector<std::pair<int32, int32> > >* split_lists) {
    split_lists->clear();
    typedef unordered_map<std::pair<int32, int32>,
                          int32, PairHasher<int32> > MapType;
    // this maps a pair not equal to -1,-1, to the number of times we've already seen it.
    MapType pair_to_count;
    int32 cur_num_lists = 0;
  
    for (int32 i = 0; i < list.size(); i++)  {
      if (list[i].first == -1)
        continue;
      MapType::iterator iter = pair_to_count.find(list[i]);
      int32 this_count;
      if (iter == pair_to_count.end())
        pair_to_count[list[i]] = this_count = 1;
      else
        this_count = (++iter->second);
      if (this_count > cur_num_lists) {
        KALDI_ASSERT(this_count == cur_num_lists + 1);
        split_lists->resize(this_count);
        split_lists->back().resize(list.size(),
                                   std::pair<int32, int32>(-1, -1));
        cur_num_lists++;
      }
      (*split_lists)[this_count-1][i] = list[i];
    }
    if (split_lists->size() == 0)
      KALDI_ERR << "Input list has just dummy pairs";
  }
  
  void SplitLocationsBackward(
      const std::vector<std::vector<std::pair<int32, int32> > > &submat_lists,
      std::vector<std::vector<std::pair<int32, int32> > > *split_lists) {
    std::vector<std::vector<std::pair<int32, int32> > > split_lists_intermediate;
    // Split the submat_lists
    SplitLocations(submat_lists, &split_lists_intermediate);
    for (size_t i = 0; i < split_lists_intermediate.size(); i++) {
      int32 first_value;
      std::vector<int32> second_values;
      if (ConvertToIndexes(split_lists_intermediate[i],
                           &first_value, &second_values)) {
        // the .first values in split_lists_intermediate[i] are all the same (or
        // equal to -1).
        if (first_value == -1) {
          // all the .first values were equal to -1.  this is like a NULL marker.
          continue;
        }
        std::vector<std::vector<int32> > second_values_split;
        EnsureContiguousProperty(second_values, &second_values_split);
        if (second_values_split.size() == 1) {
          // this branch is an optimization for speed.
          split_lists->push_back(split_lists_intermediate[i]);
        } else {
          for (size_t j = 0; j < second_values_split.size(); j++) {
            split_lists->resize(split_lists->size() + 1);
            const std::vector<int32> &input_list = second_values_split[j];
            std::vector<std::pair<int32, int32> > &output_list =
                split_lists->back();
            output_list.resize(input_list.size());
            int32 size = input_list.size();
            for (int32 k = 0; k < size; k++) {
              int32 row = input_list[k];
              if (row == -1) output_list[k].first = -1;
              else output_list[k].first = first_value;
              output_list[k].second = row;
            }
          }
        }
      } else {
        // the .first values are not the same
        // splitting the list of pairs to ensure unique pairs, unless it is
        // (-1,-1)
        std::vector<std::vector<std::pair<int32, int32> > > new_split_lists;
        SplitPairList(split_lists_intermediate[i],
                      &new_split_lists);
        for (int32 j = 0; j < new_split_lists.size(); j++)  {
          split_lists->push_back(new_split_lists[j]);
        }
      }
    }
  }
  
  // This function returns true if for each integer i != -1, all the indexes j at
  // which indexes[j] == i are consecutive with no gaps (more formally: if j1 < j2
  // < j3 and indexes[j1] == indexes[j3], then indexes[j1] == indexes[j2]).  If
  // so, it also outputs to "reverse_indexes" the begin and end of these ranges,
  // so that indexes[j] == i for all j such that (*reverse_indexes)[i].first <= j
  // && j < (*reverse_indexes)[i].second.
  bool HasContiguousProperty(
      const std::vector<int32> &indexes,
      std::vector<std::pair<int32, int32> > *reverse_indexes) {
    reverse_indexes->clear();
    int32 num_indexes = indexes.size();
    if (num_indexes == 0)
      return true;
    int32 num_input_indexes =
        *std::max_element(indexes.begin(), indexes.end()) + 1;
    KALDI_ASSERT(num_input_indexes >= 0);
    if (num_input_indexes == 0) {
      // we don't really expect this input, filled with -1's.
      KALDI_WARN << "HasContiguousProperty called on vector of -1's.";
      return true;
    }
    reverse_indexes->resize(num_input_indexes,
                            std::pair<int32,int32>(-1, -1));
    // set each pair's "first" to the min index of all elements
    // of "indexes" with that value, and the "second" to the
    // max plus one.
    for (int32 i = 0; i < num_indexes; i++) {
      int32 j = indexes[i];
      if (j == -1) continue;
      KALDI_ASSERT(j >= 0);
      std::pair<int32, int32> &pair = (*reverse_indexes)[j];
      if (pair.first == -1) {
        pair.first = i;
        pair.second = i + 1;
      } else {
        pair.first = std::min(pair.first, i);
        pair.second = std::max(pair.second, i + 1);
      }
    }
    // check that the contiguous property holds.
    for (int32 i = 0; i < num_input_indexes; i++) {
      std::pair<int32, int32> pair = (*reverse_indexes)[i];
      if (pair.first != -1) {
        for (int32 j = pair.first; j < pair.second; j++)
          if (indexes[j] != i)
            return false;
      }
    }
    return true;
  }
  
  
  // see comment in header.
  void GetNxList(const std::vector<Index> &indexes,
                 std::vector<std::pair<int32, int32> > *pairs) {
    // set of (n,x) pairs
    std::unordered_set<std::pair<int32, int32>, PairHasher<int32> > n_x_set;
  
    for (std::vector<Index>::const_iterator iter = indexes.begin();
         iter != indexes.end(); ++iter)
      n_x_set.insert(std::pair<int32, int32>(iter->n, iter->x));
    pairs->clear();
    pairs->reserve(n_x_set.size());
    for (std::unordered_set<std::pair<int32, int32>, PairHasher<int32> >::iterator
             iter = n_x_set.begin(); iter != n_x_set.end(); ++iter)
      pairs->push_back(*iter);
    std::sort(pairs->begin(), pairs->end());
  }
  
  
  // see comment in header.
  void GetTList(const std::vector<Index> &indexes,
                std::vector<int32> *t_values) {
    // set of t values
    std::unordered_set<int32> t_set;
  
    for (std::vector<Index>::const_iterator iter = indexes.begin();
         iter != indexes.end(); ++iter)
      if (iter->t != kNoTime)
        t_set.insert(iter->t);
    t_values->clear();
    t_values->reserve(t_set.size());
    for (std::unordered_set<int32>::iterator iter = t_set.begin();
         iter != t_set.end(); ++iter)
      t_values->push_back(*iter);
    std::sort(t_values->begin(), t_values->end());
  }
  
  
  
  }  // namespace nnet3
  }  // namespace kaldi