#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

namespace py = pybind11;

// 示例C++函数：向量加法
std::vector<double> vector_add(const std::vector<double>& a, const std::vector<double>& b) {
    if (a.size() != b.size()) {
        throw std::runtime_error("Vectors must have the same size");
    }
    
    std::vector<double> result(a.size());
    for (size_t i = 0; i < a.size(); ++i) {
        result[i] = a[i] + b[i];
    }
    return result;
}

// 高性能矩阵乘法（示例）
py::array_t<double> matrix_multiply(py::array_t<double> a, py::array_t<double> b) {
    py::buffer_info a_buf = a.request();
    py::buffer_info b_buf = b.request();
    
    if (a_buf.ndim != 2 || b_buf.ndim != 2) {
        throw std::runtime_error("Input must be 2D arrays");
    }
    
    if (a_buf.shape[1] != b_buf.shape[0]) {
        throw std::runtime_error("Matrix dimensions do not match");
    }
    
    auto result = py::array_t<double>({a_buf.shape[0], b_buf.shape[1]});
    py::buffer_info r_buf = result.request();
    
    double* a_ptr = static_cast<double*>(a_buf.ptr);
    double* b_ptr = static_cast<double*>(b_buf.ptr);
    double* r_ptr = static_cast<double*>(r_buf.ptr);
    
    size_t m = a_buf.shape[0];
    size_t n = a_buf.shape[1];
    size_t p = b_buf.shape[1];
    
    // 简单矩阵乘法（可优化）
    for (size_t i = 0; i < m; ++i) {
        for (size_t j = 0; j < p; ++j) {
            double sum = 0.0;
            for (size_t k = 0; k < n; ++k) {
                sum += a_ptr[i * n + k] * b_ptr[k * p + j];
            }
            r_ptr[i * p + j] = sum;
        }
    }
    
    return result;
}

// 定义Python模块
PYBIND11_MODULE(app_cpppy, m) {
    m.doc() = "C++ extension for FastAPI";
    
    m.def("vector_add", &vector_add, "Add two vectors");
    m.def("matrix_multiply", &matrix_multiply, "Multiply two matrices");
    
    // 添加更多C++函数...
}