#!/bin/bash

# ============================================================================
# YOLO数字识别系统 - 快速运行脚本
# ============================================================================
# 
# 功能: 提供交互式菜单，简化常用操作
# 作者: Gavin Chan
# 日期: 2025-10-30
#
# 使用方法:
#   ./run.sh          # 交互式菜单
#
# 注意事项:
#   - 需要先激活虚拟环境（脚本会自动尝试）
#   - 确保最佳模型文件存在
#   - 识别结果保存在 results/ 目录
# ============================================================================

set -e  # 遇到错误立即退出

# ----------------------------------------------------------------------------
# 颜色定义（用于美化输出）
# ----------------------------------------------------------------------------
GREEN='\033[0;32m'   # 绿色：成功信息
YELLOW='\033[1;33m'  # 黄色：警告或进行中
RED='\033[0;31m'     # 红色：错误信息
NC='\033[0m'         # 无颜色：重置

# ----------------------------------------------------------------------------
# 标题
# ----------------------------------------------------------------------------
echo "=================================================="
echo "🚀 YOLO数字识别系统 - 快速运行脚本"
echo "=================================================="

# ----------------------------------------------------------------------------
# 1. 激活虚拟环境
# 尝试多个可能的虚拟环境位置，找到第一个存在的就使用
# ----------------------------------------------------------------------------
echo -e "\n${GREEN}1. 激活虚拟环境${NC}"

if [ -d ~/venv ]; then
    # 方式1: 用户主目录下的 venv
    source ~/venv/bin/activate
    echo "✓ 虚拟环境已激活: ~/venv"
elif [ -d venv ]; then
    # 方式2: 当前目录下的 venv
    source venv/bin/activate
    echo "✓ 虚拟环境已激活: ./venv"
else
    # 未找到虚拟环境，使用系统Python（可能缺少依赖）
    echo "⚠️  未找到虚拟环境，使用系统Python"
fi

# 显示当前工作目录
echo "✓ 工作目录: $(pwd)"

# ----------------------------------------------------------------------------
# 2. 检查依赖包
# 验证必需的Python包是否已安装，如果缺少则自动安装
# ----------------------------------------------------------------------------
echo -e "\n${GREEN}2. 检查依赖包${NC}"

if python -c "import ultralytics, cv2, numpy" 2>/dev/null; then
    echo "✓ 所有依赖包已安装"
else
    echo -e "${RED}✗ 缺少依赖包，正在安装...${NC}"
    pip install ultralytics opencv-python numpy
fi

# ----------------------------------------------------------------------------
# 3. 显示操作菜单
# 提供常用功能的快捷入口
# ----------------------------------------------------------------------------
echo -e "\n${GREEN}3. 请选择操作${NC}"
echo "=================================="
echo "1) 识别数字 (使用最佳模型) ⭐推荐"
echo "2) 低阈值识别 (conf=0.15，更多检测)"
echo "3) 高清识别 (imgsz=640，更精确)"
echo "4) 查看已有结果和模型信息"
echo "0) 退出"
echo "=================================="

read -p "请输入选项 [0-4]: " choice

# ----------------------------------------------------------------------------
# 配置：最佳模型路径
# 这是训练好的最佳模型，使用CLAHE预处理 + 150轮训练
# mAP50 = 0.995，在训练集上接近完美
# ----------------------------------------------------------------------------
BEST_MODEL="runs/digit_yolo/exp_preprocessed_color_150/weights/best.pt"

# ----------------------------------------------------------------------------
# 4. 执行选择的操作
# ----------------------------------------------------------------------------

case $choice in
    # ------------------------------------------------------------------------
    # 选项1: 标准识别（最常用）
    # 使用最佳模型，置信度0.2，保存可视化结果
    # 适合大多数场景，速度和准确率平衡
    # ------------------------------------------------------------------------
    1)
        echo -e "\n${YELLOW}运行改进版识别（使用最佳模型）...${NC}"
        
        # 检查模型文件是否存在
        if [ ! -f "$BEST_MODEL" ]; then
            echo -e "${RED}❌ 错误: 模型文件不存在${NC}"
            echo "   路径: $BEST_MODEL"
            echo "   请先训练模型或检查路径"
            exit 1
        fi
        
        # 执行识别
        python scripts/predict_digits_improved.py \
            --model "$BEST_MODEL" \
            --source valid \
            --conf 0.2 \
            --save-vis
        
        # 显示结果位置
        echo -e "\n${GREEN}✓ 识别完成！${NC}"
        echo "📄 结果文件: results/predictions_improved.txt"
        echo "🖼️  可视化: results/visualizations_improved/"
        ;;
    
    # ------------------------------------------------------------------------
    # 选项2: 低阈值识别
    # 降低置信度到0.15，可以检测出更多数字
    # 适合图片模糊或数字不清晰的情况
    # 可能会有更多误检
    # ------------------------------------------------------------------------
    2)
        echo -e "\n${YELLOW}运行低阈值识别（更敏感）...${NC}"
        
        python scripts/predict_digits_improved.py \
            --model "$BEST_MODEL" \
            --source valid \
            --conf 0.15 \
            --output results/predictions_conf015.txt \
            --save-vis
        
        echo -e "\n${GREEN}✓ 识别完成！${NC}"
        echo "📄 结果文件: results/predictions_conf015.txt"
        ;;
    
    # ------------------------------------------------------------------------
    # 选项3: 高清识别
    # 使用更大的输入尺寸(640)，提高识别精度
    # 适合需要高精度的场景
    # 处理速度会变慢（约2倍时间）
    # ------------------------------------------------------------------------
    3)
        echo -e "\n${YELLOW}运行高清识别（更精确）...${NC}"
        
        python scripts/predict_digits_improved.py \
            --model "$BEST_MODEL" \
            --source valid \
            --imgsz 640 \
            --output results/predictions_640.txt \
            --save-vis
        
        echo -e "\n${GREEN}✓ 识别完成！${NC}"
        echo "📄 结果文件: results/predictions_640.txt"
        ;;
    
    # ------------------------------------------------------------------------
    # 选项4: 查看结果和模型信息
    # 显示已有的模型、识别结果和可视化文件
    # 不执行任何识别操作
    # ------------------------------------------------------------------------
    4)
        echo -e "\n${YELLOW}查看已有结果和模型信息...${NC}"
        echo "=================================="
        
        # 显示最佳模型信息
        echo "📦 最佳模型:"
        if [ -f "$BEST_MODEL" ]; then
            ls -lh "$BEST_MODEL" | awk '{printf "   大小: %s\n", $5}'
            echo "   路径: $BEST_MODEL"
        else
            echo "   ⚠️  模型不存在"
        fi
        
        echo ""
        
        # 显示识别结果文件
        echo "📄 识别结果文件:"
        if ls results/*.txt 1> /dev/null 2>&1; then
            ls -lh results/*.txt | awk '{printf "   %s (%s)\n", $9, $5}'
        else
            echo "   暂无结果文件"
        fi
        
        echo ""
        
        # 显示可视化目录
        echo "🖼️  可视化结果:"
        if ls -d results/visualizations* 1> /dev/null 2>&1; then
            ls -d results/visualizations* | sed 's/^/   /'
        else
            echo "   暂无可视化结果"
        fi
        
        echo ""
        
        # 显示所有训练模型
        echo "🎯 所有训练模型:"
        if ls runs/digit_yolo/*/weights/best.pt 1> /dev/null 2>&1; then
            for model in runs/digit_yolo/*/weights/best.pt; do
                exp_name=$(dirname $(dirname "$model"))
                exp_name=$(basename "$exp_name")
                size=$(ls -lh "$model" | awk '{print $5}')
                echo "   [$exp_name] $size"
            done
        else
            echo "   暂无模型文件"
        fi
        
        echo "=================================="
        ;;
    
    # ------------------------------------------------------------------------
    # 选项0: 退出程序
    # ------------------------------------------------------------------------
    0)
        echo -e "\n${GREEN}👋 再见！${NC}"
        exit 0
        ;;
    
    # ------------------------------------------------------------------------
    # 无效选项处理
    # ------------------------------------------------------------------------
    *)
        echo -e "\n${RED}❌ 无效选项: $choice${NC}"
        echo "请输入 0-4 之间的数字"
        exit 1
        ;;
esac

# ----------------------------------------------------------------------------
# 5. 显示后续操作提示
# ----------------------------------------------------------------------------

echo ""
echo "=================================================="
echo -e "${GREEN}✓ 操作完成！${NC}"
echo "=================================================="
echo ""

# 常用命令提示
echo "📊 查看结果的命令:"
echo "  cat results/predictions*.txt      # 查看识别结果"
echo "  open results/visualizations*/     # 打开可视化图片（macOS）"
echo ""

# 更多操作提示
echo "🔧 更多操作:"
echo "  python scripts/predict_digits_improved.py --help  # 查看完整参数说明"
echo "  python scripts/preprocess_images.py --help        # 图片预处理帮助"
echo ""

# 文档链接
echo "📚 项目文档:"
echo "  cat README.md           # 完整文档"
echo "  cat QUICKSTART.md       # 快速开始指南"
echo "  cat FINAL_REPORT.md     # 项目完成报告"
echo ""