diff --git a/python/digi.py b/python/digi.py
index 8db59a2..fcbe91f 100644
--- a/python/digi.py
+++ b/python/digi.py
@@ -56,8 +56,58 @@ def calculateVoltageHcal(self, PE) :
              Number of photo electrons
         """
         return PE*(5/1)
+
+
+class DigiTimeSpread():
+    '''Type representing possible time smearing/shifting that can be appied
+    during the digitization stage, either per event/spill or per hit
+    '''
+    def __init__(self, kind, parameters):
+        if kind not in [-1, 0, 1, 2]:
+            raise ValueError(
+                "Invalid kind of time spread, must be -1 (No spread), 0 (Gaussian), 1 (Uniform) or 2 (Constant)"
+            )
+        self.kind = kind
+        self.parameters = parameters
+
+
+class NoSpread(DigiTimeSpread):
+    def __init__(self):
+        super().__init__(-1, parameters=[0.])
+
+class GausianSpread(DigiTimeSpread):
+    def __init__(self, mean, sigma):
+        '''Gaussian time spread
+        Parameters:
+        mean : float
+            Mean of the Gaussian distribution
+        sigma : float
+            Standard deviation of the Gaussian distribution
+        '''
+        super().__init__(0, parameters=[mean, sigma])
+
+class UniformSpread(DigiTimeSpread):
+    def __init__(self, min_value, max_value):
+        '''Uniform time spread
+        Parameters:
+        min_value : float
+            Minimum value of the uniform distribution
+        max_value : float
+            Maximum value of the uniform distribution
+        '''
+        super().__init__(1, parameters=[min_value, max_value])
+
+class ConstantSpread(DigiTimeSpread):
+    def __init__(self, value):
+        '''Constant time spread
+        Parameters:
+        value : float
+            Value of the constant time spread
+        '''
+        super().__init__(2, parameters=[value])
     
 class HcalDigiProducer(Producer) :
+
     """Configuration for HcalDigiProducer
 
     Attributes
@@ -102,22 +152,10 @@ def __init__(self, instance_name = 'hcalDigis') :
         # Flat time shift to apply to all hits
         self.flat_time_shift = 0.
 
-        # Apply a timeshift to each event
-        self.do_time_spread_per_spill = False
-        # Apply a timeshift to each hit
-        self.do_time_spread_per_hit = False
-
-        # Type of distribution and parameters to use for time spreads. Currently
-        # implemented are
-        # 0: Gaussian -> parameters[0] = mean, parameters[1] = sigma
-        # 1: Uniform -> parameters[0] = min, parameters[1] = max
-        # 2: Constant -> parameters[0] = value
-        #
-        # Units are same as the time unit of the input hits
-        self.time_spread_per_spill_type = 1
-        self.time_spread_per_hit_type = 0
-        self.time_spread_per_spill_parameters = []
-        self.time_spread_per_hit_parameters = []
+        self.time_spread_per_hit = NoSpread()
+        self.time_spread_per_spill = NoSpread()
+
+
 
 class HcalRecProducer(Producer) :
     """Configuration for the HcalRecProducer
diff --git a/src/Hcal/HcalDigiProducer.cxx b/src/Hcal/HcalDigiProducer.cxx
index bad066d..6b0db1e 100644
--- a/src/Hcal/HcalDigiProducer.cxx
+++ b/src/Hcal/HcalDigiProducer.cxx
@@ -24,38 +24,6 @@ HcalDigiProducer::HcalDigiProducer(const std::string& name,
    */
   noiseGenerator_ = std::make_unique<ldmx::NoiseGenerator>();
 }
-HcalDigiProducer::TimeSpreadType HcalDigiProducer::validateTimeSpreadType(
-    int kind, std::vector<double>& parameters) const {
-  auto type{static_cast<TimeSpreadType>(kind)};
-  switch (type) {
-    case TimeSpreadType::GAUSSIAN:
-      if (parameters.size() != 2) {
-        EXCEPTION_RAISE(
-            "InvalidParameter",
-            "Time spread type Gaussian requires 2 parameters (mean, sigma)");
-      }
-      break;
-    case TimeSpreadType::UNIFORM:
-      if (parameters.size() != 2) {
-        EXCEPTION_RAISE(
-            "InvalidParameter",
-            "Time spread type Uniform requires 2 parameters (min, max)");
-      }
-      break;
-    case TimeSpreadType::CONSTANT:
-      if (parameters.size() != 1) {
-        EXCEPTION_RAISE(
-            "InvalidParameter",
-            "Time spread type Constant requires 1 parameter (value)");
-      }
-      break;
-    default:
-      EXCEPTION_RAISE("InvalidParameter",
-                      "Invalid time spread type, must be 0 (Gaussian), 1 "
-                      "(Uniform), or 2 (Constant))");
-  }
-  return type;
-}
 
 double HcalDigiProducer::makeTimeDelta(TimeSpreadType kind,
                                        std::vector<double>& parameters) const {
@@ -106,24 +74,22 @@ void HcalDigiProducer::configure(framework::config::Parameters& ps) {
   flatTimeShift = ps.getParameter<double>("flat_time_shift");
   // Time spread parameters
   // Per hit
-  doTimeSpreadPerHit = ps.getParameter<bool>("do_time_spread_per_hit");
-  if (doTimeSpreadPerHit) {
-    int timeSpreadType = ps.getParameter<int>("time_spread_per_hit_type");
-    timeSpreadPerHitParameters =
-        ps.getParameter<std::vector<double>>("time_spread_per_hit_parameters");
-    timeSpreadPerHitType =
-        validateTimeSpreadType(timeSpreadType, timeSpreadPerHitParameters);
-  }
-  // Per spill/event
-  doTimeSpreadPerSpill = ps.getParameter<bool>("do_time_spread_per_spill");
-  if (doTimeSpreadPerSpill) {
-    int timeSpreadType = ps.getParameter<int>("time_spread_per_spill_type");
+  const auto& timeSpreadPerHit{framework::config::Parameters};
+  int type =
+      ps.getParameter<int>(timeSpreadPerHitType.getParameter<int>("kind"));
+  timeSpreadPerHitParameters = ps.getParameter<std::vector<double>>(
+      timeSpreadPerHit.getParameter<std::vector<double>>("parameters"));
+ doTimeSpreadPerHit = type != 0;
+ timeSpreadPerHitType = static_cast<TimeSpreadType>(type);
+  // Per spill / event
+  const auto& timeSpreadPerSpill{framework::config::Parameters};
+  int type =
+      ps.getParameter<int>(timeSpreadPerSpillType.getParameter<int>("kind"));
+  timeSpreadPerSpillParameters = ps.getParameter<std::vector<double>>(
+      timeSpreadPerSpill.getParameter<std::vector<double>>("parameters"));
+ doTimeSpreadPerSpill = type != 0;
+ timeSpreadPerSpillType = static_cast<TimeSpreadType>(type);
 
-    timeSpreadPerSpillParameters = ps.getParameter<std::vector<double>>(
-        "time_spread_per_spill_parameters");
-    timeSpreadPerSpillType =
-        validateTimeSpreadType(timeSpreadType, timeSpreadPerSpillParameters);
-  }
 }
 
 void HcalDigiProducer::produce(framework::Event& event) {
@@ -294,7 +260,8 @@ void HcalDigiProducer::produce(framework::Event& event) {
         time -= position.at(2) /
                 299.702547;  // shift light-speed particle traveling along z
         if (doTimeSpreadPerHit) {
-          time += makeTimeDelta(timeSpreadPerHitType, timeSpreadPerHitParameters);
+          time +=
+              makeTimeDelta(timeSpreadPerHitType, timeSpreadPerHitParameters);
         }
         time += timeDelta;
         if (end_close == 0) {