From be8158c569f380b6f8d22ad68ff595c22cb9db3a Mon Sep 17 00:00:00 2001
From: froejdh_e <erik.frojdh@psi.ch>
Date: Tue, 22 Apr 2025 15:59:33 +0200
Subject: [PATCH] WIP

---
 python/aare/ClusterFinder.py | 17 +++++++++++++++++
 python/aare/__init__.py      |  2 +-
 python/src/cluster.hpp       |  3 +++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/python/aare/ClusterFinder.py b/python/aare/ClusterFinder.py
index f678dd1..6e7c352 100644
--- a/python/aare/ClusterFinder.py
+++ b/python/aare/ClusterFinder.py
@@ -1,5 +1,8 @@
 
 from ._aare import ClusterFinder_Cluster3x3i, ClusterFinder_Cluster2x2i, ClusterFinderMT_Cluster3x3i, ClusterFinderMT_Cluster2x2i, ClusterCollector_Cluster3x3i, ClusterCollector_Cluster2x2i
+
+
+from ._aare import ClusterFileSink_Cluster3x3i, ClusterFileSink_Cluster2x2i
 import numpy as np
 
 def ClusterFinder(image_size, cluster_size, n_sigma=5, dtype = np.int32, capacity = 1024):
@@ -48,3 +51,17 @@ def ClusterCollector(clusterfindermt, cluster_size = (3,3), dtype=np.int32):
         #TODO! add the other formats
         raise ValueError(f"Unsupported dtype: {dtype}. Only np.int32 is supported.")
 
+def ClusterFileSink(clusterfindermt, cluster_file, dtype=np.int32): 
+    """ 
+    Factory function to create a ClusterCollector object. Provides a cleaner syntax for 
+    the templated ClusterCollector in C++.
+    """
+
+    if dtype == np.int32 and clusterfindermt.cluster_size == (3,3): 
+        return ClusterFileSink_Cluster3x3i(clusterfindermt, cluster_file)
+    elif dtype == np.int32 and clusterfindermt.cluster_size == (2,2):
+        return ClusterFileSink_Cluster2x2i(clusterfindermt, cluster_file)
+     
+    else: 
+        #TODO! add the other formats
+        raise ValueError(f"Unsupported dtype: {dtype}. Only np.int32 is supported.")
\ No newline at end of file
diff --git a/python/aare/__init__.py b/python/aare/__init__.py
index 096dddf..decd560 100644
--- a/python/aare/__init__.py
+++ b/python/aare/__init__.py
@@ -11,7 +11,7 @@ from ._aare import ROI
 
 # from ._aare import ClusterFinderMT, ClusterCollector, ClusterFileSink, ClusterVector_i
 
-from .ClusterFinder import ClusterFinder, ClusterCollector, ClusterFinderMT
+from .ClusterFinder import ClusterFinder, ClusterCollector, ClusterFinderMT, ClusterFileSink
 from .ClusterVector import ClusterVector
 
 
diff --git a/python/src/cluster.hpp b/python/src/cluster.hpp
index 6dd05ad..58f137c 100644
--- a/python/src/cluster.hpp
+++ b/python/src/cluster.hpp
@@ -93,6 +93,9 @@ void define_cluster_finder_mt_bindings(py::module &m,
                 return;
             },
             py::arg(), py::arg("frame_number") = 0)
+        .def_property_readonly("cluster_size", [](ClusterFinderMT<ClusterType, uint16_t, pd_type> &self){
+            return py::make_tuple(ClusterSizeX, ClusterSizeY);
+        })
         .def("clear_pedestal",
              &ClusterFinderMT<ClusterType, uint16_t, pd_type>::clear_pedestal)
         .def("sync", &ClusterFinderMT<ClusterType, uint16_t, pd_type>::sync)