From 216be7b1e5cafe5af24214835e660551169337c1 Mon Sep 17 00:00:00 2001
From: Jan Grewe <jan.grewe@g-node.org>
Date: Tue, 28 Jan 2025 18:00:04 +0100
Subject: [PATCH] [tracks] writing results works

---
 fixtracks/utils/signals.py  |  5 +++++
 fixtracks/utils/writer.py   |  1 -
 fixtracks/widgets/tracks.py | 36 +++++++++++++++++++++++++++++++-----
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/fixtracks/utils/signals.py b/fixtracks/utils/signals.py
index 7f0e8d7..9275d6f 100644
--- a/fixtracks/utils/signals.py
+++ b/fixtracks/utils/signals.py
@@ -8,6 +8,11 @@ class ProducerSignals(QObject):
     progress = Signal(float)
     progress2  = Signal((str, float, float))
 
+class WriterSignals(QObject):
+    error = Signal(str)
+    finished = Signal(bool)
+
+
 class DetectionViewSignals(QObject):
     trackClicked = Signal(int)
     itemsSelected = Signal(list)
diff --git a/fixtracks/utils/writer.py b/fixtracks/utils/writer.py
index b775d1c..0c6ce8b 100644
--- a/fixtracks/utils/writer.py
+++ b/fixtracks/utils/writer.py
@@ -1,4 +1,3 @@
-import pickle
 import logging
 from PySide6.QtCore import QRunnable, Slot
 
diff --git a/fixtracks/widgets/tracks.py b/fixtracks/widgets/tracks.py
index 5c356b2..8472357 100644
--- a/fixtracks/widgets/tracks.py
+++ b/fixtracks/widgets/tracks.py
@@ -1,14 +1,16 @@
 import logging
 import pathlib
+import pickle
 import numpy as np
 import pandas as pd
 
-from PySide6.QtCore import Qt, QThreadPool, Signal, QAbstractTableModel, QSortFilterProxyModel, QItemSelectionModel, QSize, QObject
+from PySide6.QtCore import Qt, QThreadPool, Signal, QAbstractTableModel, QSortFilterProxyModel, QSize, QObject
 from PySide6.QtGui import QImage, QBrush, QColor, QFont
 from PySide6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QSizePolicy, QComboBox
-from PySide6.QtWidgets import QSpinBox, QSpacerItem, QProgressBar, QTableView, QSplitter, QGridLayout
+from PySide6.QtWidgets import QSpinBox, QSpacerItem, QProgressBar, QTableView, QSplitter, QGridLayout, QFileDialog
 
 from fixtracks.utils.reader import PickleLoader
+from fixtracks.utils.writer import PickleWriter
 from fixtracks.widgets.detectionview import DetectionView
 from fixtracks.widgets.detectiontimeline import DetectionTimeline
 
@@ -253,7 +255,12 @@ class DataController(QObject):
         self._data["track"][self._user_selections] = track_id
 
     def save(self, filename):
-        pass
+        export_columns = self._columns.copy()
+        export_columns.remove("index")
+        dictionary = {c: self._data[c] for c in export_columns}
+        df = pd.DataFrame(dictionary, index=self._data["index"])
+        with open(filename, 'wb') as f:
+            pickle.dump(df, f)
 
 
 class FixTracks(QWidget):
@@ -345,9 +352,11 @@ class FixTracks(QWidget):
         # tablebox.addLayout(tracktwo_box)
 
         self._saveBtn = QPushButton("Save")
+        self._saveBtn.setShortcut("Ctrl+S")
         self._saveBtn.setEnabled(False)
         self._saveBtn.clicked.connect(self.on_save)
         self._backBtn = QPushButton("Back")
+        self._backBtn.setShortcut("ESC")
         self._backBtn.clicked.connect(self.on_back)
 
         self._data_combo = QComboBox()
@@ -472,9 +481,27 @@ class FixTracks(QWidget):
             self._data.setData(self._reader.asdict)
             self._timeline.setDetectionData(self._data.data)
             self.update()
+            self._saveBtn.setEnabled(True)
 
     def on_save(self):
-        logging.debug("Save fixtracks results")
+        logging.debug("Saving fixtracks results")
+        self._tasklabel.setText("Saving results to file...")
+        file_dialog = QFileDialog(self)
+        file_dialog.setAcceptMode(QFileDialog.AcceptMode.AcceptSave)
+        file_dialog.setNameFilter("Pickle Files (*.pkl)")
+        if file_dialog.exec():
+            file_path = file_dialog.selectedFiles()[0]
+            if not file_path.endswith(".pkl"):
+                file_path += ".pkl"
+            self._progress_bar.setRange(0,0)
+            save_task = PickleWriter(self._data, file_path)
+            save_task.signals.finished.connect(self.on_dataSaved)
+            self._threadpool.start(save_task)
+
+    def on_dataSaved(self):
+        self._tasklabel.setText("")
+        self._progress_bar.setRange(0, 100)
+        self._progress_bar.setValue(0)
 
     def on_back(self):
         logging.debug("Back button pressed!")
@@ -519,7 +546,6 @@ class FixTracks(QWidget):
 
 def main():
     from PySide6.QtWidgets import QApplication
-
     app = QApplication([])
     window = QWidget()
     window.setMinimumSize(200, 200)