import numpy as np


class StabilityAnalyzer:
    """
    Measures prediction stability across perturbations.
    
    Core principle: What doesn't break is an emergent truth.
    If a classifier gives consistent predictions on perturbed versions
    of the same input, that prediction is "stable" and likely correct.
    """
    
    def __init__(self, classifier, n_perturbations=20, noise_level=0.1):
        """
        Args:
            classifier: MLPClassifier instance with predict() and forward()
            n_perturbations: Number of perturbed samples per input
            noise_level: Standard deviation of Gaussian noise
        """
        self.classifier = classifier
        self.n_perturbations = n_perturbations
        self.noise_level = noise_level
        
    def measure_stability(self, X):
        """
        Measure how stable predictions are for input X under noise.
        
        Returns:
            stability_scores: Array [batch_size] - fraction of consistent predictions
            predicted_labels: Array [batch_size] - most common prediction
            confidence: Array [batch_size, n_classes] - prediction distribution
        """
        batch_size = X.shape[0]
        all_predictions = np.zeros((batch_size, self.n_perturbations), dtype=int)
        all_probabilities = np.zeros((batch_size, self.n_perturbations, 10))
        
        for i in range(self.n_perturbations):
            X_perturbed = X + np.random.randn(*X.shape) * self.noise_level
            probs = self.classifier.forward(X_perturbed)
            all_predictions[:, i] = self.classifier.predict(X_perturbed)
            all_probabilities[:, i] = probs
        
        # Calculate stability: fraction agreeing with mode
        mode_predictions = np.zeros(batch_size, dtype=int)
        stability_scores = np.zeros(batch_size)
        
        for j in range(batch_size):
            counts = np.bincount(all_predictions[j], minlength=10)
            mode_predictions[j] = np.argmax(counts)
            stability_scores[j] = counts[mode_predictions[j]] / self.n_perturbations
        
        # Average probabilities
        confidence = np.mean(all_probabilities, axis=1)
        
        return stability_scores, mode_predictions, confidence
    
    def find_stable_samples(self, X, threshold=0.9):
        """
        Find samples where predictions are highly stable.
        
        Returns:
            stable_mask: Boolean mask of stable samples
            stable_X: Stable input samples
            stable_labels: Their emergent truth labels
        """
        stability_scores, predictions, confidence = self.measure_stability(X)
        stable_mask = stability_scores >= threshold
        
        return {
            'mask': stable_mask,
            'X': X[stable_mask],
            'labels': predictions[stable_mask],
            'stability': stability_scores[stable_mask],
            'confidence': confidence[stable_mask]
        }
    
    def analyze_decision_boundaries(self, X, n_steps=50):
        """
        Analyze decision boundary stability by interpolating between samples.
        
        Returns:
            boundary_stability: How stable predictions are near boundaries
        """
        if X.shape[0] < 2:
            return 0.0
            
        stabilities = []
        
        for i in range(min(len(X) - 1, 10)):  # Sample pairs
            x1, x2 = X[i], X[i + 1]
            interpolation = np.linspace(0, 1, n_steps)[:, np.newaxis]
            X_interp = x1[np.newaxis, :] * (1 - interpolation) + x2[np.newaxis, :] * interpolation
            
            preds = self.classifier.predict(X_interp)
            # Count number of class transitions (boundary crossings)
            transitions = np.sum(preds[:-1] != preds[1:])
            # Fewer transitions = more stable boundary
            stability = 1.0 - (transitions / (n_steps - 1))
            stabilities.append(stability)
        
        return np.mean(stabilities) if stabilities else 0.0