#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
车牌识别接口模块
预留接口，可接入各种OCR模型进行车牌号识别
"""

import cv2
import numpy as np
from typing import List, Optional, Dict, Any
from abc import ABC, abstractmethod

class PlateRecognizerInterface(ABC):
    """车牌识别接口基类"""
    
    @abstractmethod
    def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
        """
        识别车牌号
        
        Args:
            plate_image: 车牌图像 (BGR格式)
            
        Returns:
            识别结果字典，包含:
            {
                'text': str,           # 识别的车牌号
                'confidence': float,   # 置信度 (0-1)
                'success': bool        # 是否识别成功
            }
        """
        pass
        
    @abstractmethod
    def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
        """
        批量识别车牌号
        
        Args:
            plate_images: 车牌图像列表
            
        Returns:
            识别结果列表
        """
        pass

class MockPlateRecognizer(PlateRecognizerInterface):
    """模拟车牌识别器（用于测试）"""
    
    def __init__(self):
        self.mock_plates = [
            "京A12345", "沪B67890", "粤C11111", "川D22222", 
            "鲁E33333", "苏F44444", "浙G55555", "闽H66666"
        ]
        self.call_count = 0
        
    def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
        """
        模拟识别单个车牌
        
        Args:
            plate_image: 车牌图像
            
        Returns:
            模拟识别结果
        """
        # 模拟处理时间
        import time
        time.sleep(0.01)  # 10ms模拟处理时间
        
        # 简单的图像质量检查
        if plate_image is None or plate_image.size == 0:
            return {
                'text': '',
                'confidence': 0.0,
                'success': False
            }
            
        # 检查图像尺寸
        height, width = plate_image.shape[:2]
        if width < 50 or height < 20:
            return {
                'text': '',
                'confidence': 0.3,
                'success': False
            }
            
        # 模拟识别结果
        plate_text = self.mock_plates[self.call_count % len(self.mock_plates)]
        confidence = 0.85 + (self.call_count % 10) * 0.01  # 0.85-0.94
        
        self.call_count += 1
        
        return {
            'text': plate_text,
            'confidence': confidence,
            'success': True
        }
        
    def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
        """
        批量识别车牌
        
        Args:
            plate_images: 车牌图像列表
            
        Returns:
            识别结果列表
        """
        results = []
        for plate_image in plate_images:
            result = self.recognize(plate_image)
            results.append(result)
        return results

class PaddleOCRRecognizer(PlateRecognizerInterface):
    """PaddleOCR车牌识别器（示例实现）"""
    
    def __init__(self, use_gpu: bool = True):
        """
        初始化PaddleOCR识别器
        
        Args:
            use_gpu: 是否使用GPU
        """
        self.use_gpu = use_gpu
        self.ocr = None
        self._init_ocr()
        
    def _init_ocr(self):
        """初始化OCR模型"""
        try:
            # 这里可以接入PaddleOCR
            # from paddleocr import PaddleOCR
            # self.ocr = PaddleOCR(use_angle_cls=True, lang='ch', use_gpu=self.use_gpu)
            print("PaddleOCR初始化完成（示例代码，需要安装PaddleOCR）")
        except ImportError:
            print("PaddleOCR未安装，使用模拟识别器")
            self.ocr = None
            
    def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
        """
        使用PaddleOCR识别车牌
        
        Args:
            plate_image: 车牌图像
            
        Returns:
            识别结果
        """
        if self.ocr is None:
            # 回退到模拟识别
            mock_recognizer = MockPlateRecognizer()
            return mock_recognizer.recognize(plate_image)
            
        try:
            # 使用PaddleOCR进行识别
            results = self.ocr.ocr(plate_image, cls=True)
            
            if results and len(results) > 0 and results[0]:
                # 提取文本和置信度
                text_results = []
                for line in results[0]:
                    text = line[1][0]
                    confidence = line[1][1]
                    text_results.append((text, confidence))
                    
                # 选择置信度最高的结果
                if text_results:
                    best_result = max(text_results, key=lambda x: x[1])
                    return {
                        'text': best_result[0],
                        'confidence': best_result[1],
                        'success': True
                    }
                    
        except Exception as e:
            print(f"PaddleOCR识别失败: {e}")
            
        return {
            'text': '',
            'confidence': 0.0,
            'success': False
        }
        
    def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
        """
        批量识别
        
        Args:
            plate_images: 车牌图像列表
            
        Returns:
            识别结果列表
        """
        results = []
        for plate_image in plate_images:
            result = self.recognize(plate_image)
            results.append(result)
        return results

class TesseractRecognizer(PlateRecognizerInterface):
    """Tesseract车牌识别器（示例实现）"""
    
    def __init__(self, lang: str = 'chi_sim+eng'):
        """
        初始化Tesseract识别器
        
        Args:
            lang: 识别语言
        """
        self.lang = lang
        self.tesseract_available = self._check_tesseract()
        
    def _check_tesseract(self) -> bool:
        """检查Tesseract是否可用"""
        try:
            import pytesseract
            return True
        except ImportError:
            print("pytesseract未安装，使用模拟识别器")
            return False
            
    def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
        """
        使用Tesseract识别车牌
        
        Args:
            plate_image: 车牌图像
            
        Returns:
            识别结果
        """
        if not self.tesseract_available:
            # 回退到模拟识别
            mock_recognizer = MockPlateRecognizer()
            return mock_recognizer.recognize(plate_image)
            
        try:
            import pytesseract
            
            # 图像预处理
            processed_image = self._preprocess_image(plate_image)
            
            # 使用Tesseract识别
            text = pytesseract.image_to_string(
                processed_image, 
                lang=self.lang,
                config='--psm 8 --oem 3 -c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ京沪粤川鲁苏浙闽'
            )
            
            # 清理识别结果
            text = text.strip().replace(' ', '').replace('\n', '')
            
            if text and len(text) >= 5:  # 车牌号至少5位
                return {
                    'text': text,
                    'confidence': 0.8,  # Tesseract不直接提供置信度
                    'success': True
                }
                
        except Exception as e:
            print(f"Tesseract识别失败: {e}")
            
        return {
            'text': '',
            'confidence': 0.0,
            'success': False
        }
        
    def _preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """图像预处理"""
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
            
        # 调整尺寸
        height, width = gray.shape
        if width < 200:
            scale = 200 / width
            new_width = int(width * scale)
            new_height = int(height * scale)
            gray = cv2.resize(gray, (new_width, new_height))
            
        # 二值化
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        return binary
        
    def batch_recognize(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
        """
        批量识别
        
        Args:
            plate_images: 车牌图像列表
            
        Returns:
            识别结果列表
        """
        results = []
        for plate_image in plate_images:
            result = self.recognize(plate_image)
            results.append(result)
        return results

class PlateRecognizerManager:
    """车牌识别管理器"""
    
    def __init__(self, recognizer_type: str = 'mock'):
        """
        初始化识别管理器
        
        Args:
            recognizer_type: 识别器类型 ('mock', 'paddleocr', 'tesseract')
        """
        self.recognizer_type = recognizer_type
        self.recognizer = self._create_recognizer(recognizer_type)
        
    def _create_recognizer(self, recognizer_type: str) -> PlateRecognizerInterface:
        """创建识别器"""
        if recognizer_type == 'mock':
            return MockPlateRecognizer()
        elif recognizer_type == 'paddleocr':
            return PaddleOCRRecognizer()
        elif recognizer_type == 'tesseract':
            return TesseractRecognizer()
        else:
            print(f"未知的识别器类型: {recognizer_type}，使用模拟识别器")
            return MockPlateRecognizer()
            
    def recognize_plates(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
        """
        识别车牌列表
        
        Args:
            plate_images: 车牌图像列表
            
        Returns:
            识别结果列表
        """
        if not plate_images:
            return []
            
        return self.recognizer.batch_recognize(plate_images)
        
    def switch_recognizer(self, recognizer_type: str):
        """
        切换识别器
        
        Args:
            recognizer_type: 新的识别器类型
        """
        self.recognizer_type = recognizer_type
        self.recognizer = self._create_recognizer(recognizer_type)
        print(f"已切换到识别器: {recognizer_type}")
        
    def get_recognizer_info(self) -> Dict[str, Any]:
        """
        获取识别器信息
        
        Returns:
            识别器信息
        """
        return {
            'type': self.recognizer_type,
            'class': self.recognizer.__class__.__name__
        }
    
    def preprocess_blue_plate(self, plate_image: np.ndarray, original_image: np.ndarray, bbox: List[int]) -> np.ndarray:
        """
        蓝色车牌预处理：倾斜矫正
        
        Args:
            plate_image: 切割后的车牌图像
            original_image: 原始图像
            bbox: 边界框坐标 [x1, y1, x2, y2]
            
        Returns:
            矫正后的车牌图像
        """
        try:
            # 从原图中提取车牌区域
            x1, y1, x2, y2 = bbox
            roi = original_image[y1:y2, x1:x2]
            
            # 获取蓝色车牌的二值图像
            bin_img = self._get_blue_img_bin(roi)
            
            # 倾斜矫正
            corrected_img = self._deskew_plate(bin_img, roi)
            
            return corrected_img
        except Exception as e:
            print(f"蓝色车牌预处理失败: {e}")
            return plate_image
    
    def _get_blue_img_bin(self, img: np.ndarray) -> np.ndarray:
        """
        获取蓝色车牌的二值图像
        """
        # 掩膜：BGR通道，若像素B分量在 100~255 且 G分量在 0~190 且 R分量在 0~140 置255（白色），否则置0（黑色）
        mask_bgr = cv2.inRange(img, (100, 0, 0), (255, 190, 140))
        
        # 转换成 HSV 颜色空间
        img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(img_hsv)  # 分离通道 色调(H)，饱和度(S)，明度(V)
        mask_s = cv2.inRange(s, 80, 255)  # 取饱和度通道进行掩膜得到二值图像
        
        # 与操作，两个二值图像都为白色才保留，否则置黑
        rgbs = mask_bgr & mask_s
        
        # 核的横向分量大，使车牌数字尽量连在一起
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 3))
        img_rgbs_dilate = cv2.dilate(rgbs, kernel, 3)  # 膨胀，减小车牌空洞
        
        return img_rgbs_dilate
    
    def _order_points(self, pts: np.ndarray) -> np.ndarray:
        """
        将四点按 左上、右上、右下、左下 排序
        """
        rect = np.zeros((4, 2), dtype="float32")
        s = pts.sum(axis=1)
        rect[0] = pts[np.argmin(s)]  # 左上
        rect[2] = pts[np.argmax(s)]  # 右下
        
        diff = np.diff(pts, axis=1)
        rect[1] = pts[np.argmin(diff)]  # 右上
        rect[3] = pts[np.argmax(diff)]  # 左下
        
        return rect
    
    def _deskew_plate(self, bin_img: np.ndarray, original_roi: np.ndarray) -> np.ndarray:
        """
        车牌倾斜矫正
        
        Args:
            bin_img: 二值图像
            original_roi: 原始ROI区域
            
        Returns:
            矫正后的原始图像（未被掩模，但经过旋转和切割）
        """
        try:
            # 找最大轮廓
            cnts, _ = cv2.findContours(bin_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            if not cnts:
                return original_roi
                
            c = max(cnts, key=cv2.contourArea)
            
            # 最小外接矩形
            rect = cv2.minAreaRect(c)
            box = cv2.boxPoints(rect)
            box = np.array(box, dtype="float32")
            
            # 排序四个点
            pts_src = self._order_points(box)
            
            # 计算目标矩形宽高
            (tl, tr, br, bl) = pts_src
            widthA = np.linalg.norm(br - bl)
            widthB = np.linalg.norm(tr - tl)
            maxWidth = int(max(widthA, widthB))
            
            heightA = np.linalg.norm(tr - br)
            heightB = np.linalg.norm(tl - bl)
            maxHeight = int(max(heightA, heightB))
            
            # 确保尺寸合理
            if maxWidth < 10 or maxHeight < 10:
                return original_roi
            
            # 目标点集合
            pts_dst = np.array([
                [0, 0],
                [maxWidth - 1, 0],
                [maxWidth - 1, maxHeight - 1],
                [0, maxHeight - 1]], dtype="float32")
            
            # 透视变换
            M = cv2.getPerspectiveTransform(pts_src, pts_dst)
            warped = cv2.warpPerspective(original_roi, M, (maxWidth, maxHeight))
            
            return warped
        except Exception as e:
            print(f"车牌矫正失败: {e}")
            return original_roi