1import csv
2
3import os
4import numpy as np
5from PyQt5.QtCore import Qt, pyqtSlot, pyqtSignal
6from PyQt5.QtWidgets import QDialog, QInputDialog, QApplication, QCompleter, QDirModel, QFileDialog
7
8from urh.ui.ui_csv_wizard import Ui_DialogCSVImport
9from urh.util import FileOperator, util
10from urh.util.Errors import Errors
11
12
13class CSVImportDialog(QDialog):
14    data_imported = pyqtSignal(str, float)  # Complex Filename + Sample Rate
15
16
17    PREVIEW_ROWS = 100
18    COLUMNS = {"T": 0, "I": 1, "Q": 2}
19
20    def __init__(self, filename="", parent=None):
21        super().__init__(parent)
22        self.ui = Ui_DialogCSVImport()
23        self.ui.setupUi(self)
24        self.setAttribute(Qt.WA_DeleteOnClose)
25        self.setWindowFlags(Qt.Window)
26
27        self.ui.btnAutoDefault.hide()
28
29        completer = QCompleter()
30        completer.setModel(QDirModel(completer))
31        self.ui.lineEditFilename.setCompleter(completer)
32
33        self.filename = None  # type: str
34        self.ui.lineEditFilename.setText(filename)
35        self.update_file()
36
37        self.ui.tableWidgetPreview.setColumnHidden(self.COLUMNS["T"], True)
38        self.update_preview()
39
40        self.create_connects()
41
42    def create_connects(self):
43        self.accepted.connect(self.on_accepted)
44        self.ui.lineEditFilename.editingFinished.connect(self.on_line_edit_filename_editing_finished)
45        self.ui.btnChooseFile.clicked.connect(self.on_btn_choose_file_clicked)
46        self.ui.btnAddSeparator.clicked.connect(self.on_btn_add_separator_clicked)
47        self.ui.comboBoxCSVSeparator.currentIndexChanged.connect(self.on_combobox_csv_separator_current_index_changed)
48        self.ui.spinBoxIDataColumn.valueChanged.connect(self.on_spinbox_i_data_column_value_changed)
49        self.ui.spinBoxQDataColumn.valueChanged.connect(self.on_spinbox_q_data_column_value_changed)
50        self.ui.spinBoxTimestampColumn.valueChanged.connect(self.on_spinbox_timestamp_value_changed)
51
52    def update_file(self):
53        filename = self.ui.lineEditFilename.text()
54        self.filename = filename
55
56        enable = util.file_can_be_opened(filename)
57        if enable:
58            with open(self.filename, encoding="utf-8-sig") as f:
59                lines = []
60                for i, line in enumerate(f):
61                    if i >= self.PREVIEW_ROWS:
62                        break
63                    lines.append(line.strip())
64                self.ui.plainTextEditFilePreview.setPlainText("\n".join(lines))
65        else:
66            self.ui.plainTextEditFilePreview.clear()
67
68        self.ui.plainTextEditFilePreview.setEnabled(enable)
69        self.ui.comboBoxCSVSeparator.setEnabled(enable)
70        self.ui.spinBoxIDataColumn.setEnabled(enable)
71        self.ui.spinBoxQDataColumn.setEnabled(enable)
72        self.ui.spinBoxTimestampColumn.setEnabled(enable)
73        self.ui.tableWidgetPreview.setEnabled(enable)
74        self.ui.labelFileNotFound.setVisible(not enable)
75
76    def update_preview(self):
77        if not util.file_can_be_opened(self.filename):
78            self.update_file()
79            return
80
81        i_data_col = self.ui.spinBoxIDataColumn.value() - 1
82        q_data_col = self.ui.spinBoxQDataColumn.value() - 1
83        timestamp_col = self.ui.spinBoxTimestampColumn.value() - 1
84
85        self.ui.tableWidgetPreview.setRowCount(self.PREVIEW_ROWS)
86
87        with open(self.filename, encoding="utf-8-sig") as f:
88            csv_reader = csv.reader(f, delimiter=self.ui.comboBoxCSVSeparator.currentText())
89            row = -1
90
91            for line in csv_reader:
92                row += 1
93                result = self.parse_csv_line(line, i_data_col, q_data_col, timestamp_col)
94                if result is not None:
95                    for key, value in result.items():
96                        self.ui.tableWidgetPreview.setItem(row, self.COLUMNS[key], util.create_table_item(value))
97                else:
98                    for col in self.COLUMNS.values():
99                        self.ui.tableWidgetPreview.setItem(row, col, util.create_table_item("Invalid"))
100
101                if row >= self.PREVIEW_ROWS - 1:
102                    break
103
104            self.ui.tableWidgetPreview.setRowCount(row + 1)
105
106    @staticmethod
107    def parse_csv_line(csv_line: str, i_data_col: int, q_data_col: int, timestamp_col: int):
108        result = dict()
109
110        if i_data_col >= 0:
111            try:
112                result["I"] = float(csv_line[i_data_col])
113            except:
114                return None
115        else:
116            result["I"] = 0.0
117
118        if q_data_col >= 0:
119            try:
120                result["Q"] = float(csv_line[q_data_col])
121            except:
122                return None
123        else:
124            result["Q"] = 0.0
125
126        if timestamp_col >= 0:
127            try:
128                result["T"] = float(csv_line[timestamp_col])
129            except:
130                return None
131
132        return result
133
134    @staticmethod
135    def parse_csv_file(filename: str, separator: str, i_data_col: int, q_data_col=-1, t_data_col=-1):
136        iq_data = []
137        timestamps = [] if t_data_col > -1 else None
138        with open(filename, encoding="utf-8-sig") as f:
139            csv_reader = csv.reader(f, delimiter=separator)
140            for line in csv_reader:
141                parsed = CSVImportDialog.parse_csv_line(line, i_data_col, q_data_col, t_data_col)
142                if parsed is None:
143                    continue
144
145                iq_data.append(complex(parsed["I"], parsed["Q"]))
146                if timestamps is not None:
147                    timestamps.append(parsed["T"])
148
149        iq_data = np.asarray(iq_data, dtype=np.complex64)
150        sample_rate = CSVImportDialog.estimate_sample_rate(timestamps)
151        return iq_data / abs(iq_data.max()), sample_rate
152
153    @staticmethod
154    def estimate_sample_rate(timestamps):
155        if timestamps is None or len(timestamps) < 2:
156            return None
157
158        previous_timestamp = timestamps[0]
159        durations = []
160
161        for timestamp in timestamps[1:CSVImportDialog.PREVIEW_ROWS]:
162            durations.append(abs(timestamp-previous_timestamp))
163            previous_timestamp = timestamp
164
165        return 1 / (sum(durations) / len(durations))
166
167    @pyqtSlot()
168    def on_line_edit_filename_editing_finished(self):
169        self.update_file()
170        self.update_preview()
171
172    @pyqtSlot()
173    def on_btn_choose_file_clicked(self):
174        filename, _ = QFileDialog.getOpenFileName(self, self.tr("Choose file"), directory=FileOperator.RECENT_PATH,
175                                                  filter="CSV files (*.csv);;All files (*.*)")
176
177        if filename:
178            self.ui.lineEditFilename.setText(filename)
179            self.ui.lineEditFilename.editingFinished.emit()
180
181    @pyqtSlot()
182    def on_btn_add_separator_clicked(self):
183        sep, ok = QInputDialog.getText(self, "Enter Separator", "Separator:", text=",")
184        if ok and sep not in (self.ui.comboBoxCSVSeparator.itemText(i) for i in
185                              range(self.ui.comboBoxCSVSeparator.count())):
186            if len(sep) == 1:
187                self.ui.comboBoxCSVSeparator.addItem(sep)
188            else:
189                Errors.generic_error("Invalid Separator", "Separator must be exactly one character.")
190
191    @pyqtSlot(int)
192    def on_combobox_csv_separator_current_index_changed(self, index: int):
193        self.update_preview()
194
195    @pyqtSlot(int)
196    def on_spinbox_i_data_column_value_changed(self, value: int):
197        self.update_preview()
198
199    @pyqtSlot(int)
200    def on_spinbox_q_data_column_value_changed(self, value: int):
201        self.update_preview()
202
203    @pyqtSlot(int)
204    def on_spinbox_timestamp_value_changed(self, value: int):
205        self.ui.tableWidgetPreview.setColumnHidden(self.COLUMNS["T"], value == 0)
206        self.update_preview()
207
208    @pyqtSlot()
209    def on_accepted(self):
210        QApplication.setOverrideCursor(Qt.WaitCursor)
211
212        iq_data, sample_rate = self.parse_csv_file(self.filename, self.ui.comboBoxCSVSeparator.currentText(),
213                                                   self.ui.spinBoxIDataColumn.value()-1,
214                                                   self.ui.spinBoxQDataColumn.value()-1,
215                                                   self.ui.spinBoxTimestampColumn.value()-1)
216
217        target_filename = self.filename.rstrip(".csv")
218        if os.path.exists(target_filename + ".complex"):
219            i = 1
220            while os.path.exists(target_filename + "_" + str(i) + ".complex"):
221                i += 1
222        else:
223            i = None
224
225        target_filename = target_filename if not i else target_filename + "_" + str(i)
226        target_filename += ".complex"
227
228        iq_data.tofile(target_filename)
229
230        self.data_imported.emit(target_filename, sample_rate if sample_rate is not None else 0)
231        QApplication.restoreOverrideCursor()
232
233if __name__ == '__main__':
234    app = QApplication(["urh"])
235    csv_dia = CSVImportDialog()
236    csv_dia.exec_()
237