From 46e229fa5acf5aeed305e3d5d434815dbbd7f788 Mon Sep 17 00:00:00 2001
From: julia-pfarr <pfarr@staff.uni-marburg.de>
Date: Fri, 27 Oct 2023 12:45:26 +0200
Subject: [PATCH] include code for extracting eyetrack.tsv samples table

---
 eye2bids/edf2bids.py | 77 +++++++++++++++++++++++++++++++++++---------
 1 file changed, 61 insertions(+), 16 deletions(-)

diff --git a/eye2bids/edf2bids.py b/eye2bids/edf2bids.py
index 44d83f9..869ce16 100644
--- a/eye2bids/edf2bids.py
+++ b/eye2bids/edf2bids.py
@@ -97,12 +97,19 @@ def _check_edf2asc_present() -> bool:
         return False
 
 
-def _convert_edf_to_asc(input_file: str | Path) -> Path:
-    """Convert edf to asc."""
-    subprocess.run(["edf2asc", "-y", "-e", input_file])
-    return Path(input_file).with_suffix(".asc")
+def _convert_edf_to_asc_events(input_file: str | Path) -> Path:
+    """Convert edf to asc - events."""
+    events_asc_file = (input_file.parent) / (input_file.stem + "_events")
+    subprocess.run(["edf2asc", "-y", "-e", input_file, "-o", events_asc_file])
+    return Path(events_asc_file).with_suffix(".asc")
 
 
+def _convert_edf_to_asc_samples(input_file: str | Path) -> Path:
+    """Convert edf to asc - samples."""
+    samples_asc_file = (input_file.parent) / (input_file.stem + "_samples")
+    subprocess.run(["edf2asc", "-y", "-s", input_file, "-o", samples_asc_file])
+    return Path(samples_asc_file).with_suffix(".asc")
+
 def _calibrations(df: pd.DataFrame) -> pd.DataFrame:
     return df[df[3] == "CALIBRATION"]
 
@@ -271,23 +278,43 @@ def _extract_StopTime(events: list[str]) -> int:
     return StopTime
 
 
-def _load_asc_file(asc_file: str | Path) -> list[str]:
-    with open(asc_file) as f:
+def _load_asc_file(events_asc_file: str | Path) -> list[str]:
+    with open(events_asc_file) as f:
         return f.readlines()
 
 
-def _load_asc_file_as_df(asc_file: str | Path) -> pd.DataFrame:
+def _load_asc_file_as_df(events_asc_file: str | Path) -> pd.DataFrame:
     # dataframe for events, all
-    events = _load_asc_file(asc_file)
+    events = _load_asc_file(events_asc_file)
     return pd.DataFrame([ms.split() for ms in events if ms.startswith("MSG")])
 
 
-def _load_asc_file_as_reduced_df(asc_file: str | Path) -> pd.DataFrame:
+def _load_asc_file_as_reduced_df(events_asc_file: str | Path) -> pd.DataFrame:
     # reduced dataframe without MSG and sample columns
-    df_ms = _load_asc_file_as_df(asc_file)
+    df_ms = _load_asc_file_as_df(events_asc_file)
     return pd.DataFrame(df_ms.iloc[0:, 2:])
 
 
+def _samples_to_data_frame(samples_asc_file:str | Path) -> pd.DataFrame:
+
+    column_names = ['eye_timestamp',
+                    'eye1_x_coordinate', 'eye1_y_coordinate', 'eye1_pupil_size',
+                    'eye2_x_coordinate', 'eye2_y_coordinate', 'eye2_pupil_size']
+    
+    data = {name: [] for name in column_names}
+
+    with open(samples_asc_file, "r") as file:
+        for line in file:
+            columns = line.strip().split("\t")
+            for i in range(len(column_names)):
+                if i < len(columns):
+                    data[column_names[i]].append(columns[i])
+
+    data = {key: value for key, value in data.items() if any(val not in ("", "...") for val in value)}
+
+    return pd.DataFrame(data)
+
+
 def edf2bids(
     input_file: str | Path | None = None,
     metadata_file: str | Path | None = None,
@@ -303,11 +330,12 @@ def edf2bids(
     )
 
     # CONVERSION events
-    asc_file = _convert_edf_to_asc(input_file)
+    events_asc_file = _convert_edf_to_asc_events(input_file)
+    samples_asc_file = _convert_edf_to_asc_samples(input_file)
 
-    events = _load_asc_file(asc_file)
-    df_ms = _load_asc_file_as_df(asc_file)
-    df_ms_reduced = _load_asc_file_as_reduced_df(asc_file)
+    events = _load_asc_file(events_asc_file)
+    df_ms = _load_asc_file_as_df(events_asc_file)
+    df_ms_reduced = _load_asc_file_as_reduced_df(events_asc_file)
 
     if metadata_file is None:
         metadata = {}
@@ -315,12 +343,14 @@ def edf2bids(
         with open(metadata_file) as f:
             metadata = yaml.load(f, Loader=SafeLoader)
 
-    # to json
+    # file naming check    
 
     filename = os.path.splitext(input_file)[0]
     substring_eyetrack = '_eyetrack'
     substring_events = '_events'
 
+    # to json
+    # Eyetrack.json Metadata
 
     eyetrack_json = {
         "Manufacturer": "SR-Research",
@@ -381,6 +411,21 @@ def edf2bids(
         with open(output_dir / (filename + ".json"), "w") as outfile:
             json.dump(events_json, outfile, indent=4)
         e2b_log.info(f"file generated: {output_dir / (filename + '.json')}")
+   
+   
+   # Samples to eyetrack.tsv
+    eyetrack_tsv = _samples_to_data_frame(samples_asc_file)
+
+    if substring_eyetrack not in filename:
+        with open(output_dir / (filename + "_eyetrack.tsv"), "w") as outfile:
+            eyetrack_tsv.to_csv(outfile, sep='\t', index=False, compression='gzip')
+        e2b_log.info(f"file generated: {output_dir / (filename + '_eyetrack.tsv')}")
+    elif substring_eyetrack in filename: 
+        with open(output_dir / (filename + ".tsv"), "w") as outfile:
+            eyetrack_tsv.to_csv(outfile, sep='\t', index=False, compression='gzip')
+        e2b_log.info(f"file generated: {output_dir / (filename + '.tsv')}")
+
+
 
 if __name__ == "__main__":
     parser = global_parser()
@@ -389,4 +434,4 @@ def edf2bids(
         input_file=args.input_file,
         metadata_file=args.metadata_file,
         output_dir=args.output_dir,
-    )
+    )
\ No newline at end of file