1import logging
2import os
3
4import numpy as np
5from PyQt5 import uic
6from PyQt5.QtCore import QRegExp, Qt, pyqtSignal, pyqtSlot
7from PyQt5.QtGui import QBrush, QColor, QPen, QRegExpValidator
8from PyQt5.QtWidgets import QApplication, QDialog
9
10from urh.plugins.Plugin import SignalEditorPlugin
11from urh.signalprocessing.IQArray import IQArray
12from urh.ui.painting.SceneManager import SceneManager
13from urh.util.Formatter import Formatter
14from urh.util.Logger import logger
15
16
17class InsertSinePlugin(SignalEditorPlugin):
18    insert_sine_wave_clicked = pyqtSignal()
19
20    INSERT_INDICATOR_COLOR = QColor(0, 255, 0, 80)
21
22    def __init__(self):
23
24        self.__dialog_ui = None  # type: QDialog
25        self.complex_wave = None
26
27        self.__amplitude = 0.5
28        self.__frequency = 10
29        self.__phase = 0
30        self.__sample_rate = 1e6
31        self.__num_samples = int(1e6)
32
33        self.original_data = None
34        self.draw_data = None
35        self.position = 0
36
37        super().__init__(name="InsertSine")
38
39    @property
40    def dialog_ui(self) -> QDialog:
41        if self.__dialog_ui is None:
42            dir_name = os.path.dirname(os.readlink(__file__)) if os.path.islink(__file__) else os.path.dirname(__file__)
43
44            logging.getLogger().setLevel(logging.WARNING)
45            self.__dialog_ui = uic.loadUi(os.path.realpath(os.path.join(dir_name, "insert_sine_dialog.ui")))
46            logging.getLogger().setLevel(logger.level)
47
48            self.__dialog_ui.setAttribute(Qt.WA_DeleteOnClose)
49            self.__dialog_ui.setModal(True)
50            self.__dialog_ui.doubleSpinBoxAmplitude.setValue(self.__amplitude)
51            self.__dialog_ui.doubleSpinBoxFrequency.setValue(self.__frequency)
52            self.__dialog_ui.doubleSpinBoxPhase.setValue(self.__phase)
53            self.__dialog_ui.doubleSpinBoxSampleRate.setValue(self.__sample_rate)
54            self.__dialog_ui.doubleSpinBoxNSamples.setValue(self.__num_samples)
55            self.__dialog_ui.lineEditTime.setValidator(
56                QRegExpValidator(QRegExp(r"[0-9]+([nmµ]?|([\.,][0-9]{1,3}[nmµ]?))?$"))
57            )
58
59            scene_manager = SceneManager(self.dialog_ui.graphicsViewSineWave)
60            self.__dialog_ui.graphicsViewSineWave.scene_manager = scene_manager
61            self.insert_indicator = scene_manager.scene.addRect(0, -2, 0, 4,
62                                                                QPen(QColor(Qt.transparent), 0),
63                                                                QBrush(self.INSERT_INDICATOR_COLOR))
64            self.insert_indicator.stackBefore(scene_manager.scene.selection_area)
65
66            self.set_time()
67
68        return self.__dialog_ui
69
70    @property
71    def amplitude(self) -> float:
72        return self.__amplitude
73
74    @amplitude.setter
75    def amplitude(self, value: float):
76        if value != self.amplitude:
77            self.__amplitude = value
78            self.draw_sine_wave()
79
80    @property
81    def frequency(self) -> float:
82        return self.__frequency
83
84    @frequency.setter
85    def frequency(self, value: float):
86        if value != self.frequency:
87            self.__frequency = value
88            self.draw_sine_wave()
89
90    @property
91    def phase(self) -> float:
92        return self.__phase
93
94    @phase.setter
95    def phase(self, value: float):
96        if value != self.phase:
97            self.__phase = value
98            self.draw_sine_wave()
99
100    @property
101    def sample_rate(self) -> float:
102        return self.__sample_rate
103
104    @sample_rate.setter
105    def sample_rate(self, value: float):
106        if value != self.sample_rate:
107            self.__sample_rate = value
108            self.set_time()
109            self.draw_sine_wave()
110
111    @property
112    def num_samples(self) -> int:
113        return self.__num_samples
114
115    @num_samples.setter
116    def num_samples(self, value: int):
117        value = int(value)
118        if value != self.num_samples:
119            self.__num_samples = value
120            self.set_time()
121            self.draw_sine_wave()
122
123    def create_connects(self):
124        pass
125
126    def create_dialog_connects(self):
127        self.dialog_ui.doubleSpinBoxAmplitude.editingFinished.connect(
128            self.on_double_spin_box_amplitude_editing_finished)
129        self.dialog_ui.doubleSpinBoxFrequency.editingFinished.connect(
130            self.on_double_spin_box_frequency_editing_finished)
131        self.dialog_ui.doubleSpinBoxPhase.editingFinished.connect(self.on_double_spin_box_phase_editing_finished)
132        self.dialog_ui.doubleSpinBoxSampleRate.editingFinished.connect(
133            self.on_double_spin_box_sample_rate_editing_finished)
134        self.dialog_ui.doubleSpinBoxNSamples.editingFinished.connect(self.on_spin_box_n_samples_editing_finished)
135        self.dialog_ui.lineEditTime.editingFinished.connect(self.on_line_edit_time_editing_finished)
136        self.dialog_ui.buttonBox.accepted.connect(self.on_button_box_accept)
137        self.dialog_ui.buttonBox.rejected.connect(self.on_button_box_reject)
138        self.__dialog_ui.finished.connect(self.on_dialog_finished)
139
140    def get_insert_sine_dialog(self, original_data, position, sample_rate=None, num_samples=None) -> QDialog:
141        if sample_rate is not None:
142            self.__sample_rate = sample_rate
143            self.dialog_ui.doubleSpinBoxSampleRate.setValue(sample_rate)
144
145        if num_samples is not None:
146            self.__num_samples = int(num_samples)
147            self.dialog_ui.doubleSpinBoxNSamples.setValue(num_samples)
148
149        self.original_data = original_data
150        self.position = position
151
152        self.set_time()
153        self.draw_sine_wave()
154        self.create_dialog_connects()
155
156        return self.dialog_ui
157
158    def draw_sine_wave(self):
159        if self.dialog_ui.graphicsViewSineWave.scene_manager:
160            self.dialog_ui.graphicsViewSineWave.scene_manager.clear_path()
161
162        QApplication.instance().setOverrideCursor(Qt.WaitCursor)
163        self.__set_status_of_editable_elements(enabled=False)
164
165        t = np.arange(0, self.num_samples) / self.sample_rate
166        arg = 2 * np.pi * self.frequency * t + self.phase
167
168        self.complex_wave = np.empty(len(arg), dtype=np.complex64)
169        self.complex_wave.real = np.cos(arg)
170        self.complex_wave.imag = np.sin(arg)
171        self.complex_wave = IQArray(self.amplitude * self.complex_wave).convert_to(self.original_data.dtype)
172
173        self.draw_data = np.insert(self.original_data[:, 0], self.position, self.complex_wave[:, 0])
174        y, h = self.dialog_ui.graphicsViewSineWave.view_rect().y(), self.dialog_ui.graphicsViewSineWave.view_rect().height()
175        self.insert_indicator.setRect(self.position, y - h, self.num_samples, 2 * h + abs(y))
176
177        self.__set_status_of_editable_elements(enabled=True)
178        QApplication.instance().restoreOverrideCursor()
179        self.dialog_ui.graphicsViewSineWave.plot_data(self.draw_data)
180        self.dialog_ui.graphicsViewSineWave.show_full_scene()
181
182    def __set_status_of_editable_elements(self, enabled: bool):
183        for obj in ("doubleSpinBoxAmplitude", "doubleSpinBoxFrequency", "doubleSpinBoxPhase",
184                    "doubleSpinBoxSampleRate", "doubleSpinBoxNSamples", "lineEditTime", "buttonBox"):
185            getattr(self.dialog_ui, obj).setEnabled(enabled)
186
187    def set_time(self):
188        self.dialog_ui.lineEditTime.setText(Formatter.science_time(self.num_samples / self.sample_rate, decimals=3,
189                                                                   append_seconds=False, remove_spaces=True))
190
191    @pyqtSlot()
192    def on_double_spin_box_amplitude_editing_finished(self):
193        self.amplitude = self.dialog_ui.doubleSpinBoxAmplitude.value()
194
195    @pyqtSlot()
196    def on_double_spin_box_frequency_editing_finished(self):
197        self.frequency = self.dialog_ui.doubleSpinBoxFrequency.value()
198
199    @pyqtSlot()
200    def on_double_spin_box_phase_editing_finished(self):
201        self.phase = self.dialog_ui.doubleSpinBoxPhase.value()
202
203    @pyqtSlot()
204    def on_double_spin_box_sample_rate_editing_finished(self):
205        self.sample_rate = self.dialog_ui.doubleSpinBoxSampleRate.value()
206
207    @pyqtSlot()
208    def on_spin_box_n_samples_editing_finished(self):
209        self.num_samples = self.dialog_ui.doubleSpinBoxNSamples.value()
210
211    @pyqtSlot()
212    def on_line_edit_time_editing_finished(self):
213        time_str = self.dialog_ui.lineEditTime.text().replace(",", ".")
214        suffix = ""
215        try:
216            t = float(time_str)
217        except ValueError:
218            suffix = time_str[-1]
219            try:
220                t = float(time_str[:-1])
221            except ValueError:
222                return
223
224        factor = 10 ** -9 if suffix == "n" else 10 ** -6 if suffix == "µ" else 10 ** -3 if suffix == "m" else 1
225        time_val = t * factor
226
227        if self.sample_rate * time_val >= 1:
228            self.dialog_ui.doubleSpinBoxNSamples.setValue(self.sample_rate * time_val)
229            self.dialog_ui.doubleSpinBoxNSamples.editingFinished.emit()
230        else:
231            self.set_time()
232
233    @pyqtSlot()
234    def on_button_box_reject(self):
235        self.dialog_ui.reject()
236
237    @pyqtSlot()
238    def on_button_box_accept(self):
239        self.insert_sine_wave_clicked.emit()
240        self.dialog_ui.accept()
241
242    @pyqtSlot()
243    def on_dialog_finished(self):
244        self.sender().graphicsViewSineWave.eliminate()
245        self.__dialog_ui = None
246