from loguru import logger
import subprocess as sp
from ultralytics import YOLO
import time, cv2, numpy as np, math
from traceback import format_exc
from DrGraph.utils.pull_push import NetStream
from DrGraph.utils.Helper import *
from DrGraph.utils.Constant import Constant
from zipfile import ZipFile

class YOLOTracker:
    def __init__(self, model_path):
        """
        初始化YOLOv11追踪器
        """
        self.model = YOLO(model_path)
        self.tracking_config = {
            "tracker": "appIOs/configs/yolo11/bytetrack.yaml", # "/home/thsw/jcq/projects/yolov11/ultralytics-main/ultralytics/cfg/trackers/bytetrack.yaml",
            "conf": 0.25,
            "iou": 0.45,
            "persist": True,
            "verbose": False
        }
        self.frame_count = 0
        self.processing_time = 0
        
    def process_frame(self, frame):
        """
        处理单帧图像，进行目标检测和追踪
        """
        start_time = time.time()
        
        try:
            # 执行YOLOv11目标检测和追踪
            results = self.model.track(
                source=frame,
                **self.tracking_config
            )
            
            # 获取第一个结果（因为只处理单张图片）
            result = results[0]
            
            # 绘制检测结果
            processed_frame = result.plot()
            
            # 计算处理时间
            self.processing_time = (time.time() - start_time) * 1000  # 转换为毫秒
            self.frame_count += 1
            
            # 打印检测信息（可选）
            if self.frame_count % 100 == 0:
                self._print_detection_info(result)
                
            return processed_frame, result
            
        except Exception as e:
            logger.error("YOLO处理异常: {}", format_exc())
            return frame, None
    
    def _print_detection_info(self, result):
        """
        打印检测信息
        """
        boxes = result.boxes
        if boxes is not None and len(boxes) > 0:
            detection_count = len(boxes)
            unique_ids = set()
            for box in boxes:
                if box.id is not None:
                    unique_ids.add(int(box.id[0]))
            
            logger.info(f"帧 {self.frame_count}: 检测到 {detection_count} 个目标, 追踪ID数: {len(unique_ids)}, 处理时间: {self.processing_time:.2f}ms")
        else:
            logger.info(f"帧 {self.frame_count}: 未检测到目标, 处理时间: {self.processing_time:.2f}ms")
    
class YOLOTrackerManager:
    def __init__(self, model_path, pull_url, push_url, request_id):
        self.pull_url = pull_url
        self.push_url = push_url
        self.request_id = request_id
        self.tracker = YOLOTracker(model_path)
        self.stream = None
        self.videoStream = None
        self.videoType = Constant.INPUT_NONE
        self.localFile = ''
        self.localPath = ''
        self.localFiles = []
        self._currentFrame = None
        self.totalFrames = 0
        self.frameChanged = False

    def _stop(self):
        if self.videoStream is not None:
            self.videoStream.release()
            self.videoStream = None
        if self.stream is not None:
            self.stream.clear_pull_p(self.stream.pull_p, self.request_id)
            self.stream = None
        self.localFile = ''
        self.localPath = ''
        self.localFiles = []
        self._currentFrame = None
        self.totalFrames = 0
        self._frameIndex = -1
        self.videoType = Constant.INPUT_NONE
        self.frameChanged = True

    def startLocalFile(self, fileName):
        self._stop()
        self.localFile = fileName
        self._frameIndex = -1

    def startLocalDir(self, dirName):
        self._stop()
        self.localPath = dirName
        self.localFiles = [os.path.join(dirName, f) for f in os.listdir(dirName) if f.endswith(('.jpg', '.jpeg', '.png'))]
        self.totalFrames = len(self.localFiles)
        Helper.App.progressMax = self.totalFrames
        self.localFiles.sort()
        logger.info("本地目录打开: {}, 总帧数: {}", dirName, self.totalFrames)
        self._frameIndex = 0

    def startLabelledZip(self, labelledPath, categoryPath):
        self._stop()
        self.localPath = labelledPath
        localFiles = ZipFile(labelledPath).namelist()
        _, self.totalFrames = Helper.getYoloLabellingInfo(categoryPath, localFiles, '')
        imagePath = categoryPath + 'images/'
        self.localFiles = [file for file in localFiles if imagePath in file]
        logger.info(f"标注压缩文件{labelledPath}的{categoryPath}集共有{self.totalFrames}帧, 有效帧数: {len(self.localFiles)}")
        self._frameIndex = 0
        Helper.App.progressMax = self.totalFrames

    def startUsbCamera(self, index = 0):
        self._stop()
        self.videoStream = cv2.VideoCapture(index)
        self.videoType = Constant.INPUT_USB_CAMERA
        Helper.Sleep(200)
        if not self.videoStream.isOpened():
            logger.error("无法打开USB摄像头: {}", index)
            self.videoType = Constant.INPUT_NONE
            return
        self.totalFrames = 0x7FFFFFFF

    def startLocalVideo(self, fileName):
        self._stop()
        self.videoStream = cv2.VideoCapture(fileName)
        self.videoType = Constant.INPUT_LOCAL_VIDEO
        Helper.Sleep(200)
        if not self.videoStream.isOpened():
            logger.error("无法打开本地视频流: {}", fileName)
            self.videoType = Constant.INPUT_NONE
            return
        try:
            total = int(self.videoStream.get(cv2.CAP_PROP_FRAME_COUNT))
        except Exception:
            total = 0
        self.totalFrames = total if total is not None else 0
        Helper.App.progressMax = self.totalFrames
        logger.info("本地视频打开: {}, 总帧数: {}", fileName, self.totalFrames)

    def startPull(self, url = ''):
        self._stop()
        if len(url) > 0:
            self.pull_url = url
        logger.info("拉流地址: {}", self.pull_url)
        self.stream = NetStream(self.pull_url, self.push_url, self.request_id)   
        self.stream.prepare_pull() 

    def getCurrentFrame(self):
        if self._currentFrame is None:
            self._currentFrame = self.nextFrame()
        if self._currentFrame is not None:
            return self._currentFrame.copy()
        return None
    currentFrame = Property_Rw(getCurrentFrame, None)
    
    def setFrameIndex(self, index):
        if self.videoStream is None and len(self.localFiles) == 0:
            return
        if self.videoStream is not None and self.videoType != Constant.INPUT_LOCAL_VIDEO:
            return
        if index < 0:
            index = 0
        if index >= self.totalFrames:
            index = self.totalFrames - 1
        if self.videoStream:                
            self.videoStream.set(cv2.CAP_PROP_POS_FRAMES, index)
        self._frameIndex = index - 1
        self._currentFrame = self.nextFrame()
        self.frameChanged = True
    frameIndex = Property_rW(setFrameIndex, 0)

    def getLabels(self):
        with ZipFile(self.localPath, 'r') as zip_ref:
            content = zip_ref.read(self.localFile)   
            content = content.decode('utf-8')
            return content
        return ''
    # 取得待分析的图像帧
    def getAnalysisFrame(self, nextFlag):
        frameChanged = self.frameChanged
        self.frameChanged = False
        if nextFlag: # 流式媒体
            self._currentFrame = self.nextFrame()
            self.frameChanged = True
        frame = self.currentFrame
        return frame.copy() if frame is not None else None, frameChanged

    def nextFrame(self):
        frame = None
        if self.stream:
            frame = self.stream.next_pull_frame()
        elif self.videoStream:
            ret, frame = self.videoStream.read()
            self._frameIndex += 1
            if not ret:
                self._frameIndex -= 1
                frame = None
        elif len(self.localFiles) > 0:
            if self.localPath.endswith('.zip'):
                index = -1
                for img_file in self.localFiles:
                    if '/images/' in img_file:
                        if index == self._frameIndex:                            
                            # logger.warning(f'Loading image from zip file: {img_file}')
                            try:
                                with ZipFile(self.localPath, 'r') as zip_ref:
                                    image_data = zip_ref.read(img_file)     
                                    nparr = np.frombuffer(image_data, np.uint8)
                                    frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
                                    self._frameIndex += 1
                                    lable_file = img_file.replace('/images/', '/labels/').replace('.jpg', '.txt').replace('.png', '.txt')
                                    self.localFile = lable_file
                            except Exception as e:
                                # logger.error(f"读取压缩文件 {self.localPath} 中的 {img_file} 失败: {e}")
                                frame = None    
                            break   
                        index += 1                     
            else:
                if self._frameIndex < 0:
                    self._frameIndex = 0
                if self._frameIndex >= len(self.localFiles):
                    self._frameIndex = 0
                if self._frameIndex < len(self.localFiles):
                    frame = cv2.imread(self.localFiles[self._frameIndex])
                if frame is None:
                    logger.error(f"无法读取目标目录 {self.localPath}中下标为 {self._frameIndex} 的视频文件 {self.localFiles[self._frameIndex]}")
                    self._frameIndex = -1
                    return
                self._frameIndex += 1
        elif self.localFile is not None and self.localFile != '':
            frame = cv2.imread(self.localFile)
            if frame is None:
                logger.error("无法读取本地视频文件: {}", self.localFile)
                return
        if frame is not None:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        if self.totalFrames > 0:
            Helper.App.progress = self._frameIndex
        return frame

    def test_yolo11_recognize(self, frame):
        processed_frame = self.process_frame_with_yolo(frame, self.request_id)
        return processed_frame
    
    def process_frame_with_yolo(self, frame, requestId):
        """
        使用YOLOv11处理帧
        """
        try:
            # 使用YOLO进行目标检测和追踪
            processed_frame, detection_result = self.tracker.process_frame(frame)
            
            # 在帧上添加处理信息
            fps_info = f"FPS: {1000/max(self.tracker.processing_time, 1):.1f}"
            cv2.putText(processed_frame, fps_info, (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            # 添加检测目标数量信息
            if detection_result and detection_result.boxes is not None:
                obj_count = len(detection_result.boxes)
                count_info = f"Objects: {obj_count}"
                cv2.putText(processed_frame, count_info, (10, 70), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            return processed_frame
            
        except Exception as e:
            logger.error("YOLO处理异常：{}, requestId:{}", format_exc(), requestId)
            # 如果处理失败，返回原帧
            return frame

    def get_gray_mask(self, frame):
        """
        生成灰度像素的掩码图
        灰度像素定义：三颜色分量差小于20
        """
        # 创建与原图大小相同的掩码图
        maskMat = np.zeros(frame.shape[:2], dtype=np.uint8)
        
        # 获取图像的三个颜色通道
        b, g, r = cv2.split(frame)
        
        r = r.astype(np.int16)
        g = g.astype(np.int16)
        b = b.astype(np.int16)
        # 计算任意两个颜色分量之间的差值
        diff_rg = np.abs(r - g)
        is_shadow = (b > r) & (b - r < 40)
        diff_rb = np.abs(r - b)
        diff_gb = np.abs(g - b)
        
        # 判断条件：三颜色分量差都小于20
        gray_pixels = (diff_rg < 20 ) & (diff_rb < 20| is_shadow) & (diff_gb < 20)
        
        # 将满足条件的像素在掩码图中设为255（白色）
        maskMat[gray_pixels] = 255
        
        return maskMat

    def debugLine(self, line, y_intersect):
        x1, y1, x2, y2 = line
        length = np.linalg.norm([x2 - x1, y2 - y1])
        # 计算线与水平线的夹角（度数）
        # 使用atan2计算弧度，再转换为度数
        angle_rad = math.atan2(y2 - y1, x2 - x1)
        angle_deg = math.degrees(angle_rad)
        # 调整角度范围到0-180度（平面角）
        if angle_deg < 0:
            angle_deg += 180
        # angle_deg = min(angle_deg, 180 - angle_deg)
        x_intersect = (x2 - x1) * (y_intersect - y1) / (y2 - y1) + x1
        return angle_deg, length, x_intersect
    def test_highway_recognize(self, frame, debugFlag = False):
        processed_frame = frame.copy()
        
        try:
            IGNORE_HEIGHT = 100
            y_intersect = frame.shape[0] / 2
            frame[:IGNORE_HEIGHT, :] = (255, 0, 0)

            gray_mask = self.get_gray_mask(frame)
            
            kernel = np.ones((5, 5), np.uint8)  # 使用形态学开运算（先腐蚀后膨胀）去除小噪声点
            gray_mask = cv2.erode(gray_mask, kernel)
            gray_mask = cv2.erode(gray_mask, kernel)

            # 过滤掉面积小于10000的区域
            contours, _ = cv2.findContours(gray_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            # 创建新的掩码图像，只保留面积大于等于10000的区域
            filtered_mask = np.zeros_like(gray_mask)
            for contour in contours:
                area = cv2.contourArea(contour)
                if area >= 10000:  # 填充满足条件的轮廓区域
                    cv2.fillPoly(filtered_mask, [contour], 255)
                        
            gray_mask = filtered_mask                                       # 使用过滤后的掩码替换原来的gray_mask         
            edges = cv2.Canny(frame, 100, 200)                              # 边缘检测
            road_edges = cv2.bitwise_and(edges, edges, mask=filtered_mask)  # 在过滤后的路面区域内进行边缘检测
            
            # 用color_mask过滤原图，得到待处理的图
            whiteLineMat = cv2.bitwise_and(processed_frame, processed_frame, mask=filtered_mask)
            whiteLineMat = cv2.cvtColor(whiteLineMat, cv2.COLOR_RGB2GRAY)   # 灰度化
            # sobel边缘检测
            whiteLineMat = cv2.Sobel(whiteLineMat, cv2.CV_8U, 1, 0, ksize=3)
            tempMat = whiteLineMat.copy()
            # whiteLineMat = cv2.Canny(whiteLineMat, 100, 200)
            lines = cv2.HoughLinesP(tempMat, 1, np.pi/180, threshold=100, minLineLength=100, maxLineGap=10)
            whiteLineMat = cv2.cvtColor(whiteLineMat, cv2.COLOR_GRAY2RGB)

            # logger.info(f"{lines.shape[0]} lines: ")
            # if lines is not None:
            #     for line in lines:
            #         x1, y1, x2, y2 = line[0]
            #         cv2.line(whiteLineMat, (x1, y1), (x2, y2), (255, 0, 0), 2)


            # 创建彩色掩码用于叠加（使用绿色标记识别出的路面）
            color_mask = cv2.cvtColor(gray_mask, cv2.COLOR_GRAY2RGB)
            color_mask[:] = (0, 255, 0)  # 设置为绿色
            color_mask = cv2.bitwise_and(color_mask, color_mask, mask=filtered_mask)
            
            # 先叠加路面绿色标记，再叠加白色线条红色标记
            overlay = cv2.addWeighted(processed_frame, 0.7, color_mask, 0.3, 0)

            # # 在road_edges的基础上，识别其中的实线
            # lines = cv2.HoughLinesP(road_edges, 1, np.pi/180, threshold=100, minLineLength=100, maxLineGap=10)
            # logger.info(f"{lines.shape[0]} lines: ")
            # linesWithAngle = []
            # # if lines is not None:
            # for index, line in enumerate(lines):    
            #     angle_deg, length, x_intersect = self.debugLine(line[0], y_intersect)
            #     linesWithAngle.append((line, angle_deg, x_intersect))                
            #     if debugFlag:
            #         logger.info(f'line {index + 1}: {line}, 线长：{length:.2f}, 夹角：{angle_deg:.2f}°， 交点：({x_intersect:.2f}, {y_intersect:.2f})')                

            # linesWithAngle进行聚类算法，按夹角分两类即可
            # 使用自定义的简单K-means聚类实现
            # line_data = np.array([[angle, x_intersect] for line, angle, x_intersect in linesWithAngle])
            # if len(line_data) > 0:
            #     labels = self._simple_kmeans(line_data, n_clusters=2, random_state=2, random_state=0)
            #     # 输出两类线的数目
            #     logger.info(f"聚类结果：{np.bincount(labels)}")
            #     if debugFlag:
            #         lines0 = [line for idx, line in enumerate(linesWithAngle) if labels[idx] == 0]
            #         lines1 = [line for idx, line in enumerate(linesWithAngle) if labels[idx] == 1]
            #         # 取得lines0中所有线段并输出日志信息
            #         for index, line in enumerate(lines0):   
            #             angle_deg, length, x_intersect = self.debugLine(line[0][0], y_intersect)
            #             logger.info(f'聚类0: {line[0]}, 线长：{length:.2f}, 夹角：{angle_deg:.2f}°， 交点：({x_intersect:.2f}, {y_intersect:.2f})')                

            #         for index, line in enumerate(lines1):   
            #             angle_deg, length, x_intersect = self.debugLine(line[0][0], y_intersect)
            #             logger.info(f'聚类1: {line[0]}, 线长：{length:.2f}, 夹角：{angle_deg:.2f}°， 交点：({x_intersect:.2f}, {y_intersect:.2f})')                

            #     # 保留数量多的类别
            #     dominant_cluster = np.argmax(np.bincount(labels))
            #     # 绘制dominant_cluster类别的线
            #     dominant_lines = [line for idx, line in enumerate(linesWithAngle) if labels[idx] == dominant_cluster]
                
            #     for line, angle, x_intersect in dominant_lines:
            #         cv2.line(overlay, (int(line[0][0]), int(line[0][1])), (int(line[0][2]), int(line[0][3])), (255, 0, 0), 2)
              
            return overlay, color_mask, whiteLineMat # cv2.cvtColor(whiteLineMat, cv2.COLOR_GRAY2RGB) # cv2.cvtColor(road_edges, cv2.COLOR_GRAY2RGB)
            
        except Exception as e:
            logger.error("路面识别异常：{}", format_exc())
            # 如果处理失败，返回原始帧
            return processed_frame

    # def _simple_kmeans(self, data, n_clusters=2, max_iter=100, random_state=0):
    #     """
    #     使用K-means算法对数据进行聚类
        
    #     参数:
    #         data: array-like, 形状为 (n_samples, n_features) 的输入数据
    #         n_clusters: int, 聚类数量，默认为2
    #         max_iter: int, 最大迭代次数，默认为100
    #         random_state: int, 随机种子，用于初始化质心，默认为0
            
    #     返回:
    #         labels: array, 形状为 (n_samples,) 的聚类标签数组
    #     """
    #     np.random.seed(random_state)
        
    #     # 随机选择初始质心
    #     centroids_idx = np.random.choice(len(data), size=n_clusters, replace=False)
    #     centroids = data[centroids_idx].copy()
        
    #     # 迭代优化质心位置
    #     for _ in range(max_iter):
    #         # 为每个数据点分配最近的质心标签
    #         labels = np.zeros(len(data), dtype=int)
    #         for i, point in enumerate(data):                
    #             distancesi=ids - point, ax(centroids - point, axis=1)   ce置为-                # 情况如果d sfnpcnsy>d9e，则置为>180 -9dis   ini作为新质心
    #                 new_centroids[c] = data[np.random.choice(len(data))]
            
    #         # 检查收敛条件
    #         if np.allclose(centroids, new_centroids):
    #             break
            
    #         centroids = new_centroids
        
    #     return labels

