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.frontsLabel = QTW.QLabel("Fronts (us): ")
        self.frontsEdit = QTW.QSpinBox()
        self.frontsEdit.setMaximum(100)
        self.frontsEdit.setValue(2)
        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.frontsLabel)
        self.paramLayout.addWidget(self.frontsEdit)
        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.syncLayout = QTW.QHBoxLayout()
        self.syncButton = QTW.QPushButton("Sync Alone Test")
        self.syncButton.pressed.connect(self.syncAlone)

        self.loadLayout.addWidget(self.freqLabel)
        self.loadLayout.addWidget(self.freqEdit)
        self.loadLayout.addWidget(self.loadButton)
        self.loadLayout.addWidget(self.loadLabel)
        self.syncLayout.addWidget(self.syncButton)
        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.mainlayout.addLayout(self.syncLayout)

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

        self.setLayout(self.mainlayout)

    def syncAlone(self):
        self.sock.connect(('rp-f0cf82.local', 5005))
        buffer = struct.pack('<B', 0x05)
        self.sock.sendall(buffer)
        buffer = self.sock.recv(4096)

    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()
            fronts = self.frontsEdit.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
                frontSamples = fronts * 125
                for i in range(num):
                    for j in range(frontSamples):
                        array = np.append(array, np.int16(32700 * j // frontSamples))
                        array = np.append(array, np.int16(0))
                    for j in range(highSamples):
                        array = np.append(array, np.int16(32700))
                        array = np.append(array, np.int16(0))
                    for j in range(frontSamples):
                        array = np.append(array, np.int16(32700 - 32700 * j // frontSamples))
                        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()
            fronts = self.frontsEdit.value()
            freq = self.freqEdit.value()
            print(len(alts))
            print(len(phases))
            fbuf = np.array([], dtype=np.float32)
            tbuf = np.arange(0, np.float32((swidth + delay + 2 * fronts) * 10e-6), np.float32(8 * 10e-9))
            fbuf = np.float32((alts / 32767) * np.float32(np.sin(2 * np.pi * freq * 100 * tbuf + np.pi * phases / (32767 * 2))))
            print(fbuf[10000:12000])

            self.sock.connect(('rp-f0cf82.local', 5005))
            buffer = struct.pack('<BIII', 0x01, (len(fbuf) // 4096 + 1) * 4096, 1, 1)
            self.sock.sendall(buffer)
            buffer = self.sock.recv(4096)
            if buffer[0] == 0x11:
                for i in range(0, len(fbuf), 509):
                    datasize = 509
                    if len(fbuf) - i < 509:
                        datasize = len(fbuf) - i
                    buffer = struct.pack('<BI', 0x11, np.uint32(datasize)) + fbuf[i:(i+datasize)].tobytes()
                    print(np.frombuffer(fbuf[i:(i+datasize)].tobytes(), dtype=np.float32, count=1, offset=0))
                    self.sock.sendall(buffer)
                    buffer = self.sock.recv(4096)
                    if buffer[0] != 0x11:
                        self.sock.close()
                        return -1
            else:
                self.sock.close()
                return -1
            buffer = struct.pack('<B', 0x21)
            self.sock.sendall(buffer)
            buffer = self.sock.recv(1024)
            if buffer[0] == 0x01:
                 self.loadLabel.setText("Loaded")
            self.sock.close()    


            



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

sys.exit(app.exec())