"""Ensemble judge that combines multiple judges."""

from typing import List, Optional
from concurrent.futures import ThreadPoolExecutor, as_completed
import numpy as np
from collections import Counter

from openevalkit.judges.base import Judge
from openevalkit.judges.rubric import Rubric
from openevalkit.judgment import Judgment
from openevalkit.run import Run
from openevalkit.errors import JudgingError


class EnsembleJudge(Judge):
    """
    Combines multiple judges for more reliable evaluation.
    
    Uses multiple LLM judges and aggregates their judgments to reduce
    variance and improve reliability.
    
    Examples:
        >>> from openevalkit.judges import EnsembleJudge, LLMJudge, LLMConfig, Rubric
        >>> 
        >>> rubric = Rubric(criteria=["helpfulness"], scale="0-1")
        >>> 
        >>> # Sequential (default)
        >>> ensemble = EnsembleJudge(
        ...     judges=[
        ...         LLMJudge(LLMConfig(model="gpt-4o-mini"), rubric),
        ...         LLMJudge(LLMConfig(model="gpt-4o"), rubric),
        ...         LLMJudge(LLMConfig(model="claude-3-5-sonnet-20241022"), rubric)
        ...     ],
        ...     method="average"
        ... )
        >>> 
        >>> # Parallel (faster, matches sklearn pattern)
        >>> ensemble = EnsembleJudge(
        ...     judges=[...],
        ...     n_jobs=3,  # Run 3 judges in parallel
        ...     method="average"
        ... )
    """
    
    def __init__(
        self,
        judges: List[Judge],
        method: str = "average",
        min_agreement: Optional[float] = None,
        n_jobs: Optional[int] = None,  
        name: Optional[str] = None
    ):
        """
        Initialize ensemble judge.
        
        Args:
            judges: List of Judge instances to combine (2+)
            method: Aggregation method
                - "average": Mean of continuous scores
                - "majority_vote": Mode for discrete/binary
                - "median": Median (robust to outliers)
                - "unanimous": Minimum score (all must agree it's good)
            min_agreement: Warn if correlation between judges below this threshold
            n_jobs: Number of judges to run in parallel
                - None or 1: Sequential execution (default)
                - n > 1: Parallel execution with n threads
                - -1: Use all judges in parallel (len(judges) threads)
            name: Optional name (auto-generated if None)
        """
        if len(judges) < 2:
            raise ValueError("EnsembleJudge requires at least 2 judges")
        
        # Validate all judges have rubrics
        for i, judge in enumerate(judges):
            if not hasattr(judge, 'rubric') or judge.rubric is None:
                raise ValueError(
                    f"Judge {i} ({judge.name}) does not have a rubric. "
                    f"All judges in ensemble must have rubrics."
                )
        
        # Extract rubric from first judge
        self.rubric = judges[0].rubric
        
        # VALIDATE: All judges must use the SAME rubric
        for i, judge in enumerate(judges[1:], start=1):
            if judge.rubric.criteria != self.rubric.criteria:
                raise ValueError(
                    f"Judge {i} ({judge.name}) has different criteria than first judge.\n"
                    f"  First judge criteria: {self.rubric.criteria}\n"
                    f"  Judge {i} criteria: {judge.rubric.criteria}\n"
                    f"All judges in ensemble must evaluate the same criteria."
                )
            
            if judge.rubric.scale != self.rubric.scale:
                raise ValueError(
                    f"Judge {i} ({judge.name}) has different scale than first judge.\n"
                    f"  First judge scale: {self.rubric.scale}\n"
                    f"  Judge {i} scale: {judge.rubric.scale}\n"
                    f"All judges in ensemble must use the same scale."
                )
        
        valid_methods = ["average", "majority_vote", "median", "unanimous"]
        if method not in valid_methods:
            raise ValueError(
                f"Invalid method: '{method}'. Must be one of: {valid_methods}"
            )
        
        self.judges = judges
        self.method = method
        self.min_agreement = min_agreement
        
        # Set n_jobs (sklearn pattern)
        if n_jobs == -1:
            self.n_jobs = len(judges)  # All judges in parallel
        elif n_jobs is None or n_jobs == 1:
            self.n_jobs = 1  # Sequential
        else:
            self.n_jobs = min(n_jobs, len(judges))  # Cap at num judges
        
        self.name = name or f"ensemble_{len(judges)}judges_{method}"
    
    def judge(self, run: Run, rubric: Rubric) -> Judgment:
        """
        Judge using ensemble of judges.
        
        Args:
            run: Run to judge
            rubric: Evaluation criteria
            
        Returns:
            Aggregated judgment from all judges
            
        Raises:
            JudgingError: If judging fails
        """
        # Get judgment from each judge (parallel or sequential)
        if self.n_jobs == 1:
            judgments = self._judge_sequential(run, rubric)
        else:
            judgments = self._judge_parallel(run, rubric)
        
        # Combine judgments
        combined = self._combine_judgments(judgments, rubric)
        
        # Check agreement if threshold set
        if self.min_agreement:
            agreement = self._compute_agreement(judgments)
            if agreement < self.min_agreement:
                import warnings
                warnings.warn(
                    f"Low judge agreement ({agreement:.2f}) for run {run.id}. "
                    f"Expected >= {self.min_agreement}"
                )
        
        return combined
    
    def _judge_sequential(self, run: Run, rubric: Rubric) -> List[Judgment]:
        """Sequential judging (original implementation)."""
        judgments = []
        for i, judge in enumerate(self.judges):
            try:
                judgment = judge.judge(run, rubric)
                judgments.append(judgment)
            except Exception as e:
                raise JudgingError(
                    f"Judge {i} ({judge.name}) failed: {str(e)}"
                ) from e
        return judgments
    
    def _judge_parallel(self, run: Run, rubric: Rubric) -> List[Judgment]:
        """Parallel judging using ThreadPoolExecutor."""
        
        def judge_wrapper(judge_info):
            """Wrapper to catch errors per judge."""
            idx, judge = judge_info
            try:
                return (idx, judge.judge(run, rubric), None)
            except Exception as e:
                return (idx, None, e)
        
        # Execute all judges in parallel
        with ThreadPoolExecutor(max_workers=self.n_jobs) as executor:
            # Submit all judges
            futures = {
                executor.submit(judge_wrapper, (i, judge)): i
                for i, judge in enumerate(self.judges)
            }
            
            # Collect results (maintain order)
            results = [None] * len(self.judges)
            for future in as_completed(futures):
                idx, judgment, error = future.result()
                
                if error:
                    raise JudgingError(
                        f"Judge {idx} ({self.judges[idx].name}) failed: {str(error)}"
                    ) from error
                
                results[idx] = judgment
        
        return results
    
    def _combine_judgments(
        self, 
        judgments: List[Judgment], 
        rubric: Rubric
    ) -> Judgment:
        """Combine multiple judgments with full individual tracking."""
        
        # Get all criteria from first judgment
        criteria = list(judgments[0].criteria_scores.keys())
        
        # Aggregate per-criterion scores
        combined_scores = {}
        for criterion in criteria:
            scores = [j.criteria_scores[criterion] for j in judgments]
            combined_scores[criterion] = self._aggregate_values(scores)
        
        # For explanations, use placeholder pointing to metadata
        combined_explanations = {}
        for criterion in criteria:
            combined_explanations[criterion] = (
                f"[Ensemble of {len(judgments)} judges - see metadata for details]"
            )
        
        # Aggregate per-criterion confidences (if available)
        combined_confidences = None
        if all(j.criteria_confidences for j in judgments):
            combined_confidences = {}
            for criterion in criteria:
                confidences = [
                    j.criteria_confidences[criterion] 
                    for j in judgments
                ]
                combined_confidences[criterion] = float(np.mean(confidences))
        
        # Compute overall score from combined criterion scores
        overall_score = self._compute_overall_score(
            combined_scores, 
            rubric.weights
        )
        
        # Compute overall confidence from combined criterion confidences
        overall_confidence = None
        if combined_confidences:
            overall_confidence = self._compute_overall_confidence(
                combined_confidences,
                rubric.weights,
                method="weighted_avg"
            )
        
        # Compute agreement metric
        agreement = self._compute_agreement(judgments)
        
        # Build individual judgments metadata
        individual_judgments = []
        for i, judgment in enumerate(judgments):
            judge_name = judgment.metadata.get("model", f"judge_{i}")
            individual_judgments.append({
                "judge_name": judge_name,
                "score": judgment.score,
                "criteria_scores": judgment.criteria_scores,
                "criteria_explanations": judgment.criteria_explanations,
                "criteria_confidences": judgment.criteria_confidences,
                "explanation": judgment.explanation,
                "confidence": judgment.confidence
            })
        
        return Judgment(
            score=overall_score,
            criteria_scores=combined_scores,
            criteria_explanations=combined_explanations,
            explanation=f"Ensemble of {len(judgments)} judges using {self.method} aggregation",
            criteria_confidences=combined_confidences,
            confidence=overall_confidence,
            metadata={
                "num_judges": len(judgments),
                "method": self.method,
                "agreement": agreement,
                "n_jobs": self.n_jobs,
                "individual_scores": [j.score for j in judgments],
                "individual_judgments": individual_judgments
            }
        )
    
    def _aggregate_values(self, values: List[float]) -> float:
        """Aggregate values based on method."""
        if self.method == "average":
            return float(np.mean(values))
        
        elif self.method == "median":
            return float(np.median(values))
        
        elif self.method == "majority_vote":
            # Warn if using on continuous scores
            if not all(v in {0.0, 1.0} or v == round(v) for v in values):
                import warnings
                warnings.warn(
                    f"majority_vote rounding continuous scores {values}. "
                    f"Best for binary/discrete scores (0/1, pass/fail, categorical)."
                    f"This may lose precision. For continuous scores, consider 'average' or 'median'."
                )
            
            # Round and count
            rounded = [round(v) for v in values]
            counter = Counter(rounded)
            most_common = counter.most_common()  # Get ALL items
            
            # Find the top count
            top_count = most_common[0][1]
            
            # Find all values with the top count (handles N-way ties)
            tied_values = [value for value, count in most_common if count == top_count]
            
            # Check for tie (2-way, 3-way, 4-way, ... N-way)
            if len(tied_values) > 1:
                # Multi-way tie! Fall back to average of original values
                import warnings
                warnings.warn(
                    f"Tie in majority_vote: {len(tied_values)} classes tied "
                    f"with {top_count} votes each ({tied_values}). "
                    f"Falling back to average of original values."
                )
                return float(np.mean(values))  # Use original continuous values
            
            # Clear winner - one class has more votes than all others
            return float(most_common[0][0])
        
        elif self.method == "unanimous":
            # Check if data appears to be multi-class nominal
            unique_rounded = set(round(v) for v in values)
            if len(unique_rounded) > 2:
                import warnings
                warnings.warn(
                    f"unanimous (minimum) appears to be used with {len(unique_rounded)} classes "
                    f"({sorted(unique_rounded)}). This method assumes ordinal data where higher=better. "
                    f"For multi-class classification, use 'majority_vote' instead."
                )
            # Conservative: minimum (all must agree it's good)
            return float(min(values))
        else:
            raise ValueError(f"Unknown method: {self.method}")
        

    def _compute_agreement(self, judgments: List[Judgment]) -> float:
        """
        Compute inter-judge agreement using correlation.
        
        Returns correlation coefficient between judge scores (0-1).
        """
        if len(judgments) < 2:
            return 1.0
        
        # Get overall scores from all judges
        scores = [j.score for j in judgments]
        
        # If all judges gave same score, perfect agreement
        if len(set(scores)) == 1:
            return 1.0
        
        # Compute pairwise correlations
        correlations = []
        for i in range(len(scores)):
            for j in range(i + 1, len(scores)):
                # For single values, use variance-based agreement
                # Lower variance = higher agreement
                diff = abs(scores[i] - scores[j])
                agreement = 1.0 - min(diff, 1.0)  # Cap at 1.0
                correlations.append(agreement)
        
        return float(np.mean(correlations)) if correlations else 1.0