from pathlib import Path

import pandas as pd
from PIL import Image
from sentence_transformers import SentenceTransformer

from utils.utils import SRC_PATH

image_paths = sorted((SRC_PATH / "data" / "images").glob("*.jpg"))

if not image_paths:
    raise FileNotFoundError(f"No JPG images found in {SRC_PATH / 'data'}")

labels = [
    "top (tshirt, shirt, polo...)",
]

model = SentenceTransformer(
    "sentence-transformers/clip-ViT-B-32",
    device="cpu",
)

images = [Image.open(path).convert("RGB") for path in image_paths]

image_embeddings = model.encode(
    images,
    batch_size=16,
    normalize_embeddings=True,
    convert_to_numpy=True,
    show_progress_bar=True,
)

text_embeddings = model.encode(
    prompts,
    normalize_embeddings=True,
    convert_to_numpy=True,
)

similarities = image_embeddings @ text_embeddings.T

similarity_df = pd.DataFrame(
    similarities,
    index=[path.name for path in image_paths],
    columns=prompts,
)

print(similarity_df.round(3).to_string())