/*
 * SPDX-FileCopyrightText: 2017-2017 CSSlayer <wengxt@gmail.com>
 *
 * SPDX-License-Identifier: LGPL-2.1-or-later
 */

#include "decoder.h"
#include <algorithm>
#include <cassert>
#include <chrono>
#include <cstddef>
#include <functional>
#include <limits>
#include <memory>
#include <queue>
#include <ranges>
#include <string>
#include <string_view>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include <boost/container_hash/hash.hpp>
#include <boost/functional/hash.hpp>
#include <boost/ptr_container/ptr_vector.hpp>
#include <fcitx-utils/macros.h>
#include "languagemodel.h"
#include "lattice.h"
#include "lattice_p.h"
#include "segmentgraph.h"
#include "utils.h"
#include "utils_p.h"

namespace libime {

constexpr int MAX_BACKWARD_SEARCH_SIZE = 10000;

class NBestNode {
public:
    NBestNode(const LatticeNode *node,
              size_t next = std::numeric_limits<size_t>::max())
        : node_(node), next_(next) {}

    const NBestNode *next(const std::vector<NBestNode> &pool) const {
        return next_ == std::numeric_limits<size_t>::max() ? nullptr
                                                           : &pool[next_];
    }

    const LatticeNode *node() const { return node_; }

    // for nbest
    float gn_ = 0.0F;
    float fn_ = -std::numeric_limits<float>::max();

private:
    const LatticeNode *const node_;
    size_t next_ = std::numeric_limits<size_t>::max();
};

struct NBestNodeLess {
    NBestNodeLess(const std::vector<NBestNode> &pool) : pool_(pool) {}
    bool operator()(size_t lhs, size_t rhs) const {
        return pool_[lhs].fn_ < pool_[rhs].fn_;
    }
    const std::vector<NBestNode> &pool_;
};

class DecoderPrivate {
public:
    DecoderPrivate(const Dictionary *dict, const LanguageModelBase *model)
        : dict_(dict), model_(model) {}

    // Try to update lattice based on existing data.
    bool
    buildLattice(const Decoder *q, Lattice &l,
                 const std::unordered_set<const SegmentGraphNode *> &ignore,
                 const State &state, const SegmentGraph &graph,
                 size_t frameSize, void *helper) const;

    void
    forwardSearch(const Decoder *q, const SegmentGraph &graph, Lattice &lattice,
                  const std::unordered_set<const SegmentGraphNode *> &ignore,
                  size_t beamSize) const;
    void backwardSearch(const SegmentGraph &graph, Lattice &l, size_t nbest,
                        float max, float min, size_t beamSize) const;

    const Dictionary *dict_;
    const LanguageModelBase *model_;
};

bool DecoderPrivate::buildLattice(
    const Decoder *q, Lattice &l,
    const std::unordered_set<const SegmentGraphNode *> &ignore,
    const State &state, const SegmentGraph &graph, size_t frameSize,
    void *helper) const {
    LatticeMap &lattice = l.d_ptr->lattice_;

    // Create the root node.
    if (!lattice.contains(&graph.start())) {
        lattice[&graph.start()].push_back(
            q->createLatticeNode(graph, model_, "", model_->beginSentence(),
                                 {nullptr, &graph.start()}, state, 0));
    }

    // std::vector is used here to make sure std::make_heap works.
    std::unordered_map<
        std::pair<const SegmentGraphNode *, const SegmentGraphNode *>,
        std::vector<std::unique_ptr<LatticeNode>>,
        boost::hash<
            std::pair<const SegmentGraphNode *, const SegmentGraphNode *>>>
        frames;

    auto dictMatchCallback = [this, &graph, &frames, q, frameSize](
                                 const SegmentGraphPath &path, WordNode &word,
                                 float adjust,
                                 std::unique_ptr<LatticeNodeData> data) {
        if (InvalidWordIndex == word.idx()) {
            auto idx = model_->index(word.word());
            word.setIdx(idx);
        }
        assert(path.front());
        auto &frame = frames[std::make_pair(path.front(), path.back())];
        const bool applyFrameSize =
            path.front() != &graph.start() && frameSize > 0;
        auto *node = q->createLatticeNode(
            graph, model_, word.word(), word.idx(), path, model_->nullState(),
            adjust, std::move(data), frame.empty());
        if (!node) {
            return;
        }

        frame.emplace_back(node);
        if (!applyFrameSize) {
            return;
        }
        // Make a maximum heap.
        auto scoreGreaterThan = [](const auto &lhs, const auto &rhs) {
            return lhs->score() > rhs->score();
        };
        // Just reach the limit, initialize the heap.
        if (frame.size() == frameSize) {
            for (auto &n : frame) {
                // Cache the score here.
                n->setScore(model_->singleWordScore(n->word()) + n->cost());
            }
            std::make_heap(frame.begin(), frame.end(), scoreGreaterThan);
        } else if (frame.size() == frameSize + 1) {
            // Cache the score here.
            node->setScore(model_->singleWordScore(node->word()) +
                           node->cost());
            // Take a short cut, check if node score greater than minimum
            if (scoreGreaterThan(node, frame[0])) {
                std::push_heap(frame.begin(), frame.end(), scoreGreaterThan);
                std::pop_heap(frame.begin(), frame.end(), scoreGreaterThan);
            }
            frame.pop_back();
        }
    };

    dict_->matchPrefix(graph, dictMatchCallback, ignore, helper);

    for (auto &[path, nodes] : frames) {
        auto &latticeUnit = lattice[path.second];
        for (auto &node : nodes) {
            latticeUnit.push_back(node.release());
        }
    }
    if (!lattice.contains(&graph.end())) {
        return false;
    }

    // Create the node for end.
    lattice[nullptr].push_back(
        q->createLatticeNode(graph, model_, "", model_->endSentence(),
                             {&graph.end(), nullptr}, model_->nullState()));
    return true;
}

void DecoderPrivate::forwardSearch(
    const Decoder *q, const SegmentGraph &graph, Lattice &l,
    const std::unordered_set<const SegmentGraphNode *> &ignore,
    size_t beamSize) const {
    State state;
    LatticeMap &lattice = l.d_ptr->lattice_;
    std::unordered_map<const SegmentGraphNode *,
                       std::tuple<float, LatticeNode *, State>>
        unknownIdCache;
    const auto *start = &graph.start();
    // forward search
    auto updateForNode = [&](const SegmentGraphBase &,
                             const SegmentGraphNode *graphNode) {
        if (graphNode == start || !lattice.contains(graphNode) ||
            ignore.contains(graphNode)) {
            return true;
        }
        auto &latticeNodes = lattice[graphNode];
        for (auto &node : latticeNodes) {
            const auto *from = node.from();
            assert(graph.checkNodeInGraph(from));
            float maxScore = -std::numeric_limits<float>::max();
            LatticeNode *maxNode = nullptr;
            State maxState;
            bool isUnknown = model_->isNodeUnknown(node);
            if (isUnknown) {
                auto iter = unknownIdCache.find(from);
                if (iter != unknownIdCache.end()) {
                    std::tie(maxScore, maxNode, maxState) = iter->second;
                }
            }

            if (!maxNode) {
                auto iter = lattice.find(from);
                // assert(iter != lattice.end());
                if (iter == lattice.end()) {
                    continue;
                }
                auto &searchFrom = iter->second;
                auto searchSize = beamSize;
                if (searchSize) {
                    searchSize = std::min(searchSize, lattice[from].size());
                } else {
                    searchSize = lattice[from].size();
                }
                for (auto &parent : searchFrom | std::views::take(searchSize)) {
                    auto score = parent.score() +
                                 model_->score(parent.state(), node, state);
                    if (score > maxScore) {
                        maxScore = score;
                        maxNode = &parent;
                        maxState = state;
                    }
                }

                if (isUnknown) {
                    unknownIdCache.emplace(
                        std::piecewise_construct, std::forward_as_tuple(from),
                        std::forward_as_tuple(maxScore, maxNode, maxState));
                }
            }

            assert(maxNode);
            node.setScore(maxScore + node.cost());
            node.setPrev(maxNode);
            node.state() = maxState;
        }
        if (q->needSort(graph, graphNode)) {
            latticeNodes.sort(
                [](const LatticeNode &lhs, const LatticeNode &rhs) {
                    return lhs.score() > rhs.score();
                });
        }
        return true;
    };

    graph.bfs(start, updateForNode);
    updateForNode(graph, nullptr);
}

std::string concatNBest(const std::vector<NBestNode> &pool,
                        const NBestNode *node, std::string_view sep = "") {
    std::string result;
    while (node) {
        result.append(node->node()->word());
        result.append(sep.data(), sep.size());
        node = node->next(pool);
    }
    return result;
}

void DecoderPrivate::backwardSearch(const SegmentGraph &graph, Lattice &l,
                                    size_t nbest, float max, float min,
                                    size_t beamSize) const {
    auto &lattice = l.d_ptr->lattice_;
    State state;
    // backward search
    assert(lattice[&graph.start()].size() == 1);
    assert(lattice[nullptr].size() == 1);
    auto *pos = &lattice[nullptr][0];
    l.d_ptr->nbests_.push_back(pos->toSentenceResult());

    struct ModelLookupCacheHash {
        size_t operator()(const std::pair<const LatticeNode *,
                                          const LatticeNode *> &key) const {
            size_t seed = std::hash<const LatticeNode *>()(key.first);
            boost::hash_combine(seed,
                                std::hash<const LatticeNode *>()(key.second));
            return seed;
        }
    };
    std::unordered_map<std::pair<const LatticeNode *, const LatticeNode *>,
                       float, ModelLookupCacheHash>
        cache;
    if (nbest > 1) {
        std::unordered_set<std::string> dup;
        dup.insert(l.d_ptr->nbests_[0].toString());

        std::vector<NBestNode> pool;
        pool.reserve(MAX_BACKWARD_SEARCH_SIZE + 1);

        NBestNodeLess cmp(pool);
        using PriorityQueueType =
            std::priority_queue<size_t, std::vector<size_t>, NBestNodeLess>;
        PriorityQueueType q(cmp);
        PriorityQueueType result(cmp);

        auto *eos = &lattice[nullptr][0];
        auto pushNewNBestNode =
            [&](const LatticeNode *node,
                size_t next = std::numeric_limits<size_t>::max()) -> size_t {
            pool.emplace_back(node, next);
            return pool.size() - 1;
        };

        q.push(pushNewNBestNode(eos));
        auto *bos = &lattice[&graph.start()][0];
        while (!q.empty()) {
            size_t nodeIdx = q.top();
            q.pop();
            const auto &node = pool[nodeIdx];
            if (bos == node.node()) {
                auto sentence = concatNBest(pool, &node);
                if (dup.contains(sentence)) {
                    continue;
                }

                if (eos->score() - node.fn_ > max) {
                    break;
                }
                result.push(nodeIdx);
                if (result.size() >= nbest) {
                    break;
                }
                dup.insert(sentence);
            } else {
                if (pool.size() >= MAX_BACKWARD_SEARCH_SIZE) {
                    continue;
                }
                auto searchSize = beamSize;
                const auto *from_node = node.node();
                if (searchSize) {
                    searchSize =
                        std::min(searchSize, lattice[from_node->from()].size());
                } else {
                    searchSize = lattice[from_node->from()].size();
                }
                for (auto &from : lattice[from_node->from()] |
                                      std::views::take(searchSize)) {
                    float score;

                    if (node.node()->prev() == &from) {
                        // We can skip the model call if the node is the same as
                        // the previous one from forward search.
                        // Following should hold:
                        // node.score = from.score + model_score + node.cost
                        score = node.node()->score() - from.score();
                    } else {
                        auto it =
                            cache.find(std::make_pair(&from, node.node()));
                        if (it != cache.end()) {
                            score = it->second;
                        } else {
                            score = model_->score(from.state(), *node.node(),
                                                  state);
                            cache[std::make_pair(&std::as_const(from),
                                                 node.node())] = score;
                        }
                        score += node.node()->cost();
                    }
                    if (&from != bos && score < min) {
                        continue;
                    }

                    const float gn = score + node.gn_;
                    const float fn = gn + from.score();
                    // fn should be the best possible score for the current
                    // sentence, if diff is worse than the diff max, we can skip
                    // it.
                    if (eos->score() - fn <= max) {
                        size_t parentIdx = pushNewNBestNode(&from, nodeIdx);
                        pool[parentIdx].gn_ = gn;
                        pool[parentIdx].fn_ = fn;
                        q.push(parentIdx);
                        if (pool.size() >= MAX_BACKWARD_SEARCH_SIZE) {
                            break;
                        }
                    } else if (node.node()->prev() == &from) {
                        // lattice should be sorted. So this is the best
                        // possible score for the current sentence, if diff is
                        // larger than max, we can skip the rest of the parent
                        // nodes.
                        break;
                    }
                }
            }
        }

        while (!result.empty()) {
            size_t nodeIdx = result.top();
            result.pop();
            const auto &node = pool[nodeIdx];
            // loop twice to avoid problem
            size_t count = 0;
            // skip bos
            const auto *pivot = node.next(pool);
            while (pivot) {
                pivot = pivot->next(pool);
                count++;
            }
            SentenceResult::Sentence sentence;
            sentence.reserve(count);
            pivot = node.next(pool);
            while (pivot) {
                if (pivot->node()->to()) {
                    sentence.emplace_back(pivot->node());
                }
                pivot = pivot->next(pool);
            }
            l.d_ptr->nbests_.emplace_back(std::move(sentence), node.fn_);
        }
    }
}

Decoder::Decoder(const Dictionary *dict, const LanguageModelBase *model)
    : d_ptr(std::make_unique<DecoderPrivate>(dict, model)) {}

Decoder::~Decoder() {}

const Dictionary *Decoder::dict() const {
    FCITX_D();
    return d->dict_;
}

const LanguageModelBase *Decoder::model() const {
    FCITX_D();
    return d->model_;
}

bool Decoder::decode(Lattice &l, const SegmentGraph &graph, size_t nbest,
                     const State &beginState, float max, float min,
                     size_t beamSize, size_t frameSize, void *helper) const {
    FCITX_D();
    LatticeMap &lattice = l.d_ptr->lattice_;
    // Clear the result.
    l.d_ptr->nbests_.clear();
    // Remove end node.
    lattice.erase(nullptr);
    std::unordered_set<const SegmentGraphNode *> ignore;
    // Add existing SegmentGraphNode to ignore set.
    for (auto &p : lattice) {
        ignore.insert(p.first);
    }

    auto t0 = std::chrono::high_resolution_clock::now();

    if (!d->buildLattice(this, l, ignore, beginState, graph, frameSize,
                         helper)) {
        return false;
    }
    LIBIME_DEBUG() << "Build Lattice: " << millisecondsTill(t0);
    d->forwardSearch(this, graph, l, ignore, beamSize);
    LIBIME_DEBUG() << "Forward Search: " << millisecondsTill(t0);
    d->backwardSearch(graph, l, nbest, max, min, nbest * 2);
    LIBIME_DEBUG() << "Backward Search: " << millisecondsTill(t0);
    return true;
}

LatticeNode *Decoder::createLatticeNodeImpl(
    const SegmentGraphBase & /*unused*/, const LanguageModelBase * /*unused*/,
    std::string_view word, WordIndex idx, SegmentGraphPath path,
    const State &state, float cost, std::unique_ptr<LatticeNodeData> /*unused*/,
    bool /*unused*/) const {
    return new LatticeNode(word, idx, std::move(path), state, cost);
}
} // namespace libime
