#include <librapid/librapid.hpp>

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


#include <string>

namespace py = pybind11;

// Docstring for the module
std::string module_docstring = "A fast math and neural network library for Python and C++";

PYBIND11_MODULE(librapid, module)
{
    module.doc() = module_docstring;

    py::class_<librapid::ndarray::extent>(module, "extent")
        .def(py::init<>())
        .def(py::init<const std::vector<nd_int> &>())
        .def(py::init<nd_int>())
        .def(py::init<const librapid::ndarray::extent &>())
        .def(py::init<py::args>())

        .def("__getitem__", [](const librapid::ndarray::extent &e, nd_int index) { return e[index]; })
        .def("__setitem__", [](librapid::ndarray::extent &e, nd_int index, nd_int val) { e[index] = val; })

        .def("compressed", &librapid::ndarray::extent::compressed)
        .def_property_readonly("ndim", &librapid::ndarray::extent::ndim)
        .def_property_readonly("is_valid", &librapid::ndarray::extent::is_valid)
        .def("reshape", [](librapid::ndarray::extent &e, const std::vector<nd_int> &order) { e.reshape(order); })

        .def("__eq__", &librapid::ndarray::extent::operator==)
        .def("__str__", &librapid::ndarray::extent::str)
        .def("__repr__", [](const librapid::ndarray::extent &e) { return "<librapid." + e.str() + ">"; });


    py::class_<librapid::ndarray::stride>(module, "stride")
        .def(py::init<>())
        .def(py::init<std::vector<nd_int>>())
        .def(py::init<nd_int>())
        .def(py::init<const librapid::ndarray::stride &>())
        .def(py::init<py::args>())
        .def("from_extent", [](const std::vector<nd_int> &extent) { return librapid::ndarray::stride::from_extent(extent); })

        .def("__getitem__", [](const librapid::ndarray::stride &s, nd_int index) { return s[index]; })
        .def("__setitem__", [](librapid::ndarray::stride &s, nd_int index, nd_int val) { s[index] = val; })

        .def_property_readonly("ndim", &librapid::ndarray::stride::ndim)
        .def_property_readonly("is_valid", &librapid::ndarray::stride::is_valid)
        .def_property_readonly("is_trivial", &librapid::ndarray::stride::is_trivial)
        .def("set_dimensions", &librapid::ndarray::stride::set_dimensions)
        .def("reshape", [](librapid::ndarray::stride &s, const std::vector<nd_int> &order) { s.reshape(order); })

        .def("__eq__", &librapid::ndarray::stride::operator==)
        .def("__str__", &librapid::ndarray::stride::str)
        .def("__repr__", [](const librapid::ndarray::stride &s) { return "<librapid." + s.str() + ">"; });


    py::class_<librapid::ndarray::ndarray>(module, "ndarray")
        .def(py::init<>())
        .def(py::init<const librapid::ndarray::extent &>())
        .def(py::init<const librapid::ndarray::extent &, double>())
        .def(py::init<const librapid::ndarray::ndarray &>())

        .def_property_readonly("ndim", &librapid::ndarray::ndarray::ndim)
        .def_property_readonly("is_initialized", &librapid::ndarray::ndarray::is_initialized)
        .def_property_readonly("is_scalar", &librapid::ndarray::ndarray::is_scalar)

        .def("get_extent", &librapid::ndarray::ndarray::get_extent)
        .def("get_stride", &librapid::ndarray::ndarray::get_stride)
        .def_property_readonly("extent", &librapid::ndarray::ndarray::get_extent)
        .def_property_readonly("stride", &librapid::ndarray::ndarray::get_stride)

        .def("__getitem__", [](const librapid::ndarray::ndarray &arr, nd_int index) { return arr[index]; })
        .def("__setitem__", [](librapid::ndarray::ndarray &arr, nd_int index, const librapid::ndarray::ndarray &value) { arr[index] = value; })
        .def("__setitem__", [](librapid::ndarray::ndarray &arr, nd_int index, double value) { arr[index] = value; })

        .def("fill", [](librapid::ndarray::ndarray &arr, double filler) { arr.fill(filler); })
        .def("filled", [](const librapid::ndarray::ndarray &arr, double filler) { return arr.filled(filler); })
        .def("clone", &librapid::ndarray::ndarray::clone)
        .def("set_value", &librapid::ndarray::ndarray::set_value)

        .def("__add__", [](const librapid::ndarray::ndarray &lhs, const librapid::ndarray::ndarray &rhs) { return lhs + rhs; })
        .def("__sub__", [](const librapid::ndarray::ndarray &lhs, const librapid::ndarray::ndarray &rhs) { return lhs - rhs; })
        .def("__mul__", [](const librapid::ndarray::ndarray &lhs, const librapid::ndarray::ndarray &rhs) { return lhs * rhs; })
        .def("__truediv__", [](const librapid::ndarray::ndarray &lhs, const librapid::ndarray::ndarray &rhs) { return lhs / rhs; })

        .def("__add__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return lhs + rhs; })
        .def("__sub__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return lhs - rhs; })
        .def("__mul__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return lhs * rhs; })
        .def("__truediv__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return lhs / rhs; })

        .def("__add__", [](double lhs, const librapid::ndarray::ndarray &rhs) { return lhs + rhs; })
        .def("__sub__", [](double lhs, const librapid::ndarray::ndarray &rhs) { return lhs - rhs; })
        .def("__mul__", [](double lhs, const librapid::ndarray::ndarray &rhs) { return lhs * rhs; })
        .def("__truediv__", [](double lhs, const librapid::ndarray::ndarray &rhs) { return lhs / rhs; })

        .def("__radd__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return rhs + lhs; })
        .def("__rsub__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return rhs - lhs; })
        .def("__rmul__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return rhs * lhs; })
        .def("__rtruediv__", [](const librapid::ndarray::ndarray &lhs, double rhs) { return rhs / lhs; })

        .def("__neg__", [](const librapid::ndarray::ndarray &arr) { return -arr; })

        .def("dot", [](const librapid::ndarray::ndarray &lhs, const librapid::ndarray::ndarray &rhs) { return lhs.dot(rhs); })
        .def("__matmul__", [](const librapid::ndarray::ndarray &lhs, const librapid::ndarray::ndarray &rhs) { return lhs.dot(rhs); })

        .def("reshape", [](librapid::ndarray::ndarray &arr, const std::vector<nd_int> &order) { arr.reshape(order); })
        .def("reshaped", [](const librapid::ndarray::ndarray &arr, const std::vector<nd_int> &order) { return arr.reshaped(order); })
        .def("reshape", [](librapid::ndarray::ndarray &arr, py::args args) { arr.reshape(py::cast<std::vector<nd_int>>(args)); })
        .def("reshaped", [](const librapid::ndarray::ndarray &arr, py::args args) { return arr.reshaped(py::cast<std::vector<nd_int>>(args)); })
        .def("subarray", [](const librapid::ndarray::ndarray &arr, const std::vector<nd_int> &axes) { return arr.subarray(axes); })

        .def("strip_front", &librapid::ndarray::ndarray::strip_front)
        .def("strip_back", &librapid::ndarray::ndarray::strip_back)
        .def("strip", &librapid::ndarray::ndarray::strip)
        .def("stripped_front", &librapid::ndarray::ndarray::stripped_front)
        .def("stripped_back", &librapid::ndarray::ndarray::stripped_back)
        .def("stripped", &librapid::ndarray::ndarray::stripped)

        .def("transpose", [](librapid::ndarray::ndarray &arr, const std::vector<nd_int> &order) { arr.transpose(order); })
        .def("transpose", [](librapid::ndarray::ndarray &arr) { arr.transpose(); })
        .def("transposed", [](const librapid::ndarray::ndarray &arr, const std::vector<nd_int> &order) { return arr.transposed(order); })
        .def("transposed", [](librapid::ndarray::ndarray &arr) { return arr.transposed(); })

        .def("__str__", [](const librapid::ndarray::ndarray &arr) { return arr.str(0); })
        .def("__int__", [](const librapid::ndarray::ndarray &arr) { if (!arr.is_scalar()) throw py::value_error("Cannot convert non-scalar array to int"); return (nd_int) *arr.get_data_start(); })
        .def("__float__", [](const librapid::ndarray::ndarray &arr) { if (!arr.is_scalar()) throw py::value_error("Cannot convert non-scalar array to float"); return (double) *arr.get_data_start(); })
        .def("__repr__", [](const librapid::ndarray::ndarray &arr) { return "<librapid.ndarray " + arr.str(18) + ">"; });
}
