// rnnlm/sampler.cc

// Copyright 2017  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 <algorithm>
#include <numeric>
#include <queue>
#include "rnnlm/sampler.h"
#include "base/kaldi-math.h"
#include "util/stl-utils.h"

namespace kaldi {
namespace rnnlm {


void SampleWithoutReplacement(const std::vector<double> &probs,
                              std::vector<int32> *sample) {

  // This outer loop over 't' will *almost always* just run for t == 0.  The
  // loop is necessary only to handle a pathological case.
  for (int32 t = 0; t < 10; t++) {
    sample->clear();
    int32 n = probs.size();

#define DO_SHUFFLE 0

#if DO_SHUFFLE
    // Removing the random shuffling for now because it turns out that
    // with multiple threads, it has to do locking inside the call to rand(),
    // which is quite slow; and using re-entrant random number generators with
    // this is quite complicated.  Anyway this wasn't necessary.
    // Maybe at some later point we can redo the data structures.


    // We randomize the order in which we process the indexes,
    // in order to reduce correlations.. not that this will
    // matter for most applications.
    std::vector<int32> order(n);
    for (int32 i = 0; i < n; i++) order[i] = i;
    std::random_shuffle(order.begin(), order.end());
#endif

    double r = RandUniform();  // r <= 0 <= 1.

    double c = -r;  // c is a kind of counter, to which we add the probabilities
                    // we we process them..  Whenever it becomes >= 0, we add something
                    // to the sample and subtract 1 from c.
    for (int32 i = 0; i < n; i++) {
#if DO_SHUFFLE
      int32 j = order[i];
#else
      int32 j = i;
#endif
      double p = probs[j];
      c += p;
      if (c >= 0) {
        sample->push_back(j);
        c -= 1.0;
      }
    }

    // you can verify by looking at the few lines of code above that
    // 'total_prob' is the total of the 'probs' array.
    double total_prob = c + sample->size() + r;
    int32 k = std::round(total_prob);
    if (std::abs(total_prob - k) > 1.0e-04) {
      // If this happened then the preconditions for this function are
      // violated-- the probs were not well enough normalized.  In double
      // precision the relative rounding error is about 10^-16, so to get an
      // error of 1.0e-04 from rounding we'd have to have 10^12 numbers which in
      // double precision would take 8000 G of memory-- unlikely.
      KALDI_ERR << "The sum of the inputs was " << k << " + "
                << (total_prob - k)
                << " which is too far from a whole number.";
    }
    if (sample->size() == k) {
      return;
    } else {
      // The only possible situations where the sample-size != k and we didn't
      // already crash, are when
      // c + r < 0.9999 or c + r > 0.9999.
      // Since -1 <= c < 0 and 0 <= r <= 1, this is only possible when r < 0.0001
      // and c < -0.9999, or r > 0.9999 and c >= -0.0001.

      // give it a bit of extra space in the assertion.
      KALDI_ASSERT((r < 0.00011 && c < -0.99985) ||
                   (r > 0.99985 && c > -0.00011));

      // .. and continue around the loop.
      // Having 'r' take these values is extremely improbable, so
      // it will be rare to go around the loop more than once.
    }
  }
  KALDI_ERR << "Looped too many times: likely bug.";
}


void CheckDistribution(const Distribution &d) {
  Distribution::const_iterator iter = d.begin(),
      endm1 = d.end() - 1;
  if (d.empty())
    return;
  for (; iter != endm1; ++iter) {
    KALDI_ASSERT(iter->second > 0.0 &&
                 iter->first < (iter+1)->first);
  }
  KALDI_ASSERT(d.back().second > 0.0);
}

void WeightDistribution(BaseFloat weight,
                        Distribution *d) {
  Distribution::iterator iter = d->begin(),
      end = d->end();
  for (; iter != end; ++iter)
    iter->second *= weight;
}


BaseFloat TotalOfDistribution(const Distribution &d) {
  double tot = 0.0;
  Distribution::const_iterator iter = d.begin(),
      end = d.end();
  for (; iter != end; ++iter)
    tot += iter->second;
  return tot;
}


const double* SampleFromCdf(const double *cdf_start,
                            const double *cdf_end) {
  double tot_prob = *cdf_end - *cdf_start;
  KALDI_ASSERT(cdf_end > cdf_start && tot_prob > 0.0);
  double cutoff = *cdf_start + tot_prob * RandUniform();
  if (cutoff >= *cdf_end) {
    // Mathematically speaking this should not happen; if it happens it is due
    // to roundoff.  It should be extremely rare in any case.
    cutoff = *cdf_start;
  }
  // With respect to the sample where [*cdf_start.. *cdf_end]
  // is [ 0.50, 0.55, 0.65, 0.70 ], suppose the randomly
  // sampled 'cutoff' is 0.68.  The std::upper_bound call
  // below finds the first location in the range
  // [ 0.55, 0.65, 0.70 ] that is > 0.68, which
  // in this case is '0.70'.  We then return that
  // pointer minus one, which is the pointer to 0.65.
  const double *ans = std::upper_bound(cdf_start + 1,
                                       cdf_end + 1,
                                       cutoff) - 1;
  // if the following assertion fails, it means that upper_bound returned
  // cdf_end + 1, which means that *cdf_end was not > 'cutoff'.  But we
  // ensured above that *cdf_end was > cutoff, so this should not happen.
  KALDI_ASSERT(ans != cdf_end);
  // If the following fails, it would be an error in our logic or in
  // std::upper_bound.
  KALDI_ASSERT(ans[1] != ans[0]);
  return ans;
}


// Merges two distributions, summing the probabilities of any elements that
// occur in both.
void MergeDistributions(const Distribution &d1,
                        const Distribution &d2,
                        Distribution *d) {
  if (GetVerboseLevel() >= 2) {
    CheckDistribution(d1);
    CheckDistribution(d2);
  }
  d->resize(d1.size() + d2.size());
  // we could write a single function that does the jobs of the
  // two things below, which might improve speed slightly, if
  // this becomes a bottleneck.
  std::merge(d1.begin(), d1.end(), d2.begin(), d2.end(), d->begin());
  MergePairVectorSumming(d);
  if (GetVerboseLevel() >= 2) {
    CheckDistribution(*d);
  }

}


void Sampler::SampleWords(
    int32 num_words_to_sample,
    BaseFloat unigram_weight,
    const std::vector<std::pair<int32, BaseFloat> > &higher_order_probs,
    const std::vector<int32> &words_we_must_sample,
    std::vector<std::pair<int32, BaseFloat> > *sample) const {
  CheckDistribution(higher_order_probs);  // TODO: delete this.
  int32 vocab_size = unigram_cdf_.size();
  KALDI_ASSERT(IsSortedAndUniq(words_we_must_sample) &&
               num_words_to_sample > 0 && num_words_to_sample < vocab_size);

  int32 num_words_we_must_sample = words_we_must_sample.size();
  if (num_words_we_must_sample > 0) {
    KALDI_ASSERT(num_words_we_must_sample < vocab_size &&
                 num_words_we_must_sample < num_words_to_sample);
    KALDI_ASSERT(words_we_must_sample.front() >= 0 &&
                 words_we_must_sample.back() < vocab_size);
  }

  BaseFloat total_existing_weight = unigram_weight +
      TotalOfDistribution(higher_order_probs);
  // To ensure that all the words we must sample actually get sampled,
  // we need to make sure that they are sampled with probability 1.0.
  // i.e. that after computing alpha, alpha p(i) for all of those words
  // is >= 1.0.  See comment for one of the versions of SampleWords()
  // in the header for more explanation.
  //
  // We can ensure this by making sure that after adding these words
  // each with probability p, each of these words has at least
  // weight T / num_words_to_sample, where 'T' is the new total
  // of the distribution.  I.e., we require that
  //   p > T / num_words_to_sample
  // Since T = total_existing_weight + p * num_words_we_must_sample, we have:
  //   p * num_words_to_sample > total_existing_weight + p * num_words_we_must_sample.
  // i.e.
  //  p > total_existing_weight / (num_words_to_sample - num_words_we_must_sample).
  // To minimize roundoff problems we make p just a little bigger than
  // that, bigger by a factor of 0.1.  So we'll set
  // p = 1.1 * total_existing_weight / (num_words_to_sample - num_words_we_must_sample).

  BaseFloat p = 1.1 * total_existing_weight /
      (num_words_to_sample - num_words_we_must_sample);

  std::vector<std::pair<int32, BaseFloat> > words_we_must_sample_distribution(
      num_words_we_must_sample);
  for (int32 i = 0 ; i < num_words_we_must_sample; i++) {
    words_we_must_sample_distribution[i].first = words_we_must_sample[i];
    words_we_must_sample_distribution[i].second = p;
  }

  std::vector<std::pair<int32, BaseFloat> > merged_distribution;
  MergeDistributions(higher_order_probs,
                     words_we_must_sample_distribution,
                     &merged_distribution);

  SampleWords(num_words_to_sample, unigram_weight,
              merged_distribution,
              sample);
  if (GetVerboseLevel() >= 2) {
    std::vector<int32> merged_list(words_we_must_sample);
    for (size_t i = 0; i < sample->size(); i++)
      merged_list.push_back((*sample)[i].first);
    SortAndUniq(&merged_list);
    // if the following assert fails, it means that one of the words
    // that we were required to sample, was not in fact sampled.
    // This implies there was a bug somewhere, or a flaw in
    // our reasoning.
    KALDI_ASSERT(merged_list.size() == sample->size());
  }
}

Sampler::Sampler(const std::vector<BaseFloat> &unigram_probs) {
  KALDI_ASSERT(!unigram_probs.empty());
  double total = std::accumulate(unigram_probs.begin(),
                                 unigram_probs.end(),
                                 0.0);
  KALDI_ASSERT(std::abs(total - 1.0) < 1.0e-02);
  double inv_total = 1.0 / total;

  double sum = 0.0;
  size_t n = unigram_probs.size();
  unigram_cdf_.resize(n + 1);
  unigram_cdf_[0] = 0.0;
  for (size_t i = 0; i < n; i++) {
    sum += unigram_probs[i];
    unigram_cdf_[i + 1] = sum * inv_total;
  }
}



void Sampler::SampleWords(
    int32 num_words_to_sample,
    BaseFloat unigram_weight,
    const std::vector<std::pair<int32, BaseFloat> > &higher_order_probs,
    std::vector<std::pair<int32, BaseFloat> > *sample) const {
  int32 vocab_size = unigram_cdf_.size() - 1;
  KALDI_ASSERT(num_words_to_sample > 0 &&
               num_words_to_sample + 1 < unigram_cdf_.size() &&
               unigram_weight > 0.0);
  if (!higher_order_probs.empty()) {
    KALDI_ASSERT(higher_order_probs.front().first >= 0 &&
                 higher_order_probs.back().first < vocab_size);
  }
  if (GetVerboseLevel() >= 2) {
    CheckDistribution(higher_order_probs);
  }

  std::vector<Interval> intervals;
  double total_p = GetInitialIntervals(unigram_weight, higher_order_probs,
                                       &intervals);
  // you can interpret total_p as sum_i p(i), with reference to the
  // math in the header next to the declaration of SampleWords().
  if (GetVerboseLevel() >= 2) {
    AssertEqual(total_p,
                unigram_weight + TotalOfDistribution(higher_order_probs));
  }
  NormalizeIntervals(num_words_to_sample, total_p, &intervals);
  SampleFromIntervals(intervals, sample);
}



// This hacked version of std::priority_queue allows us to extract all elements
// of the priority queue to a supplied vector, in an efficient way.  It relies
// on the fact that std::priority<queue> stores the underlying container as a
// protected member 'c'.  The only way to do this using the supplied interface
// of std::priority_queue is to repeatedly pop() the element from the queue, but
// that is too slow, and it actually had an impact on the speed of the
// application.
template <typename T>
class hacked_priority_queue: public std::priority_queue<T> {
 public:
  void append_all_elements(std::vector<T> *output) const {
    output->insert(output->end(), this->c.begin(), this->c.end());
  }
  // we have to redeclare the constructor.
  template <typename InputIter> hacked_priority_queue(
      InputIter begin, const InputIter end): std::priority_queue<T>(begin, end) { }
};


// static
void Sampler::NormalizeIntervals(int32 num_words_to_sample,
                                 double total_p,
                                 std::vector<Interval> *intervals) {
  // as mentioned in the header, if the input probabilities of the words
  // (here represented as Intervals)/ are p(i), we define
  //  q(i) = min(alpha p(i), 1.0)
  // where alpha is chosen so that sum_i q(i) == num_words_to_sample.
  // 'current_alpha' is initialized to the alpha that we would
  // have if none of the quantities alpha p(i) were greater than 1.
  // This function computes q(i).
  double current_alpha = num_words_to_sample / total_p;

  // 'num_ones' is the number of times the expression min(alpha p(i), 1.0)
  // is >= 1.0.
  int32 num_ones = 0;
  // 'total_remaining_p' is total_p [which equals the sum of p(i)] minus the
  // total of the p(i) for which we already know that alpha p(i) >= 1.
  double total_remaining_p = total_p;

  // In general, we will have:
  //  current_alpha = (num_words_to_sample - num_ones) / total_remaining_p.
  // As we update 'num_ones' and 'total_remaining_p', we will continue
  // to update current_alpha, and it will keep getting larger.
  hacked_priority_queue<Interval> queue(intervals->begin(), intervals->end());

  // clear 'intervals'; we'll use the space to store the intervals that will
  // have a prob of exactly 1.0, and eventually we'll add the rest.
  intervals->clear();
  while (!queue.empty()) {  // note: we normally won't reach the condition where
                          // queue.empty() here; we'll normally break.
    Interval top = queue.top();
    if (current_alpha * top.prob < 1.0) {
      break;  // we're done.
    } else {
      queue.pop();
      size_t interval_size = top.end - top.start;
      if (interval_size > 1) {
        // it's a range containing more than one thing -> we can split the
        // range.
        size_t half_size = interval_size / 2;
        double start_cdf = top.start[0],
            mid_cdf = top.start[half_size],
            end_cdf = top.end[0],
            total_unigram_prob = end_cdf - start_cdf,
            first_half_unigram_prob = mid_cdf - start_cdf,
            second_half_unigram_prob = (total_unigram_prob -
                                        first_half_unigram_prob),
            top_prob = top.prob;
        KALDI_ASSERT(total_unigram_prob > 0.0 && top_prob > 0.0);
        if (first_half_unigram_prob > 0.0) {
          queue.push(Interval(top_prob * first_half_unigram_prob /
                              total_unigram_prob,
                              top.start,
                              top.start + half_size));
        }
        if (second_half_unigram_prob > 0.0) {
          queue.push(Interval(top_prob * second_half_unigram_prob /
                              total_unigram_prob,
                              top.start + half_size, top.end));
        }
      } else {
        // It's an interval of one thing, we can't split it.
        // It's one of those things that takes the "1.0" in the expression
        // min(alpha p(i), 1.0).
        total_remaining_p -= top.prob;
        num_ones++;
        double new_alpha = (num_words_to_sample - num_ones) / total_remaining_p;
        top.prob = 1.0;
        intervals->push_back(top);
        KALDI_ASSERT(queue.empty() ||
                     (total_remaining_p > 0 && new_alpha > current_alpha));
        current_alpha = new_alpha;
      }
    }
  }
#if 0
  // The following code is a bit slow but has the advantage of not assuming
  // anything about the internals of class std::priority_queue.
  while (!queue.empty()) {
    Interval top = queue.top();
    top.prob *= current_alpha;
    queue.pop();
    intervals->push_back(top);
  }
#else
  { // This code is faster but relies on the fact that priority_queue
    // has a protected member 'c' which is the underlying container.
    size_t cur_size = intervals->size();
    queue.append_all_elements(intervals);
    // the next loop scales the 'prob' members of the elements we just
    // added to 'intervals', by current_alpha.
    std::vector<Interval>::iterator iter = intervals->begin() + cur_size,
        end = intervals->end();
    for (; iter != end; ++iter) iter->prob *= current_alpha;
  }
#endif

  if (GetVerboseLevel() >= 2) {
    double tot_prob = 0.0;
    for (size_t i = 0; i < intervals->size(); i++) {
      double p = (*intervals)[i].prob;
      KALDI_ASSERT(p > 0.0 && p <= 1.0);
      tot_prob += p;
    }
    KALDI_ASSERT(tot_prob - num_words_to_sample < 1.0e-04);
  }
}

void Sampler::SampleFromIntervals(const std::vector<Interval> &intervals,
                                  std::vector<std::pair<int32, BaseFloat> > *samples)  const {
  size_t num_intervals = intervals.size();
  std::vector<double> probs(num_intervals);
  for (size_t i = 0; i < num_intervals; i++)
    probs[i] = intervals[i].prob;
  // 'raw_samples' will contain indexes into the 'intervals' vector,
  // which we need to convert into actual words.
  std::vector<int32> raw_samples;
  SampleWithoutReplacement(probs, &raw_samples);
  size_t num_samples = raw_samples.size();
  samples->resize(num_samples);
  const double *cdf_start = &(unigram_cdf_[0]);
  for (size_t i = 0; i < num_samples; i++) {
    int32 j = raw_samples[i];  // j is interval index.
    const Interval &interval = intervals[j];
    if (interval.end == interval.start + 1) {
      // handle this simple case simply, even though
      // the general case on the other side of the if-statement
      // would give the correct expression.
      int32 word = interval.start - cdf_start;
      (*samples)[i].first = word;
      (*samples)[i].second = interval.prob;
    } else {
      const double *word_ptr = SampleFromCdf(interval.start,
                                             interval.end);
      int32 word = word_ptr - cdf_start;
      // the probability with which this word was sampled is: the probability of
      // sampling from this interval of the unigram, times the probability of
      // this word given the unigram distribution (which equals word[1] -
      // word[0]), divided by the probability of this whole unigram interval.
      // Actually we could more simply compute this, as unigram_weight * alpha *
      // (word_ptr[1] - word_ptr[0]), but alpha and unigram_weight not passed
      // into this function and I'd rather not make linkages between different
      // parts of the code.
      BaseFloat prob = interval.prob *
          (word_ptr[1] - word_ptr[0]) / (*interval.end - *interval.start);
      (*samples)[i].first = word;
      (*samples)[i].second = prob;
    }
  }
}


double Sampler::GetInitialIntervals(
    BaseFloat unigram_weight,
    const std::vector<std::pair<int32, BaseFloat> > &higher_order_probs,
    std::vector<Interval> *intervals) const {
  double ans = 0.0;
  intervals->clear();
  intervals->reserve(higher_order_probs.size() * 2 + 1);

  std::vector<std::pair<int32, BaseFloat> >::const_iterator
      h_iter = higher_order_probs.begin(),
      h_end = higher_order_probs.end();

  size_t vocab_size = unigram_cdf_.size() - 1;
  size_t cur_start = 0;
  const double *cdf = &(unigram_cdf_[0]);

  for (; h_iter != h_end; ++h_iter) {
    int32 w = h_iter->first;
    // include the unigram part of the probability in 'p': we add the unigram,
    // we don't back off to it.
    double p = h_iter->second +
        unigram_weight * (cdf[w + 1] - cdf[w]);
    KALDI_ASSERT(p > 0);
    if (w > cur_start && cdf[w] > cdf[cur_start]) {
      // Before we can add an Interval for (w, p), we have some lower-numbered
      // words to deal with.
      double range_p = unigram_weight * (cdf[w] - cdf[cur_start]);
      intervals->push_back(Interval(range_p,
                                    cdf + cur_start,
                                    cdf + w));
      ans += range_p;
    }
    intervals->push_back(Interval(p, cdf + w, cdf + w + 1));
    ans += p;
    cur_start = w + 1;
  }
  KALDI_ASSERT(cur_start <= vocab_size);
  double range_p = unigram_weight * (cdf[vocab_size] - cdf[cur_start]);
  if (range_p > 0) {
    intervals->push_back(Interval(range_p, cdf + cur_start, cdf + vocab_size));
    ans += range_p;
  }
  return ans;
}



}  // namespace rnnlm
}  // namespace kaldi