From 1e02f22908f055daabbf8fe074067cc664c237ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CAidan?= <“abeiku17@gmail.com”>
Date: Wed, 10 Jul 2024 10:34:13 +0100
Subject: [PATCH 1/3] Get start and end times from csv if not found in info

---
 src/asleep/get_sleep.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/asleep/get_sleep.py b/src/asleep/get_sleep.py
index 675550f..6762408 100644
--- a/src/asleep/get_sleep.py
+++ b/src/asleep/get_sleep.py
@@ -35,6 +35,8 @@
 NON_WEAR_THRESHOLD = 3  # H
 NON_WEAR_PREDICTION_FLAG = -1
 
+START_TIME_IDX = 0
+END_TIME_IDX = -1
 
 def load_model(model_path, force_download=False):
     """ Load trained model. Download if not exists. """
@@ -69,8 +71,8 @@ def get_parsed_data(raw_data_path, info_data_path, resample_hz, args):
         data = data.reset_index()
 
         # apply time shift
-        start_time = pd.to_datetime(info['StartTime'])
-        end_time = pd.to_datetime(info['EndTime'])
+        start_time = pd.to_datetime(info.get("StartTime", data.index[START_TIME_IDX]))
+        end_time = pd.to_datetime(info.get("EndTime", data.index[END_TIME_IDX]))
         info['StartTime'] = start_time + datetime.timedelta(hours=time_shift)
         info['EndTime'] = end_time + datetime.timedelta(hours=time_shift)
         info['StartTime'] = info['StartTime'].strftime('%Y-%m-%d %H:%M:%S')

From 72fc96973201c171d3c1260062804ab8e39e6b9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CAidan?= <“abeiku17@gmail.com”>
Date: Wed, 10 Jul 2024 12:47:19 +0100
Subject: [PATCH 2/3] Update extracted column

---
 src/asleep/get_sleep.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/asleep/get_sleep.py b/src/asleep/get_sleep.py
index 6762408..a52ef92 100644
--- a/src/asleep/get_sleep.py
+++ b/src/asleep/get_sleep.py
@@ -71,8 +71,8 @@ def get_parsed_data(raw_data_path, info_data_path, resample_hz, args):
         data = data.reset_index()
 
         # apply time shift
-        start_time = pd.to_datetime(info.get("StartTime", data.index[START_TIME_IDX]))
-        end_time = pd.to_datetime(info.get("EndTime", data.index[END_TIME_IDX]))
+        start_time = pd.to_datetime(info.get("StartTime", data["time"].iloc[START_TIME_IDX]))
+        end_time = pd.to_datetime(info.get("EndTime", data["time"].iloc[END_TIME_IDX]))
         info['StartTime'] = start_time + datetime.timedelta(hours=time_shift)
         info['EndTime'] = end_time + datetime.timedelta(hours=time_shift)
         info['StartTime'] = info['StartTime'].strftime('%Y-%m-%d %H:%M:%S')

From 3ee232900a22115fa9b23e8bfbf8b2e5e21ca80d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CAidan?= <“abeiku17@gmail.com”>
Date: Thu, 11 Jul 2024 11:22:29 +0100
Subject: [PATCH 3/3] Add extra line to pass tests

---
 src/asleep/get_sleep.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/asleep/get_sleep.py b/src/asleep/get_sleep.py
index a52ef92..646115e 100644
--- a/src/asleep/get_sleep.py
+++ b/src/asleep/get_sleep.py
@@ -38,6 +38,7 @@
 START_TIME_IDX = 0
 END_TIME_IDX = -1
 
+
 def load_model(model_path, force_download=False):
     """ Load trained model. Download if not exists. """