from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import os
import time
import argparse

import cv2
import numpy as np
from glob import glob

import sys
sys.path.append(os.getcwd())
from core.config import cfg
from models.rknnlite_rk3588_tracker import NnoTracker_RKNNLite

# 声明全局变量（在模块级别定义，确保所有函数可访问）
drawing = False  # 是否正在绘制
ix, iy = -1, -1  # 起始点
roi = None       # 存储ROI区域
tracking = False # 是否正在跟踪
paused = False   # 视频是否暂停
frame = None     # 当前帧
tracker = None
# 鼠标回调函数
def draw_rectangle(event, x, y, flags, param):
    # 声明使用全局变量
    global ix, iy, drawing, roi, tracking, paused, frame
    
    # 左键按下开始绘制，同时暂停视频
    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y
        roi = None
        paused = True  # 开始绘制时暂停视频
    
    # 鼠标移动绘制矩形
    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing and paused and frame is not None:  # 确保frame已初始化
            # 复制当前帧用于显示临时矩形
            temp_frame = frame.copy()
            cv2.rectangle(temp_frame, (ix, iy), (x, y), (0, 255, 0), 2)
           
            cv2.imshow('Tracking', temp_frame)
    
    # 左键释放完成绘制，继续播放
    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        # 确保ROI有效（大于10x10像素）
        if abs(x - ix) > 10 and abs(y - iy) > 10 and frame is not None:
            roi = [min(ix, x), min(iy, y), abs(x - ix), abs(y - iy)]
            # 初始化跟踪器（假设tracker是全局可访问的）
            global tracker
             # 加载权重并初始化跟踪器
            Tback_weight = './weights/track_backbone_T.rknn'
            Xback_weight = './weights/track_backbone_X.rknn'
            Head_weight = './weights/head.rknn'
            tracker = NnoTracker_RKNNLite(Tback_weight, Xback_weight, Head_weight)
            tracker.init(frame, roi)
            tracking = True
            print(f"初始化跟踪，ROI: {roi}")
        paused = False  # 完成绘制后继续播放
    
    # 右键点击取消跟踪
    elif event == cv2.EVENT_RBUTTONDOWN:
        tracking = False
        roi = None
        
        print("已取消跟踪")


def main():
    # 声明使用全局变量
    global drawing, ix, iy, roi, tracking, paused, frame, tracker
    
    # 解析参数
    parser = argparse.ArgumentParser(description='tracking demo')
    parser.add_argument('--config', default='./models/config/config.yaml', type=str, help='config file')
    parser.add_argument('--save', action='store_true', help='whether visualzie result')
    args = parser.parse_args()

    # 加载配置
    cfg.merge_from_file(args.config)

    
    
    # 初始化视频捕获
    video_path = '/home/orangepi/Desktop/track1.mp4'
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return

    # 创建窗口并设置鼠标回调
    cv2.namedWindow('Tracking')
    cv2.setMouseCallback('Tracking', draw_rectangle)

    # 读取第一帧并初始化
    ret, img_src = cap.read()
    if ret:
        frame = cv2.resize(img_src, (1280, 720))
    else:
        print("无法读取视频帧")
        return

    # 主循环
    while cap.isOpened():
        # 只有在未暂停且能读取帧时，才获取新帧
        if not paused and ret:
            ret, img_src = cap.read()
            if ret:
                frame = cv2.resize(img_src, (1280, 720))
        
        # 视频结束判断
        if not ret:
            cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            ret, img_src = cap.read()
            #break
        
        # 复制帧用于显示（避免修改原始帧）
        display_frame = frame.copy()

       

        # 跟踪逻辑
        if tracking:
            t1 = time.time()
            outputs = tracker.track(frame)
            fps = 1. / (time.time() - t1)
            
            # 显示FPS
            cv2.putText(display_frame, f"FPS: {fps:.2f}", (10, 150), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            
            # 绘制跟踪框
            if 'bbox' in outputs:
                bbox = list(map(int, outputs['bbox']))
                cv2.rectangle(display_frame, (bbox[0], bbox[1]),
                             (bbox[0] + bbox[2], bbox[1] + bbox[3]),
                             (0, 255, 0), 3)
            else:
                cv2.putText(display_frame, "Tracking Failed!", (10, 90), 
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
                tracking = False
        else:
            # 未跟踪时显示提示
            cv2.putText(display_frame, "No Tracking", (10, 50), 
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
       

        # 显示帧（始终执行，确保窗口刷新）
        cv2.imshow('Tracking', display_frame)
        
        # 按键处理（始终执行，确保能响应操作）
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord(' '):  # 空格键控制暂停/继续
            paused = not paused
            print("暂停" if paused else "继续")

    # 释放资源
    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()