import numpy as np
from matplotlib.figure import Figure
import socket
import msgpack
from PySide6.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QPushButton, QHBoxLayout, QMessageBox
from PySide6.QtWidgets import QDialog, QFormLayout, QLineEdit, QLabel, QComboBox, QSpinBox, QToolButton, QFileDialog
from PySide6.QtCore import QTimer
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from PySide6.QtWidgets import QSizePolicy
import json

class MplCanvas(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi, facecolor='grey')
        super().__init__(self.fig)
        self.setParent(parent)
        self.axes = self.fig.add_subplot(111)
        # Построим простой график
        self.axes.plot([0, 1, 2, 3], [10, 1, 20, 3])
        self.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)

class SDRClient(QMainWindow):
    def __init__(self):
        super().__init__()
        # Параметры подключения к серверу
        self.server_ip = '127.0.0.1'
        self.server_port = 5003
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # Параметры для отрисовки
        self.channel_data = {
            'time': np.array([], dtype=np.int64),
            'channel_g': np.array([], dtype=np.int16),
            'channel_rf': np.array([], dtype=np.int16),
            'channel_adc_ttl': np.array([], dtype=np.int16),
            'channel_rf_ttl': np.array([], dtype=np.int16)
        }
        self.wavetypes = ["default"]
        # Настройка окна
        self.setWindowTitle("SDR Client")
        self.setGeometry(100, 100, 1280, 720)
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)

        self.layout = QHBoxLayout()

        self.plot_layout = QVBoxLayout()
        self.menu_layout = QVBoxLayout()
        self.layout.addLayout(self.menu_layout)
        self.layout.addLayout(self.plot_layout)

        self.central_widget.setLayout(self.layout)

        self.command_layout = QHBoxLayout()
        self.plot_layout.addLayout(self.command_layout)
        
        self.add_event_button = QPushButton("Add Event")
        self.command_layout.addWidget(self.add_event_button)
        self.add_rfwave_button = QPushButton("Add RF Wave")
        self.command_layout.addWidget(self.add_rfwave_button)
        self.run_event_list_button = QPushButton("Run Event List")
        self.command_layout.addWidget(self.run_event_list_button)
        self.clear_event_list_button = QPushButton("Clear Event List")
        self.command_layout.addWidget(self.clear_event_list_button)

        self.plot_canvas = MplCanvas(self, width=5, height=4, dpi=100)
        #self.plot_canvas.setSizePolicy(self.plot_widget.sizePolicy())
        self.plot_layout.addWidget(self.plot_canvas, stretch=1)

        toolbar = NavigationToolbar(self.plot_canvas, self)
        self.plot_layout.addWidget(toolbar, stretch=1)

        self.connect_button = QPushButton("Connect to SDR")
        self.connect_button.setMinimumSize(200, 50)
        self.disconnect_button = QPushButton("Disconnect")
        self.disconnect_button.setMinimumSize(200, 50)
        self.settings_button = QPushButton("Settings")
        self.settings_button.setMinimumSize(200, 50)
        self.recv_sim_data_button = QPushButton("Receive Sim Data")
        self.recv_sim_data_button.setMinimumSize(200, 50)

        self.menu_layout.addWidget(self.connect_button)
        self.menu_layout.addWidget(self.disconnect_button)
        self.menu_layout.addWidget(self.settings_button)
        self.menu_layout.addWidget(self.recv_sim_data_button)

        self.connect_button.clicked.connect(self.connect_to_server)
        self.disconnect_button.clicked.connect(self.disconnect_from_server)
        self.add_event_button.clicked.connect(self.add_event_dialog)
        self.add_rfwave_button.clicked.connect(self.add_wavetype_dialog)
        self.run_event_list_button.clicked.connect(self.run_event_list)
        self.clear_event_list_button.clicked.connect(self.clear_event_list)
        self.recv_sim_data_button.clicked.connect(self.receive_sim_data)

    def connect_to_server(self):
        try:
            self.sock.connect((self.server_ip, self.server_port))
            print(f"Connected to SDR server at {self.server_ip}:{self.server_port}")
            QMessageBox.information(self, "Connection Successful", f"Successfully connected to SDR server at {self.server_ip}:{self.server_port}")
        except Exception as e:
            print(f"Failed to connect to server: {e}")
            QMessageBox.critical(self, "Connection Error", f"Failed to connect to server: {e}")

    def disconnect_from_server(self):
        if self.sock:
            self.sock.close()
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            QMessageBox.information(self, "Disconnected", "Successfully disconnected from SDR server")
            print("Disconnected from SDR server")

    def add_event_dialog(self):
        #if not self.sock:
        #    QMessageBox.warning(self, "Not Connected", "Please connect to the SDR server first")
        #    return False
        add_event_dialog = QDialog(self)
        add_event_dialog.setWindowTitle("Add Event")
        layout = QFormLayout()

        labels = ["Start Time (ticks)", "Duration (ticks)", "Waveform Type", "time_front_rf_ttl (ticks)", "time_back_rf_ttl (ticks)", "time_front_adc_ttl (ticks)", "time_back_adc_ttl (ticks)", "time_start_rf (ticks)", "time_end_rf (ticks)", "rf_freq (Hz)", "phase_offset"]
        self.event_inputs = {}
        for label in labels:
            if label == "Waveform Type":
                combo = QComboBox()
                combo.addItems(self.wavetypes)
                layout.addRow(QLabel(label), combo)
                self.event_inputs[label] = combo
            else:
                line_edit = QLineEdit()
                layout.addRow(QLabel(label), line_edit)
                self.event_inputs[label] = line_edit
        
        submit_button = QPushButton("Submit")
        layout.addWidget(submit_button)

        load_event_list_button = QPushButton("Load Event List")
        layout.addWidget(load_event_list_button)

        submit_button.clicked.connect(self.submit_event)
        load_event_list_button.clicked.connect(self.load_event_list)

        add_event_dialog.setLayout(layout)
        add_event_dialog.exec()

    def submit_event(self):
        try:
            msg = {
                "magic": 0xAA,
                "cmd": 0x10,
                "time": int(self.event_inputs["Start Time (ticks)"].text()),
                "duration": int(self.event_inputs["Duration (ticks)"].text()),
                "wave": self.event_inputs["Waveform Type"].currentText(),
                "time_front_rf_trigger": int(self.event_inputs["time_front_rf_ttl (ticks)"].text()),
                "time_back_rf_trigger": int(self.event_inputs["time_back_rf_ttl (ticks)"].text()),
                "time_front_adc_trigger": int(self.event_inputs["time_front_adc_ttl (ticks)"].text()),
                "time_back_adc_trigger": int(self.event_inputs["time_back_adc_ttl (ticks)"].text()),
                "rf_start": int(self.event_inputs["time_start_rf (ticks)"].text()),
                "rf_end": int(self.event_inputs["time_end_rf (ticks)"].text()),
                "carrier_freq": float(self.event_inputs["rf_freq (Hz)"].text()),
                "phase_offset": int(self.event_inputs["phase_offset"].text())
            }
            msg = msgpack.packb(msg, use_bin_type=True)
            self.sock.sendall(msg)
            raw_resp = self.sock.recv(4096)
            if not raw_resp:
                raise Exception("No response from server")
            resp = msgpack.unpackb(raw_resp, raw=False)
            if resp["cmd"] != 0xFE or resp["code"] != 0x00000010:
                raise Exception(f"Server error: {resp['message']}")
            print("Event added successfully")
            QMessageBox.information(self, "Success", "Event added successfully")
        except Exception as e:
            print(f"Failed to add event: {e}")
            QMessageBox.critical(self, "Error", f"Failed to add event: {e}")

    def load_event_list(self):
        file_dialog = QFileDialog(self)
        file_dialog.setNameFilter("JSON Files (*.json)")
        if file_dialog.exec():
            selected_file = file_dialog.selectedFiles()[0]
            try:
                with open(selected_file, 'r') as f:
                    event_list = json.load(f)
                for event_data in event_list:
                    msg = {
                        "magic": 0xAA,
                        "cmd": 0x10,
                        "time": event_data["time"],
                        "duration": event_data["duration"],
                        "wave": event_data["wave"],
                        "time_front_rf_trigger": event_data["time_front_rf_trigger"],
                        "time_back_rf_trigger": event_data["time_back_rf_trigger"],
                        "time_front_adc_trigger": event_data["time_front_adc_trigger"],
                        "time_back_adc_trigger": event_data["time_back_adc_trigger"],
                        "rf_start": event_data["rf_start"],
                        "rf_end": event_data["rf_end"],
                        "carrier_freq": event_data["carrier_freq"],
                        "phase_offset": event_data["phase_offset"]
                    }
                    packed_msg = msgpack.packb(msg, use_bin_type=True)
                    self.sock.sendall(packed_msg)
                    resp = self.sock.recv(4096)
                    if not resp:
                        raise Exception("No response from server")
                    unpacked_resp = msgpack.unpackb(resp, raw=False)
                    if unpacked_resp["cmd"] != 0xFE and unpacked_resp["code"] != 0x00000010:
                        raise Exception(f"Server error: {unpacked_resp['message']}")
                print("Event list loaded successfully")
                QMessageBox.information(self, "Success", "Event list loaded successfully")
            except Exception as e:
                print(f"Failed to load event list: {e}")
                QMessageBox.critical(self, "Error", f"Failed to load event list: {e}")

    def add_wavetype_dialog(self):
        add_wave_dialog = QDialog(self)
        add_wave_dialog.setWindowTitle("Add Waveform Type")
        layout = QFormLayout()

        name_input = QLineEdit()
        layout.addRow(QLabel("Waveform Name"), name_input)

        self.file_input = QLineEdit()
        file_dialog_button = QPushButton("Browse")
        file_dialog_button.clicked.connect(lambda: self.open_wave_file_dialog(self.file_input))
        layout.addRow(QLabel("Waveform File"), self.file_input)
        layout.addWidget(file_dialog_button)

        submit_button = QPushButton("Submit")
        layout.addWidget(submit_button)

        submit_button.clicked.connect(lambda: self.submit_wavetype(name_input.text(), self.file_input.text(), add_wave_dialog))

        add_wave_dialog.setLayout(layout)
        add_wave_dialog.exec()

    def open_wave_file_dialog(self, file_input):
        file_dialog = QFileDialog(self)
        file_dialog.setNameFilter("All Files (*);;Binary Files (*.bin)")
        if file_dialog.exec():
            selected_file = file_dialog.selectedFiles()[0]
            file_input.setText(selected_file)

    def submit_wavetype(self, name, file_path, dialog):
        if not name:
            QMessageBox.warning(self, "Input Error", "Please enter a waveform name")
            return
        if name in self.wavetypes:
            QMessageBox.warning(self, "Input Error", "Waveform name already exists")
            return
        header_msg = {
            "magic": 0xAA,
            "cmd": 0x20,
            "wavetype": name,
            "packet_length": 512,  # Define packet length for sending waveform data
            "total_length": 0,  # Placeholder, will be updated after loading the waveform data
            "total_packets": 0,  # Placeholder, will be updated after loading the waveform data
        }

        data = b''
        with open(file_path, 'rb') as f:
            # Load rows point from .bin file (first 2 bytes is Q-point, next 2 bytes is I-point)
            data = f.read()

        total_wavedata_q = np.array([], dtype=np.int16)
        total_wavedata_i = np.array([], dtype=np.int16)
        for i in range(len(data) // 4):
            total_wavedata_q = np.append(total_wavedata_q, np.frombuffer(data, dtype=np.int16, count=1, offset=4*i))
            total_wavedata_i = np.append(total_wavedata_i, np.frombuffer(data, dtype=np.int16, count=1, offset=4*i+2))

        if len(total_wavedata_q) != len(total_wavedata_i):
            QMessageBox.warning(self, "Input Error", "In-phase and quadrature wave data must have the same length")
            return
        
        packet_length = 512  # Define packet length for sending waveform data
        header_msg["total_length"] = len(total_wavedata_q)
        header_msg["total_packets"] = len(total_wavedata_q) // packet_length + (1 if len(total_wavedata_q) % packet_length > 0 else 0)

        try:
            packed_header = msgpack.packb(header_msg, use_bin_type=True)
            self.sock.sendall(packed_header)
            
            raw_resp = self.sock.recv(4096)
            if not raw_resp:
                raise Exception("No response from server")
            resp = msgpack.unpackb(raw_resp, raw=False)
            if resp["cmd"] != 0xFE or resp["code"] != 0x00000001:
                raise Exception(f"Server error: {resp['message']}")
        except Exception as e:
            print(f"Failed to send waveform header: {e}")
            QMessageBox.critical(self, "Error", f"Failed to send waveform header: {e}")
            return
        
        print(f"Total packets to send: {header_msg["total_packets"]}")
        try:
            for i in range(header_msg["total_packets"]):
                print(f"Sending packet {i}")
                packet_length = min(512, len(total_wavedata_q) - i * 512)
                start_idx = i * 512
                end_idx = min((i + 1) * 512, len(total_wavedata_q))

                packet_data = {
                    "magic": 0xAA,
                    "cmd": 0x2C,
                    "packet_index": i,
                    "packet_length": packet_length,
                    "wavedata_q": total_wavedata_q[start_idx:end_idx].tobytes(),
                    "wavedata_i": total_wavedata_i[start_idx:end_idx].tobytes()
                }
                packed_packet = msgpack.packb(packet_data, use_bin_type=True)
                self.sock.sendall(packed_packet)
                raw_resp = self.sock.recv(4096)
                if not raw_resp:
                    raise Exception("No response from server")
                resp = msgpack.unpackb(raw_resp, raw=False)
                if resp["cmd"] != 0xFE or resp["code"] != 0x00000002:
                    raise Exception(f"Server error: {resp['message']}")
                self.sock.sendall(b"ACK") # Send ACK to server to indicate ready for next packet
                print(f"Packet {i} sended!")
            print("All packets sent!")
            raw_resp = self.sock.recv(4096)
            if not raw_resp:
                raise Exception("No response from server after sending all packets")
            resp = msgpack.unpackb(raw_resp, raw=False)
            if resp["cmd"] != 0xFE or resp["code"] != 0x00000003:
                raise Exception(f"Server error after sending all packets: {resp['message']}")
            self.sock.sendall(b'ACK')
            print("All packets sent!")
        except Exception as e:
            print(f"Failed to send waveform data: {e}")
            QMessageBox.critical(self, "Error", f"Failed to send waveform data: {e}")
            return

        self.wavetypes.append(name)
        print(f"Added new waveform type: {name}")
        QMessageBox.information(self, "Success", f"Added new waveform type: {name}")
        dialog.close()

    def run_event_list(self):
        msg = {
            "magic": 0xAA,
            "cmd": 0x30,
            "interpolation_method": "cubic"  # Placeholder, can be extended to allow user selection of interpolation method
        }
        try:
            packed_msg = msgpack.packb(msg, use_bin_type=True)
            self.sock.sendall(packed_msg)
            print("Sent run event list command to server")
        except Exception as e:
            print(f"Failed to send run event list command: {e}")
            QMessageBox.critical(self, "Error", f"Failed to send run event list command: {e}")
            return
        raw_resp = self.sock.recv(4096)
        if not raw_resp:
            print("No response from server after sending run event list command")
            QMessageBox.warning(self, "No Response", "No response from server after sending run event list command")
            return
        resp = msgpack.unpackb(raw_resp, raw=False)
        if resp["cmd"] != 0xFE or resp["code"] != 0x00000032:
            print(f"Server error after sending run event list command: {resp['message']}")
            QMessageBox.warning(self, "Server Error", f"Server error after sending run event list command: {resp['message']}")
            return
        QMessageBox.information(self, "Success", "Run event list command sent to server")

    def clear_event_list(self):
        msg = {
            "magic": 0xAA,
            "cmd": 0x60
        }
        try:
            packed_msg = msgpack.packb(msg, use_bin_type=True)
            self.sock.sendall(packed_msg)
            print("Sent clear event list command to server")
        except Exception as e:
            print(f"Failed to send clear event list command: {e}")
            QMessageBox.critical(self, "Error", f"Failed to send clear event list command: {e}")
            return
        raw_resp = self.sock.recv(4096)
        if not raw_resp:
            print("No response from server after sending clear event list command")
            QMessageBox.warning(self, "No Response", "No response from server after sending clear event list command")
            return
        resp = msgpack.unpackb(raw_resp, raw=False)
        if resp["cmd"] != 0xFE or resp["code"] != 0x00000040:
            print(f"Server error after sending clear event list command: {resp['message']}")
            QMessageBox.warning(self, "Server Error", f"Server error after sending clear event list command: {resp['message']}")
            return
        QMessageBox.information(self, "Success", "Clear event list command sent to server")

    def plot_sim_data(self):
        self.plot_canvas.axes.clear()
        self.plot_canvas.axes.plot(self.channel_data["time"], self.channel_data["channel_g"], label="Gradient")
        self.plot_canvas.axes.plot(self.channel_data["time"], self.channel_data["channel_rf"], label="RF")
        self.plot_canvas.axes.plot(self.channel_data["time"], self.channel_data["channel_adc_ttl"], label="ADC TTL")
        self.plot_canvas.axes.plot(self.channel_data["time"], self.channel_data["channel_rf_ttl"], label="RF TTL")
        self.plot_canvas.axes.set_xlabel("Time (ticks)")
        self.plot_canvas.axes.set_ylabel("Signal Value")
        self.plot_canvas.axes.set_title("Simulation Data from SDR Server")
        self.plot_canvas.axes.legend()
        self.plot_canvas.draw()

    def receive_sim_data(self):
        self.channel_data = {
            'time': np.array([], dtype=np.int64),
            'channel_g': np.array([], dtype=np.int16),
            'channel_rf': np.array([], dtype=np.int16),
            'channel_adc_ttl': np.array([], dtype=np.int16),
            'channel_rf_ttl': np.array([], dtype=np.int16)
        }
        msg = {
            "magic": 0xAA,
            "cmd": 0x40
        }

        try:
            packed_msg = msgpack.packb(msg, use_bin_type=True)
            self.sock.sendall(packed_msg)
            print("Sent request to receive simulation data from server")
        except Exception as e:
            print(f"Failed to send request for simulation data: {e}")
            QMessageBox.critical(self, "Error", f"Failed to send request for simulation data: {e}")
            return
        try:
            raw_resp = self.sock.recv(4096)
            if not raw_resp:
                raise Exception("No response from server after requesting simulation data")
            resp = msgpack.unpackb(raw_resp, raw=False) 
            if resp["cmd"] != 0xFE or resp["code"] != 0x00000043:
                raise Exception(f"Server error after requesting simulation data: {resp['message']}")
        except Exception as e:
            print(f"Failed to receive simulation data header: {e}")
            QMessageBox.critical(self, "Error", f"Failed to receive simulation data header: {e}")
            return
        
        print(f"Server is ready to send simulation data: {resp['message']}")
        self.sock.sendall(b"ACK") # Send ACK to server to indicate ready to receive simulation data

        while True:
            try:
                raw_packet = self.sock.recv(8192)
                if not raw_packet:
                    print("No more data from server, finished receiving simulation data")
                    break
                packet = msgpack.unpackb(raw_packet, raw=False)
                if packet["cmd"] == 0xFE and packet["code"] == 0x00000044:
                    print(f"Received end of simulation data message from server: {self.channel_data}") # Debug print to check received data
                    self.plot_sim_data()
                    QMessageBox.information(self, "Simulation Data Received", "All simulation data packets received from server")
                    break
                if packet["cmd"] != 0x33:
                    print(f"Received non-simulation data message from server: {packet}")
                    continue
                channel = packet["channel"]
                packet_index = packet["packet_index"]
                total_packets = packet["total_packets"]
                data_length = packet["data_length"]
                print(f"Received simulation data packet {packet_index}/{total_packets} for channel {channel} with length {data_length} bytes")
                if channel not in self.channel_data:
                    print(f"Unknown channel {channel} in received simulation data packet")
                    continue
                points = None
                if channel == "time":
                    points = np.frombuffer(packet["data"], dtype=np.int64)
                else:
                    # Append received data to channel data
                    points = np.frombuffer(packet["data"], dtype=np.int16)
                    print(f"Packet data preview (first 16 points): {points[:16]}") # Debug print to check packet data
                self.channel_data[channel] = np.append(self.channel_data[channel], points)
                # Send ACK to server to indicate ready for next packet
                self.sock.sendall(b"ACK")
            except Exception as e:
                print(f"Failed to receive or process simulation data packet: {e}")
                QMessageBox.critical(self, "Error", f"Failed to receive or process simulation data packet: {e}")
                return


app = QApplication([])
client = SDRClient()
client.show()
#plt.ion()  # Enable interactive mode for plotting
#plt.show()
app.exec()