"""
Neural Reranking Engine for Toxo
Advanced response reranking using neural models and multi-criteria optimization.
"""

import asyncio
import logging
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Any, Optional, Tuple, Union
from dataclasses import dataclass, asdict
import json
from datetime import datetime
from pathlib import Path

# ML libraries
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.preprocessing import StandardScaler
import joblib

# Toxo imports
from ..integrations.gemini_client import GeminiClient
from ..utils.logger import get_logger
from ..utils.exceptions import ModelError, ValidationError


@dataclass
class RerankingCandidate:
    """Represents a candidate response for reranking."""
    response_id: str
    content: str
    source: str  # e.g., 'gemini', 'cached', 'template'
    confidence: float
    metadata: Dict[str, Any]
    embeddings: Optional[np.ndarray] = None


@dataclass
class RerankingCriteria:
    """Criteria for response reranking."""
    semantic_relevance: float = 1.0
    factual_accuracy: float = 1.0
    safety_score: float = 1.0
    user_preference: float = 1.0
    coherence: float = 0.8
    completeness: float = 0.8
    style_match: float = 0.6
    brevity: float = 0.4


@dataclass
class RerankingResult:
    """Result of response reranking."""
    ranked_candidates: List[Tuple[RerankingCandidate, float]]
    best_response: RerankingCandidate
    scoring_breakdown: Dict[str, List[float]]
    confidence: float
    reasoning: str


class CrossEncoder(nn.Module):
    """
    Neural cross-encoder for response relevance scoring.
    """
    
    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.encoder = AutoModel.from_pretrained(model_name)
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(self.encoder.config.hidden_size, 1)
        
    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Forward pass through the cross-encoder."""
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        pooled_output = outputs.last_hidden_state[:, 0]  # CLS token
        pooled_output = self.dropout(pooled_output)
        relevance_score = self.classifier(pooled_output)
        return torch.sigmoid(relevance_score)
    
    def encode_pair(self, query: str, response: str) -> Tuple[torch.Tensor, torch.Tensor]:
        """Encode query-response pair for scoring."""
        encoded = self.tokenizer(
            query, response,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        )
        return encoded['input_ids'], encoded['attention_mask']


class MultiCriteriaRanker(nn.Module):
    """
    Multi-criteria neural ranker that combines multiple scoring dimensions.
    """
    
    def __init__(self, embedding_dim: int = 384, num_criteria: int = 8):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.num_criteria = num_criteria
        
        # Feature processing layers
        self.query_encoder = nn.Linear(embedding_dim, 128)
        self.response_encoder = nn.Linear(embedding_dim, 128)
        self.interaction_layer = nn.Bilinear(128, 128, 64)
        
        # Criteria-specific scoring heads
        self.criteria_heads = nn.ModuleList([
            nn.Sequential(
                nn.Linear(256, 128),
                nn.ReLU(),
                nn.Dropout(0.1),
                nn.Linear(128, 1)
            ) for _ in range(num_criteria)
        ])
        
        # Final ranking layer
        self.ranking_layer = nn.Sequential(
            nn.Linear(num_criteria, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
    
    def forward(
        self,
        query_embeddings: torch.Tensor,
        response_embeddings: torch.Tensor,
        criteria_weights: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass for multi-criteria ranking.
        
        Returns:
            overall_scores: Final ranking scores
            criteria_scores: Individual criteria scores
        """
        # Encode features
        query_features = F.relu(self.query_encoder(query_embeddings))
        response_features = F.relu(self.response_encoder(response_embeddings))
        
        # Interaction features
        interaction_features = self.interaction_layer(query_features, response_features)
        
        # Combined features
        combined_features = torch.cat([
            query_features, response_features, interaction_features
        ], dim=-1)
        
        # Criteria-specific scores
        criteria_scores = []
        for head in self.criteria_heads:
            score = torch.sigmoid(head(combined_features))
            criteria_scores.append(score)
        
        criteria_scores = torch.cat(criteria_scores, dim=-1)
        
        # Weight by criteria importance
        weighted_scores = criteria_scores * criteria_weights
        
        # Final ranking score
        overall_scores = self.ranking_layer(weighted_scores)
        
        return torch.sigmoid(overall_scores), criteria_scores


class NeuralReranker:
    """
    Advanced neural reranking system for Toxo responses.
    """
    
    def __init__(
        self,
        gemini_client: Optional[GeminiClient] = None,
        model_path: Optional[str] = None,
        device: str = "cpu"
    ):
        self.logger = get_logger(__name__)
        self.gemini_client = gemini_client
        self.device = torch.device(device)
        
        # Initialize models
        self.cross_encoder = CrossEncoder().to(self.device)
        self.multi_criteria_ranker = MultiCriteriaRanker().to(self.device)
        
        # Load pre-trained weights if available
        if model_path and Path(model_path).exists():
            self.load_models(model_path)
        
        # Initialize tokenizer for embeddings
        self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        self.embedding_model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2").to(self.device)
        
        # Initialize scaler for feature normalization
        self.feature_scaler = StandardScaler()
        
        # Safety and quality classifiers
        self.safety_keywords = {
            'harmful': ['violence', 'hate', 'discrimination', 'illegal', 'harmful'],
            'inappropriate': ['explicit', 'offensive', 'inappropriate', 'nsfw'],
            'misinformation': ['false', 'misleading', 'unverified', 'conspiracy']
        }
        
        # User preference patterns (learned from feedback)
        self.user_preferences = {
            'style': {},
            'length': {},
            'formality': {},
            'detail_level': {}
        }
    
    async def rerank_responses(
        self,
        query: str,
        candidates: List[RerankingCandidate],
        criteria: Optional[RerankingCriteria] = None,
        user_context: Optional[Dict[str, Any]] = None
    ) -> RerankingResult:
        """
        Main entry point for response reranking.
        
        Args:
            query: Original user query
            candidates: List of candidate responses
            criteria: Reranking criteria with weights
            user_context: User-specific context and preferences
            
        Returns:
            RerankingResult with ranked responses and explanations
        """
        if not candidates:
            raise ValidationError("No candidates provided for reranking")
        
        criteria = criteria or RerankingCriteria()
        user_context = user_context or {}
        
        self.logger.info(f"Reranking {len(candidates)} response candidates")
        
        try:
            # Generate embeddings for all candidates
            await self._generate_embeddings(query, candidates)
            
            # Score candidates on multiple criteria
            scoring_results = await self._score_candidates(query, candidates, criteria, user_context)
            
            # Apply neural ranking
            neural_scores = await self._apply_neural_ranking(query, candidates, criteria)
            
            # Combine scores and rank
            final_scores = self._combine_scores(scoring_results, neural_scores, criteria)
            
            # Sort by final scores
            ranked_candidates = sorted(
                zip(candidates, final_scores),
                key=lambda x: x[1],
                reverse=True
            )
            
            # Generate explanation
            reasoning = await self._generate_ranking_explanation(
                query, ranked_candidates[:3], scoring_results
            )
            
            result = RerankingResult(
                ranked_candidates=ranked_candidates,
                best_response=ranked_candidates[0][0],
                scoring_breakdown=scoring_results,
                confidence=self._calculate_ranking_confidence(final_scores),
                reasoning=reasoning
            )
            
            # Update user preferences based on selection
            await self._update_user_preferences(user_context, result)
            
            self.logger.info(f"Reranking completed. Best response: {result.best_response.response_id}")
            return result
            
        except Exception as e:
            self.logger.error(f"Reranking failed: {str(e)}")
            raise ModelError(f"Neural reranking failed: {str(e)}") from e
    
    async def _generate_embeddings(self, query: str, candidates: List[RerankingCandidate]):
        """Generate embeddings for query and all candidates."""
        texts = [query] + [candidate.content for candidate in candidates]
        
        # Tokenize all texts
        encoded = self.tokenizer(
            texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        ).to(self.device)
        
        # Generate embeddings
        with torch.no_grad():
            outputs = self.embedding_model(**encoded)
            embeddings = outputs.last_hidden_state[:, 0].cpu().numpy()  # CLS embeddings
        
        # Store embeddings
        query_embedding = embeddings[0]
        for i, candidate in enumerate(candidates):
            candidate.embeddings = embeddings[i + 1]
        
        return query_embedding
    
    async def _score_candidates(
        self,
        query: str,
        candidates: List[RerankingCandidate],
        criteria: RerankingCriteria,
        user_context: Dict[str, Any]
    ) -> Dict[str, List[float]]:
        """Score candidates on multiple criteria."""
        scoring_results = {
            'semantic_relevance': [],
            'factual_accuracy': [],
            'safety_score': [],
            'user_preference': [],
            'coherence': [],
            'completeness': [],
            'style_match': [],
            'brevity': []
        }
        
        query_embedding = await self._generate_embeddings(query, [])
        
        for candidate in candidates:
            # Semantic relevance (cosine similarity)
            relevance = cosine_similarity(
                [query_embedding], [candidate.embeddings]
            )[0][0]
            scoring_results['semantic_relevance'].append(float(relevance))
            
            # Safety score
            safety = self._assess_safety(candidate.content)
            scoring_results['safety_score'].append(safety)
            
            # Coherence and completeness
            coherence = self._assess_coherence(candidate.content)
            completeness = self._assess_completeness(candidate.content, query)
            scoring_results['coherence'].append(coherence)
            scoring_results['completeness'].append(completeness)
            
            # Style match and brevity
            style_match = self._assess_style_match(candidate.content, user_context)
            brevity = self._assess_brevity(candidate.content)
            scoring_results['style_match'].append(style_match)
            scoring_results['brevity'].append(brevity)
            
            # User preference (learned from history)
            user_pref = self._assess_user_preference(candidate.content, user_context)
            scoring_results['user_preference'].append(user_pref)
        
        # Factual accuracy (requires AI analysis)
        if self.gemini_client:
            factual_scores = await self._assess_factual_accuracy(query, candidates)
            scoring_results['factual_accuracy'] = factual_scores
        else:
            scoring_results['factual_accuracy'] = [0.5] * len(candidates)  # Neutral
        
        return scoring_results
    
    async def _apply_neural_ranking(
        self,
        query: str,
        candidates: List[RerankingCandidate],
        criteria: RerankingCriteria
    ) -> List[float]:
        """Apply neural ranking models."""
        neural_scores = []
        
        # Prepare criteria weights
        criteria_weights = torch.tensor([
            criteria.semantic_relevance,
            criteria.factual_accuracy,
            criteria.safety_score,
            criteria.user_preference,
            criteria.coherence,
            criteria.completeness,
            criteria.style_match,
            criteria.brevity
        ]).unsqueeze(0).to(self.device)
        
        query_embedding = await self._generate_embeddings(query, [])
        query_tensor = torch.tensor(query_embedding).unsqueeze(0).to(self.device)
        
        for candidate in candidates:
            # Cross-encoder score
            input_ids, attention_mask = self.cross_encoder.encode_pair(query, candidate.content)
            with torch.no_grad():
                cross_score = self.cross_encoder(input_ids.to(self.device), attention_mask.to(self.device))
            
            # Multi-criteria score
            response_tensor = torch.tensor(candidate.embeddings).unsqueeze(0).to(self.device)
            with torch.no_grad():
                overall_score, _ = self.multi_criteria_ranker(
                    query_tensor, response_tensor, criteria_weights
                )
            
            # Combine neural scores
            combined_score = (cross_score.item() + overall_score.item()) / 2
            neural_scores.append(combined_score)
        
        return neural_scores
    
    def _combine_scores(
        self,
        traditional_scores: Dict[str, List[float]],
        neural_scores: List[float],
        criteria: RerankingCriteria
    ) -> List[float]:
        """Combine traditional and neural scores with criteria weights."""
        num_candidates = len(neural_scores)
        final_scores = []
        
        # Normalize all score lists to [0, 1]
        for criterion in traditional_scores:
            scores = traditional_scores[criterion]
            if scores and max(scores) > min(scores):
                normalized = [(s - min(scores)) / (max(scores) - min(scores)) for s in scores]
                traditional_scores[criterion] = normalized
        
        neural_scores_normalized = neural_scores
        if max(neural_scores) > min(neural_scores):
            neural_scores_normalized = [
                (s - min(neural_scores)) / (max(neural_scores) - min(neural_scores))
                for s in neural_scores
            ]
        
        for i in range(num_candidates):
            # Traditional criteria scores
            traditional_score = (
                traditional_scores['semantic_relevance'][i] * criteria.semantic_relevance +
                traditional_scores['factual_accuracy'][i] * criteria.factual_accuracy +
                traditional_scores['safety_score'][i] * criteria.safety_score +
                traditional_scores['user_preference'][i] * criteria.user_preference +
                traditional_scores['coherence'][i] * criteria.coherence +
                traditional_scores['completeness'][i] * criteria.completeness +
                traditional_scores['style_match'][i] * criteria.style_match +
                traditional_scores['brevity'][i] * criteria.brevity
            )
            
            # Normalize by total weight
            total_weight = (
                criteria.semantic_relevance + criteria.factual_accuracy +
                criteria.safety_score + criteria.user_preference +
                criteria.coherence + criteria.completeness +
                criteria.style_match + criteria.brevity
            )
            traditional_score /= total_weight
            
            # Combine with neural score (weighted combination)
            final_score = 0.6 * traditional_score + 0.4 * neural_scores_normalized[i]
            final_scores.append(final_score)
        
        return final_scores
    
    def _assess_safety(self, content: str) -> float:
        """Assess safety of response content."""
        content_lower = content.lower()
        safety_score = 1.0
        
        for category, keywords in self.safety_keywords.items():
            for keyword in keywords:
                if keyword in content_lower:
                    if category == 'harmful':
                        safety_score -= 0.3
                    elif category == 'inappropriate':
                        safety_score -= 0.2
                    elif category == 'misinformation':
                        safety_score -= 0.1
        
        return max(0.0, safety_score)
    
    def _assess_coherence(self, content: str) -> float:
        """Assess coherence of response content."""
        # Simple heuristics for coherence
        sentences = content.split('.')
        if len(sentences) < 2:
            return 0.7  # Single sentence
        
        # Check for logical flow (simplified)
        coherence_score = 0.8
        
        # Penalty for very short or very repetitive responses
        if len(content.split()) < 5:
            coherence_score -= 0.2
        
        return max(0.0, min(1.0, coherence_score))
    
    def _assess_completeness(self, content: str, query: str) -> float:
        """Assess completeness of response relative to query."""
        # Simple heuristic based on length and content coverage
        query_words = set(query.lower().split())
        content_words = set(content.lower().split())
        
        # Coverage of query terms
        coverage = len(query_words.intersection(content_words)) / max(len(query_words), 1)
        
        # Length appropriateness
        word_count = len(content.split())
        if word_count < 10:
            length_score = 0.5
        elif word_count < 50:
            length_score = 0.8
        elif word_count < 200:
            length_score = 1.0
        else:
            length_score = 0.9  # Slight penalty for very long responses
        
        return (coverage + length_score) / 2
    
    def _assess_style_match(self, content: str, user_context: Dict[str, Any]) -> float:
        """Assess style match with user preferences."""
        # Default neutral score
        style_score = 0.5
        
        # Check formality level
        formal_indicators = ['therefore', 'however', 'furthermore', 'consequently']
        informal_indicators = ["that's", "you're", "it's", "can't", "won't"]
        
        formal_count = sum(1 for word in formal_indicators if word in content.lower())
        informal_count = sum(1 for word in informal_indicators if word in content.lower())
        
        # Adjust based on user preference (if known)
        preferred_formality = user_context.get('formality_preference', 'neutral')
        if preferred_formality == 'formal' and formal_count > informal_count:
            style_score += 0.3
        elif preferred_formality == 'informal' and informal_count > formal_count:
            style_score += 0.3
        
        return max(0.0, min(1.0, style_score))
    
    def _assess_brevity(self, content: str) -> float:
        """Assess brevity (shorter is better, but not too short)."""
        word_count = len(content.split())
        
        if word_count < 5:
            return 0.3  # Too short
        elif word_count < 30:
            return 1.0  # Good brevity
        elif word_count < 100:
            return 0.8  # Acceptable
        else:
            return 0.5  # Too long
    
    def _assess_user_preference(self, content: str, user_context: Dict[str, Any]) -> float:
        """Assess alignment with learned user preferences."""
        # This would be enhanced with actual user preference learning
        return 0.5  # Neutral default
    
    async def _assess_factual_accuracy(
        self,
        query: str,
        candidates: List[RerankingCandidate]
    ) -> List[float]:
        """Assess factual accuracy using AI analysis."""
        if not self.gemini_client:
            return [0.5] * len(candidates)
        
        accuracy_scores = []
        
        for candidate in candidates:
            try:
                accuracy_prompt = f"""
                Assess the factual accuracy of this response to the given query.
                Consider:
                1. Are the facts stated correct?
                2. Are there any misleading statements?
                3. Is the information up-to-date and relevant?
                
                Query: {query}
                Response: {candidate.content}
                
                Provide a factual accuracy score from 0.0 to 1.0, where:
                - 1.0 = Completely accurate and factual
                - 0.8 = Mostly accurate with minor issues
                - 0.6 = Generally accurate but some concerns
                - 0.4 = Mixed accuracy, significant issues
                - 0.2 = Mostly inaccurate
                - 0.0 = Completely inaccurate or misleading
                
                Return only the numerical score.
                """
                
                response = await self.gemini_client.generate(accuracy_prompt)
                try:
                    score = float(response.strip())
                    accuracy_scores.append(max(0.0, min(1.0, score)))
                except ValueError:
                    accuracy_scores.append(0.5)  # Default if parsing fails
                    
            except Exception as e:
                self.logger.warning(f"Factual accuracy assessment failed: {e}")
                accuracy_scores.append(0.5)
        
        return accuracy_scores
    
    async def _generate_ranking_explanation(
        self,
        query: str,
        top_candidates: List[Tuple[RerankingCandidate, float]],
        scoring_breakdown: Dict[str, List[float]]
    ) -> str:
        """Generate explanation for ranking decisions."""
        if not self.gemini_client:
            return "Ranking based on neural scoring and multi-criteria analysis."
        
        try:
            explanation_prompt = f"""
            Explain why the following response was selected as the best answer to the query.
            
            Query: {query}
            
            Best Response: {top_candidates[0][0].content}
            Score: {top_candidates[0][1]:.3f}
            
            Provide a brief, clear explanation focusing on:
            1. Why this response best addresses the query
            2. Key strengths that made it stand out
            3. How it compares to alternatives
            
            Keep the explanation concise and user-friendly.
            """
            
            explanation = await self.gemini_client.generate(explanation_prompt)
            return explanation.strip()
            
        except Exception as e:
            self.logger.warning(f"Explanation generation failed: {e}")
            return "Selected based on relevance, accuracy, and quality metrics."
    
    def _calculate_ranking_confidence(self, scores: List[float]) -> float:
        """Calculate confidence in the ranking based on score distribution."""
        if len(scores) < 2:
            return 1.0
        
        scores_sorted = sorted(scores, reverse=True)
        
        # Confidence based on gap between top scores
        top_gap = scores_sorted[0] - scores_sorted[1]
        
        # Higher gap = higher confidence
        confidence = min(1.0, 0.5 + top_gap * 2)
        
        return confidence
    
    async def _update_user_preferences(
        self,
        user_context: Dict[str, Any],
        result: RerankingResult
    ):
        """Update learned user preferences based on selection."""
        # This would implement preference learning algorithms
        # For now, just log the selection for future analysis
        self.logger.debug(f"User selected response: {result.best_response.response_id}")
    
    def train_neural_models(
        self,
        training_data: List[Dict[str, Any]],
        validation_data: Optional[List[Dict[str, Any]]] = None
    ):
        """Train the neural ranking models on provided data."""
        # Implementation for training the neural models
        # This would include data preparation, training loops, validation, etc.
        self.logger.info("Neural model training not implemented in this version")
    
    def save_models(self, path: str):
        """Save trained models to disk."""
        save_path = Path(path)
        save_path.mkdir(parents=True, exist_ok=True)
        
        torch.save(self.cross_encoder.state_dict(), save_path / "cross_encoder.pt")
        torch.save(self.multi_criteria_ranker.state_dict(), save_path / "multi_criteria_ranker.pt")
        
        # Save other components
        joblib.dump(self.feature_scaler, save_path / "feature_scaler.pkl")
        
        with open(save_path / "user_preferences.json", 'w') as f:
            json.dump(self.user_preferences, f)
    
    def load_models(self, path: str):
        """Load trained models from disk."""
        load_path = Path(path)
        
        if (load_path / "cross_encoder.pt").exists():
            self.cross_encoder.load_state_dict(torch.load(load_path / "cross_encoder.pt"))
        
        if (load_path / "multi_criteria_ranker.pt").exists():
            self.multi_criteria_ranker.load_state_dict(torch.load(load_path / "multi_criteria_ranker.pt"))
        
        if (load_path / "feature_scaler.pkl").exists():
            self.feature_scaler = joblib.load(load_path / "feature_scaler.pkl")
        
        if (load_path / "user_preferences.json").exists():
            with open(load_path / "user_preferences.json", 'r') as f:
                self.user_preferences = json.load(f) 