import cppyy
from cppyy import gbl

cppyy.cppdef(
    r"""
struct ERDataFrame {
    size_t rows;
    std::vector<size_t> elements_size;
    std::vector<std::string> column_names;
    std::vector<void*> columns;
    bool owned = true;
    ERDataFrame(size_t s) : rows(s) {}
    ~ERDataFrame() {
        if (!owned) return;
        for (void* i : columns) free(i);
    }
    template <typename R, typename... T>
    ERDataFrame Define(std::string name, R (*f)(T...)) {
        std::cout << "Adding column " << columns.size() + 1 << " " << name << std::endl;
        column_names.push_back(name);
        columns.push_back(new R[rows]);
        for (size_t i = 0; i < rows; i++) {
            R* res = static_cast<R*>(columns.back()) + i;
            size_t I = 0;
            // Collect arguments safely in sequence
            auto args = std::tuple{(*(static_cast<T*>(columns[I++]) + i))...};
            *res = std::apply(f, args);
        }
        owned = false;
        return *this;
    }

    template <typename C, size_t N>
    void print() {
        std::cout << column_names[N] << ": [";
        for (size_t i = 0; i < rows; i++) {
            std::cout << *(static_cast<C*>(columns[N]) + i) << ", ";
        }
        std::cout << "]" << std::endl;
    }
};

template<size_t N>
void print(ERDataFrame &df) {
    df.print<double, N>();
}

double get_one() { return 1.0; }
double plus_one(double x) { return x + 1.0; }
double add(double x, double y) { return x + y; }
double add3(double x, double y, double z) { return x + y + z; }
"""
)

df = (
    gbl.ERDataFrame(1)
    .Define("one", gbl.get_one)
    .Define("two", gbl.plus_one)
    .Define("add", gbl.add)
    .Define("add3", gbl.add3)
)

gbl.print[0](df)
gbl.print[1](df)
gbl.print[2](df)
gbl.print[3](df)