main.cpp

#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <sstream>
#include <cmath>
#include <random>
#include <Eigen/Dense>


// --- 辅助函数 ---

// Sigmoid 激活函数
double sigmoid(double x) {
    return 1.0 / (1.0 + std::exp(-x));
}

// Sigmoid 导数
double sigmoid_derivative(double x) {
    double s = sigmoid(x);
    return s * (1.0 - s);
}

// --- 神经网络类 ---

class SimpleNN {
private:
    int input_size;
    int hidden_size;
    int output_size;
    double learning_rate;

    // 权重和偏置
    // W1: hidden_size x input_size
    Eigen::MatrixXd W1;
    Eigen::VectorXd b1;
    // W2: output_size x hidden_size
    Eigen::MatrixXd W2;
    Eigen::VectorXd b2;

public:
    SimpleNN(int in, int hid, int out, double lr) 
        : input_size(in), hidden_size(hid), output_size(out), learning_rate(lr) {
        
        // 初始化随机数种子
        std::random_device rd;
        std::mt19937 gen(rd());
        std::normal_distribution<double> dist(0.0, 1.0);

        // Xavier 初始化
        W1 = Eigen::MatrixXd::NullaryExpr(hidden_size, input_size, [&](){ return dist(gen) * sqrt(2.0/input_size); });
        b1 = Eigen::VectorXd::Zero(hidden_size);
        W2 = Eigen::MatrixXd::NullaryExpr(output_size, hidden_size, [&](){ return dist(gen) * sqrt(2.0/hidden_size); });
        b2 = Eigen::VectorXd::Zero(output_size);
    }

    // 前向传播
    Eigen::VectorXd forward(const Eigen::VectorXd& input, Eigen::VectorXd& hidden_activation) {
        // Layer 1 (Hidden)
        Eigen::VectorXd z1 = W1 * input + b1;
        hidden_activation = z1.unaryExpr(&sigmoid); // Apply Sigmoid
        
        // Layer 2 (Output) - 线性激活用于回归任务
        Eigen::VectorXd z2 = W2 * hidden_activation + b2;
        return z2; 
    }

    // 单步训练 (SGD)
    double train_step(const Eigen::VectorXd& input, const Eigen::VectorXd& target) {
        Eigen::VectorXd hidden_out;
        Eigen::VectorXd output = forward(input, hidden_out);

        // 计算 Loss (MSE) 用于显示
        Eigen::VectorXd error = output - target;
        double loss = error.squaredNorm() / 2.0;

        // --- 反向传播 ---
        
        // Output Layer Gradients (Linear Activation derivative is 1)
        // dL/dOutput = (output - target)
        Eigen::VectorXd delta_output = error;

        // Hidden Layer Gradients
        // dL/dW2 = delta_output * hidden_out^T
        Eigen::MatrixXd dW2 = delta_output * hidden_out.transpose();
        Eigen::VectorXd db2 = delta_output;

        // dL/dHidden = W2^T * delta_output
        Eigen::VectorXd error_hidden = W2.transpose() * delta_output;
        
        // dL/dZ1 = error_hidden * sigmoid_derivative(Z1)
        // 注意：这里的 hidden_out 已经是 sigmoid(z1)
        Eigen::VectorXd delta_hidden = error_hidden.cwiseProduct(hidden_out.unaryExpr([](double a){ return a * (1.0 - a); }));

        // dL/dW1 = delta_hidden * input^T
        Eigen::MatrixXd dW1 = delta_hidden * input.transpose();
        Eigen::VectorXd db1 = delta_hidden;

        // --- 更新权重 ---
        W1 -= learning_rate * dW1;
        b1 -= learning_rate * db1;
        W2 -= learning_rate * dW2;
        b2 -= learning_rate * db2;

        return loss;
    }

    // 保存模型
    void save_model(const std::string& filename) {
        std::ofstream file(filename);
        if (file.is_open()) {
            file << W1 << std::endl << b1 << std::endl << W2 << std::endl << b2 << std::endl;
            file.close();
            std::cout << "模型已保存至: " << filename << std::endl;
        } else {
            std::cerr << "无法打开文件保存模型。" << std::endl;
        }
    }

    // 加载模型
    void load_model(const std::string& filename) {
        std::ifstream file(filename);
        if (file.is_open()) {
            for(int i=0; i<W1.rows(); ++i) for(int j=0; j<W1.cols(); ++j) file >> W1(i,j);
            for(int i=0; i<b1.size(); ++i) file >> b1(i);
            for(int i=0; i<W2.rows(); ++i) for(int j=0; j<W2.cols(); ++j) file >> W2(i,j);
            for(int i=0; i<b2.size(); ++i) file >> b2(i);
            file.close();
            std::cout << "模型已加载: " << filename << std::endl;
        } else {
            std::cerr << "无法打开模型文件，将使用随机初始化参数。" << std::endl;
        }
    }
};

// --- 数据处理 ---

struct DataPoint {
    Eigen::VectorXd input; // Size 6: [Op1, Op2, Op3, Op4, A, B]
    Eigen::VectorXd target; // Size 1: [Result]
};

// CSV 格式: op_type, a, b, result (训练时需要result，推理时忽略)
// op_type: 0=add, 1=sub, 2=mul, 3=div
std::vector<DataPoint> load_data(const std::string& filename, bool is_training) {
    std::vector<DataPoint> data;
    std::ifstream file(filename);
    std::string line;

    if (!file.is_open()) {
        std::cerr << "错误: 无法打开数据文件 " << filename << std::endl;
        exit(1);
    }

    while (getline(file, line)) {
        std::stringstream ss(line);
        std::string val;
        std::vector<double> row;
        
        while (getline(ss, val, ',')) {
            row.push_back(stod(val));
        }

        if (row.size() < 3) continue;

        DataPoint dp;
        dp.input = Eigen::VectorXd::Zero(6);
        
        // One-hot encoding for operation
        int op = (int)row[0];
        if (op >=0 && op <= 3) dp.input(op) = 1.0;
        
        // Operands
        dp.input(4) = row[1];
        dp.input(5) = row[2];

        if (is_training && row.size() >= 4) {
            dp.target = Eigen::VectorXd(1);
            dp.target(0) = row[3];
        }

        data.push_back(dp);
    }
    return data;
}

// --- 主程序 ---

int main(int argc, char* argv[]) {
    if (argc < 4) {
        std::cout << "用法:" << std::endl;
        std::cout << "  训练: ./neural_net --train <数据文件> <模型保存路径>" << std::endl;
        std::cout << "  推理: ./neural_net --infer <数据文件> <模型加载路径>" << std::endl;
        return 1;
    }

    std::string mode = argv[1];
    std::string data_file = argv[2];
    std::string model_file = argv[3];

    // 网络配置: 6 输入, 60 隐藏, 1 输出, 学习率 0.01
    SimpleNN nn(6, 60, 1, 0.08);

    if (mode == "--train") {
        std::cout << "=== 开始训练 ===" << std::endl;
        std::vector<DataPoint> training_data = load_data(data_file, true);

        // 1. 打开一个日志文件
        std::ofstream log_file("loss_log.txt");
        if (!log_file.is_open()) {
            std::cerr << "无法创建日志文件 loss_log.txt" << std::endl;
            return 1;
        }
        //  写入表头，方便查看
        log_file << "epoch,loss" << std::endl;
        
        int epochs = 5000; // 训练轮数
        for (int epoch = 0; epoch < epochs; ++epoch) {
            double total_loss = 0;
            // 简单的随机洗牌可以增加鲁棒性 (这里省略，直接遍历)
            for (const auto& point : training_data) {
                total_loss += nn.train_step(point.input, point.target);
            }
            
            // 每 100 轮记录一次
            if (epoch % 100 == 0) {
                double avg_loss = total_loss / training_data.size();
                std::cout << "Epoch " << epoch << " | Avg Loss: " << total_loss / training_data.size() << std::endl;
                log_file << epoch << "," << avg_loss << std::endl;
            }
        }
        log_file.close();
        nn.save_model(model_file);

    } else if (mode == "--infer") {
        std::cout << "=== 开始推理 ===" << std::endl;
        nn.load_model(model_file);
        std::vector<DataPoint> infer_data = load_data(data_file, false);

        std::cout << "Op\tA\tB\t预测结果" << std::endl;
        std::cout << "------------------------------------" << std::endl;

        Eigen::VectorXd dummy;
        for (const auto& point : infer_data) {
            Eigen::VectorXd result = nn.forward(point.input, dummy);
            
            std::string op_sym;
            if (point.input(0) == 1) op_sym = "+";
            else if (point.input(1) == 1) op_sym = "-";
            else if (point.input(2) == 1) op_sym = "*";
            else if (point.input(3) == 1) op_sym = "/";
            else op_sym = "?";

            std::cout << op_sym << "\t" 
                 << point.input(4) << "\t" 
                 << point.input(5) << "\t" 
                 << result(0) << std::endl;
        }
    } else {
        std::cerr << "未知模式: " << mode << std::endl;
        return 1;
    }

    return 0;
}