import PySide6.QtWidgets as QTW
import numpy as np
import struct
import socket
import sys

class Window(QTW.QWidget):
    def __init__(self):
        super().__init__()
        
        self.genLayout = QTW.QHBoxLayout()
        self.genButton = QTW.QPushButton("Generate")
        self.genButton.pressed.connect(self.genFileDialog)
        self.genLabel = QTW.QLabel("Not generated")
        self.genLayout.addWidget(self.genButton)
        self.genLayout.addWidget(self.genLabel)

        self.paramLayout = QTW.QHBoxLayout()
        self.typeGenLabel = QTW.QLabel("Type: ")
        self.typeGenCombo = QTW.QComboBox()
        self.typeGenCombo.addItem("Rect")
        self.typeGenCombo.addItem("Sinc")
        self.typeGenCombo.addItem("Triangle")
        self.sigWidthLabel = QTW.QLabel("Width (us): ")
        self.sigWidthEdit = QTW.QSpinBox()
        self.sigWidthEdit.setMaximum(1000000)
        self.sigWidthEdit.setValue(1000)
        self.delayLabel = QTW.QLabel("Delay (us): ")
        self.delayEdit = QTW.QSpinBox()
        self.delayEdit.setMaximum(1000000)
        self.delayEdit.setValue(500)
        self.numLabel = QTW.QLabel("Repeats: ")
        self.numEdit = QTW.QSpinBox()
        self.numEdit.setMaximum(1000)
        self.numEdit.setValue(5)
        self.paramLayout.addWidget(self.typeGenLabel)
        self.paramLayout.addWidget(self.typeGenCombo)
        self.paramLayout.addWidget(self.sigWidthLabel)
        self.paramLayout.addWidget(self.sigWidthEdit)
        self.paramLayout.addWidget(self.delayLabel)
        self.paramLayout.addWidget(self.delayEdit)
        self.paramLayout.addWidget(self.numLabel)
        self.paramLayout.addWidget(self.numEdit)

        self.loadLayout = QTW.QHBoxLayout()
        self.loadFileEdit = QTW.QLineEdit()
        self.loadFileButton = QTW.QToolButton()
        self.freqLabel = QTW.QLabel("Central freq (kHz): ")
        self.freqEdit = QTW.QSpinBox()
        self.freqEdit.setMaximum(50000)
        self.freqEdit.setValue(2950)
        self.loadButton = QTW.QPushButton("Load")

        self.loadLabel = QTW.QLabel("Not loaded")
        self.loadLayout.addWidget(self.freqLabel)
        self.loadLayout.addWidget(self.freqEdit)
        self.loadLayout.addWidget(self.loadButton)
        self.loadLayout.addWidget(self.loadLabel)
        self.loadButton.pressed.connect(self.loadFileDialog)
        self.mainlayout = QTW.QVBoxLayout()
        self.mainlayout.addLayout(self.genLayout)
        self.mainlayout.addLayout(self.paramLayout)
        self.mainlayout.addLayout(self.loadLayout)

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

        self.setLayout(self.mainlayout)

    def genFileDialog(self):
        filename, ok = QTW.QFileDialog.getSaveFileName(
            self,
            "Save a File", 
            "", 
            "BIN traject (*.bin)"
        )

        if filename:
            swidth = self.sigWidthEdit.value()
            delay = self.delayEdit.value()
            num = self.numEdit.value()
            sigtype = self.typeGenCombo.currentText()

            array = np.array([], dtype=np.int16)
            samples = (swidth + delay) * 2 * num

            if sigtype == "Rect":
                delaySamples = delay * 125
                highSamples = swidth * 125
                for i in range(num):
                    for j in range(highSamples):
                        array = np.append(array, np.int16(32700))
                        array = np.append(array, np.int16(0))
                    for j in range(delaySamples):
                        array = np.append(array, np.int16(0))
                        array = np.append(array, np.int16(0))
            else:
                return -1

            with open(filename, "wb") as f:
                data = array.tobytes()
                f.write(data)

            self.genLabel.setText("Generated")


    def loadFileDialog(self):
        filename, ok = QTW.QFileDialog.getOpenFileName(
            self,
            "Open a File", 
            "", 
            "BIN traject (*.bin)"
        )
        
        if filename:
            data = b''
            with open(filename, "rb") as f:
                data = f.read()

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

            swidth = self.sigWidthEdit.value()
            delay = self.delayEdit.value()
            freq = self.freqEdit.value()
            fbuf = np.array([], dtype=np.float32)
            tbuf = np.arange(0, (swidth + delay) * 10e-6, 8 * 10e-9)
            fbuf = (alts / 32767) * np.sin(freq * 10e3 * tbuf + np.pi * phases / (32767 * 2))

            self.sock.connect(('rp-f082cf.local', 5005))
            buffer = struct.pack('<BIII', 0x02, len(fbuf), 0, 1)
            self.sock.sendall(buffer)
            buffer = self.sock.recv(1024)
            if buffer[0] == 0x11:
                for i in range(0, len(fbuf), 1019):
                    datasize = 1019
                    if len(fbuf) - i < 1019:
                        datasize = len(fbuf) - i
                    buffer = struct.pack('<BI', 0x31, np.uint32(datasize)) + fbuf[i:(i+datasize)].tobytes()
                
                    self.sock.sendall(buffer)
                    buffer = self.sock.recv(1024)
                    if buffer[0] != 0x11:
                        return -1
            buffer = struct.pack('<B', 0x21)
            self.sock.sendall(buffer)
            if buffer[0] == 0x01:
                 self.loadLabel.setText("Loaded")    


            



app = QTW.QApplication(sys.argv)
window = Window()
window.show()

sys.exit(app.exec())