//! Integration tests for vocoder modules

use rtx_tensor::{Device, Tensor};
use rtx_tts::vocoder::{
    GriffinLim, GriffinLimConfig, HiFiGAN, HiFiGANConfig, MelConfig, MelSpectrogram, Vocoder,
};

#[test]
fn test_mel_spectrogram_basic() {
    let device = Device::cpu();
    let config = MelConfig::default();
    let mel_proc = MelSpectrogram::new(config.clone(), &device).unwrap();

    // Test audio to mel conversion
    let audio_len = config.sample_rate;
    let audio = Tensor::randn(&[audio_len], &device).unwrap();
    let mel_spec = mel_proc.audio_to_mel(&audio).unwrap();

    assert_eq!(mel_spec.dims()[0], config.n_mels);
}

#[test]
fn test_griffin_lim_basic() {
    let device = Device::cpu();
    let config = GriffinLimConfig {
        n_fft: 512,
        hop_length: 128,
        win_length: 512,
        n_iter: 5, // Fewer iterations for faster test
        sample_rate: 16000,
        mel_config: MelConfig {
            sample_rate: 16000,
            n_fft: 512,
            hop_length: 128,
            win_length: 512,
            n_mels: 40,
            f_min: 0.0,
            f_max: None,
        },
    };

    let vocoder = GriffinLim::new(config.clone(), &device).unwrap();
    assert_eq!(vocoder.get_sample_rate(), 16000);

    // Test synthesis
    let mel = Tensor::randn(&[40, 30], &device).unwrap();
    let audio = vocoder.synthesize(&mel).unwrap();
    assert!(audio.dims()[0] > 0);
}

#[test]
fn test_hifigan_basic() {
    let device = Device::cpu();
    let config = HiFiGANConfig {
        upsample_rates: vec![2, 2],
        upsample_kernel_sizes: vec![4, 4],
        resblock_kernel_sizes: vec![3],
        resblock_dilation_sizes: vec![vec![1]],
        initial_channel: 64,
        mel_channels: 40,
        sample_rate: 16000,
    };

    let vocoder = HiFiGAN::new(config, &device).unwrap();
    assert_eq!(vocoder.get_sample_rate(), 16000);

    // Test synthesis
    let mel = Tensor::randn(&[40, 20], &device).unwrap();
    let audio = vocoder.synthesize(&mel).unwrap();
    assert!(audio.dims()[0] > 0);
}

#[test]
fn test_mel_round_trip() {
    let device = Device::cpu();
    let config = MelConfig {
        sample_rate: 16000,
        n_fft: 512,
        hop_length: 128,
        win_length: 512,
        n_mels: 40,
        f_min: 0.0,
        f_max: None,
    };
    let mel_proc = MelSpectrogram::new(config.clone(), &device).unwrap();

    // Create a linear spectrogram
    let n_freqs = config.n_fft / 2 + 1;
    let n_frames = 50;
    let linear_spec = Tensor::randn(&[n_freqs, n_frames], &device).unwrap();

    // Convert to mel and back
    let mel_spec = mel_proc.spectrogram_to_mel(&linear_spec).unwrap();
    let reconstructed = mel_proc.mel_to_spectrogram(&mel_spec).unwrap();

    // Check shapes match
    assert_eq!(linear_spec.dims(), reconstructed.dims());
}

#[test]
fn test_full_pipeline_griffin_lim() {
    let device = Device::cpu();

    // Create mel processor
    let mel_config = MelConfig {
        sample_rate: 16000,
        n_fft: 512,
        hop_length: 128,
        win_length: 512,
        n_mels: 40,
        f_min: 0.0,
        f_max: None,
    };
    let mel_proc = MelSpectrogram::new(mel_config.clone(), &device).unwrap();

    // Create vocoder
    let gl_config = GriffinLimConfig {
        n_fft: 512,
        hop_length: 128,
        win_length: 512,
        n_iter: 10,
        sample_rate: 16000,
        mel_config: mel_config.clone(),
    };
    let vocoder = GriffinLim::new(gl_config, &device).unwrap();

    // Create audio, convert to mel, then back to audio
    let audio_len = 16000; // 1 second
    let original_audio = Tensor::randn(&[audio_len], &device).unwrap();

    let mel_spec = mel_proc.audio_to_mel(&original_audio).unwrap();
    let reconstructed_audio = vocoder.synthesize(&mel_spec).unwrap();

    // Check that audio was produced
    assert!(reconstructed_audio.dims()[0] > 0);
}

#[test]
fn test_vocoder_trait_polymorphism() {
    let device = Device::cpu();

    // Test that both vocoders implement the Vocoder trait
    let gl_config = GriffinLimConfig {
        n_fft: 256,
        hop_length: 64,
        win_length: 256,
        n_iter: 5,
        sample_rate: 16000,
        mel_config: MelConfig {
            sample_rate: 16000,
            n_fft: 256,
            hop_length: 64,
            win_length: 256,
            n_mels: 20,
            f_min: 0.0,
            f_max: None,
        },
    };
    let gl = GriffinLim::new(gl_config, &device).unwrap();

    let hg_config = HiFiGANConfig {
        upsample_rates: vec![2, 2],
        upsample_kernel_sizes: vec![4, 4],
        resblock_kernel_sizes: vec![3],
        resblock_dilation_sizes: vec![vec![1]],
        initial_channel: 64,
        mel_channels: 20,
        sample_rate: 16000,
    };
    let hg = HiFiGAN::new(hg_config, &device).unwrap();

    // Test both through the trait
    let mel = Tensor::randn(&[20, 30], &device).unwrap();

    let audio_gl = gl.synthesize(&mel).unwrap();
    let audio_hg = hg.synthesize(&mel).unwrap();

    assert!(audio_gl.dims()[0] > 0);
    assert!(audio_hg.dims()[0] > 0);
    assert_eq!(gl.get_sample_rate(), 16000);
    assert_eq!(hg.get_sample_rate(), 16000);
}