import cv2
import numpy as np
from src.distance_estimator import DistanceEstimator
from src.person_detector import PersonDetector

def test_distance_estimator():
    # 初始化距离估算器（使用实际相机参数）
    estimator = DistanceEstimator(
        focal_length_mm=35,
        sensor_width_mm=23.5,
        sensor_height_mm=15.6,
        image_width_pixels=1920,
        image_height_pixels=1080,
        camera_height_mm=1700,
        camera_tilt_angle=15
    )
    
    print("测试距离估算...")
    
    # 读取测试图像
    test_image = cv2.imread("tests/test_data/test_image.jpg")
    assert test_image is not None, "无法加载测试图像"

    # 初始化检测器
    detector = PersonDetector("yolov8n.pt", conf_threshold=0.3)
    persons, vis_path = detector.detect(test_image, save_visualization=True)

    print("len(person)", len(persons))

    for person in persons:
        assert 'bbox' in person, "检测结果缺少bbox"
        assert 'confidence' in person, "检测结果缺少confidence"
        assert len(person['bbox']) == 4, "bbox格式不正确"
        assert 0 <= person['confidence'] <= 1, "confidence值范围不正确"


    
        # 测试边界框方法
        bbox = person['bbox']  # 示例边界框
        bbox_distance = estimator.estimate_distance_from_bbox(bbox[3] - bbox[1])
        print(f"边界框方法估算距离: {bbox_distance:.2f}mm")
        assert bbox_distance > 0, "距离估算结果应该大于0"
        
        # 测试关键点方法
        keypoint_distance, keypoint_conf = estimator.estimate_distance_from_keypoints(test_image)
        if keypoint_distance is not None:
            print(f"关键点方法估算距离: {keypoint_distance:.2f}mm, 置信度: {keypoint_conf:.2f}")
            assert keypoint_distance > 0, "距离估算结果应该大于0"
            assert 0 <= keypoint_conf <= 1, "置信度应该在0-1之间"
        
        # 测试地平面投影方法
        ground_distance = estimator.estimate_distance_from_ground_plane(bbox[3])
        print(f"地平面投影方法估算距离: {ground_distance:.2f}mm")
        assert ground_distance > 0, "距离估算结果应该大于0"
        
        # 测试综合估计方法
        final_distance, final_confidence = estimator.estimate_distance(test_image, bbox)
        print(f"综合估计距离: {final_distance:.2f}mm, 置信度: {final_confidence:.2f}")
        assert final_distance > 0, "距离估算结果应该大于0"
        assert 0 <= final_confidence <= 1, "置信度应该在0-1之间"
        
        # 测试可视化
        vis_image = estimator.visualize_estimation(test_image, bbox, final_distance, final_confidence)
        assert vis_image is not None, "可视化结果不应为空"
        assert vis_image.shape == test_image.shape, "可视化图像尺寸应与原图相同"
        
        # 保存可视化结果
    cv2.imwrite("tests/test_data/distance_estimation_result.jpg", vis_image)
    
    print("距离估算测试完成!")

if __name__ == "__main__":
    test_distance_estimator()