/** @file
 * @brief Xapian::MSet class
 */
/* Copyright (C) 2017,2024,2025,2026 Olly Betts
 * Copyright (C) 2018 Uppinder Chugh
 *
 * 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 "msetinternal.h"
#include "xapian/mset.h"

// FIXME: Clustering API needs work: #include "xapian/cluster.h"

#include "backends/documentinternal.h"
#include "clamp_cast.h"
#include "net/serialise.h"
#include "matcher/msetcmp.h"
#include "omassert.h"
#include "pack.h"
#include "roundestimate.h"
#include "serialise-double.h"
#include "str.h"
#include "unicode/description_append.h"

#include <algorithm>
#include <cfloat>
#include <string>
#include <string_view>
#include <unordered_set>

using namespace std;

namespace Xapian {

MSet::MSet(const MSet&) = default;

MSet&
MSet::operator=(const MSet&) = default;

MSet::MSet(MSet&&) = default;

MSet&
MSet::operator=(MSet&&) = default;

MSet::MSet() : internal(new MSet::Internal) {}

MSet::MSet(Internal* internal_) : internal(internal_) {}

MSet::~MSet() {}

void
MSet::fetch_(Xapian::doccount first, Xapian::doccount last) const
{
    internal->fetch(first, last);
}

void
MSet::set_item_weight(Xapian::doccount i, double weight)
{
    internal->set_item_weight(i, weight);
}

#if 0 // FIXME: Diversification API needs work.
/** Evaluate a diversified mset
 *
 *  Evaluate a diversified mset using MPT algorithm
 *
 *  @param dmset        Set of points representing candidate diversifed set of
 *                      documents.
 *  @param cset         Set of clusters of given MSet.
 */
static double
evaluate_dmset(const vector<Xapian::docid>& dmset,
               const Xapian::ClusterSet& cset,
               double factor1,
               double factor2,
               const Xapian::MSet& mset,
               const vector<double>& dissimilarity)
{
    double score_1 = 0, score_2 = 0;

    // FIXME: We could compute score_1 once then adjust for each candidate
    // change.
    // Seems hard to do similar for score_2 though.
    for (auto mset_index : dmset)
        score_1 += mset[mset_index].get_weight();

    auto cset_size = cset.size();
    for (Xapian::doccount c = 0; c < cset_size; ++c) {
        double min_dist = numeric_limits<double>::max();
        unsigned int pos = 1;
        for (auto mset_index : dmset) {
            // FIXME: Pre-compute 1.0 / log(2.0 + i) for i = [0, dmset.size()) ?
            double weight = dissimilarity[mset_index * cset_size + c];
            weight /= log(1.0 + pos);
            min_dist = min(min_dist, weight);
            ++pos;
        }
        score_2 += min_dist;
    }

    return factor2 * score_2 - factor1 * score_1;
}

void
MSet::diversify_(Xapian::doccount k,
                 Xapian::doccount r,
                 double factor1,
                 double factor2)
{
    // Ensured by inlined caller.
    AssertRel(k, >=, 2);

    auto mset_size = size();
    if (mset_size <= k) {
        // Picking k documents would pick the whole MSet so nothing to do.
        //
        // Since k >= 2, this means we don't try to diversify an MSet with
        // 2 documents (for which reordering can't usefully improve diversity
        // since the only possible change is to swap the order of the 2
        // documents).
        return;
    }

    /// Store MSet indices of top k diversified documents
    std::vector<Xapian::doccount> main_dmset;
    main_dmset.reserve(k);

    Xapian::doccount count = 0;
    TermListGroup tlg(*this);
    std::vector<Xapian::Point> points;
    points.reserve(mset_size);
    for (MSetIterator it = begin(); it != end(); ++it) {
        Xapian::Document doc = it.get_document();
        doc.internal->set_index(count);
        points.push_back(Xapian::Point(tlg, doc));
        // Initial top-k diversified documents
        if (count < k) {
            // The initial diversified document set is the top-k documents from
            // the MSet.
            main_dmset.push_back(count);
        }
        ++count;
    }

    // Cluster the MSet into k clusters.
    Xapian::ClusterSet cset = Xapian::LCDClusterer(k).cluster(*this);

    /** Precompute dissimilarity scores between each document and cluster
     *  centroid.
     *
     *  These scores are:
     *
     *  1.0 - cosine_similarity(docid, cluster_index)
     *
     *  The index into dissimilarity is:
     *
     *    mset_index * number_of_clusters + cluster_index
     */

    // Pre-compute all the dissimilarity values.
    auto cset_size = cset.size();
    std::vector<double> dissimilarity;
    dissimilarity.reserve(cset_size * points.size());
    {
        Xapian::CosineDistance d;
        for (const auto& point : points) {
            for (unsigned int c = 0; c < cset_size; ++c) {
                double dist = d.similarity(point, cset[c].get_centroid());
                dissimilarity.push_back(1.0 - dist);
            }
        }
    }

    // Build topc, which contains the union of the top-r relevant documents of
    // each cluster.
    vector<Xapian::docid> topc;
    for (Xapian::doccount c = 0; c < cset_size; ++c) {
        // FIXME: This is supposed to pick the `r` most relevant documents, but
        // actually seems to pick those with the lowest docids.
        auto documents = cset[c].get_documents();
        auto limit = std::min(r, documents.size());
        for (Xapian::doccount d = 0; d < limit; ++d) {
            auto mset_index = documents[d].internal->get_index();
            topc.push_back(mset_index);
        }
    }

    vector<Xapian::doccount> curr_dmset = main_dmset;

    while (true) {
        bool found_better_dmset = false;
        for (unsigned int i = 0; i < main_dmset.size(); ++i) {
            auto curr_doc = main_dmset[i];
            double best_score = evaluate_dmset(curr_dmset, cset,
                                               factor1, factor2,
                                               *this, dissimilarity);
            bool found_better_doc = false;

            for (unsigned int j = 0; j < topc.size(); ++j) {
                // Continue if candidate document from topc already
                // exists in curr_dmset.  FIXME: Linear search!
                auto candidate_doc = find(curr_dmset.begin(), curr_dmset.end(),
                                          topc[j]);
                if (candidate_doc != curr_dmset.end()) {
                    continue;
                }

                auto temp_doc = curr_dmset[i];
                curr_dmset[i] = topc[j];
                double score = evaluate_dmset(curr_dmset, cset,
                                              factor1, factor2,
                                              *this, dissimilarity);

                if (score < best_score) {
                    curr_doc = curr_dmset[i];
                    best_score = score;
                    found_better_doc = true;
                }

                curr_dmset[i] = temp_doc;
            }
            if (found_better_doc) {
                curr_dmset[i] = curr_doc;
                found_better_dmset = true;
            }
        }

        // Terminate algorithm when there's no change in current
        // document matchset
        if (!found_better_dmset)
            break;

        main_dmset = curr_dmset;
    }

    // Reorder the results to reflect the diversification.  To do this we need
    // to partition the MSet so the promoted documents come first (in original
    // MSet order), followed by the non-promoted documents (also in original
    // MSet order).
    unordered_set<Xapian::docid> promoted{k};
    for (auto mset_index : main_dmset) {
        promoted.insert(internal->items[mset_index].get_docid());
    }

    stable_partition(internal->items.begin(), internal->items.end(),
                     [&](const Result& result) {
                         return promoted.count(result.get_docid());
                     });
}
#endif

void
MSet::sort_by_relevance()
{
    std::sort(internal->items.begin(), internal->items.end(),
              get_msetcmp_function(Enquire::Internal::REL, true, false));
}

int
MSet::convert_to_percent(double weight) const
{
    return internal->convert_to_percent(weight);
}

Xapian::doccount
MSet::get_termfreq(std::string_view term) const
{
    // Check the cached data for query terms first.
    Xapian::doccount termfreq;
    if (usual(internal->stats && internal->stats->get_stats(term, termfreq))) {
        return termfreq;
    }

    if (rare(!internal->enquire)) {
        // Consistent with get_termfreq() on an empty database which always
        // returns 0.
        return 0;
    }

    // Fall back to asking the database via enquire.
    return internal->enquire->get_termfreq(term);
}

double
MSet::get_termweight(std::string_view term) const
{
    // A term not in the query has no termweight, so 0.0 makes sense as the
    // answer in such cases.
    double weight = 0.0;
    if (usual(internal->stats)) {
        (void)internal->stats->get_termweight(term, weight);
    }
    return weight;
}

Xapian::doccount
MSet::get_firstitem() const
{
    return internal->first;
}

Xapian::doccount
MSet::get_matches_lower_bound() const
{
    return internal->matches_lower_bound;
}

Xapian::doccount
MSet::get_matches_estimated() const
{
    // Doing this here avoids calculating if the estimate is never looked at,
    // though does mean we recalculate if this method is called more than once.
    return round_estimate(internal->matches_lower_bound,
                          internal->matches_upper_bound,
                          internal->matches_estimated);
}

Xapian::doccount
MSet::get_matches_upper_bound() const
{
    return internal->matches_upper_bound;
}

Xapian::doccount
MSet::get_uncollapsed_matches_lower_bound() const
{
    return internal->uncollapsed_lower_bound;
}

Xapian::doccount
MSet::get_uncollapsed_matches_estimated() const
{
    // Doing this here avoids calculating if the estimate is never looked at,
    // though does mean we recalculate if this method is called more than once.
    return round_estimate(internal->uncollapsed_lower_bound,
                          internal->uncollapsed_upper_bound,
                          internal->uncollapsed_estimated);
}

Xapian::doccount
MSet::get_uncollapsed_matches_upper_bound() const
{
    return internal->uncollapsed_upper_bound;
}

double
MSet::get_max_attained() const
{
    return internal->max_attained;
}

double
MSet::get_max_possible() const
{
    return internal->max_possible;
}

Xapian::doccount
MSet::size() const
{
    return clamp_cast<Xapian::doccount>(internal->items.size());
}

std::string
MSet::snippet(std::string_view text,
              size_t length,
              const Xapian::Stem& stemmer,
              unsigned flags,
              std::string_view hi_start,
              std::string_view hi_end,
              std::string_view omit) const
{
    // The actual implementation is in queryparser/termgenerator_internal.cc.
    return internal->snippet(text, length, stemmer, flags,
                             hi_start, hi_end, omit);
}

std::string
MSet::get_description() const
{
    return internal->get_description();
}

Document
MSet::Internal::get_document(Xapian::doccount index) const
{
    if (index >= items.size()) {
        string msg = "Requested index ";
        msg += str(index);
        msg += " in MSet of size ";
        msg += str(items.size());
        throw Xapian::RangeError(msg);
    }
    Assert(enquire);
    return enquire->get_document(items[index].get_docid());
}

void
MSet::Internal::fetch(Xapian::doccount first_, Xapian::doccount last) const
{
    if (items.empty() || !enquire) {
        return;
    }
    if (last > items.size() - 1) {
        last = Xapian::doccount(items.size() - 1);
    }
    if (first_ <= last) {
        Xapian::doccount n = last - first_;
        for (Xapian::doccount i = 0; i <= n; ++i) {
            enquire->request_document(items[i].get_docid());
        }
    }
}

void
MSet::Internal::set_item_weight(Xapian::doccount i, double weight)
{
    // max_attained is updated assuming that set_item_weight is called on every
    // MSet item from 0 up. While assigning new weights max_attained is updated
    // as the maximum of the new weights set till Xapian::doccount i.
    if (i == 0)
        max_attained = weight;
    else
        max_attained = max(max_attained, weight);
    // Ideally the max_possible should be the maximum possible weight that
    // can be assigned by the reranking algorithm, but since it is not always
    // possible to calculate the max possible weight for a reranking algorithm
    // we use this approach.
    max_possible = max(max_possible, max_attained);
    items[i].set_weight(weight);
}

int
MSet::Internal::convert_to_percent(double weight) const
{
    int percent;
    if (percent_scale_factor == 0.0) {
        // For an unweighted search, give all matches 100%.
        percent = 100;
    } else if (weight <= 0.0) {
        // Some weighting schemes can return zero relevance while matching,
        // so give such matches 0%.
        percent = 0;
    } else {
        // Adding on 100 * DBL_EPSILON was a hack to work around excess
        // precision (e.g. on x86 when not using SSE), but this code seems like
        // it's generally asking for problems with floating point rounding
        // issues - maybe we ought to carry through the matching and total
        // number of subqueries and calculate using those instead.
        //
        // There are corresponding hacks in matcher/matcher.cc.
        percent = int(weight * percent_scale_factor + 100.0 * DBL_EPSILON);
        if (percent <= 0) {
            // Make any non-zero weight give a non-zero percentage.
            percent = 1;
        } else if (percent > 100) {
            // Make sure we don't ever exceed 100%.
            percent = 100;
        }
        // FIXME: Ideally we should also make sure any non-exact match gives
        // < 100%.
    }
    return percent;
}

void
MSet::Internal::unshard_docids(Xapian::doccount shard,
                               Xapian::doccount n_shards)
{
    for (auto& result : items) {
        result.unshard_docid(shard, n_shards);
    }
}

void
MSet::Internal::merge_stats(const Internal* o, bool collapsing)
{
    if (snippet_bg_relevance.empty()) {
        snippet_bg_relevance = o->snippet_bg_relevance;
    } else {
        Assert(snippet_bg_relevance == o->snippet_bg_relevance);
    }
    if (collapsing) {
        matches_lower_bound = max(matches_lower_bound, o->matches_lower_bound);
        // matches_estimated will get adjusted later in this case.
    } else {
        matches_lower_bound += o->matches_lower_bound;
    }
    matches_estimated += o->matches_estimated;
    matches_upper_bound += o->matches_upper_bound;
    uncollapsed_lower_bound += o->uncollapsed_lower_bound;
    uncollapsed_estimated += o->uncollapsed_estimated;
    uncollapsed_upper_bound += o->uncollapsed_upper_bound;
    max_possible = max(max_possible, o->max_possible);
    if (o->max_attained > max_attained) {
        max_attained = o->max_attained;
        percent_scale_factor = o->percent_scale_factor;
    }
}

string
MSet::Internal::serialise() const
{
    string result;

    result += serialise_double(max_possible);
    result += serialise_double(max_attained);

    result += serialise_double(percent_scale_factor);

    pack_uint(result, first);
    // Send back the raw matches_* values.  MSet::get_matches_estimated()
    // rounds the estimate lazily, but when we merge MSet objects we really
    // want to merge based on the raw estimates.
    //
    // It is also cleaner that a round-trip through serialisation gives you an
    // object which is as close to the original as possible.
    pack_uint(result, matches_lower_bound);
    pack_uint(result, matches_estimated);
    pack_uint(result, matches_upper_bound);
    pack_uint(result, uncollapsed_lower_bound);
    pack_uint(result, uncollapsed_estimated);
    pack_uint(result, uncollapsed_upper_bound);

    pack_uint(result, items.size());
    for (auto&& item : items) {
        result += serialise_double(item.get_weight());
        pack_uint(result, item.get_docid());
        pack_string(result, item.get_sort_key());
        pack_string(result, item.get_collapse_key());
        pack_uint(result, item.get_collapse_count());
    }

    if (stats)
        result += serialise_stats(*stats);

    return result;
}

void
MSet::Internal::unserialise(const char * p, const char * p_end)
{
    items.clear();

    max_possible = unserialise_double(&p, p_end);
    max_attained = unserialise_double(&p, p_end);

    percent_scale_factor = unserialise_double(&p, p_end);

    size_t msize;
    if (!unpack_uint(&p, p_end, &first) ||
        !unpack_uint(&p, p_end, &matches_lower_bound) ||
        !unpack_uint(&p, p_end, &matches_estimated) ||
        !unpack_uint(&p, p_end, &matches_upper_bound) ||
        !unpack_uint(&p, p_end, &uncollapsed_lower_bound) ||
        !unpack_uint(&p, p_end, &uncollapsed_estimated) ||
        !unpack_uint(&p, p_end, &uncollapsed_upper_bound) ||
        !unpack_uint(&p, p_end, &msize)) {
        unpack_throw_serialisation_error(p);
    }
    for ( ; msize; --msize) {
        double wt = unserialise_double(&p, p_end);
        Xapian::docid did;
        string sort_key, key;
        Xapian::doccount collapse_cnt;
        if (!unpack_uint(&p, p_end, &did) ||
            !unpack_string(&p, p_end, sort_key) ||
            !unpack_string(&p, p_end, key) ||
            !unpack_uint(&p, p_end, &collapse_cnt)) {
            unpack_throw_serialisation_error(p);
        }
        items.emplace_back(wt, did, std::move(key), collapse_cnt,
                           std::move(sort_key));
    }

    if (p != p_end) {
        stats.reset(new Xapian::Weight::Internal());
        unserialise_stats(p, p_end, *stats);
    }
}

string
MSet::Internal::get_description() const
{
    string desc = "MSet(matches_lower_bound=";
    desc += str(matches_lower_bound);
    desc += ", matches_estimated=";
    desc += str(matches_estimated);
    desc += ", matches_upper_bound=";
    desc += str(matches_upper_bound);
    if (uncollapsed_lower_bound != matches_lower_bound) {
        desc += ", uncollapsed_lower_bound=";
        desc += str(uncollapsed_lower_bound);
    }
    if (uncollapsed_estimated != matches_estimated) {
        desc += ", uncollapsed_estimated=";
        desc += str(uncollapsed_estimated);
    }
    if (uncollapsed_upper_bound != matches_upper_bound) {
        desc += ", uncollapsed_upper_bound=";
        desc += str(uncollapsed_upper_bound);
    }
    if (first != 0) {
        desc += ", first=";
        desc += str(first);
    }
    if (max_possible > 0) {
        desc += ", max_possible=";
        desc += str(max_possible);
    }
    if (max_attained > 0) {
        desc += ", max_attained=";
        desc += str(max_attained);
    }
    desc += ", [";
    bool comma = false;
    for (auto&& item : items) {
        if (comma) {
            desc += ", ";
        } else {
            comma = true;
        }
        desc += item.get_description();
    }
    desc += "])";
    return desc;
}

}
