import numpy as np
from emergent_truth import EmergentTruthLearner
from stability_analyzer import StabilityAnalyzer


class SelfLearningWrapper:
    """
    Wraps MLPClassifier to add self-learning capabilities.
    
    Usage:
        wrapper = SelfLearningWrapper(f)
        wrapper.train_with_self_learning(X_train, y_train, X_test, y_test)
        
        # Use enhanced prediction
        pred = wrapper.predict(X)
        confidence = wrapper.predict_with_confidence(X)
    """
    
    def __init__(self, classifier, 
                 discovery_interval=10,
                 self_train_interval=5,
                 stability_threshold=0.85):
        """
        Args:
            classifier: MLPClassifier instance (your f from c00.py)
            discovery_interval: Discover emergent truths every N iterations
            self_train_interval: Self-train every N iterations
            stability_threshold: Minimum stability to consider as truth
        """
        self.classifier = classifier
        self.learner = EmergentTruthLearner(classifier)
        self.analyzer = StabilityAnalyzer(classifier)
        
        self.discovery_interval = discovery_interval
        self.self_train_interval = self_train_interval
        self.stability_threshold = stability_threshold
        
        self.iteration = 0
        self.history = {
            'supervised_loss': [],
            'self_learning_loss': [],
            'stability_scores': [],
            'n_stable_found': []
        }
    
    def predict(self, X):
        """Standard prediction."""
        return self.classifier.predict(X)
    
    def predict_with_confidence(self, X):
        """
        Prediction with stability-based confidence.
        
        Returns:
            predictions: Array of predicted labels
            confidence: Array [batch_size, 10] - stability-weighted probabilities
            stability: Array [batch_size] - how stable each prediction is
        """
        stability_scores, predictions, avg_confidence = self.analyzer.measure_stability(X)
        return predictions, avg_confidence, stability_scores
    
    def train_step(self, X, y_onehot):
        """
        Single training step with optional self-learning.
        
        Args:
            X: Training batch
            y_onehot: One-hot encoded labels
            
        Returns:
            dict with losses and stats
        """
        # Standard supervised update
        self.classifier.update(X, y_onehot)
        y_pred = self.classifier.forward(X)
        sup_loss = self.classifier.compute_loss(y_onehot, y_pred)
        
        results = {
            'iteration': self.iteration,
            'supervised_loss': sup_loss,
            'self_learning_loss': None,
            'stability': None,
            'n_stable': 0
        }
        
        # Periodic self-learning
        if self.iteration % self.discovery_interval == 0:
            n_stable = self.learner.discover_from_noise(
                n_samples=100,
                stability_threshold=self.stability_threshold
            )
            results['n_stable'] = n_stable
        
        if self.iteration % self.self_train_interval == 0:
            self_loss = self.learner.self_train(batch_size=100)
            results['self_learning_loss'] = self_loss
        
        # Measure stability periodically
        if self.iteration % 50 == 0:
            X_test_noise = np.random.randn(100, 784)
            stability, _, _ = self.analyzer.measure_stability(X_test_noise)
            results['stability'] = float(np.mean(stability))
        
        self.history['supervised_loss'].append(sup_loss)
        self.history['self_learning_loss'].append(results['self_learning_loss'])
        self.history['n_stable_found'].append(results['n_stable'])
        
        self.iteration += 1
        return results
    
    def train_loop(self, X_train, y_train, n_iterations=1000, 
                   batch_size=100, verbose=True):
        """
        Full training loop combining supervised and self-learning.
        
        Args:
            X_train: Training data [n_samples, 784]
            y_train: Training labels [n_samples]
            n_iterations: Number of iterations
            batch_size: Batch size for supervised learning
            verbose: Print progress
        """
        for i in range(n_iterations):
            # Sample batch
            idx = np.random.randint(0, len(X_train), batch_size)
            X_batch = X_train[idx]
            y_batch = y_train[idx]
            
            # Training step
            results = self.train_step(X_batch, np.eye(10)[y_batch])
            
            if verbose and i % 100 == 0:
                stats = self.learner.get_truth_statistics()
                print(f"Iter {i:6d} | "
                      f"Sup Loss: {results['supervised_loss']:.4f} | "
                      f"Self Loss: {results['self_learning_loss'] if results['self_learning_loss'] else 'N/A':>8} | "
                      f"Stable: {results['n_stable']:3d} | "
                      f"Buffer: {stats['buffer_usage']:.1%}")
    
    def get_truth_report(self):
        """Get comprehensive report on discovered emergent truths."""
        stats = self.learner.get_truth_statistics()
        
        # Test on actual test data if available
        report = {
            'emergent_truths': stats,
            'training_history': {
                'iterations': self.iteration,
                'final_supervised_loss': self.history['supervised_loss'][-1] if self.history['supervised_loss'] else None,
            }
        }
        
        return report