/** @file
 * @brief ProtoMSet class
 */
/* Copyright (C) 2004-2026 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_PROTOMSET_H
#define XAPIAN_INCLUDED_PROTOMSET_H

#include "api/enquireinternal.h"
#include "api/result.h"
#include "api/smallvector.h"
#include "collapser.h"
#include "heap.h"
#include "matchtimeout.h"
#include "msetcmp.h"
#include "omassert.h"
#include "spymaster.h"

#include <algorithm>

using Xapian::Internal::intrusive_ptr;

class ProtoMSet {
    /// Adapt MSetCmp to be usable with min_heap.
    class MCmpAdaptor {
	ProtoMSet* protomset;

      public:
	explicit MCmpAdaptor(ProtoMSet* protomset_) : protomset(protomset_) {}

	bool operator()(Xapian::doccount a, Xapian::doccount b) const {
	    return protomset->mcmp(protomset->results[a],
				   protomset->results[b]);
	}
    };

    friend class MCmpAdaptor;

    /** Maximum size the ProtoMSet needs to grow to.
     *
     *  This is the maximum rank we care about.
     */
    Xapian::doccount max_size;

    Xapian::doccount check_at_least;

    Xapian::Enquire::Internal::sort_setting sort_by;

    MSetCmp mcmp;

    /** Minimum threshold on the weight.
     *
     *  If the primary result ordering is by decreasing relevance (i.e. @a
     *  sort_by is REL or REL_VAL) then once the min_heap kicks in this
     *  threshold is raised to the lowest weight in the proto-mset.
     *
     *  Enquire::set_cutoff() can also affect min_weight - an absolute
     *  threshold determines the initial value; a percentage threshold raises
     *  the threshold each time max_weight increases (unless it's already
     *  higher than the value the percentage threshold results in).
     */
    double min_weight = 0.0;

    /** The highest document weight seen.
     *
     * This weight may not actually be present in @a results if we're not
     * sorting primarily by relevance, or if min_weight > max_weight.
     */
    double max_weight = 0.0;

    bool min_weight_pending = false;

    /** Count of how many known matching documents have been processed so far.
     *
     *  Used to implement "check_at_least".
     */
    Xapian::doccount known_matching_docs = 0;

    /// The items in the proto-MSet.
    std::vector<Result> results;

    /** A heap of offsets into @a results.
     *
     *  Created lazily once we actually need it.
     */
    std::vector<Xapian::doccount> min_heap;

    /// First entry wanted in MSet.
    Xapian::doccount first;

    /** How many weighted leaf subqueries there are.
     *
     *  Used for scaling percentages when the highest weighted document doesn't
     *  "match all terms".
     */
    Xapian::termcount total_subqs;

    /// The number of subqueries which matched to give max_weight.
    Xapian::termcount max_weight_subqs_matched = 0;

    int percent_threshold;

    double percent_threshold_factor;

    double percent_scale = 0.0;

    PostListTree& pltree;

    Collapser collapser;

    double max_possible;

    bool stop_once_full;

    TimeOut timeout;

    Xapian::doccount size() const { return Xapian::doccount(results.size()); }

  public:
    ProtoMSet(Xapian::doccount first_,
	      Xapian::doccount max_items,
	      Xapian::doccount check_at_least_,
	      MSetCmp mcmp_,
	      Xapian::Enquire::Internal::sort_setting sort_by_,
	      Xapian::termcount total_subqs_,
	      PostListTree& pltree_,
	      Xapian::valueno collapse_key,
	      Xapian::doccount collapse_max,
	      int percent_threshold_,
	      double percent_threshold_factor_,
	      double max_possible_,
	      bool stop_once_full_,
	      double time_limit)
        : check_at_least(check_at_least_),
	  sort_by(sort_by_),
	  mcmp(mcmp_),
	  first(first_),
	  total_subqs(total_subqs_),
	  percent_threshold(percent_threshold_),
	  percent_threshold_factor(percent_threshold_factor_),
	  pltree(pltree_),
	  collapser(collapse_key, collapse_max, results, mcmp),
	  max_possible(max_possible_),
	  stop_once_full(stop_once_full_),
	  timeout(time_limit)
    {
        // Always track at least one result so the matcher can rely on
        // being able to look at it to see the best match so far.
        max_size = std::max(first_ + max_items, Xapian::doccount{1});
	results.reserve(max_size);
    }

    ProtoMSet(const ProtoMSet&) = delete;

    ProtoMSet& operator=(const ProtoMSet&) = delete;

    Collapser& get_collapser() { return collapser; }

    bool full() const { return size() == max_size; }

    double get_min_weight() const { return min_weight; }

    void update_max_weight(double weight) {
	if (weight <= max_weight)
	    return;

	max_weight = weight;
	max_weight_subqs_matched = pltree.count_matching_subqs();
	if (percent_threshold) {
	    set_new_min_weight(weight * percent_threshold_factor);
	}
    }

    bool checked_enough() {
	if (known_matching_docs >= check_at_least) {
	    return true;
	}
	if (known_matching_docs >= max_size && timeout.timed_out()) {
	    check_at_least = max_size;
	    return true;
	}
	return false;
    }

    /** Resolve a pending min_weight change.
     *
     *  Only called when there's a percentage weight cut-off.
     */
    bool handle_min_weight_pending(bool finalising = false) {
	// min_weight_pending shouldn't get set when unweighted.
	Assert(sort_by != Xapian::Enquire::Internal::DOCID);
	min_weight_pending = false;
	bool weight_first = (sort_by == Xapian::Enquire::Internal::REL ||
			     sort_by == Xapian::Enquire::Internal::REL_VAL);
	double new_min_weight = HUGE_VAL;
	Xapian::doccount j = 0;
	Xapian::doccount min_elt = 0;
	for (Xapian::doccount i = 0; i != size(); ++i) {
	    if (results[i].get_weight() < min_weight) {
		continue;
	    }
	    if (i != j) {
		results[j] = std::move(results[i]);
		if (collapser) {
		    collapser.result_has_moved(i, j);
		}
	    }
	    if (weight_first && results[j].get_weight() < new_min_weight) {
		new_min_weight = results[j].get_weight();
		min_elt = j;
	    }
	    ++j;
	}
	if (weight_first) {
	    if (finalising) {
		if (known_matching_docs >= check_at_least)
		    min_weight = new_min_weight;
	    } else {
		if (checked_enough())
		    min_weight = new_min_weight;
	    }
	}
	if (j != size()) {
	    results.erase(results.begin() + j, results.end());
	    if (!finalising) {
		return false;
	    }
	}
	if (!finalising && min_elt != 0 && !collapser) {
	    // Install the correct element at the tip of the heap, so
	    // that Heap::make() has less to do.  NB Breaks collapsing.
	    std::swap(results[0], results[min_elt]);
	}
	return true;
    }

    bool early_reject(Result& new_item,
		      bool calculated_weight,
		      SpyMaster& spymaster,
		      const Xapian::Document& doc) {
	if (min_heap.empty())
	    return false;

	// We're sorting by value (in part at least), so compare the item
	// against the lowest currently in the proto-mset.  If sort_by is VAL,
	// then new_item.get_weight() won't be set yet, but that doesn't matter
	// since it's not used by the sort function.
	Xapian::doccount worst_idx = min_heap.front();
	if (mcmp(new_item, results[worst_idx]))
	    return false;

	// The candidate isn't good enough to make the proto-mset, but there
	// are still things we may need to do with it.
	if (!collapser) {
	    // We're not collapsing so we can perform an early reject.
	    ++known_matching_docs;
	    double weight =
		calculated_weight ? new_item.get_weight() : pltree.get_weight();
	    spymaster(doc, weight);
	    update_max_weight(weight);
	    return true;
	}

	// We're collapsing - the question is should we increment
	// known_matching_docs?

	if (checked_enough()) {
	    // We are collapsing but known_matching_docs has already reached
	    // check_at_least so we don't need to worry about whether we can
	    // increment it further.
	    double weight =
		calculated_weight ? new_item.get_weight() : pltree.get_weight();
	    update_max_weight(weight);
	    return true;
	}

	// We can't early reject but need to continue on and check if this item
	// would be collapsed or not (and if not ProtoMSet::add() will get
	// called and known_matching_docs incremented there.
	return false;
    }

    /** Process new_item.
     *
     *  Conceptually this is "add new_item", but taking into account
     *  collapsing.
     */
    bool process(Result&& new_item,
		 ValueStreamDocument& vsdoc) {
	update_max_weight(new_item.get_weight());

	if (!collapser) {
	    // No collapsing, so just add the item.
	    add(std::move(new_item));
	} else {
	    auto res = collapser.check(new_item, vsdoc);
	    switch (res) {
		case REJECT:
		    return true;

		case REPLACE:
		    // There was a previous item in the collapse tab so the
		    // MSet can't be empty.
		    Assert(!results.empty());

		    // This is one of the best collapse_max potential MSet
		    // entries with this key which we've seen so far.  The
		    // entry with this key which it displaced is still in the
		    // proto-MSet so replace it.
		    replace(collapser.old_item, std::move(new_item));
		    return true;

		default:
		    break;
	    }

	    auto elt = add(std::move(new_item));
	    if (res != EMPTY && elt != Xapian::doccount(-1)) {
		collapser.process(res, elt);
	    }
	}

	if (stop_once_full) {
	    if (full() && checked_enough()) {
		return false;
	    }
	}

	return true;
    }

    // Returns the new item's index, or Xapian::doccount(-1) if not added.
    Xapian::doccount add(Result&& item) {
	++known_matching_docs;

	if (item.get_weight() < min_weight) {
	    return Xapian::doccount(-1);
	}

	if (item.get_weight() > max_weight) {
	    update_max_weight(item.get_weight());
	}

	if (!full()) {
	    // We're still filling, or just about to become full.
	    results.push_back(std::move(item));
	    Assert(min_heap.empty());
	    return size() - 1;
	}

	if (min_heap.empty()) {
	    // This breaks if we're collapsing because it moves elements around
	    // but can be used if we aren't (and could be for elements with
	    // no collapse key too - FIXME).
	    if (min_weight_pending) {
		if (!handle_min_weight_pending()) {
		    results.push_back(std::move(item));
		    return size() - 1;
		}
	    }

	    if (size() == 0) {
		// E.g. get_mset(0, 0, 10);
		return Xapian::doccount(-1);
	    }
	    min_heap.reserve(size());
	    for (Xapian::doccount i = 0; i != size(); ++i)
		min_heap.push_back(i);
	    Heap::make(min_heap.begin(), min_heap.end(), MCmpAdaptor(this));
	    if (sort_by == Xapian::Enquire::Internal::REL ||
		sort_by == Xapian::Enquire::Internal::REL_VAL) {
		if (checked_enough()) {
		    min_weight = results[min_heap.front()].get_weight();
		}
	    }
	}

	Xapian::doccount worst_idx = min_heap.front();
	if (!mcmp(item, results[worst_idx])) {
	    // The new item is less than what we already had.
	    return Xapian::doccount(-1);
	}

	results[worst_idx] = std::move(item);
	Heap::replace(min_heap.begin(), min_heap.end(), MCmpAdaptor(this));
	if (sort_by == Xapian::Enquire::Internal::REL ||
	    sort_by == Xapian::Enquire::Internal::REL_VAL) {
	    if (checked_enough()) {
		min_weight = results[min_heap.front()].get_weight();
	    }
	}
	return worst_idx;
    }

    void replace(Xapian::doccount old_item, Result&& b) {
	results[old_item] = std::move(b);
	if (min_heap.empty())
	    return;

	// We need to find the entry in min_heap corresponding to old_item.
	// The simplest way is just to linear-scan for it, and that's actually
	// fairly efficient as we're just searching for an integer in a
	// vector of integers.  The heap structure means that the lowest ranked
	// entry is first and lower ranked entries will tend to be nearer the
	// start, so intuitively scanning forwards for an entry which we're
	// removing because we found a higher ranking one seems sensible, but
	// I've not actually profiled this.
	auto it = std::find(min_heap.begin(), min_heap.end(), old_item);
	if (rare(it == min_heap.end())) {
	    // min_heap should contain all indices of results.
	    Assert(false);
	    return;
	}

	// siftdown() here is correct (because it's on a min-heap).
	Heap::siftdown(min_heap.begin(), min_heap.end(), it, MCmpAdaptor(this));
    }

    void set_new_min_weight(double min_wt) {
	if (min_wt <= min_weight)
	    return;

	min_weight = min_wt;

	if (results.empty()) {
	    // This method gets called before we start matching to set the
	    // fixed weight_threshold threshold.
	    return;
	}

#if 0
	// FIXME: Is this possible?  set_new_min_weight() from a percentage
	// threshold can't do this...
	if (min_wt > max_weight) {
	    // The new threshold invalidates all current entries.
	    results.resize(0);
	    min_heap.resize(0);
	    return;
	}
#endif

	if (!min_heap.empty()) {
	    // If sorting primarily by weight, we could pop the heap while the
	    // tip's weight is < min_wt, but each pop needs 2*log(n)
	    // comparisons, and then pushing replacements for each of those
	    // items needs log(n) comparisons.
	    //
	    // Instead we just discard the heap - if we need to rebuild it,
	    // that'll require 3*n comparisons.  The break even is about 3
	    // discarded items for n=10 or about 5 for n=100, but we may never
	    // need to rebuild the heap.
	    min_heap.clear();
	}

	// Note that we need to check items against min_weight at some point.
	min_weight_pending = true;
    }

    void finalise_percentages() {
	if (results.empty() || max_weight == 0.0)
	    return;

	percent_scale = max_weight_subqs_matched / double(total_subqs);
	percent_scale /= max_weight;
	Assert(percent_scale > 0);
	if (!percent_threshold) {
	    return;
	}

	// Truncate the results if necessary.
	set_new_min_weight(percent_threshold_factor / percent_scale);
	if (min_weight_pending) {
	    handle_min_weight_pending(true);
	}
    }

    Xapian::MSet
    finalise(const Xapian::MatchDecider* mdecider,
	     const std::vector<std::unique_ptr<LocalSubMatch>>& locals,
             const Xapian::VecUniquePtr<EstimateOp>& estimates,
             Xapian::doccount max_items) {
	finalise_percentages();

	Xapian::doccount matches_lower_bound;
	Xapian::doccount matches_estimated;
	Xapian::doccount matches_upper_bound;
	Xapian::doccount uncollapsed_lower_bound;
	Xapian::doccount uncollapsed_estimated;
	Xapian::doccount uncollapsed_upper_bound;

	if (!collapser && (!full() || known_matching_docs < check_at_least)) {
	    // Under these conditions we know exactly how many matching docs
	    // there are for the full match so we don't need to resolve the
	    // EstimateOp stack.
	    Xapian::doccount m;
	    if (!full()) {
		// We didn't get all the results requested, so we know that
		// we've got all there are, and the bounds and estimate are
		// all equal to that number.
		m = size();
		// And that should equal known_matching_docs, unless a percentage
		// threshold caused some matches to be excluded.
		if (!percent_threshold) {
		    AssertEq(m, known_matching_docs);
		} else {
		    AssertRel(m, <=, known_matching_docs);
		}
	    } else {
		// Otherwise we didn't reach check_at_least, so
		// known_matching_docs gives the exact size.
		m = known_matching_docs;
	    }

	    matches_lower_bound = matches_estimated = matches_upper_bound = m;

	    // When not collapsing the uncollapsed bounds are just the same.
	    uncollapsed_lower_bound = matches_lower_bound;
	    uncollapsed_estimated = matches_estimated;
	    uncollapsed_upper_bound = matches_upper_bound;
	} else {
	    matches_lower_bound = 0;
	    matches_estimated = 0;
	    matches_upper_bound = 0;
	    for (size_t i = 0; i != estimates.size(); ++i) {
		if (estimates[i]) {
		    Assert(locals[i].get());
		    Estimates e = locals[i]->resolve(estimates[i]);
		    matches_lower_bound += e.min;
		    matches_estimated += e.est;
		    matches_upper_bound += e.max;
		}
	    }

	    AssertRel(matches_estimated, >=, matches_lower_bound);
	    AssertRel(matches_estimated, <=, matches_upper_bound);

	    uncollapsed_lower_bound = matches_lower_bound;
	    uncollapsed_estimated = matches_estimated;
	    uncollapsed_upper_bound = matches_upper_bound;

	    if (!full()) {
		// We didn't get all the results requested, so we know that we've
		// got all there are, and the bounds and estimate are all equal to
		// that number.
		matches_lower_bound = size();
		matches_estimated = matches_lower_bound;
		matches_upper_bound = matches_lower_bound;

		// And that should equal known_matching_docs, unless a percentage
		// threshold caused some matches to be excluded.
		if (!percent_threshold) {
		    AssertEq(matches_estimated, known_matching_docs);
		} else {
		    AssertRel(matches_estimated, <=, known_matching_docs);
		}

		if (matches_lower_bound > uncollapsed_lower_bound) {
		    // Clamp the uncollapsed bound to be at least the collapsed
		    // one.
		    uncollapsed_lower_bound = matches_lower_bound;
		}
	    } else {
		// We can end up scaling the estimate more than once, so collect
		// the scale factors and apply them in one go to avoid rounding
		// more than once.
		double estimate_scale = 1.0;
		double unique_rate = 1.0;

		if (collapser) {
		    matches_lower_bound = collapser.get_matches_lower_bound();

		    Xapian::doccount docs_considered =
			collapser.get_docs_considered();
		    Xapian::doccount dups_ignored = collapser.get_dups_ignored();
		    if (docs_considered > 0) {
			// Scale the estimate by the rate at which we've been
			// finding unique documents.
			double unique = double(docs_considered - dups_ignored);
			unique_rate = unique / double(docs_considered);
		    }

		    // We can safely reduce the upper bound by the number of
		    // duplicates we've ignored.
		    matches_upper_bound -= dups_ignored;
		}

		if (mdecider) {
		    if (!percent_threshold && !collapser) {
			if (known_matching_docs > matches_lower_bound) {
			    // We're not collapsing or doing a percentage
			    // threshold, so known_matching_docs is a lower bound
			    // on the total number of matches.
			    matches_lower_bound = known_matching_docs;
			}
		    }
		}

		if (percent_threshold) {
		    // Scale the estimate assuming that document weights are evenly
		    // distributed from 0 to the maximum weight seen.
		    estimate_scale *= (1.0 - percent_threshold_factor);

		    // This is all we can be sure of without additional work.
		    matches_lower_bound = size();

		    if (collapser) {
			uncollapsed_lower_bound = matches_lower_bound;
		    }
		}

		if (collapser && estimate_scale != 1.0) {
		    uncollapsed_estimated =
			Xapian::doccount(uncollapsed_estimated * estimate_scale +
					 0.5);
		}

		estimate_scale *= unique_rate;

		if (estimate_scale != 1.0) {
		    matches_estimated =
			Xapian::doccount(matches_estimated * estimate_scale + 0.5);
		    if (matches_estimated < matches_lower_bound)
			matches_estimated = matches_lower_bound;
		}

		if (collapser || mdecider) {
		    // Clamp the estimate to the range given by the bounds.
		    AssertRel(matches_lower_bound, <=, matches_upper_bound);
		    matches_estimated = std::clamp(matches_estimated,
						   matches_lower_bound,
						   matches_upper_bound);
		} else if (!percent_threshold) {
		    AssertRel(known_matching_docs, <=, matches_upper_bound);
		    if (known_matching_docs > matches_lower_bound)
			matches_lower_bound = known_matching_docs;
		    if (known_matching_docs > matches_estimated)
			matches_estimated = known_matching_docs;
		}

		if (collapser) {
		    if (!mdecider && !percent_threshold) {
			AssertRel(known_matching_docs, <=, uncollapsed_upper_bound);
			if (known_matching_docs > uncollapsed_lower_bound)
			    uncollapsed_lower_bound = known_matching_docs;
		    }

		    if (matches_lower_bound > uncollapsed_lower_bound) {
			// Clamp the uncollapsed bound to be at least the collapsed
			// one.
			uncollapsed_lower_bound = matches_lower_bound;
		    }

		    // Clamp the estimate to lie within the known bounds.
		    if (uncollapsed_estimated < uncollapsed_lower_bound) {
			uncollapsed_estimated = uncollapsed_lower_bound;
		    } else if (uncollapsed_estimated > uncollapsed_upper_bound) {
			uncollapsed_estimated = uncollapsed_upper_bound;
		    }
		} else {
		    // When not collapsing the uncollapsed bounds are just the same.
		    uncollapsed_lower_bound = matches_lower_bound;
		    uncollapsed_estimated = matches_estimated;
		    uncollapsed_upper_bound = matches_upper_bound;
		}
	    }
	}

	// FIXME: Profile using min_heap here (when it's been created) to
	// handle "first" and perform the sort.
	if (first != 0) {
	    if (first > size()) {
		results.clear();
                goto no_results_to_sort;
	    }
            // We perform nth_element() on reverse iterators so that the
            // unwanted elements end up at the end of items, which means
            // that the call to erase() to remove them doesn't have to copy
            // any elements.
            auto nth = results.rbegin() + first;
            std::nth_element(results.rbegin(), nth, results.rend(), mcmp);
            // Discard the unwanted elements.
            results.erase(results.end() - first, results.end());
        } else if (max_items == 0) {
            results.clear();
            goto no_results_to_sort;
	}

	std::sort(results.begin(), results.end(), mcmp);

        // Note: finalise() is a no-op if results.empty() so no_results_to_sort
        // can safely skip over this too.
	collapser.finalise(min_weight, percent_threshold);

no_results_to_sort:

	// The estimates should lie between the bounds.
	AssertRel(matches_lower_bound, <=, matches_estimated);
	AssertRel(matches_estimated, <=, matches_upper_bound);
	AssertRel(uncollapsed_lower_bound, <=, uncollapsed_estimated);
	AssertRel(uncollapsed_estimated, <=, uncollapsed_upper_bound);

	// Collapsing should only reduce the bounds and estimate.
	AssertRel(matches_lower_bound, <=, uncollapsed_lower_bound);
	AssertRel(matches_estimated, <=, uncollapsed_estimated);
	AssertRel(matches_upper_bound, <=, uncollapsed_upper_bound);

	return Xapian::MSet(new Xapian::MSet::Internal(first,
						       matches_upper_bound,
						       matches_lower_bound,
						       matches_estimated,
						       uncollapsed_upper_bound,
						       uncollapsed_lower_bound,
						       uncollapsed_estimated,
						       max_possible,
						       max_weight,
						       std::move(results),
						       percent_scale * 100.0));
    }
};

#endif // XAPIAN_INCLUDED_PROTOMSET_H
