/** @file
 * @brief Collapse documents with the same collapse key during the match.
 */
/* Copyright (C) 2009,2011,2017 Olly Betts
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 2 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, see
 * <https://www.gnu.org/licenses/>.
 */

#ifndef XAPIAN_INCLUDED_COLLAPSER_H
#define XAPIAN_INCLUDED_COLLAPSER_H

#include "backends/documentinternal.h"
#include "msetcmp.h"
#include "omassert.h"
#include "api/result.h"

#include <unordered_map>
#include <vector>

/// Enumeration reporting how a result will be handled by the Collapser.
typedef enum {
    EMPTY,
    NEW,
    ADD,
    REJECT,
    REPLACE
} collapse_result;

/// Class tracking information for a given value of the collapse key.
class CollapseData {
    /** Currently kept MSet entries for this value of the collapse key.
     *
     *  If collapse_max > 1, then this is a min-heap once collapse_count > 0.
     *
     *  The first member of the pair is the index into proto_mset.results
     *  and the second is the docid of the entry (used to detect if the
     *  entry in proto_mset.results we were referring to has been dropped).
     *
     *  FIXME: We expect collapse_max to be small, so perhaps we should
     *  preallocate space for that many entries and/or allocate space in
     *  larger blocks to divvy up?
     */
    std::vector<std::pair<Xapian::doccount, Xapian::docid>> items;

    /// The highest weight of a document we've rejected.
    double next_best_weight = 0.0;

    /// The number of documents we've rejected.
    Xapian::doccount collapse_count = 0;

  public:
    /// Construct with the given item.
    CollapseData(Xapian::doccount item, Xapian::docid did)
        : items(1, { item, did })
    { }

    /** Check a new result with this collapse key value.
     *
     *  If this method determines the action to take is ADD or REPLACE, then
     *  the proto-mset should be updated and then add_item() called to complete
     *  the update of the CollapseData (if the result doesn't actually get
     *  added, then it's OK not to follow up with a call to add_item()).
     *
     *  @param results          The results so far.
     *  @param result           The new result.
     *  @param collapse_max     Max no. of items for each collapse key value.
     *  @param mcmp             Result comparison functor.
     *  @param[out] old_item    Item to be replaced (when REPLACE is returned).
     *
     *  @return How to handle @a result: ADD, REJECT or REPLACE.
     */
    collapse_result check_item(const std::vector<Result>& results,
                               const Result& result,
                               Xapian::doccount collapse_max,
                               MSetCmp mcmp,
                               Xapian::doccount& old_item);

    /** Set item after constructing with a placeholder.
     *
     *  @param item             The new item (index into results).
     */
    void set_item(Xapian::doccount item);

    /** Complete update of new result with this collapse key value.
     *
     *  @param results          The results so far.
     *  @param item             The new item (index into results).
     *  @param collapse_max     Max no. of items for each collapse key value.
     *  @param mcmp             Result comparison functor.
     */
    void add_item(const std::vector<Result>& results,
                  Xapian::doccount item,
                  Xapian::doccount collapse_max,
                  MSetCmp mcmp);

    /** Process relocation of entry in results.
     *
     *  @param from     The old item (index into results).
     *  @param to       The new item (index into results).
     */
    void result_has_moved(Xapian::doccount from, Xapian::doccount to) {
        // Yes, this *is* a linear search, but only through up to collapse_max
        // items and we expect collapse_max to be very small (it's the maximum
        // number of items to keep for each collapse key value).
        for (auto&& item : items) {
            if (item.first == from) {
                item.first = to;
                return;
            }
        }
        // The entry ought to be present.
        Assert(false);
    }

    /// The highest weight of a document we've rejected.
    double get_next_best_weight() const { return next_best_weight; }

    /// The number of documents we've rejected.
    Xapian::doccount get_collapse_count() const { return collapse_count; }
};

/// The Collapser class tracks collapse keys and the documents they match.
class Collapser {
    /// Map from collapse key values to the items we're keeping for them.
    std::unordered_map<std::string, CollapseData> table;

    /// How many items we're currently keeping in @a table.
    Xapian::doccount entry_count = 0;

    /** How many documents have we seen without a collapse key?
     *
     *  We use this statistic to improve matches_lower_bound.
     */
    Xapian::doccount no_collapse_key = 0;

    /** How many documents with duplicate collapse keys we have ignored.
     *
     *  We use this statistic to improve matches_estimated (by considering
     *  the rate of collapsing) and matches_upper_bound.
     */
    Xapian::doccount dups_ignored = 0;

    /** How many documents we've considered for collapsing.
     *
     *  We use this statistic to improve matches_estimated (by considering
     *  the rate of collapsing).
     */
    Xapian::doccount docs_considered = 0;

    /** The value slot we're getting collapse keys from. */
    Xapian::valueno slot;

    /** The maximum number of items to keep for each collapse key value. */
    Xapian::doccount collapse_max;

    std::vector<Result>& results;

    MSetCmp mcmp;

    /** Pointer to CollapseData when NEW or ADD is in progress. */
    CollapseData* ptr = NULL;

    /// Adapt @a mcmp to be usable with min_heap.
    bool operator()(Xapian::doccount a, Xapian::doccount b) const {
        return mcmp(results[a], results[b]);
    }

  public:
    /// Replaced item when REPLACE is returned by @a collapse().
    Xapian::doccount old_item = 0;

    Collapser(Xapian::valueno slot_,
              Xapian::doccount collapse_max_,
              std::vector<Result>& results_,
              MSetCmp mcmp_)
        : slot(slot_),
          collapse_max(collapse_max_),
          results(results_),
          mcmp(mcmp_) { }

    /// Return true if collapsing is active for this match.
    operator bool() const { return collapse_max != 0; }

    /** Check a new result.
     *
     *  If this method determines the action to take is NEW or ADD then the
     *  proto-mset should be updated and then process() called to complete the
     *  update (if the result doesn't actually get added, then it's OK not
     *  to follow up with a call to process()).
     *
     *  @param result   The new result.
     *  @param vsdoc    Document for getting values.
     *
     *  @return How to handle @a result: EMPTY, NEW, ADD, REJECT or REPLACE.
     */
    collapse_result check(Result& result,
                          Xapian::Document::Internal & vsdoc);

    /** Handle a new Result.
     *
     *  @param action   The collapse_result returned by check().
     *  @param item     The new item (index into results).
     */
    void process(collapse_result action, Xapian::doccount item);

    /** Process relocation of entry in results.
     *
     *  @param from     The old item (index into results).
     *  @param to       The new item (index into results).
     */
    void result_has_moved(Xapian::doccount from, Xapian::doccount to) {
        const std::string& collapse_key = results[to].get_collapse_key();
        if (collapse_key.empty()) {
            return;
        }
        auto it = table.find(collapse_key);
        if (rare(it == table.end())) {
            // The entry ought to be present.
            Assert(false);
            return;
        }

        CollapseData& collapse_data = it->second;
        collapse_data.result_has_moved(from, to);
    }

    Xapian::doccount get_collapse_count(const std::string & collapse_key,
                                        int percent_threshold,
                                        double min_weight) const;

    Xapian::doccount get_docs_considered() const { return docs_considered; }

    Xapian::doccount get_dups_ignored() const { return dups_ignored; }

    Xapian::doccount get_entries() const { return entry_count; }

    Xapian::doccount get_matches_lower_bound() const;

    void finalise(double min_weight, int percent_threshold);
};

/** Simpler version of Collapser used when merging MSet objects.
 *
 *  We merge results in descending rank order, so collapsing is much simpler
 *  than during the match - we just need to discard documents if we've already
 *  seen collapse_max with the same key.
 */
class CollapserLite {
    /// Map from collapse key values to collapse counts.
    std::unordered_map<std::string, Xapian::doccount> table;

    /// How many items we're currently keeping in @a table.
    Xapian::doccount entry_count = 0;

    /** How many documents have we seen without a collapse key?
     *
     *  We use this statistic to improve matches_lower_bound.
     */
    Xapian::doccount no_collapse_key = 0;

    /** How many documents with duplicate collapse keys we have ignored.
     *
     *  We use this statistic to improve matches_estimated (by considering
     *  the rate of collapsing) and matches_upper_bound.
     */
    Xapian::doccount dups_ignored = 0;

    /** How many documents we've considered for collapsing.
     *
     *  We use this statistic to improve matches_estimated (by considering
     *  the rate of collapsing).
     */
    Xapian::doccount docs_considered = 0;

    /** The maximum number of items to keep for each collapse key value. */
    Xapian::doccount collapse_max;

  public:
    CollapserLite(Xapian::doccount collapse_max_)
        : collapse_max(collapse_max_) {}

    /// Return true if collapsing is active for this match.
    operator bool() const { return collapse_max != 0; }

    /** Try to add a new key.
     *
     *  @return true if accepted; false if rejected.
     */
    bool add(const std::string& key) {
        ++docs_considered;

        if (key.empty()) {
            ++no_collapse_key;
            return true;
        }

        auto r = table.emplace(key, 1);
        if (r.second) {
            // New entry, set to 1.
        } else if (r.first->second == collapse_max) {
            // Already seen collapse_max with this key so reject.
            ++dups_ignored;
            return false;
        } else {
            // Increment count.
            ++r.first->second;
        }
        ++entry_count;
        return true;
    }

    Xapian::doccount get_docs_considered() const { return docs_considered; }

    Xapian::doccount get_dups_ignored() const { return dups_ignored; }

    Xapian::doccount get_entries() const { return entry_count; }

    Xapian::doccount get_matches_lower_bound() const {
        return no_collapse_key + entry_count;
    }

    void finalise(std::vector<Result>& results, int percent_threshold) {
        if (table.empty() || results.empty())
            return;

        // We need to fill in collapse_count values in results using the
        // information stored in table.
        Xapian::doccount todo = entry_count;
        for (Result& result : results) {
            const std::string& key = result.get_collapse_key();
            if (key.empty())
                continue;

            // Adjust collapse_count.
            if (percent_threshold) {
                // FIXME: We can probably do better here.
                result.set_collapse_count(1);
            } else {
                auto c = result.get_collapse_count() + table[key];
                result.set_collapse_count(c);
            }

            if (--todo == 0) {
                // Terminate early if we've handled all non-empty entries.
                break;
            }
        }
    }
};

#endif // XAPIAN_INCLUDED_COLLAPSER_H
