#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>  // needed for array handling
#include <pybind11/stl.h>  // Required for automatic conversion of STL containers
#include <iostream>
#include <vector>
#include <map>
#include <string>
#include <cmath>

namespace py = pybind11;

// Example 1: Calculate square of array elements
py::array_t<double> square(py::array_t<double> inp) {
    // Get information about the input array
    auto buf = inp.request();
    
    // Create output array
    auto out = py::array_t<double>(buf.size);
    auto buf_out = out.request();
    
    // Get data pointers
    double *ptr = static_cast<double*>(buf.ptr);
    double *ptr_out = static_cast<double*>(buf_out.ptr);
    
    // Perform calculation
    for (py::ssize_t i = 0; i < buf.size; i++) {
        ptr_out[i] = ptr[i] * ptr[i];
    }
    
    return out;
}

// Example 2: 2D array operation - element-wise matrix multiplication
py::array_t<double> ew_multiply(py::array_t<double> inp1, py::array_t<double> inp2) {
    auto buf1 = inp1.request();
    auto buf2 = inp2.request();
    
    // Check dimensions
    if (buf1.ndim != 2 || buf2.ndim != 2) {
        throw std::runtime_error("Input arrays must be 2-dimensional");
    }
    
    // Check if shapes match
    if (buf1.shape[0] != buf2.shape[0] || buf1.shape[1] != buf2.shape[1]) {
        throw std::runtime_error("Input arrays must have the same shape");
    }
    
    // Create output array with the same shape
    auto out = py::array_t<double>(buf1.size);
    out.resize({buf1.shape[0], buf1.shape[1]});
    auto buf_out = out.request();
    
    double *ptr1 = static_cast<double*>(buf1.ptr);
    double *ptr2 = static_cast<double*>(buf2.ptr);
    double *ptr_out = static_cast<double*>(buf_out.ptr);
    
    for (py::ssize_t i = 0; i < buf1.size; i++) {
        ptr_out[i] = ptr1[i] * ptr2[i];
    }
    
    return out;
}

// Example 3: Statistical function - compute mean(axis=0) of a 2D array
// Returns a 1-D array of column means (shape: [cols]).
py::array_t<double> mean_row(py::array_t<double> inp) {
    auto buf = inp.request();

    if (buf.ndim != 2) {
        throw std::runtime_error("Input array must be 2-dimensional");
    }

    py::ssize_t rows = buf.shape[0];
    py::ssize_t cols = buf.shape[1];

    if (rows == 0) {
        throw std::runtime_error("Input array must have at least one row");
    }

    double *ptr = static_cast<double*>(buf.ptr);

    // Create output 1-D array with length = cols
    auto out = py::array_t<double>(cols);
    auto buf_out = out.request();
    double *ptr_out = static_cast<double*>(buf_out.ptr);

    // Initialize sums to 0
    for (py::ssize_t j = 0; j < cols; ++j) {
        ptr_out[j] = 0.0;
    }

    // Accumulate column sums. Assuming row-major contiguous layout.
    for (py::ssize_t i = 0; i < rows; ++i) {
        for (py::ssize_t j = 0; j < cols; ++j) {
            ptr_out[j] += ptr[i * cols + j];
        }
    }

    // Divide by number of rows to get mean
    for (py::ssize_t j = 0; j < cols; ++j) {
        ptr_out[j] /= static_cast<double>(rows);
    }

    return out;
}

// Example 4: In-place array modification
void multiply_inplace(py::array_t<double> inp) {
    auto buf = inp.request();
    double *ptr = static_cast<double*>(buf.ptr);
    
    for (py::ssize_t i = 0; i < buf.size; i++) {
        ptr[i] *= 2;
    }
}

// Example 5: Working with lists. Double all elements in a list and return new list
std::vector<double> double_list(const std::list<double>& inp) {
    std::vector<double> out;
    for (double val : inp) {
        out.push_back(val * 2);
    }
    return out;
}

// Example 6: Working with tuples. Calculate statistics from a tuple of numbers
py::tuple calculate_statistics(const std::tuple<double, double, double>& values) {
    double a = std::get<0>(values);
    double b = std::get<1>(values);
    double c = std::get<2>(values);
    
    double sum = a + b + c;
    double mean = sum / 3.0;
    double product = a * b * c;
    
    return py::make_tuple(sum, mean, product);
}

// Example 7: Working with integers and floats. Calculate area of a rectangle
double rectangle_area(double length, double width) {
    return length * width;
}

// Example 8: Working with dictionaries. Create a student record
std::map<std::string, py::object> create_student_record(
    const std::string& name, 
    int age, 
    const std::vector<double>& grades) {
    
    std::map<std::string, py::object> record;
    
    record["name"] = py::cast(name);
    record["age"] = py::cast(age);
    record["grades"] = py::cast(grades);
    
    // Calculate average grade
    double sum = 0;
    for (double grade : grades) {
        sum += grade;
    }
    double average = grades.empty() ? 0 : sum / grades.size();
    record["average"] = py::cast(average);
    
    return record;
}

// Module definition
PYBIND11_MODULE(myext, m) {
    m.doc() = "Pybind11 numpy/python-types example plugin";
    
    m.def("square", &square, "Calculate square of each element in the array");
    m.def("ew_multiply", &ew_multiply, "Element-wise multiplication of two matrices");
    m.def("mean_row", &mean_row, "Compute mean of array elements along axis=0");
    m.def("multiply_inplace", &multiply_inplace, "Multiply array elements by 2 in-place");
    m.def("double_list", &double_list, "Double all elements in a list and return new list");
    m.def("calculate_statistics", &calculate_statistics, "Calculate statistics from a tuple of numbers");
    m.def("rectangle_area", &rectangle_area, "Calculate area of a rectangle");
    m.def("create_student_record", &create_student_record, "Create a student record dictionary");
}