# Python environment setup for FAISS serialization and advanced indexing
# Requires faiss-cpu

import faiss
import numpy as np
import time
import os

# Assume 'protein_embeddings' from previous steps is available
EMBEDDING_DIM = protein_embeddings.shape[1]
FULL_DATA_SIZE = protein_embeddings.shape[0]

# --- 1. Advanced FAISS Index Creation (IndexIVFPQ) ---
# IndexIVFPQ combines inverted file indexing with Product Quantization.
# It's highly effective for very large datasets by significantly reducing memory footprint.

print("\n--- Activating FAISS IndexIVFPQ for memory-efficient indexing ---")

# 'nlist': Number of inverted lists (centroids). Must be less than num_vectors/4.
#          A common strategy is to make it proportional to sqrt(N) for N vectors.
#          For 1M vectors, 4096 is a good starting point.
nlist = 4096
if FULL_DATA_SIZE < 1000000: # Adjust nlist for smaller datasets if testing with less than 1M
    nlist = int(FULL_DATA_SIZE / 250)
    if nlist < 64: nlist = 64

# 'm': Number of subquantizers (e.g., 8, 16, 32). This controls the compression ratio.
#      Must divide EMBEDDING_DIM. A common choice is 8 or 16.
m = 64 # Example: Split 1024 dim into 16 parts of 64 dim each (1024/16 = 64). Or m=8 (1024/8 = 128).
      # If EMBEDDING_DIM is 1024, m=8 means 8 subvectors of 128 dimensions each.
      # If EMBEDDING_DIM is 1024, m=16 means 16 subvectors of 64 dimensions each.
      # We choose m=64 for a 1024-dim embedding to demonstrate strong compression.
      # For 1024 dim, m=8 would be sub-vector size 128. m=16, sub-vector size 64.
      # Ensure EMBEDDING_DIM % m == 0
if EMBEDDING_DIM % m != 0:
    # Adjust m if it doesn't divide EMBEDDING_DIM cleanly, or choose a different m
    print(f"Warning: EMBEDDING_DIM ({EMBEDDING_DIM}) not divisible by m ({m}). Adjusting m.")
    for i in range(m, 0, -1):
        if EMBEDDING_DIM % i == 0:
            m = i
            print(f"Adjusted m to {m}.")
            break

# 'nbits': Number of bits per subquantizer. 8 bits is standard (256 clusters per subvector).
nbits = 8

# 1. Quantizer: Coarse quantizer (e.g., L2 distance)
quantizer = faiss.IndexFlatL2(EMBEDDING_DIM)

# 2. IndexIVFPQ: The main index with Product Quantization
faiss_ivfpq_index = faiss.IndexIVFPQ(quantizer, EMBEDDING_DIM, nlist, m, nbits)

# Ensure the index is trained. This is crucial for IVFFlat and IVFPQ.
# Training needs a representative subset of the data, ideally around 10*nlist.
# For 1M embeddings, 10*4096 = 40960 samples are suitable.
num_training_samples = min(FULL_DATA_SIZE, 10 * nlist)
if num_training_samples < FULL_DATA_SIZE:
    # Randomly sample for training to ensure representativeness
    train_idx = np.random.choice(FULL_DATA_SIZE, num_training_samples, replace=False)
    training_data = protein_embeddings[train_idx]
else:
    training_data = protein_embeddings

print(f"Training FAISS IndexIVFPQ with {training_data.shape[0]} samples...")
start_time = time.time()
faiss_ivfpq_index.train(training_data)
end_time = time.time()
print(f"FAISS IndexIVFPQ training time: {end_time - start_time:.2f} seconds")

# Add vectors to the index. For massive datasets, consider batching.
print("Adding vectors to FAISS IndexIVFPQ...")
start_time = time.time()
faiss_ivfpq_index.add(protein_embeddings)
end_time = time.time()
print(f"FAISS IndexIVFPQ add time: {end_time - start_time:.2f} seconds")
print(f"FAISS IndexIVFPQ contains {faiss_ivfpq_index.ntotal} vectors.")

# --- 2. Saving and Loading the Index (Persistence) ---
# Persisting the index prevents re-building it every time, critical for large datasets.
INDEX_FILE_FAISS = 'faiss_ivfpq_protein_index.faiss'

print(f"\nSaving FAISS index to {INDEX_FILE_FAISS}...")
faiss.write_index(faiss_ivfpq_index, INDEX_FILE_FAISS)
print("FAISS index saved.")

print(f"Loading FAISS index from {INDEX_FILE_FAISS}...")
loaded_faiss_index = faiss.read_index(INDEX_FILE_FAISS)
print(f"Loaded FAISS index contains {loaded_faiss_index.ntotal} vectors.")

# --- 3. Memory Footprint Estimation (Important for resource planning) ---
# For IndexIVFPQ, memory roughly (nlist * dim * sizeof(float32)) + (nbits * ntotal / 8)
# where ntotal/8 is for the PQ compressed vectors.

# Size of centroids (quantizer)
centroids_mem = nlist * EMBEDDING_DIM * 4 # 4 bytes per float32

# Size of compressed vectors
compressed_vectors_mem = (faiss_ivfpq_index.code_size * faiss_ivfpq_index.ntotal) # code_size is bytes per vector

total_memory_bytes = centroids_mem + compressed_vectors_mem
print(f"\nEstimated FAISS IndexIVFPQ memory footprint: {total_memory_bytes / (1024**2):.2f} MB")

# Compare to raw data size:
raw_data_memory_bytes = FULL_DATA_SIZE * EMBEDDING_DIM * 4
print(f"Raw data memory footprint: {raw_data_memory_bytes / (1024**2):.2f} MB")

# Key insight: Product Quantization drastically reduces memory while maintaining search efficacy.
# Parameter tuning (nlist, m, nbits) and persistence are paramount for large-scale deployments.