import pandas as pd
import torch
from sentence_transformers import SentenceTransformer
import numpy as np

def create_embeddings(
    input_file: str,
    output_file: str,
    model_name: str = "sentence-transformers/all-MiniLM-L6-v2"
):
    """
    Generate and save embeddings using SentenceTransformer.
    
    Args:
        input_file: Path to input CSV containing Title and Description columns
        output_file: Path to save embeddings
        model_name: Name of the sentence-transformer model to use
    """
    print("Loading data...")
    df = pd.read_csv(input_file)
    
    print(f"Loading model: {model_name}")
    embedder = SentenceTransformer(model_name)
    
    # Use GPU if available
    if not torch.cuda.is_available():
        print("Warning: No GPU found. Using CPU.")
    else:
        print("GPU found! Using CUDA.")
        embedder = embedder.to('cuda')
    
    # Generate combined strings for embedding
    print("Preparing text for embedding...")
    descriptions = df.apply(
        lambda row: f"Title: {row['Title']}; Description: {row['Description']}", 
        axis=1
    ).tolist()
    
    # Generate embeddings
    print("Generating embeddings...")
    embeddings = embedder.encode(
        descriptions, 
        show_progress_bar=True,
        batch_size=32 if torch.cuda.is_available() else 8
    )
    
    # Save embeddings
    print("Saving embeddings...")
    embeddings_df = pd.DataFrame(
        embeddings,
        columns=[f"dim_{i}" for i in range(embeddings.shape[1])]
    )
    embeddings_df.to_csv(output_file, index=False)
    
    print(f"Embeddings saved to {output_file}")
    print(f"Shape of embeddings: {embeddings.shape}")
    
if __name__ == "__main__":
    create_embeddings(
        input_file="packages_with_titles.csv",
        output_file="embeddings.csv"
    )