"""
Audio quality metrics for Sempress audio compression.

Provides comprehensive metrics for evaluating audio compression quality.
"""

import numpy as np
from typing import Dict, Tuple


def snr(original: np.ndarray, reconstructed: np.ndarray) -> float:
    """
    Calculate Signal-to-Noise Ratio (SNR) in dB.
    
    Higher is better. >30 dB is excellent, 20-30 dB is good.
    """
    signal_power = np.mean(original ** 2)
    noise_power = np.mean((original - reconstructed) ** 2)
    
    if noise_power == 0:
        return float('inf')
    
    snr_val = 10 * np.log10(signal_power / noise_power)
    return float(snr_val)


def rmse(original: np.ndarray, reconstructed: np.ndarray) -> float:
    """
    Calculate Root Mean Square Error.
    
    Lower is better. Measures average magnitude of error.
    """
    return float(np.sqrt(np.mean((original - reconstructed) ** 2)))


def mae(original: np.ndarray, reconstructed: np.ndarray) -> float:
    """
    Calculate Mean Absolute Error.
    
    Lower is better. Less sensitive to outliers than RMSE.
    """
    return float(np.mean(np.abs(original - reconstructed)))


def pesq_score(original: np.ndarray, reconstructed: np.ndarray, sr: int) -> float:
    """
    Calculate PESQ (Perceptual Evaluation of Speech Quality) score.
    
    Range: -0.5 to 4.5 (higher is better)
    Only works for speech audio at 8kHz or 16kHz.
    
    Note: Requires pesq library. Returns None if not available.
    """
    try:
        from pesq import pesq as pesq_calc
        
        # PESQ requires specific sample rates
        if sr not in [8000, 16000]:
            return None
        
        # Ensure same length
        min_len = min(len(original), len(reconstructed))
        original = original[:min_len]
        reconstructed = reconstructed[:min_len]
        
        # Calculate PESQ
        score = pesq_calc(sr, original, reconstructed, 'wb' if sr == 16000 else 'nb')
        return float(score)
    except ImportError:
        return None
    except Exception:
        return None


def spectral_convergence(original: np.ndarray, reconstructed: np.ndarray) -> float:
    """
    Calculate Spectral Convergence (lower is better).
    
    Measures how well the frequency content is preserved.
    """
    # Compute FFT
    orig_fft = np.fft.rfft(original)
    recon_fft = np.fft.rfft(reconstructed[:len(original)])
    
    orig_mag = np.abs(orig_fft)
    recon_mag = np.abs(recon_fft)
    
    # Spectral convergence
    numerator = np.linalg.norm(orig_mag - recon_mag)
    denominator = np.linalg.norm(orig_mag)
    
    if denominator == 0:
        return 0.0
    
    return float(numerator / denominator)


def log_spectral_distance(original: np.ndarray, reconstructed: np.ndarray) -> float:
    """
    Calculate Log Spectral Distance (LSD) in dB.
    
    Lower is better. Measures perceptual distance in frequency domain.
    """
    # Compute FFT
    orig_fft = np.fft.rfft(original)
    recon_fft = np.fft.rfft(reconstructed[:len(original)])
    
    orig_mag = np.abs(orig_fft) + 1e-10
    recon_mag = np.abs(recon_fft) + 1e-10
    
    # Log spectral distance
    lsd = np.sqrt(np.mean((np.log10(orig_mag) - np.log10(recon_mag)) ** 2))
    
    return float(lsd)


def eval_audio(original: np.ndarray, reconstructed: np.ndarray, 
               compressed_size: int, sr: int) -> Dict:
    """
    Comprehensive audio quality evaluation.
    
    Args:
        original: Original audio waveform
        reconstructed: Reconstructed audio waveform
        compressed_size: Size of compressed data in bytes
        sr: Sample rate
        
    Returns:
        Dictionary with all quality metrics
    """
    # Handle stereo by averaging across channels
    if original.ndim == 2:
        original = np.mean(original, axis=1)
    if reconstructed.ndim == 2:
        reconstructed = np.mean(reconstructed, axis=1)
    
    # Ensure same length
    min_len = min(len(original), len(reconstructed))
    original = original[:min_len]
    reconstructed = reconstructed[:min_len]
    
    # Calculate original size (uncompressed)
    bytes_per_sample = 2  # 16-bit audio
    n_channels = 1  # Already averaged if stereo
    original_size = min_len * bytes_per_sample * n_channels
    
    # Basic metrics
    metrics = {
        "original_size_bytes": original_size,
        "compressed_size_bytes": compressed_size,
        "compression_ratio": original_size / compressed_size,
        "space_saving_pct": (1 - compressed_size / original_size) * 100,
        
        # Quality metrics
        "snr_db": snr(original, reconstructed),
        "rmse": rmse(original, reconstructed),
        "mae": mae(original, reconstructed),
        "spectral_convergence": spectral_convergence(original, reconstructed),
        "log_spectral_distance": log_spectral_distance(original, reconstructed),
        
        # Bitrate
        "bitrate_kbps": (compressed_size * 8 * sr) / (min_len * 1000),
        "bits_per_sample": (compressed_size * 8) / min_len,
        
        # Audio info
        "sample_rate": sr,
        "duration_sec": min_len / sr,
        "n_samples": min_len
    }
    
    # Try to calculate PESQ if available
    pesq_val = pesq_score(original, reconstructed, sr)
    if pesq_val is not None:
        metrics["pesq_score"] = pesq_val
    
    return metrics


def compare_methods(original: np.ndarray, results: Dict[str, Tuple[np.ndarray, int]], sr: int) -> Dict:
    """
    Compare multiple compression methods.
    
    Args:
        original: Original audio waveform
        results: Dict mapping method name to (reconstructed, compressed_size)
        sr: Sample rate
        
    Returns:
        Dictionary with comparison results
    """
    comparison = {}
    
    for method_name, (reconstructed, compressed_size) in results.items():
        metrics = eval_audio(original, reconstructed, compressed_size, sr)
        comparison[method_name] = metrics
    
    return comparison


def print_audio_metrics(metrics: Dict, title: str = "Audio Metrics"):
    """Pretty print audio metrics."""
    print(f"\n{'='*60}")
    print(f"{title}")
    print(f"{'='*60}")
    print(f"Compression Ratio:     {metrics['compression_ratio']:.2f}×")
    print(f"Space Saving:          {metrics['space_saving_pct']:.1f}%")
    print(f"Bitrate:               {metrics['bitrate_kbps']:.1f} kbps")
    print(f"Bits per Sample:       {metrics['bits_per_sample']:.2f}")
    print(f"\nQuality Metrics:")
    print(f"  SNR:                 {metrics['snr_db']:.2f} dB")
    print(f"  RMSE:                {metrics['rmse']:.6f}")
    print(f"  MAE:                 {metrics['mae']:.6f}")
    print(f"  Spectral Conv:       {metrics['spectral_convergence']:.4f}")
    print(f"  Log Spectral Dist:   {metrics['log_spectral_distance']:.4f}")
    
    if 'pesq_score' in metrics:
        print(f"  PESQ Score:          {metrics['pesq_score']:.2f}")
    
    print(f"\nAudio Info:")
    print(f"  Duration:            {metrics['duration_sec']:.2f} sec")
    print(f"  Sample Rate:         {metrics['sample_rate']} Hz")
    print(f"  Samples:             {metrics['n_samples']:,}")
    print(f"{'='*60}")
