Main

from PyQt5.QtWidgets import (
    QApplication, QMainWindow)
from PyQt5.QtCore import QTimer
from qt_material import apply_stylesheet
import sys
import numpy as np
from  PhasedArray import Array
from Receiver import Receiver
import Scenarios
import UI
import matplotlib.pyplot as plt
from Scenarios import arrays_scenarios, receivers_scenarios, five_G, tumor_ablation, ultrasound

class Main(QMainWindow, UI.ui):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.receivers_ = {}
        self.arrays = {}
        x = np.linspace(-20, 20, 100)
        y = np.linspace( 0, 40, 100)
        self.meshgrid_x, self.meshgrid_y = np.meshgrid(x, y) 
        self.meshgrid = np.array([self.meshgrid_x, self.meshgrid_y])
        self.button_of_add_new_array.clicked.connect(self.save_update)
        self.save_button.clicked.connect(self.add_array)
        self.button_of_add_new_reciver.clicked.connect(self.clear_receiver_inputs)
        self.save_button_of_reciver.clicked.connect(self.add_update_receiver)
        self.remove_array_button.clicked.connect(self.remove_array)
        self.remove_button_of_reciver.clicked.connect(self.remove_receiver)
        self.current_arrays_combo_box.currentTextChanged.connect(self.change_current_array)
        self.current_recivers_combo_box.currentTextChanged.connect(self.change_current_receiver)
        self.array_info_combobox.currentTextChanged.connect(self.display_array_info)
        self.button_of_add_new_reciver.clicked.connect(self.clear_receiver_inputs)
        self.applay_the_senario.clicked.connect(self.set_scenario)
        self.beam_profile_x = np.linspace(-20, 20, 500)
        self.beam_profile_y = np.zeros_like(self.beam_profile_x)
        self.beam_profile_angles = np.linspace(0, 2 * np.pi, 360)
        self.beam_profile_power = np.zeros_like(self.beam_profile_angles)
        self.current_array = None
        self.current_receiver = None
        self.receiver_texts = []
        self.receivers_positions = []
        self.frame = 100
        self.max_frames = 100 
        self.update_cheaker = 0
        # self.timer = QTimer()
        # self.timer.timeout.connect(self.update_all)  
        # self.timer.start(1000)
        self.angles = np.linspace(0, 2 * np.pi, 360)
        self.polar_power = np.zeros_like(self.angles)
        contour = self.heat_map.axes.contourf(self.meshgrid_x, self.meshgrid_y,
                                               np.zeros_like(self.meshgrid_x), levels=np.arange(-15, 15.1, 0.1), cmap='RdBu')  
        self.color_bar = plt.colorbar(contour, ax=self.heat_map.axes, label='Amplitude')


    def save_update(self):
        if self.button_of_add_new_array.text() == "Add Array":
            self.button_of_add_new_array.setText("real time")
            self.radio_button_of_linear.setEnabled(False)
            self.radio_button_of_curve.setEnabled(False)
            self.save_button.setVisible(False)
            self.remove_array_button.setVisible(False)
            self.array_name_line_edit.setVisible(False)
            self.label_of_array_name.setVisible(False)
            self.slider_of_transmiters_number.sliderReleased.connect(self.update_array)
            self.slider_of_element_spacing.sliderReleased.connect(self.update_array)
            self.slider_of_steering_angle.sliderReleased.connect(self.update_array)
            self.slider_of_arc_angle.sliderReleased.connect(self.update_array)
            self.frequencies_line_edit.returnPressed.connect(self.update_array)
            self.array_position_x_line_edit.returnPressed.connect(self.update_array)
            self.array_position_y_line_edit.returnPressed.connect(self.update_array)
            self.radius_line_edit.returnPressed.connect(self.update_array)
            self.clear_array_inputs()
        else:
            self.button_of_add_new_array.setText("Add Array")
            self.slider_of_transmiters_number.sliderReleased.disconnect(self.update_array)
            self.slider_of_element_spacing.sliderReleased.disconnect(self.update_array)
            self.slider_of_steering_angle.sliderReleased.disconnect(self.update_array)
            self.slider_of_arc_angle.sliderReleased.disconnect(self.update_array)
            self.frequencies_line_edit.returnPressed.disconnect(self.update_array)
            self.array_position_x_line_edit.returnPressed.disconnect(self.update_array)
            self.array_position_y_line_edit.returnPressed.disconnect(self.update_array)
            self.radius_line_edit.returnPressed.disconnect(self.update_array)
            self.radio_button_of_linear.setEnabled(True)
            self.radio_button_of_curve.setEnabled(True)
            self.save_button.setVisible(True)
            self.remove_array_button.setVisible(True)
            self.array_name_line_edit.setVisible(True)
            self.label_of_array_name.setVisible(True)

    def set_scenario(self):
        scenario = self.combo_box_of_senario.currentText()
        self.remove_scenario()
        if scenario == "5G":
            five_G(self.meshgrid_x, self.meshgrid_y, self.beam_profile_x, self.beam_profile_y)
        
        elif scenario == "Ultrasound":
            ultrasound(self.meshgrid_x, self.meshgrid_y, self.beam_profile_x, self.beam_profile_y)

        elif scenario == "Tumor Ablation":
            tumor_ablation(self.meshgrid_x, self.meshgrid_y, self.beam_profile_x, self.beam_profile_y)
        
        self.arrays.update(arrays_scenarios)
        arrays_name = [k for k in arrays_scenarios.keys()]
        self.current_arrays_combo_box.addItems(arrays_name)
        self.array_info_combobox.addItems(arrays_name)
        self.current_arrays_combo_box.setCurrentText(arrays_name[-1])
        self.array_info_combobox.setCurrentText(arrays_name[-1])
        
        receivers_names = [k for k in receivers_scenarios.keys()]
        self.receivers_.update(receivers_scenarios)
        self.current_recivers_combo_box.addItems(receivers_names)
        self.current_recivers_combo_box.setCurrentText(receivers_names[-1])
        self.change_current_receiver()

        self.draw_arrays_and_receivers()
        self.update(self.frame)

    def remove_scenario(self):
        
        arrays = [array for array in self.arrays.keys()]
        receivers = [receiver for receiver in self.receivers_]
        for array in arrays:
            self.remove_array(array)

        for receiver in receivers:
            self.remove_receiver(receiver)
        self.update(self.frame)

    def update_all(self):
        if len(self.arrays) > 0:
            # self.update(self.frame)
            self.frame = (self.frame + 1) % self.max_frames

    def add_array(self):
        name = self.array_name_line_edit.text()
        number_of_elements = self.slider_of_transmiters_number.value()
        frequencies_text = self.frequencies_line_edit.text().strip()
        frequencies_text = frequencies_text[:-1] if frequencies_text[-1] == "," else frequencies_text
        frequencies_text = frequencies_text.split(",")
        frequencies = []

        for i in range (len(frequencies_text)):
            frequencies.append(int(frequencies_text[i]))

        steering_angle = self.slider_of_steering_angle.value()
        position = [int(self.array_position_x_line_edit.text()), int(self.array_position_y_line_edit.text())]

        if self.radio_button_of_linear.isChecked():
            array_type = "Linear"
            elements_spacing = self.slider_of_element_spacing.value() * 0.5

            if self.current_arrays_combo_box.currentText() in self.arrays:
                old_name = self.current_array.name
                self.current_array.start_array(name, array_type, number_of_elements, frequencies, steering_angle, 
                    position, self.meshgrid_x, self.meshgrid_y, self.beam_profile_x, self.beam_profile_y,
                    element_spacing= elements_spacing)
            else:
                new_array = Array(name, array_type, number_of_elements, frequencies, steering_angle, 
                    position, self.meshgrid_x, self.meshgrid_y, self.beam_profile_x, self.beam_profile_y,
                    element_spacing= elements_spacing)
                
        elif self.radio_button_of_curve.isChecked():
            array_type = "Curved"
            arc_angle = self.slider_of_arc_angle.value()
            radius = int(self.radius_line_edit.text())
        
            if self.current_arrays_combo_box.currentText() in self.arrays:
                old_name = self.current_array.name
                self.current_array.start_array(name, array_type, number_of_elements, frequencies, steering_angle, position, self.meshgrid_x, 
                                            self.meshgrid_y, self.beam_profile_x, self.beam_profile_y, element_spacing= None, radius= radius, 
                                            arc_angle= arc_angle)
            else:
                new_array = Array(name, array_type, number_of_elements, frequencies, steering_angle, position, self.meshgrid_x, 
                                            self.meshgrid_y, self.beam_profile_x, self.beam_profile_y, element_spacing= None, radius= radius, 
                                            arc_angle= arc_angle)
        
        self.current_arrays_combo_box.addItem(name)
        self.array_info_combobox.addItem(name)
        if self.current_arrays_combo_box.currentText() in self.arrays:
            self.arrays[name] = self.arrays.pop(old_name)
            self.remove_combobox_item(self.current_arrays_combo_box, old_name)
            self.remove_combobox_item(self.array_info_combobox, old_name)
        else:
            self.arrays[name] = new_array

        self.current_arrays_combo_box.setCurrentText(name)
        self.array_info_combobox.setCurrentText(name)
        self.change_current_array()
        self.draw_arrays_and_receivers()
        self.display_array_info()
        self.update(self.frame)

    def update_array(self):
        number_of_elements = self.slider_of_transmiters_number.value()
        frequencies_text = self.frequencies_line_edit.text().strip()
        frequencies_text = (frequencies_text[:-1] if frequencies_text[-1] == "," else frequencies_text) if len(frequencies_text) >0 else "1"
        frequencies_text = frequencies_text.split(",")
        frequencies = []
        steering_angle = self.slider_of_steering_angle.value() 

        for i in range (len(frequencies_text)):
            frequencies.append(int(frequencies_text[i]))

        pos_x = int(self.array_position_x_line_edit.text()) if len(self.array_position_x_line_edit.text()) > 0 else  0 
        pos_y = int(self.array_position_y_line_edit.text()) if len(self.array_position_y_line_edit.text()) > 0 else  0
        position = [pos_x, pos_y]

        if self.radio_button_of_linear.isChecked():
            array_type = "Linear"
            elements_spacing = self.slider_of_element_spacing.value() * 0.5
            arc_angle = None
            radius = None
        else:
            array_type = "Curved"
            elements_spacing = None
            arc_angle = self.slider_of_arc_angle.value()
            radius = int(self.radius_line_edit.text())
        
        name = self.current_array.name
        self.current_array.start_array(name, array_type, number_of_elements, frequencies, steering_angle, 
            position, self.meshgrid_x, self.meshgrid_y, self.beam_profile_x, self.beam_profile_y,
            element_spacing= elements_spacing, radius= radius, arc_angle=arc_angle)

        self.draw_arrays_and_receivers()
        self.display_array_info()
        self.update(self.frame)

    def add_update_receiver(self):
        if self.reciver_position_x.text().strip() == ""or self.reciver_position_y.text().strip() == "" or self.reciver_name.text().strip()=="" :
            return
        name = self.reciver_name.text()
        position = [int(self.reciver_position_x.text()), int(self.reciver_position_y.text())]
        if self.current_recivers_combo_box.currentText() in self.receivers_:
            old_name = self.current_receiver.name
            self.current_receiver.initialize_receiver_data(name, position)
            self.receivers_[name] = self.receivers_.pop(old_name)
            self.current_recivers_combo_box.addItem(name)
            self.remove_combobox_item(self.current_recivers_combo_box, old_name)
        else:
            new_receiver = Receiver(name, position)
            self.receivers_[name] = new_receiver
            self.current_recivers_combo_box.addItem(name)
        self.current_recivers_combo_box.setCurrentText(name)
        self.draw_arrays_and_receivers()
        self.change_current_receiver()
        self.update(self.frame)

    def draw_arrays_and_receivers(self):
        self.transmiters_recivers_plotter.axes.clear()
        self.transmiters_recivers_plotter.axes.grid(visible = True)
        for idx, array in enumerate(self.arrays.values()):
            positions = np.array(array.array_data["positions"])
            self.transmiters_recivers_plotter.axes.scatter(
                positions[:, 0], positions[:, 1], label=f"Array {idx+1}", marker='o'
            )
        self.receivers_positions = []
        self.receiver_texts = []
        for receiver in self.receivers_.values():
            self.receivers_positions.append(receiver.position)
            self.receiver_texts.append(
                self.transmiters_recivers_plotter.axes.text(receiver.position[0] + 0.1,
                                                        receiver.position[1], "", fontsize=8, color="black")
            )
        if len(self.receivers_positions) > 0:
            self.transmiters_recivers_plotter.axes.scatter(*zip(*self.receivers_positions),
                                                                        color='red', marker='x', s=100, label="Receivers")
        
    def update_phase_shifts(self, slider_value, slider_index):
        self.clear_layout(self.grid_layout_of_change_info)
        self.phase_shifts_sliders_values[slider_index] = slider_value
        self.current_array.phase_shifts = self.phase_shifts_sliders_values

    def change_current_array(self):
        if self.current_arrays_combo_box.currentText()in self.arrays:
            self.update_cheaker = len(self.arrays)+1
            self.current_array = self.arrays[self.current_arrays_combo_box.currentText()]
            self.array_name_line_edit.setText(self.current_array.name)
            self.slider_of_transmiters_number.setValue(self.current_array.number_of_elements)
            self.array_position_x_line_edit.setText(str(self.current_array.position[0]))
            self.array_position_y_line_edit.setText(str(self.current_array.position[1]))
            temp_string = "".join(","+str(i) for i in self.current_array.frequencies)
            temp_string = temp_string[1:]
            self.frequencies_line_edit.setText(temp_string)

            if self.current_array.array_type == "Linear":
                self.radio_button_of_linear.setChecked(True)
                self.slider_of_element_spacing.setValue(int(self.current_array.elements_spacing * 2))
            else:
                self.radio_button_of_curve.setChecked(True)
                self.radius_line_edit.setText(str(self.current_array.radius))
                self.slider_of_arc_angle.setValue(int(self.current_array.arc_angle))

    def remove_array(self, array_name = None):
        if not array_name:
            array_name = self.current_arrays_combo_box.currentText()
        if array_name in self.arrays:
            temp_array = self.arrays[array_name]
            del temp_array
            del self.arrays[array_name]
            self.current_arrays_combo_box.removeItem(self.current_arrays_combo_box.findText(array_name))
            self.array_info_combobox.removeItem(self.array_info_combobox.findText(array_name))
        self.draw_arrays_and_receivers()
        self.display_array_info()
        self.update(self.frame)
        
    def remove_receiver(self, receiver_name= None):
        if not receiver_name :
            receiver_name = self.current_recivers_combo_box.currentText()
        if receiver_name in self.receivers_:
            temp_receiver = self.receivers_[receiver_name]
            del temp_receiver
            del self.receivers_[receiver_name]
            self.current_recivers_combo_box.removeItem(self.current_recivers_combo_box.findText(receiver_name))
        self.draw_arrays_and_receivers()

    def clear_array_inputs(self):
        if self.current_array is not None: 
            self.slider_of_transmiters_number.setValue(self.current_array.number_of_elements)
            if self.current_array.array_type == "Linear":
                self.slider_of_element_spacing.setValue(int(self.current_array.elements_spacing))
                self.slider_of_arc_angle.setValue(0)
                self.radius_line_edit.clear()
            else: 
                self.slider_of_element_spacing.setValue(0)
                self.slider_of_arc_angle.setValue(int(self.current_array.arc_angle))
                self.radius_line_edit.setText(f"{self.current_array.radius}")
            frequencies = str(self.current_array.frequencies)
            self.frequencies_line_edit.setText(frequencies[1:-1])
            self.array_position_x_line_edit.setText(f"{self.current_array.position[0]}")
            self.array_position_y_line_edit.setText(f"{self.current_array.position[1]}")
            
            self.slider_of_steering_angle.setValue(int(self.current_array.steering_angle))

        else:
            self.current_arrays_combo_box.setCurrentIndex(0)
            self.slider_of_transmiters_number.setValue(8)
            self.slider_of_element_spacing.setValue(1)
            self.slider_of_arc_angle.setValue(0)
            self.slider_of_steering_angle.setValue(0)
            self.radius_line_edit.clear()
            self.array_name_line_edit.clear()
            self.frequencies_line_edit.clear()
            self.array_position_x_line_edit.clear()
            self.array_position_y_line_edit.clear()
    
    def clear_receiver_inputs(self):
        self.current_recivers_combo_box.setCurrentIndex(0)
        self.reciver_position_x.clear()
        self.reciver_position_y.clear()
        self.reciver_name.clear()

    def change_current_receiver(self):
        name = self.current_recivers_combo_box.currentText()
        if name != "Add Receiver":
            self.current_receiver = self.receivers_[self.current_recivers_combo_box.currentText()]
            self.reciver_name.setText(self.current_receiver.name)
            self.reciver_position_x.setText(str(self.current_receiver.position[0]))
            self.reciver_position_y.setText(str(self.current_receiver.position[1]))

    def display_array_info(self):
        if len(self.array_info_combobox) == 0:
            pass
        if self.array_info_combobox.currentText() in self.arrays:
            array = self.arrays[self.array_info_combobox.currentText()]

            self.label_of_which_type.setText(str(array.array_type))
            
            self.label_info_number_of_transmiter_value.setText(str(array.number_of_elements))
            
            temp_string = "".join(","+str(i) for i in array.frequencies)
            temp_string = temp_string[1:]
            self.label_info_frequencies_value.setText(temp_string + "Hz")
            
            self.label_info_steering_angle_value.setText(f"{int(np.degrees(array.steering_angle))}"+ "˚")

            self.label_info_positoin_value.setText(f"{array.position[0]}x {array.position[1]}y")
            
            if array.array_type == "Linear":
                self.label_info_elements_spacing.show()
                self.label_info_elements_spacing_value.show()
                self.label_info_elements_spacing_value.setText(f"{array.elements_spacing}" + "λ")

                self.label_info_raduis_value.setText(".............")
                self.label_info_raduis.hide()
                self.label_info_raduis_value.hide()
                
                self.label_info_arc_angle_value.setText(".............")
                self.label_info_arc_angle.hide()
                self.label_info_arc_angle_value.hide()

            else:
                self.label_info_elements_spacing.hide()
                self.label_info_elements_spacing_value.hide()
                self.label_info_elements_spacing_value.setText(".............")

                self.label_info_raduis_value.setText(str(array.radius) + "m")
                self.label_info_raduis.show()
                self.label_info_raduis_value.show()

                self.label_info_arc_angle_value.setText(f"{int(np.degrees(array.arc_angle))}" + "˚")
                self.label_info_arc_angle.show()
                self.label_info_arc_angle_value.show()

    def update(self, frame):
        time = frame / 10.0
        self.polar_power.fill(0)
        resultant_wave = np.zeros_like(self.meshgrid_x)
        beam_profile_amplitude = np.zeros_like(self.beam_profile_x)
        receiver_signals = []

        for array in self.arrays.values():
            positions = array.array_data["unit_placements"]
            distances = array.array_data["distances"]
            beam_distances = array.array_data["beam_distances"]
            phase_shifts = array.array_data["phase_shifts"]
            frequencies = array.array_data["frequencies"]
            wavelengths = array.array_data["wavelengths"]
            steering_angle = array.array_data["steering_angle"]

            if array.name == "Transducer":
                array.steering_angle = array.steering_angle + Scenarios.steering_angle_increment_of_ultrasound
                if array.steering_angle > 40 or array.steering_angle < -40:
                    Scenarios.steering_angle_increment_of_ultrasound *= -1
                array.update_steering_angle()

            for r, beam_r, phi in zip(distances, beam_distances, phase_shifts):
                for freq, wavelength in zip(frequencies, wavelengths):
                    k = 2 * np.pi / wavelength
                    omega = 2 * np.pi * freq
                    resultant_wave += np.sin(k * r - omega * time + phi[wavelengths.index(wavelength)])
                    beam_profile_amplitude += np.sin(k * beam_r - omega * time + phi[wavelengths.index(wavelength)])

            for wavelength in wavelengths:
                k = 2 * np.pi / wavelength

                for theta_idx, theta in enumerate(self.angles):
                    steering_vector = np.exp(1j * k * (
                        positions[:, 0] * np.cos(theta - steering_angle) +
                        positions[:, 1] * np.sin(theta - steering_angle)
                    ))
                    if self.update_cheaker != len(self.arrays) or True:
                        self.polar_power[theta_idx] += np.abs(np.sum(steering_vector))**2
        for rx in self.receivers_positions:
            signal = 0
            for array in self.arrays.values():
                for r, phi in zip(array.array_data["distances"], array.array_data["phase_shifts"]):
                    signal += np.sin(2 * np.pi * np.linalg.norm(rx) - omega * time + phi[0])
            receiver_signals.append(signal)

        for i, text in enumerate(self.receiver_texts):
            text.set_text(f"{receiver_signals[i]:.2f}")

        for c in self.heat_map.axes.collections:
            c.remove()

        # if self.heat_map.no_label:
        #     self.heat_map.no_label = False
        #     self.heat_map.vmin, self.heat_map.vmax = np.min(resultant_wave), np.max(resultant_wave)

        # self.heat_map.axes.clear()
        # cax = self.heat_map.axes.pcolormesh(self.meshgrid_x, self.meshgrid_y, resultant_wave,
        #                                      shading='gouraud', cmap='plasma', vmin=self.heat_map.vmin, vmax=self.heat_map.vmax)

        # if not hasattr(self.heat_map, 'colorbar') or self.heat_map.colorbar is None:
        #     self.heat_map.colorbar = self.heat_map.figure.colorbar(cax, ax=self.heat_map.axes)

        # else:
        #     self.heat_map.colorbar.update_normal(cax)

        res_list = np.arange(resultant_wave.min(), resultant_wave.max()+0.1, 0.1)
        if len(res_list) < 2 :
            res_list = np.arange(-15, 15.1, 0.1) 
        countour = self.heat_map.axes.contourf(self.meshgrid_x, self.meshgrid_y, resultant_wave, levels=res_list, cmap='RdBu')

        if self.update_cheaker != len(self.arrays) or True:
            self.color_bar.set_ticks([res_list[0], res_list[-1]])
            self.color_bar.update_normal(countour)
            self.update_cheaker = len(self.arrays)
            if len(self.arrays) > 0:
                polar_power_db = 10 * np.log10(self.polar_power / np.max(self.polar_power) + 1e-10)
                self.beam_profile.axes.clear()
                self.beam_profile.axes.plot(self.angles - np.pi /2, polar_power_db)
            else:
                self.beam_profile.axes.clear()
            self.beam_profile.axes.set_title("Beam Pattern")
            self.beam_profile.axes.set_theta_zero_location("N")
            self.beam_profile.axes.set_theta_direction(-1) 
            self.beam_profile.axes.set_theta_offset(np.pi / 2)  
            self.beam_profile.axes.set_thetamin(-90)            
            self.beam_profile.axes.set_thetamax(90)
            self.beam_profile.draw()
        self.heat_map.draw()
        self.transmiters_recivers_plotter.draw()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    apply_stylesheet(app, "dark_purple.xml")
    window = Main()
    window.show()
    sys.exit(app.exec_())