main.cpp

#include <stdio.h>
#include <string.h>
#include <bitset>
#include <vector>
#include <utility>
#include <algorithm>
#include "stmr.h"
#include "utils.h"
#include "global.h"
#include "getfile.h"

bool getInput(std::string &query) {
    query = "";
    std::cout << "Enter query that ends with EOF in a single line: " << std::endl;
    std::string buf;
    while( true ) {
        std::getline(std::cin, buf);
        if( buf == "EOF" ) break;
        query += buf + " ";
    }
    return true;
}

struct FileInfo {
    size_t ID;
    size_t reg;
    float reg_score;
    float matched_words;
    float tfidf_product;
    FileInfo(): ID(0), reg(0), reg_score(0.0), matched_words(0), tfidf_product(0.0) {}
    FileInfo(size_t ID, size_t reg, float reg_score, float matched_words, float tfidf_product): ID(ID), reg(reg), reg_score(reg_score), matched_words(matched_words), tfidf_product(tfidf_product) {}
    bool operator < (const FileInfo& other) const {
        return reg < other.reg;
    }
};

std::string slice_origin(const std::string& str, int &p, const int& q) {
    char ret[100] = {0}; // TODO：这里大小是最长单词的长度，可能要继续调整
    int len = 0;
    while( p <= q && valid_char(str[p]) ) ret[len++] = tolower(str[p++]);
    std::string ret_str(ret);
    return ret_str;
}

std::string getAbstract(const std::string& filename, const int& slice_pos, int word_count) {
    std::string text = parsefile(filename);
    std::string abstract = "... ";
    int p = 0, q = text.size() - 1;
    for(int p = 0, q = text.size() - 1, file_index = 0; p <= q; ++file_index) {
        std::string wd = slice_origin(text, p, q);
        if( wd == "" ) {
            while( p <= q && !valid_char(text[p]) ) ++p;
            --file_index;
            continue;
        }
        if( file_index >= slice_pos ) {
            abstract += wd + " ";
            if( --word_count == 0 ) break;
        }
    } 
    return abstract + "...";
}

int main() {

    // count_word();
    // return 0;

    std::cout << "Loading files..." << std::endl;
    getfile();
    std::cout << "Loading files completed" << std::endl;
    std::cout << "Calculating TF-IDF..." << std::endl;
    TF_IDFcalc();
    std::cout << "Calculating TF-IDF completed" << std::endl;
    // 解析文件并将其存入索引中
    // 之后只要完成询问就好

    std::string query;
    while( getInput(query) ) {
        // TODO：查询主要流程是读入所有查询值，然后先获取所有搜索到的文件，之后进行权重排序
        // 最后应该直接输出一个文件名排名表即可。
        std::bitset<TotalFiles+1> file_vis(0);
        std::vector<int> file_valid_word_cnt(TotalFiles+1, 0);
        std::vector<std::pair<float, FileInfo> > result; 
        std::vector<size_t> words_id;  
        int p = 0, q = query.size() - 1;
        std::string wd;
        File qfile;
        while( p <= q ) {
            int tmp_p = p;
            wd = slice(query, p, q);
            if( wd == "" ) {
                while( p <= q && !valid_char(query[p]) ) ++p;
                continue;
            }
            int wd_id = word_set.getIndex(wd);
            if( !wd_id ) {
                std::cout << "A";
                // std::cout << "Unrecognized or stop word: " << wd << std::endl;
                wd_id = word_trie.search(wd);
                if( wd_id == 0 ) {
                    std::cout << "Unrecognized word: " << slice_origin(query, tmp_p, q) << ", cannot find similar word." << std::endl;
                    continue;
                }
                wd = word_set.lookupTable[wd_id];
                std::cout << "Unrecognized word: " << slice_origin(query, tmp_p, q) << ", are you saying: " << wd << "?" << std::endl;
            } 
            Word* word = word_set.findVal(wd);
            std::bitset<TotalFiles+1> file_set = word->get_file_set();
            file_vis |= file_set;
            qfile.tfidf[wd_id] += 1;
            if( float_equal(qfile.tfidf[wd_id], 1) ) {
                words_id.push_back(wd_id);
                for(int i = 1; i <= TotalFiles; ++i) {
                    file_valid_word_cnt[i] += file_set[i]; // 计算每个文件中的有效词数量
                }
            }
            qfile.appendWordIndex(wd_id);
            while( p <= q && !valid_char(query[p]) ) ++p;
        }
        if( file_vis.count() == 0 )  {
            std::cout << "No valid file found by query." << std::endl;
            continue;
        }
        for(auto i: words_id) {
            qfile.tfidf[i] = qfile.tfidf[i] / qfile.word_count * word_set.findVal(word_set.lookupTable[i])->idf; 
        }
        // std::cout << "A";
        for(size_t i = 1; i <= TotalFiles; ++i) {
            if( file_vis[i] == 1 ) {
                // std::cout << i << ' ';
                File &f = file_set[i];
                if( f.word_count <= 3 ) continue; // 过滤掉长度过短的文件
                std::vector<int> file_word_pos(f.word_count, 0);
                // std::cout << qfile.word_list[0] << std::endl;
                // std::cout << "C";
                for(int j = 0; j < qfile.word_count; ++j) {
                    f.getCommit(qfile.word_list[j], -j, file_word_pos);
                }
                // std::cout << "D";
                size_t reg = 0;
                for(int j = 1; j < f.word_count-2; ++j) {
                    if( 0.15 * file_word_pos[j] + 0.7 * file_word_pos[j+1] + 0.15 * file_word_pos[j+2] > 
                        0.15 * file_word_pos[reg] + 0.7 * file_word_pos[reg+1] + 0.15 * file_word_pos[reg+2] ) {
                            reg = j;
                    }
                }
                // std::cout << "E";

                float tfidf_product = qfile.tfidf * file_set[i].tfidf;
                result.push_back(
                    std::make_pair(tfidf_product * 0.4 + file_valid_word_cnt[i] * 0.55 / words_id.size() + (0.15 * file_word_pos[reg] + 0.7 * file_word_pos[reg+1] + 0.15 * file_word_pos[reg+2]) * 0.05, 
                FileInfo(i, reg, 0.15 * file_word_pos[reg] + 0.7 * file_word_pos[reg+1] + 0.15 * file_word_pos[reg+2], (float)file_valid_word_cnt[i] / words_id.size(), tfidf_product))); 
            }   
        }
        sort(result.begin(), result.end(), std::greater<std::pair<float, FileInfo>>());
        float threshold = result[0].first * 0.01;
        for(size_t i = 0; i < result.size(); ++i) {
            if( result[i].first < threshold || i >= 20 ) break;
            std::cout << "File: " << file_set.lookupTable[result[i].second.ID] << ", slice_pos: " << result[i].second.reg << ", slice_score: " << result[i].second.reg_score << ", General score: " << result[i].first << std::endl;
            std::cout << "Words matched: " << result[i].second.matched_words << ", TF-IDF product: " << result[i].second.tfidf_product << std::endl;
            std::cout << "Abstract: " << getAbstract(file_set.lookupTable[result[i].second.ID], result[i].second.reg - 3, qfile.word_count + 10) << std::endl;
        }
    }
    return 0;
}