import cv2
import numpy as np
from ultralytics import YOLO
import os
import json
import argparse


class Target2D:
    """二维坐标下的叶尖跟踪器"""

    def __init__(self, positions=None, hub=None):
        if positions is None:
            positions = np.array([[-100, -100], [-200, -200], [-300, -300]])
        if hub is None:
            hub = np.array([400, 300])

        self.positions = np.array(positions, dtype=float)  # 叶尖位置
        self.hub = np.array(hub, dtype=float)  # 轮毂位置

    def update(self, new_positions, new_hub=None):
        """更新叶尖位置，使用最近点匹配算法"""
        if new_positions is None or len(new_positions) == 0:
            return self.positions, []

        # 确保新位置是二维数组
        new_positions = np.array(new_positions, dtype=float)
        if new_positions.ndim == 1:
            new_positions = new_positions.reshape(1, -1)

        # 更新轮毂位置
        if new_hub is not None:
            self.hub = np.array(new_hub, dtype=float)

        # 如果还没有初始化位置，使用新位置
        if self.positions.size == 0:
            self.positions = new_positions
            return self.positions
        # 创建位置列表副本
        new_pos_list = new_positions.tolist()
        update_positons = self.positions.copy()
        used = set()
        # 对于每个新位置，找到最近的点
        for i, new_pos in enumerate(new_pos_list):
            # 计算所有距离
            distances = np.linalg.norm(self.positions - new_pos, axis=1)
            available_indices = [j for j in range(len(distances)) if j not in used]
            # 找到最近点的索引
            closest_idx = min(available_indices, key=lambda j: distances[j])

            # 记录ID并更新位置
            update_positons[closest_idx] = new_pos  # new_pos

        # 更新位置
        self.positions = np.array(update_positons)

        return self.positions


class ClearanceCalculator:
    def __init__(self, model_path, scale_factor=0.02,width=544, height=960):
        self.size = (width, height)
        # 初始化模型配置
        self.model = YOLO(model_path)
        self.cap = None
        self.video_path = ""
        self.is_playing = False
        self.scale_factor = scale_factor
        self.reference_pixel = 0.0
        self.reference_actual = 0.0
        self.current_clearance = 0.0
        self.use_camera = False
        self.hub_x_range = 100  # 轮毂x坐标的左右区间（像素）
        self.hub_y_range = 50
        self.inside_region = False  # 轨迹是否在轮毂区域内
        self.entered_from_left = False  # 轨迹是否从左侧进入
        self.last_pixel_clearance = 0.0  # 保存上一次的像素净空
        self.last_actual_clearance = 0.0  # 保存上一次的实际净空
        # 轨迹相关变量
        self.trajectories = []
        # self.lowest_points = []
        self.calculate_points = []
        self.hub_position = None
        self.last_tip_positions = []
        self.frame_count = 0
        self.nearest_tip_trajectory = []  # 只存储与轮毂x坐标最近的叶尖轨迹
        self.traj_cnt = 0
        self.tip_clearances = {
            0: None,  # tip1
            1: None,  # tip2
            2: None  # tip3
        }
        # 净空历史记录
        self.clearance_history = {0: [], 1: [], 2: []}

        # 初始化轨迹颜色
        self.tip_colors = [
            (255, 0, 0),  # 红色
            (0, 255, 0),  # 绿色
            (0, 0, 255)  # 蓝色
        ]

        # 初始化2D跟踪器
        self.tracker = Target2D()

    def set_hub_x_range(self, x_value, y_value=None):
        """设置轮毂x和y坐标的区间"""
        self.hub_x_range = x_value
        if y_value is not None:
            self.hub_y_range = y_value

    def load_video(self, path):
        """加载视频文件"""
        self.use_camera = False
        self.video_path = path
        self.cap = cv2.VideoCapture(path)
        return self.cap.isOpened()

    def process_frame(self, frame):
        """处理视频帧并检测叶尖和轮毂"""
        frame = cv2.resize(frame, self.size)
        # 使用模型进行检测
        results = self.model(frame, verbose=False)

        # 初始化检测结果字典
        detected = {'tip': [], 'hub': []}

        # 解析检测结果
        for result in results:
            boxes = result.boxes
            classes = boxes.cls
            class_names = result.names

            for i in range(len(boxes)):
                class_id = int(classes[i])
                box = boxes.xyxy[i]
                x1, y1, x2, y2 = box.tolist()
                class_name = class_names[class_id]

                # 只关注叶尖(tip)和轮毂(hub)
                if class_name in detected:
                    # 使用底部中心点作为位置
                    if class_name == 'tip':
                        tip_x = (x1 + x2) / 2
                        tip_y = y2  # 使用底部作为叶尖位置
                        detected[class_name].append([tip_x, tip_y])
                    else:  # hub
                        hub_x = (x1 + x2) / 2
                        hub_y = (y1 + y2) / 2
                        detected[class_name].append([hub_x, hub_y])

        # 限制检测数量
        if len(detected['tip']) > 3:
            detected['tip'] = detected['tip'][:3]
        if len(detected['hub']) > 1:
            detected['hub'] = detected['hub'][:1]

        return frame, detected['tip'], detected['hub']

    def check_trajectory_state(self, tip, hub, thresh=20):
        """检查轨迹是否进入或离开轮毂区域"""
        hub_x = hub[0] if hub else None
        if hub_x is None or tip is None:
            return None, None

        # 检查轨迹是否在轮毂区域内
        inside = (hub_x - self.hub_x_range) < tip[0] < (hub_x + self.hub_x_range)
        # 如果轨迹进入区域，并且是从左侧进入
        if inside and not self.inside_region:
            self.inside_region = True
            self.entered_from_left = tip[0] >= (hub_x - self.hub_x_range)

        if inside and self.inside_region and self.entered_from_left and (hub_x - thresh) < tip[0] < (hub_x + thresh):
            distance, actual_distance = self.calculate_clearance()
            return distance, actual_distance
        # 如果轨迹离开区域，并且是从左侧进入且从右侧离开
        if not inside and self.inside_region and self.entered_from_left and tip[0] > (hub_x + self.hub_x_range):
            self.inside_region = False
            self.traj_cnt += 1
        return None, None

    def update_trajectories(self, tips, hub):
        """更新叶尖轨迹，只关注轮毂x和y坐标附近的叶尖"""
        # 使用2D跟踪器更新位置
        self.tracker.update(tips, hub)
        sorted_tips = self.tracker.positions

        # 如果没有轮毂位置，无法判断附近
        if hub is None or len(hub) == 0:
            return None, None

        hub_x = hub[0]  # 轮毂的x坐标
        hub_y = hub[1]  # 轮毂的y坐标

        # 初始化轨迹
        if not self.trajectories:
            for _ in range(len(sorted_tips)):
                self.trajectories.append([])

        # 找到与轮毂x和y坐标都最近的叶尖
        nearest_tip = None
        min_distance = float('inf')

        for tip in sorted_tips:
            tip_x, tip_y = tip
            # 检查是否在轮毂x坐标左右区间内
            if abs(tip_x - hub_x) > self.hub_x_range + 20:
                continue
            # 检查是否大于轮毂y坐标加y区间
            if tip_y < hub_y + self.hub_y_range:
                continue

            # 计算与轮毂的综合距离（考虑x）
            distance = abs(tip_x - hub_x)  # + (tip_y - hub_y) ** 2
            if distance < min_distance:
                min_distance = distance
                nearest_tip = tip

        # 更新最近的叶尖轨迹
        if nearest_tip is not None:
            if not self.nearest_tip_trajectory:
                self.nearest_tip_trajectory = [nearest_tip]
            else:
                self.nearest_tip_trajectory.append(nearest_tip)

            # 限制轨迹长度
            if len(self.nearest_tip_trajectory) > 200:
                self.nearest_tip_trajectory.pop(0)

        return self.check_trajectory_state(nearest_tip, hub)

    def find_closest_tip_x(self):
        """找到与轮毂x坐标最接近的叶尖x坐标的叶尖"""
        if not self.nearest_tip_trajectory:
            return None
        hub_x = self.hub_position[0] if self.hub_position else None
        if hub_x is None:
            return None
        trajectories = np.array(self.nearest_tip_trajectory[:][0])
        distances = np.abs(trajectories - hub_x)
        idx = np.argmin(distances)
        tip = self.nearest_tip_trajectory[idx]
        self.calculate_points = [tip]  # 只保留一个最低点
        return tip

    def calculate_clearance(self):
        """计算叶尖到轮毂的净空"""
        if not self.nearest_tip_trajectory or self.hub_position is None:
            return None, None

        hub_position = self.hub_position
        tip_x, tip_y = self.find_closest_tip_x()
        self.calculate_points = [[tip_x, tip_y]]  # 只保留一个最低点
        hub_x, hub_y = hub_position
        distance = np.sqrt((tip_x - hub_x) ** 2 + (tip_y - hub_y) ** 2)

        # 更新UI显示净空
        if self.scale_factor:
            actual_distance = distance * self.scale_factor
            tip_idx = self.traj_cnt % 3  # 根据实际匹配逻辑
            self.tip_clearances[tip_idx] = actual_distance
            # 记录到历史
            self.clearance_history[tip_idx].append(actual_distance)
            return distance, actual_distance

    def draw_annotations(self, frame):
        """在帧上绘制轨迹和标记"""
        # 绘制轮毂位置
        if self.hub_position:
            hub_x, hub_y = self.hub_position
            # 绘制轮毂区域标记
            cv2.rectangle(frame,
                          (int(hub_x - self.hub_x_range), int(hub_y - 10)),
                          (int(hub_x + self.hub_x_range), int(hub_y + 10)),
                          (0, 255, 255), 2)
            cv2.circle(frame, (int(hub_x), int(hub_y)), 8, (0, 255, 255), -1)
            cv2.putText(frame, "Hub", (int(hub_x) + 15, int(hub_y)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)

        # 绘制最近的叶尖轨迹
        if len(self.nearest_tip_trajectory) >= 2:
            # 绘制轨迹线（黄色）
            for j in range(1, len(self.nearest_tip_trajectory)):
                start_point = (int(self.nearest_tip_trajectory[j - 1][0]),
                               int(self.nearest_tip_trajectory[j - 1][1]))
                end_point = (int(self.nearest_tip_trajectory[j][0]),
                             int(self.nearest_tip_trajectory[j][1]))
                cv2.line(frame, start_point, end_point, (0, 255, 255), 1)
        # 绘制三个叶尖的净空值
        y_offset = 30
        for i in range(3):
            clearance = self.tip_clearances.get(i, None)
            if clearance is not None:
                text = f"tip{i + 1}: {clearance:.2f} m"
            else:
                text = f"tip{i + 1}: -- m"
            cv2.putText(frame, text, (10, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
            y_offset += 25

        # 绘制其他叶尖（不绘制轨迹）
        if self.tracker.positions.size > 0:
            for i, tip in enumerate(self.tracker.positions):
                if not any(np.array_equal(tip, self.nearest_tip_trajectory[-1]) for p in self.nearest_tip_trajectory if
                           len(self.nearest_tip_trajectory) > 0):
                    color = self.tip_colors[i % len(self.tip_colors)]
                    cv2.circle(frame, (int(tip[0]), int(tip[1])), 6, color, -1)
                    cv2.putText(frame, f"Tip {i + 1}", (int(tip[0]) + 15, int(tip[1])),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

        return frame

    def process_video(self, output_dir):
        """处理整个视频并保存结果"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        last_frame = None
        frame_count = 0

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

            # 处理帧
            processed_frame, tips, hubs = self.process_frame(frame)

            # 更新轮毂位置
            if hubs:
                self.hub_position = hubs[0]

            # 更新轨迹
            if tips and self.hub_position is not None:
                pixel_clearance, actual_clearance = self.update_trajectories(tips, self.hub_position)

            # 绘制注释
            annotated_frame = self.draw_annotations(processed_frame)
            last_frame = annotated_frame.copy()
            frame_count += 1

            # 显示处理进度
            if frame_count % 10 == 0:
                print(f"已处理 {frame_count} 帧...")

        # 保存最后一帧
        last_frame_path = os.path.join(output_dir, "last_frame.jpg")
        cv2.imwrite(last_frame_path, last_frame)
        print(f"最后一帧已保存至: {last_frame_path}")

        # 准备结果数据
        result = {
            "clearance_history": {
                "tip0": self.clearance_history[0],
                "tip1": self.clearance_history[1],
                "tip2": self.clearance_history[2]
            },
            "last_frame_clearances": {
                "tip0": self.tip_clearances[0],
                "tip1": self.tip_clearances[1],
                "tip2": self.tip_clearances[2]
            },
            "last_frame_image": last_frame_path
        }

        # 保存JSON结果
        json_path = os.path.join(output_dir, "results.json")
        with open(json_path, 'w') as f:
            json.dump(result, f, indent=4)
        print(f"结果已保存至: {json_path}")

        return result


def main():
    parser = argparse.ArgumentParser(description='叶片净空计算系统')
    parser.add_argument('--video_path', type=str, required=True, help='输入视频路径')
    parser.add_argument('--model_path', type=str, required=True, help='模型文件路径')
    parser.add_argument('--output_dir', type=str, required=True, help='结果保存目录')
    parser.add_argument('--scale_factor', type=float, default=0.02, help='比例尺因子 (实际距离/像素距离)')
    parser.add_argument('--width', type=int, default=544, help='图像宽度（像素）')
    parser.add_argument('--height', type=int, default=960, help='图像高度（像素）')
    args = parser.parse_args()

    calculator = ClearanceCalculator(
        model_path=args.model_path,
        scale_factor=args.scale_factor,
        width=args.width,
        height=args.height
    )

    if not calculator.load_video(args.video_path):
        print(f"无法加载视频: {args.video_path}")
        return

    print(f"开始处理视频: {args.video_path}")
    print(f"图像尺寸设置为: ({args.width}, {args.height})")

    result = calculator.process_video(args.output_dir)


if __name__ == "__main__":
    main()
    """
    python cli_clearance.py \
  --video test.mp4 \
  --model best.pt \
  --output ./results \
  --scale_factor 0.02 \
  --width 640 \
  --height 480
    """
    """
    {
  "clearances": {
    "0": [5.12, 5.09, 5.11],
    "1": [4.98, 5.01],
    "2": []
  },
  "last_frame_clearances": {
    "tip1": 5.11,
    "tip2": 5.01,
    "tip3": null
  }
}
    """