# Python libraries for numerical operations
import numpy as np
from sklearn.metrics.pairwise import euclidean_distances, cosine_similarity

# --- Step 1: Generate Synthetic Biological Embeddings ---
# In real-world biological pipelines, these would come from models
# like ProtTrans for proteins, RDKit for molecules, or deep learning models for genomic data.
# We forge a dataset of 1000 'protein-like' embeddings, each with 128 dimensions.
# These dimensions capture abstract features relevant to the protein's function or structure.
n_vectors = 1000  # Number of biological entities (e.g., proteins)
embedding_dim = 128 # Dimensionality of each embedding vector

# Simulate embeddings using random numbers for demonstration
# In practice, these would be rich, context-aware representations.
np.random.seed(42) # For reproducibility
database_embeddings = np.random.rand(n_vectors, embedding_dim).astype('float32')

print(f"Generated database with {n_vectors} embeddings, each of {embedding_dim} dimensions.")
print(f"Example embedding shape: {database_embeddings[0].shape}")

# --- Step 2: Define Query Embedding ---
# This is the vector representation of the biological entity we want to find similar ones to.
query_embedding = np.random.rand(1, embedding_dim).astype('float32')
print(f"Generated query embedding of shape: {query_embedding.shape}")

# --- Step 3: Implement Basic Similarity Metrics ---
# We decode the fundamental mathematical operations that quantify 'similarity' or 'distance'.
# Euclidean distance measures the straight-line distance between two points.
# Cosine similarity measures the cosine of the angle between two vectors, indicating directional similarity.

def calculate_euclidean_distance(vec1, vec2):
    return np.linalg.norm(vec1 - vec2)

def calculate_cosine_similarity(vec1, vec2):
    # Ensure vectors are 1D for dot product if coming from (1, D) shape
    vec1 = vec1.flatten()
    vec2 = vec2.flatten()
    dot_product = np.dot(vec1, vec2)
    norm_vec1 = np.linalg.norm(vec1)
    norm_vec2 = np.linalg.norm(vec2)
    if norm_vec1 == 0 or norm_vec2 == 0:
        return 0.0 # Handle division by zero for zero vectors
    return dot_product / (norm_vec1 * norm_vec2)

# Demonstrate with two arbitrary embeddings from our database
vec_a = database_embeddings[0]
vec_b = database_embeddings[1]

print(f"\nDistance between vec_a and vec_b (Euclidean): {calculate_euclidean_distance(vec_a, vec_b):.4f}")
print(f"Similarity between vec_a and vec_b (Cosine): {calculate_cosine_similarity(vec_a, vec_b):.4f}")

# Using sklearn for verification (often optimized for speed)
print(f"Distance between vec_a and vec_b (sklearn Euclidean): {euclidean_distances(vec_a.reshape(1, -1), vec_b.reshape(1, -1))[0,0]:.4f}")
print(f"Similarity between vec_a and vec_b (sklearn Cosine): {cosine_similarity(vec_a.reshape(1, -1), vec_b.reshape(1, -1))[0,0]:.4f}")

# --- Step 4: Illustrate Brute-Force Exact Search for a Single Query ---
# This is the most straightforward method: compare the query to every vector in the database.
# We will delve deeper into its implications later.

def brute_force_exact_search(query_vec, database_vectors, k=5, metric='euclidean'):
    distances = []
    for i, db_vec in enumerate(database_vectors):
        if metric == 'euclidean':
            dist = calculate_euclidean_distance(query_vec, db_vec)
            distances.append((dist, i))
        elif metric == 'cosine':
            # For cosine similarity, higher value means more similar, so we store negative for sorting
            sim = calculate_cosine_similarity(query_vec, db_vec)
            distances.append((-sim, i)) # Store negative similarity to sort by smallest 'distance'
    
    # Sort by distance/negative similarity and get the top k
    distances.sort()
    
    results = []
    for dist_val, idx in distances[:k]:
        if metric == 'euclidean':
            results.append({'index': idx, 'distance': dist_val})
        elif metric == 'cosine':
            results.append({'index': idx, 'similarity': -dist_val}) # Revert negative similarity
            
    return results

print("\n--- Brute-Force Exact Search (Top 5 Euclidean) ---")
bf_results_euclidean = brute_force_exact_search(query_embedding, database_embeddings, k=5, metric='euclidean')
for res in bf_results_euclidean:
    print(f"Index: {res['index']}, Distance: {res['distance']:.4f}")

print("\n--- Brute-Force Exact Search (Top 5 Cosine) ---")
bf_results_cosine = brute_force_exact_search(query_embedding, database_embeddings, k=5, metric='cosine')
for res in bf_results_cosine:
    print(f"Index: {res['index']}, Similarity: {res['similarity']:.4f}")

# Python libraries for numerical operations and timing
import numpy as np
from sklearn.metrics.pairwise import euclidean_distances
import time

# --- Step 1: Generate a Larger Synthetic Biological Embedding Dataset ---
# We scale up our dataset to better illustrate performance challenges.
# Let's consider 100,000 'molecular' embeddings, each 256 dimensions.
# This represents a moderately sized molecular library for screening.
n_vectors_large = 100_000 # Number of molecules
embedding_dim_large = 256 # Features per molecule

np.random.seed(42) # Ensure reproducibility
database_embeddings_large = np.random.rand(n_vectors_large, embedding_dim_large).astype('float32')
query_embedding_large = np.random.rand(1, embedding_dim_large).astype('float32')

print(f"Generated large database with {n_vectors_large} embeddings, each of {embedding_dim_large} dimensions.")

# --- Step 2: Optimized Brute-Force Exact Search Function ---
# We optimize the brute-force search by leveraging NumPy's vectorized operations.
# This is significantly faster than explicit Python loops for numerical tasks.
# For Euclidean distance, np.linalg.norm is efficient.
# For cosine similarity, sklearn.metrics.pairwise.cosine_similarity is ideal.

def optimized_brute_force_search(query_vec, database_vectors, k=5, metric='euclidean'):
    if metric == 'euclidean':
        # Reshape query_vec for sklearn compatibility (1, D)
        distances = euclidean_distances(query_vec.reshape(1, -1), database_vectors).flatten()
        # Get indices of the smallest distances
        nearest_indices = np.argsort(distances)[:k]
        results = []
        for idx in nearest_indices:
            results.append({'index': idx, 'distance': distances[idx]})
    elif metric == 'cosine':
        # Cosine similarity returns values from -1 to 1; higher is better.
        # We want the highest similarities, so we sort in descending order.
        similarities = cosine_similarity(query_vec.reshape(1, -1), database_vectors).flatten()
        # Get indices of the largest similarities (descending sort)
        nearest_indices = np.argsort(similarities)[::-1][:k]
        results = []
        for idx in nearest_indices:
            results.append({'index': idx, 'similarity': similarities[idx]})
    else:
        raise ValueError("Unsupported metric. Choose 'euclidean' or 'cosine'.")
    return results

# --- Step 3: Benchmark Brute-Force Performance ---
print("\n--- Benchmarking Optimized Brute-Force Exact Search ---")

start_time = time.time()
bf_results_large = optimized_brute_force_search(query_embedding_large, database_embeddings_large, k=10, metric='euclidean')
end_time = time.time()

print(f"Brute-force search for {n_vectors_large} vectors took {end_time - start_time:.4f} seconds.")
print("Top 10 Exact Neighbors (Euclidean):")
for res in bf_results_large:
    print(f"Index: {res['index']}, Distance: {res['distance']:.4f}")


start_time = time.time()
bf_results_large_cosine = optimized_brute_force_search(query_embedding_large, database_embeddings_large, k=10, metric='cosine')
end_time = time.time()
print(f"Brute-force search for {n_vectors_large} vectors (cosine) took {end_time - start_time:.4f} seconds.")
print("Top 10 Exact Neighbors (Cosine):")
for res in bf_results_large_cosine:
    print(f"Index: {res['index']}, Similarity: {res['similarity']:.4f}")

# --- Step 4: Discussion on Scalability ---
# Observe how the time increases significantly with dataset size.
# For N=10^6 vectors, this can easily take minutes or hours on a CPU.
# This makes real-time applications or large-scale batch processing infeasible.
print("\n--- Scalability Implications ---")
print(f"For a database of {n_vectors_large} vectors, brute-force exact search is already showing significant latency.")
print("As N grows to millions or billions, this approach becomes computationally prohibitive, demanding minutes, hours, or even days for a single query on standard hardware.")
print("Memory footprint also becomes a concern: storing all distances in memory for sorting can be problematic for very large N and k.")

# Python libraries for numerical operations, timing, and FAISS
import numpy as np
import time
import faiss

# Ensure FAISS is installed: pip install faiss-cpu or pip install faiss-gpu

# --- Step 1: Generate a Large Synthetic Biological Embedding Dataset (re-using previous one) ---
n_vectors_large = 1_000_000 # Let's scale up to 1 million for FAISS demonstration
embedding_dim_large = 256

np.random.seed(42) # For reproducibility
database_embeddings_large = np.random.rand(n_vectors_large, embedding_dim_large).astype('float32')
query_embedding_large = np.random.rand(10, embedding_dim_large).astype('float32') # 10 queries for batch search

print(f"Generated large database for FAISS with {n_vectors_large} embeddings, each of {embedding_dim_large} dimensions.")
print(f"Generated {query_embedding_large.shape[0]} query embeddings.")

# --- Step 2: Initialize a FAISS Index ---
# We decode the FAISS index types. IVF (Inverted File Index) is a common and powerful choice.
# index_factory is a flexible way to create complex indexes.
# 'IDx,Flat' for a simple index.
# 'IVF100,Flat' means Inverted File with 100 clusters, then Flat (exact) search within clusters.
# 'IVF100,PQ64' means IVF with 100 clusters, then Product Quantization with 64 bytes for compression.

# For simplicity, we start with a basic IVF index. 
# nlist = number of inverted lists (clusters). A good heuristic is sqrt(N) or 4*sqrt(N).
# nlist = 100 is chosen for this example, adjust based on N.
nlist = 100 # Number of Voronoi cells (clusters)
quantizer = faiss.IndexFlatL2(embedding_dim_large) # The quantizer is an index that stores the centroids of the Voronoi cells
index = faiss.IndexIVFFlat(quantizer, embedding_dim_large, nlist, faiss.METRIC_L2) # Use L2 for Euclidean distance

print(f"Initialized FAISS IndexIVFFlat with {nlist} clusters.")

# --- Step 3: Train the FAISS Index ---
# The index needs to be 'trained' on a subset of the data to learn the clusters.
# This step is crucial for IVF indexes. We use a sample of the database.
# If the database is small, train on all data. For large databases, use a representative subset.

# Ensure the index is not already trained if running multiple times.
if not index.is_trained:
    print("Training FAISS index...")
    start_time = time.time()
    # Using a subset for training (e.g., 10% or more, depending on data distribution)
    train_data_size = min(50_000, n_vectors_large) # Train on up to 50k vectors or full dataset if smaller
    index.train(database_embeddings_large[:train_data_size])
    end_time = time.time()
    print(f"FAISS index training took {end_time - start_time:.4f} seconds.")

# --- Step 4: Add Vectors to the FAISS Index ---
# Once trained, we add all our biological embeddings to the index.
print("Adding vectors to FAISS index...")
start_time = time.time()
index.add(database_embeddings_large)
end_time = time.time()
print(f"Added {index.ntotal} vectors to FAISS index in {end_time - start_time:.4f} seconds.")

# --- Step 5: Perform Approximate Nearest Neighbor Search ---
# nprobe determines how many clusters to search. Higher nprobe = higher accuracy, lower speed.
# It's a key parameter for accuracy-speed trade-off.
index.nprobe = 10 # Search in 10 clusters (out of 100)
k_faiss = 10 # Number of nearest neighbors to retrieve

print(f"\nPerforming FAISS ANN search with nprobe={index.nprobe} for k={k_faiss} neighbors...")
start_time = time.time()
distances_faiss, indices_faiss = index.search(query_embedding_large, k_faiss) # Search for 10 queries
end_time = time.time()

print(f"FAISS search for {query_embedding_large.shape[0]} queries took {end_time - start_time:.4f} seconds.")
print("Example FAISS Results (Query 0, Top 10):")
for i in range(k_faiss):
    print(f"Index: {indices_faiss[0, i]}, Distance: {distances_faiss[0, i]:.4f}")

# --- Step 6: Explore Trade-offs: Adjusting nprobe ---
# We illustrate how changing nprobe impacts search time and potential accuracy.
print("\n--- Exploring nprobe Trade-offs ---")
index.nprobe = 1 # Very fast, potentially lower recall
start_time = time.time()
distances_faiss_low_nprobe, indices_faiss_low_nprobe = index.search(query_embedding_large, k_faiss)
end_time = time.time()
print(f"FAISS search with nprobe=1 took {end_time - start_time:.4f} seconds.")

index.nprobe = 50 # Slower, potentially higher recall
start_time = time.time()
distances_faiss_high_nprobe, indices_faiss_high_nprobe = index.search(query_embedding_large, k_faiss)
end_time = time.time()
print(f"FAISS search with nprobe=50 took {end_time - start_time:.4f} seconds.")

# In a real scenario, we would evaluate recall by comparing FAISS results to exact results.
# This demonstrates the core speed advantage and tunable accuracy of FAISS.

# Python libraries for numerical operations, timing, FAISS, and evaluation
import numpy as np
import time
import faiss
from sklearn.metrics import recall_score

# --- Step 1: Generate a Large Synthetic Biological Embedding Dataset (consistent for comparison) ---
n_vectors_final_eval = 100_000 # Using 100k for faster demonstration of comparison
embedding_dim_final_eval = 256

np.random.seed(42) 
database_embeddings_final_eval = np.random.rand(n_vectors_final_eval, embedding_dim_final_eval).astype('float32')
query_embeddings_final_eval = np.random.rand(100, embedding_dim_final_eval).astype('float32') # Multiple queries for robust benchmarking

print(f"Generated database: {n_vectors_final_eval} embeddings, dim: {embedding_dim_final_eval}")
print(f"Generated queries: {query_embeddings_final_eval.shape[0]} queries")

# --- Step 2: Establish Brute-Force Baseline (Exact Results) ---
def get_exact_neighbors(query_vecs, database_vecs, k=5, metric='euclidean'):
    exact_indices = []
    for query_vec in query_vecs:
        if metric == 'euclidean':
            distances = faiss.pairwise_L2sqr(query_vec.reshape(1, -1), database_vecs).flatten()
            nearest_indices = np.argsort(distances)[:k]
        elif metric == 'cosine': # For cosine, FAISS works with L2 on normalized vectors
            # Normalize vectors for cosine similarity if using L2 distance
            norm_query = query_vec / np.linalg.norm(query_vec)
            norm_database = database_vecs / np.linalg.norm(database_vecs, axis=1, keepdims=True)
            # Cosine similarity is 1 - L2_distance_sq/2 for normalized vectors
            # So, min L2_distance_sq means max cosine similarity
            distances = faiss.pairwise_L2sqr(norm_query.reshape(1, -1), norm_database).flatten()
            nearest_indices = np.argsort(distances)[:k]
        else:
            raise ValueError("Metric not supported.")
        exact_indices.append(set(nearest_indices)) # Use set for easier comparison
    return exact_indices

print("\n--- Calculating Exact Nearest Neighbors (Brute-Force Baseline) ---")
start_time = time.time()
exact_results = get_exact_neighbors(query_embeddings_final_eval, database_embeddings_final_eval, k=5, metric='euclidean')
end_time = time.time()
exact_time = end_time - start_time
print(f"Brute-force exact search for {query_embeddings_final_eval.shape[0]} queries took {exact_time:.4f} seconds.")

# --- Step 3: Configure and Benchmark FAISS (ANN) Search ---
nlist_faiss = 100 # Number of clusters
quantizer_faiss = faiss.IndexFlatL2(embedding_dim_final_eval)
index_faiss = faiss.IndexIVFFlat(quantizer_faiss, embedding_dim_final_eval, nlist_faiss, faiss.METRIC_L2)

if not index_faiss.is_trained:
    print("Training FAISS index for evaluation...")
    index_faiss.train(database_embeddings_final_eval[:min(50000, n_vectors_final_eval)]) # Train on subset
    index_faiss.add(database_embeddings_final_eval)

# Test different nprobe values for trade-off analysis
nprobe_values = [1, 5, 10, 20]
k_ann = 5 # Retrieve 5 neighbors

faiss_results_by_nprobe = {}

print("\n--- Benchmarking FAISS ANN Search with Varying nprobe ---")
for nprobe in nprobe_values:
    index_faiss.nprobe = nprobe
    start_time = time.time()
    distances_ann, indices_ann = index_faiss.search(query_embeddings_final_eval, k_ann)
    end_time = time.time()
    faiss_time = end_time - start_time
    
    # Calculate recall: How many of the exact neighbors did FAISS find?
    total_recall = []
    for i in range(query_embeddings_final_eval.shape[0]):
        ann_found = set(indices_ann[i])
        true_positives = len(exact_results[i].intersection(ann_found))
        total_recall.append(true_positives / k_ann)
    
    avg_recall = np.mean(total_recall)
    faiss_results_by_nprobe[nprobe] = {
        'time': faiss_time,
        'recall': avg_recall,
        'speedup_factor': exact_time / faiss_time if faiss_time > 0 else float('inf')
    }
    print(f"nprobe={nprobe}: Time={faiss_time:.4f}s, Avg Recall={avg_recall:.4f}, Speedup vs Brute-Force={faiss_results_by_nprobe[nprobe]['speedup_factor']:.2f}x")

# --- Step 4: Analyze Trade-offs and Provide Recommendations ---
print("\n--- Analysis of Trade-offs ---")
print("We observe a clear trade-off: as nprobe increases, recall improves (closer to exact results) but search time also increases. Conversely, reducing nprobe drastically cuts search time at the cost of recall.")

print("\n--- Strategic Decisions for Biological Pipelines ---")
print("1.  **Critical Precision (e.g., drug lead validation, rare variant identification):** When recall MUST be 1.0, and the database is manageable (e.g., < 100k-1M vectors depending on D and hardware), choose Brute-Force. Validate all hits with biological assays.")
print("2.  **High-Throughput Screening (e.g., initial compound library filtering, large-scale protein similarity search):** When speed is paramount for millions/billions of vectors, and a slight recall drop (e.g., 0.8-0.95) is acceptable, FAISS is indispensable. Tune nprobe to balance speed and an acceptable level of recall. Optimize for throughput.")
print("3.  **Memory Constraints:** FAISS offers advanced index types (e.g., Product Quantization) that significantly reduce memory footprint, enabling searches on datasets that would otherwise exceed RAM capacity, especially critical for embedded systems or massive on-disk indexes.")
print("4.  **Hybrid Approaches:** Consider a two-stage approach: use FAISS for a fast, coarse pre-filtering to identify a smaller candidate set, then apply brute-force or more precise calculations on this reduced set for final ranking.")

print("\nForge your vector search strategy with these insights to optimize your biological discovery workflows.")