#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import os
import glob
import datetime

from utils.detect_dual_track_lines import auto_detect_dual_track_lines

def visualize_detection(img_path, output_dir='results'):
    """可视化轨迹线检测结果并保存"""
    # 读取图像
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图像: {img_path}")
        return False
    
    # 获取文件名
    filename = os.path.basename(img_path)
    
    # 检测轨迹线
    result = auto_detect_dual_track_lines(img, observe=False, save_log=False)
    if result[0] is None:
        print(f"检测失败: {img_path}")
        return False
    
    center_info, left_info, right_info = result
    
    # 绘制结果
    result_img = img.copy()
    height, width = img.shape[:2]
    center_x = width // 2
    
    # 获取线段坐标
    left_x1, left_y1, left_x2, left_y2 = left_info["line"]
    right_x1, right_y1, right_x2, right_y2 = right_info["line"]
    center_point = center_info["point"]
    deviation = center_info["deviation"]
    
    # 绘制左右轨迹线
    cv2.line(result_img, (int(left_x1), int(left_y1)), (int(left_x2), int(left_y2)), (255, 0, 0), 2)
    cv2.line(result_img, (int(right_x1), int(right_y1)), (int(right_x2), int(right_y2)), (0, 0, 255), 2)
    
    # 计算中心线
    center_line_x1 = (left_x1 + right_x1) // 2
    center_line_y1 = (left_y1 + right_y1) // 2
    center_line_x2 = (left_x2 + right_x2) // 2
    center_line_y2 = (left_y2 + right_y2) // 2
    
    # 绘制中心线
    cv2.line(result_img, (center_line_x1, center_line_y1), (center_line_x2, center_line_y2), (0, 255, 0), 2)
    
    # 绘制图像中心线
    cv2.line(result_img, (center_x, 0), (center_x, height), (0, 0, 255), 1)
    
    # 标记中心点
    cv2.circle(result_img, center_point, 10, (255, 0, 255), -1)
    
    # 显示偏差信息
    cv2.putText(result_img, f"Deviation: {deviation}px", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    
    # 显示是否为石板路模式
    if center_info["stone_path_mode"]:
        cv2.putText(result_img, "Stone Path Mode", (10, 60), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 保存结果图像
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    output_path = os.path.join(output_dir, f"result_{filename}")
    cv2.imwrite(output_path, result_img)
    print(f"结果已保存到: {output_path}")
    return True

def process_all_images(input_dir, output_dir='results'):
    """处理目录中的所有图像"""
    # 获取所有jpg图像
    image_paths = glob.glob(os.path.join(input_dir, "*.jpg"))
    
    total = len(image_paths)
    success = 0
    
    for img_path in sorted(image_paths):
        print(f"\n处理图像: {img_path}")
        if visualize_detection(img_path, output_dir):
            success += 1
    
    success_rate = (success / total) * 100 if total > 0 else 0
    print(f"\n处理完成: 成功率 {success_rate:.2f}% ({success}/{total})")

if __name__ == "__main__":
    input_dir = "res/path/dual"
    output_dir = "results/dual_track"
    process_all_images(input_dir, output_dir)