parallel_example.cpp

#include "ParallelKitsune.h"
#include <chrono>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <vector>

/**
 * 简单的命令行进度显示
 */
void printProgress(float progress, const std::string &prefix = "") {
    int barWidth = 70;

    std::cout << prefix << " [";
    int pos = barWidth * progress;
    for (int i = 0; i < barWidth; ++i) {
        if (i < pos)
            std::cout << "=";
        else if (i == pos)
            std::cout << ">";
        else
            std::cout << " ";
    }

    std::cout << "] " << int(progress * 100.0) << "%\r";
    std::cout.flush();
}

int main() {
    // 硬编码参数
    std::string pcap_path = "./mirai.pcap";        // 当前目录下的mirai.pcap
    unsigned int num_packets = 0xFFFFFFFF;         // 处理所有数据包
    std::string output_file = "mirai_results.csv"; // 输出文件名
    size_t num_threads = 10;                       // 10个线程
    size_t batch_size = 10;                        // 批处理大小为10

    // 参数设置
    size_t max_autoencoder_size = 10;
    size_t fm_grace_period = 5000;  // 前5000个数据包用于特征映射
    size_t ad_grace_period = 50000; // 随后15000个数据包用于训练异常检测器

    std::cout << "初始化并行Kitsune" << std::endl;
    std::cout << "输入文件: " << pcap_path << std::endl;
    std::cout << "线程数: " << num_threads << ", 批处理大小: " << batch_size
              << std::endl;
    std::cout << "学习期: 特征映射=" << fm_grace_period
              << ", 异常检测器=" << ad_grace_period << std::endl;

    // 计时
    auto start_time = std::chrono::high_resolution_clock::now();

    // 创建ParallelKitsune实例
    ParallelKitsune kitsune(pcap_path, num_packets, num_threads,
                            max_autoencoder_size, fm_grace_period,
                            ad_grace_period);

    // 显示实际使用的线程数
    std::cout << "实际使用线程数: " << kitsune.getNumThreads() << std::endl;

    // 处理所有数据包
    unsigned int processed = 0;
    std::vector<double> rmse_values;

    // 使用批处理模式
    while (true) {
        size_t batch_processed = kitsune.proc_batch(batch_size);
        if (batch_processed == 0) {
            break; // 无更多数据包
        }

        processed += batch_processed;

        if (processed % 1000 == 0) {
            printProgress(processed / 100000.0,
                          "处理进度"); // 假设大约有10万个数据包

            // 获取当前结果
            auto results = kitsune.get_and_clear_results();
            rmse_values.insert(rmse_values.end(), results.begin(),
                               results.end());
        }
    }

    // 等待所有任务完成并获取剩余结果
    kitsune.waitForAllTasks();
    auto remaining_results = kitsune.get_and_clear_results();
    rmse_values.insert(rmse_values.end(), remaining_results.begin(),
                       remaining_results.end());

    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(
                        end_time - start_time)
                        .count();

    std::cout << std::endl
              << "处理完成! 共处理 " << processed << " 个数据包, 用时 "
              << duration / 1000.0 << " 秒" << std::endl;
    std::cout << "平均每秒处理 " << (processed * 1000.0 / duration)
              << " 个数据包" << std::endl;

    // 将RMSE值保存到文件
    std::ofstream out_file(output_file);
    if (!out_file) {
        std::cerr << "无法打开输出文件: " << output_file << std::endl;
        return 1;
    }

    // 写入CSV文件头
    out_file << "index,rmse" << std::endl;

    // 写入RMSE值
    for (size_t i = 0; i < rmse_values.size(); ++i) {
        out_file << i << "," << std::fixed << std::setprecision(10)
                 << rmse_values[i] << std::endl;
    }

    std::cout << "异常分数已保存到: " << output_file << std::endl;

    return 0;
}