import os
import torch
from PIL import Image
from transformers import ChineseCLIPProcessor, ChineseCLIPModel

def clip_result(image_path):
    # 設置設備
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Get the directory where this script is located
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Construct the full path to the file in the subfolder
    model_path = os.path.join(script_dir, 'artifacts/models', 'best_clip_model.pth')

    print("model_path:", model_path)

    # 載入訓練好的模型和處理器
    model = ChineseCLIPModel.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")
    model.load_state_dict(torch.load(model_path, map_location=device))
    model = model.to(device)
    model.eval()

    processor = ChineseCLIPProcessor.from_pretrained("OFA-Sys/chinese-clip-vit-base-patch16")

    # 1. 加載圖片
    # image_path = '/content/drive/MyDrive/幽靈吉伊卡哇.png'
    image = Image.open(image_path)

    # 2. 加載中文詞彙表
    with open('./chiikawa/word_list.txt', 'r', encoding='utf-8') as f:
        vocab = [line.strip() for line in f.readlines()]

    # 3. 圖像和文本處理
    batch_size = 16  # 每次處理16個詞彙
    similarities = []

    # 釋放未使用的顯存
    torch.cuda.empty_cache()

    with torch.no_grad():
        for i in range(0, len(vocab), batch_size):
            batch_vocab = vocab[i:i + batch_size]
            inputs = processor(
                text=batch_vocab,
                images=image,
                return_tensors="pt",
                padding=True
            ).to(device)

            # 推理並進行相似度計算
            outputs = model(**inputs)
            image_embeds = outputs.image_embeds / outputs.image_embeds.norm(dim=-1, keepdim=True)
            text_embeds = outputs.text_embeds / outputs.text_embeds.norm(dim=-1, keepdim=True)
            similarity = torch.matmul(image_embeds, text_embeds.T).squeeze(0)
            similarities.append(similarity)

    # 4. 合併所有相似度
    similarity = torch.cat(similarities, dim=0)

    # 5. 找到相似度最高的詞彙
    top_k = 3
    top_k_indices = torch.topk(similarity, top_k).indices.tolist()
    top_k_words = [vocab[idx] for idx in top_k_indices]

    # 6. 輸出最接近的前3名中文詞彙

    # print("最接近的前3名中文詞彙是:")
    # for rank, word in enumerate(top_k_words, 1):
    #     print(f"{rank}. {word}")

    return top_k_words