"""
🧠 Information Bottleneck - Core Module
========================================

🎯 ELI5 EXPLANATION:
==================
Imagine your brain trying to understand a movie while being bombarded with millions of details!

Your brain is incredibly smart - it doesn't try to remember every single pixel or sound wave. Instead, it finds the perfect "compression" that keeps all the important stuff (the plot, characters, emotions) while throwing away useless details (background noise, lighting fluctuations).

The Information Bottleneck principle does exactly this for AI:
1. 📥 **Input**: Complex, noisy data (images, text, sensors)
2. 🎯 **Bottleneck**: Find the minimal representation that preserves what matters
3. 📤 **Output**: Perfect predictions using only essential information
4. 🧠 **Magic**: Discovers the same information processing principles your brain uses!

🔬 RESEARCH FOUNDATION:
======================
Implements Naftali Tishby's revolutionary information-theoretic learning principle:
- Tishby, Pereira & Bialek (1999): "The Information Bottleneck Method"
- Tishby & Zaslavsky (2015): "Deep Learning and the Information Bottleneck Principle"  
- Alemi et al. (2017): "Deep Variational Information Bottleneck"
- Schwartz-Ziv & Tishby (2017): "Opening the Black Box of Deep Neural Networks"

🧮 MATHEMATICAL PRINCIPLES:
==========================
**Core Principle - Optimal Information Trade-off:**
L_IB = I(X;Z) - βI(Z;Y)

Where:
• I(X;Z) = Compression term (minimize information about input)
• I(Z;Y) = Prediction term (maximize relevant information about output)
• β = Lagrange multiplier controlling the trade-off
• Z = Compressed representation (the "bottleneck")

**Information Measures:**
I(X;Y) = ∫∫ p(x,y) log[p(x,y)/(p(x)p(y))] dx dy

**Deterministic Annealing Schedule:**
β(t) = β_min × (β_max/β_min)^(t/T)

📊 ARCHITECTURE VISUALIZATION:
==============================
```
🔬 INFORMATION BOTTLENECK PRINCIPLE 🔬

High-Dimensional Input     Information Bottleneck      Optimal Predictions
┌─────────────────────┐    ┌─────────────────────────┐  ┌───────────────────┐
│  🌍 Raw Input X     │    │    🎯 BOTTLENECK Z      │  │  ✨ Predictions Y │
│                     │    │                         │  │                   │
│  📸 Images: 1M dims │───→│  🧠 Compress to ~100    │─→│  🏷️ Classes: 10   │
│  📝 Text: 10K words │    │                         │  │                   │
│  🔊 Audio: 100K Hz  │    │   I(X;Z) ← minimize     │  │   I(Z;Y) ← maximize │
│  🎥 Video: 1B pixels│    │   (Remove noise)        │  │   (Keep signal)   │
└─────────────────────┘    └─────────────────────────┘  └───────────────────┘
         ↓                            ↓                           ↓
    Noisy, complex              Perfect balance            Optimal decisions
    redundant data              compression ↔ prediction    using minimal info

🔄 LEARNING DYNAMICS:
   β → 0:  Z ≈ X (no compression, overfitting)
   β → ∞:  Z ≈ constant (too much compression, underfitting)  
   β* optimal: Z contains exactly the right information for Y
```

💰 SUPPORT THIS RESEARCH - PLEASE DONATE! 💰

🙏 If this library helps your research or project, please consider supporting:
💳 PayPal: https://www.paypal.com/cgi-bin/webscr?cmd=_s-xclick&hosted_button_id=WXQKYYKPHWXHS
⭐ GitHub Sponsors: https://github.com/sponsors/benedictchen

Your support enables cutting-edge AI research for everyone! 🚀

"""
"""
Information Bottleneck Core Classes - Refactored for 800-line limit
=================================================================

Main algorithm implementations importing from modular mixins.
Refactored from the original 1472-line core.py to meet the 800-line standard.

Author: Benedict Chen (benedict@benedictchen.com)
Based on: Naftali Tishby, Fernando C. Pereira & William Bialek (1999)
"""

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.base import BaseEstimator, ClassifierMixin, TransformerMixin
from sklearn.cluster import KMeans
from sklearn.metrics import mutual_info_score
import warnings
from typing import Optional, Union, Dict, List, Tuple, Any

# Import the modular mixins 
from .mixins import (
    CoreTheoryMixin,
    MutualInformationMixin, 
    OptimizationMixin,
    TransformPredictMixin,
    EvaluationMixin
)

# ============================================================================
# MAIN ALGORITHM IMPLEMENTATIONS
# ============================================================================

class InformationBottleneck(BaseEstimator, TransformerMixin, CoreTheoryMixin, 
                          MutualInformationMixin, OptimizationMixin, 
                          TransformPredictMixin, EvaluationMixin):
    """
    Classical Information Bottleneck Implementation.
    
    Based on: Tishby, Pereira & Bialek (1999)
    "The Information Bottleneck Method"
    
    The Information Bottleneck method finds representations T of data X that are:
    1. Maximally informative about target Y: maximize I(T;Y)
    2. Minimally complex: minimize I(X;T) 
    
    The trade-off is controlled by parameter β in the objective:
    L = I(T;Y) - β*I(X;T)
    
    Parameters
    ----------
    n_clusters : int, default=10
        Number of clusters T (bottleneck size)
    beta : float, default=1.0
        Trade-off parameter between compression and relevance
    max_iter : int, default=100
        Maximum iterations for optimization
    tol : float, default=1e-6
        Convergence tolerance
    method : str, default='blahut_arimoto'
        Optimization method ('blahut_arimoto', 'deterministic_annealing')
    """
    
    def __init__(self, n_clusters=10, beta=1.0, max_iter=100, tol=1e-6, 
                 method='blahut_arimoto', random_state=None):
        self.n_clusters = n_clusters
        self.beta = beta
        self.max_iter = max_iter
        self.tol = tol
        self.method = method
        self.random_state = random_state
    
    def fit(self, X, Y):
        """
        Fit Information Bottleneck model.
        
        Parameters
        ----------
        X : array-like, shape (n_samples,) or (n_samples, n_features)
            Input data
        Y : array-like, shape (n_samples,)
            Target variable
        """
        X = np.asarray(X).flatten() if X.ndim > 1 else np.asarray(X)
        Y = np.asarray(Y).flatten()
        
        if len(X) != len(Y):
            raise ValueError("X and Y must have same number of samples")
        
        self.X_fit_ = X
        self.Y_fit_ = Y
        
        if self.random_state is not None:
            np.random.seed(self.random_state)
        
        if self.method == 'blahut_arimoto':
            self.p_t_given_x_, self.p_y_given_t_, self.objective_ = \
                self._optimize_blahut_arimoto(X, Y, self.n_clusters, self.beta, 
                                            self.max_iter, self.tol)
        elif self.method == 'deterministic_annealing':
            beta_schedule = np.logspace(-2, np.log10(self.beta), 10)
            results = self._deterministic_annealing(X, Y, self.n_clusters, 
                                                  beta_schedule, self.max_iter//10)
            # Use final result
            final_result = results[-1]
            self.p_t_given_x_ = final_result['p_t_given_x']
            self.p_y_given_t_ = final_result['p_y_given_t']
            self.objective_ = final_result['objective']
        else:
            raise ValueError(f"Unknown method: {self.method}")
            
        return self
    
    # Helper methods for optimization
    def _update_conditional_probabilities(self, X, Y, p_t_given_x):
        """Update p(y|t) based on current cluster assignments."""
        unique_x = np.unique(X)
        unique_y = np.unique(Y)
        n_clusters = p_t_given_x.shape[1]
        n_y_vals = len(unique_y)
        
        p_y_given_t = np.zeros((n_clusters, n_y_vals))
        
        for t in range(n_clusters):
            # Compute p(y|t) = Σ_x p(y|x) * p(x|t)
            denominator = 0
            for i, x in enumerate(unique_x):
                x_mask = (X == x)
                p_x_given_t = p_t_given_x[i, t] * np.sum(x_mask) / len(X)
                denominator += p_x_given_t
                
                for j, y in enumerate(unique_y):
                    y_given_x = np.sum((Y == y) & x_mask) / max(np.sum(x_mask), 1)
                    p_y_given_t[t, j] += y_given_x * p_x_given_t
                    
            if denominator > 0:
                p_y_given_t[t] /= denominator
            else:
                p_y_given_t[t] = np.ones(n_y_vals) / n_y_vals
                
        return p_y_given_t
    
    def _update_cluster_assignments(self, X, Y, p_y_given_t, beta):
        """Update p(t|x) based on current conditional probabilities."""
        unique_x = np.unique(X)
        unique_y = np.unique(Y)
        n_clusters = p_y_given_t.shape[0]
        n_x_vals = len(unique_x)
        
        p_t_given_x = np.zeros((n_x_vals, n_clusters))
        
        for i, x in enumerate(unique_x):
            x_mask = (X == x)
            
            # Compute p(y|x) for this x
            p_y_given_x = np.zeros(len(unique_y))
            for j, y in enumerate(unique_y):
                p_y_given_x[j] = np.sum((Y == y) & x_mask) / max(np.sum(x_mask), 1)
                
            # Update p(t|x) ∝ p(t) * exp(-β * D_KL[p(y|x) || p(y|t)])
            for t in range(n_clusters):
                # Compute KL divergence D_KL[p(y|x) || p(y|t)]
                kl_div = 0
                for j in range(len(unique_y)):
                    if p_y_given_x[j] > 0 and p_y_given_t[t, j] > 0:
                        kl_div += p_y_given_x[j] * np.log(p_y_given_x[j] / p_y_given_t[t, j])
                
                # Assume uniform prior p(t)
                p_t = 1.0 / n_clusters
                p_t_given_x[i, t] = p_t * np.exp(-beta * kl_div)
                
            # Normalize
            p_t_given_x[i] /= np.sum(p_t_given_x[i])
            
        return p_t_given_x
    
    def _compute_I_TX(self, X, p_t_given_x):
        """Compute I(X;T) from cluster assignments."""
        unique_x = np.unique(X)
        n_clusters = p_t_given_x.shape[1]
        
        # Compute joint distribution p(x,t)
        p_xt = np.zeros((len(unique_x), n_clusters))
        for i, x in enumerate(unique_x):
            p_x = np.sum(X == x) / len(X)
            p_xt[i] = p_x * p_t_given_x[i]
            
        # Compute marginals
        p_x = np.sum(p_xt, axis=1)
        p_t = np.sum(p_xt, axis=0)
        
        # Compute MI
        mi = 0
        for i in range(len(unique_x)):
            for t in range(n_clusters):
                if p_xt[i, t] > 0:
                    mi += p_xt[i, t] * np.log(p_xt[i, t] / (p_x[i] * p_t[t]))
                    
        return mi
    
    def _compute_I_TY(self, Y, p_y_given_t, p_t_given_x):
        """Compute I(T;Y) from conditional probabilities."""
        unique_x = np.unique(self.X_fit_)
        unique_y = np.unique(Y)
        n_clusters = p_y_given_t.shape[0]
        
        # Compute p(t) from p(t|x) and p(x)
        p_t = np.zeros(n_clusters)
        for i, x in enumerate(unique_x):
            p_x = np.sum(self.X_fit_ == x) / len(self.X_fit_)
            p_t += p_x * p_t_given_x[i]
            
        # Compute p(y)
        p_y = np.zeros(len(unique_y))
        for j, y in enumerate(unique_y):
            p_y[j] = np.sum(Y == y) / len(Y)
            
        # Compute joint p(t,y) = p(y|t) * p(t)
        p_ty = p_y_given_t * p_t.reshape(-1, 1)
        
        # Compute MI
        mi = 0
        for t in range(n_clusters):
            for j in range(len(unique_y)):
                if p_ty[t, j] > 0:
                    mi += p_ty[t, j] * np.log(p_ty[t, j] / (p_t[t] * p_y[j]))
                    
        return mi


class NeuralInformationBottleneck(nn.Module, BaseEstimator, TransformerMixin):
    """
    Neural Information Bottleneck implementation using deep networks.
    
    Uses variational approximation with neural networks to optimize:
    L = I(Z;Y) - β*I(X;Z)
    
    Where Z is the learned representation.
    """
    
    def __init__(self, input_dim, latent_dim=10, hidden_dims=[50, 20], 
                 beta=1.0, learning_rate=0.001, epochs=100):
        super().__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.hidden_dims = hidden_dims
        self.beta = beta
        self.learning_rate = learning_rate
        self.epochs = epochs
        
        # Build encoder network
        encoder_layers = []
        prev_dim = input_dim
        for hidden_dim in hidden_dims:
            encoder_layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                nn.ReLU(),
                nn.Dropout(0.2)
            ])
            prev_dim = hidden_dim
        
        # Output mean and log variance for reparameterization trick
        encoder_layers.append(nn.Linear(prev_dim, latent_dim * 2))
        self.encoder = nn.Sequential(*encoder_layers)
        
        # Build decoder/predictor network
        decoder_layers = [
            nn.Linear(latent_dim, hidden_dims[-1]),
            nn.ReLU(),
            nn.Linear(hidden_dims[-1], 1),  # Assuming regression task
            nn.Sigmoid()
        ]
        self.decoder = nn.Sequential(*decoder_layers)
        
    def encode(self, x):
        """Encode input to latent distribution parameters."""
        encoded = self.encoder(x)
        mu = encoded[:, :self.latent_dim] 
        logvar = encoded[:, self.latent_dim:]
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        """Reparameterization trick for backpropagation through stochastic nodes."""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        """Decode latent representation to prediction."""
        return self.decoder(z)
    
    def forward(self, x):
        """Forward pass through the network."""
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        y_pred = self.decode(z)
        return y_pred, mu, logvar, z
    
    def fit(self, X, y):
        """Fit the neural information bottleneck."""
        X_tensor = torch.FloatTensor(X)
        y_tensor = torch.FloatTensor(y).unsqueeze(1)
        
        optimizer = optim.Adam(self.parameters(), lr=self.learning_rate)
        
        for epoch in range(self.epochs):
            optimizer.zero_grad()
            
            y_pred, mu, logvar, z = self.forward(X_tensor)
            
            # Reconstruction loss (negative log-likelihood)
            recon_loss = nn.MSELoss()(y_pred, y_tensor)
            
            # KL divergence regularization (compression term)
            kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
            kl_loss /= X_tensor.size(0)  # Average over batch
            
            # Information Bottleneck objective
            loss = recon_loss + self.beta * kl_loss
            
            loss.backward()
            optimizer.step()
            
            if epoch % 20 == 0:
                print(f'Epoch {epoch}, Loss: {loss.item():.4f}, '
                      f'Recon: {recon_loss.item():.4f}, KL: {kl_loss.item():.4f}')
        
        return self
    
    def transform(self, X):
        """Transform input data to latent representation."""
        self.eval()
        with torch.no_grad():
            X_tensor = torch.FloatTensor(X)
            mu, _ = self.encode(X_tensor)
            return mu.numpy()
    
    def predict(self, X):
        """Predict target values."""
        self.eval()
        with torch.no_grad():
            X_tensor = torch.FloatTensor(X)
            y_pred, _, _, _ = self.forward(X_tensor)
            return y_pred.squeeze().numpy()


class DeepInformationBottleneck(NeuralInformationBottleneck):
    """
    Deep Information Bottleneck with multiple bottleneck layers.
    
    Extends the neural IB with hierarchical compression stages.
    """
    
    def __init__(self, input_dim, bottleneck_dims=[50, 20, 10], **kwargs):
        # Use final bottleneck dimension as latent_dim
        super().__init__(input_dim, latent_dim=bottleneck_dims[-1], **kwargs)
        self.bottleneck_dims = bottleneck_dims
        
        # Override encoder with hierarchical bottlenecks
        encoder_layers = []
        prev_dim = input_dim
        
        for i, bottleneck_dim in enumerate(bottleneck_dims):
            # Each bottleneck layer
            encoder_layers.extend([
                nn.Linear(prev_dim, bottleneck_dim),
                nn.BatchNorm1d(bottleneck_dim),
                nn.ReLU(),
                nn.Dropout(0.2)
            ])
            prev_dim = bottleneck_dim
        
        # Final layer outputs latent parameters
        encoder_layers.append(nn.Linear(prev_dim, self.latent_dim * 2))
        self.encoder = nn.Sequential(*encoder_layers)


class InformationBottleneckClassifier(BaseEstimator, ClassifierMixin, TransformerMixin):
    """
    Information Bottleneck for classification tasks.
    
    Optimizes compression-prediction trade-off specifically for classification.
    """
    
    def __init__(self, n_clusters=10, beta=1.0, **kwargs):
        self.ib = InformationBottleneck(n_clusters=n_clusters, beta=beta, **kwargs)
    
    def fit(self, X, y):
        """Fit IB classifier."""
        self.classes_ = np.unique(y)
        self.ib.fit(X, y)
        return self
    
    def predict(self, X):
        """Predict class labels."""
        return self.ib.predict(X)
    
    def predict_proba(self, X):
        """Predict class probabilities."""
        return self.ib.predict_proba(X)
    
    def transform(self, X):
        """Transform data through information bottleneck."""
        return self.ib.transform(X)


class IBOptimizer:
    """
    Optimizer for Information Bottleneck hyperparameters.
    
    Automatically finds optimal β and cluster count using cross-validation.
    """
    
    def __init__(self, beta_range=(0.1, 10.0), n_clusters_range=(2, 20), cv=5):
        self.beta_range = beta_range
        self.n_clusters_range = n_clusters_range
        self.cv = cv
    
    def optimize(self, X, y, metric='accuracy'):
        """Find optimal hyperparameters."""
        from sklearn.model_selection import cross_val_score
        
        best_score = -np.inf
        best_params = {}
        
        beta_values = np.logspace(np.log10(self.beta_range[0]), 
                                  np.log10(self.beta_range[1]), 10)
        
        for n_clusters in range(self.n_clusters_range[0], self.n_clusters_range[1] + 1):
            for beta in beta_values:
                ib = InformationBottleneckClassifier(n_clusters=n_clusters, beta=beta)
                
                try:
                    scores = cross_val_score(ib, X, y, cv=self.cv, scoring=metric)
                    mean_score = np.mean(scores)
                    
                    if mean_score > best_score:
                        best_score = mean_score
                        best_params = {'n_clusters': n_clusters, 'beta': beta}
                except:
                    continue
        
        self.best_params_ = best_params
        self.best_score_ = best_score
        
        return best_params


class MutualInfoEstimator:
    """
    Comprehensive mutual information estimation toolkit.
    
    Provides multiple estimators for different scenarios.
    """
    
    def __init__(self, method='discrete'):
        self.method = method
    
    def estimate(self, X, Y):
        """Estimate mutual information between X and Y."""
        if self.method == 'discrete':
            return self._discrete_mi(X, Y)
        elif self.method == 'continuous':
            return self._continuous_mi(X, Y)
        elif self.method == 'mixed':
            return self._mixed_mi(X, Y)
        else:
            raise ValueError(f"Unknown method: {self.method}")
    
    def _discrete_mi(self, X, Y):
        """Discrete mutual information."""
        from sklearn.metrics import mutual_info_score
        return mutual_info_score(X, Y)
    
    def _continuous_mi(self, X, Y):
        """Continuous mutual information using k-NN."""
        from sklearn.feature_selection import mutual_info_regression
        return mutual_info_regression(X.reshape(-1, 1), Y)[0]
    
    def _mixed_mi(self, X, Y):
        """Mixed discrete-continuous MI."""
        # Discretize continuous variables
        X_discrete = np.digitize(X, bins=np.quantile(X, np.linspace(0, 1, 10)))
        return self._discrete_mi(X_discrete, Y)


class MutualInfoCore:
    """
    Core mutual information computation algorithms.
    
    High-performance implementations of MI calculation methods.
    """
    
    @staticmethod
    def compute_joint_entropy(X, Y, bins='auto'):
        """
        Compute joint entropy H(X,Y) using histogram estimation.
        
        Args:
            X, Y: Input arrays
            bins: Number of bins for histogram ('auto' uses Sturges' rule)
            
        Returns:
            float: Joint entropy H(X,Y) in bits
        """
        # Convert to numpy arrays
        X = np.asarray(X).flatten()
        Y = np.asarray(Y).flatten()
        
        # Determine bin count using Sturges' rule if auto
        if bins == 'auto':
            bins = int(np.ceil(np.log2(len(X)) + 1))
        
        # Create joint histogram
        joint_hist, _, _ = np.histogram2d(X, Y, bins=bins)
        
        # Convert to probabilities (normalize)
        joint_prob = joint_hist / np.sum(joint_hist)
        
        # Remove zero probabilities to avoid log(0)
        joint_prob = joint_prob[joint_prob > 0]
        
        # Compute joint entropy: H(X,Y) = -Σ p(x,y) log₂ p(x,y)
        joint_entropy = -np.sum(joint_prob * np.log2(joint_prob))
        
        return joint_entropy
    
    @staticmethod
    def compute_conditional_entropy(X, Y, bins='auto'):
        """
        Compute conditional entropy H(Y|X) = H(X,Y) - H(X).
        
        Args:
            X, Y: Input arrays
            bins: Number of bins for histogram estimation
            
        Returns:
            float: Conditional entropy H(Y|X) in bits
        """
        # Convert to numpy arrays
        X = np.asarray(X).flatten()
        Y = np.asarray(Y).flatten()
        
        # Determine bin count using Sturges' rule if auto
        if bins == 'auto':
            bins = int(np.ceil(np.log2(len(X)) + 1))
        
        # Compute H(X,Y)
        joint_entropy = InformationBottleneckCore.compute_joint_entropy(X, Y, bins)
        
        # Compute H(X) - marginal entropy of X
        x_hist, _ = np.histogram(X, bins=bins)
        x_prob = x_hist / np.sum(x_hist)
        x_prob = x_prob[x_prob > 0]  # Remove zeros
        marginal_entropy_x = -np.sum(x_prob * np.log2(x_prob))
        
        # H(Y|X) = H(X,Y) - H(X)
        conditional_entropy = joint_entropy - marginal_entropy_x
        
        return conditional_entropy
    
    @staticmethod
    def compute_kl_divergence(P, Q):
        """Compute KL divergence D_KL(P||Q)."""
        # Ensure no zeros
        P = P + 1e-10
        Q = Q + 1e-10
        
        return np.sum(P * np.log(P / Q))