#!/usr/bin/env python3
"""
WS2812B SPI测试脚本 - 基于rpi_ws281x库
使用SPI接口和GPIO10控制WS2812B LED灯带
rpi_ws281x库在SPI模式下只能使用GPIO10引脚
"""

import time
import subprocess
import sys

import RPi.GPIO as GPIO
from rpi_ws281x import Color, PixelStrip

# LED参数配置 - SPI模式
LED_COUNT = 10          # LED数量
LED_PIN = 10           # GPIO引脚（SPI模式只能使用GPIO10）
LED_FREQ_HZ = 800000   # LED信号频率
LED_DMA = 10           # DMA通道
LED_BRIGHTNESS = 100   # 亮度（0-255）- 恢复适中亮度
LED_INVERT = False     # 是否反转信号
LED_CHANNEL = 0        # PWM通道（SPI模式）

# SPI时序参数
SPI_RESET_DELAY = 0.1   # 重置延时（秒）
SPI_SHOW_DELAY = 0.001  # show()后延时（秒）

def color_wipe(strip, color, wait_ms=50):
    """颜色擦除效果 - 逐个点亮LED"""
    for i in range(strip.numPixels()):
        strip.setPixelColor(i, color)
        strip.show()
        time.sleep(wait_ms / 1000.0)
        time.sleep(SPI_SHOW_DELAY)  # 添加show()后延时
    
    # 颜色显示完成后添加重置延时
    time.sleep(SPI_RESET_DELAY)

def theater_chase(strip, color, wait_ms=50, iterations=10):
    """剧院追逐效果"""
    for j in range(iterations):
        for q in range(3):
            for i in range(0, strip.numPixels(), 3):
                strip.setPixelColor(i + q, color)
            strip.show()
            time.sleep(wait_ms / 1000.0)
            for i in range(0, strip.numPixels(), 3):
                strip.setPixelColor(i + q, 0)

def rainbow(strip, wait_ms=20, iterations=1):
    """彩虹效果"""
    for j in range(256 * iterations):
        for i in range(strip.numPixels()):
            strip.setPixelColor(i, wheel((i + j) & 255))
        strip.show()
        time.sleep(SPI_SHOW_DELAY)  # 添加show后延时
        time.sleep(wait_ms / 1000.0)

def rainbow_cycle(strip, wait_ms=20, iterations=5):
    """彩虹循环效果"""
    for j in range(256 * iterations):
        for i in range(strip.numPixels()):
            strip.setPixelColor(i, wheel((int(i * 256 / strip.numPixels()) + j) & 255))
        strip.show()
        time.sleep(SPI_SHOW_DELAY)  # 添加show后延时
        time.sleep(wait_ms / 1000.0)

def wheel(pos):
    """生成彩虹颜色"""
    if pos < 85:
        return Color(pos * 3, 255 - pos * 3, 0)
    elif pos < 170:
        pos -= 85
        return Color(255 - pos * 3, 0, pos * 3)
    else:
        pos -= 170
        return Color(0, pos * 3, 255 - pos * 3)

def reset_strip(strip):
    """重置LED灯带，解决颜色锁死问题"""
    # 清空所有LED
    for i in range(strip.numPixels()):
        strip.setPixelColor(i, Color(0, 0, 0))
    strip.show()
    
    # 更长的重置延时确保WS2812B完全重置
    time.sleep(SPI_RESET_DELAY)
    
    # 发送额外的重置序列
    for i in range(5):
        strip.show()
        time.sleep(SPI_SHOW_DELAY)
    
    # 最终重置延时
    time.sleep(SPI_RESET_DELAY)

def hard_reset_spi():
    """硬件级别的SPI重置"""
    print("执行硬件级SPI重置...")
    
    # 尝试通过系统命令重置SPI
    try:
        # 重启SPI内核模块
        subprocess.run(["sudo", "rmmod", "spi_bcm2835"], check=False, capture_output=True)
        subprocess.run(["sudo", "modprobe", "spi_bcm2835"], check=False, capture_output=True)
        print("SPI内核模块重置完成")
    except:
        print("SPI内核模块重置失败")
    
    # 等待SPI重新初始化
    time.sleep(1)

def complete_cleanup(strip):
    """完全清理SPI资源"""
    print("执行完全清理...")
    
    try:
        # 清空LED
        for i in range(strip.numPixels()):
            strip.setPixelColor(i, Color(0, 0, 0))
        strip.show()
        
        # 长时间重置
        time.sleep(SPI_RESET_DELAY * 2)
        
        # 尝试释放SPI资源
        # 注意：rpi_ws281x库没有直接的close方法，所以我们通过其他方式清理
        
    except Exception as e:
        print(f"清理过程中出现错误: {e}")
    
    # 硬件重置
    hard_reset_spi()
    
    print("完全清理完成")

def test_single_color_with_reset(strip, color, color_name):
    """测试单个颜色并立即重置"""
    print(f"测试{color_name}...")
    
    try:
        # 设置颜色
        for i in range(strip.numPixels()):
            strip.setPixelColor(i, color)
        strip.show()
        time.sleep(SPI_SHOW_DELAY)
        
        # 短暂显示
        time.sleep(0.5)
        
        # 立即清空
        for i in range(strip.numPixels()):
            strip.setPixelColor(i, Color(0, 0, 0))
        strip.show()
        time.sleep(SPI_RESET_DELAY)
        
        print(f"{color_name}测试完成")
        
    except Exception as e:
        print(f"{color_name}测试失败: {e}")
        return False
    
    return True

def safe_color_test(strip):
    """安全颜色测试，避免WS2812B锁死"""
    print("安全颜色测试...")
    
    # 使用新的测试方法
    success = True
    success &= test_single_color_with_reset(strip, Color(255, 0, 0), "红色")
    time.sleep(0.5)
    
    success &= test_single_color_with_reset(strip, Color(0, 255, 0), "绿色")
    time.sleep(0.5)
    
    success &= test_single_color_with_reset(strip, Color(0, 0, 255), "蓝色")
    time.sleep(0.5)
    
    success &= test_single_color_with_reset(strip, Color(255, 255, 255), "白色")
    
    if success:
        print("所有颜色测试完成")
    else:
        print("颜色测试过程中出现问题")
        complete_cleanup(strip)

def set_brightness(strip, brightness):
    """设置所有LED的亮度"""
    for i in range(strip.numPixels()):
        # 获取当前颜色并应用新的亮度
        color = strip.getPixelColor(i)
        r = (color >> 16) & 0xFF
        g = (color >> 8) & 0xFF
        b = color & 0xFF
        
        # 应用亮度
        r = int(r * brightness / 255)
        g = int(g * brightness / 255)
        b = int(b * brightness / 255)
        
        strip.setPixelColor(i, Color(r, g, b))
    strip.show()

def test_basic_colors(strip):
    """测试基本颜色"""
    print("测试基本颜色...")
    safe_color_test(strip)

def test_brightness_levels(strip):
    """测试不同亮度级别"""
    print("测试亮度级别...")
    
    # 高亮度
    print("高亮度")
    set_brightness(strip, 255)
    time.sleep(2)
    
    # 中等亮度
    print("中等亮度")
    set_brightness(strip, 128)
    time.sleep(2)
    
    # 低亮度
    print("低亮度")
    set_brightness(strip, 64)
    time.sleep(2)
    
    # 恢复高亮度
    set_brightness(strip, 255)

def test_patterns(strip):
    """测试各种模式"""
    print("测试彩虹效果...")
    rainbow(strip)
    time.sleep(1)
    
    print("测试彩虹循环...")
    rainbow_cycle(strip)
    time.sleep(1)
    
    print("测试剧院追逐...")
    theater_chase(strip, Color(255, 0, 0))
    theater_chase(strip, Color(0, 255, 0))
    theater_chase(strip, Color(0, 0, 255))
    time.sleep(1)

def test_spi_timing(strip):
    """测试SPI时序"""
    print("测试SPI时序...")
    
    # 快速刷新测试
    print("快速刷新测试")
    for i in range(10):
        color = Color((i * 25) % 256, (i * 50) % 256, (i * 75) % 256)
        color_wipe(strip, color, 10)
    
    # 精确时序测试
    print("精确时序测试")
    start_time = time.time()
    for i in range(100):
        strip.setPixelColor(0, Color(255, 0, 0))
        strip.show()
        time.sleep(SPI_SHOW_DELAY)
        strip.setPixelColor(0, Color(0, 255, 0))
        strip.show()
        time.sleep(SPI_SHOW_DELAY)
        strip.setPixelColor(0, Color(0, 0, 255))
        strip.show()
        time.sleep(SPI_SHOW_DELAY)
    end_time = time.time()
    
    print(f"300次刷新耗时: {end_time - start_time:.3f}秒")
    print(f"平均每次刷新: {(end_time - start_time) / 300 * 1000:.2f}毫秒")
    
    # SPI模式验证
    print("SPI模式验证")
    print("当前使用GPIO10 (MOSI)引脚进行SPI通信")
    print("数据通过SPI0.0通道传输到WS2812B")
    print(f"重置延时: {SPI_RESET_DELAY}秒")
    print(f"Show延时: {SPI_SHOW_DELAY}秒")

def diagnose_spi_issues(strip):
    """诊断SPI问题"""
    print("=== SPI问题诊断 ===")
    
    print("1. 测试单个LED颜色切换...")
    for color_name, color in [("红色", Color(255, 0, 0)), 
                              ("绿色", Color(0, 255, 0)), 
                              ("蓝色", Color(0, 0, 255))]:
        print(f"  {color_name}...")
        strip.setPixelColor(0, color)
        strip.show()
        time.sleep(SPI_SHOW_DELAY)
        time.sleep(0.5)
        reset_strip(strip)
        time.sleep(0.5)
    
    print("2. 测试不同延时设置...")
    for delay in [0.001, 0.005, 0.01, 0.05]:
        print(f"  延时 {delay} 秒...")
        strip.setPixelColor(0, Color(255, 0, 0))
        strip.show()
        time.sleep(delay)
        strip.setPixelColor(0, Color(0, 255, 0))
        strip.show()
        time.sleep(delay)
        strip.setPixelColor(0, Color(0, 0, 255))
        strip.show()
        time.sleep(delay)
        reset_strip(strip)
        time.sleep(0.5)
    
    print("诊断完成")

def initialize_strip():
    """初始化LED灯带"""
    print("初始化WS2812B LED灯带...")
    print("使用SPI模式和GPIO10引脚...")
    
    # 确保SPI处于干净状态
    hard_reset_spi()
    
    # 创建PixelStrip对象
    strip = PixelStrip(LED_COUNT, LED_PIN, LED_FREQ_HZ, LED_DMA, LED_INVERT, LED_BRIGHTNESS, LED_CHANNEL)
    strip.begin()
    
    # 初始化后清空
    time.sleep(0.1)
    for i in range(strip.numPixels()):
        strip.setPixelColor(i, Color(0, 0, 0))
    strip.show()
    time.sleep(SPI_RESET_DELAY)
    
    return strip

def main():
    """主测试函数"""
    # 初始化LED灯带
    strip = initialize_strip()
    
    print("WS2812B SPI测试开始...")
    print(f"LED数量: {LED_COUNT}")
    print(f"使用引脚: GPIO{LED_PIN} (SPI MOSI)")
    print(f"SPI通道: SPI0.0")
    print(f"亮度级别: {LED_BRIGHTNESS}")
    print("注意: rpi_ws281x库在SPI模式下必须使用GPIO10引脚")
    print("重要: 每次测试后都会执行完全清理以避免锁死")
    
    try:
        while True:
            print("\n=== WS2812B SPI测试菜单 ===")
            print("1. 基本颜色测试（修复版）")
            print("2. 亮度级别测试")
            print("3. 特效模式测试")
            print("4. SPI时序测试")
            print("5. SPI问题诊断")
            print("6. 全部测试")
            print("7. 强制重置LED")
            print("8. 退出")
            print("当前使用GPIO10 (SPI模式)")
            
            choice = input("请选择测试模式 (1-8): ")
            
            if choice == '1':
                test_basic_colors(strip)
            elif choice == '2':
                test_brightness_levels(strip)
            elif choice == '3':
                test_patterns(strip)
            elif choice == '4':
                test_spi_timing(strip)
            elif choice == '5':
                diagnose_spi_issues(strip)
            elif choice == '6':
                test_basic_colors(strip)
                test_brightness_levels(strip)
                test_patterns(strip)
                test_spi_timing(strip)
            elif choice == '7':
                print("强制重置LED...")
                complete_cleanup(strip)
                print("完全重置完成，需要重新初始化...")
                # 重新初始化
                strip = initialize_strip()
            elif choice == '8':
                print("执行完全清理并退出...")
                complete_cleanup(strip)
                break
            else:
                print("无效选择，请重试")
                
    except KeyboardInterrupt:
        print("\n收到中断信号...")
    except Exception as e:
        print(f"测试过程中发生错误: {e}")
    finally:
        # 执行完全清理
        print("执行完全清理...")
        complete_cleanup(strip)
        print("测试结束")

if __name__ == '__main__':
    main()