From d8759c4191391342d31f7c8264ff46e8cea278a4 Mon Sep 17 00:00:00 2001
From: Camilo Diaz <camilo_diaz@brown.edu>
Date: Fri, 3 May 2024 09:52:26 -0400
Subject: [PATCH] ENH: Added feature to load single .csv files from the gui

---
 hnn_core/__init__.py |  2 +-
 hnn_core/dipole.py   | 18 ++++++++++--------
 hnn_core/gui/gui.py  | 11 ++++++++---
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/hnn_core/__init__.py b/hnn_core/__init__.py
index 89e430a63c..40bf1bbbfb 100644
--- a/hnn_core/__init__.py
+++ b/hnn_core/__init__.py
@@ -1,4 +1,4 @@
-from .dipole import simulate_dipole, read_dipole, average_dipoles, Dipole
+from .dipole import simulate_dipole, read_dipole, average_dipoles, Dipole,_read_dipole_txt
 from .params import Params, read_params, convert_to_hdf5
 from .network import Network, pick_connection
 from .network_models import jones_2009_model, law_2021_model, calcium_model
diff --git a/hnn_core/dipole.py b/hnn_core/dipole.py
index 28fe515a47..3a2a329031 100644
--- a/hnn_core/dipole.py
+++ b/hnn_core/dipole.py
@@ -106,20 +106,26 @@ def simulate_dipole(net, tstop, dt=0.025, n_trials=None, record_vsec=False,
     return dpls
 
 
-def _read_dipole_txt(fname):
+def _read_dipole_txt(fname, extension='.txt'):
     """Read dipole values from a txt file and create a Dipole instance.
 
     Parameters
     ----------
     fname : str
-        Full path to the input file (.txt)
-
+        Full path to the input file (.txt or .csv)
+    fname: io.StringIO
+        Content of file in memory as a StringIO
     Returns
     -------
     dpl : Dipole
         The instance of Dipole class
     """
-    dpl_data = np.loadtxt(fname, dtype=float)
+    if extension == '.csv':
+        # read from a csv file ignoring the headers
+        dpl_data = np.genfromtxt(fname, delimiter=',',
+                                 skip_header=1, dtype=float)
+    else:
+        dpl_data = np.loadtxt(fname, dtype=float)
     ncols = dpl_data.shape[1]
     if ncols not in (2, 4):
         raise ValueError(
@@ -174,10 +180,6 @@ def read_dipole(fname):
         The instance of Dipole class
     """
 
-    # For supporting tests in test_gui.py
-    if isinstance(fname, StringIO):
-        return _read_dipole_txt(fname)
-
     fname = str(fname)
     if not os.path.exists(fname):
         raise FileNotFoundError('File not found at path %s.' % (fname,))
diff --git a/hnn_core/gui/gui.py b/hnn_core/gui/gui.py
index 4e39b06a70..1a5a29bcd1 100644
--- a/hnn_core/gui/gui.py
+++ b/hnn_core/gui/gui.py
@@ -234,7 +234,7 @@ def __init__(self, theme_color="#8A2BE2",
                                             description='Cores:',
                                             disabled=False)
         self.load_data_button = FileUpload(
-            accept='.txt', multiple=False,
+            accept='.txt,.csv', multiple=False,
             style={'button_color': self.layout['theme_color']},
             description='Load data',
             button_style='success')
@@ -1211,7 +1211,12 @@ def on_upload_data_change(change, data, viz_manager, log_out):
 
     data_dict = change['new'][0]
 
-    data_fname = data_dict['name'].rstrip('.txt')
+    dict_name = data_dict['name'].rsplit('.', 1)
+    data_fname = dict_name[0]
+    file_extension = f".{dict_name[1]}"
+
+    # data_fname = data_dict['name'].rstrip('.txt')
+
     if data_fname in data['simulation_data'].keys():
         logger.error(f"Found existing data: {data_fname}.")
         return
@@ -1220,7 +1225,7 @@ def on_upload_data_change(change, data, viz_manager, log_out):
     ext_content = codecs.decode(ext_content, encoding="utf-8")
     with log_out:
         data['simulation_data'][data_fname] = {'net': None, 'dpls': [
-            hnn_core.read_dipole(io.StringIO(ext_content))
+            hnn_core._read_dipole_txt(io.StringIO(ext_content), file_extension)
         ]}
         logger.info(f'External data {data_fname} loaded.')
         _template_name = "[Blank] single figure"