minicalosim/include/GeometryDescriptor.hh

#ifndef GeometryDescriptor_HH
#define GeometryDescriptor_HH

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

#include "G4ThreeVector.hh"
#include <vector>
#include <string>
#include "G4VPhysicalVolume.hh"
class G4System;

class Sensor {
public:
    Sensor() {};
    ~Sensor() {};

    const G4double getEnergy() const {
        return energy;
    }
    G4ThreeVector getPos() const {
        return position;
    }
    G4ThreeVector getSize() const {
        return size;
    }

    double getX() const {
        return position.x();
    }

    double getY() const {
        return position.y();
    }

    double getZ() const {
        return position.z();
    }

    double getdx() const {
        return size.x();
    }

    double getdy() const {
        return size.y();
    }

    double getdz() const {
        return size.z();
    }

    G4ThreeVector position;
    G4ThreeVector size;
    mutable G4double energy;

    pybind11::tuple __getstate__() const {
        return pybind11::make_tuple(position.x(), position.y(), position.z(), size.x(), size.y(), size.z(), energy);
    }

    static Sensor __setstate__(pybind11::tuple t) {
        if (t.size() != 7) throw std::runtime_error("Invalid state!");
        Sensor sensor;
        sensor.position = G4ThreeVector(t[0].cast<double>(), t[1].cast<double>(), t[2].cast<double>());
        sensor.size = G4ThreeVector(t[3].cast<double>(), t[4].cast<double>(), t[5].cast<double>());
        sensor.energy = t[6].cast<G4double>();
        return sensor;
    }
};

class Layer {
public:
    Layer() : thickness(0),sens_xwidth(0), sens_ywidth(0),  material(""), nx(1), ny(1), isActive(false), physicalVolume(nullptr), position(0,0,0) {};
    ~Layer() {};

    void setThickness(double thickness_cm);
    void setMaterial(std::string material);
    void setNx(int nx);
    void setNy(int ny);
    void setIsActive(bool isActive);
    void assignPhysicalVolume(G4VPhysicalVolume* physicalVolume);

    void unAssign() {
        physicalVolume = nullptr;
        sensors.clear();
    }

    double getX() const {
        return position.x();
    }

    double getY() const {
        return position.y();
    }

    double getZ() const {
        return position.z();
    }

    double thickness;
    double sens_xwidth;
    double sens_ywidth;
    std::string material;
    int nx;
    int ny;
    bool isActive;


    G4VPhysicalVolume* physicalVolume;
    std::string name;

    G4ThreeVector position;
    std::vector<Sensor> sensors;

    pybind11::tuple __getstate__() const {
        return pybind11::make_tuple(thickness, sens_xwidth, sens_ywidth, material, nx, ny, isActive, position.x(), position.y(), position.z(), sensors);
    }

    static Layer __setstate__(pybind11::tuple t) {
        if (t.size() != 11) throw std::runtime_error("Invalid state!");
        Layer layer;
        layer.thickness = t[0].cast<double>();
        layer.sens_xwidth = t[1].cast<double>();
        layer.sens_ywidth = t[2].cast<double>();
        layer.material = t[3].cast<std::string>();
        layer.nx = t[4].cast<int>();
        layer.ny = t[5].cast<int>();
        layer.isActive = t[6].cast<bool>();
        layer.position = G4ThreeVector(t[7].cast<double>(), t[8].cast<double>(), t[9].cast<double>());
        layer.sensors = t[10].cast<std::vector<Sensor>>();
        layer.physicalVolume = nullptr; // Reset pointer
        return layer;
    }
};

class GeometryDescriptor {
public:
    GeometryDescriptor() : xywidth(200), g4system(nullptr) {};
    ~GeometryDescriptor();

    void addLayer(double thickness_cm, std::string material, bool isActive = true, int nx = 1, int ny = -1);

    std::vector<Layer>& getLayers() {
        return layers;
    }
    const std::vector<Layer>& getLayers() const {
        return layers;
    }

    double getXYWidth() const {
        return xywidth;
    }

    void resetSensorEnergies() const; // energies are mutable
    int getNSensors() const {
        int n_sensors = 0;
        for (const auto& layer : layers) {
            n_sensors += layer.sensors.size();
        }
        return n_sensors;
    }

    Layer* getLayerByVolume(G4VPhysicalVolume* volume);
    const Layer* getLayerByVolume(G4VPhysicalVolume* volume) const;
    Sensor* getSensorByVolume(G4VPhysicalVolume* volume);
    const Sensor* getSensorByVolume(G4VPhysicalVolume* volume) const;
    void printSensorEnergies() const;

    bool isAssigned() const {
        if (layers.empty()) {
            return false;
        } else {
            return layers[0].physicalVolume != nullptr;
        }
    }

    void unAssign() {
        for (auto& layer : layers) {
            layer.unAssign();
        }
    }

    bool isEmpty() const {
        return layers.empty();
    }

    void setG4System(G4System* g4system) {
        this->g4system = g4system;
    }

    pybind11::tuple __getstate__() const {
        pybind11::list layer_list;
        for (const auto& layer : layers) {
            layer_list.append(layer.__getstate__());
        }
        return pybind11::make_tuple(xywidth, layer_list);
    }

    static GeometryDescriptor __setstate__(pybind11::tuple t) {
        if (t.size() != 2) throw std::runtime_error("Invalid state!");
        GeometryDescriptor geom;
        geom.xywidth = t[0].cast<double>();

        pybind11::list layer_list = t[1].cast<pybind11::list>();
        for (auto item : layer_list) {
            geom.layers.push_back(Layer::__setstate__(item.cast<pybind11::tuple>()));
        }
        geom.g4system = nullptr; // Reset pointer
        return geom;
    }

//private:
    double xywidth;
    std::vector<Layer> layers;
    G4System* g4system;
};

#endif