#!/usr/bin/env python3
"""
Script to compare mel spectrograms generated by Python and Rust implementations
"""
import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import pearsonr
from sklearn.metrics import mean_squared_error, mean_absolute_error

def load_data():
    """Load the mel spectrograms from both implementations"""
    base_path = os.path.join(os.path.dirname(__file__), '..', 'temp')
    
    python_path = os.path.join(base_path, 'python_mel.npy')
    rust_path = os.path.join(base_path, 'rust_mel.npy')
    
    python_mel = None
    rust_mel = None
    
    if os.path.exists(python_path):
        python_mel = np.load(python_path)
        print(f"Loaded Python mel spectrogram: {python_mel.shape}")
    else:
        print(f"Python mel spectrogram not found at: {python_path}")
    
    if os.path.exists(rust_path):
        rust_mel = np.load(rust_path)
        print(f"Loaded Rust mel spectrogram: {rust_mel.shape}")
    else:
        print(f"Rust mel spectrogram not found at: {rust_path}")
    
    return python_mel, rust_mel

def compare_shapes(python_mel, rust_mel):
    """Compare the shapes of the two mel spectrograms"""
    print("\n=== Shape Comparison ===")
    print(f"Python shape: {python_mel.shape}")
    print(f"Rust shape: {rust_mel.shape}")
    
    if python_mel.shape == rust_mel.shape:
        print("✓ Shapes match exactly")
        return True
    else:
        print("✗ Shapes do not match")
        
        # Try to find the minimum common shape
        min_dims = min(len(python_mel.shape), len(rust_mel.shape))
        common_shape = []
        for i in range(min_dims):
            common_shape.append(min(python_mel.shape[i], rust_mel.shape[i]))
        
        print(f"Common shape for comparison: {common_shape}")
        return False

def compare_statistics(python_mel, rust_mel):
    """Compare basic statistics"""
    print("\n=== Statistical Comparison ===")
    
    stats = ['min', 'max', 'mean', 'std']
    
    print(f"{'Statistic':<10} {'Python':<15} {'Rust':<15} {'Difference':<15}")
    print("-" * 60)
    
    for stat in stats:
        py_val = getattr(python_mel, stat)()
        rust_val = getattr(rust_mel, stat)()
        diff = abs(py_val - rust_val)
        
        print(f"{stat:<10} {py_val:<15.6f} {rust_val:<15.6f} {diff:<15.6f}")

def compare_values(python_mel, rust_mel, tolerance=1e-3):
    """Compare actual values with tolerance"""
    print(f"\n=== Value Comparison (tolerance: {tolerance}) ===")
    
    # Ensure same shape for comparison
    if python_mel.shape != rust_mel.shape:
        # Crop to minimum common shape
        min_shape = tuple(min(p, r) for p, r in zip(python_mel.shape, rust_mel.shape))
        
        if len(min_shape) == 2:
            python_mel = python_mel[:min_shape[0], :min_shape[1]]
            rust_mel = rust_mel[:min_shape[0], :min_shape[1]]
        elif len(min_shape) == 3:
            python_mel = python_mel[:min_shape[0], :min_shape[1], :min_shape[2]]
            rust_mel = rust_mel[:min_shape[0], :min_shape[1], :min_shape[2]]
        
        print(f"Cropped to common shape: {min_shape}")
    
    # Calculate differences
    diff = np.abs(python_mel - rust_mel)
    
    # Statistics on differences
    print(f"Absolute difference - min: {diff.min():.6f}, max: {diff.max():.6f}")
    print(f"Absolute difference - mean: {diff.mean():.6f}, std: {diff.std():.6f}")
    
    # Check tolerance
    within_tolerance = np.sum(diff <= tolerance)
    total_elements = diff.size
    percentage = (within_tolerance / total_elements) * 100
    
    print(f"Elements within tolerance: {within_tolerance}/{total_elements} ({percentage:.2f}%)")
    
    # Calculate correlation
    try:
        correlation, p_value = pearsonr(python_mel.flatten(), rust_mel.flatten())
        print(f"Pearson correlation: {correlation:.6f} (p-value: {p_value:.2e})")
    except:
        print("Could not calculate correlation")
    
    # Calculate error metrics
    try:
        mse = mean_squared_error(python_mel.flatten(), rust_mel.flatten())
        mae = mean_absolute_error(python_mel.flatten(), rust_mel.flatten())
        print(f"Mean Squared Error: {mse:.6f}")
        print(f"Mean Absolute Error: {mae:.6f}")
    except:
        print("Could not calculate error metrics")

def visualize_comparison(python_mel, rust_mel):
    """Create visualizations for comparison"""
    print("\n=== Generating Visualizations ===")
    
    # Ensure same shape for visualization
    if python_mel.shape != rust_mel.shape:
        min_shape = tuple(min(p, r) for p, r in zip(python_mel.shape, rust_mel.shape))
        
        if len(min_shape) == 2:
            python_mel = python_mel[:min_shape[0], :min_shape[1]]
            rust_mel = rust_mel[:min_shape[0], :min_shape[1]]
        elif len(min_shape) == 3:
            # Take first channel if 3D
            python_mel = python_mel[0] if len(python_mel.shape) == 3 else python_mel
            rust_mel = rust_mel[0] if len(rust_mel.shape) == 3 else rust_mel
    
    # Create comparison plots
    fig, axes = plt.subplots(2, 2, figsize=(15, 10))
    
    # Python mel spectrogram
    im1 = axes[0, 0].imshow(python_mel, aspect='auto', origin='lower')
    axes[0, 0].set_title('Python Mel Spectrogram')
    axes[0, 0].set_xlabel('Time')
    axes[0, 0].set_ylabel('Mel Bins')
    plt.colorbar(im1, ax=axes[0, 0])
    
    # Rust mel spectrogram
    im2 = axes[0, 1].imshow(rust_mel, aspect='auto', origin='lower')
    axes[0, 1].set_title('Rust Mel Spectrogram')
    axes[0, 1].set_xlabel('Time')
    axes[0, 1].set_ylabel('Mel Bins')
    plt.colorbar(im2, ax=axes[0, 1])
    
    # Difference
    diff = python_mel - rust_mel
    im3 = axes[1, 0].imshow(diff, aspect='auto', origin='lower', cmap='RdBu_r')
    axes[1, 0].set_title('Difference (Python - Rust)')
    axes[1, 0].set_xlabel('Time')
    axes[1, 0].set_ylabel('Mel Bins')
    plt.colorbar(im3, ax=axes[1, 0])
    
    # Histogram of differences
    axes[1, 1].hist(diff.flatten(), bins=50, alpha=0.7)
    axes[1, 1].set_title('Histogram of Differences')
    axes[1, 1].set_xlabel('Difference Value')
    axes[1, 1].set_ylabel('Frequency')
    
    plt.tight_layout()
    
    # Save the plot
    output_path = os.path.join(os.path.dirname(__file__), '..', 'temp', 'mel_comparison.png')
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"Saved comparison plot to: {output_path}")
    
    # Show plot if in interactive mode
    try:
        plt.show()
    except:
        print("Could not display plot (non-interactive environment)")

def main():
    print("Mel Spectrogram Comparison Tool")
    print("=" * 50)
    
    # Load data
    python_mel, rust_mel = load_data()
    
    if python_mel is None or rust_mel is None:
        print("Cannot proceed without both mel spectrograms.")
        print("Please run test_mel.py and test_mel_rust.py first.")
        return
    
    # Compare shapes
    shapes_match = compare_shapes(python_mel, rust_mel)
    
    # Compare statistics
    compare_statistics(python_mel, rust_mel)
    
    # Compare values
    compare_values(python_mel, rust_mel)
    
    # Generate visualizations
    try:
        visualize_comparison(python_mel, rust_mel)
    except ImportError:
        print("Matplotlib not available, skipping visualizations")
    except Exception as e:
        print(f"Error generating visualizations: {e}")
    
    # Save comparison report
    report_path = os.path.join(os.path.dirname(__file__), '..', 'temp', 'comparison_report.txt')
    with open(report_path, 'w') as f:
        f.write("Mel Spectrogram Comparison Report\n")
        f.write("=" * 40 + "\n\n")
        f.write(f"Python shape: {python_mel.shape}\n")
        f.write(f"Rust shape: {rust_mel.shape}\n")
        f.write(f"Shapes match: {shapes_match}\n\n")
        
        stats = ['min', 'max', 'mean', 'std']
        f.write("Statistics:\n")
        for stat in stats:
            py_val = getattr(python_mel, stat)()
            rust_val = getattr(rust_mel, stat)()
            diff = abs(py_val - rust_val)
            f.write(f"  {stat}: Python={py_val:.6f}, Rust={rust_val:.6f}, Diff={diff:.6f}\n")
    
    print(f"\nSaved comparison report to: {report_path}")

if __name__ == "__main__":
    main()
