/** @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/>.
 */

#include <config.h>

#include "collapser.h"

#include "heap.h"
#include "omassert.h"

#include <algorithm>

using namespace std;

collapse_result
CollapseData::check_item(const vector<Result>& results,
                         const Result& result,
                         Xapian::doccount collapse_max, MSetCmp mcmp,
                         Xapian::doccount& old_item)
{
    if (items.size() < collapse_max) {
        return ADD;
    }

    // We already have collapse_max items better than result so we need to
    // eliminate the lowest ranked.
    if (collapse_count == 0 && collapse_max != 1) {
        // Be lazy about building the heap - if we see <= collapse_max items
        // with a particular collapse key, we never need to use the heap.
        Heap::make(items.begin(), items.end(),
                   [&](pair<Xapian::doccount, Xapian::docid> a,
                       pair<Xapian::doccount, Xapian::docid> b) {
                       return mcmp(results[a.first], results[b.first]);
                   });
    }
    ++collapse_count;

    Xapian::doccount old_item_candidate = items.front().first;
    const Result& old_result = results[old_item_candidate];
    if (items.front().second != old_result.get_docid()) {
        // The previous result with this collapse key we were going to replace
        // has been pushed out of the protomset by higher ranking results.
        //
        // Awkwardly that means we don't know its exact weight, but we
        // only need next_best_weight to know if we can zero collapse_count
        // when there's a percentage cut-off, so setting it to an overestimate
        // is OK, and it's not critically important as it makes no difference
        // at all unless there's a percentage cut-off set.
        //
        // For now use the new result's weight.  FIXME: The lowest weight in
        // the current MSet would be a tighter bound if we're sorting primarily
        // by weight (if we're not, then is next_best_weight useful at all?)
        // We could also check other entries in items to find an upper bound
        // weight.
        next_best_weight = result.get_weight();

        return REPLACE;
    }

    if (mcmp(old_result, result)) {
        // If this is the "best runner-up", update next_best_weight.
        if (result.get_weight() > next_best_weight)
            next_best_weight = result.get_weight();
        return REJECT;
    }

    old_item = old_item_candidate;
    next_best_weight = old_result.get_weight();

    items.front() = { old_item, result.get_docid() };
    Heap::replace(items.begin(), items.end(),
                  [&](pair<Xapian::doccount, Xapian::docid> a,
                      pair<Xapian::doccount, Xapian::docid> b) {
                      return mcmp(results[a.first], results[b.first]);
                  });

    return REPLACE;
}

void
CollapseData::set_item(Xapian::doccount item)
{
    AssertEq(items.size(), 1);
    AssertEq(items[0].first, 0);
    items[0].first = item;
}

void
CollapseData::add_item(const vector<Result>& results,
                       Xapian::doccount item,
                       Xapian::doccount collapse_max,
                       MSetCmp mcmp)
{
    const Result& result = results[item];
    if (items.size() < collapse_max) {
        items.emplace_back(item, result.get_docid());
        return;
    }

    items.front() = { item, result.get_docid() };
    Heap::replace(items.begin(), items.end(),
                  [&](pair<Xapian::doccount, Xapian::docid> a,
                      pair<Xapian::doccount, Xapian::docid> b) {
                      return mcmp(results[a.first], results[b.first]);
                  });
}

collapse_result
Collapser::check(Result& result,
                 Xapian::Document::Internal& vsdoc)
{
    ptr = NULL;
    ++docs_considered;
    result.set_collapse_key(vsdoc.get_value(slot));

    if (result.get_collapse_key().empty()) {
        // We don't collapse results with an empty collapse key.
        ++no_collapse_key;
        return EMPTY;
    }

    // Use dummy value 0 for item - if process() is called, this will get
    // updated to the appropriate value, and if it isn't then the docid won't
    // match and we'll know the item isn't in the current proto-mset.
    auto r = table.emplace(result.get_collapse_key(),
                           CollapseData(0, result.get_docid()));
    ptr = &r.first->second;
    if (r.second) {
        // We've not seen this collapse key before.
        ++entry_count;
        return NEW;
    }

    collapse_result res;
    CollapseData& collapse_data = *ptr;
    res = collapse_data.check_item(results, result, collapse_max, mcmp,
                                   old_item);
    if (res == ADD) {
        ++entry_count;
    } else if (res == REJECT || res == REPLACE) {
        ++dups_ignored;
    }
    return res;
}

void
Collapser::process(collapse_result action,
                   Xapian::doccount item)
{
    switch (action) {
        case NEW:
            // We've not seen this collapse key before.
            Assert(ptr);
            ptr->set_item(item);
            return;
        case ADD: {
            Assert(ptr);
            ptr->add_item(results, item, collapse_max, mcmp);
            break;
        }
        default:
            // Shouldn't be called for other actions.
            Assert(false);
    }
}

Xapian::doccount
Collapser::get_collapse_count(const string & collapse_key,
                              int percent_threshold,
                              double min_weight) const
{
    auto key = table.find(collapse_key);
    // If a collapse key is present in the MSet, it must be in our table.
    Assert(key != table.end());

    if (!percent_threshold) {
        // The recorded collapse_count is correct.
        return key->second.get_collapse_count();
    }

    if (key->second.get_next_best_weight() < min_weight) {
        // We know for certain that all collapsed items would have failed the
        // percentage cutoff, so collapse_count should be 0.
        return 0;
    }

    // We know that the highest weighted collapsed item would have survived the
    // percentage cutoff, but it's possible all other collapsed items would
    // have failed it.  Since collapse_count is a lower bound, we must set it
    // to 1.
    return 1;
}

Xapian::doccount
Collapser::get_matches_lower_bound() const
{
    // We've seen this many matches, but all other documents matching the query
    // could be collapsed onto values already seen.
    Xapian::doccount matches_lower_bound = no_collapse_key + entry_count;
    return matches_lower_bound;
    // FIXME: *Unless* we haven't achieved collapse_max occurrences of *any*
    // collapse key value, so we can increase matches_lower_bound like the
    // code below, but it isn't quite that simple - there may not be that
    // many documents.
#if 0
    Xapian::doccount max_kept = 0;
    for (auto i = table.begin(); i != table.end(); ++i) {
        if (i->second.get_collapse_count() > max_kept) {
            max_kept = i->second.get_collapse_count();
            if (max_kept == collapse_max) {
                return matches_lower_bound;
            }
        }
    }
    return matches_lower_bound + (collapse_max - max_kept);
#endif
}

void
Collapser::finalise(double min_weight, int percent_threshold)
{
    if (entry_count == 0 || 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 string& key = result.get_collapse_key();
        if (key.empty())
            continue;

        // Fill in collapse_count.
        result.set_collapse_count(get_collapse_count(key, percent_threshold,
                                                     min_weight));
        if (--todo == 0) {
            // Terminate early if we've handled all non-empty entries.
            break;
        }
    }
}
