pq_index.go

// Package comet implements Product Quantization (PQ) for similarity search.
//
// WHAT IS PRODUCT QUANTIZATION?
// PQ is a lossy compression technique that dramatically reduces memory usage for vector
// storage while enabling approximate similarity search. It achieves compression ratios
// of 10-500x by dividing vectors into subspaces and quantizing each independently.
//
// THE CORE IDEA - DIVIDE AND COMPRESS:
// Instead of storing full high-dimensional vectors:
// 1. Divide each vector into M equal-sized subvectors (subspaces)
// 2. Learn a codebook of K centroids for each subspace via k-means
// 3. Encode each subvector with the ID of its nearest centroid
// 4. Store only these compact codes instead of original vectors
//
// COMPRESSION EXAMPLE:
// Original: 768 dims × 4 bytes = 3,072 bytes
// PQ (M=8, K=256): 8 subspaces × 1 byte = 8 bytes
// Compression: 384x smaller!
//
// TIME COMPLEXITY:
//   - Training: O(M × iterations × K × n × dsub) where dsub = dim/M
//   - Add: O(M × K × dsub) per vector
//   - Search: O(M × K × dsub + n × M) - table build + lookups
//
// WHEN TO USE PQ:
// - Dataset too large for RAM
// - Can tolerate 85-95% recall
// - L2 or inner product metric
// - Want massive compression
package comet

import (
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"sync"

	"github.com/RoaringBitmap/roaring"
)

// Compile-time checks
var _ VectorIndex = (*PQIndex)(nil)

// CalculatePQParams returns recommended PQ parameters for a given dimension.
// A neat utility function to get the recommended PQ parameters for a given dimension.
// Returns:
//   - M: Number of subquantizers (subspaces) that divides dim evenly
//   - Nbits: Bits per PQ code (default: 8, giving K=256 centroids per subspace)
func CalculatePQParams(dim int) (M int, Nbits int) {
	// Find M that divides dimension evenly
	// Prefer M=8 as good balance
	m := 8
	if dim%m != 0 {
		// Find divisor close to 8
		for m = 8; m <= 32; m++ {
			if dim%m == 0 {
				break
			}
		}
		if dim%m != 0 {
			m = 4 // Fallback
		}
	}

	return m, 8 // Standard: 256 centroids per subspace
}

// PQIndex represents a Product Quantization index.
//
// Memory layout:
//   - Codes: n × M bytes (compressed vectors)
//   - Codebooks: M × K × (dim/M) × 4 bytes
//   - Typical: 10-500x smaller than original
type PQIndex struct {
	// dim is the dimensionality of original vectors
	dim int

	// distanceKind specifies the distance metric
	distanceKind DistanceKind

	// distance is the distance calculator
	distance Distance

	// M is the number of subquantizers
	M int

	// Nbits is bits per PQ code
	Nbits int

	// Ksub is centroids per subquantizer (K = 2^Nbits)
	Ksub int

	// dsub is dimension of each subspace (dim/M)
	dsub int

	// codebooks stores M independent codebooks
	// codebooks[m][k*dsub:(k+1)*dsub] is centroid k in subspace m
	codebooks [][]float32

	// codes stores compressed representations
	// codes[i] is M bytes, one per subspace
	codes [][]uint8

	// vectorNodes stores the original VectorNode metadata
	vectorNodes []VectorNode

	// deletedNodes tracks soft-deleted IDs using roaring bitmap
	// CRITICAL OPTIMIZATION: RoaringBitmap is much more efficient than map[uint32]bool
	// - O(log n) membership test with better memory efficiency
	// - Compressed bitmap representation
	// - Fast iteration for batch operations
	deletedNodes *roaring.Bitmap

	// mu provides thread-safe access
	mu sync.RWMutex

	// trained indicates whether codebooks have been learned
	trained bool
}

// NewPQIndex creates a new Product Quantization index.
//
// Parameters:
//   - dim: Vector dimensionality (must be divisible by M)
//   - distanceKind: Distance metric
//   - M: Number of subquantizers (subspaces). Must divide dim evenly.
//   - Nbits: Bits per PQ code, determines K=2^Nbits centroids per subspace
//
// Returns:
//   - *PQIndex: New untrained PQ index
//   - error: Returns error if parameters invalid
//
// Tip: Use CalculatePQParams(dim) to get recommended M and Nbits values.
func NewPQIndex(dim int, distanceKind DistanceKind, M int, Nbits int) (*PQIndex, error) {
	// Validate dimension
	if dim <= 0 {
		return nil, fmt.Errorf("dimension must be positive")
	}

	// Validate M
	if M <= 0 {
		return nil, fmt.Errorf("parameter M must be positive")
	}

	// Critical: dimension must be evenly divisible by M
	if dim%M != 0 {
		return nil, fmt.Errorf("dimension %d must be divisible by M %d", dim, M)
	}

	// Validate Nbits
	if Nbits <= 0 || Nbits > 16 {
		return nil, fmt.Errorf("parameter Nbits must be in [1,16]")
	}

	// Create distance calculator
	distance, err := NewDistance(distanceKind)
	if err != nil {
		return nil, err
	}

	// Calculate derived parameters
	Ksub := 1 << Nbits // K = 2^Nbits
	dsub := dim / M

	return &PQIndex{
		dim:          dim,
		distanceKind: distanceKind,
		distance:     distance,
		M:            M,
		Nbits:        Nbits,
		Ksub:         Ksub,
		dsub:         dsub,
		codes:        make([][]uint8, 0),
		vectorNodes:  make([]VectorNode, 0),
		deletedNodes: roaring.New(), // Initialize empty bitmap for soft deletes
	}, nil
}

// Train learns codebooks for each subspace using k-means.
//
// Algorithm:
//  1. For each of M subspaces:
//     a. Extract that subspace from all training vectors
//     b. Run k-means to find K centroids
//     c. Store centroids as codebook for that subspace
//
// Parameters:
//   - vectors: Training vectors (need at least Ksub vectors)
//
// Returns:
//   - error: Returns error if insufficient training data
func (idx *PQIndex) Train(vectors []VectorNode) error {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	// Validate sufficient training vectors
	if len(vectors) < idx.Ksub {
		return fmt.Errorf("need at least %d vectors for training", idx.Ksub)
	}

	// Validate dimensionality
	for _, v := range vectors {
		if len(v.Vector()) != idx.dim {
			return fmt.Errorf("vector dimension mismatch: expected %d, got %d",
				idx.dim, len(v.Vector()))
		}
	}

	// Extract raw float32 slices for k-means
	rawVectors := make([][]float32, len(vectors))
	for i, v := range vectors {
		rawVectors[i] = v.Vector()
	}

	// Allocate codebooks
	idx.codebooks = make([][]float32, idx.M)

	// Train each subspace independently
	for m := 0; m < idx.M; m++ {
		// Extract subspace m from all vectors
		subVectors := make([][]float32, len(rawVectors))
		start := m * idx.dsub
		end := start + idx.dsub

		for i, v := range rawVectors {
			subVectors[i] = v[start:end]
		}

		// Run k-means on this subspace
		// Use our improved KMeansSubspace function
		centroids, _ := KMeansSubspace(subVectors, idx.Ksub, 20)

		if centroids == nil {
			return fmt.Errorf("k-means failed for subspace %d", m)
		}

		// Store centroids as flattened codebook
		idx.codebooks[m] = make([]float32, idx.Ksub*idx.dsub)
		for k := 0; k < idx.Ksub; k++ {
			copy(idx.codebooks[m][k*idx.dsub:(k+1)*idx.dsub], centroids[k])
		}
	}

	idx.trained = true
	return nil
}

// Add compresses and adds vectors to the index.
//
// Encoding process:
// 1. Divide vector into M subvectors
// 2. For each subvector, find nearest centroid in its codebook
// 3. Store centroid IDs as M-byte code
//
// Original vectors are discarded after encoding!
//
// Parameters:
//   - vector: Vector to compress and add
//
// Returns:
//   - error: Returns error if not trained or dimension mismatch
func (idx *PQIndex) Add(vector VectorNode) error {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	// Enforce training requirement
	if !idx.trained {
		return fmt.Errorf("index must be trained before adding")
	}

	// Validate dimensionality
	if len(vector.Vector()) != idx.dim {
		return fmt.Errorf("vector dimension mismatch: expected %d, got %d",
			idx.dim, len(vector.Vector()))
	}

	// Preprocess vector according to distance metric
	if err := idx.distance.PreprocessInPlace(vector.Vector()); err != nil {
		return err
	}

	// Encode vector into PQ code
	code := idx.encode(vector.Vector())

	// Store compressed code and metadata
	idx.codes = append(idx.codes, code)
	idx.vectorNodes = append(idx.vectorNodes, vector)

	return nil
}

// Remove performs soft delete using roaring bitmap.
//
// CONCURRENCY OPTIMIZATION:
// - Uses read lock first (cheaper) to check if node exists
// - Only acquires write lock for the actual bitmap modification
// - Minimizes write lock contention
//
// SOFT DELETE MECHANISM:
// Instead of immediately removing from the codes and vectorNodes slices (expensive O(n)),
// we mark as deleted in roaring bitmap. Deleted nodes are:
//   - Skipped during search
//   - Still in storage slices
//   - Not counted as active nodes
//
// Call Flush() periodically for actual cleanup and memory reclamation.
//
// Parameters:
//   - vector: Vector to remove (only the ID field is used for matching)
//
// Returns:
//   - error: Returns error if vector is not found or already deleted
//
// Time Complexity: O(n) for existence check + O(log n) for bitmap operation
//
// Thread-safety: Uses read lock for validation, write lock for modification
func (idx *PQIndex) Remove(vector VectorNode) error {
	id := vector.ID()

	// ════════════════════════════════════════════════════════════════════════
	// STEP 1: CHECK EXISTENCE (READ LOCK - CHEAPER)
	// ════════════════════════════════════════════════════════════════════════
	idx.mu.RLock()
	exists := false
	for _, v := range idx.vectorNodes {
		if v.ID() == id {
			exists = true
			break
		}
	}
	alreadyDeleted := idx.deletedNodes.Contains(id)
	idx.mu.RUnlock()

	// Fast-fail validation outside of write lock
	if !exists {
		return fmt.Errorf("vector with ID %d not found", id)
	}
	if alreadyDeleted {
		return fmt.Errorf("vector with ID %d already deleted", id)
	}

	// ════════════════════════════════════════════════════════════════════════
	// STEP 2: MARK AS DELETED (WRITE LOCK - ONLY FOR BITMAP UPDATE)
	// ════════════════════════════════════════════════════════════════════════
	idx.mu.Lock()
	idx.deletedNodes.Add(id)
	idx.mu.Unlock()

	return nil
}

// Flush performs hard delete of soft-deleted nodes.
//
// WHEN TO CALL:
//   - After multiple Remove() calls (batch cleanup)
//   - When deleted nodes are significant (e.g., > 10% of index)
//   - During off-peak hours
//
// WHAT IT DOES:
// 1. Removes all soft-deleted codes and vectorNodes from their parallel slices
// 2. Reclaims memory occupied by deleted vectors
// 3. Clears the deleted nodes bitmap
//
// COST: O(n) where n = number of vectors in the index
//
// Thread-safety: Acquires exclusive write lock
func (idx *PQIndex) Flush() error {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	// Quick exit if nothing to flush
	deletedCount := int(idx.deletedNodes.GetCardinality())
	if deletedCount == 0 {
		return nil
	}

	// ═══════════════════════════════════════════════════════════════════════
	// PHASE 1: FILTER OUT DELETED VECTORS (PARALLEL SLICES)
	// ═══════════════════════════════════════════════════════════════════════
	// Pre-allocate slices with capacity for non-deleted vectors
	activeCodes := make([][]uint8, 0, len(idx.codes)-deletedCount)
	activeVectorNodes := make([]VectorNode, 0, len(idx.vectorNodes)-deletedCount)

	// Filter both slices in parallel
	for i, v := range idx.vectorNodes {
		// Keep vector only if NOT deleted
		// RoaringBitmap Contains() is very fast - O(log n)
		if !idx.deletedNodes.Contains(v.ID()) {
			activeCodes = append(activeCodes, idx.codes[i])
			activeVectorNodes = append(activeVectorNodes, v)
		}
	}

	// Replace slices with filtered versions
	idx.codes = activeCodes
	idx.vectorNodes = activeVectorNodes

	// ═══════════════════════════════════════════════════════════════════════
	// PHASE 2: RESET DELETED TRACKING
	// ═══════════════════════════════════════════════════════════════════════
	idx.deletedNodes.Clear()

	return nil
}

// NewSearch creates a new search builder.
func (idx *PQIndex) NewSearch() VectorSearch {
	return &pqIndexSearch{
		index:  idx,
		k:      10,
		cutoff: -1, // Default no cutoff
	}
}

// Dimensions returns the dimensionality of original vectors.
func (idx *PQIndex) Dimensions() int {
	return idx.dim
}

// DistanceKind returns the distance metric.
func (idx *PQIndex) DistanceKind() DistanceKind {
	return idx.distanceKind
}

// Kind returns the index type.
func (idx *PQIndex) Kind() VectorIndexKind {
	return PQIndexKind
}

// Trained returns true if the index has been trained
func (idx *PQIndex) Trained() bool {
	return idx.trained
}

// encode converts a vector into a compact PQ code.
//
// Time Complexity: O(M × K × dsub)
func (idx *PQIndex) encode(v []float32) []uint8 {
	code := make([]uint8, idx.M)

	// Encode each subspace independently
	for m := 0; m < idx.M; m++ {
		// Extract subvector
		start := m * idx.dsub
		end := start + idx.dsub
		subVector := v[start:end]

		// Find nearest centroid
		minDist := float32(math.Inf(1))
		minIdx := 0

		for ksub := 0; ksub < idx.Ksub; ksub++ {
			centroid := idx.codebooks[m][ksub*idx.dsub : (ksub+1)*idx.dsub]

			// Use L2 squared for efficiency
			var dist float32
			for i := range subVector {
				diff := subVector[i] - centroid[i]
				dist += diff * diff
			}

			if dist < minDist {
				minDist = dist
				minIdx = ksub
			}
		}

		code[m] = uint8(minIdx)
	}

	return code
}

// WriteTo serializes the PQIndex to an io.Writer.
//
// IMPORTANT: This method calls Flush() before serialization to ensure all soft-deleted
// vectors are permanently removed from the serialized data.
//
// The serialization format is:
// 1. Magic number (4 bytes) - "PQIX" identifier for validation
// 2. Version (4 bytes) - Format version for backward compatibility
// 3. Basic parameters:
//   - Dimensionality (4 bytes)
//   - Distance kind length (4 bytes) + distance kind string
//   - M (4 bytes) - number of subquantizers
//   - Nbits (4 bytes) - bits per PQ code
//   - Ksub (4 bytes) - centroids per subquantizer
//   - dsub (4 bytes) - dimension of each subspace
//   - trained (1 byte) - whether codebooks have been learned
//
// 4. Codebooks (only if trained):
//   - For each of M subquantizers:
//   - Codebook size (4 bytes)
//   - Codebook data (Ksub * dsub * 4 bytes as float32)
//
// 5. Number of vectors (4 bytes)
// 6. For each vector:
//   - Vector ID (4 bytes)
//   - PQ code (M bytes)
//
// 7. Deleted nodes bitmap size (4 bytes) + roaring bitmap bytes
//
// Thread-safety: Acquires read lock during serialization
//
// Returns:
//   - int64: Number of bytes written
//   - error: Returns error if write fails or flush fails
func (idx *PQIndex) WriteTo(w io.Writer) (int64, error) {
	// Flush before serializing to remove soft-deleted vectors
	if err := idx.Flush(); err != nil {
		return 0, fmt.Errorf("failed to flush before serialization: %w", err)
	}

	idx.mu.RLock()
	defer idx.mu.RUnlock()

	var bytesWritten int64

	// Helper function to track writes
	write := func(data interface{}) error {
		err := binary.Write(w, binary.LittleEndian, data)
		if err == nil {
			switch v := data.(type) {
			case uint32, int32, float32:
				bytesWritten += 4
			case uint8, int8, bool:
				bytesWritten += 1
			case []byte:
				bytesWritten += int64(len(v))
			case []float32:
				bytesWritten += int64(len(v) * 4)
			}
		}
		return err
	}

	// 1. Write magic number "PQIX"
	magic := [4]byte{'P', 'Q', 'I', 'X'}
	if _, err := w.Write(magic[:]); err != nil {
		return bytesWritten, fmt.Errorf("failed to write magic number: %w", err)
	}
	bytesWritten += 4

	// 2. Write version
	version := uint32(1)
	if err := write(version); err != nil {
		return bytesWritten, fmt.Errorf("failed to write version: %w", err)
	}

	// 3. Write basic parameters
	if err := write(uint32(idx.dim)); err != nil {
		return bytesWritten, fmt.Errorf("failed to write dimensionality: %w", err)
	}

	// Write distance kind
	distanceKindBytes := []byte(idx.distanceKind)
	if err := write(uint32(len(distanceKindBytes))); err != nil {
		return bytesWritten, fmt.Errorf("failed to write distance kind length: %w", err)
	}
	if _, err := w.Write(distanceKindBytes); err != nil {
		return bytesWritten, fmt.Errorf("failed to write distance kind: %w", err)
	}
	bytesWritten += int64(len(distanceKindBytes))

	if err := write(uint32(idx.M)); err != nil {
		return bytesWritten, fmt.Errorf("failed to write M: %w", err)
	}

	if err := write(uint32(idx.Nbits)); err != nil {
		return bytesWritten, fmt.Errorf("failed to write Nbits: %w", err)
	}

	if err := write(uint32(idx.Ksub)); err != nil {
		return bytesWritten, fmt.Errorf("failed to write Ksub: %w", err)
	}

	if err := write(uint32(idx.dsub)); err != nil {
		return bytesWritten, fmt.Errorf("failed to write dsub: %w", err)
	}

	// Write trained flag
	trainedByte := uint8(0)
	if idx.trained {
		trainedByte = 1
	}
	if err := write(trainedByte); err != nil {
		return bytesWritten, fmt.Errorf("failed to write trained flag: %w", err)
	}

	// 4. Write codebooks (only if trained)
	if idx.trained {
		for m := 0; m < idx.M; m++ {
			// Write codebook size
			codebookSize := uint32(len(idx.codebooks[m]))
			if err := write(codebookSize); err != nil {
				return bytesWritten, fmt.Errorf("failed to write codebook %d size: %w", m, err)
			}

			// Write codebook data
			for _, val := range idx.codebooks[m] {
				if err := write(val); err != nil {
					return bytesWritten, fmt.Errorf("failed to write codebook %d data: %w", m, err)
				}
			}
		}
	}

	// 5. Write number of vectors
	if err := write(uint32(len(idx.vectorNodes))); err != nil {
		return bytesWritten, fmt.Errorf("failed to write vector count: %w", err)
	}

	// 6. Write each vector (ID + PQ code)
	for i, node := range idx.vectorNodes {
		// Write vector ID
		if err := write(node.ID()); err != nil {
			return bytesWritten, fmt.Errorf("failed to write vector %d ID: %w", i, err)
		}

		// Write PQ code
		if _, err := w.Write(idx.codes[i]); err != nil {
			return bytesWritten, fmt.Errorf("failed to write vector %d code: %w", i, err)
		}
		bytesWritten += int64(len(idx.codes[i]))
	}

	// 7. Write deleted nodes bitmap
	bitmapBytes, err := idx.deletedNodes.ToBytes()
	if err != nil {
		return bytesWritten, fmt.Errorf("failed to serialize deleted nodes bitmap: %w", err)
	}
	if err := write(uint32(len(bitmapBytes))); err != nil {
		return bytesWritten, fmt.Errorf("failed to write bitmap size: %w", err)
	}
	if _, err := w.Write(bitmapBytes); err != nil {
		return bytesWritten, fmt.Errorf("failed to write bitmap data: %w", err)
	}
	bytesWritten += int64(len(bitmapBytes))

	return bytesWritten, nil
}

// ReadFrom deserializes a PQIndex from an io.Reader.
//
// This method reconstructs a PQIndex from the serialized format created by WriteTo.
// The deserialized index is fully functional and ready to use for searches.
//
// Thread-safety: Acquires write lock during deserialization
//
// Returns:
//   - int64: Number of bytes read
//   - error: Returns error if read fails, format is invalid, or data is corrupted
//
// Example:
//
//	// Save index
//	file, _ := os.Create("index.bin")
//	idx.WriteTo(file)
//	file.Close()
//
//	// Load index
//	file, _ := os.Open("index.bin")
//	M, Nbits := CalculatePQParams(384)
//	idx2, _ := NewPQIndex(384, Cosine, M, Nbits)
//	idx2.ReadFrom(file)
//	file.Close()
func (idx *PQIndex) ReadFrom(r io.Reader) (int64, error) {
	idx.mu.Lock()
	defer idx.mu.Unlock()

	var bytesRead int64

	// Helper function to track reads
	read := func(data interface{}) error {
		err := binary.Read(r, binary.LittleEndian, data)
		if err == nil {
			switch data.(type) {
			case *uint32, *int32, *float32:
				bytesRead += 4
			case *uint8, *int8, *bool:
				bytesRead += 1
			}
		}
		return err
	}

	// 1. Read and validate magic number
	magic := make([]byte, 4)
	if _, err := io.ReadFull(r, magic); err != nil {
		return bytesRead, fmt.Errorf("failed to read magic number: %w", err)
	}
	bytesRead += 4
	if string(magic) != "PQIX" {
		return bytesRead, fmt.Errorf("invalid magic number: expected 'PQIX', got '%s'", string(magic))
	}

	// 2. Read version
	var version uint32
	if err := read(&version); err != nil {
		return bytesRead, fmt.Errorf("failed to read version: %w", err)
	}
	if version != 1 {
		return bytesRead, fmt.Errorf("unsupported version: %d", version)
	}

	// 3. Read basic parameters
	var dim uint32
	if err := read(&dim); err != nil {
		return bytesRead, fmt.Errorf("failed to read dimensionality: %w", err)
	}

	// Validate dimension matches
	if int(dim) != idx.dim {
		return bytesRead, fmt.Errorf("dimension mismatch: index has dim=%d, serialized data has dim=%d", idx.dim, dim)
	}

	// Read distance kind
	var distanceKindLen uint32
	if err := read(&distanceKindLen); err != nil {
		return bytesRead, fmt.Errorf("failed to read distance kind length: %w", err)
	}

	distanceKindBytes := make([]byte, distanceKindLen)
	if _, err := io.ReadFull(r, distanceKindBytes); err != nil {
		return bytesRead, fmt.Errorf("failed to read distance kind: %w", err)
	}
	bytesRead += int64(distanceKindLen)

	distanceKind := DistanceKind(distanceKindBytes)
	if distanceKind != idx.distanceKind {
		return bytesRead, fmt.Errorf("distance kind mismatch: index uses '%s', serialized data uses '%s'", idx.distanceKind, distanceKind)
	}

	// Read M, Nbits, Ksub, dsub
	var M, Nbits, Ksub, dsub uint32
	if err := read(&M); err != nil {
		return bytesRead, fmt.Errorf("failed to read M: %w", err)
	}
	if err := read(&Nbits); err != nil {
		return bytesRead, fmt.Errorf("failed to read Nbits: %w", err)
	}
	if err := read(&Ksub); err != nil {
		return bytesRead, fmt.Errorf("failed to read Ksub: %w", err)
	}
	if err := read(&dsub); err != nil {
		return bytesRead, fmt.Errorf("failed to read dsub: %w", err)
	}

	// Validate parameters match
	if int(M) != idx.M {
		return bytesRead, fmt.Errorf("parameter M mismatch: index has M=%d, serialized data has M=%d", idx.M, M)
	}
	if int(Nbits) != idx.Nbits {
		return bytesRead, fmt.Errorf("parameter Nbits mismatch: index has Nbits=%d, serialized data has Nbits=%d", idx.Nbits, Nbits)
	}
	if int(Ksub) != idx.Ksub {
		return bytesRead, fmt.Errorf("parameter Ksub mismatch: index has Ksub=%d, serialized data has Ksub=%d", idx.Ksub, Ksub)
	}
	if int(dsub) != idx.dsub {
		return bytesRead, fmt.Errorf("parameter dsub mismatch: index has dsub=%d, serialized data has dsub=%d", idx.dsub, dsub)
	}

	// Read trained flag
	var trainedByte uint8
	if err := read(&trainedByte); err != nil {
		return bytesRead, fmt.Errorf("failed to read trained flag: %w", err)
	}
	trained := trainedByte == 1

	// 4. Read codebooks (only if trained)
	var codebooks [][]float32
	if trained {
		codebooks = make([][]float32, idx.M)
		for m := 0; m < idx.M; m++ {
			// Read codebook size
			var codebookSize uint32
			if err := read(&codebookSize); err != nil {
				return bytesRead, fmt.Errorf("failed to read codebook %d size: %w", m, err)
			}

			// Read codebook data
			codebooks[m] = make([]float32, codebookSize)
			for j := uint32(0); j < codebookSize; j++ {
				if err := read(&codebooks[m][j]); err != nil {
					return bytesRead, fmt.Errorf("failed to read codebook %d data: %w", m, err)
				}
			}
		}
	}

	// 5. Read number of vectors
	var vectorCount uint32
	if err := read(&vectorCount); err != nil {
		return bytesRead, fmt.Errorf("failed to read vector count: %w", err)
	}

	// 6. Read vectors (ID + PQ code)
	vectorNodes := make([]VectorNode, vectorCount)
	codes := make([][]uint8, vectorCount)

	for i := uint32(0); i < vectorCount; i++ {
		// Read vector ID
		var id uint32
		if err := read(&id); err != nil {
			return bytesRead, fmt.Errorf("failed to read vector %d ID: %w", i, err)
		}

		// Read PQ code
		code := make([]uint8, idx.M)
		if _, err := io.ReadFull(r, code); err != nil {
			return bytesRead, fmt.Errorf("failed to read vector %d code: %w", i, err)
		}
		bytesRead += int64(idx.M)

		// Create VectorNode with empty vector (PQ doesn't store original vectors)
		vectorNodes[i] = *NewVectorNodeWithID(id, nil)
		codes[i] = code
	}

	// 7. Read deleted nodes bitmap
	var bitmapSize uint32
	if err := read(&bitmapSize); err != nil {
		return bytesRead, fmt.Errorf("failed to read bitmap size: %w", err)
	}

	bitmapBytes := make([]byte, bitmapSize)
	if _, err := io.ReadFull(r, bitmapBytes); err != nil {
		return bytesRead, fmt.Errorf("failed to read bitmap data: %w", err)
	}
	bytesRead += int64(bitmapSize)

	deletedNodes := roaring.New()
	if err := deletedNodes.UnmarshalBinary(bitmapBytes); err != nil {
		return bytesRead, fmt.Errorf("failed to deserialize deleted nodes bitmap: %w", err)
	}

	// Update index state
	idx.trained = trained
	idx.codebooks = codebooks
	idx.vectorNodes = vectorNodes
	idx.codes = codes
	idx.deletedNodes = deletedNodes

	return bytesRead, nil
}