//! HNSW Index Product Quantization integration
//!
//! Contains PQ-related methods for vector compression and accelerated search.

use std::collections::{BinaryHeap, HashSet};
use std::sync::atomic::Ordering as AtomicOrdering;

use ordered_float::OrderedFloat;

use crate::vector::distance::Distance;

use super::index::HNSWIndex;
use super::types::{
    Neighbor, MaxNeighbor, SearchResult, MemoryStats, HNSWError,
};

impl HNSWIndex {
    // ============================================================================
    // P1-1: Product Quantization Methods
    // ============================================================================
    
    /// Train the PQ quantizer on current vectors
    /// 
    /// This should be called after inserting a representative sample of vectors.
    /// After training, new vectors will be compressed using PQ codes.
    /// 
    /// Optimized: uses vector references where possible
    pub fn train_pq(&self) -> Result<(), HNSWError> {
        let mut quantizer_guard = self.quantizer.write();
        let quantizer = quantizer_guard.as_mut()
            .ok_or(HNSWError::PQNotEnabled)?;
        
        if quantizer.is_trained() {
            return Err(HNSWError::PQAlreadyTrained);
        }
        
        // Collect vectors for training
        let nodes = self.nodes.read();
        let deleted = self.deleted.read();
        
        // Clone vectors for training (necessary as PQ training takes ownership)
        let vectors: Vec<Vec<f32>> = nodes.iter()
            .enumerate()
            .filter(|(id, _)| !deleted.contains(id))
            .map(|(_, node)| (*node.vector).clone())
            .collect();
        
        drop(deleted);
        drop(nodes);
        
        if vectors.is_empty() {
            return Err(HNSWError::PQTrainingError("No vectors to train on".to_string()));
        }
        
        // Train the quantizer
        quantizer.train(&vectors)
            .map_err(|e| HNSWError::PQTrainingError(e.to_string()))?;
        
        // Encode all existing vectors
        let mut pq_codes = self.pq_codes.write();
        pq_codes.clear();
        
        let nodes = self.nodes.read();
        for node in nodes.iter() {
            let code = quantizer.encode(node.vector())
                .map_err(|e| HNSWError::PQTrainingError(e.to_string()))?;
            pq_codes.push(code);
        }
        
        Ok(())
    }
    
    /// Check if PQ is enabled
    pub fn pq_enabled(&self) -> bool {
        self.quantizer.read().is_some()
    }
    
    /// Check if PQ is trained
    pub fn pq_trained(&self) -> bool {
        self.quantizer.read()
            .as_ref()
            .map(|q| q.is_trained())
            .unwrap_or(false)
    }
    
    /// Get PQ compression ratio
    pub fn pq_compression_ratio(&self) -> Option<f32> {
        self.quantizer.read()
            .as_ref()
            .map(|q| q.compression_ratio())
    }
    
    /// Get memory usage statistics
    pub fn memory_stats(&self) -> MemoryStats {
        let nodes = self.nodes.read();
        let pq_codes = self.pq_codes.read();
        let quantizer = self.quantizer.read();
        
        let node_count = nodes.len();
        let dimension = self.config.dimension;
        
        // Original vector memory (if stored)
        let original_vectors_bytes = if self.store_original {
            node_count * dimension * std::mem::size_of::<f32>()
        } else {
            0
        };
        
        // PQ codes memory
        let pq_codes_bytes: usize = pq_codes.iter().map(|c| c.len()).sum();
        
        // Compression ratio
        let compression_ratio = if pq_codes_bytes > 0 && !self.store_original {
            (node_count * dimension * std::mem::size_of::<f32>()) as f32 / pq_codes_bytes as f32
        } else {
            1.0
        };
        
        MemoryStats {
            node_count,
            original_vectors_bytes,
            pq_codes_bytes,
            compression_ratio,
            pq_trained: quantizer.as_ref().map(|q| q.is_trained()).unwrap_or(false),
        }
    }
    
    /// Search using PQ for fast approximate distance computation
    /// 
    /// If PQ is trained, this will:
    /// 1. Use PQ distance tables for fast initial search
    /// 2. Optionally rerank top candidates using original vectors (if stored)
    pub fn search_pq(&self, query: &[f32], k: usize) -> Result<Vec<SearchResult>, HNSWError> {
        self.search_pq_with_rerank(query, k, 0)
    }
    
    /// Search using PQ with configurable reranking
    /// 
    /// - `rerank_k`: Number of candidates to rerank with original vectors (0 = no reranking)
    pub fn search_pq_with_rerank(&self, query: &[f32], k: usize, rerank_k: usize) -> Result<Vec<SearchResult>, HNSWError> {
        if query.len() != self.config.dimension {
            return Err(HNSWError::DimensionMismatch {
                expected: self.config.dimension,
                got: query.len(),
            });
        }
        
        // Fall back to regular search if PQ not trained
        if !self.pq_trained() {
            return self.search(query, k);
        }
        
        let entry = self.entry_point.read();
        let entry_id = match *entry {
            Some(id) => id,
            None => return Ok(Vec::new()),
        };
        drop(entry);
        
        let max_level = self.max_level.load(AtomicOrdering::SeqCst);
        
        // Precompute PQ distance table for fast distance computation
        let distance_table = {
            let quantizer = self.quantizer.read();
            quantizer.as_ref()
                .ok_or(HNSWError::PQNotEnabled)?
                .compute_distance_table(query)
                .map_err(|e| HNSWError::PQTrainingError(e.to_string()))?
        };
        
        // Start from entry point - use regular search for upper levels
        let mut curr_id = entry_id;
        
        // Traverse from top level to level 1 using regular greedy search
        for level in (1..=max_level).rev() {
            curr_id = self.greedy_search_level(query, curr_id, level);
        }
        
        // Search at level 0 with PQ-accelerated distance computation
        let deleted = self.deleted.read();
        let deleted_count = deleted.len();
        let ef_adjusted = (self.config.ef_search.max(k) + deleted_count).min(self.nodes.read().len());
        drop(deleted);
        
        // Determine how many candidates to retrieve for reranking
        let search_k = if rerank_k > 0 && self.store_original {
            rerank_k.max(k)
        } else {
            ef_adjusted
        };
        
        let neighbors = self.search_level_pq(query, curr_id, search_k, 0, &distance_table);
        
        // Filter out deleted nodes
        let deleted = self.deleted.read();
        let mut results: Vec<_> = neighbors.into_iter()
            .filter(|n| !deleted.contains(&n.id))
            .collect();
        drop(deleted);
        
        // Rerank with original vectors if requested and available
        if rerank_k > 0 && self.store_original {
            let nodes = self.nodes.read();
            for result in &mut results {
                // Compute exact distance using original vector (no clone!)
                let exact_dist = self.distance.calculate(query, nodes[result.id].vector());
                result.distance = OrderedFloat(exact_dist);
            }
        }
        
        // Always sort results by distance
        results.sort_by(|a, b| a.distance.cmp(&b.distance));
        results.truncate(k);
        
        Ok(results.into_iter()
            .map(|n| SearchResult { id: n.id, distance: n.distance.0 })
            .collect())
    }
    
    /// Search at a specific level using PQ distance table for fast computation
    /// 
    /// Optimized: uses node.vector() reference instead of cloning
    pub(crate) fn search_level_pq(
        &self, 
        query: &[f32], 
        entry_id: usize, 
        ef: usize, 
        level: usize,
        distance_table: &[Vec<f32>],
    ) -> Vec<Neighbor> {
        let nodes = self.nodes.read();
        let pq_codes = self.pq_codes.read();
        let quantizer = self.quantizer.read();
        
        let mut visited = HashSet::new();
        let mut candidates = BinaryHeap::new(); // min-heap
        let mut results = BinaryHeap::new(); // max-heap (using MaxNeighbor)
        
        // Calculate entry distance - use PQ if available, otherwise exact
        let entry_dist = if let Some(ref q) = *quantizer {
            if entry_id < pq_codes.len() && !pq_codes[entry_id].is_empty() {
                q.distance_with_table(distance_table, &pq_codes[entry_id])
            } else {
                self.distance.calculate(query, nodes[entry_id].vector())
            }
        } else {
            self.distance.calculate(query, nodes[entry_id].vector())
        };
        
        visited.insert(entry_id);
        candidates.push(Neighbor { id: entry_id, distance: OrderedFloat(entry_dist) });
        results.push(MaxNeighbor { id: entry_id, distance: OrderedFloat(entry_dist) });
        
        while let Some(curr) = candidates.pop() {
            // Get furthest result distance
            let furthest_dist = results.peek().map(|n| n.distance.0).unwrap_or(f32::MAX);
            
            // Stop if current candidate is further than furthest result
            if curr.distance.0 > furthest_dist {
                break;
            }
            
            // Explore neighbors
            for &neighbor_id in &nodes[curr.id].neighbors[level] {
                if visited.contains(&neighbor_id) {
                    continue;
                }
                visited.insert(neighbor_id);
                
                // Use PQ distance if available, otherwise exact distance
                let neighbor_dist = if let Some(ref q) = *quantizer {
                    if neighbor_id < pq_codes.len() && !pq_codes[neighbor_id].is_empty() {
                        q.distance_with_table(distance_table, &pq_codes[neighbor_id])
                    } else {
                        self.distance.calculate(query, nodes[neighbor_id].vector())
                    }
                } else {
                    self.distance.calculate(query, nodes[neighbor_id].vector())
                };
                
                let furthest_dist = results.peek().map(|n| n.distance.0).unwrap_or(f32::MAX);
                
                if neighbor_dist < furthest_dist || results.len() < ef {
                    candidates.push(Neighbor { id: neighbor_id, distance: OrderedFloat(neighbor_dist) });
                    results.push(MaxNeighbor { id: neighbor_id, distance: OrderedFloat(neighbor_dist) });
                    
                    if results.len() > ef {
                        results.pop();
                    }
                }
            }
        }
        
        // Convert results to Neighbor vec
        results.into_iter()
            .map(|n| Neighbor { id: n.id, distance: n.distance })
            .collect()
    }
    
    /// Hybrid search: PQ for initial candidates, exact distance for reranking
    /// 
    /// This provides a good balance between speed and accuracy:
    /// - Uses PQ distance for fast initial filtering
    /// - Reranks top candidates with exact distance
    pub fn search_hybrid(&self, query: &[f32], k: usize, rerank_multiplier: usize) -> Result<Vec<SearchResult>, HNSWError> {
        let rerank_k = k * rerank_multiplier.max(1);
        self.search_pq_with_rerank(query, k, rerank_k)
    }
}
