import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import requests
from io import BytesIO
import numpy as np

# 모델 불러오기
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
model.eval()

# 이미지 전처리 변환 정의
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 이미지 다운로드 및 전처리 함수
def download_and_preprocess_image(url):
    try:
        response = requests.get(url)
        img = Image.open(BytesIO(response.content)).convert('RGB')
        img_preprocessed = preprocess(img)
        return img_preprocessed
    except Exception as e:
        print(f"이미지 다운로드 및 전처리 실패: {e}")
        return None

# 이미지 특징 추출 함수
def extract_features(image):
    with torch.no_grad():
        features = model(image.unsqueeze(0))
    return features

# 코사인 유사도 계산 함수
def cosine_similarity(feature1, feature2):
    cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
    similarity = cos(feature1, feature2)
    return similarity

# 가장 비슷한 이미지 찾기 함수
def find_most_similar_image(source_url, target_urls):
    source_image = download_and_preprocess_image(source_url)
    source_features = extract_features(source_image)
    similarities = []
    for url in target_urls:
        target_image = download_and_preprocess_image(url)
        if target_image is None:
            continue
        target_features = extract_features(target_image)
        similarity = cosine_similarity(source_features, target_features)
        similarities.append(similarity.item())
    most_similar_index = np.argmax(similarities)
    most_similar_url = target_urls[most_similar_index]
    most_similar_score = similarities[most_similar_index]
    return most_similar_url, most_similar_score

# 실제 이미지 URL 사용
source_image_url = "https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/7a287e6a-b7cf-4d2d-97fc-5317be80ede1.jpg"
target_image_urls = [
    "https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/73709918-af96-4e51-b488-8bfb67c2eaf3.jpg",
    "https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/c6c92af6-23eb-4dec-ad33-ad4718200835.jpg",
    "https://file.percenty.co.kr/public/652bed8e865b1f32ea62bf1f/products/65e9ca5a8941d454f61ecb75/2d5148ac-082c-4539-a901-4a37b909b8c8.jpg",
]

most_similar_url, similarity_score = find_most_similar_image(source_image_url, target_image_urls)
print(f"가장 비슷한 이미지 URL: {most_similar_url}, 유사도 점수: {similarity_score}")