"""
Test audio compression and decompression roundtrip.
"""

import numpy as np
import pytest
from pathlib import Path

# Skip tests if audio libraries not available
try:
    import librosa
    import soundfile as sf
    from sempress import (
        encode_audio_from_array,
        decode_audio,
        AudioEncodeConfig,
        snr,
        eval_audio
    )
    AUDIO_AVAILABLE = True
except ImportError:
    AUDIO_AVAILABLE = False


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_roundtrip_kmeans():
    """Test basic audio compression with K-Means."""
    # Create a simple test audio (sine wave)
    duration = 1.0
    sr = 16000
    freq = 440  # A4 note
    t = np.linspace(0, duration, int(sr * duration))
    audio = np.sin(2 * np.pi * freq * t).astype(np.float32)
    
    # Configure compression
    config = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=64,  # Small for fast test
        method="kmeans",
        use_spectral=True,
        n_mfcc=13
    )
    
    # Compress
    compressed = encode_audio_from_array(audio, sr, config)
    
    # Decompress
    reconstructed, sr_recon = decode_audio(compressed)
    
    # Check sample rate matches
    assert sr_recon == sr
    
    # Check compression worked
    original_size = len(audio) * 2  # 16-bit
    assert len(compressed) < original_size
    assert len(compressed) > 0
    
    # Check quality (should be reasonable for sine wave)
    snr_val = snr(audio, reconstructed)
    assert snr_val > 10  # At least 10 dB SNR


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_roundtrip_gmm():
    """Test audio compression with GMM."""
    duration = 1.0
    sr = 16000
    t = np.linspace(0, duration, int(sr * duration))
    
    # More complex signal (chord)
    audio = (np.sin(2 * np.pi * 440 * t) + 
             np.sin(2 * np.pi * 554 * t) +
             np.sin(2 * np.pi * 659 * t)).astype(np.float32) / 3
    
    config = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=64,
        method="gmm",
        use_spectral=True,
        n_mfcc=13
    )
    
    compressed = encode_audio_from_array(audio, sr, config)
    reconstructed, sr_recon = decode_audio(compressed)
    
    assert sr_recon == sr
    assert len(compressed) > 0
    
    # Should achieve some compression
    compression_ratio = (len(audio) * 2) / len(compressed)
    assert compression_ratio > 1.0


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_roundtrip_pca():
    """Test audio compression with PCA."""
    duration = 1.0
    sr = 16000
    t = np.linspace(0, duration, int(sr * duration))
    audio = np.sin(2 * np.pi * 440 * t).astype(np.float32)
    
    config = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=64,
        method="pca",
        use_spectral=True,
        n_mfcc=13
    )
    
    compressed = encode_audio_from_array(audio, sr, config)
    reconstructed, sr_recon = decode_audio(compressed)
    
    assert sr_recon == sr
    assert len(compressed) > 0


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_stereo():
    """Test stereo audio compression."""
    duration = 1.0
    sr = 16000
    t = np.linspace(0, duration, int(sr * duration))
    
    # Create stereo signal
    left = np.sin(2 * np.pi * 440 * t).astype(np.float32)
    right = np.sin(2 * np.pi * 554 * t).astype(np.float32)
    audio_stereo = np.stack([left, right], axis=1)
    
    config = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=64,
        method="kmeans",
        use_spectral=True
    )
    
    compressed = encode_audio_from_array(audio_stereo, sr, config)
    reconstructed, sr_recon = decode_audio(compressed)
    
    # Should reconstruct as stereo
    assert reconstructed.ndim == 2
    assert reconstructed.shape[1] == 2
    assert sr_recon == sr


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_waveform_features():
    """Test compression with raw waveform features (no spectral)."""
    duration = 0.5  # Shorter for waveform mode
    sr = 16000
    t = np.linspace(0, duration, int(sr * duration))
    audio = np.sin(2 * np.pi * 440 * t).astype(np.float32)
    
    config = AudioEncodeConfig(
        frame_size=512,
        hop_length=256,
        n_components=32,
        method="kmeans",
        use_spectral=False  # Use raw waveform
    )
    
    compressed = encode_audio_from_array(audio, sr, config)
    reconstructed, sr_recon = decode_audio(compressed)
    
    assert sr_recon == sr
    assert len(compressed) > 0
    
    # Waveform mode should reconstruct length exactly
    assert len(reconstructed) == len(audio)


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_residuals():
    """Test residual storage for high-error frames."""
    duration = 1.0
    sr = 16000
    t = np.linspace(0, duration, int(sr * duration))
    
    # Noisy signal (harder to compress)
    np.random.seed(42)
    audio = (np.sin(2 * np.pi * 440 * t) + 
             np.random.normal(0, 0.2, len(t))).astype(np.float32)
    
    # Low threshold = more residuals
    config_with_residuals = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=32,
        method="kmeans",
        residual_threshold=0.05,  # Low threshold
        use_spectral=True
    )
    
    # High threshold = fewer residuals
    config_without_residuals = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=32,
        method="kmeans",
        residual_threshold=0.5,  # High threshold
        use_spectral=True
    )
    
    compressed_with = encode_audio_from_array(audio, sr, config_with_residuals)
    compressed_without = encode_audio_from_array(audio, sr, config_without_residuals)
    
    # Version with residuals should be larger
    assert len(compressed_with) >= len(compressed_without)
    
    recon_with, _ = decode_audio(compressed_with)
    recon_without, _ = decode_audio(compressed_without)
    
    # More residuals should give better or equal quality
    snr_with = snr(audio, recon_with)
    snr_without = snr(audio, recon_without)
    
    # Allow small margin
    assert snr_with >= snr_without - 2.0


@pytest.mark.skipif(not AUDIO_AVAILABLE, reason="Audio libraries not installed")
def test_audio_eval_metrics():
    """Test audio evaluation metrics."""
    duration = 1.0
    sr = 16000
    t = np.linspace(0, duration, int(sr * duration))
    audio = np.sin(2 * np.pi * 440 * t).astype(np.float32)
    
    config = AudioEncodeConfig(
        frame_size=1024,
        hop_length=256,
        n_components=64,
        method="kmeans",
        use_spectral=True
    )
    
    compressed = encode_audio_from_array(audio, sr, config)
    reconstructed, _ = decode_audio(compressed)
    
    # Evaluate metrics
    metrics = eval_audio(audio, reconstructed, len(compressed), sr)
    
    # Check all expected metrics exist
    assert "compression_ratio" in metrics
    assert "snr_db" in metrics
    assert "rmse" in metrics
    assert "mae" in metrics
    assert "spectral_convergence" in metrics
    assert "log_spectral_distance" in metrics
    assert "bitrate_kbps" in metrics
    
    # Check reasonable values
    assert metrics["compression_ratio"] > 1.0
    assert metrics["snr_db"] > 0
    assert metrics["rmse"] >= 0
    assert metrics["bitrate_kbps"] > 0


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
