# -*- coding: utf-8 -*-
"""
测试脚本 - 用训练好的模型识别学生手写数字图片

使用方法：
    python test_image.py path/to/image.png
    python test_image.py path/to/image.jpg
    python test_image.py path/to/folder/    # 识别文件夹内所有图片

依赖：
    pip install numpy pillow

图片要求：
    - 建议尺寸：28x28 或更大（程序会自动缩放）
    - 背景最好为白色，数字为黑色
    - 手写数字清晰、无遮挡
"""

import sys
import os
import numpy as np
from PIL import Image

# 尝试导入模型
try:
    from model_numpy import MLP
except ImportError:
    print("错误：请在 digit_mlp_class 目录下运行此脚本")
    print("  cd digit_mlp_class")
    print("  python test_image.py your_image.png")
    sys.exit(1)


def find_latest_model():
    """查找最新的模型文件"""
    model_files = [f for f in os.listdir('.') if f.startswith('mnist_mlp_') and f.endswith('.npy')]
    if not model_files:
        return None
    
    # 按时间戳分组
    timestamps = set()
    for f in model_files:
        # 文件格式: mnist_mlp_YYMMDD_HHMMSS_suffix.npy
        # 去掉后缀 _W1.npy, _b1.npy 等
        base = f.rsplit('_', 1)[0]
        timestamps.add(base)
    
    # 选择最新的
    latest = sorted(timestamps)[-1]
    
    # 检查完整模型
    required = ['W1', 'b1', 'W2', 'b2']
    for r in required:
        if not os.path.exists(f'{latest}_{r}.npy'):
            return None
    
    return latest


def load_model():
    """加载训练好的模型"""
    # 查找最新模型
    model_path = find_latest_model()
    
    if model_path is None:
        print("错误：未找到训练好的模型！")
        print("  请先运行: python main.py")
        print("  或确保当前目录有 mnist_mlp_*.npy 模型文件")
        return None
    
    print(f"加载模型: {model_path}")
    
    model = MLP(input_size=784, hidden_size=128, num_classes=10, learning_rate=0.1)
    
    model.W1 = np.load(f'{model_path}_W1.npy')
    model.b1 = np.load(f'{model_path}_b1.npy')
    model.W2 = np.load(f'{model_path}_W2.npy')
    model.b2 = np.load(f'{model_path}_b2.npy')
    
    print("模型加载成功！\n")
    return model


def preprocess_image(image_path):
    """
    图片预处理：将任意图片转为 28x28 归一化向量
    
    处理流程：
    1. 打开图片并转为灰度
    2. 调整大小为 28x28
    3. 转为NumPy数组并归一化到 [0, 1]
    4. 展平为 784 维向量
    """
    try:
        img = Image.open(image_path).convert('L')  # 转灰度
    except Exception as e:
        raise ValueError(f"无法打开图片: {e}")
    
    # 保存原始尺寸用于调试
    orig_width, orig_height = img.size
    
    # 缩放到 28x28
    img = img.resize((28, 28), Image.LANCZOS)
    
    # 转为NumPy数组并归一化
    img_array = np.array(img, dtype=np.float32) / 255.0
    
    # MNIST是白底黑字(0=背景, 1=数字)，如果图片是黑底白字需要反转
    # 检查是否是黑底白字：背景均值 > 0.5
    if img_array.mean() > 0.5:
        img_array = 1.0 - img_array
    
    # 展平为向量
    img_vector = img_array.flatten()
    
    print(f"  图片: {os.path.basename(image_path)}")
    print(f"  原始尺寸: {orig_width}x{orig_height}")
    print(f"  处理后尺寸: 28x28")
    
    return img_vector


def predict_image(model, image_path):
    """
    识别单张图片
    
    返回：
        predicted_digit: 预测的数字 (0-9)
        confidence: 置信度 (0-1)
        all_probs: 所有数字的概率
    """
    # 预处理
    img_vector = preprocess_image(image_path)
    
    # 预测
    probs = model.predict_proba(img_vector.reshape(1, -1))[0]
    predicted_digit = np.argmax(probs)
    confidence = probs[predicted_digit]
    
    return predicted_digit, confidence, probs


def print_results(digit, confidence, probs):
    """打印识别结果"""
    print(f"  预测结果: {digit}")
    print(f"  置信度: {confidence:.2%}")
    
    # 显示各数字概率
    print(f"\n  各数字概率:")
    print(f"  ", end="")
    for i in range(10):
        bar_len = int(probs[i] * 20)
        print(f" {i}:{'█' * bar_len}{'░' * (20-bar_len)} {probs[i]:.1%}")
    
    print()


def main():
    """主函数"""
    print("\n" + "=" * 60)
    print("手写数字识别 - 图片测试")
    print("=" * 60 + "\n")
    
    # 加载模型
    model = load_model()
    if model is None:
        sys.exit(1)
    
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("使用方法:")
        print("  python test_image.py path/to/image.png")
        print("  python test_image.py path/to/image.jpg")
        print("  python test_image.py path/to/folder/")
        print()
        print("示例:")
        print("  python test_image.py my_digit.png")
        print("  python test_image.py ./test_images/")
        sys.exit(1)
    
    target = sys.argv[1]
    
    # 收集所有要处理的图片
    image_paths = []
    
    if os.path.isdir(target):
        # 文件夹：收集所有图片
        extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.gif']
        for f in os.listdir(target):
            if any(f.lower().endswith(ext) for ext in extensions):
                image_paths.append(os.path.join(target, f))
        image_paths.sort()
    elif os.path.isfile(target):
        image_paths = [target]
    else:
        print(f"错误：文件或目录不存在: {target}")
        sys.exit(1)
    
    if not image_paths:
        print(f"错误：在 {target} 中未找到图片文件")
        sys.exit(1)
    
    print(f"找到 {len(image_paths)} 张图片，开始识别...\n")
    
    # 批量识别
    results = []
    for path in image_paths:
        try:
            digit, confidence, probs = predict_image(model, path)
            results.append((path, digit, confidence))
            print_results(digit, confidence, probs)
        except Exception as e:
            print(f"  识别失败: {e}\n")
    
    # 汇总结果
    print("=" * 60)
    print(f"识别完成！共 {len(results)} 张图片")
    print("=" * 60)
    print(f"\n{'文件名':<30} {'预测':<6} {'置信度':<10}")
    print("-" * 50)
    
    for path, digit, confidence in results:
        filename = os.path.basename(path)
        print(f"{filename:<30} {digit:<6} {confidence:.1%}")
    
    # 如果有真实标签（文件夹命名中包含），可以计算准确率
    print()


if __name__ == '__main__':
    main()