#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include "GeometryDescriptor.hh"
#include "G4System.hh"

namespace py = pybind11;

PYBIND11_MODULE(minicalo, m) {
    py::class_<Sensor>(m, "Sensor")
        .def(py::init<>())
        .def("getEnergy", &Sensor::getEnergy)
        .def("getPos", &Sensor::getPos)
        .def("getSize", &Sensor::getSize)
        .def("getX", &Sensor::getX)
        .def("getY", &Sensor::getY)
        .def("getZ", &Sensor::getZ)
        .def("getdx", &Sensor::getdx)
        .def("getdy", &Sensor::getdy)
        .def("getdz", &Sensor::getdz)
        .def(py::pickle(
            [](const Sensor &s) { // __getstate__
                return s.__getstate__();
            },
            [](py::tuple t) { // __setstate__
                return Sensor::__setstate__(t);
            }
        ));

    py::class_<Layer>(m, "Layer")
        .def(py::init<>())
        .def("setThickness", &Layer::setThickness)
        .def("setMaterial", &Layer::setMaterial)
        .def("setNx", &Layer::setNx)
        .def("setNy", &Layer::setNy)
        .def("setIsActive", &Layer::setIsActive)
        .def("assignPhysicalVolume", &Layer::assignPhysicalVolume)
        .def("unAssign", &Layer::unAssign)
        .def(py::pickle(
            [](const Layer &l) { // __getstate__
                return l.__getstate__();
            },
            [](py::tuple t) { // __setstate__
                return Layer::__setstate__(t);
            }
        ));

    py::class_<GeometryDescriptor>(m, "GeometryDescriptor")
        .def(py::init<>())
        .def("addLayer", &GeometryDescriptor::addLayer)
        .def("getLayers", py::overload_cast<>(&GeometryDescriptor::getLayers))
        .def("getLayers", py::overload_cast<>(&GeometryDescriptor::getLayers, py::const_))
        .def("getXYWidth", &GeometryDescriptor::getXYWidth)
        .def("resetSensorEnergies", &GeometryDescriptor::resetSensorEnergies)
        .def("getNSensors", &GeometryDescriptor::getNSensors)
        .def("printSensorEnergies", &GeometryDescriptor::printSensorEnergies)
        .def("isAssigned", &GeometryDescriptor::isAssigned)
        .def("unAssign", &GeometryDescriptor::unAssign)
        .def("isEmpty", &GeometryDescriptor::isEmpty)
        .def(py::pickle(
            [](const GeometryDescriptor &g) { // __getstate__
                return g.__getstate__();
            },
            [](py::tuple t) { // __setstate__
                return GeometryDescriptor::__setstate__(t);
            }
        ));

    //now for G4System
    py::class_<G4System>(m, "G4System")
        .def(py::init<>())
        .def("init", &G4System::init)
        .def("run_visualize", &G4System::run_visualize)
        .def("run_gui", &G4System::run_gui)
        .def("run_batch", &G4System::run_batch)
        .def("applyUICommand", &G4System::applyUICommand)
        .def("displayEvent", &G4System::displayEvent)
        .def("printMaterial", &G4System::printMaterial);
}