src/components/retrievers/rerankers/neural_reranker.py

"""
Enhanced Neural Reranker for Advanced Retrieval.

This module provides a sophisticated neural reranking sub-component that extends
the existing reranker capabilities with advanced features including multiple 
model support, adaptive strategies, score fusion, and performance optimization.

This is the architecture-compliant implementation in the proper rerankers/ 
sub-component location, enhanced with capabilities from the migrated utilities.
"""

import logging
import time
from typing import List, Dict, Any, Tuple, Optional, Union
import numpy as np

from src.core.interfaces import Document
from .base import Reranker
from .utils import (
    ScoreFusion, AdaptiveStrategies, CrossEncoderModels, PerformanceOptimizer,
    ModelConfig, WeightsConfig, NormalizationConfig
)

logger = logging.getLogger(__name__)


class NeuralRerankingError(Exception):
    """Raised when neural reranking operations fail."""
    pass


class NeuralReranker(Reranker):
    """
    Enhanced neural reranker with sophisticated capabilities.
    
    This reranker extends the base Reranker interface with advanced features:
    - Multiple cross-encoder model support with lazy loading
    - Query-type adaptive reranking strategies  
    - Advanced neural + retrieval score fusion
    - Performance optimization with caching and batching
    - Graceful degradation and comprehensive error handling
    - Real-time performance monitoring and adaptation
    
    The implementation follows proper architecture patterns by enhancing
    the existing rerankers sub-component with utilities from the migrated
    reranking/ module capabilities.
    
    Features:
    - ✅ Multiple model support (general + technical domains)
    - ✅ Adaptive model selection based on query type
    - ✅ Advanced score fusion with normalization strategies
    - ✅ Performance optimization (<200ms target)
    - ✅ Intelligent caching with LRU eviction
    - ✅ Batch processing with dynamic sizing
    - ✅ Comprehensive error handling and fallbacks
    - ✅ Real-time metrics collection and adaptation
    
    Example:
        config = {
            "enabled": True,
            "models": {
                "default_model": {
                    "name": "cross-encoder/ms-marco-MiniLM-L6-v2",
                    "max_length": 512,
                    "batch_size": 16
                },
                "technical_model": {
                    "name": "cross-encoder/ms-marco-electra-base",
                    "max_length": 512,
                    "batch_size": 8
                }
            },
            "adaptive": {
                "enabled": True,
                "confidence_threshold": 0.7
            },
            "score_fusion": {
                "method": "weighted",
                "weights": {
                    "neural_score": 0.7,
                    "retrieval_score": 0.3
                }
            },
            "performance": {
                "max_latency_ms": 200,
                "enable_caching": True,
                "max_cache_size": 10000
            }
        }
        reranker = NeuralReranker(config)
        results = reranker.rerank(query, documents, initial_scores)
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        Initialize enhanced neural reranker.
        
        Args:
            config: Neural reranking configuration dictionary
        """
        self.config = config
        self.enabled = config.get("enabled", True)
        
        # Parse configuration sections
        self._parse_configuration(config)
        
        # Initialize advanced components
        self.models_manager = None
        self.adaptive_strategies = None
        self.score_fusion = None
        self.performance_optimizer = None
        
        # Performance tracking
        self.stats = {
            "total_queries": 0,
            "successful_queries": 0, 
            "failed_queries": 0,
            "total_latency_ms": 0.0,
            "model_switches": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "fallback_activations": 0,
            "adaptive_adjustments": 0
        }
        
        # State management
        self._initialized = False
        self._error_count = 0
        self._last_performance_check = time.time()
        
        # Initialize immediately if enabled (remove lazy initialization)
        initialize_immediately = config.get("initialize_immediately", True)
        if self.enabled and initialize_immediately:
            try:
                self._initialize_if_needed()
            except Exception as e:
                logger.warning(f"Failed to initialize neural reranker: {e}")
                logger.warning("Disabling neural reranker and falling back to identity mode")
                self.enabled = False
                self._initialized = True  # Mark as initialized even when disabled
        else:
            self._initialized = True  # Mark as initialized when disabled
        
        logger.info(f"Enhanced NeuralReranker initialized with {len(self.models_config)} models, "
                   f"enabled={self.enabled}, initialized={self._initialized}")
    
    def _parse_configuration(self, config: Dict[str, Any]):
        """Parse and validate configuration sections."""
        # Models configuration
        self.models_config = config.get("models", {
            "default_model": {
                "name": "cross-encoder/ms-marco-MiniLM-L6-v2",
                "max_length": 512,
                "batch_size": 16
            }
        })
        
        # Convert to ModelConfig objects
        self.model_configs = {}
        for name, model_config in self.models_config.items():
            self.model_configs[name] = ModelConfig(**model_config)
        
        # Adaptive configuration
        adaptive_config = config.get("adaptive", {})
        self.adaptive_enabled = adaptive_config.get("enabled", True)
        self.confidence_threshold = adaptive_config.get("confidence_threshold", 0.7)
        
        # Score fusion configuration
        fusion_config = config.get("score_fusion", {})
        self.fusion_method = fusion_config.get("method", "weighted")
        
        weights_config = fusion_config.get("weights", {})
        self.weights = WeightsConfig(
            retrieval_score=weights_config.get("retrieval_score", 0.3),
            neural_score=weights_config.get("neural_score", 0.7),
            graph_score=weights_config.get("graph_score", 0.0),
            temporal_score=weights_config.get("temporal_score", 0.0)
        )
        
        normalization_config = fusion_config.get("normalization", {})
        self.normalization = NormalizationConfig(
            method=normalization_config.get("method", "min_max"),
            clip_outliers=normalization_config.get("clip_outliers", True),
            outlier_threshold=normalization_config.get("outlier_threshold", 3.0)
        )
        
        # Performance configuration
        perf_config = config.get("performance", {})
        self.max_latency_ms = perf_config.get("max_latency_ms", 200)
        self.target_latency_ms = perf_config.get("target_latency_ms", 150)
        self.enable_caching = perf_config.get("enable_caching", True)
        self.max_cache_size = perf_config.get("max_cache_size", 10000)
        self.cache_ttl_seconds = perf_config.get("cache_ttl_seconds", 3600)
        
        # Legacy compatibility
        self.max_candidates = config.get("max_candidates", 50)
        self.default_model = config.get("default_model", list(self.models_config.keys())[0])
    
    def _initialize_if_needed(self) -> None:
        """Initialize advanced components lazily for better startup performance."""
        if self._initialized or not self.enabled:
            return
        
        try:
            # Initialize models manager
            self.models_manager = CrossEncoderModels(self.model_configs)
            
            # Initialize adaptive strategies
            if self.adaptive_enabled:
                self.adaptive_strategies = AdaptiveStrategies(
                    enabled=True,
                    confidence_threshold=self.confidence_threshold
                )
            
            # Initialize score fusion
            self.score_fusion = ScoreFusion(
                method=self.fusion_method,
                weights=self.weights,
                normalization=self.normalization
            )
            
            # Initialize performance optimizer
            self.performance_optimizer = PerformanceOptimizer(
                max_latency_ms=self.max_latency_ms,
                target_latency_ms=self.target_latency_ms,
                enable_caching=self.enable_caching,
                cache_ttl_seconds=self.cache_ttl_seconds,
                max_cache_size=self.max_cache_size
            )
            
            self._initialized = True
            logger.info("Enhanced NeuralReranker initialization completed")
            
        except Exception as e:
            logger.error(f"Failed to initialize enhanced neural reranker: {e}")
            self.enabled = False
            raise NeuralRerankingError(f"Initialization failed: {e}")
    
    def rerank(
        self, 
        query: str, 
        documents: List[Document], 
        initial_scores: List[float]
    ) -> List[Tuple[int, float]]:
        """
        Rerank documents using enhanced neural models with advanced strategies.
        
        Args:
            query: The search query
            documents: List of candidate documents
            initial_scores: Initial relevance scores from retrieval
            
        Returns:
            List of (document_index, reranked_score) tuples sorted by score
        """
        start_time = time.time()
        self.stats["total_queries"] += 1
        
        try:
            # Check if reranking is enabled
            if not self.enabled:
                return [(i, score) for i, score in enumerate(initial_scores)]
            
            # Validate inputs
            if not documents or not query.strip():
                return []
            
            if len(initial_scores) != len(documents):
                logger.warning("Mismatch between documents and scores, using defaults")
                initial_scores = [1.0] * len(documents)
            
            # Initialize if needed
            self._initialize_if_needed()
            
            if not self._initialized:
                logger.warning("Enhanced neural reranker not initialized, using initial scores")
                return [(i, score) for i, score in enumerate(initial_scores)]
            
            # Check cache first
            cached_scores = self.performance_optimizer.get_cached_scores(
                query, documents, self.default_model
            )
            
            if cached_scores is not None:
                self.stats["cache_hits"] += 1
                return [(i, score) for i, score in enumerate(cached_scores)]
            
            self.stats["cache_misses"] += 1
            
            # Limit candidates for performance
            max_candidates = min(len(documents), self.max_candidates)
            if len(documents) > max_candidates:
                # Sort by initial scores and take top candidates
                sorted_indices = sorted(range(len(initial_scores)), 
                                      key=lambda i: initial_scores[i], reverse=True)
                top_indices = sorted_indices[:max_candidates]
                documents = [documents[i] for i in top_indices]
                initial_scores = [initial_scores[i] for i in top_indices]
            else:
                top_indices = list(range(len(documents)))
            
            # Select optimal model for this query
            selected_model = self._select_model_for_query(query)
            
            # Get neural scores
            neural_scores = self._get_neural_scores(query, documents, selected_model)
            
            # Fuse scores using advanced fusion strategies
            fused_scores = self.score_fusion.fuse_scores(
                initial_scores, neural_scores, query, documents
            )
            
            # Cache the results
            self.performance_optimizer.cache_scores(
                query, documents, selected_model, fused_scores
            )
            
            # Create final results with original indices
            results = []
            for i, score in enumerate(fused_scores):
                original_idx = top_indices[i] if len(documents) <= max_candidates else top_indices[i]
                results.append((original_idx, score))
            
            results.sort(key=lambda x: x[1], reverse=True)
            
            # Update performance statistics
            latency_ms = (time.time() - start_time) * 1000
            self._update_stats(latency_ms, success=True)
            self.performance_optimizer.record_latency(latency_ms)
            
            # Record performance for adaptive learning
            if self.adaptive_strategies:
                query_analysis = self.adaptive_strategies.detector.classify_query(query) if self.adaptive_strategies.detector else None
                query_type = query_analysis.query_type if query_analysis else "general"
                quality_score = self._estimate_quality_score(fused_scores, initial_scores)
                
                self.adaptive_strategies.record_performance(
                    selected_model, query_type, latency_ms, quality_score
                )
            
            # Log performance
            logger.debug(f"Enhanced neural reranking completed in {latency_ms:.1f}ms for {len(documents)} documents")
            
            return results
            
        except Exception as e:
            # Handle errors gracefully
            latency_ms = (time.time() - start_time) * 1000
            self._update_stats(latency_ms, success=False)
            self.stats["fallback_activations"] += 1
            logger.error(f"Enhanced neural reranking failed: {e}")
            
            # Return fallback results
            return [(i, score) for i, score in enumerate(initial_scores)]
    
    def _select_model_for_query(self, query: str) -> str:
        """
        Select the optimal model for the given query using adaptive strategies.
        
        Args:
            query: The search query
            
        Returns:
            Name of the selected model
        """
        if not self.adaptive_strategies or not self.adaptive_strategies.enabled:
            return self.default_model
        
        try:
            available_models = self.models_manager.get_available_models()
            selected_model = self.adaptive_strategies.select_model(
                query, available_models, self.default_model
            )
            
            if selected_model != self.default_model:
                self.stats["model_switches"] += 1
                self.stats["adaptive_adjustments"] += 1
            
            return selected_model
            
        except Exception as e:
            logger.warning(f"Adaptive model selection failed: {e}, using default")
            return self.default_model
    
    def _get_neural_scores(
        self, 
        query: str, 
        documents: List[Document], 
        model_name: str
    ) -> List[float]:
        """
        Get neural relevance scores for query-document pairs.
        
        Args:
            query: The search query
            documents: List of documents
            model_name: Name of the model to use
            
        Returns:
            List of neural relevance scores
        """
        try:
            # Prepare query-document pairs
            query_doc_pairs = []
            for doc in documents:
                doc_text = doc.content
                
                # Smart truncation for model max_length
                model_config = self.model_configs.get(model_name)
                if model_config:
                    max_length = model_config.max_length
                    if len(doc_text) > max_length:
                        # Try to keep complete sentences
                        truncated = doc_text[:max_length - 50]
                        last_period = truncated.rfind('.')
                        if last_period > max_length // 2:
                            doc_text = truncated[:last_period + 1]
                        else:
                            doc_text = truncated + "..."
                
                query_doc_pairs.append([query, doc_text])
            
            # Get scores using models manager
            neural_scores = self.models_manager.predict(query_doc_pairs, model_name)
            
            return neural_scores
            
        except Exception as e:
            logger.error(f"Neural scoring failed: {e}")
            return [0.0] * len(documents)
    
    def _estimate_quality_score(
        self, 
        fused_scores: List[float], 
        initial_scores: List[float]
    ) -> float:
        """
        Estimate quality improvement from neural reranking.
        
        Args:
            fused_scores: Final fused scores
            initial_scores: Initial retrieval scores
            
        Returns:
            Quality score (0-1)
        """
        try:
            # Simple heuristic: how much did the score distribution improve?
            if not fused_scores or not initial_scores:
                return 0.5
            
            # Calculate variance of scores (higher variance = better discrimination)
            fused_variance = np.var(fused_scores)
            initial_variance = np.var(initial_scores)
            
            # Calculate improvement ratio
            if initial_variance > 0:
                improvement_ratio = fused_variance / initial_variance
                # Normalize to 0-1 range
                quality_score = min(1.0, improvement_ratio / 2.0)
            else:
                quality_score = 0.5
            
            return quality_score
            
        except Exception:
            return 0.5
    
    def _update_stats(self, latency_ms: float, success: bool) -> None:
        """Update performance statistics."""
        if success:
            self.stats["successful_queries"] += 1
        else:
            self.stats["failed_queries"] += 1
            self._error_count += 1
        
        self.stats["total_latency_ms"] += latency_ms
    
    def is_enabled(self) -> bool:
        """
        Check if neural reranking is enabled.
        
        Returns:
            True if reranking should be performed
        """
        # Return True if configured to be enabled, regardless of initialization status
        # Initialization happens lazily when rerank() is called
        return self.enabled
    
    def get_reranker_info(self) -> Dict[str, Any]:
        """
        Get information about the enhanced neural reranker.
        
        Returns:
            Dictionary with reranker configuration and statistics
        """
        base_info = {
            "type": "enhanced_neural_reranker",
            "enabled": self.enabled,
            "initialized": self._initialized,
            "default_model": self.default_model,
            "total_models": len(self.models_config),
            "adaptive_enabled": self.adaptive_enabled,
            "score_fusion_method": self.fusion_method,
            "max_latency_ms": self.max_latency_ms,
            "target_latency_ms": self.target_latency_ms,
            "caching_enabled": self.enable_caching
        }
        
        # Add model information
        if self._initialized and self.models_manager:
            base_info["models"] = {
                name: config.name 
                for name, config in self.model_configs.items()
            }
            base_info["model_stats"] = self.models_manager.get_model_stats()
        
        # Add statistics
        base_info["statistics"] = self.stats.copy()
        
        # Add component statistics
        if self._initialized:
            if self.adaptive_strategies:
                base_info["adaptive_stats"] = self.adaptive_strategies.get_stats()
            
            if self.score_fusion:
                base_info["fusion_stats"] = self.score_fusion.get_stats()
            
            if self.performance_optimizer:
                base_info["performance_stats"] = self.performance_optimizer.get_stats()
        
        # Add performance metrics
        if self.stats["total_queries"] > 0:
            base_info["avg_latency_ms"] = self.stats["total_latency_ms"] / self.stats["total_queries"]
            base_info["success_rate"] = self.stats["successful_queries"] / self.stats["total_queries"]
            cache_total = self.stats["cache_hits"] + self.stats["cache_misses"]
            if cache_total > 0:
                base_info["cache_hit_rate"] = self.stats["cache_hits"] / cache_total
        
        return base_info
    
    def reset_stats(self) -> None:
        """Reset all statistics."""
        self.stats = {
            "total_queries": 0,
            "successful_queries": 0, 
            "failed_queries": 0,
            "total_latency_ms": 0.0,
            "model_switches": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "fallback_activations": 0,
            "adaptive_adjustments": 0
        }
        
        if self._initialized:
            if self.adaptive_strategies:
                self.adaptive_strategies.reset_stats()
            if self.score_fusion:
                self.score_fusion.reset_stats()
            if self.performance_optimizer:
                self.performance_optimizer.reset_stats()