From b0f411017b646e421be11fefd010e796b35d73c9 Mon Sep 17 00:00:00 2001
From: maestroque <kikasgeorg@ece.auth.gr>
Date: Sun, 11 Aug 2024 22:06:47 +0300
Subject: [PATCH] Add Metric class for easier properties accessing

---
 phys2denoise/metrics/utils.py              | 62 +++++++++++++++++++++-
 phys2denoise/tests/test_metrics_cardiac.py |  8 ++-
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/phys2denoise/metrics/utils.py b/phys2denoise/metrics/utils.py
index edd66ec..eab3e7a 100644
--- a/phys2denoise/metrics/utils.py
+++ b/phys2denoise/metrics/utils.py
@@ -353,13 +353,29 @@ def return_physio_or_metric(*, return_physio=True):
     """
 
     def determine_return_type(func):
+        metrics_with_lags = ["respiratory_variance_time"]
+        convolved_metrics = [
+            "respiratory_variance",
+            "heart_rate",
+            "heart_rate_variability",
+            "heart_beat_interval",
+        ]
+
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
             physio, metric = func(*args, **kwargs)
             if isinstance(args[0], Physio):
-                physio._computed_metrics[func.__name__] = dict(
-                    metric=metric, args=kwargs
+                has_lags = True if func.__name__ in metrics_with_lags else False
+                is_convolved = True if func.__name__ in convolved_metrics else False
+
+                physio._computed_metrics[func.__name__] = Metric(
+                    func.__name__,
+                    metric,
+                    kwargs,
+                    has_lags=has_lags,
+                    is_convolved=is_convolved,
                 )
+
                 return_physio_value = kwargs.get("return_physio", return_physio)
                 if return_physio_value:
                     return physio
@@ -371,3 +387,45 @@ def wrapper(*args, **kwargs):
         return wrapper
 
     return determine_return_type
+
+
+class Metric:
+    def __init__(self, name, data, args, has_lags=False, is_convolved=False):
+        self.name = name
+        self._data = data
+        self._args = args
+        self._has_lags = has_lags
+        self._is_convolved = is_convolved
+
+    def __array__(self):
+        return self.data
+
+    def __getitem__(self, slicer):
+        return self.data[slicer]
+
+    def __len__(self):
+        return len(self.data)
+
+    @property
+    def ndim(self):
+        return self.data.ndim
+
+    @property
+    def shape(self):
+        return self.data.shape
+
+    @property
+    def data(self):
+        return self._data
+
+    @property
+    def args(self):
+        return self._args
+
+    @property
+    def has_lags(self):
+        return self._has_lags
+
+    @property
+    def is_convolved(self):
+        return self._is_convolved
diff --git a/phys2denoise/tests/test_metrics_cardiac.py b/phys2denoise/tests/test_metrics_cardiac.py
index 839b040..2d776cf 100644
--- a/phys2denoise/tests/test_metrics_cardiac.py
+++ b/phys2denoise/tests/test_metrics_cardiac.py
@@ -58,12 +58,16 @@ def test_cardiac_phase_smoke_physio_obj():
         n_scans=n_scans,
         t_r=t_r,
     )
+
     assert phys.history[0][0] == "phys2denoise.metrics.cardiac.cardiac_phase"
-    assert phys.computed_metrics["cardiac_phase"]["metric"].ndim == 2
-    assert phys.computed_metrics["cardiac_phase"]["metric"].shape == (
+    assert phys.computed_metrics["cardiac_phase"].ndim == 2
+    assert phys.computed_metrics["cardiac_phase"].shape == (
         n_scans,
         slice_timings.size,
     )
+    assert phys.computed_metrics["cardiac_phase"].args["slice_timings"] is not None
+    assert phys.computed_metrics["cardiac_phase"].args["n_scans"] is not None
+    assert phys.computed_metrics["cardiac_phase"].args["t_r"] is not None
 
     # Test where the metric is returned
     card_phase = cardiac.cardiac_phase(