import torch
import clip
from torch.utils.data import DataLoader
from dataset import FlickrStreamer
from torchvision import transforms
from tqdm import tqdm

# --- SETTINGS ---
INDEX_LIMIT = 5000  # How many images to make searchable
BATCH_SIZE = 32
SAVE_PATH = "flickr_embeddings.pt"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ----------------

def build_index():
    print(f"Loading CLIP on {DEVICE}...")
    model, preprocess = clip.load("ViT-B/32", device=DEVICE)
    
    # We need the images resized for CLIP (224x224)
    tf = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor()
    ])
    
    # Load Dataset (Stream mode)
    print(f"Connecting to Deep Lake (Limit: {INDEX_LIMIT})...")
    dataset = FlickrStreamer(limit=INDEX_LIMIT, transform=tf)
    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
    
    all_embeddings = []
    
    print("Indexing images (calculating fingerprints)...")
    with torch.no_grad():
        for images in tqdm(loader):
            images = images.to(DEVICE)
            
            # Calculate features
            features = model.encode_image(images)
            
            # Normalize features (Crucial for Cosine Similarity later)
            features /= features.norm(dim=-1, keepdim=True)
            
            all_embeddings.append(features.cpu())
            
    # Concatenate all batches into one big list
    final_index = torch.cat(all_embeddings)
    
    # Save to file
    torch.save(final_index, SAVE_PATH)
    print(f"\nSuccess! Saved {len(final_index)} image fingerprints to '{SAVE_PATH}'")

if __name__ == "__main__":
    build_index()