"""
📋 Ib Algorithms
=================

🎯 ELI5 Summary:
This is the brain of our operation! Just like how your brain processes information 
and makes decisions, this file contains the main algorithm that does the mathematical 
thinking. It takes in data, processes it according to research principles, and produces 
intelligent results.

🧪 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.

📋 Component Integration:
========================
    ┌──────────┐
    │   This   │
    │Component │ ←→ Other Components
    └──────────┘
         ↑↓
    System Integration

"""
"""
Core Information Bottleneck algorithms and optimization methods
Based on: Tishby, Pereira & Bialek (1999) "The Information Bottleneck Method"
"""

import numpy as np
from typing import Dict, Tuple, Optional, List
from sklearn.cluster import KMeans
from mutual_info_core import MutualInfoCore
from ib_config import IBConfig, EncoderUpdateMethod, DecoderUpdateMethod, InitializationMethod


class IBAlgorithms:
    """Core algorithms for Information Bottleneck optimization"""
    
    def __init__(self, config: IBConfig):
        self.config = config
        self.mi_estimator = MutualInfoCore()
        
        # Algorithm state
        self.p_z_given_x = None  # Encoder p(z|x)
        self.p_y_given_z = None  # Decoder p(y|z)  
        self.p_z = None         # Prior p(z)
        
        # Training history
        self.history = {
            'compression': [],
            'relevance': [],
            'objective': [],
            'beta_values': []
        } if config.save_history else None
    
    def initialize_distributions(self, X: np.ndarray, Y: np.ndarray) -> None:
        """
        Initialize encoder and decoder distributions
        
        Args:
            X: Input data [n_samples, n_features]
            Y: Target data [n_samples, n_targets]
        """
        n_samples = X.shape[0]
        n_clusters = self.config.n_clusters
        
        if self.config.initialization == InitializationMethod.RANDOM:
            # Random initialization
            self.p_z_given_x = np.random.dirichlet(np.ones(n_clusters), n_samples)
            
        elif self.config.initialization == InitializationMethod.KMEANS_PLUS_PLUS:
            self._kmeans_plus_plus_initialization(X, Y)
            
        elif self.config.initialization == InitializationMethod.MUTUAL_INFO:
            self._mutual_info_initialization(X, Y)
            
        elif self.config.initialization == InitializationMethod.HIERARCHICAL:
            self._hierarchical_ib_initialization(X, Y)
        
        # Initialize decoder using Bayes' rule
        self._update_decoder_bayes_rule(X, Y)
        
        # Update prior
        self.p_z = np.mean(self.p_z_given_x, axis=0)
    
    def _kmeans_plus_plus_initialization(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Initialize using k-means++ clustering"""
        # Use both X and Y for clustering to get relevant clusters
        if Y.ndim == 1:
            Y_features = Y.reshape(-1, 1)
        else:
            Y_features = Y
            
        # Combine features (weighted towards Y since it's the target)
        # Add epsilon to avoid division by zero
        eps = 1e-8
        X_std = np.std(X, axis=0, keepdims=True) + eps
        Y_std = np.std(Y_features, axis=0, keepdims=True) + eps
        
        combined_features = np.hstack([
            X / X_std,  # Standardize X
            3 * Y_features / Y_std  # Weight Y higher
        ])
        
        # K-means clustering
        kmeans = KMeans(
            n_clusters=self.config.n_clusters,
            init='k-means++',
            n_init=self.config.n_init,
            random_state=self.config.random_state
        )
        cluster_labels = kmeans.fit_predict(combined_features)
        
        # Convert to soft assignment with temperature
        temperature = 0.5
        distances = kmeans.transform(combined_features)
        soft_assignments = np.exp(-distances / temperature)
        self.p_z_given_x = soft_assignments / np.sum(soft_assignments, axis=1, keepdims=True)
    
    def _mutual_info_initialization(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Initialize based on mutual information between features and targets"""
        n_samples, n_features = X.shape
        n_clusters = self.config.n_clusters
        
        # Compute MI between each feature and target
        feature_mi = np.zeros(n_features)
        for i in range(n_features):
            feature_mi[i] = self.mi_estimator.estimate_mutual_info_continuous(
                X[:, i:i+1], Y, method='binning'
            )
        
        # Select most informative features
        top_features = np.argsort(feature_mi)[-min(5, n_features):]
        X_selected = X[:, top_features]
        
        # Cluster using informative features
        kmeans = KMeans(n_clusters=n_clusters, random_state=self.config.random_state)
        cluster_labels = kmeans.fit_predict(X_selected)
        
        # Convert to soft assignment
        self.p_z_given_x = np.zeros((n_samples, n_clusters))
        for i, label in enumerate(cluster_labels):
            self.p_z_given_x[i, label] = 1.0
        
        # Add small noise for better convergence
        noise = np.random.dirichlet(0.1 * np.ones(n_clusters), n_samples)
        self.p_z_given_x = 0.9 * self.p_z_given_x + 0.1 * noise
    
    def _hierarchical_ib_initialization(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Hierarchical initialization for better local optima"""
        n_samples = X.shape[0]
        n_clusters = self.config.n_clusters
        
        # Start with 2 clusters and progressively split
        current_clusters = 2
        self.p_z_given_x = np.zeros((n_samples, 2))
        
        # Initial 2-cluster split using k-means
        kmeans = KMeans(n_clusters=2, random_state=self.config.random_state)
        labels = kmeans.fit_predict(X)
        for i, label in enumerate(labels):
            self.p_z_given_x[i, label] = 1.0
        
        # Progressively split clusters
        while current_clusters < n_clusters:
            # Find cluster with highest internal variance
            cluster_variances = np.zeros(current_clusters)
            for z in range(current_clusters):
                cluster_mask = self.p_z_given_x[:, z] > 0.5
                if np.sum(cluster_mask) > 1:
                    cluster_data = X[cluster_mask]
                    cluster_variances[z] = np.trace(np.cov(cluster_data.T))
            
            # Split the cluster with highest variance
            split_cluster = np.argmax(cluster_variances)
            cluster_mask = self.p_z_given_x[:, split_cluster] > 0.5
            
            if np.sum(cluster_mask) > 2:  # Need at least 2 points to split
                cluster_data = X[cluster_mask]
                sub_kmeans = KMeans(n_clusters=2, random_state=self.config.random_state)
                sub_labels = sub_kmeans.fit_predict(cluster_data)
                
                # Update assignments
                new_p_z = np.zeros((n_samples, current_clusters + 1))
                new_p_z[:, :split_cluster] = self.p_z_given_x[:, :split_cluster]
                new_p_z[:, split_cluster+1:current_clusters] = self.p_z_given_x[:, split_cluster+1:current_clusters]
                
                # Split the selected cluster
                cluster_indices = np.where(cluster_mask)[0]
                for i, idx in enumerate(cluster_indices):
                    if sub_labels[i] == 0:
                        new_p_z[idx, split_cluster] = 1.0
                    else:
                        new_p_z[idx, current_clusters] = 1.0
                
                self.p_z_given_x = new_p_z
                current_clusters += 1
            else:
                # Can't split further, add random cluster
                new_p_z = np.zeros((n_samples, current_clusters + 1))
                new_p_z[:, :-1] = self.p_z_given_x
                # Assign some points randomly to new cluster
                random_indices = np.random.choice(n_samples, size=max(1, n_samples//10), replace=False)
                for idx in random_indices:
                    new_p_z[idx, :-1] = 0.1 * new_p_z[idx, :-1]
                    new_p_z[idx, -1] = 0.9
                    new_p_z[idx] /= np.sum(new_p_z[idx])
                
                self.p_z_given_x = new_p_z
                current_clusters += 1
        
        # Ensure we have exactly n_clusters
        if current_clusters > n_clusters:
            # Merge similar clusters
            self.p_z_given_x = self.p_z_given_x[:, :n_clusters]
            self.p_z_given_x /= np.sum(self.p_z_given_x, axis=1, keepdims=True)
    
    def update_encoder(self, X: np.ndarray, Y: np.ndarray, temperature: float = 1.0) -> None:
        """
        Update encoder distribution p(z|x)
        
        Args:
            X: Input data
            Y: Target data
            temperature: Temperature parameter for annealing
        """
        if self.config.encoder_update_method == EncoderUpdateMethod.BLAHUT_ARIMOTO:
            self._blahut_arimoto_update(X, Y, temperature)
        elif self.config.encoder_update_method == EncoderUpdateMethod.NATURAL_GRADIENT:
            self._natural_gradient_update(X, Y, temperature)
        elif self.config.encoder_update_method == EncoderUpdateMethod.TEMPERATURE_SCALED:
            self._temperature_scaled_update(X, Y, temperature)
        elif self.config.encoder_update_method == EncoderUpdateMethod.DETERMINISTIC_ANNEALING:
            self._deterministic_annealing_update(X, Y, temperature)
    
    def _blahut_arimoto_update(self, X: np.ndarray, Y: np.ndarray, temperature: float = 1.0) -> None:
        """Standard Blahut-Arimoto algorithm update"""
        n_samples = X.shape[0]
        n_clusters = self.config.n_clusters
        
        # Update encoder using self-consistent equations
        new_p_z_given_x = np.zeros((n_samples, n_clusters))
        
        for i in range(n_samples):
            for z in range(n_clusters):
                # Compute KL divergence D(p(y|x_i) || p(y|z))
                kl_div = self._compute_exact_kl_divergence(i, z)
                
                # Self-consistent update
                new_p_z_given_x[i, z] = self.p_z[z] * np.exp(-self.config.beta * kl_div / temperature)
        
        # Normalize
        for i in range(n_samples):
            Z_i = np.sum(new_p_z_given_x[i, :])
            if Z_i > 1e-12:
                new_p_z_given_x[i, :] /= Z_i
            else:
                # If all probabilities are zero, use uniform distribution
                new_p_z_given_x[i, :] = 1.0 / n_clusters
        
        self.p_z_given_x = new_p_z_given_x
        
        # Update prior
        self.p_z = np.mean(self.p_z_given_x, axis=0)
    
    def _compute_exact_kl_divergence(self, x_idx: int, z: int) -> float:
        """Compute KL divergence D(p(y|x_i) || p(y|z))"""
        # For discrete case, this requires knowing p(y|x_i)
        # In practice, we approximate this using the current decoder
        kl_div = 0.0
        
        # Simple approximation: use squared difference in predictions
        # This is not exact KL but provides the right optimization direction
        y_given_x = self.p_y_given_z[z, :]  # Current prediction for z
        
        # Compute actual KL divergence D_KL(p(y|z) || p(y)) based on Tishby et al. (1999)
        # KL divergence: sum_y p(y|z) * log(p(y|z) / p(y))
        
        # Get marginal distribution p(y) 
        p_y = np.sum(self.p_y_given_z * self.p_z[:, np.newaxis], axis=0)
        p_y = np.maximum(p_y, 1e-12)  # Avoid log(0)
        
        # Ensure p(y|z) is normalized and positive
        y_given_z_normalized = np.maximum(y_given_x, 1e-12)
        y_given_z_normalized = y_given_z_normalized / np.sum(y_given_z_normalized)
        
        # Compute KL divergence: D_KL(p(y|z) || p(y))
        kl_div = 0.0
        for y_idx in range(len(y_given_z_normalized)):
            if y_given_z_normalized[y_idx] > 1e-12:
                kl_div += y_given_z_normalized[y_idx] * np.log(y_given_z_normalized[y_idx] / p_y[y_idx])
        
        return kl_div
    
    def _natural_gradient_update(self, X: np.ndarray, Y: np.ndarray, temperature: float = 1.0) -> None:
        """Natural gradient update for encoder"""
        # Compute gradient of IB objective
        gradient = self._compute_ib_gradient(X, Y, temperature)
        
        # Compute Fisher information metric (approximation)
        fisher_diag = self._compute_fisher_diagonal()
        
        # Natural gradient step
        lr = 0.01 / temperature  # Adaptive learning rate
        update = lr * gradient / (fisher_diag + 1e-8)
        
        # Update parameters in probability simplex
        self.p_z_given_x += update
        
        # Project back to probability simplex
        self._project_to_simplex()
        
        # Update prior
        self.p_z = np.mean(self.p_z_given_x, axis=0)
    
    def _compute_ib_gradient(self, X: np.ndarray, Y: np.ndarray, temperature: float) -> np.ndarray:
        """Compute gradient of IB objective w.r.t. encoder parameters"""
        n_samples, n_clusters = self.p_z_given_x.shape
        gradient = np.zeros((n_samples, n_clusters))
        
        for i in range(n_samples):
            for z in range(n_clusters):
                # Compression term gradient
                compression_grad = np.log(self.p_z_given_x[i, z] / self.p_z[z])
                
                # Relevance term gradient (approximation)
                relevance_grad = -np.sum(Y[i] * np.log(self.p_y_given_z[z, :] + 1e-12))
                
                gradient[i, z] = compression_grad - self.config.beta * relevance_grad
        
        return gradient
    
    def _compute_fisher_diagonal(self) -> np.ndarray:
        """Compute diagonal approximation of Fisher information matrix"""
        fisher = np.zeros_like(self.p_z_given_x)
        
        for i in range(self.p_z_given_x.shape[0]):
            for z in range(self.p_z_given_x.shape[1]):
                # Diagonal Fisher information
                fisher[i, z] = 1.0 / (self.p_z_given_x[i, z] + 1e-8)
        
        return fisher
    
    def _project_to_simplex(self) -> None:
        """Project encoder probabilities to probability simplex"""
        for i in range(self.p_z_given_x.shape[0]):
            # Ensure non-negative
            self.p_z_given_x[i, :] = np.maximum(self.p_z_given_x[i, :], 1e-12)
            
            # Normalize to sum to 1
            row_sum = np.sum(self.p_z_given_x[i, :])
            if row_sum > 1e-12:
                self.p_z_given_x[i, :] /= row_sum
            else:
                # If all zeros, use uniform distribution
                self.p_z_given_x[i, :] = 1.0 / self.p_z_given_x.shape[1]
    
    def _temperature_scaled_update(self, X: np.ndarray, Y: np.ndarray, temperature: float) -> None:
        """Temperature-scaled soft assignment update"""
        n_samples = X.shape[0]
        n_clusters = self.config.n_clusters
        
        # Compute assignment scores
        scores = np.zeros((n_samples, n_clusters))
        
        for i in range(n_samples):
            for z in range(n_clusters):
                # Score based on relevance to target Y
                if Y.ndim == 1:
                    target_similarity = -np.abs(Y[i] - np.dot(self.p_y_given_z[z, :], np.arange(self.p_y_given_z.shape[1])))
                else:
                    target_similarity = -np.sum((Y[i] - self.p_y_given_z[z, :])**2)
                
                # Prior preference  
                prior_score = np.log(self.p_z[z] + 1e-12)
                
                scores[i, z] = self.config.beta * target_similarity + prior_score
        
        # Soft assignment with temperature
        self.p_z_given_x = np.exp(scores / temperature)
        
        # Normalize
        for i in range(n_samples):
            row_sum = np.sum(self.p_z_given_x[i, :])
            if row_sum > 1e-12:
                self.p_z_given_x[i, :] /= row_sum
            else:
                self.p_z_given_x[i, :] = 1.0 / n_clusters
        
        # Update prior
        self.p_z = np.mean(self.p_z_given_x, axis=0)
    
    def _deterministic_annealing_update(self, X: np.ndarray, Y: np.ndarray, 
                                      temperature: float, phase: str = 'cooling') -> None:
        """Deterministic annealing with phase transitions"""
        if phase == 'cooling':
            # Standard cooling phase
            self._blahut_arimoto_update(X, Y, temperature)
        elif phase == 'heating':
            # Heating phase to escape local minima
            noise_level = 1.0 / temperature
            noise = np.random.dirichlet(noise_level * np.ones(self.config.n_clusters), X.shape[0])
            self.p_z_given_x = 0.8 * self.p_z_given_x + 0.2 * noise
            self.p_z = np.mean(self.p_z_given_x, axis=0)
    
    def update_decoder(self, X: np.ndarray, Y: np.ndarray) -> None:
        """
        Update decoder distribution p(y|z)
        
        Args:
            X: Input data
            Y: Target data
        """
        if self.config.decoder_update_method == DecoderUpdateMethod.BAYES_RULE:
            self._update_decoder_bayes_rule(X, Y)
        elif self.config.decoder_update_method == DecoderUpdateMethod.EM:
            self._em_decoder_update(X, Y)
        elif self.config.decoder_update_method == DecoderUpdateMethod.REGULARIZED:
            self._regularized_decoder_update(X, Y)
    
    def _update_decoder_bayes_rule(self, X: np.ndarray, Y: np.ndarray) -> None:
        """Update decoder using Bayes' rule: p(y|z) = Σ_x p(y|x)p(x|z)"""
        n_clusters = self.config.n_clusters
        
        if Y.ndim == 1:
            # Discrete Y
            n_y_values = len(np.unique(Y))
            self.p_y_given_z = np.zeros((n_clusters, n_y_values))
            
            for z in range(n_clusters):
                for y_val in range(n_y_values):
                    # Find samples with this y value
                    y_indices = np.where(Y == y_val)[0]
                    
                    # Weighted average using p(z|x) 
                    if len(y_indices) > 0:
                        weight_sum = np.sum(self.p_z_given_x[y_indices, z])
                        if weight_sum > 1e-12:
                            self.p_y_given_z[z, y_val] = weight_sum
                
                # Normalize to probability distribution
                row_sum = np.sum(self.p_y_given_z[z, :])
                if row_sum > 1e-12:
                    self.p_y_given_z[z, :] /= row_sum
                else:
                    self.p_y_given_z[z, :] = 1.0 / n_y_values
        else:
            # Continuous Y - use weighted average
            n_y_dim = Y.shape[1] 
            self.p_y_given_z = np.zeros((n_clusters, n_y_dim))
            
            for z in range(n_clusters):
                # Compute weighted average
                weights = self.p_z_given_x[:, z]
                weight_sum = np.sum(weights)
                
                if weight_sum > 1e-12:
                    self.p_y_given_z[z, :] = np.average(Y, axis=0, weights=weights)
                else:
                    self.p_y_given_z[z, :] = np.mean(Y, axis=0)
    
    def _em_decoder_update(self, X: np.ndarray, Y: np.ndarray) -> None:
        """EM-style decoder update with soft assignments"""
        # Similar to Bayes rule but with additional smoothing
        self._update_decoder_bayes_rule(X, Y)
        
        # Add Laplace smoothing for discrete case
        if Y.ndim == 1:
            self.p_y_given_z += self.config.decoder_regularization
            self.p_y_given_z /= np.sum(self.p_y_given_z, axis=1, keepdims=True)
    
    def _regularized_decoder_update(self, X: np.ndarray, Y: np.ndarray, alpha: float = 0.1) -> None:
        """Regularized decoder update with L2 penalty"""
        # Start with Bayes rule update
        self._update_decoder_bayes_rule(X, Y)
        
        # Add regularization towards uniform/mean
        if Y.ndim == 1:
            # Discrete case: regularize towards uniform
            n_y_values = self.p_y_given_z.shape[1]
            uniform = np.ones(n_y_values) / n_y_values
            self.p_y_given_z = (1 - alpha) * self.p_y_given_z + alpha * uniform
        else:
            # Continuous case: regularize towards global mean
            global_mean = np.mean(Y, axis=0)
            self.p_y_given_z = (1 - alpha) * self.p_y_given_z + alpha * global_mean
    
    def compute_ib_objective(self, X: np.ndarray, Y: np.ndarray) -> Dict[str, float]:
        """
        Compute Information Bottleneck objective and its components
        
        Args:
            X: Input data
            Y: Target data
            
        Returns:
            Dictionary with objective components
        """
        # Compression term: I(X;Z)
        compression = self._compute_compression_term(X)
        
        # Relevance term: I(Z;Y) 
        relevance = self._compute_relevance_term(Y)
        
        # IB objective: I(X;Z) - β*I(Z;Y)
        objective = compression - self.config.beta * relevance
        
        results = {
            'compression': compression,
            'relevance': relevance, 
            'objective': objective,
            'beta': self.config.beta
        }
        
        # Save to history if enabled
        if self.history is not None:
            self.history['compression'].append(compression)
            self.history['relevance'].append(relevance)
            self.history['objective'].append(objective)
            self.history['beta_values'].append(self.config.beta)
        
        return results
    
    def _compute_compression_term(self, X: np.ndarray) -> float:
        """Compute compression term I(X;Z)"""
        # Use current encoder distribution
        compression = 0.0
        n_samples = X.shape[0]
        
        for i in range(n_samples):
            for z in range(self.config.n_clusters):
                if self.p_z_given_x[i, z] > 1e-12 and self.p_z[z] > 1e-12:
                    compression += (1.0/n_samples) * self.p_z_given_x[i, z] * \
                                  np.log(self.p_z_given_x[i, z] / self.p_z[z])
        
        return max(0.0, compression)
    
    def _compute_relevance_term(self, Y: np.ndarray) -> float:
        """Compute relevance term I(Z;Y)"""
        # Approximate using decoder distribution
        if self.p_y_given_z is None:
            return 0.0
            
        relevance = 0.0
        n_samples = Y.shape[0]
        
        if Y.ndim == 1:
            # Discrete Y case
            y_values = np.unique(Y)
            p_y = np.bincount(Y.astype(int)) / len(Y)
            
            for z in range(self.config.n_clusters):
                if self.p_z[z] > 1e-12:
                    for y_idx, y_val in enumerate(y_values):
                        if self.p_y_given_z[z, y_idx] > 1e-12 and p_y[y_idx] > 1e-12:
                            relevance += self.p_z[z] * self.p_y_given_z[z, y_idx] * \
                                       np.log(self.p_y_given_z[z, y_idx] / p_y[y_idx])
        else:
            # Continuous Y - use differential entropy approximation
            # This is a simplified approximation
            for z in range(self.config.n_clusters):
                if self.p_z[z] > 1e-12:
                    # Estimate entropy of p(y|z)
                    cluster_mask = self.p_z_given_x[:, z] > 0.1
                    if np.sum(cluster_mask) > 1:
                        y_cluster = Y[cluster_mask]
                        cluster_entropy = -0.5 * np.log(2 * np.pi * np.e * np.var(y_cluster, axis=0))
                        relevance += self.p_z[z] * np.mean(cluster_entropy)
        
        return max(0.0, relevance)