import numpy as np
from scipy.io.wavfile import read
from scipy.signal import find_peaks


class AutoregressiveLimitTracker:
    """
    Tracks autoregressive limit estimates for each weight element W[i,j].
    Each element maintains its own converging statistics to estimate the limit as n→∞.
    """
    
    def __init__(self, shape, decay=0.99, min_samples=50):
        """
        Initialize tracker for all elements.
        
        Args:
            shape: Shape of weight matrix (e.g., (784, 100))
            decay: Exponential decay factor for moving average (closer to 1 = more history)
            min_samples: Minimum samples before limit estimate is reliable
        """
        self.shape = shape
        self.decay = decay
        self.min_samples = min_samples
        
        # Running statistics for each element
        self.ema_mean = np.zeros(shape)          # Exponential moving average (limit estimate)
        self.ema_var = np.zeros(shape)           # Exponential moving variance (oscillation strength)
        self.ema_update_sq = np.zeros(shape)     # EMA of update magnitudes
        self.sample_count = 0
        
        # Oscillation detection per element
        self.recent_values = {}  # Store recent values for peak detection
        
        # Convergence status
        self.is_converged = np.zeros(shape, dtype=bool)
        self.is_oscillating = np.zeros(shape, dtype=bool)
        
    def update(self, weight_matrix):
        """
        Update tracker with current weight values.
        
        Args:
            weight_matrix: Current weight matrix W[i,j]
        """
        self.sample_count += 1
        
        # Calculate updates (difference from previous EMA)
        if self.sample_count == 1:
            self.ema_mean = weight_matrix.copy()
            return
        
        diff = weight_matrix - self.ema_mean
        
        # Update exponential moving average (this is the limit estimate)
        self.ema_mean = self.decay * self.ema_mean + (1 - self.decay) * weight_matrix
        
        # Update exponential moving variance (measures oscillation amplitude)
        self.ema_var = self.decay * self.ema_var + (1 - self.decay) * (diff ** 2)
        
        # Update EMA of update magnitudes
        self.ema_update_sq = self.decay * self.ema_update_sq + (1 - self.decay) * (diff ** 2)
        
        # Track recent values for oscillation detection (last 200 samples)
        # Store as flattened array for memory efficiency
        flat_weights = weight_matrix.flatten()
        for idx in range(len(flat_weights)):
            if idx not in self.recent_values:
                self.recent_values[idx] = []
            self.recent_values[idx].append(flat_weights[idx])
            # Keep only last 200 values
            if len(self.recent_values[idx]) > 200:
                self.recent_values[idx] = self.recent_values[idx][-200:]
        
        # Update convergence/oscillation status
        self._update_status()
    
    def _update_status(self):
        """Update convergence and oscillation status for all elements."""
        flat_ema_var = self.ema_var.flatten()
        flat_ema_mean = self.ema_mean.flatten()
        
        for idx in range(len(flat_ema_mean)):
            if idx in self.recent_values and len(self.recent_values[idx]) >= self.min_samples:
                history = np.array(self.recent_values[idx])
                
                # Check oscillation
                peaks, _ = find_peaks(history, distance=3)
                oscillation_amplitude = np.sqrt(flat_ema_var[idx])
                
                # Oscillating if variance is significant and has peaks
                self.is_oscillating.flat[idx] = (
                    len(peaks) >= 3 and 
                    oscillation_amplitude > 1e-6 and
                    oscillation_amplitude < 1.0  # Not diverging
                )
                
                # Converged if variance is small relative to mean
                mean_abs = np.abs(flat_ema_mean[idx])
                if mean_abs > 1e-10:
                    cv = oscillation_amplitude / mean_abs  # Coefficient of variation
                    self.is_converged.flat[idx] = (cv < 0.01 and len(peaks) < 5)
                else:
                    self.is_converged.flat[idx] = (oscillation_amplitude < 1e-5)
    
    def get_limit_matrix(self):
        """
        Get the current limit estimate for all elements.
        Returns the EMA mean which is the best estimate of lim_{n→∞} E[W[i,j]].
        """
        return self.ema_mean.copy()
    
    def get_oscillation_amplitude_matrix(self):
        """Get oscillation amplitude (std) for each element."""
        return np.sqrt(self.ema_var.copy())
    
    def get_element_stats(self, i, j):
        """
        Get detailed statistics for specific element W[i,j].
        
        Args:
            i, j: Element indices
            
        Returns:
            dict with statistics
        """
        idx = i * self.shape[1] + j
        flat_mean = self.ema_mean.flat[idx]
        flat_var = self.ema_var.flat[idx]
        flat_std = np.sqrt(flat_var)
        
        history = np.array(self.recent_values.get(idx, []))
        
        # Frequency analysis if oscillating
        dominant_freq = None
        if len(history) > 10 and self.is_oscillating.flat[idx]:
            signal = history - np.mean(history)
            fft_values = np.fft.rfft(signal)
            fft_freqs = np.fft.rfftfreq(len(signal))
            fft_power = np.abs(fft_values) ** 2
            
            if len(fft_power) > 1:
                dominant_idx = np.argmax(fft_power[1:]) + 1
                dominant_freq = fft_freqs[dominant_idx]
        
        return {
            'limit_estimate': flat_mean,
            'oscillation_amplitude': flat_std,
            'is_oscillating': self.is_oscillating.flat[idx],
            'is_converged': self.is_converged.flat[idx],
            'sample_count': self.sample_count,
            'current_value': history[-1] if len(history) > 0 else None,
            'dominant_frequency': dominant_freq,
            'confidence': min(1.0, self.sample_count / self.min_samples)
        }
    
    def summary(self, top_n=10):
        """
        Get summary of convergence across all elements.
        
        Args:
            top_n: Number of top oscillating/converged elements to show
            
        Returns:
            dict with summary statistics
        """
        total_elements = self.shape[0] * self.shape[1]
        num_oscillating = np.sum(self.is_oscillating)
        num_converged = np.sum(self.is_converged)
        
        # Find most interesting elements
        flat_std = np.sqrt(self.ema_var).flatten()
        flat_mean = self.ema_mean.flatten()
        
        # Top oscillating
        top_osc_idx = np.argsort(flat_std)[-top_n:][::-1]
        top_conv_idx = np.argsort(np.sqrt(self.ema_var).flatten())[:top_n]
        
        return {
            'total_elements': total_elements,
            'num_oscillating': int(num_oscillating),
            'num_converged': int(num_converged),
            'pct_oscillating': f"{100*num_oscillating/total_elements:.1f}%",
            'pct_converged': f"{100*num_converged/total_elements:.1f}%",
            'top_oscillating': [
                (idx // self.shape[1], idx % self.shape[1], flat_std[idx])
                for idx in top_osc_idx if flat_std[idx] > 1e-6
            ],
            'most_converged': [
                (idx // self.shape[1], idx % self.shape[1], flat_std[idx])
                for idx in top_conv_idx
            ],
            'global_limit_matrix_shape': self.shape,
            'average_limit': np.mean(flat_mean),
            'average_oscillation': np.mean(flat_std)
        }


class MLPClassifierWithAutoregressiveLimits:
    """MLP Classifier with autoregressive limit tracking for all W1 elements."""
    
    def __init__(self, input_size, hidden_size, output_size, learning_rate=np.random.rand(4)):
        self.W1 = np.random.randn(input_size, hidden_size) * 0.01
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * 0.01
        self.b2 = np.zeros((1, output_size))
        self.learning_rate = learning_rate
        
        # Autoregressive limit tracker for W1
        self.w1_tracker = AutoregressiveLimitTracker((input_size, hidden_size), decay=0.99)
        
        # Training history
        self.loss_history = []
        self.accuracy_history = []
        self.iteration_count = 0
        
    def relu(self, x):
        return np.maximum(0, x)
    
    def relu_derivative(self, x):
        return np.where(x > 0, 1, 0)
    
    def softmax(self, x):
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)
    
    def forward(self, X):
        self.z1 = np.dot(X, self.W1) + self.b1
        self.a1 = self.relu(self.z1)
        self.z2 = np.dot(self.a1, self.W2) + self.b2
        output = self.softmax(self.z2)
        return output
    
    def compute_loss(self, y_true, y_pred):
        m = y_true.shape[0]
        loss = -np.sum(y_true * np.log(y_pred + 1e-9)) / m
        return loss
    
    def backward(self, X, y_true, y_pred):
        m = y_true.shape[0]
        dz2 = y_pred - y_true
        dW2 = np.dot(self.a1.T, dz2) / m
        db2 = np.sum(dz2, axis=0, keepdims=True) / m
        da1 = np.dot(dz2, self.W2.T)
        dz1 = da1 * self.relu_derivative(self.z1)
        dW1 = np.dot(X.T, dz1) / m
        db1 = np.sum(dz1, axis=0, keepdims=True) / m
        return dW1, db1, dW2, db2
    
    def update(self, X, y_true):
        y_pred = self.forward(X)
        dW1, db1, dW2, db2 = self.backward(X, y_true, y_pred)
        self.W1 -= self.learning_rate[0] * dW1
        self.b1 -= self.learning_rate[1] * db1
        self.W2 -= self.learning_rate[2] * dW2
        self.b2 -= self.learning_rate[3] * db2
        
        # Record loss
        loss = self.compute_loss(y_true, y_pred)
        self.loss_history.append(loss)
        self.iteration_count += 1
    
    def track_limits(self):
        """Update autoregressive limit tracker with current W1."""
        self.w1_tracker.update(self.W1)
    
    def get_limit_matrix(self):
        """Get limit estimate for all W1[i,j]."""
        return self.w1_tracker.get_limit_matrix()
    
    def predict(self, X):
        probabilities = self.forward(X)
        return np.argmax(probabilities, axis=1)
    
    def score(self, X, y_true):
        y = self.predict(X)
        acc = np.mean(y == y_true)
        self.accuracy_history.append(acc)
        return acc


def train_with_autoregressive_limits(num_iterations=5000, track_every=1):
    """
    Train MLP while tracking autoregressive limits for all W1[i,j].
    
    Args:
        num_iterations: Number of training iterations
        track_every: How often to update limit tracker (1 = every iteration)
        
    Returns:
        Trained MLP and limit analysis results
    """
    # Load data
    print("Loading data...")
    X_train = read('../X_train.wav')[1].reshape(-1, 784)
    y_train = (read('../y_train.wav')[1] * 9).astype(int)
    X_test = read('../X_test.wav')[1].reshape(-1, 784)
    y_test = (read('../y_test.wav')[1] * 9).astype(int)
    X20 = X_test[:1000]
    yt20 = y_test[:1000]
    
    print(f"X_train shape: {X_train.shape}")
    print(f"W1 shape will be: (784, 100) = {784*100} elements to track")
    
    # Initialize MLP
    learning_rate = np.random.rand(4)
    print(f"Learning rates: {learning_rate}")
    
    f = MLPClassifierWithAutoregressiveLimits(
        input_size=784, 
        hidden_size=100, 
        output_size=10, 
        learning_rate=learning_rate
    )
    
    # Training loop
    print(f"\nTraining for {num_iterations} iterations...")
    print(f"Tracking limits every {track_every} iterations")
    print("="*70)
    
    for i in range(num_iterations):
        # Sample batch
        idx = np.random.randint(0, 60000, 100)
        X = X_train[idx]
        yt = y_train[idx]
        f.update(X, np.eye(10)[yt])
            
        # Track limits
        if i % track_every == 0:
            f.track_limits()
        
        # Print progress
        if i % 500 == 0:
            acc = f.score(X20, yt20)
            summary = f.w1_tracker.summary()
            
            print(f"\nIteration {i:5d} | Accuracy: {acc:.4f} | Loss: {f.loss_history[-1]:.4f}")
            print(f"  Oscillating: {summary['num_oscillating']:5d} ({summary['pct_oscillating']})")
            print(f"  Converged:   {summary['num_converged']:5d} ({summary['pct_converged']})")
            print(f"  Avg limit:   {summary['average_limit']:.6e}")
            print(f"  Avg osc:     {summary['average_oscillation']:.6e}")
    
    # Final analysis
    print("\n" + "="*70)
    print("FINAL AUTOREGRESSIVE LIMIT ANALYSIS")
    print("="*70)
    
    summary = f.w1_tracker.summary(top_n=10)
    
    print(f"\nTotal W1 elements: {summary['total_elements']}")
    print(f"Oscillating: {summary['num_oscillating']} ({summary['pct_oscillating']})")
    print(f"Converged: {summary['num_converged']} ({summary['pct_converged']})")
    print(f"\nAverage limit value: {summary['average_limit']:.6e}")
    print(f"Average oscillation amplitude: {summary['average_oscillation']:.6e}")
    
    print(f"\nTop 10 most oscillating elements:")
    for i, j, amp in summary['top_oscillating']:
        stats = f.w1_tracker.get_element_stats(i, j)
        print(f"  W1[{i:3d},{j:3d}]: limit={stats['limit_estimate']:12.6e}, "
              f"osc_amp={amp:.6e}, freq={stats['dominant_frequency']}")
    
    print(f"\nTop 10 most converged elements:")
    for i, j, amp in summary['most_converged']:
        stats = f.w1_tracker.get_element_stats(i, j)
        print(f"  W1[{i:3d},{j:3d}]: limit={stats['limit_estimate']:12.6e}, "
              f"osc_amp={amp:.6e}, converged={stats['is_converged']}")
    
    # Export limit matrix
    limit_matrix = f.get_limit_matrix()
    print(f"\nLimit matrix W1_lim shape: {limit_matrix.shape}")
    print(f"Limit matrix stats: min={limit_matrix.min():.6e}, "
          f"max={limit_matrix.max():.6e}, mean={limit_matrix.mean():.6e}")
    
    return f, limit_matrix, summary


def iterative_limit_replacement(num_cycles=10, iterations_per_cycle=50):
    """
    Train for N iterations, compute limits, replace W1 with limits, repeat.
    
    Args:
        num_cycles: Number of training → limit → replacement cycles
        iterations_per_cycle: Iterations per cycle (e.g., 50)
    """
    # Load data
    print("Loading data...")
    X_train = read('../X_train.wav')[1].reshape(-1, 784)
    y_train = (read('../y_train.wav')[1] * 9).astype(int)
    X_test = read('../X_test.wav')[1].reshape(-1, 784)
    y_test = (read('../y_test.wav')[1] * 9).astype(int)
    X20 = X_test[:1000]
    yt20 = y_test[:1000]
    
    print(f"X_train shape: {X_train.shape}")
    
    # Initialize MLP
    learning_rate = np.random.rand(4)
    print(f"Learning rates: {learning_rate}")
    
    f = MLPClassifierWithAutoregressiveLimits(
        input_size=784, 
        hidden_size=100, 
        output_size=10, 
        learning_rate=learning_rate
    )
    
    all_limit_matrices = []
    all_accuracies = []
    
    print(f"\n{'='*70}")
    print(f"ITERATIVE LIMIT REPLACEMENT: {num_cycles} cycles × {iterations_per_cycle} iterations")
    print(f"{'='*70}\n")
    
    for cycle in range(num_cycles):
        print(f"\n{'='*70}")
        print(f"CYCLE {cycle + 1}/{num_cycles}")
        print(f"{'='*70}")
        
        # Train for iterations_per_cycle
        for i in range(iterations_per_cycle):
            idx = np.random.randint(0, 60000, 100)
            X0 = X_train[idx]
            y0 = y_train[idx]
        
            # Sample batch
            for _ in range(10):
                idx = np.random.randint(0, 60000, 100)
                X = X_train[idx]
                yt = y_train[idx]
                
                # Update weights #2
                #for k in range(100):
                f.update(X0, np.eye(10)[y0])
                f.update(X, np.eye(10)[yt])
            
                # Track limits every iteration
            f.track_limits()
        
        # Get accuracy before replacement
        acc = f.score(X20, yt20)
        all_accuracies.append(acc)
        
        # Get limit matrix
        limit_matrix = f.get_limit_matrix()
        all_limit_matrices.append(limit_matrix.copy())
        
        # Print cycle results
        summary = f.w1_tracker.summary()
        print(f"  Accuracy before replacement: {acc:.4f}")
        print(f"  Oscillating: {summary['num_oscillating']} ({summary['pct_oscillating']})")
        print(f"  Converged: {summary['num_converged']} ({summary['pct_converged']})")
        print(f"  Limit matrix - mean: {limit_matrix.mean():.6e}, std: {limit_matrix.std():.6e}")
        
        # Replace W1 with limit matrix
        print(f"  → Replacing W1 with limit_matrix...")
        f.W1 = np.random.uniform(f.W1,limit_matrix.copy())
        
        # Reset tracker for next cycle to detect new oscillations
        f.w1_tracker = AutoregressiveLimitTracker((784, 100), decay=0.99)
        
        # Test accuracy after replacement
        acc_after = f.score(X20, yt20)
        print(f"  Accuracy after replacement: {acc_after:.4f}")
        print(f"  Accuracy delta: {acc_after - acc:+.4f}")
    
    # Final summary
    print(f"\n{'='*70}")
    print(f"FINAL SUMMARY")
    print(f"{'='*70}")
    print(f"Total cycles completed: {num_cycles}")
    print(f"\nAccuracy progression:")
    for i, acc in enumerate(all_accuracies):
        print(f"  Cycle {i+1:2d}: {acc:.4f}")
    
    # Save results
    np.save('W1_limit_matrices_all_cycles.npy', np.array(all_limit_matrices))
    np.save('accuracies_per_cycle.npy', np.array(all_accuracies))
    print(f"\n✓ All limit matrices saved to 'W1_limit_matrices_all_cycles.npy'")
    print(f"✓ Accuracies saved to 'accuracies_per_cycle.npy'")
    
    return f, all_limit_matrices, all_accuracies


def analyze_specific_elements(mlp, elements=None):
    """
    Analyze specific W1[i,j] elements in detail.
    
    Args:
        mlp: Trained MLP with limit tracking
        elements: List of (i,j) tuples to analyze, or None for random sample
    """
    if elements is None:
        # Sample random elements
        np.random.seed(42)
        elements = [(np.random.randint(784), np.random.randint(100)) for _ in range(5)]
    
    print("\n" + "="*70)
    print("DETAILED ELEMENT ANALYSIS")
    print("="*70)
    
    for i, j in elements:
        stats = mlp.w1_tracker.get_element_stats(i, j)
        print(f"\nW1[{i},{j}]:")
        print(f"  Limit estimate:      {stats['limit_estimate']:15.8e}")
        print(f"  Oscillation amp:     {stats['oscillation_amplitude']:15.8e}")
        print(f"  Is oscillating:      {stats['is_oscillating']}")
        print(f"  Is converged:        {stats['is_converged']}")
        print(f"  Current value:       {stats['current_value']}")
        print(f"  Dominant frequency:  {stats['dominant_frequency']}")
        print(f"  Confidence:          {stats['confidence']:.4f}")


if __name__ == "__main__":
    # Iterative limit replacement: train 50 iters → replace W1 with limits → repeat
    mlp, all_limits, all_accuracies = iterative_limit_replacement(
        num_cycles=100,
        iterations_per_cycle=30
    )
    
    print("\nAnalysis complete!")