#!/usr/bin/env python3
"""
测试图像编码功能
"""

from multimodal_retrieval_multigpu import MultiGPUMultimodalRetrieval
import numpy as np
from PIL import Image

def test_image_encoding():
    print('正在初始化多GPU多模态检索系统...')
    retrieval = MultiGPUMultimodalRetrieval()

    # 测试文本编码
    print('测试文本编码...')
    text_embeddings = retrieval.encode_text_batch(['这是一个测试文本'])
    print(f'文本embedding形状: {text_embeddings.shape}')
    print(f'文本embedding数据类型: {text_embeddings.dtype}')

    # 测试图像编码
    print('测试图像编码...')
    test_images = ['sample_images/1755677101_1__.jpg']
    image_embeddings = retrieval.encode_image_batch(test_images)
    print(f'图像embedding形状: {image_embeddings.shape}')
    print(f'图像embedding数据类型: {image_embeddings.dtype}')
    
    # 测试两次相同图像的embedding是否一致
    print('测试embedding一致性...')
    image_embeddings2 = retrieval.encode_image_batch(test_images)
    consistency = np.allclose(image_embeddings, image_embeddings2, rtol=1e-5)
    print(f'相同图像embedding一致性: {consistency}')
    
    # 测试不同图像的embedding差异
    print('测试不同图像embedding差异...')
    test_images2 = ['sample_images/1755677101_2__.jpg']
    image_embeddings3 = retrieval.encode_image_batch(test_images2)
    similarity = np.dot(image_embeddings[0], image_embeddings3[0]) / (np.linalg.norm(image_embeddings[0]) * np.linalg.norm(image_embeddings3[0]))
    print(f'不同图像间相似度: {similarity:.4f}')

    # 验证维度一致性
    if text_embeddings.shape[1] == image_embeddings.shape[1]:
        print('✅ 文本和图像embedding维度一致')
    else:
        print('❌ 文本和图像embedding维度不一致')

    print('测试完成！')

if __name__ == "__main__":
    test_image_encoding()