from PyQt5 import QtCore, QtGui, QtWidgets
import cv2
import sys
from CameraStream import YOLOv8CameraStream  # Import the YOLOv8CameraStream class

class Ui_MainWindow(object):
    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(1090, 675)
        self.centralwidget = QtWidgets.QWidget(MainWindow)
        self.centralwidget.setObjectName("centralwidget")

        # Create a vertical layout
        self.layout = QtWidgets.QVBoxLayout(self.centralwidget)

        # Add your widgets to the layout
        self.AuftragsdetailsTable = QtWidgets.QTableWidget(self.centralwidget)
        self.AuftragsdetailsTable.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.AuftragsdetailsTable.setColumnCount(5)
        self.AuftragsdetailsTable.setRowCount(1)
        self.layout.addWidget(self.AuftragsdetailsTable)

        self.label = QtWidgets.QLabel(self.centralwidget)
        self.layout.addWidget(self.label)

        self.textEdit = QtWidgets.QTextEdit(self.centralwidget)
        self.layout.addWidget(self.textEdit)

        self.bauteilTypBtn = QtWidgets.QPushButton(self.centralwidget)
        self.layout.addWidget(self.bauteilTypBtn)

        self.bauteiltypTextbox = QtWidgets.QTextEdit(self.centralwidget)
        self.layout.addWidget(self.bauteiltypTextbox)

        self.label_3 = QtWidgets.QLabel(self.centralwidget)
        self.layout.addWidget(self.label_3)

        self.label_4 = QtWidgets.QLabel(self.centralwidget)
        self.layout.addWidget(self.label_4)

        self.textEdit_3 = QtWidgets.QTextEdit(self.centralwidget)
        self.layout.addWidget(self.textEdit_3)

        self.checkBox = QtWidgets.QCheckBox(self.centralwidget)
        self.layout.addWidget(self.checkBox)

        # Create the QGraphicsView
        self.graphicsView = QtWidgets.QGraphicsView(self.centralwidget)
        self.layout.addWidget(self.graphicsView)

        # Add a button to the layout
        self.btn1 = QtWidgets.QPushButton(self.centralwidget)
        self.layout.addWidget(self.btn1)

        # Set the layout for the central widget
        self.centralwidget.setLayout(self.layout)

        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        MainWindow.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        MainWindow.setStatusBar(self.statusbar)

        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
        self.btn1.setText(_translate("MainWindow", "load Auftrag"))
        self.label.setText(_translate("MainWindow", "Auftragsnummer:"))
        self.bauteilTypBtn.setText(_translate("MainWindow", "Bauteiltype erkennen"))
        self.label_3.setText(_translate("MainWindow", "Bezeichnung"))
        self.label_4.setText(_translate("MainWindow", "Pos.-Nr:"))
        self.checkBox.setText(_translate("MainWindow", "Bauteiltype locked"))

class CameraStreamApp(QtWidgets.QMainWindow):
    def __init__(self, ui):
        super().__init__()
        self.ui = ui
        self.yolo_stream = None  # Initialize YOLOv8CameraStream as None
        self.timer = QtCore.QTimer(self)
        self.timer.timeout.connect(self.update_frame)
        self.ui.btn1.clicked.connect(self.start_camera)  # start camera if button LoadAuftrag is clicked
        self.scene = QtWidgets.QGraphicsScene(self)
        self.ui.graphicsView.setScene(self.scene)

    def start_camera(self):
        # Start the YOLOv8 camera stream (only if not already started)
        if self.yolo_stream is None:
            self.yolo_stream = YOLOv8CameraStream(model_path="models/yolov8m_seg_e300.pt", logging_level="high")
            self.timer.start(30)  # Start the timer to update the frame every 30ms (about 33 FPS)

    def update_frame(self):
        # Update the frame from YOLOv8 stream
        if self.yolo_stream:
            ret, frame = self.yolo_stream.cap.read()  # Capture frame from YOLOv8 stream

            if ret:
                # Process frame via YOLO model
                processed_frame = self.yolo_stream.process_frame(frame)

                # Convert the frame from BGR (OpenCV format) to RGB
                frame = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB)

                # Convert the frame to QImage
                h, w, ch = frame.shape
                bytes_per_line = ch * w
                qt_image = QtGui.QImage(frame.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888)

                # Add the QImage to a QPixmap
                pixmap = QtGui.QPixmap.fromImage(qt_image)

                # Get the size of the graphicsView and scale the pixmap to fit
                view_size = self.ui.graphicsView.size()
                scaled_pixmap = pixmap.scaled(view_size, QtCore.Qt.KeepAspectRatio)

                # Update the scene with the scaled pixmap
                self.scene.clear()
                self.scene.addPixmap(scaled_pixmap)

    def closeEvent(self, event):
        # Release the camera when the application is closed
        if self.yolo_stream is not None:
            self.yolo_stream.cap.release()
        event.accept()

if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    MainWindow = QtWidgets.QMainWindow()
    ui = Ui_MainWindow()
    ui.setupUi(MainWindow)
    
    # Initialize the CameraStreamApp with the UI
    camera_app = CameraStreamApp(ui)
    
    MainWindow.show()
    sys.exit(app.exec_())