import numpy as np
from scipy.io.wavfile import read
from scipy.signal import find_peaks
from scipy.stats import linregress
import matplotlib.pyplot as plt


class OscillationAnalyzer:
    """Analyze oscillation behavior in weight updates during MLP training."""
    
    def __init__(self, history, sample_interval=1):
        """
        Initialize analyzer with weight history.
        
        Args:
            history: Array of weight values over time
            sample_interval: Interval between samples
        """
        self.history = np.array(history)
        self.sample_interval = sample_interval
        self.time = np.arange(len(history)) * sample_interval
        
    def detect_oscillation(self, min_amplitude=1e-6, min_peaks=3):
        """
        Detect if the weight history shows oscillatory behavior.
        
        Args:
            min_amplitude: Minimum amplitude to consider as oscillation
            min_peaks: Minimum number of peaks to confirm oscillation
            
        Returns:
            dict with oscillation characteristics
        """
        if len(self.history) < 10:
            return {'is_oscillating': False, 'reason': 'Insufficient data'}
        
        # Find peaks
        peaks, properties = find_peaks(self.history, distance=2)
        troughs, _ = find_peaks(-self.history, distance=2)
        
        num_peaks = len(peaks)
        num_troughs = len(troughs)
        
        # Calculate amplitude
        if num_peaks >= 2:
            peak_values = self.history[peaks]
            amplitude = np.std(peak_values)
        else:
            amplitude = 0
        
        # Check for oscillation
        is_oscillating = (num_peaks >= min_peaks and amplitude > min_amplitude)
        
        return {
            'is_oscillating': is_oscillating,
            'num_peaks': num_peaks,
            'num_troughs': num_troughs,
            'amplitude': amplitude,
            'mean': np.mean(self.history),
            'std': np.std(self.history)
        }
    
    def estimate_limit_mean(self, window_size=50):
        """
        Estimate the limit value using converging moving averages.
        
        Args:
            window_size: Size of the moving average window
            
        Returns:
            Estimated limit value (mean of oscillation)
        """
        if len(self.history) < window_size:
            return np.mean(self.history)
        
        # Calculate moving average
        moving_avg = np.convolve(self.history, 
                                 np.ones(window_size)/window_size, 
                                 mode='valid')
        
        # Use last portion for more stable estimate
        last_portion = moving_avg[-min(100, len(moving_avg)):]
        
        return np.mean(last_portion)
    
    def estimate_limit_extrapolation(self, degree=2):
        """
        Estimate limit by fitting polynomial and extrapolating to infinity.
        
        Args:
            degree: Polynomial degree for fitting
            
        Returns:
            Estimated limit value
        """
        if len(self.history) < degree + 1:
            return np.mean(self.history)
        
        # Use indices as x values
        x = np.arange(len(self.history))
        
        # Fit polynomial
        coeffs = np.polyfit(x, self.history, degree)
        poly = np.poly1d(coeffs)
        
        # For oscillating series, the limit is the center of oscillation
        # Use the polynomial's behavior at the end
        last_x = x[-1]
        
        # If linear or constant trend, evaluate at mean
        if degree <= 1:
            return np.mean(self.history[-100:]) if len(self.history) > 100 else np.mean(self.history)
        
        # For higher degrees, use the polynomial's asymptotic behavior
        # For oscillation, we want the mean value
        return np.mean(self.history)
    
    def frequency_analysis(self):
        """
        Perform FFT to find dominant frequency of oscillation.
        
        Returns:
            dict with frequency characteristics
        """
        if len(self.history) < 4:
            return {'dominant_freq': None, 'power': None}
        
        # Remove mean
        signal = self.history - np.mean(self.history)
        
        # FFT
        fft_values = np.fft.rfft(signal)
        fft_freqs = np.fft.rfftfreq(len(signal))
        fft_power = np.abs(fft_values) ** 2
        
        # Find dominant frequency (exclude DC component)
        if len(fft_power) > 1:
            dominant_idx = np.argmax(fft_power[1:]) + 1
            dominant_freq = fft_freqs[dominant_idx]
            dominant_power = fft_power[dominant_idx]
        else:
            dominant_freq = None
            dominant_power = None
        
        return {
            'dominant_freq': dominant_freq,
            'dominant_period': 1/dominant_freq if dominant_freq and dominant_freq > 0 else None,
            'power': dominant_power
        }
    
    def calculate_convergence_limit(self):
        """
        Comprehensive limit calculation combining multiple methods.
        
        Returns:
            dict with limit estimates and confidence metrics
        """
        oscillation_info = self.detect_oscillation()
        
        results = {
            'is_oscillating': oscillation_info['is_oscillating'],
            'oscillation_amplitude': oscillation_info['amplitude'],
            'mean_value': oscillation_info['mean'],
            'std_value': oscillation_info['std'],
        }
        
        if oscillation_info['is_oscillating']:
            # For oscillating series, limit is the center of oscillation
            results['limit_estimate'] = self.estimate_limit_mean()
            results['limit_method'] = 'converging_moving_average'
            
            # Frequency analysis
            freq_info = self.frequency_analysis()
            results['oscillation_frequency'] = freq_info['dominant_freq']
            results['oscillation_period'] = freq_info['dominant_period']
            
            # Confidence based on stability of recent values
            if len(self.history) > 100:
                recent_std = np.std(self.history[-100:])
                overall_std = np.std(self.history)
                results['confidence'] = 1 - (recent_std / (overall_std + 1e-10))
            else:
                results['confidence'] = 0.5
        else:
            # For non-oscillating, check convergence
            results['limit_estimate'] = self.estimate_limit_extrapolation()
            results['limit_method'] = 'polynomial_extrapolation'
            results['confidence'] = 0.7
        
        return results


def analyze_weight_oscillation(mlp, X_history, yt_history, layer='W1', element=(0, 0)):
    """
    Analyze oscillation of a specific weight element during training.
    
    Args:
        mlp: Trained MLP classifier
        X_history: List of X batches used for training
        yt_history: List of yt batches used for training
        layer: Which weight layer to analyze ('W1', 'W2', 'b1', 'b2')
        element: Tuple of indices for the specific element
        
    Returns:
        Analysis results
    """
    # We need to re-simulate training to get weight history
    # This requires modifying the original training loop to record weights
    
    # For now, analyze the current weight value
    weight_matrix = getattr(mlp, layer)
    current_value = weight_matrix[element]
    
    return {
        'current_value': current_value,
        'note': 'Full history analysis requires recording weights during training'
    }


class MLPClassifierWithTracking:
    """MLP Classifier that tracks weight history for oscillation analysis."""
    
    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
        
        # Tracking
        self.W1_history = []
        self.b1_history = []
        self.W2_history = []
        self.b2_history = []
        self.loss_history = []
        self.accuracy_history = []
        
    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)
    
    def record_weights(self):
        """Record current weight values to history."""
        self.W1_history.append(self.W1.copy())
        self.b1_history.append(self.b1.copy())
        self.W2_history.append(self.W2.copy())
        self.b2_history.append(self.b2.copy())
    
    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 track_and_analyze_oscillation(num_iterations=1000, track_element_W1=(0, 0)):
    """
    Main function to track weight oscillation and analyze limits.
    
    Args:
        num_iterations: Number of training iterations
        track_element_W1: Which W1 element to track (row, col)
        
    Returns:
        Analyzer and 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]
    
    # Initialize MLP with tracking
    learning_rate = np.random.rand(4)
    f = MLPClassifierWithTracking(input_size=784, hidden_size=100, 
                                   output_size=10, learning_rate=learning_rate)
    
    print(f"Training for {num_iterations} iterations...")
    print(f"Tracking W1{track_element_W1} element...")
    
    # Training loop with tracking
    for i in range(num_iterations):
        idx = np.random.randint(0, 60000, 100)
        X = X_train[idx]
        yt = y_train[idx]
        
        for k in range(10):
            f.update(X, np.eye(10)[yt])
        
        # Record weights every iteration
        f.record_weights()
        
        # Print progress
        if i % 100 == 0:
            acc = f.score(X20, yt20)
            print(f"Iteration {i}, Accuracy: {acc:.4f}, Loss: {f.loss_history[-1]:.4f}")
    
    # Extract history for the tracked element
    W1_tracked_history = np.array([w[track_element_W1] for w in f.W1_history])
    
    # Analyze oscillation
    print("\nAnalyzing oscillation...")
    analyzer = OscillationAnalyzer(W1_tracked_history)
    results = analyzer.calculate_convergence_limit()
    
    # Print results
    print("\n" + "="*60)
    print(f"OSCILLATION ANALYSIS FOR W1{track_element_W1}")
    print("="*60)
    print(f"Is oscillating: {results['is_oscillating']}")
    print(f"Oscillation amplitude: {results['oscillation_amplitude']:.6e}")
    print(f"Mean value: {results['mean_value']:.6e}")
    print(f"Std deviation: {results['std_value']:.6e}")
    print(f"Estimated limit: {results['limit_estimate']:.6e}")
    print(f"Method: {results['limit_method']}")
    print(f"Confidence: {results['confidence']:.4f}")
    
    if results['is_oscillating']:
        print(f"Oscillation frequency: {results['oscillation_frequency']}")
        if results['oscillation_period']:
            print(f"Oscillation period: {results['oscillation_period']:.2f} iterations")
    
    return f, analyzer, results, W1_tracked_history


if __name__ == "__main__":
    # Run analysis
    mlp, analyzer, results, history = track_and_analyze_oscillation(
        num_iterations=1000,
        track_element_W1=(0, 0)
    )
    
    print("\nAnalysis complete!")
    print(f"Final limit estimate for W1[0,0]: {results['limit_estimate']:.8f}")