SearchEngine.cpp

#include "SearchEngine.h"

#include <algorithm>
#include <cmath>
#include <filesystem>
#include <print>

#include "Lexer.h"
#include "lib/pugixml/pugixml.hpp"

#define SEARCH_LIMIT 10

SearchEngine::SearchEngine() {
    this->filePaths = {};
    this->termInDocCount = {};
    this->allDocTermFreq = {};
}

void SearchEngine::compile() {
    if (this->filePaths.size() == 0) {
        std::print("No document to compile");
        return;
    }

    std::println("Indexing documents...");
    this->parseAllDocAndGetTermFreq();
}

void SearchEngine::addDocument(const std::string &filePath) {
    this->filePaths.emplace_back(filePath);
}

void SearchEngine::addDir(const std::string &path) {
    std::filesystem::directory_iterator dirIter =
        std::filesystem::directory_iterator(path);
    for (const auto &filePath : dirIter) {
        this->filePaths.emplace_back(filePath.path().string());
    }
}

TermFreq SearchEngine::getTermFreq(std::string &content) {
    Lexer lexer(content);
    std::vector<std::string> terms = lexer.process();

    TermFreq result = {};
    result.reserve(int(terms.size() / 2));

    for (std::string &term : terms) {
        this->makeLowercase(term);

        if (result.find(term) == result.end())
            result[term] = 1;
        else
            result[term] += 1;
    }

    for (auto &[word, freq] : result) {
        freq /= terms.size();

        if (this->termInDocCount.find(word) == this->termInDocCount.end())
            this->termInDocCount[word] = 1;
        else
            this->termInDocCount[word]++;
    }

    return result;
}

double SearchEngine::getTermInverseDocumentFreq(const std::string &term) {
    int documentCount = 1;
    if (this->termInDocCount.find(term) != this->termInDocCount.end())
        documentCount = this->termInDocCount[term];

    return std::log(this->filePaths.size() / documentCount);
}

double SearchEngine::calculateTFIDF(
    const std::string &term, const std::string &path
) {
    return this->allDocTermFreq[path][term] *
           this->getTermInverseDocumentFreq(term);
}

void SearchEngine::parseAllDocAndGetTermFreq() {
    for (const std::string &path : this->filePaths) {
        std::string content = this->parseFile(path);
        this->allDocTermFreq[path] = this->getTermFreq(content);
    }
}

std::string SearchEngine::parseFile(const std::string &filePath) {
    pugi::xml_document doc;
    doc.load_file(filePath.data());
    return this->collectText(doc);
}

std::string SearchEngine::collectText(const pugi::xml_node &node) {
    std::string content = "";
    for (pugi::xml_node child : node.children()) {
        if (child.type() == pugi::node_pcdata ||
            child.type() == pugi::node_cdata) {
            content += " ";
            content += child.value();
        }

        content += collectText(child);
    }

    return content;
}

void SearchEngine::makeLowercase(std::string &term) {
    std::transform(term.begin(), term.end(), term.begin(), [](unsigned char c) {
        return std::tolower(c);
    });
}

void SearchEngine::find(std::string &keyword) {
    Lexer lexer(keyword);
    std::vector<std::string> terms = lexer.process();

    std::vector<std::pair<std::string, double>> ranking;
    ranking.reserve(this->allDocTermFreq.size());

    for (auto &[path, termFreq] : this->allDocTermFreq) {
        double rankingScore = 0.0;
        for (std::string &term : terms) {
            this->makeLowercase(term);
            rankingScore += this->calculateTFIDF(term, path);
        }

        ranking.emplace_back(path, rankingScore);
    }
    std::sort(
        ranking.begin(),
        ranking.end(),
        [=](std::pair<std::string, double> &first,
            std::pair<std::string, double> &second) {
            return first.second > second.second;
        }
    );

    int limit = ranking.size() > 10 ? SEARCH_LIMIT : ranking.size();
    for (int i = 0; i < limit; i++) {
        std::println("{} -> {}", ranking[i].first, ranking[i].second);
    }
}