import json
import os
import subprocess
import time

import cv2
import numpy as np
import torch
from ultralytics import YOLO

from middleware.read_srt import parse_srt_file
from touying.ImageReproject_python.cal_func import red_line_reproject
from touying.ImageReproject_python.img_types import Point
from yolo.cv_multi_model_back_video import is_point_in_polygon

# 配置参数
local_video_path = r"/mnt/mydisk1/dj/ai2/middleware/DJI_20250912152112_0001_V.mp4"
srt_path = r"/mnt/mydisk1/dj/ai2/middleware/DJI_20250912152112_0001_V.mp4.srt"
invade_file_path = r"/mnt/mydisk1/dj/ai2/test/高压线-0826.geojson"
# invade_file_path = r"C:\Users\14867\xwechat_files\wxid_lqq76m3bwzja21_94e1\msg\file\2025-08\营业线-0826.geojson"
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model_paths = [
    # r"D:\project\AI-PYTHON\Ai_tottle\pt\best.pt",  # 模型1
    r"/mnt/mydisk1/dj/ai2/pt/GDCL.pt"  # 模型2（可扩展）
]
output_dir = r"/mnt/mydisk1/dj/ai2/test/test_pic_1"
os.makedirs(output_dir, exist_ok=True)


def load_models(model_paths, device):
    """加载多个YOLO模型"""
    models = []
    for path in model_paths:
        model = YOLO(path).to(device)
        models.append(model)
    return models


def merge_results(results_list, conf_threshold=0.5, iou_threshold=0.5):
    """整合多个模型的推理结果"""
    all_boxes = []
    for results in results_list:
        for result in results:
            for box in result.boxes:
                if box.conf[0] > conf_threshold:
                    x1, y1, x2, y2 = map(float, box.xyxy[0].tolist())
                    conf = float(box.conf[0])
                    class_id = int(box.cls[0])
                    all_boxes.append([x1, y1, x2, y2, conf, class_id])

    if not all_boxes:
        return np.empty((0, 6))  # 返回空数组而不是None

    # 转换为NumPy数组
    boxes = np.array(all_boxes)

    try:
        import torchvision
        boxes_tensor = torch.tensor(boxes[:, :4])  # 坐标
        scores = torch.tensor(boxes[:, 4])  # 置信度
        keep = torchvision.ops.nms(boxes_tensor, scores, iou_threshold)
        merged_boxes = boxes[keep]
    except ImportError:
        print("Warning: torchvision not installed, skipping NMS.")
        merged_boxes = boxes

    return merged_boxes  # 确保返回的是二维数组


def cal():
    # 加载多边形数据
    list_points = []
    with open(invade_file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
        features = data.get('features', [])
        for polygon in features:
            geometry = polygon.get('geometry', {})
            coordinates = geometry.get('coordinates', [])
            points = [Point(coord[1], coord[0], coord[2], key) for key, coord in enumerate(coordinates[0])]
            list_points.append(points)

    # # 提取视频字幕（略，保持原逻辑）
    # dir_name = os.path.dirname(local_video_path)
    # srt_name = os.path.basename(local_video_path) + ".srt"
    # srt_path = os.path.join(dir_name, srt_name)
    # command = ["ffmpeg", "-i", local_video_path, "-map", "0:s:0", "-c:s", "srt", srt_path]
    # try:
    #     subprocess.run(command, check=True)
    # except Exception as e:
    #     print(f"字幕提取失败: {e}")

    srt_list = parse_srt_file(srt_path)
    cap = cv2.VideoCapture(local_video_path)
    if not cap.isOpened():
        print("Error: 无法打开视频")
        return

    # 加载多模型
    models = load_models(model_paths, device)
    frame_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break

        img_height, img_width = frame.shape[:2]
        frame_count += 1

        # 处理SRT字幕（保持原逻辑）
        if frame_count < len(srt_list):
            dji_srt = srt_list[frame_count]
            gimbal_pitch = dji_srt.gb_pitch
            gb_yaw = dji_srt.gb_yaw
            air_height = dji_srt.abs_alt
            cam_longitude = dji_srt.longitude
            cam_latitude = dji_srt.latitude

        # 多模型推理
        all_results = []
        for model in models:
            try:
                results = model(frame)
                all_results.append(results)
            except Exception as e:
                print(f"模型推理错误: {e}")

        # 修改后的结果处理部分
        merged_boxes = merge_results(all_results)

        # 确保merged_boxes是二维数组
        if len(merged_boxes.shape) == 1:
            merged_boxes = np.expand_dims(merged_boxes, 0)


        # 绘制多边形（保持原逻辑）
        for points in list_points:
            red_result = red_line_reproject(gb_yaw, gimbal_pitch, dji_srt.gb_roll, air_height,
                                            cam_longitude, cam_latitude, img_width, img_height, points)
            point_list = [[int(p["u"]), int(p["v"])] for p in red_result]
            cv2.polylines(frame, [np.array(point_list, dtype=np.int32)], isClosed=True, color=(0, 0, 255), thickness=2)
            message_point=[]
            message_out_point=[]
            # 可视化合并后的结果
            for box in merged_boxes:
                # 现在box应该是长度为6的数组
                x1, y1, x2, y2, conf, class_id = map(int, box[:6])  # 取前6个元素
                # cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                # cv2.putText(frame, f"ID:{class_id} {conf:.2f}", (x1, y1 - 10),
                #             cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                point_x = (x1 + x2) / 2
                point_y = (y1 + y2) / 2
                is_invade = is_point_in_polygon(point_x, point_y, red_result)
                if is_invade:

                    message_point.append({
                        # "u": point_x,
                        # "v": point_y,
                        # "class_name": class_name
                        "box": [x1, y1, x2, y2]
                    })
                else:
                    message_out_point.append({
                        # "u": point_x,
                        # "v": point_y,
                        # "class_name": class_name
                        "box": [x1, y1, x2, y2]
                    })

            for point in message_point:
                cv2.rectangle(frame, (point["box"][0], point["box"][1]), (point["box"][2], point["box"][3]),
                              (0, 255, 255), 2)
            # for point in message_out_point:
            #     cv2.rectangle(frame, (point["box"][0], point["box"][1]), (point["box"][2], point["box"][3]),
            #                   (0, 255, 0), 2)

        # 显示并保存结果
#        cv2.imshow("Multi-Model Detection", frame)
        cv2.imwrite(os.path.join(output_dir, f"{frame_count}.png"), frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    cal()