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

namespace py = pybind11;

// myext.cpp
// A small pybind11 module exposing two functions:
// 1) add_arrays: element-wise sum of two numpy arrays (returns a new array)
// 2) variance: compute the variance of all elements in an array (population variance)

// add_arrays
// - Accepts two double-typed numpy arrays with the same shape.
// - Returns a new numpy array (py::array_t<double>) containing element-wise sums.
py::array_t<double> add_arrays(py::array_t<double> a, py::array_t<double> b) {
    auto buf_a = a.request();
    auto buf_b = b.request();

    if (buf_a.ndim != buf_b.ndim) {
        throw std::runtime_error("Input arrays must have the same number of dimensions");
    }

    for (ssize_t i = 0; i < buf_a.ndim; ++i) {
        if (buf_a.shape[i] != buf_b.shape[i]) {
            throw std::runtime_error("Input arrays must have the same shape");
        }
    }

    // Create output array with the same number of elements and shape
    auto out = py::array_t<double>(buf_a.size);
    out.resize(buf_a.shape);
    auto buf_out = out.request();

    double *ptr_a = static_cast<double*>(buf_a.ptr);
    double *ptr_b = static_cast<double*>(buf_b.ptr);
    double *ptr_out = static_cast<double*>(buf_out.ptr);

    for (py::ssize_t i = 0; i < buf_a.size; ++i) {
        ptr_out[i] = ptr_a[i] + ptr_b[i];
    }

    return out;
}

// variance
// - Accepts a double-typed numpy array of any shape
// - Computes the population variance over all elements: variance = sum((x - mean)^2) / N
// - Returns a double

double variance(py::array_t<double> arr) {
    auto buf = arr.request();

    if (buf.size == 0) {
        throw std::runtime_error("Input array must contain at least one element");
    }

    double *ptr = static_cast<double*>(buf.ptr);
    py::ssize_t N = buf.size;

    // Compute mean
    double sum = 0.0;
    for (py::ssize_t i = 0; i < N; ++i) sum += ptr[i];
    double mean = sum / static_cast<double>(N);

    // Compute variance
    double accum = 0.0;
    for (py::ssize_t i = 0; i < N; ++i) {
        double d = ptr[i] - mean;
        accum += d * d;
    }

    return accum / static_cast<double>(N);
}

PYBIND11_MODULE(myext, m) {
    m.doc() = "Simple array operations: add_arrays and variance";

    m.def("add_arrays", &add_arrays, "Element-wise sum of two numpy arrays (returns new array)",
          py::arg("a"), py::arg("b"));

    m.def("variance", &variance, "Compute population variance of all elements in an array",
          py::arg("arr"));
}
