#!/usr/bin/env python3
"""
EZAI-Go2 优化版本测试 - 稳定性和清晰度改进
"""

import os
import sys
import time
import ezai.go2 as go2

os.environ["DISPLAY"] = ":0"

def optimized_video_demo():
    """优化的视频流演示"""
    print("=== EZAI-Go2 优化视频演示 ===")
    
    # 1. 创建Go2实例
    robot = go2.Go2()
    robot.init()
    print(f"使用接口: {robot.interface}")
    
    # 2. 获取优化的视频读取器
    read_frame, show_frame, cleanup = robot.get_video_reader(method='pipe')
    
    if read_frame is None:
        print("视频流初始化失败")
        return
    
    print("优化视频流启动成功")
    print("操作说明:")
    print("  q - 退出")
    print("  s - 保存当前帧")
    print("  f - 切换全屏")
    print("  r - 重置统计")
    
    try:
        frame_count = 0
        valid_count = 0
        start_time = time.time()
        last_stats_time = start_time
        
        while True:
            # 读取一帧
            frame = read_frame()
            
            # 显示帧
            key = show_frame(frame)
            
            # 统计
            if frame is not None:
                valid_count += 1
            
            frame_count += 1
            
            # 每5秒显示一次统计
            current_time = time.time()
            if current_time - last_stats_time >= 5.0:
                elapsed = current_time - start_time
                fps = frame_count / elapsed if elapsed > 0 else 0
                valid_rate = (valid_count / frame_count * 100) if frame_count > 0 else 0
                print(f"统计: 总帧数={frame_count}, 有效帧={valid_count}, "
                      f"有效帧率={valid_rate:.1f}%, FPS={fps:.1f}")
                last_stats_time = current_time
            
            # 处理按键
            if key == ord('q'):
                break
            elif key == ord('s') and frame is not None:
                import cv2
                filename = f"/tmp/go2_optimized_{int(time.time())}.jpg"
                cv2.imwrite(filename, frame)
                print(f"保存帧到: {filename}")
            elif key == ord('f'):
                # 全屏切换
                cv2.setWindowProperty('Go2视频流', cv2.WND_PROP_FULLSCREEN, 
                                    cv2.WINDOW_FULLSCREEN)
                print("切换到全屏")
            elif key == ord('r'):
                # 重置统计
                frame_count = 0
                valid_count = 0
                start_time = time.time()
                last_stats_time = start_time
                print("统计已重置")
                
    except KeyboardInterrupt:
        print("\n用户中断")
    finally:
        # 清理资源
        cleanup()
        
        # 最终统计
        if frame_count > 0:
            elapsed = time.time() - start_time
            fps = frame_count / elapsed if elapsed > 0 else 0
            valid_rate = (valid_count / frame_count * 100)
            print(f"\n最终统计:")
            print(f"  总帧数: {frame_count}")
            print(f"  有效帧数: {valid_count}")
            print(f"  有效帧率: {valid_rate:.1f}%")
            print(f"  平均FPS: {fps:.1f}")
        
        print("程序结束")

def quality_comparison_test():
    """质量对比测试 - 不同参数的效果"""
    print("=== 质量对比测试 ===")
    
    robot = go2.Go2()
    robot.init()
    
    # 测试不同的读取方法
    methods = ['auto', 'pipe']
    
    for method in methods:
        print(f"\n--- 测试方法: {method} ---")
        
        read_frame, show_frame, cleanup = robot.get_video_reader(method=method)
        
        if read_frame is None:
            print(f"方法 {method} 初始化失败")
            continue
        
        try:
            frame_count = 0
            valid_count = 0
            
            # 测试30秒
            start_time = time.time()
            while time.time() - start_time < 30:
                frame = read_frame()
                key = show_frame(frame)
                
                if frame is not None:
                    valid_count += 1
                
                frame_count += 1
                
                if key == ord('q'):
                    break
            
            # 统计结果
            elapsed = time.time() - start_time
            fps = frame_count / elapsed if elapsed > 0 else 0
            valid_rate = (valid_count / frame_count * 100) if frame_count > 0 else 0
            
            print(f"方法 {method} 结果:")
            print(f"  总帧数: {frame_count}")
            print(f"  有效帧数: {valid_count}")
            print(f"  有效帧率: {valid_rate:.1f}%")
            print(f"  平均FPS: {fps:.1f}")
            
        except Exception as e:
            print(f"方法 {method} 测试出错: {e}")
        finally:
            cleanup()
        
        # 方法间间隔
        if method != methods[-1]:
            print("等待3秒后测试下一个方法...")
            time.sleep(3)

if __name__ == "__main__":
    print("选择测试模式:")
    print("1. 优化视频演示")
    print("2. 质量对比测试")
    
    try:
        choice = input("请选择 (1/2): ").strip()
        
        if choice == "1":
            optimized_video_demo()
        elif choice == "2":
            quality_comparison_test()
        else:
            print("默认运行优化视频演示...")
            optimized_video_demo()
            
    except KeyboardInterrupt:
        print("\n程序被用户中断")