"""
Predictive Learning Layer

Generic predictive coding implementation that works with ANY PyTorch model.
Can be added to any layer to enable local predictive learning.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from .adaptive_compression import AdaptiveSignalLinear

class PredictiveLayer(nn.Module):
    """
    Add predictive coding to any layer.
    
    The layer learns to predict the next layer's activation,
    providing a local learning signal.
    
    Usage:
        layer = nn.Linear(256, 256)
        predictive_layer = PredictiveLayer(layer, hidden_dim=256)
        
        # Forward pass
        output, prediction = predictive_layer(input)
        
        # Compute prediction error with next layer's actual output
        pred_error = predictive_layer.compute_prediction_error(prediction, actual_next)
    """
    
    def __init__(self, base_layer, hidden_dim=None, prediction_weight=0.3):
        super().__init__()
        
        self.base_layer = base_layer
        self.prediction_weight = prediction_weight
        
        # Infer dimensions
        if hasattr(base_layer, 'out_features'):
            out_dim = base_layer.out_features
        elif hasattr(base_layer, 'd_model'):
            out_dim = base_layer.d_model
        else:
            raise ValueError("Cannot infer output dimension from base_layer")
        
        hidden_dim = hidden_dim or out_dim
        
        # Prediction network: predicts next layer's activation
        self.prediction_net = nn.Sequential(
            AdaptiveSignalLinear(out_dim, hidden_dim),
            nn.GELU(),
            AdaptiveSignalLinear(hidden_dim, out_dim)
        )
        
        # Store prediction error for analysis
        self.register_buffer('last_prediction_error', torch.tensor(0.0))
        
    def forward(self, x, return_prediction=True):
        """
        Forward pass with optional prediction.
        
        Args:
            x: Input tensor
            return_prediction: Whether to compute and return prediction
        
        Returns:
            output: Output from base layer
            prediction: Predicted next layer activation (if return_prediction=True)
        """
        # Standard forward through base layer
        output = self.base_layer(x)
        
        if return_prediction:
            # Predict next layer's activation
            prediction = self.prediction_net(output)
            return output, prediction
        else:
            return output, None
    
    def compute_prediction_error(self, prediction, actual_next):
        """
        Compute prediction error for local learning.
        
        Args:
            prediction: Predicted next layer activation
            actual_next: Actual next layer activation
        
        Returns:
            error: MSE between prediction and actual
        """
        if prediction is None:
            return torch.tensor(0.0, device=actual_next.device)
        
        # Detach actual to avoid backprop through it
        error = F.mse_loss(prediction, actual_next.detach())
        self.last_prediction_error = error.detach()
        
        return error
    
    def get_prediction_loss(self):
        """Get the last prediction error"""
        return self.last_prediction_error.item()


class PredictiveLearningMixin:
    """
    Mixin to add predictive learning to any PyTorch model.
    
    Usage:
        class MyModel(nn.Module, PredictiveLearningMixin):
            def __init__(self):
                super().__init__()
                self.layer1 = nn.Linear(128, 256)
                self.layer2 = nn.Linear(256, 256)
                self.layer3 = nn.Linear(256, 128)
                
                # Enable predictive learning
                self.enable_predictive_learning()
            
            def forward(self, x):
                x = self.layer1(x)
                x = self.layer2(x)
                x = self.layer3(x)
                return x
    """
    
    def enable_predictive_learning(self, layer_names=None, prediction_weight=0.3):
        """
        Enable predictive learning for specified layers.
        
        Args:
            layer_names: List of layer names to wrap (None = all eligible layers)
            prediction_weight: Weight for prediction loss (0-1)
        """
        self.predictive_layers = {}
        self.prediction_weight = prediction_weight
        
        # Find layers to wrap
        if layer_names is None:
            layer_names = [name for name, module in self.named_modules()
                          if isinstance(module, (nn.Linear, nn.Conv2d))]
        
        # Wrap each layer
        for name in layer_names:
            try:
                # Get the layer
                parts = name.split('.')
                parent = self
                for part in parts[:-1]:
                    parent = getattr(parent, part)
                
                layer = getattr(parent, parts[-1])
                
                # Wrap with predictive layer
                predictive_layer = PredictiveLayer(layer, prediction_weight=prediction_weight)
                setattr(parent, parts[-1], predictive_layer)
                
                self.predictive_layers[name] = predictive_layer
                
            except Exception as e:
                print(f"Warning: Could not wrap layer {name}: {e}")
        
        print(f"✅ Enabled predictive learning for {len(self.predictive_layers)} layers")
    
    def compute_prediction_losses(self, layer_activations):
        """
        Compute prediction losses for all predictive layers.
        
        Args:
            layer_activations: Dict mapping layer names to their activations
        
        Returns:
            total_pred_loss: Sum of all prediction losses
        """
        if not hasattr(self, 'predictive_layers'):
            return torch.tensor(0.0)
        
        pred_losses = []
        layer_names = list(layer_activations.keys())
        
        for i, (name, layer) in enumerate(self.predictive_layers.items()):
            if i < len(layer_names) - 1:
                # Get prediction and actual next
                current_activation = layer_activations[layer_names[i]]
                next_activation = layer_activations[layer_names[i+1]]
                
                # Compute prediction if we have one
                if hasattr(layer, 'prediction_net'):
                    prediction = layer.prediction_net(current_activation)
                    pred_error = layer.compute_prediction_error(prediction, next_activation)
                    pred_losses.append(pred_error)
        
        if len(pred_losses) > 0:
            return sum(pred_losses) / len(pred_losses)
        else:
            return torch.tensor(0.0)
    
    def get_prediction_stats(self):
        """Get statistics about prediction errors"""
        if not hasattr(self, 'predictive_layers'):
            return {}
        
        stats = {}
        for name, layer in self.predictive_layers.items():
            stats[name] = layer.get_prediction_loss()
        
        return stats


# Example usage
if __name__ == "__main__":
    print("🧪 Testing Predictive Layer\n")
    
    # Test 1: Wrap a single layer
    print("Test 1: Single Layer")
    base_layer = nn.Linear(128, 256)
    pred_layer = PredictiveLayer(base_layer)
    
    x = torch.randn(8, 128)
    output, prediction = pred_layer(x)
    
    print(f"✅ Input: {x.shape}")
    print(f"✅ Output: {output.shape}")
    print(f"✅ Prediction: {prediction.shape}")
    
    # Simulate next layer
    next_output = torch.randn(8, 256)
    pred_error = pred_layer.compute_prediction_error(prediction, next_output)
    print(f"✅ Prediction error: {pred_error.item():.4f}\n")
    
    # Test 2: Mixin with full model
    print("Test 2: Full Model with Mixin")
    
    class TestModel(nn.Module, PredictiveLearningMixin):
        def __init__(self):
            super().__init__()
            self.layer1 = nn.Linear(128, 256)
            self.layer2 = nn.Linear(256, 256)
            self.layer3 = nn.Linear(256, 128)
            
            # Enable predictive learning
            self.enable_predictive_learning()
        
        def forward(self, x):
            activations = {}
            activations['layer1'] = self.layer1(x)
            activations['layer2'] = self.layer2(activations['layer1'])
            activations['layer3'] = self.layer3(activations['layer2'])
            return activations['layer3'], activations
    
    model = TestModel()
    x = torch.randn(8, 128)
    output, activations = model(x)
    
    pred_loss = model.compute_prediction_losses(activations)
    print(f"✅ Output: {output.shape}")
    print(f"✅ Prediction loss: {pred_loss.item():.4f}")
    print(f"✅ Prediction stats: {model.get_prediction_stats()}")
    
    print("\n✅ All tests passed!")