"""Base class for all judges."""

from abc import ABC, abstractmethod
import numpy as np
from typing import List, Dict, Optional

from openevalkit.run import Run
from openevalkit.judgment import Judgment
from openevalkit.judges.rubric import Rubric


class Judge(ABC):
    """
    Base class for all judges.

    Judges evaluate runs against rubrics and return judgments.
    Subclasses must implement judge() method.
    """

    name: str
    cacheable: bool = True

    @abstractmethod
    def judge(self, run: Run, rubric: Rubric) -> Judgment:
        """
        Judge a single run against a rubric

        Args:
            run: The run to judge
            rubric: Evaluation criteria

        Returns:
            Judgment for this run with scores, explanation and confidence

        Raises:
            JudgingError: If judging fails
        """

        raise NotImplementedError()
    
    def batch_judge(self, runs: List[Run], rubric: Rubric) -> List[Judgment]:
        """
        Judge multiple runs (optional batch processing).
        
        Default implementation calls judge() for each run.
        Override for true batch processing.
        
        Args:
            runs: List of runs to judge
            rubric: Evaluation criteria
            
        Returns:
            List of judgments
        """
        return [self.judge(run, rubric) for run in runs]
    

    def aggregate(self, judgments: List[Judgment]) -> Dict[str, float]:
        """
        Aggregate judgments across multiple runs.
        
        Returns per-criterion statistics across all runs:
        {
            'helpfulness_mean': 0.85,
            'helpfulness_median': 0.87,
            'clarity_mean': 0.78,
            'clarity_median': 0.80,
            'overall_mean': 0.815
        }
        
        Args:
            judgments: List of Judgment objects from all runs
            
        Returns:
            Dict of metric_name -> aggregate_value
        """
        if not judgments:
            return {}
        
        # Get all criteria
        criteria = judgments[0].criteria_scores.keys()
        aggregates = {}
        
        # Aggregate per-criterion
        for criterion in criteria:
            scores = [j.criteria_scores[criterion] for j in judgments]
            aggregates[f"{criterion}_mean"] = float(np.mean(scores))
            aggregates[f"{criterion}_median"] = float(np.median(scores))
        
        # Aggregate overall scores
        overall_scores = [j.score for j in judgments]
        aggregates["overall_mean"] = float(np.mean(overall_scores))
        aggregates["overall_median"] = float(np.median(overall_scores))
        
        return aggregates
    
    def _compute_overall_score(
        self, 
        criteria_scores: Dict[str, float], 
        weights: Optional[Dict[str, float]]
    ) -> float:
        """
        Compute overall score from per-criterion scores.
        
        Args:
            criteria_scores: Per-criterion scores from LLM
            weights: Optional weights from rubric
            
        Returns:
            Overall weighted score
        """
        if weights is None:
            # Equal weights (default)
            return float(np.mean(list(criteria_scores.values())))
        
        # Weighted average
        weighted_sum = sum(criteria_scores[c] * weights[c] for c in criteria_scores)
        total_weight = sum(weights.values())
        
        return weighted_sum / total_weight
    
    def _compute_overall_confidence(
        self, 
        criteria_confidences: Dict[str, float],
        weights: Optional[Dict[str, float]],
        method: str = "weighted_avg"
    ) -> float:
        """
        Compute overall confidence from per-criterion confidence.
        
        Args:
            criteria_confidences: Per-criterion confidences from LLM
            weights: Optional weights from rubric
            method: Aggregation method ("weighted_avg", "min", "variance_based")
            
        Returns:
            Overall confidence score
        """
        if method == "weighted_avg":
            if weights:
                weighted_sum = sum(criteria_confidences[c] * weights[c] for c in criteria_confidences)
                total_weight = sum(weights.values())
                return weighted_sum / total_weight
            else:
                return float(np.mean(list(criteria_confidences.values())))
        
        elif method == "min":
            # Conservative: weakest link
            return min(criteria_confidences.values())
        
        elif method == "variance_based":
            # Lower confidence if criteria disagree
            values = list(criteria_confidences.values())
            avg = sum(values) / len(values)
            variance = np.var(values)
            return avg * (1 - min(variance, 1.0))  # Cap variance effect at 1.0
        
        else:
            raise ValueError(f"Unknown confidence method: {method}")