rustytorch/crates/models/rtx-tts/src/vocoder/hifigan.rs

//! HiFi-GAN: Generative Adversarial Network for Efficient and High Fidelity Speech Synthesis
//!
//! This module implements the HiFi-GAN generator architecture for neural vocoding.
//! The generator uses transposed convolutions for upsampling and multi-receptive field
//! fusion (MRF) modules for high-quality audio synthesis.
//!
//! Reference: "HiFi-GAN: Generative Adversarial Networks for Efficient and High Fidelity
//! Speech Synthesis" (Kong et al., 2020)

use super::Vocoder;
use crate::error::{Result, TtsError};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};

/// Configuration for HiFi-GAN vocoder
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HiFiGANConfig {
    /// Upsampling rates for each transposed convolution layer
    pub upsample_rates: Vec<usize>,
    /// Kernel sizes for upsampling layers
    pub upsample_kernel_sizes: Vec<usize>,
    /// Kernel sizes for residual blocks
    pub resblock_kernel_sizes: Vec<usize>,
    /// Dilation sizes for each residual block
    pub resblock_dilation_sizes: Vec<Vec<usize>>,
    /// Number of initial channels after pre-conv
    pub initial_channel: usize,
    /// Number of mel frequency channels (input)
    pub mel_channels: usize,
    /// Sample rate in Hz
    pub sample_rate: usize,
}

impl Default for HiFiGANConfig {
    fn default() -> Self {
        Self {
            upsample_rates: vec![8, 8, 2, 2],
            upsample_kernel_sizes: vec![16, 16, 4, 4],
            resblock_kernel_sizes: vec![3, 7, 11],
            resblock_dilation_sizes: vec![vec![1, 3, 5], vec![1, 3, 5], vec![1, 3, 5]],
            initial_channel: 512,
            mel_channels: 80,
            sample_rate: 22050,
        }
    }
}

impl HiFiGANConfig {
    /// Validate configuration
    pub fn validate(&self) -> Result<()> {
        if self.upsample_rates.is_empty() {
            return Err(TtsError::InvalidConfig(
                "Upsample rates cannot be empty".to_string(),
            ));
        }
        if self.upsample_rates.len() != self.upsample_kernel_sizes.len() {
            return Err(TtsError::InvalidConfig(
                "Upsample rates and kernel sizes must have same length".to_string(),
            ));
        }
        if self.resblock_kernel_sizes.is_empty() {
            return Err(TtsError::InvalidConfig(
                "Resblock kernel sizes cannot be empty".to_string(),
            ));
        }
        if self.resblock_kernel_sizes.len() != self.resblock_dilation_sizes.len() {
            return Err(TtsError::InvalidConfig(
                "Resblock kernel sizes and dilation sizes must have same length".to_string(),
            ));
        }
        if self.initial_channel == 0 {
            return Err(TtsError::InvalidConfig(
                "Initial channel must be > 0".to_string(),
            ));
        }
        if self.mel_channels == 0 {
            return Err(TtsError::InvalidConfig(
                "Mel channels must be > 0".to_string(),
            ));
        }
        if self.sample_rate == 0 {
            return Err(TtsError::InvalidConfig(
                "Sample rate must be > 0".to_string(),
            ));
        }
        Ok(())
    }
}

/// Residual block type 1 for HiFi-GAN
#[derive(Debug)]
struct ResBlock1 {
    kernel_size: usize,
    dilations: Vec<usize>,
    device: Device,
}

impl ResBlock1 {
    fn new(
        _channels: usize,
        kernel_size: usize,
        dilations: Vec<usize>,
        device: &Device,
    ) -> Result<Self> {
        if dilations.is_empty() {
            return Err(TtsError::InvalidConfig(
                "Dilations cannot be empty".to_string(),
            ));
        }

        Ok(Self {
            kernel_size,
            dilations,
            device: device.clone(),
        })
    }

    fn forward(&self, input: &Tensor, channels: usize) -> Result<Tensor> {
        let mut output = input.clone();

        for &dilation in &self.dilations {
            // Simplified residual path: input -> conv -> LeakyReLU -> conv -> LeakyReLU
            // In full implementation, this would use proper Conv1d layers with dilation

            // Apply convolution-like transformation (simplified)
            let residual = self.apply_conv_block(&output, channels, dilation)?;

            // Residual connection
            output = output
                .add(&residual)
                .map_err(|e| TtsError::TensorError(e.to_string()))?;
        }

        Ok(output)
    }

    fn apply_conv_block(
        &self,
        input: &Tensor,
        _channels: usize,
        _dilation: usize,
    ) -> Result<Tensor> {
        // Simplified convolution block
        // In a full implementation, this would use proper Conv1d with LeakyReLU

        // For now, apply a simple transformation to maintain shape
        // This is a placeholder for the actual convolution operation
        let dims = input.dims();

        // Apply LeakyReLU-like transformation (y = max(0.2x, x))
        let data = input
            .to_cpu()
            .map_err(|e| TtsError::TensorError(e.to_string()))?;

        let transformed: Vec<f32> = data
            .iter()
            .map(|&x| if x > 0.0 { x } else { 0.2 * x })
            .collect();

        Tensor::from_data(transformed, dims.to_vec(), &self.device)
            .map_err(|e| TtsError::TensorError(e.to_string()))
    }
}

/// HiFi-GAN Generator
#[derive(Debug)]
pub struct HiFiGANGenerator {
    config: HiFiGANConfig,
    device: Device,
    training: bool,
}

impl HiFiGANGenerator {
    /// Create a new HiFi-GAN generator
    pub fn new(config: HiFiGANConfig, device: &Device) -> Result<Self> {
        config.validate()?;

        Ok(Self {
            config,
            device: device.clone(),
            training: false,
        })
    }

    /// Forward pass through the generator
    pub fn forward(&self, mel: &Tensor) -> Result<Tensor> {
        let dims = mel.dims();
        if dims.len() != 2 {
            return Err(TtsError::InvalidInput(
                "Mel spectrogram must be 2D [mel_channels, time_frames]".to_string(),
            ));
        }

        if dims[0] != self.config.mel_channels {
            return Err(TtsError::InvalidInput(format!(
                "Expected {} mel channels, got {}",
                self.config.mel_channels, dims[0]
            )));
        }

        // Reshape to [1, mel_channels, time_frames] for convolution
        let mut x = mel
            .unsqueeze(0)
            .map_err(|e| TtsError::TensorError(e.to_string()))?;

        // Pre-convolution: expand mel channels to initial_channel
        x = self.pre_conv(&x)?;

        // Upsampling blocks with MRF
        let mut current_channels = self.config.initial_channel;

        for (i, &upsample_rate) in self.config.upsample_rates.iter().enumerate() {
            // Upsampling via transposed convolution
            x = self.upsample_block(&x, current_channels, upsample_rate, i)?;

            // Multi-receptive field fusion
            x = self.mrf_block(&x, current_channels / 2)?;

            current_channels /= 2;
        }

        // Post-convolution and activation
        x = self.post_conv(&x)?;

        // Remove batch dimension and return [samples]
        x = x
            .squeeze(Some(0))
            .map_err(|e| TtsError::TensorError(e.to_string()))?;
        x = x
            .squeeze(Some(0))
            .map_err(|e| TtsError::TensorError(e.to_string()))?;

        Ok(x)
    }

    fn pre_conv(&self, input: &Tensor) -> Result<Tensor> {
        // Expand from mel_channels to initial_channel
        // In full implementation: Conv1d(mel_channels, initial_channel, kernel_size=7, padding=3)

        let dims = input.dims();
        let batch_size = dims[0];
        let time_frames = dims[2];

        // Create expanded tensor with proper shape
        let output_size = batch_size * self.config.initial_channel * time_frames;
        let data = vec![0.01f32; output_size]; // Small initialization

        Tensor::from_data(
            data,
            vec![batch_size, self.config.initial_channel, time_frames],
            &self.device,
        )
        .map_err(|e| TtsError::TensorError(e.to_string()))
    }

    fn upsample_block(
        &self,
        input: &Tensor,
        channels: usize,
        rate: usize,
        _idx: usize,
    ) -> Result<Tensor> {
        // Transposed convolution for upsampling
        // In full implementation: ConvTranspose1d with proper upsampling

        let dims = input.dims();
        let batch_size = dims[0];
        let time_frames = dims[2];
        let upsampled_frames = time_frames * rate;

        // Create upsampled tensor
        let out_channels = channels / 2;
        let output_size = batch_size * out_channels * upsampled_frames;
        let data = vec![0.01f32; output_size];

        Tensor::from_data(
            data,
            vec![batch_size, out_channels, upsampled_frames],
            &self.device,
        )
        .map_err(|e| TtsError::TensorError(e.to_string()))
    }

    fn mrf_block(&self, input: &Tensor, channels: usize) -> Result<Tensor> {
        // Multi-receptive field fusion using multiple ResBlocks
        let mut outputs = Vec::new();

        for (i, &kernel_size) in self.config.resblock_kernel_sizes.iter().enumerate() {
            let dilations = self.config.resblock_dilation_sizes[i].clone();
            let resblock = ResBlock1::new(channels, kernel_size, dilations, &self.device)?;
            let output = resblock.forward(input, channels)?;
            outputs.push(output);
        }

        // Average the outputs
        if outputs.is_empty() {
            return Ok(input.clone());
        }

        let mut result = outputs[0].clone();
        for output in outputs.iter().skip(1) {
            result = result
                .add(output)
                .map_err(|e| TtsError::TensorError(e.to_string()))?;
        }

        // Divide by number of outputs for averaging
        let scale = 1.0 / outputs.len() as f32;
        result
            .mul_scalar(scale)
            .map_err(|e| TtsError::TensorError(e.to_string()))
    }

    fn post_conv(&self, input: &Tensor) -> Result<Tensor> {
        // Final convolution to get waveform
        // In full implementation: Conv1d(channels, 1, kernel_size=7, padding=3) + Tanh

        let dims = input.dims();
        let batch_size = dims[0];
        let time_frames = dims[2];

        // Create output waveform
        let data = vec![0.0f32; batch_size * time_frames];
        let mut output = Tensor::from_data(data, vec![batch_size, 1, time_frames], &self.device)
            .map_err(|e| TtsError::TensorError(e.to_string()))?;

        // Apply tanh activation for audio range [-1, 1]
        output = output
            .tanh()
            .map_err(|e| TtsError::TensorError(e.to_string()))?;

        Ok(output)
    }

    /// Set training mode
    pub fn train(&mut self, mode: bool) {
        self.training = mode;
    }

    /// Check if in training mode
    pub fn is_training(&self) -> bool {
        self.training
    }

    /// Get configuration
    pub fn config(&self) -> &HiFiGANConfig {
        &self.config
    }
}

/// HiFi-GAN vocoder wrapper
pub struct HiFiGAN {
    generator: HiFiGANGenerator,
    config: HiFiGANConfig,
}

impl HiFiGAN {
    /// Create a new HiFi-GAN vocoder
    pub fn new(config: HiFiGANConfig, device: &Device) -> Result<Self> {
        let generator = HiFiGANGenerator::new(config.clone(), device)?;

        Ok(Self { generator, config })
    }

    /// Get the generator
    pub fn generator(&self) -> &HiFiGANGenerator {
        &self.generator
    }

    /// Get mutable generator
    pub fn generator_mut(&mut self) -> &mut HiFiGANGenerator {
        &mut self.generator
    }
}

impl Vocoder for HiFiGAN {
    fn synthesize(&self, mel: &Tensor) -> Result<Tensor> {
        self.generator.forward(mel)
    }

    fn get_sample_rate(&self) -> usize {
        self.config.sample_rate
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_hifigan_config_default() {
        let config = HiFiGANConfig::default();
        assert_eq!(config.mel_channels, 80);
        assert_eq!(config.sample_rate, 22050);
        assert_eq!(config.initial_channel, 512);
        assert_eq!(config.upsample_rates.len(), 4);
    }

    #[test]
    fn test_hifigan_config_validate_valid() {
        let config = HiFiGANConfig::default();
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_hifigan_config_validate_empty_upsample() {
        let mut config = HiFiGANConfig::default();
        config.upsample_rates = vec![];
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_validate_mismatched_upsample() {
        let mut config = HiFiGANConfig::default();
        config.upsample_rates = vec![8, 8];
        config.upsample_kernel_sizes = vec![16];
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_validate_empty_resblock() {
        let mut config = HiFiGANConfig::default();
        config.resblock_kernel_sizes = vec![];
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_validate_mismatched_resblock() {
        let mut config = HiFiGANConfig::default();
        config.resblock_kernel_sizes = vec![3, 7];
        config.resblock_dilation_sizes = vec![vec![1, 3, 5]];
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_validate_invalid_channels() {
        let mut config = HiFiGANConfig::default();
        config.initial_channel = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_validate_invalid_mel_channels() {
        let mut config = HiFiGANConfig::default();
        config.mel_channels = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_validate_invalid_sample_rate() {
        let mut config = HiFiGANConfig::default();
        config.sample_rate = 0;
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_hifigan_config_serialize() {
        let config = HiFiGANConfig::default();
        let json = serde_json::to_string(&config).unwrap();
        let deserialized: HiFiGANConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(config.mel_channels, deserialized.mel_channels);
        assert_eq!(config.sample_rate, deserialized.sample_rate);
    }

    #[test]
    fn test_resblock1_creation() {
        let device = Device::cpu();
        let dilations = vec![1, 3, 5];
        let resblock = ResBlock1::new(128, 3, dilations, &device);
        assert!(resblock.is_ok());
    }

    #[test]
    fn test_resblock1_creation_empty_dilations() {
        let device = Device::cpu();
        let dilations = vec![];
        let resblock = ResBlock1::new(128, 3, dilations, &device);
        assert!(resblock.is_err());
    }

    #[test]
    fn test_resblock1_forward() {
        let device = Device::cpu();
        let dilations = vec![1, 3];
        let resblock = ResBlock1::new(64, 3, dilations, &device).unwrap();

        let input = Tensor::randn(&[1, 64, 100], &device).unwrap();
        let output = resblock.forward(&input, 64).unwrap();

        assert_eq!(output.dims(), input.dims());
    }

    #[test]
    fn test_hifigan_generator_creation() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config.clone(), &device);
        assert!(generator.is_ok());

        let generator = generator.unwrap();
        assert_eq!(generator.config().mel_channels, config.mel_channels);
        assert!(!generator.is_training());
    }

    #[test]
    fn test_hifigan_generator_creation_invalid_config() {
        let device = Device::cpu();
        let mut config = HiFiGANConfig::default();
        config.mel_channels = 0;
        let result = HiFiGANGenerator::new(config, &device);
        assert!(result.is_err());
    }

    #[test]
    fn test_hifigan_generator_training_mode() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let mut generator = HiFiGANGenerator::new(config, &device).unwrap();

        assert!(!generator.is_training());
        generator.train(true);
        assert!(generator.is_training());
        generator.train(false);
        assert!(!generator.is_training());
    }

    #[test]
    fn test_hifigan_generator_forward_shape() {
        let device = Device::cpu();
        let config = HiFiGANConfig {
            upsample_rates: vec![4, 4],
            upsample_kernel_sizes: vec![8, 8],
            resblock_kernel_sizes: vec![3, 5],
            resblock_dilation_sizes: vec![vec![1, 2], vec![1, 2]],
            initial_channel: 128,
            mel_channels: 80,
            sample_rate: 16000,
        };
        let generator = HiFiGANGenerator::new(config.clone(), &device).unwrap();

        let mel = Tensor::randn(&[80, 50], &device).unwrap();
        let audio = generator.forward(&mel).unwrap();

        let audio_dims = audio.dims();
        assert_eq!(audio_dims.len(), 1);
        assert!(audio_dims[0] > 0);
    }

    #[test]
    fn test_hifigan_generator_forward_invalid_shape() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config, &device).unwrap();

        let mel = Tensor::randn(&[100], &device).unwrap();
        let result = generator.forward(&mel);
        assert!(result.is_err());
    }

    #[test]
    fn test_hifigan_generator_forward_wrong_mel_channels() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config, &device).unwrap();

        let mel = Tensor::randn(&[40, 50], &device).unwrap();
        let result = generator.forward(&mel);
        assert!(result.is_err());
    }

    #[test]
    fn test_hifigan_creation() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let hifigan = HiFiGAN::new(config.clone(), &device);
        assert!(hifigan.is_ok());

        let hifigan = hifigan.unwrap();
        assert_eq!(hifigan.get_sample_rate(), config.sample_rate);
    }

    #[test]
    fn test_hifigan_synthesize() {
        let device = Device::cpu();
        let config = HiFiGANConfig {
            upsample_rates: vec![4, 4],
            upsample_kernel_sizes: vec![8, 8],
            resblock_kernel_sizes: vec![3],
            resblock_dilation_sizes: vec![vec![1, 2]],
            initial_channel: 128,
            mel_channels: 40,
            sample_rate: 16000,
        };
        let hifigan = HiFiGAN::new(config, &device).unwrap();

        let mel = Tensor::randn(&[40, 30], &device).unwrap();
        let audio = hifigan.synthesize(&mel).unwrap();

        assert_eq!(audio.dims().len(), 1);
        assert!(audio.dims()[0] > 0);
    }

    #[test]
    fn test_hifigan_vocoder_trait() {
        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 hifigan = HiFiGAN::new(config, &device).unwrap();

        assert_eq!(hifigan.get_sample_rate(), 16000);

        let mel = Tensor::randn(&[40, 20], &device).unwrap();
        let audio = hifigan.synthesize(&mel).unwrap();
        assert_eq!(audio.dims().len(), 1);
    }

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

        let configs = vec![
            HiFiGANConfig {
                upsample_rates: vec![4, 4, 2],
                upsample_kernel_sizes: vec![8, 8, 4],
                resblock_kernel_sizes: vec![3, 5],
                resblock_dilation_sizes: vec![vec![1, 2], vec![1, 2]],
                initial_channel: 256,
                mel_channels: 80,
                sample_rate: 22050,
            },
            HiFiGANConfig {
                upsample_rates: vec![8, 8],
                upsample_kernel_sizes: vec![16, 16],
                resblock_kernel_sizes: vec![3, 7, 11],
                resblock_dilation_sizes: vec![vec![1, 3], vec![1, 3], vec![1, 3]],
                initial_channel: 512,
                mel_channels: 40,
                sample_rate: 16000,
            },
        ];

        for config in configs {
            let hifigan = HiFiGAN::new(config.clone(), &device).unwrap();
            let mel = Tensor::randn(&[config.mel_channels, 20], &device).unwrap();
            let audio = hifigan.synthesize(&mel).unwrap();
            assert_eq!(audio.dims().len(), 1);
        }
    }

    #[test]
    fn test_hifigan_generator_accessor() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let mut hifigan = HiFiGAN::new(config, &device).unwrap();

        assert!(!hifigan.generator().is_training());
        hifigan.generator_mut().train(true);
        assert!(hifigan.generator().is_training());
    }

    #[test]
    fn test_pre_conv_output_shape() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config.clone(), &device).unwrap();

        let input = Tensor::randn(&[1, config.mel_channels, 50], &device).unwrap();
        let output = generator.pre_conv(&input).unwrap();

        let out_dims = output.dims();
        assert_eq!(out_dims[0], 1); // batch
        assert_eq!(out_dims[1], config.initial_channel);
        assert_eq!(out_dims[2], 50); // time frames
    }

    #[test]
    fn test_upsample_block_output_shape() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 256, 50], &device).unwrap();
        let output = generator.upsample_block(&input, 256, 4, 0).unwrap();

        let out_dims = output.dims();
        assert_eq!(out_dims[0], 1);
        assert_eq!(out_dims[1], 128); // channels halved
        assert_eq!(out_dims[2], 200); // 50 * 4 upsampling
    }

    #[test]
    fn test_mrf_block_output_shape() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 128, 100], &device).unwrap();
        let output = generator.mrf_block(&input, 128).unwrap();

        assert_eq!(output.dims(), input.dims());
    }

    #[test]
    fn test_post_conv_output_shape() {
        let device = Device::cpu();
        let config = HiFiGANConfig::default();
        let generator = HiFiGANGenerator::new(config, &device).unwrap();

        let input = Tensor::randn(&[1, 64, 200], &device).unwrap();
        let output = generator.post_conv(&input).unwrap();

        let out_dims = output.dims();
        assert_eq!(out_dims[0], 1);
        assert_eq!(out_dims[1], 1); // single channel audio
        assert_eq!(out_dims[2], 200);
    }
}