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")
        self.btn1 = QtWidgets.QPushButton(self.centralwidget)
        self.btn1.setGeometry(QtCore.QRect(900, 80, 111, 25))
        self.btn1.setObjectName("btn1")
        self.AuftragsdetailsTable = QtWidgets.QTableWidget(self.centralwidget)
        self.AuftragsdetailsTable.setGeometry(QtCore.QRect(10, 10, 661, 192))
        self.AuftragsdetailsTable.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
        self.AuftragsdetailsTable.setObjectName("AuftragsdetailsTable")
        self.AuftragsdetailsTable.setColumnCount(5)
        self.AuftragsdetailsTable.setRowCount(1)
        item = QtWidgets.QTableWidgetItem()
        self.AuftragsdetailsTable.setVerticalHeaderItem(0, item)
        item = QtWidgets.QTableWidgetItem()
        font = QtGui.QFont()
        font.setPointSize(8)
        item.setFont(font)
        self.AuftragsdetailsTable.setHorizontalHeaderItem(0, item)
        item = QtWidgets.QTableWidgetItem()
        font.setPointSize(8)
        item.setFont(font)
        self.AuftragsdetailsTable.setHorizontalHeaderItem(1, item)
        item = QtWidgets.QTableWidgetItem()
        font.setPointSize(8)
        item.setFont(font)
        self.AuftragsdetailsTable.setHorizontalHeaderItem(2, item)
        item = QtWidgets.QTableWidgetItem()
        font.setPointSize(8)
        item.setFont(font)
        self.AuftragsdetailsTable.setHorizontalHeaderItem(3, item)
        item = QtWidgets.QTableWidgetItem()
        font.setPointSize(8)
        item.setFont(font)
        self.AuftragsdetailsTable.setHorizontalHeaderItem(4, item)
        item = QtWidgets.QTableWidgetItem()
        self.AuftragsdetailsTable.setItem(0, 0, item)
        self.label = QtWidgets.QLabel(self.centralwidget)
        self.label.setGeometry(QtCore.QRect(900, 20, 121, 17))
        self.label.setObjectName("label")
        self.textEdit = QtWidgets.QTextEdit(self.centralwidget)
        self.textEdit.setGeometry(QtCore.QRect(860, 40, 221, 31))
        self.textEdit.setObjectName("textEdit")
        self.bauteilTypBtn = QtWidgets.QPushButton(self.centralwidget)
        self.bauteilTypBtn.setGeometry(QtCore.QRect(60, 340, 161, 25))
        self.bauteilTypBtn.setObjectName("bauteilTypBtn")
        self.bauteiltypTextbox = QtWidgets.QTextEdit(self.centralwidget)
        self.bauteiltypTextbox.setGeometry(QtCore.QRect(230, 340, 221, 31))
        self.bauteiltypTextbox.setObjectName("bauteiltypTextbox")
        self.label_3 = QtWidgets.QLabel(self.centralwidget)
        self.label_3.setGeometry(QtCore.QRect(250, 320, 181, 20))
        self.label_3.setObjectName("label_3")
        self.label_4 = QtWidgets.QLabel(self.centralwidget)
        self.label_4.setGeometry(QtCore.QRect(520, 320, 67, 17))
        self.label_4.setObjectName("label_4")
        self.textEdit_3 = QtWidgets.QTextEdit(self.centralwidget)
        self.textEdit_3.setGeometry(QtCore.QRect(460, 340, 191, 31))
        self.textEdit_3.setObjectName("textEdit_3")
        self.checkBox = QtWidgets.QCheckBox(self.centralwidget)
        self.checkBox.setGeometry(QtCore.QRect(70, 410, 151, 23))
        self.checkBox.setChecked(False)
        self.checkBox.setObjectName("checkBox")
        self.graphicsView = QtWidgets.QGraphicsView(self.centralwidget)
        self.graphicsView.setGeometry(QtCore.QRect(800, 280, 256, 192))
        self.graphicsView.setObjectName("graphicsView")
        MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 1090, 22))
        self.menubar.setObjectName("menubar")
        MainWindow.setMenuBar(self.menubar)
        self.statusbar = QtWidgets.QStatusBar(MainWindow)
        self.statusbar.setObjectName("statusbar")
        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.yolo_stream.start()  # Start the YOLOv8 stream
            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:
                # Convert the frame from BGR (OpenCV format) to RGB
                frame = cv2.cvtColor(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_())