import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QPushButton, QScrollArea, QFrame, QSizePolicy
)
from PyQt5.QtCore import QTimer, Qt, pyqtSignal, QThread
from PyQt5.QtGui import QImage, QPixmap, QFont, QPainter, QPen, QColor
import os
from yolopart.detector import LicensePlateYOLO
from OCR_part.ocr_interface import ocr_predict
#from CRNN_part.crnn_interface import crnn_predict（不使用CRNN）

class CameraThread(QThread):
    """摄像头线程类"""
    frame_ready = pyqtSignal(np.ndarray)
    
    def __init__(self):
        super().__init__()
        self.camera = None
        self.running = False
    
    def start_camera(self):
        """启动摄像头"""
        self.camera = cv2.VideoCapture(0)
        if self.camera.isOpened():
            self.running = True
            self.start()
            return True
        return False
    
    def stop_camera(self):
        """停止摄像头"""
        self.running = False
        if self.camera:
            self.camera.release()
        self.quit()
        self.wait()
    
    def run(self):
        """线程运行函数"""
        while self.running:
            if self.camera and self.camera.isOpened():
                ret, frame = self.camera.read()
                if ret:
                    self.frame_ready.emit(frame)
            self.msleep(30)  # 约30fps

class LicensePlateWidget(QWidget):
    """单个车牌结果显示组件"""
    
    def __init__(self, plate_id, class_name, corrected_image, plate_number):
        super().__init__()
        self.plate_id = plate_id
        self.init_ui(class_name, corrected_image, plate_number)
    
    def init_ui(self, class_name, corrected_image, plate_number):
        layout = QHBoxLayout()
        layout.setContentsMargins(10, 5, 10, 5)
        
        # 车牌类型标签
        type_label = QLabel(class_name)
        type_label.setFixedWidth(60)
        type_label.setAlignment(Qt.AlignCenter)
        type_label.setStyleSheet(
            "QLabel { "
            "background-color: #4CAF50 if class_name == '绿牌' else #2196F3; "
            "color: white; "
            "border-radius: 5px; "
            "padding: 5px; "
            "font-weight: bold; "
            "}"
        )
        if class_name == '绿牌':
            type_label.setStyleSheet(
                "QLabel { "
                "background-color: #4CAF50; "
                "color: white; "
                "border-radius: 5px; "
                "padding: 5px; "
                "font-weight: bold; "
                "}"
            )
        else:
            type_label.setStyleSheet(
                "QLabel { "
                "background-color: #2196F3; "
                "color: white; "
                "border-radius: 5px; "
                "padding: 5px; "
                "font-weight: bold; "
                "}"
            )
        
        # 矫正后的车牌图像
        image_label = QLabel()
        image_label.setFixedSize(120, 40)
        image_label.setStyleSheet("border: 1px solid #ddd; background-color: white;")
        
        if corrected_image is not None:
            # 转换numpy数组为QPixmap
            h, w = corrected_image.shape[:2]
            if len(corrected_image.shape) == 3:
                bytes_per_line = 3 * w
                q_image = QImage(corrected_image.data, w, h, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            else:
                bytes_per_line = w
                q_image = QImage(corrected_image.data, w, h, bytes_per_line, QImage.Format_Grayscale8)
            
            pixmap = QPixmap.fromImage(q_image)
            scaled_pixmap = pixmap.scaled(120, 40, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            image_label.setPixmap(scaled_pixmap)
        else:
            image_label.setText("车牌未完全\n进入摄像头")
            image_label.setAlignment(Qt.AlignCenter)
            image_label.setStyleSheet("border: 1px solid #ddd; background-color: #f5f5f5; color: #666;")
        
        # 车牌号标签
        number_label = QLabel(plate_number)
        number_label.setFixedWidth(150)
        number_label.setAlignment(Qt.AlignCenter)
        number_label.setStyleSheet(
            "QLabel { "
            "border: 1px solid #ddd; "
            "background-color: white; "
            "padding: 8px; "
            "font-family: 'Courier New'; "
            "font-size: 14px; "
            "font-weight: bold; "
            "}"
        )
        
        layout.addWidget(type_label)
        layout.addWidget(image_label)
        layout.addWidget(number_label)
        layout.addStretch()
        
        self.setLayout(layout)
        self.setStyleSheet(
            "QWidget { "
            "background-color: white; "
            "border: 1px solid #e0e0e0; "
            "border-radius: 8px; "
            "margin: 2px; "
            "}"
        )

class MainWindow(QMainWindow):
    """主窗口类"""
    
    def __init__(self):
        super().__init__()
        self.detector = None
        self.camera_thread = None
        self.current_frame = None
        self.detections = []
        
        self.init_ui()
        self.init_detector()
        self.init_camera()
    
    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("车牌识别系统")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中央widget
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 创建主布局
        main_layout = QHBoxLayout(central_widget)
        
        # 左侧摄像头显示区域
        left_frame = QFrame()
        left_frame.setFrameStyle(QFrame.StyledPanel)
        left_frame.setStyleSheet("QFrame { background-color: #f0f0f0; border: 2px solid #ddd; }")
        left_layout = QVBoxLayout(left_frame)
        
        # 摄像头显示标签
        self.camera_label = QLabel()
        self.camera_label.setMinimumSize(640, 480)
        self.camera_label.setStyleSheet("QLabel { background-color: black; border: 1px solid #ccc; }")
        self.camera_label.setAlignment(Qt.AlignCenter)
        self.camera_label.setText("摄像头未启动")
        self.camera_label.setScaledContents(True)
        
        # 控制按钮
        button_layout = QHBoxLayout()
        self.start_button = QPushButton("启动摄像头")
        self.stop_button = QPushButton("停止摄像头")
        self.start_button.clicked.connect(self.start_camera)
        self.stop_button.clicked.connect(self.stop_camera)
        self.stop_button.setEnabled(False)
        
        button_layout.addWidget(self.start_button)
        button_layout.addWidget(self.stop_button)
        button_layout.addStretch()
        
        left_layout.addWidget(self.camera_label)
        left_layout.addLayout(button_layout)
        
        # 右侧结果显示区域
        right_frame = QFrame()
        right_frame.setFrameStyle(QFrame.StyledPanel)
        right_frame.setFixedWidth(400)
        right_frame.setStyleSheet("QFrame { background-color: #fafafa; border: 2px solid #ddd; }")
        right_layout = QVBoxLayout(right_frame)
        
        # 标题
        title_label = QLabel("检测结果")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setFont(QFont("Arial", 16, QFont.Bold))
        title_label.setStyleSheet("QLabel { color: #333; padding: 10px; }")
        
        # 车牌数量显示
        self.count_label = QLabel("识别到的车牌数量: 0")
        self.count_label.setAlignment(Qt.AlignCenter)
        self.count_label.setFont(QFont("Arial", 12))
        self.count_label.setStyleSheet(
            "QLabel { "
            "background-color: #e3f2fd; "
            "border: 1px solid #2196f3; "
            "border-radius: 5px; "
            "padding: 8px; "
            "color: #1976d2; "
            "font-weight: bold; "
            "}"
        )
        
        # 滚动区域用于显示车牌结果
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setStyleSheet("QScrollArea { border: none; background-color: transparent; }")
        
        self.results_widget = QWidget()
        self.results_layout = QVBoxLayout(self.results_widget)
        self.results_layout.setAlignment(Qt.AlignTop)
        
        scroll_area.setWidget(self.results_widget)
        
        right_layout.addWidget(title_label)
        right_layout.addWidget(self.count_label)
        right_layout.addWidget(scroll_area)
        
        # 添加到主布局
        main_layout.addWidget(left_frame, 2)
        main_layout.addWidget(right_frame, 1)
        
        # 设置样式
        self.setStyleSheet("""
            QMainWindow {
                background-color: #f5f5f5;
            }
            QPushButton {
                background-color: #2196F3;
                color: white;
                border: none;
                padding: 8px 16px;
                border-radius: 4px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #1976D2;
            }
            QPushButton:pressed {
                background-color: #0D47A1;
            }
            QPushButton:disabled {
                background-color: #cccccc;
                color: #666666;
            }
        """)
    
    def init_detector(self):
        """初始化检测器"""
        model_path = os.path.join(os.path.dirname(__file__), "yolopart", "yolo11s-pose42.pt")
        self.detector = LicensePlateYOLO(model_path)
    
    def init_camera(self):
        """初始化摄像头线程"""
        self.camera_thread = CameraThread()
        self.camera_thread.frame_ready.connect(self.process_frame)
    
    def start_camera(self):
        """启动摄像头"""
        if self.camera_thread.start_camera():
            self.start_button.setEnabled(False)
            self.stop_button.setEnabled(True)
            self.camera_label.setText("摄像头启动中...")
        else:
            self.camera_label.setText("摄像头启动失败")
    
    def stop_camera(self):
        """停止摄像头"""
        self.camera_thread.stop_camera()
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.camera_label.setText("摄像头已停止")
        self.camera_label.clear()
    
    def process_frame(self, frame):
        """处理摄像头帧"""
        self.current_frame = frame.copy()
        
        # 进行车牌检测
        self.detections = self.detector.detect_license_plates(frame)
        
        # 在图像上绘制检测结果
        display_frame = self.draw_detections(frame.copy())
        
        # 转换为Qt格式并显示
        self.display_frame(display_frame)
        
        # 更新右侧结果显示
        self.update_results_display()
    
    def draw_detections(self, frame):
        """在图像上绘制检测结果"""
        return self.detector.draw_detections(frame, self.detections)
    
    def display_frame(self, frame):
        """显示帧到界面"""
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_frame.shape
        bytes_per_line = ch * w
        qt_image = QImage(rgb_frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
        
        pixmap = QPixmap.fromImage(qt_image)
        scaled_pixmap = pixmap.scaled(self.camera_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation)
        self.camera_label.setPixmap(scaled_pixmap)
    
    def update_results_display(self):
        """更新右侧结果显示"""
        # 更新车牌数量
        count = len(self.detections)
        self.count_label.setText(f"识别到的车牌数量: {count}")
        
        # 清除之前的结果
        for i in reversed(range(self.results_layout.count())):
            child = self.results_layout.itemAt(i).widget()
            if child:
                child.setParent(None)
        
        # 添加新的结果
        for i, detection in enumerate(self.detections):
            # 矫正车牌图像
            corrected_image = self.correct_license_plate(detection)
            
            # 获取车牌号（占位）
            plate_number = self.recognize_plate_number(corrected_image)
            
            # 创建车牌显示组件
            plate_widget = LicensePlateWidget(
                i + 1,
                detection['class_name'],
                corrected_image,
                plate_number
            )
            
            self.results_layout.addWidget(plate_widget)
    
    def correct_license_plate(self, detection):
        """矫正车牌图像"""
        if self.current_frame is None:
            return None
        
        # 检查是否为不完整检测
        if detection.get('incomplete', False):
            return None
        
        # 使用检测器的矫正方法
        return self.detector.correct_license_plate(
            self.current_frame, 
            detection['keypoints']
        )
    
    def recognize_plate_number(self, corrected_image):
        """识别车牌号"""
        if corrected_image is None:
            return "识别失败"
        
        try:
            # 使用OCR接口进行识别
            # 可以根据需要切换为CRNN: crnn_predict(corrected_image)
            result = ocr_predict(corrected_image)
            
            # 将字符列表转换为字符串
            if isinstance(result, list) and len(result) >= 7:
                return ''.join(result[:7])
            else:
                return "识别失败"
        except Exception as e:
            print(f"车牌号识别失败: {e}")
            return "识别失败"
    
    def closeEvent(self, event):
        """窗口关闭事件"""
        if self.camera_thread:
            self.camera_thread.stop_camera()
        event.accept()

def main():
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()