#include <nanobind/nanobind.h>
#include <nanobind/stl/vector.h>
#include <nanobind/stl/complex.h>
#include <nanobind/eigen/dense.h>
#include <Eigen/Dense>
#include "tdscontrol/roots.hpp"
#include "tdscontrol/tds.hpp"

namespace nb = nanobind;

NB_MODULE(pytdscontrol, m) {
    m.doc() = "tds-control python bindings";

    nb::class_<tds::tds>(m, "tds")
        .def(nb::init<std::vector<Eigen::MatrixXd>, std::vector<double>>(),
             "A", "hA")
        .def_prop_ro("mA", &tds::tds::mA)
        .def_prop_ro("n", &tds::tds::n)
        .def("A", nb::overload_cast<const std::size_t>(&tds::tds::A), nb::rv_policy::reference_internal)
        .def("hA", [](tds::tds &sys) {
            return nb::ndarray<nb::numpy, double>(sys.hA().data(), {sys.mA()} ,nb::handle()); }, nb::rv_policy::reference_internal);
    
    // TODO: return numpy array instead of list
    m.def("roots", &tds::roots);
}