import numpy as np


class EmergentTruthLearner:
    """
    Self-learning system that discovers truths without labels.
    
    Philosophy: What doesn't break is truth.
    By testing classifier predictions on random noise and measuring
    stability, we discover which patterns the classifier considers
    "real" - these are emergent truths that can guide learning.
    """
    
    def __init__(self, classifier, buffer_size=5000):
        """
        Args:
            classifier: MLPClassifier instance
            buffer_size: Size of memory buffer for stable samples
        """
        self.classifier = classifier
        
        # Memory of emergent truths
        self.truth_buffer_X = np.zeros((buffer_size, 784))
        self.truth_buffer_labels = np.zeros(buffer_size, dtype=int)
        self.truth_buffer_stability = np.zeros(buffer_size)
        self.buffer_idx = 0
        self.buffer_full = False
        self.buffer_size = buffer_size
        
        # Statistics
        self.total_stable_found = 0
        self.iteration = 0
        
    def discover_from_noise(self, n_samples=100, stability_threshold=0.85):
        """
        Core self-learning: Find emergent truths in random noise.
        
        The classifier's stable predictions on noise reveal its
        internal truth about what patterns belong to which classes.
        
        Returns:
            n_stable: Number of stable samples found
        """
        # Generate random noise (like f.predict(np.random.randn(100, 784)))
        X_noise = np.random.randn(n_samples, 784)
        
        # Measure stability
        from stability_analyzer import StabilityAnalyzer
        analyzer = StabilityAnalyzer(self.classifier, n_perturbations=15, noise_level=0.05)
        stable_info = analyzer.find_stable_samples(X_noise, threshold=stability_threshold)
        
        # Store stable samples as emergent truths
        n_stable = stable_info['X'].shape[0]
        if n_stable > 0:
            self._add_to_truth_buffer(
                stable_info['X'],
                stable_info['labels'],
                stable_info['stability']
            )
            self.total_stable_found += n_stable
        
        self.iteration += 1
        return n_stable
    
    def _add_to_truth_buffer(self, X, labels, stability):
        """Add stable samples to circular buffer."""
        n = len(X)
        for i in range(n):
            idx = self.buffer_idx % self.buffer_size
            self.truth_buffer_X[idx] = X[i]
            self.truth_buffer_labels[idx] = labels[i]
            self.truth_buffer_stability[idx] = stability[i]
            self.buffer_idx += 1
            
            if self.buffer_idx >= self.buffer_size:
                self.buffer_full = True
    
    def get_self_training_data(self, min_stability=0.8):
        """
        Get data for self-supervised training.
        
        Returns:
            X, labels, weights: Filtered by stability, weighted by confidence
        """
        if not self.buffer_full:
            mask = np.arange(self.buffer_idx)
        else:
            mask = np.arange(self.buffer_size)
        
        # Filter by stability
        stable_mask = self.truth_buffer_stability[mask] >= min_stability
        X = self.truth_buffer_X[mask][stable_mask]
        labels = self.truth_buffer_labels[mask][stable_mask]
        weights = self.truth_buffer_stability[mask][stable_mask]
        
        return X, labels, weights
    
    def self_train(self, batch_size=200, learning_rate_scale=0.5):
        """
        Train classifier on emergent truths.
        
        Args:
            batch_size: Number of stable samples to train on
            learning_rate_scale: Scale factor for self-learning rate
            
        Returns:
            loss: Training loss, or None if no stable samples
        """
        X, labels, weights = self.get_self_training_data()
        
        if len(X) < batch_size:
            return None
        
        # Sample batch
        idx = np.random.choice(len(X), batch_size, replace=False)
        X_batch = X[idx]
        labels_batch = labels[idx]
        weights_batch = weights[idx]
        
        # Forward pass
        probs = self.classifier.forward(X_batch)
        
        # Compute weighted loss
        one_hot = np.eye(10)[labels_batch]
        loss = -np.sum(weights_batch[:, np.newaxis] * one_hot * np.log(probs + 1e-9)) / batch_size
        
        # Backward pass with stability-weighted gradients
        dz2 = probs - one_hot
        dW2 = np.dot(self.classifier.a1.T, dz2 * weights_batch[:, np.newaxis]) / batch_size
        db2 = np.sum(dz2 * weights_batch[:, np.newaxis], axis=0, keepdims=True) / batch_size
        
        da1 = np.dot(dz2 * weights_batch[:, np.newaxis], self.classifier.W2.T)
        dz1 = da1 * self.classifier.relu_derivative(self.classifier.z1)
        dW1 = np.dot(X_batch.T, dz1) / batch_size
        db1 = np.sum(dz1, axis=0, keepdims=True) / batch_size
        
        # Update with scaled learning rate
        lr = self.classifier.learning_rate * learning_rate_scale
        self.classifier.W1 -= lr[0] * dW1
        self.classifier.b1 -= lr[1] * db1
        self.classifier.W2 -= lr[2] * dW2
        self.classifier.b2 -= lr[3] * db2
        
        return loss
    
    def get_truth_statistics(self):
        """Get statistics about discovered emergent truths."""
        if not self.buffer_full:
            n_total = self.buffer_idx
        else:
            n_total = self.buffer_size
        
        if n_total == 0:
            return {'buffer_usage': 0, 'mean_stability': 0, 'class_distribution': {}}
        
        stability = self.truth_buffer_stability[:n_total]
        labels = self.truth_buffer_labels[:n_total]
        
        class_dist = {}
        for c in range(10):
            class_dist[c] = int(np.sum(labels == c))
        
        return {
            'buffer_usage': n_total / self.buffer_size,
            'mean_stability': float(np.mean(stability)),
            'total_stable_found': self.total_stable_found,
            'class_distribution': class_dist
        }
    
    def combined_training_step(self, X_train, y_train, batch_size=100, 
                               self_learning_ratio=0.3):
        """
        Combine supervised learning with self-learning.
        
        Args:
            X_train, y_train: Labeled training data
            batch_size: Total batch size
            self_learning_ratio: Fraction from self-learning
            
        Returns:
            supervised_loss, self_learning_loss
        """
        # Supervised learning step
        idx_sup = np.random.randint(0, len(X_train), int(batch_size * (1 - self_learning_ratio)))
        X_sup = X_train[idx_sup]
        y_sup = y_train[idx_sup]
        self.classifier.update(X_sup, np.eye(10)[y_sup])
        sup_loss = self.classifier.compute_loss(
            np.eye(10)[y_sup],
            self.classifier.forward(X_sup)
        )
        
        # Self-learning step
        self_loss = self.self_train(batch_size=int(batch_size * self_learning_ratio))
        
        return sup_loss, self_loss
