"""
Ensemble Voting Strategies
═══════════════════════════════════════════════════════════

Pure aggregation math - no models, no opinions.

STRATEGIES:
    majority: Most common answer wins
    confidence: Weighted by confidence scores
    fitness_weighted: Weighted by fitness (for evolved ensembles)
    unanimous: Only agree if all match
    threshold: Agree if confidence exceeds threshold

USAGE:
    >>> from cascade.ensemble import EnsembleVoting
    >>> 
    >>> outputs = [
    ...     {'answer': 'yes', 'score': 0.9},
    ...     {'answer': 'yes', 'score': 0.7},
    ...     {'answer': 'no', 'score': 0.8},
    ... ]
    >>> confidences = [0.9, 0.7, 0.6]
    >>> 
    >>> result = EnsembleVoting.confidence(outputs, confidences)
    >>> # Result weighted towards 'yes' due to higher total confidence
"""

import numpy as np
from typing import Dict, Any, List, Optional, Union
from collections import Counter


class EnsembleVoting:
    """
    Static voting strategies for ensemble aggregation.
    
    These are pure functions - no state, no model dependencies.
    Pass in outputs and weights, get aggregated result.
    """
    
    @staticmethod
    def majority(outputs: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Majority voting - most common output wins.
        
        For discrete values: mode (most frequent)
        For arrays: element-wise mean
        
        Args:
            outputs: List of output dictionaries from ensemble members
            
        Returns:
            Aggregated dictionary with majority values
        """
        if not outputs:
            return {}
        
        result = {}
        all_keys = set()
        for out in outputs:
            if out is not None:
                all_keys.update(out.keys())
        
        for key in all_keys:
            values = [out.get(key) for out in outputs if out is not None and key in out]
            if not values:
                continue
            
            first = values[0]
            
            # Check if all values have the same type and shape
            try:
                if isinstance(first, np.ndarray):
                    # Verify all are arrays of same shape
                    if all(isinstance(v, np.ndarray) and v.shape == first.shape for v in values):
                        result[key] = np.mean(values, axis=0)
                    else:
                        # Heterogeneous shapes - just take first
                        result[key] = first
                elif isinstance(first, list):
                    # Check if all lists have same length
                    if all(isinstance(v, list) and len(v) == len(first) for v in values):
                        result[key] = np.mean(values, axis=0).tolist()
                    else:
                        result[key] = first
                elif isinstance(first, (int, float)):
                    # Mean for scalars (filter to only numeric)
                    numeric = [v for v in values if isinstance(v, (int, float))]
                    if numeric:
                        result[key] = np.mean(numeric)
                    else:
                        result[key] = first
                else:
                    # Mode for discrete (strings, etc.)
                    hashable = [v if not isinstance(v, dict) else str(v) for v in values]
                    result[key] = Counter(hashable).most_common(1)[0][0]
            except Exception:
                # Any error - just take the first value
                result[key] = first
        
        return result
    
    @staticmethod
    def confidence(outputs: List[Dict[str, Any]], 
                   confidences: List[float]) -> Dict[str, Any]:
        """
        Confidence-weighted voting.
        
        Higher confidence outputs have more weight in the final result.
        
        Args:
            outputs: List of output dictionaries
            confidences: List of confidence scores (0-1)
            
        Returns:
            Confidence-weighted aggregated dictionary
        """
        if not outputs or not confidences:
            return {}
        
        # Normalize confidences to sum to 1
        conf_array = np.array(confidences, dtype=float)
        total = conf_array.sum() + 1e-8
        weights = conf_array / total
        
        result = {}
        all_keys = set()
        for out in outputs:
            if out is not None:
                all_keys.update(out.keys())
        
        for key in all_keys:
            values = []
            key_weights = []
            
            for out, w in zip(outputs, weights):
                if out is not None and key in out and out[key] is not None:
                    values.append(out[key])
                    key_weights.append(w)
            
            if not values:
                continue
            
            first = values[0]
            
            try:
                if isinstance(first, np.ndarray):
                    # Verify all are arrays of same shape
                    if all(isinstance(v, np.ndarray) and v.shape == first.shape for v in values):
                        w_norm = np.array(key_weights) / (sum(key_weights) + 1e-8)
                        result[key] = np.average(values, weights=w_norm, axis=0)
                    else:
                        # Heterogeneous - use highest weighted
                        max_idx = np.argmax(key_weights)
                        result[key] = values[max_idx]
                elif isinstance(first, list):
                    # Check if all lists have same length
                    if all(isinstance(v, list) and len(v) == len(first) for v in values):
                        w_norm = np.array(key_weights) / (sum(key_weights) + 1e-8)
                        result[key] = np.average(values, weights=w_norm, axis=0).tolist()
                    else:
                        max_idx = np.argmax(key_weights)
                        result[key] = values[max_idx]
                elif isinstance(first, (int, float)):
                    # Weighted average for scalars (filter to numeric)
                    numeric_vals = []
                    numeric_weights = []
                    for v, w in zip(values, key_weights):
                        if isinstance(v, (int, float)):
                            numeric_vals.append(v)
                            numeric_weights.append(w)
                    if numeric_vals:
                        result[key] = np.average(numeric_vals, weights=numeric_weights)
                    else:
                        result[key] = first
                else:
                    # Weighted mode for discrete
                    weighted = {}
                    for v, w in zip(values, key_weights):
                        v_key = v if not isinstance(v, dict) else str(v)
                        weighted[v_key] = weighted.get(v_key, 0) + w
                    result[key] = max(weighted, key=weighted.get)
            except Exception:
                # Fallback to highest weighted value
                max_idx = np.argmax(key_weights) if key_weights else 0
                result[key] = values[max_idx] if values else first
        
        return result
    
    @staticmethod
    def fitness_weighted(outputs: List[Dict[str, Any]], 
                         fitnesses: List[float]) -> Dict[str, Any]:
        """
        Fitness-weighted voting (for evolved ensembles).
        
        Higher fitness members have more influence.
        Negative fitnesses are shifted to positive.
        
        Args:
            outputs: List of output dictionaries
            fitnesses: List of fitness scores (can be negative)
            
        Returns:
            Fitness-weighted aggregated dictionary
        """
        if not fitnesses:
            return EnsembleVoting.majority(outputs)
        
        # Shift to positive (handle negative fitness)
        min_fit = min(fitnesses)
        shifted = [f - min_fit + 1e-8 for f in fitnesses]
        
        return EnsembleVoting.confidence(outputs, shifted)
    
    @staticmethod
    def unanimous(outputs: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
        """
        Unanimous voting - only return result if all agree.
        
        Returns None if there's any disagreement.
        
        Args:
            outputs: List of output dictionaries
            
        Returns:
            The unanimous output, or None if disagreement
        """
        if not outputs:
            return None
        
        # Check if all outputs match
        first = outputs[0]
        for out in outputs[1:]:
            if out != first:
                return None
        
        return first
    
    @staticmethod
    def threshold(outputs: List[Dict[str, Any]], 
                  confidences: List[float],
                  threshold: float = 0.8) -> Optional[Dict[str, Any]]:
        """
        Threshold voting - only use outputs above confidence threshold.
        
        Filters to high-confidence outputs, then does majority vote.
        Returns None if no outputs pass threshold.
        
        Args:
            outputs: List of output dictionaries
            confidences: List of confidence scores (0-1)
            threshold: Minimum confidence to include (default 0.8)
            
        Returns:
            Aggregated high-confidence outputs, or None
        """
        if not outputs or not confidences:
            return None
        
        # Filter to above threshold
        filtered_outputs = []
        filtered_confidences = []
        
        for out, conf in zip(outputs, confidences):
            if conf >= threshold:
                filtered_outputs.append(out)
                filtered_confidences.append(conf)
        
        if not filtered_outputs:
            return None
        
        return EnsembleVoting.confidence(filtered_outputs, filtered_confidences)
    
    @staticmethod
    def softmax_temperature(outputs: List[Dict[str, Any]], 
                           confidences: List[float],
                           temperature: float = 1.0) -> Dict[str, Any]:
        """
        Softmax temperature voting.
        
        Applies temperature scaling to confidences before weighting:
        - temperature < 1: Sharper (high confidence dominates)
        - temperature = 1: Normal
        - temperature > 1: Flatter (more equal weighting)
        
        Args:
            outputs: List of output dictionaries
            confidences: List of confidence scores
            temperature: Temperature parameter (default 1.0)
            
        Returns:
            Temperature-scaled weighted aggregation
        """
        if not confidences or temperature <= 0:
            return EnsembleVoting.majority(outputs)
        
        # Apply softmax with temperature
        conf_array = np.array(confidences, dtype=float)
        scaled = conf_array / temperature
        exp_scaled = np.exp(scaled - np.max(scaled))  # Numerically stable
        softmax_weights = exp_scaled / (exp_scaled.sum() + 1e-8)
        
        return EnsembleVoting.confidence(outputs, softmax_weights.tolist())
    
    @staticmethod
    def rank_based(outputs: List[Dict[str, Any]], 
                   confidences: List[float]) -> Dict[str, Any]:
        """
        Rank-based voting.
        
        Weights by rank rather than raw confidence.
        More robust to outlier confidence values.
        
        Args:
            outputs: List of output dictionaries
            confidences: List of confidence scores
            
        Returns:
            Rank-weighted aggregation
        """
        if not confidences:
            return EnsembleVoting.majority(outputs)
        
        # Convert to ranks (1 = highest confidence)
        n = len(confidences)
        sorted_indices = np.argsort(confidences)[::-1]  # Descending
        ranks = np.zeros(n)
        for rank, idx in enumerate(sorted_indices):
            ranks[idx] = n - rank  # Higher rank = higher weight
        
        return EnsembleVoting.confidence(outputs, ranks.tolist())