# 要求保证视频不能移动，且全过程没有任何遮挡物
import os
import sys
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from utils import draw_grid

import cv2
import numpy as np
import tkinter as tk
from tkinter import simpledialog
from dotenv import load_dotenv
load_dotenv()  # 加载环境变量

# 所要切分的图块宽高
patch_w = int(os.getenv('PATCH_WIDTH', 256))
patch_h = int(os.getenv('PATCH_HEIGHT', 256))
scale = 2
# 存储标记点
clicked_points = []

def get_text_input(prompt):
    """创建弹窗获取文本输入"""
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    root.attributes('-topmost', True)  # 确保弹窗置顶
    result = simpledialog.askstring(' ', prompt)
    root.destroy()
    if result is None:
        result = ""
    if not result.strip().isdigit():
        result = -1
    return int(result)

def mouse_callback(event, x, y, flags, param):
    """
    鼠标回调函数
    """
    global clicked_points
    global patch_w
    global patch_h
    global scale

    if event == cv2.EVENT_LBUTTONDOWN:  # 左键点击
        # 在点击位置绘制红色圆点
        scale_patch_w = patch_w // scale
        scale_patch_h = patch_h // scale
        # 格子角点
        circle_x_corner = (x // scale_patch_w)*scale_patch_w
        circle_y_corner = (y // scale_patch_h)*scale_patch_h
        # 格子中心点
        circle_x_center = circle_x_corner + scale_patch_w//2
        circle_y_center = circle_y_corner + scale_patch_h//2
        cv2.circle(param, (circle_x_center, circle_y_center), 5, (0, 0, 255), -1)
        cv2.circle(param, (circle_x_corner, circle_y_corner), 5, (255, 0, 0), -1)

        # 更新显示
        cv2.imshow('img', param)

        cls = get_text_input('请输入类别:0.背景 1.浑浊 -1.不参与')
        # 显示标签文本
        cv2.putText(param, str(cls), (circle_x_center + 10, circle_y_center + 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        # 更新显示
        cv2.imshow('img', param)
        valid_cls = [0, 1, -1]
        if cls in valid_cls:
            print(f"点击网格角点: ({circle_x_corner}, {circle_y_corner}) 中心点: ({circle_x_center}, {circle_y_center}) 类别:{cls}")
            # 记录标注数据，角点 u v w h cls
            clicked_points.append([circle_x_corner, circle_y_corner, scale_patch_w, scale_patch_h, cls])
        else:
            print("请输入正确的类别！")
    elif event == cv2.EVENT_RBUTTONDOWN:  # 右键点击
        removed_point = clicked_points.pop()
        print(f"撤销标注点: ({removed_point[0]}, {removed_point[1]}) 类别: {removed_point[4]}")
        # 将撤销点标记为黑色
        x = removed_point[0]
        y = removed_point[1]
        # 在点击位置绘制黑色圆点
        scale_patch_w = patch_w // scale
        scale_patch_h = patch_h // scale
        # 格子角点
        circle_x_corner = (x // scale_patch_w)*scale_patch_w
        circle_y_corner = (y // scale_patch_h)*scale_patch_h
        # 格子中心点
        circle_x_center = circle_x_corner + scale_patch_w//2
        circle_y_center = circle_y_corner + scale_patch_h//2
        cv2.circle(param, (circle_x_center, circle_y_center), 5, (128, 128, 128), -1)
        cv2.circle(param, (circle_x_corner, circle_y_corner), 5, (128, 128, 128), -1)
        # 更新显示
        cv2.imshow('img', param)


def remove_duplicates(arr:list):
    """列表去重"""
    unique_list = []
    [unique_list.append(item) for item in arr if item not in unique_list]
    return unique_list

def main():
    """
    固定摄像头标注,只需要标注一张图像，后续图像保持一致
    1.标注过程，先将图像划分为图框，用cv2划线工具在图像上划网格线
    2.用鼠标进行交互，点击图块输入标签，按下空格键完成交互过程，保存标签
    3.标签格式：u,v,w,h,label u,v为块左上角坐标，w,h为块的宽和高，label为块的标签
    """
    global clicked_points
    global patch_w
    global patch_h
    global scale
    # TODO: 需要更改为准备标注的图像路径，使用当前目录下的000000.jpg，结果保存在当前目录下label.txt
    img_path = r'D:\code\water_turbidity_det\data\video1_20251129120104_20251129123102\000000.jpg'
    img = cv2.imread(img_path)
    # resize 图像太大了显示不全
    img = cv2.resize(img, (img.shape[1] // scale, img.shape[0] // scale))
    # 绘制网格线
    draw_grid(img, patch_w // scale, patch_h // scale)
    # 交互标注
    print("操作说明:")
    print("- 点击鼠标左键在图像上添加红色标记点: 0.其他 1.浑浊 -1.忽略，不参与训练和测试")
    print("- 按 'c' 键清除所有标记点")
    print("- 按 ESC 键退出程序")
    cv2.namedWindow('img')
    cv2.setMouseCallback('img', mouse_callback, img)
    # 交互标注
    while True:
        cv2.imshow('img', img)
        key = cv2.waitKey(1) & 0xFF

        # 按 'c' 键清除所有标记点
        if key == ord('c'):
            img = cv2.imread(img_path)
            img = cv2.resize(img, (img.shape[1] // scale, img.shape[0] // scale))
            draw_grid(img, patch_w // scale, patch_h // scale)
            clicked_points.clear()
            cv2.setMouseCallback('img', mouse_callback, img)

        # 按 ESC 键退出
        elif key == 27:  # ESC键
            break
    cv2.destroyAllWindows()
    # 输出所有点击位置
    # 列表去重
    clicked_points = remove_duplicates(clicked_points)

    print(f"总共标记了 {len(clicked_points)} 个点:")
    for i, point in enumerate(clicked_points):
        print(f"  点 {i + 1}: ({point[0]}, {point[1]}, {point[2]}, {point[3]}, {point[4]})")
    # 恢复尺寸
    clicked_points = [[p[0]*scale, p[1]*scale, p[2]*scale, p[3]*scale, p[4]] for p in clicked_points]
    # 写入txt
    if clicked_points:
        with open(os.path.join(os.path.dirname(img_path), 'label.txt'), 'w') as fw:
            for point in clicked_points:
                fw.write(f"{point[0]},{point[1]},{point[2]},{point[3]},{point[4]}\n")
        # 保存点
        print(f"保存标记点 {len(clicked_points)} 个:")
        for i, point in enumerate(clicked_points):
            print(f"  点 {i + 1}: ({point[0]}, {point[1]}, {point[2]}, {point[3]}, {point[4]})")
    else :
        print("没有标记点！不保存任何文件！")



if __name__ == '__main__':
    main()