"""
FAISS相似度检索系统快速演示脚本
一键运行完整的索引构建和检索演示
"""

import os
import time


def check_dependencies():
    """检查依赖包"""
    required_packages = [
        'torch', 'torchvision', 'faiss', 'PIL', 'numpy', 'tqdm', 'matplotlib'
    ]
    
    missing_packages = []
    
    for package in required_packages:
        try:
            if package == 'PIL':
                import PIL
            elif package == 'faiss':
                import faiss
            else:
                __import__(package)
        except ImportError:
            missing_packages.append(package)
    
    if missing_packages:
        print("缺少以下依赖包:")
        for pkg in missing_packages:
            if pkg == 'faiss':
                print(f"  {pkg} (安装命令: pip install faiss-cpu)")
            else:
                print(f"  {pkg}")
        print("\n请先安装缺少的依赖包")
        return False
    
    return True

def check_files():
    """检查必要文件"""
    required_files = [
        'net/resnet_embedding.py',
        'settings/settings.py',
        'model/embedding_20250917_145342/best_embedding_model.pth',
        'dataset/train'
    ]
    
    missing_files = []
    
    for file_path in required_files:
        if not os.path.exists(file_path):
            missing_files.append(file_path)
    
    if missing_files:
        print("缺少以下必要文件:")
        for file_path in missing_files:
            print(f"  {file_path}")
        return False
    
    return True

def run_index_building():
    """运行索引构建"""
    print("=" * 60)
    print("步骤 1: 构建FAISS索引")
    print("=" * 60)
    
    try:
        from faiss_vector_db.build_faiss_index import main as build_main
        build_main()
        return True
    except Exception as e:
        print(f"索引构建失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def run_search_demo():
    """运行检索演示"""
    print("\n" + "=" * 60)
    print("步骤 2: 相似度检索演示")
    print("=" * 60)
    
    try:
        from faiss_vector_db.build_faiss_index import FAISSSearcher
        
        # 配置参数
        index_dir = "DishClassification/faiss_index"
        model_path = "../model/embedding_20250917_145342/best_embedding_model.pth"
        
        # 检查索引是否存在
        if not os.path.exists(os.path.join(index_dir, 'faiss_index.bin')):
            print("错误: FAISS索引不存在，请先构建索引")
            return False
        
        # 创建检索器
        searcher = FAISSSearcher(index_dir, model_path)
        
        # 显示索引信息
        print(f"索引信息:")
        print(f"  总向量数: {searcher.index.ntotal}")
        print(f"  特征维度: {searcher.embedding_dim}")
        print(f"  类别数: {len(searcher.class_names)}")
        print(f"  类别: {', '.join(searcher.class_names)}")
        
        # 找一些测试图片
        test_images = []
        test_dirs = ['dataset/test', 'dataset/val', 'dataset/train']
        
        for test_dir in test_dirs:
            if os.path.exists(test_dir):
                for root, dirs, files in os.walk(test_dir):
                    for file in files:
                        if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                            test_images.append(os.path.join(root, file))
                            if len(test_images) >= 3:  # 只测试3张图片
                                break
                    if len(test_images) >= 3:
                        break
                if len(test_images) >= 3:
                    break
        
        if not test_images:
            print("未找到测试图片")
            return False
        
        # 执行相似度检索演示
        print(f"\n开始相似度检索演示 (测试 {len(test_images)} 张图片):")
        print("-" * 50)
        
        total_time = 0
        
        for i, test_image in enumerate(test_images, 1):
            print(f"\n[{i}/{len(test_images)}] 测试图片: {os.path.relpath(test_image)}")
            
            start_time = time.time()
            results = searcher.search_similar(test_image, k=5)
            search_time = time.time() - start_time
            total_time += search_time
            
            print(f"搜索耗时: {search_time:.3f}秒")
            print("最相似的5张图片:")
            
            for j, (img_path, class_name, score) in enumerate(results, 1):
                rel_path = os.path.relpath(img_path)
                print(f"  {j}. {os.path.basename(img_path):20s} "
                      f"类别: {class_name:8s} "
                      f"相似度: {score:.4f}")
            
            # 统计类别分布
            class_counts = {}
            for _, class_name, _ in results:
                class_counts[class_name] = class_counts.get(class_name, 0) + 1
            
            print(f"  类别分布: {dict(class_counts)}")
        
        print(f"\n演示完成!")
        print(f"平均搜索时间: {total_time/len(test_images):.3f}秒")
        
        return True
        
    except Exception as e:
        print(f"检索演示失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def show_usage_examples():
    """显示使用示例"""
    print("\n" + "=" * 60)
    print("使用示例")
    print("=" * 60)
    
    print("\n1. 查看索引信息:")
    print("   python faiss_similarity_search.py --info")
    
    print("\n2. 单张图片相似度检索:")
    print("   python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg")
    
    print("\n3. 显示检索结果图片:")
    print("   python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images")
    
    print("\n4. 批量测试:")
    print("   python faiss_similarity_search.py --batch_test dataset/test --max_tests 10")
    
    print("\n5. 保存结果图片:")
    print("   python faiss_similarity_search.py --query dataset/test/回锅肉/aug_000.jpg --show_images --save_result")

def main():
    """主函数"""
    print("FAISS相似度检索系统 - 快速演示")
    print("=" * 60)
    
    # 检查依赖
    print("检查依赖包...")
    if not check_dependencies():
        return
    print("✓ 依赖包检查通过")
    
    # 检查文件
    print("检查必要文件...")
    if not check_files():
        return
    print("✓ 文件检查通过")
    
    # 检查是否已有索引
    index_exists = os.path.exists("DishClassification/faiss_index/faiss_index.bin")
    
    if index_exists:
        print("✓ 发现已存在的FAISS索引")
        
        # 询问是否重新构建
        while True:
            choice = input("\n是否重新构建索引? (y/n): ").lower().strip()
            if choice in ['y', 'yes', '是']:
                if not run_index_building():
                    return
                break
            elif choice in ['n', 'no', '否']:
                print("跳过索引构建，直接进行检索演示")
                break
            else:
                print("请输入 y 或 n")
    else:
        print("未发现FAISS索引，开始构建...")
        if not run_index_building():
            return
    
    # 运行检索演示
    if not run_search_demo():
        return
    
    # 显示使用示例
    show_usage_examples()
    
    print("\n" + "=" * 60)
    print("演示完成! 🎉")
    print("您现在可以使用 faiss_similarity_search.py 进行更多测试")
    print("=" * 60)

if __name__ == "__main__":
    main()