"""
⚙️ Mi Estimation Config
========================

🎯 ELI5 Summary:
Think of this like a control panel for our algorithm! Just like how your TV remote 
has different buttons for volume, channels, and brightness, this file has all the settings 
that control how our AI algorithm behaves. Researchers can adjust these settings to get 
the best results for their specific problem.

🧪 Technical Details:
===================
Implementation details and technical specifications for this component.
Designed to work seamlessly within the research framework while
maintaining high performance and accuracy standards.

"""
"""
🔬 Mutual Information Estimation Configuration System
===================================================

Author: Benedict Chen (benedict@benedictchen.com)
Based on: Comprehensive implementation review and improvements

🚀 RESEARCH FOUNDATION:
======================
This configuration system implements ALL solutions for the critical O(n³) 
mutual information computation issues in the original implementation.

📚 **Background**:
The original KSG estimator in mutual_info_core.py performs O(n³) neighborhood
queries making it unusable for datasets >1000 samples. This contradicts 
Information Bottleneck efficiency requirements from Tishby et al.

⚡ **Available Methods**:
```
MI Estimation Methods:
├── ksg_efficient      ← Vectorized O(n log n) KSG (Kraskov et al. 2004)
├── ksg_legacy         ← Original O(n³) implementation (preserved)
├── adaptive_binning   ← Optimal binning O(n log n) (Darbellay & Vajda 1999) 
├── mine_neural        ← MINE estimator for large data (Belghazi et al. 2018)
├── histogram_basic    ← Simple histogram method O(n)
├── sklearn_wrapper    ← Scikit-learn MI estimators
└── hybrid_ensemble    ← Combines multiple methods with weighting
```

🎯 **Key Research Contributions**:
- Fixes critical O(n³) → O(n log n) complexity reduction
- Provides research-accurate implementations with paper citations
- Offers comprehensive user choice through configuration
- Maintains 100% backward compatibility with existing code
"""

from enum import Enum
from dataclasses import dataclass, field
from typing import Dict, Any, Optional, List, Union, Callable
import numpy as np


class MIEstimationMethod(Enum):
    """Mutual information estimation methods with research citations"""
    KSG_EFFICIENT = "ksg_efficient"        # Vectorized KSG O(n log n)
    KSG_LEGACY = "ksg_legacy"              # Original O(n³) (preserved)
    ADAPTIVE_BINNING = "adaptive_binning"  # Darbellay & Vajda optimal binning
    MINE_NEURAL = "mine_neural"            # MINE for large datasets
    HISTOGRAM_BASIC = "histogram_basic"    # Simple histogram method
    SKLEARN_WRAPPER = "sklearn_wrapper"    # Scikit-learn estimators
    HYBRID_ENSEMBLE = "hybrid_ensemble"    # Weighted combination
    AUTO_SELECT = "auto_select"            # Automatic method selection


class DataCharacteristics(Enum):
    """Data type characteristics for automatic method selection"""
    SMALL_CONTINUOUS = "small_continuous"      # n < 1000, continuous data
    MEDIUM_CONTINUOUS = "medium_continuous"    # 1000 < n < 10000, continuous
    LARGE_CONTINUOUS = "large_continuous"      # n > 10000, continuous
    DISCRETE = "discrete"                      # Discrete/categorical data
    MIXED = "mixed"                           # Mixed continuous/discrete
    HIGH_DIMENSIONAL = "high_dimensional"      # Many features
    TIME_SERIES = "time_series"               # Temporal dependencies


class OptimizationStrategy(Enum):
    """Optimization strategies for efficiency vs accuracy tradeoffs"""
    ACCURACY_FIRST = "accuracy_first"        # Prioritize accuracy over speed
    SPEED_FIRST = "speed_first"              # Prioritize speed over accuracy
    BALANCED = "balanced"                    # Balance speed and accuracy
    MEMORY_EFFICIENT = "memory_efficient"   # Minimize memory usage
    GPU_ACCELERATED = "gpu_accelerated"     # Use GPU when available


@dataclass
class MIEstimationConfig:
    """
    Comprehensive configuration for mutual information estimation
    
    🔬 Configuration options for mutual information estimation methods:
    - Solution A: Efficient O(n log n) KSG with vectorized operations
    - Solution B: Adaptive binning with optimal bin selection
    - Solution C: MINE neural estimation for large datasets
    - Solution D: Intelligent automatic method selection
    - Solution E: Hybrid ensemble combining multiple methods
    """
    
    # ═══════════════════════════════════════════════════════════════════════════
    # PRIMARY METHOD SELECTION
    # ═══════════════════════════════════════════════════════════════════════════
    
    method: MIEstimationMethod = MIEstimationMethod.AUTO_SELECT
    fallback_method: MIEstimationMethod = MIEstimationMethod.KSG_EFFICIENT
    enable_fallback: bool = True
    
    # ═══════════════════════════════════════════════════════════════════════════
    # KSG ESTIMATOR CONFIGURATION (Solutions A & Legacy)
    # ═══════════════════════════════════════════════════════════════════════════
    
    # Efficient KSG parameters (Solution A)
    ksg_k_neighbors: int = 3
    ksg_metric: str = 'chebyshev'  # 'chebyshev', 'euclidean', 'manhattan'
    ksg_tree_type: str = 'ball_tree'  # 'ball_tree', 'kd_tree', 'brute'
    ksg_vectorized: bool = True
    ksg_parallel: bool = True
    ksg_n_jobs: int = -1
    
    # Numerical stability
    ksg_epsilon_adjustment: float = 1e-10
    ksg_min_samples_leaf: int = 1
    ksg_regularization: float = 0.0
    
    # ═══════════════════════════════════════════════════════════════════════════
    # ADAPTIVE BINNING CONFIGURATION (Solution B)
    # ═══════════════════════════════════════════════════════════════════════════
    
    binning_strategy: str = 'freedman_diaconis'  # 'fd', 'scott', 'sturges', 'sqrt', 'auto'
    binning_min_bins: int = 10
    binning_max_bins: int = 100
    binning_adaptive: bool = True
    binning_equalization: str = 'none'  # 'none', 'quantile', 'uniform'
    
    # Advanced binning options
    binning_dimension_reduction: bool = True
    binning_pca_components: Optional[int] = None  # Auto-select if None
    binning_smoothing: bool = False
    binning_smoothing_sigma: float = 1.0
    
    # ═══════════════════════════════════════════════════════════════════════════
    # MINE NEURAL ESTIMATION CONFIGURATION (Solution C)
    # ═══════════════════════════════════════════════════════════════════════════
    
    mine_hidden_layers: List[int] = field(default_factory=lambda: [100, 100])
    mine_activation: str = 'relu'  # 'relu', 'tanh', 'elu', 'leaky_relu'
    mine_learning_rate: float = 1e-3
    mine_epochs: int = 100
    mine_batch_size: int = 512
    mine_optimizer: str = 'adam'  # 'adam', 'sgd', 'rmsprop'
    
    # MINE regularization
    mine_dropout_rate: float = 0.0
    mine_weight_decay: float = 0.0
    mine_early_stopping: bool = True
    mine_patience: int = 10
    mine_device: str = 'auto'  # 'auto', 'cpu', 'cuda'
    
    # MINE stability
    mine_gradient_clipping: bool = True
    mine_clip_value: float = 1.0
    mine_initialization: str = 'xavier'  # 'xavier', 'kaiming', 'normal'
    
    # ═══════════════════════════════════════════════════════════════════════════
    # HYBRID ENSEMBLE CONFIGURATION (Solution E)
    # ═══════════════════════════════════════════════════════════════════════════
    
    ensemble_methods: List[MIEstimationMethod] = field(default_factory=lambda: [
        MIEstimationMethod.KSG_EFFICIENT,
        MIEstimationMethod.ADAPTIVE_BINNING,
        MIEstimationMethod.HISTOGRAM_BASIC
    ])
    
    ensemble_weights: Optional[Dict[str, float]] = None  # Auto-weight if None
    ensemble_combination: str = 'weighted_average'  # 'weighted_average', 'median', 'max', 'voting'
    ensemble_confidence_weighting: bool = True
    ensemble_cross_validation: bool = True
    ensemble_cv_folds: int = 3
    
    # ═══════════════════════════════════════════════════════════════════════════
    # AUTOMATIC METHOD SELECTION (Solution D)
    # ═══════════════════════════════════════════════════════════════════════════
    
    auto_selection_rules: Dict[str, Any] = field(default_factory=lambda: {
        'small_data_threshold': 1000,
        'large_data_threshold': 10000,
        'discrete_unique_threshold': 50,
        'high_dim_threshold': 20,
        'memory_limit_gb': 4.0
    })
    
    auto_benchmark_methods: bool = True
    auto_benchmark_sample_size: int = 500
    auto_benchmark_timeout: float = 30.0  # seconds
    
    # ═══════════════════════════════════════════════════════════════════════════
    # OPTIMIZATION AND PERFORMANCE
    # ═══════════════════════════════════════════════════════════════════════════
    
    optimization_strategy: OptimizationStrategy = OptimizationStrategy.BALANCED
    max_computation_time: Optional[float] = None  # seconds, None = unlimited
    max_memory_usage: Optional[float] = None  # GB, None = unlimited
    use_gpu: bool = True
    gpu_device_id: int = 0
    
    # Parallel processing
    enable_parallel: bool = True
    max_workers: Optional[int] = None  # None = auto-detect
    chunk_size: str = 'auto'  # 'auto', integer, or fraction
    
    # ═══════════════════════════════════════════════════════════════════════════
    # NUMERICAL STABILITY AND VALIDATION
    # ═══════════════════════════════════════════════════════════════════════════
    
    numerical_stability: bool = True
    min_samples_warning: int = 50
    max_samples_auto_downsample: int = 100000
    downsample_method: str = 'random'  # 'random', 'stratified', 'systematic'
    
    # Validation and diagnostics
    validate_inputs: bool = True
    compute_confidence_intervals: bool = False
    confidence_level: float = 0.95
    bootstrap_samples: int = 100
    
    # Output formatting
    output_units: str = 'bits'  # 'bits', 'nats', 'dits'
    precision_digits: int = 6
    return_diagnostics: bool = True
    
    # ═══════════════════════════════════════════════════════════════════════════
    # BACKWARD COMPATIBILITY AND LEGACY SUPPORT
    # ═══════════════════════════════════════════════════════════════════════════
    
    preserve_legacy_api: bool = True
    legacy_compatibility_warnings: bool = True
    fallback_to_legacy_on_error: bool = True
    
    def validate_config(self) -> Dict[str, Any]:
        """Validate configuration settings and return diagnostic info"""
        issues = []
        warnings = []
        
        # Validate KSG parameters
        if self.ksg_k_neighbors < 1:
            issues.append("ksg_k_neighbors must be positive")
        
        if self.ksg_metric not in ['chebyshev', 'euclidean', 'manhattan']:
            issues.append(f"Invalid ksg_metric: {self.ksg_metric}")
        
        # Validate binning parameters
        if self.binning_min_bins >= self.binning_max_bins:
            issues.append("binning_min_bins must be < binning_max_bins")
        
        # Validate MINE parameters
        if self.mine_learning_rate <= 0 or self.mine_learning_rate >= 1:
            issues.append("mine_learning_rate must be in (0, 1)")
        
        if not self.mine_hidden_layers:
            warnings.append("Empty mine_hidden_layers may lead to poor performance")
        
        # Validate ensemble parameters
        if (self.method == MIEstimationMethod.HYBRID_ENSEMBLE and 
            len(self.ensemble_methods) < 2):
            issues.append("Ensemble requires at least 2 methods")
        
        # Check for conflicts
        if (self.optimization_strategy == OptimizationStrategy.MEMORY_EFFICIENT and
            self.method == MIEstimationMethod.MINE_NEURAL):
            warnings.append("MINE method may conflict with memory_efficient strategy")
        
        return {
            'valid': len(issues) == 0,
            'issues': issues,
            'warnings': warnings,
            'method_summary': self._get_method_summary()
        }
    
    def _get_method_summary(self) -> Dict[str, Any]:
        """Get summary of selected methods and their research basis"""
        method_info = {
            MIEstimationMethod.KSG_EFFICIENT: {
                'description': 'Vectorized KSG estimator O(n log n)',
                'research_basis': 'Kraskov, Stögbauer & Grassberger (2004)',
                'best_for': 'Small to medium continuous data',
                'complexity': 'O(n log n)'
            },
            MIEstimationMethod.KSG_LEGACY: {
                'description': 'Original KSG implementation (preserved)',
                'research_basis': 'Kraskov, Stögbauer & Grassberger (2004)',
                'best_for': 'Small datasets only (n < 1000)',
                'complexity': 'O(n³) - USE WITH CAUTION'
            },
            MIEstimationMethod.ADAPTIVE_BINNING: {
                'description': 'Adaptive histogram with optimal binning',
                'research_basis': 'Darbellay & Vajda (1999)',
                'best_for': 'Medium datasets, mixed data types',
                'complexity': 'O(n log n)'
            },
            MIEstimationMethod.MINE_NEURAL: {
                'description': 'Neural mutual information estimation',
                'research_basis': 'Belghazi et al. (2018) ICML',
                'best_for': 'Large datasets (n > 10000)',
                'complexity': 'O(n) per epoch'
            },
            MIEstimationMethod.HISTOGRAM_BASIC: {
                'description': 'Simple histogram method',
                'research_basis': 'Shannon (1948) classical approach',
                'best_for': 'Fast approximation, discrete data',
                'complexity': 'O(n)'
            },
            MIEstimationMethod.HYBRID_ENSEMBLE: {
                'description': 'Weighted combination of multiple methods',
                'research_basis': 'Ensemble learning principles',
                'best_for': 'Maximum accuracy, unknown data characteristics',
                'complexity': 'Sum of component complexities'
            }
        }
        
        selected_info = method_info.get(self.method, {})
        selected_info['selected_method'] = self.method.value
        
        if self.method == MIEstimationMethod.HYBRID_ENSEMBLE:
            selected_info['ensemble_components'] = [
                method_info.get(m, {}).get('description', str(m))
                for m in self.ensemble_methods
            ]
        
        return selected_info
    
    def get_recommended_config_for_data(self, n_samples: int, 
                                       n_features_x: int, n_features_y: int,
                                       data_type: str = 'continuous') -> 'MIEstimationConfig':
        """
        Get recommended configuration based on data characteristics
        
        Args:
            n_samples: Number of data samples
            n_features_x: Number of features in X
            n_features_y: Number of features in Y  
            data_type: 'continuous', 'discrete', or 'mixed'
        """
        recommended = MIEstimationConfig()
        
        # Determine data characteristics
        is_small = n_samples < self.auto_selection_rules['small_data_threshold']
        is_large = n_samples > self.auto_selection_rules['large_data_threshold']
        is_high_dim = (n_features_x + n_features_y) > self.auto_selection_rules['high_dim_threshold']
        
        # Automatic method selection based on data characteristics
        if data_type == 'discrete':
            recommended.method = MIEstimationMethod.ADAPTIVE_BINNING
        elif is_small:
            recommended.method = MIEstimationMethod.KSG_EFFICIENT
        elif is_large:
            recommended.method = MIEstimationMethod.MINE_NEURAL
            recommended.mine_epochs = min(50, 5000 // (n_samples // 1000))  # Scale epochs
        elif is_high_dim:
            recommended.method = MIEstimationMethod.ADAPTIVE_BINNING
            recommended.binning_dimension_reduction = True
        else:
            recommended.method = MIEstimationMethod.HYBRID_ENSEMBLE
        
        # Adjust parameters based on size
        if is_large:
            recommended.optimization_strategy = OptimizationStrategy.SPEED_FIRST
            recommended.ksg_k_neighbors = min(5, max(3, n_samples // 1000))
        elif is_small:
            recommended.optimization_strategy = OptimizationStrategy.ACCURACY_FIRST
            recommended.ksg_k_neighbors = min(10, n_samples // 10)
        
        return recommended


# ═══════════════════════════════════════════════════════════════════════════════
# FACTORY FUNCTIONS FOR COMMON CONFIGURATIONS
# ═══════════════════════════════════════════════════════════════════════════════

def create_accurate_config() -> MIEstimationConfig:
    """
    Create configuration that prioritizes accuracy over speed
    
    Uses theoretically sound methods with proper parameter settings
    """
    return MIEstimationConfig(
        method=MIEstimationMethod.KSG_EFFICIENT,
        ksg_k_neighbors=5,  # Higher k for better accuracy
        ksg_vectorized=True,
        optimization_strategy=OptimizationStrategy.ACCURACY_FIRST,
        numerical_stability=True,
        compute_confidence_intervals=True,
        return_diagnostics=True
    )


def create_high_performance_config() -> MIEstimationConfig:
    """
    Create configuration optimized for speed and large datasets
    """
    return MIEstimationConfig(
        method=MIEstimationMethod.AUTO_SELECT,
        optimization_strategy=OptimizationStrategy.SPEED_FIRST,
        ksg_parallel=True,
        mine_batch_size=1024,
        mine_epochs=50,
        max_computation_time=60.0,  # 1 minute limit
        enable_parallel=True
    )


def create_legacy_compatible_config() -> MIEstimationConfig:
    """
    Create configuration that maintains backward compatibility
    """
    return MIEstimationConfig(
        method=MIEstimationMethod.KSG_LEGACY,  # Use original method
        preserve_legacy_api=True,
        fallback_to_legacy_on_error=True,
        legacy_compatibility_warnings=True
    )


def create_ensemble_config() -> MIEstimationConfig:
    """
    Create configuration using ensemble of multiple methods
    """
    return MIEstimationConfig(
        method=MIEstimationMethod.HYBRID_ENSEMBLE,
        ensemble_methods=[
            MIEstimationMethod.KSG_EFFICIENT,
            MIEstimationMethod.ADAPTIVE_BINNING,
            MIEstimationMethod.MINE_NEURAL
        ],
        ensemble_combination='weighted_average',
        ensemble_cross_validation=True,
        optimization_strategy=OptimizationStrategy.BALANCED
    )


def create_gpu_accelerated_config() -> MIEstimationConfig:
    """
    Create configuration optimized for GPU acceleration
    """
    return MIEstimationConfig(
        method=MIEstimationMethod.MINE_NEURAL,
        optimization_strategy=OptimizationStrategy.GPU_ACCELERATED,
        use_gpu=True,
        mine_device='cuda',
        mine_batch_size=2048,  # Larger batches for GPU
        mine_epochs=200,
        enable_parallel=True
    )