#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
FAISS多模态检索系统简单测试
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from multimodal_retrieval_faiss import MultimodalRetrievalFAISS
from PIL import Image
import numpy as np

def test_text_retrieval():
    print("=== 测试文本检索 ===")
    
    # 初始化检索系统
    print("初始化检索系统...")
    retrieval = MultimodalRetrievalFAISS(
        model_name="OpenSearch-AI/Ops-MM-embedding-v1-7B",
        use_all_gpus=True,
        index_path="faiss_index_test"
    )
    
    # 测试文本
    texts = [
        "一只可爱的橘色猫咪在沙发上睡觉",
        "城市夜景中的高楼大厦和车流",
        "阳光明媚的海滩上，人们在冲浪和晒太阳",
        "美味的意大利面配红酒和沙拉",
        "雪山上滑雪的运动员"
    ]
    
    # 添加文本
    print("\n添加文本到检索系统...")
    text_ids = retrieval.add_texts(texts)
    print(f"添加了{len(text_ids)}条文本")
    print(f"当前向量数量: {retrieval.get_vector_count()}")
    
    # 测试文本搜索
    print("\n测试文本搜索...")
    queries = ["一只猫在睡觉", "都市风光", "海边的景色"]
    
    for query in queries:
        print(f"\n查询: {query}")
        results = retrieval.search_by_text(query, k=2)
        for i, result in enumerate(results):
            print(f"  结果 {i+1}: {result.get('text', 'N/A')} (分数: {result.get('score', 0):.4f})")
    
    # 保存索引
    print("\n保存索引...")
    retrieval.save_index()
    
    print("\n测试完成！")

if __name__ == "__main__":
    test_text_retrieval()